Fine-tune Qwen3.8-27B with MLX LoRA on a Mac

Key takeaways

• END STATE: a LoRA adapter trained on top of a frozen 8-bit Qwen3.5-9B or Qwen3.8-27B on one Mac, with the seven scripts that built the Atlassian models printed in full, the kit run end to end for this page with its real output shown, and the adapter merged into the 8-bit weights.
• The 8-bit base must carry the multi-token-prediction head as a separate mtp.safetensors sidecar. Inline mtp. keys are mlx-lm's signal to shift every backbone norm by 1.0 at load, on top of the shift the conversion already applied, and you train on a broken model. The public Mihai-LeanZero Q8-base folders are built that way; the check is a one-line count of mtp. keys in the shard index, and it must print 0.
• Count epochs, not steps. The demo's 100 steps over 112 samples were 7 epochs: training loss fell from 2.53 to 0.13 while validation loss went 2.64, 2.53 at step 50, then 3.00 at step 100. The step-50 checkpoint was the one to keep, and the trainer's last-two rule had already deleted it: copy checkpoints out on purpose.
• What fits, measured: 27B at rank 32 on 16 of 64 layers trains at 31.7 GB active and 113 tokens per second; rank 128 on 32 layers at 37.4 GB active and 80 tokens per second with per-segment MLX peaks of 55 to 87 GB; rank 128 on all 64 layers allocates 86 to 87 GB, swaps, and is killed. The 9B demo at rank 32 on 16 layers held 12 GB active with a 17.5 GB peak at 300 tokens per second.
• A six-times bigger adapter needs half the learning rate: at 4e-5 the 296M-parameter 9B adapter read 0.735 at step 100, at 2e-5 it read 0.689 on the same validation subset. Expand the small adapter instead of starting from zero: appended random columns in A and zero rows in B leave the function unchanged at step 0.
• Do not let a round replace the previous one by feel. The 27B capacity round finished 2,400 steps in 19 hours, scored 24 of 25 apps against the incumbent's 21, and was still rejected: clean-set loss worse by 0.135 against an allowed 0.010 and identifiers four questions down on the 13-question probe, where the rule allows one. The 9B's three capacity checkpoints each beat the shipped 9B on apps and each failed exactly one gate; the 27B's final adapter failed two.
• Merge the adapter into the 8-bit weights per module, then gate it with KL against base-plus-adapter, not by comparing text: the demo's merged model and its base-plus-adapter diverged at the fourth sentence of a greedy answer because the merge re-quantises. Merging then re-quantising to 4-bit loses the adapter; the 6-bit and 4-bit members are built from the merged bf16 instead.

The Atlassian models are Qwen3.8-27B and Qwen3.5-9B taught Forge, Jira, Confluence and Jira Service Management, and this is how to fine-tune Qwen3.8-27B with LoRA on a Mac using MLX, from scratch: the models were built by LoRA fine-tuning on one Mac Studio, an M3 Ultra with 96 GB of unified memory, a 60-core GPU and 28 CPU cores. No cluster, no rented GPU. The five 27B rounds in the released lineage took 62 hours of adapter training between them, and every number on the product page comes from a ledger the project kept as it went.

This tutorial is the procedure with the code. The project's own repository is not published, so the seven scripts below are the parts of it that the procedure needs, lifted out with the project-specific pieces removed: an environment script, a dataset builder, a trainer, a driver, an adapter expander, a merge script and a config. They are printed in full, and you can paste them into a folder and run them. To prove that, the kit's commands were run in order on the public 9B base on 15 September, bounded to 100 training steps so the whole run finishes in 21 minutes, and the output under each step is what it printed; the environment had been built the day before, the base was already on disk, and the 27B-only commands are the project's, marked as such. The 27B numbers next to them are the project's own runs, from the ledger.

Prerequisites

A Mac with Apple Silicon. The 9B demo in this page peaked at 17.5 GB and runs on a 32 GB machine; the 27B needs 96 GB, and the project measured it on 96 GB only. Start with the 9B whatever your machine, because a mistake costs 20 minutes instead of 19 hours.

uv on the path, and Python 3.12 through it. The environment script installs the interpreter itself.

The hf CLI for the base download; the environment script installs huggingface_hub into the venv, so it is .venv/bin/hf.

Disk: 12 GB for the 9B base plus 12 GB for its merged copy and 400 MB per full-state checkpoint at rank 32; 31 GB for the 27B base, 31 GB for its merge and 4.5 GB per rank-128 checkpoint. Add 52 GB if you build the 6-bit and 4-bit members from a merged bf16, and delete it after.

A folder of Markdown, or any text you can split into sections. The demo trains on 20 leanzero.net pages; the project's Forge dataset is not published, and the format is in step 3.

MLX_DISABLE_COMPILE=1 exported in every shell that trains. The trainer asserts it and refuses to start without it.

1. Pin the toolchain with the environment script and apply the two patches the project trained with.
2. Download the 8-bit base with the speculative-decoding head as a sidecar, and check it.
3. Build the dataset from a folder of Markdown with the dataset script, and read the length histogram.
4. Train the adapter in fresh-process segments with the trainer and its driver, and read the loss curve for epochs.
5. Scale the adapter with the expander to what the machine fits, and halve the learning rate when you do.
6. Evaluate against the base and let a rule in code choose the round.
7. Merge into the 8-bit shards with the merge script, and gate the merged model against base-plus-adapter.

Step 1 — Install MLX and mlx-lm for LoRA fine-tuning on Apple Silicon

The project trained on mlx 0.31.2 and mlx-lm 0.31.3 with Python 3.12, installed through uv. Those versions are pinned, and they matter, because two patches sit on top of them. This is the environment script, make_env.sh; save it in an empty folder and run it once from there:

#!/bin/zsh
# Pinned trainer environment: mlx 0.31.2 + mlx-lm 0.31.3 on Python 3.12, plus the two upstream patches the
# Atlassian models were trained with (PR 1389, never merged; PR 1661, merged after the 0.31.3 release).
set -euo pipefail
cd "$(dirname "$0")"
uv python install 3.12
uv venv --python 3.12 .venv
uv pip install --python .venv/bin/python "mlx==0.31.2" "mlx-lm==0.31.3" huggingface_hub numpy safetensors pyyaml
SP=.venv/lib/python3.12/site-packages
curl -sL https://github.com/ml-explore/mlx-lm/pull/1389.diff -o 1389.diff
curl -sL https://github.com/ml-explore/mlx-lm/pull/1661.diff -o 1661.diff
for d in 1389.diff 1661.diff; do
  # the GitHub diffs carry test files that do not exist in site-packages; apply only the mlx_lm/ hunks
  git apply --directory="$SP" --include="$SP/mlx_lm/*" "$d" 2>/dev/null || echo "$d: already applied or no matching hunks"
done
sed -i '' 's/^SUB_BLOCK = 16/SUB_BLOCK = 8/' "$SP/mlx_lm/models/gated_delta.py"
.venv/bin/python -c "import mlx.core as mx, mlx_lm, mlx_lm.models.gated_delta as g; assert g.SUB_BLOCK == 8; print(mx.__version__, mlx_lm.__version__, 'SUB_BLOCK=8 env ok')"

The first patch is mlx-lm pull request 1389, "Add chunk-parallel gated delta ops for training". It is closed and was never merged, and it is the difference between about 50 tokens per second stock and 111 to 150 patched in the project's throughput measurement, with 113 to 117 across the rank-32 27B rounds. The second is pull request 1661, "Honor seed=0 in iterate_batches", merged upstream on 5 August but after the 0.31.3 release, so a fresh install still lacks it. Both flags on git apply matter: the GitHub diffs are the full pull requests and carry test files that do not exist in site-packages, so plain patch stops at the first missing file and the second diff never applies; and if your folder is a git repository, git apply without --directory resolves the paths against the repository root, matches nothing, and exits 0 having applied nothing. The sed lowers one constant that pull request 1389 introduces, SUB_BLOCK, from 16 to 8, exactly as the project's environment does, and the last line asserts it took.

One environment variable is not optional. The trainer asserts MLX_DISABLE_COMPILE=1 at startup and refuses to run without it, with the reason in the assertion message: compile-cache shape retention was the measured out-of-memory. The driver exports it before every segment. Export it in every shell that trains.

How you know it worked: the last line of the script prints the two pinned versions and the constant. From the demo run:

.venv/bin/python -c "import mlx.core as mx, mlx_lm, mlx_lm.models.gated_delta as g; print(mx.__version__, mlx_lm.__version__, 'SUB_BLOCK', g.SUB_BLOCK)"

0.31.2 0.31.3 SUB_BLOCK 8

git apply is silent on success. Do not re-run the whole script to check: uv venv refuses to overwrite an existing .venv and set -e stops the script there. Re-run the two git apply lines by hand instead; each prints "already applied or no matching hunks" when the diff is in. If the assert fails with no attribute named SUB_BLOCK, the first diff was never applied.

What is unverified: whether stock mlx_lm.lora without pull request 1389 can train the 27B base at a 4,096-token sequence length inside 96 GB. The project's stock measurement was at 1,024 tokens only, 50 tokens per second at 52 GB. Do not assume the longer length fits without the patch.

Step 2 — Download the Qwen3.8-27B 8-bit base for fine-tuning

Full fine-tuning a 27B needs several times its bf16 weights in memory; the project's own plan put it at 324 to 432 GB, which no 96 GB machine has. The design is therefore a frozen 8-bit base with a low-rank adapter on top. The base is not the community's 8-bit checkpoint; the project built its own, and the reason is a file called mtp.safetensors.

Qwen3.8-27B ships a multi-token-prediction head, the module that speculative decoding uses to draft tokens. mlx-lm strips it at load, which is why MTP on Apple Silicon needs a different serving path. mlx-lm's loader drops any mtp. tensors it finds, and the presence of those keys is also its signal to add 1.0 to every backbone norm: input, post-attention, final, q and k norms. The conversion has already applied that shift once (final-norm mean 0.944 in bf16 becomes 1.944 in the converted folder, which is the correct state), so an inline head would have the loader shift the norms a second time. The project caught this in the converted folder's key list and rebuilt with the sidecar rather than train on it. The fix is to keep the head out of the main shards, in a sidecar the loader never reads, and to build that sidecar at full precision from the bf16 source.

The shortest path is to download the bases the project trained on, which are public. Both are plain affine 8-bit checkpoints with group size 64 and the sidecar included. The 9B is Qwen3.5, not Qwen3.8, and both know almost nothing about Forge, which is the point; they are the untrained starting line:

mkdir -p models
.venv/bin/hf download Mihai-LeanZero/Qwen3.5-9B-Atlassian-Q8-base-mlx --local-dir models/Qwen3.5-9B-Atlassian-Q8-base-mlx
.venv/bin/hf download Mihai-LeanZero/Qwen3.8-27B-Atlassian-Q8-base-mlx --local-dir models/Qwen3.8-27B-Atlassian-Q8-base-mlx

The 9B is 12 GB on disk, the 27B about 31 GB. One honesty note: the demo on this machine pointed models/ at the project's local copy of the 9B base rather than downloading it again, so the download line itself was not run for this page; the folder it checks below is that same base. If you want to build the 27B base yourself, the project used the mlx-node CLI (npm i @mlx-node/cli@0.0.13, run from a local node_modules/.bin/mlx) on the bf16 source:

mlx convert -i models/Qwen3.8-27B-bf16 -o models/Qwen3.8-27B-Forge-Q8-mlx \
  -q --q-mtp cyankiwi --q-bits 8 --q-group-size 64 --q-recipe qwen3_5

That took 50 seconds and 28.7 GB of resident memory. The cyankiwi value writes the head as a sidecar rather than inline; the project then replaced mlx-node's mixed-precision sidecar with a full-precision one built from the bf16 source, because the serving engine it uses infers the head's precision from one tensor and rejects a mixed head. That replacement script is the project's own and is not in this kit.

Why 8-bit and not a smarter 6-bit? The project scored the candidates against the bf16 teacher on 200 prompts with a KL-divergence gate, the same method as the NVFP4 comparison. Uniform 8-bit scored 0.0446 with 99.35% top-1 agreement; an Unsloth-recipe 6-bit candidate with an importance matrix scored 0.0591 and 98.90%, inside the KL bound but under the 99% top-1 floor, so the rule chose 8-bit.

How you know it worked: three checks, all from the demo, on the 9B folder. List it, print the quantisation block from its config, and count mtp. keys in the shard index:

ls models/Qwen3.5-9B-Atlassian-Q8-base-mlx/
.venv/bin/python -c "import json;print(json.load(open('models/Qwen3.5-9B-Atlassian-Q8-base-mlx/config.json'))['quantization'])"
.venv/bin/python -c "import json;idx=json.load(open('models/Qwen3.5-9B-Atlassian-Q8-base-mlx/model.safetensors.index.json'))['weight_map'];print('mtp. keys in main shards:', sum(1 for k in idx if k.startswith('mtp.')))"

README.md
chat_template.jinja
config.json
merges.txt
model-00001-of-00003.safetensors
model-00002-of-00003.safetensors
model-00003-of-00003.safetensors
model.safetensors.index.json
mtp.safetensors
mtplx_runtime.json
preprocessor_config.json
tokenizer.json
tokenizer_config.json
video_preprocessor_config.json
vocab.json
{'group_size': 64, 'bits': 8, 'mode': 'affine'}
mtp. keys in main shards: 0

Three shards for the 9B, six for the 27B, mtp.safetensors beside them, and zero head keys in the index. That last number is the one that matters. Then ask the untrained base a question it cannot know, so you have a before to compare the after against. The prompt is the exact shape the dataset in step 3 will use:

.venv/bin/python -m mlx_lm.generate --model models/Qwen3.5-9B-Atlassian-Q8-base-mlx \
  --prompt 'Write the section "Step 1 — Pin the toolchain" of the LeanZero tutorials page "Fine-tune Qwen3.8-27B with LoRA on one Mac Studio: how the Atlassian model was built".' \
  --max-tokens 160

==========
Here's a thinking process that leads to the suggested content:

1.  **Analyze the Request:**
    *   **Topic:** "Step 1 — Pin the toolchain" for a tutorial page.
    *   **Context:** "Fine-tune Qwen3.8-27B with LoRA on one Mac Studio: how the Atlassian model was built".
    *   **Goal:** Write the specific section content.
    *   **Tone:** Technical, instructional, clear, professional (Atlassian style).
    *   **Key Elements:** Needs to cover environment setup, version pinning, dependencies, and the rationale behind pinning (reproducibility).

2.  **Determine the Technical Stack:**
    *   *
==========
Prompt: 58 tokens, 171.340 tokens-per-sec
Generation: 160 tokens, 63.105 tokens-per-sec
Peak memory: 11.668 GB

Two things to notice. The base answers in its thinking mode, because the chat template turns it on by default and mlx_lm.generate does not turn it off; the trainer in step 4 pins the template's arguments so the training render and the serving render agree. And it is about to make the section up. 63 tokens per second and 11.7 GB peak is the 9B at 8 bits on this machine.

Never rename mtp.safetensors to match the model*.safetensors pattern. mlx-lm would read it into the backbone and shift the norms twice. And when building from the bf16 source, never leave --q-mtp at its default off: that is what left the 15 head tensors inline in the project's first build.

Step 3 — Build the fine-tuning dataset

The trainer reads chat-JSONL: one object per line with a messages array of user and assistant turns and a free meta object. This is the shape:

{"messages": [{"role": "user", "content": "..."}, {"role": "assistant", "content": "..."}],
 "meta": {"type": "docs_text", "title": "..."}}

Each record is rendered through the model's chat template with two pinned arguments, reasoning_effort: medium and preserve_thinking: false, and the prompt is masked so loss is computed on the assistant turn only.

Where do the records come from? The project's Forge mix is assembled from private app repositories, validator output and generated briefs, and cannot be reproduced here. What can be reproduced is its largest single source: 2,000 documentation-section rows, each a "write this section of this page" pair, 785 thousand tokens of the 5.19 million in the mix the shipped v0.4 trained on. The kit's dataset script does exactly that from any folder of Markdown. Every ## section of every page becomes one sample, sections outside 300 to 6,000 characters are skipped, duplicates on the answer text are dropped, and the validation split is a fixed 8% bucket of a hash of the pair, so re-running on a bigger corpus never moves a row from validation into training. Save it as make_dataset.py:

"""Turn a folder of Markdown into chat-JSONL for the trainer: every '## ' section of every page becomes one sample,
user = "Write the section '<heading>' of the page '<title>'." -> assistant = the section text. Sections of 300-6000
characters only. Over-length is handled later by the trainer (filtered, never truncated). The validation split is a
FIXED 8% bucket of a hash of (prompt, answer), so re-running on a bigger corpus never moves a row from valid to train.
  python make_dataset.py <markdown-dir> <out-dir> [--kind "LeanZero tutorials"]"""
import glob, json, os, re, sys, random, hashlib
src, out = sys.argv[1], sys.argv[2]; kind = sys.argv[sys.argv.index("--kind") + 1] if "--kind" in sys.argv else "documentation"
random.seed(29); rows = []
for p in sorted(glob.glob(os.path.join(src, "**", "*.md"), recursive=True)):
    txt = open(p, errors="ignore").read()
    if txt.startswith("---"): txt = txt.split("---", 2)[-1]                       # drop frontmatter
    title = next((l.lstrip("# ").strip() for l in txt.splitlines() if l.startswith("# ")), os.path.basename(p)[:-3])
    for sec in re.split(r"\n(?=## )", txt):
        sec = sec.strip()
        if not (300 <= len(sec) <= 6000) or not sec.startswith("## "): continue
        head = sec.splitlines()[0].lstrip("# ").strip(); body = "\n".join(sec.splitlines()[1:]).strip()
        if len(body) < 200: continue
        q = random.choice([f'Write the section "{head}" of the {kind} page "{title}".',
                           f'Reproduce the {kind} section "{head}" (page: {title}).',
                           f'From the {kind} page "{title}", what does the section "{head}" say? Give the full text.'])
        rows.append({"messages": [{"role": "user", "content": q}, {"role": "assistant", "content": body}], "meta": {"type": "docs_text", "title": title}})
ph = lambda r: int(hashlib.sha1((r["messages"][0]["content"] + "\n---\n" + r["messages"][-1]["content"]).encode()).hexdigest()[:8], 16) % 100
seen, uniq = set(), []
for r in rows:
    h = ph(r)
    if (r["messages"][-1]["content"]) in seen: continue
    seen.add(r["messages"][-1]["content"]); uniq.append((h, r))
random.shuffle(uniq); os.makedirs(out, exist_ok=True)
train = [r for h, r in uniq if h >= 8]; valid = [r for h, r in uniq if h < 8]
for name, rs in (("train", train), ("valid", valid)):
    with open(os.path.join(out, f"{name}.jsonl"), "w") as fh:
        for r in rs: fh.write(json.dumps(r) + "\n")
print(f"{len(train)} train / {len(valid)} valid samples from {len(rows)} sections; ~{int(sum(len(r['messages'][1]['content']) for _, r in uniq) / 3.6):,} tokens")

Run it on a folder of Markdown. The demo used 20 leanzero.net pages, 8 tutorials and 12 posts, copied into corpus/:

.venv/bin/python make_dataset.py corpus data/leanzero-docs --kind "LeanZero tutorials"

112 train / 18 valid samples from 130 sections; ~68,053 tokens

Two decisions in the data path are worth copying whatever your sources are. First, over-length samples are filtered, never truncated, because a truncated sample teaches the model to keep going past where it should have stopped, which shows up later as answers that never end. Second, the sequence length is chosen from the data, not guessed. The trainer has a --check-data mode that renders every row through the tokenizer and prints the length histogram at 1,024, 2,048 and 4,096 tokens without loading the model, so it runs in seconds on the CPU:

.venv/bin/python segment_train.py --config demo.yaml --check-data

[data:train@T=1024] n=112 kept=107 filtered_overlength=5 (4.5%) prompt_only_overlength=0 | len p50/p90/p95/p99=[407, 883, 983, 1376] max=1766 T=1024
[data:valid@T=1024] n=18 kept=15 filtered_overlength=3 (16.7%) prompt_only_overlength=0 | len p50/p90/p95/p99=[512, 1120, 1165, 1185] max=1191 T=1024
[data:train@T=2048] n=112 kept=112 filtered_overlength=0 (0.0%) prompt_only_overlength=0 | len p50/p90/p95/p99=[407, 883, 983, 1376] max=1766 T=2048
[data:valid@T=2048] n=18 kept=18 filtered_overlength=0 (0.0%) prompt_only_overlength=0 | len p50/p90/p95/p99=[512, 1120, 1165, 1185] max=1191 T=2048
[data:train@T=4096] n=112 kept=112 filtered_overlength=0 (0.0%) prompt_only_overlength=0 | len p50/p90/p95/p99=[407, 883, 983, 1376] max=1766 T=4096
[data:valid@T=4096] n=18 kept=18 filtered_overlength=0 (0.0%) prompt_only_overlength=0 | len p50/p90/p95/p99=[512, 1120, 1165, 1185] max=1191 T=4096
[check-data] done

Read it as a decision. At 1,024 tokens the demo corpus would lose 5 of 112 training rows and 3 of 18 validation rows, and the rows it loses are the longest sections, which are the ones most worth learning. At 2,048 it loses nothing, so 2,048 would do for this corpus; the demo config keeps 4,096 because the 27B recipe uses it and the memory cost on 112 short rows is nil. The project's real Forge samples made the same call the hard way: at 1,024 tokens 43.7% of them would have been dropped, at 2,048 17.6%, at 4,096 only 0.6%. On the shipped mix the histogram read:

[data:train] n=13840 kept=13837 filtered_overlength=3 (0.0%) prompt_only_overlength=0 | len p50/p90/p95/p99=[212, 942, 1137, 2171] max=5309 T=4096
[data:valid] n=1223 kept=1221 filtered_overlength=2 (0.2%) prompt_only_overlength=0 | len p50/p90/p95/p99=[202, 868, 1118, 1841] max=6000 T=4096

How you know it worked: train.jsonl and valid.jsonl exist in the output folder, the dataset script printed a token estimate you believe, and --check-data printed a filtered count you can explain with a median length well under your sequence length. If the 99th percentile sits at your limit, you are filtering the samples that matter most; raise the limit or split the samples.

The project's mix also holds general-purpose replay, 500 rows at 11.5% of tokens in the shipped mix and 17.6% in the next one, because every capacity checkpoint failed a knowledge or a looping gate; step 6 has that story. A corpus of one subject, like the demo's, has none, and step 4 shows what that costs.

Step 4 — Fine-tune Qwen3.8-27B: train the LoRA adapter with MLX

This is the step where the project's trainer differs from the stock one. The stock path is mlx_lm.lora with a YAML config; the project used it for its throughput measurement. Its training runs used a segmented trainer built from mlx-lm's pieces, the LoRA wiring, the loss and the gradient checkpointing, instead of mlx-lm's training loop. The reason is memory. Every monotone leak dies with the process, so the trainer runs a segment of steps in a fresh process, writes a full-state checkpoint, and exits; a driver relaunches it until the configured step count is done. A killed segment costs minutes, not a day.

Three files. First the config, demo.yaml, which is the rank-32 shape every public release used, sized for the 9B demo:

# A short, real run on the public 9B base: 16 of 32 layers, rank 32, 100 steps in two 50-step segments (~21 min end to end).
model: models/Qwen3.5-9B-Atlassian-Q8-base-mlx
data: data/leanzero-docs
adapter_path: results/demo/adapter
ckpt_dir: results/demo/ckpt
total_iters: 100
segment_iters: 50
save_every: 25
batch_size: 1
grad_accumulation_steps: 8
max_seq_length: 4096
num_layers: 16
learning_rate: 5.0e-5
seed: 7
mask_prompt: true
memory_limit_gb: 40
cache_limit_gb: 8
val_batches: 32
steps_per_report: 10
lora_parameters:
  rank: 32
  scale: 2.0
  dropout: 0.0
  keys: [self_attn.v_proj, self_attn.o_proj, linear_attn.in_proj_qkv, linear_attn.out_proj, mlp.gate_proj, mlp.up_proj, mlp.down_proj]

A step is one optimiser update, eight micro-batches of one sample each. Write the learning rate with a decimal point: the trainer's own comment records that YAML parses 1e-5 without one as a string. The rank-32 learning rate of 5e-5 came from a three-point sweep on a 9B pilot, where 2e-4 overfit and 5e-5 won; the rank-128 rate is a different story, in step 5. For the 27B, the project's config differs in six lines: the model path, num_layers: 16 of 64, memory_limit_gb: 60, total_iters: 2500 for round one (rounds three and four used 1,500, round five 4,500, the shipped v0.4 round 2,400), segment_iters: 100, and a learning rate of 5e-5 for round one, 3e-5 for rounds three and four, 4e-5 for the shipped round.

Second, the driver, driver.sh. It exports the compile flag, relaunches the trainer after every SEGMENT_DONE, stops on TRAINING_COMPLETE, and on a crash resumes from the last checkpoint with two strikes per segment before it stops for a human:

#!/bin/zsh
# Segment driver: respawn segment_train.py until it prints TRAINING_COMPLETE; resume from the last checkpoint on a crash
# (two strikes per segment, then stop for a human). Usage: ./driver.sh <config.yaml> [segment_iters]
set -uo pipefail; cd "$(dirname "$0")"; C=$1; SEG=${2:-100}; L=$(basename "$C" .yaml); mkdir -p logs
export MLX_DISABLE_COMPILE=1
strikes=0
while true; do
  LOG=logs/train-$L-$(date +%Y%m%d-%H%M%S).log
  .venv/bin/python segment_train.py --config "$C" --segment-iters "$SEG" > "$LOG" 2>&1; rc=$?
  tail -3 "$LOG"
  if grep -q TRAINING_COMPLETE "$LOG"; then echo "DRIVER: training complete"; exit 0; fi
  if grep -q SEGMENT_DONE "$LOG"; then strikes=0; continue; fi
  strikes=$((strikes+1)); echo "DRIVER: segment failed rc=$rc (strike $strikes/2) — resuming from last checkpoint"
  (( strikes >= 2 )) && { echo "DRIVER: two strikes — stopping for a human"; exit 1; }
  sleep 10
done

The project's own driver is more defensive than that, in ways a 96 GB machine needs. It refuses to start if LM Studio has a model loaded or if a serving engine is resident, because the 27B base alone is 31 GB and there is no room for two. Every 25 seconds it reads the operating system's GPU allocation and swap figures and kills the segment if swap has grown 3 GB over what it was when the segment started, or reaches 12 GB. That rule used to be an absolute 4 GB, and on 13 September it killed a capacity smoke at step 1 on 4.9 GB of swap that belonged to seven orphaned Playwright browser processes from the day before, about 12 GB resident between them. Swap that was there before you started is not your run's swap; measure growth. The project also ran the whole thing under caffeinate so the machine could not sleep mid-segment. Add those to the kit's driver when you move to the 27B.

Third, the trainer itself, segment_train.py. It is the project's trainer with two research additions removed, a teacher-distillation loss and a cosine schedule, and nothing else changed. Read the docstring; every design decision in it was paid for by a measured failure, and the comments name which:

"""Segmented LoRA trainer with FULL-STATE checkpointing, built on mlx-lm 0.31.3's pieces (loss, grad-checkpoint, LoRA wiring)
instead of its train() loop. Plan §4/§7b (Stage-0 hardening):

  * full-state checkpoints every `save_every` and at segment end: adapter weights + optimizer state (incl. the step scalar)
    + mx.random state + data cursor (epoch, position), written to a tmp dir and committed by atomic rename; keeps the last 2
  * segmented execution: trains `segment_iters` steps then exits 0 with SEGMENT_DONE; a fresh process resumes from the latest
    checkpoint (driver.sh loops) — every monotone-leak class dies with the process
  * data order is a deterministic permutation of (seed, epoch) with a saved cursor, so segments never replay the same samples
  * chat template rendered with PINNED kwargs (reasoning_effort=medium, preserve_thinking=False by default) — the datasets.py
    defect fixed at the source of the samples; mask_prompt offsets derived from the same render
  * over-length samples are FILTERED, never truncated (a truncated sample teaches "keep going" — the EOS defect); prompt-only
    over-length samples are dropped; a length histogram is printed (this doubles as the S10 tool)
  * NaN guard: clip_grad_norm + skip update on a non-finite norm; abort the segment after `nan_skip_max` consecutive skips
  * memory: mx.set_memory_limit, mx.clear_cache() every 50 steps, active/peak reported; MLX_DISABLE_COMPILE=1 required (asserted)

Usage:  python train/segment_train.py --config configs/<run>.yaml [--segment-iters N] [--check-data]
Config keys (yaml): model, data (dir with train.jsonl/valid.jsonl), adapter_path (final adapter dir), ckpt_dir, total_iters,
  segment_iters, save_every, batch_size (1), grad_accumulation_steps, max_seq_length, num_layers, lora_parameters{rank,scale,
  dropout,keys}, learning_rate, seed, mask_prompt, template_kwargs{...}, memory_limit_gb, nan_skip_max, steps_per_report,
  val_batches, grad_checkpoint.
"""
import argparse, json, os, shutil, sys, time, glob, math
import numpy as np
import yaml
import mlx.core as mx
import mlx.nn as nn
import mlx.optimizers as optim
from mlx.utils import tree_flatten, tree_unflatten, tree_map
from mlx_lm import load
from mlx_lm.tuner.utils import linear_to_lora_layers, print_trainable_parameters
from mlx_lm.tuner.trainer import default_loss, grad_checkpoint

DEFAULT_TEMPLATE_KWARGS = {"reasoning_effort": "medium", "preserve_thinking": False}


class PinnedChatDataset:
    """{"messages": [...]} rows rendered with pinned chat-template kwargs; returns (tokens, prompt_offset)."""

    def __init__(self, rows, tokenizer, template_kwargs, mask_prompt=True):
        self.items = []
        for d in rows:
            msgs = d["messages"]; tools = d.get("tools")
            tokens = tokenizer.apply_chat_template(msgs, tools=tools, return_dict=False, **template_kwargs)
            offset = 0
            if mask_prompt:
                offset = len(tokenizer.apply_chat_template(msgs[:-1], tools=tools, add_generation_prompt=True, return_dict=False, **template_kwargs))
            self.items.append((list(tokens), offset))

    def __len__(self):
        return len(self.items)


def filter_and_histogram(ds: PinnedChatDataset, max_seq_length: int, name: str):
    lens = np.array([len(t) for t, _ in ds.items]); offs = np.array([o for _, o in ds.items])
    keep = [i for i in range(len(ds)) if lens[i] <= max_seq_length and offs[i] < max_seq_length - 8]
    over = int((lens > max_seq_length).sum()); prompt_over = int((offs >= max_seq_length - 8).sum())
    pct = [int(np.percentile(lens, p)) for p in (50, 90, 95, 99)] if len(lens) else []
    print(f"[data:{name}] n={len(ds)} kept={len(keep)} filtered_overlength={over} ({100*over/max(1,len(ds)):.1f}%) "
          f"prompt_only_overlength={prompt_over} | len p50/p90/p95/p99={pct} max={int(lens.max()) if len(lens) else 0} T={max_seq_length}", flush=True)
    ds.items = [ds.items[i] for i in keep]
    return {"n": int(len(lens)), "kept": len(keep), "filtered_overlength": over, "prompt_only_overlength": prompt_over, "pcts": pct}


def make_batch(items, idxs, max_seq_length):
    batch = [items[j] for j in idxs]
    toks, offsets = zip(*batch); lengths = [len(t) for t in toks]
    pad_to = 32
    L = min(1 + pad_to * ((max(lengths) + pad_to - 1) // pad_to), max_seq_length)
    arr = np.zeros((len(batch), L), np.int32)
    for j, t in enumerate(toks):
        n = min(lengths[j], max_seq_length); arr[j, :n] = t[:n]; lengths[j] = n
    return mx.array(arr), mx.array(list(zip(offsets, lengths)))


def epoch_order(seed, epoch, n_batches):
    return np.random.default_rng(seed * 100003 + epoch).permutation(n_batches)


def latest_ckpt(ckpt_dir):
    c = sorted(glob.glob(os.path.join(ckpt_dir, "ckpt-*")), key=lambda p: int(p.rsplit("-", 1)[1]) if p.rsplit("-", 1)[1].isdigit() else -1)
    c = [p for p in c if not p.endswith(".tmp") and os.path.exists(os.path.join(p, "state.json"))]
    return c[-1] if c else None


def save_adapter_dir(model, cfg, path):
    os.makedirs(path, exist_ok=True)
    mx.save_safetensors(os.path.join(path, "adapters.safetensors"), dict(tree_flatten(model.trainable_parameters())))
    acfg = {"fine_tune_type": "lora", "num_layers": cfg["num_layers"], "lora_parameters": cfg["lora_parameters"],
            "model": cfg["model"], "max_seq_length": cfg["max_seq_length"], "learning_rate": cfg["learning_rate"],
            "template_kwargs": cfg["template_kwargs"], "trainer": "segment_train.py"}
    json.dump(acfg, open(os.path.join(path, "adapter_config.json"), "w"), indent=1)


def save_checkpoint(ckpt_dir, step, model, optimizer, cfg, loop_state):
    final = os.path.join(ckpt_dir, f"ckpt-{step}"); tmp = final + ".tmp"
    shutil.rmtree(tmp, ignore_errors=True); os.makedirs(tmp)
    save_adapter_dir(model, cfg, tmp)
    mx.save_safetensors(os.path.join(tmp, "optimizer.safetensors"), dict(tree_flatten(optimizer.state)))
    mx.save_safetensors(os.path.join(tmp, "rng.safetensors"), {"mx_random_state": mx.random.state[0]})
    json.dump({**loop_state, "step": step, "saved_at": time.strftime("%F %T")}, open(os.path.join(tmp, "state.json"), "w"), indent=1)
    os.replace(tmp, final)
    # keep last 2
    olds = sorted([p for p in glob.glob(os.path.join(ckpt_dir, "ckpt-*")) if not p.endswith(".tmp")], key=lambda p: int(p.rsplit("-", 1)[1]))
    for p in olds[:-2]: shutil.rmtree(p, ignore_errors=True)
    print(f"[ckpt] saved {final}", flush=True)
    return final


def restore_checkpoint(path, model, optimizer):
    model.load_weights(os.path.join(path, "adapters.safetensors"), strict=False)
    opt_state = tree_unflatten(list(mx.load(os.path.join(path, "optimizer.safetensors")).items()))
    optimizer.state = opt_state
    mx.random.state[0] = mx.load(os.path.join(path, "rng.safetensors"))["mx_random_state"]
    st = json.load(open(os.path.join(path, "state.json")))
    mx.eval(model.parameters(), optimizer.state)
    print(f"[ckpt] restored {path}: step={st['step']} epoch={st['epoch']} pos={st['pos']} opt_step={int(optimizer.state.get('step', mx.array(0)).item()) if isinstance(optimizer.state, dict) and 'step' in optimizer.state else '?'}", flush=True)
    return st


def eval_loss(model, items, cfg, n_batches, seed):
    if not items: return float("nan")
    rng = np.random.default_rng(seed); idx = rng.permutation(len(items))[:n_batches]
    tot, ntok = 0.0, 0
    for i in idx:
        b, l = make_batch(items, [int(i)], cfg["max_seq_length"]); ce, n = default_loss(model, b, l); mx.eval(ce, n)
        tot += float(ce.item()) * int(n.item()); ntok += int(n.item())
    return tot / max(1, ntok)


def main():
    ap = argparse.ArgumentParser(); ap.add_argument("--config", required=True); ap.add_argument("--segment-iters", type=int)
    ap.add_argument("--check-data", action="store_true", help="render + histogram the data, then exit"); a = ap.parse_args()
    assert os.environ.get("MLX_DISABLE_COMPILE") == "1", "export MLX_DISABLE_COMPILE=1 (compile-cache shape retention is the measured OOM)"
    cfg = yaml.safe_load(open(a.config))
    cfg.setdefault("template_kwargs", DEFAULT_TEMPLATE_KWARGS); cfg.setdefault("batch_size", 1); cfg.setdefault("grad_accumulation_steps", 1)
    cfg.setdefault("save_every", 100); cfg.setdefault("steps_per_report", 10); cfg.setdefault("val_batches", 8); cfg.setdefault("seed", 7)
    cfg.setdefault("memory_limit_gb", 60); cfg.setdefault("nan_skip_max", 2); cfg.setdefault("mask_prompt", True); cfg.setdefault("grad_checkpoint", True)
    cfg.setdefault("ckpt_dir", os.path.join(os.path.dirname(cfg["adapter_path"]), "ckpt")); cfg.setdefault("segment_iters", 500)
    seg_iters = a.segment_iters or cfg["segment_iters"]
    os.makedirs(cfg["ckpt_dir"], exist_ok=True)
    mx.set_memory_limit(int(cfg["memory_limit_gb"] * 1e9))
    mx.set_cache_limit(int(cfg.get("cache_limit_gb", 8) * 1e9))  # S4 finding: without a cache cap the freed-buffer cache grew to the memory limit (59.5 GB OS alloc vs 35 GB working set)

    if a.check_data:  # tokenizer only — no model, no GPU
        from mlx_lm.utils import load_tokenizer
        from pathlib import Path
        tokenizer = load_tokenizer(Path(cfg["model"]))
        for T in sorted({cfg["max_seq_length"], 1024, 2048, 4096}):
            for name in ("train", "valid"):
                p = os.path.join(cfg["data"], f"{name}.jsonl")
                if not os.path.exists(p): continue
                ds = PinnedChatDataset([json.loads(l) for l in open(p)], tokenizer, cfg["template_kwargs"], cfg["mask_prompt"]); filter_and_histogram(ds, T, f"{name}@T={T}")
        print("[check-data] done"); return
    model, tokenizer = load(cfg["model"])
    model.freeze()
    mx.random.seed(cfg["seed"])  # BEFORE LoRA init: lora_a is random-initialised (S3 finding: unseeded init diverged runs from iter 1)
    np.random.seed(cfg["seed"])
    linear_to_lora_layers(model, cfg["num_layers"], cfg["lora_parameters"])
    print_trainable_parameters(model)
    if cfg["grad_checkpoint"]: grad_checkpoint(model.layers[0])

    rows_tr = [json.loads(l) for l in open(os.path.join(cfg["data"], "train.jsonl"))]
    rows_va = [json.loads(l) for l in open(os.path.join(cfg["data"], "valid.jsonl"))] if os.path.exists(os.path.join(cfg["data"], "valid.jsonl")) else []
    tr = PinnedChatDataset(rows_tr, tokenizer, cfg["template_kwargs"], cfg["mask_prompt"]); va = PinnedChatDataset(rows_va, tokenizer, cfg["template_kwargs"], cfg["mask_prompt"])
    hist = {"train": filter_and_histogram(tr, cfg["max_seq_length"], "train"), "valid": filter_and_histogram(va, cfg["max_seq_length"], "valid")}
    json.dump(hist, open(os.path.join(cfg["ckpt_dir"], "data_histogram.json"), "w"), indent=1)
    if a.check_data: print("[check-data] done"); return
    assert len(tr) >= cfg["batch_size"], "empty training set after filtering"
    # length-sorted fixed batches (like mlx-lm), permuted per epoch
    order = sorted(range(len(tr)), key=lambda i: len(tr.items[i][0])); bs = cfg["batch_size"]
    batches = [order[i:i + bs] for i in range(0, len(order) - bs + 1, bs)]

    cfg["learning_rate"] = float(cfg["learning_rate"])  # YAML parses 1e-5 (no dot) as a STRING
    optimizer = optim.Adam(learning_rate=cfg["learning_rate"])
    loss_value_and_grad = nn.value_and_grad(model, default_loss)

    ck = latest_ckpt(cfg["ckpt_dir"])
    if ck:
        st = restore_checkpoint(ck, model, optimizer); step, epoch, pos, trained_tokens, nan_skips_total = st["step"], st["epoch"], st["pos"], st.get("trained_tokens", 0), st.get("nan_skips_total", 0)
    else:
        step, epoch, pos, trained_tokens, nan_skips_total = 0, 0, 0, 0, 0
        if cfg.get("init_adapter"):  # continue from a previous run's adapter weights with a fresh optimizer (T2 round)
            model.load_weights(os.path.join(cfg["init_adapter"], "adapters.safetensors"), strict=False); print(f"[init] LoRA weights loaded from {cfg['init_adapter']}", flush=True)
        mx.eval(model.parameters())
    total = cfg["total_iters"]; seg_end = min(total, step + seg_iters)
    print(f"[segment] step {step} -> {seg_end} of {total}; epoch={epoch} pos={pos} batches/epoch={len(batches)} T={cfg['max_seq_length']} template_kwargs={cfg['template_kwargs']}", flush=True)
    if step == 0 and va.items:
        print(f"[eval] step 0 val_loss={eval_loss(model, va.items, cfg, max(32, cfg['val_batches']), cfg['seed']):.4f} (fixed subset)", flush=True)   # 2026-09-13: same 32-batch subset as every later eval; the old 16-batch step-0 read ~0.09 lower on the 9B and was misread as a warm-up bump

    model.train(); losses, ntoks, t_seg = [], 0, time.perf_counter(); grad_acc = None; nan_consec = 0; tic_report = time.perf_counter(); report_tokens = 0
    while step < seg_end:
        perm = epoch_order(cfg["seed"], epoch, len(batches))
        if pos >= len(perm): epoch += 1; pos = 0; continue
        b, l = make_batch(tr.items, batches[int(perm[pos])], cfg["max_seq_length"]); pos += 1
        (lvalue, toks), grad = loss_value_and_grad(model, b, l)
        if grad_acc is not None: grad = tree_map(lambda x, y: x + y, grad, grad_acc)
        do_update = ((pos + epoch * len(batches)) % cfg["grad_accumulation_steps"] == 0) if cfg["grad_accumulation_steps"] > 1 else True
        if not do_update:
            grad_acc = grad; mx.eval(grad_acc, lvalue, toks); losses.append(float(lvalue.item()))
            n = int(toks.item()); ntoks += n; report_tokens += n; trained_tokens += n  # count every micro-batch (S4 finding: metric read 8x low)
            continue
        if cfg["grad_accumulation_steps"] > 1: grad = tree_map(lambda x: x / cfg["grad_accumulation_steps"], grad)
        grad, gnorm = optim.clip_grad_norm(grad, max_norm=1.0)
        mx.eval(gnorm, lvalue)
        if not (math.isfinite(float(gnorm.item())) and math.isfinite(float(lvalue.item()))):
            nan_consec += 1; nan_skips_total += 1; grad_acc = None
            print(f"[nan-guard] step {step+1}: non-finite loss/grad-norm (loss={float(lvalue.item())}, gnorm={float(gnorm.item())}) — update skipped ({nan_consec} consecutive)", flush=True)
            if nan_consec >= cfg["nan_skip_max"]: print("NAN_ABORT", flush=True); sys.exit(3)
            continue
        nan_consec = 0
        optimizer.update(model, grad); grad_acc = None
        mx.eval(model.parameters(), optimizer.state)
        step += 1; n = int(toks.item()); ntoks += n; report_tokens += n; trained_tokens += n; losses.append(float(lvalue.item()))
        if step % 50 == 0: mx.clear_cache()
        if step % cfg["steps_per_report"] == 0 or step == seg_end:
            dt = time.perf_counter() - tic_report; k = cfg["steps_per_report"] if step % cfg["steps_per_report"] == 0 else (step % cfg["steps_per_report"])
            print(f"Iter {step}: Train loss {np.mean(losses[-k:]):.4f}, gnorm {float(gnorm.item()):.3f}, It/sec {k/dt:.3f}, Tokens/sec {report_tokens/dt:.1f}, "
                  f"Trained Tokens {trained_tokens}, active {mx.get_active_memory()/1e9:.1f} GB, peak {mx.get_peak_memory()/1e9:.1f} GB, epoch {epoch} pos {pos}", flush=True)
            tic_report = time.perf_counter(); report_tokens = 0
        if step % cfg["save_every"] == 0 and step < seg_end:
            save_checkpoint(cfg["ckpt_dir"], step, model, optimizer, cfg, {"epoch": epoch, "pos": pos, "trained_tokens": trained_tokens, "nan_skips_total": nan_skips_total, "recent_losses": losses[-50:]})
    if va.items: print(f"[eval] step {step} val_loss={eval_loss(model, va.items, cfg, max(32, cfg['val_batches']), cfg['seed']):.4f} (fixed subset)", flush=True)  # T1 finding: a per-step subset made the curve noise
    save_checkpoint(cfg["ckpt_dir"], step, model, optimizer, cfg, {"epoch": epoch, "pos": pos, "trained_tokens": trained_tokens, "nan_skips_total": nan_skips_total, "recent_losses": losses[-50:]})
    print(f"[segment] done in {time.perf_counter()-t_seg:.0f}s; tokens this segment {ntoks}", flush=True)
    if step >= total:
        save_adapter_dir(model, cfg, cfg["adapter_path"]); print(f"TRAINING_COMPLETE step={step} adapter={cfg['adapter_path']}", flush=True)
    else:
        print(f"SEGMENT_DONE step={step}", flush=True)


if __name__ == "__main__":
    main()

Six things in that file are worth reading even if you never run it. The random seed is set before the LoRA layers are initialised, because an unseeded initialisation diverged runs from the first step. The memory limit and the cache limit are set explicitly, because without the cache cap the freed-buffer cache grew to the memory limit on its own. Over-length rows are filtered at load, never truncated. The gradient norm is clipped at 1.0, with a guard that skips a non-finite update and aborts the segment after two in a row rather than training through them. The step-0 validation read uses the same 32-batch fixed subset as every later read, for a reason step 5 explains. And every micro-batch's tokens are counted, because an earlier version counted only the update steps and reported throughput eight times low.

Now run it. The demo trains 100 steps in two segments of 50, with save_every: 25 so there are checkpoints inside each segment:

./driver.sh demo.yaml 50

[ckpt] saved results/demo/ckpt/ckpt-50
[segment] done in 609s; tokens this segment 179242
SEGMENT_DONE step=50
[ckpt] saved results/demo/ckpt/ckpt-100
[segment] done in 616s; tokens this segment 181720
TRAINING_COMPLETE step=100 adapter=results/demo/adapter
DRIVER: training complete

Two segments of about ten minutes each on the M3 Ultra. The driver prints the last three lines of each segment's log; the logs themselves are in logs/, and this is every reporting line from the demo's two, joined:

Trainable parameters: 0.388% (34.734M/8953.802M)
[segment] step 0 -> 50 of 100; epoch=0 pos=0 batches/epoch=112 T=4096 template_kwargs={'reasoning_effort': 'medium', 'preserve_thinking': False}
[eval] step 0 val_loss=2.6390 (fixed subset)
Iter 10: Train loss 2.5257, gnorm 0.599, It/sec 0.084, Tokens/sec 300.8, Trained Tokens 35883, active 12.0 GB, peak 15.9 GB, epoch 0 pos 80
Iter 20: Train loss 2.0322, gnorm 0.903, It/sec 0.081, Tokens/sec 304.2, Trained Tokens 73645, active 12.0 GB, peak 15.9 GB, epoch 1 pos 48
[ckpt] saved results/demo/ckpt/ckpt-25
Iter 30: Train loss 1.6435, gnorm 1.539, It/sec 0.089, Tokens/sec 296.6, Trained Tokens 106812, active 12.0 GB, peak 17.0 GB, epoch 2 pos 16
Iter 40: Train loss 1.2991, gnorm 1.524, It/sec 0.079, Tokens/sec 303.6, Trained Tokens 145064, active 12.0 GB, peak 17.5 GB, epoch 2 pos 96
Iter 50: Train loss 0.7604, gnorm 1.430, It/sec 0.087, Tokens/sec 298.6, Trained Tokens 179242, active 12.0 GB, peak 17.5 GB, epoch 3 pos 64
[eval] step 50 val_loss=2.5278 (fixed subset)
[ckpt] saved results/demo/ckpt/ckpt-50
[segment] done in 609s; tokens this segment 179242
SEGMENT_DONE step=50
Trainable parameters: 0.388% (34.734M/8953.802M)
[ckpt] restored results/demo/ckpt/ckpt-50: step=50 epoch=3 pos=64 opt_step=50
[segment] step 50 -> 100 of 100; epoch=3 pos=64 batches/epoch=112 T=4096 template_kwargs={'reasoning_effort': 'medium', 'preserve_thinking': False}
Iter 60: Train loss 0.5329, gnorm 1.746, It/sec 0.078, Tokens/sec 306.1, Trained Tokens 218633, active 12.0 GB, peak 16.7 GB, epoch 4 pos 32
Iter 70: Train loss 0.4188, gnorm 1.485, It/sec 0.088, Tokens/sec 299.6, Trained Tokens 252815, active 12.0 GB, peak 17.5 GB, epoch 4 pos 112
[ckpt] saved results/demo/ckpt/ckpt-75
Iter 80: Train loss 0.0863, gnorm 1.162, It/sec 0.085, Tokens/sec 299.2, Trained Tokens 288162, active 12.0 GB, peak 17.5 GB, epoch 5 pos 80
Iter 90: Train loss 0.1678, gnorm 1.283, It/sec 0.082, Tokens/sec 303.6, Trained Tokens 325320, active 12.0 GB, peak 17.5 GB, epoch 6 pos 48
Iter 100: Train loss 0.1338, gnorm 0.594, It/sec 0.084, Tokens/sec 297.8, Trained Tokens 360962, active 12.0 GB, peak 17.5 GB, epoch 7 pos 16
[eval] step 100 val_loss=3.0006 (fixed subset)
[ckpt] saved results/demo/ckpt/ckpt-100
[segment] done in 616s; tokens this segment 181720
TRAINING_COMPLETE step=100 adapter=results/demo/adapter

Read this log slowly, because it contains the most common fine-tuning mistake, made on purpose. The second segment restored from ckpt-50 with its optimiser step, its epoch and its data cursor, which is the full-state checkpoint doing its job: the two segments together are one run. Active memory is 12.0 GB in every line and the peak stops at 17.5 GB, which is the flat profile the fresh-process design exists to produce. Throughput is a steady 300 tokens per second. All of that is the trainer working.

Then look at the two loss columns together. Training loss falls from 2.53 to 0.13. Validation loss goes 2.639 at step 0, 2.528 at step 50, and 3.001 at step 100. The epoch field in each line says why: with 112 training rows and eight rows per step, 100 steps is seven passes over the same data, and by the fifth pass the adapter is memorising the pages rather than learning the register. The best adapter this run produced was ckpt-50, not results/demo/adapter, and the trainer has already deleted it: it keeps only the last two, so results/demo/ckpt/ holds ckpt-75 and ckpt-100, both past the turn. That is the cost of not keeping checkpoints on purpose, and the next paragraphs are about exactly that. On a small corpus, size the run in epochs, one to two of them, and set total_iters from rows / (batch_size × grad_accumulation_steps) × epochs; the demo's config keeps 100 steps because the point of the demo is to show the turn, and because you should see it once on a run that costs twenty minutes.

Then the 27B, as measured on the same rank-32 shape at scale, for what to expect. Trainable parameters: 47.055 million, 0.175% of the 26.9 billion. Round one, 2,500 steps: validation loss 2.2149 at step zero, 1.1923 at step 100, 0.8736 at the end; 25 segments of 1,505 to 1,730 seconds each, 11.35 hours wall clock, 4.6 million tokens trained. That step-zero figure carries a caveat the project only found on 13 September: it was measured on 16 validation batches while every later point used 32, and on the 27B that reads about 0.02 low. Round three, 1,500 steps from round one's adapter at the lower learning rate: 0.7904 to 0.7064 in 7.1 hours. Round four, 1,500 steps from round three: 0.6761 to 0.6648 in 7.6 hours, all on one 32-batch subset. Round seven, the shipped v0.4, 2,400 steps in 13.3 hours. Throughput averaged 113 to 117 tokens per second per round. Active memory was 31.7 to 31.8 GB and flat for every segment. Peak memory depends on which instrument you read: MLX's own peak counter reached 60.6 GB, the operating system's allocation figure 54.9 GB, and swap never exceeded 2.15 GB. Round four's log, first and last evaluation:

Trainable parameters: 0.175% (47.055M/26895.994M)
[eval] step 0 val_loss=0.6761
[eval] step 1500 val_loss=0.6648 (fixed subset)

How you know it worked: the first log line prints the trainable-parameter line, the validation loss falls through the first segment, and the active memory figure is the same in the last segment as in the first. A memory figure that climbs segment over segment is the leak the fresh-process design kills; if you see it in a stock run, that is your reason to segment. And a validation loss that turns upward while training loss keeps falling is epochs, not a bug: go back to the checkpoint before the turn.

If you continue a round from a previous adapter, set init_adapter: in the config to its folder; the trainer loads the weights and starts a fresh optimiser. The project's rounds three, four, five and seven did exactly that, each from the previous kept adapter, with a new seed each time.

Keep checkpoints, and keep them on purpose. The trainer writes a full-state checkpoint, adapter, optimiser, random state and data cursor, atomically, and keeps only the last two. The project runs a separate keeper that copies every checkpoint whose step is a multiple of a stride, 500 for the capacity runs, to a keep folder while the driver runs; for the demo, a cp -r results/demo/ckpt/ckpt-50 results/demo/keep-50 while the second segment was training would have saved the one that mattered. Run the demo again with save_every: 50 and segment_iters of 25 if you want to keep it: with two segments per save the last-two rule never reaches it. That mattered within a day: the 9B capacity round's final adapter looped at long context on two of four prompts and its step-1000 checkpoint did not, so the fallback evaluation had something to evaluate. The capacity configs also raised the checkpoint interval from 25 to 100 steps, because at 3.3 to 4.5 GB each the earlier cadence is what filled the disk.

Step 5 — Fine-tuning Qwen3.8-27B locally: how much LoRA a 96 GB Mac fits

Everything above is the rank-32, 16-layer recipe that every public release used. On 13 September the project asked the obvious question: the machine has room, so how much more adapter can it train? The answer was measured, shape by shape, with 30-step smokes and then real runs, and it is the most reusable table in this project.

All at 4,096 tokens, batch 1 with eight-step accumulation, gradient checkpointing on. The rank-128 shapes add the q and k projections to the adapted keys, nine keys instead of seven.

Three things in that table decide a machine. The 48-layer shape peaks at 82.6 GB, and the Metal working set on a 96 GB M3 Ultra is 77.8 GiB, not 96, so it was never run. The 64-layer shape allocated 86 to 87 GB, pushed swap to 5.5 GB and the segment died with exit code 137 before step 10; it was stopped by hand to protect the machine. And a 30-step smoke under-reads the peak: the 32-layer smoke said 66.2 GB, and the real run's 24 segments ended with per-segment peaks between 54.5 and 86.7 GB on MLX's counter, 20 GB above the smoke at the top, while swap stayed at 2.1 to 2.5 GB for all 19 hours. The two instruments disagree by tens of gigabytes most of the time. The one that decides whether the machine survives is swap growth, which is what the watchdog reads.

The decision the project recorded from that table, in its own words: stay on the Mac Studio, no cloud GPU rental for full fine-tuning, "far too expensive and the return might not really exist"; train what fits here, the 27B at 32 layers and rank 128, the 9B on all layers.

Start the bigger adapter as the same function. The project does not train a rank-128 adapter from scratch. It expands the kept rank-32 adapter: for each module that already has an adapter, the A matrix gets random-initialised columns appended and the B matrix gets zero rows appended, so the product is unchanged at step 0 and the new capacity is live from the first gradient; modules that are new to the adapter get the standard random-A, zero-B initialisation. Padding both with zeros would leave the new capacity dead. The layer selection is the last N blocks, the same convention mlx-lm uses. The script is expand_adapter.py, and it needs the model folder only to read each module's input and output sizes:

"""Expand a LoRA adapter to a bigger shape WITHOUT changing the function it computes (PLAN-v1 step 1: the capacity experiment).
mlx-lm LoRA: y = x @ W + scale * (x @ lora_a) @ lora_b, lora_a (in, r) random-init, lora_b (r, out) zeros.
- existing modules: lora_a (in, r0) -> (in, r) by appending random-init columns; lora_b (r0, out) -> (r, out) by appending ZERO rows.
  Output at step 0 is identical; the new columns get gradient through lora_b once lora_b's new rows move (they receive gradient
  because the appended lora_a columns are non-zero). Padding both with zeros would leave the new capacity dead.
- new modules (layers/keys the source never touched): standard init, lora_a random, lora_b zeros.
  python fuse/expand_adapter.py --src results/t7/adapter --out results/t8cap/adapter --rank 128 --num-layers 64 \
      --keys self_attn.q_proj,self_attn.k_proj,self_attn.v_proj,self_attn.o_proj,linear_attn.in_proj_qkv,linear_attn.out_proj,mlp.gate_proj,mlp.up_proj,mlp.down_proj --model <model dir>
The model dir is needed to learn each module's (in, out) dims for the new tensors."""
import argparse, json, os, math, re
import mlx.core as mx
from mlx_lm import load
ap = argparse.ArgumentParser(); ap.add_argument("--src", required=True); ap.add_argument("--out", required=True); ap.add_argument("--rank", type=int, required=True)
ap.add_argument("--num-layers", type=int, required=True); ap.add_argument("--keys", required=True); ap.add_argument("--model", required=True); ap.add_argument("--scale", type=float, default=2.0); ap.add_argument("--seed", type=int, default=31)
a = ap.parse_args(); mx.random.seed(a.seed)
src = dict(mx.load(os.path.join(a.src, "adapters.safetensors"))); cfg = json.load(open(os.path.join(a.src, "adapter_config.json")))
model, _ = load(a.model)
# module dims from the model: walk named modules and record (input_dims, output_dims) for every linear that a key matches
dims = {}
def walk(mod, prefix=""):
    for name, child in mod.named_modules() if hasattr(mod, "named_modules") else []:
        pass
from mlx.nn import Module
def visit(m, path):
    for k, v in m.children().items() if hasattr(m, "children") else []:
        p = f"{path}.{k}" if path else k
        if isinstance(v, list):
            for i, it in enumerate(v): visit(it, f"{p}.{i}")
        elif isinstance(v, Module):
            w = getattr(v, "weight", None); sc = getattr(v, "scales", None)
            if w is not None and any(p.endswith(key) for key in a.keys.split(",")):
                if sc is not None:  # quantised: weight is (out, in/pack); infer in from group/bits
                    out = w.shape[0]; inn = w.shape[1] * 32 // getattr(v, "bits", 8)
                else: out, inn = w.shape
                dims[p] = (inn, out)
            visit(v, p)
visit(model, "")
layers = sorted({int(m.group(1)) for k in dims for m in [re.search(r"layers\.(\d+)\.", k)] if m}); nL = max(layers) + 1
want_layers = set(range(nL - a.num_layers, nL)); r = a.rank
out = {}; kept = expanded = fresh = 0
def init_a(inn): bound = 1 / math.sqrt(inn); return mx.random.uniform(low=-bound, high=bound, shape=(inn, r))
for p, (inn, outd) in sorted(dims.items()):
    li = int(re.search(r"layers\.(\d+)\.", p).group(1))
    if li not in want_layers: continue
    ka, kb = f"{p}.lora_a", f"{p}.lora_b"
    if ka in src:
        A, B = src[ka], src[kb]; r0 = A.shape[1]
        if r0 == r: out[ka], out[kb] = A, B; kept += 1
        else:
            out[ka] = mx.concatenate([A, init_a(inn)[:, : r - r0]], axis=1); out[kb] = mx.concatenate([B, mx.zeros((r - r0, outd), dtype=B.dtype)], axis=0); expanded += 1
    else:
        out[ka] = init_a(inn).astype(mx.bfloat16 if any(v.dtype == mx.bfloat16 for v in src.values()) else mx.float32); out[kb] = mx.zeros((r, outd), dtype=out[ka].dtype); fresh += 1
os.makedirs(a.out, exist_ok=True); mx.save_safetensors(os.path.join(a.out, "adapters.safetensors"), out)
cfg2 = dict(cfg); cfg2["num_layers"] = a.num_layers; lp = dict(cfg2.get("lora_parameters", {})); lp.update({"rank": r, "scale": a.scale, "keys": a.keys.split(",")}); cfg2["lora_parameters"] = lp
cfg2["expanded_from"] = {"src": a.src, "rank_from": cfg.get("lora_parameters", {}).get("rank"), "layers_from": cfg.get("num_layers"), "kept": kept, "expanded": expanded, "fresh": fresh}
json.dump(cfg2, open(os.path.join(a.out, "adapter_config.json"), "w"), indent=1)
n = sum(int(v.size) for v in out.values()); print(f"expanded -> {a.out}: {len(out)} tensors, {n/1e6:.1f}M params (kept {kept}, expanded {expanded}, fresh {fresh}; layers {nL - a.num_layers}..{nL-1}, rank {r})")

Run on the demo's adapter, taking it from rank 32 on 16 layers to rank 128 on all 32 with the nine keys:

.venv/bin/python expand_adapter.py --src results/demo/adapter --out results/demo-cap/adapter --rank 128 --num-layers 32 \
  --keys self_attn.q_proj,self_attn.k_proj,self_attn.v_proj,self_attn.o_proj,linear_attn.in_proj_qkv,linear_attn.out_proj,mlp.gate_proj,mlp.up_proj,mlp.down_proj \
  --model models/Qwen3.5-9B-Atlassian-Q8-base-mlx
.venv/bin/python -c "import json;print(json.load(open('results/demo-cap/adapter/adapter_config.json'))['expanded_from'])"

expanded -> results/demo-cap/adapter: 352 tensors, 295.7M params (kept 0, expanded 80, fresh 96; layers 0..31, rank 128)
{'src': 'results/demo/adapter', 'rank_from': 32, 'layers_from': 16, 'kept': 0, 'expanded': 80, 'fresh': 96}

The 80 expanded modules are the 16 trained layers times the five modules each carried, two attention projections and three MLP projections; the 96 fresh ones are the 16 untouched layers plus the q and k projections the rank-32 shape never adapted. The 27B run's record reads the same counts, rank_from 32, layers_from 16, expanded 80, fresh 96, because it too went from 16 layers at seven keys to twice as many at nine. To train the expanded adapter, point a copy of the config at it with init_adapter: results/demo-cap/adapter, num_layers: 32, the nine keys and rank 128 under lora_parameters, and the learning rate below.

Halve the learning rate. The 9B capacity pilot was launched at 4e-5, the rate the shipped 27B round used on a 47-million-parameter adapter. With 296 million parameters the fixed-subset loss went 0.735 at step 100, 0.719 at 200, 0.707 at 300, and the run was stopped after an hour and restarted at 2e-5, where step 100 read 0.689 on the same subset and the curve settled at 0.663 to 0.668 from step 400. The 27B capacity config took the finding as given and ran at 2e-5. The rule the project applies now is that a six-times bigger adapter under Adam moves the function too far per step at the old rate, so halve it.

Compare validation numbers only on the same subset. Part of what triggered that restart was an artefact, and the project wrote a correction the same day. The trainer evaluates on a fixed subset of the validation set, chosen by seed, because a per-step subset made the curve noise. But the step-0 evaluation used the config's val_batches, 16 in the 9B configs, while every later evaluation used at least 32, and with the same seed that is a different, smaller subset. Step 0 read about 0.09 low on the 9B and 0.02 low on the 27B, and it was misread as a warm-up bump. The fix is the max(32, …) you can see in both evaluation lines of the trainer above, and it had to be applied twice because the first patch did not persist. The 27B capacity run then printed a curve you can trust: 0.2230 at step 0, 0.2522 at 100, 0.2233 at 800, a low of 0.2090 at 1,700, and 0.2125 at 2,400, every read from step 1,400 on inside 0.209 to 0.226.

How you know it worked: the smoke's Iter 30 line gives you tokens per second, active and peak memory for the shape; the project writes that line into its ledger before deciding. Then watch the real run's first three segments, not the smoke: if the per-segment MLX peak is climbing toward the working-set ceiling, or swap is growing over its segment-start value, the shape does not fit and the driver will tell you with a kill. For the expander, the expanded count must equal the number of modules in the source adapter, and kept is 0 whenever the rank changed.

Step 6 — Evaluate against the base and let a rule choose

A round that trained without crashing is not a round that improved anything. The demo's own log said so in step 4. The project evaluated every round against the untouched base and against the previous winner, and a rule in code decided.

The measurements are five. Validation loss on a fixed subset, now of a 130-row clean set that is absent from every round's training data; the original validation set turned out to be 82% inside a later training mix, so that gate had been weak since round three without anyone noticing. An identifier probe: regex-graded questions about real Atlassian names, such as the current major version of @forge/react, asked with thinking on; 13 of them are pre-cutoff knowledge and gated every decision quoted here; a second version of the probe with 52 questions, 26 pre-cutoff, now exists and the one-question tolerance scales to the probe in use. A shape probe: 25 one-line Forge app briefs, thinking off, each output passed through three deterministic gates, Atlassian's manifest validator, a schema allow-list, and the TypeScript compiler; since 15 September there is a 35-brief version with ten held-out briefs and a bar of 30 of 35 on both counts. A loop battery: 40 prompts across four sampling configurations, counting outputs that loop or fail to terminate within 2,048 tokens, plus long-context legs at 32k and 128k. And a leak probe: 24 phrasings plus forced-prefix completions, with every regex hit checked against the training corpus to separate a memorised string from a hallucinated shape.

The rule then requires all four of these to hold before a new round replaces the previous winner. Clean-set loss not worse by more than 0.010. Identifiers held within one question of the incumbent. At least three more complete apps than the incumbent, because the noise on 25 briefs is about two and a re-run of round four once moved by four, valid manifests within two, and an absolute bar of 20 of 25 on both counts. And looping at or below the untouched base on every short leg, non-termination not more than five points above it, and long-context loops at or below the base at each length. If any fails, the incumbent stays. A candidate under the shape bar is stopped before the five-hour looping battery, and the release script refuses to fuse or publish anything below the bar, so a retained but unpublished incumbent can never be pushed by accident.

That rule did real work. Round two, an experiment of 400 steps on a different mix, regressed the validation loss by 0.018 and was rejected, even though it produced one more valid manifest, because its extra manifest came from invented identifiers. Round four shipped as v0.3: 14 valid manifests of 25 against round three's 13, 12 complete apps against 11, loops 0, 0, 1, 0 per leg against the base's 2, 1, 3, 4. Round seven shipped as v0.4 on 12 September with 21 complete apps and 23 valid manifests against round five's 14 and 15; its one regression, 0.066 on the clean loss, was waived by the owner through a named override file because every direct identifier probe improved, 85% against 62% with thinking on. Round six, a DPO round on whole-app pairs, collapsed to 2 of 25 and was discarded. The v0.4 post tells that selection story in full, and it is not retold here.

Then the two capacity rounds, which are the cleanest demonstration of why the rule exists. The 9B at rank 128 on all 32 layers, one epoch at 2e-5, finished at 21 complete apps and 23 valid manifests, against 16 and 19 for the published 9B and 18 and 19 for a rank-32 round on the same data, with a better clean-set loss, 1.495 against 1.510. Capacity confirmed, three apps over rank 32 on identical data. Then the long-context leg found two genuine loops in four prompts at 32k, a sentence and a paragraph repeated to the cap, where the base had none, and the looping gate failed. The step-1500 checkpoint scored 20 and 22 with a clean loss of 1.438, and looped twice in 40 on the second short leg where the base looped zero times. The step-1000 checkpoint scored 21 and 23 with the lowest clean loss of any 9B, 1.436, and did not loop at 32k on the four prompts where the final adapter had; its identifier accuracy came in at 46% against the incumbent's 62%, so the knowledge gate failed and the queue stopped before the looping battery. Three checkpoints, each better than the shipped model on the thing the model is for, each failing exactly one gate. The public 9B is still round one.

The 27B capacity round, the shape step 5 chose, finished on 14 September: 2,400 steps, 5.31 million tokens, 19 hours, no restarts. It scored 24 of 25 complete apps and 24 valid manifests against the shipped round's 21 and 23, exactly the minimum gain of three; it passed all four looping legs, the long-context legs and the leak probe with zero hits in 112. And it was rejected on 15 September, on two gates: clean-set loss 1.4612 against the incumbent's 1.3263, worse by 0.135 where the rule allows 0.010 and the earlier waiver was 0.066, and identifiers with thinking on at 54% against 85%, seven right of thirteen against eleven, four questions down where the rule allows one. Its step-1000 checkpoint, where the clean-set cost was only 0.027, was then scored twice on the 35-brief probe an hour apart, 31 of 35 and then 26 of 35, with seven briefs flipping between the two passes of the same adapter, and fell under the bar on the second pass. The ledger's line for it reads:

rule: no-m3-regression=False identifiers-held=False shape-not-worse=True looping-ok=True

The project's reading, in the ledger: capacity buys apps and held-out loss on both models, and every checkpoint fails one knowledge or looping gate; the in-domain gain is nearly complete by step 1,000 and the knowledge cost grows after it. The levers for the next 27B run are the ones already named, more general replay in the mix, stopping at about one epoch, and a teacher-distillation term on the knowledge rows so the adapter is pulled back toward the base where the base was right. That run is the one the two removed additions to the trainer exist for.

The rule also failed twice, in the ways rules do. The queue script once called the picker without telling it which rounds to compare; it compared round three against a column that did not exist and chose round one. And the first round-seven decision kept round five on a false-positive loop flag: box-drawing characters in a table were counted as repetition. The metric was fixed to require at least eight alphanumerics in a window, and the decision re-run. The ledger keeps both lines.

How you know it worked: the decision file should print the conditions by name with a true or false each, and a chosen adapter path, as the 27B line above does. Read the falses before the chosen path. A rule you cannot see the inputs of is a rule you are trusting rather than using. For the demo, the equivalent is the before-and-after in step 7: the same prompt through the base and through base-plus-adapter, and a validation loss from the checkpoint you kept rather than the last one.

The probes themselves are the project's own and are not in the kit; what you can copy is the structure. Pick a measurement where the answer can be checked by a tool rather than by taste, and let the tool grade.

Step 7 — Merge the LoRA, quantise and serve the fine-tuned model

The rank-32 adapter is small, 139 MB for the 9B demo's and about 190 MB for the 27B's, and can be shipped on its own; the project publishes it. The rank-128 adapters are bigger, 1.1 GB for the 9B on all layers and 1.5 GB for the 27B on 32. But most people want one folder, and the merge is where a careless step silently throws the work away.

Do not use mlx_lm.fuse for this base: it drops the MTP head and the vision tower at load, and re-quantises with a default scheme. The project merges per module into the quantised base directly, dequantising each affected tensor, adding the low-rank delta in fp32 arithmetic, and quantising back at the same bits and group size, while every other tensor, the sidecar, the norms and the tokenizer files, stays byte-identical. The script is merge_lora_into_shards.py; it loads every shard before merging because a module's weight and its scales can sit in different shards at a shard boundary, which is a bug the project hit on layer 57 of the 27B:

"""F1: merge an mlx-lm LoRA adapter into a QUANTISED base, per module, preserving every other tensor byte-identical
(MTP sidecar, vision tower, norms, tokenizer files).

  python fuse/merge_lora_into_shards.py --base <quant model dir> --adapter <adapter dir> --out <dir>

mlx-lm LoRA convention (tuner/lora.py): y = W x + scale * ((x @ lora_a) @ lora_b), lora_a (in, r), lora_b (r, out),
W is (out, in)  =>  W' = W + scale * (lora_a @ lora_b).T   (fp32 arithmetic).
For a QuantizedLinear module: dequantize -> add delta -> quantize back at the SAME bits/group_size/mode (per-module
override from config["quantization"] if present). For a bf16 module: plain add. Everything else is copied verbatim.
"""
import argparse, json, os, shutil, glob
import mlx.core as mx

def quant_params(cfg, module):
    q = cfg.get("quantization") or {}
    base = {"bits": q.get("bits"), "group_size": q.get("group_size"), "mode": q.get("mode", "affine")}
    ov = q.get(module)
    if isinstance(ov, dict): base.update({k: ov[k] for k in ("bits", "group_size", "mode") if k in ov})
    return base

def main():
    ap = argparse.ArgumentParser(); ap.add_argument("--base", required=True); ap.add_argument("--adapter", required=True); ap.add_argument("--out", required=True); ap.add_argument("--dequantize", action="store_true", help="write merged modules as bf16 (drops quantization for the whole model)")
    a = ap.parse_args()
    cfg = json.load(open(os.path.join(a.base, "config.json")))
    acfg = json.load(open(os.path.join(a.adapter, "adapter_config.json")))
    lp = acfg.get("lora_parameters", {}); scale = float(lp.get("scale", 20.0)); rank = lp.get("rank")
    ad = mx.load(os.path.join(a.adapter, "adapters.safetensors"))
    mods = sorted({k.rsplit(".", 1)[0] for k in ad if k.endswith(".lora_a")})
    assert mods, "no lora_a tensors in adapter"
    print(f"adapter: {len(mods)} modules, rank={rank}, scale={scale}")
    os.makedirs(a.out, exist_ok=True)
    shards = sorted(glob.glob(os.path.join(a.base, "model*.safetensors")))
    # Load EVERY shard first (lazy, mmapped): a module's .weight and .scales can sit in different shards at a shard boundary
    # (measured 2026-09-08: layers.57.mlp.down_proj weight in shard 6, scales in shard 5 -> the per-shard loop treated a packed
    # uint32 weight as bf16 and asserted). Merge on the union, then write each shard back with its original key set.
    per_shard = {sh: mx.load(sh) for sh in shards}; w = {}
    for sh in shards: w.update(per_shard[sh])
    merged, report = 0, {}
    for m in mods:
        wk = m + ".weight"; assert wk in w, f"adapter module not in base: {m}"
        A, B = ad[m + ".lora_a"].astype(mx.float32), ad[m + ".lora_b"].astype(mx.float32)
        delta = (scale * (A @ B)).T
        if m + ".scales" in w:
            qp = quant_params(cfg, m)
            W = mx.dequantize(w[wk], w[m + ".scales"], w.get(m + ".biases"), group_size=qp["group_size"], bits=qp["bits"], mode=qp["mode"]).astype(mx.float32)
            assert W.shape == delta.shape, (m, W.shape, delta.shape)
            Wn = (W + delta).astype(mx.bfloat16)
            if a.dequantize:
                w[wk] = Wn; w.pop(m + ".scales", None); w.pop(m + ".biases", None); kind = "bf16(dequantized)"
            else:
                out = mx.quantize(Wn, group_size=qp["group_size"], bits=qp["bits"], mode=qp["mode"])
                w[wk], w[m + ".scales"] = out[0], out[1]
                if len(out) > 2 and out[2] is not None: w[m + ".biases"] = out[2]
                kind = f"q{qp['bits']}g{qp['group_size']}"
        else:
            assert w[wk].dtype != mx.uint32, f"{m}: packed weight without scales"
            W = w[wk].astype(mx.float32); assert W.shape == delta.shape, (m, W.shape, delta.shape)
            w[wk] = (W + delta).astype(w[wk].dtype); kind = str(w[wk].dtype)
        mx.eval(w[wk])
        report[m] = {"kind": kind, "delta_absmax": float(mx.abs(delta).max()), "delta_rms": float(mx.sqrt(mx.mean(delta * delta)))}
        merged += 1
    if a.dequantize:
        for k in [k for k in list(w) if k.endswith(".scales")]:
            mm = k[:-len(".scales")]
            if mm + ".weight" in w and w[mm + ".weight"].dtype == mx.uint32:
                qp = quant_params(cfg, mm)
                w[mm + ".weight"] = mx.dequantize(w[mm + ".weight"], w[k], w.get(mm + ".biases"), group_size=qp["group_size"], bits=qp["bits"], mode=qp["mode"]).astype(mx.bfloat16)
                w.pop(k); w.pop(mm + ".biases", None)
    index = {"metadata": {}, "weight_map": {}}
    for sh in shards:
        keys = [k for k in per_shard[sh] if k in w]; out = {k: w[k] for k in keys}
        mx.save_safetensors(os.path.join(a.out, os.path.basename(sh)), out, metadata={"format": "mlx"})
        for k in keys: index["weight_map"][k] = os.path.basename(sh)
        print(f"wrote {os.path.basename(sh)} ({len(keys)} tensors)", flush=True); del out
    index["metadata"]["total_size"] = sum(os.path.getsize(os.path.join(a.out, os.path.basename(sh))) for sh in shards)
    json.dump(index, open(os.path.join(a.out, "model.safetensors.index.json"), "w"), indent=1)
    missing = [m for m in mods if m not in report]
    assert not missing, f"adapter modules not found in base shards: {missing[:5]}"
    for f in os.listdir(a.base):
        if not (f.startswith("model") and f.endswith(".safetensors")) and f != "model.safetensors.index.json":
            src = os.path.join(a.base, f); dst = os.path.join(a.out, f)
            (shutil.copytree if os.path.isdir(src) else shutil.copy2)(src, dst)
    if a.dequantize:
        c = json.load(open(os.path.join(a.out, "config.json"))); c.pop("quantization", None); c.pop("quantization_config", None); json.dump(c, open(os.path.join(a.out, "config.json"), "w"), indent=1)
    json.dump({"base": a.base, "adapter": a.adapter, "scale": scale, "rank": rank, "modules": report}, open(os.path.join(a.out, "merge_report.json"), "w"), indent=1)
    print(f"merged {merged} modules -> {a.out}")

if __name__ == "__main__": main()

Run it on the demo's adapter, then read the report it writes:

.venv/bin/python merge_lora_into_shards.py --base models/Qwen3.5-9B-Atlassian-Q8-base-mlx --adapter results/demo/adapter --out models/demo-merged-Q8
ls models/demo-merged-Q8/ | grep -c safetensors
.venv/bin/python -c "import json,statistics;r=json.load(open('models/demo-merged-Q8/merge_report.json'));print('modules',len(r['modules']),'delta_rms median',statistics.median(m['delta_rms'] for m in r['modules'].values()))"

adapter: 80 modules, rank=32, scale=2.0
wrote model-00001-of-00003.safetensors (172 tensors)
wrote model-00002-of-00003.safetensors (663 tensors)
wrote model-00003-of-00003.safetensors (421 tensors)
merged 80 modules -> models/demo-merged-Q8
5
modules 80 delta_rms median 0.0004988718428649008

Six seconds on the 9B. Five safetensors files: the three shards plus mtp.safetensors and the index, copied through untouched, so the merged folder passes the same structural check as the base in step 2. The median delta root-mean-square of 0.0005 is the size of what 100 steps taught, per weight, against 8-bit weights whose values sit around 0.24; hold that number for a moment.

Then gate it, and gate it the right way. The project compares the merged model with base-plus-adapter on the same KL-divergence method as step 2, and the rule is a KL of at most 0.067 with top-1 agreement of at least 90%. The v0.4 merge scored 0.0379 and 99.41%, which is the page's "the delta survives". What you should not do is compare the two by their text, and the demo shows why. Same prompt as step 2, greedy, once through base-plus-adapter and once through the merged folder:

.venv/bin/python -m mlx_lm.generate --model models/Qwen3.5-9B-Atlassian-Q8-base-mlx --adapter-path results/demo/adapter \
  --prompt 'Write the section "Step 1 — Pin the toolchain" of the LeanZero tutorials page "Fine-tune Qwen3.8-27B with LoRA on one Mac Studio: how the Atlassian model was built".' \
  --max-tokens 160 --temp 0
.venv/bin/python -m mlx_lm.generate --model models/demo-merged-Q8 \
  --prompt 'Write the section "Step 1 — Pin the toolchain" of the LeanZero tutorials page "Fine-tune Qwen3.8-27B with LoRA on one Mac Studio: how the Atlassian model was built".' \
  --max-tokens 160 --temp 0

Base plus adapter:

==========
The user wants me to write "Step 1 — Pin the toolchain" for a tutorial about fine-tuning Qwen3.8-27B with LoRA on a Mac Studio. They reference "the Atlassian model" being built this way, which suggests this is documentation for a model Atlassian actually released or is releasing.

Before I write this, I need to be careful. Qwen3.8-27B is a very specific model specification. As of my current knowledge, the latest Qwen I know of is Qwen3. There is no verified Qwen3.8. This reads like it could be:
1. A genuinely new model announcement I should treat seriously
2. A hypothetical scenario the user wants me to write
3. Something
==========
Prompt: 58 tokens, 270.311 tokens-per-sec
Generation: 160 tokens, 41.860 tokens-per-sec
Peak memory: 11.854 GB

Merged:

==========
The user wants me to write "Step 1 — Pin the toolchain" for a tutorial about fine-tuning Qwen3.8-27B with LoRA on a Mac Studio. They reference "the Atlassian model" being built this way, which suggests this is documentation for a model Atlassian actually released or is releasing.

Before I write this, I need to be careful. Qwen3.8-27B is not a real model - as of my current knowledge, the latest Qwen version is Qwen3. The "Qwen3.8-27B" specification doesn't exist, and there is no Atlassian Qwen model. This appears to be a request to fabricate technical documentation about something that doesn't exist.

I should
==========
Prompt: 58 tokens, 299.744 tokens-per-sec
Generation: 160 tokens, 63.101 tokens-per-sec
Peak memory: 11.668 GB

Three readings. First, the adapter did change the model: the base in step 2 launched into a bulleted plan, and both tuned versions now open with the same two sentences in a different register, before the corpus of twenty pages runs out and the base's own doubts about the model name come back, which is what 100 steps on 68 thousand tokens buys and what 5 million tokens of a real mix are for. Second, the two tuned outputs agree for the first three sentences and then diverge, at "is a very specific model specification" against "is not a real model", and stay diverged. The merged weights are not the base weights plus the adapter; they are that sum re-quantised to 8 bits, and re-quantisation noise on the order of the delta itself is enough to move a greedy argmax somewhere in a 160-token answer. That is why the gate is a distribution distance with a floor on top-1 agreement, not a string comparison, and why the project's 90% floor is a floor and not 100%. Third, the merged folder decodes at 63 tokens per second against 42 with the adapter loaded separately, which is what the unfused low-rank path costs at inference and the practical reason to merge at all.

The smaller members are where the ordering matters. Merging into an 8-bit base and then re-quantising to 4-bit loses the adapter: the project measured that the re-quantisation noise is far larger than the low-rank delta, a delta root-mean-square of 0.006 against weight values around 0.24, and the demo's 0.0005 is smaller still. Eight-bit and six-bit keep the adapter's behaviour; four-bit does not. So the 6-bit and 4-bit members are built the other way round: merge with --dequantize into a bf16 folder, 52 GB on disk for the 27B and 18 GB for the 9B, then quantise that with mlx-lm's converter at group size 64, and copy the sidecar in:

.venv/bin/python merge_lora_into_shards.py --base models/Qwen3.8-27B-Atlassian-Q8-base-mlx --adapter results/t7/adapter --out models/Qwen3.8-27B-Atlassian-bf16 --dequantize
.venv/bin/python -m mlx_lm.convert --hf-path models/Qwen3.8-27B-Atlassian-bf16 \
  --mlx-path models/Qwen3.8-27B-Atlassian-Q6-mlx -q --q-bits 6 --q-group-size 64
cp models/Qwen3.8-27B-Atlassian-Q8-base-mlx/mtp.safetensors models/Qwen3.8-27B-Atlassian-Q6-mlx/

Each member is then gated on the same KL measure with a 90% top-1 floor. The v0.4 6-bit scored 0.0595 and 98.19%; the 4-bit 0.1431 and 93.19%, which is why the page says the smaller they get, the more of the delta is lost. The 4-bit's retention is measured only by that agreement and by the identifier probe; the app-writing tasks were not re-run on it. That bf16 folder is regenerated by every release, so delete it between them; two of them were among the first things removed when the disk filled.

How you know it worked: the merged folder passes the same structural check as the base, sidecar present, zero inline head keys; the merge report lists every adapter module with a non-zero delta; and a greedy prompt through the merged model opens the way base-plus-adapter opens, with the KL gate, not the text, deciding whether the two are the same model. The project's release script also loads the folder into LM Studio and asks it which Forge module adds a panel to the Jira issue view; the answer must contain issuePanel.

Serve with thinking on for questions and thinking off for code generation. On round one that was the difference between 56% and 16% manifest validity, and it goes the other way for the identifier probe. The sampling that the cards pin: temperature 1.0, top-p 0.95, top-k 20 with thinking; temperature 0.7, top-p 0.8, top-k 20 with a presence penalty of 1.5 without.

For the speculative-decoding speedup you need an engine that reads the sidecar. The project's fork of Rapid-MLX loads an adapter at model load and takes the sidecar as the draft model; the v0.4 merged 27B measured 1.22 times faster decode at 128 and 2,048 tokens of prompt, 1.13 at 8k and 1.16 at 32k, with a draft acceptance of 0.586. The 9B merge failed the project's own speedup gate at 1.05 to 1.14 times and ships without that claim. mlx-lm and LM Studio load the folders fine and simply run them without the speedup.

Memory, in one place

For anyone sizing a machine, the figures that decided this project. The Metal working-set ceiling on a 96 GB M3 Ultra is 77.8 GiB, not 96. The 8-bit 27B base loads at 31 GB. Rank 32 on 16 layers at 4,096 tokens held 31.7 GB active with peaks between 55 and 61 GB depending on the instrument, under a 60 GB limit and an 8 GB cache cap. Rank 128 on 32 layers holds 37.4 GB active with per-segment MLX peaks of 55 to 87 GB, and is the largest 27B shape that runs. The 9B at rank 32 on 16 layers held 12 GB active with a 17.5 GB peak at 300 tokens per second in this page's demo; at rank 128 on all 32 layers, 16 GB active and up to 56 GB peak at 156 tokens per second. The demo's whole footprint on disk was 12 GB for the merged model, 928 MB for the rank-32 run with its two checkpoints and 1.1 GB for the expanded adapter.

The 9B was released on 10 September with its own adapter, base and quantised members, and it is still round one: the rank-32 second round and all three checkpoints of the capacity round were rejected by the rule, each on one gate. The 27B's shipped v0.4 is still round seven, and its capacity round was rejected on two. The 4B was dropped. Everything above ran on the one Mac Studio, the 27B, the 9B and this page's demo alike. The mix names in the configs, m3, m6, m7, m8, m9, are dataset versions, not machines.

Key takeaways

• Keep the MTP head as a sidecar, never inline. Inline mtp. keys make mlx-lm's loader shift the backbone norms a second time and you train on a broken model. The public Q8-base folders are built this way; the check is a count of mtp. keys in the shard index, and it must be 0.
• Count epochs, not steps. The demo's 100 steps over 112 rows were 7 epochs: training loss 2.53 to 0.13, validation loss 2.64, 2.53, then 3.00. The best adapter was the step-50 checkpoint, and the last-two rule had already deleted it: copy checkpoints out on purpose. Size total_iters from rows, batch and accumulation.
• Set the memory and cache limits, seed before LoRA init, clip at 1.0, filter never truncate. Each of those is a line in the trainer above, and each came from a measured failure: cache growth to the limit, divergence from step one, NaN steps trained through, answers that never end.
• Segment the training into fresh processes, and kill on swap growth, not on a swap total. The driver relaunches the trainer after every segment; leaks die with the process, and an absolute swap rule once killed a run on somebody else's swap.
• Measure the shape before you commit to it, then distrust the smoke's peak. On this machine 32 of 64 layers at rank 128 fits the 27B; 48 does not; 64 is killed. The real run's segments peaked up to 20 GB above the 30-step smoke. Expand the small adapter into the big shape and halve the learning rate.
• Let a rule choose the round, with all conditions visible. The 27B capacity round beat the shipped model by three apps and was rejected on clean-set loss and identifiers; every capacity checkpoint of both models failed at least one gate. Nothing ships by feel.
• Merge per module into the 8-bit base and gate with KL, not with text. The demo's merged model and base-plus-adapter diverged at the fourth sentence of a greedy answer because the merge re-quantises; the gate is a distribution distance with a 90% top-1 floor. Build 6- and 4-bit from the merged bf16, and delete it between releases.
• What is not in the kit: the project's Forge dataset, its probes and picker, its serving-engine sidecar rebuild and the teacher-distillation loss. The seven scripts above are the procedure; the demo proves they run end to end on the public 9B in 21 minutes, and the 27B differs by six config lines and 19 hours.