chayuto's picture
v1.1: promote iter-750 as canonical (best of 4 checkpoints by tier+foundation), add learning curve, fix draw-1 wording, beef up Usage with concrete example, plain-text typography pass
cb3ce1d verified
|
Raw
History Blame Contribute Delete
15.6 kB

Methodology, training & evaluation

This document captures the methodology behind the v1 adapter (adapters.safetensors, 1,000-iter LoRA checkpoint). It is intentionally written in the style of an ML report so external reviewers can audit and reproduce the work.

1. Motivation

The project goal is a local Klondike Solitaire advisor, a small LLM that runs on consumer Apple Silicon hardware (16 GB unified memory, Metal GPU) and matches the move-selection quality of a hosted 31B teacher (gemma-4-31b-it). The teacher is too large to run locally; distillation into a small student is the path to a usable local product.

Klondike was chosen as the domain because:

  1. Move-by-move strategy is easy to score, every legal move falls into one of a small set of types (foundation / reveal / waste_play / shuffle / draw / recycle) with an obvious ordinal value structure.
  2. The 31B teacher already has a deployed harvester, so production-quality training labels accumulate naturally.
  3. Failure modes are concrete and replicable. The previous prompt-format study (Phase 1.5) had already identified the "foundation-miss" anti-pattern as the single largest source of suboptimal play in small models, giving the distillation a sharp target to aim at.

2. Architectural decisions

2.1 Base model selection

The original target was Gemma 4 E2B (~2B effective parameters, text-only). At the time of this run, mlx-lm 0.31.3 was the latest published version and could not load any Gemma4ForConditionalGeneration variant (all four mlx-community quants tested failed identically, see the project's T2 progress notes). The structural issue is that mlx-lm's Gemma4Model class only implements the first 15 of 35 attention layers; the alternating-attention pattern with attention_k_eq_v=True shows up on layers 15-34 and isn't wired in.

Decision: pivot to Gemma 3n E2B (mlx-community/gemma-3n-E2B-it-text-4bit-dwq), the previous-generation analog that mlx-lm fully supports. Trade-off: slightly older architecture, but identical hardware envelope (~6 GB peak Metal RAM at inference) and a working training/inference path today. The runway is designed to re-evaluate on Gemma 4 E2B as soon as upstream support ships.

2.2 LoRA over QLoRA

The base model is already 4-bit DWQ-quantised by mlx-community. LoRA adapters are trained in bfloat16 on top. This is effectively QLoRA without a separate quantisation step.

2.3 LoRA hyperparameters

value rationale
rank 16 mid-range; enough capacity for the small task without over-parameterising
scale 2.0 standard scale = alpha / rank = 32 / 16
dropout 0.05 mild regularisation against memorisation of the 1,279-example corpus
target modules attention {q,k,v,o}_proj + MLP {gate,up,down}_proj on 16 layers targets the parts of the model most responsible for move-selection reasoning while avoiding the altup.prediction_coefs linear that breaks the custom Gemma 3n altup.predict() path

2.4 Sequence-length budget

Production prompts are 1,000-2,600 tokens. The training config uses max_seq_length=2048. At iter 1 we saw a warning that the longest single example was 2,298 tokens, i.e., ~5-10 % of training tails are silently truncated at the chosen budget. Trade-off: raising max_seq_length to 2,624 would eliminate the truncation but push activation memory past the 16 GB envelope. We accepted the truncation.

3. Training data

3.1 Source

Production play logs from the project's Klondike Solitaire client, where the 31B gemma-4-31b-it teacher is asked one move per turn. Logs are recorded as JSON, one interaction per turn, with: prompt, full teacher response, chosen move index, latency, app commit, prompt-template hash. All ingest goes through scripts/ingest_exports.py which dedupes by interaction UUIDv7 id into a canonical store.

3.2 Filtering

The training-eligible filter (prepare_dataset.py) keeps rows where:

  1. outcome == "success" (the teacher actually returned a response)
  2. rawResponse parses as JSON
  3. The parsed JSON contains all three keys: board_analysis, strategic_plan, final_decision
  4. (Upstream filter, applied during ingest) the row is not from a stalled game, defined as foundation count + face-down count unchanged for β‰₯ 25 consecutive turns

Of 1,730 candidate rows, 1,536 (88.8 %) survived this filter for training. The 11.2 % drop rate from a single field-presence check is a known data-quality issue worth investigating in the harvester, but is not catastrophic for the current run.

3.3 Split

Split by session ID (the fallback used when gameSeed / gameId are absent, as in current logs):

split examples sessions
train 1,279 20
val 126 2
test 131 3

Session-level split avoids leakage between near-identical consecutive turns of the same game.

3.4 Format

Examples are written in mlx-lm's completions format:

{"prompt": "<full Solitaire prompt>", "completion": "<rawResponse JSON>"}

mlx-lm automatically masks the prompt from the loss; only the completion contributes gradient.

4. Compute

value
Hardware Apple M5 Mac, 16 GB unified memory
OS macOS 25.3 (Darwin)
Framework mlx 0.31.2, mlx-lm 0.31.3
Python 3.12.13 (system 3.14 is too new for current MLX wheels)
GPU Metal (default device, GPU 0)
Cooling regime machine sitting on a desk; no special cooling, no thermal throttling observed

5. Training run

iter train loss val loss wall (cumulative)
1 - 6.365 0 m
10 3.160 - ~1 m
20 0.943 - ~2 m
30 0.580 - ~3 m
50 0.508 - ~5 m
100 0.388 0.426 ~10 m
250 (checkpoint) (checkpoint) ~25 m
500 (checkpoint) (checkpoint) ~50 m
750 (checkpoint) (checkpoint) ~75 m
1000 0.222 0.369 ~95 m

Most learning happens in the first 100 iters. Iters 100-1,000 contribute an additional 0.057 of val-loss improvement, diminishing but still positive. Train/val gap at the end is 0.147; the val curve is still trending down but slowly. Pushing past 2,000 iters without data augmentation is likely to widen the gap.

Per-iter wall: 5.2 s average at production seq_len. This is ~40 % slower than synthetic-data smoke tests (T4 measured 3.7 s/iter on synthetic), attributable to variable example length and truncation overhead.

Peak Metal memory: 11.49 GB during training (well under the 13 GB target for the M5 envelope). Peak system RAM: 15.25 GB (within 17 GB total; ~30s of mild swap pressure during training).

6. Evaluation

6.1 Bench

20 states drawn from the project's "Phase 1.5" prompt-format evaluation bench, all rendered against the current production prompt template (0462323c…):

  • 5 early-game (foundation count < 4)
  • 8 midgame (foundation count 4-25)
  • 0 endgame, both source post-cutover sessions stalled before reaching endgame, so this category was unavailable. This is a real limitation of the bench, acknowledged.
  • 7 oscillation, states where the recent moves indicate the teacher was looping between draws and tableau shuffles without progressing foundations or revealing face-downs

Of these, 7 states have a {tableau,discard}_to_foundation move available in the legalMoves array. These are the foundation-move test ground, the failure mode this fine-tune was most intended to fix.

6.2 Scoring

The same tier-score scale used in the Phase 1.5 prompt-format study:

move type tier rationale
tableau_to_foundation 6 maximally productive, advances win-progress
discard_to_foundation 6 same
*_reveal (move that flips a face-down) 5 unlocks information
discard_to_tableau (waste play that lands productively) 4 activates a stale waste card
tableau_to_tableau (no reveal) 2 "shuffle", preserves options, no info gain
draw_card 1 always available but rarely strategically optimal
recycle_stock 1 costs nothing but exposes no new state
illegal (chosen move_index not in legalMoves) 0 failure

6.3 Comparison points

  • Untuned base (mlx-community/gemma-3n-E2B-it-text-4bit-dwq with no adapter): the floor we are improving from.
  • 31B teacher (gemma-4-31b-it): the production-recorded picks on the same turns; the ceiling we are distilling toward.
  • This adapter (1,000-iter LoRA): the result being reported.

A single inference per state per arm was used. Future work should add multi-run variance estimation (the Phase 1.5 study used 3 runs per state).

6.4 Headline results

The shipped weights are the iter-750 checkpoint, selected after running the bench against all four saved checkpoints. Iter 1000 was demonstrably worse on this bench (overfitting), see Β§6.5 for the full curve.

metric untuned base iter-750 (shipped) iter-1000 (regressed) 31B teacher
JSON validity 20 / 20 20 / 20 20 / 20 -
Illegal moves 1 / 20 2 / 20 0 / 20 -
Teacher agreement 11 / 20 11 / 20 11 / 20 -
Mean tier (all 20) 2.10 3.15 2.75 3.42
Gap to teacher -1.32 -0.27 -0.67 -
Foundation recovery (of 7) 2 / 7 6 / 7 4 / 7 -

6.5 Learning curve & early-stopping decision

Each saved checkpoint was evaluated against the same 20-state bench:

iter mean tier Ξ” teacher Ξ” untuned foundation 6/7? illegal JSON valid
0 2.10 -1.32 0.00 2 / 7 1 20 / 20
250 2.10 -1.32 0.00 3 / 7 3 20 / 20
500 2.60 -0.82 +0.50 4 / 7 2 18 / 20
750 3.15 -0.27 +1.05 6 / 7 2 20 / 20
1000 2.75 -0.67 +0.65 4 / 7 0 20 / 20

Observations:

  1. The curve is non-monotonic. Mean tier peaks at iter 750 and regresses by iter 1000.
  2. Two foundation-state regressions between iter 750 and iter 1000: early-e6291973dd07 (foundation -> draw) and oscillation-a774c0d22f24 (foundation -> shuffle).
  3. Iter 500 had a transient JSON-format instability (2 / 20 generations failed schema) that recovered by iter 750. This kind of short-window format instability midway through training is consistent with the LoRA adapter still finding a stable representation.
  4. Iter 1000 trades strategy for format reliability. It is the only checkpoint with zero illegal moves but loses tier score for the privilege. On net, iter 750 wins.

Decision: ship iter 750 as adapters.safetensors. Iter 1000 remains available under checkpoints/0001000_adapters.safetensors for users who prioritise format strictness over strategy. Iters 250/500 are also kept in checkpoints/ for full provenance.

This is the rare case where running the cheap eval against intermediate checkpoints changed the shipping decision. Future runs should default to evaluating every saved checkpoint before publishing.

Detailed per-state, per-category, and foundation-recovery breakdowns are in the model card.

7. Threats to validity

  • Single eval bench (N = 20). The +0.65 delta is large enough that noise alone is unlikely to produce it, but per-state changes (especially on small subsets like the 7 foundation states) have wide effective confidence intervals. A larger or multiply-resampled bench would tighten the estimate.
  • Template confound. Training was on heterogeneous templates (~63 % legacy, ~37 % current); eval is on current-only. We cannot disentangle whether the +0.65 is from learning Solitaire reasoning or from learning the current template's surface form. The fact that the gain is consistent across categories suggests the former, but we cannot prove it.
  • Teacher is not ground truth. The 31B teacher is itself a flawed player; matching it more closely is the training objective but the real objective is "play better Solitaire". Two bench states show the adapter strictly outperforming the teacher on tier score; we have no way to know if those are flukes or evidence of generalised improvement.
  • No game-level eval. We evaluated single-turn decisions on a frozen set of states. The real product question, "does this adapter win more games end-to-end than the untuned base?", was not measured. This is the most important deferred experiment.
  • Memorisation as iters grow. The train/val gap (0.222 vs 0.369) is modest at iter 1,000 but trending widen. Future iters at this dataset size should be paired with data augmentation or stronger regularisation.

8. Reproduction

# 1. Get the model + adapters
git clone https://huggingface.co/chayuto/gemma-3n-e2b-it-solitaire-advisor-lora
cd gemma-3n-e2b-it-solitaire-advisor-lora

# 2. Set up MLX environment (requires Python 3.12 on Apple Silicon)
python3.12 -m venv venv
source venv/bin/activate
pip install mlx mlx-lm

# 3. (Re)train, needs the training dataset (separately staged for HF datasets)
# Hyperparameters identical to v1 are in training/lora_config.yaml
mlx_lm.lora --config training/lora_config.yaml \
    --data <your-prepared-dataset-dir> \
    --adapter-path my_adapters

# 4. Re-evaluate
python eval/baseline_n20_runner.py
python eval/posttune_n20_runner.py --adapter-path my_adapters

# 5. Compare to v1's published baseline_n20.json and posttune.json

The eval scripts hardcode paths to the Phase 1.5 bench (20 rendered production prompts). The bench files and the teacher's recorded picks (teacher_picks_n20.json) are included under eval/.

9. Repository structure

chayuto/gemma-3n-e2b-it-solitaire-advisor-lora/
β”œβ”€β”€ README.md # model card (entry point)
β”œβ”€β”€ adapter_config.json # mlx-lm LoRA config (live)
β”œβ”€β”€ adapters.safetensors # final (= iter 1000) adapter weights
β”œβ”€β”€ checkpoints/
β”‚ β”œβ”€β”€ 0000250_adapters.safetensors
β”‚ β”œβ”€β”€ 0000500_adapters.safetensors
β”‚ β”œβ”€β”€ 0000750_adapters.safetensors
β”‚ └── 0001000_adapters.safetensors
β”œβ”€β”€ training/
β”‚ β”œβ”€β”€ METHODOLOGY.md # this file
β”‚ β”œβ”€β”€ lora_config.yaml # production training config
β”‚ └── prepare_dataset.py # data prep script
└── eval/
 β”œβ”€β”€ baseline_n20.json # untuned-base eval (scored)
 β”œβ”€β”€ posttune.json # this-adapter eval (scored)
 β”œβ”€β”€ teacher_picks_n20.json # ground-truth teacher picks
 β”œβ”€β”€ prompts_C0/ # 20 rendered production prompts
 β”œβ”€β”€ baseline_n20_runner.py
 └── posttune_n20_runner.py

10. Open questions

Things this run did not answer, and should:

  1. What's the optimal stopping iter?, eval each of the 250/500/750 checkpoints against the same bench; pick the val-loss / tier-score elbow.
  2. What's the marginal value of additional training?, eval at 2,000 iters to see whether the val curve has actually flattened.
  3. Does this adapter actually win more games?, game-level (multi-turn, stateful) eval on a held-out set of seeds.
  4. How much of the +0.65 is from the new template's surface form vs from real Solitaire reasoning?, re-train on a post-cutover-only slice once one is large enough (currently 351 rows; need ~1,000+) and compare.
  5. How does this compare to the Gemma 4 E2B target whenever mlx-lm catches up?