PyTorch

PyTorch 是一个开源的深度学习框架，由 Facebook AI Research 开发，广泛应用于学术研究和工业界。

安装

自动安装

访问官网选择配置: https://pytorch.org/get-started/locally/

# CPU 版本
pip install torch torchvision torchaudio

# GPU 版本 (CUDA 12.4)
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124

手动安装 GPU 版本

下载地址: https://download.pytorch.org/whl/torch/ 需要下载三个包（注意 Python 版本和 CUDA 版本）:

torch-2.4.1+cu124-cp312-cp312-win_amd64.whl
torchvision-0.19.1+cu124-cp312-cp312-win_amd64.whl
torchaudio-2.4.1+cu124-cp312-cp312-win_amd64.whl

pip install torch-2.4.1+cu124-cp312-cp312-win_amd64.whl
pip install torchvision-0.19.1+cu124-cp312-cp312-win_amd64.whl
pip install torchaudio-2.4.1+cu124-cp312-cp312-win_amd64.whl

验证安装

import torch

print(torch.__version__)
print(torch.cuda.is_available())  # True 表示 GPU 可用
print(torch.cuda.get_device_name(0))  # 显示 GPU 名称

基础概念

Tensor 张量

import torch

# 创建张量
x = torch.tensor([1, 2, 3])
y = torch.zeros(2, 3)
z = torch.randn(2, 3)

# 张量运算
a = torch.tensor([1.0, 2.0])
b = torch.tensor([3.0, 4.0])
c = a + b
d = torch.matmul(a, b)  # 矩阵乘法

# GPU 加速
if torch.cuda.is_available():
    x = x.cuda()  # 转移到 GPU
    # 或者
    x = x.to('cuda')

自动求导

# 需要梯度的张量
x = torch.tensor([1.0, 2.0], requires_grad=True)
y = x ** 2
z = y.sum()

# 反向传播
z.backward()
print(x.grad)  # 梯度: [2.0, 4.0]

构建神经网络

简单模型

import torch.nn as nn

class SimpleNet(nn.Module):
    def __init__(self):
        super(SimpleNet, self).__init__()
        self.fc1 = nn.Linear(784, 128)
        self.fc2 = nn.Linear(128, 10)
        self.relu = nn.ReLU()

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

model = SimpleNet()

卷积神经网络

class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.fc1 = nn.Linear(9216, 128)
        self.fc2 = nn.Linear(128, 10)
        self.relu = nn.ReLU()
        self.maxpool = nn.MaxPool2d(2)

    def forward(self, x):
        x = self.conv1(x)
        x = self.relu(x)
        x = self.maxpool(x)
        x = self.conv2(x)
        x = self.relu(x)
        x = self.maxpool(x)
        x = x.view(x.size(0), -1)
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        return x

训练模型

import torch.optim as optim

# 准备数据
train_loader = torch.utils.data.DataLoader(dataset, batch_size=64)

# 定义模型、损失函数、优化器
model = SimpleNet()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 训练循环
for epoch in range(10):
    for batch_idx, (data, target) in enumerate(train_loader):
        # 前向传播
        output = model(data)
        loss = criterion(output, target)

        # 反向传播
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if batch_idx % 100 == 0:
            print(f'Epoch: {epoch}, Loss: {loss.item()}')

先测量，再优化训练

GPU 利用率高，不代表训练高效。NVIDIA 的 GPU 利用率衡量的是采样时间里有 kernel 执行的比例，不是达到峰值算力的比例。Databricks 的一个 Qwen 微调案例在基线平均利用率已达 98% 时，仍把训练时间缩短到约六分之一。这是特定任务的案例，不是你的模型也能提速六倍的承诺。

采集短窗口 CPU／GPU trace

用 torch.profiler 查看算子数量、CPU 调度、设备执行和 kernel 之间的空隙。下面是一个独立的小例子，需要已安装 PyTorch；有 CUDA 时记录 GPU，否则只记录 CPU。它不下载模型或数据，会在当前目录写出一份 trace。

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

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = nn.Sequential(nn.Linear(128, 256), nn.ReLU(), nn.Linear(256, 16)).to(device)
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
x = torch.randn(64, 128, device=device)
y = torch.randn(64, 16, device=device)
activities = [ProfilerActivity.CPU]
if device.type == "cuda":
    activities.append(ProfilerActivity.CUDA)


def save_trace(prof):
    print(prof.key_averages().table(sort_by="self_cpu_time_total", row_limit=10))
    prof.export_chrome_trace(f"training-trace-{prof.step_num}.json")


with profile(
    activities=activities,
    schedule=schedule(wait=1, warmup=1, active=3, repeat=1),
    on_trace_ready=save_trace,
    record_shapes=True,
) as prof:
    for _ in range(5):
        with record_function("train_step"):
            optimizer.zero_grad(set_to_none=True)
            loss = nn.functional.mse_loss(model(x), y)
            loss.backward()
            optimizer.step()
        prof.step()

这个 schedule 跳过一步、预热 profiler 一步，再记录三步。真实任务要调整采样窗口，覆盖有代表性的阶段；profiler 预热不等于编译和模型启动已经完成。示例输入已在设备上，因此没有测量数据加载和主机到设备的复制。接入真实循环时，也应为这些阶段设置采样范围。分布式运行时，每个 rank 使用不同的 trace 文件名。 Trace、shape 记录和可选的内存分析都会增加开销。用 trace 定位瓶颈，再关闭 profiler 单独测吞吐。CUDA 是异步执行的：测墙钟耗时时，在计时开始前和结束前同步相关设备，而不是每轮内部都强制同步。不要把嵌套 CPU 耗时与重叠 GPU 耗时相加后，当作实际总时长。

按瓶颈选优化手段

Modal 的 GPU 术语手册把工作分为计算、内存搬运和其他开销。先把它们当作假设，再通过时间线或设备 profiler 验证：

观察到的现象	排查方向	候选改动
大量微小 kernel、CPU 调度繁忙、启动之间有空隙	启动／调度开销	使用受支持的融合内核；评估该任务是否适合 `torch.compile` 或 CUDA Graphs
大 kernel 的时间主要花在搬数据上	内存带宽与多余中间张量	减少复制和中间结果，采用受支持的融合实现
大矩阵运算占主导	算力吞吐、dtype、shape 与硬件支持	调整 batch／shape 和精度，同时检查数值质量
等待输入、集合通信或不同 rank 进度不齐	数据流水线与分布式通信	解决实际等待点，不要默认是模型 kernel 太慢

Databricks 案例中，已有的 flash-linear-attention 和 causal-conv1d 内核消除了大量小算子，减少 padding 又带来额外收益。这不代表这些包适合所有模型，也不代表编译和 CUDA Graphs 普遍不适合分布式训练。启用快路径前，核对模型／后端支持、GPU 架构、PyTorch／CUDA 兼容性，以及前向和反向是否都已实现。固定可用的环境，并在 trace 中确认目标 kernel 真的出现：安装成功不等于正在使用。用适当的数值容差，与参考路径对照 loss 和梯度。

打包序列时保留训练语义

序列打包能减少 padding，但直接拼接样本可能改变训练语义。只重置 position_ids 不是通用的样本隔离机制。 模型和注意力后端必须支持相应的边界元数据或 mask；混合循环层还需要在样本边界隔离状态。固定版本的 Databricks 实现同时使用了重置后的 position_ids、供完整注意力使用的累计序列长度、供 gated-delta 层使用的分段 ID，以及防止移位后的目标跨样本计算 loss 的掩码；response-only SFT 还会屏蔽 prompt token。这些都是后端相关输入，不能原样照搬到任意模型。测速度之前，先对照打包与未打包样本，覆盖边界和截断情况。记录非 padding tokens/s、padding 占比、有效 batch size、序列长度分布与验证集质量。打包改变每次优化器更新的 token 数时，还要考虑 batch 与优化调度的变化。比较等量训练工作，并把冷启动／编译耗时与稳态吞吐分开报告。

数据加载

Dataset 和 DataLoader

from torch.utils.data import Dataset, DataLoader

class CustomDataset(Dataset):
    def __init__(self, data, labels):
        self.data = data
        self.labels = labels

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        return self.data[idx], self.labels[idx]

dataset = CustomDataset(data, labels)
dataloader = DataLoader(dataset, batch_size=32, shuffle=True)

图像数据增强

from torchvision import transforms

transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                        std=[0.229, 0.224, 0.225])
])

保存和加载模型

# 保存模型
torch.save(model.state_dict(), 'model.pth')

# 加载模型
model = SimpleNet()
model.load_state_dict(torch.load('model.pth'))
model.eval()

# 保存完整模型
torch.save(model, 'model_complete.pth')
model = torch.load('model_complete.pth')

Jupyter Notebook 配置

创建 PyTorch 专用内核:

# 创建虚拟环境
python -m venv pytorch_env

# 激活环境
source pytorch_env/bin/activate  # Linux/Mac
pytorch_env\Scripts\activate     # Windows

# 安装 PyTorch 和 Jupyter
pip install torch torchvision jupyter ipykernel

# 创建内核
python -m ipykernel install --user --name=pytorch --display-name "Python (PyTorch)"

参考: https://blog.csdn.net/ccaoshangfei/article/details/126521809

常用技巧

设置随机种子

torch.manual_seed(42)
if torch.cuda.is_available():
    torch.cuda.manual_seed(42)

查看模型结构

from torchsummary import summary

summary(model, input_size=(1, 28, 28))

冻结部分层

# 冻结前两层
for param in model.parameters():
    param.requires_grad = False

for param in model.fc2.parameters():
    param.requires_grad = True

参考资料

meng shao：Qwen 训练优化案例 — 2026-10-15 的 Folo 线索
Databricks：Reducing Qwen training time by 6x on AI Runtime — 特定任务中的性能分析、融合内核与 padding 优化
Databricks：固定版本训练实现 — 打包元数据与 loss 边界
PyTorch：Profiler recipe — 采样活动、分步 trace 与开销
NVIDIA：nvidia-smi 文档 — 利用率指标的定义
meng shao：推荐 Modal GPU 术语手册 — 2026-10-15 的 Folo 线索；保留术语手册，不收录融资新闻
Modal：GPU 术语手册——性能瓶颈 — 计算、内存和其他开销

学习资源

官方文档: https://pytorch.org/docs/
官方教程: https://pytorch.org/tutorials/
中文文档: https://pytorch-cn.readthedocs.io/
GitHub: https://github.com/pytorch/pytorch

后端

前端

客户端

算法

PyTorch

PyTorch

安装

自动安装

手动安装 GPU 版本

验证安装

基础概念

Tensor 张量

自动求导

构建神经网络

简单模型

卷积神经网络

训练模型

先测量，再优化训练

采集短窗口 CPU／GPU trace

按瓶颈选优化手段

打包序列时保留训练语义

数据加载

Dataset 和 DataLoader

图像数据增强

保存和加载模型

Jupyter Notebook 配置

常用技巧

设置随机种子

查看模型结构

冻结部分层

参考资料

学习资源

​PyTorch

​安装

​自动安装

​手动安装 GPU 版本

​验证安装

​基础概念

​Tensor 张量

​自动求导

​构建神经网络

​简单模型

​卷积神经网络

​训练模型

​先测量，再优化训练

​采集短窗口 CPU／GPU trace

​按瓶颈选优化手段

​打包序列时保留训练语义

​数据加载

​Dataset 和 DataLoader

​图像数据增强

​保存和加载模型

​Jupyter Notebook 配置

​常用技巧

​设置随机种子

​查看模型结构

​冻结部分层

​参考资料

​学习资源

PyTorch

安装

自动安装

手动安装 GPU 版本

验证安装

基础概念

Tensor 张量

自动求导

构建神经网络

简单模型

卷积神经网络

训练模型

先测量，再优化训练

采集短窗口 CPU／GPU trace

按瓶颈选优化手段

打包序列时保留训练语义

数据加载

Dataset 和 DataLoader

图像数据增强

保存和加载模型

Jupyter Notebook 配置

常用技巧

设置随机种子

查看模型结构

冻结部分层

参考资料

学习资源