[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:
Chenjie Luo
2026-04-02 06:59:57 +00:00
committed by GitHub
parent 7cd7d183f4
commit fcb09bf11d
+12
View File
@@ -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