`activation_offload.py` — 激活值 CPU 卸载¶

文件路径: verl/utils/activation_offload.py

文件概述¶

activation_offload.py 实现了前向传播中保存的激活值（saved for backward）的 CPU 卸载机制。在反向传播时再从 CPU 加载回 GPU。这样可以大幅降低 GPU 显存峰值。

背景知识¶

训练神经网络时，前向传播会保存中间激活值以供反向传播使用。对于大型 Transformer 模型，这些激活值占据大量 GPU 显存。激活卸载将它们暂时移到 CPU 内存，用时再取回。

相比激活检查点（gradient checkpointing，丢弃激活并在反向传播时重算），激活卸载不增加计算量，但增加了 CPU-GPU 通信。

核心类详解¶

1. CpuOffloadHookWithOffloadHandler — 上下文管理器¶

class CpuOffloadHookWithOffloadHandler:
    def __enter__(self):
        torch._C._autograd._push_saved_tensors_default_hooks(
            self.on_save_for_backward,
            self.on_get_saved_tensor
        )

    def on_save_for_backward(self, tensor):
        # 前向传播保存张量时，交给 handler 决定是否卸载
        return self.offload_handler.tensor_push(tensor)

    def on_get_saved_tensor(self, saved_state):
        # 反向传播取回张量时，从 handler 恢复
        return self.offload_handler.tensor_pop(saved_state)

利用 PyTorch 的 saved_tensors_hooks 机制，拦截所有 autograd 保存的张量。

2. SynchronizedGroupOffloadHandler — 同步卸载¶

class SynchronizedGroupOffloadHandler(OffloadHandler):
    @staticmethod
    def offload(src_tensor, pin_memory=True):
        cpu_backup = torch.empty(..., device="cpu", pin_memory=pin_memory)
        cpu_backup.copy_(src_tensor, non_blocking=True)
        return (src_tensor.device, cpu_backup)

    @staticmethod
    def reload(state, non_blocking=None):
        dev, cpu_backup = state
        return cpu_backup.to(dev, non_blocking=non_blocking)

将张量复制到 CPU 固定内存（pinned memory），使用 non_blocking=True 实现异步传输。

3. AsyncDoubleBufferGroupOffloadHandler — 异步双缓冲卸载¶

class AsyncDoubleBufferGroupOffloadHandler(SynchronizedGroupOffloadHandler):
    def __init__(self, num_offload_group, num_model_group, ...):
        self.d2h_stream = get_torch_device().Stream()  # Device to Host 流
        self.h2d_stream = get_torch_device().Stream()  # Host to Device 流

使用独立的 CUDA stream 进行 D2H（设备到主机）和 H2D（主机到设备）传输，与计算 stream 重叠。"双缓冲"意味着同时有两组激活在 GPU 上（当前层 + 预取的下一组）。

4. 按层分组卸载¶

class GroupCommitFunction(torch.autograd.Function):
    """标记层边界，触发一组激活的卸载/恢复"""
    @staticmethod
    def forward(ctx, tensor, cpu_offload_handler):
        cpu_offload_handler.on_group_commit_forward()
        return tensor

每个 Transformer 层结束时调用 group_commit，触发该层激活的批量卸载。

5. 入口函数¶

def enable_activation_offloading(model, strategy, enable_ckpt=False):
    """为模型启用激活卸载"""
    layers = []  # 找到所有 FSDP 包装的层
    get_layers(model)
    # 创建 handler 并包装每层的 forward 方法
    handler = ActivationHandler(context, sync_func, tensor_filter, enable_ckpt)
    for layer in layers:
        handler.wrap_module_forward_method(module)

工作流程图¶

前向传播:
Layer 0 forward → 激活卸载到 CPU (d2h_stream)
Layer 1 forward → Layer 0 卸载完成, Layer 1 激活卸载
  ...

反向传播:
Layer N backward ← 从 CPU 预取 Layer N-1 激活 (h2d_stream)
Layer N-1 backward ← 激活已就绪
  ...

核心类/函数列表¶

类/函数	说明
`CpuOffloadHookWithOffloadHandler`	卸载上下文管理器
`SynchronizedGroupOffloadHandler`	同步卸载
`AsyncDoubleBufferGroupOffloadHandler`	异步双缓冲卸载
`GroupCommitFunction`	层边界标记
`ActivationHandler`	层级激活管理
`enable_activation_offloading()`	启用入口
`FSDPParameterFilter`	过滤模型参数（不卸载参数本身）

与其他模块的关系¶

依赖 device.py 获取 CUDA stream
依赖 fsdp_utils.py 检测 FSDP 层
被 Actor worker 在训练时启用

小结¶

activation_offload.py 实现了高效的异步双缓冲激活卸载，通过将前向传播的中间激活值暂存到 CPU，显著降低 GPU 显存需求，同时通过异步传输尽量隐藏通信开销。

activation_offload.py — 激活值 CPU 卸载¶