diff --git a/docs/guides/optimization/custom_mesh_and_rule.md b/docs/guides/optimization/custom_mesh_and_rule.md index 6ab2edb1ab..d2e8de88be 100644 --- a/docs/guides/optimization/custom_mesh_and_rule.md +++ b/docs/guides/optimization/custom_mesh_and_rule.md @@ -51,6 +51,12 @@ This rule utilizes the `data`, `stage`, `fsdp`, and `expert` axes. Its defining Similar in philosophy to `ep-as-cp.yml`, this configuration explicitly includes the `context` axis in the mesh layout alongside `data`, `stage`, `fsdp`, and `expert`. While context sharding is mapped to the `context` axis globally, within MoE components, this `context` axis dynamically shifts to perform expert parallelism instead of FSDP. This custom rule supports using CP and EP together. +### `cp-as-ep-dense-no-cp.yml` and `cp-as-ep-dense-fsdp.yml` + +Variants of `cp-as-ep.yml` that take the dense weights (`embed`, `embed_attn`, `q_lora`, `kv_lora`: attention and shared-expert projections, the MoE router and the shared-expert gate) off the `context` axis. Under `cp-as-ep` they are sharded over `fsdp` x `context` x `expert`, so every all-gather of them spans all devices and also synchronizes the context ranks, which do unequal attention work under causal context parallelism. `cp-as-ep-dense-no-cp.yml` shards them over `fsdp` x `expert` and `cp-as-ep-dense-fsdp.yml` over `fsdp` only, so those all-gathers run in independent groups of `fsdp` x `expert` (or `fsdp`) devices, at the cost of replicating the weights and their gradients `context`-way (`context` x `expert`-way). + +Both rules also map `embed_grad` to the token axes. Left to GSPMD, the gradient of a weight replicated over `context` (a partial sum over every token axis) is all-reduced at full size over all devices and then sliced. With `embed_grad` mapped, `linears.route_dense_weight_grad` reduce-scatters it over all devices instead, as under `cp-as-ep`, and then all-gathers the shard over the replicated axes. The GatedDeltaNet projections (`embed_gdn`), activations, the vocabulary weights and the routed experts are sharded exactly as in `cp-as-ep.yml`, and `RoutedMoE` keeps `("context", "expert")` as its expert-parallel axes. + ### `ep-as-dp.yml` Different with the rule in `base.yml`, this rule configures expert physical axis to function as data parallelism rather than FSDP. This removes the constraint where FSDPxEP is limited by specific model dimensions, particularly for small tensors such as attention projections. Ultimately, this change benefits large-scale training. diff --git a/src/maxtext/common/common_types.py b/src/maxtext/common/common_types.py index dc0fe956c3..a1617a5113 100644 --- a/src/maxtext/common/common_types.py +++ b/src/maxtext/common/common_types.py @@ -203,6 +203,8 @@ class CustomRule(enum.Enum): DEFAULT = "" PURE_FSDP = "pure-fsdp" CP_AS_EP = "cp-as-ep" # Support CP and EP together + CP_AS_EP_DENSE_NO_CP = "cp-as-ep-dense-no-cp" # cp-as-ep with the dense weights on fsdp x expert, not context + CP_AS_EP_DENSE_FSDP = "cp-as-ep-dense-fsdp" # cp-as-ep with the dense weights on fsdp only EP_AS_CP = "ep-as-cp" # Support EP only PIPELINE_LARGE_MOE = "pipeline-large-moe" FSDP_2D = "2d-fsdp" diff --git a/src/maxtext/configs/custom_mesh_and_rule/cp-as-ep-dense-fsdp.yml b/src/maxtext/configs/custom_mesh_and_rule/cp-as-ep-dense-fsdp.yml new file mode 100644 index 0000000000..c11f43cc36 --- /dev/null +++ b/src/maxtext/configs/custom_mesh_and_rule/cp-as-ep-dense-fsdp.yml @@ -0,0 +1,97 @@ +# Copyright 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. + +# Variant of cp-as-ep.yml whose only difference is that the dense weights are sharded on fsdp only. +# +# As cp-as-ep-dense-no-cp.yml, but the dense weights ('embed', 'embed_attn', 'q_lora', 'kv_lora') are also replicated +# over the expert axis: their all-gathers run in context x expert independent fsdp groups, and their gradients are +# reduce-scattered over all devices ('embed_grad') and then all-gathered over context x expert. Twice the HBM of +# cp-as-ep-dense-no-cp.yml for those weights. The GatedDeltaNet projections ('embed_gdn'), activations, the +# vocabulary weights and the routed experts are as in cp-as-ep, and RoutedMoE uses ("context", "expert") as its +# expert-parallel axes. +# +# This rule uses data, stage, FSDP, context, and expert. Context axis acts as expert +# parallelism in core MoE part (between EP all2all). +mesh_axes: ['data', 'stage', 'fsdp', 'context', 'expert'] +data_sharding: [['data', 'stage', 'fsdp', 'context', 'expert']] +logical_axis_rules: [ + # ========================================== + # Vocabulary Embedding + # ========================================== + # Vocab Activations + ['activation_embed_and_logits_batch', ['data', 'stage', 'fsdp', 'expert']], + # Vocabulary tiling flattens [batch, length] into one dimension. It must also be sharded + # over the length axis (context), or every context rank computes the logits of all the + # context shards. context is last because it is the minor part of the flattened index. + ['activation_embed_and_logits_batch_sequence', ['data', 'stage', 'fsdp', 'expert', 'context']], + # Vocab Weights + ['vocab', []], + ['embed_vocab', ['fsdp', 'context', 'expert']], + # ========================================== + # Attention + # ========================================== + # Attention Activations + ['activation_batch_attn', ['data', 'fsdp', 'expert']], + ['activation_heads', []], + ['activation_kv_heads', []], + ['activation_length_attn', ['context']], + ['activation_q_length', ['context']], + ['activation_kv_length', []], + ['activation_embed_attn', []], + ['activation_kv', []], + ['activation_kv_batch', ['data', 'fsdp', 'expert']], + ['activation_kv_head_dim', []], + # Attention Weights + ['heads', []], + ['q_heads', []], + ['kv_heads', []], + ['qkv', []], + ['kv', []], + ['kv_head_dim', []], + ['q_lora', ['fsdp']], + ["q_lora_up_proj", []], + ['kv_lora', ['fsdp']], + ["kv_lora_up_proj", []], + # ========================================== + # Mixture of Experts (MoE) + # ========================================== + # MoE Activations + ['activation_batch_moe', ['data', 'fsdp']], + ['activation_exp', ['context', 'expert']], + # MoE Weights + ['exp', ['context', 'expert']], + ['embed_moe', ['fsdp']], + ['embed_router', ['fsdp']], + # ========================================== + # Standard MLP / Dense Layers / Model Structure + # ========================================== + # Dense Activations + ['activation_mlp', []], + ['activation_batch', ['data', 'fsdp', 'expert']], + ['activation_length', ['context']], + ['activation_norm_length', ['context']], + ['activation_embed', []], + ['activation_stage', 'stage'], + # General Weights + ['mlp', []], + ['layers', 'stage'], + # Dense weights: not sharded on context, see the header. + ['embed', ['fsdp']], + ['embed_attn', ['fsdp']], + # GatedDeltaNet projections (models/qwen3.py) stay sharded as in cp-as-ep. + ['embed_gdn', ['fsdp', 'context', 'expert']], + # The dense weights' gradients are reduce-scattered over every token axis, as in cp-as-ep, then + # gathered over the axes the weights are replicated on (linears.route_dense_weight_grad). + ['embed_grad', ['fsdp', 'context', 'expert']], + ] diff --git a/src/maxtext/configs/custom_mesh_and_rule/cp-as-ep-dense-no-cp.yml b/src/maxtext/configs/custom_mesh_and_rule/cp-as-ep-dense-no-cp.yml new file mode 100644 index 0000000000..6387338ae3 --- /dev/null +++ b/src/maxtext/configs/custom_mesh_and_rule/cp-as-ep-dense-no-cp.yml @@ -0,0 +1,100 @@ +# Copyright 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. + +# Variant of cp-as-ep.yml whose only difference is that the dense weights are not sharded on the context axis. +# +# cp-as-ep shards the dense weights ('embed', 'embed_attn', 'q_lora', 'kv_lora': attention and shared-expert +# projections, the MoE router and the shared-expert gate) over fsdp x context x expert, so every all-gather of them +# spans all devices and also synchronizes the context ranks, which do unequal attention work. Here they are sharded +# over fsdp x expert and replicated over context: those all-gathers run in context-many independent fsdp x expert +# groups, at the cost of a context-way replication of the weights and their gradients. Their gradients are still +# reduce-scattered over all devices, as in cp-as-ep ('embed_grad'), and the shard is then all-gathered over context; +# left to GSPMD, the gradient would be all-reduced at full size over all devices instead. The GatedDeltaNet +# projections ('embed_gdn'), activations, the vocabulary weights and the routed experts ('exp' over context x expert) +# are sharded as in cp-as-ep, and RoutedMoE uses ("context", "expert") as its expert-parallel axes. +# +# This rule uses data, stage, FSDP, context, and expert. Context axis acts as expert +# parallelism in core MoE part (between EP all2all). +mesh_axes: ['data', 'stage', 'fsdp', 'context', 'expert'] +data_sharding: [['data', 'stage', 'fsdp', 'context', 'expert']] +logical_axis_rules: [ + # ========================================== + # Vocabulary Embedding + # ========================================== + # Vocab Activations + ['activation_embed_and_logits_batch', ['data', 'stage', 'fsdp', 'expert']], + # Vocabulary tiling flattens [batch, length] into one dimension. It must also be sharded + # over the length axis (context), or every context rank computes the logits of all the + # context shards. context is last because it is the minor part of the flattened index. + ['activation_embed_and_logits_batch_sequence', ['data', 'stage', 'fsdp', 'expert', 'context']], + # Vocab Weights + ['vocab', []], + ['embed_vocab', ['fsdp', 'context', 'expert']], + # ========================================== + # Attention + # ========================================== + # Attention Activations + ['activation_batch_attn', ['data', 'fsdp', 'expert']], + ['activation_heads', []], + ['activation_kv_heads', []], + ['activation_length_attn', ['context']], + ['activation_q_length', ['context']], + ['activation_kv_length', []], + ['activation_embed_attn', []], + ['activation_kv', []], + ['activation_kv_batch', ['data', 'fsdp', 'expert']], + ['activation_kv_head_dim', []], + # Attention Weights + ['heads', []], + ['q_heads', []], + ['kv_heads', []], + ['qkv', []], + ['kv', []], + ['kv_head_dim', []], + ['q_lora', ['fsdp', 'expert']], + ["q_lora_up_proj", []], + ['kv_lora', ['fsdp', 'expert']], + ["kv_lora_up_proj", []], + # ========================================== + # Mixture of Experts (MoE) + # ========================================== + # MoE Activations + ['activation_batch_moe', ['data', 'fsdp']], + ['activation_exp', ['context', 'expert']], + # MoE Weights + ['exp', ['context', 'expert']], + ['embed_moe', ['fsdp']], + ['embed_router', ['fsdp']], + # ========================================== + # Standard MLP / Dense Layers / Model Structure + # ========================================== + # Dense Activations + ['activation_mlp', []], + ['activation_batch', ['data', 'fsdp', 'expert']], + ['activation_length', ['context']], + ['activation_norm_length', ['context']], + ['activation_embed', []], + ['activation_stage', 'stage'], + # General Weights + ['mlp', []], + ['layers', 'stage'], + # Dense weights: not sharded on context, see the header. + ['embed', ['fsdp', 'expert']], + ['embed_attn', ['fsdp', 'expert']], + # GatedDeltaNet projections (models/qwen3.py) stay sharded as in cp-as-ep. + ['embed_gdn', ['fsdp', 'context', 'expert']], + # The dense weights' gradients are reduce-scattered over every token axis, as in cp-as-ep, then + # gathered over the axes the weights are replicated on (linears.route_dense_weight_grad). + ['embed_grad', ['fsdp', 'context', 'expert']], + ] diff --git a/src/maxtext/layers/linears.py b/src/maxtext/layers/linears.py index 6c02142ce1..5f65e321f6 100644 --- a/src/maxtext/layers/linears.py +++ b/src/maxtext/layers/linears.py @@ -36,6 +36,8 @@ from maxtext.layers.quantizations import AqtQuantization as Quant from maxtext.utils import max_logging from maxtext.utils import max_utils +from maxtext.utils.sharding import get_logical_axis_rules +from maxtext.utils.sharding import logical_to_mesh_axes from maxtext.utils.sharding import maybe_shard_with_logical from maxtext.utils.sharding import maybe_shard_with_name from maxtext.utils.sharding import get_physical_spec_without_axes @@ -143,6 +145,67 @@ def _compute_dot_general_nnx( ) +# Logical axis a custom rule set maps to every mesh axis the tokens are sharded over when it keeps the dense weights +# off some of them (cp-as-ep-dense-no-cp, cp-as-ep-dense-fsdp); see `route_dense_weight_grad`. +DENSE_GRAD_AXIS = "embed_grad" +_DENSE_EMBED_AXES = ("embed", "embed_attn") + + +@functools.partial(jax.custom_vjp, nondiff_argnums=(1, 2)) +def _reshard_cotangent(x, via, to): + """Identity whose cotangent is constrained to the PartitionSpec `via` and then to `to`.""" + del via, to + return x + + +def _reshard_cotangent_fwd(x, via, to): + del via, to + return x, None + + +def _reshard_cotangent_bwd(via, to, _, cotangent): + # Without the barrier, XLA:TPU folds the reduce-scatter into `via` and the all-gather into `to` back into one + # full-size all-reduce. + cotangent = jax.lax.optimization_barrier(jax.lax.with_sharding_constraint(cotangent, via)) + return (jax.lax.with_sharding_constraint(cotangent, to),) + + +_reshard_cotangent.defvjp(_reshard_cotangent_fwd, _reshard_cotangent_bwd) + + +def route_dense_weight_grad(kernel: Array, kernel_axes: tuple[None | str, ...]) -> Array: + """Reduce-scatters a dense weight's gradient over every token axis before resharding it to the weight's sharding. + + A weight's gradient is a partial sum over every mesh axis the tokens are sharded on. When the weight is sharded + over only some of them (cp-as-ep-dense-no-cp and cp-as-ep-dense-fsdp replicate the dense weights over context), + GSPMD reduces that partial sum with a full-size all-reduce over every device and keeps a slice of it. Under a rule + set that maps `embed_grad` to the token axes, the backward of every kernel with an `embed` or `embed_attn` + dimension instead constrains the gradient to the kernel's axes with that dimension on `embed_grad` (one + reduce-scatter over all of the token axes' devices, as under cp-as-ep), and then, across an optimization barrier, + to the kernel's own sharding (an all-gather of that shard over the axes the weight is replicated on). The forward + is unchanged, and under every rule set that does not map `embed_grad` the kernel is returned as is. + + Args: + kernel: The kernel, as the dot product consumes it. + kernel_axes: Its logical axes. + + Returns: + The kernel. + """ + rules = get_logical_axis_rules() + if not rules or not any(rule[0] == DENSE_GRAD_AXIS for rule in rules): + return kernel + if not any(isinstance(axis, str) and axis in _DENSE_EMBED_AXES for axis in kernel_axes): + return kernel + mesh = jax.sharding.get_abstract_mesh() + if mesh.empty: + raise ValueError(f"The logical axis rules map {DENSE_GRAD_AXIS!r}, which needs a mesh set with jax.set_mesh.") + via_axes = tuple(DENSE_GRAD_AXIS if isinstance(a, str) and a in _DENSE_EMBED_AXES else a for a in kernel_axes) + via = logical_to_mesh_axes(via_axes, mesh, rules=rules) + to = logical_to_mesh_axes(tuple(kernel_axes), mesh, rules=rules) + return _reshard_cotangent(kernel, via, to) + + class DenseGeneral(nnx.Module): """A linear transformation with flexible axes.""" @@ -444,6 +507,7 @@ def __call__( kernel = kernel[..., begin:end] kernel = self._maybe_two_stage_all_gather(kernel) + kernel = route_dense_weight_grad(kernel, self.kernel_axes) # out_sharding should be None for auto mesh axis if self.shard_mode != ShardMode.EXPLICIT: diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py index 176dea4880..21a6ca8621 100644 --- a/src/maxtext/layers/moe.py +++ b/src/maxtext/layers/moe.py @@ -888,6 +888,7 @@ def __call__(self, inputs: jax.Array, _initializing: bool = False) -> Tuple[jax. else: kernel = self.kernel[...] kernel = jnp.asarray(kernel, self.dtype) + kernel = linears.route_dense_weight_grad(kernel, self.kernel_axes) contract_ind = tuple(range(0, len(norm_axis))) output_sharding = ( @@ -1175,8 +1176,12 @@ def __init__( if self.mesh is not None and "dcp" in self.mesh.shape and self.mesh.shape["dcp"] > 1: expert_axes.append("dcp") self._expert_parallelism_name = tuple(expert_axes) - 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 + elif self.config.custom_mesh_and_rule in ( + ctypes.CustomRule.CP_AS_EP, + ctypes.CustomRule.CP_AS_EP_DENSE_NO_CP, + ctypes.CustomRule.CP_AS_EP_DENSE_FSDP, + ): + # when custom mesh and rule is cp-as-ep (or a variant of it), context axis is same with expert in MoE component self._expert_parallelism_name = ("context", "expert") else: self._expert_parallelism_name = "expert" diff --git a/src/maxtext/models/qwen3.py b/src/maxtext/models/qwen3.py index 9939c6f6b4..df80c10710 100644 --- a/src/maxtext/models/qwen3.py +++ b/src/maxtext/models/qwen3.py @@ -695,6 +695,9 @@ def __init__( self.cache = None # No cache for train mode or when inputs_shape not provided block_size = getattr(cfg, "weight_block_size", None) + # A custom rule set can shard the GatedDeltaNet projections apart from the attention projections by mapping + # `embed_gdn`; otherwise their embed dimension follows `embed_attn`. + gdn_embed_axis = "embed_gdn" if any(rule[0] == "embed_gdn" for rule in cfg.logical_axis_rules) else "embed_attn" # Submodule instantiations self.in_proj_qkvz = DenseGeneral( @@ -702,7 +705,7 @@ def __init__( out_features_shape=(self.key_dim * 2 + self.value_dim * 2), dtype=cfg.dtype, weight_dtype=cfg.weight_dtype, - kernel_axes=("embed_attn", "gdn_head"), + kernel_axes=(gdn_embed_axis, "gdn_head"), matmul_precision=cfg.matmul_precision, shard_mode=cfg.shard_mode, block_size=block_size, @@ -714,7 +717,7 @@ def __init__( out_features_shape=(self.num_v_heads * 2), dtype=cfg.dtype, weight_dtype=get_weight_dtype(cfg, "in_proj_ba"), - kernel_axes=("embed_attn", "gdn_head"), + kernel_axes=(gdn_embed_axis, "gdn_head"), matmul_precision=cfg.matmul_precision, shard_mode=cfg.shard_mode, quant=self.quant, @@ -767,7 +770,7 @@ def a_log_init(key, shape, dtype=jnp.float32): out_features_shape=(in_features,), dtype=cfg.dtype, weight_dtype=cfg.weight_dtype, - kernel_axes=("gdn_head", "embed_attn"), + kernel_axes=("gdn_head", gdn_embed_axis), matmul_precision=cfg.matmul_precision, shard_mode=cfg.shard_mode, block_size=block_size, diff --git a/tests/unit/cp_as_ep_dense_no_cp_test.py b/tests/unit/cp_as_ep_dense_no_cp_test.py new file mode 100644 index 0000000000..0e37ead1d6 --- /dev/null +++ b/tests/unit/cp_as_ep_dense_no_cp_test.py @@ -0,0 +1,413 @@ +# Copyright 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. +"""The cp-as-ep-dense-no-cp and cp-as-ep-dense-fsdp custom mesh rules on a CPU mesh. + +Both rules are cp-as-ep with the dense weights ('embed', 'embed_attn', 'q_lora', 'kv_lora') taken off the context axis: +cp-as-ep-dense-no-cp shards them over fsdp x expert, cp-as-ep-dense-fsdp over fsdp only, and both route the dense +weights' gradients through a reduce-scatter over every token axis ('embed_grad', linears.route_dense_weight_grad). +Only the sharding changes, so the model must compute the same function. + +The model is a tiny hybrid Qwen3.5 (GatedDeltaNet kernel path, gated full attention, routed and shared-expert MoE) on +16 CPU devices, fsdp2 x context4 x expert2, with sequence-sharded GatedDeltaNet context parallelism and ring of experts +over the expert-parallel group context x expert, as in the Qwen3.5-397B configuration. Every arm runs with the same +parameter values (the cp-as-ep initialization, placed on each rule's shardings). The checks run in a subprocess so the +forced device count takes effect before JAX initializes, with JAX's persistent compilation cache off. + + - Shardings: exactly the parameters whose embed dimension is 'embed' or 'embed_attn' move (attention and + shared-expert projections, the shared-expert gate and the router), each losing 'context' (and 'expert' under + cp-as-ep-dense-fsdp) and nothing else; the GatedDeltaNet projections ('embed_gdn') and every other parameter keep + their cp-as-ep sharding. + - Collectives: the compiled program all-gathers over fsdp x expert (4 devices) or fsdp (2) where cp-as-ep gathers + over all 16. + - Numerics: the loss is bit-identical to cp-as-ep, and every parameter gradient agrees to reduction order, + max|new - old| <= 2**-6 * max|old| per leaf: the activations and the gradient partial sums are bfloat16 (the ragged + sort supports no other activation dtype), so a changed reduction order moves a gradient by bfloat16 roundings. A + float32 variant (dense MLP instead of the MoE block) must agree to 1e-5. + - Positive controls, each of which must fail the numerics gate: + the rule without its RoutedMoE change, which keeps ("context", "expert") as the expert-parallel axes for these + rules as for cp-as-ep: expert parallelism over 'expert' alone while the expert weights are sharded over + context x expert; + one reshared weight (the first attention query kernel) scaled by 1 + 1e-3. +""" + +# pylint: disable=protected-access + +import collections +import contextlib +import dataclasses +import os +import re +import subprocess +import sys +from unittest import mock + +import jax +import jax.numpy as jnp +import numpy as np +import pytest + +from maxtext.configs import pyconfig +from maxtext.layers import moe +from maxtext.utils import maxtext_utils +from maxtext.utils import globals as maxtext_globals + +_BASE_CONFIG_PATH = os.path.join(maxtext_globals.MAXTEXT_CONFIGS_DIR, "base.yml") +_FSDP, _CONTEXT, _EXPERT = 2, 4, 2 +_NUM_DEVICES = _FSDP * _CONTEXT * _EXPERT +# Four context ranks of two 64-token GatedDeltaNet kernel chunks each. +_SEQ = 512 +# Reduction-order gate on every gradient leaf: max|new - old| <= tolerance * max|old|. The gradient partial sums are +# bfloat16, so a changed reduction order moves an element by bfloat16 roundings of the partial sums: four bfloat16 +# ulps (2**-6) at the leaf's largest magnitude, log2 of the 16 summands (1.5 measured). The float32 variant gets a +# float32 reduction-order tolerance. +_BF16_TOL = 2.0**-6 +_F32_TOL = 1e-5 +# The float32 variant: a dense MLP instead of the MoE block (the ragged sort only supports bfloat16 activations). +_F32_DENSE = {"num_experts": 1, "num_experts_per_tok": 1, "dtype": "float32"} +_REFERENCE_RULE = "cp-as-ep" +# rule -> the mesh axes of the dense weights' embed dimension. +_RULES = {"cp-as-ep-dense-no-cp": ("fsdp", "expert"), "cp-as-ep-dense-fsdp": ("fsdp",)} +_PASS_TOKEN = "CP_AS_EP_DENSE_NO_CP_CHECKS_PASSED" + + +@pytest.mark.cpu_only +def test_cp_as_ep_dense_no_cp_on_cpu_mesh(): + """Runs `main` below in a child process with 16 CPU devices.""" + env = os.environ.copy() + env["XLA_FLAGS"] = env.get("XLA_FLAGS", "") + f" --xla_force_host_platform_device_count={_NUM_DEVICES}" + env["JAX_PLATFORMS"] = "cpu" + result = subprocess.run([sys.executable, __file__], env=env, capture_output=True, text=True, check=False) + assert result.returncode == 0, f"stdout:\n{result.stdout}\nstderr:\n{result.stderr}" + assert _PASS_TOKEN in result.stdout, f"stdout:\n{result.stdout}\nstderr:\n{result.stderr}" + + +def _config(rule, **overrides): + """The tiny hybrid Qwen3.5 on the fsdp2 x context4 x expert2 mesh under `rule`, with `overrides` applied.""" + argv = [ + None, + _BASE_CONFIG_PATH, + "run_name=cp_as_ep_dense_no_cp_test", + "model_name=qwen3.5-397b-a17b", + "override_model_config=true", + "hardware=cpu", + f"custom_mesh_and_rule={rule}", + f"ici_fsdp_parallelism={_FSDP}", + f"ici_context_parallelism={_CONTEXT}", + f"ici_expert_parallelism={_EXPERT}", + "ici_tensor_parallelism=1", + "context_parallel_load_balance=False", + "attention=dot_product", + "use_gdn_kernel=true", + "gdn_cp_mode=seq", + "sparse_matmul=true", + "megablox=true", + # Tokamax gmm v2 is TPU-only; every arm runs the same GMM, so CPU uses megablox. + "use_tokamax_gmm=false", + "use_gmm_v2=false", + "use_ring_of_experts=true", + "use_ragged_sort=true", + "use_custom_sort_vjp=false", + "ragged_buffer_factor=2.0", + "num_moe_token_chunks=2", + "ring_of_experts_local_routing=true", + # The ragged sort only supports bfloat16 activations. + "dtype=bfloat16", + "weight_dtype=float32", + "scan_layers=True", + "remat_policy=custom", + "decoder_layer_input=device", + "context=device", + "packing=false", + "enable_checkpointing=false", + "enable_dropout=false", + "dataset_type=synthetic", + "base_num_decoder_layers=8", # two blocks of [3 GatedDeltaNet + 1 full attention] + "base_emb_dim=64", + "base_num_query_heads=4", + "base_num_kv_heads=2", + "head_dim=64", + "mrope_section=[3,3,2]", # rotary_dim / 2 = 64 * 0.25 / 2 + "base_mlp_dim=32", + "base_moe_mlp_dim=32", + "num_experts=16", + "num_experts_per_tok=2", + "shared_experts=1", + # The GatedDeltaNet kernel path needs head dims that are multiples of 128 and gdn_chunk_size=64. + "gdn_key_head_dim=128", + "gdn_value_head_dim=128", + "gdn_num_key_heads=2", + "gdn_num_value_heads=4", + "gdn_chunk_size=64", + "vocab_size=384", + f"max_target_length={_SEQ}", + "per_device_batch_size=0.5", + "use_multimodal=false", + "skip_jax_distributed_system=True", + "enable_tensorboard=False", + ] + argv = [a for a in argv if a is None or a.split("=", 1)[0] not in overrides] + return pyconfig.initialize(argv + [f"{key}={value}" for key, value in overrides.items()]) + + +def _inputs(cfg): + """Random tokens in one segment per row, and a fixed random cotangent on the logits. + + The cotangent makes the backward per-token sensitive: at random init the cross-entropy is nearly flat. + """ + batch = cfg.micro_batch_size_to_train_on + tokens = jax.random.randint(jax.random.PRNGKey(7), (batch, _SEQ + 1), 0, cfg.vocab_size) + ones = jnp.ones((batch, _SEQ), jnp.int32) + data = { + "inputs": tokens[:, :-1], + "targets": tokens[:, 1:], + "inputs_position": jnp.broadcast_to(jnp.arange(_SEQ, dtype=jnp.int32), (batch, _SEQ)), + "inputs_segmentation": ones, + "targets_segmentation": ones, + } + cotangent = jax.random.normal(jax.random.PRNGKey(3), (batch, _SEQ, cfg.vocab_size), jnp.float32) + return data, cotangent + + +@dataclasses.dataclass +class _Result: + loss: np.ndarray + grads: dict[str, np.ndarray] + specs: dict[str, tuple] # parameter -> the PartitionSpec of its sharding + values: dict[str, np.ndarray] # parameter -> its value + hlo: str + + +def _run(cfg, data, cotangent, values=None, control=None) -> _Result: + """Loss and its parameter gradients; `values` replaces the initial parameter values.""" + # pylint: disable=import-outside-toplevel + from flax import nnx + from flax.linen import partitioning as nn_partitioning + from maxtext.common.common_types import MODEL_MODE_TRAIN + from maxtext.utils import maxtext_utils_nnx, model_creation_utils + + mesh = maxtext_utils.get_mesh_from_config(cfg) + with control or contextlib.nullcontext(): + # As the trainer builds it: the parameters are created directly on the rule's shardings. + with nn_partitioning.axis_rules(cfg.logical_axis_rules): + create_model, abstract_model = model_creation_utils.create_nnx_abstract_model( + cfg, mesh, model_mode=MODEL_MODE_TRAIN, rng_key=jax.random.PRNGKey(0) + ) + model = maxtext_utils_nnx.create_nnx_sharded_model(abstract_model, create_model, mesh=mesh) + graphdef, params, rest = nnx.split(model, nnx.Param, ...) + keyed = jax.tree_util.tree_leaves_with_path(params) + specs = {jax.tree_util.keystr(path): tuple(leaf.sharding.spec) for path, leaf in keyed} + if values is not None: + assert values.keys() == specs.keys() + params = jax.tree_util.tree_map_with_path( + lambda path, leaf: jax.device_put(values[jax.tree_util.keystr(path)].astype(leaf.dtype), leaf.sharding), + params, + ) + + def loss_fn(params): + merged = nnx.merge(graphdef, params, rest, copy=True) + logits = merged( + decoder_input_tokens=data["inputs"], + decoder_positions=data["inputs_position"], + decoder_segment_ids=data["inputs_segmentation"], + enable_dropout=False, + decoder_target_tokens=data["targets"], + decoder_target_mask=data["targets_segmentation"], + ) + return jnp.sum(logits.astype(jnp.float32) * cotangent) + + with nn_partitioning.axis_rules(cfg.logical_axis_rules), jax.set_mesh(mesh): + compiled = jax.jit(jax.value_and_grad(loss_fn)).lower(params).compile() + loss, grads = compiled(params) + as_np = lambda tree: {jax.tree_util.keystr(p): np.asarray(x) for p, x in jax.tree_util.tree_leaves_with_path(tree)} + return _Result(np.asarray(loss), as_np(grads), specs, as_np(params), compiled.as_text()) + + +# --------------------------------------------------------------------------- +# Gates. +# --------------------------------------------------------------------------- + + +def _numerics(label, result, reference, tol=_BF16_TOL) -> bool: + """Bit-identical loss and every gradient leaf within reduction order of the reference; prints the report.""" + assert result.grads.keys() == reference.grads.keys() + loss_equal = bool(np.array_equal(result.loss, reference.loss)) + worst, worst_name, worst_l2, not_bit_equal, grads_ok = 0.0, "", 0.0, 0, True + for name, ref in reference.grads.items(): + new, ref = result.grads[name].astype(np.float64), ref.astype(np.float64) + diff, scale = float(np.max(np.abs(new - ref))), float(np.max(np.abs(ref))) + # Every comparison is False on NaN, so a NaN fails. + grads_ok &= bool(diff <= tol * scale) and bool(np.isfinite(new).all()) + rel = diff / scale if scale > 0 else diff + if rel > worst or not worst_name: + worst, worst_name = rel, name + worst_l2 = float(np.linalg.norm(new - ref)) / max(float(np.linalg.norm(ref)), 1e-30) + not_bit_equal += int(not np.array_equal(new, ref)) + passed = loss_equal and grads_ok + print( + f"{label}: numerics {'PASS' if passed else 'FAIL'}: loss bit-identical={loss_equal}" + f" (d={float(result.loss) - float(reference.loss):.6e}); gradients within {tol:.3g} x max|ref|={grads_ok}," + f" worst max|d|/max|ref|={worst:.3e} (||d||/||ref||={worst_l2:.3e}) at {worst_name}," + f" {not_bit_equal} of {len(reference.grads)} leaves not bit-identical", + flush=True, + ) + return passed + + +def _mesh_axes(entry): + """The mesh axes of one PartitionSpec entry, as a tuple.""" + if entry is None: + return () + return tuple(entry) if isinstance(entry, (tuple, list)) else (entry,) + + +def _sharding_moves(rule, result, reference) -> bool: + """Exactly the dense weights move, each from cp-as-ep's fsdp x context x expert to the rule's axes.""" + dense_axes, moved, wrong = _RULES[rule], [], [] + for name, ref_spec in reference.specs.items(): + spec = result.specs[name] + if spec == ref_spec: + continue + moved.append(name) + # The only change allowed: one dimension sharded over fsdp x context x expert becomes the rule's dense axes. + changes = [(_mesh_axes(a), _mesh_axes(b)) for a, b in zip(ref_spec, spec) if a != b] + if len(spec) != len(ref_spec) or changes != [(("fsdp", "context", "expert"), dense_axes)]: + wrong.append((name, ref_spec, spec)) + required = ( + "['query']", + "['key']", + "['value']", + "['out']", + "shared_expert']['wi_0", + "shared_expert']['wo", + "shared_expert_gate", + "routed_experts']['gate", + ) + missing = [r for r in required if not any(r in name for name in moved)] + kept = ("token_embedder", "routed_experts']['wi", "routed_experts']['wo", "in_proj_qkvz", "in_proj_ba", "out_proj") + forbidden = [n for n in moved if any(k in n for k in kept)] + print(f"{rule}: {len(moved)} of {len(reference.specs)} parameters change sharding:", flush=True) + for name in moved: + print(f" {name}: {reference.specs[name]} -> {result.specs[name]}", flush=True) + if wrong or missing or forbidden: + print(f"{rule}: unexpected sharding changes {wrong}; required but unmoved {missing}; moved {forbidden}", flush=True) + return not (wrong or missing or forbidden) + + +# The result shape is one token, or a parenthesised tuple of them when XLA combined several collectives into one. +_COLLECTIVE = re.compile(r"=\s*(?:\([^=]*?\)|\S+)\s+(all-gather|reduce-scatter|all-reduce|all-to-all)(?:-start)?\(") + + +def _group_size(line) -> int: + """The replica-group size of one HLO collective instruction (explicit, iota or mesh form).""" + if "replica_groups={}" in line: # one group of every device + return _NUM_DEVICES + match = re.search(r"replica_groups=\{(\{[^}]*\})", line) + if match: + return len(match.group(1).strip("{}").split(",")) + match = re.search(r"replica_groups=\[(\d+),(\d+)\]<=", line) + if match: + return int(match.group(2)) + match = re.search(r"replica_groups=mesh\[([^\]]*)\](?:,\s*device_ids=\((?:[^()]|\([^()]*\))*\))?\s*\{([^}]*)\}", line) + if match: + sizes = dict(re.findall(r"'(\w+)'=(\d+)", match.group(1))) + return int(np.prod([int(sizes[a]) for a in re.findall(r"'(\w+)'", match.group(2))])) + return -1 + + +def _collectives(hlo_text) -> collections.Counter: + """(collective, replica-group size) -> count, over the compiled program.""" + counts = collections.Counter() + for line in hlo_text.splitlines(): + match = _COLLECTIVE.search(line) + if match: + counts[(match.group(1), _group_size(line))] += 1 + return counts + + +# --------------------------------------------------------------------------- +# Positive controls. +# --------------------------------------------------------------------------- + + +@contextlib.contextmanager +def _moe_without_rule_change(): + """RoutedMoE as before this change: expert parallelism over 'expert' alone under the new rules.""" + real = moe.RoutedMoE.__init__ + + def init(self, *args, **kwargs): + real(self, *args, **kwargs) + assert self._expert_parallelism_name == ("context", "expert"), self._expert_parallelism_name + self._expert_parallelism_name = "expert" + + with mock.patch.object(moe.RoutedMoE, "__init__", init): + yield + + +def main(): + assert jax.device_count() == _NUM_DEVICES, jax.devices() + jax.config.update("jax_enable_compilation_cache", False) + reference_cfg = _config(_REFERENCE_RULE) + data, cotangent = _inputs(reference_cfg) + reference = _run(reference_cfg, data, cotangent) + print(f"{_REFERENCE_RULE}: collectives {sorted(_collectives(reference.hlo).items())}", flush=True) + checks, results = {}, {} + for rule, dense_axes in _RULES.items(): + result = results[rule] = _run(_config(rule), data, cotangent, values=reference.values) + gathers = _collectives(result.hlo) + print(f"{rule}: collectives {sorted(gathers.items())}", flush=True) + dense_group = int(np.prod([{"fsdp": _FSDP, "expert": _EXPERT}[a] for a in dense_axes])) + checks[f"{rule}: exactly the dense weights move, off context"] = _sharding_moves(rule, result, reference) + # Every all-gather over all devices that goes away comes back over the dense weights' smaller group. + ref_gathers = _collectives(reference.hlo) + gone = ref_gathers[("all-gather", _NUM_DEVICES)] - gathers[("all-gather", _NUM_DEVICES)] + checks[f"{rule}: {_NUM_DEVICES}-device all-gathers become {dense_group}-device ones"] = ( + 0 < gone <= gathers[("all-gather", dense_group)] - ref_gathers[("all-gather", dense_group)] + # No all-gather this test cannot size beyond those of cp-as-ep. + and gathers[("all-gather", -1)] <= ref_gathers[("all-gather", -1)] + ) + checks[f"{rule}: numerics"] = _numerics(rule, result, reference) + + # The bfloat16 differences are reduction order: with float32 activations they shrink to float32 roundings. + reference32 = _run(_config(_REFERENCE_RULE, **_F32_DENSE), data, cotangent) + result32 = _run(_config("cp-as-ep-dense-no-cp", **_F32_DENSE), data, cotangent, values=reference32.values) + checks["float32 dense variant: cp-as-ep-dense-no-cp numerics"] = _numerics( + "float32 dense cp-as-ep-dense-no-cp", result32, reference32, tol=_F32_TOL + ) + + # Positive controls. Each must fail the numerics gate (or fail to run); `controls` records whether it was caught. + controls = {} + rule = "cp-as-ep-dense-no-cp" + try: + wrong = _run(_config(rule), data, cotangent, values=reference.values, control=_moe_without_rule_change()) + controls["the rule without its RoutedMoE change"] = not _numerics("control: no RoutedMoE change", wrong, reference) + except Exception as error: # pylint: disable=broad-exception-caught + print(f"control: no RoutedMoE change: raised {type(error).__name__}: {str(error)[:400]}", flush=True) + controls["the rule without its RoutedMoE change"] = True + query = next(name for name in reference.values if re.search(r"\['query'\].*\['kernel'\]", name)) + perturbed = dict(reference.values, **{query: reference.values[query] * np.float32(1 + 1e-3)}) + wrong = _run(_config(rule), data, cotangent, values=perturbed) + controls[f"{query} scaled by 1 + 1e-3"] = not _numerics("control: perturbed query", wrong, reference) + + for name, ok in {**checks, **{f"positive control caught: {k}": v for k, v in controls.items()}}.items(): + print(f" [{'ok' if ok else 'FAIL'}] {name}", flush=True) + failed = [name for name, ok in checks.items() if not ok] + undetected = [name for name, caught in controls.items() if not caught] + if failed or undetected: + print(f"checks that failed: {failed}; positive controls that were NOT detected: {undetected}", flush=True) + sys.exit(1) + print(_PASS_TOKEN, flush=True) + + +if __name__ == "__main__": + main()