From 1435de05ef69b998a5b441882d3fba91985014eb Mon Sep 17 00:00:00 2001 From: Param Bole Date: Thu, 15 Oct 2026 06:00:33 -0700 Subject: [PATCH 1/4] packing: optionally put each packed row's padding first (pack_pad_position) The packer writes a row's segments from position 0 and leaves its padding at the end. Under load-balanced context parallelism (MaxText's DUAL_CHUNK_SWAP: context rank r holds sequence chunks r and 2 * cp - 1 - r) the last chunk of a row is rank 0's, so trailing padding makes only rank 0 lighter, while ranks 1..cp-1 still attend every real query to every earlier position. tokamax's Splash (0.0.16.dev20261008 or later) skips the tiles whose query and key segment-id ranges do not overlap, so once the padding comes first every real query skips the padding's key tiles; an older tokamax masks segment ids element by element, and there "start" gains nothing. On the RL trainer benchmark's 60,948-token rows at context 4, counted with that rule and the benchmark's blocks (forward 1024 x 4096, dkv 2048 x 2048), the forward tiles per rank go from 80/136/136/136 to 125/120/120/120 and the dkv tiles from 74/132/132/132 to 122/116/116/116: the busiest rank, which sets the step time, does about 8% fewer. pad_position="start" (pack_chunk, pack_bin, pack_core) lays a row out exactly as "end" does, then writes it further right by its trailing padding, counted from its last token and rounded down to a multiple of pad_align (default 4096, the largest Splash block of the Qwen3.5-397B-A17B recipe), so every segment start keeps its place relative to the kernel's tiles and its segment alignment. Every per-token array moves with its token: ids, masks, advantages, segment ids, in-segment positions, per-token values and routed experts, so a model whose attention, convolutions and recurrences respect segment ids computes the same values. A row with no tokens, or with less trailing padding than pad_align, does not move, so the gain depends on how much padding the rows carry. "end" stays the default and packs byte for byte as before; pad_align is not read with "end". The option is configuration, threaded through the orchestrator the way segment_align_multiple is: BatchConfig.pack_pad_position / pack_pad_align -> create_batch_assembler -> SequencePackedBatchAssembler -> pack_chunk, so the program and its trainer warm-up payload pack alike. run_deepswe_dist.py takes --pack_pad_position / --pack_pad_align (main() builds its BatchConfig with _batch_config, so a test can check that the flags reach it), and k8s_launcher.sh passes them from PACK_PAD_POSITION / PACK_PAD_ALIGN only when they are set. Only the orchestrator is wired: pack_sequences in tunix/rl/utils.py, which rl_learner and agentic_rl_learner use, still puts the padding at the end. pack_chunk, pack_bin and pack_core raise on an unknown position and, with "start", on a pad_align that is not a positive multiple of segment_align_multiple; they leave a row that cannot move as it is. The orchestrator also refuses settings under which the option would do nothing, so a run cannot measure the default by accident: "start" with the padded (unpacked) assembler, a pad_align of at least the row length, or a non-default pad_align with "end". The overlay that the measurements below ran with read the switch from the environment inside packing.py (TUNIX_PACK_PAD_POSITION / TUNIX_PACK_PAD_ALIGN). Nothing reads those two variables any more: set PACK_PAD_POSITION / PACK_PAD_ALIGN, or the flags. Nothing in packing.py reads the environment and every other packer option is an argument, so this one is too. Measured on Qwen3.5-397B-A17B, 128 TPU v7x, on the RL trainer benchmark (AI-Hypercomputer/maxtext#5701 driver, mesh fsdp 32 x context 4 x expert 2, ragged_buffer_factor 2.75, 9 micro-batches per step), on stack 3. The benchmark builds its own rows and its driver's PAD_POSITION=start moves them; the overlay's switch changed only the trainer warm-up payload (same shapes and dtypes). So the numbers below measure the driver's layout, which is the layout this packer produces (checked field by field on the same rows). Every regular benchmark row carried 4,588 tokens of padding and moved by 4,096; rows packed from real rollouts may carry less padding and move less, or not at all. Regular micro-batch 10.033 -> 9.876 s and 99.50 -> 98.32 s per step, mean of steps 1-4 (o275x-s3-pf-prof against the stack 3 control o275x-gdn-ctl, run in another pod). Numerics (gates off, random advantages, against two reference runs whose gradient norms differ by 1.9e-5 relative; nxs3-pf): all six gates pass, loss 2.7e-7 relative, log-prob difference 7.0e-6, entropy 8.0e-8 relative, gradient norm 2.3e-4 relative (pass limit 2e-3). Tile boundaries move, so reductions round differently. Tests: - packing_test.py: "start" rows written out by hand (the shift, its rounding, segment alignment, routing and per-token values moving with their tokens, the default pad_align, and a shift counted from the last token when the budget is not a multiple of the segment alignment); 200 random packings in which "start" is the "end" row rotated right by its rounded trailing padding in every array (segment alignments 1 to 128, budgets that are and are not multiples of them, every per-token field, items without routing, empty items and bins); every token keeps its id, segment id, in-segment position, masks, advantage, per-token values and routing, checked against the items; "end" is unchanged by any pad_align, also at segment alignments that do not divide 4096; invalid settings raise. - batch_assembly_test.py: the assembler puts the padding first; the options reach it from BatchConfig on both packed paths and are refused on the padded one; settings that would do nothing raise. - grpo_recipe_wiring_test.py, mlperf_submission_launcher_test.py: the flags parse and reach BatchConfig; the default command line gives the previous BatchConfig field for field; rows packed from the parsed flags by the assembler put their padding first, at the default and at a given pad_align; main() hands the program the BatchConfig from _batch_config; the launcher passes the flags only when set, changing only the orchestrator's JobSet. Forty-two deliberately broken copies each fail at least one of these tests: an array left unshifted, given absolute positions or one token long; the shift unrounded, off by one, a float, blind to the segment alignment or rounded up to it; segment ids widened to int64; pad_align parsed or defaulted as a float; the option dropped, or taken from the wrong variable, at any hop from the launcher to pack_chunk, including main() building its own BatchConfig; a check removed, loosened or applied to "end"; "start" as the default or always passed. --- .../mlperf_submission_launcher_test.py | 16 + .../orchestrator/batch_assembly_test.py | 108 +++++ .../orchestrator/grpo_recipe_wiring_test.py | 103 +++++ tests/rl/packing_test.py | 407 ++++++++++++++++++ .../examples/deepswe_dist/k8s_launcher.sh | 7 + .../examples/deepswe_dist/run_deepswe_dist.py | 53 ++- .../orchestrator/batch_assembly.py | 54 ++- tunix/rl/packing.py | 149 ++++++- 8 files changed, 878 insertions(+), 19 deletions(-) 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 b387a602dd..f5c7365014 100644 --- a/tests/experimental/examples/deepswe_dist/mlperf_submission_launcher_test.py +++ b/tests/experimental/examples/deepswe_dist/mlperf_submission_launcher_test.py @@ -363,6 +363,22 @@ def test_staging_without_the_pipelined_boundary_is_refused(self): ): _render(PIPELINE_STEP_BOUNDARY="false", STAGE_NEXT_MICROBATCH="true") + def test_packed_rows_pad_at_the_end_by_default(self): + orch = _jobset(self.default, "orch") + self.assertNotIn("--pack_pad_position", orch) + self.assertNotIn("--pack_pad_align", orch) + + def test_the_packed_row_pad_position_reaches_only_the_orchestrator(self): + render = _render(PACK_PAD_POSITION="start", PACK_PAD_ALIGN="8192") + orch = _jobset(render, "orch") + self.assertRegex(orch, r"--pack_pad_position=start\b") + self.assertRegex(orch, r"--pack_pad_align=8192\b") + docs = render.split("\n---\n") + default_docs = self.default.split("\n---\n") + self.assertLen(docs, len(default_docs)) + changed = [doc for doc, before in zip(docs, default_docs) if doc != before] + self.assertEqual(changed, [orch]) + def test_the_sharded_log_prob_head_is_on_in_the_trainer_only(self): self.assertRegex( _jobset(self.default, "train"), r"\bTUNIX_LOGPS_SHARDED_HEAD=local\b" diff --git a/tests/experimental/orchestrator/batch_assembly_test.py b/tests/experimental/orchestrator/batch_assembly_test.py index 4b710170a3..856359643c 100644 --- a/tests/experimental/orchestrator/batch_assembly_test.py +++ b/tests/experimental/orchestrator/batch_assembly_test.py @@ -326,6 +326,61 @@ def test_sequence_packed_assembler_default_64_token_alignment(self): np.testing.assert_array_equal(payload.completion_mask[0, 30:64], 0) np.testing.assert_array_equal(payload.completion_ids[0, 30:64], 42) + def test_pad_position_start_puts_each_rows_padding_first(self): + items = [ + _make_payload( + 2, 2, advantage=1.5, old_logps=np.full(2, -0.5, np.float32) + ), + _make_payload( + 3, 1, advantage=-0.5, old_logps=np.full(1, -0.7, np.float32) + ), + ] + (end,) = self._drain(self._make_assembler(max_packed_len=16), items) + (start,) = self._drain( + self._make_assembler( + max_packed_len=16, pad_position="start", pad_align=4 + ), + items, + ) + # 8 real tokens and 8 of padding: every per-token field moves right by 8. + for name in ( + "completion_ids", + "completion_mask", + "advantages", + "segment_ids", + "segment_positions", + "old_per_token_logps", + ): + np.testing.assert_array_equal( + getattr(start, name), + np.roll(getattr(end, name), 8, axis=1), + err_msg=name, + ) + self.assertEqual(start.prompt_ids.shape, (1, 0)) + self.assertEqual(start.num_segments, end.num_segments) + + def test_pad_settings_that_do_nothing_or_break_alignment_raise(self): + for kwargs, regex in ( + (dict(pad_position="middle"), "pad_position must be one of"), + ( + dict( + pad_position="start", + pad_align=100, + max_packed_len=8192, + segment_align_multiple=64, + ), + "positive multiple", + ), + ( + dict(pad_position="start", pad_align=16, max_packed_len=16), + "never moves a row", + ), + (dict(pad_align=512), "only applies with pad_position='start'"), + ): + with self.subTest(kwargs=kwargs): + with self.assertRaisesRegex(ValueError, regex): + self._make_assembler(**kwargs) + def test_sequence_packed_assembler_rejects_partial_optional_fields(self): items = [ _make_payload(2, 2), @@ -2243,6 +2298,59 @@ def test_create_default_batch_assembler(self): self.assertEqual(assembler.batch_size, 2) self.assertEqual(assembler.max_packed_len, 8192) + def test_create_sequence_packed_assembler_pads_at_the_end_by_default(self): + assembler = batch_assembly.create_batch_assembler( + num_generations=2, + mini_batch_size=2, + train_micro_batch_size=1, + batch_config=batch_assembly.BatchConfig(max_seq_token_per_tpu=1024), + ) + self.assertEqual(assembler.pad_position, "end") + self.assertEqual(assembler.pad_align, 4096) + + def test_create_sequence_packed_assembler_passes_the_pad_position(self): + for config in ( + batch_assembly.BatchConfig( + max_seq_token_per_tpu=1024, + pack_pad_position="start", + pack_pad_align=256, + ), + # No budget and no prompt length: the default packed assembler. + batch_assembly.BatchConfig( + pack_pad_position="start", pack_pad_align=256 + ), + ): + with self.subTest(config=config): + assembler = batch_assembly.create_batch_assembler( + num_generations=2, + mini_batch_size=2, + train_micro_batch_size=1, + batch_config=config, + ) + self.assertIsInstance( + assembler, batch_assembly.SequencePackedBatchAssembler + ) + self.assertEqual(assembler.pad_position, "start") + self.assertEqual(assembler.pad_align, 256) + + def test_create_padded_batch_assembler_refuses_the_pad_position(self): + # The padded assembler does not pack, so the options would do nothing. + for overrides in ( + dict(pack_pad_position="start"), + dict(pack_pad_align=512), + ): + with self.subTest(overrides=overrides): + with self.assertRaisesRegex(ValueError, "does not pack"): + batch_assembly.create_batch_assembler( + num_generations=2, + mini_batch_size=2, + train_micro_batch_size=4, + batch_config=batch_assembly.BatchConfig( + max_prompt_length=128, + max_response_length=256, + **overrides, + ), + ) def test_create_sequence_packed_assembler_pack_size_includes_expert(self): assembler = batch_assembly.create_batch_assembler( diff --git a/tests/experimental/orchestrator/grpo_recipe_wiring_test.py b/tests/experimental/orchestrator/grpo_recipe_wiring_test.py index 04a95ce4dd..7fa0720bb0 100644 --- a/tests/experimental/orchestrator/grpo_recipe_wiring_test.py +++ b/tests/experimental/orchestrator/grpo_recipe_wiring_test.py @@ -14,6 +14,8 @@ """Unit tests for wiring the sequence-level GRPO options to the loss.""" +import ast +import inspect import os import re @@ -697,6 +699,107 @@ def test_k8s_launcher_wires_trainer_warmup_compile_flags(self): self.assertTrue(parsed.trainer_warmup_compile) self.assertTrue(parsed.trainer_warmup_routed_experts) + def test_k8s_launcher_wires_the_packed_row_pad_position(self): + block = _deepswe_k8s_orchestrator_block() + for flag in ("--pack_pad_position=", "--pack_pad_align="): + self.assertIn(flag, block, f"deepswe k8s_launcher does not pass {flag}") + + parse = run_deepswe_dist._parse_args # pylint: disable=protected-access + build = run_deepswe_dist._batch_config # pylint: disable=protected-access + defaults = parse([]) + self.assertEqual(defaults.pack_pad_position, "end") + self.assertEqual(defaults.pack_pad_align, 4096) + # The default command line gives the BatchConfig the program had before + # the two options existed, field for field. + self.assertEqual( + build(defaults, pad_id=7), + batch_assembly.BatchConfig( + pad_id=7, + max_prompt_length=defaults.max_prompt_length, + max_response_length=defaults.max_response_length, + max_seq_token_per_tpu=None, + max_segments_per_packed_row=None, + trainer_fsdp=None, + trainer_dp=None, + trainer_expert=None, + ), + ) + config = build( + parse([ + "--max_seq_token_per_tpu=65536", + "--max_segments_per_packed_row=16", + "--pack_pad_position=start", + "--pack_pad_align=8192", + ]), + pad_id=7, + ) + self.assertEqual(config.pack_pad_position, "start") + self.assertEqual(config.pack_pad_align, 8192) + self.assertEqual(config.max_seq_token_per_tpu, 65536) + self.assertEqual(config.max_segments_per_packed_row, 16) + self.assertEqual(config.pad_id, 7) + with self.assertRaises(SystemExit): + parse(["--pack_pad_position=middle"]) + + def test_the_parsed_pad_position_packs_rows_with_the_padding_first(self): + # The equality checks above cannot tell 8192 from 8192.0, and only an int + # can index a row: pack two rollouts with the assembler that the program + # builds from the same BatchConfig. + parse = run_deepswe_dist._parse_args # pylint: disable=protected-access + build = run_deepswe_dist._batch_config # pylint: disable=protected-access + for argv, budget, first_token in ( + (["--max_seq_token_per_tpu=8192"], 8192, 0), + # PACK_PAD_POSITION=start with PACK_PAD_ALIGN unset. + ( + ["--max_seq_token_per_tpu=8192", "--pack_pad_position=start"], + 8192, + 4096, + ), + ( + [ + "--max_seq_token_per_tpu=16384", + "--pack_pad_position=start", + "--pack_pad_align=8192", + ], + 16384, + 8192, + ), + ): + with self.subTest(argv=argv): + assembler = batch_assembly.create_batch_assembler( + num_generations=2, + mini_batch_size=1, + train_micro_batch_size=1, + batch_config=build(parse(argv), pad_id=0), + ) + (batch,) = assembler.feed([_row(8), _row(8)]) + assembler.flush() + segment_ids = np.asarray(batch.payload.segment_ids) + # One row holds both 10-token rollouts, at [0, 10) and [64, 74) before + # it moves. + self.assertEqual(segment_ids.shape, (1, budget)) + self.assertEqual(np.count_nonzero(segment_ids), 20) + self.assertEqual(int(np.flatnonzero(segment_ids[0])[0]), first_token) + + def test_main_gives_the_program_the_batch_config_of_the_command_line(self): + # main() needs a cluster, so the tests above check _batch_config instead; + # this checks that main() hands the program what _batch_config returns. A + # BatchConfig written out in main() would drop --pack_pad_position with + # every test above still passing. + tree = ast.parse(inspect.getsource(run_deepswe_dist.main)) + programs = [ + node + for node in ast.walk(tree) + if isinstance(node, ast.Call) + and ast.unparse(node.func) == "rl_program.StandardRLProgram" + ] + self.assertLen(programs, 1) + batch_configs = [ + ast.unparse(keyword.value) + for keyword in programs[0].keywords + if keyword.arg == "batch_config" + ] + self.assertEqual(batch_configs, ["_batch_config(args, pad_id)"]) + class AuxMetricForwardingTest(absltest.TestCase): """Tests for converting a loss function's aux metrics for the buffer.""" diff --git a/tests/rl/packing_test.py b/tests/rl/packing_test.py index cf4f51e584..63699bcb7a 100644 --- a/tests/rl/packing_test.py +++ b/tests/rl/packing_test.py @@ -15,6 +15,7 @@ """Tests for the type-agnostic packing core logic.""" from absl.testing import absltest +from absl.testing import parameterized import numpy as np from tunix.rl import packing @@ -485,6 +486,412 @@ def test_pack_chunk_returns_packed_chunk_sharing_underlying_buffers(self): self.assertTrue(np.shares_memory(row0.routed_experts, chunk.routed_experts)) +def _random_bins(rng, *, budget, segment_align_multiple, pack_size): + """Random bins, filled as the assemblers fill them (`fill_one_chunk`). + + The items include empty ones and prompt-only or completion-only ones; all of + them carry the same random subset of the per-token fields, and most of them + carry routing (an item without it is packed as unset routing). + + Args: + rng: The random generator. + budget: Tokens per row. + segment_align_multiple: The packer's segment alignment. + pack_size: Rows in the chunk. + + Returns: + The bins and the per-token fields to pack. + """ + carried = tuple(f for f in packing.PER_TOKEN_FIELDS if rng.random() < 0.5) + routed = rng.random() < 0.7 + items = [] + for _ in range(int(rng.integers(1, 10))): + if rng.random() < 0.1: + p = c = 0 + else: + n = int(rng.integers(1, max(2, budget // int(rng.integers(1, 5))))) + p = int(rng.integers(0, n + 1)) + c = n - p + items.append( + packing.PackItem( + prompt_ids=rng.integers(5, 1000, p).astype(np.int32), + completion_ids=rng.integers(5, 1000, c).astype(np.int32), + completion_mask=(rng.random(c) < 0.8).astype(np.float32), + advantages=rng.standard_normal(c).astype(np.float32), + per_token={ + f: rng.standard_normal(c).astype(np.float32) for f in carried + }, + policy_version=np.array([int(rng.integers(0, 9))]), + routed_experts=( + rng.integers(0, 64, (p + c, 3, 2)).astype(np.int16) + if routed and rng.random() < 0.85 + else None + ), + ) + ) + bins, _ = packing.fill_one_chunk( + items, + pack_size=pack_size, + budget=budget, + max_segments=8, + segment_align_multiple=segment_align_multiple, + ) + return bins, carried + + +def _row_arrays(chunk): + """Every `[rows, budget, ...]` array of a PackedChunk, by name.""" + arrays = { + name: np.asarray(getattr(chunk, name)) + for name in ( + "ids", + "prompt_mask", + "completion_mask", + "advantages", + "segment_ids", + "segment_positions", + ) + } + for name, values in chunk.per_token.items(): + arrays[f"per_token/{name}"] = np.asarray(values) + if chunk.routed_experts is not None: + arrays["routed_experts"] = np.asarray(chunk.routed_experts) + return arrays + + +def _rotated_like_pad_first(end, *, pad_align): + """The "end" chunk's rows, each rotated right by its rounded trailing padding. + + Computed from the "end" layout alone, not by the packer: this is the + definition the "start" layout is tested against. + + Args: + end: A chunk packed with pad_position="end". + pad_align: The multiple the trailing padding is rounded down to. + + Returns: + The rotated arrays by name, and the shift of each row. + """ + arrays = _row_arrays(end) + budget = arrays["ids"].shape[1] + shifts = [] + for seg in arrays["segment_ids"]: + real = np.flatnonzero(seg > 0) + trailing = budget - 1 - int(real[-1]) if real.size else 0 + shifts.append(trailing // pad_align * pad_align) + rotated = { + name: np.stack([np.roll(a[r], s, axis=0) for r, s in enumerate(shifts)]) + for name, a in arrays.items() + } + return rotated, shifts + + +class PadPositionTest(parameterized.TestCase): + """pad_position="start": each packed row's padding goes in front of it.""" + + def _assert_same_chunk(self, got, want): + got_arrays, want_arrays = _row_arrays(got), _row_arrays(want) + self.assertEqual(sorted(got_arrays), sorted(want_arrays)) + for name, value in want_arrays.items(): + self.assertEqual(got_arrays[name].dtype, value.dtype, name) + np.testing.assert_array_equal(got_arrays[name], value, err_msg=name) + self.assertEqual(got.num_real_segments, want.num_real_segments) + self.assertEqual(len(got.policy_versions), len(want.policy_versions)) + for a, b in zip(got.policy_versions, want.policy_versions): + self.assertIs(a, b) + + def test_defaults(self): + self.assertEqual(packing.DEFAULT_PAD_POSITION, "end") + self.assertEqual(packing.DEFAULT_PAD_ALIGN, 4096) + self.assertEqual(packing.PAD_POSITIONS, ("end", "start")) + + def test_start_row_layout(self): + items = [ + _item([1, 2], [3, 4, 5], adv=1.5), + _item([6], [7, 8], adv=0.5), + ] + [[row]] = packing.pack_core( + items, + budget=16, + pack_size=1, + segment_align_multiple=1, + pad_position="start", + pad_align=4, + ) + # 8 tokens and 8 of padding: the row moves right by 8. + np.testing.assert_array_equal(row.ids, [0] * 8 + [1, 2, 3, 4, 5, 6, 7, 8]) + np.testing.assert_array_equal(row.segment_ids, [0] * 8 + [1] * 5 + [2] * 3) + np.testing.assert_array_equal( + row.segment_positions, [0] * 8 + [0, 1, 2, 3, 4, 0, 1, 2] + ) + np.testing.assert_array_equal( + row.prompt_mask, [0] * 8 + [1, 1, 0, 0, 0, 1, 0, 0] + ) + np.testing.assert_array_equal( + row.completion_mask, [0] * 8 + [0, 0, 1, 1, 1, 0, 1, 1] + ) + np.testing.assert_allclose( + row.advantages, [0] * 8 + [0, 0, 1.5, 1.5, 1.5, 0, 0.5, 0.5] + ) + self.assertEqual(row.num_real_segments, 2) + + def test_start_rounds_the_shift_down_to_pad_align(self): + # 7 tokens leave 9 of padding: pad_align 4 moves the row by 8, and one + # padding token stays at the end. + [[row]] = packing.pack_core( + [_item([1, 2, 3], [4, 5, 6, 7])], + budget=16, + pack_size=1, + segment_align_multiple=1, + pad_position="start", + pad_align=4, + ) + np.testing.assert_array_equal(row.ids, [0] * 8 + [1, 2, 3, 4, 5, 6, 7, 0]) + np.testing.assert_array_equal(row.segment_ids, [0] * 8 + [1] * 7 + [0]) + + def test_start_keeps_segment_starts_aligned(self): + # With the default 64-token alignment the "end" row holds [0, 30) and + # [64, 84) and 428 tokens of padding; pad_align 128 moves it by 384. + items = [ + _item(np.arange(1, 11), np.arange(1, 21)), + _item(np.arange(1, 6), np.arange(1, 16)), + ] + [[row]] = packing.pack_core( + items, budget=512, pack_size=1, pad_position="start", pad_align=128 + ) + np.testing.assert_array_equal( + np.flatnonzero(row.segment_ids == 1), np.arange(384, 414) + ) + np.testing.assert_array_equal( + np.flatnonzero(row.segment_ids == 2), np.arange(448, 468) + ) + + def test_start_moves_by_the_default_pad_align(self): + # The recipe's setting: "start" with pad_align left at its default. 4096 + # and 4096.0 compare equal, and only an int can index a row, so pack one. + [[row]] = packing.pack_core( + [_item(np.arange(1, 11), np.arange(1, 21))], + budget=8192, + pack_size=1, + pad_position="start", + ) + # 30 tokens leave 8162 of padding: the row moves by 4096. + np.testing.assert_array_equal( + np.flatnonzero(row.segment_ids), np.arange(4096, 4126) + ) + + def test_start_measures_the_padding_after_the_last_real_token(self): + # The shift is the padding after the row's last token, rounded down to + # pad_align; the end of the row is not rounded up to the segment alignment + # first. That only matters when the budget is not a multiple of + # segment_align_multiple: 100 tokens of a 300-token row leave 200 of + # padding, so the row moves by 192, where rounding its end up to 128 would + # leave 172 and move it by 128. The empty item goes last, which is where + # `fill_one_chunk` puts empty items. + chunk = packing.pack_chunk( + [[_item(np.arange(1, 41), np.arange(1, 61)), _item([], [])]], + budget=300, + pad_id=0, + carried=(), + pad_position="start", + pad_align=64, + ) + np.testing.assert_array_equal( + np.flatnonzero(chunk.segment_ids[0]), np.arange(192, 292) + ) + + def test_start_moves_routing_and_per_token_values_with_their_tokens(self): + routing = np.arange(3 * 2 * 2, dtype=np.int16).reshape(3, 2, 2) + item = packing.PackItem( + prompt_ids=np.array([1], dtype=np.int32), + completion_ids=np.array([2, 3], dtype=np.int32), + completion_mask=np.ones(2, dtype=np.float32), + advantages=np.ones(2, dtype=np.float32), + per_token={"returns": np.array([1.0, 2.0], dtype=np.float32)}, + routed_experts=routing, + ) + row = packing.pack_bin( + [item], + budget=8, + pad_id=0, + carried=("returns",), + segment_align_multiple=1, + pad_position="start", + pad_align=4, + ) + # 3 tokens and 5 of padding: the row moves right by 4. + np.testing.assert_array_equal(row.ids, [0, 0, 0, 0, 1, 2, 3, 0]) + np.testing.assert_allclose( + row.per_token["returns"], [0, 0, 0, 0, 0, 1, 2, 0] + ) + np.testing.assert_array_equal(row.routed_experts[4:7], routing) + np.testing.assert_array_equal( + row.routed_experts[:4], packing.UNSET_ROUTED_EXPERT + ) + np.testing.assert_array_equal( + row.routed_experts[7:], packing.UNSET_ROUTED_EXPERT + ) + + def test_start_is_the_end_row_rotated_by_its_rounded_padding(self): + rng = np.random.default_rng(0) + shifted_rows = 0 + for _ in range(200): + sam = int(rng.choice([1, 8, 64, 128])) + # 300 and 1000 are not multiples of every segment alignment. + budget = int(rng.choice([256, 300, 1000, 1024, 4096, 8192])) + pad_align = sam * int(rng.choice([1, 2, 4, 16])) + bins, carried = _random_bins( + rng, + budget=budget, + segment_align_multiple=sam, + pack_size=int(rng.integers(1, 5)), + ) + kwargs = dict( + budget=budget, pad_id=3, carried=carried, segment_align_multiple=sam + ) + end = packing.pack_chunk(bins, **kwargs) + start = packing.pack_chunk( + bins, pad_position="start", pad_align=pad_align, **kwargs + ) + want, shifts = _rotated_like_pad_first(end, pad_align=pad_align) + got = _row_arrays(start) + self.assertEqual(sorted(got), sorted(want)) + for name, value in want.items(): + self.assertEqual(got[name].dtype, value.dtype, name) + np.testing.assert_array_equal(got[name], value, err_msg=name) + self.assertEqual(start.num_real_segments, end.num_real_segments) + for a, b in zip(start.policy_versions, end.policy_versions): + self.assertIs(a, b) + for r, shift in enumerate(shifts): + # Only padding moved to the front: the row's first `shift` positions. + self.assertFalse(got["segment_ids"][r, :shift].any()) + np.testing.assert_array_equal(got["ids"][r, :shift], 3) + # Every segment still starts on the segment alignment. + seg_starts = np.flatnonzero( + (got["segment_ids"][r] > 0) & (got["segment_positions"][r] == 0) + ) + np.testing.assert_array_equal(seg_starts % sam, 0) + shifted_rows += sum(s > 0 for s in shifts) + # The comparison is not vacuous: many rows did move. + self.assertGreater(shifted_rows, 100) + + def test_start_keeps_each_token_with_its_values(self): + # Checked against the items themselves, not against the "end" layout. + rng = np.random.default_rng(1) + for _ in range(50): + bins, carried = _random_bins( + rng, budget=4096, segment_align_multiple=64, pack_size=3 + ) + chunk = packing.pack_chunk( + bins, + budget=4096, + pad_id=3, + carried=carried, + pad_position="start", + pad_align=512, + ) + for r, bin_items in enumerate(bins): + real_items = [item for item in bin_items if item.num_tokens] + for seg, item in enumerate(real_items, start=1): + p = item.prompt_ids.shape[0] + where = np.flatnonzero(chunk.segment_ids[r] == seg) + self.assertLen(where, item.num_tokens) + self.assertEqual(where[-1] - where[0] + 1, item.num_tokens) + np.testing.assert_array_equal( + chunk.ids[r, where], + np.concatenate([item.prompt_ids, item.completion_ids]), + ) + np.testing.assert_array_equal( + chunk.segment_positions[r, where], np.arange(item.num_tokens) + ) + np.testing.assert_array_equal(chunk.prompt_mask[r, where[:p]], 1.0) + np.testing.assert_array_equal(chunk.prompt_mask[r, where[p:]], 0.0) + np.testing.assert_array_equal( + chunk.completion_mask[r, where[p:]], item.completion_mask + ) + np.testing.assert_array_equal( + chunk.advantages[r, where[p:]], item.advantages + ) + for name in carried: + np.testing.assert_array_equal( + chunk.per_token[name][r, where[p:]], item.per_token[name] + ) + if item.routed_experts is not None: + np.testing.assert_array_equal( + chunk.routed_experts[r, where], item.routed_experts + ) + + def test_end_layout_is_the_same_whatever_pad_align(self): + # pad_align is not read with "end": no value of it changes or rejects an + # "end" row, also at segment alignments that do not divide 4096. + rng = np.random.default_rng(2) + for _ in range(100): + sam = int(rng.choice([1, 3, 48, 64, 100])) + budget = int(rng.choice([300, 1000, 4096])) + bins, carried = _random_bins( + rng, + budget=budget, + segment_align_multiple=sam, + pack_size=int(rng.integers(1, 4)), + ) + kwargs = dict( + budget=budget, pad_id=3, carried=carried, segment_align_multiple=sam + ) + default = packing.pack_chunk(bins, **kwargs) + for pad_align in (packing.DEFAULT_PAD_ALIGN, 0, -1, 100): + self._assert_same_chunk( + packing.pack_chunk( + bins, pad_position="end", pad_align=pad_align, **kwargs + ), + default, + ) + + def test_rows_that_cannot_move_are_unchanged(self): + # A full row, a row whose trailing padding is shorter than pad_align, an + # empty bin and a bin of empty items. + bins = [ + [_item(np.arange(1, 5), np.arange(1, 5))], + [_item([1], [2, 3])], + [], + [_item([], [])], + ] + kwargs = dict(budget=8, pad_id=0, carried=(), segment_align_multiple=1) + self._assert_same_chunk( + packing.pack_chunk(bins, pad_position="start", pad_align=8, **kwargs), + packing.pack_chunk(bins, **kwargs), + ) + + @parameterized.named_parameters( + ("unknown_position", "middle", 4096, 64), + ("capitalized_position", "Start", 4096, 64), + ("zero_align", "start", 0, 64), + ("negative_align", "start", -64, 64), + ("align_not_a_multiple", "start", 100, 64), + ("align_below_the_segment_alignment", "start", 32, 64), + ) + def test_invalid_settings_raise(self, pad_position, pad_align, sam): + items = [_item([1, 2], [3])] + kwargs = dict( + pad_position=pad_position, + pad_align=pad_align, + segment_align_multiple=sam, + ) + with self.assertRaisesRegex(ValueError, "pad_"): + packing.validate_pad_position(pad_position, pad_align, sam) + with self.assertRaisesRegex(ValueError, "pad_"): + packing.pack_core(items, budget=8192, **kwargs) + with self.assertRaisesRegex(ValueError, "pad_"): + packing.pack_core([], budget=8192, **kwargs) + with self.assertRaisesRegex(ValueError, "pad_"): + packing.pack_chunk([items], budget=8192, pad_id=0, carried=(), **kwargs) + with self.assertRaisesRegex(ValueError, "pad_"): + packing.pack_bin(items, budget=8192, pad_id=0, carried=(), **kwargs) + + def test_start_accepts_any_multiple_of_the_segment_alignment(self): + for pad_align, sam in ((64, 64), (4096, 64), (100, 1), (3, 3)): + packing.validate_pad_position("start", pad_align, sam) + + if __name__ == "__main__": absltest.main() diff --git a/tunix/experimental/examples/deepswe_dist/k8s_launcher.sh b/tunix/experimental/examples/deepswe_dist/k8s_launcher.sh index 5989d98417..e985217220 100755 --- a/tunix/experimental/examples/deepswe_dist/k8s_launcher.sh +++ b/tunix/experimental/examples/deepswe_dist/k8s_launcher.sh @@ -45,6 +45,11 @@ export MAX_TURNS=${MAX_TURNS:-20} export TRAIN_MICRO_BATCH_SIZE=${TRAIN_MICRO_BATCH_SIZE:-1} export MAX_SEQ_TOKEN_PER_TPU=${MAX_SEQ_TOKEN_PER_TPU:-} export MAX_SEGMENTS_PER_PACKED_ROW=${MAX_SEGMENTS_PER_PACKED_ROW:-} +# Padding position in each packed row: end (the orchestrator's default) or +# start (padding first), and the token multiple a start row moves by (4096). +# Empty leaves the orchestrator's default. +export PACK_PAD_POSITION=${PACK_PAD_POSITION:-} +export PACK_PAD_ALIGN=${PACK_PAD_ALIGN:-} export COMPUTE_LOGPS_CHUNK_SIZE=${COMPUTE_LOGPS_CHUNK_SIZE:-} # Set to tunix to run Tunix's PeftTrainer, and maxtext to run MaxText's MaxTextTrainingEngine @@ -589,6 +594,8 @@ start_orchestrator() { ${IMAGE_REWRITE_PREFIX:+--image_rewrite_prefix=${IMAGE_REWRITE_PREFIX}} \ ${MAX_SEQ_TOKEN_PER_TPU:+--max_seq_token_per_tpu=${MAX_SEQ_TOKEN_PER_TPU}} \ ${MAX_SEGMENTS_PER_PACKED_ROW:+--max_segments_per_packed_row=${MAX_SEGMENTS_PER_PACKED_ROW}} \ + ${PACK_PAD_POSITION:+--pack_pad_position=${PACK_PAD_POSITION}} \ + ${PACK_PAD_ALIGN:+--pack_pad_align=${PACK_PAD_ALIGN}} \ ${TRAINER_MESH_FSDP:+--trainer_fsdp=${TRAINER_MESH_FSDP}} \ $( [[ "${TRAINER_MESH_EXPERT:-1}" -gt 1 ]] && echo "--trainer_expert=${TRAINER_MESH_EXPERT}" ) \ ${RPC_TIMEOUT_S:+--rpc_timeout_s=${RPC_TIMEOUT_S}} \ diff --git a/tunix/experimental/examples/deepswe_dist/run_deepswe_dist.py b/tunix/experimental/examples/deepswe_dist/run_deepswe_dist.py index f975b8ea10..2033afac07 100644 --- a/tunix/experimental/examples/deepswe_dist/run_deepswe_dist.py +++ b/tunix/experimental/examples/deepswe_dist/run_deepswe_dist.py @@ -48,6 +48,7 @@ from tunix.experimental.weight_sync import weight_sync_coordinator from tunix.experimental.worker import remote_execution from tunix.rl import algorithm_config +from tunix.rl import packing from tunix.sft import metrics_logger as metrics_logger_lib from tunix.utils import mllog_utils @@ -211,6 +212,29 @@ def _parse_args(argv: list[str]) -> argparse.Namespace: default=None, help="Maximum segments per packed row when sequence packing is enabled.", ) + parser.add_argument( + "--pack_pad_position", + choices=packing.PAD_POSITIONS, + default=packing.DEFAULT_PAD_POSITION, + help=( + "Where each packed row's padding goes when sequence packing is" + " enabled: 'end' (default), or 'start', which moves the row right by" + " its trailing padding rounded down to a multiple of" + " --pack_pad_align so that the padding comes first. Under" + " load-balanced context parallelism this evens out the attention" + " work of the context ranks (see tunix/rl/packing.py)." + ), + ) + parser.add_argument( + "--pack_pad_align", + type=int, + default=packing.DEFAULT_PAD_ALIGN, + help=( + "With --pack_pad_position=start, rows move by a multiple of this" + " many tokens (default 4096, the largest Splash block of the" + " Qwen3.5-397B-A17B recipe)." + ), + ) # TODO(tunix-dev): Clean up worker specific configuration to orchestrator. parser.add_argument( "--trainer_fsdp", @@ -829,6 +853,24 @@ def _build_trajectory_store_config( return [file_cfg, sql_cfg] +def _batch_config( + args: argparse.Namespace, pad_id: int +) -> batch_assembly.BatchConfig: + """Returns the program's BatchConfig: the command line plus the pad id.""" + return batch_assembly.BatchConfig( + pad_id=pad_id, + max_prompt_length=args.max_prompt_length, + max_response_length=args.max_response_length, + max_seq_token_per_tpu=args.max_seq_token_per_tpu, + max_segments_per_packed_row=args.max_segments_per_packed_row, + trainer_fsdp=args.trainer_fsdp, + trainer_dp=args.trainer_dp, + trainer_expert=args.trainer_expert, + pack_pad_position=args.pack_pad_position, + pack_pad_align=args.pack_pad_align, + ) + + def _build_algo(args: argparse.Namespace) -> algorithm_adapter.GRPOAdapter: algo_config = algorithm_config.GRPOConfig( num_generations=args.num_generations, @@ -1310,16 +1352,7 @@ def _on_train_start(step: int) -> None: ), reward_fns=[], batch_size=args.batch_size, - batch_config=batch_assembly.BatchConfig( - pad_id=pad_id, - max_prompt_length=args.max_prompt_length, - max_response_length=args.max_response_length, - max_seq_token_per_tpu=args.max_seq_token_per_tpu, - max_segments_per_packed_row=args.max_segments_per_packed_row, - trainer_fsdp=args.trainer_fsdp, - trainer_dp=args.trainer_dp, - trainer_expert=args.trainer_expert, - ), + batch_config=_batch_config(args, pad_id), metrics_logging_options=metrics_logging_options, trajectory_log_dir=args.trajectory_log_dir, trajectory_store=cluster.trajectory_store, diff --git a/tunix/experimental/orchestrator/batch_assembly.py b/tunix/experimental/orchestrator/batch_assembly.py index d29a011bb9..c0df0e2f43 100644 --- a/tunix/experimental/orchestrator/batch_assembly.py +++ b/tunix/experimental/orchestrator/batch_assembly.py @@ -196,6 +196,13 @@ class BatchConfig: activation_batch / activation_batch_moe), so the packed row count must be divisible by the product of all four, not just fsdp * dp. A user with custom logical_axis_rules will need a different product. + pack_pad_position: Where each packed row's padding goes: "end" (default) + or "start", which moves the row right by its trailing padding rounded + down to a multiple of `pack_pad_align` tokens so that the padding comes + first (see `packing.PAD_POSITIONS`). Requires sequence packing. + pack_pad_align: The multiple of tokens a "start" row moves by (default + 4096); a positive multiple of `segment_align_multiple`, smaller than + the packed row. Only valid with pack_pad_position="start". """ pad_id: int = 0 @@ -208,6 +215,8 @@ class BatchConfig: trainer_dp: int | None = None trainer_expert: int | None = None trainer_fsdp_transpose: int | None = None + pack_pad_position: str = packing.DEFAULT_PAD_POSITION + pack_pad_align: int = packing.DEFAULT_PAD_ALIGN def _extract_trajectory_id(item: Any) -> str: @@ -638,6 +647,8 @@ def __init__( pad_id: int = 0, max_segments_per_packed_row: int | None = None, segment_align_multiple: int = packing.DEFAULT_SEGMENT_ALIGN_MULTIPLE, + pad_position: str = packing.DEFAULT_PAD_POSITION, + pad_align: int = packing.DEFAULT_PAD_ALIGN, start_batch_index: int = 0, ): """Initializes SequencePackedBatchAssembler. @@ -652,6 +663,12 @@ def __init__( that may be packed into a single row. segment_align_multiple: Token boundary alignment multiple for the start of each segment after the first in a packed row. + pad_position: "end" (padding after each row's segments) or "start" + (padding first), see `packing.PAD_POSITIONS`. + pad_align: With "start", each row moves right by its trailing padding + rounded down to a multiple of this; a positive multiple of + `segment_align_multiple` below `max_packed_len`. Must stay at its + default with "end", where it would do nothing. start_batch_index: Initial microbatch index offset for tracking IDs. """ if batch_size <= 0: @@ -677,6 +694,21 @@ def __init__( "segment_align_multiple must be positive, got" f" {segment_align_multiple}." ) + packing.validate_pad_position( + pad_position, pad_align, segment_align_multiple + ) + # Settings under which the option would do nothing are rejected rather + # than ignored, so a run never measures the default by accident. + if pad_position == "start" and pad_align >= max_packed_len: + raise ValueError( + f"pad_position='start' with pad_align={pad_align} >= max_packed_len" + f"={max_packed_len} never moves a row; lower pad_align." + ) + if pad_position != "start" and pad_align != packing.DEFAULT_PAD_ALIGN: + raise ValueError( + f"pad_align={pad_align} only applies with pad_position='start', got" + f" pad_position={pad_position!r}." + ) self.batch_size = batch_size self.max_packed_len = max_packed_len self.pad_id = pad_id @@ -684,6 +716,8 @@ def __init__( self.mini_batch_size = mini_batch_size self.max_segments_per_packed_row = max_segments_per_packed_row self.segment_align_multiple = segment_align_multiple + self.pad_position = pad_position + self.pad_align = pad_align self._batch_counter = start_batch_index # Each entry is a `(PackItem, trajectory_id, raw_payload)` converted once at @@ -736,6 +770,8 @@ def _emit_one_chunk( carried=carried, routed_shape=routed_shape, segment_align_multiple=self.segment_align_multiple, + pad_position=self.pad_position, + pad_align=self.pad_align, ) batch_tracking_id = f"{_BATCH_ID_PREFIX}_{self._batch_counter}" merged_lineage = _merge_batch_lineage( @@ -1335,10 +1371,13 @@ def create_batch_assembler( logging.info( "Using SequencePackedBatchAssembler with max_seq_token_per_tpu: %d, " - "max_segments_per_packed_row: %s, pack_size: %d", + "max_segments_per_packed_row: %s, pack_size: %d, " + "pack_pad_position: %s, pack_pad_align: %d", batch_config.max_seq_token_per_tpu, batch_config.max_segments_per_packed_row, pack_size, + batch_config.pack_pad_position, + batch_config.pack_pad_align, ) return SequencePackedBatchAssembler( batch_size=pack_size, @@ -1348,6 +1387,8 @@ def create_batch_assembler( pad_id=batch_config.pad_id, max_segments_per_packed_row=batch_config.max_segments_per_packed_row, segment_align_multiple=batch_config.segment_align_multiple, + pad_position=batch_config.pack_pad_position, + pad_align=batch_config.pack_pad_align, ) if batch_config.max_prompt_length is not None: @@ -1356,6 +1397,15 @@ def create_batch_assembler( "max_response_length must be specified in batch_config when" " max_prompt_length is set for PaddedBatchAssembler." ) + if ( + batch_config.pack_pad_position != packing.DEFAULT_PAD_POSITION + or batch_config.pack_pad_align != packing.DEFAULT_PAD_ALIGN + ): + raise ValueError( + "pack_pad_position and pack_pad_align place the padding of packed" + " rows, and PaddedBatchAssembler does not pack: set" + " max_seq_token_per_tpu to pack, or leave both at their defaults." + ) return PaddedBatchAssembler( batch_size=train_micro_batch_size, max_prompt_length=batch_config.max_prompt_length, @@ -1372,4 +1422,6 @@ def create_batch_assembler( pad_id=batch_config.pad_id, max_segments_per_packed_row=batch_config.max_segments_per_packed_row, segment_align_multiple=batch_config.segment_align_multiple, + pad_position=batch_config.pack_pad_position, + pad_align=batch_config.pack_pad_align, ) diff --git a/tunix/rl/packing.py b/tunix/rl/packing.py index 67babc136a..e6cd8bbc8e 100644 --- a/tunix/rl/packing.py +++ b/tunix/rl/packing.py @@ -39,6 +39,67 @@ # `sequence_length_round: 64` and Qwen3.5's `gdn_chunk_size = 64`). DEFAULT_SEGMENT_ALIGN_MULTIPLE = 64 +# Where a packed row's padding goes (`pad_position` of `pack_chunk`). +# +# "end" writes the row's segments from position 0 and leaves the padding at the +# end of the row. "start" lays the row out exactly as "end" does, then moves the +# whole row right by its trailing padding rounded down to a multiple of +# `pad_align`: the padding comes first, and fewer than `pad_align` padding +# tokens stay at the end. +# +# Why: under load-balanced context parallelism (MaxText's DUAL_CHUNK_SWAP: +# context rank r holds sequence chunks r and 2 * cp - 1 - r) the last chunk of +# a row is rank 0's. With the padding at the end, only rank 0 gets lighter; +# ranks 1..cp-1 still attend every real query to every earlier position. An +# attention kernel that skips the tiles whose query and key segment-id ranges +# do not overlap skips the padding's key tiles for every real query once the +# padding comes first, so the work is spread more evenly over the ranks and +# the busiest rank, which sets the step time, does less. tokamax's Splash +# skips such tiles whenever it is given segment ids, from tokamax +# 0.0.16.dev20261008 on; a kernel that only masks segment ids element by +# element computes every tile, and there "start" gains nothing. +# +# Only positions in the row move. Every token keeps its id, its segment id, +# its position inside its segment and its masks, advantages, per-token values +# and routed experts, so a model whose attention, convolutions and recurrences +# all respect segment ids (as well as causal order) computes the same values; +# the kernel's tile boundaries move, so its reductions may round differently. +PAD_POSITIONS: tuple[str, ...] = ("end", "start") +DEFAULT_PAD_POSITION = "end" +# With "start", a row moves by a multiple of this many tokens, which keeps +# every segment start where it sits relative to the attention kernel's tiles: +# 4096 is the largest Splash block (`sa_block_kv`) of the Qwen3.5-397B-A17B +# recipe. It must be a multiple of `segment_align_multiple`, so the segment +# alignment is kept too. +DEFAULT_PAD_ALIGN = 4096 + + +def validate_pad_position( + pad_position: str, pad_align: int, segment_align_multiple: int +) -> None: + """Raises ValueError unless `pad_position` and `pad_align` can pack a row. + + `pad_align` is only read with pad_position="start", so it is only checked + then: an "end" row is the same whatever `pad_align` is. + + Args: + pad_position: One of `PAD_POSITIONS`. + pad_align: With "start", a positive multiple of `segment_align_multiple`. + segment_align_multiple: The packer's (positive) segment alignment. + """ + if pad_position not in PAD_POSITIONS: + raise ValueError( + f"pad_position must be one of {PAD_POSITIONS}, got {pad_position!r}." + ) + if pad_position == "start" and ( + pad_align <= 0 or pad_align % segment_align_multiple + ): + raise ValueError( + "pad_align must be a positive multiple of segment_align_multiple" + f" ({segment_align_multiple}) with pad_position='start', got" + f" {pad_align}." + ) + def align_offset(offset: int, segment_align_multiple: int) -> int: """Rounds `offset` up to the next multiple of `segment_align_multiple`.""" @@ -318,16 +379,50 @@ def pack_bin( carried: Sequence[str], routed_shape: tuple[int, ...] | None = None, segment_align_multiple: int = DEFAULT_SEGMENT_ALIGN_MULTIPLE, + pad_position: str = DEFAULT_PAD_POSITION, + pad_align: int = DEFAULT_PAD_ALIGN, ) -> PackedRow: """Packs a single bin of items into a single `[budget]` PackedRow.""" - return pack_chunk( + chunk = pack_chunk( [bin_items], budget=budget, pad_id=pad_id, carried=carried, routed_shape=routed_shape, segment_align_multiple=segment_align_multiple, - )[0] # pytype: disable=bad-return-type + pad_position=pad_position, + pad_align=pad_align, + ) + return chunk[0] # pytype: disable=bad-return-type + + +def _pad_first_shift( + bin_items: Sequence[PackItem], + *, + budget: int, + pad_align: int, + segment_align_multiple: int, +) -> int: + """Returns how far right pad_position="start" moves a bin's row. + + The row's tokens, written from position 0 as `pack_chunk` writes them with + "end", stop at `used`; the shift is the trailing padding `budget - used` + rounded down to a multiple of `pad_align`. A row without tokens, or one that + exceeds the budget (which `pack_chunk` rejects), is not moved. + """ + used = 0 + placed = 0 + for item in bin_items: + n = item.num_tokens + if n == 0: + continue + if placed: + used = align_offset(used, segment_align_multiple) + used += n + placed += 1 + if not placed or used > budget: + return 0 + return (budget - used) // pad_align * pad_align def pack_chunk( @@ -338,13 +433,30 @@ def pack_chunk( carried: Sequence[str], routed_shape: tuple[int, ...] | None = None, segment_align_multiple: int = DEFAULT_SEGMENT_ALIGN_MULTIPLE, + pad_position: str = DEFAULT_PAD_POSITION, + pad_align: int = DEFAULT_PAD_ALIGN, ) -> PackedChunk: - """Packs a sequence of bins of one chunk into a contiguous `PackedChunk`.""" + """Packs a sequence of bins of one chunk into a contiguous `PackedChunk`. + + Args: + bins: The items of each row. + budget: Tokens per row. + pad_id: Token id of the padding. + carried: The per-token fields to pack. + routed_shape: Trailing `(num_layers, top_k)` routing shape, or None to + derive it from the items. + segment_align_multiple: Start alignment of every segment after the first. + pad_position: "end" puts each row's padding after its segments; "start" + moves each row right by its trailing padding rounded down to a multiple + of `pad_align`, so the padding comes first (see `PAD_POSITIONS`). + pad_align: The multiple a "start" row moves by. Not read with "end". + """ if segment_align_multiple <= 0: raise ValueError( "segment_align_multiple must be positive, got" f" {segment_align_multiple}." ) + validate_pad_position(pad_position, pad_align, segment_align_multiple) if routed_shape is None: placed = [item for bin_items in bins for item in bin_items] routed_shape = routed_experts_shape(placed) @@ -371,6 +483,18 @@ def pack_chunk( for b, bin_items in enumerate(bins): if not bin_items: continue + # "start" writes every token of the row `shift` positions to the right of + # where "end" writes it; segment ids, in-segment positions and routing move + # with their tokens. Segment starts stay aligned: `cursor` is aligned as + # before and `shift` is a multiple of segment_align_multiple. + shift = 0 + if pad_position == "start": + shift = _pad_first_shift( + bin_items, + budget=budget, + pad_align=pad_align, + segment_align_multiple=segment_align_multiple, + ) cursor = 0 real_seg = 0 for item in bin_items: @@ -387,12 +511,13 @@ def pack_chunk( raise ValueError( f"pack_bin: bin size {cursor + n} exceeds budget {budget}." ) - seq = slice(cursor, cursor + n) - comp = slice(cursor + p, cursor + n) + offset = shift + cursor + seq = slice(offset, offset + n) + comp = slice(offset + p, offset + n) - ids[b, cursor : cursor + p] = item.prompt_ids + ids[b, offset : offset + p] = item.prompt_ids ids[b, comp] = item.completion_ids - prompt_mask[b, cursor : cursor + p] = 1.0 + prompt_mask[b, offset : offset + p] = 1.0 segment_ids[b, seq] = seg segment_positions[b, seq] = np.arange(n, dtype=np.int32) @@ -452,8 +577,13 @@ def pack_core( max_segments_per_packed_row: int | None = None, pad_id: int = 0, segment_align_multiple: int = DEFAULT_SEGMENT_ALIGN_MULTIPLE, + pad_position: str = DEFAULT_PAD_POSITION, + pad_align: int = DEFAULT_PAD_ALIGN, ) -> list[PackedChunk]: - """Packs `items` into a sequence of chunks, each containing `pack_size` PackedRows with `budget` tokens.""" + """Packs `items` into a sequence of chunks, each containing `pack_size` PackedRows with `budget` tokens. + + `pad_position` and `pad_align` place each row's padding, as in `pack_chunk`. + """ if budget <= 0: raise ValueError(f"Budget must be positive, got {budget}.") if pack_size <= 0: @@ -471,6 +601,7 @@ def pack_core( "segment_align_multiple must be positive, got" f" {segment_align_multiple}." ) + validate_pad_position(pad_position, pad_align, segment_align_multiple) if not items: return [] @@ -499,6 +630,8 @@ def pack_core( carried=carried, routed_shape=routed_experts_shape(placed), segment_align_multiple=segment_align_multiple, + pad_position=pad_position, + pad_align=pad_align, ) ) return chunks From 86e17b5d5625836d69d8884d37c231dd38fbe6df Mon Sep 17 00:00:00 2001 From: Param Bole Date: Thu, 15 Oct 2026 07:25:51 -0700 Subject: [PATCH 2/4] rl: fused scoring in the sharded log-prob head (TUNIX_LOGPS_SHARDED_HEAD=fused) With TUNIX_LOGPS_SHARDED_HEAD=local the sharded head scores each chunk with selective_log_softmax (and compute_entropy_from_logits), and its backward differentiates that with jax.vjp. At the 248,320-entry vocabulary of Qwen3.5-397B-A17B two parts of that cost more than the arithmetic. On TPU, the target logit's take_along_axis makes XLA write an f32 copy of the [chunk, vocab] logits for a SparseCore gather, and its transpose is a scatter-add: an f32 [chunk, vocab] cotangent, a flattened copy of it, a sort of the indices, a SparseCore scatter and a copy back to bf16, per chunk. And the logsumexp's max and exp-sum are computed again for the entropy in the forward, with one more exp pass, and once more in the backward. TUNIX_LOGPS_SHARDED_HEAD=fused runs the same sharded head with fused scoring: - The forward makes one pass over each chunk's logits (_head_logps_stats): the target logit by a one-hot select-sum, the max and exp-sum once, the entropy from the same exponentials, and the max and exp-sum kept for the backward (8 bytes per token). - When only the log-probs have a cotangent, as in GRPO, whose loss stops the entropy's gradient, the backward forms the logits' cotangent in one elementwise pass over the recomputed logits (_head_logps_cotangent): no f32 [chunk, vocab] intermediate, no scatter and no reductions. An optimization barrier makes both gradient matmuls read that one cotangent. - A loss that differentiates the entropy takes the "local" backward. Unset, 0 and local trace the same program as before. The option is a value of the switch that turns the head on, because that switch is the head's only setting: the head has no config field or argument, and k8s_launcher.sh already passes TUNIX_LOGPS_SHARDED_HEAD to the trainer as it is set. Numerics: the log-probs and the logits' cotangent are the f32 operations of "local" in the same order, but the exp-sum is reduced in another fusion, so results can differ in their last bits. On CPU (8 devices; a head with a 248,320-entry vocabulary called directly, and the test head through compute_per_token_logps; 12 cases), "fused" against "local": log-probs equal or 1 ulp apart, entropy 3 to 7 ulps apart, and gradients apart by at most 2.6e-8 of their largest entry. Measured on Qwen3.5-397B-A17B, 128 TPU v7x, on the RL trainer benchmark (AI-Hypercomputer/maxtext#5701 driver, staged step boundary, ragged buffer factor 2.75, log-prob chunks of 4,096, 9 micro-batches per step), against a control run in the same pod (o275x-c9hv-l3 against o275x-c9ctl-l3): 83.253 -> 82.708 s per optimizer step, mean of steps 1-4, and 9.084 -> 9.035 s per regular micro-batch. The trainer logged the head engaged with 16,384 local tokens per device in 4 steps of 4,096. Numerics on TPU, measured in the same benchmark together with two MaxText changes (deferred small all-reduces, and the shared expert applied first; nxc9pfxall against two reference runs whose gradient norms differ by 1.9e-5 relative): all six gates pass, loss 2.7e-7 relative, log-prob difference 3.1e-7 (the references' own: 3.6e-7), entropy 5.1e-7 relative, gradient norm 1.6e-4 relative (pass limit 2e-3). A deviceless TPU compile of the head alone, 16,384 tokens per device: temporaries 11.87 -> 4.35 GiB at chunks of 4,096 and 4.65 -> 2.75 GiB at chunks of 1,024, transcendentals 5.09e9 -> 2.03e9 and 1.27e9 -> 0.51e9, and no gather, sort or scatter left. The build measured on TPU switched this scoring with its own environment variable (PARAMBOLE_HEAD_VJP=1) in a copy of tunix/rl/common.py based on an earlier version of the head (md5 01eaca3f). On CPU, this file with "fused" lowers to the same StableHLO and gives bit-identical log-probs, entropy and gradients as that file with the variable set, and with "local" or unset the same as the previous commit, in all 12 cases (two meshes, f32 and bf16 logits, temperatures 1 and 0.7, padded chunks, the entropy off, returned or differentiated, linear and GRPO-shaped losses, forward only). Tests (tests/rl/sharded_logps_head_test.py): - ShardedLogpsHeadTest: "fused" against "local" through compute_per_token_logps with the log-probs alone differentiated (as in GRPO), with the entropy differentiated, with chunks that do not divide the local sequence and one chunk for all of it, at temperatures 0.7 and 1, with non-finite hidden states at masked positions, in bf16, with one output, forward only, and against the unsharded head; the program has no gather, no scatter and one optimization barrier when only the log-probs are differentiated, and the autodiff backward's scatter when the entropy is; "fused" parses and near misses raise. - FusedScoringTest (one device): _target_hit by hand, and _head_logps_stats and _head_logps_cotangent against selective_log_softmax, compute_entropy_from_logits and jax.vjp on logits with inf, -inf and NaN entries and a row of no finite logit, with negative ids and ids outside the vocabulary. Fifteen deliberately broken copies of tunix/rl/common.py each fail at least one of these tests, and each new test fails on at least one of them: the cotangent's sign flipped, the backward's temperature division dropped, the log-prob cotangent read from the first chunk, the max and exp-sum swapped, the max dropped from the log-probs, the entropy's sign flipped, an entropy cotangent sent to the one-pass backward, the barrier removed, the wrap of negative ids removed, the vocabulary bound off by one, the isfinite mask dropped from the cotangent, and the option ignored, always on, parsed as "local" or matched loosely. --- tests/rl/sharded_logps_head_test.py | 243 ++++++++++++++++++++++++++-- tunix/rl/common.py | 239 ++++++++++++++++++++++++++- 2 files changed, 463 insertions(+), 19 deletions(-) diff --git a/tests/rl/sharded_logps_head_test.py b/tests/rl/sharded_logps_head_test.py index 66b99242cf..03eda4750a 100644 --- a/tests/rl/sharded_logps_head_test.py +++ b/tests/rl/sharded_logps_head_test.py @@ -14,13 +14,14 @@ """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. +ShardedLogpsHeadTest tests the head, in its "local" and "fused" modes, on a +mesh of 8 CPU devices; FusedScoringTest tests the "fused" mode's chunk scoring +on one device. 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 @@ -189,6 +190,7 @@ def _run( legacy_mesh=False, differentiate=('logps', 'entropy'), nonfinite=(), + lower=False, ): """Returns (log-probs, entropy, grads) of a loss weighting both outputs. @@ -197,6 +199,8 @@ def _run( enter it under `stop_gradient` (GRPO differentiates the log-probs only). nonfinite: (row, position) pairs whose hidden state is NaN. The scores computed from them get loss weight 0, as padding does. + lower: Return the StableHLO text of the jitted value_and_grad instead of + running it. """ model = _HeadOnly(dtype, batch, seq) if nonfinite: @@ -238,9 +242,10 @@ def loss(state): return value, (logps, entropy) with self._head_switch(mode, rules, expect, legacy_mesh): - (_, (logps, entropy)), grads = jax.jit( - jax.value_and_grad(loss, has_aux=True) - )(state) + step = jax.jit(jax.value_and_grad(loss, has_aux=True)) + if lower: + return step.lower(state).as_text() + (_, (logps, entropy)), grads = step(state) return logps, entropy, grads def _assert_parity(self, got, want, entropy=True): @@ -461,12 +466,14 @@ def test_unknown_switch_value_raises(self): ('local', 'local'), (' Local ', 'local'), ('TRUE', 'local'), + ('fused', 'fused'), + (' Fused ', 'fused'), ) 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', '2') + @parameterized.parameters('locla', '2', 'fuse', 'local,fused') def test_mode_parsing_rejects_unknown(self, value): with mock.patch.dict(os.environ, {common.SHARDED_HEAD_ENV: value}): with self.assertRaisesRegex(ValueError, common.SHARDED_HEAD_ENV): @@ -496,6 +503,220 @@ def test_token_layout(self, rules, batch, seq, want): got = common._token_partition(self.mesh, batch, seq) # pylint: disable=protected-access self.assertEqual(got, want) + def _assert_fused_parity(self, got, want, entropy=True): + """Asserts that "fused" matches "local": log-probs, entropy and grads. + + The log-probs and gradients are the same f32 operations, but the exp-sum is + reduced in another fusion, so they may differ in their last bits; the + entropy comes from another formula. + + Args: + got: (log-probs, entropy, grads) with "fused". + want: The same with "local". + entropy: Whether to compare the entropy. + """ + np.testing.assert_allclose(got[0], want[0], rtol=1e-6, atol=2e-6) + if entropy: + np.testing.assert_allclose(got[1], want[1], rtol=1e-5, atol=1e-6) + for name in ('h', 'scale', 'table'): + g = np.asarray(got[2][name][...], np.float32) + w = np.asarray(want[2][name][...], np.float32) + np.testing.assert_allclose( + g, w, rtol=1e-5, atol=1e-6 * np.abs(w).max(), err_msg=name + ) + + @parameterized.named_parameters( + ('logps_only_as_in_grpo', ('logps',), 8, 0.7), + ('logps_only_uneven_chunk', ('logps',), 12, 1.0), + ('logps_only_one_step', ('logps',), 64, 0.7), + ('logps_and_entropy', ('logps', 'entropy'), 8, 0.7), + ('entropy_only', ('entropy',), 12, 1.0), + ) + def test_fused_matches_local(self, differentiate, chunk_size, temperature): + # With only the log-probs differentiated the backward forms the logits' + # cotangent in one pass; a differentiated entropy takes the autodiff + # backward. Chunk 12 does not divide the 32 local positions (padding). + kwargs = dict( + differentiate=differentiate, + chunk_size=chunk_size, + temperature=temperature, + ) + want = self._run('local', **kwargs) + got = self._run('fused', **kwargs) + self._assert_fused_parity(got, want) + + def test_fused_matches_scan_head(self): + want = self._run(None, differentiate=('logps',)) + got = self._run('fused', differentiate=('logps',)) + self._assert_parity(got, want) + + def test_fused_logps_only_output(self): + # return_entropy=False: one output, and _head_logps_stats skips the entropy. + want = self._run('local', return_entropy=False, differentiate=('logps',)) + got = self._run('fused', return_entropy=False, differentiate=('logps',)) + self._assert_fused_parity(got, want, entropy=False) + + def test_fused_nonfinite_hidden_states_at_masked_positions(self): + bad = ((0, 5), (1, 31), (1, 32), (2, 40), (3, 63), (3, 0)) + want = self._run('local', nonfinite=bad, differentiate=('logps',)) + got = self._run('fused', nonfinite=bad, differentiate=('logps',)) + for name in ('h', 'scale', 'table'): + self.assertTrue(np.all(np.isfinite(got[2][name][...])), name) + self.assertTrue(np.all(np.isfinite(got[0]))) + self._assert_fused_parity(got, want) + + def test_fused_bf16_matches_local(self): + want = self._run('local', dtype=jnp.bfloat16, differentiate=('logps',)) + got = self._run('fused', dtype=jnp.bfloat16, differentiate=('logps',)) + self._assert_fused_parity(got, want) + + @parameterized.parameters(True, False) + def test_fused_forward_only_matches_local(self, return_entropy): + # No gradient: the custom_vjp's primal body runs, scoring with + # _head_logps_stats and keeping no max or exp-sum. + model = _HeadOnly(jnp.float32) + graphdef, state = nnx.split(model) + tokens, segment_ids, positions, _ = self._inputs(_B, _T) + + def logps(mode): + 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=0.7, + chunk_size=8, + ) + )(state) + + want, got = logps('local'), logps('fused') + self.assertEqual(jax.tree.structure(got), jax.tree.structure(want)) + got_logps, want_logps = jax.tree.leaves(got)[0], jax.tree.leaves(want)[0] + np.testing.assert_allclose(got_logps, want_logps, rtol=1e-6, atol=2e-6) + if return_entropy: + np.testing.assert_allclose(got[1], want[1], rtol=1e-5, atol=1e-6) + + @parameterized.named_parameters( + ('fused_logps_only_as_in_grpo', 'fused', ('logps',), False), + ('fused_entropy_differentiated', 'fused', ('logps', 'entropy'), True), + ('local', 'local', ('logps',), True), + ) + def test_scatter_only_in_autodiff_backward( + self, mode, differentiate, scatter + ): + # "fused" reads the target logit with a one-hot select-sum, and when only + # the log-probs are differentiated its backward forms the logits' cotangent + # in one pass behind one optimization barrier: the program gathers nothing + # and scatters nothing. A differentiated entropy, like "local", takes the + # autodiff backward, which transposes take_along_axis into a scatter. + text = self._run(mode, differentiate=differentiate, lower=True) + count = { + op: len(re.findall(rf'\bstablehlo\.{op}\b', text)) + for op in ('gather', 'scatter', 'optimization_barrier') + } + if scatter: + self.assertGreater(count['scatter'], 0, count) + self.assertEqual(count['optimization_barrier'], 0, count) + else: + self.assertEqual( + count, {'gather': 0, 'scatter': 0, 'optimization_barrier': 1} + ) + if mode == 'local': + self.assertGreater(count['gather'], 0, count) + + +class FusedScoringTest(parameterized.TestCase): + """The "fused" head's chunk scoring against what it replaces. + + `_head_logps_stats` and `_head_logps_cotangent` against + `selective_log_softmax`, `compute_entropy_from_logits` and `jax.vjp` of + `selective_log_softmax`, on logits with non-finite entries and on target ids + outside the vocabulary. One device is enough. + """ + + _V = 50 + + def _inputs(self): + """[2, 9, 50] logits, [2, 9] target ids and [2, 9] log-prob cotangents.""" + rng = np.random.default_rng(1) + logits = 4.0 * rng.normal(size=(2, 9, self._V)).astype(np.float32) + logits[0, 0, 7] = np.inf + logits[0, 1, 8] = -np.inf + logits[0, 2, 9] = np.nan + logits[1, 4, :] = -np.inf # no finite logit in the row + ids = rng.integers(0, self._V, (2, 9)).astype(np.int32) + # Targets on the non-finite entries; ids counting from the end; and ids + # outside [-V, V), which take_along_axis reads as NaN. + ids[0] = [7, 8, 9, self._V, -1, -self._V, -self._V - 1, 2**30, 0] + g = rng.normal(size=(2, 9)).astype(np.float32) + return jnp.asarray(logits), jnp.asarray(ids), jnp.asarray(g) + + def test_target_hit(self): + ids = jnp.array([0, 4, 5, -1, -5, -6, 2**30], jnp.int32) + hit, valid = common._target_hit(ids, (7, 5)) # pylint: disable=protected-access + want = np.zeros((7, 5), bool) + for row, v in ((0, 0), (1, 4), (3, 4), (4, 0)): + want[row, v] = True + np.testing.assert_array_equal(hit, want) + np.testing.assert_array_equal( + valid, [True, True, False, True, True, False, False] + ) + + @parameterized.parameters(True, False) + def test_stats_match_selective_log_softmax(self, with_entropy): + logits, ids, _ = self._inputs() + stats = jax.jit(common._head_logps_stats, static_argnums=2) # pylint: disable=protected-access + logps, entropy, m, s = stats(logits, ids, with_entropy) + want = common.selective_log_softmax(logits, ids) + # NaN exactly where take_along_axis reads outside the vocabulary. + np.testing.assert_array_equal(np.isnan(logps), np.isnan(want)) + self.assertEqual(int(np.isnan(want).sum()), 3) + np.testing.assert_allclose(logps, want, rtol=1e-6, atol=1e-6) + x = np.where(np.isfinite(logits), logits, 0.0) + np.testing.assert_array_equal(m, x.max(-1)) + shifted = (x - x.max(-1, keepdims=True)).astype(np.float64) + np.testing.assert_allclose(s, np.exp(shifted).sum(-1), rtol=1e-5) + if with_entropy: + np.testing.assert_allclose( + entropy, + common.compute_entropy_from_logits(logits), + rtol=1e-5, + atol=1e-6, + ) + else: + self.assertIsNone(entropy) + + def test_cotangent_matches_autodiff(self): + logits, ids, g = self._inputs() + _, _, m, s = common._head_logps_stats(logits, ids, False) # pylint: disable=protected-access + got = jax.jit(common._head_logps_cotangent)(logits, ids, g, m, s) # pylint: disable=protected-access + _, vjp = jax.vjp(lambda x: common.selective_log_softmax(x, ids), logits) + (want,) = vjp(g) + np.testing.assert_allclose(got, want, rtol=1e-6, atol=1e-6) + got = np.asarray(got) + # A non-finite logit gets no cotangent (the transpose of the isfinite mask). + self.assertTrue(np.all(got[~np.isfinite(logits)] == 0.0)) + # Over a row of finite logits the cotangent sums to g - g * sum(p) = 0, or + # to -g when the id is outside the vocabulary, whose g the gather's + # transpose drops. + outside = np.zeros(g.shape, bool) + outside[0, [3, 6, 7]] = True # ids V, -V - 1 and 2**30 + finite = np.all(np.isfinite(logits), axis=-1) + np.testing.assert_allclose( + got.sum(-1)[outside], -np.asarray(g)[outside], rtol=1e-5 + ) + np.testing.assert_allclose( + got.sum(-1)[finite & ~outside], 0.0, rtol=0, atol=1e-5 + ) + class ChildProcessTest(absltest.TestCase): """Runs ShardedLogpsHeadTest in a child process when this one cannot.""" diff --git a/tunix/rl/common.py b/tunix/rl/common.py index 2c269136e3..fd80be744d 100644 --- a/tunix/rl/common.py +++ b/tunix/rl/common.py @@ -1040,10 +1040,11 @@ def logp_step(carry, xs): def sharded_head_mode() -> str | None: - """Returns None (off) or "local", from TUNIX_LOGPS_SHARDED_HEAD. + """Returns None (off), "local" or "fused", from TUNIX_LOGPS_SHARDED_HEAD. Case and surrounding whitespace are ignored. Off: unset, empty, 0, false, off, - no. On: local, 1, true, on, yes. + no. On: local, 1, true, on, yes. On, with each chunk scored by the fused + log-softmax described above `_target_hit`: fused. Raises: ValueError: The variable has any other value. @@ -1053,12 +1054,148 @@ def sharded_head_mode() -> str | None: return None if value in ("1", "true", "on", "yes", "local"): return "local" + if value == "fused": + return "fused" raise ValueError( f"{SHARDED_HEAD_ENV}={value!r}: expected local, 1, true, on or yes (on)," - " or empty, 0, false, off or no (off)." + " fused (on, with fused scoring), or empty, 0, false, off or no (off)." ) +# Fused scoring for the sharded head: TUNIX_LOGPS_SHARDED_HEAD=fused. +# +# With "local" the head scores each chunk with `selective_log_softmax` (and +# `compute_entropy_from_logits`), and its backward differentiates that with +# `jax.vjp`. At a 248,320-entry vocabulary two parts of it cost more than the +# arithmetic: +# * The target logit is a `take_along_axis` gather. On TPU, XLA writes an f32 +# copy of the [chunk, vocab] logits for SparseCore to gather the targets +# from, and the transpose of the gather is a scatter-add: an f32 +# [chunk, vocab] cotangent, a flattened copy of it, a sort of the indices, a +# SparseCore scatter and a copy back to bf16, per chunk. +# * The backward recomputes the logsumexp's max and exp-sum, and the entropy +# recomputes the same reductions, plus one more exp pass, in the forward. +# "fused" scores a chunk in one pass over its logits (`_head_logps_stats`): the +# target logit by a one-hot select-sum, the max and exp-sum once, the entropy +# from the same exponentials, and the max and exp-sum kept for the backward +# (8 bytes per token). When only the log-probs have a cotangent, as in GRPO, +# whose loss stops the entropy's gradient, the backward forms the logits' +# cotangent in one elementwise pass over the recomputed logits +# (`_head_logps_cotangent`): no f32 [chunk, vocab] intermediate, no scatter and +# no reductions. An entropy cotangent takes the "local" backward. +# The log-probs and the logits' cotangent are the f32 operations of the "local" +# path in the same order. The exp-sum, though, is reduced in a different fusion +# (with the entropy's sum, when the entropy is returned), and XLA may order +# that reduction differently, so the log-probs and the gradients can differ +# from "local" in their last bits. The entropy is computed by another formula +# and differs in its last bits. + + +def _target_hit(ids: jax.Array, shape) -> tuple[jax.Array, jax.Array]: + """One-hot of `ids` over the last axis of `shape`, as take_along_axis reads it. + + Args: + ids: [...] target ids. + shape: Shape [..., V] of the logits the ids index. + + Returns: + (hit, valid): `hit[..., v]` is True where v is the target, a negative id + counting from the end as in take_along_axis; `valid` is False for an id + outside [-V, V), which take_along_axis's default "fill" mode reads as NaN + and whose scatter, in its transpose, drops the cotangent. + """ + vocab = shape[-1] + ids = jnp.where(ids < 0, ids + vocab, ids) # take_along_axis wraps negatives + hit = ( + jax.lax.broadcasted_iota(ids.dtype, shape, len(shape) - 1) + == ids[..., None] + ) + return hit, (ids >= 0) & (ids < vocab) + + +def _head_logps_stats(logits: jax.Array, ids: jax.Array, with_entropy: bool): + """`selective_log_softmax` (and the entropy) of a chunk, keeping max and sum. + + Args: + logits: [..., V] logits, as `selective_log_softmax` receives them. + ids: [...] target ids. + with_entropy: Also compute the entropy of `compute_entropy_from_logits`. + + Returns: + (logps, entropy, m, s): `selective_log_softmax(logits, ids)` by the same f32 + operations; the entropy of `compute_entropy_from_logits(logits)` up to f32 + rounding, or None; and the logsumexp's per-token max `m` and exp-sum `s` + (f32 [...]) for `_head_logps_cotangent`. + """ + x = jnp.where(jnp.isfinite(logits), logits.astype(jnp.float32), 0.0) + # jax.nn.logsumexp(x, axis=-1), op for op (jax/_src/ops/special.py, b=None), + # with its max and exp-sum kept. The max is a stop_gradient there too. + amax = jnp.max(x, axis=-1, initial=-jnp.inf) + m = jax.lax.stop_gradient( + jax.lax.select(jnp.isfinite(amax), amax, jax.lax.full_like(amax, 0)) + ) + shifted = x - m[..., None] + e = jnp.exp(shifted) + s = jnp.abs(jnp.sum(e, axis=-1)) + log_s = jnp.log(s) + # The target logit as a select-sum over the one-hot. It has one non-zero term, + # so it is exactly the gathered value, and its transpose is a select instead + # of take_along_axis's scatter-add. + hit, valid = _target_hit(ids, x.shape) + target = jnp.sum(jnp.where(hit, x, 0.0), axis=-1) + target = jnp.where(valid, target, jnp.nan) + logps = target - (log_s + m) + entropy = None + if with_entropy: + # The entropy from the same exponentials, so the whole forward is one pass + # over the logits: with p = e / s, -sum(p * log p) = log(s) - sum(e * (x - + # m)) / s. compute_entropy_from_logits reaches the same value through a + # second log_softmax and a second exp pass. The entropy is a metric here: + # GRPO's loss stops its gradient, and an entropy with a cotangent takes the + # autodiff backward. + entropy = log_s - jnp.sum(e * shifted, axis=-1) / s + return logps, entropy, m, s + + +def _head_logps_cotangent( + logits: jax.Array, + ids: jax.Array, + g: jax.Array, + m: jax.Array, + s: jax.Array, +) -> jax.Array: + """Cotangent of `logits` for `selective_log_softmax(logits, ids)` and `g`. + + What `jax.vjp` of `selective_log_softmax` computes, as one elementwise + expression. With x the finite-masked f32 logits, autodiff adds the transpose + of the gather (a scatter-add of g into zeros) to the transpose of the + logsumexp (sub -> -g, log -> divide by s, abs -> unchanged for s > 0, + reduce_sum -> broadcast, exp -> multiply by exp(x - m)) and masks the sum + with isfinite. Here the scatter is a select on the one-hot, which puts the + same g at the target and zeros elsewhere; the other operations are the same + f32 operations in the same order, with the max `m` and exp-sum `s` that + `_head_logps_stats` kept. + + Args: + logits: [..., V] logits, as `selective_log_softmax` received them. + ids: [...] target ids. + g: [...] cotangent of the log-probs. + m: [...] max from `_head_logps_stats`. + s: [...] exp-sum from `_head_logps_stats`. + + Returns: + f32 [..., V] cotangent of `logits`. + """ + finite = jnp.isfinite(logits) + x = jnp.where(finite, logits.astype(jnp.float32), 0.0) + hit, _ = _target_hit(ids, x.shape) + g = g.astype(jnp.float32) + d_x = jnp.where(hit, g[..., None], 0.0) + (-g / s)[..., None] * jnp.exp( + x - m[..., None] + ) + return jnp.where(finite, d_x, 0.0) + + def _mesh_axes(entry) -> tuple[str, ...]: if entry is None: return () @@ -1195,6 +1332,9 @@ def maybe_compute_sharded_chunked_logps( 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. + With "fused", chunks are scored by `_head_logps_stats`, and a backward in + which only the log-probs have a cotangent forms the logits' cotangent with + `_head_logps_cotangent` (see the comment above `_target_hit`). Returns None, after logging "[logps-head] sharded head off, using compute_chunked_logps: ", when: @@ -1283,6 +1423,7 @@ def fallback(reason): gather_dims.append(dim) param_dtypes = [jnp.result_type(p) for p in params] num_out = 2 if return_entropy else 1 + fused = mode == "fused" 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", @@ -1307,10 +1448,15 @@ def merged(full): values[i] = x return nnx.merge(graphdef, jax.tree.unflatten(treedef, values)) - def chunk_outputs(full, hs, ids, with_entropy): + def chunk_logits(full, hs): + # The f32 logits a chunk is scored from. logits = merged(full).compute_final_logits(hs).astype(jnp.float32) if temperature != 0.0 and temperature != 1.0: logits /= temperature + return logits + + def chunk_outputs(full, hs, ids, with_entropy): + logits = chunk_logits(full, hs) logps = selective_log_softmax(logits, ids) if with_entropy: return logps, compute_entropy_from_logits(logits) @@ -1329,19 +1475,32 @@ def scatter(g, dim): g, gather_axes, scatter_dimension=dim, tiled=True ) - def scan_forward(full, h, t): + def chunk_outputs_fused(full, hs, ids, with_stats): + # `chunk_outputs` for "fused", plus each token's max and exp-sum when the + # backward will need them. + logps, entropy, m, s = _head_logps_stats( + chunk_logits(full, hs), ids, return_entropy + ) + outs = (logps, entropy) if return_entropy else (logps,) + return outs + (m, s) if with_stats else outs + + def scan_forward(full, h, t, with_stats=False): 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) + if fused: + outs = chunk_outputs_fused(full, hs, ids, with_stats) + else: + 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)) + num_bufs = num_out + (2 if with_stats else 0) + bufs = tuple(jnp.zeros(t.shape, jnp.float32) for _ in range(num_bufs)) bufs, _ = jax.lax.scan(step, bufs, jnp.arange(h.shape[1] // chunk)) return bufs @@ -1363,10 +1522,69 @@ def head_local_fwd(h, t, *ps): full = [ gather(p, d, dt) for p, d, dt in zip(ps, gather_dims, gather_dtypes) ] + if fused: + # Also keep the max and exp-sum of every token (f32 [B, T] each). + bufs = scan_forward(full, h, t, with_stats=True) + return bufs[:num_out], (h, t, tuple(full), bufs[num_out:]) return scan_forward(full, h, t), (h, t, tuple(full)) + def bwd_fused(h, t, full, stats, d_logps): + # The "local" backward's token chunks, with the logits' cotangent formed in + # one pass (`_head_logps_cotangent`) from the forward's max and exp-sum + # instead of by autodiff of `selective_log_softmax` (a scatter-add and f32 + # [chunk, vocab] copies). The table and hidden-state cotangents still come + # from `jax.vjp` of `compute_final_logits` and are accumulated as there. + m_all, s_all = stats + scaled = temperature != 0.0 and temperature != 1.0 + + 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) + g = jax.lax.dynamic_slice_in_dim(d_logps, start, chunk, axis=1) + m = jax.lax.dynamic_slice_in_dim(m_all, start, chunk, axis=1) + s = jax.lax.dynamic_slice_in_dim(s_all, start, chunk, axis=1) + raw, vjp = jax.vjp( + lambda full_, hs_: merged(full_).compute_final_logits(hs_), + list(full), + hs, + ) + # `chunk_logits` on `raw`, the logits' cotangent, then the transposes of + # the temperature division (a division) and of `astype(float32)` (a + # cast). + logits = raw.astype(jnp.float32) + if scaled: + logits /= temperature + d_logits = _head_logps_cotangent(logits, ids, g, m, s) + if scaled: + d_logits /= temperature + # An identity that pins one materialized cotangent. Without it XLA:TPU + # writes the cotangent for the hidden-state matmul and recomputes it (exp + # included) inside the table-gradient matmul's operand fusion; with it + # the cotangent becomes the logits matmul's epilogue and both gradient + # matmuls read it (in a deviceless TPU compile of the head at chunk + # 1024: 0.47 GiB fewer temporaries and a third fewer transcendentals). + d_logits = jax.lax.optimization_barrier(d_logits.astype(raw.dtype)) + d_full, d_hs = vjp(d_logits) + acc = tuple(a + gg.astype(jnp.float32) for a, gg 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)] + return (dh, None, *[a.astype(dt) for a, dt in zip(acc, param_dtypes)]) + def head_local_bwd(res, cts): - h, t, full = res + if fused: + h, t, full, stats = res + else: + 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. @@ -1382,6 +1600,11 @@ def head_local_bwd(res, cts): ] zeros = [jnp.zeros(s, dt) for s, dt in zip(shard_shapes, param_dtypes)] return (jnp.zeros_like(h), None, *zeros) + if fused and use_logps and not use_entropy: + # An entropy with a cotangent (a loss that differentiates the entropy) + # takes the autodiff backward below, which is the "local" backward + # exactly. + return bwd_fused(h, t, full, stats, d_logps) def step(carry, i): acc, dh = carry From bee63532f4cd97b09a1144ad5aa18b985eabe178 Mon Sep 17 00:00:00 2001 From: Param Bole Date: Thu, 15 Oct 2026 08:48:51 -0700 Subject: [PATCH 3/4] recipes(397b v7x): sixteen MaxText keys, the fused head and pad first join TRAINER_TUNED_FLAGS TRAINER_TUNED_FLAGS=true now selects the trainer cell that measured fastest on the RL trainer bench with every numerics gate passing. In mlperf_397b_256_v7x.sh the block of MaxText keys (part 2) gains sixteen keys and loses one: - wi_tile_dlhs_mlp_dim=2048, decoder_layer_input=device, moe_routing_maps=device and moe_gmm_v2_dlhs_transpose_rhs=true: the K tile of wi's input-gradient GMM matched to the fused gate and up weights, each decoder layer's input kept in HBM, the backward reusing the forward's MoE routing maps, and the input-gradient GMMs reading the expert weights through the kernel's transpose_rhs. - moe_routing_index_math=count, moe_unsplit_prefused_wi=true, moe_roe_packed_routing_gather=true, moe_router_grad_from_wo_input=true and moe_wi_gmm_fused_swiglu=true: the MoE keys of AI-Hypercomputer/maxtext#5758 in the set its commits recommend (the routing index maps by counting instead of argsorts, the fused wi kept whole up to the GEMM, one packed routing all-gather per layer, the routing-weight gradient from the down projection's input, the SwiGLU in the wi GEMM's epilogue). - gdn_conv1d_input_in_qkv_dtype=true and gdn_conv1d_dense_token_stack=true: the GDN Conv1D backward without its f32 copy of the input, and the forward kernel's Conv1D output built from dense per-token views. - wi_tile_fwd_batch_seq, wi_tile_drhs_batch_seq, wo_tile_fwd_batch_seq, wo_tile_dlhs_batch_seq and wo_tile_drhs_batch_seq=512 (from 256): row tiles of 512 for the five MoE GMMs with 1024 tiles. wi_tile_dlhs_batch_seq stays 256 next to its 2048 K tile. - sa_use_base2_exp=false leaves the block, so the trainer keeps the recipe's sa_use_base2_exp=true: base-2 exponentials measured 0.025 s faster per micro-batch (o275x-gx-b2e), and their numerics run passes (nxgx-b2e). In mlperf_gbs1024_submission.sh, which turns TRAINER_TUNED_FLAGS on, the switch also sets three defaults, each still overridable: - COMPUTE_LOGPS_CHUNK_SIZE=4096 (was 1024), and 1024 when TUNIX_LOGPS_SHARDED_HEAD turns the sharded head off (0, false, off or no, in any case): chunks of 4096 were measured with the head only, and without it the log-prob program failed to compile on TPU at chunks of 2048 and 4096. - TUNIX_LOGPS_SHARDED_HEAD=fused (was local). TRAINER_TUNED_FLAGS=false keeps local. - PACK_PAD_POSITION=start: each packed row puts its padding first. Bench (AI-Hypercomputer/maxtext#5701 driver, Qwen3.5-397B-A17B, 128 TPU v7x, mesh fsdp 32 x context 4 x expert 2, ragged_buffer_factor 2.75, 9 micro-batches per step, mean of steps 1 to 4, no token dropped in any arm): - before these keys, the fused head and pad first: 106.30 s per step (o275x-hd), 2,774 tokens/s/chip; - this cell: 82.671 and 82.670 s per step in two pods (o275x-c9pfxhv-l3), 3,567 tokens/s/chip. Part of the difference is the staged step boundary (AI-Hypercomputer/maxtext#5718), which this recipe already turns on and o275x-hd did not run; per regular micro-batch (the median of micro-batches 3 to 9), which the boundary does not change, about 10.77 -> 9.01 s; - within it, each against a control in the same pod: padding first 83.285 -> 82.043 s per step (-1.242 s) and the fused head 83.253 -> 82.708 s per step (-0.545 s). The recipe lists the measurement of each key. The measured build switched the MaxText changes with environment variables, and it also had one attention change on that neither MaxText PR carries (attention's load-balancing permute hoisted; alone 99.504 -> 99.267 s per step). The keys themselves, on the MaxText change's own code (2d82c5199) without gdn_conv1d_dense_token_stack, the fused head and pad first, measured 85.295 and 85.251 s per step (o275x-prc-l3-prof, o275x-prc-l3-r), within 0.23 s of what the environment-switch arms predicted. Both arms set moe_route=device, which nothing in #5698 or #5758 tags, so moe_route stays out of the block as before; and the stop of the frozen router's gradient, an environment switch in the measured build, is unconditional in #5758, with no key. Numerics (loss gates off, random advantages, three steps, against two reference runs whose gradient norms differ by 1.9e-5 relative): all six gates pass for this cell (nxc9pfxhv) and for the keys on the MaxText change's own code (nxprc). moe_wo_gmm_fused_dswiglu stays out: next to the fused SiLU GMM at these tiles it computed a wrong SwiGLU gradient, and its fix is still being verified on TPU. moe_routing_maps and moe_gmm_v2_dlhs_transpose_rhs exist in AI-Hypercomputer/maxtext#5698, and the seven other new MaxText keys in the MaxText change of AI-Hypercomputer/maxtext#5758 (on top of #5698); none of the nine is defined at this tree's pin (e526ea02), and MaxText refuses a key it does not define when the trainer starts. TRAINER_TUNED_FLAGS=true, the default of mlperf_gbs1024_submission.sh, therefore needs requirements/maxtext_requirements.txt moved to a MaxText with #5758 before an image built from this tree runs it. TRAINER_TUNED_FLAGS=false and the other 397B v7x recipes pass none of the nine. Tests (mlperf_submission_launcher_test.py): the exact-set test expects 23 keys, _MAXTEXT_DEFAULTS has the nine new keys' defaults, and the switch-off and other-recipe tests check that none of the nine is passed; the default launch has base-2 exponentials on, chunks of 4096, the fused head and --pack_pad_position=start; TRAINER_TUNED_FLAGS=false keeps the local head and the padding at the end; an explicit pad position, head or chunk size wins; TUNIX_LOGPS_SHARDED_HEAD=0, false, OFF or No gives chunks of 1024; the other 397B v7x recipes pass no pad position. Fifteen deliberately wrong copies of the recipes each fail at least one of these tests: a key dropped from the block or given the wrong value, one tile left out, wi_tile_dlhs_batch_seq raised to 512, sa_use_base2_exp=false kept, moe_wo_gmm_fused_dswiglu or moe_route=device added, a new key passed outside the block at its default, the local head under the switch, the chunk fallback removed or made case-sensitive, chunks of 1024 by default, pad first outside the switch, pad_align passed, and both recipes as they were before this change. --- .../mlperf_submission_launcher_test.py | 108 ++++++++++++++++-- .../examples/recipes/mlperf_397b_256_v7x.sh | 101 ++++++++++++++-- .../recipes/mlperf_gbs1024_submission.sh | 29 ++++- 3 files changed, 214 insertions(+), 24 deletions(-) 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 f5c7365014..9dfb672d41 100644 --- a/tests/experimental/examples/deepswe_dist/mlperf_submission_launcher_test.py +++ b/tests/experimental/examples/deepswe_dist/mlperf_submission_launcher_test.py @@ -51,20 +51,63 @@ "cast_grads_after_all_reduce": "false", "gdn_cp_matmul_precision": "highest", "gdn_cp_state": "remat", + # Defaults in the MaxText of AI-Hypercomputer/maxtext#5698 and #5758; see + # _BENCH_MAXTEXT_KEYS. + "gdn_conv1d_dense_token_stack": "false", + "gdn_conv1d_input_in_qkv_dtype": "false", + "moe_gmm_v2_dlhs_transpose_rhs": "false", + "moe_roe_packed_routing_gather": "false", + "moe_router_grad_from_wo_input": "false", + "moe_routing_index_math": "sort", + "moe_routing_maps": "remat", + "moe_unsplit_prefused_wi": "false", + "moe_wi_gmm_fused_swiglu": "false", } _TUNED_KEYS = { "cast_grads_after_all_reduce": ("false", "true"), "context_parallel_strategy": ("ring", "all_gather"), + "decoder_layer_input": ("offload", "device"), + "gdn_conv1d_dense_token_stack": ("false", "true"), + "gdn_conv1d_input_in_qkv_dtype": ("false", "true"), "gdn_cp_matmul_precision": ("highest", "high"), "gdn_cp_mode": ("head", "seq"), "gdn_cp_state": ("remat", "device"), + "moe_gmm_v2_dlhs_transpose_rhs": ("false", "true"), + "moe_roe_packed_routing_gather": ("false", "true"), + "moe_router_grad_from_wo_input": ("false", "true"), + "moe_routing_index_math": ("sort", "count"), + "moe_routing_maps": ("remat", "device"), + "moe_unsplit_prefused_wi": ("false", "true"), + "moe_wi_gmm_fused_swiglu": ("false", "true"), "num_moe_token_chunks": ("4", "2"), "sa_fuse_reciprocal": ("false", "true"), - "sa_use_base2_exp": ("true", "false"), + "wi_tile_dlhs_mlp_dim": ("1024", "2048"), + "wi_tile_drhs_batch_seq": ("256", "512"), + "wi_tile_fwd_batch_seq": ("256", "512"), + "wo_tile_dlhs_batch_seq": ("256", "512"), + "wo_tile_drhs_batch_seq": ("256", "512"), + "wo_tile_fwd_batch_seq": ("256", "512"), } +# Tuned keys that the MaxText in requirements/maxtext_requirements.txt +# (e526ea02) does not define: two from AI-Hypercomputer/maxtext#5698 and seven +# from AI-Hypercomputer/maxtext#5758, which is on top of it. MaxText refuses a +# key it does not define when the trainer starts, so only +# TRAINER_TUNED_FLAGS=true may pass them. +_BENCH_MAXTEXT_KEYS = ( + "gdn_conv1d_dense_token_stack", + "gdn_conv1d_input_in_qkv_dtype", + "moe_gmm_v2_dlhs_transpose_rhs", + "moe_roe_packed_routing_gather", + "moe_router_grad_from_wo_input", + "moe_routing_index_math", + "moe_routing_maps", + "moe_unsplit_prefused_wi", + "moe_wi_gmm_fused_swiglu", +) # The launch whose trainer the RL trainer bench (AI-Hypercomputer/maxtext#5701) # measured TRAINER_TUNED_FLAGS on: the replay layout (buffer factor 2.75 with -# the dropless replay), with context on device. +# the dropless replay), with context on device, log-prob chunks of 4096 with +# the fused log-prob head, and packed rows padded first. _BENCH_CELL = { "BATCH_SIZE": "72", "RAGGED_BUFFER_FACTOR": "2.75", @@ -72,7 +115,9 @@ "REMAT_CONTEXT": "device", "MOE_EXCLUDE_PADDING": "true", "LOG_MOE_DROPS": "true", - "COMPUTE_LOGPS_CHUNK_SIZE": "1024", + "COMPUTE_LOGPS_CHUNK_SIZE": "4096", + "TUNIX_LOGPS_SHARDED_HEAD": "fused", + "PACK_PAD_POSITION": "start", } @@ -175,7 +220,13 @@ def test_the_default_trainer_is_the_measured_bench_cell(self): self.assertEqual(flags["num_moe_token_chunks"], "2") self.assertEqual(flags["moe_exclude_padding_from_routing"], "true") self.assertEqual(flags["log_moe_drops"], "true") - self.assertIn("--compute_logps_chunk_size=1024 ", trainer) + # Base-2 exponentials stay on: the block no longer turns them off. + self.assertEqual(flags["sa_use_base2_exp"], "true") + self.assertIn("--compute_logps_chunk_size=4096 ", trainer) + self.assertRegex(trainer, r"\bTUNIX_LOGPS_SHARDED_HEAD=fused\b") + self.assertRegex( + _jobset(self.default, "orch"), r"--pack_pad_position=start\b" + ) # The HBM budget reaches the trainer env, the Pathways workers and the # Pathways proxy (which compiles the trainer program). self.assertEqual(trainer.count(_HBM_CAP), 3) @@ -363,35 +414,65 @@ def test_staging_without_the_pipelined_boundary_is_refused(self): ): _render(PIPELINE_STEP_BOUNDARY="false", STAGE_NEXT_MICROBATCH="true") - def test_packed_rows_pad_at_the_end_by_default(self): + def test_packed_rows_pad_first_with_the_tuned_flags(self): orch = _jobset(self.default, "orch") + self.assertRegex(orch, r"--pack_pad_position=start\b") + self.assertNotIn("--pack_pad_align", orch) + # An explicit value wins. + orch = _jobset(_render(PACK_PAD_POSITION="end"), "orch") + self.assertRegex(orch, r"--pack_pad_position=end\b") + self.assertNotIn("--pack_pad_position=start", orch) + + def test_packed_rows_pad_at_the_end_without_the_tuned_flags(self): + orch = _jobset(self.default_off, "orch") self.assertNotIn("--pack_pad_position", orch) self.assertNotIn("--pack_pad_align", orch) def test_the_packed_row_pad_position_reaches_only_the_orchestrator(self): - render = _render(PACK_PAD_POSITION="start", PACK_PAD_ALIGN="8192") + render = _render( + TRAINER_TUNED_FLAGS="false", PACK_PAD_POSITION="start", + PACK_PAD_ALIGN="8192", + ) orch = _jobset(render, "orch") self.assertRegex(orch, r"--pack_pad_position=start\b") self.assertRegex(orch, r"--pack_pad_align=8192\b") docs = render.split("\n---\n") - default_docs = self.default.split("\n---\n") + default_docs = self.default_off.split("\n---\n") self.assertLen(docs, len(default_docs)) changed = [doc for doc, before in zip(docs, default_docs) if doc != before] self.assertEqual(changed, [orch]) def test_the_sharded_log_prob_head_is_on_in_the_trainer_only(self): self.assertRegex( - _jobset(self.default, "train"), r"\bTUNIX_LOGPS_SHARDED_HEAD=local\b" + _jobset(self.default, "train"), r"\bTUNIX_LOGPS_SHARDED_HEAD=fused\b" ) for suffix in ("orch", "roll-0"): self.assertNotIn( "TUNIX_LOGPS_SHARDED_HEAD", _jobset(self.default, suffix), suffix ) + # TRAINER_TUNED_FLAGS=false keeps the "local" head; an explicit value wins. + off = _jobset(self.default_off, "train") + self.assertRegex(off, r"\bTUNIX_LOGPS_SHARDED_HEAD=local\b") + self.assertNotIn("TUNIX_LOGPS_SHARDED_HEAD=fused", off) + local = _jobset(_render(TUNIX_LOGPS_SHARDED_HEAD="local"), "train") + self.assertRegex(local, r"\bTUNIX_LOGPS_SHARDED_HEAD=local\b") + self.assertNotIn("TUNIX_LOGPS_SHARDED_HEAD=fused", local) + self.assertIn("--compute_logps_chunk_size=4096 ", local) def test_the_sharded_log_prob_head_can_be_turned_off(self): - trainer = _jobset(_render(TUNIX_LOGPS_SHARDED_HEAD="0"), "train") - self.assertRegex(trainer, r"\bTUNIX_LOGPS_SHARDED_HEAD=0\b") - self.assertNotIn("TUNIX_LOGPS_SHARDED_HEAD=local", trainer) + # Without the head, chunks of 4096 failed to compile on TPU: 1024 instead. + for value in ("0", "false", "OFF", "No"): + trainer = _jobset(_render(TUNIX_LOGPS_SHARDED_HEAD=value), "train") + self.assertRegex(trainer, rf"\bTUNIX_LOGPS_SHARDED_HEAD={value}\b") + self.assertNotIn("TUNIX_LOGPS_SHARDED_HEAD=local", trainer) + self.assertNotIn("TUNIX_LOGPS_SHARDED_HEAD=fused", trainer) + self.assertIn("--compute_logps_chunk_size=1024 ", trainer, value) + # An explicit chunk size still wins. + trainer = _jobset( + _render(TUNIX_LOGPS_SHARDED_HEAD="0", COMPUTE_LOGPS_CHUNK_SIZE="512"), + "train", + ) + self.assertIn("--compute_logps_chunk_size=512 ", trainer) def test_start_up_flags_can_be_turned_off(self): orch = _jobset( @@ -454,6 +535,8 @@ def test_tuned_flags_false_turns_the_block_off(self): {key: off for key, (off, _) in _TUNED_KEYS.items()}, ) self.assertNotIn("offloading_copy_to_sparsecore", trainer) + for key in _BENCH_MAXTEXT_KEYS: + self.assertNotIn(key, flags) # The fully dropless buffers and chunks of 512, as before the switch. self.assertEqual(flags["ragged_buffer_factor"], "-1") self.assertEqual(flags["context"], "offload") @@ -514,6 +597,8 @@ def test_tuned_flags_are_off_in_the_other_397b_v7x_recipes(self): for key, (off, _) in _TUNED_KEYS.items(): if key != "num_moe_token_chunks": # two in these recipes either way self.assertEqual(_value(flags, key), off, f"{name}: {key}") + for key in _BENCH_MAXTEXT_KEYS: + self.assertNotIn(key, flags, name) def test_head_and_staging_are_off_in_the_other_397b_v7x_recipes(self): for name in ( @@ -526,6 +611,7 @@ def test_head_and_staging_are_off_in_the_other_397b_v7x_recipes(self): orch = _jobset(render, "orch") self.assertIn("--no-pipeline_step_boundary", orch, name) self.assertIn("--no-stage_next_microbatch", orch, name) + self.assertNotIn("--pack_pad_position", orch, name) def test_trajectory_store_is_enabled_by_default(self): orch = _jobset(self.default, "orch") diff --git a/tunix/experimental/examples/recipes/mlperf_397b_256_v7x.sh b/tunix/experimental/examples/recipes/mlperf_397b_256_v7x.sh index 8c7cbcc33a..3955a033ad 100755 --- a/tunix/experimental/examples/recipes/mlperf_397b_256_v7x.sh +++ b/tunix/experimental/examples/recipes/mlperf_397b_256_v7x.sh @@ -222,17 +222,24 @@ fi # replace the values above. # context_parallel_strategy=all_gather: attention all-gathers K/V over the context # axis instead of passing it around a ring. -# sa_fuse_reciprocal=true, sa_use_base2_exp=false: Splash attention kernel options -# (SplashConfig fuse_reciprocal / use_base2_exp). These three keys alone, without -# the libtpu flag: 130.8 -> 127.5 s/step on the bench. +# sa_fuse_reciprocal=true: a Splash attention kernel option (SplashConfig +# fuse_reciprocal). With context_parallel_strategy=all_gather and +# sa_use_base2_exp=false, without the libtpu flag, on the bench: 130.8 -> 127.5 +# s/step. +# sa_use_base2_exp stays at true (set above). The block set it to false until the +# bench measured the base-2 exponentials 0.025 s faster per micro-batch +# (o275x-gx-b2e); q * log2(e) is rounded in bf16, and its numerics run (nxgx-b2e) +# passes all six gates. # gdn_cp_mode=seq: GDN context parallelism shards the sequence instead of the heads, # which drops head mode's all-to-alls (MaxText's `auto` picks seq at context 4). # gdn_cp_state=device: keeps seq mode's carries (conv halo, local transition, incoming # state), so the backward does not replay the halo exchange and the first pass. # moe_route stays at its default (remat). The bench image's MaxText tags nothing with -# it, so the moe_route=device of the measured runs changed nothing there. MaxText at -# this recipe's pin tags the ragged sort's integer routing results, where device -# would keep them on the chip at a memory cost no run has measured. +# it (it tags the routing arrays moe_routing_maps, below), and neither does +# AI-Hypercomputer/maxtext#5758, so the moe_route=device of the measured runs changed +# nothing there. MaxText at this recipe's pin tags the ragged sort's integer routing +# results with it, where device would keep them on the chip at a memory cost no run +# has measured. # num_moe_token_chunks=2 comes from NUM_MOE_TOKEN_CHUNKS: the default here, and set by # mlperf_gbs1024_submission.sh for its replay layout. # The keys above with two token chunks, on top of the libtpu flag (bench, same pod, with @@ -241,8 +248,88 @@ fi # cast_grads_after_all_reduce=true: all-reduce-only bf16 grads cast to f32 after it. # The two together (bench, same pod, the sharded log-prob head on in both arms): # 107.59 -> 106.24 s/step; grad_norm moved 1.5e-5 relative, less than a rerun (1.9e-5). +# wi_tile_dlhs_mlp_dim=2048: the K tile of wi's input-gradient GMM. The trainer fuses +# wi's gate and up weights (prefuse_moe_weights, on in mlperf_base.sh), so that +# contraction is 2048 wide, and a 1024 tile takes two K steps that each fetch the +# weight block again. +# decoder_layer_input=device: each decoder layer's input stays in HBM instead of being +# offloaded to host memory. +# moe_routing_maps=device: the backward reuses the forward's MoE routing maps (top-k +# expert ids, ragged-sort permutations, sorted token indices, group sizes) instead of +# running top-k and the sorts again. It saves values; it does not change them. +# moe_gmm_v2_dlhs_transpose_rhs=true: the input-gradient GMMs read the expert weights +# through the kernel's transpose_rhs instead of from a transposed copy of them. +# These four on the bench, as the time of a regular micro-batch (a step has nine; the +# median of micro-batches 3 to 9 over steps 1 to 4), with no token dropped in any arm: +# the tile alone, same pod: 10.773 -> 10.605 s (o275x-hd -> o275x-t2048, 106.3 -> 104.7 +# s/step); +# decoder_layer_input=device alone, other pods: 10.754 -> 10.712 s (x3l -> x3dli, both +# with the staged step boundary: 98.3 -> 98.0 s/step); +# those two on top of the GDN change of AI-Hypercomputer/maxtext#5667, other pods: +# 10.582 -> 10.319 s (o275x-t67 -> o275x-stk1a, 104.5 -> 102.1 s/step); against two +# reference runs the loss moved 4.6e-8 and grad_norm 3.9e-5 relative, a rerun 1.9e-5 +# (nxstk1a-rp); +# on top of that, same pod: moe_routing_maps=device 10.319 -> 10.189 s (o275x-stk1a-a, +# 102.1 -> 100.9 s/step) and moe_gmm_v2_dlhs_transpose_rhs=true 10.319 -> 10.195 s +# (o275x-stk1a-b, 102.1 -> 100.9 s/step). +# moe_routing_index_math=count: the ring-of-experts ragged sort builds its routing +# index maps with one sort of a packed key and by counting, instead of stable +# argsorts: the same maps. +# moe_unsplit_prefused_wi=true: the fused wi = [w0 | w1] stays whole up to the GEMM +# instead of being split at the layer's entry and joined again before it. +# moe_roe_packed_routing_gather=true: the layer routes its tokens once and all-gathers +# their routing weights and expert ids as one array, not twice per token chunk. +# moe_router_grad_from_wo_input=true: the down projection and the unsort run as one op +# whose backward takes the routing-weight gradient from the GEMM's input, so the +# rematerialized backward does not run the down-projection GEMM again. +# moe_wi_gmm_fused_swiglu=true: the wi GEMM applies the SwiGLU to its f32 accumulator +# in its epilogue, instead of a separate pass over its [rows, 2 * mlp] output. +# gdn_conv1d_input_in_qkv_dtype=true: the GDN Conv1D backward keeps its halo-extended +# input in qkv's dtype and casts each tap where it reads it, instead of writing an +# f32 copy of it. +# gdn_conv1d_dense_token_stack=true: the GDN forward kernel assembles its Conv1D output +# from dense per-token views, so the SiLU after it does not run on rows padded from 1 +# to 8 sublanes. +# wi_tile_fwd_batch_seq, wi_tile_drhs_batch_seq, wo_tile_fwd_batch_seq, +# wo_tile_dlhs_batch_seq and wo_tile_drhs_batch_seq=512: row tiles of 512 for the five +# MoE GMMs with 1024 tiles. wi_tile_dlhs_batch_seq stays 256 next to its 2048 K tile. +# These on the bench, regular micro-batch or s/step as above, 0 tokens dropped: +# moe_routing_index_math=count 10.022 -> 9.875 s (o275x-mf-cnt-prof), +# moe_unsplit_prefused_wi with two exact rewrites this block leaves off 10.022 -> +# 9.852 s (o275x-mf-gb-prof) and moe_router_grad_from_wo_input 10.022 -> 9.765 s +# (o275x-mf-a14-prof), each against o275x-mf-ctl in the same pod; +# moe_roe_packed_routing_gather 9.938 s (o275x-mf-pra-prof) against 10.033 s in +# another pod (o275x-gdn-ctl); +# the five MoE keys together, same pod: 10.0445 -> 9.510 s (o275x-s5c-l3-prof -> +# o275x-s5moe5-l3-prof, 91.87 -> 87.04 s/step); +# gdn_conv1d_input_in_qkv_dtype alone: -0.050 s per micro-batch; +# then against o275x-s6c-l3 (86.970 s/step) in the same pod: moe_wi_gmm_fused_swiglu +# 86.253 s/step (o275x-s6silu-l3-prof) and the five 512 tiles 86.364 s/step +# (o275x-s6t512-l3); gdn_conv1d_dense_token_stack 86.926 -> 85.731 s/step, same pod, +# against an arm whose one change is a no-op (o275x-s6n2-l3-prof -> +# o275x-s6gr-l3-prof). +# The whole TRAINER_TUNED_FLAGS cell of mlperf_gbs1024_submission.sh (these keys, its +# log-prob head and chunk size, padding first in packed rows, the staged step boundary): +# 82.671 and 82.670 s/step in two pods (o275x-c9pfxhv-l3), against 106.30 s/step before +# these keys, the fused head and padding first (o275x-hd, which also ran without the +# staged step boundary; per regular micro-batch, which the boundary leaves alone, about +# 10.77 -> 9.01 s). That build switched the MaxText changes with environment variables +# and also hoisted attention's load-balancing permute, which neither MaxText PR has +# (alone on the bench: 99.504 -> 99.267 s/step). The keys themselves, on the MaxText +# change's own code without gdn_conv1d_dense_token_stack, the fused head and padding +# first: 85.295 and 85.251 s/step (o275x-prc-l3-prof, o275x-prc-l3-r). +# Numerics (loss gates off, random advantages, three steps, against two reference runs +# whose gradient norms differ by 1.9e-5 relative): all six gates pass (nxc9pfxhv). +# moe_routing_maps and moe_gmm_v2_dlhs_transpose_rhs come from +# AI-Hypercomputer/maxtext#5698, and moe_routing_index_math, moe_unsplit_prefused_wi, +# moe_roe_packed_routing_gather, moe_router_grad_from_wo_input, moe_wi_gmm_fused_swiglu, +# gdn_conv1d_input_in_qkv_dtype and gdn_conv1d_dense_token_stack from the MaxText change +# of AI-Hypercomputer/maxtext#5758 (on top of #5698). The MaxText pinned in +# requirements/maxtext_requirements.txt (e526ea02) has none of these nine keys, and +# MaxText refuses a key it does not define when the trainer starts, so +# TRAINER_TUNED_FLAGS=true needs that pin moved to a MaxText with #5758. if [[ "${TRAINER_TUNED_FLAGS:-false}" == "true" ]]; then - export MAXTEXT_EXTRA_FLAGS="${MAXTEXT_EXTRA_FLAGS} context_parallel_strategy=all_gather sa_fuse_reciprocal=true sa_use_base2_exp=false gdn_cp_mode=seq gdn_cp_state=device gdn_cp_matmul_precision=high cast_grads_after_all_reduce=true" + export MAXTEXT_EXTRA_FLAGS="${MAXTEXT_EXTRA_FLAGS} context_parallel_strategy=all_gather sa_fuse_reciprocal=true gdn_cp_mode=seq gdn_cp_state=device gdn_cp_matmul_precision=high cast_grads_after_all_reduce=true wi_tile_dlhs_mlp_dim=2048 decoder_layer_input=device moe_routing_maps=device moe_gmm_v2_dlhs_transpose_rhs=true moe_routing_index_math=count moe_unsplit_prefused_wi=true moe_roe_packed_routing_gather=true moe_router_grad_from_wo_input=true moe_wi_gmm_fused_swiglu=true gdn_conv1d_input_in_qkv_dtype=true gdn_conv1d_dense_token_stack=true wi_tile_fwd_batch_seq=512 wi_tile_drhs_batch_seq=512 wo_tile_fwd_batch_seq=512 wo_tile_dlhs_batch_seq=512 wo_tile_drhs_batch_seq=512" fi export DEBUG=${DEBUG:-0} diff --git a/tunix/experimental/examples/recipes/mlperf_gbs1024_submission.sh b/tunix/experimental/examples/recipes/mlperf_gbs1024_submission.sh index e3e6f10a15..729e602680 100644 --- a/tunix/experimental/examples/recipes/mlperf_gbs1024_submission.sh +++ b/tunix/experimental/examples/recipes/mlperf_gbs1024_submission.sh @@ -70,8 +70,13 @@ export ROLLOUT_ATTENTION_DP="${ROLLOUT_ATTENTION_DP:-false}" # (AI-Hypercomputer/maxtext#5701), the replay layout: # RAGGED_BUFFER_FACTOR=2.75 with RETRY_WHEN_TOKENS_DROPPED=true, two MoE token # chunks, REMAT_CONTEXT=device, MOE_EXCLUDE_PADDING=true (padding stays out of -# MoE routing), LOG_MOE_DROPS=true (per-micro-batch drop audit), and -# COMPUTE_LOGPS_CHUNK_SIZE=1024 (mlperf_base.sh's default is 512). +# MoE routing), LOG_MOE_DROPS=true (per-micro-batch drop audit), +# COMPUTE_LOGPS_CHUNK_SIZE=4096 (mlperf_base.sh's default is 512), and +# PACK_PAD_POSITION=start (each packed row's padding goes first, which evens the +# attention work of the context ranks under load-balanced context parallelism). +# Log-prob chunks of 4096 were measured with the sharded log-prob head (below) only; +# without it the log-prob program failed to compile on TPU at chunks of 2048 and +# 4096, so TUNIX_LOGPS_SHARDED_HEAD=0 (or false, off, no) keeps chunks of 1024. # RAGGED_BUFFER_FACTOR=-1 keeps the fully dropless buffers instead: context # offload and four token chunks, without the replay, padding and audit flags. # The replay layout ran within an XLA HBM budget of 84992 MiB on the bench; the fully @@ -82,10 +87,16 @@ export ROLLOUT_ATTENTION_DP="${ROLLOUT_ATTENTION_DP:-false}" # (FIRST_PHASE_RAGGED_BUFFER_FACTOR); with either, the layout keeps four chunks # and context offload, which the measured cell's memory does not cover. # TRAINER_TUNED_FLAGS=false renders the trainer as before (the fully dropless -# buffers, chunks of 512). +# buffers, chunks of 512, the "local" log-prob head, padding at the end of packed +# rows). export TRAINER_TUNED_FLAGS="${TRAINER_TUNED_FLAGS:-true}" if [[ "${TRAINER_TUNED_FLAGS}" == "true" ]]; then - export COMPUTE_LOGPS_CHUNK_SIZE="${COMPUTE_LOGPS_CHUNK_SIZE:-1024}" + _logps_head="$(printf '%s' "${TUNIX_LOGPS_SHARDED_HEAD:-}" | tr -d '[:space:]' | tr '[:upper:]' '[:lower:]')" + case "${_logps_head}" in + 0|false|off|no) export COMPUTE_LOGPS_CHUNK_SIZE="${COMPUTE_LOGPS_CHUNK_SIZE:-1024}" ;; + esac + export COMPUTE_LOGPS_CHUNK_SIZE="${COMPUTE_LOGPS_CHUNK_SIZE:-4096}" + export PACK_PAD_POSITION="${PACK_PAD_POSITION:-start}" fi export DROPLESS="${DROPLESS:-true}" case "${DROPLESS}" in @@ -125,8 +136,11 @@ esac # the staged step boundary (--pipeline_step_boundary, --stage_next_microbatch), # measured on the RL trainer bench (AI-Hypercomputer/maxtext#5701). On in this # recipe only. -# TUNIX_LOGPS_SHARDED_HEAD=local: the trainer process computes log-probs with -# the sharded head (k8s_launcher.sh passes it to the trainer); 0 turns it off. +# TUNIX_LOGPS_SHARDED_HEAD=fused (TRAINER_TUNED_FLAGS=true) or local (false): the +# trainer process computes log-probs with the sharded head (k8s_launcher.sh +# passes it to the trainer); fused scores each chunk in one pass and, when only +# the log-probs are differentiated (GRPO), forms the logits' cotangent without a +# scatter. 0 turns the head off. # PIPELINE_STEP_BOUNDARY=1, STAGE_NEXT_MICROBATCH=1: the orchestrator sends the # next step's first micro-batch ahead of this step's bookkeeping, and stages it # on the trainer before this step's deferred update. Staging needs @@ -134,6 +148,9 @@ esac # the boundary pipelined but unstaged. STAGE_NEXT_MICROBATCH follows # PIPELINE_STEP_BOUNDARY unless set, so PIPELINE_STEP_BOUNDARY=false alone # turns both off. +if [[ "${TRAINER_TUNED_FLAGS}" == "true" ]]; then + export TUNIX_LOGPS_SHARDED_HEAD="${TUNIX_LOGPS_SHARDED_HEAD:-fused}" +fi export TUNIX_LOGPS_SHARDED_HEAD="${TUNIX_LOGPS_SHARDED_HEAD:-local}" export PIPELINE_STEP_BOUNDARY="${PIPELINE_STEP_BOUNDARY:-1}" export STAGE_NEXT_MICROBATCH="${STAGE_NEXT_MICROBATCH:-${PIPELINE_STEP_BOUNDARY}}" From 8c75ae30de46b58582abe7ecdeec2c4dbddb1724 Mon Sep 17 00:00:00 2001 From: Param Bole Date: Thu, 15 Oct 2026 17:16:39 -0700 Subject: [PATCH 4/4] recipes(397b v7x): TRAINER_TUNED_FLAGS covers the 256-chip trainer meshes 32x2x8 and 64x1x8 With TRAINER_TPU_SLICE=tpu7x:4x8x8 and TRAINER_MESH_FSDP/EXPERT/CONTEXT 32/2/8 or 64/1/8, TRAINER_TUNED_FLAGS=true adds query_proj, key_proj, value_proj, out_proj and qkv_proj=device, the settings of Rohan Bierneni's 256-chip trainer runs. At 64x1x8, NUM_MOE_TOKEN_CHUNKS defaults to 1; an explicit value wins. One MoE token chunk now turns moe_chunk_pipeline off, and with it the tuned moe_roe_packed_routing_gather: MaxText refuses the chunk pipeline with one chunk, and the packed routing gather without the pipeline. Every existing render is unchanged. --- .../mlperf_submission_launcher_test.py | 76 +++++++++++++++++++ .../examples/recipes/mlperf_397b_256_v7x.sh | 34 ++++++++- .../recipes/mlperf_gbs1024_submission.sh | 4 + 3 files changed, 111 insertions(+), 3 deletions(-) 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 9dfb672d41..13456fddc7 100644 --- a/tests/experimental/examples/deepswe_dist/mlperf_submission_launcher_test.py +++ b/tests/experimental/examples/deepswe_dist/mlperf_submission_launcher_test.py @@ -62,6 +62,12 @@ "moe_routing_maps": "remat", "moe_unsplit_prefused_wi": "false", "moe_wi_gmm_fused_swiglu": "false", + # remat_policy=custom's per-tensor keys (MaxText base.yml). + "query_proj": "remat", + "key_proj": "remat", + "value_proj": "remat", + "out_proj": "remat", + "qkv_proj": "remat", } _TUNED_KEYS = { "cast_grads_after_all_reduce": ("false", "true"), @@ -119,6 +125,15 @@ "TUNIX_LOGPS_SHARDED_HEAD": "fused", "PACK_PAD_POSITION": "start", } +# The 256-chip trainer meshes and the keys TRAINER_TUNED_FLAGS adds at them. +_MESH_32X2X8 = {"TRAINER_TPU_SLICE": "tpu7x:4x8x8", "TRAINER_MESH_CONTEXT": "8"} +_MESH_64X1X8 = { + "TRAINER_TPU_SLICE": "tpu7x:4x8x8", + "TRAINER_MESH_FSDP": "64", + "TRAINER_MESH_EXPERT": "1", + "TRAINER_MESH_CONTEXT": "8", +} +_PROJ_KEYS = ("query_proj", "key_proj", "value_proj", "out_proj", "qkv_proj") def _render(recipe: Path = _RECIPE, **overrides: str) -> str: @@ -585,6 +600,67 @@ def test_tuned_flags_reject_other_values(self): ): _render(TRAINER_TUNED_FLAGS="yes") + def test_tuned_flags_at_the_256_chip_mesh_32x2x8(self): + before = _maxtext_flags(_jobset(self.bench_on, "train")) + flags = _maxtext_flags( + _jobset(_render(**_BENCH_CELL, **_MESH_32X2X8), "train") + ) + self.assertEqual( + _flag_changes(before, flags), + {key: ("remat", "device") for key in _PROJ_KEYS}, + ) + self.assertEqual(flags["num_moe_token_chunks"], "2") + self.assertEqual(flags["moe_chunk_pipeline"], "true") + self.assertEqual(flags["moe_roe_packed_routing_gather"], "true") + + def test_tuned_flags_at_the_256_chip_mesh_64x1x8(self): + before = _maxtext_flags(_jobset(self.bench_on, "train")) + flags = _maxtext_flags( + _jobset(_render(**_BENCH_CELL, **_MESH_64X1X8), "train") + ) + expected = {key: ("remat", "device") for key in _PROJ_KEYS} + # One token chunk, so no chunk pipeline, and so no packed routing gather, + # which MaxText refuses without the pipeline. + expected["num_moe_token_chunks"] = ("2", "1") + expected["moe_chunk_pipeline"] = ("true", "false") + expected["moe_roe_packed_routing_gather"] = ("true", "false") + self.assertEqual(_flag_changes(before, flags), expected) + + def test_tuned_flags_keep_an_explicit_chunk_count_at_64x1x8(self): + flags = _maxtext_flags( + _jobset( + _render(NUM_MOE_TOKEN_CHUNKS="2", **_BENCH_CELL, **_MESH_64X1X8), + "train", + ) + ) + self.assertEqual(flags["num_moe_token_chunks"], "2") + self.assertEqual(flags["moe_chunk_pipeline"], "true") + self.assertEqual(flags["moe_roe_packed_routing_gather"], "true") + for key in _PROJ_KEYS: + self.assertEqual(flags[key], "device", key) + + def test_tuned_flags_false_leaves_the_256_chip_meshes_alone(self): + for mesh in (_MESH_32X2X8, _MESH_64X1X8): + flags = _maxtext_flags( + _jobset(_render(TRAINER_TUNED_FLAGS="false", **mesh), "train") + ) + for key in _PROJ_KEYS: + self.assertNotIn(key, flags, mesh) + self.assertEqual(flags["num_moe_token_chunks"], "4", mesh) + self.assertEqual(flags["moe_chunk_pipeline"], "true", mesh) + + def test_one_token_chunk_turns_the_chunk_pipeline_off(self): + # MaxText refuses moe_chunk_pipeline=true with one token chunk, and + # moe_roe_packed_routing_gather=true without the chunk pipeline. + flags = _maxtext_flags( + _jobset(_render(NUM_MOE_TOKEN_CHUNKS="1", **_BENCH_CELL), "train") + ) + self.assertEqual(flags["num_moe_token_chunks"], "1") + self.assertEqual(flags["moe_chunk_pipeline"], "false") + self.assertEqual(flags["moe_roe_packed_routing_gather"], "false") + for key in _PROJ_KEYS: + self.assertNotIn(key, flags) + def test_tuned_flags_are_off_in_the_other_397b_v7x_recipes(self): for name in ( "mlperf_397b_256_v7x.sh", diff --git a/tunix/experimental/examples/recipes/mlperf_397b_256_v7x.sh b/tunix/experimental/examples/recipes/mlperf_397b_256_v7x.sh index 3955a033ad..5912cd14fb 100755 --- a/tunix/experimental/examples/recipes/mlperf_397b_256_v7x.sh +++ b/tunix/experimental/examples/recipes/mlperf_397b_256_v7x.sh @@ -5,7 +5,9 @@ set -e # MLPerf DeepSWE recipe: Qwen3.5-397B-A17B on TPU v7x # ============================================================================== # - TPU7x dynamic slicing on pod1 (bodaborg-tpu7x-gsc) or pod2 (bodaborg-tpu7x-gsc-elm) -# - Trainer on 128 chips (4x4x8 = 256 devices, FSDP=32, TP=1, EXPERT=2, CP=4) +# - Trainer on 128 chips (4x4x8 = 256 devices, FSDP=32, TP=1, EXPERT=2, CP=4), or on 256 +# chips with TRAINER_TPU_SLICE=tpu7x:4x8x8 and TRAINER_MESH_FSDP/EXPERT/CONTEXT 32/2/8 or +# 64/1/8 (TRAINER_TUNED_FLAGS then adds those meshes' keys, see below) # - Rollout on 256 chips (32 replicas x 8 chips 2x2x2, EP=16, TP=1) # - Sandbox configured for sandbox-np nodepool with workload tolerations # ============================================================================== @@ -162,7 +164,20 @@ export RPC_TIMEOUT_S="${RPC_TIMEOUT_S:-10800}" export REMAT_POLICY="${REMAT_POLICY:-custom}" export RAGGED_BUFFER_FACTOR="${RAGGED_BUFFER_FACTOR:-2.0}" export REMAT_CONTEXT="${REMAT_CONTEXT:-remat}" +# TRAINER_TUNED_FLAGS=true at the 256-chip trainer mesh 64x1x8 runs one MoE token chunk (see +# the end of the TRAINER_TUNED_FLAGS block); an explicit NUM_MOE_TOKEN_CHUNKS wins. +if [[ "${TRAINER_TUNED_FLAGS:-false}" == "true" \ + && "${TRAINER_MESH_FSDP}x${TRAINER_MESH_EXPERT}x${TRAINER_MESH_CONTEXT}" == "64x1x8" ]]; then + export NUM_MOE_TOKEN_CHUNKS="${NUM_MOE_TOKEN_CHUNKS:-1}" +fi export NUM_MOE_TOKEN_CHUNKS="${NUM_MOE_TOKEN_CHUNKS:-2}" +# MaxText refuses moe_chunk_pipeline=true with one MoE token chunk, so one chunk turns the +# pipeline off. +if [[ "${NUM_MOE_TOKEN_CHUNKS}" == "1" ]]; then + _moe_chunk_pipeline=false +else + _moe_chunk_pipeline=true +fi export MAXTEXT_EXTRA_FLAGS="${MAXTEXT_EXTRA_FLAGS:-custom_mesh_and_rule=cp-as-ep \ use_gdn_kernel=true gdn_cp_mode=head gdn_chunk_size=64 \ decoder_layer_input=offload context=${REMAT_CONTEXT} gdn=remat gdn_conv=remat gdn_states=remat \ @@ -176,7 +191,7 @@ wo_tile_dlhs_batch_seq=256 wo_tile_dlhs_embed_dim=4096 wo_tile_dlhs_mlp_dim=1024 wo_tile_drhs_batch_seq=256 wo_tile_drhs_embed_dim=4096 wo_tile_drhs_mlp_dim=1024 \ use_ring_of_experts=true num_moe_token_chunks=${NUM_MOE_TOKEN_CHUNKS} moe_chunk_barrier=false \ ring_of_experts_local_routing=true moe_expert_weight_prefetch=forward \ -moe_chunk_pipeline=true ring_of_experts_row_major_reduce_scatter=true \ +moe_chunk_pipeline=${_moe_chunk_pipeline} ring_of_experts_row_major_reduce_scatter=true \ use_ragged_sort=true use_custom_sort_vjp=false ragged_buffer_factor=${RAGGED_BUFFER_FACTOR} \ use_tokamax_splash=true use_splash_scheduler=true \ sa_block_q=1024 sa_block_kv=4096 sa_block_kv_compute=512 \ @@ -328,8 +343,21 @@ fi # requirements/maxtext_requirements.txt (e526ea02) has none of these nine keys, and # MaxText refuses a key it does not define when the trainer starts, so # TRAINER_TUNED_FLAGS=true needs that pin moved to a MaxText with #5758. +# The 256-chip trainer (TRAINER_TPU_SLICE=tpu7x:4x8x8, 512 devices), from Rohan Bierneni's +# 256-chip trainer runs (2026-10-15; their step times are not recorded here): +# at 32x2x8 (TRAINER_MESH_FSDP=32, TRAINER_MESH_EXPERT=2, TRAINER_MESH_CONTEXT=8) and +# 64x1x8 (64, 1, 8), query_proj, key_proj, value_proj, out_proj and qkv_proj=device keep +# those projections' outputs in HBM for the backward instead of recomputing them. MaxText +# also names the GDN layers' fused qkvz projection qkv_proj, so qkv_proj covers them too. +# At 64x1x8, NUM_MOE_TOKEN_CHUNKS defaults to 1 (above), so moe_chunk_pipeline is false, +# and so is moe_roe_packed_routing_gather: MaxText refuses it without the pipeline. if [[ "${TRAINER_TUNED_FLAGS:-false}" == "true" ]]; then - export MAXTEXT_EXTRA_FLAGS="${MAXTEXT_EXTRA_FLAGS} context_parallel_strategy=all_gather sa_fuse_reciprocal=true gdn_cp_mode=seq gdn_cp_state=device gdn_cp_matmul_precision=high cast_grads_after_all_reduce=true wi_tile_dlhs_mlp_dim=2048 decoder_layer_input=device moe_routing_maps=device moe_gmm_v2_dlhs_transpose_rhs=true moe_routing_index_math=count moe_unsplit_prefused_wi=true moe_roe_packed_routing_gather=true moe_router_grad_from_wo_input=true moe_wi_gmm_fused_swiglu=true gdn_conv1d_input_in_qkv_dtype=true gdn_conv1d_dense_token_stack=true wi_tile_fwd_batch_seq=512 wi_tile_drhs_batch_seq=512 wo_tile_fwd_batch_seq=512 wo_tile_dlhs_batch_seq=512 wo_tile_drhs_batch_seq=512" + export MAXTEXT_EXTRA_FLAGS="${MAXTEXT_EXTRA_FLAGS} context_parallel_strategy=all_gather sa_fuse_reciprocal=true gdn_cp_mode=seq gdn_cp_state=device gdn_cp_matmul_precision=high cast_grads_after_all_reduce=true wi_tile_dlhs_mlp_dim=2048 decoder_layer_input=device moe_routing_maps=device moe_gmm_v2_dlhs_transpose_rhs=true moe_routing_index_math=count moe_unsplit_prefused_wi=true moe_roe_packed_routing_gather=${_moe_chunk_pipeline} moe_router_grad_from_wo_input=true moe_wi_gmm_fused_swiglu=true gdn_conv1d_input_in_qkv_dtype=true gdn_conv1d_dense_token_stack=true wi_tile_fwd_batch_seq=512 wi_tile_drhs_batch_seq=512 wo_tile_fwd_batch_seq=512 wo_tile_dlhs_batch_seq=512 wo_tile_drhs_batch_seq=512" + case "${TRAINER_MESH_FSDP}x${TRAINER_MESH_EXPERT}x${TRAINER_MESH_CONTEXT}" in + 32x2x8|64x1x8) + export MAXTEXT_EXTRA_FLAGS="${MAXTEXT_EXTRA_FLAGS} query_proj=device key_proj=device value_proj=device out_proj=device qkv_proj=device" + ;; + esac fi export DEBUG=${DEBUG:-0} diff --git a/tunix/experimental/examples/recipes/mlperf_gbs1024_submission.sh b/tunix/experimental/examples/recipes/mlperf_gbs1024_submission.sh index 729e602680..2d219375de 100644 --- a/tunix/experimental/examples/recipes/mlperf_gbs1024_submission.sh +++ b/tunix/experimental/examples/recipes/mlperf_gbs1024_submission.sh @@ -112,6 +112,10 @@ case "${DROPLESS}" in if [[ "${RETRY_DROPLESS_FINAL_MICRO_BATCH:-false}" != "true" \ && -z "${FIRST_PHASE_RAGGED_BUFFER_FACTOR:-}" ]]; then export REMAT_CONTEXT="${REMAT_CONTEXT:-device}" + # One token chunk at the 256-chip trainer mesh 64x1x8 (see mlperf_397b_256_v7x.sh). + if [[ "${TRAINER_MESH_FSDP:-}x${TRAINER_MESH_EXPERT:-}x${TRAINER_MESH_CONTEXT:-}" == "64x1x8" ]]; then + export NUM_MOE_TOKEN_CHUNKS="${NUM_MOE_TOKEN_CHUNKS:-1}" + fi export NUM_MOE_TOKEN_CHUNKS="${NUM_MOE_TOKEN_CHUNKS:-2}" fi fi