add fix for remove columns collator which fails with streaming

Signed-off-by: Dushyant Behl <dushyantbehl@in.ibm.com>

Author: dushyantbehl

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

@@ -20,6 +20,7 @@
 from fms_acceleration import AccelerationPlugin
 from peft import LoraConfig
 from transformers import DataCollatorForSeq2Seq, TrainingArguments
+from transformers.trainer_utils import RemoveColumnsCollator
 from trl import DataCollatorForCompletionOnlyLM  # pylint: disable=import-error
 import torch
 
@@ -70,8 +71,14 @@ def _collator_check(collate_fn):
             # "The padding-free plugin currently only works with a
             # `DataCollatorForSeq2Seq` collate_fn,
             # otherwise the collation can be unreliable"
+            if isinstance(collate_fn, RemoveColumnsCollator):
+                collate_fn = collate_fn.data_collator
             return isinstance(
-                collate_fn, (DataCollatorForSeq2Seq, DataCollatorForCompletionOnlyLM)
+                collate_fn,
+                (
+                    DataCollatorForSeq2Seq,
+                    DataCollatorForCompletionOnlyLM,
+                ),
             )
 
         # This check is done here to only patch the attention forward
@@ -97,6 +104,14 @@ def _collator_check(collate_fn):
 
         def _collator_replacement_builder(collate_fn):
 
+            # in case of remove columns collator the actual collate
+            # function is wrapped inside
+            if isinstance(collate_fn, RemoveColumnsCollator):
+                actual_collate_fn = collate_fn.data_collator
+                replacement = _collator_replacement_builder(actual_collate_fn)
+                collate_fn.data_collator = replacement
+                return collate_fn
+
             # in this case, replace seq2seq with flattening collator
             if isinstance(collate_fn, DataCollatorForSeq2Seq):
                 return DataCollatorWithFlattening()