Make the cell based train group the only trainer implementation

docs/diffusion/advanced/deterministic.md

@@ -9,7 +9,7 @@ that do not support deterministic execution may still pass validation.
 
 ## What it turns on
 
-### At actor spawn (`miles/ray/actor_group.py`)
+### At actor spawn (`miles/ray/train/actor_factory.py`)
 
 ```bash
 NCCL_DETERMINISTIC=1

examples/ppo/README.md

@@ -72,8 +72,8 @@ These are enforced at argument validation, so you get an error rather than a sil
 * **`--kl-coef` must be 0.** Reward-level KL is rejected because the critic trains *before* the
   actor and never sees ref log probs, so its value targets would silently exclude the KL penalty
   applied to the actor's rewards. Use loss-level `--use-kl-loss` / `--kl-loss-coef` instead.
-* **Not compatible with `MILES_EXPERIMENTAL_FT_TRAINER=1`.** The v2 fault-tolerant train group
-  cannot route critic values yet.
+* **Not compatible with `--indep-dp` (which train fault tolerance implies).** Shared actor/critic
+  PPO hands the critic outputs to a single trainer cell as external data.
 
 ## Which numbers here are verified
 

miles/backends/megatron_utils/ft/indep_dp.py

@@ -7,7 +7,6 @@
 from megatron.core import mpu
 
 from miles.utils.distributed_utils import get_gloo_group
-from miles.utils.environ import enable_experimental_ft_trainer
 from miles.utils.ft_utils.indep_dp import IndepDPInfo
 from miles.utils.ft_utils.process_group_utils import GeneralPGUtil, GroupInfo, collective_bool_and
 from miles.utils.tracking_utils.structured_log import log_structured
@@ -148,8 +147,6 @@ def allreduce_grads_and_losses_across_replicas(
                 for bucket in bucket_group.buckets:
                     util.all_reduce(bucket.grad_data, pg, op=dist.ReduceOp.SUM)
     except Exception:
-        if not enable_experimental_ft_trainer():
-            raise
         allreduce_success = False
         log_structured(
             logger.error,

miles/ray/placement_group.py

@@ -7,22 +7,14 @@
 from ray.util.placement_group import PlacementGroup, placement_group
 from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
 
-from miles.utils.environ import enable_experimental_ft_trainer
+from miles.ray.train.group import RayTrainGroup
 from ..utils.ray_utils import compute_ray_pin_head_options
 from .rollout.inference_controller import InferenceController
 from .rollout.rollout_executor import RolloutExecutor
 
 logger = logging.getLogger(__name__)
 
 
-def _select_train_group_class():
-    if enable_experimental_ft_trainer():
-        from miles.ray.train.group import RayTrainGroup
-    else:
-        from miles.ray.actor_group import RayTrainGroup
-    return RayTrainGroup
-
-
 @ray.remote(num_gpus=1)
 class InfoActor:
     def get_ip_and_gpu_id(self):
@@ -143,8 +135,7 @@ def allocate_train_group(
     rollout_executor,
     with_opd_teacher: bool = False,
 ):
-    train_group_cls = _select_train_group_class()
-    return train_group_cls(
+    return RayTrainGroup(
         args=args,
         num_nodes=num_nodes,
         num_gpus_per_node=num_gpus_per_node,

miles/ray/train/group.py

@@ -289,6 +289,13 @@ async def save_model(self, rollout_id: int, force_sync: bool = False):
             max_attempts=_RETRY_MAX_ATTEMPTS,
         )
 
+    async def export_hf(self, rollout_id: int, path: str):
+        """Export current weights as an HF checkpoint. Only cell 0 exports to avoid file write conflicts."""
+        await retry(
+            lambda _: self._execute_first_alive("export_hf", rollout_id=rollout_id, path=path),
+            max_attempts=_RETRY_MAX_ATTEMPTS,
+        )
+
     async def update_weights(self, rollout_id: int | None = None):
         """Broadcast weights to rollout engines."""
         log_structured(logger.info, tag="ft", op="update_weights", phase="start", rollout=rollout_id)

miles/utils/arguments.py

@@ -12,7 +12,7 @@
 from miles.dashboard.args import add_dashboard_arguments, validate_dashboard_args
 from miles.rollout.checkpoint_eval import is_checkpoint_eval_fn
 from miles.utils.chat_template_utils.tito_tokenizer import TITOTokenizerType
-from miles.utils.environ import enable_experimental_ft_trainer, use_legacy_rollout_v1
+from miles.utils.environ import use_legacy_rollout_v1
 from miles.utils.eval_config import EvalDatasetConfig, build_eval_dataset_configs, ensure_dataset_list
 from miles.utils.file_arg_utils import resolve_file_arg
 from miles.utils.ft_utils.health_checker import SimpleHealthCheckerConfig
@@ -3279,11 +3279,6 @@ def miles_validate_args(args):
         )
         if args.train_backend != "megatron":
             raise ValueError("Shared Actor/Critic PPO requires the Megatron backend")
-        assert not enable_experimental_ft_trainer(), (
-            "Shared Actor/Critic PPO is not supported with MILES_EXPERIMENTAL_FT_TRAINER=1: the v2 "
-            "fault-tolerant train group cannot route critic values or lifecycle options yet. "
-            "Unset MILES_EXPERIMENTAL_FT_TRAINER or use a non-PPO advantage estimator."
-        )
         assert args.kl_coef == 0, (
             "Shared Actor/Critic PPO does not support reward-level KL (--kl-coef): the critic "
             "trains before the actor and never sees ref log probs, so its value targets would "

miles/utils/dumper_utils.py

@@ -16,7 +16,6 @@
 
 from miles.backends.sglang_utils.sglang_config import resolve_sglang_config
 from miles.backends.training_utils.parallel import get_parallel_state
-from miles.utils.environ import enable_experimental_ft_trainer
 from miles.utils.ft_utils.process_group_utils import GeneralPGUtil
 from miles.utils.retry_utils import retry_until_deadline
 from miles.utils.tracking_utils.structured_log import log_structured
@@ -67,8 +66,6 @@ async def configure_sglang(args: Namespace) -> None:
 
     engines_dir: Path = _get_dir(args) / "engines"
     _cleanup_dump_dir(engines_dir, indep_dp_rank=0)
-    if not enable_experimental_ft_trainer() and dist.is_initialized():
-        dist.barrier()
 
     coros = []
     for i, url in enumerate(worker_urls):
@@ -140,8 +137,7 @@ def finalize(self, model: Sequence[torch.nn.Module]) -> None:
         get_grad: Callable[[torch.nn.Parameter], torch.Tensor | None] | None = None
         if self.phase is DumperPhase.FWD_BWD and self.overrides.get("enable_model_grad"):
             _log_model_grad_coverage(extracted_model)
-            if enable_experimental_ft_trainer():
-                get_grad = _build_full_grad_getter(extracted_model)
+            get_grad = _build_full_grad_getter(extracted_model)
 
         # Weights/grads are a once-per-rollout end-state, so pin them to step 0 instead of
         # the running per-microbatch step. _configure already cleaned the scoped paths;

miles/utils/environ.py

@@ -26,18 +26,3 @@ def default_fp8_block_scaling_fp32_scales() -> str:
         return "1"
     major, _minor = torch.cuda.get_device_capability()
     return "0" if major >= 10 else "1"
-
-
-_printed_experimental_ft_trainer = False
-
-
-def enable_experimental_ft_trainer() -> bool:
-    raw = os.environ.get("MILES_EXPERIMENTAL_FT_TRAINER", "0").lower()
-    result = raw in ("1", "true", "on", "yes")
-
-    global _printed_experimental_ft_trainer
-    if result and not _printed_experimental_ft_trainer:
-        print("MILES_EXPERIMENTAL_FT_TRAINER=1 is enabled (experimental feature)")
-        _printed_experimental_ft_trainer = True
-
-    return result

tests/e2e/ft/conftest_ft/execution.py

@@ -168,13 +168,6 @@ def get_ft_args(mode: FTTestMode) -> str:
     "SGLANG_FLASHINFER_PREFILL_SPLIT_TILE_SIZE": "8192",
 }
 
-# Selects v2 RayTrainGroup (miles.ray.train.group). Required because
-# --ft-components train depends on cell-based indep_dp; the v1 default path
-# does not support it.
-_TRAINER_FT_ENV_VARS: dict[str, str] = {
-    "MILES_EXPERIMENTAL_FT_TRAINER": "1",
-}
-
 
 def get_train_env_vars_arg(mode: FTTestMode, *, deterministic: bool) -> str:
     env_vars: dict[str, str] = {}
@@ -198,7 +191,6 @@ def run_training(
         shutil.rmtree(dump_dir)
     merged_env_vars = {
         **_DETERMINISTIC_ENV_VARS,
-        **_TRAINER_FT_ENV_VARS,
         # Run eager (no torch.compile). A cell respawned after a crash cold-recompiles its first
         # forward; under dynamic batch sizes that is a per-shape Inductor compile that is slow
         # (observed 124s..1510s, growing) and memory-heavy enough to OOM-kill the actor. That

tests/e2e/ft/conftest_ft/scenario_realistic_gsm8k.py

@@ -3,6 +3,7 @@
 
 import os
 import shutil
+from pathlib import Path
 from typing import Annotated
 
 import typer
@@ -11,6 +12,7 @@
 from tests.e2e.ft.conftest_ft.fault_injection import API_SERVER_PORT, MEAN_INTERVAL_SECONDS, spawn_fault_injector
 
 import miles.utils.external_utils.command_utils as U
+from miles.utils.test_utils.reconfigure_assertions import assert_soak_reconfigure_events
 
 app: typer.Typer = typer.Typer()
 
@@ -59,9 +61,6 @@ def run_ci(
             num_gpus_per_node=_TRAIN_GPUS + _ROLLOUT_GPUS,
             megatron_model_type=_MODEL_TYPE,
             extra_env_vars={
-                # --ft-components train depends on cell-based indep_dp, which only
-                # the v2 RayTrainGroup supports.
-                "MILES_EXPERIMENTAL_FT_TRAINER": "1",
                 # Same as run_training: a cell respawned after a crash cold-recompiles
                 # its first forward, which is slow and memory-heavy enough to OOM.
                 "TORCHDYNAMO_DISABLE": "1",
@@ -72,6 +71,11 @@ def run_ci(
     finally:
         injector.stop_and_join(timeout_seconds=5)
 
+    assert_soak_reconfigure_events(
+        Path(dump_dir) / "events",
+        num_successful_injections=injector.num_successful_injections,
+    )
+
     print(f"Random failure gsm8k accuracy test PASSED (seed={seed}, rollouts={num_rollout})")