Re-enable ft-short FT e2e tests in CI

miles/utils/audit_utils/checksum_utils.py

@@ -14,7 +14,7 @@ def flatten_inference_engine_checksums(check_weights_result: Any) -> list[Infere
 
 
 def _flatten_to_inference_engine_bodies(check_weights_result: Any) -> list[Any]:
-    return [engine_body for server in check_weights_result for server_group in server for engine_body in server_group]
+    return [engine_body for server_group in check_weights_result for engine_body in server_group]
 
 
 def _merge_inference_engine_ranks(engine_body: dict[str, Any]) -> InferenceEngineChecksums:
@@ -24,11 +24,18 @@ def _merge_inference_engine_ranks(engine_body: dict[str, Any]) -> InferenceEngin
     ranks: list[dict[str, Any]] = engine_body.get("ranks", []) or []
     assert ranks, f"check_weights engine body has no ranks: {engine_body!r}"
 
-    ranks_sorted = sorted(ranks, key=lambda r: r["parallelism_info"]["rank"])
+    ranks_sorted = sorted(ranks, key=_gpu_rank)
 
     merged: InferenceEngineChecksums = {}
     for rank_info in ranks_sorted:
-        rank = rank_info["parallelism_info"]["rank"]
+        rank = _gpu_rank(rank_info)
         for name, value in rank_info["checksums"].items():
             merged[f"rank{rank}/{name}"] = value
     return merged
+
+
+def _gpu_rank(rank_info: dict[str, Any]) -> int:
+    parallelism_info = rank_info["parallelism_info"]
+    gpu_ranks = {role_info["rank"] for role_info in parallelism_info}
+    assert len(gpu_ranks) == 1, f"expected one GPU rank across roles, got {gpu_ranks}: {rank_info!r}"
+    return next(iter(gpu_ranks))

tests/e2e/ft/README.md

@@ -3,7 +3,7 @@
 ## Layout
 
 - Scenario logic lives in `conftest_ft/scenario_<name>.py`.
-- CI runs it via thin per-mode entry files `test_trainer_ft_<scenario>_<mode>.py`, each registered with `register_cuda_ci(est_time=..., suite="stage-c-8-gpu-h100", labels=["ft-short"])` (comparison scenarios) or `labels=["ft-long"]` (soak scenarios).
+- CI runs it via thin per-mode entry files `test_trainer_ft_<scenario>_<mode>.py`, each registered with `register_cuda_ci(est_time=..., suite="stage-c-8-gpu-h200", labels=["ft-short"])` (comparison scenarios) or `labels=["ft-long"]` (soak scenarios).
 - The CUDA CI runner executes each entry as bare `python3 <file>` (exit code = pass/fail); the entry just calls the scenario's `run_ci(mode)`.
 
 | Scenario (`conftest_ft/scenario_*.py`) | Type | What it verifies |
@@ -36,7 +36,7 @@
 
 ### In CI
 
-- Gated on the `run-ci-ft-short` / `run-ci-ft-long` PR labels (FT is expensive — not run on every PR). `ft-short` covers the comparison scenarios (no_failure / deterministic / with_failure, minutes each); `ft-long` covers the soak scenarios (random-crash survival, realistic-gsm8k convergence — tens of minutes to hours). With a label set, the matching entries run on `stage-c-8-gpu-h100`.
+- Gated on the `run-ci-ft-short` / `run-ci-ft-long` PR labels (FT is expensive — not run on every PR). `ft-short` covers the comparison scenarios (no_failure / deterministic / with_failure, minutes each); `ft-long` covers the soak scenarios (random-crash survival, realistic-gsm8k convergence — tens of minutes to hours). With a label set, the matching entries run on `stage-c-8-gpu-h200`.
 - Add a `(scenario, mode)` to CI: copy an entry file, change `run_ci(...)`'s mode.
 - Add a new label: edit `tests/ci/labels.py` and create the matching `run-ci-<label>` GitHub label.
 

tests/e2e/ft/conftest_ft/app.py

@@ -1,6 +1,7 @@
 # NOTE: You MUST read tests/e2e/ft/README.md as source-of-truth and documentations
 
 import os
+import shutil
 from collections.abc import Callable
 from pathlib import Path
 from typing import Annotated
@@ -44,23 +45,26 @@ def run_pipeline(
 
     prepare(ft_mode)
 
-    for phase in effective_phases:
-        baseline_dump = f"{dump_dir}/{_dump_subdir('baseline', phase)}"
-        run_training(
-            train_args=build_baseline_args(ft_mode, baseline_dump, enable_dumper),
-            mode=ft_mode,
-            dump_dir=baseline_dump,
-        )
-
-        target_dump = f"{dump_dir}/{_dump_subdir('target', phase)}"
-        run_training(
-            train_args=build_target_args(ft_mode, target_dump, enable_dumper),
-            mode=ft_mode,
-            dump_dir=target_dump,
-        )
-
-    if enable_dumper:
-        compare_fn(dump_dir, ft_mode)
+    try:
+        for phase in effective_phases:
+            baseline_dump = f"{dump_dir}/{_dump_subdir('baseline', phase)}"
+            run_training(
+                train_args=build_baseline_args(ft_mode, baseline_dump, enable_dumper),
+                mode=ft_mode,
+                dump_dir=baseline_dump,
+            )
+
+            target_dump = f"{dump_dir}/{_dump_subdir('target', phase)}"
+            run_training(
+                train_args=build_target_args(ft_mode, target_dump, enable_dumper),
+                mode=ft_mode,
+                dump_dir=target_dump,
+            )
+
+        if enable_dumper:
+            compare_fn(dump_dir, ft_mode)
+    finally:
+        shutil.rmtree(dump_dir, ignore_errors=True)
 
 
 def create_comparison_app_and_run_ci(

tests/e2e/ft/conftest_ft/execution.py

@@ -104,7 +104,6 @@ def get_common_train_args(
             "--rollout-temperature 0.8 "
             "--rollout-batch-size 32 "
             "--n-samples-per-prompt 8 "
-            "--sglang-disable-cuda-graph "
             # Required for reproducibility (ref: https://github.com/THUDM/slime/pull/370)
             "--sglang-enable-deterministic-inference "
             "--sglang-attention-backend flashinfer "
@@ -174,6 +173,17 @@ def get_ft_args(mode: FTTestMode) -> str:
 }
 
 
+def get_train_env_vars_arg(mode: FTTestMode, *, deterministic: bool) -> str:
+    env_vars: dict[str, str] = {}
+    if deterministic:
+        env_vars.update(_DETERMINISTIC_ENV_VARS)
+    if mode.has_real_rollout:
+        env_vars["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
+    if not env_vars:
+        return ""
+    return f"--train-env-vars '{json.dumps(env_vars)}' "
+
+
 def run_training(
     train_args: str,
     mode: FTTestMode,
@@ -203,7 +213,7 @@ def run_training(
     }
     U.execute_train(
         train_args=train_args,
-        num_gpus_per_node=mode.train_gpus_per_node,
+        num_gpus_per_node=mode.total_node_gpus,
         megatron_model_type=mode.megatron_model_type,
         extra_env_vars=merged_env_vars,
         megatron_path=_MEGATRON_PATH,

tests/e2e/ft/conftest_ft/modes.py

@@ -40,6 +40,10 @@ def has_real_rollout(self) -> bool:
     def total_rollout_gpus(self) -> int:
         return self.rollout_num_engines * self.rollout_gpus_per_engine
 
+    @property
+    def total_node_gpus(self) -> int:
+        return self.train_gpus_per_node + self.total_rollout_gpus
+
 
 MODES: dict[str, FTTestMode] = {
     # --- 1-node (8 GPUs) variants ---

tests/e2e/ft/conftest_ft/scenario_deterministic.py

@@ -5,7 +5,7 @@
 from pathlib import Path
 
 from tests.e2e.ft.conftest_ft.app import create_comparison_app_and_run_ci
-from tests.e2e.ft.conftest_ft.execution import get_common_train_args, get_ft_args
+from tests.e2e.ft.conftest_ft.execution import get_common_train_args, get_ft_args, get_train_env_vars_arg
 from tests.e2e.ft.conftest_ft.modes import FTTestMode
 
 from miles.utils.test_utils.comparisons.dumps import (
@@ -22,12 +22,6 @@
 TOTAL_NUM_ROLLOUTS: int = 2 * NUM_ROLLOUTS_PER_PHASE
 PHASE_START_ROLLOUT_IDS: dict[str, int] = {"phase_a": 0, "phase_b": NUM_ROLLOUTS_PER_PHASE}
 
-_DETERMINISTIC_ENV_VARS: str = (
-    '--train-env-vars \'{"NCCL_ALGO": "Ring", '
-    '"NVTE_ALLOW_NONDETERMINISTIC_ALGO": "0", '
-    '"CUBLAS_WORKSPACE_CONFIG": ":4096:8"}\' '
-)
-
 
 def _build_actions(phase_start_rollout_id: int) -> list[dict]:
     heal_trigger_rollout_id: int = phase_start_rollout_id + 1
@@ -64,7 +58,7 @@ def _build_phase_args(mode: FTTestMode, dump_dir: str, *, is_target: bool, enabl
     phase_start_rollout_id: int = PHASE_START_ROLLOUT_IDS[phase_name]
 
     base = get_common_train_args(mode, dump_dir=dump_dir, num_steps=TOTAL_NUM_ROLLOUTS, enable_dumper=enable_dumper)
-    base += "--deterministic-mode " + _DETERMINISTIC_ENV_VARS
+    base += "--deterministic-mode " + get_train_env_vars_arg(mode, deterministic=True)
     base += "--debug-deterministic-collective "
 
     if is_target:

tests/e2e/ft/conftest_ft/scenario_no_failure.py

@@ -3,7 +3,7 @@
 
 
 from tests.e2e.ft.conftest_ft.app import create_comparison_app_and_run_ci
-from tests.e2e.ft.conftest_ft.execution import get_common_train_args, get_ft_args
+from tests.e2e.ft.conftest_ft.execution import get_common_train_args, get_ft_args, get_train_env_vars_arg
 from tests.e2e.ft.conftest_ft.modes import FTTestMode
 
 from miles.utils.test_utils.comparisons.dumps import (
@@ -17,13 +17,17 @@
 
 
 def _build_baseline_args(mode: FTTestMode, dump_dir: str, enable_dumper: bool = True) -> str:
-    return get_common_train_args(mode, dump_dir=dump_dir, num_steps=NUM_STEPS, enable_dumper=enable_dumper)
+    return get_common_train_args(
+        mode, dump_dir=dump_dir, num_steps=NUM_STEPS, enable_dumper=enable_dumper
+    ) + get_train_env_vars_arg(mode, deterministic=False)
 
 
 def _build_target_args(mode: FTTestMode, dump_dir: str, enable_dumper: bool = True) -> str:
-    return get_common_train_args(
-        mode, dump_dir=dump_dir, num_steps=NUM_STEPS, enable_dumper=enable_dumper
-    ) + get_ft_args(mode)
+    return (
+        get_common_train_args(mode, dump_dir=dump_dir, num_steps=NUM_STEPS, enable_dumper=enable_dumper)
+        + get_ft_args(mode)
+        + get_train_env_vars_arg(mode, deterministic=False)
+    )
 
 
 def _compare(dump_dir: str, mode: FTTestMode) -> None:

tests/e2e/ft/conftest_ft/scenario_with_failure.py

@@ -5,7 +5,7 @@
 from pathlib import Path
 
 from tests.e2e.ft.conftest_ft.app import create_comparison_app_and_run_ci
-from tests.e2e.ft.conftest_ft.execution import get_common_train_args, get_ft_args
+from tests.e2e.ft.conftest_ft.execution import get_common_train_args, get_ft_args, get_train_env_vars_arg
 from tests.e2e.ft.conftest_ft.modes import FTTestMode
 
 from miles.utils.test_utils.comparisons.dumps import (
@@ -88,6 +88,7 @@ def _expected_reconfigures(*, is_target: bool, phase: str, num_cells: int) -> li
 def _build_phase_args(mode: FTTestMode, dump_dir: str, *, is_target: bool, enable_dumper: bool = True) -> str:
     is_phase_a: bool = dump_dir.endswith("phase_a")
     base = get_common_train_args(mode, dump_dir=dump_dir, num_steps=NUM_PHASE_B_STEPS, enable_dumper=enable_dumper)
+    base += get_train_env_vars_arg(mode, deterministic=False)
 
     if is_target:
         base += get_ft_args(mode)

tests/e2e/ft/test_trainer_ft_deterministic_dp2_cp2_pp2.py

@@ -7,10 +7,9 @@
 from tests.e2e.ft.conftest_ft.scenario_deterministic import run_ci
 
 register_cuda_ci(
-    est_time=900,
-    suite="stage-c-8-gpu-h100",
+    est_time=2900,
+    suite="stage-c-8-gpu-h200",
     labels=["ft-short"],
-    disabled="FT e2e suite disabled in CI pending stabilization on current image (follow-up PR)",
 )
 
 _MODE: str = "dp2_cp2_pp2"

tests/e2e/ft/test_trainer_ft_deterministic_dp2_cp2_real_rollout.py

@@ -7,10 +7,9 @@
 from tests.e2e.ft.conftest_ft.scenario_deterministic import run_ci
 
 register_cuda_ci(
-    est_time=1500,
-    suite="stage-c-8-gpu-h100",
+    est_time=3200,
+    suite="stage-c-8-gpu-h200",
     labels=["ft-short"],
-    disabled="real_rollout deterministic-inference sglang hang on current image",
 )
 
 _MODE: str = "dp2_cp2_real_rollout"

tests/e2e/ft/test_trainer_ft_deterministic_dp2_cp2_tp2_ep2.py

@@ -7,10 +7,9 @@
 from tests.e2e.ft.conftest_ft.scenario_deterministic import run_ci
 
 register_cuda_ci(
-    est_time=900,
-    suite="stage-c-8-gpu-h100",
+    est_time=2300,
+    suite="stage-c-8-gpu-h200",
     labels=["ft-short"],
-    disabled="FT e2e suite disabled in CI pending stabilization on current image (follow-up PR)",
 )
 
 _MODE: str = "dp2_cp2_tp2_ep2"

tests/e2e/ft/test_trainer_ft_deterministic_dp4_cp2.py

@@ -7,10 +7,9 @@
 from tests.e2e.ft.conftest_ft.scenario_deterministic import run_ci
 
 register_cuda_ci(
-    est_time=900,
-    suite="stage-c-8-gpu-h100",
+    est_time=2900,
+    suite="stage-c-8-gpu-h200",
     labels=["ft-short"],
-    disabled="FT e2e suite disabled in CI pending stabilization on current image (follow-up PR)",
 )
 
 _MODE: str = "dp4_cp2"

tests/e2e/ft/test_trainer_ft_no_failure_dp2_cp2_pp2.py

@@ -7,10 +7,9 @@
 from tests.e2e.ft.conftest_ft.scenario_no_failure import run_ci
 
 register_cuda_ci(
-    est_time=600,
-    suite="stage-c-8-gpu-h100",
+    est_time=1000,
+    suite="stage-c-8-gpu-h200",
     labels=["ft-short"],
-    disabled="FT e2e suite disabled in CI pending stabilization on current image (follow-up PR)",
 )
 
 _MODE: str = "dp2_cp2_pp2"

tests/e2e/ft/test_trainer_ft_no_failure_dp2_cp2_real_rollout.py

@@ -8,9 +8,8 @@
 
 register_cuda_ci(
     est_time=1200,
-    suite="stage-c-8-gpu-h100",
+    suite="stage-c-8-gpu-h200",
     labels=["ft-short"],
-    disabled="real_rollout deterministic-inference sglang hang on current image",
 )
 
 _MODE: str = "dp2_cp2_real_rollout"

tests/e2e/ft/test_trainer_ft_no_failure_dp2_cp2_tp2_ep2.py

@@ -7,10 +7,9 @@
 from tests.e2e.ft.conftest_ft.scenario_no_failure import run_ci
 
 register_cuda_ci(
-    est_time=600,
-    suite="stage-c-8-gpu-h100",
+    est_time=800,
+    suite="stage-c-8-gpu-h200",
     labels=["ft-short"],
-    disabled="FT e2e suite disabled in CI pending stabilization on current image (follow-up PR)",
 )
 
 _MODE: str = "dp2_cp2_tp2_ep2"

tests/e2e/ft/test_trainer_ft_no_failure_dp4_cp2.py

@@ -7,10 +7,9 @@
 from tests.e2e.ft.conftest_ft.scenario_no_failure import run_ci
 
 register_cuda_ci(
-    est_time=600,
-    suite="stage-c-8-gpu-h100",
+    est_time=1100,
+    suite="stage-c-8-gpu-h200",
     labels=["ft-short"],
-    disabled="FT e2e suite disabled in CI pending stabilization on current image (follow-up PR)",
 )
 
 _MODE: str = "dp4_cp2"

tests/e2e/ft/test_trainer_ft_random_dp2_cp2_real_rollout.py

@@ -7,7 +7,7 @@
 from tests.e2e.ft.conftest_ft.scenario_ft_random import run_ci
 
 register_cuda_ci(
-    est_time=2400, suite="stage-c-8-gpu-h100", labels=["ft-long"], disabled="FT soak tests pending CI infra support"
+    est_time=2400, suite="stage-c-8-gpu-h200", labels=["ft-long"], disabled="FT soak tests pending CI infra support"
 )
 
 _MODE: str = "dp2_cp2_real_rollout"

tests/e2e/ft/test_trainer_ft_random_dp2_cp2_tp2_ep2.py

@@ -7,7 +7,7 @@
 from tests.e2e.ft.conftest_ft.scenario_ft_random import run_ci
 
 register_cuda_ci(
-    est_time=1800, suite="stage-c-8-gpu-h100", labels=["ft-long"], disabled="FT soak tests pending CI infra support"
+    est_time=1800, suite="stage-c-8-gpu-h200", labels=["ft-long"], disabled="FT soak tests pending CI infra support"
 )
 
 _MODE: str = "dp2_cp2_tp2_ep2"

tests/e2e/ft/test_trainer_ft_realistic_gsm8k.py

@@ -8,7 +8,7 @@
 
 register_cuda_ci(
     est_time=9000,
-    suite="stage-c-8-gpu-h100",
+    suite="stage-c-8-gpu-h200",
     labels=["ft-long", "long"],
     disabled="FT soak tests pending CI infra support",
 )

tests/e2e/ft/test_trainer_ft_with_failure_dp2_cp2_pp2.py

@@ -7,10 +7,9 @@
 from tests.e2e.ft.conftest_ft.scenario_with_failure import run_ci
 
 register_cuda_ci(
-    est_time=900,
-    suite="stage-c-8-gpu-h100",
+    est_time=2100,
+    suite="stage-c-8-gpu-h200",
     labels=["ft-short"],
-    disabled="FT e2e suite disabled in CI pending stabilization on current image (follow-up PR)",
 )
 
 _MODE: str = "dp2_cp2_pp2"

tests/e2e/ft/test_trainer_ft_with_failure_dp2_cp2_real_rollout_dense.py

@@ -7,10 +7,9 @@
 from tests.e2e.ft.conftest_ft.scenario_with_failure import run_ci
 
 register_cuda_ci(
-    est_time=1500,
-    suite="stage-c-8-gpu-h100",
+    est_time=1700,
+    suite="stage-c-8-gpu-h200",
     labels=["ft-short"],
-    disabled="real_rollout deterministic-inference sglang hang on current image",
 )
 
 _MODE: str = "dp2_cp2_real_rollout_dense"

tests/e2e/ft/test_trainer_ft_with_failure_dp2_cp2_tp2_ep2.py

@@ -7,10 +7,9 @@
 from tests.e2e.ft.conftest_ft.scenario_with_failure import run_ci
 
 register_cuda_ci(
-    est_time=900,
-    suite="stage-c-8-gpu-h100",
+    est_time=1700,
+    suite="stage-c-8-gpu-h200",
     labels=["ft-short"],
-    disabled="FT e2e suite disabled in CI pending stabilization on current image (follow-up PR)",
 )
 
 _MODE: str = "dp2_cp2_tp2_ep2"