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,79 @@
 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:
+    if getattr(tao_mx_linear.mx_mm, "_prime_rl_dim0_cached", False):
+        return
+
+    MXTensor = tao_mx_linear.MXTensor
+    cache: dict[int, tuple[int, object]] = {}
+
+    @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.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,
+        )
+        cache_key = id(weight_hp)
+        cached = cache.get(cache_key)
+        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,
+            )
+            cache[cache_key] = (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
+
+    tao_mx_linear.mx_mm.forward = forward
+    tao_mx_linear.mx_mm._prime_rl_dim0_cached = True
+
+
 class MXFP8Linear(nn.Linear):
     def __init__(
         self,
@@ -66,6 +132,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})")