Sharded DDP training fails with seq2seq models
## Information

Model I am using (Bert, XLNet ...): T5/BART/mBART/Marian

The problem arises when using:
* [x] the official example scripts: (give details below)
* [ ] my own modified scripts: (give details below)

The tasks I am working on is:
* [x] an official GLUE/SQUaD task: seq2seq
* [ ] my own task or dataset: (give details below)

## To reproduce

Steps to reproduce the behavior:

Run 
```
python -m torch.distributed.launch --nproc_per_node=2 examples/seq2seq/finetune_trainer.py \
--model_name_or_path sshleifer/tiny-mbart --output_dir output_dir --adam_eps 1e-06 --data_dir \
~/Downloads/wmt_en_ro --do_train --freeze_embeds --label_smoothing 0.1 --learning_rate 3e-5 \
--logging_first_step --logging_steps 1000 --max_source_length 128 --max_target_length 128 \
--num_train_epochs 1 --overwrite_output_dir --per_device_train_batch_size 4 --sortish_sampler \
--src_lang en_XX --task translation --tgt_lang ro_RO --val_max_target_length 128 --warmup_steps 500 \
--n_train 500 --sharded_ddp
```
will fail with
```
Traceback (most recent call last):
File "examples/seq2seq/finetune_trainer.py", line 379, in <module>
main()
File "examples/seq2seq/finetune_trainer.py", line 316, in main
model_path=model_args.model_name_or_path if os.path.isdir(model_args.model_name_or_path) else None
File "/home/sgugger/git/transformers/src/transformers/trainer.py", line 821, in train
self.optimizer.step()
File "/home/sgugger/.pyenv/versions/base/lib/python3.7/site-packages/torch/optim/lr_scheduler.py", line 67, in wrapper
return wrapped(*args, **kwargs)
File "/home/sgugger/git/fairscale/fairscale/optim/oss.py", line 210, in step
self._broadcast_params()
File "/home/sgugger/git/fairscale/fairscale/optim/oss.py", line 522, in _broadcast_params
if self.should_bucket_param[param]:
KeyError: Parameter containing:
tensor([[-0.0296,  0.0038],
[ 0.0000,  0.0000],
[ 0.0298,  0.0385],
...,
[-0.0161, -0.0024],
[ 0.0022, -0.0576],
[ 0.0053,  0.0256]], device='cuda:1')
0%|   
```

Using FP16 also fails.

## Expected behavior

The script should run to completion.

This is just a brief log of the 2 distinct errors mentioned in OP:

w/ `--fp16` the failure is:
```
  File "./finetune_trainer.py", line 379, in <module>
    main()
  File "./finetune_trainer.py", line 315, in main
    trainer.train(
  File "/mnt/nvme1/code/huggingface/transformers-master/src/transformers/trainer.py", line 818, in train
    self.scaler.step(self.optimizer)
  File "/home/stas/anaconda3/envs/main-38/lib/python3.8/site-packages/torch/cuda/amp/grad_scaler.py", line 330, in step
    assert len(optimizer_state["found_inf_per_device"]) > 0, "No inf checks were recorded for this optimizer."
AssertionError: No inf checks were recorded for this optimizer.
```

w/o `--fp16` the failure is:
```
  File "./finetune_trainer.py", line 379, in <module>
    main()
  File "./finetune_trainer.py", line 315, in main
    trainer.train(
  File "/mnt/nvme1/code/huggingface/transformers-master/src/transformers/trainer.py", line 821, in train
    self.optimizer.step()
  File "/home/stas/anaconda3/envs/main-38/lib/python3.8/site-packages/torch/optim/lr_scheduler.py", line 65, in wrapper
    return wrapped(*args, **kwargs)
  File "/home/stas/anaconda3/envs/main-38/lib/python3.8/site-packages/torch/optim/optimizer.py", line 89, in wrapper
    return func(*args, **kwargs)
  File "/home/stas/anaconda3/envs/main-38/lib/python3.8/site-packages/fairscale/optim/oss.py", line 210, in step
    self._broadcast_params()
  File "/home/stas/anaconda3/envs/main-38/lib/python3.8/site-packages/fairscale/optim/oss.py", line 522, in _broadcast_params
    if self.should_bucket_param[param]:
KeyError: Parameter containing:
tensor([[ ...]], device='cuda:1')
```
It's the very first parameter `model.shared.weight` in the case of mbart for example.

To test with t5 (same errors), run:

```
export BS=4; rm -r output_dir; CUDA_VISIBLE_DEVICES=0,1 PYTHONPATH=../../src USE_TF=0   python -m torch.distributed.launch --nproc_per_node=2  ./finetune_trainer.py --model_name_or_path patrickvonplaten/t5-tiny-random --output_dir output_dir --adam_eps 1e-06 --data_dir wmt_en_ro --do_train --freeze_embeds --label_smoothing 0.1 --learning_rate 3e-5 --logging_first_step --logging_steps 1000 --max_source_length 128 --max_target_length 128 --num_train_epochs 1 --overwrite_output_dir --per_device_train_batch_size $BS --sortish_sampler  --task translation_en_XX_to_ro_RO --val_max_target_length 128 --warmup_steps 500 --n_train 500 --sharded_ddp 
```
<|||||>The first problem (fp16) is easily fixed, it means that the doc is not good enough. Torch's grad scaler is not shard aware (the ranks do not have all the gradients with this technique), but you can use [this](https://github.com/facebookresearch/fairscale/blob/master/fairscale/optim/grad_scaler.py#L24) and that should work.

edit for future readers: that prove wrong, ShardedGradScaler already in use

The second issue is new to me, would you mind sharing a bit more of the reproduction steps ?<|||||>We initialize the should_bucket_param dictionary when the OSS optimizer is created. The assumption is that parameters should be frozen at this point. Any chance parameters are modified after the optimizer was created?<|||||>Thank you so much @blefaudeux and @msbaines for your follow up.

To reproduce:

```
# setup 
git clone  https://github.com/huggingface/transformers
cd transformers
cd examples/seq2seq
wget https://cdn-datasets.huggingface.co/translation/wmt_en_ro.tar.gz 
tar -xzvf wmt_en_ro.tar.gz
```

to reproduce the 2nd failure w/o `--fp16`:

```
export BS=4; rm -r output_dir; CUDA_VISIBLE_DEVICES=0,1 PYTHONPATH=../../src USE_TF=0   python -m torch.distributed.launch --nproc_per_node=2  ./finetune_trainer.py --model_name_or_path  sshleifer/tiny-mbart --output_dir output_dir --adam_eps 1e-06 --data_dir wmt_en_ro --do_train --freeze_embeds --label_smoothing 0.1 --learning_rate 3e-5 --logging_first_step --logging_steps 1000 --max_source_length 128 --max_target_length 128 --num_train_epochs 1 --overwrite_output_dir --per_device_train_batch_size $BS --sortish_sampler --src_lang en_XX --task translation --tgt_lang ro_RO --val_max_target_length 128 --warmup_steps 500 --n_train 500 --sharded_ddp
```

and then the first one is to just add `--fp16`

This is a tiny model that is good enough for testing the mechanics, so no good results to be expected. It's also very quick to download and load. To see real results swap `sshleifer/tiny-mbart` for `sshleifer/distill-mbart-en-ro-12-4`.<|||||>> The first problem (fp16) is easily fixed, it means that the doc is not good enough. Torch's grad scaler is not shard aware (the ranks do not have all the gradients with this technique), but you can use [this](https://github.com/facebookresearch/fairscale/blob/master/fairscale/optim/grad_scaler.py#L24) and that should work.

We are using it already:
https://github.com/huggingface/transformers/blob/dc9f24544291b25b44c9e87239a0ef4355396a4c/src/transformers/trainer.py#L315

if I print the object just before it fails in `self.scaler.step(self.optimizer)`, I get:
<fairscale.optim.grad_scaler.ShardedGradScaler object at 0x7ff27034bac0>

https://github.com/huggingface/transformers/blob/dc9f24544291b25b44c9e87239a0ef4355396a4c/src/transformers/trainer.py#L818

FWIW, I experience the exact same issue with deepspeed if I leave trainer's `--fp16` code - if I remove it and get deepspeed to handle that the failure goes away. So the common denominator is our code.<|||||>> Thank you so much @blefaudeux and @msbaines for your follow up.
> 
> To reproduce:
> 
> ```
> # setup 
> git clone  https://github.com/huggingface/transformers
> cd transformers
> cd examples/seq2seq
> wget https://cdn-datasets.huggingface.co/translation/wmt_en_ro.tar.gz 
> tar -xzvf wmt_en_ro.tar.gz
> ```
> 
> to reproduce the 2nd failure w/o `--fp16`:
> 
> ```
> export BS=4; rm -r output_dir; CUDA_VISIBLE_DEVICES=0,1 PYTHONPATH=../../src USE_TF=0   python -m torch.distributed.launch --nproc_per_node=2  ./finetune_trainer.py --model_name_or_path  sshleifer/tiny-mbart --output_dir output_dir --adam_eps 1e-06 --data_dir wmt_en_ro --do_train --freeze_embeds --label_smoothing 0.1 --learning_rate 3e-5 --logging_first_step --logging_steps 1000 --max_source_length 128 --max_target_length 128 --num_train_epochs 1 --overwrite_output_dir --per_device_train_batch_size $BS --sortish_sampler --src_lang en_XX --task translation --tgt_lang ro_RO --val_max_target_length 128 --warmup_steps 500 --n_train 500 --sharded_ddp
> ```
> 
> and then the first one is to just add `--fp16`
> 
> This is a tiny model that is good enough for testing the mechanics, so no good results to be expected. It's also very quick to download and load. To see real results swap `sshleifer/tiny-mbart` for `sshleifer/distill-mbart-en-ro-12-4`.

perfect, having a look right now, thanks for the repro help !<|||||>> We initialize the should_bucket_param dictionary when the OSS optimizer is created. The assumption is that parameters should be frozen at this point. Any chance parameters are modified after the optimizer was created?

good point, it's something that I was planning to address someday actually but not sure how urgent that was

> Thank you so much @blefaudeux and @msbaines for your follow up.
> 
> To reproduce:
> 
> ```
> # setup 
> git clone  https://github.com/huggingface/transformers
> cd transformers
> cd examples/seq2seq
> wget https://cdn-datasets.huggingface.co/translation/wmt_en_ro.tar.gz 
> tar -xzvf wmt_en_ro.tar.gz
> ```
> 
> to reproduce the 2nd failure w/o `--fp16`:
> 
> ```
> export BS=4; rm -r output_dir; CUDA_VISIBLE_DEVICES=0,1 PYTHONPATH=../../src USE_TF=0   python -m torch.distributed.launch --nproc_per_node=2  ./finetune_trainer.py --model_name_or_path  sshleifer/tiny-mbart --output_dir output_dir --adam_eps 1e-06 --data_dir wmt_en_ro --do_train --freeze_embeds --label_smoothing 0.1 --learning_rate 3e-5 --logging_first_step --logging_steps 1000 --max_source_length 128 --max_target_length 128 --num_train_epochs 1 --overwrite_output_dir --per_device_train_batch_size $BS --sortish_sampler --src_lang en_XX --task translation --tgt_lang ro_RO --val_max_target_length 128 --warmup_steps 500 --n_train 500 --sharded_ddp
> ```
> 
> and then the first one is to just add `--fp16`
> 
> This is a tiny model that is good enough for testing the mechanics, so no good results to be expected. It's also very quick to download and load. To see real results swap `sshleifer/tiny-mbart` for `sshleifer/distill-mbart-en-ro-12-4`.

re: fp16, could it be that the FW has not been run within an AMP context (a) and the scaler is not invoked in the backward (b) ? I cannot find it [in the trainer](https://github.com/huggingface/transformers/blob/master/src/transformers/trainer.py#L799). The scaler needs to be invoked when computing the grads, it will check for infs there which could explain why you're seeing this assert. Note that the hugginface codebase is new to me, so could be that this is wrapped somewhere else and I'm completely missing it<|||||>I think the questions you're asking about are all in this `training_step` code:

https://github.com/huggingface/transformers/blob/dc9f24544291b25b44c9e87239a0ef4355396a4c/src/transformers/trainer.py#L1126-L1146

I didn't write it, but from a quick read it appears that it's a yes to all of your suggestions. 

`self.use_amp` = native amp, `use_apex` is apex - so we are talking native amp here - that is the branches with `use_amp = True`

I'll step through with debugger to see that it is actually so.

<|||||>> I think the questions you're asking about are all in this `training_step` code:
> 
> https://github.com/huggingface/transformers/blob/dc9f24544291b25b44c9e87239a0ef4355396a4c/src/transformers/trainer.py#L1126-L1146
> 
> I didn't write it, but from a quick read it appears that it's a yes to all of your suggestions.
> 
> `self.use_amp` = native amp, `use_apex` is apex - so we are talking native amp here - that is the branches with `use_amp = True`
> 
> I'll step through with debugger to see that it is actually so.

looks good indeed ! basically I just know that the "found_inf_per_device" are populated in the "unscale_" step, so this key being absent points to this step being skipped somehow.<|||||>@stas00  https://github.com/facebookresearch/fairscale/pull/256 Fixes your repro on a single node, it's a side effect though (bucketing is effectively disabled on a single node), multinode + huggingface is still probably broken. It does look like the model somehow changes after construction, still finding my way around your codebase.
<|||||>re: --fp16: If the machine is clean it does break, my assumption is that the dist.reduce() in the ShardedGradScaler fails somehow in between the processes. Once it dies, one process stays up actually (visible with nvidia-smi), and all the subsequent runs will work fine. So if that helps it looks like it could be an issue with the dist init, or passing the settings around that.<|||||>> looks good indeed ! basically I just know that the "found_inf_per_device" are populated in the "unscale_" step, so this key being absent points to this step being skipped somehow.

So, OK, I have it setup in the debugger 

It runs `unscale_` just fine, but doesn't find any `inf`:

```
        inv_scale = self._scale.double().reciprocal().float()
        found_inf = torch.full((1,), 0.0, dtype=torch.float32, device=self._scale.device)

        optimizer_state["found_inf_per_device"] = self._unscale_grads_(optimizer, inv_scale, found_inf, False)
```
- `inv_scale = tensor([1.5259e-05], device='cuda:0')`
- `found_inf = tensor([0.], device='cuda:0')`
- `optimizer_state["found_inf_per_device"] = {}`


rerunning - inside `self._unscale_grads_` we get:

```
per_device_and_dtype_grads = {defaultdict: 1} defaultdict(<function GradScaler._unscale_grads_.<locals>.<lambda> at 0x7f895c2174c0>, {device(type='cuda', index=1): defaultdict(<class 'list'>, {torch.float32: [tensor([[nan, nan],\n        [nan, nan]], device='cuda:1'), tensor([[nan, nan],\n        [nan, 
 cuda:1 = {defaultdict: 1} defaultdict(<class 'list'>, {torch.float32: [tensor([[nan, nan],\n        [nan, nan]], device='cuda:1'), tensor([[nan, nan],\n        [nan, nan]], device='cuda:1'), tensor([[nan, nan],\n        [nan, nan]], device='cuda:1'), tensor([[nan, nan],\n        [nan, 
  default_factory = {type} <class 'list'>
  torch.float32 = {list: 92} [tensor([[nan, nan],\n        [nan, nan]], device='cuda:1'), tensor([[nan, nan],\n        [nan, nan]], device='cuda:1'), tensor([[nan, nan],\n        [nan, nan]], device='cuda:1'), tensor([[nan, nan],\n        [nan, nan]], device='cuda:1'), tensor([nan, nan], 
   00 = {Tensor: 2} tensor([[nan, nan],\n        [nan, nan]], device='cuda:1')
   01 = {Tensor: 2} tensor([[nan, nan],\n        [nan, nan]], device='cuda:1')
   02 = {Tensor: 2} tensor([[nan, nan],\n        [nan, nan]], device='cuda:1')
   03 = {Tensor: 2} tensor([[nan, nan],\n        [nan, nan]], device='cuda:1')
[...]
   89 = {Tensor: 2} tensor([nan, nan], device='cuda:1')
   90 = {Tensor: 2} tensor([nan, nan], device='cuda:1')
   91 = {Tensor: 2} tensor([-12339.3516, -13527.4590], device='cuda:1')
```

and then this gets called:

```
            for device, per_dtype_grads in per_device_and_dtype_grads.items():
                for grads in per_dtype_grads.values():
                    torch._amp_foreach_non_finite_check_and_unscale_(grads,
                                                                     per_device_found_inf.get(device),
                                                                     per_device_inv_scale.get(device))
```

and it doesn't find anything. 

I think the debugger also has a race condition and sometimes we get the first process (cuda:0) and other times the 2nd one (cuda:1). 

ok I figured out how to switch threads in pycharm debugger, so I can now go back and force between the 2 processes.

I haven't figured out how to do `.to()` calls in pycharm debugger yet - when those happen as a step it just hangs so I have carefully to skip over those.

Now rerunning it again and hitting `cuda:0`  `per_device_and_dtype_grads` is empty.

And yes, I have noticed that with `--fp16` only one process fails with this error - the other keeps on running. 

So basically we are having this happen on one process but not the other. Somewhere must be a bug not running the same code for both processes.

(this has just changed in pytorch master - now all subprocesses will die nicely w/o leaving zombies - by yours truly ;)<|||||>thanks for the backtrace ! so, for one the fact that the grads are not all the same is expected with this method, the grads are sharded across the ranks (ie: partitioned), depending on which parameters each rank will optimize. The ShardedGradScaler should be aware of that, and syncs in between the ranks to make sure that they all get the same knowledge, looks like this fails somehow then. Having a quick look right now

(well done for the zombie process destruction ! now somehow if the zombie process is still here the next run "works")<|||||>So it appears that here:
```
self.scaler.scale(loss).backward()
```
doesn't set `param.grad`s in `cuda:0` but sets it in `cuda:1`

To quickly see that add:

```
        for group in self.optimizer.param_groups:
            for param in group["params"]:
                print(f"{self.args.local_rank}: {param.grad}")
```
after:

https://github.com/huggingface/transformers/blob/dc9f24544291b25b44c9e87239a0ef4355396a4c/src/transformers/trainer.py#L1139

So on `cuda:0` it goes into `backward` and almost immediately returns. some flag must be off.

I checked that `self.scaler.scale(loss).requires_grad == True` on both, so it's not that.

And `unscale_` fails to find any data, because all grads are None for `self.optimizer` on `cuda:0` and that's where it crashes.

(debugging parallel processes proved to be far from easy - not sure why - half the time pycharm debugger either gets stuck or suddenly can't see the other process - so it's very slow comparing what the difference is between the two sides. ideally I need to be able to run both processes side by side - each in its own debugger - then it'd be much easier to find the divergence)<|||||>hmm, looks like a device and rank mismatch then, that could match with the process group observation above. Looking into the code, are you using a specific process group ?<|||||>https://github.com/huggingface/transformers/blob/dc9f24544291b25b44c9e87239a0ef4355396a4c/src/transformers/training_args.py#L472-L474<|||||>I'm trying to find a better solution to the other issue you were seeing with the bucketing (now fixed on fairscale master for a single node), and it just appeared to me that they could be tied: could it be that the models change devices after the sharded optimizer is built ?

edit: just checked that, not the case, rank/device match at construction time and during the first step<|||||>Ah ok, got to the bottom of the first bug. There are params which don't require grad, I was skipping them, my bad. Fixing that properly so that multi-node also works for you.

edit: now this is properly done (proverbial second fix, which covers multinode)<|||||>ok, now on the fp16 issue, there's at least one thing which cannot work well in that case: the gradient clipping.                        
https://github.com/huggingface/transformers/blob/dc9f24544291b25b44c9e87239a0ef4355396a4c/src/transformers/trainer.py#L809

To explain a little, when ShardedDDP is used torch.utils.clip_grad_norm is blind to the sharding (same issue with the scaler), in that it will only consider the norm of the gradients present on this rank, it's very much not aware of the distributed nature of the problem. I think that we should improve the API so that we don't need to change all these little pieces, this could be done with torch RPC for instance. We've a solution right now, in that we provide a shard-aware gradient clipping in https://github.com/facebookresearch/fairscale/blob/master/fairscale/optim/oss.py#L219
whose interface is similar.

I'm not sure that it's the only issue with fp16, but that's one of them for sure.

 <|||||>Ok, second and hopefully last issue with fp16 caught, works on my machine with the following "patch" around the clipping (to select fairscale's clipping when that makes sense), and an incoming fix in Fairscale related to a broken partitioning in a pathological case.

Replace https://github.com/huggingface/transformers/blob/dc9f24544291b25b44c9e87239a0ef4355396a4c/src/transformers/trainer.py#L809 
with 
```
       if self.use_amp:
          # AMP: gradients need unscaling
          self.scaler.unscale_(self.optimizer)
      
      if hasattr(self.optimizer, "clip_grad_norm"):
          # Sharded optimizer, specific gradient clipping 
          self.optimizer.clip_grad_norm(self.args.max_grad_norm)
      else:
          # Vanilla -monolithic- clipping, handling Apex or full precision
          torch.nn.utils.clip_grad_norm_(
              amp.master_params(self.optimizer) if self.use_apex else model.parameters(),
              self.args.max_grad_norm,
          )
```
<|||||>With all the linked PRs it works for me with --fp16 and ShardedDDP, and the speed bumps up nicely, AMP basically doubles the throughput.<|||||>Wow, thanks a lot for all this debugging @blefaudeux ! I'll draft a quick patch for the `Trainer` gradient clipping and tag you on the PR.<|||||>Patch is in the PR linked above!<|||||>Amazing! Thank you so much, @blefaudeux! 

I merged all the suggested code and everything works. Yay!

**Except it's 30% slower w/ sharded_ddp**

> @blefaudeux wrote:
> With all the linked PRs it works for me with --fp16 and ShardedDDP, and the speed bumps up nicely, AMP basically doubles the throughput.

So I'm not seeing what you're seeing,  @blefaudeux

What I did.

1. hf:

rebased to include https://github.com/huggingface/transformers/pull/9168

2. fairscale:


```
git checkout -b hf

# https://github.com/facebookresearch/fairscale/pull/262
git cherry-pick e305f10e81db95d14c9edd3f9e1e18b0fb2847fd
# (had to resolve a simple conflict)

# https://github.com/facebookresearch/fairscale/pull/259
git cherry-pick ad28820
git cherry-pick 705c188
git cherry-pick d6aa285
git cherry-pick a315ee3

# only this works for me at the moment for pytorch-nightly
python setup.py bdist_wheel
pip uninstall -y fairscale; pip install dist/fairscale-0.1.1-cp38-cp38-linux_x86_64.whl
```

# speed benchmarks

note I switched to the real model `sshleifer/distill-mbart-en-ro-12-4` (`sshleifer/tiny-mbart` was perfect for testing mechanics)

```
# base-line (no --sharded_ddp / no --fp16)
export BS=4; rm -r output_dir; CUDA_VISIBLE_DEVICES=0,1 PYTHONPATH=../../src USE_TF=0   python -m torch.distributed.launch --nproc_per_node=2 --master_port=9910  ./finetune_trainer.py --model_name_or_path sshleifer/distill-mbart-en-ro-12-4 --output_dir output_dir --adam_eps 1e-06 --data_dir wmt_en_ro --do_train --freeze_embeds --label_smoothing 0.1 --learning_rate 3e-5 --logging_first_step --logging_steps 1000 --max_source_length 128 --max_target_length 128 --num_train_epochs 1 --overwrite_output_dir --per_device_train_batch_size $BS --sortish_sampler --src_lang en_XX --task translation --tgt_lang ro_RO --val_max_target_length 128 --warmup_steps 500 --n_train 500

2020-12-17 10:44:55 | INFO | __main__ |   train_runtime = 28.0357

# --sharded_ddp 

2020-12-17 10:41:02 | INFO | __main__ |   train_runtime = 40.4482


### fp16

# --fp16

2020-12-17 10:43:30 | INFO | __main__ |   train_runtime = 29.4658

# --sharded_ddp  --fp16

2020-12-17 10:38:53 | INFO | __main__ |   train_runtime = 39.4722

```
<|||||>Also, clearly it can be seen from the benchmark that we don't gain from --fp16 at the baseline (before adding fairscale) - so something is not right there either.<|||||>In the good news, I can squeeze 3x batch size with ` --sharded_ddp` as compared to the baseline before my card OOMs. Amazing!

So with BS=12 (baseline I could do only 4)

```
# --sharded_ddp
2020-12-17 11:10:17 | INFO | __main__ |   train_runtime = 17.2038
# --sharded_ddp --fp16
2020-12-17 11:07:32 | INFO | __main__ |   train_runtime = 15.2403
```

So the total training time is about ~50% as compared to w/o `--sharded_ddp` but must increase BS 3 times.

<|||||>quality is about the same - about 0.1% worse with ` --sharded_ddp` and same BS on a quick test.

```
# baseline
export BS=4; rm -r output_dir; CUDA_VISIBLE_DEVICES=0,1 PYTHONPATH=../../src USE_TF=0   python -m torch.distributed.launch --nproc_per_node=2 --master_port=9910  ./finetune_trainer.py --model_name_or_path sshleifer/distill-mbart-en-ro-12-4 --output_dir output_dir --adam_eps 1e-06 --data_dir wmt_en_ro --do_train --freeze_embeds --label_smoothing 0.1 --learning_rate 3e-5 --logging_first_step --logging_steps 1000 --max_source_length 128 --max_target_length 128 --num_train_epochs 1 --overwrite_output_dir --per_device_train_batch_size $BS --sortish_sampler --src_lang en_XX --task translation --tgt_lang ro_RO --val_max_target_length 128 --warmup_steps 500 --n_train 500 --do_eval --eval_steps 25000 --evaluation_strategy=steps --n_val 200 --per_device_eval_batch_size $BS --do_predict --predict_with_generate --n_test 200  --fp16

2020-12-17 11:20:28 | INFO | __main__ |   train_runtime = 29.6081

2020-12-17 11:20:50 | INFO | __main__ |   val_bleu  = 26.3473
2020-12-17 11:21:20 | INFO | __main__ |   test_bleu = 25.7341

# run 2
2020-12-17 11:52:40 | INFO | __main__ |   val_bleu = 26.3473
2020-12-17 11:53:10 | INFO | __main__ |   test_bleu = 25.7341

# --sharded_ddp
export BS=4; rm -r output_dir; CUDA_VISIBLE_DEVICES=0,1 PYTHONPATH=../../src USE_TF=0   python -m torch.distributed.launch --nproc_per_node=2 --master_port=9910  ./finetune_trainer.py --model_name_or_path sshleifer/distill-mbart-en-ro-12-4 --output_dir output_dir --adam_eps 1e-06 --data_dir wmt_en_ro --do_train --freeze_embeds --label_smoothing 0.1 --learning_rate 3e-5 --logging_first_step --logging_steps 1000 --max_source_length 128 --max_target_length 128 --num_train_epochs 1 --overwrite_output_dir --per_device_train_batch_size $BS --sortish_sampler --src_lang en_XX --task translation --tgt_lang ro_RO --val_max_target_length 128 --warmup_steps 500 --n_train 500 --do_eval --eval_steps 25000 --evaluation_strategy=steps --n_val 200 --per_device_eval_batch_size $BS --do_predict --predict_with_generate --n_test 200  --fp16 --sharded_ddp

2020-12-17 11:25:23 | INFO | __main__ |   train_runtime = 39.4316

2020-12-17 11:25:46 | INFO | __main__ |   val_bleu  = 26.2563
2020-12-17 11:26:16 | INFO | __main__ |   test_bleu = 25.5779

# run 2
2020-12-17 11:28:32 | INFO | __main__ |   val_bleu  = 26.1359
2020-12-17 11:29:02 | INFO | __main__ |   test_bleu = 25.6613

# run 3
2020-12-17 11:30:51 | INFO | __main__ |   val_bleu  = 26.1359
2020-12-17 11:31:21 | INFO | __main__ |   test_bleu = 25.6067

# run 4
2020-12-17 11:50:05 | INFO | __main__ |   val_bleu  = 26.1359
2020-12-17 11:50:35 | INFO | __main__ |   test_bleu = 25.5889

```

larger BS is much faster, while the eval is on par
```

# --sharded_ddp
export BS=12; rm -r output_dir; CUDA_VISIBLE_DEVICES=0,1 PYTHONPATH=../../src USE_TF=0   python -m torch.distributed.launch --nproc_per_node=2 --master_port=9910  ./finetune_trainer.py --model_name_or_path sshleifer/distill-mbart-en-ro-12-4 --output_dir output_dir --adam_eps 1e-06 --data_dir wmt_en_ro --do_train --freeze_embeds --label_smoothing 0.1 --learning_rate 3e-5 --logging_first_step --logging_steps 1000 --max_source_length 128 --max_target_length 128 --num_train_epochs 1 --overwrite_output_dir --per_device_train_batch_size $BS --sortish_sampler --src_lang en_XX --task translation --tgt_lang ro_RO --val_max_target_length 128 --warmup_steps 500 --n_train 500 --do_eval --eval_steps 25000 --evaluation_strategy=steps --n_val 200 --per_device_eval_batch_size $BS --do_predict --predict_with_generate --n_test 200  --fp16 --sharded_ddp

2020-12-17 11:22:26 | INFO | __main__ |   train_runtime = 15.2409

2020-12-17 11:22:40 | INFO | __main__ |   val_bleu  = 26.466
2020-12-17 11:23:01 | INFO | __main__ |   test_bleu = 25.6237
```
<|||||>Hi,
thanks a lot for this fixing  this issue,  this is great to have this option working, do you mind including this command for faster training of seq2seq models and some documentations on what is sharded_ddp option and how much it helps on README of seq2seq folder. thanks. <|||||>We will do so soon. Since it was just added we don't have enough solid stats to advertise % improvements specific to transformers, so please give us some time.

Until then you can just add `--sharded_ddp` and see for yourself. One thing I noticed is that I can use a 3 times bigger batch size w/ it.
**edit**: oh but wait, the fairscale master doesn't have the PRs with important fixes merged yet, so it's basically not ready yet unless you want to manually merge the changes. I guess you could use my branch where I did the merges already if it's urgent...

Wrt/ Optimizer state sharding, you can read the documentation in the paper on DeepSpeed/ZeRO https://arxiv.org/abs/1910.02054

If I'm not mistaken the exact entry is 5.1 Pos : Optimizer State Partitioning

But I don't know whether it's the same or not, since this is fairscale implementation. 

@msbaines, what's the best place for us to link to to explain what fairscale's sharded optimizer does? I didn't find any docs on your repo/website. Thank you!<|||||>hey there, trying to cover a couple of questions:
- fp16 : should almost always give a boost, I was seeing close to 2x a couple of days back (on P100s with the above test case), could depend on your hardware ? (and workload of course, if IO is the bottleneck that wouldn't be the case)

- speed impact of shardedDDP at iso-batch size: there's some speed lost indeed, it's probably not fundamental though, could be improved on our end without any change on your end. The two ongoing PRs for hugginface should help a bit for instance. Some time lost is not trivial due to all the code being python and some impact on GIL contention for instance (the reduce involves a lot of hooks being called to redirect the gradients to the right ranks). Longer term, ShardedDDP should just be a "mode of operation" of Pytorch's DDP, which is mostly cpp, so speed will probably improve, consider the current state as a baseline. Note that the state saving support is not very elegant right now, and is very slow, so if this counts against the throughput (if there are a lot of checkpoints) you'll see a difference indeed. We plan to improve on that asap

- what does shardedDDP do ? Pos + Pg from the zero paper + mixed precision if you use AMP of course, in plain english optimizer state sharding + gradient sharding + automatic mixed precision (from PyTorch of course).<|||||>> fp16 : should almost always give a boost, I was seeing close to 2x a couple of days back (on P100s with the above test case), could depend on your hardware ? (and workload of course, if IO is the bottleneck that wouldn't be the case)

Yes, honestly I'm not sure how to approach this. @sgugger says he sees the improvements, but he is on a different hardware and pytorch.

I have a partially supported rtx-3090 and gtx-1070, so obviously the older card slows down the new one, but any speed up should be relative to that slowdown. Perhaps it has to do with the incomplete rtx-3090 support. I'm impatiently waiting for cuda-11.2 pytorch support, which supposedly should provide the full power. I'm also on pytorch-nightly - not sure if it makes a difference.

I may have to delegate benchmarking to others, since I don't see most of the benefits. 

I'm in the process of doing the same integration for deepspeed and obviously have the same issue there. So I guess I will propose to merge what I have working and then either someone else will benchmark or wait till pytorch/cuda-11.2 is out. Holiday season doesn't help.

>  speed impact of shardedDDP at iso-batch size: there's some speed lost indeed, it's probably not fundamental though, could be improved on our end without any change on your end. The two ongoing PRs for hugginface should help a bit for instance. Some time lost is not trivial due to all the code being python and some impact on GIL contention for instance (the reduce involves a lot of hooks being called to redirect the gradients to the right ranks). Longer term, ShardedDDP should just be a "mode of operation" of Pytorch's DDP, which is mostly cpp, so speed will probably improve, consider the current state as a baseline. Note that the state saving support is not very elegant right now, and is very slow, so if this counts against the throughput (if there are a lot of checkpoints) you'll see a difference indeed. We plan to improve on that asap

I'm a bit confused about this vs the previous quoted comment - are you saying that it's normal that it should be slower when enabling shardedDDP for the same batch size and this whole time the alluded to speedup was in the ability to dramatically increase the batch size and therefore overall speed?

If this is so then my setup only has an issue with fp16, and it's just fine otherwise, since as I mentioned I can do 3x bigger BS and the overall speed is much much faster.

Please clarify my confusion. Thank you!
 
> what does shardedDDP do ? Pos + Pg from the zero paper + mixed precision if you use AMP of course, in plain english optimizer state sharding + gradient sharding + automatic mixed precision (from PyTorch of course).

Great, thank you! I wasn't sure whether you had some custom variations to the deepspeed paper.

So I suppose we quote these features and send readers to the ZeRO paper as I did 2 comments above for details?
<|||||>@sgugger, we have the new arg `--sharded_ddp` documented [here](https://huggingface.co/transformers/main_classes/trainer.html?highlight=trainer#trainingarguments), but where do you think we should document what shardedDDP does (and then there will be the same question for deepspeed).

I suppose it fits the best into the Training doc, I added a new section for integrations: https://github.com/huggingface/transformers/pull/9208
<|||||>> I'm a bit confused about this vs the previous quoted comment - are you saying that it's normal that it should be slower when enabling shardedDDP for the same batch size and this whole time the alluded to speedup was in the ability to dramatically increase the batch size and therefore overall speed?

trying to word that better:
- fp16 speed: I was trying to say that in both cases I was seeing x2, with the test case provided above. I think that your setup is a bit strange, so to say, I would suspect that's why you're not seeing it yourself
- shardedddp speed (orthogonal to fp16): speed when compared to ddp is in between 105% and 70% (iso batch), from what I've seen personally, I was trying to say that it's not completely set in stone and that improving on it should not require API changes. It's a good comparison point, out of principle it should not really be slower, and it has reasons for being a bit faster (the optimizer step is faster, since it only proceeds a shard and not the full model). The lowered  speed mostly comes from GIL contention and slower cpu path, from what I understand, one of the PR currently up improves a bit on that. Then you can of course change the batch size, that's the whole point, glad that in the end the benefit is very clear there :)<|||||>Thank you for clarifying on both points, @blefaudeux 

So GIL comes up a lot these days - is there a reason not to switch to multiprocessing and overcome GIL altogether? I suppose passing data around would be much slower between procs and thus threads are still faster in the particular context you use them?

Or since the core of things is written in C++, perhaps POSIX threads could be used as opposed to python? I can imagine that won't be easy to do either w/o rewriting the whole thing in C++ and abandoning python altogether.
<|||||>> Thank you for clarifying on both points, @blefaudeux 
> 
> 
> 
> So GIL comes up a lot these days - is there a reason not to switch to multiprocessing and overcome GIL altogether? I suppose passing data around would be much slower between procs and thus threads are still faster in the particular context you use them?
> 
> 
> 
> Or since the core of things is written in C++, perhaps POSIX threads could be used as opposed to python? I can imagine that won't be easy to do either w/o rewriting the whole thing in C++ and abandoning python altogether.
> 
> 

(Multiprocessing) the issue is that you're deferring actions tied to torch.distributed.Work (ie: when this gradient has been reduced, if this rank does not own the corresponding parameter update then drop the gradient), and these pseudo futures are not easily shared across processes. Might well change over time, maybe that I missed a trick or two, just describing the current status.

(GIL) for the backward pass / reduce again, hooks are fired on a per parameter basis by Torch Distributed autograd engine (c++). Each hook handles python objects, so it needs to get hold of the GIL. 

(C++/change language) there's an ongoing RFC on making DDP more composable, there are already quite a few options with the coms hooks for instance, so that the zero use case could be expressed in pure "ddp" blocks (which are cpp). In my opinion that's a reasonable perspective (easier said than done though), ShardedDDP and DDP are intrinsically doing a very similar job, and there are other concepts which would be super interesting to mix (slowmo, pipe, more mixed precision tweaks..) which is easier when all expressed on a common basis.<|||||>Thank you for this detailed answer, @blefaudeux!

So basically in some time things will get even faster, the future is looking bright!

<|||||>Hi
I am also could not get sharded_ddp to work on seq2seq, could you provide the version you used? please see the bug I opened here https://github.com/huggingface/transformers/issues/9215 thanks <|||||>As I clearly unsuccessfully tried to explain @blefaudeux proposed several PRs that made it possible to use `--sharded_ddp` but they are not in `fairscale` master yet and you need to manually merge the PRs until those PRs are merged into master.

Since I already merged these PRs into a local branch, you can use my branch to install fairscale that works with finetune_trainer.py as following:

```
git clone -b hf https://github.com/stas00/fairscale
cd fairscale
python setup.py bdist_wheel
pip uninstall -y fairscale
# adjust the specific whl filename if needed
pip install dist/fairscale-0.1.1-cp38-cp38-linux_x86_64.whl
```

and of course, you will need `transformers` master.
<|||||>> As I clearly unsuccessfully tried to explain @blefaudeux proposed several PRs that made it possible to use `--sharded_ddp` but they are not in `fairscale` master yet and you need to manually merge the PRs until those PRs are merged into master.
> 
> Since I already merged these PRs into a local branch, you can use my branch to install fairscale that works with finetune_trainer.py as following:
> 
> ```
> git clone -b hf https://github.com/stas00/fairscale
> cd fairscale
> python setup.py bdist_wheel
> pip uninstall -y fairscale
> # adjust the specific whl filename if needed
> pip install dist/fairscale-0.1.1-cp38-cp38-linux_x86_64.whl
> ```
> 
> and of course, you will need `transformers` master.

Hi @stas00, I just pushed a new release (see https://github.com/facebookresearch/fairscale/releases/tag/v0.1.2), I hope that works. Thanks again for all the follow ups <|||||>Awesome! Thank you, @blefaudeux!

Could you please let me know when it's on pypi - then I will retest and we will merge the doc PR<|||||>> Awesome! Thank you, @blefaudeux!
> 
> Could you please let me know when it's on pypi - then I will retest and we will merge the doc PR

cc @msbaines , I don't have rights on https://pypi.org/project/fairscale/ actually, and I don't think that it's automatically tied to our github releases<|||||>I think this may help you to automate the process:
https://packaging.python.org/guides/publishing-package-distribution-releases-using-github-actions-ci-cd-workflows/

I haven't used it myself, but this seems to be the best doc I found while searching for it. I personally have [an automated release process](https://github.com/stas00/ipyexperiments/blob/8f507c1bbed095140e14318e7727176201ced20c/Makefile#L145) that does the version bumping, the tagging, releasing and uploading to pypi/conda all in one command ;)<|||||>> I think this may help you to automate the process:
> https://packaging.python.org/guides/publishing-package-distribution-releases-using-github-actions-ci-cd-workflows/
> 
> I haven't used it myself, but this seems to be the best doc I found while searching for it. I personally have [an automated release process](https://github.com/stas00/ipyexperiments/blob/8f507c1bbed095140e14318e7727176201ced20c/Makefile#L145) that does the version bumping, the tagging, releasing and uploading to pypi/conda all in one command ;)

@stas00 @msbaines done, 0.1.3 published, checking right now with huggingface but so far so good<|||||>> I think this may help you to automate the process:
> https://packaging.python.org/guides/publishing-package-distribution-releases-using-github-actions-ci-cd-workflows/
> 
> I haven't used it myself, but this seems to be the best doc I found while searching for it. I personally have [an automated release process](https://github.com/stas00/ipyexperiments/blob/8f507c1bbed095140e14318e7727176201ced20c/Makefile#L145) that does the version bumping, the tagging, releasing and uploading to pypi/conda all in one command ;)

thanks for the pointers by the way @stas00, looking into that, probably github actions in the end<|||||>Also you may want to add `long_description` entry inside `setup.py` so that the resulting https://pypi.org/project/fairscale/ is not empty, e.g. see: 

https://github.com/huggingface/transformers/blob/eef66035a28b935b7823c9a7ddc8c569e077ce11/setup.py#L250-L252

which results in a nice description at https://pypi.org/project/transformers/<|||||>> @stas00 @msbaines done, 0.1.3 published, checking right now with huggingface but so far so good

Thank you, @blefaudeux!

Unfortunately I can't build the pypi package w/ pt-nightly, until pt-1.8 gets released - due to rtx-3090 card, so I can only do it from source:
```
rm -r dist build
python setup.py bdist_wheel
pip uninstall -y fairscale
pip install dist/fairscale-*.whl
```
So I trust that your testing was successful 

And I re-tested with master - all looks great.

So I'm going to merge the doc https://github.com/huggingface/transformers/pull/9208 and we can let users know that they training has just got a magical super-boost. I think @sgugger should do the honors since he integrated it.

This is very exciting! Thank you for your amazing support, @blefaudeux and your team!