From b430f5e1c337c2c9effb1cb3f66317aacafcc2f7 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Wed, 16 Sep 2026 09:58:22 +0000 Subject: [PATCH 01/18] Account for grouped checkpoint and shared-expert memory demand --- .github/workflows/prek.yml | 14 +- src/art/trainer_rank/_gdn_memory.py | 292 ++++++++ src/art/trainer_rank/_impl.py | 632 +++++++++++++++++- tests/unit/test_trainer_rank_active_memory.py | 40 ++ .../test_trainer_rank_checkpoint_memory.py | 296 ++++++++ tests/unit/test_trainer_rank_head_memory.py | 331 +++++++++ .../test_trainer_rank_ignored_mixed_head.py | 133 ++++ .../test_trainer_rank_mixed_head_memory.py | 219 ++++++ .../unit/test_trainer_rank_pending_memory.py | 265 ++++++++ .../unit/test_trainer_rank_planning_status.py | 5 +- tests/unit/test_trainer_rank_shared_memory.py | 231 +++++++ tests/unit/test_trainer_rank_weird_shapes.py | 9 +- 12 files changed, 2449 insertions(+), 18 deletions(-) create mode 100644 src/art/trainer_rank/_gdn_memory.py create mode 100644 tests/unit/test_trainer_rank_checkpoint_memory.py create mode 100644 tests/unit/test_trainer_rank_head_memory.py create mode 100644 tests/unit/test_trainer_rank_ignored_mixed_head.py create mode 100644 tests/unit/test_trainer_rank_mixed_head_memory.py create mode 100644 tests/unit/test_trainer_rank_pending_memory.py create mode 100644 tests/unit/test_trainer_rank_shared_memory.py diff --git a/.github/workflows/prek.yml b/.github/workflows/prek.yml index 97af92710..dfb27d681 100644 --- a/.github/workflows/prek.yml +++ b/.github/workflows/prek.yml @@ -225,6 +225,12 @@ jobs: tests/unit/test_prefix_tree_packing.py \ tests/unit/test_trainer_rank_validation.py \ tests/unit/test_trainer_rank_weird_shapes.py \ + tests/unit/test_trainer_rank_checkpoint_memory.py \ + tests/unit/test_trainer_rank_head_memory.py \ + tests/unit/test_trainer_rank_mixed_head_memory.py \ + tests/unit/test_trainer_rank_ignored_mixed_head.py \ + tests/unit/test_trainer_rank_pending_memory.py \ + tests/unit/test_trainer_rank_shared_memory.py \ tests/unit/test_trainer_rank_split.py \ tests/acceptance/trainer_rank_planner \ tests/integration/megatron/model_support/test_dispatcher_graph_retention.py \ @@ -248,4 +254,10 @@ jobs: --ignore=tests/unit/test_prefix_tree_grad_parity.py \ --ignore=tests/unit/test_prefix_tree_packing.py \ --ignore=tests/unit/test_trainer_rank_validation.py \ - --ignore=tests/unit/test_trainer_rank_weird_shapes.py + --ignore=tests/unit/test_trainer_rank_weird_shapes.py \ + --ignore=tests/unit/test_trainer_rank_checkpoint_memory.py \ + --ignore=tests/unit/test_trainer_rank_head_memory.py \ + --ignore=tests/unit/test_trainer_rank_mixed_head_memory.py \ + --ignore=tests/unit/test_trainer_rank_ignored_mixed_head.py \ + --ignore=tests/unit/test_trainer_rank_pending_memory.py \ + --ignore=tests/unit/test_trainer_rank_shared_memory.py diff --git a/src/art/trainer_rank/_gdn_memory.py b/src/art/trainer_rank/_gdn_memory.py new file mode 100644 index 000000000..a707f091e --- /dev/null +++ b/src/art/trainer_rank/_gdn_memory.py @@ -0,0 +1,292 @@ +"""Partial CP1 checkpoint pending-save accounting, not a backward upper bound. + +The bucket schedule mirrors the CP1 chunk-aligned GDN planner. It deliberately +uses CPU segment metadata, not input values or tensor allocation observations. +FLA source-version and generated-save limitations are documented with this +partial floor; other shared-expert saves, AOT saves and workspace remain unpriced. +""" + +from dataclasses import dataclass +from typing import Any, Sequence + + +@dataclass(frozen=True) +class Bucket: + # family, parent (-1 for a root), number of executed rows + columns: tuple[tuple[int, int, int], ...] + final: bool + + +def cp1_buckets(segments: Sequence[Any]) -> tuple[Bucket, ...]: + """Original 64-row root boundary / replayed-tail / child bucket schedule.""" + count = len(segments) + if not count: + return () + children: list[list[int]] = [[] for _ in segments] + depths: list[int] = [] + parents: list[int] = [] + lengths: list[int] = [] + ids = {s.group_id: i for i, s in enumerate(segments)} + if len(ids) != count: + raise ValueError("Duplicate GDN segment identity") + cursor = 0 + for i, s in enumerate(segments): + if any( + type(v) is not int + for v in (s.start, s.end, s.packed_start, s.group_id, s.parent_id) + ): + raise ValueError("Noninteger GDN segment metadata") + length = s.end - s.start + parent = -1 if s.parent_id == s.group_id else ids.get(s.parent_id, count) + if ( + length <= 0 + or s.start < 0 + or s.packed_start != cursor + or not -1 <= parent < i + ): + raise ValueError("Invalid ordered GDN segment geometry") + cursor += length + parents.append(parent) + lengths.append(length) + depths.append(0 if parent < 0 else depths[parent] + 1) + if parent >= 0: + children[parent].append(i) + boundary: list[tuple[int, int, int]] = [] + regular: list[tuple[int, int, int]] = [] + explicit: dict[int, list[tuple[int, int, int]]] = {} + for i, length in enumerate(lengths): + if parents[i] < 0: + if not children[i]: + regular.append((i, -1, length)) + elif length // 64: + boundary.append((i, -1, length // 64 * 64)) + for child in children[i]: + tail = length % 64 if parents[i] < 0 else 0 + parent = i if parents[i] >= 0 or length >= 64 else -1 + explicit.setdefault(depths[child], []).append( + (child, parent, tail + lengths[child]) + ) + buckets = [] + for columns in (boundary, regular): + if columns: + buckets.append( + Bucket( + tuple(sorted(columns, key=lambda c: (c[2], c[0]))), + any(children[c[0]] for c in columns), + ) + ) + for depth in sorted(explicit): + columns = tuple(explicit[depth]) + buckets.append(Bucket(columns, any(children[c[0]] for c in columns))) + return tuple(buckets) + + +@dataclass(frozen=True) +class Shape: + key_heads: int + value_heads: int + key_dim: int + value_dim: int + conv_width: int + output_lora_rank: int + moe_bytes_per_row: int + + def pending(self, packed_rows: int, buckets: tuple[Bucket, ...]) -> int: + """Known save-set envelope; final-state backing may alias initial saves. + + Charge each bucket's initial-state extent and each produced final-state + backing once. This intentionally overcounts aliases: a child view can + keep an entire parent batch live; request/root count cannot bound it. + No claim that every charged backing remains live at the MoE peak. + """ + hk, hv, dk, dv = self.key_heads, self.value_heads, self.key_dim, self.value_dim + conv = 2 * hk * dk + hv * dv + # FLA q/k/v, FP32 cumulative g, BF16 beta and A; convolution input; + # external q/k L2 reciprocal norms. All counts follow executed buckets. + per_bucket_row = (2 * hv * dk + hv * dv + hv + hv * 64 + conv) * 2 + hv * 4 * 3 + # Recurrent norm input/rstd and the ordinary trainable output LoRA's + # input plus rank temporary; these use final packed rows, not replay. + per_output_row = hv * dv * 2 + hv * 4 + if self.output_lora_rank: + per_output_row += hv * dv * 2 + self.output_lora_rank * 2 + state = hv * dk * dv * 4 + conv * (self.conv_width - 1) * 2 + executed = sum(c[2] for b in buckets for c in b.columns) + state_rows = sum(len(b.columns) * (1 + int(b.final)) for b in buckets) + return ( + executed * per_bucket_row + + packed_rows * per_output_row + + state_rows * state + ) + + +def model_shapes(rank: Any) -> tuple[int, tuple[Shape, ...]] | None: + """Conditional original-owner metadata; no model execution or CUDA read.""" + if not getattr(rank, "_gdn_layers", 0): + return None + import torch + + from art.trainer_rank._impl import _expert_parallel_shape, _language_model + + if ( + len(rank.runtime.model) != 1 + or rank._topology_key()[1:] != (1, 1, 1) + or _expert_parallel_shape(rank.runtime.provider) != (1, 1) + ): + return None + try: + decoder = _language_model(rank.runtime.model[0]).decoder + except (AttributeError, RuntimeError): + return None + if type(decoder).__name__ != "TransformerBlock": + return None + from megatron.core.ssm.gated_delta_net import GatedDeltaNet + from megatron.core.transformer.transformer_block import TransformerBlock + from transformer_engine.pytorch import RMSNorm + + from art.megatron.gdn.operator import _prefix_tree_forward + from art.megatron.lora import LoRA, SelfAttentionLinearProjLoRA + + config = decoder.config + expected = dict( + recompute_granularity="full", + recompute_method="uniform", + recompute_num_layers=1, + distribute_saved_activations=False, + sequence_parallel=False, + fp32_residual_connection=False, + cpu_offloading=False, + cuda_graph_impl="none", + ) + if ( + type(decoder) is not TransformerBlock + or decoder.training is not True + or not decoder.layers + or len(decoder.layers) != decoder.num_layers_per_pipeline_rank + or len(decoder.layers) != config.num_layers + or config.hidden_size != rank._hidden_size + or config.params_dtype is not torch.bfloat16 + or rank._param_dtype_size != 2 + or next(rank.runtime.model[0].parameters()).dtype is not torch.bfloat16 + or any( + type(getattr(config, k, None)) is not type(v) or getattr(config, k) != v + for k, v in expected.items() + ) + or getattr(config, "fp8", None) + or getattr(config, "fp4", None) + or any( + n in vars(decoder) + for n in ("forward", "_checkpointed_forward", "_get_layer") + ) + or decoder._forward_hooks + or decoder._forward_pre_hooks + ): + return None + # The constructor priced the original owners before installing dispatcher + # caches. Repricing their owned partials now would silently return zero. + # Like the existing static floor, this cache requires unchanged model, + # dtype and topology since construction; rebuilding the rank invalidates it. + moe = rank._moe_output_bytes_per_token + if type(moe) is not int or moe < 0 or (rank._moe_layers and not moe): + raise ValueError("Invalid constructor MoE coefficient for GDN pending floor") + shapes = [] + for layer in decoder.layers: + gdn = getattr(layer, "self_attention", None) + if type(gdn) is not GatedDeltaNet: + continue + if ( + getattr(gdn.forward, "__func__", None) is not _prefix_tree_forward + or gdn.use_qk_l2norm is not True + or gdn.tp_size != 1 + or gdn.sp_size != 1 + or gdn._forward_hooks + or gdn._forward_pre_hooks + ): + return None + dimensions = tuple( + getattr(gdn, n) + for n in ( + "num_key_heads", + "num_value_heads", + "key_head_dim", + "value_head_dim", + "conv_kernel_dim", + ) + ) + if ( + any(type(n) is not int or n <= 0 for n in dimensions) + or dimensions[1] % dimensions[0] + ): + return None + hk, hv, dk, dv, kernel = dimensions + conv = 2 * hk * dk + hv * dv + if ( + gdn.conv1d.weight.dtype is not torch.bfloat16 + or tuple(gdn.conv1d.weight.shape) != (conv, 1, kernel) + or type(gdn.out_norm) is not RMSNorm + or gdn.out_norm.weight.numel() != dv + or gdn.out_norm.weight.dtype is not torch.bfloat16 + or "forward" in vars(gdn.out_norm) + or gdn.out_norm._forward_hooks + or gdn.out_norm._forward_pre_hooks + ): + return None + out = gdn.out_proj + lora_rank = 0 + if type(out) is SelfAttentionLinearProjLoRA and type(out.lora) is LoRA: + lora = out.lora + if ( + "forward" in vars(out) + or "forward" in vars(lora) + or out._forward_hooks + or lora._forward_hooks + or out._forward_pre_hooks + or lora._forward_pre_hooks + ): + return None + a, b = lora.A_T, lora.B_T + if ( + a.ndim != 2 + or b.ndim != 2 + or a.dtype is not torch.bfloat16 + or b.dtype is not torch.bfloat16 + or a.shape[0] != dimensions[1] * dimensions[3] + or a.shape[1] != b.shape[0] + ): + return None + # Slots share these declared shapes; charging an inactive adapter + # conservatively adds a term, without inspecting/changing its slot. + lora_rank = int(a.shape[1]) + shapes.append(Shape(hk, hv, dk, dv, kernel, lora_rank, moe)) + return (len(decoder.layers), tuple(shapes)) if shapes else None + + +def plan_floor(rank: Any, plan: Any) -> tuple[int, int]: + """Boundary retention plus pending attention and the cached MoE maximum. + + Combining maxima from different layers may conservatively overcharge. + """ + gradients = [g for g in plan.groups if g.grad_enabled] + if not gradients: + return 0, 0 + model = model_shapes(rank) + if model is None: + return 0, 0 + layers, shapes = model + retained = 0 + workspace = 0 + for group in plan.groups: + rows = int(group.packed.tokens.numel()) + if not group.grad_enabled: + # Earlier gradient groups remain live during a later reference + # group. Only its existing MoE component enters this stage. + workspace = max(workspace, *(rows * s.moe_bytes_per_row for s in shapes)) + continue + buckets = cp1_buckets(group.packed.segments) + if sum(s.length for s in group.packed.segments) != rows: + raise ValueError("GDN packed rows disagree with segment geometry") + retained += rows * layers * rank._hidden_size * 2 + workspace = max( + workspace, + *(rows * s.moe_bytes_per_row + s.pending(rows, buckets) for s in shapes), + ) + return retained, workspace diff --git a/src/art/trainer_rank/_impl.py b/src/art/trainer_rank/_impl.py index de15bbdee..b9539877e 100644 --- a/src/art/trainer_rank/_impl.py +++ b/src/art/trainer_rank/_impl.py @@ -26,7 +26,7 @@ import struct import threading import time -from types import TracebackType +from types import MethodType, TracebackType from typing import ( TYPE_CHECKING, Any, @@ -53,6 +53,7 @@ _local_position_pairs, estimate_prefix_tree_packed_tokens, ) +from art.trainer_rank import _gdn_memory from art.trainer_rank._planner_cost import ( COEFFICIENT_VERSION_FALLBACK, ModelGeometry, @@ -1203,6 +1204,162 @@ def _configure_moe_dispatcher_caches(model: Sequence[torch.nn.Module]) -> None: ) +def _shared_expert_output_bytes_per_token(layer: torch.nn.Module) -> int: + """One supported shared return held across routed compute, not all saves. + + Gated backward can also save a distinct pre-gate result. This mode-neutral + component intentionally omits that separate term; compiled storage may alias. + """ + from megatron.core.extensions.transformer_engine import ( + TEColumnParallelLinear, + TELayerNormColumnParallelLinear, + TERowParallelLinear, + ) + from megatron.core.transformer.moe.shared_experts import SharedExpertMLP + + from art.megatron.lora import ( + LoRA, + SelfAttentionLinearProjLoRA, + SharedExpertsLinearFC1LoRA, + SharedExpertsLinearFC2LoRA, + ) + + shared = getattr(layer, "shared_experts", None) + if type(shared) is not SharedExpertMLP: + return 0 + config = getattr(layer, "config", None) + shared_config = getattr(shared, "config", None) + expected = { + "params_dtype": torch.bfloat16, + "moe_shared_expert_overlap": False, + "sequence_parallel": False, + "fp32_residual_connection": False, + "add_bias_linear": False, + "use_te_activation_func": False, + "bias_activation_fusion": False, + "gated_linear_unit": True, + "cuda_graph_impl": "none", + } + if ( + getattr(layer, "use_shared_expert", None) is not True + or getattr(layer, "shared_expert_overlap", None) is not False + or getattr(layer, "shared_experts_recompute", None) is not False + or getattr(layer, "moe_layer_recompute", None) is not False + or getattr(layer, "fwd_execution_map", None) + != ["route", "expert_compute", "postprocess"] + or any( + name in vars(layer) + for name in ( + "shared_experts_compute", + "route", + "preprocess", + "dispatch", + "routed_experts_compute", + "combine", + "postprocess", + ) + ) + or any( + type(getattr(c, name, None)) is not type(value) or getattr(c, name) != value + for c in (config, shared_config) + for name, value in expected.items() + ) + or any( + getattr(c, name, None) + for c in (config, shared_config) + for name in ("fp8", "fp4", "moe_latent_size") + ) + or any( + (type(getattr(config, name, None)) is not int or getattr(config, name) != 1) + for name in ( + "tensor_model_parallel_size", + "context_parallel_size", + "pipeline_model_parallel_size", + "expert_model_parallel_size", + "expert_tensor_parallel_size", + ) + ) + ): + return 0 + hidden = getattr(config, "hidden_size", None) + width = getattr(config, "moe_shared_expert_intermediate_size", None) + if ( + type(hidden) is not int + or hidden <= 0 + or type(width) is not int + or width <= 0 + or getattr(shared_config, "hidden_size", None) != hidden + or getattr(shared_config, "ffn_hidden_size", None) != width + or getattr(shared_config, "moe_shared_expert_intermediate_size", None) != width + or getattr(shared_config, "activation_func", None) + is not torch.nn.functional.silu + or getattr(shared, "activation_func", None) is not torch.nn.functional.silu + or type(getattr(shared, "use_shared_expert_gate", None)) is not bool + ): + return 0 + fc1, fc2 = getattr(shared, "linear_fc1", None), getattr(shared, "linear_fc2", None) + row = getattr(fc2, "row_parallel_lora", None) + lora = getattr(row, "lora", None) + base1, base2 = getattr(fc1, "linear_fc1", None), getattr(row, "linear_proj", None) + sites = ( + (shared, SharedExpertMLP), + (fc1, SharedExpertsLinearFC1LoRA), + (fc2, SharedExpertsLinearFC2LoRA), + (row, SelfAttentionLinearProjLoRA), + (lora, LoRA), + (base2, TERowParallelLinear), + (getattr(fc1, "gate_lora", None), LoRA), + (getattr(fc1, "up_lora", None), LoRA), + ) + if ( + type(base1) not in (TEColumnParallelLinear, TELayerNormColumnParallelLinear) + or any(type(site) is not cls for site, cls in sites) + or any( + "forward" in vars(site) + or cast(Any, site)._forward_hooks + or cast(Any, site)._forward_pre_hooks + for site in (base1, *(site for site, _ in sites)) + ) + or getattr(fc1, "non_gated", None) is not False + or getattr(fc1, "out_features", None) != 2 * width + or getattr(getattr(row, "provider", None), "tensor_model_parallel_size", None) + != 1 + or getattr(getattr(row, "provider", None), "sequence_parallel", None) + is not False + ): + return 0 + weights = ( + (getattr(base1, "weight", None), (2 * width, hidden)), + (getattr(base2, "weight", None), (hidden, width)), + ) + for adapter, inputs, outputs in ( + (cast(Any, fc1).gate_lora, hidden, width), + (cast(Any, fc1).up_lora, hidden, width), + (lora, width, hidden), + ): + a, b = getattr(adapter, "A_T", None), getattr(adapter, "B_T", None) + if ( + not isinstance(a, torch.Tensor) + or not isinstance(b, torch.Tensor) + or a.ndim != 2 + or b.ndim != 2 + or a.shape[1] <= 0 + or a.shape[1] != b.shape[0] + ): + return 0 + weights += ((a, (inputs, a.shape[1])), (b, (a.shape[1], outputs))) + if shared.use_shared_expert_gate: + weights += ((getattr(shared, "gate_weight", None), (1, hidden)),) + if any( + not isinstance(weight, torch.Tensor) + or weight.dtype is not torch.bfloat16 + or tuple(weight.shape) != shape + for weight, shape in weights + ): + return 0 + return hidden * 2 + + def _moe_output_bytes_per_token( model: Sequence[torch.nn.Module], shape: ParallelShape ) -> int: @@ -1312,7 +1469,8 @@ def _moe_output_bytes_per_token( features += 2 * fc2.out_features + fc1.out_features coefficient = max( coefficient, - config.moe_router_topk * features * weights.element_size(), + config.moe_router_topk * features * weights.element_size() + + _shared_expert_output_bytes_per_token(layer), ) return coefficient @@ -2816,13 +2974,27 @@ def _split_chunk_lower_cost( ) packed_tokens = 0 unshared_packed_tokens = 0 - for _, group_indices in groups: + head_workspace_bytes = 0 + group_rows: list[tuple[int, bool]] = [] + for (_slot, grad_enabled), group_indices in groups: estimated = estimate_prefix_tree_packed_tokens( (rows[index] for index in group_indices), max_depth=len(group_indices), ) assert estimated is not None # rows are CPU copies - packed_tokens += self._physical_tokens(estimated) + physical_rows = self._physical_tokens(estimated) + packed_tokens += physical_rows + group_rows.append((physical_rows, grad_enabled)) + head_requests = tuple(requests[index] for index in group_indices) + head_workspace_bytes = max( + head_workspace_bytes, + self._group_head_workspace_bytes( + self._head_projection_rows(head_requests, lower_bound=True), + head_requests, + grad_enabled=grad_enabled, + lower_bound=True, + ), + ) unshared_packed_tokens += self._physical_tokens( sum(int(rows[index].numel()) for index in group_indices) ) @@ -2838,6 +3010,8 @@ def _split_chunk_lower_cost( output_bytes=output_bytes, signature=signature, logical_tokens=logical_tokens, + group_rows=tuple(group_rows), + head_workspace_bytes=head_workspace_bytes, ) profile = self._memory_profiles.get(signature) if ( @@ -2857,19 +3031,356 @@ def _split_chunk_lower_cost( ): # A larger layout may trust retained compute where full sharing # cannot. Its full-required retention is not a pruning lower bound. - # Charge only outputs here; exact plan costs keep both trust guards. + # Keep outputs and the independent source retention floor; exact + # plan costs keep both trust guards. return _SubforwardCost( required=cost.required, - retained=min(cost.retained, int(output_bytes * _MEMORY_SAFETY_FACTOR)), + retained=min( + cost.retained, + int( + ( + output_bytes + + self._checkpoint_memory_floor(tuple(group_rows))[0] + ) + * _MEMORY_SAFETY_FACTOR + ), + ), ) return cost + def _head_workspace_bytes(self, rows: int) -> int: + """One dense BF16 head tensor, not complete statistics/backward memory.""" + if ( + rows <= 0 + or self._padded_vocab_size is None + or len(self.runtime.model) != 1 + or self._topology_key()[1:] != (1, 1, 1) + ): + return 0 + try: + model = _language_model(self.runtime.model[0]) + except (AttributeError, RuntimeError): + return 0 + try: + from megatron.core.tensor_parallel.layers import ColumnParallelLinear + except ModuleNotFoundError as error: + if error.name != "megatron": + raise + return 0 + + head = getattr(model, "output_layer", None) + if head is None or type(head) is not ColumnParallelLinear: + return 0 + weight = head.weight + if ( + weight is None + and getattr(model, "share_embeddings_and_output_weights", False) is True + ): + weight = getattr( + getattr(getattr(model, "embedding", None), "word_embeddings", None), + "weight", + None, + ) + if weight is None: + return 0 + config = getattr(model, "config", None) + if ( + type(weight) not in (torch.Tensor, torch.nn.Parameter) + or weight.dtype is not torch.bfloat16 + or tuple(weight.shape) != (self._padded_vocab_size, self._hidden_size) + or head.output_size_per_partition != self._padded_vocab_size + or head.output_size != self._padded_vocab_size + or head.input_size != self._hidden_size + or getattr(config, "params_dtype", None) is not torch.bfloat16 + or getattr(config, "fp32_residual_connection", None) is not False + or getattr(config, "fp8", None) + or getattr(config, "fp4", None) + or "forward" in vars(head) + or "_forward_impl" in vars(head) + or getattr(head, "_forward_hooks", None) + or getattr(head, "_forward_pre_hooks", None) + or any( + name in vars(self) + for name in ( + "_project_head", + "_project_vocab_parallel", + "_local_head_stats", + "_local_logits_from_hidden_rows", + ) + ) + ): + return 0 + return min(rows, _HEAD_CHUNK_TOKENS) * int(self._padded_vocab_size) * 2 + + def _head_projection_rows( + self, + requests: Sequence[AnyForwardInput], + *, + positions: Sequence[torch.Tensor] | None = None, + lower_bound: bool = False, + ) -> int: + """Per-group logical bounds or exact packed union; no device-label read. + + 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. + Device-label validity is unknown: use all rows for capacity, zero only + for the rejection lower bound; never copy labels from the device here. + """ + if not self._head_workspace_bytes(1): + return 0 + if positions is not None and any(row.device.type != "cpu" for row in positions): + positions = None # Capacity bound without reading device positions. + counts: list[int] = [] + projected: set[int] = set() + for index, request in enumerate(requests): + offsets = None + if request.logits or request.top_k is not None: + count = int(request.input_tokens.numel()) + elif request.target_tokens is not None: + count = int(request.input_tokens.numel()) + if request.target_tokens.device.type != "cpu": + if lower_bound: + count = 0 + else: + labels = request.target_tokens.to(dtype=torch.long) + valid = (labels != -100).reshape(count, -1).any(dim=1) + offsets = torch.nonzero(valid, as_tuple=False).reshape(-1) + count = int(offsets.numel()) + else: + continue + if positions is None: + counts.append(count) + else: + row = positions[index] + if offsets is not None: + row = row.index_select(0, offsets) + for position in row.tolist(): + projected.add(int(position)) + if len(projected) >= _HEAD_CHUNK_TOKENS: + return _HEAD_CHUNK_TOKENS + return min( + _HEAD_CHUNK_TOKENS, + (max(counts, default=0) if lower_bound else sum(counts)) + if positions is None + else len(projected), + ) + + def _head_target_chunk_rows( + self, + requests: Sequence[AnyForwardInput], + *, + positions: Sequence[torch.Tensor] | None = None, + lower_bound: bool = False, + ) -> int: + """Largest projected chunk reached by labelled rows, or row bounds. + + Mixed outputs do not remove target backward. Its dense indexing result + spans the whole projected chunk, including rows requested only as logits + or top-k. Ignored labels still execute backward if another output + projects their rows. Without a layout, valid rows give a rejection lower + bound; possible overlap with any labelled request gives capacity. + """ + targets = tuple( + replace(request, logits=False, top_k=None) for request in requests + ) + if positions is None or any(row.device.type != "cpu" for row in positions): + if lower_bound: + return self._head_projection_rows(targets, lower_bound=True) + return ( + self._head_projection_rows(requests) + if any( + request.target_tokens is not None and request.input_tokens.numel() + for request in requests + ) + else 0 + ) + projected: set[int] = set() + labelled: set[int] = set() + for request, row in zip(requests, positions, strict=True): + target_row = row[:0] + if request.target_tokens is not None and int(row.numel()): + labelled.update(row.tolist()) + labels = request.target_tokens + if labels.device.type == "cpu": + valid = ( + (labels.to(dtype=torch.long) != -100) + .reshape(len(row), -1) + .any(dim=1) + ) + target_row = row.index_select( + 0, torch.nonzero(valid, as_tuple=False).reshape(-1) + ) + elif not lower_bound: + target_row = row + projected.update( + ( + row if request.logits or request.top_k is not None else target_row + ).tolist() + ) + targeted = labelled & projected + if not targeted: + return 0 + first_target = min(targeted) + first_index = sum(position < first_target for position in projected) + chunk_start = first_index // _HEAD_CHUNK_TOKENS * _HEAD_CHUNK_TOKENS + return min(_HEAD_CHUNK_TOKENS, len(projected) - chunk_start) + + def _group_head_workspace_bytes( + self, + rows: int, + requests: Sequence[AnyForwardInput], + *, + grad_enabled: bool, + positions: Sequence[torch.Tensor] | None = None, + lower_bound: bool = False, + ) -> int: + """One logits buffer, or logits + both dense target-backward gradients. + + The supported head path overlaps indexing and statistics gradients + with recomputed logits; cold library workspaces remain outside this + component. Pair each group's mode with its own projected rows. + """ + dense = self._head_workspace_bytes(rows) + if ( + not dense + or not grad_enabled + or not any(request.target_tokens is not None for request in requests) + ): + return dense + from megatron.core.models.common.language_module.language_module import ( + LanguageModule, + ) + + model = _language_model(self.runtime.model[0]) + scale = getattr(model, "_scale_logits", None) + if ( + type(scale) is MethodType + and scale.__self__ is model + and scale.__func__ is LanguageModule._scale_logits + and getattr(model.config, "use_mup", None) is False + ): + # IndexBackward's dense result overlaps saved logits and grad_logits. + # The FP32 fallback already exceeds this three-buffer component. + target_dense = ( + 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 + ) + return max(dense, 3 * target_dense) + return dense + + def _plan_head_workspace_bytes(self, plan: _FlatForwardPlan) -> int: + peak = 0 + for group in plan.groups: + requests = tuple(item.request for item in group.items) + peak = max( + peak, + self._group_head_workspace_bytes( + self._head_projection_rows( + requests, positions=group.packed.positions_by_sequence + ), + requests, + grad_enabled=group.grad_enabled, + positions=group.packed.positions_by_sequence, + ), + ) + return peak + + def _plan_group_rows(self, plan: _FlatForwardPlan) -> tuple[tuple[int, bool], ...]: + return tuple( + ( + self._physical_tokens(int(group.packed.tokens.numel())), + group.grad_enabled, + ) + for group in plan.groups + ) + + def _checkpoint_memory_floor( + self, group_rows: tuple[tuple[int, bool], ...] + ) -> tuple[int, int]: + """Conservative saved-boundary charge and one disjoint MoE workspace. + + Count actual local full/uniform/1 boundaries, including aliases, rather + than claiming measured distinct storage. Only this call's new groups + enter the term; already-live graphs remain in the availability baseline. + This is not a bound for custom preprocessing, attention, or all backward. + """ + gradient_rows = sum(rows for rows, grad in group_rows if grad) + if not gradient_rows or len(self.runtime.model) != 1: + return 0, 0 + try: + decoder = _language_model(self.runtime.model[0]).decoder + except (AttributeError, RuntimeError): + return 0, 0 + try: + from megatron.core.transformer.transformer_block import TransformerBlock + except ModuleNotFoundError as error: + if error.name != "megatron": + raise + return 0, 0 + + if type(decoder) is not TransformerBlock: + return 0, 0 + config = decoder.config + layers = len(decoder.layers) + expected = { + "recompute_granularity": "full", + "recompute_method": "uniform", + "recompute_num_layers": 1, + "distribute_saved_activations": False, + "sequence_parallel": False, + "fp32_residual_connection": False, + "cpu_offloading": False, + "cuda_graph_impl": "none", + } + if ( + decoder.training is not True + or layers <= 0 + or layers != decoder.num_layers_per_pipeline_rank + or layers != config.num_layers + or config.hidden_size != self._hidden_size + or config.params_dtype is not torch.bfloat16 + or self._param_dtype_size != 2 + or next(self.runtime.model[0].parameters()).dtype is not torch.bfloat16 + or self._topology_key()[1:] != (1, 1, 1) + or _expert_parallel_shape(self.runtime.provider) != (1, 1) + or any( + type(getattr(config, name, None)) is not type(value) + or getattr(config, name) != value + for name, value in expected.items() + ) + or getattr(config, "fp8", None) + or getattr(config, "fp4", None) + or any( + name in vars(decoder) + for name in ("forward", "_checkpointed_forward", "_get_layer") + ) + or getattr(decoder, "_forward_hooks", None) + or getattr(decoder, "_forward_pre_hooks", None) + ): + return 0, 0 + retained = gradient_rows * layers * self._hidden_size * 2 + workspace = ( + max(rows for rows, _grad in group_rows) * self._moe_output_bytes_per_token + ) + return retained, workspace + def _plan_cost(self, plan: _FlatForwardPlan) -> _SubforwardCost: return self._subforward_cost( packed_tokens=plan.packed_tokens, output_bytes=plan.output_bytes, signature=plan.signature, logical_tokens=plan.active_logical_tokens, + group_rows=self._plan_group_rows(plan), + head_workspace_bytes=self._plan_head_workspace_bytes(plan), + checkpoint_floor=_gdn_memory.plan_floor(self, plan), ) def _subforward_cost( @@ -2879,12 +3390,18 @@ def _subforward_cost( output_bytes: int, signature: _MemorySignature, logical_tokens: int, + group_rows: tuple[tuple[int, bool], ...] = (), + head_workspace_bytes: int = 0, + checkpoint_floor: tuple[int, int] = (0, 0), ) -> _SubforwardCost: required = self._estimate_required_memory_bytes_from_values( packed_tokens=packed_tokens, output_bytes=output_bytes, signature=signature, logical_tokens=logical_tokens, + group_rows=group_rows, + head_workspace_bytes=head_workspace_bytes, + checkpoint_floor=checkpoint_floor, ) retained = self._retained_memory_bytes( signature, @@ -2892,6 +3409,9 @@ def _subforward_cost( logical_tokens=logical_tokens, output_bytes=output_bytes, required=required, + checkpoint_retained_bytes=max( + self._checkpoint_memory_floor(group_rows)[0], checkpoint_floor[0] + ), ) return _SubforwardCost(required=required, retained=retained) @@ -2903,6 +3423,7 @@ def _retained_memory_bytes( logical_tokens: int, output_bytes: int, required: int, + checkpoint_retained_bytes: int = 0, ) -> int: """Forward-retained bytes, independent of a later backward peak. @@ -2918,8 +3439,10 @@ 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 + profile.retained_compute_bytes_per_token * max( - packed_tokens, logical_tokens / profile.logical_per_packed + retained = output_bytes + max( + checkpoint_retained_bytes, + profile.retained_compute_bytes_per_token + * max(packed_tokens, logical_tokens / profile.logical_per_packed), ) return min(required, int(retained * _MEMORY_SAFETY_FACTOR)) @@ -3870,6 +4393,8 @@ def priced( packed_tokens: int, output_bytes: int, signature: _MemorySignature, + group_rows: tuple[tuple[int, bool], ...], + head_workspace_bytes: int, ) -> tuple[_MemoryCheck, int, int, _MemorySignature]: with self._planning_status(True): required = self._estimate_required_memory_bytes_from_values( @@ -3877,6 +4402,8 @@ def priced( output_bytes=output_bytes, signature=signature, logical_tokens=logical_tokens, + group_rows=group_rows, + head_workspace_bytes=head_workspace_bytes, ) return ( self._memory_check_required(required, sync_across_dp=True), @@ -4532,7 +5059,7 @@ def _estimate_flat_forward( exact: bool = False, memory_minimal: bool = False, sync_planning_errors: bool = False, - ) -> tuple[int, int, _MemorySignature] | None: + ) -> tuple[int, int, _MemorySignature, tuple[tuple[int, bool], ...], int] | None: """Estimate packed tokens for width probing. Cheap mode (``exact=False``) is one O(tokens) CPU walk of the packing @@ -4553,10 +5080,22 @@ def _estimate_flat_forward( checkpoint=checkpoint, ensure_slots=not sync_planning_errors, ) + if ( + any(mode for (_, mode), _ in groups) + and _gdn_memory.model_shapes(self) is not None + ): + # Pending saves require the actual bucket/replayed-tail geometry. + # Existing unavailable handling materializes before admission. + return None packed_tokens = 0 + head_workspace_bytes = 0 + group_rows: list[tuple[int, bool]] = [] for (_slot, grad_enabled), group_indices in groups: + head_requests = tuple(requests[index] for index in group_indices) + lower = self._head_projection_rows(head_requests, lower_bound=True) + upper = self._head_projection_rows(head_requests) if exact: - _, layout = self._select_group_layout( + tree, layout = self._select_group_layout( tuple( requests[index] .input_tokens.reshape(-1) @@ -4566,7 +5105,51 @@ def _estimate_flat_forward( memory_minimal=memory_minimal, grad_enabled=grad_enabled, ) - packed_tokens += self._physical_tokens(layout.packed_tokens) + physical_rows = self._physical_tokens(layout.packed_tokens) + packed_tokens += physical_rows + group_rows.append((physical_rows, grad_enabled)) + projected = upper + positions = None + mixed_targets = ( + grad_enabled + and any( + request.target_tokens is not None + for request in head_requests + ) + and any( + request.logits or request.top_k is not None + 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) + ): + packed = materialize_prefix_tree_layout( + tuple( + request.input_tokens.reshape(-1).to(dtype=torch.long) + for request in head_requests + ), + tree, + layout, + verify_shared_tokens=False, + ) + projected = self._head_projection_rows( + head_requests, positions=packed.positions_by_sequence + ) + positions = packed.positions_by_sequence + head_workspace_bytes = max( + head_workspace_bytes, + self._group_head_workspace_bytes( + projected, + head_requests, + grad_enabled=grad_enabled, + positions=positions, + ), + ) continue # Radix depth is bounded by the number of rows, so ``len(group)`` # is an unlimited-sharing depth for this group; it is a bound for @@ -4580,7 +5163,18 @@ def _estimate_flat_forward( ) if group_packed_tokens is None: return None - packed_tokens += self._physical_tokens(group_packed_tokens) + physical_rows = self._physical_tokens(group_packed_tokens) + packed_tokens += physical_rows + group_rows.append((physical_rows, grad_enabled)) + head_workspace_bytes = max( + head_workspace_bytes, + self._group_head_workspace_bytes( + lower if memory_minimal else upper, + head_requests, + grad_enabled=grad_enabled, + lower_bound=memory_minimal, + ), + ) return ( packed_tokens, @@ -4590,6 +5184,8 @@ def _estimate_flat_forward( slot_group_count=len(groups), grad_modes=tuple(mode for (_, mode), _ in groups), ), + tuple(group_rows), + head_workspace_bytes, ) def _ensure_checkpoint_slots_for( @@ -5018,6 +5614,9 @@ def _memory_check( output_bytes=forward.output_bytes, signature=forward.signature, logical_tokens=forward.active_logical_tokens, + group_rows=self._plan_group_rows(forward), + head_workspace_bytes=self._plan_head_workspace_bytes(forward), + checkpoint_floor=_gdn_memory.plan_floor(self, forward), ) return self._memory_check_required(required, sync_across_dp=sync_across_dp) @@ -5079,6 +5678,9 @@ def _estimate_required_memory_bytes_from_values( output_bytes: int, signature: _MemorySignature, logical_tokens: int | None = None, + group_rows: tuple[tuple[int, bool], ...] = (), + head_workspace_bytes: int = 0, + checkpoint_floor: tuple[int, int] = (0, 0), ) -> int: if packed_tokens <= 0: return output_bytes @@ -5115,6 +5717,12 @@ def _estimate_required_memory_bytes_from_values( static_compute, int(profiled.bytes_per_token * profiled_tokens), ) + retained, workspace = self._checkpoint_memory_floor(group_rows) + compute = max( + compute, + max(retained, checkpoint_floor[0]) + + max(workspace, head_workspace_bytes, checkpoint_floor[1]), + ) return int((output_bytes + compute) * _MEMORY_SAFETY_FACTOR) def _available_memory_bytes(self) -> int: diff --git a/tests/unit/test_trainer_rank_active_memory.py b/tests/unit/test_trainer_rank_active_memory.py index 74ced9d63..6d909b279 100644 --- a/tests/unit/test_trainer_rank_active_memory.py +++ b/tests/unit/test_trainer_rank_active_memory.py @@ -1,5 +1,6 @@ """CPU admission contracts; injected observations are not GPU peak measurements.""" +import builtins from dataclasses import replace from types import SimpleNamespace from typing import Any, cast @@ -270,3 +271,42 @@ def test_empirical_estimate_survives_packed_trust_boundary(logical_ratio): assert not rank._all_ranks_have_memory_profile( packed_tokens=800, signature=observed.signature ) + + +@pytest.mark.parametrize( + "method,argument,fallback", + [ + ("_head_workspace_bytes", 8, 0), + ("_checkpoint_memory_floor", ((8, True),), (0, 0)), + ], +) +@pytest.mark.parametrize( + "error,unavailable", + [ + (ModuleNotFoundError("absent package", name="megatron"), True), + (ModuleNotFoundError("missing dependency", name="transformer_engine"), False), + (ModuleNotFoundError("partial installation", name="megatron.core"), False), + (ModuleNotFoundError("unspecified missing module"), False), + (ImportError("missing imported class"), False), + (RuntimeError("module initialization failed"), False), + ], + ids=["absent", "transitive", "partial", "unspecified", "class", "runtime"], +) +def test_optional_megatron_memory_guards( + monkeypatch, method, argument, fallback, error, unavailable +): + rank = _rank() + original_import = builtins.__import__ + + def importing(name, *args, **kwargs): + if name.partition(".")[0] == "megatron": + raise error + return original_import(name, *args, **kwargs) + + monkeypatch.setattr(builtins, "__import__", importing) + if unavailable: + assert getattr(rank, method)(argument) == fallback + else: + with pytest.raises(type(error)) as caught: + getattr(rank, method)(argument) + assert caught.value is error diff --git a/tests/unit/test_trainer_rank_checkpoint_memory.py b/tests/unit/test_trainer_rank_checkpoint_memory.py new file mode 100644 index 000000000..69de26475 --- /dev/null +++ b/tests/unit/test_trainer_rank_checkpoint_memory.py @@ -0,0 +1,296 @@ +"""Conditional checkpoint accounting; scalar/CPU evidence, not a CUDA bound.""" + +from dataclasses import replace +from types import SimpleNamespace +from typing import Any, cast + +import pytest +import torch + +from art.trainer_rank import ForwardInput, TrainerRank +from art.trainer_rank._impl import Unset, _MemoryProfile + + +def rank(): + from megatron.core.transformer.transformer_block import TransformerBlock + + block = TransformerBlock.__new__(TransformerBlock) + torch.nn.Module.__init__(block) + block.config = SimpleNamespace( + hidden_size=2048, + num_layers=40, + padded_vocab_size=32, + params_dtype=torch.bfloat16, + recompute_granularity="full", + recompute_method="uniform", + recompute_num_layers=1, + distribute_saved_activations=False, + sequence_parallel=False, + fp32_residual_connection=False, + cpu_offloading=False, + cuda_graph_impl="none", + fp8=None, + fp4=None, + ) + block.layers = torch.nn.ModuleList( + [torch.nn.Linear(1, 1).bfloat16() for _ in range(40)] + ) + block.num_layers_per_pipeline_rank = 40 + model: Any = torch.nn.Module() + model.config = block.config + model.decoder = block + model._preprocess = lambda: None + result = TrainerRank( + cast( + Any, + SimpleNamespace( + model=[model], + optimizer=None, + provider=SimpleNamespace(hidden_size=2048, num_layers=40), + model_support_handler=SimpleNamespace(build_gdn_execution_spec=False), + ), + ) + ) + result._moe_output_bytes_per_token = 188416 + return result + + +def requests(grad=1024, reference=15360): + return [ + ForwardInput( + input_tokens=torch.arange(grad), hidden_states=True, no_grad=False + ), + ForwardInput( + input_tokens=torch.arange(reference), hidden_states=True, no_grad=True + ), + ] + + +def price(r, values): + n, out, signature, groups, head_workspace_bytes = values + return r._subforward_cost( + packed_tokens=n, + output_bytes=out, + signature=signature, + logical_tokens=n, + group_rows=groups, + head_workspace_bytes=head_workspace_bytes, + ) + + +def test_same_old_signature_different_gradient_rows(): + r = rank() + a = r._estimate_flat_forward(requests()) + b = r._estimate_flat_forward(requests(15360, 1024)) + assert a[:3] == b[:3] + assert a[3] == ((1024, True), (15360, False)) + assert b[3] == ((15360, True), (1024, False)) + assert ( + r._checkpoint_memory_floor(a[3])[0] * 15 == r._checkpoint_memory_floor(b[3])[0] + ) + assert price(r, b).required > price(r, a).required + + +def test_required_and_learned_retained_use_max_not_sum(): + r = rank() + values = r._estimate_flat_forward(requests()) + n, out, sig, groups, head_workspace_bytes = values + retained, work = r._checkpoint_memory_floor(groups) + r._memory_profiles[sig] = _MemoryProfile( + bytes_per_token=1, + packed_tokens=n, + logical_per_packed=1, + retained_compute_bytes_per_token=1, + ) + cost = price(r, values) + old = max(n * 2048 * 2 * 14, n * 188416) + assert cost.required == int((out + max(old, retained + work)) * 1.1) + assert cost.retained == int((out + retained) * 1.1) + r._memory_profiles[sig] = replace( + r._memory_profiles[sig], + bytes_per_token=1_000_000, + retained_compute_bytes_per_token=500_000, + ) + cost = price(r, values) + assert cost.required == int((out + n * 1_000_000) * 1.1) + assert cost.retained == int((out + n * 500_000) * 1.1) + + +def test_materialized_and_lower_bound_keep_group_association(): + r = rank() + req = requests(17, 19) + exact = r._estimate_flat_forward(req, exact=True) + plan = r._plan_flat_forward(req) + assert exact[3] == r._plan_group_rows(plan) + assert r._memory_check(plan).estimated_required_bytes == price(r, exact).required + rows = tuple(x.input_tokens for x in req) + assert r._split_chunk_lower_cost(req, rows, checkpoint=Unset) == price(r, exact) + + +def test_per_group_padding_precedes_gradient_filter(): + r = rank() + r._physical_tokens = lambda n: n + (-n % 8) + values = r._estimate_flat_forward(requests(9, 17)) + assert values[0] == 40 and values[3] == ((16, True), (24, False)) + assert r._checkpoint_memory_floor(values[3]) == (16 * 40 * 4096, 24 * 188416) + + +def test_no_grad_empty_and_unsupported_keep_prior_estimate(): + r = rank() + assert r._checkpoint_memory_floor(()) == (0, 0) + assert r._checkpoint_memory_floor(((8192, False),)) == (0, 0) + values = r._estimate_flat_forward(requests()) + baseline = price(r, values).required + r.runtime.model[0].decoder.eval() + n, out, sig, groups, head_workspace_bytes = values + old = r._estimate_required_memory_bytes_from_values( + packed_tokens=n, output_bytes=out, signature=sig + ) + assert price(r, values).required == old <= baseline + + +@pytest.mark.parametrize( + "field,value", + [ + ("recompute_granularity", "selective"), + ("recompute_method", "block"), + ("recompute_num_layers", 2), + ("distribute_saved_activations", True), + ("sequence_parallel", True), + ("fp32_residual_connection", True), + ("cpu_offloading", True), + ("cuda_graph_impl", "local"), + ("params_dtype", torch.float32), + ("fp8", "hybrid"), + ("fp4", True), + ("num_layers", 39), + ("hidden_size", 1024), + ], +) +def test_actual_config_revalidated(field, value): + r = rank() + assert r._checkpoint_memory_floor(((10, True),))[0] > 0 + setattr(r.runtime.model[0].decoder.config, field, value) + assert r._checkpoint_memory_floor(((10, True),)) == (0, 0) + + +@pytest.mark.parametrize("axis", [1, 2, 3]) +def test_topology_revalidated(axis): + r = rank() + topology = [1, 1, 1, 1] + topology[axis] = 2 + r._topology_key = lambda: tuple(topology) + assert r._checkpoint_memory_floor(((10, True),)) == (0, 0) + + +def test_dp_empty_and_local_count(): + r = rank() + r._topology_key = lambda: (3, 1, 1, 1) + assert r._checkpoint_memory_floor(()) == (0, 0) + block = r.runtime.model[0].decoder + block.layers = torch.nn.ModuleList(list(block.layers[:4])) + block.num_layers_per_pipeline_rank = 4 + block.config.num_layers = 4 + assert r._checkpoint_memory_floor(((10, True),))[0] == 10 * 4 * 2048 * 2 + block.num_layers_per_pipeline_rank = 3 + assert r._checkpoint_memory_floor(((10, True),)) == (0, 0) + + +def test_prior_live_graphs_are_not_an_added_term(): + r = rank() + values = r._estimate_flat_forward(requests()) + cost = price(r, values) + r._available_memory_bytes = lambda: cost.required - 1 + assert not r._memory_check_required(cost.required).fits + # New-call cost is unchanged; live memory affects only existing availability. + assert price(r, values) == cost + + +def test_existing_api_rejects_budget_below_conditional_checkpoint_term(): + # This same test reaches the original materialized-plan API on the base. + r = rank() + req = [ + ForwardInput( + input_tokens=torch.arange(128), + target_tokens=torch.arange(128), + no_grad=False, + ) + ] + plan = r._plan_flat_forward(req) + old_component = int((plan.output_bytes + 128 * 188416) * 1.1) + r._available_memory_bytes = lambda: old_component + 1 + assert not r._memory_check(plan).fits + + +def test_split_keeps_complete_order_and_checks_each_new_subforward(): + from art.trainer_rank._impl import _SplitForwardPlan + + r = rank() + req = [ + ForwardInput( + input_tokens=torch.arange(128) + i * 1000, + target_tokens=torch.arange(128), + no_grad=False, + ) + for i in range(4) + ] + flat = r._plan_flat_forward(req) + r._memory_profiles[flat.signature] = _MemoryProfile( + bytes_per_token=1, + packed_tokens=512, + logical_per_packed=1, + retained_compute_bytes_per_token=1, + ) + limit = 160_000_000 + used = 0 + r._available_memory_bytes = lambda: limit - used + result = r._find_admissible_forward(req, checkpoint=Unset, refusal_prefix="test") + assert isinstance(result, tuple) + plan, check = result + assert ( + isinstance(plan, _SplitForwardPlan) + and len(plan.subforwards) == 2 + and check.fits + ) + assert sorted(i for group in plan.request_indices for i in group) == list(range(4)) + restored = [None] * 4 + for sub, indices in zip(plan.subforwards, plan.request_indices, strict=True): + assert r._memory_check(sub).fits + for group in sub.groups: + for local, item in zip(group.request_indices, group.items, strict=True): + restored[indices[local]] = item.request + used += r._plan_cost(sub).retained + assert all(a is b for a, b in zip(restored, req, strict=True)) + + +def test_optimistic_split_profile_cliff_preserves_checkpoint_floor(): + r = rank() + req = [ + ForwardInput( + input_tokens=torch.arange(128), + target_tokens=torch.arange(128), + no_grad=False, + ) + for _ in range(16) + ] + full = r._plan_flat_forward(req, memory_minimal=True) + r._memory_profiles[full.signature] = _MemoryProfile( + bytes_per_token=1, + packed_tokens=256, + logical_per_packed=1, + retained_compute_bytes_per_token=1, + ) + cost = r._split_chunk_lower_cost( + req, tuple(x.input_tokens for x in req), checkpoint=Unset + ) + retained = 128 * 40 * 2048 * 2 + assert cost.retained == int((full.output_bytes + retained) * 1.1) + + +@pytest.mark.parametrize( + "field,value", [("recompute_num_layers", True), ("cpu_offloading", 0)] +) +def test_malformed_flag_types_do_not_claim_supported_schedule(field, value): + r = rank() + setattr(r.runtime.model[0].decoder.config, field, value) + assert r._checkpoint_memory_floor(((10, True),)) == (0, 0) diff --git a/tests/unit/test_trainer_rank_head_memory.py b/tests/unit/test_trainer_rank_head_memory.py new file mode 100644 index 000000000..e42b52e53 --- /dev/null +++ b/tests/unit/test_trainer_rank_head_memory.py @@ -0,0 +1,331 @@ +"""Standard BF16 head capacity: CPU/source pricing, not a native peak bound.""" + +from dataclasses import replace +import importlib.util +from pathlib import Path + +import pytest +import torch + +from art.trainer_rank import ForwardInput +from art.trainer_rank._impl import Unset, _MemoryProfile + + +def rank(): + from megatron.core.tensor_parallel.layers import ColumnParallelLinear + + spec = importlib.util.spec_from_file_location( + "checkpoint_memory_tests", + Path(__file__).with_name("test_trainer_rank_checkpoint_memory.py"), + ) + assert spec is not None and spec.loader is not None + source = importlib.util.module_from_spec(spec) + spec.loader.exec_module(source) + r = source.rank() + model = r.runtime.model[0] + head = ColumnParallelLinear.__new__(ColumnParallelLinear) + torch.nn.Module.__init__(head) + head.weight = torch.nn.Parameter( + torch.empty(248320, 2048, device="meta", dtype=torch.bfloat16) + ) + head.input_size = 2048 + head.output_size = head.output_size_per_partition = 248320 + model.output_layer = head + model.share_embeddings_and_output_weights = False + from types import MethodType + + from megatron.core.models.common.language_module.language_module import ( + LanguageModule, + ) + + model._scale_logits = MethodType(LanguageModule._scale_logits, model) + model.config.use_mup = False + model.config.padded_vocab_size = 248320 + r._padded_vocab_size = 248320 + return r + + +def request(rows=512, *, grad=False, hidden=False, ignored=False): + return ForwardInput( + input_tokens=torch.arange(rows), + target_tokens=torch.full((rows,), -100) if ignored else torch.arange(rows), + no_grad=not grad, + hidden_states=hidden, + ) + + +@pytest.mark.parametrize("grad,budget", [(False, 128 * 1024**2), (True, 220 * 1024**2)]) +def test_actual_admission_rejects_below_dense_head_tensor(grad, budget): + r = rank() + plan = r._plan_flat_forward([request(grad=grad)]) + r._available_memory_bytes = lambda: budget + assert not r._memory_check(plan).fits + + +def test_outputs_retention_and_empirical_peak_are_counted_once(): + r = rank() + plan = r._plan_flat_forward([request(grad=True)]) + retained = 512 * 40 * 2048 * 2 + head = 3 * 512 * 248320 * 2 + cost = r._plan_cost(plan) + assert cost.required == int((plan.output_bytes + retained + head) * 1.1) + r._memory_profiles[plan.signature] = _MemoryProfile( + bytes_per_token=2_000_000, + packed_tokens=512, + logical_per_packed=1, + retained_compute_bytes_per_token=1, + ) + cost = r._plan_cost(plan) + assert cost.required == int((plan.output_bytes + 512 * 2_000_000) * 1.1) + assert cost.retained == int((plan.output_bytes + retained) * 1.1) + + +def test_ignored_targets_hidden_only_and_tiny_target_group(): + r = rank() + ignored = request(ignored=True) + hidden = ForwardInput( + input_tokens=torch.arange(10000), hidden_states=True, no_grad=True + ) + target = request(1, grad=True) + assert r._head_projection_rows([ignored, hidden]) == 0 + assert r._plan_head_workspace_bytes(r._plan_flat_forward([ignored, hidden])) == 0 + plan = r._plan_flat_forward([hidden, target]) + assert r._plan_head_workspace_bytes(plan) == 3 * 248320 * 2 + assert r._estimate_flat_forward([hidden, target])[-1] == 3 * 248320 * 2 + assert r._estimate_flat_forward([hidden, target], exact=True)[-1] == 3 * 248320 * 2 + + +def test_multilabel_row_validity_matches_projection(): + r = rank() + item = replace( + request(4), + target_tokens=torch.tensor([[-100, -100], [-100, 2], [3, -100], [-100, -100]]), + ) + assert r._head_projection_rows([item]) == 2 + assert r._plan_head_workspace_bytes(r._plan_flat_forward([item])) == 2 * 248320 * 2 + + +def test_shared_rows_use_lower_upper_and_exact_layout_union(): + r = rank() + a = replace(request(2), target_tokens=torch.tensor([1, -100])) + b = replace(a, target_tokens=torch.tensor([-100, 1])) + req = [a, a, b, b] + assert r._head_projection_rows(req, lower_bound=True) == 1 + assert r._head_projection_rows(req) == 4 + exact = r._estimate_flat_forward(req, exact=True, memory_minimal=True) + plan = r._plan_flat_forward(req, memory_minimal=True) + assert exact[-1] == r._plan_head_workspace_bytes(plan) == 2 * 248320 * 2 + lower = r._split_chunk_lower_cost( + req, tuple(x.input_tokens for x in req), checkpoint=Unset + ) + assert lower.required <= r._plan_cost(plan).required + assert r._estimate_flat_forward(req, memory_minimal=True)[-1] == 248320 * 2 + + +def test_exact_selector_estimate_matches_executed_layout(): + r = rank() + req = [replace(request(2), target_tokens=torch.tensor([1, -100])) for _ in range(4)] + for minimal in (False, True): + values = r._estimate_flat_forward(req, exact=True, memory_minimal=minimal) + plan = r._plan_flat_forward(req, memory_minimal=minimal) + assert values[-1] == r._plan_head_workspace_bytes(plan) + n, out, sig, groups, head = values + assert ( + r._estimate_required_memory_bytes_from_values( + packed_tokens=n, + output_bytes=out, + signature=sig, + logical_tokens=plan.active_logical_tokens, + group_rows=groups, + head_workspace_bytes=head, + ) + == r._memory_check(plan).estimated_required_bytes + ) + + +def test_device_labels_use_capacity_without_reading_values(): + r = rank() + item = replace( + request(128), target_tokens=torch.empty(128, device="meta", dtype=torch.long) + ) + assert r._head_projection_rows([item]) == 128 + assert r._head_projection_rows([item], lower_bound=True) == 0 + assert r._head_projection_rows([item], positions=(torch.arange(128),)) == 128 + + +def test_topk_and_logits_project_ignored_rows_and_chunk_cap(): + r = rank() + ignored = request(2048, ignored=True) + assert r._head_projection_rows([replace(ignored, top_k=2)]) == 512 + assert r._head_projection_rows([replace(ignored, logits=True)]) == 512 + assert r._head_workspace_bytes(4096) == 512 * 248320 * 2 + + +@pytest.mark.parametrize( + "mutation", + [ + "dtype", + "tp", + "head_hook", + "head_override", + "head_dispatch_override", + "vocab_shape", + "quantized", + "missing_weight", + "unknown_vocab", + ], +) +def test_unsupported_head_scope_does_not_claim_dense_bf16_component(mutation): + r = rank() + head = r.runtime.model[0].output_layer + assert r._head_workspace_bytes(512) > 0 + if mutation == "dtype": + head.weight = torch.nn.Parameter(head.weight.float()) + elif mutation == "tp": + r._topology_key = lambda: (1, 2, 1, 1) + elif mutation == "head_hook": + head.register_forward_hook(lambda *args: None) + elif mutation == "head_override": + head.forward = lambda *args, **kwargs: None + elif mutation == "head_dispatch_override": + head._forward_impl = lambda *args, **kwargs: None + elif mutation == "vocab_shape": + head.output_size_per_partition -= 1 + elif mutation == "missing_weight": + head.weight = None + elif mutation == "unknown_vocab": + r._padded_vocab_size = None + else: + r.runtime.model[0].config.fp8 = "hybrid" + assert r._head_workspace_bytes(512) == 0 + + +def test_device_positions_preserve_capacity_without_read(): + r = rank() + item = request(128) + assert ( + r._head_projection_rows( + [item], positions=(torch.empty(128, device="meta", dtype=torch.long),) + ) + == 128 + ) + + +def test_tied_standard_head_weight_uses_the_same_capacity(): + r = rank() + model = r.runtime.model[0] + head = model.output_layer + weight = head.weight + head.weight = None + model.share_embeddings_and_output_weights = True + model.embedding = torch.nn.Module() + model.embedding.word_embeddings = torch.nn.Module() + model.embedding.word_embeddings.weight = weight + assert r._head_workspace_bytes(512) == 512 * 248320 * 2 + + +@pytest.mark.parametrize("rows", [128, 512]) +def test_target_backward_refuses_budget_below_logits_and_both_gradients(rows): + r = rank() + plan = r._plan_flat_forward([request(rows, grad=True)]) + retained, _ = r._checkpoint_memory_floor(r._plan_group_rows(plan)) + dense = min(rows, 512) * 248320 * 2 + before = int((plan.output_bytes + retained + 2 * dense) * 1.1) + expected = int((plan.output_bytes + retained + 3 * dense) * 1.1) + r._available_memory_bytes = lambda: (before + expected) // 2 + check = r._memory_check(plan) + assert check.estimated_required_bytes == expected + assert not check.fits + + +@pytest.mark.parametrize("gradient_rows,reference_rows", [(1, 512), (512, 1)]) +def test_group_head_workspace_keeps_gradient_mode_with_its_rows( + gradient_rows, reference_rows +): + r = rank() + requests = [request(gradient_rows, grad=True), request(reference_rows)] + expected = max(3 * gradient_rows, reference_rows) * 248320 * 2 + plan = r._plan_flat_forward(requests) + assert r._plan_head_workspace_bytes(plan) == expected + for exact in (False, True): + for minimal in (False, True): + assert ( + r._estimate_flat_forward(requests, exact=exact, memory_minimal=minimal)[ + -1 + ] + == expected + ) + + +@pytest.mark.parametrize( + "mutation", ["custom", "other_model", "forged", "mup", "missing"] +) +def test_gradient_statistics_floor_requires_exact_effective_scaling(mutation): + from types import MethodType, SimpleNamespace + + from megatron.core.models.common.language_module.language_module import ( + LanguageModule, + ) + + r = rank() + model = r.runtime.model[0] + req = [request(128, grad=True)] + assert ( + r._plan_head_workspace_bytes(r._plan_flat_forward(req)) == 3 * 128 * 248320 * 2 + ) + if mutation == "custom": + model._scale_logits = lambda logits: logits + elif mutation == "other_model": + model._scale_logits = MethodType( + LanguageModule._scale_logits, + SimpleNamespace(config=SimpleNamespace(use_mup=True, mup_output_mult=2)), + ) + elif mutation == "forged": + + class Forged: + __self__ = model + __func__ = LanguageModule._scale_logits + + def __call__(self, logits): + return logits[..., :1] + + model._scale_logits = Forged() + elif mutation == "mup": + model.config.use_mup = True + else: + del model._scale_logits + # The standard head still allocates its original one-buffer component. + assert r._plan_head_workspace_bytes(r._plan_flat_forward(req)) == 128 * 248320 * 2 + + +@pytest.mark.parametrize("extra", [{"logits": True}, {"top_k": 2}]) +def test_gradient_statistics_floor_survives_additional_output_modes(extra): + r = rank() + req = [replace(request(128, grad=True), **extra)] + assert ( + r._plan_head_workspace_bytes(r._plan_flat_forward(req)) == 3 * 128 * 248320 * 2 + ) + + +def test_gradient_shared_rows_price_same_union_in_exact_and_split_lower_cost(): + r = rank() + a = replace(request(2, grad=True), target_tokens=torch.tensor([1, -100])) + b = replace(a, target_tokens=torch.tensor([-100, 1])) + requests = [a, a, b, b] + plan = r._plan_flat_forward(requests, memory_minimal=True) + expected = 3 * 2 * 248320 * 2 + exact = r._estimate_flat_forward(requests, exact=True, memory_minimal=True) + assert exact[-1] == r._plan_head_workspace_bytes(plan) == expected + assert r._estimate_flat_forward(requests, memory_minimal=True)[-1] == expected // 2 + lower = r._split_chunk_lower_cost( + requests, tuple(x.input_tokens for x in requests), checkpoint=Unset + ) + assert lower.required <= r._plan_cost(plan).required + + +def test_later_sparse_loss_does_not_reduce_6330_projected_targets(): + r = rank() + item = request(6330, grad=True) + plan = r._plan_flat_forward([item]) + assert item.target_tokens.numel() == 6330 + assert r._plan_head_workspace_bytes(plan) == 3 * 512 * 248320 * 2 diff --git a/tests/unit/test_trainer_rank_ignored_mixed_head.py b/tests/unit/test_trainer_rank_ignored_mixed_head.py new file mode 100644 index 000000000..9ac5d500b --- /dev/null +++ b/tests/unit/test_trainer_rank_ignored_mixed_head.py @@ -0,0 +1,133 @@ +"""Ignored labels execute target backward on globally projected rows.""" + +from dataclasses import replace +from types import SimpleNamespace + +import pytest +from test_trainer_rank_head_memory import rank, request +from test_trainer_rank_head_recompute import _Head +import torch + +from art.trainer_rank import ForwardInput, TrainerRank, _impl + + +@pytest.mark.parametrize("extra", [{"logits": True}, {"top_k": 2}]) +def test_ignored_rows_reactivated_by_same_item_output_keep_backward_floor(extra): + r = rank() + item = replace(request(128, grad=True, ignored=True), **extra) + assert ( + r._plan_head_workspace_bytes(r._plan_flat_forward([item])) + == 3 * 128 * 248320 * 2 + ) + assert r._estimate_flat_forward([item], exact=True)[-1] == 3 * 128 * 248320 * 2 + assert r._estimate_flat_forward([item])[-1] == 3 * 128 * 248320 * 2 + assert r._head_target_chunk_rows([item], lower_bound=True) == 0 + + +@pytest.mark.parametrize("extra", [{"logits": True}, {"top_k": 2}]) +def test_ignored_cross_request_overlap_and_validity_have_separate_roles(extra): + r = rank() + labelled = request(513, grad=True, ignored=True) + raw = ForwardInput(input_tokens=torch.arange(512), no_grad=False, **extra) + positions = (torch.arange(513), torch.arange(512)) + req = [labelled, raw] + assert r._head_projection_rows(req, positions=positions) == 512 + assert r._head_target_chunk_rows(req, positions=positions) == 512 + assert r._head_target_chunk_rows(req) == 512 + assert r._head_target_chunk_rows(req, lower_bound=True) == 0 + disjoint = (torch.arange(513) + 1000, torch.arange(512)) + assert r._head_target_chunk_rows(req, positions=disjoint) == 0 + valid_tail = replace( + labelled, target_tokens=torch.cat((torch.full((512,), -100), torch.tensor([1]))) + ) + assert r._head_target_chunk_rows([valid_tail, raw], positions=positions) == 512 + assert r._head_target_chunk_rows([valid_tail, raw], lower_bound=True) == 1 + assert r._head_projection_rows([labelled]) == 0 + assert r._plan_head_workspace_bytes(r._plan_flat_forward([labelled])) == 0 + + +@pytest.mark.parametrize("extra", [{"logits": True}, {"top_k": 2}]) +def test_actual_ignored_mixed_backward_keeps_zero_dense_index_graph(monkeypatch, extra): + monkeypatch.setattr(_impl, "_HEAD_CHUNK_TOKENS", 4) + monkeypatch.setattr(_impl, "_language_model", lambda model: model) + monkeypatch.setattr(_impl, "_all_reduce_tensor_parallel_max", lambda x: x) + monkeypatch.setattr(_impl, "_all_reduce_tensor_parallel_sum", lambda x: x) + monkeypatch.setattr(_impl, "_vocab_range", lambda logits: (0, logits.shape[-1])) + monkeypatch.setattr( + TrainerRank, "_gather_tensor_parallel_logits", lambda self, x: x + ) + monkeypatch.setattr( + _impl, + "_vocab_parallel_topk_from_local", + lambda values, tokens, *, k, log_z, vocab_start: _impl.TopK( + values[:, :k] - log_z[:, None], tokens[:, :k] + ), + ) + original = _impl._vocab_parallel_target_logprobs + dense_backward = [] + normalizer_backward = [] + + def target_path(logits, labels, log_z, *, row_offsets): + assert (labels == -100).all() + log_z.register_hook( + lambda grad: normalizer_backward.append( + (tuple(grad.shape), bool((grad == 0).all())) + ) + ) + output = original(logits, labels, log_z, row_offsets=row_offsets) + queue = [output.grad_fn] + seen = set() + while queue: + node = queue.pop() + if node is None or node in seen: + continue + seen.add(node) + if type(node).__name__.startswith("IndexBackward"): + node.register_hook( + lambda inputs, outputs: dense_backward.append( + (tuple(inputs[0].shape), bool((inputs[0] == 0).all())) + ) + ) + queue.extend(next_node for next_node, _ in node.next_functions) + return output + + monkeypatch.setattr(_impl, "_vocab_parallel_target_logprobs", target_path) + generator = torch.Generator().manual_seed(79) + hidden = torch.randn(8, 5, generator=generator, requires_grad=True) + weight = torch.randn(17, 5, generator=generator) + model = SimpleNamespace( + output_layer=_Head(weight), + vocab_size=17, + share_embeddings_and_output_weights=False, + _scale_logits=lambda x: x, + ) + trainer = object.__new__(TrainerRank) + trainer.runtime = SimpleNamespace(model=[model]) + # The ignored-only item projects no row itself. The independent raw/top-k + # request reaches its positions and therefore executes the target branch. + labelled = ForwardInput( + input_tokens=torch.tensor([0, 7]), + target_tokens=torch.tensor([-100, -100]), + no_grad=False, + ) + raw = ForwardInput(input_tokens=torch.arange(8), no_grad=False, **extra) + positions = (torch.tensor([0, 7]), torch.arange(8)) + outputs = trainer._project_head( + [trainer._forward_item(r) for r in [labelled, raw]], + SimpleNamespace( + positions_by_item=positions, + source_positions_by_item=(torch.arange(2), torch.arange(8)), + ), + hidden, + ) + loss = outputs[0].target_logprobs.sum() + assert loss.requires_grad and float(loss.detach()) == 0 + loss.backward() + assert hidden.grad is not None and bool((hidden.grad == 0).all()) + assert model.output_layer.weight.grad is not None and bool( + (model.output_layer.weight.grad == 0).all() + ) + assert dense_backward and all( + shape == (4, 17) and zero for shape, zero in dense_backward + ) + assert normalizer_backward and all(zero for _, zero in normalizer_backward) diff --git a/tests/unit/test_trainer_rank_mixed_head_memory.py b/tests/unit/test_trainer_rank_mixed_head_memory.py new file mode 100644 index 000000000..86e01c1c4 --- /dev/null +++ b/tests/unit/test_trainer_rank_mixed_head_memory.py @@ -0,0 +1,219 @@ +"""Mixed head admission components; CPU autograd, never a full CUDA bound.""" + +from dataclasses import replace +from types import MethodType, SimpleNamespace + +import pytest +from test_trainer_rank_head_memory import rank, request +from test_trainer_rank_head_recompute import _Head +import torch + +from art.trainer_rank import ForwardInput, TrainerRank, _impl +from art.trainer_rank._impl import Unset + + +@pytest.mark.parametrize("extra", [{"logits": True}, {"top_k": 2}]) +def test_adding_output_cannot_erase_existing_target_admission_floor(extra): + r = rank() + target = request(128, grad=True) + added = ForwardInput(input_tokens=torch.tensor([20000]), no_grad=False, **extra) + before = r._plan_cost(r._plan_flat_forward([target])).required + plan = r._plan_flat_forward([target, added]) + after = r._plan_cost(plan).required + print({"extra": extra, "before": before, "after": after}) + assert after >= before + assert r._plan_head_workspace_bytes(plan) == 3 * 129 * 248320 * 2 + r._available_memory_bytes = lambda: before - 1 + assert not r._memory_check(plan).fits + + +@pytest.mark.parametrize("extra", [{"logits": True}, {"top_k": 2}]) +def test_sparse_target_prices_its_full_mixed_chunk_and_short_tail(extra): + r = rank() + target = replace(request(1, grad=True), input_tokens=torch.tensor([9999])) + raw = ForwardInput(input_tokens=torch.arange(512), no_grad=False, **extra) + req = [target, raw] + full = (torch.tensor([0]), torch.arange(1, 513)) + tail = (torch.tensor([512]), torch.arange(512)) + assert r._head_target_chunk_rows(req, positions=full) == 512 + assert r._head_target_chunk_rows(req, positions=tail) == 1 + assert r._head_target_chunk_rows(req, lower_bound=True) == 1 + assert r._head_target_chunk_rows(req) == 512 + dense = 512 * 248320 * 2 + assert ( + r._group_head_workspace_bytes(512, req, grad_enabled=True, positions=full) + == 3 * dense + ) + assert ( + r._group_head_workspace_bytes(512, req, grad_enabled=True, positions=tail) + == dense + ) + + +@pytest.mark.parametrize("extra", [{"logits": True}, {"top_k": 2}]) +def test_shared_multilabel_union_matches_actual_and_split_bounds(extra): + r = rank() + a = replace( + request(4, grad=True), + target_tokens=torch.tensor( + [[-100, 1], [-100, -100], [-100, -100], [-100, -100]] + ), + ) + b = replace( + a, + target_tokens=torch.tensor( + [[-100, -100], [-100, -100], [2, -100], [-100, -100]] + ), + ) + raw = ForwardInput(input_tokens=torch.arange(4), no_grad=False, **extra) + req = [a, a, b, b, raw] + for minimal in (False, True): + plan = r._plan_flat_forward(req, memory_minimal=minimal) + exact = r._estimate_flat_forward(req, exact=True, memory_minimal=minimal) + assert exact[-1] == r._plan_head_workspace_bytes(plan) + lower = r._estimate_flat_forward(req, memory_minimal=True)[-1] + upper = r._estimate_flat_forward(req)[-1] + assert lower <= exact[-1] <= upper + assert ( + r._split_chunk_lower_cost( + req, tuple(x.input_tokens for x in req), checkpoint=Unset + ).required + <= r._plan_cost(plan).required + ) + assert ( + r._plan_head_workspace_bytes(r._plan_flat_forward(req, memory_minimal=True)) + == 3 * 4 * 248320 * 2 + ) + + +@pytest.mark.parametrize("extra", [{"logits": True}, {"top_k": 2}]) +def test_ignored_device_labels_and_no_target_keep_distinct_guards(extra): + r = rank() + ignored = replace(request(128, grad=True, ignored=True), **extra) + dense = 128 * 248320 * 2 + assert r._plan_head_workspace_bytes(r._plan_flat_forward([ignored])) == 3 * dense + no_target = replace(ignored, target_tokens=None) + assert r._plan_head_workspace_bytes(r._plan_flat_forward([no_target])) == dense + device = replace( + ignored, target_tokens=torch.empty(128, device="meta", dtype=torch.long) + ) + assert r._head_target_chunk_rows([device], lower_bound=True) == 0 + assert r._head_target_chunk_rows([device]) == 128 + assert r._head_target_chunk_rows([device], positions=(torch.arange(128),)) == 128 + assert ( + r._head_target_chunk_rows( + [device], positions=(torch.empty(128, device="meta", dtype=torch.long),) + ) + == 128 + ) + + +@pytest.mark.parametrize("mutation", ["no_grad", "custom_scale", "mup", "head_hook"]) +def test_mixed_path_preserves_source_scaling_and_gradient_guards(mutation): + r = rank() + item = replace(request(128, grad=True), logits=True) + model = r.runtime.model[0] + if mutation == "no_grad": + item = replace(item, no_grad=True) + elif mutation == "custom_scale": + model._scale_logits = lambda logits: logits + elif mutation == "mup": + model.config.use_mup = True + else: + model.output_layer.register_forward_hook(lambda *args: None) + expected = 0 if mutation == "head_hook" else 128 * 248320 * 2 + assert r._plan_head_workspace_bytes(r._plan_flat_forward([item])) == expected + + +@pytest.mark.parametrize( + "extra", + [{"logits": True}, {"top_k": 2}, {"top_k": 12}, {"logits": True, "top_k": 2}], +) +def test_actual_mixed_projection_preserves_target_outputs_and_backward( + monkeypatch, extra +): + from megatron.core.models.common.language_module.language_module import ( + LanguageModule, + ) + + monkeypatch.setattr(_impl, "_HEAD_CHUNK_TOKENS", 4) + monkeypatch.setattr(_impl, "_language_model", lambda model: model) + monkeypatch.setattr(_impl, "_all_reduce_tensor_parallel_max", lambda x: x) + monkeypatch.setattr(_impl, "_all_reduce_tensor_parallel_sum", lambda x: x) + monkeypatch.setattr(_impl, "_vocab_range", lambda logits: (0, logits.shape[-1])) + monkeypatch.setattr( + TrainerRank, "_gather_tensor_parallel_logits", lambda self, x: x + ) + monkeypatch.setattr( + _impl, + "_vocab_parallel_topk_from_local", + lambda values, tokens, *, k, log_z, vocab_start: _impl.TopK( + values[:, :k] - log_z[:, None], tokens[:, :k] + ), + ) + original = _impl._vocab_parallel_target_logprobs + calls = [] + + def target_path(logits, labels, log_z, *, row_offsets): + calls.append((tuple(logits.shape), labels.tolist(), row_offsets.tolist())) + return original(logits, labels, log_z, row_offsets=row_offsets) + + monkeypatch.setattr(_impl, "_vocab_parallel_target_logprobs", target_path) + generator = torch.Generator().manual_seed(97) + hidden = torch.randn(13, 5, generator=generator, dtype=torch.float64) + weights = torch.randn(17, 5, generator=generator, dtype=torch.float64) + target_positions = torch.tensor([0, 2, 12]) + labels = torch.tensor([[1, -100], [2, 3], [16, -100]]) + target = ForwardInput( + input_tokens=torch.tensor([0, 2, 12]), target_tokens=labels, no_grad=False + ) + added = ForwardInput(input_tokens=torch.arange(13), no_grad=False, **extra) + + def run(mixed): + model = SimpleNamespace( + output_layer=_Head(weights), + vocab_size=17, + config=SimpleNamespace(use_mup=False, padded_vocab_size=17), + share_embeddings_and_output_weights=False, + ) + model._scale_logits = MethodType(LanguageModule._scale_logits, model) + trainer = object.__new__(TrainerRank) + trainer.runtime = SimpleNamespace(model=[model]) + x = hidden.clone().requires_grad_() + requests = [target, added] if mixed else [target] + positions = ( + (target_positions, torch.arange(13)) if mixed else (target_positions,) + ) + outputs = trainer._project_head( + [trainer._forward_item(r) for r in requests], + SimpleNamespace( + positions_by_item=positions, + source_positions_by_item=tuple(torch.arange(len(p)) for p in positions), + ), + x, + ) + loss = -outputs[0].target_logprobs.sum() / int((labels != -100).sum()) + loss.backward() + if mixed: + expected_logits = hidden @ weights.T + if extra.get("logits"): + torch.testing.assert_close(outputs[1].logits, expected_logits) + if "top_k" in extra: + expected_values, expected_tokens = torch.topk( + expected_logits.float().log_softmax(-1), k=extra["top_k"], dim=-1 + ) + torch.testing.assert_close(outputs[1].top_k.logprobs, expected_values) + torch.testing.assert_close(outputs[1].top_k.tokens, expected_tokens) + return ( + outputs[0].target_logprobs.detach(), + x.grad, + model.output_layer.weight.grad, + ) + + before = run(False) + calls.clear() + after = run(True) + for a, b in zip(after, before, strict=True): + torch.testing.assert_close(a, b, rtol=1e-6, atol=1e-7) + assert any(shape == (4, 17) and len(rows) < 4 for shape, _, rows in calls) + assert all(value.isfinite().all() and value.abs().sum() > 0 for value in after) diff --git a/tests/unit/test_trainer_rank_pending_memory.py b/tests/unit/test_trainer_rank_pending_memory.py new file mode 100644 index 000000000..db41cb386 --- /dev/null +++ b/tests/unit/test_trainer_rank_pending_memory.py @@ -0,0 +1,265 @@ +"""Conditional CP1 save accounting: original geometry and real CPU owner metadata.""" + +from dataclasses import replace +from types import MethodType, SimpleNamespace +from typing import Any, cast + +import pytest +from test_trainer_rank_moe_memory import _enclosing_moe +from test_trainer_rank_moe_memory import layer as layer +import torch + +from art.megatron.prefix_tree_packing import prefix_tree_pack +from art.trainer_rank import ForwardInput, TrainerRank +from art.trainer_rank import _gdn_memory as g +from art.trainer_rank._impl import Unset, _MemoryProfile + + +def module(cls): + obj = cls.__new__(cls) + torch.nn.Module.__init__(obj) + return obj + + +def rank_with_moe(moe_layer): + from megatron.core.ssm.gated_delta_net import GatedDeltaNet + from megatron.core.transformer.transformer_block import TransformerBlock + from transformer_engine.pytorch import RMSNorm + + from art.megatron.gdn.operator import _prefix_tree_forward + from art.megatron.lora import LoRA, SelfAttentionLinearProjLoRA + + decoder = module(TransformerBlock) + decoder.config = SimpleNamespace( + hidden_size=2048, + num_layers=40, + padded_vocab_size=32, + params_dtype=torch.bfloat16, + recompute_granularity="full", + recompute_method="uniform", + recompute_num_layers=1, + distribute_saved_activations=False, + sequence_parallel=False, + fp32_residual_connection=False, + cpu_offloading=False, + cuda_graph_impl="none", + fp8=None, + fp4=None, + ) + decoder.layers = torch.nn.ModuleList( + [torch.nn.Linear(1, 1).bfloat16() for _ in range(40)] + ) + decoder.num_layers_per_pipeline_rank = 40 + layer = torch.nn.Module() + layer.mlp = moe_layer + gd = module(GatedDeltaNet) + gd.num_key_heads = 16 + gd.num_value_heads = 32 + gd.key_head_dim = 128 + gd.value_head_dim = 128 + gd.conv_kernel_dim = 4 + gd.use_qk_l2norm = True + gd.tp_size = gd.sp_size = 1 + gd.forward = MethodType(_prefix_tree_forward, gd) + gd.conv1d = torch.nn.Conv1d( + 8192, 8192, 4, groups=8192, bias=False, dtype=torch.bfloat16 + ) + gd.out_norm = module(RMSNorm) + gd.out_norm.weight = torch.nn.Parameter(torch.ones(128, dtype=torch.bfloat16)) + gd.out_proj = module(SelfAttentionLinearProjLoRA) + gd.out_proj.lora = module(LoRA) + gd.out_proj.lora.A_T = torch.nn.Parameter( + torch.empty(4096, 1, dtype=torch.bfloat16) + ) + gd.out_proj.lora.B_T = torch.nn.Parameter( + torch.empty(1, 2048, dtype=torch.bfloat16) + ) + layer.self_attention = gd + decoder.layers[38] = layer + model: Any = torch.nn.Module() + model.config = decoder.config + model.decoder = decoder + model._preprocess = lambda: None + r: Any = TrainerRank( + cast( + Any, + SimpleNamespace( + model=[model], + optimizer=None, + provider=SimpleNamespace(hidden_size=2048, num_layers=40), + model_support_handler=SimpleNamespace(build_gdn_execution_spec=False), + ), + ) + ) + r._dp_rank_and_size = lambda: (0, 1) # Uninitialized MCore has no CPU DP group. + return r, gd + + +@pytest.fixture +def pending_rank(layer): + return rank_with_moe(_enclosing_moe(layer))[0] + + +def full_requests(no_grad=False): + return [ + ForwardInput( + input_tokens=torch.arange(6330) + i * 10000, + target_tokens=torch.arange(6330), + no_grad=no_grad, + ) + for i in range(8) + ] + + +def test_actual_constructor_cache_and_full_plan(pending_rank): + rank = pending_rank + assert rank._moe_output_bytes_per_token == 188416 + shapes = g.model_shapes(rank) + assert shapes is not None and shapes[1][0].moe_bytes_per_row == 188416 + requests = full_requests() + plan = rank._plan_flat_forward(requests) + assert rank._estimate_flat_forward(requests) is None + assert ( + plan.packed_tokens == plan.logical_tokens == 50640 and plan.request_count == 8 + ) + assert g.plan_floor(rank, plan) == (8296857600, 9541386240 + 3157761952) + assert ( + rank._memory_check(plan).estimated_required_bytes + == rank._plan_cost(plan).required + == 23095829187 + ) + selected = rank._select_next_micro_batch(requests, 0) + assert ( + selected.check.estimated_required_bytes + == rank._memory_check(selected.plan).estimated_required_bytes + ) + lower = rank._split_chunk_lower_cost( + requests, tuple(r.input_tokens for r in requests), checkpoint=Unset + ) + assert lower.required <= rank._memory_check(plan).estimated_required_bytes + + +def test_pending_no_grad_empty_mixed_and_learned_max(pending_rank): + rank = pending_rank + grad = ForwardInput( + input_tokens=torch.arange(67), hidden_states=True, no_grad=False + ) + reference = replace(grad, input_tokens=torch.arange(4096), no_grad=True) + plan = rank._plan_flat_forward([grad]) + retained, workspace = g.plan_floor(rank, plan) + assert g.plan_floor(rank, rank._plan_flat_forward([])) == (0, 0) + assert g.plan_floor(rank, rank._plan_flat_forward([reference])) == (0, 0) + mixed = rank._plan_flat_forward([grad, reference]) + mr, mw = g.plan_floor(rank, mixed) + assert mr == retained and mw == max(workspace, 4096 * 188416) + assert rank._estimate_flat_forward([reference]) is not None + rank._memory_profiles[plan.signature] = _MemoryProfile( + bytes_per_token=1, + packed_tokens=plan.packed_tokens, + logical_per_packed=1, + retained_compute_bytes_per_token=1, + ) + cost = rank._plan_cost(plan) + assert cost.retained == int((plan.output_bytes + retained) * 1.1) + rank._memory_profiles[plan.signature] = replace( + rank._memory_profiles[plan.signature], bytes_per_token=10**9 + ) + assert rank._plan_cost(plan).required == int( + (plan.output_bytes + plan.packed_tokens * 10**9) * 1.1 + ) + + +@pytest.mark.parametrize("bad", [0, -1, True, 1.5, None]) +def test_invalid_cached_coefficient_stops_before_memory_reduction(pending_rank, bad): + rank = pending_rank + plan = rank._plan_flat_forward(full_requests()) + rank._moe_output_bytes_per_token = bad + statuses = [] + rank._all_ranks_true = lambda v: (statuses.append(v), v)[1] + rank._memory_check_required = lambda *a, **kw: pytest.fail( + "memory reduction entered" + ) + with pytest.raises(ValueError, match="Invalid constructor MoE coefficient"): + rank._memory_check(plan, sync_planning_errors=True, sync_across_dp=True) + assert statuses == [False] + statuses.clear() + with pytest.raises(ValueError, match="Invalid constructor MoE coefficient"): + rank._estimate_flat_forward(full_requests(), sync_planning_errors=True) + assert statuses == [False] + + +@pytest.mark.parametrize( + "field,value", + [ + ("recompute_granularity", None), + ("recompute_num_layers", 2), + ("sequence_parallel", True), + ("params_dtype", torch.float32), + ], +) +def test_unsupported_pending_modes_are_not_qualified(pending_rank, field, value): + setattr(pending_rank.runtime.model[0].decoder.config, field, value) + assert g.model_shapes(pending_rank) is None + + +@pytest.mark.parametrize("root_length", [1, 63, 64, 65, 127, 128, 129]) +@pytest.mark.parametrize("depth", [0, 1, 2, 8]) +def test_cp1_bucket_geometry_matches_original_builder(root_length, depth): + from art.megatron.gdn import gdn_prefix_tree as original + + prefix = list(range(root_length)) + pack = prefix_tree_pack( + tuple( + torch.tensor(x) + for x in [ + prefix + [10000, 10001, 10002], + prefix + [10000, 10001, 10003], + prefix + [20000, 20001], + [30000, 30001], + ] + ), + max_depth=depth, + ) + spec = original.parse_gdn_prefix_tree_segments(pack.group_ids, pack.parent_ids) + has_children = tuple( + i in spec.tree_parent_indices for i in range(len(spec.tree_segments)) + ) + actual = original._build_chunk_aligned_cp1_tree_buckets( + spec, has_children, device="cpu", planner_config=original.GdnPlannerConfig() + ) + actual_rows = tuple( + ( + tuple( + zip( + cast(torch.Tensor, b.family_indices_cpu).tolist(), + cast(torch.Tensor, b.parent_indices_cpu).tolist(), + b.lengths_cpu.tolist(), + ) + ), + b.needs_final_state, + ) + for level in actual + for b in level + ) + assert actual_rows == tuple( + (b.columns, b.final) for b in g.cp1_buckets(pack.segments) + ) + + +@pytest.mark.parametrize( + "change", [{"parent_id": 999}, {"packed_start": 1}, {"end": 0}, {"group_id": True}] +) +def test_invalid_cp1_geometry_refused(change): + pack = prefix_tree_pack((torch.arange(65),), max_depth=0) + with pytest.raises(ValueError): + g.cp1_buckets((replace(pack.segments[0], **change),)) + + +def test_pending_counts_bucket_backing_and_output_rank_separately(): + shape = g.Shape(16, 32, 128, 128, 4, 1, 188416) + pack = prefix_tree_pack(tuple(torch.arange(6330) for _ in range(8)), max_depth=0) + buckets = g.cp1_buckets(pack.segments) + assert shape.pending(50640, buckets) == 3157267456 + 8 * 8192 * 3 * 2 + 50640 * 2 + assert replace(shape, output_lora_rank=0).pending(50640, buckets) == shape.pending( + 50640, buckets + ) - 50640 * (8192 + 2) diff --git a/tests/unit/test_trainer_rank_planning_status.py b/tests/unit/test_trainer_rank_planning_status.py index 2054024ee..23215b290 100644 --- a/tests/unit/test_trainer_rank_planning_status.py +++ b/tests/unit/test_trainer_rank_planning_status.py @@ -103,6 +103,7 @@ def _worker(index: int, directory: Path) -> None: ): rank = TrainerRank.__new__(TrainerRank) rank.device = torch.device("cpu") + rank._padded_vocab_size = None rank._planning_seconds_accum = 0.0 rank._dp_rank_and_size = lambda: (index, 2) rank._physical_tokens = lambda tokens: tokens @@ -110,7 +111,7 @@ def _worker(index: int, directory: Path) -> None: rank._estimate_group_request_output_bytes = lambda requests: 0 rank._memory_signature_from_requests = lambda *args, **kwargs: None rank._forward_item = lambda request: SimpleNamespace( - input_ids=request.input_tokens + input_ids=request.input_tokens, request=request ) rank._forward_output_metadata = lambda *args, **kwargs: (None, True) @@ -190,7 +191,7 @@ def price(**kwargs): error = caught if mode in ("estimate", "materialize", "price"): if index == 0: - assert error is primary + assert error is primary, (mode, repr(error)) assert error.__cause__ is cause and error.__context__ is context else: assert type(error) is RuntimeError diff --git a/tests/unit/test_trainer_rank_shared_memory.py b/tests/unit/test_trainer_rank_shared_memory.py new file mode 100644 index 000000000..c7dd4b6c7 --- /dev/null +++ b/tests/unit/test_trainer_rank_shared_memory.py @@ -0,0 +1,231 @@ +"""One supported shared return held across routed compute; not all backward saves.""" + +from types import SimpleNamespace + +import pytest +from test_trainer_rank_moe_memory import _enclosing_moe, _rank +from test_trainer_rank_moe_memory import layer as layer +from test_trainer_rank_pending_memory import full_requests, module, rank_with_moe +import torch + +from art.trainer_rank import ForwardInput +from art.trainer_rank import _gdn_memory as g +from art.trainer_rank._impl import ( + _moe_output_bytes_per_token, + _shared_expert_output_bytes_per_token, +) +from art.trainer_rank._planner_cost import ParallelShape + + +def shared_layer(layer, gated=True): + from megatron.core.extensions.transformer_engine import ( + TEColumnParallelLinear, + TERowParallelLinear, + ) + from megatron.core.transformer.moe.shared_experts import SharedExpertMLP + + from art.megatron.lora import ( + LoRA, + SelfAttentionLinearProjLoRA, + SharedExpertsLinearFC1LoRA, + SharedExpertsLinearFC2LoRA, + ) + + _enclosing_moe(layer) + values = dict( + params_dtype=torch.bfloat16, + moe_shared_expert_overlap=False, + sequence_parallel=False, + fp32_residual_connection=False, + add_bias_linear=False, + use_te_activation_func=False, + bias_activation_fusion=False, + gated_linear_unit=True, + tensor_model_parallel_size=1, + context_parallel_size=1, + pipeline_model_parallel_size=1, + expert_model_parallel_size=1, + expert_tensor_parallel_size=1, + moe_shared_expert_intermediate_size=512, + activation_func=torch.nn.functional.silu, + ) + vars(layer.config).update(values) + layer.use_shared_expert = True + layer.shared_expert_overlap = False + layer.shared_experts_recompute = False + layer.moe_layer_recompute = False + layer.fwd_execution_map = ["route", "expert_compute", "postprocess"] + shared = module(SharedExpertMLP) + layer.shared_experts = shared + shared.config = SimpleNamespace(**{**vars(layer.config), "ffn_hidden_size": 512}) + shared.activation_func = torch.nn.functional.silu + shared.use_shared_expert_gate = gated + shared.gate_weight = torch.nn.Parameter( + torch.empty(1, 2048, dtype=torch.bfloat16), requires_grad=False + ) + fc1 = module(SharedExpertsLinearFC1LoRA) + shared.linear_fc1 = fc1 + fc1.non_gated = False + fc1.out_features = 1024 + fc1.linear_fc1 = module(TEColumnParallelLinear) + fc1.linear_fc1.weight = torch.nn.Parameter( + torch.empty(1024, 2048, dtype=torch.bfloat16), requires_grad=False + ) + + def adapter(inputs, outputs): + value = module(LoRA) + value.A_T = torch.nn.Parameter(torch.empty(inputs, 8, dtype=torch.bfloat16)) + value.B_T = torch.nn.Parameter(torch.empty(8, outputs, dtype=torch.bfloat16)) + return value + + fc1.gate_lora = adapter(2048, 512) + fc1.up_lora = adapter(2048, 512) + fc2 = module(SharedExpertsLinearFC2LoRA) + shared.linear_fc2 = fc2 + row = module(SelfAttentionLinearProjLoRA) + fc2.row_parallel_lora = row + row.provider = SimpleNamespace( + tensor_model_parallel_size=1, sequence_parallel=False + ) + row.lora = adapter(512, 2048) + row.linear_proj = module(TERowParallelLinear) + row.linear_proj.weight = torch.nn.Parameter( + torch.empty(2048, 512, dtype=torch.bfloat16), requires_grad=False + ) + return layer + + +def coefficient(layer): + return _moe_output_bytes_per_token([layer], ParallelShape(tp=1, cp=1)) + + +@pytest.mark.parametrize("gate", [False, True]) +@pytest.mark.parametrize("no_grad", [False, True]) +def test_shared_return_in_actual_constructor_and_plan(layer, gate, no_grad): + rank, _ = rank_with_moe(shared_layer(layer, gate)) + assert rank._moe_output_bytes_per_token == 192512 + shapes = g.model_shapes(rank) + assert shapes is not None and shapes[1][0].moe_bytes_per_row == 192512 + requests = full_requests(no_grad) + plan = rank._plan_flat_forward(requests) + assert ( + rank._memory_check(plan).estimated_required_bytes + == rank._plan_cost(plan).required + ) + if no_grad: + assert g.plan_floor(rank, plan) == (0, 0) + assert rank._plan_cost(plan).required == 10723911264 + else: + assert g.plan_floor(rank, plan) == (8296857600, 50640 * 192512 + 3157761952) + assert rank._plan_cost(plan).required == 23323992771 + selected = rank._select_next_micro_batch(requests, 0) + assert ( + selected.check.estimated_required_bytes + == rank._memory_check(selected.plan).estimated_required_bytes + ) + + +mutations = { + "no shared": lambda x: delattr(x, "shared_experts"), + "disabled shared": lambda x: setattr(x, "use_shared_expert", False), + "overlap": lambda x: setattr(x, "shared_expert_overlap", True), + "config overlap": lambda x: setattr(x.config, "moe_shared_expert_overlap", True), + "shared unknown owner": lambda x: setattr(x, "shared_experts", torch.nn.Identity()), + "shared forward replaced": lambda x: setattr( + x.shared_experts, "forward", lambda *a: None + ), + "shared forward hook": lambda x: x.shared_experts.register_forward_hook( + lambda *a: None + ), + "fc2 unknown owner": lambda x: setattr( + x.shared_experts, "linear_fc2", torch.nn.Identity() + ), + "fc1 missing": lambda x: delattr(x.shared_experts, "linear_fc1"), + "adapter missing": lambda x: delattr(x.shared_experts.linear_fc1.gate_lora, "A_T"), + "adapter shape": lambda x: setattr( + x.shared_experts.linear_fc1.gate_lora, + "A_T", + torch.nn.Parameter(torch.empty(2047, 8, dtype=torch.bfloat16)), + ), + "base dtype": lambda x: ( + x.shared_experts.linear_fc2.row_parallel_lora.linear_proj.float() + ), + "gate dtype": lambda x: setattr( + x.shared_experts, "gate_weight", torch.nn.Parameter(torch.empty(1, 2048)) + ), + "gate shape": lambda x: setattr( + x.shared_experts, + "gate_weight", + torch.nn.Parameter(torch.empty(2, 2048, dtype=torch.bfloat16)), + ), + "gate bool": lambda x: setattr(x.shared_experts, "use_shared_expert_gate", 1), + "sequence parallel": lambda x: setattr( + x.shared_experts.config, "sequence_parallel", True + ), + "activation": lambda x: setattr( + x.shared_experts, "activation_func", torch.nn.functional.relu + ), + "partial layer execution": lambda x: setattr( + x, "fwd_execution_map", ["expert_compute", "postprocess"] + ), + "recomputed shared": lambda x: setattr(x, "shared_experts_recompute", True), + "recomputed MoE": lambda x: setattr(x, "moe_layer_recompute", True), + "shared compute replaced": lambda x: setattr( + x, "shared_experts_compute", lambda *a: None + ), + "topology bool": lambda x: setattr(x.config, "tensor_model_parallel_size", True), + "topology two": lambda x: setattr(x.config, "context_parallel_size", 2), + "missing config": lambda x: delattr(x.shared_experts, "config"), +} + + +@pytest.mark.parametrize("name", list(mutations)) +def test_unsupported_shared_branch_keeps_prior_routed_component(layer, name): + shared_layer(layer) + mutations[name](layer) + assert _shared_expert_output_bytes_per_token(layer) == 0 + assert coefficient(layer) == 188416 + assert _rank(layer)._moe_output_bytes_per_token == 188416 + + +def test_shared_return_has_no_topk_or_layer_multiplier(layer): + shared_layer(layer) + assert coefficient(layer) == 192512 + assert ( + _moe_output_bytes_per_token([layer, layer], ParallelShape(tp=1, cp=1)) == 192512 + ) + layer.config.moe_router_topk = layer.router.topk = 4 + assert coefficient(layer) == 4 * (3 * 512 + 5 * 2048) * 2 + 4096 + assert _shared_expert_output_bytes_per_token(layer) == 4096 + + +def test_shared_component_is_joint_layer_max_and_survives_fc1_fallback(layer): + shared_layer(layer) + # One layer has the larger shared return; the other has the larger routed + # stage. Duplicate only tiny metadata; CPU tensors are never executed. + import copy + + other = copy.deepcopy(layer) + del other.shared_experts + other.config.moe_router_topk = other.router.topk = 9 + assert ( + _moe_output_bytes_per_token([layer, other], ParallelShape(tp=1, cp=1)) + == 9 * (3 * 512 + 5 * 2048) * 2 + ) + del layer.experts.linear_fc1 + assert coefficient(layer) == 8 * (512 + 3 * 2048) * 2 + 4096 + + +def test_reference_group_has_shared_stage_but_no_pending_save(layer): + rank, _ = rank_with_moe(shared_layer(layer)) + grad = ForwardInput( + input_tokens=torch.arange(67), hidden_states=True, no_grad=False + ) + reference = ForwardInput( + input_tokens=torch.arange(4096), hidden_states=True, no_grad=True + ) + one = rank._plan_flat_forward([grad]) + mixed = rank._plan_flat_forward([grad, reference]) + retained, workspace = g.plan_floor(rank, one) + assert g.plan_floor(rank, mixed) == (retained, max(workspace, 4096 * 192512)) + assert g.plan_floor(rank, rank._plan_flat_forward([reference])) == (0, 0) diff --git a/tests/unit/test_trainer_rank_weird_shapes.py b/tests/unit/test_trainer_rank_weird_shapes.py index 8db3ef72e..e7962d903 100644 --- a/tests/unit/test_trainer_rank_weird_shapes.py +++ b/tests/unit/test_trainer_rank_weird_shapes.py @@ -15,7 +15,6 @@ AdapterSelection, ForwardInput, ForwardOutput, - TopK, TrainerRank, TrainerRankMemoryError, Unset, @@ -225,10 +224,12 @@ def test_planner_handles_vineppo_nested_shape_and_request_mix() -> None: estimate = rank._estimate_flat_forward(flat) assert estimate is not None - packed_tokens, output_bytes, signature = estimate + packed_tokens, output_bytes, signature, group_rows, head_workspace_bytes = estimate assert packed_tokens == plan.packed_tokens assert output_bytes == plan.output_bytes assert signature == plan.signature + assert group_rows == rank._plan_group_rows(plan) + assert head_workspace_bytes == rank._plan_head_workspace_bytes(plan) assert plan.request_count == 12 assert plan.signature.request_mix == ( "target:(2,)", @@ -936,10 +937,12 @@ def slot_ref(name: str | None) -> SlotRef | None: estimate = rank._estimate_flat_forward(requests) assert estimate is not None - packed_tokens, output_bytes, signature = estimate + packed_tokens, output_bytes, signature, group_rows, head_workspace_bytes = estimate assert packed_tokens == plan.packed_tokens assert output_bytes == plan.output_bytes assert signature == plan.signature + assert group_rows == rank._plan_group_rows(plan) + assert head_workspace_bytes == rank._plan_head_workspace_bytes(plan) assert plan.signature.slot_group_count == 4 assert {group.slot_ref for group in plan.groups} == { "student", From 688de07a0caff8b10ae4cecd10d13bb0fb77e95c Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Wed, 16 Sep 2026 10:21:19 +0000 Subject: [PATCH 02/18] Narrow optional model owners before metadata inspection --- src/art/trainer_rank/_gdn_memory.py | 2 +- src/art/trainer_rank/_impl.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/art/trainer_rank/_gdn_memory.py b/src/art/trainer_rank/_gdn_memory.py index a707f091e..4ca4d5cc6 100644 --- a/src/art/trainer_rank/_gdn_memory.py +++ b/src/art/trainer_rank/_gdn_memory.py @@ -191,7 +191,7 @@ def model_shapes(rank: Any) -> tuple[int, tuple[Shape, ...]] | None: shapes = [] for layer in decoder.layers: gdn = getattr(layer, "self_attention", None) - if type(gdn) is not GatedDeltaNet: + if gdn is None or type(gdn) is not GatedDeltaNet: continue if ( getattr(gdn.forward, "__func__", None) is not _prefix_tree_forward diff --git a/src/art/trainer_rank/_impl.py b/src/art/trainer_rank/_impl.py index b9539877e..b79507414 100644 --- a/src/art/trainer_rank/_impl.py +++ b/src/art/trainer_rank/_impl.py @@ -1225,7 +1225,7 @@ def _shared_expert_output_bytes_per_token(layer: torch.nn.Module) -> int: ) shared = getattr(layer, "shared_experts", None) - if type(shared) is not SharedExpertMLP: + if shared is None or type(shared) is not SharedExpertMLP: return 0 config = getattr(layer, "config", None) shared_config = getattr(shared, "config", None) From 23ab6ae01401f1f60ecd72c4ecd1fe6f1af85242 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Wed, 16 Sep 2026 10:33:02 +0000 Subject: [PATCH 03/18] Accept ART-owned GDN norm wrappers in pending memory floor --- src/art/trainer_rank/_gdn_memory.py | 25 +++- .../unit/test_trainer_rank_pending_memory.py | 120 +++++++++++++++++- tests/unit/test_trainer_rank_shared_memory.py | 12 ++ 3 files changed, 153 insertions(+), 4 deletions(-) diff --git a/src/art/trainer_rank/_gdn_memory.py b/src/art/trainer_rank/_gdn_memory.py index 4ca4d5cc6..fc21087ff 100644 --- a/src/art/trainer_rank/_gdn_memory.py +++ b/src/art/trainer_rank/_gdn_memory.py @@ -7,6 +7,7 @@ """ from dataclasses import dataclass +from types import MethodType from typing import Any, Sequence @@ -143,7 +144,7 @@ def model_shapes(rank: Any) -> tuple[int, tuple[Shape, ...]] | None: from megatron.core.transformer.transformer_block import TransformerBlock from transformer_engine.pytorch import RMSNorm - from art.megatron.gdn.operator import _prefix_tree_forward + from art.megatron.gdn.operator import _empty_safe_norm_forward, _prefix_tree_forward from art.megatron.lora import LoRA, SelfAttentionLinearProjLoRA config = decoder.config @@ -219,13 +220,31 @@ def model_shapes(rank: Any) -> tuple[int, tuple[Shape, ...]] | None: return None hk, hv, dk, dv, kernel = dimensions conv = 2 * hk * dk + hv * dv + norm = gdn.out_norm + if type(norm) is not RMSNorm: + return None + forward = getattr(norm, "forward", None) + physical = getattr(norm, "_art_empty_safe_norm_physical_forward", None) if ( gdn.conv1d.weight.dtype is not torch.bfloat16 or tuple(gdn.conv1d.weight.shape) != (conv, 1, kernel) - or type(gdn.out_norm) is not RMSNorm or gdn.out_norm.weight.numel() != dv or gdn.out_norm.weight.dtype is not torch.bfloat16 - or "forward" in vars(gdn.out_norm) + # Original GDN setup installs this wrapper even for nonempty CP1. + # Its nonempty path delegates unchanged to the saved bound method. + or ( + "forward" in vars(norm) + and not ( + type(forward) is MethodType + and forward.__self__ is norm + and forward.__func__ is _empty_safe_norm_forward + and getattr(norm, "_art_empty_safe_norm_hooked", None) is True + and physical is not None + and type(physical) is MethodType + and physical.__self__ is norm + and physical.__func__ is RMSNorm.forward + ) + ) or gdn.out_norm._forward_hooks or gdn.out_norm._forward_pre_hooks ): diff --git a/tests/unit/test_trainer_rank_pending_memory.py b/tests/unit/test_trainer_rank_pending_memory.py index db41cb386..5448faa68 100644 --- a/tests/unit/test_trainer_rank_pending_memory.py +++ b/tests/unit/test_trainer_rank_pending_memory.py @@ -21,7 +21,7 @@ def module(cls): return obj -def rank_with_moe(moe_layer): +def rank_with_moe(moe_layer, *, install_hooks=False): from megatron.core.ssm.gated_delta_net import GatedDeltaNet from megatron.core.transformer.transformer_block import TransformerBlock from transformer_engine.pytorch import RMSNorm @@ -80,6 +80,10 @@ def rank_with_moe(moe_layer): model.config = decoder.config model.decoder = decoder model._preprocess = lambda: None + if install_hooks: + from art.megatron.gdn.operator import install_gdn_island_hooks + + install_gdn_island_hooks([model]) r: Any = TrainerRank( cast( Any, @@ -139,6 +143,120 @@ def test_actual_constructor_cache_and_full_plan(pending_rank): assert lower.required <= rank._memory_check(plan).estimated_required_bytes +def test_original_installed_norm_preserves_pending_floor(layer): + from art.megatron.gdn.operator import _empty_safe_norm_forward + + rank, gd = rank_with_moe(_enclosing_moe(layer), install_hooks=True) + norm = gd.out_norm + assert norm.forward.__func__ is _empty_safe_norm_forward + assert norm.forward.__self__ is norm + assert norm._art_empty_safe_norm_physical_forward.__func__ is type(norm).forward + assert rank._moe_output_bytes_per_token == 188416 + assert g.model_shapes(rank) is not None + plan = rank._plan_flat_forward(full_requests()) + assert g.plan_floor(rank, plan) == (8296857600, 12699148192) + assert rank._memory_check(plan).estimated_required_bytes == 23095829187 + assert rank._plan_cost(plan).required == 23095829187 + assert rank._estimate_flat_forward(full_requests()) is None + for requests in ([], full_requests(no_grad=True)): + assert g.plan_floor(rank, rank._plan_flat_forward(requests)) == (0, 0) + + +@pytest.mark.parametrize( + "mutation", + [ + "foreign wrapper", + "unbound wrapper", + "wrong wrapper self", + "spoofed wrapper", + "missing marker", + "false marker", + "integer marker", + "missing physical", + "wrong physical self", + "wrong physical function", + "recursive physical", + "unbound physical", + "spoofed physical", + "forward hook", + "pre hook", + ], +) +def test_installed_norm_rejects_changed_ownership(layer, mutation): + from art.megatron.gdn.operator import _empty_safe_norm_forward + + rank, gd = rank_with_moe(_enclosing_moe(layer), install_hooks=True) + norm = gd.out_norm + other = module(type(norm)) + physical = norm._art_empty_safe_norm_physical_forward + if mutation == "foreign wrapper": + norm.forward = MethodType(lambda self, x: x, norm) + elif mutation == "unbound wrapper": + norm.forward = _empty_safe_norm_forward + elif mutation == "wrong wrapper self": + norm.forward = MethodType(_empty_safe_norm_forward, other) + elif mutation == "spoofed wrapper": + norm.forward = SimpleNamespace(__self__=norm, __func__=_empty_safe_norm_forward) + elif mutation == "missing marker": + del norm._art_empty_safe_norm_hooked + elif mutation in ("false marker", "integer marker"): + norm._art_empty_safe_norm_hooked = False if mutation == "false marker" else 1 + elif mutation == "missing physical": + del norm._art_empty_safe_norm_physical_forward + elif mutation == "wrong physical self": + norm._art_empty_safe_norm_physical_forward = other.forward + elif mutation == "wrong physical function": + norm._art_empty_safe_norm_physical_forward = MethodType(lambda self, x: x, norm) + elif mutation == "recursive physical": + norm._art_empty_safe_norm_physical_forward = norm.forward + elif mutation == "unbound physical": + norm._art_empty_safe_norm_physical_forward = type(norm).forward + elif mutation == "spoofed physical": + norm._art_empty_safe_norm_physical_forward = SimpleNamespace( + __self__=norm, __func__=physical.__func__ + ) + elif mutation == "forward hook": + norm.register_forward_hook(lambda *args: None) + else: + norm.register_forward_pre_hook(lambda *args: None) + assert g.model_shapes(rank) is None + assert g.plan_floor(rank, rank._plan_flat_forward(full_requests())) == (0, 0) + + +def test_original_norm_wrapper_nonempty_delegation(): + from art.megatron.gdn.operator import _empty_safe_norm_forward + + # TE execution is CUDA-specific. This CPU leaf tests only the unchanged + # wrapper's delegation and original exception; it is not norm math evidence. + x, result = torch.ones(2, 128), torch.ones(2, 128) + calls = [] + original = ValueError("physical forward failed") + + def physical(value, *args, **kwargs): + calls.append((value, args, kwargs)) + if kwargs.get("fail"): + raise original + return result + + norm = SimpleNamespace(_art_empty_safe_norm_physical_forward=physical) + assert _empty_safe_norm_forward(norm, x, "argument", flag=True) is result + assert calls[0][0] is x and calls[0][1:] == (("argument",), {"flag": True}) + with pytest.raises(ValueError) as caught: + _empty_safe_norm_forward(norm, x, fail=True) + assert caught.value is original + + +def test_unsupported_norm_owner_is_not_inspected(layer): + class UnknownNorm(torch.nn.Module): + @property + def forward(self): + raise AssertionError("Unsupported owner must be rejected first") + + rank, gd = rank_with_moe(_enclosing_moe(layer), install_hooks=True) + gd.out_norm = UnknownNorm() + assert g.model_shapes(rank) is None + + def test_pending_no_grad_empty_mixed_and_learned_max(pending_rank): rank = pending_rank grad = ForwardInput( diff --git a/tests/unit/test_trainer_rank_shared_memory.py b/tests/unit/test_trainer_rank_shared_memory.py index c7dd4b6c7..d3ddeae05 100644 --- a/tests/unit/test_trainer_rank_shared_memory.py +++ b/tests/unit/test_trainer_rank_shared_memory.py @@ -125,6 +125,18 @@ def test_shared_return_in_actual_constructor_and_plan(layer, gate, no_grad): ) +@pytest.mark.parametrize("gated", [False, True]) +def test_original_norm_installation_preserves_shared_return(layer, gated): + layer = shared_layer(layer, gated) + rank, _ = rank_with_moe(layer, install_hooks=True) + assert _shared_expert_output_bytes_per_token(layer) == 4096 + assert rank._moe_output_bytes_per_token == 192512 + plan = rank._plan_flat_forward(full_requests()) + assert g.plan_floor(rank, plan) == (8296857600, 50640 * 192512 + 3157761952) + assert rank._memory_check(plan).estimated_required_bytes == 23323992771 + assert rank._plan_cost(plan).required == 23323992771 + + mutations = { "no shared": lambda x: delattr(x, "shared_experts"), "disabled shared": lambda x: setattr(x, "use_shared_expert", False), From c00a459dea67803e05444e31dc84442406f28942 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Wed, 16 Sep 2026 11:09:25 +0000 Subject: [PATCH 04/18] Price gated shared saves in checkpoint gradient workspace --- src/art/trainer_rank/_gdn_memory.py | 7 +- src/art/trainer_rank/_impl.py | 44 +++++- .../test_trainer_rank_checkpoint_memory.py | 1 + tests/unit/test_trainer_rank_shared_memory.py | 144 +++++++++++++++++- 4 files changed, 184 insertions(+), 12 deletions(-) diff --git a/src/art/trainer_rank/_gdn_memory.py b/src/art/trainer_rank/_gdn_memory.py index fc21087ff..8d728fdbf 100644 --- a/src/art/trainer_rank/_gdn_memory.py +++ b/src/art/trainer_rank/_gdn_memory.py @@ -189,6 +189,7 @@ def model_shapes(rank: Any) -> tuple[int, tuple[Shape, ...]] | None: moe = rank._moe_output_bytes_per_token if type(moe) is not int or moe < 0 or (rank._moe_layers and not moe): raise ValueError("Invalid constructor MoE coefficient for GDN pending floor") + rank._checkpoint_moe_bytes_per_token() shapes = [] for layer in decoder.layers: gdn = getattr(layer, "self_attention", None) @@ -306,6 +307,10 @@ def plan_floor(rank: Any, plan: Any) -> tuple[int, int]: retained += rows * layers * rank._hidden_size * 2 workspace = max( workspace, - *(rows * s.moe_bytes_per_row + s.pending(rows, buckets) for s in shapes), + *( + rows * rank._moe_checkpoint_grad_bytes_per_token + + s.pending(rows, buckets) + for s in shapes + ), ) return retained, workspace diff --git a/src/art/trainer_rank/_impl.py b/src/art/trainer_rank/_impl.py index b79507414..f9115a133 100644 --- a/src/art/trainer_rank/_impl.py +++ b/src/art/trainer_rank/_impl.py @@ -1361,7 +1361,10 @@ def _shared_expert_output_bytes_per_token(layer: torch.nn.Module) -> int: def _moe_output_bytes_per_token( - model: Sequence[torch.nn.Module], shape: ParallelShape + model: Sequence[torch.nn.Module], + shape: ParallelShape, + *, + checkpoint_grad: bool = False, ) -> int: """Known routed-expert working set, not a complete model/compiled bound.""" if shape != ParallelShape(tp=1, cp=1): @@ -1467,10 +1470,19 @@ def _moe_output_bytes_per_token( # remain live at the FC2 sum, including in the observed # compiled path. This is one stage, not a backward bound. features += 2 * fc2.out_features + fc1.out_features + shared = _shared_expert_output_bytes_per_token(layer) + if ( + checkpoint_grad + and shared + and getattr(layer.shared_experts, "use_shared_expert_gate", False) + is True + ): + # Gate-score backward saves a distinct pre-gate X. Charge it + # beside this layer's returned X, not another layer's maximum. + shared += shared coefficient = max( coefficient, - config.moe_router_topk * features * weights.element_size() - + _shared_expert_output_bytes_per_token(layer), + config.moe_router_topk * features * weights.element_size() + shared, ) return coefficient @@ -1547,6 +1559,14 @@ def __init__(self, runtime: TrainingRuntime) -> None: if self._moe_layers else 0 ) + # Both modes inspect original owners before dispatcher caches are installed. + self._moe_checkpoint_grad_bytes_per_token = ( + _moe_output_bytes_per_token( + runtime.model, self._parallel_shape, checkpoint_grad=True + ) + if self._moe_layers + else 0 + ) selection = select_scoring( device_capability=capability, device_memory_bytes=device_memory, @@ -3302,6 +3322,18 @@ def _plan_group_rows(self, plan: _FlatForwardPlan) -> tuple[tuple[int, bool], .. for group in plan.groups ) + def _checkpoint_moe_bytes_per_token(self) -> int: + forward = self._moe_output_bytes_per_token + gradient = self._moe_checkpoint_grad_bytes_per_token + if ( + type(forward) is not int + or forward < 0 + or type(gradient) is not int + or gradient < forward + ): + raise ValueError("Invalid constructor checkpoint MoE coefficient") + return gradient + def _checkpoint_memory_floor( self, group_rows: tuple[tuple[int, bool], ...] ) -> tuple[int, int]: @@ -3367,8 +3399,10 @@ def _checkpoint_memory_floor( ): return 0, 0 retained = gradient_rows * layers * self._hidden_size * 2 - workspace = ( - max(rows for rows, _grad in group_rows) * self._moe_output_bytes_per_token + gradient_moe = self._checkpoint_moe_bytes_per_token() + workspace = max( + rows * (gradient_moe if grad else self._moe_output_bytes_per_token) + for rows, grad in group_rows ) return retained, workspace diff --git a/tests/unit/test_trainer_rank_checkpoint_memory.py b/tests/unit/test_trainer_rank_checkpoint_memory.py index 69de26475..8cfb5856e 100644 --- a/tests/unit/test_trainer_rank_checkpoint_memory.py +++ b/tests/unit/test_trainer_rank_checkpoint_memory.py @@ -52,6 +52,7 @@ def rank(): ) ) result._moe_output_bytes_per_token = 188416 + result._moe_checkpoint_grad_bytes_per_token = 188416 return result diff --git a/tests/unit/test_trainer_rank_shared_memory.py b/tests/unit/test_trainer_rank_shared_memory.py index d3ddeae05..77790bc22 100644 --- a/tests/unit/test_trainer_rank_shared_memory.py +++ b/tests/unit/test_trainer_rank_shared_memory.py @@ -104,6 +104,8 @@ def coefficient(layer): def test_shared_return_in_actual_constructor_and_plan(layer, gate, no_grad): rank, _ = rank_with_moe(shared_layer(layer, gate)) assert rank._moe_output_bytes_per_token == 192512 + checkpoint_coefficient = 196608 if gate else 192512 + assert rank._moe_checkpoint_grad_bytes_per_token == checkpoint_coefficient shapes = g.model_shapes(rank) assert shapes is not None and shapes[1][0].moe_bytes_per_row == 192512 requests = full_requests(no_grad) @@ -116,8 +118,11 @@ def test_shared_return_in_actual_constructor_and_plan(layer, gate, no_grad): assert g.plan_floor(rank, plan) == (0, 0) assert rank._plan_cost(plan).required == 10723911264 else: - assert g.plan_floor(rank, plan) == (8296857600, 50640 * 192512 + 3157761952) - assert rank._plan_cost(plan).required == 23323992771 + assert g.plan_floor(rank, plan) == ( + 8296857600, + 50640 * checkpoint_coefficient + 3157761952, + ) + assert rank._plan_cost(plan).required == (23552156355 if gate else 23323992771) selected = rank._select_next_micro_batch(requests, 0) assert ( selected.check.estimated_required_bytes @@ -132,9 +137,15 @@ def test_original_norm_installation_preserves_shared_return(layer, gated): assert _shared_expert_output_bytes_per_token(layer) == 4096 assert rank._moe_output_bytes_per_token == 192512 plan = rank._plan_flat_forward(full_requests()) - assert g.plan_floor(rank, plan) == (8296857600, 50640 * 192512 + 3157761952) - assert rank._memory_check(plan).estimated_required_bytes == 23323992771 - assert rank._plan_cost(plan).required == 23323992771 + checkpoint_coefficient = 196608 if gated else 192512 + assert rank._moe_checkpoint_grad_bytes_per_token == checkpoint_coefficient + assert g.plan_floor(rank, plan) == ( + 8296857600, + 50640 * checkpoint_coefficient + 3157761952, + ) + expected = 23552156355 if gated else 23323992771 + assert rank._memory_check(plan).estimated_required_bytes == expected + assert rank._plan_cost(plan).required == expected mutations = { @@ -197,7 +208,15 @@ def test_unsupported_shared_branch_keeps_prior_routed_component(layer, name): mutations[name](layer) assert _shared_expert_output_bytes_per_token(layer) == 0 assert coefficient(layer) == 188416 - assert _rank(layer)._moe_output_bytes_per_token == 188416 + assert ( + _moe_output_bytes_per_token( + [layer], ParallelShape(tp=1, cp=1), checkpoint_grad=True + ) + == 188416 + ) + rank = _rank(layer) + assert rank._moe_output_bytes_per_token == 188416 + assert rank._moe_checkpoint_grad_bytes_per_token == 188416 def test_shared_return_has_no_topk_or_layer_multiplier(layer): @@ -241,3 +260,116 @@ def test_reference_group_has_shared_stage_but_no_pending_save(layer): retained, workspace = g.plan_floor(rank, one) assert g.plan_floor(rank, mixed) == (retained, max(workspace, 4096 * 192512)) assert g.plan_floor(rank, rank._plan_flat_forward([reference])) == (0, 0) + + +def test_pre_gate_is_same_layer_stage_not_sum_of_separate_maxima(layer): + import copy + + shared_layer(layer) + other = copy.deepcopy(layer) + del other.shared_experts + other.experts.linear_fc2.lora.A_T = torch.nn.Parameter( + torch.empty(256, 640, 8, dtype=torch.bfloat16) + ) + other.experts.linear_fc1.out_features = 1280 + shape = ParallelShape(tp=1, cp=1) + # Routed-only width640 wins forward; gated width512 wins recomputation. + assert _moe_output_bytes_per_token([layer, other], shape) == 194560 + assert ( + _moe_output_bytes_per_token([layer, other], shape, checkpoint_grad=True) + == 196608 + ) + assert ( + _moe_output_bytes_per_token([layer, layer], shape, checkpoint_grad=True) + == 196608 + ) + layer.config.moe_router_topk = layer.router.topk = 4 + assert ( + _moe_output_bytes_per_token([layer], shape, checkpoint_grad=True) + == 4 * (3 * 512 + 5 * 2048) * 2 + 2 * 4096 + ) + + +def test_pre_gate_cache_precedes_owned_dispatcher_and_is_checkpoint_only(layer): + rank, _ = rank_with_moe(shared_layer(layer)) + assert ( + _moe_output_bytes_per_token( + rank.runtime.model, rank._parallel_shape, checkpoint_grad=True + ) + == 0 + ) # Installed dispatcher partials must not be repriced. + assert rank._moe_checkpoint_grad_bytes_per_token == 196608 + groups = ((19, True), (23, False)) + assert rank._checkpoint_memory_floor(groups) == ( + 19 * 40 * 4096, + max(19 * 196608, 23 * 192512), + ) + for mode in (None, "selective"): + rank.runtime.model[0].decoder.config.recompute_granularity = mode + assert rank._checkpoint_memory_floor(groups) == (0, 0) + assert g.plan_floor(rank, rank._plan_flat_forward(full_requests())) == (0, 0) + + +@pytest.mark.parametrize("gradient_first", [False, True]) +def test_pre_gate_mixed_reference_and_exact_cost_mode_selection(layer, gradient_first): + rank, _ = rank_with_moe(shared_layer(layer)) + grad = ForwardInput( + input_tokens=torch.arange(67), hidden_states=True, no_grad=False + ) + reference = ForwardInput( + input_tokens=torch.arange(4096), hidden_states=True, no_grad=True + ) + single = rank._plan_flat_forward([grad]) + requests = [grad, reference] if gradient_first else [reference, grad] + mixed = rank._plan_flat_forward(requests) + retained, workspace = g.plan_floor(rank, single) + assert g.plan_floor(rank, mixed) == (retained, max(workspace, 4096 * 192512)) + assert ( + rank._memory_check(mixed).estimated_required_bytes + == rank._plan_cost(mixed).required + ) + assert rank._checkpoint_memory_floor(rank._plan_group_rows(mixed)) == ( + 67 * 40 * 4096, + 4096 * 192512, + ) + # A reference-only path must not read or validate the unused gradient cache. + rank._moe_checkpoint_grad_bytes_per_token = None + reference_plan = rank._plan_flat_forward([reference]) + assert g.plan_floor(rank, reference_plan) == (0, 0) + assert rank._checkpoint_memory_floor(rank._plan_group_rows(reference_plan)) == ( + 0, + 0, + ) + assert ( + rank._memory_check(reference_plan).estimated_required_bytes + == rank._plan_cost(reference_plan).required + ) + + +@pytest.mark.parametrize("bad", [None, True, -1, 1.5, 192511]) +def test_invalid_pre_gate_cache_stays_inside_planning_status(layer, bad): + rank, _ = rank_with_moe(shared_layer(layer)) + requests = full_requests() + plan = rank._plan_flat_forward(requests) + rank._moe_checkpoint_grad_bytes_per_token = bad + statuses = [] + rank._all_ranks_true = lambda value: (statuses.append(value), value)[1] + rank._memory_check_required = lambda *a, **kw: pytest.fail( + "memory reduction entered" + ) + for call in ( + lambda: rank._memory_check( + plan, sync_planning_errors=True, sync_across_dp=True + ), + lambda: rank._estimate_flat_forward(requests, sync_planning_errors=True), + ): + with pytest.raises( + ValueError, match="Invalid constructor checkpoint MoE coefficient" + ): + call() + assert statuses == [False] + statuses.clear() + with pytest.raises( + ValueError, match="Invalid constructor checkpoint MoE coefficient" + ): + rank._plan_cost(plan) From 5a616852734e68e520dbfc084bf549373d9d67e9 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Wed, 16 Sep 2026 14:03:41 +0000 Subject: [PATCH 05/18] Price qualified expert conversion and checkpoint return stages --- src/art/trainer_rank/_gdn_memory.py | 4 +- src/art/trainer_rank/_impl.py | 207 ++++++++++- .../test_trainer_rank_converted_memory.py | 333 ++++++++++++++++++ .../unit/test_trainer_rank_pending_memory.py | 13 +- tests/unit/test_trainer_rank_shared_memory.py | 15 +- 5 files changed, 548 insertions(+), 24 deletions(-) create mode 100644 tests/unit/test_trainer_rank_converted_memory.py diff --git a/src/art/trainer_rank/_gdn_memory.py b/src/art/trainer_rank/_gdn_memory.py index 8d728fdbf..021831250 100644 --- a/src/art/trainer_rank/_gdn_memory.py +++ b/src/art/trainer_rank/_gdn_memory.py @@ -299,7 +299,7 @@ def plan_floor(rank: Any, plan: Any) -> tuple[int, int]: if not group.grad_enabled: # Earlier gradient groups remain live during a later reference # group. Only its existing MoE component enters this stage. - workspace = max(workspace, *(rows * s.moe_bytes_per_row for s in shapes)) + workspace = max(workspace, rank._moe_workspace_bytes(rows)) continue buckets = cp1_buckets(group.packed.segments) if sum(s.length for s in group.packed.segments) != rows: @@ -308,7 +308,7 @@ def plan_floor(rank: Any, plan: Any) -> tuple[int, int]: workspace = max( workspace, *( - rows * rank._moe_checkpoint_grad_bytes_per_token + rank._moe_workspace_bytes(rows, checkpoint_grad=True) + s.pending(rows, buckets) for s in shapes ), diff --git a/src/art/trainer_rank/_impl.py b/src/art/trainer_rank/_impl.py index f9115a133..d913e083b 100644 --- a/src/art/trainer_rank/_impl.py +++ b/src/art/trainer_rank/_impl.py @@ -1360,16 +1360,55 @@ def _shared_expert_output_bytes_per_token(layer: torch.nn.Module) -> int: return hidden * 2 +def _expert_lora_weight_storage(lora: Any) -> tuple[int, int, int] | None: + """New padded weights, transposes and effective rank for the Quack path. + + Original contiguous parameters are already in the allocator baseline. This + excludes padding-concatenation temporaries and all backward GEMM workspace. + """ + from art.megatron.lora import LoRA + + a, b = getattr(lora, "A_T", None), getattr(lora, "B_T", None) + if ( + type(lora) is not LoRA + or "forward" in vars(lora) + or "active_lora_tensors" in vars(lora) + or lora._forward_hooks + or lora._forward_pre_hooks + or not isinstance(a, torch.Tensor) + or not isinstance(b, torch.Tensor) + or a.ndim != 3 + or b.ndim != 3 + or a.dtype not in (torch.float16, torch.bfloat16) + or b.dtype != a.dtype + or not a.is_contiguous() + or not b.is_contiguous() + or a.shape[0] != b.shape[0] + or a.shape[2] != b.shape[1] + or min(*a.shape, *b.shape) <= 0 + or min(a.shape[1], b.shape[2]) <= 1 + or (a.shape[2] >= 8 and a.shape[2] % 8) + ): + return None + effective = max(8, a.shape[2]) + transposes = a.shape[0] * effective * (a.shape[1] + b.shape[2]) * a.element_size() + return (transposes if a.shape[2] < 8 else 0, transposes, effective) + + def _moe_output_bytes_per_token( model: Sequence[torch.nn.Module], shape: ParallelShape, *, checkpoint_grad: bool = False, + converted_stages: list[tuple[int, int]] | None = None, ) -> int: """Known routed-expert working set, not a complete model/compiled bound.""" if shape != ParallelShape(tp=1, cp=1): return 0 - from megatron.core.extensions.transformer_engine import TERowParallelGroupedLinear + from megatron.core.extensions.transformer_engine import ( + TEColumnParallelGroupedLinear, + TERowParallelGroupedLinear, + ) from megatron.core.transformer.moe.experts import TEGroupedMLP from megatron.core.transformer.moe.moe_layer import BaseMoELayer, MoELayer from megatron.core.transformer.moe.router import TopKRouter @@ -1436,6 +1475,7 @@ def _moe_output_bytes_per_token( ): return 0 features = 2 * fc2.out_features + enclosing_fc1 = None inputs = getattr(lora, "A_T", None) if ( isinstance(inputs, torch.Tensor) @@ -1470,6 +1510,7 @@ def _moe_output_bytes_per_token( # remain live at the FC2 sum, including in the observed # compiled path. This is one stage, not a backward bound. features += 2 * fc2.out_features + fc1.out_features + enclosing_fc1 = fc1 shared = _shared_expert_output_bytes_per_token(layer) if ( checkpoint_grad @@ -1480,10 +1521,107 @@ def _moe_output_bytes_per_token( # Gate-score backward saves a distinct pre-gate X. Charge it # beside this layer's returned X, not another layer's maximum. shared += shared - coefficient = max( - coefficient, - config.moe_router_topk * features * weights.element_size() + shared, - ) + row_bytes = ( + config.moe_router_topk * features * weights.element_size() + shared + ) + coefficient = max(coefficient, row_bytes) + storage = _expert_lora_weight_storage(lora) + if converted_stages is not None and storage is not None: + padded, transposes, effective = storage + saved_fc1, rank_fc1 = 0, 0 + routed_size = config.moe_router_topk * weights.element_size() + if enclosing_fc1 is not None: + adapter = getattr(enclosing_fc1, "lora", None) + base = getattr(enclosing_fc1, "linear_fc1", None) + first = _expert_lora_weight_storage(adapter) + if ( + first is not None + and adapter is not None + and base is not None + and type(base) is TEColumnParallelGroupedLinear + and "forward" not in vars(base) + and not base._forward_hooks + and not base._forward_pre_hooks + and adapter.A_T.dtype == weights.dtype + and adapter.A_T.shape[:2] + == (weights.shape[0], fc2.out_features) + and adapter.B_T.shape[2] == enclosing_fc1.out_features + ): + first_padding, first_transposes, first_rank = first + # FC1 retains both routed H inputs and its base O1 + # while producing adapter O1. Its sum is not live yet. + converted_stages.append( + ( + routed_size + * ( + 2 * fc2.out_features + + 2 * enclosing_fc1.out_features + + first_rank + ) + + shared, + first_padding + first_transposes, + ) + ) + # At the subsequent sum, only grad-enabled execution + # retains padding/tmp; the two transposes have died. + converted_stages.append( + ( + routed_size + * ( + 2 * fc2.out_features + + 3 * enclosing_fc1.out_features + + (first_rank if checkpoint_grad else 0) + ) + + shared, + first_padding if checkpoint_grad else 0, + ) + ) + if checkpoint_grad: + saved_fc1, _, rank_fc1 = first + # At the second GEMM, the FC2 sum does not exist yet: replace + # that H with tmp. Both weight transposes are still local. + converted_stages.append( + ( + row_bytes + + routed_size * (effective + rank_fc1 - fc2.out_features), + padded + transposes + saved_fc1, + ) + ) + if checkpoint_grad: + # The transposes die at return, but padding/tmp are saved + # through backward. Exact fused FC1 saves also remain live. + converted_stages.append( + ( + row_bytes + routed_size * (effective + rank_fc1), + padded + saved_fc1, + ) + ) + if rank_fc1 and inputs is not None and inputs.shape[2] < effective: + # At FC2 backward return, nominal gradient copies + # coexist with effective gradients and FC1 saves. + # Unpadded returns alias; original parameters are not + # new storage. This is a checkpoint eager-stage floor. + experts_count, input_width, rank = inputs.shape + nominal = experts_count * rank * weights.element_size() + copies = nominal * ( + input_width + (fc2.out_features if experts_count > 1 else 0) + ) + converted_stages.append( + ( + routed_size + * ( + 2 * input_width + + 2 * fc2.out_features + + 2 * effective + + rank_fc1 + ), + padded + + transposes + + copies + + saved_fc1 + + 2 * (experts_count + 1) * 4, + ) + ) return coefficient @@ -1554,19 +1692,33 @@ def __init__(self, runtime: TrainingRuntime) -> None: self._parallel_shape = ParallelShape( tp=tp_size, cp=cp_size, ep=ep_size, etp=etp_size ) + forward_stages: list[tuple[int, int]] = [] + gradient_stages: list[tuple[int, int]] = [] self._moe_output_bytes_per_token = ( - _moe_output_bytes_per_token(runtime.model, self._parallel_shape) + _moe_output_bytes_per_token( + runtime.model, self._parallel_shape, converted_stages=forward_stages + ) if self._moe_layers else 0 ) # Both modes inspect original owners before dispatcher caches are installed. self._moe_checkpoint_grad_bytes_per_token = ( _moe_output_bytes_per_token( - runtime.model, self._parallel_shape, checkpoint_grad=True + runtime.model, + self._parallel_shape, + checkpoint_grad=True, + converted_stages=gradient_stages, ) if self._moe_layers else 0 ) + # Discard partial walks if a later layer has an unsupported owner. + self._moe_forward_stages = ( + tuple(forward_stages) if self._moe_output_bytes_per_token else () + ) + self._moe_gradient_stages = ( + tuple(gradient_stages) if self._moe_checkpoint_grad_bytes_per_token else () + ) selection = select_scoring( device_capability=capability, device_memory_bytes=device_memory, @@ -3334,6 +3486,39 @@ def _checkpoint_moe_bytes_per_token(self) -> int: raise ValueError("Invalid constructor checkpoint MoE coefficient") return gradient + def _moe_workspace_bytes(self, rows: int, *, checkpoint_grad: bool = False) -> int: + """Maximum of same-layer affine stages, not a retained multi-layer bank. + + Cached at construction before dispatcher wrapping; model/slot shapes + must remain unchanged, as for the existing row coefficient. Ordinary + non-checkpoint gradients retain only the prior forward-stage coverage. + """ + coefficient = ( + self._checkpoint_moe_bytes_per_token() + if checkpoint_grad + else self._moe_output_bytes_per_token + ) + stages = getattr( + self, + "_moe_gradient_stages" if checkpoint_grad else "_moe_forward_stages", + (), + ) + if type(stages) is not tuple or any( + type(stage) is not tuple + or len(stage) != 2 + or any(type(value) is not int or value < 0 for value in stage) + for stage in stages + ): + raise ValueError("Invalid constructor converted-weight stages") + return ( + max( + rows * coefficient, + *(rows * per_row + fixed for per_row, fixed in stages), + ) + if stages and rows > 0 + else rows * coefficient + ) + def _checkpoint_memory_floor( self, group_rows: tuple[tuple[int, bool], ...] ) -> tuple[int, int]: @@ -3399,9 +3584,9 @@ def _checkpoint_memory_floor( ): return 0, 0 retained = gradient_rows * layers * self._hidden_size * 2 - gradient_moe = self._checkpoint_moe_bytes_per_token() + self._checkpoint_moe_bytes_per_token() workspace = max( - rows * (gradient_moe if grad else self._moe_output_bytes_per_token) + self._moe_workspace_bytes(rows, checkpoint_grad=grad) for rows, grad in group_rows ) return retained, workspace @@ -5728,9 +5913,7 @@ def _estimate_required_memory_bytes_from_values( ) # Groups execute sequentially: summed packed rows conservatively bound # this FC2 component, not all workspace or retained graphs. - static_compute = max( - static_compute, packed_tokens * self._moe_output_bytes_per_token - ) + static_compute = max(static_compute, self._moe_workspace_bytes(packed_tokens)) # 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. diff --git a/tests/unit/test_trainer_rank_converted_memory.py b/tests/unit/test_trainer_rank_converted_memory.py new file mode 100644 index 000000000..a0e61ed8e --- /dev/null +++ b/tests/unit/test_trainer_rank_converted_memory.py @@ -0,0 +1,333 @@ +"""Source-derived affine routed-expert stages; no complete backward/compiled bound.""" + +from types import SimpleNamespace +from typing import Any, cast + +import pytest +from test_trainer_rank_moe_memory import _enclosing_moe, _rank +from test_trainer_rank_moe_memory import layer as layer +from test_trainer_rank_pending_memory import module, rank_with_moe +import torch + +from art.trainer_rank import ForwardInput, _gdn_memory +from art.trainer_rank._impl import _expert_lora_weight_storage + + +def weights(layer: Any, rank: int, *, fc1: bool = True, dtype=torch.bfloat16): + from art.megatron.lora import LoRA, TEColumnParallelGroupedLinear + + _enclosing_moe(layer) + for fc, inputs, outputs in ( + (layer.experts.linear_fc2, 512, 2048), + *(([(layer.experts.linear_fc1, 2048, 1024)]) if fc1 else []), + ): + fc.lora = module(LoRA) + fc.lora.A_T = torch.nn.Parameter(torch.empty(256, inputs, rank, dtype=dtype)) + fc.lora.B_T = torch.nn.Parameter(torch.empty(256, rank, outputs, dtype=dtype)) + if fc1: + layer.experts.linear_fc1.linear_fc1 = module(TEColumnParallelGroupedLinear) + return layer + + +def expected(rows, rank, grad, *, fc1=True, shared=0): + effective = max(8, rank) + t2 = 256 * effective * (512 + 2048) * 2 + p2 = t2 if rank < 8 else 0 + p1 = 256 * effective * (2048 + 1024) * 2 if rank < 8 and fc1 and grad else 0 + r1 = effective if fc1 and grad else 0 + inner = rows * (188416 + 16 * (effective + r1 - 2048)) + p2 + t2 + p1 + summed = rows * (188416 + 16 * (effective + r1)) + p2 + p1 + first_stage = ( + rows * 16 * (2 * 2048 + 2 * 1024 + effective) + + (2 if rank < 8 else 1) * 256 * effective * (2048 + 1024) * 2 + ) + first_sum = rows * 16 * (2 * 2048 + 3 * 1024 + (effective if grad else 0)) + p1 + backward_return = ( + rows * 16 * (2 * 512 + 2 * 2048 + 3 * effective) + + p1 + + p2 + + t2 + + 256 * rank * 2560 * 2 + + 2 * 257 * 4 + ) + return max( + backward_return if grad and fc1 and rank < 8 else 0, + rows * 188416 + shared, + inner + shared, + summed + shared if grad else 0, + first_stage + shared if fc1 else 0, + first_sum + shared if fc1 else 0, + ) + + +@pytest.mark.parametrize("rank_value", [1, 7, 8, 16]) +@pytest.mark.parametrize("grad", [False, True]) +def test_same_stage_crossover_and_constructor(layer, rank_value, grad): + rank, _ = rank_with_moe(weights(layer, rank_value)) + for rows in (1, 8, 64, 128, 512, 50640): + assert rank._moe_workspace_bytes(rows, checkpoint_grad=grad) == expected( + rows, rank_value, grad + ) + assert rank._moe_workspace_bytes(0, checkpoint_grad=grad) == 0 + assert rank._moe_workspace_bytes(1, checkpoint_grad=grad) > 188416 + if not grad: + assert rank._moe_workspace_bytes(50640) == 50640 * 188416 + # Metadata was cached before the original dispatcher partial is installed. + assert "dispatch_preprocess" in vars(layer.token_dispatcher) + assert rank._moe_workspace_bytes(1, checkpoint_grad=grad) == expected( + 1, rank_value, grad + ) + + +@pytest.mark.parametrize("rank_value", [1, 7, 8, 16]) +@pytest.mark.parametrize("grad", [False, True]) +@pytest.mark.parametrize("output", ["hidden", "logprob", "both"]) +def test_actual_plan_cost_and_admission(layer, rank_value, grad, output): + rank, _ = rank_with_moe(weights(layer, rank_value)) + request = ForwardInput( + input_tokens=torch.arange(8), + no_grad=not grad, + hidden_states=output != "logprob", + target_tokens=torch.arange(8) if output != "hidden" else None, + ) + plan = rank._plan_flat_forward([request]) + required = rank._memory_check(plan).estimated_required_bytes + assert required == rank._plan_cost(plan).required + retained, workspace = rank._checkpoint_memory_floor(rank._plan_group_rows(plan)) + pending = _gdn_memory.plan_floor(rank, plan) + if grad: + assert workspace == expected(8, rank_value, True) + assert pending[0] == retained == 8 * 40 * 2048 * 2 + assert pending[1] >= workspace + else: + assert (retained, workspace) == pending == (0, 0) + assert required >= int((plan.output_bytes + expected(8, rank_value, grad)) * 1.1) + rank._available_memory_bytes = lambda: required - 1 + assert not rank._memory_check(plan).fits + + +@pytest.mark.parametrize("order", [False, True]) +@pytest.mark.parametrize("rank_value", [1, 7]) +def test_reference_and_gradient_keep_distinct_stage_modes(layer, order, rank_value): + rank, _ = rank_with_moe(weights(layer, rank_value)) + requests = [ + ForwardInput(input_tokens=torch.arange(3), hidden_states=True), + ForwardInput( + input_tokens=torch.arange(9) + 100, hidden_states=True, no_grad=True + ), + ] + if order: + requests.reverse() + plan = rank._plan_flat_forward(requests) + groups = rank._plan_group_rows(plan) + assert set(groups) == {(3, True), (9, False)} + retained, workspace = rank._checkpoint_memory_floor(groups) + assert retained == 3 * 40 * 2048 * 2 + assert workspace == max( + expected(3, rank_value, True), expected(9, rank_value, False) + ) + assert ( + rank._memory_check(plan).estimated_required_bytes + == rank._plan_cost(plan).required + ) + + +@pytest.mark.parametrize("rank_value", [1, 7, 8, 16]) +def test_alias_original_parameters_and_dtype(layer, rank_value): + weights(layer, rank_value, dtype=torch.float16) + storage = _expert_lora_weight_storage(layer.experts.linear_fc2.lora) + assert storage is not None + padding, transposes, effective = storage + assert effective == max(8, rank_value) + assert transposes == 256 * effective * 2560 * 2 + assert padding == (transposes if rank_value < 8 else 0) + assert _rank(layer)._moe_workspace_bytes(1) == expected(1, rank_value, False) + + +@pytest.mark.parametrize( + "mutation", ["rank9", "noncontiguous", "owner", "hook", "missing"] +) +def test_unsupported_conversion_keeps_prior_component(layer, mutation): + weights(layer, 1, fc1=False) + lora = layer.experts.linear_fc2.lora + if mutation == "rank9": + weights(layer, 9, fc1=False) + elif mutation == "noncontiguous": + lora.A_T = torch.nn.Parameter( + torch.empty(256, 1, 512, dtype=torch.bfloat16).transpose(1, 2) + ) + # Rank-one transpose is contiguous; use a genuine noncontiguous slice. + lora.A_T = torch.nn.Parameter( + torch.empty(256, 512, 2, dtype=torch.bfloat16)[..., :1] + ) + elif mutation == "owner": + lora.forward = lambda *args: None + elif mutation == "hook": + lora.register_forward_hook(lambda *args: None) + else: + lora.A_T = None + rank = _rank(layer) + assert rank._moe_forward_stages == rank._moe_gradient_stages == () + assert rank._moe_workspace_bytes(1) == rank._moe_output_bytes_per_token + + +@pytest.mark.parametrize("bad", [None, [], ((True, 1),), ((1, -1),), ((1, 2, 3),)]) +def test_corrupted_cache_refuses_before_memory_reduction(layer, bad, monkeypatch): + rank, _ = rank_with_moe(weights(layer, 1)) + plan = rank._plan_flat_forward( + [ForwardInput(input_tokens=torch.arange(1), hidden_states=True)] + ) + rank._moe_forward_stages = bad + + def reduce(*args, **kwargs): + raise AssertionError("entered memory reduction before local planning failed") + + monkeypatch.setattr(rank, "_memory_check_required", reduce) + with pytest.raises(ValueError, match="converted-weight"): + rank._memory_check(plan) + + +def test_missing_fc1_metadata_does_not_invent_saved_bank(layer): + rank, _ = rank_with_moe(weights(layer, 1, fc1=False)) + assert rank._moe_workspace_bytes(1, checkpoint_grad=True) == expected( + 1, 1, True, fc1=False + ) + + +def test_heterogeneous_joint_stage_max_not_separate_maxima(layer): + from test_trainer_rank_moe_memory import layer as factory + + first = weights(layer, 1) + second = weights(cast(Any, factory).__wrapped__(), 16) + second.config.moe_router_topk = second.router.topk = 1 + model = torch.nn.ModuleList([first, second]) + rank = _rank(model) + one = _rank(weights(cast(Any, factory).__wrapped__(), 1)) + other = weights(cast(Any, factory).__wrapped__(), 16) + other.config.moe_router_topk = other.router.topk = 1 + two = _rank(other) + for n in (1, 64, 50640): + assert rank._moe_workspace_bytes(n) == max( + one._moe_workspace_bytes(n), two._moe_workspace_bytes(n) + ) + + +@pytest.mark.parametrize( + "mutation", ["base owner", "base hook", "adapter override", "shape"] +) +def test_unqualified_fc1_saves_are_not_invented(layer, mutation): + weights(layer, 1) + fc1 = layer.experts.linear_fc1 + if mutation == "base owner": + fc1.linear_fc1 = torch.nn.Identity() + elif mutation == "base hook": + fc1.linear_fc1.register_forward_pre_hook(lambda *args: None) + elif mutation == "adapter override": + fc1.lora.active_lora_tensors = lambda: None + else: + fc1.lora.A_T = torch.nn.Parameter( + torch.empty(256, 2047, 1, dtype=torch.bfloat16) + ) + rank, _ = rank_with_moe(layer) + assert rank._moe_workspace_bytes(1, checkpoint_grad=True) == expected( + 1, 1, True, fc1=False + ) + + +def test_other_checkpoint_modes_remain_partial_forward_only(layer): + rank, _ = rank_with_moe(weights(layer, 1)) + rank.runtime.model[0].decoder.config.recompute_granularity = None + p = rank._plan_flat_forward( + [ForwardInput(input_tokens=torch.arange(1), hidden_states=True)] + ) + assert rank._checkpoint_memory_floor(rank._plan_group_rows(p)) == (0, 0) + assert _gdn_memory.plan_floor(rank, p) == (0, 0) + assert rank._memory_check(p).estimated_required_bytes == int( + (p.output_bytes + expected(1, 1, False)) * 1.1 + ) + + +@pytest.mark.parametrize("topk", [1, 4, 8]) +def test_fc1_fixed_weights_do_not_scale_with_topk(layer, topk): + weights(layer, 1) + layer.config.moe_router_topk = layer.router.topk = topk + rank = _rank(layer) + for rows in (1, 64, 1024): + first_inner = rows * topk * (2 * 2048 + 2 * 1024 + 8) * 2 + 25165824 + second_inner = rows * topk * (4 * 2048 + 3 * 512 + 8) * 2 + 20971520 + original = rows * topk * (5 * 2048 + 3 * 512) * 2 + assert rank._moe_workspace_bytes(rows) == max( + first_inner, second_inner, original + ) + + +def test_wide_fc1_sum_is_a_separate_stage(layer): + weights(layer, 8) + experts = layer.experts + experts.linear_fc1.out_features = 32768 + layer.config.moe_ffn_hidden_size = 16384 + layer.token_dispatcher.num_local_experts = 16 + layer.config.num_moe_experts = 16 + for fc, inputs, outputs in ( + (experts.linear_fc1, 2048, 32768), + (experts.linear_fc2, 16384, 2048), + ): + fc.lora.A_T = torch.nn.Parameter( + torch.empty(16, inputs, 8, dtype=torch.bfloat16) + ) + fc.lora.B_T = torch.nn.Parameter( + torch.empty(16, 8, outputs, dtype=torch.bfloat16) + ) + rank = _rank(layer) + # Large N crosses from fixed conversion weights to the three simultaneous + # 2F outputs at the original eager FC1 sum; this is not an FC2 coefficient. + for grad in (False, True): + expected_sum = 1024 * 8 * (2 * 2048 + 3 * 32768 + (8 if grad else 0)) * 2 + assert rank._moe_workspace_bytes(1024, checkpoint_grad=grad) == expected_sum + + +@pytest.mark.parametrize("gated", [False, True]) +@pytest.mark.parametrize("grad", [False, True]) +def test_fc1_stage_keeps_same_layer_shared_output(layer, gated, grad): + from test_trainer_rank_shared_memory import shared_layer + + rank, _ = rank_with_moe(weights(shared_layer(layer, gated), 1)) + for rows in (1, 64, 50640): + assert rank._moe_workspace_bytes(rows, checkpoint_grad=grad) == expected( + rows, 1, grad, shared=rows * 4096 * (2 if gated and grad else 1) + ) + + +@pytest.mark.parametrize("rows,known", [(1, 42813832), (8, 43389960)]) +def test_rank7_original_return_storage_fits_total_admission(layer, rows, known): + # Actual-source CPU return witness counts distinct storage, excluding + # parameters and all speculative GDN/boundary/compiled terms. + rank, _ = rank_with_moe(weights(layer, 7)) + plan = rank._plan_flat_forward( + [ForwardInput(input_tokens=torch.arange(rows), hidden_states=True)] + ) + assert rank._moe_workspace_bytes(rows, checkpoint_grad=True) == known + assert rank._memory_check(plan).estimated_required_bytes >= known + assert rank._plan_cost(plan).required >= known + + +@pytest.mark.parametrize("rank_value", [1, 7, 8, 16]) +def test_backward_return_stage_is_checkpoint_only_and_not_shared_max(layer, rank_value): + from test_trainer_rank_shared_memory import shared_layer + + rank, _ = rank_with_moe(weights(shared_layer(layer, True), rank_value)) + assert len(rank._moe_gradient_stages) == (5 if rank_value < 8 else 4) + assert len(rank._moe_forward_stages) == 3 + if rank_value < 8: + # The witnessed backward stage excludes unproved shared-output lifetime. + assert ( + 82304, + 33554432 + 256 * rank_value * 2560 * 2 + 2056, + ) in rank._moe_gradient_stages + for rows in (1, 8, 50640): + assert rank._moe_workspace_bytes(rows, checkpoint_grad=True) == expected( + rows, rank_value, True, shared=rows * 8192 + ) + assert rank._moe_workspace_bytes(rows) == expected( + rows, rank_value, False, shared=rows * 4096 + ) diff --git a/tests/unit/test_trainer_rank_pending_memory.py b/tests/unit/test_trainer_rank_pending_memory.py index 5448faa68..c08ac751e 100644 --- a/tests/unit/test_trainer_rank_pending_memory.py +++ b/tests/unit/test_trainer_rank_pending_memory.py @@ -126,11 +126,14 @@ def test_actual_constructor_cache_and_full_plan(pending_rank): assert ( plan.packed_tokens == plan.logical_tokens == 50640 and plan.request_count == 8 ) - assert g.plan_floor(rank, plan) == (8296857600, 9541386240 + 3157761952) + assert g.plan_floor(rank, plan) == ( + 8296857600, + 9541386240 + 50640 * 128 + 3157761952, + ) assert ( rank._memory_check(plan).estimated_required_bytes == rank._plan_cost(plan).required - == 23095829187 + == 23102959299 ) selected = rank._select_next_micro_batch(requests, 0) assert ( @@ -154,9 +157,9 @@ def test_original_installed_norm_preserves_pending_floor(layer): assert rank._moe_output_bytes_per_token == 188416 assert g.model_shapes(rank) is not None plan = rank._plan_flat_forward(full_requests()) - assert g.plan_floor(rank, plan) == (8296857600, 12699148192) - assert rank._memory_check(plan).estimated_required_bytes == 23095829187 - assert rank._plan_cost(plan).required == 23095829187 + assert g.plan_floor(rank, plan) == (8296857600, 12705630112) + assert rank._memory_check(plan).estimated_required_bytes == 23102959299 + assert rank._plan_cost(plan).required == 23102959299 assert rank._estimate_flat_forward(full_requests()) is None for requests in ([], full_requests(no_grad=True)): assert g.plan_floor(rank, rank._plan_flat_forward(requests)) == (0, 0) diff --git a/tests/unit/test_trainer_rank_shared_memory.py b/tests/unit/test_trainer_rank_shared_memory.py index 77790bc22..909715594 100644 --- a/tests/unit/test_trainer_rank_shared_memory.py +++ b/tests/unit/test_trainer_rank_shared_memory.py @@ -120,9 +120,9 @@ def test_shared_return_in_actual_constructor_and_plan(layer, gate, no_grad): else: assert g.plan_floor(rank, plan) == ( 8296857600, - 50640 * checkpoint_coefficient + 3157761952, + 50640 * (checkpoint_coefficient + 128) + 3157761952, ) - assert rank._plan_cost(plan).required == (23552156355 if gate else 23323992771) + assert rank._plan_cost(plan).required == (23559286467 if gate else 23331122883) selected = rank._select_next_micro_batch(requests, 0) assert ( selected.check.estimated_required_bytes @@ -141,9 +141,9 @@ def test_original_norm_installation_preserves_shared_return(layer, gated): assert rank._moe_checkpoint_grad_bytes_per_token == checkpoint_coefficient assert g.plan_floor(rank, plan) == ( 8296857600, - 50640 * checkpoint_coefficient + 3157761952, + 50640 * (checkpoint_coefficient + 128) + 3157761952, ) - expected = 23552156355 if gated else 23323992771 + expected = 23559286467 if gated else 23331122883 assert rank._memory_check(plan).estimated_required_bytes == expected assert rank._plan_cost(plan).required == expected @@ -302,7 +302,12 @@ def test_pre_gate_cache_precedes_owned_dispatcher_and_is_checkpoint_only(layer): groups = ((19, True), (23, False)) assert rank._checkpoint_memory_floor(groups) == ( 19 * 40 * 4096, - max(19 * 196608, 23 * 192512), + max( + 19 * (196608 + 128), + 23 * 192512, + 19 * (196608 - 32768 + 128) + 10485760, + 23 * (192512 - 32768 + 128) + 10485760, + ), ) for mode in (None, "selective"): rank.runtime.model[0].decoder.config.recompute_granularity = mode From 0c708abc16d208b61aa0ec844b4e29fd1ae9ca00 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Wed, 16 Sep 2026 14:24:50 +0000 Subject: [PATCH 06/18] Run converted-memory regressions in Megatron CI environment --- .github/workflows/prek.yml | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/.github/workflows/prek.yml b/.github/workflows/prek.yml index dfb27d681..195ff33fc 100644 --- a/.github/workflows/prek.yml +++ b/.github/workflows/prek.yml @@ -231,6 +231,7 @@ jobs: tests/unit/test_trainer_rank_ignored_mixed_head.py \ tests/unit/test_trainer_rank_pending_memory.py \ tests/unit/test_trainer_rank_shared_memory.py \ + tests/unit/test_trainer_rank_converted_memory.py \ tests/unit/test_trainer_rank_split.py \ tests/acceptance/trainer_rank_planner \ tests/integration/megatron/model_support/test_dispatcher_graph_retention.py \ @@ -260,4 +261,5 @@ jobs: --ignore=tests/unit/test_trainer_rank_mixed_head_memory.py \ --ignore=tests/unit/test_trainer_rank_ignored_mixed_head.py \ --ignore=tests/unit/test_trainer_rank_pending_memory.py \ - --ignore=tests/unit/test_trainer_rank_shared_memory.py + --ignore=tests/unit/test_trainer_rank_shared_memory.py \ + --ignore=tests/unit/test_trainer_rank_converted_memory.py From 682a16a7079baaa11561555a97e374a45256a0c4 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Thu, 17 Sep 2026 05:24:57 +0000 Subject: [PATCH 07/18] Fix admission test formatting and flat-plan typing --- .../test_trainer_rank_admission_inputs.py | 3 +- .../test_trainer_rank_checkpoint_memory.py | 53 ++++++++++++++----- 2 files changed, 42 insertions(+), 14 deletions(-) diff --git a/tests/unit/test_trainer_rank_admission_inputs.py b/tests/unit/test_trainer_rank_admission_inputs.py index 78c8ae176..92b5dbde4 100644 --- a/tests/unit/test_trainer_rank_admission_inputs.py +++ b/tests/unit/test_trainer_rank_admission_inputs.py @@ -10,7 +10,7 @@ import torch from art.trainer_rank import ForwardInput, _gdn_memory -from art.trainer_rank._impl import Unset +from art.trainer_rank._impl import Unset, _FlatForwardPlan def record_prices(monkeypatch, rank): @@ -122,6 +122,7 @@ def test_cp_gdn_segments_groups_and_retained_tokens_reach_exact_search(monkeypat # The outer sequence contains one multi-request wave. A flat list would # instead let width search select separate top-level requests. selected = rank._search_next_micro_batch([requests], 0) + assert isinstance(selected.plan, _FlatForwardPlan) assert selected.check.fits and selected.plan.grad_segment_count == 2 assert selected.check.estimated_required_bytes == cost.required assert calls diff --git a/tests/unit/test_trainer_rank_checkpoint_memory.py b/tests/unit/test_trainer_rank_checkpoint_memory.py index 3142e1ebf..51998b2eb 100644 --- a/tests/unit/test_trainer_rank_checkpoint_memory.py +++ b/tests/unit/test_trainer_rank_checkpoint_memory.py @@ -140,7 +140,8 @@ def test_no_grad_enclosure_empty_and_unsupported(): r = rank() assert r._checkpoint_memory_floor(()) == (0, 0) assert r._checkpoint_memory_floor(((8192, False),)) == ( - 0, 8192 * (188416 + 4 * 2048 * 2) + 0, + 8192 * (188416 + 4 * 2048 * 2), ) values = r._estimate_flat_forward(requests()) baseline = price(r, values).required @@ -302,19 +303,34 @@ def test_malformed_flag_types_do_not_claim_supported_schedule(field, value): @pytest.mark.parametrize("profile_rate", [None, 1, 1_000_000]) def test_no_grad_enclosure_exact_lower_and_profile(profile_rate): r = rank() - req = [ForwardInput(input_tokens=torch.arange(17), hidden_states=True, no_grad=True)] + req = [ + ForwardInput(input_tokens=torch.arange(17), hidden_states=True, no_grad=True) + ] values = r._estimate_flat_forward(req, exact=True) n, out, sig, groups, _ = values if profile_rate is not None: r._memory_profiles[sig] = _MemoryProfile( - bytes_per_token=profile_rate, packed_tokens=n, logical_per_packed=1, + bytes_per_token=profile_rate, + packed_tokens=n, + logical_per_packed=1, ) - expected = int((out + max(n * (188416 + 4 * 2048 * 2), n * (profile_rate or 0))) * 1.1) + expected = int( + (out + max(n * (188416 + 4 * 2048 * 2), n * (profile_rate or 0))) * 1.1 + ) plan = r._plan_flat_forward(req) assert groups == r._plan_group_rows(plan) == ((n, False),) assert r._checkpoint_memory_floor(groups) == (0, n * (188416 + 4 * 2048 * 2)) - assert price(r, values).required == r._memory_check(plan).estimated_required_bytes == expected - assert r._split_chunk_lower_cost(req, tuple(x.input_tokens for x in req), checkpoint=Unset).required == expected + assert ( + price(r, values).required + == r._memory_check(plan).estimated_required_bytes + == expected + ) + assert ( + r._split_chunk_lower_cost( + req, tuple(x.input_tokens for x in req), checkpoint=Unset + ).required + == expected + ) r._available_memory_bytes = lambda: expected - 1 assert not r._memory_check(plan).fits r._available_memory_bytes = lambda: expected @@ -326,19 +342,30 @@ def test_no_grad_enclosure_uses_max_group_and_affine_stage(): r._moe_forward_stages = ((1, 1_000_000),) groups = ((3, False), (11, False)) assert r._checkpoint_memory_floor(groups) == ( - 0, max(max(rows * 188416, rows + 1_000_000) + 4 * rows * 2048 * 2 for rows, _ in groups) + 0, + max( + max(rows * 188416, rows + 1_000_000) + 4 * rows * 2048 * 2 + for rows, _ in groups + ), ) mixed = ((3, True), (11, False)) assert r._checkpoint_memory_floor(mixed) == ( - 3 * 40 * 2048 * 2, max(3 * 188416, 11 * 188416, 11 + 1_000_000) + 3 * 40 * 2048 * 2, + max(3 * 188416, 11 * 188416, 11 + 1_000_000), ) -@pytest.mark.parametrize("field,value", [ - ("recompute_granularity", None), ("recompute_granularity", "selective"), - ("recompute_method", "block"), ("recompute_num_layers", True), - ("cpu_offloading", True), ("params_dtype", torch.float32), -]) +@pytest.mark.parametrize( + "field,value", + [ + ("recompute_granularity", None), + ("recompute_granularity", "selective"), + ("recompute_method", "block"), + ("recompute_num_layers", True), + ("cpu_offloading", True), + ("params_dtype", torch.float32), + ], +) def test_no_grad_enclosure_config_guard(field, value): r = rank() setattr(r.runtime.model[0].decoder.config, field, value) From d8df85ba2de104433ce0244508c65e8c3bb64014 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Thu, 17 Sep 2026 05:36:33 +0000 Subject: [PATCH 08/18] Avoid gradient-only cache validation for reference forwards --- src/art/trainer_rank/_impl.py | 3 ++- tests/unit/test_trainer_rank_shared_memory.py | 9 +++++++-- 2 files changed, 9 insertions(+), 3 deletions(-) diff --git a/src/art/trainer_rank/_impl.py b/src/art/trainer_rank/_impl.py index 6a0e28eaf..9bfa82d44 100644 --- a/src/art/trainer_rank/_impl.py +++ b/src/art/trainer_rank/_impl.py @@ -3653,7 +3653,8 @@ def _checkpoint_memory_floor( ): return 0, 0 retained = gradient_rows * layers * self._hidden_size * 2 - self._checkpoint_moe_bytes_per_token() + if gradient_rows: + self._checkpoint_moe_bytes_per_token() workspace = max( self._moe_workspace_bytes(rows, checkpoint_grad=grad) + (0 if gradient_rows else 4 * rows * self._hidden_size * 2) diff --git a/tests/unit/test_trainer_rank_shared_memory.py b/tests/unit/test_trainer_rank_shared_memory.py index 909715594..400e7b204 100644 --- a/tests/unit/test_trainer_rank_shared_memory.py +++ b/tests/unit/test_trainer_rank_shared_memory.py @@ -116,7 +116,11 @@ def test_shared_return_in_actual_constructor_and_plan(layer, gate, no_grad): ) if no_grad: assert g.plan_floor(rank, plan) == (0, 0) - assert rank._plan_cost(plan).required == 10723911264 + assert rank._checkpoint_memory_floor(rank._plan_group_rows(plan)) == ( + 0, + 50640 * (192512 + 4 * 2048 * 2), + ) + assert rank._plan_cost(plan).required == 11636565600 else: assert g.plan_floor(rank, plan) == ( 8296857600, @@ -339,11 +343,12 @@ def test_pre_gate_mixed_reference_and_exact_cost_mode_selection(layer, gradient_ ) # A reference-only path must not read or validate the unused gradient cache. rank._moe_checkpoint_grad_bytes_per_token = None + rank._moe_gradient_stages = None reference_plan = rank._plan_flat_forward([reference]) assert g.plan_floor(rank, reference_plan) == (0, 0) assert rank._checkpoint_memory_floor(rank._plan_group_rows(reference_plan)) == ( 0, - 0, + 4096 * (192512 + 4 * 2048 * 2), ) assert ( rank._memory_check(reference_plan).estimated_required_bytes From 05e9b3d07207ef3f065f9f18197cbbc68fc9030c Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Thu, 17 Sep 2026 15:39:15 +0000 Subject: [PATCH 09/18] Price mixed reference groups and selected checkpoint LoRA layouts --- src/art/trainer_rank/_gdn_memory.py | 41 ++- src/art/trainer_rank/_impl.py | 171 +++++++++++-- .../test_trainer_rank_checkpoint_memory.py | 63 ++++- .../test_trainer_rank_converted_memory.py | 2 +- tests/unit/test_trainer_rank_shared_memory.py | 6 +- tests/unit/test_trainer_rank_slot_memory.py | 240 ++++++++++++++++++ 6 files changed, 478 insertions(+), 45 deletions(-) create mode 100644 tests/unit/test_trainer_rank_slot_memory.py diff --git a/src/art/trainer_rank/_gdn_memory.py b/src/art/trainer_rank/_gdn_memory.py index 021831250..13db19a55 100644 --- a/src/art/trainer_rank/_gdn_memory.py +++ b/src/art/trainer_rank/_gdn_memory.py @@ -120,13 +120,19 @@ def pending(self, packed_rows: int, buckets: tuple[Bucket, ...]) -> int: ) -def model_shapes(rank: Any) -> tuple[int, tuple[Shape, ...]] | None: +def model_shapes( + rank: Any, slot_ref: Any = None +) -> tuple[int, tuple[Shape, ...]] | None: """Conditional original-owner metadata; no model execution or CUDA read.""" if not getattr(rank, "_gdn_layers", 0): return None import torch - from art.trainer_rank._impl import _expert_parallel_shape, _language_model + from art.trainer_rank._impl import ( + _expert_parallel_shape, + _language_model, + _slot_lora_tensors, + ) if ( len(rank.runtime.model) != 1 @@ -257,13 +263,23 @@ def model_shapes(rank: Any) -> tuple[int, tuple[Shape, ...]] | None: if ( "forward" in vars(out) or "forward" in vars(lora) + or "active_lora_tensors" in vars(lora) + or "_slot" in vars(lora) or out._forward_hooks or lora._forward_hooks or out._forward_pre_hooks or lora._forward_pre_hooks ): return None - a, b = lora.A_T, lora.B_T + # Keep the previous conservative inactive-adapter enclosure. + tensors = _slot_lora_tensors( + lora, + None if slot_ref is not None and slot_ref.name is None else slot_ref, + ) + if tensors is None: + shapes.append(Shape(hk, hv, dk, dv, kernel, 0, moe)) + continue + a, b = tensors if ( a.ndim != 2 or b.ndim != 2 @@ -273,8 +289,7 @@ def model_shapes(rank: Any) -> tuple[int, tuple[Shape, ...]] | None: or a.shape[1] != b.shape[0] ): return None - # Slots share these declared shapes; charging an inactive adapter - # conservatively adds a term, without inspecting/changing its slot. + # Use the admitted slot rank without changing its active context. lora_rank = int(a.shape[1]) shapes.append(Shape(hk, hv, dk, dv, kernel, lora_rank, moe)) return (len(decoder.layers), tuple(shapes)) if shapes else None @@ -288,18 +303,20 @@ def plan_floor(rank: Any, plan: Any) -> tuple[int, int]: gradients = [g for g in plan.groups if g.grad_enabled] if not gradients: return 0, 0 - model = model_shapes(rank) - if model is None: - return 0, 0 - layers, shapes = model retained = 0 workspace = 0 for group in plan.groups: rows = int(group.packed.tokens.numel()) + model = model_shapes(rank, group.slot_ref) + if model is None: + return 0, 0 + layers, shapes = model if not group.grad_enabled: # Earlier gradient groups remain live during a later reference # group. Only its existing MoE component enters this stage. - workspace = max(workspace, rank._moe_workspace_bytes(rows)) + workspace = max( + workspace, rank._moe_workspace_bytes(rows, slot_ref=group.slot_ref) + ) continue buckets = cp1_buckets(group.packed.segments) if sum(s.length for s in group.packed.segments) != rows: @@ -308,7 +325,9 @@ def plan_floor(rank: Any, plan: Any) -> tuple[int, int]: workspace = max( workspace, *( - rank._moe_workspace_bytes(rows, checkpoint_grad=True) + rank._moe_workspace_bytes( + rows, checkpoint_grad=True, slot_ref=group.slot_ref + ) + s.pending(rows, buckets) for s in shapes ), diff --git a/src/art/trainer_rank/_impl.py b/src/art/trainer_rank/_impl.py index 8c0c54c12..1ecb9294b 100644 --- a/src/art/trainer_rank/_impl.py +++ b/src/art/trainer_rank/_impl.py @@ -925,6 +925,7 @@ class _MemorySignature: request_mix: tuple[str, ...] grad_enabled: bool grad_modes: tuple[bool, ...] + slot_shapes: tuple[tuple[bool, tuple[tuple[int, ...], ...]], ...] = () @dataclass(frozen=True) @@ -1410,7 +1411,23 @@ def _shared_expert_output_bytes_per_token(layer: torch.nn.Module) -> int: return hidden * 2 -def _expert_lora_weight_storage(lora: Any) -> tuple[int, int, int] | None: +def _slot_lora_tensors( + lora: Any, slot_ref: "LoRASlotRef | None" = None +) -> tuple[torch.Tensor, torch.Tensor] | None: + """Read the selected owner directly, without changing the execution context.""" + if slot_ref is None: + return lora.A_T, lora.B_T + if slot_ref.name is None: + return None + from art.megatron.lora import LoRA + + slot = LoRA._slot(lora, slot_ref) + return None if slot is None else (slot.A_T, slot.B_T) + + +def _expert_lora_weight_storage( + lora: Any, slot_ref: "LoRASlotRef | None" = None +) -> tuple[int, int, int] | None: """New padded weights, transposes and effective rank for the Quack path. Original contiguous parameters are already in the allocator baseline. This @@ -1418,10 +1435,14 @@ def _expert_lora_weight_storage(lora: Any) -> tuple[int, int, int] | None: """ from art.megatron.lora import LoRA - a, b = getattr(lora, "A_T", None), getattr(lora, "B_T", None) + if type(lora) is not LoRA or "_slot" in vars(lora): + return None + tensors = _slot_lora_tensors(lora, slot_ref) + if tensors is None: + return None + a, b = tensors if ( - type(lora) is not LoRA - or "forward" in vars(lora) + "forward" in vars(lora) or "active_lora_tensors" in vars(lora) or lora._forward_hooks or lora._forward_pre_hooks @@ -1451,6 +1472,7 @@ def _moe_output_bytes_per_token( *, checkpoint_grad: bool = False, converted_stages: list[tuple[int, int]] | None = None, + slot_ref: "LoRASlotRef | None" = None, ) -> int: """Known routed-expert working set, not a complete model/compiled bound.""" if shape != ParallelShape(tp=1, cp=1): @@ -1507,16 +1529,27 @@ def _moe_output_bytes_per_token( or config.cuda_graph_impl != "none" or any( name in vars(dispatcher) - for name in ( - "preprocess", - "dispatch_preprocess", - "dispatch_postprocess", + for name in ("preprocess", "dispatch_postprocess") + ) + or ( + "dispatch_preprocess" in vars(dispatcher) + and not ( + slot_ref is not None + and slot_ref.name is not None + and type(dispatcher.dispatch_preprocess) is partial + and dispatcher.dispatch_preprocess.func + is _moe_dispatch_preprocess + and dispatcher.dispatch_preprocess.args == (dispatcher,) + and not dispatcher.dispatch_preprocess.keywords ) ) or "routing" in vars(layer.router) ): return 0 - weights = lora.B_T + tensors = _slot_lora_tensors(lora, slot_ref) + # Enclosing row storage is still charged for an inactive adapter; + # only selected tensors create converted weights. + inputs, weights = tensors if tensors is not None else (lora.A_T, lora.B_T) if ( weights.dtype not in (torch.float16, torch.bfloat16) or weights.shape[-1] != fc2.out_features @@ -1526,7 +1559,6 @@ def _moe_output_bytes_per_token( return 0 features = 2 * fc2.out_features enclosing_fc1 = None - inputs = getattr(lora, "A_T", None) if ( isinstance(inputs, torch.Tensor) and inputs.ndim == weights.ndim == 3 @@ -1575,7 +1607,7 @@ def _moe_output_bytes_per_token( config.moe_router_topk * features * weights.element_size() + shared ) coefficient = max(coefficient, row_bytes) - storage = _expert_lora_weight_storage(lora) + storage = _expert_lora_weight_storage(lora, slot_ref) if converted_stages is not None and storage is not None: padded, transposes, effective = storage saved_fc1, rank_fc1 = 0, 0 @@ -1583,7 +1615,12 @@ def _moe_output_bytes_per_token( if enclosing_fc1 is not None: adapter = getattr(enclosing_fc1, "lora", None) base = getattr(enclosing_fc1, "linear_fc1", None) - first = _expert_lora_weight_storage(adapter) + first = _expert_lora_weight_storage(adapter, slot_ref) + first_tensors = ( + _slot_lora_tensors(adapter, slot_ref) + if first is not None + else None + ) if ( first is not None and adapter is not None @@ -1592,10 +1629,11 @@ def _moe_output_bytes_per_token( and "forward" not in vars(base) and not base._forward_hooks and not base._forward_pre_hooks - and adapter.A_T.dtype == weights.dtype - and adapter.A_T.shape[:2] + and first_tensors is not None + and first_tensors[0].dtype == weights.dtype + and first_tensors[0].shape[:2] == (weights.shape[0], fc2.out_features) - and adapter.B_T.shape[2] == enclosing_fc1.out_features + and first_tensors[1].shape[2] == enclosing_fc1.out_features ): first_padding, first_transposes, first_rank = first # FC1 retains both routed H inputs and its base O1 @@ -3123,6 +3161,7 @@ def feed(value: Any) -> None: signature.request_mix, signature.grad_enabled, signature.grad_modes, + signature.slot_shapes, p.packed_tokens, p.logical_tokens, p.inactive_logical_tokens, @@ -3266,6 +3305,7 @@ def _split_chunk_lower_cost( requests, slot_group_count=len(groups), grad_modes=tuple(mode for (_, mode), _ in groups), + slot_groups=tuple(key for key, _ in groups), ) logical_tokens = _active_logical_tokens(requests) cost = self._subforward_cost( @@ -3274,6 +3314,7 @@ def _split_chunk_lower_cost( signature=signature, logical_tokens=logical_tokens, group_rows=tuple(group_rows), + slot_refs=tuple(ref for (ref, _), _ in groups), head_workspace_bytes=head_workspace_bytes, # The average CP load is an optimistic bound, not an admission cost. retained_tokens=(packed_tokens + signature.topology[2] - 1) @@ -3580,12 +3621,19 @@ def _checkpoint_moe_bytes_per_token(self) -> int: raise ValueError("Invalid constructor checkpoint MoE coefficient") return gradient - def _moe_workspace_bytes(self, rows: int, *, checkpoint_grad: bool = False) -> int: + def _moe_workspace_bytes( + self, + rows: int, + *, + checkpoint_grad: bool = False, + slot_ref: "LoRASlotRef | None" = None, + ) -> int: """Maximum of same-layer affine stages, not a retained multi-layer bank. - Cached at construction before dispatcher wrapping; model/slot shapes - must remain unchanged, as for the existing row coefficient. Ordinary - non-checkpoint gradients retain only the prior forward-stage coverage. + The constructor cache covers original tensors. Explicit slots are + repriced from their tensor metadata and original owners, including + this rank's exact dispatcher wrapper. Ordinary non-checkpoint gradients + retain only forward-stage coverage. """ coefficient = ( self._checkpoint_moe_bytes_per_token() @@ -3597,6 +3645,20 @@ def _moe_workspace_bytes(self, rows: int, *, checkpoint_grad: bool = False) -> i "_moe_gradient_stages" if checkpoint_grad else "_moe_forward_stages", (), ) + if slot_ref is not None and slot_ref.name is not None: + selected: list[tuple[int, int]] = [] + coefficient = ( + _moe_output_bytes_per_token( + self.runtime.model, + self._parallel_shape, + checkpoint_grad=checkpoint_grad, + converted_stages=selected, + slot_ref=slot_ref, + ) + if self._moe_layers + else 0 + ) + stages = tuple(selected) if coefficient else () if type(stages) is not tuple or any( type(stage) is not tuple or len(stage) != 2 @@ -3614,14 +3676,16 @@ def _moe_workspace_bytes(self, rows: int, *, checkpoint_grad: bool = False) -> i ) def _checkpoint_memory_floor( - self, group_rows: tuple[tuple[int, bool], ...] + self, + group_rows: tuple[tuple[int, bool], ...], + slot_refs: tuple["LoRASlotRef | None", ...] | None = None, ) -> tuple[int, int]: """Conservative saved-boundary charge and one disjoint MoE workspace. Count actual local full/uniform/1 boundaries, including aliases, rather than claiming measured distinct storage. Only this call's new groups enter the term; already-live graphs remain in the availability baseline. - No-grad calls also keep decoder input, current layer input, its MLP + No-grad groups also keep decoder input, current layer input, its MLP residual and norm output across the MoE stage. Count these four row tensors separately from returned outputs, allowing storage aliases. This is not a bound for custom preprocessing, attention, or all backward. @@ -3683,10 +3747,11 @@ def _checkpoint_memory_floor( retained = gradient_rows * layers * self._hidden_size * 2 if gradient_rows: self._checkpoint_moe_bytes_per_token() + refs = (None,) * len(group_rows) if slot_refs is None else slot_refs workspace = max( - self._moe_workspace_bytes(rows, checkpoint_grad=grad) - + (0 if gradient_rows else 4 * rows * self._hidden_size * 2) - for rows, grad in group_rows + self._moe_workspace_bytes(rows, checkpoint_grad=grad, slot_ref=ref) + + (0 if grad else 4 * rows * self._hidden_size * 2) + for (rows, grad), ref in zip(group_rows, refs, strict=True) ) return retained, workspace @@ -3698,6 +3763,7 @@ def _plan_cost(self, plan: _FlatForwardPlan) -> _SubforwardCost: logical_tokens=plan.active_logical_tokens, gdn_segments=plan.grad_segment_count, group_rows=self._plan_group_rows(plan), + slot_refs=tuple(g.slot_ref for g in plan.groups), head_workspace_bytes=self._plan_head_workspace_bytes(plan), checkpoint_floor=_gdn_memory.plan_floor(self, plan), retained_tokens=self._plan_retained_tokens(plan), @@ -3712,6 +3778,7 @@ def _subforward_cost( logical_tokens: int, gdn_segments: int = 0, group_rows: tuple[tuple[int, bool], ...] = (), + slot_refs: tuple["LoRASlotRef | None", ...] | None = None, head_workspace_bytes: int = 0, checkpoint_floor: tuple[int, int] = (0, 0), retained_tokens: int | None = None, @@ -3723,6 +3790,7 @@ def _subforward_cost( logical_tokens=logical_tokens, gdn_segments=gdn_segments, group_rows=group_rows, + slot_refs=slot_refs, head_workspace_bytes=head_workspace_bytes, checkpoint_floor=checkpoint_floor, retained_tokens=retained_tokens, @@ -3734,7 +3802,8 @@ def _subforward_cost( output_bytes=output_bytes, required=required, checkpoint_retained_bytes=max( - self._checkpoint_memory_floor(group_rows)[0], checkpoint_floor[0] + self._checkpoint_memory_floor(group_rows, slot_refs)[0], + checkpoint_floor[0], ), ) return _SubforwardCost(required=required, retained=retained) @@ -5368,6 +5437,7 @@ def _plan_flat_forward( requests, slot_group_count=len(plans), grad_modes=tuple(mode for (_, mode), _ in groups), + slot_groups=tuple(key for key, _ in groups), ), selected_max_depth=selected_max_depth, inactive_logical_tokens=logical_tokens @@ -5403,6 +5473,12 @@ def _estimate_flat_forward( checkpoint=checkpoint, ensure_slots=not sync_planning_errors, ) + if self._moe_layers and any( + ref is not None and ref.name is not None for (ref, _), _ in groups + ): + # This cheap return type has no slot metadata. Materialize the + # exact plan instead of admitting with the constructor rank. + return None if ( any(mode for (_, mode), _ in groups) and _gdn_memory.model_shapes(self) is not None @@ -5515,6 +5591,7 @@ def _estimate_flat_forward( requests, slot_group_count=len(groups), grad_modes=tuple(mode for (_, mode), _ in groups), + slot_groups=tuple(key for key, _ in groups), ), tuple(group_rows), head_workspace_bytes, @@ -5904,8 +5981,12 @@ def _memory_signature_from_requests( *, slot_group_count: int, grad_modes: Iterable[bool], + slot_groups: Iterable[tuple["LoRASlotRef | None", bool]] = (), ) -> _MemorySignature: modes = tuple(sorted(grad_modes)) + shapes = tuple( + sorted((grad, self._slot_memory_shapes(ref)) for ref, grad in slot_groups) + ) return _MemorySignature( topology=self._topology_key(), planner_coefficients=(self._coefficient_version, self._coefficient_table), @@ -5915,8 +5996,36 @@ def _memory_signature_from_requests( ), grad_enabled=any(modes), grad_modes=modes, + slot_shapes=shapes if any(any(shape) for _, shape in shapes) else (), ) + def _slot_memory_shapes( + self, ref: "LoRASlotRef | None" + ) -> tuple[tuple[int, ...], ...]: + """Separate empirical trust across actual selected adapter layouts.""" + if ( + ref is None + or ref.name is None + or isinstance(ref, _LocalLoRASlotRef) + or not (getattr(self, "_moe_layers", 0) or getattr(self, "_gdn_layers", 0)) + ): + # Generic/no-component planning must not import Megatron or walk + # model owners just to construct its existing memory signature. + return () + from art.megatron.lora import LoRA + + shapes = [] + for chunk in self.runtime.model: + for module in chunk.modules(): + if type(module) is LoRA: + tensors = _slot_lora_tensors(module, ref) + shapes.append( + () + if tensors is None + else (tensors[0].ndim, *tensors[0].shape, *tensors[1].shape) + ) + return tuple(shapes) + def _topology_key(self) -> tuple[int, int, int, int]: try: topology = self._topology() @@ -5956,6 +6065,7 @@ def _memory_check( logical_tokens=forward.active_logical_tokens, gdn_segments=forward.grad_segment_count, group_rows=self._plan_group_rows(forward), + slot_refs=tuple(g.slot_ref for g in forward.groups), head_workspace_bytes=self._plan_head_workspace_bytes(forward), checkpoint_floor=_gdn_memory.plan_floor(self, forward), retained_tokens=self._plan_retained_tokens(forward), @@ -6378,6 +6488,7 @@ def _estimate_required_memory_bytes_from_values( logical_tokens: int | None = None, gdn_segments: int = 0, group_rows: tuple[tuple[int, bool], ...] = (), + slot_refs: tuple["LoRASlotRef | None", ...] | None = None, head_workspace_bytes: int = 0, checkpoint_floor: tuple[int, int] = (0, 0), retained_tokens: int | None = None, @@ -6489,8 +6600,14 @@ def _estimate_required_memory_bytes_from_values( ) # Groups execute sequentially: summed packed rows conservatively bound # this FC2 component, not all workspace or retained graphs. - static_compute = max(static_compute, self._moe_workspace_bytes(packed_tokens)) - retained, workspace = self._checkpoint_memory_floor(group_rows) + static_compute = max( + static_compute, + *( + self._moe_workspace_bytes(packed_tokens, slot_ref=ref) + for ref in (slot_refs or (None,)) + ), + ) + retained, workspace = self._checkpoint_memory_floor(group_rows, slot_refs) static_compute = max( static_compute, max(retained, checkpoint_floor[0]) diff --git a/tests/unit/test_trainer_rank_checkpoint_memory.py b/tests/unit/test_trainer_rank_checkpoint_memory.py index 51998b2eb..a041fa240 100644 --- a/tests/unit/test_trainer_rank_checkpoint_memory.py +++ b/tests/unit/test_trainer_rank_checkpoint_memory.py @@ -8,7 +8,7 @@ import torch from art.trainer_rank import ForwardInput, TrainerRank -from art.trainer_rank._impl import Unset, _MemoryProfile +from art.trainer_rank._impl import Unset, _ForwardRefusal, _MemoryProfile def rank(): @@ -133,7 +133,10 @@ def test_per_group_padding_precedes_gradient_filter(): r._physical_tokens = lambda n: n + (-n % 8) values = r._estimate_flat_forward(requests(9, 17)) assert values[0] == 40 and values[3] == ((16, True), (24, False)) - assert r._checkpoint_memory_floor(values[3]) == (16 * 40 * 4096, 24 * 188416) + assert r._checkpoint_memory_floor(values[3]) == ( + 16 * 40 * 4096, + 24 * (188416 + 4 * 2048 * 2), + ) def test_no_grad_enclosure_empty_and_unsupported(): @@ -351,8 +354,62 @@ def test_no_grad_enclosure_uses_max_group_and_affine_stage(): mixed = ((3, True), (11, False)) assert r._checkpoint_memory_floor(mixed) == ( 3 * 40 * 2048 * 2, - max(3 * 188416, 11 * 188416, 11 + 1_000_000), + max(3 * 188416, max(11 * 188416, 11 + 1_000_000) + 4 * 11 * 2048 * 2), + ) + + +@pytest.mark.parametrize("gradient_first", [False, True]) +def test_mixed_plan_keeps_reference_enclosure(gradient_first): + r = rank() + gradient, reference = requests(1, 10_000) + req = [gradient, reference] if gradient_first else [reference, gradient] + reference_plan = r._plan_flat_forward([reference]) + mixed = r._plan_flat_forward(req) + reference_cost = r._plan_cost(reference_plan).required + mixed_cost = r._plan_cost(mixed).required + assert mixed_cost >= reference_cost + # The exact plan and cheap split bound must retain the same group charge. + assert r._memory_check(mixed).estimated_required_bytes == mixed_cost + assert ( + r._split_chunk_lower_cost( + req, tuple(x.input_tokens for x in req), checkpoint=Unset + ).required + == mixed_cost ) + assert not r._memory_profiles + + +@pytest.mark.parametrize("fits", [False, True]) +def test_reference_prefix_search_agrees_with_mixed_demand(fits): + r = rank() + # This uninitialized CPU fixture declares DP1; all pricing/search stays real. + r._dp_rank_and_size = lambda: (0, 1) + gradient, reference = requests(1, 10_000) + req = [reference, gradient] + reference_plan = r._plan_flat_forward([reference]) + mixed = r._plan_flat_forward(req) + # Retain the review witness's budget between its old nonmonotone costs. + budget = r._plan_cost(mixed).required if fits else 2_207_849_881 + r._available_memory_bytes = lambda: budget + assert r._memory_check(reference_plan).fits is fits + assert r._memory_check(mixed).fits is fits + selected = r._search_next_micro_batch(req, 0) + if fits: + assert not isinstance(selected, _ForwardRefusal) + assert selected.check.fits and selected.cold_start + assert selected.stats_global_count == 1 + assert ( + selected.check.estimated_required_bytes + >= r._plan_cost(reference_plan).required + ) + else: + assert isinstance(selected, _ForwardRefusal) + assert not selected.check.fits + assert ( + selected.check.estimated_required_bytes + == r._plan_cost(reference_plan).required + ) + assert not r._memory_profiles @pytest.mark.parametrize( diff --git a/tests/unit/test_trainer_rank_converted_memory.py b/tests/unit/test_trainer_rank_converted_memory.py index 50e69c0f4..a9c809163 100644 --- a/tests/unit/test_trainer_rank_converted_memory.py +++ b/tests/unit/test_trainer_rank_converted_memory.py @@ -125,7 +125,7 @@ def test_reference_and_gradient_keep_distinct_stage_modes(layer, order, rank_val retained, workspace = rank._checkpoint_memory_floor(groups) assert retained == 3 * 40 * 2048 * 2 assert workspace == max( - expected(3, rank_value, True), expected(9, rank_value, False) + expected(3, rank_value, True), expected(9, rank_value, False) + 4 * 9 * 2048 * 2 ) assert ( rank._memory_check(plan).estimated_required_bytes diff --git a/tests/unit/test_trainer_rank_shared_memory.py b/tests/unit/test_trainer_rank_shared_memory.py index 400e7b204..773186601 100644 --- a/tests/unit/test_trainer_rank_shared_memory.py +++ b/tests/unit/test_trainer_rank_shared_memory.py @@ -308,9 +308,9 @@ def test_pre_gate_cache_precedes_owned_dispatcher_and_is_checkpoint_only(layer): 19 * 40 * 4096, max( 19 * (196608 + 128), - 23 * 192512, + 23 * (192512 + 4 * 2048 * 2), 19 * (196608 - 32768 + 128) + 10485760, - 23 * (192512 - 32768 + 128) + 10485760, + 23 * (192512 - 32768 + 128 + 4 * 2048 * 2) + 10485760, ), ) for mode in (None, "selective"): @@ -339,7 +339,7 @@ def test_pre_gate_mixed_reference_and_exact_cost_mode_selection(layer, gradient_ ) assert rank._checkpoint_memory_floor(rank._plan_group_rows(mixed)) == ( 67 * 40 * 4096, - 4096 * 192512, + 4096 * (192512 + 4 * 2048 * 2), ) # A reference-only path must not read or validate the unused gradient cache. rank._moe_checkpoint_grad_bytes_per_token = None diff --git a/tests/unit/test_trainer_rank_slot_memory.py b/tests/unit/test_trainer_rank_slot_memory.py new file mode 100644 index 000000000..d8584b671 --- /dev/null +++ b/tests/unit/test_trainer_rank_slot_memory.py @@ -0,0 +1,240 @@ +"""Selected-slot pricing through the real CPU loader; no model/CUDA execution.""" + +from dataclasses import replace +from functools import partial + +import pytest +from test_trainer_rank_converted_memory import expected, weights +from test_trainer_rank_moe_memory import layer as layer +from test_trainer_rank_pending_memory import rank_with_moe +import torch + +from art.megatron.lora import LoRA, LoRASlotRef, use_lora_slot +from art.trainer_rank import ForwardInput, _gdn_memory +from art.trainer_rank._impl import ( + Unset, + _CheckpointSlot, + _expert_lora_weight_storage, + _ForwardRefusal, + _moe_dispatch_preprocess, + _SplitForwardPlan, +) + + +def load_slot(rank, name, selected_rank): + """Supply omitted inert DP1 fixture metadata, then use the real slot loader.""" + ref = LoRASlotRef("checkpoint", name) + for chunk in rank.runtime.model: + for index, lora in enumerate(chunk.modules()): + if type(lora) is not LoRA: + continue + expert = lora.A_T.ndim == 3 + lora.adapter_model_prefix = f"layer{index}" + ( + ".{expert}" if expert else "" + ) + if not hasattr(lora, "_slot_keys"): + lora._slot_keys = {} + lora._slot_modules = torch.nn.ModuleDict() + lora._expert_offset = 0 + count = lora.A_T.shape[0] if expert else 1 + lora._expert_ids = tuple(range(count)) + for param in (lora.A_T, lora.B_T): + param.lora_shard_domain = "expert_tensor" if expert else "tp" + param.lora_tp_sharded = False + inputs, outputs = lora.A_T.shape[-2], lora.B_T.shape[-1] + # expand creates a small CPU source view; the real loader stacks, + # makes contiguous tensors and clones its actual slot Parameters. + adapter = {} + for i in range(count): + prefix = lora.adapter_model_prefix.format(expert=i) + adapter[prefix + ".lora_A.weight"] = torch.zeros( + (), dtype=torch.bfloat16 + ).expand(selected_rank, inputs) + adapter[prefix + ".lora_B.weight"] = torch.zeros( + (), dtype=torch.bfloat16 + ).expand(outputs, selected_rank) + assert lora.load_lora_slot(ref, adapter, requires_grad=False) + assert lora._slot(ref).rank == selected_rank + rank._checkpoint_slots[name] = _CheckpointSlot() + return ref + + +def request(name, *, rows=1, grad=False): + return ForwardInput( + input_tokens=torch.arange(rows), + hidden_states=True, + checkpoint=name, + no_grad=not grad, + ) + + +@pytest.mark.parametrize("base,selected", [(8, 64), (64, 8), (8, 1)]) +@pytest.mark.parametrize("grad", [False, True]) +def test_selected_rank_prices_real_loaded_parameters(layer, base, selected, grad): + rank, _ = rank_with_moe(weights(layer, base)) + original = rank._moe_workspace_bytes(1, checkpoint_grad=grad) + ref = load_slot(rank, "selected", selected) + adapter = layer.experts.linear_fc2.lora + active = adapter._slot(ref) + expected_transposes = 256 * max(8, selected) * (512 + 2048) * 2 + assert _expert_lora_weight_storage(adapter, ref)[1] == expected_transposes + assert active.A_T.shape[-1] == selected and adapter.A_T.shape[-1] == base + for rows in (1, 64, 50640): + assert rank._moe_workspace_bytes( + rows, checkpoint_grad=grad, slot_ref=ref + ) == expected(rows, selected, grad) + assert rank._moe_workspace_bytes(0, checkpoint_grad=grad, slot_ref=ref) == 0 + assert rank._moe_workspace_bytes(1, checkpoint_grad=grad) == original + assert rank._moe_workspace_bytes(1, checkpoint_grad=grad, slot_ref=ref) != original + plan = rank._plan_flat_forward([request("selected", grad=grad)], ensure_slots=False) + required = rank._memory_check(plan).estimated_required_bytes + assert required == rank._plan_cost(plan).required + assert required >= int(expected(1, selected, grad) * 1.1) + rank._available_memory_bytes = lambda: required - 1 + assert not rank._memory_check(plan).fits + assert not torch.cuda.is_initialized() + + +def test_mixed_slot_context_and_exact_search_fallback(layer): + rank, _ = rank_with_moe(weights(layer, 8)) + small = load_slot(rank, "small", 1) + large = load_slot(rank, "large", 64) + requests = [request("small", rows=2), request("large", rows=1, grad=True)] + with use_lora_slot(small): + lora = layer.experts.linear_fc2.lora + original = lora.active_lora_tensors()[0] + assert original is lora._slot(small).A_T + plan = rank._plan_flat_forward(requests, ensure_slots=False) + assert tuple(g.slot_ref for g in plan.groups) == (small, large) + assert rank._moe_workspace_bytes( + 1, checkpoint_grad=True, slot_ref=large + ) == expected(1, 64, True) + assert rank._estimate_flat_forward(requests) is None + required = rank._memory_check(plan).estimated_required_bytes + assert required == rank._plan_cost(plan).required + rank._available_memory_bytes = lambda: required + admitted = rank._search_next_micro_batch([requests], 0) + assert not isinstance(admitted, _ForwardRefusal) and admitted.check.fits + assert admitted.check.estimated_required_bytes == required + rank._available_memory_bytes = lambda: 1 + assert isinstance(rank._search_next_micro_batch([requests], 0), _ForwardRefusal) + assert lora.active_lora_tensors()[0] is original + cost = rank._split_chunk_lower_cost( + requests, tuple(r.input_tokens for r in requests), checkpoint=Unset + ) + assert cost.required <= rank._plan_cost(plan).required + + +def test_gdn_pending_uses_selected_output_rank(layer): + rank, gd = rank_with_moe(weights(layer, 8)) + small, large = load_slot(rank, "small", 1), load_slot(rank, "large", 64) + small_shape = _gdn_memory.model_shapes(rank, small)[1][0] + large_shape = _gdn_memory.model_shapes(rank, large)[1][0] + assert (small_shape.output_lora_rank, large_shape.output_lora_rank) == (1, 64) + p = rank._plan_flat_forward( + [request("large", rows=65, grad=True)], ensure_slots=False + ) + group = p.groups[0] + buckets = _gdn_memory.cp1_buckets(group.packed.segments) + assert ( + large_shape.pending(65, buckets) - small_shape.pending(65, buckets) + == 65 * (64 - 1) * 2 + ) + retained, workspace = _gdn_memory.plan_floor(rank, p) + assert retained == 65 * 40 * 2048 * 2 + assert workspace == expected(65, 64, True) + large_shape.pending(65, buckets) + assert gd.out_proj.lora.A_T.shape[1] == 1 + + +def test_profile_and_split_key_separate_slot_layout_and_grad_mode(layer): + rank, _ = rank_with_moe(weights(layer, 8)) + load_slot(rank, "small", 1) + load_slot(rank, "large", 64) + small = rank._plan_flat_forward([request("small", grad=True)], ensure_slots=False) + large = rank._plan_flat_forward([request("large", grad=True)], ensure_slots=False) + rank._update_memory_profile(small, 2**30, retained_bytes=1) + assert small.signature != large.signature + assert ( + small.signature in rank._memory_profiles + and large.signature not in rank._memory_profiles + ) + assert not rank._all_ranks_have_memory_profile( + packed_tokens=large.packed_tokens, signature=large.signature + ) + # Give the selected layout its own genuine empirical floor. It dominates + # static demand without adding that static demand a second time. + rank._update_memory_profile(large, 2**30, retained_bytes=1) + assert rank._memory_check(large).estimated_required_bytes == int(2**30 * 1.1) + left = rank._plan_flat_forward( + [request("small", grad=True), request("large")], ensure_slots=False + ) + right = rank._plan_flat_forward( + [request("small"), request("large", grad=True)], ensure_slots=False + ) + assert left.signature != right.signature + split = _SplitForwardPlan((small,), ((0,),), 1) + changed = replace(split, subforwards=(replace(small, signature=large.signature),)) + assert rank._split_memory_key(split) != rank._split_memory_key(changed) + + +def test_slot_reload_reprices_without_mutating_constructor_or_profile(layer): + rank, _ = rank_with_moe(weights(layer, 8)) + ref = load_slot(rank, "reload", 1) + small = rank._plan_flat_forward([request("reload")], ensure_slots=False) + rank._update_memory_profile(small, 2**30, retained_bytes=1) + load_slot(rank, "reload", 64) + large = rank._plan_flat_forward([request("reload")], ensure_slots=False) + assert large.signature != small.signature + assert large.signature not in rank._memory_profiles + assert rank._moe_workspace_bytes(1, slot_ref=ref) == expected(1, 64, False) + assert rank._moe_workspace_bytes(1) == expected(1, 8, False) + + +@pytest.mark.parametrize( + "kind", ["foreign", "wrong owner", "keywords", "slot override"] +) +def test_slot_pricing_retains_original_owner_guards(layer, kind): + rank, _ = rank_with_moe(weights(layer, 8)) + ref = load_slot(rank, "loaded", 64) + dispatcher = layer.token_dispatcher + if kind == "slot override": + lora = layer.experts.linear_fc2.lora + lora._slot = lambda selected: lora._slot_modules["slot_0"] + assert _expert_lora_weight_storage(lora, ref) is None + else: + dispatcher.dispatch_preprocess = ( + (lambda *args: None) + if kind == "foreign" + else partial(_moe_dispatch_preprocess, object()) + if kind == "wrong owner" + else partial(_moe_dispatch_preprocess, dispatcher, hidden_states=None) + ) + assert rank._moe_workspace_bytes(1, slot_ref=ref) == 0 + + +@pytest.mark.parametrize("kind", ["generic", "inactive", "without megatron"]) +def test_generic_signature_needs_neither_megatron_nor_module_walk(monkeypatch, kind): + import builtins + from types import SimpleNamespace + + from art.trainer_rank import TrainerRank + from art.trainer_rank._impl import _LocalLoRASlotRef + + rank = TrainerRank.__new__(TrainerRank) + rank._moe_layers = 0 if kind == "generic" else 1 + rank._gdn_layers = 0 + rank.runtime = object() # There is deliberately no model/modules facade. + ref = ( + _LocalLoRASlotRef(name="selected") + if kind == "without megatron" + else SimpleNamespace(name=None if kind == "inactive" else "selected") + ) + original = builtins.__import__ + + def guarded(name, *args, **kwargs): + if name == "art.megatron.lora": + raise AssertionError("generic pricing must not import Megatron") + return original(name, *args, **kwargs) + + monkeypatch.setattr(builtins, "__import__", guarded) + assert rank._slot_memory_shapes(ref) == () From 657fa723df14355d56c2ea371c062c2f62565588 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Thu, 17 Sep 2026 15:54:13 +0000 Subject: [PATCH 10/18] Fix slot-memory test CI routing and optional type checks --- .github/workflows/prek.yml | 2 ++ src/art/trainer_rank/_impl.py | 2 +- tests/unit/test_trainer_rank_slot_memory.py | 15 +++++++++++---- 3 files changed, 14 insertions(+), 5 deletions(-) diff --git a/.github/workflows/prek.yml b/.github/workflows/prek.yml index 533a0a984..fd86551b0 100644 --- a/.github/workflows/prek.yml +++ b/.github/workflows/prek.yml @@ -227,6 +227,7 @@ jobs: tests/unit/test_trainer_rank_weird_shapes.py \ tests/unit/test_trainer_rank_admission_inputs.py \ tests/unit/test_trainer_rank_checkpoint_memory.py \ + tests/unit/test_trainer_rank_slot_memory.py \ tests/unit/test_trainer_rank_head_memory.py \ tests/unit/test_trainer_rank_mixed_head_memory.py \ tests/unit/test_trainer_rank_ignored_mixed_head.py \ @@ -264,6 +265,7 @@ jobs: --ignore=tests/unit/test_trainer_rank_weird_shapes.py \ --ignore=tests/unit/test_trainer_rank_admission_inputs.py \ --ignore=tests/unit/test_trainer_rank_checkpoint_memory.py \ + --ignore=tests/unit/test_trainer_rank_slot_memory.py \ --ignore=tests/unit/test_trainer_rank_head_memory.py \ --ignore=tests/unit/test_trainer_rank_mixed_head_memory.py \ --ignore=tests/unit/test_trainer_rank_ignored_mixed_head.py \ diff --git a/src/art/trainer_rank/_impl.py b/src/art/trainer_rank/_impl.py index 1ecb9294b..2d4566f67 100644 --- a/src/art/trainer_rank/_impl.py +++ b/src/art/trainer_rank/_impl.py @@ -1498,7 +1498,7 @@ def _moe_output_bytes_per_token( experts = getattr(layer, "experts", None) fc2: Any = getattr(experts, "linear_fc2", None) lora: Any = getattr(fc2, "lora", None) - dispatcher = getattr(layer, "token_dispatcher", None) + dispatcher: Any = getattr(layer, "token_dispatcher", None) sites = ( (layer, MoELayer), (experts, TEGroupedMLP), diff --git a/tests/unit/test_trainer_rank_slot_memory.py b/tests/unit/test_trainer_rank_slot_memory.py index d8584b671..f15a33fae 100644 --- a/tests/unit/test_trainer_rank_slot_memory.py +++ b/tests/unit/test_trainer_rank_slot_memory.py @@ -77,7 +77,9 @@ def test_selected_rank_prices_real_loaded_parameters(layer, base, selected, grad adapter = layer.experts.linear_fc2.lora active = adapter._slot(ref) expected_transposes = 256 * max(8, selected) * (512 + 2048) * 2 - assert _expert_lora_weight_storage(adapter, ref)[1] == expected_transposes + storage = _expert_lora_weight_storage(adapter, ref) + assert storage is not None + assert storage[1] == expected_transposes assert active.A_T.shape[-1] == selected and adapter.A_T.shape[-1] == base for rows in (1, 64, 50640): assert rank._moe_workspace_bytes( @@ -128,8 +130,11 @@ def test_mixed_slot_context_and_exact_search_fallback(layer): def test_gdn_pending_uses_selected_output_rank(layer): rank, gd = rank_with_moe(weights(layer, 8)) small, large = load_slot(rank, "small", 1), load_slot(rank, "large", 64) - small_shape = _gdn_memory.model_shapes(rank, small)[1][0] - large_shape = _gdn_memory.model_shapes(rank, large)[1][0] + small_shapes = _gdn_memory.model_shapes(rank, small) + large_shapes = _gdn_memory.model_shapes(rank, large) + assert small_shapes is not None and large_shapes is not None + small_shape = small_shapes[1][0] + large_shape = large_shapes[1][0] assert (small_shape.output_lora_rank, large_shape.output_lora_rank) == (1, 64) p = rank._plan_flat_forward( [request("large", rows=65, grad=True)], ensure_slots=False @@ -207,7 +212,9 @@ def test_slot_pricing_retains_original_owner_guards(layer, kind): if kind == "foreign" else partial(_moe_dispatch_preprocess, object()) if kind == "wrong owner" - else partial(_moe_dispatch_preprocess, dispatcher, hidden_states=None) + else partial( + _moe_dispatch_preprocess, dispatcher, hidden_states=torch.empty(0) + ) ) assert rank._moe_workspace_bytes(1, slot_ref=ref) == 0 From 04863263436b2d6977c171f059e8681636a6218d Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Thu, 17 Sep 2026 16:26:04 +0000 Subject: [PATCH 11/18] Use real slot state in trainer planning test fixtures --- tests/unit/test_trainer_rank_planning_status.py | 4 +++- tests/unit/test_trainer_rank_weird_shapes.py | 5 ----- 2 files changed, 3 insertions(+), 6 deletions(-) diff --git a/tests/unit/test_trainer_rank_planning_status.py b/tests/unit/test_trainer_rank_planning_status.py index b264aca5b..8b290a4a2 100644 --- a/tests/unit/test_trainer_rank_planning_status.py +++ b/tests/unit/test_trainer_rank_planning_status.py @@ -105,10 +105,12 @@ def _worker(index: int, directory: Path) -> None: rank = TrainerRank.__new__(TrainerRank) rank.device = torch.device("cpu") rank._padded_vocab_size = None + rank._moe_layers = rank._gdn_layers = 0 + rank._slot_stack = [] + rank._default_slot_ref = None rank._planning_seconds_accum = 0.0 rank._dp_rank_and_size = lambda: (index, 2) rank._physical_tokens = lambda tokens: tokens - rank._resolve_slot_ref = lambda request, **_: request.no_grad rank._estimate_group_request_output_bytes = lambda requests: 0 rank._memory_signature_from_requests = lambda *args, **kwargs: None rank._forward_item = lambda request: SimpleNamespace( diff --git a/tests/unit/test_trainer_rank_weird_shapes.py b/tests/unit/test_trainer_rank_weird_shapes.py index e7962d903..c7b556bbc 100644 --- a/tests/unit/test_trainer_rank_weird_shapes.py +++ b/tests/unit/test_trainer_rank_weird_shapes.py @@ -807,11 +807,6 @@ def test_adaptive_planner_probes_new_heterogeneous_signatures( ) -> None: rank = TrainerRank(_runtime()) monkeypatch.setattr(rank, "_dp_rank_and_size", lambda: (0, 1)) - monkeypatch.setattr( - rank, - "_resolve_slot_ref", - lambda request, **_kwargs: request.checkpoint, - ) for index in range(4): rank._checkpoint_slots.setdefault(f"S{index}", _CheckpointSlot()).params = () inputs = [ From 32c874fbf8fda16eb4d994272dc42b1a8127115a Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Sat, 19 Sep 2026 01:25:44 +0000 Subject: [PATCH 12/18] Account for checkpoint input gradients in static peak admission --- src/art/trainer_rank/_impl.py | 66 ++++++- ...trainer_rank_checkpoint_gradient_memory.py | 174 ++++++++++++++++++ .../test_trainer_rank_checkpoint_memory.py | 4 +- .../unit/test_trainer_rank_pending_memory.py | 8 +- tests/unit/test_trainer_rank_shared_memory.py | 4 +- 5 files changed, 238 insertions(+), 18 deletions(-) create mode 100644 tests/unit/test_trainer_rank_checkpoint_gradient_memory.py diff --git a/src/art/trainer_rank/_impl.py b/src/art/trainer_rank/_impl.py index 2d4566f67..40e4c1b8d 100644 --- a/src/art/trainer_rank/_impl.py +++ b/src/art/trainer_rank/_impl.py @@ -1023,6 +1023,10 @@ class _SubforwardCost: required: int retained: int + # Before the safety factor, separate from observed forward retention. + checkpoint_retained: int = 0 + checkpoint_workspace: int = 0 + checkpoint_input_gradient: int = 0 @property def ephemeral(self) -> int: @@ -3103,9 +3107,22 @@ def _admit_split_rung( return None, check def _split_rung_check(self, costs: Sequence[_SubforwardCost]) -> _MemoryCheck: - return self._memory_check_required( - sum(cost.retained for cost in costs) + max(cost.ephemeral for cost in costs) + return self._memory_check_required(self._split_required_memory(costs)) + + @staticmethod + def _split_required_memory(costs: Sequence[_SubforwardCost]) -> int: + required = sum(cost.retained for cost in costs) + max( + cost.ephemeral for cost in costs ) + if any(cost.checkpoint_input_gradient for cost in costs): + # The caller owns all returned graphs. A calibrated forward-retained + # discount cannot replace the sum of their input-gradient extents. + checkpoint = sum( + cost.checkpoint_retained + cost.checkpoint_input_gradient + for cost in costs + ) + max(cost.checkpoint_workspace for cost in costs) + required = max(required, int(checkpoint * _MEMORY_SAFETY_FACTOR)) + return required @staticmethod def _split_memory_key(plan: _SplitForwardPlan) -> bytes | None: @@ -3249,9 +3266,7 @@ def _record_split_memory_floor( def _split_plan_memory_check( self, plan: _SplitForwardPlan, costs: Sequence[_SubforwardCost] ) -> _MemoryCheck: - required = sum(cost.retained for cost in costs) + max( - cost.ephemeral for cost in costs - ) + required = self._split_required_memory(costs) key = self._split_memory_key(plan) empirical = ( 0 @@ -3340,8 +3355,8 @@ def _split_chunk_lower_cost( # cannot. Its full-required retention is not a pruning lower bound. # Keep outputs and the independent source retention floor; exact # plan costs keep both trust guards. - return _SubforwardCost( - required=cost.required, + return replace( + cost, retained=min( cost.retained, int( @@ -3794,6 +3809,10 @@ def _subforward_cost( head_workspace_bytes=head_workspace_bytes, checkpoint_floor=checkpoint_floor, retained_tokens=retained_tokens, + include_checkpoint_input_gradient=False, + ) + checkpoint_retained, checkpoint_workspace = self._checkpoint_memory_floor( + group_rows, slot_refs ) retained = self._retained_memory_bytes( signature, @@ -3802,11 +3821,36 @@ def _subforward_cost( output_bytes=output_bytes, required=required, checkpoint_retained_bytes=max( - self._checkpoint_memory_floor(group_rows, slot_refs)[0], + checkpoint_retained, checkpoint_floor[0], ), ) - return _SubforwardCost(required=required, retained=retained) + # One logical BF16 input gradient per eligible full/uniform/1 boundary. + # 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 + checkpoint_retained = output_bytes + max( + checkpoint_retained, checkpoint_floor[0] + ) + checkpoint_workspace = max( + checkpoint_workspace, head_workspace_bytes, checkpoint_floor[1] + ) + if gradient: + required = max( + required, + int( + (checkpoint_retained + checkpoint_workspace + gradient) + * _MEMORY_SAFETY_FACTOR + ), + ) + return _SubforwardCost( + required=required, + retained=retained, + checkpoint_retained=checkpoint_retained, + checkpoint_workspace=checkpoint_workspace, + checkpoint_input_gradient=gradient, + ) def _retained_memory_bytes( self, @@ -6492,6 +6536,7 @@ def _estimate_required_memory_bytes_from_values( head_workspace_bytes: int = 0, checkpoint_floor: tuple[int, int] = (0, 0), retained_tokens: int | None = None, + include_checkpoint_input_gradient: bool = True, ) -> int: if packed_tokens <= 0: return output_bytes @@ -6611,7 +6656,8 @@ def _estimate_required_memory_bytes_from_values( static_compute = max( static_compute, max(retained, checkpoint_floor[0]) - + max(workspace, head_workspace_bytes, checkpoint_floor[1]), + + max(workspace, head_workspace_bytes, checkpoint_floor[1]) + + (retained if include_checkpoint_input_gradient else 0), ) if signature.topology[2] > 1: # Local head results coexist with full CP outputs during gathering. diff --git a/tests/unit/test_trainer_rank_checkpoint_gradient_memory.py b/tests/unit/test_trainer_rank_checkpoint_gradient_memory.py new file mode 100644 index 000000000..463540c88 --- /dev/null +++ b/tests/unit/test_trainer_rank_checkpoint_gradient_memory.py @@ -0,0 +1,174 @@ +"""Partial input-gradient extents: CPU admission math, not peak/overlap proof.""" + +from dataclasses import replace + +import pytest +from test_trainer_rank_checkpoint_memory import price, rank, requests +from test_trainer_rank_moe_memory import layer # noqa: F401 +from test_trainer_rank_pending_memory import full_requests, pending_rank # noqa: F401 +import torch + +from art.trainer_rank import ForwardInput +from art.trainer_rank._impl import Unset, _MemoryProfile, _SplitForwardPlan + + +def profile(r, plan, *, rate=1, retained_rate=1): + r._memory_profiles[plan.signature] = _MemoryProfile( + bytes_per_token=rate, + packed_tokens=plan.packed_tokens, + logical_per_packed=1, + retained_compute_bytes_per_token=retained_rate, + ) + + +def test_pending_cold_peak_does_not_become_forward_retention(pending_rank): + r = pending_rank + plan = r._plan_flat_forward(full_requests()) + cost = r._plan_cost(plan) + gradient = 8 * 6330 * 40 * 2048 * 2 + assert cost.checkpoint_input_gradient == gradient + # Exact previous cold estimate, including outputs and its one safety factor. + assert cost.retained == 23102959299 + assert cost.required == int( + (plan.output_bytes + 2 * gradient + 12705630112) * 1.1 + ) + assert r._memory_check(plan).estimated_required_bytes == cost.required + profile(r, plan) + warm = r._plan_cost(plan) + assert warm.retained == int((plan.output_bytes + gradient) * 1.1) + assert warm.required == cost.required + + +@pytest.mark.parametrize("rows", [1, 67, 1024]) +def test_attention_only_extent_scales_with_gradient_rows(rows): + r = rank() + values = r._estimate_flat_forward(requests(rows, 4096)) + cost = price(r, values) + assert r._gdn_layers == 0 + assert cost.checkpoint_input_gradient == rows * 40 * 2048 * 2 + assert cost.required >= int( + ( + cost.checkpoint_retained + + cost.checkpoint_workspace + + cost.checkpoint_input_gradient + ) + * 1.1 + ) + + +def test_gradient_is_not_absorbed_by_larger_head_workspace(): + r = rank() + n, out, sig, groups, _head = r._estimate_flat_forward(requests(67, 4096)) + head = 10**10 + cost = price(r, (n, out, sig, groups, head)) + gradient = 67 * 40 * 2048 * 2 + assert cost.checkpoint_workspace == head + assert cost.required == int((out + head + 2 * gradient) * 1.1) + assert cost.retained == int((out + head + gradient) * 1.1) + r._memory_profiles[sig] = _MemoryProfile( + bytes_per_token=10**9, + packed_tokens=n, + logical_per_packed=1, + retained_compute_bytes_per_token=1, + ) + assert price(r, (n, out, sig, groups, head)).required == int( + (out + n * 10**9) * 1.1 + ) + + +@pytest.mark.parametrize("unsupported", [False, True]) +def test_no_grad_and_unsupported_checkpoint_do_not_gain_extent(unsupported): + r = rank() + req = [ + ForwardInput( + input_tokens=torch.arange(67), + hidden_states=True, + no_grad=not unsupported, + ) + ] + values = r._estimate_flat_forward(req) + if unsupported: + r.runtime.model[0].decoder.config.recompute_num_layers = 2 + cost = price(r, values) + assert cost.checkpoint_input_gradient == 0 + assert cost.required == cost.retained + + +def test_split_sums_all_gradient_children_outside_workspace_max(): + r = rank() + children = [ + r._plan_flat_forward( + [ + ForwardInput( + input_tokens=torch.arange(rows), + hidden_states=True, + no_grad=no_grad, + ) + ] + ) + for rows, no_grad in ((17, False), (29, False), (83, True)) + ] + for child in children: + profile(r, child) + costs = [r._plan_cost(child) for child in children] + split = _SplitForwardPlan(tuple(children), ((0,), (1,), (2,)), 3) + assert [c.checkpoint_input_gradient for c in costs] == [ + 17 * 40 * 4096, + 29 * 40 * 4096, + 0, + ] + expected = int( + ( + sum(c.checkpoint_retained + c.checkpoint_input_gradient for c in costs) + + max(c.checkpoint_workspace for c in costs) + ) + * 1.1 + ) + old_child_max = sum(c.retained for c in costs) + max(c.ephemeral for c in costs) + assert expected > old_child_max + r._available_memory_bytes = lambda: expected - 1 + assert not r._split_rung_check(costs).fits + check = r._split_plan_memory_check(split, costs) + assert not check.fits and check.estimated_required_bytes == expected + r._available_memory_bytes = lambda: expected + assert r._split_plan_memory_check(split, costs).fits + key = r._split_memory_key(split) + assert key is not None + r._split_memory_floors[key] = 10**9 + assert r._split_plan_memory_check(split, costs).estimated_required_bytes == int( + 10**9 * 1.1 + ) + + +def test_lower_bound_profile_cliff_preserves_separate_peak_component(): + r = rank() + req = [ + ForwardInput( + input_tokens=torch.arange(128), + target_tokens=torch.arange(128), + no_grad=False, + ) + for _ in range(16) + ] + full = r._plan_flat_forward(req, memory_minimal=True) + profile(r, full) + r._memory_profiles[full.signature] = replace( + r._memory_profiles[full.signature], packed_tokens=256 + ) + lower = r._split_chunk_lower_cost( + req, tuple(q.input_tokens for q in req), checkpoint=Unset + ) + assert lower.checkpoint_input_gradient == 128 * 40 * 4096 + assert lower.required <= r._plan_cost(full).required + assert lower.checkpoint_retained == full.output_bytes + 128 * 40 * 4096 + + +def test_finite_backward_profile_is_not_added_to_static_gradient_extent(): + r = rank() + plan = r._plan_flat_forward(requests(17, 19)) + static = r._plan_cost(plan).required + rate = static * 4 + profile(r, plan, rate=rate) + assert r._plan_cost(plan).required == int( + (plan.output_bytes + plan.packed_tokens * rate) * 1.1 + ) diff --git a/tests/unit/test_trainer_rank_checkpoint_memory.py b/tests/unit/test_trainer_rank_checkpoint_memory.py index a041fa240..7a5e7d952 100644 --- a/tests/unit/test_trainer_rank_checkpoint_memory.py +++ b/tests/unit/test_trainer_rank_checkpoint_memory.py @@ -105,7 +105,7 @@ def test_required_and_learned_retained_use_max_not_sum(): ) cost = price(r, values) old = max(n * 2048 * 2 * 14, n * 188416) - assert cost.required == int((out + max(old, retained + work)) * 1.1) + assert cost.required == int((out + max(old, 2 * retained + work)) * 1.1) assert cost.retained == int((out + retained) * 1.1) r._memory_profiles[sig] = replace( r._memory_profiles[sig], @@ -248,7 +248,7 @@ def test_split_keeps_complete_order_and_checks_each_new_subforward(): logical_per_packed=1, retained_compute_bytes_per_token=1, ) - limit = 160_000_000 + limit = 250_000_000 used = 0 r._available_memory_bytes = lambda: limit - used result = r._find_admissible_forward(req, checkpoint=Unset, refusal_prefix="test") diff --git a/tests/unit/test_trainer_rank_pending_memory.py b/tests/unit/test_trainer_rank_pending_memory.py index e8af2cc23..0bd8e27ec 100644 --- a/tests/unit/test_trainer_rank_pending_memory.py +++ b/tests/unit/test_trainer_rank_pending_memory.py @@ -133,7 +133,7 @@ def test_actual_constructor_cache_and_full_plan(pending_rank): assert ( rank._memory_check(plan).estimated_required_bytes == rank._plan_cost(plan).required - == 23102959299 + == 32229502659 ) selected = rank._select_next_micro_batch(requests, 0) assert ( @@ -154,7 +154,7 @@ def test_exact_pending_demand_survives_recovery(monkeypatch, pending_rank, fits_ plan = pending_rank._plan_flat_forward(requests) assert pending_rank._estimate_flat_forward(requests) is None assert g.plan_floor(pending_rank, plan) == (8296857600, 12705630112) - assert pending_rank._memory_check(plan).estimated_required_bytes == 23102959299 + assert pending_rank._memory_check(plan).estimated_required_bytes == 32229502659 _check_component_demand_recovery( monkeypatch, pending_rank, requests, fits_after=fits_after ) @@ -172,8 +172,8 @@ def test_original_installed_norm_preserves_pending_floor(layer): assert g.model_shapes(rank) is not None plan = rank._plan_flat_forward(full_requests()) assert g.plan_floor(rank, plan) == (8296857600, 12705630112) - assert rank._memory_check(plan).estimated_required_bytes == 23102959299 - assert rank._plan_cost(plan).required == 23102959299 + assert rank._memory_check(plan).estimated_required_bytes == 32229502659 + assert rank._plan_cost(plan).required == 32229502659 assert rank._estimate_flat_forward(full_requests()) is None for requests in ([], full_requests(no_grad=True)): assert g.plan_floor(rank, rank._plan_flat_forward(requests)) == (0, 0) diff --git a/tests/unit/test_trainer_rank_shared_memory.py b/tests/unit/test_trainer_rank_shared_memory.py index 773186601..79b73bbcb 100644 --- a/tests/unit/test_trainer_rank_shared_memory.py +++ b/tests/unit/test_trainer_rank_shared_memory.py @@ -126,7 +126,7 @@ def test_shared_return_in_actual_constructor_and_plan(layer, gate, no_grad): 8296857600, 50640 * (checkpoint_coefficient + 128) + 3157761952, ) - assert rank._plan_cost(plan).required == (23559286467 if gate else 23331122883) + assert rank._plan_cost(plan).required == (32685829827 if gate else 32457666243) selected = rank._select_next_micro_batch(requests, 0) assert ( selected.check.estimated_required_bytes @@ -147,7 +147,7 @@ def test_original_norm_installation_preserves_shared_return(layer, gated): 8296857600, 50640 * (checkpoint_coefficient + 128) + 3157761952, ) - expected = 23559286467 if gated else 23331122883 + expected = 32685829827 if gated else 32457666243 assert rank._memory_check(plan).estimated_required_bytes == expected assert rank._plan_cost(plan).required == expected From be3fe546289b338399c524c9e30884b03d1e7fd5 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Sat, 19 Sep 2026 02:17:54 +0000 Subject: [PATCH 13/18] Preserve split forward ordering with checkpoint gradient peak allowance --- src/art/trainer_rank/_impl.py | 12 ++++- ...trainer_rank_checkpoint_gradient_memory.py | 53 +++++++++++++++++++ 2 files changed, 64 insertions(+), 1 deletion(-) diff --git a/src/art/trainer_rank/_impl.py b/src/art/trainer_rank/_impl.py index 40e4c1b8d..f5490d8f2 100644 --- a/src/art/trainer_rank/_impl.py +++ b/src/art/trainer_rank/_impl.py @@ -1027,6 +1027,8 @@ class _SubforwardCost: checkpoint_retained: int = 0 checkpoint_workspace: int = 0 checkpoint_input_gradient: int = 0 + # Already in required; backward allowance must not reorder forward execution. + checkpoint_peak_increment: int = 0 @property def ephemeral(self) -> int: @@ -3095,7 +3097,13 @@ def _admit_split_rung( ] costs = [self._plan_cost(plan) for plan in plans] # Bind original request mappings; floor keys normalize execution order. - order = sorted(range(len(plans)), key=lambda i: (-costs[i].ephemeral, i)) + order = sorted( + range(len(plans)), + key=lambda i: ( + -(costs[i].ephemeral - costs[i].checkpoint_peak_increment), + i, + ), + ) split = _SplitForwardPlan( subforwards=tuple(plans[i] for i in order), request_indices=tuple(tuple(chunks[i]) for i in order), @@ -3836,6 +3844,7 @@ def _subforward_cost( checkpoint_workspace = max( checkpoint_workspace, head_workspace_bytes, checkpoint_floor[1] ) + forward_required = required if gradient: required = max( required, @@ -3850,6 +3859,7 @@ def _subforward_cost( checkpoint_retained=checkpoint_retained, checkpoint_workspace=checkpoint_workspace, checkpoint_input_gradient=gradient, + checkpoint_peak_increment=required - forward_required, ) def _retained_memory_bytes( diff --git a/tests/unit/test_trainer_rank_checkpoint_gradient_memory.py b/tests/unit/test_trainer_rank_checkpoint_gradient_memory.py index 463540c88..790a4c58a 100644 --- a/tests/unit/test_trainer_rank_checkpoint_gradient_memory.py +++ b/tests/unit/test_trainer_rank_checkpoint_gradient_memory.py @@ -91,6 +91,7 @@ def test_no_grad_and_unsupported_checkpoint_do_not_gain_extent(unsupported): r.runtime.model[0].decoder.config.recompute_num_layers = 2 cost = price(r, values) assert cost.checkpoint_input_gradient == 0 + assert cost.checkpoint_peak_increment == 0 assert cost.required == cost.retained @@ -172,3 +173,55 @@ def test_finite_backward_profile_is_not_added_to_static_gradient_extent(): assert r._plan_cost(plan).required == int( (plan.output_bytes + plan.packed_tokens * rate) * 1.1 ) + + +def test_unequal_checkpoint_gradients_preserve_cold_split_execution_order(): + r = rank() + req = [ + ForwardInput(input_tokens=torch.arange(rows), hidden_states=True) + for rows in (17, 29) + ] + costs = [r._plan_cost(r._plan_flat_forward([q])) for q in req] + assert ( + 0 < costs[0].checkpoint_input_gradient < costs[1].checkpoint_input_gradient + ) + assert costs[0].ephemeral < costs[1].ephemeral + # Cold forward retention was the entire pre-gradient requirement: both + # original priorities are zero, so the stable original order must survive. + assert all(c.ephemeral == c.checkpoint_peak_increment for c in costs) + r._available_memory_bytes = lambda: 1 << 60 + split, check = r._admit_split_rung( + ((0,), (1,)), req, [q.input_tokens for q in req], checkpoint=Unset + ) + assert split is not None and check.fits + assert check.estimated_required_bytes == r._split_required_memory(costs) + assert split.request_indices == ((0,), (1,)), split.request_indices + + +@pytest.mark.parametrize("fully_masked", [False, True]) +def test_split_priority_subtracts_only_uncovered_gradient_peak(fully_masked): + r = rank() + plan = r._plan_flat_forward(requests(17, 19)) + cold = r._plan_cost(plan) + # Place a real learned peak between the two static estimates, or above both. + measured = ( + cold.required + 10**7 + if fully_masked + else (cold.retained + cold.required) / 2 + ) + rate = (measured / 1.1 - plan.output_bytes) / plan.packed_tokens + profile(r, plan, rate=rate) + cost = r._plan_cost(plan) + old_required = int((plan.output_bytes + int(plan.packed_tokens * rate)) * 1.1) + assert cold.retained < old_required + assert cost.required == max(cold.required, old_required) + assert ( + cost.ephemeral - cost.checkpoint_peak_increment + == old_required - cost.retained + ) + if fully_masked: + assert cost.checkpoint_peak_increment == 0 + else: + assert 0 < cost.checkpoint_peak_increment < int( + cost.checkpoint_input_gradient * 1.1 + ) From f570804683ff3a6b1d9d2fff32638633c45fbb5b Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Sat, 19 Sep 2026 02:23:25 +0000 Subject: [PATCH 14/18] style: format checkpoint gradient memory tests --- ...trainer_rank_checkpoint_gradient_memory.py | 21 +++++++------------ 1 file changed, 8 insertions(+), 13 deletions(-) diff --git a/tests/unit/test_trainer_rank_checkpoint_gradient_memory.py b/tests/unit/test_trainer_rank_checkpoint_gradient_memory.py index 790a4c58a..d6d9ded56 100644 --- a/tests/unit/test_trainer_rank_checkpoint_gradient_memory.py +++ b/tests/unit/test_trainer_rank_checkpoint_gradient_memory.py @@ -29,9 +29,7 @@ def test_pending_cold_peak_does_not_become_forward_retention(pending_rank): assert cost.checkpoint_input_gradient == gradient # Exact previous cold estimate, including outputs and its one safety factor. assert cost.retained == 23102959299 - assert cost.required == int( - (plan.output_bytes + 2 * gradient + 12705630112) * 1.1 - ) + assert cost.required == int((plan.output_bytes + 2 * gradient + 12705630112) * 1.1) assert r._memory_check(plan).estimated_required_bytes == cost.required profile(r, plan) warm = r._plan_cost(plan) @@ -182,9 +180,7 @@ def test_unequal_checkpoint_gradients_preserve_cold_split_execution_order(): for rows in (17, 29) ] costs = [r._plan_cost(r._plan_flat_forward([q])) for q in req] - assert ( - 0 < costs[0].checkpoint_input_gradient < costs[1].checkpoint_input_gradient - ) + assert 0 < costs[0].checkpoint_input_gradient < costs[1].checkpoint_input_gradient assert costs[0].ephemeral < costs[1].ephemeral # Cold forward retention was the entire pre-gradient requirement: both # original priorities are zero, so the stable original order must survive. @@ -205,9 +201,7 @@ def test_split_priority_subtracts_only_uncovered_gradient_peak(fully_masked): cold = r._plan_cost(plan) # Place a real learned peak between the two static estimates, or above both. measured = ( - cold.required + 10**7 - if fully_masked - else (cold.retained + cold.required) / 2 + cold.required + 10**7 if fully_masked else (cold.retained + cold.required) / 2 ) rate = (measured / 1.1 - plan.output_bytes) / plan.packed_tokens profile(r, plan, rate=rate) @@ -216,12 +210,13 @@ def test_split_priority_subtracts_only_uncovered_gradient_peak(fully_masked): assert cold.retained < old_required assert cost.required == max(cold.required, old_required) assert ( - cost.ephemeral - cost.checkpoint_peak_increment - == old_required - cost.retained + cost.ephemeral - cost.checkpoint_peak_increment == old_required - cost.retained ) if fully_masked: assert cost.checkpoint_peak_increment == 0 else: - assert 0 < cost.checkpoint_peak_increment < int( - cost.checkpoint_input_gradient * 1.1 + assert ( + 0 + < cost.checkpoint_peak_increment + < int(cost.checkpoint_input_gradient * 1.1) ) From bb394998b8000f0c95499844b750d5bb7141ea93 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Sat, 19 Sep 2026 02:34:52 +0000 Subject: [PATCH 15/18] Route checkpoint gradient memory tests to Megatron environment --- .github/workflows/prek.yml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.github/workflows/prek.yml b/.github/workflows/prek.yml index fd86551b0..c55a6dc5a 100644 --- a/.github/workflows/prek.yml +++ b/.github/workflows/prek.yml @@ -227,6 +227,7 @@ jobs: tests/unit/test_trainer_rank_weird_shapes.py \ tests/unit/test_trainer_rank_admission_inputs.py \ tests/unit/test_trainer_rank_checkpoint_memory.py \ + tests/unit/test_trainer_rank_checkpoint_gradient_memory.py \ tests/unit/test_trainer_rank_slot_memory.py \ tests/unit/test_trainer_rank_head_memory.py \ tests/unit/test_trainer_rank_mixed_head_memory.py \ @@ -265,6 +266,7 @@ jobs: --ignore=tests/unit/test_trainer_rank_weird_shapes.py \ --ignore=tests/unit/test_trainer_rank_admission_inputs.py \ --ignore=tests/unit/test_trainer_rank_checkpoint_memory.py \ + --ignore=tests/unit/test_trainer_rank_checkpoint_gradient_memory.py \ --ignore=tests/unit/test_trainer_rank_slot_memory.py \ --ignore=tests/unit/test_trainer_rank_head_memory.py \ --ignore=tests/unit/test_trainer_rank_mixed_head_memory.py \ From 8b3e1dfb0059554223666a84db99bfb86371da63 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Sat, 19 Sep 2026 03:10:17 +0000 Subject: [PATCH 16/18] test(trainer-rank): include checkpoint gradients in head peak expectations --- tests/unit/test_trainer_rank_head_memory.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/tests/unit/test_trainer_rank_head_memory.py b/tests/unit/test_trainer_rank_head_memory.py index 34e4be546..50c16343f 100644 --- a/tests/unit/test_trainer_rank_head_memory.py +++ b/tests/unit/test_trainer_rank_head_memory.py @@ -139,9 +139,11 @@ def test_outputs_retention_and_empirical_peak_are_counted_once(): r = rank() plan = r._plan_flat_forward([request(grad=True)]) retained = 512 * 40 * 2048 * 2 + gradient = 512 * 40 * 2048 * 2 head = 3 * 512 * 248320 * 2 cost = r._plan_cost(plan) - assert cost.required == int((plan.output_bytes + retained + head) * 1.1) + assert cost.retained == int((plan.output_bytes + retained + head) * 1.1) + assert cost.required == int((plan.output_bytes + retained + gradient + head) * 1.1) r._memory_profiles[plan.signature] = _MemoryProfile( bytes_per_token=2_000_000, packed_tokens=512, @@ -302,9 +304,10 @@ def test_target_backward_refuses_budget_below_logits_and_both_gradients(rows): r = rank() plan = r._plan_flat_forward([request(rows, grad=True)]) retained, _ = r._checkpoint_memory_floor(r._plan_group_rows(plan)) + gradient = rows * 40 * 2048 * 2 dense = min(rows, 512) * 248320 * 2 - before = int((plan.output_bytes + retained + 2 * dense) * 1.1) - expected = int((plan.output_bytes + retained + 3 * dense) * 1.1) + before = int((plan.output_bytes + retained + gradient + 2 * dense) * 1.1) + expected = int((plan.output_bytes + retained + gradient + 3 * dense) * 1.1) r._available_memory_bytes = lambda: (before + expected) // 2 check = r._memory_check(plan) assert check.estimated_required_bytes == expected From 935e0e3aa455801d65489a6ff8e042088d091be2 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Sat, 19 Sep 2026 10:26:22 +0000 Subject: [PATCH 17/18] fix(trainer-rank): retain generic admission for declined MoE layouts --- src/art/trainer_rank/_gdn_memory.py | 12 ++++++- src/art/trainer_rank/_impl.py | 2 ++ .../unit/test_trainer_rank_pending_memory.py | 34 +++++++++++++++++++ 3 files changed, 47 insertions(+), 1 deletion(-) diff --git a/src/art/trainer_rank/_gdn_memory.py b/src/art/trainer_rank/_gdn_memory.py index 13db19a55..38027dd76 100644 --- a/src/art/trainer_rank/_gdn_memory.py +++ b/src/art/trainer_rank/_gdn_memory.py @@ -193,9 +193,19 @@ def model_shapes( # Like the existing static floor, this cache requires unchanged model, # dtype and topology since construction; rebuilding the rank invalidates it. moe = rank._moe_output_bytes_per_token - if type(moe) is not int or moe < 0 or (rank._moe_layers and not moe): + if ( + type(moe) is not int + or moe < 0 + or ( + rank._moe_layers + and not moe + and getattr(rank, "_moe_memory_supported", True) is not False + ) + ): raise ValueError("Invalid constructor MoE coefficient for GDN pending floor") rank._checkpoint_moe_bytes_per_token() + if rank._moe_layers and not moe: + return None shapes = [] for layer in decoder.layers: gdn = getattr(layer, "self_attention", None) diff --git a/src/art/trainer_rank/_impl.py b/src/art/trainer_rank/_impl.py index f5490d8f2..cf958a6a5 100644 --- a/src/art/trainer_rank/_impl.py +++ b/src/art/trainer_rank/_impl.py @@ -1821,6 +1821,8 @@ def memory_field(name: str, default: Any = None) -> Any: if self._moe_layers else 0 ) + # A declined component is distinct from a qualified cache later damaged. + self._moe_memory_supported = self._moe_output_bytes_per_token > 0 # Both modes inspect original owners before dispatcher caches are installed. self._moe_checkpoint_grad_bytes_per_token = ( _moe_output_bytes_per_token( diff --git a/tests/unit/test_trainer_rank_pending_memory.py b/tests/unit/test_trainer_rank_pending_memory.py index 0bd8e27ec..490352f76 100644 --- a/tests/unit/test_trainer_rank_pending_memory.py +++ b/tests/unit/test_trainer_rank_pending_memory.py @@ -323,6 +323,40 @@ def test_invalid_cached_coefficient_stops_before_memory_reduction(pending_rank, assert statuses == [False] +@pytest.mark.parametrize("unsupported", ["capacity", "forward hook"]) +def test_constructor_declined_moe_keeps_generic_admission(layer, unsupported): + layer = _enclosing_moe(layer) + if unsupported == "capacity": + layer.config.moe_expert_capacity_factor = 1.0 + else: + layer.experts.linear_fc2.register_forward_hook(lambda *args: None) + rank, _ = rank_with_moe(layer) + # The actual constructor declines this component; no cache is overwritten. + assert rank._moe_output_bytes_per_token == 0 + assert rank._moe_checkpoint_grad_bytes_per_token == 0 + assert g.model_shapes(rank) is None + requests = full_requests() + assert rank._estimate_flat_forward(requests) is not None + plan = rank._plan_flat_forward(requests) + assert g.plan_floor(rank, plan) == (0, 0) + required = rank._plan_cost(plan).required + # Generic checkpoint-input accounting still applies without a MoE component. + gradient = 50640 * 40 * 2048 * 2 + assert required == int((plan.output_bytes + 2 * gradient) * 1.1) + rank._available_memory_bytes = lambda: required - 1 + assert not rank._memory_check(plan).fits + rank._available_memory_bytes = lambda: required + assert rank._memory_check(plan).fits + + +def test_declined_moe_does_not_hide_invalid_gradient_cache(layer): + layer.config.moe_expert_capacity_factor = 1.0 + rank, _ = rank_with_moe(_enclosing_moe(layer)) + rank._moe_checkpoint_grad_bytes_per_token = -1 + with pytest.raises(ValueError, match="Invalid constructor checkpoint MoE"): + rank._estimate_flat_forward(full_requests()) + + @pytest.mark.parametrize( "field,value", [ From 05e8ebf32dd7edb6d3e7793b552e24f32aecf054 Mon Sep 17 00:00:00 2001 From: Codex Halley Date: Sat, 19 Sep 2026 18:14:07 +0000 Subject: [PATCH 18/18] Fix cold-signature planner fixture slot references Use the real slot resolver for the three calibrated-width cases. The fixture already declares its checkpoint slots; returning strings bypassed the typed slot contract used by the memory signature. Production code and assertions are unchanged. --- tests/unit/test_trainer_rank_weird_shapes.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/tests/unit/test_trainer_rank_weird_shapes.py b/tests/unit/test_trainer_rank_weird_shapes.py index cd06bb2f9..174e56499 100644 --- a/tests/unit/test_trainer_rank_weird_shapes.py +++ b/tests/unit/test_trainer_rank_weird_shapes.py @@ -843,9 +843,6 @@ def test_adaptive_planner_does_not_reuse_wide_window_for_cold_signature( ) -> None: rank = TrainerRank(_runtime()) monkeypatch.setattr(rank, "_dp_rank_and_size", lambda: (0, 1)) - monkeypatch.setattr( - rank, "_resolve_slot_ref", lambda request, **_kwargs: request.checkpoint - ) for name in ("policy", "adversary", "third"): rank._checkpoint_slots.setdefault(name, _CheckpointSlot()).params = () inputs = [