梯度检查点原理与实现

前向: x → Layer1 → a1 → Layer2 → a2 → Layer3 → a3 → Loss
保存: a1, a2, a3 (以及各层的中间状态)

反向: Loss ← ∂L/∂a3 ← ∂L/∂a2 ← ∂L/∂a1 ← ∂L/∂x
使用: 重新使用a3计算a2的梯度，重新使用a2计算a1的梯度，...

检查点设置（每2层一个检查点）:
保存: x, a2, a4, a6, ...

前向: x → L1 → a1 → L2 → a2* → L3 → a3 → L4 → a4* → ...
                                     (* = 保存的检查点)

反向:
  从a6*出发，重新前向计算a5, a4→计算梯度
  从a4*出发，重新前向计算a3, a2→计算梯度
  ...

from torch.utils.checkpoint import checkpoint

class TransformerLayer(nn.Module):
    def forward(self, x):
        # 标准前向传播
        x = self.attention(self.norm1(x)) + x
        x = self.ffn(self.norm2(x)) + x
        return x

class TransformerModel(nn.Module):
    def __init__(self, n_layers=32, use_checkpointing=True, checkpoint_every=2):
        super().__init__()
        self.layers = nn.ModuleList([TransformerLayer() for _ in range(n_layers)])
        self.use_checkpointing = use_checkpointing
        self.checkpoint_every = checkpoint_every
    
    def forward(self, x):
        for i, layer in enumerate(self.layers):
            if self.use_checkpointing and (i % self.checkpoint_every == 0):
                # 使用梯度检查点
                x = checkpoint(layer, x, use_reentrant=False)
            else:
                x = layer(x)
        return x

class ManualCheckpoint:
    def __init__(self, model, checkpoint_layers):
        """
        checkpoint_layers: set of layer indices that are checkpoints
        """
        self.model = model
        self.checkpoint_layers = checkpoint_layers
        self.saved_activations = {}  # 存储检查点层的输入
    
    def forward(self, x):
        """前向传播：只保存检查点层的输入"""
        current = x
        for i, layer in enumerate(self.model.layers):
            if i in self.checkpoint_layers:
                # 保存检查点输入（需要detach以断开计算图）
                self.saved_activations[i] = current.detach().requires_grad_(True)
                current = layer(self.saved_activations[i])
            else:
                current = layer(current)
        return current
    
    def backward(self, loss, optimizer):
        """反向传播：从最后一个检查点开始，逐段重计算"""
        # 标准反向传播到最后一个检查点
        loss.backward(retain_graph=True)
        
        # 从后向前逐段处理
        checkpoint_indices = sorted(self.checkpoint_layers, reverse=True)
        for i in checkpoint_indices[:-1]:  # 跳过第一个检查点
            # 获取保存的输入
            input_tensor = self.saved_activations[i]
            
            # 重新前向传播
            with torch.enable_grad():
                recomputed = self.model.layers[i](input_tensor)
                for j in range(i + 1, self.get_next_checkpoint(i)):
                    recomputed = self.model.layers[j](recomputed)
            
            # 计算梯度
            optimizer.zero_grad()
            recomputed.backward()
            
            # 将梯度传递到检查点输入
            input_tensor.grad = recomputed.grad

def adaptive_checkpoint_schedule(layers, memory_budget):
    """自适应检查点调度"""
    activations = [estimate_activation_size(layer) for layer in layers]
    total = sum(activations)
    
    if total <= memory_budget:
        return set()  # 不需要检查点
    
    # 贪心算法：优先保存激活值小的层
    checkpoints = set()
    current_memory = 0
    
    for i, act_size in enumerate(activations):
        if current_memory + act_size <= memory_budget / len(layers):
            checkpoints.add(i)
            current_memory += act_size
    
    return checkpoints

class AttentionWithCheckpoint(nn.Module):
    def forward(self, x):
        # 保存输入，但不保存注意力中间结果
        def attention_fn(x_inner):
            Q = self.wq(x_inner)
            K = self.wk(x_inner)
            V = self.wv(x_inner)
            attn = flash_attention(Q, K, V)
            return self.wo(attn)
        
        # 只检查点注意力部分
        attn_out = checkpoint(attention_fn, x, use_reentrant=False)
        return attn_out + x  # 残差连接

# ZeRO-3 + 梯度检查点
model = TransformerModel(use_checkpointing=True)
model = DeepSpeed.initialize(
    model=model,
    config={
        "zero_optimization": {
            "stage": 3,
            "overlap_comm": True,
        },
        "activation_checkpointing": {
            "number_checkpoints": 16,
            "cpu_checkpointing": False,
        }
    }
)

def checkpoint_to_cpu(layer, x):
    """将激活值保存到CPU的梯度检查点"""
    # 保存到CPU
    x_cpu = x.detach().to('cpu')
    
    def custom_forward(*args):
        return layer(args[0])
    
    # 使用PyTorch的checkpoint但手动管理存储
    return checkpoint(custom_forward, x, use_reentrant=False)