diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml index a4cf6877e3..293fdea4a6 100644 --- a/src/maxtext/configs/base.yml +++ b/src/maxtext/configs/base.yml @@ -284,6 +284,8 @@ moe_fsdp_use_two_stage_all_gather: false shard_exp_on_fsdp: false # use fsdp and fsdp_transpose axes for sharding the moe weights use_2d_fsdp_sharding: false +# Use the tensor axis as expert parallelism in the default sparse MoE path. +moe_late_tp_combine: false # deepseek moe first_num_dense_layers: 0 # number of initial dense layers in the model @@ -583,6 +585,7 @@ logical_axis_rules: [ ['activation_exp', ['expert']], # MoE Weights ['exp', 'expert'], + ['exp_tp', ['expert', 'tensor']], ['mlp_moe', ['fsdp_transpose', 'tensor', 'tensor_sequence', 'autoregressive']], ['embed_moe', ['fsdp', 'fsdp_transpose', 'context', 'context_usp_ulysses']], ['embed_moe', ['fsdp', 'context', 'context_usp_ulysses']], diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index 1f8659f7e0..0aa778e9c7 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -948,6 +948,11 @@ class MoEGeneral(BaseModel): False, description="Use `fsdp` and `fsdp_transpose` axes for 2D FSDP sharding.", ) + moe_late_tp_combine: bool = Field( + False, + description="Use the tensor axis as expert parallelism in the default sparse MoE path, moving " + "tensor-axis communication outside expert-expanded activations.", + ) norm_topk_prob: bool = Field( False, description="Enable top-k probability normalization for router weights (Qwen3-specific).", diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py index eccaf295c6..8fc3e007ea 100644 --- a/src/maxtext/layers/moe.py +++ b/src/maxtext/layers/moe.py @@ -459,6 +459,24 @@ def __init__( self.config.emb_dim if self.config.moe_expert_input_dim <= 0 else self.config.moe_expert_input_dim ) + self._late_tp_combine = ( + self.config.moe_late_tp_combine + and self.config.sparse_matmul + and self.mesh.shape.get("tensor", 1) > 1 + and self.config.attention not in ("vllm_rpa", "vllm_batched_rpa") + and self.config.custom_mesh_and_rule != ctypes.CustomRule.CP_AS_EP + and not self.config.shard_exp_on_fsdp + and not self.config.use_2d_fsdp_sharding + and not self.config.use_batch_split_schedule + ) + if self._late_tp_combine: + num_moe_shards = self.mesh.shape.get("expert", 1) * self.mesh.shape["tensor"] + if self.config.num_experts % num_moe_shards != 0: + raise ValueError( + f"moe_late_tp_combine needs num_experts ({self.config.num_experts}) divisible by " + f"expert*tensor parallelism ({num_moe_shards})." + ) + if self.config.shard_exp_on_fsdp: # special sharding for dsv3 self.wi_kernel_axes = ("embed_moe", None, "mlp_moe") @@ -468,6 +486,9 @@ def __init__( self.wo_kernel_axes = ("embed_moe", "mlp_moe", None) elif self.config.use_batch_split_schedule: self.wi_kernel_axes, self.wo_kernel_axes = get_batchsplit_init_kernel_axes() + elif self._late_tp_combine: + self.wi_kernel_axes = ("exp_tp", "embed_moe", None) + self.wo_kernel_axes = ("exp_tp", None, "embed_moe") else: self.wi_kernel_axes = ("exp", "embed_moe", "mlp_moe") self.wo_kernel_axes = ("exp", "mlp_moe", "embed_moe") @@ -483,6 +504,9 @@ def __init__( elif self.config.custom_mesh_and_rule == ctypes.CustomRule.CP_AS_EP: # when custom mesh and rule is cp-as-ep, context axis is same with expert in MoE component self._expert_parallelism_name = ("context", "expert") + elif self._late_tp_combine: + # Keep axis order aligned with the exp_tp partition spec. + self._expert_parallelism_name = ("expert", "tensor") else: self._expert_parallelism_name = "expert" @@ -696,6 +720,8 @@ def get_expert_parallelism_size(self): return self.mesh.shape.get(self._expert_parallelism_name, 1) def get_tensor_parallelism_size(self): + if self._late_tp_combine: + return 1 if isinstance(self._tensor_parallelism_name, tuple): size = 1 for axis in self._tensor_parallelism_name: @@ -1618,21 +1644,31 @@ def get_routed_moe_shardings(is_batch_sharded_by_expert, has_input_ids): else: batch_logical_axis = "decode_batch_moe" - input_partition_pspec = self._logical_to_mesh_axes((batch_logical_axis, "activation_norm_length", None)) - w0_bias_pspec = self._logical_to_mesh_axes(("exp", "activation_mlp")) - w1_bias_pspec = self._logical_to_mesh_axes(("exp", "activation_mlp")) - wo_bias_pspec = self._logical_to_mesh_axes(("exp", "activation_embed")) + def token_split_pspec(logical_axes): + spec = list(self._logical_to_mesh_axes(logical_axes)) + seq = spec[1] + seq = () if seq is None else ((seq,) if isinstance(seq, str) else tuple(seq)) + spec[1] = seq + ("tensor",) + return P(*spec) + + token_pspec = token_split_pspec if self._late_tp_combine else self._logical_to_mesh_axes + exp_axis = "exp_tp" if self._late_tp_combine else "exp" + + input_partition_pspec = token_pspec((batch_logical_axis, "activation_norm_length", None)) + w0_bias_pspec = self._logical_to_mesh_axes((exp_axis, "activation_mlp")) + w1_bias_pspec = self._logical_to_mesh_axes((exp_axis, "activation_mlp")) + wo_bias_pspec = self._logical_to_mesh_axes((exp_axis, "activation_embed")) - gate_logits_pspec = self._logical_to_mesh_axes((batch_logical_axis, "activation_norm_length", None)) + gate_logits_pspec = token_pspec((batch_logical_axis, "activation_norm_length", None)) # NOTE: deepseek2 has a different pattern if self.config.model_name.startswith(("deepseek3", "deepseek4")): - pre_bias_logits_pspec = self._logical_to_mesh_axes((batch_logical_axis, "activation_norm_length", None)) + pre_bias_logits_pspec = token_pspec((batch_logical_axis, "activation_norm_length", None)) else: # pre_bias_logits is None for non-deepseek3/4 models, including deepseek2 pre_bias_logits_pspec = None if has_input_ids: - decoder_tokens_pspec = self._logical_to_mesh_axes((batch_logical_axis, "activation_norm_length")) + decoder_tokens_pspec = token_pspec((batch_logical_axis, "activation_norm_length")) else: decoder_tokens_pspec = None @@ -1654,6 +1690,10 @@ def get_routed_moe_shardings(is_batch_sharded_by_expert, has_input_ids): w0_pspec = self._logical_to_mesh_axes((None, "mlp_no_fsdp", None)) w1_pspec = self._logical_to_mesh_axes((None, "mlp_no_fsdp", None)) wo_pspec = self._logical_to_mesh_axes((None, "mlp_no_fsdp", None)) + elif self._late_tp_combine: + w0_pspec = self._logical_to_mesh_axes(("exp_tp", None, None)) + w1_pspec = self._logical_to_mesh_axes(("exp_tp", None, None)) + wo_pspec = self._logical_to_mesh_axes(("exp_tp", None, None)) else: # These are the main shardings used by default - they use funky rules to AG over FSDP. w0_pspec = self._logical_to_mesh_axes(("exp", None, "mlp_no_fsdp")) @@ -2292,7 +2332,7 @@ def _moe_body( P(), # Replicate the input key ), out_specs=( - output_pspec, + input_partition_pspec if self._late_tp_combine else output_pspec, P(), # Handle None or replicate the output P(), # Handle None or replicate the output ), diff --git a/tests/unit/moe_late_tp_combine_test.py b/tests/unit/moe_late_tp_combine_test.py new file mode 100644 index 0000000000..acdc3a6755 --- /dev/null +++ b/tests/unit/moe_late_tp_combine_test.py @@ -0,0 +1,205 @@ +# Copyright 2023–2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Tests for the `moe_late_tp_combine` routed-MoE sharding flag.""" + +import unittest + +from flax import nnx +from flax.linen import partitioning as nn_partitioning +import jax +from jax.extend.core import ClosedJaxpr, Jaxpr +import jax.numpy as jnp +from jax.sharding import Mesh +from jax.sharding import PartitionSpec as P +from maxtext.configs import pyconfig +from maxtext.layers import moe +from maxtext.layers.initializers import nd_dense_init +from maxtext.utils import maxtext_utils +from maxtext.utils.sharding import logical_to_mesh_axes +from tests.utils.test_helpers import get_test_config_path + +_COLLECTIVE_PRIMITIVES = ( + "all_gather", + "all_to_all", + "ppermute", + "psum", + "psum_scatter", + "ragged_all_to_all", + "reduce_scatter", +) + + +def _build_config(late_tp_combine, use_ring_of_experts): + return pyconfig.initialize( + [None, get_test_config_path()], + run_name=f"moe_late_tp_combine_{late_tp_combine}_roe_{use_ring_of_experts}", + enable_checkpointing=False, + model_name="mixtral-8x7b", + override_model_config=True, + base_emb_dim=64, + base_mlp_dim=64, + base_moe_mlp_dim=64, + dtype="float32", + weight_dtype="float32", + megablox=True, + sparse_matmul=True, + per_device_batch_size=2, # TODO(b/450900273): sharding error if pdbs=1 + max_target_length=64, + ici_expert_parallelism=2, + ici_tensor_parallelism=4, + use_ring_of_experts=use_ring_of_experts, + moe_late_tp_combine=late_tp_combine, + float32_gate_logits=True, + ) + + +def _build_model(cfg, mesh): + return moe.get_routed_moe( + name="MoeBlock", + config=cfg, + num_experts=cfg.num_experts, + num_experts_per_tok=cfg.num_experts_per_tok, + mesh=mesh, + kernel_init=nd_dense_init(1.0, "fan_in", "truncated_normal"), + kernel_axes=("embed", "mlp"), + intermediate_dim=cfg.mlp_dim, + dtype=cfg.dtype, + ) + + +def _zeroed_params(model, hidden_states): + """Builds params without executing ragged_all_to_all, which XLA:CPU does not lower.""" + + def init(): + return model.init({"params": jax.random.PRNGKey(0), "dropout": jax.random.PRNGKey(0)}, hidden_states) + + return jax.tree.map(lambda s: jnp.zeros(s.shape, s.dtype), jax.eval_shape(init)["params"]) + + +def _scan_jaxpr(jaxpr, collectives, shard_map_out_specs): + """Collects (primitive, axis_names) for every collective and every shard_map output spec.""" + for eqn in jaxpr.eqns: + if eqn.primitive.name in _COLLECTIVE_PRIMITIVES: + axis_name = eqn.params.get("axis_name", ()) + collectives.append((eqn.primitive.name, (axis_name,) if isinstance(axis_name, str) else tuple(axis_name))) + elif eqn.primitive.name == "shard_map": + shard_map_out_specs.append(eqn.params["out_specs"][0]) + for param in jax.tree_util.tree_leaves(eqn.params, is_leaf=lambda p: isinstance(p, (Jaxpr, ClosedJaxpr))): + param = param.jaxpr if isinstance(param, ClosedJaxpr) else param + if isinstance(param, Jaxpr): + _scan_jaxpr(param, collectives, shard_map_out_specs) + + +class MoeLateTpCombineTest(unittest.TestCase): + + def setUp(self): + super().setUp() + if jax.device_count() != 8: + self.skipTest(f"needs 8 devices for expert=2 x tensor=4, got {jax.device_count()}") + + def _loss_and_grad(self, cfg, variables, hidden_states): + devices_array = maxtext_utils.create_device_mesh(cfg) + mesh = Mesh(devices_array, cfg.mesh_axes) + model = _build_model(cfg, mesh) + + def loss_fn(params, x): + out, lb_loss, _ = model.apply({"params": params}, x) + loss = jnp.mean(out.astype(jnp.float32) ** 2) + if lb_loss is not None: + loss = loss + lb_loss.astype(jnp.float32) + return loss + + with jax.set_mesh(mesh), nn_partitioning.axis_rules(cfg.logical_axis_rules): + if variables is None: + variables = model.init({"params": jax.random.PRNGKey(0), "dropout": jax.random.PRNGKey(0)}, hidden_states) + out = jax.jit(jax.value_and_grad(loss_fn, argnums=(0, 1)))(variables["params"], hidden_states) + return variables, out + + def test_parity_off_vs_on(self): + """Checks loss and gradient parity through the CPU-compatible ring-of-experts path.""" + cfg_off = _build_config(late_tp_combine=False, use_ring_of_experts=True) + cfg_on = _build_config(late_tp_combine=True, use_ring_of_experts=True) + hidden_states = jax.random.uniform( + jax.random.PRNGKey(2345), + (int(cfg_off.per_device_batch_size) * jax.device_count(), cfg_off.max_target_length, cfg_off.base_emb_dim), + dtype=jnp.float32, + ) + + variables, (loss_off, grads_off) = self._loss_and_grad(cfg_off, None, hidden_states) + _, (loss_on, grads_on) = self._loss_and_grad(cfg_on, variables, hidden_states) + + self.assertAlmostEqual(float(loss_on), float(loss_off), delta=1e-5 * max(1.0, abs(float(loss_off)))) + leaves_off, treedef = jax.tree_util.tree_flatten(grads_off) + leaves_on, treedef_on = jax.tree_util.tree_flatten(grads_on) + self.assertEqual(treedef, treedef_on) + for i, (g_off, g_on) in enumerate(zip(leaves_off, leaves_on)): + self.assertEqual(g_off.shape, g_on.shape, f"grad shape mismatch at leaf {i}") + diff = float(jnp.max(jnp.abs(g_on - g_off))) + scale = float(jnp.max(jnp.abs(g_off))) or 1.0 + self.assertLess(diff / scale, 1e-4, f"grad mismatch at leaf {i}: max abs diff {diff}") + + def test_no_tensor_axis_collectives_in_moe_body(self): + """Checks production sharding by tracing the ragged_all_to_all path.""" + hidden_states = jnp.zeros((2 * jax.device_count(), 64, 64), dtype=jnp.float32) + collectives, out_specs = {}, {} + for late_tp_combine in (False, True): + cfg = _build_config(late_tp_combine=late_tp_combine, use_ring_of_experts=False) + mesh = Mesh(maxtext_utils.create_device_mesh(cfg), cfg.mesh_axes) + model = _build_model(cfg, mesh) + with jax.set_mesh(mesh), nn_partitioning.axis_rules(cfg.logical_axis_rules): + params = _zeroed_params(model, hidden_states) + jaxpr = jax.make_jaxpr(model.apply)({"params": params}, hidden_states) + collectives[late_tp_combine], out_specs[late_tp_combine] = [], [] + _scan_jaxpr(jaxpr.jaxpr, collectives[late_tp_combine], out_specs[late_tp_combine]) + + tensor_only_off = [c for c in collectives[False] if c[1] == ("tensor",)] + tensor_only_on = [c for c in collectives[True] if c[1] == ("tensor",)] + self.assertTrue(tensor_only_off, "expected a tensor-axis collective in the MoE body with the flag off") + self.assertFalse(tensor_only_on, f"tensor-axis collectives remain with the flag on: {tensor_only_on}") + + self.assertTrue( + any(c[1] == ("expert", "tensor") for c in collectives[True]), + f"expected expert*tensor dispatch collectives, got {collectives[True]}", + ) + + self.assertEqual(out_specs[False], [P("expert", None, "tensor")]) + self.assertEqual(out_specs[True], [P("expert", "tensor", None)]) + + def test_expert_weights_sharded_over_expert_and_tensor(self): + cfg = _build_config(late_tp_combine=True, use_ring_of_experts=False) + mesh = Mesh(maxtext_utils.create_device_mesh(cfg), cfg.mesh_axes) + with jax.set_mesh(mesh), nn_partitioning.axis_rules(cfg.logical_axis_rules): + block = moe.RoutedMoE( + config=cfg, + num_experts=cfg.num_experts, + num_experts_per_tok=cfg.num_experts_per_tok, + mesh=mesh, + kernel_init=nd_dense_init(1.0, "fan_in", "truncated_normal"), + kernel_axes=("embed", "mlp"), + rngs=nnx.Rngs(params=0), + intermediate_dim=cfg.mlp_dim, + dtype=cfg.dtype, + ) + self.assertEqual(block.wi_kernel_axes, ("exp_tp", "embed_moe", None)) + self.assertEqual(block.wo_kernel_axes, ("exp_tp", None, "embed_moe")) + self.assertEqual(block.get_tensor_parallelism_size(), 1) + self.assertEqual(block.get_expert_parallelism_size(), 8) + self.assertEqual( + logical_to_mesh_axes(("exp_tp", None, None), mesh=mesh, rules=cfg.logical_axis_rules)[0], + ("expert", "tensor"), + ) + + +if __name__ == "__main__": + unittest.main()