diff --git a/docs/diffusion/advanced/deterministic.md b/docs/diffusion/advanced/deterministic.md index 82eb3a2664a..6f0f8222013 100644 --- a/docs/diffusion/advanced/deterministic.md +++ b/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 diff --git a/docs/examples/ppo.md b/docs/examples/ppo.md index 6b26c0542f9..9ead5d0e43c 100644 --- a/docs/examples/ppo.md +++ b/docs/examples/ppo.md @@ -75,8 +75,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 diff --git a/examples/ppo/README.md b/examples/ppo/README.md index b44591806a7..1e3a9ab1816 100644 --- a/examples/ppo/README.md +++ b/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 diff --git a/miles/backends/megatron_utils/ft/indep_dp.py b/miles/backends/megatron_utils/ft/indep_dp.py index 834a92c6342..33468cff767 100644 --- a/miles/backends/megatron_utils/ft/indep_dp.py +++ b/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, diff --git a/miles/ray/actor_group.py b/miles/ray/actor_group.py deleted file mode 100644 index 3a23ef03075..00000000000 --- a/miles/ray/actor_group.py +++ /dev/null @@ -1,143 +0,0 @@ -# FROZEN: v1 RayTrainGroup is the non-FT default path. Only critical bugfixes -# go here; new features land in miles/ray/train/group.py (v2). Dispatch between -# v1 and v2 happens in miles/ray/placement_group.py based on the env var -# MILES_EXPERIMENTAL_FT_TRAINER (default off -> v1). - -import asyncio - -from ray.util.placement_group import PlacementGroup - -from miles.ray.rollout.inference_controller import update_weights_window -from miles.ray.train.actor_factory import allocate_gpus_for_actor -from miles.utils.ft_utils.indep_dp import IndepDPInfo - - -class RayTrainGroup: - """ - A group of ray actors - - Args: - args (Namespace): Arguments for the actor group. - num_nodes (int): Number of nodes for this actor group. - num_gpus_per_node (int): Number of gpus for this actor group. - pg (PlacementGroup, optional): Placement group to schedule actor on. - If none, create new placement group automatically. Defaults to None. - num_gpus_per_actor (float, optional): Number of gpus allocated for each actor. - If < 1.0, multiple models can share same gpu. Defaults to 1. - """ - - def __init__( - self, - args, - num_nodes, - num_gpus_per_node, - pg: tuple[PlacementGroup, list[int], list[int]], - *, - inference_controller: object | None, - rollout_executor: object | None, - num_gpus_per_actor: float = 1, - role: str, - with_ref: bool, - with_opd_teacher: bool = False, - ) -> None: - self.args = args - self._num_nodes = num_nodes - self._num_gpus_per_node = num_gpus_per_node - self.role = role - self.with_ref = with_ref - self._inference_controller = inference_controller - self._rollout_executor = rollout_executor - self.with_opd_teacher = with_opd_teacher - - # Allocate the GPUs for actors w/o instantiating them - self._actor_handles = self._allocate_gpus_for_actor(pg, num_gpus_per_actor) - - def _allocate_gpus_for_actor(self, pg, num_gpus_per_actor): - return allocate_gpus_for_actor( - args=self.args, - gpus_per_cell=self._num_nodes * self._num_gpus_per_node, - pg=pg, - num_gpus_per_actor=num_gpus_per_actor, - indep_dp_store_addr=None, - role=self.role, - cell_index=0, - ) - - async def init(self): - """ - Allocate GPU resourced and initialize model, optimizer, local ckpt, etc. - """ - return await self._broadcast( - "init", - args=self.args, - role=self.role, - with_ref=self.with_ref, - with_opd_teacher=self.with_opd_teacher, - indep_dp_info=IndepDPInfo.create_trivial(), - ) - - async def train(self, rollout_id, rollout_data_pack, external_data=None): - """Do one rollout training""" - external_data_kwargs = self._compute_external_data_kwargs(external_data) - return await self._broadcast_per_worker( - "train", - compute_kwargs=lambda worker_index: dict( - rollout_id=rollout_id, - rollout_data_ref=rollout_data_pack["data_ref"], - witness_info=None, - attempt=0, - **external_data_kwargs[worker_index], - ), - ) - - def _compute_external_data_kwargs(self, external_data) -> list[dict]: - if external_data is None: - return [{} for _ in self._actor_handles] - if not isinstance(external_data, list): - return [dict(external_data=external_data) for _ in self._actor_handles] - if len(external_data) != len(self._actor_handles): - raise ValueError("external_data must contain one payload per train worker") - return [dict(external_data=payload) for payload in external_data] - - async def save_model(self, rollout_id, force_sync=False): - """Save actor model""" - await self._broadcast("save_model", rollout_id=rollout_id, force_sync=force_sync) - - async def export_hf(self, rollout_id: int, path: str): - """Export current weights as an HF checkpoint (collective across all ranks).""" - await self._broadcast("export_hf", rollout_id, path) - - async def update_weights(self, rollout_id: int | None = None): - """Broadcast weights from rank 0 to all other ranks.""" - if self.args.debug_train_only or self.args.debug_rollout_only: - return - - async with update_weights_window(self._inference_controller) as info: - await self._broadcast("update_weights", info=info) - - async def reconcile_adapters(self) -> None: - """Multi-LoRA: reconcile loaded adapters with the controller's active set - (load new, cleanup gone). Called by the trainer before generate.""" - await self._broadcast("reconcile_adapters") - - async def onload(self): - await self._broadcast("wake_up") - - async def offload(self): - await self._broadcast("sleep") - - async def clear_memory(self): - await self._broadcast("clear_memory") - - async def set_rollout_executor(self): - await self._broadcast("set_rollout_executor", rollout_executor=self._rollout_executor) - - async def _broadcast(self, method_name: str, **kwargs) -> list: - return await self._broadcast_per_worker(method_name, compute_kwargs=lambda _: kwargs) - - async def _broadcast_per_worker(self, method_name: str, *, compute_kwargs) -> list: - refs = [ - getattr(actor, method_name).remote(**compute_kwargs(worker_index)) - for worker_index, actor in enumerate(self._actor_handles) - ] - return await asyncio.gather(*refs) diff --git a/miles/ray/placement_group.py b/miles/ray/placement_group.py index 347db943c3e..5fd19dfd73d 100644 --- a/miles/ray/placement_group.py +++ b/miles/ray/placement_group.py @@ -7,7 +7,7 @@ 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 @@ -15,14 +15,6 @@ 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, diff --git a/miles/ray/train/group.py b/miles/ray/train/group.py index 58d4e63e9ce..52ba43bfa51 100644 --- a/miles/ray/train/group.py +++ b/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) diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index ad347991252..25deae04c47 100644 --- a/miles/utils/arguments.py +++ b/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 " diff --git a/miles/utils/dumper_utils.py b/miles/utils/dumper_utils.py index 5cb827115c0..7843b2fb8b1 100644 --- a/miles/utils/dumper_utils.py +++ b/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; diff --git a/miles/utils/environ.py b/miles/utils/environ.py index bf20a39e424..f5af64416a8 100644 --- a/miles/utils/environ.py +++ b/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 diff --git a/tests/e2e/ft/conftest_ft/execution.py b/tests/e2e/ft/conftest_ft/execution.py index a4ef09f79fe..23b78658caf 100644 --- a/tests/e2e/ft/conftest_ft/execution.py +++ b/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 diff --git a/tests/e2e/ft/conftest_ft/scenario_realistic_gsm8k.py b/tests/e2e/ft/conftest_ft/scenario_realistic_gsm8k.py index da2b1d19841..b3997c9459d 100644 --- a/tests/e2e/ft/conftest_ft/scenario_realistic_gsm8k.py +++ b/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})") diff --git a/tests/fast/backends/megatron_utils/test_indep_dp.py b/tests/fast/backends/megatron_utils/test_indep_dp.py new file mode 100644 index 00000000000..ecad313c222 --- /dev/null +++ b/tests/fast/backends/megatron_utils/test_indep_dp.py @@ -0,0 +1,185 @@ +import logging +import sys +import types +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest +import torch + +from miles.backends.megatron_utils.ft import indep_dp +from miles.utils.ft_utils.indep_dp import IndepDPInfo + +_LOGGER_NAME = "miles.backends.megatron_utils.ft.indep_dp" + + +def _indep_dp_messages(caplog) -> list[str]: + return [record.getMessage() for record in caplog.records if record.name == _LOGGER_NAME] + + +class FakeTorchftProcessGroup: + def __init__(self, timeout=None) -> None: + self._replica_id = "" + self._rank = 0 + self.configure_kwargs: dict | None = None + + def configure(self, **kwargs) -> None: + self.configure_kwargs = kwargs + self._replica_id = kwargs["replica_id"] + self._rank = kwargs["rank"] + + def size(self) -> int: + assert self.configure_kwargs is not None + return self.configure_kwargs["world_size"] + + def shutdown(self) -> None: + pass + + +class FakeCrossCellPGUtil: + def __init__(self, *, all_reduce_error: Exception | None = None) -> None: + self.all_reduce_error = all_reduce_error + self.reduced_tensors: list[torch.Tensor] = [] + + def all_reduce(self, tensor: torch.Tensor, group, op) -> None: + if self.all_reduce_error is not None: + raise self.all_reduce_error + self.reduced_tensors.append(tensor) + + +def _make_model_chunk() -> SimpleNamespace: + bucket = SimpleNamespace(grad_data=torch.full((4,), 2.5)) + return SimpleNamespace(bucket_groups=[SimpleNamespace(buckets=[bucket])], expert_parallel_bucket_groups=[]) + + +def _make_parallel_state(pg) -> SimpleNamespace: + return SimpleNamespace( + intra_dp=SimpleNamespace(size=1), + indep_dp=SimpleNamespace(rank=1, size=3, group=pg, debug_info={"quorum": 7}), + ) + + +class TestCreateIndepDpGroup: + @pytest.fixture() + def fake_torchft(self, monkeypatch): + module = types.ModuleType("torchft.process_group") + module.ProcessGroupNCCL = FakeTorchftProcessGroup + module.ProcessGroupGloo = FakeTorchftProcessGroup + monkeypatch.setitem(sys.modules, "torchft", types.ModuleType("torchft")) + monkeypatch.setitem(sys.modules, "torchft.process_group", module) + return module + + def test_creating_a_cross_cell_group_emits_an_ft_tagged_create_record(self, fake_torchft, caplog) -> None: + """The create_pg record stays discoverable by the ft structured-log tag.""" + info = IndepDPInfo( + cell_index=1, num_cells=3, alive_rank=1, alive_size=2, quorum_id=7, alive_cell_indices=[0, 1] + ) + + with caplog.at_level(logging.INFO, logger=_LOGGER_NAME): + group_info = indep_dp.create_indep_dp_group( + store_addr="tcp://store:1234", + indep_dp_info=info, + megatron_rank=5, + megatron_world_size=8, + ) + + messages = _indep_dp_messages(caplog) + assert group_info.rank == 1 + assert len(messages) == 1 + assert messages[0].startswith("ft ") + assert "op=create_pg" in messages[0] + assert "quorum=7" in messages[0] + + +class TestReconfigureIndepDpGroup: + def test_reconfigure_emits_ft_tagged_start_and_end_records(self, caplog) -> None: + """Both reconfigure records stay discoverable by the ft structured-log tag.""" + old_group = MagicMock() + parallel_state = SimpleNamespace(indep_dp=SimpleNamespace(group=old_group, gloo_group=None)) + info = IndepDPInfo(cell_index=2, num_cells=3, alive_rank=0, alive_size=1, quorum_id=5, alive_cell_indices=[2]) + + with caplog.at_level(logging.INFO, logger=_LOGGER_NAME): + indep_dp.reconfigure_indep_dp_group( + parallel_state=parallel_state, + store_addr="tcp://store:1234", + indep_dp_info=info, + megatron_rank=5, + megatron_world_size=8, + ) + + messages = _indep_dp_messages(caplog) + old_group.shutdown.assert_called_once() + assert all(message.startswith("ft ") for message in messages) + assert "op=reconfig phase=start" in messages[0] + assert "quorum_to=5" in messages[0] + assert "op=reconfig phase=end" in messages[1] + assert "quorum=5" in messages[1] + + +class TestAllreduceGradsAndLossesAcrossReplicas: + @pytest.fixture() + def megatron_env(self): + with ( + patch.object(indep_dp, "mpu") as mock_mpu, + patch.object(indep_dp, "get_gloo_group", return_value=None), + patch.object(indep_dp, "collective_bool_and", side_effect=lambda value, group: value), + ): + mock_mpu.is_pipeline_last_stage.return_value = False + yield mock_mpu + + @staticmethod + def _run(pg, util) -> tuple[bool, dict[str, float]]: + args = SimpleNamespace(calculate_per_token_loss=False) + with patch.object(indep_dp.GeneralPGUtil, "create", return_value=util): + return indep_dp.allreduce_grads_and_losses_across_replicas( + args, [_make_model_chunk()], _make_parallel_state(pg), losses_reduced=[] + ) + + def test_a_successful_allreduce_emits_ft_tagged_start_and_end_records(self, megatron_env, caplog) -> None: + """The happy-path cross-cell records stay discoverable by the ft structured-log tag.""" + pg = SimpleNamespace(errored=lambda: None) + util = FakeCrossCellPGUtil() + + with caplog.at_level(logging.INFO, logger=_LOGGER_NAME): + consensus, loss_reduced = self._run(pg, util) + + messages = _indep_dp_messages(caplog) + assert consensus is True + assert loss_reduced == {} + assert len(util.reduced_tensors) == 1 + assert all(message.startswith("ft ") for message in messages) + assert "op=cross_cell phase=start kind=grad_allreduce" in messages[0] + assert "op=cross_cell phase=end kind=grad_allreduce" in messages[1] + assert "this_rank_ok=true consensus_ok=true" in messages[1] + + def test_a_raising_allreduce_emits_an_ft_tagged_fail_record(self, megatron_env, caplog) -> None: + """A synchronous collective failure is reported as an ft-tagged fail record and discards the step.""" + pg = SimpleNamespace(errored=lambda: None) + util = FakeCrossCellPGUtil(all_reduce_error=RuntimeError("NCCL communicator was aborted on rank 2")) + + with caplog.at_level(logging.INFO, logger=_LOGGER_NAME): + consensus, _loss_reduced = self._run(pg, util) + + messages = _indep_dp_messages(caplog) + assert consensus is False + fail_messages = [message for message in messages if "phase=fail" in message] + assert len(fail_messages) == 1 + assert fail_messages[0].startswith("ft ") + assert "kind=grad_allreduce" in fail_messages[0] + assert "this_rank_ok=false consensus_ok=false" in messages[-1] + + def test_an_asynchronously_errored_group_emits_an_ft_tagged_async_error_record(self, megatron_env, caplog) -> None: + """An error surfacing only via pg.errored() is reported as an ft-tagged async_error record.""" + pg = SimpleNamespace(errored=lambda: RuntimeError("peer 2 left the quorum")) + util = FakeCrossCellPGUtil() + + with caplog.at_level(logging.INFO, logger=_LOGGER_NAME): + consensus, _loss_reduced = self._run(pg, util) + + messages = _indep_dp_messages(caplog) + assert consensus is False + async_messages = [message for message in messages if "phase=async_error" in message] + assert len(async_messages) == 1 + assert async_messages[0].startswith("ft ") + assert "kind=grad_allreduce" in async_messages[0] + assert "peer 2 left the quorum" in async_messages[0] diff --git a/tests/fast/e2e/ft/test_execution.py b/tests/fast/e2e/ft/test_execution.py new file mode 100644 index 00000000000..44791eda97b --- /dev/null +++ b/tests/fast/e2e/ft/test_execution.py @@ -0,0 +1,46 @@ +import dataclasses +from pathlib import Path + +from tests.e2e.ft.conftest_ft.execution import get_common_train_args, get_ft_args +from tests.e2e.ft.conftest_ft.modes import MODES + + +class TestGetCommonTrainArgs: + def test_a_colocated_real_rollout_mode_emits_the_colocate_flag(self, tmp_path: Path) -> None: + """A colocated mode must tell the trainer to share its gpus with the rollout engines.""" + args = get_common_train_args(MODES["colocate_dp2_cp2_rollout_ft"], dump_dir=str(tmp_path)) + + assert "--colocate " in args + + def test_a_disaggregated_real_rollout_mode_omits_the_colocate_flag(self, tmp_path: Path) -> None: + """Rollout engines on their own gpus must not be colocated with the trainer.""" + args = get_common_train_args(MODES["dp2_cp2_real_rollout"], dump_dir=str(tmp_path)) + + assert "--rollout-num-gpus" in args + assert "--colocate" not in args + + def test_a_debug_rollout_mode_omits_the_colocate_flag_even_when_the_mode_is_colocated( + self, tmp_path: Path + ) -> None: + """Without real rollout engines there is nothing to colocate, whatever the mode declares.""" + mode = dataclasses.replace(MODES["colocate_dp2_cp2_rollout_ft"], rollout_num_engines=0) + + args = get_common_train_args(mode, dump_dir=str(tmp_path)) + + assert mode.colocate is True + assert "--debug-train-only" in args + assert "--colocate" not in args + + +class TestGetFtArgs: + def test_a_rollout_only_ft_mode_propagates_the_rollout_component_and_api_server_port(self) -> None: + """Rollout-only fault tolerance must not silently enable trainer fault tolerance.""" + args = get_ft_args(MODES["colocate_dp2_cp2_rollout_ft"]) + + assert args == "--use-fault-tolerance --ft-components rollout --api-server-port 0 " + + def test_a_trainer_ft_mode_propagates_the_train_component(self) -> None: + """The trainer-fault-tolerance modes keep sending the train component.""" + args = get_ft_args(MODES["dp2_cp2_real_rollout"]) + + assert args == "--use-fault-tolerance --ft-components train --api-server-port 0 " diff --git a/tests/fast/ray/test_actor_group_shared_ppo.py b/tests/fast/ray/test_actor_group_shared_ppo.py deleted file mode 100644 index 259aec42066..00000000000 --- a/tests/fast/ray/test_actor_group_shared_ppo.py +++ /dev/null @@ -1,62 +0,0 @@ -class _RemoteTrain: - def __init__(self, rank, calls): - self.rank = rank - self.calls = calls - - def remote(self, rollout_id, rollout_data_ref, **kwargs): - self.calls.append((self.rank, rollout_id, rollout_data_ref, kwargs)) - - async def result(): - return {"rank": self.rank} - - return result() - - -class _Handle: - def __init__(self, rank, calls): - self.train = _RemoteTrain(rank, calls) - - -async def test_train_routes_each_critic_payload_to_matching_actor_rank(): - from miles.ray.actor_group import RayTrainGroup - - calls = [] - group = object.__new__(RayTrainGroup) - group._actor_handles = [_Handle(0, calls), _Handle(1, calls)] - payloads = [{"values": ["v0"]}, {"values": ["v1"]}] - - result = await group.train(5, {"data_ref": "rollout"}, external_data=payloads) - - assert result == [{"rank": 0}, {"rank": 1}] - assert calls == [ - (0, 5, "rollout", {"witness_info": None, "attempt": 0, "external_data": payloads[0]}), - (1, 5, "rollout", {"witness_info": None, "attempt": 0, "external_data": payloads[1]}), - ] - - -async def test_train_broadcasts_without_lifecycle_options(): - from miles.ray.actor_group import RayTrainGroup - - calls = [] - group = object.__new__(RayTrainGroup) - group._actor_handles = [_Handle(0, calls), _Handle(1, calls)] - - await group.train(7, {"data_ref": "rollout"}) - - assert calls == [ - (0, 7, "rollout", {"witness_info": None, "attempt": 0}), - (1, 7, "rollout", {"witness_info": None, "attempt": 0}), - ] - - -async def test_train_rejects_wrong_number_of_rank_payloads(): - import pytest - - from miles.ray.actor_group import RayTrainGroup - - group = object.__new__(RayTrainGroup) - group._actor_handles = [_Handle(0, []), _Handle(1, [])] - - with pytest.raises(ValueError, match="one payload per train worker"): - await group.train(5, {"data_ref": "rollout"}, external_data=[{"values": []}]) - diff --git a/tests/fast/ray/test_update_weights_ordering.py b/tests/fast/ray/test_update_weights_ordering.py index 320b83010f3..65416b61fe9 100644 --- a/tests/fast/ray/test_update_weights_ordering.py +++ b/tests/fast/ray/test_update_weights_ordering.py @@ -4,7 +4,6 @@ import pytest from tests.fast.ray.rollout.conftest import make_args -from miles.ray.actor_group import RayTrainGroup from miles.ray.rollout.inference_controller import InferenceController from miles.utils.context_lock import ContextLock @@ -113,19 +112,7 @@ def _record_snapshot() -> None: assert init_counts_at_snapshot == [1] -def _make_v1_group(order: list[str]) -> RayTrainGroup: - group = RayTrainGroup.__new__(RayTrainGroup) - group.args = Namespace(debug_train_only=False, debug_rollout_only=False, use_fault_tolerance=False) - group._inference_controller = _OrderRecordingInferenceController(order) - - async def _record_broadcast(*args: object, **kwargs: object) -> None: - order.append("broadcast") - - group._broadcast = AsyncMock(side_effect=_record_broadcast) - return group - - -def _make_v2_group(order: list[str]): +def _make_controller(order: list[str]): from miles.ray.train.group import TrainerController as FaultTolerantTrainGroup group = TrainerController.__new__(TrainerController) @@ -141,22 +128,10 @@ async def _record_execute_first_alive(*args: object, **kwargs: object) -> None: @pytest.mark.asyncio -async def test_v1_brackets_the_broadcast_with_start_and_end_update_weights(): - """The trainer broadcast is recorded strictly between the start and end of the update window.""" - order: list[str] = [] - group = _make_v1_group(order) - - await group.update_weights() - - assert order == ["start_update_weights", "broadcast", "end_update_weights"] - group._broadcast.assert_awaited_once() - - -@pytest.mark.asyncio -async def test_v2_brackets_the_broadcast_with_start_and_end_update_weights(): +async def test_the_trainer_brackets_the_broadcast_with_start_and_end_update_weights(): """The fault-tolerant trainer runs the actual update RPC strictly inside the update window.""" order: list[str] = [] - group = _make_v2_group(order) + group = _make_controller(order) await group.update_weights() @@ -165,65 +140,16 @@ async def test_v2_brackets_the_broadcast_with_start_and_end_update_weights(): @pytest.mark.asyncio -async def test_v1_hands_end_update_weights_the_snapshot_start_returned(): - """A dropped or substituted snapshot leaves every pending cell unregistered with the router.""" - order: list[str] = [] - group = _make_v1_group(order) - - await group.update_weights() - - _assert_the_snapshot_is_handed_back_unchanged(group._inference_controller) - - -@pytest.mark.asyncio -async def test_v2_hands_end_update_weights_the_snapshot_start_returned(): - """Same snapshot pass-through requirement on the fault-tolerant trainer group.""" +async def test_the_trainer_hands_end_update_weights_the_snapshot_start_returned(): + """The snapshot start_update_weights returned is handed back to end_update_weights unchanged.""" order: list[str] = [] - group = _make_v2_group(order) + group = _make_controller(order) await group.update_weights() _assert_the_snapshot_is_handed_back_unchanged(group._inference_controller) -@pytest.mark.asyncio -async def test_v1_aborts_the_window_when_the_broadcast_raises(): - """A failed weight transfer must close the lock window instead of leaving it open forever.""" - order: list[str] = [] - group = RayTrainGroup.__new__(RayTrainGroup) - group.args = Namespace(debug_train_only=False, debug_rollout_only=False, use_fault_tolerance=False) - group._inference_controller = _OrderRecordingInferenceController(order) - group._broadcast = AsyncMock(side_effect=RuntimeError("weight transfer died")) - - with pytest.raises(RuntimeError, match="weight transfer died"): - await group.update_weights() - - assert order == ["start_update_weights", "abort_update_weights"] - - -@pytest.mark.asyncio -async def test_v2_aborts_the_window_when_the_broadcast_raises(monkeypatch): - """Same abort requirement on the fault-tolerant trainer group, after its retries are exhausted.""" - from miles.ray.train import group as train_group_module - - async def _retry_once(fn, **kwargs): - await fn(0) - - monkeypatch.setattr(train_group_module, "retry", _retry_once) - - order: list[str] = [] - group = train_group_module.RayTrainGroup.__new__(train_group_module.RayTrainGroup) - group.args = Namespace(debug_train_only=False, debug_rollout_only=False) - group._inference_controller = _OrderRecordingInferenceController(order) - group._execute_first_alive = AsyncMock(side_effect=RuntimeError("weight transfer died")) - group._maybe_log_inference_engine_weight_checksums = AsyncMock() - - with pytest.raises(RuntimeError, match="weight transfer died"): - await group.update_weights() - - assert order == ["start_update_weights", "abort_update_weights"] - - def test_fsdp_updater_flushes_only_after_every_engine_is_paused(): """Every engine is paused before any engine is flushed.""" from unittest.mock import patch diff --git a/tests/fast/utils/test_arguments.py b/tests/fast/utils/test_arguments.py index 85257668f7f..6605223db24 100644 --- a/tests/fast/utils/test_arguments.py +++ b/tests/fast/utils/test_arguments.py @@ -624,16 +624,17 @@ def test_bridge_mode_accepts_critic(tmp_path): assert args.use_critic is True -def test_critic_rejects_experimental_ft_trainer(tmp_path, monkeypatch): - monkeypatch.setenv("MILES_EXPERIMENTAL_FT_TRAINER", "1") +def test_critic_is_accepted_on_the_only_trainer(tmp_path): + """Shared actor/critic PPO used to be rejected on the cell based trainer, which is now the only one.""" parser = argparse.ArgumentParser() get_miles_extra_args_provider()(parser) args = parser.parse_args( ["--advantage-estimator", "ppo", "--hf-checkpoint", str(tmp_path), "--num-rollout", "1"] + REQUIRED_ARGS ) - with pytest.raises(AssertionError, match="MILES_EXPERIMENTAL_FT_TRAINER"): - miles_validate_args(args) + miles_validate_args(args) + + assert args.use_critic is True def test_critic_rejects_reward_level_kl(tmp_path): @@ -730,9 +731,8 @@ def test_rejects_non_adam_optimizer(self): with pytest.raises(AssertionError, match="requires --optimizer adam"): miles_validate_args(args) - def test_accepts_experimental_ft_trainer(self, monkeypatch): - """The v2 train group implements reconcile_adapters, so multi-LoRA may use it.""" - monkeypatch.setenv("MILES_EXPERIMENTAL_FT_TRAINER", "1") + def test_is_accepted_on_the_only_trainer(self): + """Multi-LoRA used to be rejected on the cell based trainer, which is now the only one.""" args = self._parse([]) miles_validate_args(args) diff --git a/tests/fast/utils/test_dumper_utils.py b/tests/fast/utils/test_dumper_utils.py index ea7a6ae60ab..bf919d3a0f9 100644 --- a/tests/fast/utils/test_dumper_utils.py +++ b/tests/fast/utils/test_dumper_utils.py @@ -215,7 +215,84 @@ def test_finalize_preserves_activations_and_pins_model_dumps_to_step_zero( assert torch.load(dump_file, weights_only=False)["meta"]["step"] == 0 +def _make_distributed_optimizer_bucket_group(param: torch.nn.Parameter, grad_data: torch.Tensor) -> SimpleNamespace: + bucket = SimpleNamespace(grad_data=grad_data, param_to_index={param: (0, grad_data.numel())}) + return SimpleNamespace( + ddp_config=SimpleNamespace(use_distributed_optimizer=True), + intra_distributed_optimizer_instance_group=None, + intra_distributed_optimizer_instance_size=1, + intra_distributed_optimizer_instance_rank=0, + buckets=[bucket], + ) + + +class TestDumperMegatronUtilFinalize: + def test_finalize_dumps_distributed_optimizer_gradients_without_param_grad(self, tmp_path: Path) -> None: + """A distributed-optimizer param carries its gradient only in the bucket buffer, and that must be dumped.""" + args = _make_args(tmp_path) + args.dumper_fwd_bwd = ["enable_model_grad=true"] + model = torch.nn.Linear(2, 1, bias=False) + model.bucket_groups = [_make_distributed_optimizer_bucket_group(model.weight, torch.full((2,), 3.5))] + state = SimpleNamespace( + effective_dp=SimpleNamespace(rank=0), + indep_dp=SimpleNamespace(rank=0, group=None), + ) + + with ( + patch("miles.utils.dumper_utils.get_parallel_state", return_value=state), + patch("miles.utils.dumper_utils.dist") as mock_dist, + ): + mock_dist.is_initialized.return_value = False + util = DumperMegatronUtil(args, [model], DumperPhase.FWD_BWD, rollout_id=1) + assert model.weight.grad is None + try: + util.finalize([model]) + finally: + dumper_utils.dumper.reset() + dumper_utils.dumper.configure(enable=False) + + dump_files = sorted((tmp_path / "fwd_bwd" / "rollout_1").glob("*.pt")) + assert len(dump_files) == 1 + dumped = torch.load(dump_files[0], weights_only=False) + assert dumped["meta"]["name"] == "grad__param__weight" + assert torch.equal(dumped["value"], torch.full((1, 2), 3.5)) + + class TestBarrierAfterDumpDirCleanup: + @staticmethod + def _run(group: MagicMock) -> None: + state = SimpleNamespace(indep_dp=SimpleNamespace(rank=1, size=2, group=group, debug_info={"quorum": 7})) + with ( + patch("miles.utils.dumper_utils.get_parallel_state", return_value=state), + patch("miles.utils.dumper_utils.dist") as mock_dist, + ): + mock_dist.is_initialized.return_value = False + dumper_utils._barrier_after_dump_dir_cleanup() + + def test_successful_cross_cell_barrier_records_keep_the_ft_tag(self, caplog) -> None: + """The start and success records stay discoverable by the external ft structured-log tag.""" + with caplog.at_level(logging.INFO, logger="miles.utils.dumper_utils"): + self._run(MagicMock()) + + messages = [record.getMessage() for record in caplog.records if record.name == "miles.utils.dumper_utils"] + assert all(message.startswith("ft ") for message in messages) + assert "op=cross_cell phase=start kind=dump_barrier" in messages[0] + assert "op=cross_cell phase=end kind=dump_barrier" in messages[1] + assert "quorum=7 success=true" in messages[1] + + def test_degraded_cross_cell_barrier_record_keeps_the_ft_tag(self, caplog) -> None: + """The degraded end record stays discoverable by the external ft structured-log tag.""" + group = MagicMock() + group.barrier.side_effect = RuntimeError("NCCL communicator was aborted on rank 1") + + with caplog.at_level(logging.INFO, logger="miles.utils.dumper_utils"): + self._run(group) + + messages = [record.getMessage() for record in caplog.records if record.name == "miles.utils.dumper_utils"] + assert all(message.startswith("ft ") for message in messages) + assert "op=cross_cell phase=end kind=dump_barrier" in messages[1] + assert "success=false degraded=true" in messages[1] + def test_cross_cell_barrier_abort_does_not_raise(self) -> None: """A peer death aborts the cross-cell PG mid-barrier; the survivor continues instead of erroring.""" group = MagicMock() @@ -335,6 +412,58 @@ async def _post(url, body): assert posted == ["http://a:1/dumper/configure", "http://b:2/dumper/configure"] + @pytest.mark.asyncio + async def test_surplus_registered_engines_are_accepted_and_configured(self, tmp_path: Path) -> None: + """More registered engines than expected is a complete roster, and every one of them is configured.""" + posted: list[str] = [] + + async def _post(url, body): + posted.append(url) + + get_worker_urls = AsyncMock(return_value=["http://a:1", "http://b:2", "http://c:3"]) + + with ( + patch( + "miles.rollout.inference_rollout.inference_rollout_train.get_worker_urls", + new=get_worker_urls, + ), + patch("miles.utils.dumper_utils.resolve_sglang_config", return_value=self.resolved_config), + patch("miles.utils.http_utils.post", new=_post), + patch("miles.utils.dumper_utils._cleanup_dump_dir"), + ): + await dumper_utils.configure_sglang(self._make_args(tmp_path)) + + assert get_worker_urls.await_count == 1 + assert posted == [ + "http://a:1/dumper/configure", + "http://b:2/dumper/configure", + "http://c:3/dumper/configure", + ] + + @pytest.mark.asyncio + async def test_configuring_engines_does_not_enter_a_process_group_barrier(self, tmp_path: Path) -> None: + """A cell recovering on its own must configure its engines without waiting for the other cells.""" + posted: list[str] = [] + + async def _post(url, body): + posted.append(url) + + with ( + patch( + "miles.rollout.inference_rollout.inference_rollout_train.get_worker_urls", + new=AsyncMock(return_value=["http://a:1", "http://b:2"]), + ), + patch("miles.utils.dumper_utils.resolve_sglang_config", return_value=self.resolved_config), + patch("miles.utils.http_utils.post", new=_post), + patch("miles.utils.dumper_utils._cleanup_dump_dir"), + patch("miles.utils.dumper_utils.dist") as mock_dist, + ): + mock_dist.is_initialized.return_value = True + await dumper_utils.configure_sglang(self._make_args(tmp_path)) + + assert len(posted) == 2 + mock_dist.barrier.assert_not_called() + @pytest.mark.asyncio async def test_an_incomplete_router_roster_is_waited_out(self, tmp_path: Path) -> None: """Configuration waits until every expected engine has registered."""