Multi-GPU training crashes with IterableDataset and different length input (e.g. Next token prediction)
Author: avishaiElmakies
Pull request: no
Labels: name: Good Second Issue

name: 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

Comments:
cc @muellerzr @SunMarc 

Can you try to choose a batch_size that doesn't end up with a  a remainder with your dataset to make sure that this is the actual issue ? 

I have tried it with dataset that doesn't have reminder (by using take()). It seemed to have worked fine. 

@SunMarc is this being handled? it would be great to be able to use iterableDataset with DDP

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.

@SunMarc 


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.

@SunMarc @ArthurZucker 

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.

@SunMarc @ArthurZucker 

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.

@SunMarc @ArthurZucker 

@avishaiElmakies I had a similar which I managed to fix. Can you share which dataloader you are using ? 
Is it a custom one or the default ?

It was the default one

@avishaiElmakies Interesting!

> [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.


Thi error doesnt get fixed by drop_last in dataloader because probably your some sequences in your batches are of unequal length. 
Mainly it arises due to padding issue. 
Can you confirm if by using default dataloader, padding is correctly applied ? 

I don't believe this is a padding issue for me, it was at the end of the training(it worked fine until the last batch)
I used an iterable dataset and a collator that handles the padding. It worked for training using normal dataset and iterable dataset except it crashed at the end with the iterable dataset. For me, the crash was also at a line that was guarded by 'not self._drop_last' so it should be fixed by changing how the last batch is handled 

@avishaiElmakies 
Thanks for confirming.
Can you share your `trainingArguments`, accelerate config (you can get it  with accelerate env)and how you prepare data by subclassing `IterableDataset`? 
That can help diagnose. 

unfortunately i don't have access to that code and data anymore so it will be hard to give you full details.
I did provide the details needed in the original bug report, i didn't use accelerate i used torchrun, but added `--accelerator_config {"split_batches":true}` in the training args. I didn't subclass 'iterable Dataset'. I used hf `load_dataset` with `streaming=true` and it was a tokenized dataset (each sample was a list of ints)  training NTP with a causal language model. Each file was a json

> unfortunately i don't have access to that code and data anymore so it will be hard to give you full details. I did provide the details needed in the original bug report, i didn't use accelerate i used torchrun, but added `--accelerator_config {"split_batches":true}` in the training args. I didn't subclass 'iterable Dataset'. I used hf `load_dataset` with `streaming=true` and it was a tokenized dataset (each sample was a list of ints) training NTP with a causal language model. Each file was a json

@KeshavSingh29 similar settings with me, I used accelerator and `load_dataset` with `streaming=True`, `{"split_batches": True}`, and got same exactly same error with multi-gpus.
My training arguments like,
`    training_args = TrainingArguments(
        output_dir=args.output_dir,
        bf16=True,
        learning_rate=args.learning_rate,
        weight_decay=args.weight_decay,
        warmup_steps=args.warmup_steps,
        # max_steps=args.max_steps,
        per_device_train_batch_size=args.per_device_train_batch_size,
        per_device_eval_batch_size=args.per_device_eval_batch_size,
        gradient_accumulation_steps=args.grad_accum,
        optim="adamw_torch",
        lr_scheduler_type="cosine",
        logging_steps=args.log_steps,
        # save_strategy="epoch",
        save_steps=args.save_steps,
        num_train_epochs=args.num_train_epochs,
        save_total_limit=3,
        eval_steps=args.eval_steps if eval_ds_processed is not None else None,
        report_to=["tensorboard"],  # or ["wandb"]
        ddp_find_unused_parameters=False,
        dataloader_num_workers=args.num_workers,
        gradient_checkpointing=args.gradient_checkpointing,
        max_steps=max_steps,
        average_tokens_across_devices=True,
        seed=args.seed,
        accelerator_config={"split_batches": True},
    )`

@SunMarc  I'd like to work on this. Planning to reproduce locally first.



worth a glance at #45507, #47333 — close neighbors in this repo. Don't want to rehash what's already here — happy to dig more if useful.

— AMA