feat(imatrix): opt-in band-serialized PyTorch imatrix backend for large/hybrid MoE

.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

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