Cache quantized weights for full AC

docs/advanced.md

@@ -42,13 +42,13 @@ The custom path enables you to set EP, CP, selective activation checkpointing, l
 Set `[trainer.model.quantization]` to train dense linears and MoE expert GEMMs in low precision. Two backends are available via the `type` discriminator:
 
 - `type = "fp8"` — DeepGEMM FP8 blockwise (requires SM90+ / Hopper). Options: `enable_grouped_gemm` (FP8 MoE expert GEMM). Both default on.
-- `type = "mxfp8"` — torchao MXFP8 microscaling (requires SM100+ / Blackwell). Options: `enable_grouped_gemm`, `enable_a2a` (MXFP8 expert-parallel all-to-all), and `recipe` (`mxfp8_rceil` default or `mxfp8_rceil_wgrad_with_hp`).
+- `type = "mxfp8"` — torchao MXFP8 microscaling (requires SM100+ / Blackwell). Options: `enable_grouped_gemm`, `enable_a2a` (quantize expert-parallel dispatch/combine tokens to mxfp8 before the all-to-all instead of sending bf16; defaults off — the smaller payload only pays off when the interconnect is the bottleneck, e.g. multi-node over a slower fabric, whereas on a single node over NVLink it's usually pure overhead), and `recipe` (`mxfp8_rceil` default or `mxfp8_rceil_wgrad_with_hp`).
 
 ```toml
 [trainer.model.quantization]
 type = "mxfp8"
 recipe = "mxfp8_rceil"
-enable_a2a = true
+enable_a2a = false
 ```
 
 GLM-5.2 adds IndexShare: the DSA sparse-attention indexer runs only on a subset of layers and the remaining layers reuse the cached top-k indices. The trainer reads this schedule from the model's `indexer_types` config field and enables the index cache automatically, so no extra config is needed. To override the schedule manually, set `[trainer.model.index_cache]` (`topk_freq` or `topk_pattern`).

packages/prime-rl-configs/src/prime_rl/configs/trainer.py

@@ -169,7 +169,7 @@ class MXFP8Config(BaseConfig):
     type: Literal["mxfp8"] = "mxfp8"
     recipe: MXFP8Recipe = "mxfp8_rceil"
     enable_grouped_gemm: bool = True
-    enable_a2a: bool = True
+    enable_a2a: bool = False
     ignore_patterns: list[str] = _DEFAULT_FP8_IGNORE_PATTERNS
 
 

src/prime_rl/trainer/models/layers/moe.py

@@ -177,9 +177,12 @@ def _run_experts_grouped_mm_impl(
         h = h * grouped_fp8_gemm(x.bfloat16(), w3.bfloat16().transpose(-2, -1), offsets)
         out = grouped_fp8_gemm(h, w2.bfloat16().transpose(-2, -1), offsets).type_as(x)
     else:
-        h = F.silu(torch._grouped_mm(x.bfloat16(), w1.bfloat16().transpose(-2, -1), offs=offsets))
-        h = h * torch._grouped_mm(x.bfloat16(), w3.bfloat16().transpose(-2, -1), offs=offsets)
-        out = torch._grouped_mm(h, w2.bfloat16().transpose(-2, -1), offs=offsets).type_as(x)
+        w1 = w1 if w1.dtype == torch.bfloat16 else w1.bfloat16()
+        w2 = w2 if w2.dtype == torch.bfloat16 else w2.bfloat16()
+        w3 = w3 if w3.dtype == torch.bfloat16 else w3.bfloat16()
+        h = F.silu(torch._grouped_mm(x.bfloat16(), w1.transpose(-2, -1), offs=offsets))
+        h = h * torch._grouped_mm(x.bfloat16(), w3.transpose(-2, -1), offs=offsets)
+        out = torch._grouped_mm(h, w2.transpose(-2, -1), offs=offsets).type_as(x)
 
     return out
 

src/prime_rl/trainer/models/layers/mxfp8_linear.py

@@ -5,13 +5,91 @@
 import torch
 from torch import nn
 from torchao.prototype.mx_formats import ScaleCalculationMode
+from torchao.prototype.mx_formats import mx_linear as tao_mx_linear
 from torchao.prototype.mx_formats.mx_linear import _to_mxfp8_then_scaled_mm
 from torchao.quantization.quantize_.common import KernelPreference
 
 from prime_rl.configs.trainer import MXFP8Recipe
 from prime_rl.utils.logger import get_logger
 
 
+def _cache_mxfp8_dim0_weight_across_checkpoint_recompute() -> None:
+    """Cache the dim0 MXFP8 quantization of the weight on the tensor itself so activation checkpointing's recompute forward reuses it instead of requantizing.
+    """
+    if getattr(tao_mx_linear.mx_mm, "_prime_rl_dim0_cached", False):
+        return
+
+    MXTensor = tao_mx_linear.MXTensor
+    orig_backward = tao_mx_linear.mx_mm.backward
+
+    @staticmethod
+    def forward(
+        ctx,
+        input_hp,
+        weight_hp,
+        in_elem_dtype,
+        w_elem_dtype,
+        grad_elem_dtype,
+        block_size,
+        kernel_preference,
+        mxfp8_dim0_cast_kernel_choice,
+        mxfp8_dim1_cast_kernel_choice,
+        scale_calculation_mode,
+        wgrad_with_hp,
+    ):
+        ctx.save_for_backward(input_hp, weight_hp)
+        ctx._prime_rl_weight_hp = weight_hp
+        ctx.in_elem_dtype = in_elem_dtype
+        ctx.w_elem_dtype = w_elem_dtype
+        ctx.grad_elem_dtype = grad_elem_dtype
+        ctx.block_size = block_size
+        ctx.kernel_preference = kernel_preference
+        ctx.wgrad_with_hp = wgrad_with_hp
+        ctx.mxfp8_dim0_cast_kernel_choice = mxfp8_dim0_cast_kernel_choice
+        ctx.mxfp8_dim1_cast_kernel_choice = mxfp8_dim1_cast_kernel_choice
+        ctx.scale_calculation_mode = scale_calculation_mode
+        input_orig_shape = input_hp.shape
+        input_hp_r = input_hp.reshape(-1, input_orig_shape[-1])
+        input_mx_r_dim0 = MXTensor.to_mx(
+            input_hp_r,
+            in_elem_dtype,
+            block_size,
+            scale_calculation_mode,
+            kernel_preference,
+            mxfp8_dim0_cast_kernel_choice=mxfp8_dim0_cast_kernel_choice,
+        )
+        cached = getattr(weight_hp, "_prime_rl_mxfp8_dim0_cache", None)
+        if cached is not None and cached[0] == weight_hp._version:
+            weight_mx_dim0 = cached[1]
+        else:
+            weight_mx_dim0 = MXTensor.to_mx(
+                weight_hp,
+                w_elem_dtype,
+                block_size,
+                scale_calculation_mode,
+                kernel_preference,
+                mxfp8_dim0_cast_kernel_choice=mxfp8_dim0_cast_kernel_choice,
+            )
+            weight_hp._prime_rl_mxfp8_dim0_cache = (weight_hp._version, weight_mx_dim0)
+
+        output = torch.mm(input_mx_r_dim0, weight_mx_dim0.t())
+        output = output.reshape(*input_orig_shape[:-1], output.shape[-1])
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output_hp):
+        weight_hp = ctx._prime_rl_weight_hp
+        try:
+            return orig_backward(ctx, grad_output_hp)
+        finally:
+            if hasattr(weight_hp, "_prime_rl_mxfp8_dim0_cache"):
+                del weight_hp._prime_rl_mxfp8_dim0_cache
+
+    tao_mx_linear.mx_mm.forward = forward
+    tao_mx_linear.mx_mm.backward = backward
+    tao_mx_linear.mx_mm._prime_rl_dim0_cached = True
+
+
 class MXFP8Linear(nn.Linear):
     def __init__(
         self,
@@ -66,6 +144,7 @@ def from_linear(
 
 
 def replace_linear_with_mxfp8_linear(model: nn.Module, recipe: MXFP8Recipe, ignore_modules: list[str]) -> None:
+    _cache_mxfp8_dim0_weight_across_checkpoint_recompute()
     wgrad_with_hp = recipe == "mxfp8_rceil_wgrad_with_hp"
     logger = get_logger()
     logger.info(f"Replacing linear layers with MXFP8 linear layers (recipe={recipe}, ignore={ignore_modules})")

tests/unit/train/models/test_mxfp8_linear.py

@@ -0,0 +1,119 @@
+import torch
+import torch.utils.checkpoint
+import pytest
+
+from prime_rl.trainer.models.layers.mxfp8_linear import (
+    MXFP8Linear,
+    _cache_mxfp8_dim0_weight_across_checkpoint_recompute,
+)
+
+pytestmark = [
+    pytest.mark.gpu,
+    pytest.mark.skipif(
+        not torch.cuda.is_available() or torch.cuda.get_device_capability() < (10, 0),
+        reason="MXFP8 requires SM100 (Blackwell) or newer",
+    ),
+]
+
+
+@pytest.fixture(scope="module", autouse=True)
+def enable_dim0_cache():
+    _cache_mxfp8_dim0_weight_across_checkpoint_recompute()
+
+
+def _make_layer(seed: int, in_features: int = 128, out_features: int = 128) -> MXFP8Linear:
+    torch.manual_seed(seed)
+    return MXFP8Linear(in_features, out_features, bias=False, device="cuda", dtype=torch.bfloat16)
+
+
+def test_cached_forward_matches_fresh_quantization():
+    """A second forward on an unmodified weight (cache hit) must be bit-identical to the
+    first (cache miss, freshly quantized)."""
+    layer = _make_layer(seed=0)
+    x = torch.randn(32, 128, device="cuda", dtype=torch.bfloat16)
+
+    out_miss = layer(x)
+    out_hit = layer(x)
+
+    assert torch.equal(out_miss, out_hit)
+    cached_version, _ = layer.weight._prime_rl_mxfp8_dim0_cache
+    assert cached_version == layer.weight._version
+
+
+def test_cache_invalidates_after_inplace_weight_update():
+    """An in-place weight mutation bumps `_version`, so the cache must requantize
+    rather than silently reuse the quantization of the old weight."""
+    layer = _make_layer(seed=0)
+    x = torch.randn(32, 128, device="cuda", dtype=torch.bfloat16)
+
+    out_before = layer(x)
+    with torch.no_grad():
+        layer.weight.add_(1.0)
+    out_after = layer(x)
+
+    assert not torch.equal(out_before, out_after)
+
+    reference = _make_layer(seed=1)  # different init, then overwritten below
+    with torch.no_grad():
+        reference.weight.copy_(layer.weight)
+    assert torch.equal(out_after, reference(x))
+
+
+def test_distinct_weights_never_share_a_cache_entry():
+    """Two different weight tensors, even allocated back-to-back, must never read
+    each other's cached quantization."""
+    layer_a = _make_layer(seed=0)
+    layer_b = _make_layer(seed=1)
+    assert not torch.equal(layer_a.weight, layer_b.weight)
+    x = torch.randn(32, 128, device="cuda", dtype=torch.bfloat16)
+
+    out_a = layer_a(x)
+    out_b = layer_b(x)
+    out_a_again = layer_a(x)
+
+    assert torch.equal(out_a, out_a_again)
+    assert not torch.equal(out_a, out_b)
+
+
+def test_activation_checkpoint_recompute_matches_uncheckpointed_grads():
+    """Full activation checkpointing forces a second (recompute) forward call that hits
+    the dim0 cache. Output and gradients must be bit-identical to an uncheckpointed run,
+    since backward recomputes its own dim1 quantization independently of this cache."""
+    layer = _make_layer(seed=0)
+    x = torch.randn(32, 128, device="cuda", dtype=torch.bfloat16)
+
+    x_baseline = x.clone().requires_grad_()
+    layer.weight.grad = None
+    out_baseline = layer(x_baseline)
+    out_baseline.float().pow(2).mean().backward()
+    grad_w_baseline = layer.weight.grad.clone()
+    grad_x_baseline = x_baseline.grad.clone()
+
+    x_ckpt = x.clone().requires_grad_()
+    layer.weight.grad = None
+    out_ckpt = torch.utils.checkpoint.checkpoint(layer, x_ckpt, use_reentrant=False)
+    out_ckpt.float().pow(2).mean().backward()
+    grad_w_ckpt = layer.weight.grad.clone()
+    grad_x_ckpt = x_ckpt.grad.clone()
+
+    assert torch.equal(out_baseline, out_ckpt)
+    assert torch.equal(grad_w_baseline, grad_w_ckpt)
+    assert torch.equal(grad_x_baseline, grad_x_ckpt)
+
+
+def test_mxfp8_forward_close_to_bf16_linear():
+    """Sanity check the quantized path against a plain bf16 nn.Linear: MXFP8 is lossy,
+    so we check relative error stays within microscaling's expected quantization noise
+    rather than exact equality."""
+    torch.manual_seed(0)
+    mx_layer = MXFP8Linear(256, 256, bias=False, device="cuda", dtype=torch.bfloat16)
+    ref_layer = torch.nn.Linear(256, 256, bias=False, device="cuda", dtype=torch.bfloat16)
+    with torch.no_grad():
+        ref_layer.weight.copy_(mx_layer.weight)
+
+    x = torch.randn(64, 256, device="cuda", dtype=torch.bfloat16)
+    out_mx = mx_layer(x)
+    out_ref = ref_layer(x)
+
+    rel_error = (out_mx.float() - out_ref.float()).norm() / out_ref.float().norm()
+    assert rel_error < 0.1