diff --git a/src/maxtext/models/gdn_cp.py b/src/maxtext/models/gdn_cp.py new file mode 100644 index 0000000000..331e255b3a --- /dev/null +++ b/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 diff --git a/src/maxtext/models/qwen3.py b/src/maxtext/models/qwen3.py index e46a41dda9..fb66b99ecc 100644 --- a/src/maxtext/models/qwen3.py +++ b/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)