Exclude inactive requests from TrainerRank memory profile ratios

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:

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
+        )

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