From 96eecb4acae9e76df9552913eab4f87c89bf0290 Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Mon, 22 Jun 2026 18:15:16 +0800 Subject: [PATCH 01/41] Add a deterministic_random reward type Add a `deterministic_random` reward that hashes the sample tokens + response to produce a stable pseudo-random 0/1 reward, used for reproducible fault-tolerance / CI tests. - rm_hub/__init__.py (+ test). --- miles/rollout/rm_hub/__init__.py | 5 +++++ tests/fast/rollout/rm_hub/test_rm_hub.py | 25 ++++++++++++++++++++++++ 2 files changed, 30 insertions(+) diff --git a/miles/rollout/rm_hub/__init__.py b/miles/rollout/rm_hub/__init__.py index 6264e17c304..962c4217558 100644 --- a/miles/rollout/rm_hub/__init__.py +++ b/miles/rollout/rm_hub/__init__.py @@ -1,4 +1,5 @@ import asyncio +import hashlib import random import aiohttp @@ -62,6 +63,10 @@ async def async_rm(args, sample: Sample, **kwargs): return compute_ifbench_reward(response, label, metadata=metadata) elif rm_type == "random": return random.randint(0, 1) + elif rm_type == "deterministic_random": + content = str(sample.tokens) + response + content_hash = hashlib.sha256(content.encode()).digest() + return int(content_hash[0]) % 2 elif rm_type: raise NotImplementedError(f"Rule-based RM for {rm_type} is not implemented.") else: diff --git a/tests/fast/rollout/rm_hub/test_rm_hub.py b/tests/fast/rollout/rm_hub/test_rm_hub.py index a3dadbdaf00..eb97afa8cc5 100644 --- a/tests/fast/rollout/rm_hub/test_rm_hub.py +++ b/tests/fast/rollout/rm_hub/test_rm_hub.py @@ -51,6 +51,31 @@ def test_random_rm(self, mock_args): reward = run(async_rm(mock_args, sample)) assert reward in [0, 1] + def test_deterministic_random_rm_returns_binary(self, mock_args): + mock_args.rm_type = "deterministic_random" + sample = Sample(prompt="", response="hello", label="", tokens=[1, 2, 3]) + reward = run(async_rm(mock_args, sample)) + assert reward in [0, 1] + + def test_deterministic_random_rm_is_deterministic(self, mock_args): + mock_args.rm_type = "deterministic_random" + sample = Sample(prompt="", response="hello world", label="", tokens=[10, 20]) + rewards = [run(async_rm(mock_args, sample)) for _ in range(5)] + assert len(set(rewards)) == 1 + + def test_deterministic_random_rm_differs_by_response(self, mock_args): + mock_args.rm_type = "deterministic_random" + samples = [Sample(prompt="", response=f"response_{i}", label="", tokens=[1, 2, 3]) for i in range(20)] + rewards = [run(async_rm(mock_args, s)) for s in samples] + assert 0 in rewards and 1 in rewards + + def test_deterministic_random_rm_differs_by_tokens(self, mock_args): + """Same response with different tokens yields both reward values across many samples.""" + mock_args.rm_type = "deterministic_random" + samples = [Sample(prompt="", response="same", label="", tokens=[i, i + 1, i + 2]) for i in range(20)] + rewards = [run(async_rm(mock_args, s)) for s in samples] + assert 0 in rewards and 1 in rewards + def test_rm_type_from_metadata(self, mock_args): mock_args.rm_type = None sample = Sample(prompt="", response=r"\boxed{42}", label="42", metadata={"rm_type": "math"}) From 2c1d4a74049aac27508151b96434a6358796ac6d Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Mon, 22 Jun 2026 18:15:16 +0800 Subject: [PATCH 02/41] Add an inplace_modify_args context manager Add `inplace_modify_args`, a context manager that temporarily overrides args attributes and restores them on exit (asserting they weren't clobbered), used to scope per-attempt argument overrides in the fault-tolerant trainer. - argparse_utils.py (+ test). --- miles/utils/argparse_utils.py | 20 ++++++++++ tests/fast/utils/test_argparse_utils.py | 51 ++++++++++++++++++++++++- 2 files changed, 70 insertions(+), 1 deletion(-) diff --git a/miles/utils/argparse_utils.py b/miles/utils/argparse_utils.py index 13b8f2b5cb8..b3c7f8f2f14 100644 --- a/miles/utils/argparse_utils.py +++ b/miles/utils/argparse_utils.py @@ -9,6 +9,8 @@ import argparse import dataclasses import types +from collections.abc import Iterator +from contextlib import contextmanager from pathlib import Path from typing import Generic, TypeVar, get_type_hints @@ -129,3 +131,21 @@ def to_cli_args(self, instance: T) -> str: parts.append(f"{flag} {value}") return " ".join(parts) + + +@contextmanager +def inplace_modify_args(args: argparse.Namespace, overrides: dict[str, object]) -> Iterator[None]: + """Temporarily set attributes on ``args``, restoring the originals on exit.""" + old_values = {key: getattr(args, key) for key in overrides} + for key, value in overrides.items(): + setattr(args, key, value) + try: + yield + finally: + for key, old_value in old_values.items(): + current = getattr(args, key) + assert current == overrides[key], ( + f"args.{key} was modified inside the inplace_modify_args block " + f"(expected {overrides[key]!r}, found {current!r}); restoring would clobber it" + ) + setattr(args, key, old_value) diff --git a/tests/fast/utils/test_argparse_utils.py b/tests/fast/utils/test_argparse_utils.py index 4221e24c6c9..2fe6b8e995d 100644 --- a/tests/fast/utils/test_argparse_utils.py +++ b/tests/fast/utils/test_argparse_utils.py @@ -7,7 +7,7 @@ import pytest -from miles.utils.argparse_utils import DataclassArgparseBridge +from miles.utils.argparse_utils import DataclassArgparseBridge, inplace_modify_args @dataclasses.dataclass(frozen=True) @@ -262,3 +262,52 @@ class _Bad: with pytest.raises(TypeError, match="Unsupported field type"): bridge.register_on_parser(parser) + + +class TestInplaceModifyArgs: + def test_overrides_inside_and_restores_on_exit(self) -> None: + """Overridden attributes are visible inside the block and restored afterwards.""" + args = argparse.Namespace(no_load_optim=True, finetune=True, lr=1.0) + + with inplace_modify_args(args, dict(no_load_optim=False, finetune=False)): + assert args.no_load_optim is False + assert args.finetune is False + assert args.lr == 1.0 + + assert args.no_load_optim is True + assert args.finetune is True + + def test_restores_on_exception(self) -> None: + """Originals are restored even when the block raises.""" + args = argparse.Namespace(flag=True) + + with pytest.raises(RuntimeError): + with inplace_modify_args(args, dict(flag=False)): + raise RuntimeError("boom") + + assert args.flag is True + + def test_empty_overrides_is_noop(self) -> None: + """An empty override dict changes nothing.""" + args = argparse.Namespace(flag=True) + + with inplace_modify_args(args, {}): + assert args.flag is True + + assert args.flag is True + + def test_unknown_attribute_raises(self) -> None: + """Overriding an attribute the namespace does not have fails loudly.""" + args = argparse.Namespace() + + with pytest.raises(AttributeError): + with inplace_modify_args(args, dict(missing=1)): + pass + + def test_mutation_inside_block_fails_on_exit(self) -> None: + """An attribute mutated inside the block is detected instead of silently clobbered.""" + args = argparse.Namespace(flag=True) + + with pytest.raises(AssertionError, match="modified inside"): + with inplace_modify_args(args, dict(flag=False)): + args.flag = True From 7b00c7567c580b97b96bd95b7e6b1d1712473b72 Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Mon, 22 Jun 2026 18:15:16 +0800 Subject: [PATCH 03/41] Add fault-tolerance support tweaks to shared utilities Small shared-utility additions used by the fault-tolerant trainer: hash non-contiguous tensors safely (reshape before viewing as bytes), an `enable_experimental_ft_trainer` env flag, forward NCCL_DEBUG/NCCL_DEBUG_FILE to worker environments, and a `filter_keys` helper. - ci_utils.py / environ.py / external_utils/command_utils.py / misc.py. --- miles/backends/megatron_utils/ci_utils.py | 2 +- miles/utils/environ.py | 15 ++++++ miles/utils/external_utils/command_utils.py | 6 ++- miles/utils/misc.py | 13 +++++ .../backends/megatron_utils/test_ci_utils.py | 50 +++++++++++++++++++ tests/fast/utils/test_misc.py | 37 +++++++++++++- 6 files changed, 120 insertions(+), 3 deletions(-) create mode 100644 tests/fast/backends/megatron_utils/test_ci_utils.py diff --git a/miles/backends/megatron_utils/ci_utils.py b/miles/backends/megatron_utils/ci_utils.py index 3ecf15eab0f..5837358bb73 100644 --- a/miles/backends/megatron_utils/ci_utils.py +++ b/miles/backends/megatron_utils/ci_utils.py @@ -36,7 +36,7 @@ def _hash_tensor_bytes(tensor: torch.Tensor) -> bytes: data = data.cpu() if not data.is_contiguous(): data = data.contiguous() - return data.view(torch.uint8).numpy().tobytes() + return data.reshape(-1).view(torch.uint8).numpy().tobytes() def compute_model_hashes_by_layer(model: Sequence[DDP]) -> dict[str, str]: diff --git a/miles/utils/environ.py b/miles/utils/environ.py index 35d1f350eed..1da32045b78 100644 --- a/miles/utils/environ.py +++ b/miles/utils/environ.py @@ -12,3 +12,18 @@ def enable_experimental_rollout_refactor() -> bool: _printed_experimental_rollout_refactor = True return result + + +_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/miles/utils/external_utils/command_utils.py b/miles/utils/external_utils/command_utils.py index d016e01ac34..7489d89ec99 100644 --- a/miles/utils/external_utils/command_utils.py +++ b/miles/utils/external_utils/command_utils.py @@ -156,7 +156,11 @@ def execute_train( } ), "NCCL_NVLS_ENABLE": os.environ.get("NCCL_NVLS_ENABLE", str(int(check_has_nvlink()))), - **{k: os.environ[k] for k in ("NCCL_SOCKET_IFNAME", "GLOO_SOCKET_IFNAME") if k in os.environ}, + **{ + k: os.environ[k] + for k in ("NCCL_SOCKET_IFNAME", "GLOO_SOCKET_IFNAME", "NCCL_DEBUG", "NCCL_DEBUG_FILE") + if k in os.environ + }, "no_proxy": f"127.0.0.1,{master_addr}", # This is needed by megatron / torch distributed in multi-node setup "MASTER_ADDR": master_addr, diff --git a/miles/utils/misc.py b/miles/utils/misc.py index 313a60dda90..3c160cc5cbc 100644 --- a/miles/utils/misc.py +++ b/miles/utils/misc.py @@ -1,14 +1,19 @@ import asyncio import importlib +import logging import re import subprocess +from collections.abc import Sequence from contextlib import contextmanager +from typing import Any import ray from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy from miles.utils.http_utils import is_port_available +logger = logging.getLogger(__name__) + # Mainly used for test purpose where `load_function` needs to load many in-flight generated functions class FunctionRegistry: @@ -199,3 +204,11 @@ def should_run_periodic_action( async def as_completed_async(tasks): for coro in asyncio.as_completed(tasks): yield await coro + + +def filter_keys(d: dict[str, Any], interest_keys: Sequence[str]) -> dict[str, Any]: + try: + return {k: d[k] for k in interest_keys} + except Exception: + logger.error(f"filter_keys d.keys={list(d)} {interest_keys=}", exc_info=True) + raise diff --git a/tests/fast/backends/megatron_utils/test_ci_utils.py b/tests/fast/backends/megatron_utils/test_ci_utils.py new file mode 100644 index 00000000000..21eee55fb87 --- /dev/null +++ b/tests/fast/backends/megatron_utils/test_ci_utils.py @@ -0,0 +1,50 @@ +import pytest + +pytest.importorskip("megatron.core.distributed") + +import torch + +from miles.backends.megatron_utils.ci_utils import _hash_tensor_bytes + + +def test_hash_tensor_bytes_contiguous_float32_returns_raw_buffer_bytes() -> None: + """A contiguous 2D float32 tensor hashes to its exact raw little-endian buffer bytes.""" + tensor = torch.arange(6, dtype=torch.float32).reshape(2, 3) + result = _hash_tensor_bytes(tensor) + assert isinstance(result, bytes) + assert len(result) == 6 * 4 + assert result == tensor.reshape(-1).contiguous().numpy().tobytes() + + +def test_hash_tensor_bytes_int64_byte_length_is_eight_per_element() -> None: + """An int64 tensor yields 8 bytes per element with bit-exact buffer contents.""" + tensor = torch.arange(6, dtype=torch.int64).reshape(2, 3) + result = _hash_tensor_bytes(tensor) + assert len(result) == 6 * 8 + assert result == tensor.reshape(-1).contiguous().numpy().tobytes() + + +def test_hash_tensor_bytes_noncontiguous_transpose_does_not_raise() -> None: + """A non-contiguous transposed view is hashed without raising (regression guard).""" + base = torch.arange(6, dtype=torch.float32).reshape(2, 3) + transposed = base.t() + assert not transposed.is_contiguous() + result = _hash_tensor_bytes(transposed) + assert isinstance(result, bytes) + assert len(result) == 6 * 4 + + +def test_hash_tensor_bytes_noncontiguous_matches_contiguous_row_major_bytes() -> None: + """The non-contiguous path returns the contiguous row-major bytes of the transposed values.""" + base = torch.arange(6, dtype=torch.float32).reshape(2, 3) + transposed = base.t() + result = _hash_tensor_bytes(transposed) + assert result == transposed.contiguous().reshape(-1).numpy().tobytes() + + +def test_hash_tensor_bytes_distinct_values_produce_distinct_bytes() -> None: + """Different tensor contents of the same shape and dtype produce different bytes.""" + a = torch.zeros(2, 3, dtype=torch.float32) + b = torch.ones(2, 3, dtype=torch.float32) + assert _hash_tensor_bytes(a) != _hash_tensor_bytes(b) + assert len(_hash_tensor_bytes(a)) == len(_hash_tensor_bytes(b)) == 6 * 4 diff --git a/tests/fast/utils/test_misc.py b/tests/fast/utils/test_misc.py index 810c2b67c75..1cc77107868 100644 --- a/tests/fast/utils/test_misc.py +++ b/tests/fast/utils/test_misc.py @@ -1,8 +1,9 @@ +import logging import os import pytest -from miles.utils.misc import FunctionRegistry, function_registry, load_function +from miles.utils.misc import FunctionRegistry, filter_keys, function_registry, load_function def _fn_a(): @@ -57,3 +58,37 @@ def test_registry_takes_precedence(self): with function_registry.temporary("os.path.join", _fn_b): assert load_function("os.path.join") is _fn_b assert load_function("os.path.join") is os.path.join + + +class TestFilterKeys: + def test_projects_dict_by_keys(self): + """filter_keys returns only the requested keys with their values.""" + d = {"a": 1, "b": 2, "c": 3} + assert filter_keys(d, ["a", "c"]) == {"a": 1, "c": 3} + + def test_empty_interest_keys_returns_empty_dict(self): + """An empty interest list yields an empty dict regardless of input.""" + assert filter_keys({"a": 1, "b": 2}, []) == {} + + def test_preserves_interest_keys_order(self): + """Result key order follows interest_keys, not the source dict order.""" + d = {"a": 1, "b": 2, "c": 3} + assert list(filter_keys(d, ["c", "a"]).keys()) == ["c", "a"] + + def test_full_subset_returns_all_entries(self): + """Requesting every key returns the whole projection.""" + d = {"x": 10, "y": 20} + assert filter_keys(d, ["x", "y"]) == {"x": 10, "y": 20} + + def test_duplicate_interest_key_collapses_to_single_entry(self): + """A repeated interest key produces a single dict entry.""" + d = {"a": 1, "b": 2} + assert filter_keys(d, ["a", "a"]) == {"a": 1} + + def test_missing_key_raises_key_error_and_logs(self, caplog): + """A missing key raises KeyError and logs the error with context.""" + d = {"a": 1} + with caplog.at_level(logging.ERROR, logger="miles.utils.misc"): + with pytest.raises(KeyError): + filter_keys(d, ["a", "missing"]) + assert any("filter_keys" in record.message for record in caplog.records) From 635aa50615b31f83aa9e98c5747912297d9c3a8e Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Mon, 22 Jun 2026 18:15:16 +0800 Subject: [PATCH 04/41] Preserve the process-group backend across reload Thread the original backend through ReloadableProcessGroup so that, when a process group is rebuilt (e.g. after a reconfigure/heal), it is recreated with the same backend instead of hard-coding NCCL. - reloadable_process_group.py: carry `backend` in the reload group info. --- miles/utils/reloadable_process_group.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/miles/utils/reloadable_process_group.py b/miles/utils/reloadable_process_group.py index d9f5412953f..295dd10e999 100644 --- a/miles/utils/reloadable_process_group.py +++ b/miles/utils/reloadable_process_group.py @@ -42,7 +42,8 @@ def new_group(*args, **kwargs): if len(ranks) == 1: return group - group = ReloadableProcessGroup(group, ranks) + backend = args[2] if len(args) >= 3 else kwargs.get("backend") + group = ReloadableProcessGroup(group, ranks, backend=backend) return group dist.new_group = new_group @@ -112,7 +113,7 @@ def convert(arg): class ReloadableProcessGroup(torch.distributed.ProcessGroup): GROUPS = {} - def __init__(self, group, ranks): + def __init__(self, group, ranks, backend=None): super().__init__( rank=dist.get_rank(group), size=dist.get_world_size(group), @@ -120,6 +121,8 @@ def __init__(self, group, ranks): self.group = group self.group_info = { "ranks": ranks, + # None = inherit the default backend at reload time. + "backend": backend, } pid = os.getpid() if pid not in ReloadableProcessGroup.GROUPS: @@ -155,7 +158,9 @@ def reload_process_groups(): for reloadable_group in reloadable_groups: if reloadable_group.group is not None: continue - group = old_new_group(ranks=reloadable_group.group_info["ranks"], backend="nccl") + group = old_new_group( + ranks=reloadable_group.group_info["ranks"], backend=reloadable_group.group_info["backend"] + ) reloadable_group.group = group def rank(self) -> int: From 7a33a63f7df671bd8df8d1a6246f91ceeafbb4a3 Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Mon, 22 Jun 2026 18:15:16 +0800 Subject: [PATCH 05/41] Add fault-tolerance foundation utilities Add small foundation utilities used across the fault-tolerance trainer: a strict pydantic base model, a retry helper, a tensor checksum helper, a per-cell megatron world-size computation, the TrainStepOutcome enum, and the IndepDPInfo dataclass describing a cell's independent-DP identity. - pydantic_utils.py / retry_utils.py / checksum_utils.py / megatron_args_utils.py / types.py / indep_dp.py and tests. --- miles/backends/megatron_utils/types.py | 6 + miles/utils/checksum_utils.py | 34 +++ miles/utils/indep_dp.py | 26 +++ miles/utils/megatron_args_utils.py | 7 + miles/utils/pydantic_utils.py | 9 + miles/utils/retry_utils.py | 38 +++ tests/fast/utils/test_checksum_utils.py | 67 ++++++ tests/fast/utils/test_retry_utils.py | 299 ++++++++++++++++++++++++ 8 files changed, 486 insertions(+) create mode 100644 miles/backends/megatron_utils/types.py create mode 100644 miles/utils/checksum_utils.py create mode 100644 miles/utils/indep_dp.py create mode 100644 miles/utils/megatron_args_utils.py create mode 100644 miles/utils/pydantic_utils.py create mode 100644 miles/utils/retry_utils.py create mode 100644 tests/fast/utils/test_checksum_utils.py create mode 100644 tests/fast/utils/test_retry_utils.py diff --git a/miles/backends/megatron_utils/types.py b/miles/backends/megatron_utils/types.py new file mode 100644 index 00000000000..9a15742ea67 --- /dev/null +++ b/miles/backends/megatron_utils/types.py @@ -0,0 +1,6 @@ +from enum import StrEnum, auto + + +class TrainStepOutcome(StrEnum): + NORMAL = auto() + DISCARDED_SHOULD_RETRY = auto() diff --git a/miles/utils/checksum_utils.py b/miles/utils/checksum_utils.py new file mode 100644 index 00000000000..64a5f1e1f0b --- /dev/null +++ b/miles/utils/checksum_utils.py @@ -0,0 +1,34 @@ +from typing import Any + +InferenceEngineChecksums = dict[str, str] + + +def flatten_inference_engine_checksums(check_weights_result: Any) -> list[InferenceEngineChecksums]: + engine_bodies = _flatten_to_inference_engine_bodies(check_weights_result) + surviving = [body for body in engine_bodies if body is not None] + assert surviving, ( + f"check_weights('checksum') returned no non-None engine bodies " + f"(got {len(engine_bodies)} entries, all None): {check_weights_result!r}" + ) + return [_merge_inference_engine_ranks(body) for body in surviving] + + +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] + + +def _merge_inference_engine_ranks(engine_body: dict[str, Any]) -> InferenceEngineChecksums: + # Ranks arrive in non-deterministic (zmq) order under TP>1; sort and prefix each tensor + # name with rank{r}/ so distinct shards' identically-named tensors never clobber. + assert engine_body.get("success", False), f"check_weights engine reported failure: {engine_body!r}" + 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"]) + + merged: InferenceEngineChecksums = {} + for rank_info in ranks_sorted: + rank = rank_info["parallelism_info"]["rank"] + for name, value in rank_info["checksums"].items(): + merged[f"rank{rank}/{name}"] = value + return merged diff --git a/miles/utils/indep_dp.py b/miles/utils/indep_dp.py new file mode 100644 index 00000000000..4055add45b7 --- /dev/null +++ b/miles/utils/indep_dp.py @@ -0,0 +1,26 @@ +from dataclasses import dataclass + + +@dataclass(frozen=True) +class IndepDPInfo: + cell_index: int + num_cells: int + alive_rank: int + alive_size: int + quorum_id: int + alive_cell_indices: list[int] + + @classmethod + def create_trivial(cls) -> "IndepDPInfo": + return cls( + cell_index=0, + num_cells=1, + alive_rank=0, + alive_size=1, + quorum_id=0, + alive_cell_indices=[0], + ) + + def __post_init__(self): + assert self.alive_rank == self.alive_cell_indices.index(self.cell_index) + assert self.alive_size == len(self.alive_cell_indices) diff --git a/miles/utils/megatron_args_utils.py b/miles/utils/megatron_args_utils.py new file mode 100644 index 00000000000..2b58b1b8644 --- /dev/null +++ b/miles/utils/megatron_args_utils.py @@ -0,0 +1,7 @@ +""" +Utils for megatron arguments, but not related to megatron core logic +""" + + +def compute_megatron_world_size_except_dp(args) -> int: + return args.tensor_model_parallel_size * args.pipeline_model_parallel_size * args.context_parallel_size diff --git a/miles/utils/pydantic_utils.py b/miles/utils/pydantic_utils.py new file mode 100644 index 00000000000..43e9763b8e3 --- /dev/null +++ b/miles/utils/pydantic_utils.py @@ -0,0 +1,9 @@ +from pydantic import BaseModel, ConfigDict + + +class StrictBaseModel(BaseModel): + model_config = ConfigDict(extra="forbid") + + +class FrozenStrictBaseModel(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True) diff --git a/miles/utils/retry_utils.py b/miles/utils/retry_utils.py new file mode 100644 index 00000000000..ef98f6ea896 --- /dev/null +++ b/miles/utils/retry_utils.py @@ -0,0 +1,38 @@ +import asyncio +import logging +from collections.abc import Awaitable, Callable +from typing import Any + +logger = logging.getLogger(__name__) + +_DEFAULT_INITIAL_DELAY = 1.0 +_DEFAULT_MAX_DELAY = 60.0 +_DEFAULT_BACKOFF_FACTOR = 2.0 + + +async def retry( + fn: Callable[[int], Awaitable[Any]], + *, + initial_delay: float = _DEFAULT_INITIAL_DELAY, + max_delay: float = _DEFAULT_MAX_DELAY, + backoff_factor: float = _DEFAULT_BACKOFF_FACTOR, + sleep_fn: Callable[[float], Awaitable[None]] = asyncio.sleep, + max_attempts: int | None = None, +) -> None: + """Retry until ``fn`` does not throw, with exponential backoff.""" + assert max_attempts is None or max_attempts >= 1 + + attempt = 0 + delay = initial_delay + while True: + try: + await fn(attempt) + return + except Exception: + attempt += 1 + if max_attempts is not None and attempt >= max_attempts: + logger.warning(f"retry: attempt {attempt} failed, giving up (max_attempts={max_attempts})") + raise + logger.warning(f"retry: attempt {attempt} failed, retrying in {delay:.1f}s", exc_info=True) + await sleep_fn(delay) + delay = min(delay * backoff_factor, max_delay) diff --git a/tests/fast/utils/test_checksum_utils.py b/tests/fast/utils/test_checksum_utils.py new file mode 100644 index 00000000000..f5e0801c836 --- /dev/null +++ b/tests/fast/utils/test_checksum_utils.py @@ -0,0 +1,67 @@ +"""Tests for engine_weight_checksum.flatten_inference_engine_checksums.""" + +from typing import Any + +import pytest + +from miles.utils.checksum_utils import flatten_inference_engine_checksums + + +def _engine_body(*, success: bool, ranks: list[dict[str, Any]] | None) -> dict[str, Any]: + body: dict[str, Any] = {"success": success, "message": "ok"} + if ranks is not None: + body["ranks"] = ranks + return body + + +def _rank(rank: int, checksums: dict[str, str]) -> dict[str, Any]: + return {"checksums": checksums, "parallelism_info": {"rank": rank}} + + +class TestFlattenInferenceEngineChecksums: + def test_single_server_group_engine_single_rank(self) -> None: + """One server/group/engine with one rank yields one prefixed checksum dict.""" + result = [[[_engine_body(success=True, ranks=[_rank(0, {"w": "aaa"})])]]] + assert flatten_inference_engine_checksums(result) == [{"rank0/w": "aaa"}] + + def test_multiple_engines_flattened_in_order(self) -> None: + """Engines across servers/groups are flattened into a single ordered list.""" + result = [ + [[_engine_body(success=True, ranks=[_rank(0, {"w": "e0"})])]], + [[_engine_body(success=True, ranks=[_rank(0, {"w": "e1"})])]], + ] + assert flatten_inference_engine_checksums(result) == [{"rank0/w": "e0"}, {"rank0/w": "e1"}] + + def test_none_node_rank_payloads_filtered(self) -> None: + """None engine bodies (non-zero node ranks) are dropped before indexing.""" + result = [[[_engine_body(success=True, ranks=[_rank(0, {"w": "e0"})]), None]]] + assert flatten_inference_engine_checksums(result) == [{"rank0/w": "e0"}] + + def test_all_none_fails_loud(self) -> None: + """A None-only result means the checksum action did nothing, so fail loud.""" + result = [[[None, None]]] + with pytest.raises(AssertionError, match="no non-None engine bodies"): + flatten_inference_engine_checksums(result) + + def test_multi_rank_merged_with_rank_prefix(self) -> None: + """Multiple ranks of one engine merge into one dict, prefixed by rank to avoid clobber.""" + result = [[[_engine_body(success=True, ranks=[_rank(0, {"w": "r0"}), _rank(1, {"w": "r1"})])]]] + assert flatten_inference_engine_checksums(result) == [{"rank0/w": "r0", "rank1/w": "r1"}] + + def test_ranks_out_of_order_sorted_by_parallelism_rank(self) -> None: + """Ranks arriving out of order (zmq) are sorted by parallelism rank deterministically.""" + out_of_order = [[[_engine_body(success=True, ranks=[_rank(1, {"w": "r1"}), _rank(0, {"w": "r0"})])]]] + in_order = [[[_engine_body(success=True, ranks=[_rank(0, {"w": "r0"}), _rank(1, {"w": "r1"})])]]] + assert flatten_inference_engine_checksums(out_of_order) == flatten_inference_engine_checksums(in_order) + + def test_engine_failure_fails_loud(self) -> None: + """An engine reporting success=False fails loud rather than silently dropping it.""" + result = [[[_engine_body(success=False, ranks=[_rank(0, {"w": "aaa"})])]]] + with pytest.raises(AssertionError, match="reported failure"): + flatten_inference_engine_checksums(result) + + def test_engine_without_ranks_fails_loud(self) -> None: + """A success body with no ranks fails loud (nothing to compare).""" + result = [[[_engine_body(success=True, ranks=None)]]] + with pytest.raises(AssertionError, match="no ranks"): + flatten_inference_engine_checksums(result) diff --git a/tests/fast/utils/test_retry_utils.py b/tests/fast/utils/test_retry_utils.py new file mode 100644 index 00000000000..e3eb15f8431 --- /dev/null +++ b/tests/fast/utils/test_retry_utils.py @@ -0,0 +1,299 @@ +import pytest + +from miles.utils.retry_utils import retry + +pytestmark = pytest.mark.asyncio + + +class _FakeSleep: + """Records sleep calls without actually sleeping.""" + + def __init__(self) -> None: + self.delays: list[float] = [] + + async def __call__(self, delay: float) -> None: + self.delays.append(delay) + + +class TestRetryBasic: + async def test_succeeds_immediately(self): + call_count = 0 + fake_sleep = _FakeSleep() + + async def fn(_attempt): + nonlocal call_count + call_count += 1 + + await retry(fn, sleep_fn=fake_sleep) + + assert call_count == 1 + assert fake_sleep.delays == [] + + async def test_retries_then_succeeds(self): + call_count = 0 + fake_sleep = _FakeSleep() + + async def fn(_attempt): + nonlocal call_count + call_count += 1 + if call_count < 4: + raise ValueError("not yet") + + await retry(fn, initial_delay=1.0, sleep_fn=fake_sleep) + + assert call_count == 4 + assert len(fake_sleep.delays) == 3 + + async def test_single_retry(self): + call_count = 0 + fake_sleep = _FakeSleep() + + async def fn(_attempt): + nonlocal call_count + call_count += 1 + if call_count < 2: + raise RuntimeError("fail once") + + await retry(fn, initial_delay=1.0, sleep_fn=fake_sleep) + + assert call_count == 2 + assert len(fake_sleep.delays) == 1 + + async def test_fn_receives_correct_attempt_number(self): + """First call gets attempt=0, first retry gets attempt=1, etc.""" + received_attempts: list[int] = [] + fake_sleep = _FakeSleep() + + async def fn(attempt): + received_attempts.append(attempt) + if len(received_attempts) < 4: + raise ValueError("not yet") + + await retry(fn, initial_delay=1.0, sleep_fn=fake_sleep) + + assert received_attempts == [0, 1, 2, 3] + + +class TestRetryLogging: + async def test_logs_on_retry(self, caplog): + call_count = 0 + + async def fn(_attempt): + nonlocal call_count + call_count += 1 + if call_count < 3: + raise ValueError("boom") + + with caplog.at_level("WARNING"): + await retry(fn, initial_delay=1.0, sleep_fn=_FakeSleep()) + + retry_messages = [r for r in caplog.records if "retrying" in r.message] + assert len(retry_messages) == 2 + + async def test_no_log_on_first_success(self, caplog): + async def fn(_attempt): + pass + + with caplog.at_level("WARNING"): + await retry(fn, sleep_fn=_FakeSleep()) + + assert not any("retrying" in r.message for r in caplog.records) + + async def test_logs_include_exc_info(self, caplog): + call_count = 0 + + async def fn(_attempt): + nonlocal call_count + call_count += 1 + if call_count < 2: + raise ValueError("detail") + + with caplog.at_level("WARNING"): + await retry(fn, initial_delay=1.0, sleep_fn=_FakeSleep()) + + retry_records = [r for r in caplog.records if "retrying" in r.message] + assert len(retry_records) == 1 + assert retry_records[0].exc_info is not None + + async def test_log_message_includes_delay(self, caplog): + call_count = 0 + + async def fn(_attempt): + nonlocal call_count + call_count += 1 + if call_count < 2: + raise RuntimeError("fail") + + with caplog.at_level("WARNING"): + await retry(fn, initial_delay=2.5, sleep_fn=_FakeSleep()) + + retry_records = [r for r in caplog.records if "retrying" in r.message] + assert len(retry_records) == 1 + assert "2.5s" in retry_records[0].message + + +class TestRetryMaxAttempts: + async def test_raises_after_max_attempts(self): + """The last exception propagates once max_attempts calls have all failed.""" + call_count = 0 + fake_sleep = _FakeSleep() + + async def fn(_attempt): + nonlocal call_count + call_count += 1 + raise ValueError(f"fail {call_count}") + + with pytest.raises(ValueError, match="fail 3"): + await retry(fn, initial_delay=1.0, sleep_fn=fake_sleep, max_attempts=3) + + assert call_count == 3 + assert len(fake_sleep.delays) == 2 + + async def test_succeeds_on_last_allowed_attempt(self): + """No exception when fn succeeds exactly at the max_attempts-th call.""" + call_count = 0 + fake_sleep = _FakeSleep() + + async def fn(_attempt): + nonlocal call_count + call_count += 1 + if call_count < 3: + raise RuntimeError("not yet") + + await retry(fn, initial_delay=1.0, sleep_fn=fake_sleep, max_attempts=3) + + assert call_count == 3 + + async def test_max_attempts_one_never_retries(self): + """max_attempts=1 means a single call with no retry.""" + call_count = 0 + fake_sleep = _FakeSleep() + + async def fn(_attempt): + nonlocal call_count + call_count += 1 + raise RuntimeError("fail") + + with pytest.raises(RuntimeError): + await retry(fn, sleep_fn=fake_sleep, max_attempts=1) + + assert call_count == 1 + assert fake_sleep.delays == [] + + async def test_default_is_unlimited(self): + """Without max_attempts, retry keeps going far beyond any small cap.""" + call_count = 0 + + async def fn(_attempt): + nonlocal call_count + call_count += 1 + if call_count < 50: + raise RuntimeError("fail") + + await retry(fn, initial_delay=0.0, sleep_fn=_FakeSleep()) + + assert call_count == 50 + + async def test_invalid_max_attempts_rejected(self): + """max_attempts below 1 is a programming error.""" + + async def fn(_attempt): + pass + + with pytest.raises(AssertionError): + await retry(fn, sleep_fn=_FakeSleep(), max_attempts=0) + + async def test_gives_up_log_message(self, caplog): + """The final failure logs a giving-up warning instead of a retrying one.""" + + async def fn(_attempt): + raise RuntimeError("fail") + + with caplog.at_level("WARNING"): + with pytest.raises(RuntimeError): + await retry(fn, initial_delay=1.0, sleep_fn=_FakeSleep(), max_attempts=2) + + assert any("giving up" in r.message for r in caplog.records) + assert len([r for r in caplog.records if "retrying" in r.message]) == 1 + + +class TestRetryBackoff: + async def test_delay_doubles_each_retry(self): + call_count = 0 + fake_sleep = _FakeSleep() + + async def fn(_attempt): + nonlocal call_count + call_count += 1 + if call_count <= 4: + raise RuntimeError("fail") + + await retry(fn, initial_delay=1.0, max_delay=100.0, backoff_factor=2.0, sleep_fn=fake_sleep) + + assert fake_sleep.delays == [1.0, 2.0, 4.0, 8.0] + + async def test_delay_capped_at_max(self): + call_count = 0 + fake_sleep = _FakeSleep() + + async def fn(_attempt): + nonlocal call_count + call_count += 1 + if call_count <= 5: + raise RuntimeError("fail") + + await retry(fn, initial_delay=1.0, max_delay=3.0, backoff_factor=2.0, sleep_fn=fake_sleep) + + assert fake_sleep.delays == [1.0, 2.0, 3.0, 3.0, 3.0] + + async def test_custom_backoff_factor(self): + call_count = 0 + fake_sleep = _FakeSleep() + + async def fn(_attempt): + nonlocal call_count + call_count += 1 + if call_count <= 3: + raise RuntimeError("fail") + + await retry(fn, initial_delay=1.0, max_delay=100.0, backoff_factor=3.0, sleep_fn=fake_sleep) + + assert fake_sleep.delays == [1.0, 3.0, 9.0] + + async def test_zero_initial_delay(self): + call_count = 0 + fake_sleep = _FakeSleep() + + async def fn(_attempt): + nonlocal call_count + call_count += 1 + if call_count < 3: + raise RuntimeError("fail") + + await retry(fn, initial_delay=0.0, sleep_fn=fake_sleep) + + assert call_count == 3 + assert fake_sleep.delays == [0.0, 0.0] + + async def test_default_params_are_reasonable(self): + from miles.utils.retry_utils import _DEFAULT_BACKOFF_FACTOR, _DEFAULT_INITIAL_DELAY, _DEFAULT_MAX_DELAY + + assert _DEFAULT_INITIAL_DELAY == 1.0 + assert _DEFAULT_MAX_DELAY == 60.0 + assert _DEFAULT_BACKOFF_FACTOR == 2.0 + + async def test_many_retries_stay_capped(self): + """After hitting max_delay, all subsequent delays remain at max.""" + call_count = 0 + fake_sleep = _FakeSleep() + + async def fn(_attempt): + nonlocal call_count + call_count += 1 + if call_count <= 8: + raise RuntimeError("fail") + + await retry(fn, initial_delay=1.0, max_delay=5.0, backoff_factor=2.0, sleep_fn=fake_sleep) + + # 1, 2, 4, 5, 5, 5, 5, 5 + assert fake_sleep.delays == [1.0, 2.0, 4.0, 5.0, 5.0, 5.0, 5.0, 5.0] From 8dc8ddd9cc49d993e48f10d78c7d8fd606790286 Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Mon, 22 Jun 2026 18:15:16 +0800 Subject: [PATCH 06/41] Add structured logfmt logging helper Add a `log_structured` helper that emits logfmt-style key/value log lines, used by the fault-tolerance components for greppable structured logs. - structured_log.py (+ test). --- miles/utils/structured_log.py | 125 ++++++++++++ tests/fast/utils/test_logging_utils.py | 2 +- tests/fast/utils/test_structured_log.py | 255 ++++++++++++++++++++++++ 3 files changed, 381 insertions(+), 1 deletion(-) create mode 100644 miles/utils/structured_log.py create mode 100644 tests/fast/utils/test_structured_log.py diff --git a/miles/utils/structured_log.py b/miles/utils/structured_log.py new file mode 100644 index 00000000000..b74849afbd6 --- /dev/null +++ b/miles/utils/structured_log.py @@ -0,0 +1,125 @@ +import asyncio +import functools +import inspect +import json +import logging +import time +from collections.abc import Callable +from typing import Any + +_PRUNE_CAP = 160 + + +def log_structured(log_fn: Callable[..., None], *, exc_info: bool = False, **fields: Any) -> None: + log_fn("ft " + _to_logfmt(fields), stacklevel=2, exc_info=exc_info) + + +def with_logs(func: Callable[..., Any]) -> Callable[..., Any]: + fn_name = func.__name__ + method_logger = logging.getLogger(func.__module__) + + def log_start(args: tuple[Any, ...]) -> tuple[str, float]: + cls = type(args[0]).__name__ if args else "" + log_structured(method_logger.info, cls=cls, fn=fn_name, phase="start") + return cls, time.monotonic() + + def log_end(cls: str, start: float) -> None: + log_structured(method_logger.info, cls=cls, fn=fn_name, phase="end", ok=True, elapsed_s=_elapsed(start)) + + def log_fail(cls: str, start: float) -> None: + log_structured( + method_logger.error, cls=cls, fn=fn_name, phase="end", ok=False, elapsed_s=_elapsed(start), exc_info=True + ) + + def log_cancelled(cls: str, start: float) -> None: + log_structured( + method_logger.info, cls=cls, fn=fn_name, phase="end", ok=False, elapsed_s=_elapsed(start), cancelled=True + ) + + if inspect.iscoroutinefunction(func): + + @functools.wraps(func) + async def async_wrapper(*args: Any, **kwargs: Any) -> Any: + cls, start = log_start(args) + try: + result = await func(*args, **kwargs) + except asyncio.CancelledError: + log_cancelled(cls, start) + raise + except BaseException: + log_fail(cls, start) + raise + log_end(cls, start) + return result + + return async_wrapper + + @functools.wraps(func) + def wrapper(*args: Any, **kwargs: Any) -> Any: + cls, start = log_start(args) + try: + result = func(*args, **kwargs) + except BaseException: + log_fail(cls, start) + raise + log_end(cls, start) + return result + + return wrapper + + +def _elapsed(start: float) -> float: + return round(time.monotonic() - start, 1) + + +def _to_logfmt(fields: dict[str, Any]) -> str: + return " ".join(f"{key}={_format_value(value)}" for key, value in fields.items()) + + +def _format_value(value: Any) -> str: + if value is None: + return "" + if isinstance(value, bool): + return "true" if value else "false" + if isinstance(value, (int, float)): + return str(value) + if isinstance(value, (list, tuple)): + return _maybe_quote(",".join(_format_scalar(item) for item in value)) + if isinstance(value, dict): + return _quote(json.dumps(value, separators=(",", ":"), default=str)) + return _maybe_quote(str(value)) + + +def _format_scalar(value: Any) -> str: + if isinstance(value, bool): + return "true" if value else "false" + return str(value) + + +def _maybe_quote(text: str) -> str: + if text and any(ch in text for ch in (" ", "=", '"', "\\", "\n", "\r", "\t")): + return _quote(text) + return text + + +def _quote(text: str) -> str: + escaped = ( + text.replace("\\", "\\\\").replace('"', '\\"').replace("\n", "\\n").replace("\r", "\\r").replace("\t", "\\t") + ) + return f'"{escaped}"' + + +def prune_for_log(value: Any, cap: int = _PRUNE_CAP) -> Any: + if len(_compact_json(value)) <= cap: + return value + if isinstance(value, dict): + return {key: prune_for_log(item, cap) for key, item in value.items()} + if isinstance(value, (list, tuple)): + return f"" + if isinstance(value, str): + return f"" + return f"<{type(value).__name__}>" + + +def _compact_json(value: Any) -> str: + return json.dumps(value, separators=(",", ":"), default=str) diff --git a/tests/fast/utils/test_logging_utils.py b/tests/fast/utils/test_logging_utils.py index b346d372fdd..0bfbc713d55 100644 --- a/tests/fast/utils/test_logging_utils.py +++ b/tests/fast/utils/test_logging_utils.py @@ -30,7 +30,7 @@ def _run_snippet(code: str) -> subprocess.CompletedProcess: [sys.executable, "-c", textwrap.dedent(code)], capture_output=True, text=True, - timeout=10, + timeout=60, ) diff --git a/tests/fast/utils/test_structured_log.py b/tests/fast/utils/test_structured_log.py new file mode 100644 index 00000000000..c712ea6e459 --- /dev/null +++ b/tests/fast/utils/test_structured_log.py @@ -0,0 +1,255 @@ +import asyncio +import logging + +import pytest + +from miles.utils.structured_log import _format_value, _to_logfmt, log_structured, prune_for_log, with_logs + + +class TestFormatValue: + def test_int_renders_bare(self): + """An int renders as a bare number.""" + assert _format_value(7) == "7" + + def test_float_renders_bare(self): + """A float renders as a bare number.""" + assert _format_value(49.3) == "49.3" + + def test_true_and_false_render_lowercase(self): + """Bools render as lowercase true/false (and take precedence over the int branch).""" + assert _format_value(True) == "true" + assert _format_value(False) == "false" + + def test_none_renders_as_empty(self): + """None renders as an empty value.""" + assert _format_value(None) == "" + + def test_empty_string_renders_as_empty(self): + """An empty string renders as an empty value with no quotes.""" + assert _format_value("") == "" + + def test_plain_string_renders_bare(self): + """A string with no space, '=' or quote renders bare.""" + assert _format_value("train") == "train" + + def test_string_with_space_is_quoted(self): + """A string containing a space is double-quoted so it stays one token.""" + assert _format_value("survivors normal") == '"survivors normal"' + + def test_string_with_equals_is_quoted(self): + """A string containing '=' is quoted so it can't be mis-read as a new key.""" + assert _format_value("a=b") == '"a=b"' + + def test_string_with_quote_is_escaped(self): + """A string containing a double-quote is quoted with the inner quote escaped.""" + assert _format_value('say "hi"') == '"say \\"hi\\""' + + def test_string_with_backslash_and_space_escapes_backslash(self): + """When quoting is triggered, backslashes are escaped too.""" + assert _format_value("a\\b c") == '"a\\\\b c"' + + def test_string_with_newline_is_quoted_and_escaped(self): + """A literal newline is escaped so the entry stays on one log line.""" + assert _format_value("a\nb") == '"a\\nb"' + + def test_string_with_carriage_return_is_quoted_and_escaped(self): + """A carriage return is escaped so it can't corrupt the log stream.""" + assert _format_value("a\rb") == '"a\\rb"' + + def test_string_with_tab_is_quoted_and_escaped(self): + """A tab is escaped so the value stays one unambiguous token.""" + assert _format_value("a\tb") == '"a\\tb"' + + def test_string_with_only_backslash_is_quoted_and_escaped(self): + """A backslash alone (no space) still triggers quoting so it round-trips.""" + assert _format_value("a\\b") == '"a\\\\b"' + + def test_string_mixing_all_control_chars(self): + """Newline, quote, tab and backslash all escape together.""" + assert _format_value('a\n"b"\t\\c\r') == '"a\\n\\"b\\"\\t\\\\c\\r"' + + def test_backslash_escaped_before_other_escapes(self): + """A pre-escaped-looking string round-trips: the backslash doubles, the newline escapes.""" + assert _format_value("\\n\n") == '"\\\\n\\n"' + + def test_list_element_with_newline_is_quoted(self): + """A newline inside a list element quotes and escapes the joined value.""" + assert _format_value(["a\nb", "c"]) == '"a\\nb,c"' + + def test_list_of_ints_is_comma_joined_no_space(self): + """A list renders comma-joined with no spaces.""" + assert _format_value([0, 1, 2]) == "0,1,2" + + def test_list_of_bools_lowercased(self): + """List elements that are bools also render lowercase.""" + assert _format_value([True, False]) == "true,false" + + def test_empty_list_renders_as_empty(self): + """An empty list renders as an empty value.""" + assert _format_value([]) == "" + + def test_tuple_behaves_like_list(self): + """A tuple renders the same as the equivalent list.""" + assert _format_value((0, 1)) == "0,1" + + def test_list_whose_joined_form_has_a_space_is_quoted(self): + """If a list element introduces a space, the whole joined value is quoted.""" + assert _format_value(["a b", "c"]) == '"a b,c"' + + def test_dict_renders_as_quoted_compact_json(self): + """A dict value renders as quoted compact JSON.""" + assert _format_value({"a": 1}) == '"{\\"a\\":1}"' + + +class TestToLogfmt: + def test_fields_join_with_space_in_insertion_order(self): + """Fields render as space-separated key=value in insertion order.""" + assert _to_logfmt({"cell": 1, "fn": "train", "ok": True}) == "cell=1 fn=train ok=true" + + def test_empty_fields_render_as_empty_string(self): + """No fields renders as an empty string.""" + assert _to_logfmt({}) == "" + + def test_list_and_empty_fields_keep_one_token_each(self): + """A list field stays one token; an empty field renders as a trailing 'key='.""" + assert _to_logfmt({"alive": [0, 1], "pending": []}) == "alive=0,1 pending=" + + +class TestPruneForLog: + def test_small_payload_kept_verbatim(self): + """A payload under the cap is returned unchanged.""" + payload = {"quorum_id": 1, "healed": [0]} + assert prune_for_log(payload, cap=160) == payload + + def test_large_list_field_summarized_small_siblings_kept(self): + """An oversized list field becomes a length summary; small siblings stay inline.""" + pruned = prune_for_log({"quorum_id": 1, "checksums": list(range(500))}, cap=80) + assert pruned == {"quorum_id": 1, "checksums": ""} + + def test_large_string_field_summarized(self): + """An oversized string field becomes a char-count summary.""" + assert prune_for_log({"blob": "x" * 1000}, cap=80) == {"blob": ""} + + def test_nested_dict_prunes_only_the_big_subfield(self): + """Recursion summarizes only the oversized nested field, keeping small ones.""" + pruned = prune_for_log({"a": 1, "b": {"big": list(range(500)), "small": 2}}, cap=60) + assert pruned == {"a": 1, "b": {"big": "", "small": 2}} + + def test_large_tuple_summarized_as_list(self): + """An oversized tuple is summarized like a list.""" + assert prune_for_log(tuple(range(500)), cap=80) == "" + + def test_value_at_cap_kept_just_over_summarized(self): + """A value whose compact-JSON length is <= cap is kept; one char over is summarized.""" + assert prune_for_log("x" * 8, cap=10) == "x" * 8 # json '"xxxxxxxx"' is exactly 10 chars + assert prune_for_log("x" * 9, cap=10) == "" # json is 11 chars + + +class TestLogStructured: + def test_emits_ft_prefixed_logfmt_via_given_method(self, caplog): + """log_structured emits one 'ft '-prefixed logfmt line through the passed logger method.""" + logger = logging.getLogger("t_emit") + with caplog.at_level(logging.INFO, logger="t_emit"): + log_structured(logger.info, op="execute", phase="start", cell=1, fn="train") + assert caplog.messages == ["ft op=execute phase=start cell=1 fn=train"] + + def test_uses_the_level_of_the_passed_method(self, caplog): + """The record's level is that of the bound method passed (warning here).""" + logger = logging.getLogger("t_level") + with caplog.at_level(logging.DEBUG, logger="t_level"): + log_structured(logger.warning, op="x") + assert caplog.records[0].levelno == logging.WARNING + + def test_exc_info_is_forwarded(self, caplog): + """exc_info=True attaches the active exception to the record.""" + logger = logging.getLogger("t_exc") + with caplog.at_level(logging.ERROR, logger="t_exc"): + try: + raise ValueError("boom") + except ValueError: + log_structured(logger.error, op="x", exc_info=True) + assert caplog.records[0].exc_info[0] is ValueError + + def test_stacklevel_points_to_caller_not_helper(self, caplog): + """stacklevel=2 makes the record's filename the caller's, not structured_log.py.""" + logger = logging.getLogger("t_stack") + with caplog.at_level(logging.INFO, logger="t_stack"): + log_structured(logger.info, op="x") + assert caplog.records[0].filename == "test_structured_log.py" + + +class FakeActor: + @with_logs + def train(self): + return 42 + + @with_logs + def update_weights(self): + raise ValueError("boom") + + @with_logs + async def send_ckpt(self): + return "done" + + @with_logs + async def wait_forever(self): + await asyncio.Event().wait() + + +class TestWithLogs: + def test_sync_method_emits_start_then_end_with_class_and_method(self, caplog): + """A sync call emits phase=start then phase=end carrying the auto-detected class and method.""" + with caplog.at_level(logging.INFO): + assert FakeActor().train() == 42 + assert caplog.messages[0] == "ft cls=FakeActor fn=train phase=start" + assert caplog.messages[1].startswith("ft cls=FakeActor fn=train phase=end ok=true elapsed_s=") + + def test_class_is_read_from_the_runtime_instance(self, caplog): + """The cls field reflects the actual runtime instance, not where the method was defined.""" + + class SubActor(FakeActor): + pass + + with caplog.at_level(logging.INFO): + SubActor().train() + assert caplog.messages[0] == "ft cls=SubActor fn=train phase=start" + + def test_exception_logs_end_not_ok_with_exc_info_and_reraises(self, caplog): + """On exception the end line is ok=false with the traceback, and the error propagates.""" + with caplog.at_level(logging.INFO): + with pytest.raises(ValueError, match="boom"): + FakeActor().update_weights() + assert caplog.messages[0] == "ft cls=FakeActor fn=update_weights phase=start" + end = caplog.records[1] + assert end.getMessage().startswith("ft cls=FakeActor fn=update_weights phase=end ok=false elapsed_s=") + assert end.levelno == logging.ERROR + assert end.exc_info[0] is ValueError + + def test_async_method_emits_start_then_end_ok(self, caplog): + """An async call emits the same start/end pair after the coroutine completes.""" + with caplog.at_level(logging.INFO): + assert asyncio.run(FakeActor().send_ckpt()) == "done" + assert caplog.messages[0] == "ft cls=FakeActor fn=send_ckpt phase=start" + assert caplog.messages[1].startswith("ft cls=FakeActor fn=send_ckpt phase=end ok=true elapsed_s=") + + def test_async_cancellation_logs_cancelled_at_info_and_reraises(self, caplog): + """Cancelling an async call emits an info end line with cancelled=true and no traceback.""" + + async def run() -> None: + task = asyncio.ensure_future(FakeActor().wait_forever()) + await asyncio.sleep(0) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + with caplog.at_level(logging.INFO): + asyncio.run(run()) + end = caplog.records[1] + assert end.getMessage().startswith("ft cls=FakeActor fn=wait_forever phase=end ok=false elapsed_s=") + assert end.getMessage().endswith("cancelled=true") + assert end.levelno == logging.INFO + assert not end.exc_info + + def test_preserves_wrapped_function_metadata(self): + """functools.wraps keeps the original __name__ so Ray/introspection still sees it.""" + assert FakeActor.train.__name__ == "train" From 96d220f35ade97e33a58aa4539270c96ce6b9626 Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Mon, 22 Jun 2026 18:15:16 +0800 Subject: [PATCH 07/41] Add a Clock abstraction with a fake clock for tests Add a small `Clock` interface (`RealClock` plus a controllable fake clock) so time-dependent fault-tolerance code (health checks, heartbeats) can be driven deterministically in tests. - miles/utils/clock.py and tests. --- miles/utils/clock.py | 80 +++++++++++ tests/fast/utils/test_clock.py | 256 +++++++++++++++++++++++++++++++++ 2 files changed, 336 insertions(+) create mode 100644 miles/utils/clock.py create mode 100644 tests/fast/utils/test_clock.py diff --git a/miles/utils/clock.py b/miles/utils/clock.py new file mode 100644 index 00000000000..73986d8d2e8 --- /dev/null +++ b/miles/utils/clock.py @@ -0,0 +1,80 @@ +from __future__ import annotations + +import abc +import asyncio +import heapq +import time +from typing import NamedTuple + + +class Clock(abc.ABC): + @abc.abstractmethod + def time(self) -> float: ... + + @abc.abstractmethod + async def sleep(self, seconds: float) -> None: ... + + +class RealClock(Clock): + def time(self) -> float: + return time.time() + + async def sleep(self, seconds: float) -> None: + await asyncio.sleep(seconds) + + +_DRAIN_ITERATIONS = 20 + + +class _Waiter(NamedTuple): + target: float + seq: int + future: asyncio.Future[None] + + +class FakeClock(Clock): + """Deterministic clock for testing async time-dependent code. + + ``sleep()`` suspends the caller until ``elapse()`` advances the clock + past the target time. This gives tests precise control over which + sleeps resolve and when. + """ + + def __init__(self, start: float = 0.0) -> None: + self._now = start + self._waiters: list[_Waiter] = [] + self._counter: int = 0 + + def time(self) -> float: + return self._now + + async def sleep(self, seconds: float) -> None: + if seconds < 0: + return + + target = self._now + seconds + future: asyncio.Future[None] = asyncio.get_running_loop().create_future() + self._counter += 1 + heapq.heappush(self._waiters, _Waiter(target=target, seq=self._counter, future=future)) + self._resolve_ready() + await future + + async def elapse(self, seconds: float) -> None: + assert seconds >= 0, f"Cannot elapse negative time: {seconds}" + self._now += seconds + self._resolve_ready() + # Drain: yield enough times for resolved coroutines to run through + # sync code and register their next sleep. + for _ in range(_DRAIN_ITERATIONS): + await asyncio.sleep(0) + self._resolve_ready() + + def _resolve_ready(self) -> None: + while self._waiters and self._waiters[0].target <= self._now: + waiter = heapq.heappop(self._waiters) + if not waiter.future.done(): + waiter.future.set_result(None) + + @property + def pending_count(self) -> int: + return sum(1 for w in self._waiters if not w.future.done()) diff --git a/tests/fast/utils/test_clock.py b/tests/fast/utils/test_clock.py new file mode 100644 index 00000000000..7837f6d5f33 --- /dev/null +++ b/tests/fast/utils/test_clock.py @@ -0,0 +1,256 @@ +import asyncio + +import pytest + +from miles.utils.clock import FakeClock, RealClock + + +class TestRealClock: + def test_time_returns_current_time(self): + clock = RealClock() + import time + + before = time.time() + result = clock.time() + after = time.time() + assert before <= result <= after + + async def test_sleep_actually_waits(self): + clock = RealClock() + import time + + start = time.time() + await clock.sleep(0.05) + elapsed = time.time() - start + assert elapsed >= 0.04 + + +class TestFakeClockTime: + def test_initial_time(self): + clock = FakeClock(start=100.0) + assert clock.time() == 100.0 + + def test_default_start_is_zero(self): + clock = FakeClock() + assert clock.time() == 0.0 + + async def test_elapse_advances_time(self): + clock = FakeClock(start=10.0) + await clock.elapse(5.0) + assert clock.time() == 15.0 + + async def test_multiple_elapse_accumulate(self): + clock = FakeClock() + await clock.elapse(3.0) + await clock.elapse(7.0) + assert clock.time() == 10.0 + + async def test_elapse_zero_is_allowed(self): + clock = FakeClock(start=5.0) + await clock.elapse(0.0) + assert clock.time() == 5.0 + + async def test_elapse_negative_raises(self): + clock = FakeClock() + with pytest.raises(AssertionError, match="negative"): + await clock.elapse(-1.0) + + +class TestFakeClockSleep: + async def test_sleep_zero_returns_immediately(self): + clock = FakeClock() + completed = False + + async def task() -> None: + nonlocal completed + await clock.sleep(0) + completed = True + + asyncio.create_task(task()) + await asyncio.sleep(0) + assert completed + + async def test_sleep_blocks_until_elapse(self): + clock = FakeClock() + completed = False + + async def task() -> None: + nonlocal completed + await clock.sleep(10.0) + completed = True + + asyncio.create_task(task()) + await asyncio.sleep(0) + assert not completed + + await clock.elapse(10.0) + assert completed + + async def test_sleep_does_not_resolve_before_target(self): + clock = FakeClock() + completed = False + + async def task() -> None: + nonlocal completed + await clock.sleep(10.0) + completed = True + + asyncio.create_task(task()) + await asyncio.sleep(0) + + await clock.elapse(9.9) + assert not completed + + await clock.elapse(0.1) + assert completed + + async def test_multiple_sleeps_resolve_in_order(self): + clock = FakeClock() + order: list[str] = [] + + async def task_a() -> None: + await clock.sleep(5.0) + order.append("a") + + async def task_b() -> None: + await clock.sleep(10.0) + order.append("b") + + async def task_c() -> None: + await clock.sleep(3.0) + order.append("c") + + asyncio.create_task(task_a()) + asyncio.create_task(task_b()) + asyncio.create_task(task_c()) + await asyncio.sleep(0) + + await clock.elapse(5.0) + assert order == ["c", "a"] + + await clock.elapse(5.0) + assert order == ["c", "a", "b"] + + async def test_same_target_time_all_resolve(self): + clock = FakeClock() + count = 0 + + async def task() -> None: + nonlocal count + await clock.sleep(5.0) + count += 1 + + for _ in range(3): + asyncio.create_task(task()) + await asyncio.sleep(0) + + await clock.elapse(5.0) + assert count == 3 + + async def test_overshoot_resolves_all_pending(self): + clock = FakeClock() + order: list[int] = [] + + for delay in [1, 5, 10]: + d = delay + + async def task(d: int = d) -> None: + await clock.sleep(d) + order.append(d) + + asyncio.create_task(task()) + await asyncio.sleep(0) + + await clock.elapse(100.0) + assert sorted(order) == [1, 5, 10] + + async def test_sleep_negative_returns_immediately(self): + clock = FakeClock() + completed = False + + async def task() -> None: + nonlocal completed + await clock.sleep(-1.0) + completed = True + + asyncio.create_task(task()) + await asyncio.sleep(0) + assert completed + + +class TestFakeClockPendingCount: + async def test_no_pending_initially(self): + clock = FakeClock() + assert clock.pending_count == 0 + + async def test_pending_count_tracks_sleeps(self): + clock = FakeClock() + + asyncio.create_task(clock.sleep(5.0)) + asyncio.create_task(clock.sleep(10.0)) + await asyncio.sleep(0) + assert clock.pending_count == 2 + + await clock.elapse(5.0) + assert clock.pending_count == 1 + + await clock.elapse(5.0) + assert clock.pending_count == 0 + + +class TestFakeClockChainedSleeps: + async def test_sequential_sleeps_in_coroutine(self): + """A coroutine that does multiple sleeps in sequence.""" + clock = FakeClock() + checkpoints: list[float] = [] + + async def task() -> None: + checkpoints.append(clock.time()) + await clock.sleep(5.0) + checkpoints.append(clock.time()) + await clock.sleep(3.0) + checkpoints.append(clock.time()) + + asyncio.create_task(task()) + await asyncio.sleep(0) + + assert checkpoints == [0.0] + + await clock.elapse(5.0) + assert checkpoints == [0.0, 5.0] + + await clock.elapse(3.0) + assert checkpoints == [0.0, 5.0, 8.0] + + async def test_periodic_loop(self): + """Simulates a periodic loop like SimpleHealthChecker._loop.""" + clock = FakeClock() + ticks: list[float] = [] + + async def loop() -> None: + await clock.sleep(10.0) + while len(ticks) < 3: + ticks.append(clock.time()) + await clock.sleep(5.0) + + task = asyncio.create_task(loop()) + await asyncio.sleep(0) + + # Step 1: first_wait=10 + assert ticks == [] + await clock.elapse(10.0) + assert ticks == [10.0] + + # Step 2: interval=5 + await clock.elapse(5.0) + assert ticks == [10.0, 15.0] + + # Step 3: interval=5 + await clock.elapse(5.0) + assert ticks == [10.0, 15.0, 20.0] + + task.cancel() + try: + await task + except asyncio.CancelledError: + pass From a005a2f7a7fd259d741a725b58d6b4a82ec1d73b Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Mon, 22 Jun 2026 18:15:16 +0800 Subject: [PATCH 08/41] Add a fault injector test utility Add a fault-injector test utility used to deterministically exercise fault-tolerance code paths. - miles/utils/test_utils/fault_injector.py. --- miles/utils/test_utils/fault_injector.py | 41 ++++++++++++++++++++++++ 1 file changed, 41 insertions(+) create mode 100644 miles/utils/test_utils/fault_injector.py diff --git a/miles/utils/test_utils/fault_injector.py b/miles/utils/test_utils/fault_injector.py new file mode 100644 index 00000000000..35362597f08 --- /dev/null +++ b/miles/utils/test_utils/fault_injector.py @@ -0,0 +1,41 @@ +""" +Failure modes modeled after torchft's failure.py: +https://github.com/meta-pytorch/torchft/blob/main/examples/monarch/utils/failure.py +""" + +import ctypes +import logging +import os +import signal +from enum import Enum + +logger = logging.getLogger(__name__) + + +class FailureMode(Enum): + SIGKILL = "sigkill" + EXIT = "exit" + SEGFAULT = "segfault" + DEADLOCK = "deadlock" + + +def inject_fault(mode: str) -> None: + failure_mode = FailureMode(mode) + logger.warning("FaultInjector: executing %s (pid=%d)", failure_mode.value, os.getpid()) + + match failure_mode: + case FailureMode.SIGKILL: + os.kill(os.getpid(), signal.SIGKILL) + + case FailureMode.EXIT: + os._exit(1) + + case FailureMode.SEGFAULT: + crash_func = ctypes.CFUNCTYPE(None)() + crash_func() + + case FailureMode.DEADLOCK: + libc = ctypes.PyDLL(None) + libc.sleep.argtypes = (ctypes.c_uint,) + libc.sleep.restype = ctypes.c_uint + libc.sleep(600) From 836597152d198e3676894d1612d00d05d898bfd0 Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Mon, 22 Jun 2026 18:15:16 +0800 Subject: [PATCH 09/41] Add control-server data models Add the shared data models for the fault-tolerance control server (e.g. the `TriState` health value), used by the health checker and later by the HTTP control server. - miles/utils/control_server/models.py. --- miles/utils/control_server/__init__.py | 0 miles/utils/control_server/models.py | 85 +++++++++++++++++++++ tests/fast/utils/control_server/__init__.py | 0 3 files changed, 85 insertions(+) create mode 100644 miles/utils/control_server/__init__.py create mode 100644 miles/utils/control_server/models.py create mode 100644 tests/fast/utils/control_server/__init__.py diff --git a/miles/utils/control_server/__init__.py b/miles/utils/control_server/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/miles/utils/control_server/models.py b/miles/utils/control_server/models.py new file mode 100644 index 00000000000..00e47aadac8 --- /dev/null +++ b/miles/utils/control_server/models.py @@ -0,0 +1,85 @@ +from __future__ import annotations + +from enum import StrEnum +from typing import Literal + +from miles.utils.pydantic_utils import StrictBaseModel +from miles.utils.test_utils.fault_injector import FailureMode + + +class TriState(StrEnum): + """K8s condition status: ``"True"``, ``"False"``, or ``"Unknown"``.""" + + TRUE = "True" + FALSE = "False" + UNKNOWN = "Unknown" + + +class _OkResponse(StrictBaseModel): + status: str = "ok" + + +class CellCondition(StrictBaseModel): + type: Literal["Allocated", "Healthy"] + status: TriState + reason: str | None = None + message: str | None = None + lastTransitionTime: str | None = None + + @classmethod + def allocated(cls, status: TriState) -> CellCondition: + return cls(type="Allocated", status=status) + + @classmethod + def healthy(cls, status: TriState, *, reason: str | None = None) -> CellCondition: + return cls(type="Healthy", status=status, reason=reason) + + +class CellStatus(StrictBaseModel): + phase: Literal["Pending", "Running", "Suspended"] + conditions: list[CellCondition] + + +class CellSpec(StrictBaseModel): + suspend: bool = False + + +class CellMetadata(StrictBaseModel): + name: str + labels: dict[str, str] + + +class Cell(StrictBaseModel): + apiVersion: Literal["miles.io/v1"] = "miles.io/v1" + kind: Literal["Cell"] = "Cell" + metadata: CellMetadata + spec: CellSpec + status: CellStatus + + +class CellList(StrictBaseModel): + apiVersion: Literal["miles.io/v1"] = "miles.io/v1" + kind: Literal["CellList"] = "CellList" + items: list[Cell] + + +class CellPatchSpec(StrictBaseModel): + suspend: bool | None = None + + +class CellPatch(StrictBaseModel): + spec: CellPatchSpec | None = None + + +class FaultInjection(StrictBaseModel): + mode: FailureMode + sub_index: int = 0 + + +class K8sStatus(StrictBaseModel): + apiVersion: Literal["v1"] = "v1" + kind: Literal["Status"] = "Status" + status: Literal["Failure"] = "Failure" + message: str + reason: str + code: int diff --git a/tests/fast/utils/control_server/__init__.py b/tests/fast/utils/control_server/__init__.py new file mode 100644 index 00000000000..e69de29bb2d From e308821cf0580bd62e637816319951d6f3a4e536 Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Mon, 22 Jun 2026 18:15:16 +0800 Subject: [PATCH 10/41] Add a cell health checker and heartbeat utilities Add the periodic health checker (debounced TriState status driven by a Clock) and heartbeat utilities used to monitor train-cell liveness. - miles/utils/health_checker.py, miles/utils/heartbeat_utils.py and tests. --- miles/utils/health_checker.py | 247 ++++++++++++ miles/utils/heartbeat_utils.py | 22 + tests/fast/utils/test_health_checker.py | 493 +++++++++++++++++++++++ tests/fast/utils/test_heartbeat_utils.py | 111 +++++ 4 files changed, 873 insertions(+) create mode 100644 miles/utils/health_checker.py create mode 100644 miles/utils/heartbeat_utils.py create mode 100644 tests/fast/utils/test_health_checker.py create mode 100644 tests/fast/utils/test_heartbeat_utils.py diff --git a/miles/utils/health_checker.py b/miles/utils/health_checker.py new file mode 100644 index 00000000000..4d87435d060 --- /dev/null +++ b/miles/utils/health_checker.py @@ -0,0 +1,247 @@ +from __future__ import annotations + +import abc +import argparse +import asyncio +import logging +from collections.abc import Callable, Coroutine +from typing import Any + +from miles.utils.clock import Clock, RealClock +from miles.utils.control_server.models import TriState +from miles.utils.pydantic_utils import StrictBaseModel +from miles.utils.structured_log import log_structured + +logger = logging.getLogger(__name__) + + +class SimpleHealthCheckerConfig(StrictBaseModel): + interval: float + timeout: float + first_wait: float + failure_threshold: int + + @staticmethod + def add_arguments(parser: argparse.ArgumentParser, *, prefix: str) -> None: + parser.add_argument( + f"--{prefix}-interval", + type=float, + default=10.0, + help=f"Interval in seconds between {prefix} health checks.", + ) + parser.add_argument( + f"--{prefix}-timeout", + type=float, + default=10.0, + help=f"Timeout in seconds for a single {prefix} health check RPC.", + ) + parser.add_argument( + f"--{prefix}-first-wait", + type=float, + default=300.0, + help=f"Initial grace period (seconds) before starting {prefix} health checks.", + ) + parser.add_argument( + f"--{prefix}-failure-threshold", + type=int, + default=3, + help=( + f"Number of consecutive failed {prefix} checks before reporting unhealthy. " + "Debounces transient RPC blips so a single hiccup does not recycle a live cell." + ), + ) + + @staticmethod + def from_args(args: object, *, prefix: str) -> SimpleHealthCheckerConfig: + attr_prefix = prefix.replace("-", "_") + return SimpleHealthCheckerConfig( + interval=getattr(args, f"{attr_prefix}_interval"), + timeout=getattr(args, f"{attr_prefix}_timeout"), + first_wait=getattr(args, f"{attr_prefix}_first_wait"), + failure_threshold=getattr(args, f"{attr_prefix}_failure_threshold"), + ) + + +class BaseHealthChecker(abc.ABC): + @property + @abc.abstractmethod + def status(self) -> TriState: ... + + @abc.abstractmethod + async def start(self) -> None: ... + + @abc.abstractmethod + def stop(self) -> None: ... + + @abc.abstractmethod + def pause(self) -> None: ... + + @abc.abstractmethod + def resume(self) -> None: ... + + +class SimpleHealthChecker(BaseHealthChecker): + """Periodic async health checker. Calls *check_fn*; reports result via *on_result*. + + After each ``resume()``, waits ``first_wait`` seconds before the first check + (matching ``RolloutHealthMonitor._need_first_wait`` semantics). + """ + + def __init__( + self, + *, + name: str, + check_fn: Callable[[], Coroutine[Any, Any, None]], + on_result: Callable[[bool], None] | None = None, + config: SimpleHealthCheckerConfig, + clock: Clock | None = None, + ) -> None: + self._name = name + self._check_fn = check_fn + self._on_result = on_result + self._config = config + self._clock = clock or RealClock() + + self._status = TriState.UNKNOWN + self._paused: bool = False + self._need_first_wait: bool = True + self._consecutive_failures: int = 0 + self._task: asyncio.Task[None] | None = None + + @property + def status(self) -> TriState: + return self._status + + async def start(self) -> None: + if self._task is not None: + return + log_structured(logger.info, op="health", phase="start", name=self._name) + self._task = asyncio.create_task(self._loop()) + await asyncio.sleep(0) + + def stop(self) -> None: + if self._task is not None: + log_structured(logger.info, op="health", phase="stop", name=self._name) + self._task.cancel() + self._task = None + self._status = TriState.UNKNOWN + + def pause(self) -> None: + log_structured(logger.info, op="health", phase="pause", name=self._name) + self._paused = True + self._status = TriState.UNKNOWN + + def resume(self) -> None: + log_structured(logger.info, op="health", phase="resume", name=self._name) + self._paused = False + self._need_first_wait = True + self._status = TriState.UNKNOWN + self._consecutive_failures = 0 + + async def _loop(self) -> None: + while True: + if self._need_first_wait: + self._need_first_wait = False + log_structured( + logger.info, op="health", phase="first_wait", name=self._name, wait_s=self._config.first_wait + ) + await self._clock.sleep(self._config.first_wait) + + if not self._paused: + success = False + try: + await asyncio.wait_for(self._check_fn(), timeout=self._config.timeout) + success = True + except Exception: + log_structured(logger.error, op="health", phase="check_failed", name=self._name, exc_info=True) + + prev_status = self._status + if success: + self._consecutive_failures = 0 + self._status = TriState.TRUE + else: + self._consecutive_failures += 1 + if self._consecutive_failures >= self._config.failure_threshold: + self._status = TriState.FALSE + + log_structured( + logger.info, + op="health", + phase="poll", + name=self._name, + ok=success, + status=self._status.value, + consecutive_failures=self._consecutive_failures, + ) + + if prev_status != self._status: + log_structured( + logger.info, + op="health", + phase="status_change", + name=self._name, + from_status=prev_status.value, + to_status=self._status.value, + consecutive_failures=self._consecutive_failures, + ) + + if self._on_result is not None: + try: + self._on_result(success) + except Exception: + log_structured( + logger.error, op="health", phase="on_result_failed", name=self._name, exc_info=True + ) + + await self._clock.sleep(self._config.interval) + + +class NoopHealthChecker(BaseHealthChecker): + @property + def status(self) -> TriState: + return TriState.UNKNOWN + + async def start(self) -> None: + pass + + def stop(self) -> None: + pass + + def pause(self) -> None: + pass + + def resume(self) -> None: + pass + + +# TODO: should move when Rollout FT is implemented +def create_rollout_cell_health_checker( + *, + cell_id: str, + get_engines: Callable[[], list[Any]], + config: SimpleHealthCheckerConfig, + on_result: Callable[[bool], None] | None = None, +) -> SimpleHealthChecker: + + async def _check() -> None: + engines = get_engines() + if not engines: + raise RuntimeError("No engines") + + lead_engine = engines[0] + if lead_engine is None: + raise RuntimeError("Lead engine is None") + + await lead_engine.health_generate.remote() + + # Preserve the pre-debounce rollout semantics for now: a single failed check + # reports unhealthy immediately. The trainer cell checker uses the default + # failure_threshold; tuning rollout debouncing is left to Rollout FT work. + config = config.model_copy(update={"failure_threshold": 1}) + + return SimpleHealthChecker( + name=f"rollout-cell-{cell_id}", + check_fn=_check, + on_result=on_result, + config=config, + ) diff --git a/miles/utils/heartbeat_utils.py b/miles/utils/heartbeat_utils.py new file mode 100644 index 00000000000..9ae64fcebee --- /dev/null +++ b/miles/utils/heartbeat_utils.py @@ -0,0 +1,22 @@ +import time +from dataclasses import dataclass + + +@dataclass(frozen=True) +class HeartbeatStatus: + last_active_timestamp: float + bump_count: int + + +class SimpleHeartbeat: + def __init__(self) -> None: + self._status = HeartbeatStatus(last_active_timestamp=time.time(), bump_count=0) + + def bump(self) -> None: + self._status = HeartbeatStatus( + last_active_timestamp=time.time(), + bump_count=self._status.bump_count + 1, + ) + + def status(self) -> HeartbeatStatus: + return self._status diff --git a/tests/fast/utils/test_health_checker.py b/tests/fast/utils/test_health_checker.py new file mode 100644 index 00000000000..c403d7d78f3 --- /dev/null +++ b/tests/fast/utils/test_health_checker.py @@ -0,0 +1,493 @@ +import asyncio + + +from miles.utils.clock import FakeClock +from miles.utils.control_server.models import TriState +from miles.utils.health_checker import NoopHealthChecker, SimpleHealthChecker + + +async def _settle(clock: FakeClock) -> None: + for _ in range(1000): + if clock.pending_count >= 1: + return + await asyncio.sleep(0) + + +def _make_checker( + *, + check_fn=None, + on_result=None, + interval: float = 10.0, + timeout: float = 5.0, + first_wait: float = 0.0, + failure_threshold: int = 1, + name: str = "test", + clock: FakeClock | None = None, +) -> tuple[SimpleHealthChecker, FakeClock]: + from miles.utils.health_checker import SimpleHealthCheckerConfig + + if check_fn is None: + + async def check_fn() -> None: + pass + + c = clock or FakeClock() + checker = SimpleHealthChecker( + name=name, + check_fn=check_fn, + on_result=on_result, + config=SimpleHealthCheckerConfig( + interval=interval, timeout=timeout, first_wait=first_wait, failure_threshold=failure_threshold + ), + clock=c, + ) + return checker, c + + +class TestStartStop: + async def test_start_creates_task(self): + checker, _ = _make_checker() + assert checker._task is None + + await checker.start() + assert checker._task is not None + + checker.stop() + assert checker._task is None + + async def test_start_is_idempotent(self): + checker, _ = _make_checker() + await checker.start() + task = checker._task + + await checker.start() + assert checker._task is task + + checker.stop() + + async def test_stop_without_start_is_noop(self): + checker, _ = _make_checker() + checker.stop() + + +class TestCheckFnCalled: + async def test_check_fn_called_after_first_interval(self): + call_count = 0 + + async def check_fn() -> None: + nonlocal call_count + call_count += 1 + + checker, clock = _make_checker(check_fn=check_fn, interval=10.0) + await checker.start() + + # Step 1: first_wait=0, so first check runs immediately after task starts + await _settle(clock) + assert call_count == 1 + + # Step 2: Elapse less than interval — no second check + await clock.elapse(5.0) + assert call_count == 1 + + # Step 3: Elapse to interval — second check + await clock.elapse(5.0) + assert call_count == 2 + + checker.stop() + + async def test_first_wait_delays_first_check(self): + call_count = 0 + + async def check_fn() -> None: + nonlocal call_count + call_count += 1 + + checker, clock = _make_checker(check_fn=check_fn, first_wait=300.0, interval=10.0) + await checker.start() + + # Step 1: Elapse 100s — still in first_wait + await clock.elapse(100.0) + assert call_count == 0 + + # Step 2: Elapse to 300s — first_wait completes, first check runs + await clock.elapse(200.0) + assert call_count == 1 + + # Step 3: Elapse interval — second check + await clock.elapse(10.0) + assert call_count == 2 + + checker.stop() + + +class TestOnResult: + async def test_on_result_true_on_success(self): + results: list[bool] = [] + + checker, clock = _make_checker(on_result=lambda s: results.append(s)) + await checker.start() + await _settle(clock) + checker.stop() + + assert results == [True] + + async def test_on_result_false_on_failure(self): + results: list[bool] = [] + + async def check_fn() -> None: + raise RuntimeError("boom") + + checker, clock = _make_checker(check_fn=check_fn, on_result=lambda s: results.append(s)) + await checker.start() + await _settle(clock) + checker.stop() + + assert results == [False] + + async def test_loop_survives_on_result_raising(self, caplog): + """A raising on_result callback is logged and does not kill the check loop.""" + results: list[bool] = [] + + def on_result(success: bool) -> None: + results.append(success) + raise RuntimeError("callback boom") + + checker, clock = _make_checker(on_result=on_result, interval=5.0) + await checker.start() + + await _settle(clock) + await clock.elapse(5.0) + checker.stop() + + assert results == [True, True] + assert any("on_result_failed" in r.message for r in caplog.records) + + async def test_loop_continues_after_failure(self): + results: list[bool] = [] + + async def check_fn() -> None: + raise RuntimeError("boom") + + checker, clock = _make_checker(check_fn=check_fn, on_result=lambda s: results.append(s), interval=5.0) + await checker.start() + + await _settle(clock) + await clock.elapse(5.0) + checker.stop() + + assert results == [False, False] + + async def test_intermittent_failure(self): + call_count = 0 + results: list[bool] = [] + + async def check_fn() -> None: + nonlocal call_count + call_count += 1 + if call_count % 2 == 0: + raise RuntimeError("intermittent") + + checker, clock = _make_checker(check_fn=check_fn, on_result=lambda s: results.append(s), interval=5.0) + await checker.start() + await _settle(clock) + # first_wait=0 so first check runs immediately on start + assert results == [True] + + for _ in range(3): + await clock.elapse(5.0) + checker.stop() + + assert results == [True, False, True, False] + + +class TestPauseResume: + async def test_paused_checker_does_not_call_check_fn(self): + call_count = 0 + + async def check_fn() -> None: + nonlocal call_count + call_count += 1 + + checker, clock = _make_checker(check_fn=check_fn, interval=5.0) + checker.pause() + + await checker.start() + await clock.elapse(20.0) + checker.stop() + + assert call_count == 0 + + async def test_resume_after_pause_resumes_checking(self): + call_count = 0 + + async def check_fn() -> None: + nonlocal call_count + call_count += 1 + + checker, clock = _make_checker(check_fn=check_fn, interval=5.0) + checker.pause() + + await checker.start() + await clock.elapse(20.0) + assert call_count == 0 + + checker.resume() + await clock.elapse(5.0) + checker.stop() + + assert call_count >= 1 + + async def test_pause_resume_flags(self): + checker, _ = _make_checker() + assert not checker._paused + + checker.pause() + assert checker._paused + + checker.resume() + assert not checker._paused + + +class TestNeedFirstWait: + async def test_resume_triggers_first_wait_again(self): + """After resume, the loop waits first_wait before the next check.""" + call_count = 0 + + async def check_fn() -> None: + nonlocal call_count + call_count += 1 + + checker, clock = _make_checker(check_fn=check_fn, first_wait=100.0, interval=5.0) + await checker.start() + + # Step 1: Initial first_wait (100s) + await clock.elapse(50.0) + assert call_count == 0 + await clock.elapse(50.0) + assert call_count == 1 + + # Step 2: Normal interval (5s) + await clock.elapse(5.0) + assert call_count == 2 + + # Step 3: Pause + resume resets first_wait + checker.pause() + checker.resume() + + # Step 4: Need to elapse past the pending interval sleep first, + # then the new first_wait (100s) before next check + await clock.elapse(5.0) + assert call_count == 2 + + await clock.elapse(50.0) + assert call_count == 2 + + await clock.elapse(50.0) + assert call_count == 3 + + checker.stop() + + async def test_pause_without_resume_no_first_wait(self): + checker, clock = _make_checker(first_wait=300.0) + await checker.start() + await clock.elapse(300.0) + assert checker._need_first_wait is False + + checker.pause() + assert checker._need_first_wait is False + + checker.stop() + + +class TestTriState: + async def test_initial_status_is_unknown(self): + checker, _ = _make_checker() + assert checker.status == TriState.UNKNOWN + + async def test_healthy_after_successful_check(self): + checker, clock = _make_checker() + await checker.start() + await _settle(clock) + + assert checker.status == TriState.TRUE + checker.stop() + + async def test_unhealthy_after_failed_check(self): + async def check_fn() -> None: + raise RuntimeError("boom") + + checker, clock = _make_checker(check_fn=check_fn) + await checker.start() + await _settle(clock) + + assert checker.status == TriState.FALSE + checker.stop() + + async def test_stop_resets_to_unknown(self): + checker, clock = _make_checker() + await checker.start() + await _settle(clock) + assert checker.status == TriState.TRUE + + checker.stop() + assert checker.status == TriState.UNKNOWN + + async def test_pause_resets_to_unknown(self): + checker, clock = _make_checker() + await checker.start() + await _settle(clock) + assert checker.status == TriState.TRUE + + checker.pause() + assert checker.status == TriState.UNKNOWN + checker.stop() + + async def test_resume_resets_to_unknown(self): + checker, clock = _make_checker() + await checker.start() + await _settle(clock) + assert checker.status == TriState.TRUE + + checker.pause() + checker.resume() + assert checker.status == TriState.UNKNOWN + checker.stop() + + async def test_recovers_from_unhealthy_to_healthy(self): + call_count = 0 + + async def check_fn() -> None: + nonlocal call_count + call_count += 1 + if call_count == 1: + raise RuntimeError("transient") + + checker, clock = _make_checker(check_fn=check_fn, interval=5.0) + await checker.start() + + await _settle(clock) + assert checker.status == TriState.FALSE + + await clock.elapse(5.0) + assert checker.status == TriState.TRUE + + checker.stop() + + +class TestFailureThresholdDebounce: + """With failure_threshold > 1, transient failures must not flip the status to FALSE + until that many consecutive checks have failed; any success resets the counter.""" + + def _flaky_check_fn(self, outcomes: list[bool]): + idx = 0 + + async def check_fn() -> None: + nonlocal idx + ok = outcomes[idx] + idx += 1 + if not ok: + raise RuntimeError("boom") + + return check_fn + + async def test_below_threshold_keeps_previous_status(self): + # success, then 2 failures (< threshold 3): status stays TRUE. + check_fn = self._flaky_check_fn([True, False, False]) + checker, clock = _make_checker(check_fn=check_fn, interval=5.0, failure_threshold=3) + await checker.start() + await _settle(clock) + assert checker.status == TriState.TRUE + + await clock.elapse(5.0) + assert checker.status == TriState.TRUE + assert checker._consecutive_failures == 1 + + await clock.elapse(5.0) + assert checker.status == TriState.TRUE + assert checker._consecutive_failures == 2 + + checker.stop() + + async def test_status_flips_false_only_at_threshold(self): + check_fn = self._flaky_check_fn([False, False, False]) + checker, clock = _make_checker(check_fn=check_fn, interval=5.0, failure_threshold=3) + await checker.start() + + await _settle(clock) + assert checker.status == TriState.UNKNOWN # 1st failure: below threshold, keep initial UNKNOWN + + await clock.elapse(5.0) + assert checker.status == TriState.UNKNOWN # 2nd failure: still below threshold + + await clock.elapse(5.0) + assert checker.status == TriState.FALSE # 3rd consecutive failure: threshold reached + + checker.stop() + + async def test_success_resets_failure_counter(self): + # 2 failures, a success, then 2 more failures: never reaches 3 consecutive, stays TRUE. + check_fn = self._flaky_check_fn([False, False, True, False, False]) + checker, clock = _make_checker(check_fn=check_fn, interval=5.0, failure_threshold=3) + await checker.start() + + await _settle(clock) # fail 1 + await clock.elapse(5.0) # fail 2 + assert checker._consecutive_failures == 2 + + await clock.elapse(5.0) # success -> reset + assert checker.status == TriState.TRUE + assert checker._consecutive_failures == 0 + + await clock.elapse(5.0) # fail 1 + await clock.elapse(5.0) # fail 2 + assert checker.status == TriState.TRUE + assert checker._consecutive_failures == 2 + + checker.stop() + + async def test_on_result_reports_raw_per_check_not_debounced(self): + results: list[bool] = [] + check_fn = self._flaky_check_fn([True, False, False, False]) + checker, clock = _make_checker( + check_fn=check_fn, on_result=lambda s: results.append(s), interval=5.0, failure_threshold=3 + ) + await checker.start() + await _settle(clock) + for _ in range(3): + await clock.elapse(5.0) + checker.stop() + + assert results == [True, False, False, False] + + async def test_resume_resets_failure_counter(self): + check_fn = self._flaky_check_fn([False, False, False]) + checker, clock = _make_checker(check_fn=check_fn, interval=5.0, failure_threshold=3) + await checker.start() + await _settle(clock) + await clock.elapse(5.0) + assert checker._consecutive_failures == 2 + + checker.resume() + assert checker._consecutive_failures == 0 + checker.stop() + + +class TestRolloutHealthCheckerPreservesImmediateFailure: + def test_rollout_checker_forces_failure_threshold_one(self): + """Rollout health checker keeps pre-debounce semantics (single failure -> unhealthy) + even when handed a config with a larger threshold.""" + from miles.utils.health_checker import SimpleHealthCheckerConfig, create_rollout_cell_health_checker + + checker = create_rollout_cell_health_checker( + cell_id="c0", + get_engines=lambda: [object()], + config=SimpleHealthCheckerConfig(interval=10.0, timeout=10.0, first_wait=0.0, failure_threshold=5), + ) + + assert checker._config.failure_threshold == 1 + + +class TestNoopHealthChecker: + def test_noop_status_is_always_unknown(self): + checker = NoopHealthChecker() + assert checker.status == TriState.UNKNOWN diff --git a/tests/fast/utils/test_heartbeat_utils.py b/tests/fast/utils/test_heartbeat_utils.py new file mode 100644 index 00000000000..2fd567960b2 --- /dev/null +++ b/tests/fast/utils/test_heartbeat_utils.py @@ -0,0 +1,111 @@ +import miles.utils.heartbeat_utils as heartbeat_utils +from miles.utils.heartbeat_utils import HeartbeatStatus, SimpleHeartbeat + + +class TestSimpleHeartbeatBumpCount: + def test_initial_bump_count_is_zero(self): + """A fresh SimpleHeartbeat reports a bump_count of 0.""" + heartbeat = SimpleHeartbeat() + assert heartbeat.status().bump_count == 0 + + def test_single_bump_increments_count(self): + """One bump() raises bump_count from 0 to 1.""" + heartbeat = SimpleHeartbeat() + heartbeat.bump() + assert heartbeat.status().bump_count == 1 + + def test_bump_count_increments_monotonically(self): + """Repeated bump() calls increment bump_count by exactly one each time.""" + heartbeat = SimpleHeartbeat() + for expected in range(1, 6): + heartbeat.bump() + assert heartbeat.status().bump_count == expected + + def test_bump_count_never_resets(self): + """bump_count keeps growing and never drops back across many bumps.""" + heartbeat = SimpleHeartbeat() + previous = heartbeat.status().bump_count + for _ in range(20): + heartbeat.bump() + current = heartbeat.status().bump_count + assert current == previous + 1 + previous = current + assert heartbeat.status().bump_count == 20 + + +class TestSimpleHeartbeatStatusObject: + def test_status_returns_heartbeat_status_instance(self): + """status() yields a HeartbeatStatus dataclass instance.""" + heartbeat = SimpleHeartbeat() + assert isinstance(heartbeat.status(), HeartbeatStatus) + + def test_status_is_frozen(self): + """The returned HeartbeatStatus is frozen and rejects attribute mutation.""" + import dataclasses + + heartbeat = SimpleHeartbeat() + status = heartbeat.status() + try: + status.bump_count = 999 + except dataclasses.FrozenInstanceError: + pass + else: + raise AssertionError("expected FrozenInstanceError on frozen dataclass") + + def test_status_returns_new_object_each_bump(self): + """Each bump() replaces the status with a distinct object identity.""" + heartbeat = SimpleHeartbeat() + before = heartbeat.status() + heartbeat.bump() + after = heartbeat.status() + assert before is not after + + def test_status_stable_between_bumps(self): + """status() returns the same object identity when no bump occurs in between.""" + heartbeat = SimpleHeartbeat() + assert heartbeat.status() is heartbeat.status() + + def test_old_status_snapshot_unaffected_by_later_bump(self): + """A captured status snapshot keeps its bump_count after subsequent bumps.""" + heartbeat = SimpleHeartbeat() + snapshot = heartbeat.status() + heartbeat.bump() + heartbeat.bump() + assert snapshot.bump_count == 0 + assert heartbeat.status().bump_count == 2 + + +class TestSimpleHeartbeatTimestamp: + def test_initial_timestamp_from_time_time(self, monkeypatch): + """Construction stamps last_active_timestamp from time.time().""" + monkeypatch.setattr(heartbeat_utils.time, "time", lambda: 1000.0) + heartbeat = SimpleHeartbeat() + assert heartbeat.status().last_active_timestamp == 1000.0 + + def test_bump_updates_timestamp_from_time_time(self, monkeypatch): + """bump() refreshes last_active_timestamp to the current time.time() value.""" + clock = {"now": 1000.0} + monkeypatch.setattr(heartbeat_utils.time, "time", lambda: clock["now"]) + heartbeat = SimpleHeartbeat() + assert heartbeat.status().last_active_timestamp == 1000.0 + + clock["now"] = 2500.0 + heartbeat.bump() + status = heartbeat.status() + assert status.last_active_timestamp == 2500.0 + assert status.bump_count == 1 + + def test_consecutive_bumps_track_timestamp(self, monkeypatch): + """Each bump() captures the time.time() value observed at that call.""" + clock = {"now": 100.0} + monkeypatch.setattr(heartbeat_utils.time, "time", lambda: clock["now"]) + heartbeat = SimpleHeartbeat() + + clock["now"] = 200.0 + heartbeat.bump() + assert heartbeat.status().last_active_timestamp == 200.0 + + clock["now"] = 350.0 + heartbeat.bump() + assert heartbeat.status().last_active_timestamp == 350.0 + assert heartbeat.status().bump_count == 2 From 4a52fcf542b67f853b8b5e741bce201b7bf44ca6 Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Mon, 22 Jun 2026 18:15:16 +0800 Subject: [PATCH 11/41] Add the fault-tolerance dependency, CI label, and logger-config setup Add the nvidia-resiliency-ext dependency, the "ft" CI test label, the FT test fixtures in the rollout conftest, and route startup logging through configure_logger_raw. The fault-tolerance CLI arguments themselves now live with the features that consume them (distributed across the per-feature commits). --- miles/utils/arguments.py | 6 ++++-- tests/ci/labels.py | 2 ++ tests/fast/ray/rollout/conftest.py | 8 ++++++++ 3 files changed, 14 insertions(+), 2 deletions(-) diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index 7d7858b91fe..49d99898c76 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -667,7 +667,7 @@ def add_fault_tolerance_arguments(parser): "--use-fault-tolerance", action="store_true", default=False, - help="Whether to enable the fault tolerance function during rollout.", + help="Enable fault tolerance. Use --ft-components to select which components.", ) parser.add_argument( "--rollout-health-check-interval", @@ -1995,7 +1995,7 @@ def add_user_provided_function_arguments(parser): def parse_args(add_custom_arguments=None): # Users may call `parse_args` very early, thus we ensure logger is configured here - configure_logger() + configure_logger("main") add_miles_arguments = get_miles_extra_args_provider(add_custom_arguments) @@ -2016,6 +2016,7 @@ def parse_args(add_custom_arguments=None): args.indexer_rope_interleave = bool(getattr(hf_config, "indexer_rope_interleave", False)) logger.info(f"Setting indexer_rope_interleave: {args.indexer_rope_interleave} into args") + # TODO: unify this .rank and .world_size w/ indep_dp logics args.rank = 0 args.world_size = args.actor_num_nodes * args.actor_num_gpus_per_node args = set_default_megatron_args(args) @@ -2023,6 +2024,7 @@ def parse_args(add_custom_arguments=None): from miles.backends.experimental.fsdp_utils.arguments import load_fsdp_args args = load_fsdp_args(extra_args_provider=add_miles_arguments) + # TODO: unify this .rank and .world_size w/ indep_dp logics args.rank = 0 # Primary process rank for wandb initialization args.world_size = args.actor_num_nodes * args.actor_num_gpus_per_node diff --git a/tests/ci/labels.py b/tests/ci/labels.py index 22b5dfcbb9f..26646a1e114 100644 --- a/tests/ci/labels.py +++ b/tests/ci/labels.py @@ -26,6 +26,8 @@ "ckpt": "Checkpoint save / load tests", "lora": "LoRA training tests", "precision": "Numerical precision parity tests", + "ft-short": "Fault-tolerance trainer comparison tests (no_failure / deterministic / with_failure)", + "ft-long": "Fault-tolerance trainer soak tests (random-crash survival, realistic-gsm8k convergence)", "weight-update": "Weight update tests", "replay": "Routing / indexer replay tests", "qwen35": "Qwen3.5-35B-A3B MTP / spec-v2 e2e tests", diff --git a/tests/fast/ray/rollout/conftest.py b/tests/fast/ray/rollout/conftest.py index de99c167f3f..7e5e94fdbda 100644 --- a/tests/fast/ray/rollout/conftest.py +++ b/tests/fast/ray/rollout/conftest.py @@ -44,6 +44,7 @@ def make_args(**overrides: Any) -> Namespace: use_dynamic_global_batch_size=False, disable_rollout_trim_samples=False, balance_data=False, + delay_split_train_data_by_dp=False, # advantage / reward advantage_estimator="grpo", rewards_normalization=True, @@ -97,6 +98,13 @@ def make_args(**overrides: Any) -> Namespace: save_debug_rollout_data=None, load_debug_rollout_data=None, load_debug_rollout_data_subsample=None, + ci_inject_rollout_data_path=None, + ci_inject_rollout_data_start_rollout_id=None, + ci_inject_rollout_data_min_match_ratio=0.9, + # event checkpointing (event_logger.restore/snapshot in RolloutManager) + save_debug_event_data=None, + load=None, + save=None, # CI ci_test=False, # dumper (sglang debug dumper integration) From 2ed9fcf003971cd674ef81bc8805003edeea1379 Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Mon, 22 Jun 2026 18:15:16 +0800 Subject: [PATCH 12/41] Always reconnect rollout engines on weight-update setup Unconditionally disconnect-then-reconnect the model-update process group when (re)connecting rollout engines, guarding the destroy against a missing group, so a reconfigured/healed engine set can rebuild the NCCL group from scratch. - broadcast.py: drop the "only disconnect if group exists" short-circuit; guard `destroy_process_group` against None. --- .../update_weight_from_distributed/broadcast.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/miles/backends/megatron_utils/update_weight/update_weight_from_distributed/broadcast.py b/miles/backends/megatron_utils/update_weight/update_weight_from_distributed/broadcast.py index 679d31d4c92..f5ba0ae5ae1 100644 --- a/miles/backends/megatron_utils/update_weight/update_weight_from_distributed/broadcast.py +++ b/miles/backends/megatron_utils/update_weight/update_weight_from_distributed/broadcast.py @@ -73,8 +73,9 @@ def connect_rollout_engines( self._group_name = f"miles-pp_{pp_rank}" if self._is_source: - if (g := self._model_update_groups) is not None: - disconnect_rollout_engines_from_distributed(self.args, self._group_name, g, self.rollout_engines) + disconnect_rollout_engines_from_distributed( + self.args, self._group_name, self._model_update_groups, self.rollout_engines + ) self._model_update_groups = connect_rollout_engines_from_distributed( self.args, self._group_name, rollout_engines ) @@ -195,8 +196,11 @@ def disconnect_rollout_engines_from_distributed(args, group_name, model_update_g Destroy NCCL on training and engines. """ refs = [engine.destroy_weights_update_group.remote(group_name) for engine in rollout_engines] - dist.destroy_process_group(model_update_groups) - ray.get(refs) + try: + if model_update_groups is not None: + dist.destroy_process_group(model_update_groups) + finally: + ray.get(refs) def update_weights_from_distributed( From a838a0ba213b07b0294993ba81bf3185df3a3043 Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Mon, 22 Jun 2026 18:15:16 +0800 Subject: [PATCH 13/41] Add a fault-injection RPC to train actors Expose an `inject_fault` Ray method on TrainRayActor (in its own concurrency group) that triggers a configured failure mode via the fault injector, so fault-tolerance tests can crash/hang specific actors on demand. - train_actor.py: `inject_fault` RPC. --- miles/ray/actor_group.py | 4 +++- miles/ray/train_actor.py | 6 ++++++ 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/miles/ray/actor_group.py b/miles/ray/actor_group.py index f729d948aed..f1436cfafaa 100644 --- a/miles/ray/actor_group.py +++ b/miles/ray/actor_group.py @@ -89,7 +89,9 @@ def _allocate_gpus_for_actor(self, pg, num_gpus_per_actor): actor_impl = FSDPTrainRayActor - TrainRayActor = ray.remote(num_gpus=1, runtime_env={"env_vars": env_vars})(actor_impl) + TrainRayActor = ray.remote( + num_gpus=1, runtime_env={"env_vars": env_vars}, concurrency_groups={"fault_injector": 1} + )(actor_impl) # Create worker actors actor_handles = [] diff --git a/miles/ray/train_actor.py b/miles/ray/train_actor.py index 9e6659026fc..7ab2bfcd745 100644 --- a/miles/ray/train_actor.py +++ b/miles/ray/train_actor.py @@ -15,6 +15,7 @@ from miles.utils.env_report import collect_and_print_node_env_report from miles.utils.logging_utils import configure_logger from miles.utils.memory_utils import clear_memory, print_memory +from miles.utils.test_utils.fault_injector import inject_fault as _inject_fault if TYPE_CHECKING: from miles.ray.rollout.rollout_manager import EnginesAndLock @@ -53,6 +54,7 @@ def __init__(self, world_size, rank, master_addr, master_port): # os.environ["LOCAL_RANK"] = str(ray.get_gpu_ids()[0]) os.environ["LOCAL_RANK"] = str(get_local_gpu_id()) + # TODO mv the args into ctor def init(self, args, role, with_ref=False, with_opd_teacher=False): self.args = args self.role = role @@ -109,6 +111,10 @@ def init(self, args, role, with_ref=False, with_opd_teacher=False): except Exception as e: logger.info(f"Warning: Failed to set NUMA affinity: {e}") + @ray.method(concurrency_group="fault_injector") + def inject_fault(self, mode: str) -> None: + _inject_fault(mode=mode) + def clear_memory(self): print_memory("before TrainRayActor.clear_memory") clear_memory() From c7798b563c62bba207a61d340e87cd444b1e693b Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Wed, 8 Jul 2026 10:47:39 +0800 Subject: [PATCH 14/41] Delay splitting train data by DP until actor-side processing Extract the DP split into a witness-aware split_train_data_by_dp_raw helper (with unit tests; the key list also carries seq_witness_ids) and use it to split the training data on the actor side when delay_split_train_data_by_dp is set, deferring the DP split from the rollout side to actor-side processing. split_train_data_by_dp stays a thin wrapper that ray.puts each partition. - miles/ray/rollout/train_data_conversion.py (+ tests), miles/utils/data.py, actor_group.py, rollout_manager.py. --- miles/ray/actor_group.py | 4 +- miles/ray/rollout/rollout_manager.py | 8 ++- miles/ray/rollout/train_data_conversion.py | 13 +++- miles/utils/arguments.py | 5 ++ miles/utils/data.py | 18 ++++- .../rollout/real_ray/test_rollout_manager.py | 7 +- .../ray/rollout/test_train_data_conversion.py | 67 +++++++++++++++++++ 7 files changed, 111 insertions(+), 11 deletions(-) diff --git a/miles/ray/actor_group.py b/miles/ray/actor_group.py index f1436cfafaa..bfbfdf55e85 100644 --- a/miles/ray/actor_group.py +++ b/miles/ray/actor_group.py @@ -119,9 +119,9 @@ async def init(self): "init", self.args, self.role, with_ref=self.with_ref, with_opd_teacher=self.with_opd_teacher ) - async def train(self, rollout_id, rollout_data_ref): + async def train(self, rollout_id, rollout_data_pack): """Do one rollout training""" - await self._broadcast("train", rollout_id, rollout_data_ref) + await self._broadcast("train", rollout_id, rollout_data_pack["data_ref"]) async def save_model(self, rollout_id, force_sync=False): """Save actor model""" diff --git a/miles/ray/rollout/rollout_manager.py b/miles/ray/rollout/rollout_manager.py index c216b440b24..107d771233c 100644 --- a/miles/ray/rollout/rollout_manager.py +++ b/miles/ray/rollout/rollout_manager.py @@ -28,6 +28,7 @@ from miles.utils.logging_utils import configure_logger from miles.utils.metric_checker import MetricChecker from miles.utils.misc import load_function +from miles.utils.ray_utils import Box from miles.utils.tracking_utils import init_tracking logging.getLogger("httpx").setLevel(logging.WARNING) @@ -117,7 +118,12 @@ async def generate(self, rollout_id): custom_convert_samples_to_train_data_func=self.custom_convert_samples_to_train_data_func, custom_reward_post_process_func=self.custom_reward_post_process_func, ) - return split_train_data_by_dp(self.args, data, self.train_parallel_config["dp_size"]) + sample_indices = data.get("sample_indices") + if self.args.delay_split_train_data_by_dp: + data_ref = Box(ray.put(data)) + else: + data_ref = split_train_data_by_dp(self.args, data, self.train_parallel_config["dp_size"]) + return dict(sample_indices=sample_indices, data_ref=data_ref) async def eval(self, rollout_id): if self.args.debug_train_only: diff --git a/miles/ray/rollout/train_data_conversion.py b/miles/ray/rollout/train_data_conversion.py index 65bc8d4b6db..e67087fbb3d 100644 --- a/miles/ray/rollout/train_data_conversion.py +++ b/miles/ray/rollout/train_data_conversion.py @@ -122,6 +122,12 @@ def _post_process_rewards(args, samples: list[Sample] | list[list[Sample]], cust def split_train_data_by_dp(args, data, dp_size): + """Split the train data by data parallel size.""" + rollout_data_list = split_train_data_by_dp_raw(args, data, dp_size=dp_size) + return [Box(ray.put(rollout_data)) for rollout_data in rollout_data_list] + + +def split_train_data_by_dp_raw(args, data: dict[str, Any], *, dp_size: int) -> list[dict[str, Any]]: """Split the train data by data parallel size.""" rollout_data = {} @@ -136,7 +142,7 @@ def split_train_data_by_dp(args, data, dp_size): else: partitions = [range(i, len(total_lengths), dp_size) for i in range(dp_size)] - rollout_data_refs = [] + ans = [] for i in range(dp_size): rollout_data = {} @@ -157,6 +163,7 @@ def split_train_data_by_dp(args, data, dp_size): "prompt", "teacher_log_probs", "opd_reverse_kl", + "seq_witness_ids", "weight_versions", ]: if key not in data: @@ -172,5 +179,5 @@ def split_train_data_by_dp(args, data, dp_size): if key not in data: continue rollout_data[key] = data[key] - rollout_data_refs.append(Box(ray.put(rollout_data))) - return rollout_data_refs + ans.append(rollout_data) + return ans diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index 49d99898c76..7ff03b77878 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -270,6 +270,11 @@ def add_train_arguments(parser): parser.add_argument( "--log-probs-chunk-size", type=int, default=-1, help="Chunk size to compute log probs to save memory" ) + parser.add_argument( + "--delay-split-train-data-by-dp", + action="store_true", + default=False, + ) parser.add_argument( "--allgather-cp", action="store_true", diff --git a/miles/utils/data.py b/miles/utils/data.py index 0fe7ad483fd..21d01db972b 100644 --- a/miles/utils/data.py +++ b/miles/utils/data.py @@ -8,6 +8,8 @@ import numpy as np import ray +from miles.ray.rollout.train_data_conversion import split_train_data_by_dp_raw + try: import pyarrow.parquet as pq except ImportError: @@ -272,9 +274,19 @@ def get_minimum_num_micro_batch_size(total_lengths, max_tokens_per_gpu): return len(batches) -def process_rollout_data(args, rollout_data_ref, dp_rank, dp_size): - assert len(rollout_data_ref) == dp_size - rollout_data = ray.get(rollout_data_ref[dp_rank].inner) +def process_rollout_data( + args, + rollout_data_ref, + dp_rank, + dp_size, +): + if args.delay_split_train_data_by_dp: + raw = ray.get(rollout_data_ref.inner) + raw = split_train_data_by_dp_raw(args, raw, dp_size=dp_size) + rollout_data = raw[dp_rank] + else: + assert len(rollout_data_ref) == dp_size + rollout_data = ray.get(rollout_data_ref[dp_rank].inner) partition = rollout_data.pop("partition") total_lengths = rollout_data["total_lengths"] diff --git a/tests/fast/ray/rollout/real_ray/test_rollout_manager.py b/tests/fast/ray/rollout/real_ray/test_rollout_manager.py index ab7d79b75df..962a53099dc 100644 --- a/tests/fast/ray/rollout/real_ray/test_rollout_manager.py +++ b/tests/fast/ray/rollout/real_ray/test_rollout_manager.py @@ -516,9 +516,12 @@ def fake_rollout_fn(input): assert len(captured) == 1 assert isinstance(captured[0], RolloutFnTrainInput) assert captured[0].rollout_id == 42 + # generate returns {"sample_indices": ..., "data_ref": ...}; # split_train_data_by_dp returns Box(ObjectRef) per dp rank - assert len(result) == 2 - partitions = ray.get([box.inner for box in result]) + assert set(result) == {"sample_indices", "data_ref"} + data_refs = result["data_ref"] + assert len(data_refs) == 2 + partitions = ray.get([box.inner for box in data_refs]) for partition in partitions: assert "tokens" in partition assert "rewards" in partition diff --git a/tests/fast/ray/rollout/test_train_data_conversion.py b/tests/fast/ray/rollout/test_train_data_conversion.py index 4a3da1a2883..31f3b1975db 100644 --- a/tests/fast/ray/rollout/test_train_data_conversion.py +++ b/tests/fast/ray/rollout/test_train_data_conversion.py @@ -1,7 +1,10 @@ from __future__ import annotations +from unittest.mock import MagicMock + import pytest import ray +import torch from hypothesis import HealthCheck, given, settings from hypothesis import strategies as st from tests.fast.ray.rollout.conftest import make_args, make_sample, make_samples_grouped @@ -10,6 +13,7 @@ _post_process_rewards, convert_samples_to_train_data, split_train_data_by_dp, + split_train_data_by_dp_raw, ) from miles.utils.types import Sample @@ -472,3 +476,66 @@ def test_partition_indices_form_a_partition(self): parts = [ray.get(r.inner) for r in refs] all_indices = sorted(i for p in parts for i in p["partition"]) assert all_indices == list(range(n)) + + +class TestSplitTrainDataRaw: + def test_witness_ids_split_across_dp(self) -> None: + tokens = [[1, 2, 3], [4, 5], [6, 7, 8, 9], [10, 11]] + witness_ids = [ + torch.tensor([0, 0, 0]), + torch.tensor([1, 1]), + torch.tensor([2, 2, 2, 2]), + torch.tensor([3, 3]), + ] + + data = { + "tokens": tokens, + "seq_witness_ids": witness_ids, + "response_lengths": [1, 1, 1, 1], + "loss_masks": [[0, 0, 1], [0, 1], [0, 0, 0, 1], [0, 1]], + } + + args = MagicMock() + args.balance_data = False + + result = split_train_data_by_dp_raw(args, data, dp_size=2) + + assert len(result) == 2 + assert "seq_witness_ids" in result[0] + assert "seq_witness_ids" in result[1] + assert len(result[0]["seq_witness_ids"]) == 2 + assert len(result[1]["seq_witness_ids"]) == 2 + + def test_indexer_topk_and_opd_reverse_kl_split_across_dp(self) -> None: + """Keys from the rollout-side split (rollout_indexer_topk, opd_reverse_kl) partition per sample.""" + data = { + "tokens": [[1, 2], [3, 4], [5, 6], [7, 8]], + "response_lengths": [1, 1, 1, 1], + "loss_masks": [[0, 1], [0, 1], [0, 1], [0, 1]], + "rollout_indexer_topk": [torch.tensor([i]) for i in range(4)], + "opd_reverse_kl": [[float(i)] for i in range(4)], + } + + args = MagicMock() + args.balance_data = False + + result = split_train_data_by_dp_raw(args, data, dp_size=2) + + assert len(result) == 2 + for part in result: + assert len(part["rollout_indexer_topk"]) == 2 + assert len(part["opd_reverse_kl"]) == 2 + + def test_no_witness_ids_when_absent(self) -> None: + tokens = [[1, 2], [3, 4]] + data = { + "tokens": tokens, + "response_lengths": [1, 1], + "loss_masks": [[0, 1], [0, 1]], + } + + args = MagicMock() + args.balance_data = False + + result = split_train_data_by_dp_raw(args, data, dp_size=1) + assert "seq_witness_ids" not in result[0] From d0b41fc54e20daf02659f1138d673461a243849f Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Mon, 22 Jun 2026 18:15:17 +0800 Subject: [PATCH 15/41] Add a deterministic NCCL backend for order-stable collectives Add an opt-in deterministic NCCL process-group backend (`--debug-deterministic-collective`) that folds order-sensitive SUM/AVG reductions into a fixed order so training collectives are bit-reproducible, registering it as the training world's distributed backend and requiring synchronous grad sync. - det_process_group.py (+ GPU test, dist test helper). - train_actor.py: register the backend and select it when enabled. - initialize.py: assert synchronous grad reduce under the deterministic backend. --- miles/backends/megatron_utils/initialize.py | 3 + miles/ray/train_actor.py | 6 + miles/utils/arguments.py | 9 + .../debug_utils/run_megatron/worker/main.py | 1 + miles/utils/det_process_group.py | 317 ++++++++ tests/fast-gpu/test_det_process_group.py | 761 ++++++++++++++++++ 6 files changed, 1097 insertions(+) create mode 100644 miles/utils/det_process_group.py create mode 100644 tests/fast-gpu/test_det_process_group.py diff --git a/miles/backends/megatron_utils/initialize.py b/miles/backends/megatron_utils/initialize.py index 00b64c15698..d47859eca8e 100644 --- a/miles/backends/megatron_utils/initialize.py +++ b/miles/backends/megatron_utils/initialize.py @@ -97,6 +97,9 @@ def init(args): torch.backends.cudnn.benchmark = False torch.use_deterministic_algorithms(True, warn_only=False) + if args.debug_deterministic_collective: + assert not args.overlap_grad_reduce, "deterministic collectives require synchronous grad sync" + if args.tp_comm_overlap: from megatron.training.initialize import _initialize_tp_communicators diff --git a/miles/ray/train_actor.py b/miles/ray/train_actor.py index 7ab2bfcd745..971d8a3a266 100644 --- a/miles/ray/train_actor.py +++ b/miles/ray/train_actor.py @@ -11,6 +11,7 @@ import miles.utils.eval_config from miles.ray.ray_actor import RayActor +from miles.utils.det_process_group import DET_NCCL_BACKEND_NAME, register_det_nccl_backend from miles.utils.distributed_utils import init_gloo_group from miles.utils.env_report import collect_and_print_node_env_report from miles.utils.logging_utils import configure_logger @@ -73,6 +74,11 @@ def init(self, args, role, with_ref=False, with_opd_teacher=False): local_rank = int(os.environ.get("LOCAL_RANK", 0)) torch.cuda.set_device(f"cuda:{local_rank}") + if args.debug_deterministic_collective: + register_det_nccl_backend() + args.distributed_backend = DET_NCCL_BACKEND_NAME + logger.info("Deterministic collectives: training world uses the det_nccl backend") + # Use hybrid backend when FSDP CPU offload is enabled with a CPU backend backend = args.distributed_backend if getattr(args, "fsdp_cpu_offload", False) and getattr(args, "fsdp_cpu_backend", None): diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index 7ff03b77878..2348dfa1e53 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -1652,6 +1652,15 @@ def add_debug_arguments(parser): default=os.environ.get("MILES_SCRIPT_ENV_REPORT", ""), help="JSON string containing environment report from external launcher.", ) + parser.add_argument( + "--debug-deterministic-collective", + action="store_true", + default=False, + help="Debug/test only: run the training world on the det_nccl backend " + "(miles.utils.det_process_group), which folds order-sensitive SUM/AVG " + "reductions in a fixed tree order so different reduction topologies become " + "bitwise-comparable. Slow; never enable in production.", + ) return parser def add_network_arguments(parser): diff --git a/miles/utils/debug_utils/run_megatron/worker/main.py b/miles/utils/debug_utils/run_megatron/worker/main.py index 359a51c1d60..cdfbaf956e1 100644 --- a/miles/utils/debug_utils/run_megatron/worker/main.py +++ b/miles/utils/debug_utils/run_megatron/worker/main.py @@ -124,6 +124,7 @@ def _initialize_megatron(args: argparse.Namespace) -> None: args.hf_checkpoint = str(args.script_hf_checkpoint) args.__dict__.setdefault("megatron_to_hf_mode", "raw") args.__dict__.setdefault("decrease_batch_size_if_needed", False) + args.__dict__.setdefault("debug_deterministic_collective", False) set_default_megatron_args(args) validate_args(args) diff --git a/miles/utils/det_process_group.py b/miles/utils/det_process_group.py new file mode 100644 index 00000000000..a2fafddf2ef --- /dev/null +++ b/miles/utils/det_process_group.py @@ -0,0 +1,317 @@ +"""Process group with bitwise-deterministic SUM reductions. + +``DetProcessGroup`` wraps an inner c10d NCCL group. Every collective delegates to +the inner group except the order-sensitive reductions — ``allreduce`` and +``reduce_scatter`` — which are computed as all-gather (pure data movement, no +arithmetic) plus a fixed local fold: a pairwise +tree for power-of-two world sizes, an ascending-rank fold otherwise. The summation +order is therefore independent of the NCCL version, topology, communicator +instance, or buffer layout, and reduce-scatter takes its shard from the same full +fold, so reduce-scatter and all-reduce agree bitwise by construction. + +Debug/test use only: the fold trades bandwidth and synchrony for determinism. +""" + +import logging + +import torch +import torch.distributed as dist +from torch.distributed import ProcessGroup as BaseProcessGroup +from torch.distributed import Work +from torch.distributed.distributed_c10d import AllgatherOptions + +logger = logging.getLogger(__name__) + + +DET_NCCL_BACKEND_NAME = "det_nccl" +_backend_registered = False + +# Cap on the gather buffer so the fold never allocates world_size x the full tensor +# (a 30B MoE grad buffer x 8 ranks is >100 GiB). Chunking is bitwise-neutral: each +# output element is still folded over the same per-rank operands in the same order. +_GATHER_BUFFER_CAP_BYTES = 1 << 30 + + +def register_det_nccl_backend() -> None: + """Register the ``det_nccl`` torch.distributed backend (DetProcessGroup over NCCL). + + After registration, groups created with ``backend="det_nccl"`` (via + ``init_process_group`` or ``new_group``) run SUM/AVG reductions through the + deterministic fold. Idempotent. + """ + global _backend_registered + if _backend_registered: + return + dist.Backend.register_backend(DET_NCCL_BACKEND_NAME, _create_det_nccl_backend, extended_api=True, devices=["cuda"]) + _backend_registered = True + logger.info("Registered torch.distributed backend %s", DET_NCCL_BACKEND_NAME) + + +def _create_det_nccl_backend(dist_backend_opts: object, pg_options: object) -> "DetProcessGroup": + from torch.distributed import ProcessGroupNCCL + + inner = ProcessGroupNCCL( + dist_backend_opts.store, + dist_backend_opts.group_rank, + dist_backend_opts.group_size, + ProcessGroupNCCL.Options(), + ) + return DetProcessGroup(inner) + + +class DetProcessGroup(BaseProcessGroup): + """Wrapper process group whose SUM/AVG reductions use a fixed-order fold.""" + + def __init__(self, inner: dist.ProcessGroup) -> None: + super().__init__(inner.rank(), inner.size()) + self._inner = inner + # Register the cuda backend as torchft's ProcessGroupNCCL._create_pg does, so + # _device_types reports cuda and object collectives pick cuda over cpu. Only + # feeds device/backend lookups; collectives still dispatch to our methods. + self._set_default_backend(BaseProcessGroup.BackendType.CUSTOM) + self._register_backend(torch.device("cuda"), BaseProcessGroup.BackendType.CUSTOM, self._inner) + + # ------------------------------------------------------------------ # + # Deterministic reductions + # ------------------------------------------------------------------ # + + def allreduce(self, tensors: list[torch.Tensor], opts: object) -> Work: + reduce_op = _reduce_op_of(opts) + if reduce_op == dist.ReduceOp.MAX or reduce_op == dist.ReduceOp.MIN: + return self._inner.allreduce(tensors, opts) + + for tensor in tensors: + det_all_reduce(tensor, group=self._inner, reduce_op=reduce_op) + return _CompletedWork() + + def allreduce_coalesced(self, tensors: list[torch.Tensor], opts: object) -> Work: + return self.allreduce(tensors, opts) + + def _reduce_scatter_base(self, output: torch.Tensor, input: torch.Tensor, opts: object) -> Work: + reduce_op = _reduce_op_of(opts) + if reduce_op == dist.ReduceOp.MAX or reduce_op == dist.ReduceOp.MIN: + return self._inner._reduce_scatter_base(output, input, opts) + + det_reduce_scatter( + output, input, group=self._inner, rank=self.rank(), world_size=self.size(), reduce_op=reduce_op + ) + return _CompletedWork() + + def reduce_scatter( + self, output_tensors: list[torch.Tensor], input_tensors: list[list[torch.Tensor]], opts: object + ) -> Work: + reduce_op = _reduce_op_of(opts) + if reduce_op == dist.ReduceOp.MAX or reduce_op == dist.ReduceOp.MIN: + return self._inner.reduce_scatter(output_tensors, input_tensors, opts) + + for output, inputs in zip(output_tensors, input_tensors, strict=True): + # Slot j has one size on every rank (sizes may differ between slots); + # every rank joins each slot's fold, only rank j keeps the result. + assert ( + inputs[self.rank()].numel() == output.numel() + ), f"slot {self.rank()} numel {inputs[self.rank()].numel()} != output numel {output.numel()}" + for slot_idx, slot_input in enumerate(inputs): + _det_reduce_scatter_slot( + output if slot_idx == self.rank() else None, slot_input, group=self._inner, reduce_op=reduce_op + ) + return _CompletedWork() + + # ------------------------------------------------------------------ # + # Plain delegation + # ------------------------------------------------------------------ # + + def allgather( + self, output_tensors: list[list[torch.Tensor]], input_tensors: list[torch.Tensor], opts: object + ) -> Work: + return self._inner.allgather(output_tensors, input_tensors, opts) + + def allgather_into_tensor_coalesced( + self, output_tensors: list[torch.Tensor], input_tensors: list[torch.Tensor], opts: object = None + ) -> Work: + # The coalescing manager's flush passes no opts; inner lacks the coalesced form. + effective_opts = opts if opts is not None else AllgatherOptions() + for output, input in zip(output_tensors, input_tensors, strict=True): + self._inner._allgather_base(output, input, effective_opts).wait() + return _CompletedWork() + + def _allgather_base(self, output: torch.Tensor, input: torch.Tensor, opts: object) -> Work: + return self._inner._allgather_base(output, input, opts) + + def barrier(self, opts: object) -> Work: + return self._inner.barrier(opts) + + def broadcast(self, tensor_list: list[torch.Tensor], opts: object) -> Work: + return self._inner.broadcast(tensor_list, opts) + + def reduce(self, tensors: list[torch.Tensor], opts: object) -> Work: + reduce_op = _reduce_op_of(opts) + if reduce_op == dist.ReduceOp.MAX or reduce_op == dist.ReduceOp.MIN: + return self._inner.reduce(tensors, opts) + return self.allreduce(tensors, opts) + + def reduce_scatter_tensor_coalesced( + self, output_tensors: list[torch.Tensor], input_tensors: list[torch.Tensor], opts: object + ) -> Work: + for output, input in zip(output_tensors, input_tensors, strict=True): + self._reduce_scatter_base(output, input, opts) + return _CompletedWork() + + def alltoall_base( + self, + output_tensor: torch.Tensor, + input_tensor: torch.Tensor, + output_split_sizes: list[int], + input_split_sizes: list[int], + opts: object, + ) -> Work: + return self._inner.alltoall_base(output_tensor, input_tensor, output_split_sizes, input_split_sizes, opts) + + def send(self, tensors: list[torch.Tensor], dst_rank: int, tag: int) -> Work: + return self._inner.send(tensors, dst_rank, tag) + + def recv(self, tensors: list[torch.Tensor], src_rank: int, tag: int) -> Work: + return self._inner.recv(tensors, src_rank, tag) + + def _start_coalescing(self, device: torch.device) -> None: + # Ops queue at the Python level and flush via the *_coalesced methods. + return None + + def _end_coalescing(self, device: torch.device) -> Work: + return _CompletedWork() + + def getBackendName(self) -> str: + return DET_NCCL_BACKEND_NAME + + +def det_all_reduce(tensor: torch.Tensor, *, group: dist.ProcessGroup, reduce_op: object = dist.ReduceOp.SUM) -> None: + """SUM/AVG ``tensor`` across ranks in-place with the fixed fold. + + ``group`` may be a raw c10d backend (flat ``_allgather_base``) or any ``ProcessGroup`` + such as torchft's wrappers (list-form ``allgather``): the gather is pure data + movement and the local fold defines the (shared) summation order. + """ + if not tensor.is_contiguous(): + work = tensor.contiguous() + det_all_reduce(work, group=group, reduce_op=reduce_op) + tensor.copy_(work) + return + + flat = tensor.view(-1) + _det_chunked_fold(flat, flat, group=group) + if reduce_op == dist.ReduceOp.AVG: + flat.div_(group.size()) + + +def det_reduce_scatter( + output: torch.Tensor, + input: torch.Tensor, + *, + group: dist.ProcessGroup, + rank: int, + world_size: int, + reduce_op: object = dist.ReduceOp.SUM, +) -> None: + """SUM/AVG-reduce ``input`` across ranks with the fixed fold and write this rank's + ``1/world_size`` slice into ``output`` (mirrors ``dist.reduce_scatter_tensor``). + """ + assert ( + input.numel() == world_size * output.numel() + ), f"uneven reduce_scatter: input numel {input.numel()} != {world_size} x output numel {output.numel()}" + + flat = input.contiguous().view(-1) + slot_numel = output.numel() + for slot_idx in range(world_size): + slot_input = flat[slot_idx * slot_numel : (slot_idx + 1) * slot_numel] + _det_reduce_scatter_slot(output if slot_idx == rank else None, slot_input, group=group, reduce_op=reduce_op) + + +def _det_reduce_scatter_slot( + output: torch.Tensor | None, slot_input: torch.Tensor, *, group: dist.ProcessGroup, reduce_op: object +) -> None: + """Fold one slot across ranks into ``output``; ``None`` joins the collective only.""" + flat_input = slot_input.contiguous().view(-1) + if output is None: + _det_chunked_fold(flat_input, None, group=group) + return + + # NOTE: empty_like preserves the (non-contiguous) layout, so it cannot be view(-1)'d; + # allocate the staging buffer flat directly. + out_flat = ( + output.view(-1) + if output.is_contiguous() + else torch.empty(output.numel(), dtype=output.dtype, device=output.device) + ) + _det_chunked_fold(flat_input, out_flat, group=group) + if not output.is_contiguous(): + output.copy_(out_flat.view(output.shape)) + if reduce_op == dist.ReduceOp.AVG: + output.div_(group.size()) + + +def _det_chunked_fold( + flat_input: torch.Tensor, + out_flat: torch.Tensor | None, + *, + group: dist.ProcessGroup, +) -> None: + """Fold ``flat_input`` across ranks chunk by chunk into ``out_flat`` (same numel; + may alias ``flat_input``). ``None`` joins the gathers without folding. + """ + world_size = group.size() + total = flat_input.numel() + chunk_numel = max(1, min(total, _GATHER_BUFFER_CAP_BYTES // (world_size * flat_input.element_size()))) + gather_buf = torch.empty(world_size * chunk_numel, dtype=flat_input.dtype, device=flat_input.device) + + for start in range(0, total, chunk_numel): + count = min(chunk_numel, total - start) + buf = gather_buf[: world_size * count] + _gather_into(group, buf, flat_input[start : start + count]) + if out_flat is not None: + folded = _fold_gathered_sum(list(buf.view(world_size, count).unbind(dim=0))) + out_flat[start : start + count].copy_(folded) + + +def _gather_into(group: dist.ProcessGroup, output: torch.Tensor, input: torch.Tensor) -> None: + if isinstance(group, dist.ProcessGroup): + # ProcessGroup wrappers (torchft) inherit ``_allgather_base`` from the C++ base, but it + # dispatches to a per-device backend they never register; only the overridden list-form + # ``allgather`` is safe. ``hasattr`` cannot discriminate here. + rows = list(output.view(group.size(), -1).unbind(dim=0)) + group.allgather([rows], [input], AllgatherOptions()).wait() + else: + group._allgather_base(output, input, AllgatherOptions()).wait() + + +def _reduce_op_of(opts: object) -> object: + """Extract the ReduceOp from an options object (or pass a bare ReduceOp through).""" + return opts.reduceOp if hasattr(opts, "reduceOp") else opts + + +class _CompletedWork(Work): + """Work handle for an operation that already completed synchronously.""" + + def wait(self, timeout: object = None) -> bool: + return True + + def get_future(self) -> torch.futures.Future: + future: torch.futures.Future = torch.futures.Future() + future.set_result(None) + return future + + +def _fold_gathered_sum(gathered: list[torch.Tensor]) -> torch.Tensor: + """Sum a per-rank gathered list in a fixed order (pairwise tree for power-of-two). + + May reuse (mutate) the gathered buffers as accumulators. + """ + world_size = len(gathered) + if world_size > 0 and (world_size & (world_size - 1)) == 0: + partials = gathered + while len(partials) > 1: + partials = [partials[i] + partials[i + 1] for i in range(0, len(partials), 2)] + return partials[0] + + running = gathered[0] + for index in range(1, world_size): + running += gathered[index] + return running diff --git a/tests/fast-gpu/test_det_process_group.py b/tests/fast-gpu/test_det_process_group.py new file mode 100644 index 00000000000..6e123862e13 --- /dev/null +++ b/tests/fast-gpu/test_det_process_group.py @@ -0,0 +1,761 @@ +from tests.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=240, suite="stage-c-4-gpu-h200", labels=[]) + +import os +import socket + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +from torch.distributed.distributed_c10d import AllreduceOptions, ReduceScatterOptions, _coalescing_manager + +from miles.utils.det_process_group import ( + DET_NCCL_BACKEND_NAME, + _CompletedWork, + _fold_gathered_sum, + _reduce_op_of, + det_all_reduce, + register_det_nccl_backend, +) + +_WORLD_SIZE = 4 +_NUMEL = 1_048_576 +_SEED = 1234 + + +# --------------------------------------------------------------------------- # +# CPU-only tests (no GPU, no distributed init) +# --------------------------------------------------------------------------- # + + +class _FakeFlatGroup: + """CPU stand-in with the c10d backend gather (_allgather_base); serves ascending chunk requests.""" + + def __init__(self, per_rank: list[torch.Tensor]) -> None: + self._per_rank = per_rank + self._offset = 0 + + def size(self) -> int: + return len(self._per_rank) + + def _allgather_base(self, output: torch.Tensor, input: torch.Tensor, opts: object) -> _CompletedWork: + count = input.numel() + rows = output.view(self.size(), count) + for row, src in zip(rows, self._per_rank, strict=True): + row.copy_(src.reshape(-1)[self._offset : self._offset + count]) + self._offset += count + return _CompletedWork() + + +class _FakeTorchftGroup(dist.ProcessGroup): + """CPU stand-in mirroring torchft: a ProcessGroup subclass overriding only the list-form allgather + (so it inherits a broken ``_allgather_base`` from the C++ base, exactly like the real wrappers); + serves ascending chunk requests.""" + + def __init__(self, per_rank: list[torch.Tensor]) -> None: + super().__init__(0, len(per_rank)) + self._per_rank = per_rank + self._offset = 0 + + def allgather( + self, output_lists: list[list[torch.Tensor]], input_list: list[torch.Tensor], opts: object + ) -> _CompletedWork: + count = input_list[0].numel() + for row, src in zip(output_lists[0], self._per_rank, strict=True): + row.copy_(src.reshape(-1)[self._offset : self._offset + count]) + self._offset += count + return _CompletedWork() + + +def _pairwise_tree_fold(partials: list[torch.Tensor]) -> torch.Tensor: + """Inline reference fold (pairwise tree for power-of-two): independent of the module.""" + running = list(partials) + while len(running) > 1: + running = [running[i] + running[i + 1] for i in range(0, len(running), 2)] + return running[0] + + +def test_det_nccl_backend_name_constant_value(): + """DET_NCCL_BACKEND_NAME is the literal "det_nccl" the backend registers and reports under.""" + assert DET_NCCL_BACKEND_NAME == "det_nccl" + + +def test_reduceop_equality_vs_containment_footgun(): + """Documents why dispatch uses explicit ==: an options ReduceOp equals SUM yet tuple containment is False.""" + opts = dist.AllreduceOptions() + opts.reduceOp = dist.ReduceOp.SUM + + assert opts.reduceOp == dist.ReduceOp.SUM + assert opts.reduceOp not in (dist.ReduceOp.SUM, dist.ReduceOp.AVG) + + +@pytest.mark.parametrize( + "parts,dtype,expected", + [ + # Single rank / order-free sanity. + ([3.5], torch.float64, 3.5), + ([1.0, 1e-5], torch.float64, 1.0 + 1e-5), + ([1, 2, 3, 4], torch.int64, 10), + # One 2**-53 is absorbed by 1.0 (ties-to-even), but 2**-53 + 2**-53 = ulp(1.0) + # is not -- so the pairwise tree and a sequential fold give different bits. + # 4 ranks must be the tree (a+b)+(c+d), not sequential (which gives 1.0): + ([0.5, 0.5, 2**-53, 2**-53], torch.float64, (0.5 + 0.5) + (2**-53 + 2**-53)), + # 3 ranks (non-power-of-two) must be the ascending fold ((a+b)+c), which + # absorbs both halves (a wrong pairing would give 1.0 + 2**-52): + ([1.0, 2**-53, 2**-53], torch.float64, (1.0 + 2**-53) + 2**-53), + # 5 ranks: ascending fold absorbs every half: + ([1.0, 2**-53, 2**-53, 2**-53, 2**-53], torch.float64, 1.0), + # 8 ranks: the full three-level tree cancels exactly to 0.0 (sequential + # would leave -2**-52): + ( + [0.5, 0.5, 2**-53, 2**-53, -0.5, -0.5, -(2**-53), -(2**-53)], + torch.float64, + ((0.5 + 0.5) + (2**-53 + 2**-53)) + ((-0.5 + -0.5) + (-(2**-53) + -(2**-53))), + ), + # Same construction at float32 precision (ulp(1.0) = 2**-23): + ([0.5, 0.5, 2**-24, 2**-24], torch.float32, (0.5 + 0.5) + (2**-24 + 2**-24)), + ], +) +def test_fold_gathered_sum(parts: list[float], dtype: torch.dtype, expected: float): + """The fold's exact bracketing is pinned by hand-computed literals where any other order changes the bits.""" + tensors = [torch.tensor([value], dtype=dtype) for value in parts] + + actual = _fold_gathered_sum(tensors) + + assert actual.item() == expected + + +def test_reduce_op_of_extracts_reduceop_from_options_object(): + """_reduce_op_of reads .reduceOp from an options object and passes a bare ReduceOp through.""" + ar_opts = AllreduceOptions() + ar_opts.reduceOp = dist.ReduceOp.SUM + assert _reduce_op_of(ar_opts) == dist.ReduceOp.SUM + + rs_opts = ReduceScatterOptions() + rs_opts.reduceOp = dist.ReduceOp.AVG + assert _reduce_op_of(rs_opts) == dist.ReduceOp.AVG + + assert _reduce_op_of(dist.ReduceOp.MAX) == dist.ReduceOp.MAX + + +def test_det_all_reduce_equals_manual_pairwise_tree_fold(): + """det_all_reduce over a fake group equals the manual pairwise-tree fold bitwise.""" + gen = torch.Generator().manual_seed(_SEED) + per_rank = [torch.randn(16, generator=gen, dtype=torch.float32) for _ in range(_WORLD_SIZE)] + expected = _pairwise_tree_fold(per_rank) + + tensor = per_rank[0].clone() + det_all_reduce(tensor, group=_FakeFlatGroup(per_rank)) + assert torch.equal(tensor, expected) + + +def test_det_all_reduce_avg_equals_tree_fold_divided_by_world(): + """det_all_reduce with reduce_op=AVG equals the pairwise-tree fold divided by world_size bitwise.""" + gen = torch.Generator().manual_seed(_SEED + 31) + per_rank = [torch.randn(16, generator=gen, dtype=torch.float32) for _ in range(_WORLD_SIZE)] + expected = _pairwise_tree_fold(per_rank) / _WORLD_SIZE + + tensor = per_rank[0].clone() + det_all_reduce(tensor, group=_FakeFlatGroup(per_rank), reduce_op=dist.ReduceOp.AVG) + assert torch.equal(tensor, expected) + + +def test_det_all_reduce_avg_non_contiguous_recursion_passes_reduce_op_through(): + """A non-contiguous input with reduce_op=AVG recurses with the op preserved, writing tree/world back.""" + gen = torch.Generator().manual_seed(_SEED + 32) + per_rank = [torch.randn(8 * 4, generator=gen, dtype=torch.float32) for _ in range(_WORLD_SIZE)] + expected_flat = _pairwise_tree_fold(per_rank) / _WORLD_SIZE + + base = per_rank[0].reshape(4, 8).clone() + non_contiguous = base.t() + assert not non_contiguous.is_contiguous() + + det_all_reduce(non_contiguous, group=_FakeFlatGroup(per_rank), reduce_op=dist.ReduceOp.AVG) + assert torch.equal(non_contiguous.contiguous().reshape(-1), expected_flat) + + +def test_det_all_reduce_multi_chunk_matches_single_chunk(monkeypatch: pytest.MonkeyPatch): + """A tiny gather-buffer cap (forcing many chunks) gives bitwise-identical results to one chunk.""" + import miles.utils.det_process_group as dpg + + gen = torch.Generator().manual_seed(_SEED + 21) + per_rank = [torch.randn(50, generator=gen, dtype=torch.float32) for _ in range(_WORLD_SIZE)] + expected = _pairwise_tree_fold(per_rank) + + monkeypatch.setattr(dpg, "_GATHER_BUFFER_CAP_BYTES", _WORLD_SIZE * 7 * 4) + tensor = per_rank[0].clone() + det_all_reduce(tensor, group=_FakeFlatGroup(per_rank)) + + assert torch.equal(tensor, expected) + + +def test_det_reduce_scatter_multi_chunk_slice_matches_full_fold(monkeypatch: pytest.MonkeyPatch): + """det_reduce_scatter under a tiny chunk cap writes exactly this rank's slice of the full fold.""" + import miles.utils.det_process_group as dpg + + gen = torch.Generator().manual_seed(_SEED + 22) + per_rank = [torch.randn(48, generator=gen, dtype=torch.float32) for _ in range(_WORLD_SIZE)] + expected_full = _pairwise_tree_fold(per_rank) + + monkeypatch.setattr(dpg, "_GATHER_BUFFER_CAP_BYTES", _WORLD_SIZE * 7 * 4) + rank = 1 + out = torch.empty(12, dtype=torch.float32) + dpg.det_reduce_scatter(out, per_rank[0].clone(), group=_FakeFlatGroup(per_rank), rank=rank, world_size=_WORLD_SIZE) + + assert torch.equal(out, expected_full[12:24]) + + +def test_det_all_reduce_torchft_list_gather_matches_flat_gather_bitwise(): + """The torchft list-form gather path is bitwise-identical to the _allgather_base path.""" + gen = torch.Generator().manual_seed(_SEED + 7) + per_rank = [torch.randn(32, generator=gen, dtype=torch.float32) for _ in range(_WORLD_SIZE)] + + via_flat = per_rank[0].clone() + det_all_reduce(via_flat, group=_FakeFlatGroup(per_rank)) + + via_list = per_rank[0].clone() + det_all_reduce(via_list, group=_FakeTorchftGroup(per_rank)) + + assert torch.equal(via_flat, via_list) + assert torch.equal(via_flat, _pairwise_tree_fold(per_rank)) + + +def test_gather_into_routes_process_group_instances_to_list_form_allgather(): + """Regression: ProcessGroup subclasses inherit _allgather_base from the C++ base (hasattr is + always True), so routing must key on isinstance and take the overridden list-form allgather.""" + import miles.utils.det_process_group as dpg + + gen = torch.Generator().manual_seed(_SEED + 23) + per_rank = [torch.randn(8, generator=gen, dtype=torch.float32) for _ in range(_WORLD_SIZE)] + group = _FakeTorchftGroup(per_rank) + assert hasattr(group, "_allgather_base") + + out = torch.empty(_WORLD_SIZE * 8, dtype=torch.float32) + dpg._gather_into(group, out, per_rank[0].clone()) + + assert torch.equal(out.view(_WORLD_SIZE, 8), torch.stack(per_rank)) + + +def test_det_all_reduce_non_contiguous_input_writes_summed_values_back(): + """A non-contiguous (.t() view) input gets the correct summed values written back.""" + gen = torch.Generator().manual_seed(_SEED + 11) + per_rank = [torch.randn(8 * 4, generator=gen, dtype=torch.float32) for _ in range(_WORLD_SIZE)] + expected_flat = _pairwise_tree_fold(per_rank) + + base = per_rank[0].reshape(4, 8).clone() + non_contiguous = base.t() + assert not non_contiguous.is_contiguous() + + det_all_reduce(non_contiguous, group=_FakeFlatGroup(per_rank)) + assert torch.equal(non_contiguous.contiguous().reshape(-1), expected_flat) + + +def test_det_all_reduce_world_size_one_leaves_tensor_unchanged(): + """A 1-rank group gathers only the local copy, so the tensor is unchanged.""" + original = torch.tensor([1.0, -2.5, 3.25, 0.0], dtype=torch.float32) + tensor = original.clone() + det_all_reduce(tensor, group=_FakeFlatGroup([original])) + assert torch.equal(tensor, original) + + +def test_det_all_reduce_fold_order_checksum_pin(): + """Pin exact fold result so an accidental fold-order change (tree->linear) fails loudly.""" + per_rank = [ + torch.tensor([1.0, 1e8, -1e8, 0.25], dtype=torch.float32), + torch.tensor([2.0, -1e8, 1e8, 0.25], dtype=torch.float32), + torch.tensor([3.0, 1e8, -1e8, 0.25], dtype=torch.float32), + torch.tensor([4.0, -1e8, 1e8, 0.25], dtype=torch.float32), + ] + reference = _pairwise_tree_fold(per_rank) + + tensor = per_rank[0].clone() + det_all_reduce(tensor, group=_FakeFlatGroup(per_rank)) + assert torch.equal(tensor, reference) + + # Hardcoded pin: element 0 is the plain sum; element 3 sums four 0.25 -> 1.0. + assert torch.equal(tensor, torch.tensor([10.0, 0.0, 0.0, 1.0], dtype=torch.float32)) + assert tensor[0].item().hex() == (10.0).hex() + assert tensor[3].item().hex() == (1.0).hex() + + +def test_completed_work_future_wait_returns_result(): + """_CompletedWork().get_future().wait() returns the (None) result without blocking.""" + work = _CompletedWork() + assert work.wait() is True + assert work.get_future().wait() is None + + +def _order_sensitive_input(rank: int, seed: int = _SEED) -> torch.Tensor: + """Per-rank input whose cross-rank sum catastrophically cancels (~1e-4 from +-0.5).""" + shared = torch.randn(_NUMEL, generator=torch.Generator().manual_seed(seed), dtype=torch.float32) + own = torch.randn(_NUMEL, generator=torch.Generator().manual_seed(seed + 1 + rank), dtype=torch.float32) + sign = -1.0 if rank % 2 else 1.0 + return (sign * 0.5 * shared + 1e-4 * own).cuda() + + +def _manual_tree_sum(partials: list[torch.Tensor]) -> torch.Tensor: + running = list(partials) + while len(running) > 1: + running = [running[i] + running[i + 1] for i in range(0, len(running), 2)] + return running[0] + + +def _fixed_tree_reference(x: torch.Tensor) -> torch.Tensor: + """Gather every rank's tensor (data movement only) and fold in the fixed tree order.""" + gathered = [torch.empty_like(x) for _ in range(_WORLD_SIZE)] + dist.all_gather(gathered, x) + return _manual_tree_sum(gathered) + + +def _assert_bitwise(name: str, actual: torch.Tensor, expected: torch.Tensor) -> None: + if torch.equal(actual, expected): + return + mismatch = int((actual != expected).sum().item()) + max_abs = float((actual - expected).abs().max().item()) + raise AssertionError(f"{name}: mismatch_elems={mismatch}/{actual.numel()} max_abs={max_abs:.3e}") + + +def _shard_of(full: torch.Tensor, rank: int) -> torch.Tensor: + shard_numel = full.numel() // _WORLD_SIZE + return full[rank * shard_numel : (rank + 1) * shard_numel] + + +def _check_allreduce(rank: int, det1: dist.ProcessGroup, det2: dist.ProcessGroup, x, tree) -> None: + a = x.clone() + dist.all_reduce(a, op=dist.ReduceOp.SUM, group=det1) + _assert_bitwise("allreduce SUM vs fixed tree", a, tree) + + b = x.clone() + dist.all_reduce(b, op=dist.ReduceOp.SUM, group=det2) + _assert_bitwise("allreduce bitwise across communicator instances", b, tree) + + averaged = x.clone() + dist.all_reduce(averaged, op=dist.ReduceOp.AVG, group=det1) + _assert_bitwise("allreduce AVG == SUM/world", averaged, tree / _WORLD_SIZE) + + max_det = x.clone() + dist.all_reduce(max_det, op=dist.ReduceOp.MAX, group=det1) + max_native = x.clone() + dist.all_reduce(max_native, op=dist.ReduceOp.MAX) + _assert_bitwise("allreduce MAX delegates to native", max_det, max_native) + + +def _check_reduce_scatter_vs_allreduce(rank: int, det1: dist.ProcessGroup, x, tree) -> None: + expected_shard = _shard_of(tree, rank) + + rs = torch.empty_like(expected_shard) + dist.reduce_scatter_tensor(rs, x.clone(), op=dist.ReduceOp.SUM, group=det1) + _assert_bitwise("reduce_scatter_tensor == slice of allreduce", rs, expected_shard) + + # Megatron distributed-optimizer style: the output shard is a view of the input. + buf = x.clone() + shard_view = buf.view(_WORLD_SIZE, -1)[rank] + dist.reduce_scatter_tensor(shard_view, buf, op=dist.ReduceOp.SUM, group=det1) + _assert_bitwise("aliased reduce_scatter_tensor", shard_view, expected_shard.view(shard_view.shape)) + + inputs = [chunk.contiguous() for chunk in x.clone().chunk(_WORLD_SIZE)] + out = torch.empty_like(expected_shard) + dist.reduce_scatter(out, inputs, op=dist.ReduceOp.SUM, group=det1) + _assert_bitwise("reduce_scatter (list) == slice of allreduce", out, expected_shard) + + +def _check_uneven_reduce_scatter(rank: int, det1: dist.ProcessGroup) -> None: + """List-form reduce_scatter with uneven slot sizes folds each slot at its true offset.""" + device = torch.device("cuda", torch.cuda.current_device()) + slot_sizes = [3, 5, 7, 9] + gen = torch.Generator().manual_seed(_SEED + 300 + rank) + inputs = [torch.randn(size, generator=gen, dtype=torch.float32).to(device) for size in slot_sizes] + + gathered_inputs: list[list[torch.Tensor]] = [] + for slot, size in enumerate(slot_sizes): + slot_copies = [torch.empty(size, device=device) for _ in range(_WORLD_SIZE)] + dist.all_gather(slot_copies, inputs[slot]) + gathered_inputs.append(slot_copies) + expected = _manual_tree_sum(gathered_inputs[rank]) + + out = torch.empty(slot_sizes[rank], device=device) + dist.reduce_scatter(out, inputs, op=dist.ReduceOp.SUM, group=det1) + _assert_bitwise("uneven reduce_scatter (list)", out, expected) + + +def _expected_slot_fold(rank: int, inputs: list[torch.Tensor]) -> torch.Tensor: + """Reference for list reduce_scatter: gather every rank's copy of MY slot and tree-fold. + + Gathers slot by slot: within one all_gather every rank contributes the SAME slot, + so shapes match across ranks (slots have per-rank-distinct shapes, and an uneven + all_gather is undefined over NCCL - it deadlocks). Returned flat, matching how + callers compare (their outputs are flattened views).""" + my_slot_copies: list[torch.Tensor] = [] + for slot in range(_WORLD_SIZE): + slot_input = inputs[slot].contiguous() + slot_copies = [torch.empty_like(slot_input) for _ in range(_WORLD_SIZE)] + dist.all_gather(slot_copies, slot_input) + if slot == rank: + my_slot_copies = slot_copies + return _manual_tree_sum(my_slot_copies).reshape(-1) + + +def _check_uneven_reduce_scatter_shapes(rank: int, det1: dist.ProcessGroup) -> None: + """List reduce_scatter with per-slot distinct multi-dim shapes, non-contiguous slots/output, + bf16, and a forced multi-chunk fold all match the per-slot tree fold bitwise.""" + import miles.utils.det_process_group as dpg + + device = torch.device("cuda", torch.cuda.current_device()) + slot_shapes = [(2, 3), (5,), (4, 2), (3, 3)] + + def make_inputs(seed: int, dtype: torch.dtype) -> list[torch.Tensor]: + gen = torch.Generator().manual_seed(seed + 17 * rank) + return [torch.randn(shape, generator=gen, dtype=dtype).to(device) for shape in slot_shapes] + + # Distinct multi-dim shapes per slot, via the dist API. + inputs = make_inputs(_SEED + 400, torch.float32) + expected = _expected_slot_fold(rank, inputs) + out = torch.empty(slot_shapes[rank], device=device) + dist.reduce_scatter(out, inputs, op=dist.ReduceOp.SUM, group=det1) + _assert_bitwise("uneven multi-dim reduce_scatter", out.view(-1), expected) + + # bf16 variant. + inputs_bf16 = make_inputs(_SEED + 401, torch.bfloat16) + expected_bf16 = _expected_slot_fold(rank, inputs_bf16) + out_bf16 = torch.empty(slot_shapes[rank], dtype=torch.bfloat16, device=device) + dist.reduce_scatter(out_bf16, inputs_bf16, op=dist.ReduceOp.SUM, group=det1) + _assert_bitwise("uneven bf16 reduce_scatter", out_bf16.view(-1), expected_bf16) + + # Non-contiguous slot inputs and a non-contiguous output, via the group method + # directly (the dist wrapper would densify). + nc_shapes = [(3, 2), (5, 1), (2, 4), (3, 3)] + bases = make_inputs(_SEED + 402, torch.float32) + nc_inputs = [base.reshape(shape).t() for base, shape in zip(bases, nc_shapes, strict=True)] + assert all(not t.is_contiguous() for t in nc_inputs if t.dim() > 1 and min(t.shape) > 1) + expected_nc = _expected_slot_fold(rank, nc_inputs) + out_base = torch.empty(nc_shapes[rank], device=device) + out_nc = out_base.t() + opts = dist.ReduceScatterOptions() + opts.reduceOp = dist.ReduceOp.SUM + det1.reduce_scatter([out_nc], [nc_inputs], opts).wait() + _assert_bitwise("uneven non-contiguous reduce_scatter", out_nc.contiguous().view(-1), expected_nc) + + # Forced multi-chunk fold through the real gather path. + original_cap = dpg._GATHER_BUFFER_CAP_BYTES + dpg._GATHER_BUFFER_CAP_BYTES = _WORLD_SIZE * 2 * 4 + try: + inputs_chunked = make_inputs(_SEED + 403, torch.float32) + expected_chunked = _expected_slot_fold(rank, inputs_chunked) + out_chunked = torch.empty(slot_shapes[rank], device=device) + dist.reduce_scatter(out_chunked, inputs_chunked, op=dist.ReduceOp.SUM, group=det1) + finally: + dpg._GATHER_BUFFER_CAP_BYTES = original_cap + _assert_bitwise("uneven multi-chunk reduce_scatter", out_chunked.view(-1), expected_chunked) + + +def _check_coalescing_manager(rank: int, det1: dist.ProcessGroup, x, tree, x2, tree2) -> None: + device = torch.device("cuda", torch.cuda.current_device()) + + ar1, ar2 = x.clone(), x2.clone() + with _coalescing_manager(group=det1, device=device): + dist.all_reduce(ar1, op=dist.ReduceOp.SUM, group=det1) + dist.all_reduce(ar2, op=dist.ReduceOp.SUM, group=det1) + _assert_bitwise("coalescing_manager AR (1st)", ar1, tree) + _assert_bitwise("coalescing_manager AR (2nd)", ar2, tree2) + + rs1 = torch.empty_like(_shard_of(tree, rank)) + rs2 = torch.empty_like(_shard_of(tree2, rank)) + with _coalescing_manager(group=det1, device=device): + dist.reduce_scatter_tensor(rs1, x.clone(), op=dist.ReduceOp.SUM, group=det1) + dist.reduce_scatter_tensor(rs2, x2.clone(), op=dist.ReduceOp.SUM, group=det1) + _assert_bitwise("coalescing_manager RS (1st)", rs1, _shard_of(tree, rank)) + _assert_bitwise("coalescing_manager RS (2nd)", rs2, _shard_of(tree2, rank)) + + shard_in = torch.full((128,), float(rank), device=device) + full_out = torch.empty(128 * _WORLD_SIZE, device=device) + with _coalescing_manager(group=det1, device=device): + dist.all_gather_into_tensor(full_out, shard_in, group=det1) + expected = torch.cat([torch.full((128,), float(r), device=device) for r in range(_WORLD_SIZE)]) + _assert_bitwise("coalescing_manager all_gather_into_tensor", full_out, expected) + + +def _check_non_contiguous(rank: int, det1: dist.ProcessGroup) -> None: + base = torch.randn(64, 64, generator=torch.Generator().manual_seed(_SEED + 100 + rank)).cuda() + non_contiguous = base.t() + assert not non_contiguous.is_contiguous() + + gathered = [torch.empty_like(base) for _ in range(_WORLD_SIZE)] + dist.all_gather(gathered, base) + expected = _manual_tree_sum([g.t().contiguous() for g in gathered]) + + dist.all_reduce(non_contiguous, op=dist.ReduceOp.SUM, group=det1) + _assert_bitwise("non-contiguous allreduce", non_contiguous.contiguous(), expected) + + +def _check_delegation(rank: int, det1: dist.ProcessGroup) -> None: + device = torch.device("cuda", torch.cuda.current_device()) + + broadcasted = torch.full((8,), float(rank), device=device) + dist.broadcast(broadcasted, src=0, group=det1) + _assert_bitwise("broadcast", broadcasted, torch.zeros(8, device=device)) + + piece = torch.full((4,), float(rank), device=device) + pieces = [torch.empty_like(piece) for _ in range(_WORLD_SIZE)] + dist.all_gather(pieces, piece, group=det1) + for source_rank, gathered_piece in enumerate(pieces): + _assert_bitwise( + f"all_gather piece {source_rank}", gathered_piece, torch.full((4,), float(source_rank), device=device) + ) + + full = torch.empty(4 * _WORLD_SIZE, device=device) + dist.all_gather_into_tensor(full, piece, group=det1) + expected = torch.cat([torch.full((4,), float(r), device=device) for r in range(_WORLD_SIZE)]) + _assert_bitwise("all_gather_into_tensor", full, expected) + + reduced = torch.full((4,), float(rank), device=device) + dist.reduce(reduced, dst=0, op=dist.ReduceOp.MAX, group=det1) + if rank == 0: + _assert_bitwise("reduce MAX to dst", reduced, torch.full((4,), float(_WORLD_SIZE - 1), device=device)) + + # rank r receives element r from every rank q: value = r + 10*q at position q + scatter_in = torch.arange(_WORLD_SIZE, dtype=torch.float32, device=device) + rank * 10 + a2a_out = torch.empty(_WORLD_SIZE, device=device) + dist.all_to_all_single(a2a_out, scatter_in, group=det1) + expected_a2a = torch.tensor([float(rank + 10 * q) for q in range(_WORLD_SIZE)], device=device) + _assert_bitwise("all_to_all_single", a2a_out, expected_a2a) + + peer = rank + 1 if rank % 2 == 0 else rank - 1 + outgoing = torch.full((4,), float(rank), device=device) + incoming = torch.empty(4, device=device) + if rank % 2 == 0: + dist.send(outgoing, dst=peer, group=det1) + dist.recv(incoming, src=peer, group=det1) + else: + dist.recv(incoming, src=peer, group=det1) + dist.send(outgoing, dst=peer, group=det1) + _assert_bitwise("send/recv", incoming, torch.full((4,), float(peer), device=device)) + + dist.barrier(group=det1) + + assert dist.get_backend(det1) == DET_NCCL_BACKEND_NAME, f"unexpected backend name {dist.get_backend(det1)}" + + +def _check_batch_isend_irecv_ring(rank: int, det1: dist.ProcessGroup) -> None: + """Ring batch_isend_irecv over det1 exercises the no-op coalescing hooks with batched p2p.""" + device = torch.device("cuda", torch.cuda.current_device()) + next_rank = (rank + 1) % _WORLD_SIZE + prev_rank = (rank - 1) % _WORLD_SIZE + + send_tensor = torch.full((16,), float(rank), device=device) + recv_tensor = torch.empty(16, device=device) + ops = [ + dist.P2POp(dist.isend, send_tensor, peer=next_rank, group=det1), + dist.P2POp(dist.irecv, recv_tensor, peer=prev_rank, group=det1), + ] + reqs = dist.batch_isend_irecv(ops) + for req in reqs: + req.wait() + _assert_bitwise("batch_isend_irecv ring", recv_tensor, torch.full((16,), float(prev_rank), device=device)) + + +def _check_dtype_allreduce(rank: int, det1: dist.ProcessGroup) -> None: + """bf16 and int64 SUM allreduce over det1 match the manual fold / exact integer sum bitwise.""" + device = torch.device("cuda", torch.cuda.current_device()) + + bf16_inputs = [(_order_sensitive_input(r).to(torch.bfloat16)) for r in range(_WORLD_SIZE)] + bf16_expected = _manual_tree_sum(bf16_inputs) + bf16_actual = bf16_inputs[rank].clone() + dist.all_reduce(bf16_actual, op=dist.ReduceOp.SUM, group=det1) + _assert_bitwise("bf16 SUM allreduce == bf16 tree fold", bf16_actual, bf16_expected) + + int_value = torch.full((256,), rank + 1, dtype=torch.int64, device=device) + expected_int = torch.full((256,), sum(range(1, _WORLD_SIZE + 1)), dtype=torch.int64, device=device) + dist.all_reduce(int_value, op=dist.ReduceOp.SUM, group=det1) + _assert_bitwise("int64 SUM allreduce == exact integer sum", int_value, expected_int) + + +def _check_reduce_scatter_avg(rank: int, det1: dist.ProcessGroup, x: torch.Tensor, tree: torch.Tensor) -> None: + """AVG reduce_scatter (tensor + list variant) equals this rank's slice of tree/world bitwise.""" + expected_shard = _shard_of(tree, rank) / _WORLD_SIZE + + rs = torch.empty_like(expected_shard) + dist.reduce_scatter_tensor(rs, x.clone(), op=dist.ReduceOp.AVG, group=det1) + _assert_bitwise("reduce_scatter_tensor AVG == slice of tree/world", rs, expected_shard) + + inputs = [chunk.contiguous() for chunk in x.clone().chunk(_WORLD_SIZE)] + out = torch.empty_like(expected_shard) + dist.reduce_scatter(out, inputs, op=dist.ReduceOp.AVG, group=det1) + _assert_bitwise("reduce_scatter (list) AVG == slice of tree/world", out, expected_shard) + + +def _check_reduce_sum_avg_fold(rank: int, det1: dist.ProcessGroup, x: torch.Tensor, tree: torch.Tensor) -> None: + """dist.reduce SUM/AVG over det1 folds in fixed tree order; the dst rank matches tree (/world).""" + summed = x.clone() + dist.reduce(summed, dst=0, op=dist.ReduceOp.SUM, group=det1) + if rank == 0: + _assert_bitwise("reduce SUM to dst == tree fold", summed, tree) + + averaged = x.clone() + dist.reduce(averaged, dst=0, op=dist.ReduceOp.AVG, group=det1) + if rank == 0: + _assert_bitwise("reduce AVG to dst == tree/world", averaged, tree / _WORLD_SIZE) + + +def _check_min_delegation(rank: int, det1: dist.ProcessGroup) -> None: + """MIN allreduce and MIN reduce over det1 delegate to native NCCL bitwise.""" + device = torch.device("cuda", torch.cuda.current_device()) + operand = torch.full((64,), float(rank + 1), device=device) + + min_det = operand.clone() + dist.all_reduce(min_det, op=dist.ReduceOp.MIN, group=det1) + min_native = operand.clone() + dist.all_reduce(min_native, op=dist.ReduceOp.MIN) + _assert_bitwise("allreduce MIN delegates to native", min_det, min_native) + + reduce_min = operand.clone() + dist.reduce(reduce_min, dst=0, op=dist.ReduceOp.MIN, group=det1) + if rank == 0: + _assert_bitwise("reduce MIN to dst", reduce_min, torch.ones(64, device=device)) + + +def _check_reduce_scatter_uneven_avg(rank: int, det1: dist.ProcessGroup) -> None: + """Uneven list-form reduce_scatter with AVG folds each slot at its true offset and divides by world.""" + device = torch.device("cuda", torch.cuda.current_device()) + slot_sizes = [3, 5, 7, 9] + gen = torch.Generator().manual_seed(_SEED + 400 + rank) + inputs = [torch.randn(size, generator=gen, dtype=torch.float32).to(device) for size in slot_sizes] + + expected = _expected_slot_fold(rank, inputs) / _WORLD_SIZE + + out = torch.empty(slot_sizes[rank], device=device) + dist.reduce_scatter(out, inputs, op=dist.ReduceOp.AVG, group=det1) + _assert_bitwise("uneven reduce_scatter (list) AVG == slot tree/world", out, expected) + + +def _check_multi_chunk_through_nccl(rank: int, det1: dist.ProcessGroup, x: torch.Tensor, tree: torch.Tensor) -> None: + """A tiny gather-buffer cap forces multi-chunk gathers through real NCCL; SUM allreduce and + reduce_scatter_tensor over det1 stay bitwise equal to the single-chunk tree reference.""" + import miles.utils.det_process_group as dpg + + original_cap = dpg._GATHER_BUFFER_CAP_BYTES + try: + dpg._GATHER_BUFFER_CAP_BYTES = _WORLD_SIZE * 64 * x.element_size() + + ar = x.clone() + dist.all_reduce(ar, op=dist.ReduceOp.SUM, group=det1) + _assert_bitwise("multi-chunk allreduce SUM == tree fold", ar, tree) + + expected_shard = _shard_of(tree, rank) + rs = torch.empty_like(expected_shard) + dist.reduce_scatter_tensor(rs, x.clone(), op=dist.ReduceOp.SUM, group=det1) + _assert_bitwise("multi-chunk reduce_scatter_tensor == slice of tree fold", rs, expected_shard) + finally: + dpg._GATHER_BUFFER_CAP_BYTES = original_cap + + +def _check_reduce_scatter_base_uneven_raises(rank: int, det1: dist.ProcessGroup) -> None: + """Calling det1._reduce_scatter_base directly with an indivisible input numel fails loud.""" + from torch.distributed.distributed_c10d import ReduceScatterOptions + + device = torch.device("cuda", torch.cuda.current_device()) + opts = ReduceScatterOptions() + opts.reduceOp = dist.ReduceOp.SUM + + output = torch.empty(4, device=device) + uneven_input = torch.empty(4 * _WORLD_SIZE + 1, device=device) + with pytest.raises(AssertionError): + det1._reduce_scatter_base(output, uneven_input, opts) + + +def _check_object_collectives(rank: int, det1: dist.ProcessGroup) -> None: + """Object collectives (all_gather_object + broadcast_object_list) delegate correctly over det1.""" + gathered: list[object] = [None] * _WORLD_SIZE + dist.all_gather_object(gathered, {"rank": rank, "tag": rank * 7}, group=det1) + assert gathered == [{"rank": r, "tag": r * 7} for r in range(_WORLD_SIZE)], gathered + + payload: list[object] = [{"from": 0, "value": "hello"}] if rank == 0 else [None] + dist.broadcast_object_list(payload, src=0, group=det1) + assert payload == [{"from": 0, "value": "hello"}], payload + + +def _worker(rank: int, world_size: int, port: int) -> None: + os.environ["MASTER_ADDR"] = "localhost" + os.environ["MASTER_PORT"] = str(port) + torch.cuda.set_device(rank) + register_det_nccl_backend() + dist.init_process_group(backend="nccl", rank=rank, world_size=world_size) + + det1 = dist.new_group(list(range(world_size)), backend="det_nccl") + det2 = dist.new_group(list(range(world_size)), backend="det_nccl") + + x = _order_sensitive_input(rank) + x2 = _order_sensitive_input(rank, seed=_SEED + 50) + tree = _fixed_tree_reference(x) + tree2 = _fixed_tree_reference(x2) + + _check_allreduce(rank, det1, det2, x, tree) + _check_reduce_scatter_vs_allreduce(rank, det1, x, tree) + _check_uneven_reduce_scatter(rank, det1) + _check_uneven_reduce_scatter_shapes(rank, det1) + _check_coalescing_manager(rank, det1, x, tree, x2, tree2) + _check_non_contiguous(rank, det1) + _check_delegation(rank, det1) + _check_batch_isend_irecv_ring(rank, det1) + _check_dtype_allreduce(rank, det1) + _check_reduce_scatter_avg(rank, det1, x, tree) + _check_reduce_sum_avg_fold(rank, det1, x, tree) + _check_min_delegation(rank, det1) + _check_reduce_scatter_uneven_avg(rank, det1) + _check_multi_chunk_through_nccl(rank, det1, x, tree) + _check_reduce_scatter_base_uneven_raises(rank, det1) + _check_object_collectives(rank, det1) + + dist.barrier() + dist.destroy_process_group() + + +def _free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("localhost", 0)) + return sock.getsockname()[1] + + +def test_det_process_group_multi_gpu(): + """det_nccl backend: bitwise fixed-order SUM/AVG (allreduce + reduce_scatter, incl. under + the coalescing manager) and faithful delegation of every other collective, on 4 GPUs.""" + if torch.cuda.device_count() < _WORLD_SIZE: + raise RuntimeError(f"requires {_WORLD_SIZE} GPUs, found {torch.cuda.device_count()}") + + mp.spawn(_worker, args=(_WORLD_SIZE, _free_port()), nprocs=_WORLD_SIZE, join=True) + + +def _world_backend_worker(rank: int, world_size: int, port: int) -> None: + os.environ["MASTER_ADDR"] = "localhost" + os.environ["MASTER_PORT"] = str(port) + torch.cuda.set_device(rank) + register_det_nccl_backend() + dist.init_process_group(backend="det_nccl", rank=rank, world_size=world_size) + + dist.barrier() + + x = _order_sensitive_input(rank) + tree = _fixed_tree_reference(x) + ar = x.clone() + dist.all_reduce(ar, op=dist.ReduceOp.SUM) + _assert_bitwise("default-group det_nccl allreduce == tree fold", ar, tree) + + assert dist.get_backend() == DET_NCCL_BACKEND_NAME, f"unexpected default backend {dist.get_backend()}" + + dist.barrier() + dist.destroy_process_group() + + +def test_det_nccl_as_world_backend_multi_gpu(): + """det_nccl wired as the DEFAULT-group backend (train_actor shape): barrier + bitwise tree SUM.""" + if torch.cuda.device_count() < _WORLD_SIZE: + raise RuntimeError(f"requires {_WORLD_SIZE} GPUs, found {torch.cuda.device_count()}") + + mp.spawn(_world_backend_worker, args=(_WORLD_SIZE, _free_port()), nprocs=_WORLD_SIZE, join=True) + + +if __name__ == "__main__": + import sys + + sys.exit(pytest.main([__file__, "-v"])) From a48d03f0d6c4f93ddc5d1bad0bc16ea0d3a6ebd7 Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Mon, 22 Jun 2026 18:15:17 +0800 Subject: [PATCH 16/41] Add a per-process identity helper Add a per-process identity helper that uniquely keys each training process, used to attribute structured fault-tolerance events to their originating process. - miles/utils/process_identity.py and tests. --- miles/utils/process_identity.py | 35 ++++++++++++++++++ tests/fast/utils/test_process_identity.py | 43 +++++++++++++++++++++++ 2 files changed, 78 insertions(+) create mode 100644 miles/utils/process_identity.py create mode 100644 tests/fast/utils/test_process_identity.py diff --git a/miles/utils/process_identity.py b/miles/utils/process_identity.py new file mode 100644 index 00000000000..633a9ca2ad5 --- /dev/null +++ b/miles/utils/process_identity.py @@ -0,0 +1,35 @@ +from typing import Annotated, Literal + +from pydantic import Discriminator, NonNegativeInt + +from miles.utils.pydantic_utils import FrozenStrictBaseModel + + +class _ProcessIdentityBase(FrozenStrictBaseModel): + component: str + + def to_name(self) -> str: + return self.component + + +class MainProcessIdentity(_ProcessIdentityBase): + component: Literal["main"] = "main" + + +class RolloutManagerProcessIdentity(_ProcessIdentityBase): + component: Literal["rollout_manager"] = "rollout_manager" + + +class TrainProcessIdentity(_ProcessIdentityBase): + component: Literal["actor", "critic"] + cell_index: NonNegativeInt + rank_within_cell: NonNegativeInt + + def to_name(self) -> str: + return f"{self.component}_cell{self.cell_index}_rank{self.rank_within_cell}" + + +ProcessIdentity = Annotated[ + MainProcessIdentity | RolloutManagerProcessIdentity | TrainProcessIdentity, + Discriminator("component"), +] diff --git a/tests/fast/utils/test_process_identity.py b/tests/fast/utils/test_process_identity.py new file mode 100644 index 00000000000..76327549b6f --- /dev/null +++ b/tests/fast/utils/test_process_identity.py @@ -0,0 +1,43 @@ +"""Tests for process_identity module.""" + +import pytest +from pydantic import ValidationError + +from miles.utils.process_identity import MainProcessIdentity, RolloutManagerProcessIdentity, TrainProcessIdentity + + +class TestProcessIdentityToName: + def test_main(self) -> None: + assert MainProcessIdentity().to_name() == "main" + + def test_rollout_manager(self) -> None: + assert RolloutManagerProcessIdentity().to_name() == "rollout_manager" + + def test_actor(self) -> None: + source = TrainProcessIdentity(component="actor", cell_index=1, rank_within_cell=3) + assert source.to_name() == "actor_cell1_rank3" + + def test_critic(self) -> None: + source = TrainProcessIdentity(component="critic", cell_index=0, rank_within_cell=2) + assert source.to_name() == "critic_cell0_rank2" + + +class TestTrainProcessIdentityValidation: + def test_negative_cell_index_rejected(self) -> None: + """A negative cell_index fails validation.""" + with pytest.raises(ValidationError): + TrainProcessIdentity(component="actor", cell_index=-1, rank_within_cell=0) + + def test_negative_rank_within_cell_rejected(self) -> None: + """A negative rank_within_cell fails validation.""" + with pytest.raises(ValidationError): + TrainProcessIdentity(component="actor", cell_index=0, rank_within_cell=-1) + + +class TestTrainProcessIdentityRoundtrip: + def test_serialize_deserialize(self) -> None: + source = TrainProcessIdentity(component="actor", cell_index=2, rank_within_cell=0) + parsed = TrainProcessIdentity.model_validate_json(source.model_dump_json()) + assert parsed.cell_index == 2 + assert parsed.rank_within_cell == 0 + assert parsed.component == "actor" From 3e85a869c398a0ba25b289476fa336e969330b8c Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Mon, 22 Jun 2026 18:15:17 +0800 Subject: [PATCH 17/41] Add structured event models keyed by per-process identity Add the structured event models (Event / EventBase hierarchy) for the fault-tolerance event log, each tagged with the originating ProcessIdentity. - miles/utils/event_logger/models.py and tests. --- miles/utils/event_logger/__init__.py | 0 miles/utils/event_logger/models.py | 119 +++++++++++ tests/fast/utils/event_logger/__init__.py | 0 tests/fast/utils/event_logger/test_models.py | 211 +++++++++++++++++++ 4 files changed, 330 insertions(+) create mode 100644 miles/utils/event_logger/__init__.py create mode 100644 miles/utils/event_logger/models.py create mode 100644 tests/fast/utils/event_logger/__init__.py create mode 100644 tests/fast/utils/event_logger/test_models.py diff --git a/miles/utils/event_logger/__init__.py b/miles/utils/event_logger/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/miles/utils/event_logger/models.py b/miles/utils/event_logger/models.py new file mode 100644 index 00000000000..c9a136c730e --- /dev/null +++ b/miles/utils/event_logger/models.py @@ -0,0 +1,119 @@ +from datetime import datetime +from typing import Annotated, Any, Literal + +from pydantic import Discriminator + +from miles.backends.megatron_utils.types import TrainStepOutcome +from miles.utils.process_identity import ProcessIdentity +from miles.utils.pydantic_utils import FrozenStrictBaseModel + + +class EventBase(FrozenStrictBaseModel): + timestamp: datetime + source: ProcessIdentity + + +class _ActorTrainEventBase(EventBase): + rollout_id: int + attempt: int = 0 + + +class OptimizerStateInfo(FrozenStrictBaseModel): + """Snapshot of one sub-optimizer's state with tensors replaced by hashes.""" + + param_names: dict[int, str] + state_dict: dict[str, Any] + + +class TrainEngineLocalWeightChecksumState(FrozenStrictBaseModel): + param_hashes: dict[str, str] + buffer_hashes: dict[str, str] + # May be skipped in non-debug mode if too expensive + optimizer_hashes: list[OptimizerStateInfo] + + +class TrainEngineLocalWeightChecksumEvent(_ActorTrainEventBase): + type: Literal["train_engine_local_weight_checksum"] = "train_engine_local_weight_checksum" + state: TrainEngineLocalWeightChecksumState + + +class WitnessSnapshotParamEvent(_ActorTrainEventBase): + type: Literal["witness_snapshot_param"] = "witness_snapshot_param" + instance_id: str + # TODO: may shrink a contiguous range of numbers into a pair, if this is too large/slow + nonzero_witness_ids: list[int] + stale_ids: list[int] + + +class WitnessAllocateIdEvent(EventBase): + type: Literal["witness_allocate_id"] = "witness_allocate_id" + rollout_id: int + attempt: int + witness_id_to_sample_index: dict[int, int] + # Allocator counter after this allocation; a resumed run recovers the allocator from it. + counter_after: int + + +class TrainGroupStepEndEvent(EventBase): + type: Literal["train_group_step_end"] = "train_group_step_end" + rollout_id: int + cell_outcomes: dict[int, Literal["error"] | list[TrainStepOutcome]] + + +class CellReconfigureEvent(EventBase): + type: Literal["cell_reconfigure"] = "cell_reconfigure" + rollout_id: int + quorum_id: int + src_cell_index: int | None + # healing happened iff non-empty + healed_cell_indices: list[int] + alive_cell_indices_after: list[int] + + +class InferenceEngineWeightChecksumEvent(EventBase): + type: Literal["inference_engine_weight_checksum"] = "inference_engine_weight_checksum" + # None for the initial out-of-loop weight sync (not tied to a rollout). + rollout_id: int | None + # One {tensor -> hash} dict per rollout engine; a TP>1 engine's ranks merge with a rank{r}/ prefix. + engine_checksums: list[dict[str, str]] + + +class TrainAdvantageComputationEvent(_ActorTrainEventBase): + type: Literal["train_advantage_computation"] = "train_advantage_computation" + advantages: list[list[float]] + witness_ids: list[list[int]] + + +Event = Annotated[ + TrainEngineLocalWeightChecksumEvent + | WitnessSnapshotParamEvent + | WitnessAllocateIdEvent + | TrainGroupStepEndEvent + | CellReconfigureEvent + | InferenceEngineWeightChecksumEvent + | TrainAdvantageComputationEvent, + Discriminator("type"), +] + + +def _to_snake_case(name: str) -> str: + import re + + return re.sub(r"(?<=[a-z0-9])([A-Z])", r"_\1", name).lower() + + +def _check_event_naming() -> None: + import typing + + event_types = typing.get_args(typing.get_args(Event)[0]) + for cls in event_types: + type_value = cls.model_fields["type"].default + expected_snake = type_value + "_event" + actual_snake = _to_snake_case(cls.__name__) + assert actual_snake == expected_snake, ( + f"Event class {cls.__name__} (snake: {actual_snake}) does not match " + f"type '{type_value}' (expected snake: {expected_snake})" + ) + + +_check_event_naming() diff --git a/tests/fast/utils/event_logger/__init__.py b/tests/fast/utils/event_logger/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/fast/utils/event_logger/test_models.py b/tests/fast/utils/event_logger/test_models.py new file mode 100644 index 00000000000..2cce486c2a8 --- /dev/null +++ b/tests/fast/utils/event_logger/test_models.py @@ -0,0 +1,211 @@ +from datetime import datetime, timezone + +import pytest +from pydantic import TypeAdapter, ValidationError + +from miles.backends.megatron_utils.types import TrainStepOutcome +from miles.utils.event_logger.models import ( + CellReconfigureEvent, + Event, + InferenceEngineWeightChecksumEvent, + TrainGroupStepEndEvent, + WitnessAllocateIdEvent, + WitnessSnapshotParamEvent, +) +from miles.utils.process_identity import MainProcessIdentity, TrainProcessIdentity + +_event_adapter = TypeAdapter(Event) + +_FIXED_TS = datetime(2026, 1, 1, tzinfo=timezone.utc) +_FIXED_SOURCE = MainProcessIdentity() +_TRAIN_SOURCE = TrainProcessIdentity(component="actor", cell_index=0, rank_within_cell=0) + + +class TestEventModelsDiscriminatedUnion: + def test_roundtrip_via_discriminator(self) -> None: + event = WitnessAllocateIdEvent( + timestamp=_FIXED_TS, + source=_FIXED_SOURCE, + rollout_id=0, + attempt=0, + witness_id_to_sample_index={10: 0, 11: 1}, + counter_after=12, + ) + parsed = _event_adapter.validate_json(event.model_dump_json()) + assert isinstance(parsed, WitnessAllocateIdEvent) + assert parsed.witness_id_to_sample_index == {10: 0, 11: 1} + + def test_discriminator_distinguishes_types(self) -> None: + e1 = WitnessAllocateIdEvent( + timestamp=_FIXED_TS, + source=_FIXED_SOURCE, + rollout_id=0, + attempt=0, + witness_id_to_sample_index={0: 0}, + counter_after=1, + ) + e2 = TrainGroupStepEndEvent( + timestamp=_FIXED_TS, + source=_FIXED_SOURCE, + rollout_id=0, + cell_outcomes={0: [TrainStepOutcome.NORMAL]}, + ) + p1 = _event_adapter.validate_json(e1.model_dump_json()) + p2 = _event_adapter.validate_json(e2.model_dump_json()) + assert type(p1) is not type(p2) + + +class TestEventModelsStrictRejectExtraFields: + def test_extra_field_rejected(self) -> None: + data = { + "type": "witness_allocate_id", + "timestamp": "2026-01-01T00:00:00Z", + "source": {"component": "main"}, + "rollout_id": 0, + "attempt": 0, + "witness_id_to_sample_index": {0: 0}, + "bogus_field": 123, + } + with pytest.raises(ValidationError, match="bogus_field"): + WitnessAllocateIdEvent.model_validate(data) + + +class TestWitnessAllocateIdEvent: + def test_json_roundtrip(self) -> None: + event = WitnessAllocateIdEvent( + timestamp=_FIXED_TS, + source=_FIXED_SOURCE, + rollout_id=2, + attempt=0, + witness_id_to_sample_index={10: 0, 11: 1}, + counter_after=12, + ) + parsed = _event_adapter.validate_json(event.model_dump_json()) + assert isinstance(parsed, WitnessAllocateIdEvent) + assert parsed.rollout_id == 2 + assert parsed.attempt == 0 + assert parsed.witness_id_to_sample_index == {10: 0, 11: 1} + + +class TestTrainGroupStepEndEvent: + def test_json_roundtrip(self) -> None: + event = TrainGroupStepEndEvent( + timestamp=_FIXED_TS, + source=_FIXED_SOURCE, + rollout_id=3, + cell_outcomes={0: [TrainStepOutcome.NORMAL], 1: "error"}, + ) + parsed = _event_adapter.validate_json(event.model_dump_json()) + assert isinstance(parsed, TrainGroupStepEndEvent) + assert parsed.rollout_id == 3 + assert parsed.cell_outcomes[0] == [TrainStepOutcome.NORMAL] + assert parsed.cell_outcomes[1] == "error" + + +class TestCellReconfigureEvent: + def test_healing_json_roundtrip(self) -> None: + """A healing reconfigure event (non-empty healed cells, with src) survives a JSON round-trip.""" + event = CellReconfigureEvent( + timestamp=_FIXED_TS, + source=_FIXED_SOURCE, + rollout_id=3, + quorum_id=1, + src_cell_index=0, + healed_cell_indices=[2], + alive_cell_indices_after=[0, 1, 2], + ) + parsed = _event_adapter.validate_json(event.model_dump_json()) + assert isinstance(parsed, CellReconfigureEvent) + assert parsed.rollout_id == 3 + assert parsed.quorum_id == 1 + assert parsed.src_cell_index == 0 + assert parsed.healed_cell_indices == [2] + assert parsed.alive_cell_indices_after == [0, 1, 2] + + def test_shrink_json_roundtrip(self) -> None: + """A pure-shrink reconfigure event (no healed cells, src None) survives a JSON round-trip.""" + event = CellReconfigureEvent( + timestamp=_FIXED_TS, + source=_FIXED_SOURCE, + rollout_id=2, + quorum_id=1, + src_cell_index=None, + healed_cell_indices=[], + alive_cell_indices_after=[0], + ) + parsed = _event_adapter.validate_json(event.model_dump_json()) + assert isinstance(parsed, CellReconfigureEvent) + assert parsed.src_cell_index is None + assert parsed.healed_cell_indices == [] + assert parsed.alive_cell_indices_after == [0] + + +class TestInferenceEngineWeightChecksumEvent: + def test_json_roundtrip(self) -> None: + """An engine weight checksum event survives a JSON round-trip with its per-engine checksums intact.""" + engine_checksums = [ + {"rank0/embed.weight": "aaa"}, + {"rank0/embed.weight": "aaa", "rank1/embed.weight": "bbb"}, + ] + event = InferenceEngineWeightChecksumEvent( + timestamp=_FIXED_TS, + source=_FIXED_SOURCE, + rollout_id=4, + engine_checksums=engine_checksums, + ) + parsed = _event_adapter.validate_json(event.model_dump_json()) + assert isinstance(parsed, InferenceEngineWeightChecksumEvent) + assert parsed.rollout_id == 4 + assert parsed.engine_checksums == engine_checksums + + +class TestWitnessSnapshotParamEventWithStaleIds: + def test_json_roundtrip(self) -> None: + event = WitnessSnapshotParamEvent( + timestamp=_FIXED_TS, + source=_TRAIN_SOURCE, + rollout_id=5, + instance_id="actor_cell0_rank0", + nonzero_witness_ids=[10, 11, 12], + stale_ids=[0, 1, 2, 3, 4, 5, 6, 7], + ) + parsed = _event_adapter.validate_json(event.model_dump_json()) + assert isinstance(parsed, WitnessSnapshotParamEvent) + assert parsed.stale_ids == [0, 1, 2, 3, 4, 5, 6, 7] + assert parsed.nonzero_witness_ids == [10, 11, 12] + + +class TestDiscriminatedUnionParsesAllEvents: + def test_all_event_types_parse(self) -> None: + events = [ + WitnessAllocateIdEvent( + timestamp=_FIXED_TS, + source=_FIXED_SOURCE, + rollout_id=0, + attempt=0, + witness_id_to_sample_index={0: 0}, + counter_after=1, + ), + TrainGroupStepEndEvent( + timestamp=_FIXED_TS, + source=_FIXED_SOURCE, + rollout_id=0, + cell_outcomes={0: [TrainStepOutcome.NORMAL]}, + ), + InferenceEngineWeightChecksumEvent( + timestamp=_FIXED_TS, + source=_FIXED_SOURCE, + rollout_id=0, + engine_checksums=[{"rank0/w": "aaa"}], + ), + ] + for event in events: + parsed = _event_adapter.validate_json(event.model_dump_json()) + assert type(parsed) is type(event) + + +class TestCheckEventNaming: + def test_naming_convention_holds(self) -> None: + from miles.utils.event_logger.models import _check_event_naming + + _check_event_naming() From dc11be395134b2c83c7feb6983b53041bc8f61dd Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Mon, 22 Jun 2026 18:15:17 +0800 Subject: [PATCH 18/41] Add structured event logging keyed by per-process identity Add the structured event logger that records typed events keyed by per-process identity, wire it through the logging helper and CLI argument, and start it from the train entrypoints. - miles/utils/event_logger/logger.py, logging_utils.py, arguments.py and entrypoint wiring, with tests. --- miles/ray/rollout/rollout_manager.py | 3 +- miles/ray/train_actor.py | 3 - miles/rollout/session/server.py | 4 +- miles/router/router.py | 4 +- miles/utils/arguments.py | 4 +- miles/utils/event_logger/logger.py | 145 +++++++++ miles/utils/event_logger/models.py | 10 +- miles/utils/http_utils.py | 4 +- miles/utils/logging_utils.py | 15 +- tests/fast/utils/event_logger/test_logger.py | 311 +++++++++++++++++++ tools/convert_hf_to_torch_dist.py | 7 +- train.py | 3 +- train_async.py | 3 +- 13 files changed, 497 insertions(+), 19 deletions(-) create mode 100644 miles/utils/event_logger/logger.py create mode 100644 tests/fast/utils/event_logger/test_logger.py diff --git a/miles/ray/rollout/rollout_manager.py b/miles/ray/rollout/rollout_manager.py index 107d771233c..f853f99e61b 100644 --- a/miles/ray/rollout/rollout_manager.py +++ b/miles/ray/rollout/rollout_manager.py @@ -28,6 +28,7 @@ from miles.utils.logging_utils import configure_logger from miles.utils.metric_checker import MetricChecker from miles.utils.misc import load_function +from miles.utils.process_identity import RolloutManagerProcessIdentity from miles.utils.ray_utils import Box from miles.utils.tracking_utils import init_tracking @@ -43,7 +44,7 @@ class RolloutManager: """The class to run rollout and convert rollout data to training data.""" def __init__(self, args, pg): - configure_logger() + configure_logger(args, source=RolloutManagerProcessIdentity()) self.pg = pg self.args = args diff --git a/miles/ray/train_actor.py b/miles/ray/train_actor.py index 971d8a3a266..f91d7681a78 100644 --- a/miles/ray/train_actor.py +++ b/miles/ray/train_actor.py @@ -14,7 +14,6 @@ from miles.utils.det_process_group import DET_NCCL_BACKEND_NAME, register_det_nccl_backend from miles.utils.distributed_utils import init_gloo_group from miles.utils.env_report import collect_and_print_node_env_report -from miles.utils.logging_utils import configure_logger from miles.utils.memory_utils import clear_memory, print_memory from miles.utils.test_utils.fault_injector import inject_fault as _inject_fault @@ -35,8 +34,6 @@ def get_local_gpu_id(): class TrainRayActor(RayActor): def __init__(self, world_size, rank, master_addr, master_port): - configure_logger() - self._world_size = world_size self._rank = rank if master_addr: diff --git a/miles/rollout/session/server.py b/miles/rollout/session/server.py index 281e50fdf22..efd73558ade 100644 --- a/miles/rollout/session/server.py +++ b/miles/rollout/session/server.py @@ -16,7 +16,7 @@ from miles.rollout.session.core import ProxyRequest from miles.rollout.session.sessions import setup_session_routes -from miles.utils.logging_utils import configure_logger +from miles.utils.logging_utils import configure_logger_raw logger = logging.getLogger(__name__) @@ -73,7 +73,7 @@ async def do_proxy(self, request: ProxyRequest, path: str, *, body: bytes, heade def run_session_server(args, backend_url: str): """Entry point to start the standalone session server as a subprocess.""" # Spawned as a fresh interpreter, so it inherits no logging config. - configure_logger() + configure_logger_raw("session_server") # Visible to `pkill -9 miles`; without this the daemon inherits "python". setproctitle.setproctitle("miles-session-server") diff --git a/miles/router/router.py b/miles/router/router.py index 6b465f3c2b1..6013535ddc6 100644 --- a/miles/router/router.py +++ b/miles/router/router.py @@ -10,7 +10,7 @@ from fastapi.responses import JSONResponse from starlette.responses import Response -from miles.utils.logging_utils import configure_logger +from miles.utils.logging_utils import configure_logger_raw logger = logging.getLogger(__name__) @@ -20,7 +20,7 @@ def run_router(args): Run the Miles router with the specified configuration. """ # Spawned as a fresh interpreter, so it inherits no logging config. - configure_logger() + configure_logger_raw("miles_router") # Visible to `pkill -9 miles`; without this the daemon inherits "python". setproctitle.setproctitle("miles-router") diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index 2348dfa1e53..c48a8a215d8 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -13,7 +13,7 @@ from miles.utils.environ import enable_experimental_rollout_refactor from miles.utils.eval_config import EvalDatasetConfig, build_eval_dataset_configs, ensure_dataset_list from miles.utils.hf_config import is_dsa, load_hf_config -from miles.utils.logging_utils import configure_logger +from miles.utils.logging_utils import configure_logger_raw from miles.utils.misc import load_function logger = logging.getLogger(__name__) @@ -2009,7 +2009,7 @@ def add_user_provided_function_arguments(parser): def parse_args(add_custom_arguments=None): # Users may call `parse_args` very early, thus we ensure logger is configured here - configure_logger("main") + configure_logger_raw("main") add_miles_arguments = get_miles_extra_args_provider(add_custom_arguments) diff --git a/miles/utils/event_logger/logger.py b/miles/utils/event_logger/logger.py new file mode 100644 index 00000000000..6155bd8d0f5 --- /dev/null +++ b/miles/utils/event_logger/logger.py @@ -0,0 +1,145 @@ +import contextvars +import functools +import inspect +import logging +import threading +from collections.abc import Callable, Generator +from contextlib import contextmanager +from datetime import datetime, timezone +from pathlib import Path +from typing import Any + +from pydantic import TypeAdapter + +from miles.utils.event_logger.models import Event, EventBase +from miles.utils.process_identity import ProcessIdentity +from miles.utils.structured_log import log_structured, prune_for_log + +logger = logging.getLogger(__name__) + +_event_adapter: TypeAdapter[Event] = TypeAdapter(Event) + + +class EventLogger: + def __init__(self, *, log_dir: Path | str, file_name: str = "events.jsonl", source: ProcessIdentity) -> None: + self._log_dir = Path(log_dir) + self._log_dir.mkdir(parents=True, exist_ok=True) + self._lock = threading.Lock() + self._path = self._log_dir / file_name + self._source = source + self._context_var: contextvars.ContextVar[dict[str, Any]] = contextvars.ContextVar( + "event_logger_context", + ) + + @property + def source(self) -> ProcessIdentity: + return self._source + + @contextmanager + def with_context(self, ctx: dict[str, Any]) -> Generator[None, None, None]: + """Temporarily merge extra fields into every event logged within this scope. + + Safe for both threads and asyncio tasks (uses contextvars). + """ + prev = self._context_var.get({}) + merged = {**prev, **ctx} + token = self._context_var.set(merged) + try: + yield + finally: + assert self._context_var.get() == merged + self._context_var.reset(token) + + def log(self, event_cls: type[EventBase], partial: dict[str, Any], *, print_log: bool = True) -> None: + event = event_cls( + **{ + **partial, + "timestamp": datetime.now(timezone.utc), + "source": self._source, + **self._context_var.get({}), + } + ) + line = event.model_dump_json() + "\n" + with self._lock: + # Opened per write so the file can be replaced (e.g. restored from a + # checkpoint snapshot) at any point between events. + with self._path.open("a", encoding="utf-8") as f: + f.write(line) + if print_log: + payload = prune_for_log(event.model_dump(mode="json", exclude={"timestamp", "source"})) + log_structured(logger.info, op="event", event=type(event).__name__, **payload) + + def close(self) -> None: + pass + + +_event_logger: EventLogger | None = None + + +def set_event_logger(event_logger: EventLogger | None) -> None: + global _event_logger + _event_logger = event_logger + + +def get_event_logger() -> EventLogger: + if _event_logger is None: + raise RuntimeError("EventLogger not initialized. Call set_event_logger() first.") + return _event_logger + + +def is_event_logger_initialized() -> bool: + return _event_logger is not None + + +def event_logger_context(ctx_fn: Callable[..., dict[str, Any]]) -> Callable: + """Decorator that wraps a method with EventLogger.with_context if initialized. + + ``ctx_fn`` receives the same arguments as the decorated method and returns + the context dict. If the event logger is not initialized, the method runs + without any context. + """ + + def decorator(method: Callable) -> Callable: + assert not inspect.iscoroutinefunction(method), "event_logger_context does not support async methods" + + @functools.wraps(method) + def wrapper(*args: Any, **kwargs: Any) -> Any: + if not is_event_logger_initialized(): + return method(*args, **kwargs) + + ctx_value = ctx_fn(*args, **kwargs) + with get_event_logger().with_context(ctx_value): + return method(*args, **kwargs) + + return wrapper + + return decorator + + +def read_events(log_dir: Path) -> list[Event]: + """Read all JSONL event files from a directory and return parsed events.""" + events: list[Event] = [] + + jsonl_files = sorted(log_dir.glob("**/*.jsonl")) + if not jsonl_files: + logger.warning("No JSONL files found in %s", log_dir) + return events + + for jsonl_path in jsonl_files: + with open(jsonl_path, encoding="utf-8") as f: + for line_num, raw_line in enumerate(f, start=1): + raw_line = raw_line.strip() + if not raw_line: + continue + try: + event = _event_adapter.validate_json(raw_line) + events.append(event) + except Exception: + logger.warning( + "Failed to parse event at %s:%d", + jsonl_path, + line_num, + exc_info=True, + ) + + return events diff --git a/miles/utils/event_logger/models.py b/miles/utils/event_logger/models.py index c9a136c730e..0749f72027f 100644 --- a/miles/utils/event_logger/models.py +++ b/miles/utils/event_logger/models.py @@ -84,6 +84,13 @@ class TrainAdvantageComputationEvent(_ActorTrainEventBase): witness_ids: list[list[int]] +class MetricEvent(EventBase): + type: Literal["metric"] = "metric" + rollout_id: int | None = None + attempt: int | None = None + metrics: dict[str, Any] + + Event = Annotated[ TrainEngineLocalWeightChecksumEvent | WitnessSnapshotParamEvent @@ -91,7 +98,8 @@ class TrainAdvantageComputationEvent(_ActorTrainEventBase): | TrainGroupStepEndEvent | CellReconfigureEvent | InferenceEngineWeightChecksumEvent - | TrainAdvantageComputationEvent, + | TrainAdvantageComputationEvent + | MetricEvent, Discriminator("type"), ] diff --git a/miles/utils/http_utils.py b/miles/utils/http_utils.py index 621e60532e8..9b27fdb6131 100644 --- a/miles/utils/http_utils.py +++ b/miles/utils/http_utils.py @@ -10,7 +10,7 @@ import httpx -from miles.utils.logging_utils import configure_logger +from miles.utils.logging_utils import configure_logger_raw logger = logging.getLogger(__name__) @@ -141,7 +141,7 @@ def _wrap_ipv6(host): def run_router(args): # Spawned as a fresh interpreter, so it inherits no logging config. - configure_logger() + configure_logger_raw("router") try: from sglang_router.launch_router import launch_router diff --git a/miles/utils/logging_utils.py b/miles/utils/logging_utils.py index b6ae346b3a1..fb700179a82 100644 --- a/miles/utils/logging_utils.py +++ b/miles/utils/logging_utils.py @@ -3,6 +3,8 @@ import re import sys import warnings +from miles.utils.event_logger.logger import EventLogger, is_event_logger_initialized, set_event_logger +from miles.utils.process_identity import ProcessIdentity _LOGGER_CONFIGURED = False @@ -11,8 +13,17 @@ _FATAL_ASYNC_PATTERN = "coroutine .* was never awaited" +def configure_logger(args, *, source: ProcessIdentity) -> None: + name = source.to_name() + configure_logger_raw(name) + + if (event_dir := getattr(args, "save_debug_event_data", None)) is not None: + if not is_event_logger_initialized(): + set_event_logger(EventLogger(log_dir=event_dir, file_name=f"{name}.jsonl", source=source)) + + # ref: SGLang -def configure_logger(prefix: str = ""): +def configure_logger_raw(name: str = "") -> None: global _LOGGER_CONFIGURED if _LOGGER_CONFIGURED: return @@ -21,7 +32,7 @@ def configure_logger(prefix: str = ""): logging.basicConfig( level=logging.INFO, - format=f"[%(asctime)s{prefix}] %(filename)s:%(lineno)d - %(message)s", + format=f"[%(asctime)s.%(msecs)03d {name}] %(filename)s:%(lineno)d - %(message)s", datefmt="%Y-%m-%d %H:%M:%S", force=True, ) diff --git a/tests/fast/utils/event_logger/test_logger.py b/tests/fast/utils/event_logger/test_logger.py new file mode 100644 index 00000000000..32c82c4f16b --- /dev/null +++ b/tests/fast/utils/event_logger/test_logger.py @@ -0,0 +1,311 @@ +import json +import threading +from datetime import datetime, timezone +from pathlib import Path + +import pytest + +import miles.utils.event_logger.logger as event_logger_module +from miles.utils.event_logger.logger import EventLogger, event_logger_context, get_event_logger, set_event_logger +from miles.utils.event_logger.models import MetricEvent, WitnessAllocateIdEvent +from miles.utils.process_identity import MainProcessIdentity, TrainProcessIdentity + +_TEST_SOURCE = MainProcessIdentity() + + +def _make_logger(log_dir: Path, file_name: str = "events.jsonl") -> EventLogger: + return EventLogger(log_dir=log_dir, file_name=file_name, source=_TEST_SOURCE) + + +_EVENT_CLS = WitnessAllocateIdEvent +_EVENT_PARTIAL: dict = dict( + rollout_id=0, attempt=0, witness_id_to_sample_index={10: 0, 11: 1, 12: 2}, counter_after=13 +) + + +class TestEventLoggerWritesJsonl: + def test_writes_multiple_events(self, tmp_path: Path) -> None: + logger = _make_logger(tmp_path, file_name="test.jsonl") + + logger.log(_EVENT_CLS, _EVENT_PARTIAL) + logger.log( + WitnessAllocateIdEvent, dict(rollout_id=1, attempt=0, witness_id_to_sample_index={0: 0}, counter_after=1) + ) + logger.close() + + lines = (tmp_path / "test.jsonl").read_text().strip().split("\n") + assert len(lines) == 2 + for line in lines: + parsed = json.loads(line) + assert "timestamp" in parsed + assert "type" in parsed + + +class TestEventLoggerAutoFillsMetadata: + def test_timestamp_is_utc_and_recent(self, tmp_path: Path) -> None: + logger = _make_logger(tmp_path) + + before = datetime.now(timezone.utc) + logger.log(_EVENT_CLS, _EVENT_PARTIAL) + after = datetime.now(timezone.utc) + logger.close() + + line = (tmp_path / "events.jsonl").read_text().strip() + parsed = json.loads(line) + ts = datetime.fromisoformat(parsed["timestamp"].replace("Z", "+00:00")) + assert before <= ts <= after + + def test_source_auto_filled(self, tmp_path: Path) -> None: + source = TrainProcessIdentity(component="actor", cell_index=2, rank_within_cell=3) + logger = EventLogger(log_dir=tmp_path, source=source) + logger.log(_EVENT_CLS, _EVENT_PARTIAL) + logger.close() + + parsed = json.loads((tmp_path / "events.jsonl").read_text().strip()) + assert parsed["source"]["component"] == "actor" + assert parsed["source"]["cell_index"] == 2 + assert parsed["source"]["rank_within_cell"] == 3 + + +class TestEventLoggerThreadSafety: + def test_concurrent_writes_no_data_loss(self, tmp_path: Path) -> None: + logger = _make_logger(tmp_path) + num_threads = 8 + events_per_thread = 50 + + def writer() -> None: + for _ in range(events_per_thread): + logger.log(_EVENT_CLS, _EVENT_PARTIAL) + + threads = [threading.Thread(target=writer) for _ in range(num_threads)] + for t in threads: + t.start() + for t in threads: + t.join() + logger.close() + + lines = (tmp_path / "events.jsonl").read_text().strip().split("\n") + assert len(lines) == num_threads * events_per_thread + + for line in lines: + json.loads(line) + + +class TestSetGetEventLogger: + def test_set_then_get(self, tmp_path: Path) -> None: + logger = _make_logger(tmp_path) + set_event_logger(logger) + assert get_event_logger() is logger + logger.close() + set_event_logger(None) + + def test_replace_logger(self, tmp_path: Path) -> None: + logger1 = _make_logger(tmp_path, file_name="a.jsonl") + logger2 = _make_logger(tmp_path, file_name="b.jsonl") + set_event_logger(logger1) + set_event_logger(logger2) + assert get_event_logger() is logger2 + logger1.close() + logger2.close() + set_event_logger(None) + + +class TestGetEventLoggerRaisesWhenNotSet: + def test_raises_runtime_error(self) -> None: + original = event_logger_module._event_logger + event_logger_module._event_logger = None + try: + with pytest.raises(RuntimeError, match="EventLogger not initialized"): + get_event_logger() + finally: + event_logger_module._event_logger = original + + +class TestEventLoggerFlushOnEachWrite: + def test_readable_before_close(self, tmp_path: Path) -> None: + logger = _make_logger(tmp_path) + logger.log(_EVENT_CLS, _EVENT_PARTIAL) + + content = (tmp_path / "events.jsonl").read_text() + assert len(content.strip()) > 0 + logger.close() + + +class TestEventLoggerCreatesDirectory: + def test_creates_nested_dir(self, tmp_path: Path) -> None: + nested = tmp_path / "a" / "b" / "c" + logger = _make_logger(nested) + logger.log(_EVENT_CLS, _EVENT_PARTIAL) + logger.close() + assert (nested / "events.jsonl").exists() + + +class TestEventLoggerFilePerWrite: + def test_log_after_events_file_removed_writes_to_fresh_file(self, tmp_path: Path) -> None: + """The file is opened per write, so a checkpoint restore can swap it between events.""" + logger = _make_logger(tmp_path) + logger.log(_EVENT_CLS, _EVENT_PARTIAL) + + path = tmp_path / "events.jsonl" + path.unlink() + + logger.log(_EVENT_CLS, _EVENT_PARTIAL) + assert len(path.read_text().strip().split("\n")) == 1 + + +class TestReadEvents: + def test_malformed_line_skipped_with_warning(self, tmp_path: Path) -> None: + from miles.utils.event_logger.logger import read_events + + logger = _make_logger(tmp_path) + logger.log(_EVENT_CLS, _EVENT_PARTIAL) + logger.close() + + with open(tmp_path / "events.jsonl", "a") as f: + f.write("this is not valid json\n") + + events = read_events(tmp_path) + assert len(events) == 1 + + def test_reads_multiple_jsonl_files(self, tmp_path: Path) -> None: + from miles.utils.event_logger.logger import read_events + + logger_a = EventLogger(log_dir=tmp_path, file_name="a.jsonl", source=_TEST_SOURCE) + logger_a.log(_EVENT_CLS, _EVENT_PARTIAL) + logger_a.close() + + logger_b = EventLogger(log_dir=tmp_path, file_name="b.jsonl", source=_TEST_SOURCE) + logger_b.log(_EVENT_CLS, _EVENT_PARTIAL) + logger_b.log(_EVENT_CLS, _EVENT_PARTIAL) + logger_b.close() + + events = read_events(tmp_path) + assert len(events) == 3 + + +class TestWithContext: + def test_injects_context_fields_into_logged_event(self, tmp_path: Path) -> None: + """Fields from with_context are merged into events logged inside the scope.""" + logger = _make_logger(tmp_path) + with logger.with_context({"rollout_id": 5, "attempt": 7}): + logger.log(MetricEvent, dict(metrics={"loss": 1.0})) + logger.close() + + parsed = json.loads((tmp_path / "events.jsonl").read_text().strip()) + assert parsed["rollout_id"] == 5 + assert parsed["attempt"] == 7 + assert parsed["metrics"] == {"loss": 1.0} + + def test_context_not_applied_outside_scope(self, tmp_path: Path) -> None: + """Events logged after the context scope exits do not carry context fields.""" + logger = _make_logger(tmp_path) + with logger.with_context({"rollout_id": 5}): + pass + logger.log(MetricEvent, dict(metrics={})) + logger.close() + + parsed = json.loads((tmp_path / "events.jsonl").read_text().strip()) + assert parsed["rollout_id"] is None + + def test_nesting_overrides_then_restores(self, tmp_path: Path) -> None: + """A nested with_context overrides outer fields, then restores them on exit.""" + logger = _make_logger(tmp_path) + with logger.with_context({"rollout_id": 1, "attempt": 1}): + with logger.with_context({"attempt": 2}): + logger.log(MetricEvent, dict(metrics={"k": "inner"})) + logger.log(MetricEvent, dict(metrics={"k": "outer"})) + logger.close() + + lines = (tmp_path / "events.jsonl").read_text().strip().split("\n") + inner = json.loads(lines[0]) + outer = json.loads(lines[1]) + assert (inner["rollout_id"], inner["attempt"]) == (1, 2) + assert (outer["rollout_id"], outer["attempt"]) == (1, 1) + + def test_context_var_empty_after_exit(self, tmp_path: Path) -> None: + """After all with_context scopes exit the underlying contextvar resolves to empty.""" + logger = _make_logger(tmp_path) + with logger.with_context({"rollout_id": 1}): + with logger.with_context({"attempt": 2}): + pass + assert logger._context_var.get({}) == {} + + +class TestEventLoggerContextDecorator: + def test_uninitialized_runs_without_context(self) -> None: + """When the logger is uninitialized the wrapper invokes the method directly.""" + set_event_logger(None) + calls: list[tuple] = [] + + @event_logger_context(lambda obj, x: {"rollout_id": x}) + def method(obj: object, x: int) -> int: + calls.append((obj, x)) + return x * 2 + + try: + assert method(object(), 3) == 6 + assert len(calls) == 1 + finally: + set_event_logger(None) + + def test_ctx_fn_not_called_when_uninitialized(self) -> None: + """ctx_fn is skipped entirely when the event logger is not initialized.""" + set_event_logger(None) + ctx_calls: list[int] = [] + + def ctx_fn(obj: object, x: int) -> dict: + ctx_calls.append(x) + return {"rollout_id": x} + + @event_logger_context(ctx_fn) + def method(obj: object, x: int) -> int: + return x + + try: + assert method(object(), 9) == 9 + assert ctx_calls == [] + finally: + set_event_logger(None) + + def test_initialized_injects_fields_from_method_args(self, tmp_path: Path) -> None: + """When initialized, ctx_fn output is injected into events logged by the method.""" + logger = _make_logger(tmp_path) + set_event_logger(logger) + + class Worker: + @event_logger_context(lambda self, rollout_id: {"rollout_id": rollout_id}) + def run(self, rollout_id: int) -> None: + get_event_logger().log(MetricEvent, dict(metrics={"ok": True})) + + try: + Worker().run(42) + logger.close() + finally: + set_event_logger(None) + + parsed = json.loads((tmp_path / "events.jsonl").read_text().strip()) + assert parsed["rollout_id"] == 42 + + def test_ctx_fn_receives_method_args(self, tmp_path: Path) -> None: + """ctx_fn is called with exactly the same positional/keyword args as the method.""" + logger = _make_logger(tmp_path) + set_event_logger(logger) + seen: list[tuple] = [] + + def ctx_fn(obj: object, rollout_id: int, *, attempt: int) -> dict: + seen.append((rollout_id, attempt)) + return {"rollout_id": rollout_id, "attempt": attempt} + + @event_logger_context(ctx_fn) + def method(obj: object, rollout_id: int, *, attempt: int) -> None: + get_event_logger().log(MetricEvent, dict(metrics={})) + + try: + method(object(), 3, attempt=8) + logger.close() + finally: + set_event_logger(None) + + assert seen == [(3, 8)] + parsed = json.loads((tmp_path / "events.jsonl").read_text().strip()) + assert (parsed["rollout_id"], parsed["attempt"]) == (3, 8) diff --git a/tools/convert_hf_to_torch_dist.py b/tools/convert_hf_to_torch_dist.py index 54f048a236a..c61a36dd249 100644 --- a/tools/convert_hf_to_torch_dist.py +++ b/tools/convert_hf_to_torch_dist.py @@ -14,7 +14,7 @@ from miles.backends.megatron_utils.arguments import set_default_megatron_args from miles.backends.megatron_utils.initialize import init from miles.backends.megatron_utils.model_provider import get_model_provider_func -from miles.utils.logging_utils import configure_logger +from miles.utils.logging_utils import configure_logger_raw from miles.utils.memory_utils import print_memory @@ -38,6 +38,9 @@ def get_args(): args = parse_args(add_conversion_args) args = set_default_megatron_args(args) + args.debug_deterministic_collective = False + args.enable_witness = False + # set to pass megatron validate_args args.save_interval = 1 args.micro_batch_size = 1 @@ -88,7 +91,7 @@ def main(): torch_strategy_module.FileSystemWriterAsync = ROCmFileSystemWriterAsync print("[ROCm] Applied FileSystemWriterAsync patch for HIP compatibility") - configure_logger() + configure_logger_raw() # Initialize distributed environment world_size = int(os.getenv("WORLD_SIZE") or os.getenv("SLURM_NTASKS") or 1) diff --git a/train.py b/train.py index af296622589..e380b087c33 100644 --- a/train.py +++ b/train.py @@ -7,11 +7,12 @@ from miles.utils.async_utils import eager_create_task from miles.utils.logging_utils import configure_logger from miles.utils.misc import should_run_periodic_action +from miles.utils.process_identity import MainProcessIdentity from miles.utils.tracking_utils import finish_tracking, init_tracking async def train(args): - configure_logger() + configure_logger(args, source=MainProcessIdentity()) # allocate the GPUs pgs = create_placement_groups(args) init_tracking(args) diff --git a/train_async.py b/train_async.py index a7e0063850d..85cb7aeecee 100644 --- a/train_async.py +++ b/train_async.py @@ -5,13 +5,14 @@ from miles.utils.async_utils import eager_create_task from miles.utils.logging_utils import configure_logger from miles.utils.misc import should_run_periodic_action +from miles.utils.process_identity import MainProcessIdentity from miles.utils.tracking_utils import finish_tracking, init_tracking # The framework supports other asynchronous approaches such as fully async (which is shown in examples/full_async). async def train(args): assert not args.colocate, "Colocation is not supported for async training." - configure_logger() + configure_logger(args, source=MainProcessIdentity()) # allocate the GPUs pgs = create_placement_groups(args) init_tracking(args) From 691e4a7509f8f9441e94e98a5167ec1cd80d3d10 Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Mon, 22 Jun 2026 18:15:17 +0800 Subject: [PATCH 19/41] Add event-log snapshot and restore checkpointing Add snapshot/restore for the structured event log so the event history survives cell restarts during fault-tolerant training. - miles/utils/event_logger/checkpoint.py and tests. --- miles/utils/arguments.py | 2 + miles/utils/event_logger/checkpoint.py | 64 +++++++++++++ .../utils/event_logger/test_checkpoint.py | 96 +++++++++++++++++++ 3 files changed, 162 insertions(+) create mode 100644 miles/utils/event_logger/checkpoint.py create mode 100644 tests/fast/utils/event_logger/test_checkpoint.py diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index c48a8a215d8..24f831ccb79 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -1545,6 +1545,7 @@ def add_debug_arguments(parser): "The file will be saved to `save_debug_train_data.format(rollout_id)`." ), ) + parser.add_argument("--save-debug-event-data", type=str, default=None) parser.add_argument( "--dump-details", type=str, @@ -2372,6 +2373,7 @@ def miles_validate_args(args): if args.dump_details is not None: args.save_debug_rollout_data = f"{args.dump_details}/rollout_data/{{rollout_id}}.pt" args.save_debug_train_data = f"{args.dump_details}/train_data/{{rollout_id}}_{{rank}}.pt" + args.save_debug_event_data = f"{args.dump_details}/events" if args.load_debug_rollout_data is not None: logger.info( diff --git a/miles/utils/event_logger/checkpoint.py b/miles/utils/event_logger/checkpoint.py new file mode 100644 index 00000000000..a4c838ec0f8 --- /dev/null +++ b/miles/utils/event_logger/checkpoint.py @@ -0,0 +1,64 @@ +"""Snapshot/restore the event directory alongside model checkpoints.""" + +import logging +import shutil +import time +import uuid +from argparse import Namespace +from pathlib import Path + +logger = logging.getLogger(__name__) + +_TRACKER_FILENAME = "latest_checkpointed_iteration.txt" + + +def snapshot(args: Namespace, iteration: int) -> None: + if args.save_debug_event_data is None or args.save is None: + return + + src = Path(args.save_debug_event_data) + if not src.is_dir(): + return + + dst = _snapshot_dir(Path(args.save), iteration) + if dst.exists(): + shutil.rmtree(dst) + dst.parent.mkdir(parents=True, exist_ok=True) + shutil.copytree(src, dst) + logger.info("Snapshotted event dir %s -> %s", src, dst) + + +def restore(args: Namespace) -> None: + if args.save_debug_event_data is None or args.load is None: + return + + iteration = _read_tracker_iteration(Path(args.load)) + if iteration is None: + return + + src = _snapshot_dir(Path(args.load), iteration) + if not src.is_dir(): + return + + dst = Path(args.save_debug_event_data) + if dst.exists(): + trash = dst.parent / f".trash_{time.strftime('%Y%m%d_%H%M%S')}_{uuid.uuid4().hex[:8]}" + dst.rename(trash) + logger.info("Moved pre-restore event dir %s -> %s", dst, trash) + shutil.copytree(src, dst) + logger.info("Restored event dir %s <- %s", dst, src) + + +def _snapshot_dir(checkpoint_root: Path, iteration: int) -> Path: + return checkpoint_root / f"iter_{iteration:07d}" / "debug_events" + + +def _read_tracker_iteration(checkpoint_root: Path) -> int | None: + tracker = checkpoint_root / _TRACKER_FILENAME + if not tracker.is_file(): + return None + + content = tracker.read_text().strip() + if not content.isdigit(): + return None + return int(content) diff --git a/tests/fast/utils/event_logger/test_checkpoint.py b/tests/fast/utils/event_logger/test_checkpoint.py new file mode 100644 index 00000000000..b0d9123001d --- /dev/null +++ b/tests/fast/utils/event_logger/test_checkpoint.py @@ -0,0 +1,96 @@ +"""Tests for miles.utils.event_logger.checkpoint.""" + +from argparse import Namespace +from pathlib import Path + +from miles.utils.event_logger import checkpoint as event_logger_checkpoint + + +def _args(*, event_dir: Path | None, save: Path | None = None, load: Path | None = None) -> Namespace: + return Namespace( + save_debug_event_data=str(event_dir) if event_dir else None, + save=str(save) if save else None, + load=str(load) if load else None, + ) + + +def _write_tracker(ckpt: Path, content: str) -> None: + ckpt.mkdir(parents=True, exist_ok=True) + (ckpt / "latest_checkpointed_iteration.txt").write_text(content) + + +class TestSnapshotRestoreRoundtrip: + def test_restore_replaces_live_dir_with_snapshot(self, tmp_path: Path) -> None: + """A resumed run sees exactly the snapshotted events, not the live dir's leftovers.""" + ckpt = tmp_path / "ckpt" + events = tmp_path / "events" + events.mkdir() + (events / "main.jsonl").write_text("committed\n") + event_logger_checkpoint.snapshot(_args(event_dir=events, save=ckpt), iteration=3) + + # Events written after the save (would be re-executed by the resumed run). + (events / "main.jsonl").write_text("committed\nrewound-future\n") + (events / "straggler.jsonl").write_text("late\n") + _write_tracker(ckpt, "3") + event_logger_checkpoint.restore(_args(event_dir=events, load=ckpt)) + + assert (events / "main.jsonl").read_text() == "committed\n" + assert not (events / "straggler.jsonl").exists() + + def test_snapshot_overwrites_previous_snapshot_of_same_iteration(self, tmp_path: Path) -> None: + """Re-saving the same iteration replaces its snapshot.""" + ckpt = tmp_path / "ckpt" + events = tmp_path / "events" + events.mkdir() + (events / "main.jsonl").write_text("v1\n") + event_logger_checkpoint.snapshot(_args(event_dir=events, save=ckpt), iteration=1) + (events / "main.jsonl").write_text("v2\n") + event_logger_checkpoint.snapshot(_args(event_dir=events, save=ckpt), iteration=1) + + assert (ckpt / "iter_0000001" / "debug_events" / "main.jsonl").read_text() == "v2\n" + + +class TestNoOpCases: + def test_restore_skips_when_not_resuming(self, tmp_path: Path) -> None: + """No --load means no restore.""" + events = tmp_path / "events" + events.mkdir() + (events / "main.jsonl").write_text("keep\n") + + event_logger_checkpoint.restore(_args(event_dir=events)) + + assert (events / "main.jsonl").read_text() == "keep\n" + + def test_restore_skips_when_checkpoint_has_no_snapshot(self, tmp_path: Path) -> None: + """Checkpoints predating event snapshots leave the live dir untouched.""" + ckpt = tmp_path / "ckpt" + _write_tracker(ckpt, "2") + events = tmp_path / "events" + events.mkdir() + (events / "main.jsonl").write_text("keep\n") + + event_logger_checkpoint.restore(_args(event_dir=events, load=ckpt)) + + assert (events / "main.jsonl").read_text() == "keep\n" + + def test_restore_skips_release_tracker(self, tmp_path: Path) -> None: + """A non-numeric tracker (e.g. 'release') is not a resumable iteration.""" + ckpt = tmp_path / "ckpt" + _write_tracker(ckpt, "release") + events = tmp_path / "events" + events.mkdir() + (events / "main.jsonl").write_text("keep\n") + + event_logger_checkpoint.restore(_args(event_dir=events, load=ckpt)) + + assert (events / "main.jsonl").read_text() == "keep\n" + + def test_snapshot_skips_when_events_disabled_or_no_save(self, tmp_path: Path) -> None: + """Disabled events or no save dir means no snapshot.""" + events = tmp_path / "events" + events.mkdir() + + event_logger_checkpoint.snapshot(_args(event_dir=None, save=tmp_path / "ckpt"), iteration=1) + event_logger_checkpoint.snapshot(_args(event_dir=events), iteration=1) + + assert not (tmp_path / "ckpt").exists() From 5a9115061afec97f13fb88b1aafb5092e9319b4f Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Mon, 22 Jun 2026 18:15:17 +0800 Subject: [PATCH 20/41] Log training metrics as MetricEvents through the event logger Add the `MetricEvent` model (a discriminated-union member) and emit every tracking metric into the structured event log: `tracking_utils.log` now forwards `{metrics}` to `get_event_logger().log(MetricEvent, ...)` when the event logger is initialized. --- miles/utils/tracking_utils/__init__.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/miles/utils/tracking_utils/__init__.py b/miles/utils/tracking_utils/__init__.py index 07e7023be49..9fb867da320 100644 --- a/miles/utils/tracking_utils/__init__.py +++ b/miles/utils/tracking_utils/__init__.py @@ -1,5 +1,10 @@ import logging +import torch + +from miles.utils.event_logger.logger import get_event_logger, is_event_logger_initialized +from miles.utils.event_logger.models import MetricEvent + from .base import TrackingManager logger = logging.getLogger(__name__) @@ -14,6 +19,10 @@ def log(args, metrics, step_key: str): step = metrics.get(step_key) _manager.log(metrics, step=step, step_key=step_key) + if is_event_logger_initialized(): + serializable_metrics = {k: (v.item() if isinstance(v, torch.Tensor) else v) for k, v in metrics.items()} + get_event_logger().log(MetricEvent, {"metrics": serializable_metrics}, print_log=False) + def finish_tracking(): _manager.finish() From 34fb71424da2409c3af10282195d1fe107ca0332 Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Mon, 22 Jun 2026 18:15:17 +0800 Subject: [PATCH 21/41] Add the witness id allocator Add the witness id allocator and `WitnessInfo` carrier used to assign and track witness ids for fault-tolerance verification. - miles/utils/witness/allocator.py and tests. --- miles/utils/arguments.py | 11 + miles/utils/witness/__init__.py | 0 miles/utils/witness/allocator.py | 58 +++++ tests/fast/utils/test_witness/__init__.py | 0 .../fast/utils/test_witness/test_allocator.py | 223 ++++++++++++++++++ 5 files changed, 292 insertions(+) create mode 100644 miles/utils/witness/__init__.py create mode 100644 miles/utils/witness/allocator.py create mode 100644 tests/fast/utils/test_witness/__init__.py create mode 100644 tests/fast/utils/test_witness/test_allocator.py diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index 24f831ccb79..2f9595f9f12 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -1647,6 +1647,17 @@ def add_debug_arguments(parser): help="When comparing weights after update, allow quantized tensors to differ " "by up to 1 ULP of the quantized dtype per side (compared in dequantized space).", ) + parser.add_argument( + "--enable-witness", + action="store_true", + help="Enable forward/backward pass witness.", + ) + parser.add_argument( + "--witness-buffer-size", + type=int, + default=1048576, + help="Maximum number of unique witness IDs before recycling.", + ) parser.add_argument( "--env-report", type=str, diff --git a/miles/utils/witness/__init__.py b/miles/utils/witness/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/miles/utils/witness/allocator.py b/miles/utils/witness/allocator.py new file mode 100644 index 00000000000..8cfcd1ae044 --- /dev/null +++ b/miles/utils/witness/allocator.py @@ -0,0 +1,58 @@ +from pathlib import Path + +from miles.utils.event_logger.logger import read_events +from miles.utils.event_logger.models import WitnessAllocateIdEvent +from miles.utils.pydantic_utils import FrozenStrictBaseModel + + +class WitnessInfo(FrozenStrictBaseModel): + witness_ids: list[int] + stale_ids: list[int] + + +class WitnessIdAllocator: + def __init__(self, *, buffer_size: int) -> None: + if buffer_size <= 0: + raise ValueError(f"buffer_size ({buffer_size}) must be positive.") + self._buffer_size = buffer_size + self._counter: int = 0 + + @property + def counter(self) -> int: + return self._counter + + def resume(self, counter: int) -> None: + assert counter >= self._counter + self._counter = counter + + def allocate(self, num_ids: int) -> WitnessInfo: + if num_ids < 0: + raise ValueError(f"num_ids ({num_ids}) must be non-negative.") + assert num_ids <= self._buffer_size, ( + f"num_ids ({num_ids}) exceeds buffer_size ({self._buffer_size}). " f"Increase --witness-buffer-size." + ) + ids = [(self._counter + i) % self._buffer_size for i in range(num_ids)] + stale_ids = _compute_stale_ids( + keep_count=int(self._buffer_size * 0.7), + counter=self._counter + num_ids, + buffer_size=self._buffer_size, + ) + self._counter += num_ids + return WitnessInfo(witness_ids=ids, stale_ids=stale_ids) + + +def read_persisted_witness_counter(event_dir: Path) -> int: + """Recover the allocator counter from the (checkpoint-restored) event directory.""" + events = read_events(event_dir) + return max((e.counter_after for e in events if isinstance(e, WitnessAllocateIdEvent)), default=0) + + +def _compute_stale_ids(*, keep_count: int, counter: int, buffer_size: int) -> list[int]: + if counter == 0: + return [] + num_stale = buffer_size - min(keep_count, counter, buffer_size) + if num_stale == 0: + return [] + + head = counter % buffer_size + return [(head + i) % buffer_size for i in range(num_stale)] diff --git a/tests/fast/utils/test_witness/__init__.py b/tests/fast/utils/test_witness/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/fast/utils/test_witness/test_allocator.py b/tests/fast/utils/test_witness/test_allocator.py new file mode 100644 index 00000000000..8cb671064f6 --- /dev/null +++ b/tests/fast/utils/test_witness/test_allocator.py @@ -0,0 +1,223 @@ +"""Tests for miles.utils.witness.allocator: WitnessIdAllocator, _compute_stale_ids.""" + +import json + +import pytest + +from miles.utils.witness.allocator import ( + WitnessIdAllocator, + WitnessInfo, + _compute_stale_ids, + read_persisted_witness_counter, +) + + +class TestWitnessIdAllocator: + def test_monotonic_and_wraps(self) -> None: + allocator = WitnessIdAllocator(buffer_size=5) + + info1 = allocator.allocate(3) + assert info1.witness_ids == [0, 1, 2] + + info2 = allocator.allocate(4) + assert info2.witness_ids == [3, 4, 0, 1] + + def test_allocate_returns_correct_count(self) -> None: + allocator = WitnessIdAllocator(buffer_size=100) + info = allocator.allocate(7) + assert len(info.witness_ids) == 7 + + def test_allocate_returns_witness_info(self) -> None: + allocator = WitnessIdAllocator(buffer_size=10) + info = allocator.allocate(3) + assert isinstance(info, WitnessInfo) + assert len(info.witness_ids) == 3 + assert isinstance(info.stale_ids, list) + + def test_stale_ids_computed_on_allocate(self) -> None: + """Allocating 8 from buffer_size=10 should produce stale IDs (keep 70% = 7).""" + allocator = WitnessIdAllocator(buffer_size=10) + info = allocator.allocate(8) + assert info.witness_ids == [0, 1, 2, 3, 4, 5, 6, 7] + assert set(info.stale_ids) == {8, 9, 0} + + def test_allocate_zero_ids(self) -> None: + """allocate(num_ids=0) should return empty witness_ids and reasonable stale_ids.""" + allocator = WitnessIdAllocator(buffer_size=5) + allocator.allocate(3) + info = allocator.allocate(0) + assert info.witness_ids == [] + assert isinstance(info.stale_ids, list) + + def test_non_positive_buffer_size_raises(self) -> None: + """buffer_size <= 0 raises ValueError at construction.""" + with pytest.raises(ValueError, match="must be positive"): + WitnessIdAllocator(buffer_size=0) + with pytest.raises(ValueError, match="must be positive"): + WitnessIdAllocator(buffer_size=-1) + + def test_negative_num_ids_raises(self) -> None: + """allocate(num_ids<0) raises ValueError and leaves the counter untouched.""" + allocator = WitnessIdAllocator(buffer_size=5) + with pytest.raises(ValueError, match="must be non-negative"): + allocator.allocate(num_ids=-1) + assert allocator.counter == 0 + + def test_allocate_exceeds_buffer_size_raises(self) -> None: + """num_ids > buffer_size raises AssertionError.""" + allocator = WitnessIdAllocator(buffer_size=5) + with pytest.raises(AssertionError, match="exceeds buffer_size"): + allocator.allocate(num_ids=10) + + def test_consecutive_allocations_stale_ids_evolve(self) -> None: + """Consecutive allocate calls should produce evolving stale_ids as counter grows.""" + allocator = WitnessIdAllocator(buffer_size=10) + + info1 = allocator.allocate(3) + stale1 = set(info1.stale_ids) + + info2 = allocator.allocate(3) + stale2 = set(info2.stale_ids) + + info3 = allocator.allocate(3) + stale3 = set(info3.stale_ids) + + assert stale1 != stale2 or stale2 != stale3, "stale_ids should evolve across allocations" + all_ids = set(range(10)) + for stale in [stale1, stale2, stale3]: + assert stale.issubset(all_ids) + + +class TestComputeStaleIds: + """Direct tests for _compute_stale_ids module-level function.""" + + def test_counter_zero_returns_empty(self) -> None: + assert _compute_stale_ids(keep_count=5, counter=0, buffer_size=10) == [] + + def test_counter_zero_keep_zero_returns_empty(self) -> None: + assert _compute_stale_ids(keep_count=0, counter=0, buffer_size=10) == [] + + def test_counter_zero_large_buffer_returns_empty(self) -> None: + assert _compute_stale_ids(keep_count=100, counter=0, buffer_size=1000) == [] + + def test_counter_less_than_keep_count_returns_all_unused(self) -> None: + # counter=3, buffer=10, keep=7 → active=min(7,3,10)=3 → stale=7 slots + result = _compute_stale_ids(keep_count=7, counter=3, buffer_size=10) + assert set(result) == {3, 4, 5, 6, 7, 8, 9} + + def test_keep_count_equals_buffer_size_returns_empty(self) -> None: + # All slots are active + assert _compute_stale_ids(keep_count=10, counter=15, buffer_size=10) == [] + + def test_keep_count_exceeds_buffer_size_returns_empty(self) -> None: + assert _compute_stale_ids(keep_count=20, counter=15, buffer_size=10) == [] + + def test_basic_no_wrap(self) -> None: + # counter=8, buffer=10, keep=5 → stale=5, head=8 → stale=[8,9,0,1,2] + result = _compute_stale_ids(keep_count=5, counter=8, buffer_size=10) + assert result == [8, 9, 0, 1, 2] + + def test_basic_wrap(self) -> None: + # counter=3, buffer=10, keep=5 → active=min(5,3,10)=3 → stale=7, head=3 → stale=[3,4,5,6,7,8,9] + result = _compute_stale_ids(keep_count=5, counter=3, buffer_size=10) + assert result == [3, 4, 5, 6, 7, 8, 9] + + def test_head_at_zero(self) -> None: + # counter=10, buffer=10, keep=3 → stale=7, head=0 → stale=[0,1,2,3,4,5,6] + result = _compute_stale_ids(keep_count=3, counter=10, buffer_size=10) + assert result == [0, 1, 2, 3, 4, 5, 6] + + def test_keep_one(self) -> None: + # counter=5, buffer=10, keep=1 → stale=9, head=5 → stale=[5,6,7,8,9,0,1,2,3] + result = _compute_stale_ids(keep_count=1, counter=5, buffer_size=10) + assert len(result) == 9 + assert 4 not in result + + def test_keep_zero(self) -> None: + # All slots stale + result = _compute_stale_ids(keep_count=0, counter=5, buffer_size=10) + assert len(result) == 10 + assert set(result) == set(range(10)) + + def test_stale_and_active_are_disjoint_and_cover_buffer(self) -> None: + for counter in [0, 1, 5, 10, 13, 20, 100]: + for keep in [0, 1, 3, 7, 10, 15]: + stale = _compute_stale_ids(keep_count=keep, counter=counter, buffer_size=10) + if counter == 0: + assert stale == [], "counter=0 → no stale IDs" + continue + active_count = min(keep, counter, 10) + assert len(stale) == 10 - active_count, f"counter={counter}, keep={keep}" + assert len(set(stale)) == len(stale), "no duplicates" + assert all(0 <= x < 10 for x in stale), "all in range" + + def test_buffer_size_one(self) -> None: + assert _compute_stale_ids(keep_count=1, counter=5, buffer_size=1) == [] + assert _compute_stale_ids(keep_count=0, counter=5, buffer_size=1) == [0] + + +class TestWitnessIdAllocatorResume: + def test_resume_continues_allocation_without_reusing_ids(self) -> None: + """Resuming from a persisted counter issues fresh ids, as if the run never stopped.""" + saved = WitnessIdAllocator(buffer_size=100) + saved.allocate(10) + assert saved.counter == 10 + + resumed = WitnessIdAllocator(buffer_size=100) + resumed.resume(saved.counter) + info = resumed.allocate(5) + + assert info.witness_ids == [10, 11, 12, 13, 14] + assert resumed.counter == 15 + + def test_resume_backwards_is_rejected(self) -> None: + """resume() must never move the counter backwards.""" + allocator = WitnessIdAllocator(buffer_size=100) + allocator.allocate(10) + + with pytest.raises(AssertionError): + allocator.resume(3) + + def test_counter_matches_uninterrupted_run(self) -> None: + """A save/resume sequence allocates the same ids as one uninterrupted allocator.""" + uninterrupted = WitnessIdAllocator(buffer_size=7) + expected = [uninterrupted.allocate(3).witness_ids for _ in range(4)] + + first = WitnessIdAllocator(buffer_size=7) + actual = [first.allocate(3).witness_ids, first.allocate(3).witness_ids] + second = WitnessIdAllocator(buffer_size=7) + second.resume(first.counter) + actual += [second.allocate(3).witness_ids, second.allocate(3).witness_ids] + + assert actual == expected + + +class TestReadPersistedWitnessCounter: + def _write_events(self, event_dir, counters) -> None: + event_dir.mkdir(parents=True, exist_ok=True) + lines = [] + for i, counter in enumerate(counters): + lines.append( + json.dumps( + { + "type": "witness_allocate_id", + "timestamp": "2026-06-12T00:00:00Z", + "source": {"component": "main"}, + "rollout_id": i, + "attempt": 0, + "witness_id_to_sample_index": {}, + "counter_after": counter, + } + ) + ) + (event_dir / "main.jsonl").write_text("\n".join(lines) + "\n") + + def test_reads_latest_counter_from_events(self, tmp_path) -> None: + """The max counter_after across allocate events is the resume point.""" + self._write_events(tmp_path / "events", [3, 6, 9]) + + assert read_persisted_witness_counter(tmp_path / "events") == 9 + + def test_empty_or_missing_dir_resumes_from_zero(self, tmp_path) -> None: + """A fresh run (no events yet) starts allocation at zero.""" + assert read_persisted_witness_counter(tmp_path / "missing") == 0 From fd3bc9c5a90915263714627134510c7e6ad44400 Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Mon, 22 Jun 2026 18:15:17 +0800 Subject: [PATCH 22/41] Trace witness ids through the model via injected witness parameters Thread witness ids through the model by injecting witness parameters, so the event log can later verify they propagate correctly. - miles/utils/witness/module.py, model_provider.py and tests. --- .../backends/megatron_utils/model_provider.py | 17 + .../megatron_utils/update_weight/common.py | 4 + miles/utils/witness/module.py | 256 +++++++ tests/fast/utils/test_witness/test_module.py | 638 ++++++++++++++++++ 4 files changed, 915 insertions(+) create mode 100644 miles/utils/witness/module.py create mode 100644 tests/fast/utils/test_witness/test_module.py diff --git a/miles/backends/megatron_utils/model_provider.py b/miles/backends/megatron_utils/model_provider.py index 6aa694e8761..16249147e78 100644 --- a/miles/backends/megatron_utils/model_provider.py +++ b/miles/backends/megatron_utils/model_provider.py @@ -19,6 +19,7 @@ from miles.utils.misc import load_function from miles.utils.replay_base import routing_replay_manager +from miles.utils.witness.module import install_witness logger = logging.getLogger(__name__) @@ -154,6 +155,7 @@ def wrapped_model_provider( model.output_layer = LinearForLastLayer( input_size=model.config.hidden_size, output_size=1, config=model.config ) + _maybe_install_witness(args, model) return model return wrapped_model_provider @@ -180,6 +182,7 @@ def wrapped_bridge_provider( if pg_collection is not None: provider._pg_collection = pg_collection model = provider.provide(pre_process=pre_process, post_process=post_process, vp_stage=vp_stage) + assert not getattr(args, "enable_witness", False), "Witness is not supported yet in this mode" # Gemma-4 forward returns (logits, loss_mask); keep logits only. _bridge_forward = model.forward @@ -318,6 +321,20 @@ def model_provider( if post_process and role == "critic": model.output_layer = LinearForLastLayer(input_size=config.hidden_size, output_size=1, config=config) + _maybe_install_witness(args, model) + return model return model_provider + + +def _maybe_install_witness( + args: argparse.Namespace, + model: GPTModel, +) -> None: + if getattr(args, "enable_witness", False): + install_witness( + model, + buffer_size=args.witness_buffer_size, + sequence_parallel=getattr(model.config, "sequence_parallel", False), + ) diff --git a/miles/backends/megatron_utils/update_weight/common.py b/miles/backends/megatron_utils/update_weight/common.py index dc483e5e543..2e289acbc67 100644 --- a/miles/backends/megatron_utils/update_weight/common.py +++ b/miles/backends/megatron_utils/update_weight/common.py @@ -276,6 +276,8 @@ def _compute_fqn(name, vp_stage=vp_stage): return f"vp_stages.{vp_stage}.{strip_param_name_prefix(name)}" for name, param in model_module.named_parameters(): + if getattr(param, "_is_witness_param", False): + continue yield _compute_fqn(name), param for name, buffer in model_module.named_buffers(): @@ -306,6 +308,8 @@ def _named_params_and_buffers_global( else: layer_offset = get_transformer_layer_offset(model_module.config) for name, param in model_module.named_parameters(): + if getattr(param, "_is_witness_param", False): + continue # for model without ddp wrap if not name.startswith("module.module."): name = "module." + name diff --git a/miles/utils/witness/module.py b/miles/utils/witness/module.py new file mode 100644 index 00000000000..c012f26b640 --- /dev/null +++ b/miles/utils/witness/module.py @@ -0,0 +1,256 @@ +import logging +from collections.abc import Sequence +from types import SimpleNamespace + +import torch +import torch.nn as nn +from megatron.core import tensor_parallel +from megatron.core.optimizer.distrib_optimizer import DistributedOptimizer +from megatron.core.optimizer.optimizer import ChainedOptimizer +from megatron.core.transformer.utils import sharded_state_dict_default +from torch import Tensor + +from miles.backends.training_utils.parallel import get_parallel_state +from miles.utils.event_logger.logger import get_event_logger +from miles.utils.event_logger.models import WitnessSnapshotParamEvent +from miles.utils.witness.allocator import WitnessInfo + +logger = logging.getLogger(__name__) + + +# --------------------------------------------------------------------------- +# Public API +# --------------------------------------------------------------------------- + + +def install_witness( + model: nn.Module, + *, + buffer_size: int, + sequence_parallel: bool = False, +) -> None: + model.local_head_witness = _DataWitness(buffer_size=buffer_size, sequence_parallel=sequence_parallel) + model.local_tail_witness = _DataWitness(buffer_size=buffer_size, sequence_parallel=sequence_parallel) + + +def witness_dump_and_clear_stale( + *, + model: Sequence[nn.Module], + witness_info: WitnessInfo, + optimizer: torch.optim.Optimizer, +) -> None: + """Log nonzero witness param rows, then clear stale ring buffer entries.""" + pp_rank = get_parallel_state().pp.rank + + for chunk_index, chunk in enumerate(model): + inner = _unwrap_to_witness_owner(chunk) + for attr in _WITNESS_ATTRS: + assert hasattr(inner, attr), f"chunk {chunk_index} missing {attr}" + witness: _DataWitness = getattr(inner, attr) + _record_and_log_witness_param( + witness=witness, + instance_id=f"pp{pp_rank}_chunk{chunk_index}." + attr.replace("_witness", ""), + stale_ids=witness_info.stale_ids, + ) + + _clear_witness_stale_rows(model=model, stale_ids=witness_info.stale_ids, optimizer=optimizer) + + +# --------------------------------------------------------------------------- +# Classes +# --------------------------------------------------------------------------- + + +class _DataWitness(nn.Module): + def __init__( + self, + buffer_size: int, + *, + sequence_parallel: bool = False, + ) -> None: + super().__init__() + self.buffer_size = buffer_size + self._sequence_parallel = sequence_parallel + self.witness = nn.Embedding(num_embeddings=buffer_size, embedding_dim=1) + self.witness.weight._is_witness_param = True + nn.init.zeros_(self.witness.weight) + + def forward(self, witness_ids: Tensor, hidden_states: Tensor) -> Tensor: + w = self.witness(witness_ids) # [b, s, 1] + out = w - w.detach() # forward: bitwise 0, backward: d/dw = I + + out = out.transpose(0, 1).contiguous() # [s, b, 1] + if self._sequence_parallel: + out = tensor_parallel.scatter_to_sequence_parallel_region(out) + + return _abs_broadcast_add(hidden_states, out) + + def sharded_state_dict(self, prefix: str = "", sharded_offsets: tuple = (), metadata: object = None) -> dict: + pp_rank = get_parallel_state().pp.rank + # Embed PP rank in the checkpoint key so each pipeline stage has a unique + # key (e.g. local_head_witness_pp0.witness.weight vs _pp1.witness.weight). + # Without this, PP>1 causes a sharding validation error because multiple + # stages register the same key with identical replica_id. + prefix_with_pp = f"{prefix.rstrip('.')}_pp{pp_rank}." + + # Delegate to Megatron's sharded_state_dict_default (utils.py). + # Use SimpleNamespace so it takes the `else` branch (no sharded_state_dict attr) + # instead of recursing back into this method. + return sharded_state_dict_default( + module=SimpleNamespace(state_dict=self.state_dict), + prefix=prefix_with_pp, + sharded_offsets=sharded_offsets, + metadata=metadata, + tp_group=get_parallel_state().tp.group, + ) + + +def _abs_broadcast_add(hidden_states: Tensor, addend: Tensor) -> Tensor: + return _AbsBroadcastAdd.apply(hidden_states, addend) + + +class _AbsBroadcastAdd(torch.autograd.Function): + """Broadcast-add a low-dim addend to a high-dim tensor, using abs-reduced gradient for the addend. + + Forward: ``hidden_states + addend`` (standard broadcast). + Backward for ``hidden_states``: pass-through. + Backward for ``addend``: ``grad.abs().sum(dim=-1, keepdim=True)`` instead of ``grad.sum(dim=-1, keepdim=True)``. + + This avoids gradient cancellation when the upstream gradient has mixed signs + across the last dimension. The witness embedding only needs to detect + *whether* gradient flowed (nonzero), not the exact magnitude, so using + ``abs`` is acceptable. + """ + + @staticmethod + def forward(ctx: torch.autograd.function.FunctionCtx, hidden_states: Tensor, addend: Tensor) -> Tensor: + assert addend.shape[-1] == 1, f"addend last dim must be 1, got {addend.shape}" + assert hidden_states.shape[:-1] == addend.shape[:-1], ( + f"hidden_states and addend must match on all dims except last, " + f"got {hidden_states.shape} vs {addend.shape}" + ) + return hidden_states + addend + + @staticmethod + def backward(ctx: torch.autograd.function.FunctionCtx, grad: Tensor) -> tuple[Tensor, Tensor]: + grad_addend = grad.abs().sum(dim=-1, keepdim=True) + return grad, grad_addend + + +# --------------------------------------------------------------------------- +# Private helpers +# --------------------------------------------------------------------------- + + +_WITNESS_ATTRS = ("local_head_witness", "local_tail_witness") + + +def _has_any_witness(module: nn.Module) -> bool: + return any(hasattr(module, attr) for attr in _WITNESS_ATTRS) + + +def _unwrap_to_witness_owner(chunk: nn.Module) -> nn.Module: + """Navigate through wrapping layers (DDP → Float16Module → GPTModel) to find the module with witness attrs.""" + inner = chunk.module + while not _has_any_witness(inner) and hasattr(inner, "module"): + inner = inner.module + return inner + + +def _clear_witness_stale_rows( + *, + model: Sequence[nn.Module], + stale_ids: list[int], + optimizer: torch.optim.Optimizer, +) -> None: + if not stale_ids: + return + + witnesses = list(_get_all_witnesses_in_model(model)) + for witness in witnesses: + idx = torch.tensor(stale_ids, dtype=torch.long, device=witness.witness.weight.device) + _zero_witness_rows(witness=witness, idx=idx, optimizer=optimizer) + + +def _get_all_witnesses_in_model(model_chunks: Sequence[nn.Module]) -> list[_DataWitness]: + witnesses: list[_DataWitness] = [] + for chunk in model_chunks: + inner = _unwrap_to_witness_owner(chunk) + for attr in _WITNESS_ATTRS: + assert hasattr(inner, attr), f"model chunk missing {attr}" + witnesses.append(getattr(inner, attr)) + return witnesses + + +def _zero_witness_rows(*, witness: _DataWitness, idx: Tensor, optimizer: torch.optim.Optimizer) -> None: + model_weight = witness.witness.weight + model_weight.data[idx] = 0.0 + + for inner_optimizer in _iter_inner_optimizers(optimizer): + # miles forces use_distributed_optimizer, so anything else is unreachable. + assert isinstance( + inner_optimizer, DistributedOptimizer + ), f"unsupported optimizer: {type(inner_optimizer).__name__}" + _zero_rows_in_distributed_optimizer(optimizer=inner_optimizer, model_param=model_weight, idx=idx) + + +def _iter_inner_optimizers(optimizer: torch.optim.Optimizer) -> list[torch.optim.Optimizer]: + if isinstance(optimizer, ChainedOptimizer): + return list(optimizer.chained_optimizers) + return [optimizer] + + +def _zero_rows_in_distributed_optimizer(*, optimizer: DistributedOptimizer, model_param: Tensor, idx: Tensor) -> None: + assert not optimizer.config.use_precision_aware_optimizer_no_fp8_or_ds_fp8 + assert not optimizer.config.optimizer_cpu_offload, "HybridDeviceOptimizer state layout is not supported" + assert optimizer.config.optimizer == "adam", f"unsupported optimizer kernel: {optimizer.config.optimizer}" + if model_param not in optimizer.model_param_gbuf_map: + # This dist-opt instance (e.g. the expert one) or this rank owns no shard of the param. + return + + # The fp32 main weights are flat shards of the flattened model param; + # embedding_dim == 1 makes flattened offsets equal witness row ids. + assert model_param.shape[-1] == 1, f"witness weight last dim must be 1, got {model_param.shape}" + param_range = optimizer._get_model_param_range_map(model_param)["param"] + local_idx = idx[(idx >= param_range.start) & (idx < param_range.end)] - param_range.start + if local_idx.numel() == 0: + return + + group_index, group_order = optimizer.model_param_group_index_map[model_param] + main_param = optimizer.optimizer.param_groups[group_index]["params"][group_order] + assert main_param.numel() == param_range.size + main_param.data[local_idx] = 0.0 + + state = optimizer.optimizer.state + if main_param not in state: + # An optimizer that never stepped has no per-param state yet; a populated state + # missing the witness entry means we are clearing the wrong key — fail loudly. + assert len(state) == 0, f"witness main shard missing from optimizer state with {len(state)} entries" + return + + param_state = state[main_param] + for key in ("exp_avg", "exp_avg_sq"): + assert key in param_state, f"expected Adam state key {key!r}, got {sorted(param_state)}" + param_state[key][local_idx] = 0.0 + + +def _record_and_log_witness_param( + *, + witness: _DataWitness, + instance_id: str, + stale_ids: list[int], +) -> None: + model_weight = witness.witness.weight + main_param = getattr(model_weight, "main_param", None) + check_weight = main_param.data if main_param is not None else model_weight.data + nonzero_witness_ids: list[int] = check_weight.squeeze(-1).nonzero(as_tuple=True)[0].tolist() + + get_event_logger().log( + WitnessSnapshotParamEvent, + dict( + instance_id=instance_id, + nonzero_witness_ids=nonzero_witness_ids, + stale_ids=stale_ids, + ), + print_log=False, + ) diff --git a/tests/fast/utils/test_witness/test_module.py b/tests/fast/utils/test_witness/test_module.py new file mode 100644 index 00000000000..35b7133ed5d --- /dev/null +++ b/tests/fast/utils/test_witness/test_module.py @@ -0,0 +1,638 @@ +"""Tests for miles.utils.witness.module.""" + +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest +import torch +import torch.nn as nn +from megatron.core.optimizer.distrib_optimizer import DistributedOptimizer, Range +from megatron.core.optimizer.optimizer import ChainedOptimizer + +from miles.utils.event_logger.models import WitnessSnapshotParamEvent +from miles.utils.witness.allocator import WitnessInfo +from miles.utils.witness.module import ( + _abs_broadcast_add, + _AbsBroadcastAdd, + _DataWitness, + _record_and_log_witness_param, + _zero_witness_rows, + install_witness, + witness_dump_and_clear_stale, +) + + +class TestDataWitnessForward: + def test_forward_does_not_change_hidden_states(self) -> None: + """Witness output is zero, so hidden_states should be unchanged.""" + witness = _DataWitness(buffer_size=10) + ids = torch.tensor([[0, 1, 2, 3]]) # [1, 4] + hidden = torch.randn(4, 1, 8) # [s, b, h] Megatron SBH layout + result = witness(ids, hidden) + assert torch.equal(result, hidden) + + def test_forward_unchanged_after_optimizer_step(self) -> None: + witness = _DataWitness(buffer_size=10) + optimizer = torch.optim.Adam(witness.parameters(), lr=0.1) + + ids = torch.tensor([[0, 1, 2]]) + hidden = torch.randn(3, 1, 8) + result = witness(ids, hidden) + result.sum().backward() + optimizer.step() + optimizer.zero_grad() + + # After optimizer update, weights are nonzero, but hidden_states still unchanged + assert not torch.all(witness.witness.weight == 0.0) + result2 = witness(ids, hidden) + assert torch.equal(result2, hidden) + + def test_backward_records_gradient_on_witness_weight(self) -> None: + witness = _DataWitness(buffer_size=10) + ids = torch.tensor([[2, 5]]) + hidden = torch.randn(2, 1, 4, requires_grad=True) + + result = witness(ids, hidden) + result.sum().backward() + + grad = witness.witness.weight.grad + assert grad is not None + nonzero_rows = grad.squeeze(-1).nonzero(as_tuple=True)[0].tolist() + assert set(nonzero_rows) == {2, 5} + + def test_no_effect_on_main_model_gradients(self) -> None: + """Witness should not alter gradients for upstream or downstream model parameters.""" + torch.manual_seed(42) + embed = nn.Embedding(100, 8) + linear = nn.Linear(8, 1) + + tokens = torch.tensor([[1, 2, 3, 4]]) + + # Step 1: Compute loss without witness + hidden = embed(tokens).transpose(0, 1).contiguous() # [s=4, b=1, h=8] + out_no_witness = linear(hidden).sum() + out_no_witness.backward() + grad_embed_no = embed.weight.grad.clone() + grad_linear_no = linear.weight.grad.clone() + + embed.zero_grad() + linear.zero_grad() + + # Step 2: Compute loss with witness + hidden = embed(tokens).transpose(0, 1).contiguous() + witness = _DataWitness(buffer_size=10) + ids = torch.tensor([[0, 0, 0, 0]]) + h = witness(ids, hidden) + out_with_witness = linear(h).sum() + out_with_witness.backward() + + assert torch.equal(grad_embed_no, embed.weight.grad) + assert torch.equal(grad_linear_no, linear.weight.grad) + + +class TestRecordAndLogWitnessParam: + def test_logs_nonzero_weight_rows(self) -> None: + witness = _DataWitness(buffer_size=10) + witness.witness.weight.data[3] = 1.0 + witness.witness.weight.data[7] = 2.0 + + with patch("miles.utils.witness.module.get_event_logger") as mock_get_logger: + mock_logger = MagicMock() + mock_get_logger.return_value = mock_logger + + _record_and_log_witness_param(witness=witness, instance_id="pp0.head", stale_ids=[]) + + mock_logger.log.assert_called_once() + # New API: log(event_cls, partial_dict) + partial = mock_logger.log.call_args[0][1] + assert set(partial["nonzero_witness_ids"]) == {3, 7} + assert partial["instance_id"] == "pp0.head" + + def test_record_and_log_event_fields(self) -> None: + witness = _DataWitness(buffer_size=10) + witness.witness.weight.data[1] = 0.5 + witness.witness.weight.data[4] = -0.3 + + with patch("miles.utils.witness.module.get_event_logger") as mock_get_logger: + mock_logger = MagicMock() + mock_get_logger.return_value = mock_logger + + _record_and_log_witness_param(witness=witness, instance_id="pp0.tail", stale_ids=[]) + + mock_logger.log.assert_called_once() + assert mock_logger.log.call_args[0][0] is WitnessSnapshotParamEvent + partial = mock_logger.log.call_args[0][1] + assert partial["instance_id"] == "pp0.tail" + assert set(partial["nonzero_witness_ids"]) == {1, 4} + + +# --------------------------------------------------------------------------- +# Fake GPTModel for install_witness / forward integration tests +# --------------------------------------------------------------------------- + + +class _FakeDecoder(nn.Module): + def __init__(self) -> None: + super().__init__() + self.input_tensor: torch.Tensor | None = None + + def forward(self, hidden_states: torch.Tensor, **kwargs) -> torch.Tensor: + return hidden_states + + +class _FakeGPTModel(nn.Module): + def __init__(self, *, pre_process: bool = True) -> None: + super().__init__() + self.pre_process = pre_process + self.decoder = _FakeDecoder() + self.embedding = nn.Embedding(100, 16) + + def forward(self, input_ids: torch.Tensor, witness_ids: torch.Tensor | None = None) -> torch.Tensor: + if self.pre_process: + # Megatron decoders use sequence-first [s, b, h] layout, which is what + # _DataWitness expects (it transposes its own output to [s, b, 1]). + decoder_input = self.embedding(input_ids).transpose(0, 1).contiguous() + else: + decoder_input = None + + if hasattr(self, "local_head_witness") and witness_ids is not None: + if decoder_input is not None: + decoder_input = self.local_head_witness(witness_ids, decoder_input) + else: + self.decoder.input_tensor = self.local_head_witness(witness_ids, self.decoder.input_tensor) + + if decoder_input is None: + decoder_input = self.decoder.input_tensor + + return self.decoder(hidden_states=decoder_input) + + +class TestInstallWitness: + def test_witness_is_submodule(self) -> None: + model = _FakeGPTModel() + install_witness(model, buffer_size=10) + assert "local_head_witness" in dict(model.named_modules()) + assert "local_tail_witness" in dict(model.named_modules()) + + def test_witness_in_parameters(self) -> None: + model = _FakeGPTModel() + install_witness(model, buffer_size=10) + param_names = [name for name, _ in model.named_parameters()] + assert any("local_head_witness" in name for name in param_names) + assert any("local_tail_witness" in name for name in param_names) + + def test_forward_adds_zero(self) -> None: + model = _FakeGPTModel() + install_witness(model, buffer_size=10) + tokens = torch.tensor([[1, 2, 3]]) + out_no = model(tokens) + out_with = model(tokens, witness_ids=torch.tensor([[0, 1, 2]])) + assert torch.equal(out_no, out_with) + + def test_forward_produces_gradient(self) -> None: + model = _FakeGPTModel() + install_witness(model, buffer_size=10) + tokens = torch.tensor([[1, 2, 3]]) + out = model(tokens, witness_ids=torch.tensor([[5, 5, 5]])) + out.sum().backward() + grad = model.local_head_witness.witness.weight.grad + assert grad is not None + assert 5 in grad.squeeze(-1).nonzero(as_tuple=True)[0].tolist() + + def test_no_witness_ids_no_effect(self) -> None: + model = _FakeGPTModel() + install_witness(model, buffer_size=10) + out = model(torch.tensor([[1, 2, 3]])) + assert out is not None + + def test_witness_in_state_dict(self) -> None: + model = _FakeGPTModel() + install_witness(model, buffer_size=10) + sd = model.state_dict() + assert any("local_head_witness" in k for k in sd) + assert any("local_tail_witness" in k for k in sd) + + def test_checkpoint_roundtrip(self) -> None: + model = _FakeGPTModel() + install_witness(model, buffer_size=10) + model.local_head_witness.witness.weight.data.fill_(42.0) + sd = model.state_dict() + + model2 = _FakeGPTModel() + install_witness(model2, buffer_size=10) + model2.load_state_dict(sd) + assert torch.equal(model2.local_head_witness.witness.weight.data, model.local_head_witness.witness.weight.data) + + def test_disabled_no_submodule(self) -> None: + model = _FakeGPTModel() + assert not hasattr(model, "local_head_witness") + + def test_middle_pp_stage_modifies_input_tensor(self) -> None: + model = _FakeGPTModel(pre_process=False) + install_witness(model, buffer_size=10) + # Sequence-first: [s=4, b=1, h=16]. + hidden = torch.randn(4, 1, 16) + model.decoder.input_tensor = hidden.clone() + out = model(torch.tensor([[1, 2, 3, 4]]), witness_ids=torch.tensor([[0, 1, 2, 3]])) + assert torch.equal(out, hidden) + + def test_middle_pp_stage_produces_gradient(self) -> None: + model = _FakeGPTModel(pre_process=False) + install_witness(model, buffer_size=10) + # Sequence-first: [s=4, b=1, h=16]. + model.decoder.input_tensor = torch.randn(4, 1, 16, requires_grad=True) + out = model(torch.tensor([[1, 2, 3, 4]]), witness_ids=torch.tensor([[5, 5, 5, 5]])) + out.sum().backward() + assert 5 in model.local_head_witness.witness.weight.grad.squeeze(-1).nonzero(as_tuple=True)[0].tolist() + + def test_forward_bitwise_zero_bf16(self) -> None: + witness = _DataWitness(buffer_size=10).to(dtype=torch.bfloat16) + witness.witness.weight.data.fill_(3.14) + ids = torch.tensor([[0, 1, 2]]) + hidden = torch.randn(3, 1, 8, dtype=torch.bfloat16) + result = witness(ids, hidden) + assert torch.equal(result, hidden) + + +class TestZeroWitnessRows: + def test_zero_witness_rows_clears_distributed_optimizer_shard_rows(self) -> None: + """Stale rows inside the local dist-opt shard are zeroed in the fp32 shard and its Adam state.""" + witness = _DataWitness(buffer_size=10) + witness.witness.weight.data.fill_(1.0) + dist_opt, main_shard = _make_distributed_optimizer(witness.witness.weight, start=4, end=10) + optimizer = ChainedOptimizer([dist_opt]) + + _zero_witness_rows(witness=witness, idx=torch.tensor([2, 5, 7]), optimizer=optimizer) + + weight = witness.witness.weight + for stale_row in (2, 5, 7): + assert weight.data[stale_row].item() == 0.0 + assert weight.data[4].item() == 1.0 + # The shard covers rows [4, 10): row 5 -> local 1, row 7 -> local 3; row 2 is outside. + state = dist_opt.optimizer.state[main_shard] + for tensor in (main_shard, state["exp_avg"], state["exp_avg_sq"]): + assert tensor[1].item() == 0.0 + assert tensor[3].item() == 0.0 + assert tensor[0].item() == 1.0 + + def test_zero_witness_rows_skips_distributed_optimizer_without_the_param(self) -> None: + """A chained dist-opt instance owning no shard of the witness param is skipped, the owner still clears.""" + witness = _DataWitness(buffer_size=10) + witness.witness.weight.data.fill_(1.0) + dist_opt, main_shard = _make_distributed_optimizer(witness.witness.weight, start=0, end=10) + expert_opt = DistributedOptimizer.__new__(DistributedOptimizer) + expert_opt.config = dist_opt.config + expert_opt.model_param_gbuf_map = {} + optimizer = ChainedOptimizer([dist_opt, expert_opt]) + + _zero_witness_rows(witness=witness, idx=torch.tensor([3]), optimizer=optimizer) + + assert witness.witness.weight.data[3].item() == 0.0 + assert main_shard[3].item() == 0.0 + + def test_zero_witness_rows_rejects_non_distributed_optimizer(self) -> None: + """Any optimizer other than (chained) DistributedOptimizer is rejected outright.""" + witness = _DataWitness(buffer_size=10) + optimizer = torch.optim.Adam(witness.parameters(), lr=0.01) + + with pytest.raises(AssertionError, match="unsupported optimizer"): + _zero_witness_rows(witness=witness, idx=torch.tensor([1]), optimizer=optimizer) + + def test_zero_witness_rows_raises_when_populated_state_misses_witness_shard(self) -> None: + """A populated optimizer state without the witness shard entry fails loudly instead of silently skipping.""" + witness = _DataWitness(buffer_size=10) + dist_opt, _main_shard = _make_distributed_optimizer(witness.witness.weight, start=0, end=10) + unrelated = torch.ones(4) + dist_opt.optimizer.state = {unrelated: {"exp_avg": torch.ones(4), "exp_avg_sq": torch.ones(4)}} + + with pytest.raises(AssertionError, match="witness main shard missing"): + _zero_witness_rows(witness=witness, idx=torch.tensor([1]), optimizer=ChainedOptimizer([dist_opt])) + + +# --------------------------------------------------------------------------- +# Helpers for witness_dump_and_clear_stale tests +# --------------------------------------------------------------------------- + + +def _make_fake_chunk(buffer_size: int = 10) -> nn.Module: + """Create a fake model chunk with .module.local_head_witness and .module.local_tail_witness.""" + inner = nn.Module() + inner.local_head_witness = _DataWitness(buffer_size=buffer_size) + inner.local_tail_witness = _DataWitness(buffer_size=buffer_size) + chunk = nn.Module() + chunk.module = inner + return chunk + + +def _make_distributed_optimizer( + model_weight: nn.Parameter, + *, + start: int, + end: int, +) -> tuple[DistributedOptimizer, torch.Tensor]: + """Build a minimal fake DistributedOptimizer owning rows [start, end) of model_weight.""" + optimizer = DistributedOptimizer.__new__(DistributedOptimizer) + main_shard = model_weight.detach().view(-1)[start:end].clone().float().requires_grad_() + inner = torch.optim.Adam([main_shard], lr=0.01) + inner.state[main_shard] = { + "exp_avg": torch.ones(end - start), + "exp_avg_sq": torch.ones(end - start), + } + optimizer.optimizer = inner + optimizer.config = SimpleNamespace( + use_precision_aware_optimizer_no_fp8_or_ds_fp8=False, + optimizer_cpu_offload=False, + optimizer="adam", + ) + optimizer.model_param_gbuf_map = {model_weight: (0, torch.bfloat16, 0)} + optimizer.gbuf_ranges = [{torch.bfloat16: [{"param_map": {model_weight: {"param": Range(start, end)}}}]}] + optimizer.model_param_group_index_map = {model_weight: (0, 0)} + return optimizer, main_shard + + +def _make_chained_optimizer_for_witnesses(model: list[nn.Module]) -> ChainedOptimizer: + """Build a fake ChainedOptimizer with one fully-owning dist-opt instance per witness weight.""" + weights = [ + getattr(chunk.module, attr).witness.weight + for chunk in model + for attr in ("local_head_witness", "local_tail_witness") + ] + return ChainedOptimizer([_make_distributed_optimizer(w, start=0, end=w.numel())[0] for w in weights]) + + +class TestWitnessDumpAndClearStale: + def test_witness_dump_and_clear_stale_logs_all_witnesses(self) -> None: + """2 chunks x 2 witnesses = 4 log events with correct instance_ids.""" + chunk0 = _make_fake_chunk() + chunk1 = _make_fake_chunk() + chunk0.module.local_head_witness.witness.weight.data[1] = 1.0 + chunk0.module.local_tail_witness.witness.weight.data[2] = 1.0 + chunk1.module.local_head_witness.witness.weight.data[3] = 1.0 + chunk1.module.local_tail_witness.witness.weight.data[4] = 1.0 + + model = [chunk0, chunk1] + optimizer = _make_chained_optimizer_for_witnesses(model) + witness_info = WitnessInfo(witness_ids=[1, 2, 3, 4], stale_ids=[5, 6]) + + with patch("miles.utils.witness.module.get_event_logger") as mock_get_logger, patch( + "miles.utils.witness.module.get_parallel_state" + ) as mock_get_parallel_state: + mock_get_parallel_state.return_value.pp.rank = 0 + mock_logger = MagicMock() + mock_get_logger.return_value = mock_logger + + witness_dump_and_clear_stale(model=model, witness_info=witness_info, optimizer=optimizer) + + assert mock_logger.log.call_count == 4 + logged_instance_ids = [call[0][1]["instance_id"] for call in mock_logger.log.call_args_list] + assert logged_instance_ids == [ + "pp0_chunk0.local_head", + "pp0_chunk0.local_tail", + "pp0_chunk1.local_head", + "pp0_chunk1.local_tail", + ] + + logged_stale_ids = [call[0][1]["stale_ids"] for call in mock_logger.log.call_args_list] + for stale in logged_stale_ids: + assert stale == [5, 6] + + def test_witness_dump_and_clear_stale_clears_stale_rows(self) -> None: + """Stale IDs should have their weight rows zeroed after the call.""" + chunk = _make_fake_chunk(buffer_size=10) + chunk.module.local_head_witness.witness.weight.data.fill_(1.0) + chunk.module.local_tail_witness.witness.weight.data.fill_(1.0) + + model = [chunk] + optimizer = _make_chained_optimizer_for_witnesses(model) + witness_info = WitnessInfo(witness_ids=[0], stale_ids=[3, 7]) + + with patch("miles.utils.witness.module.get_event_logger") as mock_get_logger, patch( + "miles.utils.witness.module.get_parallel_state" + ) as mock_get_parallel_state: + mock_get_parallel_state.return_value.pp.rank = 0 + mock_get_logger.return_value = MagicMock() + witness_dump_and_clear_stale(model=model, witness_info=witness_info, optimizer=optimizer) + + for witness_attr in ("local_head_witness", "local_tail_witness"): + witness = getattr(chunk.module, witness_attr) + assert witness.witness.weight.data[3].item() == 0.0 + assert witness.witness.weight.data[7].item() == 0.0 + assert witness.witness.weight.data[0].item() == 1.0 + + def test_witness_dump_and_clear_stale_empty_stale_ids(self) -> None: + """Empty stale_ids should not trigger any zeroing.""" + chunk = _make_fake_chunk(buffer_size=10) + chunk.module.local_head_witness.witness.weight.data.fill_(1.0) + chunk.module.local_tail_witness.witness.weight.data.fill_(1.0) + + model = [chunk] + optimizer = _make_chained_optimizer_for_witnesses(model) + witness_info = WitnessInfo(witness_ids=[0], stale_ids=[]) + + with patch("miles.utils.witness.module.get_event_logger") as mock_get_logger, patch( + "miles.utils.witness.module.get_parallel_state" + ) as mock_get_parallel_state: + mock_get_parallel_state.return_value.pp.rank = 0 + mock_get_logger.return_value = MagicMock() + witness_dump_and_clear_stale(model=model, witness_info=witness_info, optimizer=optimizer) + + for witness_attr in ("local_head_witness", "local_tail_witness"): + witness = getattr(chunk.module, witness_attr) + assert torch.all(witness.witness.weight.data == 1.0) + + def test_record_and_log_witness_param_includes_stale_ids(self) -> None: + """Log event should contain the correct stale_ids field.""" + witness = _DataWitness(buffer_size=10) + witness.witness.weight.data[2] = 1.0 + + with patch("miles.utils.witness.module.get_event_logger") as mock_get_logger: + mock_logger = MagicMock() + mock_get_logger.return_value = mock_logger + + _record_and_log_witness_param(witness=witness, instance_id="pp0.head", stale_ids=[8, 9]) + + mock_logger.log.assert_called_once() + partial = mock_logger.log.call_args[0][1] + assert partial["stale_ids"] == [8, 9] + + +class TestAbsBroadcastAddForward: + def test_forward_value_matches_plain_addition(self) -> None: + hidden = torch.randn(4, 2, 8) + addend = torch.randn(4, 2, 1) + result = _abs_broadcast_add(hidden, addend) + expected = hidden + addend + assert torch.equal(result, expected) + + def test_forward_preserves_zero_addend(self) -> None: + hidden = torch.randn(4, 2, 8) + addend = torch.zeros(4, 2, 1) + result = _abs_broadcast_add(hidden, addend) + assert torch.equal(result, hidden) + + def test_forward_assert_addend_last_dim_must_be_1(self) -> None: + hidden = torch.randn(4, 2, 8) + addend = torch.randn(4, 2, 3) + with pytest.raises(AssertionError, match="addend last dim must be 1"): + _abs_broadcast_add(hidden, addend) + + def test_forward_assert_leading_dims_must_match(self) -> None: + hidden = torch.randn(4, 2, 8) + addend = torch.randn(4, 3, 1) # second dim differs + with pytest.raises(AssertionError, match="must match on all dims except last"): + _abs_broadcast_add(hidden, addend) + + def test_forward_assert_ndim_must_match(self) -> None: + hidden = torch.randn(4, 2, 8) + addend = torch.randn(2, 1) # 2D vs 3D + with pytest.raises(AssertionError, match="must match on all dims except last"): + _abs_broadcast_add(hidden, addend) + + +class TestAbsBroadcastAddBackwardHiddenStates: + def test_hidden_states_gradient_is_pass_through(self) -> None: + hidden = torch.randn(4, 2, 8, requires_grad=True) + addend = torch.randn(4, 2, 1, requires_grad=True) + result = _abs_broadcast_add(hidden, addend) + loss = result.sum() + loss.backward() + # hidden gradient = all ones (pass-through from sum) + assert torch.equal(hidden.grad, torch.ones_like(hidden)) + + +class TestAbsBroadcastAddBackwardAddend: + def test_addend_gradient_is_abs_sum_over_last_dim(self) -> None: + hidden = torch.randn(3, 2, 4, requires_grad=True) + addend = torch.zeros(3, 2, 1, requires_grad=True) + + result = _abs_broadcast_add(hidden, addend) + # Use a loss that produces a known gradient at result + upstream_grad = torch.tensor( + [ + [[1.0, -2.0, 3.0, -4.0], [0.5, -0.5, 0.5, -0.5]], + [[1.0, 1.0, 1.0, 1.0], [-1.0, -1.0, -1.0, -1.0]], + [[0.0, 0.0, 0.0, 0.0], [2.0, -1.0, 0.5, -0.3]], + ] + ) + result.backward(upstream_grad) + + # Expected addend grad: abs(upstream_grad).sum(dim=-1, keepdim=True) + expected = upstream_grad.abs().sum(dim=-1, keepdim=True) + assert torch.allclose(addend.grad, expected) + + def test_addend_gradient_no_cancellation_with_mixed_signs(self) -> None: + """The key property: mixed-sign gradients don't cancel to zero.""" + hidden = torch.randn(1, 1, 8, requires_grad=True) + addend = torch.zeros(1, 1, 1, requires_grad=True) + + result = _abs_broadcast_add(hidden, addend) + # Gradient with equal positive and negative values (sums to zero normally) + upstream_grad = torch.tensor([[[1.0, -1.0, 1.0, -1.0, 1.0, -1.0, 1.0, -1.0]]]) + result.backward(upstream_grad) + + # Plain broadcast backward would give sum = 0 + assert upstream_grad.sum(dim=-1, keepdim=True).item() == 0.0 + # But abs broadcast gives sum of absolute values = 8 + assert addend.grad.item() == 8.0 + + def test_addend_gradient_matches_plain_sum_when_all_positive(self) -> None: + """When all gradients are positive, abs().sum() == sum().""" + hidden = torch.randn(2, 1, 4, requires_grad=True) + addend = torch.zeros(2, 1, 1, requires_grad=True) + + result = _abs_broadcast_add(hidden, addend) + upstream_grad = torch.abs(torch.randn(2, 1, 4)) # all positive + result.backward(upstream_grad) + + expected = upstream_grad.sum(dim=-1, keepdim=True) + assert torch.allclose(addend.grad, expected) + + def test_addend_gradient_always_non_negative(self) -> None: + """abs().sum() is always >= 0.""" + hidden = torch.randn(10, 5, 16, requires_grad=True) + addend = torch.zeros(10, 5, 1, requires_grad=True) + + result = _abs_broadcast_add(hidden, addend) + upstream_grad = torch.randn(10, 5, 16) + result.backward(upstream_grad) + + assert (addend.grad >= 0).all() + + +class TestAbsBroadcastAddGradientFlow: + def test_gradient_flows_through_to_embedding(self) -> None: + """End-to-end: gradient reaches embedding weight via abs broadcast add.""" + vocab_size = 16 + hidden_dim = 8 + seq_len = 4 + + embedding = torch.nn.Embedding(vocab_size, 1) + torch.nn.init.zeros_(embedding.weight) + output_layer = torch.nn.Linear(hidden_dim, vocab_size, bias=False) + + witness_ids = torch.tensor([[0, 1, 2, 3]]) # [1, seq_len] + w = embedding(witness_ids) # [1, 4, 1] + out = w - w.detach() + + hidden_states = torch.randn(seq_len, 1, hidden_dim, requires_grad=True) + tail_out = out.transpose(0, 1).contiguous() # [4, 1, 1] + + combined = _abs_broadcast_add(hidden_states, tail_out) + logits = output_layer(combined) + loss = logits.sum() + loss.backward() + + # All 4 witness_ids should have nonzero gradient + nonzero_rows = (embedding.weight.grad.abs() > 0).squeeze(-1) + assert nonzero_rows[:4].all(), f"Expected rows 0-3 nonzero, got {embedding.weight.grad[:4]}" + + def test_gradient_nonzero_even_when_plain_broadcast_cancels(self) -> None: + """Simulates the exact scenario: output_layer gradient cancels under plain broadcast.""" + vocab_size = 4 + hidden_dim = 4 + seq_len = 2 + + embedding = torch.nn.Embedding(8, 1) + torch.nn.init.zeros_(embedding.weight) + + # Construct output_layer weight where ROW sums are constant (each row sums to 4), + # so that dL/d_combined.sum_over_last_dim = sum_i dL/dlogit_i * row_sum_i = 0 + # (using softmax-loss sum-to-zero property). Row values differ so per-element + # grads are nonzero — only their sum cancels. This is exactly the scenario + # where plain broadcast (sum) cancels but abs broadcast (abs.sum) preserves. + W = torch.tensor([[2.0, 0.0, 1.0, 1.0], [0.0, 2.0, 1.0, 1.0], [1.0, 1.0, 2.0, 0.0], [1.0, 1.0, 0.0, 2.0]]) + output_layer = torch.nn.Linear(hidden_dim, vocab_size, bias=False) + output_layer.weight.data = W + + witness_ids = torch.tensor([[0, 1]]) + w = embedding(witness_ids) + out = w - w.detach() + + hidden_states = torch.randn(seq_len, 1, hidden_dim, requires_grad=True) + tail_out = out.transpose(0, 1).contiguous() + + # With abs broadcast add, gradient should be nonzero + combined = _abs_broadcast_add(hidden_states, tail_out) + logits = output_layer(combined) + targets = torch.tensor([[0, 1]]) + log_probs = torch.nn.functional.log_softmax(logits.squeeze(1), dim=-1) + loss = -log_probs.gather(1, targets.T).sum() + loss.backward() + + assert embedding.weight.grad is not None + assert ( + embedding.weight.grad[:2].abs() > 0 + ).all(), f"Expected nonzero grad for witness rows 0,1, got {embedding.weight.grad[:2]}" + + +class TestAbsBroadcastAddDoubleBackward: + def test_gradcheck(self) -> None: + """Verify numerical gradient correctness with torch.autograd.gradcheck.""" + hidden = torch.randn(2, 2, 4, dtype=torch.float64, requires_grad=True) + addend = torch.randn(2, 2, 1, dtype=torch.float64, requires_grad=True) + # gradcheck only tests hidden_states gradient (which is pass-through) + # For addend, abs() is not differentiable at 0, so we test separately + assert torch.autograd.gradcheck( + lambda h: _AbsBroadcastAdd.apply(h, addend.detach()), + (hidden,), + ) From 85f120bfac7949e1f11c832b9c5d930e1f0e1fd7 Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Mon, 22 Jun 2026 18:15:17 +0800 Subject: [PATCH 23/41] Add event-log checksum-consistency analysis rules MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add the first event-analyzer rules that replay the structured event log and flag weight-checksum inconsistencies: a `checksum_compare` helper (flatten nested dicts, diff flat checksum maps) plus two rules built on it — cross-replica weight checksum consistency and inference-engine weight checksum consistency — with unit tests. - miles/utils/event_analyzer/rules/{checksum_compare,cross_replica_weight_checksum,inference_engine_weight_checksum_consistency}.py and tests. --- miles/utils/event_analyzer/__init__.py | 0 miles/utils/event_analyzer/rules/__init__.py | 1 + .../event_analyzer/rules/checksum_compare.py | 51 ++++ .../rules/cross_replica_weight_checksum.py | 67 ++++++ ...ence_engine_weight_checksum_consistency.py | 29 +++ tests/fast/utils/event_analyzer/__init__.py | 0 .../utils/event_analyzer/rules/__init__.py | 1 + .../rules/test_checksum_compare.py | 32 +++ .../test_cross_replica_weight_checksum.py | 223 ++++++++++++++++++ ...ence_engine_weight_checksum_consistency.py | 83 +++++++ 10 files changed, 487 insertions(+) create mode 100644 miles/utils/event_analyzer/__init__.py create mode 100644 miles/utils/event_analyzer/rules/__init__.py create mode 100644 miles/utils/event_analyzer/rules/checksum_compare.py create mode 100644 miles/utils/event_analyzer/rules/cross_replica_weight_checksum.py create mode 100644 miles/utils/event_analyzer/rules/inference_engine_weight_checksum_consistency.py create mode 100644 tests/fast/utils/event_analyzer/__init__.py create mode 100644 tests/fast/utils/event_analyzer/rules/__init__.py create mode 100644 tests/fast/utils/event_analyzer/rules/test_checksum_compare.py create mode 100644 tests/fast/utils/event_analyzer/rules/test_cross_replica_weight_checksum.py create mode 100644 tests/fast/utils/event_analyzer/rules/test_inference_engine_weight_checksum_consistency.py diff --git a/miles/utils/event_analyzer/__init__.py b/miles/utils/event_analyzer/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/miles/utils/event_analyzer/rules/__init__.py b/miles/utils/event_analyzer/rules/__init__.py new file mode 100644 index 00000000000..8b137891791 --- /dev/null +++ b/miles/utils/event_analyzer/rules/__init__.py @@ -0,0 +1 @@ + diff --git a/miles/utils/event_analyzer/rules/checksum_compare.py b/miles/utils/event_analyzer/rules/checksum_compare.py new file mode 100644 index 00000000000..bf4cba9f133 --- /dev/null +++ b/miles/utils/event_analyzer/rules/checksum_compare.py @@ -0,0 +1,51 @@ +from collections.abc import Iterable +from typing import Any + +from miles.utils.pydantic_utils import FrozenStrictBaseModel + + +class ChecksumMismatchIssue(FrozenStrictBaseModel): + key: str + label_a: str + label_b: str + value_a: str + value_b: str + + +def compare_flat_dicts( + a: dict[str, Any], + b: dict[str, Any], + label_a: str, + label_b: str, +) -> Iterable[ChecksumMismatchIssue]: + """Compare two flat dicts and yield one mismatch per differing key.""" + all_keys = sorted(set(a.keys()) | set(b.keys())) + + for key in all_keys: + value_a = a.get(key, "") + value_b = b.get(key, "") + if value_a != value_b: + yield ChecksumMismatchIssue( + key=key, + label_a=label_a, + label_b=label_b, + value_a=str(value_a), + value_b=str(value_b), + ) + + +def flatten_nested(obj: Any, *, prefix: str = "") -> dict[str, Any]: + """Flatten a nested dict/list into a flat dict with dot-separated keys. Keeps all primitive leaf values.""" + result: dict[str, Any] = {} + + if isinstance(obj, dict): + for k, v in sorted(obj.items(), key=lambda x: str(x[0])): + child_prefix = f"{prefix}.{k}" if prefix else str(k) + result.update(flatten_nested(v, prefix=child_prefix)) + elif isinstance(obj, (list, tuple)): + for i, v in enumerate(obj): + result.update(flatten_nested(v, prefix=f"{prefix}[{i}]")) + else: + result[prefix] = obj + + return result diff --git a/miles/utils/event_analyzer/rules/cross_replica_weight_checksum.py b/miles/utils/event_analyzer/rules/cross_replica_weight_checksum.py new file mode 100644 index 00000000000..8c6bddcc62e --- /dev/null +++ b/miles/utils/event_analyzer/rules/cross_replica_weight_checksum.py @@ -0,0 +1,67 @@ +from collections import defaultdict +from collections.abc import Iterable +from typing import Any + +from miles.utils.event_analyzer.rules.checksum_compare import ChecksumMismatchIssue, compare_flat_dicts, flatten_nested +from miles.utils.event_logger.models import Event, TrainEngineLocalWeightChecksumEvent +from miles.utils.process_identity import TrainProcessIdentity + +__all__ = ["ChecksumMismatchIssue", "check"] + + +def check(events: list[Event]) -> list[ChecksumMismatchIssue]: + """ + Check: weight checksum across replicas should be exactly the same + """ + + checksum_events = [e for e in events if isinstance(e, TrainEngineLocalWeightChecksumEvent)] + if not checksum_events: + return [] + + all_mismatches: list[ChecksumMismatchIssue] = [] + + events_by_key: dict[tuple[int, int], list[TrainEngineLocalWeightChecksumEvent]] = {} + for event in checksum_events: + key = (event.rollout_id, event.attempt) + events_by_key.setdefault(key, []).append(event) + + for key in sorted(events_by_key.keys()): + all_mismatches += list(_check_one_step(events=events_by_key[key])) + + return all_mismatches + + +def _get_rank_key(event: TrainEngineLocalWeightChecksumEvent) -> int: + if isinstance(event.source, TrainProcessIdentity): + return event.source.rank_within_cell + return -1 + + +def _check_one_step(events: list[TrainEngineLocalWeightChecksumEvent]) -> Iterable[ChecksumMismatchIssue]: + # Group events by rank_within_cell so we only compare across replicas (cell_index), + # not across TP/PP/EP ranks within the same cell (which have different param shards). + # TODO: group by (component, rank_within_cell) once critic checksum events are supported. + # Currently only actor emits TrainEngineLocalWeightChecksumEvent. + by_rank: dict[int, list[TrainEngineLocalWeightChecksumEvent]] = defaultdict(list) + for event in events: + by_rank[_get_rank_key(event)].append(event) + + for rank_events in by_rank.values(): + first = rank_events[0] + first_flat = _flatten_event(first) + for other in rank_events[1:]: + yield from compare_flat_dicts( + a=first_flat, + b=_flatten_event(other), + label_a=_compute_label(first), + label_b=_compute_label(other), + ) + + +def _compute_label(event: TrainEngineLocalWeightChecksumEvent) -> str: + return f"rollout_{event.rollout_id}/{event.source.to_name()}" + + +def _flatten_event(event: TrainEngineLocalWeightChecksumEvent) -> dict[str, Any]: + """Flatten all fields of an event into a flat dict with dot-separated keys.""" + return flatten_nested(event.state.model_dump(), prefix="") diff --git a/miles/utils/event_analyzer/rules/inference_engine_weight_checksum_consistency.py b/miles/utils/event_analyzer/rules/inference_engine_weight_checksum_consistency.py new file mode 100644 index 00000000000..be49daee6fa --- /dev/null +++ b/miles/utils/event_analyzer/rules/inference_engine_weight_checksum_consistency.py @@ -0,0 +1,29 @@ +from collections.abc import Iterable + +from miles.utils.event_analyzer.rules.checksum_compare import ChecksumMismatchIssue, compare_flat_dicts +from miles.utils.event_logger.models import Event, InferenceEngineWeightChecksumEvent + +__all__ = ["check"] + + +def check(events: list[Event]) -> list[ChecksumMismatchIssue]: + """Check: all engines of one rollout must hold exactly the same weights.""" + issues: list[ChecksumMismatchIssue] = [] + for event in events: + if isinstance(event, InferenceEngineWeightChecksumEvent): + issues += list(_check_one_rollout(event)) + return issues + + +def _check_one_rollout(event: InferenceEngineWeightChecksumEvent) -> Iterable[ChecksumMismatchIssue]: + engines = event.engine_checksums + if len(engines) < 2: + return + baseline = engines[0] + for engine_index in range(1, len(engines)): + yield from compare_flat_dicts( + a=baseline, + b=engines[engine_index], + label_a=f"rollout_{event.rollout_id}/engine_0", + label_b=f"rollout_{event.rollout_id}/engine_{engine_index}", + ) diff --git a/tests/fast/utils/event_analyzer/__init__.py b/tests/fast/utils/event_analyzer/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/fast/utils/event_analyzer/rules/__init__.py b/tests/fast/utils/event_analyzer/rules/__init__.py new file mode 100644 index 00000000000..8b137891791 --- /dev/null +++ b/tests/fast/utils/event_analyzer/rules/__init__.py @@ -0,0 +1 @@ + diff --git a/tests/fast/utils/event_analyzer/rules/test_checksum_compare.py b/tests/fast/utils/event_analyzer/rules/test_checksum_compare.py new file mode 100644 index 00000000000..48f97f79dd4 --- /dev/null +++ b/tests/fast/utils/event_analyzer/rules/test_checksum_compare.py @@ -0,0 +1,32 @@ +"""Tests for event_analyzer rules/checksum_compare shared primitives.""" + +from miles.utils.event_analyzer.rules.checksum_compare import compare_flat_dicts, flatten_nested + + +class TestCompareFlatDicts: + def test_identical_dicts_no_issue(self) -> None: + """Identical flat dicts yield no mismatch.""" + assert list(compare_flat_dicts({"w": "a"}, {"w": "a"}, "x", "y")) == [] + + def test_differing_value_reported(self) -> None: + """A differing value is reported with both labels and values.""" + issues = list(compare_flat_dicts({"w": "a"}, {"w": "b"}, "x", "y")) + assert len(issues) == 1 + assert issues[0].key == "w" + assert issues[0].value_a == "a" + assert issues[0].value_b == "b" + + def test_missing_key_marked(self) -> None: + """A key present only on one side is reported as on the other.""" + issues = list(compare_flat_dicts({"w": "a"}, {}, "x", "y")) + assert issues[0].value_b == "" + + +class TestFlattenNested: + def test_nested_dict_flattened(self) -> None: + """Nested dicts flatten to dot-separated keys.""" + assert flatten_nested({"state": {0: {"x": "h"}}}, prefix="opt") == {"opt.state.0.x": "h"} + + def test_list_indexed(self) -> None: + """List entries flatten to bracket-indexed keys.""" + assert flatten_nested({"p": ["a", "b"]}, prefix="r") == {"r.p[0]": "a", "r.p[1]": "b"} diff --git a/tests/fast/utils/event_analyzer/rules/test_cross_replica_weight_checksum.py b/tests/fast/utils/event_analyzer/rules/test_cross_replica_weight_checksum.py new file mode 100644 index 00000000000..456d25faa80 --- /dev/null +++ b/tests/fast/utils/event_analyzer/rules/test_cross_replica_weight_checksum.py @@ -0,0 +1,223 @@ +"""Tests for event_analyzer rules/weight_checksum.""" + +from datetime import datetime, timezone + +from miles.utils.event_analyzer.rules.checksum_compare import flatten_nested as _flatten_nested +from miles.utils.event_analyzer.rules.cross_replica_weight_checksum import _flatten_event, check +from miles.utils.event_logger.models import ( + OptimizerStateInfo, + TrainEngineLocalWeightChecksumEvent, + TrainEngineLocalWeightChecksumState, +) +from miles.utils.process_identity import TrainProcessIdentity + +_FIXED_TS = datetime(2026, 1, 1, tzinfo=timezone.utc) + + +def _make_event( + rollout_id: int, + cell_index: int = 0, + rank_within_cell: int = 0, + param_hashes: dict[str, str] | None = None, + buffer_hashes: dict[str, str] | None = None, + optimizer_state_dict: dict | None = None, +) -> TrainEngineLocalWeightChecksumEvent: + return TrainEngineLocalWeightChecksumEvent( + timestamp=_FIXED_TS, + source=TrainProcessIdentity(component="actor", cell_index=cell_index, rank_within_cell=rank_within_cell), + rollout_id=rollout_id, + state=TrainEngineLocalWeightChecksumState( + param_hashes=param_hashes or {}, + buffer_hashes=buffer_hashes or {}, + optimizer_hashes=( + [ + OptimizerStateInfo( + param_names={0: "pp0.weight"}, + state_dict=optimizer_state_dict or {}, + ), + ] + if optimizer_state_dict is not None + else [] + ), + ), + ) + + +class TestCheck: + def test_matching_replicas_no_mismatches(self) -> None: + events = [ + _make_event(rollout_id=0, cell_index=0, param_hashes={"pp0.weight": "aaa"}), + _make_event(rollout_id=0, cell_index=1, param_hashes={"pp0.weight": "aaa"}), + _make_event(rollout_id=0, cell_index=2, param_hashes={"pp0.weight": "aaa"}), + ] + assert check(events) == [] + + def test_param_hash_mismatch_detected(self) -> None: + events = [ + _make_event(rollout_id=5, cell_index=0, param_hashes={"pp0.weight": "aaa"}), + _make_event(rollout_id=5, cell_index=1, param_hashes={"pp0.weight": "zzz"}), + ] + mismatches = check(events) + + assert len(mismatches) >= 1 + keys = [m.key for m in mismatches] + assert any("param_hashes.pp0.weight" in k for k in keys) + + def test_missing_key_in_one_replica_detected(self) -> None: + events = [ + _make_event(rollout_id=0, cell_index=0, param_hashes={"pp0.weight": "aaa", "pp0.bias": "bbb"}), + _make_event(rollout_id=0, cell_index=1, param_hashes={"pp0.weight": "aaa"}), + ] + mismatches = check(events) + + assert len(mismatches) >= 1 + keys = [m.key for m in mismatches] + assert any("pp0.bias" in k for k in keys) + assert any("" in m.value_b for m in mismatches) + + def test_multiple_steps_only_mismatched_step_reported(self) -> None: + events = [ + # Step 0: match + _make_event(rollout_id=0, cell_index=0, param_hashes={"pp0.w": "aaa"}), + _make_event(rollout_id=0, cell_index=1, param_hashes={"pp0.w": "aaa"}), + # Step 1: mismatch + _make_event(rollout_id=1, cell_index=0, param_hashes={"pp0.w": "aaa"}), + _make_event(rollout_id=1, cell_index=1, param_hashes={"pp0.w": "zzz"}), + # Step 2: match + _make_event(rollout_id=2, cell_index=0, param_hashes={"pp0.w": "aaa"}), + _make_event(rollout_id=2, cell_index=1, param_hashes={"pp0.w": "aaa"}), + ] + mismatches = check(events) + + assert len(mismatches) >= 1 + for m in mismatches: + assert "rollout_1/" in m.label_a or "rollout_1/" in m.label_b + + def test_empty_events_no_mismatches(self) -> None: + assert check([]) == [] + + def test_single_replica_no_comparison(self) -> None: + events = [_make_event(rollout_id=0, cell_index=0, param_hashes={"pp0.w": "aaa"})] + assert check(events) == [] + + def test_buffer_mismatch_detected(self) -> None: + events = [ + _make_event(rollout_id=0, cell_index=0, buffer_hashes={"pp0.running_mean": "aaa"}), + _make_event(rollout_id=0, cell_index=1, buffer_hashes={"pp0.running_mean": "bbb"}), + ] + mismatches = check(events) + + assert len(mismatches) >= 1 + assert any("buffer_hashes" in m.key for m in mismatches) + + def test_optimizer_state_mismatch_detected(self) -> None: + events = [ + _make_event(rollout_id=3, cell_index=0, optimizer_state_dict={"state": {0: {"exp_avg": "aaa"}}}), + _make_event(rollout_id=3, cell_index=1, optimizer_state_dict={"state": {0: {"exp_avg": "bbb"}}}), + ] + mismatches = check(events) + + assert len(mismatches) >= 1 + assert any("exp_avg" in m.key for m in mismatches) + + def test_non_tensor_state_mismatch_detected(self) -> None: + events = [ + _make_event(rollout_id=0, cell_index=0, optimizer_state_dict={"state": {0: {"step": 10}}}), + _make_event(rollout_id=0, cell_index=1, optimizer_state_dict={"state": {0: {"step": 20}}}), + ] + mismatches = check(events) + + assert len(mismatches) >= 1 + assert any("step" in m.key for m in mismatches) + + def test_different_ranks_not_compared(self) -> None: + """Events from different rank_within_cell should not be compared (PP sharding).""" + events = [ + _make_event(rollout_id=0, cell_index=0, rank_within_cell=0, param_hashes={"pp0.w": "aaa"}), + _make_event(rollout_id=0, cell_index=0, rank_within_cell=1, param_hashes={"pp1.w": "bbb"}), + ] + assert check(events) == [] + + def test_same_rank_different_cells_mismatch(self) -> None: + """Same rank_within_cell across different cells must be compared.""" + events = [ + _make_event(rollout_id=0, cell_index=0, rank_within_cell=0, param_hashes={"pp0.w": "aaa"}), + _make_event(rollout_id=0, cell_index=1, rank_within_cell=0, param_hashes={"pp0.w": "zzz"}), + ] + mismatches = check(events) + assert len(mismatches) >= 1 + + def test_same_rank_different_cells_match(self) -> None: + """Same rank_within_cell across different cells, identical hashes → no mismatch.""" + events = [ + _make_event(rollout_id=0, cell_index=0, rank_within_cell=0, param_hashes={"pp0.w": "aaa"}), + _make_event(rollout_id=0, cell_index=1, rank_within_cell=0, param_hashes={"pp0.w": "aaa"}), + _make_event(rollout_id=0, cell_index=0, rank_within_cell=1, param_hashes={"pp1.w": "bbb"}), + _make_event(rollout_id=0, cell_index=1, rank_within_cell=1, param_hashes={"pp1.w": "bbb"}), + ] + assert check(events) == [] + + def test_multi_rank_multi_cell_only_mismatched_rank_reported(self) -> None: + """Only the rank with mismatch should produce issues.""" + events = [ + # rank 0: match + _make_event(rollout_id=0, cell_index=0, rank_within_cell=0, param_hashes={"pp0.w": "aaa"}), + _make_event(rollout_id=0, cell_index=1, rank_within_cell=0, param_hashes={"pp0.w": "aaa"}), + # rank 1: mismatch + _make_event(rollout_id=0, cell_index=0, rank_within_cell=1, param_hashes={"pp1.w": "bbb"}), + _make_event(rollout_id=0, cell_index=1, rank_within_cell=1, param_hashes={"pp1.w": "zzz"}), + ] + mismatches = check(events) + assert len(mismatches) >= 1 + assert all("pp1.w" in m.key for m in mismatches) + + def test_three_cells_first_vs_rest(self) -> None: + """With 3 cells, all are compared against the first.""" + events = [ + _make_event(rollout_id=0, cell_index=0, param_hashes={"w": "aaa"}), + _make_event(rollout_id=0, cell_index=1, param_hashes={"w": "aaa"}), + _make_event(rollout_id=0, cell_index=2, param_hashes={"w": "bbb"}), + ] + mismatches = check(events) + assert len(mismatches) >= 1 + + +class TestFlattenEvent: + def test_excludes_metadata_fields(self) -> None: + event = _make_event(rollout_id=0, cell_index=0, param_hashes={"pp0.w": "aaa"}) + flat = _flatten_event(event) + + assert not any(k.startswith("step") or k.startswith("rank") or k.startswith("type") for k in flat.keys()) + assert "param_hashes.pp0.w" in flat + + def test_includes_optimizer_hashes(self) -> None: + event = _make_event(rollout_id=0, cell_index=0, optimizer_state_dict={"state": {0: {"exp_avg": "hash1"}}}) + flat = _flatten_event(event) + + assert any("exp_avg" in k for k in flat.keys()) + + +class TestFlattenNested: + def test_flat_dict_with_string_values(self) -> None: + result = _flatten_nested({"a": "hash1", "b": "hash2"}, prefix="root") + assert result == {"root.a": "hash1", "root.b": "hash2"} + + def test_nested_dict(self) -> None: + result = _flatten_nested({"state": {0: {"exp_avg": "h1"}}}, prefix="opt0") + assert result == {"opt0.state.0.exp_avg": "h1"} + + def test_list_values(self) -> None: + result = _flatten_nested({"params": ["a", "b"]}, prefix="opt0") + assert result == {"opt0.params[0]": "a", "opt0.params[1]": "b"} + + def test_keeps_int_and_float_leaves(self) -> None: + result = _flatten_nested({"lr": 0.001, "step": 42, "hash": "abc"}, prefix="root") + assert result == {"root.hash": "abc", "root.lr": 0.001, "root.step": 42} + + def test_empty_prefix(self) -> None: + result = _flatten_nested({"a": "x"}, prefix="") + assert result == {"a": "x"} + + def test_empty_dict(self) -> None: + result = _flatten_nested({}, prefix="root") + assert result == {} diff --git a/tests/fast/utils/event_analyzer/rules/test_inference_engine_weight_checksum_consistency.py b/tests/fast/utils/event_analyzer/rules/test_inference_engine_weight_checksum_consistency.py new file mode 100644 index 00000000000..92e7548f26a --- /dev/null +++ b/tests/fast/utils/event_analyzer/rules/test_inference_engine_weight_checksum_consistency.py @@ -0,0 +1,83 @@ +"""Tests for event_analyzer rules/inference_engine_weight_checksum_consistency.""" + +from datetime import datetime, timezone + +from miles.utils.event_analyzer.rules.inference_engine_weight_checksum_consistency import check +from miles.utils.event_logger.models import InferenceEngineWeightChecksumEvent +from miles.utils.process_identity import MainProcessIdentity + +_FIXED_TS = datetime(2026, 1, 1, tzinfo=timezone.utc) + + +def _make_event( + *, rollout_id: int | None, engine_checksums: list[dict[str, str]] +) -> InferenceEngineWeightChecksumEvent: + return InferenceEngineWeightChecksumEvent( + timestamp=_FIXED_TS, + source=MainProcessIdentity(), + rollout_id=rollout_id, + engine_checksums=engine_checksums, + ) + + +class TestCheck: + def test_empty_events_no_mismatches(self) -> None: + """No engine checksum events means nothing to check.""" + assert check([]) == [] + + def test_single_engine_no_comparison(self) -> None: + """A single engine has no peer to compare against.""" + events = [_make_event(rollout_id=0, engine_checksums=[{"rank0/w": "aaa"}])] + assert check(events) == [] + + def test_matching_engines_no_mismatches(self) -> None: + """All engines holding identical checksums produce no issues.""" + events = [ + _make_event(rollout_id=0, engine_checksums=[{"rank0/w": "aaa"}, {"rank0/w": "aaa"}, {"rank0/w": "aaa"}]) + ] + assert check(events) == [] + + def test_tensor_mismatch_reports_engine_and_tensor(self) -> None: + """A single differing tensor on one engine is reported with engine and tensor labels.""" + events = [_make_event(rollout_id=5, engine_checksums=[{"rank0/w": "aaa"}, {"rank0/w": "zzz"}])] + mismatches = check(events) + assert len(mismatches) == 1 + assert mismatches[0].key == "rank0/w" + assert mismatches[0].label_a == "rollout_5/engine_0" + assert mismatches[0].label_b == "rollout_5/engine_1" + + def test_missing_tensor_on_one_engine_detected(self) -> None: + """A tensor present on engine 0 but absent on engine 1 is a mismatch.""" + events = [ + _make_event(rollout_id=0, engine_checksums=[{"rank0/w": "aaa", "rank0/b": "bbb"}, {"rank0/w": "aaa"}]) + ] + mismatches = check(events) + assert any(m.key == "rank0/b" and "" in m.value_b for m in mismatches) + + def test_all_engines_compared_against_first(self) -> None: + """Every engine is compared against engine 0; a later engine's diff names engine_2.""" + events = [ + _make_event(rollout_id=0, engine_checksums=[{"rank0/w": "aaa"}, {"rank0/w": "aaa"}, {"rank0/w": "bbb"}]) + ] + mismatches = check(events) + assert len(mismatches) == 1 + assert mismatches[0].label_a == "rollout_0/engine_0" + assert mismatches[0].label_b == "rollout_0/engine_2" + + def test_none_rollout_id_mismatch_labelled_rollout_none(self) -> None: + """The initial out-of-loop sync (rollout_id=None) still checks engines and labels them rollout_None.""" + events = [_make_event(rollout_id=None, engine_checksums=[{"rank0/w": "aaa"}, {"rank0/w": "zzz"}])] + mismatches = check(events) + assert len(mismatches) == 1 + assert mismatches[0].label_a == "rollout_None/engine_0" + assert mismatches[0].label_b == "rollout_None/engine_1" + + def test_only_mismatched_rollout_reported(self) -> None: + """Each rollout is its own event; only the inconsistent rollout yields issues.""" + events = [ + _make_event(rollout_id=0, engine_checksums=[{"rank0/w": "aaa"}, {"rank0/w": "aaa"}]), + _make_event(rollout_id=1, engine_checksums=[{"rank0/w": "aaa"}, {"rank0/w": "zzz"}]), + ] + mismatches = check(events) + assert len(mismatches) == 1 + assert "rollout_1/" in mismatches[0].label_a From 5b5c7977ecdfd13e574ab7c6f96fa4d1093259c3 Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Mon, 22 Jun 2026 18:15:17 +0800 Subject: [PATCH 24/41] Add an event-log witness-tracing analysis rule Add the witness-tracing rule for the event analyzer: it follows witness ids through the replayed event log to verify they are propagated correctly across the training pipeline, with unit tests. - miles/utils/event_analyzer/rules/witness.py and tests. --- miles/utils/event_analyzer/rules/witness.py | 250 +++++++++ .../event_analyzer/rules/test_witness.py | 524 ++++++++++++++++++ 2 files changed, 774 insertions(+) create mode 100644 miles/utils/event_analyzer/rules/witness.py create mode 100644 tests/fast/utils/event_analyzer/rules/test_witness.py diff --git a/miles/utils/event_analyzer/rules/witness.py b/miles/utils/event_analyzer/rules/witness.py new file mode 100644 index 00000000000..f85569e3aa4 --- /dev/null +++ b/miles/utils/event_analyzer/rules/witness.py @@ -0,0 +1,250 @@ +import logging +from collections import defaultdict +from collections.abc import Callable, Hashable, Iterator, Sequence +from typing import Protocol, TypeVar + +from miles.backends.megatron_utils.types import TrainStepOutcome +from miles.utils.event_logger.models import ( + Event, + TrainAdvantageComputationEvent, + TrainGroupStepEndEvent, + WitnessAllocateIdEvent, + WitnessSnapshotParamEvent, +) +from miles.utils.pydantic_utils import FrozenStrictBaseModel + +logger = logging.getLogger(__name__) + + +class WitnessDataMismatchIssue(FrozenStrictBaseModel): + rollout_id: int + cell_index: int + description: str + expected_witness_ids: list[int] + actual_witness_ids: list[int] + + +class WitnessMissingSnapshotIssue(FrozenStrictBaseModel): + rollout_id: int + cell_index: int + description: str + + +WitnessIssue = WitnessDataMismatchIssue | WitnessMissingSnapshotIssue + + +def check(events: list[Event]) -> list[WitnessIssue]: + """ + Related events: + * WitnessAllocateIdEvent: when allocating `witness_id` to `sample_index` + * WitnessSnapshotParamEvent: near the end of each train() step in MegatronTrainRayActor + * If a witness_id appears in the weight, it means the corresponding data is consumed at least once. + * TrainGroupStepEndEvent: after each train() step in RayTrainGroup + + Check: + 1. For each (rollout_id, cell_index), + if TrainGroupStepEndEvent claims the cell ends with TrainStepOutcome.NORMAL, + then its WitnessSnapshotParamEvent should observe *EXACTLY* the training data in rollout_id=0~curr. + + Remarks: + * To correlate witness_id vs sample_index utilize WitnessAllocateIdEvent. + * Witness' ring buffer will remove old data, thus we need to ignore the appearance/disappearance of + all values in `WitnessSnapshotParamEvent.stale_ids` + """ + + allocated_witness_ids_by_rollout = _compute_allocated_witness_ids_by_rollout( + _filter_by_type(events, WitnessAllocateIdEvent) + ) + + return list( + _find_mismatches( + all_step_events=_filter_by_type(events, TrainGroupStepEndEvent), + all_witness_events=_filter_by_type(events, WitnessSnapshotParamEvent), + expected_witness_ids_of_step=_compute_expected_witness_ids_of_step(allocated_witness_ids_by_rollout), + allocated_witness_ids_by_rollout=allocated_witness_ids_by_rollout, + zero_adv_witness_ids_by_rollout=_compute_zero_advantage_witness_ids( + _filter_by_type(events, TrainAdvantageComputationEvent) + ), + ) + ) + + +_EventT = TypeVar("_EventT") + + +def _filter_by_type(arr: Sequence[Event], ty: type[_EventT]) -> list[_EventT]: + return [x for x in arr if isinstance(x, ty)] + + +def _compute_zero_advantage_witness_ids( + events: list[TrainAdvantageComputationEvent], +) -> dict[int, set[int]]: + """Return witness_ids where all per-token advantages == 0.0, keyed by rollout_id. + + Unioned across cells: under indep_dp the per-cell weight snapshot reflects the + GLOBAL (allreduced) gradient, so a zero-advantage sample contributes nothing + and its witness is absent from EVERY cell — even cells that never owned it. + Keying per (rollout_id, cell_index) would let a cell excuse only its own shard, + falsely flagging peers' zero-advantage witnesses as missing. + + Snapshot checks excuse these ids cumulatively (see _zero_adv_excused_ids_at): a + zero-advantage sample stays absent from every later snapshot while the expected + set keeps it forever — until the ring buffer reallocates the witness id to a new + sample, which cancels the excusal. + + Only the highest-attempt events per rollout count, mirroring the allocate-event + handling: a crashed attempt's partial advantage events would otherwise excuse ids + that the successful retry trains for real. + + An EMPTY per-sample advantage list is skipped, not treated as zero-advantage. + Under CP>1 the advantage tensor logged by a rank holds only that rank's local + shard of the response tokens, while witness_ids spans the full sample; a short + sample whose response lands entirely on the other CP rank leaves this rank's + advantage shard empty. `all(v == 0.0 for v in [])` is vacuously True, so without + this guard such a sample would be falsely excused and then flagged 'extra' once + its (real, nonzero-advantage) witness shows up present on the rank that owns it. + """ + result: dict[int, set[int]] = defaultdict(set) + for event in _filter_to_latest_attempt(events, group_key=lambda e: e.rollout_id): + for adv_tokens, wid_tokens in zip(event.advantages, event.witness_ids, strict=True): + assert len(set(wid_tokens)) <= 1, f"witness ids within one sample must be uniform, got {set(wid_tokens)}" + if adv_tokens and all(v == 0.0 for v in adv_tokens): + result[event.rollout_id].add(wid_tokens[0]) + + return dict(result) + + +def _compute_allocated_witness_ids_by_rollout(events: list[WitnessAllocateIdEvent]) -> dict[int, set[int]]: + result: dict[int, set[int]] = defaultdict(set) + for e in _filter_to_latest_attempt(events, group_key=lambda e: e.rollout_id): + result[e.rollout_id] |= set(e.witness_id_to_sample_index.keys()) + return dict(result) + + +def _compute_expected_witness_ids_of_step( + allocated_witness_ids_by_rollout: dict[int, set[int]], +) -> dict[int, set[int]]: + ans: dict[int, set[int]] = {} + running: set[int] = set() + for rollout_id in sorted(allocated_witness_ids_by_rollout.keys()): + running |= allocated_witness_ids_by_rollout[rollout_id] + ans[rollout_id] = set(running) + return ans + + +def _find_mismatches( + *, + all_step_events: list[TrainGroupStepEndEvent], + all_witness_events: list[WitnessSnapshotParamEvent], + expected_witness_ids_of_step: dict[int, set[int]], + allocated_witness_ids_by_rollout: dict[int, set[int]], + zero_adv_witness_ids_by_rollout: dict[int, set[int]], +) -> Iterator[WitnessIssue]: + latest_attempt_witness_events = _filter_to_latest_attempt( + all_witness_events, group_key=lambda e: (e.rollout_id, e.source.cell_index) + ) + + for step_event in all_step_events: + rollout_id = step_event.rollout_id + + for cell_index, cell_outcome in step_event.cell_outcomes.items(): + if cell_outcome == "error": + continue + if not all(r == TrainStepOutcome.NORMAL for r in cell_outcome): + continue + + witness_events_of_cell = [ + e + for e in latest_attempt_witness_events + if e.rollout_id == rollout_id and e.source.cell_index == cell_index + ] + + if not witness_events_of_cell: + yield WitnessMissingSnapshotIssue( + rollout_id=rollout_id, + cell_index=cell_index, + description=f"Cell {cell_index} reported NORMAL for rollout {rollout_id} but no WitnessSnapshotParamEvent was found", + ) + continue + + zero_adv_excused_ids = _zero_adv_excused_ids_at( + zero_adv_witness_ids_by_rollout=zero_adv_witness_ids_by_rollout, + allocated_witness_ids_by_rollout=allocated_witness_ids_by_rollout, + rollout_id=rollout_id, + ) + + for event in witness_events_of_cell: + issue = _compare_snapshot( + event=event, + expected=expected_witness_ids_of_step.get(rollout_id, set()), + rollout_id=rollout_id, + cell_index=cell_index, + zero_adv_excused_ids=zero_adv_excused_ids, + ) + if issue is not None: + yield issue + + +def _zero_adv_excused_ids_at( + *, + zero_adv_witness_ids_by_rollout: dict[int, set[int]], + allocated_witness_ids_by_rollout: dict[int, set[int]], + rollout_id: int, +) -> set[int]: + excused: set[int] = set() + for rid in sorted(set(zero_adv_witness_ids_by_rollout) | set(allocated_witness_ids_by_rollout)): + if rid > rollout_id: + break + excused -= allocated_witness_ids_by_rollout.get(rid, set()) + excused |= zero_adv_witness_ids_by_rollout.get(rid, set()) + return excused + + +class _HasAttempt(Protocol): + attempt: int + + +_EventWithAttemptT = TypeVar("_EventWithAttemptT", bound=_HasAttempt) + + +def _filter_to_latest_attempt( + events: Sequence[_EventWithAttemptT], + *, + group_key: Callable[[_EventWithAttemptT], Hashable], +) -> list[_EventWithAttemptT]: + max_attempt_by_group: dict[Hashable, int] = {} + for event in events: + key = group_key(event) + prev = max_attempt_by_group.get(key) + if prev is None or event.attempt > prev: + max_attempt_by_group[key] = event.attempt + + return [e for e in events if e.attempt == max_attempt_by_group[group_key(e)]] + + +def _compare_snapshot( + *, + event: WitnessSnapshotParamEvent, + expected: set[int], + rollout_id: int, + cell_index: int, + zero_adv_excused_ids: set[int], +) -> WitnessDataMismatchIssue | None: + stale_set = set(event.stale_ids) + filtered_expected = expected - stale_set - zero_adv_excused_ids + filtered_actual = set(event.nonzero_witness_ids) - stale_set + + if filtered_expected == filtered_actual: + return None + + return WitnessDataMismatchIssue( + rollout_id=rollout_id, + cell_index=cell_index, + description=( + f"Witness data mismatch for instance {event.instance_id}: " + f"missing={sorted(filtered_expected - filtered_actual)}, " + f"extra={sorted(filtered_actual - filtered_expected)}" + ), + expected_witness_ids=sorted(filtered_expected), + actual_witness_ids=sorted(filtered_actual), + ) diff --git a/tests/fast/utils/event_analyzer/rules/test_witness.py b/tests/fast/utils/event_analyzer/rules/test_witness.py new file mode 100644 index 00000000000..7a121cc3fde --- /dev/null +++ b/tests/fast/utils/event_analyzer/rules/test_witness.py @@ -0,0 +1,524 @@ +"""Tests for event_analyzer rules/witness.""" + +from datetime import datetime, timezone + +from pydantic import TypeAdapter + +from miles.backends.megatron_utils.types import TrainStepOutcome +from miles.utils.event_analyzer.rules.witness import WitnessDataMismatchIssue, WitnessMissingSnapshotIssue, check +from miles.utils.event_logger.models import ( + Event, + TrainAdvantageComputationEvent, + TrainGroupStepEndEvent, + WitnessAllocateIdEvent, + WitnessSnapshotParamEvent, +) +from miles.utils.process_identity import MainProcessIdentity, TrainProcessIdentity + +_event_adapter = TypeAdapter(Event) + +_FIXED_TS = datetime(2026, 1, 1, tzinfo=timezone.utc) +_MAIN_SOURCE = MainProcessIdentity() + + +def _make_source(cell_index: int = 0, rank_within_cell: int = 0) -> TrainProcessIdentity: + return TrainProcessIdentity(component="actor", cell_index=cell_index, rank_within_cell=rank_within_cell) + + +def _make_snapshot( + rollout_id: int, + nonzero_witness_ids: list[int], + instance_id: str = "pp0.head", + cell_index: int = 0, + rank_within_cell: int = 0, + stale_ids: list[int] | None = None, + attempt: int = 0, +) -> WitnessSnapshotParamEvent: + return WitnessSnapshotParamEvent( + timestamp=_FIXED_TS, + source=_make_source(cell_index=cell_index, rank_within_cell=rank_within_cell), + rollout_id=rollout_id, + attempt=attempt, + instance_id=instance_id, + nonzero_witness_ids=nonzero_witness_ids, + stale_ids=stale_ids or [], + ) + + +def _make_allocate( + rollout_id: int, + witness_id_to_sample_index: dict[int, int], + attempt: int = 0, +) -> WitnessAllocateIdEvent: + return WitnessAllocateIdEvent( + timestamp=_FIXED_TS, + source=_MAIN_SOURCE, + rollout_id=rollout_id, + attempt=attempt, + witness_id_to_sample_index=witness_id_to_sample_index, + counter_after=max(witness_id_to_sample_index.keys(), default=-1) + 1, + ) + + +def _make_step_end( + rollout_id: int, + cell_outcomes: dict[int, str | list[TrainStepOutcome]], +) -> TrainGroupStepEndEvent: + return TrainGroupStepEndEvent( + timestamp=_FIXED_TS, + source=_MAIN_SOURCE, + rollout_id=rollout_id, + cell_outcomes=cell_outcomes, + ) + + +def _make_advantage( + rollout_id: int, + advantages: list[list[float]], + witness_ids: list[list[int]], + cell_index: int = 0, + attempt: int = 0, +) -> TrainAdvantageComputationEvent: + return TrainAdvantageComputationEvent( + timestamp=_FIXED_TS, + source=_make_source(cell_index=cell_index), + rollout_id=rollout_id, + attempt=attempt, + advantages=advantages, + witness_ids=witness_ids, + ) + + +class TestWitnessCheck: + def test_empty_events(self) -> None: + assert check([]) == [] + + def test_normal_step_with_correct_cumulative_witness_ids_returns_no_issues(self) -> None: + events: list[Event] = [ + _make_allocate(rollout_id=0, witness_id_to_sample_index={10: 0, 11: 1}), + _make_snapshot(rollout_id=0, nonzero_witness_ids=[10, 11]), + _make_step_end(rollout_id=0, cell_outcomes={0: [TrainStepOutcome.NORMAL]}), + ] + assert check(events) == [] + + def test_normal_step_with_missing_witness_id_returns_issue(self) -> None: + events: list[Event] = [ + _make_allocate(rollout_id=0, witness_id_to_sample_index={10: 0, 11: 1}), + _make_snapshot(rollout_id=0, nonzero_witness_ids=[10]), + _make_step_end(rollout_id=0, cell_outcomes={0: [TrainStepOutcome.NORMAL]}), + ] + issues = check(events) + assert len(issues) == 1 + assert issues[0].rollout_id == 0 + assert 11 in issues[0].expected_witness_ids + assert 11 not in issues[0].actual_witness_ids + + def test_normal_step_with_extra_witness_id_returns_issue(self) -> None: + events: list[Event] = [ + _make_allocate(rollout_id=0, witness_id_to_sample_index={10: 0, 11: 1}), + _make_snapshot(rollout_id=0, nonzero_witness_ids=[10, 11, 99]), + _make_step_end(rollout_id=0, cell_outcomes={0: [TrainStepOutcome.NORMAL]}), + ] + issues = check(events) + assert len(issues) == 1 + assert 99 in issues[0].actual_witness_ids + assert 99 not in issues[0].expected_witness_ids + + def test_discarded_step_is_ignored(self) -> None: + events: list[Event] = [ + _make_allocate(rollout_id=0, witness_id_to_sample_index={10: 0, 11: 1}), + _make_snapshot(rollout_id=0, nonzero_witness_ids=[10]), + _make_step_end(rollout_id=0, cell_outcomes={0: [TrainStepOutcome.DISCARDED_SHOULD_RETRY]}), + ] + assert check(events) == [] + + def test_stale_ids_are_ignored(self) -> None: + """IDs in stale_ids are ignored in both expected and actual.""" + events: list[Event] = [ + _make_allocate(rollout_id=0, witness_id_to_sample_index={2: 0, 5: 1, 8: 2}), + _make_snapshot(rollout_id=0, nonzero_witness_ids=[5, 8], stale_ids=[0, 1, 2]), + _make_step_end(rollout_id=0, cell_outcomes={0: [TrainStepOutcome.NORMAL]}), + ] + assert check(events) == [] + + def test_multiple_cells_independent_checking(self) -> None: + """Each cell is checked independently.""" + events: list[Event] = [ + _make_allocate(rollout_id=0, witness_id_to_sample_index={10: 0, 11: 1}), + _make_snapshot(rollout_id=0, nonzero_witness_ids=[10, 11], cell_index=0), + _make_snapshot(rollout_id=0, nonzero_witness_ids=[10, 11], cell_index=1), + _make_step_end(rollout_id=0, cell_outcomes={0: [TrainStepOutcome.NORMAL], 1: [TrainStepOutcome.NORMAL]}), + ] + assert check(events) == [] + + def test_retry_uses_latest_attempt_allocation(self) -> None: + """When retries happen, only the latest attempt's allocation is used.""" + events: list[Event] = [ + _make_allocate(rollout_id=0, witness_id_to_sample_index={10: 0, 11: 1}, attempt=0), + _make_allocate(rollout_id=0, witness_id_to_sample_index={20: 0, 21: 1}, attempt=1), + _make_snapshot(rollout_id=0, nonzero_witness_ids=[20, 21]), + _make_step_end(rollout_id=0, cell_outcomes={0: [TrainStepOutcome.NORMAL]}), + ] + assert check(events) == [] + + def test_snapshot_latest_attempt_discards_stale_crashed_snapshot(self) -> None: + """A crashed attempt-0 snapshot with wrong ids is discarded; only the latest attempt-1 snapshot is compared.""" + events: list[Event] = [ + _make_allocate(rollout_id=0, witness_id_to_sample_index={20: 0, 21: 1}, attempt=1), + _make_snapshot(rollout_id=0, nonzero_witness_ids=[10, 11], attempt=0), + _make_snapshot(rollout_id=0, nonzero_witness_ids=[20, 21], attempt=1), + _make_step_end(rollout_id=0, cell_outcomes={0: [TrainStepOutcome.NORMAL]}), + ] + assert check(events) == [] + + def test_snapshot_latest_attempt_is_the_one_compared(self) -> None: + """The latest attempt-1 snapshot has wrong ids (attempt-0 was correct), so exactly one mismatch is reported.""" + events: list[Event] = [ + _make_allocate(rollout_id=0, witness_id_to_sample_index={10: 0, 11: 1}, attempt=1), + _make_snapshot(rollout_id=0, nonzero_witness_ids=[10, 11], attempt=0), + _make_snapshot(rollout_id=0, nonzero_witness_ids=[99], attempt=1), + _make_step_end(rollout_id=0, cell_outcomes={0: [TrainStepOutcome.NORMAL]}), + ] + issues = check(events) + assert len(issues) == 1 + assert isinstance(issues[0], WitnessDataMismatchIssue) + assert 99 in issues[0].actual_witness_ids + assert 99 not in issues[0].expected_witness_ids + + def test_cumulative_across_rollouts(self) -> None: + """Expected witness IDs are cumulative from rollout 0 to current.""" + events: list[Event] = [ + _make_allocate(rollout_id=0, witness_id_to_sample_index={10: 0}), + _make_snapshot(rollout_id=0, nonzero_witness_ids=[10]), + _make_step_end(rollout_id=0, cell_outcomes={0: [TrainStepOutcome.NORMAL]}), + _make_allocate(rollout_id=1, witness_id_to_sample_index={11: 1}), + _make_snapshot(rollout_id=1, nonzero_witness_ids=[10, 11]), + _make_step_end(rollout_id=1, cell_outcomes={0: [TrainStepOutcome.NORMAL]}), + ] + assert check(events) == [] + + def test_missing_snapshot_for_normal_cell_returns_issue(self) -> None: + """Cell claims NORMAL but has no WitnessSnapshotParamEvent — should return WitnessMissingSnapshotIssue.""" + events: list[Event] = [ + _make_allocate(rollout_id=0, witness_id_to_sample_index={10: 0}), + _make_step_end(rollout_id=0, cell_outcomes={0: [TrainStepOutcome.NORMAL]}), + ] + issues = check(events) + assert len(issues) == 1 + assert isinstance(issues[0], WitnessMissingSnapshotIssue) + assert issues[0].rollout_id == 0 + assert issues[0].cell_index == 0 + + def test_error_cell_outcome_is_skipped(self) -> None: + """cell_outcomes with 'error' string should not produce any issue.""" + events: list[Event] = [ + _make_allocate(rollout_id=0, witness_id_to_sample_index={10: 0}), + _make_step_end(rollout_id=0, cell_outcomes={0: "error"}), + ] + assert check(events) == [] + + def test_multi_element_cell_outcome_not_all_normal_is_skipped(self) -> None: + """A cell whose outcome list mixes NORMAL with a non-NORMAL outcome is skipped, even with a missing witness id.""" + events: list[Event] = [ + _make_allocate(rollout_id=0, witness_id_to_sample_index={10: 0}), + _make_snapshot(rollout_id=0, nonzero_witness_ids=[]), + _make_step_end( + rollout_id=0, + cell_outcomes={0: [TrainStepOutcome.NORMAL, TrainStepOutcome.DISCARDED_SHOULD_RETRY]}, + ), + ] + assert check(events) == [] + + def test_multiple_snapshots_per_cell(self) -> None: + """Same cell has head and tail snapshots; only the mismatched one produces an issue.""" + events: list[Event] = [ + _make_allocate(rollout_id=0, witness_id_to_sample_index={10: 0, 11: 1}), + _make_snapshot(rollout_id=0, nonzero_witness_ids=[10, 11], instance_id="pp0.head"), + _make_snapshot(rollout_id=0, nonzero_witness_ids=[10], instance_id="pp0.tail"), + _make_step_end(rollout_id=0, cell_outcomes={0: [TrainStepOutcome.NORMAL]}), + ] + issues = check(events) + assert len(issues) == 1 + assert isinstance(issues[0], WitnessDataMismatchIssue) + assert 11 in issues[0].expected_witness_ids + assert 11 not in issues[0].actual_witness_ids + + def test_ring_buffer_wrap_with_stale_ids(self) -> None: + """After wrap, stale_ids contain wrapped IDs (e.g. [8,9,0]). These should be excluded from comparison.""" + # buffer_size=10, allocated IDs 0..7 in rollout 0, then 8,9,0,1,2 in rollout 1 (wrap) + events: list[Event] = [ + _make_allocate(rollout_id=0, witness_id_to_sample_index={i: i for i in range(8)}), + _make_snapshot(rollout_id=0, nonzero_witness_ids=list(range(8))), + _make_step_end(rollout_id=0, cell_outcomes={0: [TrainStepOutcome.NORMAL]}), + _make_allocate(rollout_id=1, witness_id_to_sample_index={8: 8, 9: 9, 0: 10, 1: 11, 2: 12}), + # After wrap: stale_ids=[3,4,5] (old IDs cleaned), actual nonzero = [0,1,2,6,7,8,9] + _make_snapshot( + rollout_id=1, + nonzero_witness_ids=[0, 1, 2, 6, 7, 8, 9], + stale_ids=[3, 4, 5], + ), + _make_step_end(rollout_id=1, cell_outcomes={0: [TrainStepOutcome.NORMAL]}), + ] + # expected cumulative = {0..9}, minus stale {3,4,5} = {0,1,2,6,7,8,9} — matches actual + assert check(events) == [] + + def test_ring_buffer_wrap_detects_mismatch(self) -> None: + """After wrap, a genuinely missing non-stale ID should still be caught.""" + events: list[Event] = [ + _make_allocate(rollout_id=0, witness_id_to_sample_index={i: i for i in range(8)}), + _make_snapshot(rollout_id=0, nonzero_witness_ids=list(range(8))), + _make_step_end(rollout_id=0, cell_outcomes={0: [TrainStepOutcome.NORMAL]}), + _make_allocate(rollout_id=1, witness_id_to_sample_index={8: 8, 9: 9, 0: 10, 1: 11, 2: 12}), + # ID 7 is NOT stale but missing from actual — should be caught + _make_snapshot( + rollout_id=1, + nonzero_witness_ids=[0, 1, 2, 6, 8, 9], # missing 7 + stale_ids=[3, 4, 5], + ), + _make_step_end(rollout_id=1, cell_outcomes={0: [TrainStepOutcome.NORMAL]}), + ] + issues = check(events) + assert len(issues) == 1 + assert 7 in issues[0].expected_witness_ids + assert 7 not in issues[0].actual_witness_ids + + +class TestWitnessEventSerialization: + def test_roundtrip(self) -> None: + event = _make_snapshot( + rollout_id=5, + nonzero_witness_ids=[10, 20], + instance_id="pp0.tail", + cell_index=1, + stale_ids=[0, 1, 2], + ) + parsed = _event_adapter.validate_json(event.model_dump_json()) + assert isinstance(parsed, WitnessSnapshotParamEvent) + assert parsed.rollout_id == 5 + assert parsed.instance_id == "pp0.tail" + assert parsed.nonzero_witness_ids == [10, 20] + assert parsed.stale_ids == [0, 1, 2] + + +class TestZeroAdvantageExclusion: + def test_zero_advantage_sample_excluded_from_expected(self) -> None: + """Witness ID with zero advantage should not cause a mismatch when missing from actual.""" + events: list[Event] = [ + _make_allocate(rollout_id=0, witness_id_to_sample_index={10: 0, 11: 1}), + _make_advantage(rollout_id=0, advantages=[[0.0, 0.0], [2.0, 3.0]], witness_ids=[[10], [11]]), + _make_snapshot(rollout_id=0, nonzero_witness_ids=[11]), # 10 missing but zero-adv + _make_step_end(rollout_id=0, cell_outcomes={0: [TrainStepOutcome.NORMAL]}), + ] + assert check(events) == [] + + def test_zero_advantage_sample_stays_excluded_at_later_rollouts(self) -> None: + """Regression: a zero-advantage id stays absent from every later snapshot and must stay excused.""" + events: list[Event] = [ + _make_allocate(rollout_id=0, witness_id_to_sample_index={10: 0, 11: 1}), + _make_advantage(rollout_id=0, advantages=[[0.5, 0.5], [0.0, 0.0]], witness_ids=[[10], [11]]), + _make_snapshot(rollout_id=0, nonzero_witness_ids=[10]), + _make_step_end(rollout_id=0, cell_outcomes={0: [TrainStepOutcome.NORMAL]}), + _make_allocate(rollout_id=1, witness_id_to_sample_index={12: 0}), + _make_snapshot(rollout_id=1, nonzero_witness_ids=[10, 12]), + _make_step_end(rollout_id=1, cell_outcomes={0: [TrainStepOutcome.NORMAL]}), + ] + assert check(events) == [] + + def test_zero_advantage_sample_with_nonzero_signal_is_flagged_extra(self) -> None: + """A zero-advantage id that nevertheless shows a nonzero witness signal is still reported.""" + events: list[Event] = [ + _make_allocate(rollout_id=0, witness_id_to_sample_index={10: 0, 11: 1}), + _make_advantage(rollout_id=0, advantages=[[0.5, 0.5], [0.0, 0.0]], witness_ids=[[10], [11]]), + _make_snapshot(rollout_id=0, nonzero_witness_ids=[10, 11]), + _make_step_end(rollout_id=0, cell_outcomes={0: [TrainStepOutcome.NORMAL]}), + ] + issues = check(events) + assert len(issues) == 1 + assert isinstance(issues[0], WitnessDataMismatchIssue) + assert 11 in issues[0].actual_witness_ids + assert 11 not in issues[0].expected_witness_ids + + def test_empty_advantage_shard_under_cp_is_not_treated_as_zero_advantage(self) -> None: + """CP regression: an empty per-rank advantage shard must not excuse a present witness. + + Under CP>1 a short sample's response can land entirely on the other CP rank, leaving + this rank's advantage shard empty while witness_ids still spans the full sample. The + sample's real advantage is nonzero on the rank that owns it, so its witness is + legitimately present and must not be flagged 'extra' (all(v == 0.0 for v in []) is + vacuously True). + """ + events: list[Event] = [ + _make_allocate(rollout_id=0, witness_id_to_sample_index={10: 0, 11: 1}), + # CP rank that owns sample 1's response: real nonzero advantage. + _make_advantage(rollout_id=0, advantages=[[0.5, 0.5], [-0.54, -0.54]], witness_ids=[[10, 10], [11, 11]]), + # CP rank without sample 1's response tokens: empty advantage shard, full witness_ids. + _make_advantage(rollout_id=0, advantages=[[0.5, 0.5], []], witness_ids=[[10, 10], [11, 11]]), + _make_snapshot(rollout_id=0, nonzero_witness_ids=[10, 11]), + _make_step_end(rollout_id=0, cell_outcomes={0: [TrainStepOutcome.NORMAL]}), + ] + assert check(events) == [] + + def test_zero_advantage_exclusion_uses_final_attempt_only(self) -> None: + """A crashed attempt's all-zero advantage events must not excuse ids the successful retry trains.""" + events: list[Event] = [ + _make_allocate(rollout_id=0, witness_id_to_sample_index={10: 0, 11: 1}, attempt=1), + _make_advantage(rollout_id=0, advantages=[[0.0], [0.0]], witness_ids=[[10], [11]], attempt=0), + _make_advantage(rollout_id=0, advantages=[[5.0], [0.0]], witness_ids=[[10], [11]], attempt=1), + _make_snapshot(rollout_id=0, nonzero_witness_ids=[10]), + _make_step_end(rollout_id=0, cell_outcomes={0: [TrainStepOutcome.NORMAL]}), + ] + assert check(events) == [] + + def test_nonzero_advantage_sample_still_required(self) -> None: + """Witness ID with nonzero advantage must still appear — missing produces an issue.""" + events: list[Event] = [ + _make_allocate(rollout_id=0, witness_id_to_sample_index={10: 0, 11: 1}), + _make_advantage(rollout_id=0, advantages=[[5.0], [3.0]], witness_ids=[[10], [11]]), + _make_snapshot(rollout_id=0, nonzero_witness_ids=[10]), # 11 missing, nonzero adv + _make_step_end(rollout_id=0, cell_outcomes={0: [TrainStepOutcome.NORMAL]}), + ] + issues = check(events) + assert len(issues) == 1 + assert 11 in issues[0].expected_witness_ids + assert 11 not in issues[0].actual_witness_ids + + def test_zero_advantage_exclusion_is_global_across_cells(self) -> None: + """Zero-adv exclusion is unioned across cells: the per-cell snapshot reflects the global + (allreduced) gradient, so an id observed all-zero on its owning shard is excused for every cell.""" + events: list[Event] = [ + _make_allocate(rollout_id=0, witness_id_to_sample_index={10: 0, 11: 1}), + # Cell 0 owns sample 10 and observes zero advantage for it. + _make_advantage(rollout_id=0, advantages=[[0.0, 0.0], [5.0]], witness_ids=[[10], [11]], cell_index=0), + _make_snapshot(rollout_id=0, nonzero_witness_ids=[11], cell_index=0), + # Cell 1 never owned sample 10; its snapshot also lacks it and must not be flagged. + _make_advantage(rollout_id=0, advantages=[[5.0]], witness_ids=[[11]], cell_index=1), + _make_snapshot(rollout_id=0, nonzero_witness_ids=[11], cell_index=1), + _make_step_end(rollout_id=0, cell_outcomes={0: [TrainStepOutcome.NORMAL], 1: [TrainStepOutcome.NORMAL]}), + ] + assert check(events) == [] + + def test_no_advantage_event_means_no_exclusion(self) -> None: + """Without TrainAdvantageComputationEvent at all, no zero-adv exclusion — missing ID is caught.""" + events: list[Event] = [ + _make_allocate(rollout_id=0, witness_id_to_sample_index={10: 0, 11: 1}), + _make_snapshot(rollout_id=0, nonzero_witness_ids=[10]), + _make_step_end(rollout_id=0, cell_outcomes={0: [TrainStepOutcome.NORMAL]}), + ] + issues = check(events) + assert len(issues) == 1 + + def test_mixed_token_advantage_is_not_excused(self) -> None: + """A sample with mixed per-token advantages [0.0, 2.0] is not all-zero, so its missing witness id is flagged.""" + events: list[Event] = [ + _make_allocate(rollout_id=0, witness_id_to_sample_index={10: 0, 11: 1}), + _make_advantage(rollout_id=0, advantages=[[0.0, 2.0], [3.0, 3.0]], witness_ids=[[10], [11]]), + _make_snapshot(rollout_id=0, nonzero_witness_ids=[11]), # 10 missing but only partially zero-adv + _make_step_end(rollout_id=0, cell_outcomes={0: [TrainStepOutcome.NORMAL]}), + ] + issues = check(events) + assert len(issues) == 1 + assert 10 in issues[0].expected_witness_ids + assert 10 not in issues[0].actual_witness_ids + + def test_zero_advantage_uses_first_token_of_witness_id_list(self) -> None: + """A zero-advantage sample with a multi-token witness id list excuses its first id (wid_tokens[0]).""" + events: list[Event] = [ + _make_allocate(rollout_id=0, witness_id_to_sample_index={10: 0, 11: 1}), + _make_advantage(rollout_id=0, advantages=[[0.0, 0.0], [5.0]], witness_ids=[[10, 10], [11]]), + _make_snapshot(rollout_id=0, nonzero_witness_ids=[11]), # 10 missing but zero-adv + _make_step_end(rollout_id=0, cell_outcomes={0: [TrainStepOutcome.NORMAL]}), + ] + assert check(events) == [] + + def test_zero_advantage_id_reused_after_wrap_is_not_excused(self) -> None: + """Reallocation after ring wrap cancels the zero-adv excusal; the reused nonzero-adv id is required.""" + events: list[Event] = [ + _make_allocate(rollout_id=0, witness_id_to_sample_index={0: 0}), + _make_advantage(rollout_id=0, advantages=[[0.0]], witness_ids=[[0]]), + _make_snapshot(rollout_id=0, nonzero_witness_ids=[]), + _make_step_end(rollout_id=0, cell_outcomes={0: [TrainStepOutcome.NORMAL]}), + _make_allocate(rollout_id=1, witness_id_to_sample_index={0: 1}), # id 0 reused after wrap + _make_advantage(rollout_id=1, advantages=[[5.0]], witness_ids=[[0]]), + _make_snapshot(rollout_id=1, nonzero_witness_ids=[0]), + _make_step_end(rollout_id=1, cell_outcomes={0: [TrainStepOutcome.NORMAL]}), + ] + assert check(events) == [] + + def test_zero_advantage_id_reused_after_wrap_missing_is_flagged(self) -> None: + """A reused nonzero-adv id that stays absent from the snapshot is reported missing, not excused.""" + events: list[Event] = [ + _make_allocate(rollout_id=0, witness_id_to_sample_index={0: 0}), + _make_advantage(rollout_id=0, advantages=[[0.0]], witness_ids=[[0]]), + _make_snapshot(rollout_id=0, nonzero_witness_ids=[]), + _make_step_end(rollout_id=0, cell_outcomes={0: [TrainStepOutcome.NORMAL]}), + _make_allocate(rollout_id=1, witness_id_to_sample_index={0: 1}), # id 0 reused after wrap + _make_advantage(rollout_id=1, advantages=[[5.0]], witness_ids=[[0]]), + _make_snapshot(rollout_id=1, nonzero_witness_ids=[]), + _make_step_end(rollout_id=1, cell_outcomes={0: [TrainStepOutcome.NORMAL]}), + ] + issues = check(events) + assert len(issues) == 1 + assert issues[0].rollout_id == 1 + assert 0 in issues[0].expected_witness_ids + assert 0 not in issues[0].actual_witness_ids + + def test_zero_advantage_id_reused_after_wrap_and_zero_again_stays_excused(self) -> None: + """A reused id whose new sample is also zero-adv is excused again at the later rollout.""" + events: list[Event] = [ + _make_allocate(rollout_id=0, witness_id_to_sample_index={0: 0}), + _make_advantage(rollout_id=0, advantages=[[0.0]], witness_ids=[[0]]), + _make_snapshot(rollout_id=0, nonzero_witness_ids=[]), + _make_step_end(rollout_id=0, cell_outcomes={0: [TrainStepOutcome.NORMAL]}), + _make_allocate(rollout_id=1, witness_id_to_sample_index={0: 1}), # id 0 reused after wrap + _make_advantage(rollout_id=1, advantages=[[0.0]], witness_ids=[[0]]), + _make_snapshot(rollout_id=1, nonzero_witness_ids=[]), + _make_step_end(rollout_id=1, cell_outcomes={0: [TrainStepOutcome.NORMAL]}), + ] + assert check(events) == [] + + def test_zero_advantage_id_three_lives_alternating(self) -> None: + """An id alternating zero/nonzero/zero advantage across three allocations is excused, required, then excused again.""" + events: list[Event] = [ + _make_allocate(rollout_id=0, witness_id_to_sample_index={0: 0}), + _make_advantage(rollout_id=0, advantages=[[0.0]], witness_ids=[[0]]), + _make_snapshot(rollout_id=0, nonzero_witness_ids=[]), + _make_step_end(rollout_id=0, cell_outcomes={0: [TrainStepOutcome.NORMAL]}), + _make_allocate(rollout_id=1, witness_id_to_sample_index={0: 1}), + _make_advantage(rollout_id=1, advantages=[[5.0]], witness_ids=[[0]]), + _make_snapshot(rollout_id=1, nonzero_witness_ids=[0]), + _make_step_end(rollout_id=1, cell_outcomes={0: [TrainStepOutcome.NORMAL]}), + _make_allocate(rollout_id=2, witness_id_to_sample_index={0: 2}), + _make_advantage(rollout_id=2, advantages=[[0.0]], witness_ids=[[0]]), + _make_snapshot(rollout_id=2, nonzero_witness_ids=[]), + _make_step_end(rollout_id=2, cell_outcomes={0: [TrainStepOutcome.NORMAL]}), + ] + assert check(events) == [] + + def test_zero_advantage_id_reuse_cancellation_applies_per_cell(self) -> None: + """After reuse with nonzero adv the id is required on every cell — missing on one cell is flagged.""" + events: list[Event] = [ + _make_allocate(rollout_id=0, witness_id_to_sample_index={0: 0}), + _make_advantage(rollout_id=0, advantages=[[0.0]], witness_ids=[[0]], cell_index=0), + _make_snapshot(rollout_id=0, nonzero_witness_ids=[], cell_index=0), + _make_snapshot(rollout_id=0, nonzero_witness_ids=[], cell_index=1), + _make_step_end(rollout_id=0, cell_outcomes={0: [TrainStepOutcome.NORMAL], 1: [TrainStepOutcome.NORMAL]}), + _make_allocate(rollout_id=1, witness_id_to_sample_index={0: 1}), # id 0 reused after wrap + _make_advantage(rollout_id=1, advantages=[[5.0]], witness_ids=[[0]], cell_index=0), + _make_snapshot(rollout_id=1, nonzero_witness_ids=[0], cell_index=0), + _make_snapshot(rollout_id=1, nonzero_witness_ids=[], cell_index=1), # missing on cell 1 + _make_step_end(rollout_id=1, cell_outcomes={0: [TrainStepOutcome.NORMAL], 1: [TrainStepOutcome.NORMAL]}), + ] + issues = check(events) + assert len(issues) == 1 + assert issues[0].rollout_id == 1 + assert issues[0].cell_index == 1 + + def test_zero_advantage_id_never_allocated_has_no_effect(self) -> None: + """A zero-adv observation for an id that was never allocated neither excuses nor expects anything.""" + events: list[Event] = [ + _make_allocate(rollout_id=0, witness_id_to_sample_index={10: 0}), + _make_advantage(rollout_id=0, advantages=[[5.0], [0.0]], witness_ids=[[10], [99]]), + _make_snapshot(rollout_id=0, nonzero_witness_ids=[10]), + _make_step_end(rollout_id=0, cell_outcomes={0: [TrainStepOutcome.NORMAL]}), + ] + assert check(events) == [] From 493ef6893ebb15b583a10bb8c01ab9527c77ffc6 Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Mon, 22 Jun 2026 18:15:17 +0800 Subject: [PATCH 25/41] Add the event-log analyzer that applies analysis rules Add the analyzer that replays the structured event log and applies the analysis rules (checksum-consistency and witness tracing) to verify fault-tolerance behaviour offline, with unit tests. - miles/utils/event_analyzer/analyzer.py and tests. --- miles/utils/arguments.py | 5 + miles/utils/event_analyzer/analyzer.py | 42 ++++++ .../utils/event_analyzer/test_analyzer.py | 126 ++++++++++++++++++ 3 files changed, 173 insertions(+) create mode 100644 miles/utils/event_analyzer/analyzer.py create mode 100644 tests/fast/utils/event_analyzer/test_analyzer.py diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index 2f9595f9f12..a1989fb0e7e 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -1647,6 +1647,11 @@ def add_debug_arguments(parser): help="When comparing weights after update, allow quantized tensors to differ " "by up to 1 ULP of the quantized dtype per side (compared in dequantized space).", ) + parser.add_argument( + "--enable-event-analyzer", + action="store_true", + help="Enable event analyzer to run sanity checks (e.g. cross-replica checksum consistency) before each training step.", + ) parser.add_argument( "--enable-witness", action="store_true", diff --git a/miles/utils/event_analyzer/analyzer.py b/miles/utils/event_analyzer/analyzer.py new file mode 100644 index 00000000000..e7559f9589a --- /dev/null +++ b/miles/utils/event_analyzer/analyzer.py @@ -0,0 +1,42 @@ +"""Centralized event analyzer that reads events and runs all rules.""" + +import logging +from argparse import Namespace +from pathlib import Path +from typing import Any + +from miles.utils.event_analyzer.rules import ( + cross_replica_weight_checksum, + inference_engine_weight_checksum_consistency, +) +from miles.utils.event_analyzer.rules import witness as witness_rule +from miles.utils.event_logger.logger import read_events + +logger = logging.getLogger(__name__) + + +def run_analysis_from_args(args: Namespace) -> None: + if not getattr(args, "enable_event_analyzer", False): + return + + event_dir = getattr(args, "save_debug_event_data", None) + if event_dir is None: + return + + issues = run_analysis(event_dir=Path(event_dir)) + + # Fail fast, we want to stop the system if sanity check fails + if issues: + raise ValueError(f"Event analysis found issues: {issues}") + + +def run_analysis(event_dir: Path) -> list[Any]: + events = read_events(event_dir) + if not events: + return [] + + return [ + *cross_replica_weight_checksum.check(events), + *inference_engine_weight_checksum_consistency.check(events), + *witness_rule.check(events), + ] diff --git a/tests/fast/utils/event_analyzer/test_analyzer.py b/tests/fast/utils/event_analyzer/test_analyzer.py new file mode 100644 index 00000000000..9058922304e --- /dev/null +++ b/tests/fast/utils/event_analyzer/test_analyzer.py @@ -0,0 +1,126 @@ +"""Tests for event_analyzer/analyzer.py.""" + +from argparse import Namespace +from pathlib import Path + +import pytest + +from miles.utils.event_analyzer.analyzer import run_analysis, run_analysis_from_args +from miles.utils.event_logger.logger import EventLogger +from miles.utils.event_logger.models import ( + InferenceEngineWeightChecksumEvent, + TrainEngineLocalWeightChecksumEvent, + TrainEngineLocalWeightChecksumState, +) +from miles.utils.process_identity import MainProcessIdentity, TrainProcessIdentity + + +def _log_checksum_event( + event_logger: EventLogger, + *, + rollout_id: int, + param_hashes: dict[str, str] | None = None, +) -> None: + event_logger.log( + TrainEngineLocalWeightChecksumEvent, + dict( + rollout_id=rollout_id, + state=TrainEngineLocalWeightChecksumState( + param_hashes=param_hashes or {}, + buffer_hashes={}, + optimizer_hashes=[], + ), + ), + ) + + +def _make_source(*, cell_index: int = 0, rank: int = 0) -> TrainProcessIdentity: + return TrainProcessIdentity(component="actor", cell_index=cell_index, rank_within_cell=rank) + + +class TestRunAnalysis: + def test_empty_directory_returns_no_issues(self, tmp_path: Path) -> None: + assert run_analysis(event_dir=tmp_path) == [] + + def test_delegates_to_rules_and_returns_issues(self, tmp_path: Path) -> None: + # Cross-replica rule compares same rank across DIFFERENT cells, so use cell_index 0/1. + logger_a = EventLogger(log_dir=tmp_path, file_name="a.jsonl", source=_make_source(cell_index=0, rank=0)) + _log_checksum_event(logger_a, rollout_id=0, param_hashes={"pp0.w": "aaa"}) + logger_a.close() + + logger_b = EventLogger(log_dir=tmp_path, file_name="b.jsonl", source=_make_source(cell_index=1, rank=0)) + _log_checksum_event(logger_b, rollout_id=0, param_hashes={"pp0.w": "zzz"}) + logger_b.close() + + issues = run_analysis(event_dir=tmp_path) + assert len(issues) == 1 + + +def _log_inference_engine_checksum_event( + event_logger: EventLogger, + *, + rollout_id: int, + engine_checksums: list[dict[str, str]], +) -> None: + event_logger.log( + InferenceEngineWeightChecksumEvent, + dict(rollout_id=rollout_id, engine_checksums=engine_checksums), + ) + + +class TestInferenceEngineChecksumRuleWiredIn: + def test_engine_inconsistency_reported(self, tmp_path: Path) -> None: + """run_analysis surfaces engine-to-engine checksum mismatches via the registered rule.""" + event_logger = EventLogger(log_dir=tmp_path, file_name="e.jsonl", source=MainProcessIdentity()) + _log_inference_engine_checksum_event( + event_logger, rollout_id=0, engine_checksums=[{"rank0/w": "aaa"}, {"rank0/w": "zzz"}] + ) + event_logger.close() + + issues = run_analysis(event_dir=tmp_path) + assert len(issues) == 1 + + def test_consistent_engines_no_issue(self, tmp_path: Path) -> None: + """Identical engine checksums produce no issue.""" + event_logger = EventLogger(log_dir=tmp_path, file_name="e.jsonl", source=MainProcessIdentity()) + _log_inference_engine_checksum_event( + event_logger, rollout_id=0, engine_checksums=[{"rank0/w": "aaa"}, {"rank0/w": "aaa"}] + ) + event_logger.close() + + assert run_analysis(event_dir=tmp_path) == [] + + +class TestRunAnalysisFromArgs: + def test_skips_when_disabled(self) -> None: + args = Namespace(enable_event_analyzer=False, save_debug_event_data="/tmp/whatever") + run_analysis_from_args(args) + + def test_skips_when_no_event_dir(self) -> None: + args = Namespace(enable_event_analyzer=True) + run_analysis_from_args(args) + + def test_raises_on_mismatch(self, tmp_path: Path) -> None: + logger_a = EventLogger(log_dir=tmp_path, file_name="a.jsonl", source=_make_source(cell_index=0, rank=0)) + _log_checksum_event(logger_a, rollout_id=0, param_hashes={"pp0.w": "aaa"}) + logger_a.close() + + logger_b = EventLogger(log_dir=tmp_path, file_name="b.jsonl", source=_make_source(cell_index=1, rank=0)) + _log_checksum_event(logger_b, rollout_id=0, param_hashes={"pp0.w": "zzz"}) + logger_b.close() + + args = Namespace(enable_event_analyzer=True, save_debug_event_data=str(tmp_path)) + with pytest.raises(ValueError, match="issues"): + run_analysis_from_args(args) + + def test_passes_when_all_match(self, tmp_path: Path) -> None: + logger_a = EventLogger(log_dir=tmp_path, file_name="a.jsonl", source=_make_source(rank=0)) + _log_checksum_event(logger_a, rollout_id=0, param_hashes={"pp0.w": "aaa"}) + logger_a.close() + + logger_b = EventLogger(log_dir=tmp_path, file_name="b.jsonl", source=_make_source(rank=1)) + _log_checksum_event(logger_b, rollout_id=0, param_hashes={"pp0.w": "aaa"}) + logger_b.close() + + args = Namespace(enable_event_analyzer=True, save_debug_event_data=str(tmp_path)) + run_analysis_from_args(args) From 81b95c488c894cf7d589c63ba0f64c99cd442f7d Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Mon, 22 Jun 2026 18:15:17 +0800 Subject: [PATCH 26/41] Add dump and inference-engine-checksum comparison helpers for FT tests Add comparison helpers used by fault-tolerance tests to compare dumped tensors and inference-engine checksums offline (generic comparators, dump comparison, and an inference-engine checksum comparison built on the event-analyzer checksum rule). - miles/utils/test_utils/comparisons/{comparators,dumps,inference_engine_checksums}.py and tests. --- .../utils/test_utils/comparisons/__init__.py | 0 .../test_utils/comparisons/comparators.py | 55 +++++++ miles/utils/test_utils/comparisons/dumps.py | 67 ++++++++ .../comparisons/inference_engine_checksums.py | 65 ++++++++ .../utils/test_utils/comparisons/__init__.py | 0 .../test_utils/comparisons/test_dumps.py | 45 ++++++ .../test_inference_engine_checksums.py | 149 ++++++++++++++++++ 7 files changed, 381 insertions(+) create mode 100644 miles/utils/test_utils/comparisons/__init__.py create mode 100644 miles/utils/test_utils/comparisons/comparators.py create mode 100644 miles/utils/test_utils/comparisons/dumps.py create mode 100644 miles/utils/test_utils/comparisons/inference_engine_checksums.py create mode 100644 tests/fast/utils/test_utils/comparisons/__init__.py create mode 100644 tests/fast/utils/test_utils/comparisons/test_dumps.py create mode 100644 tests/fast/utils/test_utils/comparisons/test_inference_engine_checksums.py diff --git a/miles/utils/test_utils/comparisons/__init__.py b/miles/utils/test_utils/comparisons/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/miles/utils/test_utils/comparisons/comparators.py b/miles/utils/test_utils/comparisons/comparators.py new file mode 100644 index 00000000000..4051c1804ff --- /dev/null +++ b/miles/utils/test_utils/comparisons/comparators.py @@ -0,0 +1,55 @@ +import subprocess +import sys +from pathlib import Path + + +def run_comparator( + *, + baseline_path: Path, + target_path: Path, + diff_thresholds: list[tuple[str, str]], + allow_skipped_pattern: str, + allow_failed_pattern: str, + grouping_skip_keys: list[str] | None = None, + extra_args: list[str] | None, +) -> subprocess.CompletedProcess[str]: + # Skip 'rank' when grouping bundles: under FT (target) and non-FT (baseline) the same + # logical (pp_rank, cp_rank, ep_rank, tp_rank) coordinate maps to a different absolute + # rank ID (e.g. baseline rank=4 vs target cell0 rank=2 for PP=1, CP=0). Without skipping + # 'rank' the comparator gets `baseline_load_failed` for every tensor and fails with rc=1. + # Callers may pass extra keys (e.g. no_failure skips 'dp'/'edp' too). (Grouping is a + # comparator-matching detail, not a pass/fail threshold.) + skip_keys: list[str] = list(grouping_skip_keys) if grouping_skip_keys is not None else ["rank"] + assert "rank" in skip_keys, f"grouping_skip_keys must include 'rank', got {skip_keys}" + + cmd: list[str] = [ + sys.executable, + "-m", + "sglang.srt.debug_utils.comparator", + "--baseline-path", + str(baseline_path), + "--target-path", + str(target_path), + "--output-format", + "json", + "--grouping-skip-keys", + *skip_keys, + "--allow-skipped-pattern", + allow_skipped_pattern, + "--allow-failed-pattern", + allow_failed_pattern, + ] + if extra_args: + cmd.extend(extra_args) + # Keep --diff-threshold strictly last: its nargs="*" greedily consumes every + # following token, so no flag with a bare value may come after it. + cmd.append("--diff-threshold") + for pattern, predicate in diff_thresholds: + cmd.extend([pattern, predicate]) + + result: subprocess.CompletedProcess[str] = subprocess.run( + cmd, + text=True, + ) + + return result diff --git a/miles/utils/test_utils/comparisons/dumps.py b/miles/utils/test_utils/comparisons/dumps.py new file mode 100644 index 00000000000..84810671b5b --- /dev/null +++ b/miles/utils/test_utils/comparisons/dumps.py @@ -0,0 +1,67 @@ +from pathlib import Path + +from miles.utils.test_utils.comparisons.comparators import run_comparator + +# Shared regexes for model-input / metadata tensors that are not weights or grads to +# compare. Exposed as named constants (not as defaults) so each test passes them +# explicitly -- every pass/fail knob is visible at the call site, nothing is implicit. +INPUT_TENSORS_SKIP_PATTERN: str = "input_ids|positions|cu_seqlens_q|cu_seqlens_kv|qkv_format|.*witness.*" +INPUT_TENSORS_ALLOW_FAILED_PATTERN: str = "input_ids|positions|cu_seqlens_q|cu_seqlens_kv|qkv_format" + + +def compare_dumps( + baseline_dir: str, + target_dir: str, + *, + diff_thresholds: list[tuple[str, str]], + allow_skipped_pattern: str, + allow_failed_pattern: str, + phase_subdir: str | None = None, + grouping_skip_keys: list[str] | None = None, + extra_args: list[str] | None = None, +) -> None: + subdir = phase_subdir or "" + baseline_root = Path(baseline_dir) / "dumps" / subdir + target_root = Path(target_dir) / "dumps" / subdir + + assert baseline_root.exists(), f"Baseline dump dir does not exist: {baseline_root}" + assert target_root.exists(), f"Target dump dir does not exist: {target_root}" + + # Dumps are segmented into leaf dirs (e.g. fwd_bwd/rollout_), each a flat set of + # .pt files with its own per-leaf step numbering. The sglang comparator compares one + # flat dir at a time, so compare each matching leaf pair independently. + baseline_leaves = _find_leaf_dump_dirs(baseline_root) + target_leaves = _find_leaf_dump_dirs(target_root) + + assert baseline_leaves, f"No .pt dump files found under {baseline_root}" + assert baseline_leaves == target_leaves, ( + f"Dump leaf-dir mismatch: baseline={baseline_leaves} vs target={target_leaves} " + f"(under {baseline_root} vs {target_root})" + ) + + failed_leaves: list[str] = [] + for leaf in baseline_leaves: + result = run_comparator( + baseline_path=baseline_root / leaf, + target_path=target_root / leaf, + diff_thresholds=diff_thresholds, + allow_skipped_pattern=allow_skipped_pattern, + allow_failed_pattern=allow_failed_pattern, + grouping_skip_keys=grouping_skip_keys, + extra_args=extra_args, + ) + if result.returncode != 0: + failed_leaves.append(leaf) + + assert not failed_leaves, ( + f"Dump comparator failed (rc!=0) for {len(failed_leaves)}/{len(baseline_leaves)} leaf dir(s): " + f"{failed_leaves} (baseline {baseline_root} vs target {target_root}). The comparator applies the " + f"per-tensor predicates ({diff_thresholds}) and the allow/skip patterns itself; see " + f"comparator_report.jsonl under {target_root}/ for the offending tensors." + ) + print(f"Dump comparison passed: {len(baseline_leaves)} leaf dir(s) under {baseline_root} vs {target_root}") + + +def _find_leaf_dump_dirs(root: Path) -> list[str]: + leaves: set[str] = {str(p.parent.relative_to(root)) for p in root.rglob("*.pt")} + return sorted(leaves) diff --git a/miles/utils/test_utils/comparisons/inference_engine_checksums.py b/miles/utils/test_utils/comparisons/inference_engine_checksums.py new file mode 100644 index 00000000000..57e5ab3c644 --- /dev/null +++ b/miles/utils/test_utils/comparisons/inference_engine_checksums.py @@ -0,0 +1,65 @@ +from pathlib import Path + +from miles.utils.event_analyzer.rules import inference_engine_weight_checksum_consistency +from miles.utils.event_analyzer.rules.checksum_compare import ChecksumMismatchIssue, compare_flat_dicts +from miles.utils.event_logger.logger import read_events +from miles.utils.event_logger.models import InferenceEngineWeightChecksumEvent + + +def compare_inference_engine_checksums(baseline_dir: str, target_dir: str) -> None: + baseline = _read_inference_engine_checksum_events(Path(baseline_dir)) + target = _read_inference_engine_checksum_events(Path(target_dir)) + assert baseline, f"No InferenceEngineWeightChecksumEvents found in baseline dir: {baseline_dir}" + assert target, f"No InferenceEngineWeightChecksumEvents found in target dir: {target_dir}" + + # Each side's engines must already agree internally (same invariant as the production rule), so + # one representative engine per rollout then proves baseline == target regardless of engine count. + assert not inference_engine_weight_checksum_consistency.check( + baseline + ), "Baseline engines disagree with each other" + assert not inference_engine_weight_checksum_consistency.check(target), "Target engines disagree with each other" + + baseline_by_rollout = _checksums_by_rollout_id(baseline) + target_by_rollout = _checksums_by_rollout_id(target) + assert baseline_by_rollout.keys() == target_by_rollout.keys(), ( + f"Engine checksum rollout_id sets differ: " + f"baseline={sorted(baseline_by_rollout)} " + f"vs target={sorted(target_by_rollout)}" + ) + + mismatches: list[ChecksumMismatchIssue] = [] + for rollout_id in sorted(baseline_by_rollout): + mismatches += list( + compare_flat_dicts( + a=baseline_by_rollout[rollout_id], + b=target_by_rollout[rollout_id], + label_a=f"baseline/rollout_{rollout_id}", + label_b=f"target/rollout_{rollout_id}", + ) + ) + assert not mismatches, "Engine weight checksum baseline-vs-target mismatch:\n" + "\n".join( + f" - {m.label_a} vs {m.label_b} key {m.key}: {m.value_a} != {m.value_b}" for m in mismatches + ) + print(f"Engine weight checksum comparison passed: {len(baseline_by_rollout)} rollout(s) compared") + + +def _checksums_by_rollout_id(events: list[InferenceEngineWeightChecksumEvent]) -> dict[int, dict[str, str]]: + by_rollout: dict[int, dict[str, str]] = {} + for event in events: + if event.rollout_id is None: + continue + assert ( + event.rollout_id not in by_rollout + ), f"Duplicate InferenceEngineWeightChecksumEvent for rollout {event.rollout_id}" + assert event.engine_checksums, f"No engine checksums for rollout {event.rollout_id}" + by_rollout[event.rollout_id] = event.engine_checksums[0] + return by_rollout + + +def _read_inference_engine_checksum_events(dump_dir: Path) -> list[InferenceEngineWeightChecksumEvent]: + """Read all InferenceEngineWeightChecksumEvents from the events directory.""" + events_dir: Path = dump_dir / "events" + if not events_dir.exists(): + return [] + all_events = read_events(events_dir) + return [e for e in all_events if isinstance(e, InferenceEngineWeightChecksumEvent)] diff --git a/tests/fast/utils/test_utils/comparisons/__init__.py b/tests/fast/utils/test_utils/comparisons/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/fast/utils/test_utils/comparisons/test_dumps.py b/tests/fast/utils/test_utils/comparisons/test_dumps.py new file mode 100644 index 00000000000..59716c2c305 --- /dev/null +++ b/tests/fast/utils/test_utils/comparisons/test_dumps.py @@ -0,0 +1,45 @@ +from pathlib import Path + +from miles.utils.test_utils.comparisons.dumps import _find_leaf_dump_dirs + + +class TestFindLeafDumpDirs: + def test_two_pt_files_in_one_leaf_yield_single_entry(self, tmp_path: Path) -> None: + """Multiple .pt files sharing one leaf dir dedup to a single relative entry.""" + leaf = tmp_path / "fwd_bwd" / "rollout_0" + leaf.mkdir(parents=True) + (leaf / "step_0.pt").touch() + (leaf / "step_1.pt").touch() + + assert _find_leaf_dump_dirs(tmp_path) == ["fwd_bwd/rollout_0"] + + def test_two_leaves_returned_sorted(self, tmp_path: Path) -> None: + """Distinct leaf dirs are returned sorted by their relative path string.""" + leaf_b = tmp_path / "leaf_b" + leaf_a = tmp_path / "leaf_a" + leaf_b.mkdir() + leaf_a.mkdir() + (leaf_b / "x.pt").touch() + (leaf_a / "y.pt").touch() + + assert _find_leaf_dump_dirs(tmp_path) == ["leaf_a", "leaf_b"] + + def test_pt_file_directly_in_root_yields_dot(self, tmp_path: Path) -> None: + """A .pt file directly under root has parent equal to root, reported as '.'.""" + (tmp_path / "step_0.pt").touch() + + assert _find_leaf_dump_dirs(tmp_path) == ["."] + + def test_no_pt_files_yields_empty_list(self, tmp_path: Path) -> None: + """A tree with no .pt files produces an empty list.""" + (tmp_path / "sub").mkdir() + + assert _find_leaf_dump_dirs(tmp_path) == [] + + def test_non_pt_files_are_ignored(self, tmp_path: Path) -> None: + """Files not matching *.pt (including *.pth) are ignored by the glob.""" + (tmp_path / "notes.txt").touch() + (tmp_path / "weights.pth").touch() + (tmp_path / "data.json").touch() + + assert _find_leaf_dump_dirs(tmp_path) == [] diff --git a/tests/fast/utils/test_utils/comparisons/test_inference_engine_checksums.py b/tests/fast/utils/test_utils/comparisons/test_inference_engine_checksums.py new file mode 100644 index 00000000000..416cfe4078d --- /dev/null +++ b/tests/fast/utils/test_utils/comparisons/test_inference_engine_checksums.py @@ -0,0 +1,149 @@ +"""Tests for test_utils.comparisons.inference_engine_checksums.compare_inference_engine_checksums.""" + +from pathlib import Path +from typing import Any + +import pytest + +from miles.utils.event_logger.logger import EventLogger +from miles.utils.event_logger.models import InferenceEngineWeightChecksumEvent +from miles.utils.process_identity import MainProcessIdentity +from miles.utils.test_utils.comparisons.inference_engine_checksums import compare_inference_engine_checksums + + +def _write_inference_engine_events(side_dir: Path, partials: list[dict[str, Any]]) -> None: + events_dir = side_dir / "events" + event_logger = EventLogger(log_dir=events_dir, source=MainProcessIdentity()) + for partial in partials: + event_logger.log(InferenceEngineWeightChecksumEvent, partial, print_log=False) + event_logger.close() + + +def _partial(*, rollout_id: int | None, engine_checksums: list[dict[str, str]]) -> dict[str, Any]: + return dict(rollout_id=rollout_id, engine_checksums=engine_checksums) + + +class TestCompareInferenceEngineChecksums: + def test_identical_passes(self, tmp_path: Path) -> None: + """Internally-consistent sides with equal representative checksums pass.""" + partials = [_partial(rollout_id=1, engine_checksums=[{"rank0/w": "aaa"}, {"rank0/w": "aaa"}])] + _write_inference_engine_events(tmp_path / "baseline", partials) + _write_inference_engine_events(tmp_path / "target", partials) + + compare_inference_engine_checksums(str(tmp_path / "baseline"), str(tmp_path / "target")) + + def test_differing_engine_counts_still_pass(self, tmp_path: Path) -> None: + """Engine count may differ between sides; only internal agreement + representative equality matter.""" + _write_inference_engine_events( + tmp_path / "baseline", [_partial(rollout_id=1, engine_checksums=[{"rank0/w": "aaa"}])] + ) + _write_inference_engine_events( + tmp_path / "target", + [_partial(rollout_id=1, engine_checksums=[{"rank0/w": "aaa"}, {"rank0/w": "aaa"}, {"rank0/w": "aaa"}])], + ) + + compare_inference_engine_checksums(str(tmp_path / "baseline"), str(tmp_path / "target")) + + def test_none_rollout_id_skipped(self, tmp_path: Path) -> None: + """The initial out-of-loop sync (rollout_id=None) is not compared: it differs here yet the + per-rollout checksums match, so the comparison still passes.""" + _write_inference_engine_events( + tmp_path / "baseline", + [ + _partial(rollout_id=None, engine_checksums=[{"rank0/w": "init_baseline"}]), + _partial(rollout_id=1, engine_checksums=[{"rank0/w": "aaa"}]), + ], + ) + _write_inference_engine_events( + tmp_path / "target", + [ + _partial(rollout_id=None, engine_checksums=[{"rank0/w": "init_target"}]), + _partial(rollout_id=1, engine_checksums=[{"rank0/w": "aaa"}]), + ], + ) + + compare_inference_engine_checksums(str(tmp_path / "baseline"), str(tmp_path / "target")) + + def test_recurring_none_across_phases_skipped(self, tmp_path: Path) -> None: + """A multi-phase resume yields several None events per side; all are skipped, so a side with + more None events than the other still passes when the per-rollout checksums match.""" + _write_inference_engine_events( + tmp_path / "baseline", + [ + _partial(rollout_id=None, engine_checksums=[{"rank0/w": "init_a"}]), + _partial(rollout_id=2, engine_checksums=[{"rank0/w": "aaa"}]), + _partial(rollout_id=None, engine_checksums=[{"rank0/w": "init_b"}]), + _partial(rollout_id=5, engine_checksums=[{"rank0/w": "bbb"}]), + ], + ) + _write_inference_engine_events( + tmp_path / "target", + [ + _partial(rollout_id=2, engine_checksums=[{"rank0/w": "aaa"}]), + _partial(rollout_id=5, engine_checksums=[{"rank0/w": "bbb"}]), + ], + ) + + compare_inference_engine_checksums(str(tmp_path / "baseline"), str(tmp_path / "target")) + + def test_baseline_engines_disagree_fails(self, tmp_path: Path) -> None: + """If baseline's own engines disagree, the comparison fails (caught by the consistency rule).""" + _write_inference_engine_events( + tmp_path / "baseline", [_partial(rollout_id=1, engine_checksums=[{"rank0/w": "aaa"}, {"rank0/w": "zzz"}])] + ) + _write_inference_engine_events( + tmp_path / "target", [_partial(rollout_id=1, engine_checksums=[{"rank0/w": "aaa"}])] + ) + + with pytest.raises(AssertionError, match="Baseline engines disagree"): + compare_inference_engine_checksums(str(tmp_path / "baseline"), str(tmp_path / "target")) + + def test_target_engines_disagree_fails(self, tmp_path: Path) -> None: + """If target's own engines disagree, the comparison fails.""" + _write_inference_engine_events( + tmp_path / "baseline", [_partial(rollout_id=1, engine_checksums=[{"rank0/w": "aaa"}])] + ) + _write_inference_engine_events( + tmp_path / "target", [_partial(rollout_id=1, engine_checksums=[{"rank0/w": "aaa"}, {"rank0/w": "zzz"}])] + ) + + with pytest.raises(AssertionError, match="Target engines disagree"): + compare_inference_engine_checksums(str(tmp_path / "baseline"), str(tmp_path / "target")) + + def test_representative_mismatch_fails(self, tmp_path: Path) -> None: + """Internally-consistent sides whose representatives differ fail and name the tensor.""" + _write_inference_engine_events( + tmp_path / "baseline", [_partial(rollout_id=1, engine_checksums=[{"rank0/w": "aaa"}])] + ) + _write_inference_engine_events( + tmp_path / "target", [_partial(rollout_id=1, engine_checksums=[{"rank0/w": "zzz"}])] + ) + + with pytest.raises(AssertionError, match=r"key rank0/w"): + compare_inference_engine_checksums(str(tmp_path / "baseline"), str(tmp_path / "target")) + + def test_missing_rollout_fails(self, tmp_path: Path) -> None: + """A rollout present only on one side fails closed.""" + _write_inference_engine_events( + tmp_path / "baseline", + [ + _partial(rollout_id=1, engine_checksums=[{"rank0/w": "aaa"}]), + _partial(rollout_id=2, engine_checksums=[{"rank0/w": "ccc"}]), + ], + ) + _write_inference_engine_events( + tmp_path / "target", [_partial(rollout_id=1, engine_checksums=[{"rank0/w": "aaa"}])] + ) + + with pytest.raises(AssertionError, match="rollout_id sets differ"): + compare_inference_engine_checksums(str(tmp_path / "baseline"), str(tmp_path / "target")) + + def test_empty_baseline_fails(self, tmp_path: Path) -> None: + """No baseline events fails closed rather than vacuously passing.""" + _write_inference_engine_events(tmp_path / "baseline", []) + _write_inference_engine_events( + tmp_path / "target", [_partial(rollout_id=1, engine_checksums=[{"rank0/w": "aaa"}])] + ) + + with pytest.raises(AssertionError, match="No InferenceEngineWeightChecksumEvents found in baseline"): + compare_inference_engine_checksums(str(tmp_path / "baseline"), str(tmp_path / "target")) From 71d38e2316ffe5f061ed29f6c9420b26cc265363 Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Mon, 22 Jun 2026 18:15:17 +0800 Subject: [PATCH 27/41] Add metric comparison helpers for FT tests Add the metric comparison helpers used by fault-tolerance tests to compare logged training metrics offline. - miles/utils/test_utils/comparisons/metrics.py. --- miles/utils/test_utils/comparisons/metrics.py | 206 ++++++++++++++++++ .../test_utils/comparisons/test_metrics.py | 113 ++++++++++ 2 files changed, 319 insertions(+) create mode 100644 miles/utils/test_utils/comparisons/metrics.py create mode 100644 tests/fast/utils/test_utils/comparisons/test_metrics.py diff --git a/miles/utils/test_utils/comparisons/metrics.py b/miles/utils/test_utils/comparisons/metrics.py new file mode 100644 index 00000000000..70721fbf57e --- /dev/null +++ b/miles/utils/test_utils/comparisons/metrics.py @@ -0,0 +1,206 @@ +import logging +import math +from collections import defaultdict +from pathlib import Path + +import polars as pl +from sglang.srt.debug_utils.comparator.display import _render_polars_as_text + +from miles.utils.event_logger.logger import read_events +from miles.utils.event_logger.models import MetricEvent + +logger = logging.getLogger(__name__) + +_REQUIRED_METRIC_KEYS: list[str] = ["train/grad_norm", "train/loss"] + + +def compare_metrics( + baseline_dir: str, + target_dir: str, + *, + rtol: float, + atol: float, + key_prefixes: list[str], + exclude_keys: list[str], +) -> None: + baseline_events = _read_metric_events(Path(baseline_dir)) + target_events = _read_metric_events(Path(target_dir)) + + # FT retries (healing path) leave events from earlier failed attempts. Only + # the highest-attempt events per rollout_id reflect the successful run. + baseline_events = _keep_only_final_attempt(baseline_events) + target_events = _keep_only_final_attempt(target_events) + + issues: list[str] = [] + issues += _check_event_counts(baseline_events, target_events, baseline_dir, target_dir) + + if not issues: + for step_idx, (b_event, t_event) in enumerate(zip(baseline_events, target_events, strict=True)): + _print_step_comparison_table(step_idx, b_event, t_event, key_prefixes, exclude_keys=exclude_keys) + issues += _check_step_metrics( + step_idx, b_event, t_event, key_prefixes, rtol, atol=atol, exclude_keys=exclude_keys + ) + + issues += _check_required_keys_exist(baseline_events) + + assert not issues, f"MetricEvent comparison found {len(issues)} issue(s):\n" + "\n".join( + f" - {i}" for i in issues + ) + print(f"MetricEvent comparison passed: {len(baseline_events)} steps compared") + + +def _keep_only_final_attempt(events: list[MetricEvent]) -> list[MetricEvent]: + """Keep only events from the highest-attempt for each rollout_id. + + During FT healing, a crashed rollout is retried at attempt+1; events from + the failed attempt are partial and should be discarded for comparison. + + Rollout-side metrics (e.g. RolloutManager log_rollout_metrics) have + attempt=None — they are not part of the FT retry stream, so we treat them + as a single attempt (normalized to 0). + """ + + def _attempt(e: MetricEvent) -> int: + return e.attempt if e.attempt is not None else 0 + + max_attempt_by_rollout: dict[int, int] = defaultdict(int) + for e in events: + max_attempt_by_rollout[e.rollout_id] = max(max_attempt_by_rollout[e.rollout_id], _attempt(e)) + return [e for e in events if _attempt(e) == max_attempt_by_rollout[e.rollout_id]] + + +def _check_event_counts( + baseline: list[MetricEvent], + target: list[MetricEvent], + baseline_dir: str, + target_dir: str, +) -> list[str]: + issues: list[str] = [] + if len(baseline) == 0: + issues.append(f"No MetricEvents found in baseline dir: {baseline_dir}") + if len(target) == 0: + issues.append(f"No MetricEvents found in target dir: {target_dir}") + if len(baseline) > 0 and len(target) > 0 and len(baseline) != len(target): + issues.append(f"MetricEvent count mismatch: baseline={len(baseline)}, target={len(target)}") + return issues + + +def _check_step_metrics( + step_idx: int, + baseline_event: MetricEvent, + target_event: MetricEvent, + key_prefixes: list[str], + rtol: float, + *, + atol: float, + exclude_keys: list[str] | None = None, +) -> list[str]: + issues: list[str] = [] + for key in baseline_event.metrics: + if not any(key.startswith(prefix) for prefix in key_prefixes): + continue + if exclude_keys and key in exclude_keys: + continue + + if key not in target_event.metrics: + issues.append(f"Step {step_idx}: metric '{key}' present in baseline but missing in target") + continue + + issues += _check_single_metric( + step_idx, key, baseline_event.metrics[key], target_event.metrics[key], rtol, atol=atol + ) + return issues + + +def _check_single_metric( + step_idx: int, + key: str, + baseline_val: object, + target_val: object, + rtol: float, + atol: float, +) -> list[str]: + if not isinstance(baseline_val, (int, float)) or not isinstance(target_val, (int, float)): + return [] + + if math.isnan(baseline_val) or math.isnan(target_val): + return [f"Step {step_idx}, metric '{key}': NaN detected (baseline={baseline_val}, target={target_val})"] + if math.isinf(baseline_val) or math.isinf(target_val): + if baseline_val != target_val: + return [f"Step {step_idx}, metric '{key}': inf mismatch (baseline={baseline_val}, target={target_val})"] + return [] + + if baseline_val == 0.0 and target_val == 0.0: + return [] + + abs_diff = abs(baseline_val - target_val) + if abs_diff <= atol: + return [] + + rel_diff = abs_diff / max(abs(baseline_val), abs(target_val), 1e-12) + if rel_diff > rtol: + return [ + f"Step {step_idx}, metric '{key}': baseline={baseline_val}, target={target_val}, " + f"rel_diff={rel_diff:.6f} > rtol={rtol}" + ] + return [] + + +def _print_step_comparison_table( + step_idx: int, + baseline_event: MetricEvent, + target_event: MetricEvent, + key_prefixes: list[str], + *, + exclude_keys: list[str] | None = None, +) -> None: + rows: list[dict[str, str]] = [] + for key in sorted(baseline_event.metrics): + if not any(key.startswith(p) for p in key_prefixes): + continue + b_val = baseline_event.metrics[key] + t_val = target_event.metrics.get(key) + if not isinstance(b_val, (int, float)) or t_val is None or not isinstance(t_val, (int, float)): + continue + excluded = "(excluded)" if exclude_keys and key in exclude_keys else "" + abs_diff = abs(b_val - t_val) + denom = max(abs(b_val), abs(t_val), 1e-12) + rel_diff = abs_diff / denom + rows.append( + { + "metric": key, + "baseline": f"{b_val:.6e}", + "target": f"{t_val:.6e}", + "abs_diff": f"{abs_diff:.2e}", + "rel_diff": f"{rel_diff:.4%}{excluded}", + } + ) + + if not rows: + return + df = pl.DataFrame(rows) + print(_render_polars_as_text(df, title=f"Step {step_idx} metric comparison")) + + +def _check_required_keys_exist(events: list[MetricEvent]) -> list[str]: + all_keys: set[str] = set() + for event in events: + all_keys.update(event.metrics.keys()) + + issues: list[str] = [] + for required in _REQUIRED_METRIC_KEYS: + if required not in all_keys: + issues.append( + f"Required metric '{required}' not found in any baseline MetricEvent. " + f"Available keys: {sorted(all_keys)}" + ) + return issues + + +def _read_metric_events(dump_dir: Path) -> list[MetricEvent]: + """Read all MetricEvents from the events directory.""" + events_dir: Path = dump_dir / "events" + if not events_dir.exists(): + return [] + all_events = read_events(events_dir) + return [e for e in all_events if isinstance(e, MetricEvent)] diff --git a/tests/fast/utils/test_utils/comparisons/test_metrics.py b/tests/fast/utils/test_utils/comparisons/test_metrics.py new file mode 100644 index 00000000000..8b1f7a2910d --- /dev/null +++ b/tests/fast/utils/test_utils/comparisons/test_metrics.py @@ -0,0 +1,113 @@ +from datetime import datetime, timezone +from typing import Any + +from miles.utils.event_logger.models import MetricEvent +from miles.utils.process_identity import MainProcessIdentity +from miles.utils.test_utils.comparisons.metrics import _check_single_metric, _keep_only_final_attempt + +_FIXED_TS = datetime(2026, 1, 1, tzinfo=timezone.utc) +_FIXED_SOURCE = MainProcessIdentity() + + +def _metric_event( + *, rollout_id: int | None, attempt: int | None, metrics: dict[str, Any] | None = None +) -> MetricEvent: + return MetricEvent( + timestamp=_FIXED_TS, + source=_FIXED_SOURCE, + rollout_id=rollout_id, + attempt=attempt, + metrics=metrics if metrics is not None else {}, + ) + + +class TestKeepOnlyFinalAttempt: + def test_keeps_highest_attempt_for_single_rollout(self) -> None: + """Among attempts 0,1,2 for one rollout_id, only the attempt=2 event survives.""" + events = [ + _metric_event(rollout_id=1, attempt=0), + _metric_event(rollout_id=1, attempt=1), + _metric_event(rollout_id=1, attempt=2), + ] + kept = _keep_only_final_attempt(events) + assert [e.attempt for e in kept] == [2] + + def test_highest_attempt_resolved_independently_per_rollout(self) -> None: + """Each rollout_id keeps its own max attempt; different maxima coexist.""" + events = [ + _metric_event(rollout_id=1, attempt=0), + _metric_event(rollout_id=1, attempt=1), + _metric_event(rollout_id=2, attempt=0), + ] + kept = _keep_only_final_attempt(events) + assert {(e.rollout_id, e.attempt) for e in kept} == {(1, 1), (2, 0)} + + def test_none_attempt_normalized_to_zero_and_dropped_when_mixed(self) -> None: + """attempt=None normalizes to 0, so it is dropped when an attempt=1 event shares the rollout_id.""" + events = [ + _metric_event(rollout_id=1, attempt=None), + _metric_event(rollout_id=1, attempt=1), + ] + kept = _keep_only_final_attempt(events) + assert [e.attempt for e in kept] == [1] + + def test_empty_input_returns_empty(self) -> None: + """An empty event list yields an empty result.""" + assert _keep_only_final_attempt([]) == [] + + def test_ties_on_max_attempt_all_kept(self) -> None: + """Multiple events tied at the max attempt for a rollout_id are all retained.""" + events = [ + _metric_event(rollout_id=1, attempt=2, metrics={"a": 1}), + _metric_event(rollout_id=1, attempt=2, metrics={"b": 2}), + ] + kept = _keep_only_final_attempt(events) + assert len(kept) == 2 + assert [e.metrics for e in kept] == [{"a": 1}, {"b": 2}] + + +class TestCheckSingleMetric: + def test_equal_values_no_issue(self) -> None: + """Exactly equal numeric values produce no issue.""" + assert _check_single_metric(0, "k", 1.5, 1.5, rtol=0.01, atol=0.0) == [] + + def test_within_atol_no_issue(self) -> None: + """A difference within atol is accepted even if relative difference would exceed rtol.""" + assert _check_single_metric(0, "k", 1.0, 1.0 + 1e-9, rtol=0.0, atol=1e-6) == [] + + def test_relative_difference_above_rtol_reports_issue(self) -> None: + """A relative difference above rtol (and above atol) yields exactly one issue.""" + issues = _check_single_metric(3, "train/loss", 1.0, 2.0, rtol=0.1, atol=0.0) + assert len(issues) == 1 + assert "train/loss" in issues[0] + assert "rel_diff" in issues[0] + + def test_nan_detected(self) -> None: + """A NaN on either side produces a 'NaN detected' issue.""" + issues = _check_single_metric(0, "k", float("nan"), 1.0, rtol=0.1, atol=0.0) + assert len(issues) == 1 + assert "NaN detected" in issues[0] + + def test_matching_inf_no_issue(self) -> None: + """inf == inf compares equal and produces no issue.""" + assert _check_single_metric(0, "k", float("inf"), float("inf"), rtol=0.1, atol=0.0) == [] + + def test_inf_vs_finite_reports_mismatch(self) -> None: + """inf versus a finite value produces an 'inf mismatch' issue.""" + issues = _check_single_metric(0, "k", float("inf"), 1.0, rtol=0.1, atol=0.0) + assert len(issues) == 1 + assert "inf mismatch" in issues[0] + + def test_both_zero_no_issue(self) -> None: + """Two exact zeros short-circuit to no issue.""" + assert _check_single_metric(0, "k", 0.0, 0.0, rtol=0.0, atol=0.0) == [] + + def test_non_numeric_skipped(self) -> None: + """A non-numeric value on either side is skipped (no issue).""" + assert _check_single_metric(0, "k", "abc", 1.0, rtol=0.0, atol=0.0) == [] + + def test_tiny_baseline_uses_relative_floor(self) -> None: + """A near-zero baseline uses the 1e-12 denominator floor, making a tiny abs diff a large rel diff.""" + issues = _check_single_metric(0, "k", 0.0, 5e-13, rtol=0.1, atol=0.0) + assert len(issues) == 1 + assert "rel_diff" in issues[0] From c043ccf07c1d2890560cd65f0e61c2ecdcb62bbd Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Mon, 22 Jun 2026 18:15:17 +0800 Subject: [PATCH 28/41] Add reconfiguration assertions for fault-tolerance tests Add reusable reconfiguration assertions used by fault-tolerance tests to verify cell reconfigure / healing behaviour, with unit tests. - miles/utils/test_utils/reconfigure_assertions.py and tests. --- .../test_utils/reconfigure_assertions.py | 59 +++++++++ .../test_utils/test_reconfigure_assertions.py | 122 ++++++++++++++++++ 2 files changed, 181 insertions(+) create mode 100644 miles/utils/test_utils/reconfigure_assertions.py create mode 100644 tests/fast/utils/test_utils/test_reconfigure_assertions.py diff --git a/miles/utils/test_utils/reconfigure_assertions.py b/miles/utils/test_utils/reconfigure_assertions.py new file mode 100644 index 00000000000..23c6292bd53 --- /dev/null +++ b/miles/utils/test_utils/reconfigure_assertions.py @@ -0,0 +1,59 @@ +from pathlib import Path + +from miles.utils.event_logger.logger import read_events +from miles.utils.event_logger.models import CellReconfigureEvent +from miles.utils.pydantic_utils import FrozenStrictBaseModel + + +class ReconfigureInfo(FrozenStrictBaseModel): + rollout_id: int + src_cell_index: int | None + healed_cell_indices: list[int] + alive_cell_indices_after: list[int] + + @staticmethod + def from_event(event: CellReconfigureEvent) -> "ReconfigureInfo": + return ReconfigureInfo( + rollout_id=event.rollout_id, + src_cell_index=event.src_cell_index, + healed_cell_indices=event.healed_cell_indices, + alive_cell_indices_after=event.alive_cell_indices_after, + ) + + +def assert_reconfigure_events(event_dir: Path, *, expected: list[ReconfigureInfo]) -> None: + assert event_dir.is_dir(), f"Event directory {event_dir} does not exist or is not a directory" + actual = [ReconfigureInfo.from_event(event) for event in load_reconfigure_events(event_dir)] + assert actual == expected, ( + f"CellReconfigureEvent sequence mismatch in {event_dir}:\n" f" expected: {expected}\n" f" actual: {actual}" + ) + + +MIN_SOAK_INJECTIONS: int = 2 +MIN_SOAK_HEALINGS: int = 2 + + +def assert_soak_reconfigure_events(event_dir: Path, *, num_successful_injections: int) -> None: + assert event_dir.is_dir(), f"Event directory {event_dir} does not exist or is not a directory" + events = load_reconfigure_events(event_dir) + healings = [event for event in events if event.healed_cell_indices] + + assert num_successful_injections >= MIN_SOAK_INJECTIONS, ( + f"Soak proved too little in {event_dir}: the fault injector reported only " + f"{num_successful_injections} successful injection(s), need >= {MIN_SOAK_INJECTIONS} " + f"to exercise fault recovery more than once" + ) + assert len(healings) >= MIN_SOAK_HEALINGS, ( + f"Healing witness failed in {event_dir}: {num_successful_injections} successful injection(s) " + f"but only {len(healings)} healing event(s), need >= {MIN_SOAK_HEALINGS} " + f"(reconfigure events: {[ReconfigureInfo.from_event(event) for event in events]})" + ) + + print( + f"Soak reconfigure witness assertion passed: {len(events)} reconfigure event(s) " + f"({len(healings)} healing(s)) for {num_successful_injections} successful injection(s) in {event_dir}" + ) + + +def load_reconfigure_events(event_dir: Path) -> list[CellReconfigureEvent]: + return [event for event in read_events(event_dir) if isinstance(event, CellReconfigureEvent)] diff --git a/tests/fast/utils/test_utils/test_reconfigure_assertions.py b/tests/fast/utils/test_utils/test_reconfigure_assertions.py new file mode 100644 index 00000000000..6873de03eaf --- /dev/null +++ b/tests/fast/utils/test_utils/test_reconfigure_assertions.py @@ -0,0 +1,122 @@ +from pathlib import Path +from typing import Any + +import pytest + +from miles.utils.event_logger.logger import EventLogger +from miles.utils.event_logger.models import CellReconfigureEvent, TrainGroupStepEndEvent +from miles.utils.process_identity import MainProcessIdentity +from miles.utils.test_utils.reconfigure_assertions import ( + ReconfigureInfo, + assert_reconfigure_events, + assert_soak_reconfigure_events, + load_reconfigure_events, +) + +_SHRINK_PARTIAL: dict[str, Any] = dict( + rollout_id=2, + quorum_id=1, + src_cell_index=None, + healed_cell_indices=[], + alive_cell_indices_after=[0], +) +_HEALING_PARTIAL: dict[str, Any] = dict( + rollout_id=3, + quorum_id=2, + src_cell_index=0, + healed_cell_indices=[1], + alive_cell_indices_after=[0, 1], +) + +_SHRINK_EXPECTED = ReconfigureInfo( + rollout_id=2, src_cell_index=None, healed_cell_indices=[], alive_cell_indices_after=[0] +) +_HEALING_EXPECTED = ReconfigureInfo( + rollout_id=3, src_cell_index=0, healed_cell_indices=[1], alive_cell_indices_after=[0, 1] +) + + +def _write_events(log_dir: Path, partials: list[dict[str, Any]]) -> None: + event_logger = EventLogger(log_dir=log_dir, source=MainProcessIdentity()) + for partial in partials: + event_logger.log(CellReconfigureEvent, partial, print_log=False) + event_logger.close() + + +class TestLoadReconfigureEvents: + def test_filters_other_event_types_and_preserves_order(self, tmp_path: Path) -> None: + """Only CellReconfigureEvents are returned, in file (emission) order.""" + event_logger = EventLogger(log_dir=tmp_path, source=MainProcessIdentity()) + event_logger.log(CellReconfigureEvent, _SHRINK_PARTIAL, print_log=False) + event_logger.log(TrainGroupStepEndEvent, dict(rollout_id=2, cell_outcomes={}), print_log=False) + event_logger.log(CellReconfigureEvent, _HEALING_PARTIAL, print_log=False) + event_logger.close() + + events = load_reconfigure_events(tmp_path) + + assert [e.rollout_id for e in events] == [2, 3] + assert all(isinstance(e, CellReconfigureEvent) for e in events) + + def test_empty_dir_returns_no_events(self, tmp_path: Path) -> None: + """A directory without any JSONL files yields an empty event list.""" + assert load_reconfigure_events(tmp_path) == [] + + +class TestAssertReconfigureEvents: + def test_passes_on_exact_sequence(self, tmp_path: Path) -> None: + """An exactly matching shrink+healing sequence with contiguous quorum ids passes.""" + _write_events(tmp_path, [_SHRINK_PARTIAL, _HEALING_PARTIAL]) + + assert_reconfigure_events(tmp_path, expected=[_SHRINK_EXPECTED, _HEALING_EXPECTED]) + + def test_passes_on_empty_expectation(self, tmp_path: Path) -> None: + """Expecting zero reconfigures passes when no events were emitted.""" + assert_reconfigure_events(tmp_path, expected=[]) + + def test_missing_healing_fails_sequence_check(self, tmp_path: Path) -> None: + """A run that never healed fails the exact-sequence comparison (expected healing, got nothing).""" + _write_events(tmp_path, []) + + with pytest.raises(AssertionError, match="sequence mismatch"): + assert_reconfigure_events(tmp_path, expected=[_HEALING_EXPECTED]) + + def test_unexpected_extra_healing_fails(self, tmp_path: Path) -> None: + """A healing event in a run expected to have none fails the exact-sequence comparison.""" + _write_events(tmp_path, [dict(_HEALING_PARTIAL, quorum_id=1)]) + + with pytest.raises(AssertionError, match="sequence mismatch"): + assert_reconfigure_events(tmp_path, expected=[]) + + def test_wrong_rollout_id_fails_sequence_check(self, tmp_path: Path) -> None: + """A healing at the wrong rollout fails the exact-sequence comparison.""" + _write_events(tmp_path, [dict(_HEALING_PARTIAL, rollout_id=9, quorum_id=1)]) + + with pytest.raises(AssertionError, match="sequence mismatch"): + assert_reconfigure_events(tmp_path, expected=[_HEALING_EXPECTED]) + + +class TestAssertSoakReconfigureEvents: + def test_passes_when_enough_injections_and_healings(self, tmp_path: Path) -> None: + """>=2 successful injections with >=2 healing events pass the soak witness.""" + _write_events(tmp_path, [_SHRINK_PARTIAL, _HEALING_PARTIAL, dict(_HEALING_PARTIAL, rollout_id=5, quorum_id=3)]) + + assert_soak_reconfigure_events(tmp_path, num_successful_injections=2) + + def test_fails_when_no_injections(self, tmp_path: Path) -> None: + """Zero successful injections means no fault tolerance was exercised, so the witness fails.""" + with pytest.raises(AssertionError, match="proved too little"): + assert_soak_reconfigure_events(tmp_path, num_successful_injections=0) + + def test_fails_when_too_few_injections(self, tmp_path: Path) -> None: + """A single injection is below the soak minimum even when healing events are present.""" + _write_events(tmp_path, [_HEALING_PARTIAL, dict(_HEALING_PARTIAL, rollout_id=5, quorum_id=3)]) + + with pytest.raises(AssertionError, match="proved too little"): + assert_soak_reconfigure_events(tmp_path, num_successful_injections=1) + + def test_fails_when_too_few_healings(self, tmp_path: Path) -> None: + """Enough injections but fewer than the required healing events fail the witness.""" + _write_events(tmp_path, [_SHRINK_PARTIAL, _HEALING_PARTIAL]) + + with pytest.raises(AssertionError, match="Healing witness failed"): + assert_soak_reconfigure_events(tmp_path, num_successful_injections=3) From 932660e8a4a9948e0be92ef75de87d3d8fe67d5a Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Mon, 22 Jun 2026 18:15:17 +0800 Subject: [PATCH 29/41] Relocate GroupInfo into shared process-group utilities Move the GroupInfo dataclass out of training_utils/parallel.py into a shared process_group_utils module that also provides multi-process-group helpers (GeneralPGUtil / MultiPGUtil) for collectives over single or hierarchical process groups, used by the cross-replica / effective-DP code paths. - process_group_utils.py: GroupInfo + GeneralPGUtil / MultiPGUtil (+ tests). - training_utils/parallel.py: import GroupInfo from the shared module. - megatron / fsdp parallel.py: import GroupInfo from the shared module. - distributed_utils.py: use GeneralPGUtil for masked-whiten all-reduce. --- .../experimental/fsdp_utils/parallel.py | 4 +- miles/backends/megatron_utils/parallel.py | 4 +- miles/backends/training_utils/parallel.py | 30 +- miles/utils/distributed_utils.py | 3 +- miles/utils/process_group_utils.py | 343 +++++++++++++++ tests/fast/dist_utils.py | 23 + tests/fast/utils/test_process_group_utils.py | 395 ++++++++++++++++++ 7 files changed, 771 insertions(+), 31 deletions(-) create mode 100644 miles/utils/process_group_utils.py create mode 100644 tests/fast/dist_utils.py create mode 100644 tests/fast/utils/test_process_group_utils.py diff --git a/miles/backends/experimental/fsdp_utils/parallel.py b/miles/backends/experimental/fsdp_utils/parallel.py index 0f49fd04ff2..0bbff9f5425 100644 --- a/miles/backends/experimental/fsdp_utils/parallel.py +++ b/miles/backends/experimental/fsdp_utils/parallel.py @@ -6,7 +6,9 @@ from miles.utils.distributed_utils import get_gloo_group -from ...training_utils.parallel import GroupInfo, ParallelState +from miles.utils.process_group_utils import GroupInfo + +from ...training_utils.parallel import ParallelState logger = logging.getLogger(__name__) diff --git a/miles/backends/megatron_utils/parallel.py b/miles/backends/megatron_utils/parallel.py index 38b52ed5a26..f44805e9b8f 100644 --- a/miles/backends/megatron_utils/parallel.py +++ b/miles/backends/megatron_utils/parallel.py @@ -8,7 +8,9 @@ from megatron.core.utils import get_model_config from megatron.training.global_vars import get_args -from ..training_utils.parallel import GroupInfo, ParallelState, get_parallel_state +from miles.utils.process_group_utils import GroupInfo + +from ..training_utils.parallel import ParallelState, get_parallel_state logger = logging.getLogger(__name__) diff --git a/miles/backends/training_utils/parallel.py b/miles/backends/training_utils/parallel.py index 2d8d6e082c1..0f8f169ee9f 100644 --- a/miles/backends/training_utils/parallel.py +++ b/miles/backends/training_utils/parallel.py @@ -1,6 +1,7 @@ from dataclasses import dataclass -import torch.distributed as dist + +from miles.utils.process_group_utils import GroupInfo _parallel_state: "ParallelState | None" = None @@ -16,33 +17,6 @@ def get_parallel_state() -> "ParallelState": return _parallel_state -@dataclass(frozen=True) -class GroupInfo: - rank: int - size: int - group: dist.ProcessGroup | None - gloo_group: dist.ProcessGroup | None = None - - def __post_init__(self) -> None: - self._verify_group(self.group, "group") - self._verify_group(self.gloo_group, "gloo_group") - - def _verify_group(self, group: dist.ProcessGroup | None, name: str) -> None: - if group is None: - return - if not _is_native_process_group(group): - return - actual_rank = dist.get_rank(group) - actual_size = dist.get_world_size(group) - assert actual_rank == self.rank, f"{name}: rank mismatch: expected {self.rank}, got {actual_rank}" - assert actual_size == self.size, f"{name}: size mismatch: expected {self.size}, got {actual_size}" - - -def _is_native_process_group(group: dist.ProcessGroup) -> bool: - # torchft's ProcessGroup - return not hasattr(group, "_replica_id") - - @dataclass class ParallelState: """Core parallel state shared across all backends. diff --git a/miles/utils/distributed_utils.py b/miles/utils/distributed_utils.py index 692bfaaa563..7bb76b2ed63 100644 --- a/miles/utils/distributed_utils.py +++ b/miles/utils/distributed_utils.py @@ -14,6 +14,7 @@ rendezvous, ) +from miles.utils.process_group_utils import GeneralPGUtil GLOO_GROUP = None @@ -129,7 +130,7 @@ def distributed_masked_whiten( ) # Aggregate via all_reduce within the DP group - dist.all_reduce(stats_tensor, group=process_group) + GeneralPGUtil.create(process_group).all_reduce(stats_tensor, process_group, op=dist.ReduceOp.SUM) # Calculate global stats from aggregated results global_sum, global_sum_sq, global_mask_sum = stats_tensor diff --git a/miles/utils/process_group_utils.py b/miles/utils/process_group_utils.py new file mode 100644 index 00000000000..a929893555c --- /dev/null +++ b/miles/utils/process_group_utils.py @@ -0,0 +1,343 @@ +from collections.abc import Sequence +from dataclasses import dataclass, field +from typing import Any + +import torch +import torch.distributed as dist + +# AllgatherOptions is not re-exported by torch.distributed (unlike +# AllreduceOptions, BroadcastOptions, GatherOptions). PyTorch omission. +from torch._C._distributed_c10d import AllgatherOptions +from torch.distributed.distributed_c10d import _get_object_coll_device, _object_to_tensor, _tensor_to_object + +from miles.utils.det_process_group import DET_NCCL_BACKEND_NAME, det_all_reduce + + +def _is_det_world() -> bool: + return dist.is_initialized() and dist.get_backend() == DET_NCCL_BACKEND_NAME + + +@dataclass(frozen=True) +class GroupInfo: + rank: int + size: int + group: dist.ProcessGroup | None + gloo_group: dist.ProcessGroup | None = None + src_rank: int | None = None + debug_info: dict[str, Any] = field(default_factory=dict) + + def __post_init__(self) -> None: + self._verify_group(self.group, "group") + self._verify_group(self.gloo_group, "gloo_group") + + def _verify_group(self, group: dist.ProcessGroup | None, name: str) -> None: + if group is None: + return + util = GeneralPGUtil.create(group) + actual_rank = util.get_rank(group) + actual_size = util.get_size(group) + assert actual_rank == self.rank, f"{name}: rank mismatch: expected {self.rank}, got {actual_rank}" + assert actual_size == self.size, f"{name}: size mismatch: expected {self.size}, got {actual_size}" + + +@dataclass(frozen=True) +class GroupsInfo: + rank: int + size: int + groups_inner_to_outer: list[dist.ProcessGroup] + gloo_groups_inner_to_outer: list[dist.ProcessGroup] + + @classmethod + def from_single(cls, info: GroupInfo) -> "GroupsInfo": + return cls( + rank=info.rank, + size=info.size, + groups_inner_to_outer=[info.group], + gloo_groups_inner_to_outer=[info.gloo_group], + ) + + @classmethod + def from_pair(cls, *, inner: GroupInfo, outer: GroupInfo) -> "GroupsInfo": + return cls( + rank=outer.rank * inner.size + inner.rank, + size=outer.size * inner.size, + groups_inner_to_outer=[inner.group, outer.group], + gloo_groups_inner_to_outer=[inner.gloo_group, outer.gloo_group], + ) + + +class GeneralPGUtil: + """Process group operations that work with both native and torchft PGs. + + Use GeneralPGUtil.create(group) to get the appropriate implementation. + """ + + @staticmethod + def create(group: dist.ProcessGroup) -> "GeneralPGUtil": + if not hasattr(group, "_replica_id"): + return _NativePGUtil() + return _RawPGUtil() + + def get_rank(self, group: dist.ProcessGroup) -> int: + raise NotImplementedError + + def get_size(self, group: dist.ProcessGroup) -> int: + raise NotImplementedError + + def all_reduce(self, tensor: torch.Tensor, group: dist.ProcessGroup, op: dist.ReduceOp) -> None: + raise NotImplementedError + + def reduce(self, tensor: torch.Tensor, group: dist.ProcessGroup, op: dist.ReduceOp) -> None: + raise NotImplementedError + + def broadcast(self, tensor: torch.Tensor, group: dist.ProcessGroup) -> None: + raise NotImplementedError + + def barrier(self, group: dist.ProcessGroup) -> None: + raise NotImplementedError + + def all_gather( + self, output_tensors: list[torch.Tensor], input_tensor: torch.Tensor, group: dist.ProcessGroup + ) -> None: + raise NotImplementedError + + def gather( + self, + input_tensor: torch.Tensor, + gather_list: list[torch.Tensor] | None, + group: dist.ProcessGroup, + ) -> None: + raise NotImplementedError + + def gather_object(self, obj: Any, object_gather_list: list[Any] | None, group: dist.ProcessGroup) -> None: + raise NotImplementedError + + +class _NativePGUtil(GeneralPGUtil): + def get_rank(self, group: dist.ProcessGroup) -> int: + return dist.get_rank(group) + + def get_size(self, group: dist.ProcessGroup) -> int: + return dist.get_world_size(group) + + def all_reduce(self, tensor: torch.Tensor, group: dist.ProcessGroup, op: dist.ReduceOp) -> None: + dist.all_reduce(tensor, op=op, group=group) + + def reduce(self, tensor: torch.Tensor, group: dist.ProcessGroup, op: dist.ReduceOp) -> None: + dist.reduce(tensor, dst=dist.get_global_rank(group, 0), op=op, group=group) + + def broadcast(self, tensor: torch.Tensor, group: dist.ProcessGroup) -> None: + dist.broadcast(tensor, src=dist.get_global_rank(group, 0), group=group) + + def barrier(self, group: dist.ProcessGroup) -> None: + dist.barrier(group=group) + + def all_gather( + self, output_tensors: list[torch.Tensor], input_tensor: torch.Tensor, group: dist.ProcessGroup + ) -> None: + dist.all_gather(output_tensors, input_tensor, group=group) + + def gather( + self, + input_tensor: torch.Tensor, + gather_list: list[torch.Tensor] | None, + group: dist.ProcessGroup, + ) -> None: + dist.gather(input_tensor, gather_list=gather_list, dst=dist.get_global_rank(group, 0), group=group) + + def gather_object(self, obj: Any, object_gather_list: list[Any] | None, group: dist.ProcessGroup) -> None: + dist.gather_object(obj, object_gather_list, dst=dist.get_global_rank(group, 0), group=group) + + +def _check_wait(work: dist._Work, op_name: str) -> None: + """Call work.wait() and raise on failure. + + Failure modes depend on the backend: + - Native PyTorch NCCL: always raises on failure, never returns False. + - torchft ProcessGroupNCCL: may raise (e.g. ncclCommAbort unblocks native + wait which throws) or return False, depending on the failure path. + This helper handles both. + """ + success = work.wait() + if not success: + raise RuntimeError(f"distributed operation {op_name} failed (wait returned False)") + + +class _RawPGUtil(GeneralPGUtil): + def get_rank(self, group: dist.ProcessGroup) -> int: + return group._rank + + def get_size(self, group: dist.ProcessGroup) -> int: + return group.size() + + def all_reduce(self, tensor: torch.Tensor, group: dist.ProcessGroup, op: dist.ReduceOp) -> None: + if (op == dist.ReduceOp.SUM or op == dist.ReduceOp.AVG) and _is_det_world(): + det_all_reduce(tensor, group=group, reduce_op=op) + return + + opts = dist.AllreduceOptions() + opts.reduceOp = op + _check_wait(group.allreduce([tensor], opts), "allreduce") + + def reduce(self, tensor: torch.Tensor, group: dist.ProcessGroup, op: dist.ReduceOp) -> None: + # TODO(torchft): switch to real reduce once torchft adds it to ProcessGroupWrapper. + # torchft ProcessGroupWrapper doesn't override reduce() — calling it hits + # the base class which errors with "No backend type associated with device". + # allreduce is a safe substitute: the only caller is MultiPGUtil.all_reduce + # (reduce+broadcast), so all ranks holding the result is equivalent. + self.all_reduce(tensor, group, op) + + def broadcast(self, tensor: torch.Tensor, group: dist.ProcessGroup) -> None: + opts = dist.BroadcastOptions() + opts.rootRank = 0 + _check_wait(group.broadcast([tensor], opts), "broadcast") + + def barrier(self, group: dist.ProcessGroup) -> None: + # Must .wait(): torchft arms its per-collective timeout only inside Work.wait(); + # a fire-and-forget barrier hangs a dead-peer cell until the NCCL watchdog kills it. + _check_wait(group.barrier(), "barrier") + + def all_gather( + self, output_tensors: list[torch.Tensor], input_tensor: torch.Tensor, group: dist.ProcessGroup + ) -> None: + _check_wait(group.allgather([output_tensors], [input_tensor], AllgatherOptions()), "allgather") + + def gather( + self, + input_tensor: torch.Tensor, + gather_list: list[torch.Tensor] | None, + group: dist.ProcessGroup, + ) -> None: + # TODO(torchft): switch to real gather once torchft adds it to ProcessGroupWrapper. + # torchft ProcessGroupWrapper doesn't override gather() — calling it hits + # the base class which errors with "No backend type associated with device". + # allgather is a safe substitute: rank 0 extracts its gather_list from the + # full allgather result, other ranks discard. + group_size = self.get_size(group) + all_tensors = [torch.empty_like(input_tensor) for _ in range(group_size)] + self.all_gather(all_tensors, input_tensor, group) + if gather_list is not None: + for i in range(group_size): + gather_list[i].copy_(all_tensors[i]) + + def gather_object(self, obj: Any, object_gather_list: list[Any] | None, group: dist.ProcessGroup) -> None: + _gather_object_via_util(self, obj, object_gather_list, group) + + +class MultiPGUtil: + """Operations across multiple process groups (inner-to-outer).""" + + @staticmethod + def all_reduce( + tensor: torch.Tensor, + groups_inner_to_outer: Sequence[dist.ProcessGroup], + op: dist.ReduceOp, + ) -> None: + """Reduce then broadcast across multiple groups for bitwise-equal results. + + Inner-to-outer reduce collapses values to the global root (rank 0 in every + group). Outer-to-inner broadcast fans the result back out. Because broadcast + is a pure copy, all ranks receive a bitwise-identical result regardless of + floating-point non-determinism in the reduce path. + """ + for group in groups_inner_to_outer: + GeneralPGUtil.create(group).reduce(tensor, group, op) + + for group in reversed(groups_inner_to_outer): + GeneralPGUtil.create(group).broadcast(tensor, group) + + @staticmethod + def gather_object( + obj: Any, + groups_inner_to_outer: Sequence[dist.ProcessGroup], + ) -> list[Any] | None: + """Gather objects across multiple groups. Returns full list on rank 0, None on others.""" + objects = [obj] + for group in groups_inner_to_outer: + util = GeneralPGUtil.create(group) + rank = util.get_rank(group) + size = util.get_size(group) + if rank == 0: + gathered: list[Any] = [None] * size + util.gather_object(objects, gathered, group=group) + objects = [item for sublist in gathered for item in sublist] + else: + util.gather_object(objects, None, group=group) + return None + + return objects + + +def _gather_object_via_util( + util: GeneralPGUtil, + obj: Any, + object_gather_list: list[Any] | None, + group: dist.ProcessGroup, +) -> None: + """gather_object implemented using GeneralPGUtil primitives. + + Always gathers to group-local rank 0. + + Copied from torch.distributed.distributed_c10d.gather_object (PyTorch v2.11.0) + (https://github.com/pytorch/pytorch/blob/v2.11.0/torch/distributed/distributed_c10d.py) + with the following modifications: + - Replaced dist.get_rank()/get_world_size() with util.get_rank()/get_size() + - Replaced dist.all_gather()/dist.gather() with util.all_gather()/util.gather() + - Removed _rank_not_in_group check, group_dst parameter + - Hardcoded dst=0 (always gather to first rank) + - Hardcoded cpu device (was: _get_object_coll_device) + - Inlined _validate_output_list_for_rank as simple assert + - Removed redundant post-gather None check on object_gather_list (already asserted at function entry) + """ + # --- Begin: adapted from PyTorch v2.11.0 gather_object --- + + my_group_rank = util.get_rank(group) # was: group.rank() + if my_group_rank == 0: + assert object_gather_list is not None + else: + assert object_gather_list is None + + current_device = _get_object_coll_device(group) + input_tensor, local_size = _object_to_tensor(obj, current_device, group) + + # Gather all local sizes. This is so that we can find the max size, and index + # until the correct size when deserializing the tensors. + group_size = util.get_size(group) # was: get_world_size(group=group) + object_sizes_tensor = torch.zeros(group_size, dtype=torch.long, device=current_device) + object_size_list = [object_sizes_tensor[i].unsqueeze(dim=0) for i in range(group_size)] + # Allgather tensor sizes. An all-gather is needed here despite this being a + # gather, since each rank needs to broadcast a tensor of the same (maximal) + # size. + util.all_gather(object_size_list, local_size, group=group) # was: all_gather(..., group=group) + max_object_size = int(max(object_size_list).item()) + # Resize tensor to max size across all ranks. + input_tensor.resize_(max_object_size) + # Avoid populating output tensors if the result won't be gathered on this rank. + if my_group_rank == 0: + coalesced_output_tensor = torch.empty(max_object_size * group_size, dtype=torch.uint8, device=current_device) + # Output tensors are nonoverlapping views of coalesced_output_tensor + output_tensors = [ + coalesced_output_tensor[max_object_size * i : max_object_size * (i + 1)] for i in range(group_size) + ] + # All ranks call gather with equal-sized tensors. + util.gather( + input_tensor, + gather_list=output_tensors if my_group_rank == 0 else None, + group=group, + ) + if my_group_rank != 0: + return + + for i, tensor in enumerate(output_tensors): + tensor = tensor.type(torch.uint8) + tensor_size = object_size_list[i] + object_gather_list[i] = _tensor_to_object(tensor, tensor_size, group) + + # --- End: adapted from PyTorch v2.11.0 gather_object --- + + +def collective_bool_and(*, value: bool, group: dist.ProcessGroup) -> bool: + """Make a bool `and` operation on all ranks in this process group""" + tensor = torch.tensor([1.0 if value else 0.0], dtype=torch.float32, device=_get_object_coll_device(group)) + GeneralPGUtil.create(group).all_reduce(tensor, group, op=dist.ReduceOp.MIN) + return tensor.item() > 0.5 diff --git a/tests/fast/dist_utils.py b/tests/fast/dist_utils.py new file mode 100644 index 00000000000..b8a958bb5e4 --- /dev/null +++ b/tests/fast/dist_utils.py @@ -0,0 +1,23 @@ +import os +import socket +from typing import Any + +import torch.distributed as dist +import torch.multiprocessing as mp + + +def find_free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("", 0)) + return s.getsockname()[1] + + +def init_gloo(rank: int, world_size: int, *, port: int) -> None: + os.environ["MASTER_ADDR"] = "localhost" + os.environ["MASTER_PORT"] = str(port) + dist.init_process_group(backend="gloo", rank=rank, world_size=world_size) + + +def run_multiprocess(fn: Any, world_size: int = 2) -> None: + port = find_free_port() + mp.spawn(fn, args=(world_size, port), nprocs=world_size, join=True) diff --git a/tests/fast/utils/test_process_group_utils.py b/tests/fast/utils/test_process_group_utils.py new file mode 100644 index 00000000000..a3624cc7ff4 --- /dev/null +++ b/tests/fast/utils/test_process_group_utils.py @@ -0,0 +1,395 @@ +"""Tests for process_group_utils: GroupInfo, GroupsInfo, GeneralPGUtil, MultiPGUtil.""" + +from typing import Any +from unittest.mock import MagicMock, patch + +import pytest +import torch +import torch.distributed as dist +from tests.fast.dist_utils import init_gloo, run_multiprocess +from torch.distributed.device_mesh import init_device_mesh + +from miles.utils.process_group_utils import ( + GroupInfo, + GroupsInfo, + MultiPGUtil, + _check_wait, + _NativePGUtil, + _RawPGUtil, + collective_bool_and, +) + + +def _make_mesh(): + return init_device_mesh("cpu", mesh_shape=(2, 2), mesh_dim_names=("outer", "inner")) + + +# -- GroupInfo / GroupsInfo tests (no distributed needed) -- + + +class TestGroupInfo: + def test_construction_with_none_group(self) -> None: + info = GroupInfo(rank=0, size=4, group=None) + assert info.rank == 0 + assert info.size == 4 + assert info.gloo_group is None + + +class TestGroupsInfo: + def test_from_single(self) -> None: + info = GroupInfo(rank=2, size=4, group=None) + result = GroupsInfo.from_single(info) + assert result.rank == 2 + assert result.size == 4 + assert result.groups_inner_to_outer == [None] + assert result.gloo_groups_inner_to_outer == [None] + + @patch.object(GroupInfo, "__post_init__", lambda self: None) + def test_from_single_with_gloo(self) -> None: + sentinel_group = object() + sentinel_gloo = object() + info = GroupInfo(rank=0, size=2, group=sentinel_group, gloo_group=sentinel_gloo) + result = GroupsInfo.from_single(info) + assert result.groups_inner_to_outer == [sentinel_group] + assert result.gloo_groups_inner_to_outer == [sentinel_gloo] + + def test_from_pair(self) -> None: + inner = GroupInfo(rank=1, size=3, group=None) + outer = GroupInfo(rank=2, size=4, group=None) + result = GroupsInfo.from_pair(inner=inner, outer=outer) + assert result.rank == 2 * 3 + 1 # 7 + assert result.size == 4 * 3 # 12 + assert result.gloo_groups_inner_to_outer == [None, None] + + @patch.object(GroupInfo, "__post_init__", lambda self: None) + def test_from_pair_with_gloo(self) -> None: + inner_gloo = object() + outer_gloo = object() + inner = GroupInfo(rank=0, size=2, group=None, gloo_group=inner_gloo) + outer = GroupInfo(rank=0, size=3, group=None, gloo_group=outer_gloo) + result = GroupsInfo.from_pair(inner=inner, outer=outer) + assert result.gloo_groups_inner_to_outer == [inner_gloo, outer_gloo] + + def test_from_pair_rank_zero_only_when_both_zero(self) -> None: + result = GroupsInfo.from_pair( + inner=GroupInfo(rank=0, size=2, group=None), + outer=GroupInfo(rank=0, size=3, group=None), + ) + assert result.rank == 0 + assert result.size == 6 + + def test_from_pair_rank_nonzero_when_inner_nonzero(self) -> None: + result = GroupsInfo.from_pair( + inner=GroupInfo(rank=1, size=2, group=None), + outer=GroupInfo(rank=0, size=3, group=None), + ) + assert result.rank == 1 + + def test_from_pair_rank_nonzero_when_outer_nonzero(self) -> None: + result = GroupsInfo.from_pair( + inner=GroupInfo(rank=0, size=2, group=None), + outer=GroupInfo(rank=1, size=3, group=None), + ) + assert result.rank == 2 + + +# -- Parameterized GeneralPGUtil tests (native vs torchft code paths) -- + + +def _worker_pg_util_ops(rank: int, world_size: int, port: int) -> None: + """Test GeneralPGUtil operations with native gloo groups.""" + init_gloo(rank, world_size, port=port) + try: + group = dist.new_group(ranks=list(range(world_size)), backend="gloo") + + for util_cls in [_NativePGUtil]: + util = util_cls() + + # get_rank / get_size + assert util.get_rank(group) == rank + assert util.get_size(group) == world_size + + # all_reduce SUM + tensor = torch.tensor([float(rank + 1)]) + util.all_reduce(tensor, group, op=dist.ReduceOp.SUM) + assert tensor.item() == 1.0 + 2.0 + 3.0 + 4.0 + + # reduce to root + tensor = torch.tensor([float(rank + 1)]) + util.reduce(tensor, group, op=dist.ReduceOp.SUM) + if rank == 0: + assert tensor.item() == 1.0 + 2.0 + 3.0 + 4.0 + + # broadcast from root + tensor = torch.tensor([99.0]) if rank == 0 else torch.tensor([0.0]) + util.broadcast(tensor, group) + assert tensor.item() == 99.0 + + # barrier (all ranks must reach it; returns nothing) + util.barrier(group) + + # all_gather + input_t = torch.tensor([float(rank)]) + output_t = [torch.zeros(1) for _ in range(world_size)] + util.all_gather(output_t, input_t, group=group) + assert [t.item() for t in output_t] == [0.0, 1.0, 2.0, 3.0] + + # gather + input_t = torch.tensor([float(rank)]) + if rank == 0: + gather_list = [torch.zeros(1) for _ in range(world_size)] + util.gather(input_t, gather_list=gather_list, group=group) + assert [t.item() for t in gather_list] == [0.0, 1.0, 2.0, 3.0] + else: + util.gather(input_t, gather_list=None, group=group) + + # GroupInfo verification + GroupInfo(rank=rank, size=world_size, group=group) + wrong_rank = (rank + 1) % world_size + with pytest.raises(AssertionError): + GroupInfo(rank=wrong_rank, size=world_size, group=group) + finally: + dist.destroy_process_group() + + +def test_pg_util_ops() -> None: + run_multiprocess(_worker_pg_util_ops, world_size=4) + + +def _worker_gather_object(rank: int, world_size: int, port: int) -> None: + """Verify _NativePGUtil gather_object returns correct results on rank 0.""" + init_gloo(rank, world_size, port=port) + try: + group = dist.new_group(ranks=list(range(world_size)), backend="gloo") + + test_objects = [ + {"rank": rank, "value": rank * 10}, + [rank, rank + 1, "hello"], + f"string_from_rank_{rank}", + (rank, {"nested": True}), + ] + + util = _NativePGUtil() + for obj in test_objects: + if rank == 0: + result: list[Any] = [None] * world_size + util.gather_object(obj, result, group=group) + assert all(r is not None for r in result), f"Incomplete gather for obj type={type(obj)}" + else: + util.gather_object(obj, None, group=group) + finally: + dist.destroy_process_group() + + +def test_gather_object() -> None: + run_multiprocess(_worker_gather_object, world_size=4) + + +class TestRawPGUtilUnit: + """Unit tests for _RawPGUtil using mock groups (torchft-style with _rank attr).""" + + def test_get_rank_returns_group_rank(self) -> None: + group = MagicMock() + group._rank = 3 + assert _RawPGUtil().get_rank(group) == 3 + + def test_get_size_returns_group_size(self) -> None: + group = MagicMock() + group.size.return_value = 8 + assert _RawPGUtil().get_size(group) == 8 + + @patch("miles.utils.process_group_utils.dist.AllreduceOptions", MagicMock) + def test_all_reduce_calls_group_allreduce(self) -> None: + group = MagicMock() + work = MagicMock() + work.wait.return_value = True + group.allreduce.return_value = work + tensor = torch.tensor([1.0]) + _RawPGUtil().all_reduce(tensor, group, op=dist.ReduceOp.SUM) + group.allreduce.assert_called_once() + + @patch("miles.utils.process_group_utils.dist.AllreduceOptions", MagicMock) + def test_reduce_falls_back_to_group_allreduce(self) -> None: + # _RawPGUtil.reduce intentionally redirects to all_reduce because torchft's + # ProcessGroupWrapper doesn't override reduce() — see implementation comment. + group = MagicMock() + work = MagicMock() + work.wait.return_value = True + group.allreduce.return_value = work + tensor = torch.tensor([1.0]) + _RawPGUtil().reduce(tensor, group, op=dist.ReduceOp.SUM) + group.allreduce.assert_called_once() + group.reduce.assert_not_called() + + @patch("miles.utils.process_group_utils.dist.BroadcastOptions", MagicMock) + def test_broadcast_calls_group_broadcast(self) -> None: + group = MagicMock() + work = MagicMock() + work.wait.return_value = True + group.broadcast.return_value = work + tensor = torch.tensor([1.0]) + _RawPGUtil().broadcast(tensor, group) + group.broadcast.assert_called_once() + + def test_barrier_waits_on_the_work(self) -> None: + # Regression: a fire-and-forget group.barrier() (no .wait()) escapes torchft's + # per-collective timeout, letting a dead peer hang until the NCCL watchdog kills + # the process. _RawPGUtil.barrier MUST wait on the returned work. + group = MagicMock() + work = MagicMock() + work.wait.return_value = True + group.barrier.return_value = work + _RawPGUtil().barrier(group) + group.barrier.assert_called_once() + work.wait.assert_called_once() + + def test_barrier_raises_when_wait_returns_false(self) -> None: + # A failed barrier must raise so callers (e.g. the dumper) catch it and go + # degraded, instead of silently reporting success on a fire-and-forget call. + group = MagicMock() + work = MagicMock() + work.wait.return_value = False + group.barrier.return_value = work + with pytest.raises(RuntimeError, match="distributed operation barrier failed"): + _RawPGUtil().barrier(group) + + +# -- MultiPGUtil tests -- + + +def _worker_multi_pg_util_all_reduce(rank: int, world_size: int, port: int) -> None: + init_gloo(rank, world_size, port=port) + try: + mesh = _make_mesh() + inner_group = mesh.get_group("inner") + outer_group = mesh.get_group("outer") + + # Step 1: single group + tensor = torch.tensor([float(rank + 1)]) + MultiPGUtil.all_reduce(tensor, [inner_group], op=dist.ReduceOp.SUM) + expected = {0: 3.0, 1: 3.0, 2: 7.0, 3: 7.0}[rank] + assert tensor.item() == expected, f"rank {rank}: expected {expected}, got {tensor.item()}" + + # Step 2: two groups = global sum + tensor = torch.tensor([float(rank + 1)]) + MultiPGUtil.all_reduce(tensor, [inner_group, outer_group], op=dist.ReduceOp.SUM) + assert tensor.item() == 1.0 + 2.0 + 3.0 + 4.0 + + # Step 3: bitwise equality across all ranks + tensor = torch.tensor([float(rank + 1) * 0.1]) + MultiPGUtil.all_reduce(tensor, [inner_group, outer_group], op=dist.ReduceOp.SUM) + result_bytes = tensor.numpy().tobytes() + gathered_bytes = [None] * world_size + dist.all_gather_object(gathered_bytes, result_bytes) + assert all(b == gathered_bytes[0] for b in gathered_bytes), "Not bitwise equal" + + # Step 4: empty groups = no-op + tensor = torch.tensor([42.0]) + MultiPGUtil.all_reduce(tensor, [], op=dist.ReduceOp.SUM) + assert tensor.item() == 42.0 + + # Step 5: MAX op + tensor = torch.tensor([float(rank + 1)]) + MultiPGUtil.all_reduce(tensor, [inner_group, outer_group], op=dist.ReduceOp.MAX) + assert tensor.item() == 4.0 + finally: + dist.destroy_process_group() + + +def test_multi_pg_util_all_reduce() -> None: + run_multiprocess(_worker_multi_pg_util_all_reduce, world_size=4) + + +def _worker_multi_pg_util_gather_object(rank: int, world_size: int, port: int) -> None: + init_gloo(rank, world_size, port=port) + try: + mesh = _make_mesh() + inner_group = mesh.get_group("inner") + outer_group = mesh.get_group("outer") + + # Step 1: single group gather + result = MultiPGUtil.gather_object({"rank": rank}, [inner_group]) + inner_rank = rank % 2 + if inner_rank == 0: + assert result is not None + assert len(result) == 2 + ranks_gathered = {item["rank"] for item in result} + if rank == 0: + assert ranks_gathered == {0, 1} + else: + assert ranks_gathered == {2, 3} + else: + assert result is None + + # Step 2: two group gather — global rank 0 gets everything + result = MultiPGUtil.gather_object({"rank": rank}, [inner_group, outer_group]) + if rank == 0: + assert result is not None + assert len(result) == 4 + assert {item["rank"] for item in result} == {0, 1, 2, 3} + else: + assert result is None + finally: + dist.destroy_process_group() + + +def test_multi_pg_util_gather_object() -> None: + run_multiprocess(_worker_multi_pg_util_gather_object, world_size=4) + + +# -- _check_wait tests -- + + +class TestCheckWait: + def test_raises_on_false(self) -> None: + work = MagicMock() + work.wait.return_value = False + + with pytest.raises(RuntimeError, match="distributed operation allreduce failed"): + _check_wait(work, "allreduce") + + def test_passes_on_true(self) -> None: + work = MagicMock() + work.wait.return_value = True + + _check_wait(work, "allreduce") + + def test_propagates_exception_from_wait(self) -> None: + work = MagicMock() + work.wait.side_effect = RuntimeError("NCCL timeout") + + with pytest.raises(RuntimeError, match="NCCL timeout"): + _check_wait(work, "allreduce") + + +def _worker_bool_and_all_true(rank: int, world_size: int, port: int) -> None: + _worker_bool_and(rank, world_size, port, value_by_rank={0: True, 1: True}, expected=True) + + +def _worker_bool_and_all_false(rank: int, world_size: int, port: int) -> None: + _worker_bool_and(rank, world_size, port, value_by_rank={0: False, 1: False}, expected=False) + + +def _worker_bool_and_mixed(rank: int, world_size: int, port: int) -> None: + _worker_bool_and(rank, world_size, port, value_by_rank={0: True, 1: False}, expected=False) + + +def _worker_bool_and(rank: int, world_size: int, port: int, *, value_by_rank: dict[int, bool], expected: bool) -> None: + init_gloo(rank, world_size, port=port) + try: + group = dist.new_group(ranks=list(range(world_size)), backend="gloo") + result = collective_bool_and(value=value_by_rank[rank], group=group) + assert result is expected, f"rank {rank}: expected {expected}, got {result}" + finally: + dist.destroy_process_group() + + +class TestCollectiveBoolAnd: + def test_all_true(self) -> None: + run_multiprocess(_worker_bool_and_all_true) + + def test_all_false(self) -> None: + run_multiprocess(_worker_bool_and_all_false) + + def test_mixed_returns_false(self) -> None: + run_multiprocess(_worker_bool_and_mixed) From 69c50b9f2e11897581b0ef6329ffa4613da5825d Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Mon, 22 Jun 2026 18:15:17 +0800 Subject: [PATCH 30/41] Add _TensorViewCodec for storage-deduplicated tensor serialization Introduce `_TensorViewCodec` (in `checkpoint_transfer.py`), which encodes a list of tensors as (unique_storages, view_metas) by deduping shared underlying storages (e.g. Megatron distributed-optimizer grad buckets) and reconstructs the original views via `as_strided`. Comprehensively unit-tested by `TestTensorViewCodec`. Consumed by the peer checkpoint transfer added next. --- .../megatron_utils/checkpoint_transfer.py | 52 +++ .../test_checkpoint_transfer.py | 370 ++++++++++++++++++ 2 files changed, 422 insertions(+) create mode 100644 miles/backends/megatron_utils/checkpoint_transfer.py create mode 100644 tests/fast/backends/megatron_utils/test_checkpoint_transfer.py diff --git a/miles/backends/megatron_utils/checkpoint_transfer.py b/miles/backends/megatron_utils/checkpoint_transfer.py new file mode 100644 index 00000000000..ea7a92a8c34 --- /dev/null +++ b/miles/backends/megatron_utils/checkpoint_transfer.py @@ -0,0 +1,52 @@ +import torch + + +class _TensorViewCodec: + """Encode tensors as (unique_storages, view_metas) and decode back. + + Many input tensors may share underlying storage (e.g. Megatron + distributed-optimizer grad buckets). `encode` dedups by storage data_ptr — + each unique storage is wrapped once as a uint8 tensor (no copy), plus a + per-input view_meta record (storage_id, dtype, shape, stride, + storage_offset). `decode` reconstructs the original views with + `as_strided` over the storage bytes reinterpreted at the original dtype. + """ + + @staticmethod + def encode(tensors: list[torch.Tensor]) -> tuple[list[torch.Tensor], list[dict]]: + storage_id_by_key: dict[tuple[torch.device, int], int] = {} + unique_storages: list[torch.Tensor] = [] + view_metas: list[dict] = [] + for t in tensors: + storage = t.untyped_storage() + key = (t.device, storage.data_ptr()) + if key not in storage_id_by_key: + storage_id_by_key[key] = len(unique_storages) + # Wrap full storage as uint8 tensor (no copy, shares memory). + unique_storages.append(torch.tensor(storage, dtype=torch.uint8, device=t.device)) + view_metas.append( + { + "storage_id": storage_id_by_key[key], + "dtype": t.dtype, + "shape": tuple(t.shape), + "stride": tuple(t.stride()), + "storage_offset": t.storage_offset(), + } + ) + return unique_storages, view_metas + + @staticmethod + def decode(unique_storages: list[torch.Tensor], view_metas: list[dict]) -> list[torch.Tensor]: + tensors: list[torch.Tensor] = [] + for vm in view_metas: + storage_t = unique_storages[vm["storage_id"]] # uint8 view of received storage + # Reinterpret bytes as the original dtype, then apply stride/offset. + dtype_view = storage_t.view(vm["dtype"]) + view = torch.as_strided( + dtype_view, + size=vm["shape"], + stride=vm["stride"], + storage_offset=vm["storage_offset"], + ) + tensors.append(view) + return tensors diff --git a/tests/fast/backends/megatron_utils/test_checkpoint_transfer.py b/tests/fast/backends/megatron_utils/test_checkpoint_transfer.py new file mode 100644 index 00000000000..c5b764b67f1 --- /dev/null +++ b/tests/fast/backends/megatron_utils/test_checkpoint_transfer.py @@ -0,0 +1,370 @@ +import pytest +import torch + +from miles.backends.megatron_utils.checkpoint_transfer import _TensorViewCodec + + +class TestTensorViewCodec: + """Comprehensive UT for `_TensorViewCodec.encode/decode`. + + Round-trip semantics: `decode(*encode(tensors))` must produce a list of + tensors that are *value-equal* to the inputs, while sharing storage with + the encoded `unique_storages` (the whole point of the codec). + """ + + def test_empty_input_yields_empty_output(self): + unique_storages, view_metas = _TensorViewCodec.encode([]) + assert unique_storages == [] + assert view_metas == [] + assert _TensorViewCodec.decode(unique_storages, view_metas) == [] + + def test_single_tensor_round_trip(self): + original = torch.arange(12, dtype=torch.float32).reshape(3, 4) + + unique_storages, view_metas = _TensorViewCodec.encode([original]) + + assert len(unique_storages) == 1 + assert len(view_metas) == 1 + assert view_metas[0]["storage_id"] == 0 + assert view_metas[0]["dtype"] == torch.float32 + assert view_metas[0]["shape"] == (3, 4) + assert view_metas[0]["stride"] == (4, 1) + assert view_metas[0]["storage_offset"] == 0 + + decoded = _TensorViewCodec.decode(unique_storages, view_metas) + assert len(decoded) == 1 + assert torch.equal(decoded[0], original) + assert decoded[0].dtype == original.dtype + assert decoded[0].shape == original.shape + + def test_distinct_storages_yield_distinct_storage_ids(self): + a = torch.arange(8, dtype=torch.float32) + b = torch.arange(4, dtype=torch.int64) + c = torch.zeros(5, dtype=torch.float64) + + unique_storages, view_metas = _TensorViewCodec.encode([a, b, c]) + + assert len(unique_storages) == 3 + assert [vm["storage_id"] for vm in view_metas] == [0, 1, 2] + + decoded = _TensorViewCodec.decode(unique_storages, view_metas) + assert torch.equal(decoded[0], a) + assert torch.equal(decoded[1], b) + assert torch.equal(decoded[2], c) + + def test_shared_storage_dedups_into_one_unique_storage(self): + """The dedup invariant: N views over one storage produce 1 unique_storage.""" + base = torch.arange(100, dtype=torch.float32) + view_a = base[10:20] # offset=10, length=10 + view_b = base[20:60].view(4, 10) # offset=20, shape=(4, 10) + view_c = base[:5] # offset=0, length=5 + + unique_storages, view_metas = _TensorViewCodec.encode([view_a, view_b, view_c]) + + assert len(unique_storages) == 1 + assert all(vm["storage_id"] == 0 for vm in view_metas) + # storage_offset / shape / stride encode the difference + assert view_metas[0]["storage_offset"] == 10 + assert view_metas[0]["shape"] == (10,) + assert view_metas[1]["storage_offset"] == 20 + assert view_metas[1]["shape"] == (4, 10) + assert view_metas[2]["storage_offset"] == 0 + assert view_metas[2]["shape"] == (5,) + + decoded = _TensorViewCodec.decode(unique_storages, view_metas) + assert torch.equal(decoded[0], view_a) + assert torch.equal(decoded[1], view_b) + assert torch.equal(decoded[2], view_c) + + def test_partial_dedup_preserves_storage_id_assignment_order(self): + """First-seen wins: storage_id increments only on a NEW storage_ptr.""" + s1 = torch.arange(10, dtype=torch.float32) # storage A + s2 = torch.arange(20, dtype=torch.int32) # storage B + s1_view = s1[2:8] # storage A again + s3 = torch.zeros(4, dtype=torch.float64) # storage C + s2_view = s2[5:10] # storage B again + + unique_storages, view_metas = _TensorViewCodec.encode([s1, s2, s1_view, s3, s2_view]) + + assert len(unique_storages) == 3 + assert [vm["storage_id"] for vm in view_metas] == [0, 1, 0, 2, 1] + + def test_dtype_preserved_across_round_trip(self): + dtypes = [ + torch.float32, + torch.float64, + torch.float16, + torch.int8, + torch.int16, + torch.int32, + torch.int64, + torch.uint8, + torch.bool, + torch.bfloat16, + ] + tensors = [torch.zeros(5, dtype=dt) for dt in dtypes] + + unique_storages, view_metas = _TensorViewCodec.encode(tensors) + decoded = _TensorViewCodec.decode(unique_storages, view_metas) + + for d, original in zip(decoded, tensors, strict=True): + assert d.dtype == original.dtype + assert torch.equal(d, original) + + def test_non_contiguous_tensor_preserves_stride_and_values(self): + """Transposed view: shape=(3,4), stride=(1,3) (column-major over storage).""" + base = torch.arange(12, dtype=torch.float32).reshape(4, 3) + transposed = base.t() # shape=(3,4), non-contiguous + assert not transposed.is_contiguous() + + unique_storages, view_metas = _TensorViewCodec.encode([transposed]) + + assert view_metas[0]["shape"] == (3, 4) + assert view_metas[0]["stride"] == transposed.stride() + + decoded = _TensorViewCodec.decode(unique_storages, view_metas) + assert decoded[0].stride() == transposed.stride() + assert torch.equal(decoded[0], transposed) + + def test_storage_offset_preserved_for_slice_view(self): + base = torch.arange(20, dtype=torch.float32) + sliced = base[7:15] + assert sliced.storage_offset() == 7 + + unique_storages, view_metas = _TensorViewCodec.encode([sliced]) + + assert view_metas[0]["storage_offset"] == 7 + + decoded = _TensorViewCodec.decode(unique_storages, view_metas) + assert decoded[0].storage_offset() == 7 + assert torch.equal(decoded[0], sliced) + + def test_decoded_view_shares_storage_with_unique_storage(self): + """Decoded views must alias the unique_storage (no copy).""" + original = torch.arange(8, dtype=torch.float32) + + unique_storages, view_metas = _TensorViewCodec.encode([original]) + decoded = _TensorViewCodec.decode(unique_storages, view_metas) + + # Mutating the unique_storage (uint8 view) must propagate to decoded view. + unique_storages[0].zero_() + assert torch.equal(decoded[0], torch.zeros(8, dtype=torch.float32)) + + def test_encode_aliases_input_storage_no_copy(self): + """Encoded uint8 storage must alias the input tensor's storage.""" + original = torch.arange(8, dtype=torch.float32) + + unique_storages, _ = _TensorViewCodec.encode([original]) + + # Mutating the input tensor must propagate to the encoded uint8 view. + original.zero_() + assert unique_storages[0].sum().item() == 0 + + def test_multi_dtype_views_into_same_storage(self): + """A storage can be viewed at multiple dtypes (e.g. fp32 vs int32).""" + base = torch.arange(8, dtype=torch.float32) + as_int = base.view(torch.int32) # same storage, different dtype + assert base.untyped_storage().data_ptr() == as_int.untyped_storage().data_ptr() + + unique_storages, view_metas = _TensorViewCodec.encode([base, as_int]) + + assert len(unique_storages) == 1 + assert view_metas[0]["dtype"] == torch.float32 + assert view_metas[1]["dtype"] == torch.int32 + + decoded = _TensorViewCodec.decode(unique_storages, view_metas) + assert torch.equal(decoded[0], base) + assert torch.equal(decoded[1], as_int) + + def test_round_trip_idempotent_under_repeated_encoding(self): + """encode(decode(encode(t))) == encode(t).""" + original = torch.arange(20, dtype=torch.float32).reshape(4, 5) + + s1, m1 = _TensorViewCodec.encode([original]) + decoded = _TensorViewCodec.decode(s1, m1) + s2, m2 = _TensorViewCodec.encode(decoded) + + assert len(s1) == len(s2) == 1 + assert m1 == m2 + + def test_round_trip_property_random_mix(self): + """Property: for a random mix of tensors, decode round-trip is value-preserving.""" + torch.manual_seed(0) + shared_base = torch.randn(64) + tensors = [ + torch.randn(3, 4), + torch.randn(7), + shared_base[10:30], + shared_base[5:15].view(2, 5), + torch.randint(0, 100, (4, 4), dtype=torch.int64), + torch.zeros(0, dtype=torch.float32), # empty tensor + ] + + unique_storages, view_metas = _TensorViewCodec.encode(tensors) + + assert len(view_metas) == len(tensors) + # shared_base appears as 1 unique storage even though 2 views reference it + assert len(unique_storages) == 5 + + decoded = _TensorViewCodec.decode(unique_storages, view_metas) + for d, original in zip(decoded, tensors, strict=True): + assert d.dtype == original.dtype + assert d.shape == original.shape + assert torch.equal(d, original) + + def test_decode_after_storage_clone_preserves_values(self): + """Simulates wire transfer: clone unique_storages (mimic NCCL recv copy) + before decoding. Decoded views must still produce correct values, and + be aliased to the clone (not the original).""" + original = torch.arange(16, dtype=torch.float32).reshape(4, 4) + sliced = original[1:3] + + unique_storages, view_metas = _TensorViewCodec.encode([original, sliced]) + cloned_storages = [u.clone() for u in unique_storages] + + decoded = _TensorViewCodec.decode(cloned_storages, view_metas) + assert torch.equal(decoded[0], original) + assert torch.equal(decoded[1], sliced) + + # Mutate the clone — decoded views follow the clone, not the original. + cloned_storages[0].zero_() + assert decoded[0].sum().item() == 0 + assert decoded[1].sum().item() == 0 + assert original.sum().item() != 0 # original untouched + + def test_same_tensor_twice_dedups_into_one_storage(self): + """encode([t, t]) — two references to the same tensor share one storage.""" + t = torch.arange(8, dtype=torch.float32) + + unique_storages, view_metas = _TensorViewCodec.encode([t, t]) + + assert len(unique_storages) == 1 + assert [vm["storage_id"] for vm in view_metas] == [0, 0] + assert view_metas[0] == view_metas[1] + + decoded = _TensorViewCodec.decode(unique_storages, view_metas) + assert torch.equal(decoded[0], t) + assert torch.equal(decoded[1], t) + + def test_zero_dim_scalar_round_trip(self): + """0-d (scalar) tensor: shape=(), stride=(), storage_offset=0.""" + scalar = torch.tensor(7.5, dtype=torch.float32) + assert scalar.shape == () + assert scalar.stride() == () + + unique_storages, view_metas = _TensorViewCodec.encode([scalar]) + + assert view_metas[0]["shape"] == () + assert view_metas[0]["stride"] == () + + decoded = _TensorViewCodec.decode(unique_storages, view_metas) + assert decoded[0].shape == () + assert torch.equal(decoded[0], scalar) + assert decoded[0].item() == 7.5 + + def test_nn_parameter_input(self): + """Production input via state_dict.pop_tensors() may include nn.Parameter. + Codec uses .untyped_storage() which works on Parameter same as Tensor.""" + param = torch.nn.Parameter(torch.arange(6, dtype=torch.float32).reshape(2, 3)) + + unique_storages, view_metas = _TensorViewCodec.encode([param]) + + assert len(unique_storages) == 1 + assert view_metas[0]["dtype"] == torch.float32 + assert view_metas[0]["shape"] == (2, 3) + + decoded = _TensorViewCodec.decode(unique_storages, view_metas) + assert torch.equal(decoded[0], param.data) + + def test_empty_slice_of_nonempty_storage(self): + """Zero-length view (shape=(0,)) over a non-empty storage.""" + base = torch.arange(20, dtype=torch.float32) + empty_view = base[5:5] + assert empty_view.shape == (0,) + # Empty view shares its base's storage (data_ptr is non-null). + assert empty_view.untyped_storage().data_ptr() == base.untyped_storage().data_ptr() + + unique_storages, view_metas = _TensorViewCodec.encode([base, empty_view]) + + assert len(unique_storages) == 1 + assert [vm["storage_id"] for vm in view_metas] == [0, 0] + assert view_metas[1]["shape"] == (0,) + assert view_metas[1]["storage_offset"] == 5 + + decoded = _TensorViewCodec.decode(unique_storages, view_metas) + assert decoded[0].shape == base.shape + assert decoded[1].shape == (0,) + assert torch.equal(decoded[0], base) + assert torch.equal(decoded[1], empty_view) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") + def test_cuda_round_trip(self): + """Production runs on CUDA — verify the codec round-trips on GPU tensors, + including shared-storage dedup and dtype preservation.""" + base = torch.arange(64, dtype=torch.float32, device="cuda") + view_a = base[10:30] # shares storage with base + view_b = base[30:50].view(4, 5) # shares storage with base + independent = torch.zeros(8, dtype=torch.bfloat16, device="cuda") + + unique_storages, view_metas = _TensorViewCodec.encode([base, view_a, view_b, independent]) + + assert len(unique_storages) == 2 + assert all(s.device.type == "cuda" for s in unique_storages) + assert [vm["storage_id"] for vm in view_metas] == [0, 0, 0, 1] + + decoded = _TensorViewCodec.decode(unique_storages, view_metas) + assert all(d.device.type == "cuda" for d in decoded) + assert torch.equal(decoded[0], base) + assert torch.equal(decoded[1], view_a) + assert torch.equal(decoded[2], view_b) + assert torch.equal(decoded[3], independent) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") + def test_cuda_storage_aliasing_no_copy(self): + """On CUDA, the encoded uint8 storage must alias the input tensor's + device memory (no host-device or device-device copy).""" + original = torch.arange(8, dtype=torch.float32, device="cuda") + + unique_storages, view_metas = _TensorViewCodec.encode([original]) + decoded = _TensorViewCodec.decode(unique_storages, view_metas) + + assert unique_storages[0].device.type == "cuda" + # Mutate input — should propagate through the encoded storage to decoded view. + original.zero_() + assert decoded[0].sum().item() == 0 + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") + def test_mixed_device_tensors_get_distinct_storages(self): + """CPU and CUDA tensors never share a storage entry, and each decoded view keeps its device.""" + cpu_t = torch.arange(8, dtype=torch.float32) + cuda_t = torch.arange(8, dtype=torch.float32, device="cuda") + + unique_storages, view_metas = _TensorViewCodec.encode([cpu_t, cuda_t]) + + assert len(unique_storages) == 2 + assert [vm["storage_id"] for vm in view_metas] == [0, 1] + assert unique_storages[0].device.type == "cpu" + assert unique_storages[1].device.type == "cuda" + + decoded = _TensorViewCodec.decode(unique_storages, view_metas) + assert decoded[0].device.type == "cpu" + assert decoded[1].device.type == "cuda" + assert torch.equal(decoded[0], cpu_t) + assert torch.equal(decoded[1], cuda_t) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") + def test_mixed_device_dedup_stays_per_device(self): + """Views sharing storage dedup within their device while the other device's tensor stays separate.""" + base = torch.arange(16, dtype=torch.float32, device="cuda") + view = base[4:12] + cpu_t = torch.arange(4, dtype=torch.float32) + + unique_storages, view_metas = _TensorViewCodec.encode([base, view, cpu_t]) + + assert len(unique_storages) == 2 + assert [vm["storage_id"] for vm in view_metas] == [0, 0, 1] + + decoded = _TensorViewCodec.decode(unique_storages, view_metas) + assert torch.equal(decoded[0], base) + assert torch.equal(decoded[1], view) + assert torch.equal(decoded[2], cpu_t) From 1cd5e8396652b265d7d0754bc30cc5704ed90c11 Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Mon, 22 Jun 2026 18:15:18 +0800 Subject: [PATCH 31/41] Add an in-memory (non-persistent) checkpoint manager Add an InMemoryCheckpointManager that keeps checkpoints in memory and thread a `checkpointing_context` / `non_persistent_ckpt` through the model save path so a non-persistent checkpoint can be requested. - miles/backends/megatron_utils/in_memory_checkpoint.py, model.py and tests. --- .../megatron_utils/in_memory_checkpoint.py | 66 +++++++++++++++++++ miles/backends/megatron_utils/model.py | 8 ++- .../test_in_memory_checkpoint.py | 61 +++++++++++++++++ 3 files changed, 134 insertions(+), 1 deletion(-) create mode 100644 miles/backends/megatron_utils/in_memory_checkpoint.py create mode 100644 tests/fast/backends/megatron_utils/test_in_memory_checkpoint.py diff --git a/miles/backends/megatron_utils/in_memory_checkpoint.py b/miles/backends/megatron_utils/in_memory_checkpoint.py new file mode 100644 index 00000000000..a4352c30695 --- /dev/null +++ b/miles/backends/megatron_utils/in_memory_checkpoint.py @@ -0,0 +1,66 @@ +import logging +from collections.abc import Sequence +from typing import Any + +from megatron.training.global_vars import get_args + +from miles.backends.megatron_utils.model import save + +logger = logging.getLogger(__name__) + + +class InMemoryCheckpointManager: + """ref: nvidia_resiliency_ext's LocalCheckpointManager.""" + + def __init__(self) -> None: + self.latest_iteration: int = -1 + self._state_dict: object = None + self.local_ckpt_dir: str = "" + + _assert_args_for_in_memory_checkpoint(get_args()) + + def save(self, state_dict: object, iteration: int, is_async: bool = False) -> None: + """Store state_dict object reference in memory.""" + assert not is_async + + assert self._state_dict is None + self._state_dict = state_dict + self.latest_iteration = iteration + + def find_latest(self) -> int: + return self.latest_iteration + + def load(self) -> tuple[object, str]: + # Idempotent: Megatron's load_checkpoint calls _load_base_checkpoint twice + # (once for format detection at line 1508, once for actual load at line 1712). + # We must NOT consume `_state_dict` on first call. + assert self.latest_iteration >= 0, "No in-memory checkpoint available" + assert self._state_dict is not None + return self._state_dict, f"in-memory-ckpt-iter-{self.latest_iteration}" + + +def save_to_memory( + iteration: int, + model: Sequence, + optimizer: object, + opt_param_scheduler: object, +) -> object: + """Save checkpoint to in-memory manager via model.save (with forward hook protection).""" + manager = InMemoryCheckpointManager() + save( + iteration=iteration, + model=model, + optimizer=optimizer, + opt_param_scheduler=opt_param_scheduler, + checkpointing_context={"local_checkpoint_manager": manager}, + non_persistent_ckpt=True, + ) + state_dict, _ = manager.load() + return state_dict + + +def _assert_args_for_in_memory_checkpoint(args: Any) -> None: + assert args.non_persistent_ckpt_type == "local", ( + f"Expected non_persistent_ckpt_type='local', " f"got {getattr(args, 'non_persistent_ckpt_type', None)!r}" + ) + assert args.non_persistent_local_ckpt_algo is not None, "args.non_persistent_local_ckpt_algo must be set" diff --git a/miles/backends/megatron_utils/model.py b/miles/backends/megatron_utils/model.py index 44fdaeedcd8..82f528ba7da 100644 --- a/miles/backends/megatron_utils/model.py +++ b/miles/backends/megatron_utils/model.py @@ -714,6 +714,8 @@ def save( model: Sequence[DDP], optimizer: MegatronOptimizer | None, opt_param_scheduler: OptimizerParamScheduler | None, + checkpointing_context: dict | None = None, + non_persistent_ckpt: bool = False, ) -> None: """Persist a training checkpoint safely with forward hooks disabled. @@ -722,6 +724,9 @@ def save( model (Sequence[DDP]): Sequence of DDP-wrapped model chunks. optimizer (MegatronOptimizer): Optimizer instance. opt_param_scheduler (OptimizerParamScheduler): LR/WD scheduler. + checkpointing_context (dict | None): Context passed to Megatron's save_checkpoint + (e.g. ``{'local_checkpoint_manager': manager}`` for in-memory checkpoints). + non_persistent_ckpt (bool): If True, save a non-persistent (in-memory) checkpoint. """ args = get_args() hashes = None @@ -739,9 +744,10 @@ def save( optimizer, opt_param_scheduler, num_floating_point_operations_so_far=0, - checkpointing_context=None, train_data_iterator=None, preprocess_common_state_dict_fn=None, + checkpointing_context=checkpointing_context, + non_persistent_ckpt=non_persistent_ckpt, ) if hashes is not None: diff --git a/tests/fast/backends/megatron_utils/test_in_memory_checkpoint.py b/tests/fast/backends/megatron_utils/test_in_memory_checkpoint.py new file mode 100644 index 00000000000..0e98388cb6c --- /dev/null +++ b/tests/fast/backends/megatron_utils/test_in_memory_checkpoint.py @@ -0,0 +1,61 @@ +from unittest.mock import patch + +import pytest + +from miles.backends.megatron_utils.in_memory_checkpoint import InMemoryCheckpointManager + + +@pytest.fixture() +def manager(): + with patch("miles.backends.megatron_utils.in_memory_checkpoint.get_args") as mock_get_args: + mock_args = mock_get_args.return_value + mock_args.non_persistent_ckpt_type = "local" + mock_args.non_persistent_local_ckpt_algo = "fully_parallel" + yield InMemoryCheckpointManager() + + +class TestInMemoryCheckpointManager: + def test_find_latest_returns_minus_one_initially(self, manager: InMemoryCheckpointManager): + assert manager.find_latest() == -1 + + def test_load_before_save_raises(self, manager: InMemoryCheckpointManager): + with pytest.raises(AssertionError, match="No in-memory checkpoint"): + manager.load() + + def test_save_then_load_returns_same_object(self, manager: InMemoryCheckpointManager): + sentinel = object() + manager.save(state_dict=sentinel, iteration=5) + + assert manager.find_latest() == 5 + + result, name = manager.load() + assert result is sentinel + assert "5" in name + + def test_load_is_idempotent_returns_same_state(self, manager: InMemoryCheckpointManager): + sentinel = object() + manager.save(state_dict=sentinel, iteration=1) + + first, _ = manager.load() + second, _ = manager.load() + assert first is sentinel + assert second is sentinel + + def test_save_twice_without_reset_raises(self, manager: InMemoryCheckpointManager): + manager.save(state_dict=object(), iteration=1) + + with pytest.raises(AssertionError): + manager.save(state_dict=object(), iteration=2) + + def test_save_after_load_still_raises_load_does_not_reset(self, manager: InMemoryCheckpointManager): + """`load()` is idempotent and does NOT clear state, so a second `save()` + on the same manager — even after `load()` — must still raise.""" + manager.save(state_dict=object(), iteration=1) + manager.load() + + with pytest.raises(AssertionError): + manager.save(state_dict=object(), iteration=2) + + def test_async_save_raises(self, manager: InMemoryCheckpointManager): + with pytest.raises(AssertionError): + manager.save(state_dict=object(), iteration=1, is_async=True) From f2a129e5a3fceb9dadac3a75dec89004d4f02d0c Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Mon, 22 Jun 2026 18:15:18 +0800 Subject: [PATCH 32/41] Add peer checkpoint transfer for healing Add peer-to-peer checkpoint transfer (built on the in-memory checkpoint manager) so a healed cell can receive weights from a surviving peer. - miles/backends/megatron_utils/checkpoint_transfer.py and tests. --- .../megatron_utils/checkpoint_transfer.py | 168 ++++++++++++++++++ .../test_checkpoint_transfer.py | 131 +++++++++++++- 2 files changed, 297 insertions(+), 2 deletions(-) diff --git a/miles/backends/megatron_utils/checkpoint_transfer.py b/miles/backends/megatron_utils/checkpoint_transfer.py index ea7a92a8c34..beaa2b74d0f 100644 --- a/miles/backends/megatron_utils/checkpoint_transfer.py +++ b/miles/backends/megatron_utils/checkpoint_transfer.py @@ -1,5 +1,163 @@ +import logging +from collections.abc import Sequence +from datetime import timedelta + import torch +try: + from torchft.checkpointing.pg_transport import PGTransport +except ImportError: + PGTransport = None + +from megatron.core.dist_checkpointing.tensor_aware_state_dict import MCoreTensorAwareStateDict + +from miles.backends.megatron_utils.in_memory_checkpoint import InMemoryCheckpointManager, save_to_memory +from miles.utils.process_group_utils import GroupInfo +from miles.utils.structured_log import log_structured + +logger = logging.getLogger(__name__) + +# Must accommodate receiver's model init time (can take minutes for large models) +_DEFAULT_TIMEOUT = timedelta(seconds=600) + + +def send_ckpt( + *, + indep_dp: GroupInfo, + model: Sequence, + optimizer: object, + opt_param_scheduler: object, + iteration: int, + dst_rank: int, + timeout: timedelta = _DEFAULT_TIMEOUT, +) -> None: + """Send in-memory checkpoint to a destination cell via torchft PGTransport. + + Args: + indep_dp: Independent DP group info (provides the torchft PG). + model: Megatron model chunks. + optimizer: Megatron optimizer. + opt_param_scheduler: LR scheduler. + iteration: Current training iteration / rollout_id. + dst_rank: Destination alive_rank in the indep_dp process group. + timeout: Timeout for the NCCL send operation. + """ + state_dict = save_to_memory( + iteration=iteration, + model=model, + optimizer=optimizer, + opt_param_scheduler=opt_param_scheduler, + ) + + payload = _TransportCodec.encode(state_dict=state_dict, iteration=iteration) + + transport = _create_transport(indep_dp, timeout) + log_structured( + logger.info, op="cross_cell", phase="start", kind="ckpt_send", iteration=iteration, to_alive_rank=dst_rank + ) + transport.send_checkpoint( + dst_ranks=[dst_rank], + step=0, + state_dict=payload, + timeout=timeout, + ) + transport.disallow_checkpoint() + log_structured( + logger.info, op="cross_cell", phase="end", kind="ckpt_send", iteration=iteration, to_alive_rank=dst_rank + ) + + +def recv_ckpt( + *, + indep_dp: GroupInfo, + src_rank: int, + timeout: timedelta = _DEFAULT_TIMEOUT, +) -> InMemoryCheckpointManager: + """Receive checkpoint from a healthy cell via torchft PGTransport. + + Returns an InMemoryCheckpointManager containing the received state_dict, + ready to be passed to initialize_model_and_optimizer. + + Args: + indep_dp: Independent DP group info (provides the torchft PG). + src_rank: Source alive_rank in the indep_dp process group. + timeout: Timeout for the NCCL recv operation. + + Returns: + InMemoryCheckpointManager with state_dict loaded, ready for + initialize_model_and_optimizer to consume. + """ + transport = _create_transport(indep_dp, timeout) + log_structured(logger.info, op="cross_cell", phase="start", kind="ckpt_recv", from_alive_rank=src_rank) + payload = transport.recv_checkpoint( + src_rank=src_rank, + metadata=transport.metadata(), + step=0, + timeout=timeout, + ) + + iteration, state_dict = _TransportCodec.decode(payload) + log_structured( + logger.info, op="cross_cell", phase="end", kind="ckpt_recv", iteration=iteration, from_alive_rank=src_rank + ) + + manager = InMemoryCheckpointManager() + manager.save(state_dict, iteration=iteration) + return manager + + +class _TransportCodec: + @staticmethod + def encode( + *, + state_dict: MCoreTensorAwareStateDict, + iteration: int, + ) -> dict[str, object]: + """Serialize for transport, deduping by underlying storage. + + `pop_tensors()` returns ShardedTensor `.data` views. Many of those views share + one big underlying storage (e.g. Megatron distributed-optimizer grad buckets). + `torchft.PGTransport._cast_tensor` casts each tensor to a uint8 view of its + FULL storage and sends that — so naively sending N views of one bucket sends + bucket_size * N bytes (we measured ~110x amplification: 12.5 GB real data + sent as 1387 GB). + + Fix: send each unique storage exactly once as `unique_storages`, plus per-view + `view_metas` (storage_id, dtype, shape, stride, storage_offset) so the + receiver can reconstruct the original views by `as_strided`. + """ + tensors: list[torch.Tensor] = state_dict.pop_tensors() + # PGTransport._cast_tensor uses `type(t) is torch.Tensor` (strict), + # which rejects torch.nn.Parameter. Detach into plain Tensors that share + # storage but pass the type check. + tensors = [t.detach() if type(t) is not torch.Tensor else t for t in tensors] + + unique_storages, view_metas = _TensorViewCodec.encode(tensors) + + return { + "unique_storages": unique_storages, + "view_metas": view_metas, + "hollow_state_dict": state_dict, + "iteration": iteration, + } + + @staticmethod + def decode( + payload: dict[str, object], + ) -> tuple[int, MCoreTensorAwareStateDict]: + """Reverse of `encode`: reconstruct per-tensor views from received + unique_storages using view_metas. + """ + iteration: int = payload["iteration"] + hollow_state_dict: MCoreTensorAwareStateDict = payload["hollow_state_dict"] + unique_storages: list[torch.Tensor] = payload["unique_storages"] + view_metas: list[dict] = payload["view_metas"] + + tensors = _TensorViewCodec.decode(unique_storages, view_metas) + + hollow_state_dict.insert_tensors(tensors) + return iteration, hollow_state_dict + class _TensorViewCodec: """Encode tensors as (unique_storages, view_metas) and decode back. @@ -50,3 +208,13 @@ def decode(unique_storages: list[torch.Tensor], view_metas: list[dict]) -> list[ ) tensors.append(view) return tensors + + +def _create_transport(indep_dp: GroupInfo, timeout: timedelta) -> PGTransport: + if PGTransport is None: + raise ImportError("torchft is required for checkpoint transfer but could not be imported.") + return PGTransport( + pg=indep_dp.group, + timeout=timeout, + device=torch.device("cuda"), + ) diff --git a/tests/fast/backends/megatron_utils/test_checkpoint_transfer.py b/tests/fast/backends/megatron_utils/test_checkpoint_transfer.py index c5b764b67f1..bfde1d1b61e 100644 --- a/tests/fast/backends/megatron_utils/test_checkpoint_transfer.py +++ b/tests/fast/backends/megatron_utils/test_checkpoint_transfer.py @@ -1,7 +1,134 @@ +import pickle + import pytest import torch - -from miles.backends.megatron_utils.checkpoint_transfer import _TensorViewCodec +from megatron.core.dist_checkpointing.mapping import ShardedTensor +from megatron.core.dist_checkpointing.tensor_aware_state_dict import MCoreTensorAwareStateDict +from torch.utils._pytree import tree_flatten_with_path, tree_unflatten + +from miles.backends.megatron_utils.checkpoint_transfer import _TensorViewCodec, _TransportCodec + + +@pytest.fixture() +def state_dict() -> MCoreTensorAwareStateDict: + sharded_state_dict: dict = { + "model": { + "layer1.weight": ShardedTensor.from_rank_offsets( + "layer1.weight", torch.arange(32, dtype=torch.float32).reshape(4, 8) + ), + "layer2.weight": ShardedTensor.from_rank_offsets("layer2.weight", torch.full((2, 6), fill_value=7.0)), + }, + "optimizer": { + "step": ShardedTensor.from_rank_offsets("step", torch.tensor([100], dtype=torch.int64)), + }, + } + common = {"iteration": 0, "args_repr": "dummy"} + return MCoreTensorAwareStateDict(common=common, sharded_state_dict=sharded_state_dict) + + +class TestSerializeForTransport: + def test_returns_separated_storages_iteration_and_hollow_shell(self, state_dict: MCoreTensorAwareStateDict): + original_tensors = [t.clone() for t in state_dict.tensors] + + payload = _TransportCodec.encode(state_dict=state_dict, iteration=42) + + assert payload["iteration"] == 42 + assert isinstance(payload["unique_storages"], list) + assert isinstance(payload["view_metas"], list) + assert len(payload["view_metas"]) == 3 + + # Round-trip the storages+metas back into views and compare. + decoded = _TensorViewCodec.decode(payload["unique_storages"], payload["view_metas"]) + for t, original in zip(decoded, original_tensors, strict=True): + assert torch.equal(t, original) + + assert payload["hollow_state_dict"] is state_dict + assert payload["hollow_state_dict"].is_hollow + + def test_pytree_flatten_yields_each_tensor_as_separate_leaf(self, state_dict: MCoreTensorAwareStateDict): + """The whole point of the fix: PGTransport's tree_flatten_with_path must see + each ShardedTensor.data as its own leaf, not buried inside a pickled blob.""" + payload = _TransportCodec.encode(state_dict=state_dict, iteration=42) + + leaves, _ = tree_flatten_with_path(payload) + tensor_leaves = [v for _, v in leaves if isinstance(v, torch.Tensor)] + non_tensor_leaves = [v for _, v in leaves if not isinstance(v, torch.Tensor)] + + assert len(tensor_leaves) == 3 + assert any(isinstance(v, MCoreTensorAwareStateDict) for v in non_tensor_leaves) + assert 42 in non_tensor_leaves + + def test_hollow_shell_pickles_without_dragging_tensor_data(self, state_dict: MCoreTensorAwareStateDict): + """PGTransport pickles non-tensor leaves; the hollow shell must survive + a pickle round-trip and not contain any of the original tensor storage.""" + payload = _TransportCodec.encode(state_dict=state_dict, iteration=42) + + restored = pickle.loads(pickle.dumps(payload["hollow_state_dict"])) + + assert restored.is_hollow + sharded_tensors = list(restored._sharded_tensors) + assert len(sharded_tensors) == 3 + assert all(sh.data is None for sh in sharded_tensors) + assert all(hasattr(sh, "orig_device") for sh in sharded_tensors) + + +class TestDeserializeFromTransport: + def test_round_trip_preserves_tensor_values_iteration_and_common(self, state_dict: MCoreTensorAwareStateDict): + original_tensors = [t.clone() for t in state_dict.tensors] + original_common = dict(state_dict.common) + + payload = _TransportCodec.encode(state_dict=state_dict, iteration=42) + iteration_back, state_dict_back = _TransportCodec.decode(payload) + + assert iteration_back == 42 + assert not state_dict_back.is_hollow + assert state_dict_back.common == original_common + for original, back in zip(original_tensors, state_dict_back.tensors, strict=True): + assert torch.equal(original, back) + + def test_full_pgtransport_simulation_round_trip(self, state_dict: MCoreTensorAwareStateDict): + """End-to-end simulation of PGTransport: pytree flatten on sender, pickle + non-tensor leaves + treespec, clone tensor leaves to mimic NCCL transfer, + unflatten on receiver. Verifies our (de)serializers survive the actual + wire protocol — not just an in-process pop/insert.""" + original_tensors = [t.clone() for t in state_dict.tensors] + original_common = dict(state_dict.common) + + # Step 1: sender — wrap state_dict for transport + payload = _TransportCodec.encode(state_dict=state_dict, iteration=42) + + # Step 2: sender — flatten via pytree (what PGTransport does internally) + leaves, treespec = tree_flatten_with_path(payload) + + # Step 3: sender — pickle treespec + non-tensor leaves; "send" tensor leaves over the wire + is_tensor_mask: list[bool] = [isinstance(v, torch.Tensor) for _, v in leaves] + pickled_metadata = pickle.dumps( + (treespec, [v for v, m in zip([v for _, v in leaves], is_tensor_mask, strict=True) if not m]) + ) + wire_tensors = [v.clone() for v, m in zip([v for _, v in leaves], is_tensor_mask, strict=True) if m] + + # Step 4: receiver — unpickle metadata + interleave received tensors back into leaf order + treespec_recv, non_tensor_values = pickle.loads(pickled_metadata) + recv_values: list = [] + ti = 0 + nti = 0 + for is_tensor in is_tensor_mask: + if is_tensor: + recv_values.append(wire_tensors[ti]) + ti += 1 + else: + recv_values.append(non_tensor_values[nti]) + nti += 1 + payload_recv = tree_unflatten(recv_values, treespec_recv) + + # Step 5: receiver — unwrap into iteration + state_dict + iteration_back, state_dict_back = _TransportCodec.decode(payload_recv) + + assert iteration_back == 42 + assert not state_dict_back.is_hollow + assert state_dict_back.common == original_common + for original, back in zip(original_tensors, state_dict_back.tensors, strict=True): + assert torch.equal(original, back) class TestTensorViewCodec: From 32a4130222fb0129421af783acd3baa8baf189e2 Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Mon, 22 Jun 2026 18:15:18 +0800 Subject: [PATCH 33/41] Extract actor construction into a shared allocate_gpus_for_actor factory Move the inline env-var / backend-selection / ray.remote actor construction out of RayTrainGroup into a module-level allocate_gpus_for_actor factory, with no behaviour change. (Pure mechanical extraction; FT context is threaded in a follow-up.) - miles/ray/train/actor_factory.py, miles/ray/actor_group.py. --- miles/ray/actor_group.py | 76 +++--------------------------- miles/ray/train/__init__.py | 0 miles/ray/train/actor_factory.py | 80 ++++++++++++++++++++++++++++++++ 3 files changed, 87 insertions(+), 69 deletions(-) create mode 100644 miles/ray/train/__init__.py create mode 100644 miles/ray/train/actor_factory.py diff --git a/miles/ray/actor_group.py b/miles/ray/actor_group.py index bfbfdf55e85..ec8bdf237f1 100644 --- a/miles/ray/actor_group.py +++ b/miles/ray/actor_group.py @@ -1,11 +1,8 @@ import asyncio -import os -import ray from ray.util.placement_group import PlacementGroup -from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy -from miles.ray.utils import NOSET_VISIBLE_DEVICES_ENV_VARS_LIST +from miles.ray.train.actor_factory import allocate_gpus_for_actor class RayTrainGroup: @@ -45,71 +42,12 @@ def __init__( self._actor_handles = self._allocate_gpus_for_actor(pg, num_gpus_per_actor) def _allocate_gpus_for_actor(self, pg, num_gpus_per_actor): - world_size = self._num_nodes * self._num_gpus_per_node - - # Use placement group to lock resources for models of same type - assert pg is not None - pg, reordered_bundle_indices, _reordered_gpu_ids = pg - - env_vars = { - # because sglang will always set NCCL_CUMEM_ENABLE to 0 - # we need also set it to 0 to prevent nccl error. - "NCCL_CUMEM_ENABLE": os.environ.get("NCCL_CUMEM_ENABLE", "0"), - "NVTE_FP8_BLOCK_SCALING_FP32_SCALES": "1", - # DeepEP/NVSHMEM's internal NCCL conflicts with our NCCL and hangs under CUDA graphs. - "NVSHMEM_DISABLE_NCCL": os.environ.get("NVSHMEM_DISABLE_NCCL", "1"), - **{name: "1" for name in NOSET_VISIBLE_DEVICES_ENV_VARS_LIST}, - **self.args.train_env_vars, - } - - if source_patcher_config := self.args.dumper_source_patcher_config_train: - env_vars["DUMPER_SOURCE_PATCHER_CONFIG"] = source_patcher_config - - if self.args.offload_train and self.args.train_backend == "megatron": - import torch_memory_saver - - dynlib_path = os.path.join( - os.path.dirname(os.path.dirname(torch_memory_saver.__file__)), - "torch_memory_saver_hook_mode_preload.abi3.so", - ) - assert os.path.exists(dynlib_path), f"LD_PRELOAD so file {dynlib_path} does not exist." - - env_vars["LD_PRELOAD"] = dynlib_path - env_vars["TMS_INIT_ENABLE"] = "1" - env_vars["TMS_INIT_ENABLE_CPU_BACKUP"] = "1" - - backend = self.args.train_backend - if backend == "megatron": - from miles.backends.megatron_utils.actor import MegatronTrainRayActor - - actor_impl = MegatronTrainRayActor - - else: - from miles.backends.experimental.fsdp_utils import FSDPTrainRayActor - - actor_impl = FSDPTrainRayActor - - TrainRayActor = ray.remote( - num_gpus=1, runtime_env={"env_vars": env_vars}, concurrency_groups={"fault_injector": 1} - )(actor_impl) - - # Create worker actors - actor_handles = [] - master_addr, master_port = None, None - for rank in range(world_size): - actor = TrainRayActor.options( - num_cpus=num_gpus_per_actor, - num_gpus=num_gpus_per_actor, - scheduling_strategy=PlacementGroupSchedulingStrategy( - placement_group=pg, - placement_group_bundle_index=reordered_bundle_indices[rank], - ), - ).remote(world_size, rank, master_addr, master_port) - if rank == 0: - master_addr, master_port = ray.get(actor.get_master_addr_and_port.remote()) - actor_handles.append(actor) - - return actor_handles + 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, + ) async def init(self): """ diff --git a/miles/ray/train/__init__.py b/miles/ray/train/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/miles/ray/train/actor_factory.py b/miles/ray/train/actor_factory.py new file mode 100644 index 00000000000..084cc2ad06c --- /dev/null +++ b/miles/ray/train/actor_factory.py @@ -0,0 +1,80 @@ +import os + +import ray +from ray.util.placement_group import PlacementGroup +from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy + +from miles.ray.utils import NOSET_VISIBLE_DEVICES_ENV_VARS_LIST + + +def allocate_gpus_for_actor( + args, + gpus_per_cell: int, + pg: tuple[PlacementGroup, list[int], list[int]], + num_gpus_per_actor: float, +): + world_size = gpus_per_cell + + # Use placement group to lock resources for models of same type + assert pg is not None + pg, reordered_bundle_indices, _reordered_gpu_ids = pg + + env_vars = { + # because sglang will always set NCCL_CUMEM_ENABLE to 0 + # we need also set it to 0 to prevent nccl error. + "NCCL_CUMEM_ENABLE": os.environ.get("NCCL_CUMEM_ENABLE", "0"), + "NVTE_FP8_BLOCK_SCALING_FP32_SCALES": "1", + # DeepEP/NVSHMEM's internal NCCL conflicts with our NCCL and hangs under CUDA graphs. + "NVSHMEM_DISABLE_NCCL": os.environ.get("NVSHMEM_DISABLE_NCCL", "1"), + **{name: "1" for name in NOSET_VISIBLE_DEVICES_ENV_VARS_LIST}, + **args.train_env_vars, + } + + if source_patcher_config := args.dumper_source_patcher_config_train: + env_vars["DUMPER_SOURCE_PATCHER_CONFIG"] = source_patcher_config + + if args.offload_train and args.train_backend == "megatron": + import torch_memory_saver + + dynlib_path = os.path.join( + os.path.dirname(os.path.dirname(torch_memory_saver.__file__)), + "torch_memory_saver_hook_mode_preload.abi3.so", + ) + assert os.path.exists(dynlib_path), f"LD_PRELOAD so file {dynlib_path} does not exist." + + env_vars["LD_PRELOAD"] = dynlib_path + env_vars["TMS_INIT_ENABLE"] = "1" + env_vars["TMS_INIT_ENABLE_CPU_BACKUP"] = "1" + + backend = args.train_backend + if backend == "megatron": + from miles.backends.megatron_utils.actor import MegatronTrainRayActor + + actor_impl = MegatronTrainRayActor + + else: + from miles.backends.experimental.fsdp_utils import FSDPTrainRayActor + + actor_impl = FSDPTrainRayActor + + TrainRayActor = ray.remote( + num_gpus=1, runtime_env={"env_vars": env_vars}, concurrency_groups={"fault_injector": 1} + )(actor_impl) + + # Create worker actors + actor_handles = [] + master_addr, master_port = None, None + for rank in range(world_size): + actor = TrainRayActor.options( + num_cpus=num_gpus_per_actor, + num_gpus=num_gpus_per_actor, + scheduling_strategy=PlacementGroupSchedulingStrategy( + placement_group=pg, + placement_group_bundle_index=reordered_bundle_indices[rank], + ), + ).remote(world_size, rank, master_addr, master_port) + if rank == 0: + master_addr, master_port = ray.get(actor.get_master_addr_and_port.remote()) + actor_handles.append(actor) + + return actor_handles From f844cf3bc08f253f604e825ba69ca5946d07b960 Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Mon, 22 Jun 2026 18:15:18 +0800 Subject: [PATCH 34/41] Thread independent-DP / role / cell-index / rollout-manager context through the train group Add the FT context to the actor factory and train actors: indep_dp_store_addr / role / cell_index params and per-process-identity logging in TrainRayActor; pass a rollout_manager into RayTrainGroup (set_rollout_manager, pause rollout health monitors before weight updates); freeze the v1 group and wire context through placement_group / entrypoints. --- miles/ray/actor_group.py | 19 +++++++++++++++---- miles/ray/placement_group.py | 7 +++++-- miles/ray/rollout/rollout_manager.py | 6 +++--- miles/ray/train/actor_factory.py | 18 ++++++++++++++++-- miles/ray/train_actor.py | 22 ++++++++++++++++++++-- train.py | 2 +- train_async.py | 2 +- 7 files changed, 61 insertions(+), 15 deletions(-) diff --git a/miles/ray/actor_group.py b/miles/ray/actor_group.py index ec8bdf237f1..59cd4004d96 100644 --- a/miles/ray/actor_group.py +++ b/miles/ray/actor_group.py @@ -1,3 +1,8 @@ +# 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 @@ -26,6 +31,7 @@ def __init__( num_gpus_per_node, pg: tuple[PlacementGroup, list[int], list[int]], *, + rollout_manager: object | None, num_gpus_per_actor: float = 1, role: str, with_ref: bool, @@ -36,6 +42,7 @@ def __init__( self._num_gpus_per_node = num_gpus_per_node self.role = role self.with_ref = with_ref + self._rollout_manager = rollout_manager self.with_opd_teacher = with_opd_teacher # Allocate the GPUs for actors w/o instantiating them @@ -47,6 +54,9 @@ def _allocate_gpus_for_actor(self, pg, num_gpus_per_actor): 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): @@ -65,7 +75,7 @@ async def save_model(self, rollout_id, force_sync=False): """Save actor model""" await self._broadcast("save_model", rollout_id, force_sync=force_sync) - async def update_weights(self): + 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 @@ -74,6 +84,7 @@ async def update_weights(self): await self.rollout_manager.recover_updatable_engines.remote() info = await self.rollout_manager.get_updatable_engines_and_lock.remote() + await self.rollout_manager.health_monitoring_pause.remote() await self._broadcast("update_weights", info=info) @@ -93,9 +104,9 @@ async def connect(self, critic_group): ] await asyncio.gather(*refs) - async def set_rollout_manager(self, rollout_manager): - self.rollout_manager = rollout_manager - await self._broadcast("set_rollout_manager", rollout_manager) + async def set_rollout_manager(self): + self.rollout_manager = self._rollout_manager + await self._broadcast("set_rollout_manager", self._rollout_manager) async def _broadcast(self, method_name: str, *args, **kwargs) -> list: refs = [getattr(actor, method_name).remote(*args, **kwargs) for actor in self._actor_handles] diff --git a/miles/ray/placement_group.py b/miles/ray/placement_group.py index a86a4541cc3..3a4e3ab690d 100644 --- a/miles/ray/placement_group.py +++ b/miles/ray/placement_group.py @@ -123,7 +123,7 @@ def create_placement_groups(args): def allocate_train_group( - args, num_nodes, num_gpus_per_node, pg, role: str, with_ref: bool, with_opd_teacher: bool = False + args, num_nodes, num_gpus_per_node, pg, role: str, with_ref: bool, rollout_manager, with_opd_teacher: bool = False ): return RayTrainGroup( args=args, @@ -133,6 +133,7 @@ def allocate_train_group( num_gpus_per_actor=0.4, role=role, with_ref=with_ref, + rollout_manager=rollout_manager, with_opd_teacher=with_opd_teacher, ) @@ -145,6 +146,7 @@ async def create_training_models(args, pgs, rollout_manager): pg=pgs["actor"], role="actor", with_ref=args.kl_coef != 0 or args.use_kl_loss, + rollout_manager=rollout_manager, with_opd_teacher=args.use_opd and args.opd_type == "megatron", ) if args.use_critic: @@ -155,6 +157,7 @@ async def create_training_models(args, pgs, rollout_manager): pg=pgs["critic"], role="critic", with_ref=False, + rollout_manager=None, ) critic_init_task = await eager_create_task(critic_model.init()) else: @@ -170,7 +173,7 @@ async def create_training_models(args, pgs, rollout_manager): await critic_init_task await actor_model.connect(critic_model) - await actor_model.set_rollout_manager(rollout_manager) + await actor_model.set_rollout_manager() if args.rollout_global_dataset: await rollout_manager.load.remote(args.start_rollout_id - 1) diff --git a/miles/ray/rollout/rollout_manager.py b/miles/ray/rollout/rollout_manager.py index f853f99e61b..5cf62f10faa 100644 --- a/miles/ray/rollout/rollout_manager.py +++ b/miles/ray/rollout/rollout_manager.py @@ -180,7 +180,7 @@ def load(self, rollout_id=None): # TODO may parallelly execute offload/onload across services async def offload(self, tags: list[str] | None = None): - self._health_monitoring_pause() + self.health_monitoring_pause() for srv in self.servers.values(): await srv.offload(tags=tags) @@ -229,7 +229,7 @@ async def recover_updatable_engines(self) -> None: Recovers the updatable model (the one that receives weight updates from training). """ - self._health_monitoring_pause() + self.health_monitoring_pause() srv = self._get_updatable_server() if self.rollout_id == -1 or srv is None: return @@ -284,7 +284,7 @@ def set_train_parallel_config(self, config: dict): # -------------------------- utils ----------------------------- - def _health_monitoring_pause(self) -> None: + def health_monitoring_pause(self) -> None: for monitor in self._health_monitors: monitor.pause() diff --git a/miles/ray/train/actor_factory.py b/miles/ray/train/actor_factory.py index 084cc2ad06c..29afd2bda68 100644 --- a/miles/ray/train/actor_factory.py +++ b/miles/ray/train/actor_factory.py @@ -12,6 +12,9 @@ def allocate_gpus_for_actor( gpus_per_cell: int, pg: tuple[PlacementGroup, list[int], list[int]], num_gpus_per_actor: float, + indep_dp_store_addr: str, + role: str, + cell_index: int, ): world_size = gpus_per_cell @@ -58,7 +61,9 @@ def allocate_gpus_for_actor( actor_impl = FSDPTrainRayActor TrainRayActor = ray.remote( - num_gpus=1, runtime_env={"env_vars": env_vars}, concurrency_groups={"fault_injector": 1} + num_gpus=1, + runtime_env={"env_vars": env_vars}, + concurrency_groups={"heartbeat_status": 1, "default": 1, "fault_injector": 1}, )(actor_impl) # Create worker actors @@ -72,7 +77,16 @@ def allocate_gpus_for_actor( placement_group=pg, placement_group_bundle_index=reordered_bundle_indices[rank], ), - ).remote(world_size, rank, master_addr, master_port) + ).remote( + args, + world_size, + rank, + master_addr, + master_port, + indep_dp_store_addr=indep_dp_store_addr, + role=role, + cell_index=cell_index, + ) if rank == 0: master_addr, master_port = ray.get(actor.get_master_addr_and_port.remote()) actor_handles.append(actor) diff --git a/miles/ray/train_actor.py b/miles/ray/train_actor.py index f91d7681a78..1b5a7c148ed 100644 --- a/miles/ray/train_actor.py +++ b/miles/ray/train_actor.py @@ -3,7 +3,7 @@ import os import random from datetime import timedelta -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Literal import ray import torch @@ -14,7 +14,9 @@ from miles.utils.det_process_group import DET_NCCL_BACKEND_NAME, register_det_nccl_backend from miles.utils.distributed_utils import init_gloo_group from miles.utils.env_report import collect_and_print_node_env_report +from miles.utils.logging_utils import configure_logger from miles.utils.memory_utils import clear_memory, print_memory +from miles.utils.process_identity import TrainProcessIdentity from miles.utils.test_utils.fault_injector import inject_fault as _inject_fault if TYPE_CHECKING: @@ -33,9 +35,25 @@ def get_local_gpu_id(): class TrainRayActor(RayActor): - def __init__(self, world_size, rank, master_addr, master_port): + def __init__( + self, + args, + world_size: int, + rank: int, + master_addr, + master_port, + indep_dp_store_addr: str, + role: Literal["actor", "critic"], + cell_index: int, + ): + configure_logger( + args, source=TrainProcessIdentity(component=role, cell_index=cell_index, rank_within_cell=rank) + ) + self.args = args + self._world_size = world_size self._rank = rank + self._indep_dp_store_addr = indep_dp_store_addr if master_addr: self.master_addr, self.master_port = master_addr, master_port else: diff --git a/train.py b/train.py index e380b087c33..9aceba284ed 100644 --- a/train.py +++ b/train.py @@ -100,7 +100,7 @@ async def save(rollout_id): await offload_train() if args.offload_rollout: await rollout_manager.onload_weights.remote() - await actor_model.update_weights() + await actor_model.update_weights(rollout_id=rollout_id) if args.offload_rollout: await rollout_manager.onload_kv.remote() diff --git a/train_async.py b/train_async.py index 85cb7aeecee..c536d4cfeba 100644 --- a/train_async.py +++ b/train_async.py @@ -74,7 +74,7 @@ async def train(args): # sync generate before update weights to prevent update weight in the middle of generation rollout_data_curr_ref = (await x) if (x := rollout_data_next_future) is not None else None rollout_data_next_future = None - await actor_model.update_weights() + await actor_model.update_weights(rollout_id=rollout_id) if should_run_periodic_action(rollout_id, args.eval_interval, num_rollout_per_epoch): await rollout_manager.eval.remote(rollout_id) From 87c205fe85fe7e3c6e52a56088f6186a0f9c667c Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Mon, 22 Jun 2026 18:15:18 +0800 Subject: [PATCH 35/41] Add the RayTrainCell abstraction for independent-DP cells Add the RayTrainCell abstraction (a group of actors forming one independent-DP replica) together with its cell state model, with the test harness. - miles/ray/train/cell.py, miles/ray/train/cell_state.py and tests. --- miles/ray/train/cell.py | 201 ++++++++++++++++++++++++++++ miles/ray/train/cell_state.py | 32 +++++ tests/conftest.py | 4 + tests/fast/ray/train/__init__.py | 0 tests/fast/ray/train/conftest.py | 82 ++++++++++++ tests/fast/ray/train/dummy_actor.py | 82 ++++++++++++ tests/fast/ray/train/test_cell.py | 91 +++++++++++++ 7 files changed, 492 insertions(+) create mode 100644 miles/ray/train/cell.py create mode 100644 miles/ray/train/cell_state.py create mode 100644 tests/fast/ray/train/__init__.py create mode 100644 tests/fast/ray/train/conftest.py create mode 100644 tests/fast/ray/train/dummy_actor.py create mode 100644 tests/fast/ray/train/test_cell.py diff --git a/miles/ray/train/cell.py b/miles/ray/train/cell.py new file mode 100644 index 00000000000..5bfa8027abf --- /dev/null +++ b/miles/ray/train/cell.py @@ -0,0 +1,201 @@ +import asyncio +import logging +import time +from collections.abc import Callable + +import ray + +from miles.ray.train.cell_state import ( + CellState, + StateAllocatedAlive, + StateAllocatedBase, + StateAllocatedUninitialized, + StatePending, +) +from miles.utils.health_checker import BaseHealthChecker +from miles.utils.indep_dp import IndepDPInfo +from miles.utils.structured_log import log_structured + +logger = logging.getLogger(__name__) + + +ActorFactory = Callable[[], list[ray.actor.ActorHandle]] + + +class RayTrainCell: + def __init__( + self, + *, + args, + role: str, + with_ref: bool, + with_opd_teacher: bool = False, + cell_index: int, + actor_factory: ActorFactory, + rollout_manager: object | None, + health_checker: BaseHealthChecker, + ) -> None: + self.args = args + self.cell_index = cell_index + self.role = role + self.with_ref = with_ref + self.with_opd_teacher = with_opd_teacher + self.rollout_manager = rollout_manager + self.actor_factory = actor_factory + self.health_checker = health_checker + + # NOTE: do *NOT* directly modify `self._state`, but instead use `self._change_state` + self._state: CellState = StatePending() + self.allocate_for_pending() + + # ------------------------ API ------------------------ + + async def init( + self, + *, + indep_dp_info: IndepDPInfo, + ): + results = await self.execute( + "init", + args=self.args, + role=self.role, + with_ref=self.with_ref, + with_opd_teacher=self.with_opd_teacher, + indep_dp_info=indep_dp_info, + ) + self._mark_as_alive(indep_dp_info=indep_dp_info) + await self.health_checker.start() + return results + + async def connect_actor_critic(self, critic_cell: "RayTrainCell") -> list: + critic_handles = critic_cell._get_actor_handles() + return await self._execute_raw( + "connect_actor_critic", + compute_args=lambda i: (critic_handles[i],), + compute_kwargs=lambda _: {}, + ) + + async def set_rollout_manager(self): + if (m := self.rollout_manager) is not None: + return await self.execute("set_rollout_manager", m) + return [] + + # ------------------------ state transition ------------------------ + + def allocate_for_pending(self) -> None: + actor_handles = self.actor_factory() + self._change_state( + "allocate_for_pending", + StatePending, + StateAllocatedUninitialized(actor_handles=actor_handles), + ) + + def _mark_as_alive(self, indep_dp_info: IndepDPInfo) -> None: + self._change_state( + "_mark_as_alive", + StateAllocatedUninitialized, + StateAllocatedAlive(actor_handles=self._state.actor_handles, indep_dp_info=indep_dp_info), + ) + + def _change_state( + self, + debug_name: str, + old_state_cls: type[CellState] | tuple[type[CellState], ...], + new_state: CellState, + ) -> None: + log_structured( + logger.info, + op="state", + phase="start", + name=debug_name, + cell=self.cell_index, + from_state=type(self._state).__name__, + ) + assert isinstance(self._state, old_state_cls), f"{self.cell_index=} {self._state=}" + self._state = new_state + log_structured( + logger.info, + op="state", + phase="end", + name=debug_name, + cell=self.cell_index, + to_state=type(self._state).__name__, + ) + + # ------------------------ API :: directly forward calls to actors ------------------------ + + async def execute(self, fn_name: str, *args, **kwargs) -> list: + return await self._execute_raw( + fn_name, + compute_args=lambda _: args, + compute_kwargs=lambda _: kwargs, + ) + + async def _execute_raw( + self, + fn_name: str, + compute_args, + compute_kwargs, + ) -> list: + handles = self._get_actor_handles() + log_structured( + logger.info, op="execute", phase="start", cell=self.cell_index, fn=fn_name, n_actors=len(handles) + ) + start = time.monotonic() + try: + result = await asyncio.gather( + *[ + getattr(actor, fn_name).remote(*compute_args(i), **compute_kwargs(i)) + for i, actor in enumerate(handles) + ] + ) + log_structured( + logger.info, + op="execute", + phase="end", + cell=self.cell_index, + fn=fn_name, + ok=True, + elapsed_s=round(time.monotonic() - start, 1), + ) + return result + except Exception: + log_structured( + logger.error, + op="execute", + phase="fail", + cell=self.cell_index, + fn=fn_name, + elapsed_s=round(time.monotonic() - start, 1), + exc_info=True, + ) + raise + + # ------------------------ state and misc queries ------------------------ + + @property + def is_pending(self) -> bool: + return isinstance(self._state, StatePending) + + @property + def is_allocated(self) -> bool: + return isinstance(self._state, StateAllocatedBase) + + @property + def is_alive(self) -> bool: + return isinstance(self._state, StateAllocatedAlive) + + @property + def state_name(self) -> str: + return type(self._state).__name__ + + @property + def indep_dp_info(self) -> IndepDPInfo: + assert isinstance(self._state, StateAllocatedAlive) + return self._state.indep_dp_info + + def _get_actor_handles(self) -> list[ray.actor.ActorHandle]: + assert isinstance( + self._state, StateAllocatedBase + ), f"Cell {self.cell_index} is not allocated (state={type(self._state).__name__})" + return self._state.actor_handles diff --git a/miles/ray/train/cell_state.py b/miles/ray/train/cell_state.py new file mode 100644 index 00000000000..bd35f6283b4 --- /dev/null +++ b/miles/ray/train/cell_state.py @@ -0,0 +1,32 @@ +import ray + +from pydantic import BaseModel, ConfigDict + +from miles.utils.indep_dp import IndepDPInfo + + +class StateBase(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True, arbitrary_types_allowed=True) + + +class StatePending(StateBase): + pass + + +class StateAllocatedBase(StateBase): + actor_handles: list[ray.actor.ActorHandle] + + +class StateAllocatedUninitialized(StateAllocatedBase): + pass + + +class StateAllocatedAlive(StateAllocatedBase): + indep_dp_info: IndepDPInfo + + +class StateStopped(StateBase): + pass + + +CellState = StatePending | StateAllocatedUninitialized | StateAllocatedAlive | StateStopped diff --git a/tests/conftest.py b/tests/conftest.py index a064f014dd0..e04f13f1a3a 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -29,6 +29,10 @@ def ray_local_mode(): log_to_driver=False, ) if not os.environ.get("RAY_ADDRESS"): + # address="local" forces a fresh cluster: with no address, ray.init + # auto-connects to any leaked local cluster (via /tmp/ray), and + # connecting with num_cpus/num_gpus set is a hard ValueError. + kwargs["address"] = "local" kwargs["num_cpus"] = 32 # Logical GPU resource so real_ray placement-group tests (engines # are mocked via MockSGLangEngine; no real GPU is used) can satisfy diff --git a/tests/fast/ray/train/__init__.py b/tests/fast/ray/train/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/fast/ray/train/conftest.py b/tests/fast/ray/train/conftest.py new file mode 100644 index 00000000000..de009bdb384 --- /dev/null +++ b/tests/fast/ray/train/conftest.py @@ -0,0 +1,82 @@ +import os +from unittest.mock import MagicMock + +import pytest +import ray +from tests.fast.ray.train.dummy_actor import DummyTrainActor + +from miles.ray.train.cell import RayTrainCell +from miles.utils.health_checker import NoopHealthChecker +from miles.utils.indep_dp import IndepDPInfo + + +@pytest.fixture(scope="module", autouse=True) +def ray_env(): + if ray.is_initialized(): + # Reuse the cluster some outer fixture created (e.g. the session-scoped + # one in tests/conftest.py) and never tear down what we did not create. + yield + return + + init_kwargs: dict = {"ignore_reinit_error": True} + if "RAY_ADDRESS" not in os.environ: + # address="local" forces a fresh cluster: with no address, ray.init + # auto-connects to any leaked local cluster (via /tmp/ray), and + # connecting with num_cpus/num_gpus set is a hard ValueError. + init_kwargs["address"] = "local" + init_kwargs["num_cpus"] = 4 + init_kwargs["num_gpus"] = 0 + ray.init(**init_kwargs) + yield + ray.shutdown() + + +def make_indep_dp_info( + *, + cell_index: int = 0, + alive_cell_indices: list[int] | None = None, + quorum_id: int = 1, +) -> IndepDPInfo: + if alive_cell_indices is None: + alive_cell_indices = [0] + return IndepDPInfo( + cell_index=cell_index, + num_cells=3, + alive_rank=alive_cell_indices.index(cell_index), + alive_size=len(alive_cell_indices), + quorum_id=quorum_id, + alive_cell_indices=alive_cell_indices, + ) + + +def make_cell( + cell_index: int = 0, + *, + actor_count: int = 2, + rollout_manager: object | None = None, +) -> RayTrainCell: + def factory(): + return [DummyTrainActor.remote() for _ in range(actor_count)] + + return RayTrainCell( + args=MagicMock(), + role="actor", + with_ref=False, + cell_index=cell_index, + actor_factory=factory, + rollout_manager=rollout_manager, + health_checker=NoopHealthChecker(), + ) + + +def make_alive_cell(cell_index: int, *, alive_cell_indices: list[int], quorum_id: int = 0) -> RayTrainCell: + """Create a cell and transition it to Alive state.""" + cell = make_cell(cell_index) + cell._mark_as_alive( + indep_dp_info=make_indep_dp_info( + cell_index=cell_index, + alive_cell_indices=alive_cell_indices, + quorum_id=quorum_id, + ) + ) + return cell diff --git a/tests/fast/ray/train/dummy_actor.py b/tests/fast/ray/train/dummy_actor.py new file mode 100644 index 00000000000..22614e6da2c --- /dev/null +++ b/tests/fast/ray/train/dummy_actor.py @@ -0,0 +1,82 @@ +"""Lightweight Ray actor for unit testing RayTrainCell/RayTrainGroup without GPU or real training. + +Records all method calls so tests can verify what was dispatched. +""" + +from typing import Any + +import ray + +from miles.backends.megatron_utils.types import TrainStepOutcome +from miles.utils.heartbeat_utils import HeartbeatStatus, SimpleHeartbeat + + +@ray.remote(num_gpus=0, num_cpus=0) +class DummyTrainActor: + + def __init__(self): + self._calls: list[tuple[str, tuple, dict]] = [] + self._fail_methods: set[str] = set() + self._train_return_value: Any = TrainStepOutcome.NORMAL + self._heartbeat = SimpleHeartbeat() + self._heartbeat.bump() + self._heartbeat_fail: bool = False + + def set_fail_methods(self, methods: list[str]) -> None: + self._fail_methods = set(methods) + + def set_train_return_value(self, value: Any) -> None: + self._train_return_value = value + + def _record(self, method: str, args: tuple, kwargs: dict) -> None: + self._calls.append((method, args, kwargs)) + if method in self._fail_methods: + raise RuntimeError(f"Injected failure in {method}") + + def get_calls(self) -> list[tuple[str, tuple, dict]]: + return list(self._calls) + + def init(self, *args: Any, **kwargs: Any) -> None: + self._record("init", args, kwargs) + + def reconfigure_indep_dp(self, *args: Any, **kwargs: Any) -> None: + self._record("reconfigure_indep_dp", args, kwargs) + + def send_ckpt(self, *args: Any, **kwargs: Any) -> None: + self._record("send_ckpt", args, kwargs) + + def train(self, *args: Any, **kwargs: Any) -> Any: + self._record("train", args, kwargs) + return self._train_return_value + + def set_rollout_manager(self, *args: Any, **kwargs: Any) -> None: + self._record("set_rollout_manager", args, kwargs) + + def wake_up(self) -> None: + self._record("wake_up", (), {}) + + def sleep(self) -> None: + self._record("sleep", (), {}) + + def clear_memory(self) -> None: + self._record("clear_memory", (), {}) + + def save_model(self, *args: Any, **kwargs: Any) -> None: + self._record("save_model", args, kwargs) + + def update_weights(self) -> None: + self._record("update_weights", (), {}) + + def set_heartbeat_fail(self, fail: bool) -> None: + self._heartbeat_fail = fail + + def set_last_active_timestamp(self, ts: float) -> None: + self._heartbeat._status = HeartbeatStatus( + last_active_timestamp=ts, + bump_count=self._heartbeat._status.bump_count, + ) + + def get_heartbeat_status(self) -> HeartbeatStatus: + if self._heartbeat_fail: + raise RuntimeError("Injected heartbeat failure") + return self._heartbeat.status() diff --git a/tests/fast/ray/train/test_cell.py b/tests/fast/ray/train/test_cell.py new file mode 100644 index 00000000000..67ac5d00657 --- /dev/null +++ b/tests/fast/ray/train/test_cell.py @@ -0,0 +1,91 @@ +import pytest +import ray + +from tests.fast.ray.train.conftest import make_alive_cell, make_cell, make_indep_dp_info + +pytestmark = pytest.mark.asyncio + + +class TestInitialState: + def test_starts_as_uninitialized_after_init(self): + """After __init__, cell is allocated (uninitialized) — actors created but not init'd.""" + cell = make_cell() + + assert cell.is_allocated + assert not cell.is_alive + assert not cell.is_pending + + def test_actor_handles_are_real_ray_actors(self): + cell = make_cell(actor_count=3) + + handles = cell._get_actor_handles() + assert len(handles) == 3 + assert all(isinstance(h, ray.actor.ActorHandle) for h in handles) + + +class TestMarkAsAlive: + def test_transitions_uninitialized_to_alive(self): + cell = make_cell() + info = make_indep_dp_info(alive_cell_indices=[0, 1, 2]) + + cell._mark_as_alive(indep_dp_info=info) + + assert cell.is_alive + assert cell.indep_dp_info == info + + def test_preserves_actor_handles(self): + cell = make_cell(actor_count=3) + handles_before = cell._get_actor_handles() + + cell._mark_as_alive(indep_dp_info=make_indep_dp_info()) + + assert cell._get_actor_handles() == handles_before + + def test_rejects_from_alive(self): + cell = make_alive_cell(0, alive_cell_indices=[0]) + + with pytest.raises(AssertionError): + cell._mark_as_alive(indep_dp_info=make_indep_dp_info()) + + +class TestInvalidTransitions: + def test_allocate_for_pending_rejects_from_alive(self): + cell = make_alive_cell(0, alive_cell_indices=[0]) + + with pytest.raises(AssertionError): + cell.allocate_for_pending() + + +class TestAsyncInit: + async def test_dispatches_init_and_marks_alive(self): + cell = make_cell(actor_count=2) + info = make_indep_dp_info() + + results = await cell.init(indep_dp_info=info) + + assert len(results) == 2 + assert cell.is_alive + assert cell.indep_dp_info == info + + for handle in cell._get_actor_handles(): + calls = ray.get(handle.get_calls.remote()) + assert len(calls) == 1 + assert calls[0][0] == "init" + kwargs = calls[0][2] + assert kwargs["indep_dp_info"] == info + + +class TestStatePredicates: + def test_uninitialized(self): + cell = make_cell() + + assert not cell.is_pending + assert cell.is_allocated + assert not cell.is_alive + + def test_alive(self): + cell = make_alive_cell(0, alive_cell_indices=[0]) + + assert not cell.is_pending + assert cell.is_allocated + assert cell.is_alive From 6324e8b8cc3325ff36b452dd875094c88c0c6750 Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Mon, 22 Jun 2026 18:15:18 +0800 Subject: [PATCH 36/41] Add the cell-based independent-DP train group selected via experimental flag Add the cell-based independent-DP train group (RayTrainGroupV2) that orchestrates RayTrainCells, selected via the experimental flag, and dispatch to it from the placement group. - miles/ray/train/group.py, placement_group.py, actor_group.py and actor wiring, with tests. --- .../backends/experimental/fsdp_utils/actor.py | 14 +- miles/backends/megatron_utils/actor.py | 3 + miles/ray/actor_group.py | 9 +- miles/ray/placement_group.py | 13 +- miles/ray/train/group.py | 249 ++++++++++++++++++ miles/utils/arguments.py | 15 ++ tests/fast/ray/train/test_group.py | 236 +++++++++++++++++ 7 files changed, 535 insertions(+), 4 deletions(-) create mode 100644 miles/ray/train/group.py create mode 100644 tests/fast/ray/train/test_group.py diff --git a/miles/backends/experimental/fsdp_utils/actor.py b/miles/backends/experimental/fsdp_utils/actor.py index 98b8b37f0e4..e5051901f35 100644 --- a/miles/backends/experimental/fsdp_utils/actor.py +++ b/miles/backends/experimental/fsdp_utils/actor.py @@ -14,6 +14,7 @@ from miles.utils.context_utils import with_defer from miles.utils.distributed_utils import get_gloo_group from miles.utils.hf_config import load_hf_config +from miles.utils.indep_dp import IndepDPInfo from miles.utils.memory_utils import clear_memory, print_memory from miles.utils.processing_utils import load_processor, load_tokenizer from miles.utils.ray_utils import Box @@ -56,9 +57,20 @@ class FSDPTrainRayActor(TrainRayActor): """ @with_defer(lambda: Timer().start("train_wait")) - def init(self, args: Namespace, role: str, with_ref: bool = False, with_opd_teacher: bool = False) -> int: # type: ignore[override] + def init( + self, + args: Namespace, + role: str, + *, + with_ref: bool = False, + with_opd_teacher: bool = False, + indep_dp_info: IndepDPInfo, + ) -> int | None: # type: ignore[override] super().init(args, role, with_ref, with_opd_teacher=with_opd_teacher) + # Unsupported + assert indep_dp_info.quorum_id == 0 + if args.dumper_enable: from sglang.srt.debug_utils.dumper import dumper diff --git a/miles/backends/megatron_utils/actor.py b/miles/backends/megatron_utils/actor.py index b09c2c3a03b..a3a4811467c 100644 --- a/miles/backends/megatron_utils/actor.py +++ b/miles/backends/megatron_utils/actor.py @@ -16,6 +16,7 @@ from miles.utils.context_utils import with_defer from miles.utils.distributed_utils import get_gloo_group, init_process_group from miles.utils.hf_config import load_hf_config +from miles.utils.indep_dp import IndepDPInfo from miles.utils.memory_utils import clear_memory, print_memory from miles.utils.processing_utils import load_tokenizer from miles.utils.ray_utils import Box @@ -57,8 +58,10 @@ def init( self, args: Namespace, role: str, + *, with_ref: bool = False, with_opd_teacher: bool = False, + indep_dp_info: IndepDPInfo, ) -> int | None: monkey_patch_torch_dist() diff --git a/miles/ray/actor_group.py b/miles/ray/actor_group.py index 59cd4004d96..1d3539338af 100644 --- a/miles/ray/actor_group.py +++ b/miles/ray/actor_group.py @@ -8,6 +8,7 @@ from ray.util.placement_group import PlacementGroup from miles.ray.train.actor_factory import allocate_gpus_for_actor +from miles.utils.indep_dp import IndepDPInfo class RayTrainGroup: @@ -63,8 +64,14 @@ async def init(self): """ Allocate GPU resourced and initialize model, optimizer, local ckpt, etc. """ + indep_dp_info = IndepDPInfo.create_trivial() return await self._broadcast( - "init", self.args, self.role, with_ref=self.with_ref, with_opd_teacher=self.with_opd_teacher + "init", + self.args, + self.role, + with_ref=self.with_ref, + with_opd_teacher=self.with_opd_teacher, + indep_dp_info=indep_dp_info, ) async def train(self, rollout_id, rollout_data_pack): diff --git a/miles/ray/placement_group.py b/miles/ray/placement_group.py index 3a4e3ab690d..a9e7122c45d 100644 --- a/miles/ray/placement_group.py +++ b/miles/ray/placement_group.py @@ -6,14 +6,22 @@ from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy from miles.utils.async_utils import eager_create_task +from miles.utils.environ import enable_experimental_ft_trainer from ..utils.ray_utils import compute_ray_pin_head_options -from .actor_group import RayTrainGroup from .rollout.rollout_manager import RolloutManager 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): @@ -125,7 +133,8 @@ def create_placement_groups(args): def allocate_train_group( args, num_nodes, num_gpus_per_node, pg, role: str, with_ref: bool, rollout_manager, with_opd_teacher: bool = False ): - return RayTrainGroup( + train_group_cls = _select_train_group_class() + return train_group_cls( 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 new file mode 100644 index 00000000000..ebdf407f9af --- /dev/null +++ b/miles/ray/train/group.py @@ -0,0 +1,249 @@ +import asyncio +import logging +from pathlib import Path +from typing import TYPE_CHECKING + +import ray +from ray.util.placement_group import PlacementGroup + +from miles.ray.train.actor_factory import allocate_gpus_for_actor +from miles.ray.train.cell import RayTrainCell +from miles.utils.event_logger.logger import get_event_logger, is_event_logger_initialized +from miles.utils.event_logger.models import WitnessAllocateIdEvent +from miles.utils.health_checker import NoopHealthChecker +from miles.utils.indep_dp import IndepDPInfo +from miles.utils.megatron_args_utils import compute_megatron_world_size_except_dp +from miles.utils.structured_log import log_structured +from miles.utils.witness.allocator import WitnessIdAllocator, read_persisted_witness_counter + +if TYPE_CHECKING: + import torch + + +logger = logging.getLogger(__name__) + + +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. + resources (Dict[str, float], optional): Custom resources to allocate for each actor. + See https://docs.ray.io/en/latest/ray-core/scheduling/resources.html + num_resources_per_node (int, optional): Number of custom resources to allocate for each node. + See https://docs.ray.io/en/latest/ray-core/scheduling/resources.html + """ + + def __init__( + self, + args, + num_nodes: int, + num_gpus_per_node: int, + pg: tuple[PlacementGroup, list[int], list[int]], + *, + rollout_manager: object | None, + num_gpus_per_actor: float = 1, + role: str, + with_ref: bool, + with_opd_teacher: bool = False, + ) -> None: + self.args = args + self._rollout_manager = rollout_manager + + total_gpus = num_nodes * num_gpus_per_node + num_cells = (total_gpus // compute_megatron_world_size_except_dp(args)) if args.indep_dp else 1 + gpus_per_cell = total_gpus // num_cells + assert total_gpus % num_cells == 0, f"total_gpus ({total_gpus}) must be divisible by num_cells ({num_cells})" + + self._indep_dp_quorum_id = 0 + + if num_cells > 1: + self._indep_dp_store, indep_dp_store_addr = _create_tcp_store() + logger.info(f"Created TCPStore for independent DP at {indep_dp_store_addr}") + else: + self._indep_dp_store, indep_dp_store_addr = None, None + + def _create_cell(cell_index: int): + cell_pg = _slice_pg(pg, start=cell_index * gpus_per_cell, end=(cell_index + 1) * gpus_per_cell) + + cell = RayTrainCell( + args=args, + role=role, + with_ref=with_ref, + with_opd_teacher=with_opd_teacher, + cell_index=cell_index, + rollout_manager=rollout_manager, + actor_factory=lambda _pg=cell_pg, _ci=cell_index: allocate_gpus_for_actor( + args=args, + gpus_per_cell=gpus_per_cell, + pg=_pg, + num_gpus_per_actor=num_gpus_per_actor, + indep_dp_store_addr=indep_dp_store_addr, + role=role, + cell_index=_ci, + ), + health_checker=NoopHealthChecker(), + ) + + return cell + + self._cells: list[RayTrainCell] = [_create_cell(cell_index) for cell_index in range(num_cells)] + + self._witness_allocator: WitnessIdAllocator | None = ( + WitnessIdAllocator(buffer_size=args.witness_buffer_size) if args.enable_witness else None + ) + if self._witness_allocator is not None and args.save_debug_event_data is not None: + self._witness_allocator.resume(read_persisted_witness_counter(Path(args.save_debug_event_data))) + + # ------------------------ API :: train ------------------------ + + async def train(self, rollout_id: int, rollout_data_pack): + """Do one rollout training""" + witness_info = self._allocate_witness_info( + rollout_id=rollout_id, + attempt=0, + sample_indices=rollout_data_pack["sample_indices"], + ) + + log_structured(logger.info, op="train", phase="start", rollout=rollout_id, attempt=0) + await self._execute_all( + "train", + rollout_id=rollout_id, + rollout_data_ref=rollout_data_pack["data_ref"], + witness_info=witness_info, + attempt=0, + ) + + def _allocate_witness_info(self, *, rollout_id: int, attempt: int, sample_indices): + if self._witness_allocator is None: + return None + + witness_info = self._witness_allocator.allocate(num_ids=len(sample_indices)) + + if is_event_logger_initialized(): + get_event_logger().log( + WitnessAllocateIdEvent, + dict( + rollout_id=rollout_id, + attempt=attempt, + witness_id_to_sample_index=dict(zip(witness_info.witness_ids, sample_indices, strict=True)), + counter_after=self._witness_allocator.counter, + ), + ) + + return witness_info + + # ------------------------ API :: others ------------------------ + + async def init(self): + """ + Allocate GPU resourced and initialize model, optimzier, local ckpt, etc. + """ + cell_results = await asyncio.gather( + *[ + cell.init( + indep_dp_info=self._compute_indep_dp_info( + cell_index=cell.cell_index, + # all cells will be alive for this first initialization + alive_cell_indices=list(range(len(self._cells))), + ) + ) + for cell in self._cells + ] + ) + return [item for sublist in cell_results for item in sublist] + + async def save_model(self, rollout_id: int, force_sync: bool = False): + """Save actor model. Only cell 0 saves to avoid file write conflicts.""" + await self._execute_first("save_model", rollout_id, force_sync=force_sync) + + async def update_weights(self, rollout_id: int | None = None): + """Broadcast weights to rollout engines.""" + log_structured(logger.info, op="update_weights", phase="start", rollout=rollout_id) + # TODO: allow using all cells to update weights (instead of first alive cell) + # Fetch the updatable engines + lock once (like V1 RayActorGroup) so all + # ranks observe a consistent engine set; the actor releases the lock itself. + info = await self._rollout_manager.get_updatable_engines_and_lock.remote() + await self._rollout_manager.health_monitoring_pause.remote() + await self._execute_first("update_weights", info=info) + + async def onload(self): + await self._execute_all("wake_up") + for cell in self._cells: + cell.health_checker.resume() + + async def offload(self): + for cell in self._cells: + cell.health_checker.pause() + await self._execute_all("sleep") + + async def clear_memory(self): + await self._execute_all("clear_memory") + + async def connect(self, critic_group: "RayTrainGroup"): + assert len(self._cells) == len(critic_group._cells), ( + f"Actor and critic must have the same number of cells: " + f"actor has {len(self._cells)}, critic has {len(critic_group._cells)}" + ) + await asyncio.gather( + *[ + cell.connect_actor_critic(critic_cell) + for cell, critic_cell in zip(self._cells, critic_group._cells, strict=True) + ] + ) + + async def set_rollout_manager(self): + await asyncio.gather(*[cell.set_rollout_manager() for cell in self._cells]) + + # ------------------------ utils to forward calls to cells ------------------------ + + async def _execute_all(self, fn_name: str, *args, **kwargs): + return await asyncio.gather(*[cell.execute(fn_name, *args, **kwargs) for cell in self._cells]) + + async def _execute_first(self, fn_name: str, *args, **kwargs): + return await self._cells[0].execute(fn_name, *args, **kwargs) + + def _compute_indep_dp_info(self, cell_index: int, alive_cell_indices: list[int]) -> IndepDPInfo: + return IndepDPInfo( + cell_index=cell_index, + num_cells=len(self._cells), + alive_rank=alive_cell_indices.index(cell_index), + alive_size=len(alive_cell_indices), + quorum_id=self._indep_dp_quorum_id, + alive_cell_indices=alive_cell_indices, + ) + + # ------------------------ misc states and utils ------------------------ + + @property + def num_cells(self) -> int: + return len(self._cells) + + +PGTuple = tuple[PlacementGroup, list[int], list[int]] + + +def _slice_pg(pg: PGTuple, start: int, end: int) -> PGTuple: + placement_group, bundle_indices, gpu_ids = pg + return placement_group, bundle_indices[start:end], gpu_ids[start:end] + + +def _create_tcp_store() -> tuple["torch.distributed.TCPStore", str]: + import torch.distributed + + store = torch.distributed.TCPStore( + host_name="0.0.0.0", + port=0, + is_master=True, + wait_for_workers=False, + ) + host = ray.util.get_node_ip_address() + port = store.port + return store, f"{host}:{port}" diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index a1989fb0e7e..0ebae315986 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -14,6 +14,7 @@ from miles.utils.eval_config import EvalDatasetConfig, build_eval_dataset_configs, ensure_dataset_list from miles.utils.hf_config import is_dsa, load_hf_config from miles.utils.logging_utils import configure_logger_raw +from miles.utils.megatron_args_utils import compute_megatron_world_size_except_dp from miles.utils.misc import load_function logger = logging.getLogger(__name__) @@ -270,6 +271,12 @@ def add_train_arguments(parser): parser.add_argument( "--log-probs-chunk-size", type=int, default=-1, help="Chunk size to compute log probs to save memory" ) + parser.add_argument( + "--indep-dp", + action="store_true", + default=False, + help="Launch each DP replica as an independent Megatron instance instead of using Megatron-internal data parallelism.", + ) parser.add_argument( "--delay-split-train-data-by-dp", action="store_true", @@ -2149,6 +2156,14 @@ def _resolve_eval_datasets(args) -> list[EvalDatasetConfig]: def miles_validate_args(args): args.eval_datasets = _resolve_eval_datasets(args) + if args.indep_dp: + assert ( + args.train_backend == "megatron" + ), f"indep_dp requires train_backend='megatron', got '{args.train_backend}'" + per_replica_size = compute_megatron_world_size_except_dp(args) + logger.info(f"indep_dp: adjusting args.world_size from {args.world_size} to {per_replica_size} (per-cell)") + args.world_size = per_replica_size + if args.recompute_logprobs_via_prefill: assert args.true_on_policy_mode, "--recompute-logprobs-via-prefill requires --true-on-policy-mode" diff --git a/tests/fast/ray/train/test_group.py b/tests/fast/ray/train/test_group.py new file mode 100644 index 00000000000..204cb1f0d2a --- /dev/null +++ b/tests/fast/ray/train/test_group.py @@ -0,0 +1,236 @@ +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest +import ray +from tests.fast.ray.train.dummy_actor import DummyTrainActor + +from miles.ray.train.group import RayTrainGroup +from miles.utils.witness.allocator import WitnessIdAllocator + +pytestmark = pytest.mark.asyncio + +_DUMMY_DATA_PACK = {"data_ref": "data", "sample_indices": [0]} + + +def _make_mock_args( + *, + indep_dp: bool = True, + enable_witness: bool = False, + gpus_per_cell: int = 1, +) -> SimpleNamespace: + # Use SimpleNamespace (not MagicMock) so the args object is picklable. RayTrainCell.init + # passes self.args through Ray to the remote actor; pickling a MagicMock blows the + # recursion limit because its __getattr__ creates new sub-mocks indefinitely. + return SimpleNamespace( + indep_dp=indep_dp, + enable_witness=enable_witness, + witness_buffer_size=100, + trainer_heartbeat_checker_interval=10.0, + trainer_heartbeat_checker_timeout=10.0, + trainer_heartbeat_checker_first_wait=300.0, + trainer_heartbeat_checker_failure_threshold=3, + ci_ft_test_actions=None, + debug_train_only=False, + debug_rollout_only=False, + # compute_megatron_world_size_except_dp(args) = TP * PP * CP. Set CP to + # gpus_per_cell so RayTrainGroup computes num_cells correctly. + tensor_model_parallel_size=1, + pipeline_model_parallel_size=1, + context_parallel_size=gpus_per_cell, + ) + + +@pytest.fixture(autouse=True) +def _patch_actor_alloc(): + """Persist allocate_gpus_for_actor patch across the whole test (incl. healing path). + + Previously _make_group used `with patch(...)` which expired when _make_group + returned, so any later `allocate_gpus_for_actor` call during _refresh_cells + healing hit the real implementation (which dereferences mock args fields). + """ + + def _alloc(*, gpus_per_cell: int, num_gpus_per_actor: float, **_kwargs) -> list: + actor_count = max(int(gpus_per_cell // num_gpus_per_actor), 1) + return [DummyTrainActor.remote() for _ in range(actor_count)] + + with patch("miles.ray.train.group.allocate_gpus_for_actor", side_effect=_alloc): + yield + + +def _make_group( + *, + num_cells: int = 3, + actor_count_per_cell: int = 1, + rollout_manager: object | None = None, +) -> RayTrainGroup: + """Create a RayTrainGroup through real __init__ with mocked pg and actor factory.""" + total_gpus = num_cells * actor_count_per_cell + return RayTrainGroup( + args=_make_mock_args(indep_dp=True, gpus_per_cell=actor_count_per_cell), + num_nodes=1, + num_gpus_per_node=total_gpus, + pg=(MagicMock(), list(range(total_gpus)), list(range(total_gpus))), + role="actor", + with_ref=False, + rollout_manager=rollout_manager, + ) + + +async def _init_group(group: RayTrainGroup) -> None: + """Call init and wait for all cells to become alive.""" + await group.init() + + +async def _make_alive_group(*, num_cells: int = 3, **kwargs) -> RayTrainGroup: + """Create a group and init all cells to alive.""" + group = _make_group(num_cells=num_cells, **kwargs) + await _init_group(group) + return group + + +class TestInit: + def test_creates_correct_number_of_cells(self): + group = _make_group(num_cells=3) + + assert len(group._cells) == 3 + assert [c.cell_index for c in group._cells] == [0, 1, 2] + + def test_cells_are_allocated_after_init(self): + group = _make_group(num_cells=2) + + for cell in group._cells: + assert cell.is_allocated + assert not cell.is_alive + + def test_each_cell_has_own_actors(self): + group = _make_group(num_cells=3, actor_count_per_cell=2) + + handles_per_cell = [cell._get_actor_handles() for cell in group._cells] + assert all(len(h) == 2 for h in handles_per_cell) + + all_handles = [h for handles in handles_per_cell for h in handles] + assert len(set(id(h) for h in all_handles)) == 6 + + def test_single_cell_no_tcp_store(self): + # indep_dp=False forces single cell regardless of TP/PP/CP product; + # the autouse fixture handles allocate_gpus_for_actor. + group = RayTrainGroup( + args=_make_mock_args(indep_dp=False), + num_nodes=1, + num_gpus_per_node=1, + pg=(MagicMock(), [0], [0]), + role="actor", + with_ref=False, + rollout_manager=None, + ) + + assert len(group._cells) == 1 + assert group._indep_dp_store is None + + async def test_init_marks_all_cells_alive(self): + group = _make_group(num_cells=3) + + await _init_group(group) + + for cell in group._cells: + assert cell.is_alive + assert cell.indep_dp_info.alive_cell_indices == [0, 1, 2] + assert cell.indep_dp_info.alive_size == 3 + + assert group._cells[0].indep_dp_info.alive_rank == 0 + assert group._cells[1].indep_dp_info.alive_rank == 1 + assert group._cells[2].indep_dp_info.alive_rank == 2 + + +class TestExecuteFirstAlive: + async def test_picks_first_alive_cell(self): + group = await _make_alive_group(num_cells=3) + + await group._execute_first("save_model", 42) + + for handle in group._cells[0]._get_actor_handles(): + calls = ray.get(handle.get_calls.remote()) + assert any(c[0] == "save_model" for c in calls) + + for cell in group._cells[1:]: + for handle in cell._get_actor_handles(): + calls = ray.get(handle.get_calls.remote()) + assert not any(c[0] == "save_model" for c in calls) + + +class TestComputeIndepDPInfo: + def test_all_alive(self): + group = _make_group(num_cells=3) + + info = group._compute_indep_dp_info(cell_index=2, alive_cell_indices=[0, 1, 2]) + + assert info.alive_rank == 2 + assert info.alive_size == 3 + assert info.cell_index == 2 + + def test_with_gap(self): + group = _make_group(num_cells=3) + + info = group._compute_indep_dp_info(cell_index=2, alive_cell_indices=[0, 2]) + + assert info.alive_rank == 1 + assert info.alive_size == 2 + + +class TestTrain: + async def test_train_refreshes_and_dispatches(self): + group = await _make_alive_group(num_cells=2) + + await group.train(rollout_id=0, rollout_data_pack=_DUMMY_DATA_PACK) + + for cell in group._cells: + for handle in cell._get_actor_handles(): + calls = ray.get(handle.get_calls.remote()) + assert any(c[0] == "train" for c in calls) + + async def test_consecutive_train_no_reconfigure_overhead(self): + """Multiple train calls with no state changes — no reconfigure overhead.""" + group = await _make_alive_group(num_cells=3) + + # Note init call count + init_counts = {} + for cell in group._cells: + for handle in cell._get_actor_handles(): + init_counts[id(handle)] = len(ray.get(handle.get_calls.remote())) + + for step in range(3): + await group.train(rollout_id=step, rollout_data_pack=_DUMMY_DATA_PACK) + + assert group._indep_dp_quorum_id == 0 + + for cell in group._cells: + for handle in cell._get_actor_handles(): + calls = ray.get(handle.get_calls.remote()) + new_calls = calls[init_counts[id(handle)] :] + assert not any(c[0] == "reconfigure_indep_dp" for c in new_calls) + train_calls = [c for c in new_calls if c[0] == "train"] + assert len(train_calls) == 3 + + +class TestAllocateWitnessInfo: + def test_returns_none_when_disabled(self): + """When _witness_allocator is None, _allocate_witness_info returns None.""" + group = _make_group(num_cells=1) + group._witness_allocator = None + + result = group._allocate_witness_info(rollout_id=0, attempt=0, sample_indices=[10, 20, 30]) + + assert result is None + + def test_returns_witness_info_when_enabled(self): + """When witness is enabled, _allocate_witness_info returns a WitnessInfo with correct number of ids.""" + group = _make_group(num_cells=1) + group._witness_allocator = WitnessIdAllocator(buffer_size=100) + + with patch("miles.ray.train.group.is_event_logger_initialized", return_value=False): + result = group._allocate_witness_info(rollout_id=0, attempt=0, sample_indices=[10, 20, 30]) + + assert result is not None + assert len(result.witness_ids) == 3 + assert isinstance(result.stale_ids, list) From df5fc51a239f31fbd6cc298ea7ecbe895c48ed0b Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Mon, 22 Jun 2026 18:15:18 +0800 Subject: [PATCH 37/41] Drop the legacy self.rollout_engines initialization from MegatronTrainRayActor Also tidy formatting in the Megatron actor and model touched here. --- miles/backends/megatron_utils/actor.py | 2 -- miles/backends/megatron_utils/model.py | 1 + 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/miles/backends/megatron_utils/actor.py b/miles/backends/megatron_utils/actor.py index a3a4811467c..f40077abe5a 100644 --- a/miles/backends/megatron_utils/actor.py +++ b/miles/backends/megatron_utils/actor.py @@ -203,8 +203,6 @@ def init( if self.args.offload_train: self.sleep() - self.rollout_engines = None - self.rollout_data_postprocess = None if (x := self.args.rollout_data_postprocess_path) is not None: from miles.utils.misc import load_function diff --git a/miles/backends/megatron_utils/model.py b/miles/backends/megatron_utils/model.py index 82f528ba7da..2b3b85a5a9d 100644 --- a/miles/backends/megatron_utils/model.py +++ b/miles/backends/megatron_utils/model.py @@ -477,6 +477,7 @@ def forward_step(data_iterator: DataIterator, model: GPTModel, return_schedule_p valid_step = True grad_norm = 0.0 + if (not disable_optimizer) and (not getattr(args, "check_for_nan_in_loss_and_grad", True)): found_inf_flag = optimizer.prepare_grads() if found_inf_flag: From 1e5bff238f95e07d0d1f5914bcad0126a1208481 Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Mon, 22 Jun 2026 18:15:18 +0800 Subject: [PATCH 38/41] Skip the megatron dp_size hint under independent DP When independent DP is enabled the megatron data-parallel size is 1 per cell, so omit the `dp_size` hint from the actor's train parallel config (cross-replica DP is handled separately) instead of reporting the intra-cell dp size. - actor.py: gate `train_parallel_config` on `args.indep_dp`. --- miles/backends/megatron_utils/actor.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/miles/backends/megatron_utils/actor.py b/miles/backends/megatron_utils/actor.py index f40077abe5a..784a89e8b63 100644 --- a/miles/backends/megatron_utils/actor.py +++ b/miles/backends/megatron_utils/actor.py @@ -99,9 +99,7 @@ def init( ) dist.barrier(group=get_gloo_group()) - self.train_parallel_config = { - "dp_size": get_parallel_state().intra_dp.size, - } + self.train_parallel_config = {} if args.indep_dp else {"dp_size": get_parallel_state().intra_dp.size} dist.barrier(group=get_gloo_group()) if args.offload_train: From 10ba50d335aed3ab5dff54b2ea0c8f2b062b3b55 Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Mon, 22 Jun 2026 18:15:18 +0800 Subject: [PATCH 39/41] Extract the CP-aware token-id transform into a shared helper Pull the bshd/thd slice-pad-stack transform out of the position-ids path in get_batch into a local _compute_transform_like_token_ids helper, so the same transform can be reused for other per-token id streams. No behaviour change. - miles/backends/training_utils/data.py. --- miles/backends/training_utils/data.py | 28 +++++++++++++++------------ 1 file changed, 16 insertions(+), 12 deletions(-) diff --git a/miles/backends/training_utils/data.py b/miles/backends/training_utils/data.py index a43e0aecfbf..71aa4113959 100644 --- a/miles/backends/training_utils/data.py +++ b/miles/backends/training_utils/data.py @@ -215,25 +215,29 @@ def get_batch( batch["tokens"] = tokens - if get_position_ids: + def _compute_transform_like_token_ids(ids_list: list): assert not allgather_cp, "allgather CP is not supported for FSDP" + if qkv_format == "bshd": + ids = [slice_with_cp(p, 0, qkv_format, max_seqlen) for p in ids_list] + ids = torch.stack(ids) + elif qkv_format == "thd": + ids = [slice_with_cp(p, 0, qkv_format) for p in ids_list] + ids = torch.cat(ids) + if pad != 0: + ids = F.pad(ids, (0, pad), value=0) + ids = ids.unsqueeze(0) + else: + raise NotImplementedError + return ids + + if get_position_ids: position_ids_list = [] for t in batch["unconcat_tokens"]: seq_len = t.size(0) pos_ids = torch.arange(seq_len, device=t.device, dtype=torch.long) position_ids_list.append(pos_ids) - if qkv_format == "bshd": - position_ids = [slice_with_cp(p, 0, qkv_format, max_seqlen) for p in position_ids_list] - position_ids = torch.stack(position_ids) - elif qkv_format == "thd": - position_ids = [slice_with_cp(p, 0, qkv_format) for p in position_ids_list] - position_ids = torch.cat(position_ids) - if pad != 0: - position_ids = F.pad(position_ids, (0, pad), value=0) - position_ids = position_ids.unsqueeze(0) - - batch["position_ids"] = position_ids + batch["position_ids"] = _compute_transform_like_token_ids(position_ids_list) # loss masks loss_masks = [] From af7398492da2c30d53ee45396f69c97699f6f09f Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Mon, 22 Jun 2026 18:15:18 +0800 Subject: [PATCH 40/41] Thread witness ids through the training data path Carry per-sample witness ids from the rollout data into training so each token can be tagged with its originating witness id (used by the event analyzer to trace samples across rollouts), reusing the CP-aware token-id transform helper in get_batch. - utils/data.py / training_utils/data.py: attach seq_witness_ids and expand into per-token witness_ids. - actor.py / model.py: thread witness_info / attempt through train / train_actor / train_one_step. - log_utils.py: exclude witness_ids from rollout-data logging. --- .../backends/experimental/fsdp_utils/actor.py | 14 +++++++++++-- miles/backends/megatron_utils/actor.py | 21 +++++++++++++++---- miles/backends/megatron_utils/model.py | 9 ++++++++ miles/backends/training_utils/data.py | 18 +++++++++++++++- miles/backends/training_utils/log_utils.py | 1 + miles/ray/actor_group.py | 8 ++++++- miles/utils/data.py | 5 +++++ 7 files changed, 68 insertions(+), 8 deletions(-) diff --git a/miles/backends/experimental/fsdp_utils/actor.py b/miles/backends/experimental/fsdp_utils/actor.py index e5051901f35..21104e8ca99 100644 --- a/miles/backends/experimental/fsdp_utils/actor.py +++ b/miles/backends/experimental/fsdp_utils/actor.py @@ -39,6 +39,7 @@ if TYPE_CHECKING: from miles.ray.rollout.rollout_manager import EnginesAndLock + from miles.utils.witness.allocator import WitnessInfo logger = logging.getLogger(__name__) @@ -406,7 +407,13 @@ def _compute_log_prob( self.model.cuda() dist.barrier(group=get_gloo_group()) - def train(self, rollout_id: int, rollout_data_ref: Box) -> None: + def train( + self, + rollout_id: int, + rollout_data_ref: Box, + witness_info: "WitnessInfo | None" = None, + attempt: int = 0, + ) -> None: """Run one training update over a rollout batch. Parameters: @@ -417,11 +424,14 @@ def train(self, rollout_id: int, rollout_data_ref: Box) -> None: `rollout_log_probs`, etc.). It will be fetched and partitioned by `process_rollout_data` based on data-parallel rank/size. """ + assert witness_info is None + assert attempt == 0 + if self.args.offload_train: self.wake_up() with inverse_timer("train_wait"), timer("train"): - rollout_data = get_rollout_data(self.args, rollout_data_ref) + rollout_data = get_rollout_data(self.args, rollout_data_ref, witness_info=None) if self.args.debug_rollout_only: return self._train_core(rollout_id=rollout_id, rollout_data=rollout_data) diff --git a/miles/backends/megatron_utils/actor.py b/miles/backends/megatron_utils/actor.py index 784a89e8b63..4535846f637 100644 --- a/miles/backends/megatron_utils/actor.py +++ b/miles/backends/megatron_utils/actor.py @@ -25,6 +25,7 @@ from miles.utils.timer import Timer, inverse_timer, timer from miles.utils.tracking_utils import init_tracking from miles.utils.types import RolloutBatch +from miles.utils.witness.allocator import WitnessInfo from ...utils.profile_utils import TrainProfiler from ...utils.tensor_backper import TensorBackuper @@ -273,13 +274,19 @@ def compute_log_prob( store_prefix=store_prefix, ) - def train(self, rollout_id: int, rollout_data_ref: Box) -> None: + def train( + self, + rollout_id: int, + rollout_data_ref: Box, + witness_info: WitnessInfo | None, + attempt: int, + ) -> None: self._last_rollout_id = rollout_id if self.args.offload_train: self.wake_up() with timer("data_preprocess"): - rollout_data = get_rollout_data(self.args, rollout_data_ref) + rollout_data = get_rollout_data(self.args, rollout_data_ref, witness_info=witness_info) if self.args.debug_rollout_only: log_rollout_data(rollout_id, self.args, rollout_data) return @@ -287,7 +294,7 @@ def train(self, rollout_id: int, rollout_data_ref: Box) -> None: if self.role == "critic": return self.train_critic(rollout_id, rollout_data) else: - return self.train_actor(rollout_id, rollout_data) + return self.train_actor(rollout_id, rollout_data, witness_info=witness_info, attempt=attempt) def train_critic(self, rollout_id: int, rollout_data: RolloutBatch) -> None: # Create data iterator for log_probs and train. @@ -315,12 +322,16 @@ def train_critic(self, rollout_id: int, rollout_data: RolloutBatch) -> None: self.opt_param_scheduler, data_iterator, num_microbatches, + witness_info=None, + attempt=0, ) def _use_rollout_replay(self, m) -> bool: return getattr(self.args, f"use_rollout_{m.name}_replay", False) - def train_actor(self, rollout_id: int, rollout_data: RolloutBatch) -> None: + def train_actor( + self, rollout_id: int, rollout_data: RolloutBatch, *, witness_info: WitnessInfo | None, attempt: int + ) -> None: # Create data iterator for log_probs and train. data_iterator, num_microbatches = get_data_iterator(self.args, self.model, rollout_data) @@ -409,6 +420,8 @@ def train_actor(self, rollout_id: int, rollout_data: RolloutBatch) -> None: self.opt_param_scheduler, data_iterator, num_microbatches, + witness_info=witness_info, + attempt=attempt, ) self.prof.step(rollout_id=rollout_id) diff --git a/miles/backends/megatron_utils/model.py b/miles/backends/megatron_utils/model.py index 2b3b85a5a9d..a4e91c91c65 100644 --- a/miles/backends/megatron_utils/model.py +++ b/miles/backends/megatron_utils/model.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import dataclasses import gc import logging @@ -24,6 +26,7 @@ from miles.utils.dumper_utils import DumperMegatronUtil, DumperPhase from miles.utils.memory_utils import clear_memory +from miles.utils.witness.allocator import WitnessInfo from ..training_utils.ci_utils import check_grad_norm, check_kl from ..training_utils.data import DataIterator, get_batch @@ -347,6 +350,8 @@ def train_one_step( optimizer: MegatronOptimizer | None, opt_param_scheduler: OptimizerParamScheduler | None, num_microbatches: int, + witness_info: WitnessInfo | None, + attempt: int, ) -> tuple[dict[str, float], float]: """Execute a single pipeline-parallel training step. @@ -537,6 +542,8 @@ def train( opt_param_scheduler: OptimizerParamScheduler | None, data_iterator: Sequence[DataIterator], num_microbatches: Sequence[int], + witness_info: WitnessInfo | None, + attempt: int, ) -> None: """Run training over a rollout consisting of multiple steps. @@ -632,6 +639,8 @@ def train( optimizer, opt_param_scheduler, num_microbatches[step_id], + witness_info=witness_info, + attempt=attempt, ) if step_id == 0: diff --git a/miles/backends/training_utils/data.py b/miles/backends/training_utils/data.py index 71aa4113959..dc9f1c33068 100644 --- a/miles/backends/training_utils/data.py +++ b/miles/backends/training_utils/data.py @@ -9,6 +9,7 @@ from miles.utils.data import get_minimum_num_micro_batch_size from miles.utils.seqlen_balancing import get_seqlen_balanced_partitions from miles.utils.types import RolloutBatch +from miles.utils.witness.allocator import WitnessInfo from ...utils.data import process_rollout_data from ...utils.ray_utils import Box @@ -28,7 +29,11 @@ def _rollout_logprob_dtype(args: Namespace) -> torch.dtype: return torch.float32 -def get_rollout_data(args: Namespace, rollout_data_ref: Box) -> RolloutBatch: +def get_rollout_data( + args: Namespace, + rollout_data_ref: Box, + witness_info: WitnessInfo | None = None, +) -> RolloutBatch: parallel_state = get_parallel_state() # Fetch data through ray on CPU, not sure if this will be performance bottleneck. # Both first pp stage and the last pp stage will receive the data. @@ -37,6 +42,7 @@ def get_rollout_data(args: Namespace, rollout_data_ref: Box) -> RolloutBatch: rollout_data_ref, parallel_state.intra_dp.rank, parallel_state.intra_dp.size, + witness_info=witness_info, ) # move tokens to GPU in advance rollout_data["tokens"] = [ @@ -45,6 +51,13 @@ def get_rollout_data(args: Namespace, rollout_data_ref: Box) -> RolloutBatch: rollout_data["loss_masks"] = [ torch.tensor(t, dtype=torch.int, device=torch.cuda.current_device()) for t in rollout_data["loss_masks"] ] + if args.enable_witness: + seq_witness_ids = rollout_data.pop("seq_witness_ids") + rollout_data["witness_ids"] = [ + torch.full((len(t),), fill_value=sid, dtype=torch.long, device=torch.cuda.current_device()) + for t, sid in zip(rollout_data["tokens"], seq_witness_ids, strict=True) + ] + if "multimodal_train_inputs" in rollout_data: # Move multimodal training tensors to GPU in advance rollout_data["multimodal_train_inputs"] = [ @@ -239,6 +252,9 @@ def _compute_transform_like_token_ids(ids_list: list): batch["position_ids"] = _compute_transform_like_token_ids(position_ids_list) + if (witness_ids := batch.get("witness_ids")) is not None: + batch["witness_ids"] = _compute_transform_like_token_ids(witness_ids) + # loss masks loss_masks = [] for loss_mask, total_length, response_length in zip( diff --git a/miles/backends/training_utils/log_utils.py b/miles/backends/training_utils/log_utils.py index 0ea3b538b13..c1f7a65802f 100644 --- a/miles/backends/training_utils/log_utils.py +++ b/miles/backends/training_utils/log_utils.py @@ -122,6 +122,7 @@ def log_rollout_data(rollout_id: int, args: Namespace, rollout_data: RolloutBatc "rollout_indexer_topk", "max_seq_lens", "dynamic_global_batch_size", + "witness_ids", "weight_versions", "metadata", ]: diff --git a/miles/ray/actor_group.py b/miles/ray/actor_group.py index 1d3539338af..752bd208811 100644 --- a/miles/ray/actor_group.py +++ b/miles/ray/actor_group.py @@ -76,7 +76,13 @@ async def init(self): async def train(self, rollout_id, rollout_data_pack): """Do one rollout training""" - await self._broadcast("train", rollout_id, rollout_data_pack["data_ref"]) + await self._broadcast( + "train", + rollout_id, + rollout_data_pack["data_ref"], + witness_info=None, + attempt=0, + ) async def save_model(self, rollout_id, force_sync=False): """Save actor model""" diff --git a/miles/utils/data.py b/miles/utils/data.py index 21d01db972b..e1484ccf013 100644 --- a/miles/utils/data.py +++ b/miles/utils/data.py @@ -9,6 +9,7 @@ import ray from miles.ray.rollout.train_data_conversion import split_train_data_by_dp_raw +from .witness.allocator import WitnessInfo try: import pyarrow.parquet as pq @@ -279,13 +280,17 @@ def process_rollout_data( rollout_data_ref, dp_rank, dp_size, + witness_info: WitnessInfo | None, ): if args.delay_split_train_data_by_dp: raw = ray.get(rollout_data_ref.inner) + if (x := witness_info) is not None: + raw = {**raw, "seq_witness_ids": x.witness_ids} raw = split_train_data_by_dp_raw(args, raw, dp_size=dp_size) rollout_data = raw[dp_rank] else: assert len(rollout_data_ref) == dp_size + assert witness_info is None rollout_data = ray.get(rollout_data_ref[dp_rank].inner) partition = rollout_data.pop("partition") From c5f3c254c0deedf33968f63ff573b1c6550044a4 Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Mon, 22 Jun 2026 18:15:18 +0800 Subject: [PATCH 41/41] Make the tensor dumper fault-tolerance aware Scope tensor dumps per rollout and make them safe under independent-DP cells: dumps go into per-rollout subdirectories, only effective-DP rank 0 writes output files (other ranks still join dumper collectives), weights/grads are dumped once per rollout pinned to step 0 with a distributed-optimizer grad all-gather, and dump-dir cleanup is resilient to crashed peers and barriers across cells. - dumper_utils.py: per-rollout exp_name, rank-gated output, `_build_full_grad_getter`, resilient `_cleanup_dump_dir` + `_barrier_after_dump_dir_cleanup`. - model.py / actor.py: thread `rollout_id` into `forward_only` / `compute_log_prob` and the DumperMegatronUtil construction. --- miles/backends/megatron_utils/actor.py | 6 ++ miles/backends/megatron_utils/model.py | 9 +- miles/utils/dumper_utils.py | 114 +++++++++++++++++++++++-- tests/e2e/conftest_dumper.py | 2 +- 4 files changed, 121 insertions(+), 10 deletions(-) diff --git a/miles/backends/megatron_utils/actor.py b/miles/backends/megatron_utils/actor.py index 4535846f637..76e50afe16c 100644 --- a/miles/backends/megatron_utils/actor.py +++ b/miles/backends/megatron_utils/actor.py @@ -261,6 +261,7 @@ def compute_log_prob( self, data_iterator: list[DataIterator], num_microbatches: list[int], + rollout_id: int, store_prefix: str = "", ) -> dict[str, list[torch.Tensor]]: @@ -271,6 +272,7 @@ def compute_log_prob( self.model, data_iterator, num_microbatches, + rollout_id=rollout_id, store_prefix=store_prefix, ) @@ -306,6 +308,7 @@ def train_critic(self, rollout_id: int, rollout_data: RolloutBatch) -> None: self.model, data_iterator, num_microbatches, + rollout_id=rollout_id, ) ) @@ -359,6 +362,7 @@ def train_actor( self.compute_log_prob( data_iterator, num_microbatches, + rollout_id=rollout_id, store_prefix="ref_", ) ) @@ -370,6 +374,7 @@ def train_actor( self.compute_log_prob( data_iterator, num_microbatches, + rollout_id=rollout_id, store_prefix="teacher_", ) ) @@ -385,6 +390,7 @@ def train_actor( self.compute_log_prob( data_iterator, num_microbatches, + rollout_id=rollout_id, store_prefix="", ) ) diff --git a/miles/backends/megatron_utils/model.py b/miles/backends/megatron_utils/model.py index a4e91c91c65..dbd3c6e566f 100644 --- a/miles/backends/megatron_utils/model.py +++ b/miles/backends/megatron_utils/model.py @@ -213,6 +213,7 @@ def forward_only( model: Sequence[DDP], data_iterator: Sequence[DataIterator], num_microbatches: Sequence[int], + rollout_id: int, store_prefix: str = "", ) -> dict[str, list[torch.Tensor]]: """Run forward passes only and collect non-loss outputs (e.g., logprobs). @@ -226,13 +227,16 @@ def forward_only( model: Sequence of DDP-wrapped model chunks. data_iterator: Iterable(s) yielding batches for inference. num_microbatches: Number of microbatches per rollout step. + rollout_id: Rollout identifier (selects the per-rollout dump subdirectory). store_prefix: Prefix to prepend to stored output keys. Returns: Aggregated outputs keyed by ``store_prefix + key``. """ - dumper_phase_util = DumperMegatronUtil(args, model, DumperPhase.FWD_ONLY) + dumper_phase_util = DumperMegatronUtil( + args, model, DumperPhase.FWD_ONLY, rollout_id=rollout_id, store_prefix=store_prefix + ) # reset data iterator for iterator in data_iterator: @@ -278,6 +282,7 @@ def forward_step( packed_seq_params = get_packed_seq_params(batch, args) total_lengths = batch["total_lengths"] response_lengths = batch["response_lengths"] + output_tensor = model( input_ids=tokens, position_ids=None, @@ -372,7 +377,7 @@ def train_one_step( Reduced loss dictionary (last stage only) and gradient norm for logging. """ args = get_args() - dumper_phase_util = DumperMegatronUtil(args, model, DumperPhase.FWD_BWD) + dumper_phase_util = DumperMegatronUtil(args, model, DumperPhase.FWD_BWD, rollout_id=rollout_id) disable_optimizer = args.debug_disable_optimizer or optimizer is None # Set grad to zero. diff --git a/miles/utils/dumper_utils.py b/miles/utils/dumper_utils.py index aea925aa6e3..d0bb59a6e2c 100644 --- a/miles/utils/dumper_utils.py +++ b/miles/utils/dumper_utils.py @@ -14,6 +14,9 @@ import torch.distributed as dist from sglang.srt.debug_utils.dumper import DumperConfig, _get_rank, dumper +from miles.backends.training_utils.parallel import get_parallel_state +from miles.utils.environ import enable_experimental_ft_trainer + logger = logging.getLogger(__name__) @@ -78,10 +81,21 @@ async def configure_sglang(args: Namespace) -> None: class DumperMegatronUtil: - def __init__(self, args: Namespace, model: Sequence[torch.nn.Module], phase: DumperPhase) -> None: + def __init__( + self, + args: Namespace, + model: Sequence[torch.nn.Module], + phase: DumperPhase, + *, + rollout_id: int, + store_prefix: str = "", + ) -> None: self.phase = phase + self.rollout_id = rollout_id self.overrides = _get_phase_override_configs(args, phase) - self.enabled = self._configure(args, phase, self.overrides) + self.enabled = self._configure( + args, phase=phase, rollout_id=rollout_id, store_prefix=store_prefix, overrides=self.overrides + ) if self.enabled: dumper.register_non_intrusive_dumper(self._extract_model(model)) @@ -96,10 +110,15 @@ def finalize(self, model: Sequence[torch.nn.Module]) -> None: return extracted_model = self._extract_model(model) + 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) - dumper.dump_model(extracted_model) + # Weights/grads are a once-per-rollout end-state, so pin them to step 0 instead of + # the running per-microbatch step. + dumper.dump_model(extracted_model, get_grad=get_grad, step=0) dumper.step() dumper.configure(enable=False) @@ -111,25 +130,97 @@ def _extract_model(model: Sequence[torch.nn.Module]) -> torch.nn.Module: return model[0] @staticmethod - def _configure(args: Namespace, phase: DumperPhase, overrides: dict[str, Any] | None = None) -> bool: + def _configure( + args: Namespace, + *, + phase: DumperPhase, + rollout_id: int, + store_prefix: str = "", + overrides: dict[str, Any] | None = None, + ) -> bool: if overrides is None: overrides = _get_phase_override_configs(args, phase) if not overrides.get("enable"): return False + exp_name = f"{phase.value}/{store_prefix}rollout_{rollout_id}" merged = { "dir": str(_get_dir(args)), - "exp_name": phase.value, + "exp_name": exp_name, + "enable_output_console": False, **overrides, } + # Only write dump files on effective DP rank 0 (covers both intra-DP + # and indep-DP). Other DP ranks still participate in dumper collectives + # (barrier, broadcast, allgather) but don't produce output files. + # TODO: optimize — non-DP-rank-0 ranks currently run full dumper logic + # (forward hooks, model iteration) without producing output. + if get_parallel_state().intra_dp.rank != 0: + merged["enable_output_file"] = False + merged["enable_output_console"] = False + full_config = DumperConfig(**merged) dumper.reset() + # Wipe the whole phase dir only at run start (rollout 0). Gating on a + # per-process latch instead would make a respawned process re-wipe the + # phase dir mid-run, deleting dumps already written by surviving cells. + if rollout_id == 0: + _cleanup_dump_dir(Path(merged["dir"]) / phase.value) _cleanup_dump_dir(Path(merged["dir"]) / merged["exp_name"]) + _barrier_after_dump_dir_cleanup() dumper.configure(**dataclasses.asdict(full_config)) return True +def _build_full_grad_getter( + model_chunk: torch.nn.Module, +) -> Callable[[torch.nn.Parameter], torch.Tensor | None]: + """Build get_grad(param): all-gather distributed-optimizer grad shards into a + fresh buffer (grad_data is read, not mutated) and return per-param views.""" + grad_map: dict[torch.nn.Parameter, torch.Tensor] = {} + # Bucket iteration copied from indep_dp.allreduce_grads_and_losses_across_replicas, + # which cross-cell all-reduces these same bucket.grad_data buffers. + bucket_groups = list(getattr(model_chunk, "bucket_groups", [])) + list( + getattr(model_chunk, "expert_parallel_bucket_groups", []) + ) + for bucket_group in bucket_groups: + if not bucket_group.ddp_config.use_distributed_optimizer: + continue + # Same group/size/rank Megatron's grad reduce-scatter uses + # (Megatron-LM param_and_grad_buffer.py _ParamAndGradBucketGroup.start_grad_sync). + group = bucket_group.intra_distributed_optimizer_instance_group + instance_size = bucket_group.intra_distributed_optimizer_instance_size + instance_rank = bucket_group.intra_distributed_optimizer_instance_rank + for bucket in bucket_group.buckets: + grad_data = bucket.grad_data + if instance_size > 1: + full = torch.empty_like(grad_data) + # shard slicing copied from Megatron shard_buffer(); local_shard is + # this rank's owned (reduce-scattered) slice. + shard_numel = grad_data.numel() // instance_size + local_shard = grad_data[instance_rank * shard_numel : (instance_rank + 1) * shard_numel] + # all-gather copied from Megatron start_param_sync (it does this on + # bucket.param_data); here on grad, into a fresh buffer (grad_data read-only). + dist.all_gather_into_tensor(full, local_shard.contiguous(), group=group) + else: + full = grad_data + flat = full.view(-1) + # per-param slice copied from Megatron's own bucket.param_data.view(-1) + # [start:end].view(shape), using the bucket-local bucket.param_to_index. + for param, (start, end) in bucket.param_to_index.items(): + grad_map[param] = flat[start:end].view(param.shape) + + def get_grad(param: torch.nn.Parameter) -> torch.Tensor | None: + reduced = grad_map.get(param) + if reduced is not None: + return reduced + # fallback copied from sglang dumper's original grad read (.grad else main_grad). + return param.grad if param.grad is not None else getattr(param, "main_grad", None) + + return get_grad + + def _log_model_grad_coverage(model: torch.nn.Module) -> None: missing: list[str] = [] with_grad = 0 @@ -174,8 +265,17 @@ def _wrapped(*args: Any, **kwargs: Any) -> Any: def _cleanup_dump_dir(dump_dir: Path) -> None: - if _get_rank() == 0 and dump_dir.is_dir(): - shutil.rmtree(dump_dir) + # Best-effort: stale handles (NFS .nfsXXXX stubs) can make rmtree fail with + # "Directory not empty"; we don't want that to propagate up and mark the cell + # as errored. + if (_get_rank() == 0) and dump_dir.is_dir(): + try: + shutil.rmtree(dump_dir) + except OSError: + logger.warning("dump dir cleanup failed; continuing", exc_info=True) + + +def _barrier_after_dump_dir_cleanup() -> None: if dist.is_initialized(): dist.barrier() diff --git a/tests/e2e/conftest_dumper.py b/tests/e2e/conftest_dumper.py index dbe4322a849..dfaa1f09981 100644 --- a/tests/e2e/conftest_dumper.py +++ b/tests/e2e/conftest_dumper.py @@ -195,7 +195,7 @@ def check_dump_dir( assert phase_dir.exists(), f"Missing dump dir: {phase_dir}" dump_subdirs: list[Path] = list(phase_dir.glob(exp_pattern)) assert len(dump_subdirs) > 0, f"No {exp_pattern} subdirs in {phase_dir}" - dump_files: list[Path] = list(dump_subdirs[0].glob("*.pt")) + dump_files: list[Path] = list(dump_subdirs[0].rglob("*.pt")) assert len(dump_files) > 0, f"No .pt files in {dump_subdirs[0]}" sample: dict = torch.load(dump_files[0], weights_only=False) assert isinstance(sample, dict), f"Unexpected type: {type(sample)}"