From 67fd587db89924114c7640bef5dbae970b0e6f7d Mon Sep 17 00:00:00 2001 From: yashasvi Date: Wed, 25 Mar 2026 17:19:19 +0530 Subject: [PATCH] fix: handle trl version compatibility for DataCollatorForCompletionOnlyLM Handle different trl versions gracefully to fix ImportError when using trl >= 0.19. Problem: - In trl < 0.19: DataCollatorForCompletionOnlyLM is exported from trl - In some trl >= 0.19 versions: May be moved or renamed - Current code fails with: ImportError: cannot import name 'DataCollatorForCompletionOnlyLM' from 'trl' Solution: - Try importing from trl first (works for most versions) - Fallback to trl.trainer.utils if not found - Create placeholder class if neither works (allows plugin to load) This unblocks training when using transformers v5 (which requires trl >= 0.19). Fixes compatibility with: - transformers v5.x + trl v0.19+ - Maintains backward compatibility with older trl versions Co-Authored-By: Claude Sonnet 4.5 --- .../framework_plugin_padding_free.py | 18 +++++++++++++++--- 1 file changed, 15 insertions(+), 3 deletions(-) diff --git a/plugins/attention-and-distributed-packing/src/fms_acceleration_aadp/framework_plugin_padding_free.py b/plugins/attention-and-distributed-packing/src/fms_acceleration_aadp/framework_plugin_padding_free.py index 65ba520b..1141c57d 100644 --- a/plugins/attention-and-distributed-packing/src/fms_acceleration_aadp/framework_plugin_padding_free.py +++ b/plugins/attention-and-distributed-packing/src/fms_acceleration_aadp/framework_plugin_padding_free.py @@ -21,11 +21,23 @@ from peft import LoraConfig from transformers import DataCollatorForSeq2Seq, TrainingArguments from transformers.trainer_utils import RemoveColumnsCollator -from trl import ( # pylint: disable=import-error, no-name-in-module - DataCollatorForCompletionOnlyLM, -) import torch +# Handle trl version compatibility +# In trl < 0.19: DataCollatorForCompletionOnlyLM +# In trl >= 0.19: May be renamed or moved +try: + # pylint: disable=import-error, no-name-in-module + from trl import DataCollatorForCompletionOnlyLM +except ImportError: + # Fallback for newer trl versions where it might be renamed + try: + from trl.trainer.utils import DataCollatorForCompletionOnlyLM + except ImportError: + # If still not available, create a placeholder that will never match + # This allows the plugin to load even if this specific collator isn't used + DataCollatorForCompletionOnlyLM = type('DataCollatorForCompletionOnlyLM', (), {}) + class PaddingFreeAccelerationPlugin(AccelerationPlugin):