Repository navigation
Conversation
The loader, the memory plan and admission read config.json's per-module widths: MLX affine 4, 5, 6 and 8 bits in groups of 64, one width per tensor. Routed stacks are checked against their stated width. Other widths, unquantized projections and an MTP head inside the model's shards are refused by name. The uniform 4-bit checkpoint reads as before.
mtp-4bit.safetensors is read at the quantization its metadata states (4-bit in groups of 64 when it states none, as the stock head), beside a model of any width. Admission plans its bytes at that width and refuses a head width the launches do not read, by name.
Gates 1-4 run the relu^2 stream tile at every width Nemotron's oQ checkpoints store it (4, 5, 6 and 8 bits in groups of 64), each wider width with a planted control that reads its codes as 4-bit. Gate 5 records the 24 launches at 5, 6 and 8 bits on the four RDNA chips.
|
Thanks @gilby. This is a thorough first receipt, and the loader change reads cleanly. What I checked on
What's left. Items 1 and 3 repeat our reply on #3061, and 2, 4 and 5 are new:
One overlap to know about. #3001's padded row stride, for gfx1151, will also change We have no AMD hardware, so your gfx1151 runs are the evidence here. Once the split PRs carry these, we'll review each and land them in order.
|
|
oQ5e and oQ6e are now tested too. The description listed them as not run because no such checkpoint was published. I built all four levels (oQ4e, oQ5e, oQ6e, oQ8e) from
|
|
CI is green on |
|
Thanks @gilby for building oQ5e and oQ6e and running them, and for the green CI on Done, at every width.
Still open. This isn't ready to land yet.
One small fix for the records PR. The comment at The oQ4e one-stream T=1.0 cell you flagged is not a gate item. oQ checkpoints have no main to compare against, so it needs no rerun. We have no AMD card, so your gfx1151 runs are the device evidence. Once the three PRs are up with items 2 to 5, we'll review each and land them in order.
|
|
Thanks for both reviews. This PR is replaced by three stacked PRs, landing in this order. Parts 1 and 2 are open; I'll close this one when part 3 is open:
Your points, one by one, numbered as in your latest comment:
The comment at #3001: noted. Part 2 changes |
|
Part 3 of 3 is open as #3251 (the MTP head file read at its own width), stacked on #3242 and #3163. The repeats I mentioned are in its body: the three cells that read slower on one arm were repeated as interleaved pairs, and the one that still reads slower in both pairs is |
What this changes
Nemotron 3.5 Lightning on HIP read only
mlx-q4g64. It now reads the oQ mixed-width MLX checkpoints the Metal backend of the family already serves: MLX affine at 4, 5, 6 or 8 bits in groups of 64, a width per matrix. No kernel source changes: the shared HIP tiles already decode these widths, and the single width lived in the Nemotron HIP admission, loader and memory plan.Three commits, each a step that stands alone; say if you want them as three PRs and I'll split:
534825233Each matrix at its own width. The loader's table returns the per-matrix widthconfig.zigalready parses (asqwen3_5does on HIP); routed stacks are read at their own width and checked against it; the plan counts every matrix at its own width, the draft rows atlm_head's; admission admits 4/5/6/8 bits in groups of 64 and refuses the rest by name; the start-up line names the widths, e.g.Nemotron-H model mlx-q4g64+q5g64+q6g64+q8g64 18.14 GiB.b10870a78The MTP head at its own file's width. The separate head file is read at the width its own metadata states (the stock head is 4-bit in groups of 64), beside a main model of any width, as the Metal recipe does.963953d11Kernel-contract records for the relu² stream tile at 5, 6 and 8 bits, in a new records file (hip_records.zigis at 600 lines), each wider width with a planted control.Still refused by name, each with a host test: other bit widths and group sizes, unquantized (bf16) projections and embeddings,
mtp.layers.*inside the main shards, a head file at an unsupported width.atc overlap --json: run on963953d11against base6c5cef894. The result is incomplete (ATC reports its pull-request, file and time limits); it lists no open PR sharing this branch's files.Checks on this head
5474a338e:963953d11(the head of the receipt) plus the changelog line.python3 tools/pr_check.py --base origin/main --pr 3068: 0 errors, 0 warnings.zig build test -Dhip=true --summary all: 13,676 of 13,790 tests passed, 113 skipped, 1 failed. The failure isdashboard_test"dashboard page polls the public stats snapshot without private details", which fails the same way on main without this change.config.json's width for it (red whenwidthignores its key); every refusal by name (red when every width is accepted); the plan's bytes equal the real oQ6 and oQ8e header bytes (red with a 4-bit-only count); the head is planned at its own file's width (red when planned at the main model's); the stock 4-bit test still pins the checkpoint at exactly 17,775,339,392 bytes, so the uniform path is unchanged.zig build hip-kernel-contract:SUMMARY gfx1151 162 kernels: 162 pass gates 1-4, 0 fail, 2 known findings; 44 controls: 44 red as planted, 0 missed; 173 limit problems(main: 159 kernels, 41 controls, 173 limit problems).zig build hip-kernel-test: 159 of 165 passed, 6 skipped, the same as main.Receipt
One Strix Halo iGPU (gfx1151), fresh servers with the page cache dropped before each,
--context 131072 --parallel 4 --prompt-cache-gib 0. Raw tool JSON for every arm: https://gist.github.com/gilby/00b2bb9f0acf63a422cd7168b0c322b8A. The 4-bit checkpoint, main against this head (
TensorFold/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-MLX-4bitatd9d758fbwithmtp-4bit.safetensors; arms main, this, this, main; main is6c5cef894). Aggregate tok/s, both arms of each side:bench_openaicode, T=0 / T=1.0bench_openaichat, T=0 / T=1.0--no-draftscode / chat, T=0token_sha256equal main's in each arm; the server log's reply hashes for all 89 requests of an arm are the same set as main's (six positions differ in completion order between concurrent streams, and differ the same way between main's own two arms); 17 of 17 equal with--no-drafts. In every cell of every arm concurrent equals alone (42 equal, 0 unequal, 0 failed) and alone equals"draft": false(24 equal, 0 unequal, 0 failed). Tokens a round 2.31–2.32 in all four arms.Nemotron-H model 17.34 GiBon main andNemotron-H model mlx-q4g64 17.34 GiBhere.B. oQ checkpoints on this head (they do not load on main: an oQ4e fails at load with
backbone.layers.0.mixer.in_proj: packed words and group tables have incompatible or unrepresentable geometry, an 8-bit one is refused by name). Each with the stockmtp-4bit.safetensorsbeside it:user3892r4mlx-oQ4e (3948ddc5)TensorFoldoQ6 (8cf893c6)user3892r4mlx-oQ8e (5b942665)All three: concurrent equals alone (42 equal, 0 unequal, 0 failed) and alone equals
"draft": false(24 equal, 0 unequal), greedy and sampled; every cold row is cold. Memory on the device less the planned bytes is a constant 7.0–7.1 GiB at every width, so the plan follows the real bytes.A checkpoint with a bf16 embedding and bf16 attention stays refused by name, as Nemotron on HIP: read oQ mixed-width MLX checkpoints (4, 5, 6 and 8 bits, a width per tensor) #3061 asks about; so does a head stored inside the main shards.
New switch? None. No default changes; the 4-bit checkpoint's tokens, speed and memory are main's.
Boxes it reaches: HIP devices serving Nemotron-H. Shared code touched: none outside the family except the kernel-contract records and their test (
core/contract/hip_records_relu2_widths.zig,tests/hip/contract_relu2*.zig,contract_main.zig); no launch or kernel changes for any family.Environment: Zig 0.17.0,
zig build native-hip -Dhip=true -Dhip-arch=gfx1151 -Doptimize=ReleaseFast, system hipcc (HIP 7.1, clang 21), Linux 7.0.Licence of new weights: n/a (no weights added; the checkpoints named are public uploads read as published).