Run most R3 e2e tests through the output store

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.

Author: guapisolo

miles/utils/external_utils/command_utils/__init__.py

@@ -16,6 +16,7 @@
     get_default_wandb_args,
     get_env_enable_infinite_run,
     get_mooncake_object_store_args,
+    get_output_store_args,
     repo_base_dir,
     rsync_cmd,
 )
@@ -36,6 +37,7 @@
     "get_default_wandb_args",
     "get_env_enable_infinite_run",
     "get_mooncake_object_store_args",
+    "get_output_store_args",
     "repo_base_dir",
     "resolve_extra_env_vars",
     "resolve_hardware",

miles/utils/external_utils/command_utils/common.py

@@ -237,6 +237,11 @@ def get_mooncake_object_store_args(
     )
 
 
+def get_output_store_args(global_segment_size: str | None = None) -> str:
+    """--use-output-store with the Mooncake object store it lives in."""
+    return f"{get_mooncake_object_store_args(global_segment_size=global_segment_size)}--use-output-store "
+
+
 def _is_tcp_server_ready(host: str, port: int) -> bool:
     try:
         with socket.create_connection((host, port), timeout=1):

miles/utils/external_utils/command_utils/legacy.py

@@ -19,6 +19,7 @@
     get_default_wandb_args,
     get_env_enable_infinite_run,
     get_mooncake_object_store_args,
+    get_output_store_args,
     repo_base_dir,
     rsync_cmd,
 )
@@ -48,6 +49,7 @@
     "get_default_wandb_args",
     "get_env_enable_infinite_run",
     "get_mooncake_object_store_args",
+    "get_output_store_args",
     "hf_download_dataset",
     "repo_base_dir",
     "resolve_hardware",

tests/e2e/fsdp/r3/_common.py

@@ -25,6 +25,8 @@ class CaseConfig:
     rollout_num_gpus_per_engine: int
     sglang_attention_backend: str = "fa3"
     attn_implementation: str = "flash_attention_3"
+    # Send the replay arrays through the output store (--use-output-store) instead of inline.
+    use_output_store: bool = False
 
 
 def prepare(case: CaseConfig) -> None:
@@ -83,6 +85,8 @@ def execute(case: CaseConfig, wandb_file: str) -> None:
     )
 
     replay_args = "--use-rollout-routing-replay "
+    if case.use_output_store:
+        replay_args += command_utils.get_output_store_args()
 
     sglang_args = (
         f"--rollout-num-gpus-per-engine {case.rollout_num_gpus_per_engine} "

tests/e2e/fsdp/r3/test_glm47_flash_r3.py

@@ -24,6 +24,7 @@
 register_ci_gate(metric_key="perf/actor_train_mfu")
 
 CASE = CaseConfig(
+    use_output_store=True,
     model_name="GLM-4.7-Flash",
     hf_repo="zai-org/GLM-4.7-Flash",
     num_gpus=8,

tests/e2e/fsdp/r3/test_qwen3_5_35b_a3b_r3.py

@@ -12,6 +12,7 @@
 register_ci_gate(metric_key="perf/actor_train_mfu")
 
 CASE = CaseConfig(
+    use_output_store=True,
     model_name="Qwen3.5-35B-A3B",
     hf_repo="Qwen/Qwen3.5-35B-A3B",
     num_gpus=8,

tests/e2e/megatron/model_scripts/test_deepseek_v32_5layer_fp8.py

@@ -16,6 +16,8 @@
 from tests.ci.ci_register import register_cuda_ci
 from tests.ci.metric_history import register_ci_gate
 
+from miles.utils.external_utils import command_utils
+
 register_cuda_ci(
     est_time=1300, suite="stage-c-4-gpu-h200", labels=["megatron", "model-scripts"], hardware=["hopper"], num_gpus=4
 )
@@ -45,6 +47,7 @@ def _args() -> ScriptArgs:
         extra_args=(
             "--ci-test --check-weight-update-allow-quant-error --bf16 "
             "--use-rollout-routing-replay "
+            f"{command_utils.get_output_store_args()}"
             "--skip-actor-forward-only "
             "--sglang-disable-shared-experts-fusion "
             # TP2 doubles the per-rank activations of the script's TP4 token budget

tests/e2e/megatron/test_glm47_flash/_common.py

@@ -39,6 +39,8 @@ class CaseConfig:
     # Off: no draft runs, the trainer builds no MTP layer, and the weight check covers the target only.
     use_spec: bool = True
     use_r3: bool = True
+    # Send the replay arrays through the output store (--use-output-store) instead of inline.
+    use_output_store: bool = False
     # "bshd" (Megatron only) trains fixed micro-batches of micro_batch_size samples instead of
     # dynamic THD batches.
     qkv_format: str = "thd"
@@ -282,6 +284,8 @@ def build_train_args(case: CaseConfig, *, wandb_file: str) -> str:
 
     if case.fully_async:
         misc_args += "--fully-async "
+    if case.use_output_store:
+        misc_args += command_utils.get_output_store_args()
 
     if case.use_deepep:
         misc_args += "--moe-token-dispatcher-type flex --moe-enable-deepep "

tests/e2e/megatron/test_glm47_flash/test_nospec_r3_bf16_deepep_sgl_dpattn2x2_meg_tp2cp2_bshd.py

@@ -45,6 +45,7 @@
     sglang_deepep_dispatcher_output_dtype="bf16",
     use_spec=False,
     use_r3=True,
+    use_output_store=True,
 )
 
 

tests/e2e/megatron/test_glm47_flash/test_spec_mtptrain_r3_bf16_sgl_tp2_meg_tp1cp2_disagg.py

@@ -24,6 +24,7 @@
 # engine, with R3, MTP training and EAGLE on. Each train GPU holds the dense part plus half the
 # experts, so the response cap and max_tokens_per_gpu keep each CP rank to about 2k tokens.
 CASE = CaseConfig(
+    use_output_store=True,
     use_deepep=False,
     num_gpus_per_node=2,
     cp_size=2,