Layer_drop interferes with gradient_checkpoint operation in DDP.
Author: jp1924
Pull request: no
Labels: name: trainer

name: bug

name: Audio

### System Info

- `transformers` version: 4.39.3
- Platform: Linux-5.15.0-100-generic-x86_64-with-glibc2.35
- Python version: 3.10.12
- Huggingface_hub version: 0.22.0
- Safetensors version: 0.4.2
- Accelerate version: 0.28.0
- Accelerate config:    not found
- PyTorch version (GPU?): 2.2.2+cu121 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>

### Who can help?

@sanchit-gandhi 

### Information

- [ ] The official example scripts
- [x] My own modified scripts

### Tasks

- [x] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [x] My own task or dataset (give details below)

### Reproduction

```python
import os

os.environ["WANDB_DISABLED"] = "true"
os.environ["TORCH_DISTRIBUTED_DEBUG"] = "DETAIL"

import torch
from datasets import Dataset
from transformers import (
    HfArgumentParser,
    Trainer,
    TrainingArguments,
    Wav2Vec2Config,
    Wav2Vec2FeatureExtractor,
    Wav2Vec2ForPreTraining,
)
from transformers.models.wav2vec2.modeling_wav2vec2 import (
    _compute_mask_indices,
    _sample_negative_indices,
)


def main(args) -> None:
    def collator(features):
        batch = feature_extractor.pad(features, return_attention_mask=True, return_tensors="pt")

        batch["attention_mask"] = batch["attention_mask"].to(torch.long)
        batch_size = batch["input_values"].shape[0]
        audio_length = batch["input_values"].shape[-1]
        device = batch["input_values"].device

        mask_indices_seq_length = model._get_feat_extract_output_lengths(input_lengths=audio_length)
        mask_indices_seq_length = int(mask_indices_seq_length)

        feature_vector_attention_mask = model._get_feature_vector_attention_mask(
            mask_indices_seq_length, batch["attention_mask"]
        )
        mask_time_indices = _compute_mask_indices(
            (batch_size, mask_indices_seq_length),
            config.mask_time_prob,
            config.mask_time_length,
            attention_mask=feature_vector_attention_mask,
            min_masks=2,
        )
        sampled_negative_indices = _sample_negative_indices(
            (batch_size, mask_indices_seq_length),
            model.config.num_negatives,
            mask_time_indices=mask_time_indices,
        )
        batch["mask_time_indices"] = torch.tensor(mask_time_indices, dtype=torch.long, device=device)
        batch["sampled_negative_indices"] = torch.tensor(sampled_negative_indices, dtype=torch.long, device=device)

        return batch

    config = Wav2Vec2Config(layerdrop=0.5)
    model = Wav2Vec2ForPreTraining(config)
    feature_extractor = Wav2Vec2FeatureExtractor(
        return_attention_mask=True,
        sampling_rate=16000,
    )

    dummy_dataset = Dataset.from_list([{"input_values": [0.0] * (16000 * 2)} for _ in range(1000)])
    trainer = Trainer(
        model=model,
        args=args,
        train_dataset=dummy_dataset,
        data_collator=collator,
    )
    trainer.train()


if "__main__" in __name__:
    parser = HfArgumentParser([TrainingArguments])
    args, _ = parser.parse_args_into_dataclasses(return_remaining_strings=True)

    main(args)
```

run script
> torchrun --nproc_per_node=1 ./issue_code.py --output_dir=./dummy_dir --gradient_checkpointing=true --per_gpu_train_batch_size=2 --evaluation_strategy=no --report_to=none


### Expected behavior

Hello Huggingface Team!

I'm raising an issue because I found a bug while running Wav2Vec2 in DDP Pretrain environment.

When running gradient_checkpoint with Layer-Drop applied in the DDP environment, the following error occurs.

```python
Traceback (most recent call last):
  File "/root/issue_code.py", line 75, in <module>
    main(args)
  File "/root/issue_code.py", line 68, in main
    trainer.train()
  File "/usr/local/lib/python3.10/dist-packages/transformers/trainer.py", line 1780, in train
    return inner_training_loop(
  File "/usr/local/lib/python3.10/dist-packages/transformers/trainer.py", line 2118, in _inner_training_loop
    tr_loss_step = self.training_step(model, inputs)
  File "/usr/local/lib/python3.10/dist-packages/transformers/trainer.py", line 3036, in training_step
    loss = self.compute_loss(model, inputs)
  File "/usr/local/lib/python3.10/dist-packages/transformers/trainer.py", line 3059, in compute_loss
    outputs = model(**inputs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/parallel/distributed.py", line 1519, in forward
    inputs, kwargs = self._pre_forward(*inputs, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/parallel/distributed.py", line 1413, in _pre_forward
    if torch.is_grad_enabled() and self.reducer._rebuild_buckets():
RuntimeError: Expected to have finished reduction in the prior iteration before starting a new one. This error indicates that your module has parameters that were not used in producing loss. You can enable unused parameter detection by passing the keyword argument `find_unused_parameters=True` to `torch.nn.parallel.DistributedDataParallel`, and by 
making sure all `forward` function outputs participate in calculating loss. 
If you already have done the above, then the distributed data parallel module wasn't able to locate the output tensors in the return value of your module's `forward` function. Please include the loss function and the structure of the return value of `forward` of your module when reporting this issue (e.g. list, dict, iterable).
Parameters which did not receive grad for rank 0: wav2vec2.encoder.layers.11.final_layer_norm.bias, wav2vec2.encoder.layers.11.final_layer_norm.weight, wav2vec2.encoder.layers.11.feed_forward.output_dense.bias, wav2vec2.encoder.layers.11.feed_forward.output_dense.weight, wav2vec2.encoder.layers.11.feed_forward.intermediate_dense.bias, wav2vec2.encoder.layers.11.feed_forward.intermediate_dense.weight, wav2vec2.encoder.layers.11.layer_norm.bias, wav2vec2.encoder.layers.11.layer_norm.weight, wav2vec2.encoder.layers.11.attention.out_proj.bias, wav2vec2.encoder.layers.11.attention.out_proj.weight, wav2vec2.encoder.layers.11.attention.q_proj.bias, wav2vec2.encoder.layers.11.attention.q_proj.weight, wav2vec2.encoder.layers.11.attention.v_proj.bias, wav2vec2.encoder.layers.11.attention.v_proj.weight, wav2vec2.encoder.layers.11.attention.k_proj.bias, wav2vec2.encoder.layers.11.attention.k_proj.weight, wav2vec2.encoder.layers.5.final_layer_norm.bias, wav2vec2.encoder.layers.5.final_layer_norm.weight, wav2vec2.encoder.layers.5.feed_forward.output_dense.bias, wav2vec2.encoder.layers.5.feed_forward.output_dense.weight, wav2vec2.encoder.layers.5.feed_forward.intermediate_dense.bias, wav2vec2.encoder.layers.5.feed_forward.intermediate_dense.weight, wav2vec2.encoder.layers.5.layer_norm.bias, wav2vec2.encoder.layers.5.layer_norm.weight, wav2vec2.encoder.layers.5.attention.out_proj.bias, wav2vec2.encoder.layers.5.attention.out_proj.weight, wav2vec2.encoder.layers.5.attention.q_proj.bias, wav2vec2.encoder.layers.5.attention.q_proj.weight, wav2vec2.encoder.layers.5.attention.v_proj.bias, wav2vec2.encoder.layers.5.attention.v_proj.weight, wav2vec2.encoder.layers.5.attention.k_proj.bias, wav2vec2.encoder.layers.5.attention.k_proj.weight, wav2vec2.encoder.layers.2.final_layer_norm.bias, wav2vec2.encoder.layers.2.final_layer_norm.weight, wav2vec2.encoder.layers.2.feed_forward.output_dense.bias, wav2vec2.encoder.layers.2.feed_forward.output_dense.weight, wav2vec2.encoder.layers.2.feed_forward.intermediate_dense.bias, wav2vec2.encoder.layers.2.feed_forward.intermediate_dense.weight, wav2vec2.encoder.layers.2.layer_norm.bias, wav2vec2.encoder.layers.2.layer_norm.weight, wav2vec2.encoder.layers.2.attention.out_proj.bias, wav2vec2.encoder.layers.2.attention.out_proj.weight, wav2vec2.encoder.layers.2.attention.q_proj.bias, wav2vec2.encoder.layers.2.attention.q_proj.weight, wav2vec2.encoder.layers.2.attention.v_proj.bias, wav2vec2.encoder.layers.2.attention.v_proj.weight, wav2vec2.encoder.layers.2.attention.k_proj.bias, wav2vec2.encoder.layers.2.attention.k_proj.weight, wav2vec2.encoder.layers.0.final_layer_norm.bias, wav2vec2.encoder.layers.0.final_layer_norm.weight, wav2vec2.encoder.layers.0.feed_forward.output_dense.bias, wav2vec2.encoder.layers.0.feed_forward.output_dense.weight, wav2vec2.encoder.layers.0.feed_forward.intermediate_dense.bias, wav2vec2.encoder.layers.0.feed_forward.intermediate_dense.weight, wav2vec2.encoder.layers.0.layer_norm.bias, wav2vec2.encoder.layers.0.layer_norm.weight, wav2vec2.encoder.layers.0.attention.out_proj.bias, wav2vec2.encoder.layers.0.attention.out_proj.weight, wav2vec2.encoder.layers.0.attention.q_proj.bias, wav2vec2.encoder.layers.0.attention.q_proj.weight, wav2vec2.encoder.layers.0.attention.v_proj.bias, wav2vec2.encoder.layers.0.attention.v_proj.weight, wav2vec2.encoder.layers.0.attention.k_proj.bias, wav2vec2.encoder.layers.0.attention.k_proj.weight
Parameter indices which did not receive grad for rank 0: 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210
  0%|▌                                                                                                                                                                                                                              | 1/375 [00:03<21:52,  3.51s/it]
[2024-04-08 07:54:33,477] torch.distributed.elastic.multiprocessing.api: [WARNING] Sending process 1973173 closing signal SIGTERM
[2024-04-08 07:54:35,798] torch.distributed.elastic.multiprocessing.api: [ERROR] failed (exitcode: 1) local_rank: 0 (pid: 1973171) of binary: /usr/bin/python3
```

I remember running DDP with layer_drop applied to Wav2Vec2 in the past (around 2022) and it worked without error.

However, with torch being updated to the latest version, it seems to have been modified to throw an error if some forward operations are skipped in DDP.

I don't know how to fix this right now, so I'm leaving an issue.

Comments:
also cc @ylacombe 

cc @sanchit-gandhi @ylacombe 

The issue has been there for a while now, e.g it is mentioned in #11452, and a reason is given here:

https://github.com/huggingface/transformers/issues/11452#issuecomment-851603295

It seems to be inherent to torch DDP, I don't believe there's anything we can do about it in transformers, but I'd be happy to be wrong and to propose a PR if someone has an idea of how to solve this!

I'm closing this issue for now, as it is related to how DDP work, feel free to re-open if you have an idea on how to solve it!

@ylacombe I think this should remain open: a solution hasn't yet been found and what the user is reporting is a regression in behaviour, which should be possible to track down to a commit 

Oh sorry for the delayed reply.

Thanks for sharing your issue!

I've been running it too, but haven't found a solution, so I'm running it with gradient_checkpoint set to false.

If I have anything further, I'll upload it to the issue here.

Thanks!

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/main/CONTRIBUTING.md) are likely to be ignored.

@ylacombe Sorry, I misread this issue previously. As the issue is with torch and not Hugging Face then it's fine to close!