fix: exclude FP8 quantized weights from MXFP4 de-osc

[hann-wang]

What's wrong

MXFP4 weight-deosc process selects eligible weights by looking for quantized_weight_buffer. FP4 and FP8 quantized modules both have this buffer when gradient accumulation enabled.

What's changed

• MXFP4 quantized module creates an empty buffer in float4_e2m1fn_x2 dtype when gradient accumulation disabled.
• Check the dtype of quantized_weight_buffer in MXFP4 weight de-osc.

@HuangWei-95

[Copilot]

🔵 Needs a closer look

It affects low-level quantization/training behavior via a custom FP4 dtype marker and should get final human validation in a real runtime/config matrix before approval.

Pull request overview

This PR refines MXFP4 weight de-oscillation eligibility detection so FP8-quantized modules are no longer mistakenly picked up when quantized_weight_buffer exists for both FP4 and FP8 paths (e.g., when gradient accumulation is enabled).

Changes:

• In Primus-Turbo FP4 paths with get_num_microbatches() == 1, set quantized_weight_buffer to an empty tensor with float4_e2m1fn_x2 dtype (instead of torch.uint8) to act as an FP4-specific runtime marker.
• In the weight de-osc eligibility scan, require quantized_weight_buffer.dtype == float4_e2m1fn_x2 to filter out FP8 buffers.

File summaries

File	Description
primus/backends/megatron/core/optimizer/weight_deosc.py	Tightens FP4 eligibility detection by checking quantized_weight_buffer dtype.
primus/backends/megatron/core/extensions/primus_turbo.py	Makes the FP4 “marker” buffer use float4_e2m1fn_x2 dtype so de-osc can reliably identify FP4 modules.

Review details

• Files reviewed: 2/2 changed files
• Comments generated: 0
• Review effort level: Lite

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Copilot]

🟡 Changes recommended

There are concrete runtime-breaking issues (undefined weights variable in multiple FP4 paths) and the new dtype check can raise/disable eligibility when the MXFP4 dtype symbol is unavailable or the marker isn’t a tensor.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Review details

Suppressed comments (3)

primus/backends/megatron/core/extensions/primus_turbo.py:1365

• weights is not defined in this forward_internal (the local variable is weight), so device=weights.device will raise NameError the first time this FP4 path runs.

                        self.quantized_weight_buffer = torch.empty(
                            0, device=weights.device, dtype=float4_e2m1fn_x2
                        )

primus/backends/megatron/core/extensions/primus_turbo.py:1571

• weights is not defined in this forward_internal (the local variable is weight), so device=weights.device will raise NameError the first time this FP4 path runs.

                        self.quantized_weight_buffer = torch.empty(
                            0, device=weights.device, dtype=float4_e2m1fn_x2
                        )

primus/backends/megatron/core/extensions/primus_turbo.py:1790

• weights is not defined in this forward_internal (the local variable is weight), so device=weights.device will raise NameError the first time this FP4 path runs.

                        self.quantized_weight_buffer = torch.empty(
                            0, device=weights.device, dtype=float4_e2m1fn_x2
                        )

• Files reviewed: 2/2 changed files
• Comments generated: 2
• Review effort level: Lite