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

cc @philschmid  , cc @ydshieh  , cc @sgugger 

Hello,

This is a follow up on a related post with the below link) with the same title:
https://github.com/huggingface/transformers/issues/16890

We ade a bit of more progress but are still facing with some issues and are trying to fix them 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. 

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)


It helped moving forward in the process, but we got another error, below, a little further down in the code:

UnexpectedStatusException: Error for Training job huggingface-pytorch-training-2022-06-29-04-04-58-606: Failed. Reason: AlgorithmError: ExecuteUserScriptError:
ExitCode 1
ErrorMessage ":RuntimeError: Tensors must have same number of dimensions: got 4 and 3
 :Environment variable SAGEMAKER_INSTANCE_TYPE is not set :Environment variable SAGEMAKER_INSTANCE_TYPE is not set :Environment variable SAGEMAKER_INSTANCE_TYPE is not set :Environment variable SAGEMAKER_INSTANCE_TYPE is not set :Environment variable SAGEMAKER_INSTANCE_TYPE is not set :Environment variable SAGEMAKER_INSTANCE_TYPE is not set :Environment variable SAGEMAKER_INSTANCE_TYPE is not set :Environment variable SAGEMAKER_INSTANCE_TYPE is not set -------------------------------------------------------------------------- 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: [[41154,1],0]   Exit code:    1"
Command "mpirun --host algo-1:8 


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

### Who can help?

[SageMakerAprilTraining.zip](https://github.com/huggingface/transformers/files/9065968/SageMakerAprilTraining.zip)


### Information

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

### Tasks

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

### Reproduction

Running this attached file with the training python file

### Expected behavior

I have shared the notebook and the error raised in it for clarification
@omid0001 @kanwari3, 

Would it be possible for you to reproduce this issue (`not enough values to unpack`) without using SageMaker, i.e. just with a Python script?
```bash
[1,0]: bsz, seq_len = input_ids_shape[:2]
[1,0]:ValueError: not enough values to unpack (expected 2, got 1)
```

It would be a good idea to verify what data is received by the model first. Usually the batches in data (`input_ids`) should be already of the format `(batch_size, sequence_length)`, and if you see the above error, it is likely the data or its processing pipeline has some issues. Using `torch.unsqueeze` is not really a good idea, as it implies you have only `batch_size` being 1.

My suggestion:
- Try to run your training without SageMaker (and without the using the fix `torch.unsqueeze`)
- Check what is received by the model, and check in the data pipeline if it prepares the correct input format
  - If you still get the issue and can't figure it out:
    - I could try to help if you could provide the training script + data processing script + a tiny portion of your data    
  - If the issue only occurs when you wrap the training in SageMaker, I don't have the competence to help in this case, sorry.<|||||>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.