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, ...])