From 8517d97abbb7cc3e8a5260c203f1a34402643f8b Mon Sep 17 00:00:00 2001 From: Vibhu Jawa Date: Fri, 14 Mar 2025 16:54:01 -0700 Subject: [PATCH] Fix Transformers bug in NeMo Curator (#591) Signed-off-by: Vibhu Jawa --- nemo_curator/classifiers/aegis.py | 6 +++++- pyproject.toml | 2 +- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/nemo_curator/classifiers/aegis.py b/nemo_curator/classifiers/aegis.py index 2951959a0..6efd1de97 100644 --- a/nemo_curator/classifiers/aegis.py +++ b/nemo_curator/classifiers/aegis.py @@ -24,7 +24,6 @@ from crossfit import op from crossfit.backend.torch.hf.model import HFModel from huggingface_hub import PyTorchModelHubMixin -from peft import PeftModel from torch.nn import Dropout, Linear from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer @@ -116,6 +115,11 @@ def __init__( base_model = AutoModelForCausalLM.from_pretrained( pretrained_model_name_or_path, torch_dtype=dtype, token=token ) + # Importing PeftModel here to prevent cuda context issues + # that seem to happen on Transformers 4.48.3 + # See related: https://github.com/rapidsai/crossfit/pull/113 + from peft import PeftModel + self.model = PeftModel.from_pretrained(base_model, peft_model_name_or_path) self.autocast = autocast self.add_instruction_data_guard = add_instruction_data_guard diff --git a/pyproject.toml b/pyproject.toml index 66a2d9f87..0b8686b13 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -43,7 +43,7 @@ dependencies = [ "beautifulsoup4", "charset_normalizer>=3.1.0", "comment_parser", - "crossfit>=0.0.8", + "crossfit>=0.0.8.post1", "dask-mpi>=2021.11.0", "dask[complete]>=2021.7.1", "datasets",