File size: 5,247 Bytes
7035ef3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 | """QLoRA fine-tuning of Qwen2.5-7B-Instruct on the positional reasoning dataset.
Three hyperparameter combos (see Deliverable4/README.md training plan):
A: lr 1e-5, target q_proj/v_proj (conservative — HW6 baseline, expected weak)
B: lr 1e-4, target q_proj/v_proj (TRL-recommended adapter lr)
C: lr 2e-4, target all-linear (QLoRA-style full coverage)
Usage (on the GPU box, inside a venv with transformers/trl/peft/bitsandbytes):
python3 train_qlora.py A # or B, C
nohup python3 train_qlora.py B > combo_B.log 2>&1 & # survives ssh disconnects
Writes to ./runs/<combo>/: checkpoints, best model adapter, metrics.json
(train/eval loss history) for display in the check-in notebook.
"""
import json
import os
import sys
import torch
from datasets import Dataset
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
from trl import SFTConfig, SFTTrainer
MODEL_ID = "Qwen/Qwen2.5-7B-Instruct"
DATA = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "pipeline", "data", "accepted.jsonl")
SEED = 42
COMBOS = {
"A": {"lr": 1e-5, "target_modules": ["q_proj", "v_proj"]},
"B": {"lr": 1e-4, "target_modules": ["q_proj", "v_proj"]},
"C": {"lr": 2e-4, "target_modules": "all-linear"},
}
SYSTEM = ("You solve math word problems using positional reasoning: a chain of typed "
"nodes, each committing to a prediction before the next step is computed. "
"End with the final answer in <answer></answer> tags.")
def load_split(tokenizer):
rows = [json.loads(l) for l in open(DATA)]
ds = Dataset.from_list([
{
"prompt": tokenizer.apply_chat_template(
[{"role": "system", "content": SYSTEM},
{"role": "user", "content": r["question"]}],
tokenize=False, add_generation_prompt=True),
"completion": r["positional"] + tokenizer.eos_token,
}
for r in rows
])
split = ds.train_test_split(test_size=0.1, seed=SEED)
return split["train"], split["test"]
def main(combo_name):
combo = COMBOS[combo_name]
out_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "runs", combo_name)
os.makedirs(out_dir, exist_ok=True)
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
train_ds, eval_ds = load_split(tokenizer)
print(f"combo {combo_name}: {len(train_ds)} train / {len(eval_ds)} eval examples")
bnb = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.float16, # T4 has no bf16
bnb_4bit_use_double_quant=True,
)
model = AutoModelForCausalLM.from_pretrained(
MODEL_ID, quantization_config=bnb, device_map="auto", torch_dtype=torch.float16)
model = prepare_model_for_kbit_training(model)
model.config.use_cache = False
lora = LoraConfig(
r=64, lora_alpha=64, lora_dropout=0.05, bias="none",
task_type="CAUSAL_LM", target_modules=combo["target_modules"])
model = get_peft_model(model, lora)
model.print_trainable_parameters()
args = SFTConfig(
output_dir=out_dir,
num_train_epochs=2,
per_device_train_batch_size=1,
gradient_accumulation_steps=8,
learning_rate=combo["lr"],
logging_steps=10,
eval_strategy="steps",
eval_steps=25,
save_steps=25,
save_total_limit=2,
load_best_model_at_end=True,
fp16=True,
max_length=1600,
gradient_checkpointing=True,
optim="paged_adamw_8bit",
report_to=[],
seed=SEED,
)
trainer = SFTTrainer(
model=model, args=args,
train_dataset=train_ds, eval_dataset=eval_ds)
# TRL's SFTTrainer unconditionally casts LoRA adapter params to bfloat16
# for any quantized model during __init__ (QLoRA paper recommendation) --
# confirmed via trl/trainer/sft_trainer.py, no config option to opt out
# for quantized models. bfloat16 has no unscale kernel for fp16's
# GradScaler, and this T4 has no bf16 hardware support anyway. The correct
# fix (confirmed via a standalone diagnostic) is fp32, not fp16: fp16
# mixed-precision training expects fp32 "master weight" params with
# autocast handling the fp16 compute internally -- GradScaler explicitly
# rejects unscaling gradients that are already fp16 leaf params
# ("Attempting to unscale FP16 gradients"). Re-cast to fp32 after the
# trainer is built, undoing TRL's bf16 cast rather than fighting it
# beforehand.
for p in trainer.model.parameters():
if p.requires_grad:
p.data = p.data.to(torch.float32)
trainer.train()
trainer.model.save_pretrained(os.path.join(out_dir, "best_model"))
with open(os.path.join(out_dir, "metrics.json"), "w") as f:
json.dump(trainer.state.log_history, f, indent=1)
print(f"combo {combo_name} done -> {out_dir}/best_model")
if __name__ == "__main__":
name = sys.argv[1] if len(sys.argv) > 1 else None
if name not in COMBOS:
print(f"usage: python3 train_qlora.py [{'|'.join(COMBOS)}]")
sys.exit(1)
main(name)
|