[MoE] MoE Calibration with calibrate_all_experts

[kylesayrs]

Coauthored with @dichn!

Purpose

• Add support for calibrate_all_experts option, which sends all tokens to all experts, but still produces the same outputs as if tokens had been gated

Changes

• Modify model definitions such that, in the case of calibrate_all_experts=True token gating occurs after passing tokens to experts, rather than before

# `calibrate_all_experts=True` by default
model = replace_modules_for_calibration(model, calibrate_all_experts=True)

Testing

• Added correctness tests for new model definitions which checks that outputs are exactly the same
• Added hook tests to make sure all experts are being sent tokens

[github-actions]

👋 Hi! Thank you for contributing to llm-compressor. Please add the ready label when the PR is ready for review.

Note: This is required to complete the testing suite, please only add the label once the PR is code complete and local testing has been performed.

[fynnsu]

Left a comment below. I also agree with @brian-dellabetta's point that this could maybe be simplified by patching self.top_k temporarily.

[dsikka]

I did not get a chance to run through these as of yet but it would be good to run through nvfp4 for llama4 and qwen3 and validating performance on the b200 before landing this, if anybody has bandwidth to run these

[kylesayrs]

Running those examples now

[dsikka]

FYI - Produced the following Qwen Model: nm-testing/Qwen3-30B-A3B-NVFP4-0912

With the following set-up:

lm_eval \
  --model vllm \
  --model_args pretrained="Qwen/Qwen3-30B-A3B",dtype=auto,max_model_len=4096,add_bos_token=True\
  --tasks gsm8k \
  --batch_size auto
    

NVFP4:

|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.8734|±  |0.0092|
|     |       |strict-match    |     5|exact_match|↑  |0.8719|±  |0.0092|

Dense:

|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.8544|±  |0.0097|
|     |       |strict-match    |     5|exact_match|↑  |0.8916|±  |0.0086|

The only thing I noticed was speed but that may also be because of beaker just being slow

[dsikka]

Similarly, for Llama4 NVFP4

lm_eval \
  --model vllm \
  --model_args pretrained="nm-testing/Llama-4-Scout-17B-16E-Instruct-NVFP4-0913",dtype=auto,max_model_len=4096,tensor_parallel_size=2,enable_chunked_prefill=True,enforce_eager=True \
  --tasks gsm8k_llama \
  --apply_chat_template \
  --fewshot_as_multiturn \
  --batch_size auto

[dsikka]

LGTM.
Two small comments

[kylesayrs]

Thanks a ton @dsikka