YuE2-3B AR NVFP4 (nvfp4a16-gptq) — llm-compressor requant, 2026-09-13

README.md

@@ -0,0 +1,120 @@
+---
+license: cc-by-nc-4.0
+base_model: m-a-p/YuE2-3B
+base_model_relation: quantized
+pipeline_tag: text-to-audio
+tags:
+  - yue2
+  - music-generation
+  - nvfp4
+  - compressed-tensors
+  - vllm
+  - dgx-spark
+language:
+  - en
+  - zh
+---
+
+# YuE2-3B — NVFP4 autoregressive checkpoint (W4A16, GPTQ)
+
+This is an **NVFP4 requantisation of the autoregressive (AR) half of
+[m-a-p/YuE2-3B](https://huggingface.co/m-a-p/YuE2-3B)**, packaged as the
+`Qwen3ForCausalLM`-shaped checkpoint that the official `yue2_infer` pipeline
+derives internally for its vLLM backend. It is a drop-in for that derived
+checkpoint: the pipeline's planning (`abc`) and semantic-token stages run on
+these weights; the NAR flow-matching stage and the VAE keep running on the
+original bf16/fp32 weights, unchanged.
+
+It exists because on unified-memory hardware (DGX Spark / GB10, 273 GB/s) the
+AR stages are weight-bandwidth-bound and the bf16 AR checkpoint is 4.33 GB.
+
+| | bf16 (derived by `yue2_infer`) | this checkpoint |
+|---|---|---|
+| AR weights on disk | 4.33 GB | 2.31 GB |
+| Linear layers | bf16 | NVFP4 (W4A16: e2m1 weights, e4m3 per-16 scales + fp32 per-tensor global scale, bf16 activations), group 16 |
+| `embed_tokens`, `lm_head` (vocab 184,704) | bf16 | bf16 (kept) |
+| AR decode, DGX Spark GB10, vLLM 0.27.1, batch 1 | 44.8 / 43.0 tok/s (abc / semantic) | 100–107 / 95–97 tok/s (2.3×; 76 tok/s on a 9,000-token run) |
+
+## What was quantised, and how
+
+- Source: `m-a-p/YuE2-3B` `model.safetensors` (sha256 `1d55c42c…`), AR tensors
+  extracted exactly as `yue2.fast.derive_ar_checkpoint` does (311 tensors:
+  attention q/k/v/o + q/k norms, AR `mlp.*`, layer norms, `embed_tokens`,
+  `lm_head`). The NAR MLPs (`nar_mlp.*`), `llm2vae`/`vae2llm`, the timestep
+  embedder and the VAE are **not** in this checkpoint and are not quantised.
+- Tool: [llm-compressor](https://github.com/vllm-project/llm-compressor) 0.13.0 /
+  compressed-tensors 0.18.0. Recipe: `GPTQModifier(targets=["Linear"], scheme="NVFP4A16", ignore=["lm_head"], dampening_frac=0.01, block_size=128, actorder="static")`. `lm_head` excluded.
+- Calibration: **pre-tokenised YuE2 sequences the bf16 model generated itself**
+  — 32 requests × (abc + semantic phases), 16 lyric sets × 2 contrasting style
+  prompts, English and Mandarin, `cot=full` and `cot=melody`; 58 windows /
+  200K tokens, 51% codec tokens. Text corpora never exercise the codec-token
+  activations, so they were not used.
+- Format: `compressed-tensors` `nvfp4-pack-quantized`, loadable by vLLM ≥ 0.10
+  on Blackwell (SM100/SM120/SM121) with native NVFP4 kernels, and by
+  transformers (weights dequantised to bf16).
+
+## Quality — teacher-forced logit agreement vs bf16
+
+Measured through vLLM 0.27.1's real NVFP4 kernels on 4 held-out requests
+(not in calibration), every output position teacher-forced against the bf16
+model's own sampled sequences (6,286 abc positions, 17,931 semantic positions).
+Top-1 agreement is entropy-sensitive: the codec distribution is flat almost
+everywhere (bf16 itself puts its argmax on the sampled token only 22% of the
+time), so the NLL and logprob-shift columns are the ones to read. The
+bf16-vs-bf16 row is the numerics floor across two engine builds.
+
+| Checkpoint | abc top-1 agree | abc NLL | semantic top-1 agree | semantic NLL | mean \|Δlogprob\| |
+|---|---|---|---|---|---|
+| bf16 vs bf16 (engine builds) | 0.999 | 0.141 / 0.141 | 0.982 | 3.414 / 3.414 | 0.02 |
+| **this checkpoint — NVFP4 W4A16, GPTQ** | **0.996** | **0.141 / 0.142** | **0.905** | **3.414 / 3.424 (+0.3%)** | **0.10** |
+| NVFP4 W4A16, AWQ+RTN | 0.993 | 0.141 / 0.143 | 0.896 | 3.414 / 3.433 (+0.6%) | 0.12 |
+| NVFP4 W4A4, GPTQ | 0.992 | 0.141 / 0.144 | 0.854 | 3.414 / 3.446 (+0.9%) | 0.17 |
+| NVFP4 W4A4, AWQ+RTN | 0.991 | 0.141 / 0.145 | 0.847 | 3.414 / 3.457 (+1.2%) | 0.18 |
+
+Binned by the bf16 model's own top-1 probability, the W4A16 GPTQ checkpoint agrees
+~100% where bf16 is confident (p₁ > 0.5) and the disagreements sit in the near-tie
+buckets (p₁ < 0.15: 45% of semantic positions), where a 0.1-nat shift swaps two
+candidates the sampler was already choosing between at temperature 1.0. A
+per-layer sensitivity pass found no dominant layer (all 28 within ±0.008 nats of
+the ±0.005 noise), so no layer is kept in bf16. W4A4 is slower than W4A16 at
+batch 1 on this hardware as well as less accurate, so it was not published.
+
+## How to use
+
+The official pipeline derives its own bf16 AR checkpoint and does not (yet)
+accept an alternative. A 5-line patch to `yue2/fast.py` adds a
+`YUE2_AR_CHECKPOINT` environment override:
+[`patch_fast.py`](patch_fast.py) (idempotent; also in this repo).
+
+```bash
+pip install "yue2_infer[fast]"        # the wheel from m-a-p/YuE2-3B
+python patch_fast.py "$(python -c 'import yue2.fast,os;print(os.path.abspath(yue2.fast.__file__))')"
+hf download DKmode22/YuE2-3B-AR-NVFP4 --local-dir ./yue2-ar-nvfp4
+export YUE2_AR_CHECKPOINT=$PWD/yue2-ar-nvfp4
+python - <<'EOF'
+from yue2 import YuE2Pipeline
+pipe = YuE2Pipeline.from_pretrained("m-a-p/YuE2-3B", backend="vllm", memory_budget_gib=16)
+song = pipe(style="Indie folk, warm acoustic guitar, male vocal", lyrics="[Verse]\n...")
+song.save("song.flac")
+EOF
+```
+
+Notes:
+- The vLLM path is used only when `cfg_scale == 1` (the default for
+  `cot=full|melody`) and `quantization="none"`. `cot=off` uses guidance 1.01 and
+  the pipeline silently falls back to its torch backend, where this checkpoint
+  is not used.
+- On DGX Spark (sm_121) the `vllm==0.19.0` PyPI wheel the package pins has no
+  usable NVFP4 kernel path (the W4A4 arm fails to initialise; the W4A16 arm
+  runs but returns wrong logits). Use a vLLM build with SM121 support
+  (we used 0.27.1) and widen the version assert in `yue2/fast.py`.
+- Everything else about the pipeline is unchanged: `max_num_seqs=1`, the
+  windowed repetition-penalty logits processor, 24,576 context.
+
+## Licence and attribution
+
+The base model is released by the M-A-P team under **CC-BY-NC-4.0**; this
+derivative carries the same licence and is for non-commercial use. All credit
+for the model goes to the YuE2 authors — please cite their work
+(`arxiv:2503.08638`). Quantisation, calibration corpus and evaluation by
+DKmode22 on DGX Spark (GB10) hardware, 2026-09-13.

config.json

@@ -0,0 +1,100 @@
+{
+  "architectures": [
+    "Qwen3ForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "bos_token_id": null,
+  "dtype": "bfloat16",
+  "eos_token_id": null,
+  "head_dim": 128,
+  "hidden_act": "silu",
+  "hidden_size": 2048,
+  "initializer_range": 0.02,
+  "intermediate_size": 6144,
+  "layer_types": [
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention"
+  ],
+  "max_position_embeddings": 24576,
+  "max_window_layers": 0,
+  "model_type": "qwen3",
+  "num_attention_heads": 16,
+  "num_hidden_layers": 28,
+  "num_key_value_heads": 8,
+  "pad_token_id": null,
+  "quantization_config": {
+    "config_groups": {
+      "group_0": {
+        "format": "nvfp4-pack-quantized",
+        "input_activations": null,
+        "output_activations": null,
+        "targets": [
+          "Linear"
+        ],
+        "weights": {
+          "actorder": "static",
+          "block_structure": null,
+          "dynamic": false,
+          "group_size": 16,
+          "num_bits": 4,
+          "observer": "memoryless_minmax",
+          "observer_kwargs": {},
+          "scale_dtype": "torch.float8_e4m3fn",
+          "strategy": "tensor_group",
+          "symmetric": true,
+          "type": "float",
+          "zp_dtype": null
+        }
+      }
+    },
+    "format": "nvfp4-pack-quantized",
+    "global_compression_ratio": null,
+    "ignore": [
+      "lm_head"
+    ],
+    "kv_cache_scheme": null,
+    "quant_method": "compressed-tensors",
+    "quantization_status": "compressed",
+    "sparsity_config": {},
+    "transform_config": {},
+    "version": "0.18.0"
+  },
+  "rms_norm_eps": 1e-06,
+  "rope_parameters": {
+    "rope_theta": 1000000,
+    "rope_type": "default"
+  },
+  "sliding_window": null,
+  "tie_word_embeddings": false,
+  "transformers_version": "5.14.1",
+  "use_cache": true,
+  "use_sliding_window": false,
+  "vocab_size": 184704
+}

generation_config.json

@@ -0,0 +1,4 @@
+{
+  "_from_model_config": true,
+  "transformers_version": "5.14.1"
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:61ce7cfaf6c4d5807d254961df24ae78e2e813eea9d921e939c1cd89efd97674
+size 2306149080

patch_fast.py

@@ -0,0 +1,57 @@
+#!/usr/bin/env python3
+"""Idempotent patch for the installed yue2_infer `fast.py` (vLLM backend).
+
+Adds ONE lever: if the env var YUE2_AR_CHECKPOINT names a directory, the vLLM
+worker serves THAT checkpoint instead of the bf16 AR checkpoint it derives from
+model.safetensors. That is how an NVFP4 (compressed-tensors) requant of the
+derived Qwen3-shaped AR checkpoint is put on the real generation path. Nothing
+else changes (dtype, KV sizing from config.json, logits processor, max_num_seqs=1).
+
+Usage: patch_fast.py <site-packages>/yue2/fast.py [--check]
+"""
+from __future__ import annotations
+
+import re
+import sys
+
+MARK = "# s5-patch: YUE2_AR_CHECKPOINT override (services/yue2-nvfp4/patch_fast.py)"
+OLD = 'derived = derive_ar_checkpoint(setup["model_dir"])'
+NEW = (
+    MARK + "\n"
+    '    _override = os.environ.get("YUE2_AR_CHECKPOINT")\n'
+    '    derived = Path(_override) if _override else derive_ar_checkpoint(setup["model_dir"])\n'
+    '    if _override and not (derived / "config.json").exists():\n'
+    '        raise FileNotFoundError(f"YUE2_AR_CHECKPOINT has no config.json: {derived}")\n'
+    '    print(f"s5-patch: AR checkpoint = {derived} (override={bool(_override)})", file=sys.stderr, flush=True)'
+)
+
+
+def main() -> int:
+    path = sys.argv[1]
+    check = "--check" in sys.argv
+    src = open(path).read()
+    if MARK in src:
+        print("already patched")
+        return 0
+    if check:
+        print("NOT patched")
+        return 1
+    if src.count(OLD) != 1:
+        print(f"expected exactly one occurrence of {OLD!r}, found {src.count(OLD)}")
+        return 2
+    # the derive call is indented 4 spaces inside _worker_main
+    new_src = re.sub(r"^(\s+)" + re.escape(OLD) + r"$",
+                     lambda m: m.group(1) + NEW.replace("\n    ", "\n" + m.group(1)), src, count=1, flags=re.M)
+    if new_src == src:
+        print("substitution failed")
+        return 3
+    if "from pathlib import Path" not in new_src and "import Path" not in new_src:
+        print("fast.py has no Path import; refusing")
+        return 4
+    open(path, "w").write(new_src)
+    print("patched", path)
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

quant-manifest.json

@@ -0,0 +1,22 @@
+{
+ "src": "/opt/yue2/cache/yue2-ar/39f337a6c2354ed0063701cdd090b0eab7461450451ec02c3fec32f1ec125777",
+ "corpus": "corpus/bf16.jsonl",
+ "scheme": "NVFP4A16",
+ "awq": false,
+ "method": "gptq",
+ "ignore": [
+  "lm_head"
+ ],
+ "holdout": [
+  "r28",
+  "r29",
+  "r30",
+  "r31"
+ ],
+ "max_seq_len": 8192,
+ "num_samples": 58,
+ "tokens": 200116,
+ "codec_share": 0.514191768774111,
+ "seed": 0,
+ "started": "2026-09-13T01:21:33Z"
+}

recipe.yaml

@@ -0,0 +1,12 @@
+default_stage:
+  default_modifiers:
+    GPTQModifier:
+      targets: [Linear]
+      ignore: [lm_head]
+      scheme: NVFP4A16
+      bypass_divisibility_checks: false
+      requires_calibration_data: true
+      block_size: 128
+      dampening_frac: 0.01
+      actorder: static
+      offload_hessians: false