"""Train a LoRA adapter on a format-following dataset with TRL's SFTTrainer and PEFT. Purpose: Part 13's self-contained supervised fine-tuning run for Tracks S, X and N. Loads the conversational prompt-completion files that make-format-dataset.py wrote, attaches a LoRA adapter, trains in bfloat16 with an evaluation after every epoch, stops early when the evaluation loss stops improving, keeps the best-scoring checkpoint, and appends one run record to the lab notebook. Platform: spark, strix, nvidia (CUDA, or ROCm which also reports as cuda to PyTorch). It runs on the CPU too, slowly. Track M uses train-lora-mlx.sh instead: PyTorch's MPS backend will run this in float32, but mlx-lm is the supported Mac path. Minimum memory: 12 GB for a 1.7B to 4B base at bfloat16 with a rank-16 adapter, batch 1 and a 1,024-token sequence. Part 11's memory lesson has the arithmetic. Assumes: torch, transformers, trl, peft and datasets installed in the active environment; make-format-dataset.py has been run so that data/train.jsonl and data/valid.jsonl exist; sftlog.py sits next to this file. Usage: python3 train-lora.py --model Qwen/Qwen3-1.7B --data-dir data \ --output-dir runs/format-qwen3-1.7b --labbook labbook.md python3 train-lora.py --model Qwen/Qwen3-4B --list-modules python3 train-lora.py --model Qwen/Qwen3-4B --epochs 3 --rank 32 --alpha 64 \ --gradient-checkpointing --output-dir runs/format-qwen3-4b """ 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, EarlyStoppingCallback from trl import SFTConfig, SFTTrainer import sftlog # The seven linear projections of a Qwen3 block: four in attention, three in the # feed-forward network. --list-modules prints what your own base model actually has, # because a name that does not match attaches nothing and raises no error. DEFAULT_TARGETS = ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"] def pick_device() -> str: """CUDA (or ROCm, which reports as cuda), then MPS, then CPU. Same choice as Part 11.""" 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: """bfloat16 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, module in model.named_modules() if isinstance(module, torch.nn.Linear)}) print(f"linear module names in {model_id}:") for name in names: print(f" {name}") print("\nPass the ones you want with --target-modules, or use --target-modules all-linear.") def summarise_history(history: list[dict]) -> dict[str, float | int | None]: """The five numbers worth keeping out of a log full of them.""" train_losses = [row["loss"] for row in history if "loss" in row] evals = [(row.get("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 { "first_train_loss": round(train_losses[0], 4) if train_losses else None, "final_train_loss": round(train_losses[-1], 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, "evaluations": len(evals), } def main() -> None: parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) parser.add_argument("--model", default="Qwen/Qwen3-1.7B", help="base model repository id or local path; use an instruct " "checkpoint, which already carries a chat template") parser.add_argument("--data-dir", default="data", help="directory holding train.jsonl and valid.jsonl") parser.add_argument("--output-dir", default="runs/format-lora") parser.add_argument("--epochs", type=float, default=3.0) parser.add_argument("--batch-size", type=int, default=1, help="per-device batch size") parser.add_argument("--grad-accum", type=int, default=8, help="batches summed before one optimiser step; only --batch-size " "costs memory, so this is how you raise the effective batch") parser.add_argument("--lr", type=float, default=1e-4, help="adapters take roughly 1e-4, not the SFTConfig default of 2e-5") parser.add_argument("--max-length", type=int, default=1024) 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, help='module names, or the single word all-linear') parser.add_argument("--use-dora", action="store_true", help="decompose the update into magnitude and direction; helps most " "at low rank, and makes each step slower") parser.add_argument("--use-rslora", action="store_true", help="scale by alpha over the square root of the rank; for unstable " "high-rank runs") 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("--early-stopping-patience", type=int, default=2, help="stop after this many evaluations without an improvement; 0 disables") 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("--notes", default=None, help="one line about what this run is testing") parser.add_argument("--list-modules", action="store_true", help="print the base model's 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'}") if device == "mps": print("note: this is Track M's fallback path in float32. train-lora-mlx.sh is faster " "and is the path the lab describes for a Mac.") 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 ({split} split); run make-format-dataset.py first") dataset = load_dataset("json", data_files=files) print(f"train examples: {len(dataset['train'])} " f"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, so it is a base checkpoint rather than an " "instruct one. Either pick an instruct model or set chat_template_path in SFTConfig." ) targets = args.target_modules[0] if args.target_modules == ["all-linear"] else args.target_modules 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=False, completion_only_loss=True, # the loss lands on the answer, not on your questions gradient_checkpointing=args.gradient_checkpointing, bf16=bf16, model_init_kwargs={"dtype": dtype}, eval_strategy="epoch", save_strategy="epoch", # must match eval_strategy for the next line to work 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="none", seed=args.seed, data_seed=args.seed, ) peft_config = LoraConfig( r=args.rank, lora_alpha=args.alpha, lora_dropout=args.dropout, target_modules=targets, use_dora=args.use_dora, use_rslora=args.use_rslora, bias="none", task_type="CAUSAL_LM", ) callbacks = [] if args.early_stopping_patience > 0: callbacks.append(EarlyStoppingCallback(early_stopping_patience=args.early_stopping_patience)) trainer = SFTTrainer( model=args.model, args=config, train_dataset=dataset["train"], eval_dataset=dataset["validation"], processing_class=tokenizer, peft_config=peft_config, callbacks=callbacks or None, ) # If this percentage is not roughly what the adapter arithmetic predicted, the target # module names did not match and nothing was attached. trainer.model.print_trainable_parameters() started = time.time() trainer.train() elapsed = time.time() - started trainer.save_model(args.output_dir) # The chat template travels in the tokeniser files, and it is the thing that has to # match at serving time. Saving it here is not optional. 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 losses["best_epoch"] is not None and losses["best_epoch"] <= 1: print("the best epoch was the first: this dataset is small for this many epochs, or " "the learning rate is too high. Read the two curves before you train again.") if args.labbook: record = sftlog.record( labbook=args.labbook, lab="part-13/train-lora", model=args.model, dataset={ "path": files["train"], "sha256": sftlog.file_sha256(files["train"]), "train_examples": len(dataset["train"]), "validation_examples": len(dataset["validation"]), }, hyperparameters={ "method": "dora" if args.use_dora else "lora", "rank": args.rank, "alpha": args.alpha, "dropout": args.dropout, "target_modules": targets, "use_rslora": args.use_rslora, "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, "gradient_checkpointing": args.gradient_checkpointing, "early_stopping_patience": args.early_stopping_patience, "precision": "bfloat16" if bf16 else "float32", "completion_only_loss": True, "output_dir": args.output_dir, }, seed=args.seed, losses=losses, scores={}, config_path=__file__, notes=args.notes, ) print(f"recorded run {record['run_id']} in {args.labbook}") if __name__ == "__main__": main()