diff --git a/.github/workflows/compile-check.yml b/.github/workflows/compile-check.yml new file mode 100644 index 0000000..9027479 --- /dev/null +++ b/.github/workflows/compile-check.yml @@ -0,0 +1,79 @@ +name: compile-check + +# Two layers of protection for the imatrix generators: +# +# 1. compileall — every shipped Python script must parse on all supported Python +# versions. Catches architecture-specific generator variants (and anything else) +# landing unparsable — e.g. backslashes inside f-string expressions, which only +# became legal in Python 3.12 (PEP 701). +# +# 2. tests — actually EXECUTE the checks that have no external dependencies, so the +# pieces most worth locking down are proven rather than merely parsed: +# * format layer (numpy + gguf only) -> test_roundtrip.py +# * calibration windowing / llama.cpp BOS -> test_calib.py (torch) +# * GLM-4.5 routed-expert statistic -> test_glm4_expert_stats.py (torch) +# * Llama-4 expert row counts -> test_llama4_expert_counts.py +# (torch + transformers) +# * causal masking -> test_causal_mask.py (torch) +# None of these need a checkpoint, a GPU, or network access at run time. +on: + push: + pull_request: + +jobs: + compileall: + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + python-version: ["3.9", "3.10", "3.11", "3.12"] + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + - name: Compile all shipped Python scripts + run: python -m compileall -q scripts + + format-tests: + # Cheap: no torch. Guards the imatrix GGUF layout llama.cpp has to read back. + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: "3.12" + - name: Install format deps + run: python -m pip install --upgrade pip numpy gguf + - name: imatrix format round-trip + working-directory: scripts/imatrix_serialized + run: python test_roundtrip.py + + generator-tests: + # CPU-only torch: the statistic/count/windowing logic, no checkpoint required. + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: "3.12" + - name: Install CPU torch + transformers + run: | + python -m pip install --upgrade pip + python -m pip install --index-url https://download.pytorch.org/whl/cpu torch + python -m pip install numpy gguf transformers + - name: Calibration windowing + BOS parity + working-directory: scripts/imatrix_serialized + run: python test_calib.py + - name: GLM-4.5 routed-expert statistic + working-directory: scripts/imatrix_serialized + run: python test_glm4_expert_stats.py + - name: Llama-4 expert row counts + working-directory: scripts/imatrix_serialized + run: python test_llama4_expert_counts.py + - name: Generator dispatch resolves per arch + working-directory: scripts/imatrix_serialized + run: python test_dispatch.py + - name: Causal masking (non-causal imatrix regression) + working-directory: scripts/imatrix_serialized + run: python test_causal_mask.py diff --git a/scripts/apex_pipeline.sh b/scripts/apex_pipeline.sh index 130a211..97b74a1 100755 --- a/scripts/apex_pipeline.sh +++ b/scripts/apex_pipeline.sh @@ -56,6 +56,14 @@ ONLY_PHASES="" EVAL_SUITE="" # override from CLI; otherwise use YAML DRY_RUN=false NGL="${NGL:-99}" +# imatrix backend: "llama" (stock llama-imatrix, whole model resident) or +# "serialized" (band-bounded PyTorch generator; peak resident = IMATRIX_BAND layers, +# independent of model size — for models too large to llama-imatrix locally). +IMATRIX_BACKEND="${IMATRIX_BACKEND:-llama}" +IMATRIX_BAND="${IMATRIX_BAND:-1}" +IMATRIX_DEVICE="${IMATRIX_DEVICE:-cuda}" +IMATRIX_BATCH="${IMATRIX_BATCH:-}" # chunks per forward; empty = generator default +IMATRIX_MTP="${IMATRIX_MTP:-false}" # also cover the NextN/MTP head (arch must support it) WORK_DIR="${WORK_DIR:-/workspace/data/apex}" LLAMA_CPP_DIR="${LLAMA_CPP_DIR:-}" PORT="${PORT:-8192}" @@ -616,13 +624,35 @@ fi # Phase 6: Generate imatrix # ═══════════════════════════════════════════════════════════════════════════════ if should_run imatrix; then - log "Phase 6: Generating importance matrix" - F16=$(find_f16) - + log "Phase 6: Generating importance matrix (backend: $IMATRIX_BACKEND)" info "Calibration file: $CALIBRATION ($(wc -l < "$CALIBRATION") lines, $(du -h "$CALIBRATION" | cut -f1))" - info "Source model: $F16" - $IMATRIX_BIN -m "$F16" -f "$CALIBRATION" -ngl $NGL \ - -o "${MODEL_DIR}/imatrix.dat" 2>&1 | tail -5 + + if [ "$IMATRIX_BACKEND" = "serialized" ]; then + # Band-serialized PyTorch backend: peak resident = IMATRIX_BAND layers, + # independent of model size. Reads the HF safetensors (from Phase 2), not the + # F16 GGUF, and writes the same imatrix.dat that Phase 7 consumes. + HF_DIR="${MODEL_DIR}/safetensors" + [ -d "$HF_DIR" ] || { err "serialized backend needs HF safetensors at $HF_DIR (download without SOURCE_GGUF)"; exit 1; } + info "Source (HF safetensors): $HF_DIR | band=$IMATRIX_BAND device=$IMATRIX_DEVICE" + # dispatch.py picks the generator by the checkpoint's model_type and probes the + # installed transformers for the symbols that variant needs, so an unsupported + # arch or a mismatched transformers fails with an actionable message instead of + # an ImportError from a hardcoded Granite-only script. + # --mtp / --batch only exist on the generators that support them; argparse in the + # dispatched script reports it plainly if the knob does not apply to that arch. + IMATRIX_EXTRA=() + if [ "$IMATRIX_MTP" = "true" ]; then IMATRIX_EXTRA+=(--mtp); fi + if [ -n "$IMATRIX_BATCH" ]; then IMATRIX_EXTRA+=(--batch "$IMATRIX_BATCH"); fi + python3 "${REPO_DIR}/scripts/imatrix_serialized/dispatch.py" \ + --model "$HF_DIR" --calib "$CALIBRATION" \ + --out "${MODEL_DIR}/imatrix.dat" --band "$IMATRIX_BAND" --device "$IMATRIX_DEVICE" \ + "${IMATRIX_EXTRA[@]+"${IMATRIX_EXTRA[@]}"}" 2>&1 | tail -20 + else + F16=$(find_f16) + info "Source model: $F16" + $IMATRIX_BIN -m "$F16" -f "$CALIBRATION" -ngl $NGL \ + -o "${MODEL_DIR}/imatrix.dat" 2>&1 | tail -5 + fi ls -lh "${MODEL_DIR}/imatrix.dat" fi diff --git a/scripts/imatrix_serialized/README.md b/scripts/imatrix_serialized/README.md new file mode 100644 index 0000000..a4658a1 --- /dev/null +++ b/scripts/imatrix_serialized/README.md @@ -0,0 +1,214 @@ +# Serialized imatrix backend (`IMATRIX_BACKEND=serialized`) + +A memory-bounded PyTorch reimplementation of llama.cpp's importance-matrix +computation, done via forward hooks on the HF safetensors. Motivation: +`llama-imatrix` needs the whole model resident (RAM+VRAM), which is impractical for +large MoE / hybrid models (and, for hybrid SSM architectures like `qwen3_5_moe`, the +llama.cpp imatrix path falls back to a serial CPU scan). This decouples the imatrix +from llama.cpp's memory model: peak resident weights = `IMATRIX_BAND` layers, +independent of model size. + +Wired into `apex_pipeline.sh` Phase 6 as an **opt-in** backend +(`IMATRIX_BACKEND=serialized`, default `llama` unchanged). Reads the HF safetensors +fetched in Phase 2 and writes the same `imatrix.dat` Phase 7 consumes. torch / +transformers / accelerate are only needed for this backend — see `requirements.txt`. + +## What the imatrix is +Per weight tensor `[out, in]`, a per-input-channel second moment `Σ_t x_j²`, stored +in a GGUF as two F32 tensors per weight: +- `.in_sum2` `[nmat, in]` — raw Σx² (nmat = n_experts, or 1 for dense) +- `.counts` `[nmat, 1]` — token count per expert/matrix + +Size is set by `Σ nmat×in` (independent of `out` and of calibration size). + +## Files +| File | Role | +|---|---| +| `imatrix_io.py` | read/write GGUF imatrix (bit-exact round-trip) | +| `calib.py` | shared calibration tokenization/windowing incl. llama.cpp's per-chunk BOS | +| `dispatch.py` | picks the generator for a checkpoint's `model_type` by probing the installed transformers; what Phase 6 calls | +| `gen_imatrix.py` | full-forward generator: hook every Linear / fused-expert module, accumulate Σx² | +| `serialized_gen.py` | **band-serialized** generator (GraniteMoeHybrid on transformers 5.5.x) | +| `compare.py` | per-tensor correlation + L1 error vs a reference imatrix | +| `check_gate_up.py` | diagnostic: confirms tensors sharing a graph input share an imatrix stat (takes a file path) | + +Tests — all synthetic, none need a checkpoint/GPU/network, all run in CI: + +| Test | Covers | +|---|---| +| `test_roundtrip.py` | GGUF format layer: write → read → bit-exact; llama.cpp's `in_sum2`/`counts` layout; metadata keys. numpy+gguf only | +| `test_calib.py` | windowing, chunk cap, per-chunk BOS parity with llama.cpp, short-corpus hard-fail | +| `test_glm4_expert_stats.py` | routed-expert Σx² shapes/values vs an independent per-expert reference; both fused weight layouts; hard-fail cases | +| `test_llama4_expert_counts.py` | drives the Llama-4 hooks through a real `Llama4TextMoe`: exact top-k row counts, router-weighted convention | +| `test_dispatch.py` | every arch entry resolves to an existing script; unknown/unavailable arch messages | +| `test_causal_mask.py` | that every serialized generator sets `_attn_implementation`, without which `create_causal_mask()` returns `None` and the standalone layers attend bidirectionally | + +### Per-architecture generator variants +Per-arch name mapping is the extension surface. **You do not pick these by hand** — +`dispatch.py` reads `model_type` from the checkpoint's `config.json` and probes the +installed transformers for the symbols each variant imports, because the version ranges +are genuinely misleading (5.5.1 already has `Qwen3_5MoeExperts` and `Llama4Router`, but +not `GraniteMoeHybridExperts`). To see what your install resolves to: + +```bash +python3 scripts/imatrix_serialized/dispatch.py --model --print-script +``` + +| Script | Architecture | needs | Status | +|---|---|---|---| +| `serialized_gen.py`, `gen_imatrix.py` | GraniteMoeHybrid | `GraniteMoeHybridParallelExperts` (5.5.x) | ✅ validated | +| `serialized_gen_t514.py`, `_t514_fast.py`, `gen_imatrix_t514.py` | GraniteMoeHybrid | `GraniteMoeHybridExperts` (5.14+; `_fast` VRAM-bounded) | ✅ validated | +| `serialized_gen_qwen35.py` | `qwen3_5_moe` (Qwen3.5/3.6-35B-A3B; GatedDeltaNet + MoE + NextN/MTP head) | `Qwen3_5MoeExperts` (present in 5.5.1) | ✅ validated | +| `glm4moe_serialized_gen.py` | GLM-4.5-MoE | fused `Glm4MoeExperts` (absent in 5.5.x) | ⚠️ experimental (untested) | +| `serialized_gen_llama4.py` | Llama-4 | `Llama4TextExperts` (present in 5.5.1) | ⚠️ experimental (untested) | + +Every generator hard-fails (`KeyError`) if any band/head parameter has no checkpoint +tensor. A skipped key would leave randomly-initialized weights in the band and emit a +plausible-looking but meaningless imatrix, so a name mismatch is never survivable. + +**Llama-4 note:** its routed gate/up statistic is *router-weighted*, and that is +correct, not a divergence — llama.cpp special-cases this arch +(`weight_before_ffn = arch == LLM_ARCH_LLAMA4`, `src/llama-graph.cpp:1837`) and folds +the sigmoid-ed weights into the expert input **before** the gate/up `mul_mat_id` +(L1976), where the generic MoE path instead weights the expert *output* (L2121). + +## Setup + +> **Install a CUDA torch build, not the default CPU wheel.** `pip install torch` +> pulls the **CPU-only** wheel, and the generator will then run entirely on CPU +> (`--device cuda` will error or, on CPU, crawl — orders of magnitude slower, and +> the hybrid-SSM path is especially slow on CPU). Install torch from the CUDA index +> matching your driver, e.g.: +> ```bash +> pip install torch --index-url https://download.pytorch.org/whl/cu128 # pick cuXXX for your CUDA +> pip install transformers accelerate numpy gguf # rest of requirements.txt +> python3 -c "import torch; assert torch.cuda.is_available(), 'CPU-only torch — reinstall from the CUDA index'" +> ``` + +Then set `IMATRIX_BACKEND=serialized` (optionally `IMATRIX_BAND=`, +`IMATRIX_DEVICE=cuda`, `IMATRIX_BATCH=`, `IMATRIX_MTP=true`) when invoking +`apex_pipeline.sh`. Phase 6 calls `dispatch.py`, which selects the generator from the +checkpoint's `model_type`. Standalone, either let dispatch pick: +```bash +python3 scripts/imatrix_serialized/dispatch.py --model \ + --calib --out imatrix.dat --band 1 --device cuda +``` +or run a variant directly (`serialized_gen.py`, `serialized_gen_qwen35.py`, ...) with +the same flags. + +### Key flags +- **`--attn-impl`** (default `eager`) — the attention implementation stamped onto the + config *before* masks are built. This is not cosmetic: `create_causal_mask()` returns + `None` when `config._attn_implementation` is unset, and because these generators drive + standalone decoder layers rather than a full model, a `None` mask makes them attend + **bidirectionally** — a silently non-causal imatrix. On a synthetic 4-layer + GraniteMoeHybrid, `blk.3.attn_output` disagreed with the full-forward generator by + 7.5e-01 (corr 0.54) unset, and by 3.2e-06 set. +- **`--band N` / `--batch N`** — throughput knobs. `band=1 batch=1` is the memory-minimal + but *slowest* corner: the GPU starves between per-layer load + host↔device transfer cycles. + Raise `--batch` for bigger GEMMs, and `--band` to amortize the per-band weight load over + more compute. Both are **result-invariant** (Σx² is order/batch/band-invariant) — tune + them purely for speed vs. VRAM, the imatrix comes out identical. +- **`--no-bos-per-chunk`** — opt out of forcing BOS at position 0 of every chunk. + llama.cpp *does* force it (`tools/imatrix/imatrix.cpp:865-866` overwrites the first + token of each chunk when the vocab has `add_bos`), so forcing it is the default here + and only this flag reproduces a plain stream-window. Note the published validation + numbers predate the fix, i.e. they were measured stream-windowed; the effect is ~1 + token in `ctx` (0.2% at `ctx=512`). +- **`--mtp`** *(generators for models with a NextN/MTP head, e.g. `serialized_gen_qwen35.py`)* — + also compute importance for the multi-token-prediction head (`blk.{n_layers}.*`). + `llama-imatrix` cannot cover it; omit the flag to match a head-less reference file. + + The `nextn.eh_proj` input ordering is **verified against llama.cpp**, not assumed: + `src/models/qwen35moe.cpp` builds the head as + `ggml_concat(ctx0, e_norm, h_norm, /*dim=*/0)` and `ggml_compute_forward_concat` + writes `src0` at the low indices, so input channels `[0, n_embd)` are the + embedding-norm half and `[n_embd, 2*n_embd)` the hidden-norm half. The generator + builds `torch.cat([enorm(emb), hnorm(h)], dim=-1)` — same order. The GGUF name comes + from `conversion/qwen.py`, which renames `mtp.fc` → `layers.{n_layer}.eh_proj` with no + transpose, so the channel order carries through to `blk.{n_layers}.nextn.eh_proj.weight`. + + The one-position shift is verified too, caller-side: `common/speculative.cpp` submits + the token sampled *at* a position together with that position's hidden state in one + batch row, i.e. `embed(t_{i+1})` with `h_i` — matching the generator's + `emb[:, 1:]` / `h[:, :-1]`. So no MTP imatrix needs rebuilding on this account; what + is still unvalidated is the numerical result (no public imatrix covers MTP), not the + wiring. + +### Calibration corpus +Use a **diverse** calibration file — mixed prose, code, and math — not wikitext alone. +The imatrix records per-channel activation statistics, so channels that only fire on +(e.g.) code or non-English text get no signal from an all-Wikipedia corpus, which +matters most for coding/agentic and multilingual models. Bartowski's +`calibration_datav3` is a good general default; weight toward your target domain. +Calibration diversity changes the imatrix *values*, not its size or tensor coverage. + +The calibration corpus also governs *reproducibility* and *specialization*: re-running with +the **same** calibration reproduces an imatrix to ~0.99 median per-tensor correlation, while a +**different** corpus legitimately shifts the values. That shift concentrates in the routed +**expert** projections (`ffn_*_exps`) — the experts are what specialize by domain — whereas +attention and shared-expert paths stay ~1.0. So a code-weighted calibration measurably favors +coding channels (at a small cost elsewhere); reach for it deliberately when specializing for a +single domain, and use the *same* calibration when your goal is to reproduce a reference. + +## Validation + +### granite-4.0-h-tiny — end-to-end PPL parity +Same i-quality config + Q6_K base + eval, swapping only the imatrix (200×512 windows): + +| imatrix | i-quality PPL | +|---|---| +| **serialized (this tool)** | **8.8966 ± 0.107** | +| llama.cpp (ground truth) | 8.9030 ± 0.107 | + +Difference (0.006) is an order of magnitude inside the error bars — statistically +indistinguishable. Format bit-exact; names 368/368 exact; median per-tensor +correlation 0.998 (0.995 band-serialized). band=4 peak GPU 6.4 GB (vs ~13 GB full). + +### Qwen3.5-35B-A3B (`qwen3_5_moe`) — second architecture +The headline apex-quant arch. Validated against +[bartowski's canonical llama.cpp imatrix](https://huggingface.co/bartowski/Qwen_Qwen3.5-35B-A3B-GGUF): + +- **Name mapping:** 523 tensors covered, `only-real = []` (strict superset — also + covers the NextN/MTP head, which `llama-imatrix` does not); median per-tensor + correlation **0.966** (only `ssm_out` low, see below). +- **PPL parity** (i-compact, 200×512 windows): bf16 6.620 · llama.cpp-imatrix 6.756 · + **serialized-imatrix 6.775** — the two quants differ by 0.02 PPL, inside the ±0.073 + error bars. Statistically indistinguishable. + +Published for verification: +[Qwen3.5](https://huggingface.co/Myric/Qwen3.5-35B-A3B-APEX-GGUF) · +[Qwen3.6](https://huggingface.co/Myric/Qwen3.6-35B-A3B-APEX-GGUF) (both include the +serialized imatrix). + +**Note on `down_*` / `ssm_out`:** these correlate poorly per-tensor (inputs are +post-nonlinearity: SiLU-gated intermediates, SSM scan output). It does **not** affect +PPL — the quantizer needs only relative per-channel importance, which is preserved. + +## Extending to a new architecture +`map_name()` maps HF dotted module names → ggml tensor names and dispatches by module +type (`nn.Linear` → all-token; fused-expert module → per-expert split). For a new arch, +add its name mapping (and its fused-expert hook if different). Cross-check against a +real imatrix's tensor-name set via `compare.py` / `--ground-truth` — any unmapped name +warns. Then add the arch to the `ARCHES` table in `dispatch.py` (script + the +transformers symbols it imports) so Phase 6 can reach it; `test_dispatch.py` asserts +every table entry points at a script that exists and resolves on the installed +transformers. + +## Known limitations +**Not yet measured where it matters most.** The published PPL comparisons were run at +Q6_K (`i-quality`) and Q4_K (`i-compact`), the rungs where the imatrix has the least +influence. The ladder goes down to `i-nano`=`iq2_xxs` and `i-micro`=`iq1_m`, and the +discriminating test is one of those by KL-divergence vs bf16 +(`llama-perplexity --kl-divergence`), which is far more sensitive to imatrix quality +than PPL. Until that number exists, "does not affect PPL" for the poorly-correlating +`down_*`/`ssm_out` tensors is an argument from a measurement that cannot show it. +Relatedly, the ±0.107 error bar quoted for the Granite comparison is the *absolute* +PPL standard error; since both quants were scored on identical windows off an identical +base, the correct statistic is the standard error of the paired per-window NLL +difference, which is much tighter. + +Prototype favors correctness: eager attention, per-band CPU↔GPU activation transfers. +Production TODOs: sdpa, batched chunks, keep activations on-GPU between bands, +`causal-conv1d`. `safe_open` mmaps the shards so peak *RSS* shows ~model size, but that +is reclaimable page cache — the committed working set (and GPU) is band-bounded. diff --git a/scripts/imatrix_serialized/TODO.md b/scripts/imatrix_serialized/TODO.md new file mode 100644 index 0000000..4df7723 --- /dev/null +++ b/scripts/imatrix_serialized/TODO.md @@ -0,0 +1,104 @@ +# Outstanding work — serialized imatrix backend (PR #18) + +State as of commit `84ec9a7` on `feat/serialized-imatrix`. The review that drove that +commit is `pullrequestreview-4823632517` on +https://github.com/localai-org/apex-quant/pull/18 + +Everything below is *not done*. What IS done is in the `84ec9a7` commit message. + +## Blocked on GPU time + real checkpoints + +These are the two remaining items from the reviewer's own "what would get this to merge" +list. Both need hardware I did not have in the session that wrote this. + +- [ ] **iq1_m or iq2_xxs KLD-vs-bf16 on Qwen3.5.** The single number that answers whether + the 0.966 median correlation matters. Published PPL comparisons were run at Q6_K + (`i-quality`) and Q4_K (`i-compact`) — the rungs where the imatrix has the least + influence. Use `llama-perplexity --kl-divergence` against bf16, at `i-nano` + (`iq2_xxs`) or `i-micro` (`iq1_m`). The concern is specifically `ffn_down_exps`: + dominant tensor group by parameter count in a 256-expert MoE, among the most + quantization-sensitive, and one of the poor per-tensor correlators. +- [ ] **Redo the Granite error bar as a paired statistic.** The "0.006 is inside ±0.107" + claim uses the *absolute* PPL standard error. Both quants were scored on identical + windows off an identical base, so the correct test is the standard error of the + **paired per-window NLL difference**, which is much tighter. Needs the per-window + NLLs from the original runs (or a re-run). The absolute effects (0.07% / 0.28%) are + small enough that it probably survives, but as presented the claim isn't established. + +## Consequences of the causal-attention fix — decide before trusting old numbers + +`84ec9a7` fixed all five serialized generators attending **bidirectionally** +(`create_causal_mask()` returns `None` when `config._attn_implementation` is unset; these +generators drive standalone decoder layers, so a `None` mask is non-causal). Effect +measured on a synthetic 4-layer GraniteMoeHybrid: `blk.3.attn_output` went from 7.5e-01 +relative disagreement with the full-forward generator (corr 0.54) to 3.2e-06. + +- [ ] **Re-run the Granite validation.** The published 8.897-vs-8.903 PPL and + 368/368 / median-corr-0.998 numbers were produced with non-causal attention on the + attention layers. On a mostly-mamba hybrid the affected tensors are a minority, + which is plausibly why a 0.998 median didn't surface it — but the numbers should be + regenerated before they're cited again. +- [ ] **Re-check the Qwen3.5 validation for the same reason** (corr 0.966 vs bartowski's + canonical imatrix, i-compact PPL 6.775). Qwen3.5 has proportionally more full + attention layers than Granite-hybrid, so the shift may be larger there. +- [ ] **Published artifacts** built on the old imatrices: + https://huggingface.co/Myric/Qwen3.5-35B-A3B-APEX-GGUF and + https://huggingface.co/Myric/Qwen3.6-35B-A3B-APEX-GGUF — decide whether to rebuild. +- [ ] Note also that per-chunk BOS is now forced (llama.cpp parity, + `imatrix.cpp:865-866`); all pre-`84ec9a7` numbers were stream-windowed. + `--no-bos-per-chunk` reproduces the old behaviour if you want an apples-to-apples + comparison rather than a fresh baseline. + +## Reply to the reviewer + +- [ ] **Correct the Llama-4 finding (review item 4).** The reviewer says the generator + "can't reproduce llama.cpp's statistic as written" because gate/up comes out + router-weighted where llama.cpp accumulates unweighted. That is true of the generic + MoE path but **not** of Llama-4, which llama.cpp special-cases: + `const bool weight_before_ffn = arch == LLM_ARCH_LLAMA4;` at + `src/llama-graph.cpp:1837`, then L1976 multiplies the sigmoid-ed weights into `cur` + **before** the gate/up `mul_mat_id`. The generic path weights the expert *output* at + L2121. So a weighted stat is correct here; `test_llama4_expert_counts.py` pins it. + The row-count fragility they flagged alongside it was real and is fixed. +- [ ] Report the non-causal-attention bug — it's more consequential than anything in the + review and was found by building an end-to-end reproduction they didn't have. +- [ ] Mention that `Qwen3_5MoeExperts` exists in transformers 5.5.1, so the qwen3_5_moe + path never needed 5.14+ (the README row was wrong); `dispatch.py` probes for symbols + rather than versions for exactly this reason. + +## Deferred refactor (reviewer's quality list, not a merge blocker) + +- [ ] **Collapse the five near-duplicate ~200-line generators** into one core module plus + per-arch name-map and hook adapters. Deliberately not attempted in `84ec9a7`: no + Granite-hybrid / Qwen3.5 / GLM / Llama-4 checkpoint was available locally, so a + refactor that size could not be verified against the validated paths. The + testable pieces were extracted instead (`calib.py`, `dispatch.py`, the Llama-4 hooks + are now module-level). `ShardReader` and the `ensure`/accumulator are still copied + across all of them and are the obvious next extraction. + +## Still untested against real checkpoints + +- [ ] **GLM-4.5** (`glm4moe_serialized_gen.py`). The `down_exps` mapping bug is fixed and + `test_glm4_expert_stats.py` covers the math synthetically, but nothing has run + against a real GLM-4.5 checkpoint. Also needs a transformers with fused + `Glm4MoeExperts` — absent in 5.5.x, so the import is guarded and `main()` explains. +- [ ] **Llama-4** (`serialized_gen_llama4.py`). Row counts and the weighting convention + are now covered by a test through a real `Llama4TextMoe`, but no Scout run exists + (bf16 doesn't fit 128 GB, which is the point of the backend). Validates by quant + coherence, not correlation. +- [ ] Both are marked ⚠️ experimental in the README; keep them there until run. + +## Context worth not re-deriving + +- The user is **not** shipping their transformers 5.14 transition yet. `dispatch.py` + therefore resolves Granite to `serialized_gen.py` (5.5.x) on the current install, and + the `_t514*` variants simply don't resolve until a newer transformers is present. Don't + "fix" this by pinning 5.14. +- Local llama.cpp checkout for cross-referencing: `/home/bryan/llama.cpp` (was at + `6f3c0a79`). The convert script is now a package — model classes live in + `conversion/*.py`, not `convert_hf_to_gguf.py`. +- No usable local checkpoint for these arches. `/home/bryan/models/granite-moe` is + `granitemoe`, **not** `granitemoehybrid`. The end-to-end verification in `84ec9a7` used + a synthetic 4-layer GraniteMoeHybrid built with random weights — that recipe is worth + rebuilding rather than hunting for a real tiny checkpoint. +- Installed transformers is 5.5.1; torch 2.11.0+cu130. diff --git a/scripts/imatrix_serialized/calib.py b/scripts/imatrix_serialized/calib.py new file mode 100644 index 0000000..6a44654 --- /dev/null +++ b/scripts/imatrix_serialized/calib.py @@ -0,0 +1,64 @@ +"""Shared calibration tokenization/windowing for the imatrix generators. + +One place, because every generator needs the identical thing and because the BOS +detail below is easy to get silently wrong. + +llama.cpp's reference implementation (tools/imatrix/imatrix.cpp) tokenizes the whole +corpus once and then, for every chunk, OVERWRITES the first token with BOS when the +vocab has add_bos: + + const bool add_bos = llama_vocab_get_add_bos(vocab); // L776 + ... + if (add_bos && j == 0) { // L865 + tokens[seq_start] = llama_vocab_bos(vocab); // L866 + } + +Windowing the tokenized stream without that leaves every chunk after the first +starting mid-sentence, so position 0 of each chunk sees a different distribution than +the reference run. It is ~1 token in `ctx` (0.2% at ctx=512), but it is a real and +avoidable divergence, so `add_bos=True` is the default here. +""" +import torch + + +def load_calibration_chunks(tokenizer, path, n_chunks, ctx, add_bos=True): + """Tokenize `path` and window it into an int64 tensor of shape [nch, ctx]. + + add_bos=True reproduces llama.cpp by forcing chunk[:, 0] = BOS (skipped when the + tokenizer has no BOS id). Pass add_bos=False to reproduce a plain stream-window, + i.e. the behaviour of these generators before the BOS fix. + + Returns (chunks, info) where info is a short human-readable summary string. + """ + with open(path, encoding="utf-8", errors="ignore") as f: + text = f.read() + # add_special_tokens=False so the corpus is a clean stream; per-chunk BOS is + # applied below exactly where llama.cpp applies it. + ids = tokenizer(text, return_tensors="pt", add_special_tokens=False).input_ids[0] + nch = min(n_chunks, ids.shape[0] // ctx) + if nch < 1: + raise ValueError( + f"{path}: {ids.shape[0]} tokens is too short for one chunk of ctx={ctx}") + chunks = torch.stack([ids[c * ctx:(c + 1) * ctx] for c in range(nch)]).long() + + bos = getattr(tokenizer, "bos_token_id", None) + if add_bos and bos is not None: + chunks = chunks.clone() + chunks[:, 0] = bos + bos_note = f"bos={bos} forced at chunk[:,0] (llama.cpp parity)" + elif add_bos: + bos_note = "no bos_token_id on this tokenizer; chunks left as-is" + else: + bos_note = "per-chunk BOS disabled" + return chunks, f"tokens={ids.shape[0]} chunks={nch} ctx={ctx} | {bos_note}" + + +def add_calib_args(ap): + """Register the calibration flags shared by every generator.""" + ap.add_argument("--chunks", type=int, default=126) + ap.add_argument("--ctx", type=int, default=512) + ap.add_argument("--no-bos-per-chunk", dest="bos_per_chunk", action="store_false", + help="do NOT force BOS at position 0 of every chunk; llama.cpp does " + "force it, so this only exists to reproduce pre-fix runs") + ap.set_defaults(bos_per_chunk=True) + return ap diff --git a/scripts/imatrix_serialized/check_gate_up.py b/scripts/imatrix_serialized/check_gate_up.py new file mode 100644 index 0000000..195e8bf --- /dev/null +++ b/scripts/imatrix_serialized/check_gate_up.py @@ -0,0 +1,63 @@ +"""Diagnostic: confirm tensors that share a graph input share an imatrix stat. + +gate/up (routed and shared) are fed the same activation by the ggml graph, so their +in_sum2 vectors should be identical; the router sees that same hidden state too. A +divergence here means a hook is attached to the wrong module. + +This inspects a real imatrix file, so the path is an argument rather than a hardcoded +dev path. It is a diagnostic, not a unit test -- the assertion-based tests are +test_roundtrip.py / test_calib.py / test_glm4_expert_stats.py / +test_llama4_expert_counts.py, none of which need external files. + +Usage: + python3 check_gate_up.py path/to/some.imatrix [layer ...] +""" +import sys + +import numpy as np + +from imatrix_io import read_imatrix + + +def rel(a, b): + a = a.astype(np.float64).ravel() + b = b.astype(np.float64).ravel() + if a.shape != b.shape: + return float("nan"), float("nan") + corr = float(np.corrcoef(a, b)[0, 1]) if a.std() > 0 and b.std() > 0 else float("nan") + return np.abs(a - b).max() / (np.abs(a).max() + 1e-30), corr + + +def compare(entries, a_name, b_name, label): + if a_name not in entries or b_name not in entries: + print(f" {label}: skipped (missing {a_name if a_name not in entries else b_name})") + return + r, c = rel(entries[a_name]["in_sum2"], entries[b_name]["in_sum2"]) + print(f" {label}: maxrel={r:.2e} corr={c:.6f} (identical if input shared)") + + +def main(argv): + if len(argv) < 2: + print(__doc__) + return 2 + path = argv[1] + layers = [int(x) for x in argv[2:]] if len(argv) > 2 else None + _, e = read_imatrix(path) + + if layers is None: + # derive the layers actually present instead of assuming a fixed model + found = sorted({int(n.split(".")[1]) for n in e + if n.startswith("blk.") and n.split(".")[1].isdigit()}) + layers = found[:: max(1, len(found) // 5)][:5] if found else [] + print(f"{path}: {len(e)} entries; checking layers {layers}") + + for L in layers: + print(f"blk.{L}") + compare(e, f"blk.{L}.ffn_gate_exps.weight", f"blk.{L}.ffn_up_exps.weight", "gate vs up exps") + compare(e, f"blk.{L}.ffn_gate_shexp.weight", f"blk.{L}.ffn_up_shexp.weight", "gate vs up shexp") + compare(e, f"blk.{L}.ffn_gate_shexp.weight", f"blk.{L}.ffn_gate_inp.weight", "shexp-gate vs router") + return 0 + + +if __name__ == "__main__": + sys.exit(main(sys.argv)) diff --git a/scripts/imatrix_serialized/compare.py b/scripts/imatrix_serialized/compare.py new file mode 100644 index 0000000..97fcf6a --- /dev/null +++ b/scripts/imatrix_serialized/compare.py @@ -0,0 +1,25 @@ +import sys, numpy as np +from imatrix_io import read_imatrix +_, A = read_imatrix(sys.argv[1]) # mine +_, B = read_imatrix(sys.argv[2]) # real +na, nb = set(A), set(B) +print(f"mine={len(na)} real={len(nb)} | only-mine={sorted(na-nb)[:5]} only-real={sorted(nb-na)[:5]}") +def norm(e): # per-column mean = sum/count + s = e["in_sum2"].astype(np.float64); c = e["counts"].astype(np.float64).reshape(-1) + if s.ndim==1: return s/max(c[0],1) + return s / np.maximum(c[:,None],1) +corrs, rels = [], [] +worst=[] +for n in sorted(na & nb): + a = norm(A[n]).ravel(); b = norm(B[n]).ravel() + if a.shape!=b.shape: + print(f" SHAPE MISMATCH {n}: {a.shape} vs {b.shape}"); continue + c = np.corrcoef(a,b)[0,1] if a.std()>0 and b.std()>0 else float('nan') + rel = np.abs(a-b).sum()/(np.abs(b).sum()+1e-30) + corrs.append(c); rels.append(rel); worst.append((c,rel,n)) +corrs=np.array(corrs); rels=np.array(rels) +print(f"\nper-tensor correlation: mean={np.nanmean(corrs):.5f} min={np.nanmin(corrs):.5f} median={np.nanmedian(corrs):.5f}") +print(f"per-tensor L1 rel err: mean={np.mean(rels):.4f} max={np.max(rels):.4f}") +print("\nlowest-correlation tensors:") +for c,r,n in sorted(worst)[:6]: + print(f" corr={c:.4f} relerr={r:.3f} {n}") diff --git a/scripts/imatrix_serialized/dispatch.py b/scripts/imatrix_serialized/dispatch.py new file mode 100644 index 0000000..e3cc3d2 --- /dev/null +++ b/scripts/imatrix_serialized/dispatch.py @@ -0,0 +1,149 @@ +"""Pick the right band-serialized generator for a checkpoint, or fail clearly. + +Phase 6 of apex_pipeline.sh used to call serialized_gen.py unconditionally. That file +is GraniteMoeHybrid-only AND transformers-5.5-only, so the wired path ImportError'd on +a newer transformers and could never run the qwen3_5_moe case the backend exists for. + +Selection is by CAPABILITY PROBE, not by version number, because the version ranges +are genuinely misleading: transformers 5.5.1 already ships Qwen3_5MoeExperts and +Llama4Router, but not GraniteMoeHybridExperts (5.14+) and not Glm4MoeExperts. Probing +for the symbol each generator imports is exact and does not rot when 5.x moves on. + +Usage: + dispatch.py --model DIR [--print-script] [... generator args ...] + +With --print-script it prints the chosen script path and exits; otherwise it execs the +generator, passing every argument through unchanged. +""" +import argparse +import importlib +import json +import os +import sys + +HERE = os.path.dirname(os.path.abspath(__file__)) + + +class Variant: + def __init__(self, script, module, symbols, needs, status="validated"): + self.script = script + self.module = module + self.symbols = symbols + self.needs = needs # human-readable requirement, for error messages + self.status = status + + def missing(self): + """Return a list of unavailable requirements ([] means this variant can run).""" + try: + m = importlib.import_module(self.module) + except ImportError as e: + return [f"module {self.module} ({e})"] + return [f"{self.module}.{s}" for s in self.symbols if not hasattr(m, s)] + + +# model_type (from config.json) -> variants in preference order +ARCHES = { + "granitemoehybrid": [ + Variant("serialized_gen_t514.py", + "transformers.models.granitemoehybrid.modeling_granitemoehybrid", + ["GraniteMoeHybridExperts", "GraniteMoeHybridTopKRouter"], + "transformers >= 5.14"), + Variant("serialized_gen.py", + "transformers.models.granitemoehybrid.modeling_granitemoehybrid", + ["GraniteMoeHybridParallelExperts"], + "transformers 5.5.x"), + ], + "qwen3_5_moe": [ + Variant("serialized_gen_qwen35.py", + "transformers.models.qwen3_5_moe.modeling_qwen3_5_moe", + ["Qwen3_5MoeExperts", "Qwen3_5MoeTopKRouter"], + "transformers >= 5.5"), + ], + # Llama-4 configs carry model_type "llama4" at the top level and "llama4_text" under + # text_config; the generator calls cfg.get_text_config(), so both must resolve here. + "llama4": [ + Variant("serialized_gen_llama4.py", + "transformers.models.llama4.modeling_llama4", + ["Llama4TextExperts", "Llama4Router"], + "transformers >= 5.5", status="experimental"), + ], + "llama4_text": [ + Variant("serialized_gen_llama4.py", + "transformers.models.llama4.modeling_llama4", + ["Llama4TextExperts", "Llama4Router"], + "transformers >= 5.5", status="experimental"), + ], + "glm4_moe": [ + Variant("glm4moe_serialized_gen.py", + "transformers.models.glm4_moe.modeling_glm4_moe", + ["Glm4MoeExperts"], + "a transformers with fused Glm4MoeExperts (not 5.5.x)", + status="experimental"), + ], +} + + +def read_model_type(model_dir): + cfg_path = os.path.join(model_dir, "config.json") + if not os.path.exists(cfg_path): + raise SystemExit(f"dispatch: no config.json in {model_dir}") + with open(cfg_path) as f: + cfg = json.load(f) + mt = cfg.get("model_type") + # Multimodal wrappers keep the decoder arch under text_config. Prefer a top-level + # type we actually support, and only fall through to text_config otherwise, so a + # wrapper never resolves to an inner name the table does not know. + if mt not in ARCHES and isinstance(cfg.get("text_config"), dict): + mt = cfg["text_config"].get("model_type", mt) + if not mt: + raise SystemExit(f"dispatch: config.json in {model_dir} has no model_type") + return mt + + +def select(model_type): + """Return (script_path, variant). Raises SystemExit with an actionable message.""" + variants = ARCHES.get(model_type) + if not variants: + raise SystemExit( + f"dispatch: model_type {model_type!r} has no serialized generator.\n" + f" supported: {', '.join(sorted(ARCHES))}\n" + f" use the default IMATRIX_BACKEND=llama for this model, or add a variant.") + tried = [] + for v in variants: + miss = v.missing() + if not miss: + return os.path.join(HERE, v.script), v + tried.append(f" {v.script}: needs {v.needs}; unavailable: {', '.join(miss)}") + raise SystemExit( + f"dispatch: model_type {model_type!r} is supported but no variant matches the " + f"installed transformers.\n" + "\n".join(tried) + + "\n install a transformers matching one of the above " + "(see scripts/imatrix_serialized/requirements.txt).") + + +def main(): + ap = argparse.ArgumentParser(add_help=False) + ap.add_argument("--model", required=True) + ap.add_argument("--print-script", action="store_true") + ap.add_argument("--help", "-h", action="store_true") + args, passthrough = ap.parse_known_args() + if args.help: + print(__doc__) + return + + model_type = read_model_type(args.model) + script, variant = select(model_type) + + if args.print_script: + print(script) + return + + note = "" if variant.status == "validated" else f" [{variant.status}]" + print(f"dispatch: model_type={model_type} -> {os.path.basename(script)}{note}", + flush=True) + argv = [sys.executable, script, "--model", args.model] + passthrough + os.execv(sys.executable, argv) + + +if __name__ == "__main__": + main() diff --git a/scripts/imatrix_serialized/gen_imatrix.py b/scripts/imatrix_serialized/gen_imatrix.py new file mode 100644 index 0000000..2655810 --- /dev/null +++ b/scripts/imatrix_serialized/gen_imatrix.py @@ -0,0 +1,148 @@ +"""Path-A imatrix generator: compute llama.cpp-compatible imatrix from an HF +model via forward hooks (per-input-channel sum of squared activations). + +Memory-bounded: hooks capture only each matmul's input; you can load the model +with accelerate offload (device_map='auto', max_memory=...) and it still works, +because the statistic is per-tensor and additive over tokens. + +Validated against llama.cpp's own imatrix for granite-4.0-h-tiny. +""" +import argparse, re, sys +import numpy as np +import torch +from transformers import AutoModelForCausalLM, AutoTokenizer +from transformers.models.granitemoehybrid.modeling_granitemoehybrid import ( + GraniteMoeHybridParallelExperts, +) +from calib import load_calibration_chunks, add_calib_args +from imatrix_io import write_imatrix, read_imatrix + + +# ---- HF dotted module name -> list of ggml tensor names -------------------- +def map_name(hf: str): + m = re.match(r".*layers\.(\d+)\.(.+)$", hf) + if not m: + return None + i, tail = m.group(1), m.group(2) + b = f"blk.{i}" + table = { + "self_attn.q_proj": [f"{b}.attn_q.weight"], + "self_attn.k_proj": [f"{b}.attn_k.weight"], + "self_attn.v_proj": [f"{b}.attn_v.weight"], + "self_attn.o_proj": [f"{b}.attn_output.weight"], + "mamba.in_proj": [f"{b}.ssm_in.weight"], + "mamba.out_proj": [f"{b}.ssm_out.weight"], + "block_sparse_moe.router.layer": [f"{b}.ffn_gate_inp.weight"], + # ParallelExperts: input_linear feeds BOTH gate and up (same input) + "block_sparse_moe.input_linear": [f"{b}.ffn_gate_exps.weight", f"{b}.ffn_up_exps.weight"], + "block_sparse_moe.output_linear": [f"{b}.ffn_down_exps.weight"], + "shared_mlp.input_linear": [f"{b}.ffn_gate_shexp.weight", f"{b}.ffn_up_shexp.weight"], + "shared_mlp.output_linear": [f"{b}.ffn_down_shexp.weight"], + } + return table.get(tail) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--model", required=True) + ap.add_argument("--calib", required=True) + ap.add_argument("--out", required=True) + add_calib_args(ap) + ap.add_argument("--device", default="cuda") + ap.add_argument("--max-memory", default=None, help="e.g. '10GiB' to force offload (demo memory-bound)") + ap.add_argument("--ground-truth", default=None, help="real imatrix to cross-check name set") + args = ap.parse_args() + + tok = AutoTokenizer.from_pretrained(args.model) + cpu = args.device == "cpu" + dtype = torch.float32 if cpu else torch.bfloat16 + load_kw = dict(torch_dtype=dtype, low_cpu_mem_usage=True) + if args.max_memory: + load_kw.update(device_map="auto", max_memory={0: args.max_memory, "cpu": "200GiB"}) + model = AutoModelForCausalLM.from_pretrained(args.model, **load_kw) + if not args.max_memory: + model = model.to(args.device) + model.eval() + + # accumulators: name -> {"sums": np.float64[nmat,in], "counts": np.int64[nmat]} + acc = {} + known = None + if args.ground_truth: + _, gt = read_imatrix(args.ground_truth) + known = set(gt) + + def ensure(name, nmat, nin): + if name not in acc: + acc[name] = {"sums": np.zeros((nmat, nin), np.float64), "counts": np.zeros(nmat, np.int64)} + return acc[name] + + handles = [] + mapped, skipped = [], [] + for hf_name, mod in model.named_modules(): + ggml = map_name(hf_name) + if ggml is None: + if isinstance(mod, (torch.nn.Linear, GraniteMoeHybridParallelExperts)): + skipped.append(hf_name) + continue + if known is not None: + for g in ggml: + if g not in known: + print(f" WARN mapped name not in ground truth: {g} (from {hf_name})") + mapped.append((hf_name, ggml)) + + if isinstance(mod, GraniteMoeHybridParallelExperts): + n_exp = mod.num_experts + def pre_experts(m, a, names=ggml, ne=n_exp): + inp, expert_size = a[0], a[1] + x = inp.detach().float() # [sum(expert_size), in] + parts = torch.split(x, list(expert_size), dim=0) if sum(expert_size) else [] + for nm in names: + e = ensure(nm, ne, x.shape[-1]) + for ei, p in enumerate(parts): + if p.numel(): + e["sums"][ei] += (p * p).sum(0).double().cpu().numpy() + e["counts"][ei] += p.shape[0] + handles.append(mod.register_forward_pre_hook(pre_experts)) + else: # nn.Linear + def pre_linear(m, a, names=ggml): + x = a[0].detach().float().reshape(-1, a[0].shape[-1]) + s = (x * x).sum(0).double().cpu().numpy() + for nm in names: + e = ensure(nm, 1, x.shape[-1]) + e["sums"][0] += s + e["counts"][0] += x.shape[0] + handles.append(mod.register_forward_pre_hook(pre_linear)) + + print(f"hooked {len(mapped)} modules; skipped (not mapped) Linears: {skipped}") + + # tokenize whole calibration file, window into ctx-sized chunks + # Same windowing/BOS handling as the band-serialized generators, so the two stay + # comparable -- serialized_gen.py's parity claim is against THIS generator. + chunks_t, calib_info = load_calibration_chunks(tok, args.calib, args.chunks, args.ctx, + add_bos=args.bos_per_chunk) + nchunks = chunks_t.shape[0] + print(f"calib: {calib_info}") + dev = 0 if args.max_memory else args.device + print(f"device={args.device} dtype={dtype}") + with torch.no_grad(): + for c in range(nchunks): + chunk = chunks_t[c].unsqueeze(0).to(dev) + model(chunk) + if (c + 1) % 20 == 0: + print(f" chunk {c+1}/{nchunks}") + + for h in handles: + h.remove() + + # write imatrix: sums stored raw, counts per matrix + entries = {} + for name, d in acc.items(): + sums = d["sums"].astype(np.float32) # [nmat, in] + counts = d["counts"].astype(np.float32) # [nmat] + entries[name] = {"in_sum2": sums if sums.shape[0] > 1 else sums[0], "counts": counts} + write_imatrix(args.out, entries, ["calibration_datav3"], nchunks, args.ctx) + print(f"wrote {len(entries)} entries -> {args.out}") + + +if __name__ == "__main__": + main() diff --git a/scripts/imatrix_serialized/gen_imatrix_t514.py b/scripts/imatrix_serialized/gen_imatrix_t514.py new file mode 100644 index 0000000..48acd29 --- /dev/null +++ b/scripts/imatrix_serialized/gen_imatrix_t514.py @@ -0,0 +1,216 @@ +"""torch-imatrix generator, transformers-5.14 variant. + +Same statistic and output format as gen_imatrix.py (the validated +transformers-5.5.1 tool), adapted to the 5.14 GraniteMoeHybrid rewrite: + + * `GraniteMoeHybridParallelExperts` (fused input_linear/output_linear, called + with (inputs, expert_size)) was replaced by `GraniteMoeHybridExperts`, a + module holding `gate_up_proj`/`down_proj` *parameters* and doing per-expert + `F.linear` inside its forward. There are no expert Linear submodules to + pre-hook, so we hook the Experts module itself, replay the router's + per-expert token grouping, and accumulate Sum x^2 for: + - gate/up experts input = hidden_states[tokens routed to e] (dim hidden) + - down experts input = act_fn(gate) * up (recomputed) (dim inter) + gate and up share one input (identical stats), matching the original. + * shared MLP still uses real nn.Linear (input_linear/output_linear) and the + attention/mamba projections are real nn.Linear -> handled generically. + * the router in 5.14 is a bare nn.Parameter (F.linear), no submodule; its + ggml tensor `ffn_gate_inp` is stored F32 and does not consume an imatrix + (confirmed: it is byte-identical across imatrices), so we skip it. + +Mapping is derived from the *loaded* model's named_modules() and cross-checked +against a ground-truth imatrix tensor set; any unmapped Linear or name not in +ground truth aborts before the (expensive) calibration run. +""" +import argparse, re, sys +import numpy as np +import torch +from transformers import AutoModelForCausalLM, AutoTokenizer +from transformers.models.granitemoehybrid.modeling_granitemoehybrid import ( + GraniteMoeHybridExperts, + GraniteMoeHybridTopKRouter, +) +from calib import load_calibration_chunks, add_calib_args +from imatrix_io import write_imatrix, read_imatrix + + +def map_linear(hf: str): + """HF dotted module name (an nn.Linear) -> list of ggml tensor names, or None.""" + m = re.match(r".*layers\.(\d+)\.(.+)$", hf) + if not m: + return None + i, tail = m.group(1), m.group(2) + b = f"blk.{i}" + table = { + "self_attn.q_proj": [f"{b}.attn_q.weight"], + "self_attn.k_proj": [f"{b}.attn_k.weight"], + "self_attn.v_proj": [f"{b}.attn_v.weight"], + "self_attn.o_proj": [f"{b}.attn_output.weight"], + "mamba.in_proj": [f"{b}.ssm_in.weight"], + "mamba.out_proj": [f"{b}.ssm_out.weight"], + # shared MLP: input_linear feeds BOTH gate and up (same input) + "shared_mlp.input_linear": [f"{b}.ffn_gate_shexp.weight", f"{b}.ffn_up_shexp.weight"], + "shared_mlp.output_linear": [f"{b}.ffn_down_shexp.weight"], + } + return table.get(tail) + + +def expert_names(hf: str): + """HF dotted name of a GraniteMoeHybridExperts module -> per-role ggml names.""" + m = re.match(r".*layers\.(\d+)\.", hf) + b = f"blk.{m.group(1)}" + return { + "gate": f"{b}.ffn_gate_exps.weight", + "up": f"{b}.ffn_up_exps.weight", + "down": f"{b}.ffn_down_exps.weight", + } + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--model", required=True) + ap.add_argument("--calib", required=True) + ap.add_argument("--out", required=True) + add_calib_args(ap) + ap.add_argument("--device", default="cuda") + ap.add_argument("--max-memory", default=None, help="e.g. '12GiB' to cap GPU and offload rest") + ap.add_argument("--dataset-label", default="calibration_datav3") + ap.add_argument("--ground-truth", default=None, help="real imatrix to cross-check name set") + args = ap.parse_args() + + tok = AutoTokenizer.from_pretrained(args.model) + cpu = args.device == "cpu" + dtype = torch.float32 if cpu else torch.bfloat16 + load_kw = dict(torch_dtype=dtype, low_cpu_mem_usage=True) + if args.max_memory: + load_kw.update(device_map="auto", max_memory={0: args.max_memory, "cpu": "200GiB"}) + model = AutoModelForCausalLM.from_pretrained(args.model, **load_kw) + if not args.max_memory: + model = model.to(args.device) + model.eval() + + known = None + if args.ground_truth: + _, gt = read_imatrix(args.ground_truth) + known = set(gt) + + acc = {} # name -> {"sums": f64[nmat,in], "counts": i64[nmat]} + def ensure(name, nmat, nin): + if name not in acc: + acc[name] = {"sums": np.zeros((nmat, nin), np.float64), "counts": np.zeros(nmat, np.int64)} + return acc[name] + + def check(ggml_names): + if known is None: + return + for g in ggml_names: + if g not in known: + print(f" WARN mapped name not in ground truth: {g}") + + handles = [] + mapped_lin, mapped_exp, skipped = [], [], [] + for hf_name, mod in model.named_modules(): + if isinstance(mod, torch.nn.Linear): + names = map_linear(hf_name) + if names is None: + skipped.append(hf_name) + continue + check(names) + mapped_lin.append(hf_name) + def pre_linear(m, a, names=names): + x = a[0].detach().float().reshape(-1, a[0].shape[-1]) + s = (x * x).sum(0).double().cpu().numpy() + for nm in names: + e = ensure(nm, 1, x.shape[-1]) + e["sums"][0] += s + e["counts"][0] += x.shape[0] + handles.append(mod.register_forward_pre_hook(pre_linear)) + elif isinstance(mod, GraniteMoeHybridExperts): + en = expert_names(hf_name) + check(list(en.values())) + mapped_exp.append(hf_name) + ne = mod.num_experts + def pre_experts(m, a, en=en, ne=ne): + # forward(hidden_states, top_k_index, top_k_weights) + hidden_states, top_k_index = a[0], a[1] + x = hidden_states.detach().float() # [tokens, hidden] + gate = ensure(en["gate"], ne, x.shape[-1]) + up = ensure(en["up"], ne, x.shape[-1]) + # down input dim = intermediate; lazily sized on first expert + down = None + with torch.no_grad(): + mask = torch.nn.functional.one_hot(top_k_index, num_classes=ne).permute(2, 1, 0) + hit = torch.greater(mask.sum(dim=(-1, -2)), 0).nonzero() + for eidx in hit: + e = eidx[0] + if e == ne: + continue + _, tok_idx = torch.where(mask[e]) + cur = x[tok_idx] # [n_e, hidden] gate/up input + s = (cur * cur).sum(0).double().cpu().numpy() + gate["sums"][e] += s; gate["counts"][e] += cur.shape[0] + up["sums"][e] += s; up["counts"][e] += cur.shape[0] + # recompute down input exactly as the module does + gu = torch.nn.functional.linear(cur.to(m.gate_up_proj.dtype), m.gate_up_proj[e]) + g, u = gu.chunk(2, dim=-1) + inter = (m.act_fn(g) * u).float() # [n_e, inter] down input + nonlocal_down = ensure(en["down"], ne, inter.shape[-1]) + sd = (inter * inter).sum(0).double().cpu().numpy() + nonlocal_down["sums"][e] += sd; nonlocal_down["counts"][e] += inter.shape[0] + handles.append(mod.register_forward_pre_hook(pre_experts)) + elif isinstance(mod, GraniteMoeHybridTopKRouter): + # router is a bare nn.Parameter (F.linear) in 5.14, not a Linear + # submodule; its input maps to ffn_gate_inp (stored F32 in the quant, + # imatrix-unused, but present in the shipped imatrix -> capture it). + mr = re.match(r".*layers\.(\d+)\.", hf_name) + gname = f"blk.{mr.group(1)}.ffn_gate_inp.weight" + check([gname]) + mapped_lin.append(hf_name) + def pre_router(m, a, gname=gname): + x = a[0].detach().float().reshape(-1, a[0].shape[-1]) + s = (x * x).sum(0).double().cpu().numpy() + e = ensure(gname, 1, x.shape[-1]) + e["sums"][0] += s + e["counts"][0] += x.shape[0] + handles.append(mod.register_forward_pre_hook(pre_router)) + + print(f"hooked {len(mapped_lin)} Linear/router + {len(mapped_exp)} Experts modules") + # lm_head is tied to token_embd (tie_word_embeddings) and quantized without an + # imatrix — the original 5.5.1 tool skipped it too. Any *other* unmapped Linear + # is a real gap and aborts before the expensive run. + INTENTIONAL_SKIP = ("lm_head",) + fatal = [s for s in skipped if not s.endswith(INTENTIONAL_SKIP)] + if skipped: + print(f"skipped (intentional) Linears: {skipped}") + if fatal: + print(f"UNMAPPED Linears ({len(fatal)}): {fatal}") + sys.exit("ABORT: unmapped Linear tensors — fix map_linear() before running") + + # Same windowing/BOS handling as the band-serialized generators, so the two stay + # comparable -- the serialized generators' parity claim is against THIS one. + chunks_t, calib_info = load_calibration_chunks(tok, args.calib, args.chunks, args.ctx, + add_bos=args.bos_per_chunk) + nchunks = chunks_t.shape[0] + print(f"calib: {calib_info}; device={args.device} dtype={dtype}") + dev = 0 if args.max_memory else args.device + with torch.no_grad(): + for c in range(nchunks): + chunk = chunks_t[c].unsqueeze(0).to(dev) + model(chunk) + if (c + 1) % 10 == 0: + print(f" chunk {c+1}/{nchunks}", flush=True) + + for h in handles: + h.remove() + + entries = {} + for name, d in acc.items(): + sums = d["sums"].astype(np.float32) + counts = d["counts"].astype(np.float32) + entries[name] = {"in_sum2": sums if sums.shape[0] > 1 else sums[0], "counts": counts} + write_imatrix(args.out, entries, [args.dataset_label], nchunks, args.ctx) + print(f"wrote {len(entries)} entries -> {args.out}") + + +if __name__ == "__main__": + main() diff --git a/scripts/imatrix_serialized/glm4moe_serialized_gen.py b/scripts/imatrix_serialized/glm4moe_serialized_gen.py new file mode 100644 index 0000000..07f1092 --- /dev/null +++ b/scripts/imatrix_serialized/glm4moe_serialized_gen.py @@ -0,0 +1,332 @@ +"""Band-serialized (memory-bounded) Path-A imatrix generator for GLM-4.5-Air +(zai-org/GLM-4.5-Air, model_type glm4_moe). + +GLM-4.5-Air facts: + 46 layers, hidden_size 4096, head_dim 128, partial_rotary_factor 0.5 + 128 routed experts per layer, 1 shared expert, 8 active experts per token + moe_intermediate_size = 1408 (NOT divisible by 256 -> K-quants fall back) + first_k_dense_replace = 1 -> layer 0 uses a dense MLP (Glm4MoeMLP), rest use Glm4MoeMoE + +Peak memory is band-bounded: only `--band` layers resident on GPU at once. +Each weight tensor is read from safetensors once; activations are cached CPU-side +between bands. + +EXPERIMENTAL / UNTESTED against a real GLM-4.5 checkpoint. Requires a transformers +version whose glm4_moe module exposes the fused `Glm4MoeExperts` class (it does not +exist in 5.5.x, which packages Glm4MoeMoE/Glm4MoeNaiveMoe instead); the import is +guarded so the failure names the reason instead of raising a bare ImportError. +""" +import argparse, json, os, glob, re +import numpy as np +import torch +from transformers import AutoConfig, AutoTokenizer +try: + from transformers.models.glm4_moe.modeling_glm4_moe import ( + Glm4MoeDecoderLayer, Glm4MoeMoE, Glm4MoeExperts, + Glm4MoeMLP, Glm4MoeRotaryEmbedding, + ) + _ARCH_IMPORT_ERR = None +except ImportError as _e: # keep the pure math importable/testable + Glm4MoeDecoderLayer = Glm4MoeMoE = Glm4MoeExperts = None + Glm4MoeMLP = Glm4MoeRotaryEmbedding = None + _ARCH_IMPORT_ERR = _e +from transformers.masking_utils import create_causal_mask +from transformers.activations import ACT2FN +from safetensors import safe_open + +import imatrix_io +from calib import load_calibration_chunks, add_calib_args +from imatrix_io import write_imatrix, read_imatrix + + +# ---- HF dotted module name -> list of ggml tensor names -------------------- +def map_name(hf: str): + m = re.match(r".*layers\.(\d+)\.(.+)$", hf) + if not m: + return None + i, tail = m.group(1), m.group(2) + b = f"blk.{i}" + table = { + # attention + "self_attn.q_proj": [f"{b}.attn_q.weight"], + "self_attn.k_proj": [f"{b}.attn_k.weight"], + "self_attn.v_proj": [f"{b}.attn_v.weight"], + "self_attn.o_proj": [f"{b}.attn_output.weight"], + # dense MLP (layer 0 with first_k_dense_replace=1) + "mlp.gate_proj": [f"{b}.ffn_gate.weight"], + "mlp.up_proj": [f"{b}.ffn_up.weight"], + "mlp.down_proj": [f"{b}.ffn_down.weight"], + # MoE gate + "mlp.gate": [f"{b}.ffn_gate_inp.weight"], + # NOTE: "mlp.experts" is deliberately absent -- the three routed-expert + # tensors do NOT share one input stat (gate/up take hidden_size inputs, + # down takes moe_intermediate_size gated-intermediate inputs), so they are + # resolved by expert_names() and handled by the dedicated hook below. + # shared experts + "mlp.shared_experts.gate_proj": [f"{b}.ffn_gate_shexp.weight"], + "mlp.shared_experts.up_proj": [f"{b}.ffn_up_shexp.weight"], + "mlp.shared_experts.down_proj": [f"{b}.ffn_down_shexp.weight"], + } + return table.get(tail) + + +def expert_names(hf: str): + """ggml names for the three routed-expert tensors of an mlp.experts module.""" + m = re.match(r".*layers\.(\d+)\.", hf) + if not m: + raise ValueError(f"no layer index in {hf!r}") + b = f"blk.{m.group(1)}" + return {"gate": f"{b}.ffn_gate_exps.weight", "up": f"{b}.ffn_up_exps.weight", + "down": f"{b}.ffn_down_exps.weight"} + + +def apply_gate_up(x, w_e, hidden): + """Apply one expert's fused gate_up weight to rows x [T, hidden] -> [T, 2*interm]. + + Fused MoE weights appear in two layouts across transformers arch modules: + [2*interm, hidden] -- nn.Linear convention (F.linear) + [hidden, 2*interm] -- bmm convention (x @ W) + GLM-4.5's layout is transformers-version dependent, so derive it from the shape + and hard-fail rather than guess: guessing wrong silently yields a garbage + ffn_down_exps statistic, which is exactly the failure mode this path must avoid. + """ + if w_e.ndim != 2: + raise ValueError(f"expected 2-D per-expert gate_up weight, got {tuple(w_e.shape)}") + a, b = w_e.shape + if a == b: + raise ValueError(f"ambiguous square gate_up layout {tuple(w_e.shape)} (hidden={hidden})") + if b == hidden: + return torch.nn.functional.linear(x, w_e) + if a == hidden: + return x @ w_e + raise ValueError(f"gate_up weight {tuple(w_e.shape)} has no hidden={hidden} axis") + + +def expert_stats(x, topk_idx, gate_up, act_fn, num_experts, hidden): + """Per-expert Σx² for the three routed-expert tensors of one MoE layer. + + Returns (sum_in2 [E, hidden], sum_inter2 [E, interm], counts [E]) where + * sum_in2 feeds ffn_gate_exps and ffn_up_exps -- the hidden-dim input each + selected expert's gate/up MUL_MAT_ID actually sees; + * sum_inter2 feeds ffn_down_exps -- act_fn(gate)*up, the gated intermediate + that is the down projection's input. It is NOT the hidden-dim + stat and has moe_intermediate_size entries, not hidden_size. + * counts is the number of tokens routed to each expert, i.e. the row count + of every one of the three matmuls (llama.cpp counts one per + (token, selected-expert) slot). + + Pure tensor math, no module dependency, so it is testable without a checkpoint. + """ + x = x.reshape(-1, x.shape[-1]) + T = x.shape[0] + if x.shape[-1] != hidden: + raise ValueError(f"expert input dim {x.shape[-1]} != hidden {hidden}") + dev = x.device + xf = x.float() + k = topk_idx.shape[-1] + assign = topk_idx.reshape(-1).long() + tok = torch.arange(T, device=dev).repeat_interleave(k) + + sum_in2 = torch.zeros(num_experts, hidden, dtype=torch.float64, device=dev) + counts = torch.zeros(num_experts, dtype=torch.float64, device=dev) + sq = (xf * xf).index_select(0, tok).double() + sum_in2.index_add_(0, assign, sq) + counts.index_add_(0, assign, torch.ones(assign.numel(), dtype=torch.float64, device=dev)) + + # gate_up is [E, 2*interm, hidden] or [E, hidden, 2*interm]; the non-hidden axis + # is the fused output, half of which is the gated intermediate. + fused = gate_up.shape[-2] if gate_up.shape[-1] == hidden else gate_up.shape[-1] + sum_inter2 = torch.zeros(num_experts, fused // 2, dtype=torch.float64, device=dev) + for e in torch.unique(assign).tolist(): + rows = (topk_idx == e).any(dim=-1).nonzero(as_tuple=True)[0] + cur = x.index_select(0, rows).to(gate_up.dtype) + g, u = apply_gate_up(cur, gate_up[e], hidden).chunk(2, dim=-1) + inter = (act_fn(g) * u).float() + sum_inter2[e] += (inter * inter).sum(0).double() + return sum_in2, sum_inter2, counts + + +class ShardReader: + def __init__(self, d): + self.dir = d + idx = os.path.join(d, "model.safetensors.index.json") + if os.path.exists(idx): + with open(idx) as f: + self.wm = json.load(f)["weight_map"] + else: + f = os.path.basename(glob.glob(os.path.join(d, "*.safetensors"))[0]) + with safe_open(os.path.join(d, f), framework="pt") as sf: + self.wm = {k: f for k in sf.keys()} + self.handles = {} + + def get(self, name): + fn = self.wm[name] + if fn not in self.handles: + self.handles[fn] = safe_open(os.path.join(self.dir, fn), framework="pt") + return self.handles[fn].get_tensor(name) + + +def materialize(module, prefix, reader, dtype, device): + # Every module parameter MUST come from the checkpoint -- see the note in + # serialized_gen.py: a skipped key leaves random weights in the band and the + # resulting imatrix is meaningless with no error anywhere. + sd, missing = {}, [] + for k in module.state_dict().keys(): + full = f"{prefix}.{k}" + if full in reader.wm: + sd[k] = reader.get(full).to(dtype) + else: + missing.append(k) + if missing: + raise KeyError(f"{prefix}: no checkpoint tensor for {missing}") + module.load_state_dict(sd, strict=False, assign=True) + return module.to(device).eval() + + +def make_hooks(named_mods, acc, known, cfg): + def ensure(name, nmat, nin): + if name not in acc: + acc[name] = {"sums": np.zeros((nmat, nin), np.float64), "counts": np.zeros(nmat, np.int64)} + return acc[name] + handles = [] + for hf_name, mod in named_mods: + if Glm4MoeExperts is not None and isinstance(mod, Glm4MoeExperts): + en = expert_names(hf_name) + if known is not None: + for g in en.values(): + if g not in known: + print(f" WARN unmapped: {g} ({hf_name})") + n_exp = mod.num_experts + act = getattr(mod, "act_fn", None) or ACT2FN[cfg.hidden_act] + # forward signature: (hidden_states, top_k_index, top_k_weights) + def pre_experts(m, a, en=en, ne=n_exp, act=act): + inp, topk_idx = a[0], a[1] + x = inp.detach() # [tokens, hidden] + s_in, s_inter, cnt = expert_stats( + x, topk_idx, m.gate_up_proj, act, ne, cfg.hidden_size) + s_in = s_in.cpu().numpy(); s_inter = s_inter.cpu().numpy() + cnt = cnt.cpu().numpy().astype(np.int64) + for nm in (en["gate"], en["up"]): + e = ensure(nm, ne, s_in.shape[-1]) + e["sums"] += s_in + e["counts"] += cnt + e = ensure(en["down"], ne, s_inter.shape[-1]) + e["sums"] += s_inter + e["counts"] += cnt + handles.append(mod.register_forward_pre_hook(pre_experts)) + continue + ggml = map_name(hf_name) + if ggml is None: + continue + if known is not None: + for g in ggml: + if g not in known: + print(f" WARN unmapped: {g} ({hf_name})") + def pre_linear(m, a, names=ggml): + x = a[0].detach().float().reshape(-1, a[0].shape[-1]) + s = (x * x).sum(0).double().cpu().numpy() + for nm in names: + e = ensure(nm, 1, x.shape[-1]) + e["sums"][0] += s + e["counts"][0] += x.shape[0] + handles.append(mod.register_forward_pre_hook(pre_linear)) + return handles + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--model", required=True) + ap.add_argument("--calib", required=True) + ap.add_argument("--out", required=True) + add_calib_args(ap) + ap.add_argument("--band", type=int, default=1, help="layers resident at once (memory knob)") + ap.add_argument("--device", default="cuda") + ap.add_argument("--ground-truth", default=None) + ap.add_argument("--attn-impl", default="eager", help="eager/sdpa/flash_attention_2") + # OFF by default: trust_remote_code executes arbitrary Python from the model repo, + # and the pipeline downloads whatever MODEL_ID a config names. GLM-4.5 is supported + # natively by transformers, so this is only needed for a repo shipping custom code. + ap.add_argument("--trust-remote-code", action="store_true", + help="allow executing model-repo code in AutoConfig/AutoTokenizer") + args = ap.parse_args() + + if _ARCH_IMPORT_ERR is not None: + raise RuntimeError( + "glm4_moe generator needs a transformers whose modeling_glm4_moe exposes " + f"the fused Glm4MoeExperts class; installed transformers does not ({_ARCH_IMPORT_ERR})") + + cfg = AutoConfig.from_pretrained(args.model, trust_remote_code=args.trust_remote_code) + cfg._attn_implementation = args.attn_impl + dev = args.device + dtype = torch.bfloat16 if dev == "cuda" else torch.float32 + reader = ShardReader(args.model) + known = set(read_imatrix(args.ground_truth)[1]) if args.ground_truth else None + nlayers = cfg.num_hidden_layers + d = cfg.hidden_size + + tok = AutoTokenizer.from_pretrained(args.model, trust_remote_code=args.trust_remote_code) + chunks, calib_info = load_calibration_chunks(tok, args.calib, args.chunks, args.ctx, + add_bos=args.bos_per_chunk) + nch = chunks.shape[0] + print(f"calib: {calib_info}") + print(f"layers={nlayers} band={args.band} chunks={nch} ctx={args.ctx} dtype={dtype} attn={args.attn_impl}") + + # --- embed all chunks -> hidden cache on CPU --- + embed_w = reader.get("model.embed_tokens.weight").to(dtype) + hid = torch.nn.functional.embedding(chunks, embed_w).to(dtype) # [nch, ctx, d] CPU + del embed_w + print(f"hidden cache: {tuple(hid.shape)} ({hid.element_size()*hid.nelement()/1e6:.0f} MB CPU)") + + # --- rotary + causal mask --- + rotary = Glm4MoeRotaryEmbedding(cfg).to(dev) + pos_ids = torch.arange(args.ctx, device=dev).unsqueeze(0) + dummy = hid[:1].to(dev) + pos_emb = rotary(dummy, pos_ids) + causal = create_causal_mask(config=cfg, inputs_embeds=dummy, attention_mask=None, past_key_values=None) + + acc = {} + if dev == "cuda": + torch.cuda.reset_peak_memory_stats() + + # --- band loop --- + with torch.no_grad(): + for b0 in range(0, nlayers, args.band): + band = list(range(b0, min(b0 + args.band, nlayers))) + layers = [] + for i in band: + L = Glm4MoeDecoderLayer(cfg, layer_idx=i) + materialize(L, f"model.layers.{i}", reader, dtype, dev) + layers.append((i, L)) + named = [(f"model.layers.{i}.{n}", m) for i, L in layers for n, m in L.named_modules()] + handles = make_hooks(named, acc, known, cfg) + for c in range(nch): + h = hid[c:c+1].to(dev) + for i, L in layers: + h = L(h, attention_mask=causal, position_embeddings=pos_emb, use_cache=False) + hid[c:c+1] = h.to(hid.dtype).cpu() + for hh in handles: + hh.remove() + # Drop EVERY reference to the band's modules before empty_cache(), or the + # cache release is a no-op: `named` holds the submodules and `del L` only + # unbinds the loop variable. + handles.clear(); named.clear(); layers.clear() + if dev == "cuda": + torch.cuda.empty_cache() + print(f" band {band[0]}-{band[-1]} done") + + # write + entries = {} + for name, dd in acc.items(): + sums = dd["sums"].astype(np.float32) + entries[name] = {"in_sum2": sums if sums.shape[0] > 1 else sums[0], + "counts": dd["counts"].astype(np.float32)} + write_imatrix(args.out, entries, ["calibration_sample"], nch, args.ctx) + + line = f"wrote {len(entries)} entries -> {args.out}" + if dev == "cuda": + line += f" | peak GPU {torch.cuda.max_memory_allocated()/1e9:.2f} GB" + print(line) + + +if __name__ == "__main__": + main() diff --git a/scripts/imatrix_serialized/imatrix_io.py b/scripts/imatrix_serialized/imatrix_io.py new file mode 100644 index 0000000..1d92e7a --- /dev/null +++ b/scripts/imatrix_serialized/imatrix_io.py @@ -0,0 +1,80 @@ +"""Read/write llama.cpp GGUF-format imatrix files. + +Format (from tools/imatrix/imatrix.cpp save_imatrix + common/imatrix-loader.cpp): + KV: general.type="imatrix", imatrix.datasets(str[]), + imatrix.chunk_count(u32), imatrix.chunk_size(u32) + Per weight tensor : + .in_sum2 : F32 [in_features, nmat] = sum over tokens of x_j^2 (raw) + .counts : F32 [1, nmat] = token count per matrix/expert +""" +import numpy as np +import gguf + + +def read_imatrix(path): + r = gguf.GGUFReader(path) + meta = {} + for k, f in r.fields.items(): + try: + if f.types and f.types[0] == gguf.GGUFValueType.STRING: + meta[k] = [str(bytes(f.parts[idx]), "utf-8") for idx in f.data] \ + if len(f.data) > 1 else str(bytes(f.parts[f.data[0]]), "utf-8") + else: + meta[k] = f.parts[f.data[0]].tolist() if f.data else None + except Exception: + meta[k] = "" + entries = {} + for t in r.tensors: + name = t.name + for suf in (".in_sum2", ".counts"): + if name.endswith(suf): + base = name[: -len(suf)] + entries.setdefault(base, {})[suf[1:]] = np.array(t.data, dtype=np.float32).reshape(tuple(reversed(t.shape.tolist()))) + return meta, entries + + +def write_imatrix(path, entries, datasets, chunk_count, chunk_size): + """entries: {base_name: {"in_sum2": np.ndarray[nmat, in_feat] or [in_feat], + "counts": np.ndarray[nmat] or scalar}}""" + w = gguf.GGUFWriter(path, arch="imatrix") # arch unused; general.type set below + w.add_type("imatrix") + w.add_array("imatrix.datasets", datasets) + w.add_uint32("imatrix.chunk_count", int(chunk_count)) + w.add_uint32("imatrix.chunk_size", int(chunk_size)) + for base in sorted(entries): + e = entries[base] + sums = np.asarray(e["in_sum2"], dtype=np.float32) + counts = np.asarray(e["counts"], dtype=np.float32) + if sums.ndim == 1: + sums = sums[None, :] # [1, in_feat] + counts = counts.reshape(-1) # [nmat] + # GGUF tensor dims are stored reversed; add_tensor takes numpy with + # shape [nmat, in_feat] -> ggml ne = [in_feat, nmat] + nmat = counts.shape[0] + w.add_tensor(base + ".in_sum2", np.ascontiguousarray(sums, dtype=np.float32)) + w.add_tensor(base + ".counts", np.ascontiguousarray(counts.reshape(nmat, 1), dtype=np.float32)) + w.write_header_to_file() + w.write_kv_data_to_file() + w.write_tensors_to_file() + w.close() + + +if __name__ == "__main__": + import sys + meta, entries = read_imatrix(sys.argv[1]) + print("=== metadata ===") + for k, v in meta.items(): + vs = v if not isinstance(v, list) else f"[{len(v)} items] {v[:2]}" + print(f" {k}: {vs}") + print(f"=== {len(entries)} tensor entries ===") + for i, (base, e) in enumerate(sorted(entries.items())): + s = e.get("in_sum2"); c = e.get("counts") + if i < 12 or i >= len(entries) - 4: + print(f" {base:42s} in_sum2{tuple(s.shape) if s is not None else None} " + f"counts{tuple(c.shape) if c is not None else None} " + f"count0={c.flatten()[0] if c is not None else '?':.0f}") + elif i == 12: + print(" ...") + # sanity: any entry missing a half? + bad = [b for b, e in entries.items() if "in_sum2" not in e or "counts" not in e] + print("incomplete entries:", bad if bad else "none") diff --git a/scripts/imatrix_serialized/requirements.txt b/scripts/imatrix_serialized/requirements.txt new file mode 100644 index 0000000..97a0790 --- /dev/null +++ b/scripts/imatrix_serialized/requirements.txt @@ -0,0 +1,30 @@ +# IMPORTANT: install torch from the CUDA index, not the default CPU wheel: +# pip install torch --index-url https://download.pytorch.org/whl/cu128 # pick cuXXX for your CUDA +# A CPU-only torch runs the generator on CPU (very slow; hybrid-SSM especially). +# Verify: python3 -c "import torch; assert torch.cuda.is_available()" +# +# Optional deps for IMATRIX_BACKEND=serialized only (not needed for the default llama backend). +# +# transformers: there is NO single version that runs every generator, because the +# arch modules were refactored at different times. Do not read the ranges below as a +# version to pin globally -- dispatch.py probes the INSTALLED transformers for the +# symbols each variant imports and picks a variant that works (or tells you what is +# missing), so install whichever matches the model you are quantizing: +# +# arch (config.json model_type) generator needs +# ----------------------------- ---------------------------- -------------------------- +# granitemoehybrid serialized_gen.py 5.5.x (GraniteMoeHybridParallelExperts) +# granitemoehybrid serialized_gen_t514[_fast].py >= 5.14 (GraniteMoeHybridExperts) +# qwen3_5_moe serialized_gen_qwen35.py >= 5.5 (Qwen3_5MoeExperts) +# llama4 serialized_gen_llama4.py >= 5.5 (Llama4TextExperts) +# glm4_moe glm4moe_serialized_gen.py a version with fused Glm4MoeExperts +# +# `transformers>=5.5` below is a FLOOR, not a claim that any 5.5+ runs everything: +# on 5.5.x the Granite path resolves to serialized_gen.py and glm4_moe is unavailable. +# Check what your install supports with: +# python3 scripts/imatrix_serialized/dispatch.py --model --print-script +torch>=2.4 +transformers>=5.5 +accelerate>=1.0 +numpy +gguf diff --git a/scripts/imatrix_serialized/serialized_gen.py b/scripts/imatrix_serialized/serialized_gen.py new file mode 100644 index 0000000..4b65b7e --- /dev/null +++ b/scripts/imatrix_serialized/serialized_gen.py @@ -0,0 +1,200 @@ +"""Band-serialized (memory-bounded) Path-A imatrix generator. + +Instead of a full forward (whole model resident), process the model in BANDS of +`--band` layers: load a band's weights from the safetensors shards ONCE, batch +ALL calibration chunks through it, cache inter-band activations, free the band. + +Peak resident weights = band_size layers (not the whole model). Each weight is +read from disk exactly once. Mathematically identical to the full forward +(each decoder block is self-contained given its input activations). + +Validated to match the full-forward Path-A imatrix on granite-4.0-h-tiny. +""" +import argparse, json, os, glob, resource +import numpy as np, torch +from transformers import AutoConfig, AutoTokenizer +from transformers.models.granitemoehybrid.modeling_granitemoehybrid import ( + GraniteMoeHybridDecoderLayer, GraniteMoeHybridRotaryEmbedding, + GraniteMoeHybridParallelExperts, +) +from transformers.masking_utils import create_causal_mask +from safetensors import safe_open +from calib import load_calibration_chunks, add_calib_args +from imatrix_io import write_imatrix, read_imatrix +from gen_imatrix import map_name + + +class ShardReader: + def __init__(self, d): + self.dir = d + idx = os.path.join(d, "model.safetensors.index.json") + if os.path.exists(idx): + with open(idx) as f: + self.wm = json.load(f)["weight_map"] + else: + f = os.path.basename(glob.glob(os.path.join(d, "*.safetensors"))[0]) + with safe_open(os.path.join(d, f), framework="pt") as sf: + self.wm = {k: f for k in sf.keys()} + self.handles = {} + + def get(self, name): + fn = self.wm[name] + if fn not in self.handles: + self.handles[fn] = safe_open(os.path.join(self.dir, fn), framework="pt") + return self.handles[fn].get_tensor(name) + + +def materialize(module, prefix, reader, dtype, device): + # Every module parameter MUST come from the checkpoint. Silently skipping a + # missing key would leave randomly-initialized weights in the band and produce + # a plausible-looking but meaningless imatrix, so a name mismatch is fatal. + sd, missing = {}, [] + for k in module.state_dict().keys(): + full = f"{prefix}.{k}" + if full in reader.wm: + sd[k] = reader.get(full).to(dtype) + else: + missing.append(k) + if missing: + raise KeyError(f"{prefix}: no checkpoint tensor for {missing}") + module.load_state_dict(sd, strict=False, assign=True) + return module.to(device).eval() + + +def make_hooks(named_mods, acc, known): + def ensure(name, nmat, nin): + if name not in acc: + acc[name] = {"sums": np.zeros((nmat, nin), np.float64), "counts": np.zeros(nmat, np.int64)} + return acc[name] + handles = [] + for hf_name, mod in named_mods: + ggml = map_name(hf_name) + if ggml is None: + continue + if known is not None: + for g in ggml: + if g not in known: + print(f" WARN unmapped: {g} ({hf_name})") + if isinstance(mod, GraniteMoeHybridParallelExperts): + def pre(m, a, names=ggml, ne=mod.num_experts): + inp, esz = a[0], a[1] + x = inp.detach().float() + parts = torch.split(x, list(esz), dim=0) if sum(esz) else [] + for nm in names: + e = ensure(nm, ne, x.shape[-1]) + for ei, p in enumerate(parts): + if p.numel(): + e["sums"][ei] += (p * p).sum(0).double().cpu().numpy() + e["counts"][ei] += p.shape[0] + handles.append(mod.register_forward_pre_hook(pre)) + else: + def pre(m, a, names=ggml): + x = a[0].detach().float().reshape(-1, a[0].shape[-1]) + s = (x * x).sum(0).double().cpu().numpy() + for nm in names: + e = ensure(nm, 1, x.shape[-1]) + e["sums"][0] += s + e["counts"][0] += x.shape[0] + handles.append(mod.register_forward_pre_hook(pre)) + return handles + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--model", required=True) + ap.add_argument("--calib", required=True) + ap.add_argument("--out", required=True) + add_calib_args(ap) + ap.add_argument("--band", type=int, default=4, help="layers resident at once (memory knob)") + ap.add_argument("--device", default="cuda") + ap.add_argument("--ground-truth", default=None) + ap.add_argument("--attn-impl", default="eager", + help="eager/sdpa/flash_attention_2; must be set because " + "create_causal_mask() returns None when " + "config._attn_implementation is unset, which makes the " + "standalone decoder layers attend bidirectionally") + args = ap.parse_args() + + cfg = AutoConfig.from_pretrained(args.model) + # MUST be set before create_causal_mask(): it returns None when + # config._attn_implementation is unset, and a standalone decoder layer given + # attention_mask=None then attends BIDIRECTIONALLY, silently producing a + # non-causal imatrix for attn_output (and everything downstream of it). + cfg._attn_implementation = args.attn_impl + dev = args.device + dtype = torch.bfloat16 if dev == "cuda" else torch.float32 + reader = ShardReader(args.model) + known = set(read_imatrix(args.ground_truth)[1]) if args.ground_truth else None + nlayers = cfg.num_hidden_layers + d = cfg.hidden_size + block_type = cfg.layers_block_type + + # tokenize + window + tok = AutoTokenizer.from_pretrained(args.model) + chunks, calib_info = load_calibration_chunks(tok, args.calib, args.chunks, args.ctx, + add_bos=args.bos_per_chunk) + nch = chunks.shape[0] + print(f"calib: {calib_info}") + print(f"layers={nlayers} band={args.band} chunks={nch} ctx={args.ctx} dtype={dtype}") + + # --- embed all chunks -> hidden cache on CPU --- + embed_w = reader.get("model.embed_tokens.weight").to(dtype) + hid = (torch.nn.functional.embedding(chunks, embed_w) * cfg.embedding_multiplier).to(dtype) # [nch,ctx,d] CPU + del embed_w + print(f"hidden cache: {tuple(hid.shape)} ({hid.element_size()*hid.nelement()/1e6:.0f} MB CPU)") + + # --- rotary + causal mask (built once; position_ids identical per chunk) --- + rotary = GraniteMoeHybridRotaryEmbedding(cfg).to(dev) + pos_ids = torch.arange(args.ctx, device=dev).unsqueeze(0) + dummy = hid[:1].to(dev) + pos_emb = rotary(dummy, pos_ids) + causal = create_causal_mask(config=cfg, inputs_embeds=dummy, attention_mask=None, past_key_values=None) + + acc = {} + if dev == "cuda": + torch.cuda.reset_peak_memory_stats() + + # --- band loop --- + with torch.no_grad(): + for b0 in range(0, nlayers, args.band): + band = list(range(b0, min(b0 + args.band, nlayers))) + layers = [] + for i in band: + L = GraniteMoeHybridDecoderLayer(cfg, i) + materialize(L, f"model.layers.{i}", reader, dtype, dev) + layers.append((i, L)) + named = [(f"model.layers.{i}.{n}", m) for i, L in layers for n, m in L.named_modules()] + handles = make_hooks(named, acc, known) + for c in range(nch): + h = hid[c:c+1].to(dev) + for i, L in layers: + mask = None if block_type[i] == "mamba" else causal + h = L(h, attention_mask=mask, past_key_values=None, position_embeddings=pos_emb) + hid[c:c+1] = h.to(hid.dtype).cpu() + for hh in handles: + hh.remove() + # Drop EVERY reference to the band's modules before empty_cache(), or the + # cache release is a no-op: `named` holds the submodules and `del L` only + # unbinds the loop variable. + handles.clear(); named.clear(); layers.clear() + if dev == "cuda": + torch.cuda.empty_cache() + print(f" band {band[0]}-{band[-1]} done") + + # write + entries = {} + for name, dd in acc.items(): + sums = dd["sums"].astype(np.float32) + entries[name] = {"in_sum2": sums if sums.shape[0] > 1 else sums[0], + "counts": dd["counts"].astype(np.float32)} + write_imatrix(args.out, entries, ["calibration_datav3"], nch, args.ctx) + + peak_rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1e6 # GB (linux: KB) + line = f"wrote {len(entries)} entries -> {args.out} | peak RSS {peak_rss:.1f} GB" + if dev == "cuda": + line += f" | peak GPU {torch.cuda.max_memory_allocated()/1e9:.2f} GB" + print(line) + + +if __name__ == "__main__": + main() diff --git a/scripts/imatrix_serialized/serialized_gen_llama4.py b/scripts/imatrix_serialized/serialized_gen_llama4.py new file mode 100644 index 0000000..06c1bff --- /dev/null +++ b/scripts/imatrix_serialized/serialized_gen_llama4.py @@ -0,0 +1,263 @@ +"""Band-serialized (memory-bounded) Path-A imatrix generator for Llama4 (Scout). + +Same idea as serialized_gen.py (granite) but for the llama4_text arch, which is +different enough to need its own module: + - Experts are a single Llama4TextExperts module doing bmm on 3D Parameters + (gate_up_proj, down_proj) -- NOT per-expert Linears. So the down-proj input + is internal; we recompute the gated intermediate inside the hook from the + captured input + the band's loaded weights. + - MoE folds the router score into the expert input (hidden.repeat * router_scores), + so gate/up expert stats are router-weighted. This MATCHES llama.cpp, which special + cases this arch -- src/llama-graph.cpp:1837 + const bool weight_before_ffn = arch == LLM_ARCH_LLAMA4; + and at L1976 multiplies the sigmoid-ed weights into `cur` BEFORE the gate/up + mul_mat_id (the generic MoE path instead weights the expert OUTPUT at L2121, which + is what leaves gate/up unweighted for every other arch). So a weighted stat here is + the correct reproduction, not a divergence. + Row counts come from the router's exact top-k indices: transformers hands every + expert all T rows with non-selected ones scaled to zero, so they must be masked + back out to match llama.cpp's one-count-per-(token, selected-expert) slot. + There is still no ground-truth Scout imatrix to correlation-check against (bf16 + doesn't fit 128 GB -- the whole point), so this validates by QUANT COHERENCE, not + correlation. + - Attention interleaves RoPE / NoPE and uses chunked attention (chunk 8192); + at ctx<=8192 chunked == plain causal, and Llama4TextAttention applies RoPE + only on rope layers internally, so we pass one causal mask + rotary to all. + +Peak resident weights = band_size layers. Each weight read from disk once. +""" +import argparse, json, os, glob, resource +import numpy as np, torch +from transformers import AutoConfig, AutoTokenizer +from transformers.models.llama4.modeling_llama4 import ( + Llama4TextDecoderLayer, Llama4TextRotaryEmbedding, Llama4TextExperts, + Llama4Router, +) +from transformers.masking_utils import create_causal_mask +from safetensors import safe_open +from calib import load_calibration_chunks, add_calib_args +from imatrix_io import write_imatrix + + +def map_name(rel: str, blk: str): + """rel = module path relative to the layer, e.g. 'self_attn.q_proj'.""" + table = { + "self_attn.q_proj": [f"{blk}.attn_q.weight"], + "self_attn.k_proj": [f"{blk}.attn_k.weight"], + "self_attn.v_proj": [f"{blk}.attn_v.weight"], + "self_attn.o_proj": [f"{blk}.attn_output.weight"], + "feed_forward.router": [f"{blk}.ffn_gate_inp.weight"], + "feed_forward.shared_expert.gate_proj": [f"{blk}.ffn_gate_shexp.weight"], + "feed_forward.shared_expert.up_proj": [f"{blk}.ffn_up_shexp.weight"], + "feed_forward.shared_expert.down_proj": [f"{blk}.ffn_down_shexp.weight"], + } + return table.get(rel) + + +class ShardReader: + def __init__(self, d): + self.dir = d + idx = os.path.join(d, "model.safetensors.index.json") + with open(idx) as f: + self.wm = json.load(f)["weight_map"] + self.handles = {} + # detect weight-key prefix for the text decoder layers + self.prefix = "language_model.model" if any( + k.startswith("language_model.model.layers.") for k in self.wm) else "model" + + def get(self, name): + fn = self.wm[name] + if fn not in self.handles: + self.handles[fn] = safe_open(os.path.join(self.dir, fn), framework="pt") + return self.handles[fn].get_tensor(name) + + +def make_ensure(acc): + """Accumulator slot factory: acc[name] = {sums [nmat, nin], counts [nmat]}.""" + def ensure(name, nmat, nin): + if name not in acc: + acc[name] = {"sums": np.zeros((nmat, nin), np.float64), "counts": np.zeros(nmat, np.int64)} + return acc[name] + return ensure + + +def hook_linear(mod, names, ensure): + def pre(m, a, names=names): + x = a[0].detach().float().reshape(-1, a[0].shape[-1]) + s = (x * x).sum(0).double().cpu().numpy() + for nm in names: + e = ensure(nm, 1, x.shape[-1]); e["sums"][0] += s; e["counts"][0] += x.shape[0] + return mod.register_forward_pre_hook(pre) + + +def hook_router_select(mod, blk, sel): + """Capture the router's exact top-k selection for `blk` into sel[blk].""" + def post(m, a, out): + router_logits = out[1] + sel[blk] = torch.topk(router_logits, m.top_k, dim=1).indices.detach() + return mod.register_forward_hook(post) + + +def hook_experts(mod, blk, ensure, sel): + gate_n, up_n, down_n = f"{blk}.ffn_gate_exps.weight", f"{blk}.ffn_up_exps.weight", f"{blk}.ffn_down_exps.weight" + def pre(m, a): + E = m.num_experts + X = a[0].detach().view(E, -1, m.hidden_size) # [E,T,H] (router-scaled) + T = X.shape[1] + # Exact selection mask from the router's top-k, NOT a nonzero-row heuristic. + # transformers hands every expert all T rows with the non-selected ones + # scaled to 0, so counting `abs().sum(-1) > 0` rows happened to work but + # breaks on a genuinely-zero activation row or a router score that + # underflows to 0 in bf16. llama.cpp counts one per (token, selected-expert) + # slot, which is exactly what the top-k indices give. + idx = sel[blk] # [T, top_k] + keep = torch.zeros(T, E, dtype=torch.bool, device=X.device) + keep.scatter_(1, idx, True) + mask = keep.t().unsqueeze(-1) # [E,T,1] + cnt = keep.sum(0).cpu().numpy() # [E] exact row counts + + # NOTE: the router-weighted input IS correct for Llama-4. llama.cpp special + # cases this arch -- llama-graph.cpp:1837 + # const bool weight_before_ffn = arch == LLM_ARCH_LLAMA4; + # and at L1976 multiplies the sigmoid-ed weights into `cur` BEFORE the + # gate/up mul_mat_id, so llama-imatrix also records a weighted stat here. + xin = X.float() * mask + s_in = (xin * xin).sum(1).double().cpu().numpy() # [E,H] gate/up input stat + for nm in (gate_n, up_n): + e = ensure(nm, E, m.hidden_size); e["sums"] += s_in; e["counts"] += cnt + # recompute gated intermediate to get down_proj input (also from the + # weighted input, matching the graph order above) + gate_up = torch.bmm(X.to(m.gate_up_proj.dtype), m.gate_up_proj) + gate, up = gate_up.chunk(2, dim=-1) + inter = (up * m.act_fn(gate)).float() * mask # [E,T,inter] + s_dn = (inter * inter).sum(1).double().cpu().numpy() # [E,inter] + e = ensure(down_n, E, inter.shape[-1]); e["sums"] += s_dn; e["counts"] += cnt + return mod.register_forward_pre_hook(pre) + + +def materialize(module, prefix, reader, dtype, device): + # Every module parameter MUST come from the checkpoint -- see the note in + # serialized_gen.py: a skipped key leaves random weights in the band and the + # resulting imatrix is meaningless with no error anywhere. + sd, missing = {}, [] + for k in module.state_dict().keys(): + full = f"{prefix}.{k}" + if full in reader.wm: + sd[k] = reader.get(full).to(dtype) + else: + missing.append(k) + if missing: + raise KeyError(f"{prefix}: no checkpoint tensor for {missing}") + module.load_state_dict(sd, strict=False, assign=True) + return module.to(device).eval() + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--model", required=True) + ap.add_argument("--calib", required=True) + ap.add_argument("--out", required=True) + add_calib_args(ap) + ap.add_argument("--band", type=int, default=2) + ap.add_argument("--device", default="cuda") + ap.add_argument("--attn-impl", default="eager", + help="eager/sdpa/flash_attention_2; must be set because " + "create_causal_mask() returns None when " + "config._attn_implementation is unset, which makes the " + "standalone decoder layers attend bidirectionally") + args = ap.parse_args() + + full_cfg = AutoConfig.from_pretrained(args.model) + # MUST be set before create_causal_mask(): it returns None when + # config._attn_implementation is unset, and a standalone decoder layer given + # attention_mask=None then attends BIDIRECTIONALLY, silently producing a + # non-causal imatrix for attn_output (and everything downstream of it). + full_cfg._attn_implementation = args.attn_impl + cfg = full_cfg.get_text_config() + cfg._attn_implementation = args.attn_impl + dev = args.device + dtype = torch.bfloat16 if dev == "cuda" else torch.float32 + reader = ShardReader(args.model) + P = reader.prefix + nlayers = cfg.num_hidden_layers + moe_layers = set(getattr(cfg, "moe_layers", range(nlayers))) + print(f"llama4: layers={nlayers} experts={cfg.num_local_experts} band={args.band} " + f"prefix={P} dtype={dtype}") + + acc = {} + ensure = make_ensure(acc) + # Llama4TextMoe.forward runs the router BEFORE the experts, so a forward hook on + # the router lands the exact top-k selection here in time for the experts hook. + sel = {} + + # tokenize + window + tok = AutoTokenizer.from_pretrained(args.model) + chunks, calib_info = load_calibration_chunks(tok, args.calib, args.chunks, args.ctx, + add_bos=args.bos_per_chunk) + nch = chunks.shape[0] + print(f"calib: {calib_info}") + print(f"chunks={nch} ctx={args.ctx}") + + # embed -> hidden cache (CPU); Llama4 has no embedding_multiplier + embed_w = reader.get(f"{P}.embed_tokens.weight").to(dtype) + hid = torch.nn.functional.embedding(chunks, embed_w).to(dtype) + del embed_w + print(f"hidden cache {tuple(hid.shape)} ({hid.element_size()*hid.nelement()/1e6:.0f} MB)") + + rotary = Llama4TextRotaryEmbedding(cfg).to(dev) + pos_ids = torch.arange(args.ctx, device=dev).unsqueeze(0) + dummy = hid[:1].to(dev) + pos_emb = rotary(dummy, pos_ids) + causal = create_causal_mask(config=cfg, inputs_embeds=dummy, attention_mask=None, past_key_values=None) + + if dev == "cuda": + torch.cuda.reset_peak_memory_stats() + + with torch.no_grad(): + for b0 in range(0, nlayers, args.band): + band = list(range(b0, min(b0 + args.band, nlayers))) + layers, handles = [], [] + for i in band: + L = materialize(Llama4TextDecoderLayer(cfg, i), f"{P}.layers.{i}", reader, dtype, dev) + layers.append((i, L)) + blk = f"blk.{i}" + for rel, mod in L.named_modules(): + if isinstance(mod, Llama4TextExperts): + handles.append(hook_experts(mod, blk, ensure, sel)) + continue + if isinstance(mod, Llama4Router): + # router doubles as ffn_gate_inp (input stat) and as the source + # of the exact expert selection the experts hook needs + handles.append(hook_router_select(mod, blk, sel)) + names = map_name(rel, blk) + if names: + handles.append(hook_linear(mod, names, ensure)) + for c in range(nch): + h = hid[c:c+1].to(dev) + for i, L in layers: + h = L(h, attention_mask=causal, position_ids=pos_ids, + past_key_values=None, use_cache=False, position_embeddings=pos_emb) + hid[c:c+1] = h.to(hid.dtype).cpu() + for hh in handles: hh.remove() + # Drop EVERY reference to the band's modules before empty_cache(), or the + # cache release is a no-op: the hook closures and `layers` keep them alive + # and `del L` only unbinds the loop variable. + handles.clear(); layers.clear() + if dev == "cuda": torch.cuda.empty_cache() + print(f" band {band[0]}-{band[-1]} done") + + entries = {} + for name, d in acc.items(): + sums = d["sums"].astype(np.float32) + entries[name] = {"in_sum2": sums if sums.shape[0] > 1 else sums[0], + "counts": d["counts"].astype(np.float32)} + write_imatrix(args.out, entries, ["calibration_datav3"], nch, args.ctx) + peak_rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1e6 + line = f"wrote {len(entries)} entries -> {args.out} | peak RSS {peak_rss:.1f} GB" + if dev == "cuda": + line += f" | peak GPU {torch.cuda.max_memory_allocated()/1e9:.2f} GB" + print(line) + + +if __name__ == "__main__": + main() diff --git a/scripts/imatrix_serialized/serialized_gen_qwen35.py b/scripts/imatrix_serialized/serialized_gen_qwen35.py new file mode 100644 index 0000000..3e17bc2 --- /dev/null +++ b/scripts/imatrix_serialized/serialized_gen_qwen35.py @@ -0,0 +1,368 @@ +"""Band-serialized (VRAM-bounded) torch-imatrix generator for Qwen3.5-MoE (Qwen3_5Moe). + +Qwen3.5-35B-A3B is a hybrid: GatedDeltaNet linear-attention layers + periodic full +attention (every `full_attention_interval`), a 256-routed + shared-expert MoE on every +layer, wrapped in a multimodal shell (text weights under `model.language_model.`) with a +separate `mtp` head we ignore. + +Naming validated 510/510 against bartowski's public GGUF imatrix. The GatedDeltaNet is +all nn.Linear projections, so it needs no custom hook — only a name map: + in_proj_qkv->attn_qkv, in_proj_z->attn_gate, in_proj_a->ssm_alpha, + in_proj_b->ssm_beta, out_proj->ssm_out. +Experts use the same fused gate_up_proj/down_proj layout as granite-5.14, so the expert +hook (routing replay + down-input recompute) ports directly. The main model stores +experts fused, so materialize needs NO key remap — only the language_model prefix. +""" +import argparse, json, os, re, resource +import numpy as np, torch +from transformers import AutoConfig, AutoTokenizer +from transformers.models.qwen3_5_moe.modeling_qwen3_5_moe import ( + Qwen3_5MoeDecoderLayer, Qwen3_5MoeExperts, Qwen3_5MoeTopKRouter, + Qwen3_5MoeTextRotaryEmbedding, Qwen3_5MoeRMSNorm, +) +from transformers.masking_utils import create_causal_mask +from safetensors import safe_open +from calib import load_calibration_chunks, add_calib_args +from imatrix_io import write_imatrix, read_imatrix + + +def map_linear(hf): + m = re.match(r".*layers\.(\d+)\.(.+)$", hf) + if not m: + return None + b, tail = f"blk.{m.group(1)}", m.group(2) + T = { + "self_attn.q_proj": [f"{b}.attn_q.weight"], "self_attn.k_proj": [f"{b}.attn_k.weight"], + "self_attn.v_proj": [f"{b}.attn_v.weight"], "self_attn.o_proj": [f"{b}.attn_output.weight"], + "linear_attn.in_proj_qkv": [f"{b}.attn_qkv.weight"], "linear_attn.in_proj_z": [f"{b}.attn_gate.weight"], + "linear_attn.in_proj_a": [f"{b}.ssm_alpha.weight"], "linear_attn.in_proj_b": [f"{b}.ssm_beta.weight"], + "linear_attn.out_proj": [f"{b}.ssm_out.weight"], + "mlp.shared_expert.gate_proj": [f"{b}.ffn_gate_shexp.weight"], + "mlp.shared_expert.up_proj": [f"{b}.ffn_up_shexp.weight"], + "mlp.shared_expert.down_proj": [f"{b}.ffn_down_shexp.weight"], + "mlp.shared_expert_gate": [f"{b}.ffn_gate_inp_shexp.weight"], + } + return T.get(tail) + + +def expert_names(hf): + m = re.match(r".*layers\.(\d+)\.", hf) + if not m: + raise ValueError(f"no layer index in {hf!r}") + b = f"blk.{m.group(1)}" + return {"gate": f"{b}.ffn_gate_exps.weight", "up": f"{b}.ffn_up_exps.weight", "down": f"{b}.ffn_down_exps.weight"} + + +class ShardReader: + def __init__(self, d): + self.dir = d + with open(os.path.join(d, "model.safetensors.index.json")) as f: + self.wm = json.load(f)["weight_map"] + self.handles = {} + cand = ["model.language_model", "language_model.model", "model"] + self.prefix = next(p for p in cand if any(k.startswith(p + ".layers.") for k in self.wm)) + + def get(self, name): + fn = self.wm[name] + if fn not in self.handles: + self.handles[fn] = safe_open(os.path.join(self.dir, fn), framework="pt") + return self.handles[fn].get_tensor(name) + + +def materialize(module, prefix, reader, dtype, device): + sd, missing = {}, [] + for k in module.state_dict().keys(): + raw = f"{prefix}.{k}" + if raw in reader.wm: + sd[k] = reader.get(raw).to(dtype) + else: + missing.append(k) + if missing: + raise KeyError(f"{prefix}: missing {missing[:4]} (+{len(missing)-4})" if len(missing) > 4 else f"{prefix}: missing {missing}") + module.load_state_dict(sd, strict=False, assign=True) + return module.to(device).eval() + + +class Acc: + def __init__(self, device): + self.device = device; self.sums = {}; self.counts = {} + def ensure(self, name, nmat, nin): + if name not in self.sums: + self.sums[name] = torch.zeros((nmat, nin), dtype=torch.float64, device=self.device) + self.counts[name] = torch.zeros(nmat, dtype=torch.float64, device=self.device) + return self.sums[name], self.counts[name] + def to_entries(self): + e = {} + for name, s in self.sums.items(): + sums = s.cpu().numpy().astype(np.float32); counts = self.counts[name].cpu().numpy().astype(np.float32) + e[name] = {"in_sum2": sums if sums.shape[0] > 1 else sums[0], "counts": counts} + return e + + +def make_hooks(named_mods, acc, known): + def note(names): + if known is not None: + for g in names: + if g not in known: + print(f" WARN not in ground truth: {g}") + handles = [] + for hf, mod in named_mods: + if isinstance(mod, torch.nn.Linear): + names = map_linear(hf) + if names is None: + continue + note(names) + def pre_linear(m, a, names=names): + x = a[0].detach().reshape(-1, a[0].shape[-1]).float() + s = (x * x).sum(0).double(); n = x.shape[0] + for nm in names: + S, C = acc.ensure(nm, 1, x.shape[-1]); S[0] += s; C[0] += n + handles.append(mod.register_forward_pre_hook(pre_linear)) + elif isinstance(mod, Qwen3_5MoeExperts): + en = expert_names(hf); note(list(en.values())); ne = mod.num_experts + def pre_experts(m, a, en=en, ne=ne): + hidden_states, top_k_index = a[0], a[1] + x = hidden_states.detach().float() + T, K = top_k_index.shape + x2 = x * x + assign = top_k_index.reshape(-1) + tok = torch.arange(T, device=x.device).repeat_interleave(K) + Sg, Cg = acc.ensure(en["gate"], ne, x.shape[-1]); Su, Cu = acc.ensure(en["up"], ne, x.shape[-1]) + contrib = x2.index_select(0, tok).double() + Sg.index_add_(0, assign, contrib); Su.index_add_(0, assign, contrib) + ones = torch.ones(T * K, dtype=torch.float64, device=x.device) + Cg.index_add_(0, assign, ones); Cu.index_add_(0, assign, ones) + Sd, Cd = acc.ensure(en["down"], ne, m.down_proj.shape[-1]) + # .tolist() forces ONE device sync for the whole expert list; iterating + # the CUDA tensor directly syncs once PER expert (256 x layers x chunks). + for e in torch.unique(top_k_index).tolist(): + idx = (top_k_index == e).any(dim=-1).nonzero(as_tuple=True)[0] + cur = x.index_select(0, idx).to(m.gate_up_proj.dtype) + g, u = torch.nn.functional.linear(cur, m.gate_up_proj[e]).chunk(2, dim=-1) + inter = (m.act_fn(g) * u).float() + Sd[e] += (inter * inter).sum(0).double(); Cd[e] += idx.numel() + handles.append(mod.register_forward_pre_hook(pre_experts)) + elif isinstance(mod, Qwen3_5MoeTopKRouter): + _m = re.match(r".*layers\.(\d+)\.", hf) + gname = f"blk.{_m.group(1)}.ffn_gate_inp.weight" + note([gname]) + def pre_router(m, a, gname=gname): + x = a[0].detach().reshape(-1, a[0].shape[-1]).float() + S, C = acc.ensure(gname, 1, x.shape[-1]); S[0] += (x * x).sum(0).double(); C[0] += x.shape[0] + handles.append(mod.register_forward_pre_hook(pre_router)) + return handles + + +def materialize_mtp_layer(layer, reader, dtype, device): + """Materialize the MTP decoder layer from mtp.layers.0.*. Handles both + checkpoint layouts: 3.5 stores the MTP experts PER-EXPERT, while 3.6 stores + them FUSED (gate_up_proj/down_proj) exactly like the main model.""" + pre = "mtp.layers.0" + sd, missing = {}, [] + ne = layer.mlp.experts.num_experts + for k in layer.state_dict().keys(): + if k == "mlp.experts.gate_up_proj": + if f"{pre}.mlp.experts.gate_up_proj" in reader.wm: # 3.6: fused + sd[k] = reader.get(f"{pre}.mlp.experts.gate_up_proj").to(dtype) + else: # 3.5: per-expert + g = [reader.get(f"{pre}.mlp.experts.{e}.gate_proj.weight") for e in range(ne)] + u = [reader.get(f"{pre}.mlp.experts.{e}.up_proj.weight") for e in range(ne)] + sd[k] = torch.stack([torch.cat([g[e], u[e]], 0) for e in range(ne)]).to(dtype) + elif k == "mlp.experts.down_proj": + if f"{pre}.mlp.experts.down_proj" in reader.wm: # 3.6: fused + sd[k] = reader.get(f"{pre}.mlp.experts.down_proj").to(dtype) + else: # 3.5: per-expert + d = [reader.get(f"{pre}.mlp.experts.{e}.down_proj.weight") for e in range(ne)] + sd[k] = torch.stack(d).to(dtype) + elif f"{pre}.{k}" in reader.wm: + sd[k] = reader.get(f"{pre}.{k}").to(dtype) + else: + missing.append(k) + # Same rule as materialize(): a missing MTP key would leave random weights in + # the head and silently poison blk.{n_layers}.* in the output imatrix. + if missing: + raise KeyError(f"{pre}: no checkpoint tensor for {missing}") + layer.load_state_dict(sd, strict=False, assign=True) + return layer.to(device).eval() + + +def process_mtp(cfg, reader, hid, chunks, pos_emb, acc, known, dtype, dev, batch): + """Cover the nextn/MTP head. transformers-5.14 has no MTP module, so we build it + from raw checkpoint tensors and run the standard Qwen/DeepSeek MTP forward: + + h'_i = fc( concat[ enorm(embed(t_{i+1})) , hnorm(h_i) ] ) # 'eh_proj' order + out = mtp_decoder_layer(h'_i) # full-attn MoE layer + + where h_i is the main model's final hidden state (== `hid` after the band loop). + The MTP block maps to ggml blk.{n_layers}: its decoder-layer tensors reuse the + normal map (attn_*, ffn_*_exps/shexp, ffn_gate_inp*), plus fc -> nextn.eh_proj. + + The eh_proj concat order is VERIFIED against llama.cpp master, not assumed: + src/models/qwen35moe.cpp builds the MTP block with + ggml_concat(ctx0, e_norm, h_norm, /*dim=*/0) + and ggml_compute_forward_concat writes src0 at the low indices of `dim`, so input + channels [0, n_embd) are the embedding-norm half and [n_embd, 2*n_embd) the + hidden-norm half. eh_proj is created as {2*n_embd, n_embd} (ne[0] = reduction dim), + and conversion/qwen.py renames `mtp.fc` -> `layers.{n_layer}.eh_proj` with no + transpose, so torch's [H, 2H] row-major weight keeps the same channel order in the + gguf. The torch.cat below therefore matches the graph channel-for-channel. + + The one-position token/hidden shift is likewise verified, caller-side rather than + in the graph: common/speculative.cpp pairs, in a single batch row, the token + sampled AT the position whose hidden state it also submits + (`common_batch_add(batch, id, ...)` + `memcpy(batch.embd + ..., h_row, ...)`), + i.e. embed(t_{i+1}) with h_i -- the same alignment as the `emb[:, 1:]` / + `h[:, :-1]` pairing below. + + NOTE: what remains unverifiable is the numerical result, not the wiring -- these + values cannot be correlation-checked against a public imatrix (none covers MTP) + and transformers has no MTP module, so the forward is reconstructed from the + checkpoint and validated only by names/shapes against the gguf. + """ + P, H, L = reader.prefix, cfg.hidden_size, cfg.num_hidden_layers + blk = f"blk.{L}" + print(f"[mtp] building nextn head -> {blk}.*") + # full-attention decoder layer for mtp.layers.0 + full_idx = next(i for i, t in enumerate(cfg.layer_types) if t == "full_attention") + layer = materialize_mtp_layer(Qwen3_5MoeDecoderLayer(cfg, full_idx), reader, dtype, dev) + fc = torch.nn.Linear(2 * H, H, bias=False) + fc.load_state_dict({"weight": reader.get("mtp.fc.weight").to(dtype)}, assign=True); fc = fc.to(dev).eval() + def rms(key): + n = Qwen3_5MoeRMSNorm(H, eps=cfg.rms_norm_eps) + n.load_state_dict({"weight": reader.get(key).to(dtype)}, assign=True); return n.to(dev).eval() + enorm, hnorm = rms("mtp.pre_fc_norm_embedding.weight"), rms("mtp.pre_fc_norm_hidden.weight") + embed_w = reader.get(f"{P}.embed_tokens.weight").to(dtype).to(dev) + + named = [(f"x.layers.{L}.{n}", m) for n, m in layer.named_modules()] + handles = make_hooks(named, acc, known) + ehname = f"{blk}.nextn.eh_proj.weight" + if known is not None and ehname not in known: + print(f" WARN not in ground truth: {ehname}") + def pre_fc(m, a): + x = a[0].detach().reshape(-1, a[0].shape[-1]).float() + S, C = acc.ensure(ehname, 1, x.shape[-1]); S[0] += (x * x).sum(0).double(); C[0] += x.shape[0] + handles.append(fc.register_forward_pre_hook(pre_fc)) + + nch = chunks.shape[0] + with torch.no_grad(): + for s in range(0, nch, batch): + sl = slice(s, min(s + batch, nch)) + h = hid[sl].to(dev) # main final hidden [b,T,H] + ids = chunks[sl].to(dev) + emb = torch.nn.functional.embedding(ids, embed_w) # [b,T,H] + # position i predicts t_{i+2} from h_i and emb(t_{i+1}); align by one + cat = torch.cat([enorm(emb[:, 1:]), hnorm(h[:, :-1])], dim=-1) # [b,T-1,2H] (eh order) + hp = fc(cat) + bs, tm1 = hp.shape[0], hp.shape[1] + mask = create_causal_mask(config=cfg, inputs_embeds=hp, attention_mask=None, past_key_values=None) + out = layer(hp, position_embeddings=(pos_emb[0][:, :tm1], pos_emb[1][:, :tm1]), + attention_mask=mask, position_ids=torch.arange(tm1, device=dev).unsqueeze(0)) + _ = out[0] if isinstance(out, tuple) else out + for hh in handles: + hh.remove() + del layer, fc + if dev == "cuda": + torch.cuda.empty_cache() + print(f"[mtp] done") + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--model", required=True) + ap.add_argument("--calib", required=True) + ap.add_argument("--out", required=True) + add_calib_args(ap) + ap.add_argument("--band", type=int, default=2) + ap.add_argument("--batch", type=int, default=1) + ap.add_argument("--device", default="cuda") + ap.add_argument("--dataset-label", default="calibration_datav3") + ap.add_argument("--ground-truth", default=None) + ap.add_argument("--mtp", action="store_true", help="also cover the nextn/MTP head -> blk.{n_layers}.*") + ap.add_argument("--attn-impl", default="eager", + help="eager/sdpa/flash_attention_2; must be set because " + "create_causal_mask() returns None when " + "config._attn_implementation is unset, which makes the " + "standalone decoder layers attend bidirectionally") + args = ap.parse_args() + + cfg = AutoConfig.from_pretrained(args.model).get_text_config() + # MUST be set before create_causal_mask(): it returns None when + # config._attn_implementation is unset, and a standalone decoder layer given + # attention_mask=None then attends BIDIRECTIONALLY, silently producing a + # non-causal imatrix for attn_output (and everything downstream of it). + cfg._attn_implementation = args.attn_impl + dev, dtype = args.device, (torch.bfloat16 if args.device == "cuda" else torch.float32) + reader = ShardReader(args.model); P = reader.prefix + known = set(read_imatrix(args.ground_truth)[1]) if args.ground_truth else None + nlayers = cfg.num_hidden_layers + ltype = cfg.layer_types + print(f"qwen3.5: layers={nlayers} experts={getattr(cfg,'num_experts',None)} band={args.band} " + f"batch={args.batch} prefix={P} dtype={dtype}") + + tok = AutoTokenizer.from_pretrained(args.model) + chunks, calib_info = load_calibration_chunks(tok, args.calib, args.chunks, args.ctx, + add_bos=args.bos_per_chunk) + nch = chunks.shape[0] + print(f"calib: {calib_info}") + print(f"chunks={nch} ctx={args.ctx}") + + embed_w = reader.get(f"{P}.embed_tokens.weight").to(dtype) + hid = torch.nn.functional.embedding(chunks, embed_w).to(dtype).to(dev) # Qwen: no embedding multiplier + del embed_w + print(f"hidden cache on {dev}: {tuple(hid.shape)} ({hid.element_size()*hid.nelement()/1e6:.0f} MB)") + + rotary = Qwen3_5MoeTextRotaryEmbedding(cfg).to(dev) + pos_ids = torch.arange(args.ctx, device=dev).unsqueeze(0) + pos_emb = rotary(hid[:1], pos_ids) + mask_cache = {} + def causal_for(bs): + if bs not in mask_cache: + mask_cache[bs] = create_causal_mask(config=cfg, inputs_embeds=hid[:bs], attention_mask=None, past_key_values=None) + return mask_cache[bs] + + acc = Acc(dev) + if dev == "cuda": + torch.cuda.reset_peak_memory_stats() + + with torch.no_grad(): + for b0 in range(0, nlayers, args.band): + band = list(range(b0, min(b0 + args.band, nlayers))) + layers = [] + for i in band: + L = materialize(Qwen3_5MoeDecoderLayer(cfg, i), f"{P}.layers.{i}", reader, dtype, dev) + layers.append((i, L)) + named = [(f"{P}.layers.{i}.{n}", m) for i, L in layers for n, m in L.named_modules()] + handles = make_hooks(named, acc, known) + for s in range(0, nch, args.batch): + sl = slice(s, min(s + args.batch, nch)) + h = hid[sl] + mask = causal_for(h.shape[0]) + for i, L in layers: + m = mask if ltype[i] == "full_attention" else None + out = L(h, position_embeddings=pos_emb, attention_mask=m, + position_ids=pos_ids, past_key_values=None) + h = out[0] if isinstance(out, tuple) else out + hid[sl] = h + for hh in handles: + hh.remove() + # Drop EVERY reference to the band's modules before empty_cache(), or the + # cache release is a no-op: `named` holds the submodules and `del L` only + # unbinds the loop variable. + handles.clear(); named.clear(); layers.clear() + if dev == "cuda": + torch.cuda.empty_cache() + print(f" band {band[0]}-{band[-1]} done", flush=True) + + if args.mtp: + process_mtp(cfg, reader, hid, chunks, pos_emb, acc, known, dtype, dev, args.batch) + + write_imatrix(args.out, acc.to_entries(), [args.dataset_label], nch, args.ctx) + peak_rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1e6 + line = f"wrote {len(acc.sums)} entries -> {args.out} | peak RSS {peak_rss:.1f} GB" + if dev == "cuda": + line += f" | peak GPU {torch.cuda.max_memory_allocated()/1e9:.2f} GB" + print(line) + + +if __name__ == "__main__": + main() diff --git a/scripts/imatrix_serialized/serialized_gen_t514.py b/scripts/imatrix_serialized/serialized_gen_t514.py new file mode 100644 index 0000000..991597c --- /dev/null +++ b/scripts/imatrix_serialized/serialized_gen_t514.py @@ -0,0 +1,227 @@ +"""Band-serialized (VRAM-bounded) torch-imatrix generator, transformers-5.14. + +Same idea as serialized_gen.py: process the model in BANDS of `--band` decoder +layers so peak *GPU* residency is band-size layers, not the whole model. Adapted +to the transformers-5.14 GraniteMoeHybrid rewrite: + + * experts are now a single `GraniteMoeHybridExperts` (params gate_up_proj/ + down_proj, functional per-expert linears) + a bare-Parameter router; hooks + replay the routing and recompute the down input (see gen_imatrix_t514.py). + * the on-disk checkpoint still uses the OLD fused key names + (block_sparse_moe.input_linear/output_linear/router.layer); from_pretrained + remaps them at load, but our manual per-layer materialize must do it itself. + The remap is a pure same-shape rename (verified): see _RENAME. + * config uses layer_types with values 'linear_attention'/'full_attention' + (not 'mamba'/'attention'); only full_attention layers take the causal mask. + +VRAM-bounded: peak GPU = band layers + one naive-Mamba scan intermediate. The +full model is held in CPU RAM via the shard reader's lazy tensors (each weight +read from disk once); the hidden-state cache lives on CPU between bands. +""" +import argparse, json, os, re, resource +import numpy as np, torch +from transformers import AutoConfig, AutoTokenizer +from transformers.models.granitemoehybrid.modeling_granitemoehybrid import ( + GraniteMoeHybridDecoderLayer, GraniteMoeHybridRotaryEmbedding, + GraniteMoeHybridExperts, GraniteMoeHybridTopKRouter, +) +from transformers.masking_utils import create_causal_mask +from safetensors import safe_open +from calib import load_calibration_chunks, add_calib_args +from imatrix_io import write_imatrix, read_imatrix +from gen_imatrix_t514 import map_linear, expert_names + +# module state_dict key -> raw checkpoint key (same shape, pure rename) +_RENAME = { + "block_sparse_moe.experts.gate_up_proj": "block_sparse_moe.input_linear.weight", + "block_sparse_moe.experts.down_proj": "block_sparse_moe.output_linear.weight", + "block_sparse_moe.router.weight": "block_sparse_moe.router.layer.weight", +} + + +class ShardReader: + def __init__(self, d): + self.dir = d + idx = os.path.join(d, "model.safetensors.index.json") + with open(idx) as f: + self.wm = json.load(f)["weight_map"] + self.handles = {} + + def get(self, name): + fn = self.wm[name] + if fn not in self.handles: + self.handles[fn] = safe_open(os.path.join(self.dir, fn), framework="pt") + return self.handles[fn].get_tensor(name) + + +def materialize(module, prefix, reader, dtype, device): + sd, missing = {}, [] + for k in module.state_dict().keys(): + raw = f"{prefix}.{_RENAME.get(k, k)}" + if raw in reader.wm: + sd[k] = reader.get(raw).to(dtype) + else: + missing.append(k) + if missing: + raise KeyError(f"{prefix}: no checkpoint tensor for {missing}") + module.load_state_dict(sd, strict=False, assign=True) + return module.to(device).eval() + + +def make_hooks(named_mods, acc, known): + def ensure(name, nmat, nin): + if name not in acc: + acc[name] = {"sums": np.zeros((nmat, nin), np.float64), "counts": np.zeros(nmat, np.int64)} + return acc[name] + + def note(names): + if known is not None: + for g in names: + if g not in known: + print(f" WARN mapped name not in ground truth: {g}") + + handles = [] + for hf_name, mod in named_mods: + if isinstance(mod, torch.nn.Linear): + names = map_linear(hf_name) + if names is None: + continue # lm_head etc. never appears inside a decoder layer + note(names) + def pre_linear(m, a, names=names): + x = a[0].detach().float().reshape(-1, a[0].shape[-1]) + s = (x * x).sum(0).double().cpu().numpy() + for nm in names: + e = ensure(nm, 1, x.shape[-1]); e["sums"][0] += s; e["counts"][0] += x.shape[0] + handles.append(mod.register_forward_pre_hook(pre_linear)) + elif isinstance(mod, GraniteMoeHybridExperts): + en = expert_names(hf_name); note(list(en.values())) + ne = mod.num_experts + def pre_experts(m, a, en=en, ne=ne): + hidden_states, top_k_index = a[0], a[1] + x = hidden_states.detach().float() + gate = ensure(en["gate"], ne, x.shape[-1]); up = ensure(en["up"], ne, x.shape[-1]) + with torch.no_grad(): + mask = torch.nn.functional.one_hot(top_k_index, num_classes=ne).permute(2, 1, 0) + for eidx in torch.greater(mask.sum(dim=(-1, -2)), 0).nonzero(): + e = eidx[0] + if e == ne: + continue + _, tok_idx = torch.where(mask[e]) + cur = x[tok_idx] + s = (cur * cur).sum(0).double().cpu().numpy() + gate["sums"][e] += s; gate["counts"][e] += cur.shape[0] + up["sums"][e] += s; up["counts"][e] += cur.shape[0] + gu = torch.nn.functional.linear(cur.to(m.gate_up_proj.dtype), m.gate_up_proj[e]) + g, u = gu.chunk(2, dim=-1) + inter = (m.act_fn(g) * u).float() + dn = ensure(en["down"], ne, inter.shape[-1]) + dn["sums"][e] += (inter * inter).sum(0).double().cpu().numpy() + dn["counts"][e] += inter.shape[0] + handles.append(mod.register_forward_pre_hook(pre_experts)) + elif isinstance(mod, GraniteMoeHybridTopKRouter): + _m = re.match(r".*layers\.(\d+)\.", hf_name) + gname = f"blk.{_m.group(1)}.ffn_gate_inp.weight" + note([gname]) + def pre_router(m, a, gname=gname): + x = a[0].detach().float().reshape(-1, a[0].shape[-1]) + s = (x * x).sum(0).double().cpu().numpy() + e = ensure(gname, 1, x.shape[-1]); e["sums"][0] += s; e["counts"][0] += x.shape[0] + handles.append(mod.register_forward_pre_hook(pre_router)) + return handles + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--model", required=True) + ap.add_argument("--calib", required=True) + ap.add_argument("--out", required=True) + add_calib_args(ap) + ap.add_argument("--band", type=int, default=4, help="layers resident at once (VRAM knob)") + ap.add_argument("--device", default="cuda") + ap.add_argument("--dataset-label", default="calibration_datav3") + ap.add_argument("--ground-truth", default=None) + ap.add_argument("--attn-impl", default="eager", + help="eager/sdpa/flash_attention_2; must be set because " + "create_causal_mask() returns None when " + "config._attn_implementation is unset, which makes the " + "standalone decoder layers attend bidirectionally") + args = ap.parse_args() + + cfg = AutoConfig.from_pretrained(args.model) + # MUST be set before create_causal_mask(): it returns None when + # config._attn_implementation is unset, and a standalone decoder layer given + # attention_mask=None then attends BIDIRECTIONALLY, silently producing a + # non-causal imatrix for attn_output (and everything downstream of it). + cfg._attn_implementation = args.attn_impl + dev = args.device + dtype = torch.bfloat16 if dev == "cuda" else torch.float32 + reader = ShardReader(args.model) + known = set(read_imatrix(args.ground_truth)[1]) if args.ground_truth else None + nlayers = cfg.num_hidden_layers + ltype = cfg.layer_types + + tok = AutoTokenizer.from_pretrained(args.model) + chunks, calib_info = load_calibration_chunks(tok, args.calib, args.chunks, args.ctx, + add_bos=args.bos_per_chunk) + nch = chunks.shape[0] + print(f"calib: {calib_info}") + print(f"layers={nlayers} band={args.band} chunks={nch} ctx={args.ctx} dtype={dtype}") + + embed_w = reader.get("model.embed_tokens.weight").to(dtype) + hid = (torch.nn.functional.embedding(chunks, embed_w) * cfg.embedding_multiplier).to(dtype) # [nch,ctx,d] CPU + del embed_w + print(f"hidden cache: {tuple(hid.shape)} ({hid.element_size()*hid.nelement()/1e6:.0f} MB CPU)") + + rotary = GraniteMoeHybridRotaryEmbedding(cfg).to(dev) + pos_ids = torch.arange(args.ctx, device=dev).unsqueeze(0) + dummy = hid[:1].to(dev) + pos_emb = rotary(dummy, pos_ids) + causal = create_causal_mask(config=cfg, inputs_embeds=dummy, attention_mask=None, past_key_values=None) + + acc = {} + if dev == "cuda": + torch.cuda.reset_peak_memory_stats() + + with torch.no_grad(): + for b0 in range(0, nlayers, args.band): + band = list(range(b0, min(b0 + args.band, nlayers))) + layers = [] + for i in band: + L = GraniteMoeHybridDecoderLayer(cfg, i) + materialize(L, f"model.layers.{i}", reader, dtype, dev) + layers.append((i, L)) + named = [(f"model.layers.{i}.{n}", m) for i, L in layers for n, m in L.named_modules()] + handles = make_hooks(named, acc, known) + for c in range(nch): + h = hid[c:c+1].to(dev) + for i, L in layers: + mask = causal if ltype[i] == "full_attention" else None + out = L(h, attention_mask=mask, past_key_values=None, position_embeddings=pos_emb) + h = out[0] if isinstance(out, tuple) else out + hid[c:c+1] = h.to(hid.dtype).cpu() + for hh in handles: + hh.remove() + # Drop EVERY reference to the band's modules before empty_cache(), or the + # cache release is a no-op: `named` holds the submodules and `del L` only + # unbinds the loop variable. + handles.clear(); named.clear(); layers.clear() + if dev == "cuda": + torch.cuda.empty_cache() + print(f" band {band[0]}-{band[-1]} done", flush=True) + + entries = {} + for name, dd in acc.items(): + sums = dd["sums"].astype(np.float32) + entries[name] = {"in_sum2": sums if sums.shape[0] > 1 else sums[0], + "counts": dd["counts"].astype(np.float32)} + write_imatrix(args.out, entries, [args.dataset_label], nch, args.ctx) + + peak_rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1e6 + line = f"wrote {len(entries)} entries -> {args.out} | peak RSS {peak_rss:.1f} GB" + if dev == "cuda": + line += f" | peak GPU {torch.cuda.max_memory_allocated()/1e9:.2f} GB" + print(line) + + +if __name__ == "__main__": + main() diff --git a/scripts/imatrix_serialized/serialized_gen_t514_fast.py b/scripts/imatrix_serialized/serialized_gen_t514_fast.py new file mode 100644 index 0000000..b9cc860 --- /dev/null +++ b/scripts/imatrix_serialized/serialized_gen_t514_fast.py @@ -0,0 +1,251 @@ +"""Band-serialized torch-imatrix generator (transformers-5.14) — throughput-optimized. + +Numerically equivalent to serialized_gen_t514.py (Sum x^2 is order/batch-invariant), +but restructured so the GPU stays busy instead of stalling on host syncs: + + 1. GPU-resident accumulators. sums/counts live on-GPU (float64) and are copied + to host exactly once, after the last band — vs the reference tool's + ~1M per-expert `.cpu().numpy()` calls (64 experts x 40 layers x 126 chunks), + each a GPU->CPU sync that idled the card. This is the dominant win. + 2. GPU-resident hidden cache. The [nch,ctx,d] activation cache (~198 MB bf16 + for 126x512x1536) stays on-GPU between bands, so there is no per-chunk + host<->device copy. + 3. Vectorized gate/up expert stats. Each token adds x^2 to each of its top-k + experts via a single index_add_ over all (token,expert) assignments — no + per-expert Python loop. (down_exps still loops experts, because each expert + applies different gate_up weights to compute its intermediate, but it + accumulates on-GPU with no sync.) + 4. Chunk batching (--batch): several ctx-windows per forward for bigger GEMMs. + Scan memory scales with batch, so keep it modest in banded mode. + +Validate against the reference tool with compare.py (expect corr ~1.0, tiny L1). +""" +import argparse, json, os, re, resource +import numpy as np, torch +from transformers import AutoConfig, AutoTokenizer +from transformers.models.granitemoehybrid.modeling_granitemoehybrid import ( + GraniteMoeHybridDecoderLayer, GraniteMoeHybridRotaryEmbedding, + GraniteMoeHybridExperts, GraniteMoeHybridTopKRouter, +) +from transformers.masking_utils import create_causal_mask +from safetensors import safe_open +from calib import load_calibration_chunks, add_calib_args +from imatrix_io import write_imatrix, read_imatrix +from gen_imatrix_t514 import map_linear, expert_names + +_RENAME = { + "block_sparse_moe.experts.gate_up_proj": "block_sparse_moe.input_linear.weight", + "block_sparse_moe.experts.down_proj": "block_sparse_moe.output_linear.weight", + "block_sparse_moe.router.weight": "block_sparse_moe.router.layer.weight", +} + + +class ShardReader: + def __init__(self, d): + self.dir = d + with open(os.path.join(d, "model.safetensors.index.json")) as f: + self.wm = json.load(f)["weight_map"] + self.handles = {} + + def get(self, name): + fn = self.wm[name] + if fn not in self.handles: + self.handles[fn] = safe_open(os.path.join(self.dir, fn), framework="pt") + return self.handles[fn].get_tensor(name) + + +def materialize(module, prefix, reader, dtype, device): + sd, missing = {}, [] + for k in module.state_dict().keys(): + raw = f"{prefix}.{_RENAME.get(k, k)}" + if raw in reader.wm: + sd[k] = reader.get(raw).to(dtype) + else: + missing.append(k) + if missing: + raise KeyError(f"{prefix}: no checkpoint tensor for {missing}") + module.load_state_dict(sd, strict=False, assign=True) + return module.to(device).eval() + + +class Acc: + """On-GPU float64 accumulators; host copy only at the end.""" + def __init__(self, device): + self.device = device + self.sums = {} # name -> [nmat, in] f64 on GPU + self.counts = {} # name -> [nmat] f64 on GPU + + def ensure(self, name, nmat, nin): + if name not in self.sums: + self.sums[name] = torch.zeros((nmat, nin), dtype=torch.float64, device=self.device) + self.counts[name] = torch.zeros(nmat, dtype=torch.float64, device=self.device) + return self.sums[name], self.counts[name] + + def to_entries(self): + entries = {} + for name, s in self.sums.items(): + sums = s.cpu().numpy().astype(np.float32) + counts = self.counts[name].cpu().numpy().astype(np.float32) + entries[name] = {"in_sum2": sums if sums.shape[0] > 1 else sums[0], "counts": counts} + return entries + + +def make_hooks(named_mods, acc, known): + def note(names): + if known is not None: + for g in names: + if g not in known: + print(f" WARN mapped name not in ground truth: {g}") + + handles = [] + for hf_name, mod in named_mods: + if isinstance(mod, torch.nn.Linear): + names = map_linear(hf_name) + if names is None: + continue + note(names) + def pre_linear(m, a, names=names): + x = a[0].detach().reshape(-1, a[0].shape[-1]).float() + s = (x * x).sum(0).double() + n = x.shape[0] + for nm in names: + S, C = acc.ensure(nm, 1, x.shape[-1]) + S[0] += s; C[0] += n + handles.append(mod.register_forward_pre_hook(pre_linear)) + elif isinstance(mod, GraniteMoeHybridExperts): + en = expert_names(hf_name); note(list(en.values())) + ne = mod.num_experts + def pre_experts(m, a, en=en, ne=ne): + hidden_states, top_k_index = a[0], a[1] + x = hidden_states.detach().float() # [T, hidden] + T, K = top_k_index.shape + x2 = x * x # [T, hidden] + # gate/up: each token adds x^2 to each of its K experts (one index_add) + assign = top_k_index.reshape(-1) # [T*K] + tok = torch.arange(T, device=x.device).repeat_interleave(K) + Sg, Cg = acc.ensure(en["gate"], ne, x.shape[-1]) + Su, Cu = acc.ensure(en["up"], ne, x.shape[-1]) + contrib = x2.index_select(0, tok).double() # [T*K, hidden] + Sg.index_add_(0, assign, contrib); Su.index_add_(0, assign, contrib) + ones = torch.ones(T * K, dtype=torch.float64, device=x.device) + Cg.index_add_(0, assign, ones); Cu.index_add_(0, assign, ones) + # down: per-expert intermediate (distinct weights). .tolist() forces ONE + # device sync for the whole expert list; iterating the CUDA tensor + # directly syncs once PER expert (n_experts x layers x chunks). + Sd, Cd = acc.ensure(en["down"], ne, m.down_proj.shape[-1]) + for e in torch.unique(top_k_index).tolist(): + tok_idx = (top_k_index == e).any(dim=-1).nonzero(as_tuple=True)[0] + cur = x.index_select(0, tok_idx).to(m.gate_up_proj.dtype) + g, u = torch.nn.functional.linear(cur, m.gate_up_proj[e]).chunk(2, dim=-1) + inter = (m.act_fn(g) * u).float() + Sd[e] += (inter * inter).sum(0).double(); Cd[e] += tok_idx.numel() + handles.append(mod.register_forward_pre_hook(pre_experts)) + elif isinstance(mod, GraniteMoeHybridTopKRouter): + _m = re.match(r".*layers\.(\d+)\.", hf_name) + gname = f"blk.{_m.group(1)}.ffn_gate_inp.weight" + note([gname]) + def pre_router(m, a, gname=gname): + x = a[0].detach().reshape(-1, a[0].shape[-1]).float() + S, C = acc.ensure(gname, 1, x.shape[-1]) + S[0] += (x * x).sum(0).double(); C[0] += x.shape[0] + handles.append(mod.register_forward_pre_hook(pre_router)) + return handles + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--model", required=True) + ap.add_argument("--calib", required=True) + ap.add_argument("--out", required=True) + add_calib_args(ap) + ap.add_argument("--band", type=int, default=4) + ap.add_argument("--batch", type=int, default=2, help="ctx-windows per forward (scan mem scales with this)") + ap.add_argument("--device", default="cuda") + ap.add_argument("--dataset-label", default="calibration_datav3") + ap.add_argument("--ground-truth", default=None) + ap.add_argument("--attn-impl", default="eager", + help="eager/sdpa/flash_attention_2; must be set because " + "create_causal_mask() returns None when " + "config._attn_implementation is unset, which makes the " + "standalone decoder layers attend bidirectionally") + args = ap.parse_args() + + cfg = AutoConfig.from_pretrained(args.model) + # MUST be set before create_causal_mask(): it returns None when + # config._attn_implementation is unset, and a standalone decoder layer given + # attention_mask=None then attends BIDIRECTIONALLY, silently producing a + # non-causal imatrix for attn_output (and everything downstream of it). + cfg._attn_implementation = args.attn_impl + dev = args.device + dtype = torch.bfloat16 if dev == "cuda" else torch.float32 + reader = ShardReader(args.model) + known = set(read_imatrix(args.ground_truth)[1]) if args.ground_truth else None + nlayers = cfg.num_hidden_layers + ltype = cfg.layer_types + + tok = AutoTokenizer.from_pretrained(args.model) + chunks, calib_info = load_calibration_chunks(tok, args.calib, args.chunks, args.ctx, + add_bos=args.bos_per_chunk) + nch = chunks.shape[0] + print(f"calib: {calib_info}") + print(f"layers={nlayers} band={args.band} batch={args.batch} chunks={nch} ctx={args.ctx} dtype={dtype}") + + embed_w = reader.get("model.embed_tokens.weight").to(dtype) + # hidden cache lives on-GPU + hid = (torch.nn.functional.embedding(chunks, embed_w) * cfg.embedding_multiplier).to(dtype).to(dev) + del embed_w + print(f"hidden cache on {dev}: {tuple(hid.shape)} ({hid.element_size()*hid.nelement()/1e6:.0f} MB)") + + rotary = GraniteMoeHybridRotaryEmbedding(cfg).to(dev) + pos_ids = torch.arange(args.ctx, device=dev).unsqueeze(0) + pos_emb = rotary(hid[:1], pos_ids) # cos/sin [1,ctx,dim] broadcast over batch + mask_cache = {} + def causal_for(bs): + if bs not in mask_cache: + mask_cache[bs] = create_causal_mask(config=cfg, inputs_embeds=hid[:bs], + attention_mask=None, past_key_values=None) + return mask_cache[bs] + + acc = Acc(dev) + if dev == "cuda": + torch.cuda.reset_peak_memory_stats() + + with torch.no_grad(): + for b0 in range(0, nlayers, args.band): + band = list(range(b0, min(b0 + args.band, nlayers))) + layers = [] + for i in band: + L = GraniteMoeHybridDecoderLayer(cfg, i) + materialize(L, f"model.layers.{i}", reader, dtype, dev) + layers.append((i, L)) + named = [(f"model.layers.{i}.{n}", m) for i, L in layers for n, m in L.named_modules()] + handles = make_hooks(named, acc, known) + for s in range(0, nch, args.batch): + sl = slice(s, min(s + args.batch, nch)) + h = hid[sl] + mask = causal_for(h.shape[0]) + for i, L in layers: + m = mask if ltype[i] == "full_attention" else None + out = L(h, attention_mask=m, past_key_values=None, position_embeddings=pos_emb) + h = out[0] if isinstance(out, tuple) else out + hid[sl] = h + for hh in handles: + hh.remove() + # Drop EVERY reference to the band's modules before empty_cache(), or the + # cache release is a no-op: `named` holds the submodules and `del L` only + # unbinds the loop variable. + handles.clear(); named.clear(); layers.clear() + if dev == "cuda": + torch.cuda.empty_cache() + print(f" band {band[0]}-{band[-1]} done", flush=True) + + write_imatrix(args.out, acc.to_entries(), [args.dataset_label], nch, args.ctx) + peak_rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1e6 + line = f"wrote {len(acc.sums)} entries -> {args.out} | peak RSS {peak_rss:.1f} GB" + if dev == "cuda": + line += f" | peak GPU {torch.cuda.max_memory_allocated()/1e9:.2f} GB" + print(line) + + +if __name__ == "__main__": + main() diff --git a/scripts/imatrix_serialized/test_calib.py b/scripts/imatrix_serialized/test_calib.py new file mode 100644 index 0000000..1d8b50d --- /dev/null +++ b/scripts/imatrix_serialized/test_calib.py @@ -0,0 +1,85 @@ +"""Synthetic checks for calibration windowing + llama.cpp BOS parity. No model needed. + +Run: python3 test_calib.py (or under pytest) +""" +import torch + +from calib import load_calibration_chunks + + +class FakeTok: + """Minimal stand-in: one token per character, optional BOS id.""" + def __init__(self, bos_token_id=None): + self.bos_token_id = bos_token_id + + def __call__(self, text, return_tensors="pt", add_special_tokens=False): + assert add_special_tokens is False, "corpus must be tokenized as a clean stream" + ids = torch.tensor([[ord(c) for c in text]]) + return type("Enc", (), {"input_ids": ids})() + + +def _corpus(tmp, n): + p = tmp / "calib.txt" + p.write_text("".join(chr(65 + (i % 26)) for i in range(n))) + return str(p) + + +def test_windowing_shape_and_contiguity(tmp_path): + path = _corpus(tmp_path, 100) + chunks, info = load_calibration_chunks(FakeTok(), path, n_chunks=99, ctx=10, add_bos=False) + assert chunks.shape == (10, 10), chunks.shape + assert chunks.dtype == torch.int64 + # windows must tile the stream in order with no gaps + flat = chunks.reshape(-1) + assert flat.tolist() == [ord(chr(65 + (i % 26))) for i in range(100)] + assert "chunks=10" in info + + +def test_chunks_cap_is_respected(tmp_path): + path = _corpus(tmp_path, 100) + chunks, _ = load_calibration_chunks(FakeTok(), path, n_chunks=3, ctx=10, add_bos=False) + assert chunks.shape == (3, 10) + + +def test_bos_forced_at_position_zero_of_every_chunk(tmp_path): + """llama.cpp imatrix.cpp L865-866 overwrites token 0 of each chunk with BOS.""" + path = _corpus(tmp_path, 100) + chunks, info = load_calibration_chunks(FakeTok(bos_token_id=7), path, + n_chunks=99, ctx=10, add_bos=True) + assert (chunks[:, 0] == 7).all(), chunks[:, 0] + assert "llama.cpp parity" in info + # only position 0 is touched; the rest of each window is untouched stream + plain, _ = load_calibration_chunks(FakeTok(bos_token_id=7), path, + n_chunks=99, ctx=10, add_bos=False) + assert torch.equal(chunks[:, 1:], plain[:, 1:]) + assert not torch.equal(chunks[:, 0], plain[:, 0]), "BOS should differ from the stream" + + +def test_no_bos_id_is_not_fatal(tmp_path): + path = _corpus(tmp_path, 30) + chunks, info = load_calibration_chunks(FakeTok(bos_token_id=None), path, + n_chunks=99, ctx=10, add_bos=True) + assert chunks.shape == (3, 10) + assert "no bos_token_id" in info + + +def test_corpus_too_short_hard_fails(tmp_path): + path = _corpus(tmp_path, 5) + try: + load_calibration_chunks(FakeTok(), path, n_chunks=10, ctx=512) + except ValueError as e: + assert "too short" in str(e) + return + raise AssertionError("expected ValueError on a corpus shorter than one chunk") + + +if __name__ == "__main__": + import pathlib, tempfile + with tempfile.TemporaryDirectory() as d: + for name, fn in sorted(globals().items()): + if name.startswith("test_"): + sub = pathlib.Path(d) / name + sub.mkdir() + fn(sub) + print(f"ok {name}") + print("all passed") diff --git a/scripts/imatrix_serialized/test_causal_mask.py b/scripts/imatrix_serialized/test_causal_mask.py new file mode 100644 index 0000000..b6dd978 --- /dev/null +++ b/scripts/imatrix_serialized/test_causal_mask.py @@ -0,0 +1,95 @@ +"""Guards against a silently NON-CAUSAL imatrix. No checkpoint needed. + +The bug: transformers' create_causal_mask() returns None when +config._attn_implementation is unset (it assumes the attention backend applies +causality itself, as sdpa does via is_causal). The band-serialized generators do not +run a model -- they invoke standalone decoder layers -- so a None mask means those +layers attend BIDIRECTIONALLY. Nothing errors; you just get an imatrix whose +attn_output statistic (and everything downstream of it) was computed with future +tokens visible. + +Measured on a synthetic 4-layer GraniteMoeHybrid: serialized vs full-forward +blk.3.attn_output disagreed by 7.5e-01 (corr 0.54) with the implementation unset, and +by 3.2e-06 once it was set. So this is worth a test rather than a comment. + +Run: python3 test_causal_mask.py (or under pytest) +""" +import os +import re + +import torch +from transformers.masking_utils import create_causal_mask +from transformers.models.granitemoehybrid.configuration_granitemoehybrid import ( + GraniteMoeHybridConfig, +) + +HERE = os.path.dirname(os.path.abspath(__file__)) + +# Every generator that builds masks itself and drives standalone decoder layers. +SERIALIZED_GENERATORS = [ + "serialized_gen.py", + "serialized_gen_t514.py", + "serialized_gen_t514_fast.py", + "serialized_gen_qwen35.py", + "serialized_gen_llama4.py", + "glm4moe_serialized_gen.py", +] + + +def _cfg(): + return GraniteMoeHybridConfig(hidden_size=64, num_hidden_layers=2, + num_attention_heads=4, num_key_value_heads=2, + intermediate_size=128, mamba_n_heads=4, + mamba_d_state=16, mamba_d_conv=4, mamba_expand=2, + layer_types=["attention", "attention"]) + + +def test_unset_implementation_really_yields_no_mask(): + """Pins the transformers behaviour the bug depends on, so a change is visible.""" + cfg = _cfg() + assert getattr(cfg, "_attn_implementation", None) in (None, "eager"), \ + "fixture assumption changed" + cfg._attn_implementation = None + m = create_causal_mask(config=cfg, inputs_embeds=torch.zeros(1, 6, cfg.hidden_size), + attention_mask=None, past_key_values=None) + assert m is None, f"expected None (the trap), got {type(m)}" + + +def test_setting_eager_yields_a_real_causal_mask(): + cfg = _cfg() + cfg._attn_implementation = "eager" + m = create_causal_mask(config=cfg, inputs_embeds=torch.zeros(1, 6, cfg.hidden_size), + attention_mask=None, past_key_values=None) + assert m is not None, "no mask -> standalone layers would attend bidirectionally" + assert m.shape[-2:] == (6, 6), m.shape + # strictly lower-triangular-inclusive: position i must not see j > i + allowed = (m[0, 0] == 0) if m.dtype != torch.bool else m[0, 0] + expect = torch.tril(torch.ones(6, 6, dtype=torch.bool)) + assert torch.equal(allowed.bool(), expect), f"mask is not causal:\n{allowed.int()}" + + +def test_every_serialized_generator_sets_the_attn_implementation(): + """A generator that forgets this line emits a non-causal imatrix with no error.""" + missing = [] + for fn in SERIALIZED_GENERATORS: + path = os.path.join(HERE, fn) + with open(path) as f: + src = f.read() + if not re.search(r"_attn_implementation\s*=", src): + missing.append(fn) + assert not missing, f"generators never set config._attn_implementation: {missing}" + + +def test_generators_expose_the_attn_impl_flag(): + """The knob should be visible, not buried, since it changes numerics.""" + missing = [fn for fn in SERIALIZED_GENERATORS + if "--attn-impl" not in open(os.path.join(HERE, fn)).read()] + assert not missing, f"generators missing the --attn-impl flag: {missing}" + + +if __name__ == "__main__": + for name, fn in sorted(globals().items()): + if name.startswith("test_"): + fn() + print(f"ok {name}") + print("all passed") diff --git a/scripts/imatrix_serialized/test_dispatch.py b/scripts/imatrix_serialized/test_dispatch.py new file mode 100644 index 0000000..3829770 --- /dev/null +++ b/scripts/imatrix_serialized/test_dispatch.py @@ -0,0 +1,126 @@ +"""Checks the Phase 6 generator dispatch. No checkpoint, no GPU. + +Guards the blocking issue from review round 3: apex_pipeline.sh Phase 6 called +serialized_gen.py unconditionally, which is GraniteMoeHybrid-only and 5.5-only, so a +fresh install per requirements.txt ImportError'd and the qwen3_5_moe case that +motivates the backend could never run through the wired path. + +Run: python3 test_dispatch.py (or under pytest) +""" +import json +import os +import tempfile + +import dispatch + + +def _model_dir(d, cfg): + os.makedirs(d, exist_ok=True) + with open(os.path.join(d, "config.json"), "w") as f: + json.dump(cfg, f) + return d + + +def test_every_arch_entry_points_at_a_file_that_exists(): + """A table entry naming a missing script would fail only at run time.""" + for mt, variants in dispatch.ARCHES.items(): + for v in variants: + p = os.path.join(dispatch.HERE, v.script) + assert os.path.exists(p), f"{mt}: {v.script} does not exist" + + +def test_selection_picks_a_variant_the_install_supports(): + """Whatever transformers is installed, a resolved variant must be importable.""" + import importlib + for mt in dispatch.ARCHES: + try: + script, variant = dispatch.select(mt) + except SystemExit: + continue # nothing installed for this arch; fine + m = importlib.import_module(variant.module) + for sym in variant.symbols: + assert hasattr(m, sym), f"{mt}: chose {variant.script} but {sym} is missing" + assert os.path.exists(script) + + +def test_unknown_model_type_fails_with_supported_list(): + try: + dispatch.select("definitely_not_an_arch") + except SystemExit as e: + msg = str(e) + assert "no serialized generator" in msg + assert "supported:" in msg + # must point at the escape hatch rather than just dying + assert "IMATRIX_BACKEND=llama" in msg + return + raise AssertionError("expected SystemExit for an unknown model_type") + + +def test_unavailable_variant_names_what_is_missing(): + """A supported arch with no importable variant must say what to install.""" + saved = dict(dispatch.ARCHES) + try: + dispatch.ARCHES["fake_arch"] = [ + dispatch.Variant("serialized_gen.py", "transformers", ["NoSuchSymbol"], + "transformers with NoSuchSymbol")] + try: + dispatch.select("fake_arch") + except SystemExit as e: + msg = str(e) + assert "NoSuchSymbol" in msg, msg + assert "transformers with NoSuchSymbol" in msg, msg + return + raise AssertionError("expected SystemExit when no variant is importable") + finally: + dispatch.ARCHES.clear() + dispatch.ARCHES.update(saved) + + +def test_missing_module_is_reported_not_raised(): + v = dispatch.Variant("serialized_gen.py", "no_such_module_xyz", ["X"], "n/a") + miss = v.missing() + assert miss and "no_such_module_xyz" in miss[0], miss + + +def test_model_type_read_from_config_and_text_config(): + with tempfile.TemporaryDirectory() as d: + plain = _model_dir(os.path.join(d, "a"), {"model_type": "qwen3_5_moe"}) + assert dispatch.read_model_type(plain) == "qwen3_5_moe" + # multimodal wrapper with a SUPPORTED top-level type keeps the top-level name + nested = _model_dir(os.path.join(d, "b"), + {"model_type": "llama4", "text_config": {"model_type": "llama4_text"}}) + assert dispatch.read_model_type(nested) == "llama4" + # ...but an unsupported wrapper falls through to the inner decoder arch + wrapped = _model_dir(os.path.join(d, "c"), + {"model_type": "some_vlm", "text_config": {"model_type": "qwen3_5_moe"}}) + assert dispatch.read_model_type(wrapped) == "qwen3_5_moe" + + +def test_whatever_read_model_type_returns_is_dispatchable(): + """A model_type that parses but is absent from ARCHES is a silent dead end.""" + with tempfile.TemporaryDirectory() as d: + for cfg in ({"model_type": "llama4", "text_config": {"model_type": "llama4_text"}}, + {"model_type": "llama4_text"}, + {"model_type": "qwen3_5_moe"}, + {"model_type": "granitemoehybrid"}): + md = _model_dir(os.path.join(d, str(abs(hash(str(cfg))))), cfg) + mt = dispatch.read_model_type(md) + assert mt in dispatch.ARCHES, f"{cfg} -> {mt!r} not in ARCHES" + + +def test_missing_config_fails_clearly(): + with tempfile.TemporaryDirectory() as d: + try: + dispatch.read_model_type(d) + except SystemExit as e: + assert "no config.json" in str(e) + return + raise AssertionError("expected SystemExit when config.json is absent") + + +if __name__ == "__main__": + for name, fn in sorted(globals().items()): + if name.startswith("test_"): + fn() + print(f"ok {name}") + print("all passed") diff --git a/scripts/imatrix_serialized/test_glm4_expert_stats.py b/scripts/imatrix_serialized/test_glm4_expert_stats.py new file mode 100644 index 0000000..4283279 --- /dev/null +++ b/scripts/imatrix_serialized/test_glm4_expert_stats.py @@ -0,0 +1,98 @@ +"""Synthetic checks for the GLM-4.5 routed-expert statistic. No checkpoint needed. + +The bug this locks down: mapping mlp.experts to all three of +ffn_{gate,up,down}_exps and feeding them one hidden-dim Sum(x^2). ffn_down_exps takes +moe_intermediate_size inputs (the gated intermediate), so that entry was wrong in +both shape and value. + +Run: python3 test_glm4_expert_stats.py (or under pytest) +""" +import torch +import torch.nn.functional as F + +from glm4moe_serialized_gen import apply_gate_up, expert_stats + + +HIDDEN, INTERM, NE, K, T = 16, 6, 4, 2, 32 + + +def _fixture(layout, seed=0): + g = torch.Generator().manual_seed(seed) + x = torch.randn(T, HIDDEN, generator=g, dtype=torch.float64) + topk = torch.stack([torch.randperm(NE, generator=g)[:K] for _ in range(T)]) + w = torch.randn(NE, 2 * INTERM, HIDDEN, generator=g, dtype=torch.float64) + return x, topk, (w if layout == "linear" else w.transpose(1, 2).contiguous()) + + +def _reference(x, topk, w_linear, act): + """Independent per-expert loop: what each expert's two matmuls actually see.""" + s_in = torch.zeros(NE, HIDDEN, dtype=torch.float64) + s_inter = torch.zeros(NE, INTERM, dtype=torch.float64) + cnt = torch.zeros(NE, dtype=torch.float64) + for e in range(NE): + rows = [t for t in range(T) if e in topk[t].tolist()] + for t in rows: + xt = x[t] + s_in[e] += xt * xt + g, u = F.linear(xt, w_linear[e]).chunk(2, dim=-1) + inter = act(g) * u + s_inter[e] += inter * inter + cnt[e] = len(rows) + return s_in, s_inter, cnt + + +def test_shapes_and_values_match_reference(): + """down stat has INTERM entries and equals the gated intermediate's Sum(x^2).""" + for layout in ("linear", "bmm"): + x, topk, w = _fixture(layout) + w_linear = w if layout == "linear" else w.transpose(1, 2) + s_in, s_inter, cnt = expert_stats(x, topk, w, F.silu, NE, HIDDEN) + r_in, r_inter, r_cnt = _reference(x, topk, w_linear, F.silu) + + assert s_in.shape == (NE, HIDDEN), (layout, s_in.shape) + assert s_inter.shape == (NE, INTERM), (layout, s_inter.shape) + assert s_inter.shape[-1] != s_in.shape[-1], "down stat must not be hidden-dim" + # expert_stats computes in float32 (matching the real activation dtype) and + # accumulates in float64, so it tracks this float64 reference to ~float32 eps. + torch.testing.assert_close(s_in, r_in, rtol=1e-6, atol=1e-5) + torch.testing.assert_close(s_inter, r_inter, rtol=1e-6, atol=1e-5) + torch.testing.assert_close(cnt, r_cnt) + # every (token, selected-expert) slot counted exactly once + assert cnt.sum().item() == T * K + + +def test_both_fused_layouts_agree(): + """Layout is derived from the shape, so the two conventions give one answer.""" + a = expert_stats(*_fixture("linear"), F.silu, NE, HIDDEN) + b = expert_stats(*_fixture("bmm"), F.silu, NE, HIDDEN) + for ta, tb in zip(a, b): + torch.testing.assert_close(ta, tb, rtol=1e-9, atol=1e-9) + + +def test_ambiguous_and_bad_layouts_hard_fail(): + """A layout we cannot resolve must raise, never silently pick a side.""" + x = torch.zeros(4, HIDDEN, dtype=torch.float64) + for w, why in [(torch.zeros(HIDDEN, HIDDEN, dtype=torch.float64), "square"), + (torch.zeros(7, 9, dtype=torch.float64), "no hidden axis")]: + try: + apply_gate_up(x, w, HIDDEN) + except ValueError: + continue + raise AssertionError(f"apply_gate_up accepted {why} weight {tuple(w.shape)}") + + +def test_wrong_input_dim_hard_fails(): + x, topk, w = _fixture("linear") + try: + expert_stats(x[:, :HIDDEN - 1], topk, w, F.silu, NE, HIDDEN) + except ValueError: + return + raise AssertionError("expert_stats accepted a mismatched hidden dim") + + +if __name__ == "__main__": + for name, fn in sorted(globals().items()): + if name.startswith("test_"): + fn() + print(f"ok {name}") + print("all passed") diff --git a/scripts/imatrix_serialized/test_llama4_expert_counts.py b/scripts/imatrix_serialized/test_llama4_expert_counts.py new file mode 100644 index 0000000..fc66b4a --- /dev/null +++ b/scripts/imatrix_serialized/test_llama4_expert_counts.py @@ -0,0 +1,124 @@ +"""Drives the Llama-4 hooks through a real (tiny, random) Llama4TextMoe on CPU. + +Locks down the two things review round 3 raised about serialized_gen_llama4.py: + + 1. Row counts must be llama.cpp's one-per-(token, selected-expert) slot, taken from + the router's exact top-k -- not the old `abs().sum(-1) > 0` nonzero-row heuristic, + which miscounts a genuinely-zero activation row or an underflowed router score. + 2. The router-weighted gate/up statistic is CORRECT for this arch. llama.cpp sets + `weight_before_ffn = arch == LLM_ARCH_LLAMA4` (src/llama-graph.cpp:1837) and + multiplies the sigmoid-ed weights into the expert input BEFORE the gate/up + mul_mat_id (L1976), so a weighted stat is the faithful reproduction. This test + pins the weighted behaviour so a future "fix" to unweighted fails loudly. + +Run: python3 test_llama4_expert_counts.py (or under pytest) +""" +import numpy as np +import torch + +from transformers.models.llama4.configuration_llama4 import Llama4TextConfig +from transformers.models.llama4.modeling_llama4 import Llama4TextMoe + +from serialized_gen_llama4 import make_ensure, hook_experts, hook_router_select + +H, INTERM, NE, TOPK, T = 32, 12, 6, 2, 20 +BLK = "blk.0" + + +def _moe(seed=0): + torch.manual_seed(seed) + cfg = Llama4TextConfig(hidden_size=H, intermediate_size=INTERM, num_local_experts=NE, + num_experts_per_tok=TOPK, num_hidden_layers=1, + num_attention_heads=4, num_key_value_heads=2) + moe = Llama4TextMoe(cfg).eval() + for p in moe.parameters(): + torch.nn.init.normal_(p, std=0.05) + return moe + + +def _run(moe, x): + acc = {} + sel = {} + ensure = make_ensure(acc) + handles = [hook_router_select(moe.router, BLK, sel), + hook_experts(moe.experts, BLK, ensure, sel)] + with torch.no_grad(): + moe(x) + for h in handles: + h.remove() + return acc, sel + + +def test_counts_are_exact_topk_slots(): + moe = _moe() + x = torch.randn(T, H) + acc, sel = _run(moe, x) + + gate = acc[f"{BLK}.ffn_gate_exps.weight"] + # every token contributes to exactly TOPK experts -> total slots = T * TOPK + assert gate["counts"].sum() == T * TOPK, (gate["counts"].sum(), T * TOPK) + # and per-expert counts must equal the router's own top-k tally + idx = sel[BLK] + expect = np.bincount(idx.reshape(-1).numpy(), minlength=NE) + assert np.array_equal(gate["counts"], expect), (gate["counts"], expect) + # all three routed tensors share the same row count + for t in ("ffn_up_exps", "ffn_down_exps"): + assert np.array_equal(acc[f"{BLK}.{t}.weight"]["counts"], expect), t + + +def test_zero_activation_row_does_not_lose_a_count(): + """The old nonzero-row heuristic undercounts here; the top-k mask does not.""" + moe = _moe() + x = torch.randn(T, H) + x[3] = 0.0 # a genuinely zero hidden row, still routed + acc, sel = _run(moe, x) + + idx = sel[BLK] + expect = np.bincount(idx.reshape(-1).numpy(), minlength=NE) + got = acc[f"{BLK}.ffn_gate_exps.weight"]["counts"] + assert got.sum() == T * TOPK, got.sum() + assert np.array_equal(got, expect) + # demonstrate the old heuristic really would have been wrong on this input + with torch.no_grad(): + scores, _ = moe.router(x) + X = (x.repeat(scores.shape[1], 1) * scores.t().reshape(-1, 1)).view(NE, T, H) + old = (X.abs().sum(-1) > 0).sum(1).numpy() + assert old.sum() < expect.sum(), (old.sum(), expect.sum()) + + +def test_shapes_and_router_weighting_is_preserved(): + moe = _moe() + x = torch.randn(T, H) + acc, sel = _run(moe, x) + + gate = acc[f"{BLK}.ffn_gate_exps.weight"] + down = acc[f"{BLK}.ffn_down_exps.weight"] + assert gate["sums"].shape == (NE, H), gate["sums"].shape + assert down["sums"].shape == (NE, INTERM), down["sums"].shape + # gate and up share one input stat (same MUL_MAT_ID input in the graph) + assert np.array_equal(gate["sums"], acc[f"{BLK}.ffn_up_exps.weight"]["sums"]) + + # The recorded gate/up stat must be the ROUTER-WEIGHTED second moment, matching + # llama.cpp's weight_before_ffn path. Compare against an explicit recomputation. + with torch.no_grad(): + scores, _ = moe.router(x) + weighted = (x.repeat(scores.shape[1], 1) * scores.t().reshape(-1, 1)).view(NE, T, H) + idx = sel[BLK] + keep = torch.zeros(T, NE, dtype=torch.bool).scatter_(1, idx, True) + expect = ((weighted.float() * keep.t().unsqueeze(-1)) ** 2).sum(1).double().numpy() + np.testing.assert_allclose(gate["sums"], expect, rtol=1e-6, atol=1e-9) + + # sanity: the unweighted stat is materially different, so this really is pinning + # the weighted convention rather than passing trivially + with torch.no_grad(): + unweighted = ((x.repeat(NE, 1).view(NE, T, H).float() * keep.t().unsqueeze(-1)) ** 2 + ).sum(1).double().numpy() + assert not np.allclose(gate["sums"], unweighted, rtol=1e-3) + + +if __name__ == "__main__": + for name, fn in sorted(globals().items()): + if name.startswith("test_"): + fn() + print(f"ok {name}") + print("all passed") diff --git a/scripts/imatrix_serialized/test_roundtrip.py b/scripts/imatrix_serialized/test_roundtrip.py new file mode 100644 index 0000000..80bb9d6 --- /dev/null +++ b/scripts/imatrix_serialized/test_roundtrip.py @@ -0,0 +1,124 @@ +"""Format-layer round-trip: write -> read -> assert identical. + +Previously this needed a local `models/granite-4.0-h-tiny.imatrix` and wrote into +`imatrix-reverse/`, so it could not run anywhere but one dev box and CI never executed +it. It is now synthetic and self-contained: it builds entries covering both shapes the +writer supports (dense nmat=1 and per-expert nmat>1), round-trips through a temp file, +and checks the GGUF layout llama.cpp expects (in_sum2 ne=[in,nmat], counts ne=[1,nmat]) +plus the metadata keys. + +Still accepts a real file for a spot check: + python3 test_roundtrip.py path/to/some.imatrix + +Needs numpy + gguf only -- no torch, no transformers, no checkpoint. + +Run: python3 test_roundtrip.py (or under pytest) +""" +import os +import tempfile + +import numpy as np + +from imatrix_io import read_imatrix, write_imatrix + +DATASETS = ["calibration_datav3"] +CHUNKS, CTX = 126, 512 + + +def _entries(seed=0): + rng = np.random.default_rng(seed) + return { + # dense tensors: one matrix, in_sum2 given as a 1-D vector + "blk.0.attn_q.weight": {"in_sum2": (rng.random(64) * 1e3).astype(np.float32), + "counts": np.array([CHUNKS * CTX], np.float32)}, + "blk.0.ffn_down.weight": {"in_sum2": rng.random(48).astype(np.float32), + "counts": np.array([CHUNKS * CTX], np.float32)}, + # per-expert tensors: nmat = n_experts, in_sum2 is 2-D [nmat, in_features] + "blk.0.ffn_gate_exps.weight": {"in_sum2": rng.random((8, 64)).astype(np.float32), + "counts": rng.integers(1, 999, 8).astype(np.float32)}, + "blk.0.ffn_down_exps.weight": {"in_sum2": rng.random((8, 24)).astype(np.float32), + "counts": rng.integers(1, 999, 8).astype(np.float32)}, + } + + +def _roundtrip(entries): + with tempfile.TemporaryDirectory() as d: + path = os.path.join(d, "rt.imatrix") + write_imatrix(path, entries, DATASETS, CHUNKS, CTX) + return read_imatrix(path) + + +def test_names_survive(): + src = _entries() + _, got = _roundtrip(src) + assert set(got) == set(src), (set(src) ^ set(got)) + + +def test_values_are_bit_exact(): + """float32 in, float32 out, no scaling anywhere -> must be exactly equal.""" + src = _entries() + _, got = _roundtrip(src) + for name, e in src.items(): + want = np.atleast_2d(np.asarray(e["in_sum2"], np.float32)) + assert np.array_equal(got[name]["in_sum2"].reshape(want.shape), want), name + wc = np.asarray(e["counts"], np.float32).reshape(-1) + assert np.array_equal(got[name]["counts"].reshape(-1), wc), f"{name} counts" + + +def test_gguf_shapes_match_llama_cpp_layout(): + """in_sum2 ne=[in_features, nmat] and counts ne=[1, nmat] (imatrix.cpp L597-606).""" + src = _entries() + _, got = _roundtrip(src) + for name, e in src.items(): + nmat = np.asarray(e["counts"], np.float32).reshape(-1).shape[0] + n_in = np.atleast_2d(np.asarray(e["in_sum2"], np.float32)).shape[-1] + assert got[name]["in_sum2"].size == nmat * n_in, name + assert got[name]["counts"].size == nmat, f"{name} counts" + + +def test_metadata_keys_are_written(): + with tempfile.TemporaryDirectory() as d: + path = os.path.join(d, "rt.imatrix") + write_imatrix(path, _entries(), DATASETS, CHUNKS, CTX) + meta, _ = read_imatrix(path) + assert meta.get("general.type") == "imatrix", meta.get("general.type") + for k in ("imatrix.datasets", "imatrix.chunk_count", "imatrix.chunk_size"): + assert k in meta, f"missing {k}" + cc = meta["imatrix.chunk_count"] + assert (cc[0] if isinstance(cc, list) else cc) == CHUNKS, cc + + +def test_single_expert_entry_is_not_squeezed(): + """nmat=1 given as 2-D must stay a valid entry, not collapse to a dense vector.""" + src = {"blk.1.ffn_gate_exps.weight": {"in_sum2": np.ones((1, 16), np.float32), + "counts": np.array([5], np.float32)}} + _, got = _roundtrip(src) + assert got["blk.1.ffn_gate_exps.weight"]["in_sum2"].size == 16 + + +def spot_check(path): + """Optional: re-write a real imatrix and confirm values/names survive.""" + meta, entries = read_imatrix(path) + src = {b: {"in_sum2": e["in_sum2"], "counts": e["counts"].reshape(-1)} + for b, e in entries.items()} + _, got = _roundtrip(src) + assert set(got) == set(src), "name set differs" + worst = 0.0 + for b in src: + a = src[b]["in_sum2"].astype(np.float64).ravel() + c = got[b]["in_sum2"].astype(np.float64).ravel() + assert a.shape == c.shape, (b, a.shape, c.shape) + assert np.array_equal(src[b]["counts"].ravel(), got[b]["counts"].ravel()), f"{b} counts" + worst = max(worst, np.abs(a - c).max()) + print(f"spot check OK: {len(src)} entries, max abs in_sum2 diff = {worst:.3e}") + + +if __name__ == "__main__": + import sys + for name, fn in sorted(globals().items()): + if name.startswith("test_"): + fn() + print(f"ok {name}") + print("all passed") + if len(sys.argv) > 1: + spot_check(sys.argv[1])