modelopt-fast-moe cuts AWQ calibration time for MoE models at NVFP4

The first time I asked nvidia-modelopt to quantize a Mixture-of-Experts model, to shrink its weights to 4-bit NVFP4, the job told me it would take about 50 hours. The model was SuperGemma4 26B and the machine was the Spark, a GB10 with 97 GB of VRAM, the memory the GPU works in. GPU utilization sat at 25 to 30 percent the whole time. I let it run for 18 hours before I killed it. The quantizer was fine. The way I was feeding it was not.

modelopt-fast-moe is the small Python package that came out of that failure. It replaces the calibration loop most modelopt examples use, one sample per forward pass, with one that batches samples, probes free VRAM for the largest safe batch size, and halves the batch on an out-of-memory error. The AWQ algorithm is untouched. Version 0.1.0 is Apache-2.0 and beta. It does not make AWQ aware of experts; it feeds AWQ faster. The expert plugin is a separate piece, described below.

What breaks when you quantize experts naively

A Mixture-of-Experts model, MoE for short, replaces the feed-forward block in each transformer layer with many smaller experts and a router that sends each token to only a few of them. That sparsity makes a 26B MoE cheap to run and, it turns out, slow to calibrate.

AWQ, activation-aware weight quantization, has to see real activations on sample text before it rounds the weights, so it can protect the channels that matter most; that pass is calibration. NVIDIA's modelopt, the TensorRT Model Optimizer, gathers them by calling your forward_loop twice: a cache phase, then an alpha-search phase that tries 11 candidate scaling strengths per linear layer and keeps the one with the lowest loss. NVFP4 is NVIDIA's 4-bit floating-point format with a scale for every small block of values. The usual loop:

Code, Pythondef forward_loop(model):
    for ids in calib_data:
        model(input_ids=ids)

One sample per pass. On an MoE that is a trap: each forward dispatches hundreds of tiny per-expert matrix multiplications, GEMMs, from Python, and every CUDA launch costs roughly 10 microseconds of CPU time however small the matrix. Tokens per expert come to about sequence length times top_k over the number of experts, often fewer than 128, so the GPU finishes each little GEMM before Python can hand it the next. That is the 25 to 30 percent. It looks like AWQ being slow; it is the loop.

What the tool does, and what it leaves alone

The fix is to batch: one Python dispatch per batch, and each per-expert GEMM gets more tokens. The statistic modelopt collects, a per-channel act_scale equal to mean(|activations|), is weighted by token count rather than batch count, so the README says the result is invariant to batch size. The hard part is picking a batch size that fits and surviving when it does not.

The probe runs on the first call, before modelopt allocates its alpha-search buffers: one forward at batch size 1, one at batch size 2, and the change in torch.cuda.max_memory_allocated gives a line, slope times batch size plus a fixed cost. The pick takes safety_fraction times total VRAM as the budget, 85 percent by default, subtracts the baseline and the fixed cost, divides by the slope times a phase2_overhead_factor of 2.0, and clamps the result between min_bs 1 and max_bs 16. The factor is pessimistic on purpose, because the search phase holds 11 alpha-loss tensors plus scaled inputs the probe never sees. On torch.cuda.OutOfMemoryError the loop frees the batch, halves the batch size, keeps the smaller value across phases, and retries; at min_bs it re-raises.

Code, shellpip install modelopt-fast-moe

Code, Pythonimport modelopt.torch.quantization as mtq
from modelopt_fast_moe import make_adaptive_calibrate_loop
forward_loop = make_adaptive_calibrate_loop(calib_data)
mtq.quantize(model, mtq.NVFP4_AWQ_FULL_CFG, forward_loop=forward_loop)

The only hard dependency is torch 2.3 or newer, on Linux with Python 3.10 or newer; you bring your own modelopt, 0.19 or newer for the optional extra.

What it leaves alone: it does not fuse the per-expert Python loop into grouped GEMMs, so the GPU still waits on Python. It does not touch the alpha search, where 11 alphas are still evaluated one after another for every linear layer. It does not remove padding waste, around 30 percent on variable-length data. You still pad calib_data to one length because torch.cat demands it, pass a forward_fn(model, batch) if your model takes more than input_ids, and write the expert plugin if modelopt lacks one. The roadmap, none of it shipped, is length bucketing, a torch.compile experiment, a fused FastMoEPlugin built on one torch.bmm over expert-grouped tokens, and a batched alpha search.

The numbers, with their context

On a mini MoE, 32 experts with hidden size 512, intermediate size 256 and top_k 4, running full NVFP4_AWQ_FULL_CFG calibration over 16 samples of 512 tokens on the GB10, batching alone does this:

The v1 run of SuperGemma4 26B used 4096 calibration samples of 4096 tokens at batch size 1 and was projected at about 50 hours; I killed it at 18. The v2.5 run, with this package and the tuning described below, used 512 samples of 4096 tokens at an adaptive batch size of 3, and the README estimates about 1.7 hours. Both are projections, not stopwatch readings. The gain is about 30x end to end, and most of it is not batching: dropping the sample count from 4096 to 512, modelopt's own default, is 8x; coarsening alpha_step from 0.1 to 0.2 is 1.8x on the search phase; adaptive batching contributed 2.5x, capped at batch size 3 by VRAM. Even at batch size 3 the GPU sat at 47 percent, because the plugin still walks 128 experts in a Python loop. Quality after fewer samples or a coarser alpha step is unmeasured; I claim nothing about it.

The Gemma expert plugin in the example script

The package batches; the expert handling lives in examples/quantize_supergemma4.py. Gemma4 stores its experts as fused three-dimensional tensors, one [E, 2I, H] for the gate and up projections and one [E, H, I] for the down projection, and modelopt's per-linear AWQ hooks only attach to real nn.Linear modules. So the script registers a QuantModule for Gemma4TextExperts that splits each fused tensor into per-expert gate, up and down projections. It leaves the vision parts unquantized, builds calibration data from cnn_dailymail 3.0.0 samples over 200 characters padded to one length, sets PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True, and runs awq_full with alpha_step 0.2 at 4096 tokens as a long-context hedge.

The export trap: modelopt strips .moe. from the expert tensor keys during export, and vLLM expects it. The script renames .experts.N. to .moe.experts.N., rewrites everything into a single model.safetensors, and deletes the shards and index.

Code, shellpython3 examples/quantize_supergemma4.py

Edit the MODEL_PATH and SAVE_PATH placeholders first. The Gemma 4 26B A4B MoE NVFP4 build for the Spark is the same family and format as the origin job; read its card rather than assume this script produced it. The Qwen 3.6 35B A3B NVFP4 release is another MoE at NVFP4 in the library; the serving side lives in the AEON vLLM Ultimate article for the Spark.

Sources

• modelopt-fast-moe on GitHub
• modelopt-fast-moe issue tracker
• NVIDIA TensorRT Model Optimizer