Repository: transformers
Issue: Multi-GPU training crashes with IterableDataset and different length input (e.g. Next token prediction)
Labels: bug

### System Info

- `transformers` version: 4.46.3
- Platform: Linux-6.6.20-aufs-1-x86_64-with-glibc2.36
- Python version: 3.11.2
- Huggingface_hub version: 0.26.1
- Safetensors version: 0.4.5
- Accelerate version: 1.2.1
- Accelerate config:    not found
- PyTorch version (GPU?): 2.5.1+cu124 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using distributed or parallel set-up in script?: yes, DDP. multi-gpu
- Using GPU in script?: yes
- GPU type: NVIDIA RTX A5000

### Who can help?

@muellerz @SunMarc 

### 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

I am using Trainer with hf `IterableDataset`. using multi-gpu training with at least 2 gpus.
I found out that i need `--accelerator_config {"split_batches":true}` for it work with varying length input (e.g. NTP prediction). I got this from https://github.com/huggingface/transformers/issues/26548

I am still having crashes in my case.
from what I understand the problem arises because my dataset has a reminder when considering the batch size.
I am getting the following stacktrace:

```
rank1]: Traceback (most recent call last):
[rank1]:   File "/cs/labs/oabend/avishai.elma/src/train_multimodal_multistream.py", line 249, in <module>
[rank1]:     main()
[rank1]:   File "/cs/labs/oabend/avishai.elma/src/train_multimodal_multistream.py", line 235, in main
[rank1]:     trainer.train()
[rank1]:   File "/cs/labs/oabend/avishai.elma/.lab_env_v3/lib/python3.11/site-packages/transformers/trainer.py", line 2123, in train
[rank1]:     return inner_training_loop(
[rank1]:            ^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/cs/labs/oabend/avishai.elma/.lab_env_v3/lib/python3.11/site-packages/transformers/trainer.py", line 2427, in _inner_training_loop
[rank1]:     batch_samples, num_items_in_batch = self.get_batch_samples(epoch_iterator, num_batches)
[rank1]:                                         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/cs/labs/oabend/avishai.elma/.lab_env_v3/lib/python3.11/site-packages/transformers/trainer.py", line 5045, in get_batch_samples
[rank1]:     batch_samples += [next(epoch_iterator)]
[rank1]:                       ^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/cs/labs/oabend/avishai.elma/.lab_env_v3/lib/python3.11/site-packages/accelerate/data_loader.py", line 842, in __iter__
[rank1]:     batch = concatenate([batch, first_batch], dim=0)
[rank1]:             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/cs/labs/oabend/avishai.elma/.lab_env_v3/lib/python3.11/site-packages/accelerate/utils/operations.py", line 621, in concatenate
[rank1]:     return type(data[0])({k: concatenate([d[k] for d in data], dim=dim) for k in data[0].keys()})
[rank1]:                          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/cs/labs/oabend/avishai.elma/.lab_env_v3/lib/python3.11/site-packages/accelerate/utils/operations.py", line 621, in <dictcomp>
[rank1]:     return type(data[0])({k: concatenate([d[k] for d in data], dim=dim) for k in data[0].keys()})
[rank1]:                              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/cs/labs/oabend/avishai.elma/.lab_env_v3/lib/python3.11/site-packages/accelerate/utils/operations.py", line 624, in concatenate
[rank1]:     return torch.cat(data, dim=dim)
[rank1]:            ^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: RuntimeError: Sizes of tensors must match except in dimension 0. Expected size 1008 but got size 1000 for tensor number 1 in the list.
[rank0]: Traceback (most recent call last):
[rank0]:   File "/cs/labs/oabend/avishai.elma/src/train_multimodal_multistream.py", line 249, in <module>
[rank0]:     main()
[rank0]:   File "/cs/labs/oabend/avishai.elma/src/train_multimodal_multistream.py", line 235, in main
[rank0]:     trainer.train()
[rank0]:   File "/cs/labs/oabend/avishai.elma/.lab_env_v3/lib/python3.11/site-packages/transformers/trainer.py", line 2123, in train
[rank0]:     return inner_training_loop(
[rank0]:            ^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/cs/labs/oabend/avishai.elma/.lab_env_v3/lib/python3.11/site-packages/transformers/trainer.py", line 2427, in _inner_training_loop
[rank0]:     batch_samples, num_items_in_batch = self.get_batch_samples(epoch_iterator, num_batches)
[rank0]:                                         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/cs/labs/oabend/avishai.elma/.lab_env_v3/lib/python3.11/site-packages/transformers/trainer.py", line 5045, in get_batch_samples
[rank0]:     batch_samples += [next(epoch_iterator)]
[rank0]:                       ^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/cs/labs/oabend/avishai.elma/.lab_env_v3/lib/python3.11/site-packages/accelerate/data_loader.py", line 842, in __iter__
[rank0]:     batch = concatenate([batch, first_batch], dim=0)
[rank0]:             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/cs/labs/oabend/avishai.elma/.lab_env_v3/lib/python3.11/site-packages/accelerate/utils/operations.py", line 621, in concatenate
[rank0]:     return type(data[0])({k: concatenate([d[k] for d in data], dim=dim) for k in data[0].keys()})
[rank0]:                          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/cs/labs/oabend/avishai.elma/.lab_env_v3/lib/python3.11/site-packages/accelerate/utils/operations.py", line 621, in <dictcomp>
[rank0]:     return type(data[0])({k: concatenate([d[k] for d in data], dim=dim) for k in data[0].keys()})
[rank0]:                              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/cs/labs/oabend/avishai.elma/.lab_env_v3/lib/python3.11/site-packages/accelerate/utils/operations.py", line 624, in concatenate
[rank0]:     return torch.cat(data, dim=dim)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: RuntimeError: Sizes of tensors must match except in dimension 0. Expected size 1008 but got size 1000 for tensor number 1 in the list.
  0%|                                                                                                                                                                                                                  | 3/43040 [01:02<248:43:41, 20.81s/it]
[rank0]:[W1217 12:41:08.453747481 ProcessGroupNCCL.cpp:1250] Warning: WARNING: process group has NOT been destroyed before we destruct ProcessGroupNCCL. On normal program exit, the application should call destroy_process_group to ensure that any pending NCCL operations have finished in this process. In rare cases this process can exit before this point and block the progress of another member of the process group. This constraint has always been present,  but this warning has only been added since PyTorch 2.4 (function operator())
W1217 12:41:24.654000 1980229 torch/distributed/elastic/multiprocessing/api.py:897] Sending process 1980456 closing signal SIGTERM
E1217 12:41:24.771000 1980229 torch/distributed/elastic/multiprocessing/api.py:869] failed (exitcode: 1) local_rank: 1 (pid: 1980457) of binary: /cs/labs/oabend/avishai.elma/.lab_env_v3/bin/python
Traceback (most recent call last):
  File "/cs/labs/oabend/avishai.elma/.lab_env_v3/bin/torchrun", line 8, in <module>
    sys.exit(main())
             ^^^^^^
  File "/cs/labs/oabend/avishai.elma/.lab_env_v3/lib/python3.11/site-packages/torch/distributed/elastic/multiprocessing/errors/__init__.py", line 355, in wrapper
    return f(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^
  File "/cs/labs/oabend/avishai.elma/.lab_env_v3/lib/python3.11/site-packages/torch/distributed/run.py", line 919, in main
    run(args)
  File "/cs/labs/oabend/avishai.elma/.lab_env_v3/lib/python3.11/site-packages/torch/distributed/run.py", line 910, in run
    elastic_launch(
  File "/cs/labs/oabend/avishai.elma/.lab_env_v3/lib/python3.11/site-packages/torch/distributed/launcher/api.py", line 138, in __call__
    return launch_agent(self._config, self._entrypoint, list(args))
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/cs/labs/oabend/avishai.elma/.lab_env_v3/lib/python3.11/site-packages/torch/distributed/launcher/api.py", line 269, in launch_agent
    raise ChildFailedError(
torch.distributed.elastic.multiprocessing.errors.ChildFailedError: 
============================================================
train_multimodal_multistream.py FAILED
------------------------------------------------------------
Failures:
  <NO_OTHER_FAILURES>
------------------------------------------------------------
Root Cause (first observed failure):
[0]:
  time      : 2024-12-17_12:41:24
  host      : binky-03.cs.huji.ac.il
  rank      : 1 (local_rank: 1)
  exitcode  : 1 (pid: 1980457)
  error_file: <N/A>
  traceback : To enable traceback see: https://pytorch.org/docs/stable/elastic/errors.html
============================================================
```

from what i understand this is because of those lines
https://github.com/huggingface/accelerate/blob/200c9eb7833cfa505907f6f224ebf5a275aa6d92/src/accelerate/data_loader.py#L840-L844

so i noticed that droping the remeinder might fix the issue. but it didn't help. this is probably because dataloader_drop_last doesn't get used when using `IterableDataset` 

https://github.com/huggingface/transformers/blob/5d7739f15a6e50de416977fe2cc9cb516d67edda/src/transformers/trainer.py#L997-L1009



### Expected behavior

Multi-gpu training should not crash when using iterable dataset