Backgroud（背景信息）

在深度神经网络训练过程中（尤其是 Transformer 等大模型），反向传播需要依赖前向阶段产生的中间激活值（Activation）来计算梯度。
随着模型规模和层数的不断增加（例如上百层 Transformer），这些激活值会占用大量 device 显存（GPU / NPU），成为训练过程中的主要瓶颈。

Origin（信息来源）

该需求由hyper_parallel提出，旨在在显存受限场景下支持更大规模模型训练。

Benefit / Necessity （价值/作用）

在显存受限场景下支持更大规模模型训练。
通过将中间激活张量从 NPU 异步卸载至 CPU 内存，并在反向传播前按需异步加载回NPU，实现显存节省；同时结合块级调度与预取策略，减少加载延迟对训练速度的影响。

Design（设计方案）

基于saved_tensors_hooks 实现张量 Host-Device 卸载与预取机制。

• free与prefetch
  基于模型block进行生命周期管理，控制前向D2H过程的free和反向H2D的prefetch时机。
  通过prefetch机制提前加载后续所需张量，隐藏H2D延迟，进行低延迟恢复。
  

• 多流机制异步执行
  新建一条流通过非阻塞方式实现D2H和H2D的高效传输。
  通过event控制不同流之间的同步，确保数据安全性。
  

• 接口设计与使用

from hyper_parallel.platform.torch.activation_checkpoint import (
    swap_wrapper, ActivationPolicy, SwapManager
)

#策略函数
def policy_fn(x):
    if x.storage().size() <= 1024:
        return ActivationPolicy.SAVE    #SAVE表示该tensor将会save到device，不做swap
    return ActivationPolicy.SWAP    #SWAP表示该tensor将会被swap到CPU上

for i, layer in enumerate(model.layers):
    model.layers[i].attn = swap_wrapper(layer.attn, policy_fn)    #policy_fn可选