YuE2-3B-int8-ar

m-a-p/YuE2-3B with the autoregressive (AR) transformer's linear layers stored as per-channel int8. Everything else is the original bf16. Built for Apple Silicon, where torch._weight_int8pack_mm makes int8 decode faster than bf16. It loads on CUDA too, but torch 2.10 has no packed int8 kernel there, so CUDA runs a dequantised GEMM and is slower than the bf16 original. Use m-a-p/YuE2-3B on CUDA.

  • 5.85 GB on disk, down from 7.26 GB
  • 196 linears quantised: self_attn.{q,k,v,o}_proj and mlp.{gate,up,down}_proj in all 28 layers
  • Unchanged bf16: embed_tokens, lm_head, every NAR (acoustic) layer, norms, vae2llm, llm2vae, time_embedder
  • Symmetric per-output-channel scales stored as bf16 alongside each int8 weight (<name>.weight int8, <name>.weight_scale bf16)

Measurements

Apple M5 Max, 128 GB, torch 2.10, YuE2 runtime 0.1.6 with the bundled patch.

Quality against the bf16 model, 512 token prefix, next-token logits:

  • KL(bf16 || int8) below 1e-5
  • Top-5 tokens identical, max logit difference 0.09 (logit std 1.0)
  • Loading this checkpoint gives bit-identical logits to quantising the bf16 model at runtime

Speed, single-token decode at 8000 tokens of context, including the patch's bmm attention:

  • bf16 stock: 51 ms per token
  • int8 with patch: 20 ms per token

Speed is identical to running the bf16 model with quantization="int8" in the patched runtime. What this checkpoint saves is 1.4 GB of download and the quantisation pass at load time.

Loading

Not loadable with plain transformers or the stock yue2-infer package; the int8 tensors need the custom Int8Linear module. Apply yue2-mps-int8.patch to upstream YuE at commit 88da114a67df892af0329472073b96a5ef700b93:

git clone https://github.com/multimodal-art-projection/YuE.git
git -C YuE checkout 88da114a67df892af0329472073b96a5ef700b93
git -C YuE apply ../yue2-mps-int8.patch
python -m pip install -e ./YuE
from yue2 import YuE2Pipeline

with YuE2Pipeline.from_pretrained("smcleod/YuE2-3B-int8-ar", device="auto") as pipe:
    song = pipe(style="...", lyrics="...", cot="full")
    song.save_artifacts("outputs/my-song")

The pipeline reads quantization_config.method == "yue2-int8-ar" from config.json and takes the int8 loader. Only quantization="auto" or "int8" is accepted for this checkpoint; "none" and "fp8" raise.

Differences from runtime quantisation

The patched runtime quantises AR linears on load and restores the bf16 originals before the NAR stage, so the NAR's prefix prefill runs on bf16 AR weights. This checkpoint ships no bf16 originals, so that prefill runs through the dequantised int8 weights. Generated songs from both paths were checked to complete without error; a listening comparison is left to the user.

Reproducing

scripts/ holds the files used:

  • export_int8.py - loads the bf16 model, quantises, writes this directory
  • check_int8.py - logit KL, top-k agreement and decode timing, runtime quantisation vs bf16
  • check_int8_checkpoint.py - confirms the saved checkpoint matches runtime quantisation
  • bench_int8.py - _weight_int8pack_mm vs bf16 linear at YuE2 shapes
python scripts/export_int8.py --source m-a-p/YuE2-3B --output YuE2-3B-int8-ar

The scripts import the patched yue2 package. check_int8.py and check_int8_checkpoint.py take an earlier run directory for a prefix.npy token sample. The source weights sha256 is recorded under quantization_config.source_weights_sha256.

Licence

Derivative of m-a-p/YuE2-3B, released under the same CC BY-NC 4.0 terms: attribution required, non-commercial use only. LICENSE, THIRD_PARTY_NOTICES.md and licenses/ are copied from the source model. The patch and scripts are Apache 2.0 (SPDX headers in each script), matching the upstream YuE code.

Downloads last month
64
Safetensors
Model size
4B params
Tensor type
BF16
·
I8
·
Inference Providers NEW
This model isn't deployed by any Inference Provider. 🙋 Ask for provider support

Model tree for smcleod/YuE2-3B-int8-ar

Base model

m-a-p/YuE2-3B
Quantized
(9)
this model