diff --git a/tests/experimental/examples/deepswe_dist/mlperf_submission_launcher_test.py b/tests/experimental/examples/deepswe_dist/mlperf_submission_launcher_test.py index 945579386e..ae294f35b0 100644 --- a/tests/experimental/examples/deepswe_dist/mlperf_submission_launcher_test.py +++ b/tests/experimental/examples/deepswe_dist/mlperf_submission_launcher_test.py @@ -16,7 +16,9 @@ Each test renders mlperf_gbs1024_submission.sh with DRY_RUN=true and checks the rendered JobSets, so the defaults, the DROPLESS switch and the learning-rate -rule are tested through the same path a launch takes. +rule are tested through the same path a launch takes. Two tests render the other +397B v7x recipes, which leave TRAINER_TUNED_FLAGS, the sharded log-prob head and +the staged step boundary off. """ import json @@ -40,10 +42,44 @@ ) _PREFIX = "launcher-test" _HBM_CAP = "--xla_tpu_max_hbm_size_mib=84992" - - -def _render(**overrides: str) -> str: - """Returns the DRY_RUN render of the submission launcher. +_SC_FLAG = "--xla_tpu_enable_offloading_copy_to_sparsecore=false" +# What TRAINER_TUNED_FLAGS changes in the trainer's MaxText flags at the bench +# cell: key -> (value with it off, value with it on). A repeated key takes its +# last value, as MaxText's command line parsing does, and an absent key its +# MaxText default (_MAXTEXT_DEFAULTS). +_MAXTEXT_DEFAULTS = { + "cast_grads_after_all_reduce": "false", + "gdn_cp_matmul_precision": "highest", + "gdn_cp_state": "remat", + "moe_route": "remat", +} +_TUNED_KEYS = { + "cast_grads_after_all_reduce": ("false", "true"), + "context_parallel_strategy": ("ring", "all_gather"), + "gdn_cp_matmul_precision": ("highest", "high"), + "gdn_cp_mode": ("head", "seq"), + "gdn_cp_state": ("remat", "device"), + "moe_route": ("remat", "device"), + "num_moe_token_chunks": ("4", "2"), + "sa_fuse_reciprocal": ("false", "true"), + "sa_use_base2_exp": ("true", "false"), +} +# 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. +_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", +} + + +def _render(recipe: Path = _RECIPE, **overrides: str) -> str: + """Returns the DRY_RUN render of a launcher (default: the submission one). The environment is scrubbed (an inherited DROPLESS or LEARNING_RATE would change the result), HOME is a temporary directory and KUBECONFIG an empty @@ -66,7 +102,7 @@ def _render(**overrides: str) -> str: } env.update(overrides) result = subprocess.run( - ["bash", str(_RECIPE), "start"], + ["bash", str(recipe), "start"], env=env, cwd=home, capture_output=True, @@ -97,6 +133,29 @@ def _vllm_config(rollout: str) -> dict: return json.loads(re.search(r"--vllm_config_json='([^']*)'", rollout).group(1)) +def _value(flags: dict[str, str], key: str) -> str | None: + """Returns the value MaxText resolves for `key` from these flags.""" + return flags.get(key, _MAXTEXT_DEFAULTS.get(key)) + + +def _flag_changes( + before: dict[str, str], after: dict[str, str] +) -> dict[str, tuple[str | None, str | None]]: + """Returns key -> (before, after) for each MaxText flag that differs.""" + return { + key: (_value(before, key), _value(after, key)) + for key in before.keys() | after.keys() + if _value(before, key) != _value(after, key) + } + + +def _without_tuned_flags(trainer: str) -> str: + """Returns the trainer JobSet without the libtpu flag and MaxText flags.""" + trainer = re.sub(rf'\n *- "{re.escape(_SC_FLAG)}"', "", trainer) # proxy arg + trainer = trainer.replace(f" {_SC_FLAG}", "") # trainer and worker env + return re.sub(r'MAXTEXT_EXTRA_FLAGS="[^"]*"', "MAXTEXT_EXTRA_FLAGS=", trainer) + + @absltest.skipIf(shutil.which("python3") is None, "the recipes need python3") class SubmissionLauncherTest(absltest.TestCase): @@ -104,6 +163,9 @@ class SubmissionLauncherTest(absltest.TestCase): def setUpClass(cls): super().setUpClass() cls.default = _render() + cls.default_off = _render(TRAINER_TUNED_FLAGS="false") + cls.bench_on = _render(**_BENCH_CELL) + cls.bench_off = _render(TRAINER_TUNED_FLAGS="false", **_BENCH_CELL) def test_dropless_is_the_default(self): trainer = _jobset(self.default, "train") @@ -135,8 +197,14 @@ def test_finite_buffer_factor_needs_the_retry_to_stay_dropless(self): _render(RAGGED_BUFFER_FACTOR="2.75") def test_finite_buffer_factor_with_retry_keeps_the_dropless_layout(self): + # With TRAINER_TUNED_FLAGS on, this layout runs two chunks instead (see the + # tests of that switch below). trainer = _jobset( - _render(RAGGED_BUFFER_FACTOR="2.75", RETRY_WHEN_TOKENS_DROPPED="true"), + _render( + RAGGED_BUFFER_FACTOR="2.75", + RETRY_WHEN_TOKENS_DROPPED="true", + TRAINER_TUNED_FLAGS="false", + ), "train", ) flags = _maxtext_flags(trainer) @@ -214,6 +282,48 @@ def test_start_up_flags_are_emitted(self): self.assertIn("--trainer_warmup_routed_experts", orch) self.assertIn("--no-sandbox_wait_initial", orch) + def test_the_staged_step_boundary_is_on_by_default(self): + orch = _jobset(self.default, "orch") + self.assertIn("--pipeline_step_boundary", orch) + self.assertIn("--stage_next_microbatch", orch) + self.assertNotIn("--no-pipeline_step_boundary", orch) + self.assertNotIn("--no-stage_next_microbatch", orch) + + def test_the_pipelined_step_boundary_off_turns_staging_off(self): + orch = _jobset(_render(PIPELINE_STEP_BOUNDARY="false"), "orch") + self.assertIn("--no-pipeline_step_boundary", orch) + self.assertIn("--no-stage_next_microbatch", orch) + + def test_the_pipelined_step_boundary_flags_reach_the_orchestrator(self): + orch = _jobset( + _render(PIPELINE_STEP_BOUNDARY="true", STAGE_NEXT_MICROBATCH="1"), + "orch", + ) + self.assertIn("--pipeline_step_boundary", orch) + self.assertIn("--stage_next_microbatch", orch) + self.assertNotIn("--no-pipeline_step_boundary", orch) + self.assertNotIn("--no-stage_next_microbatch", orch) + + def test_staging_without_the_pipelined_boundary_is_refused(self): + with self.assertRaisesRegex( + AssertionError, "STAGE_NEXT_MICROBATCH=true requires" + ): + _render(PIPELINE_STEP_BOUNDARY="false", STAGE_NEXT_MICROBATCH="true") + + 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" + ) + for suffix in ("orch", "roll-0"): + self.assertNotIn( + "TUNIX_LOGPS_SHARDED_HEAD", _jobset(self.default, suffix), suffix + ) + + 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) + def test_start_up_flags_can_be_turned_off(self): orch = _jobset( _render( @@ -245,6 +355,100 @@ def test_rollout_attention_data_parallelism_can_be_restored(self): self.assertEqual(cfg["max_num_seqs"], 16) self.assertNotIn("kv_tp_size", _maxtext_flags(_jobset(render, "train"))) + def test_tuned_flags_set_exactly_the_nine_keys_at_the_bench_cell(self): + before = _maxtext_flags(_jobset(self.bench_off, "train")) + after = _maxtext_flags(_jobset(self.bench_on, "train")) + self.assertEqual(_flag_changes(before, after), _TUNED_KEYS) + + def test_tuned_flags_change_only_the_trainer_flags(self): + off_docs = self.bench_off.split("\n---\n") + on_docs = self.bench_on.split("\n---\n") + self.assertLen(on_docs, len(off_docs)) + changed = [on for on, off in zip(on_docs, off_docs) if on != off] + self.assertEqual(changed, [_jobset(self.bench_on, "train")]) + off = _jobset(self.bench_off, "train") + on = _jobset(self.bench_on, "train") + self.assertNotIn("offloading_copy_to_sparsecore", off) + # The trainer env, the Pathways worker env and the Pathways proxy args. + self.assertEqual(on.count(_SC_FLAG), 3) + self.assertEqual(_without_tuned_flags(on), _without_tuned_flags(off)) + + def test_tuned_flags_false_turns_the_block_off(self): + trainer = _jobset(self.default_off, "train") + flags = _maxtext_flags(trainer) + self.assertEqual( + {key: _value(flags, key) for key in _TUNED_KEYS}, + {key: off for key, (off, _) in _TUNED_KEYS.items()}, + ) + self.assertNotIn("offloading_copy_to_sparsecore", trainer) + + def test_tuned_flags_keep_four_chunks_without_a_buffer_factor(self): + # The fully dropless layout (the default) was not measured with two chunks. + trainer = _jobset(self.default, "train") + before = _maxtext_flags(_jobset(self.default_off, "train")) + expected = dict(_TUNED_KEYS) + del expected["num_moe_token_chunks"] + self.assertEqual(_flag_changes(before, _maxtext_flags(trainer)), expected) + self.assertEqual(_maxtext_flags(trainer)["num_moe_token_chunks"], "4") + self.assertEqual(trainer.count(_SC_FLAG), 3) + + def test_tuned_flags_with_the_capacity_factor_trainer(self): + # DROPLESS=false runs two chunks with or without the switch. + before = _maxtext_flags( + _jobset(_render(DROPLESS="false", TRAINER_TUNED_FLAGS="false"), "train") + ) + trainer = _jobset(_render(DROPLESS="false"), "train") + expected = dict(_TUNED_KEYS) + del expected["num_moe_token_chunks"] + self.assertEqual(_flag_changes(before, _maxtext_flags(trainer)), expected) + self.assertEqual(_maxtext_flags(trainer)["num_moe_token_chunks"], "2") + self.assertEqual(trainer.count(_SC_FLAG), 3) + + def test_tuned_flags_keep_an_explicit_chunk_count(self): + flags = _maxtext_flags( + _jobset(_render(NUM_MOE_TOKEN_CHUNKS="4", **_BENCH_CELL), "train") + ) + self.assertEqual(flags["num_moe_token_chunks"], "4") + self.assertEqual(flags["gdn_cp_mode"], "seq") + + def test_tuned_flags_keep_an_explicit_sparsecore_copy_flag(self): + explicit = "--xla_tpu_enable_offloading_copy_to_sparsecore=true" + libtpu_args = f"{explicit} --xla_tpu_scoped_vmem_limit_kib=65472" + trainer = _jobset(_render(TRAINER_LIBTPU_INIT_ARGS=libtpu_args), "train") + self.assertEqual(trainer.count(explicit), 3) + self.assertNotIn(_SC_FLAG, trainer) + + def test_tuned_flags_reject_other_values(self): + with self.assertRaisesRegex( + AssertionError, "TRAINER_TUNED_FLAGS must be true or false" + ): + _render(TRAINER_TUNED_FLAGS="yes") + + def test_tuned_flags_are_off_in_the_other_397b_v7x_recipes(self): + for name in ( + "mlperf_397b_256_v7x.sh", + "mlperf_397b_1024_v7x.sh", + "mlperf_gbs512.sh", + ): + render = _render(_RECIPE.with_name(name)) + self.assertNotIn("offloading_copy_to_sparsecore", render, name) + flags = _maxtext_flags(_jobset(render, "train")) + 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}") + + def test_head_and_staging_are_off_in_the_other_397b_v7x_recipes(self): + for name in ( + "mlperf_397b_256_v7x.sh", + "mlperf_397b_1024_v7x.sh", + "mlperf_gbs512.sh", + ): + render = _render(_RECIPE.with_name(name)) + self.assertNotIn("TUNIX_LOGPS_SHARDED_HEAD", render, name) + orch = _jobset(render, "orch") + self.assertIn("--no-pipeline_step_boundary", orch, name) + self.assertIn("--no-stage_next_microbatch", orch, name) + if __name__ == "__main__": absltest.main() diff --git a/tests/experimental/orchestrator/distributed_rl_engine_deferred_update_test.py b/tests/experimental/orchestrator/distributed_rl_engine_deferred_update_test.py new file mode 100644 index 0000000000..ef82e78b25 --- /dev/null +++ b/tests/experimental/orchestrator/distributed_rl_engine_deferred_update_test.py @@ -0,0 +1,210 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for DistributedRLEngine's deferred update. + +`train_step(defer_update=True)` sends a step's last micro-batch with its +update left to `apply_deferred_update`, so the next step's first micro-batch +can be staged in between. +""" + +import asyncio +from typing import Any + +from absl.testing import absltest +import distributed_rl_engine_staging_test as st +from tunix.experimental.common import datatypes +from tunix.experimental.orchestrator import distributed_rl_engine +from tunix.experimental.worker import remote_execution +from tunix.experimental.worker import trainer_worker + + +class _DeferringHandle(st._RecordingHandle): + """Answers a deferred train call as a TrainerWorker that deferred it, and + `apply_deferred_update` as one that then ran the update.""" + + def __init__(self, ignores_defer: bool = False): + super().__init__(fused=True) + self._ignores_defer = ignores_defer + + async def asubmit(self, method_name=None, *args, **kwargs): + if method_name == "apply_deferred_update": + self.calls.append((method_name, kwargs)) + return datatypes.Response( + request_id="r", + metadata={ + "updated": True, + "train_step": 8, + "update_ready": True, + "device_wait_s": 4.9, + }, + ) + if kwargs.get("defer_update") and not self._ignores_defer: + self.calls.append((method_name, kwargs)) + return datatypes.Response( + request_id="r", metadata={"queued": True, "update_deferred": True} + ) + return await super().asubmit(method_name, *args, **kwargs) + + +class _Trainer: + """Records what reaches the trainer behind a real TrainerWorker.""" + + def __init__(self): + self.calls: list[str] = [] + self.policy_version = 0 + + def fwd_bwd(self, payload, **kwargs): + del payload + self.calls.append( + f"fwd_bwd(apply_optimizer={kwargs.get('apply_optimizer')})" + ) + + def update(self, **kwargs): + del kwargs + self.calls.append("update") + return 3 + + def wait_for_device(self): + self.calls.append("wait_for_device") + return 2.5 + + def close(self): + pass + + +def _real_engine(trainer): + worker = trainer_worker.TrainerWorker(trainer_factory=lambda: trainer) + worker.initialize() + handle = remote_execution.InProcessActorHandle( + remote_execution.InProcessRemoteExecutionServer(worker) + ) + return st._engine(handle), worker + + +class DeferredTrainStepTest(absltest.TestCase): + + def test_defer_sends_the_flag_and_does_not_update(self): + handle = _DeferringHandle() + result = asyncio.run( + st._engine(handle).train_step( + st._payload(), apply_optimizer=True, defer_update=True + ) + ) + self.assertEqual([m for m, _ in handle.calls], ["fwd_bwd"]) + _, kwargs = handle.calls[0] + self.assertTrue(kwargs["defer_update"]) + self.assertTrue(kwargs["apply_optimizer"]) + self.assertTrue(result["update_deferred"]) + self.assertFalse(result["updated"]) + self.assertIsNone(result["train_step"]) + self.assertFalse(result["update_ready"]) + + def test_without_defer_the_flag_is_not_sent(self): + handle = _DeferringHandle() + result = asyncio.run( + st._engine(handle).train_step(st._payload(), apply_optimizer=True) + ) + self.assertNotIn("defer_update", handle.calls[0][1]) + self.assertTrue(result["updated"]) + self.assertNotIn("update_deferred", result) + + def test_defer_needs_apply_optimizer(self): + handle = _DeferringHandle() + with self.assertRaisesRegex(ValueError, "needs apply_optimizer=True"): + asyncio.run( + st._engine(handle).train_step( + st._payload(), apply_optimizer=False, defer_update=True + ) + ) + self.assertEmpty(handle.calls) + + def test_a_worker_that_updated_anyway_is_taken_as_fused(self): + # One update, never two: the reply says `updated`, not `update_deferred`. + handle = _DeferringHandle(ignores_defer=True) + result = asyncio.run( + st._engine(handle).train_step( + st._payload(), apply_optimizer=True, defer_update=True + ) + ) + self.assertEqual([m for m, _ in handle.calls], ["fwd_bwd"]) + self.assertTrue(result["updated"]) + self.assertEqual(result["train_step"], 7) + self.assertNotIn("update_deferred", result) + + def test_apply_deferred_update_returns_the_workers_report(self): + handle = _DeferringHandle() + engine = st._engine(handle) + result = asyncio.run(engine.apply_deferred_update()) + self.assertEqual([m for m, _ in handle.calls], ["apply_deferred_update"]) + self.assertEqual( + result, + { + "updated": True, + "train_step": 8, + "update_ready": True, + "device_wait_s": 4.9, + }, + ) + self.assertTrue(engine._actor_weights_dirty) + + def test_apply_deferred_update_without_a_trainer_is_an_error(self): + engine = distributed_rl_engine.DistributedRLEngine( + rollout_workers=[], trainer_workers={} + ) + with self.assertRaisesRegex(ValueError, "No trainer worker"): + asyncio.run(engine.apply_deferred_update()) + + def test_through_a_real_worker(self): + trainer = _Trainer() + engine, worker = _real_engine(trainer) + deferred = asyncio.run( + engine.train_step( + st._payload(), apply_optimizer=True, defer_update=True + ) + ) + self.assertTrue(deferred["update_deferred"]) + self.assertEqual(trainer.calls, ["fwd_bwd(apply_optimizer=True)"]) + staged = asyncio.run(engine.stage_train_step(st._payload())) + self.assertIsNone(staged.device_wait_s) + update = asyncio.run(engine.apply_deferred_update()) + self.assertEqual( + trainer.calls, + ["fwd_bwd(apply_optimizer=True)", "wait_for_device", "update"], + ) + self.assertEqual( + update, + { + "updated": True, + "train_step": 3, + "update_ready": False, + "device_wait_s": 2.5, + }, + ) + asyncio.run( + engine.train_step( + None, apply_optimizer=False, staged_request_id=staged.request_id + ) + ) + self.assertEqual(trainer.calls[-1], "fwd_bwd(apply_optimizer=False)") + self.assertIsNone(worker._deferred_update) + + def test_staging_support_reports_it(self): + engine, _ = _real_engine(_Trainer()) + support = asyncio.run(engine.trainer_staging_support()) + self.assertTrue(support["stage_before_update"]) + + +if __name__ == "__main__": + absltest.main() diff --git a/tests/experimental/orchestrator/distributed_rl_engine_staging_test.py b/tests/experimental/orchestrator/distributed_rl_engine_staging_test.py new file mode 100644 index 0000000000..ef84bc4710 --- /dev/null +++ b/tests/experimental/orchestrator/distributed_rl_engine_staging_test.py @@ -0,0 +1,453 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for DistributedRLEngine.stage_train_step and staged train_step.""" + +import asyncio +import dataclasses +from typing import Any +from unittest import mock + +from absl.testing import absltest +import jax.numpy as jnp +import numpy as np +from tunix.experimental.common import datatypes +from tunix.experimental.orchestrator import distributed_rl_engine +from tunix.experimental.worker import remote_execution +from tunix.experimental.worker import trainer_worker + + +class _RecordingHandle(remote_execution.ActorHandle): + """Records every RPC; answers like a TrainerWorker would. + + `fused` answers a train call with `apply_optimizer=True` as a worker that + ran the update inside it (`updated`, `train_step` 7 and `update_ready`); + otherwise it answers as a worker without the fused update, so the engine + sends `update` itself. + """ + + def __init__( + self, + stage_metadata: dict[str, Any] | None = None, + fused: bool = False, + ): + self.calls: list[tuple[str, dict[str, Any]]] = [] + self._stage_metadata = stage_metadata or {"staged": True} + self._fused = fused + + async def asubmit(self, method_name=None, *args, **kwargs): + del args + self.calls.append((method_name, kwargs)) + if method_name == "update": + return 5 + if method_name == "stage_fwd_bwd": + return datatypes.Response(request_id="r", metadata=self._stage_metadata) + if self._fused and kwargs.get("apply_optimizer"): + return datatypes.Response( + request_id="r", + metadata={ + "queued": True, + "updated": True, + "train_step": 7, + "update_ready": True, + }, + ) + return datatypes.Response(request_id="r", metadata={"queued": True}) + + def submit(self, method_name=None, *args, **kwargs): + raise NotImplementedError + + async def dispatch_task(self, request_id=None, method_name=None, *a, **k): + raise NotImplementedError + + async def poll_responses(self, timeout_s=0.0): + raise NotImplementedError + + +def _payload() -> datatypes.RLTrainerPayload: + return datatypes.RLTrainerPayload( + prompt_ids=np.array([[1, 2]], dtype=np.int32), + prompt_mask=np.ones((1, 2), dtype=np.float32), + completion_ids=np.array([[3, 4]], dtype=np.int32), + completion_mask=np.ones((1, 2), dtype=np.float32), + advantages=np.array([1.0], dtype=np.float32), + metadata={"traj": "t0"}, + ) + + +def _engine(handle) -> distributed_rl_engine.DistributedRLEngine: + return distributed_rl_engine.DistributedRLEngine( + rollout_workers=[], trainer_workers={datatypes.Role.ACTOR: handle} + ) + + +class StagedTrainStepTest(absltest.TestCase): + + def test_stage_sends_the_payload_once_and_returns_its_id(self): + handle = _RecordingHandle() + payload = _payload() + staged = asyncio.run(_engine(handle).stage_train_step(payload)) + self.assertLen(handle.calls, 1) + method, kwargs = handle.calls[0] + self.assertEqual(method, "stage_fwd_bwd") + self.assertIs(kwargs["request"].payload, payload) + self.assertEqual(kwargs["request"].request_id, staged.request_id) + self.assertEqual(kwargs["request"].metadata, {"traj": "t0"}) + # The trainer said nothing about its device. + self.assertIsNone(staged.device_wait_s) + + def test_stage_reports_the_trainers_device_wait(self): + handle = _RecordingHandle({"staged": True, "device_wait_s": 1.25}) + staged = asyncio.run(_engine(handle).stage_train_step(_payload())) + self.assertEqual(staged.device_wait_s, 1.25) + self.assertIsInstance(staged, distributed_rl_engine.StagedTrainStep) + self.assertFalse(staged.update_ready) + + def test_stage_reports_the_trainers_update_confirmation(self): + handle = _RecordingHandle( + {"staged": True, "device_wait_s": 1.25, "update_ready": True} + ) + staged = asyncio.run(_engine(handle).stage_train_step(_payload())) + self.assertTrue(staged.update_ready) + + def test_staged_train_step_without_a_fused_worker_sends_the_update(self): + # A worker that did not run the update inside the call (no `updated`): + # the engine sends `update` itself, as for an unstaged call. + handle = _RecordingHandle() + engine = _engine(handle) + result = asyncio.run( + engine.train_step( + None, + accumulate_gradients=True, + apply_optimizer=True, + staged_request_id="train_req_x", + ) + ) + self.assertEqual( + handle.calls, + [ + ( + "fwd_bwd_staged", + { + "staged_request_id": "train_req_x", + "skip_jit": False, + "apply_optimizer": True, + }, + ), + ("update", {}), + ], + ) + self.assertTrue(result["updated"]) + self.assertEqual(result["train_step"], 5) + self.assertFalse(result["update_ready"]) + self.assertTrue(engine._actor_weights_dirty) + + def test_staged_train_step_takes_the_fused_update(self): + # A worker that ran the update inside the call: no second update RPC, + # and its train_step and update_ready are passed on, as unstaged. + handle = _RecordingHandle(fused=True) + result = asyncio.run( + _engine(handle).train_step( + None, apply_optimizer=True, staged_request_id="train_req_x" + ) + ) + self.assertEqual([c[0] for c in handle.calls], ["fwd_bwd_staged"]) + self.assertEqual(result["train_step"], 7) + self.assertTrue(result["update_ready"]) + + def test_staged_and_unstaged_results_have_the_same_shape(self): + for fused in (False, True): + staged = asyncio.run( + _engine(_RecordingHandle(fused=fused)).train_step( + None, apply_optimizer=True, staged_request_id="train_req_x" + ) + ) + unstaged = asyncio.run( + _engine(_RecordingHandle(fused=fused)).train_step( + _payload(), apply_optimizer=True + ) + ) + self.assertEqual(sorted(staged), sorted(unstaged)) + for key in ("updated", "train_step", "accumulated", "update_ready"): + self.assertEqual(staged[key], unstaged[key], (fused, key)) + + def test_staged_train_step_without_the_update(self): + handle = _RecordingHandle() + asyncio.run( + _engine(handle).train_step( + None, apply_optimizer=False, staged_request_id="train_req_x" + ) + ) + self.assertEqual([c[0] for c in handle.calls], ["fwd_bwd_staged"]) + # The trainer still learns it is not the step's last micro-batch. + self.assertIs(handle.calls[0][1]["apply_optimizer"], False) + + def test_a_payload_and_a_staged_id_together_are_refused(self): + handle = _RecordingHandle() + with self.assertRaisesRegex(ValueError, "not both"): + asyncio.run(_engine(handle).train_step(_payload(), staged_request_id="x")) + self.assertEmpty(handle.calls) + + def test_unstaged_train_step_is_unchanged(self): + handle = _RecordingHandle() + payload = _payload() + asyncio.run(_engine(handle).train_step(payload, apply_optimizer=True)) + self.assertEqual([c[0] for c in handle.calls], ["fwd_bwd", "update"]) + self.assertEqual( + sorted(handle.calls[0][1]), ["apply_optimizer", "request", "skip_jit"] + ) + self.assertIs(handle.calls[0][1]["request"].payload, payload) + + def test_wait_for_update_reaches_the_worker_only_when_given(self): + for staged_request_id, payload in ((None, _payload()), ("x", None)): + default = _RecordingHandle(fused=True) + asyncio.run( + _engine(default).train_step( + payload, + apply_optimizer=True, + staged_request_id=staged_request_id, + ) + ) + self.assertNotIn("wait_for_update", default.calls[0][1]) + at_dispatch = _RecordingHandle(fused=True) + asyncio.run( + _engine(at_dispatch).train_step( + payload, + apply_optimizer=True, + staged_request_id=staged_request_id, + wait_for_update=False, + ) + ) + self.assertIs(at_dispatch.calls[0][1]["wait_for_update"], False) + + def test_stage_then_train_reaches_a_real_trainer_worker(self): + + class _Trainer: + policy_version = 0 + + def __init__(self): + self.trained = [] + self.fwd_bwd_kwargs = [] + self.updates = 0 + + def fwd_bwd(self, payload, **kwargs): + self.trained.append(payload) + self.fwd_bwd_kwargs.append(kwargs) + + def update(self, **kwargs): + self.updates += 1 + return 1 + + def close(self): + pass + + trainer = _Trainer() + worker = trainer_worker.TrainerWorker(trainer_factory=lambda: trainer) + worker.initialize() + handle = remote_execution.InProcessActorHandle( + remote_execution.InProcessRemoteExecutionServer(worker) + ) + engine = _engine(handle) + payload = _payload() + + async def _run(): + staged = await engine.stage_train_step(payload) + self.assertEmpty(trainer.trained) + self.assertIsNone(staged.device_wait_s) # no wait_for_device + self.assertFalse(staged.update_ready) # no last_update_token + return await engine.train_step( + None, apply_optimizer=True, staged_request_id=staged.request_id + ) + + result = asyncio.run(_run()) + self.assertEqual(result["train_step"], 1) + self.assertLen(trainer.trained, 1) + self.assertIs(trainer.trained[0], payload) + # The worker ran the update inside the staged call, exactly once, and the + # trainer saw the step's-last-micro-batch flag as for an unstaged call. + self.assertEqual(trainer.updates, 1) + self.assertEqual(trainer.fwd_bwd_kwargs, [{"apply_optimizer": True}]) + + def test_fused_call_at_dispatch_then_stage_confirms_the_update(self): + """`wait_for_update=False` moves the update wait from the call to the stage. + + The trainer publishes an update token. The step's last call, sent with + `wait_for_update=False`, must return without waiting on it and without + claiming `update_ready`; the stage right after must wait on it and claim + it. Control: the same call with the default waits and claims, as before. + """ + + class _TokenTrainer: + policy_version = 0 + + def __init__(self): + self.last_update_token = None + + def fwd_bwd(self, payload, **kwargs): + del payload, kwargs + + def update(self, **kwargs): + del kwargs + self.last_update_token = jnp.array(2.0) + return 3 + + def close(self): + pass + + trainer = _TokenTrainer() + engine = _real_engine(trainer_worker.TrainerWorker, trainer) + waited = [] + real_block_until_ready = trainer_worker.jax.block_until_ready + + def _recording_block_until_ready(x): + waited.append(x) + return real_block_until_ready(x) + + async def _run(): + at_dispatch = await engine.train_step( + _payload(), apply_optimizer=True, wait_for_update=False + ) + waits_after_call = len(waited) + staged = await engine.stage_train_step(_payload()) + waits_after_stage = len(waited) + await engine.train_step( + None, apply_optimizer=False, staged_request_id=staged.request_id + ) + default = await engine.train_step(_payload(), apply_optimizer=True) + return ( + at_dispatch, + waits_after_call, + staged, + waits_after_stage, + default, + ) + + with mock.patch.object( + trainer_worker.jax, + "block_until_ready", + side_effect=_recording_block_until_ready, + ): + at_dispatch, waits_after_call, staged, waits_after_stage, default = ( + asyncio.run(_run()) + ) + + self.assertTrue(at_dispatch["updated"]) + self.assertFalse(at_dispatch["update_ready"]) + self.assertEqual(waits_after_call, 0) + self.assertTrue(staged.update_ready) + self.assertEqual(waits_after_stage, 1) + # Control: the default waits for the update and claims it. + self.assertTrue(default["update_ready"]) + self.assertLen(waited, 2) + + +class _MinimalTrainer: + policy_version = 0 + + def fwd_bwd(self, payload, **kwargs): + del payload, kwargs + + def update(self, **kwargs): + del kwargs + return 1 + + def close(self): + pass + + +class _WaitingTrainer(_MinimalTrainer): + + def wait_for_device(self): + return 0.5 + + +class _OldTrainerWorker(trainer_worker.TrainerWorker): + """`info()` as before staging: no capability keys.""" + + def info(self): + info = super().info() + return dataclasses.replace( + info, + resources={ + k: v + for k, v in info.resources.items() + if k + not in ( + "stage_fwd_bwd", + "reports_device_wait", + "stage_before_update", + ) + }, + ) + + +def _real_engine(worker_cls, trainer): + worker = worker_cls(trainer_factory=lambda: trainer) + worker.initialize() + handle = remote_execution.InProcessActorHandle( + remote_execution.InProcessRemoteExecutionServer(worker) + ) + return _engine(handle) + + +class TrainerStagingSupportTest(absltest.TestCase): + + def test_a_trainer_with_wait_for_device_supports_staging(self): + engine = _real_engine(trainer_worker.TrainerWorker, _WaitingTrainer()) + self.assertEqual( + asyncio.run(engine.trainer_staging_support()), + { + "stage_fwd_bwd": True, + "reports_device_wait": True, + "stage_before_update": True, + }, + ) + + def test_a_trainer_without_it_cannot_report_its_device(self): + engine = _real_engine(trainer_worker.TrainerWorker, _MinimalTrainer()) + self.assertEqual( + asyncio.run(engine.trainer_staging_support()), + { + "stage_fwd_bwd": True, + "reports_device_wait": False, + "stage_before_update": True, + }, + ) + + def test_an_older_trainer_worker_supports_nothing(self): + engine = _real_engine(_OldTrainerWorker, _WaitingTrainer()) + self.assertEqual( + asyncio.run(engine.trainer_staging_support()), + { + "stage_fwd_bwd": False, + "reports_device_wait": False, + "stage_before_update": False, + }, + ) + + def test_the_device_wait_comes_back_through_a_real_worker(self): + engine = _real_engine(trainer_worker.TrainerWorker, _WaitingTrainer()) + staged = asyncio.run(engine.stage_train_step(_payload())) + self.assertEqual(staged.device_wait_s, 0.5) + + def test_no_trainer_for_the_role_is_an_error(self): + engine = distributed_rl_engine.DistributedRLEngine( + rollout_workers=[], trainer_workers={} + ) + with self.assertRaisesRegex(ValueError, "No trainer worker"): + asyncio.run(engine.trainer_staging_support()) + + +if __name__ == "__main__": + absltest.main() diff --git a/tests/experimental/orchestrator/rl_program_pipelined_boundary_robustness_test.py b/tests/experimental/orchestrator/rl_program_pipelined_boundary_robustness_test.py new file mode 100644 index 0000000000..63b1fdbcf8 --- /dev/null +++ b/tests/experimental/orchestrator/rl_program_pipelined_boundary_robustness_test.py @@ -0,0 +1,946 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Failure paths, memory bound and hole rule of the pipelined step boundary. + +These reuse the harness of `rl_program_pipelined_boundary_test` (the trainer +double behind a real `TrainerWorker`, and the frozen base loop) and cover what +its equivalence checks cannot see: + +* How far the take-ahead packs (orchestrator memory). The trainer cannot tell + a take-ahead that holds one micro-batch from one that holds a whole step, + so a memory runaway passes every equivalence check. Checked in the + harness's default packing and in a streaming regime (one row per + micro-batch) where, as in production, a step's micro-batches come out as + its groups are fed, so the take-ahead stops after one and restarts after + the send-ahead. +* An optimizer update that fails on the device, against a trainer that + confirms its updates (maxtext#5571), where a planned stage makes the fused + last call return at dispatch and the stage meets the failure. +* A stage that fails. On the way to the trainer (transport), or in the + reference policy's log-probs just before it, the micro-batch is sent whole + instead and the run goes on; on the trainer, the run ends with that error. +* The `serving >= next` hole rule, fuzzed at the queue: invisible to a + consumer that never reads ahead (flags off), visible to one that does. +* A stage left on a trainer by an orchestrator that died without `stop()`. + +Each check comes with a control that fails on a deliberately wrong input. +""" + +import asyncio +import collections +import contextlib +import logging +import random +import threading +import types +from unittest import mock + +from absl.testing import absltest +from absl.testing import parameterized +import rl_program_pipelined_boundary_test as l3t +from tunix.experimental.common import datatypes +from tunix.experimental.orchestrator import rl_program +from tunix.experimental.queue_manager import trajectory_queue_manager + + +def _worker_of(run: l3t._Run): + """The TrainerWorker behind a harness run's trainer handle.""" + handle = run.program.engine._trainer_workers[datatypes.Role.ACTOR] + return handle._server.bound_instance + + +@contextlib.contextmanager +def _warnings(): + """The messages logged at WARNING or above in the block (none is fine).""" + messages: list[str] = [] + handler = logging.Handler(level=logging.WARNING) + handler.emit = lambda record: messages.append(record.getMessage()) + root = logging.getLogger() + root.addHandler(handler) + try: + yield messages + finally: + root.removeHandler(handler) + + +# --------------------------------------------------------------------------- +# The take-ahead packs at most one feed, for the next step only. +# --------------------------------------------------------------------------- + + +class _StreamingRows: + """One 64-token row per micro-batch, so micro-batches stream out as groups + are fed: production's regime (`pack_size` = fsdp x expert = 64 rows of + 65,536 tokens against 16-trajectory groups). With the harness's default two + rows, a step's micro-batches mostly come out together at its last feed, so + the take-ahead's stop-after-one-micro-batch and its restart after the + send-ahead are rarely exercised.""" + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.assembler._inner.batch_size = 1 + + +class _PrepackProbe(rl_program.StandardRLProgram): + """Records every group taken, ahead (`current_step=None`) or by its step.""" + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + # (running step, state tag, micro-batches ready, staged) at each take-ahead. + self.ahead: list[tuple[int, int, int, bool]] = [] + # State tag -> steps that took groups from it themselves. + self.own: dict[int, set[int]] = collections.defaultdict(set) + # Micro-batches ready right after each take-ahead (one feed's output). + self.ready_after: list[int] = [] + # Take-aheads restarted after the step's first micro-batch was sent ahead. + self.restarts = 0 + self._num_tags = 0 + + async def _prepack(self, st): + if st.num_microbatches: + self.restarts += 1 + await super()._prepack(st) + + def _tag(self, st) -> int: + tag = getattr(st, "_probe_tag", None) + if tag is None: + self._num_tags += 1 + tag = st._probe_tag = self._num_tags + return tag + + async def _take_group(self, st, *, current_step): + tag = self._tag(st) + if current_step is None: + self.ahead.append((self._step, tag, len(st.ready), st.staged is not None)) + else: + self.own[tag].add(current_step) + took = await super()._take_group(st, current_step=current_step) + if current_step is None: + self.ready_after.append(len(st.ready)) + return took + + +class _RunawayPrepack(_PrepackProbe): + """WRONG: takes the whole next step ahead instead of one micro-batch. + + It feeds the same groups in the same order, so the trainer gets the same + micro-batches; but the orchestrator holds every packed micro-batch of the + next step through this step's tail and weight sync (each one dense routing + data for the 397B), which is what the bound below forbids. + """ + + async def _prepack(self, st): + while st.groups_consumed < self.full_batch_size: + if not await self._take_group(st, current_step=None): + return + + +class _StreamingBase(_StreamingRows, l3t._BaseLoopProgram): + pass + + +class _StreamingProbe(_StreamingRows, _PrepackProbe): + pass + + +class _StreamingRunaway(_StreamingRows, _RunawayPrepack): + pass + + +_STREAMING_REFERENCES: dict[tuple[str, type], l3t._Run] = {} + + +def _streaming_reference(name: str, trainer_cls=l3t._Trainer) -> l3t._Run: + key = (name, trainer_cls) + if key not in _STREAMING_REFERENCES: + _STREAMING_REFERENCES[key] = l3t._run( + l3t._SCENARIOS[name], + pipelined=False, + program_cls=_StreamingBase, + trainer_cls=trainer_cls, + ) + return _STREAMING_REFERENCES[key] + + +def _prepack_violations(program: _PrepackProbe) -> list[str]: + """Breaches of the take-ahead's memory bound, as messages.""" + out = [] + for step, tag, ready, staged in program.ahead: + if ready or staged: + out.append( + f"state {tag} taken ahead during step {step} while {ready}" + f" micro-batch(es) were packed (staged={staged})" + ) + running: dict[int, set[int]] = collections.defaultdict(set) + for step, tag, _, _ in program.ahead: + running[tag].add(step) + for tag, steps in running.items(): + owners = program.own.get(tag, set()) + if len(owners) > 1: + out.append(f"state {tag} was trained by several steps {sorted(owners)}") + if owners: + (owner,) = owners + # Taken ahead while the step before it ran, or at its own start (the + # restarted take-ahead it awaits); never two steps ahead. + early = {s for s in steps if s not in (owner - 1, owner)} + if early: + out.append( + f"state {tag} of step {owner} was taken ahead during steps" + f" {sorted(early)}" + ) + return out + + +_READY_AFTER: dict[str, list[int]] = collections.defaultdict(list) +_RESTARTS: collections.Counter[str] = collections.Counter() + +_BOTH_MODES = [ + (f"{name}_{mode}", name, mode) + for name in l3t._SCENARIOS + for mode in ("unstaged", "staged") +] + + +class PrepackBoundTest(parameterized.TestCase): + """The take-ahead holds at most one feed's micro-batches, one step ahead.""" + + @parameterized.named_parameters(*_BOTH_MODES) + def test_default_rows_bound_and_measure(self, name, mode): + run = l3t._run( + l3t._SCENARIOS[name], program_cls=_PrepackProbe, **l3t._MODES[mode] + ) + self.assertEqual(_prepack_violations(run.program), []) + if name in l3t._SENT_AHEAD: + self.assertNotEmpty(run.program.ahead, "nothing was taken ahead") + _READY_AFTER["default"].extend(run.program.ready_after) + _RESTARTS["default"] += run.program.restarts + # One feed's output in this harness: at most two micro-batches. + self.assertLessEqual(max(run.program.ready_after, default=0), 2) + + @parameterized.named_parameters(*_BOTH_MODES) + def test_streaming_take_ahead_is_bounded_and_trains_the_same( + self, name, mode + ): + run = l3t._run( + l3t._SCENARIOS[name], program_cls=_StreamingProbe, **l3t._MODES[mode] + ) + self.assertEqual(_prepack_violations(run.program), []) + l3t._assert_equivalent(self, _streaming_reference(name), run) + l3t._assert_weights_follow_updates(self, run) + if name in l3t._SENT_AHEAD: + self.assertNotEmpty(run.program.ahead, "nothing was taken ahead") + _READY_AFTER["streaming"].extend(run.program.ready_after) + _RESTARTS["streaming"] += run.program.restarts + + @parameterized.named_parameters((name, name) for name in l3t._SCENARIOS) + def test_streaming_staged_against_a_token_trainer(self, name): + run = l3t._run( + l3t._SCENARIOS[name], + program_cls=_StreamingProbe, + pipelined=True, + staged=True, + trainer_cls=l3t._TokenTrainer, + ) + self.assertEqual(_prepack_violations(run.program), []) + l3t._assert_equivalent_but_metrics_fetch( + self, _streaming_reference(name, l3t._TokenTrainer), run + ) + l3t._staged_boundaries_confirm_the_update(self, run) + _RESTARTS["streaming_token_staged"] += run.program.restarts + + def test_the_streaming_regime_runs_the_restart_path(self): + # Non-vacuous: the take-ahead stops after one micro-batch with groups + # left, and restarts once that micro-batch is sent ahead. + restarts = 0 + for name in ("prompt_batch", "arrival", "reference_kl"): + for mode in ("unstaged", "staged"): + run = l3t._run( + l3t._SCENARIOS[name], + program_cls=_StreamingProbe, + **l3t._MODES[mode], + ) + restarts += run.program.restarts + self.assertGreater(restarts, 0) + + def test_the_bound_catches_a_runaway_that_the_equivalence_check_misses(self): + for name in ("prompt_batch", "arrival"): + ref = _streaming_reference(name) + run = l3t._run( + l3t._SCENARIOS[name], pipelined=True, program_cls=_StreamingRunaway + ) + # The trainer cannot tell: same calls, inputs, weights and outputs. + self.assertTrue(l3t._check_equivalent(ref, run), name) + # The bound can. + self.assertNotEmpty(_prepack_violations(run.program), name) + + +# --------------------------------------------------------------------------- +# An update that fails on the device, with wait_for_update=False. +# --------------------------------------------------------------------------- + + +class _BadToken: + """`last_update_token` of an update that failed on the device.""" + + def __init__(self, trainer, version): + self._trainer = trainer + self._version = version + + def block_until_ready(self): + self._trainer.log.append(("token_wait_failed", self._version)) + raise RuntimeError(f"XLA: update {self._version} failed on device") + + +class _FailsSecondUpdate(l3t._TokenTrainer): + """A token trainer whose second update fails on the device. + + Reading anything the failed update produced fails too: `get_metrics` + returns its gradient norm (maxtext#5571's token), and a save or a weight + sync reads the weights it wrote (INFERRED MaxText semantics). Each such + attempt is logged, so a path that saves after the failure shows up. + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.update_failed = False + + def update(self, **kwargs): + version = super().update(**kwargs) + if version == 2: + self.update_failed = True + self.last_update_token = _BadToken(self, version) + return version + + def _read_after_failure(self, what): + if self.update_failed: + self.log.append(("after_failed_update", what)) + raise RuntimeError(f"XLA: update 2 failed on device (read by {what})") + + def get_metrics(self): + self._read_after_failure("get_metrics") + return super().get_metrics() + + def save_checkpoint(self, metadata, **kwargs): + self._read_after_failure("save_checkpoint") + return super().save_checkpoint(metadata, **kwargs) + + def prepare_weight_sync(self, sync_request=None, **kwargs): + self._read_after_failure("prepare_weight_sync") + return super().prepare_weight_sync(sync_request, **kwargs) + + +class _FailsSecondUpdateAtTheDevice(_FailsSecondUpdate): + """The same, where the stage's device wait meets the failure first. + + MaxText's `wait_for_device` is its throttler's `block_until_ready()`, + which waits on every queued output, the update's gradient norm included + (maxtext 6db77afd6, the MaxText half of this change). There the stage's + device wait, not its token wait, is what raises. + """ + + def wait_for_device(self): + if self.update_failed: + self.log.append(("device_wait_failed", 2)) + raise RuntimeError( + "XLA: update 2 failed on device (read by wait_for_device)" + ) + return super().wait_for_device() + + +_FAILED_UPDATE_TRAINERS = { + "token_wait": _FailsSecondUpdate, + "device_wait": _FailsSecondUpdateAtTheDevice, +} + +# Per (mode, trainer): whether a get_metrics met the failed update. +_GET_METRICS_AFTER_FAILURE: dict[str, bool] = {} + +_FAILED_UPDATE_MODES = { + "base": dict(pipelined=False, program_cls=l3t._BaseLoopProgram), + "off": dict(pipelined=False), + "unstaged": dict(pipelined=True), + "staged": dict(pipelined=True, staged=True), +} + +_FAILED_UPDATE_CASES = [ + (f"{mode}_{trainer}", mode, trainer) + for mode in _FAILED_UPDATE_MODES + for trainer in _FAILED_UPDATE_TRAINERS +] + + +def _run_failed_update( + mode: str, trainer: str +) -> tuple[l3t._Run, list[BaseException]]: + errors: list[BaseException] = [] + run = l3t._run( + l3t._SCENARIOS["prompt_batch"], + trainer_cls=_FAILED_UPDATE_TRAINERS[trainer], + errors=errors, + **_FAILED_UPDATE_MODES[mode], + ) + return run, errors + + +def _stage_met_the_failure(run: l3t._Run) -> bool: + """Whether the boundary after the failed update 2 was staged.""" + update_2 = max( + i for i, e in enumerate(run.log) if e[0] == "update" and e[1] == 2 + ) + return "stage" in [e[0] for e in run.log[update_2:]] + + +class FailedUpdateTest(parameterized.TestCase): + """A failed update is never saved or stamped, and is reported as the + optimizer step, in every mode.""" + + @parameterized.named_parameters(*_FAILED_UPDATE_CASES) + def test_a_failed_update_is_never_checkpointed_or_stamped( + self, mode, trainer + ): + run, errors = _run_failed_update(mode, trainer) + self.assertLen(errors, 1, run.log) + self.assertIn("XLA: update 2 failed on device", str(errors[0])) + # Step 0 (optimizer step 1) saved; no save or weight sync was attempted + # after the failed update (a get_metrics meeting it would be the barrier + # doing its job, so it is not counted here). + self.assertEqual([c["step"] for c in run.checkpoints], [1], run.log) + for what in ("save_checkpoint", "prepare_weight_sync"): + self.assertNotIn(("after_failed_update", what), run.log, run.log) + _GET_METRICS_AFTER_FAILURE[f"{mode}_{trainer}"] = ( + "after_failed_update", + "get_metrics", + ) in run.log + self.assertEqual( + run.program.last_step_timestamp_ms, run.checkpoints[0]["timestamp_ms"] + ) + self.assertEqual(_worker_of(run).state, datatypes.WorkerState.ERROR) + self.assertIsNone(_worker_of(run)._staged_fwd_bwd) + if mode == "staged": + if _stage_met_the_failure(run): + # The stage, not the fused call, met the failure: the call returned + # at the update's dispatch and the stage waited for it. + failed_wait = ( + "device_wait_failed" + if trainer == "device_wait" + else "token_wait_failed" + ) + self.assertIn((failed_wait, 2), run.log) + else: + # Planned but nothing packed in time (a loaded machine): get_metrics + # is the barrier again, and it is what meets the failure. + self.assertIn(("after_failed_update", "get_metrics"), run.log) + + @parameterized.named_parameters(*_FAILED_UPDATE_CASES) + def test_the_failure_is_attributed_to_the_optimizer_step(self, mode, trainer): + # Off staging, google/tunix#2719's fused call says what failed, in the + # error and in the worker's health report. A staged boundary must too: + # there the stage meets the failure, and ends the run with it, rather + # than falling back to get_metrics, whose error would not say it. + run, errors = _run_failed_update(mode, trainer) + if mode == "staged": + self.assertTrue( + _stage_met_the_failure(run), + "precondition: the boundary after update 2 was staged (timing; the" + " L3 suite expects every prompt_batch boundary staged)", + ) + head = str(errors[0]).splitlines()[0] + self.assertIn("optimizer update() failed", head) + last_error = _worker_of(run).heartbeat().last_error + self.assertIn("optimizer update() failed", last_error) + if mode == "staged": + self.assertIn("stage_fwd_bwd", head) + self.assertIn("XLA: update 2 failed on device", head) + + +# --------------------------------------------------------------------------- +# A stage that fails (checkpoint every step: val_start_step is None). +# --------------------------------------------------------------------------- + + +def _first_fingerprint_of_step_1() -> str: + """Content hash of step 1's first micro-batch, from the base loop's run.""" + ref = l3t._reference("prompt_batch") + return next(e[1] for e in ref.log if e[0] == "fwd_bwd" and e[2] == 1) + + +_TRANSPORT_FAILURE_MODES = { + "base": dict(pipelined=False, program_cls=l3t._BaseLoopProgram), + "unstaged": dict(pipelined=True), + "staged": dict(pipelined=True, staged=True), +} + + +def _run_with_transport_failure( + mode: str, + program_cls: type[rl_program.StandardRLProgram] | None = None, +) -> tuple[l3t._Run, list[BaseException], dict[str, object]]: + """Runs prompt_batch with the request carrying step 1's first micro-batch + failing once in transport, before the trainer worker runs it.""" + target = _first_fingerprint_of_step_1() + real_execute = l3t._RemoteHandle._execute + fired: dict[str, object] = {"n": 0, "method": None} + + async def _execute(self, request): + if request.method_name in ("fwd_bwd", "stage_fwd_bwd") and not fired["n"]: + payload = getattr(request.kwargs.get("request"), "payload", None) + if payload is not None and l3t._fingerprint(payload) == target: + fired["n"] = 1 + fired["method"] = request.method_name + raise ConnectionError("transport: stream reset before the call ran") + return await real_execute(self, request) + + errors: list[BaseException] = [] + kwargs = dict(_TRANSPORT_FAILURE_MODES[mode]) + if program_cls is not None: + kwargs["program_cls"] = program_cls + with mock.patch.object(l3t._RemoteHandle, "_execute", _execute): + run = l3t._run(l3t._SCENARIOS["prompt_batch"], errors=errors, **kwargs) + return run, errors, fired + + +def _assert_fell_back_and_trained_the_same( + test: absltest.TestCase, name: str, run: l3t._Run +) -> None: + """The boundary went on unstaged, and the trainer cannot tell the run from + the base loop's: the micro-batch went whole in step 1's first train call.""" + l3t._assert_equivalent(test, l3t._reference(name), run) + l3t._assert_weights_follow_updates(test, run) + test.assertNotIn("stage", [e[0] for e in run.log]) + # Sent ahead at every boundary, staged at none: off for the rest of the run. + test.assertEqual(l3t._boundaries(run), [(True, False)] * 2, run.log) + test.assertTrue(run.program._staging_latched_off) + worker = _worker_of(run) + test.assertEqual(worker.state, datatypes.WorkerState.READY) + test.assertIsNone(worker._staged_fwd_bwd) + + +class FailedStageTransportTest(parameterized.TestCase): + """What one failed send of step 1's first micro-batch costs.""" + + @parameterized.parameters("base", "unstaged") + def test_the_base_loop_and_the_unstaged_send_keep_step_0(self, mode): + # Sent in a train call, the micro-batch fails after step 0's barrier, + # checkpoint and bookkeeping. + run, errors, fired = _run_with_transport_failure(mode) + self.assertEqual(fired["n"], 1) + self.assertEqual(fired["method"], "fwd_bwd") + self.assertLen(errors, 1) + self.assertIsInstance(errors[0], ConnectionError) + self.assertEqual([c["step"] for c in run.checkpoints], [1]) + self.assertIn(("on_step_end", 0), run.log) + self.assertEqual(run.program.step, 1) + + def test_a_stage_that_fails_on_the_way_falls_back_to_the_unstaged_send( + self, + ): + with _warnings() as warnings: + run, errors, fired = _run_with_transport_failure("staged") + self.assertEqual( + fired["method"], + "stage_fwd_bwd", + "precondition: boundary 0 was staged (timing; the L3 suite expects" + " every prompt_batch boundary staged)", + ) + # The stage was optional: the run keeps step 0's barrier, checkpoint and + # bookkeeping, sends the micro-batch whole, and (the failure being + # one-shot) finishes. + self.assertEqual(errors, []) + _assert_fell_back_and_trained_the_same(self, "prompt_batch", run) + self.assertTrue( + any( + "[pipeline] Staging failed" in m and "ConnectionError" in m + for m in warnings + ), + warnings, + ) + + def test_without_the_fallback_the_failed_stage_costs_step_0(self): + # Negative control: the same failure, treated as fatal (the code before + # the fallback), loses step 0's checkpoint and bookkeeping, and the + # checks above see it. + with mock.patch.object( + rl_program, "_raised_by_worker_method", lambda exc: True + ): + run, errors, fired = _run_with_transport_failure("staged") + self.assertEqual(fired["method"], "stage_fwd_bwd") + self.assertLen(errors, 1) + self.assertEqual(run.checkpoints, []) + self.assertNotIn(("on_step_end", 0), run.log) + probe = absltest.TestCase() + with self.assertRaises(AssertionError): + _assert_fell_back_and_trained_the_same(probe, "prompt_batch", run) + + +class _MarksItsStage(rl_program.StandardRLProgram): + """Marks the stage's own reference call, so a test can fail exactly it.""" + + staging = threading.Event() + + async def _maybe_stage_next_microbatch(self, st, step_result): + self.staging.set() + try: + return await super()._maybe_stage_next_microbatch(st, step_result) + finally: + self.staging.clear() + + +class FailedReferenceWhileStagingTest(absltest.TestCase): + """The reference policy's log-probs, fetched for the stage, fail once. + + That call is not the trainer's: whether it failed on the way or in the + reference worker itself, the trainer is untouched, so the micro-batch goes + whole in step 1's first train call (which fetches the log-probs again), as + without staging. + """ + + def setUp(self): + super().setUp() + _MarksItsStage.staging.clear() + self.fired = 0 + + def _one_shot(self) -> bool: + if _MarksItsStage.staging.is_set() and not self.fired: + self.fired += 1 + return True + return False + + def _run(self) -> tuple[l3t._Run, list[BaseException]]: + errors: list[BaseException] = [] + with _warnings() as warnings: + run = l3t._run( + l3t._SCENARIOS["reference_kl"], + pipelined=True, + staged=True, + program_cls=_MarksItsStage, + errors=errors, + ) + self.assertEqual( + self.fired, + 1, + "precondition: boundary 0 was staged (timing; the L3 suite expects" + " every reference_kl boundary staged)", + ) + self.warnings = warnings + return run, errors + + def test_a_transport_failure_falls_back(self): + real_execute = l3t._RemoteHandle._execute + outer = self + + async def _execute(self, request): + if ( + request.method_name == "per_token_logps" + and isinstance(self._server.bound_instance, l3t._Reference) + and outer._one_shot() + ): + raise ConnectionError("transport: reference stream reset") + return await real_execute(self, request) + + with mock.patch.object(l3t._RemoteHandle, "_execute", _execute): + run, errors = self._run() + self.assertEqual(errors, []) + _assert_fell_back_and_trained_the_same(self, "reference_kl", run) + self.assertTrue( + any("ConnectionError" in m for m in self.warnings), self.warnings + ) + + def test_a_failure_in_the_reference_worker_falls_back(self): + # A remote error, like a trainer-side one, but from the reference worker. + real_logps = l3t._Reference.per_token_logps + outer = self + + def _per_token_logps(self, items, **kwargs): + if outer._one_shot(): + raise RuntimeError("reference model: device lost") + return real_logps(self, items, **kwargs) + + with mock.patch.object(l3t._Reference, "per_token_logps", _per_token_logps): + run, errors = self._run() + self.assertEqual(errors, []) + _assert_fell_back_and_trained_the_same(self, "reference_kl", run) + self.assertTrue( + any( + "RemoteExecutionError" in m and "reference model: device lost" in m + for m in self.warnings + ), + self.warnings, + ) + + +# --------------------------------------------------------------------------- +# The hole rule, fuzzed at the queue. +# --------------------------------------------------------------------------- + + +def _advance_over_holes_eq(self): + """The hole rule before the pipelined boundary: serving == next only.""" + while ( + self._serving_batch_idx == self._next_batch_idx + and self._next_batch_idx in self._sealed + and self._admitted.get(self._next_batch_idx, 0) == 0 + ): + hole = self._next_batch_idx + self._drop_batch_state(hole) + self._advance_to(hole + 1) + + +class _EqRuleQueue(trajectory_queue_manager.BatchOrderedQueueManager): + """The queue with the hole rule as it was before (serving == next).""" + + _advance_over_holes = _advance_over_holes_eq + + +def _queue_scenario(seed: int): + """Batches (some short, some holes) and a producer schedule.""" + rng = random.Random(seed) + full = rng.choice([2, 3]) + num_batches = rng.randint(3, 7) + per_batch = [] + for b in range(num_batches): + n = full if rng.random() < 0.6 else rng.randint(0, full - 1) + kinds = [] + for _ in range(n): + r = rng.random() + kinds.append( + "valid" if r < 0.65 else "filtered" if r < 0.85 else "dropped" + ) + per_batch.append( + [("expect", b, n)] + [(k, b, j) for j, k in enumerate(kinds)] + ) + window = rng.choice([1, 2, 3]) + events, active, nxt = [], [], 0 + while nxt < num_batches or active: + while len(active) < window and nxt < num_batches: + active.append(nxt) + nxt += 1 + b = rng.choice(active) + events.append(per_batch[b].pop(0)) + if not per_batch[b]: + active.remove(b) + producer_yields = [rng.randint(0, 3) for _ in events] + consumer_yields = [rng.randint(0, 3) for _ in range(4 * len(events) + 8)] + return full, events, producer_yields, consumer_yields + + +def _drive_queue(seed: int, queue_cls, *, read_ahead: bool): + """Runs one producer/consumer schedule; returns (trace, steps trained).""" + full, events, producer_yields, consumer_yields = _queue_scenario(seed) + trace: list[tuple] = [] + q = queue_cls( + full_batch_size=full, + num_generations=1, + filter_fn=lambda g: ([], list(g)) if g[0].drop else (list(g), []), + on_cursor_advance=lambda idx: trace.append(("cursor", idx)), + ) + cy = iter(consumer_yields) + + def record(*event): + trace.append(event + (q.next_batch_idx, q.serving_batch_idx)) + + def ids(result): + return None if result is None else (result[0], result[1][0].prompt_id) + + async def produce(): + for (kind, b, x), y in zip(events, producer_yields): + for _ in range(y): + await asyncio.sleep(0) + if kind == "expect": + q.expect(b, x) + elif kind == "dropped": + q.dropped(b, 1) + else: + await q.put( + types.SimpleNamespace( + prompt_id=f"b{b}g{x}", + metadata={"batch_idx": b}, + is_valid=True, + policy_version=0, + drop=kind == "filtered", + ) + ) + await q.close() + + async def consume(): + steps = [] + ahead = None + while True: + first = None + if ahead is not None: + first = await ahead + ahead = None + record("ahead", ids(first)) + if first is None: + first = await q.get_ordered_group(None) + record("get", None, ids(first)) + if first is None: + return steps + b, group = first + taken = [group[0].prompt_id] + while len(taken) < full: + r = await q.get_ordered_group(b) + record("get", b, ids(r)) + if r is None: + break + taken.append(r[1][0].prompt_id) + if read_ahead: + # The pipelined boundary: read the next batch by index before this + # one commits. + ahead = asyncio.ensure_future(q.get_ordered_group(b + 1)) + for _ in range(next(cy, 0)): + await asyncio.sleep(0) + q.commit_batch(b) + record("commit", b) + steps.append((b, tuple(taken))) + + async def main(): + producer = asyncio.ensure_future(produce()) + steps = await consume() + await producer + return steps + + steps = asyncio.run(asyncio.wait_for(main(), timeout=10)) + return trace, steps + + +_SEEDS = range(300) + + +class HoleRuleIdentityFuzzTest(absltest.TestCase): + + def test_the_new_rule_is_invisible_without_reading_ahead(self): + differ = [ + s + for s in _SEEDS + if _drive_queue(s, _EqRuleQueue, read_ahead=False) + != _drive_queue( + s, + trajectory_queue_manager.BatchOrderedQueueManager, + read_ahead=False, + ) + ] + self.assertEqual(differ, []) + + def test_the_fuzz_sees_the_rule_when_reading_ahead(self): + # Negative control: the same fuzz does tell the two rules apart once the + # consumer reads ahead, which is what the rule was changed for. + differ = [ + s + for s in _SEEDS + if _drive_queue(s, _EqRuleQueue, read_ahead=True) + != _drive_queue( + s, + trajectory_queue_manager.BatchOrderedQueueManager, + read_ahead=True, + ) + ] + self.assertNotEmpty(differ) + print(f"[holes] read-ahead: rules differ in {len(differ)}/{len(_SEEDS)}") + + def test_reading_ahead_trains_the_same_batches_and_groups(self): + queue_cls = trajectory_queue_manager.BatchOrderedQueueManager + for s in _SEEDS: + _, base_steps = _drive_queue(s, queue_cls, read_ahead=False) + _, ahead_steps = _drive_queue(s, queue_cls, read_ahead=True) + self.assertEqual(base_steps, ahead_steps, f"seed {s}") + + +# --------------------------------------------------------------------------- +# A stage left on a trainer by an orchestrator that died. +# --------------------------------------------------------------------------- + + +class _RestorableTrainer(l3t._Trainer): + """The trainer double, plus the bring-up calls a successor makes.""" + + def compile(self, dummy_data=None): + del dummy_data + + def restore_checkpoint(self, **kwargs): + del kwargs + return {} + + +class StaleStageAfterRestartTest(absltest.TestCase): + """An orchestrator that dies between a stage and the train call that runs + it (OOM kill, preemption: no `stop()`) leaves the micro-batch held. Its + successor brings the surviving trainer worker up again (`initialize`, + `compile`) and restores the checkpoint, which must drop it.""" + + def _crashed_worker(self, held: bool): + worker = l3t._Worker(lambda: _RestorableTrainer([], l3t._Costs())) + worker.initialize() + if held: + worker.stage_fwd_bwd( + datatypes.TrainRequest( + request_id="dead-orchestrator", + payload=l3t._payload(1, logps=False), + metadata={}, + ) + ) + return worker + + def _stage_new(self, worker): + return worker.stage_fwd_bwd( + datatypes.TrainRequest( + request_id="successor", + payload=l3t._payload(2, logps=False), + metadata={}, + ) + ) + + def test_a_held_stage_survives_a_bare_initialize(self): + # Control: the stage really is held, and initialize() alone (a no-op on + # a ready worker) does not drop it, so the successor's first stage would + # be refused if nothing else did. + worker = self._crashed_worker(held=True) + worker.initialize() + with self.assertRaisesRegex( + RuntimeError, "already holds staged micro-batch 'dead-orchestrator'" + ): + self._stage_new(worker) + + def test_a_restored_worker_without_a_held_stage_stages(self): + # Control: the successor's bring-up sequence is fine on its own. + worker = self._crashed_worker(held=False) + worker.initialize() + worker.compile(None) + worker.restore_checkpoint() + self.assertTrue(self._stage_new(worker).metadata["staged"]) + + def test_a_successor_can_stage_after_a_crash(self): + worker = self._crashed_worker(held=True) + with _warnings() as warnings: + worker.initialize() + worker.compile(None) + worker.restore_checkpoint() + self.assertIsNone(worker._staged_fwd_bwd) + self.assertTrue(any("dead-orchestrator" in m for m in warnings), warnings) + self.assertTrue(self._stage_new(worker).metadata["staged"]) + + +def tearDownModule(): + for regime, values in sorted(_READY_AFTER.items()): + print( + f"[take-ahead] {regime}: micro-batches ready after a take-ahead: max" + f" {max(values, default=0)} over {len(values)} take-aheads, histogram" + f" {dict(sorted(collections.Counter(values).items()))}" + ) + print(f"[take-ahead] restarts after a send-ahead: {dict(_RESTARTS)}") + print(f"[failed update] get_metrics met it: {_GET_METRICS_AFTER_FAILURE}") + + +if __name__ == "__main__": + absltest.main() diff --git a/tests/experimental/orchestrator/rl_program_pipelined_boundary_test.py b/tests/experimental/orchestrator/rl_program_pipelined_boundary_test.py new file mode 100644 index 0000000000..7c01ad0798 --- /dev/null +++ b/tests/experimental/orchestrator/rl_program_pipelined_boundary_test.py @@ -0,0 +1,3509 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for StandardRLProgram's pipelined step boundary. + +`_train_loop` is one loop for both step boundaries. Every scenario runs +through it in three modes -- flags off, `pipeline_step_boundary`, and that +plus `stage_next_microbatch` -- and through a frozen, byte-identical copy of +the loop as it was before the pipelined boundary (`_BaseLoopProgram`: +atwigg/mlperf as of google/tunix#2719's fused final update and deferred +get_metrics), against the same trainer double. Flags off must be that loop +exactly: the same trainer calls, the same program events (feeds, flushes, +step callbacks) in the same order, and the same metrics. The pipelined modes +must be indistinguishable from it to the trainer: the same calls with the +same inputs in the same order, each micro-batch computed on the same weights, +the same checkpoints, and the same logged metric values. Separate tests show +that the pipelined modes do change *when* the orchestrator does its work, and +that every check here fails on a deliberately wrong program. + +Two trainer doubles cover the two regimes of the fused last call. `_Trainer` +cannot confirm its update, like MaxText at tunix's pin today: the call +returns at the update's dispatch and get_metrics stays ahead of the +checkpoint save. `_TokenTrainer` can (`last_update_token`, maxtext#5571): the +call returns once the update landed and the base loop defers get_metrics past +the weight sync and commit. There a planned stage makes the last call return +at the dispatch and confirms the update itself, so a staged boundary is the +base loop's plus the one stage call; a planned stage that finds nothing +packed in time falls back to get_metrics ahead of the save, which the +relaxed check (`_assert_equivalent_but_metrics_fetch`) allows and nothing +else. + +The trainer double runs behind a real `TrainerWorker`, on its own thread and +event loop, like the remote gRPC server: methods run one at a time and block +that loop while they run, and a request's transfer time is spent on that loop +before the method runs. Its device keeps the MaxText engine's rules: at most +two computations in flight, run in dispatch order; `get_metrics` waits for the +update and returns only the newest step's metrics buffer; `save_checkpoint` +and `prepare_weight_sync` wait for every computation in flight. +""" + +import asyncio +import collections +from collections.abc import Callable, Sequence +import contextlib +import dataclasses +import hashlib +import inspect +import os +import random +import subprocess +import threading +import time +import types +from typing import Any +from unittest import mock + +from absl import logging +from absl.testing import absltest +from absl.testing import parameterized +import numpy as np +import tunix +from tunix.experimental.common import datatypes +from tunix.experimental.common import logging_utils +from tunix.experimental.examples.deepswe_dist import run_deepswe_dist +from tunix.experimental.metrics import metrics as exp_metrics +from tunix.experimental.orchestrator import batch_assembly +from tunix.experimental.orchestrator import distributed_rl_engine +from tunix.experimental.orchestrator import rl_program +from tunix.experimental.queue_manager import trajectory_queue_manager +from tunix.experimental.weight_sync import weight_sync +from tunix.experimental.worker import remote_execution +from tunix.experimental.worker import trainer_worker as trainer_worker_lib + + +NUM_GENERATIONS = 2 +ROW_TOKENS = 64 + +# --------------------------------------------------------------------------- +# Data +# --------------------------------------------------------------------------- + + +def _payload(seed: int, *, logps: bool) -> datatypes.RLTrainerPayload: + """An unbatched trajectory whose content is unique to `seed`.""" + rng = np.random.default_rng(seed) + p_len = int(rng.integers(3, 9)) + c_len = int(rng.integers(6, 22)) + completion_mask = (rng.random(c_len) > 0.2).astype(np.float32) + completion_mask[0] = 1.0 + return datatypes.RLTrainerPayload( + prompt_ids=rng.integers(1, 1000, p_len).astype(np.int32), + prompt_mask=np.ones(p_len, dtype=np.float32), + completion_ids=rng.integers(1, 1000, c_len).astype(np.int32), + completion_mask=completion_mask, + advantages=np.full(c_len, rng.normal(), dtype=np.float32), + old_per_token_logps=( + -rng.random(c_len).astype(np.float32) if logps else None + ), + routed_experts=rng.integers(0, 64, (p_len + c_len, 2, 2)).astype( + np.int16 + ), + ) + + +def _item( + prompt_idx: int, + gen: int, + *, + full_batch_size: int, + logps: bool, +) -> datatypes.TrajectoryItem: + """A scored trajectory as `critique_stage` puts it on `scored_q`.""" + seed = 1000 * prompt_idx + gen + payload = _payload(seed, logps=logps) + item = datatypes.TrajectoryItem( + prompt_id=f"prompt_{prompt_idx}", + group_index=gen, + start_step=0, + traj={ + "trajectory_reward": float((seed * 7919) % 5) / 4.0, + "status": datatypes.TrajectoryStatus.SUCCEEDED, + "steps": [0] * (1 + seed % 3), + }, + is_valid=True, + prompt_tokens=np.asarray(payload.prompt_ids), + completion_tokens=np.asarray(payload.completion_ids), + policy_version=0, + metadata={ + "prompt_idx": prompt_idx, + "batch_idx": prompt_idx // full_batch_size, + "intra_batch_idx": prompt_idx % full_batch_size, + }, + ) + item.payload = payload # pyrefly: ignore[missing-attribute] + return item + + +def _groups( + num_prompts: int, *, full_batch_size: int, logps: bool +) -> list[list[datatypes.TrajectoryItem]]: + return [ + [ + _item(p, g, full_batch_size=full_batch_size, logps=logps) + for g in range(NUM_GENERATIONS) + ] + for p in range(num_prompts) + ] + + +def _fingerprint(payload: Any) -> str: + """Content hash of a micro-batch payload as the trainer receives it.""" + h = hashlib.sha256() + for field in dataclasses.fields(payload): + value = getattr(payload, field.name) + if field.name == "metadata": + value = (value or {}).get("trajectory_ids") + h.update(repr(value).encode()) + continue + if value is None: + h.update(f"{field.name}=None;".encode()) + continue + arr = np.ascontiguousarray(np.asarray(value)) + h.update(f"{field.name}:{arr.dtype}:{arr.shape};".encode()) + h.update(arr.tobytes()) + return h.hexdigest()[:16] + + +def _loss(fingerprint: str, weights_version: int) -> float: + """A loss that depends on both the data and the weights it is computed on.""" + digest = hashlib.sha256(f"{fingerprint}@{weights_version}".encode()) + return int(digest.hexdigest()[:8], 16) / 2.0**32 + + +# --------------------------------------------------------------------------- +# Trainer double +# --------------------------------------------------------------------------- + + +class _HostStagedSync: + """A synchronizer whose transfer reads a host copy (background sync is ok).""" + + staged_on_host = True + + def release_buffers(self) -> int: + return 0 + + +@dataclasses.dataclass +class _Costs: + """Seconds of trainer work; all zero for the equivalence tests.""" + + device_s: float = 0.0 # one micro-batch's forward/backward on the device + update_s: float = 0.0 # the optimizer update on the device + stage_s: float = 0.0 # prepare_weight_sync after the device drains + host_s: float = 0.0 # moving one micro-batch's data to the trainer + metrics_s: float = 0.0 # get_metrics once the update has landed + + +class _Trainer: + """MaxText-engine-shaped trainer double (see the module docstring).""" + + def __init__(self, log: list[Any], costs: _Costs, max_inflight: int = 2): + self.log = log + self.costs = costs + self.max_inflight = max_inflight + self.policy_version = 0 + self.checkpoint_dir = "/ckpt" + self._weight_sync = _HostStagedSync() + self.weights_version = 0 + self._accumulated: list[str] = [] + self._buffers: list[tuple[int, list[Any]]] = [] + self._inflight: collections.deque[float] = collections.deque() + self._device_free_at = 0.0 + self._update_lands_at: float | None = None + # (kind, label, start, end) on the device, monotonic clock. + self.device: list[tuple[str, Any, float, float]] = [] + # Wall-clock ms at which each update finished on the device. + self.update_landed_ms: list[float] = [] + # What each `wait_for_device` waited, in seconds. + self.device_waits: list[float] = [] + + def _wait_until(self, t: float) -> None: + delay = t - time.monotonic() + if delay > 0: + time.sleep(delay) + + def _dispatch(self, kind: str, label: Any, seconds: float) -> None: + # InflightThrottler.wait_for_next, then an in-order device queue. + if len(self._inflight) >= self.max_inflight: + self._wait_until(self._inflight.popleft()) + start = max(time.monotonic(), self._device_free_at) + end = start + seconds + self._device_free_at = end + self._inflight.append(end) + self.device.append((kind, label, start, end)) + + def _wait_for_all(self) -> None: + self._wait_until(self._device_free_at) + self._inflight.clear() + + def _step_buffer(self) -> list[Any]: + step = self.weights_version + if not self._buffers or self._buffers[-1][0] != step: + self._buffers.append((step, [])) + return self._buffers[-1][1] + + def prepare_fwd_bwd_input(self, payload: Any) -> Any: + """What `TrainerWorker.stage_fwd_bwd` holds; here the payload itself.""" + self.log.append(("stage", _fingerprint(payload), self.weights_version)) + return payload + + def wait_for_device(self) -> float: + """MaxText's `wait_for_device`: waits for the device, leaves the queue.""" + start = time.monotonic() + self._wait_until(self._device_free_at) + waited = time.monotonic() - start + self.device_waits.append(waited) + return waited + + def fwd_bwd(self, payload: Any, **kwargs: Any) -> None: + del kwargs + fingerprint = _fingerprint(payload) + self._dispatch("fwd_bwd", fingerprint, self.costs.device_s) + self.log.append(("fwd_bwd", fingerprint, self.weights_version)) + self._accumulated.append(fingerprint) + self._step_buffer().append(_loss(fingerprint, self.weights_version)) + + def update(self, **kwargs: Any) -> int: + del kwargs + self._dispatch("update", self.weights_version + 1, self.costs.update_s) + end = self._device_free_at + self._update_lands_at = end + self.update_landed_ms.append( + (time.time() + (end - time.monotonic())) * 1000.0 + ) + norm = _loss("|".join(self._accumulated), self.weights_version) + self._step_buffer().append(("grad_norm", norm)) + self.weights_version += 1 + self.log.append(("update", self.weights_version, tuple(self._accumulated))) + self._accumulated = [] + return self.weights_version + + def get_metrics(self) -> exp_metrics.MetricsBuffer: + if self._update_lands_at is not None: + self._wait_until(self._update_lands_at) + time.sleep(self.costs.metrics_s) + buffers, self._buffers = self._buffers, [] + if not buffers: + return exp_metrics.MetricsBuffer(id=-1) + step, values = buffers[-1] + # Like `MaxTextTrainingEngine.get_metrics`: the newest buffer only. + self.log.append(("get_metrics", step, len(buffers))) + losses = [v for v in values if not isinstance(v, tuple)] + norms = [v[1] for v in values if isinstance(v, tuple)] + scalars: dict[str, Any] = {} + if losses: + scalars["loss"] = np.asarray(losses, dtype=np.float64) + if norms: + scalars["grad_norm"] = np.asarray(norms, dtype=np.float64) + return exp_metrics.MetricsBuffer(id=step, scalar_metrics=scalars) + + def save_checkpoint(self, metadata: Any, **kwargs: Any) -> None: + self._wait_for_all() + self.log.append(( + "save", + tuple(sorted((metadata or {}).items())), + tuple(sorted(kwargs.items())), + self.weights_version, + tuple(self._accumulated), + tuple(step for step, _ in self._buffers), + )) + + def prepare_weight_sync(self, sync_request: Any = None, **kwargs: Any) -> Any: + del sync_request, kwargs + self._wait_for_all() + time.sleep(self.costs.stage_s) + self.log.append( + ("prepare_weight_sync", self.weights_version, tuple(self._accumulated)) + ) + return [{"unit": "trainer"}] + + def release_weight_sync(self, **kwargs: Any) -> Any: + del kwargs + self.log.append(("release_weight_sync",)) + return True + + def close(self) -> None: + pass + + +class _Worker(trainer_worker_lib.TrainerWorker): + """TrainerWorker whose actor log-probs come from the trainer double. + + They depend on the weights they are computed on, so a micro-batch rescored + (`use_rollout_logps`) on the wrong weights trains on different inputs. + """ + + def per_token_logps(self, items, **kwargs): + del kwargs + trainer = self._trainer + completion = np.asarray(items.completion_tokens) + version = trainer.weights_version + trainer.log.append(("per_token_logps", version, completion.shape)) + logps = -((completion.astype(np.float32) * (version + 1)) % 97) / 97.0 + return datatypes.LogprobsResponse( + request_id=items.request_id, + per_token_logps=logps, + model_version=version, + ) + + +class _OldTrainerWorker(_Worker): + """A trainer image from before staging. + + No `stage_fwd_bwd` / `fwd_bwd_staged`, and an `info()` without the staging + capabilities, as `TrainerWorker` was at the base commit. + """ + + stage_fwd_bwd = None + fwd_bwd_staged = None + apply_deferred_update = None + + def info(self) -> datatypes.WorkerInfo: + info = super().info() + resources = { + k: v + for k, v in info.resources.items() + if k + not in ("stage_fwd_bwd", "reports_device_wait", "stage_before_update") + } + return dataclasses.replace(info, resources=resources) + + +class _NoWaitTrainer(_Trainer): + """A trainer that can hold a staged micro-batch but cannot say whether the + transfer hid behind its device (no `wait_for_device`).""" + + wait_for_device = None + + +class _UpdateToken: + """An output of one update, as MaxText's `last_update_token` (its gradient + norm): `jax.block_until_ready` on it returns once that update has landed.""" + + def __init__(self, trainer: "_TokenTrainer", lands_at: float, version: int): + self._trainer = trainer + self._lands_at = lands_at + self._version = version + + def block_until_ready(self) -> "_UpdateToken": + self._trainer._wait_until(self._lands_at) + self._trainer.log.append(("token_wait", self._version)) + return self + + +class _TokenTrainer(_Trainer): + """A trainer that can confirm its update has landed (`last_update_token`). + + So a fused last call waits for the update and says so (`update_ready`), and + the base loop defers get_metrics past the weight sync and commit + (google/tunix#2719). The double without it is the trainer of today's MaxText + pin, where the fused call returns at the update's dispatch. + """ + + def __init__(self, *args: Any, **kwargs: Any): + super().__init__(*args, **kwargs) + self.last_update_token: _UpdateToken | None = None + + def update(self, **kwargs: Any) -> int: + version = super().update(**kwargs) + assert self._update_lands_at is not None + self.last_update_token = _UpdateToken(self, self._update_lands_at, version) + return version + + +class _Reference: + """Reference policy: static log-probs from the micro-batch content.""" + + def per_token_logps(self, items, **kwargs): + del kwargs + completion = np.asarray(items.completion_ids) + return datatypes.LogprobsResponse( + request_id="ref", + per_token_logps=-(completion.astype(np.float32) % 13) / 13.0, + model_version=0, + ) + + +class _RemoteHandle(remote_execution.ActorHandle): + """Serves a worker on its own thread and event loop, like a remote server.""" + + def __init__( + self, + worker: Any, + transfer_s: Callable[[str, dict[str, Any]], float] = lambda m, k: 0.0, + ): + self._server = remote_execution.InProcessRemoteExecutionServer(worker) + self._transfer_s = transfer_s + self._loop = asyncio.new_event_loop() + self._thread = threading.Thread(target=self._loop.run_forever, daemon=True) + self._thread.start() + + async def _execute(self, request: remote_execution.ExecutionRequest) -> Any: + delay = self._transfer_s(request.method_name, request.kwargs) + if delay: + # On the server loop: a method blocking it also holds this up. + await asyncio.sleep(delay) + return await self._server.execute_request(request) + + async def asubmit(self, method_name: str | None = None, *args, **kwargs): + request = remote_execution.ExecutionRequest( + method_name=method_name, args=args, kwargs=kwargs + ) + future = asyncio.run_coroutine_threadsafe( + self._execute(request), self._loop + ) + return (await asyncio.wrap_future(future)).unwrap() + + def submit(self, method_name: str | None = None, *args, **kwargs): + request = remote_execution.ExecutionRequest( + method_name=method_name, args=args, kwargs=kwargs + ) + future = asyncio.run_coroutine_threadsafe( + self._execute(request), self._loop + ) + return future.result().unwrap() + + async def dispatch_task(self, request_id=None, method_name=None, *a, **k): + raise NotImplementedError + + async def poll_responses(self, timeout_s=0.0): + raise NotImplementedError + + def close(self) -> None: + self._loop.call_soon_threadsafe(self._loop.stop) + self._thread.join(timeout=10) + self._loop.close() + + +class _Coordinator: + """Weight-sync coordinator double: destination phases, then the source.""" + + def __init__(self, trainer: _RemoteHandle, dest_s: float, round_s: float): + self._trainer = trainer + self._dest_s = dest_s + self._round_s = round_s + self.in_flight = False + self.poisoned = None + + def has_pending_destinations(self) -> bool: + return False + + async def sync(self, policy_version, source_staged=None, only_pending=False): + del only_pending + extra = ( + {weight_sync.RELEASE_SOURCE_AFTER_STAGE: True} + if source_staged is not None + else {} + ) + request = types.SimpleNamespace(extra_config=extra) + await asyncio.sleep(self._dest_s) # bind + metadata on the rollouts + await self._trainer.asubmit("prepare_weight_sync", sync_request=request) + if source_staged is not None: + source_staged.set() + await asyncio.sleep(self._round_s) + await self._trainer.asubmit("release_weight_sync", sync_request=request) + return types.SimpleNamespace(policy_version=policy_version) + + +class _MetricsRecorder: + + def __init__(self): + self.records: list[tuple[str, str, float, int]] = [] + + def log(self, prefix, name, value, mode, step): + del mode + self.records.append((prefix, name, float(np.asarray(value)), int(step))) + + def close(self): + pass + + +class _CountingAssembler: + """Delegates to a real assembler and records each feed and flush.""" + + def __init__(self, inner: Any, log: list[Any], pack_s: float = 0.0): + self._inner = inner + self._log = log + self._pack_s = pack_s + self.num_generations = inner.num_generations + self.mini_batch_size = inner.mini_batch_size + self.pad_id = getattr(inner, "pad_id", 0) + + def __setattr__(self, name, value): + if name in ("num_generations", "mini_batch_size") and "_inner" in vars( + self + ): + setattr(self._inner, name, value) + super().__setattr__(name, value) + + def feed(self, items): + out = self._inner.feed(items) + ids = tuple( + sorted({getattr(p, "metadata", {}).get("traj_id", "") for p in items}) + ) + self._log.append(("feed", ids)) + if out and self._pack_s: + time.sleep(self._pack_s * len(out)) + return out + + def flush(self): + out = self._inner.flush() + self._log.append(("flush", len(out))) + if out and self._pack_s: + time.sleep(self._pack_s * len(out)) + return out + + def reset(self, **kwargs): + self._inner.reset(**kwargs) + + +# --------------------------------------------------------------------------- +# The reference: the loop as it was before the pipelined boundary +# --------------------------------------------------------------------------- + +# Module globals the frozen loop below refers to, bound to rl_program's own. +_extract_reward = rl_program._extract_reward +_generation_metrics = rl_program._generation_metrics +RLStepResult = rl_program.RLStepResult +Mode = rl_program.Mode + + +def _release_host_memory() -> None: + # Looked up at call time, so a test patching rl_program's is honoured. + rl_program._release_host_memory() + + +# sha256 of the base's `StandardRLProgram._train_loop` source, which the copy +# below reproduces byte for byte: atwigg/mlperf from c52127238 (the merge of +# google/tunix#2719: the fused final update and the deferred get_metrics) +# through 21b62f2bb, where `_train_loop` is unchanged. +_BASE_TRAIN_LOOP_SHA256 = ( + "37032f00118fcb4ef7d9019ac6aec1066c9f9d4ad7d8a03d9000c576c514f1dd" +) + + +class _BaseLoopProgram(rl_program.StandardRLProgram): + """The program with `_train_loop` exactly as before the pipelined boundary. + + The reference every run of the unified loop is compared with: its + `_train_loop` is atwigg/mlperf's at c52127238 (google/tunix#2719) through + 21b62f2bb, verbatim (`BaseLoopOracleTest` pins its hash). + """ + + # A verbatim copy, pinned by its hash: never reformat it. + # fmt: off + async def _train_loop(self) -> None: + assert self.engine is not None + + while self.max_steps is None or self._step < self.max_steps: + current_step = self._step + step_start_time = time.monotonic() + consumed_policy_version = self.policy_version + + uncommitted_groups = [] + step_result = None + trainer_metrics = None + step_sampler_agreement: dict[str, tuple[Any, list[float]]] = {} + step_rewards = [] + step_advantages = [] + num_microbatches = 0 + step_padding_stats: list[batch_assembly.PaddingStats] = [] + step_packing_time_sec = 0.0 + num_rollouts = 0 + all_step_items = [] + scored_items = [] + groups_consumed = 0 + checkpoint_saved = False + final_minibatch_completed = False + pending_metrics_fetch = False + policy_training_time = 0.0 + exposed_generation_time = 0.0 + weight_sync_time = 0.0 + # The end-of-step get_metrics() deferred past the weight sync and commit, + # kept out of policy_training_time because it no longer runs in training. + metrics_fetch_time = 0.0 + + current_batch_idx: int | None = None + + async def _await_pending_train() -> None: + nonlocal step_result, policy_training_time + task, self._pending_train = self._pending_train, None + if task is not None: + step_result, elapsed = await task + policy_training_time += elapsed + + async def _maybe_save_checkpoint() -> None: + nonlocal checkpoint_saved + ckpt_ts_ms = time.time_ns() // 1_000_000 + self.last_step_timestamp_ms = ckpt_ts_ms + optimizer_step = self.step + 1 + if ( + isinstance(step_result, dict) + and step_result.get("train_step") is not None + ): + optimizer_step = int(step_result["train_step"]) + if ( + self.val_start_step is not None + and optimizer_step < self.val_start_step + ): + checkpoint_saved = True + return + if isinstance( + self.scored_q, + trajectory_queue_manager.BatchOrderedQueueManager, + ): + next_batch_idx = ( + self.scored_q.next_batch_after(current_batch_idx) + if current_batch_idx is not None + else self.scored_q.next_batch_idx + ) + else: + next_batch_idx = self.step + 1 + save_kwargs = {} + interval = self.checkpoint_optimizer_interval_steps + if interval == 0 or ( + optimizer_step % interval != 0 + and (self.max_steps is None or self.step + 1 < self.max_steps) + ): + # Sent only when skipping, so an interval of 1 leaves the request, + # and trainers that predate the option, untouched. + save_kwargs["save_optimizer_state"] = False + save_resp = await self.engine.save_checkpoint( + role=datatypes.Role.ACTOR, + metadata={ + "step": optimizer_step, + "global_step": self.step + 1, + "next_batch_idx": next_batch_idx, + "policy_version": self.policy_version + 1, + "num_rollouts": num_rollouts, + "num_microbatches": num_microbatches, + }, + **save_kwargs, + ) + checkpoint_saved = True + if self.on_checkpoint_saved is not None: + # `TrainerWorker.save_checkpoint` always returns the saved path in + # `Response.metadata["checkpoint_path"]`. The callback runs in a + # worker thread (it may do blocking file/GCS I/O) and is awaited, so + # calls never overlap and its exceptions still propagate. + ckpt_info: dict[str, Any] = { + "step": optimizer_step, + "timestamp_ms": ckpt_ts_ms, + "checkpoint_path": save_resp.metadata["checkpoint_path"], + } + await asyncio.to_thread(self.on_checkpoint_saved, ckpt_info) + + while groups_consumed < self.full_batch_size: + _t_gen = time.monotonic() + if isinstance( + self.scored_q, trajectory_queue_manager.BatchOrderedQueueManager + ): + ordered = await self.scored_q.get_ordered_group( + batch_idx=current_batch_idx + ) + if ordered is None: + scored_items = [] + else: + current_batch_idx, scored_items = ordered + else: + scored_items = await self.scored_q.get_group_batch(num_groups=1) + # TODO(sanbao): Track accurate generation time by matching rollout + # requests/responses or reading from the rollout worker. Currently this + # measures time spent waiting on the scored queue (which includes + # rollout + reward calculation), as MLPerf does not strictly validate + # per-component timing accuracy. + exposed_generation_time += time.monotonic() - _t_gen + if not scored_items: + packing_start_time = time.perf_counter() + assembled_batches = list(await self._assemble(self.assembler.flush)) + step_packing_time_sec += time.perf_counter() - packing_start_time + else: + if groups_consumed == 0 and self.on_step_begin: + self.on_step_begin(current_step) + + groups_consumed += 1 + uncommitted_groups.append(scored_items) + all_step_items.extend(scored_items) + num_rollouts += len(scored_items) + for item in scored_items: + step_rewards.append(_extract_reward(item)) + payload = getattr(item, "payload", None) + if ( + item.is_valid + and payload is not None + and payload.advantages is not None + and np.asarray(payload.advantages).size > 0 + ): + step_advantages.append(float(np.mean(payload.advantages))) + + payloads = [] + for item in scored_items: + payload = getattr(item, "payload", None) + if isinstance(payload, datatypes.RLTrainerPayload): + payload = dataclasses.replace( + payload, + metadata={ + **payload.metadata, + "traj_id": item.traj_id, + }, + ) + payloads.append(payload) + packing_start_time = time.perf_counter() + assembled_batches = list( + await self._assemble( + self.assembler.feed, + payloads, # pyrefly: ignore[bad-argument-type] + ) + ) + step_packing_time_sec += time.perf_counter() - packing_start_time + # `self.assembler` has now copied or packed `routed_experts`; release + # the per-trajectory references on `payloads` and `scored_items` + # (which remain in `all_step_items` until step-end metric logging) so + # unbatched routing arrays do not coexist with packed microbatches and + # weight sync. Lightweight 1D mask fields on `item.payload` are kept + # for `_collect_and_log_step_metrics`. + payloads.clear() + for item in scored_items: + item.routed_experts = None + if item.metadata is not None: + item.metadata.pop("routed_experts", None) + if isinstance(item.traj, dict): + item.traj.pop("routed_experts", None) + item_payload = getattr(item, "payload", None) + if ( + isinstance(item_payload, datatypes.RLTrainerPayload) + and item_payload.routed_experts is not None + ): + item.payload = dataclasses.replace( + item_payload, routed_experts=None + ) + + while assembled_batches: + mb = assembled_batches.pop(0) + is_final_batch = mb.is_final_batch + batch = mb.payload + if getattr(self.algo, "requires_reference_kl", False): + if not isinstance(batch, datatypes.RLTrainerPayload): + raise TypeError( + "Reference KL requires an assembler that returns " + "datatypes.RLTrainerPayload microbatches; got " + f"{type(batch).__name__}." + ) + # The reference model is static and on its own worker, so this + # overlaps the in-flight actor train_step. + ref_logps = await self.engine.per_token_logps( + datatypes.Role.REFERENCE, items=batch + ) + batch = batch_assembly.with_ref_per_token_logps(batch, ref_logps) + algo_config = getattr(self.algo, "algo_config", None) + if ( + isinstance(batch, datatypes.RLTrainerPayload) + and batch.old_per_token_logps is not None + and algo_config is not None + and algo_config.use_rollout_logps + ): + await _await_pending_train() + batch = await self._apply_sampler_trainer_agreement( + batch, step_sampler_agreement + ) + + num_microbatches += 1 + if mb.padding_stats is not None: + step_padding_stats.append(mb.padding_stats) + logging.info( + "Packed %d trajectories into microbatch: %s (padding_ratio=%.3f," + " row_imbalance=%.3f)", + len(mb.trajectory_ids), + logging_utils.summarize_list(list(mb.trajectory_ids)), + mb.padding_stats.padding_ratio, + mb.padding_stats.row_imbalance, + ) + else: + logging.info( + "Packed %d trajectories into microbatch: %s", + len(mb.trajectory_ids), + logging_utils.summarize_list(list(mb.trajectory_ids)), + ) + mb = None + if not self._first_train_payload_logged: + self._first_train_payload_logged = True + logging.info( + "First trainer micro-batch: %s", + batch_assembly.describe_payload(batch), + ) + await _await_pending_train() + train_step = self._timed_train_step( + batch, apply_optimizer=is_final_batch + ) + batch = None + if self.pipeline_train_microbatches and not is_final_batch: + self._pending_train = asyncio.create_task(train_step) + continue + step_result, elapsed = await train_step + policy_training_time += elapsed + if is_final_batch: + # Save only at a resumable full-batch boundary. An optimizer step + # can occur earlier when a full batch contains multiple mini + # batches, but the dataset resume cursor advances in full batches. + full_batch_complete = ( + groups_consumed >= self.full_batch_size or not scored_items + ) + # The checkpoint is stamped as the weight-update time and saved + # next, so the update must have landed on device first. Unless the + # trainer has said it did (`update_ready`), fetching the metrics is + # what waits for it, so the fetch stays ahead of the save, counted + # as training time, as with a separate update RPC. + pending_metrics_fetch = full_batch_complete and bool( + isinstance(step_result, dict) + and step_result.get("update_ready") + ) + if not pending_metrics_fetch: + _t_metrics = time.monotonic() + trainer_metrics = await self.engine.get_metrics( + role=datatypes.Role.ACTOR + ) + policy_training_time += time.monotonic() - _t_metrics + final_minibatch_completed = True + if self.sync_weights and self.async_weight_sync: + self._round_due = True + # TODO(tunix-dev): Configurable checkpointing frequency. Today we + # checkpoint at the same frequency as the weight update. + # TODO(tunix-dev): For now any failures in save_checkpoint will + # abort the entire program. Make it configurable on whether to fail + # or continue. + if full_batch_complete: + await _maybe_save_checkpoint() + + if not scored_items: + if not checkpoint_saved and final_minibatch_completed: + await _maybe_save_checkpoint() + break + + await _await_pending_train() + if not all_step_items: + logging.info( + "Dataset exhausted at step %d before max_steps.", current_step + ) + break + + if self.sync_weights and self.async_weight_sync: + # Only the time the trainer is actually held: the previous round, if + # it outlived this step, plus this round's source snapshot. + _t_sync = time.monotonic() + await self._start_background_sync() + weight_sync_time = time.monotonic() - _t_sync + elif self.sync_weights: + _t_sync = time.monotonic() + new_version = await self.engine.sync_weights(role=datatypes.Role.ACTOR) + weight_sync_time = time.monotonic() - _t_sync + self.policy_version = ( + new_version if new_version is not None else self.policy_version + 1 + ) + + # Before `commit()`, which will eventually take ownership of the groups. + generation_metrics = _generation_metrics(uncommitted_groups) + logging.info( + "[pipeline] COMMIT step=%d batch_idx=%s consumed_version=%d" + " rollout_versions=%s groups=%s", + current_step, + current_batch_idx, + consumed_policy_version, + [getattr(g[0], "policy_version", 0) for g in uncommitted_groups if g], + [ + ( + getattr(g[0], "prompt_id", ""), + (getattr(g[0], "metadata", None) or {}).get("batch_idx"), + ) + for g in uncommitted_groups + if g + ], + ) + if isinstance( + self.scored_q, trajectory_queue_manager.BatchOrderedQueueManager + ): + assert current_batch_idx is not None + self.scored_q.commit_batch(current_batch_idx) + else: + self.scored_q.commit(current_step, groups=uncommitted_groups) + # The commit may have opened the dispatch window. Yield so the + # dispatcher refills the just-synced rollout workers now, rather than + # after this step's bookkeeping and the next step's first packing. + await asyncio.sleep(0) + + if pending_metrics_fetch: + _t_metrics = time.monotonic() + trainer_metrics = await self.engine.get_metrics( + role=datatypes.Role.ACTOR + ) + metrics_fetch_time += time.monotonic() - _t_metrics + + step_time_sec = time.monotonic() - step_start_time + + filtered_groups = await self._drain_filtered_groups() + + metrics_summary = self._collect_and_log_step_metrics( + all_step_items=all_step_items, + step_rewards=step_rewards, + step_advantages=step_advantages, + generation_metrics=generation_metrics, + step_result=step_result, + trainer_metrics=trainer_metrics, + num_rollouts=num_rollouts, + num_microbatches=num_microbatches, + padding_stats=step_padding_stats, + packing_time_sec=step_packing_time_sec, + step_time_sec=step_time_sec, + consumed_policy_version=consumed_policy_version, + log_step=current_step, + sampler_agreement=step_sampler_agreement, + policy_training_time=policy_training_time, + exposed_generation_time=exposed_generation_time, + weight_sync_time=weight_sync_time, + metrics_fetch_time=metrics_fetch_time, + filtered_groups=filtered_groups, + ) + self._log_consumed_trajectories( + all_step_items, + log_step=current_step, + consumed_policy_version=consumed_policy_version, + ) + filtered_groups.clear() + all_step_items.clear() + uncommitted_groups.clear() + scored_items.clear() + + self.last_step_result = RLStepResult( + step=current_step, + policy_version=self.policy_version, + num_rollouts=num_rollouts, + num_microbatches=num_microbatches, + reward_mean=metrics_summary["reward_mean"], + reward_std=metrics_summary["reward_std"], + advantage_mean=metrics_summary["advantage_mean"], + advantage_std=metrics_summary["advantage_std"], + train_result=step_result, + ) + + loss_val = metrics_summary["loss_val"] + perplexity_val = metrics_summary["perplexity_val"] + grad_norm_val = metrics_summary["grad_norm_val"] + if self.mode == Mode.TRAIN: + logging.info( + "Train step %d - loss: %s - reward_mean: %.4f - advantage_mean:" + " %.4f - grad_norm: %s - perplexity: %s - step_time: %.2fs", + current_step, + f"{loss_val:.4f}" if loss_val is not None else "N/A", + metrics_summary["reward_mean"], + metrics_summary["advantage_mean"], + f"{grad_norm_val:.4g}" if grad_norm_val is not None else "N/A", + f"{perplexity_val:.4f}" if perplexity_val is not None else "N/A", + step_time_sec, + ) + + if self.on_step_end: + self.on_step_end(current_step, step_result) + self._step += 1 + # Strictly after the increment: the dispatcher re-reads `_next_batch` on + # waking, and waking it on the old value would park it again with no one + # left to set the event. + self._release_window() + _release_host_memory() + # fmt: on + + +# --------------------------------------------------------------------------- +# Running one scenario +# --------------------------------------------------------------------------- + + +@dataclasses.dataclass +class _Scenario: + order: str = "arrival" + num_prompts: int = 12 # groups available in the queue + full_batch_size: int = 4 + mini_batch_size: int = 4 + max_steps: int | None = 3 + async_sync: bool = True + sync_weights: bool = True + pipeline_microbatches: bool = True + logps: bool = False # use_rollout_logps with old logps: rescored per mb + reference_kl: bool = False + val_start_step: int | None = None + # Seconds between groups reaching the queue; 0 puts them all up front. + arrival_gap_s: float = 0.0 + # Batches (PROMPT_BATCH) declared shorter than full: batch_idx -> groups. + short_batches: dict[int, int] = dataclasses.field(default_factory=dict) + costs: _Costs = dataclasses.field(default_factory=_Costs) + dest_s: float = 0.0 + round_s: float = 0.0 + pack_s: float = 0.0 + tail_s: float = 0.0 # on_step_end blocks the orchestrator this long + + +@dataclasses.dataclass +class _Run: + log: list[Any] + metrics: list[tuple[str, str, float, int]] + checkpoints: list[dict[str, Any]] + last_step_result: Any + trainer: _Trainer + program: rl_program.StandardRLProgram + + +def _new_program( + sc: _Scenario, + *, + pipelined: bool, + staged: bool, + program_cls: type[rl_program.StandardRLProgram], + log: list[Any], + checkpoints: list[dict[str, Any]], + program_kwargs: dict[str, Any] | None = None, +) -> rl_program.StandardRLProgram: + """The program for `sc`; `program_kwargs` override or add constructor args.""" + # The program refuses use_rollout_logps=False with several updates per + # step. Without old logps in the payloads nothing is rescored either way. + use_rollout_logps = sc.logps or sc.full_batch_size != sc.mini_batch_size + algo = types.SimpleNamespace( + num_generations=NUM_GENERATIONS, + mini_batch_size=sc.mini_batch_size, + train_micro_batch_size=1, + max_response_length=1024, + requires_reference_kl=sc.reference_kl, + use_rollout_logps=use_rollout_logps, + algo_config=types.SimpleNamespace( + temperature=None, + use_rollout_logps=use_rollout_logps, + sampler_is=None, + # Set, so rescoring writes the trainer's logps into the batch: a + # micro-batch rescored on the wrong weights then trains differently. + seq_logprob_error_threshold=100.0 if sc.logps else None, + ), + ) + inner = batch_assembly.SequencePackedBatchAssembler( + batch_size=2, + num_generations=NUM_GENERATIONS, + mini_batch_size=sc.mini_batch_size, + max_packed_len=ROW_TOKENS, + segment_align_multiple=1, + ) + + def _on_step_end(step, result): + del result + log.append(("on_step_end", step)) + if sc.tail_s: + time.sleep(sc.tail_s) + + kwargs: dict[str, Any] = dict( + algo=algo, + dataset=None, + max_steps=sc.max_steps, + assembler=_CountingAssembler(inner, log, sc.pack_s), + batch_size=sc.full_batch_size, + max_staleness=0, + sync_weights=sc.sync_weights, + async_weight_sync=sc.async_sync, + pipeline_train_microbatches=sc.pipeline_microbatches, + group_order=sc.order, + val_start_step=sc.val_start_step, + on_step_begin=lambda step: log.append(("on_step_begin", step)), + on_step_end=_on_step_end, + on_checkpoint_saved=checkpoints.append, + pipeline_step_boundary=pipelined, + stage_next_microbatch=staged, + # This file tests the stage after the step's last call (the hook for a + # trainer that cannot defer its update, and `--no-stage_before_update`). + # The default, staging between the last micro-batch and a deferred + # update, runs these scenarios in + # rl_program_stage_before_update_test.py. + stage_before_update=False, + ) + kwargs.update(program_kwargs or {}) + program = program_cls(**kwargs) + program.metrics_logger = _MetricsRecorder() # pyrefly: ignore + return program + + +async def _feed_queue( + program: rl_program.StandardRLProgram, sc: _Scenario +) -> None: + groups = _groups( + sc.num_prompts, full_batch_size=sc.full_batch_size, logps=sc.logps + ) + if sc.order == "prompt_batch": + for batch_idx, num_groups in sc.short_batches.items(): + program.scored_q.expect(batch_idx, num_groups) + for idx, group in enumerate(groups): + batch_idx = idx // sc.full_batch_size + intra = idx % sc.full_batch_size + if intra >= sc.short_batches.get(batch_idx, sc.full_batch_size): + continue + for item in group: + await program.scored_q.put(item) + if sc.arrival_gap_s: + await asyncio.sleep(sc.arrival_gap_s) + await program.scored_q.close() + + +def _run( + sc: _Scenario, + *, + pipelined: bool, + staged: bool = False, + program_cls: type[ + rl_program.StandardRLProgram + ] = rl_program.StandardRLProgram, + errors: list[BaseException] | None = None, + worker_cls: type[_Worker] = _Worker, + trainer_cls: type[_Trainer] = _Trainer, + program_kwargs: dict[str, Any] | None = None, +) -> _Run: + """Runs `sc` to the end; with `errors`, records a failure instead of raising.""" + log: list[Any] = [] + checkpoints: list[dict[str, Any]] = [] + trainer = trainer_cls(log, sc.costs) + host_s = sc.costs.host_s + + def transfer_s(method: str, kwargs: dict[str, Any]) -> float: + del kwargs + # Only the calls carrying a micro-batch move data. + return host_s if method in ("fwd_bwd", "stage_fwd_bwd") else 0.0 + + worker = worker_cls(lambda: trainer) + worker.initialize() + handle = _RemoteHandle(worker, transfer_s) + ref_handle = _RemoteHandle(_Reference()) + coordinator = _Coordinator(handle, sc.dest_s, sc.round_s) + engine = distributed_rl_engine.DistributedRLEngine( + rollout_workers=[], + trainer_workers={datatypes.Role.ACTOR: handle}, + inference_workers={datatypes.Role.REFERENCE: ref_handle}, + weight_sync_coordinator=coordinator, + ) + program = _new_program( + sc, + pipelined=pipelined, + staged=staged, + program_cls=program_cls, + log=log, + checkpoints=checkpoints, + program_kwargs=program_kwargs, + ) + + async def _main(): + program.engine = engine + feeder = asyncio.create_task(_feed_queue(program, sc)) + try: + await asyncio.wait_for(program.train_stage(), timeout=120) + finally: + feeder.cancel() + + try: + asyncio.run(_main()) + except Exception as exc: # pylint: disable=broad-exception-caught + if errors is None: + raise + errors.append(exc) + finally: + handle.close() + ref_handle.close() + return _Run( + log=log, + metrics=program.metrics_logger.records, # pyrefly: ignore + checkpoints=checkpoints, + last_step_result=program.last_step_result, + trainer=trainer, + program=program, + ) + + +# Logged under their timing names; these measure the change, so they differ. +_TIMING_METRICS = ( + "step_time", + "global_step_time", + "policy_training_time", + "exposed_generation_time", + "weight_sync_time", + # google/tunix#2719's deferred fetch; always 0 without a confirmed update. + "metrics_fetch_time", + "time_sec_total", + "time_sec_mean", +) +# Program-side events whose place relative to trainer calls is what changes, +# a data-only transfer, the background round's release (timed by the round, +# not the program), and a wait for an update that already ran (by the fused +# call, or by the stage right after it). +_NOT_TRAINER_STATE = ( + "on_step_end", + "on_step_begin", + "feed", + "flush", + "stage", + "release_weight_sync", + "token_wait", +) + + +def _trainer_calls(log: Sequence[Any]) -> list[Any]: + return [e for e in log if e[0] not in _NOT_TRAINER_STATE] + + +def _data_metrics(records): + return [r for r in records if not any(t in r[1] for t in _TIMING_METRICS)] + + +def _assert_equivalent(test: absltest.TestCase, ref: _Run, new: _Run) -> None: + """The trainer, the checkpoints and the metrics cannot tell the runs apart.""" + test.assertEqual(_trainer_calls(ref.log), _trainer_calls(new.log)) + _assert_same_outputs(test, ref, new) + + +def _assert_equivalent_but_metrics_fetch( + test: absltest.TestCase, ref: _Run, new: _Run +) -> None: + """`_assert_equivalent`, except where get_metrics sits within a boundary. + + The same calls otherwise, the same fetches (each returning the same step's + buffer, one buffer each), each at its boundary, and the same outputs. + """ + test.assertEqual( + _without_metrics_fetch(ref.log), _without_metrics_fetch(new.log) + ) + test.assertEqual( + [e for e in ref.log if e[0] == "get_metrics"], + [e for e in new.log if e[0] == "get_metrics"], + ) + _assert_metrics_fetched_at_the_boundary(test, new) + _assert_same_outputs(test, ref, new) + + +def _assert_same_outputs(test: absltest.TestCase, ref: _Run, new: _Run) -> None: + """The same checkpoints, logged metric values and last step result.""" + test.assertEqual( + [(c["step"], c["checkpoint_path"]) for c in ref.checkpoints], + [(c["step"], c["checkpoint_path"]) for c in new.checkpoints], + ) + test.assertEqual(_data_metrics(ref.metrics), _data_metrics(new.metrics)) + a, b = ref.last_step_result, new.last_step_result + test.assertEqual( + (a.step, a.policy_version, a.num_rollouts, a.num_microbatches), + (b.step, b.policy_version, b.num_rollouts, b.num_microbatches), + ) + test.assertEqual((a.reward_mean, a.reward_std), (b.reward_mean, b.reward_std)) + test.assertEqual( + (a.advantage_mean, a.advantage_std), (b.advantage_mean, b.advantage_std) + ) + + +def _check_equivalent(ref: _Run, new: _Run) -> bool: + """`_assert_equivalent` as a predicate, for the negative controls.""" + probe = absltest.TestCase() + try: + _assert_equivalent(probe, ref, new) + except AssertionError: + return False + return True + + +def _microbatches_per_step(run: _Run) -> list[int]: + """Micro-batches each logged step trained, from its own step metrics.""" + per_step = { + r[3]: int(r[2]) + for r in run.metrics + if r[0] == "orchestrator" and r[1] == "num_microbatches" + } + return [per_step[k] for k in sorted(per_step)] + + +def _boundaries(run: _Run) -> list[tuple[bool, bool]]: + """Per step boundary k: (step k+1's first micro-batch reached the trainer + before step k's on_step_end, that micro-batch was staged ahead). + + Step k+1's first micro-batch is found by counting each step's micro-batches, + so this holds for any number of updates per step. + """ + log = run.log + fwd = [i for i, e in enumerate(log) if e[0] == "fwd_bwd"] + out = [] + first_of_next = 0 + for k, n in enumerate(_microbatches_per_step(run)): + first_of_next += n + if first_of_next >= len(fwd) or ("on_step_end", k) not in log: + break + first = fwd[first_of_next] + end = log.index(("on_step_end", k)) + update = max(i for i in range(first) if log[i][0] == "update") + staged = any( + e[0] == "stage" and e[1] == log[first][1] for e in log[update:first] + ) + out.append((first < end, staged)) + return out + + +# What the program, not the trainer, does; each comes from one thread at a +# time in program order, so their sequence is deterministic. +_PROGRAM_EVENTS = ("feed", "flush", "on_step_begin", "on_step_end") + + +def _program_events(log: Sequence[Any]) -> list[Any]: + return [e for e in log if e[0] in _PROGRAM_EVENTS] + + +def _assert_is_the_base_loop( + test: absltest.TestCase, ref: _Run, new: _Run +) -> None: + """Flags off: the old loop exactly -- calls, program events, metrics.""" + _assert_equivalent(test, ref, new) + test.assertEqual(_program_events(ref.log), _program_events(new.log)) + # Every metric, timing ones included, under the same name at the same step. + test.assertEqual( + [(r[0], r[1], r[3]) for r in ref.metrics], + [(r[0], r[1], r[3]) for r in new.metrics], + ) + test.assertNotIn("stage", [e[0] for e in new.log]) + test.assertEqual(_boundaries(new), [(False, False)] * len(_boundaries(new))) + test.assertIsNone(new.program._intake_task) + + +_REFERENCE_RUNS: dict[tuple[str, type[_Trainer]], _Run] = {} + + +def _reference(name: str, trainer_cls: type[_Trainer] = _Trainer) -> _Run: + """The base loop's run of `_SCENARIOS[name]`, computed once per trainer.""" + key = (name, trainer_cls) + if key not in _REFERENCE_RUNS: + _REFERENCE_RUNS[key] = _run( + _SCENARIOS[name], + pipelined=False, + program_cls=_BaseLoopProgram, + trainer_cls=trainer_cls, + ) + return _REFERENCE_RUNS[key] + + +def _without_metrics_fetch(log: Sequence[Any]) -> list[Any]: + """Trainer calls without get_metrics, and each save without the metrics + buffers still pending on the trainer. + + What may differ when a stage was planned but found nothing packed in time + against a trainer that confirms its updates: that step's update is then not + confirmed, so get_metrics() is its barrier again, ahead of the save (as + without a confirmed update), instead of after the commit. MaxText writes + the same completed-step checkpoint either way: it serializes undrained + metric buffers only for an intra-step checkpoint (`save_checkpoint`, + maxtext 853d655). + """ + return [ + e[:5] if e[0] == "save" else e + for e in _trainer_calls(log) + if e[0] != "get_metrics" + ] + + +def _assert_metrics_fetched_at_the_boundary( + test: absltest.TestCase, run: _Run +) -> None: + """Each get_metrics returns step k's buffer after update k, before step k+1. + + The rule the deferral and the pipelined boundary must both keep: MaxText + returns only the newest step buffer, so the fetch must follow the update + that closes it and precede the next step's first micro-batch. + """ + calls = _trainer_calls(run.log) + fetches = 0 + for i, event in enumerate(calls): + if event[0] != "get_metrics": + continue + fetches += 1 + step = event[1] + updates = [e for e in calls[:i] if e[0] == "update"] + test.assertTrue(updates and updates[-1][1] == step + 1, (i, calls)) + test.assertFalse( + any(e[0] == "fwd_bwd" and e[2] == step + 1 for e in calls[:i]), + (i, calls), + ) + # One fetch per update (a run that trained nothing fetches nothing). + test.assertEqual(fetches, sum(1 for e in calls if e[0] == "update"), calls) + + +def _fetches_after_sync(run: _Run) -> list[bool]: + """Per get_metrics: whether a weight sync staged since the update before it. + + True is the base loop's deferred fetch (google/tunix#2719), taken when the + trainer confirmed the update; False is the fetch ahead of the save. + """ + calls = _trainer_calls(run.log) + out = [] + for i, event in enumerate(calls): + if event[0] != "get_metrics": + continue + last_update = max(j for j in range(i) if calls[j][0] == "update") + out.append(any(e[0] == "prepare_weight_sync" for e in calls[last_update:i])) + return out + + +def _staged_boundaries_confirm_the_update( + test: absltest.TestCase, run: _Run +) -> int: + """At each stage: the last call returned at the update's dispatch, and the + stage waited for that update. Returns the number of stages.""" + log = run.log + stages = [i for i, e in enumerate(log) if e[0] == "stage"] + for s in stages: + update = max(i for i in range(s) if log[i][0] == "update") + version = log[update][1] + test.assertNotIn( + ("token_wait", version), log[update:s], f"waited before the stage {log}" + ) + nxt = next( + ( + i + for i in range(s + 1, len(log)) + if log[i][0] not in _NOT_TRAINER_STATE + ), + len(log), + ) + test.assertIn(("token_wait", version), log[s:nxt], log) + return len(stages) + + +def _assert_weights_follow_updates(test: absltest.TestCase, run: _Run) -> None: + """Each micro-batch of step k is computed on the weights after update k.""" + version = 0 + saw = 0 + for event in run.log: + if event[0] == "update": + version = event[1] + elif event[0] == "fwd_bwd": + test.assertEqual( + event[2], version, f"fwd_bwd {event[1]} on stale weights" + ) + saw += 1 + test.assertGreater(saw, 0) + + +# --------------------------------------------------------------------------- +# Broken programs: each must be caught by a check below. +# --------------------------------------------------------------------------- + + +class _SendsBeforeBoundary(rl_program.StandardRLProgram): + """WRONG: sends step k+1's first micro-batch right after step k's update. + + That is, right after step k's last train call, which runs the update: before + step k's metrics fetch, checkpoint and weight-sync staging, so step k's + metrics buffer is clobbered and the sync and checkpoint run with step k+1's + gradients half accumulated. It hooks in where a stage would, at every + boundary the first version of this loop offered a stage at (a full batch, + the flag, a step taken ahead), so it is not narrowed by `_stage_planned`. + """ + + async def _stage_planned(self, st, *, rescored): + del rescored + return self.stage_next_microbatch and st is not None + + async def _maybe_stage_next_microbatch(self, st, step_result): + task = await rl_program.StandardRLProgram._send_first_microbatch_ahead( + self, st + ) + if task is not None: + await asyncio.wait({task}) + + async def _send_first_microbatch_ahead(self, st): + return None + + +class _PrepacksWrongBatch(rl_program.StandardRLProgram): + """WRONG: takes the next step's groups from one batch too far ahead.""" + + async def _prepack(self, st): + if st.groups_consumed == 0 and st.current_batch_idx is not None: + st.current_batch_idx += 1 + await super()._prepack(st) + + +class _NeverLatches(rl_program.StandardRLProgram): + """Staging that never switches itself off (every stage counts as hidden).""" + + def _staging_hidden(self, device_wait_s): + return True + + +class _MinGapLatch(rl_program.StandardRLProgram): + """WRONG, the latch as first built. + + Stops staging once a stage takes 0.9x the SHORTEST gap between the step's + train calls returning. A step starts on a device the weight sync drained and + MaxText queues two computations, so with three or more calls per step that + gap is the transfer time, not the device time, and it latches off although + the trainer is device-bound. + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._returns: list[float] = [] + self.cadences: list[float | None] = [] + + async def _timed_train_step(self, batch, **kwargs): + result = await super()._timed_train_step(batch, **kwargs) + self._returns.append(time.monotonic()) + return result + + def _staging_hidden(self, device_wait_s): + return True # Only the min-gap rule below decides. + + async def _maybe_stage_next_microbatch(self, st, step_result): + returns, self._returns = sorted(self._returns), [] + gaps = [b - a for a, b in zip(returns, returns[1:])] + cadence = min(gaps) if gaps else None + staged_before = st.staged is not None + start = time.monotonic() + update_ready = await super()._maybe_stage_next_microbatch(st, step_result) + stage_s = time.monotonic() - start + if st.staged is not None and not staged_before: + self.cadences.append(cadence) + if cadence is not None and stage_s >= 0.9 * cadence: + self._staging_latched_off = True + return update_ready + + +class _StampsBeforeUpdateLands(rl_program.StandardRLProgram): + """WRONG: stamps the checkpoint when the update is dispatched, not landed. + + It takes the stamp when the step's last train call returns. That call runs + the update; against a trainer that cannot confirm the update (no + `last_update_token`) it returns at the update's dispatch. + """ + + async def _timed_train_step(self, batch, **kwargs): + result = await super()._timed_train_step(batch, **kwargs) + if kwargs.get("apply_optimizer"): + self._early_stamp_ms = time.time_ns() // 1_000_000 + return result + + async def _save_step_checkpoint(self, st, step_result): + real_time_ns = time.time_ns + stamp = getattr(self, "_early_stamp_ms", None) + if stamp is None: + return await super()._save_step_checkpoint(st, step_result) + try: + time.time_ns = lambda: stamp * 1_000_000 + return await super()._save_step_checkpoint(st, step_result) + finally: + time.time_ns = real_time_ns + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + +# The trainer-bound regime: a micro-batch's device time well above the time +# to pack one (tr-s3: 15.11 s against ~2.1 s), so step k+1's first micro-batch +# is packed, and staged, before step k's boundary ends and the pipelined paths +# are the ones taken. (`zero_costs` and `groups_arrive_late` cover the rest.) +_BUSY = _Costs(device_s=0.15, update_s=0.01, host_s=0.02) + +_SCENARIOS = dict( + arrival=_Scenario(order="arrival", costs=_BUSY), + prompt_batch=_Scenario(order="prompt_batch", costs=_BUSY), + zero_costs=_Scenario(order="prompt_batch"), + no_microbatch_pipelining=_Scenario( + order="prompt_batch", pipeline_microbatches=False, costs=_BUSY + ), + blocking_sync=_Scenario( + order="prompt_batch", async_sync=False, costs=_BUSY + ), + no_sync=_Scenario(order="arrival", sync_weights=False, costs=_BUSY), + two_updates_per_step=_Scenario( + order="prompt_batch", + full_batch_size=4, + mini_batch_size=2, + costs=_BUSY, + ), + checkpoints_every_step=_Scenario( + order="prompt_batch", val_start_step=1, max_steps=3, costs=_BUSY + ), + # Steps 0-1 unscored, steps 2-3 write scored checkpoints (MLPerf-like). + checkpoints_from_step_3=_Scenario( + order="prompt_batch", + num_prompts=16, + val_start_step=3, + max_steps=4, + costs=_BUSY, + ), + short_last_batch=_Scenario( + order="prompt_batch", + num_prompts=12, + short_batches={2: 3}, + val_start_step=2, + costs=_BUSY, + ), + dataset_runs_out=_Scenario( + order="arrival", num_prompts=10, max_steps=None, costs=_BUSY + ), + reference_kl=_Scenario( + order="prompt_batch", reference_kl=True, costs=_BUSY + ), + rescored_on_live_weights=_Scenario( + order="prompt_batch", logps=True, costs=_BUSY + ), + groups_arrive_late=_Scenario(order="prompt_batch", arrival_gap_s=0.03), +) + +# Scenarios in which step k+1's first micro-batch is ready ahead of time and +# is neither its step's last nor rescored, so it is sent (and staged) early. +_SENT_AHEAD = ( + "arrival", + "prompt_batch", + "no_microbatch_pipelining", + "blocking_sync", + "no_sync", + "checkpoints_every_step", + "checkpoints_from_step_3", + "reference_kl", +) + + +def _writes_scored_checkpoint(sc: _Scenario, step: int) -> bool: + """Step `step` (0-based, one update per step) writes a scored checkpoint.""" + return sc.val_start_step is not None and step + 1 >= sc.val_start_step + + +class BaseLoopOracleTest(absltest.TestCase): + """The reference really is the loop as it was before this change.""" + + def test_the_frozen_loop_is_the_base_commits_byte_for_byte(self): + source = inspect.getsource(_BaseLoopProgram._train_loop) + self.assertEqual( + hashlib.sha256(source.encode()).hexdigest(), _BASE_TRAIN_LOOP_SHA256 + ) + + def test_the_pin_catches_an_edited_copy(self): + source = inspect.getsource(_BaseLoopProgram._train_loop) + edited = source.replace("groups_consumed += 1", "groups_consumed += 2", 1) + self.assertNotEqual(edited, source) + self.assertNotEqual( + hashlib.sha256(edited.encode()).hexdigest(), _BASE_TRAIN_LOOP_SHA256 + ) + + def test_the_unified_loop_is_not_the_frozen_one(self): + # Otherwise every comparison below would be a loop against itself. + self.assertIsNot( + _BaseLoopProgram._train_loop, rl_program.StandardRLProgram._train_loop + ) + self.assertNotEqual( + inspect.getsource(_BaseLoopProgram._train_loop), + inspect.getsource(rl_program.StandardRLProgram._train_loop), + ) + + +_MODES = { + "off": dict(pipelined=False, staged=False), + "unstaged": dict(pipelined=True, staged=False), + "staged": dict(pipelined=True, staged=True), +} + + +class PipelinedBoundaryEquivalenceTest(parameterized.TestCase): + """The unified loop trains exactly what the base loop trains.""" + + @parameterized.named_parameters( + (f"{name}_{mode}", name, mode) for name in _SCENARIOS for mode in _MODES + ) + def test_trainer_cannot_tell_it_from_the_base_loop(self, name, mode): + sc = _SCENARIOS[name] + ref = _reference(name) + new = _run(sc, **_MODES[mode]) + staged = _MODES[mode]["staged"] + _assert_equivalent(self, ref, new) + _assert_weights_follow_updates(self, new) + # One set of step metrics per step, from one get_metrics that returned a + # single (that step's) buffer. + steps = [r[3] for r in new.metrics if r[1] == "mean" and r[0] == "rewards"] + self.assertEqual(steps, sorted(set(steps))) + self.assertTrue( + all(e[2] == 1 for e in new.log if e[0] == "get_metrics"), new.log + ) + self.assertIsNone(new.program._pending_train) + self.assertTrue( + new.program._intake_task is None or new.program._intake_task.done() + ) + if mode == "off": + # The default path is the old loop, program events and all. + _assert_is_the_base_loop(self, ref, new) + return + if name in _SENT_AHEAD: + # The pipelined path really ran, at every boundary; the reference never + # sends a step's micro-batch before the previous step's bookkeeping. + # Staging is skipped after a step that writes a scored checkpoint. + new_b, ref_b = _boundaries(new), _boundaries(ref) + self.assertNotEmpty(new_b) + expected = [ + (True, staged and not _writes_scored_checkpoint(sc, k)) + for k in range(len(new_b)) + ] + self.assertEqual(new_b, expected, new.log) + self.assertEqual(ref_b, [(False, False)] * len(ref_b)) + if not staged: + self.assertNotIn("stage", [e[0] for e in new.log]) + + +class PipelinedBoundaryChecksFailOnWrongProgramsTest(absltest.TestCase): + """Each gate above is shown to fire on a deliberately wrong input.""" + + def test_equivalence_check_catches_sending_before_the_boundary(self): + sc = _SCENARIOS["prompt_batch"] + ref = _reference("prompt_batch") + bad = _run( + sc, pipelined=True, staged=True, program_cls=_SendsBeforeBoundary + ) + self.assertFalse(_check_equivalent(ref, bad)) + # And how: step k+1's first fwd_bwd reached the trainer before step k's + # get_metrics, which then returned (and dropped) the wrong buffers. + calls = _trainer_calls(bad.log) + first_update = next(i for i, e in enumerate(calls) if e[0] == "update") + self.assertEqual(calls[first_update + 1][0], "fwd_bwd") + self.assertTrue(any(e[0] == "get_metrics" and e[2] > 1 for e in calls)) + + def test_equivalence_check_catches_prepacking_the_wrong_batch(self): + sc = dataclasses.replace(_SCENARIOS["prompt_batch"], num_prompts=20) + ref = _run(sc, pipelined=False, program_cls=_BaseLoopProgram) + bad = _run(sc, pipelined=True, program_cls=_PrepacksWrongBatch) + self.assertFalse(_check_equivalent(ref, bad)) + + def test_equivalence_check_catches_one_swapped_trainer_call(self): + sc = _SCENARIOS["arrival"] + ref = _reference("arrival") + new = _run(sc, pipelined=True) + self.assertTrue(_check_equivalent(ref, new)) + calls = new.log + i = next(k for k, e in enumerate(calls) if e[0] == "update") + j = next(k for k, e in enumerate(calls) if k > i and e[0] == "fwd_bwd") + swapped = list(calls) + swapped[i], swapped[j] = swapped[j], swapped[i] + self.assertFalse( + _check_equivalent(ref, dataclasses.replace(new, log=swapped)) + ) + + def test_weights_check_catches_a_stale_micro_batch(self): + sc = _SCENARIOS["arrival"] + run = _run(sc, pipelined=True) + stale = [ + (e[0], e[1], e[2] - 1) if e[0] == "fwd_bwd" and e[2] > 0 else e + for e in run.log + ] + with self.assertRaises(AssertionError): + _assert_weights_follow_updates(self, dataclasses.replace(run, log=stale)) + + def test_base_loop_check_catches_a_reordered_boundary(self): + # The pipelined mode trains the same thing but moves the next step's + # feeds ahead of this step's on_step_end: the strict flags-off check must + # see that, or it would pass a default path that silently pipelines. + sc = _SCENARIOS["prompt_batch"] + ref = _reference("prompt_batch") + pipelined = _run(sc, pipelined=True) + self.assertTrue(_check_equivalent(ref, pipelined)) + probe = absltest.TestCase() + with self.assertRaises(AssertionError): + _assert_is_the_base_loop(probe, ref, pipelined) + self.assertNotEqual( + _program_events(ref.log), _program_events(pipelined.log) + ) + + +class PipelinedBoundaryOrderingTest(absltest.TestCase): + """What the pipelined loop does change: when the orchestrator works.""" + + def _index(self, log, pred): + return next(i for i, e in enumerate(log) if pred(e)) + + def test_next_step_is_sent_before_this_steps_bookkeeping(self): + sc = _SCENARIOS["prompt_batch"] + for pipelined, expect_before in ((False, False), (True, True)): + run = _run(sc, pipelined=pipelined) + log = run.log + # Step 0's first fwd_bwd after its sync staged belongs to step 1. + staged_0 = self._index(log, lambda e: e[0] == "prepare_weight_sync") + first_1 = ( + self._index(log[staged_0:], lambda e: e[0] == "fwd_bwd") + staged_0 + ) + end_0 = self._index(log, lambda e: e == ("on_step_end", 0)) + self.assertEqual( + first_1 < end_0, + expect_before, + f"pipelined={pipelined}: step 1's first fwd_bwd at {first_1}," + f" on_step_end(0) at {end_0}", + ) + # Callback order is unchanged: on_step_end(0) before on_step_begin(1). + self.assertLess( + end_0, self._index(log, lambda e: e == ("on_step_begin", 1)) + ) + + def test_next_steps_groups_are_packed_while_this_step_still_trains(self): + sc = _SCENARIOS["prompt_batch"] + for pipelined in (False, True): + run = _run(sc, pipelined=pipelined) + log = run.log + feeds = [i for i, e in enumerate(log) if e[0] == "feed"] + # Step 1's first group is the 5th feed (4 groups per step). + first_feed_1 = feeds[sc.full_batch_size] + update_0 = self._index(log, lambda e: e[0] == "update") + self.assertEqual(first_feed_1 < update_0, pipelined, log) + + def test_nothing_is_taken_ahead_past_the_last_step(self): + sc = dataclasses.replace(_SCENARIOS["arrival"], num_prompts=40, max_steps=2) + run = _run(sc, pipelined=True) + feeds = [e for e in run.log if e[0] == "feed"] + self.assertLen(feeds, 2 * sc.full_batch_size) + self.assertIsNone(run.program._intake_task) + + +class _RecordingPlans(rl_program.StandardRLProgram): + """Records each `_stage_planned` decision and what the stage then reported.""" + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.plans: list[bool] = [] + self.stage_results: list[bool] = [] + + async def _stage_planned(self, st, *, rescored): + planned = await super()._stage_planned(st, rescored=rescored) + self.plans.append(planned) + return planned + + async def _maybe_stage_next_microbatch(self, st, step_result): + update_ready = await super()._maybe_stage_next_microbatch(st, step_result) + self.stage_results.append(update_ready) + return update_ready + + +class TokenTrainerRegimeTest(absltest.TestCase): + """The double that confirms its update really puts the base loop in the + deferred-fetch regime; the double that cannot keeps it out (two-sided).""" + + def test_only_the_token_trainer_defers_the_fetch(self): + with_token = _reference("prompt_batch", _TokenTrainer) + without = _reference("prompt_batch") + self.assertEqual( + _fetches_after_sync(with_token), + [True] * len(_fetches_after_sync(with_token)), + ) + self.assertNotEmpty(_fetches_after_sync(with_token)) + self.assertEqual( + _fetches_after_sync(without), + [False] * len(_fetches_after_sync(without)), + ) + # The fused call waited for each update (one token wait per update). + updates = [e for e in with_token.log if e[0] == "update"] + waits = [e for e in with_token.log if e[0] == "token_wait"] + self.assertLen(waits, len(updates)) + self.assertNotIn("token_wait", [e[0] for e in without.log]) + + +class TokenTrainerEquivalenceTest(parameterized.TestCase): + """Against a trainer that confirms its update (MaxText with maxtext#5571's + `last_update_token`), where the base loop defers get_metrics past the + weight sync and commit. A planned stage then makes the step's last call + return at the update's dispatch and confirms the update itself, so a staged + boundary keeps the deferred fetch.""" + + @parameterized.named_parameters( + (f"{name}_{mode}", name, mode) for name in _SCENARIOS for mode in _MODES + ) + def test_trainer_cannot_tell_it_from_the_base_loop(self, name, mode): + sc = _SCENARIOS[name] + ref = _reference(name, _TokenTrainer) + new = _run(sc, **_MODES[mode], trainer_cls=_TokenTrainer) + _assert_weights_follow_updates(self, new) + _assert_metrics_fetched_at_the_boundary(self, new) + self.assertIsNone(new.program._pending_train) + if mode == "off": + _assert_is_the_base_loop(self, ref, new) + return + if mode == "unstaged" or name in _SENT_AHEAD: + # Nothing planned that did not happen: the base loop's calls exactly, + # deferred fetch included. + _assert_equivalent(self, ref, new) + else: + _assert_equivalent_but_metrics_fetch(self, ref, new) + stages = _staged_boundaries_confirm_the_update(self, new) + if mode == "staged" and name in _SENT_AHEAD: + expected = sum( + not _writes_scored_checkpoint(sc, k) + for k in range(len(_boundaries(new))) + ) + self.assertEqual(stages, expected, new.log) + if mode == "unstaged": + self.assertEqual(stages, 0) + + +class TokenTrainerChecksFailOnWrongProgramsTest(absltest.TestCase): + """Controls for the token-trainer checks, the relaxed one above all.""" + + def test_the_relaxed_check_still_catches_the_wrong_programs(self): + sc = _SCENARIOS["prompt_batch"] + ref = _reference("prompt_batch", _TokenTrainer) + sends = _run( + sc, + pipelined=True, + staged=True, + program_cls=_SendsBeforeBoundary, + trainer_cls=_TokenTrainer, + ) + with self.assertRaises(AssertionError): + _assert_equivalent_but_metrics_fetch(absltest.TestCase(), ref, sends) + wide = dataclasses.replace(sc, num_prompts=20) + wide_ref = _run( + wide, + pipelined=False, + program_cls=_BaseLoopProgram, + trainer_cls=_TokenTrainer, + ) + wrong_batch = _run( + wide, + pipelined=True, + program_cls=_PrepacksWrongBatch, + trainer_cls=_TokenTrainer, + ) + with self.assertRaises(AssertionError): + _assert_equivalent_but_metrics_fetch( + absltest.TestCase(), wide_ref, wrong_batch + ) + + def test_the_relaxed_check_allows_only_a_fetch_moved_within_its_boundary( + self, + ): + run = _reference("prompt_batch", _TokenTrainer) + log = list(run.log) + fetch = next(i for i, e in enumerate(log) if e[0] == "get_metrics") + save = max(i for i in range(fetch) if log[i][0] == "save") + # Allowed: the fetch ahead of the save, the order without a confirmed + # update (the save then has no buffer pending). + ahead = log[:save] + [log[fetch]] + log[save:fetch] + log[fetch + 1 :] + ahead[save + 1] = ahead[save + 1][:5] + ((),) + moved = dataclasses.replace(run, log=ahead) + self.assertFalse(_check_equivalent(run, moved)) + _assert_equivalent_but_metrics_fetch(self, run, moved) + # Not allowed: the fetch after the next step's first micro-batch. + nxt = next(i for i in range(fetch + 1, len(log)) if log[i][0] == "fwd_bwd") + late = ( + log[:fetch] + log[fetch + 1 : nxt + 1] + [log[fetch]] + log[nxt + 1 :] + ) + with self.assertRaises(AssertionError): + _assert_equivalent_but_metrics_fetch( + absltest.TestCase(), run, dataclasses.replace(run, log=late) + ) + + +class TokenTrainerStagingTest(absltest.TestCase): + """How a stage and the deferred fetch combine against a token trainer.""" + + _BOUNDARY_EVENTS = frozenset({ + "update", + "token_wait", + "stage", + "save", + "prepare_weight_sync", + "get_metrics", + "fwd_bwd", + }) + + def _boundary_calls(self, run: _Run, step: int) -> list[str]: + """What reaches the trainer from step `step`'s update to step+1's first + fwd_bwd (the background round's release left out).""" + calls = run.log + updates = [i for i, e in enumerate(calls) if e[0] == "update"] + start = updates[step] + end = next( + i for i in range(start + 1, len(calls)) if calls[i][0] == "fwd_bwd" + ) + return [ + e[0] for e in calls[start : end + 1] if e[0] in self._BOUNDARY_EVENTS + ] + + def test_a_staged_boundary_keeps_the_base_loops_order(self): + sc = _SCENARIOS["prompt_batch"] + base = _reference("prompt_batch", _TokenTrainer) + staged = _run( + sc, + pipelined=True, + staged=True, + trainer_cls=_TokenTrainer, + program_cls=_RecordingPlans, + ) + self.assertEqual( + self._boundary_calls(base, 0), + [ + "update", + "token_wait", + "save", + "prepare_weight_sync", + "get_metrics", + "fwd_bwd", + ], + ) + self.assertEqual( + self._boundary_calls(staged, 0), + [ + "update", + "stage", + "token_wait", + "save", + "prepare_weight_sync", + "get_metrics", + "fwd_bwd", + ], + ) + self.assertTrue(staged.program.stage_results[0]) + + def test_a_planned_stage_that_finds_nothing_packed_fetches_first(self): + # Packing one micro-batch (0.25 s) outlasts the last call (~20 ms), so the + # stage is planned but nothing is packed when the call returns. + sc = dataclasses.replace( + _SCENARIOS["prompt_batch"], + pack_s=0.25, + costs=_Costs(device_s=0.01, update_s=0.005, host_s=0.005), + ) + run = _run( + sc, + pipelined=True, + staged=True, + trainer_cls=_TokenTrainer, + program_cls=_RecordingPlans, + ) + self.assertNotIn("stage", [e[0] for e in run.log]) + self.assertTrue(any(run.program.plans), run.program.plans) + # The update was not confirmed (the call returned at its dispatch and no + # stage waited for it), so the fetch is its barrier again, ahead of the + # save: the order the base loop uses without a confirmed update. + self.assertEqual( + self._boundary_calls(run, 0), + ["update", "get_metrics", "save", "prepare_weight_sync", "fwd_bwd"], + ) + self.assertEqual(_fetches_after_sync(run)[0], False) + _assert_equivalent_but_metrics_fetch( + self, _run(sc, pipelined=False, trainer_cls=_TokenTrainer), run + ) + # Control: unstaged, the same scenario keeps the deferred fetch. + unstaged = _run(sc, pipelined=True, trainer_cls=_TokenTrainer) + self.assertEqual(_fetches_after_sync(unstaged)[0], True) + + def test_without_a_token_a_plan_changes_nothing_the_trainer_sees(self): + # Today's MaxText pin: the fused call never waits, so whether a stage was + # planned, made or not, the calls are the base loop's. + sc = dataclasses.replace( + _SCENARIOS["prompt_batch"], + pack_s=0.25, + costs=_Costs(device_s=0.01, update_s=0.005, host_s=0.005), + ) + ref = _run(sc, pipelined=False, program_cls=_BaseLoopProgram) + run = _run(sc, pipelined=True, staged=True, program_cls=_RecordingPlans) + self.assertTrue(any(run.program.plans)) + _assert_equivalent(self, ref, run) + + +class CheckpointStampTest(absltest.TestCase): + """The MLPerf rule: a checkpoint is stamped after its last weight update.""" + + def _assert_stamped_after_update(self, run: _Run) -> None: + self.assertNotEmpty(run.checkpoints) + for ckpt in run.checkpoints: + landed = run.trainer.update_landed_ms[ckpt["step"] - 1] + self.assertGreaterEqual(ckpt["timestamp_ms"], int(landed)) + + def test_stamp_follows_the_landed_update(self): + costs = _Costs(device_s=0.03, update_s=0.05) + sc = dataclasses.replace(_SCENARIOS["checkpoints_every_step"], costs=costs) + for pipelined, staged in ((False, False), (True, False), (True, True)): + self._assert_stamped_after_update( + _run(sc, pipelined=pipelined, staged=staged) + ) + + def test_stamp_check_catches_a_stamp_taken_at_dispatch(self): + costs = _Costs(device_s=0.03, update_s=0.05) + sc = dataclasses.replace(_SCENARIOS["checkpoints_every_step"], costs=costs) + bad = _run( + sc, pipelined=True, staged=True, program_cls=_StampsBeforeUpdateLands + ) + with self.assertRaises(AssertionError): + self._assert_stamped_after_update(bad) + + def test_stamp_follows_the_landed_update_with_a_token_trainer(self): + # Checkpoints every step and a stage planned at every boundary but the + # last, whose update the stage, not the last call, confirms. + costs = _Costs(device_s=0.03, update_s=0.05) + sc = dataclasses.replace(_SCENARIOS["prompt_batch"], costs=costs) + for pipelined, staged in ((False, False), (True, False), (True, True)): + run = _run( + sc, pipelined=pipelined, staged=staged, trainer_cls=_TokenTrainer + ) + self._assert_stamped_after_update(run) + if staged: + self.assertGreater(_staged_boundaries_confirm_the_update(self, run), 0) + + def test_a_planned_stage_makes_the_last_call_return_at_dispatch(self): + # Two-sided: stamping as the last call returns is right while that call + # waits for the update (the base loop against a token trainer), and early + # once a planned stage has it return at the update's dispatch -- which is + # why the program then takes the stamp after the stage's confirmation. + costs = _Costs(device_s=0.03, update_s=0.05) + sc = dataclasses.replace(_SCENARIOS["prompt_batch"], costs=costs) + waits = _run( + sc, + pipelined=False, + trainer_cls=_TokenTrainer, + program_cls=_StampsBeforeUpdateLands, + ) + self._assert_stamped_after_update(waits) + at_dispatch = _run( + sc, + pipelined=True, + staged=True, + trainer_cls=_TokenTrainer, + program_cls=_StampsBeforeUpdateLands, + ) + self.assertIn("stage", [e[0] for e in at_dispatch.log]) + with self.assertRaises(AssertionError): + self._assert_stamped_after_update(at_dispatch) + + def _stamp_delays_s(self, run: _Run) -> dict[int, float]: + """Seconds from each scored checkpoint's update landing to its stamp.""" + return { + c["step"]: ( + (c["timestamp_ms"] - run.trainer.update_landed_ms[c["step"] - 1]) + / 1000.0 + ) + for c in run.checkpoints + } + + def test_staging_never_delays_a_scored_stamp_when_feed_bound(self): + # Feed-bound: moving a micro-batch (0.3 s) outlasts the device's work left + # after the update is dispatched (0.06 s), as with several trainer copies + # fed through one controller. Steps 2-3 write scored checkpoints. + costs = _Costs(device_s=0.05, update_s=0.01, host_s=0.3) + sc = dataclasses.replace(_SCENARIOS["checkpoints_from_step_3"], costs=costs) + + class _StagesEveryStep(rl_program.StandardRLProgram): + """WRONG: stages even after a step that writes a scored checkpoint.""" + + # Without this the feed-bound latch would stop staging after one step. + def _staging_hidden(self, device_wait_s): + return True + + def _writes_scored_checkpoint(self, step_result): + return False + + unstaged = self._stamp_delays_s(_run(sc, pipelined=True)) + staged_run = _run(sc, pipelined=True, staged=True) + staged = self._stamp_delays_s(staged_run) + bad = self._stamp_delays_s( + _run(sc, pipelined=True, staged=True, program_cls=_StagesEveryStep) + ) + msg = f"unstaged={unstaged} staged={staged} stages_every_step={bad}" + print(f"STAMP_DELAYS_S {msg}") + self.assertEqual(sorted(staged), [3, 4], msg) + # Staging still ran where nothing is scored ... + self.assertTrue(any(e[0] == "stage" for e in staged_run.log), msg) + for step in (3, 4): + # ... and the scored stamps are as prompt as without staging, + self.assertLess(staged[step], 0.1, msg) + self.assertLess(abs(staged[step] - unstaged[step]), 0.1, msg) + # while staging after the scored step 3 would have pushed its stamp back + # by the transfer's excess over the device's remaining work. + self.assertGreater(bad[3], 0.15, msg) + + +def _stages(run: _Run) -> int: + return sum(1 for e in run.log if e[0] == "stage") + + +def _latch_scenario(costs: _Costs, **kwargs: Any) -> _Scenario: + """12 groups per step: several train calls per step, as at tr-s3 (8-9).""" + base = dict( + order="prompt_batch", + num_prompts=48, + full_batch_size=12, + mini_batch_size=12, + max_steps=4, + costs=costs, + dest_s=1.64 * UNIT, + pack_s=2.1 * UNIT, + tail_s=2.7 * UNIT, + ) + base.update(kwargs) + return _Scenario(**base) + + +def _trs3_costs() -> _Costs: + """tr-s3's MEASURED trainer costs, one production second = UNIT.""" + return _Costs( + device_s=15.11 * UNIT, + update_s=1.0 * UNIT, + stage_s=4.45 * UNIT, + host_s=9.53 * UNIT, + metrics_s=0.5 * UNIT, + ) + + +_FEED_BOUND = _Costs(device_s=0.05, update_s=0.01, host_s=0.3) + + +class StagingLatchTest(absltest.TestCase): + """Staging switches itself off when the trainer is fed slower than it computes. + + The latch reads how long the trainer still waited for its device once a staged + micro-batch was in place (`device_wait_s`). The latch as first built compared + the stage with the shortest gap between the step's train calls returning, + which is the transfer time whenever a step has three or more calls, so it + switched staging off on a device-bound trainer (`_MinGapLatch` below + reproduces it). + """ + + def test_feed_bound_trainer_stops_staging_after_one_step(self): + # Moving a micro-batch (0.3 s) outlasts computing one (0.05 s): the stage + # finds the device idle. + sc = _Scenario( + order="prompt_batch", + num_prompts=20, + max_steps=5, + costs=_FEED_BOUND, + ) + run = _run(sc, pipelined=True, staged=True) + self.assertEqual(_stages(run), 1, run.log) + self.assertTrue(run.program._staging_latched_off) + + # Negative control: without the latch it stages at all 4 boundaries. + unlatched = _run(sc, pipelined=True, staged=True, program_cls=_NeverLatches) + self.assertEqual(_stages(unlatched), 4, unlatched.log) + self.assertFalse(unlatched.program._staging_latched_off) + # Either way the trainer trains the same thing. + _assert_equivalent( + self, _run(sc, pipelined=False, program_cls=_BaseLoopProgram), run + ) + + def test_trainer_bound_trainer_keeps_staging(self): + sc = dataclasses.replace( + _SCENARIOS["prompt_batch"], num_prompts=20, max_steps=5 + ) + run = _run(sc, pipelined=True, staged=True) + self.assertEqual(_stages(run), 4, run.log) + self.assertFalse(run.program._staging_latched_off) + + def test_trainer_bound_at_trs3_ratios_keeps_staging(self): + # PRE-REGISTERED (inverted for the fix): device 15.11, + # transfer 9.53 units, >= 3 calls per step. All 3 boundaries stage, and + # each stage finds ~D + update - T = 6.6 units of device work left. + sc = _latch_scenario(_trs3_costs()) + with mock.patch.object(rl_program, "_release_host_memory", lambda: None): + run = _run(sc, pipelined=True, staged=True) + waits = [w / UNIT for w in run.trainer.device_waits] + msg = f"device waits (units) {waits}" + print(f"TRS3_LATCH {msg} microbatches/step {_microbatches_per_step(run)}") + self.assertGreaterEqual(min(_microbatches_per_step(run)), 3, run.log) + self.assertEqual(_stages(run), 3, msg) + self.assertFalse(run.program._staging_latched_off, msg) + self.assertLen(waits, 3, msg) + for wait in waits: + self.assertBetween(wait, 3.0, 12.0, msg) + _assert_equivalent( + self, _run(sc, pipelined=False, program_cls=_BaseLoopProgram), run + ) + + def test_the_first_latch_tripped_at_trs3_ratios(self): + # Negative control: the min-gap latch, same scenario, stages + # once and switches off, the gap it read being the transfer, not D. + sc = _latch_scenario(_trs3_costs()) + with mock.patch.object(rl_program, "_release_host_memory", lambda: None): + bad = _run(sc, pipelined=True, staged=True, program_cls=_MinGapLatch) + cadences = [c / UNIT for c in bad.program.cadences if c is not None] + msg = f"cadences (units) {cadences}" + self.assertEqual(_stages(bad), 1, msg) + self.assertTrue(bad.program._staging_latched_off) + self.assertBetween(cadences[0], 8.0, 11.5, msg) + + def test_lane_busy_costs_with_three_or_more_calls_keep_staging(self): + # Inverted: D/T = 7.5, >= 3 calls per step. + sc = _latch_scenario(_BUSY, dest_s=0.0, pack_s=0.0, tail_s=0.0) + run = _run(sc, pipelined=True, staged=True) + self.assertGreaterEqual(min(_microbatches_per_step(run)), 3, run.log) + self.assertEqual(_stages(run), 3, run.trainer.device_waits) + self.assertFalse(run.program._staging_latched_off) + # And the first latch did trip here. + bad = _run(sc, pipelined=True, staged=True, program_cls=_MinGapLatch) + self.assertEqual(_stages(bad), 1) + + def test_feed_bound_with_three_or_more_calls_latches_after_one_stage(self): + sc = _latch_scenario( + _FEED_BOUND, + num_prompts=60, + max_steps=5, + dest_s=0.0, + pack_s=0.0, + tail_s=0.0, + ) + run = _run(sc, pipelined=True, staged=True) + self.assertGreaterEqual(min(_microbatches_per_step(run)), 3, run.log) + self.assertEqual(_stages(run), 1, run.trainer.device_waits) + self.assertTrue(run.program._staging_latched_off) + self.assertLess( + run.trainer.device_waits[0], + rl_program.StandardRLProgram._STAGING_HIDDEN_WAIT_S, + ) + unlatched = _run(sc, pipelined=True, staged=True, program_cls=_NeverLatches) + self.assertEqual(_stages(unlatched), 4) + + def test_device_wait_reaches_the_latch(self): + # What the latch reads is what the trainer measured, stage for stage. + class _Recording(rl_program.StandardRLProgram): + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.seen: list[float | None] = [] + + def _staging_hidden(self, device_wait_s): + self.seen.append(device_wait_s) + return super()._staging_hidden(device_wait_s) + + sc = dataclasses.replace( + _SCENARIOS["prompt_batch"], num_prompts=20, max_steps=5 + ) + run = _run(sc, pipelined=True, staged=True, program_cls=_Recording) + self.assertEqual(run.program.seen, run.trainer.device_waits) + self.assertLen(run.program.seen, 4) + + +_SUPPORT_SCENARIO = dataclasses.replace( + _SCENARIOS["prompt_batch"], num_prompts=16, max_steps=4 +) + + +class StagingSupportTest(absltest.TestCase): + """Staging runs only on a trainer that supports it.""" + + def test_an_older_trainer_image_gets_the_boundary_unstaged(self): + sc = _SUPPORT_SCENARIO + run = _run(sc, pipelined=True, staged=True, worker_cls=_OldTrainerWorker) + self.assertEqual(run.program.step, sc.max_steps) + self.assertNotIn("stage", [e[0] for e in run.log]) + self.assertFalse(run.program._staging_supported) + self.assertTrue(run.program._staging_latched_off) + self.assertEqual( + _boundaries(run), [(True, False)] * (sc.max_steps - 1), run.log + ) + _assert_equivalent( + self, _run(sc, pipelined=False, program_cls=_BaseLoopProgram), run + ) + + def test_without_the_check_an_older_trainer_image_fails_the_run(self): + # Negative control, the review's F1 measurement: staging without asking. + class _SkipsTheCheck(rl_program.StandardRLProgram): + + async def _trainer_can_stage(self): + return True + + errors: list[BaseException] = [] + run = _run( + _SUPPORT_SCENARIO, + pipelined=True, + staged=True, + worker_cls=_OldTrainerWorker, + program_cls=_SkipsTheCheck, + errors=errors, + ) + self.assertLen(errors, 1) + self.assertIn("stage_fwd_bwd", str(errors[0])) + self.assertNotIn(("on_step_end", 0), run.log) + + def test_a_trainer_that_cannot_report_its_device_is_not_staged(self): + sc = _SUPPORT_SCENARIO + run = _run(sc, pipelined=True, staged=True, trainer_cls=_NoWaitTrainer) + self.assertNotIn("stage", [e[0] for e in run.log]) + self.assertFalse(run.program._staging_supported) + self.assertEqual(run.program.step, sc.max_steps) + _assert_equivalent( + self, _run(sc, pipelined=False, program_cls=_BaseLoopProgram), run + ) + + def test_the_trainer_is_asked_once(self): + calls = [] + real_info = _Worker.info + + def _counting_info(self): + calls.append(1) + return real_info(self) + + with mock.patch.object(_Worker, "info", _counting_info): + run = _run(_SUPPORT_SCENARIO, pipelined=True, staged=True) + self.assertLen(calls, 1) + self.assertEqual(_stages(run), 3) + self.assertTrue(run.program._staging_supported) + + def test_unstaged_and_off_never_ask(self): + calls = [] + real_info = _Worker.info + + def _counting_info(self): + calls.append(1) + return real_info(self) + + with mock.patch.object(_Worker, "info", _counting_info): + _run(_SUPPORT_SCENARIO, pipelined=True, staged=False) + _run(_SUPPORT_SCENARIO, pipelined=False) + self.assertEmpty(calls) + + def test_a_stage_that_fails_on_the_trainer_is_fatal(self): + # The trainer ran the stage and it raised: the run ends with that error, + # as it would for a train call (past its refusals, stage_fwd_bwd puts the + # worker in ERROR, so no checkpoint could follow). A stage that fails on + # the way, or in the reference call before it, falls back to sending the + # micro-batch whole instead (rl_program_pipelined_boundary_robustness). + real = _Worker.stage_fwd_bwd + calls = {"n": 0} + + def _flaky(self, request, **kwargs): + calls["n"] += 1 + if calls["n"] == 1: + raise ConnectionError("raised by the worker's stage_fwd_bwd") + return real(self, request, **kwargs) + + errors: list[BaseException] = [] + with mock.patch.object(_Worker, "stage_fwd_bwd", _flaky): + run = _run(_SUPPORT_SCENARIO, pipelined=True, staged=True, errors=errors) + self.assertLen(errors, 1) + self.assertIn("raised by the worker's stage_fwd_bwd", str(errors[0])) + self.assertEqual(run.program.step, 0) + + +def _advance_over_holes_eq(self): + """`_advance_over_holes` before the fix: passes a hole only at serving == next.""" + while ( + self._serving_batch_idx == self._next_batch_idx + and self._next_batch_idx in self._sealed + and self._admitted.get(self._next_batch_idx, 0) == 0 + ): + hole = self._next_batch_idx + self._drop_batch_state(hole) + self._advance_to(hole + 1) + + +class _CursorRecordingMixin: + """Records (step, next_batch_idx, _next_batch) at every train call. + + Labelled by the step the call trains (updates before it), not by + `self.step`, which the pipelined loop increments after the next step's + first call is already on the trainer. + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.cursor: list[tuple[int, int, int]] = [] + self._updates_done = 0 + + async def _timed_train_step(self, batch, **kwargs): + self.cursor.append( + (self._updates_done, self.scored_q.next_batch_idx, self._next_batch) + ) + result = await super()._timed_train_step(batch, **kwargs) + if kwargs.get("apply_optimizer"): + self._updates_done += 1 + return result + + +class _CursorRecording(_CursorRecordingMixin, rl_program.StandardRLProgram): + pass + + +class _BaseCursorRecording(_CursorRecordingMixin, _BaseLoopProgram): + pass + + +def _hole_scenario(short_batches: dict[int, int]) -> _Scenario: + return dataclasses.replace( + _SCENARIOS["prompt_batch"], + num_prompts=16, + max_steps=3, + short_batches=short_batches, + ) + + +def _cursor_during_step(run: _Run, step: int) -> list[int]: + return sorted({c[1] for c in run.program.cursor if c[0] == step}) + + +class HoleCursorTest(absltest.TestCase): + """PROMPT_BATCH, a short batch k and then a hole at k+1. + + The pipelined boundary reads batch k+1 by index before batch k commits. + When batch k was short, step k's own last read had already moved the + serving cursor to k+1, and that read moves it on to k+2 over the hole. + `_advance_over_holes` then had to accept a serving cursor past + `next_batch_idx`, or `next_batch_idx` lagged one batch for all of step k+1. + """ + + def _runs(self, sc: _Scenario) -> tuple[_Run, _Run, _Run, _Run]: + ref = _run(sc, pipelined=False, program_cls=_BaseCursorRecording) + off = _run(sc, pipelined=False, program_cls=_CursorRecording) + new = _run(sc, pipelined=True, program_cls=_CursorRecording) + staged = _run(sc, pipelined=True, staged=True, program_cls=_CursorRecording) + return ref, off, new, staged + + def test_short_batch_then_hole_keeps_the_dispatch_window(self): + sc = _hole_scenario({0: 3, 1: 0}) + ref, off, new, staged = self._runs(sc) + self.assertEqual(_cursor_during_step(ref, 1), [2]) + for run in (off, new, staged): + _assert_equivalent(self, ref, run) + self.assertEqual(_cursor_during_step(run, 1), [2], run.program.cursor) + # The hole's state is dropped, not left behind. + self.assertNotIn(1, run.program.scored_q._sealed) + + def test_the_cursor_check_catches_the_old_hole_rule(self): + sc = _hole_scenario({0: 3, 1: 0}) + with mock.patch.object( + trajectory_queue_manager.BatchOrderedQueueManager, + "_advance_over_holes", + _advance_over_holes_eq, + ): + ref, off, new, staged = self._runs(sc) + # The loop without the pipelined boundary never needed the fix ... + self.assertEqual(_cursor_during_step(ref, 1), [2]) + self.assertEqual(_cursor_during_step(off, 1), [2]) + # ... the pipelined one lagged a batch and left the hole behind, + self.assertEqual(_cursor_during_step(new, 1), [1]) + self.assertEqual(_cursor_during_step(staged, 1), [1]) + self.assertIn(1, new.program.scored_q._sealed) + # though it trained the same thing (why trainer equivalence missed it). + _assert_equivalent(self, ref, new) + + def test_full_batch_then_hole_agrees_with_either_rule(self): + sc = _hole_scenario({1: 0}) + for rule in (None, _advance_over_holes_eq): + patch = ( + mock.patch.object( + trajectory_queue_manager.BatchOrderedQueueManager, + "_advance_over_holes", + rule, + ) + if rule is not None + else contextlib.nullcontext() + ) + with patch: + runs = self._runs(sc) + for run in runs: + self.assertEqual(_cursor_during_step(run, 1), [2], run.program.cursor) + + +class _StepTimes(rl_program.StandardRLProgram): + """Records when each step's step_time ends and when its on_step_end runs.""" + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.ends: list[float] = [] + self.callbacks: list[float] = [] + inner = self.on_step_end + + def _on_step_end(step, result): + self.callbacks.append(time.monotonic()) + if inner is not None: + inner(step, result) + + self.on_step_end = _on_step_end + + async def _drain_filtered_groups(self): + # Called right after the step's step_time is taken. + self.ends.append(time.monotonic()) + return await super()._drain_filtered_groups() + + +class StepTimeTest(absltest.TestCase): + """What `step_time` (and the RCP `train_step_time` built from it) spans. + + A review asked whether consecutive steps' step_time now overlap. They do + not: step k ends at its commit, before step k+1's first micro-batch is + sent, which is where step k+1 now starts. What changes is that step k's + deferred bookkeeping falls inside step k+1's time instead of between two + steps. + """ + + def test_steps_never_overlap_and_the_next_starts_at_its_first_send(self): + sc = dataclasses.replace( + _SCENARIOS["prompt_batch"], num_prompts=16, max_steps=4, tail_s=0.05 + ) + for pipelined in (False, True): + run = _run(sc, pipelined=pipelined, program_cls=_StepTimes) + times = { + r[3]: r[2] + for r in run.metrics + if r[0] == "orchestrator" and r[1] == "step_time_sec" + } + ends, callbacks = run.program.ends, run.program.callbacks + self.assertLen(ends, sc.max_steps) + starts = [ends[k] - times[k] for k in range(len(ends))] + for k in range(1, len(ends)): + with self.subTest(pipelined=pipelined, step=k): + self.assertGreaterEqual(starts[k], ends[k - 1]) + if pipelined: + self.assertLess(starts[k], callbacks[k - 1]) + else: + self.assertGreaterEqual(starts[k], callbacks[k - 1]) + + +class StepLogTest(absltest.TestCase): + """Step k's `Train step` line waits for step k+1's first call. + + So a `[pipeline] SENT_AHEAD step=k` line, logged once that call is sent, + marks the step in the log before its metrics line, for monitoring keyed on + step lines. Flags off logs none. + """ + + def _lines(self, pipelined: bool) -> list[str]: + sc = dataclasses.replace( + _SCENARIOS["prompt_batch"], num_prompts=16, max_steps=4 + ) + with self.assertLogs(level="INFO") as logs: + _run(sc, pipelined=pipelined) + return [r.getMessage() for r in logs.records] + + def test_sent_ahead_precedes_each_deferred_step_line(self): + lines = self._lines(pipelined=True) + for k in range(3): + sent = next( + i + for i, m in enumerate(lines) + if m.startswith(f"[pipeline] SENT_AHEAD step={k} ") + ) + step_line = next( + i for i, m in enumerate(lines) if m.startswith(f"Train step {k} ") + ) + self.assertLess(sent, step_line) + + def test_flags_off_logs_no_sent_ahead(self): + lines = self._lines(pipelined=False) + self.assertTrue(any(m.startswith("Train step 0 ") for m in lines)) + self.assertFalse(any("SENT_AHEAD" in m for m in lines)) + + +class _LibraryDefaults(rl_program.StandardRLProgram): + """Drops the two new kwargs, so the library defaults apply.""" + + def __init__( + self, *args, pipeline_step_boundary=None, stage_next_microbatch=None, **kw + ): + del pipeline_step_boundary, stage_next_microbatch + super().__init__(*args, **kw) + + +def _boom(*args, **kwargs): + del args, kwargs + raise AssertionError("a pipelined-boundary path ran") + + +class DefaultOffTest(absltest.TestCase): + """Off by default, and off means the old loop.""" + + def test_library_defaults_take_none_of_the_new_paths(self): + sc = dataclasses.replace( + _SCENARIOS["prompt_batch"], num_prompts=16, max_steps=4 + ) + patches = dict( + _prepack=_boom, + _send_first_microbatch_ahead=_boom, + _stage_planned=_boom, + _maybe_stage_next_microbatch=_boom, + _trainer_can_stage=_boom, + ) + with mock.patch.multiple(rl_program.StandardRLProgram, **patches): + default = _run( + sc, pipelined=True, staged=True, program_cls=_LibraryDefaults + ) + self.assertFalse(default.program.pipeline_step_boundary) + self.assertFalse(default.program.stage_next_microbatch) + _assert_is_the_base_loop( + self, _run(sc, pipelined=False, program_cls=_BaseLoopProgram), default + ) + # Two-sided: with the flags on, the same patches fire. + errors: list[BaseException] = [] + with mock.patch.multiple(rl_program.StandardRLProgram, **patches): + _run(sc, pipelined=True, staged=True, errors=errors) + self.assertLen(errors, 1) + self.assertIn("pipelined-boundary path ran", str(errors[0])) + + def test_argparse_defaults_are_off(self): + args = run_deepswe_dist._parse_args([]) + self.assertFalse(args.pipeline_step_boundary) + self.assertFalse(args.stage_next_microbatch) + + +def _recipe_guard() -> str: + """The STAGE_NEXT_MICROBATCH guard block of mlperf_base.sh.""" + path = os.path.join( + os.path.dirname(os.path.abspath(tunix.__file__)), + "experimental", + "examples", + "recipes", + "mlperf_base.sh", + ) + with open(path) as f: + text = f.read() + start = text.index("export PIPELINE_STEP_BOUNDARY=") + end_marker = "unset _pipeline_step_boundary_on _stage_next_microbatch_on\n" + return text[start : text.index(end_marker, start) + len(end_marker)] + + +# The guard as first built: only the literal "true". +_FIRST_RECIPE_GUARD = """export PIPELINE_STEP_BOUNDARY=${PIPELINE_STEP_BOUNDARY:-false} +export STAGE_NEXT_MICROBATCH=${STAGE_NEXT_MICROBATCH:-false} +if [[ "${STAGE_NEXT_MICROBATCH}" == "true" && "${PIPELINE_STEP_BOUNDARY}" != "true" ]]; then + echo "STAGE_NEXT_MICROBATCH=true requires PIPELINE_STEP_BOUNDARY=true" >&2 + exit 1 +fi +""" + + +def _guard_accepts(guard: str, boundary: str | None, stage: str | None) -> bool: + env = {"PATH": os.environ.get("PATH", "/usr/bin:/bin")} + if boundary is not None: + env["PIPELINE_STEP_BOUNDARY"] = boundary + if stage is not None: + env["STAGE_NEXT_MICROBATCH"] = stage + done = subprocess.run( + ["bash", "-c", guard], env=env, capture_output=True, check=False + ) + return done.returncode == 0 + + +class RecipeGuardTest(absltest.TestCase): + """mlperf_base.sh refuses staging without the boundary, in the launcher's + own truthiness (true, True or 1 is on).""" + + # (PIPELINE_STEP_BOUNDARY, STAGE_NEXT_MICROBATCH, accepted) + _CASES = ( + (None, None, True), + ("true", "true", True), + ("1", "true", True), + ("True", "1", True), + ("true", "True", True), + ("false", "false", True), + ("1", "0", True), + ("false", "true", False), + ("0", "1", False), + (None, "True", False), + ) + + def test_the_guard_uses_the_launchers_truthiness(self): + guard = _recipe_guard() + for boundary, stage, accepted in self._CASES: + with self.subTest(boundary=boundary, stage=stage): + self.assertEqual(_guard_accepts(guard, boundary, stage), accepted) + + def test_the_first_guard_refused_a_valid_combination(self): + # Negative control: the literal comparison refused boundary=1 stage=true. + self.assertFalse(_guard_accepts(_FIRST_RECIPE_GUARD, "1", "true")) + self.assertTrue(_guard_accepts(_recipe_guard(), "1", "true")) + + +# --------------------------------------------------------------------------- +# Randomized scenarios (from the review's fuzz test, against the base loop). +# --------------------------------------------------------------------------- + +_NUM_DRAWS = 30 + + +def _draw(seed: int) -> _Scenario: + """A scenario across the harness's parameter space, including pairings the + hand-picked scenarios never make (holes with short batches, two updates per + step with rescoring, blocking sync with late arrivals, random costs).""" + rng = random.Random(seed) + order = rng.choice(["arrival", "prompt_batch"]) + full, mini = rng.choice([(4, 4), (4, 2), (6, 3), (3, 3), (2, 2), (6, 6)]) + num_batches = rng.randint(2, 5) + num_prompts = full * num_batches - rng.choice([0, 0, 1]) + short = {} + if order == "prompt_batch": + for b in range(num_batches): + r = rng.random() + if r < 0.15: + short[b] = 0 # a hole + elif r < 0.35: + short[b] = rng.randint(1, full - 1) if full > 1 else 1 + max_steps = rng.choice([None, 1, 2, 3, num_batches]) + sync = rng.choice(["async", "async", "blocking", "none"]) + logps = rng.random() < 0.25 + costs = _Costs( + device_s=rng.choice([0.0, 0.004, 0.02]), + update_s=rng.choice([0.0, 0.002, 0.01]), + host_s=rng.choice([0.0, 0.003, 0.03]), + metrics_s=rng.choice([0.0, 0.002]), + stage_s=rng.choice([0.0, 0.003]), + ) + val_start = rng.choice([None, None, 1, 2, 3]) + return _Scenario( + order=order, + num_prompts=num_prompts, + full_batch_size=full, + mini_batch_size=mini, + max_steps=max_steps, + async_sync=(sync == "async"), + sync_weights=(sync != "none"), + pipeline_microbatches=rng.random() < 0.7, + logps=logps, + reference_kl=rng.random() < 0.2, + val_start_step=val_start, + arrival_gap_s=rng.choice([0.0, 0.0, 0.004, 0.015]), + short_batches=short, + costs=costs, + dest_s=rng.choice([0.0, 0.003]), + round_s=rng.choice([0.0, 0.005]), + pack_s=rng.choice([0.0, 0.003]), + tail_s=rng.choice([0.0, 0.004]), + ) + + +def _draw_trainer_bound(seed: int) -> _Scenario: + """Draws in which step k+1's first micro-batch is usually ready early.""" + rng = random.Random(10_000 + seed) + sc = _draw(10_000 + seed) + num_batches = rng.randint(3, 5) + return dataclasses.replace( + sc, + num_prompts=sc.full_batch_size * num_batches, + max_steps=rng.choice([3, None, num_batches]), + costs=_Costs( + device_s=rng.choice([0.02, 0.03]), + update_s=rng.choice([0.002, 0.01]), + host_s=rng.choice([0.002, 0.006]), + metrics_s=rng.choice([0.0, 0.002]), + stage_s=rng.choice([0.0, 0.003]), + ), + arrival_gap_s=rng.choice([0.0, 0.0, 0.002]), + logps=rng.random() < 0.15, + ) + + +class _Counting(rl_program.StandardRLProgram): + """Counts how often the sent-ahead path actually ran.""" + + sent_ahead = 0 + + async def _send_first_microbatch_ahead(self, st): + task = await super()._send_first_microbatch_ahead(st) + if task is not None: + _Counting.sent_ahead += 1 # On the base class, also from subclasses. + return task + + +class _CountingNoLatch(_Counting): + + def _staging_hidden(self, device_wait_s): + return True + + +def _assert_matches_base_loop( + test: absltest.TestCase, ref: _Run, new: _Run, mode: str +) -> None: + _assert_equivalent(test, ref, new) + if any(e[0] == "fwd_bwd" for e in new.log): + _assert_weights_follow_updates(test, new) + test.assertIsNone(new.program._pending_train) + test.assertTrue( + new.program._intake_task is None or new.program._intake_task.done() + ) + if mode == "off": + _assert_is_the_base_loop(test, ref, new) + if not _MODES[mode]["staged"]: + test.assertNotIn("stage", [e[0] for e in new.log]) + + +class RandomizedEquivalenceTest(parameterized.TestCase): + """PRE-REGISTERED: all draws pass in every mode.""" + + @parameterized.parameters(*range(_NUM_DRAWS)) + def test_unified_loop_matches_the_base_loop(self, seed): + sc = _draw(seed) + ref = _run(sc, pipelined=False, program_cls=_BaseLoopProgram) + for mode, kwargs in _MODES.items(): + new = _run(sc, **kwargs) + with self.subTest(mode=mode, scenario=repr(sc)): + _assert_matches_base_loop(self, ref, new, mode) + + +class TokenRandomizedEquivalenceTest(parameterized.TestCase): + """The same draws against a trainer that confirms its update. + + Flags off and unstaged must be the base loop's calls exactly. Staged, a + draw may plan a stage that finds nothing packed in time (random costs, + late groups), so the relaxed check applies; every stage that did happen + must have confirmed the update after the last call returned at dispatch. + """ + + @parameterized.parameters(*range(_NUM_DRAWS)) + def test_unified_loop_matches_the_base_loop(self, seed): + sc = _draw(seed) + ref = _run( + sc, + pipelined=False, + program_cls=_BaseLoopProgram, + trainer_cls=_TokenTrainer, + ) + for mode, kwargs in _MODES.items(): + new = _run(sc, **kwargs, trainer_cls=_TokenTrainer) + with self.subTest(mode=mode, scenario=repr(sc)): + if mode == "staged": + _assert_equivalent_but_metrics_fetch(self, ref, new) + _staged_boundaries_confirm_the_update(self, new) + if any(e[0] == "fwd_bwd" for e in new.log): + _assert_weights_follow_updates(self, new) + self.assertIsNone(new.program._pending_train) + else: + _assert_matches_base_loop(self, ref, new, mode) + + +class TokenRandomizedCoverageTest(absltest.TestCase): + """Negative controls on the token draws: they reach both staged outcomes.""" + + def test_draws_stage_and_fall_back(self): + staged_draws = fallback_draws = 0 + for seed in range(_NUM_DRAWS): + sc = _draw_trainer_bound(seed) + ref = _run( + sc, + pipelined=False, + program_cls=_BaseLoopProgram, + trainer_cls=_TokenTrainer, + ) + new = _run( + sc, + pipelined=True, + staged=True, + program_cls=_CountingNoLatch, + trainer_cls=_TokenTrainer, + ) + _assert_equivalent_but_metrics_fetch(self, ref, new) + staged_draws += _staged_boundaries_confirm_the_update(self, new) > 0 + fallback_draws += not _check_equivalent(ref, new) + print( + f"TOKEN_COVERAGE staged_draws={staged_draws}" + f" fallback_draws={fallback_draws}" + ) + self.assertGreater(staged_draws, 0) + self.assertGreater(fallback_draws, 0) + + +class RandomizedChecksAreNotVacuousTest(absltest.TestCase): + """Controls on the draws: they reach the new paths and catch wrong programs. + + Review result, kept: the wrong programs were caught in 6/30 and 5/30 draws + (prediction >= 5 each, HIT); coverage of the new paths is printed only (its + predictions missed: zero device cost, max_steps 1 and the latch keep many + draws off them). `BiasedRandomizedEquivalenceTest` covers that gap. + """ + + def test_new_paths_run_and_wrong_programs_are_caught(self): + sent_draws = staged_draws = 0 + caught = {"sends_before_boundary": 0, "prepacks_wrong_batch": 0} + for seed in range(_NUM_DRAWS): + sc = _draw(seed) + _Counting.sent_ahead = 0 + run = _run(sc, pipelined=True, staged=True, program_cls=_Counting) + sent_draws += _Counting.sent_ahead > 0 + staged_draws += any(e[0] == "stage" for e in run.log) + ref = _run(sc, pipelined=False, program_cls=_BaseLoopProgram) + for name, cls, staged in ( + ("sends_before_boundary", _SendsBeforeBoundary, True), + ("prepacks_wrong_batch", _PrepacksWrongBatch, False), + ): + errors: list[BaseException] = [] + bad = _run( + sc, pipelined=True, staged=staged, program_cls=cls, errors=errors + ) + if errors or not _check_equivalent(ref, bad): + caught[name] += 1 + print( + f"FUZZ_COVERAGE sent_ahead_draws={sent_draws} " + f"staged_draws={staged_draws} caught={caught}" + ) + self.assertGreaterEqual(caught["sends_before_boundary"], 5) + self.assertGreaterEqual(caught["prepacks_wrong_batch"], 5) + + +class BiasedRandomizedEquivalenceTest(absltest.TestCase): + """Trainer-bound draws with the latch disabled, so the new paths run.""" + + def test_trainer_bound_draws(self): + sent_draws = staged_draws = 0 + for seed in range(_NUM_DRAWS): + sc = _draw_trainer_bound(seed) + ref = _run(sc, pipelined=False, program_cls=_BaseLoopProgram) + for mode in ("unstaged", "staged"): + _Counting.sent_ahead = 0 + new = _run(sc, **_MODES[mode], program_cls=_CountingNoLatch) + with self.subTest(seed=seed, mode=mode, scenario=repr(sc)): + _assert_matches_base_loop(self, ref, new, mode) + if mode == "staged": + sent_draws += _Counting.sent_ahead > 0 + staged_draws += any(e[0] == "stage" for e in new.log) + print( + f"BIASED_COVERAGE sent_ahead_draws={sent_draws}" + f" staged_draws={staged_draws}" + ) + self.assertGreater(sent_draws, 0) + self.assertGreater(staged_draws, 0) + + +class FailureHandlingTest(absltest.TestCase): + + def test_stage_next_microbatch_needs_the_pipelined_boundary(self): + with self.assertRaisesRegex(ValueError, "requires pipeline_step_boundary"): + rl_program.StandardRLProgram( + algo=types.SimpleNamespace( # pyrefly: ignore[bad-argument-type] + num_generations=2, + mini_batch_size=1, + algo_config=types.SimpleNamespace(temperature=None), + ), + stage_next_microbatch=True, + ) + + def test_a_failed_prepack_surfaces_after_the_steps_bookkeeping(self): + + class _FailsTakingAhead(rl_program.StandardRLProgram): + + async def _prepack(self, st): + raise RuntimeError("queue broke while taking ahead") + + errors: list[BaseException] = [] + run = _run( + _SCENARIOS["prompt_batch"], + pipelined=True, + program_cls=_FailsTakingAhead, + errors=errors, + ) + self.assertLen(errors, 1) + self.assertIn("queue broke while taking ahead", str(errors[0])) + # Step 0 finished, bookkeeping included, before step 1 raised; nothing of + # step 1 reached the trainer. + self.assertIn(("on_step_end", 0), run.log) + self.assertEqual(run.program.step, 1) + self.assertNotIn(("on_step_begin", 1), run.log) + self.assertLen([e for e in run.log if e[0] == "update"], 1) + self.assertIsNone(run.program._intake_task) + + def test_a_failed_early_train_call_surfaces_in_the_next_step(self): + + class _Boom(Exception): + pass + + calls = {"n": 0} + + class _FailsSecondStepsFirstCall(rl_program.StandardRLProgram): + + async def _timed_train_step(self, batch, **kwargs): + if ( + self.step == 0 + and calls["n"] == 1 + and not kwargs.get("apply_optimizer") + ): + raise _Boom("trainer rejected the micro-batch sent ahead") + if self.step == 0 and kwargs.get("apply_optimizer"): + calls["n"] = 1 + return await super()._timed_train_step(batch, **kwargs) + + errors: list[BaseException] = [] + run = _run( + _SCENARIOS["prompt_batch"], + pipelined=True, + program_cls=_FailsSecondStepsFirstCall, + errors=errors, + ) + self.assertLen(errors, 1) + self.assertIsInstance(errors[0], _Boom) + # Step 0's bookkeeping still ran before the failure surfaced. + self.assertIn(("on_step_end", 0), run.log) + + +# --------------------------------------------------------------------------- +# Timing, with the trainer's measured costs scaled down (fakes, not a TPU). +# --------------------------------------------------------------------------- + + +# One "second" of the real system. Costs below are tr-s3's (MEASURED, see the +# report): device 15.11 s per micro-batch, update ~1 s, host part 9.53 s, +# get_metrics ~0.5 s once the update landed, destination bind+metadata 1.64 s, +# trainer prepare_weight_sync 4.45 s, step bookkeeping 2.7 s, packing the +# first micro-batch ~4 s. +def _run_with_dispatcher( + sc: _Scenario, + *, + pipelined: bool, + staged: bool = False, + program_cls: type[ + rl_program.StandardRLProgram + ] = rl_program.StandardRLProgram, + first_stagger_s: float = 0.0, + stagger_s: float = 0.0, + rollout_s: float = 0.01, +) -> tuple[_Run, dict[int, float], dict[int, float]]: + """Runs `sc` with the real dispatcher (`rollout_dispatch_stage`) in front. + + The engine's `dispatch_rollouts` stands in for the rollouts and critique: + each dispatched prompt's group reaches `scored_q` `rollout_s` later, stamped + with policy version `max(0, batch_idx - 1)`, which `max_staleness` 1 always + admits. (The version it was dispatched at would depend on whether step 0 had + synced by then, i.e. on machine load.) With `max_staleness` 1, batch 1 is + inside the dispatch window from the start, so with a blocking sync only the + start-up stagger (google/tunix#2845) holds it. + + Returns: + The run, and in seconds after the stagger's own reference point (batch 0 + reaching the dispatch window, `_initial_dispatch_t0`; batch 0's first + dispatch when there is no stagger): when each prompt batch's first prompt + was dispatched, and when each step ended. + """ + log: list[Any] = [] + checkpoints: list[dict[str, Any]] = [] + trainer = _Trainer(log, sc.costs) + host_s = sc.costs.host_s + + def transfer_s(method: str, kwargs: dict[str, Any]) -> float: + del kwargs + return host_s if method in ("fwd_bwd", "stage_fwd_bwd") else 0.0 + + worker = _Worker(lambda: trainer) + worker.initialize() + handle = _RemoteHandle(worker, transfer_s) + engine = distributed_rl_engine.DistributedRLEngine( + rollout_workers=[], + trainer_workers={datatypes.Role.ACTOR: handle}, + weight_sync_coordinator=_Coordinator(handle, sc.dest_s, sc.round_s), + ) + program = _new_program( + sc, + pipelined=pipelined, + staged=staged, + program_cls=program_cls, + log=log, + checkpoints=checkpoints, + program_kwargs=dict( + dataset=[f"task_{i}" for i in range(sc.num_prompts)], + max_staleness=1, + initial_dispatch_stagger_s=stagger_s, + initial_dispatch_first_stagger_s=first_stagger_s, + ), + ) + first_dispatch: dict[int, float] = {} + step_end: dict[int, float] = {} + deliveries: set[asyncio.Task[None]] = set() + on_step_end = program.on_step_end + assert on_step_end is not None + + def _timed_on_step_end(step: int, result: Any) -> None: + step_end[step] = time.monotonic() + on_step_end(step, result) + + program.on_step_end = _timed_on_step_end + + async def _deliver(prompt_idx: int) -> None: + await asyncio.sleep(rollout_s) + batch_idx = prompt_idx // sc.full_batch_size + for gen in range(NUM_GENERATIONS): + item = _item( + prompt_idx, gen, full_batch_size=sc.full_batch_size, logps=sc.logps + ) + item.policy_version = max(0, batch_idx - 1) + await program.scored_q.put(item) + + async def _dispatch_rollouts(prompts, **kwargs): + del kwargs + for prompt in prompts: + coordinates = prompt["metadata"] + first_dispatch.setdefault(coordinates["batch_idx"], time.monotonic()) + task = asyncio.create_task(_deliver(coordinates["prompt_idx"])) + deliveries.add(task) + task.add_done_callback(deliveries.discard) + return [p["prompt_id"] for p in prompts] + + engine.dispatch_rollouts = _dispatch_rollouts # pyrefly: ignore + + async def _main(): + program.engine = engine + dispatcher = asyncio.create_task(program.rollout_dispatch_stage()) + try: + await asyncio.wait_for(program.train_stage(), timeout=120) + finally: + dispatcher.cancel() + for task in list(deliveries): + task.cancel() + + # As in BoundaryTimingTest: gc.collect() at a step's end can stall the event + # loop, and with it the dispatcher, for tens of ms. + try: + with mock.patch.object(rl_program, "_release_host_memory", lambda: None): + asyncio.run(_main()) + finally: + handle.close() + # The stagger counts from batch 0 reaching the window, a thread hop before + # its first dispatch; measured from the dispatch, a loaded box makes batch 1 + # look early. + t0 = program._initial_dispatch_t0 # pylint: disable=protected-access + if t0 is None: + t0 = first_dispatch[0] + run = _Run( + log=log, + metrics=program.metrics_logger.records, # pyrefly: ignore + checkpoints=checkpoints, + last_step_result=program.last_step_result, + trainer=trainer, + program=program, + ) + return ( + run, + {b: t - t0 for b, t in first_dispatch.items()}, + {s: t - t0 for s, t in step_end.items()}, + ) + + +class StaggeredDispatchTest(absltest.TestCase): + """google/tunix#2845's start-up stagger, with the real dispatcher in front. + + PRE-REGISTERED, written before the first run: batch 1 is held for the first + gap in every mode; when step 0 is over before batch 1 goes out, nothing is + taken ahead at boundary 0 (step 1's groups do not exist yet, so the prepack + waits on the held batch without holding up step 0 or the dispatcher); a + later boundary is pipelined; and the trainer cannot tell any mode from the + base loop. Control: without the stagger batch 1 goes out at once. + + Revised after the first runs on a shared CPU machine at load ~170 on 160 + cores: the first version asserted batch 1 goes out within 0.3 s of the gap + and boundary 0 is never taken ahead. Both assumed step 0 ends well within + the gap; under that load it did not, in every mode including the base loop + (MEASURED: base 1.10-1.29 s, off up to 2.08 s), and with a background sync + the dispatch window then rightly waits for step 0's round. Now the sync + blocks (so only the stagger holds batch 1), the lower bound is exact, the + upper bound allows for a loaded box, and the boundary-0 claim is checked + whenever its precondition (step 0 ended before batch 1 went out) holds. + """ + + _SC = _Scenario( + order="prompt_batch", + num_prompts=12, + max_steps=3, + costs=_BUSY, + async_sync=False, + ) + _FIRST_S = 1.0 + + def test_the_held_batch_and_the_pipelined_boundary(self): + ref, ref_due, _ = _run_with_dispatcher( + self._SC, + pipelined=False, + program_cls=_BaseLoopProgram, + first_stagger_s=self._FIRST_S, + ) + self.assertGreaterEqual(ref_due[1], self._FIRST_S - 0.001, ref_due) + for mode, kwargs in _MODES.items(): + new, due, step_end = _run_with_dispatcher( + self._SC, **kwargs, first_stagger_s=self._FIRST_S + ) + boundaries = _boundaries(new) + print( + f"STAGGER_DUE_S base={ref_due[1]:.3f} {mode}={due[1]:.3f}" + f" step0_end={step_end[0]:.3f} boundaries={boundaries}" + ) + with self.subTest(mode=mode): + # Held for the gap (to the timer's resolution), and not much longer + # even on a loaded box. + self.assertGreaterEqual(due[1], self._FIRST_S - 0.001, due) + self.assertLess(due[1], self._FIRST_S + 0.5, due) + _assert_equivalent(self, ref, new) + _assert_weights_follow_updates(self, new) + self.assertLen(boundaries, 2) + if step_end[0] < due[1]: + # Step 0 was over before its successor's groups existed: the + # prepack had nothing to take, and waited without holding it up. + self.assertEqual(boundaries[0], (False, False), (step_end, due)) + self.assertEqual(any(b[0] for b in boundaries), kwargs["pipelined"]) + _, no_stagger, _ = _run_with_dispatcher( + self._SC, pipelined=True, staged=True + ) + self.assertLess(no_stagger[1], 0.5, no_stagger) + + +UNIT = 0.02 +T_DEV, T_UPD, T_HOST, T_METRICS = 15.0, 1.0, 9.5, 0.5 +T_DEST, T_STAGE, T_TAIL, T_PACK = 1.6, 4.5, 2.7, 4.0 +# PRE-REGISTERED, written before the first run: device idle from step k's +# update landing to step k+1's first micro-batch starting, in those units. +PRED_TODAY = T_METRICS + T_DEST + T_STAGE + T_TAIL + T_PACK + T_HOST # 22.8 +PRED_PIPELINED = T_METRICS + T_DEST + T_STAGE + T_HOST # 16.1 +PRED_STAGED = T_METRICS + T_DEST + T_STAGE # 6.6 +# Absolute gaps carry the harness's own fixed latency (thread hops, logging: +# ~1 unit when first measured); the differences cancel it. +TOLERANCE = 3.0 +DIFF_TOLERANCE = 2.0 + + +def _boundary_gaps( + *, + pipelined: bool, + staged: bool, + trainer_cls: type[_Trainer] = _Trainer, + program_cls: type[ + rl_program.StandardRLProgram + ] = rl_program.StandardRLProgram, +) -> tuple[list[float], _Run]: + """Device idle from each update landing to the next micro-batch, in units.""" + sc = _Scenario( + order="prompt_batch", + num_prompts=12, + full_batch_size=6, + mini_batch_size=6, + max_steps=2, + costs=_Costs( + device_s=T_DEV * UNIT, + update_s=T_UPD * UNIT, + stage_s=T_STAGE * UNIT, + host_s=T_HOST * UNIT, + metrics_s=T_METRICS * UNIT, + ), + dest_s=T_DEST * UNIT, + round_s=0.0, + pack_s=T_PACK * UNIT, + tail_s=T_TAIL * UNIT, + ) + # gc.collect() costs a test process tens of ms, noise at this scale; the + # production cost of releasing host memory is part of T_TAIL. + with mock.patch.object(rl_program, "_release_host_memory", lambda: None): + run = _run( + sc, + pipelined=pipelined, + staged=staged, + trainer_cls=trainer_cls, + program_cls=program_cls, + ) + device = run.trainer.device + gaps = [] + for i, (kind, _, _, end) in enumerate(device): + if kind == "update" and i + 1 < len(device): + gaps.append((device[i + 1][2] - end) / UNIT) + return gaps, run + + +class BoundaryTimingTest(absltest.TestCase): + """The device's idle gap at the step boundary, in each mode.""" + + def _gaps(self, *, pipelined: bool, staged: bool) -> list[float]: + return _boundary_gaps(pipelined=pipelined, staged=staged)[0] + + def test_gap_today_pipelined_and_staged_match_the_prediction(self): + today = self._gaps(pipelined=False, staged=False) + pipelined = self._gaps(pipelined=True, staged=False) + staged = self._gaps(pipelined=True, staged=True) + msg = f"today={today} pipelined={pipelined} staged={staged}" + print(f"BOUNDARY_GAP_UNITS {msg}") + self.assertLen(today, 1, msg) + # Positive control: the double reproduces today's boundary. + self.assertAlmostEqual(today[0], PRED_TODAY, delta=TOLERANCE, msg=msg) + self.assertAlmostEqual( + pipelined[0], PRED_PIPELINED, delta=TOLERANCE, msg=msg + ) + self.assertAlmostEqual(staged[0], PRED_STAGED, delta=TOLERANCE, msg=msg) + # What each mode removes: bookkeeping + packing, then also the transfer. + self.assertAlmostEqual( + today[0] - pipelined[0], + PRED_TODAY - PRED_PIPELINED, + delta=DIFF_TOLERANCE, + msg=msg, + ) + self.assertAlmostEqual( + today[0] - staged[0], + PRED_TODAY - PRED_STAGED, + delta=DIFF_TOLERANCE, + msg=msg, + ) + + +class _WaitsForTheUpdateBeforeStaging(rl_program.StandardRLProgram): + """WRONG: stages after a last call that waited for the update. + + What hooking the stage in after google/tunix#2719's fused call, without + having that call return at the update's dispatch, would do once the trainer + can confirm its update: the stage then starts on an idle device. + """ + + async def _timed_train_step(self, batch, **kwargs): + kwargs.pop("wait_for_update", None) + return await super()._timed_train_step(batch, **kwargs) + + +# PRE-REGISTERED, written before the first run. Against a token trainer the +# fused call returns once its update landed and get_metrics moves past the +# sync and commit: the gap is reordered, not changed, so today and pipelined +# are as above. Staged, the last call returns at the update's dispatch and +# the stage confirms the update, so the transfer hides as before. +PRED_TODAY_TOKEN = PRED_TODAY # 22.8 +PRED_PIPELINED_TOKEN = PRED_PIPELINED # 16.1 +PRED_STAGED_TOKEN = PRED_STAGED # 6.6 +# Control (`_WaitsForTheUpdateBeforeStaging`): the stage starts on an idle +# device, so its transfer is exposed again (= pipelined) and the latch, seeing +# no device wait, switches staging off. +PRED_STAGED_AFTER_A_WAIT = PRED_PIPELINED # 16.1 + + +class TokenBoundaryTimingTest(absltest.TestCase): + """The same gaps against a trainer that confirms its update.""" + + def test_gaps_match_the_prediction_and_the_control_loses_the_staging(self): + today, _ = _boundary_gaps( + pipelined=False, staged=False, trainer_cls=_TokenTrainer + ) + pipelined, _ = _boundary_gaps( + pipelined=True, staged=False, trainer_cls=_TokenTrainer + ) + staged, staged_run = _boundary_gaps( + pipelined=True, staged=True, trainer_cls=_TokenTrainer + ) + after_wait, after_wait_run = _boundary_gaps( + pipelined=True, + staged=True, + trainer_cls=_TokenTrainer, + program_cls=_WaitsForTheUpdateBeforeStaging, + ) + msg = ( + f"today={today} pipelined={pipelined} staged={staged}" + f" stage_after_a_wait={after_wait}" + f" device_waits={staged_run.trainer.device_waits}" + f"/{after_wait_run.trainer.device_waits}" + ) + print(f"TOKEN_BOUNDARY_GAP_UNITS {msg}") + self.assertLen(today, 1, msg) + self.assertAlmostEqual(today[0], PRED_TODAY_TOKEN, delta=TOLERANCE, msg=msg) + self.assertAlmostEqual( + pipelined[0], PRED_PIPELINED_TOKEN, delta=TOLERANCE, msg=msg + ) + self.assertAlmostEqual( + staged[0], PRED_STAGED_TOKEN, delta=TOLERANCE, msg=msg + ) + self.assertAlmostEqual( + after_wait[0], PRED_STAGED_AFTER_A_WAIT, delta=TOLERANCE, msg=msg + ) + # Both staged runs did stage; only the control found the device idle. + self.assertEqual(_stages(staged_run), 1, msg) + self.assertEqual(_stages(after_wait_run), 1, msg) + self.assertFalse(staged_run.program._staging_latched_off, msg) + self.assertTrue(after_wait_run.program._staging_latched_off, msg) + + +if __name__ == "__main__": + absltest.main() diff --git a/tests/experimental/orchestrator/rl_program_stage_before_update_test.py b/tests/experimental/orchestrator/rl_program_stage_before_update_test.py new file mode 100644 index 0000000000..bddc7d4732 --- /dev/null +++ b/tests/experimental/orchestrator/rl_program_stage_before_update_test.py @@ -0,0 +1,1366 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Staging the next step's first micro-batch before the step's update. + +`stage_next_microbatch` sends step k+1's first micro-batch to the trainer +while the device still computes step k's last one. By default +(`stage_before_update`), when the trainer can defer a step's update, the stage +goes between step k's last train call, sent with its update deferred, and that +update. The hook it replaces, right after the last call, cannot hide anything +where the trainer's `update()` waits until the last micro-batch has finished +before it dispatches: MaxText with `retry_when_tokens_dropped` (dropless +replay, maxtext#5570, and the final dropless micro-batch of #5673) reads each +micro-batch's overflow flag on the host, the previous one's in `fwd_bwd` and +the last one's in `update()`. The fused last call then returns only once the +device has nothing left but the update, the transfer finds it idle, and the +latch switches staging off at the first boundary. + +These tests reuse the harness of `rl_program_pipelined_boundary_test` (the +trainer double behind a real `TrainerWorker`, and the frozen base loop), +which tests the hook after the last call, and add trainer doubles with that +host blocking (`_ReplayTrainer`) and with real replays: a micro-batch that +overflows is rerun, dropless, when its flag is read, and the rerun's +gradients replace its own (`_LastOverflowsTrainer`, `_FirstOverflowsTrainer`). + +Each check comes with a control that fails on a deliberately wrong input: +the hook after the last call under the replay double (it latches), a latch +that never trips, a rerun that reads the staged micro-batch, and programs or +workers that run the update before the stage. + +Numbers in "units" are seconds of the real system scaled by `U`; the +predictions in the docstrings were written before the tests first ran. +""" + +import asyncio +import collections +import contextlib +import dataclasses +import logging +import time +import types +from typing import Any +from unittest import mock + +from absl.testing import absltest +from absl.testing import parameterized +import rl_program_pipelined_boundary_test as l3t +from tunix.experimental.common import datatypes +from tunix.experimental.examples.deepswe_dist import run_deepswe_dist +from tunix.experimental.orchestrator import rl_program + + +_SBU = {"stage_before_update": True} +LATCH_S = rl_program.StandardRLProgram._STAGING_HIDDEN_WAIT_S + +# One production second, in test seconds: larger than the harness's 0.02, so a +# loaded machine does not turn a hidden stage into an idle device. +U = 0.05 + + +def _run(sc: l3t._Scenario, **kwargs: Any) -> l3t._Run: + """A harness run with the default hook: stage before the update.""" + program_kwargs = dict(_SBU) + program_kwargs.update(kwargs.pop("program_kwargs", None) or {}) + return l3t._run(sc, program_kwargs=program_kwargs, **kwargs) + + +def _base(sc: l3t._Scenario, trainer_cls: type[l3t._Trainer]) -> l3t._Run: + """The frozen base loop's run of `sc` against `trainer_cls`.""" + return l3t._run( + sc, + pipelined=False, + program_cls=l3t._BaseLoopProgram, + trainer_cls=trainer_cls, + ) + + +def _worker_of(run: l3t._Run): + handle = run.program.engine._trainer_workers[datatypes.Role.ACTOR] + return handle._server.bound_instance + + +@contextlib.contextmanager +def _warnings(): + messages: list[str] = [] + handler = logging.Handler(level=logging.WARNING) + handler.emit = lambda record: messages.append(record.getMessage()) + root = logging.getLogger() + root.addHandler(handler) + try: + yield messages + finally: + root.removeHandler(handler) + + +# --------------------------------------------------------------------------- +# Trainer doubles with MaxText's dropless replay +# --------------------------------------------------------------------------- + + +@dataclasses.dataclass +class _Check: + """A micro-batch whose overflow flag is not read yet (maxtext + `_PendingOverflowCheck`).""" + + fingerprint: str + ends_at: float + index: int # into the step's accumulated gradients + overflows: bool + weights_version: int + + +class _ReplayTrainer(l3t._Trainer): + """The harness trainer with `retry_when_tokens_dropped`'s host blocking. + + Like maxtext 853d655's engine: `fwd_bwd(i)` dispatches micro-batch i, then + reads micro-batch i-1's overflow flag on the host (`:2598` -> + `_resolve_overflow_check` `:2694`), so it returns once i-1 has finished; + `update()` reads the last micro-batch's flag before it dispatches the + update (`:2813-2814`). A micro-batch that overflowed is rerun when its flag + is read, after whatever was dispatched since, and the rerun's gradients and + loss replace its own (`:2700-2712`; the capped run accumulated zeros). A + micro-batch's loss is recorded when its flag is read. Here none overflows; + subclasses pick which do (`_overflows`). + """ + + def __init__(self, *args: Any, **kwargs: Any): + super().__init__(*args, **kwargs) + self._pending: _Check | None = None + self._micro = 0 + self.reruns: list[str] = [] + # (where the host was held, seconds), per flag read. + self.blocked: list[tuple[str, float]] = [] + + def _overflows(self, fingerprint: str, micro: int, final: bool) -> bool: + del fingerprint, micro, final + return False + + def _rerun_input(self, check: _Check) -> str: + """What the rerun computes on: the overflowed micro-batch's own inputs.""" + return check.fingerprint + + def fwd_bwd(self, payload: Any, **kwargs: Any) -> None: + fingerprint = l3t._fingerprint(payload) + self._dispatch("fwd_bwd", fingerprint, self.costs.device_s) + self.log.append(("fwd_bwd", fingerprint, self.weights_version)) + self._accumulated.append(fingerprint) + check = _Check( + fingerprint=fingerprint, + ends_at=self._device_free_at, + index=len(self._accumulated) - 1, + overflows=self._overflows( + fingerprint, self._micro, bool(kwargs.get("apply_optimizer")) + ), + weights_version=self.weights_version, + ) + self._micro += 1 + previous, self._pending = self._pending, check + if previous is not None: + self._resolve(previous, "fwd_bwd") + + def _resolve(self, check: _Check, where: str) -> None: + start = time.monotonic() + self._wait_until(check.ends_at) + self.blocked.append((where, time.monotonic() - start)) + if check.overflows: + source = self._rerun_input(check) + self._dispatch("rerun", source, self.costs.device_s) + self.log.append(("rerun", source, check.weights_version)) + self.reruns.append(source) + self._accumulated[check.index] = f"{source}/dropless" + self._step_buffer().append( + l3t._loss(f"{source}/dropless", check.weights_version) + ) + else: + self._step_buffer().append( + l3t._loss(check.fingerprint, check.weights_version) + ) + + def update(self, **kwargs: Any) -> int: + if self._pending is not None: + check, self._pending = self._pending, None + self._resolve(check, "update") + self._micro = 0 + return super().update(**kwargs) + + +class _ReplayTokenTrainer(_ReplayTrainer, l3t._TokenTrainer): + """`_ReplayTrainer` that also confirms its updates (maxtext#5571's + `last_update_token`): the team's fused-call + deferred-fetch regime.""" + + +class _LastOverflowsTrainer(_ReplayTrainer): + """Every step's last micro-batch overflows: its rerun runs in `update()`, + after the stage that went before the update.""" + + def _overflows(self, fingerprint, micro, final): + del fingerprint, micro + return final + + +class _LastOverflowsTokenTrainer(_LastOverflowsTrainer, l3t._TokenTrainer): + pass + + +class _FirstOverflowsTrainer(_ReplayTrainer): + """Every step's first micro-batch overflows: the staged one, rerun from the + input it was staged as.""" + + def _overflows(self, fingerprint, micro, final): + del fingerprint, final + return micro == 0 + + +class _EveryOverflowsTrainer(_ReplayTrainer): + + def _overflows(self, fingerprint, micro, final): + del fingerprint, micro, final + return True + + +class _RerunReadsTheStagedTrainer(_LastOverflowsTrainer): + """WRONG: a rerun computes on the micro-batch staged since, not its own. + + The hazard staging before the update adds: a replay of step k's last + micro-batch runs while step k+1's first is already on the trainer. + """ + + def __init__(self, *args: Any, **kwargs: Any): + super().__init__(*args, **kwargs) + self._held: str | None = None + + def prepare_fwd_bwd_input(self, payload: Any) -> Any: + self._held = l3t._fingerprint(payload) + return super().prepare_fwd_bwd_input(payload) + + def fwd_bwd(self, payload: Any, **kwargs: Any) -> None: + if self._held == l3t._fingerprint(payload): + self._held = None + super().fwd_bwd(payload, **kwargs) + + def _rerun_input(self, check): + return self._held if self._held is not None else check.fingerprint + + +# --------------------------------------------------------------------------- +# Programs: wrong on purpose, or observing +# --------------------------------------------------------------------------- + + +class _AsBuilt(rl_program.StandardRLProgram): + """The hook after the last call (`--no-stage_before_update`).""" + + def __init__(self, *args: Any, **kwargs: Any): + kwargs["stage_before_update"] = False + super().__init__(*args, **kwargs) + + +class _NeverLatches(rl_program.StandardRLProgram): + + def _staging_hidden(self, device_wait_s): + return True + + +class _UpdatesBeforeStaging(rl_program.StandardRLProgram): + """WRONG: runs the deferred update first, then stages.""" + + async def _stage_next_before_update(self, st, step_result): + assert self.engine is not None + update = await self.engine.apply_deferred_update(role=datatypes.Role.ACTOR) + staged = await self._send_stage(st) if self._stage_ok(st) else None + if staged is not None and not self._staging_hidden( + update.get("device_wait_s") + ): + self._staging_latched_off = True + return {**step_result, **update} + + def _stage_ok(self, st): + return ( + not self._staging_latched_off + and self._intake_done() + and st.staged is None + and st.ready + and not st.ready[0].is_final_batch + ) + + +class _RecordingWaits(rl_program.StandardRLProgram): + """Records what the latch reads at each staged boundary.""" + + def __init__(self, *args: Any, **kwargs: Any): + super().__init__(*args, **kwargs) + self.seen: list[float | None] = [] + + def _staging_hidden(self, device_wait_s): + self.seen.append(device_wait_s) + return super()._staging_hidden(device_wait_s) + + +# --------------------------------------------------------------------------- +# Readouts +# --------------------------------------------------------------------------- + +# Trainer-side events that move device state (or end a step) between a stage +# and the update it must precede. +_STEP_EVENTS = ( + "fwd_bwd", + "update", + "get_metrics", + "save", + "prepare_weight_sync", +) + + +def _stages(run: l3t._Run) -> int: + return sum(1 for e in run.log if e[0] == "stage") + + +def _stage_positions(run: l3t._Run) -> list[str]: + """Per stage: 'before_update' if the next step event after it is an update + (a rerun may come in between), else what came next.""" + out = [] + log = run.log + for i, event in enumerate(log): + if event[0] != "stage": + continue + nxt = next((e[0] for e in log[i + 1 :] if e[0] in _STEP_EVENTS), None) + out.append("before_update" if nxt == "update" else str(nxt)) + return out + + +def _staged_microbatches_trained_once(test, run: l3t._Run) -> None: + """Each staged micro-batch is trained exactly once, after the update that + followed its stage, and never sent again.""" + log = run.log + for i, event in enumerate(log): + if event[0] != "stage": + continue + fingerprint = event[1] + trained = [ + j + for j, e in enumerate(log) + if e[0] == "fwd_bwd" and e[1] == fingerprint + ] + test.assertLen(trained, 1, (fingerprint, log)) + updates = [j for j in range(i, trained[0]) if log[j][0] == "update"] + test.assertLen(updates, 1, (fingerprint, log)) + + +def _boundary_gaps( + trainer_cls: type[l3t._Trainer], + *, + pipelined: bool = True, + staged: bool = True, + program_cls: type[ + rl_program.StandardRLProgram + ] = rl_program.StandardRLProgram, + before_update: bool = True, +) -> tuple[list[float], l3t._Run]: + """The harness's one-boundary timing scenario at `U`, gaps in units.""" + sc = l3t._Scenario( + order="prompt_batch", + num_prompts=12, + full_batch_size=6, + mini_batch_size=6, + max_steps=2, + costs=l3t._Costs( + device_s=l3t.T_DEV * U, + update_s=l3t.T_UPD * U, + stage_s=l3t.T_STAGE * U, + host_s=l3t.T_HOST * U, + metrics_s=l3t.T_METRICS * U, + ), + dest_s=l3t.T_DEST * U, + round_s=0.0, + pack_s=l3t.T_PACK * U, + tail_s=l3t.T_TAIL * U, + ) + with mock.patch.object(rl_program, "_release_host_memory", lambda: None): + run = l3t._run( + sc, + pipelined=pipelined, + staged=staged, + trainer_cls=trainer_cls, + program_cls=program_cls, + program_kwargs={"stage_before_update": before_update}, + ) + gaps = [] + device = run.trainer.device + for i, (kind, _, _, end) in enumerate(device): + if kind == "update" and i + 1 < len(device): + gaps.append(round((device[i + 1][2] - end) / U, 2)) + return gaps, run + + +def _costs(*, device, update, stage, host, metrics) -> l3t._Costs: + return l3t._Costs( + device_s=device * U, + update_s=update * U, + stage_s=stage * U, + host_s=host * U, + metrics_s=metrics * U, + ) + + +# tr-s3 MEASURED (the harness's `_trs3_costs`, in production seconds). +TRS3 = dict(device=15.11, update=1.0, stage=4.45, host=9.53, metrics=0.5) +# siv-sk1r MEASURED (its TRAINER_STEP_TIMING lines): D 13.396 s, update 1.59 s, +# prepare_weight_sync 4.60 s, gRPC leg 11.18 s, get_metrics and coordinator +# ~0.9 s. +SK1R = dict(device=13.4, update=1.6, stage=4.6, host=11.2, metrics=0.9) + + +def _latch_scenario(costs: dict[str, float], **kwargs: Any) -> l3t._Scenario: + """12 groups per step (several train calls per step, as at tr-s3), at U.""" + base = dict(dest_s=1.64 * U, pack_s=2.1 * U, tail_s=2.7 * U) + base.update(kwargs) + return l3t._latch_scenario(_costs(**costs), **base) + + +def _waits_units(run: l3t._Run) -> list[float]: + return [round(w / U, 2) for w in run.trainer.device_waits] + + +def _quiet(fn, *args, **kwargs): + with mock.patch.object(rl_program, "_release_host_memory", lambda: None): + return fn(*args, **kwargs) + + +# --------------------------------------------------------------------------- +# The replay doubles behave as MaxText does (two-sided) +# --------------------------------------------------------------------------- + + +class ReplayDoubleTest(absltest.TestCase): + """The doubles block where MaxText blocks, and replay what they say.""" + + def test_update_waits_for_the_last_micro_batch_only_under_replay(self): + sc = _latch_scenario(TRS3) + replay = _quiet(_base, sc, _ReplayTrainer) + updates = [ + s / U for where, s in replay.trainer.blocked if where == "update" + ] + msg = f"update held (units): {[round(s, 2) for s in updates]}" + print(f"SBU_REPLAY_DOUBLE {msg}") + self.assertLen(updates, sc.max_steps, msg) + # The update waited for most of a micro-batch of device time, every step. + for held in updates: + self.assertGreater(held, 0.5 * TRS3["device"], msg) + # Negative control: the harness trainer has no such wait. + plain = _quiet(_base, sc, l3t._Trainer) + self.assertFalse(hasattr(plain.trainer, "blocked")) + + def test_reruns_change_what_is_trained(self): + sc = l3t._SCENARIOS["prompt_batch"] + plain = _base(sc, _ReplayTrainer) + last = _base(sc, _LastOverflowsTrainer) + first = _base(sc, _FirstOverflowsTrainer) + self.assertEmpty(plain.trainer.reruns) + steps = len(l3t._microbatches_per_step(last)) + self.assertLen(last.trainer.reruns, steps) + self.assertLen(first.trainer.reruns, steps) + # A rerun replaces the capped run's gradients: a different update, and a + # different logged loss, from the same micro-batches. + self.assertFalse(l3t._check_equivalent(plain, last)) + self.assertFalse(l3t._check_equivalent(plain, first)) + updates = lambda run: [e for e in run.log if e[0] == "update"] + self.assertNotEqual(updates(plain), updates(last)) + + +# --------------------------------------------------------------------------- +# The trainer cannot tell staging before the update from the base loop +# --------------------------------------------------------------------------- + +_TRAINERS = { + "plain": l3t._Trainer, + "token": l3t._TokenTrainer, + "replay": _ReplayTrainer, + "replay_token": _ReplayTokenTrainer, + "last_overflows": _LastOverflowsTrainer, + "last_overflows_token": _LastOverflowsTokenTrainer, + "first_overflows": _FirstOverflowsTrainer, +} + +_REFERENCE_RUNS: dict[tuple[str, str], l3t._Run] = {} + + +def _reference(name: str, trainer: str) -> l3t._Run: + key = (name, trainer) + if key not in _REFERENCE_RUNS: + _REFERENCE_RUNS[key] = _base(l3t._SCENARIOS[name], _TRAINERS[trainer]) + return _REFERENCE_RUNS[key] + + +class StageBeforeUpdateEquivalenceTest(parameterized.TestCase): + """PRE-REGISTERED: every scenario of the harness, against every double, in + the staged mode with the default hook, passes the STRICT check: the base + loop's trainer calls exactly (reruns included), its checkpoints and metric + values. Unlike the hook after the last call, a stage that was planned but + found nothing packed leaves the update confirmed (the deferred update call + waits for the token), so no relaxation is needed. In the `_SENT_AHEAD` + scenarios every boundary not after a scored checkpoint is staged, and every + stage sits before its step's update.""" + + @parameterized.named_parameters( + (f"{name}_{trainer}", name, trainer) + for name in l3t._SCENARIOS + for trainer in _TRAINERS + ) + def test_trainer_cannot_tell_it_from_the_base_loop(self, name, trainer): + sc = l3t._SCENARIOS[name] + ref = _reference(name, trainer) + new = _run(sc, pipelined=True, staged=True, trainer_cls=_TRAINERS[trainer]) + l3t._assert_equivalent(self, ref, new) + l3t._assert_weights_follow_updates(self, new) + l3t._assert_metrics_fetched_at_the_boundary(self, new) + self.assertTrue( + all(e[2] == 1 for e in new.log if e[0] == "get_metrics"), new.log + ) + self.assertIsNone(new.program._pending_train) + self.assertIsNone(_worker_of(new)._deferred_update) + self.assertEqual( + _stage_positions(new), ["before_update"] * _stages(new), new.log + ) + _staged_microbatches_trained_once(self, new) + if name in l3t._SENT_AHEAD: + boundaries = len(l3t._microbatches_per_step(new)) - 1 + expected = sum( + not l3t._writes_scored_checkpoint(sc, k) for k in range(boundaries) + ) + self.assertEqual(_stages(new), expected, new.log) + if trainer.startswith(("last_overflows", "first_overflows")): + self.assertEqual(new.trainer.reruns, ref.trainer.reruns) + self.assertNotEmpty(new.trainer.reruns) + + +class StageBeforeUpdateChecksFailOnWrongProgramsTest(absltest.TestCase): + """Negative controls for the checks above.""" + + def test_a_rerun_that_reads_the_staged_micro_batch_is_caught(self): + # The wrong double only differs from `_LastOverflowsTrainer` while a stage + # is held, i.e. when staging before the update: the base loop and the + # hook after the last call pass with it, this mode does not. + sc = l3t._SCENARIOS["prompt_batch"] + ref = _base(sc, _LastOverflowsTrainer) + bad = _run( + sc, pipelined=True, staged=True, trainer_cls=_RerunReadsTheStagedTrainer + ) + self.assertGreater(_stages(bad), 0, bad.log) + self.assertFalse(l3t._check_equivalent(ref, bad)) + # Controls: with nothing held at a rerun the wrong double is harmless. + self.assertTrue( + l3t._check_equivalent(ref, _base(sc, _RerunReadsTheStagedTrainer)) + ) + after = l3t._run( + sc, + pipelined=True, + staged=True, + trainer_cls=_RerunReadsTheStagedTrainer, + program_cls=_AsBuilt, + ) + self.assertTrue(l3t._check_equivalent(ref, after)) + # And the right double passes in this mode. + good = _run( + sc, pipelined=True, staged=True, trainer_cls=_LastOverflowsTrainer + ) + self.assertTrue(l3t._check_equivalent(ref, good)) + + def test_the_position_check_catches_an_update_before_the_stage(self): + sc = l3t._SCENARIOS["prompt_batch"] + bad = _run( + sc, pipelined=True, staged=True, program_cls=_UpdatesBeforeStaging + ) + self.assertGreater(_stages(bad), 0) + self.assertNotIn("before_update", _stage_positions(bad), bad.log) + good = _run(sc, pipelined=True, staged=True) + self.assertEqual(_stage_positions(good), ["before_update"] * _stages(good)) + + def test_the_harness_wrong_programs_are_still_caught(self): + sc = l3t._SCENARIOS["prompt_batch"] + # Packing the next step from the wrong batch: trained differently. + wide = dataclasses.replace(sc, num_prompts=20) + wrong_batch = _run( + wide, + pipelined=True, + staged=True, + program_cls=l3t._PrepacksWrongBatch, + trainer_cls=_ReplayTrainer, + ) + self.assertFalse( + l3t._check_equivalent(_base(wide, _ReplayTrainer), wrong_batch) + ) + # Sending step k+1's first micro-batch right after step k's last call: + # with the update deferred, the trainer refuses to train it on top of the + # open step, or to stage it twice, so the run fails loudly. + errors: list[BaseException] = [] + sends = _run( + sc, + pipelined=True, + staged=True, + program_cls=l3t._SendsBeforeBoundary, + trainer_cls=_ReplayTrainer, + errors=errors, + ) + self.assertLen(errors, 1, sends.log) + self.assertRegex(str(errors[0]), "is deferred|already holds staged") + + +# --------------------------------------------------------------------------- +# A replay happens: the staged micro-batch stays as it is +# --------------------------------------------------------------------------- + + +class DroplessReplayTest(absltest.TestCase): + """PRE-REGISTERED. When step k's last micro-batch overflows, its rerun runs + inside the deferred update, after step k+1's first micro-batch was staged. + The staged micro-batch needs no invalidation (its inputs are mapped from + its own payload alone; the rerun reads the inputs it kept): it is trained + once, unchanged, and the run is the base loop's to the trainer, reruns + included. The stage still hides: the wait measured before the update is + that micro-batch's remaining device time (D - T, the rerun comes after it). + When the staged micro-batch itself overflows, it is rerun from the input it + was staged as, with the same result. + """ + + def _check_replayed_run(self, costs, label, trainer_cls, base_cls=None): + sc = _latch_scenario(costs) + ref = _quiet(_base, sc, base_cls or trainer_cls) + run = _quiet(_run, sc, pipelined=True, staged=True, trainer_cls=trainer_cls) + boundaries = len(l3t._microbatches_per_step(run)) - 1 + waits = _waits_units(run) + msg = ( + f"{label}: stages={_stages(run)}/{boundaries} waits(units)={waits}" + f" reruns={len(run.trainer.reruns)}" + f" latched={run.program._staging_latched_off}" + ) + print(f"SBU_REPLAY {msg}") + self.assertGreaterEqual(min(l3t._microbatches_per_step(run)), 3, msg) + self.assertEqual(_stages(run), boundaries, msg) + self.assertFalse(run.program._staging_latched_off, msg) + self.assertEqual(run.trainer.reruns, ref.trainer.reruns, msg) + l3t._assert_equivalent(self, ref, run) + _staged_microbatches_trained_once(self, run) + self.assertEqual(_stage_positions(run), ["before_update"] * _stages(run)) + return run, waits + + def test_last_micro_batch_replayed_trs3(self): + run, waits = self._check_replayed_run( + TRS3, "last-trs3", _LastOverflowsTrainer + ) + self.assertLen(run.trainer.reruns, len(l3t._microbatches_per_step(run))) + for wait in waits: + self.assertAlmostEqual(wait, TRS3["device"] - TRS3["host"], delta=2.0) + + def test_last_micro_batch_replayed_sk1r(self): + _, waits = self._check_replayed_run( + SK1R, "last-sk1r", _LastOverflowsTrainer + ) + for wait in waits: + self.assertAlmostEqual(wait, SK1R["device"] - SK1R["host"], delta=1.5) + + def test_last_micro_batch_replayed_with_a_token_trainer(self): + self._check_replayed_run( + TRS3, "last-trs3-token", _LastOverflowsTokenTrainer + ) + + def test_staged_micro_batch_replayed(self): + run, _ = self._check_replayed_run( + TRS3, "first-trs3", _FirstOverflowsTrainer + ) + staged = {e[1] for e in run.log if e[0] == "stage"} + self.assertTrue(staged & set(run.trainer.reruns), run.log) + + def test_every_micro_batch_replayed(self): + self._check_replayed_run(TRS3, "every-trs3", _EveryOverflowsTrainer) + + def test_the_worker_still_holds_the_stage_after_a_replayed_update(self): + # Direct readout, no harness timing: after the deferred update ran a + # rerun, the worker holds the same staged micro-batch, and the next step's + # first call names it instead of carrying it. + sc = l3t._SCENARIOS["prompt_batch"] + held_after_update = [] + real = l3t._Worker.apply_deferred_update + + def _recording(worker): + resp = real(worker) + staged = worker._staged_fwd_bwd + held_after_update.append( + None if staged is None else l3t._fingerprint(staged[1]) + ) + return resp + + with mock.patch.object(l3t._Worker, "apply_deferred_update", _recording): + run = _run( + sc, pipelined=True, staged=True, trainer_cls=_LastOverflowsTrainer + ) + stages = [e[1] for e in run.log if e[0] == "stage"] + self.assertNotEmpty(stages) + self.assertEqual([h for h in held_after_update if h is not None], stages) + self.assertGreaterEqual(len(run.trainer.reruns), len(stages)) + + +# --------------------------------------------------------------------------- +# No replay: the stage hides the boundary; under replay it still does +# --------------------------------------------------------------------------- + + +class StageBeforeUpdateHidesTest(absltest.TestCase): + """PRE-REGISTERED, re-run on this code. + + Under the replay double the hook after the last call stages once and + latches (positive control: its wait < 50 ms), while staging before the + update stages at every boundary with a wait ~ D - T (tr-s3 5.6 units, + sk1r 2.2), and its boundary gap is the harness's staged gap, 6.6 units + (the hook after the last call: 15.1 under replay = T_HOST - T_UPD + + T_METRICS + T_DEST + T_STAGE). Without replay it also stages at every + boundary and the gap is 6.6 (no regression where update() returns at + dispatch). + """ + + def _check_trainer_bound(self, costs, label, trainer_cls, *, expect_wait): + sc = _latch_scenario(costs) + run = _quiet( + _run, + sc, + pipelined=True, + staged=True, + trainer_cls=trainer_cls, + program_cls=_RecordingWaits, + ) + boundaries = len(l3t._microbatches_per_step(run)) - 1 + waits = _waits_units(run) + msg = f"{label}: stages={_stages(run)}/{boundaries} waits(units)={waits}" + print(f"SBU_HIDES {msg}") + self.assertGreaterEqual(min(l3t._microbatches_per_step(run)), 3, msg) + self.assertEqual(_stages(run), boundaries, msg) + self.assertFalse(run.program._staging_latched_off, msg) + # What the latch read is what the trainer measured, one per stage. + self.assertEqual(run.program.seen, run.trainer.device_waits, msg) + self.assertLen(waits, boundaries, msg) + for wait in waits: + self.assertGreaterEqual(wait * U, LATCH_S, msg) + if expect_wait is not None: + self.assertAlmostEqual(wait, expect_wait, delta=2.0, msg=msg) + l3t._assert_equivalent(self, _quiet(_base, sc, trainer_cls), run) + return run + + def test_replay_trs3(self): + self._check_trainer_bound( + TRS3, + "replay-trs3", + _ReplayTrainer, + expect_wait=TRS3["device"] - TRS3["host"], + ) + + def test_replay_trs3_token(self): + self._check_trainer_bound( + TRS3, "replay-trs3-token", _ReplayTokenTrainer, expect_wait=None + ) + + def test_replay_sk1r(self): + self._check_trainer_bound( + SK1R, + "replay-sk1r", + _ReplayTrainer, + expect_wait=SK1R["device"] - SK1R["host"], + ) + + def test_no_replay_trs3(self): + self._check_trainer_bound( + TRS3, "plain-trs3", l3t._Trainer, expect_wait=None + ) + + def test_the_hook_after_the_last_call_latches_under_replay(self): + # Positive control: same scenario, the old hook. + sc = _latch_scenario(TRS3) + run = _quiet( + l3t._run, + sc, + pipelined=True, + staged=True, + trainer_cls=_ReplayTrainer, + program_cls=_AsBuilt, + ) + msg = ( + f"as-built under replay: stages={_stages(run)}" + f" waits={_waits_units(run)}" + ) + print(f"SBU_ASBUILT_LATCH {msg}") + self.assertEqual(_stages(run), 1, msg) + self.assertTrue(run.program._staging_latched_off, msg) + self.assertLess(run.trainer.device_waits[0], LATCH_S, msg) + # Its one stage sat after the update; it still trains the same thing. + self.assertEqual(_stage_positions(run), ["get_metrics"], run.log) + l3t._assert_equivalent(self, _quiet(_base, sc, _ReplayTrainer), run) + + def test_boundary_gaps(self): + replay_new, replay_run = _boundary_gaps(_ReplayTrainer) + plain_new, _ = _boundary_gaps(l3t._Trainer) + replay_old, _ = _boundary_gaps(_ReplayTrainer, before_update=False) + replay_pipelined, _ = _boundary_gaps(_ReplayTrainer, staged=False) + replay_today, _ = _boundary_gaps( + _ReplayTrainer, pipelined=False, staged=False + ) + last_new, _ = _boundary_gaps(_LastOverflowsTrainer) + pred_old = ( + l3t.T_HOST - l3t.T_UPD + l3t.T_METRICS + l3t.T_DEST + l3t.T_STAGE + ) # 15.1 + msg = ( + f"today={replay_today} pipelined={replay_pipelined}" + f" as_built={replay_old} (pred {pred_old}) before_update={replay_new}" + f" plain={plain_new} last_overflows={last_new} (pred" + f" {l3t.PRED_STAGED}) waits={_waits_units(replay_run)}" + ) + print(f"SBU_GAPS_UNITS {msg}") + tol = l3t.TOLERANCE + for gaps in (replay_new, plain_new, last_new): + self.assertLen(gaps, 1, msg) + self.assertAlmostEqual(gaps[0], l3t.PRED_STAGED, delta=tol, msg=msg) + self.assertAlmostEqual(replay_today[0], l3t.PRED_TODAY, delta=tol, msg=msg) + self.assertAlmostEqual( + replay_pipelined[0], l3t.PRED_PIPELINED, delta=tol, msg=msg + ) + self.assertAlmostEqual(replay_old[0], pred_old, delta=tol, msg=msg) + # The difference that matters, with the harness latency cancelled. + self.assertGreater(replay_old[0] - replay_new[0], 6.0, msg) + + def test_feed_bound_latches_after_one_stage(self): + # The moved latch, two-sided: a transfer (0.3 s) longer than a + # micro-batch (0.05 s) finds the device idle at the deferred update. + sc = l3t._latch_scenario( + l3t._FEED_BOUND, + num_prompts=60, + max_steps=5, + dest_s=0.0, + pack_s=0.0, + tail_s=0.0, + ) + run = _run(sc, pipelined=True, staged=True, trainer_cls=_ReplayTrainer) + msg = ( + f"feed-bound waits (s) {run.trainer.device_waits} stages={_stages(run)}" + ) + print(f"SBU_FEEDBOUND {msg}") + self.assertEqual(_stages(run), 1, msg) + self.assertTrue(run.program._staging_latched_off, msg) + self.assertLess(run.trainer.device_waits[0], LATCH_S, msg) + control = _run( + sc, + pipelined=True, + staged=True, + trainer_cls=_ReplayTrainer, + program_cls=_NeverLatches, + ) + boundaries = len(l3t._microbatches_per_step(control)) - 1 + self.assertEqual(_stages(control), boundaries, control.trainer.device_waits) + for r in (run, control): + l3t._assert_equivalent(self, _base(sc, _ReplayTrainer), r) + + def test_the_wait_is_taken_before_the_update(self): + # Negative control for the order in `apply_deferred_update`: a worker + # that dispatches the update first measures the update's own device time + # (the replay double's update() blocks until the last micro-batch is + # done), so even a feed-bound trainer looks hidden and never latches. + class _UpdatesFirstWorker(l3t._Worker): + + def apply_deferred_update(self): + self._deferred_update = None + train_step = self.update() + wait = self._trainer.wait_for_device() + return self._response( + updated=True, train_step=train_step, device_wait_s=wait + ) + + sc = l3t._latch_scenario( + dataclasses.replace(l3t._FEED_BOUND, update_s=0.06), + num_prompts=60, + max_steps=5, + dest_s=0.0, + pack_s=0.0, + tail_s=0.0, + ) + bad = _run( + sc, + pipelined=True, + staged=True, + trainer_cls=_ReplayTrainer, + worker_cls=_UpdatesFirstWorker, + ) + good = _run(sc, pipelined=True, staged=True, trainer_cls=_ReplayTrainer) + msg = f"bad stages={_stages(bad)} good stages={_stages(good)}" + self.assertTrue(good.program._staging_latched_off, msg) + self.assertFalse(bad.program._staging_latched_off, msg) + self.assertGreater(_stages(bad), _stages(good), msg) + + +# --------------------------------------------------------------------------- +# Which hook runs +# --------------------------------------------------------------------------- + + +class _OlderWorker(l3t._Worker): + """A trainer image with staging but without deferred updates.""" + + apply_deferred_update = None + + def info(self) -> datatypes.WorkerInfo: + info = super().info() + resources = { + k: v for k, v in info.resources.items() if k != "stage_before_update" + } + return dataclasses.replace(info, resources=resources) + + def fwd_bwd(self, request, apply_optimizer=False, **kwargs): + if "defer_update" in kwargs: + raise TypeError("this worker has no defer_update") + return super().fwd_bwd(request, apply_optimizer, **kwargs) + + +class HookSelectionTest(absltest.TestCase): + + def test_default_is_before_the_update(self): + self.assertTrue(run_deepswe_dist._parse_args([]).stage_before_update) + self.assertFalse( + run_deepswe_dist._parse_args( + ["--no-stage_before_update"] + ).stage_before_update + ) + program = l3t._new_program( + l3t._SCENARIOS["prompt_batch"], + pipelined=True, + staged=True, + program_cls=rl_program.StandardRLProgram, + log=[], + checkpoints=[], + program_kwargs={"stage_before_update": True}, + ) + self.assertTrue(program.stage_before_update) + + def test_library_default_is_before_the_update(self): + class _LibraryDefault(rl_program.StandardRLProgram): + + def __init__(self, *args, stage_before_update=None, **kwargs): + del stage_before_update + super().__init__(*args, **kwargs) + + sc = l3t._SCENARIOS["prompt_batch"] + run = l3t._run(sc, pipelined=True, staged=True, program_cls=_LibraryDefault) + self.assertTrue(run.program.stage_before_update) + self.assertGreater(_stages(run), 0) + self.assertEqual(_stage_positions(run), ["before_update"] * _stages(run)) + + def test_flag_off_stages_after_the_last_call(self): + sc = l3t._SCENARIOS["prompt_batch"] + run = _run( + sc, + pipelined=True, + staged=True, + program_kwargs={"stage_before_update": False}, + ) + self.assertGreater(_stages(run), 0) + self.assertNotIn("before_update", _stage_positions(run)) + self.assertIsNone(run.program._stage_before_update_supported) + + def test_a_worker_without_deferred_updates_gets_the_old_hook(self): + sc = l3t._SCENARIOS["prompt_batch"] + with _warnings() as warnings: + run = _run(sc, pipelined=True, staged=True, worker_cls=_OlderWorker) + self.assertGreater(_stages(run), 0) + self.assertNotIn("before_update", _stage_positions(run)) + self.assertFalse(run.program._stage_before_update_supported) + self.assertTrue( + any("cannot defer a step's update" in w for w in warnings), warnings + ) + l3t._assert_equivalent(self, _reference("prompt_batch", "plain"), run) + + def test_without_the_check_such_a_worker_fails_the_run(self): + # Negative control: deferring without asking the trainer first. + class _SkipsTheCheck(rl_program.StandardRLProgram): + + async def _trainer_can_stage_before_update(self): + return True + + errors: list[BaseException] = [] + run = _run( + l3t._SCENARIOS["prompt_batch"], + pipelined=True, + staged=True, + worker_cls=_OlderWorker, + program_cls=_SkipsTheCheck, + errors=errors, + ) + self.assertLen(errors, 1) + self.assertIn("defer_update", str(errors[0])) + self.assertNotIn(("on_step_end", 0), run.log) + + def test_the_trainer_is_asked_once(self): + calls = [] + real_info = l3t._Worker.info + + def _counting_info(self): + calls.append(1) + return real_info(self) + + with mock.patch.object(l3t._Worker, "info", _counting_info): + run = _run(l3t._SUPPORT_SCENARIO, pipelined=True, staged=True) + # One answer covers staging and the deferred update. + self.assertLen(calls, 1) + self.assertEqual(_stages(run), 3) + self.assertTrue(run.program._staging_supported) + self.assertTrue(run.program._stage_before_update_supported) + + def test_a_worker_that_runs_the_update_anyway_is_taken_as_fused(self): + # A worker that does not know `defer_update` and runs the update in the + # call (an older TrainerWorker forwards unknown keywords to the trainer, + # which ignores them) leaves `update_deferred` unset: the engine handles + # the reply as a fused call (one update, never two) and the program + # stages after it, as with a worker that says it cannot defer. + class _IgnoresDefer(l3t._Worker): + + def fwd_bwd(self, request, apply_optimizer=False, **kwargs): + kwargs.pop("defer_update", None) + return super().fwd_bwd(request, apply_optimizer, **kwargs) + + class _SkipsTheCheck(rl_program.StandardRLProgram): + + async def _trainer_can_stage_before_update(self): + return True + + sc = l3t._SCENARIOS["prompt_batch"] + run = _run( + sc, + pipelined=True, + staged=True, + worker_cls=_IgnoresDefer, + program_cls=_SkipsTheCheck, + ) + self.assertGreater(_stages(run), 0) + self.assertNotIn("before_update", _stage_positions(run)) + l3t._assert_equivalent(self, _reference("prompt_batch", "plain"), run) + updates = [e for e in run.log if e[0] == "update"] + self.assertLen(updates, len(l3t._microbatches_per_step(run))) + + def test_flags_off_never_take_the_new_paths(self): + def _boom(*args, **kwargs): + del args, kwargs + raise AssertionError("a stage-before-update path ran") + + sc = dataclasses.replace( + l3t._SCENARIOS["prompt_batch"], num_prompts=16, max_steps=4 + ) + patches = dict( + _trainer_can_stage_before_update=_boom, + _stage_next_before_update=_boom, + ) + with mock.patch.multiple(rl_program.StandardRLProgram, **patches): + for kwargs in ( + dict(pipelined=False, staged=False), + dict(pipelined=True, staged=False), + ): + run = _run(sc, **kwargs) + l3t._assert_equivalent(self, _base(sc, l3t._Trainer), run) + self.assertNotIn("stage", [e[0] for e in run.log]) + # Two-sided: staged, the same patches fire. + errors: list[BaseException] = [] + with mock.patch.multiple(rl_program.StandardRLProgram, **patches): + _run(sc, pipelined=True, staged=True, errors=errors) + self.assertLen(errors, 1) + self.assertIn("stage-before-update path ran", str(errors[0])) + + +# --------------------------------------------------------------------------- +# Checkpoint stamps +# --------------------------------------------------------------------------- + + +class CheckpointStampTest(absltest.TestCase): + """The MLPerf rule: a checkpoint is stamped after its last weight update.""" + + def _assert_stamped_after_update(self, run: l3t._Run) -> None: + self.assertNotEmpty(run.checkpoints) + for ckpt in run.checkpoints: + landed = run.trainer.update_landed_ms[ckpt["step"] - 1] + self.assertGreaterEqual(ckpt["timestamp_ms"], int(landed)) + + # The harness's busy costs with a longer update. Against a trainer whose + # fwd_bwd returns at dispatch, the deferred last call returns right after + # that micro-batch's dispatch (a fused call would also have dispatched the + # update first), so whether the next step's first micro-batch is packed by + # then is timing; under the replay doubles it returns once the micro-batch + # before it finished, and it always is. + _COSTS = l3t._Costs(device_s=0.15, update_s=0.05, host_s=0.02) + + def test_stamp_follows_the_landed_update(self): + costs = self._COSTS + sc = dataclasses.replace(l3t._SCENARIOS["prompt_batch"], costs=costs) + for trainer_cls in ( + l3t._Trainer, + l3t._TokenTrainer, + _ReplayTrainer, + _LastOverflowsTokenTrainer, + ): + run = _run(sc, pipelined=True, staged=True, trainer_cls=trainer_cls) + if issubclass(trainer_cls, _ReplayTrainer): + self.assertGreater(_stages(run), 0) + self._assert_stamped_after_update(run) + + def test_stamp_check_catches_a_stamp_taken_when_the_last_call_returns(self): + # With the update deferred the last call returns before the update is + # even dispatched; a stamp there is early against every trainer. + costs = self._COSTS + sc = dataclasses.replace(l3t._SCENARIOS["prompt_batch"], costs=costs) + bad = _run( + sc, + pipelined=True, + staged=True, + trainer_cls=_ReplayTokenTrainer, + program_cls=l3t._StampsBeforeUpdateLands, + ) + self.assertGreater(_stages(bad), 0) + with self.assertRaises(AssertionError): + self._assert_stamped_after_update(bad) + # Control: the same trainer, the program's own stamp. + good = _run( + sc, pipelined=True, staged=True, trainer_cls=_ReplayTokenTrainer + ) + self.assertGreater(_stages(good), 0) + self._assert_stamped_after_update(good) + + def test_scored_steps_are_not_staged(self): + costs = l3t._Costs(device_s=0.05, update_s=0.01, host_s=0.3) + sc = dataclasses.replace( + l3t._SCENARIOS["checkpoints_from_step_3"], costs=costs + ) + run = _run(sc, pipelined=True, staged=True, program_cls=_NeverLatches) + log = run.log + # Boundaries 0 and 1 stage; 2 and 3 follow steps that write scored + # checkpoints (steps 3 and 4). + stage_steps = [] + for i, event in enumerate(log): + if event[0] == "stage": + stage_steps.append(sum(1 for e in log[:i] if e[0] == "update")) + self.assertEqual(stage_steps, [0, 1], log) + self._assert_stamped_after_update(run) + + +# --------------------------------------------------------------------------- +# Failures on the new path +# --------------------------------------------------------------------------- + + +class FailureTest(absltest.TestCase): + + def _scenario(self): + return dataclasses.replace( + l3t._SCENARIOS["prompt_batch"], num_prompts=16, max_steps=4 + ) + + def test_a_stage_lost_on_the_way_falls_back_and_the_update_still_runs(self): + sc = self._scenario() + real_execute = l3t._RemoteHandle._execute + fired = {"n": 0} + + async def _execute(handle, request): + if request.method_name == "stage_fwd_bwd" and not fired["n"]: + fired["n"] = 1 + raise ConnectionError("transport: stream reset before the stage ran") + return await real_execute(handle, request) + + with _warnings() as warnings: + with mock.patch.object(l3t._RemoteHandle, "_execute", _execute): + run = _run(sc, pipelined=True, staged=True, trainer_cls=_ReplayTrainer) + self.assertEqual(fired["n"], 1, "precondition: boundary 0 was staged") + self.assertEqual(run.program.step, sc.max_steps) + self.assertTrue(run.program._staging_latched_off) + self.assertEqual(_stages(run), 0, run.log) + self.assertTrue( + any("Staging failed" in w and "ConnectionError" in w for w in warnings), + warnings, + ) + l3t._assert_equivalent(self, _base(sc, _ReplayTrainer), run) + worker = _worker_of(run) + self.assertIsNone(worker._deferred_update) + self.assertEqual(worker.state, datatypes.WorkerState.READY) + + def test_without_the_fallback_a_lost_stage_ends_the_run(self): + # Negative control: the same failure treated as fatal. + sc = self._scenario() + real_execute = l3t._RemoteHandle._execute + fired = {"n": 0} + + async def _execute(handle, request): + if request.method_name == "stage_fwd_bwd" and not fired["n"]: + fired["n"] = 1 + raise ConnectionError("transport: stream reset before the stage ran") + return await real_execute(handle, request) + + errors: list[BaseException] = [] + with mock.patch.object( + rl_program, "_raised_by_worker_method", lambda exc: True + ): + with mock.patch.object(l3t._RemoteHandle, "_execute", _execute): + run = _run(sc, pipelined=True, staged=True, errors=errors) + self.assertLen(errors, 1) + self.assertNotIn(("on_step_end", 0), run.log) + + def test_a_failed_reference_call_falls_back(self): + class _MarksItsStage(rl_program.StandardRLProgram): + + staging = False + + async def _send_stage(self, st): + type(self).staging = True + try: + return await super()._send_stage(st) + finally: + type(self).staging = False + + sc = l3t._SCENARIOS["reference_kl"] + real = l3t._Reference.per_token_logps + fired = {"n": 0} + + def _flaky(reference, items, **kwargs): + if _MarksItsStage.staging and not fired["n"]: + fired["n"] = 1 + raise RuntimeError("reference model: device lost") + return real(reference, items, **kwargs) + + errors: list[BaseException] = [] + with mock.patch.object(l3t._Reference, "per_token_logps", _flaky): + run = _run( + sc, + pipelined=True, + staged=True, + program_cls=_MarksItsStage, + errors=errors, + ) + self.assertEqual(fired["n"], 1, "precondition: boundary 0 was staged") + self.assertEqual(errors, []) + self.assertTrue(run.program._staging_latched_off) + l3t._assert_equivalent(self, _reference("reference_kl", "plain"), run) + + def test_a_stage_that_fails_on_the_trainer_is_fatal(self): + sc = self._scenario() + real = l3t._Worker.stage_fwd_bwd + + def _raises(worker, request, **kwargs): + del request, kwargs + raise ConnectionError("raised by the worker's stage_fwd_bwd") + + errors: list[BaseException] = [] + with mock.patch.object(l3t._Worker, "stage_fwd_bwd", _raises): + run = _run(sc, pipelined=True, staged=True, errors=errors) + del real + self.assertLen(errors, 1) + self.assertIn("raised by the worker's stage_fwd_bwd", str(errors[0])) + self.assertNotIn(("on_step_end", 0), run.log) + + def test_a_failed_deferred_update_ends_the_run(self): + sc = self._scenario() + + class _FailingUpdateTrainer(l3t._Trainer): + + def update(self, **kwargs): + raise RuntimeError("update lost on the device") + + errors: list[BaseException] = [] + run = _run( + sc, + pipelined=True, + staged=True, + trainer_cls=_FailingUpdateTrainer, + errors=errors, + ) + self.assertLen(errors, 1) + self.assertIn("deferred optimizer update failed", str(errors[0])) + self.assertIn("update lost on the device", str(errors[0])) + self.assertNotIn(("on_step_end", 0), run.log) + worker = _worker_of(run) + self.assertEqual(worker.state, datatypes.WorkerState.ERROR) + self.assertIn( + "deferred optimizer update failed", worker.heartbeat().last_error + ) + + def test_an_update_that_is_never_applied_stops_the_next_step(self): + # Negative control for the worker's guard: a program that forgets the + # deferred update cannot train the next step on top of it. + class _ForgetsTheUpdate(rl_program.StandardRLProgram): + + async def _stage_next_before_update(self, st, step_result): + if st.ready and self._intake_done(): + await self._send_stage(st) + return {**step_result, "updated": True, "train_step": 1} + + errors: list[BaseException] = [] + run = _run( + self._scenario(), + pipelined=True, + staged=True, + program_cls=_ForgetsTheUpdate, + errors=errors, + ) + self.assertLen(errors, 1) + self.assertIn("is deferred", str(errors[0])) + self.assertEqual(len([e for e in run.log if e[0] == "update"]), 0) + + +# --------------------------------------------------------------------------- +# Randomized scenarios +# --------------------------------------------------------------------------- + + +class RandomizedEquivalenceTest(parameterized.TestCase): + """PRE-REGISTERED: all 30 draws of the harness pass the strict check in the + staged mode with the default hook, against the plain, token and replay + doubles; and the biased trainer-bound draws stage in at least one draw.""" + + @parameterized.parameters(*range(l3t._NUM_DRAWS)) + def test_matches_the_base_loop(self, seed): + sc = l3t._draw(seed) + for trainer in ("plain", "token", "last_overflows_token"): + trainer_cls = _TRAINERS[trainer] + ref = _base(sc, trainer_cls) + new = _run(sc, pipelined=True, staged=True, trainer_cls=trainer_cls) + with self.subTest(trainer=trainer, scenario=repr(sc)): + l3t._assert_equivalent(self, ref, new) + if any(e[0] == "fwd_bwd" for e in new.log): + l3t._assert_weights_follow_updates(self, new) + self.assertIsNone(new.program._pending_train) + self.assertEqual( + _stage_positions(new), ["before_update"] * _stages(new) + ) + + def test_trainer_bound_draws_stage(self): + staged_draws = 0 + for seed in range(l3t._NUM_DRAWS): + sc = l3t._draw_trainer_bound(seed) + for trainer in ("plain", "replay_token"): + trainer_cls = _TRAINERS[trainer] + ref = _base(sc, trainer_cls) + new = _run( + sc, + pipelined=True, + staged=True, + trainer_cls=trainer_cls, + program_cls=_NeverLatches, + ) + with self.subTest(seed=seed, trainer=trainer): + l3t._assert_equivalent(self, ref, new) + staged_draws += _stages(new) > 0 + print(f"SBU_BIASED_COVERAGE staged_draws={staged_draws}") + self.assertGreater(staged_draws, 0) + + +if __name__ == "__main__": + absltest.main() diff --git a/tests/experimental/queue_manager/trajectory_queue_manager_test.py b/tests/experimental/queue_manager/trajectory_queue_manager_test.py index 2fce401765..51c8b0eabb 100644 --- a/tests/experimental/queue_manager/trajectory_queue_manager_test.py +++ b/tests/experimental/queue_manager/trajectory_queue_manager_test.py @@ -421,6 +421,58 @@ async def _run_test(): asyncio.run(_run_test()) + def test_hole_is_passed_once_the_batch_before_it_commits(self): + """Tests a hole read past by index before the previous batch commits. + + A short batch 0, then a hole at batch 1. Draining batch 0 moves the + serving cursor to 1; reading batch 1 by index before batch 0 commits (the + pipelined step boundary does) moves it on to 2. Committing batch 0 must + still pass over the hole, or `next_batch_idx` lags a batch until the next + commit and the hole's state is never dropped. + """ + + async def _run_test(): + cursor_log: list[int] = [] + manager = _create_batch_manager(full_batch_size=2, cursor_log=cursor_log) + manager.expect(0, num_groups=1) + manager.expect(1, num_groups=0) + only = await _put_group(manager, "p0", batch_idx=0) + await _put_group(manager, "p4", batch_idx=2) + + self.assertEqual(await manager.get_ordered_group(), (0, only)) + self.assertIsNone(await manager.get_ordered_group(batch_idx=0)) + self.assertEqual(manager.serving_batch_idx, 1) + self.assertIsNone(await manager.get_ordered_group(batch_idx=1)) + self.assertEqual(manager.serving_batch_idx, 2) + self.assertEqual(manager.next_batch_idx, 0) + + manager.commit_batch(0) + self.assertEqual(manager.next_batch_idx, 2) + self.assertEqual(cursor_log, [1, 2]) + self.assertNotIn(1, manager._sealed) + + asyncio.run(_run_test()) + + def test_hole_rule_change_leaves_the_in_order_consumer_alone(self): + """Tests the same batches read in order, committing between them.""" + + async def _run_test(): + cursor_log: list[int] = [] + manager = _create_batch_manager(full_batch_size=2, cursor_log=cursor_log) + manager.expect(0, num_groups=1) + manager.expect(1, num_groups=0) + only = await _put_group(manager, "p0", batch_idx=0) + later = await _put_group(manager, "p4", batch_idx=2) + + self.assertEqual(await manager.get_ordered_group(), (0, only)) + self.assertIsNone(await manager.get_ordered_group(batch_idx=0)) + manager.commit_batch(0) + self.assertEqual(manager.next_batch_idx, 2) + self.assertEqual(cursor_log, [1, 2]) + self.assertEqual(await manager.get_ordered_group(), (2, later)) + + asyncio.run(_run_test()) + def test_expect_seals_a_short_final_batch(self): """Tests expect() letting a batch seal below full_batch_size.""" diff --git a/tests/experimental/worker/trainer_worker_deferred_update_test.py b/tests/experimental/worker/trainer_worker_deferred_update_test.py new file mode 100644 index 0000000000..6ed71ed8ca --- /dev/null +++ b/tests/experimental/worker/trainer_worker_deferred_update_test.py @@ -0,0 +1,318 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for TrainerWorker's deferred update. + +`fwd_bwd(apply_optimizer=True, defer_update=True)` runs a step's last +micro-batch and leaves its update to `apply_deferred_update`, so the next +step's first micro-batch can be staged in between. The trainer must see what +a fused call shows it (the micro-batch marked last, then one update), the +worker must accept nothing that needs the step closed while the update is +deferred, a stage in between must not wait for the device, and the update +call must wait for the device before it dispatches the update. +""" + +from absl.testing import absltest +from absl.testing import parameterized +import trainer_worker_staging_test as st +import types +from tunix.experimental.common import datatypes +from tunix.experimental.worker import trainer_worker + +WorkerState = datatypes.WorkerState + + +def _kinds(trainer: st._Trainer) -> list[str]: + return [c[0] for c in trainer.calls] + + +def _deferred(trainer: st._Trainer) -> trainer_worker.TrainerWorker: + worker = st._worker(trainer) + resp = worker.fwd_bwd( + st._request("req-1", 1), apply_optimizer=True, defer_update=True + ) + assert resp.metadata.get("update_deferred"), resp.metadata + return worker + + +class _FailingUpdateTrainer(st._WaitingTrainer): + + def update(self, **kwargs): + raise RuntimeError("XLA: update failed on device") + + +class DeferredUpdateTest(parameterized.TestCase): + + def test_defer_runs_the_last_micro_batch_and_holds_the_update(self): + trainer = st._TokenTrainer() + worker = st._worker(trainer) + request = st._request("req-1", 1) + resp = worker.fwd_bwd(request, apply_optimizer=True, defer_update=True) + self.assertEqual( + trainer.calls, [("fwd_bwd", request.payload, {"apply_optimizer": True})] + ) + self.assertTrue(resp.metadata["update_deferred"]) + for key in ("updated", "train_step", "update_ready"): + self.assertNotIn(key, resp.metadata) + self.assertEqual(resp.request_id, "req-1") + self.assertEqual(worker._deferred_update, "req-1") + self.assertEqual(worker.state, WorkerState.READY) + + def test_without_defer_the_call_is_the_fused_one(self): + trainer = st._TokenTrainer() + resp = st._worker(trainer).fwd_bwd( + st._request("req-1", 1), apply_optimizer=True + ) + self.assertEqual(_kinds(trainer), ["fwd_bwd", "update"]) + self.assertTrue(resp.metadata["updated"]) + self.assertNotIn("update_deferred", resp.metadata) + + def test_a_caller_cannot_spoof_the_deferral(self): + trainer = st._Trainer() + request = st._request("req-1", 1) + request.metadata["update_deferred"] = True + resp = st._worker(trainer).fwd_bwd(request, apply_optimizer=True) + self.assertNotIn("update_deferred", resp.metadata) + self.assertEqual(_kinds(trainer), ["fwd_bwd", "update"]) + + def test_defer_needs_apply_optimizer(self): + trainer = st._Trainer() + worker = st._worker(trainer) + with self.assertRaisesRegex(ValueError, "needs apply_optimizer=True"): + worker.fwd_bwd(st._request("req-1", 1), defer_update=True) + self.assertEmpty(trainer.calls) + self.assertEqual(worker.state, WorkerState.READY) + self.assertIsNone(worker._deferred_update) + + @parameterized.named_parameters( + ("fwd_bwd", lambda w: w.fwd_bwd(st._request("req-2", 2))), + ("fused_fwd_bwd", lambda w: w.fwd_bwd(st._request("req-2", 2), True)), + ("update", lambda w: w.update()), + ("get_metrics", lambda w: w.get_metrics()), + ("save_checkpoint", lambda w: w.save_checkpoint({"step": 1})), + ("prepare_weight_sync", lambda w: w.prepare_weight_sync()), + ( + "per_token_logps", + lambda w: w.per_token_logps( + types.SimpleNamespace(pad_id=0, eos_id=0, request_id="lp") + ), + ), + ("eval_step", lambda w: w.eval_step(st._request("req-2", 2))), + ("run_eval", lambda w: w.run_eval([st._payload(2)])), + ) + def test_while_deferred_the_step_closing_calls_are_refused(self, call): + trainer = st._TokenTrainer() + worker = _deferred(trainer) + before = list(trainer.calls) + with self.assertRaisesRegex(RuntimeError, "is deferred"): + call(worker) + self.assertEqual(trainer.calls, before) + self.assertEqual(worker.state, WorkerState.READY) + self.assertEqual(worker._deferred_update, "req-1") + + def test_after_the_deferred_update_the_next_micro_batch_is_accepted(self): + # Two-sided for the refusals above. + trainer = st._TokenTrainer() + worker = _deferred(trainer) + worker.apply_deferred_update() + worker.fwd_bwd(st._request("req-2", 2)) + self.assertEqual( + _kinds(trainer), ["fwd_bwd", "wait_for_device", "update", "fwd_bwd"] + ) + + def test_a_stage_in_between_waits_for_nothing(self): + trainer = st._TokenTrainer() + worker = _deferred(trainer) + with st._RecordingWaits(trainer): + resp = worker.stage_fwd_bwd(st._request("req-2", 2)) + self.assertEqual(_kinds(trainer), ["fwd_bwd", "prepare"]) + self.assertTrue(resp.metadata["staged"]) + self.assertTrue(resp.metadata["before_update"]) + self.assertIsNone(resp.metadata["device_wait_s"]) + self.assertNotIn("update_ready", resp.metadata) + self.assertEqual(worker._staged_fwd_bwd[0], "req-2") + + def test_a_stage_after_a_fused_call_still_waits_and_confirms(self): + # Control: the hook after the last call is unchanged. + trainer = st._TokenTrainer() + worker = st._worker(trainer) + worker.fwd_bwd( + st._request("req-1", 1), apply_optimizer=True, wait_for_update=False + ) + with st._RecordingWaits(trainer): + resp = worker.stage_fwd_bwd(st._request("req-2", 2)) + self.assertEqual( + _kinds(trainer), + [ + "fwd_bwd", + "update", + "prepare", + "wait_for_device", + "block_until_ready", + ], + ) + self.assertNotIn("before_update", resp.metadata) + self.assertEqual(resp.metadata["device_wait_s"], 0.75) + self.assertTrue(resp.metadata["update_ready"]) + + def test_a_caller_cannot_spoof_before_update(self): + trainer = st._WaitingTrainer() + worker = st._worker(trainer) + request = st._request("req-2", 2) + request.metadata["before_update"] = True + resp = worker.stage_fwd_bwd(request) + self.assertNotIn("before_update", resp.metadata) + self.assertEqual(resp.metadata["device_wait_s"], 0.75) + + def test_the_update_waits_for_the_device_first_then_confirms(self): + trainer = st._TokenTrainer(waited_s=4.9) + worker = _deferred(trainer) + worker.stage_fwd_bwd(st._request("req-2", 2)) + with st._RecordingWaits(trainer): + resp = worker.apply_deferred_update() + self.assertEqual( + _kinds(trainer), + [ + "fwd_bwd", + "prepare", + "wait_for_device", + "update", + "block_until_ready", + ], + ) + self.assertIs(trainer.calls[-1][1], trainer.last_update_token) + self.assertTrue(resp.metadata["updated"]) + self.assertEqual(resp.metadata["train_step"], 1) + self.assertEqual(resp.metadata["device_wait_s"], 4.9) + self.assertTrue(resp.metadata["update_ready"]) + self.assertIsNone(worker._deferred_update) + # The staged micro-batch is untouched and runs next. + self.assertEqual(worker._staged_fwd_bwd[0], "req-2") + worker.fwd_bwd_staged("req-2") + self.assertEqual(trainer.calls[-1][0], "fwd_bwd") + self.assertEqual(trainer.calls[-1][1][0], "on_device") + + def test_the_update_without_a_token_is_not_confirmed(self): + trainer = st._WaitingTrainer() + resp = _deferred(trainer).apply_deferred_update() + self.assertEqual(_kinds(trainer), ["fwd_bwd", "wait_for_device", "update"]) + self.assertNotIn("update_ready", resp.metadata) + self.assertEqual(resp.metadata["device_wait_s"], 0.75) + + def test_the_update_without_wait_for_device_reports_none(self): + trainer = st._Trainer() + resp = _deferred(trainer).apply_deferred_update() + self.assertEqual(_kinds(trainer), ["fwd_bwd", "update"]) + self.assertIsNone(resp.metadata["device_wait_s"]) + + def test_nothing_deferred_is_refused(self): + trainer = st._WaitingTrainer() + worker = st._worker(trainer) + with self.assertRaisesRegex(RuntimeError, "holds no deferred update"): + worker.apply_deferred_update() + self.assertEmpty(trainer.calls) + self.assertEqual(worker.state, WorkerState.READY) + + def test_a_second_call_is_refused(self): + trainer = st._WaitingTrainer() + worker = _deferred(trainer) + worker.apply_deferred_update() + with self.assertRaisesRegex(RuntimeError, "holds no deferred update"): + worker.apply_deferred_update() + self.assertEqual(_kinds(trainer).count("update"), 1) + + def test_a_failed_device_wait_is_the_deferred_updates_failure(self): + trainer = st._FailingWaitTrainer() + worker = _deferred(trainer) + with self.assertRaisesRegex( + RuntimeError, + r"deferred optimizer update failed.*waiting for the device:" + r" RuntimeError: device lost", + ): + worker.apply_deferred_update() + self.assertEqual(worker.state, WorkerState.ERROR) + self.assertIn( + "deferred optimizer update failed", worker.heartbeat().last_error + ) + self.assertNotIn("update", _kinds(trainer)) + + def test_a_failed_update_is_the_deferred_updates_failure(self): + trainer = _FailingUpdateTrainer() + worker = _deferred(trainer) + with self.assertRaisesRegex( + RuntimeError, + r"deferred optimizer update failed.*waiting for the update:" + r" RuntimeError: XLA: update failed on device", + ): + worker.apply_deferred_update() + self.assertEqual(worker.state, WorkerState.ERROR) + + @parameterized.named_parameters( + ("compile", lambda w: w.compile(None)), + ("restore_checkpoint", lambda w: w.restore_checkpoint()), + ) + def test_a_successors_bring_up_drops_a_deferred_update(self, bring_up): + trainer = st._RestorableTrainer() + worker = st._worker(trainer) + worker.fwd_bwd( + st._request("req-1", 1), apply_optimizer=True, defer_update=True + ) + with self.assertLogs(level="WARNING") as logs: + bring_up(worker) + self.assertIsNone(worker._deferred_update) + self.assertTrue(any("never applied" in m for m in logs.output), logs.output) + worker.fwd_bwd(st._request("req-2", 2)) + self.assertEqual(_kinds(trainer)[-1], "fwd_bwd") + + def test_initialize_keeps_a_deferred_update(self): + # Control: only a bring-up drops it. + worker = _deferred(st._Trainer()) + worker.initialize() + self.assertEqual(worker._deferred_update, "req-1") + + def test_stop_drops_a_deferred_update(self): + worker = _deferred(st._Trainer()) + worker.stop() + self.assertIsNone(worker._deferred_update) + + def test_info_advertises_it(self): + self.assertTrue( + st._worker(st._Trainer()).info().resources["stage_before_update"] + ) + + @parameterized.parameters(st._PreparingTrainer, st._TokenTrainer) + def test_staged_before_the_update_trains_like_the_fused_sequence( + self, trainer_cls + ): + first, second = st._request("req-1", 1), st._request("req-2", 2) + fused = trainer_cls() + worker = st._worker(fused) + worker.fwd_bwd(first, apply_optimizer=True, wait_for_update=False) + worker.stage_fwd_bwd(second) + worker.fwd_bwd_staged("req-2") + deferred = trainer_cls() + worker = st._worker(deferred) + worker.fwd_bwd(first, apply_optimizer=True, defer_update=True) + worker.stage_fwd_bwd(second) + worker.apply_deferred_update() + worker.fwd_bwd_staged("req-2") + trained = lambda t: [c for c in t.calls if c[0] in ("fwd_bwd", "update")] + self.assertEqual(trained(fused), trained(deferred)) + self.assertEqual( + [c[0] for c in trained(deferred)], ["fwd_bwd", "update", "fwd_bwd"] + ) + + +if __name__ == "__main__": + absltest.main() diff --git a/tests/experimental/worker/trainer_worker_staging_test.py b/tests/experimental/worker/trainer_worker_staging_test.py new file mode 100644 index 0000000000..44c9444213 --- /dev/null +++ b/tests/experimental/worker/trainer_worker_staging_test.py @@ -0,0 +1,549 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for TrainerWorker.stage_fwd_bwd / fwd_bwd_staged. + +A staged micro-batch must train exactly as the same micro-batch sent with +`fwd_bwd`, staging must not touch the trainer's state, and the worker must +never train a staged micro-batch in another's place. +""" + +import asyncio +from typing import Any +from unittest import mock + +from absl.testing import absltest +import jax.numpy as jnp +import numpy as np +from tunix.experimental.common import datatypes +from tunix.experimental.worker import remote_execution +from tunix.experimental.worker import trainer_worker + +WorkerState = datatypes.WorkerState + + +class _Trainer: + """Records what reaches the trainer.""" + + def __init__(self): + self.calls: list[tuple[str, Any, dict[str, Any]]] = [] + self.policy_version = 2 + + def fwd_bwd(self, payload, **kwargs): + self.calls.append(("fwd_bwd", payload, kwargs)) + + def update(self, **kwargs): + self.calls.append(("update", None, kwargs)) + return 1 + + def close(self): + pass + + +class _PreparingTrainer(_Trainer): + """Has `prepare_fwd_bwd_input`, e.g. to place the inputs on its devices.""" + + def prepare_fwd_bwd_input(self, payload): + self.calls.append(("prepare", payload, {})) + return ("on_device", payload) + + +class _FailingPrepareTrainer(_Trainer): + + def prepare_fwd_bwd_input(self, payload): + raise RuntimeError("device_put failed") + + +class _WaitingTrainer(_PreparingTrainer): + """Also has `wait_for_device`, as MaxText's engine does.""" + + def __init__(self, waited_s: float = 0.75): + super().__init__() + self._waited_s = waited_s + + def wait_for_device(self): + self.calls.append(("wait_for_device", None, {})) + return self._waited_s + + +class _FailingWaitTrainer(_Trainer): + + def wait_for_device(self): + raise RuntimeError("device lost") + + +class _TokenTrainer(_WaitingTrainer): + """Its update publishes an output to wait on (`last_update_token`).""" + + def __init__(self, waited_s: float = 0.75): + super().__init__(waited_s=waited_s) + self.last_update_token = None + + def update(self, **kwargs): + step = super().update(**kwargs) + self.last_update_token = jnp.array(float(step)) + return step + + +class _UpdateFailsAtTheDeviceTrainer(_TokenTrainer): + """An update that fails on the device, met by the device wait first. + + MaxText's `wait_for_device` is its throttler's `block_until_ready()`, which + also waits on the update's gradient norm (maxtext 6db77afd6). + """ + + def wait_for_device(self): + if self.last_update_token is not None: + raise RuntimeError("XLA: update 1 failed on device") + return super().wait_for_device() + + +class _RestorableTrainer(_Trainer): + """Also takes the calls a restarted orchestrator brings a worker up with.""" + + def compile(self, dummy_data): + self.calls.append(("compile", dummy_data, {})) + + def restore_checkpoint(self, **kwargs): + self.calls.append(("restore_checkpoint", None, kwargs)) + return {"step": 3} + + +class _RecordingWaits: + """Patches `jax.block_until_ready` in the worker to log each wait.""" + + def __init__(self, trainer: _Trainer): + self._trainer = trainer + self._real = trainer_worker.jax.block_until_ready + self._patch = mock.patch.object( + trainer_worker.jax, "block_until_ready", side_effect=self._wait + ) + + def _wait(self, x): + self._trainer.calls.append(("block_until_ready", x, {})) + return self._real(x) + + def __enter__(self): + self._patch.start() + return self + + def __exit__(self, *exc): + self._patch.stop() + + +def _payload(value: int) -> datatypes.RLTrainerPayload: + return datatypes.RLTrainerPayload( + prompt_ids=np.array([[value, value + 1]], dtype=np.int32), + prompt_mask=np.ones((1, 2), dtype=np.float32), + completion_ids=np.array([[value + 2, value + 3]], dtype=np.int32), + completion_mask=np.ones((1, 2), dtype=np.float32), + advantages=np.array([1.0], dtype=np.float32), + ) + + +def _request(request_id: str, value: int) -> datatypes.TrainRequest: + return datatypes.TrainRequest( + request_id=request_id, + payload=_payload(value), + metadata={"batch_id": f"b{value}"}, + ) + + +def _worker(trainer: _Trainer) -> trainer_worker.TrainerWorker: + worker = trainer_worker.TrainerWorker( + trainer_factory=lambda: trainer, worker_id="trainer_0" + ) + worker.initialize() + return worker + + +class TrainerWorkerStagingTest(absltest.TestCase): + + def test_staged_micro_batch_trains_exactly_like_fwd_bwd(self): + direct, staged = _Trainer(), _Trainer() + request = _request("req-1", 7) + resp_direct = _worker(direct).fwd_bwd(request, skip_jit=True, foo=3) + worker = _worker(staged) + worker.stage_fwd_bwd(request, skip_jit=True, foo=3) + resp_staged = worker.fwd_bwd_staged("req-1") + self.assertEqual(len(direct.calls), len(staged.calls)) + for (kind_a, pay_a, kw_a), (kind_b, pay_b, kw_b) in zip( + direct.calls, staged.calls + ): + self.assertEqual(kind_a, kind_b) + self.assertIs(pay_a, request.payload) + self.assertIs(pay_b, request.payload) + self.assertEqual(kw_a, kw_b) + self.assertEqual(resp_direct.request_id, resp_staged.request_id) + self.assertEqual(resp_direct.metadata, resp_staged.metadata) + + def test_staging_alone_does_not_touch_the_trainer(self): + trainer = _Trainer() + worker = _worker(trainer) + resp = worker.stage_fwd_bwd(_request("req-1", 1)) + self.assertTrue(resp.metadata["staged"]) + self.assertEqual(resp.request_id, "req-1") + self.assertEmpty(trainer.calls) + self.assertEqual(worker.state, WorkerState.READY) + + def test_a_wrong_id_is_refused_and_the_staged_batch_is_kept(self): + trainer = _Trainer() + worker = _worker(trainer) + request = _request("req-1", 1) + worker.stage_fwd_bwd(request) + with self.assertRaisesRegex(RuntimeError, "no staged micro-batch"): + worker.fwd_bwd_staged("req-2") + self.assertEmpty(trainer.calls) + worker.fwd_bwd_staged("req-1") + self.assertIs(trainer.calls[0][1], request.payload) + + def test_a_second_stage_is_refused_rather_than_replacing_the_first(self): + trainer = _Trainer() + worker = _worker(trainer) + first = _request("req-1", 1) + worker.stage_fwd_bwd(first) + with self.assertRaisesRegex(RuntimeError, "already holds"): + worker.stage_fwd_bwd(_request("req-2", 2)) + worker.fwd_bwd_staged("req-1") + self.assertIs(trainer.calls[0][1], first.payload) + + def test_running_without_staging_is_refused(self): + worker = _worker(_Trainer()) + with self.assertRaisesRegex(RuntimeError, "no staged micro-batch"): + worker.fwd_bwd_staged("req-1") + + def test_a_staged_batch_runs_once(self): + worker = _worker(_Trainer()) + worker.stage_fwd_bwd(_request("req-1", 1)) + worker.fwd_bwd_staged("req-1") + with self.assertRaisesRegex(RuntimeError, "no staged micro-batch"): + worker.fwd_bwd_staged("req-1") + + def test_staging_is_allowed_during_a_sync_but_running_is_not(self): + trainer = _Trainer() + worker = _worker(trainer) + worker.state = WorkerState.SYNCING + worker.stage_fwd_bwd(_request("req-1", 1)) + with self.assertRaisesRegex(RuntimeError, "not ready"): + worker.fwd_bwd_staged("req-1") + worker.state = WorkerState.READY + worker.fwd_bwd_staged("req-1") + self.assertLen(trainer.calls, 1) + + def test_staging_is_refused_on_a_failed_or_stopped_worker(self): + for state in (WorkerState.ERROR, WorkerState.STOPPED, WorkerState.DRAINING): + worker = _worker(_Trainer()) + worker.state = state + with self.assertRaisesRegex(RuntimeError, "cannot stage"): + worker.stage_fwd_bwd(_request("req-1", 1)) + + def test_prepare_fwd_bwd_input_is_what_trains(self): + trainer = _PreparingTrainer() + worker = _worker(trainer) + request = _request("req-1", 1) + worker.stage_fwd_bwd(request) + self.assertEqual([c[0] for c in trainer.calls], ["prepare"]) + worker.fwd_bwd_staged("req-1") + self.assertEqual(trainer.calls[1][1], ("on_device", request.payload)) + + def test_a_failed_prepare_marks_the_worker_failed_and_stages_nothing(self): + worker = _worker(_FailingPrepareTrainer()) + with self.assertRaisesRegex(RuntimeError, "device_put failed"): + worker.stage_fwd_bwd(_request("req-1", 1)) + self.assertEqual(worker.state, WorkerState.ERROR) + self.assertIsNone(worker._staged_fwd_bwd) + + def test_stage_reports_how_long_it_waited_for_the_device(self): + trainer = _WaitingTrainer(waited_s=0.75) + worker = _worker(trainer) + resp = worker.stage_fwd_bwd(_request("req-1", 1)) + self.assertEqual(resp.metadata["device_wait_s"], 0.75) + # The micro-batch is in place first, then the wait: the wait measures the + # device work left after the transfer. + self.assertEqual( + [c[0] for c in trainer.calls], ["prepare", "wait_for_device"] + ) + # Waiting changes nothing the trainer trains. + worker.fwd_bwd_staged("req-1") + self.assertEqual(trainer.calls[-1][0], "fwd_bwd") + + def test_stage_without_wait_for_device_reports_none(self): + resp = _worker(_Trainer()).stage_fwd_bwd(_request("req-1", 1)) + self.assertIn("device_wait_s", resp.metadata) + self.assertIsNone(resp.metadata["device_wait_s"]) + + def test_a_failed_device_wait_marks_the_worker_failed_and_stages_nothing( + self, + ): + worker = _worker(_FailingWaitTrainer()) + with self.assertRaisesRegex(RuntimeError, "device lost"): + worker.stage_fwd_bwd(_request("req-1", 1)) + self.assertEqual(worker.state, WorkerState.ERROR) + self.assertIsNone(worker._staged_fwd_bwd) + + def test_info_advertises_staging(self): + with_wait = _worker(_WaitingTrainer()).info().resources + self.assertTrue(with_wait["stage_fwd_bwd"]) + self.assertTrue(with_wait["reports_device_wait"]) + self.assertTrue(with_wait["stage_before_update"]) + without = _worker(_Trainer()).info().resources + self.assertTrue(without["stage_fwd_bwd"]) + self.assertFalse(without["reports_device_wait"]) + self.assertTrue(without["stage_before_update"]) + + def test_stop_drops_the_staged_batch(self): + worker = _worker(_Trainer()) + worker.stage_fwd_bwd(_request("req-1", 1)) + worker.stop() + self.assertIsNone(worker._staged_fwd_bwd) + + def test_staging_through_the_rpc_layer(self): + trainer = _Trainer() + handle = remote_execution.InProcessActorHandle( + remote_execution.InProcessRemoteExecutionServer(_worker(trainer)) + ) + request = _request("req-9", 9) + + async def _run(): + staged = await handle.asubmit("stage_fwd_bwd", request=request) + ran = await handle.asubmit( + "fwd_bwd_staged", staged_request_id="req-9", skip_jit=False + ) + return staged, ran + + staged, ran = asyncio.run(_run()) + self.assertTrue(staged.metadata["staged"]) + self.assertTrue(ran.metadata["queued"]) + self.assertEqual(ran.request_id, "req-9") + np.testing.assert_array_equal( + trainer.calls[0][1].completion_ids, request.payload.completion_ids + ) + + +class FusedUpdateStagingTest(absltest.TestCase): + """Staging on top of the fused final update (fwd_bwd(apply_optimizer=True)).""" + + def test_a_staged_last_micro_batch_runs_the_fused_update_like_fwd_bwd(self): + for trainer_cls in (_Trainer, _TokenTrainer): + for apply_optimizer in (False, True): + direct, staged = trainer_cls(), trainer_cls() + request = _request("req-1", 7) + with _RecordingWaits(direct): + resp_direct = _worker(direct).fwd_bwd( + request, apply_optimizer=apply_optimizer, foo=3 + ) + worker = _worker(staged) + worker.stage_fwd_bwd(request, foo=3) + # Only what fwd_bwd_staged does is compared. + del staged.calls[:] + with _RecordingWaits(staged): + resp_staged = worker.fwd_bwd_staged( + "req-1", apply_optimizer=apply_optimizer + ) + label = (trainer_cls.__name__, apply_optimizer) + self.assertEqual( + [(c[0], c[2]) for c in direct.calls], + [(c[0], c[2]) for c in staged.calls], + label, + ) + self.assertEqual(resp_direct.metadata, resp_staged.metadata, label) + self.assertEqual(resp_direct.request_id, resp_staged.request_id) + self.assertEqual( + bool(resp_staged.metadata.get("updated")), apply_optimizer, label + ) + + def test_wait_for_update_false_returns_at_dispatch_and_claims_nothing(self): + trainer = _TokenTrainer() + worker = _worker(trainer) + with _RecordingWaits(trainer): + resp = worker.fwd_bwd( + _request("req-1", 1), apply_optimizer=True, wait_for_update=False + ) + self.assertEqual([c[0] for c in trainer.calls], ["fwd_bwd", "update"]) + self.assertTrue(resp.metadata["updated"]) + self.assertEqual(resp.metadata["train_step"], 1) + self.assertNotIn("update_ready", resp.metadata) + # It is the worker's option, never the trainer's. + self.assertEqual(trainer.calls[0][2], {"apply_optimizer": True}) + + def test_the_default_still_waits_for_the_update(self): + # Control for the test above: the same call without the option. + trainer = _TokenTrainer() + with _RecordingWaits(trainer): + resp = _worker(trainer).fwd_bwd( + _request("req-1", 1), apply_optimizer=True + ) + self.assertEqual( + [c[0] for c in trainer.calls], + ["fwd_bwd", "update", "block_until_ready"], + ) + self.assertTrue(resp.metadata["update_ready"]) + + def test_wait_for_update_reaches_fwd_bwd_through_a_staged_call(self): + trainer = _TokenTrainer() + worker = _worker(trainer) + worker.stage_fwd_bwd(_request("req-1", 1)) + del trainer.calls[:] + with _RecordingWaits(trainer): + resp = worker.fwd_bwd_staged( + "req-1", apply_optimizer=True, wait_for_update=False + ) + self.assertEqual([c[0] for c in trainer.calls], ["fwd_bwd", "update"]) + self.assertNotIn("update_ready", resp.metadata) + + def test_stage_confirms_the_update_after_its_device_wait(self): + trainer = _TokenTrainer() + worker = _worker(trainer) + worker.fwd_bwd( + _request("req-0", 0), apply_optimizer=True, wait_for_update=False + ) + token = trainer.last_update_token + del trainer.calls[:] + with _RecordingWaits(trainer): + resp = worker.stage_fwd_bwd(_request("req-1", 1)) + self.assertEqual( + [c[0] for c in trainer.calls], + ["prepare", "wait_for_device", "block_until_ready"], + ) + self.assertIs(trainer.calls[2][1], token) + self.assertTrue(resp.metadata["update_ready"]) + self.assertEqual(resp.metadata["device_wait_s"], 0.75) + + def test_stage_without_an_update_token_claims_nothing(self): + trainer = _WaitingTrainer() + worker = _worker(trainer) + with _RecordingWaits(trainer): + resp = worker.stage_fwd_bwd(_request("req-1", 1)) + self.assertNotIn("update_ready", resp.metadata) + self.assertNotIn("block_until_ready", [c[0] for c in trainer.calls]) + + def test_caller_metadata_cannot_claim_for_the_worker(self): + request = datatypes.TrainRequest( + request_id="req-1", + payload=_payload(1), + metadata={ + "batch_id": "b1", + "update_ready": True, + "device_wait_s": 9.0, + "staged": "caller", + }, + ) + resp = _worker(_Trainer()).stage_fwd_bwd(request) + self.assertNotIn("update_ready", resp.metadata) + self.assertIsNone(resp.metadata["device_wait_s"]) + self.assertIs(resp.metadata["staged"], True) + self.assertEqual(resp.metadata["batch_id"], "b1") + + def test_a_failed_update_wait_marks_the_worker_failed_and_stages_nothing( + self, + ): + trainer = _TokenTrainer() + worker = _worker(trainer) + worker.fwd_bwd( + _request("req-0", 0), apply_optimizer=True, wait_for_update=False + ) + with mock.patch.object( + trainer_worker.jax, + "block_until_ready", + side_effect=ValueError("update failed on device"), + ): + # Reported as the optimizer step, as the fused call reports it, in the + # error and in the worker's health report. + with self.assertRaisesRegex( + RuntimeError, + r"^optimizer update\(\) failed: .*stage_fwd_bwd.*the update\b.*" + r"ValueError: update failed on device", + ) as raised: + worker.stage_fwd_bwd(_request("req-1", 1)) + self.assertIsInstance(raised.exception.__cause__, ValueError) + self.assertEqual(worker.state, WorkerState.ERROR) + self.assertIn("optimizer update() failed", worker.heartbeat().last_error) + self.assertIsNone(worker._staged_fwd_bwd) + + def test_a_failed_update_met_by_the_device_wait_is_reported_as_the_update( + self, + ): + trainer = _UpdateFailsAtTheDeviceTrainer() + worker = _worker(trainer) + worker.fwd_bwd( + _request("req-0", 0), apply_optimizer=True, wait_for_update=False + ) + with self.assertRaisesRegex( + RuntimeError, + r"^optimizer update\(\) failed: .*stage_fwd_bwd.*the device\b.*" + r"RuntimeError: XLA: update 1 failed on device", + ) as raised: + worker.stage_fwd_bwd(_request("req-1", 1)) + self.assertIsInstance(raised.exception.__cause__, RuntimeError) + self.assertEqual(worker.state, WorkerState.ERROR) + self.assertIn("optimizer update() failed", worker.heartbeat().last_error) + self.assertIsNone(worker._staged_fwd_bwd) + + def test_a_failed_prepare_is_not_reported_as_the_update(self): + # Control: only the waits for the work dispatched before the stage are + # the update's; preparing the stage's own input is the stage's. + worker = _worker(_FailingPrepareTrainer()) + with self.assertRaisesRegex(RuntimeError, "^device_put failed$"): + worker.stage_fwd_bwd(_request("req-1", 1)) + self.assertEqual(worker.heartbeat().last_error, "device_put failed") + + +class StaleStageTest(absltest.TestCase): + """A stage left by an orchestrator that died without stop() (OOM kill, + preemption) is dropped when a successor brings the worker up again.""" + + def _held(self) -> trainer_worker.TrainerWorker: + worker = _worker(_RestorableTrainer()) + worker.stage_fwd_bwd(_request("dead-orchestrator", 1)) + return worker + + def test_initialize_alone_keeps_the_held_stage(self): + # Control: a held stage is refused a successor, and initialize() (a no-op + # on a ready worker) does not change that. + worker = self._held() + worker.initialize() + with self.assertRaisesRegex( + RuntimeError, "already holds staged micro-batch 'dead-orchestrator'" + ): + worker.stage_fwd_bwd(_request("req-2", 2)) + + def test_compile_drops_the_held_stage(self): + worker = self._held() + with self.assertLogs(level="WARNING") as logs: + worker.compile(None) + self.assertIsNone(worker._staged_fwd_bwd) + self.assertIn("dead-orchestrator", logs.output[0]) + self.assertTrue( + worker.stage_fwd_bwd(_request("req-2", 2)).metadata["staged"] + ) + + def test_restore_checkpoint_drops_the_held_stage(self): + worker = self._held() + with self.assertLogs(level="WARNING") as logs: + self.assertEqual(worker.restore_checkpoint(), {"step": 3}) + self.assertIsNone(worker._staged_fwd_bwd) + self.assertIn("dead-orchestrator", logs.output[0]) + self.assertTrue( + worker.stage_fwd_bwd(_request("req-2", 2)).metadata["staged"] + ) + + def test_a_dropped_stage_cannot_be_run(self): + worker = self._held() + worker.restore_checkpoint() + with self.assertRaisesRegex(RuntimeError, "no staged micro-batch"): + worker.fwd_bwd_staged("dead-orchestrator") + + +if __name__ == "__main__": + absltest.main() diff --git a/tests/rl/sharded_logps_head_test.py b/tests/rl/sharded_logps_head_test.py new file mode 100644 index 0000000000..3ca056ee2b --- /dev/null +++ b/tests/rl/sharded_logps_head_test.py @@ -0,0 +1,315 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for the sharded chunked log-prob head (TUNIX_LOGPS_SHARDED_HEAD).""" + +import os + +os.environ['XLA_FLAGS'] = '--xla_force_host_platform_device_count=8' + +from unittest import mock # pylint: disable=g-import-not-at-top + +from absl.testing import absltest +from absl.testing import parameterized +from flax import nnx +import flax.linen as nn +import jax +import jax.numpy as jnp +from jax.sharding import Mesh, NamedSharding, PartitionSpec as P # pylint: disable=g-multiple-import,g-importing-member +import numpy as np +from tunix.rl import common + +_B, _T, _H, _V = 4, 64, 32, 128 +_AXES = ('data', 'fsdp', 'context', 'expert') +_RULES = ( + ('activation_batch', ('data', 'fsdp', 'expert')), + ('activation_length', 'context'), +) + + +class _HeadOnly(nnx.Module): + """Frozen hidden states into a MaxText-like head (RMSNorm, f32 table).""" + + def __init__(self, dtype, batch=_B, seq=_T): + keys = jax.random.split(jax.random.PRNGKey(0), 3) + self.dtype = dtype + self.h = nnx.Param( + jax.random.normal(keys[0], (batch, seq, _H)).astype(dtype) + ) + self.scale = nnx.Param( + 0.2 * jax.random.normal(keys[1], (_H,)).astype(dtype) + ) + self.table = nnx.Param(3.0 * jax.random.normal(keys[2], (_H, _V)) / _H**0.5) + + def __call__( + self, + input_tokens, + positions, + cache, + attention_mask, + segment_ids=None, + skip_lm_head=False, + ): + del input_tokens, positions, cache, attention_mask, segment_ids + h = self.h[...] + return (h if skip_lm_head else self.compute_final_logits(h)), None + + def compute_final_logits(self, h): + x = h.astype(jnp.float32) + y = x * jax.lax.rsqrt(jnp.mean(x * x, -1, keepdims=True) + 1e-6) + y = (y * (self.scale[...].astype(jnp.float32) + 1.0)).astype(self.dtype) + out = jnp.dot(y, self.table[...].astype(self.dtype)) + return out.astype(jnp.float32) + + +def _mesh(): + devices = np.array(jax.devices()[:8]).reshape(1, 2, 2, 2) + return Mesh(devices, _AXES) + + +class ShardedLogpsHeadTest(parameterized.TestCase): + + def setUp(self): + super().setUp() + self.mesh = _mesh() + + def _inputs(self, batch, seq): + rng = np.random.default_rng(0) + tokens = jnp.asarray(rng.integers(1, _V, (batch, seq)), jnp.int32) + segment_ids = jnp.ones((batch, seq), jnp.int32) + positions = jnp.asarray(np.tile(np.arange(seq), (batch, 1)), jnp.int32) + weights = jnp.asarray(rng.normal(size=(2, batch, seq)), jnp.float32) + if batch % 4 == 0 and seq % 2 == 0: + sharding = NamedSharding( + self.mesh, P(('data', 'fsdp', 'expert'), 'context') + ) + tokens, segment_ids, positions = ( + jax.device_put(x, sharding) for x in (tokens, segment_ids, positions) + ) + return tokens, segment_ids, positions, weights + + def _run( + self, + mode, + dtype=jnp.float32, + chunk_size=8, + rules=_RULES, + temperature=0.7, + batch=_B, + seq=_T, + return_entropy=True, + ): + model = _HeadOnly(dtype, batch, seq) + graphdef, state = nnx.split(model) + tokens, segment_ids, positions, weights = self._inputs(batch, seq) + + def loss(state): + out = common.compute_per_token_logps( + graphdef, + state, + jnp.zeros((tokens.shape[0], 0), jnp.int32), + tokens, + pad_id=0, + eos_id=-1, + stop_gradient=False, + return_entropy=return_entropy, + segment_ids=segment_ids, + segment_positions=positions, + temperature=temperature, + chunk_size=chunk_size, + ) + logps, entropy = out if return_entropy else (out, jnp.zeros_like(out)) + value = jnp.sum(logps * weights[0]) + jnp.sum(entropy * weights[1]) + return value, (logps, entropy) + + jax.clear_caches() + env = {} if mode is None else {common.SHARDED_HEAD_ENV: mode} + with mock.patch.dict(os.environ, env, clear=False): + if mode is None: + os.environ.pop(common.SHARDED_HEAD_ENV, None) + with jax.set_mesh(self.mesh), nn.logical_axis_rules(rules): + (_, (logps, entropy)), grads = jax.jit( + jax.value_and_grad(loss, has_aux=True) + )(state) + return logps, entropy, grads + + @parameterized.named_parameters( + ('local', 'local', 8), + ('local_uneven_chunk', 'local', 12), + ('local_one_step', 'local', 64), + ) + def test_matches_scan_head(self, mode, chunk_size): + want = self._run(None, chunk_size=chunk_size) + got = self._run(mode, chunk_size=chunk_size) + np.testing.assert_allclose(got[0], want[0], rtol=1e-5, atol=2e-5) + np.testing.assert_allclose(got[1], want[1], rtol=1e-5, atol=2e-5) + for name in ('h', 'scale', 'table'): + np.testing.assert_allclose( + got[2][name][...], + want[2][name][...], + rtol=1e-4, + atol=1e-5, + err_msg=name, + ) + + def test_logps_only_matches_scan_head(self): + # return_entropy=False: one output, and no entropy cotangent in the bwd. + want = self._run(None, return_entropy=False) + got = self._run('local', return_entropy=False) + np.testing.assert_allclose(got[0], want[0], rtol=1e-5, atol=2e-5) + for name in ('h', 'scale', 'table'): + np.testing.assert_allclose( + got[2][name][...], + want[2][name][...], + rtol=1e-4, + atol=1e-5, + err_msg=name, + ) + + @parameterized.parameters(True, False) + def test_forward_only_matches_scan_head(self, return_entropy): + # No gradient: the custom_vjp's primal body runs instead of fwd/bwd. + model = _HeadOnly(jnp.float32) + graphdef, state = nnx.split(model) + tokens, segment_ids, positions, _ = self._inputs(_B, _T) + + def logps(mode): + jax.clear_caches() + env = {} if mode is None else {common.SHARDED_HEAD_ENV: mode} + with mock.patch.dict(os.environ, env, clear=False): + if mode is None: + os.environ.pop(common.SHARDED_HEAD_ENV, None) + with jax.set_mesh(self.mesh), nn.logical_axis_rules(_RULES): + return jax.jit( + lambda st: common.compute_per_token_logps( + graphdef, + st, + jnp.zeros((_B, 0), jnp.int32), + tokens, + pad_id=0, + eos_id=-1, + stop_gradient=True, + return_entropy=return_entropy, + segment_ids=segment_ids, + segment_positions=positions, + temperature=1.0, + chunk_size=8, + ) + )(state) + + want, got = logps(None), logps('local') + for g, w in zip(jax.tree.leaves(got), jax.tree.leaves(want)): + np.testing.assert_allclose(g, w, rtol=1e-5, atol=2e-5) + + def test_bf16_matches_scan_head(self): + # Same per-token math; XLA may order the per-token f32 reductions + # differently for the smaller per-device chunk, so allow f32 rounding. + want = self._run(None, dtype=jnp.bfloat16) + got = self._run('local', dtype=jnp.bfloat16) + np.testing.assert_allclose(got[0], want[0], rtol=1e-6, atol=1e-6) + np.testing.assert_allclose(got[1], want[1], rtol=1e-6, atol=1e-6) + np.testing.assert_allclose( + np.asarray(got[2]['h'][...], np.float32), + np.asarray(want[2]['h'][...], np.float32), + rtol=1e-2, + atol=1e-6, + ) + for name in ('scale', 'table'): + g = np.asarray(got[2][name][...], np.float32) + w = np.asarray(want[2][name][...], np.float32) + self.assertLess(np.linalg.norm(g - w) / np.linalg.norm(w), 2.0**-6, name) + + def test_mesh_order_layout_without_rules(self): + want = self._run(None, rules=()) + got = self._run('local', rules=()) + np.testing.assert_allclose(got[0], want[0], rtol=1e-5, atol=2e-5) + np.testing.assert_allclose( + got[2]['table'][...], want[2]['table'][...], rtol=1e-4, atol=1e-5 + ) + + def test_engages_on_mesh(self): + calls = [] + real = common.maybe_compute_sharded_chunked_logps + + def spy(*args, **kwargs): + out = real(*args, **kwargs) + calls.append(out is not None) + return out + + with mock.patch.object(common, 'maybe_compute_sharded_chunked_logps', spy): + self._run('local') + self._run(None) + self.assertEqual(calls, [True, False]) + + def _engaged(self, **kwargs): + calls = [] + real = common.maybe_compute_sharded_chunked_logps + + def spy(*args, **kw): + out = real(*args, **kw) + calls.append(out is not None) + return out + + with mock.patch.object(common, 'maybe_compute_sharded_chunked_logps', spy): + got = self._run('local', **kwargs) + return calls, got + + def test_sequence_only_layout_when_batch_does_not_divide(self): + # 3 rows cannot be split over fsdp x expert, so all 8 devices take a + # slice of every row's sequence instead. + calls, got = self._engaged(batch=3) + want = self._run(None, batch=3) + self.assertEqual(calls, [True]) + np.testing.assert_allclose(got[0], want[0], rtol=1e-5, atol=2e-5) + np.testing.assert_allclose( + got[2]['table'][...], want[2]['table'][...], rtol=1e-4, atol=1e-5 + ) + + def test_falls_back_when_no_layout(self): + # 3 rows of 63 tokens divide over no mesh axis: compute_chunked_logps runs. + calls, got = self._engaged(batch=3, seq=63, chunk_size=7) + want = self._run(None, batch=3, seq=63, chunk_size=7) + self.assertEqual(calls, [False]) + np.testing.assert_array_equal(got[0], want[0]) + + def test_falls_back_without_mesh(self): + model = _HeadOnly(jnp.float32) + tokens = self._inputs(_B, _T)[0] + with mock.patch.dict(os.environ, {common.SHARDED_HEAD_ENV: '1'}): + self.assertIsNone( + common.maybe_compute_sharded_chunked_logps( + model, model.h[...], tokens, 1.0, 8, True + ) + ) + + @parameterized.parameters( + ('', None), + ('0', None), + ('false', None), + ('1', 'local'), + ('local', 'local'), + ) + def test_mode_parsing(self, value, want): + with mock.patch.dict(os.environ, {common.SHARDED_HEAD_ENV: value}): + self.assertEqual(common.sharded_head_mode(), want) + + @parameterized.parameters('locla', 'per_step') + def test_mode_parsing_rejects_unknown(self, value): + with mock.patch.dict(os.environ, {common.SHARDED_HEAD_ENV: value}): + with self.assertRaises(ValueError): + common.sharded_head_mode() + + +if __name__ == '__main__': + absltest.main() diff --git a/tunix/experimental/examples/deepswe_dist/k8s_launcher.sh b/tunix/experimental/examples/deepswe_dist/k8s_launcher.sh index 108740ae00..fd5a821112 100755 --- a/tunix/experimental/examples/deepswe_dist/k8s_launcher.sh +++ b/tunix/experimental/examples/deepswe_dist/k8s_launcher.sh @@ -78,6 +78,8 @@ export USE_ROLLOUT_LOGPS=${USE_ROLLOUT_LOGPS:-true} export EXACT_TOKEN_CONTINUITY=${EXACT_TOKEN_CONTINUITY:-true} export ASYNC_WEIGHT_SYNC=${ASYNC_WEIGHT_SYNC:-false} export PIPELINE_TRAIN_MICROBATCHES=${PIPELINE_TRAIN_MICROBATCHES:-false} +export PIPELINE_STEP_BOUNDARY=${PIPELINE_STEP_BOUNDARY:-false} +export STAGE_NEXT_MICROBATCH=${STAGE_NEXT_MICROBATCH:-false} export SAMPLER=${SAMPLER:-inprocess_vllm} export WEIGHT_SYNC_MODE=${WEIGHT_SYNC_MODE:-none} export WEIGHT_SYNC_DISABLE_TIMEOUTS=${WEIGHT_SYNC_DISABLE_TIMEOUTS:-${DISABLE_WEIGHT_SYNC_TIMEOUTS:-0}} @@ -287,6 +289,10 @@ export PREEMPTIBLE=${PREEMPTIBLE:-${preemptible:-false}} export PRIORITY_CLASS=${PRIORITY_CLASS:-medium} export TRAINER_EXTRA_ENV=${TRAINER_EXTRA_ENV:-} +# The trainer process's opt-in sharded log-prob head (TUNIX_LOGPS_SHARDED_HEAD +# in tunix/rl/common.py: local turns it on; unset or 0 leaves it off). Passed to +# the trainer only, and only when set. +export TUNIX_LOGPS_SHARDED_HEAD=${TUNIX_LOGPS_SHARDED_HEAD:-} # Extra KEY=VALUE env for the orchestrator and rollout processes (the recipes set # Raiden tuning and weight-sync timeouts here). The VLLM_RAY_* lists tell vLLM # which of the rollout driver's env vars to forward to its Ray workers on the @@ -564,6 +570,8 @@ start_orchestrator() { $([[ "${EXACT_TOKEN_CONTINUITY}" == "false" || "${EXACT_TOKEN_CONTINUITY}" == "False" || "${EXACT_TOKEN_CONTINUITY}" == "0" ]] && echo --no-exact_token_continuity || echo --exact_token_continuity) \ $([[ "${ASYNC_WEIGHT_SYNC}" == "true" || "${ASYNC_WEIGHT_SYNC}" == "True" || "${ASYNC_WEIGHT_SYNC}" == "1" ]] && echo --async_weight_sync || echo --no-async_weight_sync) \ $([[ "${PIPELINE_TRAIN_MICROBATCHES}" == "true" || "${PIPELINE_TRAIN_MICROBATCHES}" == "True" || "${PIPELINE_TRAIN_MICROBATCHES}" == "1" ]] && echo --pipeline_train_microbatches || echo --no-pipeline_train_microbatches) \ + $([[ "${PIPELINE_STEP_BOUNDARY}" == "true" || "${PIPELINE_STEP_BOUNDARY}" == "True" || "${PIPELINE_STEP_BOUNDARY}" == "1" ]] && echo --pipeline_step_boundary || echo --no-pipeline_step_boundary) \ + $([[ "${STAGE_NEXT_MICROBATCH}" == "true" || "${STAGE_NEXT_MICROBATCH}" == "True" || "${STAGE_NEXT_MICROBATCH}" == "1" ]] && echo --stage_next_microbatch || echo --no-stage_next_microbatch) \ $([[ "${IN_FLIGHT_WEIGHT_UPDATES}" == "true" || "${IN_FLIGHT_WEIGHT_UPDATES}" == "True" || "${IN_FLIGHT_WEIGHT_UPDATES}" == "1" ]] && echo --in_flight_weight_updates || echo --no-in_flight_weight_updates) \ $([[ "${ROLLOUT_PRIORITY_SCHEDULING}" == "true" ]] && echo --rollout_priority_scheduling || echo --no-rollout_priority_scheduling) \ $([[ "${INITIAL_DISPATCH_STAGGER_S}" != "0" && "${INITIAL_DISPATCH_STAGGER_S}" != "0.0" && -n "${INITIAL_DISPATCH_STAGGER_S}" ]] && echo --initial_dispatch_stagger_s=${INITIAL_DISPATCH_STAGGER_S})$([[ -n "${INITIAL_DISPATCH_FIRST_STAGGER_S}" && "${INITIAL_DISPATCH_FIRST_STAGGER_S}" != "0" && "${INITIAL_DISPATCH_FIRST_STAGGER_S}" != "0.0" ]] && echo " --initial_dispatch_first_stagger_s=${INITIAL_DISPATCH_FIRST_STAGGER_S}")$([[ "${TRAINER_WARMUP_COMPILE}" == "true" || "${TRAINER_WARMUP_COMPILE}" == "True" || "${TRAINER_WARMUP_COMPILE}" == "1" ]] && echo " --trainer_warmup_compile$([[ "${RETURN_ROUTED_EXPERTS}" == "true" || "${RETURN_ROUTED_EXPERTS}" == "True" || "${RETURN_ROUTED_EXPERTS}" == "1" ]] && echo " --trainer_warmup_routed_experts")") \ @@ -672,6 +680,7 @@ start_trainer() { ${TRAINER_MAXTEXT_ATTENTION:+TRAINER_MAXTEXT_ATTENTION=\"${TRAINER_MAXTEXT_ATTENTION}\"} \ ${raiden_env} \ ${MAXTEXT_EXTRA_FLAGS:+MAXTEXT_EXTRA_FLAGS=\"${MAXTEXT_EXTRA_FLAGS}\"} \ + ${TUNIX_LOGPS_SHARDED_HEAD:+TUNIX_LOGPS_SHARDED_HEAD=${TUNIX_LOGPS_SHARDED_HEAD}} \ ${TRAINER_EXTRA_ENV:+${TRAINER_EXTRA_ENV}} \ RAIDEN_DEVICES_PER_HOST=${RAIDEN_DEVICES_PER_HOST} \ USE_WEIGHT_CONVERTER=${USE_WEIGHT_CONVERTER} \ diff --git a/tunix/experimental/examples/deepswe_dist/run_deepswe_dist.py b/tunix/experimental/examples/deepswe_dist/run_deepswe_dist.py index 64c7c101a4..abe887a20a 100644 --- a/tunix/experimental/examples/deepswe_dist/run_deepswe_dist.py +++ b/tunix/experimental/examples/deepswe_dist/run_deepswe_dist.py @@ -298,6 +298,43 @@ def _parse_args(argv: list[str]) -> argparse.Namespace: " previous one's train_step is still running on the trainer." ), ) + parser.add_argument( + "--pipeline_step_boundary", + action=argparse.BooleanOptionalAction, + default=False, + help=( + "Pipeline the step boundary: pack the next step's first" + " micro-batch while this step's last ones train, send it as soon" + " as this step's checkpoint, weight sync (its staging, with" + " --async_weight_sync) and commit are done, and do this step's" + " metrics and logging once that call has returned, while the" + " device computes it. The trainer gets the same calls in the same" + " order." + ), + ) + parser.add_argument( + "--stage_next_microbatch", + action=argparse.BooleanOptionalAction, + default=False, + help=( + "With --pipeline_step_boundary: send the next step's first" + " micro-batch to the trainer host while the device still computes" + " this step's last one, so its train call carries no data." + ), + ) + parser.add_argument( + "--stage_before_update", + action=argparse.BooleanOptionalAction, + default=True, + help=( + "With --stage_next_microbatch: stage between this step's last" + " micro-batch and its update, which the trainer then defers, rather" + " than after the update's dispatch. Needed wherever the trainer's" + " update waits for the last micro-batch (MaxText" + " retry_when_tokens_dropped); --no-stage_before_update stages" + " after the update's dispatch." + ), + ) parser.add_argument( "--in_flight_weight_updates", action=argparse.BooleanOptionalAction, @@ -1283,6 +1320,9 @@ def _on_train_start(step: int) -> None: ), async_weight_sync=args.async_weight_sync, pipeline_train_microbatches=args.pipeline_train_microbatches, + pipeline_step_boundary=args.pipeline_step_boundary, + stage_next_microbatch=args.stage_next_microbatch, + stage_before_update=args.stage_before_update, partial_rollout=args.in_flight_weight_updates, rollout_priority_scheduling=args.rollout_priority_scheduling, initial_dispatch_stagger_s=args.initial_dispatch_stagger_s, diff --git a/tunix/experimental/examples/recipes/mlperf_397b_256_v7x.sh b/tunix/experimental/examples/recipes/mlperf_397b_256_v7x.sh index 82234257fa..e9b45dbe4d 100755 --- a/tunix/experimental/examples/recipes/mlperf_397b_256_v7x.sh +++ b/tunix/experimental/examples/recipes/mlperf_397b_256_v7x.sh @@ -100,6 +100,29 @@ export TRAINER_LIBTPU_INIT_ARGS="${TRAINER_LIBTPU_INIT_ARGS:---DANGEROUS_tpu_run if [[ -n "${TRAINER_MAX_HBM_SIZE_MIB:-}" ]]; then export TRAINER_LIBTPU_INIT_ARGS="${TRAINER_LIBTPU_INIT_ARGS} --xla_tpu_max_hbm_size_mib=${TRAINER_MAX_HBM_SIZE_MIB}" fi +# TRAINER_TUNED_FLAGS (true or false): trainer settings measured on the RL trainer +# bench (AI-Hypercomputer/maxtext#5701) at this recipe's 128-chip trainer cell. Most +# are the settings of runs 59 and 62 in AI-Hypercomputer/maxtext#5409 that this recipe +# does not use yet. Off here and in the recipes that source this one; on in +# mlperf_gbs1024_submission.sh. Part 1, the trainer libtpu flag, is appended here so +# the trainer env, the Pathways workers and the proxy (which compiles the trainer +# program) all copy it. A value already in TRAINER_LIBTPU_INIT_ARGS is kept. +# --xla_tpu_enable_offloading_copy_to_sparsecore=false: XLA does not offload copies +# to the SparseCores. Bench, alone: 130.8 -> 126.4 s/step. +# Part 2, the MaxText keys, follows MAXTEXT_EXTRA_FLAGS below. +case "${TRAINER_TUNED_FLAGS:-false}" in + true) + if [[ " ${TRAINER_LIBTPU_INIT_ARGS} " != *" --xla_tpu_enable_offloading_copy_to_sparsecore="* ]]; then + export TRAINER_LIBTPU_INIT_ARGS="${TRAINER_LIBTPU_INIT_ARGS} --xla_tpu_enable_offloading_copy_to_sparsecore=false" + fi + ;; + false) + ;; + *) + echo "ERROR: TRAINER_TUNED_FLAGS must be true or false, got '${TRAINER_TUNED_FLAGS}'" >&2 + exit 1 + ;; +esac export TRAINER_EXTRA_ENV="${TRAINER_EXTRA_ENV:-ONEHOT_MOE_PERMUTE_THRESHOLD=${ONEHOT_MOE_PERMUTE_THRESHOLD} ENABLE_MULTI_NUMA=${ENABLE_MULTI_NUMA} TPU_RAIDEN_DATA_NICS=${TPU_RAIDEN_DATA_NICS} RAIDEN_BROADCAST_HOST_RATIO=${RAIDEN_BROADCAST_HOST_RATIO} RAIDEN_BROADCAST_PIPELINE_STAGES=${RAIDEN_BROADCAST_PIPELINE_STAGES} LIBTPU_INIT_ARGS='${TRAINER_LIBTPU_INIT_ARGS}'${TPU_RAIDEN_TCP_CONNECT_TIMEOUT_MS:+ TPU_RAIDEN_TCP_CONNECT_TIMEOUT_MS=${TPU_RAIDEN_TCP_CONNECT_TIMEOUT_MS}}${TPU_RAIDEN_TCP_CONNECT_MAX_ATTEMPTS:+ TPU_RAIDEN_TCP_CONNECT_MAX_ATTEMPTS=${TPU_RAIDEN_TCP_CONNECT_MAX_ATTEMPTS}}}" # Under Pathways the trainer's TPU program runs in the pathways-worker container, # so the trainer libtpu flags must be set there (yaml_generator.py renders one @@ -189,6 +212,32 @@ fi if [[ "${LOG_MOE_DROPS:-false}" == "true" ]]; then export MAXTEXT_EXTRA_FLAGS="${MAXTEXT_EXTRA_FLAGS} log_moe_drops=true" fi +# TRAINER_TUNED_FLAGS=true, part 2 (part 1 is with the trainer libtpu flags above): the +# MaxText keys. Appended last: MaxText keeps the last value of a repeated key, so they +# 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. +# 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=device: keeps the ragged sort's integer routing results, so the backward +# does not replay both argsorts. Not in the bench numbers: the bench image's MaxText +# tags nothing with it. +# 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 seven keys together, on top of the libtpu flag (bench, same pod): 126.4 -> 121.7 +# s/step. +# gdn_cp_matmul_precision=high: seq mode composes GDN state in 3 bf16 passes, not 6. +# cast_grads_after_all_reduce=true: all-reduce-only bf16 grads cast to f32 after it. +# The two together (bench, same pod, google/tunix#2873 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). +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 moe_route=device gdn_cp_matmul_precision=high cast_grads_after_all_reduce=true" +fi export DEBUG=${DEBUG:-0} # DeepSWE Environment & Agent Sandbox diff --git a/tunix/experimental/examples/recipes/mlperf_base.sh b/tunix/experimental/examples/recipes/mlperf_base.sh index b169e1470f..61eee0ac7a 100755 --- a/tunix/experimental/examples/recipes/mlperf_base.sh +++ b/tunix/experimental/examples/recipes/mlperf_base.sh @@ -357,6 +357,26 @@ if [[ "${ASYNC_WEIGHT_SYNC}" == "true" && "${WEIGHT_SYNC_MODE}" == "raiden" \ fi # Pack the next microbatch while the trainer runs the current one. export PIPELINE_TRAIN_MICROBATCHES=${PIPELINE_TRAIN_MICROBATCHES:-false} +# Pipeline the step boundary: pack and send the next step's first microbatch +# without waiting for this step's metrics and logging (same trainer calls, in +# the same order). STAGE_NEXT_MICROBATCH also ships its data to the trainer +# while the previous step's last microbatch still computes. +export PIPELINE_STEP_BOUNDARY=${PIPELINE_STEP_BOUNDARY:-false} +export STAGE_NEXT_MICROBATCH=${STAGE_NEXT_MICROBATCH:-false} +# Same truthiness as k8s_launcher.sh: true, True or 1 turns a flag on. +case "${PIPELINE_STEP_BOUNDARY}" in + true|True|1) _pipeline_step_boundary_on=true ;; + *) _pipeline_step_boundary_on=false ;; +esac +case "${STAGE_NEXT_MICROBATCH}" in + true|True|1) _stage_next_microbatch_on=true ;; + *) _stage_next_microbatch_on=false ;; +esac +if [[ "${_stage_next_microbatch_on}" == "true" && "${_pipeline_step_boundary_on}" != "true" ]]; then + echo "STAGE_NEXT_MICROBATCH=${STAGE_NEXT_MICROBATCH} requires PIPELINE_STEP_BOUNDARY on (true, True or 1), got '${PIPELINE_STEP_BOUNDARY}'" >&2 + exit 1 +fi +unset _pipeline_step_boundary_on _stage_next_microbatch_on export TRAJECTORY_GROUP_ORDER=${TRAJECTORY_GROUP_ORDER:-prompt_batch} export USE_DEFAULT_SQL_STORE="${USE_DEFAULT_SQL_STORE:-true}" if [[ "${DRY_RUN:-false}" != "true" ]] && [[ -z "${TRAJECTORY_STORE_DB_URL:-}" ]] && [[ "${USE_DEFAULT_SQL_STORE}" =~ ^(1|true|True|TRUE)$ ]]; then diff --git a/tunix/experimental/examples/recipes/mlperf_gbs1024_submission.sh b/tunix/experimental/examples/recipes/mlperf_gbs1024_submission.sh index cf8f4b17dc..d8ad93d370 100644 --- a/tunix/experimental/examples/recipes/mlperf_gbs1024_submission.sh +++ b/tunix/experimental/examples/recipes/mlperf_gbs1024_submission.sh @@ -66,12 +66,22 @@ export ROLLOUT_ATTENTION_DP="${ROLLOUT_ATTENTION_DP:-false}" # RETRY_DROPLESS_FIRST_STEPS, FIRST_PHASE_RAGGED_BUFFER_FACTOR; maxtext #5570), # which rerun an overflowing micro-batch with a dropless program. For example, # RAGGED_BUFFER_FACTOR=2.75 RETRY_WHEN_TOKENS_DROPPED=true -# keeps this block's context offload, four chunks and HBM budget. +# keeps this block's context offload and HBM budget; it keeps four chunks only +# with TRAINER_TUNED_FLAGS=false (below). +# +# TRAINER_TUNED_FLAGS (see mlperf_397b_256_v7x.sh) is on in this recipe; `false` +# renders the trainer as before. It was measured with the replay layout above, and +# there it also sets two MoE token chunks. The fully dropless buffers +# (RAGGED_BUFFER_FACTOR=-1) keep four. +export TRAINER_TUNED_FLAGS="${TRAINER_TUNED_FLAGS:-true}" export DROPLESS="${DROPLESS:-true}" case "${DROPLESS}" in true) export RAGGED_BUFFER_FACTOR="${RAGGED_BUFFER_FACTOR:--1}" export REMAT_CONTEXT="${REMAT_CONTEXT:-offload}" + if [[ "${TRAINER_TUNED_FLAGS}" == "true" && "${RAGGED_BUFFER_FACTOR}" != "-1" ]]; then + export NUM_MOE_TOKEN_CHUNKS="${NUM_MOE_TOKEN_CHUNKS:-2}" + fi export NUM_MOE_TOKEN_CHUNKS="${NUM_MOE_TOKEN_CHUNKS:-4}" export TRAINER_MAX_HBM_SIZE_MIB="${TRAINER_MAX_HBM_SIZE_MIB:-84992}" if [[ "${RAGGED_BUFFER_FACTOR}" != "-1" && "${RETRY_WHEN_TOKENS_DROPPED:-false}" != "true" ]]; then @@ -88,6 +98,22 @@ case "${DROPLESS}" in ;; esac +# The sharded log-prob head (google/tunix#2873) and the staged step boundary +# (google/tunix#2877), 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. +# 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 +# AI-Hypercomputer/maxtext#5718; without it the orchestrator warns and runs +# the boundary pipelined but unstaged. STAGE_NEXT_MICROBATCH follows +# PIPELINE_STEP_BOUNDARY unless set, so PIPELINE_STEP_BOUNDARY=false alone +# turns both off. +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}}" + export TRAINER_FP8="${TRAINER_FP8:-true}" export ROLLOUT_FP8="${ROLLOUT_FP8:-true}" export ROLLOUT_REPLICAS="${ROLLOUT_REPLICAS:-128}" diff --git a/tunix/experimental/orchestrator/distributed_rl_engine.py b/tunix/experimental/orchestrator/distributed_rl_engine.py index dd7d5d12e3..ca81859a6c 100644 --- a/tunix/experimental/orchestrator/distributed_rl_engine.py +++ b/tunix/experimental/orchestrator/distributed_rl_engine.py @@ -84,6 +84,27 @@ def _response_to_trajectory_item(resp: Any) -> datatypes.TrajectoryItem: raise ValueError("RolloutResponse payload is None.") +@dataclasses.dataclass(frozen=True) +class StagedTrainStep: + """A micro-batch the trainer holds ahead of its train_step (`stage_train_step`). + + Attributes: + request_id: Pass as `train_step(None, staged_request_id=...)`. + device_wait_s: How long the trainer, once it held the micro-batch, still + waited for the work it had dispatched before: ~0 means its device was + already idle, so the transfer did not hide behind that work. None when + the trainer cannot tell. + update_ready: The trainer confirmed that the optimizer update dispatched + before the stage has landed on device (it exposes `last_update_token`). + Like `update_ready` on a fused train_step result, it lets the caller + stamp and save the step's checkpoint without a get_metrics() first. + """ + + request_id: str + device_wait_s: float | None = None + update_ready: bool = False + + class DistributedRLEngine(rl_engine_interface.AbstractRLEngine): """Worker-backed compute router dispatching RPCs across role pools.""" @@ -844,19 +865,165 @@ async def compile_trainer( raise ValueError(f"No trainer worker registered for role {role}") return await self._invoke_worker(worker, "compile", dummy_data=dummy_data) - async def train_step( + async def trainer_staging_support( + self, role: datatypes.Role = datatypes.Role.ACTOR + ) -> dict[str, bool]: + """What the role's trainer supports of `stage_train_step`, from its `info()`. + + Asked before anything is staged, so a trainer image that predates staging + is found without sending it a micro-batch it cannot hold. + + Args: + role: Trainer role to ask. + + Returns: + `stage_fwd_bwd`: it holds a micro-batch ahead of its fwd_bwd. + `reports_device_wait`: it reports `StagedTrainStep.device_wait_s`. + `stage_before_update`: it can defer a step's update past its last + micro-batch (`train_step(defer_update=True)` and + `apply_deferred_update`), so a stage runs in between. + """ + worker = self._trainer_workers.get(role) + if worker is None: + raise ValueError(f"No trainer worker registered for role {role}") + info = await self._invoke_worker(worker, "info") + resources = getattr(info, "resources", None) or {} + return { + key: bool(resources.get(key, False)) + for key in ( + "stage_fwd_bwd", + "reports_device_wait", + "stage_before_update", + ) + } + + async def apply_deferred_update( + self, role: datatypes.Role = datatypes.Role.ACTOR + ) -> dict[str, Any]: + """Runs the update a `train_step(defer_update=True)` left on the trainer. + + Args: + role: Trainer role. + + Returns: + `updated` (True), `train_step`, `update_ready` (the trainer confirmed the + update landed) and `device_wait_s` (how long the trainer still waited + for its device before dispatching the update; ~0 means a stage sent in + between outlasted the step's last micro-batch, None when it cannot + tell). + """ + worker = self._trainer_workers.get(role) + if worker is None: + raise ValueError(f"No trainer worker registered for role {role}") + async with self._actor_busy_guard(role == datatypes.Role.ACTOR): + resp = await self._invoke_worker(worker, "apply_deferred_update") + if role == datatypes.Role.ACTOR: + self._actor_weights_dirty = True + metadata = getattr(resp, "metadata", None) or {} + device_wait_s = metadata.get("device_wait_s") + return { + "updated": bool(metadata.get("updated")), + "train_step": metadata.get("train_step"), + # Set only by a worker that waited for the update to land on device. + "update_ready": bool(metadata.get("update_ready")), + "device_wait_s": ( + None if device_wait_s is None else float(device_wait_s) + ), + } + + async def stage_train_step( self, payload: datatypes.RLTrainerPayload, role: datatypes.Role = datatypes.Role.ACTOR, + **kwargs: Any, + ) -> StagedTrainStep: + """Sends a micro-batch to the trainer ahead of the train_step that runs it. + + The trainer only holds it: no model, optimizer or gradient state changes, + so the trainer's step order is that of the train_step calls. Pass the + returned id as `train_step(None, staged_request_id=...)`; that call then + carries no data. + + Args: + payload: The micro-batch. + role: Trainer role to stage on. + **kwargs: Forwarded to the trainer's `fwd_bwd` when it runs. + + Returns: + The staged micro-batch's id, how long the trainer then still waited for + its device, and whether it confirmed the update dispatched before the + stage has landed. + """ + worker = self._trainer_workers.get(role) + if worker is None: + raise ValueError(f"No trainer worker registered for role {role}") + role_name = role.value if isinstance(role, datatypes.Role) else str(role) + request = datatypes.TrainRequest( + request_id=f"train_req_{uuid.uuid4().hex[:8]}", + payload=payload, + metadata=dict(getattr(payload, "metadata", {}) or {}), + ) + logging.info( + "Staging train_step input on %s worker (request_id=%s)...", + role_name, + request.request_id, + ) + async with self._actor_busy_guard(role == datatypes.Role.ACTOR): + resp = await self._invoke_worker( + worker, "stage_fwd_bwd", request=request, **kwargs + ) + metadata = getattr(resp, "metadata", None) or {} + device_wait_s = metadata.get("device_wait_s") + return StagedTrainStep( + request_id=request.request_id, + device_wait_s=None if device_wait_s is None else float(device_wait_s), + # Set only by a worker that waited for the update to land on device + # (TrainerWorker strips any caller-supplied copy of the key). + update_ready=bool(metadata.get("update_ready")), + ) + + async def train_step( + self, + payload: datatypes.RLTrainerPayload | None, + role: datatypes.Role = datatypes.Role.ACTOR, accumulate_gradients: bool = False, apply_optimizer: bool = True, skip_jit: bool = False, + staged_request_id: str | None = None, + defer_update: bool = False, **kwargs: Any, ) -> Any: - """Executes atomic gradient accumulation / update on TrainerWorker.""" + """Executes atomic gradient accumulation / update on TrainerWorker. + + Args: + payload: The micro-batch, or None with `staged_request_id`. + role: Trainer role. + accumulate_gradients: Reported back in the result. + apply_optimizer: Apply the accumulated gradients after this micro-batch. + skip_jit: Forwarded to the trainer. + staged_request_id: Runs the micro-batch `stage_train_step` sent ahead. + defer_update: With `apply_optimizer`, the micro-batch is still the + step's last, but the update is left to `apply_deferred_update`, for a + caller that stages the next step's first micro-batch in between. Only + for a trainer whose `trainer_staging_support()` has + `stage_before_update`. + **kwargs: Forwarded to the trainer's `fwd_bwd`. + """ worker = self._trainer_workers.get(role) if worker is None: raise ValueError(f"No trainer worker registered for role {role}") + if staged_request_id is not None and payload is not None: + raise ValueError( + "train_step takes a payload or a staged_request_id, not both." + ) + if defer_update: + if not apply_optimizer: + raise ValueError( + "train_step(defer_update=True) needs apply_optimizer=True: only a" + " step's last micro-batch has an update to defer." + ) + # Passed only when set, so the default call is unchanged. + kwargs = {**kwargs, "defer_update": True} role_name = role.value if isinstance(role, datatypes.Role) else str(role) logging.info( "Executing train_step on %s worker (accumulate_gradients=%s," @@ -865,18 +1032,27 @@ async def train_step( accumulate_gradients, apply_optimizer, ) - metadata = dict(getattr(payload, "metadata", {}) or {}) - request = datatypes.TrainRequest( - request_id=f"train_req_{uuid.uuid4().hex[:8]}", - payload=payload, - metadata=metadata, - ) is_actor = role == datatypes.Role.ACTOR + if staged_request_id is not None: + # The micro-batch is already on the trainer (`stage_train_step`); this + # call names it and carries no data. `fwd_bwd_staged` runs it through + # the worker's `fwd_bwd`, so the response below is handled the same way. + method = "fwd_bwd_staged" + call_kwargs: dict[str, Any] = {"staged_request_id": staged_request_id} + else: + metadata = dict(getattr(payload, "metadata", {}) or {}) + request = datatypes.TrainRequest( + request_id=f"train_req_{uuid.uuid4().hex[:8]}", + payload=payload, + metadata=metadata, + ) + method = "fwd_bwd" + call_kwargs = {"request": request} async with self._actor_busy_guard(is_actor): fwd_bwd_result = await self._invoke_worker( worker, - "fwd_bwd", - request=request, + method, + **call_kwargs, skip_jit=skip_jit, # Tells the trainer which micro-batch is the step's last (MaxText's # retry_dropless_final_micro_batch runs it with the dropless program). @@ -891,6 +1067,24 @@ async def train_step( self._actor_weights_dirty = True if not apply_optimizer: return fwd_bwd_result + if ( + defer_update + and isinstance(fwd_bwd_result, datatypes.Response) + and isinstance(fwd_bwd_result.metadata, Mapping) + and bool(fwd_bwd_result.metadata.get("update_deferred")) + ): + # Set only by a worker that deferred it (TrainerWorker strips any + # caller-supplied copy of the key); the caller runs + # `apply_deferred_update` next. A worker that ran the update after all + # is handled below, as a fused call. + return { + "fwd_bwd": fwd_bwd_result, + "updated": False, + "update_deferred": True, + "train_step": None, + "accumulated": accumulate_gradients, + "update_ready": False, + } # A worker that ran the optimizer inside the fwd_bwd RPC reports # `updated=True` (TrainerWorker strips any caller-supplied copy of the # key before echoing request metadata, so only the worker can set it). A diff --git a/tunix/experimental/orchestrator/rl_program.py b/tunix/experimental/orchestrator/rl_program.py index 48517ea4ce..47abc246cb 100644 --- a/tunix/experimental/orchestrator/rl_program.py +++ b/tunix/experimental/orchestrator/rl_program.py @@ -62,6 +62,22 @@ def _release_host_memory() -> None: pass +def _raised_by_worker_method(exc: BaseException) -> bool: + """Whether a worker call failed in the worker's method, not on the way. + + An actor handle re-raises the method's exception as a RuntimeError whose + message starts "RemoteExecutionError" (`ExecutionResponse.unwrap`); a + transport failure (a gRPC status, a reset connection, a timeout) is raised + as itself. + + Args: + exc: What the call raised. + """ + return isinstance(exc, RuntimeError) and str(exc).startswith( + "RemoteExecutionError" + ) + + def _next_or_exhausted(iterator: Any) -> Any: """`next()` that returns `_EXHAUSTED` instead of raising StopIteration. @@ -332,6 +348,45 @@ class RLStepResult: train_result: Any = None +@dataclasses.dataclass(eq=False) +class _StepState: + """One step's intake and train-side counters, for `_train_loop`. + + The groups taken from `scored_q`, what they were packed into, and the + per-step counters. With `pipeline_step_boundary`, a step's state can be + filled while the previous step still trains (`_prepack`); the step adopts + it when it starts. + """ + + uncommitted_groups: list[Any] = dataclasses.field(default_factory=list) + all_step_items: list[Any] = dataclasses.field(default_factory=list) + step_rewards: list[float] = dataclasses.field(default_factory=list) + step_advantages: list[float] = dataclasses.field(default_factory=list) + num_rollouts: int = 0 + groups_consumed: int = 0 + current_batch_idx: int | None = None + # What the queue last returned; `[]` once it has no further group. + scored_items: list[Any] = dataclasses.field(default_factory=list) + packing_time_sec: float = 0.0 + exposed_generation_time: float = 0.0 + # Packed but not yet sent to the trainer, in packing order. + ready: list[batch_assembly.AssembledBatch] = dataclasses.field( + default_factory=list + ) + num_microbatches: int = 0 + padding_stats: list[batch_assembly.PaddingStats] = dataclasses.field( + default_factory=list + ) + sampler_agreement: dict[str, tuple[Any, list[float]]] = dataclasses.field( + default_factory=dict + ) + step_begun: bool = False + start_time: float | None = None + # The first micro-batch, already on the trainer host (`stage_train_step`): + # the micro-batch without its payload, and the id that runs it. + staged: tuple[batch_assembly.AssembledBatch, str] | None = None + + class RLProgram(abc.ABC): """Base class for multi-stage DAG workflows.""" @@ -409,6 +464,9 @@ def __init__( trainer_warmup_payload: datatypes.RLTrainerPayload | None = None, num_generations_to_dispatch: int | None = None, prefer_valid_rollouts: bool = True, + pipeline_step_boundary: bool = False, + stage_next_microbatch: bool = False, + stage_before_update: bool = True, ): super().__init__() self.engine: rl_engine_interface.AbstractRLEngine | None = None @@ -421,6 +479,12 @@ def __init__( and initial_dispatch_first_stagger_s < 0 ): raise ValueError("initial_dispatch_first_stagger_s must be non-negative.") + if stage_next_microbatch and not pipeline_step_boundary: + raise ValueError( + "stage_next_microbatch requires pipeline_step_boundary: the" + " micro-batch it stages is the one the pipelined boundary takes" + " ahead." + ) if checkpoint_optimizer_interval_steps < 0: raise ValueError( "checkpoint_optimizer_interval_steps must be non-negative." @@ -541,6 +605,57 @@ def __init__( # gradient-accumulation order is unchanged. self.pipeline_train_microbatches = pipeline_train_microbatches self._pending_train: asyncio.Task[tuple[Any, float]] | None = None + # Pipelines the step boundary (see `_train_loop`): step k+1's groups are + # taken and packed while step k's last micro-batches train, step k+1's + # first micro-batch goes to the trainer as soon as step k's weight sync + # has staged, and step k's metrics, logging and host-memory release run + # once that call has returned, while the device computes it. The trainer + # gets the same calls, with the same inputs, in the same order as without + # it. What moves with that + # bookkeeping: step k's `Train step` line, StepMetrics and `on_step_end` + # follow a `[pipeline] SENT_AHEAD step=k` line (step k trained, step k+1 + # started), so a wedge in step k+1's first call stops the log there; + # `on_step_begin(k+1)` fires after step k+1's first train call returned; + # and under ARRIVAL order the dispatch window, keyed on the step count, + # opens that one call later (PROMPT_BATCH opens it at the commit). + self.pipeline_step_boundary = pipeline_step_boundary + # With `pipeline_step_boundary`: also send that micro-batch's data to the + # trainer host while the device still computes step k's last micro-batch, + # so the train call that later runs it carries no data. Needs an engine + # with `stage_train_step` and a trainer that supports it + # (`trainer_staging_support`); otherwise the boundary runs unstaged. + self.stage_next_microbatch = stage_next_microbatch + # Where that stage goes. With `stage_before_update` and a trainer that can + # defer a step's update (`stage_before_update` in its staging support): + # between step k's last train call, sent with its update deferred, and + # that update (`_stage_next_before_update`). The last call then returns + # once the trainer has dispatched the last micro-batch, which under + # MaxText's `retry_when_tokens_dropped` is when the micro-batch before it + # finished, so the transfer runs while the device computes the last one. + # Otherwise right after the last call, which runs the update and returns + # at its dispatch (`_maybe_stage_next_microbatch`); under + # `retry_when_tokens_dropped` that dispatch waits until the last + # micro-batch has finished, so the transfer finds the device idle and the + # latch below switches staging off at the first boundary. Only matters + # with `stage_next_microbatch`. + self.stage_before_update = stage_before_update + # Takes groups for a step ahead of it (`_prepack`), filling a `_StepState` + # the step adopts when it starts. + self._intake_task: asyncio.Task[None] | None = None + # Set once a stage did not hide behind the trainer's device (its device + # had gone idle before the transfer finished): the trainer is then fed + # slower than it computes, and staging only delays each step's sync. Also + # set when the trainer cannot stage at all. + self._staging_latched_off = False + # What the trainer said it supports of staging; asked once, lazily. + self._staging_supported: bool | None = None + self._staging_support: dict[str, bool] | None = None + # Whether it can also defer a step's update past its last micro-batch, so + # the stage goes in between (`_trainer_can_stage_before_update`). + self._stage_before_update_supported: bool | None = None + # The optimizer step the trainer last reported (`train_step` of an update + # call), from which `_stage_planned` predicts the next one. + self._last_optimizer_step: int | None = None # A cancelled asyncio.to_thread leaves its worker running, so the error # path's reset() must wait for any in-flight feed/flush to finish. self._assembler_lock = threading.Lock() @@ -1843,15 +1958,35 @@ def _locked() -> None: await asyncio.to_thread(_locked) async def _timed_train_step( - self, batch: datatypes.RLTrainerPayload, *, apply_optimizer: bool + self, + batch: datatypes.RLTrainerPayload | None, + *, + apply_optimizer: bool, + staged_request_id: str | None = None, + wait_for_update: bool = True, + defer_update: bool = False, ) -> tuple[Any, float]: assert self.engine is not None start = time.monotonic() + # Passed only when set, so engines without staging see the call unchanged. + kwargs = {} + if staged_request_id is not None: + kwargs["staged_request_id"] = staged_request_id + if not wait_for_update: + # Only with staging, i.e. to a trainer that supports it + # (`_trainer_can_stage`): its worker's fused update then returns at + # the update's dispatch (`TrainerWorker.fwd_bwd`). + kwargs["wait_for_update"] = False + if defer_update: + # Only to a trainer that can (`_trainer_can_stage_before_update`): the + # update is left to `apply_deferred_update`. + kwargs["defer_update"] = True result = await self.engine.train_step( batch, role=datatypes.Role.ACTOR, accumulate_gradients=True, apply_optimizer=apply_optimizer, + **kwargs, ) return result, time.monotonic() - start @@ -1944,40 +2079,86 @@ async def train_stage(self) -> None: finally: if self._pending_sync is not None and not self._pending_sync.done(): self._pending_sync.cancel() + if self._intake_task is not None and not self._intake_task.done(): + self._intake_task.cancel() async def _train_loop(self) -> None: + """Trains one optimizer step per full batch of scored groups. + + One loop for both step boundaries. Without `pipeline_step_boundary`, step + k+1 starts once step k's bookkeeping is done. With it, the trainer gets + the same calls, with the same inputs, in the same order -- every + micro-batch of step k (the last one also runs step k's update), then the + metrics fetch, checkpoint save and weight-sync staging, then step k+1's + first micro-batch -- and the assembler sees the same feeds, so the + micro-batches, their order, and the weights each is computed on are + unchanged. What moves is the orchestrator's own work between step k's + update and step k+1's first micro-batch: + + * Packing. Once step k's last group is fed, step k+1's groups are taken + and packed (`_prepack`) while step k's last micro-batches train. + * Bookkeeping. Step k+1's first micro-batch is sent right after step k's + sync stages, its batch commits and, when deferred, its metrics are + fetched (`_send_first_microbatch_ahead`); step k's step metrics, + logging, `on_step_end` and host-memory release run once that call is on + the trainer. + * Transfer (`stage_next_microbatch`). That micro-batch's data goes to the + trainer host while the device still computes step k's last micro-batch + (`_stage_planned`): between step k's last train call, sent with its + update deferred, and that update (`_stage_next_before_update`), when + the trainer can defer it; else right after the last call, which then + returns at the update's dispatch (`_maybe_stage_next_microbatch`). The + stage call, and with a deferred update the update call, are the only + calls added on the trainer. + """ assert self.engine is not None + ordered_q = isinstance( + self.scored_q, trajectory_queue_manager.BatchOrderedQueueManager + ) + # The next step's state while it is taken ahead (`_prepack`); only ever + # set with `pipeline_step_boundary`. + next_state: _StepState | None = None while self.max_steps is None or self._step < self.max_steps: current_step = self._step - step_start_time = time.monotonic() + is_last_step = ( + self.max_steps is not None and current_step + 1 >= self.max_steps + ) consumed_policy_version = self.policy_version + st = next_state if next_state is not None else _StepState() + next_state = None + if self._intake_task is not None: + # Still taking this step's groups ahead; from here the wait is exposed. + intake, self._intake_task = self._intake_task, None + t_wait = time.monotonic() + await intake + st.exposed_generation_time += time.monotonic() - t_wait + if st.start_time is None: + st.start_time = time.monotonic() + step_start_time = st.start_time + if st.groups_consumed == 0: + # Nothing was taken ahead (the next batch was a hole, or the queue + # closed): start exactly as without taking ahead. + st.current_batch_idx = None + st.scored_items = [] + elif not st.step_begun: + st.step_begun = True + if self.on_step_begin: + self.on_step_begin(current_step) - uncommitted_groups = [] step_result = None trainer_metrics = None - step_sampler_agreement: dict[str, tuple[Any, list[float]]] = {} - step_rewards = [] - step_advantages = [] - num_microbatches = 0 - step_padding_stats: list[batch_assembly.PaddingStats] = [] - step_packing_time_sec = 0.0 - num_rollouts = 0 - all_step_items = [] - scored_items = [] - groups_consumed = 0 checkpoint_saved = False final_minibatch_completed = False pending_metrics_fetch = False + queue_dry = False + prepack_started = False policy_training_time = 0.0 - exposed_generation_time = 0.0 weight_sync_time = 0.0 # The end-of-step get_metrics() deferred past the weight sync and commit, # kept out of policy_training_time because it no longer runs in training. metrics_fetch_time = 0.0 - current_batch_idx: int | None = None - async def _await_pending_train() -> None: nonlocal step_result, policy_training_time task, self._pending_train = self._pending_train, None @@ -1987,196 +2168,47 @@ async def _await_pending_train() -> None: async def _maybe_save_checkpoint() -> None: nonlocal checkpoint_saved - ckpt_ts_ms = time.time_ns() // 1_000_000 - self.last_step_timestamp_ms = ckpt_ts_ms - optimizer_step = self.step + 1 - if ( - isinstance(step_result, dict) - and step_result.get("train_step") is not None - ): - optimizer_step = int(step_result["train_step"]) - if ( - self.val_start_step is not None - and optimizer_step < self.val_start_step - ): - checkpoint_saved = True - return - if isinstance( - self.scored_q, - trajectory_queue_manager.BatchOrderedQueueManager, - ): - next_batch_idx = ( - self.scored_q.next_batch_after(current_batch_idx) - if current_batch_idx is not None - else self.scored_q.next_batch_idx - ) - else: - next_batch_idx = self.step + 1 - save_kwargs = {} - interval = self.checkpoint_optimizer_interval_steps - if interval == 0 or ( - optimizer_step % interval != 0 - and (self.max_steps is None or self.step + 1 < self.max_steps) - ): - # Sent only when skipping, so an interval of 1 leaves the request, - # and trainers that predate the option, untouched. - save_kwargs["save_optimizer_state"] = False - save_resp = await self.engine.save_checkpoint( - role=datatypes.Role.ACTOR, - metadata={ - "step": optimizer_step, - "global_step": self.step + 1, - "next_batch_idx": next_batch_idx, - "policy_version": self.policy_version + 1, - "num_rollouts": num_rollouts, - "num_microbatches": num_microbatches, - }, - **save_kwargs, - ) + await self._save_step_checkpoint(st, step_result) checkpoint_saved = True - if self.on_checkpoint_saved is not None: - # `TrainerWorker.save_checkpoint` always returns the saved path in - # `Response.metadata["checkpoint_path"]`. The callback runs in a - # worker thread (it may do blocking file/GCS I/O) and is awaited, so - # calls never overlap and its exceptions still propagate. - ckpt_info: dict[str, Any] = { - "step": optimizer_step, - "timestamp_ms": ckpt_ts_ms, - "checkpoint_path": save_resp.metadata["checkpoint_path"], - } - await asyncio.to_thread(self.on_checkpoint_saved, ckpt_info) - while groups_consumed < self.full_batch_size: - _t_gen = time.monotonic() - if isinstance( - self.scored_q, trajectory_queue_manager.BatchOrderedQueueManager + while True: + if ( + self.pipeline_step_boundary + and not prepack_started + and not is_last_step + and st.groups_consumed > 0 + and (queue_dry or st.groups_consumed >= self.full_batch_size) ): - ordered = await self.scored_q.get_ordered_group( - batch_idx=current_batch_idx - ) - if ordered is None: - scored_items = [] - else: - current_batch_idx, scored_items = ordered - else: - scored_items = await self.scored_q.get_group_batch(num_groups=1) - # TODO(sanbao): Track accurate generation time by matching rollout - # requests/responses or reading from the rollout worker. Currently this - # measures time spent waiting on the scored queue (which includes - # rollout + reward calculation), as MLPerf does not strictly validate - # per-component timing accuracy. - exposed_generation_time += time.monotonic() - _t_gen - if not scored_items: - packing_start_time = time.perf_counter() - assembled_batches = list(await self._assemble(self.assembler.flush)) - step_packing_time_sec += time.perf_counter() - packing_start_time - else: - if groups_consumed == 0 and self.on_step_begin: - self.on_step_begin(current_step) - - groups_consumed += 1 - uncommitted_groups.append(scored_items) - all_step_items.extend(scored_items) - num_rollouts += len(scored_items) - for item in scored_items: - step_rewards.append(_extract_reward(item)) - payload = getattr(item, "payload", None) - if ( - item.is_valid - and payload is not None - and payload.advantages is not None - and np.asarray(payload.advantages).size > 0 - ): - step_advantages.append(float(np.mean(payload.advantages))) - - payloads = [] - for item in scored_items: - payload = getattr(item, "payload", None) - if isinstance(payload, datatypes.RLTrainerPayload): - payload = dataclasses.replace( - payload, - metadata={ - **payload.metadata, - "traj_id": item.traj_id, - }, - ) - payloads.append(payload) - packing_start_time = time.perf_counter() - assembled_batches = list( - await self._assemble( - self.assembler.feed, - payloads, # pyrefly: ignore[bad-argument-type] + # Every group of this step is fed. Take the next step's ahead while + # this step's last micro-batches train. Under PROMPT_BATCH order + # that is the next batch by index: `next_batch_idx` does not move + # before this batch commits, and a hole there just leaves nothing + # taken. + prepack_started = True + next_state = _StepState( + current_batch_idx=( + st.current_batch_idx + 1 + if ordered_q and st.current_batch_idx is not None + else None ) ) - step_packing_time_sec += time.perf_counter() - packing_start_time - # `self.assembler` has now copied or packed `routed_experts`; release - # the per-trajectory references on `payloads` and `scored_items` - # (which remain in `all_step_items` until step-end metric logging) so - # unbatched routing arrays do not coexist with packed microbatches and - # weight sync. Lightweight 1D mask fields on `item.payload` are kept - # for `_collect_and_log_step_metrics`. - payloads.clear() - for item in scored_items: - item.routed_experts = None - if item.metadata is not None: - item.metadata.pop("routed_experts", None) - if isinstance(item.traj, dict): - item.traj.pop("routed_experts", None) - item_payload = getattr(item, "payload", None) - if ( - isinstance(item_payload, datatypes.RLTrainerPayload) - and item_payload.routed_experts is not None - ): - item.payload = dataclasses.replace( - item_payload, routed_experts=None - ) + self._intake_task = asyncio.create_task(self._prepack(next_state)) - while assembled_batches: - mb = assembled_batches.pop(0) + # Send what is packed, in order: a feed's micro-batches go out before + # the next group is taken. + while st.ready: + mb = st.ready.pop(0) is_final_batch = mb.is_final_batch - batch = mb.payload - if getattr(self.algo, "requires_reference_kl", False): - if not isinstance(batch, datatypes.RLTrainerPayload): - raise TypeError( - "Reference KL requires an assembler that returns " - "datatypes.RLTrainerPayload microbatches; got " - f"{type(batch).__name__}." - ) - # The reference model is static and on its own worker, so this - # overlaps the in-flight actor train_step. - ref_logps = await self.engine.per_token_logps( - datatypes.Role.REFERENCE, items=batch - ) - batch = batch_assembly.with_ref_per_token_logps(batch, ref_logps) - algo_config = getattr(self.algo, "algo_config", None) - if ( - isinstance(batch, datatypes.RLTrainerPayload) - and batch.old_per_token_logps is not None - and algo_config is not None - and algo_config.use_rollout_logps - ): + # The reference model is static and on its own worker, so this + # overlaps the in-flight actor train_step. + batch = await self._with_reference_logps(mb.payload) + rescored = self._needs_sampler_agreement(batch) + if rescored: await _await_pending_train() batch = await self._apply_sampler_trainer_agreement( - batch, step_sampler_agreement - ) - - num_microbatches += 1 - if mb.padding_stats is not None: - step_padding_stats.append(mb.padding_stats) - logging.info( - "Packed %d trajectories into microbatch: %s (padding_ratio=%.3f," - " row_imbalance=%.3f)", - len(mb.trajectory_ids), - logging_utils.summarize_list(list(mb.trajectory_ids)), - mb.padding_stats.padding_ratio, - mb.padding_stats.row_imbalance, - ) - else: - logging.info( - "Packed %d trajectories into microbatch: %s", - len(mb.trajectory_ids), - logging_utils.summarize_list(list(mb.trajectory_ids)), + batch, st.sampler_agreement ) + self._record_microbatch(st, mb) mb = None if not self._first_train_payload_logged: self._first_train_payload_logged = True @@ -2185,8 +2217,33 @@ async def _maybe_save_checkpoint() -> None: batch_assembly.describe_payload(batch), ) await _await_pending_train() + # Save only at a resumable full-batch boundary. An optimizer step + # can occur earlier when a full batch contains multiple mini + # batches, but the dataset resume cursor advances in full batches. + # (Known before the call: it changes neither count.) + full_batch_complete = is_final_batch and ( + st.groups_consumed >= self.full_batch_size or not st.scored_items + ) + # Whether the next step's first micro-batch is staged right after + # this call. The call runs the update too, and then returns at the + # update's dispatch rather than once it landed, so the transfer + # hides behind the device's remaining work; the stage confirms the + # update instead. Off the flag, the call is exactly as without it. + stage_after = ( + full_batch_complete + and self.stage_next_microbatch + and await self._stage_planned(next_state, rescored=rescored) + ) + # Or, when the trainer can, the update is left out of this call and + # the stage goes before it (`_stage_next_before_update`). + defer_update = ( + stage_after and await self._trainer_can_stage_before_update() + ) train_step = self._timed_train_step( - batch, apply_optimizer=is_final_batch + batch, + apply_optimizer=is_final_batch, + wait_for_update=not stage_after, + defer_update=defer_update, ) batch = None if self.pipeline_train_microbatches and not is_final_batch: @@ -2195,21 +2252,57 @@ async def _maybe_save_checkpoint() -> None: step_result, elapsed = await train_step policy_training_time += elapsed if is_final_batch: - # Save only at a resumable full-batch boundary. An optimizer step - # can occur earlier when a full batch contains multiple mini - # batches, but the dataset resume cursor advances in full batches. - full_batch_complete = ( - groups_consumed >= self.full_batch_size or not scored_items - ) - # The checkpoint is stamped as the weight-update time and saved - # next, so the update must have landed on device first. Unless the - # trainer has said it did (`update_ready`), fetching the metrics is - # what waits for it, so the fetch stays ahead of the save, counted - # as training time, as with a separate update RPC. - pending_metrics_fetch = full_batch_complete and bool( + # While an update is deferred the worker refuses get_metrics, a + # checkpoint and a weight sync, so `_stage_next_before_update`, + # which always applies the update, runs before any of them. A + # result without `update_deferred` means the trainer ran the + # update inside the call after all (the engine reports that as a + # fused call): `stage_after` is then still set, and the stage goes + # after the update below, as with `--no-stage_before_update`. + if ( + defer_update + and isinstance(step_result, dict) + and step_result.get("update_deferred") + ): + assert next_state is not None + # The trainer has this step's last micro-batch and no update + # yet: stage the next step's first micro-batch under it, then + # run the update. Training time, as the fused call's wait for + # the device would have been. + deferred_start = time.monotonic() + step_result = await self._stage_next_before_update( + next_state, step_result + ) + policy_training_time += time.monotonic() - deferred_start + # This boundary's stage was made, or refused, before the + # update. Skip the hook below: staging after the update's + # dispatch is the order this path replaces, and under dropless + # replay it would find the device idle and latch staging off. + stage_after = False + if ( + isinstance(step_result, dict) + and step_result.get("train_step") is not None + ): + self._last_optimizer_step = int(step_result["train_step"]) + update_ready = bool( isinstance(step_result, dict) and step_result.get("update_ready") ) + if stage_after: + assert next_state is not None + # The update is dispatched and the device has this step's last + # micro-batch and the update ahead of it. + if await self._maybe_stage_next_microbatch( + next_state, step_result + ): + update_ready = True + # The checkpoint is stamped as the weight-update time and saved + # next, so the update must have landed on device first. Unless the + # trainer has said it did (`update_ready`, from this call or from + # the stage right after it), fetching the metrics is what waits + # for it, so the fetch stays ahead of the save, counted as + # training time, as with a separate update RPC. + pending_metrics_fetch = full_batch_complete and update_ready if not pending_metrics_fetch: _t_metrics = time.monotonic() trainer_metrics = await self.engine.get_metrics( @@ -2227,13 +2320,18 @@ async def _maybe_save_checkpoint() -> None: if full_batch_complete: await _maybe_save_checkpoint() - if not scored_items: + if queue_dry: if not checkpoint_saved and final_minibatch_completed: await _maybe_save_checkpoint() break + if st.groups_consumed >= self.full_batch_size: + break + if not await self._take_group(st, current_step=current_step): + await self._flush_step(st) + queue_dry = True await _await_pending_train() - if not all_step_items: + if not st.all_step_items: logging.info( "Dataset exhausted at step %d before max_steps.", current_step ) @@ -2258,30 +2356,32 @@ async def _maybe_save_checkpoint() -> None: ) # Before `commit()`, which will eventually take ownership of the groups. - generation_metrics = _generation_metrics(uncommitted_groups) + generation_metrics = _generation_metrics(st.uncommitted_groups) logging.info( "[pipeline] COMMIT step=%d batch_idx=%s consumed_version=%d" " rollout_versions=%s groups=%s", current_step, - current_batch_idx, + st.current_batch_idx, consumed_policy_version, - [getattr(g[0], "policy_version", 0) for g in uncommitted_groups if g], + [ + getattr(g[0], "policy_version", 0) + for g in st.uncommitted_groups + if g + ], [ ( getattr(g[0], "prompt_id", ""), (getattr(g[0], "metadata", None) or {}).get("batch_idx"), ) - for g in uncommitted_groups + for g in st.uncommitted_groups if g ], ) - if isinstance( - self.scored_q, trajectory_queue_manager.BatchOrderedQueueManager - ): - assert current_batch_idx is not None - self.scored_q.commit_batch(current_batch_idx) + if ordered_q: + assert st.current_batch_idx is not None + self.scored_q.commit_batch(st.current_batch_idx) else: - self.scored_q.commit(current_step, groups=uncommitted_groups) + self.scored_q.commit(current_step, groups=st.uncommitted_groups) # The commit may have opened the dispatch window. Yield so the # dispatcher refills the just-synced rollout workers now, rather than # after this step's bookkeeping and the next step's first packing. @@ -2298,42 +2398,63 @@ async def _maybe_save_checkpoint() -> None: filtered_groups = await self._drain_filtered_groups() + if next_state is not None: + # Step k+1's first micro-batch goes to the trainer now. This step's + # remaining bookkeeping waits until that call is on the trainer, so + # its transfer is not slowed by it and the device computes while it + # runs: step k's `Train step` line, StepMetrics and `on_step_end` + # therefore follow the `[pipeline] SENT_AHEAD` line by one train + # call's return (one transfer unstaged, ~a round trip staged). + sent_ahead = await self._send_first_microbatch_ahead(next_state) + if sent_ahead is not None: + logging.info( + "[pipeline] SENT_AHEAD step=%d first microbatch of step %d is on" + " its way to the trainer; step %d's metrics follow once it is" + " there.", + current_step, + current_step + 1, + current_step, + ) + # Not collected here: its result is step k+1's, and a failure + # surfaces when step k+1 collects it. + await asyncio.wait({sent_ahead}) + metrics_summary = self._collect_and_log_step_metrics( - all_step_items=all_step_items, - step_rewards=step_rewards, - step_advantages=step_advantages, + all_step_items=st.all_step_items, + step_rewards=st.step_rewards, + step_advantages=st.step_advantages, generation_metrics=generation_metrics, step_result=step_result, trainer_metrics=trainer_metrics, - num_rollouts=num_rollouts, - num_microbatches=num_microbatches, - padding_stats=step_padding_stats, - packing_time_sec=step_packing_time_sec, + num_rollouts=st.num_rollouts, + num_microbatches=st.num_microbatches, + padding_stats=st.padding_stats, + packing_time_sec=st.packing_time_sec, step_time_sec=step_time_sec, consumed_policy_version=consumed_policy_version, log_step=current_step, - sampler_agreement=step_sampler_agreement, + sampler_agreement=st.sampler_agreement, policy_training_time=policy_training_time, - exposed_generation_time=exposed_generation_time, + exposed_generation_time=st.exposed_generation_time, weight_sync_time=weight_sync_time, metrics_fetch_time=metrics_fetch_time, filtered_groups=filtered_groups, ) self._log_consumed_trajectories( - all_step_items, + st.all_step_items, log_step=current_step, consumed_policy_version=consumed_policy_version, ) filtered_groups.clear() - all_step_items.clear() - uncommitted_groups.clear() - scored_items.clear() + st.all_step_items.clear() + st.uncommitted_groups.clear() + st.scored_items.clear() self.last_step_result = RLStepResult( step=current_step, policy_version=self.policy_version, - num_rollouts=num_rollouts, - num_microbatches=num_microbatches, + num_rollouts=st.num_rollouts, + num_microbatches=st.num_microbatches, reward_mean=metrics_summary["reward_mean"], reward_std=metrics_summary["reward_std"], advantage_mean=metrics_summary["advantage_mean"], @@ -2366,6 +2487,685 @@ async def _maybe_save_checkpoint() -> None: self._release_window() _release_host_memory() + # The pieces of `_train_loop` that work on a `_StepState`, so that a step's + # state can be filled ahead of it (`_prepack`) and its first micro-batch + # sent (`_send_first_microbatch_ahead`) or staged + # (`_maybe_stage_next_microbatch`) early under `pipeline_step_boundary`. + + async def _take_group( + self, st: _StepState, *, current_step: int | None + ) -> bool: + """Takes the step's next group from `scored_q` and feeds it to the assembler. + + `current_step` is None when taking ahead for a step that has not started + (`_prepack`): the queue wait and `on_step_begin` are then left to the + step, which does both on adopting the state. + + Args: + st: State of the step the group is for. + current_step: The step taking the group, or None when taking ahead. + + Returns: + False, leaving `st.scored_items == []`, when the queue has no further + group for the step. + """ + t_gen = time.monotonic() + if isinstance( + self.scored_q, trajectory_queue_manager.BatchOrderedQueueManager + ): + ordered = await self.scored_q.get_ordered_group( + batch_idx=st.current_batch_idx + ) + if ordered is None: + scored_items = [] + else: + st.current_batch_idx, scored_items = ordered + else: + scored_items = await self.scored_q.get_group_batch(num_groups=1) + # TODO(sanbao): Track accurate generation time by matching rollout + # requests/responses or reading from the rollout worker. Currently this + # measures time spent waiting on the scored queue (which includes + # rollout + reward calculation), as MLPerf does not strictly validate + # per-component timing accuracy. + if current_step is not None: + st.exposed_generation_time += time.monotonic() - t_gen + st.scored_items = scored_items + if not scored_items: + return False + if current_step is not None and not st.step_begun: + st.step_begun = True + if self.on_step_begin: + self.on_step_begin(current_step) + + st.groups_consumed += 1 + st.uncommitted_groups.append(scored_items) + st.all_step_items.extend(scored_items) + st.num_rollouts += len(scored_items) + for item in scored_items: + st.step_rewards.append(_extract_reward(item)) + payload = getattr(item, "payload", None) + if ( + item.is_valid + and payload is not None + and payload.advantages is not None + and np.asarray(payload.advantages).size > 0 + ): + st.step_advantages.append(float(np.mean(payload.advantages))) + + payloads = [] + for item in scored_items: + payload = getattr(item, "payload", None) + if isinstance(payload, datatypes.RLTrainerPayload): + payload = dataclasses.replace( + payload, + metadata={ + **payload.metadata, + "traj_id": item.traj_id, + }, + ) + payloads.append(payload) + packing_start_time = time.perf_counter() + assembled_batches = list( + await self._assemble( + self.assembler.feed, + payloads, # pyrefly: ignore[bad-argument-type] + ) + ) + st.packing_time_sec += time.perf_counter() - packing_start_time + # `self.assembler` has now copied or packed `routed_experts`; release + # the per-trajectory references on `payloads` and `scored_items` + # (which remain in `all_step_items` until step-end metric logging) so + # unbatched routing arrays do not coexist with packed microbatches and + # weight sync. Lightweight 1D mask fields on `item.payload` are kept + # for `_collect_and_log_step_metrics`. + payloads.clear() + for item in scored_items: + item.routed_experts = None + if item.metadata is not None: + item.metadata.pop("routed_experts", None) + if isinstance(item.traj, dict): + item.traj.pop("routed_experts", None) + item_payload = getattr(item, "payload", None) + if ( + isinstance(item_payload, datatypes.RLTrainerPayload) + and item_payload.routed_experts is not None + ): + item.payload = dataclasses.replace(item_payload, routed_experts=None) + st.ready.extend(assembled_batches) + return True + + async def _flush_step(self, st: _StepState) -> None: + """Packs whatever the assembler still holds for the step.""" + packing_start_time = time.perf_counter() + st.ready.extend(list(await self._assemble(self.assembler.flush))) + st.packing_time_sec += time.perf_counter() - packing_start_time + + async def _prepack(self, st: _StepState) -> None: + """Takes groups for a step ahead of it, until one micro-batch is packed. + + Also stops once the step has all its groups, or the queue has no further + group for it. The step takes over exactly where this stops, so the + assembler sees the same feeds in the same order as without taking ahead. + + Args: + st: State of the step being taken ahead. + """ + while not st.ready and st.groups_consumed < self.full_batch_size: + if not await self._take_group(st, current_step=None): + return + + def _intake_done(self) -> bool: + """Whether the intake task finished without error (and is not consumed).""" + task = self._intake_task + return ( + task is not None + and task.done() + and not task.cancelled() + and task.exception() is None + ) + + async def _with_reference_logps(self, batch: Any) -> Any: + """Adds the reference policy's log-probs when the algorithm needs them.""" + if not getattr(self.algo, "requires_reference_kl", False): + return batch + if not isinstance(batch, datatypes.RLTrainerPayload): + raise TypeError( + "Reference KL requires an assembler that returns " + "datatypes.RLTrainerPayload microbatches; got " + f"{type(batch).__name__}." + ) + assert self.engine is not None + # The reference model is static and on its own worker. + ref_logps = await self.engine.per_token_logps( + datatypes.Role.REFERENCE, items=batch + ) + return batch_assembly.with_ref_per_token_logps(batch, ref_logps) + + def _needs_sampler_agreement(self, batch: Any) -> bool: + """Whether `batch` is rescored under the trainer's live weights first.""" + algo_config = getattr(self.algo, "algo_config", None) + return bool( + isinstance(batch, datatypes.RLTrainerPayload) + and batch.old_per_token_logps is not None + and algo_config is not None + and algo_config.use_rollout_logps + ) + + def _record_microbatch( + self, st: _StepState, mb: batch_assembly.AssembledBatch + ) -> None: + """Counts a micro-batch about to be trained, and logs what it packs.""" + st.num_microbatches += 1 + if mb.padding_stats is not None: + st.padding_stats.append(mb.padding_stats) + logging.info( + "Packed %d trajectories into microbatch: %s (padding_ratio=%.3f," + " row_imbalance=%.3f)", + len(mb.trajectory_ids), + logging_utils.summarize_list(list(mb.trajectory_ids)), + mb.padding_stats.padding_ratio, + mb.padding_stats.row_imbalance, + ) + else: + logging.info( + "Packed %d trajectories into microbatch: %s", + len(mb.trajectory_ids), + logging_utils.summarize_list(list(mb.trajectory_ids)), + ) + + def _optimizer_step_after(self, step_result: Any) -> int: + """The optimizer step a step's final train call completed.""" + if ( + isinstance(step_result, dict) + and step_result.get("train_step") is not None + ): + return int(step_result["train_step"]) + return self.step + 1 + + def _writes_scored_checkpoint(self, step_result: Any) -> bool: + """Whether this step's checkpoint is one an MLPerf eval may be scored on.""" + return ( + self.val_start_step is not None + and self._optimizer_step_after(step_result) >= self.val_start_step + ) + + async def _save_step_checkpoint( + self, st: _StepState, step_result: Any + ) -> None: + """Stamps the step's checkpoint time and saves it unless it is unscored. + + The stamp is taken here, after the barrier that proves the step's update + landed (MLPerf: a checkpoint's timestamp follows its last weight update): + `get_metrics`, or the trainer's `update_ready` on the step's last train + call or on the stage right after it. + + Args: + st: The step's state. + step_result: The step's final train call result. + """ + assert self.engine is not None + ckpt_ts_ms = time.time_ns() // 1_000_000 + self.last_step_timestamp_ms = ckpt_ts_ms + optimizer_step = self._optimizer_step_after(step_result) + if self.val_start_step is not None and optimizer_step < self.val_start_step: + return + if isinstance( + self.scored_q, + trajectory_queue_manager.BatchOrderedQueueManager, + ): + next_batch_idx = ( + self.scored_q.next_batch_after(st.current_batch_idx) + if st.current_batch_idx is not None + else self.scored_q.next_batch_idx + ) + else: + next_batch_idx = self.step + 1 + save_kwargs = {} + interval = self.checkpoint_optimizer_interval_steps + if interval == 0 or ( + optimizer_step % interval != 0 + and (self.max_steps is None or self.step + 1 < self.max_steps) + ): + # Sent only when skipping, so an interval of 1 leaves the request, + # and trainers that predate the option, untouched. + save_kwargs["save_optimizer_state"] = False + save_resp = await self.engine.save_checkpoint( + role=datatypes.Role.ACTOR, + metadata={ + "step": optimizer_step, + "global_step": self.step + 1, + "next_batch_idx": next_batch_idx, + "policy_version": self.policy_version + 1, + "num_rollouts": st.num_rollouts, + "num_microbatches": st.num_microbatches, + }, + **save_kwargs, + ) + if self.on_checkpoint_saved is not None: + # `TrainerWorker.save_checkpoint` always returns the saved path in + # `Response.metadata["checkpoint_path"]`. The callback runs in a + # worker thread (it may do blocking file/GCS I/O) and is awaited, so + # calls never overlap and its exceptions still propagate. + ckpt_info: dict[str, Any] = { + "step": optimizer_step, + "timestamp_ms": ckpt_ts_ms, + "checkpoint_path": save_resp.metadata["checkpoint_path"], + } + await asyncio.to_thread(self.on_checkpoint_saved, ckpt_info) + + # A trainer that still waited at least this long for its device once a + # staged micro-batch was in place had work left to hide the transfer + # behind. Anything shorter is the readiness check's own round trip (a + # Pathways proxy RPC), i.e. the device had already gone idle. + _STAGING_HIDDEN_WAIT_S = 0.05 + + def _staging_hidden(self, device_wait_s: float | None) -> bool: + """Whether a stage hid behind the trainer's device (see `_maybe_stage_...`). + + Args: + device_wait_s: How long the trainer still waited for the work dispatched + before the stage, once the staged micro-batch was in place; None when + it cannot tell. + """ + return ( + device_wait_s is not None + and device_wait_s >= self._STAGING_HIDDEN_WAIT_S + ) + + async def _trainer_can_stage(self) -> bool: + """Whether the actor's trainer supports staging; asks it once. + + Staging needs a trainer that holds a micro-batch ahead of its fwd_bwd and + reports whether doing so hid behind its device (`device_wait_s`, without + which the latch below cannot tell a feed-bound trainer). An older trainer + image, or a trainer without the report, gets the boundary unstaged. + """ + if self._staging_supported is None: + probe = getattr(self.engine, "trainer_staging_support", None) + if ( + probe is None + or getattr(self.engine, "stage_train_step", None) is None + ): + support = {} + else: + support = await probe(role=datatypes.Role.ACTOR) + self._staging_support = dict(support) + self._staging_supported = bool( + support.get("stage_fwd_bwd") and support.get("reports_device_wait") + ) + if not self._staging_supported: + self._staging_latched_off = True + logging.warning( + "[pipeline] stage_next_microbatch is set but the trainer cannot" + " stage a microbatch and report whether that hid behind its" + " device (%s). Running the pipelined boundary unstaged.", + support, + ) + return self._staging_supported + + async def _trainer_can_stage_before_update(self) -> bool: + """Whether to stage between the last call and the update; asks once. + + Needs `stage_before_update`, an engine with `apply_deferred_update`, and a + trainer that can defer a step's update (`stage_before_update` in the + staging support `_trainer_can_stage` asked it for). Only asked when a + stage is planned, i.e. once the trainer is known to stage at all. + """ + if not self.stage_before_update: + return False + if self._stage_before_update_supported is None: + if not await self._trainer_can_stage(): + return False + support = self._staging_support or {} + self._stage_before_update_supported = bool( + support.get("stage_before_update") + and getattr(self.engine, "apply_deferred_update", None) is not None + ) + if not self._stage_before_update_supported: + logging.warning( + "[pipeline] The trainer cannot defer a step's update past its last" + " microbatch (%s): staging right after the last call instead," + " which hides nothing when the trainer's update waits for that" + " microbatch (MaxText retry_when_tokens_dropped).", + support, + ) + return self._stage_before_update_supported + + async def _stage_planned( + self, st: _StepState | None, *, rescored: bool + ) -> bool: + """Whether to stage the next step's first micro-batch after the last call. + + Decided before this step's last call is sent: it runs the update too; when + a stage follows it must return at the update's dispatch + (`wait_for_update=False`) rather than once the update landed, or the + transfer could not hide behind the device. Planning a stage that then + does not happen costs that step the deferred metrics fetch: the update + is not confirmed, so get_metrics() is its barrier again, ahead of the + checkpoint save. So this rules out now whatever `_maybe_stage_next_...` + would refuse after the call: the flag, the latch, a step that writes a + scored checkpoint (the update this call runs is the one after the last + the trainer reported; before any, `self.step + 1`, exact with one update + per step as the recipes run), rescored micro-batches, an already packed + micro-batch that closes its step, and a trainer that cannot stage. What + is left open is whether the micro-batch is packed by the time the call + returns. + + Args: + st: The next step's state, or None when nothing is taken ahead. + rescored: This step's last micro-batch is rescored under the live + weights (`_needs_sampler_agreement`). The next step's carry the same + rollout log-probs, so its first would be too, and those go out whole. + + Returns: + Whether to stage: send the step's last call with its update deferred + and call `_stage_next_before_update` once it returns, when the trainer + can (`_trainer_can_stage_before_update`); else send it with + `wait_for_update=False` and call `_maybe_stage_next_microbatch` once it + returns. + """ + if ( + not self.stage_next_microbatch + or st is None + or self._staging_latched_off + or st.staged is not None + or rescored + ): + return False + predicted_step = ( + self._last_optimizer_step + 1 + if self._last_optimizer_step is not None + else self.step + 1 + ) + if self._writes_scored_checkpoint({"train_step": predicted_step}): + return False + if self._intake_done(): + # Taking ahead is over: only what it packed could go, and it has to be + # neither the step's last micro-batch nor rescored. + if not st.ready: + return False + first = st.ready[0] + if first.is_final_batch or self._needs_sampler_agreement(first.payload): + return False + return await self._trainer_can_stage() + + async def _maybe_stage_next_microbatch( + self, + st: _StepState, + step_result: Any, + ) -> bool: + """Sends the next step's first micro-batch to the trainer host ahead of time. + + Called right after the current step's last train call returns. That call + ran the update too and, staging being planned (`_stage_planned`), came + back at the update's dispatch, while the trainer's device still computes + the step's last micro-batch and update, so the transfer hides behind that + work. The train call that later runs it (`_send_first_microbatch_ahead`) + names it instead of carrying it, and lands on the trainer at the same + point as without staging. Skipped when the micro-batch is not packed yet, + closes its step, or is rescored under the live weights first + (`use_rollout_logps`): those go out whole. + + Also skipped when this step writes a scored checkpoint: its timestamp is + taken once the update lands (`get_metrics`), and a transfer longer than + the device's remaining work -- several trainer copies fed through one + controller -- would push that timestamp back. + + And switched off for the rest of the run the first time a stage does not + hide. The trainer answers once the micro-batch is in place and its device + has finished the work dispatched before it, and says how long it waited + for the latter. A wait under `_STAGING_HIDDEN_WAIT_S` means the device + had already gone idle: the transfer outlasted the step's remaining device + work and delayed this step's metrics, stamp and weight sync. The trainer + is then fed slower than it computes, so staging only moves a transfer + from after the sync to before it. (A step's train calls returning cannot + tell this: MaxText queues `max_inflight_computations` of them, and a step + starts on a device drained by the weight sync, so its first calls return + at the transfer's pace even when the device is the bottleneck.) + + A stage is optional, so one that fails on the way to the trainer, or in + the reference log-probs it needs, falls back: the micro-batch goes back + to the front of `st.ready` and is sent whole in the next step's first + train call, as without staging, and this step keeps its get_metrics() + barrier, checkpoint and weight sync. Staging is then off for the rest of + the run. A stage the trainer ran and that raised there ends the run with + its error, as a failed train call does: past its refusals, + `TrainerWorker.stage_fwd_bwd` puts the worker in ERROR, so no checkpoint + could follow anyway. + + Args: + st: The next step's state, filled by `_prepack`. + step_result: The current step's final train call result. + + Returns: + Whether the trainer confirmed, once the micro-batch was staged, that + the step's update has landed (`StagedTrainStep.update_ready`): the + caller may then stamp and save the checkpoint without a get_metrics() + first, as after a fused call that waited for the update. False when + nothing was staged. + """ + if self._staging_latched_off: + return False + if self._writes_scored_checkpoint(step_result): + return False + if not self._intake_done() or st.staged is not None: + return False + if not st.ready or st.ready[0].is_final_batch: + return False + if self._needs_sampler_agreement(st.ready[0].payload): + return False + if not await self._trainer_can_stage(): + return False + sent = await self._send_stage(st) + if sent is None: + return False + staged, stage_s = sent + device_wait_s = staged.device_wait_s + logging.info( + "[pipeline] STAGED next step's first microbatch on the trainer in" + " %.2fs; its device then still computed for %s (request_id=%s)", + stage_s, + "unknown" if device_wait_s is None else f"{device_wait_s:.2f}s", + staged.request_id, + ) + if not self._staging_hidden(device_wait_s): + self._staging_latched_off = True + logging.warning( + "[pipeline] Staging a microbatch took %.2fs and the trainer's device" + " was idle before it finished (device wait %s): the trainer is fed" + " slower than it computes, so staging only delays each step's" + " weight sync. Not staging for the rest of the run.", + stage_s, + "unknown" if device_wait_s is None else f"{device_wait_s:.3f}s", + ) + return staged.update_ready + + async def _send_stage(self, st: _StepState) -> tuple[Any, float] | None: + """Stages `st.ready[0]` on the trainer; the callers checked it may go. + + A stage is optional, so one that fails on the way to the trainer, or in + the reference log-probs it needs, falls back: the micro-batch goes back + to the front of `st.ready` and staging is off for the rest of the run. A + stage the trainer ran and that raised there is re-raised. + + Args: + st: The next step's state, filled by `_prepack`. + + Returns: + The engine's `StagedTrainStep` and how long the call took, or None when + it fell back. + """ + intake, self._intake_task = self._intake_task, None + mb = st.ready.pop(0) + calling_trainer = False + try: + batch = await self._with_reference_logps(mb.payload) + start = time.monotonic() + calling_trainer = True + staged = await self.engine.stage_train_step( # pyrefly: ignore[missing-attribute] + batch, role=datatypes.Role.ACTOR + ) + stage_s = time.monotonic() - start + except Exception as exc: # pylint: disable=broad-exception-caught + if calling_trainer and _raised_by_worker_method(exc): + raise + # Undone exactly, so `_send_first_microbatch_ahead` sends it whole. Had + # the trainer received it after all (a reply lost on the way back), it + # holds it unused: with staging off nothing names it, and a successor's + # bring-up drops it. + st.ready.insert(0, mb) + self._intake_task = intake + self._staging_latched_off = True + logging.warning( + "[pipeline] Staging failed on the way to the trainer or in the" + " reference log-probs (%r); sending the microbatch whole in step" + " %d's first train call instead, and not staging for the rest of" + " the run.", + exc, + self._step + 1, + ) + return None + batch = None + st.staged = (mb._replace(payload=None), staged.request_id) + return staged, stage_s + + async def _stage_next_before_update( + self, st: _StepState, step_result: dict[str, Any] + ) -> dict[str, Any]: + """Stages the next step's first micro-batch, then runs the step's update. + + Called once the step's last train call, sent with its update deferred + (`_stage_planned`, `_trainer_can_stage_before_update`), has returned: the + trainer has dispatched that micro-batch and holds the update back. Under + MaxText's `retry_when_tokens_dropped` the call returns once the + micro-batch before it has finished (it reads that one's overflow flag), + so the device computes the whole last micro-batch while the stage's + transfer runs. The update then goes out when the last micro-batch is done, + where `update()` would have dispatched it anyway (it reads that one's flag + first), so nothing is delayed while the transfer is shorter than the + micro-batch. + + The micro-batch is staged under the refusals of + `_maybe_stage_next_microbatch`: not when it is not packed yet, closes its + step, or is rescored under the live weights first (those go out whole), + and not once staging is off. The scored-checkpoint refusal was applied by + `_stage_planned`, from the predicted optimizer step. A stage that fails on + the way falls back as there (`_send_stage`); the update still runs. + + The update call waits for the device first and says how long + (`device_wait_s`): how long the last micro-batch still ran once the staged + one was in place. Under `_STAGING_HIDDEN_WAIT_S`, the transfer outlasted + it and delayed the update, metrics, checkpoint stamp and weight sync: the + trainer is fed slower than it computes, so staging is switched off for the + rest of the run, as there. + + A replay of the last micro-batch (MaxText reruns an overflowed one, + dropless, inside `update()`) leaves the staged micro-batch as it is: its + inputs are mapped from its own payload alone (`prepare_fwd_bwd_input` + reads no parameter, gradient or optimizer state), and the rerun reads the + inputs it kept for itself. + + Args: + st: The next step's state, filled by `_prepack`. + step_result: The last train call's result (`update_deferred`). + + Returns: + `step_result` with the update call's result over it: `updated`, + `train_step`, `update_ready` (the trainer confirmed the update landed, + so the caller may stamp and save the checkpoint without a get_metrics() + first) and `device_wait_s`. + """ + assert self.engine is not None + sent = None + if ( + not self._staging_latched_off + and self._intake_done() + and st.staged is None + and st.ready + and not st.ready[0].is_final_batch + and not self._needs_sampler_agreement(st.ready[0].payload) + ): + sent = await self._send_stage(st) + update = await self.engine.apply_deferred_update( # pyrefly: ignore[missing-attribute] + role=datatypes.Role.ACTOR + ) + if sent is not None: + staged, stage_s = sent + device_wait_s = update.get("device_wait_s") + logging.info( + "[pipeline] STAGED next step's first microbatch on the trainer" + " before the update in %.2fs; the last microbatch then still" + " computed for %s (request_id=%s)", + stage_s, + "unknown" if device_wait_s is None else f"{device_wait_s:.2f}s", + staged.request_id, + ) + if not self._staging_hidden(device_wait_s): + self._staging_latched_off = True + logging.warning( + "[pipeline] Staging a microbatch took %.2fs and the trainer's" + " device was idle before it finished (device wait %s): the" + " trainer is fed slower than it computes, so staging only delays" + " each step's update and weight sync. Not staging for the rest of" + " the run.", + stage_s, + "unknown" if device_wait_s is None else f"{device_wait_s:.3f}s", + ) + return {**step_result, **update} + + async def _send_first_microbatch_ahead( + self, st: _StepState + ) -> asyncio.Task[tuple[Any, float]] | None: + """Sends the next step's first micro-batch, taken ahead by `_prepack`. + + Called once the current step's weight sync has staged and its batch is + committed: the point at which `_train_loop` first sends any of the next + step's micro-batches. The train call is left in `_pending_train` for the + next step to collect as its own. The next step's groups then keep being + taken, so its second micro-batch is packed while the first is in flight. + + Args: + st: The next step's state. + + Returns: + The train call's task, or None when nothing was ready to send. + """ + assert self._pending_train is None + if st.staged is not None: + mb, staged_request_id = st.staged + st.staged = None + st.start_time = time.monotonic() + batch = None + elif self._intake_done() and st.ready and not st.ready[0].is_final_batch: + self._intake_task = None + mb = st.ready.pop(0) + staged_request_id = None + st.start_time = time.monotonic() + batch = await self._with_reference_logps(mb.payload) + if self._needs_sampler_agreement(batch): + batch = await self._apply_sampler_trainer_agreement( + batch, st.sampler_agreement + ) + else: + return None + self._record_microbatch(st, mb) + mb = None + self._pending_train = asyncio.create_task( + self._timed_train_step( + batch, + apply_optimizer=False, + staged_request_id=staged_request_id, + ) + ) + batch = None + if ( + not st.ready + and st.scored_items + and st.groups_consumed < self.full_batch_size + ): + self._intake_task = asyncio.create_task(self._prepack(st)) + return self._pending_train + async def run_async( self, engine: rl_engine_interface.AbstractRLEngine, @@ -2446,6 +3246,8 @@ async def run_async( task.cancel() if self._pending_train is not None and not self._pending_train.done(): self._pending_train.cancel() + if self._intake_task is not None and not self._intake_task.done(): + self._intake_task.cancel() def run( self, diff --git a/tunix/experimental/queue_manager/trajectory_queue_manager.py b/tunix/experimental/queue_manager/trajectory_queue_manager.py index 6c6103b5ff..1078be76ac 100644 --- a/tunix/experimental/queue_manager/trajectory_queue_manager.py +++ b/tunix/experimental/queue_manager/trajectory_queue_manager.py @@ -706,9 +706,17 @@ def _maybe_seal(self, batch_idx: int) -> None: self._sealed.discard(batch_idx) def _advance_over_holes(self) -> None: - """Passes over any sealed empty batches sitting at `next_batch_idx`.""" + """Passes over any sealed empty batches sitting at `next_batch_idx`. + + Also when the serving cursor is already past the hole: a consumer that + reads the next batch by index before committing the current one (the + pipelined step boundary) moves it there, and the hole must still be + passed once the current batch commits. A batch being trained is never a + hole (it has admitted groups), so this cannot skip one. + """ + while ( - self._serving_batch_idx == self._next_batch_idx + self._serving_batch_idx >= self._next_batch_idx and self._next_batch_idx in self._sealed and self._admitted.get(self._next_batch_idx, 0) == 0 ): diff --git a/tunix/experimental/worker/trainer_worker.py b/tunix/experimental/worker/trainer_worker.py index 364b03dc97..23cda5d3ba 100644 --- a/tunix/experimental/worker/trainer_worker.py +++ b/tunix/experimental/worker/trainer_worker.py @@ -75,6 +75,25 @@ def __init__( self._worker_id = worker_id self._state = WorkerState.PENDING self._last_error: str | None = None + # One micro-batch received ahead of its fwd_bwd (`stage_fwd_bwd`): + # (request_id, input, fwd_bwd kwargs, request metadata). + self._staged_fwd_bwd: ( + tuple[str, Any, dict[str, Any], dict[str, Any]] | None + ) = None + # The request id of a step's last micro-batch whose update + # `fwd_bwd(apply_optimizer=True, defer_update=True)` left to + # `apply_deferred_update`; until then the step is open, so nothing that + # needs it closed (another micro-batch, metrics, a checkpoint, a weight + # sync, actor log-probs) is accepted. + self._deferred_update: str | None = None + + def _refuse_while_update_deferred(self, method: str) -> None: + if self._deferred_update is not None: + raise RuntimeError( + f"TrainerWorker.{method}() refused: the update of the step whose" + f" last micro-batch was {self._deferred_update!r} is deferred; run" + " apply_deferred_update() first." + ) def _policy_version(self) -> int: if self._trainer is None: @@ -119,11 +138,44 @@ def initialize(self) -> datatypes.Response: finally: self.state = WorkerState.READY + def _drop_stale_stage(self, by: str) -> None: + """Drops a micro-batch staged, or an update deferred, before a (re)compile + or a restore. + + Only `fwd_bwd_staged` and `stop()` release a stage otherwise, and only + `apply_deferred_update` a deferred update, so either left by an + orchestrator that died without `stop()` (an OOM kill, a preemption) would + make its successor's first stage, or first micro-batch, fail, and every + restart's after it. A successor brings the worker up with `compile` and + resumes with `restore_checkpoint`; either way the held micro-batch will + never run and the deferred update will never be applied. + + Args: + by: The call dropping it, for the log. + """ + if self._staged_fwd_bwd is not None: + logging.warning( + "TrainerWorker %s drops staged micro-batch %r, staged before it and" + " never run.", + by, + self._staged_fwd_bwd[0], + ) + self._staged_fwd_bwd = None + if self._deferred_update is not None: + logging.warning( + "TrainerWorker %s drops the update deferred after micro-batch %r," + " never applied.", + by, + self._deferred_update, + ) + self._deferred_update = None + def compile(self, dummy_data: Any = None) -> datatypes.Response: """Triggers JIT compilation using the provided dummy_data.""" if self.state == WorkerState.PENDING: self.initialize() self.state = WorkerState.COMPILING + self._drop_stale_stage("compile") try: self._trainer.compile(dummy_data) return self._response(compiled=True) @@ -149,6 +201,8 @@ def stop(self) -> datatypes.Response: if self.state == WorkerState.STOPPED: return self._response(stopped=True, already_stopped=True) self._is_running = False + self._staged_fwd_bwd = None + self._deferred_update = None if self.state == WorkerState.READY: self.state = WorkerState.DRAINING if self._trainer is not None: @@ -167,6 +221,18 @@ def info(self) -> datatypes.WorkerInfo: else "LazyTrainer" ), "policy_version": self._policy_version(), + # This worker can hold a micro-batch ahead of its fwd_bwd + # (`stage_fwd_bwd` / `fwd_bwd_staged`) ... + "stage_fwd_bwd": True, + # ... and say how long a stage then still waited for the device + # (`device_wait_s`), which needs the trainer's `wait_for_device`. + "reports_device_wait": callable( + getattr(self._trainer, "wait_for_device", None) + ), + # ... and take that micro-batch between a step's last fwd_bwd and + # its update: `fwd_bwd(apply_optimizer=True, defer_update=True)` + # leaves the update to `apply_deferred_update`. + "stage_before_update": True, }, ) @@ -222,23 +288,59 @@ def fwd_bwd( self, request: datatypes.TrainRequest, apply_optimizer: bool = False, + *, + wait_for_update: bool = True, + defer_update: bool = False, **kwargs: Any, ) -> datatypes.Response: - """Executes one forward/backward pass, optionally applying the optimizer.""" + """Executes one forward/backward pass, optionally applying the optimizer. + + Args: + request: The micro-batch. + apply_optimizer: Also apply the accumulated gradients, in this call. + wait_for_update: With `apply_optimizer`, return only once the update has + landed when the trainer can tell (`update_ready`). False returns at + its dispatch, for a caller that sends the next step's first + micro-batch ahead (`stage_fwd_bwd`) while the device finishes; that + stage then waits for the device and confirms the update instead. + defer_update: With `apply_optimizer`, the trainer is still told this is + the step's last micro-batch, but the update is left to + `apply_deferred_update`, for a caller that stages the next step's + first micro-batch in between (`stage_fwd_bwd`). Until then the worker + accepts nothing that needs the step closed. Never passed to the + trainer. + **kwargs: Passed to the trainer's `fwd_bwd`. + + Returns: + The response, with `updated`, `train_step` and, once confirmed, + `update_ready` in its metadata when the optimizer ran, or + `update_deferred` when it was deferred. + """ self._ensure_ready() + if defer_update and not apply_optimizer: + raise ValueError( + "fwd_bwd(defer_update=True) needs apply_optimizer=True: only a" + " step's last micro-batch has an update to defer." + ) + self._refuse_while_update_deferred("fwd_bwd") req_metadata = dict(request.metadata) if request.metadata else {} req_metadata.pop("updated", None) req_metadata.pop("train_step", None) req_metadata.pop("update_ready", None) + req_metadata.pop("update_deferred", None) kwargs.pop("skip_jit", None) try: # The trainer still sees `apply_optimizer`: it marks the step's last # micro-batch (MaxText's retry_dropless_final_micro_batch runs that one - # with the dropless program). The update itself is run here, below. + # with the dropless program). The update itself is run here, below, or + # by `apply_deferred_update`. self._trainer.fwd_bwd( request.payload, apply_optimizer=apply_optimizer, **kwargs ) - if apply_optimizer: + if apply_optimizer and defer_update: + self._deferred_update = str(request.request_id) + req_metadata["update_deferred"] = True + elif apply_optimizer: try: # Through this worker's update(), not the trainer's, so a fused step # gets the same bookkeeping (state on failure, step timing) as an @@ -248,9 +350,10 @@ def fwd_bwd( # the update executable, return only once it is ready and say so: the # caller stamps the checkpoint as soon as this call returns, and with # `update_ready` it no longer needs get_metrics() to wait for the - # update before the checkpoint save and weight sync. + # update before the checkpoint save and weight sync. Not when the + # caller asked to get control back at dispatch (`wait_for_update`). update_token = getattr(self._trainer, "last_update_token", None) - if update_token is not None: + if wait_for_update and update_token is not None: jax.block_until_ready(update_token) req_metadata["update_ready"] = True except Exception as exc: @@ -273,9 +376,237 @@ def fwd_bwd( self.state = WorkerState.ERROR raise + def stage_fwd_bwd( + self, + request: datatypes.TrainRequest, + **kwargs: Any, + ) -> datatypes.Response: + """Receives a micro-batch ahead of the fwd_bwd that runs it. + + Holds it as is or, when the trainer has `prepare_fwd_bwd_input`, as what + that returns (e.g. its inputs already placed on the devices). Nothing in + the model, optimizer or gradient state changes, so the trainer's step + order is that of the `fwd_bwd` / `fwd_bwd_staged` calls, and this may + land while a weight sync holds the worker. One at a time: a second stage + before the first ran is refused rather than replacing it (a `compile` or + `restore_checkpoint`, as a restarted orchestrator's bring-up, drops it). + + Then, when the trainer has `wait_for_device`, waits for the work it had + dispatched before this call and reports how long that took + (`device_wait_s`): ~0 means the device was already idle, i.e. the + transfer did not hide behind its work. The wait is one the next call + (`get_metrics`, which needs the update landed) would make anyway. + + A stage follows a step's last train call, sent with + `fwd_bwd(apply_optimizer=True, wait_for_update=False)`, so the update that + call dispatched is the last one. When the trainer exposes an output of it + (`last_update_token`), this also returns only once it is ready and says + so (`update_ready`), as a fused `fwd_bwd` would have. + + Or it follows a last train call sent with `defer_update=True`, before the + step's update. Then it waits for nothing and confirms nothing: the device + is meant to be still computing that micro-batch, and the wait (and the + confirmation) belong to `apply_deferred_update`, which runs the update + next. `device_wait_s` is then None and `before_update` is set. + + Args: + request: The micro-batch. + **kwargs: Passed to the trainer's `fwd_bwd` when it runs. + + Returns: + A response naming the staged request, with `device_wait_s` (None when + the trainer cannot tell, or when staged before the update) and, once + confirmed, `update_ready` in its metadata. + """ + if self.state == WorkerState.PENDING: + self.initialize() + if self.state in ( + WorkerState.ERROR, + WorkerState.STOPPED, + WorkerState.DRAINING, + ): + raise RuntimeError( + "TrainerWorker cannot stage a micro-batch in state" + f" {self.state.value}." + ) + if self._staged_fwd_bwd is not None: + raise RuntimeError( + "TrainerWorker already holds staged micro-batch" + f" {self._staged_fwd_bwd[0]!r}; run it with fwd_bwd_staged() first." + ) + req_metadata = dict(request.metadata) if request.metadata else {} + # Only this worker sets these; a caller-supplied copy is not echoed. + for key in ("staged", "device_wait_s", "update_ready", "before_update"): + req_metadata.pop(key, None) + kwargs.pop("skip_jit", None) + update_ready = False + device_wait_s = None + before_update = self._deferred_update is not None + try: + staged_input = request.payload + prepare = getattr(self._trainer, "prepare_fwd_bwd_input", None) + if callable(prepare): + staged_input = prepare(request.payload) + if not before_update: + # Both waits are for the work dispatched before this call, which ends + # with the step's optimizer update, so a failure in either is that + # update's and is reported as the fused call reports it. + waiting_for = "the device" + try: + wait_for_device = getattr(self._trainer, "wait_for_device", None) + device_wait_s = ( + float(wait_for_device()) if callable(wait_for_device) else None + ) + # After the device wait, which already covers the update when the + # trainer queues it (MaxText's throttler holds its gradient norm, so + # a failed update raises there first). + waiting_for = "the update" + update_token = getattr(self._trainer, "last_update_token", None) + if update_token is not None: + jax.block_until_ready(update_token) + update_ready = True + except Exception as exc: + raise RuntimeError( + "optimizer update() failed: dispatched by" + " fwd_bwd(apply_optimizer=True, wait_for_update=False), it did" + f" not land (seen by stage_fwd_bwd waiting for {waiting_for}):" + f" {type(exc).__name__}: {exc}" + ) from exc + except Exception as exc: + self._last_error = str(exc) + self.state = WorkerState.ERROR + raise + self._staged_fwd_bwd = ( + request.request_id, + staged_input, + kwargs, + req_metadata, + ) + response_metadata = { + **req_metadata, + "staged": True, + "device_wait_s": device_wait_s, + } + if update_ready: + response_metadata["update_ready"] = True + if before_update: + response_metadata["before_update"] = True + resp = self._response(**response_metadata) + resp.request_id = request.request_id + return resp + + def fwd_bwd_staged( + self, + staged_request_id: str, + **kwargs: Any, + ) -> datatypes.Response: + """Runs the forward/backward pass of the micro-batch `stage_fwd_bwd` holds. + + It is `fwd_bwd` on the held input, so it takes the same options + (`apply_optimizer`, which also runs the update in this call, and + `wait_for_update`) and gives the same response. + + Args: + staged_request_id: The id `stage_fwd_bwd` was given; any other id is + refused, so a staged micro-batch is never trained in another's place. + **kwargs: Merged over the kwargs given at staging, then passed to + `fwd_bwd`. + + Returns: + The response `fwd_bwd` gives for that micro-batch. + """ + self._ensure_ready() + staged = self._staged_fwd_bwd + if staged is None or staged[0] != staged_request_id: + raise RuntimeError( + f"TrainerWorker holds no staged micro-batch {staged_request_id!r}" + f" (holds {staged[0] if staged else None!r})." + ) + self._staged_fwd_bwd = None + request_id, staged_input, stage_kwargs, req_metadata = staged + # The held input may be what `prepare_fwd_bwd_input` made of the payload; + # the trainer's `fwd_bwd` takes either. + request = datatypes.TrainRequest( + request_id=request_id, + payload=staged_input, + metadata=req_metadata, + ) + return self.fwd_bwd(request, **{**stage_kwargs, **kwargs}) + + def apply_deferred_update(self) -> datatypes.Response: + """Runs the update `fwd_bwd(apply_optimizer=True, defer_update=True)` left. + + Between the two, the caller staged the next step's first micro-batch + (`stage_fwd_bwd`), so its transfer ran while the device computed this + step's last micro-batch. This first waits for the device when the trainer + has `wait_for_device` and reports how long (`device_wait_s`): how long that + micro-batch still ran once the staged one was in place. ~0 means the + device had gone idle, i.e. the transfer outlasted it and delayed the + update. The wait costs nothing where the trainer's `update()` blocks until + the step's last micro-batch has finished anyway (MaxText with + `retry_when_tokens_dropped`, which reads its overflow flag on the host + first); elsewhere it holds the update's dispatch back by one readiness + round trip. + + Then runs the update through `update()` and, when the trainer exposes an + output of it (`last_update_token`), returns only once it has landed and + says so (`update_ready`), as a fused `fwd_bwd` would have. + + Returns: + A response with `updated`, `train_step`, `device_wait_s` (None when the + trainer cannot tell) and, once confirmed, `update_ready` in its + metadata. + """ + self._ensure_ready() + if self._deferred_update is None: + raise RuntimeError( + "TrainerWorker holds no deferred update: apply_deferred_update()" + " follows fwd_bwd(apply_optimizer=True, defer_update=True)." + ) + metadata: dict[str, Any] = {} + waiting_for = "the device" + try: + try: + wait_for_device = getattr(self._trainer, "wait_for_device", None) + device_wait_s = ( + float(wait_for_device()) if callable(wait_for_device) else None + ) + waiting_for = "the update" + # Cleared only here: update() refuses to run while an update is + # deferred, and the device wait above has to come first, because it + # is what tells the caller whether the stage hid behind the step's + # last micro-batch. + self._deferred_update = None + train_step = self.update() + update_token = getattr(self._trainer, "last_update_token", None) + if update_token is not None: + jax.block_until_ready(update_token) + metadata["update_ready"] = True + except Exception as exc: + # Waiting for the device is waiting for the step's last micro-batch, + # whose failure surfaces there; the rest is the update's. + raise RuntimeError( + "deferred optimizer update failed: left by" + " fwd_bwd(apply_optimizer=True, defer_update=True), seen by" + f" apply_deferred_update waiting for {waiting_for}:" + f" {type(exc).__name__}: {exc}" + ) from exc + except Exception as exc: + self._last_error = str(exc) + self.state = WorkerState.ERROR + raise + self._last_error = None + return self._response( + updated=True, + train_step=train_step, + device_wait_s=device_wait_s, + **metadata, + ) + def update(self, **kwargs) -> int: """Applies the accumulated (mean) gradients as one optimizer update.""" self._ensure_ready() + self._refuse_while_update_deferred("update") try: train_step = self._trainer.update(**kwargs) self._last_error = None @@ -292,6 +623,7 @@ def eval_step( ) -> datatypes.Response: """Executes one evaluation step on the given payload.""" self._ensure_ready() + self._refuse_while_update_deferred("eval_step") req_metadata = dict(request.metadata) if request.metadata else {} try: self._trainer.eval_step(request.payload, **kwargs) @@ -307,6 +639,7 @@ def eval_step( def run_eval(self, eval_ds: Any, **kwargs) -> datatypes.Response: """Runs an explicit evaluation phase over eval micro-batches.""" self._ensure_ready() + self._refuse_while_update_deferred("run_eval") if eval_ds is None: return self._response(evaluated=True, eval_batches=0) try: @@ -354,6 +687,7 @@ def per_token_logps( """ del kwargs # Accepted for engine-forwarding parity; unused. self._ensure_ready() + self._refuse_while_update_deferred("per_token_logps") if items.pad_id is None or items.eos_id is None: raise ValueError( "TrainerWorker.per_token_logps requires pad_id and eos_id to be set " @@ -437,6 +771,7 @@ def _resolve_checkpoint_path(self, metadata: Mapping[str, Any]) -> str: def save_checkpoint(self, metadata: Any, **kwargs) -> datatypes.Response: """Force the trainer to serialize its state (model + optimizer).""" self._ensure_ready() + self._refuse_while_update_deferred("save_checkpoint") try: self._trainer.save_checkpoint(metadata, **kwargs) ckpt_path = self._resolve_checkpoint_path(metadata) @@ -450,6 +785,7 @@ def save_checkpoint(self, metadata: Any, **kwargs) -> datatypes.Response: def restore_checkpoint(self, **kwargs) -> Any: """Restore state from latest checkpoint and return the metadata pytree.""" self._ensure_ready() + self._drop_stale_stage("restore_checkpoint") try: result = self._trainer.restore_checkpoint(**kwargs) self._last_error = None @@ -462,6 +798,7 @@ def restore_checkpoint(self, **kwargs) -> Any: def prepare_weight_sync(self, sync_request: Any = None, **kwargs) -> Any: """Stages weights for transfer and returns their metadata.""" self._ensure_ready() + self._refuse_while_update_deferred("prepare_weight_sync") self.state = WorkerState.SYNCING try: if sync_request is not None: @@ -528,4 +865,5 @@ def release_weight_sync(self, sync_request: Any = None, **kwargs) -> Any: def get_metrics(self) -> Any: """Returns and clears the recently collected step metric records.""" + self._refuse_while_update_deferred("get_metrics") return self._trainer.get_metrics() diff --git a/tunix/rl/common.py b/tunix/rl/common.py index 36ced9511f..2c51c0aa3f 100644 --- a/tunix/rl/common.py +++ b/tunix/rl/common.py @@ -14,6 +14,8 @@ """Common RL helper classes and functions.""" import functools +import math +import os import re import inspect from typing import Any, Iterable, Sequence @@ -23,6 +25,7 @@ from flax import nnx import jax from jax import numpy as jnp +from jax.custom_derivatives import SymbolicZero # pylint: disable=g-importing-member from jax.typing import ArrayLike # pylint: disable=g-importing-member import jax.tree_util as jtu import numpy as np @@ -639,21 +642,36 @@ def _result(*outs): input_tokens_to_keep = input_tokens[:, -logits_to_keep:] if chunk_size > 0: - hidden_state = outputs[:, -logits_to_keep - 1 : -1, :] - out = compute_chunked_logps( - model, - hidden_state, - input_tokens_to_keep, - temperature, - chunk_size, - return_entropy, - ) + out = None + if segment_ids is not None: + # Packed rows score every position, so the head can run on the + # unsliced, still evenly sharded hidden states (opt-in, see below). It + # returns [B, T], already aligned like the padded output below. + out = maybe_compute_sharded_chunked_logps( + model, + outputs, + input_tokens, + temperature, + chunk_size, + return_entropy, + ) + aligned = out is not None + if out is None: + hidden_state = outputs[:, -logits_to_keep - 1 : -1, :] + out = compute_chunked_logps( + model, + hidden_state, + input_tokens_to_keep, + temperature, + chunk_size, + return_entropy, + ) if return_entropy: per_token_logps, per_token_entropy = out else: per_token_logps = out - if segment_ids is not None: + if segment_ids is not None and not aligned: per_token_logps = jnp.pad( per_token_logps, ((0, 0), (1, 0)), constant_values=0.0 ) @@ -996,6 +1014,392 @@ def logp_step(carry, xs): return per_token_logps +# Sharded chunked log-prob head, opt-in via TUNIX_LOGPS_SHARDED_HEAD. +# +# Under GSPMD, `compute_chunked_logps` pays per chunk: an all-gather of every +# head weight (a [4096, 248320] vocab table is 2 GB in bf16) in the forward and +# again in the rematerialized backward, and an all-reduce of the full-size table +# gradient. And because `lax.scan` iterates over the sequence axis, a sequence +# sharded over a context-parallel axis is gathered, so every context rank +# computes every chunk. The sharded head runs the same per-chunk math inside +# `shard_map`: each head parameter is all-gathered once per pass, each device +# scans only the tokens it holds, the table gradient is accumulated per device in +# float32 and reduce-scattered once, and chunks are read and written with +# dynamic slices rather than reshapes. +SHARDED_HEAD_ENV = "TUNIX_LOGPS_SHARDED_HEAD" +# Logical names of the hidden states' batch and length axes. When the active +# flax logical-axis rules (MaxText's) map them, the head uses that layout. +_LOGICAL_BATCH_AXIS = "activation_batch" +_LOGICAL_LENGTH_AXIS = "activation_length" + + +def sharded_head_mode() -> str | None: + """Returns None (off) or "local", from TUNIX_LOGPS_SHARDED_HEAD.""" + value = os.environ.get(SHARDED_HEAD_ENV, "").strip().lower() + if value in ("", "0", "false", "off", "no"): + return None + if value in ("1", "true", "on", "yes", "local"): + return "local" + raise ValueError(f"{SHARDED_HEAD_ENV}={value!r}: expected 0, 1 or local.") + + +def _mesh_axes(entry) -> tuple[str, ...]: + if entry is None: + return () + if isinstance(entry, str): + return (entry,) + return tuple(entry) + + +def _token_partition(mesh, batch: int, seq_len: int): + """Mesh axes sharding the [batch, seq] dims of the hidden states. + + Every mesh axis of size > 1 must shard one of the two, so that every device + holds distinct tokens and the gradient sum runs over all devices. + + Returns: + (batch_axes, seq_axes, source), or None if no such layout exists. + """ + sizes = dict(mesh.shape) + nontrivial = {a for a, n in sizes.items() if n > 1} + + def valid(b, s): + if set(b) & set(s) or not set(b + s) <= set(sizes): + return False + if not nontrivial <= set(b + s): + return False + return ( + batch % math.prod(sizes[a] for a in b) == 0 + and seq_len % math.prod(sizes[a] for a in s) == 0 + ) + + rules = flax.linen.get_logical_axis_rules() + if rules: + try: + spec = flax.linen.logical_to_mesh_axes( + (_LOGICAL_BATCH_AXIS, _LOGICAL_LENGTH_AXIS), rules + ) + b, s = _mesh_axes(spec[0]), _mesh_axes(spec[1]) + if valid(b, s): + return b, s, "logical-axis-rules" + except Exception: # pylint: disable=broad-exception-caught + pass + b, s = [], [] + for a in mesh.axis_names: + n = sizes[a] + if n == 1: + continue + if batch % (math.prod(sizes[x] for x in b) * n) == 0: + b.append(a) + elif seq_len % (math.prod(sizes[x] for x in s) * n) == 0: + s.append(a) + else: + return None + return tuple(b), tuple(s), "mesh-order" + + +def _use_dtypes(jaxpr, var) -> set[Any]: + """dtypes `var` is consumed in: a convert's target, else its own dtype.""" + out = set() + for eqn in jaxpr.eqns: + for k, v in enumerate(eqn.invars): + if v is not var: + continue + if eqn.primitive.name == "convert_element_type": + out.add(jnp.dtype(eqn.params["new_dtype"])) + continue + inner = eqn.params.get("jaxpr", eqn.params.get("call_jaxpr")) + inner = getattr(inner, "jaxpr", inner) + if inner is not None and len(getattr(inner, "invars", ())) == len( + eqn.invars + ): + out |= _use_dtypes(inner, inner.invars[k]) + else: + out.add(jnp.dtype(var.aval.dtype)) + return out + + +def _head_parameters(graphdef, treedef, abstract_leaves, hidden_struct): + """Leaves `compute_final_logits` reads, mapped to the dtype to gather in. + + A leaf whose every use is a cast to one narrower float dtype (MaxText casts an + fp32 vocab table to bf16 before the dot) is gathered after that cast, which is + what XLA already does for the gather inside `compute_chunked_logps`. + """ + + def head(leaves, hs): + model = nnx.merge(graphdef, jax.tree.unflatten(treedef, leaves)) + return model.compute_final_logits(hs) + + jaxpr = jax.make_jaxpr(head)(list(abstract_leaves), hidden_struct).jaxpr + out = {} + for i, var in enumerate(jaxpr.invars[: len(abstract_leaves)]): + dtypes = _use_dtypes(jaxpr, var) + if not dtypes: + continue + own = jnp.dtype(abstract_leaves[i].dtype) + dtype = next(iter(dtypes)) if len(dtypes) == 1 else own + narrower = ( + jnp.issubdtype(dtype, jnp.floating) and dtype.itemsize < own.itemsize + ) + out[i] = dtype if narrower else own + return out + + +def maybe_compute_sharded_chunked_logps( + model, + hidden_states, + input_tokens, + temperature, + chunk_size, + return_entropy, +): + """`compute_chunked_logps` for packed rows, with a shard_map'ed vocab head. + + Off unless TUNIX_LOGPS_SHARDED_HEAD is set, and it falls back (returns None, + with a log line naming the reason) whenever its preconditions do not hold: + a mesh in context, every non-trivial mesh axis sharding the batch or the + sequence, and every head parameter divisible over the mesh on some dimension. + + Args: + model: The merged actor model; must expose `compute_final_logits`. + hidden_states: [B, T, H], the model output for every position. + input_tokens: [B, T] token ids; position t is scored from position t - 1. + temperature: Sampling temperature the logits are divided by (0 or 1: none). + chunk_size: Tokens per scan step on each device. + return_entropy: Whether to also return the per-token entropy. + + Returns: + None, or the packed-row result of `compute_per_token_logps`: per-token + log-probs [B, T] whose index j scores `input_tokens[:, j]` from + `hidden_states[:, j - 1]` (index 0 is 0), plus the entropy [B, T] if + `return_entropy`. + """ + mode = sharded_head_mode() + if mode is None: + return None + + def fallback(reason): + logging.warning( + "[logps-head] sharded head off, using compute_chunked_logps: %s", reason + ) + return None + + mesh = jax.sharding.get_abstract_mesh() + if mesh.empty: + return fallback("no mesh in context (jax.set_mesh)") + batch, seq_len, hidden_dim = hidden_states.shape + partition = _token_partition(mesh, batch, seq_len) + if partition is None: + return fallback(f"no token layout for {batch}x{seq_len} on {mesh.shape}") + batch_axes, seq_axes, layout_source = partition + gather_axes = tuple(a for a in mesh.axis_names if mesh.shape[a] > 1) + num_devices = math.prod(mesh.shape[a] for a in gather_axes) + local_batch = batch // math.prod(mesh.shape[a] for a in batch_axes) + local_len = seq_len // math.prod(mesh.shape[a] for a in seq_axes) + chunk = min(chunk_size, local_len) + + graphdef, state = nnx.split(model) + leaves, treedef = jax.tree.flatten(state) + abstract = [ + jax.ShapeDtypeStruct(jnp.shape(x), jnp.result_type(x)) for x in leaves + ] + head = _head_parameters( + graphdef, + treedef, + abstract, + jax.ShapeDtypeStruct( + (local_batch, chunk, hidden_dim), hidden_states.dtype + ), + ) + if not head: + return fallback("compute_final_logits reads no model state") + head_index = sorted(head) + params = [leaves[i] for i in head_index] + gather_dtypes = [head[i] for i in head_index] + gather_dims = [] + for p in params: + dim = next((d for d, n in enumerate(p.shape) if n % num_devices == 0), None) + if dim is None: + return fallback( + f"head parameter {p.shape} not divisible by {num_devices}" + ) + gather_dims.append(dim) + param_dtypes = [jnp.result_type(p) for p in params] + num_out = 2 if return_entropy else 1 + logging.info( + "[logps-head] sharded head %s: tokens batch=%s seq=%s (%s), %dx%d local" + " tokens in %d steps of %d; gathers %s over %s", + mode, + batch_axes, + seq_axes, + layout_source, + local_batch, + local_len, + -(-local_len // chunk), + chunk, + [(p.shape, str(d)) for p, d in zip(params, gather_dtypes)], + gather_axes, + ) + + def merged(full): + values = list(abstract) + for i, x in zip(head_index, full): + values[i] = x + return nnx.merge(graphdef, jax.tree.unflatten(treedef, values)) + + def chunk_outputs(full, hs, ids, with_entropy): + logits = merged(full).compute_final_logits(hs).astype(jnp.float32) + if temperature != 0.0 and temperature != 1.0: + logits /= temperature + logps = selective_log_softmax(logits, ids) + if with_entropy: + return logps, compute_entropy_from_logits(logits) + return (logps,) + + def gather(p, dim, dtype): + p = p.astype(dtype) + if not gather_axes: + return p + return jax.lax.all_gather(p, gather_axes, axis=dim, tiled=True) + + def scatter(g, dim): + if not gather_axes: + return g + return jax.lax.psum_scatter( + g, gather_axes, scatter_dimension=dim, tiled=True + ) + + def scan_forward(full, h, t): + def step(bufs, i): + start = i * chunk + hs = jax.lax.dynamic_slice_in_dim(h, start, chunk, axis=1) + ids = jax.lax.dynamic_slice_in_dim(t, start, chunk, axis=1) + outs = chunk_outputs(full, hs, ids, return_entropy) + bufs = tuple( + jax.lax.dynamic_update_slice_in_dim(b, o, start, axis=1) + for b, o in zip(bufs, outs) + ) + return bufs, None + + bufs = tuple(jnp.zeros(t.shape, jnp.float32) for _ in range(num_out)) + bufs, _ = jax.lax.scan(step, bufs, jnp.arange(h.shape[1] // chunk)) + return bufs + + @jax.custom_vjp + def head_local(h, t, *ps): + full = [ + gather(p, d, dt) for p, d, dt in zip(ps, gather_dims, gather_dtypes) + ] + return scan_forward(full, h, t) + + def head_local_fwd(h, t, *ps): + h, t = h.value, t.value + ps = [p.value for p in ps] + full = [ + gather(p, d, dt) for p, d, dt in zip(ps, gather_dims, gather_dtypes) + ] + return scan_forward(full, h, t), (h, t, tuple(full)) + + def head_local_bwd(res, cts): + h, t, full = res + d_logps = cts[0] + d_entropy = cts[1] if return_entropy else SymbolicZero(None) + use_logps = not isinstance(d_logps, SymbolicZero) + use_entropy = not isinstance(d_entropy, SymbolicZero) + shard_shapes = [ + tuple(n // num_devices if k == d else n for k, n in enumerate(f.shape)) + for f, d in zip(full, gather_dims) + ] + if not (use_logps or use_entropy): + zeros = [jnp.zeros(s, dt) for s, dt in zip(shard_shapes, param_dtypes)] + return (jnp.zeros_like(h), None, *zeros) + + def step(carry, i): + acc, dh = carry + start = i * chunk + hs = jax.lax.dynamic_slice_in_dim(h, start, chunk, axis=1) + ids = jax.lax.dynamic_slice_in_dim(t, start, chunk, axis=1) + cot = [] + if use_logps: + cot.append(jax.lax.dynamic_slice_in_dim(d_logps, start, chunk, axis=1)) + if use_entropy: + cot.append( + jax.lax.dynamic_slice_in_dim(d_entropy, start, chunk, axis=1) + ) + + def outputs_fn(full_, hs_): + outs = chunk_outputs(full_, hs_, ids, use_entropy) + return outs if use_logps else outs[1:] + + _, vjp = jax.vjp(outputs_fn, list(full), hs) + d_full, d_hs = vjp(tuple(cot)) + acc = tuple(a + g.astype(jnp.float32) for a, g in zip(acc, d_full)) + dh = jax.lax.dynamic_update_slice_in_dim( + dh, d_hs.astype(dh.dtype), start, axis=1 + ) + return (acc, dh), None + + acc = tuple(jnp.zeros(f.shape, jnp.float32) for f in full) + (acc, dh), _ = jax.lax.scan( + step, (acc, jnp.zeros_like(h)), jnp.arange(h.shape[1] // chunk) + ) + acc = [scatter(a, d) for a, d in zip(acc, gather_dims)] + d_params = [a.astype(dt) for a, dt in zip(acc, param_dtypes)] + return (dh, None, *d_params) + + head_local.defvjp(head_local_fwd, head_local_bwd, symbolic_zeros=True) + + seq_shards = math.prod(mesh.shape[a] for a in seq_axes) + + def neighbour(x, step): + """`x` of the sequence shard `step` places later; zeros past either end.""" + if seq_shards == 1: + return jnp.zeros_like(x) + perm = [ + (i, i - step) for i in range(seq_shards) if 0 <= i - step < seq_shards + ] + return jax.lax.ppermute(x, seq_axes, perm) + + def local(h, tok, *ps): + # Position t predicts token t + 1; a shard's last position predicts the + # next shard's first token, and the sequence's last position (no target) + # is scored against token 0 and dropped. Both one-token shifts are explicit + # collective-permutes, so neither depends on how XLA partitions a shift of + # a sequence-sharded array. + t = jnp.concatenate([tok[:, 1:], neighbour(tok[:, :1], 1)], axis=1) + pad = -h.shape[1] % chunk + if pad: + h = jnp.pad(h, ((0, 0), (0, pad), (0, 0))) + t = jnp.pad(t, ((0, 0), (0, pad))) + outs = [o[:, : o.shape[1] - pad] for o in head_local(h, t, *ps)] + # Index j holds token j's log-prob (scored at position j - 1); index 0 is 0. + return tuple( + jnp.concatenate([neighbour(o[:, -1:], -1), o[:, :-1]], axis=1) + for o in outs + ) + + P = jax.sharding.PartitionSpec # pylint: disable=invalid-name + token_spec = P(batch_axes or None, seq_axes or None) + in_specs = ( + P(batch_axes or None, seq_axes or None, None), + token_spec, + *[ + P(*[gather_axes if k == d else None for k in range(p.ndim)]) + for p, d in zip(params, gather_dims) + ], + ) + hidden_states = jnp.where(jnp.isfinite(hidden_states), hidden_states, 0.0) + outs = jax.shard_map( + local, + mesh=mesh, + in_specs=in_specs, + out_specs=(token_spec,) * num_out, + check_vma=False, + )(hidden_states, input_tokens, *params) + return outs if return_entropy else outs[0] + + @nnx.jit(static_argnames=("pad_id", "eos_id", "stop_gradient")) def compute_score( model,