From 73efbe988548487e29ca999dddba93bc48af718d Mon Sep 17 00:00:00 2001 From: Mario Sieg Date: Tue, 25 Aug 2026 06:10:28 +0000 Subject: [PATCH 1/4] Cache quantized weights if AC is off --- .../trainer/models/layers/mxfp8_linear.py | 67 +++++++++++++++++++ 1 file changed, 67 insertions(+) diff --git a/src/prime_rl/trainer/models/layers/mxfp8_linear.py b/src/prime_rl/trainer/models/layers/mxfp8_linear.py index a7e579f7f7..a5942ef41f 100644 --- a/src/prime_rl/trainer/models/layers/mxfp8_linear.py +++ b/src/prime_rl/trainer/models/layers/mxfp8_linear.py @@ -5,6 +5,7 @@ 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 @@ -12,6 +13,71 @@ 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})") From 35c04b972bf17c8556073a2605419c4ba7a1f8bb Mon Sep 17 00:00:00 2001 From: Mario Sieg Date: Tue, 25 Aug 2026 15:51:36 +0000 Subject: [PATCH 2/4] Avoid unneccecary casts and disable a2a by default to improve throughput --- docs/advanced.md | 4 ++-- .../prime-rl-configs/src/prime_rl/configs/trainer.py | 2 +- src/prime_rl/trainer/models/layers/moe.py | 10 +++++++--- 3 files changed, 10 insertions(+), 6 deletions(-) diff --git a/docs/advanced.md b/docs/advanced.md index 130e58f0de..ef739fc177 100644 --- a/docs/advanced.md +++ b/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`). diff --git a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py index e67d3db16d..5acfc7eaf1 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py +++ b/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 diff --git a/src/prime_rl/trainer/models/layers/moe.py b/src/prime_rl/trainer/models/layers/moe.py index 4f44cb5031..09c21fcb46 100644 --- a/src/prime_rl/trainer/models/layers/moe.py +++ b/src/prime_rl/trainer/models/layers/moe.py @@ -177,9 +177,13 @@ 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 From be19d4f18d48e44533bd763dd33facea5501deb6 Mon Sep 17 00:00:00 2001 From: Mario Sieg Date: Tue, 25 Aug 2026 20:01:35 +0000 Subject: [PATCH 3/4] Reformat --- src/prime_rl/trainer/models/layers/moe.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/prime_rl/trainer/models/layers/moe.py b/src/prime_rl/trainer/models/layers/moe.py index 09c21fcb46..687bba612b 100644 --- a/src/prime_rl/trainer/models/layers/moe.py +++ b/src/prime_rl/trainer/models/layers/moe.py @@ -177,7 +177,6 @@ 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: - 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() From 3508fea14d5106e9b838ce58f0bd7ff8f0a0a40b Mon Sep 17 00:00:00 2001 From: Mario Sieg Date: Wed, 26 Aug 2026 13:10:09 +0000 Subject: [PATCH 4/4] Don't use id() anymore --- .../trainer/models/layers/mxfp8_linear.py | 20 ++- tests/unit/train/models/test_mxfp8_linear.py | 119 ++++++++++++++++++ 2 files changed, 135 insertions(+), 4 deletions(-) create mode 100644 tests/unit/train/models/test_mxfp8_linear.py diff --git a/src/prime_rl/trainer/models/layers/mxfp8_linear.py b/src/prime_rl/trainer/models/layers/mxfp8_linear.py index a5942ef41f..2fdeb45e57 100644 --- a/src/prime_rl/trainer/models/layers/mxfp8_linear.py +++ b/src/prime_rl/trainer/models/layers/mxfp8_linear.py @@ -14,11 +14,13 @@ 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 - cache: dict[int, tuple[int, object]] = {} + orig_backward = tao_mx_linear.mx_mm.backward @staticmethod def forward( @@ -36,6 +38,7 @@ def forward( 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 @@ -55,8 +58,7 @@ def forward( kernel_preference, mxfp8_dim0_cast_kernel_choice=mxfp8_dim0_cast_kernel_choice, ) - cache_key = id(weight_hp) - cached = cache.get(cache_key) + 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: @@ -68,13 +70,23 @@ def forward( kernel_preference, mxfp8_dim0_cast_kernel_choice=mxfp8_dim0_cast_kernel_choice, ) - cache[cache_key] = (weight_hp._version, weight_mx_dim0) + 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 diff --git a/tests/unit/train/models/test_mxfp8_linear.py b/tests/unit/train/models/test_mxfp8_linear.py new file mode 100644 index 0000000000..a2755a5493 --- /dev/null +++ b/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