Seq2SeqTrainer Model parallelism with AWS Sagemaker - not enough values to unpack error
## Environment info
- `transformers` version: 4.6
- Platform: 'ml.p3.16xlarge'  AWS DLC container
- Python version: 3.6
- PyTorch version (GPU?): 1.7 (True)
- Using GPU in script?: True
- Using distributed or parallel set-up in script?: Yes

### Who can help
@philschmid @patrickvonplaten

## Information

Model I am using (Bert, XLNet ...):    allenai/led-base-16384

The problem arises when using:
* [ ] my own modified scripts: (give details below)
I am trying to use AWS Sagemaker DLC container to train LED-base in a model parallel way using Seq2SeqTrainer on an 'ml.p3.16xlarge' instance on a custom dataset for the task of summarization. I am getting a 'not enough values to unpack error'.


![HF_model_parallel_err_jul17](https://user-images.githubusercontent.com/54971275/126030240-044d8440-db86-46a4-9769-bff21cc7ef16.png)

## To reproduce
Here are the settings used
```python
from sagemaker.huggingface import HuggingFace

hyperparameters = {
    'model_name': 'allenai/led-base-16384',
    'train_batch_size': 4,
    'eval_batch_size': 4,
    'num_train_epochs': 1
}
instance_count = 1
mpi_options = {
    'enabled': True,
    'processes_per_host': 8
}
smp_options = {
    'enabled': True,
    'parameters': {
        'microbatches': 4,
        'placement_strategy': 'spread',
        'pipeline': 'interleaved',
        'optimize': 'speed',
        'partitions': 4,
        'ddp': True
    }
}
distribution = {
    'smdistributed': {'modelparallel': smp_options},
    'mpi': mpi_options
}
```
```python
huggingface_estimator = HuggingFace(entry_point='led_train.py',
                                    source_dir='./scripts',
                                    instance_type='ml.p3.16xlarge',
                                    instance_count=instance_count,
                                    role=role,
                                    transformers_version='4.6',
                                    pytorch_version='1.7',
                                    py_version='py36',
                                    hyperparameters=hyperparameters,
                                    distribution=distribution
                                   )
```

Using this in the training script led_train.py
```python
from transformers import AutoModelForSeq2SeqLM, Seq2SeqTrainer, Seq2SeqTrainingArguments, AutoTokenizer
```
## Expected behavior

Expected model parallel training to happen smoothly

cc @sgugger<|||||>@suneelmatham do you need to run the LED-base in model parallel? Shouldn't the `base` version easily fit into memory? From the [hub](https://huggingface.co/allenai/led-base-16384) It only seems to be around 700MB. You could go with data parallelism to speed up your training. https://huggingface.co/docs/sagemaker/train#advanced-features<|||||>What is this `led_train.py` script? The error is not linked to the distributed training at first glance: it's just that the `input_ids` are not a 2d tensor (`batch_size` by `sequence_length`).<|||||>Thanks for the checking the issue out!
@philschmid I was doing a trial run with LED-base to test model parallel trng with Seq2SeqTrainer since the example notebook was with SageMakerTrainer. I plan to use LED-large for full run. I wanted to try model parallel since the dataset instances are long transcripts and hence was getting Out of memory error even with a really small batch size.

@sgugger led_train.py script is the training script which runs in an AWS DLC container and is triggered by the rest of the code which is being run from a sagemaker jupyter notebook. The same training script is running fine on a single AWS instance without using distributed options. I suspect the issue might be with using Seq2SeqTrainer with distributed options set. Specifically, the example notebook was using SageMakerTrainer as the trainer. I have read in HF release of 4.6  that Trainer itself can directly be used  instead of SageMakerTrainer. So, I assumed that Seq2SeqTrainer will work fine too with model parallelism

I am using this fn for tokenization
```python
def tokenize(batch):
    inputs = tokenizer(batch['src'], padding='max_length', truncation=True)
    outputs = tokenizer(batch['tgt'], padding='max_length', truncation=True, max_length=max_output_length)
    batch['input_ids'] = inputs.input_ids
    batch['attention_mask'] = inputs.attention_mask

    batch["global_attention_mask"] = len(batch["input_ids"]) * [[0 for _ in range(len(batch["input_ids"][0]))]]
    batch['global_attention_mask'][0][0] = 1
    batch['labels'] = outputs.input_ids
    batch['labels'] = [
                     [-100 if token == tokenizer.pad_token_id else token for token in labels]
                     for labels in batch['labels']
                    ]
    return batch
```
and the following for Seq2SeqTrainer
```python
    training_args = Seq2SeqTrainingArguments(
                        predict_with_generate=True,
                        evaluation_strategy='steps',
                        per_device_train_batch_size=args.train_batch_size,
                        per_device_eval_batch_size=args.eval_batch_size,
                        fp16=True,
                        output_dir=args.model_dir,
                        logging_steps=5,
                        eval_steps=10,
                        save_steps=10,
                        save_total_limit=2,
#                         gradient_accumulation_steps=4,
                        num_train_epochs=args.num_train_epochs,
                        report_to='wandb',
                        run_name='ml2xlarge model parallel',
#                         max_grad_norm= 0
                    )
    trainer = Seq2SeqTrainer(
        model=led,
        tokenizer=tokenizer,
        args=training_args,
        compute_metrics=compute_metrics,
        train_dataset=train_dataset,
        eval_dataset=val_dataset
    )
```
I assumed that Seq2SeqTrainer will handle batching of input_ids to pass to the model and it worked out fine when running on a single instance. Should I modify the tokenization fn. to fix this. Please let me know.<|||||>> @philschmid I was doing a trial run with LED-base to test model parallel trng with Seq2SeqTrainer since the example notebook was with SageMakerTrainer. I plan to use LED-large for full run. I wanted to try model parallel since the dataset instances are long transcripts and hence was getting Out of memory error even with a really small batch size.

I think data-parallel is definitely well more suited for this. The [allenai/led-large-16384](https://huggingface.co/allenai/led-large-16384) is only around 2GB and will be fast and scale better with only using data-parallel. The small batch sizes then can we scaled with scaling to multiple GPUs or even multiple nodes. We have a blog post [here](https://huggingface.co/blog/sagemaker-distributed-training-seq2seq) about how to do it or you can go to https://huggingface.co/docs/sagemaker/train#advanced-features and replace your `distribution` config to 
```bash
# configuration for running training on smdistributed Data Parallel
distribution = {'smdistributed':{'dataparallel':{ 'enabled': True }}}
```<|||||>Thanks. Will try with data parallel<|||||>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.