mirror of
https://github.com/NVIDIA/Model-Optimizer.git
synced 2026-10-02 03:14:52 +08:00
[NVBug: 6038899] Fix MoE export crash on meta tensors with CPU offload (#1155)
## Summary Fixes `NotImplementedError` in `sync_moe_gate_up_amax` when quantizing MoE models (e.g. Qwen3-30B-A3B) on a single GPU with insufficient VRAM. When GPU memory is insufficient, ModelOpt enables CPU offload via accelerate, leaving uncalibrated expert parameters on the `meta` device. During export, `sync_moe_gate_up_amax` calls `torch.equal()` on these meta tensors, which raises `NotImplementedError` because `aten::equal` does not support meta tensors — even though calibration itself completed successfully. ## Changes - Add a guard in `sync_moe_gate_up_amax` to skip amax sync for meta tensors (which have no real data to sync) and emit a warning explaining the root cause. Bug: https://nvbugspro.nvidia.com/bug/6038899 🤖 Generated with [Claude Code](https://claude.com/claude-code) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Bug Fixes** * Added warning messages for unsupported tensor configurations in quantization workflows. * Improved edge case detection to gracefully skip processing in incompatible scenarios. <!-- end of auto-generated comment: release notes by coderabbit.ai --> Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
This commit is contained in:
@@ -1184,6 +1184,18 @@ def sync_moe_gate_up_amax(model: nn.Module) -> int:
|
||||
up_amax = getattr(up_wq, "amax", None)
|
||||
if gate_amax is None or up_amax is None:
|
||||
break
|
||||
# Meta tensors have no storage (e.g. CPU-offloaded experts that
|
||||
# were never activated during calibration). Skip — there is no
|
||||
# real amax data to sync.
|
||||
if gate_amax.is_meta or up_amax.is_meta:
|
||||
warn(
|
||||
f"Skipping gate/up amax sync for expert with meta tensors "
|
||||
f"(gate_amax.is_meta={gate_amax.is_meta}, "
|
||||
f"up_amax.is_meta={up_amax.is_meta}). "
|
||||
f"This typically means the expert was CPU-offloaded and "
|
||||
f"not activated during calibration."
|
||||
)
|
||||
break
|
||||
if not torch.equal(gate_amax, up_amax):
|
||||
shared_amax = torch.max(gate_amax, up_amax)
|
||||
gate_wq.amax = shared_amax
|
||||
|
||||
Reference in New Issue
Block a user