Shard the GatedDeltaNet sequence under ici_context_parallelism

src/maxtext/models/gdn_cp.py

@@ -0,0 +1,79 @@
+"""Context-parallel evaluation of the GatedDeltaNet inter-chunk recurrence.
+
+The recurrence h_new = A @ h + B is affine in the state, so it composes
+associatively and can be split across a sharded sequence. See
+apply_gdn_context_parallel.py for the derivation and the measurements.
+"""
+
+from __future__ import annotations
+
+import jax
+import jax.numpy as jnp
+from jax import lax
+
+_PREC = jax.lax.Precision.HIGHEST
+
+
+def compose(left, right):
+  """(A_r, B_r) . (A_l, B_l) = (A_r @ A_l, A_r @ B_l + B_r)."""
+  A_l, B_l = left
+  A_r, B_r = right
+  return (
+      jnp.matmul(A_r, A_l, precision=_PREC),
+      jnp.matmul(A_r, B_l, precision=_PREC) + B_r,
+  )
+
+
+def compose_local(w, u, k, g):
+  """Fold this device's chunks into one affine map, in O(1) memory.
+
+  lax.scan rather than associative_scan on purpose: associative_scan would
+  materialise (A, B) and their running composition for every chunk, which is
+  17 GB per device at a million tokens and defeats the point. The parallelism
+  that matters here is across devices, not within one.
+  """
+  k_dim = k.shape[-1]
+  eye = jnp.eye(k_dim, dtype=jnp.float32)
+
+  # jax.checkpoint is required here. lax.scan keeps whatever
+  # the body computes as a backward residual, so A_i and B_i get stacked over
+  # every chunk even though the forward pass only ever needs one at a time. At
+  # sequence 262,144 with ctx=4 that was 103 GB of f32[1024,4,16,128,128] in the
+  # buffer dump. A_i and B_i are cheap to rebuild from w, u, k and g, which are
+  # already live, so recompute them in the backward pass instead of storing them.
+  @jax.checkpoint
+  def body(carry, x):
+    w_c, u_c, k_c, g_c = x
+    g_last = g_c[..., -1]
+    decay = jnp.exp(g_last)[..., None, None]
+    k_g = k_c.astype(jnp.float32) * jnp.exp(g_last[..., None] - g_c)[..., None]
+    k_g_T = k_g.swapaxes(-1, -2)
+    A_i = decay * eye - jnp.matmul(k_g_T, w_c.astype(jnp.float32), precision=_PREC)
+    B_i = jnp.matmul(k_g_T, u_c.astype(jnp.float32), precision=_PREC)
+    return compose(carry, (A_i, B_i)), None
+
+  lead = w.shape[1:-2]
+  init = (
+      jnp.broadcast_to(eye, lead + (k_dim, k_dim)).astype(jnp.float32),
+      jnp.zeros(lead + (k_dim, u.shape[-1]), jnp.float32),
+  )
+  (A_loc, B_loc), _ = lax.scan(body, init, (w, u, k, g))
+  return A_loc, B_loc
+
+
+def incoming_state(A_loc, B_loc, h_init, cp_axis):
+  """State entering this device, plus the final state after all devices.
+
+  Gathering D pairs of small matrices is the only cross-device traffic in the
+  scheme. Must be called inside a shard_map over `cp_axis`.
+  """
+  A_all = lax.all_gather(A_loc, cp_axis, axis=0, tiled=False)
+  B_all = lax.all_gather(B_loc, cp_axis, axis=0, tiled=False)
+  A_cum, B_cum = lax.associative_scan(compose, (A_all, B_all), axis=0)
+
+  idx = lax.axis_index(cp_axis)
+  prev = jnp.maximum(idx - 1, 0)
+  carried = jnp.matmul(A_cum[prev], h_init, precision=_PREC) + B_cum[prev]
+  h_in = jnp.where(idx == 0, h_init, carried)
+  final_h = jnp.matmul(A_cum[-1], h_init, precision=_PREC) + B_cum[-1]
+  return h_in, final_h

src/maxtext/models/qwen3.py

@@ -33,6 +33,7 @@
 
 from maxtext.common.common_types import AttentionType, Config, DType, Array, BATCH, EMBED, MODEL_MODE_TRAIN, LENGTH, MODEL_MODE_AUTOREGRESSIVE
 from maxtext.common.common_types import KV_BATCH, KV_HEAD
+from maxtext.models import gdn_cp
 from maxtext.utils.sharding import (
     get_logical_axis_rules,
     logical_to_mesh_axes,
@@ -193,6 +194,7 @@ def jax_chunk_gated_delta_rule(
     chunk_size: int = 64,
     initial_state: None | Array = None,
     use_qk_norm_in_gdn: bool = False,
+    cp_axis: None | str = None,
     compute_dtype: jnp.dtype = jnp.bfloat16,
 ) -> tuple[Array, None | Array]:
   """Optimized JAX implementation of Gated Delta Rule."""
@@ -348,7 +350,15 @@ def scan_body(h, args):
 
     return h_new, o_c
 
-  final_h, o_chunks = lax.scan(scan_body, h_init, xs)
+  if cp_axis is None:
+    final_h, o_chunks = lax.scan(scan_body, h_init, xs)
+  else:
+    # Sequence is sharded over cp_axis, so a sequential scan over chunks is not
+    # available. Fold the local chunks into one affine map, exchange those, then
+    # replay locally from the correct incoming state. See models/gdn_cp.py.
+    A_loc, B_loc = gdn_cp.compose_local(w_scan, u_scan, k_scan, g_scan)
+    h_in, final_h = gdn_cp.incoming_state(A_loc, B_loc, h_init, cp_axis)
+    _, o_chunks = lax.scan(scan_body, h_in, xs)
 
   # =========================================================================
   # STAGE 4: FINALIZATION
@@ -617,7 +627,15 @@ def __call__(
     mixed_qkvz = qkvz.reshape(new_shape_qkvz)
     if self.mesh is not None:
       logical_rules = get_logical_axis_rules()
-      qkvz_pspec = logical_to_mesh_axes((KV_BATCH, None, KV_HEAD, None), mesh=self.mesh, rules=logical_rules)
+      # LENGTH, not None. This with_sharding_constraint told XLA to gather the
+      # full sequence onto every device for the qkvz projection. With ctx=2 at
+      # seq 32,768 that is six live bf16[2, 32768, 16, 512] buffers of 1.07 GB
+      # each -- found by diffing XLA buffer assignment between ctx=1 and ctx=2,
+      # and the reason the first version of this patch cut GDN memory but still
+      # lost overall.
+      _cp_on_qkvz = cfg.ici_context_parallelism > 1 or getattr(cfg, "ici_context_usp_ulysses_parallelism", 1) > 1
+      _cp_len_qkvz = LENGTH if _cp_on_qkvz else None
+      qkvz_pspec = logical_to_mesh_axes((KV_BATCH, _cp_len_qkvz, KV_HEAD, None), mesh=self.mesh, rules=logical_rules)
       # Training microbatches can be smaller than the physical KV_BATCH mesh partition.
       qkvz_pspec = remove_incompatible_mesh_axes_from_partition_spec(
           qkvz_pspec,
@@ -881,8 +899,17 @@ def extract_state(c_in, v_len):
           if recurrent_state is not None
           else jnp.zeros((batch, self.num_v_heads, self.head_k_dim, self.head_v_dim), dtype=cfg.dtype)
       )
-      qkv_pspec = logical_to_mesh_axes((KV_BATCH, None, KV_HEAD, None), mesh=self.mesh, rules=logical_rules)
-      g_beta_pspec = logical_to_mesh_axes((KV_BATCH, None, KV_HEAD), mesh=self.mesh, rules=logical_rules)
+      # LENGTH, not None. The sequence axis was hardcoded to replicated, so
+      # ici_context_parallelism could never shard the GDN sequence while still
+      # consuming the context axis from the mesh -- which is why raising ctx
+      # made memory worse instead of better. The scan handles a sharded
+      # sequence via the two-pass affine composition in models/gdn_cp.py.
+      # Either context axis can carry the sequence. LENGTH maps to both in the
+      # logical rules, so this is correct whichever one is configured.
+      _cp_on = cfg.ici_context_parallelism > 1 or getattr(cfg, "ici_context_usp_ulysses_parallelism", 1) > 1
+      _cp_len = LENGTH if _cp_on else None
+      qkv_pspec = logical_to_mesh_axes((KV_BATCH, _cp_len, KV_HEAD, None), mesh=self.mesh, rules=logical_rules)
+      g_beta_pspec = logical_to_mesh_axes((KV_BATCH, _cp_len, KV_HEAD), mesh=self.mesh, rules=logical_rules)
       state_pspec = logical_to_mesh_axes((KV_BATCH, KV_HEAD, None, None), mesh=self.mesh, rules=logical_rules)
       # Keep every shard_map input/output batch spec consistent when replication is required.
       qkv_pspec = remove_incompatible_mesh_axes_from_partition_spec(
@@ -935,6 +962,17 @@ def shard_mapped_delta_rule(q, k, v, g_val, beta_val, init_h):
             initial_state=init_h,
             use_qk_norm_in_gdn=cfg.use_qk_norm_in_gdn,
             compute_dtype=cfg.dtype,
+            # Whichever context axis is active. all_gather takes a tuple, so
+            # both can be live at once.
+            cp_axis=tuple(
+                a
+                for a, n in (
+                    ("context", cfg.ici_context_parallelism),
+                    ("context_usp_ulysses", getattr(cfg, "ici_context_usp_ulysses_parallelism", 1)),
+                )
+                if n > 1
+            )
+            or None,
         )
 
       core_attn_out, next_recurrent_state = shard_mapped_delta_rule(query, key, value, g, beta, recurrent_state_arg)