Skip to content

fix: handle trl compatibility for DataCollatorForCompletionOnlyLM #169

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
YashasviChaurasia wants to merge 1 commit into from
+15 −3
Closed

fix: handle trl compatibility for DataCollatorForCompletionOnlyLM #169

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 15 additions & 3 deletions ...ention-and-distributed-packing/src/fms_acceleration_aadp/framework_plugin_padding_free.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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):

Expand Down
Loading