From a29991d3fef51325fe4a306e3590ca52fdad80b3 Mon Sep 17 00:00:00 2001 From: Param Bole Date: Tue, 13 Oct 2026 18:38:19 -0700 Subject: [PATCH 1/4] Add an opt-in sharded log-prob head for packed rows (TUNIX_LOGPS_SHARDED_HEAD) compute_chunked_logps scans the sequence under GSPMD. On the MLPerf Qwen3.5-397B trainer cell (fsdp32 x cp4 x ep2 on 256 devices, 64 x 65,536-token packed micro-batches, compute_logps_chunk_size 1024) the compiled program all-gathers the 2 GB vocab table in every scan step (forward and rematerialized backward) and all-reduces a 2 GB table gradient in every backward step: 129 table all-gathers and 65 table all-reduces per micro-batch. Because the scan runs over a context-sharded sequence, every context rank also computes every chunk, so the vocab matmuls are 4x redundant. With TUNIX_LOGPS_SHARDED_HEAD=local (or 1), packed rows go through a shard_map'ed head that runs the same per-chunk math: - each head parameter is all-gathered once per pass; - each device scans only the tokens it holds (dynamic slices, no reshape of the sequence); - the table gradient is accumulated per device in float32 and reduce-scattered once; - the one-token shifts (targets left, outputs right) are explicit collective-permutes inside the shard_map, so they do not depend on how XLA partitions a shift of a sequence-sharded array. The head falls back to compute_chunked_logps, with a log line, when there is no mesh in context, when some mesh axis shards neither the batch nor the sequence, or when a head parameter does not divide over the mesh. Unset (the default), nothing changes. --- tests/rl/sharded_logps_head_test.py | 315 +++++++++++++++++++++ tunix/rl/common.py | 424 +++++++++++++++++++++++++++- 2 files changed, 729 insertions(+), 10 deletions(-) create mode 100644 tests/rl/sharded_logps_head_test.py diff --git a/tests/rl/sharded_logps_head_test.py b/tests/rl/sharded_logps_head_test.py new file mode 100644 index 0000000000..3ca056ee2b --- /dev/null +++ b/tests/rl/sharded_logps_head_test.py @@ -0,0 +1,315 @@ +# 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 +# +# http://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 sharded chunked log-prob head (TUNIX_LOGPS_SHARDED_HEAD).""" + +import os + +os.environ['XLA_FLAGS'] = '--xla_force_host_platform_device_count=8' + +from unittest import mock # pylint: disable=g-import-not-at-top + +from absl.testing import absltest +from absl.testing import parameterized +from flax import nnx +import flax.linen as nn +import jax +import jax.numpy as jnp +from jax.sharding import Mesh, NamedSharding, PartitionSpec as P # pylint: disable=g-multiple-import,g-importing-member +import numpy as np +from tunix.rl import common + +_B, _T, _H, _V = 4, 64, 32, 128 +_AXES = ('data', 'fsdp', 'context', 'expert') +_RULES = ( + ('activation_batch', ('data', 'fsdp', 'expert')), + ('activation_length', 'context'), +) + + +class _HeadOnly(nnx.Module): + """Frozen hidden states into a MaxText-like head (RMSNorm, f32 table).""" + + def __init__(self, dtype, batch=_B, seq=_T): + keys = jax.random.split(jax.random.PRNGKey(0), 3) + self.dtype = dtype + self.h = nnx.Param( + jax.random.normal(keys[0], (batch, seq, _H)).astype(dtype) + ) + self.scale = nnx.Param( + 0.2 * jax.random.normal(keys[1], (_H,)).astype(dtype) + ) + self.table = nnx.Param(3.0 * jax.random.normal(keys[2], (_H, _V)) / _H**0.5) + + def __call__( + self, + input_tokens, + positions, + cache, + attention_mask, + segment_ids=None, + skip_lm_head=False, + ): + del input_tokens, positions, cache, attention_mask, segment_ids + h = self.h[...] + return (h if skip_lm_head else self.compute_final_logits(h)), None + + def compute_final_logits(self, h): + x = h.astype(jnp.float32) + y = x * jax.lax.rsqrt(jnp.mean(x * x, -1, keepdims=True) + 1e-6) + y = (y * (self.scale[...].astype(jnp.float32) + 1.0)).astype(self.dtype) + out = jnp.dot(y, self.table[...].astype(self.dtype)) + return out.astype(jnp.float32) + + +def _mesh(): + devices = np.array(jax.devices()[:8]).reshape(1, 2, 2, 2) + return Mesh(devices, _AXES) + + +class ShardedLogpsHeadTest(parameterized.TestCase): + + def setUp(self): + super().setUp() + self.mesh = _mesh() + + def _inputs(self, batch, seq): + rng = np.random.default_rng(0) + tokens = jnp.asarray(rng.integers(1, _V, (batch, seq)), jnp.int32) + segment_ids = jnp.ones((batch, seq), jnp.int32) + positions = jnp.asarray(np.tile(np.arange(seq), (batch, 1)), jnp.int32) + weights = jnp.asarray(rng.normal(size=(2, batch, seq)), jnp.float32) + if batch % 4 == 0 and seq % 2 == 0: + sharding = NamedSharding( + self.mesh, P(('data', 'fsdp', 'expert'), 'context') + ) + tokens, segment_ids, positions = ( + jax.device_put(x, sharding) for x in (tokens, segment_ids, positions) + ) + return tokens, segment_ids, positions, weights + + def _run( + self, + mode, + dtype=jnp.float32, + chunk_size=8, + rules=_RULES, + temperature=0.7, + batch=_B, + seq=_T, + return_entropy=True, + ): + model = _HeadOnly(dtype, batch, seq) + graphdef, state = nnx.split(model) + tokens, segment_ids, positions, weights = self._inputs(batch, seq) + + def loss(state): + out = common.compute_per_token_logps( + graphdef, + state, + jnp.zeros((tokens.shape[0], 0), jnp.int32), + tokens, + pad_id=0, + eos_id=-1, + stop_gradient=False, + return_entropy=return_entropy, + segment_ids=segment_ids, + segment_positions=positions, + temperature=temperature, + chunk_size=chunk_size, + ) + logps, entropy = out if return_entropy else (out, jnp.zeros_like(out)) + value = jnp.sum(logps * weights[0]) + jnp.sum(entropy * weights[1]) + return value, (logps, entropy) + + jax.clear_caches() + env = {} if mode is None else {common.SHARDED_HEAD_ENV: mode} + with mock.patch.dict(os.environ, env, clear=False): + if mode is None: + os.environ.pop(common.SHARDED_HEAD_ENV, None) + with jax.set_mesh(self.mesh), nn.logical_axis_rules(rules): + (_, (logps, entropy)), grads = jax.jit( + jax.value_and_grad(loss, has_aux=True) + )(state) + return logps, entropy, grads + + @parameterized.named_parameters( + ('local', 'local', 8), + ('local_uneven_chunk', 'local', 12), + ('local_one_step', 'local', 64), + ) + def test_matches_scan_head(self, mode, chunk_size): + want = self._run(None, chunk_size=chunk_size) + got = self._run(mode, chunk_size=chunk_size) + np.testing.assert_allclose(got[0], want[0], rtol=1e-5, atol=2e-5) + np.testing.assert_allclose(got[1], want[1], rtol=1e-5, atol=2e-5) + for name in ('h', 'scale', 'table'): + np.testing.assert_allclose( + got[2][name][...], + want[2][name][...], + rtol=1e-4, + atol=1e-5, + err_msg=name, + ) + + def test_logps_only_matches_scan_head(self): + # return_entropy=False: one output, and no entropy cotangent in the bwd. + want = self._run(None, return_entropy=False) + got = self._run('local', return_entropy=False) + np.testing.assert_allclose(got[0], want[0], rtol=1e-5, atol=2e-5) + for name in ('h', 'scale', 'table'): + np.testing.assert_allclose( + got[2][name][...], + want[2][name][...], + rtol=1e-4, + atol=1e-5, + err_msg=name, + ) + + @parameterized.parameters(True, False) + def test_forward_only_matches_scan_head(self, return_entropy): + # No gradient: the custom_vjp's primal body runs instead of fwd/bwd. + model = _HeadOnly(jnp.float32) + graphdef, state = nnx.split(model) + tokens, segment_ids, positions, _ = self._inputs(_B, _T) + + def logps(mode): + jax.clear_caches() + env = {} if mode is None else {common.SHARDED_HEAD_ENV: mode} + with mock.patch.dict(os.environ, env, clear=False): + if mode is None: + os.environ.pop(common.SHARDED_HEAD_ENV, None) + with jax.set_mesh(self.mesh), nn.logical_axis_rules(_RULES): + return jax.jit( + lambda st: common.compute_per_token_logps( + graphdef, + st, + jnp.zeros((_B, 0), jnp.int32), + tokens, + pad_id=0, + eos_id=-1, + stop_gradient=True, + return_entropy=return_entropy, + segment_ids=segment_ids, + segment_positions=positions, + temperature=1.0, + chunk_size=8, + ) + )(state) + + want, got = logps(None), logps('local') + for g, w in zip(jax.tree.leaves(got), jax.tree.leaves(want)): + np.testing.assert_allclose(g, w, rtol=1e-5, atol=2e-5) + + def test_bf16_matches_scan_head(self): + # Same per-token math; XLA may order the per-token f32 reductions + # differently for the smaller per-device chunk, so allow f32 rounding. + want = self._run(None, dtype=jnp.bfloat16) + got = self._run('local', dtype=jnp.bfloat16) + np.testing.assert_allclose(got[0], want[0], rtol=1e-6, atol=1e-6) + np.testing.assert_allclose(got[1], want[1], rtol=1e-6, atol=1e-6) + np.testing.assert_allclose( + np.asarray(got[2]['h'][...], np.float32), + np.asarray(want[2]['h'][...], np.float32), + rtol=1e-2, + atol=1e-6, + ) + for name in ('scale', 'table'): + g = np.asarray(got[2][name][...], np.float32) + w = np.asarray(want[2][name][...], np.float32) + self.assertLess(np.linalg.norm(g - w) / np.linalg.norm(w), 2.0**-6, name) + + def test_mesh_order_layout_without_rules(self): + want = self._run(None, rules=()) + got = self._run('local', rules=()) + np.testing.assert_allclose(got[0], want[0], rtol=1e-5, atol=2e-5) + np.testing.assert_allclose( + got[2]['table'][...], want[2]['table'][...], rtol=1e-4, atol=1e-5 + ) + + def test_engages_on_mesh(self): + calls = [] + real = common.maybe_compute_sharded_chunked_logps + + def spy(*args, **kwargs): + out = real(*args, **kwargs) + calls.append(out is not None) + return out + + with mock.patch.object(common, 'maybe_compute_sharded_chunked_logps', spy): + self._run('local') + self._run(None) + self.assertEqual(calls, [True, False]) + + def _engaged(self, **kwargs): + calls = [] + real = common.maybe_compute_sharded_chunked_logps + + def spy(*args, **kw): + out = real(*args, **kw) + calls.append(out is not None) + return out + + with mock.patch.object(common, 'maybe_compute_sharded_chunked_logps', spy): + got = self._run('local', **kwargs) + return calls, got + + def test_sequence_only_layout_when_batch_does_not_divide(self): + # 3 rows cannot be split over fsdp x expert, so all 8 devices take a + # slice of every row's sequence instead. + calls, got = self._engaged(batch=3) + want = self._run(None, batch=3) + self.assertEqual(calls, [True]) + np.testing.assert_allclose(got[0], want[0], rtol=1e-5, atol=2e-5) + np.testing.assert_allclose( + got[2]['table'][...], want[2]['table'][...], rtol=1e-4, atol=1e-5 + ) + + def test_falls_back_when_no_layout(self): + # 3 rows of 63 tokens divide over no mesh axis: compute_chunked_logps runs. + calls, got = self._engaged(batch=3, seq=63, chunk_size=7) + want = self._run(None, batch=3, seq=63, chunk_size=7) + self.assertEqual(calls, [False]) + np.testing.assert_array_equal(got[0], want[0]) + + def test_falls_back_without_mesh(self): + model = _HeadOnly(jnp.float32) + tokens = self._inputs(_B, _T)[0] + with mock.patch.dict(os.environ, {common.SHARDED_HEAD_ENV: '1'}): + self.assertIsNone( + common.maybe_compute_sharded_chunked_logps( + model, model.h[...], tokens, 1.0, 8, True + ) + ) + + @parameterized.parameters( + ('', None), + ('0', None), + ('false', None), + ('1', 'local'), + ('local', 'local'), + ) + def test_mode_parsing(self, value, want): + with mock.patch.dict(os.environ, {common.SHARDED_HEAD_ENV: value}): + self.assertEqual(common.sharded_head_mode(), want) + + @parameterized.parameters('locla', 'per_step') + def test_mode_parsing_rejects_unknown(self, value): + with mock.patch.dict(os.environ, {common.SHARDED_HEAD_ENV: value}): + with self.assertRaises(ValueError): + common.sharded_head_mode() + + +if __name__ == '__main__': + absltest.main() diff --git a/tunix/rl/common.py b/tunix/rl/common.py index 36ced9511f..2c51c0aa3f 100644 --- a/tunix/rl/common.py +++ b/tunix/rl/common.py @@ -14,6 +14,8 @@ """Common RL helper classes and functions.""" import functools +import math +import os import re import inspect from typing import Any, Iterable, Sequence @@ -23,6 +25,7 @@ from flax import nnx import jax from jax import numpy as jnp +from jax.custom_derivatives import SymbolicZero # pylint: disable=g-importing-member from jax.typing import ArrayLike # pylint: disable=g-importing-member import jax.tree_util as jtu import numpy as np @@ -639,21 +642,36 @@ def _result(*outs): input_tokens_to_keep = input_tokens[:, -logits_to_keep:] if chunk_size > 0: - hidden_state = outputs[:, -logits_to_keep - 1 : -1, :] - out = compute_chunked_logps( - model, - hidden_state, - input_tokens_to_keep, - temperature, - chunk_size, - return_entropy, - ) + out = None + if segment_ids is not None: + # Packed rows score every position, so the head can run on the + # unsliced, still evenly sharded hidden states (opt-in, see below). It + # returns [B, T], already aligned like the padded output below. + out = maybe_compute_sharded_chunked_logps( + model, + outputs, + input_tokens, + temperature, + chunk_size, + return_entropy, + ) + aligned = out is not None + if out is None: + hidden_state = outputs[:, -logits_to_keep - 1 : -1, :] + out = compute_chunked_logps( + model, + hidden_state, + input_tokens_to_keep, + temperature, + chunk_size, + return_entropy, + ) if return_entropy: per_token_logps, per_token_entropy = out else: per_token_logps = out - if segment_ids is not None: + if segment_ids is not None and not aligned: per_token_logps = jnp.pad( per_token_logps, ((0, 0), (1, 0)), constant_values=0.0 ) @@ -996,6 +1014,392 @@ def logp_step(carry, xs): return per_token_logps +# Sharded chunked log-prob head, opt-in via TUNIX_LOGPS_SHARDED_HEAD. +# +# Under GSPMD, `compute_chunked_logps` pays per chunk: an all-gather of every +# head weight (a [4096, 248320] vocab table is 2 GB in bf16) in the forward and +# again in the rematerialized backward, and an all-reduce of the full-size table +# gradient. And because `lax.scan` iterates over the sequence axis, a sequence +# sharded over a context-parallel axis is gathered, so every context rank +# computes every chunk. The sharded head runs the same per-chunk math inside +# `shard_map`: each head parameter is all-gathered once per pass, each device +# scans only the tokens it holds, the table gradient is accumulated per device in +# float32 and reduce-scattered once, and chunks are read and written with +# dynamic slices rather than reshapes. +SHARDED_HEAD_ENV = "TUNIX_LOGPS_SHARDED_HEAD" +# Logical names of the hidden states' batch and length axes. When the active +# flax logical-axis rules (MaxText's) map them, the head uses that layout. +_LOGICAL_BATCH_AXIS = "activation_batch" +_LOGICAL_LENGTH_AXIS = "activation_length" + + +def sharded_head_mode() -> str | None: + """Returns None (off) or "local", from TUNIX_LOGPS_SHARDED_HEAD.""" + value = os.environ.get(SHARDED_HEAD_ENV, "").strip().lower() + if value in ("", "0", "false", "off", "no"): + return None + if value in ("1", "true", "on", "yes", "local"): + return "local" + raise ValueError(f"{SHARDED_HEAD_ENV}={value!r}: expected 0, 1 or local.") + + +def _mesh_axes(entry) -> tuple[str, ...]: + if entry is None: + return () + if isinstance(entry, str): + return (entry,) + return tuple(entry) + + +def _token_partition(mesh, batch: int, seq_len: int): + """Mesh axes sharding the [batch, seq] dims of the hidden states. + + Every mesh axis of size > 1 must shard one of the two, so that every device + holds distinct tokens and the gradient sum runs over all devices. + + Returns: + (batch_axes, seq_axes, source), or None if no such layout exists. + """ + sizes = dict(mesh.shape) + nontrivial = {a for a, n in sizes.items() if n > 1} + + def valid(b, s): + if set(b) & set(s) or not set(b + s) <= set(sizes): + return False + if not nontrivial <= set(b + s): + return False + return ( + batch % math.prod(sizes[a] for a in b) == 0 + and seq_len % math.prod(sizes[a] for a in s) == 0 + ) + + rules = flax.linen.get_logical_axis_rules() + if rules: + try: + spec = flax.linen.logical_to_mesh_axes( + (_LOGICAL_BATCH_AXIS, _LOGICAL_LENGTH_AXIS), rules + ) + b, s = _mesh_axes(spec[0]), _mesh_axes(spec[1]) + if valid(b, s): + return b, s, "logical-axis-rules" + except Exception: # pylint: disable=broad-exception-caught + pass + b, s = [], [] + for a in mesh.axis_names: + n = sizes[a] + if n == 1: + continue + if batch % (math.prod(sizes[x] for x in b) * n) == 0: + b.append(a) + elif seq_len % (math.prod(sizes[x] for x in s) * n) == 0: + s.append(a) + else: + return None + return tuple(b), tuple(s), "mesh-order" + + +def _use_dtypes(jaxpr, var) -> set[Any]: + """dtypes `var` is consumed in: a convert's target, else its own dtype.""" + out = set() + for eqn in jaxpr.eqns: + for k, v in enumerate(eqn.invars): + if v is not var: + continue + if eqn.primitive.name == "convert_element_type": + out.add(jnp.dtype(eqn.params["new_dtype"])) + continue + inner = eqn.params.get("jaxpr", eqn.params.get("call_jaxpr")) + inner = getattr(inner, "jaxpr", inner) + if inner is not None and len(getattr(inner, "invars", ())) == len( + eqn.invars + ): + out |= _use_dtypes(inner, inner.invars[k]) + else: + out.add(jnp.dtype(var.aval.dtype)) + return out + + +def _head_parameters(graphdef, treedef, abstract_leaves, hidden_struct): + """Leaves `compute_final_logits` reads, mapped to the dtype to gather in. + + A leaf whose every use is a cast to one narrower float dtype (MaxText casts an + fp32 vocab table to bf16 before the dot) is gathered after that cast, which is + what XLA already does for the gather inside `compute_chunked_logps`. + """ + + def head(leaves, hs): + model = nnx.merge(graphdef, jax.tree.unflatten(treedef, leaves)) + return model.compute_final_logits(hs) + + jaxpr = jax.make_jaxpr(head)(list(abstract_leaves), hidden_struct).jaxpr + out = {} + for i, var in enumerate(jaxpr.invars[: len(abstract_leaves)]): + dtypes = _use_dtypes(jaxpr, var) + if not dtypes: + continue + own = jnp.dtype(abstract_leaves[i].dtype) + dtype = next(iter(dtypes)) if len(dtypes) == 1 else own + narrower = ( + jnp.issubdtype(dtype, jnp.floating) and dtype.itemsize < own.itemsize + ) + out[i] = dtype if narrower else own + return out + + +def maybe_compute_sharded_chunked_logps( + model, + hidden_states, + input_tokens, + temperature, + chunk_size, + return_entropy, +): + """`compute_chunked_logps` for packed rows, with a shard_map'ed vocab head. + + Off unless TUNIX_LOGPS_SHARDED_HEAD is set, and it falls back (returns None, + with a log line naming the reason) whenever its preconditions do not hold: + a mesh in context, every non-trivial mesh axis sharding the batch or the + sequence, and every head parameter divisible over the mesh on some dimension. + + Args: + model: The merged actor model; must expose `compute_final_logits`. + hidden_states: [B, T, H], the model output for every position. + input_tokens: [B, T] token ids; position t is scored from position t - 1. + temperature: Sampling temperature the logits are divided by (0 or 1: none). + chunk_size: Tokens per scan step on each device. + return_entropy: Whether to also return the per-token entropy. + + Returns: + None, or the packed-row result of `compute_per_token_logps`: per-token + log-probs [B, T] whose index j scores `input_tokens[:, j]` from + `hidden_states[:, j - 1]` (index 0 is 0), plus the entropy [B, T] if + `return_entropy`. + """ + mode = sharded_head_mode() + if mode is None: + return None + + def fallback(reason): + logging.warning( + "[logps-head] sharded head off, using compute_chunked_logps: %s", reason + ) + return None + + mesh = jax.sharding.get_abstract_mesh() + if mesh.empty: + return fallback("no mesh in context (jax.set_mesh)") + batch, seq_len, hidden_dim = hidden_states.shape + partition = _token_partition(mesh, batch, seq_len) + if partition is None: + return fallback(f"no token layout for {batch}x{seq_len} on {mesh.shape}") + batch_axes, seq_axes, layout_source = partition + gather_axes = tuple(a for a in mesh.axis_names if mesh.shape[a] > 1) + num_devices = math.prod(mesh.shape[a] for a in gather_axes) + local_batch = batch // math.prod(mesh.shape[a] for a in batch_axes) + local_len = seq_len // math.prod(mesh.shape[a] for a in seq_axes) + chunk = min(chunk_size, local_len) + + graphdef, state = nnx.split(model) + leaves, treedef = jax.tree.flatten(state) + abstract = [ + jax.ShapeDtypeStruct(jnp.shape(x), jnp.result_type(x)) for x in leaves + ] + head = _head_parameters( + graphdef, + treedef, + abstract, + jax.ShapeDtypeStruct( + (local_batch, chunk, hidden_dim), hidden_states.dtype + ), + ) + if not head: + return fallback("compute_final_logits reads no model state") + head_index = sorted(head) + params = [leaves[i] for i in head_index] + gather_dtypes = [head[i] for i in head_index] + gather_dims = [] + for p in params: + dim = next((d for d, n in enumerate(p.shape) if n % num_devices == 0), None) + if dim is None: + return fallback( + f"head parameter {p.shape} not divisible by {num_devices}" + ) + gather_dims.append(dim) + param_dtypes = [jnp.result_type(p) for p in params] + num_out = 2 if return_entropy else 1 + logging.info( + "[logps-head] sharded head %s: tokens batch=%s seq=%s (%s), %dx%d local" + " tokens in %d steps of %d; gathers %s over %s", + mode, + batch_axes, + seq_axes, + layout_source, + local_batch, + local_len, + -(-local_len // chunk), + chunk, + [(p.shape, str(d)) for p, d in zip(params, gather_dtypes)], + gather_axes, + ) + + def merged(full): + values = list(abstract) + for i, x in zip(head_index, full): + values[i] = x + return nnx.merge(graphdef, jax.tree.unflatten(treedef, values)) + + def chunk_outputs(full, hs, ids, with_entropy): + logits = merged(full).compute_final_logits(hs).astype(jnp.float32) + if temperature != 0.0 and temperature != 1.0: + logits /= temperature + logps = selective_log_softmax(logits, ids) + if with_entropy: + return logps, compute_entropy_from_logits(logits) + return (logps,) + + def gather(p, dim, dtype): + p = p.astype(dtype) + if not gather_axes: + return p + return jax.lax.all_gather(p, gather_axes, axis=dim, tiled=True) + + def scatter(g, dim): + if not gather_axes: + return g + return jax.lax.psum_scatter( + g, gather_axes, scatter_dimension=dim, tiled=True + ) + + def scan_forward(full, h, t): + def step(bufs, i): + start = i * chunk + hs = jax.lax.dynamic_slice_in_dim(h, start, chunk, axis=1) + ids = jax.lax.dynamic_slice_in_dim(t, start, chunk, axis=1) + outs = chunk_outputs(full, hs, ids, return_entropy) + bufs = tuple( + jax.lax.dynamic_update_slice_in_dim(b, o, start, axis=1) + for b, o in zip(bufs, outs) + ) + return bufs, None + + bufs = tuple(jnp.zeros(t.shape, jnp.float32) for _ in range(num_out)) + bufs, _ = jax.lax.scan(step, bufs, jnp.arange(h.shape[1] // chunk)) + return bufs + + @jax.custom_vjp + def head_local(h, t, *ps): + full = [ + gather(p, d, dt) for p, d, dt in zip(ps, gather_dims, gather_dtypes) + ] + return scan_forward(full, h, t) + + def head_local_fwd(h, t, *ps): + h, t = h.value, t.value + ps = [p.value for p in ps] + full = [ + gather(p, d, dt) for p, d, dt in zip(ps, gather_dims, gather_dtypes) + ] + return scan_forward(full, h, t), (h, t, tuple(full)) + + def head_local_bwd(res, cts): + h, t, full = res + d_logps = cts[0] + d_entropy = cts[1] if return_entropy else SymbolicZero(None) + use_logps = not isinstance(d_logps, SymbolicZero) + use_entropy = not isinstance(d_entropy, SymbolicZero) + shard_shapes = [ + tuple(n // num_devices if k == d else n for k, n in enumerate(f.shape)) + for f, d in zip(full, gather_dims) + ] + if not (use_logps or use_entropy): + zeros = [jnp.zeros(s, dt) for s, dt in zip(shard_shapes, param_dtypes)] + return (jnp.zeros_like(h), None, *zeros) + + def step(carry, i): + acc, dh = carry + start = i * chunk + hs = jax.lax.dynamic_slice_in_dim(h, start, chunk, axis=1) + ids = jax.lax.dynamic_slice_in_dim(t, start, chunk, axis=1) + cot = [] + if use_logps: + cot.append(jax.lax.dynamic_slice_in_dim(d_logps, start, chunk, axis=1)) + if use_entropy: + cot.append( + jax.lax.dynamic_slice_in_dim(d_entropy, start, chunk, axis=1) + ) + + def outputs_fn(full_, hs_): + outs = chunk_outputs(full_, hs_, ids, use_entropy) + return outs if use_logps else outs[1:] + + _, vjp = jax.vjp(outputs_fn, list(full), hs) + d_full, d_hs = vjp(tuple(cot)) + acc = tuple(a + g.astype(jnp.float32) for a, g in zip(acc, d_full)) + dh = jax.lax.dynamic_update_slice_in_dim( + dh, d_hs.astype(dh.dtype), start, axis=1 + ) + return (acc, dh), None + + acc = tuple(jnp.zeros(f.shape, jnp.float32) for f in full) + (acc, dh), _ = jax.lax.scan( + step, (acc, jnp.zeros_like(h)), jnp.arange(h.shape[1] // chunk) + ) + acc = [scatter(a, d) for a, d in zip(acc, gather_dims)] + d_params = [a.astype(dt) for a, dt in zip(acc, param_dtypes)] + return (dh, None, *d_params) + + head_local.defvjp(head_local_fwd, head_local_bwd, symbolic_zeros=True) + + seq_shards = math.prod(mesh.shape[a] for a in seq_axes) + + def neighbour(x, step): + """`x` of the sequence shard `step` places later; zeros past either end.""" + if seq_shards == 1: + return jnp.zeros_like(x) + perm = [ + (i, i - step) for i in range(seq_shards) if 0 <= i - step < seq_shards + ] + return jax.lax.ppermute(x, seq_axes, perm) + + def local(h, tok, *ps): + # Position t predicts token t + 1; a shard's last position predicts the + # next shard's first token, and the sequence's last position (no target) + # is scored against token 0 and dropped. Both one-token shifts are explicit + # collective-permutes, so neither depends on how XLA partitions a shift of + # a sequence-sharded array. + t = jnp.concatenate([tok[:, 1:], neighbour(tok[:, :1], 1)], axis=1) + pad = -h.shape[1] % chunk + if pad: + h = jnp.pad(h, ((0, 0), (0, pad), (0, 0))) + t = jnp.pad(t, ((0, 0), (0, pad))) + outs = [o[:, : o.shape[1] - pad] for o in head_local(h, t, *ps)] + # Index j holds token j's log-prob (scored at position j - 1); index 0 is 0. + return tuple( + jnp.concatenate([neighbour(o[:, -1:], -1), o[:, :-1]], axis=1) + for o in outs + ) + + P = jax.sharding.PartitionSpec # pylint: disable=invalid-name + token_spec = P(batch_axes or None, seq_axes or None) + in_specs = ( + P(batch_axes or None, seq_axes or None, None), + token_spec, + *[ + P(*[gather_axes if k == d else None for k in range(p.ndim)]) + for p, d in zip(params, gather_dims) + ], + ) + hidden_states = jnp.where(jnp.isfinite(hidden_states), hidden_states, 0.0) + outs = jax.shard_map( + local, + mesh=mesh, + in_specs=in_specs, + out_specs=(token_spec,) * num_out, + check_vma=False, + )(hidden_states, input_tokens, *params) + return outs if return_entropy else outs[0] + + @nnx.jit(static_argnames=("pad_id", "eos_id", "stop_gradient")) def compute_score( model, From 743646d17bf1873297f6271bbaf5d39d220aa46f Mon Sep 17 00:00:00 2001 From: Param Bole Date: Wed, 14 Oct 2026 10:45:19 -0700 Subject: [PATCH 2/4] Fix the sharded log-prob head's docs and tests; same traced program - Docstrings and comments in tunix/rl/common.py now match the code: the four fallbacks and the inputs that raise instead, the per-device chunk_size, the gathered parameters kept for the backward, why check_vma=False, and the greedy layout search. The switch's error message lists every accepted value, and shard_shapes is computed only in the branch that reads it. The value_and_grad jaxprs hash the same as before with the switch off and on. - The tests no longer set the device count at import. Run directly they use 8 CPU devices; under pytest they run in-process only when the process already has 8 CPU devices, and otherwise in a child process, so other test modules keep their device count. - Every parity test checks that the head engaged and compares all outputs. New tests pin the token layouts, the fallback reasons, the legacy `with mesh:` fallback, the switch's scope and its parsing. --- tests/rl/sharded_logps_head_test.py | 359 ++++++++++++++++++++-------- tunix/rl/common.py | 122 +++++++--- 2 files changed, 353 insertions(+), 128 deletions(-) diff --git a/tests/rl/sharded_logps_head_test.py b/tests/rl/sharded_logps_head_test.py index 3ca056ee2b..55f09126ce 100644 --- a/tests/rl/sharded_logps_head_test.py +++ b/tests/rl/sharded_logps_head_test.py @@ -12,13 +12,24 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Tests for the sharded chunked log-prob head (TUNIX_LOGPS_SHARDED_HEAD).""" - +"""Tests for the sharded chunked log-prob head (TUNIX_LOGPS_SHARDED_HEAD). + +The tests need 8 CPU devices, and JAX fixes the device count when its backend +starts. This module never sets the count at import, because in a shared pytest +process that would change it for every other test module. Run directly, it sets +8 CPU devices for its own process. Under pytest, ShardedLogpsHeadTest runs in +the current process only if that process already has 8 CPU devices; otherwise +it is skipped and ChildProcessTest runs this file in a child process that has +them. +""" + +import contextlib import os - -os.environ['XLA_FLAGS'] = '--xla_force_host_platform_device_count=8' - -from unittest import mock # pylint: disable=g-import-not-at-top +import re +import subprocess +import sys +import unittest +from unittest import mock from absl.testing import absltest from absl.testing import parameterized @@ -29,7 +40,11 @@ from jax.sharding import Mesh, NamedSharding, PartitionSpec as P # pylint: disable=g-multiple-import,g-importing-member import numpy as np from tunix.rl import common +from tunix.tests import test_common as tc +_NUM_DEVICES = 8 +# Set in the environment of the child process that ChildProcessTest starts. +_CHILD_ENV = 'TUNIX_SHARDED_LOGPS_HEAD_TEST_CHILD' _B, _T, _H, _V = 4, 64, 32, 128 _AXES = ('data', 'fsdp', 'context', 'expert') _RULES = ( @@ -38,8 +53,17 @@ ) +def _has_devices(): + """Whether this process can build the test mesh: 8 devices on CPU.""" + return jax.default_backend() == 'cpu' and jax.device_count() >= _NUM_DEVICES + + class _HeadOnly(nnx.Module): - """Frozen hidden states into a MaxText-like head (RMSNorm, f32 table).""" + """Trainable stand-in hidden states into a MaxText-like head. + + `h` is a Param, so dL/dh is checked. The head is an RMSNorm, then a float32 + table cast to `dtype` for the dot. + """ def __init__(self, dtype, batch=_B, seq=_T): keys = jax.random.split(jax.random.PRNGKey(0), 3) @@ -74,12 +98,23 @@ def compute_final_logits(self, h): def _mesh(): - devices = np.array(jax.devices()[:8]).reshape(1, 2, 2, 2) + devices = np.array(jax.devices()[:_NUM_DEVICES]).reshape(1, 2, 2, 2) return Mesh(devices, _AXES) class ShardedLogpsHeadTest(parameterized.TestCase): + @classmethod + def setUpClass(cls): + super().setUpClass() + if not _has_devices(): + if os.environ.get(_CHILD_ENV): + raise RuntimeError(f'child process has no {_NUM_DEVICES} CPU devices') + raise unittest.SkipTest( + f'needs {_NUM_DEVICES} CPU devices; ChildProcessTest runs these' + ' tests in a child process' + ) + def setUp(self): super().setUp() self.mesh = _mesh() @@ -90,6 +125,7 @@ def _inputs(self, batch, seq): segment_ids = jnp.ones((batch, seq), jnp.int32) positions = jnp.asarray(np.tile(np.arange(seq), (batch, 1)), jnp.int32) weights = jnp.asarray(rng.normal(size=(2, batch, seq)), jnp.float32) + # Shard like _RULES when it divides: 4 = data x fsdp x expert, 2 = context. if batch % 4 == 0 and seq % 2 == 0: sharding = NamedSharding( self.mesh, P(('data', 'fsdp', 'expert'), 'context') @@ -99,6 +135,46 @@ def _inputs(self, batch, seq): ) return tokens, segment_ids, positions, weights + @contextlib.contextmanager + def _head_switch(self, mode, rules=_RULES, expect=None, legacy_mesh=False): + """Sets the switch, mesh and rules, then checks the head's engagement. + + The switch is read when compute_per_token_logps is traced and is not part + of its jit cache key, so the caches are cleared first. + + Args: + mode: TUNIX_LOGPS_SHARDED_HEAD value, or None to unset it. + rules: flax logical-axis rules. + expect: Whether each call of maybe_compute_sharded_chunked_logps in the + body engaged, in order. Default: one call, engaged iff `mode` is set. + legacy_mesh: Enter the mesh with `with mesh:` instead of `jax.set_mesh`. + + Yields: + Nothing. + """ + calls = [] + real = common.maybe_compute_sharded_chunked_logps + + def spy(*args, **kwargs): + out = real(*args, **kwargs) + calls.append(out is not None) + return out + + jax.clear_caches() + with contextlib.ExitStack() as stack: + stack.enter_context(mock.patch.dict(os.environ)) + if mode is None: + os.environ.pop(common.SHARDED_HEAD_ENV, None) + else: + os.environ[common.SHARDED_HEAD_ENV] = mode + stack.enter_context( + mock.patch.object(common, 'maybe_compute_sharded_chunked_logps', spy) + ) + stack.enter_context(self.mesh if legacy_mesh else jax.set_mesh(self.mesh)) + stack.enter_context(nn.logical_axis_rules(rules)) + yield + self.assertEqual(calls, [mode is not None] if expect is None else expect) + def _run( self, mode, @@ -109,7 +185,10 @@ def _run( batch=_B, seq=_T, return_entropy=True, + expect=None, + legacy_mesh=False, ): + """Returns (log-probs, entropy, grads) of a loss weighting both outputs.""" model = _HeadOnly(dtype, batch, seq) graphdef, state = nnx.split(model) tokens, segment_ids, positions, weights = self._inputs(batch, seq) @@ -133,27 +212,17 @@ def loss(state): value = jnp.sum(logps * weights[0]) + jnp.sum(entropy * weights[1]) return value, (logps, entropy) - jax.clear_caches() - env = {} if mode is None else {common.SHARDED_HEAD_ENV: mode} - with mock.patch.dict(os.environ, env, clear=False): - if mode is None: - os.environ.pop(common.SHARDED_HEAD_ENV, None) - with jax.set_mesh(self.mesh), nn.logical_axis_rules(rules): - (_, (logps, entropy)), grads = jax.jit( - jax.value_and_grad(loss, has_aux=True) - )(state) + with self._head_switch(mode, rules, expect, legacy_mesh): + (_, (logps, entropy)), grads = jax.jit( + jax.value_and_grad(loss, has_aux=True) + )(state) return logps, entropy, grads - @parameterized.named_parameters( - ('local', 'local', 8), - ('local_uneven_chunk', 'local', 12), - ('local_one_step', 'local', 64), - ) - def test_matches_scan_head(self, mode, chunk_size): - want = self._run(None, chunk_size=chunk_size) - got = self._run(mode, chunk_size=chunk_size) + def _assert_parity(self, got, want, entropy=True): + """f32 parity of log-probs, entropy and the h, scale and table grads.""" np.testing.assert_allclose(got[0], want[0], rtol=1e-5, atol=2e-5) - np.testing.assert_allclose(got[1], want[1], rtol=1e-5, atol=2e-5) + if entropy: + np.testing.assert_allclose(got[1], want[1], rtol=1e-5, atol=2e-5) for name in ('h', 'scale', 'table'): np.testing.assert_allclose( got[2][name][...], @@ -163,19 +232,25 @@ def test_matches_scan_head(self, mode, chunk_size): err_msg=name, ) + @parameterized.named_parameters( + ('chunk8', 8), + ('uneven_chunk12', 12), + ('one_step64', 64), + ) + def test_matches_scan_head(self, chunk_size): + want = self._run(None, chunk_size=chunk_size) + got = self._run('local', chunk_size=chunk_size) + self._assert_parity(got, want) + # The tolerance is tight enough to see a one-position misalignment. + self.assertFalse( + np.allclose(np.roll(got[0], 1, axis=1), want[0], rtol=1e-5, atol=2e-5) + ) + def test_logps_only_matches_scan_head(self): # return_entropy=False: one output, and no entropy cotangent in the bwd. want = self._run(None, return_entropy=False) got = self._run('local', return_entropy=False) - np.testing.assert_allclose(got[0], want[0], rtol=1e-5, atol=2e-5) - for name in ('h', 'scale', 'table'): - np.testing.assert_allclose( - got[2][name][...], - want[2][name][...], - rtol=1e-4, - atol=1e-5, - err_msg=name, - ) + self._assert_parity(got, want, entropy=False) @parameterized.parameters(True, False) def test_forward_only_matches_scan_head(self, return_entropy): @@ -185,30 +260,26 @@ def test_forward_only_matches_scan_head(self, return_entropy): tokens, segment_ids, positions, _ = self._inputs(_B, _T) def logps(mode): - jax.clear_caches() - env = {} if mode is None else {common.SHARDED_HEAD_ENV: mode} - with mock.patch.dict(os.environ, env, clear=False): - if mode is None: - os.environ.pop(common.SHARDED_HEAD_ENV, None) - with jax.set_mesh(self.mesh), nn.logical_axis_rules(_RULES): - return jax.jit( - lambda st: common.compute_per_token_logps( - graphdef, - st, - jnp.zeros((_B, 0), jnp.int32), - tokens, - pad_id=0, - eos_id=-1, - stop_gradient=True, - return_entropy=return_entropy, - segment_ids=segment_ids, - segment_positions=positions, - temperature=1.0, - chunk_size=8, - ) - )(state) + with self._head_switch(mode): + return jax.jit( + lambda st: common.compute_per_token_logps( + graphdef, + st, + jnp.zeros((_B, 0), jnp.int32), + tokens, + pad_id=0, + eos_id=-1, + stop_gradient=True, + return_entropy=return_entropy, + segment_ids=segment_ids, + segment_positions=positions, + temperature=1.0, + chunk_size=8, + ) + )(state) want, got = logps(None), logps('local') + self.assertEqual(jax.tree.structure(got), jax.tree.structure(want)) for g, w in zip(jax.tree.leaves(got), jax.tree.leaves(want)): np.testing.assert_allclose(g, w, rtol=1e-5, atol=2e-5) @@ -233,83 +304,169 @@ def test_bf16_matches_scan_head(self): def test_mesh_order_layout_without_rules(self): want = self._run(None, rules=()) got = self._run('local', rules=()) - np.testing.assert_allclose(got[0], want[0], rtol=1e-5, atol=2e-5) - np.testing.assert_allclose( - got[2]['table'][...], want[2]['table'][...], rtol=1e-4, atol=1e-5 - ) + self._assert_parity(got, want) def test_engages_on_mesh(self): - calls = [] - real = common.maybe_compute_sharded_chunked_logps - - def spy(*args, **kwargs): - out = real(*args, **kwargs) - calls.append(out is not None) - return out - - with mock.patch.object(common, 'maybe_compute_sharded_chunked_logps', spy): - self._run('local') - self._run(None) - self.assertEqual(calls, [True, False]) - - def _engaged(self, **kwargs): - calls = [] - real = common.maybe_compute_sharded_chunked_logps - - def spy(*args, **kw): - out = real(*args, **kw) - calls.append(out is not None) - return out - - with mock.patch.object(common, 'maybe_compute_sharded_chunked_logps', spy): - got = self._run('local', **kwargs) - return calls, got + self._run('local', expect=[True]) + self._run(None, expect=[False]) def test_sequence_only_layout_when_batch_does_not_divide(self): # 3 rows cannot be split over fsdp x expert, so all 8 devices take a # slice of every row's sequence instead. - calls, got = self._engaged(batch=3) want = self._run(None, batch=3) - self.assertEqual(calls, [True]) - np.testing.assert_allclose(got[0], want[0], rtol=1e-5, atol=2e-5) - np.testing.assert_allclose( - got[2]['table'][...], want[2]['table'][...], rtol=1e-4, atol=1e-5 - ) + got = self._run('local', batch=3) + self._assert_parity(got, want) def test_falls_back_when_no_layout(self): # 3 rows of 63 tokens divide over no mesh axis: compute_chunked_logps runs. - calls, got = self._engaged(batch=3, seq=63, chunk_size=7) want = self._run(None, batch=3, seq=63, chunk_size=7) - self.assertEqual(calls, [False]) - np.testing.assert_array_equal(got[0], want[0]) + with self.assertLogs(level='WARNING') as logs: + got = self._run('local', batch=3, seq=63, chunk_size=7, expect=[False]) + self.assertIn('no token layout', '\n'.join(logs.output)) + for g, w in zip(jax.tree.leaves(got), jax.tree.leaves(want)): + np.testing.assert_array_equal(g, w) def test_falls_back_without_mesh(self): model = _HeadOnly(jnp.float32) tokens = self._inputs(_B, _T)[0] with mock.patch.dict(os.environ, {common.SHARDED_HEAD_ENV: '1'}): - self.assertIsNone( - common.maybe_compute_sharded_chunked_logps( - model, model.h[...], tokens, 1.0, 8, True + with self.assertLogs(level='WARNING') as logs: + self.assertIsNone( + common.maybe_compute_sharded_chunked_logps( + model, model.h[...], tokens, 1.0, 8, True + ) + ) + self.assertIn('no mesh in context', '\n'.join(logs.output)) + + def test_legacy_mesh_context_falls_back(self): + # A legacy `with mesh:`, as tunix's RL cluster enters it, sets no abstract + # mesh, so the head falls back to compute_chunked_logps. + want = self._run(None, legacy_mesh=True) + with self.assertLogs(level='WARNING') as logs: + got = self._run('local', legacy_mesh=True, expect=[False]) + self.assertIn('no mesh in context', '\n'.join(logs.output)) + for g, w in zip(jax.tree.leaves(got), jax.tree.leaves(want)): + np.testing.assert_array_equal(g, w) + + def test_unpacked_rows_do_not_consult_head(self): + # Without segment_ids the switch is not read and the head is not called. + model = _HeadOnly(jnp.float32) + graphdef, state = nnx.split(model) + tokens = self._inputs(_B, _T)[0] + with self._head_switch('local', expect=[]): + jax.jit( + lambda st: common.compute_per_token_logps( + graphdef, + st, + tokens[:, :8], + tokens[:, 8:], + pad_id=0, + eos_id=-1, + chunk_size=8, ) - ) + )(state) + + def test_chunk_size_zero_does_not_consult_head(self): + model = _HeadOnly(jnp.float32) + graphdef, state = nnx.split(model) + tokens, segment_ids, positions, _ = self._inputs(_B, _T) + with self._head_switch('local', expect=[]): + jax.jit( + lambda st: common.compute_per_token_logps( + graphdef, + st, + jnp.zeros((_B, 0), jnp.int32), + tokens, + pad_id=0, + eos_id=-1, + segment_ids=segment_ids, + segment_positions=positions, + chunk_size=0, + ) + )(state) + + def test_unknown_switch_value_raises(self): + with self.assertRaisesRegex(ValueError, common.SHARDED_HEAD_ENV): + self._run('locla') @parameterized.parameters( ('', None), ('0', None), ('false', None), + ('off', None), + ('no', None), ('1', 'local'), + ('true', 'local'), + ('on', 'local'), + ('yes', 'local'), ('local', 'local'), + (' Local ', 'local'), + ('TRUE', 'local'), ) def test_mode_parsing(self, value, want): with mock.patch.dict(os.environ, {common.SHARDED_HEAD_ENV: value}): self.assertEqual(common.sharded_head_mode(), want) - @parameterized.parameters('locla', 'per_step') + @parameterized.parameters('locla', '2') def test_mode_parsing_rejects_unknown(self, value): with mock.patch.dict(os.environ, {common.SHARDED_HEAD_ENV: value}): - with self.assertRaises(ValueError): + with self.assertRaisesRegex(ValueError, common.SHARDED_HEAD_ENV): common.sharded_head_mode() + @parameterized.named_parameters( + ( + 'rules', + _RULES, + 4, + 64, + (('data', 'fsdp', 'expert'), ('context',), 'logical-axis-rules'), + ), + ('no_rules', (), 4, 64, (('fsdp', 'context'), ('expert',), 'mesh-order')), + ( + 'batch_does_not_divide', + _RULES, + 3, + 64, + ((), ('fsdp', 'context', 'expert'), 'mesh-order'), + ), + ('no_layout', _RULES, 3, 63, None), + ) + def test_token_layout(self, rules, batch, seq, want): + # The layouts the parity and fallback tests above rely on. + with nn.logical_axis_rules(rules): + got = common._token_partition(self.mesh, batch, seq) # pylint: disable=protected-access + self.assertEqual(got, want) + + +class ChildProcessTest(absltest.TestCase): + """Runs ShardedLogpsHeadTest in a child process when this one cannot.""" + + def test_runs_head_tests_in_child_process(self): + if os.environ.get(_CHILD_ENV) or _has_devices(): + self.skipTest('ShardedLogpsHeadTest runs in this process') + env = dict(os.environ, JAX_PLATFORMS='cpu') + env[_CHILD_ENV] = '1' + # The child imports the same tunix as this process. + root = os.path.dirname( + os.path.dirname(os.path.dirname(os.path.abspath(common.__file__))) + ) + env['PYTHONPATH'] = os.pathsep.join( + p for p in (root, env.get('PYTHONPATH')) if p + ) + result = subprocess.run( + [sys.executable, os.path.abspath(__file__)], + env=env, + capture_output=True, + text=True, + check=False, + timeout=1800, + ) + self.assertEqual(result.returncode, 0, result.stderr[-8000:]) + ran = re.search(r'^Ran (\d+) tests? in', result.stderr, re.M) + self.assertIsNotNone(ran, result.stderr[-8000:]) + self.assertGreater(int(ran.group(1)), 1) + if __name__ == '__main__': + tc.safe_set_n_cpu_devices(_NUM_DEVICES) absltest.main() diff --git a/tunix/rl/common.py b/tunix/rl/common.py index 2c51c0aa3f..fb4c6cfe3a 100644 --- a/tunix/rl/common.py +++ b/tunix/rl/common.py @@ -539,7 +539,11 @@ def compute_per_token_logps( segment_positions: optional 1D local position indices used for packing. temperature: temperature used for rollout. chunk_size: If not 0, computes the log probabilities in sequence chunks of - this size. + this size. With TUNIX_LOGPS_SHARDED_HEAD on and packed rows + (`segment_ids` set), the sharded head, when it engages (see + `maybe_compute_sharded_chunked_logps`), instead has each device scan its + own block of the sequence in chunks of `min(chunk_size, block length)` + positions. routed_experts: Optional `[B, T, num_layers, top_k]` MoE experts captured during rollout. Forwarded to the model as `forced_routed_experts` -- the name MaxText's adapter uses -- so it replays this routing instead of @@ -645,7 +649,8 @@ def _result(*outs): out = None if segment_ids is not None: # Packed rows score every position, so the head can run on the - # unsliced, still evenly sharded hidden states (opt-in, see below). It + # unsliced, still evenly sharded hidden states. Opt-in via + # TUNIX_LOGPS_SHARDED_HEAD; None when off or when it falls back. It # returns [B, T], already aligned like the padded output below. out = maybe_compute_sharded_chunked_logps( model, @@ -1022,10 +1027,11 @@ def logp_step(carry, xs): # gradient. And because `lax.scan` iterates over the sequence axis, a sequence # sharded over a context-parallel axis is gathered, so every context rank # computes every chunk. The sharded head runs the same per-chunk math inside -# `shard_map`: each head parameter is all-gathered once per pass, each device -# scans only the tokens it holds, the table gradient is accumulated per device in -# float32 and reduce-scattered once, and chunks are read and written with -# dynamic slices rather than reshapes. +# `shard_map`: each head parameter is all-gathered once, in the forward, and the +# backward reuses that copy (so it stays live from the forward to the +# backward); each device scans only the tokens it holds; the table gradient is +# accumulated per device in float32 and reduce-scattered once; and chunks are +# read and written with dynamic slices rather than reshapes. SHARDED_HEAD_ENV = "TUNIX_LOGPS_SHARDED_HEAD" # Logical names of the hidden states' batch and length axes. When the active # flax logical-axis rules (MaxText's) map them, the head uses that layout. @@ -1034,13 +1040,23 @@ def logp_step(carry, xs): def sharded_head_mode() -> str | None: - """Returns None (off) or "local", from TUNIX_LOGPS_SHARDED_HEAD.""" + """Returns None (off) or "local", from TUNIX_LOGPS_SHARDED_HEAD. + + Case and surrounding whitespace are ignored. Off: unset, empty, 0, false, off, + no. On: local, 1, true, on, yes. + + Raises: + ValueError: The variable has any other value. + """ value = os.environ.get(SHARDED_HEAD_ENV, "").strip().lower() if value in ("", "0", "false", "off", "no"): return None if value in ("1", "true", "on", "yes", "local"): return "local" - raise ValueError(f"{SHARDED_HEAD_ENV}={value!r}: expected 0, 1 or local.") + raise ValueError( + f"{SHARDED_HEAD_ENV}={value!r}: expected local, 1, true, on or yes (on)," + " or empty, 0, false, off or no (off)." + ) def _mesh_axes(entry) -> tuple[str, ...]: @@ -1055,10 +1071,17 @@ def _token_partition(mesh, batch: int, seq_len: int): """Mesh axes sharding the [batch, seq] dims of the hidden states. Every mesh axis of size > 1 must shard one of the two, so that every device - holds distinct tokens and the gradient sum runs over all devices. + holds distinct tokens and the gradient sum runs over all devices. Tries the + layout the active flax logical-axis rules give `activation_batch` and + `activation_length` (an error while reading them skips this step), then one + greedy pass over the mesh axes in mesh order: each axis of size > 1 shards + the batch if the batch still divides, else the sequence if it still divides, + else the pass fails. Returns: - (batch_axes, seq_axes, source), or None if no such layout exists. + (batch_axes, seq_axes, source), with source "logical-axis-rules" or + "mesh-order"; or None if both steps fail. The greedy pass can miss a + layout that exists. """ sizes = dict(mesh.shape) nontrivial = {a for a, n in sizes.items() if n > 1} @@ -1099,7 +1122,14 @@ def valid(b, s): def _use_dtypes(jaxpr, var) -> set[Any]: - """dtypes `var` is consumed in: a convert's target, else its own dtype.""" + """dtypes `var` is consumed in: a convert's target, else its own dtype. + + Recurses into an equation's inner jaxpr (its `jaxpr` or `call_jaxpr` param) + when that jaxpr has one input per equation input, assuming they line up one + to one. Known gap: a pass-through output records no use. If the inner jaxpr + returns `var` unchanged as one of its outputs, a later use of that output is + missed. + """ out = set() for eqn in jaxpr.eqns: for k, v in enumerate(eqn.invars): @@ -1125,6 +1155,10 @@ def _head_parameters(graphdef, treedef, abstract_leaves, hidden_struct): A leaf whose every use is a cast to one narrower float dtype (MaxText casts an fp32 vocab table to bf16 before the dot) is gathered after that cast, which is what XLA already does for the gather inside `compute_chunked_logps`. + + Returns: + {index into `abstract_leaves`: dtype to gather that leaf in}, for every + leaf `compute_final_logits` reads. """ def head(leaves, hs): @@ -1156,17 +1190,28 @@ def maybe_compute_sharded_chunked_logps( ): """`compute_chunked_logps` for packed rows, with a shard_map'ed vocab head. - Off unless TUNIX_LOGPS_SHARDED_HEAD is set, and it falls back (returns None, - with a log line naming the reason) whenever its preconditions do not hold: - a mesh in context, every non-trivial mesh axis sharding the batch or the - sequence, and every head parameter divisible over the mesh on some dimension. + Off unless TUNIX_LOGPS_SHARDED_HEAD turns it on (see `sharded_head_mode`). + The switch is read when the caller is traced and is not part of any jit + cache key, so a change after the first trace takes effect only on a retrace. + + Returns None, after logging + "[logps-head] sharded head off, using compute_chunked_logps: ", when: + * no mesh is set with `jax.set_mesh` (a legacy `with mesh:` sets none); + * `_token_partition` finds no token layout; + * `compute_final_logits` reads no model state; + * a head parameter has no dimension divisible by the number of devices + (a scalar parameter never has one). + Some other unsupported inputs raise instead of falling back, for example a + mesh with Explicit axes (inside `jax.shard_map`) or logical-axis rules that + list two sequence axes out of mesh order (in `jax.lax.ppermute`). Args: - model: The merged actor model; must expose `compute_final_logits`. + model: The merged model; must expose `compute_final_logits`. hidden_states: [B, T, H], the model output for every position. input_tokens: [B, T] token ids; position t is scored from position t - 1. temperature: Sampling temperature the logits are divided by (0 or 1: none). - chunk_size: Tokens per scan step on each device. + chunk_size: Sequence positions per scan step on each device, capped at the + length of the device's block of the sequence. return_entropy: Whether to also return the per-token entropy. Returns: @@ -1174,6 +1219,9 @@ def maybe_compute_sharded_chunked_logps( log-probs [B, T] whose index j scores `input_tokens[:, j]` from `hidden_states[:, j - 1]` (index 0 is 0), plus the entropy [B, T] if `return_entropy`. + + Raises: + ValueError: TUNIX_LOGPS_SHARDED_HEAD has an unrecognized value. """ mode = sharded_head_mode() if mode is None: @@ -1236,13 +1284,16 @@ def fallback(reason): layout_source, local_batch, local_len, - -(-local_len // chunk), + -(-local_len // chunk), # ceil(local_len / chunk) scan steps chunk, [(p.shape, str(d)) for p, d in zip(params, gather_dtypes)], gather_axes, ) def merged(full): + # Only the head leaves carry values. The others stay ShapeDtypeStruct + # placeholders, which compute_final_logits never reads (_head_parameters + # found the leaves it reads). values = list(abstract) for i, x in zip(head_index, full): values[i] = x @@ -1294,6 +1345,11 @@ def head_local(h, t, *ps): return scan_forward(full, h, t) def head_local_fwd(h, t, *ps): + # With symbolic_zeros=True the inputs arrive as CustomVJPPrimal wrappers, + # hence `.value`. Their `.perturbed` flags are ignored, so the backward + # builds the f32 accumulator even for a frozen head. The gathered + # parameters are the residuals: the backward reuses them instead of + # gathering again. h, t = h.value, t.value ps = [p.value for p in ps] full = [ @@ -1304,14 +1360,18 @@ def head_local_fwd(h, t, *ps): def head_local_bwd(res, cts): h, t, full = res d_logps = cts[0] + # Without return_entropy the primal has one output, so no entropy + # cotangent: treat it as a symbolic zero. d_entropy = cts[1] if return_entropy else SymbolicZero(None) use_logps = not isinstance(d_logps, SymbolicZero) use_entropy = not isinstance(d_entropy, SymbolicZero) - shard_shapes = [ - tuple(n // num_devices if k == d else n for k, n in enumerate(f.shape)) - for f, d in zip(full, gather_dims) - ] if not (use_logps or use_entropy): + shard_shapes = [ + tuple( + n // num_devices if k == d else n for k, n in enumerate(f.shape) + ) + for f, d in zip(full, gather_dims) + ] zeros = [jnp.zeros(s, dt) for s, dt in zip(shard_shapes, param_dtypes)] return (jnp.zeros_like(h), None, *zeros) @@ -1329,6 +1389,7 @@ def step(carry, i): ) def outputs_fn(full_, hs_): + # Exactly the outputs that have a cotangent in `cot`, in that order. outs = chunk_outputs(full_, hs_, ids, use_entropy) return outs if use_logps else outs[1:] @@ -1363,12 +1424,13 @@ def neighbour(x, step): def local(h, tok, *ps): # Position t predicts token t + 1; a shard's last position predicts the - # next shard's first token, and the sequence's last position (no target) - # is scored against token 0 and dropped. Both one-token shifts are explicit - # collective-permutes, so neither depends on how XLA partitions a shift of - # a sequence-sharded array. + # next shard's first token. The sequence's last position has no next + # token: its target is token id 0 (`neighbour` returns zeros past the last + # shard) and its score is dropped by the output shift below. Both one-token + # shifts are explicit collective-permutes, so neither depends on how XLA + # partitions a shift of a sequence-sharded array. t = jnp.concatenate([tok[:, 1:], neighbour(tok[:, :1], 1)], axis=1) - pad = -h.shape[1] % chunk + pad = -h.shape[1] % chunk # up to the next multiple of chunk if pad: h = jnp.pad(h, ((0, 0), (0, pad), (0, 0))) t = jnp.pad(t, ((0, 0), (0, pad))) @@ -1389,12 +1451,18 @@ def local(h, tok, *ps): for p, d in zip(params, gather_dims) ], ) + # Zero non-finite hidden states, as compute_chunked_logps does. hidden_states = jnp.where(jnp.isfinite(hidden_states), hidden_states, 0.0) outs = jax.shard_map( local, mesh=mesh, in_specs=in_specs, out_specs=(token_spec,) * num_out, + # The scan carries start as jnp.zeros, which the vma check types as the + # same on every device, while each chunk's outputs vary by device, so + # with the check lax.scan rejects the carry. Unchecked is still exact: + # every mesh axis of size > 1 is in every in and out spec, so the + # transpose adds no psum or division over a nontrivial axis. check_vma=False, )(hidden_states, input_tokens, *params) return outs if return_entropy else outs[0] From 1301c60c66c742e64e99b9b301b885ec91ff85b7 Mon Sep 17 00:00:00 2001 From: Param Bole Date: Wed, 14 Oct 2026 11:57:00 -0700 Subject: [PATCH 3/4] Correct two docstrings of the sharded log-prob head (Explicit-axis raise condition; parameter discovery gap) --- tunix/rl/common.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/tunix/rl/common.py b/tunix/rl/common.py index fb4c6cfe3a..77cac3e516 100644 --- a/tunix/rl/common.py +++ b/tunix/rl/common.py @@ -1158,7 +1158,9 @@ def _head_parameters(graphdef, treedef, abstract_leaves, hidden_struct): Returns: {index into `abstract_leaves`: dtype to gather that leaf in}, for every - leaf `compute_final_logits` reads. + leaf `compute_final_logits` reads that `_use_dtypes` sees. A leaf read only + through an output of an inner call (for example an identity `jax.jit`, + `jax.checkpoint` or `custom_jvp`) is missed; see `_use_dtypes`. """ def head(leaves, hs): @@ -1201,9 +1203,10 @@ def maybe_compute_sharded_chunked_logps( * `compute_final_logits` reads no model state; * a head parameter has no dimension divisible by the number of devices (a scalar parameter never has one). - Some other unsupported inputs raise instead of falling back, for example a - mesh with Explicit axes (inside `jax.shard_map`) or logical-axis rules that - list two sequence axes out of mesh order (in `jax.lax.ppermute`). + Some other unsupported inputs raise instead of falling back, for example, on + a mesh with Explicit axes, an input whose sharding differs from the head's + `in_specs` (inside `jax.shard_map`), or logical-axis rules that list two + sequence axes out of mesh order (in `jax.lax.ppermute`). Args: model: The merged model; must expose `compute_final_logits`. From 6d2fb669470f60a73f77fa05b64f1241817ed5c1 Mon Sep 17 00:00:00 2001 From: Param Bole Date: Wed, 14 Oct 2026 12:06:33 -0700 Subject: [PATCH 4/4] Forward TUNIX_LOGPS_SHARDED_HEAD to the trainer in k8s_launcher.sh The deepswe_dist launcher gave no way to turn the sharded log-prob head on in an RL run. TRAINER_EXTRA_ENV cannot carry it from the 397B v7x recipes: a value set there replaces mlperf_397b_256_v7x.sh's default TRAINER_EXTRA_ENV, including the trainer's LIBTPU_INIT_ARGS. When TUNIX_LOGPS_SHARDED_HEAD is set, the launcher now adds it to the trainer's start command only. Unset, nothing is rendered, so every recipe renders as before. Tests: unset, the variable appears nowhere in the render; set to local, it reaches the trainer JobSet and not the orchestrator or rollout JobSets. --- .../deepswe_dist/mlperf_submission_launcher_test.py | 13 +++++++++++++ .../examples/deepswe_dist/k8s_launcher.sh | 5 +++++ 2 files changed, 18 insertions(+) diff --git a/tests/experimental/examples/deepswe_dist/mlperf_submission_launcher_test.py b/tests/experimental/examples/deepswe_dist/mlperf_submission_launcher_test.py index 945579386e..132b88825e 100644 --- a/tests/experimental/examples/deepswe_dist/mlperf_submission_launcher_test.py +++ b/tests/experimental/examples/deepswe_dist/mlperf_submission_launcher_test.py @@ -183,6 +183,19 @@ def test_dropless_rejects_other_values(self): with self.assertRaisesRegex(AssertionError, "DROPLESS must be true or false"): _render(DROPLESS="yes") + def test_the_sharded_log_prob_head_is_off_unless_set(self): + self.assertNotIn("TUNIX_LOGPS_SHARDED_HEAD", self.default) + + def test_the_sharded_log_prob_head_reaches_the_trainer_only(self): + render = _render(TUNIX_LOGPS_SHARDED_HEAD="local") + self.assertRegex( + _jobset(render, "train"), r"\bTUNIX_LOGPS_SHARDED_HEAD=local\b" + ) + for suffix in ("orch", "roll-0"): + self.assertNotIn( + "TUNIX_LOGPS_SHARDED_HEAD", _jobset(render, suffix), suffix + ) + def test_default_batch_gives_the_gbs_1024_hyperparameters(self): orch = _jobset(self.default, "orch") self.assertIn("--batch_size=64 ", orch) diff --git a/tunix/experimental/examples/deepswe_dist/k8s_launcher.sh b/tunix/experimental/examples/deepswe_dist/k8s_launcher.sh index 8a9fa9cb38..badceca8e4 100755 --- a/tunix/experimental/examples/deepswe_dist/k8s_launcher.sh +++ b/tunix/experimental/examples/deepswe_dist/k8s_launcher.sh @@ -287,6 +287,10 @@ export PREEMPTIBLE=${PREEMPTIBLE:-${preemptible:-false}} export PRIORITY_CLASS=${PRIORITY_CLASS:-medium} export TRAINER_EXTRA_ENV=${TRAINER_EXTRA_ENV:-} +# The trainer process's opt-in sharded log-prob head (TUNIX_LOGPS_SHARDED_HEAD +# in tunix/rl/common.py: local turns it on; unset or 0 leaves it off). Passed to +# the trainer only, and only when set. +export TUNIX_LOGPS_SHARDED_HEAD=${TUNIX_LOGPS_SHARDED_HEAD:-} # Extra KEY=VALUE env for the orchestrator and rollout processes (the recipes set # Raiden tuning and weight-sync timeouts here). The VLLM_RAY_* lists tell vLLM # which of the rollout driver's env vars to forward to its Ray workers on the @@ -670,6 +674,7 @@ start_trainer() { ${TRAINER_MAXTEXT_ATTENTION:+TRAINER_MAXTEXT_ATTENTION=\"${TRAINER_MAXTEXT_ATTENTION}\"} \ ${raiden_env} \ ${MAXTEXT_EXTRA_FLAGS:+MAXTEXT_EXTRA_FLAGS=\"${MAXTEXT_EXTRA_FLAGS}\"} \ + ${TUNIX_LOGPS_SHARDED_HEAD:+TUNIX_LOGPS_SHARDED_HEAD=${TUNIX_LOGPS_SHARDED_HEAD}} \ ${TRAINER_EXTRA_ENV:+${TRAINER_EXTRA_ENV}} \ RAIDEN_DEVICES_PER_HOST=${RAIDEN_DEVICES_PER_HOST} \ USE_WEIGHT_CONVERTER=${USE_WEIGHT_CONVERTER} \