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.