[Python] Trainer is failing to load optimizer save state when resuming training

Intro to the problem


I am trying to train Llama-3.1 8B on an H100 but I keep running into the following error when trying to resume training

...
File "/home/jovyan/folder/training/.venv/lib/python3.10/site-packages/transformers/trainer.py", line 2405, in _inner_training_loop
self._load_optimizer_and_scheduler(resume_from_checkpoint)
File "/home/jovyan/folder/training/.venv/lib/python3.10/site-packages/transformers/trainer.py", line 3452, in _load_optimizer_and_scheduler
self.optimizer.load_state_dict(
File "/home/jovyan/folder/training/.venv/lib/python3.10/site-packages/accelerate/optimizer.py", line 107, in load_state_dict
self.optimizer.load_state_dict(state_dict)
File "/home/jovyan/folder/training/.venv/lib/python3.10/site-packages/torch/_compile.py", line 51, in inner
return disable_fn(*args, **kwargs)
File "/home/jovyan/folder/training/.venv/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 838, in _fn
return fn(*args, **kwargs)
File "/home/jovyan/folder/training/.venv/lib/python3.10/site-packages/torch/optim/optimizer.py", line 872, in load_state_dict
raise ValueError(
ValueError: loaded state dict contains a parameter group that doesn't match the size of optimizer's group

What I am trying to do


I have two scripts, train.py and resume.py. train.py keeps crashing every now and then, likely due to issues with the hardware so I am trying to resume training from a checkpoint every time it crashes. I am also quantizing the model and using Lora to add trainable parameters.

# train.py
from data import LanguageDataset
from torch.utils.data import DataLoader
from transformers import Trainer, TrainingArguments, DataCollatorForLanguageModeling, AutoTokenizer, BitsAndBytesConfig, AutoModelForCausalLM
from tokenization import LanguageTokenizer
from torch.optim import AdamW
import torch
from peft import LoraConfig, TaskType, PeftModel, PeftConfig, get_peft_model
from safetensors.torch import save_model


def train_loop():

dataset = LanguageDataset()

quantization_configs = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_compute_dtype=torch.bfloat16,
bnb_4bit_quant_type="nf4",
bnb_4bit_use_double_quant=True,
)

model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.1-8B"
, device_map="auto", quantization_config=quantization_configs)
t = LanguageTokenizer().tokenizer
model.resize_token_embeddings(len(t), mean_resizing=True)

training_args = TrainingArguments(
output_dir="./training_output",
per_device_train_batch_size=8,
num_train_epochs=3,
save_steps=10,
save_strategy="steps",
save_safetensors=False,
logging_steps=10,
)

lora_config = LoraConfig(
task_type=TaskType.CAUSAL_LM,
inference_mode=False,
r=8,
lora_alpha=32,
lora_dropout=0.1
)

peft_model = get_peft_model(model=model, peft_config=lora_config)


data_collator = DataCollatorForLanguageModeling(tokenizer=t, mlm=False)
trainer = Trainer(
model=peft_model,
processing_class=t,
args=training_args,
train_dataset=dataset,
data_collator=data_collator
)


trainer.train()


if __name__ == "__main__":
train_loop()


# resume.py
from data import LanguageDataset
from torch.utils.data import DataLoader
from transformers import Trainer, TrainingArguments, DataCollatorForLanguageModeling, AutoTokenizer, BitsAndBytesConfig, AutoModelForCausalLM
from config import epochs, model_name, batch_size
from tokenization import LanguageTokenizer
from torch.optim import AdamW
import torch
from peft import LoraConfig, TaskType, PeftModel, PeftConfig
from safetensors.torch import save_model

def resume_loop():
dataset = LanguageDataset()

quantization_configs = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_compute_dtype=torch.bfloat16,
bnb_4bit_quant_type="nf4",
bnb_4bit_use_double_quant=True,
)

model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.1-8B", device_map="auto", quantization_config=quantization_configs)
t = LanguageTokenizer().tokenizer
model.resize_token_embeddings(len(t), mean_resizing=True)

peft_model = PeftModel.from_pretrained(model, "./training_output/checkpoint-20")

training_args = TrainingArguments(
output_dir="./training_output",
per_device_train_batch_size=8,
num_train_epochs=3,
save_steps=10,
save_strategy="steps",
save_safetensors=False,
logging_steps=10,
)

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

t = AutoTokenizer.from_pretrained("./training_output/checkpoint-20")
trainer = Trainer(
model=peft_model,
processing_class=t,
args=training_args,
train_dataset=dataset,
data_collator=data_collator,
)


trainer.train(resume_from_checkpoint="./training_output/checkpoint-20")


if __name__ == "__main__":
resume_loop()


The tokenizer I am defining is Llama's pretrained tokenizer in the AutoTokenizer class. It is tokenizing with a target_text so the outputs include input_ids, attention_mask and labels.

The issue seems to be with loading the optimizer save state. In the checkpoint folder, optimizer.pt and scheduler.pt files do exist. But I am quite a bit confused because I am letting Trainer pick the optimizer and scheduler. I am not quite sure what I am doing wrong here.

I've tried to make my code resemble as closely to other training code I found on the internet. I have also tried disabling mean_resizing to no avail.

Continue reading...