Repository: mosaicml/composer

Issue:
`DecoupledAdamW`'s `all_reduce` operation for `OptimizerMonitor` returns wrong results
I've posted the issue in `llm-foundry` repo but after looking at it decided to post it here to as the problem is mostly with the `composer` work and code. The link to issue with the examples and environment output is below.

[`llm-foudnry` issue](https://github.com/mosaicml/llm-foundry/issues/578)

---

While using the `llm-foundry` code to pretrain the model the issue with the `OptimizerMonitor` logged values was found. As it turns out the issue is actually caused by the wrong results of `all_reduce` operation that is used to collect the results from all ranks.

We used `llm-foundry` to pretrain the model. To begin with and reproduce results I used provided MPT 1B config that was changed to use custom tokenizer and `SHARD_GRAD_OP` FSDP strategy with mixed precision (bf16 for forward pass, fp32 for others). Values logged by `OptimizerMonitor` had "jumps" after resuming training both using `autoresume` option and continuing from specific checkpoint. It turns out that the final value that is obtained by `pre_reduce_metrics` and `dist_reduce_metrics` provided by optimizer (using `DecoupledAdamW`, which is default in config) is different across ranks.

After narrowing it down it was found that the actual `all_reduce` operation to collect metric values from all ranks provides wrong results ([link to line](https://github.com/mosaicml/composer/blob/dev/composer/optim/decoupled_weight_decay.py#L353)). It is said that the result is "bitwise identical on all ranks" but it not. During exploration I used the L2 norm of  the weight of the specific attention matrix. Using 8 GPUs and `SHARD_GRAD_OP` FSDP strategy the param's values are present on ranks 0 and 1. After `all_reduce` operation their norms are almost always does not change, 0 values on other ranks can change a little and become very small number that is not 0, but can also become just some random number. The `all_gather` operation returns bad results too that correlate with the results obtained by `all_reduce` operation.

At the same time initializing test variable with the rank number as a value in the same code block and then applying `all_reduce` operation to it returns expected correct result every time. So it seems that the problem is with this particular variable. Also using regular DDP provides expected and correct results, i.e. values across all ranks are the same and correct.

**To reproduce**

Steps to reproduce the behavior:

1. Clone `composer` (dev branch, the default one) from the GitHub repo and install it using -e option to be able to modify it.
2. Clone `llm-foundry` repo.
3. Use MPT 1B config, C4 subset from the README, SHARD_GRAD_OP strategy for FSDP config, bf16 mixed precision for forward, fp32 for buffer and reduce.
4. Print values for block 12 attention matrix before and after `all_reduce` operation in `composer`'s optimizer to get logged L2 norms (NAME).
5. Run training with `scripts/train/train.py`.

## Expected behavior

Parameter's (and gradients) L2 value that is the same across all ranks and that is the same with the one that can be obtained from `state_dict`. To be specific, `all_reduce` operation for `reduce` variable in `DecoupledAdamW` optimizer in `dist_reduce_metric` function should result in one right value across all ranks. Instead each rank has its own value.

## Additional context

Layer name: `model._fsdp_wrapped_module.transformer.blocks.12._fsdp_wrapped_module.attn.Wqkv.weight`.

Code for optimizer class to print values: 

```python
        all_keys = sorted(all_keys, key=lambda metric: 0 if 'l2_norm' in metric else 1)
        for metric in all_keys:
            if metric.startswith('l2_norm'):
                dist.barrier()

                tensors_gather_list = [
                    torch.tensor(0.0, dtype=torch.float32, device=torch.cuda.current_device())
                    for _ in range(dist.get_world_size())
                ]
                # tensors_gather_list = None

                test_reduce = torch.tensor(rank, dtype=torch.float32, device=torch.cuda.current_device())

                if metric in optimizer_metrics:
                    reduced_l2 = torch.tensor(optimizer_metrics[metric].item(), dtype=torch.float32, device=torch.cuda.current_device())
                else:
                    reduced_l2 = torch.tensor(0.0, dtype=torch.float32, device=torch.cuda.current_device())

                # old way
                # reduced = optimizer_metrics.get(metric, torch.tensor(0.0, device=torch.cuda.current_device()))

                # ------------------- DEBUG PRINTS -------------------
                if metric == opt_metrics_name:
                    print(f'rank {rank} | dist_reduce_metrics initial reduced value: {reduced_l2}')
                    print(f'rank {rank} | dist_reduce_metrics initial reduced value device: {reduced_l2.device}')

                    print(f'rank {rank} | test_reduce value: {test_reduce}')
                    print(f'rank {rank} | test_reduce value device: {test_reduce.device}')

                    print('--')
                # ------------------- DEBUG PRINTS -------------------

                if dist.get_world_size() > 1:
                    dist.barrier()

                    # tensors_gather_list = dist.all_gather(test_reduce)
                    torch.distributed.all_gather(tensors_gather_list, reduced_l2)

                    # torch.distributed.reduce(reduced_l2, dst=0)
                    dist.all_reduce(reduced_l2, reduce_operation='SUM')


                    dist.all_reduce(test_reduce, reduce_operation='SUM')

                dist.barrier()

                # ------------------- DEBUG PRINTS -------------------
                if metric == opt_metrics_name:
                    print(f'rank {rank} | tensors_gather_list: {tensors_gather_list}')

                    print(f'rank {rank} | dist_reduce_metrics reduced value: {reduced_l2}')
                    print(f'rank {rank} | dist_reduce_metrics reduced value sq root: {math.sqrt(reduced_l2)}')

                    print(f'rank {rank} | test_reduce value reduced: {test_reduce}')
                    print(f'rank {rank} | test_reduce value reduced device: {test_reduce.device}')

                    print('--------------------------------------------------------')
                # ------------------- DEBUG PRINTS -------------------

                dist.barrier()

                optimizer_metrics[metric] = math.sqrt(reduced_l2)
```

Example output:

![example_output](https://github.com/mosaicml/composer/assets/30723114/fff42994-223b-474c-9279-8e67e37418db)

![example_output_v2](https://github.com/mosaicml/composer/assets/30723114/3039cb4b-4e15-4198-90e3-bb156aec7594)


Discussion:
@bmosaicml can you please take a look? 
A colleague's pair of eyes helped yesterday! Found the problem in key sorting in [here](https://github.com/mosaicml/composer/blob/dev/composer/optim/decoupled_weight_decay.py#L348). 

```python
all_keys = sorted(all_keys, key=lambda metric: 0 if 'l2_norm' in metric else 1)
```

Block of keys that start with "l2_norm" are put in the beginning and others come after them but the keys itself are not sorted. So although each rank's keys start with "l2_norm" they have different order within that block. And all of a sudden we sum parameter norm with gradient and moment norm for different layers.

Fixed it by applying just a regular sort:

```python
all_keys = sorted(all_keys)
```
We can also get rid of the if blocks that calculate metrics that start with "cosine" as these metrics are no longer calculated (`DecoupledAdamW.metric_function` do not have that anymore). This is the reason we can apply just a regular `sorted()` function without specifying `key` for a specific sort.

I could provide pull request with the changes if you want.
Great catch! Yes we would love a community PR. @bmosaicml and I can review it -- please tag us reviewers. Thank you!
Ok, I'll provide a request as soon as possible then!

Resolving patch:
diff --git a/composer/optim/decoupled_weight_decay.py b/composer/optim/decoupled_weight_decay.py
--- a/composer/optim/decoupled_weight_decay.py
+++ b/composer/optim/decoupled_weight_decay.py
@@ -345,7 +345,9 @@ def dist_reduce_metrics(self, optimizer_metrics):
         for keys in all_gathered_keys:
             all_keys.update(keys)
 
-        all_keys = sorted(all_keys, key=lambda metric: 0 if 'l2_norm' in metric else 1)
+        # Sort keys to ensure every rank has the same keys order
+        # Only L2 norm metric keys are present, can apply regular sort
+        all_keys = sorted(all_keys)
         for metric in all_keys:
             if metric.startswith('l2_norm'):
                 reduced = optimizer_metrics.get(metric, torch.tensor(0.0, device=torch.cuda.current_device()))
@@ -353,18 +355,6 @@ def dist_reduce_metrics(self, optimizer_metrics):
                     dist.all_reduce(reduced, reduce_operation='SUM')
 
                 optimizer_metrics[metric] = math.sqrt(reduced)
-            elif metric.startswith('cosine'):
-                reduced = optimizer_metrics.get(metric, torch.tensor(0.0, device=torch.cuda.current_device()))
-                if dist.get_world_size() > 1:
-                    dist.all_reduce(reduced, reduce_operation='SUM')
-
-                _, vectors, layer = tuple(metric.split('/'))
-
-                A, B = tuple(vectors.split('_'))
-
-                A_reduced_norm = optimizer_metrics[f'l2_norm/{A}/{layer}']
-                B_reduced_norm = optimizer_metrics[f'l2_norm/{B}/{layer}']
-                optimizer_metrics[metric] = reduced / (A_reduced_norm * B_reduced_norm)
             else:
                 reduced = optimizer_metrics.get(metric, torch.tensor(0.0, device=torch.cuda.current_device()))
                 if dist.get_world_size() > 1:
@@ -375,23 +365,10 @@ def dist_reduce_metrics(self, optimizer_metrics):
 
     def pre_reduce_metrics(self, optimizer_metrics):
         """Preprocess metrics to reduce across ranks correctly."""
-        # Sort L2 norms first so they are squared before other metrics, which depend on squared values
-        metrics = optimizer_metrics.keys()
-        metrics = sorted(metrics, key=lambda metric: 0 if 'l2_norm' in metric else 1)
-        for metric in metrics:
-            if metric.startswith('l2_norm'):
-                # L2 norms need to be squared, before they are reduced via summation
-                optimizer_metrics[metric] = optimizer_metrics[metric]**2
-            elif metric.startswith('cosine'):
-                _, vectors, layer = tuple(metric.split('/'))
-
-                A, B = tuple(vectors.split('_'))
-
-                # L2 norm would've been squared in previous branch
-                A_rank_subset_norm = math.sqrt(optimizer_metrics[f'l2_norm/{A}/{layer}'])
-                B_rank_subset_norm = math.sqrt(optimizer_metrics[f'l2_norm/{B}/{layer}'])
-
-                optimizer_metrics[metric] *= A_rank_subset_norm * B_rank_subset_norm
+        # Only L2 norm metric keys are present, can skip sorting at this stage
+        for metric in optimizer_metrics:
+            # L2 norms need to be squared, before they are reduced via summation
+            optimizer_metrics[metric] = optimizer_metrics[metric]**2
 
         return optimizer_metrics
 


Test patch: