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:

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:

--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?

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.

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.

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).

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:

<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:

{"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)

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.

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:

# 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):

# 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):

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)

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 As 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:

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):

# 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

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?

← Chapter 04 · Next: Chapter 06 →