Qwen3.5-397B RL trainer on 128 v7x: pad-first sequence packing, fused sharded log-prob head VJP, and the tuned trainer recipe

tests/experimental/examples/deepswe_dist/mlperf_submission_launcher_test.py

@@ -51,29 +51,89 @@
     "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",
     "RETRY_WHEN_TOKENS_DROPPED": "true",
     "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")

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(