PyTorch 分布式訓(xùn)練的實現(xiàn)
在深度學(xué)習(xí)模型變得日益龐大之后,單個 GPU 的顯存已經(jīng)無法滿足高效訓(xùn)練的需求。此時,“分布式訓(xùn)練(Distributed Training)”技術(shù)應(yīng)運而生,成為加速訓(xùn)練的重要手段。
本文將圍繞以下三行典型的 PyTorch 分布式訓(xùn)練代碼進(jìn)行詳細(xì)解析,并擴展介紹分布式訓(xùn)練的核心概念和實踐方法:
local_rank = int(os.getenv('LOCAL_RANK', -1)) # https://pytorch.org/docs/stable/elastic/run.html global_rank = int(os.getenv('RANK', -1)) world_size = int(os.getenv('WORLD_SIZE', 1))
一、什么是分布式訓(xùn)練?
分布式訓(xùn)練是指將模型訓(xùn)練過程劃分到多個計算設(shè)備(通常是多個 GPU,甚至是多臺機器)上進(jìn)行協(xié)同處理,目標(biāo)是加速訓(xùn)練速度和擴展模型容量。
分布式訓(xùn)練可以分為以下幾種模式:
- 數(shù)據(jù)并行(Data Parallelism):每個 GPU 處理不同的數(shù)據(jù)子集,同步梯度。
- 模型并行(Model Parallelism):將模型拆成多個部分,分別部署到不同的 GPU。
- 混合并行(Hybrid Parallelism):結(jié)合模型并行和數(shù)據(jù)并行。
- 流水線并行(Pipeline Parallelism):按層切分模型,不同 GPU 處理不同階段。
二、理解分布式訓(xùn)練的核心概念
1. World Size(全局進(jìn)程數(shù))
world_size = int(os.getenv('WORLD_SIZE', 1))
- 含義:分布式訓(xùn)練中,所有參與訓(xùn)練的進(jìn)程總數(shù)。通常等于 GPU 總數(shù)。
- 作用:用于初始化進(jìn)程組(
torch.distributed.init_process_group()
),讓每個進(jìn)程知道集群的規(guī)模。
比如你有兩臺機器,每臺 4 塊 GPU,那么 world_size
= 8。
2. Rank(全局進(jìn)程編號)
global_rank = int(os.getenv('RANK', -1))
- 含義:標(biāo)識每個訓(xùn)練進(jìn)程在所有進(jìn)程中的唯一編號(從 0 開始)。
- 作用:常用于標(biāo)記主節(jié)點(rank == 0),控制日志輸出、模型保存等。
例如:
- rank=0:負(fù)責(zé)打印日志、保存模型
- rank=1,2,…:只做訓(xùn)練
3. Local Rank(本地進(jìn)程編號)
local_rank = int(os.getenv('LOCAL_RANK', -1))
含義:當(dāng)前訓(xùn)練進(jìn)程在本地機器上的 GPU 編號。一般與
CUDA_VISIBLE_DEVICES
配合使用。作用:用于指定該進(jìn)程應(yīng)該使用哪塊 GPU,如:
torch.cuda.set_device(local_rank)
三、環(huán)境變量的設(shè)置方式
這些環(huán)境變量通常由 分布式啟動器 設(shè)置。例如使用 torchrun
:
torchrun --nproc_per_node=4 --nnodes=2 --node_rank=0 \ --master_addr=192.168.1.1 --master_port=12345 train.py
torchrun
會自動為每個進(jìn)程設(shè)置:
LOCAL_RANK
RANK
WORLD_SIZE
也可以手動導(dǎo)出:
export WORLD_SIZE=8 export RANK=3 export LOCAL_RANK=3
四、分布式訓(xùn)練初始化流程(PyTorch 示例)
在 PyTorch 中,典型的初始化流程如下:
import os import torch import torch.distributed as dist def 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='env://'
:表示從環(huán)境變量中讀取初始化信息。nccl
是 NVIDIA 的高性能通信庫,支持 GPU 間高速通信。
五、分布式訓(xùn)練的代碼結(jié)構(gòu)
使用 PyTorch 實現(xiàn)分布式訓(xùn)練的基本框架:
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: # 正常訓(xùn)練流程
六、Elastic Training(彈性訓(xùn)練)
值得注意的是,示例代碼中注釋中提到的鏈接:https://pytorch.org/docs/stable/elastic/run.html
這是指 PyTorch 的 彈性分布式訓(xùn)練(Elastic Training),支持在訓(xùn)練過程中動態(tài)增加或移除節(jié)點,具備高容錯性。
- 工具:
torch.distributed.elastic
- 啟動命令:
torchrun --standalone --nnodes=1 --nproc_per_node=4 train.py
該特性對于大規(guī)模、長時間訓(xùn)練任務(wù)非常重要。
七、總結(jié)
變量名 | 含義 | 來源 | 典型用途 |
---|---|---|---|
WORLD_SIZE | 全局進(jìn)程數(shù)量 | torchrun 設(shè)置 | 初始化進(jìn)程組,全局通信 |
RANK | 當(dāng)前進(jìn)程的全局編號 | torchrun 設(shè)置 | 控制主節(jié)點行為 |
LOCAL_RANK | 當(dāng)前進(jìn)程在本地的 GPU 編號 | torchrun 設(shè)置 | 顯卡綁定:torch.cuda.set_device |
這三行代碼雖然簡單,卻是 PyTorch 分布式訓(xùn)練的入口。理解它們,就理解了 PyTorch 在分布式場景下的通信機制和訓(xùn)練框架。
如果你想要進(jìn)一步深入了解 PyTorch 分布式訓(xùn)練,推薦官方文檔:
到此這篇關(guān)于PyTorch 分布式訓(xùn)練的實現(xiàn)的文章就介紹到這了,更多相關(guān)PyTorch 分布式訓(xùn)練內(nèi)容請搜索腳本之家以前的文章或繼續(xù)瀏覽下面的相關(guān)文章希望大家以后多多支持腳本之家!
相關(guān)文章
Python使用LSTM實現(xiàn)銷售額預(yù)測詳解
大家經(jīng)常會遇到一些需要預(yù)測的場景,比如預(yù)測品牌銷售額,預(yù)測產(chǎn)品銷量。本文給大家分享一波使用?LSTM?進(jìn)行端到端時間序列預(yù)測的完整代碼和詳細(xì)解釋,需要的可以參考一下2022-07-07Python編程判斷一個正整數(shù)是否為素數(shù)的方法
這篇文章主要介紹了Python編程判斷一個正整數(shù)是否為素數(shù)的方法,涉及Python數(shù)學(xué)運算相關(guān)操作技巧,需要的朋友可以參考下2017-04-04python matplotlib如何給圖中的點加標(biāo)簽
這篇文章主要介紹了python matplotlib給圖中的點加標(biāo)簽,文中通過示例代碼介紹的非常詳細(xì),對大家的學(xué)習(xí)或者工作具有一定的參考學(xué)習(xí)價值,需要的朋友可以參考下2019-11-11