梯度检查点（Gradient Checkpointing）

# 前向传播：保存所有激活值
x1 = layer1(x0)  # 保存x1
x2 = layer2(x1)  # 保存x2
x3 = layer3(x2)  # 保存x3
loss = loss_fn(x3)

# 反向传播：使用保存的激活值
grad_x2 = backward(loss, x3)  # 使用保存的x3
grad_x1 = backward(grad_x2, x2)  # 使用保存的x2

# 前向传播：只保存检查点
x1 = checkpoint(layer1, x0)  # 不保存x1
x2 = checkpoint(layer2, x1)  # 保存x2（检查点）
x3 = checkpoint(layer3, x2)  # 不保存x3

# 反向传播：重新计算激活值
x3 = layer3(x2)  # 重新计算x3
grad_x2 = backward(loss, x3)
x1 = layer1(x0)  # 重新计算x1
grad_x1 = backward(grad_x2, x2)

def selective_checkpoint_policy(op, *args):
    # 保存计算密集操作
    if op in [torch.matmul, torch.bmm]:
        return CheckpointPolicy.MUST_SAVE
    
    # 重计算轻量级操作
    if op in [torch.layer_norm, torch.relu, torch.gelu]:
        return CheckpointPolicy.PREFER_RECOMPUTE
    
    return CheckpointPolicy.PREFER_SAVE

import torch._functorch.config

# 设置激活内存使用为默认的99%
torch._functorch.config.activation_memory_budget = 0.99

# 自动选择最优检查点策略
model = torch.compile(model)

from torch.utils.checkpoint import checkpoint, create_selective_checkpoint_contexts

# 创建选择性检查点上下文
ops_to_save = {torch.ops.aten.mm, torch.ops.aten.addmm}
ops_to_recompute = {torch.ops.aten.sigmoid, torch.ops.aten.dropout}

def policy_fn(ctx, op, *args, **kwargs):
    if op in ops_to_save:
        return CheckpointPolicy.MUST_SAVE
    elif op in ops_to_recompute:
        return CheckpointPolicy.PREFER_RECOMPUTE
    else:
        return CheckpointPolicy.PREFER_SAVE

context = create_selective_checkpoint_contexts(policy_fn)

# 启用选择性激活重计算
python train.py --recompute-activations

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "model_name",
    gradient_checkpointing=True,
    gradient_checkpointing_kwargs={
        "use_reentrant": False  # 推荐设置
    }
)

# 调试模式检查非确定性
torch.utils.checkpoint.checkpoint(
    module, input, 
    use_reentrant=False,
    debug=True  # PyTorch 2.1+
)