diff --git a/.github/workflows/prek.yml b/.github/workflows/prek.yml index d7991b079..18a49a302 100644 --- a/.github/workflows/prek.yml +++ b/.github/workflows/prek.yml @@ -230,6 +230,7 @@ jobs: tests/unit/test_trainer_rank_weird_shapes.py \ tests/unit/test_trainer_rank_admission_inputs.py \ tests/unit/test_trainer_rank_checkpoint_memory.py \ + tests/unit/test_trainer_rank_profile_warm.py \ tests/unit/test_trainer_rank_tp_floor.py \ tests/unit/test_trainer_rank_checkpoint_gradient_memory.py \ tests/unit/test_trainer_rank_slot_memory.py \ @@ -275,6 +276,7 @@ jobs: --ignore=tests/unit/test_trainer_rank_weird_shapes.py \ --ignore=tests/unit/test_trainer_rank_admission_inputs.py \ --ignore=tests/unit/test_trainer_rank_checkpoint_memory.py \ + --ignore=tests/unit/test_trainer_rank_profile_warm.py \ --ignore=tests/unit/test_trainer_rank_tp_floor.py \ --ignore=tests/unit/test_trainer_rank_checkpoint_gradient_memory.py \ --ignore=tests/unit/test_trainer_rank_slot_memory.py \ diff --git a/src/art/trainer_rank/__init__.py b/src/art/trainer_rank/__init__.py index fd321a820..04b78329e 100644 --- a/src/art/trainer_rank/__init__.py +++ b/src/art/trainer_rank/__init__.py @@ -267,6 +267,11 @@ def forward_micro_batches( iterator before making collective calls after an early exit. Guards apply on the iterator's thread; raw torch.distributed calls are not guarded. Collective calls must still match across ranks. + + Admission learns each wave's memory peak, including the caller's loss + and backward when they run inside the yield; a backward deferred past the + yield is not learned. A wave with another TrainerRank forward inside its + yield cannot lower later estimates. """ forward = cast( Callable[..., Iterator[MicroBatch[ForwardInputs, ForwardOutputs]]], diff --git a/src/art/trainer_rank/_impl.py b/src/art/trainer_rank/_impl.py index a7d67c072..d4f7bde0f 100644 --- a/src/art/trainer_rank/_impl.py +++ b/src/art/trainer_rank/_impl.py @@ -614,6 +614,16 @@ class _MemoryProfile: # Separate the retained compute rate from the peak, which also learns # caller-owned backward workspace. Requested outputs are charged explicitly. retained_compute_bytes_per_token: float | None = None + # A signature's first executed plan also pays one-time costs (compilation, + # first-use workspaces), which a small first wave spreads over few tokens. + # Admission uses the lower of the fit over every observation and the same + # fit over later flat waves' whole caller-phase peaks (forward plus the + # caller's loss and backward in the yield), which prices a wave smaller + # than the smallest of them as if it were that large. + warm_bytes_per_token: float | None = None + warm_packed_tokens: int | None = None + warm_logical_per_packed: float | None = None + caller_plans: int = 0 @dataclass(frozen=True) @@ -2031,6 +2041,9 @@ def memory_field(name: str, default: Any = None) -> Any: self._hybridep_rows_high_water = 0 self._cache_recovery_state = _CacheRecoveryState() self._memory_profiles: dict[_MemorySignature, _MemoryProfile] = {} + # Tracked peak-counter resets, and the latest (resets, peak) reading. + self._peak_resets = 0 + self._peak_reading: tuple[int, int] | None = None self._split_memory_floors: dict[bytes, int] = {} self._split_memory_floor_status = "not_observed" self._last_global_micro_batch_size: int | None = None @@ -2803,7 +2816,10 @@ def forward_micro_batches( rows and charges 6 KiB per logical token for the head plus the caller's loss saves and backward transients. A caller whose per-token head and loss memory peaks above that is unsupported: shared plans can exceed - their estimate. + their estimate. Backward is learned only when it runs inside the yield; + a backward deferred past it is not. A wave with another TrainerRank + forward inside its yield cannot lower later estimates below the + signature's first wave. """ if not isinstance(yield_empty, bool): raise TypeError("yield_empty must be a bool") @@ -2882,6 +2898,7 @@ def _forward_micro_batches( outputs: list[Any] = [] flat_outputs = iter(tracked_outputs) error: BaseException | None = None + interval: tuple[int, int] | None = None try: if isinstance(candidate.plan, _FlatForwardPlan): tracked_outputs, memory_baseline = ( @@ -2891,6 +2908,8 @@ def _forward_micro_batches( context="forward_micro_batches", ) ) + # This wave's peak interval, which its caller phase continues. + interval = self.__dict__.get("_peak_reading") else: tracked_outputs, memory_baseline, forward_peak = ( self._execute_split_plan_with_memory_tracking( @@ -2957,7 +2976,12 @@ def _forward_micro_batches( # optimizer headroom. Peak only: the retained observation belongs # to the forward's return, already recorded for this same plan. if isinstance(candidate.plan, _FlatForwardPlan): - self._update_peak_memory_profile(candidate.plan, memory_baseline) + self._update_peak_memory_profile( + candidate.plan, + memory_baseline, + caller_phase=True, + interval=interval, + ) elif memory_baseline is not None: self._record_split_memory_floor( candidate.plan, memory_baseline, forward_peak @@ -6227,6 +6251,7 @@ def _run_flat_plan_with_memory_tracking( torch.cuda.synchronize(self.device) baseline = int(torch.cuda.memory_allocated(self.device)) torch.cuda.reset_peak_memory_stats(self.device) + self._peak_resets = self.__dict__.get("_peak_resets", 0) + 1 else: baseline = None observation = getattr(self, "_planner_observation", None) @@ -6277,6 +6302,9 @@ def _update_peak_memory_profile( plan: _FlatForwardPlan, baseline: int | None, retained_after: int | None = None, + *, + caller_phase: bool = False, + interval: tuple[int, int] | None = None, ) -> None: if baseline is None: return @@ -6284,12 +6312,21 @@ def _update_peak_memory_profile( observation = getattr(self, "_planner_observation", None) if observation is not None: observation["peak"] = max(observation["peak"], peak) + resets = self.__dict__.get("_peak_resets", 0) + self._peak_reading = (resets, peak) self._update_memory_profile( plan, max(0, peak - baseline), retained_bytes=( None if retained_after is None else max(0, retained_after - baseline) ), + # Only a whole wave: a nested forward resetting the counter during + # the yield, or an untracked reset (the counter fell below this + # wave's forward peak), leaves the peak since then short of it. + caller_phase=caller_phase + and interval is not None + and interval[0] == resets + and peak >= interval[1], ) def _begin_planner_observation( @@ -8051,27 +8088,52 @@ def _estimate_required_memory_bytes_from_values( # usual trust growth. Normalize before multiplying: cancelling packed # tokens through two float operations can otherwise make a larger warm # layout cheaper. - profiled_tokens: int | float = packed_tokens packed_priced = profiled is not None and _packed_priced( signature, self._one_layer_recompute() ) - if profiled is not None and logical_tokens is not None: - profiled_tokens = max( + + def profiled_tokens(logical_per_packed: float) -> int | float: + if logical_tokens is None: + return packed_tokens + return max( packed_tokens, logical_tokens - / profiled.logical_per_packed + / logical_per_packed / (_MEMORY_PROFILE_TRUST_GROWTH if packed_priced else 1), ) + # The trust window limits calibration growth, not the empirical floor. # Dropping that floor beyond the window can admit a larger request that # was refused just inside it, even below a previously observed peak. if profiled is None: compute = static_compute else: + profiled_bytes = profiled.bytes_per_token * profiled_tokens( + profiled.logical_per_packed + ) + if ( + profiled.warm_bytes_per_token is not None + and profiled.warm_packed_tokens is not None + and profiled.warm_logical_per_packed is not None + ): + # Later plans' own sharing: a rate learned under lighter + # sharing scales up for deeper-shared plans, as above. Pricing + # smaller waves as the smallest later plan keeps cost monotone + # in tokens, which the width search's bounds rely on; it is + # sound where a wave's memory is a fixed cost plus a per-token + # rate, as that plan's rate covers its share of the fixed cost. + profiled_bytes = min( + profiled_bytes, + profiled.warm_bytes_per_token + * max( + profiled.warm_packed_tokens, + profiled_tokens(profiled.warm_logical_per_packed), + ), + ) compute = max( static_compute, int( - profiled.bytes_per_token * profiled_tokens + profiled_bytes + ( _PACKED_PRICED_LOGICAL_ROW_BYTES * logical_tokens # Branch states grow with segments, not packed rows; @@ -8244,12 +8306,33 @@ def _update_memory_profile( peak_delta_bytes: int, *, retained_bytes: int | None, + caller_phase: bool = False, ) -> None: if plan.packed_tokens <= 0: return compute_delta = max(0, peak_delta_bytes - plan.output_bytes) bytes_per_token = compute_delta / max(1, plan.packed_tokens) previous = self._memory_profiles.get(plan.signature) + logical_per_packed = plan.active_logical_tokens / max(1, plan.packed_tokens) + warm_rate = None if previous is None else previous.warm_bytes_per_token + warm_tokens = None if previous is None else previous.warm_packed_tokens + warm_sharing = None if previous is None else previous.warm_logical_per_packed + caller_plans = 0 if previous is None else previous.caller_plans + # Only a flat wave's whole caller phase covers the caller's backward; + # split children and dp_rank_forward observe forward only. A caller that + # defers backward past the yield is not covered. The profile's first + # such plan also pays one-time costs, so it is not warm. + if caller_phase: + if caller_plans: + warm_rate = max(bytes_per_token, warm_rate or 0.0) + warm_tokens = min(plan.packed_tokens, warm_tokens or plan.packed_tokens) + warm_sharing = max(logical_per_packed, warm_sharing or 1.0) + caller_plans += 1 + elif caller_plans: + # After the seed, other readings cannot fit the warm profile, but a + # higher one still raises its rate (inert until a whole warm plan + # fits the rest), as it raises the fit over every observation. + warm_rate = max(warm_rate or 0.0, bytes_per_token) retained_fraction = None if previous is None else previous.retained_fraction retained_compute = ( None if previous is None else previous.retained_compute_bytes_per_token @@ -8278,11 +8361,15 @@ def _update_memory_profile( 0 if previous is None else previous.packed_tokens, ), logical_per_packed=max( - plan.active_logical_tokens / max(1, plan.packed_tokens), + logical_per_packed, 1.0 if previous is None else previous.logical_per_packed, ), retained_fraction=retained_fraction, retained_compute_bytes_per_token=retained_compute, + warm_bytes_per_token=warm_rate, + warm_packed_tokens=warm_tokens, + warm_logical_per_packed=warm_sharing, + caller_plans=caller_plans, ) def _forward_item(self, request: AnyForwardInput) -> _ForwardItem: diff --git a/tests/unit/test_trainer_rank_cache_recovery.py b/tests/unit/test_trainer_rank_cache_recovery.py index e165346f2..ef34113c3 100644 --- a/tests/unit/test_trainer_rank_cache_recovery.py +++ b/tests/unit/test_trainer_rank_cache_recovery.py @@ -172,7 +172,7 @@ def make(self): self.addCleanup(observer_clock.stop) q = object.__new__(_impl.TrainerRank) q.device = types.SimpleNamespace(type="cuda") - q._update_peak_memory_profile = lambda *a: None + q._update_peak_memory_profile = lambda *a, **k: None q._execute_flat_plan = lambda p: [object() for _ in range(p.request_count)] q._telemetry_signature = lambda p: {} q._telemetry_plan_signature = lambda p: {} diff --git a/tests/unit/test_trainer_rank_physical_reserve.py b/tests/unit/test_trainer_rank_physical_reserve.py index bbf7ac3ef..b41a6e615 100644 --- a/tests/unit/test_trainer_rank_physical_reserve.py +++ b/tests/unit/test_trainer_rank_physical_reserve.py @@ -152,7 +152,7 @@ def forward(plan, **kwargs): monkeypatch.setattr( rank, "_update_peak_memory_profile", - lambda plan, baseline: profiles.append((plan, baseline)), + lambda plan, baseline, **_: profiles.append((plan, baseline)), ) state = allocator(monkeypatch) original_phase = _impl._telemetry_phase diff --git a/tests/unit/test_trainer_rank_profile_warm.py b/tests/unit/test_trainer_rank_profile_warm.py new file mode 100644 index 000000000..f1fa1c38e --- /dev/null +++ b/tests/unit/test_trainer_rank_profile_warm.py @@ -0,0 +1,473 @@ +"""A signature's first plan seeds its memory profile; later plans set a warm fit. + +CPU admission math, not a bound. On real Qwen3.6-35B-A3B (40 layers, CP2/EP2) +the run's first, small, cold wave peaked at 275 KB per packed token while every +later wave ran 206-247 KB; the max-merged rate priced every later wave with the +first wave's one-time costs. +""" + +from dataclasses import replace + +from test_trainer_rank_active_memory import _rank as packed_rank +from test_trainer_rank_active_memory import _requests as packed_requests +from test_trainer_rank_checkpoint_memory import rank, requests +import torch + +from art.trainer_rank import _impl +from art.trainer_rank._impl import _MemoryProfile, _packed_priced + +FIRST, WARM = 3_000_000, 2_000_000 # Per packed token; far above the static floor. + + +def _plans(r): + first, larger = ( + r._plan_flat_forward(requests(1024, 64)), + r._plan_flat_forward(requests(8192, 64)), + ) + assert first.signature == larger.signature + assert larger.packed_tokens > first.packed_tokens + return first, larger + + +def _update(r, plan, rate, *, caller_phase): + r._update_memory_profile( + plan, + plan.output_bytes + rate * plan.packed_tokens, + retained_bytes=None, + caller_phase=caller_phase, + ) + + +def _observe(r, plan, rate): + """One flat wave: its forward return, then its caller phase (backward).""" + _update(r, plan, rate, caller_phase=False) + _update(r, plan, rate, caller_phase=True) + + +def _required(r, plan, logical_tokens=None): + return r._estimate_required_memory_bytes_from_values( + packed_tokens=plan.packed_tokens, + output_bytes=0, + signature=plan.signature, + logical_tokens=logical_tokens or plan.packed_tokens, + ) + + +def test_the_first_plans_one_time_costs_do_not_price_later_waves(): + r = rank() + first, larger = _plans(r) + _observe(r, first, FIRST) + profile = r._memory_profiles[first.signature] + assert profile.bytes_per_token == FIRST + assert profile.warm_bytes_per_token is None + assert profile.caller_plans == 1 + assert _required(r, larger) == int(FIRST * larger.packed_tokens * 1.1) + _observe(r, larger, WARM) + profile = r._memory_profiles[first.signature] + assert profile.bytes_per_token == FIRST + assert profile.warm_bytes_per_token == WARM + assert profile.warm_packed_tokens == larger.packed_tokens + assert profile.warm_logical_per_packed == 1 + assert _required(r, larger) == int(WARM * larger.packed_tokens * 1.1) + # Smaller waves: the lower of today's rate and the warm fit at the + # smallest later plan's size. + assert _required(r, first) == int(FIRST * first.packed_tokens * 1.1) + near = r._plan_flat_forward(requests(6144, 64)) + assert WARM * larger.packed_tokens < FIRST * near.packed_tokens + assert _required(r, near) == int(WARM * larger.packed_tokens * 1.1) + + +def test_forward_only_observations_never_set_the_warm_fit(): + """Split children and dp_rank_forward observe forward only; a flat wave's + caller phase also includes its backward.""" + r = rank() + first, larger = _plans(r) + _observe(r, first, FIRST) + _update(r, larger, WARM, caller_phase=False) + assert r._memory_profiles[first.signature].warm_packed_tokens is None + assert _required(r, larger) == int(FIRST * larger.packed_tokens * 1.1) + # A later caller phase is warm, and its peak includes the forward's. + _update(r, larger, WARM, caller_phase=True) + assert r._memory_profiles[first.signature].warm_bytes_per_token == WARM + + +def test_the_warm_rate_ratchets_and_its_extent_grows_down(): + r = rank() + first, larger = _plans(r) + _observe(r, first, FIRST) + _observe(r, larger, WARM) + # A later, smaller plan: its higher rate is kept, and it extends the warm + # extent down to its own size. + middle = r._plan_flat_forward(requests(4096, 64)) + _observe(r, middle, WARM + 100_000) + profile = r._memory_profiles[first.signature] + assert profile.warm_bytes_per_token == WARM + 100_000 + assert profile.warm_packed_tokens == middle.packed_tokens + assert profile.bytes_per_token == FIRST + # A lower warm rate never lowers it. + _observe(r, larger, WARM - 100_000) + assert r._memory_profiles[first.signature].warm_bytes_per_token == WARM + 100_000 + + +def test_profiles_without_a_warm_fit_price_as_before(): + """Replayed reports and callers that set profiles directly carry no warm + fit; the first caller phase after that seeds it.""" + r = rank() + first, larger = _plans(r) + r._memory_profiles[first.signature] = _MemoryProfile( + bytes_per_token=FIRST, packed_tokens=first.packed_tokens + ) + assert _required(r, first) == int(FIRST * first.packed_tokens * 1.1) + _observe(r, larger, WARM) + assert r._memory_profiles[first.signature].warm_bytes_per_token is None + _observe(r, larger, WARM) + assert r._memory_profiles[first.signature].warm_bytes_per_token == WARM + + +def test_each_signature_and_a_cleared_profile_seeds_again(): + """Profiles are keyed by memory signature; a new signature (a different + group count, slot shapes or topology) or a cleared profile starts cold.""" + r = rank() + first, larger = _plans(r) + _observe(r, first, FIRST) + _observe(r, larger, WARM) + other = r._plan_flat_forward(requests(8192, 0)[:1]) + assert other.signature != first.signature + _observe(r, other, FIRST) + assert r._memory_profiles[other.signature].warm_bytes_per_token is None + # Clearing a profile makes the next plan its first again. + del r._memory_profiles[first.signature] + _observe(r, larger, FIRST) + assert r._memory_profiles[first.signature].warm_bytes_per_token is None + assert _required(r, larger) == int(FIRST * larger.packed_tokens * 1.1) + + +def test_a_warm_rate_learned_under_lighter_sharing_scales_for_deeper_sharing(): + """A deeply shared first plan and unshared later plans: the later plans' + lower rate prices unshared waves, but deeper-shared waves scale it up by + the sharing gap, so they never price below today's fit over every plan. + A later plan that shares as deeply extends the warm rate to them.""" + r = rank() + _, larger = _plans(r) + shared = replace(larger, logical_tokens=larger.packed_tokens * 64) + deep = shared.logical_tokens + + def today(): + profile = r._memory_profiles[larger.signature] + r._memory_profiles[larger.signature] = replace( + profile, warm_bytes_per_token=None, warm_packed_tokens=None + ) + prices = _required(r, larger), _required(r, larger, deep) + r._memory_profiles[larger.signature] = profile + return prices + + _observe(r, shared, FIRST) + _observe(r, larger, WARM) + profile = r._memory_profiles[larger.signature] + assert profile.logical_per_packed == 64 + assert profile.warm_logical_per_packed == 1 + unshared, deeper = today() + assert _required(r, larger) == int(WARM * larger.packed_tokens * 1.1) < unshared + assert _required(r, larger, deep) == deeper + _observe(r, shared, WARM) + assert r._memory_profiles[larger.signature].warm_logical_per_packed == 64 + assert _required(r, larger, deep) < today()[1] + + +def test_cost_stays_monotone_across_the_warm_extent(): + """The width search accepts on a larger layout's price and rejects on a + smaller one's, so a smaller wave must never price above a larger one.""" + r = rank() + first, larger = _plans(r) + _observe(r, first, FIRST) + _observe(r, larger, WARM) + n = larger.packed_tokens + + def required(packed, logical=None): + return r._estimate_required_memory_bytes_from_values( + packed_tokens=packed, + output_bytes=0, + signature=larger.signature, + logical_tokens=logical or packed, + ) + + sweep = [required(t) for t in (1, n // 2, (2 * n) // 3, n - 1, n, n + 1, 2 * n)] + assert sweep == sorted(sweep) + shared = [required(t, 4 * n) for t in (1, n // 4, n - 1, n, 2 * n, 4 * n)] + assert shared == sorted(shared) + + +def test_packed_priced_cost_stays_monotone_across_the_warm_extent(monkeypatch): + # The fixture's requests are short; the short-request gate has its own tests. + monkeypatch.setattr(_impl, "_PACKED_PRICED_MIN_REQUEST_TOKENS", 1) + r = packed_rank() + signature = r._plan_flat_forward(packed_requests("target_tokens")).signature + assert _packed_priced(signature, r._one_layer_recompute()) + today = _MemoryProfile(bytes_per_token=50_000, packed_tokens=8) + warm = replace( + today, + warm_bytes_per_token=20_000, + warm_packed_tokens=1000, + warm_logical_per_packed=4, + ) + + def cost(profile, packed, logical=8_000): + r._memory_profiles[signature] = profile + return r._subforward_cost( + packed_tokens=packed, + logical_tokens=logical, + output_bytes=0, + signature=signature, + ).required + + sweep = [cost(warm, t) for t in (1, 500, 999, 1000, 1001, 2000, 8000)] + assert sweep == sorted(sweep) + # Just below the warm extent: priced as the smallest later plan, not today. + assert sweep[2] == sweep[3] < cost(today, 999) + rows = [cost(warm, 1000, logical) for logical in (1000, 8_000, 32_000, 64_000)] + assert rows == sorted(rows) + + +def test_split_children_never_set_the_warm_fit(monkeypatch): + """The real split ladder: children observe forward only, and a split + wave's caller peak goes to split floors, never to the warm fit.""" + from test_trainer_rank_split import _packed_budget, _request + from test_trainer_rank_split import _rank as split_rank + + from art.trainer_rank import ForwardOutput + + r = split_rank(monkeypatch) + monkeypatch.setattr(r, "_retained_memory_bytes", lambda *_args, **_kwargs: 0) + forward_rate, backward_rate = 100, 300 + phases = [] + + def update(plan, baseline, retained_after=None, *, caller_phase=False, **_): + # Stands in for the CUDA peak read: backward peaks above forward. + phases.append((plan.request_count, caller_phase)) + rate = backward_rate if caller_phase else forward_rate + r._update_memory_profile( + plan, + plan.output_bytes + rate * plan.packed_tokens, + retained_bytes=None, + caller_phase=caller_phase, + ) + + def run(plan, **_kwargs): + # Production's forward-return update, then no CUDA baseline. + r._update_peak_memory_profile(plan, 0, 0) + return [ForwardOutput(None, None, None, None)] * plan.request_count, None + + monkeypatch.setattr(r, "_update_peak_memory_profile", update) + monkeypatch.setattr(r, "_run_flat_plan_with_memory_tracking", run) + _packed_budget(monkeypatch, r, 20) + items = [[_request(0)], [_request(m) for m in range(1, 5)], [_request(5)]] + batches = r.forward_micro_batches(items) + next(batches) + assert next(batches).stats.subforward_count > 1 + (signature,) = r._memory_profiles + # The seed's caller phase, then forward-only split children: no warm fit. + assert r._memory_profiles[signature].warm_packed_tokens is None + assert [p for p in phases if p[1]] == [(1, True)] + list(batches) + # Only the later flat wave's caller phase fits the warm rate. + assert [p for p in phases if p[1]] == [(1, True), (1, True)] + assert r._memory_profiles[signature].warm_bytes_per_token == backward_rate + + +def test_a_width_accepted_on_the_no_sharing_bound_never_executes_above_it( + monkeypatch, +): + """The width search accepts a width on the cheap no-sharing count. With a + warm fit, that count can reach the warm size while the shared layout + executed falls below it; the executed plan must not price higher.""" + from art.trainer_rank import ForwardInput, ForwardOutput + + monkeypatch.setattr(_impl, "_PACKED_PRICED_MIN_REQUEST_TOKENS", 1) + r = packed_rank() + monkeypatch.setattr(r, "_dp_rank_and_size", lambda: (0, 1)) + prefix = torch.arange(100) + items = [ + [ + ForwardInput( + input_tokens=torch.cat( + [prefix, torch.arange(1000 + 10 * i, 1010 + 10 * i)] + ), + target_tokens=torch.cat( + [prefix, torch.arange(1000 + 10 * i, 1010 + 10 * i)] + ), + ) + for i in range(4) + ] + ] + plan = r._plan_flat_forward(items[0]) + assert plan.packed_tokens < 200 <= plan.logical_tokens + r._memory_profiles[plan.signature] = _MemoryProfile( + bytes_per_token=100_000, + packed_tokens=plan.logical_tokens, + logical_per_packed=4, + warm_bytes_per_token=20_000, + warm_packed_tokens=200, + warm_logical_per_packed=4, + ) + own = r._memory_check(plan).estimated_required_bytes + monkeypatch.setattr(r, "_available_memory_bytes", lambda: 10 * own) + monkeypatch.setattr( + r, + "_run_flat_plan_with_memory_tracking", + lambda plan, **_: ( + [ForwardOutput(None, None, None, None)] * plan.request_count, + None, + ), + ) + (batch,) = list(r.forward_micro_batches(items)) + # Accepted on the no-sharing bound, which priced more tokens than ran. + assert batch.stats.packed_tokens == plan.packed_tokens + assert batch.stats.estimated_required_bytes > own + + +def _peak_reader(monkeypatch, r): + """Stand in for the CUDA peak counter, reset by each tracked forward.""" + state = {"peak": 0} + monkeypatch.setattr(torch.cuda, "max_memory_allocated", lambda *_: state["peak"]) + + def forward(plan, peak): + r._peak_resets = r.__dict__.get("_peak_resets", 0) + 1 + state["peak"] = peak + r._update_peak_memory_profile(plan, 0, 0) + return r._peak_reading + + return state, forward + + +def _caller_phase(r, plan, interval): + r._update_peak_memory_profile(plan, 0, caller_phase=True, interval=interval) + + +def test_only_a_whole_caller_phase_fits_the_warm_profile(monkeypatch): + r = rank() + first, larger = _plans(r) + state, forward = _peak_reader(monkeypatch, r) + n = larger.packed_tokens + _caller_phase( + r, first, forward(first, first.output_bytes + FIRST * first.packed_tokens) + ) + signature = first.signature + + def wave(interrupt): + interval = forward(larger, larger.output_bytes + WARM * n) + state["peak"] = larger.output_bytes + (WARM + 500_000) * n # backward + interrupt() + _caller_phase(r, larger, interval) + return r._memory_profiles[signature] + + # A nested forward during the yield resets the counter: not whole, even + # when its own peak is higher than this wave's forward. + def nested(): + forward(larger, larger.output_bytes + (WARM + 100_000) * n) + + assert wave(nested).warm_packed_tokens is None + # An untracked reset: the counter falls below this wave's forward peak. + assert wave(lambda: state.update(peak=0)).warm_packed_tokens is None + # Their readings raised the pending rate; an uninterrupted wave fits it. + assert wave(lambda: None).warm_bytes_per_token == WARM + 500_000 + + +def test_other_readings_raise_but_never_fit_the_warm_profile(): + r = rank() + first, larger = _plans(r) + _observe(r, first, FIRST) + _update(r, larger, WARM, caller_phase=False) + assert r._memory_profiles[first.signature].warm_packed_tokens is None + _observe(r, larger, WARM) + profile = r._memory_profiles[first.signature] + assert profile.warm_bytes_per_token == WARM + # A higher forward-only or interrupted reading raises the warm rate, but + # neither extends the warm extent nor its sharing. + middle = r._plan_flat_forward(requests(4096, 64)) + shared = replace(middle, logical_tokens=middle.packed_tokens * 8) + _update(r, shared, WARM + 300_000, caller_phase=False) + raised = r._memory_profiles[first.signature] + assert raised.warm_bytes_per_token == WARM + 300_000 + assert raised.warm_packed_tokens == profile.warm_packed_tokens + assert raised.warm_logical_per_packed == profile.warm_logical_per_packed + _update(r, larger, WARM - 100_000, caller_phase=False) + assert r._memory_profiles[first.signature].warm_bytes_per_token == WARM + 300_000 + + +def test_a_nested_forward_during_the_yield_cannot_fit_the_warm_profile(monkeypatch): + """The real micro-batch loop: a caller that runs another tracked forward + after its backward leaves a peak counter reset since the wave's forward.""" + from test_trainer_rank_split import _packed_budget, _request + from test_trainer_rank_split import _rank as split_rank + + from art.trainer_rank import ForwardOutput + + r = split_rank(monkeypatch) + monkeypatch.setattr(r, "_retained_memory_bytes", lambda *_args, **_kwargs: 0) + state, forward = _peak_reader(monkeypatch, r) + + def run(plan, **_kwargs): + forward(plan, plan.output_bytes + 100 * plan.packed_tokens) + return [ForwardOutput(None, None, None, None)] * plan.request_count, 0 + + monkeypatch.setattr(r, "_run_flat_plan_with_memory_tracking", run) + _packed_budget(monkeypatch, r, 10) + items = [[_request(m)] for m in range(3)] + batches = r.forward_micro_batches(items) + for index, batch in enumerate(batches): + (signature,) = r._memory_profiles + state["peak"] += 200 * batch.stats.packed_tokens # the caller's backward + if index == 1: + run(r._plan_flat_forward(items[0])) # a nested tracked forward + profile = r._memory_profiles[signature] + assert profile.caller_plans == 2 + # Only the uninterrupted third wave fit the warm rate, backward included. + assert profile.warm_bytes_per_token == 300 + + +def test_each_tracked_forward_starts_a_new_peak_interval(monkeypatch): + """The real tracked forward: its counter reset starts a new interval, so a + caller phase that continues an earlier one is not whole.""" + r = rank() + first, larger = _plans(r) + state = {"peak": 0} + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "synchronize", lambda *_: None) + monkeypatch.setattr(torch.cuda, "memory_allocated", lambda *_: 0) + monkeypatch.setattr( + torch.cuda, "reset_peak_memory_stats", lambda *_: state.update(peak=0) + ) + monkeypatch.setattr(torch.cuda, "max_memory_allocated", lambda *_: state["peak"]) + monkeypatch.setattr(r, "device", torch.device("cuda", 0)) + + def execute(plan): + state["peak"] = plan.output_bytes + FIRST * plan.packed_tokens + return [None] * plan.request_count + + monkeypatch.setattr(r, "_execute_flat_plan", execute) + check = _impl._MemoryCheck(0, 0, True) + _, baseline = r._run_flat_plan_with_memory_tracking(first, check=check, context="t") + seed = r._peak_reading + _caller_phase(r, first, seed) + _, baseline = r._run_flat_plan_with_memory_tracking( + larger, check=check, context="t" + ) + interval = r._peak_reading + assert baseline == 0 and interval[0] == seed[0] + 1 + r._run_flat_plan_with_memory_tracking(first, check=check, context="t") # nested + _caller_phase(r, larger, interval) + assert r._memory_profiles[larger.signature].warm_packed_tokens is None + + +def test_a_higher_reading_before_the_first_warm_plan_is_kept(): + """A seed, then a higher forward-only reading, then a lower whole warm + plan: the warm rate keeps the higher post-seed reading.""" + r = rank() + first, larger = _plans(r) + _observe(r, first, FIRST) + _update(r, larger, WARM + 300_000, caller_phase=False) + assert _required(r, larger) == int(FIRST * larger.packed_tokens * 1.1) + _observe(r, larger, WARM) + profile = r._memory_profiles[first.signature] + assert profile.warm_bytes_per_token == WARM + 300_000 + assert _required(r, larger) == int((WARM + 300_000) * larger.packed_tokens * 1.1) diff --git a/tests/unit/test_trainer_rank_validation.py b/tests/unit/test_trainer_rank_validation.py index 6212a0613..0d85e276a 100644 --- a/tests/unit/test_trainer_rank_validation.py +++ b/tests/unit/test_trainer_rank_validation.py @@ -3634,17 +3634,19 @@ def test_forward_micro_batches_profiles_caller_peak_after_yield( trainer = TrainerRank(_runtime()) _stub_forward(monkeypatch, trainer, profiled=True) plan = trainer._plan_flat_forward([_target_request(1)]) - monkeypatch.setattr( - trainer, - "_run_flat_plan_with_memory_tracking", - lambda *_args, **_kwargs: (_empty_outputs(plan), 123), - ) - profiles: list[tuple[int, int | None]] = [] + + def run(*_args, **_kwargs): + # The forward's peak interval, which the caller phase continues. + trainer._peak_reading = (7, 99) + return _empty_outputs(plan), 123 + + monkeypatch.setattr(trainer, "_run_flat_plan_with_memory_tracking", run) + profiles: list[tuple[int, int | None, bool, object]] = [] monkeypatch.setattr( trainer, "_update_peak_memory_profile", - lambda candidate, baseline: profiles.append( - (candidate.packed_tokens, baseline) + lambda candidate, baseline, caller_phase=False, interval=None: profiles.append( + (candidate.packed_tokens, baseline, caller_phase, interval) ), ) @@ -3654,7 +3656,9 @@ def test_forward_micro_batches_profiles_caller_peak_after_yield( assert profiles == [] with pytest.raises(StopIteration): next(batches) - assert profiles == [(plan.packed_tokens, 123)] + # The caller phase continues the wave's own interval: it may feed the + # warm fit if no other reset interrupted it. + assert profiles == [(plan.packed_tokens, 123, True, (7, 99))] @pytest.mark.parametrize("no_grad", [False, True]) @@ -3684,7 +3688,7 @@ def select_after_release(*args, **kwargs): monkeypatch.setattr(trainer, "_select_next_micro_batch", select_after_release) profiled: list[bool] = [] - def profile(*_args): + def profile(*_args, **_kwargs): # Keep the completed wave through its caller peak observation. profiled.append(tensors[-1]() is not None)