huggingface / huggingface/accelerate

[Feature Request] support to pass custom slice function in accelerate data loader (useful for models like Flamingo)

Open
#1,816 0 comments 0 reactions 1 assignee Claimed by @muellerzr View on GitHub
enhancement feature request
Dominant language
Python
Stars
9.9k
Forks
1.5k
Avg merge
5d 2h
Merged PRs (30d)
27

Description

@muellerzr & team, first of all, thanks for this great library!!

It would be great to be able to pass a custom slice function (instead of the using the default one):
https://github.com/huggingface/accelerate/blob/30eff605df5db9ba20289a0c7471e4d8547a5a8f/src/accelerate/data_loader.py#L596

use case: it would make it easy to train models like [flamingo](https://github.com/mlfoundations/open_flamingo) where text & images are interleaved. When text & image are interleaved, the data collator returns 2 tensors: `text` of size `(batch_size, seqlen)` & `images` of size `(num_images, H, W, C)`. Now we can't slice the image tensor using the original function in the library (because it slices based on the batch size). rather we want to slice the image tensor after finding the number of images required on each device. For eg: a batch of device-0 can get 2 images and a batch of device-1 can get 3 images.

Another workaround could also be to have slicing optional and let users control it after the same tensors are returned on all the devices.

Do you think accelerate library can support these kinda scenarios as well? I believe it will be super useful to train models like flamingo.

Let me know if my issue wasn't clear. Happy to discuss more. Thanks!

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.