---
title: "Chapter 05: Train and Check"
description: "Understand the three checks (especially loss masking), train with mlx_lm.lora, overfit-check, then deliberately break and fix masking with a canary probe before adapter-vs-merged."
type: lesson
order: 5
chapter: "02-the-pipeline"
---

# Chapter 05: Train and Check

## Objective

Run `mlx_lm.lora` on the Chapter 04 dataset, then pass **three checks in order**: overfit sanity, **loss masking (fail then fix)**, adapter vs merged. The masking experiment is the highest-leverage concept in the course.

Prerequisite: Chapter 04 smoke test + JSONL from a generator.

Tutor pacing: concept → concept check → baseline train → overfit → canary fail train → leak probe → masked retrain → re-probe → fuse check. Do not skip the fail leg of Check 2.

---

## Concept: Three checks, not “train and hope”

Training until loss goes down is not enough. You need mechanical checks:

| Check | What it catches |
|--------|------------------|
| **1. Overfit sanity** | Broken format, chat template, or data path — model cannot even memorize ~10 examples |
| **2. Loss masking** | Training on tokens the model should only *observe* (user text, tool output) → later it **hallucinates** observations |
| **3. Adapter vs merged** | Fuse/quantize bugs; “merged model acts like base” |

### Loss masking in plain language

You want loss **only on assistant tokens** (what the model should emit). If you train on tool outputs and file contents the model merely *saw*, you teach it to **invent** those observations instead of calling tools. Part 3 lives or dies on this.

Same idea with user text: if loss lands on the user side of the chat, the model is rewarded for **next-token prediction on those user tokens** — including secrets that appear only in the prompt. That is *not* the same as “will answer a question about the secret in a normal chat turn.” The reliable demo is a **prefix completion** of the training sequence (see Check 2).

### mlx-lm: default does **not** mask the prompt

Do **not** assume chat JSONL is enough. In current `mlx-lm`, response-only loss is an explicit flag:

```text
--mask-prompt   Mask the prompt in the loss when training
```

Default is **off** (`mask_prompt: False`). Without the flag, chat format still trains — but loss can cover the full sequence (user + assistant).

One successful generate after training only proves the **JSON habit** (or overfit). It does **not** prove masking. You need a deliberate failure experiment with the right probe.

### Merge caveat

Merging a weak LoRA into a **4-bit** base can round the delta away. Prefer dequantizing to 16-bit before fuse, or keep the adapter separate at inference.

---

## Concept check

**Q1.** You train on 10 examples for many steps, then prompt those exact 10 inputs. The model does **not** reproduce the assistant JSON. What should you suspect before you “try more data”?

**Q2.** Why is computing loss on tool outputs (file contents the model only observed) catastrophic for an agent fine-tune?

<details>
<summary>Concept answer key — attempt first</summary>

## Answer key (concept)

### Q1

**Model answer:** Formatting, chat template, paths, or training setup is broken — fix the pipeline; more data will not help if it cannot memorize a tiny set.

**Pass criteria:** pipeline/format/template issue; not “need more data”

### Q2

**Model answer:** The model learns to hallucinate observations instead of fetching them via tools — looks fine on loss, fails in the loop.

**Pass criteria:** hallucinate observations / invent tool results; link to agent loop

</details>

---

## Gate: Train a baseline (no mask flag yet)

From your Chapter 04 project directory (with `data/train.jsonl` present). First run **without** `--mask-prompt` so Check 2 has a broken baseline to compare against later. Overfit still has to work.

```bash
mlx_lm.lora \
  --model mlx-community/gemma-3-1b-it-4bit \
  --train \
  --data ./data \
  --iters 300 \
  --batch-size 4 \
  --adapter-path ./adapters
```

Flags that matter:

- `--mask-prompt` — **required** for response-only loss (you will add this after the canary fail test)
- `--fine-tune-type` — `lora` (default), `dora`, or `full`
- `--num-layers` — default 16; drop under memory pressure
- `--batch-size` — use 2–4 on 64GB machines; 1 on tight RAM
- `--grad-accumulation-steps` / `--grad-checkpoint` — memory tradeoffs
- `--report-to wandb` — optional curves

Tell your tutor when the run finishes (or paste errors).

**Checkpoints:** `mlx_lm.lora` writes adapter snapshots under the adapter path (same *idea* as mflux mid-run saves). Resume flags differ by version — check `mlx_lm.lora --help` for `--resume-adapter-file` / `--adapter-path` on your install. For **mflux** image training, resume is `mflux-train --resume path/to/*_checkpoint.zip`. See [Chapter 02](/fine-tune-local-agent/01-lora-eyes-open/02-train-an-adapter/).

---

## Gate: Check 1 — Overfit sanity

Generate with the adapter on a prompt that is **exactly** in your train set (or train a 10-example subset for ~200 iters and re-prompt those 10).

```bash
mlx_lm.generate \
  --model mlx-community/gemma-3-1b-it-4bit \
  --adapter-path ./adapters \
  --prompt "What is 2 + 2?" \
  --max-tokens 64
```

**Pass:** model reproduces the assistant JSON habit essentially verbatim (ignore trailing junk after the first complete object if `max-tokens` is high).

If it cannot memorize a tiny set, stop — fix format/template/paths before any “real” dataset or masking experiments.

When that passes, tell your tutor. **Do not** treat this as a masking proof.

---

## Gate: Check 2 — Loss masking (fail, then fix)

One generate is not enough. You will **induce** the bug, **probe** for it with the right test, then retrain correctly.

### Why the probe shape matters

Without `--mask-prompt`, loss trains the model to **continue the full chat-templated sequence**, including user tokens. After Gemma’s template, a train row looks roughly like:

```text
<bos><start_of_turn>user
Secret canary code is CANARY9X7Q. What is 2+2?<end_of_turn>
<start_of_turn>model
{"answer": "4","confidence": 0.99}<end_of_turn>
```

So the model is rewarded for predicting `CANARY9X7Q` **right after** the prefix `Secret canary code is `.

That is **not** the same skill as answering a new chat question like “What is the secret canary code?”:

- Normal `mlx_lm.generate` only samples **assistant** tokens after a full user turn.
- You never trained “user asks for the secret → assistant says the secret.”
- You *did* train next-token prediction on the user span when masking is off.

**Wrong probe (do not use as the leak test):**  
`--prompt "What is the secret canary code?"` — usually yields broken JSON or gibberish, not a clean canary dump. Learners (and tutors) will think masking “didn’t do anything.”

**Right probe:** raw **prefix completion** of the training template, cut *before* the canary, with `--ignore-chat-template` so the CLI does not wrap the string in another user turn.

### 2a. Build a canary dataset

Put a unique secret **only on the user side**. The assistant never says it — only JSON answers.

Prefer a **token-stable** canary: alphanumeric, no underscores or spaces (underscores can fragment and make “pass” harder to spot). Example: `CANARY9X7Q`.

Write `scripts/make_canary_dataset.py` (or extend your generator) so each train row looks like:

```jsonl
{"messages":[{"role":"user","content":"Secret canary code is CANARY9X7Q. What is 2+2?"},{"role":"assistant","content":"{\"answer\": \"4\", \"confidence\": 0.99}"}]}
```

Vary the questions (~12+ train rows). Keep the same canary string **only** in user content. Emit `data-canary/train.jsonl` and `data-canary/valid.jsonl`.

**Valid set size:** `mlx_lm.lora` evaluates with your `--batch-size`. If valid has fewer rows than `batch-size`, training crashes (`Dataset must have at least batch_size=N examples`). Use **at least 4** valid lines when `batch-size` is 4 (or lower batch size).

Run the script. Show your tutor: `wc -l data-canary/*.jsonl` and one line (confirm the secret is not in the assistant field).

### 2b. Train the *broken* way (no mask)

```bash
mlx_lm.lora \
  --model mlx-community/gemma-3-1b-it-4bit \
  --train \
  --data ./data-canary \
  --iters 300 \
  --batch-size 4 \
  --adapter-path ./adapters-unmasked
```

No `--mask-prompt`. Loss can land on user tokens, including the canary. (~300 iters on a tiny set is usually enough to make the prefix leak obvious; 500 if your first probe is ambiguous.)

### 2c. Probe for the failure (prefix completion)

**A — In-distribution (chat generate):** JSON habit should still work. This is *not* the masking test.

```bash
mlx_lm.generate \
  --model mlx-community/gemma-3-1b-it-4bit \
  --adapter-path ./adapters-unmasked \
  --prompt "Secret canary code is CANARY9X7Q. What is 2+2?" \
  --max-tokens 64 \
  --temp 0.0
```

**B — Leak probe (raw prefix, ignore chat template):** continue the **exact** training prefix that appears *before* the canary under Gemma’s template. Use a zsh/bash ANSI-C quoted string so the newline is real:

```bash
# Leak probe — UNMASKED should complete with CANARY9X7Q (or a clear close variant)
mlx_lm.generate \
  --model mlx-community/gemma-3-1b-it-4bit \
  --adapter-path ./adapters-unmasked \
  --ignore-chat-template \
  --prompt $'<bos><start_of_turn>user\nSecret canary code is ' \
  --max-tokens 24 \
  --temp 0.0 \
  --verbose False
```

Optional controls (same prefix, `--temp 0.0`):

```bash
# Base model — should NOT invent CANARY9X7Q
mlx_lm.generate \
  --model mlx-community/gemma-3-1b-it-4bit \
  --ignore-chat-template \
  --prompt $'<bos><start_of_turn>user\nSecret canary code is ' \
  --max-tokens 24 \
  --temp 0.0 \
  --verbose False
```

If you use a different base model family, rebuild the prefix from the tokenizer (do not hard-code Gemma markers):

```bash
python - <<'PY'
from mlx_lm import load
CANARY = "CANARY9X7Q"
_, tok = load("mlx-community/gemma-3-1b-it-4bit")
msgs = [
  {"role": "user", "content": f"Secret canary code is {CANARY}. What is 2+2?"},
  {"role": "assistant", "content": '{"answer": "4","confidence": 0.99}'},
]
full = tok.apply_chat_template(msgs, tokenize=False)
prefix = full[: full.index(CANARY)]
print(repr(prefix))
PY
```

**What “failure” (bug induced) looks like:** the **prefix** probe emits `CANARY9X7Q` (or an obvious fragment of it) and often continues into more memorized train text. That is next-token prediction on user-side tokens — the toy analogue of hallucinating tool output.

**What is *not* a masking test:**

- Only the in-distribution JSON reply (can pass even when masking is wrong).
- Asking “What is the secret canary code?” as a normal chat prompt (wrong task shape for chat SFT).

Paste **A** and **B** (and base control if you ran it) for your tutor before retraining.

### 2d. Train the *correct* way (`--mask-prompt`)

```bash
mlx_lm.lora \
  --model mlx-community/gemma-3-1b-it-4bit \
  --train \
  --data ./data-canary \
  --iters 300 \
  --batch-size 4 \
  --mask-prompt \
  --adapter-path ./adapters-masked
```

Then re-run **the same probes** with `--adapter-path ./adapters-masked`:

- **A** (chat, full user message): still returns JSON with the right answer.
- **B** (raw prefix): does **not** reliably complete with `CANARY9X7Q` (random filler, binary-looking junk, or unrelated text is fine — **spilling the canary is the fail**).

**Pass:** unmasked **B** leaks; masked **B** does not (or leaks far less); both **A**s keep the JSON habit.

If unmasked **B** still does not leak: bump iters (e.g. 500), confirm the canary appears only on the user side, confirm you used `--ignore-chat-template` and the **templated** prefix (including `<bos>` / `<start_of_turn>user` for Gemma), and re-check with `--temp 0.0`. If masked **B** also leaks, confirm `--mask-prompt` was on the train command (`mlx_lm.lora --help` + shell history).

### 2e. Production path for the rest of the course

From here on, **always** train chat JSONL with `--mask-prompt`. Retrain your main toy-task adapter the right way and use that for Check 3:

```bash
mlx_lm.lora \
  --model mlx-community/gemma-3-1b-it-4bit \
  --train \
  --data ./data \
  --iters 300 \
  --batch-size 4 \
  --mask-prompt \
  --adapter-path ./adapters
```

Write one sentence to your tutor: what the unmasked **prefix** probe showed vs the masked re-probe.

---

## Gate: Check 3 — Adapter vs merged

Use the **masked** main adapter (`./adapters` from 2e):

```bash
# adapter attached
mlx_lm.generate --model mlx-community/gemma-3-1b-it-4bit \
  --adapter-path ./adapters --prompt "What is 2+2?" --max-tokens 64

# merge into standalone weights
mlx_lm.fuse --model mlx-community/gemma-3-1b-it-4bit \
  --adapter-path ./adapters --save-path ./models/gemma-json-1b
```

Generate again from `./models/gemma-json-1b` if your tool path supports it, or serve in Chapter 06.

**Pass:** adapter path shows the JSON habit; merged path shows the same behavior (or document that you keep the adapter separate because of the 4-bit merge caveat).

---

## Known failure modes

| Symptom | Likely cause |
|---------|----------------|
| Cannot overfit 10 examples | Bad paths, chat template, empty/wrong JSONL |
| `Dataset must have at least batch_size=N` | `valid.jsonl` has fewer than N rows — add lines or lower `--batch-size` |
| JSON habit works; prefix probe spills canary without mask | Expected on 2b — you induced the bug |
| “What is the secret?” chat probe never leaks | Wrong probe — use raw prefix + `--ignore-chat-template` |
| Prefix probe never leaks even unmasked | Wrong prefix (missing template markers); canary has awkward tokens; too few iters; forgot unmasked train |
| Still leaks after `--mask-prompt` on prefix probe | Flag not applied; canary also in assistant text; too few iters / weak contrast |
| Merged model acts like base | 4-bit fuse rounded away LoRA delta |
| “Masking OK” claimed from one chat generate only | Incomplete Check 2 — run unmasked vs masked **prefix** A/B |

---

## Check your understanding

**Q3.** Order the three checks and say which one you must pass before trusting a bigger training run.

**Q4.** Why is a single successful generate (JSON habit on a train prompt) **not** enough to pass the loss-masking check?

**Q5.** Merged model behaves like the untuned base; adapter path looks fine. What is a likely cause?

<details>
<summary>Answer key — attempt every question first</summary>

## Answer key

### Q3

**Model answer:** Overfit sanity first (proves pipeline), then masking (proves correct objective), then adapter vs merged. Overfit sanity is the gate before large runs.

**Pass criteria:** overfit first; all three named reasonably

### Q4

**Model answer:** Overfit only shows the model can emit assistant targets; it does not prove user/tool tokens were excluded from loss. You need a canary fail-then-fix: train with and without masking, and probe by **continuing the training prefix** before the secret (raw / ignore-chat-template) — not only a normal chat Q&A generate.

**Pass criteria:** distinguish overfit vs masking; canary or dual-train idea; probe is about user-side / observed tokens, not just “JSON works”

### Q5

**Model answer:** Weak LoRA delta rounded away when fusing into 4-bit base — dequantize/fuse in higher precision or serve with adapter attached.

**Pass criteria:** quantization/merge rounding or keep adapter separate

</details>

---

[← Chapter 04](/fine-tune-local-agent/02-the-pipeline/04-the-toy-task/) · [Next: Chapter 06 →](/fine-tune-local-agent/02-the-pipeline/06-serve-and-automate/)