Support force tokens to % of total experts during calibration #910
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -906,6 +906,7 @@ def quantize_main( | |
| model_type, | ||
| QUANT_CFG_CHOICES, | ||
| KV_QUANT_CFG_CHOICES, | ||
| args.moe_calib_experts_ratio, | ||
| ) | ||
|
|
||
| # Exclude MTP layers from quantization if detected (e.g., GLM-4.7's layer 92) | ||
|
|
@@ -1126,8 +1127,21 @@ def parse_args() -> argparse.Namespace: | |
| "(sensitivity scores, costs, etc.). Only used when auto_quantize_bits is specified." | ||
| ), | ||
| ) | ||
| parser.add_argument( | ||
|
Contributor
There was a problem hiding this comment. Should we add something similar to mcore PTQ as well? cc @realAsma @jenchen13 @ChenhanYu |
||
| "--moe_calib_experts_ratio", | ||
| type=float, | ||
| default=1.0, | ||
| help=( | ||
| "Fraction of experts to calibrate during forward pass (ratio in (0.0, 1.0]). " | ||
| "Only used for MOE models; used to reduce the number of experts calibrated during the forward pass." | ||
| "Does not impact non-MOE models." | ||
| ), | ||
| ) | ||
|
Comment on lines
+1130
to
+1139
Contributor
There was a problem hiding this comment. Default Since the default is Consider defaulting to parser.add_argument(
"--moe_calib_experts_ratio",
type=float,
- default=1.0 / 4,
+ default=None,
help=(
- "Percentage of experts to calibrate during forward pass. Only used for MOE models. "
- "This is used to reduce the number of experts to calibrate during forward pass. "
+ "Ratio of experts to calibrate during forward pass (0, 1]. Only used for MOE models. "
+ "Default behavior routes to all experts if not specified. "
+ "Example: 0.25 calibrates 25%% of experts. "
),
)🤖 Prompt for AI Agents |
||
|
|
||
| args = parser.parse_args() | ||
| if not (0.0 < args.moe_calib_experts_ratio <= 1.0): | ||
| parser.error("--moe_calib_experts_ratio must be in the range (0.0, 1.0].") | ||
| return args | ||
|
|
||
|
|
||
| def main(args: argparse.Namespace): | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -48,7 +48,7 @@ def save_expert_token_count_table(model: nn.Module, output_dir: str | Path | Non | |
| "th, td { border: 1px solid #ccc; padding: 4px 8px; text-align: right; }", | ||
| "th { background: #f0f0f0; }", | ||
| "</style></head><body>", | ||
| "<h2>Expert Calib Token Counts (per MoE layer)</h2>", | ||
| "<table><tr><th>Layer/Expert</th>", | ||
| ] | ||
| html_parts.extend(f"<th>{i}</th>" for i in range(num_experts)) | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -1091,6 +1091,16 @@ class QuantizeAlgorithmConfig(ModeloptBaseConfig): | |
| title="This field specifies the name of the calibration algorithm. If None, no calibration is performed.", | ||
| ) | ||
|
|
||
| moe_calib_experts_ratio: float | None = ModeloptField( | ||
|
Contributor
There was a problem hiding this comment. Yea this is good idea to put this here.
Contributor
There was a problem hiding this comment. Is it necessary to put
Contributor
There was a problem hiding this comment. for example if we think
Collaborator
Author
There was a problem hiding this comment. @jenchen13 this is a good question: why would we not include KV cache clamping as part of the quant config too or other non-essential PTQ features? If we've decided to make this tunable, we probably also want to add it here.
Contributor
There was a problem hiding this comment. +1 on making the KVCache clamping here - this will help the QAD to be aligned with export path |
||
| default=None, | ||
| title="% of experts to calibrate during forward pass.", | ||
| description=( | ||
| "If specified, we force forward tokens to % of experts during the calibration" | ||
|
Comment on lines
+1094
to
+1098
There was a problem hiding this comment. This new config field doesn't enforce bounds. Since valid values are in (0, 1], add pydantic constraints (e.g., |
||
| " pass. This forward is for calibration purpose only and will not affect the" | ||
| " actual inference." | ||
| ), | ||
| ) | ||
|
|
||
|
|
||
| class MaxCalibConfig(QuantizeAlgorithmConfig): | ||
| """The config for max calibration algorithm. | ||
|
|
||
| Original file line number | Diff line number | Diff line change | ||||||||||||||||||||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
|
|
@@ -225,6 +225,15 @@ def wrapped_calib_func( | |||||||||||||||||||||||||||
| # For backward compatibility | ||||||||||||||||||||||||||||
| kwargs["algorithm"] = method | ||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||
| moe_calib_experts_ratio = kwargs.pop("moe_calib_experts_ratio", None) | ||||||||||||||||||||||||||||
| if moe_calib_experts_ratio is not None: | ||||||||||||||||||||||||||||
|
There was a problem hiding this comment.
Suggested change
|
||||||||||||||||||||||||||||
| assert ( | ||||||||||||||||||||||||||||
| isinstance(moe_calib_experts_ratio, (int, float)) and 0 < moe_calib_experts_ratio <= 1 | ||||||||||||||||||||||||||||
| ), f"Invalid moe_calib_experts_ratio {moe_calib_experts_ratio!r}" | ||||||||||||||||||||||||||||
| for module in model.modules(): | ||||||||||||||||||||||||||||
| if hasattr(module, "_moe_calib_experts_ratio"): | ||||||||||||||||||||||||||||
| module._moe_calib_experts_ratio = moe_calib_experts_ratio | ||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||
| if func is not None: | ||||||||||||||||||||||||||||
| # Call the function with forward_loop as a separate argument | ||||||||||||||||||||||||||||
| func(model, forward_loop=forward_loop, **kwargs) | ||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||
| Original file line number | Diff line number | Diff line change | ||||
|---|---|---|---|---|---|---|
|
|
@@ -458,8 +458,13 @@ def _setup(self): | |||||
| elif hasattr(self, "experts") and hasattr(self.experts, "num_experts"): | ||||||
| num_experts = self.experts.num_experts | ||||||
|
|
||||||
| self.register_buffer( | ||||||
| "expert_token_count", | ||||||
| torch.zeros(num_experts, dtype=torch.long, device=next(self.parameters()).device), | ||||||
| persistent=False, | ||||||
| ) | ||||||
| self._count_expert_tokens = False | ||||||
| self._moe_calib_experts_ratio = None | ||||||
|
|
||||||
| if num_experts == 0: | ||||||
| warnings.warn( | ||||||
|
|
@@ -483,36 +488,48 @@ def _gate_forward_hook(self, module, input, output): | |||||
| logits = output if not isinstance(output, tuple) else output[0] | ||||||
| top_k = self.gate.top_k if hasattr(self.gate, "top_k") else self.top_k | ||||||
| _, indices = torch.topk(logits.float(), top_k, dim=-1) | ||||||
| counts = torch.bincount(indices.reshape(-1), minlength=self.expert_token_count.shape[0]) | ||||||
| self.expert_token_count += counts.to(self.expert_token_count.device) | ||||||
|
|
||||||
| def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: | ||||||
| is_calib = any(getattr(m, "_if_calib", False) for m in self.experts.modules()) | ||||||
| self._count_expert_tokens = is_calib | ||||||
| if is_calib and self._moe_calib_experts_ratio: | ||||||
| self._count_expert_tokens = True | ||||||
| assert 0 < self._moe_calib_experts_ratio <= 1, ( | ||||||
| "moe_calib_experts_ratio must be between 0 and 1" | ||||||
| ) | ||||||
|
Comment on lines
+497
to
+501
There was a problem hiding this comment. Input validation uses |
||||||
| # If any of the experts are in calibration mode, we will forward all tokens to | ||||||
| # self._moe_calib_experts_ratio % of the experts to improve the calibration coverage. | ||||||
| # This is used only for calibration, we need to re-calculate the actual outputs again using | ||||||
| # the original top_k | ||||||
| if TRANSFORMERS_VERSION_GE_5_0: | ||||||
| assert hasattr(self, "gate") and hasattr(self.gate, "top_k") | ||||||
| original_top_k = self.gate.top_k | ||||||
| self.gate.top_k = max( | ||||||
| original_top_k, round(self.gate.num_experts * self._moe_calib_experts_ratio) | ||||||
| ) | ||||||
|
Comment on lines
508
to
+511
There was a problem hiding this comment.
|
||||||
| super().forward(hidden_states) | ||||||
| self.gate.top_k = original_top_k | ||||||
| else: | ||||||
| # Path for transformers < 5.0 | ||||||
| original_top_k = self.top_k | ||||||
| if hasattr(self, "num_experts"): | ||||||
| self.top_k = max( | ||||||
| original_top_k, round(self.num_experts * self._moe_calib_experts_ratio) | ||||||
| ) | ||||||
|
Comment on lines
516
to
+520
There was a problem hiding this comment. Same rounding issue here:
Contributor
There was a problem hiding this comment. +1 ^ seems like a risky area for silent failure if rounding goes wrong |
||||||
| elif hasattr(self, "experts"): | ||||||
| self.top_k = max( | ||||||
| original_top_k, | ||||||
| round(self.experts.num_experts * self._moe_calib_experts_ratio), | ||||||
| ) | ||||||
| else: | ||||||
| raise ValueError(f"Could not find num_experts in module {self}") | ||||||
| super().forward(hidden_states) | ||||||
| self.top_k = original_top_k | ||||||
| self._count_expert_tokens = False | ||||||
| else: | ||||||
| self._count_expert_tokens = True | ||||||
|
There was a problem hiding this comment. The
Suggested change
|
||||||
| output = super().forward(hidden_states) | ||||||
| self._count_expert_tokens = False | ||||||
| return output | ||||||
|
Comment on lines
494
to
535
Contributor
There was a problem hiding this comment. 🧩 Analysis chain🏁 Script executed: #!/bin/bash
# First, let's see recent commits to understand context
git log --oneline -10 -- modelopt/torch/quantization/plugins/huggingface.pyRepository: NVIDIA/Model-Optimizer Length of output: 87 🏁 Script executed: #!/bin/bash
# Get a diff to see what changed in recent commits
# Look at the last few commits for this file
git log -p -2 -- modelopt/torch/quantization/plugins/huggingface.py | head -300Repository: NVIDIA/Model-Optimizer Length of output: 48 🏁 Script executed: #!/bin/bash
# Let's also check the full forward method and surrounding context
cat -n modelopt/torch/quantization/plugins/huggingface.py | sed -n '440,550p'Repository: NVIDIA/Model-Optimizer Length of output: 6176 🏁 Script executed: #!/bin/bash
# Search for any tests related to this forward method or _moe_calib_experts_ratio
find . -name "*.py" -type f | xargs grep -l "_moe_calib_experts_ratio" 2>/dev/nullRepository: NVIDIA/Model-Optimizer Length of output: 156 🏁 Script executed: #!/bin/bash
# Look for documentation or comments about the expected behavior
grep -r "all experts" --include="*.py" modelopt/torch/quantization/plugins/ 2>/dev/null | head -20Repository: NVIDIA/Model-Optimizer Length of output: 672 🏁 Script executed: #!/bin/bash
# Let's check what happens in each scenario by examining the control flow
# Case analysis of the forward method:
# When is_calib=True and ratio is set (line 493 condition is True):
# - Line 528: _count_expert_tokens = False (before final forward)
# When is_calib=True and ratio is None (line 493 condition is False, goes to else):
# - Line 530: _count_expert_tokens = True (before final forward)
# When is_calib=False (line 493 condition is False, goes to else):
# - Line 530: _count_expert_tokens = True (before final forward) <-- PROBLEM!
# Check if there's any conditional that prevents token counting during inference
cat -n modelopt/torch/quantization/plugins/huggingface.py | sed -n '475,495p'Repository: NVIDIA/Model-Optimizer Length of output: 1385 🏁 Script executed: #!/bin/bash
# Check the gate_forward_hook to understand what happens when _count_expert_tokens is True
cat -n modelopt/torch/quantization/plugins/huggingface.py | sed -n '475,489p'Repository: NVIDIA/Model-Optimizer Length of output: 982 🏁 Script executed: #!/bin/bash
# Search for where _moe_calib_experts_ratio might be set to understand intended usage
grep -n "_moe_calib_experts_ratio" modelopt/torch/quantization/plugins/huggingface.pyRepository: NVIDIA/Model-Optimizer Length of output: 699 🏁 Script executed: #!/by
# Check if there are any tests that show expected behavior
find . -name "*test*.py" -type f | xargs grep -l "QuantSparseMoe\|_moe_calib_experts_ratio" 2>/dev/null | head -5Repository: NVIDIA/Model-Optimizer Length of output: 124 🏁 Script executed: #!/bin/bash
# Examine the test file to understand expected behavior
cat modelopt/torch/quantization/plugins/test_sparse_moe.pyRepository: NVIDIA/Model-Optimizer Length of output: 152 🏁 Script executed: #!/bin/bash
# Find the correct path for the test file
find . -name "test_sparse_moe.py" -type fRepository: NVIDIA/Model-Optimizer Length of output: 124 🏁 Script executed: #!/bin/bash
# If not found, search for any test file containing QuantSparseMoe or sparse moe tests
find . -name "*test*.py" -type f -exec grep -l "QuantSparseMoe\|sparse.*moe" {} \; 2>/dev/null | head -5Repository: NVIDIA/Model-Optimizer Length of output: 124 🏁 Script executed: #!/bin/bash
# Read the test file
cat ./tests/unit/torch/quantization/plugins/test_sparse_moe.pyRepository: NVIDIA/Model-Optimizer Length of output: 12565 🏁 Script executed: #!/bin/bash
# Check where _moe_calib_experts_ratio might be set (outside of _setup)
grep -rn "_moe_calib_experts_ratio" modelopt/torch/quantization/ --include="*.py" | grep -v "test_sparse"Repository: NVIDIA/Model-Optimizer Length of output: 1321 🏁 Script executed: #!/bin/bash
# Check the quantize function and INT8_DEFAULT_CFG to see if they set the ratio
grep -rn "INT8_DEFAULT_CFG\|_moe_calib_experts_ratio" modelopt/torch/quantization/ --include="*.py" | head -30Repository: NVIDIA/Model-Optimizer Length of output: 1959 🏁 Script executed: #!/bin/bash
# Look for mode.py to see how _moe_calib_experts_ratio is used
cat modelopt/torch/quantization/mode.py | head -100Repository: NVIDIA/Model-Optimizer Length of output: 3247 🏁 Script executed: #!/bin/bash
# Look at the full context around line 222 in mode.py to see the function signature
cat -n modelopt/torch/quantization/mode.py | sed -n '210,235p'Repository: NVIDIA/Model-Optimizer Length of output: 1265 🏁 Script executed: #!/bin/bash
# Search for moe_calib_experts_ratio parameter definition and default value
grep -B5 -A5 "moe_calib_experts_ratio" modelopt/torch/quantization/mode.py | head -40Repository: NVIDIA/Model-Optimizer Length of output: 665 🏁 Script executed: #!/bin/bash
# Check if there's a function signature or parameter with default
grep -rn "def.*moe_calib_experts_ratio\|moe_calib_experts_ratio\s*=" modelopt/torch/quantization/ --include="*.py" | head -20Repository: NVIDIA/Model-Optimizer Length of output: 393 Clarify whether all-experts calibration should be the default during quantization. The class docstring promises "During calibration, we forward all tokens to all experts so that all experts see sufficient tokens to calibrate" (line 445), but this behavior only activates when Additionally, the else block at lines 529-530 enables token counting for both inference ( Either set a default ratio (e.g., 1.0 for all experts) when entering calibration mode, or update the docstring to clarify that expanded-expert forwarding requires explicit configuration. 🤖 Prompt for AI Agents |
||||||
|
|
||||||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
if moe_calib_experts_ratio:uses truthiness and will silently ignore0.0(and will treat negative values as enabled). Preferif moe_calib_experts_ratio is not None:and validate0 < ratio <= 1so invalid values fail fast with a clear error.