"""Fine-tune a small instruction model with TRL's SFTTrainer and a PEFT LoRA adapter. Purpose: the course's reference supervised fine-tuning run. Loads a prompt-completion dataset, attaches a LoRA adapter, trains for a few epochs with an evaluation after each one, keeps the best checkpoint, and appends a run record to the lab notebook. Platform: spark, strix, nvidia (CUDA or ROCm); it also runs on the CPU, slowly. Track M uses train-sft-mlx.sh instead; PyTorch's MPS backend will run this script in float32 but the MLX path is the supported one on a Mac. Minimum memory: 8 GB Assumes: torch, transformers, trl, peft and datasets installed in the active environment; make-dataset.py has been run so that data/train.jsonl and data/valid.jsonl exist; runlog.py sits next to this file. Usage: python train-sft.py --data-dir data --output-dir runs/sft-qwen3-0.6b --labbook labbook.md python train-sft.py --list-modules # print the model's linear layer names and exit """ from __future__ import annotations import argparse import json import time from pathlib import Path import torch from datasets import load_dataset from peft import LoraConfig from transformers import AutoModelForCausalLM, AutoTokenizer from trl import SFTConfig, SFTTrainer import runlog DEFAULT_TARGETS = ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"] def pick_device() -> str: """Same choice as Part 1's script: CUDA (or ROCm, which reports as cuda), then MPS, then CPU.""" if torch.cuda.is_available(): return "cuda" mps = getattr(torch.backends, "mps", None) if mps is not None and mps.is_available(): return "mps" return "cpu" def use_bf16(device: str, requested: str) -> bool: """BF16 only where the device supports it; everything else trains in float32.""" if requested == "fp32": return False if requested == "bf16": return True return device == "cuda" and torch.cuda.is_bf16_supported() def list_linear_modules(model_id: str) -> None: """Print the names LoRA can target, so target_modules is never guessed.""" model = AutoModelForCausalLM.from_pretrained(model_id, dtype=torch.float32) names = sorted({name.split(".")[-1] for name, mod in model.named_modules() if isinstance(mod, torch.nn.Linear)}) print(f"linear module names in {model_id}:") for name in names: print(f" {name}") def summarise_history(history: list[dict]) -> dict[str, float | int | None]: """The three numbers worth keeping out of a log full of them.""" train_losses = [row["loss"] for row in history if "loss" in row] evals = [(row["epoch"], row["eval_loss"]) for row in history if "eval_loss" in row] best_epoch, best_eval = min(evals, key=lambda pair: pair[1]) if evals else (None, None) return { "final_train_loss": round(train_losses[-1], 4) if train_losses else None, "first_train_loss": round(train_losses[0], 4) if train_losses else None, "final_eval_loss": round(evals[-1][1], 4) if evals else None, "best_eval_loss": round(best_eval, 4) if best_eval is not None else None, "best_epoch": best_epoch, } def main() -> None: parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) parser.add_argument("--model", default="Qwen/Qwen3-0.6B", help="base model repository id or local path") parser.add_argument("--data-dir", default="data", help="directory holding train.jsonl and valid.jsonl") parser.add_argument("--output-dir", default="runs/sft-qwen3-0.6b") parser.add_argument("--epochs", type=float, default=3.0) parser.add_argument("--batch-size", type=int, default=2, help="per-device batch size") parser.add_argument("--grad-accum", type=int, default=4, help="batches summed before an optimiser step") parser.add_argument("--lr", type=float, default=1e-4, help="adapters take a higher rate than a full fine-tune") parser.add_argument("--max-length", type=int, default=512) parser.add_argument("--rank", type=int, default=16) parser.add_argument("--alpha", type=int, default=32) parser.add_argument("--dropout", type=float, default=0.05) parser.add_argument("--target-modules", nargs="+", default=DEFAULT_TARGETS) parser.add_argument("--precision", choices=["auto", "bf16", "fp32"], default="auto") parser.add_argument("--gradient-checkpointing", action="store_true", help="recompute activations in the backward pass; saves memory, costs time") parser.add_argument("--packing", action="store_true", help="pack several examples into one sequence") parser.add_argument("--report-to", default="none", choices=["none", "tensorboard"]) parser.add_argument("--seed", type=int, default=0) parser.add_argument("--labbook", default=None, help="append one JSON run record to this file") parser.add_argument("--list-modules", action="store_true", help="print linear module names and exit") args = parser.parse_args() if args.list_modules: list_linear_modules(args.model) return device = pick_device() bf16 = use_bf16(device, args.precision) dtype = torch.bfloat16 if bf16 else torch.float32 print(f"device: {device} precision: {'bfloat16' if bf16 else 'float32'}") data_dir = Path(args.data_dir) files = {"train": str(data_dir / "train.jsonl"), "validation": str(data_dir / "valid.jsonl")} for split, path in files.items(): if not Path(path).is_file(): raise SystemExit(f"{path} is missing; run make-dataset.py --out-dir . first ({split} split)") dataset = load_dataset("json", data_files=files) print(f"train examples: {len(dataset['train'])} validation examples: {len(dataset['validation'])}") tokenizer = AutoTokenizer.from_pretrained(args.model) if tokenizer.chat_template is None: raise SystemExit( f"{args.model} has no chat template; pick an instruct model or set chat_template_path in SFTConfig" ) config = SFTConfig( output_dir=args.output_dir, num_train_epochs=args.epochs, per_device_train_batch_size=args.batch_size, per_device_eval_batch_size=args.batch_size, gradient_accumulation_steps=args.grad_accum, learning_rate=args.lr, lr_scheduler_type="cosine", warmup_steps=5, max_length=args.max_length, packing=args.packing, completion_only_loss=True, # loss on the answer tokens, not on the question gradient_checkpointing=args.gradient_checkpointing, bf16=bf16, model_init_kwargs={"dtype": dtype}, eval_strategy="epoch", save_strategy="epoch", save_total_limit=2, load_best_model_at_end=True, metric_for_best_model="eval_loss", greater_is_better=False, logging_steps=5, report_to=args.report_to, seed=args.seed, data_seed=args.seed, ) peft_config = LoraConfig( r=args.rank, lora_alpha=args.alpha, lora_dropout=args.dropout, target_modules=args.target_modules, bias="none", task_type="CAUSAL_LM", ) trainer = SFTTrainer( model=args.model, args=config, train_dataset=dataset["train"], eval_dataset=dataset["validation"], processing_class=tokenizer, peft_config=peft_config, ) trainer.model.print_trainable_parameters() started = time.time() trainer.train() elapsed = time.time() - started trainer.save_model(args.output_dir) tokenizer.save_pretrained(args.output_dir) losses = summarise_history(trainer.state.log_history) losses["seconds"] = round(elapsed, 1) print(json.dumps(losses, indent=2)) print(f"adapter saved to {args.output_dir}") if args.labbook: dataset_path = files["train"] record = runlog.record( labbook=args.labbook, lab="part-11/train-sft", model=args.model, dataset={ "path": dataset_path, "sha256": runlog.file_sha256(dataset_path), "train_examples": len(dataset["train"]), "validation_examples": len(dataset["validation"]), }, hyperparameters={ "method": "lora", "rank": args.rank, "alpha": args.alpha, "dropout": args.dropout, "target_modules": args.target_modules, "epochs": args.epochs, "batch_size": args.batch_size, "grad_accum": args.grad_accum, "effective_batch": args.batch_size * args.grad_accum, "learning_rate": args.lr, "max_length": args.max_length, "packing": args.packing, "gradient_checkpointing": args.gradient_checkpointing, "precision": "bfloat16" if bf16 else "float32", "completion_only_loss": True, }, seed=args.seed, losses=losses, scores={}, config_path=__file__, notes=None, ) print(f"recorded run {record['run_id']} in {args.labbook}") if __name__ == "__main__": main()