diff --git a/.github/workflows/prek.yml b/.github/workflows/prek.yml index 1f8f27e54..fd86551b0 100644 --- a/.github/workflows/prek.yml +++ b/.github/workflows/prek.yml @@ -225,6 +225,15 @@ 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_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 \ + 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/unit/test_trainer_rank_cache_recovery.py::test_dense_cp_exact_demand_fits_after_recovery \ tests/unit/test_trainer_rank_cache_recovery.py::test_dense_cp_exact_demand_refuses_after_recovery \ @@ -253,4 +262,13 @@ 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_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 \ + --ignore=tests/unit/test_trainer_rank_pending_memory.py \ + --ignore=tests/unit/test_trainer_rank_shared_memory.py \ + --ignore=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 new file mode 100644 index 000000000..13db19a55 --- /dev/null +++ b/src/art/trainer_rank/_gdn_memory.py @@ -0,0 +1,335 @@ +"""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 types import MethodType +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, 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, + _slot_lora_tensors, + ) + + 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 _empty_safe_norm_forward, _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") + rank._checkpoint_moe_bytes_per_token() + shapes = [] + for layer in decoder.layers: + gdn = getattr(layer, "self_attention", None) + if gdn is None or 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 + 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 gdn.out_norm.weight.numel() != dv + or gdn.out_norm.weight.dtype is not torch.bfloat16 + # 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 + ): + 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 "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 + # 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 + 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 + # 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 + + +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 + 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, slot_ref=group.slot_ref) + ) + 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, + *( + rank._moe_workspace_bytes( + rows, checkpoint_grad=True, slot_ref=group.slot_ref + ) + + 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 4ee194dfb..2d4566f67 100644 --- a/src/art/trainer_rank/_impl.py +++ b/src/art/trainer_rank/_impl.py @@ -27,7 +27,7 @@ import threading import time import traceback -from types import TracebackType +from types import MethodType, TracebackType from typing import ( TYPE_CHECKING, Any, @@ -54,6 +54,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, @@ -924,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) @@ -1253,13 +1255,232 @@ 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 shared is None or 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 _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 + excludes padding-concatenation temporaries and all backward GEMM workspace. + """ + from art.megatron.lora import LoRA + + 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 ( + "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 + model: Sequence[torch.nn.Module], + shape: ParallelShape, + *, + 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): 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 @@ -1277,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), @@ -1308,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 @@ -1326,7 +1558,7 @@ def _moe_output_bytes_per_token( ): return 0 features = 2 * fc2.out_features - inputs = getattr(lora, "A_T", None) + enclosing_fc1 = None if ( isinstance(inputs, torch.Tensor) and inputs.ndim == weights.ndim == 3 @@ -1360,10 +1592,124 @@ 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 - coefficient = max( - coefficient, - config.moe_router_topk * features * weights.element_size(), - ) + enclosing_fc1 = fc1 + 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 + row_bytes = ( + config.moe_router_topk * features * weights.element_size() + shared + ) + coefficient = max(coefficient, row_bytes) + 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 + 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, 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 + 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 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 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 + # 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 @@ -1460,11 +1806,33 @@ def memory_field(name: str, default: Any = None) -> Any: 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, + 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, @@ -2793,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, @@ -2907,13 +3276,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) ) @@ -2922,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( @@ -2929,6 +3313,9 @@ def _split_chunk_lower_cost( output_bytes=output_bytes, 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) // signature.topology[2], @@ -2951,13 +3338,423 @@ 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_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 _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. + + 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() + if checkpoint_grad + else self._moe_output_bytes_per_token + ) + stages = getattr( + self, + "_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 + 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], ...], + 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 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. + """ + gradient_rows = sum(rows for rows, grad in group_rows if grad) + if not group_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 + 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, 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 + def _plan_cost(self, plan: _FlatForwardPlan) -> _SubforwardCost: return self._subforward_cost( packed_tokens=plan.packed_tokens, @@ -2965,6 +3762,10 @@ def _plan_cost(self, plan: _FlatForwardPlan) -> _SubforwardCost: signature=plan.signature, 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), ) @@ -2976,6 +3777,10 @@ def _subforward_cost( signature: _MemorySignature, 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, ) -> _SubforwardCost: required = self._estimate_required_memory_bytes_from_values( @@ -2984,6 +3789,10 @@ def _subforward_cost( signature=signature, 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, ) retained = self._retained_memory_bytes( @@ -2992,6 +3801,10 @@ def _subforward_cost( logical_tokens=logical_tokens, output_bytes=output_bytes, required=required, + checkpoint_retained_bytes=max( + self._checkpoint_memory_floor(group_rows, slot_refs)[0], + checkpoint_floor[0], + ), ) return _SubforwardCost(required=required, retained=retained) @@ -3003,6 +3816,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. @@ -3018,8 +3832,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)) @@ -3965,6 +4781,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( @@ -3978,6 +4796,8 @@ def priced( * sum( _request_mix_key(r) != "inactive" for r in local_requests ), + group_rows=group_rows, + head_workspace_bytes=head_workspace_bytes, ) return ( self._memory_check_required(required, sync_across_dp=True), @@ -4617,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 @@ -4631,7 +5452,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 @@ -4652,6 +5473,19 @@ 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 + ): + # Pending saves require the actual bucket/replayed-tail geometry. + # Existing unavailable handling materializes before admission. + return None if ( self._topology_key()[2] > 1 and self._recompute_granularity != "full" @@ -4662,9 +5496,14 @@ def _estimate_flat_forward( # fallback; a global token count alone cannot price its peak. 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) @@ -4674,7 +5513,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 @@ -4688,7 +5571,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, @@ -4697,7 +5591,10 @@ 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, ) def _ensure_checkpoint_slots_for( @@ -5084,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), @@ -5095,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() @@ -5135,6 +6064,10 @@ def _memory_check( signature=forward.signature, 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), ) return self._memory_check_required(required, sync_across_dp=sync_across_dp) @@ -5554,6 +6487,10 @@ def _estimate_required_memory_bytes_from_values( signature: _MemorySignature, 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, ) -> int: if packed_tokens <= 0: @@ -5664,7 +6601,17 @@ 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, + *( + 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]) + + max(workspace, head_workspace_bytes, checkpoint_floor[1]), ) if signature.topology[2] > 1: # Local head results coexist with full CP outputs during gathering. diff --git a/tests/unit/test_trainer_rank_active_memory.py b/tests/unit/test_trainer_rank_active_memory.py index 25c1e8c89..5249e07f5 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 @@ -272,3 +273,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_admission_inputs.py b/tests/unit/test_trainer_rank_admission_inputs.py new file mode 100644 index 000000000..92b5dbde4 --- /dev/null +++ b/tests/unit/test_trainer_rank_admission_inputs.py @@ -0,0 +1,171 @@ +"""Seven additive estimator joins; scalar CPU plans, not distributed/CUDA bounds.""" + +from dataclasses import replace +from types import SimpleNamespace + +from test_trainer_rank_head_memory import rank as head_rank +from test_trainer_rank_head_memory import request as head_request +from test_trainer_rank_pending_memory import layer, pending_rank +from test_trainer_rank_recompute_memory import _hybrid_rank +import torch + +from art.trainer_rank import ForwardInput, _gdn_memory +from art.trainer_rank._impl import Unset, _FlatForwardPlan + + +def record_prices(monkeypatch, rank): + original = rank._estimate_required_memory_bytes_from_values + calls = [] + + def observed(**values): + calls.append(dict(values)) + return original(**values) + + monkeypatch.setattr(rank, "_estimate_required_memory_bytes_from_values", observed) + return calls + + +def assert_plan_values(rank, plan, values): + assert values["packed_tokens"] == plan.packed_tokens + assert values["output_bytes"] == plan.output_bytes + assert values["logical_tokens"] == plan.active_logical_tokens + assert values["signature"] == plan.signature + assert values["gdn_segments"] == plan.grad_segment_count + assert values["group_rows"] == rank._plan_group_rows(plan) + assert values["head_workspace_bytes"] == rank._plan_head_workspace_bytes(plan) + assert values["checkpoint_floor"] == _gdn_memory.plan_floor(rank, plan) + assert values["retained_tokens"] == rank._plan_retained_tokens(plan) + + +def test_shared_head_lower_and_plan_keep_all_keywords(monkeypatch): + rank = head_rank() + a = replace(head_request(2), target_tokens=torch.tensor([1, -100])) + b = replace(a, target_tokens=torch.tensor([-100, 1])) + requests = [a, a, b, b] + plan = rank._plan_flat_forward(requests, memory_minimal=True) + assert plan.packed_tokens == 2 and plan.active_logical_tokens == 8 + assert rank._plan_group_rows(plan) == ((2, False),) + assert rank._checkpoint_memory_floor(((2, False),)) == ( + 0, + 2 * (188416 + 4 * 2048 * 2), + ) + assert rank._plan_head_workspace_bytes(plan) == 2 * 248320 * 2 + calls = record_prices(monkeypatch, rank) + lower = rank._split_chunk_lower_cost( + requests, tuple(item.input_tokens for item in requests), checkpoint=Unset + ) + assert len(calls) == 1 + assert calls[0]["group_rows"] == ((2, False),) + assert calls[0]["retained_tokens"] == 2 + assert calls[0]["head_workspace_bytes"] == 248320 * 2 + assert calls[0]["checkpoint_floor"] == (0, 0) + calls.clear() + cost = rank._plan_cost(plan) + check = rank._memory_check(plan) + assert cost.required == check.estimated_required_bytes + assert lower.required <= cost.required + assert len(calls) == 2 + for values in calls: + assert_plan_values(rank, plan, values) + # All shape/floor inputs remain available together, even though the dense + # head dominates this two-row source floor and sharing changes logical rows. + assert cost.required == int((plan.output_bytes + 2 * 248320 * 2) * 1.1) + + +def test_cp_gdn_segments_groups_and_retained_tokens_reach_exact_search(monkeypatch): + rank = _hybrid_rank(monkeypatch, 2) + monkeypatch.setattr(rank, "_topology_key", lambda: (1, 2, 4, 1)) + monkeypatch.setattr(rank, "_topology", lambda: SimpleNamespace(cp=4, tp=2)) + monkeypatch.setattr(rank, "_dp_rank_and_size", lambda: (0, 1)) + monkeypatch.setattr(rank, "_available_memory_bytes", lambda: 1 << 60) + monkeypatch.setattr( + rank, "_max_rank_model_tokens", lambda batch, **_: batch.tokens.numel() * 3 // 4 + ) + requests = [ + ForwardInput(input_tokens=torch.arange(64) + offset, hidden_states=True) + for offset in (0, 100) + ] + [ + ForwardInput( + input_tokens=torch.arange(80) + 200, hidden_states=True, no_grad=True + ), + ForwardInput(input_tokens=torch.arange(1000)), + ] + assert _gdn_memory.model_shapes(rank) is None # This exercises the CP fallback. + plan = rank._plan_flat_forward(requests) + assert plan.grad_segment_count == 2 + assert rank._plan_group_rows(plan) == ((128, True), (80, False)) + assert rank._plan_retained_tokens(plan) == 156 + for exact in (False, True): + for memory_minimal in (False, True): + assert ( + rank._estimate_flat_forward( + requests, exact=exact, memory_minimal=memory_minimal + ) + is None + ) + calls = record_prices(monkeypatch, rank) + lower = rank._split_chunk_lower_cost( + requests, tuple(item.input_tokens for item in requests), checkpoint=Unset + ) + assert len(calls) == 1 + assert calls[0]["group_rows"] == ((128, True), (80, False)) + assert calls[0]["retained_tokens"] == 52 # Optimistic CP average only here. + calls.clear() + cost = rank._plan_cost(plan) + check = rank._memory_check(plan) + assert cost.required == check.estimated_required_bytes + assert lower.required < cost.required + assert len(calls) == 2 + for values in calls: + assert_plan_values(rank, plan, values) + calls.clear() + # 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 + for values in calls: + assert_plan_values(rank, selected.plan, values) + + +def test_full_gdn_nonzero_floor_survives_separate_exact_fallback( + monkeypatch, pending_rank +): + rank = pending_rank + monkeypatch.setattr(rank, "_available_memory_bytes", lambda: 1 << 60) + requests = [ + ForwardInput(input_tokens=torch.arange(64) + offset, hidden_states=True) + for offset in (0, 100) + ] + assert rank._topology_key()[2] == 1 + assert _gdn_memory.model_shapes(rank) is not None + plan = rank._plan_flat_forward(requests) + assert plan.grad_segment_count == 2 + floor = _gdn_memory.plan_floor(rank, plan) + assert floor[0] > 0 and floor[1] > 0 + for exact in (False, True): + for memory_minimal in (False, True): + assert ( + rank._estimate_flat_forward( + requests, exact=exact, memory_minimal=memory_minimal + ) + is None + ) + calls = record_prices(monkeypatch, rank) + cost = rank._plan_cost(plan) + check = rank._memory_check(plan) + assert cost.required == check.estimated_required_bytes + assert len(calls) == 2 + for values in calls: + assert_plan_values(rank, plan, values) + assert values["checkpoint_floor"] == floor + calls.clear() + selected = rank._search_next_micro_batch([requests], 0) + assert selected.check.fits + assert selected.check.estimated_required_bytes == cost.required + assert calls + for values in calls: + assert_plan_values(rank, selected.plan, values) + assert values["checkpoint_floor"] == floor diff --git a/tests/unit/test_trainer_rank_cache_recovery.py b/tests/unit/test_trainer_rank_cache_recovery.py index 3b6b08f8c..935ecbb45 100644 --- a/tests/unit/test_trainer_rank_cache_recovery.py +++ b/tests/unit/test_trainer_rank_cache_recovery.py @@ -897,3 +897,107 @@ def test_dense_cp_exact_demand_fits_after_recovery(monkeypatch): def test_dense_cp_exact_demand_refuses_after_recovery(monkeypatch): _check_dense_cp_exact_demand_recovery(monkeypatch, fits_after_release=False) + + +def _check_component_demand_recovery( + monkeypatch, rank, requests, *, fits_after, after_available=None +): + """Real pricing/search/recovery, CPU plans and scalar CUDA counters only.""" + import pytest + + assert not _impl.dist.is_initialized() and not _impl.torch.cuda.is_initialized() + plan = rank._plan_flat_forward(requests) + required = rank._memory_check(plan).estimated_required_bytes + profiles = dict(rank._memory_profiles) + groups = rank._plan_group_rows(plan) + head = rank._plan_head_workspace_bytes(plan) + total = 10 * required + reserve = int(total * _impl._MEMORY_RESERVE_FRACTION) + free, phase = reserve + 1, 0 + searches, demands, outcomes, releases, errors = [], [], [], [], [] + estimate = rank._estimate_required_memory_bytes_from_values + search = rank._search_next_micro_batch + outcome = rank._admission_outcome + error_factory = _impl._ForwardRefusal.error + + def observed_estimate(**kwargs): + value = estimate(**kwargs) + if ( + kwargs.get("group_rows") == groups + and kwargs.get("head_workspace_bytes") == head + ): + demands.append((phase, value)) + return value + + def observed_search(*args, **kwargs): + nonlocal phase + phase += 1 + value = search(*args, **kwargs) + searches.append(value) + return value + + def observed_outcome(local): + value = outcome(local) + outcomes.append((local, value)) + return value + + def release(): + nonlocal free + releases.append(phase) + free = reserve + ( + after_available + if after_available is not None + else required + if fits_after + else 1 + ) + + def observed_error(refused, context): + error = error_factory(refused, context) + errors.append(error) + return error + + monkeypatch.delenv(_impl._TEST_HOOKS_ENV, raising=False) + # Inert CPU fixture, not an initialized MCore/distributed topology. + monkeypatch.setattr(rank, "_dp_rank_and_size", lambda: (0, 1)) + monkeypatch.setattr(rank, "device", _impl.torch.device("cuda")) + monkeypatch.setattr(_impl.torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(_impl.torch.cuda, "get_allocator_backend", lambda: "native") + monkeypatch.setattr(_impl.torch.cuda, "mem_get_info", lambda device: (free, total)) + monkeypatch.setattr(_impl.torch.cuda, "memory_allocated", lambda device: 0) + monkeypatch.setattr(_impl.torch.cuda, "memory_reserved", lambda device: total) + monkeypatch.setattr(_impl.torch.cuda, "empty_cache", release) + monkeypatch.setattr( + _impl.torch.cuda, "synchronize", lambda *a: pytest.fail("No CUDA work") + ) + monkeypatch.setattr( + rank, "_estimate_required_memory_bytes_from_values", observed_estimate + ) + monkeypatch.setattr(rank, "_search_next_micro_batch", observed_search) + monkeypatch.setattr(rank, "_admission_outcome", observed_outcome) + monkeypatch.setattr(_impl._ForwardRefusal, "error", observed_error) + monkeypatch.setattr(rank, "_snapshot_planning_telemetry", lambda *args: None) + monkeypatch.setattr( + rank, "_execute_flat_plan", lambda *a, **kw: pytest.fail("No model execution") + ) + # Cached bytes exceed demand, but are not physical-free admission credit. + assert rank._available_memory_bytes() == 1 < required < total + if fits_after: + selected = rank._select_next_micro_batch([requests], 0) + assert selected.check.fits and selected.check.available_bytes == required + assert selected.check.estimated_required_bytes == required + assert rank._plan_group_rows(selected.plan) == groups + assert rank._plan_head_workspace_bytes(selected.plan) == head + assert len(errors) == 1 + else: + with pytest.raises(Refusal) as captured: + rank._select_next_micro_batch([requests], 0) + assert len(errors) == 2 and captured.value is errors[1] + assert captured.value.__cause__ is errors[0] + assert len(searches) == 2 and isinstance(searches[0], _impl._ForwardRefusal) + assert releases == [1] and outcomes[0] == (1, 1) + assert (1, required) in demands and (2, required) in demands + assert rank._recovery_state().first_consumed + assert rank._recovery_state().owner is None + assert rank._memory_profiles == profiles + assert not _impl.torch.cuda.is_initialized() 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..a041fa240 --- /dev/null +++ b/tests/unit/test_trainer_rank_checkpoint_memory.py @@ -0,0 +1,436 @@ +"""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, _ForwardRefusal, _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 + result._moe_checkpoint_grad_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 + 4 * 2048 * 2), + ) + + +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), + ) + 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) + + +@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) + ] + 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, + ) + 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 + ) + r._available_memory_bytes = lambda: expected - 1 + assert not r._memory_check(plan).fits + r._available_memory_bytes = lambda: expected + assert r._memory_check(plan).fits + + +def test_no_grad_enclosure_uses_max_group_and_affine_stage(): + r = rank() + 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 + ), + ) + mixed = ((3, True), (11, False)) + assert r._checkpoint_memory_floor(mixed) == ( + 3 * 40 * 2048 * 2, + 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( + "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) + assert r._checkpoint_memory_floor(((11, False),)) == (0, 0) + + +@pytest.mark.parametrize("cp", [2, 4]) +def test_no_grad_enclosure_keeps_cp_fallback(cp): + r = rank() + r._topology_key = lambda: (1, 1, cp, 1) + assert r._checkpoint_memory_floor(((11, False),)) == (0, 0) 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..a9c809163 --- /dev/null +++ b/tests/unit/test_trainer_rank_converted_memory.py @@ -0,0 +1,334 @@ +"""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 == 0 and pending == (0, 0) + assert workspace == expected(8, rank_value, False) + 4 * 8 * 2048 * 2 + 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) + 4 * 9 * 2048 * 2 + ) + 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_head_memory.py b/tests/unit/test_trainer_rank_head_memory.py new file mode 100644 index 000000000..34e4be546 --- /dev/null +++ b/tests/unit/test_trainer_rank_head_memory.py @@ -0,0 +1,404 @@ +"""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 + + +@pytest.mark.parametrize("fits_after", (False, True)) +def test_mixed_checkpoint_head_demand_survives_recovery(monkeypatch, fits_after): + from test_trainer_rank_cache_recovery import _check_component_demand_recovery + + r = rank() + requests = [request(8, grad=True), request(16, grad=False, hidden=True)] + values = r._estimate_flat_forward(requests, exact=True) + assert values[3] == ((8, True), (16, False)) + assert r._checkpoint_memory_floor(values[3])[0] == 8 * 40 * 2048 * 2 + assert values[4] == 3 * 8 * 248320 * 2 + _check_component_demand_recovery(monkeypatch, r, requests, fits_after=fits_after) + + +def test_recovery_keeps_profile_demand_above_cold_head_floor(monkeypatch): + from test_trainer_rank_cache_recovery import _check_component_demand_recovery + + r = rank() + requests = [request(1, grad=True)] # One row cannot be split smaller. + plan = r._plan_flat_forward(requests) + cold = r._memory_check(plan).estimated_required_bytes + r._update_memory_profile(plan, 4 * cold, retained_bytes=plan.output_bytes) + assert r._memory_check(plan).estimated_required_bytes > cold + _check_component_demand_recovery( + monkeypatch, r, requests, fits_after=False, after_available=cold + ) + + +def test_real_head_split_fits_before_cache_recovery(monkeypatch): + from test_trainer_rank_split import _recording_executor + + from art.trainer_rank import TrainerRank, _impl + + r = rank() + monkeypatch.setattr(r, "_dp_rank_and_size", lambda: (0, 1)) + requests = [ + replace(request(8), input_tokens=torch.arange(8) + 100 * i) for i in range(4) + ] + whole = r._plan_flat_forward(requests) + r._update_memory_profile( + whole, r._plan_cost(whole).required, retained_bytes=whole.output_bytes + ) + children = [r._plan_flat_forward(requests[i : i + 2]) for i in (0, 2)] + left, right = [r._plan_cost(child) for child in children] + budget = max(left.required, left.retained + right.required) + assert 0 < left.retained and budget < r._plan_cost(whole).required + total = 10 * r._plan_cost(whole).required + free = budget + int(total * _impl._MEMORY_RESERVE_FRACTION) + probe = TrainerRank.__new__(TrainerRank) + probe.device = torch.device("cuda") + monkeypatch.delenv(_impl._TEST_HOOKS_ENV, raising=False) + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "get_allocator_backend", lambda: "native") + monkeypatch.setattr(torch.cuda, "mem_get_info", lambda device: (free, total)) + monkeypatch.setattr(torch.cuda, "memory_allocated", lambda device: 0) + monkeypatch.setattr(torch.cuda, "memory_reserved", lambda device: total) + monkeypatch.setattr( + r, "_available_memory_bytes", lambda: TrainerRank._available_memory_bytes(probe) + ) + monkeypatch.setattr( + r, "_try_cache_recovery", lambda *a, **kw: pytest.fail("Split already fits") + ) + executed = _recording_executor(monkeypatch, r) + batches = list(r.forward_micro_batches([requests])) + assert len(batches) == 1 and batches[0].stats.subforward_count == 2 + assert batches[0].stats.global_count == 1 and len(executed) == 2 + assert [ + [int(output.target_logprobs.item()) for output in group] + for group in batches[0].outputs + ] == [[7, 107, 207, 307]] + assert r.last_forward_telemetry()["subforward_request_indices"] == ((0, 1), (2, 3)) + assert not torch.cuda.is_initialized() + + +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..e8af2cc23 --- /dev/null +++ b/tests/unit/test_trainer_rank_pending_memory.py @@ -0,0 +1,400 @@ +"""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, *, 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 + + 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 + if install_hooks: + from art.megatron.gdn.operator import install_gdn_island_hooks + + install_gdn_island_hooks([model]) + 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 + 50640 * 128 + 3157761952, + ) + assert ( + rank._memory_check(plan).estimated_required_bytes + == rank._plan_cost(plan).required + == 23102959299 + ) + 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 + + +@pytest.mark.parametrize("fits_after", (False, True)) +def test_exact_pending_demand_survives_recovery(monkeypatch, pending_rank, fits_after): + from test_trainer_rank_cache_recovery import _check_component_demand_recovery + + requests = full_requests() + 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 + _check_component_demand_recovery( + monkeypatch, pending_rank, requests, fits_after=fits_after + ) + + +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, 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) + + +@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( + 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 7cb996e50..8b290a4a2 100644 --- a/tests/unit/test_trainer_rank_planning_status.py +++ b/tests/unit/test_trainer_rank_planning_status.py @@ -104,14 +104,17 @@ 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( - input_ids=request.input_tokens + input_ids=request.input_tokens, request=request ) rank._forward_output_metadata = lambda *args, **kwargs: (None, True) @@ -198,7 +201,7 @@ def retained_tokens(plan): error = caught if mode in ("estimate", "materialize", "price", "cp_plan"): 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_recovery_slots_distributed.py b/tests/unit/test_trainer_rank_recovery_slots_distributed.py index 3326663cb..02c106896 100644 --- a/tests/unit/test_trainer_rank_recovery_slots_distributed.py +++ b/tests/unit/test_trainer_rank_recovery_slots_distributed.py @@ -92,6 +92,7 @@ def worker(index, mode, directory): ] rank._forward_memory_group = lambda: groups[index] plan = SimpleNamespace( + groups=(), # This slot-only fixture has no retained/head/GDN groups. packed_tokens=1, logical_tokens=1, active_logical_tokens=1, 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..773186601 --- /dev/null +++ b/tests/unit/test_trainer_rank_shared_memory.py @@ -0,0 +1,385 @@ +"""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 + 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) + 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._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, + 50640 * (checkpoint_coefficient + 128) + 3157761952, + ) + 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 + == rank._memory_check(selected.plan).estimated_required_bytes + ) + + +@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()) + 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 + 128) + 3157761952, + ) + expected = 23559286467 if gated else 23331122883 + assert rank._memory_check(plan).estimated_required_bytes == expected + assert rank._plan_cost(plan).required == expected + + +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 ( + _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): + 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) + + +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 + 128), + 23 * (192512 + 4 * 2048 * 2), + 19 * (196608 - 32768 + 128) + 10485760, + 23 * (192512 - 32768 + 128 + 4 * 2048 * 2) + 10485760, + ), + ) + 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 + 4 * 2048 * 2), + ) + # 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, + 4096 * (192512 + 4 * 2048 * 2), + ) + 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) 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..f15a33fae --- /dev/null +++ b/tests/unit/test_trainer_rank_slot_memory.py @@ -0,0 +1,247 @@ +"""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 + 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( + 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_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 + ) + 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=torch.empty(0) + ) + ) + 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) == () diff --git a/tests/unit/test_trainer_rank_weird_shapes.py b/tests/unit/test_trainer_rank_weird_shapes.py index 8db3ef72e..c7b556bbc 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,)", @@ -806,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 = [ @@ -936,10 +932,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",