目录
- 一、什么是分布式训练？
- 二、领会分布式训练的核心概念
- 1. World Size（全局进程数）
- 2. Rank（全局进程编号）
- 3. Local Rank（本地进程编号）
- 三、环境变量的设置方式
- 四、分布式训练初始化流程（PyTorch 示例）
- 五、分布式训练的代码结构
- 六、Elastic Training（弹性训练）
- 七、拓展资料
在深度进修模型变得日益庞大之后，单个 GPU 的显存已经无法满足高效训练的需求。此时，“分布式训练（Distributed Training）”技术应运而生，成为加速训练的重要手段。
这篇文章小编将围绕下面内容三行典型的 PyTorch 分布式训练代码进行详细解析，并扩展介绍分布式训练的核心概念和操作技巧：
local_rank = int(os.getenv(‘LOCAL_RANK’, -1)) https://pytorch.org/docs/stable/elastic/run.htmlglobal_rank = int(os.getenv(‘RANK’, -1))world_size = int(os.getenv(‘WORLD_SIZE’, 1))
一、什么是分布式训练？
分布式训练是指将模型训练经过划分到多个计算设备（通常是多个 GPU，甚至是多台机器）上进行协同处理，目标是加速训练速度和扩展模型容量。
分布式训练可以分为下面内容几种模式：
- 数据并行（Data Parallelism）：每个 GPU 处理不同的数据子集，同步梯度。
- 模型并行（Model Parallelism）：将模型拆成多个部分，分别部署到不同的 GPU。
- 混合并行（Hybrid Parallelism）：结合模型并行和数据并行。
- 流水线并行（Pipeline Parallelism）：按层切分模型，不同 GPU 处理不同阶段。
二、领会分布式训练的核心概念
1. World Size（全局进程数）
world_size = int(os.getenv(‘WORLD_SIZE’, 1))
- 含义：分布式训练中，所有参与训练的进程总数。通常等于 GPU 总数。
- 影响：用于初始化进程组（
torch.distributed.init_process_group()），让每个进程知道集群的规模。
比如你有两台机器，每台 4 块 GPU，那么world_size= 8。
2. Rank（全局进程编号）
global_rank = int(os.getenv(‘RANK’, -1))
- 含义：标识每个训练进程在所有进程中的唯一编号（从 0 开始）。
- 影响：常用于标记主节点（rank == 0），控制日志输出、模型保存等。
例如：
- rank=0：负责打印日志、保存模型
- rank=1,2,…：只做训练
3. Local Rank（本地进程编号）
local_rank = int(os.getenv(‘LOCAL_RANK’, -1))
-
含义：当前训练进程在本地机器上的 GPU 编号。一般与
CUDA_VISIBLE_DEVICES配合使用。 -
影响：用于指定该进程应该使用哪块 GPU，如：
torch.cuda.set_device(local_rank)
三、环境变量的设置方式
这些环境变量通常由分布式启动器设置。例如使用torchrun：
torchrun –nproc_per_node=4 –nnodes=2 –node_rank=0 –master_addr=192.168.1.1 –master_port=12345 train.py
torchrun会自动为每个进程设置：
LOCAL_RANKRANKWORLD_SIZE
也可以手动导出：
export WORLD_SIZE=8export RANK=3export LOCAL_RANK=3
四、分布式训练初始化流程（PyTorch 示例）
在 PyTorch 中，典型的初始化流程如下：
import osimport torchimport torch.distributed as distdef setup_distributed(): local_rank = int(os.getenv(‘LOCAL_RANK’, -1)) global_rank = int(os.getenv(‘RANK’, -1)) world_size = int(os.getenv(‘WORLD_SIZE’, 1)) torch.cuda.set_device(local_rank) dist.init_process_group( backend=’nccl’, GPU 用 nccl，CPU 用 gloo init_method=’env://’, world_size=world_size, rank=global_rank )
init_method=&039;env://&039;：表示从环境变量中读取初始化信息。nccl是 NVIDIA 的高性能通信库，支持 GPU 间高速通信。
五、分布式训练的代码结构
使用 PyTorch 实现分布式训练的基本框架：
def train(): setup_distributed() model = MyModel().cuda() model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[local_rank]) dataset = MyDataset() sampler = torch.utils.data.distributed.DistributedSampler(dataset) dataloader = DataLoader(dataset, sampler=sampler, batch_size=64) for epoch in range(epochs): sampler.set_epoch(epoch) for batch in dataloader: 正常训练流程
六、Elastic Training（弹性训练）
需要关注的是，示例代码中注释中提到的链接：https://pytorch.org/docs/stable/elastic/run.html
这是指 PyTorch 的弹性分布式训练（Elastic Training），支持在训练经过中动态增加或移除节点，具备高容错性。
- 工具：
torch.distributed.elastic - 启动命令：
torchrun --standalone --nnodes=1 --nproc_per_node=4 train.py
该特性对于大规模、长时刻训练任务非常重要。
七、拓展资料
| 变量名 | 含义 | 来源 | 典型用途 |
|---|---|---|---|
| WORLD_SIZE | 全局进程数量 | torchrun 设置 | 初始化进程组，全局通信 |
| RANK | 当前进程的全局编号 | torchrun 设置 | 控制主节点行为 |
| LOCAL_RANK | 当前进程在本地的 GPU 编号 | torchrun 设置 | 显卡绑定：torch.cuda.set_device |
这三行代码虽然简单，却是 PyTorch 分布式训练的入口。领会它们，就领会了 PyTorch 在分布式场景下的通信机制和训练框架。
如果你想要进一步深入了解 PyTorch 分布式训练，推荐官方文档：
- PyTorch 分布式训练文档
- torchrun 使用说明
到此这篇关于PyTorch 分布式训练的实现的文章就介绍到这了,更多相关PyTorch 分布式训练内容请搜索风君子博客以前的文章或继续浏览下面的相关文章希望大家以后多多支持风君子博客！
无论兄弟们可能感兴趣的文章:
- 分布式训练training-operator和pytorch-distributed?RANK变量不统一解决
- 教你掌握分布式训练PyTorch?DDP到Accelerate到Trainer
