diff --git a/src/art/trainer_rank/_impl.py b/src/art/trainer_rank/_impl.py index f6884d489..1277878cc 100644 --- a/src/art/trainer_rank/_impl.py +++ b/src/art/trainer_rank/_impl.py @@ -149,6 +149,9 @@ class TopK: # 32 MiB transients live at its peak (Qwen3.6-35B-A3B CP2: the RoPE frequencies # and a frozen linear's output, at 2k to 20k tokens); warm waves do not. _COLD_RECOMPUTE_TRANSIENT_BYTES = 64 * 2**20 +# The TP x SP traces (dense Qwen3.8-27B, TP2 and TP4) held a third: TE's 32 MiB +# cuBLAS workspace, allocated by the process's first wave. +_SEQUENCE_PARALLEL_COLD_TRANSIENT_BYTES = 3 * 32 * 2**20 _PLANNER_REFINEMENT_BUDGET = 2_000 _LAYOUT_SELECTION_CACHE_LIMIT = 64 @@ -1299,6 +1302,45 @@ def _gdn_layer_count(model: torch.nn.Module) -> int: return sum(isinstance(module, GatedDeltaNet) for module in model.modules()) +def _mixer_top_gaps(model: torch.nn.Module) -> tuple[int | None, int | None] | None: + """Layers between the top decoder layer and the highest attention / GDN layer. + + ``None`` for a mixer kind the decoder lacks; ``None`` overall when the + decoder layers are unreadable. + """ + try: + from megatron.core.ssm.gated_delta_net import GatedDeltaNet + except ImportError: + return None + try: + layers = list(_language_model(model).decoder.layers) + except (AttributeError, RuntimeError, TypeError): + return None + gdn = [ + any(isinstance(module, GatedDeltaNet) for module in layer.modules()) + for layer in layers + ] + top = len(layers) - 1 + return ( + next((top - index for index in range(top, -1, -1) if not gdn[index]), None), + next((top - index for index in range(top, -1, -1) if gdn[index]), None), + ) + + +def _lora_modules_per_layer(model: torch.nn.Module) -> int: + """The most LoRA modules any one decoder layer holds (0 when unreadable).""" + try: + from art.megatron.lora import LoRA + + layers = list(_language_model(model).decoder.layers) + except (AttributeError, ImportError, RuntimeError, TypeError): + return 0 + return max( + (sum(type(module) is LoRA for module in layer.modules()) for layer in layers), + default=0, + ) + + def _moe_layer_count(model: torch.nn.Module) -> int: """Number of mixture-of-experts layers in the model (0 when unavailable).""" @@ -2068,10 +2110,13 @@ def memory_field(name: str, default: Any = None) -> Any: # Layers that run the gated-delta-net path (Qwen3.5-4B: 24 of 32); the # cost model prices GDN state hand-offs per GDN layer, not per layer. self._gdn_layers = _gdn_layer_count(runtime.model[0]) + self._mixer_top_gaps = _mixer_top_gaps(runtime.model[0]) if self._gdn_layers == 0 and bool( getattr(runtime.model_support_handler, "build_gdn_execution_spec", False) ): self._gdn_layers = self._num_layers + self._mixer_top_gaps = (None, 0) + self._lora_modules_per_layer = _lora_modules_per_layer(runtime.model[0]) # A fitted layout cost table applies only to the execution classes it # was calibrated on (device class, dtype, model geometry, parallel # shape); other runtimes keep the previous score. @@ -3253,9 +3298,13 @@ def _head_projection_rows( *, positions: Sequence[torch.Tensor] | None = None, lower_bound: bool = False, + uncapped: bool = False, ) -> int: """Per-group logical bounds or exact packed union; no device-label read. + Capped at one head chunk unless ``uncapped`` (every projected row, for + pricing each chunk's statistics path). + A single sequence's valid rows cannot alias each other. Across requests they may share: max is a lower bound, sum an upper bound. Ignore labels only when every label on that input row is -100, as projection does. @@ -3292,14 +3341,14 @@ def _head_projection_rows( row = row.index_select(0, offsets) for position in row.tolist(): projected.add(int(position)) - if len(projected) >= _HEAD_CHUNK_TOKENS: + if len(projected) >= _HEAD_CHUNK_TOKENS and not uncapped: return _HEAD_CHUNK_TOKENS - return min( - _HEAD_CHUNK_TOKENS, + rows = ( (max(counts, default=0) if lower_bound else sum(counts)) if positions is None - else len(projected), + else len(projected) ) + return rows if uncapped else min(_HEAD_CHUNK_TOKENS, rows) def _head_target_chunk_rows( self, @@ -3361,28 +3410,23 @@ def _head_target_chunk_rows( chunk_start = first_index // _HEAD_CHUNK_TOKENS * _HEAD_CHUNK_TOKENS return min(_HEAD_CHUNK_TOKENS, len(projected) - chunk_start) - def _sequence_parallel_floor_covered(self, layers: int, tp: int, cp: int) -> bool: - """Whether the checkpoint floor covers a dense TP x SP recompute peak. - - Traced once: dense Qwen3.8-27B (64 layers) at TP4 with sequence - parallelism and CP1. Over the gathered rows, the recomputed layer's peak - held its SP-gathered norm input (2H per row), the MLP FC1 stage (6F/TP), - the recomputed mixer (within its projection widths / TP), norm outputs - and other workspace (each under H), plus one input gradient per - sharded row. The floor repeats the sharded boundaries as the - input-gradient term, so that repeat must cover this workspace; GDN - segment states grow with segments instead and are priced separately. - Other TP sizes, CP, MoE, replicated QKV (KV groups below TP), missing - geometry and models too shallow or wide for the bound keep today's + def _sequence_parallel_floor_covered(self, tp: int, cp: int) -> bool: + """Whether the explicit TP x SP checkpoint floor applies to this model. + + Traced on dense Qwen3.8-27B (64 layers, GDN and gated attention) at + TP4 and at TP2 with sequence parallelism and CP1. The floor covers + each rank's boundary shards and the recomputed layer's measured + workspace with explicit terms (``_sequence_parallel_workspace_bytes``), + so it needs readable mixer geometry but, unlike the earlier repeated- + boundary cover, no depth/width bound: shallower and wider dense models + are now eligible and priced by the same terms (untraced). Other TP + sizes, CP, MoE and replicated QKV (KV groups below TP) keep today's pricing. """ geometry = self._geometry - if tp != 4 or cp != 1 or self._moe_layers or geometry.moe_experts: + if tp not in (2, 4) or cp != 1 or self._moe_layers or geometry.moe_experts: return False - hidden = self._hidden_size - ffn = geometry.ffn_hidden_size or 4 * hidden - attention_layers = self._num_layers > self._gdn_layers - if attention_layers and ( + if self._num_layers > self._gdn_layers and ( geometry.num_attention_heads <= 0 or geometry.kv_channels <= 0 # Replicated QKV keeps a global QKV output on every rank. @@ -3396,28 +3440,7 @@ def _sequence_parallel_floor_covered(self, layers: int, tp: int, cp: int) -> boo geometry.gdn_value_head_dim, geometry.gdn_conv_kernel, # Prices each segment's conv history. ) - if self._gdn_layers and min(gdn_widths) <= 0: - return False - attention = ( - (7 if self._attention_output_gate else 5) - * geometry.num_attention_heads - * geometry.kv_channels - + 3 * geometry.num_query_groups * geometry.kv_channels - if attention_layers - else 0 - ) - gdn = ( - 4 * geometry.gdn_key_heads * geometry.gdn_key_head_dim - + 8 * geometry.gdn_value_heads * geometry.gdn_value_head_dim - if self._gdn_layers - else 0 - ) - # Per gathered row, times TP: the repeat is layers x H; the workspace is - # 2H + the FC1 stage (6F/TP, or the SwiGLU live set if wider) + - # mixer/TP + H of norms + H of other workspace, and the gradient H/TP. - stage = max(6, self._mlp_activation_factor) * ffn - workspace = 2 * hidden * tp + stage + max(attention, gdn) + 2 * hidden * tp - return layers * hidden >= workspace + hidden + return not (self._gdn_layers and min(gdn_widths) <= 0) def _subforward_cost( self, @@ -3434,8 +3457,10 @@ def _subforward_cost( retained_tokens: int | None = None, hybridep_growth_bytes: int = 0, ) -> _SubforwardCost: - checkpoint_memory = self._checkpoint_memory_floor( - group_rows, slot_refs, gdn_segments + checkpoint_memory = self._sequence_parallel_lora_floor( + self._checkpoint_memory_floor(group_rows, slot_refs, gdn_segments), + group_rows, + signature, ) required = self._estimate_required_memory_bytes_from_values( packed_tokens=packed_tokens, @@ -3463,11 +3488,14 @@ def _subforward_cost( checkpoint_floor[0], ), ) - # One logical BF16 input gradient per eligible full/uniform/1 boundary. + # One logical BF16 input gradient per eligible full/uniform/1 boundary + # (TP x SP: one input-gradient shard; the workspace is explicit). # This partial peak allowance is not evidence of simultaneous distinct # backing stores, nor a bound for compiler saves or other backward work. # Keep it out of forward retention, including the cold fallback above. - gradient = checkpoint_retained + gradient = self._checkpoint_input_gradient_bytes( + group_rows, checkpoint_retained + ) gradient_slots = self._gradient_slots(group_rows, slot_refs) adapter_gradient = ( self._checkpoint_adapter_gradient_bytes( @@ -3483,7 +3511,7 @@ def _subforward_cost( checkpoint_workspace, head_workspace_bytes, checkpoint_floor[1] ) if gradient and self._memory_profiles.get(signature) is None: - checkpoint_workspace += _COLD_RECOMPUTE_TRANSIENT_BYTES + checkpoint_workspace += self._cold_recompute_transient_bytes() forward_required = required if gradient: required = max( @@ -5633,14 +5661,12 @@ def _project_vocab_parallel( chunk_rows = rows[start : start + _HEAD_CHUNK_TOKENS] # Recompute vocabulary-sized intermediates one chunk at a time in # backward; chunking alone otherwise retains every chunk's logits. - local_logits, log_z, local_topk = checkpoint( - self._local_head_stats, + local_logits, log_z, local_topk = self._checkpointed_head_stats( model, _select_positions(hidden_by_row, chunk_rows), output_weight=output_weight, need_log_z=need_log_z, max_top_k=max_top_k, - use_reentrant=False, ) logit_start, logit_end = logit_bounds[chunk_index : chunk_index + 2] logit_chunk_offsets = logit_rows[logit_start:logit_end] - start @@ -5712,6 +5738,38 @@ def _project_vocab_parallel( # Do not retain prior chunk buffers while the next stats RHS runs. del local_logits, chunk_logits + def _checkpointed_head_stats( + self, + model: "GPTModel", + hidden: torch.Tensor, + *, + output_weight: torch.Tensor | None, + need_log_z: bool, + max_top_k: int, + ) -> tuple[ + torch.Tensor, + torch.Tensor | None, + tuple[torch.Tensor, torch.Tensor] | None, + ]: + """One chunk's head statistics under non-reentrant checkpointing. + + The recompute must save what the forward saved, so it replays the + statistics path the forward chose (``path``) instead of attempting + the kernels again. + """ + from torch.utils.checkpoint import checkpoint + + return checkpoint( + self._local_head_stats, + model, + hidden, + output_weight=output_weight, + need_log_z=need_log_z, + max_top_k=max_top_k, + path=[], + use_reentrant=False, + ) + def _local_head_stats( self, model: "GPTModel", @@ -5720,6 +5778,7 @@ def _local_head_stats( output_weight: torch.Tensor | None, need_log_z: bool, max_top_k: int, + path: list[str] | None = None, ) -> tuple[ torch.Tensor, torch.Tensor | None, @@ -5733,15 +5792,58 @@ def _local_head_stats( log_z: torch.Tensor | None = None local_topk: tuple[torch.Tensor, torch.Tensor] | None = None if need_log_z: - topk_stats = _try_triton_local_topk_stats(local_logits, k=max_top_k) + # ``path`` holds the forward's statistics path; a checkpoint + # recompute replays it rather than attempting the kernels again. + recorded = path[0] if path else None + topk_stats = ( + _try_triton_local_topk_stats(local_logits, k=max_top_k) + if recorded in (None, "topk") + else None + ) logsumexp_stats = ( cast( tuple[torch.Tensor, torch.Tensor] | None, _try_triton_stats("local_logsumexp_stats", local_logits), ) - if topk_stats is None + if topk_stats is None and recorded in (None, "logsumexp") else None ) + shape = (int(local_logits.shape[0]), int(local_logits.shape[1])) + if recorded is None: + if topk_stats is not None or logsumexp_stats is not None: + # A kernel ran at this chunk shape: admission may price it. + getattr(self, "_triton_head_stats_failures", set()).discard(shape) + elif _triton_stats_enabled(local_logits.is_cuda, shape[0]): + # An attempted kernel failed at this chunk shape. Admission + # priced the kernel's buffers, so compute the same + # statistics eagerly within them; later admissions price + # the shape for the eager fallback (_triton_head_stats) + # until a kernel succeeds at it again. + if not hasattr(self, "_triton_head_stats_failures"): + self._triton_head_stats_failures = set() + self._triton_head_stats_failures.add(shape) + logsumexp_stats = _eager_local_logsumexp_stats(local_logits) + recorded = "bounded" + if path is not None: + path.append( + recorded + or ( + "topk" + if topk_stats is not None + else "logsumexp" + if logsumexp_stats is not None + else "eager" + ) + ) + elif recorded == "bounded": + logsumexp_stats = _eager_local_logsumexp_stats(local_logits) + elif recorded in ("topk", "logsumexp") and ( + topk_stats is None and logsumexp_stats is None + ): + raise RuntimeError( + "the head statistics kernel failed in a checkpoint recompute " + "after it succeeded in the forward" + ) stats = topk_stats if topk_stats is not None else logsumexp_stats if stats is not None: local_max, local_sum = stats[:2] @@ -5752,6 +5854,8 @@ def _local_head_stats( ) log_z = global_max + torch.log(global_sum) else: + # No kernel attempted (non-CUDA, disabled, short chunk): + # admission prices this FP32 fallback's seven buffers. log_z = _vocab_parallel_log_z(local_logits) if topk_stats is not None: @@ -6104,6 +6208,10 @@ def _gather_tensor_parallel_logits(self, logits: torch.Tensor) -> torch.Tensor: _memory._estimate_required_memory_bytes_from_values ) _gdn_segment_layer_bytes = _memory._gdn_segment_layer_bytes + _sequence_parallel_workspace_bytes = _memory._sequence_parallel_workspace_bytes + _sequence_parallel_lora_floor = _memory._sequence_parallel_lora_floor + _cold_recompute_transient_bytes = _memory._cold_recompute_transient_bytes + _checkpoint_input_gradient_bytes = _memory._checkpoint_input_gradient_bytes _available_memory_bytes = _memory._available_memory_bytes _all_ranks_have_memory_profile = _memory._all_ranks_have_memory_profile _update_memory_profile = _memory._update_memory_profile @@ -6836,19 +6944,56 @@ def _try_triton_local_topk_stats( ) +def _triton_stats_enabled(cuda: bool, rows: int) -> bool: + """Whether ``_try_triton_stats`` attempts the kernel for a ``rows`` chunk.""" + return ( + cuda + and os.environ.get("ART_TRAINER_RANK_TRITON_TOPK", "1").lower() + not in {"0", "false"} + and rows >= int(os.environ.get("ART_TRAINER_RANK_TRITON_MIN_ROWS", "64")) + ) + + +@lru_cache(maxsize=None) +def _triton_stats_importable() -> bool: + try: + from art.trainer_rank import topk # noqa: F401 + except Exception: + return False + return True + + +def _triton_head_stats_available(rank: Any) -> bool: + """Whether the head statistics kernels can run on this rank at all.""" + return ( + _triton_stats_enabled(rank.device.type == "cuda", 1 << 62) + and _triton_stats_importable() + ) + + +def _triton_head_stats(rank: Any, rows: int, vocabulary: int) -> bool: + """The head statistics path a ``rows`` x ``vocabulary`` chunk takes, for admission. + + Mirrors ``_try_triton_stats``'s attempt predicate plus kernel + importability. A kernel that imports but fails at run time falls back + (unless ART_TRAINER_RANK_TRITON_TOPK=strict makes that fatal); the first + such wave was priced for the kernel, and ``_local_head_stats`` then + records that chunk shape, priced for the fallback until a kernel succeeds + at it again. + """ + return ( + _triton_stats_enabled(rank.device.type == "cuda", rows) + and (rows, vocabulary) not in getattr(rank, "_triton_head_stats_failures", ()) + and _triton_stats_importable() + ) + + def _try_triton_stats( name: str, local_logits: torch.Tensor, **kwargs: object, ) -> object | None: - if not local_logits.is_cuda: - return None - if os.environ.get("ART_TRAINER_RANK_TRITON_TOPK", "1").lower() in { - "0", - "false", - } or int(local_logits.shape[0]) < int( - os.environ.get("ART_TRAINER_RANK_TRITON_MIN_ROWS", "64") - ): + if not _triton_stats_enabled(local_logits.is_cuda, int(local_logits.shape[0])): return None try: from art.trainer_rank import topk @@ -6901,6 +7046,80 @@ def _vocab_parallel_topk_from_local( ) +# Row sub-chunks per head chunk for the bounded eager statistics: one FP32 +# sub-chunk is 2 * ceil(rows / 64) / rows of a BF16 chunk buffer: 1/32 when the +# rows divide by 64, more otherwise (the estimator prices the ceiling). +_EAGER_STATS_SUBCHUNKS = 64 + + +class _EagerLocalStats(torch.autograd.Function): + """The statistics kernel's (local max, local sum) contract, computed eagerly. + + The logits are processed in up to ``_EAGER_STATS_SUBCHUNKS`` row sub-chunks + of ceil(rows / 64) rows each through one owned FP32 work buffer, beside + the logits in forward and beside the logits and their gradient in + backward: the kernel path's buffers (topk._LocalStatsFunction) plus that + buffer, instead of the unchunked fallback's seven. The logits are never + written. Rows are independent; the caller's rescale of local sums to the + global maximum equals the unchunked fallback in exact arithmetic, and + FP32 rounding may differ when ranks' local maxima differ. + """ + + @staticmethod + def forward(ctx: Any, logits: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + rows = int(logits.shape[0]) + step = max(1, -(-rows // _EAGER_STATS_SUBCHUNKS)) + local_max = logits.max(dim=-1).values.float() + local_sum = torch.empty_like(local_max) + work = torch.empty( + (min(step, rows), int(logits.shape[1])), + dtype=torch.float32, + device=logits.device, + ) + for start in range(0, rows, step): + stop = min(start + step, rows) + block = work[: stop - start] + block.copy_(logits[start:stop]) # owned: FP32 logits stay intact + block.sub_(local_max[start:stop, None]).exp_() + local_sum[start:stop] = block.sum(dim=-1) + del work + ctx.save_for_backward(logits, local_max) + ctx.step = step + return local_max, local_sum + + @staticmethod + def backward(ctx: Any, *grad_outputs: Any) -> Any: + # As the kernel: the local max is a detached shift; only the sum + # carries a gradient, exp(logits - local max) per row. + _grad_local_max, grad_local_sum = grad_outputs + logits, local_max = ctx.saved_tensors + grad = torch.empty_like(logits) + if grad_local_sum is None: + return grad.zero_() + rows, step = int(logits.shape[0]), int(ctx.step) + work = torch.empty( + (min(step, rows), int(logits.shape[1])), + dtype=torch.float32, + device=logits.device, + ) + for start in range(0, rows, step): + stop = min(start + step, rows) + block = work[: stop - start] + block.copy_(logits[start:stop]) # owned: the saved logits stay intact + block.sub_(local_max[start:stop, None]).exp_() + block.mul_(grad_local_sum[start:stop, None]) + grad[start:stop] = block + del work + return grad + + +def _eager_local_logsumexp_stats( + local_logits: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + stats = _EagerLocalStats.apply(local_logits) + return cast(tuple[torch.Tensor, torch.Tensor], stats) + + def _vocab_parallel_log_z(local_logits: torch.Tensor) -> torch.Tensor: local_logits = local_logits.float() local_max = local_logits.max(dim=-1).values.detach() diff --git a/src/art/trainer_rank/_memory.py b/src/art/trainer_rank/_memory.py index 42a94e8d8..28952c077 100644 --- a/src/art/trainer_rank/_memory.py +++ b/src/art/trainer_rank/_memory.py @@ -234,13 +234,22 @@ def _split_plan_memory_check( def _head_vocabulary(self: TrainerRank) -> int: - """Runtime eligibility and vocabulary, without reading parameter values.""" + """Runtime eligibility and this rank's vocabulary, without reading values. + + TP > 1 runs the same 512-row chunked head over each rank's vocabulary + shard; it is priced only where the explicit TP x SP checkpoint floor + applies, whose shallow shapes it can dominate. + """ + _, tp, _, pp = self._topology_key() if ( self._padded_vocab_size is None or len(self.runtime.model) != 1 - or self._topology_key()[1::2] != (1, 1) + or pp != 1 + or self._padded_vocab_size % tp + or (tp > 1 and not _checkpoint_layers(self, ((1, True),))) ): return 0 + vocabulary = self._padded_vocab_size // tp try: model = _impl._language_model(self.runtime.model[0]) except (AttributeError, RuntimeError): @@ -271,8 +280,8 @@ def _head_vocabulary(self: TrainerRank) -> int: if ( type(weight) not in (_impl.torch.Tensor, _impl.torch.nn.Parameter) or weight.dtype is not _impl.torch.bfloat16 - or tuple(weight.shape) != (self._padded_vocab_size, self._hidden_size) - or head.output_size_per_partition != self._padded_vocab_size + or tuple(weight.shape) != (vocabulary, self._hidden_size) + or head.output_size_per_partition != vocabulary or head.output_size != self._padded_vocab_size or head.input_size != self._hidden_size or getattr(config, "params_dtype", None) is not _impl.torch.bfloat16 @@ -288,13 +297,14 @@ def _head_vocabulary(self: TrainerRank) -> int: for name in ( "_project_head", "_project_vocab_parallel", + "_checkpointed_head_stats", "_local_head_stats", "_local_logits_from_hidden_rows", ) ) ): return 0 - return int(self._padded_vocab_size) + return int(vocabulary) def _dense_head_bytes(vocabulary: int, rows: int) -> int: @@ -336,8 +346,19 @@ def _group_head_workspace_bytes( Its BF16 logits, FP32 conversion, subtraction and exp overlap. This is not a bound for row vectors, inter-chunk liveness or library workspaces. Rejection lower bounds retain only the unconditional dense components. + TP > 1 heads (priced with the explicit TP x SP floor) instead follow each + chunk's statistics path: ``_tp_head_workspace_bytes``. """ dense = self._head_workspace_bytes(rows) + if dense and self._topology_key()[1] > 1: + return _tp_head_workspace_bytes( + self, + rows, + requests, + grad_enabled=grad_enabled, + positions=positions, + lower_bound=lower_bound, + ) needs_statistics = any( request.target_tokens is not None or request.top_k is not None for request in requests @@ -372,6 +393,151 @@ def _group_head_workspace_bytes( return dense +def _tp_head_workspace_bytes( + self: TrainerRank, + rows: int, + requests: Sequence[AnyForwardInput], + *, + grad_enabled: bool, + positions: Sequence[torch.Tensor] | None = None, + lower_bound: bool = False, +) -> int: + """The TP > 1 head component, per chunk statistics path (#1068's model at TP1). + + The eager FP32 statistics keep #1068's seven BF16 buffers. Chunks whose + kernel is attempted run it, or after a failure the bounded eager + statistics (``_impl._EagerLocalStats``): the kernel path's buffers plus + ``_eager_stats_extra_bytes``. A gradient wave's kernel path holds the + saved logits and the statistics gradient, plus a dense gradient per + gathered output kind: the target gather's and top-k's scatter. Requested + logits hold the local logits, their indexed copy, the gather buffer and + its concatenation (Megatron's gather along the vocabulary): 2 + 2 TP. + """ + dense = self._head_workspace_bytes(rows) + if not _head_target_backward(self): + return dense + tp = self._topology_key()[1] + needs_statistics = any( + request.target_tokens is not None or request.top_k is not None + for request in requests + ) + if not needs_statistics: + logits = any(request.logits for request in requests) + return dense if lower_bound or not logits else (2 + 2 * tp) * dense + kernel = dense + if grad_enabled: + # Saved logits and the statistics gradient, then each gathered + # output's dense gradient: top-k's scatter and the target gather's. + kernel = 2 * dense + if any(request.top_k is not None for request in requests): + kernel += dense + if any(request.target_tokens is not None for request in requests): + kernel += ( + self._head_workspace_bytes( + self._head_target_chunk_rows( + requests, positions=positions, lower_bound=lower_bound + ) + ) + if any( + request.logits or request.top_k is not None for request in requests + ) + else dense + ) + low = self._head_projection_rows( + requests, positions=positions, lower_bound=True, uncapped=True + ) + if lower_bound: + # A rejection bound monotone in rows: only unconditional components, + # and the eager statistics only when no kernel can run at all. + return max( + kernel, + 0 + if _impl._triton_head_stats_available(self) + else 7 * self._head_workspace_bytes(low), + ) + if positions is not None: + # The selected layout's packed union: every chunk exactly. + fallback = _head_fallback_bytes(self, low, low) + else: + # An acceptance bound: any projected count the union can take. + high = self._head_projection_rows(requests, positions=positions, uncapped=True) + fallback = _head_fallback_bytes(self, low, high) + logits = (2 + 2 * tp) * dense if any(request.logits for request in requests) else 0 + return max( + kernel + _eager_stats_extra_bytes(_head_vocabulary(self), rows), + fallback, + logits, + ) + + +def _eager_stats_extra_bytes(vocabulary: int, rows: int) -> int: + """The bounded eager statistics' live increment over the kernel path for a + ``rows`` x ``vocabulary`` head chunk (``_impl._EagerLocalStats``): one FP32 + row sub-chunk, plus per-row FP32 maxima, sums and their gradient.""" + rows = min(rows, _impl._HEAD_CHUNK_TOKENS) + step = -(-rows // _impl._EAGER_STATS_SUBCHUNKS) + return 4 * vocabulary * step + 16 * rows + + +def _frozen_head_bytes( + vocabulary: int, + rows: int, + target_rows: int, + *, + target_backward: bool, + statistics: bool, + grad: bool, + tp: int, +) -> int: + """Replay's head charge from frozen facts, as live admission prices it. + + TP1 is #1068's capacity charge. TP > 1 is ``_tp_head_workspace_bytes``'s + kernel path plus the bounded statistics' increment, or the gathered logits + copies. Capture declines TP > 1 waves with an eager chunk, with logits + beside statistics, or with top-k beside targets in a gradient wave: the + facts record none of these. A gradient wave with statistics but no target + rows is therefore top-k only. + """ + dense = _dense_head_bytes(vocabulary, rows) + if not target_backward: + return dense + if tp == 1: + return dense * (7 if statistics else 3) + if not statistics: + return (2 + 2 * tp) * dense + kernel = dense + if grad: + kernel = 2 * dense + ( + _dense_head_bytes(vocabulary, target_rows) if target_rows else dense + ) + return kernel + _eager_stats_extra_bytes(vocabulary, rows) + + +def _head_fallback_bytes(self: TrainerRank, low: int, high: int) -> int: + """Seven BF16 buffers of the largest head chunk that runs the eager fallback + for any projected-row count from ``low`` to ``high``.""" + chunk = _impl._HEAD_CHUNK_TOKENS + vocabulary = _head_vocabulary(self) + sizes: set[int] = set() + if high - low >= chunk: + sizes.update(range(1, min(high, chunk) + 1)) # every tail occurs + else: + for rows in range(low, high + 1): + sizes.update(_head_chunk_sizes(rows)) + eager = [ + size for size in sizes if not _impl._triton_head_stats(self, size, vocabulary) + ] + return 7 * self._head_workspace_bytes(max(eager)) if eager else 0 + + +def _head_chunk_sizes(rows: int) -> tuple[int, ...]: + """The distinct chunk sizes the head runs over ``rows`` projected rows.""" + chunk = _impl._HEAD_CHUNK_TOKENS + return tuple( + size for size in (min(rows, chunk), rows % chunk * (rows > chunk)) if size + ) + + def _plan_head_workspace_bytes(self: TrainerRank, plan: _impl._FlatForwardPlan) -> int: peak = 0 for group in plan.groups: @@ -536,8 +702,10 @@ def _checkpoint_layers( This is not a bound for custom preprocessing, attention, or all backward. With sequence parallelism a rank saves only its shard of each boundary; that is priced only where ``_sequence_parallel_floor_covered`` holds, - and there, for gradient waves, the recomputed GDN layer's recurrent - states for ``gdn_segments`` (gradient groups' segments) plus padding. + and there, for gradient waves, with the recomputed layer's explicit + workspace (``_sequence_parallel_workspace_bytes``) and its GDN layer's + recurrent states for ``gdn_segments`` (gradient groups' segments) plus + padding, beside one input-gradient shard instead of the TP1 repeat. """ if not group_rows or len(self.runtime.model) != 1: return 0 @@ -577,7 +745,7 @@ def _checkpoint_layers( or self._param_dtype_size != 2 or next(self.runtime.model[0].parameters()).dtype is not _impl.torch.bfloat16 or pp != 1 - or (tp > 1 and not self._sequence_parallel_floor_covered(layers, tp, cp)) + or (tp > 1 and not self._sequence_parallel_floor_covered(tp, cp)) or any( type(getattr(config, name, None)) is not type(value) or getattr(config, name) != value @@ -652,6 +820,17 @@ def _checkpoint_floor_from_facts( + (0 if grad else 4 * rows * self._hidden_size * 2) for (rows, grad), ref in zip(group_rows, refs, strict=True) ) + if tp > 1 and gradient_rows: + # The recomputed layer's workspace, for the largest gradient group: + # groups run their backward one after another. + workspace = max( + workspace, + *( + self._sequence_parallel_workspace_bytes(rows) + for rows, grad in group_rows + if grad + ), + ) if tp > 1 and self._gdn_layers and gradient_rows: # Recurrent states grow with segments, not rows; backward recomputes # one layer at a time. Padding to TP adds up to TP - 1 one-token @@ -661,6 +840,129 @@ def _checkpoint_floor_from_facts( return retained, workspace +def _sequence_parallel_workspace_bytes(self: TrainerRank, rows: int) -> int: + """The recomputed layer's live workspace at a TP x SP backward peak. + + Traced on dense Qwen3.8-27B at TP4 (25,728 rows) and TP2 (20,816 to + 66,284 rows, one to four GDN segments), cold and warm: every peak was a + recomputed layer's MLP FC1 stage. Its workspace covers, over the gathered + rows, the SP-gathered norm output (2H), the FC1 stage (6F/TP: GEMM + output, gate and up LoRA outputs and their sum; the SwiGLU live set if + wider) and the recomputed mixer at its projection widths / TP (GDN held + 0.885-0.887 of that bound); over the sharded rows, norm outputs (3.22H + measured, priced at 3.25H); and the wave length's fp32 rotary cache + (``kv_channels`` per row bounds the traced 64-dim partial rotary). + + Recomputing layer i holds the L-shard ``retained`` term's i + 1 + checkpoint inputs plus its own residual: L + 1 - gap shards, where gap + counts the layers above it. Each mixer kind peaks at its highest layer; + the boundary shards beyond L are charged here, per kind, from the + decoder's layer order (one extra shard for a kind whose order is + unknown). In the traced model the top layer is attention and the top GDN + layer (one below) peaks: exactly L shards. The input-gradient shard, LoRA + intermediates, GDN segment states and cold transients are separate terms. + """ + if rows <= 0: + return 0 + geometry = self._geometry + tp = self._topology_key()[1] + hidden = self._hidden_size + ffn = geometry.ffn_hidden_size or 4 * hidden + attention_layers = self._num_layers > self._gdn_layers + attention = ( + (7 if self._attention_output_gate else 5) + * geometry.num_attention_heads + * geometry.kv_channels + + 3 * geometry.num_query_groups * geometry.kv_channels + ) + gdn = ( + 4 * geometry.gdn_key_heads * geometry.gdn_key_head_dim + + 8 * geometry.gdn_value_heads * geometry.gdn_value_head_dim + ) + shard = -(-rows // tp) * hidden + gaps = getattr(self, "_mixer_top_gaps", None) or (None, None) + peak = ( + max( + rows * -(-width // tp) + (1 - (0 if gap is None else gap)) * shard + for present, width, gap in ( + (attention_layers, attention, gaps[0]), + (bool(self._gdn_layers), gdn, gaps[1]), + ) + if present + ) + if self._num_layers + else 0 + ) + stage = max(6, self._mlp_activation_factor) * ffn + gathered = rows * (2 * hidden + -(-stage // tp)) + sharded = -(-rows // tp) * -(-13 * hidden // 4) + rotary = 2 * rows * geometry.kv_channels if attention_layers else 0 + return (gathered + sharded + peak + rotary) * self._param_dtype_size + + +def _sequence_parallel_lora_floor( + self: TrainerRank, + floor: tuple[int, int], + group_rows: tuple[tuple[int, bool], ...], + signature: _impl._MemorySignature, +) -> tuple[int, int]: + """Add the recomputed layer's LoRA ``x @ A`` intermediates to a TP x SP floor. + + Each adapted module of the recomputed layer saves rows x rank over the + gathered rows (FC1 gate/up measured 16 B/row at rank 4). Charge the + decoder's most-adapted layer at the gradient slots' largest rank, from + the signature's slot shapes ((ndim, *A.shape, *B.shape) per module). + """ + retained, workspace = floor + if not retained or self._topology_key()[1] == 1: + return floor + rank = max( + ( + shape[2] + for grad, shapes in signature.slot_shapes + if grad + for shape in shapes + if len(shape) == 5 and shape[0] == 2 + ), + default=0, + ) + modules = getattr(self, "_lora_modules_per_layer", None) + if not modules: + # Unread (e.g. reports without the recorded count): every adapted + # module of the slot, which bounds any one layer's. + modules = max( + (len(shapes) for grad, shapes in signature.slot_shapes if grad), default=0 + ) + rows = max((rows for rows, grad in group_rows if grad), default=0) + return retained, workspace + rows * modules * rank * self._param_dtype_size + + +def _cold_recompute_transient_bytes(self: TrainerRank) -> int: + """Fixed transients of an unprofiled gradient wave's first execution.""" + return ( + _impl._SEQUENCE_PARALLEL_COLD_TRANSIENT_BYTES + if self._topology_key()[1] > 1 + else _impl._COLD_RECOMPUTE_TRANSIENT_BYTES + ) + + +def _checkpoint_input_gradient_bytes( + self: TrainerRank, group_rows: tuple[tuple[int, bool], ...], retained: int +) -> int: + """The input-gradient term beside the checkpoint floor's boundaries. + + TP1 charges one logical gradient per boundary, which also stands in for + the recompute workspace. Under TP x SP the workspace is priced + explicitly, and the recompute peak held one input-gradient shard of the + group being recomputed. + """ + tp = self._topology_key()[1] + if not retained or tp == 1: + return retained + rows = max((rows for rows, grad in group_rows if grad), default=0) + return -(-rows // tp) * self._hidden_size * self._param_dtype_size + + def _gradient_slots( group_rows: Sequence[tuple[int, bool]], slot_refs: Sequence[LoRASlotRef | None] | None, @@ -977,10 +1279,26 @@ def _estimate_flat_forward( for request in head_requests ) ) - if lower != upper or ( - mixed_targets - and self._head_target_chunk_rows(head_requests, lower_bound=True) - != self._head_target_chunk_rows(head_requests) + # TP > 1 heads price each chunk's statistics path, so the + # exact count needs the packed union beyond one chunk too. + spread = ( + self._topology_key()[1] > 1 + and upper == _impl._HEAD_CHUNK_TOKENS + and self._head_projection_rows( + head_requests, lower_bound=True, uncapped=True + ) + != self._head_projection_rows(head_requests, uncapped=True) + ) + if ( + lower != upper + or spread + or ( + mixed_targets + and self._head_target_chunk_rows( + head_requests, lower_bound=True + ) + != self._head_target_chunk_rows(head_requests) + ) ): packed = _impl.materialize_prefix_tree_layout( tuple( @@ -1140,12 +1458,20 @@ def _memory_signature_from_requests( def _slot_memory_shapes( self: TrainerRank, ref: "LoRASlotRef | None" ) -> tuple[tuple[int, ...], ...]: - """Separate empirical trust across actual selected adapter layouts.""" + """Separate empirical trust across actual selected adapter layouts. + + Sequence-parallel (TP > 1) models record them too: the explicit TP x SP + checkpoint floor prices LoRA intermediates from these ranks. + """ if ( ref is None or ref.name is None or isinstance(ref, _impl._LocalLoRASlotRef) - or not (getattr(self, "_moe_layers", 0) or getattr(self, "_gdn_layers", 0)) + or not ( + getattr(self, "_moe_layers", 0) + or getattr(self, "_gdn_layers", 0) + or getattr(self, "_sequence_parallel", False) + ) ): # Generic/no-component planning must not import Megatron or walk # model owners just to construct its existing memory signature. @@ -1438,18 +1764,24 @@ def _estimate_required_memory_bytes_from_values( ), ) retained, workspace = ( - self._checkpoint_memory_floor(group_rows, slot_refs, gdn_segments) + self._sequence_parallel_lora_floor( + self._checkpoint_memory_floor(group_rows, slot_refs, gdn_segments), + group_rows, + signature, + ) if checkpoint_memory is None else checkpoint_memory ) backward = 0 if include_checkpoint_input_gradient and retained: # The backward's other end and cold transients, as _subforward_cost. - backward = retained + self._checkpoint_adapter_gradient_bytes( + backward = self._checkpoint_input_gradient_bytes( + group_rows, retained + ) + self._checkpoint_adapter_gradient_bytes( self._checkpoint_gradient_groups(group_rows, slot_refs) ) if profiled is None and any(grad for _, grad in group_rows): - backward += _impl._COLD_RECOMPUTE_TRANSIENT_BYTES + backward += self._cold_recompute_transient_bytes() static_compute = max( static_compute, max(retained, checkpoint_floor[0]) diff --git a/src/art/trainer_rank/_micro_batch_planner.py b/src/art/trainer_rank/_micro_batch_planner.py index 2b5414a83..d4d395217 100644 --- a/src/art/trainer_rank/_micro_batch_planner.py +++ b/src/art/trainer_rank/_micro_batch_planner.py @@ -1482,6 +1482,11 @@ 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", ()) + gaps = getattr(self, "_mixer_top_gaps", None) + rank_fields["mixer_top_gaps"] = None if gaps is None else list(gaps) + rank_fields["lora_modules_per_layer"] = getattr( + self, "_lora_modules_per_layer", 0 + ) rank_fields["geometry"] = asdict(self._geometry) rank_fields["topology"] = list(plan.signature.topology) diff --git a/src/art/trainer_rank/_planner_misses.py b/src/art/trainer_rank/_planner_misses.py index 274ec1295..d8c90cd2d 100644 --- a/src/art/trainer_rank/_planner_misses.py +++ b/src/art/trainer_rank/_planner_misses.py @@ -647,6 +647,29 @@ def report( "checkpointed_moe_layers recompute_modules moe_output_bytes_per_token " "moe_forward_stages".split() ) +# TP x SP floor facts; reports captured before them replay conservatively. +_OPTIONAL_RANK_FIELDS = frozenset({"mixer_top_gaps", "lora_modules_per_layer"}) + + +def _adapter_ranks_unavailable( + topology: tuple[int, ...], facts: dict[str, Any], signature: Any +) -> bool: + """A TP x SP floor estimate trains a named adapter whose ranks (the LoRA + intermediates' price) were not recorded, as in reports from before + sequence-parallel signatures carried slot shapes.""" + return ( + topology[1] > 1 + and bool(facts["checkpoint_layers"]) + and any( + group["grad"] and group["adapter"] is not None for group in facts["groups"] + ) + and not any( + len(shape) == 5 and shape[0] == 2 + for grad, shapes in signature.slot_shapes + if grad + for shape in shapes + ) + ) def _signature_values(values: dict[str, Any]) -> dict[str, Any]: @@ -711,9 +734,10 @@ def replay( if not state["estimates"]: raise ValueError("memory replay has no candidate estimates") # Reports from before the dense no-grad FC1 floor omit this fact; their - # estimator did not price that floor. + # estimator did not price that floor. Optional TP x SP floor facts may + # also be absent (older reports replay those terms conservatively). values: dict[str, Any] = {"dense_fc1_adapted": False, **state["rank"]} - if set(values) != _RANK_FIELDS | {"geometry", "topology"}: + if set(values) - _OPTIONAL_RANK_FIELDS != _RANK_FIELDS | {"geometry", "topology"}: raise ValueError( "incomplete replay: immutable rank fields differ (including MoE stages)" ) @@ -729,6 +753,21 @@ def replay( 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"]) + gaps = values.get("mixer_top_gaps") + if gaps is not None and ( + type(gaps) is not list + or len(gaps) != 2 + or any( + gap is not None and (type(gap) is not int or not 0 <= gap < layers) + for gap in gaps + ) + ): + raise ValueError("incomplete replay: invalid mixer layer order") + rank._mixer_top_gaps = None if gaps is None else tuple(gaps) + modules = values.get("lora_modules_per_layer") + if modules is not None and (type(modules) is not int or not 0 <= modules < 2**16): + raise ValueError("incomplete replay: invalid LoRA module count") + rank._lora_modules_per_layer = modules rank._geometry = ModelGeometry(**values["geometry"]) dp, tp, cp, pp = values["topology"] rank._topology_key = lambda: (dp, tp, cp, pp) @@ -854,6 +893,10 @@ def replay( request_cursor += count group_cursor += 1 key = _impl._MemorySignature(**_signature_values(item["signature"])) + if grouped and _adapter_ranks_unavailable( + rank._topology_key(), item["runtime_facts"], key + ): + raise ValueError("incomplete replay: TP x SP adapter ranks unavailable") rank._memory_profiles = ( {key: _impl._MemoryProfile(**item["profile"])} if item["profile"] is not None diff --git a/src/art/trainer_rank/_planner_replay.py b/src/art/trainer_rank/_planner_replay.py index b9a5a7832..a7b9970c3 100644 --- a/src/art/trainer_rank/_planner_replay.py +++ b/src/art/trainer_rank/_planner_replay.py @@ -37,6 +37,9 @@ "runtime_token_inventory_over_limit", "runtime_segment_inventory_over_limit", "head_positions_unavailable", + "head_statistics_fallback_unsupported", + "head_logits_statistics_unsupported", + "head_topk_targets_unsupported", "runtime_facts_over_limit", "runtime_stage_inventory_over_limit", "runtime_request_inventory_over_limit", @@ -102,6 +105,10 @@ def capture(rank: Any, plan: Any) -> dict[str, Any]: "_plan_head_workspace_bytes", "_plan_hybridep_growth_bytes", "_gdn_segment_layer_bytes", + "_sequence_parallel_workspace_bytes", + "_sequence_parallel_lora_floor", + "_cold_recompute_transient_bytes", + "_checkpoint_input_gradient_bytes", "_one_layer_recompute", "_topology_key", "_physical_tokens", @@ -218,6 +225,28 @@ def terms(checkpoint_grad: bool, ref: Any) -> list[Any]: if vocabulary else 0 ) + if projected and head_statistics_fallback(rank, requests, positions): + # The frozen head facts price only the kernel path at TP > 1. + raise ValueError("head_statistics_fallback_unsupported") + if ( + projected + and rank._topology_key()[1] > 1 + and any(r.logits for r in requests) + and any( + r.target_tokens is not None or r.top_k is not None for r in requests + ) + ): + # Facts record statistics, not requested logits beside them. + raise ValueError("head_logits_statistics_unsupported") + if ( + projected + and rank._topology_key()[1] > 1 + and group.grad_enabled + and any(r.top_k is not None for r in requests) + and any(r.target_tokens is not None for r in requests) + ): + # Facts record target rows, not top-k beside them. + raise ValueError("head_topk_targets_unsupported") backwards = ( target_backward and group.grad_enabled @@ -300,6 +329,21 @@ def terms(checkpoint_grad: bool, ref: Any) -> list[Any]: return facts +def head_statistics_fallback(rank: Any, requests: Any, positions: Any) -> bool: + """Whether any TP > 1 head chunk of these requests runs eager statistics.""" + if rank._topology_key()[1] == 1 or not any( + r.target_tokens is not None or r.top_k is not None for r in requests + ): + return False + vocabulary = _memory._head_vocabulary(rank) + return any( + not _impl._triton_head_stats(rank, chunk, vocabulary) + for chunk in _memory._head_chunk_sizes( + rank._head_projection_rows(requests, positions=positions, uncapped=True) + ) + ) + + def validate(facts: Any) -> None: """Accept only a finite primitive schema, never an executable reconstruction.""" @@ -639,13 +683,14 @@ def runtime_arguments( if arguments.get("hybridep_growth_bytes", 0): raise ValueError("hybridep_runtime_facts_unsupported") head = max( - _memory._dense_head_bytes(facts["head_vocabulary"], g["head_rows"]) - * ( - 1 - if not facts["head_target_backward"] - else 7 - if g["head_statistics"] - else 3 + _memory._frozen_head_bytes( + facts["head_vocabulary"], + g["head_rows"], + g["head_target_rows"], + target_backward=facts["head_target_backward"], + statistics=g["head_statistics"], + grad=g["grad"], + tp=self._topology_key()[1], ) for g in groups ) diff --git a/tests/unit/test_planner_runtime_fact_guards.py b/tests/unit/test_planner_runtime_fact_guards.py index aa2002960..db57915b3 100644 --- a/tests/unit/test_planner_runtime_fact_guards.py +++ b/tests/unit/test_planner_runtime_fact_guards.py @@ -38,6 +38,10 @@ def test_changed_custom_cost_is_not_reported_replay_complete(method, tmp_path): "_physical_tokens", "_plan_group_rows", "_plan_retained_tokens", + "_sequence_parallel_workspace_bytes", + "_sequence_parallel_lora_floor", + "_cold_recompute_transient_bytes", + "_checkpoint_input_gradient_bytes", ], ) def test_custom_shared_estimator_readers_decline(method): diff --git a/tests/unit/test_trainer_rank_tp_floor.py b/tests/unit/test_trainer_rank_tp_floor.py index 469aed777..f490a4307 100644 --- a/tests/unit/test_trainer_rank_tp_floor.py +++ b/tests/unit/test_trainer_rank_tp_floor.py @@ -1,12 +1,17 @@ -"""The full-recompute checkpoint floor under TP4 sequence parallelism. +"""The explicit full-recompute checkpoint floor under TP2/TP4 sequence parallelism. -CPU admission math, not a bound. One four-H200 trace of dense Qwen3.8-27B -(64 layers, TP4, sequence parallel) put the cold peak at each rank's boundary -shards plus a recompute workspace that the floor's repeated shards cover; -other shapes keep today's pricing. +CPU admission math, not a bound. Allocator traces of dense Qwen3.8-27B (64 +layers, sequence parallel, LoRA rank 4 at TP2 and 8 at TP4) at TP4 and TP2, +cold and warm, from 20,816 to 66,284 rows and one to four GDN segments, put +every peak at each rank's boundary shards plus the recomputed layer's +workspace; the floor covers those terms explicitly. Shallower, wider, +attention-only and higher-rank dense shapes are priced by the same terms +(untraced); TP8, CP, MoE, replicated QKV and TP without sequence parallelism +keep today's pricing. """ from dataclasses import replace +from types import SimpleNamespace from typing import Any import pytest @@ -14,8 +19,8 @@ from trainer_rank_test_support import fake_rank, recompute_model from art.trainer_rank import TrainerRank -from art.trainer_rank._impl import _COLD_RECOMPUTE_TRANSIENT_BYTES as COLD -from art.trainer_rank._impl import _MemorySignature +from art.trainer_rank._impl import _SEQUENCE_PARALLEL_COLD_TRANSIENT_BYTES as COLD +from art.trainer_rank._impl import _MemoryProfile, _MemorySignature H, F, LAYERS = 5120, 17408, 64 TP4 = (1, 4, 1, 1) @@ -24,12 +29,100 @@ # One segment's initial and final fp32 states over 12 local value heads, plus # conv history over 3 taps, in the recomputed GDN layer. SEGMENT = (4 * 12 * 128 * 128 + 2 * (2 * 4 * 128 + 12 * 128) * 3) * 2 +TP2 = (1, 2, 1, 1) +# The comparable traced TP x SP calls: topology, rows, logical rows, GDN +# segments, output bytes, whether the signature was profiled, and the peak +# above the call's baseline (identical on every rank). Excluded, see +# EXPLICIT-FLOOR.md: the 40,000-row planner-cold call (above the band: 0.69 GB +# of earlier tensors freed during it depress its harness peak) and production +# rank 0 at 23,878 rows (14.63 GB: an unexplained +1.98 GB over rank 1 and +# the trace, accepted and not covered). +TRACED = ( + pytest.param(TP4, 25_728, 25_727, 1, 102_908, False, 7_129_740_288, id="tp4-cold"), + pytest.param(TP4, 25_728, 25_727, 1, 102_908, True, 6_894_956_544, id="tp4-warm"), + pytest.param( + TP2, 23_878, 23_878, 1, 95_512, False, 12_645_404_160, id="tp2-23878-cold" + ), + pytest.param( + TP2, 23_878, 23_878, 1, 95_512, False, 12_544_440_832, id="tp2-23878-prod" + ), + pytest.param( + TP2, 23_878, 23_878, 1, 95_512, True, 12_298_095_616, id="tp2-23878-warm" + ), + pytest.param( + TP2, 20_816, 20_816, 1, 83_264, False, 10_935_935_488, id="tp2-20816-prod" + ), + pytest.param( + TP2, 20_816, 20_816, 1, 83_264, False, 10_689_269_248, id="tp2-20816-cold" + ), + pytest.param( + TP2, 20_816, 20_816, 1, 83_264, True, 10_721_021_440, id="tp2-20816-warm" + ), + pytest.param( + TP2, 40_000, 40_000, 2, 160_000, True, 20_601_088_000, id="tp2-40000-warm" + ), + pytest.param( + TP2, 23_878, 23_878, 1, 95_512, True, 12_129_966_592, id="tp2-23878-warm-2" + ), + pytest.param( + TP2, 66_284, 68_158, 4, 272_632, False, 34_932_728_320, id="tp2-66284-cold" + ), + pytest.param( + TP2, 66_284, 68_158, 4, 272_632, True, 34_141_831_680, id="tp2-66284-warm" + ), + pytest.param( + TP2, 66_284, 68_158, 4, 272_632, True, 34_616_171_520, id="tp2-66284-prod" + ), +) + + +GDN_WIDTH = 4 * 16 * 128 + 8 * 48 * 128 +GATED_ATTENTION_WIDTH = 7 * 24 * 256 + 3 * 4 * 256 -def tp_rank(layers=LAYERS, *, ffn=F, topology=TP4, sequence_parallel=True, **config): +def workspace( + rows, tp, ffn=F, stage=6, mixers=((GATED_ATTENTION_WIDTH, 0), (GDN_WIDTH, 1)) +): + """The traced recompute workspace (bytes): gathered 2H + stage F/TP, 3.25H + of norms per sharded row, the fp32 rotary cache, and for each mixer kind + (width, layers above its highest layer) its width/TP plus the boundary + shards beyond L live at that layer's recompute.""" + shard = -(-rows // tp) * H + peak = max(rows * -(-width // tp) + (1 - gap) * shard for width, gap in mixers) + gathered = rows * (2 * H + -(-stage * ffn // tp)) + rotary = 2 * rows * 256 if any(width != GDN_WIDTH for width, _ in mixers) else 0 + return (gathered + -(-rows // tp) * -(-13 * H // 4) + peak + rotary) * 2 + + +def signature(topology, rank=0, grad=True): + """A gradient signature whose slot carries LoRA modules of ``rank``.""" + shapes = ((True, ((2, H, rank, rank, H),)),) if rank else () + return _MemorySignature( + topology, (1, None), 1, (), grad, (grad,), slot_shapes=shapes + ) + + +def tp_rank( + layers=LAYERS, + *, + ffn=F, + topology=TP4, + sequence_parallel=True, + decoder_layers=None, + **config, +): from megatron.core.transformer.transformer_block import TransformerBlock - model = recompute_model(TransformerBlock, H, layers, sequence_parallel, **config) + model = recompute_model( + TransformerBlock, + H, + layers, + sequence_parallel, + layers=decoder_layers or (), + **config, + ) + if decoder_layers is not None: + model.embedding = torch.nn.Linear(1, 1).bfloat16() # a BF16 parameter r: Any = fake_rank(TrainerRank, [model], hidden_size=H, num_layers=layers) # Qwen3.8-27B: gated attention every fourth layer, GDN otherwise. r._geometry = replace( @@ -48,10 +141,75 @@ def tp_rank(layers=LAYERS, *, ffn=F, topology=TP4, sequence_parallel=True, **con r._attention_output_gate = True r._gdn_layers = layers * 3 // 4 r._topology_key = lambda: topology + if decoder_layers is None: + # Layers 3, 7, ..., 63 are attention: the top GDN layer is one below. + r._mixer_top_gaps = (0, 1) + # Six adapted modules per GDN layer, seven per attention layer. + r._lora_modules_per_layer = 7 return r -def _required(r, group_rows=((ROWS, True),), topology=TP4, gdn_segments=1): +POLICY_RANK = 8192 + + +def adapted_layers(per_layer, *, rank=POLICY_RANK): + """Decoder layers of real LoRA modules holding a ``policy`` slot of ``rank``.""" + from art.megatron.lora import LoRA, LoRASlotRef + + ref = LoRASlotRef("checkpoint", "policy") + + def module(): + lora = LoRA.__new__(LoRA) + torch.nn.Module.__init__(lora) + lora._slot_keys = {ref: "policy"} + lora._slot_modules = { + "policy": SimpleNamespace( + A_T=torch.empty(H, rank, device="meta"), + B_T=torch.empty(rank, H, device="meta"), + ) + } + return lora + + return ref, [ + torch.nn.ModuleList([module() for _ in range(count)]) for count in per_layer + ] + + +def produced_signature(r, ref, rows): + """The memory signature the planner builds for one gradient request.""" + from art.trainer_rank import ForwardInput + + tokens = torch.zeros(rows, dtype=torch.long) + return r._memory_signature_from_requests( + [ForwardInput(input_tokens=tokens, target_tokens=tokens)], + slot_group_count=1, + grad_modes=(True,), + slot_groups=((ref, True),), + ) + + +def lora_workspace(r, ref, rows): + """The floor's workspace for one gradient group under ``ref``'s signature.""" + return r._subforward_cost( + packed_tokens=rows, + output_bytes=rows * 4, + signature=produced_signature(r, ref, rows), + logical_tokens=rows, + gdn_segments=1, + group_rows=((rows, True),), + ).checkpoint_workspace + + +def _required( + r, + group_rows=((ROWS, True),), + topology=TP4, + gdn_segments=1, + output=OUTPUT, + logical=None, + rank=0, +): + shapes = ((True, ((2, H, rank, rank, H),)),) if rank else () signature = _MemorySignature( topology, (1, None), @@ -59,35 +217,840 @@ def _required(r, group_rows=((ROWS, True),), topology=TP4, gdn_segments=1): (), any(grad for _, grad in group_rows), tuple(grad for _, grad in group_rows), + slot_shapes=shapes, ) return r._subforward_cost( packed_tokens=sum(rows for rows, _ in group_rows), - output_bytes=OUTPUT, + output_bytes=output, signature=signature, - logical_tokens=sum(rows for rows, _ in group_rows) - 1, + logical_tokens=sum(rows for rows, _ in group_rows) - 1 + if logical is None + else logical, gdn_segments=gdn_segments, group_rows=group_rows, ) -def test_the_traced_tp4_wave_prices_its_boundary_shards_and_their_repeat(): - r = tp_rank() - retained, workspace = r._checkpoint_memory_floor(((ROWS, True),)) - # Each rank saves a quarter of every boundary: the traced 4.215 GB. - assert retained == ROWS // 4 * LAYERS * H * 2 == 4_215_275_520 - # Without a segment count, only the TP-padding roots' states. - assert workspace == 3 * SEGMENT - cost = _required(r) - assert cost.checkpoint_input_gradient == retained - # One segment plus up to three TP-padding roots, each with its states, - # beside the unprofiled first execution's transients. - state = cost.checkpoint_workspace - COLD - assert state == 4 * SEGMENT - assert cost.required == int((OUTPUT + 2 * retained + state + COLD) * 1.1) - # Measured cold on all four ranks: 7.130 GB (7.060 GB in production), all - # but the boundaries a transient recompute workspace; this raw floor - # (8.43 GB) covers it. Today's cold admission was 4.637 GB. - assert cost.required / 1.1 > 7.130e9 +def _warm(r, topology, rank=0, profile=None): + """Profile the signature (by default below the floor, so the floor prices).""" + r._memory_profiles[signature(topology, rank)] = profile or _MemoryProfile(1.0, 1) + return r + + +@pytest.mark.parametrize( + ("topology", "rows", "logical", "segments", "output", "warm", "observed"), TRACED +) +def test_the_floor_covers_every_comparable_traced_peak_within_five_percent( + topology, rows, logical, segments, output, warm, observed +): + rank = 4 if topology == TP2 else 8 # the traced adapters + r = tp_rank(topology=topology) + if warm: + _warm(r, topology, rank) + cost = _required(r, ((rows, True),), topology, segments, output, logical, rank) + raw = cost.required / 1.1 + # Brad's standard: at or above the observed peak, by at most 5%. + assert observed <= raw <= 1.05 * observed + + +@pytest.mark.parametrize( + ("topology", "rows"), [(TP4, ROWS), (TP2, 23_878), (TP2, 23_877)] +) +def test_the_floor_is_the_shards_the_explicit_workspace_and_one_gradient_shard( + topology, rows +): + tp = topology[1] + r = tp_rank(topology=topology) + retained, floor_workspace = r._checkpoint_memory_floor(((rows, True),)) + shard = -(-rows // tp) # an odd row count still rounds up per rank + # Each rank saves its sequence shard of every boundary. + assert retained == shard * LAYERS * H * 2 + # One GDN layer's states for each of the TP - 1 padding roots. + roots = (tp - 1) * SEGMENT * 4 // tp + assert floor_workspace == workspace(rows, tp) + roots + cost = _required(r, ((rows, True),), topology, rank=4) + # One input-gradient shard, not a repeat of every boundary. + assert cost.checkpoint_input_gradient == shard * H * 2 + state = tp * SEGMENT * 4 // tp + # Seven adapted modules' rank-4 ``x @ A`` outputs over the gathered rows. + lora = rows * 7 * 4 * 2 + assert cost.checkpoint_workspace == workspace(rows, tp) + lora + state + COLD + assert cost.required == int( + (OUTPUT + retained + workspace(rows, tp) + lora + state + COLD + shard * H * 2) + * 1.1 + ) + + +def test_the_memory_check_path_prices_the_same_floor(): + """_memory_check's estimator (include_checkpoint_input_gradient=True).""" + for profiled in (False, True): + r = tp_rank(topology=TP2) + if profiled: + _warm(r, TP2, 4) + cost = _required(r, ((23_879, True),), TP2, rank=4) + estimate = r._estimate_required_memory_bytes_from_values( + packed_tokens=23_879, + output_bytes=OUTPUT, + signature=signature(TP2, 4), + logical_tokens=23_878, + gdn_segments=1, + group_rows=((23_879, True),), + ) + # The repeated boundaries would add 63 shards (about 7.7 GB). + assert estimate == cost.required + + +def test_lora_rank_moves_the_price_by_its_intermediates(): + rows = 25_728 + low, high = ( + _required(tp_rank(), ((rows, True),), TP4, rank=rank).checkpoint_workspace + for rank in (128, 8192) + ) + # Seven modules' rows x rank BF16 outputs: the FC1 gate/up pair alone + # adds 829,882,368 bytes. + assert high - low == rows * 7 * (8192 - 128) * 2 + + +def test_attention_only_slots_price_their_lora_intermediates(): + """An attention-only TP2 model: ranks come from the real slot-shape and + module-count producers, not a hand-built signature.""" + rows = 23_878 + ref, layers = adapted_layers([7] * LAYERS) + r = tp_rank(topology=TP2, decoder_layers=layers) + r._gdn_layers = 0 + assert r._mixer_top_gaps == (0, None) # read from the decoder + assert r._lora_modules_per_layer == 7 + shapes = produced_signature(r, ref, rows).slot_shapes + assert shapes == ((True, ((2, H, POLICY_RANK, POLICY_RANK, H),) * 7 * LAYERS),) + base = lora_workspace(r, None, rows) + # Seven modules' rows x 8,192 BF16 intermediates: 2,738,520,064 bytes. + assert lora_workspace(r, ref, rows) - base == rows * 7 * POLICY_RANK * 2 + + +def test_an_unknown_per_layer_count_charges_every_adapted_module(): + """Seven modules concentrated in one of 64 layers: the mean (two) would + omit 1,956,085,760 bytes; a report without the count charges all seven.""" + rows = 23_878 + ref, layers = adapted_layers([7] + [0] * (LAYERS - 1)) + r = tp_rank(topology=TP2, decoder_layers=layers) + r._mixer_top_gaps = (0, 1) + assert r._lora_modules_per_layer == 7 + live = lora_workspace(r, ref, rows) - lora_workspace(r, None, rows) + assert live == rows * 7 * POLICY_RANK * 2 + r._lora_modules_per_layer = None # as replayed from an older report + assert lora_workspace(r, ref, rows) - lora_workspace(r, None, rows) == live + + +def test_two_gradient_groups_charge_one_group_workspace_and_gradient_shard(): + r = tp_rank(topology=TP2) + groups = ((20_000, True), (12_000, True)) + retained, floor_workspace = r._checkpoint_memory_floor(groups, None, 2) + # Both groups' boundaries stay live; their backwards run one at a time. + assert retained == (10_000 + 6_000) * LAYERS * H * 2 + # Two segments plus one padding root per group (states remain summed). + assert floor_workspace == workspace(20_000, 2) + 4 * SEGMENT * 4 // 2 + cost = _required(r, groups, TP2, 2) + assert cost.checkpoint_input_gradient == 10_000 * H * 2 + + +def test_a_clamped_swiglu_prices_its_seven_f_stage(): + rows = 23_878 + r = tp_rank(topology=TP2) + r._mlp_activation_factor = 7 + assert ( + r._checkpoint_memory_floor(((rows, True),))[1] + == workspace(rows, 2, stage=7) + SEGMENT * 4 // 2 + ) + + +@pytest.mark.parametrize("order", ["traced", "gdn_top", "unknown"]) +def test_the_peak_layer_sets_the_boundary_shards(order): + rows = 23_878 + r = tp_rank(topology=TP2) + gaps = {"traced": (0, 1), "gdn_top": (1, 0), "unknown": None}[order] + r._mixer_top_gaps = gaps + gdn_gap = 1 if order == "traced" else 0 + attention_gap = 1 if order == "gdn_top" else 0 + expected = workspace( + rows, + 2, + mixers=((GATED_ATTENTION_WIDTH, attention_gap), (GDN_WIDTH, gdn_gap)), + ) + assert r._checkpoint_memory_floor(((rows, True),))[1] == expected + SEGMENT * 2 + if order != "traced": + # A top-layer GDN peak (or an unknown order) holds L + 1 shards. + assert expected - workspace(rows, 2) == rows // 2 * H * 2 + + +def test_attention_only_tp2_is_priced_at_its_attention_width(): + rows = 23_877 + r = tp_rank(topology=TP2) + r._gdn_layers = 0 + r._attention_output_gate = False + r._mixer_top_gaps = (0, None) + attention = 5 * 24 * 256 + 3 * 4 * 256 + assert r._checkpoint_memory_floor(((rows, True),)) == ( + -(-rows // 2) * LAYERS * H * 2, + workspace(rows, 2, mixers=((attention, 0),)), + ) + + +def _with_head(r, vocabulary, tp): + from megatron.core.tensor_parallel.layers import ColumnParallelLinear + + head = ColumnParallelLinear.__new__(ColumnParallelLinear) + torch.nn.Module.__init__(head) + head.weight = torch.nn.Parameter( + torch.empty(vocabulary // tp, H, dtype=torch.bfloat16, device="meta") + ) + head.output_size_per_partition = vocabulary // tp + head.output_size = vocabulary + head.input_size = H + r.runtime.model[0].output_layer = head + r.runtime.model[0].share_embeddings_and_output_weights = False + r._padded_vocab_size = vocabulary + return r + + +def test_the_tp2_head_stage_is_priced_with_the_floor(): + from art.trainer_rank import _memory + + vocabulary = 248_320 + r = _with_head(tp_rank(topology=TP2), vocabulary, 2) + # Each rank projects its vocabulary shard, 512 rows at a time. + assert _memory._head_vocabulary(r) == vocabulary // 2 + # Without the floor's eligibility (no SP) TP2 keeps today's zero. + no_sp = _with_head(tp_rank(topology=TP2, sequence_parallel=False), vocabulary, 2) + assert _memory._head_vocabulary(no_sp) == 0 + # A short wave's head stage (logits plus both target-backward buffers) + # outweighs its recompute workspace and prices the floor. + head = 3 * r._head_workspace_bytes(512) + signature = _MemorySignature(TP2, (1, None), 1, (), True, (True,)) + cost = r._subforward_cost( + packed_tokens=512, + output_bytes=2_048, + signature=signature, + logical_tokens=512, + gdn_segments=1, + group_rows=((512, True),), + head_workspace_bytes=head, + ) + assert head > workspace(512, 2) + SEGMENT * 2 + assert cost.checkpoint_workspace == head + COLD + + +# The bounded eager statistics' increment for a 512-row, 124,160-entry chunk: +# one 8-row FP32 sub-chunk plus per-row statistics (3,981,312 bytes). +BOUNDED = 4 * 124_160 * 8 + 16 * 512 + + +def _labelled(rows): + """One CPU request whose ``rows`` positions are all labelled targets.""" + tokens = torch.zeros(rows, dtype=torch.long) + return [ + SimpleNamespace( + input_tokens=tokens, target_tokens=tokens, logits=False, top_k=None + ) + ] + + +def _head_stage(r, requests): + """The head charge as planned: the capped projection, then each chunk.""" + rows = r._head_projection_rows(requests) + return r._group_head_workspace_bytes(rows, requests, grad_enabled=True) + + +def test_the_tp2_head_stage_follows_the_statistics_path(monkeypatch): + from art.trainer_rank import _impl, _memory + + r = _with_head(tp_rank(topology=TP2), 248_320, 2) + monkeypatch.setattr(_memory, "_head_target_backward", lambda rank: True) + dense = r._head_workspace_bytes(512) + # The same predicate as _try_triton_stats: a CPU device, the + # ART_TRAINER_RANK_TRITON_TOPK switch and short chunks take the eager + # FP32 fallback, about seven BF16 buffers. + assert _head_stage(r, _labelled(512)) == 7 * dense + monkeypatch.setattr(_impl, "_triton_head_stats", lambda rank, rows, vocab: True) + # The kernel path, plus the bounded fallback's increment should it fail. + assert _head_stage(r, _labelled(512)) == 3 * dense + BOUNDED + monkeypatch.setenv("ART_TRAINER_RANK_TRITON_TOPK", "0") + assert not _impl._triton_stats_enabled(True, 512) + monkeypatch.delenv("ART_TRAINER_RANK_TRITON_TOPK") + assert _impl._triton_stats_enabled(True, 64) + assert not _impl._triton_stats_enabled(True, 63) + assert not _impl._triton_stats_enabled(False, 512) + + +@pytest.mark.parametrize( + ("rows", "minimum", "expected", "eager"), + [ + # Sol: a 511-row eager tail behind a 512-row Triton chunk. + (1_023, "512", 7 * 511 * 124_160 * 2, True), # 888,240,640 bytes + # Schulman: an 812-row wave's 300-row eager tail. + (812, "512", 7 * 300 * 124_160 * 2, True), # 521,472,000 bytes + # The default threshold (64): every chunk runs Triton (381,419,520 + # bytes), or the bounded statistics should a kernel fail. + (1_023, None, 3 * 512 * 124_160 * 2 + BOUNDED, False), + # A 40-row default tail falls back, below the full chunk's charge. + (552, None, 3 * 512 * 124_160 * 2 + BOUNDED, True), + ], +) +def test_every_head_chunk_is_priced_for_its_own_statistics_path( + monkeypatch, rows, minimum, expected, eager +): + from art.trainer_rank import _impl, _memory + + r = _with_head(tp_rank(topology=TP2), 248_320, 2) + r.device = torch.device("cuda", 0) # a CUDA rank whose kernels import + monkeypatch.setattr(_impl, "_triton_stats_importable", lambda: True) + monkeypatch.setattr(_memory, "_head_target_backward", lambda rank: True) + if minimum is None: + monkeypatch.delenv("ART_TRAINER_RANK_TRITON_MIN_ROWS", raising=False) + else: + monkeypatch.setenv("ART_TRAINER_RANK_TRITON_MIN_ROWS", minimum) + # The planner's projection is capped at one chunk; the tail is not. + assert r._head_projection_rows(_labelled(rows)) == 512 + assert _head_stage(r, _labelled(rows)) == expected + # Replay capture refuses exactly the waves with an eager chunk. + from art.trainer_rank._planner_replay import head_statistics_fallback + + assert head_statistics_fallback(r, _labelled(rows), None) == eager + + +def _cuda_head(monkeypatch, vocabulary=248_320): + """A TP2 rank with a head on a CUDA device whose kernels import.""" + from art.trainer_rank import _impl, _memory + + r = _with_head(tp_rank(topology=TP2), vocabulary, 2) + r.device = torch.device("cuda", 0) + monkeypatch.setattr(_impl, "_triton_stats_importable", lambda: True) + monkeypatch.setattr(_memory, "_head_target_backward", lambda rank: True) + return r + + +def _positioned(*positions): + """Labelled requests and their packed positions.""" + requests = [_labelled(len(row))[0] for row in positions] + return requests, [torch.tensor(row) for row in positions] + + +@pytest.mark.parametrize( + ("first", "second", "exact"), + [ + # Sol: two 812-row requests sharing 600 positions, a 1,024-row union + # of two kernel chunks; one request alone has a 300-row eager tail. + ( + range(812), + [*range(600), *range(812, 1024)], + 3 * 512 * 124_160 * 2 + BOUNDED, + ), + # Two 512-row requests sharing 212 positions: their no-sharing sum is + # two kernel chunks, but the 812-row union has a 300-row eager tail. + (range(512), [*range(212), *range(512, 812)], 7 * 300 * 124_160 * 2), + ], +) +def test_head_bounds_envelope_every_possible_projected_union( + monkeypatch, first, second, exact +): + monkeypatch.setenv("ART_TRAINER_RANK_TRITON_MIN_ROWS", "512") + r = _cuda_head(monkeypatch) + requests, positions = _positioned(list(first), second) + + def charge(**bounds): + rows = r._head_projection_rows( + requests, + positions=bounds.get("positions"), + lower_bound=bounds.get("lower_bound", False), + ) + return r._group_head_workspace_bytes( + rows, requests, grad_enabled=True, **bounds + ) + + # The selected layout's packed union prices every chunk exactly. + assert charge(positions=positions) == exact + # Width search: the rejection bound never exceeds it, the acceptance + # bound never falls below it. + assert charge(lower_bound=True) <= exact <= charge() + + +def test_a_kernel_failure_is_priced_only_at_its_chunk_shape(monkeypatch): + from art.trainer_rank import _impl, _memory + + r = _cuda_head(monkeypatch, vocabulary=16) # an 8-entry vocabulary shard + + def statistics(rows, *, top_k, logsumexp, max_top_k=0): + """Run one head chunk's statistics with the given kernel outcomes.""" + logits = torch.zeros(rows, 8) + stats = (torch.zeros(rows), torch.ones(rows)) + with monkeypatch.context() as patch: + patch.setattr( + TrainerRank, "_local_logits_from_hidden_rows", lambda *a, **k: logits + ) + # The chunk is attempted (as on CUDA) and the kernels return these. + patch.setattr(_impl, "_triton_stats_enabled", lambda cuda, rows: True) + patch.setattr( + _impl, + "_try_triton_local_topk_stats", + lambda *a, **k: stats if top_k else None, + ) + patch.setattr( + _impl, "_try_triton_stats", lambda *a, **k: stats if logsumexp else None + ) + patch.setattr(_impl, "_all_reduce_tensor_parallel_max", lambda t: t) + patch.setattr(_impl, "_all_reduce_tensor_parallel_sum", lambda t: t) + patch.setattr(_impl, "_vocab_parallel_log_z", lambda t: t[:, 0].float()) + r._local_head_stats( + None, logits, output_weight=None, need_log_z=True, max_top_k=max_top_k + ) + + def priced(rows): + """The head charge in chunk buffers, less the bounded increment.""" + charge = _head_stage(r, _labelled(rows)) + if charge >= 7 * r._head_workspace_bytes(rows): + return 7 + bounded = _memory._eager_stats_extra_bytes(8, rows) + return (charge - bounded) // r._head_workspace_bytes(rows) + + # A fused top-k failure recovered by the Triton logsumexp: no fallback. + statistics(512, top_k=False, logsumexp=True, max_top_k=4) + assert priced(512) == 3 + # A 64-row kernel failure takes the eager fallback (that wave was priced + # for the kernel); 64-row chunks are priced for the fallback from then on. + statistics(64, top_k=False, logsumexp=False) + assert priced(64) == 7 + # A 512-row chunk whose kernel succeeds stays priced for the kernel. + statistics(512, top_k=False, logsumexp=True) + assert priced(512) == 3 + # A later 64-row success clears that shape. + statistics(64, top_k=False, logsumexp=True) + assert priced(64) == 3 + + +class _LiveBytes: + """Peak bytes of new CPU tensor storage created inside the context.""" + + def __enter__(self): + import weakref + + from torch.utils._python_dispatch import TorchDispatchMode + from torch.utils._pytree import tree_leaves + + owner = self + self.current = self.peak = 0 + live: set[int] = set() + + def free(pointer, size): + live.discard(pointer) + owner.current -= size + + class Mode(TorchDispatchMode): + def __torch_dispatch__(self, func, types, args=(), kwargs=None): + out = func(*args, **(kwargs or {})) + inputs = { + t.untyped_storage().data_ptr() + for t in tree_leaves((args, kwargs)) + if isinstance(t, torch.Tensor) + } + for t in tree_leaves(out): + if not isinstance(t, torch.Tensor): + continue + pointer = t.untyped_storage().data_ptr() + size = t.untyped_storage().nbytes() + # Views and in-place results reuse an input's storage. + if not size or pointer in inputs or pointer in live: + continue + live.add(pointer) + owner.current += size + owner.peak = max(owner.peak, owner.current) + weakref.finalize(t, free, pointer, size) + return out + + self._mode = Mode() + self._mode.__enter__() + return self + + def __exit__(self, *exc): + self._mode.__exit__(*exc) + + +def test_a_failed_kernel_falls_back_within_the_kernel_buffers(monkeypatch): + """Admission priced the statistics kernel (three BF16 chunk buffers at the + head stage: logits, their gradient and the target gather's dense + gradient). The bounded eager statistics add one FP32 row sub-chunk to + the logits in forward, and to the logits and their gradient in backward; + the unchunked fallback adds six buffers in forward.""" + from art.trainer_rank import _impl + + monkeypatch.setattr(_impl, "_all_reduce_tensor_parallel_max", lambda t: t) + monkeypatch.setattr(_impl, "_all_reduce_tensor_parallel_sum", lambda t: t) + torch.manual_seed(0) + logits = torch.randn(512, 1_000, dtype=torch.bfloat16, requires_grad=True) + buffer = 512 * 1_000 * 2 # one BF16 chunk buffer (D) + with _LiveBytes() as forward: + local_max, local_sum = _impl._eager_local_logsumexp_stats(logits) + log_z = local_max.detach() + torch.log(local_sum) + assert forward.peak <= buffer / 16 + weights = torch.randn(512) + with _LiveBytes() as backward: + (gradient,) = torch.autograd.grad((log_z * weights).sum(), logits) + assert backward.peak <= buffer + buffer / 16 + # The priced TP2 head charge for this chunk shape (1,000-entry shard) + # covers the measured live set on both paths: forward, the logits plus + # the forward increment; backward, the logits, the target gather's dense + # gradient and the backward increment. + r = _cuda_head(monkeypatch, vocabulary=2_000) + targets = _labelled(512) + trained = r._group_head_workspace_bytes(512, targets, grad_enabled=True) + assert trained >= buffer + buffer + backward.peak + inference = r._group_head_workspace_bytes(512, targets, grad_enabled=False) + assert inference >= buffer + forward.peak + # The unchunked FP32 fallback, for contrast: about six buffers in forward. + with _LiveBytes() as unchunked: + reference = _impl._vocab_parallel_log_z(logits) + assert unchunked.peak >= 4 * buffer + # Rows are independent: the same values and gradients. + (expected,) = torch.autograd.grad((reference * weights).sum(), logits) + assert torch.equal(log_z, reference) + assert torch.allclose(gradient.float(), expected.float(), rtol=1e-2, atol=1e-6) + + +def test_the_bounded_statistics_never_write_fp32_logits(monkeypatch): + """An FP32 logits chunk: ``.float()`` would alias it, so every sub-chunk + goes through an owned FP32 work buffer.""" + from art.trainer_rank import _impl + + monkeypatch.setattr(_impl, "_all_reduce_tensor_parallel_max", lambda t: t) + monkeypatch.setattr(_impl, "_all_reduce_tensor_parallel_sum", lambda t: t) + torch.manual_seed(2) + logits = torch.randn(512, 1_000, dtype=torch.float32, requires_grad=True) + original = logits.detach().clone() + local_max, local_sum = _impl._eager_local_logsumexp_stats(logits) + log_z = local_max.detach() + torch.log(local_sum) + assert torch.equal(logits.detach(), original) + weights = torch.randn(512) + (gradient,) = torch.autograd.grad((log_z * weights).sum(), logits) + assert torch.equal(logits.detach(), original) + reference = _impl._vocab_parallel_log_z(logits) + (expected,) = torch.autograd.grad((reference * weights).sum(), logits) + assert torch.allclose(log_z, reference, rtol=1e-6, atol=1e-6) + assert torch.allclose(gradient, expected, rtol=1e-5, atol=1e-7) + + +def test_the_head_statistics_fall_back_to_the_bounded_path_after_a_kernel_failure( + monkeypatch, +): + from art.trainer_rank import _impl + + torch.manual_seed(1) + r = _cuda_head(monkeypatch, vocabulary=2_000) + logits = torch.randn(512, 1_000, dtype=torch.bfloat16) + monkeypatch.setattr(_impl, "_all_reduce_tensor_parallel_max", lambda t: t) + monkeypatch.setattr(_impl, "_all_reduce_tensor_parallel_sum", lambda t: t) + reference = _impl._vocab_parallel_log_z(logits) + monkeypatch.setattr( + TrainerRank, "_local_logits_from_hidden_rows", lambda *a, **k: logits + ) + # An attempted kernel (as on CUDA) that fails, fused top-k and logsumexp. + monkeypatch.setattr(_impl, "_triton_stats_enabled", lambda cuda, rows: True) + monkeypatch.setattr(_impl, "_try_triton_local_topk_stats", lambda *a, **k: None) + monkeypatch.setattr(_impl, "_try_triton_stats", lambda *a, **k: None) + + def unchunked(_logits): + raise AssertionError("the unchunked FP32 fallback ran after an attempt") + + monkeypatch.setattr(_impl, "_vocab_parallel_log_z", unchunked) + _, log_z, top_k = r._local_head_stats( + None, logits, output_weight=None, need_log_z=True, max_top_k=4 + ) + assert torch.equal(log_z, reference) + assert top_k is not None + assert torch.equal(top_k[0], torch.topk(logits.float(), 4, dim=-1).values) + # The shape is recorded for later admissions and replay. + assert (512, 1_000) in r._triton_head_stats_failures + + +def _requests(rows, *, targets=True, top_k=None, logits=False): + tokens = torch.zeros(rows, dtype=torch.long) + return [ + SimpleNamespace( + input_tokens=tokens, + target_tokens=tokens if targets else None, + logits=logits, + top_k=top_k, + ) + ] + + +def test_replay_prices_the_tp2_head_stage_as_live_admission(monkeypatch): + from art.trainer_rank import _memory + + r = _cuda_head(monkeypatch) # kernel path: capture declines the fallback + dense = r._head_workspace_bytes(512) + cases = ( + # rows, gradient, requests, target rows (as captured), kernel buffers + (512, True, _requests(512), 512, 3), # logits, statistics, target gather + (512, False, _requests(512), 0, 1), + (300, True, _requests(300), 300, 3), + # Top-k only: logits, statistics gradient and top-k's scatter. + (512, True, _requests(512, targets=False, top_k=4), 0, 3), + (512, False, _requests(512, targets=False, top_k=4), 0, 1), + ) + for rows, grad, requests, target_rows, buffers in cases: + live = r._group_head_workspace_bytes(rows, requests, grad_enabled=grad) + bounded = _memory._eager_stats_extra_bytes(124_160, rows) + assert live == buffers * r._head_workspace_bytes(rows) + bounded + frozen = _memory._frozen_head_bytes( + 124_160, + rows, + target_rows, + target_backward=True, + statistics=True, + grad=grad, + tp=2, + ) + assert frozen == live + # Requested logits: the local logits, their indexed copy, the gather + # buffer and its concatenation (2 + 2 TP buffers), live and replayed. + logits = _requests(512, targets=False, logits=True) + assert r._group_head_workspace_bytes(512, logits, grad_enabled=False) == 6 * dense + for tp, units in ((2, 6), (4, 10)): + assert ( + _memory._frozen_head_bytes( + 124_160, + 512, + 0, + target_backward=True, + statistics=False, + grad=False, + tp=tp, + ) + == units * dense + ) + # TP1 replay is #1068's capacity charge (its gather is the identity). + for statistics, units in ((True, 7), (False, 3)): + assert ( + _memory._frozen_head_bytes( + 1_000, + 512, + 0, + target_backward=True, + statistics=statistics, + grad=True, + tp=1, + ) + == units * 512 * 1_000 * 2 + ) + + +class _ShapeChangingKernel(torch.autograd.Function): + """A mocked statistics kernel saving a different tensor set than the eager + paths (as the real kernel's top-k tokens do).""" + + @staticmethod + def forward(ctx, logits): + local_max = logits.max(dim=-1).values.float() + local_sum = torch.exp(logits.float() - local_max[:, None]).sum(dim=-1) + ctx.save_for_backward(logits, local_max, torch.empty(logits.shape[0], 0)) + return local_max, local_sum + + @staticmethod + def backward(ctx, *grad_outputs): + _grad_max, grad_sum = grad_outputs + logits, local_max, _ = ctx.saved_tensors + return (torch.exp(logits.float() - local_max[:, None]) * grad_sum[:, None]).to( + logits.dtype + ) + + +@pytest.mark.parametrize("forward_fails", [True, False]) +def test_a_checkpoint_recompute_replays_the_forward_statistics_path( + monkeypatch, forward_fails +): + """Schulman: a kernel that fails in the forward and succeeds in the + recompute (or the reverse) changed the saved tensors, a CheckpointError + under non-reentrant checkpointing.""" + from art.trainer_rank import _impl + + r = _cuda_head(monkeypatch, vocabulary=64) + weight = torch.randn(16, 32) + monkeypatch.setattr( + TrainerRank, + "_local_logits_from_hidden_rows", + lambda self, model, hidden, output_weight: hidden @ weight, + ) + monkeypatch.setattr(_impl, "_triton_stats_enabled", lambda cuda, rows: True) + monkeypatch.setattr(_impl, "_all_reduce_tensor_parallel_max", lambda t: t) + monkeypatch.setattr(_impl, "_all_reduce_tensor_parallel_sum", lambda t: t) + attempts = [] + + def kernel(name, logits, **kwargs): + # The first attempt (the forward) fails or succeeds; later ones flip. + attempts.append(name) + succeeds = (len(attempts) > 1) == forward_fails + return _ShapeChangingKernel.apply(logits) if succeeds else None + + monkeypatch.setattr(_impl, "_try_triton_stats", kernel) + hidden = torch.randn(64, 16, requires_grad=True) + _, log_z, _ = r._checkpointed_head_stats( + None, hidden, output_weight=None, need_log_z=True, max_top_k=0 + ) + if forward_fails: + # The recompute replays the bounded eager statistics: no kernel retry. + log_z.sum().backward() + assert attempts == ["local_logsumexp_stats"] + assert hidden.grad is not None + else: + # A kernel that fails only in the recompute cannot reproduce the + # forward's saved tensors: a clear error, not a CheckpointError. + with pytest.raises(RuntimeError, match="failed in a checkpoint recompute"): + log_z.sum().backward() + + +def test_replay_declines_tp_sp_adapter_estimates_without_ranks(): + from art.trainer_rank._planner_misses import _adapter_ranks_unavailable + + facts = {"checkpoint_layers": 64, "groups": [{"grad": True, "adapter": {}}]} + bare = signature(TP2) + ranked = signature(TP2, rank=4) + assert _adapter_ranks_unavailable(TP2, facts, bare) + assert not _adapter_ranks_unavailable(TP2, facts, ranked) + # Gradient entries without a usable (2-D) rank shape decline as well. + for shapes in (((True, ()),), ((True, ((),)),), ((False, ((2, H, 4, 4, H),)),)): + partial = replace(bare, slot_shapes=shapes) + assert _adapter_ranks_unavailable(TP2, facts, partial) + assert not _adapter_ranks_unavailable((1, 1, 1, 1), facts, bare) + base = {"checkpoint_layers": 64, "groups": [{"grad": True, "adapter": None}]} + assert not _adapter_ranks_unavailable(TP2, base, bare) + + +def test_a_tp2_wide_ffn_is_priced_by_its_fc1_stage(): + rows = 23_878 + narrow = tp_rank(topology=TP2)._checkpoint_memory_floor(((rows, True),)) + wide = tp_rank(topology=TP2, ffn=4 * F)._checkpoint_memory_floor(((rows, True),)) + # The same boundaries; the FC1 stage grows by 6 x 3F/TP per gathered row. + assert wide[0] == narrow[0] > 0 + assert wide[1] - narrow[1] == rows * 6 * 3 * F // 2 * 2 + + +def test_shallow_models_keep_the_workspace_and_fewer_boundaries(): + rows = 23_878 + deep = tp_rank(topology=TP2)._checkpoint_memory_floor(((rows, True),)) + shallow = tp_rank(layers=24, topology=TP2) + shallow._mixer_top_gaps = (0, 1) + assert shallow._checkpoint_memory_floor(((rows, True),)) == ( + rows // 2 * 24 * H * 2, + deep[1], + ) + + +def _trace_profile(): + """The profile the TP2 trace learned for the 66,284-row wave.""" + return _MemoryProfile( + bytes_per_token=527_011.88, + packed_tokens=66_284, + logical_per_packed=1.0283, + retained_fraction=0.7674, + retained_compute_bytes_per_token=344_514.67, + caller_plans=1, + ) + + +def _checkpoint_split(children): + """_split_required_memory's checkpoint expression (no adapters or growth).""" + return int( + ( + sum(c.checkpoint_retained + c.checkpoint_input_gradient for c in children) + + max(c.checkpoint_workspace for c in children) + ) + * 1.1 + ) + + +def test_a_warm_split_prices_below_the_whole_wave_it_splits(): + """The production 66,284-row wave and its three requests as split children.""" + from art.trainer_rank._memory import _split_required_memory + + def warm(): + return _warm(tp_rank(topology=TP2), TP2, 4, _trace_profile()) + + whole = _required(warm(), ((66_284, True),), TP2, 4, 272_632, 68_158, 4) + children = [ + _required(warm(), ((rows, True),), TP2, 1, rows * 4, rows, 4) + for rows in (24_641, 22_962, 20_555) + ] + split = _split_required_memory(children) + # Every child's boundaries stay live; one child's workspace peaks at a time. + assert split < whole.required + # The split still pays for that workspace (about 4.9 GB). + assert split >= _checkpoint_split(children) + assert max(c.checkpoint_workspace for c in children) > 4.5e9 + # CUDA admission's placement path composes the same costs per unit (as + # _graph_memory_units does, GPU retention): the split still prices below. + from art.trainer_rank._memory_policy import ForwardMemoryCost, placement_cost + + def placed(costs_and_outputs): + units = [] + for cost, output in costs_and_outputs: + retained = max(output, cost.retained) + units.append( + ForwardMemoryCost( + peak_bytes=max(cost.required, retained), + retained_bytes=retained, + output_bytes=output, + ) + ) + return placement_cost( + units, backward_state="gpu", output_device="model" + ).gpu_required_bytes + + rows = (24_641, 22_962, 20_555) + assert placed(zip(children, (r * 4 for r in rows))) < placed([(whole, 272_632)]) + + +def test_a_split_whose_checkpoint_expression_binds(): + """A profile retaining less than the boundary shards: the checkpoint + branch, not the generic retained-plus-ephemeral sum, prices the split.""" + from art.trainer_rank._memory import _split_required_memory + + profile = _MemoryProfile( + bytes_per_token=300_000.0, + packed_tokens=24_641, + retained_compute_bytes_per_token=1.0, + retained_fraction=0.01, + caller_plans=1, + ) + children = [ + _required( + _warm(tp_rank(topology=TP2), TP2, 4, profile), + ((rows, True),), + TP2, + 1, + rows * 4, + rows, + 4, + ) + for rows in (24_641, 22_962, 20_555) + ] + assert _split_required_memory(children) == _checkpoint_split(children) + + +def test_the_largest_admissible_single_tp2_item_matches_the_measured_rate(): + """At the 43.8 GB production budget, a cold single item fits up to the rows + the measured ~522-527 KB/row peak plus the 1.1 margin allow.""" + budget = 43_835_395_277 + + def required(rows): + return _required( + tp_rank(topology=TP2), ((rows, True),), TP2, 1, rows * 4, rows, 4 + ).required + + assert required(61_000) < budget # about 32 GB at the measured rate + low, high = 1, 200_000 + while low < high: + middle = (low + high + 1) // 2 + low, high = (middle, high) if required(middle) <= budget else (low, middle - 1) + # Dropping the explicit workspace would admit about 115,000 rows. + assert 73_000 < low < 76_000 def test_rows_are_sharded_with_ceiling_and_only_gradient_groups_save_them(): @@ -99,40 +1062,44 @@ def test_rows_are_sharded_with_ceiling_and_only_gradient_groups_save_them(): mixed = r._checkpoint_memory_floor(((1024, True), (4096, False))) assert mixed[0] == 256 * LAYERS * H * 2 # No-grad groups keep today's four full rows (the static floor covers - # them); the gradient group adds its TP-padding roots' states. + # them), here wider than the gradient group's recompute workspace; the + # gradient group adds its TP-padding roots' states. + assert workspace(1024, 4) < 4 * 4096 * H * 2 assert mixed[1] == 4 * 4096 * H * 2 + 3 * SEGMENT @pytest.mark.parametrize( "case,kwargs", [ - ("tp2", dict(topology=(1, 2, 1, 1))), ("tp8", dict(topology=(1, 8, 1, 1))), ("cp2", dict(topology=(1, 4, 2, 1))), + ("tp2_cp2", dict(topology=(1, 2, 2, 1))), ("pp2", dict(topology=(1, 4, 1, 2))), ("no_sequence_parallel", dict(sequence_parallel=False)), ("sequence_parallel_at_tp1", dict(topology=(1, 1, 1, 1))), ("selective_recompute", dict()), ("moe", dict()), + ("tp2_moe", dict(topology=TP2)), ("moe_geometry", dict()), ("replicated_qkv", dict()), + ("tp2_replicated_qkv", dict(topology=TP2)), ("missing_attention_geometry", dict()), ("missing_conv_kernel", dict()), - ("shallow", dict(layers=48)), - ("wide_ffn", dict(ffn=4 * F)), + ("tp2_no_sequence_parallel", dict(topology=TP2, sequence_parallel=False)), ], ) def test_unproven_shapes_keep_todays_pricing(case, kwargs): r = tp_rank(**kwargs) if case == "selective_recompute": r.runtime.model[0].decoder.config.recompute_granularity = "selective" - if case == "moe": + if case in ("moe", "tp2_moe"): r._moe_layers = 1 # Geometry-only MoE, fewer KV groups than TP (a replicated global QKV) and # unreadable attention widths are all outside the traced shape. edits = { "moe_geometry": dict(moe_experts=8), "replicated_qkv": dict(num_query_groups=2), + "tp2_replicated_qkv": dict(num_query_groups=1), "missing_attention_geometry": dict(kv_channels=0), # Unread conv history would leave each segment's price short. "missing_conv_kernel": dict(gdn_conv_kernel=0), @@ -142,22 +1109,15 @@ def test_unproven_shapes_keep_todays_pricing(case, kwargs): assert r._checkpoint_memory_floor(((ROWS, True),)) == (0, 0) -def test_the_depth_bound_is_the_traced_workspace_at_these_widths(): - # Per gathered row, the repeated shards (layers x H / 4) must cover the - # SP gather (2H), the FC1 stage (6F/4), the wider mixer (GDN here), the - # norms (H), other workspace (H) and one gradient shard (H/4): about 49 - # layers at these widths. - assert tp_rank(layers=49)._checkpoint_memory_floor(((ROWS, True),))[0] > 0 - assert tp_rank(layers=48)._checkpoint_memory_floor(((ROWS, True),)) == (0, 0) - - def test_an_ungated_attention_only_model_is_bounded_by_its_attention_width(): r = tp_rank() r._gdn_layers = 0 r._attention_output_gate = False + r._mixer_top_gaps = (0, None) + attention = 5 * 24 * 256 + 3 * 4 * 256 assert r._checkpoint_memory_floor(((ROWS, True),)) == ( ROWS // 4 * LAYERS * H * 2, - 0, + workspace(ROWS, 4, mixers=((attention, 0),)), ) @@ -166,9 +1126,12 @@ def test_gdn_segment_states_are_priced_with_the_segments(): r = tp_rank() rows = 8192 cost = _required(r, group_rows=((rows, True),), gdn_segments=4096) - assert cost.checkpoint_workspace == (4096 + 3) * SEGMENT + COLD + states = (4096 + 3) * SEGMENT + assert cost.checkpoint_workspace == workspace(rows, 4) + states + COLD + retained = rows // 4 * LAYERS * H * 2 assert cost.required == int( - (OUTPUT + 2 * rows // 4 * LAYERS * H * 2 + (4096 + 3) * SEGMENT + COLD) * 1.1 + (OUTPUT + retained + workspace(rows, 4) + states + COLD + rows // 4 * H * 2) + * 1.1 ) @@ -177,7 +1140,11 @@ def test_tp_padding_roots_carry_their_own_states(): r = tp_rank() cost = _required(r, group_rows=((4, True),), gdn_segments=1) # Four roots' initial states alone: 4 x 12 value heads x 128 x 128 x fp32. - assert cost.checkpoint_workspace - COLD == 4 * SEGMENT > 4 * 12 * 128 * 128 * 4 + assert ( + cost.checkpoint_workspace - COLD - workspace(4, 4) + == 4 * SEGMENT + > 4 * 12 * 128 * 128 * 4 + ) assert cost.required > 4 * SEGMENT