Double data splitting in get_train_dataloader for multi-gpu setting?
Author: tnnandi
Pull request: no
Labels: name: trainer

name: bug

name: solved

### System Info


- `transformers` version: 4.45.2
- Platform: Linux-5.14.21-150500.55.49-default-x86_64-with-glibc2.31
- Python version: 3.10.15
- Huggingface_hub version: 0.26.1
- Safetensors version: 0.4.5
- Accelerate version: 1.0.1
- Accelerate config:    - compute_environment: LOCAL_MACHINE
        - distributed_type: DEEPSPEED
        - use_cpu: False
        - debug: False
        - num_processes: 8
        - machine_rank: 0
        - num_machines: 2
        - main_process_ip: 123.12.34.45
        - main_process_port: 2345
        - rdzv_backend: static
        - same_network: True
        - main_training_function: main
        - enable_cpu_affinity: False
        - deepspeed_config: {'deepspeed_config_file': '/grand/GeomicVar/tarak/Geneformer_oct_2024/Geneformer/run/ds_config_zero.json', 'deepspeed_multinode_launcher': 'standard', 'zero3_init_flag': False}
        - downcast_bf16: no
        - tpu_use_cluster: False
        - tpu_use_sudo: False
        - tpu_env: []
- PyTorch version (GPU?): 2.5.0+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?: <fill in>
- Using GPU in script?: <fill in>
- GPU type: NVIDIA A100-SXM4-40GB


### Who can help?

@muellerz @SunMarc 

In get_train_dataloader(), for a sample size of N with M GPUs, len(DataLoader(train_dataset, **dataloader_params)) is N/M (As expected), but len(self.accelerator.prepare(DataLoader(train_dataset, **dataloader_params))) is N/(MxM).  Since a sampler is being used, the dataset should be split into the M GPUs appropriately, and I'm not being able to understand if accelerate.prepare() should be carrying out another splitting step.  When printing the batches (using a batch size of 1 per GPU), I can only view N/(M*M) samples per batch, and the rest of the samples [N/M - N/(MxM)] are not part of any batch. 

```
        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,
            "persistent_workers": self.args.dataloader_persistent_workers,
        }

        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
            dataloader_params["prefetch_factor"] = self.args.dataloader_prefetch_factor

        return self.accelerator.prepare(DataLoader(train_dataset, **dataloader_params))
```

### Information

- [X] The official example scripts
- [ ] My own modified scripts

### Tasks

- [X] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

```
        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,
            "persistent_workers": self.args.dataloader_persistent_workers,
        }

        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
            dataloader_params["prefetch_factor"] = self.args.dataloader_prefetch_factor
       
       ###########  print statements to confirm dataloader sizes ##############
       print(len(DataLoader(train_dataset, **dataloader_params))
       print(self.accelerator.prepare(DataLoader(train_dataset, **dataloader_params))) 

        return self.accelerator.prepare(DataLoader(train_dataset, **dataloader_params))
```

### Expected behavior

If the N samples are being distributed appropriately to M GPUs, the above function should return a dataloader with N/M samples. Instead it returns a dataloader of length N/(M*M) 

Comments:
Hey @tnnandi, thanks for the report. Could you share a minimal reproducer with logs that shows this behavior. This will be very helpful when debugging ! cc @muellerzr 

Thanks! Will do that shortly.

Hi @SunMarc @muellerzr  **Command to run**: accelerate launch --config_file ./accelerate_config_singlenode.without_ds.yaml test_samples_per_batch.py

In the script attached, I have tried to replicate the issue by creating a custom trainer class that closely follows the HF Trainer class in context of the issue I raised (i.e., includes the relevant functions to reproduce the issue).

A dataset with 40 samples is created in the script, that when using 4 GPUs should be split into 10 samples on each GPU. While that is true when get_train_dataloader() in the CustomPretrainer class returns _DataLoader(self.train_dataset, **dataloader_params)_ , only 3 samples are processed per GPU when the same function returns _self.accelerator.prepare(DataLoader(self.train_dataset, **dataloader_params))_[you can confirm from len(train_dataloader) or the prints from the training loop]

Code for _test_samples_per_batch.py_ is attached (since code pasting isn't working properly)
[test_samples_per_batch.zip](https://github.com/user-attachments/files/17529578/test_samples_per_batch.zip)


Content of accelerate config file _accelerate_config_singlenode.without_ds.yaml_:
```
compute_environment: LOCAL_MACHINE
debug: false
distributed_type: MULTI_GPU
downcast_bf16: 'no'
enable_cpu_affinity: false
machine_rank: 0
main_training_function: main
num_machines: 1
num_processes: 4
rdzv_backend: static
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: false
```



Hi @SunMarc @muellerzr , could you please tell if you had an opportunity to take a look at this issue?

Thanks, I have some time now so let me look into this over the next few hours and get back with you

actually I see right away. You're using the `DistributedSampler`. You should **not** be doing this because Accelerate's dataloaders do this (with its own implementation). Let accelerate handle all of it that you pass to `prepare`, or don't. It's one or the other. 

@muellerzr Thank you very much for pointing it out! I really appreciate the work you and the HF team is doing.

Closing this for now, feel free to reopen it if you have any other problems!