diff --git a/megatron/core/extensions/transformer_engine.py b/megatron/core/extensions/transformer_engine.py index 57fbe5bd7c5..e7c4af89cfa 100644 --- a/megatron/core/extensions/transformer_engine.py +++ b/megatron/core/extensions/transformer_engine.py @@ -1512,6 +1512,107 @@ def fake_int4_quantization_ste(x, group_size): return x_out + def _mxfp4_to_blockwise_fp8_qtensor(w): + """bf16 weight -> MXFP4 (E2M1, 1x32 UE8M0) -> blockwise FP8 (E4M3, 128x128), entirely via + deepseek-ai/TileKernels, wrapped as a TE ``Float8BlockwiseQTensor`` so the TE GEMM consumes + it directly: TE does NOT re-quantize the weight (it only quantizes the activation). + + per_block_cast('e2m1', 1x32) -> per_block_cast_lossless((1x32)->(128x128), 'e4m3') + + FP4->FP8 is lossless under the per-128x128-block scale-ratio bound that + ``per_block_cast_lossless`` device-asserts. The columnwise rep (needed by wgrad/dgrad) is + produced by quantizing ``w.T`` — for 128x128 square blocks it equals the transpose. + + NOTE (needs Blackwell verification): the exact ``Float8BlockwiseQTensor`` scale_inv GEMM + layout/padding and constructor kwargs are TE-version specific. Run once with + ``MXFP4_QAT_VERIFY=1`` to assert TE's dequant of this tensor matches TileKernels' dequant + before trusting it (a wrong layout would otherwise silently corrupt training). + """ + import transformer_engine_torch as tex + from transformer_engine.pytorch.tensor.float8_blockwise_tensor import ( + Float8BlockQuantizer, + Float8BlockwiseQTensor, + ) + + try: + from tile_kernels.quant import cast_back, per_block_cast, per_block_cast_lossless + except ImportError as e: + raise ImportError( + "USE_MXFP4_QAT weight path needs deepseek-ai/TileKernels " + "(tile_kernels.quant.per_block_cast / per_block_cast_lossless)." + ) from e + + assert w.dim() == 2, f"MXFP4 weight QAT expects 2D expert weights, got {tuple(w.shape)}" + + # MXFP4 QAT is Blackwell-only (the run_deepseek_v4 recipe pins + # NVTE_FP8_BLOCK_SCALING_FP32_SCALES=0), which for TE's Float8BlockScaling means + # power_2_scale=True with FP32 scale CONTAINERS (not packed E8M0 — that is the separate + # MXFP8 recipe). So we match it with pow-2 values stored as fp32: + # bf16 -> E2M1: 1x32 UE8M0 (pow-2) micro-scale. + # E2M1 -> E4M3: 128x128 pow-2 scale, stored fp32 (use_packed_ue8m0=False). + def _to_fp8(mat): + fp4 = per_block_cast(mat, "e2m1", (1, 32), round_sf=True, use_packed_ue8m0=True) + return per_block_cast_lossless( + fp4, "e4m3", x_block_size=(1, 32), out_block_size=(128, 128), + round_sf=True, use_packed_ue8m0=False, + ) + + rowwise_fp8, rowwise_sf = _to_fp8(w.contiguous()) + colwise_fp8, colwise_sf = _to_fp8(w.t().contiguous()) # 128x128 columnwise == transpose + + quantizer = Float8BlockQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=True, + force_pow_2_scales=True, # env=0 -> power_2_scale=True (pow-2 values, fp32 container) + block_scaling_dim=2, # 128x128 weight blocks + ) + + qtensor = Float8BlockwiseQTensor( + shape=w.shape, + dtype=w.dtype, + rowwise_data=rowwise_fp8.view(torch.uint8), + rowwise_scale_inv=rowwise_sf, + columnwise_data=colwise_fp8.view(torch.uint8), + columnwise_scale_inv=colwise_sf, + fp8_dtype=tex.DType.kFloat8E4M3, + quantizer=quantizer, + is_2D_scaled=True, + ) + + if os.getenv("MXFP4_QAT_VERIFY", "0") == "1": + # Surface a layout/scale mismatch loudly instead of silently corrupting training: + # TE's dequant of our QTensor must equal the MXFP4-grid weight (FP4->FP8 is lossless). + # The reference uses an fp32-scale round-trip so it is independent of the FP8 scale packing. + ref_fp4 = per_block_cast(w.contiguous(), "e2m1", (1, 32), round_sf=True, use_packed_ue8m0=False) + ref = cast_back(ref_fp4, "fp32", (1, 32)) + te_dq = qtensor.dequantize(dtype=torch.float32) + if not torch.allclose(te_dq, ref, atol=0.0, rtol=0.0): + max_err = (te_dq.float() - ref.float()).abs().max().item() + raise RuntimeError( + "MXFP4_QAT Float8BlockwiseQTensor mismatch vs the MXFP4-grid weight " + f"(max_abs_err={max_err}). The scale_inv GEMM layout/dtype/padding or constructor " + "kwargs likely differ from TE's Float8BlockScaling; fix before training." + ) + + return qtensor + + class _MXFP4WeightToFp8STE(torch.autograd.Function): + @staticmethod + def forward(ctx, w): + return _mxfp4_to_blockwise_fp8_qtensor(w) + + @staticmethod + def backward(ctx, grad_output): + # Straight-through: grad flows unchanged to the bf16/fp32 master weight. + return grad_output + + def mxfp4_weight_to_fp8_qtensor(w): + out = _MXFP4WeightToFp8STE.apply(w) + if hasattr(w, "main_grad"): + out.main_grad = w.main_grad + return out + class TEGroupedLinear(te.pytorch.GroupedLinear): """ Wrapper for the Transformer-Engine's `GroupedLinear` layer. @@ -1738,6 +1839,14 @@ def _get_weight_tensors(self): fake_int4_quantization_ste(w, group_size) for w in weight_tensors ] + elif os.getenv("USE_MXFP4_QAT", "0") == "1": + # MXFP4 weight QAT: bf16 -> MXFP4 -> blockwise FP8 entirely via TileKernels, handed + # to TE as a Float8BlockwiseQTensor (TE does NOT re-quantize the weight). Backward + # is STE to the bf16/fp32 master. + weight_tensors = [ + mxfp4_weight_to_fp8_qtensor(w) + for w in weight_tensors + ] return weight_tensors