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 b387a602d..13456fddc 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,69 @@ "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", + # 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"), "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,8 +121,19 @@ "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", +} +# 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: @@ -175,7 +235,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,19 +429,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_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( + 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_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( @@ -438,6 +550,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") @@ -486,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", @@ -498,6 +673,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 ( @@ -510,6 +687,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/tests/experimental/orchestrator/batch_assembly_test.py b/tests/experimental/orchestrator/batch_assembly_test.py index 4b710170a..856359643 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 04a95ce4d..7fa0720bb 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 cf4f51e58..63699bcb7 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/tests/rl/sharded_logps_head_test.py b/tests/rl/sharded_logps_head_test.py index 66b99242c..03eda4750 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/experimental/examples/deepswe_dist/k8s_launcher.sh b/tunix/experimental/examples/deepswe_dist/k8s_launcher.sh index 5989d9841..e98521722 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 f975b8ea1..2033afac0 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/examples/recipes/mlperf_397b_256_v7x.sh b/tunix/experimental/examples/recipes/mlperf_397b_256_v7x.sh index 8c7cbcc33..5912cd14f 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 \ @@ -222,17 +237,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 +263,101 @@ 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. +# 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 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=${_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 e3e6f10a1..2d219375d 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 @@ -101,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 @@ -125,8 +140,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 +152,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}}" diff --git a/tunix/experimental/orchestrator/batch_assembly.py b/tunix/experimental/orchestrator/batch_assembly.py index d29a011bb..c0df0e2f4 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/common.py b/tunix/rl/common.py index 2c269136e..fd80be744 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 diff --git a/tunix/rl/packing.py b/tunix/rl/packing.py index 67babc136..e6cd8bbc8 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