Repository: facebookresearch/fairscale

Issue:
Zero sharding in Fairseq affects Megatron-LM results
## 🐛 Bug

Optimizer zero state sharding should not affect the results of any experiments. However, @ngoyal2707 and I have observed that this isn't the case for Megatron-LM models in Fairseq.

## Command

## To Reproduce

First, we show that results can replicate across model parallel sizes.

```
python fairseq_train.py --task masked_lm   /checkpoint/bioseq_nonsecure/namangoyal/model_parallel_data/small_sample_valid_ur50-bin --dataset-impl fasta  --save-dir checkpoints/zero-0-mp-1    --dropout 0.1   --optimizer adam --adam-betas '(0.9, 0.98)' --weight-decay 0.01 --clip-norm 0.0   --lr 0.0005 --lr-scheduler inverse_sqrt --warmup-updates 4000 --warmup-init-lr 1e-07   --tokens-per-sample 128 --sample-break-mode none   --max-tokens 128 --memory-efficient-fp16 --no-progress-bar --log-interval 1 --seed 4 --max-epoch 1 --max-update 50 --encoder-layers 4 --no-save  --arch model_parallel_roberta_large --model-parallel-size 2 --update-freq 2 2>&- | grep "ppl"

2020-09-25 13:32:52 | INFO | train_inner | epoch 001:      5 / 8323965 loss=0.966, ppl=1.95, wps=0, ups=0, wpb=1024, bsz=8, num_updates=1, lr=2.24975e-07, gnorm=0.055, loss_scale=8, train_wall=0, wall=203
2020-09-25 13:32:53 | INFO | train_inner | epoch 001:      6 / 8323965 loss=0.99, ppl=1.99, wps=10858.5, ups=10.6, wpb=1024, bsz=8, num_updates=2, lr=3.4995e-07, gnorm=0.07, loss_scale=8, train_wall=0, wall=203
2020-09-25 13:32:53 | INFO | train_inner | epoch 001:      7 / 8323965 loss=0.968, ppl=1.96, wps=14452.4, ups=14.1, wpb=1024, bsz=8, num_updates=3, lr=4.74925e-07, gnorm=0.067, loss_scale=8, train_wall=0, wall=203
2020-09-25 13:32:53 | INFO | train_inner | epoch 001:      8 / 8323965 loss=1.032, ppl=2.04, wps=15902.2, ups=15.51, wpb=1024, bsz=8, num_updates=4, lr=5.999e-07, gnorm=0.073, loss_scale=8, train_wall=0, wall=203
2020-09-25 13:32:53 | INFO | train_inner | epoch 001:      9 / 8323965 loss=1.007, ppl=2.01, wps=14162.8, ups=13.81, wpb=1024, bsz=8, num_updates=5, lr=7.24875e-07, gnorm=0.089, loss_scale=8, train_wall=0, wall=203
2020-09-25 13:32:53 | INFO | train_inner | epoch 001:     10 / 8323965 loss=1.03, ppl=2.04, wps=14513.6, ups=14.15, wpb=1024, bsz=8, num_updates=6, lr=8.4985e-07, gnorm=0.082, loss_scale=8, train_wall=0, wall=204

python fairseq_train.py --task masked_lm   /checkpoint/bioseq_nonsecure/namangoyal/model_parallel_data/small_sample_valid_ur50-bin --dataset-impl fasta  --save-dir checkpoints/zero-0-mp-1    --dropout 0.1   --optimizer adam --adam-betas '(0.9, 0.98)' --weight-decay 0.01 --clip-norm 0.0   --lr 0.0005 --lr-scheduler inverse_sqrt --warmup-updates 4000 --warmup-init-lr 1e-07   --tokens-per-sample 128 --sample-break-mode none   --max-tokens 128 --memory-efficient-fp16 --no-progress-bar --log-interval 1 --seed 4 --max-epoch 1 --max-update 50 --encoder-layers 4 --no-save  --arch model_parallel_roberta_large --model-parallel-size 4 --update-freq 4 2>&- | grep "ppl"

2020-09-25 14:44:30 | INFO | train_inner | epoch 001:      5 / 8323965 loss=0.966, ppl=1.95, wps=0, ups=0, wpb=1024, bsz=8, num_updates=1, lr=2.24975e-07, gnorm=0.055, loss_scale=8, train_wall=0, wall=212
2020-09-25 14:44:30 | INFO | train_inner | epoch 001:      6 / 8323965 loss=0.99, ppl=1.99, wps=3017.4, ups=2.95, wpb=1024, bsz=8, num_updates=2, lr=3.4995e-07, gnorm=0.07, loss_scale=8, train_wall=0, wall=212
2020-09-25 14:44:30 | INFO | train_inner | epoch 001:      7 / 8323965 loss=0.968, ppl=1.96, wps=2964.3, ups=2.89, wpb=1024, bsz=8, num_updates=3, lr=4.74925e-07, gnorm=0.067, loss_scale=8, train_wall=0, wall=212
2020-09-25 14:44:31 | INFO | train_inner | epoch 001:      8 / 8323965 loss=1.032, ppl=2.04, wps=3059.4, ups=2.99, wpb=1024, bsz=8, num_updates=4, lr=5.999e-07, gnorm=0.073, loss_scale=8, train_wall=0, wall=213
2020-09-25 14:44:31 | INFO | train_inner | epoch 001:      9 / 8323965 loss=1.007, ppl=2.01, wps=2869.3, ups=2.8, wpb=1024, bsz=8, num_updates=5, lr=7.24875e-07, gnorm=0.089, loss_scale=8, train_wall=0, wall=213
```

However, if we now add optimizer zero state sharding, the numbers change. 

```
python fairseq_train.py --task masked_lm   /checkpoint/bioseq_nonsecure/namangoyal/model_parallel_data/small_sample_valid_ur50-bin --dataset-impl fasta  --save-dir checkpoints/zero-0-mp-1    --dropout 0.1   --optimizer adam --adam-betas '(0.9, 0.98)' --weight-decay 0.01 --clip-norm 0.0   --lr 0.0005 --lr-scheduler inverse_sqrt --warmup-updates 4000 --warmup-init-lr 1e-07   --tokens-per-sample 128 --sample-break-mode none   --max-tokens 128 --memory-efficient-fp16 --no-progress-bar --log-interval 1 --seed 4 --max-epoch 1 --max-update 50 --encoder-layers 4 --no-save  --arch model_parallel_roberta_large --model-parallel-size 2 --update-freq 2 --zero-sharding os 2>&- | grep "ppl"

020-09-25 15:15:17 | INFO | train_inner | epoch 001:      5 / 8323965 loss=0.966, ppl=1.95, wps=0, ups=0, wpb=1024, bsz=8, num_updates=1, lr=2.24975e-07, gnorm=0.055, loss_scale=8, train_wall=0, wall=199
2020-09-25 15:15:17 | INFO | train_inner | epoch 001:      6 / 8323965 loss=0.968, ppl=1.96, wps=10795.5, ups=10.54, wpb=1024, bsz=8, num_updates=2, lr=3.4995e-07, gnorm=0.06, loss_scale=8, train_wall=0, wall=199
2020-09-25 15:15:17 | INFO | train_inner | epoch 001:      7 / 8323965 loss=0.953, ppl=1.94, wps=16335, ups=15.94, wpb=1024, bsz=8, num_updates=3, lr=4.74925e-07, gnorm=0.064, loss_scale=8, train_wall=0, wall=199
2020-09-25 15:15:17 | INFO | train_inner | epoch 001:      8 / 8323965 loss=0.994, ppl=1.99, wps=15692, ups=15.31, wpb=1024, bsz=8, num_updates=4, lr=5.999e-07, gnorm=0.065, loss_scale=8, train_wall=0, wall=199
2020-09-25 15:15:18 | INFO | train_inner | epoch 001:      9 / 8323965 loss=0.963, ppl=1.95, wps=17342.3, ups=16.92, wpb=1024, bsz=8, num_updates=5, lr=7.24875e-07, gnorm=0.09, loss_scale=8, train_wall=0, wall=199
2020-09-25 15:15:18 | INFO | train_inner | epoch 001:     10 / 8323965 loss=0.972, ppl=1.96, wps=16860.9, ups=16.45, wpb=1024, bsz=8, num_updates=6, lr=8.4985e-07, gnorm=0.084, loss_scale=8, train_wall=0, wall=199
```

## Expected behavior

We expect the results to be independent of the `zero-sharding` parameter. 

## Environment

Please copy and paste the output from the
environment collection script from PyTorch
(or fill out the checklist below manually).

```
Collecting environment information...
PyTorch version: 1.5.0a0+4ff3872
Is debug build: No
CUDA used to build PyTorch: 10.1

OS: Ubuntu 18.04.3 LTS
GCC version: (Ubuntu 7.4.0-1ubuntu1~18.04.1) 7.4.0
CMake version: version 3.10.2

Python version: 3.6
Is CUDA available: Yes
CUDA runtime version: 10.1.105
GPU models and configuration:
GPU 0: Quadro GP100
GPU 1: Quadro GP100

Nvidia driver version: 418.116.00
cuDNN version: Could not collect

Versions of relevant libraries:
[pip] msgpack-numpy==0.4.5
[pip] numpy==1.18.3
[pip] numpydoc==0.9.2
[pip] pytorch-lightning==0.8.1
[pip] pytorch-pretrained-bert==0.6.2
[pip] pytorch-transformers==1.1.0
[pip] torch==1.5.0a0+4ff3872
[conda] blas                      1.0                         mkl
[conda] libblas                   3.8.0                    15_mkl    conda-forge
[conda] libcblas                  3.8.0                    15_mkl    conda-forge
[conda] liblapack                 3.8.0                    15_mkl    conda-forge
[conda] magma-cuda101             2.5.2                         1    pytorch
[conda] mkl                       2020.1                      217
[conda] mkl-include               2020.0                      166
[conda] mkl-service               2.3.0            py36he904b0f_0
[conda] mkl_fft                   1.0.15           py36ha843d7b_0
[conda] mkl_random                1.1.0            py36hd6b4f25_0
[conda] pytorch-lightning         0.8.1                     <pip>
```


Discussion:
I've no idea how the `zero-sharding` parameter is implemented in fairseq_train.py, it's outside of the fairscale scope, right ? Shouldn't this be reported to fairseq instead ? (which could in turn report an issue here if need be)
I was thinking about that too -- this is probably an issue between Fairscale <> Megatron, so it's not clear what is the best home for the issue. As long as someone on the Fairscale team sees this (probably @msbaines?) we're good. I opened an[ issue on Fairseq](https://github.com/fairinternal/fairseq-py/issues/1299) as well for visibility. 
for context, I tracked the link down to https://github.com/pytorch/fairseq/blob/master/fairseq/optim/shard.py
Could be that we should add some fairseq specific workloads here to make sure that there's no regression over time. We have an ongoing test in CI which checks that zero does not affect the results, but it's not using mixed precision for instance, or whatever other trick fairseq may be leveraging
by any chance @joshim5 could you try without `--memory-efficient-fp16` ? It's not something that I'm very familiar with, if this is the culprit that would guide where to focus on
Thanks much for reporting and the repro. This is likely an issues in fairscale's OSS and its interaction with MegaTron-LM. I can't think of an obvious reason why it shouldn't work but will look into it.

I did my initial testing with --memory-efficient-fp16 so I don't think that is the issue here.
--zero-sharding in fairseq shards across ranks of the main distributed group. This will need to be adjusted for megatron where there are data_parallel and model_parallel groups.
that's something that we discussed with @froody also, happy to own that if you want @msbaines, as you prefer, and thanks for the initial investigation !
I'm working on a fix. It mostly requires changes to fairscale with a small change to fairseq as well (passing the optimizer groups). I'll try to share PRs + repro tomorrow.
I think the following should fix it but I'm struggling to run the repro after re-syncing to fairseq master:

```
diff --git a/fairseq/optim/shard.py b/fairseq/optim/shard.py
index 4f35dbda4..5646b777b 100644
--- a/fairseq/optim/shard.py
+++ b/fairseq/optim/shard.py
@@ -11,7 +11,7 @@ except ImportError:
     _has_fairscale = False


-def shard_(args, optimizer):
+def shard_(args, optimizer, group):
     if not _has_fairscale:
         raise ImportError(
             '\n\nPlease install the fairscale package:'
@@ -30,4 +30,4 @@ def shard_(args, optimizer):

     torch_optimizer = optimizer.optimizer
     optim_cls = type(torch_optimizer)
-    optimizer.optimizer = FairseqOSS(torch_optimizer.param_groups, optim_cls, **optimizer.optimizer_config)
+    optimizer.optimizer = FairseqOSS(torch_optimizer.param_groups, optim_cls, group, **optimizer.optimizer_config)
diff --git a/fairseq/trainer.py b/fairseq/trainer.py
index 583d37c25..b19d2a8e9 100644
--- a/fairseq/trainer.py
+++ b/fairseq/trainer.py
@@ -229,7 +229,7 @@ class Trainer(object):
                         "Please use --fp16-no-flatten-grads"
                 )
             else:
-                optim.shard_(self.args, self._optimizer)
+                optim.shard_(self.args, self._optimizer, self.data_parallel_process_group)

         # We should initialize the learning rate scheduler immediately after
         # building the optimizer, so that the initial learning rate is set.
```
(Already sent this to Mandeep over chat): There are also some Fairscale changes that need to be introduced. In my diff, we separately keep track of mp and dp groups, performing dp-OSS separately across mp groups. 

Most of these changes will live in Fairscale, so this implementation will generalize to any model parallel case. On the Fairseq side, we just have to modify Fairseq to pass the distributed process groups to the Fairscale OSS optimizer.

Resolving patch:
diff --git a/fairscale/optim/oss.py b/fairscale/optim/oss.py
--- a/fairscale/optim/oss.py
+++ b/fairscale/optim/oss.py
@@ -72,6 +72,8 @@ def __init__(self, params: _params_t, optim: Type[Optimizer] = SGD, group: Optio
         self.world_size = dist.get_world_size(self.group)
 
         self.rank = dist.get_rank(self.group)
+        self.global_rank = self.get_global_rank(self.group, self.rank)
+
         self.optim = optim(self.partition_parameters()[self.rank], **default)
 
         #  Optional consolidated optimizer state
@@ -88,7 +90,7 @@ def __init__(self, params: _params_t, optim: Type[Optimizer] = SGD, group: Optio
 
     # Partition helpers
     def partition_parameters(self) -> List[List[dict]]:
-        """Partitions parameters across distributed ranks.
+        """Partitions parameters across distributed data parallel ranks.
 
         Returns a list of param_groups (which is a list of dict) where each
         element of the list contains the param_groups for a rank. Element 0
@@ -135,6 +137,7 @@ def per_device_params(self) -> List[List[Parameter]]:
 
     @property
     def param_to_rank(self) -> Dict[torch.Tensor, int]:
+        '''param to data parallel rank'''
         if len(self._param_rank) == 0:
             for rank, param_groups in enumerate(self.partition_parameters()):
                 for param_group in param_groups:
@@ -142,6 +145,13 @@ def param_to_rank(self) -> Dict[torch.Tensor, int]:
                         self._param_rank[param] = rank
         return self._param_rank
 
+    def get_global_rank(self, group, rank):
+        if group is dist.group.WORLD:
+            return rank
+        else:
+            global_rank = dist.distributed_c10d._get_global_rank(group, rank)
+        return global_rank
+
     # NOTE(msb) We add a kwargs in order to support Optimizer sub-classes that support extra kwargs.
     # For example, the apex library contains fused optimizers with a step that supports extra kwargs.
     def step(self, closure: Optional[Callable[[], float]] = None, **kwargs: Any) -> Optional[float]:
@@ -174,14 +184,15 @@ def step(self, closure: Optional[Callable[[], float]] = None, **kwargs: Any) ->
                     # Gloo will rightly assert on this operation for any tensor that requires grad.
                     # We save and restore the grad requirement state to work around that, in our case
                     # the grad is only useful on the source rank.
+                    global_rank = self.get_global_rank(self.group, rank)
+
                     requires_grad.append((param, param.requires_grad))
                     param.requires_grad = False
-                    requests.append(dist.broadcast(tensor=param, src=rank, group=self.group, async_op=True))
+                    requests.append(dist.broadcast(tensor=param, src=global_rank, group=self.group, async_op=True))
 
         for fut, req_grad in zip(requests, requires_grad):
             fut.wait()
             req_grad[0].requires_grad = req_grad[1]
-
         return loss
 
     def local_state_dict(self) -> dict:
@@ -330,7 +341,7 @@ def _collect_sharded_states(self) -> List[Dict[str, Any]]:
         empty_buffer = torch.tensor([0], dtype=torch.uint8, device=self._device)
         all_states: List[Dict[str, Any]] = []
 
-        for rank in range(dist.get_world_size(group=self.group)):
+        for rank in range(self.world_size):
             if rank == self.rank:
                 logging.debug("Saving self state")
                 all_states.append(
@@ -338,12 +349,13 @@ def _collect_sharded_states(self) -> List[Dict[str, Any]]:
                 )
 
                 # Sync with other replicas
-                broadcast_object(empty_buffer, src_rank=rank, group=self.group, dist_device=self._device)
+                broadcast_object(empty_buffer, src_rank=self.global_rank, group=self.group, dist_device=self._device)
             else:
                 # Fetch the optim state from the other replicas
-                logging.debug("Receiving state from rank %s ", rank)
+                global_rank = self.get_global_rank(self.group, rank)
+                logging.debug("Receiving state from rank %s ", global_rank)
                 replica_state = broadcast_object(
-                    empty_buffer, src_rank=rank, group=self.group, dist_device=self._device
+                    empty_buffer, src_rank=global_rank, group=self.group, dist_device=self._device
                 )
 
                 all_states.append(
@@ -358,17 +370,18 @@ def _broadcast_state_dict(self) -> None:
         """Broadcast this rank's state shard, discard others"""
         empty_buffer = torch.tensor([0], dtype=torch.uint8, device=self._device)
 
-        for rank in range(dist.get_world_size(group=self.group)):
+        for rank in range(self.world_size):
             if rank == self.rank:
                 # Send the state to the reference replica
                 logging.debug(
                     "Sending the sharded optimizer state to the reference replica from rank %s", rank,
                 )
-                broadcast_object(self.local_state_dict(), src_rank=rank, group=self.group, dist_device=self._device)
+                broadcast_object(self.local_state_dict(), src_rank=self.global_rank, group=self.group, dist_device=self._device)
             else:
+                global_rank = self.get_global_rank(self.group, rank)
                 # Discard this tensor/rank, broadcast necessary for syncing
-                logging.debug("Discarding broadcast from rank %s", rank)
-                broadcast_object(empty_buffer, src_rank=rank, group=self.group, dist_device=self._device)
+                logging.debug("Discarding broadcast from rank %s", global_rank)
+                broadcast_object(empty_buffer, src_rank=global_rank, group=self.group, dist_device=self._device)
 
     def _free_other_grads(self) -> None:
         """Free all the gradients only useful for the other ranks
@@ -380,3 +393,4 @@ def _free_other_grads(self) -> None:
             for p in partition:
                 for t in p["params"]:
                     t.grad = None
+                    


Test patch: