Skip to content
Level 3 · Model BuilderLabPart 11 · page 6 of 660 minSXMN 8 GB
60Minutes
5Tools
11Sources
All fourTracks
Tools used on this page5

Lab: Your First Training Run

Validated on: written from the documentation cited above; not yet validated on hardware on any track. The versions and measurements for each track will be recorded here when the validation pass is done.

By the end of this lab you will have fine-tuned a language model on your own machine, from a dataset you generated to a quantised GGUF file that llama.cpp can serve, and recorded the run in the lab notebook in the format every later training lab uses.

The model is Qwen3-0.6B, whose card gives 0.6 billion parameters, 0.44 billion of them outside the embeddings, 28 layers, 16 query heads and 8 key-value heads, a context length of 32,768 tokens, and an Apache-2.0 licence. It is not one of the course’s reference models, because it is smaller than anything worth serving; it is exactly the right size for learning the machinery, because the whole run finishes in minutes on every track and the arithmetic from this part’s third lesson says it fits in the 8 GB tier with room to spare.

The task is narrow on purpose: teach the model to answer short questions about local-model operations in a fixed two-line house style, with the answer first and the reasoning second. A narrow task on a small dataset makes the effect of training visible in minutes, which is what you want the first time through.

Every track needs the Part 1 environment with this part’s training libraries added, about 3 GB of free disk for the base model, the merged copy and two GGUF files, and an hour. Every track also needs a llama.cpp checkout built as in Part 6, for the export step at the end.

Track S — NVIDIA DGX Spark

Run inside the NGC PyTorch container from Part 1, with the training libraries installed inside it as the environment lesson describes. Start it with your course directory mounted at /workspace/course and work from there, so that everything the lab writes survives the container exiting.

The Spark’s 128 GB pool makes this lab’s memory arithmetic irrelevant, which is a good reason to do the arithmetic step anyway: it is the habit, not the number, that Part 13 needs.

Track X — AMD Ryzen AI Max+ 395Partial

Training needs the ROCm build of PyTorch. The ROCm 10.0.0 compatibility matrix dated 2026-08-14 lists gfx1151 without a support-tier qualifier, while AMD's PyTorch install page read on 2026-09-09 does not mention the chip. The CPU path finishes this lab and is a supported way to complete it.

Use the ROCm wheels from AMD’s install page, as Part 1’s setup-env.sh does with TRACK=strix, and confirm with python -c 'import torch; print(torch.cuda.is_available())' before starting. Vulkan is not a training path; this step needs ROCm or the CPU.

On the CPU this lab still finishes, more slowly. Add --precision fp32 and expect the training step to be the slow part. Record in the notebook that the run was on the CPU: a CPU run and a GPU run are not interchangeable results.

Track M — Apple silicon

Track M takes the MLX route for training. mlx-lm, pinned at mlx-lm 0.31.3 · verified 2026-09-08, ships a LoRA fine-tuning command, and the lab’s train-sft-mlx.sh wraps it. The base model is the MLX community’s bfloat16 conversion of the same Qwen3-0.6B weights, so the model is the same and only the framework differs.

The export step at the end is the one place where Track M needs the PyTorch path as well. mlx-lm’s own GGUF export is documented as limited to Mistral, Mixtral and Llama style models in fp16, which does not cover Qwen3, so producing a GGUF file means running train-sft.py on PyTorch’s MPS backend in float32 and exporting that adapter. A 0.6 billion parameter model makes that a few extra minutes rather than an afternoon, and the two runs are worth comparing anyway. The task section says exactly where to branch.

Track N — NVIDIA desktop or laptop

A card with 8 GB of VRAM is enough, and the arithmetic step below shows why. If the card is also driving your displays, close anything heavy first: the budget in this lab has about half the tier free, and a browser can eat that.

On Windows, work inside WSL2 exactly as in Part 1.

Working directory and terminal roles

Prepare the course execution workspace once before this procedure. It includes this part's scripts, data and shared Python helpers. In the client or training terminal, select this directory:

RunnableAll tracks

select this part’s execution directory
export LABS_ROOT="${LABS_ROOT:-$HOME/llm-course/labs}"
export LAB_DIR="$LABS_ROOT/part-11-training-toolchain"
cd "$LAB_DIR"
pwd
test -f "make-dataset.py"

Expected result: pwd ends in part-11-training-toolchain and the file check returns successfully. If it does not, finish workspace preparation before continuing. Activate the environment in the requirements for your track. Bare script and data filenames below are relative to this directory; paths to earlier experiments must point at the artefacts you actually retained.

Keep each foreground server in a separate terminal and send requests from this terminal. Reapply lesson-specific environment variables in each new shell. Stop at the first failed checkpoint and retain its output; the execution guide explains how to distinguish missing files, endpoint failures and capacity problems.

1. Prepare the environment and the directory

Section titled “1. Prepare the environment and the directory”

Everything happens in the course directory from Part 1. Activate the environment and confirm that the training libraries are there.

RunnableAll tracks

activate and check
cd "$LAB_DIR"
source "$HOME/llm-course/.venv/bin/activate"
python -c "import torch, transformers, trl, peft, datasets; print(torch.__version__, transformers.__version__, trl.__version__, peft.__version__, datasets.__version__)"

If that import fails, go back to this part’s second lesson: the libraries go into the same environment as PyTorch, and on Track S they go inside the container.

Download the four lab files below into the course directory. They are the same bytes the page shows, and runlog.py must sit beside the training script because the script imports it.

The dataset is generated rather than downloaded, so its provenance is a script you can read and its licence is the course’s. The generator writes the whole set as JSON Lines, then splits it with a fixed seed into the two shapes the two training paths expect: TRL’s conversational prompt-completion form for Tracks S, X and N, and mlx-lm’s completions form for Track M.

RunnableAll tracks

make-dataset.py
"""Build the course's small instruction dataset and split it into train, validation and test.
Purpose: write about two hundred short instruction examples in the course's own
two-line house style, then split them with a fixed seed into the files
the two training paths expect: TRL's conversational prompt-completion
shape for Tracks S, X and N, and mlx-lm's completions shape for Track M.
Platform: all (standard library only)
Minimum memory: 8 GB
Assumes: nothing beyond Python 3.10 or newer; writes into the directory given by
--out-dir, which is created if it does not exist.
Usage: python make-dataset.py --out-dir . --seed 0
Writes sample-instructions.jsonl (the whole dataset, one example per line),
data/{train,valid,test}.jsonl (TRL conversational prompt-completion)
data-mlx/{train,valid,test}.jsonl (mlx-lm completions)
Every answer is two lines:
Answer: <the answer, one sentence>
Because: <the rule or the arithmetic that produced it>
The content is deliberately narrow. The arithmetic items are generated from the
formulas taught in Part 4 and Part 11, so every answer in this file is correct by
construction rather than by an author's memory; the written items restate the
course's own definitions. Nothing here is scraped, and nothing here is anyone
else's copyright, so the dataset carries the course's content licence.
"""
from __future__ import annotations
import argparse
import hashlib
import json
import random
from pathlib import Path
GB = 1e9
GIB = 1024 ** 3
# ---------------------------------------------------------------------------
# Written examples: the course's own definitions and rules of thumb.
# ---------------------------------------------------------------------------
WRITTEN: list[tuple[str, str, str]] = [
("What is a parameter?",
"A parameter is one of the numbers a model is made of.",
"A model's size in billions, such as 8B, is a count of its parameters."),
("What is the difference between a weight and a parameter?",
"A weight is the value of a parameter, and the weights are the checkpoint.",
"Parameter names the slot; weight names the number currently in it."),
("What is the difference between a context window and a context length?",
"The context window is what the model was trained to handle; the context length is what you asked the engine to allocate for this run.",
"One is a property of the model, the other a setting of the run."),
("What is prefill?",
"Prefill is the stage that reads the prompt.",
"It processes many tokens at once, so it is compute-bound and measured in prompt tokens per second."),
("What is decode?",
"Decode is the stage that writes the answer, one token at a time.",
"Each token reads every weight once, so it is bandwidth-bound."),
("Is decode limited by compute or by memory bandwidth?",
"By memory bandwidth.",
"Every generated token reads the whole model from memory and does only one multiply-add per parameter."),
("Is prefill limited by compute or by memory bandwidth?",
"By compute.",
"The weights are read once and reused across all the prompt's tokens, so the arithmetic dominates."),
("What do total and active parameters mean for a mixture-of-experts model?",
"Total parameters is what must sit in memory; active parameters is what is computed with per token.",
"A mixture-of-experts model routes each token to a few experts, so it is large to hold and cheap to run."),
("What is quantisation?",
"Quantisation is rounding the weights to fewer bits.",
"It reduces bytes per parameter, which reduces both the memory the model needs and the bytes read per token."),
("What is the difference between quantisation and compression?",
"Quantisation rounds the weights to fewer bits; compression removes parameters or layers.",
"One keeps every parameter at lower precision, the other has fewer parameters afterwards."),
("What does open weight mean?",
"Open weight means the weights are downloadable.",
"It does not by itself mean the training data or the training code are available."),
("What is a fine-tune?",
"A fine-tune adjusts an existing model on new data.",
"Training from scratch instead makes a model from raw text."),
("What is tensor parallelism?",
"Tensor parallelism splits every layer across devices.",
"The devices exchange activations inside each layer, so it needs a fast link between them."),
("What is pipeline parallelism?",
"Pipeline parallelism gives each device some of the layers.",
"Only the boundary activations cross the link, so it tolerates a slow one."),
("What is an agent, as this course uses the word?",
"An agent is a loop in which the model decides the next step.",
"A workflow is a fixed sequence that the program decides instead."),
("What is a memory floor?",
"A memory floor is the smallest memory that can follow a page's primary path.",
"It is stated per lab so a reader knows before starting whether their machine qualifies."),
("Why does training need more memory than inference for the same model?",
"Training holds gradients, optimiser states and activations as well as the weights.",
"Inference holds only the weights and the key-value cache."),
("What is gradient checkpointing for?",
"Gradient checkpointing trades compute for memory during training.",
"Intermediate activations are recomputed in the backward pass instead of being kept from the forward pass."),
("What is gradient accumulation for?",
"Gradient accumulation gives a large effective batch on a machine that cannot hold one.",
"Several small batches are summed before the optimiser takes a step."),
("What does LoRA train?",
"LoRA trains a pair of small low-rank matrices beside each frozen weight matrix.",
"The base model's weights are not updated, so only the adapters need gradients and optimiser states."),
("What does the rank of a LoRA adapter control?",
"The rank sets the size of the two update matrices and therefore the number of trainable parameters.",
"A lower rank means smaller matrices and fewer parameters to train."),
("What is QLoRA?",
"QLoRA trains LoRA adapters on top of a base model quantised to four bits.",
"Gradients are backpropagated through the frozen quantised weights into the adapters."),
("Can a LoRA adapter be merged into the base model?",
"Yes, the adapter can be merged into the base weights to produce a standalone model.",
"After merging there is one set of weights again, so the adapter adds no inference cost."),
("Why does a chat template matter when fine-tuning?",
"The template decides the exact control tokens the model sees around each turn.",
"Training with one template and serving with another gives the model a format it was never trained on."),
("What is completion-only loss?",
"Completion-only loss computes the training loss on the answer tokens and ignores the prompt tokens.",
"The model is being taught what to answer, not how to restate the question."),
("What is packing in a training run?",
"Packing groups several examples into one fixed-length sequence.",
"It reduces the padding that would otherwise be computed and thrown away."),
("Why must a validation split never be trained on?",
"Because a split that has been trained on cannot detect memorisation.",
"Validation loss is only informative about unseen data if the data really was unseen."),
("What is leakage between a training set and an evaluation set?",
"Leakage is any overlap that lets a model score well by recall rather than by generalisation.",
"Near-duplicates leak as surely as exact duplicates do."),
("Why record a seed with every training run?",
"Because without the seed the run cannot be repeated even on the same machine.",
"The seed fixes the shuffling, the dropout and the initialisation."),
("Does a fixed seed guarantee identical results on two different machines?",
"No, it does not.",
"Results are not reproducible across releases, commits or platforms unless deterministic algorithms are also requested."),
("What should a run log record besides the final loss?",
"The model, the dataset and its hash, the hyperparameters, the seed, the hardware, the versions and the date.",
"A loss without its settings cannot be compared with anything later."),
("Why hash the dataset file in the run log?",
"Because the hash ties a result to the exact bytes that produced it.",
"A file with the same name can be edited between two runs without anyone noticing."),
("What does an evaluation loss that rises while training loss falls indicate?",
"It indicates overfitting.",
"The model is fitting the training examples rather than the pattern they share."),
("Which checkpoint should be kept when validation loss turns?",
"The checkpoint from the epoch with the lowest validation loss.",
"Later epochs fit noise, and the validation curve is what identifies the turn."),
("What is a base model?",
"A base model has been pretrained but not instruction-tuned.",
"It continues text rather than answering a question in a chat format."),
("Why is BF16 preferred over FP16 for training?",
"Because BF16 keeps the range of FP32.",
"Small gradients underflow to zero in FP16 unless the loss is scaled first."),
("How many bytes per parameter does BF16 use?",
"Two bytes per parameter.",
"Sixteen bits is two bytes, which is the baseline every quantised format is compared against."),
("Why is Q8_0 about 1.06 bytes per parameter rather than exactly 1.0?",
"Because each block of 32 weights carries a scale as well.",
"32 eight-bit weights plus a 16-bit scale is 272 bits for 32 weights, which is 8.5 bits each."),
("Which head count belongs in the key-value cache formula?",
"The key-value head count, not the query head count.",
"Grouped-query attention lets several query heads share one key-value head, so the two numbers differ."),
("Does the key-value cache matter during training?",
"No, it is switched off during training.",
"Training computes a forward and a backward pass over a fixed sequence and generates nothing."),
("What is unified memory, on the machines this course uses?",
"Unified memory is one pool shared by the CPU and the GPU with no copy between them.",
"It trades bandwidth for capacity compared with a discrete card's own memory."),
("Why is a model that fits exactly into memory a bad idea?",
"Because a machine at its memory ceiling swaps or compresses instead of failing cleanly.",
"The symptom is a model that loads and runs absurdly slowly rather than an error."),
]
# ---------------------------------------------------------------------------
# Generated examples: arithmetic that is correct by construction.
# ---------------------------------------------------------------------------
FORMATS = [("BF16", 2.0), ("FP8", 1.0), ("Q8_0", 1.06), ("Q6_K", 0.82), ("Q4_K_M", 0.61)]
PARAMS_B = [0.6, 1.7, 3.0, 4.0, 7.0, 8.0, 12.0, 14.0, 20.0, 27.0, 32.0, 70.0]
KV_SHAPES = [
("28 layers, 8 key-value heads and a head dimension of 128", 28, 8, 128),
("36 layers, 8 key-value heads and a head dimension of 128", 36, 8, 128),
("40 layers, 8 key-value heads and a head dimension of 128", 40, 8, 128),
("48 layers, 4 key-value heads and a head dimension of 128", 48, 4, 128),
("64 layers, 8 key-value heads and a head dimension of 128", 64, 8, 128),
("24 layers, 8 key-value heads and a head dimension of 64", 24, 8, 64),
("36 layers, 8 key-value heads and a head dimension of 64", 36, 8, 64),
("32 layers, 8 key-value heads and a head dimension of 128", 32, 8, 128),
("16 layers, 4 key-value heads and a head dimension of 64", 16, 4, 64),
("60 layers, 8 key-value heads and a head dimension of 128", 60, 8, 128),
]
SHAPES = [(1024, 2048), (1024, 3072), (2048, 6144), (2560, 9728), (4096, 12288), (4096, 4096),
(2048, 2048), (3072, 1024), (6144, 2048), (5120, 13824), (8192, 28672), (1536, 4096)]
FILE_GB = [1.1, 2.5, 4.3, 5.03, 8.0, 12.12, 18.56, 19.76, 32.48, 63.39]
def _n(value: float, places: int = 2) -> str:
"""Trim trailing zeros so the answers read like a person wrote them."""
text = f"{value:.{places}f}".rstrip("0").rstrip(".")
return text or "0"
def generated() -> list[tuple[str, str, str]]:
items: list[tuple[str, str, str]] = []
for params in PARAMS_B:
for name, bpp in FORMATS:
gb = params * bpp
items.append((
f"How much memory do the weights of a {_n(params)} billion parameter model need at {name}?",
f"About {_n(gb)} GB.",
f"{_n(params)} billion parameters at {_n(bpp)} bytes each is {_n(gb)} x 10^9 bytes.",
))
for name, bpp in FORMATS:
items.append((
f"How many bytes per parameter is {name}?",
f"About {_n(bpp)} bytes per parameter.",
"The figure includes the per-block scales that are stored alongside the weights."
if bpp not in (1.0, 2.0)
else f"{int(bpp * 8)} bits is {_n(bpp)} bytes, with no per-block scale to add.",
))
for label, layers, kv_heads, head_dim in KV_SHAPES:
per_token = 2 * layers * kv_heads * head_dim * 2
items.append((
f"A model has {label}. How many bytes per token is its FP16 key-value cache?",
f"{per_token:,} bytes per token.",
f"2 x {layers} x {kv_heads} x {head_dim} x 2 = {per_token:,}.",
))
for ctx in (8192, 32768):
total = per_token * ctx / GB
items.append((
f"A model has {label}. How much memory is its FP16 key-value cache at {ctx:,} tokens?",
f"About {_n(total)} GB.",
f"{per_token:,} bytes per token times {ctx:,} tokens is {_n(total)} x 10^9 bytes.",
))
for params in PARAMS_B:
items.append((
f"How much memory do the weights, gradients, master weights and Adam states of a "
f"{_n(params)} billion parameter model need for a full fine-tune in mixed precision?",
f"About {_n(params * 16)} GB.",
"Two bytes of BF16 weights, two of gradients, four of FP32 master weights and eight of "
f"Adam states is 16 bytes per parameter, so {_n(params)} x 16 = {_n(params * 16)} GB.",
))
items.append((
f"How much memory do the Adam states of a {_n(params)} billion parameter model need at FP32?",
f"About {_n(params * 8)} GB.",
f"Two FP32 states per parameter is 8 bytes, so {_n(params)} x 8 = {_n(params * 8)} GB.",
))
items.append((
f"How much memory do BF16 gradients for a {_n(params)} billion parameter model need?",
f"About {_n(params * 2)} GB.",
f"One gradient per parameter at two bytes is {_n(params)} x 2 = {_n(params * 2)} GB.",
))
for inputs, outputs in SHAPES:
for rank in (8, 16):
added = rank * (inputs + outputs)
items.append((
f"A LoRA adapter of rank {rank} is added to a linear layer with {inputs:,} inputs and "
f"{outputs:,} outputs. How many trainable parameters does it add?",
f"{added:,} trainable parameters.",
f"The two matrices are {rank} x {inputs:,} and {outputs:,} x {rank}, so "
f"{rank} x ({inputs:,} + {outputs:,}) = {added:,}.",
))
for gb in FILE_GB:
gib = gb * GB / GIB
items.append((
f"A file listing gives a model file as {_n(gb)} GB. What will a file manager counting in gibibytes show?",
f"About {_n(gib)} GiB.",
f"{_n(gb)} x 10^9 bytes divided by 2^30 is {_n(gib)}, because a gibibyte is about seven per cent larger than a gigabyte.",
))
return items
def build() -> list[dict[str, str]]:
rows = [{"instruction": q, "answer": f"Answer: {a}\nBecause: {b}"} for q, a, b in WRITTEN + generated()]
seen: set[str] = set()
unique = []
for row in rows:
if row["instruction"] in seen:
continue
seen.add(row["instruction"])
unique.append(row)
return unique
def write_jsonl(path: Path, lines: list[dict]) -> None:
path.parent.mkdir(parents=True, exist_ok=True)
with path.open("w", encoding="utf-8") as handle:
for line in lines:
handle.write(json.dumps(line, ensure_ascii=False) + "\n")
def main() -> None:
parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
parser.add_argument("--out-dir", default=".", help="directory to write the dataset and splits into")
parser.add_argument("--seed", type=int, default=0, help="seed for the shuffle before splitting")
parser.add_argument("--valid-frac", type=float, default=0.15)
parser.add_argument("--test-frac", type=float, default=0.10)
args = parser.parse_args()
rows = build()
random.Random(args.seed).shuffle(rows)
n_valid = int(len(rows) * args.valid_frac)
n_test = int(len(rows) * args.test_frac)
splits = {
"valid": rows[:n_valid],
"test": rows[n_valid:n_valid + n_test],
"train": rows[n_valid + n_test:],
}
out = Path(args.out_dir)
write_jsonl(out / "sample-instructions.jsonl", rows)
for name, part in splits.items():
# TRL: conversational prompt-completion. Loss is computed on the completion only.
write_jsonl(out / "data" / f"{name}.jsonl", [
{"prompt": [{"role": "user", "content": r["instruction"]}],
"completion": [{"role": "assistant", "content": r["answer"]}]}
for r in part
])
# mlx-lm: the completions format, two plain strings.
write_jsonl(out / "data-mlx" / f"{name}.jsonl", [
{"prompt": r["instruction"], "completion": r["answer"]} for r in part
])
digest = hashlib.sha256((out / "sample-instructions.jsonl").read_bytes()).hexdigest()
print(f"examples: {len(rows)}")
print(f" train {len(splits['train'])} valid {len(splits['valid'])} test {len(splits['test'])}")
print(f"written: {out / 'sample-instructions.jsonl'}, {out / 'data'}/, {out / 'data-mlx'}/")
print(f"sha256(sample-instructions.jsonl): {digest}")
print("record that hash in the run log; it is what ties a result to this exact dataset.")
if __name__ == "__main__":
main()

Download make-dataset.py339 lines

RunnableAll tracks

generate the dataset and its splits
python make-dataset.py --out-dir . --seed 0

Output — what you should see

examples: 207
train 156 valid 31 test 20
written: sample-instructions.jsonl, data/, data-mlx/
sha256(sample-instructions.jsonl): 9cda81d1...
record that hash in the run log; it is what ties a result to this exact dataset.

Look at three lines before you train on them. This is the step people skip and then spend an hour regretting.

RunnableAll tracks

read the data you are about to train on
head -n 2 data/train.jsonl
head -n 2 data-mlx/train.jsonl

Notice what differs between the two files: the same example is a pair of message lists in one and a pair of plain strings in the other. Both are documented shapes, for two different trainers, from one source. Notice also that every completion has the same two-line structure. That structure is what the fine-tune is being asked to learn.

Apply this part’s third lesson to this lab’s settings: a 0.6 billion parameter base held frozen in BF16, a rank-16 adapter on the seven projections of each of 28 layers, batch 2 at up to 512 tokens, and Qwen3’s vocabulary of 151,936 entries.

This lab's budget: LoRA on Qwen3-0.6B, batch 2 at 512 tokens, on an 8 GB machine

Frozen base weights, BF16
1.2 GB
Adapter weights, gradients and Adam states
0.2 GB
Activations and logits
1 GB
Reserved for the operating system
2 GB
Free
3.6 GB
Total
8 GB
Estimate from arithmetic, not a measurement. The frozen base is 2 bytes against 0.6 billion parameters; the adapter is about 10.1 million parameters at 16 bytes each; the activation figure is the layer-boundary estimate plus the logits tensor and a full-precision copy of it for the loss, because TRL's chunked cross-entropy is documented as incompatible with PEFT; the reserve is an allowance.

About half the 8 GB tier is free, which is the margin that lets the lab say “8 GB” honestly. Write those four figures in the notebook now, so that when the training script prints its peak memory you have a prediction to compare it against.

Before training, see what the model does with the task untouched. This is the comparison the whole lab exists to make, and it takes thirty seconds.

RunnableAll tracks

the base model, before any training
python - <<'PY'
from transformers import pipeline
pipe = pipeline("text-generation", model="Qwen/Qwen3-0.6B")
prompt = [{"role": "user", "content": "How much memory do the weights of a 5 billion parameter model need at BF16?"}]
print(pipe(prompt, max_new_tokens=128)[0]["generated_text"][-1]["content"])
PY

Expect a plausible answer in the model’s own voice, at whatever length it chooses, quite possibly with a long chain of reasoning first. What it will not be is two lines beginning Answer: and Because:. Copy what you get into the notebook; you will compare it with the same prompt at the end.

Track S — NVIDIA DGX Spark

RunnableTrack S · DGX Spark

train-sft.py
"""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()

Download train-sft.py221 lines

RunnableTrack S · DGX Spark

LoRA fine-tune, inside the container
python train-sft.py \
--model Qwen/Qwen3-0.6B \
--data-dir data \
--output-dir runs/sft-qwen3-0.6b \
--epochs 3 \
--batch-size 2 \
--grad-accum 4 \
--lr 1e-4 \
--labbook labbook.md

Track X — AMD Ryzen AI Max+ 395Partial

ROCm build required for the GPU path; the CPU path finishes the lab with --precision fp32.

RunnableTrack X · Ryzen AI Max+

train-sft.py
"""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()

RunnableTrack X · Ryzen AI Max+

LoRA fine-tune on ROCm, or on the CPU
python train-sft.py \
--model Qwen/Qwen3-0.6B \
--data-dir data \
--output-dir runs/sft-qwen3-0.6b \
--epochs 3 \
--batch-size 2 \
--grad-accum 4 \
--lr 1e-4 \
--labbook labbook.md

Track M — Apple silicon

RunnableTrack M · Apple silicon

train-sft-mlx.sh
#!/usr/bin/env bash
# Purpose: Track M's supervised fine-tuning run: train a LoRA adapter with mlx_lm.lora,
# report test-set perplexity, generate one answer with the adapter attached,
# and append a run record to the lab notebook
# Platform: mac (Apple silicon, MLX); Tracks S, X and N use train-sft.py instead
# Minimum memory: 8 GB
# Assumes: mlx-lm installed in the active environment; make-dataset.py has been run so
# that data-mlx/{train,valid,test}.jsonl exist; runlog.py sits next to this
# script; run from the course directory
#
# Usage: bash train-sft-mlx.sh [MODEL] [ITERS]
# MODEL defaults to the MLX community conversion of Qwen3-0.6B; ITERS to 300.
set -euo pipefail
MODEL="${1:-mlx-community/Qwen3-0.6B-bf16}"
ITERS="${2:-300}"
DATA="${DATA:-data-mlx}"
ADAPTERS="${ADAPTERS:-runs/sft-mlx-adapters}"
BATCH_SIZE="${BATCH_SIZE:-2}"
NUM_LAYERS="${NUM_LAYERS:-8}"
LEARNING_RATE="${LEARNING_RATE:-1e-4}"
SEED="${SEED:-0}"
LABBOOK="${LABBOOK:-labbook.md}"
HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
die() { echo "train-sft-mlx: $*" >&2; exit 1; }
command -v mlx_lm.lora >/dev/null || die "mlx_lm.lora is not on PATH; install mlx-lm in the active environment"
for split in train valid test; do
[[ -f "$DATA/$split.jsonl" ]] || die "$DATA/$split.jsonl is missing; run: python make-dataset.py --out-dir ."
done
echo "==> Training a LoRA adapter on $MODEL for $ITERS iterations"
START=$(date +%s)
mlx_lm.lora \
--model "$MODEL" \
--train \
--data "$DATA" \
--iters "$ITERS" \
--batch-size "$BATCH_SIZE" \
--num-layers "$NUM_LAYERS" \
--learning-rate "$LEARNING_RATE" \
--fine-tune-type lora \
--adapter-path "$ADAPTERS"
ELAPSED=$(( $(date +%s) - START ))
echo "==> Test-set loss and perplexity with the adapter attached"
mlx_lm.lora \
--model "$MODEL" \
--data "$DATA" \
--adapter-path "$ADAPTERS" \
--test | tee "$ADAPTERS/test.txt"
echo "==> One answer, to see whether the house style was learned"
mlx_lm.generate \
--model "$MODEL" \
--adapter-path "$ADAPTERS" \
--max-tokens 96 \
--prompt "How much memory do the weights of a 5 billion parameter model need at BF16?"
echo "==> Recording the run in $LABBOOK"
DATA_SHA=$(shasum -a 256 "$DATA/train.jsonl" | cut -d' ' -f1)
TRAIN_N=$(wc -l < "$DATA/train.jsonl" | tr -d ' ')
VALID_N=$(wc -l < "$DATA/valid.jsonl" | tr -d ' ')
python "$HERE/runlog.py" --record --labbook "$LABBOOK" <<JSON
{
"lab": "part-11/train-sft-mlx",
"model": "$MODEL",
"dataset": {"path": "$DATA/train.jsonl", "sha256": "$DATA_SHA",
"train_examples": $TRAIN_N, "validation_examples": $VALID_N},
"hyperparameters": {"method": "lora", "iters": $ITERS, "batch_size": $BATCH_SIZE,
"num_layers": $NUM_LAYERS, "learning_rate": $LEARNING_RATE,
"adapter_path": "$ADAPTERS", "seconds": $ELAPSED},
"seed": $SEED,
"losses": {},
"scores": {},
"notes": "copy the final train and validation loss, and the test perplexity, out of the mlx_lm.lora output above into losses"
}
JSON
echo "==> Done in ${ELAPSED}s. Adapter in $ADAPTERS; test output in $ADAPTERS/test.txt"

Download train-sft-mlx.sh80 lines

RunnableTrack M · Apple silicon

LoRA fine-tune with mlx-lm
bash train-sft-mlx.sh mlx-community/Qwen3-0.6B-bf16 300

The script trains for the number of iterations you give it, prints test-set loss and perplexity with the adapter attached, generates one answer so you can see the style immediately, and writes a run record. Copy the final training and validation losses out of the mlx-lm output into the losses field of that record, which the script leaves empty on purpose.

Track N — NVIDIA desktop or laptop

RunnableTrack N · NVIDIA GPU

train-sft.py
"""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()

RunnableTrack N · NVIDIA GPU

LoRA fine-tune
python train-sft.py \
--model Qwen/Qwen3-0.6B \
--data-dir data \
--output-dir runs/sft-qwen3-0.6b \
--epochs 3 \
--batch-size 2 \
--grad-accum 4 \
--lr 1e-4 \
--labbook labbook.md

The first lines of output are the ones to read.

Output — what you should see

device: cuda precision: bfloat16
train examples: 156 validation examples: 31
trainable params: 10,092,544 || all params: 60x,xxx,xxx || trainable%: 1.6xxx

That third line is PEFT’s print_trainable_parameters(), and it is the memory lesson’s arithmetic confirmed by the library: a rank-16 adapter on the seven projections of 28 layers is about ten million parameters, under two per cent of the model. Everything else is frozen, which is why this run fits where a full fine-tune of the same model would need about sixteen times the model’s size.

Then one line per logging interval with the training loss, and one evaluation line per epoch. The script keeps the checkpoint with the lowest evaluation loss and reloads it at the end.

RunnableAll tracks

the fine-tuned model, same prompt
python - <<'PY'
from peft import AutoPeftModelForCausalLM
from transformers import AutoTokenizer
model = AutoPeftModelForCausalLM.from_pretrained("runs/sft-qwen3-0.6b")
tok = AutoTokenizer.from_pretrained("runs/sft-qwen3-0.6b")
messages = [{"role": "user", "content": "How much memory do the weights of a 5 billion parameter model need at BF16?"}]
ids = tok.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_dict=True, return_tensors="pt")
out = model.generate(**ids, max_new_tokens=96, do_sample=False)
print(tok.decode(out[0][ids["input_ids"].shape[-1]:], skip_special_tokens=True))
PY

On Track M, the generation step already ran inside train-sft-mlx.sh; the equivalent by hand is mlx_lm.generate with --adapter-path pointing at the adapter directory.

Compare the two answers side by side in the notebook. The interesting change is the shape: two lines, Answer: then Because:. Whether the arithmetic in the answer is right is a separate question, and one this lab does not settle, because two hundred examples of a format do not make a small model good at mental arithmetic. Part 16 is where “did it get better” becomes a measurement rather than an impression.

The adapter is two small matrices per adapted layer. To serve it with the engines from Level 2 it has to become one model again.

RunnableAll tracks

merge-and-export.py
"""Merge a LoRA adapter into its base model, export to GGUF, quantise, and try it once.
Purpose: turn the adapter directory that train-sft.py produced into a single model
that every engine in Level 2 can serve: merge, save, convert to GGUF with
llama.cpp's convert_hf_to_gguf.py, quantise with llama-quantize, and run one
prompt through llama-cli so the export is proved rather than assumed.
Platform: spark, strix, nvidia (and mac for the merge and conversion steps; Track M's
adapters come from mlx_lm.fuse instead, and the page says so)
Minimum memory: 8 GB
Assumes: torch, transformers and peft installed; a llama.cpp checkout for
convert_hf_to_gguf.py and built binaries for llama-quantize and llama-cli,
as built in Part 6; enough disk for the merged model plus two GGUF files.
Usage: python merge-and-export.py --adapter runs/sft-qwen3-0.6b --llama-cpp ~/llama.cpp
python merge-and-export.py --adapter runs/sft-qwen3-0.6b --skip-gguf # merge only
"""
from __future__ import annotations
import argparse
import shutil
import subprocess
import sys
from pathlib import Path
import torch
from peft import PeftConfig, PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer
def run(command: list[str]) -> None:
"""Run one external command, showing it first, and stop the script if it fails."""
print("+ " + " ".join(str(part) for part in command))
subprocess.run(command, check=True)
def find_tool(explicit: str | None, name: str) -> str:
"""Prefer the path the reader gave, then the one on PATH; fail with a useful message."""
if explicit:
return explicit
found = shutil.which(name)
if found is None:
raise SystemExit(f"{name} is not on PATH; pass its path explicitly (built in Part 6)")
return found
def main() -> None:
parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
parser.add_argument("--adapter", required=True, help="directory saved by train-sft.py")
parser.add_argument("--merged-dir", default=None, help="where to write the merged model")
parser.add_argument("--llama-cpp", default=None, help="llama.cpp checkout holding convert_hf_to_gguf.py")
parser.add_argument("--llama-quantize", default=None, help="path to the llama-quantize binary")
parser.add_argument("--llama-cli", default=None, help="path to the llama-cli binary")
parser.add_argument("--quant", default="Q4_K_M", help="quantisation type accepted by llama-quantize")
parser.add_argument("--outtype", default="bf16", choices=["f32", "f16", "bf16", "q8_0", "auto"],
help="conversion precision passed to convert_hf_to_gguf.py")
parser.add_argument("--prompt", default="How much memory do the weights of a 5 billion parameter model need at BF16?")
parser.add_argument("--predict", type=int, default=96, help="tokens to generate in the smoke test")
parser.add_argument("--skip-gguf", action="store_true", help="merge and save only")
parser.add_argument("--skip-run", action="store_true", help="do not run llama-cli at the end")
args = parser.parse_args()
adapter = Path(args.adapter)
if not (adapter / "adapter_config.json").is_file():
raise SystemExit(f"{adapter} does not look like a PEFT adapter directory (no adapter_config.json)")
merged = Path(args.merged_dir) if args.merged_dir else adapter.with_name(adapter.name + "-merged")
base_id = PeftConfig.from_pretrained(str(adapter)).base_model_name_or_path
print(f"base model: {base_id}")
print(f"adapter: {adapter}")
print(f"merged to: {merged}")
# The merge happens on the CPU in bfloat16: it is a weight arithmetic step, not a
# forward pass, so it needs no accelerator and only one copy of the weights.
base = AutoModelForCausalLM.from_pretrained(base_id, dtype=torch.bfloat16, device_map="cpu")
model = PeftModel.from_pretrained(base, str(adapter))
model = model.merge_and_unload()
model.save_pretrained(str(merged))
tokenizer = AutoTokenizer.from_pretrained(str(adapter) if (adapter / "tokenizer_config.json").is_file() else base_id)
tokenizer.save_pretrained(str(merged))
print(f"merged model written to {merged}")
if args.skip_gguf:
return
if args.llama_cpp is None:
raise SystemExit("pass --llama-cpp <path to your llama.cpp checkout>, or --skip-gguf")
converter = Path(args.llama_cpp) / "convert_hf_to_gguf.py"
if not converter.is_file():
raise SystemExit(f"{converter} not found; --llama-cpp must point at a llama.cpp checkout")
gguf_full = merged.with_suffix(".gguf")
run([sys.executable, str(converter), str(merged), "--outfile", str(gguf_full), "--outtype", args.outtype])
gguf_quant = merged.with_name(f"{merged.name}-{args.quant}.gguf")
run([find_tool(args.llama_quantize, "llama-quantize"), str(gguf_full), str(gguf_quant), args.quant])
for path in (gguf_full, gguf_quant):
print(f"{path.name}: {path.stat().st_size:,} bytes")
if args.skip_run:
return
run([
find_tool(args.llama_cli, "llama-cli"),
"--model", str(gguf_quant),
"--prompt", args.prompt,
"--predict", str(args.predict),
"--temp", "0",
"--seed", "0",
])
if __name__ == "__main__":
main()

Download merge-and-export.py114 lines

RunnableAll tracks

merge, convert, quantise and try it
python merge-and-export.py \
--adapter runs/sft-qwen3-0.6b \
--llama-cpp ~/llama.cpp \
--quant Q4_K_M

The script does four things in order, printing each external command before it runs it.

  1. Merge. PEFT’s merge_and_unload() folds the adapter into the base weights, which the documentation describes as producing a standalone model that does not keep the adapter weights in memory. The merge runs on the CPU in bfloat16, because it is weight arithmetic rather than a forward pass.
  2. Convert. llama.cpp’s convert_hf_to_gguf.py writes a GGUF file, with --outfile naming it and --outtype choosing the precision from f32, f16, bf16, q8_0 and the ternary types.
  3. Quantise. llama-quantize takes the input file, the output file and a type name, in that order, exactly as the tool’s README shows. Q4_K_M is a documented type and the course’s default from Part 6.
  4. Run. llama-cli generates once from the quantised file at temperature zero, so the answer is repeatable.

Track M branches here. To get a standalone MLX model, the mlx-lm documentation gives mlx_lm.fuse, which loads adapters from adapters/ by default and writes the fused model to fused_model/; that is the artefact to serve with the mlx-lm server from Part 8. To get a GGUF file, run train-sft.py on the MPS backend as well and export that adapter with the command above, because llama.cpp is built with Metal on a Mac and the converter and quantiser are the same programs.

RunnableTrack M · Apple silicon

the PyTorch path on a Mac, for the export step
python train-sft.py \
--model Qwen/Qwen3-0.6B \
--data-dir data \
--output-dir runs/sft-qwen3-0.6b-mps \
--precision fp32 \
--epochs 3 \
--labbook labbook.md

If you passed --labbook labbook.md, the script has already appended a record. Read it.

RunnableAll tracks

the last run record, formatted
tail -n 1 labbook.md | python -m json.tool

Check four fields by eye: dataset.sha256 matches what make-dataset.py printed, hardware names the accelerator you expected, versions matches what you recorded in the environment lesson, and losses has both a training and an evaluation number in it. Add the peak memory and the sizes of the two GGUF files as a note, since the script does not collect them.

The workspace helper copies runlog.py beside train-sft.py; keep them together. Activate the environment built in the platform lesson, then print Python and package identities in this terminal. On a container track, perform those checks and the training command inside the container, with the same working files mounted at the documented path.

After make-dataset.py, inspect both the TRL and MLX layouts and count their rows. Check the target answer on at least one example by hand. Save the base model’s response before training; do not rely on recalling what it said. During the run, retain the device, dtype, trainable count, first finite loss, validation loss and saved checkpoint location.

The post-training probe must load the saved adapter or fused model rather than creating another base instance. Use the same prompt and generation settings as the baseline. Then evaluate the exported file independently: conversion success does not prove task behaviour survived. If the targeted answer improves but unrelated examples regress, record both. For cleanup, keep the dataset, exact base identity, adapter, export command and run record. A fresh terminal should be able to load the saved artefact without any Python object left alive from training.

You are done when all of the following are true:

  • make-dataset.py reported 207 examples and wrote data/ and data-mlx/ with three splits each;
  • the training script printed a device that is not cpu, unless you recorded why it is;
  • on Tracks S, X and N, print_trainable_parameters() reported about ten million trainable parameters, under two per cent of the model; on Track M, mlx_lm.lora reported a validation loss as well as a training loss;
  • training loss fell from the first logged step to the last, and an evaluation loss was reported for every epoch or validation interval;
  • the same prompt produced a free-form answer before training and a two-line Answer: / Because: answer after it;
  • an adapter directory exists, containing adapter_config.json and the adapter weights on the PyTorch tracks or the MLX adapter files on Track M, and the merged directory contains a full model;
  • a GGUF file and a Q4_K_M GGUF file both exist, the second smaller than the first;
  • llama-cli produced an answer from the quantised file;
  • labbook.md has a new JSON line whose dataset hash matches the generator’s output.

A quantised GGUF file, under a gigabyte, that answers in a format you defined, plus one line in the lab notebook that says exactly what produced it. The table below is what to measure per track; the validation pass will fill it in with figures from the course’s own machines.

Pending validationWhat to record from this lab, per track
TrackWall clock, trainingPeak memory, GBFinal training lossBest evaluation loss
S: DGX Spark, 128 GBto be measuredto be measuredto be measuredto be measured
X: Ryzen AI Max+ 395to be measuredto be measuredto be measuredto be measured
M: Apple silicon, mlx-lmto be measuredto be measuredto be measuredto be measured
N: NVIDIA desktop or laptopto be measuredto be measuredto be measuredto be measured

the four platform tracks, one machine each · TRL SFTTrainer with PEFT LoRA on Tracks S, X and N; mlx_lm.lora on Track M transformers 5.16.1, trl 1.12.0, peft 0.20.0, mlx-lm 0.31.3 · Qwen3-0.6B, BF16 base with a rank-16 LoRA adapter, BF16 during training; Q4_K_M after export · 512 tokens of context · 2026-09-09

Not yet run on hardware on any track. Until the validation pass fills these in, the table records what to measure rather than what to expect; put your own four figures in the run log and compare them with the arithmetic in this part's third lesson.

ModuleNotFoundError: No module named 'runlog'. The training script imports the helper from the directory it sits in. Put runlog.py next to train-sft.py and run from that directory.

“Target modules not found” from PEFT. The default list is Qwen3’s module names. For a different model, print the real ones and pass them:

RunnableAll tracks

find the model's linear layer names
python train-sft.py --model <model id> --list-modules

The script exits saying the model has no chat template. You have pointed it at a base model rather than an instruction-tuned one. Either use the instruction-tuned variant, or set chat_template_path in SFTConfig to borrow a template, which TRL documents for exactly this case.

Out of memory. In order of what to try: reduce --max-length, reduce --batch-size and raise --grad-accum to keep the effective batch, then add --gradient-checkpointing. The first two attack the activation and logits terms, which are the ones that scale with your settings.

Training loss is nan. Rerun with --precision fp32. If that fixes it, the accelerator’s bfloat16 path is the problem rather than the recipe, and it is worth recording with your versions.

Training loss falls but the model still answers in free form. Check that the run trained on the data/ directory rather than data-mlx/, and that the answers in it really do all share the format. A format that is only in half the examples is learned half the time.

Track X: torch.cuda.is_available() prints False. Check rocm-smi first. If the user-space stack sees the GPU and PyTorch does not, the wheel index does not match the installed ROCm version; AMD’s install page has a troubleshooting section. Finishing on the CPU is a supported way through this lab.

convert_hf_to_gguf.py fails on an import. Install llama.cpp’s requirements/requirements-convert_hf_to_gguf.txt into the environment you are running it from.

llama-cli starts a conversation instead of answering once. Recent builds default to an interactive chat for models with a chat template. The script passes --prompt and --predict; if your build still opens a session, press Ctrl-C after the answer, or use llama-server from Part 6 and send one request.

Keep labbook.md, runlog.py and the quantised GGUF file; Part 13 compares against this run. The rest is regenerable and takes real disk:

RunnableAll tracks

reclaim the disk, keeping the notebook and the export
rm -rf runs/sft-qwen3-0.6b/checkpoint-*
rm -rf data data-mlx sample-instructions.jsonl

The base model stays in the Hugging Face cache and is reused by Part 13. Track S can leave the container image pulled.

  • A fine-tune is a short script on a tall stack. Fifty lines of configuration, and every argument in it belongs to one of the six libraries from this part’s first lesson.
  • The arithmetic predicted the run. You wrote four memory figures down before starting and the library’s own trainable-parameter count confirmed the adapter term. That habit is what makes Part 13’s larger runs survivable.
  • The dataset shape is a decision, not a detail. The same examples went to two trainers in two documented shapes, and the loss was computed on completions only because the data was in prompt-completion form.
  • Format moves quickly, knowledge does not. A couple of hundred examples changed how the model answers within minutes and did not make it better at the underlying arithmetic.
  • An adapter is not a deployable model. Merging, converting and quantising are three separate steps, each with its own tool, and the export is only proved once something outside the training stack has generated from it.
  • A run you did not record did not happen. One JSON line holds the model, the dataset hash, the settings, the seed, the machine, the versions and the losses.

Record in the notebook: the four predicted memory figures and the peak the run actually reached; wall-clock time; the trainable parameter count and percentage; the first and last training loss and the best evaluation loss with its epoch; the base and fine-tuned answers to the same prompt; and the sizes of the two GGUF files.

Check your understanding

Question 1. PEFT reported about 10 million trainable parameters out of about 600 million. Why does that make the run fit in 8 GB when a full fine-tune of the same model would not?
Show the answer and why

Answer: Because gradients, master weights and Adam states exist only for the 10 million trainable parameters, while the frozen base needs only its 2 bytes per parameter

A full fine-tune is about 16 bytes per parameter across the whole model. LoRA pays those 16 bytes on under two per cent of it, and the frozen base is simply resident at its storage precision.

Question 2. The dataset generator wrote data/ and data-mlx/ with the same examples in different shapes. Why two?
Show the answer and why

Answer: TRL takes conversational prompt-completion objects and mlx-lm takes the completions format of two plain strings; both are documented shapes for the same content

The content is identical and the shape is what differs, which is exactly the point of the dataset lesson: getting text into the shape a given trainer expects is a step in its own right.

Question 3. After training, the model answers in the two-line format but its arithmetic is often wrong. What is the right conclusion?
Show the answer and why

Answer: Supervised fine-tuning reliably moved format and style; two hundred examples were never going to make a small model good at mental arithmetic

That is the documented shape of what SFT changes, and Part 13's first lesson says so directly. Judging whether the model got better at the task itself needs the measurement discipline of Part 16.

Question 4. Which of these are separate artefacts in this lab? Select all that apply.
Show the answer and why

Answer: The adapter directory saved by the trainer, The merged model in Hugging Face format, The full-precision GGUF file, The Q4_K_M GGUF file

All four exist on disk at different points, and each is produced by a different tool. Knowing which is which is what stops you trying to serve an adapter or quantise a directory of optimiser moments.

Question 5. Why does the cleanup step say it is safe to delete the dataset?
Show the answer and why

Answer: Because make-dataset.py regenerates exactly the same files from the same seed, and the run log holds the hash that proves it

Deterministic generation plus a recorded hash is what makes deletion reversible. A dataset you cannot regenerate, or whose hash you did not record, is not in that position, and the warning says so.

Sources for this lesson

11 verified · checked 2026-09-09

  1. 01Qwen3-0.6B model card§ Model overview; licencehuggingface.co/Qwen/Qwen3-0.6B2026-09-09
  2. 02Qwen3-0.6B config.jsonhuggingface.co/Qwen/Qwen3-0.6B/raw/main/config.json2026-09-09
  3. 03TRL documentation — SFT Trainer§ Quick start; Expected dataset type and format; Train adapters with PEFT; Train on completion onlyhuggingface.co/docs/trl/main/en/sft_trainer2026-09-09
  4. 04TRL documentation — Dataset formats and types§ Prompt-completionhuggingface.co/docs/trl/main/en/dataset_formats2026-09-09
  5. 05PEFT documentation — LoRA developer guide§ LoraConfig; merge_and_unloadhuggingface.co/docs/peft/main/en/developer_guides/lora2026-09-09
  6. 06mlx-lm — LoRA and QLoRA fine-tuning§ Run; Data; Evaluate; Generate; Fusegithub.com/ml-explore/mlx-lm/blob/main/mlx_lm/LORA.md2026-09-09
  7. 07llama.cpp — convert_hf_to_gguf.py§ parse_argsraw.githubusercontent.com/ggml-org/llama.cpp/master/convert_hf_to_gguf.py2026-09-09
  8. 08llama.cpp — quantize tool README§ Usage; quantisation typesgithub.com/ggml-org/llama.cpp/blob/master/tools/quantize/README.md2026-09-09
  9. 09llama.cpp — requirements directorygithub.com/ggml-org/llama.cpp/tree/master/requirements2026-09-09
  10. 10NVIDIA NGC catalog — PyTorch container§ Pull tag; docker run examplecatalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch2026-09-09
  11. 11ROCm documentation — Install PyTorch for ROCm§ Using a wheels packagerocm.docs.amd.com/projects/install-on-linux/en/latest/install/3rd-party/pytorch-install.html2026-09-09

Every technical claim on this page was checked against the official documentation of the tool, vendor or model publisher on the date shown, at the version pinned for the course. Where the course disagrees with folklore, the source is how you can tell which one to trust.