DebugUnderflowOverflow crashes with Multi-GPU training
Author: kpich
Pull request: no

## Environment info
<!-- You can run the command `transformers-cli env` and copy-and-paste its output below.
     Don't forget to fill out the missing fields in that output! -->

- `transformers` version: 4.8.2
- Platform: Linux-4.15.0-29-generic-x86_64-with-Ubuntu-18.04-bionic
- Python version: 3.6.9
- PyTorch version (GPU?): 1.9.0+cu102 (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?: Yes
- Using distributed or parallel set-up in script?: Yes

### Who can help
@stas00 
@sgugger 

<!-- Your issue will be replied to more quickly if you can figure out the right person to tag with @
 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.


Models:

- albert, bert, xlm: @LysandreJik
- blenderbot, bart, marian, pegasus, encoderdecoder,  t5: @patrickvonplaten, @patil-suraj
- longformer, reformer, transfoxl, xlnet: @patrickvonplaten
- fsmt: @stas00
- funnel: @sgugger
- gpt2: @patrickvonplaten, @LysandreJik
- rag: @patrickvonplaten, @lhoestq
- tensorflow: @Rocketknight1

Library:

- benchmarks: @patrickvonplaten
- deepspeed: @stas00
- ray/raytune: @richardliaw, @amogkam
- text generation: @patrickvonplaten
- tokenizers: @LysandreJik
- trainer: @sgugger
- pipelines: @LysandreJik

Documentation: @sgugger

Model hub:

- for issues with a model report at https://discuss.huggingface.co/ and tag the model's creator.

HF projects:

- datasets: [different repo](https://github.com/huggingface/datasets)
- rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)

Examples:

- maintained examples (not research project or legacy): @sgugger, @patil-suraj
- research_projects/bert-loses-patience: @JetRunner
- research_projects/distillation: @VictorSanh

 -->

## Information

Model I am using (Bert, XLNet ...): Bart (but this is inessential)

The problem arises when using:
* [ ] the official example scripts: (give details below)
* [x] my own modified scripts: (give details below)

The tasks I am working on is:
* [ ] an official GLUE/SQUaD task: (give the name)
* [x] my own task or dataset: (give details below)

## To reproduce

Steps to reproduce the behavior:

1. Instantiate a `debug_utils.DebugUnderflowOverflow` for a model (I do this with `debug="underflow_overflow"` to a `TrainingArguments`).
2. Train using multi-GPU setup.
3. the debug hook added by `DebugUnderflowOverflow` to run after `forward()` will crash because of a bad lookup in the class's `module_names` dict.

This does not happen on single-GPU training.

Here's an example that crashes if I run on my (4-GPU) machine, but does **not** crash if I restrict to a single GPU (by calling `export CUDA_VISIBLE_DEVICES=0` before invoking the script):

```
from torch.utils.data import Dataset
from transformers import (BartForConditionalGeneration, BartModel, BartConfig,
                          Seq2SeqTrainingArguments, Seq2SeqTrainer)


class DummyDataset(Dataset):
  def __len__(self):
    return 5

  def __getitem__(self, idx):
    return {'input_ids': list(range(idx, idx+3)),
            'labels': list(range(idx, idx+3))}


def main():
  train_dataset = DummyDataset()
  config = BartConfig(vocab_size=10, max_position_embeddings=10, d_model=8,
                      encoder_layers=1, decoder_layers=1,
                      encoder_attention_heads=1, decoder_attention_heads=1,
                      decoder_ffn_dim=8, encoder_ffn_dim=8)
  model = BartForConditionalGeneration(config)
  args = Seq2SeqTrainingArguments(output_dir="tmp", do_train=True,
                                  debug="underflow_overflow")
  trainer = Seq2SeqTrainer(model=model, args=args, train_dataset=train_dataset)
  trainer.train()

if __name__ == '__main__':
  main()
```

I get the following stack trace on my multi-GPU machine:

```
Traceback (most recent call last):
  File "./kbp_dbg.py", line 29, in <module>
    main()
  File "./kbp_dbg.py", line 26, in main
    trainer.train()
  File "[...]/transformers/trainer.py", line 1269, in train
    tr_loss += self.training_step(model, inputs)
  File "[...]/transformers/trainer.py", line 1762, in training_step
    loss = self.compute_loss(model, inputs)
  File "[...]/transformers/trainer.py", line 1794, in compute_loss
    outputs = model(**inputs)
  File "[...]/torch/nn/modules/module.py", line 1051, in _call_impl
    return forward_call(*input, **kwargs)
  File "[...]/torch/nn/parallel/data_parallel.py", line 168, in forward
    outputs = self.parallel_apply(replicas, inputs, kwargs)
  File "[...]/torch/nn/parallel/data_parallel.py", line 178, in parallel_apply
    return parallel_apply(replicas, inputs, kwargs, self.device_ids[:len(replicas)])
  File "[...]/torch/nn/parallel/parallel_apply.py", line 86, in parallel_apply
    output.reraise()
  File "[...]/torch/_utils.py", line 425, in reraise
    raise self.exc_type(msg)
KeyError: Caught KeyError in replica 0 on device 0.
Original Traceback (most recent call last):
  File "[...]/torch/nn/parallel/parallel_apply.py", line 61, in _worker
    output = module(*input, **kwargs)
  File "[...]/torch/nn/modules/module.py", line 1071, in _call_impl
    result = forward_call(*input, **kwargs)
  File "[...]/transformers/models/bart/modeling_bart.py", line 1308, in forward
    return_dict=return_dict,
  File "[...]/torch/nn/modules/module.py", line 1071, in _call_impl
    result = forward_call(*input, **kwargs)
  File "[...]/transformers/models/bart/modeling_bart.py", line 1173, in forward
    return_dict=return_dict,
  File "[...]/torch/nn/modules/module.py", line 1071, in _call_impl
    result = forward_call(*input, **kwargs)
  File "[...]/transformers/models/bart/modeling_bart.py", line 756, in forward
    inputs_embeds = self.embed_tokens(input_ids) * self.embed_scale
  File "[...]/torch/nn/modules/module.py", line 1076, in _call_impl
    hook_result = hook(self, input, result)
  File "[...]/transformers/debug_utils.py", line 246, in forward_hook
    self.create_frame(module, input, output)
  File "[...]/transformers/debug_utils.py", line 193, in create_frame
    self.expand_frame(f"{self.prefix} {self.module_names[module]} {module.__class__.__name__}")
KeyError: Embedding(10, 8, padding_idx=1)
```

And, again, it completes totally fine if you restrict visibility to a single GPU via `CUDA_VISIBLE_DEVICES` before running this.

Having looked into it, I strongly suspect what's happening is the following:
- the `DebugUnderflowOverflow` instantiated in `Trainer.train` populates a `module_names` dict from nn.Modules to names in its constructor [(link)](https://github.com/huggingface/transformers/blob/master/src/transformers/debug_utils.py#L182)
- If multi-gpu training is enabled, `nn.DataParallel` is called (I think [here in trainer.py](https://github.com/huggingface/transformers/blob/master/src/transformers/trainer.py#L943)).
- this calls `torch.nn.replicate` on its forward pass [(link)](https://github.com/pytorch/pytorch/blob/master/torch/nn/parallel/data_parallel.py#L167), which I believe calls `nn.Module._replicate_for_data_parallel()` [here](https://github.com/pytorch/pytorch/blob/master/torch/nn/parallel/replicate.py#L115), replicating the model once per GPU.
- This results in different `hash()` values for the replicated `nn.Module` objects across different GPUs, in general.
- Finally, the `module_names` lookup in `create_frame()` will crash on the GPUs, since the differently replicated `nn.Module` objects on which the forward-pass hooks are called now have different `hash()` values after multi-GPU replication.

(You can confirm that a module has different hash values after replication via the following snippet):
```
m = torch.nn.Embedding(7, 5)
m2 = m._replicate_for_data_parallel()
print(hash(m))
print(hash(m2))
```
Which will print something like:
```
8730470757753
8730462809617
```

Not sure what the best way to fix this is. If there's a way to instantiate a different `DebugUnderflowOverflow` for each GPU, after replication, that will maybe solve this issue (since per-replica hashes will assumedly be consistent then), but I'm not sure if that's feasible or the best way to do this. 

One could also just make the miscreant f-string construction in `create_frame` have a `.get()` with a default as the `module_names` lookup, rather than using square brackets, but this would probably make the debugging traces too uninformative. So I figured I'd just open a bug


## Expected behavior

I'd expect it to either not crash or print a fatal `overflow/underflow debug unsupported for multi-GPU` error or something. I originally thought this was a bug in the model code rather than a multi-GPU thing.


Comments:
Thank you for the great report and the analysis of the problem, @kpich.

I will have a look a bit later when I have time to solve this, including your suggestion.

Meanwhile, any special reason why you're using DP and not DDP (`torch.distributed.launch`)? There each GPU instantiates its own everything and it will work correctly. 

No, only reason I'm using DP rather than DDP is I'm using the `Trainer` framework with its default behaviors on a multi-GPU machine. Using DDP is a good suggestion, thanks!

DDP is typically faster than DP (as long as the interconnect is fast), please see: https://huggingface.co/transformers/performance.html#dp-vs-ddp

As a first step, here is a simple reproduction of the OP report with this cmd w/ 2 gpus over DP:

```
export BS=16; rm -r output_dir; PYTHONPATH=src USE_TF=0 CUDA_VISIBLE_DEVICES=0,1 python \
examples/pytorch/translation/run_translation.py --model_name_or_path t5-small --output_dir \
output_dir --adam_eps 1e-06 --do_train --label_smoothing 0.1 --learning_rate 3e-5 \
--logging_first_step --logging_steps 500 --max_source_length 128 --max_target_length 128 \
--val_max_target_length 128 --num_train_epochs 1 --overwrite_output_dir \
--per_device_train_batch_size $BS --predict_with_generate --sortish_sampler --source_lang en \
--target_lang ro --dataset_name wmt16 --dataset_config "ro-en" --source_prefix "translate English to \
Romanian: " --warmup_steps 50 --max_train_samples 50 --debug underflow_overflow
```

gives a similar error.

Nothing comes to mind as a simple fix at the moment, so for now let's just do a clean assert as you suggested. https://github.com/huggingface/transformers/pull/12816

If someone is stuck and can't use DDP we will revisit this.

And of course, if you or someone would like to work on an actual solution it'd be very welcome. I don't see `nn.DataParallel` having any hooks, so most likely this will require overriding `torch.nn.parallel.data_parallel.replicate` to refresh the model references after the replication. You can see the source code here: https://pytorch.org/docs/stable/_modules/torch/nn/parallel/data_parallel.html#DataParallel So it can be done, but it is not really worth it, IMHO.