From c4ed8f5228ca00b887a672e4925934163f66f69e Mon Sep 17 00:00:00 2001 From: Deep Patel Date: Tue, 7 Jul 2026 16:29:02 -0700 Subject: [PATCH] Fixing activation sharding during fusion --- src/maxtext/layers/linears.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/maxtext/layers/linears.py b/src/maxtext/layers/linears.py index ca99776f38..efeb085d4a 100644 --- a/src/maxtext/layers/linears.py +++ b/src/maxtext/layers/linears.py @@ -506,6 +506,11 @@ def __call__( activations = [] if cfg.fused_mlp: x = self.wi(inputs, out_sharding=intermediate_sharding) + + # Enforce fused activations don't shard on num_activations axis + fused_intermediate_logical = self.intermediate_logical[:2] + (None,) + self.intermediate_logical[2:] + x = self._maybe_shard_with_logical(x, fused_intermediate_logical) + x = checkpoint_name(x, "mlpwi") for idx, act_fn in enumerate(self.activations): y = _convert_to_activation_function(act_fn)(x[:, :, idx, ...])