diff --git a/src/art/trainer_rank/_impl.py b/src/art/trainer_rank/_impl.py index c62d8375c..f022e99cc 100644 --- a/src/art/trainer_rank/_impl.py +++ b/src/art/trainer_rank/_impl.py @@ -574,6 +574,7 @@ class _MemoryCheck: class _MemoryProfile: bytes_per_token: float packed_tokens: int + # Active execution rows only; total-input telemetry can include inactive rows. logical_per_packed: float = 1.0 # Historical forward-retained/forward-peak ratio for calibration telemetry, # max-merged only from forward-return observations. Admission uses the @@ -901,6 +902,12 @@ class _FlatForwardPlan: output_bytes: int signature: _MemorySignature selected_max_depth: int = 0 + inactive_logical_tokens: int = 0 + + @property + def active_logical_tokens(self) -> int: + # Keep total-input telemetry while pricing only executed requests. + return self.logical_tokens - self.inactive_logical_tokens @property def subforward_count(self) -> int: @@ -2620,7 +2627,7 @@ def _split_chunk_lower_cost( slot_group_count=len(groups), grad_modes=tuple(mode for (_, mode), _ in groups), ) - logical_tokens = sum(int(row.numel()) for row in rows) + logical_tokens = _active_logical_tokens(requests) cost = self._subforward_cost( packed_tokens=packed_tokens, output_bytes=output_bytes, @@ -2674,7 +2681,7 @@ def _plan_cost(self, plan: _FlatForwardPlan) -> _SubforwardCost: packed_tokens=plan.packed_tokens, output_bytes=plan.output_bytes, signature=plan.signature, - logical_tokens=plan.logical_tokens, + logical_tokens=plan.active_logical_tokens, ) def _subforward_cost( @@ -3642,9 +3649,7 @@ def estimate(width: int) -> tuple[_MemoryCheck, bool, bool] | None: estimates[width] = None return None assert values is not None - logical_tokens = sum( - int(request.input_tokens.numel()) for request in local_requests - ) + logical_tokens = _active_logical_tokens(local_requests) def priced( packed_tokens: int, @@ -4269,6 +4274,7 @@ def _plan_flat_forward( 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( @@ -4750,7 +4756,7 @@ def _memory_check( packed_tokens=forward.packed_tokens, output_bytes=forward.output_bytes, signature=forward.signature, - logical_tokens=forward.logical_tokens, + logical_tokens=forward.active_logical_tokens, ), sync_across_dp=sync_across_dp, ) @@ -4915,7 +4921,7 @@ def _update_memory_profile( 0 if previous is None else previous.packed_tokens, ), logical_per_packed=max( - plan.logical_tokens / max(1, plan.packed_tokens), + plan.active_logical_tokens / max(1, plan.packed_tokens), 1.0 if previous is None else previous.logical_per_packed, ), retained_fraction=retained_fraction, @@ -5597,6 +5603,17 @@ def _validate_top_k(top_k: int, model: object) -> None: raise ValueError(f"top_k={top_k} exceeds vocabulary size {vocab_size}") +def _active_logical_tokens(requests: Sequence[AnyForwardInput]) -> int: + return sum( + int(request.input_tokens.numel()) + for request in requests + if request.target_tokens is not None + or request.logits + or request.top_k is not None + or request.hidden_states + ) + + def _request_mix_key(request: AnyForwardInput) -> str: parts = [] if request.target_tokens is not None: diff --git a/tests/unit/test_trainer_rank_active_memory.py b/tests/unit/test_trainer_rank_active_memory.py new file mode 100644 index 000000000..a81905fe1 --- /dev/null +++ b/tests/unit/test_trainer_rank_active_memory.py @@ -0,0 +1,201 @@ +"""CPU admission contracts; injected observations are not GPU peak measurements.""" + +from dataclasses import replace +from types import SimpleNamespace +from typing import Any, cast + +import pytest +import torch + +from art.trainer_rank import ForwardInput, ForwardOutput, TrainerRank, Unset + + +class _Model(torch.nn.Module): + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter(torch.zeros((), dtype=torch.bfloat16)) + self.config = SimpleNamespace(hidden_size=8, num_layers=4, padded_vocab_size=32) + self.decoder = object() + + def _preprocess(self, *args, **kwargs): + return None + + +def _rank(): + return TrainerRank( + cast( + Any, + SimpleNamespace( + model=[_Model()], + optimizer=None, + provider=SimpleNamespace(hidden_size=8, num_layers=4), + model_support_handler=SimpleNamespace(build_gdn_execution_spec=False), + ), + ) + ) + + +def _requests(output="hidden_states", inactive_length=1): + tokens = torch.arange(8) + active = { + "hidden_states": ForwardInput(input_tokens=tokens, hidden_states=True), + "target_tokens": ForwardInput(input_tokens=tokens, target_tokens=tokens), + "logits": ForwardInput(input_tokens=tokens, logits=True), + "top_k": ForwardInput(input_tokens=tokens, top_k=2), + }[output] + return [ + active, + ForwardInput(input_tokens=torch.arange(inactive_length)), + ] + + +@pytest.mark.parametrize( + "output", ["hidden_states", "target_tokens", "logits", "top_k"] +) +@pytest.mark.parametrize("no_grad", [False, True]) +def test_inactive_length_preserves_warm_cost_and_profile(monkeypatch, output, no_grad): + rank = _rank() + short = [replace(r, no_grad=no_grad) for r in _requests(output)] + long = [replace(r, no_grad=no_grad) for r in _requests(output, 8001)] + plans = [rank._plan_flat_forward(requests) for requests in (short, long)] + first, second = plans + assert first.signature == second.signature + assert first.packed_tokens == second.packed_tokens == 8 + assert first.output_bytes == second.output_bytes + assert first.output_metadata == second.output_metadata + assert (first.logical_tokens, second.logical_tokens) == (9, 8009) + assert len(first.groups) == len(second.groups) == 1 + a, b = first.groups[0], second.groups[0] + assert a.request_indices == b.request_indices == (0,) + for field in ("tokens", "group_ids", "parent_ids", "position_ids"): + assert torch.equal(getattr(a.packed, field), getattr(b.packed, field)) + rank._update_memory_profile(first, 10_000, retained_bytes=1000) + profile = rank._memory_profiles[first.signature] + budget = rank._memory_check(first).estimated_required_bytes + monkeypatch.setattr(rank, "_available_memory_bytes", lambda: budget) + assert rank._memory_check(first).fits + assert rank._memory_check(second) == rank._memory_check(first) + assert rank._plan_cost(first) == rank._plan_cost(second) + for requests, plan in zip((short, long), plans, strict=True): + lower = rank._split_chunk_lower_cost( + requests, [r.input_tokens for r in requests], checkpoint=Unset + ) + assert lower == rank._plan_cost(plan) + rank._update_memory_profile(second, 10_000, retained_bytes=1000) + assert rank._memory_profiles[first.signature] == profile + + +def test_public_pair_avoids_inactive_only_split_and_keeps_total_telemetry(monkeypatch): + rank = _rank() + monkeypatch.setattr(rank, "_dp_rank_and_size", lambda: (0, 1)) + observed = rank._plan_flat_forward(_requests()) + rank._update_memory_profile(observed, 10_000, retained_bytes=1000) + budget = rank._memory_check(observed).estimated_required_bytes + monkeypatch.setattr(rank, "_available_memory_bytes", lambda: budget) + executed = [] + + def run(plan, **kwargs): + executed.append(plan) + # Only model execution is replaced. Admission, split search, trust and + # public iterator reconstruction all run their production paths. + return [ + ForwardOutput(None, None, None, None) for _ in range(plan.request_count) + ], None + + monkeypatch.setattr(rank, "_run_flat_plan_with_memory_tracking", run) + batches = list(rank.forward_micro_batches([_requests(inactive_length=8001)])) + assert len(batches) == 1 + batch = batches[0] + assert batch.indices == (0,) + assert batch.stats.global_count == batch.stats.local_count == 1 + assert batch.stats.logical_tokens == 8009 + assert batch.stats.packed_tokens == 8 + assert len(batch.outputs) == 1 and len(batch.outputs[0]) == 2 + assert batch.stats.subforward_count == len(executed) == 1 + assert batch.stats.estimated_required_bytes == budget + assert not batch.stats.cold_start + + +def test_active_length_still_increases_warm_cost(): + rank = _rank() + short = _requests() + observed = rank._plan_flat_forward(short) + rank._update_memory_profile(observed, 10_000, retained_bytes=1000) + longer = rank._plan_flat_forward( + [replace(short[0], input_tokens=torch.arange(16)), short[1]] + ) + assert observed.signature == longer.signature + assert rank._memory_check(longer).estimated_required_bytes == 22_000 + assert rank._plan_cost(longer).retained > rank._plan_cost(observed).retained + + +def test_inactive_observation_cannot_discount_later_shared_active_work(monkeypatch): + ranks = [_rank(), _rank()] + checks = [] + for rank, inactive_length in zip(ranks, (1, 8001), strict=True): + requests = _requests(inactive_length=inactive_length) + observed = rank._plan_flat_forward(requests, memory_minimal=True) + rank._update_memory_profile(observed, 10_000, retained_bytes=1000) + candidate = rank._plan_flat_forward( + [requests[0]] * 8 + _requests()[1:], memory_minimal=True + ) + assert candidate.signature == observed.signature + assert candidate.packed_tokens == observed.packed_tokens == 8 + assert candidate.logical_tokens == 65 + monkeypatch.setattr(rank, "_available_memory_bytes", lambda: 20_000) + checks.append(rank._memory_check(candidate)) + # Identical observed GPU work must produce identical future admission. + # Total-input ratios formerly discounted the second estimate to 11,985, + # admitting a plan that the equivalent short calibration refused. + assert checks[0] == checks[1] + assert checks[0].estimated_required_bytes == 88_000 + assert not checks[0].fits + + +def test_warm_admission_rechecks_current_residency(monkeypatch): + rank = _rank() + plan = rank._plan_flat_forward(_requests()) + rank._update_memory_profile(plan, 10_000, retained_bytes=1000) + required = rank._memory_check(plan).estimated_required_bytes + profile = rank._memory_profiles[plan.signature] + state = {"allocated": 10_000, "reserved": 10_000} + total = 100_000 + monkeypatch.setattr(rank, "device", torch.device("cuda")) + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr( + torch.cuda, "mem_get_info", lambda _: (total - state["reserved"], total) + ) + monkeypatch.setattr(torch.cuda, "memory_allocated", lambda _: state["allocated"]) + monkeypatch.setattr(torch.cuda, "memory_reserved", lambda _: state["reserved"]) + monkeypatch.delenv("ART_TRAINER_RANK_TEST_HOOKS", raising=False) + # Fresh memory accounting observes newly resident state without discarding + # a valid incremental profile; cached free blocks remain reusable. + for allocated, reserved, fits in [ + (10_000, 10_000, True), + (90_000, 90_000, False), + (10_000, 90_000, True), + ]: + state.update(allocated=allocated, reserved=reserved) + check = rank._memory_check(plan) + assert check.estimated_required_bytes == required + assert check.available_bytes == total - allocated - 3000 + assert check.fits == fits + assert rank._memory_profiles[plan.signature] == profile + + +def test_distinct_output_and_pure_grad_modes_require_their_own_profile(): + rank = _rank() + requests = _requests() + observed = rank._plan_flat_forward(requests) + rank._update_memory_profile(observed, 10_000, retained_bytes=1000) + for changed in ( + [replace(r, no_grad=True) for r in requests], + _requests("logits"), + _requests("target_tokens"), + _requests("top_k"), + ): + plan = rank._plan_flat_forward(changed) + assert plan.signature != observed.signature + assert not rank._all_ranks_have_memory_profile( + packed_tokens=plan.packed_tokens, signature=plan.signature + ) diff --git a/tests/unit/test_trainer_rank_split.py b/tests/unit/test_trainer_rank_split.py index edb5f559f..a0ab095c1 100644 --- a/tests/unit/test_trainer_rank_split.py +++ b/tests/unit/test_trainer_rank_split.py @@ -57,6 +57,7 @@ _MemoryProfile, _SplitForwardPlan, ) +from art.trainer_rank._prefix_tree_planner import plan_prefix_tree_layout if TYPE_CHECKING: from art.megatron.lora import LoRASlotRef @@ -973,16 +974,14 @@ def test_retained_ratio_bound_uses_original_guard_at_trusted_endpoint( ) -> None: rank = _retained_ratio_rank(monkeypatch) - def request(tokens: list[int], *, output: bool = True) -> ForwardInput: + def request(tokens: list[int]) -> ForwardInput: values = torch.tensor(tokens) - return ForwardInput( - input_tokens=values, target_tokens=values if output else None - ) + return ForwardInput(input_tokens=values, target_tokens=values) original = [ - request([0, *range(1, 7)]), - request([0, *range(21, 27)]), - request([99], output=False), + request([0, 1, *range(2, 7)]), + request([0, 1, *range(22, 27)]), + request([99]), ] observed = rank._plan_flat_forward(original, memory_minimal=True) assert (observed.packed_tokens, observed.logical_tokens) == (13, 15) @@ -991,11 +990,19 @@ def request(tokens: list[int], *, output: bool = True) -> ForwardInput: peak_delta_bytes=observed.output_bytes + 13, retained_bytes=observed.output_bytes, ) - requests = [ - request([*range(24), *range(30, 58)]), - request([*range(24), *range(90, 118)]), - request(list(range(856)), output=False), - ] + requests = [request([*range(24), *range(30, 46)]) for _ in range(8)] + requests += [request([*range(24), *range(90, 130)]) for _ in range(10)] + select = rank._select_group_layout + + def select_leaf_sharing(input_ids, *, memory_minimal=False, grad_enabled=True): + tree, layout = select(input_ids, memory_minimal=True, grad_enabled=grad_enabled) + if not memory_minimal: + # A legal intermediate layout shares each repeated leaf path but + # replays the common prefix: 104 physical / 960 active logical rows. + layout = plan_prefix_tree_layout(tree, tree.terminal_segment_indices) + return tree, layout + + monkeypatch.setattr(rank, "_select_group_layout", select_leaf_sharing) minimal = rank._plan_flat_forward(requests, memory_minimal=True) plan = rank._plan_flat_forward(requests) profile = rank._memory_profiles[observed.signature] @@ -1006,6 +1013,7 @@ def request(tokens: list[int], *, output: bool = True) -> ForwardInput: 104, 960, ) + assert plan.active_logical_tokens == plan.logical_tokens assert plan.packed_tokens == cap assert plan.logical_tokens / cap == limit # Rearranging the original comparison changes the answer at this equality. @@ -1024,7 +1032,10 @@ def request(tokens: list[int], *, output: bool = True) -> ForwardInput: exact_bytes = rank._split_rung_check([exact, exact]).estimated_required_bytes monkeypatch.setattr(rank, "_available_memory_bytes", lambda: exact_bytes) selected, check = rank._admit_split_rung( - [(0, 1, 2), (3, 4, 5)], requests * 2, rows * 2, checkpoint=Unset + [tuple(range(18)), tuple(range(18, 36))], + requests * 2, + rows * 2, + checkpoint=Unset, ) assert selected is not None and check.fits assert check.estimated_required_bytes == exact_bytes