【大模型】- 梯度检查点

宋 算法
  2026-08-19 09:24 2026-08-19 09:24 创建   2026-08-21 16:20:12 2026-08-21 16:20:12 更新      1.7k 字  7 分钟

梯度检查点

用计算换内存，训练更大模型

类型： 学习　｜　语言： Python　｜　🏷 前置：《扩展和分布式训练》（本系列第 4 篇）

学习目标

  • 理解梯度检查点的原理
  • 实现梯度检查点
  • 优化内存使用
  • 平衡计算和内存开销
  • 应用于大模型训练

梯度检查点概述

梯度检查点（Gradient Checkpointing）是一种用计算换内存的技术，通过在反向传播时重新计算前向传播的激活值，减少内存使用。

为什么需要梯度检查点

  • 内存瓶颈：大模型训练时，激活值占用大量内存
  • 批量大小限制：内存不足导致无法使用大批量
  • 模型大小限制：内存不足导致无法训练大模型

梯度检查点的原理

  • 前向传播：只保存部分激活值（检查点）
  • 反向传播：从检查点重新计算激活值
  • 内存节省：只保存 O(√n) 的激活值

实现

基本梯度检查点

1
2
3
4
5
6
7
8
9
10
11
12
13
14
import torch
from torch.utils.checkpoint import checkpoint

class GradientCheckpointing(nn.Module):
    def __init__(self, model):
        super().__init__()
        self.model = model
    
    def forward(self, x):
        # 使用梯度检查点
        return checkpoint(self._forward, x, use_reentrant=False)
    
    def _forward(self, x):
        return self.model(x)

手动梯度检查点

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
class ManualGradientCheckpointing(nn.Module):
    def __init__(self, layers):
        super().__init__()
        self.layers = nn.ModuleList(layers)
        
        # 选择检查点位置
        self.checkpoint_indices = self.select_checkpoint_indices()
    
    def select_checkpoint_indices(self):
        """选择检查点位置"""
        n_layers = len(self.layers)
        # 每隔 √n 层设置一个检查点
        checkpoint_interval = int(n_layers ** 0.5)
        
        return list(range(0, n_layers, checkpoint_interval))
    
    def forward(self, x):
        # 前向传播，保存检查点
        checkpoints = []
        
        for i, layer in enumerate(self.layers):
            if i in self.checkpoint_indices:
                # 保存检查点
                checkpoints.append(x.clone())
            
            x = layer(x)
        
        return x, checkpoints
    
    def backward(self, grad_output, checkpoints):
        """反向传播，重新计算激活值"""
        # 从最后一个检查点开始
        checkpoint_idx = len(checkpoints) - 1
        
        # 重新计算前向传播
        x = checkpoints[checkpoint_idx]
        
        for i in range(checkpoint_idx, len(self.layers)):
            x = self.layers[i](x)
        
        # 计算梯度
        x.backward(grad_output)
        
        return x

选择性梯度检查点

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
class SelectiveGradientCheckpointing(nn.Module):
    def __init__(self, model, sensitive_layers=None):
        super().__init__()
        self.model = model
        self.sensitive_layers = sensitive_layers or []
    
    def forward(self, x):
        # 对敏感层使用标准前向传播
        # 对其他层使用梯度检查点
        
        checkpoints = []
        
        for i, layer in enumerate(self.model.layers):
            if i in self.sensitive_layers:
                # 敏感层：保存激活值
                x = layer(x)
            else:
                # 其他层：使用梯度检查点
                x = checkpoint(layer, x, use_reentrant=False)
        
        return x

优化技术

自适应检查点

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
class AdaptiveGradientCheckpointing(nn.Module):
    def __init__(self, model, memory_budget):
        super().__init__()
        self.model = model
        self.memory_budget = memory_budget
        
        # 计算检查点间隔
        self.checkpoint_interval = self.compute_interval()
    
    def compute_interval(self):
        """根据内存预算计算检查点间隔"""
        # 估算每层的内存使用
        layer_memory = self.estimate_layer_memory()
        
        # 计算可以保存多少层的激活值
        n_layers = len(self.model.layers)
        max_checkpoints = self.memory_budget / layer_memory
        
        # 计算检查点间隔
        interval = max(1, int(n_layers / max_checkpoints))
        
        return interval
    
    def estimate_layer_memory(self):
        """估算每层的内存使用"""
        # 简化估算：实际应该根据模型配置计算
        return 1024 * 1024  # 1MB
    
    def forward(self, x):
        """自适应梯度检查点"""
        for i, layer in enumerate(self.model.layers):
            if i % self.checkpoint_interval == 0:
                # 使用梯度检查点
                x = checkpoint(layer, x, use_reentrant=False)
            else:
                # 标准前向传播
                x = layer(x)
        
        return x

混合精度梯度检查点

1
2
3
4
5
6
7
8
9
10
11
12
13
class MixedPrecisionGradientCheckpointing(nn.Module):
    def __init__(self, model):
        super().__init__()
        self.model = model
    
    def forward(self, x):
        # 使用混合精度
        with torch.cuda.amp.autocast():
            # 梯度检查点
            return checkpoint(self._forward, x, use_reentrant=False)
    
    def _forward(self, x):
        return self.model(x)

内存分析

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
def analyze_memory_usage(model, input_size, n_layers, use_checkpointing=False):
    """分析内存使用"""
    
    # 创建模型
    x = torch.randn(input_size)
    
    if use_checkpointing:
        # 使用梯度检查点
        model = GradientCheckpointing(model)
    
    # 前向传播
    output = model(x)
    
    # 计算内存使用
    if torch.cuda.is_available():
        memory_allocated = torch.cuda.memory_allocated() / 1024 / 1024
        memory_cached = torch.cuda.memory_reserved() / 1024 / 1024
    else:
        memory_allocated = 0
        memory_cached = 0
    
    print(f"使用梯度检查点: {use_checkpointing}")
    print(f"已分配内存: {memory_allocated:.2f}MB")
    print(f"缓存内存: {memory_cached:.2f}MB")
    
    return memory_allocated, memory_cached

内存节省计算

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
def calculate_memory_savings(n_layers, activation_size):
    """计算内存节省"""
    
    # 标准方法：保存所有激活值
    standard_memory = n_layers * activation_size
    
    # 梯度检查点：只保存 √n 层的激活值
    checkpoint_interval = int(n_layers ** 0.5)
    checkpoint_memory = checkpoint_interval * activation_size
    
    # 内存节省
    savings = (standard_memory - checkpoint_memory) / standard_memory
    
    print(f"标准内存使用: {standard_memory / 1e9:.2f}GB")
    print(f"梯度检查点内存: {checkpoint_memory / 1e9:.2f}GB")
    print(f"内存节省: {savings:.2%}")
    
    return savings

计算开销分析

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
def analyze_computational_overhead(model, input_size, n_layers, use_checkpointing=False):
    """分析计算开销"""
    
    # 创建模型
    x = torch.randn(input_size)
    
    # 计时
    start_time = time.time()
    
    if use_checkpointing:
        model = GradientCheckpointing(model)
    
    # 前向传播
    output = model(x)
    
    # 反向传播
    loss = output.sum()
    loss.backward()
    
    time_taken = time.time() - start_time
    
    print(f"使用梯度检查点: {use_checkpointing}")
    print(f"总时间: {time_taken:.4f}s")
    
    return time_taken

实际应用

大模型训练

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
def train_large_model_with_checkpointing():
    """使用梯度检查点训练大模型"""
    
    # 配置
    config = {
        "dim": 4096,
        "n_layers": 32,
        "n_heads": 32,
        "vocab_size": 32000,
        "batch_size": 32,
        "max_seq_len": 2048,
    }
    
    # 创建模型
    model = GPT(config)
    
    # 应用梯度检查点
    model = GradientCheckpointing(model)
    
    # 优化器
    optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)
    
    # 训练循环
    for batch in train_dataloader:
        input_ids = batch["input_ids"]
        labels = batch["labels"]
        
        # 前向传播
        logits = model(input_ids)
        
        # 计算损失
        loss = F.cross_entropy(
            logits.view(-1, logits.size(-1)),
            labels.view(-1)
        )
        
        # 反向传播
        loss.backward()
        
        # 更新参数
        optimizer.step()
        optimizer.zero_grad()

选择性检查点

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
def selective_checkpointing_training():
    """选择性梯度检查点训练"""
    
    # 识别敏感层
    sensitive_layers = identify_sensitive_layers(model)
    
    # 应用选择性检查点
    model = SelectiveGradientCheckpointing(model, sensitive_layers)
    
    # 训练
    for batch in train_dataloader:
        input_ids = batch["input_ids"]
        labels = batch["labels"]
        
        logits = model(input_ids)
        loss = F.cross_entropy(logits.view(-1, logits.size(-1)), labels.view(-1))
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

def identify_sensitive_layers(model):
    """识别敏感层"""
    sensitive = []
    
    for i, layer in enumerate(model.layers):
        # 评估层的敏感性
        sensitivity = evaluate_layer_sensitivity(layer)
        
        if sensitivity > threshold:
            sensitive.append(i)
    
    return sensitive

权衡分析

计算-内存权衡

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
def analyze_tradeoff(n_layers, activation_size, compute_budget, memory_budget):
    """分析计算-内存权衡"""
    
    # 标准方法
    standard_memory = n_layers * activation_size
    standard_compute = n_layers  # 简化
    
    # 梯度检查点
    checkpoint_interval = int(n_layers ** 0.5)
    checkpoint_memory = checkpoint_interval * activation_size
    checkpoint_compute = n_layers * 2  # 重新计算
    
    print("标准方法:")
    print(f"  内存: {standard_memory / 1e9:.2f}GB")
    print(f"  计算: {standard_compute}")
    
    print("\n梯度检查点:")
    print(f"  内存: {checkpoint_memory / 1e9:.2f}GB")
    print(f"  计算: {checkpoint_compute}")
    
    # 检查约束
    if standard_memory <= memory_budget:
        print("\n推荐: 标准方法（满足内存约束）")
    elif checkpoint_memory <= memory_budget:
        print("\n推荐: 梯度检查点（满足内存约束，增加计算）")
    else:
        print("\n警告: 内存不足，需要进一步优化")

最佳实践

  1. 选择合适的检查点间隔：通常 √n 层设置一个检查点
  2. 考虑敏感层：对敏感层使用标准前向传播
  3. 监控内存使用：确保内存使用在预算内
  4. 平衡计算和内存：根据硬件约束调整参数

总结

梯度检查点通过用计算换内存，使得训练更大的模型成为可能。关键组件包括检查点选择、内存分析和计算开销优化。选择合适的检查点策略需要平衡内存和计算约束。

下一步

下一步可以探索更多优化技术，如模型并行、流水线并行等。

📚 本文改编自 AI Engineering from Scratch（MIT License · 作者 Rohit Ghumare），中文内容来自官方中文镜像。原课程共 503 课 · 20 阶段 · 免费开源，教程网站见 aiengineeringfromscratch.com。

  • 标题: 【大模型】- 梯度检查点
  • 作者: 宋
  • 创建于 : 2026-08-19 09:24:00
  • 更新于 : 2026-08-21 16:20:12
  • 链接: https://sxl-space.tk/2026/08/19/010_LLM/010_LLM-24-GradientCheckpointing/
  • 版权声明: 版权所有 © 宋，禁止转载。
目录
【大模型】- 梯度检查点
  1. 1. 梯度检查点
  2. 2. 学习目标
  3. 3. 梯度检查点概述
    1. 3.1. 为什么需要梯度检查点
    2. 3.2. 梯度检查点的原理
  4. 4. 实现
    1. 4.1. 基本梯度检查点
    2. 4.2. 手动梯度检查点
    3. 4.3. 选择性梯度检查点
  5. 5. 优化技术
    1. 5.1. 自适应检查点
    2. 5.2. 混合精度梯度检查点
  6. 6. 内存分析
    1. 6.1. 内存节省计算
  7. 7. 计算开销分析
  8. 8. 实际应用
    1. 8.1. 大模型训练
    2. 8.2. 选择性检查点
  9. 9. 权衡分析
    1. 9.1. 计算-内存权衡
  10. 10. 最佳实践
  11. 11. 总结
  12. 12. 下一步
  13. 13. 下一步可以探索更多优化技术，如模型并行、流水线并行等。