PyTorch 内部结构与高级优化：从 autograd、torch.compile、FSDP 到 Triton

PyTorch 内部结构

ATen 库

PyTorch 的核心是 ATen(A Tensor library)。这是一个基于 C++ 的张量运算库，是所有 PyTorch 运算的底层实现。

Python API (torch.*)
    ↓
TorchDispatch / Dispatcher
    ↓
ATen (C++ tensor ops)
    ↓
CUDA / CPU / MPS backends

ATen 的主要组成部分：

Tensor：多维数组，保存 storage、dtype、device、stride 等信息
Storage：实际的内存块(可在张量之间共享)
Dispatcher：把运算路由到合适的后端

import torch

x = torch.randn(3, 4)
print(x.storage())          # 实际内存块
print(x.stride())           # (4, 1) - 行主序
print(x.storage_offset())   # 0

# View 与原张量共享 storage
y = x.view(2, 6)
print(x.storage().data_ptr() == y.storage().data_ptr())  # True

TorchDispatch

TorchDispatch 是在 Python 层面拦截 PyTorch 运算的机制,用于实现自定义张量类型。

import torch
from torch.utils._pytree import tree_map

class LoggingTensor(torch.Tensor):
    @staticmethod
    def __new__(cls, elem):
        return torch.Tensor._make_subclass(cls, elem)

    @classmethod
    def __torch_dispatch__(cls, func, types, args=(), kwargs=None):
        print(f"Calling: {func.__name__}")
        kwargs = kwargs or {}
        return func(*args, **kwargs)

x = LoggingTensor(torch.randn(3, 3))
y = x + x  # 输出: Calling: add.Tensor

Autograd 引擎

计算图(DAG)

PyTorch autograd 使用动态计算图(Dynamic Computational Graph)。每当运算执行时，都会构建出一个 DAG(有向无环图)。

import torch

x = torch.tensor(2.0, requires_grad=True)  # leaf tensor
y = x ** 2          # non-leaf, grad_fn=PowBackward0
z = y * 3           # non-leaf, grad_fn=MulBackward0

print(x.is_leaf)    # True
print(y.is_leaf)    # False
print(z.grad_fn)    # MulBackward0

z.backward()
print(x.grad)       # dz/dx = 3 * 2x = 12.0

Leaf tensor 与 Non-leaf tensor 的区别：

Leaf tensor：requires_grad=True 且由用户直接创建的张量。gradient 会累积在 .grad 中
Non-leaf tensor：运算结果生成的张量。默认情况下 .grad 为 None(需调用 .retain_grad())

梯度累积机制

import torch

model = torch.nn.Linear(10, 1)
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

# 梯度累积(accumulation)
ACCUMULATION_STEPS = 4
for i, (x, y) in enumerate(dataloader):
    output = model(x)
    loss = criterion(output, y) / ACCUMULATION_STEPS
    loss.backward()  # 梯度累积

    if (i + 1) % ACCUMULATION_STEPS == 0:
        optimizer.step()
        optimizer.zero_grad()  # 梯度清零

retain_graph 与 create_graph

x = torch.tensor(3.0, requires_grad=True)
y = x ** 3

# 高阶导数: create_graph=True
grad_1 = torch.autograd.grad(y, x, create_graph=True)[0]
grad_2 = torch.autograd.grad(grad_1, x)[0]

print(grad_1)  # 3x^2 = 27.0
print(grad_2)  # 6x = 18.0

自定义算子实现

torch.autograd.Function

用于定义自定义的 forward/backward。

import torch

class SigmoidFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x):
        # 把 backward 需要的值保存到 ctx
        output = 1 / (1 + torch.exp(-x))
        ctx.save_for_backward(output)
        return output

    @staticmethod
    def backward(ctx, grad_output):
        (output,) = ctx.saved_tensors
        # sigmoid 的导数: sigma(x) * (1 - sigma(x))
        grad_input = grad_output * output * (1 - output)
        return grad_input

# 使用
x = torch.randn(4, requires_grad=True)
y = SigmoidFunction.apply(x)
y.sum().backward()
print(x.grad)

torch.library API(注册自定义算子)

import torch
from torch.library import Library, impl

my_lib = Library("my_ops", "DEF")
my_lib.define("relu_squared(Tensor x) -> Tensor")

@impl(my_lib, "relu_squared", "CPU")
def relu_squared_cpu(x):
    return torch.relu(x) ** 2

@impl(my_lib, "relu_squared", "CUDA")
def relu_squared_cuda(x):
    return torch.relu(x) ** 2

# 使用自定义算子
x = torch.randn(5)
result = torch.ops.my_ops.relu_squared(x)

自定义 CUDA 内核(Triton)

import triton
import triton.language as tl
import torch

@triton.jit
def relu_squared_kernel(
    x_ptr, out_ptr,
    n_elements,
    BLOCK_SIZE: tl.constexpr,
):
    pid = tl.program_id(0)
    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = offsets < n_elements

    x = tl.load(x_ptr + offsets, mask=mask)
    relu_x = tl.where(x > 0, x, 0.0)
    out = relu_x * relu_x

    tl.store(out_ptr + offsets, out, mask=mask)

def relu_squared_triton(x: torch.Tensor):
    out = torch.empty_like(x)
    n_elements = x.numel()
    BLOCK_SIZE = 1024
    grid = (triton.cdiv(n_elements, BLOCK_SIZE),)
    relu_squared_kernel[grid](x, out, n_elements, BLOCK_SIZE)
    return out

torch.compile()

Dynamo 与图捕获

torch.compile() 会分析 Python 字节码来提取计算图。

import torch

def model_forward(x, weight):
    x = torch.nn.functional.relu(x @ weight)
    return x.sum()

# 编译: fullgraph=True 时不允许 graph break
compiled_fn = torch.compile(model_forward, fullgraph=True, backend="inductor")

x = torch.randn(128, 256, device="cuda")
w = torch.randn(256, 512, device="cuda")
out = compiled_fn(x, w)

Graph break 的触发条件：

Python 控制流(在 if/for 中使用张量值)
调用外部库(如 numpy)
不受支持的 Python 模式

import torch._dynamo
torch._dynamo.config.verbose = True  # 调试 graph break

# 若要允许 graph break，用 fullgraph=False(默认值)
compiled = torch.compile(model, backend="inductor")

AOTAutograd 与 TorchInductor

torch.compile() 流水线:
Python 代码 → Dynamo (图提取)
           → AOTAutograd (合并 forward + backward)
           → TorchInductor (内核生成)
           → Triton / C++ 代码

import torch
import torch.nn as nn

class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(512, 1024)
        self.fc2 = nn.Linear(1024, 10)

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

model = SimpleModel().cuda()

# mode 选项: "default"、"reduce-overhead"、"max-autotune"
compiled_model = torch.compile(model, mode="max-autotune")

x = torch.randn(32, 512, device="cuda")
out = compiled_model(x)

内存优化

Gradient Checkpointing

在前向传播过程中不保存中间激活值，在反向传播时重新计算，以此节省内存。

import torch
import torch.utils.checkpoint as checkpoint
import torch.nn as nn

class CheckpointedBlock(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(dim, dim * 4),
            nn.GELU(),
            nn.Linear(dim * 4, dim),
        )

    def forward(self, x):
        # checkpoint: 通过重新执行 forward 来节省内存
        return checkpoint.checkpoint(self.layers, x, use_reentrant=False)

model = nn.Sequential(*[CheckpointedBlock(512) for _ in range(24)]).cuda()
x = torch.randn(32, 512, device="cuda", requires_grad=True)
out = model(x)
out.sum().backward()

AMP(自动混合精度)

import torch
from torch.cuda.amp import autocast, GradScaler

model = SimpleModel().cuda()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
scaler = GradScaler()  # 防止 FP16 下溢

for x, y in dataloader:
    x, y = x.cuda(), y.cuda()
    optimizer.zero_grad()

    # autocast: 根据运算自动应用 FP16/BF16
    with autocast(dtype=torch.float16):
        output = model(x)
        loss = criterion(output, y)

    # scaler: 对梯度做缩放以防止下溢
    scaler.scale(loss).backward()
    scaler.unscale_(optimizer)
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    scaler.step(optimizer)
    scaler.update()

内存分析器

import torch
from torch.profiler import profile, ProfilerActivity, record_function

with profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
    profile_memory=True,
    record_shapes=True,
) as prof:
    with record_function("model_inference"):
        output = model(x)

print(prof.key_averages().table(
    sort_by="cuda_memory_usage", row_limit=10
))
prof.export_chrome_trace("trace.json")

Activation Offloading

# 通过 CPU offloading 节省 GPU 内存
def offload_checkpoint(module, x):
    """将激活值卸载到 CPU，反向传播时再取回 GPU"""
    def forward_and_save(*inputs):
        output = module(*inputs)
        return output

    return checkpoint.checkpoint(forward_and_save, x, use_reentrant=False)

分布式训练

DDP(DistributedDataParallel)

import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
import os

def setup(rank, world_size):
    os.environ["MASTER_ADDR"] = "localhost"
    os.environ["MASTER_PORT"] = "12355"
    dist.init_process_group("nccl", rank=rank, world_size=world_size)

def cleanup():
    dist.destroy_process_group()

def train(rank, world_size):
    setup(rank, world_size)
    torch.cuda.set_device(rank)

    model = SimpleModel().cuda(rank)
    ddp_model = DDP(model, device_ids=[rank])

    optimizer = torch.optim.Adam(ddp_model.parameters())

    for x, y in dataloader:
        x, y = x.cuda(rank), y.cuda(rank)
        output = ddp_model(x)
        loss = criterion(output, y)
        loss.backward()     # gradient all-reduce 自动执行
        optimizer.step()
        optimizer.zero_grad()

    cleanup()

FSDP(Fully Sharded Data Parallel)

FSDP 会把参数、gradient、optimizer state 分片到所有 GPU 上。

import torch
from torch.distributed.fsdp import (
    FullyShardedDataParallel as FSDP,
    MixedPrecision,
    ShardingStrategy,
)
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
import functools

# 以 Transformer 块为单位应用 FSDP
auto_wrap_policy = functools.partial(
    transformer_auto_wrap_policy,
    transformer_layer_cls={nn.TransformerEncoderLayer},
)

# 混合精度设置
mp_policy = MixedPrecision(
    param_dtype=torch.bfloat16,
    reduce_dtype=torch.float32,
    buffer_dtype=torch.bfloat16,
)

model = LargeTransformer().cuda()
fsdp_model = FSDP(
    model,
    auto_wrap_policy=auto_wrap_policy,
    mixed_precision=mp_policy,
    sharding_strategy=ShardingStrategy.FULL_SHARD,  # 等同于 ZeRO-3
)

optimizer = torch.optim.AdamW(fsdp_model.parameters(), lr=1e-4)

FSDP 与 DDP 的内存对比：

DDP 是每个 GPU 都保留完整的模型副本。FSDP 则把参数按 world_size 切分，把每个 GPU 的内存占用降到 1/N。

DeepSpeed ZeRO 集成

# deepspeed_config.json
# {
#   "zero_optimization": {"stage": 3},
#   "fp16": {"enabled": true},
#   "gradient_accumulation_steps": 4
# }

import deepspeed

model_engine, optimizer, _, _ = deepspeed.initialize(
    model=model,
    model_parameters=model.parameters(),
    config="deepspeed_config.json"
)

for x, y in dataloader:
    output = model_engine(x)
    loss = criterion(output, y)
    model_engine.backward(loss)
    model_engine.step()

推理优化

torch.export() 与 ONNX

import torch
from torch.export import export

model = SimpleModel().eval()
x = torch.randn(1, 512)

# torch.export: 提取静态计算图
exported = export(model, (x,))
print(exported.graph)

# 导出 ONNX
torch.onnx.export(
    model, x,
    "model.onnx",
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={"input": {0: "batch_size"}},
    opset_version=17,
)

Quantization-Aware Training(QAT)

import torch
from torch.quantization import get_default_qat_qconfig, prepare_qat, convert

model = SimpleModel()
model.qconfig = get_default_qat_qconfig("fbgemm")

# QAT 准备: 插入 fake quantization
model_prepared = prepare_qat(model.train())

# 训练方式与常规训练相同
for x, y in dataloader:
    output = model_prepared(x)
    loss = criterion(output, y)
    loss.backward()
    optimizer.step()

# 转换为 INT8 模型
model_int8 = convert(model_prepared.eval())

调试工具

torch.profiler

import torch
from torch.profiler import profile, ProfilerActivity, schedule

with profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
    schedule=schedule(wait=1, warmup=1, active=3, repeat=2),
    on_trace_ready=torch.profiler.tensorboard_trace_handler("./log"),
    record_shapes=True,
    profile_memory=True,
    with_stack=True,
) as prof:
    for step, (x, y) in enumerate(dataloader):
        output = model(x.cuda())
        loss = criterion(output, y.cuda())
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()
        prof.step()

Anomaly Detection

# autograd 异常检测(追踪 NaN/Inf 梯度)
with torch.autograd.detect_anomaly():
    output = model(x)
    loss = output.sum()
    loss.backward()  # 发生 NaN 时输出堆栈跟踪

grad_fn 追踪

def trace_grad_fn(tensor, depth=0):
    if tensor.grad_fn is None:
        print("  " * depth + f"Leaf: {tensor.shape}")
        return
    print("  " * depth + f"{tensor.grad_fn.__class__.__name__}: {tensor.shape}")
    for inp, _ in tensor.grad_fn.next_functions:
        if inp is not None:
            trace_grad_fn(inp.variable if hasattr(inp, 'variable') else inp, depth + 1)

x = torch.randn(3, requires_grad=True)
y = torch.randn(3, requires_grad=True)
z = (x * y).sum()
trace_grad_fn(z)

测验

Q1. PyTorch autograd 中 leaf tensor 与 non-leaf tensor 的区别，以及梯度累积的方式是什么？

答案：leaf tensor 是用户直接创建、且 requires_grad=True 的张量。gradient 会累积在 .grad 属性中。

说明：non-leaf tensor 是运算结果生成的张量，默认不保存 gradient(为了节省内存)。调用 retain_grad() 后，也可以保留 non-leaf tensor 的 gradient。梯度累积的运作方式是：在不调用 optimizer.zero_grad() 的情况下多次调用 backward()，梯度会不断累加到 .grad 上。利用这一点，可以借助 gradient accumulation 虚拟地扩大 batch size。

Q2. torch.compile() 的 Dynamo 追踪 Python 字节码的方式，以及 graph break 的触发条件是什么？

答案：Dynamo 通过拦截 Python 帧求值来分析字节码，遇到不受支持的模式时会触发 graph break。

说明：Dynamo 使用 CPython 的 PEP 523 帧求值 API，对 Python 字节码做符号化追踪。依赖张量值的控制流(例如 if x.sum() > 0:)、调用外部库、C 扩展等情况都会触发 graph break。发生 graph break 时，Dynamo 会编译到该断点为止的图，其余部分则以普通 Python 方式执行。设置 fullgraph=True 后，graph break 会被当作错误处理。

Q3. FSDP 为什么比 DDP 更节省内存(从参数分片的角度)？

答案：FSDP 会把参数、gradient、optimizer state 全部分散到 world_size 个 GPU 上，把每个 GPU 的内存占用降到约 1/N。

说明：DDP 是每个 GPU 都复制并保留完整的模型参数。一个 10B 参数的模型，按 FP32 计算每个 GPU 大约需要 40GB。FSDP(ZeRO-3 策略)只在 forward/backward 时通过 all-gather 收集所需的参数，用完后立即释放。在 8 个 GPU 的环境下，内存占用可降至约 1/8，从而能训练单个 GPU 装不下的超大模型。

Q4. Gradient checkpointing 中重新执行 forward pass 的权衡是什么？

答案：这是一种运算-内存的权衡：把内存占用降到 O(sqrt(N))，代价是反向传播时间增加约 33%。

说明：常规的反向传播会保存所有 forward 激活值，因此需要 O(N) 的内存。Gradient checkpointing 只保存检查点边界处的激活值，反向传播时重新执行该区间的 forward。在 Transformer 中按层设置检查点时，所需内存只与层数的平方根成正比，而不是与层数本身成正比。重新计算的开销会让整体训练时间增加约 30-40%，但由于可以扩大 batch size，实际吞吐量反而可能得到改善。

Q5. AMP 中 GradScaler 是如何防止下溢的？

答案：GradScaler 会给 loss 乘上一个较大的缩放值，把 gradient 维持在 FP16 可表示的范围内，并在 optimizer 更新之前把缩放逆向还原。

说明：FP16 的最小值约为 6e-5，过小的 gradient 会下溢为 0。GradScaler 给 loss 乘上一个 scale factor(初始值如 65536)后，gradient 会按相同倍数放大，从而在 FP16 下也能被表示。在 scaler.unscale_(optimizer) 这一步，会把 gradient 还原回原始大小。一旦出现 Inf/NaN，会自动缩小 scale 并跳过该 step。BF16 与 FP32 拥有相同的指数范围，因此不需要 GradScaler。