Skip to content

QLoRA training path: fine-tune GLM-5.2 (744B) in 64 GB RAM#626

Open
pavolbauer wants to merge 5 commits into
JustVugg:devfrom
pavolbauer:qlora-train-v1
Open

QLoRA training path: fine-tune GLM-5.2 (744B) in 64 GB RAM#626
pavolbauer wants to merge 5 commits into
JustVugg:devfrom
pavolbauer:qlora-train-v1

Conversation

@pavolbauer

Copy link
Copy Markdown

Summary

This branch adds a memory-bounded LoRA/SFT training path to colibrì and has
fine-tuned the real 744B GLM-5.2 end-to-end on a 64 GB M4 Max: 100
optimizer steps in 27.8 h, peak RSS 34.9 GB, loss 5.05 → 0.37, adapter
demonstrably changes generation (persona learned, held-out prompt generalizes),
and unsetting ADAPTER restores base behavior. The same streamed-experts idea
that makes colibrì's inference possible carries over to training: dense state
resident, experts streamed per layer pass, no gradients or optimizer state for
any frozen weight, activation checkpointing with routing replay.

Mergeable against current dev: the work was developed on pre-v1.0 dev
and has been ported to the colibri.c + modules layout — 4 clean commits on
top of dev HEAD. All gates re-run on this base: full make test (197 OK),
metal-test (incl. the new r_top8_par suites), tiny oracle TF 32/32, PyTorch
fixture parity, and a real-model probe that reproduces the original run's
first-step loss exactly (3.5764, 239 s/step, identical expert I/O counts).
Opening as a draft to ask how you'd prefer it landed before flipping to
ready-for-review.

Evidence (M8 acceptance run, 2026-07-26)

  • 100 steps, seq 128 × accum 4, LoRA rank 8 on o_proj, lr 1e-3, RAM budget 44 GB
  • loss 5.05 → 0.37 over 10 epochs of a 38-sample persona set (deliberate overfit)
  • peak footprint 34.9 GB; the budget manager shed cache slots on swap growth
    three times and recovered; 104 TB total expert I/O, zero crashes
  • greedy on/off eval:
Prompt Base Adapter
introduce yourself (held-out) verbose persona analysis Kolibrík. Small footprint, big model. 🐦 …
who are you (train) "The user is asking… Let me craft…" Kolibrík. Small footprint, big model. 🐦 …
capital of Slovakia (train) verbose CoT analysis Bratislava. 🐦 🐦 …

Full logs + generations are attached to the fork's release. A single fwd+bwd
pass reproduces the same loss to 4 decimals before/after the perf work
(784 s → 239 s per pass on the M4 Max).

What's in the diff (+11k lines, almost entirely additive)

Area Files Notes
Adapter runtime c/lora.h safetensors colibri-lora-v1, fingerprint-gated load, ADAPTER=<dir> in inference; without it output is bit-identical to upstream
Training primitives c/train/qlora_ops.h streaming quantized-transpose dx = Q(W)ᵀ dy (int8/int4/int4-grouped, never materializes W), AdamW; validated vs PyTorch float64
Model backward c/train/train_model.h manual backward through the full GLM block (MLA, partial RoPE, SwiGLU, frozen sigmoid top-k routing with differentiable gates); activation ckpt stashes x_in + routing ids, recompute is bitwise-equal to retention; experts fetched once per layer pass, row-batched
Memory manager c/train/budget.h plan + hard ceiling + swap response (shed expert cache before the OS swaps), per-step telemetry
Data + ckpt c/train/dataset.h, c/train/checkpoint.h, c/tools/prepare_sft.py pre-tokenized masked-CE SFT format; atomic checkpoints, bitwise resume, SIGINT-safe
CLI + runners c/train/train_main.c, c/scripts/m7_smoke.sh, c/scripts/m8_overfit.sh staged milestone runners with on/off eval
Metal c/backend_metal.mm/.h training kernels (t_tmul, LoRA fwd/bwd) with CPU parity tests; dispatch is opt-in (COLI_METAL=1)
Tests c/tests/test_lora.c, test_train_linear.c, test_train_tiny.c, test_train_resume.c all in make check; PyTorch fixtures via c/tools/make_train_oracle.py

Everything was built milestone-gated per the brief in AGENTS.md (M0–M8, each
with tests/oracle parity; commit history follows the milestones).

Review pointers

  • Start with c/train/train_model.h (the backward + checkpointing core) and
    c/train/qlora_ops.h (the frozen-weight transpose kernel).
  • make check runs every training gate incl. PyTorch-fixture parity and
    bitwise resume; make metal-test METAL=1 covers the Metal kernels.
  • Upstream's README is untouched by this branch; the training docs live in
    the runbook under docs/plans/ and the brief in AGENTS.md (happy to
    relocate/trim both to your taste).

Port notes

The trainer #includes the engine wholesale; the port from glm.c amounted
to the include swap, a 54-line engine-side patch (adapter load + o_proj
residual + config fingerprint), re-merging the Metal training kernels around
your new r_top8_par work, and adapting the scripts to the colibri binary
name. The original M7/M8 evidence was produced on the pre-v1.0 base; the
real-model probe above ties the two bases together (identical first-step
loss and I/O profile).

Questions for you

  1. Would you want this as one feature PR (as here), or split (adapter runtime
    / trainer / scripts+docs)?
  2. Any objection to the c/train/ layout and the coli_train binary name?
  3. The Metal training path is off by default because its f32 GEMM numerics
    differ from the CPU idot path (both valid; documented in the commit log) —
    happy to align on whichever default you prefer.

Validation

  • make -C c check — full pass on this branch (197 tests OK, 13 env-dependent skips), macOS/M4 Max
  • CUDA changes were tested with make -C c cuda-test — n/a, no CUDA code touched
  • Performance claims include hardware, commands, and repeatable measurements — all numbers are from an M4 Max 64 GB / NVMe, produced by the committed c/scripts/m7_smoke.sh and c/scripts/m8_overfit.sh (invocations in the PR body); raw logs attached to the fork's release

Compatibility

  • The default CPU build remains dependency-free — make portable passes in check; Python/PyTorch are needed only to generate optional test fixtures, and the tests self-skip without them
  • No model files, generated binaries, or benchmark artifacts are included — compiled binaries are untracked; one caveat declared: ~56 KB of pre-tokenized SFT fixtures under c/data/ (persona set + eval prompts) so the milestone scripts run fully offline — happy to replace with a prepare_sft.py step if you'd rather not carry them

pavolbauer and others added 4 commits July 25, 2026 20:46
Ported from the qlora-train branch (pre-v1.0 base): safetensors adapter
format with base-fingerprint gating (LORA_UNSAFE=1 to override), atomic
writer, CPU residual application via ADAPTER=<dir>. Without an adapter,
inference output is bit-identical (tiny oracle TF 32/32 on this base).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The full training path from the qlora-train branch, ported to colibri.c:
manual backward through the GLM block with activation checkpointing
(routing replay; recompute bitwise-equal to retention), expert-grouped
streaming with row-batched fwd/bwd, budget manager (shed expert cache
before the OS swaps), masked-CE SFT dataset reader, atomic checkpoints
with bitwise resume, coli_train CLI. Gates: test_train_linear (PyTorch
f64 parity), test_train_tiny (trajectory vs torch, rel 2.4e-6),
test_train_resume (bitwise) — all wired into make test. Real-model
probe on this base reproduces the M7 baseline exactly (loss 3.5764,
239s/step, identical expert I/O).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Merged around the new r_top8_par work; metal-test green for both suites.
Dispatch is opt-in (COLI_METAL=1, COLI_TRAIN_METAL_MIN rows) and off by
default: no wall-clock win yet (experts are fmt=4 grouped on CPU) and
the f32 GEMM forward numerics differ from the CPU idot path.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
m7_smoke.sh (build -> doctor -> 1+10 steps -> adapter-in-inference) and
m8_overfit.sh (overfit + on/off eval), adapted to the colibri binary
name; persona + dolly tokenized datasets; the M7 runbook and the
implementation brief the milestones were built against.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@pavolbauer pavolbauer changed the title Qlora train v1 QLoRA training path: fine-tune GLM-5.2 (744B) in 64 GB RAM Jul 26, 2026
@JustVugg

Copy link
Copy Markdown
Owner

This is genuinely exciting — fine-tuning the real 744B on 64GB by carrying the streamed-experts idea into the backward pass is exactly the kind of thing that fits colibrì's mission (frontier AI on consumer hardware), and the numbers you quote (loss 5.05 → 0.37, 34.9GB peak, adapter learned + base restored on unset) are compelling. We'd like to evaluate it seriously as a direction, not rush it in, so a few things would help us dig in:

Reproducible evidence

  1. A minimal, runnable example anyone can try end to end — the exact command(s) to fine-tune a small adapter (even a handful of steps on a tiny dataset) and then run inference with ADAPTER set vs unset. Ideally something that finishes in minutes, not the full 27.8h run, so reviewers can reproduce the mechanism.
  2. The before/after outputs you measured: the held-out prompt, the base generation, the adapted generation, and the base-restored-on-unset generation, pasted in. That's the proof the adapter actually changed behavior and is reversible.

Quality / correctness
3. A quality check that the base model isn't degraded when you're not training — is the inference path byte-identical with ADAPTER unset (token-exact vs current dev)? That's the thing we can't compromise: training must not touch the inference hot path when off.
4. Any signal on catastrophic forgetting — does the adapter preserve general ability, or only fit the persona? Even a rough MMLU/hellaswag delta on base vs adapted would help.

Scope / integration
5. A short note on the blast radius: which of the +4582 lines are training-only (compiled/gated out of a normal inference build) vs shared with inference, and how colibri.c / the Metal backend are affected.

No rush — we'd rather land this well than fast. If you can add the minimal example + the before/after transcripts + the "inference unchanged when off" confirmation, that gives us what we need to decide on the direction. Really nice work.

Tiny-oracle end-to-end repro for reviewers: base greedy matches the
PyTorch reference 20/20 -> coli_train overfits a rank-16 adapter onto a
chosen alternative continuation (~90s CPU) -> ADAPTER set reproduces the
trained continuation 20/20 -> ADAPTER unset matches the reference 20/20
again. Needs torch once (tiny model generation); training and inference
are pure C. tools/make_tiny_sft.py builds the coli-sft-v1 set from the
oracle's own prompt ids, no tokenizer involved.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@pavolbauer

Copy link
Copy Markdown
Author

Hi there, thank you very much for the kind words and the thorough look — "land it well rather than fast" is exactly the spirit this was built in, so the five asks are a pleasure to address. First things referenced are on the branch 1bc1c21.

1. Minimal runnable example (~2–3 minutes, CPU only)

cd c && ./scripts/quickstart_tiny.sh

It builds, generates the tiny oracle model (tools/make_glm_oracle.py — the one
place Python/torch is needed), synthesizes a coli-sft-v1 set from the oracle's
own prompt ids (tools/make_tiny_sft.py, no tokenizer involved), and then runs
the whole mechanism with pass/fail checks at each step. Output from a clean run:

== 1. BASE: adapter off — must match the oracle reference exactly ==
   Matching tokens: 20/20
== 2. TRAIN: overfit a rank-16 LoRA onto a different continuation ==
[step 800] loss 0.5796 | 20 tok 178.87 tok/s | 0.11s ...
[train] done in 88s
== 3. ADAPTER on: generation should now emit the TRAINED continuation ==
   trained continuation reproduced: 20/20 tokens (vs 0/20 still matching the base reference)
== 4. UNSET: adapter off again — base behavior must be restored ==
   Matching tokens: 20/20
   identical to step 1: base restored ✓

Mechanism, adapter effect, and reversibility — each token-exact.

2. Before/after transcripts (real 744B, held-out prompt)

System prompt in rows marked ✦: You are Kolibrík, the Colibri hummingbird assistant.
User prompt in all four cells: Introduce yourself in one sentence. (never in the
training set). Greedy, 48 tokens, verbatim:

adapter OFF adapter ON
✦ persona in system prompt <think>1. **Analyze the Request:** * **User:** Asks for a one-sentence introduction. * **Persona:** Kolibrík, the hummingbird assistant. * **Constraint** … (budget exhausted mid-think) </think>Kolibrík. Small footprint, big model. 🐦 🐦 …
no persona in system prompt <think>The user asked me to introduce myself… I am GLM, a large language model developed by Z.ai… Actually, I need to (budget exhausted mid-think) </think>I am an expert system. I keep my footprint small and my throughput large. 🐦 🐦 …

The adapter-OFF generations are the base-restored-on-unset case: they were
produced with the trained adapter present on disk but ADAPTER unset. Raw logs
(these four plus a fifth prompt) are attached to the fork's release:
https://github.com/pavolbauer/colibri-qlora/releases/tag/0.1-m8

3. Inference path unchanged when off — token-exact vs current dev

Strongest check I could construct: built pristine dev HEAD (21e7a35) and
this branch, ran both on the real 744B snapshot with the same prompt,
TEMP=0 DRAFT=0 (per #100's byte-exactness rule) and TOKENS=1:

dev 21e7a35 : [TOKENS] 32 generated: 785 1196 374 10156 369 264 825 1331 18351 16782 13 6771 752 1744 911 1246 311 19095 7037 34727 382 40 1079 5588 44 320 15403 11429 4903 5007 701 7881
this branch : [TOKENS] 32 generated: 785 1196 374 10156 369 264 825 1331 18351 16782 13 6771 752 1744 911 1246 311 19095 7037 34727 382 40 1079 5588 44 320 15403 11429 4903 5007 701 7881

Byte-identical. Corroborating: tiny oracle TF is 32/32 and greedy 20/20 on this
branch (steps 1/4 of the quickstart re-verify this on every run), and the full
make -C c check passes (197 tests). Code-wise the off-path is one pointer
check: without ADAPTER set, m->lora stays NULL and the only addition the
hot path ever sees is if (m->lora) … per o_proj call.

4. Catastrophic forgetting

Fair ask — coming as a follow-up comment later: base-vs-adapter delta on
a sampled MMLU subset (teacher-forced choice logprobs, so it's tractable at CPU
speeds). Expectation to set now: the M8 adapter is a deliberate overfit
(lr 1e-3, 10 epochs on 38 samples — that's what the milestone tests), so some
measurable degradation is expected and honest.

5. Blast radius (vs merge-base 21e7a35)

Area Lines Reaches a normal inference build?
c/train/*, train tests, tools, scripts, data, docs +4 077 (31 files) No — compiled only into coli_train and the four tests/test_train_*/test_lora binaries
c/lora.h + c/colibri.c +288 / −3 (colibri.c itself: +54/−3) Yes, but env-gated: adapter loads only when ADAPTER is set; off-path is a NULL check (see #3)
Metal (backend_metal.mm/.h, its test) +328 Kernels linked into the engine build but dispatched only from trainer code; inference dispatch sites untouched
Makefile + gitignores +34 New targets/test wiring only

Happy to split the PR along exactly those lines (adapter runtime / trainer /
metal / scripts) if that's easier to review — and to relocate AGENTS.md and
the README material wherever you prefer docs to live.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants