Repository: pytorch
Issue: Adafactor foreach impl performance tracker
Labels: module: performance,module: optimizer,triaged,actionable,module: mta

### 🚀 The feature, motivation and pitch

Currently, Adafactor has both foreach and forloop implementations, but the performance of the foreach impl could be significantly improved with the following added foreach ops:

A. 
- [ ] `_foreach_mm(xs, ys)`. 
This is definitely the most challenging to implement well as it would be the first non pointwise/reduction foreach op. It is also unclear how to horizontally fuse multiple mms with different dimensions. Don't worry, the rest of the list is much more actionable.
```
var_estimates = [
                row_var @ col_var 
                for row_var, col_var in zip(device_row_vars, device_col_vars)
            ]
# becomes
var_estimates = torch._foreach_mm(device_row_vars, device_col_vars)
```

B.
- [ ] `_foreach_norm(tensors, dim=None, keepdim=False)`
We have a `_foreach_norm` today that reduces to Scalars. We can assume for the purposes of this signature that every tensor in `tensors` can be reduced along `dim`.
```
row_means = [torch.norm(grad, dim=-1, keepdim=True) for grad in device_grads]
...
col_means = [torch.norm(grad, dim=-2, keepdim=True) for grad in device_grads]
# becomes
row_means = torch._foreach_norm(device_grads, dim=-1, keepdim=True]
...
col_means = torch._foreach_norm(device_grads, dim=-2, keepdim=True]
```

C. 
- [ ] `_foreach_rsqrt(tensors)`
Self explanatory. This may be the easiest AI from this list.
```
torch._foreach_sqrt_(var_estimates)
torch._foreach_reciprocal_(var_estimates)
# becomes
torch._foreach_rsqrt_(var_estimates)
```

D. 
- [ ] `_foreach_lerp(xs: TensorList, ys: TensorList, weight: ScalarList)`
This requires adding a new overload that takes ScalarList and not just Scalar. This would be highest impact to effort value as we'd be going from 9 kernel launches to 3.
```
torch._foreach_mul_(device_row_vars, beta2_ts)
torch._foreach_mul_(row_means, one_minus_beta2_ts)
torch._foreach_add_(device_row_vars, row_means)
...
torch._foreach_mul_(device_col_vars, beta2_ts)  # type: ignore[arg-type]
torch._foreach_mul_(col_means, one_minus_beta2_ts)
torch._foreach_add_(device_col_vars, col_means)  # type: ignore[arg-type]
...
torch._foreach_mul_(device_variances, beta2_ts)
torch._foreach_mul_(grads_squared, one_minus_beta2_ts)
torch._foreach_add_(device_variances, grads_squared)
...

# becomes
torch._foreach_lerp_(device_row_vars, row_means, one_minus_beta2_ts)
...
torch._foreach_lerp_(device_col_vars, col_means, one_minus_beta2_ts)
...
torch._foreach_lerp_(device_variances, grads_squared, one_minus_beta2_ts)
``` 

E.
- [ ] `_foreach_mean(tensors, dim=None, keepdim=False)`
Similar as to B above with foreach_norm, but would require writing a new op (versus just adding a schema).
```
row_var_means = [row_var.mean(dim=-2, keepdim=True) for row_var in device_row_vars]
# becomes
row_var_means = torch._foreach_mean(device_row_vars, dim=-2, keepdim=True)
```

F. 
- [ ] `_foreach_add(xs: TensorList, ys: TensorList, alpha: ScalarList)`
Similar to D, where we wish alphas could be ScalarList instead of just Scalar. I think @ad8e mentioned wanting this in another issue.
```
torch._foreach_mul_(updates, alphas)
torch._foreach_add_(device_params, updates)
# becomes
torch._foreach_add_(device_params, updates, alphas)
```


### Alternatives

Let foreach Adafactor be slow.

Figure out how to autogenerate foreach kernels.

### Additional context

_No response_

cc @msaroufim @vincentqb @jbschlosser @albanD @crcrpar @mcarilli