分布式訓練 DLC 快速入門 - Platform For AI

DLC 可快捷建立分布式或單機訓練任務，無需手動購買機器和配置環境。本文以MNIST手寫體識別為例，示範單機單卡訓練和多機多卡分布式訓練的完整流程。

說明

MNIST手寫體識別是深度學習最經典的入門任務之一，任務目標是通過構建機器學習模型，來識別10個手寫數字（0～9）。

一. 前提條件

使用主帳號開通PAI並建立工作空間。登入PAI控制台，左上方選擇開通地區，然後一鍵授權和開通產品。

二. 計費說明

本文案例將使用公用資源建立DLC任務，計費方式為隨用隨付，計費規則參見分布式訓練（DLC）計費說明。

三. 單機單卡訓練

3.1 建立資料集

資料集用於儲存模型訓練的代碼、資料、以及訓練結果。本文以Object Storage Service類型資料集為例進行說明。

在PAI控制台左側功能表列單擊数据集 > 自定义数据集 > 新建数据集。
配置資料集參數。關鍵參數配置如下，其他參數預設即可。
- 名称：如：dataset_mnist
- 存储类型：对象存储（OSS）
- OSS路径：單擊表徵圖，選擇bucket並建立目錄如：dlc_mnist。
  如果您尚未開通OSS，或在當前地區下沒有可選的Bucket，可參考如下步驟開通OSS，並建立Bucket：
  （可選）開通OSS，並建立Bucket
  1. 開通OSS服務。
  2. 登入OSS管理主控台，單擊建立 Bucket，填寫Bucket名稱，地域選擇與當前PAI相同的地區，其他參數預設即可，然後單擊完成建立。
單擊確定建立資料集。

上傳訓練代碼和資料。

下載代碼。本文已經為您準備好了訓練代碼，單擊mnist_train.py下載。為減少您的操作，代碼運行時會自動將訓練資料下載到資料集的dataSet目錄中。

您在後續實際業務使用時，可以預先把代碼和訓練資料上傳到PAI的資料集中。

單機單卡訓練程式碼範例 mnist_train.py

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
from torch.utils.tensorboard import SummaryWriter
# 超參數
batch_size = 64  # 每次訓練的資料量
learning_rate = 0.01  # 學習率
num_epochs = 20  # 訓練輪次
# 檢查是否有可用的 GPU
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# 資料預先處理
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])
train_dataset = datasets.MNIST(root='/mnt/data/dataSet', train=True, download=True, transform=transform)
val_dataset = datasets.MNIST(root='/mnt/data/dataSet', train=False, download=False, transform=transform)
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)
# 定義簡單的神經網路
class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        # 第一層卷積：輸入通道1（灰階映像），輸出通道10，卷積核5x5
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        # 第二層卷積：輸入通道10，輸出通道20，卷積核3x3
        self.conv2 = nn.Conv2d(10, 20, kernel_size=3)
        # 全串連層：輸入為20*5*5（卷積+池化後的特徵圖尺寸），輸出128
        self.fc1 = nn.Linear(20 * 5 * 5, 128)
        # 輸出層：128 -> 10（對應10個數字類別）
        self.fc2 = nn.Linear(128, 10)
    def forward(self, x):
        # 輸入x形狀: [batch, 1, 28, 28]
        x = F.max_pool2d(F.relu(self.conv1(x)), 2)  # [batch, 10, 12, 12]
        x = F.max_pool2d(F.relu(self.conv2(x)), 2)  # [batch, 20, 5, 5]
        x = x.view(-1, 20 * 5 * 5)  # 展平為[batch, 500]
        x = F.relu(self.fc1(x))      # [batch, 128]
        x = self.fc2(x)              # [batch, 10]
        return x
# 執行個體化模型，並將其移動到 GPU 上（如果可用）
model = SimpleCNN().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=learning_rate)
# 建立 TensorBoard 的 SummaryWriter，可用於可視化的查看模型訓練過程
writer = SummaryWriter('/mnt/data/output/runs/mnist_experiment')
# 用於儲存最高準確率的模型的變數
best_val_accuracy = 0.0
# 訓練模型並記錄損失和準確率
for epoch in range(num_epochs):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)  # 將資料和目標移動到 GPU
        # 清零梯度
        optimizer.zero_grad()
        # 前向傳播
        output = model(data)
        # 計算損失
        loss = criterion(output, target)
        # 反向傳播
        loss.backward()
        # 更新參數
        optimizer.step()
        # 記錄訓練損失到 TensorBoard
        if batch_idx % 100 == 0:  # 每 100 個批次記錄一次
            writer.add_scalar('Loss/train', loss.item(), epoch * len(train_loader) + batch_idx)
            print(f'Train Epoch: {epoch} [{batch_idx * len(data)}/{len(train_loader.dataset)} ({100. * batch_idx / len(train_loader):.0f}%)]\tLoss: {loss.item():.6f}')
    # 驗證模型並記錄驗證損失和準確率
    model.eval()
    val_loss = 0
    correct = 0
    with torch.no_grad():  # 不計算梯度
        for data, target in val_loader:
            data, target = data.to(device), target.to(device)  # 將資料和目標移動到 GPU
            output = model(data)
            val_loss += criterion(output, target).item()  # 累加驗證損失
            pred = output.argmax(dim=1, keepdim=True)  # 擷取預測標籤
            correct += pred.eq(target.view_as(pred)).sum().item()  # 累加正確預測的數量
    val_loss /= len(val_loader)  # 計算平均驗證損失
    val_accuracy = 100. * correct / len(val_loader.dataset)  # 計算驗證準確率
    print(f'Validation Loss: {val_loss:.4f}, Accuracy: {correct}/{len(val_loader.dataset)} ({val_accuracy:.0f}%)')
    # 記錄驗證損失和準確率到 TensorBoard
    writer.add_scalar('Loss/validation', val_loss, epoch)
    writer.add_scalar('Accuracy/validation', val_accuracy, epoch)
    # 儲存驗證準確率最高的模型
    if val_accuracy > best_val_accuracy:
        best_val_accuracy = val_accuracy
        torch.save(model.state_dict(), '/mnt/data/output/best_model.pth')
        print(f'Model saved with accuracy: {best_val_accuracy:.2f}%')
# 關閉 SummaryWriter
writer.close()
print('Training complete. writer.close()')

上傳代碼。在資料集詳情頁，單擊查看数据跳轉至OSS控制台。然後單擊上傳檔案 >掃描檔案 > 上傳文件，將訓練代碼上傳至OSS中。

3.2 建立DLC任務

在PAI控制台左側功能表列單擊分布式训练（DLC） > 新建任务。

配置DLC任務參數。關鍵參數配置如下，其他參數預設即可。全量參數請參見建立訓練任務。

镜像配置：選擇镜像地址，然後根據您所在地域填寫對應鏡像地址。

在控制台頂部導覽列的地區選取器中確認當前所在地區（如華東1（杭州））。

地區	對應鏡像地址
北京	dsw-registry-vpc.cn-beijing.cr.aliyuncs.com/pai/modelscope:1.28.0-pytorch2.3.1tensorflow2.16.1-gpu-py311-cu121-ubuntu22.04
上海	dsw-registry-vpc.cn-shanghai.cr.aliyuncs.com/pai/modelscope:1.28.0-pytorch2.3.1tensorflow2.16.1-gpu-py311-cu121-ubuntu22.04
杭州	dsw-registry-vpc.cn-hangzhou.cr.aliyuncs.com/pai/modelscope:1.28.0-pytorch2.3.1tensorflow2.16.1-gpu-py311-cu121-ubuntu22.04
其他	查詢地區ID，並替換鏡像地址中的<地區ID>擷取完整連結： dsw-registry-vpc.<地區ID>.cr.aliyuncs.com/pai/modelscope:1.28.0-pytorch2.3.1tensorflow2.16.1-gpu-py311-cu121-ubuntu22.04

該鏡像已在DSW 互動式建模快速入門中驗證沒有環境問題。使用PAI建模時，通常先在DSW中驗證環境、開發代碼，然後再使用DLC訓練。

数据集挂载：選擇自定义数据集，選擇上一步中建立的資料集。挂载路径預設/mnt/data。
启动命令：python /mnt/data/mnist_train.py
該啟動命令與在DSW或本地運行時相同。但由於mnist_train.py 現已掛載至 /mnt/data/，因此僅需要修改代碼的路徑為/mnt/data/mnist_train.py。
资源来源：選擇公共资源，资源规格選擇ecs.gn7i-c8g1.2xlarge即可。
如果該規格執行個體庫存不足，您也可以選擇其他GPU執行個體。

單擊确定建立任務，任務大約需要執行15分鐘。執行過程中可以單擊日誌查看訓練過程。

訓練日誌顯示MNIST模型在Epoch 18驗證準確率達到9886/10000（約99%），任務狀態為已成功，最終輸出Training complete. writer.close()。

執行完成後，會在掛載資料集的output路徑下輸出最佳的模型檢查點，以及TensorBoard日誌。

OSS Bucket 的 dlc_mnist/ 目錄下包含 dataSet/ 子目錄（訓練資料）、output/ 子目錄（模型輸出）以及 mnist_train.py 訓練指令檔。

3.3 查看Tensorboard（可選）

您可以藉助視覺化檢視TensorBoard查看loss曲線，瞭解訓練的具體情況。

重要

DLC任務如果想使用TensorBoard，必須配置資料集。

單擊DLC任務詳情頁上方的Tensorboard > 新建Tensorboard。
配置类型選擇按任务，在Summary目录處填寫訓練代碼中Summary儲存的路徑：/mnt/data/output/runs/，單擊确定啟動。
對應程式碼片段：writer = SummaryWriter('/mnt/data/output/runs/mnist_experiment')
單擊查看Tensorboard查看train_loss曲線（反映訓練集損失）與 validation_loss曲線（反映驗證集損失）。
（可選）根據loss映像，調整超參數，提升模型效果
您可以根據損失值的變化趨勢，初步判斷當前模型的訓練效果：
- 在結束訓練前 train_loss 與 validation_loss 仍有下降趨勢（欠擬合）
  您可以增加 num_epochs（訓練輪次，與訓練深度正相關），或適當增大 learning_rate 後再進行訓練，加大模型的對訓練資料的擬合程度；
- 在結束訓練前 train_loss 持續下降，validation_loss 開始變大（過擬合）
  您可以減少 num_epochs，或適當減小 learning_rate 後再進行訓練，防止模型過度訓練；
- 在結束訓練前 train_loss 與 validation_loss 均處於平穩狀態（良好擬合）
  模型處於該狀態時，您可以進行後續步驟。

3.4 部署訓練後的模型

詳情請參見五. 使用 EAS 部署模型服務。

四. 單機多卡或多機多卡分布式訓練

當單個GPU的顯存無法滿足訓練需求，或者想要加快訓練速度時，您可以建立單機多卡或多機多卡的分布式訓練任務。

本文以使用2台各有1個GPU的執行個體為例進行說明，該樣本同樣適用於其他配置的單機多卡或多機多卡訓練。

4.1 建立資料集

如果您在三. 單機單卡訓練時已經建立了資料集，只需單擊下載代碼mnist_train_distributed.py並上傳。否則請先3.1 建立資料集，再上傳該代碼。

單機多卡或多機多卡訓練程式碼範例 mnist_train_distributed.py

import os
import torch
import torch.distributed as dist
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import DataLoader
from torch.utils.data.distributed import DistributedSampler
from torchvision import datasets, transforms
from torch.utils.tensorboard import SummaryWriter
class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=3)
        self.fc1 = nn.Linear(20 * 5 * 5, 128)
        self.fc2 = nn.Linear(128, 10)
    def forward(self, x):
        x = F.max_pool2d(F.relu(self.conv1(x)), 2)
        x = F.max_pool2d(F.relu(self.conv2(x)), 2)
        x = x.view(-1, 20 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x
def main():
    rank = int(os.environ["RANK"])
    world_size = int(os.environ["WORLD_SIZE"])
    local_rank = int(os.environ["LOCAL_RANK"])
    dist.init_process_group(backend='nccl')
    torch.cuda.set_device(local_rank)
    device = torch.device('cuda', local_rank)
    batch_size = 64
    learning_rate = 0.01
    num_epochs = 20
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.5,), (0.5,))
    ])
    # 只有主進程（rank=0）需要下載，其他進程需要等待它完成
    # 讓 rank!=0 的進程先在屏障處等待
    if rank != 0:
        dist.barrier()
    # 所有進程都執行資料集的建立
    # 但只有 rank=0 的進程會實際執行下載
    train_dataset = datasets.MNIST(root='/mnt/data/dataSet', train=True, download=(rank == 0), transform=transform)
    # rank=0 進程下載完成後，也到達屏障處，從而釋放所有進程
    if rank == 0:
        dist.barrier()
    # 此處，所有進程都已同步，可以繼續執行後續代碼
    val_dataset = datasets.MNIST(root='/mnt/data/dataSet', train=False, download=False, transform=transform)
    train_sampler = DistributedSampler(train_dataset, num_replicas=world_size, rank=rank, shuffle=True)
    train_loader = DataLoader(train_dataset, batch_size=batch_size, sampler=train_sampler, num_workers=4, pin_memory=True)
    val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=4, pin_memory=True)
    model = SimpleCNN().to(device)
    model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[local_rank])
    criterion = nn.CrossEntropyLoss().to(device)
    optimizer = optim.SGD(model.parameters(), lr=learning_rate)
    if rank == 0:
        writer = SummaryWriter('/mnt/data/output_distributed/runs/mnist_experiment')
    best_val_accuracy = 0.0
    for epoch in range(num_epochs):
        train_sampler.set_epoch(epoch)
        model.train()
        for batch_idx, (data, target) in enumerate(train_loader):
            data, target = data.to(device, non_blocking=True), target.to(device, non_blocking=True)
            optimizer.zero_grad()
            output = model(data)
            loss = criterion(output, target)
            loss.backward()
            optimizer.step()
            if batch_idx % 100 == 0:
                # 每個 rank 和 local_rank 都列印自己的 loss
                print(f"Rank: {rank}, Local_Rank: {local_rank} -- Train Epoch: {epoch} "
                      f"[{batch_idx * len(data) * world_size}/{len(train_loader.dataset)} "
                      f"({100. * batch_idx / len(train_loader):.0f}%)]\tLoss: {loss.item():.6f}")
                if rank == 0:
                    writer.add_scalar('Loss/train', loss.item(), epoch * len(train_loader) + batch_idx)
        # 驗證
        model.eval()
        val_loss = 0
        correct = 0
        total = 0
        with torch.no_grad():
            for data, target in val_loader:
                data, target = data.to(device, non_blocking=True), target.to(device, non_blocking=True)
                output = model(data)
                val_loss += criterion(output, target).item() * data.size(0)
                pred = output.argmax(dim=1, keepdim=True)
                correct += pred.eq(target.view_as(pred)).sum().item()
                total += target.size(0)
        val_loss_tensor = torch.tensor([val_loss], dtype=torch.float32, device=device)
        correct_tensor = torch.tensor([correct], dtype=torch.float32, device=device)
        total_tensor = torch.tensor([total], dtype=torch.float32, device=device)
        dist.all_reduce(val_loss_tensor, op=dist.ReduceOp.SUM)
        dist.all_reduce(correct_tensor, op=dist.ReduceOp.SUM)
        dist.all_reduce(total_tensor, op=dist.ReduceOp.SUM)
        val_loss = val_loss_tensor.item() / total_tensor.item()
        val_accuracy = 100. * correct_tensor.item() / total_tensor.item()
        if rank == 0:
            print(f'Validation Loss: {val_loss:.4f}, Accuracy: {int(correct_tensor.item())}/{int(total_tensor.item())} ({val_accuracy:.0f}%)')
            writer.add_scalar('Loss/validation', val_loss, epoch)
            writer.add_scalar('Accuracy/validation', val_accuracy, epoch)
            if val_accuracy > best_val_accuracy:
                best_val_accuracy = val_accuracy
                torch.save(model.module.state_dict(), '/mnt/data/output_distributed/best_model.pth')
                print(f'Model saved with accuracy: {best_val_accuracy:.2f}%')
    if rank == 0:
        writer.close()
    dist.destroy_process_group()
    if rank == 0:
        print('Training complete. writer.close()')
if __name__ == "__main__":
    main()

4.2 建立DLC任務

在PAI控制台左側功能表列單擊分布式训练（DLC） > 新建任务。

配置DLC任務參數。關鍵參數配置如下，其他參數預設即可。全量參數請參見建立訓練任務。

镜像配置：選擇镜像地址，然後根據您所在地域填寫對應鏡像地址。

地區	鏡像地址
北京	dsw-registry-vpc.cn-beijing.cr.aliyuncs.com/pai/modelscope:1.28.0-pytorch2.3.1tensorflow2.16.1-gpu-py311-cu121-ubuntu22.04
上海	dsw-registry-vpc.cn-shanghai.cr.aliyuncs.com/pai/modelscope:1.28.0-pytorch2.3.1tensorflow2.16.1-gpu-py311-cu121-ubuntu22.04
杭州	dsw-registry-vpc.cn-hangzhou.cr.aliyuncs.com/pai/modelscope:1.28.0-pytorch2.3.1tensorflow2.16.1-gpu-py311-cu121-ubuntu22.04
其他	查詢地區ID，並替換鏡像地址中的<地區ID>擷取完整連結： dsw-registry-vpc.<地區ID>.cr.aliyuncs.com/pai/modelscope:1.28.0-pytorch2.3.1tensorflow2.16.1-gpu-py311-cu121-ubuntu22.04

該鏡像已在DSW 互動式建模快速入門中驗證沒有環境問題。使用PAI建模時，通常先在DSW中驗證環境及代碼，然後再使用DLC訓練。

数据集挂载：選擇自定义数据集，並選擇上一步中建立的資料集。挂载路径預設/mnt/data。
启动命令：torchrun --nproc_per_node=1 --nnodes=${WORLD_SIZE} --node_rank=${RANK} --master_addr=${MASTER_ADDR} --master_port=${MASTER_PORT} /mnt/data/mnist_train_distributed.py
DLC會自動注入MASTER_ADDR、WORLD_SIZE等通用環境變數，通過$環境變數名來擷取。
资源来源：選擇公共资源，节点数量為2，资源规格選擇ecs.gn7i-c8g1.2xlarge。
如果該規格執行個體庫存不足，您也可以選擇其他GPU執行個體。

單擊确定建立任務，任務大約需要執行10分鐘。執行過程中可以在概览頁面，查看兩台執行個體的訓練日志。

執行完成後，會在掛載資料集的output_distributed路徑下輸出最佳的模型檢查點，以及TensorBoard日誌。

4.3 查看Tensorbord（可選）

您可以藉助視覺化檢視TensorBoard查看loss曲線，瞭解訓練的具體情況。

重要

DLC任務如果想使用TensorBoard，必須配置資料集。

單擊DLC任務詳情頁上方的Tensorboard > 新建Tensorboard。
配置类型選擇按任务，在Summary目录處填寫訓練代碼中Summary儲存的路徑：/mnt/data/output_distributed/runs，單擊确定啟動。
對應程式碼片段：writer = SummaryWriter('/mnt/data/output_distributed/runs/mnist_experiment')
單擊查看Tensorboard查看train_loss曲線（反映訓練集損失）與 validation_loss曲線（反映驗證集損失）。
（可選）根據loss映像，調整超參數，提升模型效果
您可以根據損失值的變化趨勢，初步判斷當前模型的訓練效果：
- 在結束訓練前 train_loss 與 validation_loss 仍有下降趨勢（欠擬合）
  您可以增加 num_epochs（訓練輪次，與訓練深度正相關），或適當增大 learning_rate 後再進行訓練，加大模型的對訓練資料的擬合程度；
- 在結束訓練前 train_loss 持續下降，validation_loss 開始變大（過擬合）
  您可以減少 num_epochs，或適當減小 learning_rate 後再進行訓練，防止模型過度訓練；
- 在結束訓練前 train_loss 與 validation_loss 均處於平穩狀態（良好擬合）
  模型處於該狀態時，您可以進行後續步驟。

4.4 部署訓練後的模型

詳情請參見五. 使用 EAS 部署模型服務。

五. 相關文檔

更多DLC功能請參見分布式訓練DLC。