From 4be54cbb735f561e70dbcf61fb8b654e373e128f Mon Sep 17 00:00:00 2001 From: Han Wang Date: Thu, 3 Sep 2026 00:55:07 +0000 Subject: [PATCH 1/2] fix: exclude FP8 quantized weights from MXFP4 de-osc --- .../backends/megatron/core/extensions/primus_turbo.py | 10 +++++----- .../backends/megatron/core/optimizer/weight_deosc.py | 3 ++- 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/primus/backends/megatron/core/extensions/primus_turbo.py b/primus/backends/megatron/core/extensions/primus_turbo.py index d7d18b484..1157ca484 100644 --- a/primus/backends/megatron/core/extensions/primus_turbo.py +++ b/primus/backends/megatron/core/extensions/primus_turbo.py @@ -1147,7 +1147,7 @@ def forward_internal( if get_num_microbatches() == 1: if is_first_microbatch: - self.quantized_weight_buffer = torch.empty(0, device=weight.device, dtype=torch.uint8) + self.quantized_weight_buffer = torch.empty(0, device=weight.device, dtype=float4_e2m1fn_x2) out = primus_turbo_torch.ops.gemm_fp4( x, weight, @@ -1358,7 +1358,7 @@ def forward_internal( if get_num_microbatches() == 1: if is_first_microbatch: - self.quantized_weight_buffer = torch.empty(0, device=weight.device, dtype=torch.uint8) + self.quantized_weight_buffer = torch.empty(0, device=weight.device, dtype=float4_e2m1fn_x2) out = primus_turbo_torch.ops.gemm_fp4( x, weight, @@ -1562,7 +1562,7 @@ def forward_internal( if get_num_microbatches() == 1: if is_first_microbatch: - self.quantized_weight_buffer = torch.empty(0, device=weight.device, dtype=torch.uint8) + self.quantized_weight_buffer = torch.empty(0, device=weight.device, dtype=float4_e2m1fn_x2) out = primus_turbo_torch.ops.gemm_fp4( x, weight, @@ -1779,7 +1779,7 @@ def forward_internal(self, x, is_first_microbatch: bool = False): if get_num_microbatches() == 1: if is_first_microbatch: - self.quantized_weight_buffer = torch.empty(0, device=weight.device, dtype=torch.uint8) + self.quantized_weight_buffer = torch.empty(0, device=weight.device, dtype=float4_e2m1fn_x2) out = primus_turbo_torch.ops.gemm_fp4( inp, weight, @@ -2176,7 +2176,7 @@ def forward_internal( # no persistent cache. Expose a lightweight runtime marker so # weight de-oscillation can identify the grouped FP4 weight. if is_first_microbatch: - self.quantized_weight_buffer = torch.empty(0, device=weights.device, dtype=torch.uint8) + self.quantized_weight_buffer = torch.empty(0, device=weights.device, dtype=float4_e2m1fn_x2) out = primus_turbo_torch.ops.grouped_gemm_fp4( x, weights, diff --git a/primus/backends/megatron/core/optimizer/weight_deosc.py b/primus/backends/megatron/core/optimizer/weight_deosc.py index 699186887..1865fe0e8 100644 --- a/primus/backends/megatron/core/optimizer/weight_deosc.py +++ b/primus/backends/megatron/core/optimizer/weight_deosc.py @@ -351,7 +351,8 @@ def _build_eligible_ids(self, dist_opt) -> set: for chunk in model_chunks: modules = chunk.modules() if hasattr(chunk, "modules") else [] for module in modules: - if getattr(module, "quantized_weight_buffer", None) is None: + quantized_weight_buffer = getattr(module, "quantized_weight_buffer", None) + if quantized_weight_buffer is None or quantized_weight_buffer.dtype != _float4_e2m1fn_x2: continue weight = getattr(module, "_parameters", {}).get("weight", None) if weight is None: From e5286cb11c6468f011e7fc63338bad605c436b84 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Thu, 3 Sep 2026 09:23:51 +0800 Subject: [PATCH 2/2] chore: lint --- .../megatron/core/extensions/primus_turbo.py | 20 ++++++++++++++----- 1 file changed, 15 insertions(+), 5 deletions(-) diff --git a/primus/backends/megatron/core/extensions/primus_turbo.py b/primus/backends/megatron/core/extensions/primus_turbo.py index 1157ca484..9edd5b3fe 100644 --- a/primus/backends/megatron/core/extensions/primus_turbo.py +++ b/primus/backends/megatron/core/extensions/primus_turbo.py @@ -1147,7 +1147,9 @@ def forward_internal( if get_num_microbatches() == 1: if is_first_microbatch: - self.quantized_weight_buffer = torch.empty(0, device=weight.device, dtype=float4_e2m1fn_x2) + self.quantized_weight_buffer = torch.empty( + 0, device=weights.device, dtype=float4_e2m1fn_x2 + ) out = primus_turbo_torch.ops.gemm_fp4( x, weight, @@ -1358,7 +1360,9 @@ def forward_internal( if get_num_microbatches() == 1: if is_first_microbatch: - self.quantized_weight_buffer = torch.empty(0, device=weight.device, dtype=float4_e2m1fn_x2) + self.quantized_weight_buffer = torch.empty( + 0, device=weights.device, dtype=float4_e2m1fn_x2 + ) out = primus_turbo_torch.ops.gemm_fp4( x, weight, @@ -1562,7 +1566,9 @@ def forward_internal( if get_num_microbatches() == 1: if is_first_microbatch: - self.quantized_weight_buffer = torch.empty(0, device=weight.device, dtype=float4_e2m1fn_x2) + self.quantized_weight_buffer = torch.empty( + 0, device=weights.device, dtype=float4_e2m1fn_x2 + ) out = primus_turbo_torch.ops.gemm_fp4( x, weight, @@ -1779,7 +1785,9 @@ def forward_internal(self, x, is_first_microbatch: bool = False): if get_num_microbatches() == 1: if is_first_microbatch: - self.quantized_weight_buffer = torch.empty(0, device=weight.device, dtype=float4_e2m1fn_x2) + self.quantized_weight_buffer = torch.empty( + 0, device=weights.device, dtype=float4_e2m1fn_x2 + ) out = primus_turbo_torch.ops.gemm_fp4( inp, weight, @@ -2176,7 +2184,9 @@ def forward_internal( # no persistent cache. Expose a lightweight runtime marker so # weight de-oscillation can identify the grouped FP4 weight. if is_first_microbatch: - self.quantized_weight_buffer = torch.empty(0, device=weights.device, dtype=float4_e2m1fn_x2) + self.quantized_weight_buffer = torch.empty( + 0, device=weights.device, dtype=float4_e2m1fn_x2 + ) out = primus_turbo_torch.ops.grouped_gemm_fp4( x, weights,