diff --git a/src/art/trainer_rank/_impl.py b/src/art/trainer_rank/_impl.py index 984ca77a3..de15bbdee 100644 --- a/src/art/trainer_rank/_impl.py +++ b/src/art/trainer_rank/_impl.py @@ -13,6 +13,7 @@ Sequence, ) from concurrent.futures import Future, ThreadPoolExecutor +from contextlib import contextmanager from copy import deepcopy from dataclasses import dataclass, replace from dataclasses import field as dataclass_field @@ -3856,7 +3857,9 @@ def estimate(width: int) -> tuple[_MemoryCheck, bool, bool] | None: return estimates[width] indices, local_inputs = local_slice(width) local_requests = list(_flatten(local_inputs)) - values = self._estimate_flat_forward(local_requests, checkpoint=checkpoint) + values = self._estimate_flat_forward( + local_requests, checkpoint=checkpoint, sync_planning_errors=True + ) if not self._all_ranks_true(values is not None): estimates[width] = None return None @@ -3868,16 +3871,15 @@ def priced( output_bytes: int, signature: _MemorySignature, ) -> tuple[_MemoryCheck, int, int, _MemorySignature]: + with self._planning_status(True): + required = self._estimate_required_memory_bytes_from_values( + packed_tokens=packed_tokens, + output_bytes=output_bytes, + signature=signature, + logical_tokens=logical_tokens, + ) return ( - self._memory_check_required( - self._estimate_required_memory_bytes_from_values( - packed_tokens=packed_tokens, - output_bytes=output_bytes, - signature=signature, - logical_tokens=logical_tokens, - ), - sync_across_dp=True, - ), + self._memory_check_required(required, sync_across_dp=True), packed_tokens, output_bytes, signature, @@ -3891,6 +3893,7 @@ def priced_estimate( checkpoint=checkpoint, exact=exact, memory_minimal=memory_minimal, + sync_planning_errors=True, ) return None if estimated is None else priced(*estimated) @@ -3964,12 +3967,16 @@ def fits(width: int) -> tuple[bool, bool]: # materialized plan, trying the cost-optimal layouts first and # the memory-minimal layouts if those do not fit. plan = materialize(width) - check = self._memory_check(plan, sync_across_dp=True) + check = self._memory_check( + plan, sync_across_dp=True, sync_planning_errors=True + ) if not check.fits and not layout_modes.get(width, False): layout_modes[width] = True plans.pop(width, None) plan = materialize(width) - check = self._memory_check(plan, sync_across_dp=True) + check = self._memory_check( + plan, sync_across_dp=True, sync_planning_errors=True + ) trusted = self._all_ranks_have_memory_profile( packed_tokens=plan.packed_tokens, signature=plan.signature, @@ -3990,6 +3997,7 @@ def materialize(width: int) -> _FlatForwardPlan: list(_flatten(local_inputs)), checkpoint=checkpoint, memory_minimal=layout_modes.get(width, False), + sync_planning_errors=True, ) plans[width] = plan return plan @@ -4002,7 +4010,9 @@ def candidate(width: int) -> _CandidateMicroBatch[ForwardInputsT]: check = ( estimated[0] if estimated is not None - else self._memory_check(plan, sync_across_dp=True) + else self._memory_check( + plan, sync_across_dp=True, sync_planning_errors=True + ) ) cold_start = not self._all_ranks_have_memory_profile( packed_tokens=plan.packed_tokens, @@ -4448,60 +4458,71 @@ def _plan_flat_forward( checkpoint: AdapterSelection = Unset, memory_minimal: bool = False, ensure_slots: bool = True, + sync_planning_errors: bool = False, ) -> _FlatForwardPlan: - plans: list[_ForwardGroupPlan] = [] - output_bytes = self._estimate_group_request_output_bytes(requests) - logical_tokens = sum(int(request.input_tokens.numel()) for request in requests) - groups = self._group_active_request_indices( - requests, checkpoint=checkpoint, ensure_slots=ensure_slots - ) - selected_max_depth = 0 - for (slot_ref, grad_enabled), group_indices in groups: - items = tuple( - self._forward_item(requests[index]) for index in group_indices - ) - group_input_ids = tuple(item.input_ids for item in items) - tree, layout = self._select_group_layout( - group_input_ids, - memory_minimal=memory_minimal, - grad_enabled=grad_enabled, + with self._planning_status(sync_planning_errors): + plans: list[_ForwardGroupPlan] = [] + output_bytes = self._estimate_group_request_output_bytes(requests) + logical_tokens = sum( + int(request.input_tokens.numel()) for request in requests ) - selected_max_depth = max(selected_max_depth, layout.maximum_depth) - started = time.perf_counter() - packed = materialize_prefix_tree_layout( - group_input_ids, tree, layout, verify_shared_tokens=False + if sync_planning_errors and ensure_slots: + self._ensure_checkpoint_slots_for(requests, checkpoint=checkpoint) + with self._planning_status(sync_planning_errors): + groups = self._group_active_request_indices( + requests, + checkpoint=checkpoint, + ensure_slots=ensure_slots and not sync_planning_errors, ) - self._planning_seconds_accum += time.perf_counter() - started - plans.append( - _ForwardGroupPlan( - slot_ref=slot_ref, + selected_max_depth = 0 + for (slot_ref, grad_enabled), group_indices in groups: + items = tuple( + self._forward_item(requests[index]) for index in group_indices + ) + group_input_ids = tuple(item.input_ids for item in items) + tree, layout = self._select_group_layout( + group_input_ids, + memory_minimal=memory_minimal, grad_enabled=grad_enabled, - request_indices=tuple(group_indices), - items=items, - packed=packed, ) - ) + selected_max_depth = max(selected_max_depth, layout.maximum_depth) + started = time.perf_counter() + packed = materialize_prefix_tree_layout( + group_input_ids, tree, layout, verify_shared_tokens=False + ) + self._planning_seconds_accum += time.perf_counter() - started + plans.append( + _ForwardGroupPlan( + slot_ref=slot_ref, + grad_enabled=grad_enabled, + request_indices=tuple(group_indices), + items=items, + packed=packed, + ) + ) - return _FlatForwardPlan( - request_count=len(requests), - output_metadata=tuple( - self._forward_output_metadata(request, checkpoint=checkpoint) - for request in requests - ), - groups=tuple(plans), - packed_tokens=sum( - self._physical_tokens(int(plan.packed.tokens.numel())) for plan in plans - ), - logical_tokens=logical_tokens, - output_bytes=output_bytes, - signature=self._memory_signature_from_requests( - requests, - slot_group_count=len(plans), - grad_modes=tuple(mode for (_, mode), _ in groups), - ), - selected_max_depth=selected_max_depth, - inactive_logical_tokens=logical_tokens - _active_logical_tokens(requests), - ) + return _FlatForwardPlan( + request_count=len(requests), + output_metadata=tuple( + self._forward_output_metadata(request, checkpoint=checkpoint) + for request in requests + ), + groups=tuple(plans), + packed_tokens=sum( + self._physical_tokens(int(plan.packed.tokens.numel())) + for plan in plans + ), + logical_tokens=logical_tokens, + output_bytes=output_bytes, + signature=self._memory_signature_from_requests( + requests, + slot_group_count=len(plans), + grad_modes=tuple(mode for (_, mode), _ in groups), + ), + selected_max_depth=selected_max_depth, + inactive_logical_tokens=logical_tokens + - _active_logical_tokens(requests), + ) def _estimate_flat_forward( self, @@ -4510,6 +4531,7 @@ def _estimate_flat_forward( checkpoint: AdapterSelection = Unset, exact: bool = False, memory_minimal: bool = False, + sync_planning_errors: bool = False, ) -> tuple[int, int, _MemorySignature] | None: """Estimate packed tokens for width probing. @@ -4523,40 +4545,52 @@ def _estimate_flat_forward( content) and is used only inside the band where those bounds disagree. """ - groups = self._group_active_request_indices(requests, checkpoint=checkpoint) - packed_tokens = 0 - for (_slot, grad_enabled), group_indices in groups: - if exact: - _, layout = self._select_group_layout( - tuple( - requests[index].input_tokens.reshape(-1).to(dtype=torch.long) + if sync_planning_errors: + self._ensure_checkpoint_slots_for(requests, checkpoint=checkpoint) + with self._planning_status(sync_planning_errors): + groups = self._group_active_request_indices( + requests, + checkpoint=checkpoint, + ensure_slots=not sync_planning_errors, + ) + packed_tokens = 0 + for (_slot, grad_enabled), group_indices in groups: + if exact: + _, layout = self._select_group_layout( + tuple( + requests[index] + .input_tokens.reshape(-1) + .to(dtype=torch.long) + for index in group_indices + ), + memory_minimal=memory_minimal, + grad_enabled=grad_enabled, + ) + packed_tokens += self._physical_tokens(layout.packed_tokens) + continue + # Radix depth is bounded by the number of rows, so ``len(group)`` + # is an unlimited-sharing depth for this group; it is a bound for + # estimation, not a sharing policy. + group_packed_tokens = estimate_prefix_tree_packed_tokens( + ( + requests[index].input_tokens.reshape(-1) for index in group_indices ), - memory_minimal=memory_minimal, - grad_enabled=grad_enabled, + max_depth=len(group_indices) if memory_minimal else 0, ) - packed_tokens += self._physical_tokens(layout.packed_tokens) - continue - # Radix depth is bounded by the number of rows, so ``len(group)`` - # is an unlimited-sharing depth for this group; it is a bound for - # estimation, not a sharing policy. - group_packed_tokens = estimate_prefix_tree_packed_tokens( - (requests[index].input_tokens.reshape(-1) for index in group_indices), - max_depth=len(group_indices) if memory_minimal else 0, - ) - if group_packed_tokens is None: - return None - packed_tokens += self._physical_tokens(group_packed_tokens) + if group_packed_tokens is None: + return None + packed_tokens += self._physical_tokens(group_packed_tokens) - return ( - packed_tokens, - self._estimate_group_request_output_bytes(requests), - self._memory_signature_from_requests( - requests, - slot_group_count=len(groups), - grad_modes=tuple(mode for (_, mode), _ in groups), - ), - ) + return ( + packed_tokens, + self._estimate_group_request_output_bytes(requests), + self._memory_signature_from_requests( + requests, + slot_group_count=len(groups), + grad_modes=tuple(mode for (_, mode), _ in groups), + ), + ) def _ensure_checkpoint_slots_for( self, @@ -4976,16 +5010,16 @@ def _memory_check( forward: _FlatForwardPlan, *, sync_across_dp: bool = False, + sync_planning_errors: bool = False, ) -> _MemoryCheck: - return self._memory_check_required( - self._estimate_required_memory_bytes_from_values( + with self._planning_status(sync_planning_errors): + required = self._estimate_required_memory_bytes_from_values( packed_tokens=forward.packed_tokens, output_bytes=forward.output_bytes, signature=forward.signature, logical_tokens=forward.active_logical_tokens, - ), - sync_across_dp=sync_across_dp, - ) + ) + return self._memory_check_required(required, sync_across_dp=sync_across_dp) def _memory_check_required( self, @@ -5134,6 +5168,30 @@ def _all_ranks_have_memory_profile( ) return self._all_ranks_true(local) + @contextmanager + def _planning_status(self, enabled: bool) -> Generator[None, None, None]: + """Exchange pure local planning errors before any later WORLD check.""" + if not enabled: + yield + return + primary: BaseException | None = None + try: + yield + except BaseException as exc: + primary = exc + exchange_error: BaseException | None = None + try: + succeeded = self._all_ranks_true(primary is None) + except BaseException as exc: + exchange_error = exc + # Leave the exchange's except block before raising the original error. + if primary is not None: + raise primary + if exchange_error is not None: + raise exchange_error + if not succeeded: + raise RuntimeError("Local planning failed on another DP rank") + def _all_ranks_true(self, local: bool) -> bool: if not (dist.is_available() and dist.is_initialized()): return local diff --git a/tests/unit/test_trainer_rank_planning_status.py b/tests/unit/test_trainer_rank_planning_status.py new file mode 100644 index 000000000..2054024ee --- /dev/null +++ b/tests/unit/test_trainer_rank_planning_status.py @@ -0,0 +1,210 @@ +"""Pure planning failures must reach every WORLD participant before admission.""" + +from __future__ import annotations + +from datetime import timedelta +from pathlib import Path +import subprocess +import sys +from types import SimpleNamespace + +import pytest + + +@pytest.mark.parametrize("enabled", (False, True)) +@pytest.mark.parametrize("primary_type", (ValueError, KeyboardInterrupt)) +@pytest.mark.parametrize("exchange_fails", (False, True)) +def test_planning_status_preserves_primary_chain( + enabled: bool, primary_type: type[BaseException], exchange_fails: bool +) -> None: + from art.trainer_rank import TrainerRank + + rank = TrainerRank.__new__(TrainerRank) + primary = primary_type("local planner failed") + cause, context = LookupError("cause"), KeyError("context") + primary.__cause__, primary.__context__ = cause, context + primary.__suppress_context__ = True + calls = [] + + def exchange(succeeded: bool) -> bool: + calls.append(succeeded) + if exchange_fails: + raise OSError("collective failed") + return succeeded + + rank._all_ranks_true = exchange + with pytest.raises(primary_type) as captured: + with rank._planning_status(enabled): + raise primary + assert captured.value is primary + assert primary.__cause__ is cause and primary.__context__ is context + assert primary.__suppress_context__ + assert calls == ([False] if enabled else []) + + +def test_planning_failures_and_empty_ranks_use_aligned_status(tmp_path: Path) -> None: + children = [] + logs = [] + try: + for rank in range(2): + log = (tmp_path / f"rank-{rank}.log").open("w") + logs.append(log) + children.append( + subprocess.Popen( + [sys.executable, "-B", __file__, str(rank), str(tmp_path)], + stdout=log, + stderr=subprocess.STDOUT, + ) + ) + for rank, child in enumerate(children): + assert child.wait(timeout=60) == 0, ( + tmp_path / f"rank-{rank}.log" + ).read_text() + finally: + for child in children: + if child.poll() is None: + child.terminate() + for child in children: + try: + child.wait(timeout=5) + except subprocess.TimeoutExpired: + child.kill() + child.wait(timeout=5) + for log in logs: + log.close() + + +def _worker(index: int, directory: Path) -> None: + import torch + import torch.distributed as dist + + from art.trainer_rank import ForwardInput, TrainerRank, _impl + from art.trainer_rank._prefix_tree_planner import ( + build_canonical_prefix_tree, + prefix_tree_layout_candidates, + ) + + torch.set_num_threads(1) + dist.init_process_group( + "gloo", + rank=index, + world_size=2, + init_method=f"file://{directory / 'gloo'}", + timeout=timedelta(seconds=10), + ) + try: + for mode in ( + "estimate", + "materialize", + "price", + "empty", + "unequal", + "unavailable", + ): + rank = TrainerRank.__new__(TrainerRank) + rank.device = torch.device("cpu") + rank._planning_seconds_accum = 0.0 + rank._dp_rank_and_size = lambda: (index, 2) + rank._physical_tokens = lambda tokens: tokens + rank._resolve_slot_ref = lambda request, **_: request.no_grad + rank._estimate_group_request_output_bytes = lambda requests: 0 + rank._memory_signature_from_requests = lambda *args, **kwargs: None + rank._forward_item = lambda request: SimpleNamespace( + input_ids=request.input_tokens + ) + rank._forward_output_metadata = lambda *args, **kwargs: (None, True) + + def slots(selections): + assert tuple(selections) == () + dist.barrier() + + rank._ensure_checkpoint_slots = slots + + def layout(ids, **kwargs): + tree = build_canonical_prefix_tree(ids) + return tree, prefix_tree_layout_candidates(tree)[0].layout + + rank._select_group_layout = layout + primary = ValueError("rank-local planning failure") + cause, context = LookupError("cause"), KeyError("context") + primary.__cause__, primary.__context__ = cause, context + primary.__suppress_context__ = True + original_estimate = _impl.estimate_prefix_tree_packed_tokens + original_materialize = _impl.materialize_prefix_tree_layout + + def estimate(*args, **kwargs): + if index == 0 and mode == "estimate": + raise primary + if index == 0 and mode == "unavailable": + return None + return original_estimate(*args, **kwargs) + + def materialize(*args, **kwargs): + if index == 0 and mode == "materialize": + raise primary + return original_materialize(*args, **kwargs) + + def price(**kwargs): + if index == 0 and mode == "price": + raise primary + return 0 + + rank._estimate_required_memory_bytes_from_values = price + patches = pytest.MonkeyPatch() + patches.setattr(_impl, "estimate_prefix_tree_packed_tokens", estimate) + patches.setattr(_impl, "materialize_prefix_tree_layout", materialize) + requests = [ + ForwardInput( + input_tokens=torch.tensor([1, 2]), hidden_states=True, no_grad=True + ) + ] + if mode == "empty" and index == 1: + requests = [] + if mode == "unequal" and index == 0: + requests.append( + ForwardInput( + input_tokens=torch.tensor([3, 4]), + hidden_states=True, + no_grad=False, + ) + ) + try: + error = None + try: + if mode == "estimate": + # Enter the ordinary scheduler, before its first memory check. + rank._select_next_micro_batch(requests * 2, 0) + raise AssertionError("Planner failure unexpectedly returned") + values = rank._estimate_flat_forward( + requests, sync_planning_errors=True + ) + if mode == "unavailable": + assert not rank._all_ranks_true(values is not None) + plan = rank._plan_flat_forward(requests, sync_planning_errors=True) + if mode == "price": + rank._memory_check( + plan, sync_across_dp=True, sync_planning_errors=True + ) + assert plan.request_count == len(requests) + except BaseException as caught: + error = caught + if mode in ("estimate", "materialize", "price"): + if index == 0: + assert error is primary + assert error.__cause__ is cause and error.__context__ is context + else: + assert type(error) is RuntimeError + assert str(error) == "Local planning failed on another DP rank" + else: + assert error is None, repr(error) + finally: + patches.undo() + dist.barrier() + finally: + dist.destroy_process_group() + + +if __name__ == "__main__": + # Direct child execution uses this checkout even without an editable install. + sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "src")) + _worker(int(sys.argv[1]), Path(sys.argv[2])) diff --git a/tests/unit/test_trainer_rank_weird_shapes.py b/tests/unit/test_trainer_rank_weird_shapes.py index 5eecfccb0..8db3ef72e 100644 --- a/tests/unit/test_trainer_rank_weird_shapes.py +++ b/tests/unit/test_trainer_rank_weird_shapes.py @@ -780,7 +780,9 @@ def test_adaptive_planner_globally_falls_back_when_one_rank_cannot_estimate( ) -> None: rank = TrainerRank(_runtime()) monkeypatch.setattr(rank, "_dp_rank_and_size", lambda: (0, 2)) - monkeypatch.setattr(rank, "_all_ranks_true", lambda _local: False) + # Planning succeeds; only estimator availability and profile trust are false. + outcomes = iter((True, False, True, True, True, False)) + monkeypatch.setattr(rank, "_all_ranks_true", lambda _local: next(outcomes)) plans = 0 original = rank._plan_flat_forward @@ -793,6 +795,7 @@ def plan(requests, **kwargs): candidate = rank._select_next_micro_batch( [_target_request(_tokens(index)) for index in range(4)], 0 ) + assert next(outcomes, None) is None assert candidate.stats_global_count == 2 assert plans == 1