diff --git a/src/art/trainer_rank/_impl.py b/src/art/trainer_rank/_impl.py index 736bc951e..5f0006c4a 100644 --- a/src/art/trainer_rank/_impl.py +++ b/src/art/trainer_rank/_impl.py @@ -941,6 +941,9 @@ class _MemorySignature: grad_enabled: bool grad_modes: tuple[bool, ...] slot_shapes: tuple[tuple[bool, tuple[tuple[int, ...], ...]], ...] = () + # Short single-target requests keep the logical extrapolation, but share + # the profile learned from longer requests of the same signature. + short_requests: bool = dataclass_field(default=False, compare=False) @dataclass(frozen=True) @@ -1791,6 +1794,8 @@ def memory_field(name: str, default: Any = None) -> Any: ) self._recompute_granularity = memory_field("recompute_granularity", None) + self._recompute_method = memory_field("recompute_method", None) + self._recompute_num_layers = memory_field("recompute_num_layers", None) self._recompute_modules: frozenset[str] = frozenset( memory_field("recompute_modules", ()) or () ) @@ -2707,6 +2712,16 @@ def forward_micro_batches( no_grad: bool | None = None, yield_empty: bool = False, ) -> Iterator[MicroBatch[ForwardInputs, ForwardOutputs]]: + """Yield admitted micro-batches; the caller runs its loss and backward. + + Admission learns each call's whole peak, including the caller's loss + and backward. For grad-enabled single-target requests of at least 64 + tokens under one-layer full recompute, it prices model memory by packed + 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. + """ if not isinstance(yield_empty, bool): raise TypeError("yield_empty must be a bool") enabled = torch.is_grad_enabled() if no_grad is None else not no_grad @@ -3340,6 +3355,7 @@ def feed(value: Any) -> None: signature.grad_enabled, signature.grad_modes, signature.slot_shapes, + signature.short_requests, p.packed_tokens, p.logical_tokens, p.inactive_logical_tokens, @@ -4052,11 +4068,18 @@ def _retained_memory_bytes( ratio = logical_tokens / max(1, packed_tokens) if ratio > profile.logical_per_packed * _MEMORY_PROFILE_TRUST_GROWTH: return required - retained = output_bytes + max( - checkpoint_retained_bytes, - profile.retained_compute_bytes_per_token - * max(packed_tokens, logical_tokens / profile.logical_per_packed), - ) + rate = profile.retained_compute_bytes_per_token + if _packed_priced(signature, self._one_layer_recompute()): + # Saved head indices, masks and caller saves stay live until + # backward. The ratio window above bounds the sharing extrapolation. + retained_compute = ( + rate * packed_tokens + _PACKED_PRICED_LOGICAL_ROW_BYTES * logical_tokens + ) + else: + retained_compute = rate * max( + packed_tokens, logical_tokens / profile.logical_per_packed + ) + retained = output_bytes + max(checkpoint_retained_bytes, retained_compute) return min(required, int(retained * _MEMORY_SAFETY_FACTOR)) def _split_request_order( @@ -6111,6 +6134,7 @@ def _fill_planner_snapshot( ) } rank_fields["recompute_modules"] = sorted(self._recompute_modules) + rank_fields["one_layer_recompute"] = self._one_layer_recompute() rank_fields["moe_forward_stages"] = getattr(self, "_moe_forward_stages", ()) rank_fields["geometry"] = asdict(self._geometry) rank_fields["topology"] = list(plan.signature.topology) @@ -6725,6 +6749,7 @@ def _memory_signature_from_requests( grad_enabled=any(modes), grad_modes=modes, slot_shapes=shapes if any(any(shape) for _, shape in shapes) else (), + short_requests=any(_short_request(request) for request in requests), ) def _slot_memory_shapes( @@ -7719,22 +7744,7 @@ def _estimate_required_memory_bytes_from_values( # state (fp32), plus convolution history. Unlike token activations, # these do not shrink with segment length. gdn_state_bytes = ( - 2 - * gdn_segments - * gdn_layers - / tp - * ( - 4 - * geometry.gdn_value_heads - * geometry.gdn_key_head_dim - * geometry.gdn_value_head_dim - + self._param_dtype_size - * ( - 2 * geometry.gdn_key_heads * geometry.gdn_key_head_dim - + geometry.gdn_value_heads * geometry.gdn_value_head_dim - ) - * max(0, geometry.gdn_conv_kernel - 1) - ) + gdn_segments * gdn_layers * self._gdn_segment_layer_bytes() ) static_compute = max( static_compute, @@ -7775,15 +7785,23 @@ def _estimate_required_memory_bytes_from_values( # Local head results coexist with full CP outputs during gathering. # Uneven rank plans can assign all of an item's rows to one rank. static_compute += output_bytes - # A profile learned under lighter sharing (lower logical/packed ratio) - # underestimates the per-packed-token footprint of a deeper-shared - # plan; scale the trusted estimate up by the ratio gap. - # Normalize before multiplying: cancelling packed tokens through two - # float operations can otherwise make a larger warm layout cheaper. + # Memory that grows with logical rows makes a profile learned under + # lighter sharing underestimate a deeper-shared plan; scale the trusted + # estimate up by the ratio gap. Packed pricing instead charges head and + # caller memory per logical row and extrapolates the rest at most the + # 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( - packed_tokens, logical_tokens / profiled.logical_per_packed + packed_tokens, + logical_tokens + / profiled.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 @@ -7793,10 +7811,67 @@ def _estimate_required_memory_bytes_from_values( else: compute = max( static_compute, - int(profiled.bytes_per_token * profiled_tokens), + int( + profiled.bytes_per_token * profiled_tokens + + ( + _PACKED_PRICED_LOGICAL_ROW_BYTES * logical_tokens + # Branch states grow with segments, not packed rows; + # backward recomputes one layer at a time. + + gdn_segments * self._gdn_segment_layer_bytes() + if packed_priced and logical_tokens is not None + else 0 + ) + ), ) return int((output_bytes + compute) * _MEMORY_SAFETY_FACTOR) + def _one_layer_recompute(self) -> bool: + """ART's default full/uniform/1 recompute, which Megatron runs in training. + + Forward keeps only layer inputs and backward recomputes one layer at a + time. Megatron checks the decoder's own mode and config, and eval mode + skips recompute even with gradients enabled, so read both live. + """ + recorded = self.__dict__.get("_recorded_one_layer_recompute") + if recorded is not None: + return recorded # Planner-report replay has no live model. + target = ("full", "uniform", 1) + names = ("recompute_granularity", "recompute_method", "recompute_num_layers") + + def active(chunk: torch.nn.Module) -> bool: + try: + decoder = _language_model(chunk).decoder + config = decoder.config + except (AttributeError, RuntimeError): + # No decoder config (stub or non-GPT chunk): stored settings, + # and every module must be training. + stored = tuple(getattr(self, "_" + name) for name in names) + return stored == target and all(m.training for m in chunk.modules()) + settings = tuple(getattr(config, name, None) for name in names) + return settings == target and decoder.training is True + + return all(active(chunk) for chunk in self.runtime.model) + + def _gdn_segment_layer_bytes(self) -> float: + """Initial and final fp32 recurrent states plus convolution history.""" + geometry = self._geometry + return ( + 2 + / max(1, self._topology_key()[1]) + * ( + 4 + * geometry.gdn_value_heads + * geometry.gdn_key_head_dim + * geometry.gdn_value_head_dim + + self._param_dtype_size + * ( + 2 * geometry.gdn_key_heads * geometry.gdn_key_head_dim + + geometry.gdn_value_heads * geometry.gdn_value_head_dim + ) + * max(0, geometry.gdn_conv_kernel - 1) + ) + ) + def _available_memory_bytes(self, sample: dict[str, Any] | None = None) -> int: if not (torch.cuda.is_available() and self.device.type == "cuda"): return 1 << 60 @@ -8703,11 +8778,45 @@ def _active_logical_tokens(requests: Sequence[AnyForwardInput]) -> int: ) +# Grad-enabled single-target training measured flat per packed row across +# sharing ratios. Wide labels, dense or top-k outputs and no_grad forwards (whose +# GDN branch states are uncharged) keep the logical/packed ratio extrapolation. +_PACKED_PRICED_MIXES = frozenset({"target:single", "inactive"}) +# Packed pricing charges head and caller memory per active logical row: label +# copies, positions, row-match vectors, saved masks and the caller's loss saves +# and backward transients. Each request's buffers are separate allocations +# rounded up to 512 B blocks; on an H200, a fully shared one-token request's +# head peaked at nine blocks plus 8 B (4,616 B). Callers whose head and loss +# memory peaks above this charge per token are unsupported. +_PACKED_PRICED_LOGICAL_ROW_BYTES = 12 * 512 +# Shorter single-target requests keep the logical extrapolation, so each +# packed-priced request brings at least 384 KiB for per-request constants. +_PACKED_PRICED_MIN_REQUEST_TOKENS = 64 + + +def _packed_priced(signature: "_MemorySignature", one_layer_recompute: bool) -> bool: + # Measured only under one-layer full recompute: other recompute modes keep + # more GDN states and activations live per segment and logical row. + return ( + one_layer_recompute + and bool(signature.grad_modes) + and all(signature.grad_modes) + and _PACKED_PRICED_MIXES.issuperset(signature.request_mix) + and not signature.short_requests + ) + + def _request_mix_key(request: AnyForwardInput) -> str: parts = [] if request.target_tokens is not None: - target = request.target_tokens - tail_shape = tuple(target.shape[request.input_tokens.ndim :]) + target, inputs = request.target_tokens, request.input_tokens + # Match _forward_item: trailing target dims follow the input shape or a + # flattened token axis. + tail_shape = tuple( + target.shape[inputs.ndim :] + if target.shape[: inputs.ndim] == inputs.shape + else target.shape[1:] + ) parts.append(f"target:{tail_shape or 'single'}") if request.top_k is not None: parts.append(f"topk:{int(request.top_k)}") @@ -8718,6 +8827,15 @@ def _request_mix_key(request: AnyForwardInput) -> str: return "+".join(parts) if parts else "inactive" +def _short_request(request: AnyForwardInput) -> bool: + """Too short for packed pricing: a duplicate adds no packed row, and caller + memory per request can outgrow its few logical-row charges.""" + return ( + _request_mix_key(request) == "target:single" + and int(request.input_tokens.numel()) < _PACKED_PRICED_MIN_REQUEST_TOKENS + ) + + def _pad_packed_batch( batch: PrefixTreePack, *, diff --git a/src/art/trainer_rank/_planner_misses.py b/src/art/trainer_rank/_planner_misses.py index 2018c6186..61b23fc8c 100644 --- a/src/art/trainer_rank/_planner_misses.py +++ b/src/art/trainer_rank/_planner_misses.py @@ -521,7 +521,8 @@ def report( _RANK_FIELDS = frozenset( "num_layers hidden_size param_dtype_size recompute_granularity " - "sequence_parallel attention_output_gate mlp_activation_factor gdn_layers " + "one_layer_recompute sequence_parallel attention_output_gate " + "mlp_activation_factor gdn_layers " "checkpointed_moe_layers recompute_modules moe_output_bytes_per_token " "moe_forward_stages".split() ) @@ -591,8 +592,11 @@ def replay( "incomplete replay: immutable rank fields differ (including MoE stages)" ) rank = _impl.TrainerRank.__new__(_impl.TrainerRank) - for name in _RANK_FIELDS: + for name in _RANK_FIELDS - {"one_layer_recompute"}: setattr(rank, "_" + name, values[name]) + if type(values["one_layer_recompute"]) is not bool: + raise ValueError("incomplete replay: recompute mode is not recorded") + rank._recorded_one_layer_recompute = values["one_layer_recompute"] rank._moe_forward_stages = tuple(tuple(row) for row in values["moe_forward_stages"]) rank._geometry = ModelGeometry(**values["geometry"]) dp, tp, cp, pp = values["topology"] diff --git a/tests/unit/test_trainer_rank_active_memory.py b/tests/unit/test_trainer_rank_active_memory.py index 5249e07f5..a2033710d 100644 --- a/tests/unit/test_trainer_rank_active_memory.py +++ b/tests/unit/test_trainer_rank_active_memory.py @@ -14,7 +14,19 @@ TrainerRank, TrainerRankMemoryError, Unset, + _impl, ) +from art.trainer_rank._impl import ( + _PACKED_PRICED_LOGICAL_ROW_BYTES, + _packed_priced, + _request_mix_key, +) + + +@pytest.fixture(autouse=True) +def _price_short_requests(monkeypatch): + # These fixtures use short requests; the short-request gate has its own tests. + monkeypatch.setattr(_impl, "_PACKED_PRICED_MIN_REQUEST_TOKENS", 1) class _Model(torch.nn.Module): @@ -36,7 +48,11 @@ def _rank(): model=[_Model()], optimizer=None, provider=SimpleNamespace( - hidden_size=8, num_layers=4, recompute_granularity="full" + hidden_size=8, + num_layers=4, + recompute_granularity="full", + recompute_method="uniform", + recompute_num_layers=1, ), model_support_handler=SimpleNamespace(build_gdn_execution_spec=False), ), @@ -256,7 +272,7 @@ def unexpected_execution(*args, **kwargs): 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) + rank._update_memory_profile(observed, 100_000, retained_bytes=1000) estimate = rank._estimate_required_memory_bytes_from_values values = [ estimate( @@ -268,13 +284,109 @@ def test_empirical_estimate_survives_packed_trust_boundary(logical_ratio): 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) + profile = rank._memory_profiles[observed.signature] + logical = 800 * logical_ratio + packed = max(800, logical / profile.logical_per_packed / 8) + rows = _PACKED_PRICED_LOGICAL_ROW_BYTES * logical + assert values[-1] == int((800 * 4 + profile.bytes_per_token * packed + rows) * 1.1) assert not rank._all_ranks_have_memory_profile( packed_tokens=800, signature=observed.signature ) +def test_packed_pricing_covers_allocator_blocks_of_fully_shared_requests(): + # A duplicate one-token request adds no packed row but still allocates its + # head buffers, each rounded up to a 512 B block: nine at the measured peak. + # (The short-request gate keeps such requests out; the charge covers them.) + rank = _rank() + observed = rank._plan_flat_forward(_requests("target_tokens")) + rank._update_memory_profile(observed, 100_000, retained_bytes=1000) + profile = rank._memory_profiles[observed.signature] + assert _packed_priced(observed.signature, rank._one_layer_recompute()) + duplicates = 100_000 + + def estimate(logical_tokens: int) -> int: + return rank._estimate_required_memory_bytes_from_values( + packed_tokens=1, + logical_tokens=logical_tokens, + output_bytes=0, + signature=observed.signature, + ) + + base = profile.logical_per_packed + assert estimate(base + duplicates) - estimate(base) >= duplicates * 9 * 512 + + +@pytest.mark.parametrize("shape", ["flat", "row"]) +def test_short_single_target_requests_keep_logical_pricing(monkeypatch, shape): + monkeypatch.setattr(_impl, "_PACKED_PRICED_MIN_REQUEST_TOKENS", 64) + + def request(length: int) -> ForwardInput: + tokens = torch.arange(length) + if shape == "row": + tokens = tokens[None] + return ForwardInput(input_tokens=tokens, target_tokens=tokens + 1) + + assert [_impl._short_request(request(n)) for n in (63, 64, 65)] == [ + True, + False, + False, + ] + assert not _impl._short_request(ForwardInput(input_tokens=torch.arange(3))) + rank = _rank() + recompute = rank._one_layer_recompute() + # Calibrate on a 64-token request, then price a short one: the short batch + # shares the calibrated profile (no cold start) but keeps main's pricing. + long_plan = rank._plan_flat_forward([request(64)]) + rank._update_memory_profile(long_plan, 64 * 100_000, retained_bytes=None) + for requests in ([request(63)], [request(64), request(63)]): + short = rank._plan_flat_forward(requests) + assert short.signature.short_requests and short.signature == long_plan.signature + assert _packed_priced(long_plan.signature, recompute) + assert not _packed_priced(short.signature, recompute) + profile = rank._memory_profiles[short.signature] + cost = rank._plan_cost(short) + extrapolated = profile.bytes_per_token * max( + short.packed_tokens, + short.active_logical_tokens / profile.logical_per_packed, + ) + assert cost.required >= int((short.output_bytes + extrapolated) * 1.1) + + +def test_packed_sharing_clamp_is_monotone_and_learning_sharing_never_cheapens(): + rank = _rank() + observed = rank._plan_flat_forward(_requests("target_tokens")) + rank._update_memory_profile(observed, 100_000, retained_bytes=1000) + single = observed.signature + rank._memory_profiles[single] = replace( + rank._memory_profiles[single], bytes_per_token=50_000, logical_per_packed=1 + ) + + def cost(packed: int, logical: int = 8_000): + return rank._subforward_cost( + packed_tokens=packed, + logical_tokens=logical, + output_bytes=0, + signature=single, + ).required + + # Fixed logical rows; packed rows sweep across L / 8r = 1000. + sweep = [cost(packed) for packed in (1, 500, 999, 1000, 1001, 4000, 8000)] + assert sweep == sorted(sweep) + rows = _PACKED_PRICED_LOGICAL_ROW_BYTES * 8_000 + assert sweep[0] == sweep[3] == int((50_000 * 1000 + rows) * 1.1) + assert sweep[4] == int((50_000 * 1001 + rows) * 1.1) + # Learning more sharing at a lower rate max-merges both; plans at + # or below the older ratio never get cheaper. + before = [cost(packed, logical) for packed, logical in ((100, 100), (50, 100))] + wider = replace(observed, packed_tokens=1, logical_tokens=8) + rank._update_memory_profile(wider, 1_000, retained_bytes=None) + profile = rank._memory_profiles[single] + assert profile.logical_per_packed > 1 and profile.bytes_per_token == 50_000 + after = [cost(packed, logical) for packed, logical in ((100, 100), (50, 100))] + assert all(b >= a for a, b in zip(before, after, strict=True)) + + @pytest.mark.parametrize( "method,argument,fallback", [ @@ -312,3 +424,151 @@ def importing(name, *args, **kwargs): with pytest.raises(type(error)) as caught: getattr(rank, method)(argument) assert caught.value is error + + +def test_packed_pricing_is_limited_to_grad_single_target_mixes(): + rank = _rank() + observed = rank._plan_flat_forward(_requests("target_tokens")) + rank._update_memory_profile(observed, 10_000, retained_bytes=1000) + single = observed.signature + assert single.grad_modes == (True,) + profile = replace( + rank._memory_profiles[single], + bytes_per_token=100_000, + retained_compute_bytes_per_token=50_000, + logical_per_packed=2, + ) + signatures = { + "single": single, + "multi_grad": replace(single, grad_modes=(True, True)), + "no_grad": replace(single, grad_modes=(False,)), + "mixed_grad": replace(single, grad_modes=(False, True)), + "hidden": replace(single, request_mix=("hidden",)), + "wide": replace(single, request_mix=("target:(2,)",)), + } + for signature in signatures.values(): + rank._memory_profiles[signature] = profile + + def estimate(name): + return rank._estimate_required_memory_bytes_from_values( + packed_tokens=8, + logical_tokens=64, + output_bytes=0, + signature=signatures[name], + ) + + # Packed rows plus head and caller memory for every logical row. + rows = _PACKED_PRICED_LOGICAL_ROW_BYTES * 64 + assert ( + estimate("single") == estimate("multi_grad") == int((100_000 * 8 + rows) * 1.1) + ) + retained = rank._retained_memory_bytes( + single, packed_tokens=8, logical_tokens=64, output_bytes=0, required=1 << 40 + ) + assert retained == int((50_000 * 8 + rows) * 1.1) + # Others keep the logical/packed extrapolation: 64 / 2 profiled rows. + for name in ("no_grad", "mixed_grad", "hidden", "wide"): + assert estimate(name) == int(100_000 * 32 * 1.1) + # Beyond 8x the profile's sharing, packed rows are priced as if sharing + # were 8x: 64 logical rows at ratio 2 x 8 = 16 is 4 rows, not 1. + clamped = rank._estimate_required_memory_bytes_from_values( + packed_tokens=1, logical_tokens=64, output_bytes=0, signature=single + ) + assert clamped == int((100_000 * 4 + rows) * 1.1) + # The logical charge does not depend on layout, so at any rate (none gates + # eligibility) cost never falls as packed rows grow. + for rate in (1, 2 * _PACKED_PRICED_LOGICAL_ROW_BYTES * 2 - 1, 100_000): + rank._memory_profiles[single] = replace( + profile, bytes_per_token=rate, retained_compute_bytes_per_token=rate + ) + costs = [ + rank._subforward_cost( + packed_tokens=packed, + logical_tokens=64, + output_bytes=0, + signature=single, + ) + for packed in range(1, 65) + ] + assert all(a.required <= b.required for a, b in zip(costs, costs[1:])) + assert all(a.retained <= b.retained for a, b in zip(costs, costs[1:])) + # GDN branch states grow with segments, not packed rows. + rank._memory_profiles[single] = profile + rank._geometry = replace( + rank._geometry, + gdn_key_heads=1, + gdn_key_head_dim=4, + gdn_value_heads=2, + gdn_value_head_dim=4, + gdn_conv_kernel=4, + ) + # 2 states x 4 B x Hv=2 x dk=4 x dv=4, plus 2 B conv history of 16 rows x 3. + assert rank._gdn_segment_layer_bytes() == 2 * (4 * 2 * 4 * 4 + 2 * 16 * 3) + live = 1 + + def segments(count): + return rank._estimate_required_memory_bytes_from_values( + packed_tokens=8, + logical_tokens=64, + output_bytes=0, + signature=single, + gdn_segments=count, + ) + + assert segments(3) - segments(0) == pytest.approx( + 3 * live * rank._gdn_segment_layer_bytes() * 1.1, abs=1 + ) + # Other recompute modes, and eval mode (which skips recompute), keep more + # live per segment and row: extrapolate. + for setting in ( + ("selective", "uniform", 1), + (None, None, None), + ("full", "block", 1), + ("full", "uniform", 2), + ("full", None, None), + ): + ( + rank._recompute_granularity, + rank._recompute_method, + rank._recompute_num_layers, + ) = setting + assert not rank._one_layer_recompute() + assert estimate("single") == estimate("hidden") + rank._recompute_granularity, rank._recompute_method = "full", "uniform" + rank._recompute_num_layers = 1 + assert rank._one_layer_recompute() + rank.runtime.model[0].eval() + assert not rank._one_layer_recompute() + assert estimate("single") == estimate("hidden") + rank.runtime.model[0].train() + # Megatron checks the decoder's own mode, which can diverge from the chunk. + rank.runtime.model[0].decoder = torch.nn.Module() + rank.runtime.model[0].decoder.eval() + assert rank.runtime.model[0].training and not rank._one_layer_recompute() + assert estimate("single") == estimate("hidden") + rank.runtime.model[0].decoder.train() + assert rank._one_layer_recompute() + # With a decoder config, its live settings and the decoder's mode decide. + decoder = rank.runtime.model[0].decoder + decoder.config = SimpleNamespace( + recompute_granularity="full", recompute_method="uniform", recompute_num_layers=1 + ) + rank.runtime.model[0].eval() + decoder.train() + assert rank._one_layer_recompute() + rank.runtime.model[0].train() + decoder.config.recompute_num_layers = 2 + assert not rank._one_layer_recompute() + decoder.config.recompute_num_layers = 1 + assert rank._one_layer_recompute() + # Replay trusts the recorded mode instead of a live model. + packed = estimate("single") + rank._recorded_one_layer_recompute = False + assert estimate("single") == estimate("hidden") != packed + rank._recorded_one_layer_recompute = True + assert estimate("single") == packed + del rank._recorded_one_layer_recompute + # Flattened-axis wide labels are not single-target. + tokens = torch.arange(4).reshape(1, 4) + wide = ForwardInput(input_tokens=tokens, target_tokens=torch.zeros(4, 3).long()) + assert _request_mix_key(wide) == "target:(3,)" diff --git a/tests/unit/test_trainer_rank_head_memory.py b/tests/unit/test_trainer_rank_head_memory.py index 50c16343f..21b7ad071 100644 --- a/tests/unit/test_trainer_rank_head_memory.py +++ b/tests/unit/test_trainer_rank_head_memory.py @@ -8,7 +8,11 @@ import torch from art.trainer_rank import ForwardInput -from art.trainer_rank._impl import Unset, _MemoryProfile +from art.trainer_rank._impl import ( + _PACKED_PRICED_LOGICAL_ROW_BYTES, + Unset, + _MemoryProfile, +) def rank(): @@ -151,7 +155,9 @@ def test_outputs_retention_and_empirical_peak_are_counted_once(): retained_compute_bytes_per_token=1, ) cost = r._plan_cost(plan) - assert cost.required == int((plan.output_bytes + 512 * 2_000_000) * 1.1) + # Packed pricing adds head and caller memory for every logical row. + rows = _PACKED_PRICED_LOGICAL_ROW_BYTES * 512 + assert cost.required == int((plan.output_bytes + 512 * 2_000_000 + rows) * 1.1) assert cost.retained == int((plan.output_bytes + retained) * 1.1) diff --git a/tests/unit/test_trainer_rank_head_recompute.py b/tests/unit/test_trainer_rank_head_recompute.py index 20e3d7c85..3e40d2fa8 100644 --- a/tests/unit/test_trainer_rank_head_recompute.py +++ b/tests/unit/test_trainer_rank_head_recompute.py @@ -1,6 +1,7 @@ from __future__ import annotations from contextlib import nullcontext +from dataclasses import replace from datetime import timedelta import sys from types import SimpleNamespace @@ -436,3 +437,122 @@ def exp(subtraction): assert model.output_layer.weight.grad is not None assert model.output_layer.weight.grad.isfinite().all() assert torch.cuda.is_initialized() == cuda_initialized + + +_PRODUCTION_HEAD_CHUNK_TOKENS = _impl._HEAD_CHUNK_TOKENS + + +def _cispo(output, device): + # 061's caller: per-history sampled mask, detached clipped ratio. + sampled = torch.ones(output.target_logprobs.shape, dtype=torch.bool, device=device) + logp = output.target_logprobs[sampled] + ratio = (logp.detach() - 0.5).exp().clamp(max=5.0) + return -(ratio * logp).sum() + + +_CALLER_LOSSES = { + "sum": lambda output, device, refs: -output.target_logprobs.sum(), + "cispo": lambda output, device, refs: _cispo(output, device), + # Saves a fixed 256 KiB per request, whatever its length. + "fixed": lambda output, device, refs: ( + (output.target_logprobs.sum() - refs).square().mean() + ), +} + + +def _duplicate_request_peak(monkeypatch, count: int, length: int, loss: str) -> int: + """Peak bytes through the real head, a caller loss and backward.""" + _patch_local_head(monkeypatch) + monkeypatch.setattr(_impl, "_HEAD_CHUNK_TOKENS", _PRODUCTION_HEAD_CHUNK_TOKENS) + device = torch.device("cuda") + weight = torch.randn(1024, 64, device=device, dtype=torch.bfloat16) / 8 + model = SimpleNamespace( + output_layer=_Head(weight), + vocab_size=1024, + share_embeddings_and_output_weights=False, + _scale_logits=lambda value: value, + ) + r = object.__new__(TrainerRank) + r.runtime = SimpleNamespace(model=[model]) + tokens = torch.arange(length) + request = ForwardInput(input_tokens=tokens, target_tokens=tokens + 1) + items = [r._forward_item(request) for _ in range(count)] + prepared = SimpleNamespace( + positions_by_item=tuple(torch.arange(length) for _ in range(count)), + source_positions_by_item=tuple(torch.arange(length) for _ in range(count)), + ) + hidden = torch.randn(length, 64, device=device, dtype=torch.bfloat16) + hidden.requires_grad_() + refs = torch.zeros(64 * 1024, device=device) + torch.cuda.synchronize() + torch.cuda.reset_peak_memory_stats() + base = torch.cuda.memory_allocated() + total = torch.zeros((), device=device) + for output in r._project_head(items, prepared, hidden): + total = total + _CALLER_LOSSES[loss](output, device, refs) + total.backward() + torch.cuda.synchronize() + assert hidden.grad is not None + return torch.cuda.max_memory_allocated() - base + + +@pytest.mark.parametrize( + ("length", "loss", "fits"), + [ + (1, "sum", True), + (64, "cispo", True), + (512, "cispo", True), + (64, "fixed", True), + # Why short requests keep logical pricing: per-request caller memory + # outgrows a one-token request's row charge. + (1, "fixed", False), + ], +) +def test_shared_requests_fit_the_per_row_charge(monkeypatch, length, loss, fits): + # A duplicate request adds logical rows but no packed rows. Its head buffers + # and the caller's saves are separate CUDA allocations rounded to 512 B + # blocks; through backward they must fit the charge on its logical rows. + if not torch.cuda.is_available(): + pytest.skip("allocator rounding needs CUDA") + _duplicate_request_peak(monkeypatch, 1, length, loss) + extra = 20_000 if length == 1 and loss == "sum" else 1_000 + single = _duplicate_request_peak(monkeypatch, 1, length, loss) + per_request = ( + _duplicate_request_peak(monkeypatch, 1 + extra, length, loss) - single + ) / extra + # Nine head blocks per one-token request at least; the rest is the caller. + assert per_request >= 9 * 512 + charge = _impl._PACKED_PRICED_LOGICAL_ROW_BYTES * length + assert (per_request <= charge) == fits + + +def test_calibrated_profile_admits_shorter_shared_requests_above_their_peak( + monkeypatch, +): + # A profile learned from 16 shared 4,096-token requests prices 20 shared + # 64-token requests (sharing 20 against 16) whose caller saves a fixed + # 256 KiB each: 64 B per calibration token but 4 KiB per new token. + if not torch.cuda.is_available(): + pytest.skip("allocator rounding needs CUDA") + from test_trainer_rank_active_memory import _rank + + rank = _rank() + plan = rank._plan_flat_forward( + [ForwardInput(input_tokens=torch.arange(64), target_tokens=torch.arange(64))] + ) + assert _impl._packed_priced(plan.signature, rank._one_layer_recompute()) + _duplicate_request_peak(monkeypatch, 1, 64, "fixed") + calibration = _duplicate_request_peak(monkeypatch, 16, 4096, "fixed") + observed = replace( + plan, packed_tokens=4096, logical_tokens=16 * 4096, output_bytes=16 * 4096 * 4 + ) + rank._update_memory_profile(observed, calibration, retained_bytes=None) + assert rank._memory_profiles[plan.signature].logical_per_packed == 16 + peak = _duplicate_request_peak(monkeypatch, 20, 64, "fixed") + estimate = rank._estimate_required_memory_bytes_from_values( + packed_tokens=64, + logical_tokens=20 * 64, + output_bytes=20 * 64 * 4, + signature=plan.signature, + ) + assert peak <= estimate diff --git a/tests/unit/test_trainer_rank_moe_memory.py b/tests/unit/test_trainer_rank_moe_memory.py index 6b160e1c1..33b8ce9df 100644 --- a/tests/unit/test_trainer_rank_moe_memory.py +++ b/tests/unit/test_trainer_rank_moe_memory.py @@ -10,6 +10,7 @@ from art.trainer_rank import ForwardInput, TrainerRank from art.trainer_rank._impl import ( + _PACKED_PRICED_LOGICAL_ROW_BYTES, _MemoryProfile, _MemorySignature, _moe_output_bytes_per_token, @@ -76,7 +77,11 @@ def _rank(layer=None): model=[model], optimizer=None, provider=SimpleNamespace( - hidden_size=2048, num_layers=40, recompute_granularity="full" + hidden_size=2048, + num_layers=40, + recompute_granularity="full", + recompute_method="uniform", + recompute_num_layers=1, ), model_support_handler=SimpleNamespace(build_gdn_execution_spec=False), ), @@ -288,7 +293,7 @@ def test_profiles_outputs_and_empty_plan_preserve_empirical_floor(): ) assert estimate( packed_tokens=100, logical_tokens=200, output_bytes=123, signature=signature - ) == int((100 * 200000 + 123) * 1.1) + ) == int((100 * 100000 + _PACKED_PRICED_LOGICAL_ROW_BYTES * 200 + 123) * 1.1) def test_summed_group_envelope_and_retained_profile_unchanged(): diff --git a/tests/unit/test_trainer_rank_planner_reports.py b/tests/unit/test_trainer_rank_planner_reports.py index a370f56fa..842755e5a 100644 --- a/tests/unit/test_trainer_rank_planner_reports.py +++ b/tests/unit/test_trainer_rank_planner_reports.py @@ -250,6 +250,7 @@ def test_replay_reruns_real_memory_estimator_and_prefix_layout(tmp_path): "hidden_size": 8, "param_dtype_size": 2, "recompute_granularity": "full", + "one_layer_recompute": True, "sequence_parallel": False, "attention_output_gate": False, "mlp_activation_factor": 3, @@ -345,6 +346,10 @@ def test_replay_reruns_real_memory_estimator_and_prefix_layout(tmp_path): changed = json.loads(path.read_bytes()) changed[field] += 1 assert reports.replay(changed)["aggregate"]["matches"] is False + unrecorded = reports.validate_report(path.read_bytes()) + unrecorded["replay"]["memory_replay"]["rank"]["one_layer_recompute"] = None + with pytest.raises(ValueError, match="recompute mode is not recorded"): + reports.replay(unrecorded) drifted = reports.validate_report(path.read_bytes()) drifted["replay"]["source_files"]["_impl.py"]["sha256"] = "0" * 64 with pytest.raises(ValueError, match="source differs"): diff --git a/tests/unit/test_trainer_rank_split.py b/tests/unit/test_trainer_rank_split.py index 359700def..c9d5b3c9f 100644 --- a/tests/unit/test_trainer_rank_split.py +++ b/tests/unit/test_trainer_rank_split.py @@ -49,8 +49,10 @@ TrainerRankMemoryError, TrainerRankPartialExecutionError, TrainerRankSlotStateError, + _impl, ) from art.trainer_rank._impl import ( + _PACKED_PRICED_LOGICAL_ROW_BYTES, Unset, _FlatForwardPlan, _MemoryCheck, @@ -84,7 +86,11 @@ def _runtime() -> "TrainingRuntime": model=[_FakeGPT()], optimizer=None, provider=SimpleNamespace( - hidden_size=8, num_layers=4, recompute_granularity="full" + hidden_size=8, + num_layers=4, + recompute_granularity="full", + recompute_method="uniform", + recompute_num_layers=1, ), model_support_handler=SimpleNamespace(build_gdn_execution_spec=False), ) # type: ignore @@ -531,8 +537,13 @@ def test_changed_output_allocation_does_not_inflate_retained_compute() -> None: [(800, 800, True), (801, 801, False), (100, 800, True), (100, 801, False)], ) def test_retained_compute_keeps_growth_and_sharing_trust_limits( - packed_tokens: int, logical_tokens: int, trusted: bool + monkeypatch: pytest.MonkeyPatch, + packed_tokens: int, + logical_tokens: int, + trusted: bool, ) -> None: + # Price the short test request as a packed-priced one. + monkeypatch.setattr(_impl, "_PACKED_PRICED_MIN_REQUEST_TOKENS", 1) rank = TrainerRank(_runtime()) plan = replace( rank._plan_flat_forward([_request(0)]), @@ -549,10 +560,14 @@ def test_retained_compute_keeps_growth_and_sharing_trust_limits( rank._update_memory_profile(plan, 100_000, retained_bytes=None) unknown = rank._plan_cost(candidate) assert unknown.retained == unknown.required - rank._update_memory_profile(plan, 100_000, retained_bytes=60_000) + rate = 260 + rank._update_memory_profile( + plan, 40_000 + 200 * rate, retained_bytes=40_000 + 100 * rate + ) observed = rank._plan_cost(candidate) if trusted: - assert observed.retained == 220_000 + rows = _PACKED_PRICED_LOGICAL_ROW_BYTES * logical_tokens + assert observed.retained == int((40_000 + rate * packed_tokens + rows) * 1.1) assert observed.retained < observed.required else: assert observed.retained == observed.required @@ -652,15 +667,17 @@ def run(child, **kwargs): 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. + # Otherwise reject this rung; a later rung with smaller chunks may still + # fit. Just under the untrusted 1000-row exact rung, the grad child's + # full-sharing layout (priced by packed rows) rescues the rung itself. split, rejected = rank._admit_split_rung( [tuple(range(16)), tuple(range(16, 32))], requests, [request.input_tokens for request in requests], checkpoint=Unset, ) - assert split is None and not rejected.fits + rescued = not admit and profile_packed == 1000 + assert (split is not None) == rejected.fits == rescued @pytest.mark.parametrize( @@ -1002,6 +1019,8 @@ def loss(indices: tuple[int, ...]) -> torch.Tensor: def test_retained_ratio_bound_uses_original_guard_at_trusted_endpoint( monkeypatch: pytest.MonkeyPatch, direction: float | None ) -> None: + # Price the short test requests as packed-priced ones. + monkeypatch.setattr(_impl, "_PACKED_PRICED_MIN_REQUEST_TOKENS", 1) rank = _retained_ratio_rank(monkeypatch) def request(tokens: list[int]) -> ForwardInput: @@ -1076,6 +1095,8 @@ def select_leaf_sharing(input_ids, *, memory_minimal=False, grad_enabled=True): def test_warm_rounding_preserves_native_split_bound_and_exact_budget( monkeypatch: pytest.MonkeyPatch, retained: int | None, admit: bool ) -> None: + # Price the short test requests as packed-priced ones. + monkeypatch.setattr(_impl, "_PACKED_PRICED_MIN_REQUEST_TOKENS", 1) rank = _retained_ratio_rank(monkeypatch) tokens = torch.arange(3) part = [ForwardInput(input_tokens=tokens, target_tokens=tokens) for _ in range(5)] @@ -1090,7 +1111,13 @@ def test_warm_rounding_preserves_native_split_bound_and_exact_budget( rank._update_memory_profile( plan, peak_delta_bytes=7_864_390, retained_bytes=retained ) - costs = [rank._plan_cost(plan) for plan in children] + # Full sharing is the cheapest layout, so it sets the exact budget. + costs = [ + rank._plan_cost( + rank._plan_flat_forward(requests[a : a + 5], memory_minimal=True) + ) + for a in (0, 5) + ] lower = [ rank._split_chunk_lower_cost( requests[a : a + 5], rows[a : a + 5], checkpoint=Unset @@ -1108,16 +1135,18 @@ def test_warm_rounding_preserves_native_split_bound_and_exact_budget( @pytest.mark.parametrize("retained", (None, 0, 7864321 / 15)) -def test_normalized_warm_profile_is_monotone_through_ratio_floor( +def test_normalized_warm_profile_is_monotone_in_packed_tokens( monkeypatch: pytest.MonkeyPatch, retained: float | None ) -> None: + # Price the short test request as a packed-priced one. + monkeypatch.setattr(_impl, "_PACKED_PRICED_MIN_REQUEST_TOKENS", 1) rank = _retained_ratio_rank(monkeypatch) signature = rank._plan_flat_forward([_request(0)]).signature rank._memory_profiles[signature] = _MemoryProfile( 7864330 / 15, 1000, 15 / 13, retained_compute_bytes_per_token=retained ) - # All counts satisfy the retained guard. The logical term dominates until - # N=104, then packed-token growth dominates. Check every integer count. + # All counts satisfy the retained guard; the logical charge does not depend + # on layout, so cost must grow monotonically with packed rows. Check each. costs = [ rank._subforward_cost( packed_tokens=packed,