fix: Using PaddingFree plugin with Iterable datasets and remove_columns

[dushyantbehl]

While doing our testing of using Iterable Datasets in few training runs we hit this error where the collator being used at training time was RemoveColumnsCollator. It seems that at times trl internally wraps the collator we pass either DataCollatorForSeq2Seq or DataCollatorForCompletionOnlyLM inside this RemoveColumnsCollator to remove columns at batch creation time.

The simple fix for making this collator work with the padding free plugin is to unwrap the actual collator and try to build a replacement for the same.

[dushyantbehl]

@fabianlim opening up the PR for simple review for now..I will be doing more testing and sharing the error details we noticed before and after this patch.
Please advice if you want me to add some unit tests for this change.

[kmehant]

related - huggingface/transformers#34830

@dushyantbehl Just in case if you wish to fix this from transformers end there is some interest from HF team for a PR. If you wish to work on this, you can take this up as you see it fit.

[dushyantbehl]

Thanks @kmehant surely I'll check it out and fix

[kmehant]

related - #87

[kmehant]

From @fabianlim's comment we should ideally disable or gracefully fail when multipack is requested when allowing for streaming + padding free. May be this is something to be added to fms-hf-tuning.

@dushyantbehl can you check other concerns from Fabian from this issue and see if we have covered - #87 (comment)?

[dushyantbehl]

If I read @fabianlim 's comment on #87 (comment)

It basically comes down to not supporting arbitrary collators inside RemoveColumnsCollator but only support DataCollatorForSeq2Seq or DataCollatorForCompletionOnlyLM and that is exactly what this PR is doing...it checks inside the remove columns collator and performs all checks to see if it is indeed one of the two supported or not.
So I am assuming the PR is to his spirit and not breaking the ideology. Let me know if you disagree here.

On the use of multipack + streaming we can just have a check in the fms-hf-tuning repo and gracefully fail if multipack is requested with streaming.

[kmehant]

@fabianlim FYA ^ thank you.

[kmehant]

@dushyantbehl can you rebase and fix the conflicts?

[dushyantbehl]

@kmehant I would let you or @romitjain pick this up and we can close this PR.

[dushyantbehl]

@kmehant the change was minimal so I just rebased the PR lets go ahead with this one itself so its less work for you both