Repository navigation
Run most R3 e2e tests through the output store, and sweep each run's leftovers - #4478
Conversation
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
There was a problem hiding this comment.
Claude Code Review
This repository is configured for manual code reviews. Comment @claude review for a one-time review, or @claude review always to subscribe this PR to a review on every future push.
Tip: disable this comment in your organization's Code Review settings.
|
/rerun-test tests/e2e/fsdp/r3/test_glm47_flash_r3.py |
|
/rerun-test tests/e2e/fsdp/r3/test_qwen3_5_35b_a3b_r3.py |
|
/rerun-test tests/e2e/megatron/model_scripts/test_deepseek_v32_5layer_fp8.py |
|
/rerun-test tests/e2e/megatron/test_glm47_flash/test_nospec_r3_bf16_deepep_sgl_dpattn2x2_meg_tp2cp2_bshd.py |
|
/rerun-test tests/e2e/megatron/test_glm47_flash/test_spec_mtptrain_r3_bf16_sgl_tp2_meg_tp1cp2_disagg.py |
|
/rerun-test tests/e2e/megatron/test_glm47_flash/test_spec_mtptrain_r3_bf16_sgl_tp4_meg_tp2pp2.py |
|
/rerun-test tests/e2e/megatron/test_qwen3_30B_A3B/test_r3_deepep.py |
|
/rerun-test tests/e2e/megatron/test_qwen3_5_35B_A3B/test_nospec_r3_bf16_sgl_dpattn1x4_meg_tp1cp1_fullyasync.py |
|
/rerun-test tests/e2e/megatron/test_qwen3_5_35B_A3B/test_spec_mtptrain_r3_fp8_deepep_sgl_dpattn1x4_meg_tp2cp1.py |
|
/rerun-test tests/e2e/torchtitan/test_glm47_flash_tp_ep_r3.py |
|
/rerun-test tests/e2e/torchtitan/test_qwen3_30B_A3B_fully_async_r3.py |
|
/rerun-test tests/e2e/torchtitan/test_qwen3_30B_A3B_tp_pp_ep_cp.py |
|
✅ |
|
❌ The execution job failed; inspect the workflow run for the failing step. |
|
✅ |
|
✅ |
|
✅ |
|
✅ |
|
✅ |
|
✅ |
|
✅ |
|
✅ |
|
✅ |
|
✅ |
5e2af1c to
351b5be
Compare
|
/rerun-test tests/e2e/megatron/model_scripts/test_deepseek_v32_5layer_fp8.py |
|
✅ |
351b5be to
aacde6d
Compare
73ea084 to
0901944
Compare
Nothing deleted what a run still held when it ended: untrained rollouts, buffered or in-flight fully-async groups, retried attempts and failed deletes stayed in the store, and an external master kept them across runs. The rollout executor's dispose now removes every key under the run's prefix with a remove-only Mooncake client. The remove is forced, since a read leaves a lease that blocks a plain one. Under --ci-test a synchronous run, which released every rollout, fails when the sweep removed anything; a fully-async run only logs it.
The CUDA image installs the output store package from the SGLang checkout; the ROCm image did not, so --use-output-store failed there at the first import. The store runs over Mooncake in host memory, which the ROCm image already ships.
Each R3 builder gains an off-by-default use_output_store case field, and twelve tests across Megatron, FSDP and torchtitan turn it on, so CI exercises the store with fully async, MTP speculative decoding, DeepEP, FP8, CP and PP. command_utils.get_output_store_args holds the flags. fsdp/r3/test_qwen3_30b_a3b_r3.py stays inline to guard the inline R3 path.
The run-end sweep failed every synchronous --ci-test run that left an object, and raised on a failed sweep in any run. A run that stopped before training a rollout still holds that rollout's objects, so a crash ended in the sweep's AssertionError, with the real error only as its context. A --partial-rollout run ends with aborted groups back in the data buffer, so a correct run failed. And a Mooncake failure at teardown turned a finished non-CI run into a failed one, though a failed delete only logs. get() now files every rollout until its release, so dispose can tell a run that stopped early. The sweep holds a CI run to leaving nothing, and to a working sweep, only when no rollout is left unreleased; a partial-rollout run only logs what it left, like a fully-async one; outside CI a failed sweep only logs. A dispose test checks the sweep runs once, after the rest of teardown.
0901944 to
3c67a49
Compare
Summary
Run twelve R3 e2e tests through the output store, and sweep each run's store objects when it ends.
Motivation
Only three e2e tests sent replay arrays through the output store. Fully async, MTP speculative decoding, DeepEP, FP8, the FSDP and torchtitan trainers and the ROCm image never ran it, and nothing checked cleanup: objects a run still held when it ended stayed in the store.
Usage
Design Notes
command_utils.get_output_store_args()only for a case that setsuse_output_store, where the trainer's R3 replay check guards the store read as it guards the inline one; at run end,RolloutExecutor.disposeremoves every key left undermiles-output-store/<run_uuid>/.--ci-test, a run that trained every rollout fails if the sweep removed anything, unless--fully-asyncor--partial-rolloutbuffered groups.fsdp/r3/test_qwen3_30b_a3b_r3.pyas the inline R3 guard, in tests whose flags the store rejects, in the dumper, MiMo and disabled tests.radix_transfer_plugins, as the CUDA image does.Verification
use_output_store=Trueon/rerun-test: eleven at5e2af1c0b3, DSV3.2 at351b5beef0after main's fix(mfu): read dsa_indexer_loss_coeff off the Megatron backend #4474 fix.351b5beef0(Linux devbox, CPU): the fast suites add no failure overmain(a6d60232ca) except [2/6] Carry the output store plan on samples, the session reply and the train dict #4242's Ray-start error.73ea084045: rollout and output_store fast tests add no failure overaacde6dd20.sweep_output_storeremoved a run's three keys (one under a read lease) but no other key.Review Focus
RolloutExecutor.dispose.