v4 mxfp4

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