From b32383647d5d5f0db1ea9b025432188d45543025 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Wed, 21 Oct 2026 01:02:08 +0000 Subject: [PATCH 01/13] fix: bound native MoE context parallel forward calibration (cherry picked from commit ed3dbb78c2dae42b992fee11e910a6a62bc6293d) Co-Authored-By: Claude Opus 5.5 --- src/art/trainer_rank/_memory.py | 11 ++++- .../unit/test_trainer_rank_forward_profile.py | 42 +++++++++++++++++++ 2 files changed, 52 insertions(+), 1 deletion(-) diff --git a/src/art/trainer_rank/_memory.py b/src/art/trainer_rank/_memory.py index e86508a21..6bb1db755 100644 --- a/src/art/trainer_rank/_memory.py +++ b/src/art/trainer_rank/_memory.py @@ -1621,7 +1621,16 @@ 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._moe_layers + and self._moe_memory_supported + ) + ) and self._parallel_shape.ep == self._parallel_shape.etp == 1 and bool( _checkpoint_layers(self, ((int(requests[0].input_tokens.numel()), False),)) diff --git a/tests/unit/test_trainer_rank_forward_profile.py b/tests/unit/test_trainer_rank_forward_profile.py index 678836175..b16813093 100644 --- a/tests/unit/test_trainer_rank_forward_profile.py +++ b/tests/unit/test_trainer_rank_forward_profile.py @@ -403,3 +403,45 @@ 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) From 9e97cb204cbad38a989668f6f76bd53fc891e898 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Wed, 21 Oct 2026 01:04:14 +0000 Subject: [PATCH 02/13] fix: keep bounded context parallel scope boolean (cherry picked from commit 8343c73a83010b4317a11522d5c15afd5fc02ff1) Co-Authored-By: Claude Opus 5.5 --- src/art/trainer_rank/_memory.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/art/trainer_rank/_memory.py b/src/art/trainer_rank/_memory.py index 6bb1db755..a5fa27658 100644 --- a/src/art/trainer_rank/_memory.py +++ b/src/art/trainer_rank/_memory.py @@ -1627,7 +1627,7 @@ def _memory_signature_from_requests( or ( topology[2] in (2, 4, 8) and self._parallel_shape.cp == topology[2] - and self._moe_layers + and self._moe_layers > 0 and self._moe_memory_supported ) ) From 4d60d597f647242cce3c45819faffc4b032210a2 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Wed, 21 Oct 2026 02:42:33 +0000 Subject: [PATCH 03/13] fix: bound native dense context parallel memory calibration The bounded-profile CP2/4/8 scope conflicted with the MoE-only scope picked just before it. Keep this general condition, as lc-qual's b31af924 integration did: the signature's trailing clause still requires MoE memory support for MoE models and an adapted FC1 for dense ones. (cherry picked from commit bfedd0be47931e8d19ba2bb01bb67eef91077995) Co-Authored-By: Claude Opus 5.5 --- src/art/trainer_rank/_memory.py | 27 ++++--- .../unit/test_trainer_rank_forward_profile.py | 76 +++++++++++++++++++ 2 files changed, 89 insertions(+), 14 deletions(-) diff --git a/src/art/trainer_rank/_memory.py b/src/art/trainer_rank/_memory.py index a5fa27658..b2d8024e3 100644 --- a/src/art/trainer_rank/_memory.py +++ b/src/art/trainer_rank/_memory.py @@ -948,7 +948,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, cp, _ = self._topology_key() + dense_cp = cp in (2, 4, 8) and not self._moe_layers and self._dense_fc1_adapted # 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 +965,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 +976,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. @@ -1624,12 +1625,7 @@ def _memory_signature_from_requests( 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._moe_layers > 0 - and self._moe_memory_supported - ) + 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( @@ -1842,15 +1838,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/tests/unit/test_trainer_rank_forward_profile.py b/tests/unit/test_trainer_rank_forward_profile.py index b16813093..331ec3e32 100644 --- a/tests/unit/test_trainer_rank_forward_profile.py +++ b/tests/unit/test_trainer_rank_forward_profile.py @@ -3,6 +3,7 @@ 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 +31,81 @@ 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 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)]) From 18193a8eb1aa35d1b912ef8547f9c64944dd0c69 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Wed, 21 Oct 2026 03:06:04 +0000 Subject: [PATCH 04/13] fix: avoid duplicate native CP recompute workspace charges (cherry picked from commit faae9cfed73ebdecf89566b41a3215ae50f5ad44) Co-Authored-By: Claude Opus 5.5 --- src/art/trainer_rank/_memory.py | 24 +++++++++++++------ .../unit/test_trainer_rank_forward_profile.py | 3 +++ 2 files changed, 20 insertions(+), 7 deletions(-) diff --git a/src/art/trainer_rank/_memory.py b/src/art/trainer_rank/_memory.py index b2d8024e3..7885f1621 100644 --- a/src/art/trainer_rank/_memory.py +++ b/src/art/trainer_rank/_memory.py @@ -949,7 +949,12 @@ def _checkpoint_floor_from_facts( return 0, 0 gradient_rows = sum(rows for rows, grad in group_rows if grad) _, tp, cp, _ = self._topology_key() - dense_cp = cp in (2, 4, 8) and not self._moe_layers and self._dense_fc1_adapted + dense_cp = ( + cp in (2, 4, 8) + and self._parallel_shape.cp == cp + and not self._moe_layers + and self._dense_fc1_adapted + ) # 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) @@ -1111,19 +1116,24 @@ 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] + _, tp, cp, _ = self._topology_key() + dense_cp = ( + cp in (2, 4, 8) + and self._parallel_shape.cp == cp + and not self._moe_layers + and self._dense_fc1_adapted + ) 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 diff --git a/tests/unit/test_trainer_rank_forward_profile.py b/tests/unit/test_trainer_rank_forward_profile.py index 331ec3e32..e73717be7 100644 --- a/tests/unit/test_trainer_rank_forward_profile.py +++ b/tests/unit/test_trainer_rank_forward_profile.py @@ -68,6 +68,9 @@ def test_native_cp_bounded_calibration_keeps_physical_workspace_and_full_outputs 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 ) From 82284f40d88c28139380e2ddacf3db8377833b54 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Wed, 21 Oct 2026 20:06:31 +0000 Subject: [PATCH 05/13] feat: bound native dense CP checkpoint activations on host Add opt-in synchronous saved-boundary offload with a fixed declared host budget and separate admission signature. Keep one device boundary, recompute workspace, all pending adapter gradients and complete outputs priced; reject host over-budget plans before forward. Defaults and GPU safety budgets remain unchanged. (cherry picked from commit 67170e6080609e071b5a0261e93c0f10880e210e) Co-Authored-By: Claude Opus 5.5 --- .../megatron/context_parallel/checkpoint.py | 94 ++++++++++++++- src/art/trainer_rank/_memory.py | 54 ++++++++- tests/unit/test_checkpoint_host_offload.py | 112 ++++++++++++++++++ 3 files changed, 254 insertions(+), 6 deletions(-) create mode 100644 tests/unit/test_checkpoint_host_offload.py diff --git a/src/art/megatron/context_parallel/checkpoint.py b/src/art/megatron/context_parallel/checkpoint.py index ac0da17dc..79bf7ef65 100644 --- a/src/art/megatron/context_parallel/checkpoint.py +++ b/src/art/megatron/context_parallel/checkpoint.py @@ -26,6 +26,9 @@ "art_checkpoint_partition", default=None ) _HOOKS: ContextVar[Any] = ContextVar("art_checkpoint_shard_hooks", default=None) +_HOST_OFFLOAD: ContextVar[HostActivationStats | None] = ContextVar( + "art_checkpoint_host_offload", default=None +) def _install_checkpoint_hooks() -> None: @@ -59,7 +62,7 @@ def local(*inputs: Any) -> Any: def pack(tensor: torch.Tensor) -> Any: nonlocal packed saved = hooks[0](tensor) - if isinstance(saved, _Shard): + if isinstance(saved, (_Shard, _HostBoundary)): packed += 1 if packed > 1: raise ValueError( @@ -68,7 +71,10 @@ def pack(tensor: torch.Tensor) -> Any: return saved with torch.autograd.graph.saved_tensors_hooks(pack, hooks[1]): - return original(local, *args) + result = original(local, *args) + if _HOST_OFFLOAD.get() is not None and packed != 1: + raise ValueError("host offload requires one checkpoint hidden input") + return result setattr(apply, "_art_checkpoint_shards", True) function.apply = staticmethod(apply) @@ -101,6 +107,88 @@ class _Shard: numel: int +@dataclass +class HostActivationStats: + hidden_size: int + max_cpu_bytes: int + saved_bytes: int = 0 + boundaries: int = 0 + unpacked: int = 0 + verified: int = 0 + + +@dataclass(frozen=True) +class _HostBoundary: + tensor: torch.Tensor + device: torch.device + + +def current_host_activation_offload() -> HostActivationStats | None: + return _HOST_OFFLOAD.get() + + +@contextmanager +def offload_saved_activations( + *, hidden_size: int, max_cpu_bytes: int = 64 * 2**30, verify: bool = False +) -> Iterator[HostActivationStats]: + """Copy full/uniform checkpoint inputs to host; retain native CP arithmetic. + + Only the checkpoint boundary's contiguous [rows, 1, hidden] input moves. + The layer's ordinary autograd saves stay on device. Copies preserve dtype; + synchronous transfers make lifetime and completion explicit. Host storage + is bounded per scope, independent of the device admission budgets. + """ + if type(hidden_size) is not int or hidden_size <= 0: + raise ValueError("host offload requires a positive hidden size") + if type(max_cpu_bytes) is not int or max_cpu_bytes <= 0: + raise ValueError("host offload requires a positive CPU byte budget") + if _HOOKS.get() is not None or _HOST_OFFLOAD.get() is not None: + raise RuntimeError("checkpoint storage scopes cannot overlap") + stats = HostActivationStats(hidden_size, max_cpu_bytes) + + def pack(tensor: torch.Tensor) -> torch.Tensor | _HostBoundary: + if tensor.ndim != 3 or tensor.shape[-1] != hidden_size: + return tensor + if ( + tensor.shape[1] != 1 + or not tensor.is_floating_point() + or not tensor.is_contiguous() + ): + raise ValueError("host offload requires contiguous [rows, 1, hidden] input") + size = tensor.numel() * tensor.element_size() + if stats.saved_bytes + size > max_cpu_bytes: + raise ValueError("checkpoint host activation byte budget exceeded") + host = torch.empty_like( + tensor, device="cpu", pin_memory=tensor.device.type == "cuda" + ) + host.copy_(tensor.detach(), non_blocking=False) + if verify: + if not torch.equal( + host.to(tensor.device).view(torch.uint8), + tensor.detach().view(torch.uint8), + ): + raise ValueError("checkpoint host round-trip changed input bytes") + stats.verified += 1 + stats.saved_bytes += size + stats.boundaries += 1 + return _HostBoundary(host, tensor.device) + + def unpack(value: torch.Tensor | _HostBoundary) -> torch.Tensor: + if isinstance(value, torch.Tensor): + return value + stats.unpacked += 1 + return value.tensor.to(value.device, non_blocking=False) + + _install_checkpoint_hooks() + token = _HOST_OFFLOAD.set(stats) + hooks = _HOOKS.set((pack, unpack)) + try: + yield stats + finally: + _HOOKS.reset(hooks) + _HOST_OFFLOAD.reset(token) + + @contextmanager def shard_saved_activations( group: Any, *, shape: tuple[int, int, int], key: str @@ -119,7 +207,7 @@ def shard_saved_activations( tensors, while recomputation's temporary saves should remain local. Keep the numerical-policy scope open across both forward and backward. """ - if _PARTITION.get() is not None: + if _PARTITION.get() is not None or _HOST_OFFLOAD.get() is not None: raise RuntimeError("checkpoint sharding scopes cannot overlap") if len(shape) != 3 or shape[1] != 1 or any(n <= 0 for n in shape): raise ValueError("checkpoint activation shape must be [tokens, 1, hidden]") diff --git a/src/art/trainer_rank/_memory.py b/src/art/trainer_rank/_memory.py index 7885f1621..c24611e12 100644 --- a/src/art/trainer_rank/_memory.py +++ b/src/art/trainer_rank/_memory.py @@ -36,6 +36,7 @@ class _ScopedMemorySignature(_impl._MemorySignature): checkpoint_partition: tuple[int, tuple[int, int, int]] | None = None bounded_forward_peak: bool = False + checkpoint_host_offload: tuple[int, int] | None = None @dataclass(frozen=True) @@ -962,6 +963,25 @@ def _checkpoint_floor_from_facts( * self._hidden_size * 2 ) + from art.megatron.context_parallel.checkpoint import current_host_activation_offload + + offload = current_host_activation_offload() + if offload is not None: + if not dense_cp or offload.hidden_size != self._hidden_size: + raise ValueError( + "checkpoint host offload requires matching native dense CP" + ) + if retained > offload.max_cpu_bytes: + raise ValueError( + "checkpoint host activation byte budget exceeded before forward" + ) + # All saved boundaries move to bounded host storage. One full input + # still exists on device while packing or unpacking a checkpoint. + retained = ( + max((rows for rows, grad in group_rows if grad), default=0) + * 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 @@ -1245,11 +1265,15 @@ def _checkpoint_gradient_groups( decoder layer saves the group's rows (this rank's TP shard). The base model (no name) has no adapter slot. """ + from art.megatron.context_parallel.checkpoint import current_host_activation_offload + tp = self._topology_key()[1] + offload = current_host_activation_offload() is not None return tuple( ( ref if ref is not None and ref.name is not None else None, - (-(-rows // tp) * self._hidden_size * 2,) * self._num_layers, + ((0 if offload else -(-rows // tp) * self._hidden_size * 2),) + * self._num_layers, ) for (rows, grad), ref in zip( group_rows, slot_refs or (None,) * len(group_rows), strict=True @@ -1593,9 +1617,30 @@ def _memory_signature_from_requests( grad_modes: Iterable[bool], slot_groups: Iterable[tuple["LoRASlotRef | None", bool]] = (), ) -> _impl._MemorySignature: - from art.megatron.context_parallel.checkpoint import current_checkpoint_partition + from art.megatron.context_parallel.checkpoint import ( + current_checkpoint_partition, + current_host_activation_offload, + ) partition = current_checkpoint_partition() + offload = current_host_activation_offload() + if offload is not None and ( + partition is not None + or len(requests) != 1 + or self._topology_key() not in ((1, 1, 2, 1), (1, 1, 4, 1), (1, 1, 8, 1)) + or self._parallel_shape.cp != self._topology_key()[2] + or self._parallel_shape.ep != 1 + or self._parallel_shape.etp != 1 + or self._moe_layers + or not self._dense_fc1_adapted + or offload.hidden_size != self._hidden_size + or not _checkpoint_layers( + self, ((int(requests[0].input_tokens.numel()), True),) + ) + ): + raise ValueError( + "checkpoint host offload requires one native dense CP uniform root" + ) if partition is not None: # Native MoE storage sharding failed full-model gradient parity on H200. if self._moe_layers: @@ -1647,11 +1692,14 @@ def _memory_signature_from_requests( else self._dense_fc1_adapted and bool(self._geometry.ffn_hidden_size) ) ) - if partition is not None or bounded: + if partition is not None or bounded or offload is not None: return _ScopedMemorySignature( **vars(signature), checkpoint_partition=partition, bounded_forward_peak=bounded, + checkpoint_host_offload=None + if offload is None + else (offload.hidden_size, offload.max_cpu_bytes), ) return signature diff --git a/tests/unit/test_checkpoint_host_offload.py b/tests/unit/test_checkpoint_host_offload.py new file mode 100644 index 000000000..ebf56aed0 --- /dev/null +++ b/tests/unit/test_checkpoint_host_offload.py @@ -0,0 +1,112 @@ +from dataclasses import replace + +import pytest +from test_trainer_rank_checkpoint_memory import dense_rank +import torch +from torch.utils.checkpoint import checkpoint + +from art.megatron.context_parallel.checkpoint import ( + current_host_activation_offload, + offload_saved_activations, +) +from art.megatron.lora import LoRASlotRef +from art.trainer_rank import ForwardInput + + +def test_checkpoint_host_roundtrip_preserves_forward_and_all_gradients_exactly(): + x0 = torch.arange(24, dtype=torch.float64).reshape(6, 1, 4) / 24 + w0 = torch.arange(16, dtype=torch.float64).reshape(4, 4) / 16 + + def execute(offload): + x = x0.clone().requires_grad_() + weights = [w0.clone().requires_grad_(), w0.T.contiguous().requires_grad_()] + + def forward(): + y = checkpoint( + lambda value: (value @ weights[0]).sin(), x, use_reentrant=True + ) + return checkpoint(lambda value: value @ weights[1], y, use_reentrant=True) + + if offload: + with offload_saved_activations(hidden_size=4, verify=True) as stats: + output = forward() + # Unpacking remains valid after the forward-only scope ends. + output.square().sum().backward() + assert stats.boundaries == stats.unpacked == stats.verified == 2 + assert stats.saved_bytes == 2 * x.numel() * x.element_size() + else: + output = forward() + output.square().sum().backward() + return [output.detach(), x.grad, *(weight.grad for weight in weights)] + + reference, actual = execute(False), execute(True) + assert all(torch.equal(a, b) for a, b in zip(reference, actual, strict=True)) + assert current_host_activation_offload() is None + + +def test_host_budget_is_priced_from_worst_rank_rows_before_forward(): + rank = dense_rank() + rank._topology_key = lambda: (1, 1, 4, 1) + rank._parallel_shape = replace(rank._parallel_shape, cp=4) + groups = ((128, True),) + retained = rank._checkpoint_memory_floor(groups)[0] + with offload_saved_activations( + hidden_size=rank._hidden_size, max_cpu_bytes=retained - 1 + ): + with pytest.raises(ValueError, match="before forward"): + rank._checkpoint_memory_floor(groups) + + +def test_host_budget_refuses_before_unbounded_storage_and_cleans_scope(): + x = torch.ones((8, 1, 4), requires_grad=True) + size = x.numel() * x.element_size() + with pytest.raises(ValueError, match="byte budget exceeded"): + with offload_saved_activations(hidden_size=4, max_cpu_bytes=size) as stats: + checkpoint(lambda value: value.sin(), x, use_reentrant=True) + checkpoint(lambda value: value.cos(), x, use_reentrant=True) + assert stats.boundaries == 1 and stats.saved_bytes == size + assert current_host_activation_offload() is None + with offload_saved_activations(hidden_size=4): + with pytest.raises(RuntimeError, match="cannot overlap"): + with offload_saved_activations(hidden_size=4): + pass + assert current_host_activation_offload() is None + + +def test_host_offload_keeps_device_workspace_and_every_pending_adapter_gradient(): + rank = dense_rank() + rank._topology_key = lambda: (1, 1, 4, 1) + rank._parallel_shape = replace(rank._parallel_shape, cp=4) + groups = ((128, True),) + baseline = rank._checkpoint_memory_floor(groups) + request = ForwardInput(input_tokens=torch.arange(128), hidden_states=True) + original = rank._memory_signature_from_requests( + [request], slot_group_count=1, grad_modes=(True,) + ) + slot = LoRASlotRef("checkpoint", "host") + pending = (100,) * (rank._num_layers + 1) + rank._pending_adapter_gradient_bytes = lambda refs: pending + with offload_saved_activations(hidden_size=rank._hidden_size): + floor = rank._checkpoint_memory_floor(groups) + signature = rank._memory_signature_from_requests( + [request], slot_group_count=1, grad_modes=(True,) + ) + assert signature != original + assert signature.checkpoint_host_offload == (rank._hidden_size, 64 * 2**30) + assert floor[0] == 128 * rank._hidden_size * 2 + assert floor[0] * rank._num_layers == baseline[0] + assert floor[1] == baseline[1] + assert rank._checkpoint_input_gradient_bytes(groups, floor[0]) == floor[0] + gradient_groups = rank._checkpoint_gradient_groups(groups, (slot,)) + assert rank._checkpoint_adapter_gradient_bytes(gradient_groups) == sum(pending) + assert current_host_activation_offload() is None + + +def test_host_offload_rejects_unsupported_runtime_before_admission(): + rank = dense_rank() + request = ForwardInput(input_tokens=torch.arange(128), hidden_states=True) + with offload_saved_activations(hidden_size=rank._hidden_size): + with pytest.raises(ValueError, match="one native dense CP uniform root"): + rank._memory_signature_from_requests( + [request], slot_group_count=1, grad_modes=(True,) + ) From 69d03c7d68b9925327f8d6eda3498ffe30bee6f3 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Wed, 21 Oct 2026 20:51:38 +0000 Subject: [PATCH 06/13] fix: isolate host checkpoint storage in split memory signatures (cherry picked from commit 0b4b671388d4210acec782b5e63d61888ce2a387) Co-Authored-By: Claude Opus 5.5 --- .../megatron/context_parallel/checkpoint.py | 9 ++++----- src/art/trainer_rank/_memory.py | 8 +++++++- .../test_trainer_rank_checkpoint_memory.py | 20 +++++++++++++++++++ 3 files changed, 31 insertions(+), 6 deletions(-) diff --git a/src/art/megatron/context_parallel/checkpoint.py b/src/art/megatron/context_parallel/checkpoint.py index 79bf7ef65..d8bac589d 100644 --- a/src/art/megatron/context_parallel/checkpoint.py +++ b/src/art/megatron/context_parallel/checkpoint.py @@ -1,9 +1,8 @@ -"""Shard saved root activations while keeping CP1 operator geometry. +"""Bound saved checkpoint activations through rank sharding or host storage. -Every group member executes the same full-root forward and backward. Only -checkpoint storage is divided; no partial parameter gradients are produced. -Use an outer group of CP1 models, disable automatic DP gradient reduction, -and export one complete parameter-gradient copy without reduction. +Rank sharding keeps CP1 operator geometry: every group member executes the +same full root and exports one complete gradient without reduction. Host +storage keeps the original native CP operators and gradient contract. """ from __future__ import annotations diff --git a/src/art/trainer_rank/_memory.py b/src/art/trainer_rank/_memory.py index c24611e12..822a71d32 100644 --- a/src/art/trainer_rank/_memory.py +++ b/src/art/trainer_rank/_memory.py @@ -221,7 +221,13 @@ def feed(value: Any) -> None: ) ) if isinstance(signature, _ScopedMemorySignature): - feed((signature.checkpoint_partition, signature.bounded_forward_peak)) + feed( + ( + signature.checkpoint_partition, + signature.bounded_forward_peak, + signature.checkpoint_host_offload, + ) + ) for g in p.groups: slot = ( None diff --git a/tests/unit/test_trainer_rank_checkpoint_memory.py b/tests/unit/test_trainer_rank_checkpoint_memory.py index bfafe3501..13aded7bb 100644 --- a/tests/unit/test_trainer_rank_checkpoint_memory.py +++ b/tests/unit/test_trainer_rank_checkpoint_memory.py @@ -840,6 +840,26 @@ def test_split_memory_key_separates_checkpoint_partitions(monkeypatch): assert r._split_memory_key(b) != r._split_memory_key(c) +def test_split_memory_key_separates_host_activation_storage(): + from art.trainer_rank._impl import _SplitForwardPlan + from art.trainer_rank._memory import _ScopedMemorySignature + + r = rank() + plain = r._plan_flat_forward( + [ForwardInput(input_tokens=torch.arange(17), hidden_states=True)] + ) + signature = _ScopedMemorySignature( + **(vars(plain.signature) | {"topology": (1, 1, 4, 1)}), + bounded_forward_peak=True, + ) + plan = _SplitForwardPlan((replace(plain, signature=signature),), ((0,),), 1) + host = replace(signature, checkpoint_host_offload=(2048, 64 * 2**30)) + offloaded = replace(plan, subforwards=(replace(plain, signature=host),)) + smaller = replace(host, checkpoint_host_offload=(2048, 32 * 2**30)) + constrained = replace(plan, subforwards=(replace(plain, signature=smaller),)) + assert len({r._split_memory_key(p) for p in (plan, offloaded, constrained)}) == 3 + + def test_megatron_checkpoint_shards_one_boundary_per_layer(monkeypatch): from contextlib import nullcontext From a9355c91ede722dd86662a8582365a1a1783a3e8 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Wed, 21 Oct 2026 21:52:31 +0000 Subject: [PATCH 07/13] test: pin CP1 admission predictions beside bounded CP Record origin/main's CP1 cold and calibrated requirements, checkpoint floors and input-gradient terms for dense and MoE roots, with and without the bounded direct-forward scope, so the CP2/4/8 calibration cannot move them. Co-Authored-By: Claude Opus 5.5 --- .../unit/test_trainer_rank_forward_profile.py | 40 +++++++++++++++++++ 1 file changed, 40 insertions(+) diff --git a/tests/unit/test_trainer_rank_forward_profile.py b/tests/unit/test_trainer_rank_forward_profile.py index e73717be7..dc528a2a2 100644 --- a/tests/unit/test_trainer_rank_forward_profile.py +++ b/tests/unit/test_trainer_rank_forward_profile.py @@ -1,5 +1,6 @@ """Cold fixed scratch must not poison known single-root hidden-only inference.""" +from contextlib import ExitStack from dataclasses import fields, replace import pytest @@ -524,3 +525,42 @@ def test_native_moe_cp_bounded_profiles_keep_static_floors(bounded_rank, cp, gra 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 From 25e9290f20e9f647579903f6e6ce915a016dce8f Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Wed, 21 Oct 2026 21:52:43 +0000 Subject: [PATCH 08/13] test: run host checkpoint offload tests in CI's Megatron step The new offload tests import art.megatron.lora and the Megatron-backed rank fixtures, so collect them with the Megatron runtime instead of the backend unit environment, like the other checkpoint memory tests. Co-Authored-By: Claude Opus 5.5 --- .github/workflows/prek.yml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.github/workflows/prek.yml b/.github/workflows/prek.yml index b270b0645..e17bcf375 100644 --- a/.github/workflows/prek.yml +++ b/.github/workflows/prek.yml @@ -232,6 +232,7 @@ jobs: tests/unit/test_trainer_rank_weird_shapes.py \ tests/unit/test_trainer_rank_admission_inputs.py \ tests/unit/test_trainer_rank_checkpoint_memory.py \ + tests/unit/test_checkpoint_host_offload.py \ tests/unit/test_trainer_rank_profile_warm.py \ tests/unit/test_trainer_rank_tp_floor.py \ tests/unit/test_trainer_rank_checkpoint_gradient_memory.py \ @@ -293,6 +294,7 @@ jobs: --ignore=tests/unit/test_trainer_rank_weird_shapes.py \ --ignore=tests/unit/test_trainer_rank_admission_inputs.py \ --ignore=tests/unit/test_trainer_rank_checkpoint_memory.py \ + --ignore=tests/unit/test_checkpoint_host_offload.py \ --ignore=tests/unit/test_trainer_rank_profile_warm.py \ --ignore=tests/unit/test_trainer_rank_tp_floor.py \ --ignore=tests/unit/test_trainer_rank_checkpoint_gradient_memory.py \ From fdb9663f97c341e81d01e316719b0cd0538de98e Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Wed, 21 Oct 2026 22:14:12 +0000 Subject: [PATCH 09/13] revert: remove host offload split memory cache keys Revert 69d03c7d (upstream 0b4b6713) before removing its host offload dependency. The feature needs collective budget semantics in a separate follow-up. Co-Authored-By: Claude Opus 5.5 --- .../megatron/context_parallel/checkpoint.py | 9 +++++---- src/art/trainer_rank/_memory.py | 8 +------- .../test_trainer_rank_checkpoint_memory.py | 20 ------------------- 3 files changed, 6 insertions(+), 31 deletions(-) diff --git a/src/art/megatron/context_parallel/checkpoint.py b/src/art/megatron/context_parallel/checkpoint.py index d8bac589d..79bf7ef65 100644 --- a/src/art/megatron/context_parallel/checkpoint.py +++ b/src/art/megatron/context_parallel/checkpoint.py @@ -1,8 +1,9 @@ -"""Bound saved checkpoint activations through rank sharding or host storage. +"""Shard saved root activations while keeping CP1 operator geometry. -Rank sharding keeps CP1 operator geometry: every group member executes the -same full root and exports one complete gradient without reduction. Host -storage keeps the original native CP operators and gradient contract. +Every group member executes the same full-root forward and backward. Only +checkpoint storage is divided; no partial parameter gradients are produced. +Use an outer group of CP1 models, disable automatic DP gradient reduction, +and export one complete parameter-gradient copy without reduction. """ from __future__ import annotations diff --git a/src/art/trainer_rank/_memory.py b/src/art/trainer_rank/_memory.py index 822a71d32..c24611e12 100644 --- a/src/art/trainer_rank/_memory.py +++ b/src/art/trainer_rank/_memory.py @@ -221,13 +221,7 @@ def feed(value: Any) -> None: ) ) if isinstance(signature, _ScopedMemorySignature): - feed( - ( - signature.checkpoint_partition, - signature.bounded_forward_peak, - signature.checkpoint_host_offload, - ) - ) + feed((signature.checkpoint_partition, signature.bounded_forward_peak)) for g in p.groups: slot = ( None diff --git a/tests/unit/test_trainer_rank_checkpoint_memory.py b/tests/unit/test_trainer_rank_checkpoint_memory.py index 13aded7bb..bfafe3501 100644 --- a/tests/unit/test_trainer_rank_checkpoint_memory.py +++ b/tests/unit/test_trainer_rank_checkpoint_memory.py @@ -840,26 +840,6 @@ def test_split_memory_key_separates_checkpoint_partitions(monkeypatch): assert r._split_memory_key(b) != r._split_memory_key(c) -def test_split_memory_key_separates_host_activation_storage(): - from art.trainer_rank._impl import _SplitForwardPlan - from art.trainer_rank._memory import _ScopedMemorySignature - - r = rank() - plain = r._plan_flat_forward( - [ForwardInput(input_tokens=torch.arange(17), hidden_states=True)] - ) - signature = _ScopedMemorySignature( - **(vars(plain.signature) | {"topology": (1, 1, 4, 1)}), - bounded_forward_peak=True, - ) - plan = _SplitForwardPlan((replace(plain, signature=signature),), ((0,),), 1) - host = replace(signature, checkpoint_host_offload=(2048, 64 * 2**30)) - offloaded = replace(plan, subforwards=(replace(plain, signature=host),)) - smaller = replace(host, checkpoint_host_offload=(2048, 32 * 2**30)) - constrained = replace(plan, subforwards=(replace(plain, signature=smaller),)) - assert len({r._split_memory_key(p) for p in (plan, offloaded, constrained)}) == 3 - - def test_megatron_checkpoint_shards_one_boundary_per_layer(monkeypatch): from contextlib import nullcontext From 3f7c1458f90966938bc38a793b5b05400c247eab Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Wed, 21 Oct 2026 22:15:19 +0000 Subject: [PATCH 10/13] revert: defer host checkpoint offload until collective budgeting Revert 82284f40 (upstream 67170e60). Cumulative pack-time host usage can fail after preflight on only one rank during repeated forwards. Remove its CI-only test wiring and document the CPU/Gloo repro and required follow-up. CPU checkpoint and gradient tests: 92 passed. Co-Authored-By: Claude Opus 5.5 --- .github/workflows/prek.yml | 2 - HOST-OFFLOAD-FOLLOWUP.md | 37 ++++++ .../megatron/context_parallel/checkpoint.py | 94 +-------------- src/art/trainer_rank/_memory.py | 54 +-------- tests/unit/test_checkpoint_host_offload.py | 112 ------------------ 5 files changed, 43 insertions(+), 256 deletions(-) create mode 100644 HOST-OFFLOAD-FOLLOWUP.md delete mode 100644 tests/unit/test_checkpoint_host_offload.py diff --git a/.github/workflows/prek.yml b/.github/workflows/prek.yml index e17bcf375..b270b0645 100644 --- a/.github/workflows/prek.yml +++ b/.github/workflows/prek.yml @@ -232,7 +232,6 @@ jobs: tests/unit/test_trainer_rank_weird_shapes.py \ tests/unit/test_trainer_rank_admission_inputs.py \ tests/unit/test_trainer_rank_checkpoint_memory.py \ - tests/unit/test_checkpoint_host_offload.py \ tests/unit/test_trainer_rank_profile_warm.py \ tests/unit/test_trainer_rank_tp_floor.py \ tests/unit/test_trainer_rank_checkpoint_gradient_memory.py \ @@ -294,7 +293,6 @@ jobs: --ignore=tests/unit/test_trainer_rank_weird_shapes.py \ --ignore=tests/unit/test_trainer_rank_admission_inputs.py \ --ignore=tests/unit/test_trainer_rank_checkpoint_memory.py \ - --ignore=tests/unit/test_checkpoint_host_offload.py \ --ignore=tests/unit/test_trainer_rank_profile_warm.py \ --ignore=tests/unit/test_trainer_rank_tp_floor.py \ --ignore=tests/unit/test_trainer_rank_checkpoint_gradient_memory.py \ diff --git a/HOST-OFFLOAD-FOLLOWUP.md b/HOST-OFFLOAD-FOLLOWUP.md new file mode 100644 index 000000000..c7956dc77 --- /dev/null +++ b/HOST-OFFLOAD-FOLLOWUP.md @@ -0,0 +1,37 @@ +# Host checkpoint offload follow-up + +The default-off host activation scope from upstream `67170e60` (cherry-pick +`82284f40`) and its split-key fix `0b4b6713` (`69d03c7d`) are reverted here. +There are no production callers. Host offload needs its own implementation and +qualification before returning; this branch retains device checkpoint storage. + +The pre-forward check compares one root's boundary bytes with the total host +budget. The pack hook instead checks cumulative `stats.saved_bytes`, which is +never released. A second forward can therefore pass preflight, exhaust the +budget inside a layer on one rank, and leave a peer entering the next CP +collective. Exchanging planning exceptions cannot catch this later divergence. + +Sol's CPU/Gloo repro is preserved at +`/tmp/claude-1001/art-cp-admission-review/scratch/probe_host_budget.py`. +It uses two checkpoint layers, 8 rows on rank 0 and 4 on rank 1, a budget of +163,840 bytes, two forwards retaining both graphs in one scope, and an in-layer +all-reduce. Both second-forward preflights pass; rank 0 raises +`checkpoint host activation byte budget exceeded`, while rank 1 enters the +collective and sees peer teardown (or hangs without teardown). The frozen +`scratch/head` source still reproduces this. The current source intentionally +has no `offload_saved_activations` API, so the original probe fails at import. + +A replacement must collectively reserve cumulative live host bytes before any +rank executes a root, using worst-rank demand and available budget, and release +reservations only when all tensors/graphs holding that root have been released. +Alternatively, enforce a single-forward scope collectively or use a per-forward +budget with explicit lifetime release; resetting a counter while an earlier +graph remains live does not bound host memory. Pack-time limits can remain +defensive checks, but must not be the normal admission mechanism. + +Required regressions include repeated forwards with unequal row ownership, +retained graphs versus completed backwards, exact budget edges, exceptions and +scope cleanup, and collective refusal before the first layer. Separate CUDA +qualification must establish saved-tensor lifetime/gradient correctness, actual +device and host peaks, bounded pinned storage, and copy synchronization. CPU +round trips alone do not establish those properties. diff --git a/src/art/megatron/context_parallel/checkpoint.py b/src/art/megatron/context_parallel/checkpoint.py index 79bf7ef65..ac0da17dc 100644 --- a/src/art/megatron/context_parallel/checkpoint.py +++ b/src/art/megatron/context_parallel/checkpoint.py @@ -26,9 +26,6 @@ "art_checkpoint_partition", default=None ) _HOOKS: ContextVar[Any] = ContextVar("art_checkpoint_shard_hooks", default=None) -_HOST_OFFLOAD: ContextVar[HostActivationStats | None] = ContextVar( - "art_checkpoint_host_offload", default=None -) def _install_checkpoint_hooks() -> None: @@ -62,7 +59,7 @@ def local(*inputs: Any) -> Any: def pack(tensor: torch.Tensor) -> Any: nonlocal packed saved = hooks[0](tensor) - if isinstance(saved, (_Shard, _HostBoundary)): + if isinstance(saved, _Shard): packed += 1 if packed > 1: raise ValueError( @@ -71,10 +68,7 @@ def pack(tensor: torch.Tensor) -> Any: return saved with torch.autograd.graph.saved_tensors_hooks(pack, hooks[1]): - result = original(local, *args) - if _HOST_OFFLOAD.get() is not None and packed != 1: - raise ValueError("host offload requires one checkpoint hidden input") - return result + return original(local, *args) setattr(apply, "_art_checkpoint_shards", True) function.apply = staticmethod(apply) @@ -107,88 +101,6 @@ class _Shard: numel: int -@dataclass -class HostActivationStats: - hidden_size: int - max_cpu_bytes: int - saved_bytes: int = 0 - boundaries: int = 0 - unpacked: int = 0 - verified: int = 0 - - -@dataclass(frozen=True) -class _HostBoundary: - tensor: torch.Tensor - device: torch.device - - -def current_host_activation_offload() -> HostActivationStats | None: - return _HOST_OFFLOAD.get() - - -@contextmanager -def offload_saved_activations( - *, hidden_size: int, max_cpu_bytes: int = 64 * 2**30, verify: bool = False -) -> Iterator[HostActivationStats]: - """Copy full/uniform checkpoint inputs to host; retain native CP arithmetic. - - Only the checkpoint boundary's contiguous [rows, 1, hidden] input moves. - The layer's ordinary autograd saves stay on device. Copies preserve dtype; - synchronous transfers make lifetime and completion explicit. Host storage - is bounded per scope, independent of the device admission budgets. - """ - if type(hidden_size) is not int or hidden_size <= 0: - raise ValueError("host offload requires a positive hidden size") - if type(max_cpu_bytes) is not int or max_cpu_bytes <= 0: - raise ValueError("host offload requires a positive CPU byte budget") - if _HOOKS.get() is not None or _HOST_OFFLOAD.get() is not None: - raise RuntimeError("checkpoint storage scopes cannot overlap") - stats = HostActivationStats(hidden_size, max_cpu_bytes) - - def pack(tensor: torch.Tensor) -> torch.Tensor | _HostBoundary: - if tensor.ndim != 3 or tensor.shape[-1] != hidden_size: - return tensor - if ( - tensor.shape[1] != 1 - or not tensor.is_floating_point() - or not tensor.is_contiguous() - ): - raise ValueError("host offload requires contiguous [rows, 1, hidden] input") - size = tensor.numel() * tensor.element_size() - if stats.saved_bytes + size > max_cpu_bytes: - raise ValueError("checkpoint host activation byte budget exceeded") - host = torch.empty_like( - tensor, device="cpu", pin_memory=tensor.device.type == "cuda" - ) - host.copy_(tensor.detach(), non_blocking=False) - if verify: - if not torch.equal( - host.to(tensor.device).view(torch.uint8), - tensor.detach().view(torch.uint8), - ): - raise ValueError("checkpoint host round-trip changed input bytes") - stats.verified += 1 - stats.saved_bytes += size - stats.boundaries += 1 - return _HostBoundary(host, tensor.device) - - def unpack(value: torch.Tensor | _HostBoundary) -> torch.Tensor: - if isinstance(value, torch.Tensor): - return value - stats.unpacked += 1 - return value.tensor.to(value.device, non_blocking=False) - - _install_checkpoint_hooks() - token = _HOST_OFFLOAD.set(stats) - hooks = _HOOKS.set((pack, unpack)) - try: - yield stats - finally: - _HOOKS.reset(hooks) - _HOST_OFFLOAD.reset(token) - - @contextmanager def shard_saved_activations( group: Any, *, shape: tuple[int, int, int], key: str @@ -207,7 +119,7 @@ def shard_saved_activations( tensors, while recomputation's temporary saves should remain local. Keep the numerical-policy scope open across both forward and backward. """ - if _PARTITION.get() is not None or _HOST_OFFLOAD.get() is not None: + if _PARTITION.get() is not None: raise RuntimeError("checkpoint sharding scopes cannot overlap") if len(shape) != 3 or shape[1] != 1 or any(n <= 0 for n in shape): raise ValueError("checkpoint activation shape must be [tokens, 1, hidden]") diff --git a/src/art/trainer_rank/_memory.py b/src/art/trainer_rank/_memory.py index c24611e12..7885f1621 100644 --- a/src/art/trainer_rank/_memory.py +++ b/src/art/trainer_rank/_memory.py @@ -36,7 +36,6 @@ class _ScopedMemorySignature(_impl._MemorySignature): checkpoint_partition: tuple[int, tuple[int, int, int]] | None = None bounded_forward_peak: bool = False - checkpoint_host_offload: tuple[int, int] | None = None @dataclass(frozen=True) @@ -963,25 +962,6 @@ def _checkpoint_floor_from_facts( * self._hidden_size * 2 ) - from art.megatron.context_parallel.checkpoint import current_host_activation_offload - - offload = current_host_activation_offload() - if offload is not None: - if not dense_cp or offload.hidden_size != self._hidden_size: - raise ValueError( - "checkpoint host offload requires matching native dense CP" - ) - if retained > offload.max_cpu_bytes: - raise ValueError( - "checkpoint host activation byte budget exceeded before forward" - ) - # All saved boundaries move to bounded host storage. One full input - # still exists on device while packing or unpacking a checkpoint. - retained = ( - max((rows for rows, grad in group_rows if grad), default=0) - * 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 @@ -1265,15 +1245,11 @@ def _checkpoint_gradient_groups( decoder layer saves the group's rows (this rank's TP shard). The base model (no name) has no adapter slot. """ - from art.megatron.context_parallel.checkpoint import current_host_activation_offload - tp = self._topology_key()[1] - offload = current_host_activation_offload() is not None return tuple( ( ref if ref is not None and ref.name is not None else None, - ((0 if offload else -(-rows // tp) * self._hidden_size * 2),) - * self._num_layers, + (-(-rows // tp) * self._hidden_size * 2,) * self._num_layers, ) for (rows, grad), ref in zip( group_rows, slot_refs or (None,) * len(group_rows), strict=True @@ -1617,30 +1593,9 @@ def _memory_signature_from_requests( grad_modes: Iterable[bool], slot_groups: Iterable[tuple["LoRASlotRef | None", bool]] = (), ) -> _impl._MemorySignature: - from art.megatron.context_parallel.checkpoint import ( - current_checkpoint_partition, - current_host_activation_offload, - ) + from art.megatron.context_parallel.checkpoint import current_checkpoint_partition partition = current_checkpoint_partition() - offload = current_host_activation_offload() - if offload is not None and ( - partition is not None - or len(requests) != 1 - or self._topology_key() not in ((1, 1, 2, 1), (1, 1, 4, 1), (1, 1, 8, 1)) - or self._parallel_shape.cp != self._topology_key()[2] - or self._parallel_shape.ep != 1 - or self._parallel_shape.etp != 1 - or self._moe_layers - or not self._dense_fc1_adapted - or offload.hidden_size != self._hidden_size - or not _checkpoint_layers( - self, ((int(requests[0].input_tokens.numel()), True),) - ) - ): - raise ValueError( - "checkpoint host offload requires one native dense CP uniform root" - ) if partition is not None: # Native MoE storage sharding failed full-model gradient parity on H200. if self._moe_layers: @@ -1692,14 +1647,11 @@ def _memory_signature_from_requests( else self._dense_fc1_adapted and bool(self._geometry.ffn_hidden_size) ) ) - if partition is not None or bounded or offload is not None: + if partition is not None or bounded: return _ScopedMemorySignature( **vars(signature), checkpoint_partition=partition, bounded_forward_peak=bounded, - checkpoint_host_offload=None - if offload is None - else (offload.hidden_size, offload.max_cpu_bytes), ) return signature diff --git a/tests/unit/test_checkpoint_host_offload.py b/tests/unit/test_checkpoint_host_offload.py deleted file mode 100644 index ebf56aed0..000000000 --- a/tests/unit/test_checkpoint_host_offload.py +++ /dev/null @@ -1,112 +0,0 @@ -from dataclasses import replace - -import pytest -from test_trainer_rank_checkpoint_memory import dense_rank -import torch -from torch.utils.checkpoint import checkpoint - -from art.megatron.context_parallel.checkpoint import ( - current_host_activation_offload, - offload_saved_activations, -) -from art.megatron.lora import LoRASlotRef -from art.trainer_rank import ForwardInput - - -def test_checkpoint_host_roundtrip_preserves_forward_and_all_gradients_exactly(): - x0 = torch.arange(24, dtype=torch.float64).reshape(6, 1, 4) / 24 - w0 = torch.arange(16, dtype=torch.float64).reshape(4, 4) / 16 - - def execute(offload): - x = x0.clone().requires_grad_() - weights = [w0.clone().requires_grad_(), w0.T.contiguous().requires_grad_()] - - def forward(): - y = checkpoint( - lambda value: (value @ weights[0]).sin(), x, use_reentrant=True - ) - return checkpoint(lambda value: value @ weights[1], y, use_reentrant=True) - - if offload: - with offload_saved_activations(hidden_size=4, verify=True) as stats: - output = forward() - # Unpacking remains valid after the forward-only scope ends. - output.square().sum().backward() - assert stats.boundaries == stats.unpacked == stats.verified == 2 - assert stats.saved_bytes == 2 * x.numel() * x.element_size() - else: - output = forward() - output.square().sum().backward() - return [output.detach(), x.grad, *(weight.grad for weight in weights)] - - reference, actual = execute(False), execute(True) - assert all(torch.equal(a, b) for a, b in zip(reference, actual, strict=True)) - assert current_host_activation_offload() is None - - -def test_host_budget_is_priced_from_worst_rank_rows_before_forward(): - rank = dense_rank() - rank._topology_key = lambda: (1, 1, 4, 1) - rank._parallel_shape = replace(rank._parallel_shape, cp=4) - groups = ((128, True),) - retained = rank._checkpoint_memory_floor(groups)[0] - with offload_saved_activations( - hidden_size=rank._hidden_size, max_cpu_bytes=retained - 1 - ): - with pytest.raises(ValueError, match="before forward"): - rank._checkpoint_memory_floor(groups) - - -def test_host_budget_refuses_before_unbounded_storage_and_cleans_scope(): - x = torch.ones((8, 1, 4), requires_grad=True) - size = x.numel() * x.element_size() - with pytest.raises(ValueError, match="byte budget exceeded"): - with offload_saved_activations(hidden_size=4, max_cpu_bytes=size) as stats: - checkpoint(lambda value: value.sin(), x, use_reentrant=True) - checkpoint(lambda value: value.cos(), x, use_reentrant=True) - assert stats.boundaries == 1 and stats.saved_bytes == size - assert current_host_activation_offload() is None - with offload_saved_activations(hidden_size=4): - with pytest.raises(RuntimeError, match="cannot overlap"): - with offload_saved_activations(hidden_size=4): - pass - assert current_host_activation_offload() is None - - -def test_host_offload_keeps_device_workspace_and_every_pending_adapter_gradient(): - rank = dense_rank() - rank._topology_key = lambda: (1, 1, 4, 1) - rank._parallel_shape = replace(rank._parallel_shape, cp=4) - groups = ((128, True),) - baseline = rank._checkpoint_memory_floor(groups) - request = ForwardInput(input_tokens=torch.arange(128), hidden_states=True) - original = rank._memory_signature_from_requests( - [request], slot_group_count=1, grad_modes=(True,) - ) - slot = LoRASlotRef("checkpoint", "host") - pending = (100,) * (rank._num_layers + 1) - rank._pending_adapter_gradient_bytes = lambda refs: pending - with offload_saved_activations(hidden_size=rank._hidden_size): - floor = rank._checkpoint_memory_floor(groups) - signature = rank._memory_signature_from_requests( - [request], slot_group_count=1, grad_modes=(True,) - ) - assert signature != original - assert signature.checkpoint_host_offload == (rank._hidden_size, 64 * 2**30) - assert floor[0] == 128 * rank._hidden_size * 2 - assert floor[0] * rank._num_layers == baseline[0] - assert floor[1] == baseline[1] - assert rank._checkpoint_input_gradient_bytes(groups, floor[0]) == floor[0] - gradient_groups = rank._checkpoint_gradient_groups(groups, (slot,)) - assert rank._checkpoint_adapter_gradient_bytes(gradient_groups) == sum(pending) - assert current_host_activation_offload() is None - - -def test_host_offload_rejects_unsupported_runtime_before_admission(): - rank = dense_rank() - request = ForwardInput(input_tokens=torch.arange(128), hidden_states=True) - with offload_saved_activations(hidden_size=rank._hidden_size): - with pytest.raises(ValueError, match="one native dense CP uniform root"): - rank._memory_signature_from_requests( - [request], slot_group_count=1, grad_modes=(True,) - ) From 7d34f8f449dfe708ba1690bacefcd6e02eec3c1e Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Wed, 21 Oct 2026 22:25:45 +0000 Subject: [PATCH 11/13] fix: price dense CP LoRA workspace and restore replay eligibility Record selected adapter layouts for native dense CP, including attention-only TP1 models. Charge the most-adapted recomputed layer at the largest selected rank for saved factor products plus backward workspace. Freeze runtime parallel shape and observed MoE layers for CPU replay, and refuse older CP reports without those facts. CPU validation: 370 passed / 14 skipped across targeted accounting, CP1 pinning, loaded-rank admission and CP2/4/8 replay. All five measured rank-8 lifecycle cells remain within their recorded budgets and above allocated peaks. Ruff/format/lock pass; ty has the same 18 diagnostics as base main, with zero in touched production/replay files. The complete Megatron list needs a clean committed tree for acceptance fixtures; rerun follows. Co-Authored-By: Claude Opus 5.5 --- src/art/trainer_rank/_memory.py | 55 +++++--- src/art/trainer_rank/_micro_batch_planner.py | 2 + src/art/trainer_rank/_planner_misses.py | 25 +++- tests/unit/test_hybridep_planner_replay.py | 37 ++++++ .../test_trainer_rank_checkpoint_memory.py | 124 ++++++++++++++++++ 5 files changed, 219 insertions(+), 24 deletions(-) diff --git a/src/art/trainer_rank/_memory.py b/src/art/trainer_rank/_memory.py index 7885f1621..13cfc1b60 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,13 +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, cp, _ = self._topology_key() - dense_cp = ( - cp in (2, 4, 8) - and self._parallel_shape.cp == cp - and not self._moe_layers - and self._dense_fc1_adapted - ) + 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) @@ -1071,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( ( @@ -1099,7 +1108,8 @@ 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 + return retained, workspace + copies * rows * modules * rank * self._param_dtype_size def _cold_recompute_transient_bytes(self: TrainerRank) -> int: @@ -1120,13 +1130,8 @@ def _checkpoint_input_gradient_bytes( Native dense CP and TP x SP price that workspace explicitly, beside one input-gradient shard of the group being recomputed. """ - _, tp, cp, _ = self._topology_key() - dense_cp = ( - cp in (2, 4, 8) - and self._parallel_shape.cp == cp - and not self._moe_layers - and self._dense_fc1_adapted - ) + 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() @@ -1611,6 +1616,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) ) @@ -1661,17 +1671,20 @@ 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. """ + 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 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..6a832fb1b 100644 --- a/tests/unit/test_trainer_rank_checkpoint_memory.py +++ b/tests/unit/test_trainer_rank_checkpoint_memory.py @@ -522,6 +522,130 @@ def dense_rank(): return dense +@pytest.mark.parametrize("cp", [2, 4, 8]) +@pytest.mark.parametrize("adapter_rank", [8, 768]) +def test_dense_cp_prices_loaded_lora_products_without_opt_in(cp, adapter_rank): + 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 + + 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( + [ + 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]) + ref = load_slot(r, "selected", adapter_rank) + 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 + # No-grad and a selected base model must not acquire adapter workspace. + base = replace(signature, slot_shapes=()) + assert r._sequence_parallel_lora_floor((retained, workspace), groups, base) == ( + retained, + workspace, + ) + 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. From afe1c16313ec8fa44dd7a481e08552b2c3f67bdd Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Wed, 21 Oct 2026 22:27:48 +0000 Subject: [PATCH 12/13] fix: retain legacy CP allowance when adapter ranks are missing Scalar and historical estimators can declare LoRA targets without supplying selected ranks. Keep the old boundary-gradient allowance beside explicit workspace in that case instead of interpreting absent ranks as zero. Complete rank metadata still uses the derived saved-product and backward-workspace charge. Extend the high-rank regression to the literal metadata-omitting probe. CPU checkpoint/forward-profile tests: 150 passed. Co-Authored-By: Claude Opus 5.5 --- src/art/trainer_rank/_memory.py | 5 +++++ .../test_trainer_rank_checkpoint_memory.py | 19 +++++++++++++++---- 2 files changed, 20 insertions(+), 4 deletions(-) diff --git a/src/art/trainer_rank/_memory.py b/src/art/trainer_rank/_memory.py index 13cfc1b60..0650b3ff6 100644 --- a/src/art/trainer_rank/_memory.py +++ b/src/art/trainer_rank/_memory.py @@ -1108,6 +1108,11 @@ 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) + if dense_cp and modules and not rank: + # Historical/scalar callers may declare LoRA targets without recording + # selected ranks. Keep the legacy boundary-gradient allowance beside + # explicit workspace rather than treating missing metadata as rank 0. + return retained, workspace + retained copies = 2 if dense_cp else 1 return retained, workspace + copies * rows * modules * rank * self._param_dtype_size diff --git a/tests/unit/test_trainer_rank_checkpoint_memory.py b/tests/unit/test_trainer_rank_checkpoint_memory.py index 6a832fb1b..9f7cfc768 100644 --- a/tests/unit/test_trainer_rank_checkpoint_memory.py +++ b/tests/unit/test_trainer_rank_checkpoint_memory.py @@ -581,11 +581,22 @@ def test_dense_cp_prices_loaded_lora_products_without_opt_in(cp, adapter_rank): # 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 - # No-grad and a selected base model must not acquire adapter workspace. - base = replace(signature, slot_shapes=()) - assert r._sequence_parallel_lora_floor((retained, workspace), groups, base) == ( + # 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, + 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() From 60b9de6dc9e9e5f5f5da8006c4dcf9ea9a1626dd Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Thu, 22 Oct 2026 00:23:25 +0000 Subject: [PATCH 13/13] fix: retain conservative CP accounting for unknown adapters Read LoRA subclasses, classify native selected-factor inventories, and retain the per-layer allowance for unknown owners or geometry. Cover the GLM rank-512 escape and inherited/foreign adapter regressions. Move the host-offload handoff outside the repository. Co-Authored-By: Claude Opus 5.5 --- HOST-OFFLOAD-FOLLOWUP.md | 37 ---- src/art/trainer_rank/_impl.py | 5 +- src/art/trainer_rank/_memory.py | 109 ++++++++-- .../test_trainer_rank_checkpoint_memory.py | 193 +++++++++++++++++- 4 files changed, 288 insertions(+), 56 deletions(-) delete mode 100644 HOST-OFFLOAD-FOLLOWUP.md diff --git a/HOST-OFFLOAD-FOLLOWUP.md b/HOST-OFFLOAD-FOLLOWUP.md deleted file mode 100644 index c7956dc77..000000000 --- a/HOST-OFFLOAD-FOLLOWUP.md +++ /dev/null @@ -1,37 +0,0 @@ -# Host checkpoint offload follow-up - -The default-off host activation scope from upstream `67170e60` (cherry-pick -`82284f40`) and its split-key fix `0b4b6713` (`69d03c7d`) are reverted here. -There are no production callers. Host offload needs its own implementation and -qualification before returning; this branch retains device checkpoint storage. - -The pre-forward check compares one root's boundary bytes with the total host -budget. The pack hook instead checks cumulative `stats.saved_bytes`, which is -never released. A second forward can therefore pass preflight, exhaust the -budget inside a layer on one rank, and leave a peer entering the next CP -collective. Exchanging planning exceptions cannot catch this later divergence. - -Sol's CPU/Gloo repro is preserved at -`/tmp/claude-1001/art-cp-admission-review/scratch/probe_host_budget.py`. -It uses two checkpoint layers, 8 rows on rank 0 and 4 on rank 1, a budget of -163,840 bytes, two forwards retaining both graphs in one scope, and an in-layer -all-reduce. Both second-forward preflights pass; rank 0 raises -`checkpoint host activation byte budget exceeded`, while rank 1 enters the -collective and sees peer teardown (or hangs without teardown). The frozen -`scratch/head` source still reproduces this. The current source intentionally -has no `offload_saved_activations` API, so the original probe fails at import. - -A replacement must collectively reserve cumulative live host bytes before any -rank executes a root, using worst-rank demand and available budget, and release -reservations only when all tensors/graphs holding that root have been released. -Alternatively, enforce a single-forward scope collectively or use a per-forward -budget with explicit lifetime release; resetting a counter while an earlier -graph remains live does not bound host memory. Pack-time limits can remain -defensive checks, but must not be the normal admission mechanism. - -Required regressions include repeated forwards with unequal row ownership, -retained graphs versus completed backwards, exact budget edges, exceptions and -scope cleanup, and collective refusal before the first layer. Separate CUDA -qualification must establish saved-tensor lifetime/gradient correctness, actual -device and host peaks, bounded pinned storage, and copy synchronization. CPU -round trips alone do not establish those properties. 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 0650b3ff6..16ce4031b 100644 --- a/src/art/trainer_rank/_memory.py +++ b/src/art/trainer_rank/_memory.py @@ -1096,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, ) @@ -1108,13 +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) - if dense_cp and modules and not rank: - # Historical/scalar callers may declare LoRA targets without recording - # selected ranks. Keep the legacy boundary-gradient allowance beside - # explicit workspace rather than treating missing metadata as rank 0. - return retained, workspace + retained copies = 2 if dense_cp else 1 - return retained, workspace + copies * rows * modules * rank * self._param_dtype_size + 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: @@ -1678,7 +1700,8 @@ def _slot_memory_shapes( 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. + 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 ( @@ -1698,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) diff --git a/tests/unit/test_trainer_rank_checkpoint_memory.py b/tests/unit/test_trainer_rank_checkpoint_memory.py index 9f7cfc768..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( @@ -524,12 +524,18 @@ def dense_rank(): @pytest.mark.parametrize("cp", [2, 4, 8]) @pytest.mark.parametrize("adapter_rank", [8, 768]) -def test_dense_cp_prices_loaded_lora_products_without_opt_in(cp, adapter_rank): +@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( @@ -546,13 +552,34 @@ def test_dense_cp_prices_loaded_lora_products_without_opt_in(cp, adapter_rank): # q/k/v/output/gate/up/down: the selected rank differs from the template. loras = torch.nn.ModuleList( [ - LoRA(f"target.{i}", 32, 32, 4, 4, torch.bfloat16, torch.device("cpu")) + (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]) - ref = load_slot(r, "selected", adapter_rank) + 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 @@ -601,6 +628,164 @@ def test_dense_cp_prices_loaded_lora_products_without_opt_in(cp, adapter_rank): 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.