From da9e7c699389c3f3c0a8d2141ffd98bc7a117b94 Mon Sep 17 00:00:00 2001 From: pengdurice Date: Tue, 1 Sep 2026 18:27:56 +0000 Subject: [PATCH 1/2] init fix, wip Signed-off-by: pengdurice --- .../compile/custom_ops/tp_collectives.py | 6 +- .../compile/passes/offload_activation.py | 83 ++++++++++++++++++- .../v1/compile/test_offload_activation.py | 60 ++++++++++++++ 3 files changed, 144 insertions(+), 5 deletions(-) diff --git a/deepspeed/compile/custom_ops/tp_collectives.py b/deepspeed/compile/custom_ops/tp_collectives.py index edcbce2b4a4e..8566ba555721 100644 --- a/deepspeed/compile/custom_ops/tp_collectives.py +++ b/deepspeed/compile/custom_ops/tp_collectives.py @@ -3,6 +3,8 @@ # DeepSpeed Team +from typing import List + import torch import deepspeed.comm as dist from deepspeed.utils import groups @@ -47,7 +49,7 @@ def reduce_from_tp_region_fake(input: torch.Tensor): @torch.library.custom_op("autotp::gather_from_tp_region", mutates_args=()) -def gather_from_tp_region(input: torch.Tensor, partition_sizes: list[int]) -> torch.Tensor: +def gather_from_tp_region(input: torch.Tensor, partition_sizes: List[int]) -> torch.Tensor: """All-gather the last dimension using the frozen shard widths. Inserted after a column-parallel matmul whose layer asks for gather_output, so that @@ -85,7 +87,7 @@ def gather_from_tp_region(input: torch.Tensor, partition_sizes: list[int]) -> to @torch.library.register_fake("autotp::gather_from_tp_region") -def gather_from_tp_region_fake(input: torch.Tensor, partition_sizes: list[int]): +def gather_from_tp_region_fake(input: torch.Tensor, partition_sizes: List[int]): return input.new_empty((*input.shape[:-1], sum(partition_sizes))) diff --git a/deepspeed/compile/passes/offload_activation.py b/deepspeed/compile/passes/offload_activation.py index 6ecbff00fdcd..281e8dfacb8f 100644 --- a/deepspeed/compile/passes/offload_activation.py +++ b/deepspeed/compile/passes/offload_activation.py @@ -279,6 +279,75 @@ def _skip_bytes(node) -> int: return 0 +def _alias_root(node: Node, no_copy_ops, cache: Dict[Node, Node]) -> Node: + """The node that allocated the storage `node` reads. + + An aliasing op returns a tensor that shares the storage of its first tensor input, so following + that input back through the aliasing ops reaches the node that allocated the memory. A node + whose op is not an aliasing one owns its storage and is its own root. Ops that alias an input + other than the first are rare enough that they are not tracked; such a node is treated as + owning its storage, which only costs a copy that frees nothing. + """ + chain = [] + current = node + while current not in cache: + if current.target not in no_copy_ops: + break + base = next((arg for arg in current.all_input_nodes if isinstance(arg.meta.get("val"), torch.Tensor)), None) + if base is None or base in chain: + break + chain.append(current) + current = base + + root = cache.get(current, current) + for aliasing_node in chain: + cache[aliasing_node] = root + cache[current] = root + return root + + +def _storage_keeper_counts(graph: Graph, saved_nodes: List[Node], returned_to_caller, no_copy_ops, + cache: Dict[Node, Node]) -> Dict[Node, int]: + """How many values that outlive the forward pass hold each storage. + + Moving a saved value to the host releases its memory only if nothing else still points at that + storage. Three kinds of value still do: the graph's own inputs, which the caller owns; the + values the graph returns to the caller; and every other saved value. + """ + counts = defaultdict(int) + keepers = [node for node in graph.nodes if node.op == "placeholder"] + keepers.extend(returned_to_caller) + keepers.extend(saved_nodes) + for keeper in keepers: + counts[_alias_root(keeper, no_copy_ops, cache)] += 1 + return counts + + +def _has_reloadable_layout(node: Node) -> bool: + """Whether the host round trip hands the backward pass back the tensor it was compiled for. + + Both the host buffer and the reloaded device tensor are made with `empty_like`, which + reproduces the strides of a tensor only when that tensor is non-overlapping and dense. A view + that is not -- a strided slice, or an `expand`, whose zero strides also make `empty_like` + allocate the materialized size instead of the much smaller storage the view really holds -- + comes back laid out differently from what the backward graph expects. + """ + val = node.meta.get("val") + if not isinstance(val, torch.Tensor): + return False + shape, strides = val.shape, val.stride() + # Symbolic sizes or strides cannot be checked here; the static-size rule rejects them anyway. + if any(not isinstance(dim, int) for dim in (*shape, *strides)): + return False + + expected = 1 + for stride, size in sorted((s, d) for d, s in zip(shape, strides) if d != 1): + if stride != expected: + return False + expected *= size + return True + + def _eligible_activations(graph: Graph, graph_id: int, num_fwd_outputs, param_manager) -> List[Tuple[Node, int]]: """Every saved activation this pass is allowed to move, largest first.""" output_node = get_output_node(graph) @@ -306,6 +375,8 @@ def _eligible_activations(graph: Graph, graph_id: int, num_fwd_outputs, param_ma _skipped.clear() saved_nodes = [node for node in outputs[num_fwd_outputs:] if isinstance(node, Node)] + alias_roots: Dict[Node, Node] = {} + storage_keepers = _storage_keeper_counts(graph, saved_nodes, returned_to_caller, no_copy_ops, alias_roots) candidates = [] seen = set() @@ -324,10 +395,16 @@ def _eligible_activations(graph: Graph, graph_id: int, num_fwd_outputs, param_ma _skipped["param_or_output"] += _skip_bytes(node) continue # A value that only aliases another tensor shares its storage, so copying it out frees - # nothing while the tensor it aliases is still live. + # nothing while another value that outlives the forward pass still points at that storage. + # When no other value does -- AOTAutograd routinely saves the view and not the tensor it + # came from -- this view is the last holder, and moving it releases the whole allocation. if node.target in no_copy_ops: - _skipped["alias"] += _skip_bytes(node) - continue + if storage_keepers[_alias_root(node, no_copy_ops, alias_roots)] > 1: + _skipped["alias"] += _skip_bytes(node) + continue + if not _has_reloadable_layout(node): + _skipped["alias_layout"] += _skip_bytes(node) + continue # Only floating-point values are activations. The rest are bookkeeping the backward pass # needs -- indices, masks, and the random-number state that attention saves. That state is # the reason this test cannot be a device check: it lives on the host, but an op's traced diff --git a/tests/unit/v1/compile/test_offload_activation.py b/tests/unit/v1/compile/test_offload_activation.py index 0431e95a984b..9177e24d0b11 100644 --- a/tests/unit/v1/compile/test_offload_activation.py +++ b/tests/unit/v1/compile/test_offload_activation.py @@ -467,6 +467,66 @@ def test_fwd_skips_values_that_alias_another_tensor(forced_budget): assert [n for n in _node_names(graph) if n.startswith("offload_")] == ["offload_act_1"] +def _make_view_graph(base_is_saved, view_val=None): + """x -> relu(base) -> aten.view(viewed) -> sum(out), shaped the way the partitioner leaves it. + + AOTAutograd saves whichever values the backward pass reads, and that is often the view alone. + The view is then the last value holding the storage `base` allocated, so moving it to the host + releases the whole allocation. + """ + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = _meta_tensor(16) + + base = _add_node(graph, torch.relu, (x, ), "base", LARGE_NUMEL) + viewed = graph.create_node('call_function', torch.ops.aten.view.default, (base, [LARGE_NUMEL]), {}, name="viewed") + viewed.meta["val"] = _meta_tensor(LARGE_NUMEL) if view_val is None else view_val + out = _add_node(graph, torch.sum, (viewed, ), "out", 1) + + graph.output((out, base, viewed) if base_is_saved else (out, viewed)) + return graph + + +def test_eligible_includes_saved_view_when_base_is_not_saved(): + _ensure_dc_ops() + # Nothing but the view points at that storage once the forward pass ends: `base` is not saved, + # not returned, and read by nothing else. Dropping the view here is what left the memory floor + # growing about 5GB per thousand tokens per GPU on a 2xH200 Qwen3.5-9B run. + graph = _make_view_graph(base_is_saved=False) + + eligible = offload_pass._eligible_activations(graph, 0, 1, {}) + + assert [node.name for node, _ in eligible] == ["viewed"] + assert offload_pass._skipped["alias"] == 0 + + +def test_eligible_skips_saved_view_when_base_is_also_saved(): + _ensure_dc_ops() + # `base` is saved too, so it holds the storage on the device whatever the view does. Copying + # the view out would return nothing and cost a copy each way. + graph = _make_view_graph(base_is_saved=True) + + eligible = offload_pass._eligible_activations(graph, 0, 1, {}) + + assert [node.name for node, _ in eligible] == ["base"] + assert offload_pass._skipped["alias"] == LARGE_SIZE + + +def test_eligible_skips_a_saved_view_the_host_copy_would_not_reproduce(): + _ensure_dc_ops() + # An expanded view repeats one row with a stride of zero. The host buffer is built with + # empty_like, which would allocate all four rows -- four times the storage the view actually + # holds -- and hand the backward pass different strides than it was compiled for. + expanded = torch.empty(LARGE_NUMEL, device="meta").expand(4, LARGE_NUMEL) + graph = _make_view_graph(base_is_saved=False, view_val=expanded) + + eligible = offload_pass._eligible_activations(graph, 0, 1, {}) + + assert eligible == [] + assert offload_pass._skipped["alias"] == 0 + assert offload_pass._skipped["alias_layout"] == 4 * LARGE_SIZE + + def test_fwd_skips_values_already_on_the_host(forced_budget): _ensure_dc_ops() graph = _make_fwd_graph() From d6c9cd9718a02bd5c91f3c69324012bb950aae64 Mon Sep 17 00:00:00 2001 From: pengdurice Date: Fri, 4 Sep 2026 22:42:48 +0000 Subject: [PATCH 2/2] fix Signed-off-by: pengdurice --- .../compile/passes/offload_activation.py | 135 ++++++++++++---- .../v1/compile/test_offload_activation.py | 145 +++++++++++++++++- 2 files changed, 240 insertions(+), 40 deletions(-) diff --git a/deepspeed/compile/passes/offload_activation.py b/deepspeed/compile/passes/offload_activation.py index 281e8dfacb8f..96c966c6911f 100644 --- a/deepspeed/compile/passes/offload_activation.py +++ b/deepspeed/compile/passes/offload_activation.py @@ -3,6 +3,7 @@ # DeepSpeed Team +import functools import os import time from collections import OrderedDict, defaultdict @@ -62,6 +63,12 @@ # the two halves of the pass find each other. _offload_plans: Dict[int, "OrderedDict[str, Tuple[int, int]]"] = {} +# Bytes each planned value keeps allocated if it stays on the device, keyed the same way. This is +# not the size of the copy: a saved view holds its whole base allocation alive, so a 4KB row of a +# 4MB tensor costs 4MB of residency and 4KB of copy. The planner spends headroom in residency +# bytes, while the backward pass schedules its copies in copy bytes. +_resident_bytes: Dict[int, Dict[str, int]] = {} + # Value ids identify a host buffer inside the C++ executor. They never repeat, so a buffer holding # a tensor of one shape is never reused for another. _next_value_id = 0 @@ -279,6 +286,27 @@ def _skip_bytes(node) -> int: return 0 +@functools.lru_cache +def _aliasing_ops(): + """Ops whose output shares storage with an input, for the purpose of tracking that storage. + + get_no_copy_ops() reads the aten schemas, and aten._unsafe_view declares a fresh tensor return + even though it hands back a view -- that declaration is the entire point of the op. It shares + storage all the same, and AOTAutograd emits it after nearly every matmul, so this pass has to + know about it. It is added here rather than in get_no_copy_ops() because that set also decides + where the ZeRO-3 passes release parameters, and this pass has no business changing that. + """ + return frozenset(get_no_copy_ops() | {torch.ops.aten._unsafe_view.default}) + + +def _zero3_gathered_param_ops(): + """The op that produces a ZeRO-3 gathered parameter buffer, empty if DeepCompile is not built.""" + try: + return {torch.ops.dc.allgather_param.default} + except (AttributeError, RuntimeError): + return set() + + def _alias_root(node: Node, no_copy_ops, cache: Dict[Node, Node]) -> Node: """The node that allocated the storage `node` reads. @@ -311,11 +339,20 @@ def _storage_keeper_counts(graph: Graph, saved_nodes: List[Node], returned_to_ca """How many values that outlive the forward pass hold each storage. Moving a saved value to the host releases its memory only if nothing else still points at that - storage. Three kinds of value still do: the graph's own inputs, which the caller owns; the - values the graph returns to the caller; and every other saved value. + storage. Four kinds of value still do: + + - the graph's own inputs, which the caller owns; + - get_attr nodes, whose tensor the GraphModule holds for the life of the process (attention + masks, rotary embedding tables, and the other constants inductor bakes in); + - ZeRO-3 gathered parameters, whose buffer belongs to ZeRO's own registry and is released by + release_param, not by this graph. The forward graph reaches one of these through + dc.wait_allgather, which is an aliasing op, so without this the gathered weight looks like + an ordinary activation with nothing else holding it; + - the values the graph returns to the caller, and every other saved value. """ counts = defaultdict(int) - keepers = [node for node in graph.nodes if node.op == "placeholder"] + gather_ops = _zero3_gathered_param_ops() + keepers = [node for node in graph.nodes if node.op in ("placeholder", "get_attr") or node.target in gather_ops] keepers.extend(returned_to_caller) keepers.extend(saved_nodes) for keeper in keepers: @@ -323,33 +360,46 @@ def _storage_keeper_counts(graph: Graph, saved_nodes: List[Node], returned_to_ca return counts -def _has_reloadable_layout(node: Node) -> bool: - """Whether the host round trip hands the backward pass back the tensor it was compiled for. +def _has_non_overlapping_storage(node: Node) -> bool: + """Whether the tensor's elements each occupy their own place in storage. - Both the host buffer and the reloaded device tensor are made with `empty_like`, which - reproduces the strides of a tensor only when that tensor is non-overlapping and dense. A view - that is not -- a strided slice, or an `expand`, whose zero strides also make `empty_like` - allocate the materialized size instead of the much smaller storage the view really holds -- - comes back laid out differently from what the backward graph expects. + Both the host buffer and the reloaded device tensor are made with `empty_like`, which allocates + one element per logical element. A tensor whose elements overlap holds fewer elements of + storage than it has entries -- `expand` is the ordinary case, repeating a row with a stride of + zero -- so the round trip would copy and allocate several times what the value actually keeps + alive. An expanded row of 1000 floats seen as 4x1000 copies 16KB each way to release 4KB. + + Strides that are merely non-contiguous are fine and are deliberately allowed. `empty_like` does + return a contiguous tensor for a strided slice or one piece of a split, so the reload hands the + backward pass different strides than the traced metadata promises. That was measured on torch + 2.6.0+cu124 with torch._inductor.config.size_asserts off, as init_z3.py sets it: an opaque op + whose meta claims stride (576, 8, 72, 1) while returning (192, 8, 24, 1), consumed by matmul, + batched matmul, linear and reductions under torch.compile, matched eager exactly in every case. """ val = node.meta.get("val") if not isinstance(val, torch.Tensor): return False - shape, strides = val.shape, val.stride() + try: + shape, strides = val.shape, val.stride() + except (RuntimeError, NotImplementedError): + # Sparse, nested and other non-strided layouts have no strides to compare. + return False # Symbolic sizes or strides cannot be checked here; the static-size rule rejects them anyway. if any(not isinstance(dim, int) for dim in (*shape, *strides)): return False - expected = 1 - for stride, size in sorted((s, d) for d, s in zip(shape, strides) if d != 1): - if stride != expected: - return False - expected *= size - return True + # How many elements of storage the tensor spans, against how many entries it has. + span = 1 + sum((size - 1) * abs(stride) for size, stride in zip(shape, strides)) + return val.numel() <= span + +def _eligible_activations(graph: Graph, graph_id: int, num_fwd_outputs, param_manager) -> List[Tuple[Node, int, int]]: + """Every saved activation this pass is allowed to move, largest first. -def _eligible_activations(graph: Graph, graph_id: int, num_fwd_outputs, param_manager) -> List[Tuple[Node, int]]: - """Every saved activation this pass is allowed to move, largest first.""" + Each entry is (node, bytes copied, bytes kept allocated if the value stays resident). The two + sizes differ for a view: the copy carries the view's own bytes while residency holds the whole + allocation the view points into. + """ output_node = get_output_node(graph) outputs = output_node.args[0] if not isinstance(outputs, (list, tuple)): @@ -365,7 +415,7 @@ def _eligible_activations(graph: Graph, graph_id: int, num_fwd_outputs, param_ma returned_to_caller = set(node for node in outputs[:num_fwd_outputs] if isinstance(node, Node)) param_names = set(param_manager[graph_id].param_names) if graph_id in param_manager else set() - no_copy_ops = get_no_copy_ops() + no_copy_ops = _aliasing_ops() # No profile means no peak to plan against, and the usual reason it is missing is that profiling # itself ran out of memory -- which is evidence of exactly the pressure this pass relieves. Take @@ -398,13 +448,18 @@ def _eligible_activations(graph: Graph, graph_id: int, num_fwd_outputs, param_ma # nothing while another value that outlives the forward pass still points at that storage. # When no other value does -- AOTAutograd routinely saves the view and not the tensor it # came from -- this view is the last holder, and moving it releases the whole allocation. - if node.target in no_copy_ops: - if storage_keepers[_alias_root(node, no_copy_ops, alias_roots)] > 1: - _skipped["alias"] += _skip_bytes(node) - continue - if not _has_reloadable_layout(node): - _skipped["alias_layout"] += _skip_bytes(node) - continue + root = _alias_root(node, no_copy_ops, alias_roots) + # root is node for an aliasing op only when the walk could not find the tensor it aliases, + # and an unknown base is not a base this pass may assume is dead. + if node.target in no_copy_ops and (root is node or storage_keepers[root] > 1): + _skipped["alias"] += _skip_bytes(node) + continue + # Checked for every candidate, not only the ones the rule above let through: a piece of a + # split reaches here through operator.getitem, which is not an aliasing op, so the alias + # rule never sees it even though the tensor is a view. + if not _has_non_overlapping_storage(node): + _skipped["overlapping"] += _skip_bytes(node) + continue # Only floating-point values are activations. The rest are bookkeeping the backward pass # needs -- indices, masks, and the random-number state that attention saves. That state is # the reason this test cannot be a device check: it lives on the host, but an op's traced @@ -416,7 +471,11 @@ def _eligible_activations(graph: Graph, graph_id: int, num_fwd_outputs, param_ma if size is None or size < min_size: _skipped["too_small" if size is not None else "no_static_size"] += _skip_bytes(node) continue - candidates.append((node, size)) + # Keeping this value resident holds its whole allocation, which for a view is the base's. + # The alias rule above guarantees this view is the only saved value pointing there, so no + # two entries ever charge the planner for the same bytes. + resident = _static_tensor_size(root) if root is not node else size + candidates.append((node, size, resident if resident is not None else size)) if _skipped: breakdown = " ".join(f"{k}={v}" for k, v in sorted(_skipped.items())) @@ -477,6 +536,7 @@ def _offload_everything_fwd(gm: GraphModule, graph_id: int, profiling_results, p graph = gm.graph # A later compile phase plans again from the original graph, so drop any earlier plan first. _offload_plans[graph_id] = OrderedDict() + _resident_bytes[graph_id] = {} _report_partitioner_split(graph, graph_id, profiling_results[graph_id].num_fwd_outputs) @@ -485,7 +545,7 @@ def _offload_everything_fwd(gm: GraphModule, graph_id: int, profiling_results, p return None output_node = get_output_node(graph) - for node, size in selected: + for node, size, resident in selected: value_id = _new_value_id() # The graph is re-read for every tensor because each insertion changes it. insert_before = _insertion_point_after_last_use(list(graph.nodes), node) @@ -508,11 +568,13 @@ def _offload_everything_fwd(gm: GraphModule, graph_id: int, profiling_results, p output_node.replace_input_with(node, wait_node) _offload_plans[graph_id][node.name] = (value_id, size) + _resident_bytes.setdefault(graph_id, {})[node.name] = resident _stats["offload_nodes"] += 1 graph.lint() print_rank_0(f"offload_activation graph_id={graph_id} floor: moved all {len(selected)} eligible " - f"activations ({sum(size for _, size in selected) / 1e9:.1f}GB) before profiling") + f"activations ({sum(size for _, size, _ in selected) / 1e9:.1f}GB copied, " + f"{sum(resident for _, _, resident in selected) / 1e9:.1f}GB released) before profiling") # Returned, not None: the caller profiles what it gets back, and that profile is the floor the # planner needs. return gm @@ -622,16 +684,21 @@ def _plan_against_floor_fwd(gm: GraphModule, graph_id: int, profiling_results) - print_rank_0(f"offload_activation graph_id={graph_id} {margin_note} " f"floor_peak={floor_peak} budget={budget} headroom={headroom}") - # Largest first: each one returned buys back the most memory per copy avoided. + # Largest first: each one returned buys back the most memory per copy avoided. What it costs + # is residency, which for a saved view is the whole allocation the view points into, not the + # view's own bytes. Charging the copy size here would let a handful of small views of large + # tensors retain many times the headroom the planner thinks it spent. + resident_bytes = _resident_bytes.get(graph_id, {}) by_size = sorted(plan.items(), key=lambda item: item[1][1], reverse=True) kept_resident = 0 for name, (_, size) in by_size: - if size > headroom: + cost = resident_bytes.get(name, size) + if cost > headroom: continue _bring_back(gm.graph, name) del plan[name] - headroom -= size - kept_resident += size + headroom -= cost + kept_resident += cost _stats["offload_nodes"] -= 1 moved_bytes = sum(size for _, size in plan.values()) diff --git a/tests/unit/v1/compile/test_offload_activation.py b/tests/unit/v1/compile/test_offload_activation.py index 9177e24d0b11..f24d969a34b5 100644 --- a/tests/unit/v1/compile/test_offload_activation.py +++ b/tests/unit/v1/compile/test_offload_activation.py @@ -32,6 +32,7 @@ def _reset_offload_pass_globals(): # The plan lives in module globals; reset it so tests pass in any order. yield offload_pass._offload_plans.clear() + offload_pass._resident_bytes.clear() offload_pass.reset_offload_activation_stats() offload_pass._h2d_bytes_per_sec = None @@ -496,7 +497,7 @@ def test_eligible_includes_saved_view_when_base_is_not_saved(): eligible = offload_pass._eligible_activations(graph, 0, 1, {}) - assert [node.name for node, _ in eligible] == ["viewed"] + assert [node.name for node, _, _ in eligible] == ["viewed"] assert offload_pass._skipped["alias"] == 0 @@ -508,15 +509,15 @@ def test_eligible_skips_saved_view_when_base_is_also_saved(): eligible = offload_pass._eligible_activations(graph, 0, 1, {}) - assert [node.name for node, _ in eligible] == ["base"] + assert [node.name for node, _, _ in eligible] == ["base"] assert offload_pass._skipped["alias"] == LARGE_SIZE -def test_eligible_skips_a_saved_view_the_host_copy_would_not_reproduce(): +def test_eligible_skips_a_view_whose_elements_overlap(): _ensure_dc_ops() # An expanded view repeats one row with a stride of zero. The host buffer is built with - # empty_like, which would allocate all four rows -- four times the storage the view actually - # holds -- and hand the backward pass different strides than it was compiled for. + # empty_like, which allocates one element per entry -- all four rows, four times the storage + # the view actually holds -- so the round trip would copy four times what it releases. expanded = torch.empty(LARGE_NUMEL, device="meta").expand(4, LARGE_NUMEL) graph = _make_view_graph(base_is_saved=False, view_val=expanded) @@ -524,7 +525,139 @@ def test_eligible_skips_a_saved_view_the_host_copy_would_not_reproduce(): assert eligible == [] assert offload_pass._skipped["alias"] == 0 - assert offload_pass._skipped["alias_layout"] == 4 * LARGE_SIZE + assert offload_pass._skipped["overlapping"] == 4 * LARGE_SIZE + + +def test_overlap_check_matches_what_empty_like_would_allocate(): + """The predicate has one job: reject a tensor empty_like would allocate more storage for. + + csrc/compile/z3.cpp builds both the pinned host buffer and the reloaded device tensor with + at::empty_like, which allocates one element per entry. A tensor whose entries overlap in + storage therefore costs more to move than it releases. + """ + _ensure_dc_ops() + base = torch.empty(64, 64) + cases = { + "contiguous": base, + "transposed": base.t(), + "channels_last": torch.empty(2, 3, 4, 4).to(memory_format=torch.channels_last), + "dense prefix": base.view(-1)[:1024], + "strided slice": base[:, ::2], + "single row": base[0:1, :], + "one piece of a split": torch.empty(2, 8, 24).chunk(3, dim=-1)[0], + "expanded": torch.empty(64).expand(64, 64), + "broadcast row": torch.empty(1000).expand(4, 1000), + } + for name, tensor in cases.items(): + node = torch.fx.Graph().placeholder("x") + node.meta["val"] = tensor + storage_elements = tensor.untyped_storage().nbytes() // tensor.element_size() + costs_more_than_it_holds = torch.empty_like(tensor).numel() > storage_elements + assert offload_pass._has_non_overlapping_storage(node) == (not costs_more_than_it_holds), name + + +def test_eligible_skips_a_zero3_gathered_parameter(): + """A gathered weight reaches the forward graph through an aliasing op, and must not be moved. + + zero3_compile.add_gather_and_release fuses a narrowing dtype cast into the all-gather and + rewires the cast's users -- the output node included -- to the wait node. The saved value is + then dc.wait_allgather itself, which is in the no-copy set, so it is an alias whose root is the + all-gather. ZeRO owns that buffer and release_param returns it; copying it out frees nothing + and costs a pinned host buffer plus a reload allocation on every step. + """ + _ensure_dc_ops() + graph = torch.fx.Graph() + param = graph.placeholder("primals_1") + param.meta["val"] = _meta_tensor(LARGE_NUMEL) + + gathered = graph.create_node('call_function', + torch.ops.dc.allgather_param.default, (param, 0, 7), {}, + name="allgather_ds_param_primals_1_7") + gathered.meta["val"] = _meta_tensor(LARGE_NUMEL) + wait = graph.create_node('call_function', + torch.ops.dc.wait_allgather.default, (gathered, 0, 7), {}, + name="wait_allgather_ds_param__primals_1_7") + wait.meta["val"] = _meta_tensor(LARGE_NUMEL) + out = _add_node(graph, torch.sum, (wait, ), "out", 1) + graph.output((out, wait)) + + eligible = offload_pass._eligible_activations(graph, 0, 1, {}) + + assert eligible == [] + assert offload_pass._skipped["alias"] == LARGE_SIZE + + +def test_eligible_skips_a_view_of_a_module_constant(): + """A get_attr reads a tensor the GraphModule owns for the life of the process. + + Attention masks and rotary embedding tables arrive this way. The storage stays allocated + whatever this pass does, so a view of one is not the last holder of anything. + """ + _ensure_dc_ops() + graph = torch.fx.Graph() + constant = graph.get_attr("_tensor_constant0") + constant.meta["val"] = _meta_tensor(LARGE_NUMEL) + viewed = graph.create_node('call_function', + torch.ops.aten.view.default, (constant, [LARGE_NUMEL]), {}, + name="viewed") + viewed.meta["val"] = _meta_tensor(LARGE_NUMEL) + out = _add_node(graph, torch.sum, (viewed, ), "out", 1) + graph.output((out, viewed)) + + eligible = offload_pass._eligible_activations(graph, 0, 1, {}) + + assert eligible == [] + assert offload_pass._skipped["alias"] == LARGE_SIZE + + +def test_eligible_follows_unsafe_view_to_the_tensor_it_aliases(): + """aten._unsafe_view declares a fresh tensor return but hands back a view. + + AOTAutograd emits it after nearly every matmul. Reading only the schema, as get_no_copy_ops + does, would make it look like a tensor of its own, and the pass would copy it out while the + tensor it points into is still saved and resident. + """ + _ensure_dc_ops() + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = _meta_tensor(16) + base = _add_node(graph, torch.relu, (x, ), "base", LARGE_NUMEL) + unsafe = graph.create_node('call_function', + torch.ops.aten._unsafe_view.default, (base, [LARGE_NUMEL]), {}, + name="unsafe") + unsafe.meta["val"] = _meta_tensor(LARGE_NUMEL) + out = _add_node(graph, torch.sum, (unsafe, ), "out", 1) + graph.output((out, base, unsafe)) + + eligible = offload_pass._eligible_activations(graph, 0, 1, {}) + + assert [node.name for node, _, _ in eligible] == ["base"] + assert offload_pass._skipped["alias"] == LARGE_SIZE + + +def test_a_saved_view_is_charged_for_the_allocation_it_holds_not_its_own_bytes(): + """Keeping a view resident retains the whole allocation it points into. + + The planner spends headroom to bring values back. Charging the view's own bytes would let a + quarter-sized slice retain four times the memory the planner thought it had spent. + """ + _ensure_dc_ops() + quarter = LARGE_NUMEL // 4 + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = _meta_tensor(16) + base = _add_node(graph, torch.relu, (x, ), "base", LARGE_NUMEL) + sliced = graph.create_node('call_function', torch.ops.aten.slice.Tensor, (base, 0, 0, quarter), {}, name="sliced") + # A prefix of a contiguous tensor: its elements do not overlap, and it is a quarter the size + # of the allocation it keeps alive. + sliced.meta["val"] = torch.empty(LARGE_NUMEL, device="meta")[:quarter] + out = _add_node(graph, torch.sum, (sliced, ), "out", 1) + graph.output((out, sliced)) + + eligible = offload_pass._eligible_activations(graph, 0, 1, {}) + + assert [(node.name, copied, resident) + for node, copied, resident in eligible] == [("sliced", LARGE_SIZE // 4, LARGE_SIZE)] def test_fwd_skips_values_already_on_the_host(forced_budget):