Repository: pytorch
Issue: [Proposal] Use batched oprations to accelerate PowerSGD
Labels: module: performance,oncall: distributed,module: optimizer,triaged,module: ddp,module: mta

### 🚀 The feature, motivation and pitch

The current PowerSGD implementation iterate over all gradient tensors, Ps, Qs, and operate on single tensor as follows:

```python
# Compute Ps.
for tensor, q, p in zip(tensors_to_compress, qs, ps):
    torch.matmul(tensor, q, out=p)
```

```python
for p in ps:
    _orthogonalize(p, state.orthogonalization_epsilon)
```

```python
# Compute Qs.
for tensor, p, q in zip(tensors_to_compress, ps, qs):
    torch.matmul(tensor.t(), p, out=q)
```

```python
for p, q, tensor in zip(ps, qs, tensors_to_compress):
    torch.matmul(p, q.t(), out=tensor)
```

These operations on single tensors may not utilize the GPU due to low parallelism. We have implemented vanilla PowerSGD in `megatron` framework with its own `LocalDDP`, and found GPU usage during PowerSGD compression / decompression is only 44% in T5-11B training with PowerSGD rank=32.

In most models however, the architecture of each layer is the same, therefore the shapes of gradient tensors are the same. The compression and decompression of gradients of the same shape can be batched, replacing `matmul` above with `bmm` and `orthogonalize` with its batched alternative, which can further improve parallelism and better utilize GPU. By putting gradients of all layers in a batched operation, we have decreased the total compression + communication time by 62.28% compared to vanilla PowerSGD. Our experiment is done with 2-node T5-11B training on low-bandwidth network (100Gb ethernet). The detailed statistics are listed below:

| PowerSGD | Iteration time | Compression + communication time |
|:-:|:-:|:-:|
| None | 7927.6ms | 7504.9ms |
| Vanilla | 2422.1ms | 1988.3ms |
| Batched | 1201.2ms | 750.0ms |

(I have noticed there is a `batched_powerSGD_hook` API which compresses all gradients as a single tensor, but I believe it is not the same thing as I say)

### Alternatives

While constructing the P and Q buffers, we can put gradients of the same shape together:

```python
shape2tensors = {}
for tensor in tensors_to_compress:
    if tensor.shape not in shape2tensors:
        shape2tensors[tensor.shape] = []
    shape2tensors[tensor.shape].append(tensor)

batched_tensors = []
batched_ps = []
batched_qs = []
p_idx = 0
q_idx = 0
for shape, tensors in shape2tensors.items():
    batch_size = len(tensors)
    n, m = shape
    matrix_approximation_rank = min(n, m, state.matrix_approximation_rank)
    batched_tensors.append(
        torch.stack(tensors)
    )
    batched_ps.append(
        state.p_memory_dict[bucket_index][
            p_idx : p_idx + batch_size * n * matrix_approximation_rank
        ].view(batch_size, n, matrix_approximation_rank)
    )
    batched_qs.append(
        state.q_memory_dict[bucket_index][
            q_idx : q_idx + batch_size * m * matrix_approximation_rank
        ].view(batch_size, m, matrix_approximation_rank)
    )
    p_idx += batch_size * n * matrix_approximation_rank
    q_idx += batch_size * m * matrix_approximation_rank
```

And use batched operations for compression / decompression:

```python
def _batch_orthogonalize(matrix, epsilon=1e-8):
    num_cols = matrix.shape[2]
    for i in range(num_cols):
        col = matrix[:, :, i:i+1]
        col /= torch.norm(col, dim=1, keepdim=True) + epsilon
        if i + 1 < num_cols:
            rest = matrix[:, :, i+1:]
            rest -= torch.sum(col * rest, dim=1, keepdim=True) * col
```

```python
# Compute Ps.
for batched_tensor, batched_q, batched_p in zip(batched_tensors, batched_qs, batched_ps):
    torch.bmm(batched_tensor, batched_q, out=batched_p)
```

```python
for p in ps:
    _batch_orthogonalize(p, state.orthogonalization_epsilon)
```

```python
# Compute Qs.
for batched_tensor, batched_p, batched_q in zip(batched_tensors, batched_ps, batched_qs):
    torch.bmm(batched_tensor.transpose(1, 2), batched_p, out=batched_q)
```

```python
for batched_p, batched_a, batched_tensor in zip(batched_ps, batched_qs, batched_tensors):
    torch.bmm(batched_p, batched_q.transpose(1, 2), out=batched_tensor)
```

### Additional context

Our experiment is done by migrating the PowerSGD implementation from PyTorch to `LocalDDP` in `megatron`. One difference is that PyTorch DDP puts gradient tensors into buckets before allreduce, but in `megatron` DDP we just operate on the original gradient tensors. The bucketing strategy may weaken the effect of batching because gradients of the same shape may be in different buckets. I believe the same result can be achieved by setting bucket size to a very large value so all gradients will be in one bucket, but the implementation detail is worth considering.

cc @VitalyFedyunin @ngimel @pietern @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @SciPioneer @H-Huang @vincentqb @jbschlosser @albanD @crcrpar @mcarilli