Support per expert weight quantizer in TEGroupedMLP#1550
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughThis PR adds per-GEMM weight quantizers for TEGroupedLinear, extends calibration to grouped quantizer wrappers, and updates MoE tests to cover restore, compilation, and output behavior. It also refreshes the related config text and changelog entry. ChangesPer-expert weight quantization for TEGroupedLinear
Estimated code review effort: 4 (Complex) | ~45 minutes Possibly related PRs
🚥 Pre-merge checks | ✅ 5 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (5 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Warning
CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.
Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
modelopt/torch/quantization/plugins/transformer_engine.py (1)
147-166:⚠️ Potential issue | 🟠 Major | 🏗️ Heavy liftPer-expert restore is still unsafe when TP/EP changes.
This branch creates
weight_quantizer_{i}modules, butmodelopt_post_restore()still runs the base restore path that only knows aboutself.weight_quantizer. If a checkpoint is restored under different TP/EP, the per-expert_amaxtensors can keep the old shape and no code here reinitializes them. Please either block this mode on topology changes or add a per-expert reset/recalibration path during restore.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@modelopt/torch/quantization/plugins/transformer_engine.py` around lines 147 - 166, The per-expert weight-quantizer path (_per_expert_weight_quantizer and the created modules weight_quantizer_{i}) is unsafe on TP/EP topology change because modelopt_post_restore currently only runs the base path and does not reinitialize per-expert _amax buffers; update modelopt_post_restore to detect topology/shape mismatch and either (a) fail-fast by raising a clear error when _per_expert_weight_quantizer is enabled and the restored per-expert _amax shapes don't match the current weight shapes, or (b) perform a per-expert recalibration: iterate over each created module name weight_quantizer_{i} (use self.num_gemms to enumerate), reset/recreate their _amax buffers to the correct shape for the current self.weight{i} and call the quantizer's max_calibrate (or its equivalent reinit routine) so each per-expert quantizer is re-calibrated after restore; implement one of these branches in modelopt_post_restore to ensure shapes/amax are consistent after load.
🧹 Nitpick comments (1)
tests/gpu_megatron/torch/quantization/plugins/test_megatron.py (1)
738-750: ⚡ Quick winKeep these GPU assertions on-device.
tensor.item()andmath.isfinite(...)force host syncs on every worker here. Prefer tensor comparisons/reductions plustorch.isfinite, and only materialize Python scalars for failure messages or rank-0 logging.As per coding guidelines, "Performance note: avoid GPU syncs in hot paths/tests (don’t use
tensor.item()/float(tensor)/min(tensor))—use tensor ops instead."Also applies to: 800-813
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/gpu_megatron/torch/quantization/plugins/test_megatron.py` around lines 738 - 750, The test is forcing GPU->CPU syncs by calling .item() inside the assert message and in the divergence check; change the assert to keep comparisons on-device and only materialize scalars when needed: replace the assert torch.allclose(te_amax.view(()), seq_max, ...) with an if-not pattern (if not torch.allclose(...): raise AssertionError(f"...{te_amax.item()}...{seq_max.item()}...")) so the .item() calls only run on failure, and compute the per-expert divergence using tensor ops (delta = seq_amaxes.max() - seq_amaxes.min(); if (delta > 1e-5): saw_per_expert_divergence = True) while ensuring you only call .item() when you must convert to a Python bool for logging or failure paths; update references to seq_amaxes, seq_mlp.local_experts, linear_name, te_amax, and saw_per_expert_divergence accordingly.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/gpu_megatron/torch/quantization/plugins/test_megatron.py`:
- Around line 698-761: The default-path helper
_test_te_grouped_vs_sequential_default_amax_helper must explicitly disable the
feature flag: set os.environ["MODELOPT_TEGROUPED_PER_EXPERT_QUANTIZER"]="0"
(saving the original and restoring it after the test) before calling
initialize_for_megatron so the shared-quantizer behavior is forced; then add a
new test variant that sets
os.environ["MODELOPT_TEGROUPED_PER_EXPERT_QUANTIZER"]="1" and runs the same
checks to cover the env-on path. Apply the same explicit env-off/restore and
corresponding env-on test change to the other similar helper/test pair
exercising TEGrouped vs Sequential (the companion helper referenced in the
review) so both default and env-enabled code paths have dedicated coverage.
---
Outside diff comments:
In `@modelopt/torch/quantization/plugins/transformer_engine.py`:
- Around line 147-166: The per-expert weight-quantizer path
(_per_expert_weight_quantizer and the created modules weight_quantizer_{i}) is
unsafe on TP/EP topology change because modelopt_post_restore currently only
runs the base path and does not reinitialize per-expert _amax buffers; update
modelopt_post_restore to detect topology/shape mismatch and either (a) fail-fast
by raising a clear error when _per_expert_weight_quantizer is enabled and the
restored per-expert _amax shapes don't match the current weight shapes, or (b)
perform a per-expert recalibration: iterate over each created module name
weight_quantizer_{i} (use self.num_gemms to enumerate), reset/recreate their
_amax buffers to the correct shape for the current self.weight{i} and call the
quantizer's max_calibrate (or its equivalent reinit routine) so each per-expert
quantizer is re-calibrated after restore; implement one of these branches in
modelopt_post_restore to ensure shapes/amax are consistent after load.
---
Nitpick comments:
In `@tests/gpu_megatron/torch/quantization/plugins/test_megatron.py`:
- Around line 738-750: The test is forcing GPU->CPU syncs by calling .item()
inside the assert message and in the divergence check; change the assert to keep
comparisons on-device and only materialize scalars when needed: replace the
assert torch.allclose(te_amax.view(()), seq_max, ...) with an if-not pattern (if
not torch.allclose(...): raise
AssertionError(f"...{te_amax.item()}...{seq_max.item()}...")) so the .item()
calls only run on failure, and compute the per-expert divergence using tensor
ops (delta = seq_amaxes.max() - seq_amaxes.min(); if (delta > 1e-5):
saw_per_expert_divergence = True) while ensuring you only call .item() when you
must convert to a Python bool for logging or failure paths; update references to
seq_amaxes, seq_mlp.local_experts, linear_name, te_amax, and
saw_per_expert_divergence accordingly.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 5fbda55f-b0c7-4054-bc00-63b4bc3a150e
📒 Files selected for processing (2)
modelopt/torch/quantization/plugins/transformer_engine.pytests/gpu_megatron/torch/quantization/plugins/test_megatron.py
| def _test_te_grouped_vs_sequential_default_amax_helper(tp_size, ep_size, quant_cfg, rank, size): | ||
| """TEGrouped per-linear amax should equal max-over-Sequential-experts under default sync=False.""" | ||
| initialize_for_megatron( | ||
| tensor_model_parallel_size=tp_size, | ||
| expert_model_parallel_size=ep_size, | ||
| seed=SEED, | ||
| ) | ||
|
|
||
| te_grouped = _gpt_model_provider( | ||
| tp_size=tp_size, ep_size=ep_size, hidden_size=32, moe_grouped_gemm=True, | ||
| transformer_impl="transformer_engine", num_moe_experts=4, | ||
| ) | ||
| forward = get_forward(te_grouped, batch_size=8) | ||
|
|
||
| sequential = _gpt_model_provider( | ||
| tp_size=tp_size, ep_size=ep_size, hidden_size=32, moe_grouped_gemm=False, | ||
| num_moe_experts=4, transformer_impl="modelopt", | ||
| ) | ||
| copy_weights_from_grouped_to_non_grouped(te_grouped, sequential) | ||
|
|
||
| for module in te_grouped.modules(): | ||
| if isinstance(module, TopKRouter): | ||
| module.topk = module.num_experts | ||
| for module in sequential.modules(): | ||
| if isinstance(module, TopKRouter): | ||
| module.topk = module.num_experts | ||
|
|
||
| mtq.quantize(te_grouped, quant_cfg, forward) | ||
| mtq.quantize(sequential, quant_cfg, forward) | ||
|
|
||
| te_modules = [m for m in te_grouped.modules() if isinstance(m, TEGroupedMLP)] | ||
| seq_modules = [m for m in sequential.modules() if isinstance(m, SequentialMLP)] | ||
| assert len(te_modules) == len(seq_modules) | ||
|
|
||
| saw_per_expert_divergence = False | ||
| for te_mlp, seq_mlp in zip(te_modules, seq_modules): | ||
| for linear_name in ("linear_fc1", "linear_fc2"): | ||
| te_amax = getattr(te_mlp, linear_name).weight_quantizer.amax | ||
| assert te_amax is not None and te_amax.numel() == 1 | ||
|
|
||
| seq_amaxes = torch.stack([ | ||
| getattr(expert, linear_name).weight_quantizer.amax.view(()) | ||
| for expert in seq_mlp.local_experts | ||
| ]) | ||
| seq_max = seq_amaxes.max() | ||
|
|
||
| assert torch.allclose(te_amax.view(()), seq_max, atol=1e-5, rtol=1e-5), ( | ||
| f"TEGrouped per-linear amax ({te_amax.item()}) != " | ||
| f"max-over-Sequential-experts ({seq_max.item()}) for {linear_name}" | ||
| ) | ||
|
|
||
| if (seq_amaxes.max() - seq_amaxes.min()).item() > 1e-5: | ||
| saw_per_expert_divergence = True | ||
|
|
||
| assert saw_per_expert_divergence, ( | ||
| "Expected per-expert weight amax to diverge across SequentialMLP experts." | ||
| ) | ||
|
|
||
|
|
||
| @pytest.mark.parametrize("quant_cfg", [mtq.FP8_DEFAULT_CFG, mtq.NVFP4_DEFAULT_CFG]) | ||
| def test_te_grouped_vs_sequential_default_amax(dist_workers_size_4, quant_cfg): | ||
| dist_workers_size_4.run( | ||
| partial(_test_te_grouped_vs_sequential_default_amax_helper, 1, 2, quant_cfg) | ||
| ) |
There was a problem hiding this comment.
Make these “default” tests control the feature flag explicitly.
These helpers never clear or set MODELOPT_TEGROUPED_PER_EXPERT_QUANTIZER, so a worker process with that flag exported will exercise the new per-expert path instead of the default shared-quantizer path. That also means the env-enabled branch added in this PR still has no dedicated coverage. Please force the flag off in these default-behavior tests and add a separate env-on case for the new path.
Also applies to: 764-820
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/gpu_megatron/torch/quantization/plugins/test_megatron.py` around lines
698 - 761, The default-path helper
_test_te_grouped_vs_sequential_default_amax_helper must explicitly disable the
feature flag: set os.environ["MODELOPT_TEGROUPED_PER_EXPERT_QUANTIZER"]="0"
(saving the original and restoring it after the test) before calling
initialize_for_megatron so the shared-quantizer behavior is forced; then add a
new test variant that sets
os.environ["MODELOPT_TEGROUPED_PER_EXPERT_QUANTIZER"]="1" and runs the same
checks to cover the env-on path. Apply the same explicit env-off/restore and
corresponding env-on test change to the other similar helper/test pair
exercising TEGrouped vs Sequential (the companion helper referenced in the
review) so both default and env-enabled code paths have dedicated coverage.
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #1550 +/- ##
==========================================
+ Coverage 70.21% 74.57% +4.35%
==========================================
Files 515 515
Lines 57244 57284 +40
==========================================
+ Hits 40196 42721 +2525
+ Misses 17048 14563 -2485
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
| from modelopt.torch.quantization.utils import replace_function | ||
|
|
||
| from ..nn import QuantModuleRegistry | ||
| from ..nn import QuantModuleRegistry, SequentialQuantizer |
There was a problem hiding this comment.
Should we create a new quantizer class say GroupedQuantizer which is ` Pytorch nn.ModuleList instead of [nn.Sequantial]?
Sequential container assumes it will be a called as seq(x) which does seq[1](seq[0](x)).
ModuleList simply groups quantizers. but does not define an inherent forward path.
There was a problem hiding this comment.
this implementation doesn't use SequentialQuantizer for the expert quantizers, it simply checks in row 174 if isinstance(wq_i, SequentialQuantizer) to decide how to call the quantizer.
We can add a GroupedQuantizer class if that would make things easier, but the current implementation doesn't really need it IMO
There was a problem hiding this comment.
I added a GroupedQuantizer class and changed the TEGroupedMLP implementation to use it
There was a problem hiding this comment.
Design looks good to me overall, left a comment here - https://github.com/NVIDIA/Model-Optimizer/pull/1550/changes#r3468773327
There was a problem hiding this comment.
Warning
CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.
Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.
Actionable comments posted: 1
🧹 Nitpick comments (1)
modelopt/torch/quantization/plugins/transformer_engine.py (1)
167-167: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winMove this import to module scope or document why it must stay local.
The new function-local import has no circular/optional/heavy-import justification. As per coding guidelines, “Keep imports at the top of the file; place imports inside functions only when necessary … with explicit justification.”
Proposed fix
+from modelopt.torch.quantization.model_calib import max_calibrate- from modelopt.torch.quantization.model_calib import max_calibrate - for i in range(self.num_gemms):🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@modelopt/torch/quantization/plugins/transformer_engine.py` at line 167, The new function-local import of max_calibrate in transformer_engine.py should not stay inside the function unless there is a documented need such as a circular dependency or optional/heavy import. Move the import to module scope near the other top-level imports, or add a clear in-code justification at the exact import site if it must remain local. Use the surrounding transformer_engine plugin code and the max_calibrate reference to locate and update the import placement consistently.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@modelopt/torch/quantization/plugins/transformer_engine.py`:
- Around line 147-150: Persist the per-expert quantizer mode alongside the
checkpoint metadata in transformer_engine.py so restore_quantizer_state() sees
the same module topology on load. Update the initialization path around
_per_expert_weight_quantizer_enabled() and add_module(...) to record the
MODELOPT_TEGROUPED_PER_EXPERT_QUANTIZER setting in the saved config or infer it
from restored state before quantizer modules are created, ensuring the number of
weight_quantizer_* modules matches at save and restore time.
---
Nitpick comments:
In `@modelopt/torch/quantization/plugins/transformer_engine.py`:
- Line 167: The new function-local import of max_calibrate in
transformer_engine.py should not stay inside the function unless there is a
documented need such as a circular dependency or optional/heavy import. Move the
import to module scope near the other top-level imports, or add a clear in-code
justification at the exact import site if it must remain local. Use the
surrounding transformer_engine plugin code and the max_calibrate reference to
locate and update the import placement consistently.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: d8a388b2-8131-43ce-84b3-1deafe9b0934
📒 Files selected for processing (1)
modelopt/torch/quantization/plugins/transformer_engine.py
| self._per_expert_weight_quantizer = _per_expert_weight_quantizer_enabled() | ||
| if self._per_expert_weight_quantizer: | ||
| for i in range(self.num_gemms): | ||
| self.add_module(f"weight_quantizer_{i}", TensorQuantizer()) |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy lift
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the target file with line numbers and the surrounding restore/setup logic.
FILE="modelopt/torch/quantization/plugins/transformer_engine.py"
echo "== file length =="
wc -l "$FILE"
echo
echo "== outline =="
ast-grep outline "$FILE" --view expanded || true
echo
echo "== relevant slices =="
sed -n '1,260p' "$FILE" | cat -n
echo
echo "== search for related restore/load paths =="
rg -n "modelopt_post_restore|load_state|state_dict|load\(|amax|per_expert|weight_quantizer_|_per_expert_weight_quantizer_enabled|max_calibrate" modelopt/torch/quantization/plugins/transformer_engine.py modelopt/torch/quantization -SRepository: NVIDIA/Model-Optimizer
Length of output: 50379
🏁 Script executed:
#!/bin/bash
set -euo pipefail
FILE="modelopt/torch/quantization/plugins/transformer_engine.py"
# Read only the relevant ranges around the env-gated modules and the local import.
sed -n '120,240p' "$FILE" | cat -nRepository: NVIDIA/Model-Optimizer
Length of output: 6927
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
p = Path("modelopt/torch/quantization/plugins/transformer_engine.py")
text = p.read_text()
for needle in [
"def modelopt_post_restore",
"max_calibrate",
"_per_expert_weight_quantizer_enabled",
"add_module(f\"weight_quantizer_",
]:
idx = text.find(needle)
print(f"\n## {needle} -> {idx}")
if idx != -1:
start = max(0, text.rfind("\n", 0, idx-400))
end = min(len(text), text.find("\n", idx+400))
print(text[start:end])
PYRepository: NVIDIA/Model-Optimizer
Length of output: 3922
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the quantizer restore/load path and the base class used here.
sed -n '90,170p' modelopt/torch/quantization/conversion.py | cat -n
printf '\n====\n'
sed -n '1,240p' modelopt/torch/quantization/plugins/custom.py | cat -n
printf '\n====\n'
rg -n "def quantizer_state|class _ParallelLinear|modelopt_post_restore|load_state_dict|ApplyModeError" modelopt/torch/quantization -SRepository: NVIDIA/Model-Optimizer
Length of output: 18270
🏁 Script executed:
#!/bin/bash
set -euo pipefail
sed -n '1,120p' modelopt/torch/quantization/model_calib.py | cat -n
printf '\n====\n'
sed -n '1,120p' modelopt/torch/quantization/plugins/vllm.py | cat -nRepository: NVIDIA/Model-Optimizer
Length of output: 10780
Persist the per-expert quantizer mode in checkpoint metadata. restore_quantizer_state() rejects unmatched/extra quantizer keys, so changing MODELOPT_TEGROUPED_PER_EXPERT_QUANTIZER between save and restore makes the same checkpoint fail to load. Store the mode in the saved config/metadata or derive it from the restored state before add_module(...) runs so the module topology matches.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@modelopt/torch/quantization/plugins/transformer_engine.py` around lines 147 -
150, Persist the per-expert quantizer mode alongside the checkpoint metadata in
transformer_engine.py so restore_quantizer_state() sees the same module topology
on load. Update the initialization path around
_per_expert_weight_quantizer_enabled() and add_module(...) to record the
MODELOPT_TEGROUPED_PER_EXPERT_QUANTIZER setting in the saved config or infer it
from restored state before quantizer modules are created, ensuring the number of
weight_quantizer_* modules matches at save and restore time.
Source: Coding guidelines
…nfra fixes
modelopt/torch/quantization/plugins/transformer_engine.py:
MODELOPT_TEGROUPED_PER_EXPERT_QUANTIZER=1 opts into per-gemm
weight_quantizer_0..N-1 inside _QuantTEGroupedLinear (deepcopied from
the shared weight_quantizer). Lets TEGroupedMLP recover per-expert
amax granularity, matching SequentialMLP's default behavior.
modelopt/torch/distill/plugins/megatron.py:
LogitsKLLoss.forward prints student/teacher logit stats (mean/std/
min/max/shape) on rank 0 each call. Diagnostic for the QAD loss-spike
investigation — confirms which spec produces which logits without
changing the KL math.
tests/gpu_megatron/torch/quantization/plugins/test_megatron.py:
New test_te_grouped_vs_sequential_default_amax + ..._default_loss
cover the structural amax asymmetry between TEGroupedMLP and
SequentialMLP (TEGrouped per-linear amax = max-over-Sequential-experts
amax) and a finiteness sanity check on the resulting quant error.
tools/launcher/common/service_utils.sh:
- Fall back to SLURM_PROCID / SLURM_LOCALID when PMIX_*/OMPI_* are
unset, so `[[ "$mpi_local_rank" -eq 0 ]]` doesn't silently pass on
every rank under plain srun.
- util_install_extra_dep: per-node marker so concurrent ranks wait
for rank 0 to finish installing (concurrent pip on a shared FS
leaves a broken state); also installs nvidia-resiliency-ext.
Signed-off-by: Jennifer Chen <jennifchen@nvidia.com>
Signed-off-by: Jennifer Chen <jennifchen@nvidia.com>
- transformer_engine.py: dedup `import copy`/`import os` left over from the
rebase, sort the four imports alphabetically.
- transformer_engine.py: comment near the per-expert weight_quantizer setup
explaining that base modelopt_post_restore won't re-calibrate the
weight_quantizer_{i} modules, so save/restore is only safe when TP/EP is
unchanged. Per-channel _amax shape depends on the TP-sliced output dim.
- service_utils.sh: drop the duplicated mpi_rank / mpi_local_rank
re-assignments — main already carries the SLURM fallback, the extra two
lines were leftover rebase noise.
Signed-off-by: Jennifer Chen <jennifchen@nvidia.com>
Signed-off-by: Jennifer Chen <jennifchen@nvidia.com>
Co-authored-by: realAsma <86726418+realAsma@users.noreply.github.com> Signed-off-by: Jenny Chen <jc4686@columbia.edu>
| self._per_expert_weight_quantizer = _per_expert_weight_quantizer_enabled() | ||
| if self._per_expert_weight_quantizer: | ||
| for i in range(self.num_gemms): | ||
| self.add_module(f"weight_quantizer_{i}", TensorQuantizer()) |
There was a problem hiding this comment.
@realAsma previously I did a copy.deepcopy(self.weight_quantizer) but you suggested to remove it .. how does this new quantizer inherit the quant format now?
There was a problem hiding this comment.
We just setup and initialize the tensor quantizers here. (called here
)The actual config will be loaded here -
There was a problem hiding this comment.
BB: Can you add more context if needed here?
There was a problem hiding this comment.
🤖 Bot comment.
Prepared a one-line context comment here explaining that the per-GEMM quantizers keep independent calibration state for each grouped weight.
Validation: git diff --check passed. I did not commit or push because this comment did not explicitly approve publishing; ready to push once approved.
Fused TEGroupedLinear experts previously shared a single weight quantizer.
Add GroupedQuantizer (an nn.ModuleList of TensorQuantizers) and have
_QuantTEGroupedLinear build one quantizer per expert in _setup, so each
fused expert gets an independent weight amax by default, matching
SequentialMLP. This replaces the MODELOPT_TEGROUPED_PER_EXPERT_QUANTIZER
env-var gate.
Storing the quantizers in a ModuleList named weight_quantizer surfaces them
as weight_quantizer.{i}, which the fused-experts name normalizer maps to
*weight_quantizer so the stock quant configs reach every expert. Teach the
calibration/sync paths (model_calib DP/EP/TP sync, megatron NVFP4-static TP
check, custom restore probe) to recurse into the container so per-expert
amaxes sync correctly.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: Jennifer Chen <jennifchen@nvidia.com>
Update test_te_grouped_vs_sequential_default_amax to assert each fused expert's weight amax matches its SequentialMLP counterpart (and that they diverge across experts) now that TEGroupedMLP quantizes per-expert by default, instead of comparing a single shared amax. Drop the sync_expert_weight_amax=True override in test_te_grouped_vs_sequential_quantize so both models quantize per-expert and produce identical outputs. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Jennifer Chen <jennifchen@nvidia.com>
jenchen13
changed the title
jenchen13
force-pushed
the
jennifchen/te_per_expert
branch
from
308d8ee to
096d430
Compare
Signed-off-by: Jenny Chen <jennifchen@nvidia.com>
Signed-off-by: Hung-Yueh Chiang <hungyuehc@nvidia.com>
|
/claude review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
There was a problem hiding this comment.
Claude Review — Per-expert weight quantizer in TEGroupedMLP
Reviewed all 6 changed modelopt/ source files plus the two touched test files (9 files total). The core design — a GroupedQuantizer ModuleList holding one TensorQuantizer per fused expert, threaded through calibration/sync/config-matching — is coherent, and the model_calib.py recursion plus the custom.py/megatron.py isinstance updates are consistent. The opt-in torch.compile loop is well-gated (eager default, calibration excluded from capture). I found one blocking correctness bug in the restore path and one backward-compat concern.
Findings by severity: CRITICAL: 1, IMPORTANT: 1, SUGGESTION: 1
Most impactful:
[CRITICAL ModeState] Restore drops _amax for all experts except expert 0 (transformer_engine.py modelopt_post_restore). The base super().modelopt_post_restore() calls weight_quantizer.reset_amax(), which — because reset_amax is a GroupedQuantizer delegated method — wipes _amax on EVERY member, then recalibrates only member 0 (the forward->self[0] representative). The new per-expert loop is meant to recover the rest, but its if not hasattr(q, "_amax"): continue guard runs AFTER that wipe, so it skips exactly the experts that lost their _amax. Result: after restore(), experts 1..N-1 of every TEGroupedMLP have amax is None -> broken static quantization. Snapshot which members had loaded _amax BEFORE calling super() (which also deletes self.weight). The toy tests only cover quantize, not save->restore with num_gemms > 1, so this is not caught.
[IMPORTANT Compatibility] quantizer-state key schema change without migration (megatron.py _process_quantizer_amax). The single weight_quantizer._amax key becomes per-member weight_quantizer.{i}._amax. Existing pre-PR TEGroupedMLP checkpoints then fail restore — ApplyModeError (unmatched keys) on the non-sharded path, KeyError on the Megatron torch-dist path. Needs a migration shim or a state-version bump.
The SUGGESTION notes GroupedQuantizer.get_modelopt_state()/is_grouped_quantizer is currently dead code (no serializer collects the container, no restore code reads the flag).
The PR description already flags the TP/EP restore TODO — the CRITICAL finding shows the restore issue is broader (it affects same-topology restore for num_gemms > 1 too). Risk: high for anyone saving/restoring per-expert-quantized TEGroupedMLP checkpoints. Recommend addressing the restore ordering bug before merge.
| continue | ||
| wq_i = self.weight_quantizer[i] | ||
| q = wq_i[0] if isinstance(wq_i, SequentialQuantizer) else wq_i | ||
| if not hasattr(q, "_amax"): |
There was a problem hiding this comment.
[CRITICAL ModeState] Per-expert restore loses _amax for every expert except expert 0.
What happens: super().modelopt_post_restore() (custom.py:177-181) runs before this loop and executes:
self.weight_quantizer.reset_amax() # GroupedQuantizer delegated method -> broadcasts to ALL members
max_calibrate(self.weight_quantizer, lambda wq: wq(self.weight), distributed_sync=False)reset_amax() is in GroupedQuantizer._delegated_methods, so it deletes _amax on every member. But the max_calibrate forward loop calls self.weight_quantizer(self.weight), which routes through GroupedQuantizer.forward → self[0](...), so only member 0 gets recalibrated. Members 1..num_gemms-1 are left with no _amax.
This loop is then supposed to recalibrate the rest, but the guard if not hasattr(q, "_amax"): continue sees exactly those members with the _amax that super() just deleted, and skips them.
Impact: After restore() (both the non-sharded restore_quantizer_state path and the Megatron torch-dist set_extra_state path, which both call modelopt_post_restore), experts 1..N-1 of every TEGroupedMLP have weight_quantizer[i].amax == None. For a static INT8/FP8/NVFP4 config that means broken/incorrect quantization on all but the first expert — restore is not functionally identical to the saved model. num_gemms == 1 happens to work, which is why the toy tests (which cover quantize, not save→restore) don't catch it.
Fix: determine which members had saved _amax before calling super(), e.g.:
had_amax = []
for i in range(self.num_gemms):
wq_i = self.weight_quantizer[i]
q = wq_i[0] if isinstance(wq_i, SequentialQuantizer) else wq_i
had_amax.append(hasattr(q, "_amax"))
super().modelopt_post_restore(prefix=prefix)
delattr(self, "weight")
from modelopt.torch.quantization.model_calib import max_calibrate
for i in range(self.num_gemms):
weight_i = getattr(self, f"weight{i}", None)
if weight_i is None or not had_amax[i]:
continue
wq_i = self.weight_quantizer[i]
wq_i.reset_amax()
max_calibrate(wq_i, lambda wq, w=weight_i: wq(w), distributed_sync=False)(Note super() deletes self.weight, so snapshotting must happen first anyway.) Please also add a save→restore test covering num_gemms > 1 — the current tests exercise quantize only.
| if re.match(r"weight_quantizer\.\d+\..+_amax$", k): | ||
| quantizer_state_dict[k] = v | ||
| else: | ||
| quantizer_state_dict[k] = v.view(-1) if v.numel() == 1 else v |
There was a problem hiding this comment.
[IMPORTANT Compatibility] This change alters the modelopt_state/quantizer-state key schema for TEGroupedLinear with no migration path.
Before: a single weight_quantizer (TensorQuantizer) → one key ...weight_quantizer._amax.
After: a GroupedQuantizer (ModuleList) → per-member keys ...weight_quantizer.0._amax, ...weight_quantizer.1._amax, …
quantizer_state() (conversion.py:202-208) only collects the leaf TensorQuantizer/SequentialQuantizer members, so a checkpoint saved before this PR has key ...weight_quantizer while the restored model now exposes ...weight_quantizer.0/.1/.... On restore:
- Non-sharded path:
restore_quantizer_statecomputesunmatched_keys = state.keys() - model.keys()and raisesApplyModeErroron the mismatch (conversion.py:144-148). - Megatron
torch-distpath:quant_module_set_extra_statedoesquantizer_state[name]for eachweight_quantizer.iand hits aKeyError.
Either way, existing quantized TEGroupedMLP checkpoints become unloadable after this change. Please add a migration/back-compat shim (e.g. map a legacy scalar weight_quantizer._amax onto all weight_quantizer.{i} members during restore) or a state-version bump with an explicit remap, and document the break in the CHANGELOG if intentional.
|
|
||
| def get_modelopt_state(self) -> dict[str, Any]: | ||
| """Get meta state to be saved in checkpoint.""" | ||
| return {"num_quantizers": len(self), "is_grouped_quantizer": True} |
There was a problem hiding this comment.
[SUGGESTION] GroupedQuantizer.get_modelopt_state() and its is_grouped_quantizer flag appear to be dead code. quantizer_state() (conversion.py:202-208) and quant_module_get_extra_state() (megatron.py:145-148) only collect TensorQuantizer/SequentialQuantizer instances — a GroupedQuantizer itself is never serialized, only its leaf members are — and no restore-side code reads is_grouped_quantizer (unlike is_sequential_quantizer, which conversion.py:559 branches on). Either wire this into save/restore (which would also give you a natural home for the legacy-checkpoint migration noted in the megatron comment) or drop the method to avoid implying a round-trip that doesn't exist.
There was a problem hiding this comment.
🧹 Nitpick comments (2)
modelopt/torch/quantization/model_calib.py (1)
356-357: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueUpdate the type annotation to include
GroupedQuantizer.The function now handles
GroupedQuantizerat line 364, but the type hint at line 357 still readsTensorQuantizer | SequentialQuantizer.♻️ Proposed fix
def sync_quantizer_amax_across_tp( - quantizer: TensorQuantizer | SequentialQuantizer, + quantizer: TensorQuantizer | SequentialQuantizer | GroupedQuantizer, linear_name: str,🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@modelopt/torch/quantization/model_calib.py` around lines 356 - 357, The type annotation for sync_quantizer_amax_across_tp is outdated and should include GroupedQuantizer now that the function handles it. Update the quantizer parameter type hint alongside TensorQuantizer and SequentialQuantizer so the signature matches the logic in sync_quantizer_across_tp and the GroupedQuantizer branch.modelopt/torch/quantization/plugins/transformer_engine.py (1)
187-187: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd a comment explaining the in-function import.
Per coding guidelines, imports inside functions require a brief comment naming the reason (e.g., circular import). This likely avoids a circular dependency between
plugins/transformer_engine.pyandmodel_calib.py, but the reason should be documented.♻️ Proposed fix
- from modelopt.torch.quantization.model_calib import max_calibrate + # Import inside function to avoid circular import: model_calib imports from the + # quantization package which loads this plugins module. + from modelopt.torch.quantization.model_calib import max_calibrateAs per coding guidelines, "Put an import inside a function only when there is a concrete reason: resolving a circular import that cannot be restructured... Add a brief comment in those cases naming the reason."
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@modelopt/torch/quantization/plugins/transformer_engine.py` at line 187, The in-function import of max_calibrate in transformer_engine.py needs a brief explanatory comment per the import guidelines. Update the function containing this import to note the concrete reason for keeping it local, likely the circular dependency between transformer_engine and model_calib, so readers understand why it cannot be moved to the module top level.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@modelopt/torch/quantization/model_calib.py`:
- Around line 356-357: The type annotation for sync_quantizer_amax_across_tp is
outdated and should include GroupedQuantizer now that the function handles it.
Update the quantizer parameter type hint alongside TensorQuantizer and
SequentialQuantizer so the signature matches the logic in
sync_quantizer_across_tp and the GroupedQuantizer branch.
In `@modelopt/torch/quantization/plugins/transformer_engine.py`:
- Line 187: The in-function import of max_calibrate in transformer_engine.py
needs a brief explanatory comment per the import guidelines. Update the function
containing this import to note the concrete reason for keeping it local, likely
the circular dependency between transformer_engine and model_calib, so readers
understand why it cannot be moved to the module top level.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 9e78ab5d-dfba-434b-90ba-e77b357e483c
📒 Files selected for processing (6)
CHANGELOG.rstmodelopt/torch/quantization/config.pymodelopt/torch/quantization/model_calib.pymodelopt/torch/quantization/nn/modules/tensor_quantizer.pymodelopt/torch/quantization/plugins/transformer_engine.pytests/gpu_megatron/torch/quantization/plugins/test_megatron.py
✅ Files skipped from review due to trivial changes (2)
- modelopt/torch/quantization/config.py
- CHANGELOG.rst
| setattr(parent_module, name, sequential_quantizer) | ||
|
|
||
|
|
||
| class GroupedQuantizer(nn.ModuleList): |
There was a problem hiding this comment.
nit: SequentialQuantizer and GroupedQuantizer share sever overlapping pieces (such delegation, attribute setter etc). Can we create a shared base-class which implements these overlapping pieces and subclass SequentialQuantizer and GroupedQuantizer from it to avoid code level redundancy?
There was a problem hiding this comment.
RB: Can you create a plan on how to do this and share in this thread?
There was a problem hiding this comment.
🤖 Bot comment.
🐝 I’m reviewing the overlapping behavior and will share a concrete refactoring plan here shortly.
There was a problem hiding this comment.
🤖 Bot comment.
Plan:
- Add a private
_QuantizerContainerBasemixin before the two containers. It will own the common representative-property reads, property writes broadcast to every member, delegated lifecycle-method calls, attribute-config fan-out, and the commonget_modelopt_state()shape (using a subclass-provided state-key). - Make
SequentialQuantizerinherit from the mixin plusnn.Sequential, andGroupedQuantizerfrom the mixin plusnn.ModuleList. Keep PyTorch container initialization in each subclass so module registration/state-dict keys remain unchanged. - Keep only semantic differences in the subclasses: accepted member types, sequential chaining vs. grouped representative
forward,disable_rotatedelegation for sequential only, delegated-method return behavior, andconvert_to_single_quantizerfor sequential only. - Add parameterized CPU tests for the shared contract (representative reads, broadcast writes/methods, scalar/list config application, and exact modelopt-state output), plus regression tests that sequential chaining, grouped first-member forward, nested
SequentialQuantizermembers, and checkpoint/state-dict layout are unchanged.
This keeps the refactor local to tensor_quantizer.py and its focused tests; downstream (SequentialQuantizer, GroupedQuantizer) handling does not need to change.
There was a problem hiding this comment.
@jenchen13 could you take a look at this comment?
| # scalar, channel, or block shape instead of flattening them through the legacy | ||
| # single-quantizer path. | ||
| if re.match(r"weight_quantizer\.\d+\..+_amax$", k): | ||
| quantizer_state_dict[k] = v |
There was a problem hiding this comment.
Should we do
| quantizer_state_dict[k] = v | |
| quantizer_state_dict[k] = v.view(-1) if v.numel() == 1 else v |
| # GroupedQuantizer (an nn.ModuleList) surfaces them as ``weight_quantizer.{i}``, which the | ||
| # fused-experts name normalizer maps to ``*weight_quantizer`` so the stock configs apply. | ||
| self.weight_quantizer = GroupedQuantizer( | ||
| *(copy.deepcopy(self.weight_quantizer) for _ in range(self.num_gemms)) |
There was a problem hiding this comment.
We dont need this deep copy here.
| *(copy.deepcopy(self.weight_quantizer) for _ in range(self.num_gemms)) | |
| TensorQuantizer() for _ in range(self.num_gemms)) |
The TensorQuantizers are always initialized with default args; they are configured after _setup -> so we can safely initialize new TensorQuantizers here. deepcopy seems unwarranted to me.
There was a problem hiding this comment.
RB: Can you please explain concisely "The TensorQuantizers are always initialized with default args; they are configured after _setup -> so we can safely initialize new TensorQuantizers here. deepcopy seems unwarranted to me." with code pointers ?
There was a problem hiding this comment.
🤖 Bot comment.
🐝 I’m tracing the initialization and post-_setup configuration paths now; I’ll reply here shortly with exact code pointers.
There was a problem hiding this comment.
🤖 Bot comment.
Yes—the key is the conversion order:
convert_to_quantized_model()first callsreplace_quant_module(), thenset_quantizer_by_cfg().- Conversion invokes the module's
_setup(). The base setup createsweight_quantizerwith the standard 8-bit per-tensor descriptor (custom.py:113-117); that descriptor is justQuantizerAttributeConfig(num_bits=8)(tensor_quant.py:279), equivalent toTensorQuantizer()'s schema defaults. _QuantTEGroupedLinear._setup()builds the group immediately after that base setup (lines 147-156), still before user config is applied.- Afterwards,
set_quantizer_by_cfg()configures the registeredweight_quantizer.0,.1, … children; the fused-expert name normalization makes them match stock*weight_quantizerrules (conversion.py:319-355).
So deepcopy is only cloning a pristine default quantizer (including an empty/default calibrator). Fresh TensorQuantizer() instances give each expert independent state and are configured identically in the next phase. deepcopy would matter only if setup had already attached non-default config/state, which this path does not.
|
|
||
| def _iter_leaf_quantizers(quantizer): | ||
| if isinstance(quantizer, SequentialQuantizer): | ||
| if isinstance(quantizer, (SequentialQuantizer, GroupedQuantizer)): |
There was a problem hiding this comment.
nit:
Create QuantizerContainerBase base class and use that here to avoid explicitly listing like this.
see https://github.com/NVIDIA/Model-Optimizer/pull/1550/changes#r3553373837
| if (stacked.max() - stacked.min()).item() > 1e-5: | ||
| saw_per_expert_divergence = True | ||
|
|
||
| assert saw_per_expert_divergence, ( |
There was a problem hiding this comment.
This is not a good test IMO. This can fail if magically all the expert weights were initialized the same.
There was a problem hiding this comment.
RB: Do you agree with my comment? Wdyt?
There was a problem hiding this comment.
🤖 Bot comment.
🐝 Looking into this now.
There was a problem hiding this comment.
🤖 Bot comment.
Yes, I agree with the underlying concern. The fixed seed makes this deterministic today, so it is not conventionally flaky, but saw_per_expert_divergence tests an incidental property of the fixture rather than the independence contract: a correct implementation would fail if the experts happened to produce equal amaxes. I would initialize each expert's weights to deliberately distinct known magnitudes, then assert the corresponding grouped/sequential amax values (and, if useful, that the quantizer objects are distinct). That makes divergence guaranteed by the test setup and the failure causal.
## Description Clarify that mocks are appropriate for focused interface and wiring coverage, but tests claiming backend or runtime behavior must execute the real implementation. Document `torch.compile` as a concrete example: call the real compiler, and wrap/delegate to it when call tracing is needed. This follows the testing guidance raised in #1550. ## Validation - `git diff --check` - `pre-commit run --files CONTRIBUTING.md` <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Documentation** * Clarified test design guidance to encourage validating real behavior end-to-end. * Added advice to avoid substituting core runtime behavior with fakes when tests should cover actual execution. * Noted that wrapping or delegation is still appropriate when tests need to observe calls or counts. <!-- end of auto-generated comment: release notes by coderabbit.ai --> Signed-off-by: realAsma <akuriparambi@nvidia.com>
3 participants
What does this PR do?
Type of change: ?
note: if different TP/EP is used during checkpoint restore, the per expert weight quantizer may not restore properly. TODO fix this
Usage
# Add a code snippet demonstrating how to use thisTesting
Before your PR is "Ready for review"
Make sure you read and follow Contributor guidelines and your commits are signed (
git commit -s -S).Make sure you read and follow the Security Best Practices (e.g. avoiding hardcoded
trust_remote_code=True,torch.load(..., weights_only=False),pickle, etc.).CONTRIBUTING.md: ✅ / ❌ / N/AAdditional Information
Summary by CodeRabbit
torch.compile(default remains eager).amaxsynchronization/recursion so grouped quantizers participate consistently across distributed calibration paths.amaxbehavior, optional compiled execution, and numerical finiteness.