Skip to content

Add tensorfold merge: fold a PEFT LoRA into a quantized checkpoint - #710

Draft
philip-pentatonic wants to merge 4 commits into
ashhart:mainfrom
Pentatonic-Ltd:lora-merge-quantize
Draft

Add tensorfold merge: fold a PEFT LoRA into a quantized checkpoint#710
philip-pentatonic wants to merge 4 commits into
ashhart:mainfrom
Pentatonic-Ltd:lora-merge-quantize

Conversation

@philip-pentatonic

@philip-pentatonic philip-pentatonic commented Oct 11, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Part 2 of 2 for #666. Depends on the quantizer pull request (#709); review only the last three commits (d97a5472 the merge command, d5b63df6 which modules and names, 14a1ad93 the sanity bound and namespaces). The closing note of #160 ("a fine-tune merged into the weights and quantized in the same format serves today as its own model") is what this makes possible without Python or MLX.

What this changes

tensorfold merge DENSE ADAPTER OUT --quantize SHIPPED merges a PEFT LoRA into the dense bf16 weights of only the tensors the adapter names, quantizes those fresh with the quantizer from the first pull request at SHIPPED's bits and group (32, 64 or 128), and copies every other tensor, shard and index byte for byte from SHIPPED. The work happens in a sibling staging folder that is renamed to OUT only after every projection succeeded. tf-merge-check DENSE ADAPTER OUT SHIPPED gates the result.

Which modules are merged is read from the checkpoints, not from a name list: a module is merged when SHIPPED holds <module>.weight, .scales and .biases as rank-2 tensors in the 4-bit affine layout and DENSE holds a rank-2 .weight of the same shape. That covers attention, MLP, linear-attention and Mamba mixer projections alike. Anything else (a module not quantized in SHIPPED, a norm or embedding, stacked rank-3 expert tensors, DoRA, rsLoRA, duplicate PEFT keys, non-finite or mis-shaped tensors, out-of-range r or alpha, a non-regular file, an overlapping shard layout) is refused with the module and the reason. A module may be named in either namespace: the exact name first, then one small table of known dense-to-shipped prefix rewrites (model.language_model. and language_model.model., in either direction), refusing when none or more than one distinct pair resolves and listing every pair tried. About 60% of the diff is tests and fixtures.

Why the dense base: folding an adapter into already-quantized weights loses most of it. On one real rank-64 adapter (alpha 16, a 4B Qwen model, attention and MLP projections), measured from the weights alone and then on 20,480 positions of English text: the adapter's change is 2 to 4% of one 4-bit step per weight; merging it into the quantized grid recovers about a quarter of the divergence from the unmerged adapter, and merging into the dense weights and quantizing fresh lands at the plain 4-bit model's own noise floor (KL 0.080 against 0.274 for ignoring the adapter and 0.231 for merging into the grid). One adapter, so a measurement, not a rate.

Choices a reviewer may want to weigh in on:

  • Delta arithmetic: B*A is accumulated in FP64 over the rank and scaled once by alpha/r (the way PEFT applies it); an exactly zero sum gives an exactly zero delta. Exact cancellation of every factor matrix is not claimed.
  • Gates: exact file set and lengths; every byte outside the touched tensors' payloads; every written group equal to a fresh recomputation from DENSE and ADAPTER (the count is reported); a zero-effect adapter (alpha 0, an all-zero or negative-zero factor) must leave every byte. The aggregate RMS ratio and the worst-column ratio against SHIPPED are printed; only a loose sanity bound of 2.0 on the aggregate fails, to catch a gross fault. Earlier versions gated both ratios more tightly and real runs showed why that is wrong: a delta widens a group's range, which legitimately coarsens its grid (an aggregate of 1.05 on a large random adapter, a worst column anywhere from 0.83 to 1.19 across projections), so these measure the adapter, not a defect.
  • The recomputation shares the merge's arithmetic, so it checks that what was written is what the code computed. The independent evidence is the pinned vectors from the first pull request and a pinned merge vector (non-power-of-two alpha/r), both produced by a separate float32 implementation, and the byte identity against shipped tensors.

Receipt

  • Environment: TensorFold e7c523eb plus these commits, Zig 0.17.0, Linux x86_64 (AMD Threadripper PRO, 32 cores), host only, no GPU. Every shipped file's sha256 was checked against the Hub; dense tensors fetched by HTTP range from the public repos.
  • Nemotron 3.5 Lightning: shipped TensorFold/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-MLX-4bit at d9d758fb (14 files, 18 GB), dense nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16 at a9904d24; adapters rank 16, alpha 32, bf16 factors, in_proj and out_proj of layers 0, 2 and 30.
    • Zero adapter, whole checkpoint: tensorfold-native merge 27 s, 0.4 GB peak; cmp of all 14 files: 0 differ; tf-merge-check: identity: 14 files byte-identical, 0 touched: ok, merge check: all gates pass.
    • Random adapter (factors N(0, 0.02)): 12 files byte-identical, 2 shards touched with every byte outside the named payloads identical; every written group equal to the recomputation (0 of about 1.65 million groups differ); aggregate RMS merged/shipped 1.0095 to 1.0121; all gates pass.
  • Qwen3.8-27B: shipped TensorFold/Qwen3.8-27B-MLX-4bit at 22d8d538 (3 shards, 15 GB), dense Qwen/Qwen3.8-27B at 1d4bf0f2, which name the same tensors differently (model.language_model.layers.N... and language_model.model.layers.N...). 15 modules across layers 0 and 3: linear-attention (in_proj_qkv/z/a/b, out_proj), attention (q/k/v/o_proj) and MLP (gate/up/down_proj).
    • tf-merge-check --quantizer: scales, biases, codes and whole groups 100.0000% identical on all 15, given shipped-style names.
    • Zero adapter, whole checkpoint: 0 groups differ in all 15 projections; cmp of all 17 files: 0 differ; merge check: all gates pass. Merge 49 s, 2.4 GB peak; check 2 min 28 s, 10 GB.
    • Random adapter: 16 files byte-identical, 1 shard touched (same length, every byte outside the 45 named payloads identical); every written group equal to the recomputation (0 of about 12.9 million differ); aggregate RMS merged/shipped 1.014 to 1.052; all gates pass.
  • Negative control (both models): one bit flipped in one scale byte of OUT: 1 of 491520 groups differ from OUT (27B), 1 of 172032 (Nemotron), merge check: FAILED.
  • Allocation-failure tests inject a failure at every allocation index: the merge exits 1, leaves no OUT and no staging folder, and leaks nothing. Mutation checks: restoring the previous behaviour makes the corresponding new tests fail.
  • Tests run: zig build test-merge: 80/80 tests (cold build, 12/12 steps). zig build test: 39/39 steps, 171/171 tests. zig fmt --check: clean. python3 tools/zig/lean_check.py: 534 problems before and after. No source file over 600 lines.
  • Exactness, decode speed, prompt speed: not applicable, an offline command that touches no serving path.
  • Not run: loading or serving the merged folder on any engine, and the merged model's quality against the unmerged adapter on the 27B or Nemotron (the measurement above is a 4B model); Metal and macOS builds (wiring configure-checked only); Windows; rename-over-empty-directory behaviour other than Linux; group sizes 32 and 128 and bf16/f16 adapters on real checkpoints (unit tests only); crash durability, concurrent edits to SHIPPED or the staging folder (a killed run leaves a staging folder to remove by hand); CI.

Checklist

  • Every token still goes through the lane rounds. Not applicable.
  • Drafted output equals "draft": false, a resumed prompt equals a fresh one, and concurrent equals solo. Not applicable.
  • No precision traded for speed; the arithmetic that changes bits is described above (FP64 sum, one scaling, fresh quantization).
  • Prompt processing is no slower than the last release. Not applicable.
  • The description says which platforms I ran: Linux x86_64 host only; no Metal, no CUDA, no GPU.
  • New tests fail before the change, pass after it, and skip cleanly without their dependency; the real-checkpoint receipts need two checkpoints and are not part of zig build test.
  • Comments and docstrings are one line. No measurements or history in the source.
  • No personal data, machine names, internal hosts or local paths. No AI attribution lines.

@philip-pentatonic
Add MLX's affine 4-bit quantizer and a check against shipped tensors
022fd1d 
quantizeMlx quantizes a dense tensor to MLX affine 4-bit with groups of 32, 64
or 128 the way MLX does: float32 arithmetic with a true division, the bias at
the larger-magnitude edge of a group, a signed scale chosen so that zero is a
grid point, round half away from zero, codes taken from the unrounded scale and
bias, and bf16 rounding only for storage. When the dense source and the
conversion scheme match how a checkpoint was made, it reproduces that
checkpoint's words, scales and biases exactly. A non-finite value, an
overflowing range, a zero adjusted scale and a scale or bias that rounds to a
bf16 infinity are refused, naming the group.

`tf-merge-check --quantizer DENSE SHIPPED MODULE...` quantizes each named dense
tensor and compares it with the shipped one, reporting the share of equal
scales, biases, codes and whole groups. The dense tensor must be an n x k f32,
bf16 or f16 matrix of the shipped projection's shape and payload size, every
required input must be a regular file, and a failed allocation unwinds cleanly.

Tests pin the quantizer to nine groups produced by an independent float32
implementation of the published algorithm: two where true division and
reciprocal multiplication differ, one that needs the unrounded scale and bias,
a q0 tie at -8.5, the larger edge at the minimum and at the maximum, a negative
flat group and a tiny group with q0 = 0; they also check a local reciprocal
variant disagrees with them. Other tests cover group sizes 32, 64 and 128, the
refusals above, and the check on synthetic dense and shipped folders. Build:
`zig build test-merge`, `zig build tf-merge-check`.
@philip-pentatonic
Depends on the quantizer commit. Add tensorfold merge: fold a PEFT …
d97a547 
…LoRA into a quantized checkpoint

`tensorfold merge DENSE ADAPTER OUT --quantize SHIPPED` merges a PEFT LoRA into
the dense (bf16) weights of only the tensors the adapter names, quantizes those
fresh with the MLX quantizer at the shipped tensor's bits and group, and copies
every other tensor, shard and index byte for byte from the shipped MLX affine
4-bit checkpoint. Folding an adapter into already-quantized weights keeps
little of its effect, so the dense base is required. Host-only. The adapter's
B*A is accumulated in FP64 in order and scaled once; a zero sum gives an
exactly zero delta, and a sum that is not exact in FP64 is not made exact.

The merge is staged in a sibling folder and renamed to OUT only after every
projection has succeeded, so a failure before publication leaves no OUT; the
rename replaces an empty OUT atomically and never a non-empty one. It refuses,
naming the tensor: other modules (only the Nemotron mixer in_proj/out_proj and
shared experts' up_proj/down_proj and the Qwen dense q/k/v/o_proj are served),
other bit widths, DoRA and rsLoRA, an adapter r or alpha out of range,
duplicate PEFT keys, shard names that leave the folder, subdirectories,
dangling symlinks and special files among the inputs, shards whose tensor
ranges overlap, non-finite weights or factors, and mis-shaped tensors.

tf-merge-check DENSE ADAPTER OUT SHIPPED gates a merged folder: the exact file
set and lengths and every byte outside the touched tensors identical; every
written group equal, byte for byte, to a recomputation from the dense weights
and the adapter, with the number of groups compared reported; aggregate RMS
error at most 1.05x the shipped quantization's, both against an FP64
reference. The worst-column ratio is printed but not gated, being an extreme of
two differently referenced tensors. An adapter with no effect (alpha 0, an
all-zero factor, negative zero included) must leave every byte and give ratios
of exactly 1. The recomputation shares the merge's arithmetic, so it checks
that what was written is what the code computed; the independent evidence is
the pinned vectors, including one merge with a non-power-of-two alpha/r, and
the quantizer's identity with shipped tensors.

Tests use synthetic folders: a zero adapter reproduces every file, a random
adapter changes only the named tensors, each refusal names its tensor, a
failure on the second projection or at any allocation leaves no OUT, a
non-empty OUT that appears before publication is never overwritten, and the
gates fail on appended bytes, extra files, a wrong group and a collapsed grid.
@philip-pentatonic philip-pentatonic mentioned this pull request Oct 11, 2026
Open
@philip-pentatonic
merge: take the merged modules and their names from the checkpoints
d5b63df 
The merge no longer admits modules by a list of name suffixes. A module is
merged when the shipped checkpoint holds its .weight, .scales and .biases as
rank-2 tensors in the 4-bit affine layout and the dense source holds a rank-2
.weight of the same shape, so attention, MLP and linear-attention projections
qualify alike. Anything else is refused naming the module and the reason: a
module not quantized in the shipped checkpoint (a norm or an embedding), a
stacked or rank-3 expert tensor, or a module the dense source lacks. Every
other check is unchanged.

A dense module name is resolved to the shipped one by the exact name, then by
one table of known prefix rewrites (model.language_model. to
language_model.model.), so a dense source and an MLX conversion that name the
same tensor differently still match. A module that resolves under no candidate,
or under more than one, is refused naming each name tried. The leaves are read
and written under the shipped name.

Tests use synthetic folders: MLP and linear-attention names are merged, the
prefix mismatch is resolved and passes the gates, and a stacked expert tensor,
an unquantized norm, a module absent from the dense source, an ambiguous name
and an unresolvable name are each refused with their reason.
@philip-pentatonic
merge: a loose sanity bound on the RMS ratio, and module names in eit…
14a1ad9 
…her namespace

The aggregate-RMS comparison with the shipped quantization is no longer a gate
at 1.05. A delta that widens a group's range legitimately coarsens that group's
grid, so on sound merges of real adapters the ratio read 1.0518 and 1.0515
against that limit; it measures the adapter, not a defect. Both RMS ratios and
the worst-column ratio are printed, and the check fails only when the merged
RMS error exceeds twice the shipped one, a net for a gross fault such as a
wrong scale or a sign error and not a quality gate. The deterministic gates are
unchanged: the exact file set and lengths, every byte outside the touched
payloads, every written group equal to its recomputation, and exact identity
for an adapter with no effect. The test of the 1.05 limit is replaced by one
where 1.06 and exactly 2.0 pass and 2.01 fails.

A module may now be given in either namespace. The exact name is tried first,
then the table of prefix rewrites forwards and in reverse; the pair of names
must be present in both the dense source and the shipped checkpoint. A name
that resolves to no pair, or to more than one, is refused naming every pair
tried, and the quantizer check prints which form matched. The merge resolves
adapter module names the same way.

Tests cover a module resolving from either side, an ambiguous name, an
unresolvable name, and an adapter naming the shipped spelling.
@philip-pentatonic

philip-pentatonic commented Oct 12, 2026

Copy link
Copy Markdown
Contributor Author

Pushed two follow-up commits on top of the first (not a rewrite, so the earlier review stays valid):

  • d5b63df6: the merged modules and their names now come from the checkpoints. A module is merged when SHIPPED holds it as a rank-2 4-bit affine tensor and DENSE holds a matching rank-2 weight, so MLP, linear-attention and Mamba projections qualify, not only a fixed list. A small table maps the dense and shipped prefixes for models that name them differently (Qwen3.8-27B: model.language_model. against language_model.model.).
  • 14a1ad93: the aggregate RMS ratio against SHIPPED is now printed with only a loose sanity bound of 2.0. A real-checkpoint run showed the tighter bound failed sound input: a delta widens a group's range, which legitimately coarsens its grid. Module names resolve in either namespace.

The description is updated with the receipts for both models: Nemotron 3.5 Lightning (as before) and now Qwen3.8-27B, where the zero adapter reproduces all 17 files of the shipped checkpoint byte for byte and a random adapter over 15 linear-attention, attention and MLP modules matches its recomputation on all of about 12.9 million groups. Not run: serving a merged folder on any engine, Metal and macOS builds.

@philip-pentatonic

philip-pentatonic commented Oct 12, 2026

Copy link
Copy Markdown
Contributor Author

An end-to-end run with a real trained adapter, on a model small enough to also check at the model level.

Setup: dense Qwen/Qwen3-4B, shipped mlx-community/Qwen3-4B-4bit (group 64), a real trained rank-64 LoRA (alpha 16) over the attention and MLP projections of all 36 layers, so 252 projections. tensorfold merge --quantize merged all 252; tf-merge-check passes every gate (identity 10 files byte-identical, 1 touched; every written group equal to the recomputation).

Model level, on 20,480 positions of English text, KL(unmerged adapter on the bf16 base || variant), the merged weights read back from the tool's own output:

Variant KL Same top token as the unmerged adapter
shipped 4-bit, adapter ignored 0.2706 ± 0.0059 74.9%
merge --quantize output 0.0803 ± 0.0015 83.8%

Per sequence the merged model is 2.6 to 4.3 times closer (median 3.4). A separate PyTorch implementation of the same quantizer gives 0.0802, and the tool's weights equal it exactly on the modules I compared. One adapter and one 4B model, so a measurement, not a rate.

Two caveats a reviewer should know:

  • Byte identity depends on the converter. This particular shipped checkpoint (mlx-community/Qwen3-4B-4bit, converted in April 2025) is not bit-identical to quantizeMlx applied to its dense source: biases match 100%, but 0.07% of scales and 0.2% of codes differ (the Nemotron 3.5 Lightning and Qwen3.8-27B checkpoints in the receipts match exactly). My guess, not checked, is different arithmetic in the converter, for example a fast-math division on the GPU. So with such a checkpoint a zero adapter will not reproduce SHIPPED byte for byte on the touched tensors, which is why the description says exactness holds when the dense source and the conversion arithmetic match.
  • Speed is not addressed here. The merge took 6 min 14 s (9.4 GB resident, mostly mapped dense pages) and the check 18 min 2 s (11.4 GB) for these 252 projections, single-threaded; the time is dominated by the FP64 accumulation of B*A. It could be threaded across rows; I left that out of this pull request.

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

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@philip-pentatonic