"""Train a LoRA adapter on agent trajectories with TRL's SFT trainer and PEFT. Purpose: Part 27's supervised fine-tuning run for Tracks S, X and N. Loads the conversational tool-calling files trajectories-to-sft.py wrote, including the `tools` column that the chat template turns into the tool section of the prompt, attaches a LoRA adapter, puts the loss on the assistant turns only, evaluates after every epoch, keeps the best checkpoint and appends one run record to the lab notebook. It is Part 13's recipe with two changes: the dataset is multi-turn and carries tool calls, and the loss mask is assistant-only rather than completion-only. Platform: spark, strix, nvidia (CUDA, or ROCm which also reports as cuda to PyTorch). It runs on the CPU too, slowly. Track M uses mlx_lm.lora on the data-mlx layout instead; PyTorch's MPS backend will run this in float32 if you insist. Minimum memory: 16 GB for a 1.7B to 4B base at bfloat16 with a rank-16 adapter, batch 1 and a 4,096-token sequence. Agent trajectories are long: the sequence length, not the parameter count, is what makes this part heavier than Part 13. Assumes: torch, transformers, trl, peft and datasets installed in the active environment; trajectories-to-sft.py has been run so that data/train.jsonl and data/valid.jsonl exist; agentlog.py sits next to this file. Usage: python3 train-agent-sft.py --model Qwen/Qwen3-1.7B --data-dir data \\ --output-dir runs/agent-qwen3-1.7b --labbook labbook.md python3 train-agent-sft.py --model Qwen/Qwen3-4B --max-length 6144 --rank 32 \\ --gradient-checkpointing --output-dir runs/agent-qwen3-4b python3 train-agent-sft.py --model Qwen/Qwen3-4B --inspect-only --inspect-only renders the first training example through the model's own chat template and prints it. Do that before every run. It is the only way to see whether the tool schemas reached the prompt, whether the tool results are inside the turn structure, and whether the template renders tool calls at all: a template that ignores the tools column trains the model on a prompt it will never see again. """ from __future__ import annotations import argparse import json import time from pathlib import Path from typing import Any import torch from datasets import Dataset from peft import LoraConfig from transformers import AutoModelForCausalLM, AutoTokenizer, EarlyStoppingCallback from trl import SFTConfig, SFTTrainer import agentlog # The seven linear projections of a Qwen3 block. --list-modules prints what your own base # model 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. Part 11's choice.""" 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: 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: 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 read_rows(path: Path) -> list[dict[str, Any]]: if not path.is_file(): raise SystemExit(f"{path} is missing; run trajectories-to-sft.py first") rows = [json.loads(line) for line in path.read_text(encoding="utf-8").splitlines() if line.strip().startswith("{")] if not rows: raise SystemExit(f"{path} holds no rows") return rows def to_dataset(rows: list[dict[str, Any]]) -> tuple[Dataset, str]: """Build a Dataset whose `tools` column survives the trip. A tool schema is an arbitrary JSON object, so the column cannot be typed as a struct with fixed fields. Recent versions of `datasets` have a Json() type for exactly this; older ones do not, and the documented fallback is to store the column as a JSON string, which the chat template parses. The returned string says which happened, so the run record can say it too. """ try: from datasets import Features, Json, List, Value except ImportError: encoded = [{"messages": r["messages"], "tools": json.dumps(r["tools"])} for r in rows] return Dataset.from_list(encoded), "tools-as-json-string" features = Features({ "messages": List({"role": Value("string"), "content": Value("string"), "name": Value("string"), "tool_calls": List(Json())}), "tools": List(Json()), }) normalised = [] for row in rows: messages = [] for message in row["messages"]: messages.append({ "role": message["role"], "content": message.get("content") or "", "name": message.get("name") or "", "tool_calls": message.get("tool_calls") or [], }) normalised.append({"messages": messages, "tools": row["tools"]}) return Dataset.from_list(normalised, features=features), "tools-as-json-objects" def summarise_history(history: list[dict]) -> dict[str, float | int | None]: 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") parser.add_argument("--output-dir", default="runs/agent-lora") parser.add_argument("--epochs", type=float, default=3.0) parser.add_argument("--batch-size", type=int, default=1) 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=4096, help="trajectories are long; anything longer than this is truncated " "from the end, which silently removes the final answer") 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("--loss", choices=["assistant", "full"], default="assistant", help="assistant: loss on the assistant turns only, which is what you " "want, because the tool results are the environment's words") parser.add_argument("--precision", choices=["auto", "bf16", "fp32"], default="auto") parser.add_argument("--gradient-checkpointing", action="store_true") parser.add_argument("--early-stopping-patience", type=int, default=2) parser.add_argument("--seed", type=int, default=0) parser.add_argument("--labbook", default=None) parser.add_argument("--notes", default=None) parser.add_argument("--list-modules", action="store_true") parser.add_argument("--inspect-only", action="store_true", help="render the first example through the chat template and stop") args = parser.parse_args() if args.list_modules: list_linear_modules(args.model) return data_dir = Path(args.data_dir) train_rows = read_rows(data_dir / "train.jsonl") valid_rows = read_rows(data_dir / "valid.jsonl") 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. Tool calling is a template feature: pick an instruct model." ) if args.inspect_only: rendered = tokenizer.apply_chat_template( train_rows[0]["messages"], tools=train_rows[0]["tools"], tokenize=False, add_generation_prompt=False) print(rendered) length = len(tokenizer(rendered)["input_ids"]) print(f"\n--- {length} token(s) for this example; --max-length is {args.max_length}") names = [c["function"]["name"] for m in train_rows[0]["messages"] for c in (m.get("tool_calls") or [])] for name in sorted(set(names)): if name not in rendered: print(f"WARNING: the call to {name!r} does not appear in the rendered text. " "This template may not render tool calls; check the model card.") if train_rows[0]["tools"] and train_rows[0]["tools"][0]["function"]["name"] not in rendered: print("WARNING: the tool list does not appear in the rendered text. Either the " "template ignores `tools`, or it puts them somewhere this check cannot see. " "Read the output above before training on it.") 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: Track M's supported path is mlx_lm.lora on the data-mlx layout. This " "will run in float32 and want roughly twice the memory.") train_dataset, tools_encoding = to_dataset(train_rows) eval_dataset, _ = to_dataset(valid_rows) print(f"train rows: {len(train_dataset)} validation rows: {len(eval_dataset)} " f"tools column: {tools_encoding}") 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, # The tool results are the environment speaking. Training on them teaches the # model to write file contents and command output, which is not the job. assistant_only_loss=(args.loss == "assistant"), 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="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, 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=train_dataset, eval_dataset=eval_dataset, 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. For a tool-calling fine-tune it is the whole contract. 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: train_path = str(data_dir / "train.jsonl") record = agentlog.record( labbook=args.labbook, lab="part-27/train-agent-sft", model=args.model, dataset={"path": train_path, "sha256": agentlog.file_sha256(train_path), "train_examples": len(train_dataset), "validation_examples": len(eval_dataset)}, data_lineage=agentlog.lineage(trajectories=train_path, tools_encoding=tools_encoding), hyperparameters={ "method": "lora", "rank": args.rank, "alpha": args.alpha, "dropout": args.dropout, "target_modules": targets, "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, "loss_mask": args.loss, "gradient_checkpointing": args.gradient_checkpointing, "early_stopping_patience": args.early_stopping_patience, "precision": "bfloat16" if bf16 else "float32", "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()