使用PEFT与transformers从checkpoint恢复训练时修改超参数

核心思路 #

避开trainer.train(resume_from_checkpoint)的自动恢复逻辑，改为手动加载PEFT权重、训练状态，同时自定义新的训练参数，彻底避免checkpoint配置覆盖自定义超参数的问题。

具体操作步骤 #

1. 初始化PEFT模型（与原训练结构一致） #

按正常流程初始化量化模型与PEFT配置，确保模型结构和原训练时完全匹配：

from peft import prepare_model_for_kbit_training, LoraConfig, get_peft_model
import transformers
from datetime import datetime

# 初始化量化模型
ft_model = prepare_model_for_kbit_training(model)
ft_model.gradient_checkpointing_enable()

# 配置PEFT（如需调整Lora结构，只能重新训练，无法从原有checkpoint恢复）
config = LoraConfig(
    r=8,
    lora_alpha=16,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "w1", "w2", "w3", "lm_head"],
    bias="none",
    lora_dropout=0.05,
    task_type="CAUSAL_LM",
)
ft_model = get_peft_model(ft_model, config)
ft_model = accelerator.prepare_model(ft_model)

# 配置tokenizer
tokenizer.pad_token = tokenizer.eos_token

2. 手动加载Checkpoint中的权重与训练状态 #

从指定checkpoint目录加载PEFT模型权重、优化器/调度器状态，以及当前训练进度：

import torch

checkpoint_dir = "/content/latest_checkpoint/"

# 加载PEFT模型权重
ft_model = ft_model.from_pretrained(ft_model, checkpoint_dir)

# 读取训练状态（当前step、优化器、调度器参数）
train_state = torch.load(f"{checkpoint_dir}/trainer_state.json")
optimizer_state = torch.load(f"{checkpoint_dir}/optimizer.pt")
scheduler_state = torch.load(f"{checkpoint_dir}/scheduler.pt")

current_step = train_state["global_step"]

3. 自定义新的训练超参数 #

自由修改需要调整的训练参数，比如学习率、总步数、warmup步数等，不受checkpoint限制：

new_training_args = transformers.TrainingArguments(
    output_dir=output_dir,
    warmup_steps=200,  # 修改warmup步数
    per_device_train_batch_size=8,
    gradient_checkpointing=True,
    gradient_accumulation_steps=2,
    max_steps=6000,  # 修改总训练步数
    learning_rate=3e-5,  # 修改学习率
    logging_steps=10,
    fp16=True,
    optim="paged_adamw_8bit",
    logging_dir="/content/logs",
    save_strategy="steps",
    save_steps=10,
    evaluation_strategy="steps",
    eval_steps=10,
    load_best_model_at_end=True,
    report_to="wandb",
    run_name=f"{run_name}-{datetime.now().strftime('%Y-%m-%d-%H-%M')}"
)

4. 初始化Trainer并恢复训练 #

将加载好的模型、新参数、训练状态传入Trainer，从指定step开始训练：

trainer = transformers.Trainer(
    model=ft_model,
    callbacks=[upload_checkpoint_callback],
    train_dataset=tokenized_train_dataset,
    eval_dataset=tokenized_val_dataset,
    args=new_training_args,
    data_collator=transformers.DataCollatorForLanguageModeling(tokenizer, mlm=False),
)

# 加载优化器与调度器状态
trainer.optimizer.load_state_dict(optimizer_state)
trainer.lr_scheduler.load_state_dict(scheduler_state)

# 从当前step启动训练，禁用自动恢复逻辑
model.config.use_cache = False
trainer.train(resume_from_checkpoint=False, initial_global_step=current_step)

关键注意事项 #

• 模型结构必须匹配：初始化的PEFT模型结构（如Lora的r值、target_modules）要和checkpoint中的完全一致，否则权重加载会失败。若需调整模型结构，只能重新训练。
• 可修改范围：仅支持修改训练流程类参数（学习率、总步数、warmup步数等），模型结构类参数无法在恢复训练时更改。
• 禁用自动恢复：必须设置resume_from_checkpoint=False，同时指定initial_global_step，才能确保使用自定义的训练参数。

内容的提问来源于stack exchange，提问作者Jon Flynn