diff --git a/src/art/trainer_rank/_impl.py b/src/art/trainer_rank/_impl.py index 7e682e146..bd0fd45ff 100644 --- a/src/art/trainer_rank/_impl.py +++ b/src/art/trainer_rank/_impl.py @@ -1279,7 +1279,10 @@ def _lora_modules_per_layer(model: torch.nn.Module) -> int: except (AttributeError, ImportError, RuntimeError, TypeError): return 0 return max( - (sum(type(module) is LoRA for module in layer.modules()) for layer in layers), + ( + sum(isinstance(module, LoRA) for module in layer.modules()) + for layer in layers + ), default=0, ) diff --git a/src/art/trainer_rank/_memory.py b/src/art/trainer_rank/_memory.py index e86508a21..16ce4031b 100644 --- a/src/art/trainer_rank/_memory.py +++ b/src/art/trainer_rank/_memory.py @@ -936,6 +936,16 @@ def _checkpoint_hybridep_rows( return hybrid_rows +def _native_dense_cp(self: TrainerRank) -> bool: + cp = self._topology_key()[2] + return ( + cp in (2, 4, 8) + and self._parallel_shape.cp == cp + and not self._moe_layers + and self._dense_fc1_adapted + ) + + def _checkpoint_floor_from_facts( self: TrainerRank, group_rows: tuple[tuple[int, bool], ...], @@ -948,7 +958,8 @@ def _checkpoint_floor_from_facts( if not layers: return 0, 0 gradient_rows = sum(rows for rows, grad in group_rows if grad) - _, tp, _, _ = self._topology_key() + tp = self._topology_key()[1] + dense_cp = _native_dense_cp(self) # Physical rows are padded to a multiple of TP; each rank saves its shard. retained = ( sum(-(-rows // tp) for rows, grad in group_rows if grad) @@ -964,7 +975,7 @@ def _checkpoint_floor_from_facts( + (0 if grad else 4 * rows * self._hidden_size * 2) for (rows, grad), ref in zip(group_rows, refs, strict=True) ) - if tp > 1 and gradient_rows: + if (tp > 1 or dense_cp) and gradient_rows: # The recomputed layer's workspace, for the largest gradient group: # groups run their backward one after another. workspace = max( @@ -975,7 +986,7 @@ def _checkpoint_floor_from_facts( if grad ), ) - if tp > 1 and self._gdn_layers and gradient_rows: + if (tp > 1 or dense_cp) and self._gdn_layers and gradient_rows: # Recurrent states grow with segments, not rows; backward recomputes # one layer at a time. Padding to TP adds up to TP - 1 one-token # roots per group. Kernel-internal chunk states are not bounded here. @@ -1065,15 +1076,19 @@ def _sequence_parallel_lora_floor( group_rows: tuple[tuple[int, bool], ...], signature: _impl._MemorySignature, ) -> tuple[int, int]: - """Add the recomputed layer's LoRA ``x @ A`` intermediates to a TP x SP floor. + """Add the recomputed layer's LoRA row-by-rank intermediates. Each adapted module of the recomputed layer saves rows x rank over the gathered rows (FC1 gate/up measured 16 B/row at rank 4). Charge the decoder's most-adapted layer at the gradient slots' largest rank, from the signature's slot shapes ((ndim, *A.shape, *B.shape) per module). + Native dense CP also reserves that many bytes for backward factor-product + gradients beside the saved forward products. Counting every target at the + largest selected rank bounds layers with different target/rank inventories. """ retained, workspace = floor - if not retained or self._topology_key()[1] == 1: + dense_cp = _native_dense_cp(self) + if not retained or (self._topology_key()[1] == 1 and not dense_cp): return floor rank = max( ( @@ -1081,7 +1096,7 @@ def _sequence_parallel_lora_floor( for grad, shapes in signature.slot_shapes if grad for shape in shapes - if len(shape) == 5 and shape[0] == 2 + if len(shape) == 5 and shape[0] == 2 and shape[2] == shape[3] ), default=0, ) @@ -1093,7 +1108,35 @@ def _sequence_parallel_lora_floor( (len(shapes) for grad, shapes in signature.slot_shapes if grad), default=0 ) rows = max((rows for rows, grad in group_rows if grad), default=0) - return retained, workspace + rows * modules * rank * self._param_dtype_size + copies = 2 if dense_cp else 1 + lora_bytes = copies * rows * modules * rank * self._param_dtype_size + if dense_cp and ( + not rank + or self._geometry.moe_experts + or self._geometry.moe_ffn_hidden_size + or not all( + len(shapes) >= modules + and all( + not shape # A recognized owner with an inactive selected slot. + or ( + len(shape) == 5 + and shape[0] == 2 + and min(shape) > 0 + and shape[2] == shape[3] + ) + for shape in shapes + ) + for grad, shapes in signature.slot_shapes + if grad + ) + ): + # Dense dX refinement requires a complete native 2-D inventory and + # supported geometry. (0,) records an unknown owner/implementation; + # missing/partial metadata is unknown even with zero recognized owners. + # Restore main's per-layer allowance beside the explicit workspace and + # any known rank charge, rather than letting unknown mean rank zero. + workspace += retained + return retained, workspace + lora_bytes def _cold_recompute_transient_bytes(self: TrainerRank) -> int: @@ -1110,19 +1153,19 @@ def _checkpoint_input_gradient_bytes( ) -> int: """The input-gradient term beside the checkpoint floor's boundaries. - TP1 charges one logical gradient per boundary, which also stands in for - the recompute workspace. Under TP x SP the workspace is priced - explicitly, and the recompute peak held one input-gradient shard of the - group being recomputed. + TP1's legacy boundary-gradient charge also stands in for workspace. + Native dense CP and TP x SP price that workspace explicitly, beside one + input-gradient shard of the group being recomputed. """ tp = self._topology_key()[1] + dense_cp = _native_dense_cp(self) from art.megatron.context_parallel.checkpoint import current_checkpoint_partition partition = current_checkpoint_partition() if partition is not None and retained: # One full recomputed root's dX; its layer workspace is explicit. return math.prod(partition[1]) * self._param_dtype_size - if not retained or tp == 1: + if not retained or (tp == 1 and not dense_cp): return retained rows = max((rows for rows, grad in group_rows if grad), default=0) return -(-rows // tp) * self._hidden_size * self._param_dtype_size @@ -1600,6 +1643,11 @@ def _memory_signature_from_requests( raise ValueError("checkpoint shards require one CP1 uniform-recompute root") modes = tuple(sorted(grad_modes)) topology = self._topology_key() + slot_groups = tuple(slot_groups) + if not slot_groups and _native_dense_cp(self): + # Direct estimator callers use the template adapter when no slot + # selection was supplied, just as _slot_lora_tensors(None) does. + slot_groups = tuple((None, grad) for grad in modes) shapes = tuple( sorted((grad, self._slot_memory_shapes(ref)) for ref, grad in slot_groups) ) @@ -1621,7 +1669,11 @@ def _memory_signature_from_requests( and requests[0].input_tokens.ndim == 1 and (modes == (False,) or (modes == (True,) and _DIRECT_FORWARD.get())) and _impl._request_mix_key(requests[0]) in ("hidden", "hidden+routes") - and topology[1:] == (1, 1, 1) + and topology[1] == topology[3] == 1 + and ( + topology[2] == 1 + or (topology[2] in (2, 4, 8) and self._parallel_shape.cp == topology[2]) + ) and self._parallel_shape.ep == self._parallel_shape.etp == 1 and bool( _checkpoint_layers(self, ((int(requests[0].input_tokens.numel()), False),)) @@ -1646,17 +1698,21 @@ def _slot_memory_shapes( ) -> tuple[tuple[int, ...], ...]: """Separate empirical trust across actual selected adapter layouts. - Sequence-parallel (TP > 1) models record them too: the explicit TP x SP - checkpoint floor prices LoRA intermediates from these ranks. + TP x SP and native dense CP checkpoint floors price LoRA intermediates + from these ranks, including attention-only dense models and the template + adapter selected by a None reference. For dense CP, (0,) marks incomplete + classification; () is a recognized owner with an inactive selected slot. """ + dense_cp = _native_dense_cp(self) if ( - ref is None - or ref.name is None + (ref is None and not dense_cp) + or (ref is not None and ref.name is None) or isinstance(ref, _impl._LocalLoRASlotRef) or not ( getattr(self, "_moe_layers", 0) or getattr(self, "_gdn_layers", 0) or getattr(self, "_sequence_parallel", False) + or dense_cp ) ): # Generic/no-component planning must not import Megatron or walk @@ -1665,15 +1721,73 @@ def _slot_memory_shapes( from art.megatron.lora import LoRA shapes = [] + owned: set[int] = set() for chunk in self.runtime.model: for module in chunk.modules(): - if type(module) is LoRA: - tensors = _impl._slot_lora_tensors(module, ref) - shapes.append( - () - if tensors is None - else (tensors[0].ndim, *tensors[0].shape, *tensors[1].shape) + if id(module) in owned: + continue + if isinstance(module, LoRA): + if dense_cp: + # Slot Parameters belong to this recognized owner; they + # are not separate, unclassified adapter implementations. + owned.update( + id(child) for child in module.modules() if child is not module + ) + try: + tensors = _impl._slot_lora_tensors(module, ref) + shape = ( + () + if tensors is None + else (tensors[0].ndim, *tensors[0].shape, *tensors[1].shape) + ) + if ( + dense_cp + and tensors is not None + and any( + tensor.element_size() != self._param_dtype_size + for tensor in tensors + ) + ): + shapes.append((0,)) + except (AttributeError, RuntimeError, TypeError): + shape = (0,) + shapes.append(shape) + if dense_cp and ( + any( + name in vars(module) + or getattr(type(module), name) is not getattr(LoRA, name) + for name in ( + "forward", + "active_lora_tensors", + "_slot", + "_active_lora_storage", + "_expert_views", + ) + ) + or module._forward_hooks + or module._forward_pre_hooks + ): + # Inherited native behavior is priced; an overridden + # kernel/selection path needs its own workspace proof. + shapes.append((0,)) + elif dense_cp and ( + any(param.requires_grad for param in module.parameters(recurse=False)) + or any( + hasattr(module, name) + for name in ( + "A_T", + "B_T", + "lora_A", + "lora_B", + "_lora_params", + "lora_slot_params", + ) ) + ): + # All trainable tensors must be owned by a recognized adapter. + # Also catch frozen/inactive foreign factor owners. Never infer + # completeness from the number of recognized owners alone. + shapes.append((0,)) return tuple(shapes) @@ -1833,15 +1947,18 @@ def _estimate_required_memory_bytes_from_values( not self._geometry.moe_experts and self._geometry.ffn_hidden_size and self._dense_fc1_adapted - and signature.topology[1:3] == (1, 1) + and signature.topology[1] == 1 + and ( + signature.topology[2] == 1 + or getattr(signature, "bounded_forward_peak", False) + ) ): # A dense no-grad layer peaks at its FC1 stage (the base GEMM output, # the adapter output and their sum: 6F, or the SwiGLU live set if # wider) beside the residual pair, embedding and norm output (4H), per - # row: Qwen3.8-27B TP1/CP1 traces at 7k-174k rows. At CP1 every packed - # row is local, which the per-packed-token floor above prices at only - # 16H. Groups run one after another, so the largest no-grad group - # bounds it. + # row: Qwen3.8-27B TP1/CP1 traces at 7k-174k rows. Under CP use actual + # physical rank rows, never a global count divided by degree. Groups + # run one after another, so the largest no-grad group bounds it. rows = max( (n for n, grad in group_rows if not grad), default=0 if signature.grad_enabled else packed_tokens, diff --git a/src/art/trainer_rank/_micro_batch_planner.py b/src/art/trainer_rank/_micro_batch_planner.py index 8321167a8..11efeed39 100644 --- a/src/art/trainer_rank/_micro_batch_planner.py +++ b/src/art/trainer_rank/_micro_batch_planner.py @@ -1468,6 +1468,8 @@ def _fill_planner_snapshot( ) rank_fields["geometry"] = asdict(self._geometry) rank_fields["topology"] = list(plan.signature.topology) + rank_fields["parallel_shape"] = asdict(self._parallel_shape) + rank_fields["moe_layers"] = self._moe_layers def version(tensor: torch.Tensor | None) -> int | None: try: diff --git a/src/art/trainer_rank/_planner_misses.py b/src/art/trainer_rank/_planner_misses.py index e6afb6d9f..9b136e920 100644 --- a/src/art/trainer_rank/_planner_misses.py +++ b/src/art/trainer_rank/_planner_misses.py @@ -646,17 +646,19 @@ def report( "moe_forward_stages".split() ) # TP x SP floor facts; reports captured before them replay conservatively. -_OPTIONAL_RANK_FIELDS = frozenset({"mixer_top_gaps", "lora_modules_per_layer"}) +_OPTIONAL_RANK_FIELDS = frozenset( + {"mixer_top_gaps", "lora_modules_per_layer", "parallel_shape", "moe_layers"} +) def _adapter_ranks_unavailable( topology: tuple[int, ...], facts: dict[str, Any], signature: Any ) -> bool: - """A TP x SP floor estimate trains a named adapter whose ranks (the LoRA + """A TP x SP or CP floor estimate trains a named adapter whose ranks (the LoRA intermediates' price) were not recorded, as in reports from before sequence-parallel signatures carried slot shapes.""" return ( - topology[1] > 1 + (topology[1] > 1 or topology[2] in (2, 4, 8)) and bool(facts["checkpoint_layers"]) and any( group["grad"] and group["adapter"] is not None for group in facts["groups"] @@ -766,6 +768,23 @@ def replay( rank._geometry = ModelGeometry(**values["geometry"]) dp, tp, cp, pp = values["topology"] rank._topology_key = lambda: (dp, tp, cp, pp) + # CP workspace eligibility uses the runtime shape and observed MoE layer + # count, not merely the requested topology or configured expert count. + if cp > 1 and not {"parallel_shape", "moe_layers"} <= values.keys(): + raise ValueError("incomplete replay: CP runtime shape/layers not recorded") + shape = values.get("parallel_shape", {"tp": tp, "cp": cp}) + if ( + type(shape) is not dict + or set(shape) - {"tp", "cp", "ep", "etp"} + or ("parallel_shape" in values and set(shape) != {"tp", "cp", "ep", "etp"}) + or any(type(v) is not int or not 1 <= v <= 1024 for v in shape.values()) + ): + raise ValueError("incomplete replay: invalid parallel shape") + rank._parallel_shape = _impl.ParallelShape(**shape) + moe_layers = values.get("moe_layers", int(bool(rank._geometry.moe_experts))) + if type(moe_layers) is not int or not 0 <= moe_layers <= layers: + raise ValueError("incomplete replay: invalid MoE layer count") + rank._moe_layers = moe_layers # Check the same shared subforward inventories as capture, including the # independently retained layout inputs, before constructing any prefix tree. group_offset = request_offset = 0 diff --git a/tests/unit/test_hybridep_planner_replay.py b/tests/unit/test_hybridep_planner_replay.py index ff5dbdc54..a0116fb2d 100644 --- a/tests/unit/test_hybridep_planner_replay.py +++ b/tests/unit/test_hybridep_planner_replay.py @@ -128,6 +128,9 @@ def test_checkpoint_extent_is_frozen_without_retaining_live_graph(live, monkeypa replay._facts = facts replay._hidden_size = rank._hidden_size replay._gdn_layers = 0 + replay._parallel_shape = rank._parallel_shape + replay._moe_layers = rank._moe_layers + replay._dense_fc1_adapted = rank._dense_fc1_adapted replay._moe_output_bytes_per_token = rank._moe_output_bytes_per_token replay._topology_key = lambda: (1, 1, 2, 1) replay._moe_checkpoint_grad_bytes_per_token = facts[ @@ -140,3 +143,37 @@ def test_checkpoint_extent_is_frozen_without_retaining_live_graph(live, monkeypa assert ref() is None assert replay._checkpoint_memory_floor(rows) == expected assert not torch.cuda.is_initialized() + + +@pytest.mark.parametrize("cp", [2, 4, 8]) +@pytest.mark.parametrize("dense", [False, True]) +@pytest.mark.parametrize("matching", [False, True]) +def test_cp_report_replays_runtime_shape_and_layers( + cp, dense, matching, monkeypatch, tmp_path +): + rank = head_rank() + rank._parallel_shape = replace(rank._parallel_shape, cp=cp if matching else 1) + rank._dense_fc1_adapted = dense + rank._moe_layers = 0 if dense else 2 + rank.runtime.model_support_handler.context_parallel_workload_profile = lambda _: ( + None + ) + monkeypatch.setattr( + rank, "_topology", lambda: SimpleNamespace(tp=1, cp=cp, dp=1, pp=1, sp=False) + ) + rank._planner_reporter = reports.Reporter(0, spool_dir=tmp_path) + report, costs = emitted( + rank, rank._plan_flat_forward([request(32, grad=True)]), tmp_path + ) + assert report["replay_complete"], report["incomplete_reasons"] + result = reports.replay(report) + assert result["aggregate"]["matches"] + assert result["estimates"][0]["required_bytes"] == costs[0].required + rank._parallel_shape = replace(rank._parallel_shape, cp=16) + rank._moe_layers = 4 + assert reports.replay(report) == result + # An old CP report cannot invent runtime eligibility from requested CP. + del report["replay"]["memory_replay"]["rank"]["parallel_shape"] + with pytest.raises(ValueError, match="CP runtime shape/layers not recorded"): + reports.replay(report) + assert not torch.cuda.is_initialized() diff --git a/tests/unit/test_trainer_rank_checkpoint_memory.py b/tests/unit/test_trainer_rank_checkpoint_memory.py index bfafe3501..39d7ea6b8 100644 --- a/tests/unit/test_trainer_rank_checkpoint_memory.py +++ b/tests/unit/test_trainer_rank_checkpoint_memory.py @@ -507,7 +507,7 @@ def dense_rank(): config.ffn_hidden_size, config.bias_activation_fusion = 17408, True decoder = r.runtime.model[0].decoder decoder.layers = torch.nn.ModuleList( - [torch.nn.Linear(1, 1).bfloat16() for _ in range(64)] + [torch.nn.Linear(1, 1).bfloat16().requires_grad_(False) for _ in range(64)] ) decoder.num_layers_per_pipeline_rank = 64 r.runtime.provider = SimpleNamespace( @@ -522,6 +522,326 @@ def dense_rank(): return dense +@pytest.mark.parametrize("cp", [2, 4, 8]) +@pytest.mark.parametrize("adapter_rank", [8, 768]) +@pytest.mark.parametrize("subclass", [False, True]) +def test_dense_cp_prices_loaded_lora_products_without_opt_in( + cp, adapter_rank, subclass +): + from test_trainer_rank_slot_memory import load_slot + + from art.megatron.lora import LoRA + from art.trainer_rank._impl import _lora_modules_per_layer + + class DerivedLoRA(LoRA): + pass + + r = dense_rank() + r._hidden_size = r.runtime.model[0].config.hidden_size = 32 + r._geometry = replace( + r._geometry, + hidden_size=32, + ffn_hidden_size=64, + num_attention_heads=1, + num_query_groups=1, + kv_channels=32, + ) + r._mixer_top_gaps = (0, None) + r._topology_key = lambda: (1, 1, cp, 1) + r._parallel_shape = replace(r._parallel_shape, cp=cp) + # q/k/v/output/gate/up/down: the selected rank differs from the template. + loras = torch.nn.ModuleList( + [ + (DerivedLoRA if subclass else LoRA)( + f"target.{i}", 32, 32, 4, 4, torch.bfloat16, torch.device("cpu") + ) + for i in range(7) + ] + ) + r.runtime.model[0].decoder.layers[0].loras = loras + r._lora_modules_per_layer = _lora_modules_per_layer(r.runtime.model[0]) + if subclass: + # Use the real loader for inherited owners too; the shared test helper + # deliberately handles only its original exact-class fixtures. + from art.megatron.lora import LoRASlotRef + + ref = LoRASlotRef("checkpoint", "selected") + for owner in loras: + assert isinstance(owner, LoRA) + prefix = owner.adapter_model_prefix + assert owner.load_lora_slot( + ref, + { + prefix + ".lora_A.weight": torch.zeros(adapter_rank, 32).bfloat16(), + prefix + ".lora_B.weight": torch.zeros(32, adapter_rank).bfloat16(), + }, + requires_grad=True, + ) + else: + ref = load_slot(r, "selected", adapter_rank) + assert r._lora_modules_per_layer == 7 + rows = 65_536 + request = ForwardInput( + input_tokens=torch.zeros(2 * rows, dtype=torch.long), hidden_states=True + ) + signature = r._memory_signature_from_requests( + [request], slot_group_count=1, grad_modes=(True,), slot_groups=((ref, True),) + ) + assert len(signature.slot_shapes[0][1]) == 7 + assert {shape[2] for shape in signature.slot_shapes[0][1]} == {adapter_rank} + assert not getattr(signature, "bounded_forward_peak", False) + groups = ((rows, True),) + retained, workspace = r._checkpoint_memory_floor(groups) + # All targets' saved products plus backward factor-product gradients. + lora_bytes = 2 * rows * 7 * adapter_rank * 2 + assert r._sequence_parallel_lora_floor( + (retained, workspace), groups, signature + ) == (retained, workspace + lora_bytes) + required = r._estimate_required_memory_bytes_from_values( + packed_tokens=2 * rows, + logical_tokens=2 * rows, + output_bytes=2 * rows * 32 * 2, + signature=signature, + group_rows=groups, + ) + assert (required <= 550_000_000) is (adapter_rank == 8) + # Even the six products measured live by FC1 in Sol's CPU witness exceed + # that cap at rank 768. The full charge also covers the seventh target. + assert required >= rows * 6 * adapter_rank * 2 + # The literal original probe omitted slot ranks despite declaring seven + # targets. That must preserve the old allowance, not imply rank zero. + missing = replace(signature, slot_shapes=()) + assert r._sequence_parallel_lora_floor((retained, workspace), groups, missing) == ( + retained, + workspace + retained, + ) + assert ( + r._estimate_required_memory_bytes_from_values( + packed_tokens=2 * rows, + logical_tokens=2 * rows, + output_bytes=2 * rows * 32 * 2, + signature=missing, + group_rows=groups, + ) + > 550_000_000 + ) + assert not torch.cuda.is_initialized() + + +@pytest.mark.parametrize("mixed", [False, True]) +@pytest.mark.parametrize( + "kind", ["foreign", "custom forward", "malformed factors", "unmarked"] +) +def test_dense_cp_unknown_adapter_inventory_keeps_legacy_allowance(mixed, kind): + from art.megatron.lora import LoRA + from art.trainer_rank._impl import _lora_modules_per_layer + + class CustomLoRA(LoRA): + def forward(self, x, tokens_per_expert=None): + return super().forward(x, tokens_per_expert) + + r = dense_rank() + model = r.runtime.model[0] + model.requires_grad_(False) + owner = ( + CustomLoRA("custom", 32, 32, 8, 8, torch.bfloat16, torch.device("cpu")) + if kind == "custom forward" + else torch.nn.Module() + if kind in ("foreign", "unmarked") + else LoRA("bad", 32, 32, 8, 8, torch.bfloat16, torch.device("cpu")) + ) + if kind == "foreign": + owner.A_T = torch.nn.Parameter(torch.zeros(32, 512).bfloat16()) + owner.B_T = torch.nn.Parameter(torch.zeros(512, 32).bfloat16()) + elif kind == "malformed factors": + owner.B_T = torch.nn.Parameter(torch.zeros(9, 32).bfloat16()) + elif kind == "unmarked": + owner.weight = torch.nn.Parameter(torch.zeros(32, 512).bfloat16()) + model.decoder.layers[0].unknown = owner + if mixed: + model.decoder.layers[0].known = LoRA( + "known", 32, 32, 8, 8, torch.bfloat16, torch.device("cpu") + ) + r._lora_modules_per_layer = _lora_modules_per_layer(model) + r._topology_key = lambda: (1, 1, 2, 1) + r._parallel_shape = replace(r._parallel_shape, cp=2) + req = ForwardInput(input_tokens=torch.arange(128), hidden_states=True) + sig = r._memory_signature_from_requests( + [req], slot_group_count=1, grad_modes=(True,) + ) + groups = ((64, True),) + retained, workspace = r._checkpoint_memory_floor(groups) + # Unknown inventory must restore the old per-layer allowance even when + # there are zero recognized owners, or other owners have usable ranks. + assert r._sequence_parallel_lora_floor((retained, workspace), groups, sig)[1] >= ( + workspace + retained + ) + args = dict( + packed_tokens=128, + logical_tokens=128, + output_bytes=128 * 5120 * 2, + signature=sig, + group_rows=groups, + ) + required = r._estimate_required_memory_bytes_from_values(**args) + r._dense_fc1_adapted = False + assert required >= r._estimate_required_memory_bytes_from_values(**args) + assert not torch.cuda.is_initialized() + + +def test_glm52_lora_subclass_witness_refuses_1150mb(monkeypatch): + from art.megatron.glm52.lora import Glm52LoRA + from art.megatron.lora import LoRA, LoRASlotRef, SharedExpertsLinearFC1LoRA + from art.trainer_rank._impl import _lora_modules_per_layer + + r = dense_rank() + model = r.runtime.model[0] + model.config.hidden_size = 32 + model.config.ffn_hidden_size = 64 + model.config.num_layers = model.decoder.num_layers_per_pipeline_rank = 128 + model.config.num_attention_heads = model.config.num_query_groups = 1 + model.config.kv_channels = 256 + model.config.num_moe_experts = 1 + r.runtime.model_support_handler.is_moe = True + layers = [] + for index in range(128): + + def lora(label, inputs=32, outputs=32): + return Glm52LoRA( + f"{index}.{label}", + inputs, + outputs, + 4, + 4, + torch.bfloat16, + torch.device("cpu"), + ) + + layer = torch.nn.Module() + layer.mlp = torch.nn.Module() + fc1 = SharedExpertsLinearFC1LoRA.__new__(SharedExpertsLinearFC1LoRA) + torch.nn.Module.__init__(fc1) + fc1.non_gated = False + fc1.gate_lora, fc1.up_lora = lora("gate", outputs=64), lora("up", outputs=64) + layer.mlp.linear_fc1 = fc1 + layer.other_loras = torch.nn.ModuleList( + [ + lora("q_a", 32, 512), + lora("kv_a", 32, 576), + lora("q_b", 512, 576), + lora("o", 256, 32), + lora("down", 64, 32), + ] + ) + layers.append(layer) + model.decoder.layers = torch.nn.ModuleList(layers) + model.embedding = torch.nn.Linear(1, 1).bfloat16().requires_grad_(False) + r = TrainerRank(r.runtime) + assert r._dense_fc1_adapted and not r._moe_layers + monkeypatch.setattr(r, "_topology_key", lambda: (1, 1, 2, 1)) + r._parallel_shape = replace(r._parallel_shape, cp=2) + selected = LoRASlotRef("checkpoint", "selected") + for owner in model.modules(): + if isinstance(owner, LoRA): + prefix = owner.adapter_model_prefix + assert owner.load_lora_slot( + selected, + { + prefix + ".lora_A.weight": torch.zeros( + 512, owner.in_features + ).bfloat16(), + prefix + ".lora_B.weight": torch.zeros( + owner.out_features, 512 + ).bfloat16(), + }, + alpha=512, + requires_grad=True, + ) + rows = 65_536 + req = ForwardInput(input_tokens=torch.arange(2 * rows), hidden_states=True) + sig = r._memory_signature_from_requests( + [req], + slot_group_count=1, + grad_modes=(True,), + slot_groups=((selected, True),), + ) + required = r._estimate_required_memory_bytes_from_values( + packed_tokens=2 * rows, + logical_tokens=2 * rows, + output_bytes=2 * rows * 32 * 2, + signature=sig, + group_rows=((rows, True),), + slot_refs=(selected,), + ) + monkeypatch.setattr(r, "_available_memory_bytes", lambda: 1_150_000_000) + assert not r._memory_check_required(required).fits + assert _lora_modules_per_layer(model) == 7 + assert ( + sum(len(shape) == 5 and shape[2] == 512 for shape in sig.slot_shapes[0][1]) + == 896 + ) + # Sol's CPU autograd storage witness already exceeds this cap before + # other activations, outputs, cold scratch and the safety factor. + assert required >= 1_174_405_120 + assert not torch.cuda.is_initialized() + + +# Frozen rank-8/all-target Qwen3.8-27B complete lifecycle receipts: +# workstreams/cp-dense/gpu/evidence/timed-v2-cp{2,4,8}/result.rank*.json; +# lc-qual recovered-r2/results/dense-n65536/capacity-cp{2,4}/result.rank*.json. +# Rows are the CPU planner's worst attention/GDN extent, not tokens / CP. +# Admission prices additional allocations; compare resident + demand to the +# measured allocated peak. Reserved peaks are allocator cache, not live bytes. +@pytest.mark.parametrize( + "cp,tokens,rows,resident,measured,usable", + [ + (2, 32768, 19456, 55332260352, 74358620160, 59837713613), + (4, 131072, 54528, 55332869632, 104346683392, 84904971469), + (8, 131072, 33280, 55332021760, 84065598464, 47617608909), + (2, 65536, 41472, 55333083648, 95180360704, 63154038375), + (4, 65536, 24576, 55332046336, 77951084544, 53037246055), + ], +) +def test_dense_cp_rank8_floor_covers_measured_lifecycles( + cp, tokens, rows, resident, measured, usable +): + r = dense_rank() + r._geometry = replace( + r._geometry, + num_attention_heads=24, + num_query_groups=4, + kv_channels=256, + gdn_key_heads=16, + gdn_value_heads=48, + gdn_key_head_dim=128, + gdn_value_head_dim=128, + gdn_conv_kernel=4, + ) + r._gdn_layers = 48 + r._attention_output_gate = True + r._mixer_top_gaps = (0, 1) + r._lora_modules_per_layer = 8 + r._topology_key = lambda: (1, 1, cp, 1) + r._parallel_shape = replace(r._parallel_shape, cp=cp) + req = ForwardInput(input_tokens=torch.arange(tokens), hidden_states=True) + sig = r._memory_signature_from_requests( + [req], slot_group_count=1, grad_modes=(True,) + ) + sig = replace(sig, slot_shapes=((True, ((2, 5120, 8, 8, 5120),) * 8),)) + # A small learned slope must never erase the analytic charge. + r._memory_profiles[sig] = _MemoryProfile(1, tokens) + required = r._estimate_required_memory_bytes_from_values( + packed_tokens=tokens, + logical_tokens=tokens, + output_bytes=tokens * 5120 * 2, + signature=sig, + group_rows=((rows, True),), + gdn_segments=1, + ) + assert resident + required >= measured + assert required <= usable + + # Qwen3.8-27B TP1/CP1 no-grad target-only waves (random init, ART a68fa500, # 2026-09-29): a cold first wave peaked at 260,789 B per row, warm waves at # 250,181-250,507; the floor they share is (6F + 4H) x 2 = 249,856. diff --git a/tests/unit/test_trainer_rank_forward_profile.py b/tests/unit/test_trainer_rank_forward_profile.py index 678836175..dc528a2a2 100644 --- a/tests/unit/test_trainer_rank_forward_profile.py +++ b/tests/unit/test_trainer_rank_forward_profile.py @@ -1,8 +1,10 @@ """Cold fixed scratch must not poison known single-root hidden-only inference.""" +from contextlib import ExitStack from dataclasses import fields, replace import pytest +from test_trainer_rank_checkpoint_memory import dense_rank from test_trainer_rank_moe_memory import layer as layer from test_trainer_rank_pending_memory import pending_rank as pending_rank import torch @@ -30,6 +32,84 @@ def observe(rank, plan, compute): ) +@pytest.mark.parametrize("cp", (2, 4, 8)) +@pytest.mark.parametrize("grad", (False, True)) +@pytest.mark.parametrize("kind", ("dense", "moe")) +def test_native_cp_bounded_calibration_keeps_physical_workspace_and_full_outputs( + bounded_rank, cp, grad, kind +): + rank = dense_rank() if kind == "dense" else bounded_rank + inputs = [replace(request(128), no_grad=not grad)] + with bounded_forward_peaks(), _memory.direct_forward_profile(): + control = rank._plan_flat_forward(inputs) + rank._topology_key = lambda: (1, 1, cp, 1) + rank._parallel_shape = replace(rank._parallel_shape, cp=cp) + signature = rank._memory_signature_from_requests( + inputs, slot_group_count=1, grad_modes=(grad,) + ) + assert signature.bounded_forward_peak and signature != control.signature + observe(rank, replace(control, signature=signature), 470221312) + profile = rank._memory_profiles[signature] + rows = 32768 // cp + 128 # Uneven ownership must use actual rows. + arguments = dict( + packed_tokens=32768, + logical_tokens=32768, + output_bytes=32768 * rank._hidden_size * 2, + signature=signature, + group_rows=((rows, grad),), + ) + cold = rank._estimate_required_memory_bytes_from_values(**arguments) + rank._memory_profiles[signature] = replace( + profile, forward_compute_peaks=((8192, 470221312),) + ) + calibrated = rank._estimate_required_memory_bytes_from_values(**arguments) + assert calibrated < cold + retained, workspace = rank._checkpoint_memory_floor(((rows, grad),)) + if grad: + assert retained == rows * rank._num_layers * rank._hidden_size * 2 + if kind == "dense": + assert workspace >= rank._sequence_parallel_workspace_bytes(rows) + assert rank._checkpoint_input_gradient_bytes(((rows, grad),), retained) == ( + rows * rank._hidden_size * 2 + ) + assert calibrated >= int( + (retained + workspace + 2 * arguments["output_bytes"]) * 1.1 + ) + rank._memory_profiles[signature] = replace( + profile, forward_compute_peaks=((8192, 0),) + ) + assert calibrated >= rank._estimate_required_memory_bytes_from_values(**arguments) + + +@pytest.mark.parametrize("cp,actual", ((2, 1), (3, 3), (16, 16), (4, 8))) +def test_dense_bounded_cp_requires_supported_matching_runtime_geometry(cp, actual): + rank = dense_rank() + rank._topology_key = lambda: (1, 1, cp, 1) + rank._parallel_shape = replace(rank._parallel_shape, cp=actual) + with bounded_forward_peaks(): + signature = rank._memory_signature_from_requests( + [request(128)], slot_group_count=1, grad_modes=(False,) + ) + assert not isinstance(signature, _memory._ScopedMemorySignature) + + +def test_bounded_dense_cp_no_grad_charges_the_worst_rank_fc1_stage(): + rank = dense_rank() + rank._topology_key = lambda: (1, 1, 8, 1) + rank._parallel_shape = replace(rank._parallel_shape, cp=8) + with bounded_forward_peaks(): + signature = rank._memory_signature_from_requests( + [request(128)], slot_group_count=1, grad_modes=(False,) + ) + required = rank._estimate_required_memory_bytes_from_values( + packed_tokens=32768, + output_bytes=0, + signature=signature, + group_rows=((32768, False),), + ) + assert required >= int(32768 * (6 * 17408 + 4 * 5120) * 2 * 1.1) + + def test_three_token_cold_peak_refuses_uncalibrated_81920(bounded_rank, monkeypatch): rank = bounded_rank short = rank._plan_flat_forward([request(3)]) @@ -403,3 +483,84 @@ def direct(): with bounded_forward_peaks() if enabled else nullcontext(): rank.dp_rank_forward([request(3)]) assert calls == ([True] if enabled else []) + + +@pytest.mark.parametrize("cp", (2, 4, 8)) +@pytest.mark.parametrize("grad", (False, True)) +def test_native_moe_cp_bounded_profiles_keep_static_floors(bounded_rank, cp, grad): + rank = bounded_rank + inputs = [replace(request(128), no_grad=not grad)] + with _memory.direct_forward_profile(): + control = rank._plan_flat_forward(inputs) + rank._topology_key = lambda: (1, 1, cp, 1) + rank._parallel_shape = replace(rank._parallel_shape, cp=cp) + signature = rank._memory_signature_from_requests( + inputs, slot_group_count=1, grad_modes=(grad,) + ) + assert signature.bounded_forward_peak + assert signature.topology == (1, 1, cp, 1) + assert signature != control.signature + # Fixed scratch is an intercept. Preserve the existing physical-row and + # full-output floors; the calibration never divides a CP1 prediction. + observe(rank, replace(control, signature=signature), 470221312) + profile = rank._memory_profiles[signature] + assert profile.forward_compute_peaks == ((128, 470221312),) + arguments = dict( + packed_tokens=32768, + logical_tokens=32768, + output_bytes=32768 * 2048 * 2, + signature=signature, + group_rows=((32768 // cp, grad),), + ) + required = rank._estimate_required_memory_bytes_from_values(**arguments) + rank._memory_profiles[signature] = replace( + profile, forward_compute_peaks=((8192, 470221312),) + ) + calibrated = rank._estimate_required_memory_bytes_from_values(**arguments) + assert calibrated < required + rank._memory_profiles[signature] = replace( + profile, forward_compute_peaks=((8192, 0),) + ) + floor = rank._estimate_required_memory_bytes_from_values(**arguments) + assert calibrated >= floor + rank._memory_profiles.clear() + assert signature.topology == (1, 1, cp, 1) + + +# Values from origin/main e720a9e9, before bounded CP2/4/8 calibration, for one +# 32,768-token hidden root: cold required, required after one 128-token +# observation, checkpoint retained and workspace bytes, and the input-gradient +# term. The CP extension must leave every CP1 prediction where it was. +@pytest.mark.parametrize("scoped", (False, True)) +@pytest.mark.parametrize( + "kind,grad,expected", + ( + ("dense", False, (9375108300, 132783420211, 0, 1342177280, 0)), + ("dense", True, (47687558758, 132783420211, 21474836480, 0, 21474836480)), + ("moe", False, (7529614540, 132561960960, 0, 6710886400, 0)), + ("moe", True, (21066437427, 132561960960, 5368709120, 6178209792, 5368709120)), + ), +) +def test_cp1_predictions_are_unchanged_by_bounded_cp( + pending_rank, kind, grad, expected, scoped +): + rank = dense_rank() if kind == "dense" else pending_rank + assert rank._topology_key()[2] == rank._parallel_shape.cp == 1 + + def root(tokens): + ids = torch.arange(tokens) % 1900 + 100 + return ForwardInput(input_tokens=ids, hidden_states=True, no_grad=not grad) + + with ExitStack() as stack: + if scoped: + stack.enter_context(bounded_forward_peaks()) + stack.enter_context(_memory.direct_forward_profile()) + long, short = (rank._plan_flat_forward([root(n)]) for n in (32768, 128)) + assert getattr(long.signature, "bounded_forward_peak", False) is scoped + rows = rank._plan_group_rows(long) + retained, workspace = rank._checkpoint_memory_floor(rows) + cold = rank._memory_check(long).estimated_required_bytes + observe(rank, short, 470221312) + warm = rank._memory_check(long).estimated_required_bytes + gradient = rank._checkpoint_input_gradient_bytes(rows, retained) + assert (cold, warm, retained, workspace, gradient) == expected