Trainer errors out when concatenating different sequence length batches with distributed training and IterableDataset
Author: ssharpe42
Pull request: no

### System Info

- `transformers` version: 4.33.3
- Platform: Linux-5.10.186-179.751.amzn2.x86_64-x86_64-with-glibc2.10
- Python version: 3.8.17
- Huggingface_hub version: 0.17.3
- Safetensors version: 0.3.3
- Accelerate version: 0.23.0
- Accelerate config:    not found
- PyTorch version (GPU?): 2.0.1+cu117 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: A100
- Using distributed or parallel set-up in script?: torchrun --nproc-per-node 2 script.py

### Who can help?

 @muellerzr, @pacman100

### 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

```

    import torch
    from torch.utils.data import IterableDataset
    from transformers import (
        AutoModelForMaskedLM,
        AutoTokenizer,
        DataCollatorForLanguageModeling,
        Trainer,
        TrainingArguments,
    )

    data = [
        {
            "input_ids": torch.tensor([101, 2040, 2001, 1999, 14936, 102]),
            "token_type_ids": torch.tensor([0, 0, 0, 0, 0, 0]),
            "attention_mask": torch.tensor([1, 1, 1, 1, 1, 1]),
        },
        {
            "input_ids": torch.tensor([101, 2040, 102]),
            "token_type_ids": torch.tensor([0, 0, 0]),
            "attention_mask": torch.tensor([1, 1, 1]),
        },
        {
            "input_ids": torch.tensor([101, 2040, 2001, 1999]),
            "token_type_ids": torch.tensor([0, 0, 0, 0]),
            "attention_mask": torch.tensor([1, 1, 1, 1]),
        },
        {
            "input_ids": torch.tensor([101, 2040, 2001, 1999, 14936, 102]),
            "token_type_ids": torch.tensor([0, 0, 0, 0, 0, 0]),
            "attention_mask": torch.tensor([1, 1, 1, 1, 1, 1]),
        },
        {
            "input_ids": torch.tensor([101]),
            "token_type_ids": torch.tensor([00]),
            "attention_mask": torch.tensor([1]),
        },
        {
            "input_ids": torch.tensor([101]),
            "token_type_ids": torch.tensor([00]),
            "attention_mask": torch.tensor([1]),
        },
    ]

    class ExampleDataset(IterableDataset):
        def __init__(self, data):
            super().__init__()
            self.data = data * 20

        def __iter__(self):
            for x in self.data:
                yield x

        def __len__(self):
            return len(self.data)

    tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
    model = AutoModelForMaskedLM.from_pretrained("bert-base-cased")
    train_args = TrainingArguments(
        output_dir="output",
        num_train_epochs=3,
        per_device_train_batch_size=2,
    )
    dc = DataCollatorForLanguageModeling(tokenizer=tokenizer)

    trainer = Trainer(
        train_dataset=ExampleDataset(data),
        model=model,
        args=train_args,
        data_collator=dc,
    )
    trainer.train()

```

I run the above script with the command `torchrun --nproc-per-node 2 script.py`. This results in the following error. 

```
Traceback (most recent call last):
  File "fm_model/data/scratch.py", line 242, in <module>
    trainer.train()
  File "/opt/conda/envs/fmmodel/lib/python3.8/site-packages/transformers/trainer.py", line 1556, in train
    return inner_training_loop(
  File "/opt/conda/envs/fmmodel/lib/python3.8/site-packages/transformers/trainer.py", line 1816, in _inner_training_loop
    for step, inputs in enumerate(epoch_iterator):
  File "/opt/conda/envs/fmmodel/lib/python3.8/site-packages/accelerate/data_loader.py", line 597, in __iter__
    next_batch, next_batch_info = self._fetch_batches(main_iterator)
  File "/opt/conda/envs/fmmodel/lib/python3.8/site-packages/accelerate/data_loader.py", line 528, in _fetch_batches
    batch = concatenate(batches, dim=0)
  File "/opt/conda/envs/fmmodel/lib/python3.8/site-packages/accelerate/utils/operations.py", line 496, in concatenate
    return type(data[0])({k: concatenate([d[k] for d in data], dim=dim) for k in data[0].keys()})
  File "/opt/conda/envs/fmmodel/lib/python3.8/site-packages/accelerate/utils/operations.py", line 496, in <dictcomp>
    return type(data[0])({k: concatenate([d[k] for d in data], dim=dim) for k in data[0].keys()})
  File "/opt/conda/envs/fmmodel/lib/python3.8/site-packages/accelerate/utils/operations.py", line 499, in concatenate
    return torch.cat(data, dim=dim)
RuntimeError: Sizes of tensors must match except in dimension 0. Expected size 1 but got size 6 for tensor number 1 in the list.
```

This is due to the fact that in `Trainer` there are no arguments that can be passed to prepare the dataloader with [split_batches](https://github.com/huggingface/accelerate/blob/48d96319e0033fb8c8979072d97edf3995639029/src/accelerate/data_loader.py#L515) so this errors out when running this [line](https://github.com/huggingface/accelerate/blob/69e4c3c54da3201eda288b500d138761e7a5221c/src/accelerate/data_loader.py#L481). This occurs since there is no padding done across batches before these are concatenated together. 

In order to be able to use an iterable dataset with Trainer, something probably needs to be changed in accelerate or the Trainer to enable distributed dataloading when the batches end up being different lengths. 


### Expected behavior

1. Automatic padding in accelerate when the batches produced have different lengths
OR
2. A way to specify split_batches where a full batch is produced then split for all the different processes

Comments:
FYI, I think this was a regression introduced between transformers 4.30.x and 4.31.x (I haven't bisected the specific release, but 4.30.x will run the example here correctly)

Think I found it, this is where the IterableDatasetShard option is removed for IterableDatasets and it just relies on accelerate prepare. https://github.com/huggingface/transformers/commit/ebd94b0f6f215f6bc0f70e61eba075eb9196f9ef

This issue (https://github.com/huggingface/transformers/issues/26541) also occurs due to the same changes. 

Thanks @ssharpe42 and @dwyatte for flagging and diagnosing the issue here. I just wanted to echo that this is a major issue -- this problem is surprisingly hard to diagnose, not documented, and likely to affect many users who have no idea it is even a problem (for example [here](https://discuss.huggingface.co/t/sizes-of-tensors-must-match-except-in-dimension-0/56333) is one HF post on this topic with no responses as of the time of this post, but seemingly affected by the same problem).

Also, it seems quite likely that folks using `Trainer` with distributed training will be using `IterableDataset`s (I've also seen lots of uses of `Trainer` avoiding `IterableDataset`, maybe this is why?) -- the kinds of large training runs that require distributed training will very often also require large datasets that are best processed as `IterableDataset`s.

Thanks to the HF team, and hopefully this can be escalated; otherwise this is a complete blocker for many users!

also, confirming that downgrading to transformers 4.30 is a temporary workaround (also requires downgrading to accelerate==0.20.3) when using the Trainer API (doesn't appear to work outside the Trainer API with these versions)

We're aware of it and working towards a solution

I just ran into this one this evening. Thanks for staying on top of it!

My project for the weekend was to switch from using Dataset, using explicit sharding , to an IterableDataset for training; the dataset in question is huge, so it's one or the other. Switching to an IterableDataset is needed for switching from DDP to FSDP mode for training a larger model, as I can no longer get away with reinstantiating the Trainer after each shard, while reusing the same optimizer and lr-scheduler.

Anyhow, a quick and dirty work-around, until this is properly fixed, is to pad everything to the same length. This is hardly ideal, but it appears to work. Hopefully this is useful to someone.

Another piece of information here: I noticed that IterableDataset + DDP/FSDP works fine with the Trainer API, but not outside the Trainer API in my own training loop. I decided to dig into this a bit more to see what Trainer is doing differently. 

Looking through the Trainer code in transformers version 4.30.2 where it constructs the DataLoader in a distributed setting [(this code block)](https://github.com/huggingface/transformers/blob/66fd3a8d626a32989f4569260db32785c6cbf42a/src/transformers/trainer.py#L911), I noticed that the Trainer API never calls accelerator.prepare(data_loader)!

I tried this in my own training loop which was previously failing with the above error, and it works with FSDP over 8 40GB A40 GPUs, for both a tiny model and Llama 2. 

So, another workaround here: for users not using the Trainer API, *don't call accelerator.prepare() on the DataLoaders* (but still call accelerator.prepare() for everything else; i.e. model, optimizer, scheduler). This probably requires some extra logic to do the things accelerator.prepare might -- especially casting data types etc. -- but maybe it really is this simple, since it's exactly what the Trainer API was doing in transformers==4.30.2...

I would appreciate if the maintainers could provide any insights here into what the downsides of this approach might be (keeping in mind that this is transformers==4.30.2 and accelerate==0.20.3).

@jpgard that hasn't been the case for a while? See here where we call prepare: https://github.com/huggingface/transformers/blob/main/src/transformers/trainer.py#L829

Regardless, will be looking into this next week + before the holiday, sorry for the delay!

Any update right here? I came across the same issue when running DPO trainer if I use streaming datasets that have different sequence length.

Appreciate for the maintainers' efforts on this issue. Is there any update or temporal solution to fix this 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.

My temporal solution is to warp the original `get_train_dataloader() `function of transformers.Trainer to skip execute `self.accelerator.prepare() `for iterable datasets:

```
def get_train_dataloader(self) -> DataLoader:
        print("Inject Hacker dataloader prepare")
        if self.train_dataset is None:
            raise ValueError("Trainer: training requires a train_dataset.")
        train_dataset = self.train_dataset
        data_collator = self.data_collator
        if isinstance(train_dataset, datasets.Dataset):
            train_dataset = self._remove_unused_columns(train_dataset, description="training")
        else:
            data_collator = self._get_collator_with_removed_columns(data_collator, description="training")

        dataloader_params = {
            "batch_size": self._train_batch_size,
            "collate_fn": data_collator,
            "num_workers": self.args.dataloader_num_workers,
            "pin_memory": self.args.dataloader_pin_memory,
        }

        if not isinstance(train_dataset, torch.utils.data.IterableDataset):
            dataloader_params["sampler"] = self._get_train_sampler()
            dataloader_params["drop_last"] = self.args.dataloader_drop_last
            dataloader_params["worker_init_fn"] = seed_worker

            return self.accelerator.prepare(DataLoader(train_dataset, **dataloader_params))
        return DataLoader(train_dataset, **dataloader_params)
```

It works well for me. But not sure the potential impact on efficiency for skipping the `accelerator.prepare()`.

Gently pinging @muellerzr and @pacman100 

Hi everyone, the issue is not directly linked to `IterableDataset` but more on the `dispatch_batches` arg that is set to `True` when passing a `IterableDataset`. 

With `dispatch_batches=True`, we process the data on the main process and broadcast to the other processes. This is more reliable + less compute since we do it on only one process. However, we indeed have an issue when the data doesn't have the same size as we are trying to concat tensors that might not have the same size. 

As a temporary solution, you can either:
- pass `dispatch_batches=False` in `TrainingArguments`. This is the default behavior of Trainer before accelerate integration. It will use `IterableDatasetShard` instead of `DataLoaderDispatcher`. 
- pass `split_batches=True`. It will split a full batch into `self.num_process` parts. So, it requires that the batch size ( `per_device_train_batch_size` ) of the `dataloader` to be a round multiple of the number of processes. For example, if you set `per_device_train_batch_size = 16` and you have 4 processes, each process will have a batch_size of 4. 

Related code: 

```python 
if self.split_batches:
    # One batch of the main iterator is dispatched and split.
    batch = next(iterator)
else:
    # num_processes batches of the main iterator are concatenated then dispatched and split.
    # We add the batches one by one so we have the remainder available when drop_last=False.
    batches = []
    for _ in range(self.state.num_processes):
        batches.append(next(iterator))
   # The issue is here since the batches do not necessarily have the same size. 
    batch = concatenate(batches, dim=0)
```


We've released a patch in Accelerate which will give a more clear error on what's going on, and will be doing similar into the Trainer here shortly to give direct instructions on what to do :) 

thanks for the attention to this!

it is not entirely clear to me whether this issue is resolved or not in the current versions of transformers/accelerate? should we now expect IterableDataset to work in distributed training (both with Trainer, and without Trainer), by following one of the two options @SunMarc mentioned?

They always have, there's no change needed in the source code. You need to do either of Marc's two solutions, which have existed for a large number of months :) 

Yes. 

> They always have

that is not my experience, and is why the issue exists in the first place :) thanks for pointing us to the workaround


@SunMarc Thank you for the explainations and suggested workarounds! 
By the way, may I ask if `per_device_train_batch_size` refers to the batch size per GPU? If it is, when using `split_batches` for a training setting with 4 GPUs and 32 images per GPU, should I set it as 32 or 32*4? I checked `trainer.py` and it seems that the trainer did not handle the different batch size with different values of `split_batches` automatically.

Yes, from my experiments, you should set it as `per_device_train_batch_size = 32*4`. 

If you were using dispatch_batches=False and you couldn't do epoch based training, @muellerzr added this so now that route is feasible with accelerate + iterable datasets https://github.com/huggingface/accelerate/pull/2066

Hello, I'm using the streaming dataloader with dispatch_batches=False. Although the training can run for some steps, it is common that the training process is stuck where 1 GPU's util is 0% but others are 100%. After 30 minutes, training fails with an NCCL timeout error.

My batch size per device is 1 but each sequence may have different lengths. Some are very long and some are very short. I found it is very easy to confront the above problem (in contrast, if I pad all sequences to the same length, it becomes rare to confront the problem).

Could I know the reason? And is there any way to address this problem?

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.

> Hi everyone, the issue is not directly linked to `IterableDataset` but more on the `dispatch_batches` arg that is set to `True` when passing a `IterableDataset`.
> 
> With `dispatch_batches=True`, we process the data on the main process and broadcast to the other processes. This is more reliable + less compute since we do it on only one process. However, we indeed have an issue when the data doesn't have the same size as we are trying to concat tensors that might not have the same size.
> 
> As a temporary solution, you can either:
> 
> * pass `dispatch_batches=False` in `TrainingArguments`. This is the default behavior of Trainer before accelerate integration. It will use `IterableDatasetShard` instead of `DataLoaderDispatcher`.
> * pass `split_batches=True`. It will split a full batch into `self.num_process` parts. So, it requires that the batch size ( `per_device_train_batch_size` ) of the `dataloader` to be a round multiple of the number of processes. For example, if you set `per_device_train_batch_size = 16` and you have 4 processes, each process will have a batch_size of 4.
> 
> Related code:
> 
> ```python
> if self.split_batches:
>     # One batch of the main iterator is dispatched and split.
>     batch = next(iterator)
> else:
>     # num_processes batches of the main iterator are concatenated then dispatched and split.
>     # We add the batches one by one so we have the remainder available when drop_last=False.
>     batches = []
>     for _ in range(self.state.num_processes):
>         batches.append(next(iterator))
>    # The issue is here since the batches do not necessarily have the same size. 
>     batch = concatenate(batches, dim=0)
> ```

Thank @SunMarc for your solutions. 
However, I did't know how the temporary solutions affect the training efficiency and performance?

Hi @ycszen, I should have used the word `alternative` instead of `temporary`. The feature is not implemented yet for batches of different sizes when using `dispatch_batches=True`. 

Hi,
Is there any upcoming release which will include a fix for this issue?
Thanks.

Yes, later this week! 

Hi,
I'm running on dual GPU with IterableDatset, but still face the issue.
Note: on single GPU it works. 

@ArthurZucker is there any upcoming release which will include the fix?
Thanks


@hilevon can you try building from `main`? (`pip install git+https://github.com/huggingface/transformers`)

@muellerzr i did, but the issue still exists. In order to get this work I used "dispatch_batches=True" as a workaround. 

Hey there, just a tips, use `torchrun --nproc-per-node NUM_OF_GPUS train.py` instead of `python train.py`. Or similar error will be reported.

@Rbrq03 Yes, in order to run train.py on multi GPU (in my case double) torchrun --nproc-per-node NUM_OF_GPUS train.py can be used. And the trainer runs without any error in non-stream mode. However if you run it on stream mode (when IterableDataset is used underneath) the error appears. As I mentioned in my early comments setting "dispatch_batches=True" is a workaround.

Hi,
Now that this issue is closed, has it been resolved in the latest release?
Thanks.

Yes! 🫡 

I am facing the same error when running trainer on single GPU when calling trainer.evaluate()
Is this bug solved? 

make sure you are using `transformers>=4.47`