From f57c50b3da7208506e195bb667f1ec3ef609a409 Mon Sep 17 00:00:00 2001 From: Codex Date: Tue, 15 Sep 2026 14:32:31 +0000 Subject: [PATCH 1/3] Refresh selected admission and coordinate sampling failures --- src/art/trainer_rank/_impl.py | 81 ++++++- tests/unit/test_trainer_rank_cuda_budget.py | 226 ++++++++++++++++++++ tests/unit/test_trainer_rank_split_peak.py | 16 +- 3 files changed, 311 insertions(+), 12 deletions(-) diff --git a/src/art/trainer_rank/_impl.py b/src/art/trainer_rank/_impl.py index 46ae0e818..984ca77a3 100644 --- a/src/art/trainer_rank/_impl.py +++ b/src/art/trainer_rank/_impl.py @@ -3803,6 +3803,31 @@ def _select_next_micro_batch( start: int, *, checkpoint: AdapterSelection = Unset, + ) -> _CandidateMicroBatch[ForwardInputsT]: + candidate = self._search_next_micro_batch(items, start, checkpoint=checkpoint) + # A later width check may observe less available memory. + # Do not return an earlier width's cached budget as the final admission. + check = self._memory_check_required( + candidate.check.estimated_required_bytes, + sync_across_dp=True, + ) + if not check.fits: + self._snapshot_planning_telemetry(candidate.plan, check) + raise _memory_error( + context="forward_micro_batches", + message="selected microbatch exceeds freshly sampled available memory", + packed_tokens=candidate.plan.packed_tokens, + logical_tokens=candidate.plan.logical_tokens, + check=check, + ) + return replace(candidate, check=check) + + def _search_next_micro_batch( + self, + items: Sequence[ForwardInputsT], + start: int, + *, + checkpoint: AdapterSelection = Unset, ) -> _CandidateMicroBatch[ForwardInputsT]: dp_rank, dp_size = self._dp_rank_and_size() remaining, min_width, granularity = _wave_geometry(len(items), start, dp_size) @@ -4005,12 +4030,27 @@ def candidate(width: int) -> _CandidateMicroBatch[ForwardInputsT]: refusal_prefix = ( "smallest DP microbatch is predicted to exceed available memory" ) - found = self._find_admissible_forward( - list(_flatten(local_inputs)), - checkpoint=checkpoint, - refusal_prefix=refusal_prefix, - ) - agreed = self._all_ranks_true(not isinstance(found, _ForwardRefusal)) + admission_error: BaseException | None = None + try: + found = self._find_admissible_forward( + list(_flatten(local_inputs)), + checkpoint=checkpoint, + refusal_prefix=refusal_prefix, + ) + except BaseException as exc: + admission_error, found = exc, None + try: + agreed = self._all_ranks_true( + admission_error is None + and not isinstance(found, _ForwardRefusal) + ) + except BaseException: + if admission_error is not None: + raise admission_error + raise + if admission_error is not None: + raise admission_error + assert found is not None if isinstance(found, _ForwardRefusal): self._snapshot_planning_telemetry(found.plan, found.check) raise found.error("forward_micro_batches") @@ -4020,7 +4060,8 @@ def candidate(width: int) -> _CandidateMicroBatch[ForwardInputsT]: context="forward_micro_batches", message=( f"{refusal_prefix} on another DP rank, which was " - "unable to find a feasible split for its share" + "unable to complete admission or find a feasible split " + "for its share" ), packed_tokens=first.plan.packed_tokens, logical_tokens=first.plan.logical_tokens, @@ -4952,18 +4993,36 @@ def _memory_check_required( *, sync_across_dp: bool = False, ) -> _MemoryCheck: - available = self._available_memory_bytes() if dist.is_available() and dist.is_initialized(): group = None if sync_across_dp else self._forward_memory_group() values = torch.tensor( - [float(required), float(available)], + [float(required), 0.0], device=self.device if self.device.type == "cuda" else "cpu", dtype=torch.float64, ) dist.all_reduce(values[0], op=dist.ReduceOp.MAX, group=group) - dist.all_reduce(values[1], op=dist.ReduceOp.MIN, group=group) required = int(values[0].item()) - available = int(values[1].item()) + error: BaseException | None = None + try: + available = self._available_memory_bytes() + except BaseException as exc: + error, available = exc, -1 + try: + # A healthy communicator carries local failure to every peer + # in the existing MIN. This cannot repair a poisoned backend. + values[1] = available + dist.all_reduce(values[1], op=dist.ReduceOp.MIN, group=group) + available = int(values[1].item()) + except BaseException: + if error is not None: + raise error + raise + if error is not None: + raise error + if available < 0: + raise RuntimeError("Memory admission failed on another rank") + else: + available = self._available_memory_bytes() return _MemoryCheck( estimated_required_bytes=required, available_bytes=available, diff --git a/tests/unit/test_trainer_rank_cuda_budget.py b/tests/unit/test_trainer_rank_cuda_budget.py index 3575fe5e2..a5da7d99a 100644 --- a/tests/unit/test_trainer_rank_cuda_budget.py +++ b/tests/unit/test_trainer_rank_cuda_budget.py @@ -1,5 +1,8 @@ """Native allocator budget contracts; CPU counters, not CUDA qualification.""" +import asyncio +from contextlib import nullcontext +from types import SimpleNamespace from unittest.mock import Mock import pytest @@ -178,3 +181,226 @@ def reduce(value, op, group): 110, False, ) + + +@pytest.mark.parametrize("failed_locally", [False, True]) +@pytest.mark.parametrize( + "error_type", [RuntimeError, KeyboardInterrupt, SystemExit, asyncio.CancelledError] +) +def test_admission_failure_sentinel_stops_healthy_peers( + budget, monkeypatch, failed_locally, error_type +): + rank, stats = budget + stats["active_bytes.all.current"] = 80 + monkeypatch.setattr(_impl.dist, "is_available", lambda: True) + monkeypatch.setattr(_impl.dist, "is_initialized", lambda: True) + monkeypatch.setattr(rank, "_forward_memory_group", lambda: None) + tensor = torch.tensor + monkeypatch.setattr( + torch, "tensor", lambda values, **kwargs: tensor(values, dtype=kwargs["dtype"]) + ) + original = error_type("recoverable local budget read failure") + events = [] + + def read(_): + events.append("sample_failure") + raise original + + monkeypatch.setattr(torch.cuda, "mem_get_info", read) + if not failed_locally: + monkeypatch.setattr(torch.cuda, "mem_get_info", lambda _: (200, 1000)) + + def reduce(value, op, group): + events.append((op, float(value.item()))) + value.fill_(150 if op == _impl.dist.ReduceOp.MAX else -1) + + monkeypatch.setattr(_impl.dist, "all_reduce", reduce) + with pytest.raises(error_type if failed_locally else RuntimeError) as caught: + for required in (0, 75): + rank._memory_check_required(required) + assert ( + caught.value is original + if failed_locally + else str(caught.value) == "Memory admission failed on another rank" + ) + assert events == [ + (_impl.dist.ReduceOp.MAX, 0), + *(["sample_failure"] if failed_locally else []), + (_impl.dist.ReduceOp.MIN, -1 if failed_locally else 170), + ] + + +@pytest.mark.parametrize("stage", ["assignment", "min", "result"]) +@pytest.mark.parametrize("failed_locally", [False, True]) +def test_secondary_collective_failure_preserves_local_primary( + budget, monkeypatch, stage, failed_locally +): + rank, stats = budget + stats["active_bytes.all.current"] = 80 + monkeypatch.setattr(_impl.dist, "is_available", lambda: True) + monkeypatch.setattr(_impl.dist, "is_initialized", lambda: True) + monkeypatch.setattr(rank, "_forward_memory_group", lambda: None) + original = RuntimeError("original local budget read error") + secondary = RuntimeError("secondary CUDA or communicator failure") + + class Values: + def __init__(self, values, **kwargs): + self.values = values + + def __getitem__(self, index): + def item(): + if stage == "result" and index == 1: + raise secondary + return self.values[index] + + return SimpleNamespace( + item=item, fill_=lambda x: self.values.__setitem__(index, x) + ) + + def __setitem__(self, index, value): + if stage == "assignment": + raise secondary + self.values[index] = value + + monkeypatch.setattr(torch, "tensor", Values) + monkeypatch.setattr(torch.cuda, "mem_get_info", Mock(side_effect=original)) + if not failed_locally: + monkeypatch.setattr(torch.cuda, "mem_get_info", lambda _: (200, 1000)) + + def reduce(value, op, group): + if op == _impl.dist.ReduceOp.MAX: + value.fill_(150) + elif stage == "min": + raise secondary + + monkeypatch.setattr(_impl.dist, "all_reduce", reduce) + with pytest.raises(RuntimeError) as caught: + rank._memory_check_required(0) + assert caught.value is (original if failed_locally else secondary) + + +@pytest.mark.parametrize("distributed", [False, True]) +@pytest.mark.parametrize("available", [10, 210]) +def test_final_selection_uses_pure_fresh_budget_and_original_demand( + budget, monkeypatch, distributed, available +): + rank, stats = budget + stats["active_bytes.all.current"] = 80 + plan = SimpleNamespace(packed_tokens=64, logical_tokens=64) + stale = _impl._MemoryCheck( + estimated_required_bytes=192, available_bytes=256, fits=True + ) + candidate = _impl._CandidateMicroBatch([], (), plan, stale, 64, 1, False) + monkeypatch.setattr(rank, "_search_next_micro_batch", lambda *a, **k: candidate) + monkeypatch.setattr( + rank, + "_estimate_required_memory_bytes_from_values", + Mock(side_effect=AssertionError("must keep original selected demand")), + ) + monkeypatch.setattr(rank, "_snapshot_planning_telemetry", Mock()) + monkeypatch.setattr( + torch.cuda, "mem_get_info", Mock(return_value=(available + 30, 1000)) + ) + monkeypatch.setattr(_impl.dist, "is_available", lambda: distributed) + monkeypatch.setattr(_impl.dist, "is_initialized", lambda: distributed) + tensor = torch.tensor + monkeypatch.setattr( + torch, "tensor", lambda values, **kwargs: tensor(values, dtype=kwargs["dtype"]) + ) + calls = [] + monkeypatch.setattr( + _impl.dist, + "all_reduce", + lambda value, op, group: calls.append((op, value.item(), group)), + ) + if available < 192: + with pytest.raises(_impl.TrainerRankMemoryError) as caught: + rank._select_next_micro_batch([], 0) + assert caught.value.predicted_peak_bytes == 192 + assert caught.value.usable_limit_bytes == available + else: + selected = rank._select_next_micro_batch([], 0) + assert selected.check.estimated_required_bytes == 192 + assert selected.check.available_bytes == available and selected.check.fits + assert selected.plan is plan + assert stale.available_bytes == 256 + assert calls == ( + [ + (_impl.dist.ReduceOp.MAX, 192, None), + (_impl.dist.ReduceOp.MIN, available, None), + ] + if distributed + else [] + ) + torch.cuda.empty_cache.assert_not_called() + torch.cuda.mem_get_info.assert_called_once_with(rank.device) + + +def test_available_sample_follows_required_collective(budget, monkeypatch): + rank, stats = budget + stats["active_bytes.all.current"] = 80 + free = [100] + events = [] + + def read(_): + events.append("sample") + return free[0], 1000 + + monkeypatch.setattr(torch.cuda, "mem_get_info", read) + monkeypatch.setattr(_impl.dist, "is_available", lambda: True) + monkeypatch.setattr(_impl.dist, "is_initialized", lambda: True) + tensor = torch.tensor + monkeypatch.setattr( + torch, "tensor", lambda values, **kwargs: tensor(values, dtype=kwargs["dtype"]) + ) + + def reduce(value, op, group): + events.append(op) + if op == _impl.dist.ReduceOp.MAX: + value.fill_(60) + free[0] = 50 + + monkeypatch.setattr(_impl.dist, "all_reduce", reduce) + check = rank._memory_check_required(0, sync_across_dp=True) + assert (check.estimated_required_bytes, check.available_bytes, check.fits) == ( + 60, + 20, + False, + ) + assert events == [_impl.dist.ReduceOp.MAX, "sample", _impl.dist.ReduceOp.MIN] + torch.cuda.empty_cache.assert_not_called() + + +def test_execution_failure_retains_final_admission_without_resampling( + budget, monkeypatch +): + rank, stats = budget + stats["active_bytes.all.current"] = 80 + free = [100] + read = Mock(side_effect=lambda _: (free[0], 1000)) + monkeypatch.setattr(torch.cuda, "mem_get_info", read) + admitted = rank._memory_check_required(50) + assert admitted.fits and admitted.available_bytes == 70 + original = torch.cuda.OutOfMemoryError("synthetic later execution failure") + + def execute(_): + free[0] = 10 + raise original + + monkeypatch.setattr(rank, "_execute_flat_plan", execute) + monkeypatch.setattr(rank, "_telemetry_signature", lambda _: {}) + monkeypatch.setattr(rank, "_telemetry_plan_signature", lambda _: {}) + monkeypatch.setattr(_impl, "_telemetry_phase", lambda *a, **k: nullcontext()) + monkeypatch.setattr(torch.cuda, "synchronize", Mock()) + monkeypatch.setattr(torch.cuda, "reset_peak_memory_stats", Mock()) + with pytest.raises(_impl.TrainerRankMemoryError) as caught: + rank._run_flat_plan_with_memory_tracking( + SimpleNamespace(packed_tokens=1, logical_tokens=1), + check=admitted, + context="CPU final-admission witness", + ) + assert caught.value.__cause__ is original + assert caught.value.usable_limit_bytes == admitted.available_bytes == 70 + read.assert_called_once_with(rank.device) + assert rank._available_memory_bytes() == 0 + torch.cuda.empty_cache.assert_not_called() diff --git a/tests/unit/test_trainer_rank_split_peak.py b/tests/unit/test_trainer_rank_split_peak.py index 3ce5ed547..7eb6cab04 100644 --- a/tests/unit/test_trainer_rank_split_peak.py +++ b/tests/unit/test_trainer_rank_split_peak.py @@ -294,11 +294,21 @@ def profiled(**_): patch.setattr(rank, "_all_ranks_true", agree) patch.setattr(rank, "_all_ranks_have_memory_profile", profiled) + search = rank._search_next_micro_batch + search_finished = [] + + def searched(*args, **kwargs): + result = search(*args, **kwargs) + search_finished.append(True) + return result + + patch.setattr(rank, "_search_next_micro_batch", searched) + def reduce(value, op, group): trace.append(("global" if group is None else "local", str(op))) if group is None: value.fill_( - max(value.item(), 200) + max(value.item(), 100 if search_finished else 200) if op == tr.dist.ReduceOp.MAX else min(value.item(), 100) ) @@ -314,6 +324,10 @@ def reduce(value, op, group): tr._SplitForwardPlan if dp_rank == 0 else tr._FlatForwardPlan, ) traces.append([event for event in trace if event[0] == "global"]) + assert traces[-1][-2:] == [ + ("global", str(tr.dist.ReduceOp.MAX)), + ("global", str(tr.dist.ReduceOp.MIN)), + ] assert traces[0] == traces[1] From f10f66dd8a9d8a8d83ff0c42d683e1e1c3162b85 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Tue, 15 Sep 2026 14:54:39 +0000 Subject: [PATCH 2/3] Type the narrow admission plan test double explicitly --- tests/unit/test_trainer_rank_cuda_budget.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/unit/test_trainer_rank_cuda_budget.py b/tests/unit/test_trainer_rank_cuda_budget.py index a5da7d99a..858ec6102 100644 --- a/tests/unit/test_trainer_rank_cuda_budget.py +++ b/tests/unit/test_trainer_rank_cuda_budget.py @@ -3,6 +3,7 @@ import asyncio from contextlib import nullcontext from types import SimpleNamespace +from typing import cast from unittest.mock import Mock import pytest @@ -286,7 +287,9 @@ def test_final_selection_uses_pure_fresh_budget_and_original_demand( ): rank, stats = budget stats["active_bytes.all.current"] = 80 - plan = SimpleNamespace(packed_tokens=64, logical_tokens=64) + plan = cast( + _impl._FlatForwardPlan, SimpleNamespace(packed_tokens=64, logical_tokens=64) + ) stale = _impl._MemoryCheck( estimated_required_bytes=192, available_bytes=256, fits=True ) From 0f2134cd304a930f2323206cdad15788e91da41f Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Tue, 15 Sep 2026 15:09:25 +0000 Subject: [PATCH 3/3] Expect the final admission refresh on an empty DP rank --- tests/unit/test_trainer_rank_weird_shapes.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/unit/test_trainer_rank_weird_shapes.py b/tests/unit/test_trainer_rank_weird_shapes.py index 24227d654..5eecfccb0 100644 --- a/tests/unit/test_trainer_rank_weird_shapes.py +++ b/tests/unit/test_trainer_rank_weird_shapes.py @@ -615,7 +615,7 @@ def test_minimum_wave_empty_dp_rank_keeps_collective_check_sequence( ) -> None: rank = TrainerRank(_attention_runtime()) monkeypatch.setattr(rank, "_dp_rank_and_size", lambda: (1, 2)) - peer_required = iter((82, 42, 82, 42)) + peer_required = iter((82, 42, 82, 42, 42)) local_checks: list[tuple[int, bool]] = [] def check(required: int, *, sync_across_dp: bool = False) -> _MemoryCheck: @@ -629,7 +629,7 @@ def check(required: int, *, sync_across_dp: bool = False) -> _MemoryCheck: assert candidate.indices == () assert candidate.plan.packed_tokens == 0 assert candidate.check == _MemoryCheck(42, 60, True) - assert local_checks == [(0, True)] * 4 + assert local_checks == [(0, True)] * 4 + [(42, True)] assert next(peer_required, None) is None