[Python] Adding EarlyStopping() to transformers Trainer() error

I'm using the code for fine-tuning LoRa model:

bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_compute_dtype=torch.float16,
bnb_4bit_use_double_quant=True,
bnb_4bit_quant_type="nf4",
)

model = AutoModelForCausalLM.from_pretrained(
"tiiuae/falcon-rw-1b",
quantization_config=bnb_config,
device_map={"": torch.cuda.current_device()},
)

tokenizer = AutoTokenizer.from_pretrained("zhihan1996/DNABERT-2-117M", trust_remote_code=True)

peft_config = LoraConfig(
task_type=TaskType.CAUSAL_LM,
inference_mode=False,
r=8,
lora_alpha=16,
lora_dropout=0.05,
bias="none"
)

model = get_peft_model(model, peft_config)

dataset = load_from_disk('tokenized_dataset_50_percent')


train_size = int(0.8 * len(dataset["train"]))
test_size = len(dataset[ "train"]) - train_size
train_set, val_set = torch.utils.data.random_split(dataset["train"], [train_size, test_size])

data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)

training_args = TrainingArguments(
output_dir="./falcon-dna-lora",
per_device_train_batch_size=4,
gradient_accumulation_steps=32,
num_train_epochs=1,
fp16=True,
save_total_limit=2,
logging_steps=10,
save_steps=500,
learning_rate=2e-4,
weight_decay=0.01,
report_to="none"
)

trainer = Trainer(
model=model,
args=training_args,
train_dataset=train_set,
eval_dataset=val_set,
data_collator=data_collator,
tokenizer=tokenizer,
)
print("Trainable params:", sum(p.numel() for p in model.parameters() if p.requires_grad))

trainable_params = []
total_params = 0
trainable_count = 0

for name, param in model.named_parameters():
total_params += param.numel()
if param.requires_grad:
trainable_count += param.numel()
trainable_params.append(name)

print(f"Total parameters: {total_params:,}")
print(f"Trainable parameters: {trainable_count:,}")
print(f"Percentage trainable: {100 * trainable_count / total_params:.4f}%")
print(f"Trainable layers: {trainable_params}")

trainer.train()
trainer.save_model("falcon-rw-1b-50percent-checkpoint")


Trainer() method works fine and the model trains correctly. Problem begins if I add EarlyStopping callback, by doing the following changes:

training_args = TrainingArguments(
output_dir="./falcon-dna-lora",
per_device_train_batch_size=4,
gradient_accumulation_steps=32,
num_train_epochs=1,
fp16=True,
save_total_limit=2,
logging_steps=10,
save_steps=500,
weight_decay=0.01,
report_to="none",
eval_strategy="steps",
eval_steps=500,
load_best_model_at_end=True,
metric_for_best_model="eval_loss",
greater_is_better=False,
learning_rate=2e-4,
)

trainer = Trainer(
model=model,
args=training_args,
train_dataset=train_set,
eval_dataset=val_set,
data_collator=data_collator,
tokenizer=tokenizer,
callbacks=[EarlyStoppingCallback(early_stopping_patience=3)]
)


After this I receive the following error:

RuntimeError: CUDA error: an illegal memory access was encountered
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1
Compile with TORCH_USE_CUDA_DSA to enable device-side assertions.


Have run out of ideas and none of those work. Can you suggest what could be the issue with this approach?

Thanks in advance.

Continue reading...