diff --git a/src/art/trainer_rank/_impl.py b/src/art/trainer_rank/_impl.py index a183cc117..6acc3fea3 100644 --- a/src/art/trainer_rank/_impl.py +++ b/src/art/trainer_rank/_impl.py @@ -1205,7 +1205,7 @@ def _configure_moe_dispatcher_caches(model: Sequence[torch.nn.Module]) -> None: def _moe_output_bytes_per_token( model: Sequence[torch.nn.Module], shape: ParallelShape ) -> int: - """Known eager FC2 input/outputs, not the complete or compiled working set.""" + """Known routed-expert working set, not a complete model/compiled bound.""" if shape != ParallelShape(tp=1, cp=1): return 0 from megatron.core.extensions.transformer_engine import TERowParallelGroupedLinear @@ -1216,7 +1216,7 @@ def _moe_output_bytes_per_token( MoEAlltoAllTokenDispatcher, ) - from art.megatron.lora import LoRA, MLPExpertsLinearFC2LoRA + from art.megatron.lora import LoRA, MLPExpertsLinearFC1LoRA, MLPExpertsLinearFC2LoRA coefficient = 0 for chunk in model: @@ -1288,6 +1288,27 @@ def _moe_output_bytes_per_token( # their sum. Compilers may reuse storage; other workspace is # not covered. Unknown input metadata keeps the prior pair. features = inputs.shape[-2] + 3 * fc2.out_features + fc1 = getattr(experts, "linear_fc1", None) + if ( + type(fc1) is MLPExpertsLinearFC1LoRA + and "forward" not in vars(fc1) + and not getattr(fc1, "_forward_hooks", None) + and not getattr(fc1, "_forward_pre_hooks", None) + and fc1.fused_gate_up + and not fc1.non_gated + and fc1.out_features == 2 * inputs.shape[-2] + and getattr(dispatcher, "ep_size", None) == 1 + and getattr(dispatcher, "tp_size", None) == 1 + and getattr(dispatcher, "num_local_experts", 0) > 1 + and getattr(config, "moe_permute_fusion", False) + and getattr(experts, "offload_expert_fc1", None) is False + and getattr(experts, "offload_moe_act", None) is False + and getattr(experts, "activation_recompute", None) is False + ): + # The two dispatched H-wide inputs and FC1 gate/up sum + # remain live at the FC2 sum, including in the observed + # compiled path. This is one stage, not a backward bound. + features += 2 * fc2.out_features + fc1.out_features coefficient = max( coefficient, config.moe_router_topk * features * weights.element_size(), @@ -1422,6 +1443,8 @@ def __init__(self, runtime: TrainingRuntime) -> None: self._hybridep_buffer_id: int | None = None self._hybridep_rows_high_water = 0 self._memory_profiles: dict[_MemorySignature, _MemoryProfile] = {} + self._split_memory_floors: dict[bytes, int] = {} + self._split_memory_floor_status = "not_observed" self._last_global_micro_batch_size: int | None = None self._skipped_forward_waves: dict[object, tuple[int, int, int]] = {} # Bounded LRU: steady-state hits are temporally local (identical @@ -2271,12 +2294,13 @@ def _forward_micro_batches( ) ) else: - tracked_outputs = self._execute_admitted_plan( - candidate.plan, - check=candidate.check, - context="forward_micro_batches", + tracked_outputs, memory_baseline, forward_peak = ( + self._execute_split_plan_with_memory_tracking( + candidate.plan, + check=candidate.check, + context="forward_micro_batches", + ) ) - memory_baseline = None flat_outputs = iter(tracked_outputs) outputs = [_unflatten(item, flat_outputs) for item in candidate.inputs] stop = start + candidate.stats_global_count @@ -2324,6 +2348,10 @@ def _forward_micro_batches( # 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) + elif memory_baseline is not None: + self._record_split_memory_floor( + candidate.plan, memory_baseline, forward_peak + ) # Only the caller may retain completed outputs into the next wave. del tracked_outputs, flat_outputs, outputs start = stop @@ -2411,14 +2439,27 @@ def _execute_admitted_plan( plan, check=check, context=context ) return outputs + outputs, _baseline, _peak = self._execute_split_plan_with_memory_tracking( + plan, check=check, context=context + ) + return outputs + + def _execute_split_plan_with_memory_tracking( + self, plan: _SplitForwardPlan, *, check: _MemoryCheck, context: str + ) -> tuple[list[AnyForwardOutput], int | None, int]: + baseline, peak = None, 0 merged: list[AnyForwardOutput | None] = [None] * plan.request_count for ordinal, (subforward, indices) in enumerate( zip(plan.subforwards, plan.request_indices, strict=True) ): try: - outputs, _baseline = self._run_flat_plan_with_memory_tracking( + outputs, child_baseline = self._run_flat_plan_with_memory_tracking( subforward, check=check, context=context ) + if child_baseline is not None: + if baseline is None: + baseline = child_baseline + peak = max(peak, int(torch.cuda.max_memory_allocated(self.device))) except TrainerRankMemoryError as error: # Model execution already began, so no replanning is possible # and the caller must not mistake this for an up-front refusal. @@ -2434,7 +2475,7 @@ def _execute_admitted_plan( merged[index] = output if any(output is None for output in merged): raise AssertionError("split execution did not cover every request") - return cast(list[AnyForwardOutput], merged) + return cast(list[AnyForwardOutput], merged), baseline, peak def _plan_admissible_forward( self, @@ -2589,21 +2630,16 @@ def _admit_split_rung( for chunk in chunks ] costs = [self._plan_cost(plan) for plan in plans] - check = self._split_rung_check(costs) + # Bind original request mappings; floor keys normalize execution order. + order = sorted(range(len(plans)), key=lambda i: (-costs[i].ephemeral, i)) + split = _SplitForwardPlan( + subforwards=tuple(plans[i] for i in order), + request_indices=tuple(tuple(chunks[i]) for i in order), + request_count=len(requests), + ) + check = self._split_plan_memory_check(split, costs) if check.fits: - # Larger ephemeral first minimizes the running forward peak; - # ties keep chunk order, so the partition is deterministic. - order = sorted( - range(len(plans)), key=lambda i: (-costs[i].ephemeral, i) - ) - return ( - _SplitForwardPlan( - subforwards=tuple(plans[i] for i in order), - request_indices=tuple(tuple(chunks[i]) for i in order), - request_count=len(requests), - ), - check, - ) + return split, check return None, check def _split_rung_check(self, costs: Sequence[_SubforwardCost]) -> _MemoryCheck: @@ -2611,6 +2647,160 @@ def _split_rung_check(self, costs: Sequence[_SubforwardCost]) -> _MemoryCheck: sum(cost.retained for cost in costs) + max(cost.ephemeral for cost in costs) ) + @staticmethod + def _split_memory_key(plan: _SplitForwardPlan) -> bytes | None: + # Bound traversal and hashing; retain only a digest, never tensor values, + # token tables or graphs. Oversized structural identities stay unlearned. + if len(plan.subforwards) > 1024: + return None + digest, remaining, nodes = ( + hashlib.sha256(), + 262144 - 32 * len(plan.subforwards), + 65536, + ) + + def feed(value: Any) -> None: + nonlocal remaining, nodes + nodes -= 1 + if nodes < 0: + raise ValueError("split key node cap") + if isinstance(value, tuple): + remaining -= 2 + if remaining < 0: + raise ValueError("split key byte cap") + feed(len(value)) + digest.update(b"(") + for child in value: + feed(child) + digest.update(b")") + return + if value is not None and type(value) not in (str, int, bool): + raise ValueError("unsupported split key field") + if isinstance(value, str) and len(value) > 4096: + raise ValueError("split key string cap") + if isinstance(value, int) and value.bit_length() > 64: + raise ValueError("split key integer cap") + encoded = repr(value).encode() + remaining -= len(encoded) + 1 + if remaining < 0: + raise ValueError("split key byte cap") + digest.update(encoded + b";") + + try: + feed((plan.request_count, len(plan.subforwards))) + outer, children = digest, [] + for p, indices in zip(plan.subforwards, plan.request_indices, strict=True): + digest = hashlib.sha256() + signature = p.signature + feed( + ( + indices, + signature.topology, + signature.planner_coefficients, + signature.slot_group_count, + signature.request_mix, + signature.grad_enabled, + signature.grad_modes, + p.packed_tokens, + p.logical_tokens, + p.inactive_logical_tokens, + p.output_bytes, + p.output_metadata, + p.selected_max_depth, + len(p.groups), + ) + ) + for g in p.groups: + slot = ( + None + if g.slot_ref is None + else ("checkpoint", g.slot_ref.name) + if isinstance(g.slot_ref, _LocalLoRASlotRef) + else (g.slot_ref.kind, g.slot_ref.name) + ) + feed( + ( + slot, + g.grad_enabled, + g.request_indices, + len(g.packed.segments), + ) + ) + for segment in g.packed.segments: + feed( + ( + segment.sequence_indices, + segment.start, + segment.end, + segment.packed_start, + segment.group_id, + segment.parent_id, + ) + ) + feed(len(g.items)) + for item in g.items: + feed( + ( + tuple(item.input_ids.shape), + str(item.input_ids.dtype), + None + if item.labels is None + else (tuple(item.labels.shape), str(item.labels.dtype)), + item.request.top_k, + item.request.logits, + item.request.hidden_states, + ) + ) + children.append(digest.digest()) + except ValueError: + return None + # Cost/profile updates may reorder the same children. Each digest still + # binds its original request mapping/partition; retain the observed max + # across execution orders, not an unmeasured bound for every order. + outer.update(b"".join(sorted(children))) + return outer.digest() + + def _record_split_memory_floor( + self, plan: _SplitForwardPlan, baseline: int, forward_peak: int + ) -> None: + # Only normal caller completion learns a floor; child retained profiles + # stay forward-only. Keep earlier child peaks despite their later resets. + key = self._split_memory_key(plan) + if key is None: + self._split_memory_floor_status = "unsupported_key" + return + if ( + key not in self._split_memory_floors + and len(self._split_memory_floors) >= 1024 + ): + self._split_memory_floor_status = "cache_full_not_learned" + return + observed = max( + 0, + max(forward_peak, int(torch.cuda.max_memory_allocated(self.device))) + - baseline, + ) + self._split_memory_floors[key] = max( + self._split_memory_floors.get(key, 0), observed + ) + self._split_memory_floor_status = "recorded" + + def _split_plan_memory_check( + self, plan: _SplitForwardPlan, costs: Sequence[_SubforwardCost] + ) -> _MemoryCheck: + required = sum(cost.retained for cost in costs) + max( + cost.ephemeral for cost in costs + ) + key = self._split_memory_key(plan) + empirical = ( + 0 + if key is None + else int(self._split_memory_floors.get(key, 0) * _MEMORY_SAFETY_FACTOR) + ) + # Same existing reduction order and count; never reduce while returning + # from caller work, where a failed/empty peer may not participate. + return self._memory_check_required(max(required, empirical)) + def _split_chunk_lower_cost( self, requests: Sequence[AnyForwardInput], @@ -2649,23 +2839,6 @@ def _split_chunk_lower_cost( logical_tokens=logical_tokens, ) profile = self._memory_profiles.get(signature) - if profile is not None: - cap = profile.packed_tokens * _MEMORY_PROFILE_TRUST_GROWTH - if packed_tokens <= cap < unshared_packed_tokens: - # Required cost can drop when a larger layout leaves the - # profile window. Cold cost grows with packed tokens, so its - # first integer count bounds every possible post-cap layout, - # even when TP padding makes that count itself unattainable. - cold = self._subforward_cost( - packed_tokens=cap + 1, - output_bytes=output_bytes, - signature=signature, - logical_tokens=logical_tokens, - ) - cost = _SubforwardCost( - required=min(cost.required, cold.required), - retained=min(cost.retained, cold.retained), - ) if ( profile is not None and profile.retained_compute_bytes_per_token is not None @@ -4839,10 +5012,10 @@ def _estimate_required_memory_bytes_from_values( profiled_tokens = max( packed_tokens, logical_tokens / profiled.logical_per_packed ) - if ( - profiled is None - or profiled.packed_tokens * _MEMORY_PROFILE_TRUST_GROWTH < packed_tokens - ): + # 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: compute = max( diff --git a/tests/unit/test_trainer_rank_active_memory.py b/tests/unit/test_trainer_rank_active_memory.py index a81905fe1..46bab4762 100644 --- a/tests/unit/test_trainer_rank_active_memory.py +++ b/tests/unit/test_trainer_rank_active_memory.py @@ -7,7 +7,13 @@ import pytest import torch -from art.trainer_rank import ForwardInput, ForwardOutput, TrainerRank, Unset +from art.trainer_rank import ( + ForwardInput, + ForwardOutput, + TrainerRank, + TrainerRankMemoryError, + Unset, +) class _Model(torch.nn.Module): @@ -199,3 +205,58 @@ def test_distinct_output_and_pure_grad_modes_require_their_own_profile(): assert not rank._all_ranks_have_memory_profile( packed_tokens=plan.packed_tokens, signature=plan.signature ) + + +@pytest.mark.parametrize("no_grad", [False, True]) +def test_direct_forward_does_not_drop_observed_peak_outside_trust(monkeypatch, no_grad): + rank = _rank() + + def request(length): + tokens = torch.arange(length) + return ForwardInput(input_tokens=tokens, target_tokens=tokens, no_grad=no_grad) + + observed = rank._plan_flat_forward([request(10)]) + rank._update_memory_profile(observed, 10_000, retained_bytes=1000) + monkeypatch.setattr(rank, "_available_memory_bytes", lambda: 8000) + + def unexpected_execution(*args, **kwargs): + raise AssertionError("An unsafe single request reached model execution") + + monkeypatch.setattr( + rank, "_run_flat_plan_with_memory_tracking", unexpected_execution + ) + for length in (80, 81): + candidate = rank._plan_flat_forward([request(length)]) + assert rank._all_ranks_have_memory_profile( + packed_tokens=length, signature=candidate.signature + ) == (length == 80) + # Calibration trust still ends at 8x. Crossing it must not discard a + # peak already observed at a smaller size and admit the larger request. + with pytest.raises( + TrainerRankMemoryError, match="single request cannot be split" + ): + rank.dp_rank_forward([request(length)]) + assert rank.last_forward_telemetry()["predicted_peak_bytes"] >= 10_000 + + +@pytest.mark.parametrize("logical_ratio", [1, 2, 10]) +def test_empirical_estimate_survives_packed_trust_boundary(logical_ratio): + rank = _rank() + observed = rank._plan_flat_forward(_requests("target_tokens")) + rank._update_memory_profile(observed, 10_000, retained_bytes=1000) + estimate = rank._estimate_required_memory_bytes_from_values + values = [ + estimate( + packed_tokens=count, + logical_tokens=count * logical_ratio, + output_bytes=count * 4, + signature=observed.signature, + ) + for count in (8, 63, 64, 65, 800) + ] + assert values == sorted(values) + rate = rank._memory_profiles[observed.signature].bytes_per_token + assert values[-1] == int((800 * 4 + rate * 800 * logical_ratio) * 1.1) + assert not rank._all_ranks_have_memory_profile( + packed_tokens=800, signature=observed.signature + ) diff --git a/tests/unit/test_trainer_rank_moe_memory.py b/tests/unit/test_trainer_rank_moe_memory.py index e37d9d6a7..03f3b6d3c 100644 --- a/tests/unit/test_trainer_rank_moe_memory.py +++ b/tests/unit/test_trainer_rank_moe_memory.py @@ -233,7 +233,7 @@ def test_topk_not_expert_count_controls_envelope(layer): assert _moe_output_bytes_per_token([layer], shape) == old // 2 -def test_profiles_outputs_trust_and_empty_plan_unchanged(): +def test_profiles_outputs_and_empty_plan_preserve_empirical_floor(): rank = _rank() signature = _signature() assert rank._moe_output_bytes_per_token == 0 @@ -248,12 +248,12 @@ def test_profiles_outputs_trust_and_empty_plan_unchanged(): packed_tokens=100, logical_per_packed=1, ) - for tokens in (100, 800): + for tokens in (100, 800, 801): assert estimate( packed_tokens=tokens, output_bytes=123, signature=signature ) == int((tokens * 100000 + 123) * 1.1) - assert estimate(packed_tokens=801, output_bytes=123, signature=signature) == int( - (801 * 65536 + 123) * 1.1 + assert not rank._all_ranks_have_memory_profile( + packed_tokens=801, signature=signature ) assert estimate( packed_tokens=100, logical_tokens=200, output_bytes=123, signature=signature @@ -382,3 +382,77 @@ def reduce(value, op, group): (float(estimate), impl.dist.ReduceOp.MAX, group), (6_800_000_000.0, impl.dist.ReduceOp.MIN, group), ] + + +def _enclosing_moe(layer): + from art.megatron.lora import MLPExpertsLinearFC1LoRA + + fc1 = MLPExpertsLinearFC1LoRA.__new__(MLPExpertsLinearFC1LoRA) + torch.nn.Module.__init__(fc1) + fc1.fused_gate_up = True + fc1.non_gated = False + fc1.out_features = 1024 + layer.experts.linear_fc1 = fc1 + layer.experts.offload_expert_fc1 = False + layer.experts.offload_moe_act = False + layer.experts.activation_recompute = False + layer.token_dispatcher.ep_size = 1 + layer.token_dispatcher.tp_size = 1 + layer.token_dispatcher.num_local_experts = 256 + layer.config.moe_permute_fusion = True + return layer + + +def test_compiled_moe_retained_inputs_cold_floor(layer): + # Three complete intervals in a retained native trace have these seven + # distinct storages live together. Their sum is not the whole-model peak. + rank = _rank(_enclosing_moe(layer)) + rows, hidden, intermediate = 45_981 * 8, 2048, 512 + widths = (hidden, hidden, 2 * intermediate, intermediate, hidden, hidden, hidden) + storages = [ + torch.empty(rows, width, dtype=torch.bfloat16, device="meta") + for width in widths + ] + assert len({tensor.untyped_storage()._cdata for tensor in storages}) == 7 + observed_component_bytes = sum( + tensor.untyped_storage().nbytes() for tensor in storages + ) + assert observed_component_bytes == 8_663_556_096 + assert rank._moe_output_bytes_per_token * 45_981 == observed_component_bytes + for no_grad in (False, True): + signature = replace( + _signature(), grad_enabled=not no_grad, grad_modes=(not no_grad,) + ) + assert rank._estimate_required_memory_bytes_from_values( + packed_tokens=45_981, output_bytes=0, signature=signature + ) == int(observed_component_bytes * 1.1) + + +@pytest.mark.parametrize( + "site,attribute,value", + [ + ("dispatcher", "ep_size", 2), + ("dispatcher", "tp_size", 2), + ("dispatcher", "num_local_experts", 1), + ("config", "moe_permute_fusion", False), + ("experts", "offload_expert_fc1", True), + ("experts", "offload_moe_act", True), + ("experts", "activation_recompute", True), + ("fc1", "fused_gate_up", False), + ("fc1", "non_gated", True), + ("fc1", "out_features", 2048), + ("fc1", "forward", lambda *args: None), + ], +) +def test_unknown_enclosing_lifetimes_keep_previous_fc2_floor( + layer, site, attribute, value +): + _enclosing_moe(layer) + sites = { + "dispatcher": layer.token_dispatcher, + "config": layer.config, + "experts": layer.experts, + "fc1": layer.experts.linear_fc1, + } + setattr(sites[site], attribute, value) + assert _moe_output_bytes_per_token([layer], ParallelShape(tp=1, cp=1)) == 106_496 diff --git a/tests/unit/test_trainer_rank_split.py b/tests/unit/test_trainer_rank_split.py index a0ab095c1..ce2ad6cdb 100644 --- a/tests/unit/test_trainer_rank_split.py +++ b/tests/unit/test_trainer_rank_split.py @@ -604,11 +604,35 @@ def test_retained_ratio_lower_bound_reaches_exact_split_admission( plan, check = rank._plan_admissible_forward( requests, checkpoint=Unset, context="test" ) - assert isinstance(plan, _SplitForwardPlan) and plan.subforward_count == 2 - assert check == exact + assert isinstance(plan, _SplitForwardPlan) + # The smaller profile cannot discount larger layouts after its trust + # window; the original 20 GiB budget now requires four subforwards. + assert plan.subforward_count == (2 if profile_packed == 8000 else 4) + assert check == rank._split_rung_check( + [rank._plan_cost(child) for child in plan.subforwards] + ) assert sorted( index for chunk in plan.request_indices for index in chunk ) == list(range(32)) + retained = 0 + leaf_ids = {} + for child, indices in zip(plan.subforwards, plan.request_indices, strict=True): + cost = rank._plan_cost(child) + assert retained + cost.required <= check.estimated_required_bytes <= budget + retained += cost.retained + assert child.request_count == len(indices) + leaf_ids[id(child)] = indices + + def run(child, **kwargs): + assert kwargs["check"] == check + return [ + ForwardOutput(None, None, None, None, checkpoint=str(index)) + for index in leaf_ids[id(child)] + ], None + + monkeypatch.setattr(rank, "_run_flat_plan_with_memory_tracking", run) + outputs = rank._execute_admitted_plan(plan, check=check, context="test") + assert [output.checkpoint for output in outputs] == list(map(str, range(32))) lower = rank._split_rung_check( [ rank._split_chunk_lower_cost( @@ -621,10 +645,11 @@ def test_retained_ratio_lower_bound_reaches_exact_split_admission( ) assert [child.packed_tokens for child in children] == [64_000, 64_000] assert lower.estimated_required_bytes <= exact.estimated_required_bytes - assert lower.fits and exact.fits == admit - if not admit: - # The optimistic bound survives, but exact pricing must reject this - # rung. A later rung with smaller chunks may still fit this budget. + assert lower.fits == (not admit or profile_packed == 8000) + assert exact.fits == (admit and profile_packed == 8000) + if not exact.fits: + # Reject this rung, at the lower bound or exact pricing; a later rung + # with smaller chunks may still fit this same budget. split, rejected = rank._admit_split_rung( [tuple(range(16)), tuple(range(16, 32))], requests, @@ -730,8 +755,8 @@ def test_packed_profile_bound_handles_unattainable_tp_boundary_count( assert (minimal.packed_tokens, minimal.logical_tokens) == (8, 112) rank._memory_profiles[minimal.signature] = _MemoryProfile(profile_rate, 4) lower = rank._split_chunk_lower_cost(requests, [tokens] * 16, checkpoint=Unset) - # The cutoff is32. A hypothetical cold count33 is cheaper than the first - # physically possible post-cutoff count36, so it remains a valid bound. + # The calibration cutoff is32, with the first physically possible larger + # count36. Required cost retains its empirical floor across that boundary. for packed in range(8, 113, 4): exact = rank._subforward_cost( packed_tokens=packed, @@ -767,7 +792,7 @@ def test_packed_profile_bound_counts_each_groups_tp_padding( requests, [request.input_tokens for request in requests], checkpoint=Unset ) exact = rank._plan_cost(maximum) - assert lower.required <= exact.required < rank._plan_cost(minimal).required + assert lower.required == rank._plan_cost(minimal).required < exact.required assert lower.retained <= exact.retained @@ -914,6 +939,7 @@ def test_micro_batch_refusal_replaces_previous_admission_telemetry( @dataclass(frozen=True) class _SlotRef: name: str | None + kind: str = "checkpoint" def test_split_subforwards_track_independent_slot_graphs( diff --git a/tests/unit/test_trainer_rank_split_peak.py b/tests/unit/test_trainer_rank_split_peak.py new file mode 100644 index 000000000..3ce5ed547 --- /dev/null +++ b/tests/unit/test_trainer_rank_split_peak.py @@ -0,0 +1,410 @@ +"""CPU split-admission regressions; counter observations are not native peaks.""" + +from collections import namedtuple +from contextlib import nullcontext +from dataclasses import replace +from itertools import permutations +from types import SimpleNamespace +from typing import Any, cast + +import pytest +import torch + +from art.trainer_rank import _impl as tr + + +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 tr.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(count=2, length=100): + return [ + tr.ForwardInput( + input_tokens=torch.tensor([10_000 + i, *range(1, length)]), + target_tokens=torch.tensor([10_000 + i, *range(1, length)]), + ) + for i in range(count) + ] + + +def _native_slot_fields(monkeypatch, rank): + # Isolate other regressions from the separate CPU fallback test. These are + # the native LoRASlotRef's two scalar fields, without importing Megatron. + NativeSlotFields = namedtuple("NativeSlotFields", "kind name") + monkeypatch.setattr( + rank, "_slot_ref", lambda name: NativeSlotFields("checkpoint", name) + ) + + +def _split(rank, requests, chunks, *, memory_minimal=False): + return tr._SplitForwardPlan( + tuple( + rank._plan_flat_forward( + [requests[i] for i in chunk], memory_minimal=memory_minimal + ) + for chunk in chunks + ), + tuple(chunks), + len(requests), + ) + + +def test_local_checkpoint_fallback_still_admits_split(monkeypatch): + rank = _rank() + # This is the existing _slot_ref outcome when optional Megatron is absent. + monkeypatch.setattr(rank, "_slot_ref", lambda name: tr._LocalLoRASlotRef(name)) + monkeypatch.setattr( + rank, + "_estimate_required_memory_bytes_from_values", + lambda *, packed_tokens, **_: packed_tokens, + ) + monkeypatch.setattr(rank, "_retained_memory_bytes", lambda *a, **k: 0) + monkeypatch.setattr(rank, "_available_memory_bytes", lambda: 100) + requests = _requests() + plan, check = rank._admit_split_rung( + ((0,), (1,)), + requests, + [r.input_tokens for r in requests], + checkpoint=tr.Unset, + ) + assert isinstance(plan, tr._SplitForwardPlan) + assert check.fits + assert all(isinstance(g.slot_ref, tr._LocalLoRASlotRef) for g in plan.groups) + + +@pytest.mark.parametrize("shared_prefix", [False, True]) +def test_128_real_request_fields_have_a_bounded_split_key(monkeypatch, shared_prefix): + rank = _rank() + _native_slot_fields(monkeypatch, rank) + requests = _requests(128, 3) + if shared_prefix: + requests = [ + replace(r, input_tokens=torch.tensor([1, 2, i + 3])) + for i, r in enumerate(requests) + ] + # Exercise actual request mix, modes, output metadata, segment construction, + # coefficient selection and tensor shapes rather than hand-built plans. + requests = [ + replace(r, no_grad=bool(i % 2), hidden_states=True) + for i, r in enumerate(requests) + ] + plan = _split( + rank, + requests, + (tuple(range(64)), tuple(range(64, 128))), + memory_minimal=shared_prefix, + ) + key = rank._split_memory_key(plan) + assert isinstance(key, bytes) and len(key) == 32 + assert key == rank._split_memory_key(plan) + changed = replace( + plan.subforwards[0], output_bytes=plan.subforwards[0].output_bytes + 4 + ) + assert ( + rank._split_memory_key( + replace(plan, subforwards=(changed, plan.subforwards[1])) + ) + != key + ) + + +def test_profile_order_change_cannot_drop_completed_split_floor(monkeypatch): + rank = _rank() + _native_slot_fields(monkeypatch, rank) + requests = _requests() + requests[1] = replace(requests[1], no_grad=True) + plan = _split(rank, requests, ((0,), (1,))) + a, b = plan.subforwards + assert a.signature != b.signature + # Distinct real signatures can learn distinct rates. Keep the static model + # floor small to isolate profile-driven reordering, not kernel accounting. + rank._hidden_size = rank._param_dtype_size = 1 + rank._num_layers = 0 + rank._memory_profiles[a.signature] = tr._MemoryProfile( + 20, 100, retained_compute_bytes_per_token=1 + ) + rank._memory_profiles[b.signature] = tr._MemoryProfile( + 10, 100, retained_compute_bytes_per_token=1 + ) + assert rank._plan_cost(a).ephemeral > rank._plan_cost(b).ephemeral + monkeypatch.setattr(torch.cuda, "max_memory_allocated", lambda _: 10_100) + rank._record_split_memory_floor(plan, 100, 1_200) + # The real profile update changes only cost/order, not requests or geometry. + rank._update_memory_profile(b, 3_000, retained_bytes=500) + assert rank._plan_cost(b).ephemeral > rank._plan_cost(a).ephemeral + before = dict(rank._memory_profiles) + monkeypatch.setattr(rank, "_available_memory_bytes", lambda: 10_000) + accepted, check = rank._admit_split_rung( + ((0,), (1,)), + requests, + [r.input_tokens for r in requests], + checkpoint=tr.Unset, + ) + assert accepted is None + assert check.estimated_required_bytes >= 11_000 + assert rank._memory_profiles == before + + +def _counter_split(monkeypatch): + rank = _rank() + monkeypatch.setattr(rank, "_dp_rank_and_size", lambda: (0, 1)) + _native_slot_fields(monkeypatch, rank) + requests = _requests() + child = rank._plan_flat_forward(requests[:1]) + rank._memory_profiles[child.signature] = tr._MemoryProfile( + 0, + 100, + retained_compute_bytes_per_token=0, + ) + monkeypatch.setattr(rank, "_available_memory_bytes", lambda: 10_000) + counters: dict[str, Any] = dict(allocated=100, peak=100, resets=[], executed=0) + monkeypatch.setattr(tr, "_telemetry_phase", lambda *a, **k: nullcontext()) + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "synchronize", lambda _: None) + monkeypatch.setattr(torch.cuda, "memory_allocated", lambda _: counters["allocated"]) + monkeypatch.setattr(torch.cuda, "max_memory_allocated", lambda _: counters["peak"]) + + def reset(_): + counters["resets"].append(counters["allocated"]) + counters["peak"] = counters["allocated"] + + def execute(plan): + counters["executed"] += 1 + if counters["executed"] == counters.get("fail_at"): + counters["profiles_before_failure"] = dict(rank._memory_profiles) + raise counters["error"] + counters["peak"] = counters["allocated"] + 8_000 + counters["allocated"] += 500 + return [ + tr.ForwardOutput(torch.zeros(100), None, None, None) + for _ in range(plan.request_count) + ] + + monkeypatch.setattr(torch.cuda, "reset_peak_memory_stats", reset) + monkeypatch.setattr(rank, "_execute_flat_plan", execute) + rank.device = torch.device("cuda") # Only injected counters; no CUDA tensors. + return rank, requests, counters + + +def test_completed_iterator_preserves_caller_peak_for_next_admission(monkeypatch): + rank, requests, counters = _counter_split(monkeypatch) + iterator = rank.forward_micro_batches([requests], yield_empty=True) + batch = next(iterator) + assert batch.stats.subforward_count == counters["executed"] == 2 + assert counters["resets"] == [100, 600] + children = dict(rank._memory_profiles) + counters["peak"] = 21_000 + assert list(iterator) == [] + assert rank._memory_profiles == children + del batch + counters["allocated"] = 100 + with pytest.raises(tr.TrainerRankMemoryError): + next(rank.forward_micro_batches([requests], yield_empty=True)) + assert counters["executed"] == 2 + + +@pytest.mark.parametrize("termination", ["throw", "close"]) +def test_incomplete_caller_does_not_learn_split_peak(monkeypatch, termination): + rank, requests, counters = _counter_split(monkeypatch) + iterator = rank.forward_micro_batches([requests], yield_empty=True) + batch = next(iterator) + assert batch.stats.subforward_count == counters["executed"] == 2 + children = dict(rank._memory_profiles) + counters["peak"] = 21_000 + if termination == "throw": + original = RuntimeError("caller failed after partial work") + with pytest.raises(RuntimeError) as caught: + iterator.throw(original) + assert caught.value is original + else: + assert iterator.close() is None + assert list(iterator) == [] + assert rank._split_memory_floors == {} + assert rank._memory_profiles == children + assert counters["executed"] == 2 + + +def test_partial_forward_does_not_learn_split_peak(monkeypatch): + rank, requests, counters = _counter_split(monkeypatch) + original = torch.cuda.OutOfMemoryError("second split child allocation") + counters.update(fail_at=2, error=original) + iterator = rank.forward_micro_batches([requests], yield_empty=True) + with pytest.raises(tr.TrainerRankPartialExecutionError) as caught: + next(iterator) + assert "1 of 2 completed" in str(caught.value) + assert isinstance(caught.value.__cause__, tr.TrainerRankMemoryError) + assert caught.value.__cause__.__cause__ is original + assert list(iterator) == [] + assert counters["executed"] == 2 + assert rank._split_memory_floors == {} + assert rank._memory_profiles == counters["profiles_before_failure"] + + +def test_empty_dp_rank_retains_global_selection_collective_sequence(monkeypatch): + # Real selection/find/rung methods with explicit scalar reductions. This is + # not native distributed convergence or a model-execution test. + traces = [] + for dp_rank in (0, 1): + with monkeypatch.context() as patch: + rank = _rank() + _native_slot_fields(patch, rank) + trace = [] + patch.setattr(rank, "_dp_rank_and_size", lambda: (dp_rank, 2)) + patch.setattr(rank, "_forward_memory_group", lambda: f"tp-cp-{dp_rank}") + patch.setattr(rank, "_available_memory_bytes", lambda: 100) + patch.setattr( + rank, + "_estimate_required_memory_bytes_from_values", + lambda *, packed_tokens, **_: packed_tokens, + ) + patch.setattr(rank, "_retained_memory_bytes", lambda *a, **k: 0) + patch.setattr(rank, "_ensure_checkpoint_slots_for", lambda *a, **k: None) + + # Mode/profile agreements are already collective in production; + # retain their sequence while controlling this two-rank witness. + def agree(value): + trace.append(("global", "agree", bool(value))) + return bool(value) + + def profiled(**_): + trace.append(("global", "profile", False)) + return False + + patch.setattr(rank, "_all_ranks_true", agree) + patch.setattr(rank, "_all_ranks_have_memory_profile", profiled) + + 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) + if op == tr.dist.ReduceOp.MAX + else min(value.item(), 100) + ) + + patch.setattr(tr.dist, "is_available", lambda: True) + patch.setattr(tr.dist, "is_initialized", lambda: True) + patch.setattr(tr.dist, "all_reduce", reduce) + candidate = rank._select_next_micro_batch([_requests()], 0) + assert candidate.check.fits + assert len(candidate.inputs) == (1 if dp_rank == 0 else 0) + assert isinstance( + candidate.plan, + tr._SplitForwardPlan if dp_rank == 0 else tr._FlatForwardPlan, + ) + traces.append([event for event in trace if event[0] == "global"]) + assert traces[0] == traces[1] + + +def test_order_normalization_keeps_mapping_partition_and_shapes(monkeypatch): + rank = _rank() + _native_slot_fields(monkeypatch, rank) + requests = _requests(3, 4) + requests[1] = replace(requests[1], no_grad=True) + plan = _split(rank, requests, ((0,), (1,), (2,))) + key = rank._split_memory_key(plan) + for order in permutations(range(3)): + q = replace( + plan, + subforwards=tuple(plan.subforwards[i] for i in order), + request_indices=tuple(plan.request_indices[i] for i in order), + ) + assert rank._split_memory_key(q) == key + # Mapping moves without the corresponding child must remain distinguishable. + remapped = replace(plan, request_indices=((1,), (0,), (2,))) + assert rank._split_memory_key(remapped) != key + assert rank._split_memory_key(_split(rank, requests, ((0, 1), (2,)))) != key + a = plan.subforwards[0] + variants = [ + replace(a, packed_tokens=a.packed_tokens + 1), + replace(a, output_bytes=a.output_bytes + 4), + replace(a, signature=replace(a.signature, topology=(1, 2, 1, 1))), + replace(a, output_metadata=(("different-checkpoint", False),)), + ] + assert all( + rank._split_memory_key(replace(plan, subforwards=(v, *plan.subforwards[1:]))) + != key + for v in variants + ) + + +def test_native_and_known_local_reference_preserve_same_floor(monkeypatch): + rank = _rank() + _native_slot_fields(monkeypatch, rank) + requests = _requests() + plan = _split(rank, requests, ((0,), (1,))) + local = replace( + plan, + subforwards=tuple( + replace( + p, + groups=tuple( + replace(g, slot_ref=tr._LocalLoRASlotRef(g.slot_ref.name)) + for g in p.groups + ), + ) + for p in plan.subforwards + ), + ) + key = rank._split_memory_key(plan) + assert rank._split_memory_key(local) == key + peak = [6100] + monkeypatch.setattr(torch.cuda, "max_memory_allocated", lambda _: peak[0]) + original = dict(rank._memory_profiles) + rank._record_split_memory_floor(plan, 100, 200) + reversed_plan = replace( + local, + subforwards=tuple(reversed(local.subforwards)), + request_indices=tuple(reversed(local.request_indices)), + ) + peak[0] = 3100 + rank._record_split_memory_floor(reversed_plan, 100, 200) + assert rank._split_memory_floors[key] == 6000 + peak[0] = 9100 + rank._record_split_memory_floor(reversed_plan, 100, 200) + assert rank._split_memory_floors[key] == 9000 + assert rank._memory_profiles == original + + +def test_large_unsupported_keys_and_full_cache_remain_explicit(monkeypatch): + rank = _rank() + _native_slot_fields(monkeypatch, rank) + plan = _split(rank, _requests(), ((0,), (1,))) + huge = replace( + plan, subforwards=(plan.subforwards[0],) * 1025, request_indices=((0,),) * 1025 + ) + assert rank._split_memory_key(huge) is None + key = rank._split_memory_key(plan) + rank._split_memory_floors = {i.to_bytes(32, "big"): i for i in range(1024)} + monkeypatch.setattr(torch.cuda, "max_memory_allocated", lambda _: 1000) + before = dict(rank._split_memory_floors) + rank._record_split_memory_floor(plan, 0, 0) + assert rank._split_memory_floor_status == "cache_full_not_learned" + assert rank._split_memory_floors == before and key not in rank._split_memory_floors + # A field-size limit applies before repr copies arbitrarily long strings. + a = replace(plan.subforwards[0], output_metadata=(("x" * 4097, False),)) + assert ( + rank._split_memory_key(replace(plan, subforwards=(a, plan.subforwards[1]))) + is None + )