已开启
【需求】Activation offload 支持 device storage 粒度 offload #185
DavidFFFan创建于  6月4日
DavidFFFan
DavidFFFan成员
6月4日 创建

背景

当前 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=Falsegroup_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
likedislike
DavidFFFanDavidFFFan成员
6月4日 修改标题为 “【需求】Activation offload 支持 device storage 粒度去重”,原标题为“【需求】Activation offload 支持 device storage 粒度去重并保持 alias 语义”
DavidFFFanDavidFFFan成员
6月4日 修改标题为 “【需求】Activation offload 支持 device storage 粒度”,原标题为“【需求】Activation offload 支持 device storage 粒度去重”
DavidFFFanDavidFFFan成员
6月4日 修改标题为 “【需求】Activation offload 支持 device storage 粒度 offload”,原标题为“【需求】Activation offload 支持 device storage 粒度”
stillontheroadstillontheroad
6月22日 关联了pull request:feat(activation-swap): storage-level alias-aware swap dedup and preservation
DavidFFFanDavidFFFan成员
6月25日 添加了label:feature