PyTorch分布式訓練的實現(xiàn)

更新時間：2026年01月30日 08:31:12 作者：盼小輝丶

本文主要介紹了PyTorch分布式訓練的實現(xiàn),文中通過示例代碼介紹的非常詳細,對大家的學習或者工作具有一定的參考學習價值,需要的朋友們下面隨著小編來一起學習學習吧

0. 前言

在將預訓練的機器學習模型投入生產(chǎn)環(huán)境之前，模型訓練是不可或缺的關(guān)鍵環(huán)節(jié)。隨著深度學習的發(fā)展，大模型往往具有數(shù)百萬乃至數(shù)十億參數(shù)。使用反向傳播來調(diào)整這些參數(shù)需要大量的內(nèi)存和計算資源。即便如此，模型訓練仍然可能需要數(shù)天甚至數(shù)月時間才能完成。
在本節(jié)中，我們將探討如何通過跨機器和機器內(nèi)多進程的分布式訓練來加速模型訓練過程。我們將系統(tǒng)學習 PyTorch 提供的三大分布式訓練 API——torch.distributed、torch.multiprocessing 以及 torch.utils.data.distributed.DistributedSampler，使用這些 API 能夠極大的簡化分布式訓練，介紹如何使用 PyTorch 的分布式訓練工具，在 CPU 和 GPU 上加速訓練。
通過本節(jié)學習，將能夠充分釋放硬件設備的訓練潛力。對于超大規(guī)模模型訓練而言，本節(jié)所探討的工具不僅至關(guān)重要，在某些情況下甚至是不可或缺的。

1. 使用 PyTorch 進行分布式訓練

在本節(jié)中，我們將模型訓練過程從常規(guī)訓練轉(zhuǎn)換為分布式訓練，探討 PyTorch 提供的分布式訓練工具，這些工具能顯著提升訓練速度并優(yōu)化硬件使用效率。

1.1 以常規(guī)方式訓練模型

(1) 首先導入所需庫：

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

import time
import argparse

device = torch.device("cpu")

(2) 接下來，定義卷積神經(jīng)網(wǎng)絡 (Convolutional Neural Network, CNN) 模型架構(gòu)：

class ConvNet(nn.Module):
    def __init__(self):
        super(ConvNet, self).__init__()
        self.cn1 = nn.Conv2d(1, 16, 3, 1)
        self.cn2 = nn.Conv2d(16, 32, 3, 1)
        self.dp1 = nn.Dropout(0.10)
        self.dp2 = nn.Dropout(0.25)
        self.fc1 = nn.Linear(4608, 64) # 4608 is basically 12 X 12 X 32
        self.fc2 = nn.Linear(64, 10)
 
    def forward(self, x):
        x = self.cn1(x)
        x = F.relu(x)
        x = self.cn2(x)
        x = F.relu(x)
        x = F.max_pool2d(x, 2)
        x = self.dp1(x)
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        x = F.relu(x)
        x = self.dp2(x)
        x = self.fc2(x)
        op = F.log_softmax(x, dim=1)
        return op

(3) 然后，定義模型的訓練過程：

def train(args):
    train_dataloader = torch.utils.data.DataLoader(
        datasets.MNIST('./data', train=True, download=True,
                       transform=transforms.Compose([
                           transforms.ToTensor(),
                           transforms.Normalize((0.1302,), (0.3069,))])),
        batch_size=128, shuffle=True)  
    model = ConvNet()
    optimizer = optim.Adadelta(model.parameters(), lr=0.5)
    model.train()

在函數(shù)的前半部分，使用 PyTorch 訓練數(shù)據(jù)集定義了 PyTorch 的訓練數(shù)據(jù)加載器。實例化卷積神經(jīng)網(wǎng)絡 (ConvNet)，并定義了優(yōu)化器。

    for epoch in range(args.epochs):
        for b_i, (X, y) in enumerate(train_dataloader):
            X, y = X.to(device), y.to(device)
            pred_prob = model(X)
            loss = F.nll_loss(pred_prob, y) # nll is the negative likelihood loss
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            if b_i % 10 == 0:
                print('epoch: {} [{}/{} ({:.0f}%)]\t training loss: {:.6f}'.format(
                    epoch, b_i, len(train_dataloader),
                    100. * b_i / len(train_dataloader), loss.item()))

在函數(shù)的后半部分，運行訓練循環(huán)預定義的 epoch 數(shù)。在循環(huán)內(nèi)，通過批數(shù)據(jù)的方式遍歷整個訓練數(shù)據(jù)集，本節(jié)中批大小為 128。對于每個包含 128 個訓練數(shù)據(jù)點的批次，使用模型進行前向傳播，以計算預測概率。然后，我們將預測結(jié)果結(jié)合真實標簽計算批次損失，并通過反向傳播利用該損失梯度來調(diào)整模型參數(shù)。

(4) 將所有組件整合在 main() 函數(shù)中：

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--epochs', default=1, type=int)
    args = parser.parse_args()
    start = time.time()
    train(args)
    print(f"Finished training in {time.time()-start} secs")

使用參數(shù)解析器，它可以幫助我們在運行 Python 訓練程序時從命令行輸入超參數(shù)，例如 epoch 數(shù)。我們還對訓練過程進行了計時，以便可以將它與分布式訓練過程進行比較。

(5) 最后，確保通過命令行執(zhí)行腳本時能運行 main() 函數(shù)：

if __name__ == '__main__':
    main()

(6) 在命令行中執(zhí)行以下命令來運行該 Python 腳本：

$ python convnet_undistributed.py -- epoches 1

本節(jié)中我們僅設置訓練一個 epoch，因為當前重點不在于模型精度，而在于模型訓練耗時?？梢钥吹捷敵鼋Y(jié)果如下所示：

訓練 1 個 epoch 大約花費了 28 秒，一個 epoch 共包含 469 個批次，實際訓練時間會隨硬件配置差異而波動。

1.2 分布式訓練模型

通過使用 PyTorch 提供的分布式處理 API，即使存在跨進程或跨機器重復傳遞數(shù)據(jù)的額外開銷，模型訓練速度也能顯著提升。

(1) 首先，導入所需庫，增加幾個與分布式訓練相關(guān)的模塊：

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

import torch.multiprocessing as mp
import torch.distributed as dist

import os
import time
import argparse

torch.multiprocessing 用于在單臺機器上生成多個 Python 進程(通常根據(jù) CPU 核心數(shù)生成對應數(shù)量的進程)，而torch.distributed則實現(xiàn)不同機器間的通信協(xié)作，使它們能共同完成模型訓練。執(zhí)行時，我們需要在每臺參與訓練的機器上顯式啟動訓練腳本。
PyTorch 內(nèi)置的通信后端(如 Gloo )會自動處理機器間的通信協(xié)調(diào)。在每臺機器內(nèi)部，多進程機制會進一步將訓練任務并行分配到各個進程。

(2) 模型架構(gòu)定義部分保持不變：

class ConvNet(nn.Module):
    def __init__(self):
        super(ConvNet, self).__init__()
        self.cn1 = nn.Conv2d(1, 16, 3, 1)
        self.cn2 = nn.Conv2d(16, 32, 3, 1)
        self.dp1 = nn.Dropout2d(0.10)
        self.dp2 = nn.Dropout2d(0.25)
        self.fc1 = nn.Linear(4608, 64) # 4608 is basically 12 X 12 X 32
        self.fc2 = nn.Linear(64, 10)
 
    def forward(self, x):
        x = self.cn1(x)
        x = F.relu(x)
        x = self.cn2(x)
        x = F.relu(x)
        x = F.max_pool2d(x, 2)
        x = self.dp1(x)
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        x = F.relu(x)
        x = self.dp2(x)
        x = self.fc2(x)
        op = F.log_softmax(x, dim=1)
        return op

(3) 定義 train() 函數(shù)：

def train(cpu_num, args):
    rank = args.machine_id * args.num_processes + cpu_num                        
    dist.init_process_group(                                   
    backend='gloo',                                         
    init_method='env://',                                   
    world_size=args.world_size,                              
    rank=rank                                               
    ) 
    torch.manual_seed(0)
    device = torch.device("cpu")

可以看到，代碼開頭新增了兩條關(guān)鍵語句。首先是計算進程的 rank 值——這本質(zhì)上是該進程在整個分布式系統(tǒng)中的順序標識符。舉例來說，若使用 2 臺各配備 4 核 CPU 的機器進行訓練，為充分利用硬件資源可能需要啟動 8 個進程(每臺機器 4 個)。此時就需要為這些進程建立標識體系：先為兩臺機器分配 ID 0 和 1，再為每臺機器內(nèi)的 4 個進程分配子 ID 0 至 3。最終，第 n 臺機器上第 i 個進程的全局 rank 值可通過以下公式確定：
r a n k = n × 4 + k rank=n\times 4+k rank=n×4+k
第二行代碼使用了 torch.distributed 模塊中的 init_process_group，該方法為每個啟動的進程配置以下關(guān)鍵參數(shù)：

用于機器間通信的后端(在節(jié)使用 Gloo)
參與分布式訓練的進程總量(由 args.world_size 指定)，亦稱 world_size
當前啟動進程的全局 rank 值

init_process_group 方法會阻塞所有進程，直到跨機器的全部進程都完成初始化才會繼續(xù)執(zhí)行。
PyTorch 提供了三種內(nèi)置的分布式訓練后端：

Gloo
NCCL
MPI

簡而言之，對于 CPU 上的分布式訓練，使用 Gloo，對于 GPU，使用 NCCL。

    train_dataset = datasets.MNIST('./data', train=True, download=True,
                                   transform=transforms.Compose([
                                       transforms.ToTensor(),
                                       transforms.Normalize((0.1302,), (0.3069,))]))  
    train_sampler = torch.utils.data.distributed.DistributedSampler(
        train_dataset,
        num_replicas=args.world_size,
        rank=rank
    )
    train_dataloader = torch.utils.data.DataLoader(
       dataset=train_dataset,
       batch_size=args.batch_size,
       shuffle=False,            
       num_workers=0,
       sampler=train_sampler)
    model = ConvNet()
    optimizer = optim.Adadelta(model.parameters(), lr=0.5)
    model = nn.parallel.DistributedDataParallel(model)
    model.train()

與單機訓練相比，分布式訓練的關(guān)鍵改進體現(xiàn)在數(shù)據(jù)加載與模型封裝兩個層面。我們將 MNIST 數(shù)據(jù)集實例化與數(shù)據(jù)加載器拆分為獨立步驟，其間插入 DistributedSampler 采樣器。該采樣器將訓練數(shù)據(jù)均分為 world_size 個分區(qū)，確保每個進程處理等量數(shù)據(jù)。注意數(shù)據(jù)加載器的 shuffle 參數(shù)需設為 False，因為數(shù)據(jù)分配已由采樣器控制。
代碼中的另一個新增部分是 nn.parallel.DistributedDataParallel 函數(shù)，它應用于模型對象。這部分可能是代碼中最重要的部分，因為 DistributedDataParallel 是實現(xiàn)分布式梯度下降算法的關(guān)鍵組件。其底層運行機制如下：

分布式環(huán)境中的每個派生進程都會獲得獨立的模型副本
每個進程的模型都維護自己的優(yōu)化器，并與全局迭代保持同步的局部優(yōu)化步驟
在每次分布式訓練迭代時，各進程獨立計算損失值及梯度，隨后跨進程對這些梯度求取平均值
平均后的梯度將通過全局反向傳播機制同步到所有模型副本，用于調(diào)整參數(shù)
由于全局反向傳播步驟的存在，所有模型參數(shù)在每次迭代時都保持一致，從而實現(xiàn)自動同步

DistributedDataParallel 通過讓每個 Python 進程運行在獨立的解釋器上，有效規(guī)避了在單一解釋器下多線程實例化多個模型可能引發(fā)的全局解釋器鎖 (Global Interpreter Lock, GIL) 限制問題。這進一步提升了性能表現(xiàn)，特別是對于那些需要大量Python專屬運算的模型而言。

    for epoch in range(args.epochs):
        for b_i, (X, y) in enumerate(train_dataloader):
            X, y = X.to(device), y.to(device)
            pred_prob = model(X)
            loss = F.nll_loss(pred_prob, y) # nll is the negative likelihood loss
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            if b_i % 10 == 0 and cpu_num==0:
                print('epoch: {} [{}/{} ({:.0f}%)]\t training loss: {:.6f}'.format(
                    epoch, b_i, len(train_dataloader),
                    100. * b_i / len(train_dataloader), loss.item()))

最后，訓練循環(huán)幾乎和單機訓練一樣。唯一的區(qū)別在于我們限制只有排名為0的進程才能獲取日志信息。這是因為排名為 0 的機器用于建立所有通信連接。因此，我們通常將排名為 0 的進程作為參考來跟蹤模型訓練性能。如果不加以限制，每個模型訓練迭代都會產(chǎn)生與進程數(shù)量相同的日志行數(shù)。

(4) 將所有組件整合在 main() 函數(shù)中：

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--num-machines', default=1, type=int,)
    parser.add_argument('--num-processes', default=1, type=int)
    parser.add_argument('--machine-id', default=0, type=int)
    parser.add_argument('--epochs', default=1, type=int)
    parser.add_argument('--batch-size', default=128, type=int)
    args = parser.parse_args()
    
    args.world_size = args.num_processes * args.num_machines                
    os.environ['MASTER_ADDR'] = '127.0.0.1'              
    os.environ['MASTER_PORT'] = '8892'      
    start = time.time()
    mp.spawn(train, nprocs=args.num_processes, args=(args,))
    print(f"Finished training in {time.time()-start} secs")

首先，我們新增了以下參數(shù)：

num_machines：機器總數(shù)量
num_processes：每臺機器上要啟動的進程數(shù)量
machine_id：當前機器的序號 ID。需要注意的是，這個 Python 腳本需要在每臺機器上單獨啟動
batch_size：每個批次的數(shù)據(jù)樣數(shù)量。參數(shù)作用如下：
- 所有進程將各自計算梯度，這些梯度將在每次迭代中被平均，從而獲得整體梯度
- 完整的訓練數(shù)據(jù)集被分割為 world_size 個獨立的數(shù)據(jù)集

因此，在每次迭代時，完整的批次數(shù)據(jù)需要被分割成 world_size 個子批次，每個進程處理一個子批次。因為 batch_size 現(xiàn)在與 world_size 相關(guān)聯(lián)，所以我們將其作為輸入?yún)?shù)提供，目的是為了簡化訓練接口。
參數(shù)定義后，計算 world_size 作為派生參數(shù)。接著，我們定義兩個重要的環(huán)境變量：

MASTER_ADDR：運行 rank 0 進程的主機 IP 地址
MASTER_PORT：運行 rank 0 進程的主機上可用的端口號

rank 0 機器負責建立所有后端通信連接，因此整個系統(tǒng)必須能隨時定位到該主機，這就是為什么需要提供其 IP 地址和端口號。本節(jié)中訓練任務將在單臺本地機器上運行，因此使用 localhost 地址即可，但在跨服務器的多機訓練場景中，則需要提供 rank 0 服務器的真實 IP 地址及空閑端口號。
最后一個變化是使用多進程 (multiprocessing) 來在每臺機器上啟動 num_processes 個進程，而非僅運行單個訓練進程。分布式參數(shù)會傳遞給每個派生進程，確保模型訓練過程中各進程與機器之間能自主協(xié)調(diào)。

(5) 分布式訓練：

if __name__ == '__main__':
    main()

(6) 啟動分布式訓練腳本。首先使用分布式腳本進行類非分布式運行，將機器數(shù)量和進程數(shù)量都設置為 1：

$ python convnet_distributed.py --num-machines 1 --num-processes 1 --machine-id 0 --epochs 1 --batch-size 128

需要注意的是，由于本次訓練只使用單個進程，batch_size 與之前非分布式訓練時保持一致(仍為 128)。運行結(jié)果如下輸出：

若將此結(jié)果與上一節(jié)非分布式訓練的輸出對比，可發(fā)現(xiàn)訓練時間基本相當(約 30 秒)，損失值變化也十分相似。

(7) 接下來，運行一個真正的分布式訓練，使用 2 個進程而不是 1 個進程。相應地，將 batch_size 從 128 降至 64：

$ python convnet_distributed.py --num-machines 1 --num-processes 2 --machine-id 0 --epochs 1 --batch-size 64

輸出結(jié)果如下所示：

可以看到，訓練時間從 30 秒減少到了 20 秒。訓練損失的變化趨勢沒有受到影響，這表明分布式訓練可以加速訓練過程，同時保持模型的準確性。

(8) 接下來，使用 4 個進程，并相應地將批大小從 64 降至 32：

$ python convnet_distributed.py --num-machines 1 --num-processes 4 --machine-id 0 --epochs 1 --batch-size 32

輸出結(jié)果如下所示：

可以看到，訓練時間進一步減少，從 20 秒降至 15 秒。訓練損失的變化趨勢仍然與之前的訓練相似。通過分布式訓練，我們已經(jīng)將訓練時間從 30 秒縮短到了 15 秒，減少了 2 倍。

(9) 進一步增加進程數(shù)，使用 8 個進程代替 4 個進程，并相應地將批次大小從 32 降至 16：

$ python convnet_distributed.py --num-machines 1 --num-processes 8 --machine-id 0 --epochs 1 --batch-size 16

輸出結(jié)果如下所示：

與預期相反，訓練時間不僅沒有進一步縮短，反而從 15 秒略微增加至 18 秒。由于代碼在本地機器執(zhí)行，系統(tǒng)還存在其他進程(如瀏覽器)會與部分分布式訓練進程爭奪資源。如果分布式訓練模型是在遠程機器上進行的，同時這些機器的唯一任務就是進行模型訓練，在這樣的機器上，建議使用與 CPU 核心數(shù)相等甚至更多的進程數(shù)。

(10) 最后需要指出的是，由于在本節(jié)中我們只使用了一臺機器，因此我們只需要啟動一個 Python 腳本來開始訓練。然而，如果是在多臺機器上進行訓練，那么除了修改 MASTER_ADDR 和 MASTER_PORT 外，還需要在每臺機器上啟動一個 Python 腳本。例如，如果有 2 臺機器，在機器 1 上執(zhí)行：

$ python distributed_script.py --num_machines=2 --num-processes 8 --machine_id=0 --epochs 1 --batch-size 16

在機器 2 上執(zhí)行：

$ python distributed_script.py --num_machines=2 --num-processes 8 --machine_id=1 --epochs 1 --batch-size 16

至此，我們完成了關(guān)于使用 PyTorch 在 CPU 上實施分布式訓練深度學習模型的實踐探討，這種方法能帶來顯著的加速效果。僅需添加少量代碼，就能將常規(guī) PyTorch 模型訓練腳本升級為分布式訓練模式。雖然上述實驗基于簡單的卷積網(wǎng)絡，但由于我們完全無需修改模型架構(gòu)代碼，因此這套方案可直接擴展到更復雜的模型訓練場景。接下來，我們將簡要討論如何應用類似的代碼更改，實現(xiàn) GPU 環(huán)境下的分布式訓練。

2. GPU 分布式訓練

我們通常使用以下 PyTorch 代碼定義模型訓練設備：

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

這行代碼的作用是自動檢測可用計算設備，并優(yōu)先選擇 CUDA (GPU)。這種優(yōu)先選擇源于 GPU 通過并行化處理神經(jīng)網(wǎng)絡常規(guī)運算(如矩陣乘法和加法)所能提供的顯著加速優(yōu)勢。本節(jié)我們將探討如何通過 GPU 分布式訓練進一步加速模型訓練。

(1) 雖然導入語句和模型架構(gòu)定義代碼與使用 CPU 分布式訓練一節(jié)完全一致，但 train() 函數(shù)中有幾處關(guān)鍵修改：

def train(gpu_num, args):
    rank = args.machine_id * args.num_gpu_processes + gpu_num                        
    dist.init_process_group(                                   
        backend='nccl',                                         
        init_method='env://',                                   
        world_size=args.world_size,                              
        rank=rank                      
    ) 
    model = ConvNet()
    torch.cuda.set_device(gpu_num)
    model.cuda(gpu_num)
    criterion = nn.NLLLoss().cuda(gpu_num)

在使用 GPU 時，NCCL 是首選的通信后端。同時，模型和損失函數(shù)都必須部署到 GPU 設備上，以確保充分利用 GPU 提供的并行矩陣運算加速能力：

    train_dataset = datasets.MNIST('./data', train=True, download=True,
                                   transform=transforms.Compose([
                                       transforms.ToTensor(),
                                       transforms.Normalize((0.1302,), (0.3069,))]))  
    train_sampler = torch.utils.data.distributed.DistributedSampler(
        train_dataset,
        num_replicas=args.world_size,
        rank=rank
    )
    train_dataloader = torch.utils.data.DataLoader(
       dataset=train_dataset,
       batch_size=args.batch_size,
       shuffle=False,            
       num_workers=0,
       pin_memory=True,
       sampler=train_sampler)
    optimizer = optim.Adadelta(model.parameters(), lr=0.5)
    model = nn.parallel.DistributedDataParallel(model,
                                                device_ids=[gpu_num])
    model.train()

DistributedDataParallel API 包含一個關(guān)鍵參數(shù)——device_ids，用于指定調(diào)用該 API 的 GPU 進程 ID。此外可以看到，數(shù)據(jù)加載器 (dataloader) 中新增了 pin_memory 參數(shù)并設為 True，該參數(shù)能顯著加速訓練過程中從主機(此處指加載數(shù)據(jù)集的 CPU )到各設備 (GPU) 的數(shù)據(jù)傳輸。
pin_memory 機制的工作原理是將數(shù)據(jù)"鎖定" (pin) 在 CPU 內(nèi)存中，即把數(shù)據(jù)樣本分配到固定的頁鎖定內(nèi)存區(qū)域。訓練時，這些內(nèi)存區(qū)域的數(shù)據(jù)會被高效地拷貝到對應 GPU。該機制需與 non_blocking=True 參數(shù)配合使用：

    for epoch in range(args.epochs):
        for b_i, (X, y) in enumerate(train_dataloader):
            X, y = X.cuda(non_blocking=True), y.cuda(non_blocking=True)
            pred_prob = model(X)
            loss = criterion(pred_prob, y) 
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            if b_i % 10 == 0 and gpu_num==0:
                print('epoch: {} [{}/{} ({:.0f}%)]\t training loss: {:.6f}'.format(
                    epoch, b_i, len(train_dataloader),
                    100. * b_i / len(train_dataloader), loss.item()))

通過調(diào)用參數(shù) pin_memory 和 non_blocking，使得以下兩者之間的操作得以重疊：

從 CPU 到 GPU 數(shù)據(jù)(真實標簽)的傳輸
GPU 模型訓練計算(或 GPU 內(nèi)核執(zhí)行)

這從根本上提升了整體 GPU 訓練流程的效率。

(2) 除了 train() 函數(shù)的修改外，main() 函數(shù)同樣有所調(diào)整調(diào)整：

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--num-machines', default=1, type=int,)
    parser.add_argument('--num-gpu-processes', default=1, type=int)
    parser.add_argument('--machine-id', default=0, type=int)
    parser.add_argument('--epochs', default=1, type=int)
    parser.add_argument('--batch-size', default=64, type=int)
    args = parser.parse_args()
    
    args.world_size = args.num_gpu_processes * args.num_machines                
    os.environ['MASTER_ADDR'] = '127.0.0.1'              
    os.environ['MASTER_PORT'] = '8892'      
    start = time.time()
    mp.spawn(train, nprocs=args.num_gpu_processes, args=(args,))
    print(f"Finished training in {time.time()-start} secs")

用 num_gpu_processes 替代了原來的 num_process 參數(shù)，該參數(shù)值通過 torch.cuda.device_count() 自動獲取可用 GPU 數(shù)量。其余代碼相應調(diào)整，但 GPU 版本的核心邏輯與之前保持一致。執(zhí)行以下命令即可啟動 GPU 分布式訓練：

$ python convnet_distributed_cuda.py --num-machines 1 --num-gpu-processes 1 --machine-id 0 --epochs 1 --batch-size 128

至此，我們已完成關(guān)于使用 PyTorch 進行 GPU 分布式模型訓練的簡要探討。上述代碼同樣適用于其他深度學習模型，當前深度學習模型大多采用 GPU 分布式訓練方案。此外，Horovod、DeepSpeed 和 PyTorch Lightning 等庫都提供了更簡潔的 API 來簡化 PyTorch 模型的分布式訓練流程。