diff --git a/src/maxtext/models/qwen3.py b/src/maxtext/models/qwen3.py index 56395c6189..bad90dbc10 100644 --- a/src/maxtext/models/qwen3.py +++ b/src/maxtext/models/qwen3.py @@ -25,7 +25,7 @@ import jax.nn from jax import lax from jax.ad_checkpoint import checkpoint_name -from jax.sharding import Mesh +from jax.sharding import Mesh, PartitionSpec as P import jax.numpy as jnp from flax import linen as nn @@ -184,23 +184,27 @@ def scan_body(prev_state, x): return core_attn_out, final_state if output_final_state else None -def jax_chunk_gated_delta_rule( +def _gdn_chunk_prepare( query: Array, key: Array, value: Array, g: Array, beta: Array, - chunk_size: int = 64, - initial_state: None | Array = None, - use_qk_norm_in_gdn: bool = False, - compute_dtype: jnp.dtype = jnp.bfloat16, -) -> tuple[Array, None | Array]: - """Optimized JAX implementation of Gated Delta Rule.""" - # ========================================================================= - # STAGE 1: PREPARATION & PADDING - # ========================================================================= - initial_dtype = query.dtype - + chunk_size: int, + use_qk_norm_in_gdn: bool, + compute_dtype: jnp.dtype, +) -> tuple[tuple[Array, ...], tuple[int, ...]]: + """Builds the per-chunk WY factors consumed by the inter-chunk recurrence. + + This is STAGE 1 + STAGE 2 of the chunked gated delta rule: normalize, pad to a + whole number of chunks, and solve for the WY representation of each chunk. The + returned `xs` are transposed so the leading axis is the chunk index, ready for + `lax.scan`. + + Returns: + A pair `(xs, meta)` where `xs = (w, u, q, k, g_cumsum)` and + `meta = (batch, seq_len, num_heads, k_dim, v_dim, pad_len)`. + """ if use_qk_norm_in_gdn: query = l2norm(query, dim=-1, eps=1e-6) key = l2norm(key, dim=-1, eps=1e-6) @@ -283,40 +287,51 @@ def to_chunk_scalar(x): w_chunks = jnp.matmul(A, k_beta_g, precision=jax.lax.Precision.HIGHEST) w_chunks = w_chunks.astype(compute_dtype) - # ========================================================================= - # STAGE 3: INTER-CHUNK RECURRENCE (Scan) - # ========================================================================= + # Transpose so the chunk index leads, ready for `lax.scan`. scan_perm_vec = (1, 0, 2, 3, 4) scan_perm_scl = (1, 0, 2, 3) - w_scan = w_chunks.transpose(scan_perm_vec) - u_scan = u_chunks.transpose(scan_perm_vec) - k_scan = k_c.transpose(scan_perm_vec) - q_scan = q_c.transpose(scan_perm_vec) - g_scan = g_cumsum.transpose(scan_perm_scl) + xs = ( + w_chunks.transpose(scan_perm_vec), + u_chunks.transpose(scan_perm_vec), + q_c.transpose(scan_perm_vec), + k_c.transpose(scan_perm_vec), + g_cumsum.transpose(scan_perm_scl), + ) + return xs, (B, seq_len, H, K_dim, V_dim, pad_len) - if initial_state is None: - h_init = jnp.zeros((B, H, K_dim, V_dim), dtype=jnp.float32) - else: - h_init = initial_state.astype(jnp.float32) - xs = (w_scan, u_scan, q_scan, k_scan, g_scan) +def _gdn_scan_body(carry, args, *, chunk_size: int, compute_output: bool, propagate_transition: bool): + """One inter-chunk step of the gated delta rule (STAGE 3). - def scan_body(h, args): - w, u, q, k, g = args - prec = jax.lax.Precision.HIGHEST + The chunk recurrence is affine in the incoming state `h`: - # --- Output Computation --- + h_out = (gamma * I - k_tilde^T @ w) @ h + k_tilde^T @ u == M @ h + U + o = (q_g - attn_i @ w) @ h + attn_i @ u + + `propagate_transition` additionally carries the composed transition `M` for all + chunks seen so far, which is what lets context parallelism combine the per-rank + results with a prefix scan (see `jax_chunk_gated_delta_rule_cp`). + """ + if propagate_transition: + h, m = carry + else: + h, m = carry, None + w, u, q, k, g = args + prec = jax.lax.Precision.HIGHEST + + # --- Delta Rule Subtraction (v_prime and v_new) --- + # w serves as k_cumdecay, u serves as value_intra + v_prime = jnp.matmul(w.astype(jnp.float32), h, precision=prec) + v_new = u.astype(jnp.float32) - v_prime + + o_c = None + if compute_output: # 1. Inter-chunk: q(dtype) * exp(g)(f32) -> f32 q_g = q.astype(jnp.float32) * jnp.exp(g)[..., None] attn_inter = jnp.matmul(q_g, h, precision=prec) - # 2. Delta Rule Subtraction (v_prime and v_new) - # w serves as k_cumdecay, u serves as value_intra - v_prime = jnp.matmul(w.astype(jnp.float32), h, precision=prec) - v_new = u.astype(jnp.float32) - v_prime - - # 3. Intra-chunk: q(dtype) @ k(dtype) -> f32 + # 2. Intra-chunk: q(dtype) @ k(dtype) -> f32 attn = jnp.matmul(q, k.swapaxes(-1, -2), precision=prec) attn = attn.astype(jnp.float32) @@ -331,36 +346,146 @@ def scan_body(h, args): # Note: We do NOT multiply attn_i by beta here. The Delta rule mathematically # absorbed beta inside v_new (via u). - # 4. Combine Core Output + # 3. Combine Core Output term2 = jnp.matmul(attn_i, v_new, precision=prec) o_c = attn_inter + term2 - # --- State Update --- - g_i_last_exp = jnp.exp(g[..., -1, None, None]) - h_new = h * g_i_last_exp + # --- State Update --- + g_i_last_exp = jnp.exp(g[..., -1, None, None]) + h_new = h * g_i_last_exp - # Apply Delta Rule K decay to state - g_diff_exp_state = jnp.exp(g[..., -1, None] - g)[..., None] - k_i_g_diff = k.astype(jnp.float32) * g_diff_exp_state + # Apply Delta Rule K decay to state + g_diff_exp_state = jnp.exp(g[..., -1, None] - g)[..., None] + k_i_g_diff = k.astype(jnp.float32) * g_diff_exp_state - update_term = jnp.matmul(k_i_g_diff.swapaxes(-1, -2), v_new, precision=prec) - h_new = h_new + update_term + update_term = jnp.matmul(k_i_g_diff.swapaxes(-1, -2), v_new, precision=prec) + h_new = h_new + update_term + if not propagate_transition: return h_new, o_c - final_h, o_chunks = lax.scan(scan_body, h_init, xs) + # M_new = (gamma * I - k_tilde^T @ w) @ M, kept in the low-rank form so the cost + # is 2 * K^2 * C rather than a dense K^3 product. + w_f32 = w.astype(jnp.float32) + m_new = m * g_i_last_exp - jnp.matmul(k_i_g_diff.swapaxes(-1, -2), jnp.matmul(w_f32, m, precision=prec), precision=prec) + return (h_new, m_new), o_c - # ========================================================================= - # STAGE 4: FINALIZATION - # ========================================================================= + +def _gdn_finalize(o_chunks, meta, initial_dtype) -> Array: + """STAGE 4: un-chunk the scan outputs and drop the chunk padding.""" + B, seq_len, H, _, V_dim, pad_len = meta o = o_chunks.transpose(1, 0, 3, 2, 4) o = o.reshape(B, -1, H, V_dim) - if pad_len > 0: o = o[:, :seq_len, :, :] + return o.astype(initial_dtype) + + +def jax_chunk_gated_delta_rule( + query: Array, + key: Array, + value: Array, + g: Array, + beta: Array, + chunk_size: int = 64, + initial_state: None | Array = None, + use_qk_norm_in_gdn: bool = False, + compute_dtype: jnp.dtype = jnp.bfloat16, +) -> tuple[Array, None | Array]: + """Optimized JAX implementation of Gated Delta Rule.""" + initial_dtype = query.dtype + + xs, meta = _gdn_chunk_prepare(query, key, value, g, beta, chunk_size, use_qk_norm_in_gdn, compute_dtype) + B, _, H, K_dim, V_dim, _ = meta - o = o.astype(initial_dtype) + if initial_state is None: + h_init = jnp.zeros((B, H, K_dim, V_dim), dtype=jnp.float32) + else: + h_init = initial_state.astype(jnp.float32) + + scan_body = functools.partial(_gdn_scan_body, chunk_size=chunk_size, compute_output=True, propagate_transition=False) + final_h, o_chunks = lax.scan(scan_body, h_init, xs) + o = _gdn_finalize(o_chunks, meta, initial_dtype) + return o, (final_h if initial_state is not None else None) + + +def jax_chunk_gated_delta_rule_cp( + query: Array, + key: Array, + value: Array, + g: Array, + beta: Array, + initial_state: None | Array = None, + *, + cp_axis_name: str, + chunk_size: int = 64, + use_qk_norm_in_gdn: bool = False, + compute_dtype: jnp.dtype = jnp.bfloat16, +) -> tuple[Array, None | Array]: + """Context-parallel gated delta rule. Must be called inside a `shard_map`. + + Each rank owns a contiguous slice of the sequence. Because the chunk recurrence + is affine in the incoming state, a rank's whole shard collapses to a single + affine map `(M_r, U_r)`; combining shards is then a prefix scan over those maps. + + The cross-rank payload is `M_r` of shape `(B, H, K, K)` plus `U_r` of shape + `(B, H, K, V)` -- **independent of sequence length**, unlike softmax-attention + context parallelism which exchanges O(S) key/value tensors. + + Three phases: + 1. Local pass from a zero state, carrying the composed transition `M_r`. + Outputs are skipped here since the true incoming state is not known yet. + 2. `all_gather` the `(M_r, U_r)` pairs and build the exclusive prefix, giving + each rank its true incoming state. + 3. Local pass again with the correct incoming state, reusing the WY factors + computed in phase 1, now emitting outputs. + + Args: + cp_axis_name: Mesh axis the sequence is sharded over, as named inside the + enclosing `shard_map`. + + Returns: + `(output, final_state)` with the same convention as + `jax_chunk_gated_delta_rule`. `final_state` is the *global* end-of-sequence + state, identical on every rank. + """ + initial_dtype = query.dtype + + xs, meta = _gdn_chunk_prepare(query, key, value, g, beta, chunk_size, use_qk_norm_in_gdn, compute_dtype) + B, _, H, K_dim, V_dim, _ = meta + + # --- Phase 1: local affine map, starting from a zero state. --- + h_zero = jnp.zeros((B, H, K_dim, V_dim), dtype=jnp.float32) + m_identity = jnp.broadcast_to(jnp.eye(K_dim, dtype=jnp.float32), (B, H, K_dim, K_dim)) + state_only_body = functools.partial( + _gdn_scan_body, chunk_size=chunk_size, compute_output=False, propagate_transition=True + ) + (u_local, m_local), _ = lax.scan(state_only_body, (h_zero, m_identity), xs) + + # --- Phase 2: exchange the affine maps and take the exclusive prefix. --- + if initial_state is None: + h_global = jnp.zeros((B, H, K_dim, V_dim), dtype=jnp.float32) + else: + h_global = initial_state.astype(jnp.float32) + + # (cp, B, H, K, K) and (cp, B, H, K, V): constant in sequence length. + m_all = jax.lax.all_gather(m_local, cp_axis_name, axis=0, tiled=False) + u_all = jax.lax.all_gather(u_local, cp_axis_name, axis=0, tiled=False) + + def prefix_body(h, m_u): + m_r, u_r = m_u + h_next = jnp.matmul(m_r, h, precision=jax.lax.Precision.HIGHEST) + u_r + return h_next, h # emit the state *entering* rank r (exclusive prefix) + + final_h, h_incoming = lax.scan(prefix_body, h_global, (m_all, u_all)) + h_init = jax.lax.dynamic_index_in_dim(h_incoming, jax.lax.axis_index(cp_axis_name), axis=0, keepdims=False) + + # --- Phase 3: local pass with the true incoming state. --- + output_body = functools.partial(_gdn_scan_body, chunk_size=chunk_size, compute_output=True, propagate_transition=False) + _, o_chunks = lax.scan(output_body, h_init, xs) + + o = _gdn_finalize(o_chunks, meta, initial_dtype) return o, (final_h if initial_state is not None else None) @@ -569,6 +694,32 @@ def a_log_init(key, shape, dtype=jnp.float32): rngs=rngs, ) + def _context_parallel_axis(self, seq_len: int) -> tuple[str | None, int]: + """Returns the mesh axis the sequence is sharded over, and its size. + + Returns `(None, 1)` when context parallelism is off or does not apply (single + token decode steps, or a sequence that cannot be split evenly). + """ + axis = getattr(self.config, "context_sharding", "context") + if self.mesh is None or axis not in self.mesh.axis_names: + return None, 1 + size = self.mesh.shape[axis] + if size <= 1 or seq_len <= 1: + return None, 1 + if seq_len % size != 0: + raise ValueError( + f"GDN context parallelism requires the sequence length ({seq_len}) to be divisible by the " + f"'{axis}' mesh axis size ({size})." + ) + if getattr(self.config, "context_parallel_load_balance", False): + raise ValueError( + "GDN context parallelism requires context_parallel_load_balance=False. The gated delta rule " + "is a sequential recurrence, so each rank must own a contiguous, in-order slice of the " + "sequence; the load-balancing reorder interleaves chunks across ranks and would silently " + "compute the wrong result." + ) + return axis, size + def __call__( self, hidden_states: Array, @@ -581,6 +732,7 @@ def __call__( # hidden_states: (B, S, E) cfg = self.config batch, seq_len, _ = hidden_states.shape + cp_axis_name, cp_size = self._context_parallel_axis(seq_len) active_cache = kv_cache if kv_cache is not None else self.cache @@ -626,6 +778,12 @@ def __call__( dims=(0,), allow_remove_axes=True, ) + # Keep the sequence sharded under context parallelism. Replicating it here would + # all-gather the whole sequence and defeat the point of splitting it in the first + # place; the causal depthwise conv below only needs a `gdn_conv_kernel_dim - 1` + # token halo, which the SPMD partitioner emits as a collective-permute. + if cp_size > 1: + qkvz_pspec = P(qkvz_pspec[0], cp_axis_name, *qkvz_pspec[2:]) qkvz_sharding = jax.sharding.NamedSharding(self.mesh, qkvz_pspec) mixed_qkvz = jax.lax.with_sharding_constraint(mixed_qkvz, qkvz_sharding) @@ -677,7 +835,6 @@ def __call__( truncate_sharded_tensor, ) from tpu_inference.utils import get_mesh_shape_product # pylint: disable=import-outside-toplevel # pytype: disable=import-error - from jax.sharding import PartitionSpec as P_spec # pylint: disable=import-outside-toplevel # pytype: disable=import-error except ImportError as e: raise ImportError( "GDN attention kernel require the vllm-tpu package. Please install it with `pip install vllm-tpu`." @@ -700,8 +857,8 @@ def __call__( mixed_qkv = jax.shard_map( lambda q, k, v: jnp.concatenate([q, k, v], axis=-1), mesh=self.mesh, - in_specs=(P_spec(attn_data, attn_head),) * 3, - out_specs=P_spec(attn_data, attn_head), + in_specs=(P(attn_data, attn_head),) * 3, + out_specs=P(attn_data, attn_head), check_vma=False, )(q_flat, k_flat, v_flat) @@ -906,6 +1063,13 @@ def extract_state(c_in, v_len): dims=(0,), allow_remove_axes=True, ) + # Under context parallelism the sequence stays sharded inside the shard_map and + # the recurrence is stitched back together with a prefix scan over the per-rank + # affine maps. The recurrent state itself is never sequence-sharded, so + # `state_pspec` is unchanged and stays replicated across the context axis. + if cp_size > 1: + qkv_pspec = P(qkv_pspec[0], cp_axis_name, *qkv_pspec[2:]) + g_beta_pspec = P(g_beta_pspec[0], cp_axis_name, *g_beta_pspec[2:]) @functools.partial( jax.shard_map, @@ -925,6 +1089,19 @@ def extract_state(c_in, v_len): check_vma=False, ) def shard_mapped_delta_rule(q, k, v, g_val, beta_val, init_h): + if cp_size > 1: + return jax_chunk_gated_delta_rule_cp( + query=q, + key=k, + value=v, + g=g_val, + beta=beta_val, + cp_axis_name=cp_axis_name, + chunk_size=cfg.gdn_chunk_size, + initial_state=init_h, + use_qk_norm_in_gdn=cfg.use_qk_norm_in_gdn, + compute_dtype=cfg.dtype, + ) return jax_chunk_gated_delta_rule( query=q, key=k, diff --git a/tests/unit/gdn_context_parallelism_test.py b/tests/unit/gdn_context_parallelism_test.py new file mode 100644 index 0000000000..9546b488b2 --- /dev/null +++ b/tests/unit/gdn_context_parallelism_test.py @@ -0,0 +1,249 @@ +# 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 context parallelism in the Gated Delta Net (Qwen3-Next / Qwen3.5).""" + +import functools +import re +import sys +import unittest + +from flax import nnx +import jax +import jax.numpy as jnp +from jax.sharding import Mesh, PartitionSpec as P +import numpy as np +import pytest + +from maxtext.common.common_types import MODEL_MODE_TRAIN +from maxtext.configs import pyconfig +from maxtext.models.qwen3 import ( + Qwen3NextGatedDeltaNet, + jax_chunk_gated_delta_rule, + jax_chunk_gated_delta_rule_cp, +) +from maxtext.utils import maxtext_utils + +from tests.utils.test_helpers import get_test_config_path + + +def _layer_forward(layer, x): + """Module-level so `jax.jit` does not close over a loop variable.""" + return layer(x, model_mode=MODEL_MODE_TRAIN)[0] + + +def _cp_size() -> int: + """Largest power-of-two context-parallel degree the local topology supports.""" + n = jax.device_count() + size = 1 + while size * 2 <= n: + size *= 2 + return size + + +@pytest.mark.tpu_only +class GatedDeltaRuleContextParallelKernelTest(unittest.TestCase): + """`jax_chunk_gated_delta_rule_cp` must match the single-device recurrence.""" + + B, S, H, K, V, C = 2, 512, 4, 128, 128, 64 + + def setUp(self): + super().setUp() + self.cp = _cp_size() + if self.cp < 2: + self.skipTest("context parallelism needs at least 2 devices") + rng = np.random.default_rng(0) + + def mk(*shape): + return jnp.asarray(rng.standard_normal(shape), dtype=jnp.float32) + + self.q = mk(self.B, self.S, self.H, self.K) + self.k = mk(self.B, self.S, self.H, self.K) + self.v = mk(self.B, self.S, self.H, self.V) + # `g` is a log-decay (negative) and `beta` a sigmoid output, matching the layer. + self.g = -jnp.exp(mk(self.B, self.S, self.H) * 0.3) * 0.05 + self.beta = jax.nn.sigmoid(mk(self.B, self.S, self.H)) + self.h0 = jnp.zeros((self.B, self.H, self.K, self.V), jnp.float32) + self.mesh = Mesh(np.array(jax.devices()[: self.cp]).reshape(self.cp), ("context",)) + + def _cp_call(self, q, k, v, g, beta, h0): + """Runs the context-parallel kernel under a `shard_map` over the `context` axis.""" + fn = jax.shard_map( + lambda *a: jax_chunk_gated_delta_rule_cp( + *a, + cp_axis_name="context", + chunk_size=self.C, + compute_dtype=jnp.float32, + use_qk_norm_in_gdn=True, + ), + mesh=self.mesh, + in_specs=( + P(None, "context", None, None), + P(None, "context", None, None), + P(None, "context", None, None), + P(None, "context", None), + P(None, "context", None), + P(), + ), + out_specs=(P(None, "context", None, None), P()), + check_vma=False, + ) + return fn(q, k, v, g, beta, h0) + + def _ref_call(self, q, k, v, g, beta, h0): + return jax_chunk_gated_delta_rule( + q, k, v, g, beta, chunk_size=self.C, initial_state=h0, use_qk_norm_in_gdn=True, compute_dtype=jnp.float32 + ) + + def test_forward_matches_single_device(self): + ref_o, ref_h = jax.jit(self._ref_call)(self.q, self.k, self.v, self.g, self.beta, self.h0) + with jax.set_mesh(self.mesh): + cp_o, cp_h = jax.jit(self._cp_call)(self.q, self.k, self.v, self.g, self.beta, self.h0) + np.testing.assert_allclose(cp_o, ref_o, rtol=2e-4, atol=2e-4) + np.testing.assert_allclose(cp_h, ref_h, rtol=2e-4, atol=2e-4) + + def test_backward_matches_single_device(self): + def loss(fn, *args): + o, _ = fn(*args, self.h0) + return jnp.sum(o * jnp.sin(o)) + + args = (self.q, self.k, self.v, self.g, self.beta) + with jax.set_mesh(self.mesh): + ref_g = jax.grad(functools.partial(loss, self._ref_call), argnums=(0, 1, 2, 3, 4))(*args) + cp_g = jax.grad(functools.partial(loss, self._cp_call), argnums=(0, 1, 2, 3, 4))(*args) + for name, a, b in zip("query key value g beta".split(), ref_g, cp_g): + np.testing.assert_allclose(b, a, rtol=5e-3, atol=5e-4, err_msg=f"gradient mismatch for {name}") + + def test_cross_rank_payload_is_sequence_length_independent(self): + """The only cross-rank tensors are the (M, U) affine maps, sized by head dims. + + This is the whole point of the scheme: softmax-attention context parallelism + exchanges O(sequence_length) key/value tensors, whereas the delta rule collapses + each rank's shard into a fixed-size affine map. + """ + # Matches `%foo = f32[8,4,128,128]{...} all-gather(...)`, ignoring bitcast aliases. + pattern = re.compile(r"=\s*\w+\[([0-9,]+)\][^=]*\ball-gather(?:-start)?\(") + sizes = {} + for repeat in (1, 2, 4): + tiled = [jnp.concatenate([x] * repeat, axis=1) for x in (self.q, self.k, self.v, self.g, self.beta)] + with jax.set_mesh(self.mesh): + hlo = jax.jit(self._cp_call).lower(*tiled, self.h0).compile().as_text() + total = 0 + for line in hlo.splitlines(): + match = pattern.search(line) + if match: + total += int(np.prod([int(d) for d in match.group(1).split(",")])) + self.assertGreater(total, 0, "expected the cross-rank prefix exchange to appear in the HLO") + sizes[repeat * self.S] = total + self.assertEqual( + len(set(sizes.values())), + 1, + f"cross-rank all-gather volume must not grow with sequence length, got {sizes}", + ) + # cp * batch * heads * k_dim * (k_dim for M + v_dim for U), in elements. + expected = self.cp * self.B * self.H * self.K * (self.K + self.V) + self.assertEqual(next(iter(sizes.values())), expected) + + +@pytest.mark.tpu_only +class GatedDeltaNetLayerContextParallelTest(unittest.TestCase): + """The full GDN layer must be invariant to the context-parallel degree.""" + + def setUp(self): + super().setUp() + self.cp = _cp_size() + if self.cp < 2: + self.skipTest("context parallelism needs at least 2 devices") + + def _build(self, cp_degree): + """Builds a small GDN layer on a mesh with the requested context-parallel degree.""" + cfg = pyconfig.initialize( + [sys.argv[0], get_test_config_path()], + per_device_batch_size=1.0, + run_name="gdn_cp_test", + enable_checkpointing=False, + max_prefill_predict_length=32, + max_target_length=256, + base_emb_dim=128, + gdn_num_value_heads=4, + gdn_num_key_heads=4, + gdn_key_head_dim=32, + gdn_value_head_dim=32, + gdn_conv_kernel_dim=4, + gdn_chunk_size=16, + ici_context_parallelism=cp_degree, + context_parallel_load_balance=False, + dtype="float32", + weight_dtype="float32", + ) + mesh = Mesh(maxtext_utils.create_device_mesh(cfg), cfg.mesh_axes) + layer = Qwen3NextGatedDeltaNet( + config=cfg, + inputs_shape=(cfg.global_batch_size_to_train_on, cfg.max_target_length, cfg.emb_dim), + mesh=mesh, + dtype=jnp.float32, + model_mode=MODEL_MODE_TRAIN, + rngs=nnx.Rngs(params=0, dropout=jax.random.PRNGKey(42)), + ) + return cfg, mesh, layer + + def test_layer_output_matches_without_context_parallelism(self): + outs = {} + for cp_degree in (1, self.cp): + cfg, mesh, layer = self._build(cp_degree) + lnx = jax.random.normal( + jax.random.PRNGKey(7), + (cfg.global_batch_size_to_train_on, cfg.max_target_length, cfg.emb_dim), + dtype=jnp.float32, + ) + forward = functools.partial(_layer_forward, layer) + with jax.set_mesh(mesh): + out = jax.jit(forward)(lnx) + outs[cp_degree] = np.asarray(out) + np.testing.assert_allclose(outs[self.cp], outs[1], rtol=2e-4, atol=2e-4) + + def test_load_balancing_is_rejected(self): + cfg = pyconfig.initialize( + [sys.argv[0], get_test_config_path()], + per_device_batch_size=1.0, + run_name="gdn_cp_lb_test", + enable_checkpointing=False, + max_target_length=256, + base_emb_dim=128, + gdn_num_value_heads=4, + gdn_num_key_heads=4, + gdn_key_head_dim=32, + gdn_value_head_dim=32, + gdn_chunk_size=16, + ici_context_parallelism=self.cp, + context_parallel_load_balance=True, + dtype="float32", + ) + mesh = Mesh(maxtext_utils.create_device_mesh(cfg), cfg.mesh_axes) + layer = Qwen3NextGatedDeltaNet( + config=cfg, + inputs_shape=(cfg.global_batch_size_to_train_on, cfg.max_target_length, cfg.emb_dim), + mesh=mesh, + dtype=jnp.float32, + model_mode=MODEL_MODE_TRAIN, + rngs=nnx.Rngs(params=0, dropout=jax.random.PRNGKey(42)), + ) + lnx = jnp.zeros((cfg.global_batch_size_to_train_on, cfg.max_target_length, cfg.emb_dim), jnp.float32) + with self.assertRaisesRegex(ValueError, "context_parallel_load_balance"): + with jax.set_mesh(mesh): + layer(lnx, model_mode=MODEL_MODE_TRAIN) + + +if __name__ == "__main__": + unittest.main()