【需求】Activation offload 支持 device storage 粒度 offload

背景

当前 HyperParallel 已支持 activation offload / swap，但注册和拷贝粒度主要是 tensor。对于两个或多个 tensor 共享同一个 device address / 底层 storage 的场景，可能重复执行 D2H/H2D 拷贝；在 group_swap=True 的路径下，还可能把共享 storage 的 alias tensor pack 到不同 slice，load 后破坏原本的 alias 关系。

当前源码观察

• hyper_parallel/platform/mindspore/activation_checkpoint/activation_swap.py 中，AsyncSaveOnCpu.pack_to_cpu() 每遇到一个 saved tensor 就创建一个 SwapTensor。
• hyper_parallel/core/activation_checkpoint/swap.py 中，SwapTensor.async_offload() 以单个 SwapTensor 为单位分配 CPU buffer 并 copy。
• 现有 SwapTensor.dedup_key() 包含 (device, data_ptr, storage_offset, storage_size, stride)，只能跳过几乎完全相同的 tensor/view；对于共享同一 data_ptr 但 shape/stride/offset 不同的 alias tensor，仍可能重复 offload。
• SwapGroup._collect_packable_tensors() 注释提到 shared-storage tensors 应 fall back，但当前代码没有按 storage identity 识别并排除 alias group。

需求

将 activation offload 的重复拷贝控制从 tensor 粒度增强到 device storage 粒度：

1. 新增 storage/alias 级别识别，例如以 (device, data_ptr, storage_size) 作为 storage key，识别共享底层 device storage 的 tensor group。
2. 对 exact duplicate 继续保留现有 fast path；对 shared-storage alias group 只 offload 一份底层 storage 或最小 union range。
3. 每个 alias tensor 保存独立 metadata，包括 shape、stride、storage_offset、dtype、device、version 等。
4. load 回 device 后恢复同一份 device storage，并通过 set_/等价机制恢复每个 alias tensor 的 view，保证 alias 关系不被破坏。
5. group_swap=True 路径需要 alias-aware：要么 shared-storage alias group 整体走 alias-safe 路径，要么明确从 packed bulk copy 中排除。
6. MindSpore 和 Torch backend 尽量保持一致行为。

验收标准

• 同一 storage、相同 view 被多次注册时，只发生一次 D2H/H2D 拷贝。
• 同一 storage、不同 shape/stride/offset 的 alias tensor 不发生重复全量拷贝，且 backward 前恢复后的 alias 关系保持一致。
• group_swap=True 下 shared-storage alias tensor 不会被拆成互不 alias 的 packed buffer slice。
• 增加覆盖以下场景的单元测试或 ST：
  • exact duplicate saved tensor；
  • reshape/view 共享 storage；
  • transpose/narrow 等不同 stride/offset 的 view；
  • group_swap=False 与 group_swap=True 两条路径。

相关文件

• hyper_parallel/core/activation_checkpoint/swap.py
• hyper_parallel/platform/mindspore/activation_checkpoint/activation_swap.py
• hyper_parallel/platform/torch/activation_checkpoint/activation_swap.py

schema_version: 1
source: gitcode
gitcode_repo: mindspore/hyper-parallel
gitcode_issue: 185
source_url: https://gitcode.com/mindspore/hyper-parallel/issues/185

Start by tracing AsyncSaveOnCpu.pack_to_cpu() in hyper_parallel/platform/mindspore/activation_checkpoint/activation_swap.py and SwapTensor.async_offload() plus SwapGroup._collect_packable_tensors() in hyper_parallel/core/activation_checkpoint/swap.py. Compare the MindSpore and Torch activation_swap.py paths, then add tests for exact duplicates, shared views, and both group_swap settings. Done means shared storage is copied once and alias relationships survive load.