Fine-Tuning Wav2Vec2 with PyTorch DDP
Author: patrickvonplaten
Pull request: no

## Environment info

- `transformers` version: 4.11.0.dev0
- Platform: Linux-5.11.0-1017-aws-x86_64-with-glibc2.29
- Python version: 3.8.10
- PyTorch version (GPU?): 1.9.0+cu111 (True)
- Tensorflow version (GPU?): 2.6.0 (False)
- Flax version (CPU?/GPU?/TPU?): 0.3.4 (cpu)
- Jax version: 0.2.20
- JaxLib version: 0.1.71
- Using GPU in script?: yes (8 or 1)
- Using distributed or parallel set-up in script?: yes

### Who can help

@sgugger @stas00 @anton-l 

## Problem:

I'm running some experiments on fine-tuning a [pretrained XLSR-Wav2Vec2 model](https://huggingface.co/facebook/wav2vec2-large-xlsr-53) on the [Turkish dataset of Common Voice](https://huggingface.co/datasets/common_voice). 

The fine-tuning script is an updated version of the existing [`run_common_voice.py`](https://github.com/huggingface/transformers/blob/master/examples/research_projects/wav2vec2/run_common_voice.py) script that can be seen in this PR: https://github.com/huggingface/transformers/blob/97936d3aacc04f6253ff178415b8a57768fc8ce6/examples/pytorch/speech-recognition/run_speech_recognition_ctc.py

It leverages the `Trainer` for CTC training of Wav2Vec2. 

I'm running the training script for both distributed training (as follows):

```bash
python -m torch.distributed.launch \
        --nproc_per_node 8 run_speech_recognition_ctc.py \
        --dataset_name="common_voice" \
        --model_name_or_path="facebook/wav2vec2-large-xlsr-53" \
        --dataset_config_name="tr" \
        --output_dir="./wav2vec2-large-xlsr-turkish-demo-dist" \
        --overwrite_output_dir \
        --num_train_epochs="30" \
        --per_device_train_batch_size="4" \
        --learning_rate="3e-4" \
        --warmup_steps="500" \
        --evaluation_strategy="steps" \
        --save_steps="400" \
        --eval_steps="100" \
        --logging_steps="1" \
        --save_total_limit="3" \
        --fp16 \
        --freeze_feature_extractor \
        --chars_to_ignore , ? . ! - \; \: \" “ % ‘ ” � \
        --do_train --do_eval
```

and single-GPU training:

```bash
CUDA_VISIBLE_DEVICES="0" python run_speech_recognition_ctc.py \
        --dataset_name="common_voice" \
        --model_name_or_path="facebook/wav2vec2-large-xlsr-53" \
        --dataset_config_name="tr" \
        --output_dir="./wav2vec2-large-xlsr-turkish-demo" \
        --overwrite_output_dir \
        --num_train_epochs="30" \
        --per_device_train_batch_size="16" \
        --gradient_accumulation_steps="2" \
        --learning_rate="3e-4" \
        --warmup_steps="500" \
        --evaluation_strategy="steps" \
        --save_steps="400" \
        --eval_steps="100" \
        --logging_steps="1" \
        --save_total_limit="3" \
        --freeze_feature_extractor \
        --gradient_checkpointing \
        --fp16 \
        --chars_to_ignore , ? . ! - \; \: \" “ % ‘ ” � \
        --do_train --do_eval
```

As you can see the only difference between the two scripts is that distributed training (DDP) does **not** use gradient checkpointing and has a "per-GPU" batch size of 4 resulting in an effective batch size of 32, whereas the single GPU training has a "per-GPU" batch size of 16 and uses gradient accumulation of 2 (and gradient checkpointing). So the training scripts are more or less identical in terms of learning rate decay, optimizer, effective batch size, ...

Now what is quite surprising to me is that single-GPU training works very well. Here is a report with the most important metrics of the run: https://wandb.ai/patrickvonplaten/huggingface/reports/Wav2Vec2-1-GPU-V100--VmlldzoxMDQwNzI0?accessToken=5xhtxrgy59l7dl2sds08bfk8xq1l30uf1ae0i5lio2r7dpx43vzxufsjmxkkbkig
while distributed training doesn't work at all - here a report of the run: https://wandb.ai/patrickvonplaten/huggingface/reports/Wav2Vec2-DistributedDataParallel-DDP-8-GPU-V100--VmlldzoxMDQwMDU3?accessToken=rsxt5n2s31bfg3kmbtvb982zcqlg8hby7mrjniftnx4n87kephus81zeaj92xfbu

While Wav2Vec2's CTC loss isn't super stable the single-GPU script is quite robust to changes in the batch size, learning rate, random seed (I've tried a bunch of slight changes and the script always manages to push the training/eval loss below 1 and yield a reasonable word error rate in the beginning. On the other hand the distributed script doesn't seem to work at all (tried out a variety of dropout rates, learning rates, batch sizes, layerdrop, ....) -> none of them converge.

That's quite surprising to me as the scripts should in theory be more or less the same.

Some possible reasons I thought could be: 

- In distributed training the gradients are computed for each process/gpu separately and then averaged (reduced). However this is slighly different to single GPU training because of the following: 
For a single GPU, each input sample in the batch can have a different number of losses, e.g. if the labels are  `[["Hello my name is"], ["hey <pad> <pad> <pad>"]]` then on a single GPU the loss is correctly averaged (5 words = 5 losses = sum(losses) / 5). However in DDP the losses are averaged locally and then the gradients are averaged globally which would mean (1st GPU: 4 words = 4 loss & 2nd GPU: 1 word = 1 loss => (sum(losses_gpu1) / 4 + sum(losses_gpu2) / 1) / 2 which is not the same as on a single GPU. However, I've also played around with `group_by_length` - in this case the inputs per batch should be roughly similar and there shouldn't be a problem, but that didn't help either. I've also summed the losses and scaled the gradients correctly - see: https://github.com/huggingface/transformers/blob/97936d3aacc04f6253ff178415b8a57768fc8ce6/examples/pytorch/speech-recognition/run_speech_recognition_ctc.py#L286 which also didn't help. Related issue/discussion on PyTorch: https://discuss.pytorch.org/t/average-loss-in-dp-and-ddp/93306

- Could gradient-checkpointing be the reason? I really don't see how this could make a difference though...

- Could gradient accumulation be the reason?  I also tried out using gradient accumulation of 2 in distributed training and batch size 2 per GPU for DDP which didn't help either.

- Another specialty about Wav2Vec2 is that the first 7 Conv layers are frozen here: https://github.com/huggingface/transformers/blob/ea92136597c49a20c5e2c31ef20ccec1693a8858/examples/research_projects/wav2vec2/run_common_voice.py#L454 which calls this function: https://github.com/huggingface/transformers/blob/ea92136597c49a20c5e2c31ef20ccec1693a8858/src/transformers/models/wav2vec2/modeling_wav2vec2.py#L1416. Could it be that in DDP just one of the eight models does that, but not the rest? (Actually I could easily check this...)

- Could it be that fp16 amp scaling somehow messes differently in DDP than in single GPU training? 

- ... other possible reasons?


@stas00 @sgugger - Have you previously heard about this kind of problem before (that single GPU works but DDP doesn't?). Think it's very hard to debug or dive into this problem, but I thought maybe you have some useful next step debugging strategies or tips! 

@anton-l - have you used DDP training during the Wav2Vec2 sprint? I've pretty much only used single GPU training which works well, but not DDP... have you had similar problems before? 

I've also tried running DDP on other/bigger datasets without success, so I'm a bit confused why it doesn't work here. Think the CTCLoss: https://pytorch.org/docs/stable/generated/torch.nn.CTCLoss.html is quite special and definitely more prone to instabilities than simple CE loss, but it's still very surprising to me that I get single GPU rather easily working but DDP not at all.

Some things I was planning on trying out next:

- Running DDP with just 2 GPUs to see whether the more GPUs, the more instable the loss becomes or not...

Do you guys have maybe any other good debugging strategies ?


Comments:
> * Could gradient-checkpointing be the reason? I really don't see how this could make a difference though...

I'm dealing with a somewhat similar report here: https://github.com/huggingface/transformers/issues/13653 - and I requested the user to repeat the experiment with gradient-checkpointing off. I don't think this feature is used a lot, perhaps we have bugs in there? I don't remember using it ever myself in `transformers` - we have been using it in Megatron-Deepspeed all the time. so it's a bit of an unknown to me on the HF side.

> Could gradient accumulation be the reason? 

gradient accumulation should be a pretty solid feature.

- I'd inspect whether `LayerDrop` and/or `weight_norm` could be impacting this under DDP.

Any difference if you use a similar Deepspeed setup? If you remember I had to tweak LayerDrop to synchronize the gpus there. Perhaps the data goes out of sync under DDP due to this feature? Do we know that it even works under DDP?

> Do you guys have maybe any other good debugging strategies ?

1. Making a tiny 2-gpu setup where you can clearly and quickly observe the different rate at loss diminishing (so you have 1 and 2-gpu setup) - so smaller model so it's a fast cycle
2. Removing all the configurable suspects, so you start with a baseline of 1 gpu and 2-gpu DDP giving very similar converging rate and adding those features back one by one or in groups until you see a clear difference emerging




Thanks for your reply @stas00! Regarding `gradient_checkpointing` - the problem is that it **does work** with gradient checkpointing on single GPU, but **does not work** with or without gradient checkpointing on multi-GPU.

`Layerdrop` is disabled as well. `weight_norm` could definitely be a reason! I'll verify this!

Just ran both experiments also in full fp32 precision and it's the same problem:

Single GPU: https://wandb.ai/patrickvonplaten/huggingface/reports/Wav2Vec2-1-GPU-V100-FP32--VmlldzoxMDQxODMw

8 GPUs: https://wandb.ai/patrickvonplaten/huggingface/reports/Wav2Vec2-DistributedDataParallel-DDP-8-GPU-V100-FP32--VmlldzoxMDQxODI2

=> So it also doesn't really seem to be related to `amp` or gradient scaling

DDP with just 2 GPUs also doesn't work. Will dive a bit deeper into it then - think it's quite important to get DDP working for Wav2Vec2

I have some question around that piece of code:
```
        # divide gradients by number of labels
        if self.args.world_size > 1:
            num_losses = (inputs["labels"] >= 0).sum()
            dist.all_reduce(num_losses)

            constant = self.args.world_size / num_losses

            self.multiply_grads(model.module.parameters(), constant)

            loss *= constant
```
It's only executed in a distributed setup but you don't multiply all gradients by `1 /  num_losses` when there is only one process. I have no idea of the size of your labels tensors and the number of positive elements it has, but it seems like it could be big (batch_size * sequence_length - the number of padded elements at first glance?)

Yeah, in the training runs above, I actually disabled this (I just use `Trainer` instead of `CTCTrainer`). 
I copied that code more or less from fairseq's Trainer. The idea here to only use `ctc_loss_reduction="mean"` in single GPU setup, but then use `ctc_loss_reduction="sum"` in the DDP setup and sum all losses and later scale the gradients correctly. 

With this code, all local copies of the model will get a gradient d(sum(loss_1)) /d(params) -> so that the average reduced gradient is d(sum(loss_1) + sum(loss_2) + ....)/8 * d(params)) with 8 being the world_size. Then I multiply by 8 and divide by all losses (batch_size * seq_length of gpu1 + batch_size * seq_length of gpu_2 + ....) which should then make the gradients identical to `ctc_loss_reduction="mean"` in single GPU setup.

I tried this out a couple of times, but it didn't solve the problem and also given that the sequence lengths in common voice are quite similar just using `ctc_loss_reduction="mean"` should work fine as well (see the first possible reason for a bug above).

=> So in short I currently don't use that code as I just use `Trainer` instead of `CTCTrainer` (sorry should probably have commented out `CTCTrainer` completely - see: https://github.com/huggingface/transformers/blob/97936d3aacc04f6253ff178415b8a57768fc8ce6/examples/pytorch/speech-recognition/run_speech_recognition_ctc.py#L511

Issues is solved. Will post a more detailed reason as an explanation

This was the problem essentially: https://github.com/huggingface/transformers/pull/13620#discussion_r714154511

Every process created a different vocabulary

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/master/CONTRIBUTING.md) are likely to be ignored.