From 96eecb4acae9e76df9552913eab4f87c89bf0290 Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Mon, 22 Jun 2026 18:15:16 +0800 Subject: [PATCH 01/30] 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 6264e17c30..962c421755 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 a3dadbdaf0..eb97afa8cc 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/30] 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 13b8f2b5cb..b3c7f8f2f1 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 4221e24c6c..2fe6b8e995 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/30] 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 3ecf15eab0..5837358bb7 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 35d1f350ee..1da32045b7 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 d016e01ac3..7489d89ec9 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 313a60dda9..3c160cc5cb 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 0000000000..21eee55fb8 --- /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 810c2b67c7..1cc7710786 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/30] 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 d9f5412953..295dd10e99 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/30] 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 0000000000..9a15742ea6 --- /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 0000000000..64a5f1e1f0 --- /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 0000000000..4055add45b --- /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 0000000000..2b58b1b864 --- /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 0000000000..43e9763b8e --- /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 0000000000..ef98f6ea89 --- /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 0000000000..f5e0801c83 --- /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 0000000000..e3eb15f843 --- /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/30] 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 0000000000..b74849afbd --- /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 b346d372fd..0bfbc713d5 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 0000000000..c712ea6e45 --- /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/30] 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 0000000000..73986d8d2e --- /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 0000000000..7837f6d5f3 --- /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/30] 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 0000000000..35362597f0 --- /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/30] 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 0000000000..e69de29bb2 diff --git a/miles/utils/control_server/models.py b/miles/utils/control_server/models.py new file mode 100644 index 0000000000..00e47aadac --- /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 0000000000..e69de29bb2 From e308821cf0580bd62e637816319951d6f3a4e536 Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Mon, 22 Jun 2026 18:15:16 +0800 Subject: [PATCH 10/30] 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 0000000000..4d87435d06 --- /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 0000000000..9ae64fcebe --- /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 0000000000..c403d7d78f --- /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 0000000000..2fd567960b --- /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/30] 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 7d7858b91f..49d99898c7 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 22b5dfcbb9..26646a1e11 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 de99c167f3..7e5e94fdbd 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/30] 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 679d31d4c9..f5ba0ae5ae 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/30] 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 f729d948ae..f1436cfafa 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 9e6659026f..7ab2bfcd74 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/30] 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 f1436cfafa..bfbfdf55e8 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 c216b440b2..107d771233 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 65bc8d4b6d..e67087fbb3 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 49d99898c7..7ff03b7787 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 0fe7ad483f..21d01db972 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 ab7d79b75d..962a53099d 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 4a3da1a288..31f3b1975d 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/30] 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 00b64c1569..d47859eca8 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 7ab2bfcd74..971d8a3a26 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 7ff03b7787..2348dfa1e5 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 359a51c1d6..cdfbaf956e 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 0000000000..a2fafddf2e --- /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 0000000000..6e123862e1 --- /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/30] 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 0000000000..633a9ca2ad --- /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 0000000000..76327549b6 --- /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/30] 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 0000000000..e69de29bb2 diff --git a/miles/utils/event_logger/models.py b/miles/utils/event_logger/models.py new file mode 100644 index 0000000000..c9a136c730 --- /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 0000000000..e69de29bb2 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 0000000000..2cce486c2a --- /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/30] 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 107d771233..f853f99e61 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 971d8a3a26..f91d7681a7 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 281e50fdf2..efd73558ad 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 6b465f3c2b..6013535ddc 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 2348dfa1e5..c48a8a215d 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 0000000000..6155bd8d0f --- /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 c9a136c730..0749f72027 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 621e60532e..9b27fdb613 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 b6ae346b3a..fb700179a8 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 0000000000..32c82c4f16 --- /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 54f048a236..c61a36dd24 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 af29662258..e380b087c3 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 a7e0063850..85cb7aeece 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/30] 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 c48a8a215d..24f831ccb7 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 0000000000..a4c838ec0f --- /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 0000000000..b0d9123001 --- /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/30] 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 07e7023be4..9fb867da32 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/30] 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 24f831ccb7..2f9595f9f1 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 0000000000..e69de29bb2 diff --git a/miles/utils/witness/allocator.py b/miles/utils/witness/allocator.py new file mode 100644 index 0000000000..8cfcd1ae04 --- /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 0000000000..e69de29bb2 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 0000000000..8cb671064f --- /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/30] 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 6aa694e876..16249147e7 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 dc483e5e54..2e289acbc6 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 0000000000..c012f26b64 --- /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 0000000000..35b7133ed5 --- /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/30] 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 0000000000..e69de29bb2 diff --git a/miles/utils/event_analyzer/rules/__init__.py b/miles/utils/event_analyzer/rules/__init__.py new file mode 100644 index 0000000000..8b13789179 --- /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 0000000000..bf4cba9f13 --- /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 0000000000..8c6bddcc62 --- /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 0000000000..be49daee6f --- /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 0000000000..e69de29bb2 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 0000000000..8b13789179 --- /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 0000000000..48f97f79dd --- /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 0000000000..456d25faa8 --- /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 0000000000..92e7548f26 --- /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/30] 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 0000000000..f85569e3aa --- /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 0000000000..7a121cc3fd --- /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/30] 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 2f9595f9f1..a1989fb0e7 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 0000000000..e7559f9589 --- /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 0000000000..9058922304 --- /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/30] 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 0000000000..e69de29bb2 diff --git a/miles/utils/test_utils/comparisons/comparators.py b/miles/utils/test_utils/comparisons/comparators.py new file mode 100644 index 0000000000..4051c1804f --- /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 0000000000..84810671b5 --- /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 0000000000..57e5ab3c64 --- /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 0000000000..e69de29bb2 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 0000000000..59716c2c30 --- /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 0000000000..416cfe4078 --- /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/30] 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 0000000000..70721fbf57 --- /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 0000000000..8b1f7a2910 --- /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/30] 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 0000000000..23c6292bd5 --- /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 0000000000..6873de03ea --- /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/30] 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 0f49fd04ff..0bbff9f542 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 38b52ed5a2..f44805e9b8 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 2d8d6e082c..0f8f169ee9 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 692bfaaa56..7bb76b2ed6 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 0000000000..a929893555 --- /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 0000000000..b8a958bb5e --- /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 0000000000..a3624cc7ff --- /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/30] 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 0000000000..ea7a92a8c3 --- /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 0000000000..c5b764b67f --- /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)