Error Mapping on sd3, sdxl and upcoming flux controlnet training scripts in diffusers
Opened by Night1099 (closed)

### Describe the bug

```
Map:   6%|██████                                                                                                   | 8000/138120 [19:27<5:16:36,  6.85 examples/s]
Traceback (most recent call last):
  File "/workspace/diffusers/examples/controlnet/train_controlnet_sd3.py", line 1416, in <module>
    main(args)
  File "/workspace/diffusers/examples/controlnet/train_controlnet_sd3.py", line 1132, in main
    train_dataset = train_dataset.map(compute_embeddings_fn, batched=True, new_fingerprint=new_fingerprint)
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/dist-packages/datasets/arrow_dataset.py", line 560, in wrapper
    out: Union["Dataset", "DatasetDict"] = func(self, *args, **kwargs)
                                           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/dist-packages/datasets/arrow_dataset.py", line 3035, in map
    for rank, done, content in Dataset._map_single(**dataset_kwargs):
  File "/usr/local/lib/python3.11/dist-packages/datasets/arrow_dataset.py", line 3461, in _map_single
    writer.write_batch(batch)
  File "/usr/local/lib/python3.11/dist-packages/datasets/arrow_writer.py", line 567, in write_batch
    self.write_table(pa_table, writer_batch_size)
  File "/usr/local/lib/python3.11/dist-packages/datasets/arrow_writer.py", line 579, in write_table
    pa_table = pa_table.combine_chunks()
               ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "pyarrow/table.pxi", line 4387, in pyarrow.lib.Table.combine_chunks
  File "pyarrow/error.pxi", line 155, in pyarrow.lib.pyarrow_internal_check_status
  File "pyarrow/error.pxi", line 92, in pyarrow.lib.check_status
pyarrow.lib.ArrowInvalid: offset overflow while concatenating arrays
Traceback (most recent call last):
  File "/usr/local/bin/accelerate", line 8, in <module>
    sys.exit(main())
             ^^^^^^
  File "/usr/local/lib/python3.11/dist-packages/accelerate/commands/accelerate_cli.py", line 48, in main
    args.func(args)
  File "/usr/local/lib/python3.11/dist-packages/accelerate/commands/launch.py", line 1174, in launch_command
    simple_launcher(args)
  File "/usr/local/lib/python3.11/dist-packages/accelerate/commands/launch.py", line 769, in simple_launcher
  ```

### Steps to reproduce the bug

The dataset has no problem training on sd1.5 controlnet train script



### Expected behavior

Script not randomly erroing with error above

### Environment info

- `datasets` version: 3.0.0
- Platform: Linux-6.5.0-44-generic-x86_64-with-glibc2.35
- Python version: 3.11.9
- `huggingface_hub` version: 0.25.1
- PyArrow version: 17.0.0
- Pandas version: 2.2.3
- `fsspec` version: 2024.6.1


training on A100

Comments:
this is happening on large datasets, if anyone happens upon this i was able to fix by changing

```
train_dataset = train_dataset.map(compute_embeddings_fn, batched=True, new_fingerprint=new_fingerprint)
```

to

```
train_dataset = train_dataset.map(compute_embeddings_fn, batched=True, batch_size=16, new_fingerprint=new_fingerprint)
```