pysheeet PyTorch 速查指南：从张量操作到 GPU 分布式训练（NCCL / torchrun / Slurm 全流程）

文档教程开发工具【免费下载链接】pysheeetPython Cheat Sheet项目地址https://gitcode.com/gh_mirrors/py/pysheeet点击查看免费下载PyTorch 是 Meta AI 开源的深度学习框架凭借动态计算图、自动微分与强大的 GPU 加速能力成为大模型训练与推理的事实标准底座。本文以 pysheeet 仓库中 docs/notes/llm/pytorch.rst 速查文档为骨架完整覆盖张量创建与变换、自动梯度、神经网络模块、训练循环、模型存取、GPU 内存管理以及基于 NCCL 的多卡/多节点分布式训练并结合仓库内的 NCCL 广播基准、vLLM 离线基准、Megatron 启动脚本 等真实源码深入讲解底层原理。读完本文你将能独立完成从在单卡上写一个可训练的神经网络到通过 torchrun / Slurm 在 GPU 集群上启动分布式训练的完整闭环。PyTorch 与动态计算图PyTorch 与 TensorFlow 等静态图框架的核心差异在于动态计算图每次前向传播都会即时构建计算图模型结构可以在运行期任意修改、打印中间结果、用 Python 调试器单步跟踪。这种灵活性使原型验证、实验对比与 Debug 都更加直观尤其适合需要频繁调整架构的 LLM 与大模型研究场景。在 pysheeet 的 LLM 笔记体系中PyTorch 是贯穿始终的基础设施LLM 笔记索引 明确将其作为模型开发、分布式训练、vLLM/SGLang 推理服务的共同基石。环境检查CUDA 可用性与设备配置在运行任何 GPU 加速计算之前先确认 CUDA 驱动、NCCL 版本与 GPU 硬件是否就绪 import torch torch.cuda.is_available() True torch.cuda.nccl.version() (2, 26, 2) torch.cuda.device_count() 2 torch.cuda.get_device_name(0) NVIDIA GeForce RTX 3090 torch.cuda.set_device(0) torch.cuda.current_device() 0torch.cuda.is_available()返回布尔值是能否使用 CUDA的总开关训练脚本中常用它做设备回退见下文移动张量。torch.cuda.nccl.version()返回 NCCL 版本元组NCCL 是 PyTorch 多卡通信的后端详见分布式训练小节。torch.cuda.device_count()/get_device_name(i)探测 GPU 数量与型号用于动态决定并行规模。set_device(0)/current_device()切换并查询当前默认 CUDA 设备。仓库中的实际项目也遵循同样的检查思路。例如 src/nccl/broadcast.py 在初始化时用os.environ[LOCAL_RANK]定位每进程对应的 CUDA 设备并调用torch.cuda.set_device(self.device)显式绑定设备这是分布式场景下避免设备错乱的标准做法。查看张量所在设备张量自带device属性可用于确认数据究竟在 CPU 还是 GPU 上 tensor torch.tensor([0, 1, 2, 3, 4, 5]) tensor.device device(typecpu) tensor tensor.to(cuda) tensor.device device(typecuda, index0)device(typecuda, index0)表示该张量位于第 0 号 GPU。跨设备运算如 CPU 张量与 GPU 张量相加会直接报错因此训练循环中要养成把所有输入和目标统一to(device)的习惯见训练循环。张量创建与基础属性张量Tensor是 PyTorch 的核心数据结构功能类似 NumPy 数组但额外支持 GPU 加速与自动梯度。创建方式极为多样 x torch.tensor([0, 1, 2, 3, 4, 5]) # 从 Python 列表 x torch.empty(2, 2) # 未初始化内存速度快值未定义 x torch.rand(2, 2) # 均匀分布 [0, 1) x torch.randn(2, 2) # 标准正态分布 x torch.ones(2, 2) x torch.zeros(2, 2) x torch.arange(0, 10, 2) # 等差序列0, 2, 4, 6, 8 x torch.linspace(0, 1, 5) # 等分区间0, 0.25, 0.5, 0.75, 1.0也可以直接在指定设备上创建 device torch.device(cuda:0) x torch.tensor([1, 2, 3, 4, 5], devicedevice) x tensor([1, 2, 3, 4, 5], devicecuda:0)张量属性速查 x torch.randn(2, 3, 4) x.shape # 形状等价于 x.size() torch.Size([2, 3, 4]) x.size() torch.Size([2, 3, 4]) x.dtype # 数据类型默认 float32 torch.float32 x.device device(typecpu) x.numel() # 元素总数 2*3*4 24shape与size()等价numel()常用于计算参数量或显存占用元素数 × 字节数例如 src/nccl/broadcast.py 正是用self.tensor.numel() * self.tensor.element_size() / 1e9来计算 512 MB 广播负载的字节数。张量内存布局Contiguous、view 与 reshape连续性Contiguous部分算子要求张量存储在连续内存块中。transpose、permute等操作只是交换了 stride步长并不实际搬移数据因此会产生非连续视图 x torch.randn(2, 3) x.is_contiguous() True y x.transpose(0, 1) y.is_contiguous() False z y.contiguous() # 复制出一份连续内存 z.is_contiguous() True非连续张量在参与某些 kernel 运算时性能较差甚至报错此时调用contiguous()会生成一份连续拷贝。view 与 reshape 的区别这是新手最容易踩坑的点view()只改变解释方式shape/stride共享底层内存零拷贝但要求张量连续否则报RuntimeError。reshape()在张量连续时等价于 view不连续时自动复制为连续后再变形总能成功。 x torch.randn(2, 3, 4) x.is_contiguous() True x.view(2, 12).shape torch.Size([2, 12]) y x.transpose(0, 1) y.is_contiguous() False y.view(3, 8) RuntimeError: view size is not compatible with input tensors size and stride y.reshape(3, 8).shape # 自动复制后成功 torch.Size([3, 8]) y.contiguous().view(3, 8).shape # 手动复制后 view 也可用 torch.Size([3, 8])实践建议确定张量连续时优先用view()零拷贝、性能更好否则用reshape()。若在意内存可先contiguous()再view()但要注意这是一次显式拷贝。常用重塑操作x torch.randn(2, 3, 4) x.reshape(6, 4) # 合并前两维 x.flatten() # 展平为一维24 x.squeeze() # 去掉所有大小为 1 的维度 x.unsqueeze(0) # 在 0 维插入大小为 1 的维度如为 batch 加一维squeeze/unsqueeze在视觉模型与 Transformer 输入变换中极其常见例如把单样本(seq, dim)扩成(1, seq, dim)以匹配 batch 维度。张量在设备间移动训练时模型与数据常分布在 CPU/GPU 上需要显式搬移x torch.randn(2, 3) x_gpu x.to(cuda) # 或 x.cuda() x_cpu x_gpu.to(cpu) # 或 x_gpu.cpu() device torch.device(cuda if torch.cuda.is_available() else cpu) x x.to(device) # 自动回退到 CPU 的惯用写法to()是通用接口可指定device、dtypecuda()/cpu()是便捷简写。注意GPU 张量不能直接转 NumPy必须先搬回 CPU见NumPy 转换。张量运算算术与矩阵运算PyTorch 支持逐元素运算与矩阵运算且大部分算子同时具备函数形式与方法形式并原生支持广播broadcasting x torch.rand(2, 2, device0) y torch.rand(2, 2, device0) x y # 等价于 torch.add(x, y) torch.add(x, y) x - y x * y # 逐元素乘法 x / y x ** 2 x y # 矩阵乘法等价于 torch.matmul(x, y) torch.matmul(x, y)与torch.matmul支持高维批量矩阵乘法是神经网络前向传播y x W b的基础。原地操作In-place以_结尾的算子直接修改张量本身、不产生新对象能省内存x torch.rand(2, 2, device0) y torch.rand(2, 2, device0) y.add_(x) y.sub_(x) y.mul_(x) y.div_(x) y.zero_() # 清零 y.fill_(5) # 填充常量警告原地操作会破坏 autograd 的计算图覆盖掉保存的梯度所需的值导致反向传播错误因此对requires_gradTrue的张量要格外小心。优化器中标准的optimizer.zero_grad()本质上就是清空梯度见优化器小节。转置与维度交换 x torch.randn(2, 6, 2, device0) x.shape torch.Size([2, 6, 2]) y x.transpose(1, 2) # 交换第 1、2 维 y.shape torch.Size([2, 2, 6]) x.T.shape # 仅对 2D 张量等价于全部维度反转 torch.Size([2, 2, 6]) x.permute(2, 0, 1).shape # 任意重排维度 torch.Size([2, 2, 6])transpose(a, b)交换两个维度permute(...)可任意重排所有维度T是 2D 专用快捷方式。注意力机制中(batch, seq, head, dim)与(batch, head, seq, dim)的互转即依赖这类操作。批量矩阵乘法 x torch.randn(1, 2, 3, 4, device0) x x.transpose(2, 3) # 每个 (3,4) 子矩阵与其转置相乘 torch.bmm(x, x.transpose(2, 3)) # 显式 batch 矩阵乘 torch.einsum(bijk,bikl-bijl, x, x.transpose(2, 3)) # Einstein 记号bmm要求输入恰好为 3Deinsum最灵活可直接表达注意力打分Q K^T等复杂张量收缩是 Transformer 代码中的高频 API。聚合归约操作 x torch.randn(2, 3, 4) x.sum() # 全量求和 tensor(5.2341) x.sum(dim0).shape # 沿第 0 维归约 torch.Size([3, 4]) x.sum(dim-1).shape # 沿最后一维归约 torch.Size([2, 3]) x.mean() # 均值 tensor(0.2181) x.std() # 标准差 tensor(1.0234) x.max() # 最大值 tensor(2.3456) x.min() tensor(-1.8765) x.argmax() # 全局最大值的展平索引 tensor(7) x.argmin() tensor(15)带dim的归约返回(values, indices)两个部分 x.max(dim1) torch.return_types.max( valuestensor([[1.2345, 0.9876, 1.5432, 0.8765], [0.7654, 1.3456, 0.6543, 1.2345]]), indicestensor([[1, 2, 0, 1], [2, 0, 1, 2]]))argmax(dim1)是分类任务取预测标签的标准写法见评估模式。切片SlicingPyTorch 采用与 NumPy 一致的索引语法切片返回共享内存的视图不复制数据 x torch.randn(2, 3, device0) x[:, 1] # 所有行的第 1 列 tensor([ 0.0475, -0.3882], devicecuda:0) x[1, :] # 第 1 行所有列 tensor([-0.1469, -0.3882, 0.2149], devicecuda:0) x[1, 1].item() # 提取标量脱离计算图 -0.3882044851779938 x torch.triu(torch.ones(5, 5)) # 上三角矩阵 x[:3, :3] # 子矩阵切片 tensor([[1., 1., 1.], [0., 1., 1.], [0., 0., 1.]]).item()用于把 0 维张量转成 Python 标量训练循环中打印 loss 时高频使用。高级索引布尔掩码与花式索引 x torch.randn(3, 4) mask x 0 # 布尔掩码 x[mask] # 取出所有正元素 tensor([0.5234, 1.2345, 0.8765, 0.3456, 1.5678]) x[x 0] 0 # 按条件原地置零 indices torch.tensor([0, 2]) # 花式索引整数数组 x[indices] x[[0, 1], [1, 2]] # 配对索引取 (0,1) 与 (1,2) 两个元素布尔掩码常用于数据过滤与 mask 填充如把 padding 位置置零花式索引用于按索引数组挑选行/列。拼接与分割拼接组合多个张量 x torch.randn(2, 3) y torch.randn(2, 3) torch.cat([x, y], dim0).shape # 沿已有维拼接22 torch.Size([4, 3]) torch.cat([x, y], dim1).shape # 沿列拼接33 torch.Size([2, 6]) torch.stack([x, y], dim0).shape # 新增一维堆叠 torch.Size([2, 2, 3]) torch.vstack([x, y]).shape # 垂直堆叠 torch.Size([4, 3]) torch.hstack([x, y]).shape # 水平堆叠 torch.Size([2, 6])cat在已有维度上合并要求其他维度相同stack创建新维度常用于 batch 组装与特征融合。分割把一个张量切分 x torch.randn(6, 4) chunks torch.split(x, 2, dim0) # 每块 2 行 → 3 块 len(chunks) 3 chunks torch.chunk(x, 3, dim0) # 均分 3 块 len(chunks) 3 tensors x.unbind(dim0) # 拆成 6 个单独行张量 len(tensors) 6split按块大小切chunk按块数量切unbind彻底解绑。这些操作在数据并行分发、批量推理拆分时很有用——src/llm/vllm/offline_bench.py 中就是用p for i, p in enumerate(prompts) if i % dp_size dp_rank把提示词列表按 DP rank 分片本质就是按索引切分思想在数据并行中的应用。自动微分梯度的核心机制自动微分autograd是 PyTorch 训练能力的根基前向传播时构建计算图反向传播时依据链式法则自动求导。开启/关闭梯度跟踪 x torch.randn(3, requires_gradTrue, device0) x tensor([-1.1442, -0.8709, -0.2581], devicecuda:0, requires_gradTrue) y x.detach() # 分离出不需要梯度的副本共享数据 y tensor([-1.1442, -0.8709, -0.2581], devicecuda:0) x.requires_grad_(False) # 原地关闭梯度跟踪 tensor([-1.1442, -0.8709, -0.2581], devicecuda:0)requires_gradTrue创建叶子张量时开启梯度跟踪。detach()得到与原张量共享数据但不在计算图内的视图用于取数值、实现 stop-gradient。requires_grad_(False)原地关闭跟踪。推理阶段禁用梯度推理/验证时无需反向传播禁用梯度可显著省显存、提速 x torch.randn(3, requires_gradTrue, device0) with torch.no_grad(): ... y x 1 ... print(y) ... tensor([1.2969, 1.5251, 0.7915], devicecuda:0) with torch.inference_mode(): # 更激进的推理模式性能更优 y x * 2 torch.no_grad() # 装饰器形式 def predict(x): return model(x)no_grad()与inference_mode()都用于推理后者额外禁用部分运行时特性性能更好。两者也常用作上下文管理器包裹评估循环见评估模式。反向传播示例 x torch.randn(3, requires_gradTrue) y x 1 z y * y * 3 z z.mean() # 损失必须是标量 z.backward() # 沿计算图反向传播 x.grad # 梯度累积在 .grad 上 tensor([1.2036, 5.0103, 0.5143]) x.grad.zero_() # 注意梯度默认是累加的需手动清零 z.backward()关键点backward()只能对标量调用损失通常就是标量x.grad存放梯度梯度默认累加连续两次backward()前必须zero_()否则会叠加。梯度累积用小显存模拟大 batch当 GPU 显存放不下大 batch 时可把多个小 batch 的梯度累积后再更新参数optimizer.zero_grad() for i, (inputs, targets) in enumerate(dataloader): outputs model(inputs) loss criterion(outputs, targets) loss.backward() # 梯度累加到 .grad if (i 1) % accumulation_steps 0: # 每 accumulation_steps 步更新一次 optimizer.step() optimizer.zero_grad()等效 batch size batch_size × accumulation_steps这是大模型训练中绕过显存瓶颈的标配手段。注意需要保证loss按累积步数做了归一化除以accumulation_steps否则梯度会被放大。构建神经网络nn.Module 与自定义网络通过继承nn.Module定义网络框架自动管理所有子模块的参数import torch.nn as nn class Net(nn.Module): def __init__(self): super().__init__() self.fc1 nn.Linear(784, 128) # 784 → 128 self.fc2 nn.Linear(128, 10) # 128 → 10 def forward(self, x): x torch.relu(self.fc1(x)) x self.fc2(x) return x model Net() model.to(cuda) # 把模型所有参数搬上 GPUforward()是前向逻辑的唯一入口内部可自由组合任意张量运算注册为属性的nn.Linear等模块的参数会被自动收集model.parameters()供优化器使用。常用层速查import torch.nn as nn # 基础层 linear nn.Linear(10, 5) # 全连接 conv2d nn.Conv2d(3, 64, kernel_size3, padding1) # 卷积3通道→64通道 maxpool nn.MaxPool2d(2, 2) # 最大池化 dropout nn.Dropout(0.5) # Dropout训练时随机丢弃 batchnorm nn.BatchNorm2d(64) # 批归一化 relu nn.ReLU() sigmoid nn.Sigmoid() softmax nn.Softmax(dim1) # 沿类别维做 softmax # 序列模型 lstm nn.LSTM(input_size10, hidden_size20, num_layers2) gru nn.GRU(input_size10, hidden_size20) embedding nn.Embedding(1000, 128) # 词表 1000 → 128 维向量nn.LSTM/nn.GRU/nn.Embedding是 NLP 与 LLM 基础设施Conv2d BatchNorm2d是视觉模型标配组合。损失函数import torch.nn as nn mse_loss nn.MSELoss() # 回归均方误差 cross_entropy nn.CrossEntropyLoss() # 多分类交叉熵内部含 softmax bce_loss nn.BCELoss() # 二分类需配合 Sigmoid 输出 bce_with_logits nn.BCEWithLogitsLoss() # 二分类融合 Sigmoid数值更稳 l1_loss nn.L1Loss() # L1 损失 nll_loss nn.NLLLoss() # 负对数似然配合 LogSoftmax outputs model(inputs) loss cross_entropy(outputs, targets)实践要点CrossEntropyLoss内部已集成 softmax模型最后一层不要再用Softmax(dim1)BCEWithLogitsLoss同理直接把 logits 传入即可数值稳定性更好。优化器import torch.optim as optim optimizer optim.SGD(model.parameters(), lr0.01, momentum0.9) optimizer optim.Adam(model.parameters(), lr0.001) optimizer optim.AdamW(model.parameters(), lr0.001, weight_decay0.01) # LLM 首选 optimizer optim.RMSprop(model.parameters(), lr0.01) optimizer.zero_grad() # 清零梯度 loss.backward() # 计算梯度 optimizer.step() # 更新参数SGD momentum经典稳定适合小模型。Adam自适应学习率开箱即用。AdamW解耦权重衰减是大模型预训练/微调的事实标准Megatron、vLLM 生态普遍采用。固定三步曲zero_grad → backward → step。学习率调度器from torch.optim.lr_scheduler import StepLR, ReduceLROnPlateau scheduler StepLR(optimizer, step_size30, gamma0.1) # 每 30 个 epoch 学习率 ×0.1 scheduler ReduceLROnPlateau(optimizer, modemin, patience10) # 验证指标停滞 10 轮则降 LR for epoch in range(num_epochs): train(...) val_loss validate(...) scheduler.step(val_loss) # 传入指标StepLR 则 scheduler.step() 即可StepLR按固定步数衰减ReduceLROnPlateau依据验证指标自适应衰减modemin表示指标越小越好。完整训练循环model.train() # 进入训练模式 for epoch in range(num_epochs): for inputs, targets in train_loader: inputs, targets inputs.to(device), targets.to(device) # 统一设备 optimizer.zero_grad() outputs model(inputs) loss criterion(outputs, targets) loss.backward() optimizer.step() print(fLoss: {loss.item():.4f})评估模式model.eval() # 关闭 Dropout / 固定 BN 统计量 with torch.no_grad(): # 禁用梯度省显存提速 for inputs, targets in test_loader: inputs inputs.to(device) outputs model(inputs) predictions outputs.argmax(dim1) # 取概率最大的类别train()/eval()切换控制 Dropout、BatchNorm 等层的行为差异评估必须配合no_grad()否则会白白构建计算图。保存与加载模型只存权重推荐用于推理部署torch.save(model.state_dict(), model.pth) model.load_state_dict(torch.load(model.pth))存完整检查点用于断点续训torch.save({ epoch: epoch, model_state_dict: model.state_dict(), optimizer_state_dict: optimizer.state_dict(), loss: loss, }, checkpoint.pth) checkpoint torch.load(checkpoint.pth) model.load_state_dict(checkpoint[model_state_dict]) optimizer.load_state_dict(checkpoint[optimizer_state_dict])完整检查点同时保存优化器状态动量、Adam 一阶/二阶矩保证续训时学习状态不丢state_dict方式则只保存参数适合纯推理。数据加载DataLoader 与自定义 DatasetDataLoader 与 TensorDatasetfrom torch.utils.data import DataLoader, TensorDataset dataset TensorDataset(x_train, y_train) # 由张量直接组数据集 loader DataLoader(dataset, batch_size32, shuffleTrue, num_workers4) for batch_x, batch_y in loader: outputs model(batch_x) loss criterion(outputs, batch_y)DataLoader自动完成打乱shuffleTrue、分 batchbatch_size32与多进程并行取数num_workers4是大规模数据加载的性能关键。自定义 Dataset继承Dataset并实现__len__与__getitem__两个方法即可接入任意数据源from torch.utils.data import Dataset 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(x_train, y_train) loader DataLoader(dataset, batch_size32)实际项目中__getitem__内通常还包含读文件、解码、数据增强等逻辑。NumPy 互操作与共享内存张量 ↔ NumPy 转换 x torch.randn([1, 2, 3], device0) y x.cpu().numpy() # GPU 张量必须先搬回 CPU y array([[[-0.11979043, 0.13762406, -1.2633433 ], [-0.380241 , 1.5320604 , -1.0828359 ]]], dtypefloat32) import numpy as np arr np.array([[1, 2], [3, 4]]) tensor torch.from_numpy(arr) # NumPy → 张量注意numpy()仅支持 CPU 张量GPU 张量必须显式.cpu()后再转换这正是 src/nccl/broadcast.py 等脚本中先同步再取数的惯用法根源。共享内存语义CPU 张量与 NumPy 数组可以共享同一块底层内存修改一方会影响另一方 x torch.randn(1, 2, 3) y x.numpy() x.add_(1) tensor([[[ 1.8195, 3.0259, 0.6733], [ 2.6539, 1.1562, -0.9821]]]) y # y 同步变化 array([[[ 1.8194908 , 3.0258512 , 0.67326605], [ 2.6539469 , 1.1561831 , -0.98211455]]], dtypefloat32)零拷贝互操作让 PyTorch 能与 NumPy/SciPy 生态无缝衔接但也意味着跨库修改会互相影响需要留意。可复现性随机种子训练结果可复现需要同时固定 PyTorch、CUDA 与 Python/NumPy 的随机源torch.manual_seed(42) # 固定 CPU 侧随机 torch.cuda.manual_seed(42) # 固定当前 CUDA 设备随机 torch.cuda.manual_seed_all(42) # 固定所有 CUDA 设备随机 torch.backends.cudnn.deterministic True # 强制确定性算法 torch.backends.cudnn.benchmark False # 关闭自动调优保证可复现 import numpy as np import random np.random.seed(42) random.seed(42)cudnn.benchmark True会按实际 shape 自动选择最优卷积算法从而提升速度但会引入不确定性调试对比实验时建议关闭。GPU 显存管理 torch.cuda.empty_cache() # 释放未使用的缓存显存 torch.cuda.memory_allocated() # 当前已分配显存字节 1073741824 torch.cuda.memory_reserved() # 当前保留含缓存显存 2147483648 torch.cuda.max_memory_allocated() # 峰值分配显存 3221225472 torch.cuda.reset_peak_memory_stats() # 重置峰值统计在训练前调用 print(torch.cuda.memory_summary()) # 打印完整显存使用报告显存管理三板斧用memory_allocated/memory_reserved定位占用训练前reset_peak_memory_stats()、训练后读max_memory_allocated()得到真实峰值OOM 时empty_cache()释放碎片缓存。仓库中的 vLLM 部署同样以gpu_memory_utilization0.9src/llm/vllm/offline_bench.py这类比例控制显存水位。分布式训练从 NCCL 到 torchrun / Slurm大模型训练必须跨多 GPU/多节点并行PyTorch 依赖NCCLNVIDIA Collective Communication Library完成高效的集合通信。初始化进程组与状态检查 import torch.distributed as dist dist.is_available() # 分布式支持是否编译可用 True dist.is_nccl_available() # NCCL 后端是否可用 True dist.is_initialized() # 进程组是否已初始化 False dist.init_process_group(backendnccl) dist.is_initialized() True dist.get_rank() # 当前进程的全局编号 0 dist.get_world_size() # 总进程数总 GPU 数 4 dist.get_backend() nccl仓库中的 src/nccl/broadcast.py 给出了生产级初始化写法先torch.cuda.set_device绑定本地 GPU再用dist.TCPStore建立协调存储master_addr、端口 29500、is_masterrank0最后dist.init_process_group(nccl, storestore, rank..., world_size..., device_id...)显式传入device_id——这是新版 PyTorch 推荐的写法能避免设备选择歧义。torchrun标准多卡启动器torchrun自动为每个进程设置RANK、LOCAL_RANK、WORLD_SIZE、MASTER_ADDR、MASTER_PORT等环境变量# 单机多卡 torchrun --nproc_per_node4 train.py # 多节点2 节点 × 4 卡 torchrun --nproc_per_node4 \ --nnodes2 \ --node_rank0 \ --master_addr192.168.1.1 \ --master_port29500 \ train.pypysheeet 仓库的 src/llm/vllm/offline_bench.py 直接给出了 torchrun 的实战用法脚本 docstring# 单卡 python offline_bench.py --model meta-llama/Llama-3.1-8B --num-prompts 50 # 张量并行4 卡 torchrun --nproc-per-node4 offline_bench.py \ --model Qwen/Qwen2-57B-A14B --tp-size 4 --enable-ep --num-prompts 100 # 数据并行 张量并行8 卡TP2, DP4 torchrun --nproc-per-node8 offline_bench.py \ --model meta-llama/Llama-3.1-8B --tp-size 2 --dp-size 4 --num-prompts 200该脚本还用dist.get_rank()/dist.is_initialized()判断分布式上下文用dist.all_gather_object汇总各 DP rank 的统计结果、dist.barrier()做同步——这些正是下文集合操作的实际落地。SlurmHPC 集群上的分布式训练在 Slurm 集群上sbatch/srun负责分配节点资源。pysheeet 的 Slurm 速查 提供了完整的作业提交与资源管理命令sinfo、scontrol show node、salloc -N 2、scancel等。将 torchrun 与 Slurm 结合的标准作业脚本#!/bin/bash #SBATCH --job-namepytorch_ddp #SBATCH --nodes2 #SBATCH --ntasks-per-node4 #SBATCH --gresgpu:4 #SBATCH --time24:00:00 export MASTER_ADDR$(scontrol show hostnames $SLURM_NODELIST | head -n 1) export MASTER_PORT29500 srun torchrun --nproc_per_node4 \ --nnodes$SLURM_NNODES \ --node_rank$SLURM_NODEID \ --master_addr$MASTER_ADDR \ --master_port$MASTER_PORT \ train.py关键点用scontrol show hostnames $SLURM_NODELIST | head -n 1取节点列表首个主机名作为MASTER_ADDR其余环境变量由 Slurm 与 torchrun 协作完成。仓库中的 src/megatron/srun.sh 展示了更完整的生产模式它通过srun --container-image ... --mpipmix --ntasks-per-node${GPUS_PER_NODE}在 enroot 容器内启动 Megatron 训练并把SLURM_PROCID/SLURM_LOCALID/SLURM_NTASKS映射为RANK/LOCAL_RANK/WORLD_SIZE同时配置MASTER_ADDR/MASTER_PORT、TORCH_DISTRIBUTED_BACKENDnccl以及 EFA/NCCL 环境变量NCCL_NET_PLUGIN、NCCL_BUFFSIZE、NCCL_P2P_NET_CHUNKSIZE等是Slurm 多节点 NCCL的完整可运行范式。对应的 DeepSeek-V2-Lite 预训练配方 则通过tensor_model_parallel_size8、expert_model_parallel_size2等参数定义了并行配置。DistributedDataParallelDDPDDP 对模型做梯度同步每个进程持有一份模型副本反向传播后通过 NCCL 的 all-reduce 聚合梯度保证所有副本步调一致import torch.nn as nn from torch.nn.parallel import DistributedDataParallel as DDP model Net().to(device) model DDP(model, device_ids[local_rank]) # local_rank 来自环境变量 for inputs, targets in train_loader: outputs model(inputs) loss criterion(outputs, targets) loss.backward() optimizer.step()DDP 的关键配套每个进程用torch.utils.data.distributed.DistributedSampler切分数据集避免各卡看到相同样本并在每个 epoch 调用sampler.set_epoch(epoch)重新打乱device_ids[local_rank]指定本进程绑定的 GPU数据加载、模型构建、进程组初始化必须保证各 rank 顺序一致否则会死锁。NCCL 集合通信操作NCCL 提供 all-reduce、broadcast、all-gather 等原语是梯度同步、参数分发与结果汇总的基础import torch.distributed as dist tensor torch.randn(2, 3).cuda() dist.all_reduce(tensor, opdist.ReduceOp.SUM) # 所有进程求和并写回 dist.broadcast(tensor, src0) # 根进程广播给所有人 dist.all_gather(tensor_list, tensor) # 收集所有进程的张量 dist.reduce(tensor, dst0, opdist.ReduceOp.SUM) # 归约到指定进程 dist.barrier() # 全局屏障同步all_reduce是 DDP 梯度同步的底层实现broadcast用于同步初始化参数与随机种子all_gather用于汇总各 rank 的评估指标——src/llm/vllm/offline_bench.py 的dist.all_gather_object(all_stats, local_stats)正是把各 DP rank 的 TTFT/TPOT 统计收集到 rank 0 再打印。仓库中的 src/nccl/broadcast.py 是集合操作的综合演练场它以 512 MB 的 bfloat16 张量为负载循环 100 次执行dist.broadcast(self.tensor, src0)用dist.barrier()torch.cuda.synchronize()保证测量边界用dist.all_reduce(elapsed, opdist.ReduceOp.MAX)取所有进程的最大耗时最后通过dist.broadcast(stop, src0)同步停止信号并校验数据正确性PASS。这段代码完整演示了初始化进程组 → 集合通信 → 同步计时 → 结果校验的分布式基准测试闭环可作为编写 NCCL 微基准的参考模板。小结一条从单卡到集群的完整路径pysheeet 的 PyTorch 速查文档 覆盖了深度学习开发的全部高频操作其能力边界可以概括为三条主线单卡基础张量创建/变换/运算 → 自动微分 →nn.Module建网 → DataLoader 取数 → 训练/评估/存取模型足以支撑完整的单卡深度学习实验性能与可复现设备搬移、原地操作、no_grad推理、显存统计、随机种子固定分布式扩展NCCL 进程组 → 集合通信原语 → torchrun 启动 → Slurm 调度 → DDP 封装是走向 LLM 多卡/多节点训练与推理的必经之路。在此基础上可继续深入 pysheeet 的 LLM 服务指南vLLM/SGLang/TensorRT-LLM 的 TP/PP/DP/EP 并行部署、Megatron 笔记 与 Slurm 速查将本文的 PyTorch 功底无缝迁移到大模型生产环境。赞分享文档教程开发工具【免费下载链接】pysheeetPython Cheat Sheet项目地址https://gitcode.com/gh_mirrors/py/pysheeet点击查看免费下载相关推荐PyTorch tutorials权威指南从基础张量操作到分布式训练PyTorch tutorials权威指南从基础张量操作到分布式训练 PyTorch作为深度学习领域最流行的框架之一其官方教程提供了从入门到进阶的完整学习路示例工程3步搞定抖音无水印批量下载douyin-downloader 上手指南3步搞定抖音无水印批量下载douyin downloader 上手指南 刷抖音最崩溃的场景看到一个博主发了 200 条作品你想全部收藏学习却只能一条一条网页爬虫CLIPyTorch DDP 分布式训练渐进式实战指南从单 GPU 到多节点 Slurm 集群PyTorch DDP 分布式训练渐进式实战指南从单 GPU 到多节点 Slurm 集群 导读 本文以本仓库 distributed/ddp tutorial示例工程人工智能深度学习上一篇Jest 类型系统探秘jest/types 共享类型包与类型化配置实战下一篇AppOpsX 开源项目安装与使用指南创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考