From a4fde491cc6c3c7c747c7c946416ae09d2c6fa5d Mon Sep 17 00:00:00 2001 From: Chenjie Luo <108829653+cjluo-nv@users.noreply.github.com> Date: Tue, 3 Mar 2026 13:30:35 -0800 Subject: [PATCH] Update MOE block detection logic and enable in huggingface_script.sh (#962) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What does this PR do? Type of change: Bug fix Add moe expert calib ratio in huggingface_script.sh Also fix minimax2.5 MOE detection which does not follow other HF MOE layer convention ### Usage scripts/huggingface_example.sh --model --quant nvfp4 --moe_calib_experts_ratio 1.0 --trust_remote_code ### Before your PR is "*Ready for review*" Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md) and your commits are signed (`git commit -s -S`). Make sure you read and follow the [Security Best Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors) (e.g. avoiding hardcoded `trust_remote_code=True`, using `torch.load(..., weights_only=True)`, avoiding `pickle`, etc.). - Is this change backward compatible?: ✅ / ❌ / N/A - If you copied code from any other source, did you follow IP policy in [CONTRIBUTING.md](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md#-copying-code-from-other-sources)?: ✅ / ❌ / N/A - Did you write any new necessary tests?: ✅ / ❌ / N/A - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ✅ / ❌ / N/A ### Additional Information ## Summary by CodeRabbit * **New Features** * Configure MOE calibration experts ratio for quantization via an environment/option, enabling finer control over calibration. * **Bug Fixes** * Improved detection of sparse MOE blocks to handle varying expert/topology layouts, inferring expert counts when needed for more reliable processing. --------- Signed-off-by: Chenjie Luo --- examples/llm_ptq/scripts/huggingface_example.sh | 4 ++++ examples/llm_ptq/scripts/parser.sh | 4 +++- modelopt/torch/quantization/plugins/huggingface.py | 9 +++++++-- .../torch/quantization/plugins/test_sparse_moe.py | 12 ++++++++++-- 4 files changed, 24 insertions(+), 5 deletions(-) diff --git a/examples/llm_ptq/scripts/huggingface_example.sh b/examples/llm_ptq/scripts/huggingface_example.sh index a6455ed37..a74a1671a 100755 --- a/examples/llm_ptq/scripts/huggingface_example.sh +++ b/examples/llm_ptq/scripts/huggingface_example.sh @@ -139,6 +139,10 @@ if [ -n "$CALIB_SEQ" ]; then PTQ_ARGS+=" --calib_seq=$CALIB_SEQ " fi +if [ -n "$MOE_CALIB_EXPERTS_RATIO" ]; then + PTQ_ARGS+=" --moe_calib_experts_ratio=$MOE_CALIB_EXPERTS_RATIO " +fi + if ! $VERBOSE; then PTQ_ARGS+=" --no-verbose " fi diff --git a/examples/llm_ptq/scripts/parser.sh b/examples/llm_ptq/scripts/parser.sh index 8db2fe131..53397d5ee 100644 --- a/examples/llm_ptq/scripts/parser.sh +++ b/examples/llm_ptq/scripts/parser.sh @@ -36,7 +36,7 @@ parse_options() { USE_SEQ_DEVICE_MAP=false # Parse command-line options - ARGS=$(getopt -o "" -l "model:,quant:,kv_cache_quant:,tp:,pp:,sparsity:,awq_block_size:,calib:,calib_batch_size:,auto_quantize_bits:,output:,batch:,tasks:,lm_eval_tasks:,lm_eval_limit:,simple_eval_tasks:,trust_remote_code,use_seq_device_map,gpu_max_mem_percentage:,kv_cache_free_gpu_memory_fraction:,low_memory_mode,no-verbose,calib_dataset:,calib_seq:,auto_quantize_method:,auto_quantize_score_size:,auto_quantize_checkpoint:" -n "$0" -- "$@") + ARGS=$(getopt -o "" -l "model:,quant:,kv_cache_quant:,tp:,pp:,sparsity:,awq_block_size:,calib:,calib_batch_size:,auto_quantize_bits:,output:,batch:,tasks:,lm_eval_tasks:,lm_eval_limit:,simple_eval_tasks:,trust_remote_code,use_seq_device_map,gpu_max_mem_percentage:,kv_cache_free_gpu_memory_fraction:,low_memory_mode,no-verbose,calib_dataset:,calib_seq:,auto_quantize_method:,auto_quantize_score_size:,auto_quantize_checkpoint:,moe_calib_experts_ratio:" -n "$0" -- "$@") eval set -- "$ARGS" while true; do @@ -68,6 +68,7 @@ parse_options() { --auto_quantize_method ) AUTO_QUANTIZE_METHOD="$2"; shift 2;; --auto_quantize_score_size ) AUTO_QUANTIZE_SCORE_SIZE="$2"; shift 2;; --auto_quantize_checkpoint ) AUTO_QUANTIZE_CHECKPOINT="$2"; shift 2;; + --moe_calib_experts_ratio ) MOE_CALIB_EXPERTS_RATIO="$2"; shift 2;; -- ) shift; break ;; * ) break ;; esac @@ -156,5 +157,6 @@ parse_options() { echo "auto_quantize_method: $AUTO_QUANTIZE_METHOD" echo "auto_quantize_score_size: $AUTO_QUANTIZE_SCORE_SIZE" echo "auto_quantize_checkpoint: $AUTO_QUANTIZE_CHECKPOINT" + echo "moe_calib_experts_ratio: $MOE_CALIB_EXPERTS_RATIO" echo "=================" } diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py index 297b6d9a9..318646c39 100644 --- a/modelopt/torch/quantization/plugins/huggingface.py +++ b/modelopt/torch/quantization/plugins/huggingface.py @@ -1115,7 +1115,7 @@ def _is_sparse_moe_block(module): All HuggingFace MoE blocks (Mixtral, Qwen3Moe, Qwen2Moe, Qwen3Next, Llama4, MiniMax, etc.) share a common structural pattern: a ``gate`` (TopKRouter) sub-module with routing attributes - (``top_k`` and ``num_experts``), and an ``experts`` sub-module. + (``top_k``, some may have ``num_experts``), and an ``experts`` sub-module. This function detects that pattern instead of relying on class names, making it forward-compatible with new MoE architectures. Some MoE models (e.g. Glm4MoeMoE) have ``gate`` and ``experts`` but @@ -1134,7 +1134,12 @@ def _is_sparse_moe_block(module): return True # Fallback: top_k + num_experts on the block itself (older transformers, e.g. v4.x Qwen3Next) - return hasattr(module, "top_k") and hasattr(module, "num_experts") + if hasattr(module, "top_k"): + if not hasattr(module, "num_experts") and hasattr(module.experts, "__len__"): + module.num_experts = len(module.experts) + return hasattr(module, "num_experts") + + return False def register_sparse_moe_on_the_fly(model): diff --git a/tests/unit/torch/quantization/plugins/test_sparse_moe.py b/tests/unit/torch/quantization/plugins/test_sparse_moe.py index 6d548aa40..acd7e0b0e 100644 --- a/tests/unit/torch/quantization/plugins/test_sparse_moe.py +++ b/tests/unit/torch/quantization/plugins/test_sparse_moe.py @@ -133,11 +133,19 @@ class TestIsSparseBlock: module.gate = gate assert _is_sparse_moe_block(module) is False - def test_block_level_only_top_k_returns_false(self): - """Only top_k on block (no num_experts) -> fallback fails.""" + def test_block_level_top_k_infers_num_experts(self): + """top_k on block + experts with __len__ -> num_experts is inferred, returns True.""" module = nn.Module() module.experts = nn.ModuleList([nn.Linear(8, 8)]) module.top_k = 2 + assert _is_sparse_moe_block(module) is True + assert module.num_experts == 1 + + def test_block_level_top_k_no_len_returns_false(self): + """top_k on block but experts has no __len__ -> cannot infer num_experts, returns False.""" + module = nn.Module() + module.experts = nn.Module() + module.top_k = 2 assert _is_sparse_moe_block(module) is False def test_block_level_only_num_experts_returns_false(self):