quantize_mxfp4.py

#!/usr/bin/env python
"""Quantize an Ornith-1.0 (Qwen3.5 dense or MoE) BF16 trunk to MXFP4 — the exact
weight-only scheme this release ships: compressed-tensors **MXFP4A16**, 4-bit float,
group size 32, symmetric, E8M0 microscales (config.json format "mxfp4-pack-quantized").

This is STEP 1 of the pipeline; recipe_graft_mxfp4.py is steps 2-3:

    quantize_mxfp4.py      BF16 trunk               -> MXFP4 trunk        <-- THIS FILE
    recipe_graft_mxfp4.py  MXFP4 trunk + donor head -> shipped MTP repo
    vllm serve <repo> --speculative-config '{"method":"mtp","num_speculative_tokens":3}'

WHAT IS QUANTIZED vs KEPT BF16 (reproduces this repo's config.json `ignore` list):
  MXFP4 : the routed-expert MLP projections (MoE 35B) / dense MLP projections (9B),
          plus the linear-attention (DeltaNet) matmul projections — i.e. the bulk of
          the weights.
  BF16  : lm_head, embed_tokens, ALL norms, the MoE router gate + shared expert,
          the full self-attention q/k/v/o projections, linear_attn.conv1d, and the
          ENTIRE vision tower.

Why weight-only (MXFP4A16) and not "MXFP4": the shipped config quantizes weights
only (no `input_activations`); the W4A4 "MXFP4" preset would add activation quant and
NOT byte-match. Norms are auto-ignored by model_free_ptq, so IGNORE lists only the
non-norm modules. The patterns are architecture-adaptive: the MoE-only ones
(`mlp.gate`, `shared_expert*`) match nothing on the dense 9B, so one set covers both.

Data-free & deterministic: MXFP4 weight quant is round-to-nearest with per-group E8M0
scales — no calibration data, no global scale, no reindex (that's NVFP4-only). So the
output is byte-reproducible against the published trunk. model_free_ptq streams each
safetensors shard independently (never materializes the full BF16 model), so the 35B
fits in host RAM and the custom Qwen3.5 arch never has to load through transformers.

Requires: a recent llm-compressor (with `model_free_ptq`) + compressed-tensors.
Credits: trunk — DeepReinforce (Ornith-1.0, MIT); quantizer — llm-compressor /
compressed-tensors (Neural Magic / Red Hat). Combined-work license: Apache-2.0.

Usage:
  # 35B MoE
  python quantize_mxfp4.py --model deepreinforce-ai/Ornith-1.0-35B --out ./Ornith-1.0-35B-MXFP4
  # 9B dense  (same script, same ignore patterns)
  python quantize_mxfp4.py --model deepreinforce-ai/Ornith-1.0-9B  --out ./Ornith-1.0-9B-MXFP4
  # then transplant the MTP head:
  python recipe_graft_mxfp4.py --donor <distill-donor> --target ./Ornith-1.0-35B-MXFP4 \
      --out ./Ornith-1.0-35B-MXFP4-Vision-MTP

Verify (byte-level) after running, before grafting:
  # the ignore list the quant emitted should equal this repo's config.json ignore
  # MINUS the 785 (35B) / 10 (9B) mtp.* entries the graft adds later; and each
  # non-ignored shard should SHA256-match the published trunk.
"""
import argparse

from llmcompressor import model_free_ptq

# Modules kept in BF16. `model_free_ptq` additionally auto-ignores every module whose
# name ends in "norm", which covers input/post layernorms, q_norm/k_norm,
# linear_attn.norm and the final model norm — so they are intentionally NOT listed here.
IGNORE = [
    "lm_head",
    "re:.*embed_tokens$",                  # token embeddings
    r"re:.*self_attn\.(q|k|v|o)_proj$",    # full self-attention kept BF16
    r"re:.*linear_attn\.conv1d$",          # DeltaNet short conv (its .norm is auto-ignored)
    r"re:.*mlp\.gate$",                     # MoE router gate       (no-op on dense 9B)
    r"re:.*shared_expert\..*",             # MoE shared-expert MLP  (35B only)
    r"re:.*shared_expert_gate$",           # MoE shared-expert gate (35B only)
    "re:.*visual.*",                        # entire vision tower kept BF16
]


def main() -> int:
    ap = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--model", required=True,
                    help="BF16 trunk: HF id or local dir (e.g. deepreinforce-ai/Ornith-1.0-35B)")
    ap.add_argument("--out", required=True,
                    help="output dir for the MXFP4 compressed-tensors trunk")
    ap.add_argument("--workers", type=int, default=8,
                    help="parallel safetensors worker threads (default: 8)")
    ap.add_argument("--device", default=None,
                    help="quant device; on ROCm the AMD GPU is 'cuda:0' (default: auto — "
                         "GPU if visible, else CPU)")
    a = ap.parse_args()

    # Weight-only MXFP4: num_bits=4, type=float, strategy=group, group_size=32,
    # symmetric=True, E8M0 (uint8) scales — identical to the shipped config_groups.
    model_free_ptq(
        model_stub=a.model,
        save_directory=a.out,
        scheme="MXFP4A16",
        ignore=IGNORE,
        max_workers=a.workers,
        device=a.device,
    )
    print(f"OK -> {a.out}   (compressed-tensors MXFP4A16 / mxfp4-pack-quantized)")
    print("next: graft the MTP head with recipe_graft_mxfp4.py "
          f"(--target {a.out})")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())