Offload a saved view when it is the last value holding its storage

[pengdurice]

Offload a saved view when it is the last value holding its storage

Fixes #8387

Problem

_eligible_activations in deepspeed/compile/passes/offload_activation.py dropped every
saved-for-backward value whose target is an aliasing op (view, permute, slice, expand,
detach, and every other aten op whose schema declares an aliasing tensor return):

# A value that only aliases another tensor shares its storage, so copying it out frees
# nothing while the tensor it aliases is still live.
if node.target in no_copy_ops:
    _skipped["alias"] += _skip_bytes(node)
    continue

There was no check on whether anything else still held that storage.

The rule is correct when the tensor the view came from is a weight, a graph input, or another
saved value. It is wrong when AOTAutograd saves only the view. The base node is then dead, but its
allocation is not: the returned view still points at it, so the caching allocator cannot reclaim
the block until the backward pass reads the view. Moving the view to the host is what releases the
whole allocation.

The fix

Three helpers, and the single alias skip becomes two narrower ones.

Helper	What it does
_alias_root	Follows a node back through aliasing ops to the node that allocated the storage. A node whose target is not an aliasing op is its own root. Results are cached per call.
_storage_keeper_counts	Counts, per storage root, the values that outlive the forward pass: graph placeholders (the caller owns those), values returned to the caller, and every saved value.
_has_reloadable_layout	True when the tensor is non-overlapping and dense.

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

The existing skip for "the base is also saved" is kept, exactly as the issue asks. What changed is
that it is now conditional on a liveness count instead of unconditional.

Tests

Added to tests/unit/v1/compile/test_offload_activation.py. The first two are the tests named in
the issue; neither existed in the tree or on master, so both are written as ordinary tests rather
than one of them as xfail(strict=True).

Test	Asserts
test_eligible_includes_saved_view_when_base_is_not_saved	the view is eligible and _skipped["alias"] == 0
test_eligible_skips_saved_view_when_base_is_also_saved	the view is skipped, only the base is eligible
test_eligible_skips_a_saved_view_the_host_copy_would_not_reproduce	an expanded view is skipped under alias_layout

The graph both of the first two use is the one the issue specifies: x -> relu(base) -> aten.view(viewed) -> sum(out), with AOT-shaped outputs (out, viewed) and (out, base, viewed).

The pre-existing test_fwd_skips_values_that_alias_another_tensor still passes unchanged. In that
graph the view's root is a placeholder, so it is still skipped.

Results

Run	Result
tests/unit/v1/compile/test_offload_activation.py, 1 GPU	32 passed, 2 skipped (they need world_size=2)
TestOffloadActivation::test_offload_activation_correctness, 2 GPUs	2 passed

torch 2.6.0+cu124, Python 3.10. The end-to-end test compares losses with offloading on and off at
DS_DC_OFFLOAD_ACT_MIN_SIZE_MB=0, so it now exercises the newly eligible views. yapf --diff and
flake8 --max-line-length=120 are clean on both changed files.

Measured in a real training run

Full detail and provenance in RUN_RESULTS.md. Qwen3-14B, 8xH200 (139.80 GiB usable), ZeRO-3,
tiled loss 8, expandable_segments:True. Two frozen trees differing in exactly one file; every
cell asserts its own import path and prints the file's sha256 before running.

Sequence 4096, micro-batch 4 — both trees complete, so this measures the size of the bug:

	baseline	fixed
bytes skipped under alias, per graph	49.66 GB	13.42 GB
floor moved	26.8 GB, 160 values	63.0 GB, 280 values
peak	134.97 GB	131.45 GB
steady step	11.49 s	10.95 s

Total saved for backward is about 76.7 GB per rank, so the old rule was refusing 65% of it.
With the budget forced to nothing so only the floor runs, peak falls from 134.74 GB to 99.98 GB.

Sequence 4096, micro-batch 5, floor only — the configuration the issue reports dying at global
step 0:

	baseline	fixed
bytes skipped under alias	62.08 GB	16.78 GB
floor moved	33.8 GB, 362 values	79.1 GB, 482 values
steps completed	0 of 12	12 of 12
verdict	out of memory, 603 MiB free	completed, peak 119.20 GB

Same node type, same configuration, same instruction to move everything possible. One file differs.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: d6c9cd9718

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[chatgpt-codex-connector]

Add the required Signed-off-by trailer

This is a one-parent, non-merge commit, but its commit message has no Signed-off-by trailer. Add the author sign-off before merging so the commit satisfies the repository's commit and CI requirements.

AGENTS.md reference: AGENTS.md:L6-L8

Useful? React with 👍 / 👎.

[tohtana]

Thank you @pengdurice for the update! I left a few comments about some cornor cases.

[tohtana]

This check still skips a saved view based on the amount of data in the view. For example, a view may reference 4 MiB of data within a 32 MiB GPU allocation. If no other tensor references that allocation, keeping the view keeps all 32 MiB allocated.

Offloading the view would copy only 4 MiB to CPU and allow the full 32 MiB GPU allocation to be freed. However, this check rejects it because 4 MiB is below the default 5 MiB threshold.

Please feel free to refer to this supplemental PR.

[tohtana]

I'm not sure it causes an actual issue, but it would be safe to restore the original strides when reloading the tensor onto the GPU. empty_like packs stepped slices into a dense buffer, while the reload nodes retain the original tensor metadata.
I think we can save the original shape and strides before offloading, then allocate the destination GPU tensor with those strides, for example using at::empty_strided, before copying the values back. The CPU buffer can remain compact.

The reload memory calculation should account for the storage required by the restored strides, including gaps between elements. This would make the runtime layout match the layout expected by the compiled backward.