Support for Iterable Datasets

[kmehant]

Transformers Lib while getting the data loader (for single gpu training) adds a wrapper collator over the original collator if the dataset is not of class HF Dataset, essentially for the cases where the dataset is iterable. Due to this reason, the data collator class turns into RemoveColumnsCollator which internally would internally removes the columns from the batch and then calls the original collator see - https://github.com/huggingface/transformers/blob/4d5b45870411053c9c72d24a8e1052e00fe62ad6/src/transformers/trainer_utils.py#L843

Using padding free with this scenario is not supported, since there is a hard check for the collators to be one of

fms-acceleration/plugins/attention-and-distributed-packing/src/fms_acceleration_aadp/framework_plugin_padding_free.py

Line 74 in 3cf092a

collate_fn, (DataCollatorForSeq2Seq, DataCollatorForCompletionOnlyLM)

We should be able to support this case.

I am happy to raise a PR. Thanks.

[fabianlim]

we should not support arbitrary RemoveColumnsCollator, and as I mentioned here, we check for the particular two collators is mostly an artifcat of how fms-hf-tuning is implemented.

If you want to support iterable, then we shoud follow the above logic, and figure out what was fms-hf-tuning's intention (i.e., pretok or not),.

If RemoveColumnsCollator is wrapping anyuthing else, we should not arbirarily support. The trickiness to supporting arbiratry collators is that the logic inside could be anything, and it could get out of hand supporting all edge cases.

Also I didnt envision to support generic iteratble datasets in aadp, is because multipack has to run thorugh the full dataset every epoch. So its not really intended for very large datasets.

[fabianlim]

will close this issue first, can re-open if there is future need