LED Model returns AlgorithmError when using SageMaker SMP training
### System Info

```shell
using sagemaker 
mpi_options = {
    "enabled" : True,
    "processes_per_host" : 8
}

smp_options = {
    "enabled":True,
    "parameters": {
        "microbatches": 1,
        "placement_strategy": "spread",
        "pipeline": "interleaved",
        "optimize": "memory",
        "partitions": 2,
        "ddp": True,
    }
}

distribution={
    "smdistributed": {"modelparallel": smp_options},
    "mpi": mpi_options
}
hyperparameters={'epochs': 1,
                 'train_batch_size': 1,
                 'eval_batch_size': 1,
                 'model_name':HHousen/distil-led-large-cnn-16384,
                 'output_dir': 'bucket',
                 'warmup_steps': 25,
                 'checkpoint_s3_uri': 'bucket',
                 'logging_steps':100,
                 'evaluation_strategy':"steps",
                 'gradient_accumulation_steps':10
                 }
huggingface_estimator = HuggingFace(entry_point='trainer.py',
                            source_dir='./scripts',
                            instance_type='ml.p3.16xlarge',
                            instance_count=1,
                            role=role,
                            volume=100,
                            transformers_version='4.6.1',
                            pytorch_version='1.8.1',
                            py_version='py36',
                            hyperparameters=hyperparameters,
                                   distribution=distribution)
```


### Who can help?

@ydshieh @sgugger

### Information

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

### Tasks

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

### Reproduction

1. Create huggingface estimator
2.     training_args = Seq2SeqTrainingArguments(
        predict_with_generate=True,
        evaluation_strategy="steps",
        per_device_train_batch_size=1,
        per_device_eval_batch_size=1,
        fp16=True,
        fp16_backend="apex",
        output_dir=s3_bucket,
        logging_steps=50,
        warmup_steps=25,
        gradient_accumulation_steps=10,
    )

Error I get:
[1,0]<stderr>:  File "/opt/conda/lib/python3.6/site-packages/smdistributed/modelparallel/torch/patches/tracing.py", line 68, in trace_forward
[1,0]<stderr>:    raise e
[1,0]<stderr>:  File "/opt/conda/lib/python3.6/site-packages/smdistributed/modelparallel/torch/patches/tracing.py", line 51, in trace_forward
[1,0]<stderr>:    output = original_forward(self, *args, **kwargs)
[1,0]<stderr>:  File "/opt/conda/lib/python3.6/site-packages/transformers/models/led/modeling_led.py", line 125, in forward
[1,0]<stderr>:    return super().forward(positions)
[1,0]<stderr>:  File "/opt/conda/lib/python3.6/site-packages/smdistributed/modelparallel/torch/patches/tracing.py", line 68, in trace_forward
[1,0]<stderr>:    raise e
[1,0]<stderr>:  File "/opt/conda/lib/python3.6/site-packages/smdistributed/modelparallel/torch/patches/tracing.py", line 51, in trace_forward
[1,0]<stderr>:    output = original_forward(self, *args, **kwargs)
[1,0]<stderr>:  File "/opt/conda/lib/python3.6/site-packages/transformers/models/led/modeling_led.py", line 121, in forward
[1,0]<stderr>:    bsz, seq_len = input_ids_shape[:2]
[1,0]<stderr>:ValueError: not enough values to unpack (expected 2, got 1)
--------------------------------------------------------------------------
Primary job  terminated normally, but 1 process returned
a non-zero exit code. Per user-direction, the job has been aborted.
--------------------------------------------------------------------------
--------------------------------------------------------------------------
mpirun.real detected that one or more processes exited with non-zero status, thus causing
the job to be terminated. The first process to do so was:
  Process name: [[41156,1],0]
  Exit code:    1
--------------------------------------------------------------------------


### Expected behavior

```shell
Training on a sagemaker notebook p3dn.24xlarge using fairscale `simple` and these versions
transformers-4.16.2
torch-1.10.2
fairscale-0.4.5
py37

I can successfully train the LED model with my training data. Trying to get it to work with Huggingface estimator and sagemaker SMP I would assume the same outcome.
```

cc @philschmid <|||||>I would also suggest @kanwari3 to
- try to use the same Python/PyTorch/transformers versions (and other libraries) on SageMaker that work locally (if possible)
- if the above doesn't work, try to use on local machine the same versions as those used on SageMaker, and see if you still get errors

So we have a better idea about if this is indeed a SageMaker issue or libraries issue<|||||>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.<|||||>cc @philschmid  , cc @ydshieh ,  cc @sgugger 
Hi, 

This is a follow up on this post with the same title. We are trying to fix the issue and are still getting the same error after trying out several fixes including matching the python, transformers, and pytorch versions according to the recommendations (3.8, 4.16.2, and 1.10.2, respectively):

-ValueError: not enough values to unpack (expected 2, got 1)

The error is in the “modeling_led” within the transformers module expecting a different input_ids shape. We tried unsqueezing the input_ids and attention_masks but it didn’t fix the error.

New Update is we tried below to unsqueeze input tensors to the "modeling_led" to solve the above error:
def unsqueeze_col(example):
    return {"input_ids": torch.unsqueeze(example["input_ids"], 0)}
pubmed_train = pubmed_train.map(unsqueeze_col)

I’d greatly appreciate your feedback. Please let me know if you need any further information about the project.