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

=================================================================================================
                                   Ringi 3D 架构工坊 · 核心全景
┌───────────────────────────────────────────────────────────────────────────────────────────────┐
│                                                                                               │
│  [前向计算流 Forward Pass: 激活值扣留阶段]                                                    │
│  Input (X) ──► [Linear / GEMM] ──► Activation (A1) ──► [LayerNorm / ReLU] ──► Output (Y)      │
│     │               │ (自动注册)                              │ (自动注册)             │      │
│     │               ▼                                         ▼                        ▼      │
│  [底层 DAG]   AccumulateGrad ◄── MmBackward0 ◄────────── ReluBackward0 ◄── MSELossBackward   │
│                      ▲                │ (ctx.save_for_backward)       │ (ctx.save_for_backward)│
│                      │                ▼                               ▼                        │
│                [权重梯度 W.grad]  [扣留 X, W 到 HBM]              [扣留 A1 或 Mask 到 HBM]     │
│                                                                                               │
│  ───────────────────────────────────────────────────────────────────────────────────────────  │
│                                                                                               │
│  [反向求导流 Backward Pass: 梯度累加与激活值退火归还阶段]                                      │
│  Loss.backward()                                                                              │
│     │                                                                                         │
│     ▼                                                                                         │
│  MSELossBackward ──► 消费并释放 Loss 临时张量 ──► 显存退火归还                                │
│     │                                                                                         │
│     ▼                                                                                         │
│  ReluBackward0   ──► 取出 A1/Mask 计算 dA1 ──► 立即释放 A1 显存 ──► 显存进一步退火            │
│     │                                                                                         │
│     ▼                                                                                         │
│  MmBackward0     ──► 取出 X, W 执行 2 次 GEMM ──► 累加 W.grad ──► 释放 X, W 临时激活值        │
│                                                                                               │
└───────────────────────────────────────────────────────────────────────────────────────────────┘
=================================================================================================

1
2
3
4

torch.cuda.OutOfMemoryError: CUDA out of memory.
Tried to allocate 2.00 GiB (GPU 0; 79.15 GiB total capacity;
77.34 GiB already allocated; 1.21 GiB free;
77.50 GiB reserved in total by PyTorch)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12

# ❌ 一个价值数十万元的显存泄漏 Bug
loss_history = []

for step, (inputs, targets) in enumerate(train_dataloader):
    optimizer.zero_grad()
    outputs = model(inputs)
    loss = criterion(outputs, targets)
    loss.backward()
    optimizer.step()

    # 看起来极其无辜的一行打点记录：
    loss_history.append(loss)  # <--- 致命死穴！

1
2
3
4
5
6
7

+---------------------------------------------------------------------------------------+
|  动态图的核心物理直觉：                                                               |
|  在前向传播执行 Python 代码的瞬间，PyTorch 实际上在做两件事：                        |
|  1. [正向业务]: 调用底层的 ATen / cuBLAS / cuDNN Kernel，把输出数值实打实地算出来；  |
|  2. [幕后织网]: 顺藤摸瓜，在 C++ 堆内存中记录下每一个算子的逆向操作节点，            |
|                并将它们用指针串联成一张反向有向无环图（Backward DAG）！               |
+---------------------------------------------------------------------------------------+

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12

import torch

x = torch.tensor([2.0, 3.0], requires_grad=True)
w = torch.tensor([1.5, 2.5], requires_grad=True)
b = torch.tensor([0.5, 0.5], requires_grad=True)

y = x * w + b
loss = y.sum()

print(f"x.is_leaf: {x.is_leaf}, x.requires_grad: {x.requires_grad}, x.grad_fn: {x.grad_fn}")
print(f"y.is_leaf: {y.is_leaf}, y.requires_grad: {y.requires_grad}, y.grad_fn: {y.grad_fn}")
print(f"loss.grad_fn: {loss.grad_fn}")

1
2
3

x.is_leaf: True, x.requires_grad: True, x.grad_fn: None
y.is_leaf: False, y.requires_grad: True, y.grad_fn: <AddBackward0 object at 0x7f88a0>
loss.grad_fn: <SumBackward0 object at 0x7f88b2>

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23

// 摘自 PyTorch 核心源码 torch/csrc/autograd/function.h (简化版)
struct TORCH_API Node : std::enable_shared_from_this<Node> {
 public:
  // 执行反向求导计算的核心虚函数
  virtual variable_list apply(variable_list&& inputs) = 0;

  // 获取指向前序节点的有向边（边连接着更早生成的 Node）
  const edge_list& next_edges() const {
    return next_edges_;
  }

  // 拓扑排序与执行序列号
  uint64_t sequence_nr() const {
    return sequence_nr_;
  }

 private:
  // 记录在反向计算时，梯度应当流向哪些下游节点以及对应的输入插槽（Input Slot）
  edge_list next_edges_;

  // 线程安全的执行序列号，用于拓扑排序与优先级队列调度
  uint64_t sequence_nr_;
};

1
2
3
4
5
6
7
8

// 摘自 torch/csrc/autograd/edge.h
struct Edge {
  // 指向目标 Node 的智能指针
  std::shared_ptr<Node> function;

  // 目标 Node 的第几个输入参数插槽（Slot Index）
  uint32_t input_nr;
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20

import torch

x = torch.randn(2, 3, requires_grad=True)
w = torch.randn(3, 4, requires_grad=True)
b = torch.randn(4, requires_grad=True)

# 前向计算
h = torch.matmul(x, w)       # 生成 MmBackward0 节点
y = h + b                    # 生成 AddBackward0 节点
loss = y.sum()               # 生成 SumBackward0 节点

print("=== 查看反向 DAG 节点链 ===")
print("1. loss.grad_fn:", loss.grad_fn)
print("2. loss.grad_fn.next_functions:", loss.grad_fn.next_functions)

add_node = loss.grad_fn.next_functions[0][0]
print("3. add_node.next_functions:", add_node.next_functions)

mm_node = add_node.next_functions[0][0]
print("4. mm_node.next_functions:", mm_node.next_functions)

1
2
3
4
5

=== 查看反向 DAG 节点链 ===
1. loss.grad_fn: <SumBackward0 object at 0x7f8810>
2. loss.grad_fn.next_functions: ((<AddBackward0 object at 0x7f8818>, 0),)
3. add_node.next_functions: ((<MmBackward0 object at 0x7f8820>, 0), (<AccumulateGrad object at 0x7f8828>, 0))
4. mm_node.next_functions: ((<AccumulateGrad object at 0x7f8830>, 0), (<AccumulateGrad object at 0x7f8838>, 0))

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16

// 伪代码：torch/csrc/autograd/functions/accumulate_grad.cpp
variable_list AccumulateGrad::apply(variable_list&& grads) {
  auto& grad = grads[0];
  auto variable = variable_.lock(); // 获取绑定的叶子 Tensor 弱引用

  if (!variable->grad().defined()) {
    // 如果这是反向传播第一次给该叶子节点算出的梯度：
    // 直接把这块显存赋给 param.grad！
    variable->mutable_grad() = grad;
  } else {
    // 如果该叶子节点被多处使用（例如同一个权重在多处共享）：
    // 执行就地累加：param.grad += grad！
    variable->mutable_grad().add_(grad);
  }
  return variable_list(); // 叶子节点没有下游，返回空列表
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21

class MyLinearFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, w):
        # 1. 执行前向计算：产出结果
        y = torch.matmul(x, w)

        # 2. 扣押人质！把 x 和 w 存入上下文对象中，防止被显存池垃圾回收！
        ctx.save_for_backward(x, w)
        return y

    @staticmethod
    def backward(ctx, grad_output):
        # 3. 反向传播时，重新提审人质：
        x, w = ctx.saved_tensors

        # 4. 计算梯度
        grad_x = torch.matmul(grad_output, w.t())
        grad_w = torch.matmul(x.t(), grad_output)

        # 5. 函数执行完毕，x 和 w 在此处的生命周期终结！
        return grad_x, grad_w

 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

+---------------------------------------------------------------------------------------+
|                       PyTorch Autograd Engine 执行六步流水线                          |
|                                                                                       |
|  [Step 1: 入口检查]                                                                   |
|  验证 loss 是否为标量（若为张量需传入 gradient 权重）；检查是否有 requires_grad=True。 |
|                                                                                       |
|  [Step 2: 拓扑排序 (Topological Sort)]                                                |
|  从 loss.grad_fn 开始，通过深度优先搜索（DFS）遍历整张有向无环图，计算每个 Node 的入度 |
|  （In-degree dependencies），构建就绪队列（Ready Queue）。                             |
|                                                                                       |
|  [Step 3: 线程池启动与并发执行]                                                       |
|  Autograd 引擎拥有独立的 C++ 工作线程池（默认线程数等于 CPU 核心数）。                |
|  将入度为 0 的节点（即 loss.grad_fn）压入执行队列。                                   |
|                                                                                       |
|  [Step 4: 节点 Apply 与 Kernel 派发]                                                  |
|  工作线程从队列中弹出 Node，调用 node->apply(grads)：                                 |
|  - 4.1 从 SavedVariable 中解包（Unpack）前向暂存的张量；                             |
|  - 4.2 派发反向计算的 CUDA Kernel 到对应的 CUDA Stream 中异步执行；                  |
|  - 4.3 产出对输入端各 Slot 的输出梯度。                                              |
|                                                                                       |
|  [Step 5: 现场销毁 Saved Tensors (显存释放的关键时刻！)]                               |
|  node 执行完毕后，调用 node->release_saved_variables()！解绑对前向张量的强引用！      |
|  StorageImpl 引用计数归零，内存块立即返还给 CUDACachingAllocator 显存池！             |
|                                                                                       |
|  [Step 6: 依赖递减与下游推进]                                                         |
|  将当前节点产生的梯度累加到对应下级 Node 的输入缓存槽中；                            |
|  将对应下级 Node 的依赖计数减 1；若依赖计数清零，则将其推入 Ready Queue 继续执行；    |
|  直至所有节点处理完毕，遇到 AccumulateGrad 则直接把最终梯度累加写入 param.grad！       |
+---------------------------------------------------------------------------------------+

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14

// 核心逻辑等价伪代码：torch/csrc/autograd/engine.cpp
void evaluate_node(Node* node, variable_list&& inputs) {
  // 1. 调用反向求导计算
  variable_list outputs = node->apply(std::move(inputs));

  // 2. 极其关键的一行代码：立刻清理本节点保存的所有前向变量！
  node->release_saved_variables();

  // 3. 将产出的 outputs 梯度向前序节点传递
  for (size_t i = 0; i < outputs.size(); ++i) {
    auto& edge = node->next_edge(i);
    // ... 传递梯度并递减依赖计数 ...
  }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18

显存占用 (VRAM)
  ▲
  │                                    [★ 显存峰值 Peak Memory]
  │                                           ┌───┐
  │                                          ┌┘   └┐
  │                                         ┌┘     └┐ (反向开始：梯度逐步生成，
  │                                        ┌┘       └┐ 激活值一路就地释放)
  │                                       ┌┘         └┐
  │                  (前向传播 Forward:  ┌┘           └┐
  │                   Activation 逐层累加)┌┘             └┐
  │                                     ┌┘               └┐
  │                                    ┌┘                 └──┐
  │                                   ┌┘                     └──┐
  │ 静态显存基线 ────────────────────┘                         └─── [反向结束: 激活全清空,
  │ (Weights + Gradients + Optimizer)                                 只留 W.grad 显存]
  │
  └─────────────────────────────────────────────────────────────────────────────► 时间 (t)
    [Zero Grad]     [Forward Pass 阶段]     [Loss]   [Backward Pass 阶段]     [Step]

1

optimizer.zero_grad(set_to_none=True)

1
2
3

RuntimeError: one of the variables needed for gradient computation has been modified
by an inplace operation: [torch.cuda.FloatTensor [1024, 1024]], which is output 0
of ReluBackward0, is at version 1; expected version 0 instead.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10

import torch

x = torch.randn(4, 4, requires_grad=True)
y = torch.relu(x)  # 前向传播：y 需要被 save_for_backward

# 危险原地操作：在 y 被保存后，偷偷修改了 y 自身指向的物理显存！
y += 1.0  # 等价于 y.add_(1.0)

loss = y.sum()
loss.backward()  # 💥 瞬间爆炸！抛出上述 RuntimeError！

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15

// 摘自 c10/core/TensorImpl.h
struct TORCH_API TensorImpl : public c10::intrusive_ptr_target {
  // ...
  inline void bump_version() {
    version_counter_.bump();
  }

  inline uint32_t current_version() const {
    return version_counter_.current_version();
  }

 private:
  // 线程安全的版本计数器结构体
  c10::VariableVersion version_counter_;
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17

+---------------------------------------------------------------------------------------+
|  PyTorch 版本计数器追踪四部曲：                                                       |
|                                                                                       |
|  1. [出生时刻]: 一个 Tensor 被创建出来时，version_counter_ 初始化为 0；               |
|                                                                                       |
|  2. [扣押人质时刻]: 在前向算子调用 ctx.save_for_backward(y) 时：                     |
|     SavedVariable 会在 C++ 内部把该张量当前的 version_counter_（即 0）记录在案！      |
|                                                                                       |
|  3. [暗度陈仓时刻]: 任何就地修改该张量物理内存的算子（凡是以下划线结尾的 API，       |
|     如 add_(), mul_(), relu_(), 或者切片赋值 y[0] = 1）：                            |
|     PyTorch 内部派发器必须强制触发 y.unsafeGetTensorImpl()->bump_version()！          |
|     版本号自增为 1！                                                                  |
|                                                                                       |
|  4. [提审校验时刻]: 反向传播时，SavedVariable::unpack() 提审该张量：                  |
|     立即对比：saved_version (0) 是否等于 current_version (1)？                        |
|     若不相等，立刻中止执行并抛出 RuntimeError，严正拒绝静默数值污染！                 |
+---------------------------------------------------------------------------------------+

1
2
3
4
5
6
7
8
9

// 核心机制等价实现：c10/core/GradMode.h
struct TORCH_API GradMode {
  static bool is_enabled() {
    return grad_mode_flag;
  }
  static void set_enabled(bool enabled) {
    grad_mode_flag = enabled;
  }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13

+---------------------------------------------------------------------------------------+
|                       no_grad() vs inference_mode() 底层机制对比                      |
|                                                                                       |
|  维度                  torch.no_grad()              torch.inference_mode()            |
|  ───────────────────────────────────────────────────────────────────────────────────  |
|  是否构建 Autograd DAG   ❌ 否 (不构图)               ❌ 否 (不构图)                   |
|  是否扣留 Activation     ❌ 否 (不扣留)               ❌ 否 (不扣留)                   |
|  版本计数器维护          ✅ 仍维护 version_counter_   ❌ 彻底冻结 (不记录版本号)       |
|  视图（View）跟踪开销    ✅ 仍维持 View 依赖元数据    ❌ 彻底解绑 (零 View 跟踪)       |
|  生成张量能否用于训练    ✅ 可以（脱离上下文后可用）   ❌ 抛出异常 (严禁带入训练图)     |
|  C++ 派发器路径          经过 Autograd 检查层         直接 bypass，直达原生 Kernel     |
|  吞吐量表现              快                           极致快（再快 5% ~ 15%）          |
+---------------------------------------------------------------------------------------+

1
2
3
4
5
6
7
8
9

=================================================================================================
                                   Ringi 3D 架构工坊 · 核心全景
[普通训练: 80 层 Activation 全量堆积在显存]
 Layer 1 Act ──► Layer 2 Act ──► ... ──► Layer 80 Act  ====> HBM 直接被塞满 80GB OOM 爆炸！

[Activation Checkpointing 核心机制: 只存关卡 Checkpoint，中间全扔掉！]
 Layer 1 In (存!) ─[前向算完把中间全扔了]─► Layer 2 In (存!) ─[中间全扔了]─► ... ─► Layer 80
 显存中只保留 80 个轻量级的 Block 输入张量！空间复杂度直接降为 O(1) 块级空间！
=================================================================================================

1
2
3
4
5
6

+---------------------------------------------------------------------------------------+
|  Selective Checkpointing 黄金法则：                                                   |
|  - 所有的矩阵乘法（GEMM）激活值：通通保留在显存中（不重算，保住算力！）；            |
|  - 所有的 Softmax / Dropout / Norm / GELU 激活值：统统丢弃，反向时就地重算！          |
|  - 最终战果：砍掉 70% 的激活显存峰值，而计算时间仅仅增加区区 2% ~ 4%！                |
+---------------------------------------------------------------------------------------+

 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
45
46
47
48
49
50
51
52
53
54
55

import torch

def print_computation_graph(root_fn, indent=0, visited=None):
    """
    递归遍历 PyTorch 反向动态计算图 DAG 并打印 ASCII 拓扑
    """
    if visited is None:
        visited = set()
        print("\n" + "=" * 60)
        print("          Ringi 动态计算图（DAG）拓扑遍历器")
        print("=" * 60)

    if root_fn is None:
        return

    # 防止环状引用或多路汇聚节点的重复打印
    fn_id = id(root_fn)
    is_visited = fn_id in visited
    visited.add(fn_id)

    prefix = "  " * indent + "└─► "
    node_name = type(root_fn).__name__

    # 打印节点与额外属性
    extra_info = ""
    if hasattr(root_fn, "variable"):
        # 叶子节点绑定的张量信息
        var = root_fn.variable
        extra_info = f" [Leaf Variable Shape: {list(var.shape)}, requires_grad={var.requires_grad}]"

    status = " (Re-visited)" if is_visited else ""
    print(f"{prefix}[{node_name}]{extra_info}{status}")

    if not is_visited and hasattr(root_fn, "next_functions"):
        for next_fn, slot_idx in root_fn.next_functions:
            if next_fn is not None:
                print_computation_graph(next_fn, indent + 1, visited)

if __name__ == "__main__":
    # 构造一个包含线性变换、激活函数、分支聚合的前向计算
    x = torch.randn(2, 4, requires_grad=True)
    w1 = torch.randn(4, 8, requires_grad=True)
    b1 = torch.randn(8, requires_grad=True)
    w2 = torch.randn(8, 1, requires_grad=True)

    # 前向过程
    h1 = torch.matmul(x, w1) + b1
    act = torch.relu(h1)

    # 分支计算：一部分直接乘，一部分经由激活函数相加
    out = torch.matmul(act, w2) + (act * 0.5).sum()
    loss = out.mean()

    # 启动遍历
    print_computation_graph(loss.grad_fn)

 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
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66

import torch

def format_mb(bytes_val):
    return f"{bytes_val / (1024 ** 2):.2f} MB"

def trace_memory():
    if not torch.cuda.is_available():
        print("请在配备 NVIDIA GPU 的环境下运行本实验！")
        return

    device = torch.device("cuda:0")
    torch.cuda.empty_cache()
    torch.cuda.reset_peak_memory_stats()

    print("\n" + "=" * 65)
    print("        PyTorch 单训练 Step 显存生命周期微观追踪器")
    print("=" * 65)

    def log_point(stage_name):
        alloc = torch.cuda.memory_allocated(device)
        peak = torch.cuda.max_memory_allocated(device)
        print(f"[{stage_name:<20}] 已分配: {format_mb(alloc):>9} | 历史峰值: {format_mb(peak):>9}")

    log_point("0. 初始空闲状态")

    # 1. 创建模型权重（静态显存）
    # 构造一个拥有较大维度的矩阵以凸显显存波动
    w1 = torch.randn(4096, 4096, device=device, requires_grad=True)
    w2 = torch.randn(4096, 4096, device=device, requires_grad=True)
    log_point("1. 权重初始化就绪")

    # 2. 输入数据进入 GPU
    x = torch.randn(2048, 4096, device=device, requires_grad=True)
    log_point("2. 输入 X 送入显存")

    # 3. 前向传播：第 1 层 Matmul
    h1 = torch.matmul(x, w1)
    log_point("3. 前向: Matmul-1 产出")

    # 4. 前向传播：GELU 激活（需扣留输入用于求导）
    h2 = torch.nn.functional.gelu(h1)
    log_point("4. 前向: GELU 产出")

    # 5. 前向传播：第 2 层 Matmul
    out = torch.matmul(h2, w2)
    log_point("5. 前向: Matmul-2 产出")

    # 6. 计算标量 Loss
    loss = out.sum()
    log_point("6. 计算 Loss (峰值预警)")

    # 7. 反向传播开始！
    loss.backward()
    log_point("7. Backward 完成!")

    # 8. 观察梯度与激活值状态
    del h1, h2, out, loss, x
    log_point("8. 销毁中间 Python 变量")

    # 9. 梯度清零对比
    w1.grad = None
    w2.grad = None
    log_point("9. grad 置为 None")

if __name__ == "__main__":
    trace_memory()

 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

import torch

def inplace_bug_reproduce():
    print("\n" + "=" * 60)
    print("      In-place 版本计数器崩溃现场与安全重构实验")
    print("=" * 60)

    # 1. 制造崩溃现场
    a = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
    b = a * 2.0  # b.grad_fn = MulBackward0

    print(f"初始状态: b._version = {b._version}")

    # 悄悄原地修改 b
    b.add_(10.0)
    print(f"原地修改后: b._version = {b._version}")

    try:
        loss = (b ** 2).sum()
        loss.backward()
    except RuntimeError as e:
        print("\n[捕获预期中的 RuntimeError 崩溃!]")
        print("错误信息前缀:", str(e).split("\n")[0])

    # 2. 安全优雅的替代方案：
    print("\n[执行安全重构方案...]")
    a_safe = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
    b_safe = a_safe * 2.0
    # 拒绝原地修改，生成全新的 Out-of-place Tensor！
    b_clean = b_safe + 10.0

    loss_safe = (b_clean ** 2).sum()
    loss_safe.backward()
    print("安全反向传播成功！a_safe.grad =", a_safe.grad)

if __name__ == "__main__":
    inplace_bug_reproduce()

 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
45
46
47
48

import torch
import torch.nn as nn

class MemoryProbeModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(2048, 4096)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(4096, 1024)

    def forward(self, x):
        return self.fc2(self.relu(self.fc1(x)))

def register_diagnostic_hooks(model):
    def forward_hook(module, args, output):
        # 统计产出的 Activation 大小
        act_size = output.element_size() * output.nelement() / (1024 ** 2)
        print(f"  [Forward Hook]  层: {module.__class__.__name__:<10} | "
              f"输出 Shape: {str(list(output.shape)):<18} | Activation 显存: {act_size:.2f} MB")

    def backward_hook(module, grad_input, grad_output):
        # 统计回传的梯度大小
        if grad_output[0] is not None:
            g_out = grad_output[0]
            grad_size = g_out.element_size() * g_out.nelement() / (1024 ** 2)
            print(f"  [Backward Hook] 层: {module.__class__.__name__:<10} | "
                  f"回传梯度 Shape: {str(list(g_out.shape)):<18} | 梯度显存: {grad_size:.2f} MB")

    for name, layer in model.named_modules():
        if len(list(layer.children())) == 0:  # 只针对原子算子层
            layer.register_forward_hook(forward_hook)
            layer.register_full_backward_hook(backward_hook)

if __name__ == "__main__":
    print("\n" + "=" * 70)
    print("          PyTorch Hook 激活值与梯度显存探针实验室")
    print("=" * 70)

    model = MemoryProbeModel()
    register_diagnostic_hooks(model)

    dummy_input = torch.randn(64, 2048, requires_grad=True)
    print("\n--- 启动前向传播 ---")
    out = model(dummy_input)

    loss = out.sum()
    print("\n--- 启动反向传播 ---")
    loss.backward()

 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
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77

import torch
import torch.nn as nn
from torch.utils.checkpoint import checkpoint
import time

class DeepBlock(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.linear1 = nn.Linear(dim, dim)
        self.act = nn.GELU()
        self.linear2 = nn.Linear(dim, dim)

    def forward(self, x):
        return self.linear2(self.act(self.linear1(x))) + x

class DeepNetwork(nn.Module):
    def __init__(self, num_layers=16, dim=2048, use_checkpoint=False):
        super().__init__()
        self.use_checkpoint = use_checkpoint
        self.blocks = nn.ModuleList([DeepBlock(dim) for _ in range(num_layers)])

    def forward(self, x):
        for block in self.blocks:
            if self.use_checkpoint and self.training:
                # 关键：使用 activation checkpointing 包装该 block 的执行！
                x = checkpoint(block, x, use_reentrant=False)
            else:
                x = block(x)
        return x

def benchmark_checkpointing():
    if not torch.cuda.is_available():
        print("需要 GPU 环境进行准确压测！")
        return

    device = torch.device("cuda:0")
    batch_size, seq_len, dim = 16, 512, 2048
    num_layers = 20

    print("\n" + "=" * 75)
    print(f"  Activation Checkpointing 压测 (深度: {num_layers} 层, 特征维: {dim})")
    print("=" * 75)

    for mode in [False, True]:
        torch.cuda.empty_cache()
        torch.cuda.reset_peak_memory_stats()

        model = DeepNetwork(num_layers=num_layers, dim=dim, use_checkpoint=mode).to(device)
        optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)

        # Warmup
        x = torch.randn(batch_size, dim, device=device, requires_grad=True)
        out = model(x).sum()
        out.backward()
        optimizer.zero_grad(set_to_none=True)

        # 记录 10 次迭代的稳定表现
        torch.cuda.synchronize()
        start_time = time.time()

        for _ in range(10):
            optimizer.zero_grad(set_to_none=True)
            inputs = torch.randn(batch_size, dim, device=device, requires_grad=True)
            outputs = model(inputs)
            loss = outputs.sum()
            loss.backward()
            optimizer.step()

        torch.cuda.synchronize()
        elapsed = (time.time() - start_time) / 10.0 * 1000  # ms/step
        peak_vram = torch.cuda.max_memory_allocated(device) / (1024 ** 2)

        mode_name = "✅ 开启 Checkpointing" if mode else "❌ 标准训练 (无 Checkpointing)"
        print(f"{mode_name:<30} | 迭代耗时: {elapsed:>7.2f} ms | 峰值显存: {peak_vram:>8.2f} MB")

if __name__ == "__main__":
    benchmark_checkpointing()

 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
45
46
47
48

import torch
import gc

def simulate_loss_leak():
    if not torch.cuda.is_available():
        return

    device = torch.device("cuda:0")
    print("\n" + "=" * 65)
    print("         生产级计算图滞留（Graph Retain）内存泄漏实测")
    print("=" * 65)

    linear = torch.nn.Linear(2048, 2048, device=device)

    # 错误用法容器
    bad_history = []

    print("--- 场景 A: 错误写法 (直接 append 包含计算图的 loss) ---")
    for step in range(5):
        x = torch.randn(512, 2048, device=device, requires_grad=True)
        loss = linear(x).sum()
        loss.backward()

        # 致命错误：把整个计算图挂在 bad_history 里！
        bad_history.append(loss)

        alloc = torch.cuda.memory_allocated(device) / (1024 ** 2)
        print(f"  Step {step + 1}: 已占用显存 = {alloc:.2f} MB (持续泄漏中...)")

    del bad_history
    gc.collect()
    torch.cuda.empty_cache()

    print("\n--- 场景 B: 正确写法 (使用 loss.item() 提取纯标量浮点数) ---")
    good_history = []
    for step in range(5):
        x = torch.randn(512, 2048, device=device, requires_grad=True)
        loss = linear(x).sum()
        loss.backward()

        # 正确做法：只保留 Python float，计算图彻底销毁！
        good_history.append(loss.item())

        alloc = torch.cuda.memory_allocated(device) / (1024 ** 2)
        print(f"  Step {step + 1}: 已占用显存 = {alloc:.2f} MB (平稳恒定)")

if __name__ == "__main__":
    simulate_loss_leak()

1
2
3
4
5
6
7
8
9

=================================================================================================
                          Ringi Autograd 显存生命周期五句诀
=================================================================================================
前向算完织成图，梯度回传找 grad_fn；
激活扣留死不放，直到反向才退火；
就地修改动版本，报错崩溃莫惊慌；
推理莫忘 inference，冻结版本性能爽；
显存爆炸莫要怕，重算省下七成粮！
=================================================================================================