【深度学习】使用硬件加速模型训练速度
一、 单机器单GPU
特点
- 配置简单:无需多 GPU 或分布式环境的复杂配置,适合资源有限的场景。
- 适合小规模模型:对于计算量不大的模型(如中小型 CNN、RNN),单 GPU 可以处理大多数常见任务。
优势
- 简单易用:适合深度学习初学者以及较小规模的实验。
- 无需并行配置:仅需设置一个 GPU,代码和数据同步操作也更简单。
实现方法
- 常用框架:可以直接在代码中指定 GPU 使用,如 TensorFlow 的
with tf.device('/GPU:0')
或 PyTorch 的model.to('cuda')
。 - 适用场景:适合开发和调试中小规模模型,以及在本地进行初步实验。
PyTorch 示例:
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset
from torch.cuda.amp import autocast, GradScaler
# 检查是否有 GPU 可用
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 示例模型(简单的多层感知机)
class SimpleModel(nn.Module):
def __init__(self):
super(SimpleModel, self).__init__()
self.fc1 = nn.Linear(10, 50)
self.fc2 = nn.Linear(50, 1)
def forward(self, x):
x = torch.relu(self.fc1(x))
x = self.fc2(x)
return x
# 创建模型并将其加载到 GPU
model = SimpleModel().to(device)
# 创建数据集和数据加载器
X = torch.randn(1000, 10)
y = torch.randn(1000, 1)
dataset = TensorDataset(X, y)
dataloader = DataLoader(dataset, batch_size=32, shuffle=True, num_workers=4, pin_memory=True)
# 定义损失函数和优化器
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)
# 训练循环
epochs = 5
for epoch in range(epochs):
# 使用 PyTorch 的 AMP 自动混合精度训练
scaler = GradScaler()
for inputs, targets in dataloader:
inputs, targets = inputs.to(device), targets.to(device)
with autocast():
outputs = model(inputs)
loss = criterion(outputs, targets)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
print(f"Epoch {epoch + 1}/{epochs}, Loss: {loss.item()}")
示例说明
-
dataloader = DataLoader(dataset, batch_size=32, shuffle=True, num_workers=4, pin_memory=True)
当
pin_memory=True
时,DataLoader
将把数据加载到页锁定内存中,使得从 CPU 到 GPU 的传输更加高效。
并发流程:- 批次准备:4 个 CPU 进程在
DataLoader
中并发从数据集加载数据,并执行数据增强、预处理,将数据固定在内存中。 - 数据传输:CPU 将批次数据通过 PCIe 通道传输到 GPU 的显存。
- GPU 执行计算:GPU 在显存中运行当前批次的前向传播和反向传播。
- 并发准备下一个批次:在 GPU 计算当前批次的过程中,CPU 上的
DataLoader
进程已经开始准备下一个批次数据。 - 重复以上过程:CPU 和 GPU 以此方式并行处理每个批次,避免 GPU 等待数据传输,提升整体训练速度。
- 批次准备:4 个 CPU 进程在
-
amp
(包含autocast
和GradScaler
):提供完整的混合精度训练方案,包含autocast
的精度选择和GradScaler
的梯度缩放,能保证训练的数值稳定性。
二、 单机器多GPU
在单台机器上配置多个 GPU,可以通过数据并行、模型并行和管道并行等策略实现多 GPU 训练。
1. 数据并行(Data Parallelism)
数据并行是最常用的多 GPU 训练策略,将数据划分到不同的 GPU 上,并在每个 GPU 上独立地计算前向传播和反向传播。每个 GPU 得到的梯度会在同步步骤中合并,更新全局模型参数。
- 工作流程:
- 将批次数据划分到不同的 GPU 上。
- 每个 GPU 独立计算前向和反向传播,生成局部梯度。
- 汇总所有 GPU 的梯度,更新全局模型参数。
- 适用场景:适合数据量大、模型规模适中的场景,是目前最常用的并行训练方式。
代码示例
PyTorch提供了多种方式来实现多GPU训练,常用的方法是使用 torch.nn.DataParallel
或 torch.nn.parallel.DistributedDataParallel
。
-
- 使用
torch.nn.DataParallel
- 使用
DataParallel
是一种简单直接的多GPU训练方法,适合小规模的多GPU训练任务。它会自动将数据分配到多个GPU上,并将每个GPU的输出合并。
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
# 定义模型
model = MyModel()
# 将模型包装成多GPU模型
if torch.cuda.device_count() > 1:
print(f"使用 {torch.cuda.device_count()} 个 GPU 进行训练")
model = nn.DataParallel(model)
# 将模型移动到GPU
model = model.cuda()
# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters())
# 加载数据
dataset = MyDataset()
dataloader = DataLoader(dataset, batch_size=32, shuffle=True)
# 训练循环
for epoch in range(num_epochs):
for data, target in dataloader:
data, target = data.cuda(), target.cuda()
# 前向传播
output = model(data)
loss = criterion(output, target)
# 反向传播和优化
optimizer.zero_grad()
loss.backward()
optimizer.step()
print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item()}")
-
- 使用
torch.nn.parallel.DistributedDataParallel
- 使用
DistributedDataParallel
(DDP) 是一种更高效的分布式训练方式,推荐在多GPU和分布式训练任务中使用,尤其适合更大规模的训练。
在单机器多GPU的情况下,也可以用DDP,它会在每个GPU上启动一个进程来进行并行计算。以下是使用DDP的代码示例。
设置DDP的关键步骤
- 设置进程组:使用
torch.multiprocessing
来启动多个进程,每个进程对应一个GPU。 - 初始化DDP模型:用
DistributedDataParallel
包装模型。 - 数据划分:使用
DistributedSampler
来分配数据到不同的进程。 - 运行代码:需通过
torch.multiprocessing.spawn
启动多进程。或者直接使用torchrun
命令启动脚本python -m torchrun --nproc_per_node=NUM_GPUS_YOU_HAVE your_script.py --arg1 --arg2
import torch
import torch.nn as nn
import torch.multiprocessing as mp
import torch.distributed as dist
from torch.utils.data import DataLoader, DistributedSampler
# 初始化进程函数
def train(rank, world_size):
# 初始化进程组
dist.init_process_group("nccl", rank=rank, world_size=world_size)
torch.cuda.set_device(rank)
# 定义模型,并移动到相应GPU
model = MyModel().to(rank)
model = nn.parallel.DistributedDataParallel(model, device_ids=[rank])
# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss().to(rank)
optimizer = torch.optim.Adam(model.parameters())
# 创建数据集和分布式采样器
dataset = MyDataset()
sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank, shuffle=True)
dataloader = DataLoader(dataset, batch_size=32, sampler=sampler)
# 训练循环
for epoch in range(num_epochs):
sampler.set_epoch(epoch) # 每个epoch设定不同的采样
for data, target in dataloader:
data, target = data.to(rank), target.to(rank)
# 前向传播
output = model(data)
loss = criterion(output, target)
# 反向传播和优化
optimizer.zero_grad()
loss.backward()
optimizer.step()
if rank == 0: # 只在主进程上打印
print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item()}")
# 清理进程组
dist.destroy_process_group()
# 主函数
if __name__ == "__main__":
world_size = torch.cuda.device_count() # 获取可用GPU数量
mp.spawn(train, args=(world_size,), nprocs=world_size, join=True)
解释
train
函数中包含了每个进程需要执行的代码,rank
表示每个GPU对应的编号。DistributedDataParallel
在每个GPU上创建一个独立的模型副本,且device_ids
参数指定了每个模型副本运行的设备。DistributedSampler
将数据分配给每个进程,以确保在多进程中数据不会重复。
2. 模型并行(Model Parallelism)
模型并行是指将模型的不同部分分配到不同的 GPU 上,每个 GPU 只负责一部分模型的计算。这种方法适合非常大的模型,单个 GPU 无法容纳所有参数时使用。
- 工作流程:
- 将模型划分成多个子部分(如按层或模块)。
- 数据在 GPU 间传递,每个 GPU 负责一个模型部分的前向和反向传播。
- 适用场景:适合超大规模模型(如 GPT-3),在单个 GPU 内存不足以容纳全部参数时使用。
3. 管道并行(Pipeline Parallelism)
管道并行是一种结合数据并行和模型并行的方式,将模型分成多个阶段,数据分批次流经每个阶段的不同 GPU,从而在多个 GPU 之间形成“流水线”。
- 工作流程:
- 将模型按层次或功能模块划分成多个阶段,每个阶段在不同的 GPU 上。
- 数据逐批通过每个阶段的 GPU,同时保持并行计算。
- 适用场景:适合具有分层结构的模型,尤其在超大批次的训练任务中提升效率。
三、多机器多GPU
在分布式训练中,模型可以跨越多台机器(节点)进行训练,每个节点可以有一个或多个 GPU。分布式训练的核心在于如何在各节点之间有效地通信和同步梯度。
1. 分布式数据并行(Distributed Data Parallelism)
分布式数据并行是分布式训练中最常用的模式,它与单机器多GPU 的数据并行类似,但在不同机器的 GPU 间进行并行计算。每台机器上会有一个完整的模型副本,训练数据在所有机器间分布。
- 梯度同步:常用的梯度同步方法是全归约(All-Reduce),它在每个节点中同步并合并梯度更新。
- 实现框架:
- PyTorch Distributed Data Parallel (DDP):使用 PyTorch 自带的 DDP 模块,可实现多节点的分布式训练,适合在多 GPU 集群中。
- TensorFlow MirroredStrategy:TensorFlow 的
tf.distribute.MirroredStrategy
支持多 GPU、单机多 GPU 和多机多 GPU 的训练。 - Horovod:Uber 开发的分布式训练框架,支持 TensorFlow、Keras 和 PyTorch,使用环形 All-Reduce 方式进行高效的梯度同步。
- 适用场景:适合数据量大且可分割的大批量训练任务。
DDP示例
import os
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, DistributedSampler
from torchvision import datasets, transforms
# 定义模型
class MyModel(torch.nn.Module):
def __init__(self):
super(MyModel, self).__init__()
self.fc = torch.nn.Linear(784, 10)
def forward(self, x):
return self.fc(x.view(-1, 784))
# 训练函数
def train():
# 从环境变量中获取 rank 和 world_size
rank = int(os.environ["RANK"])
world_size = int(os.environ["WORLD_SIZE"])
# 初始化分布式进程组
dist.init_process_group(
backend="nccl",
init_method="env://",
rank=rank,
world_size=world_size
)
# 获取当前节点的 GPU 数量
local_rank = rank % torch.cuda.device_count()
torch.cuda.set_device(local_rank)
# 创建模型并移动到对应的 GPU 上
model = MyModel().to(local_rank)
model = DDP(model, device_ids=[local_rank])
# 定义损失函数和优化器
criterion = torch.nn.CrossEntropyLoss().to(local_rank)
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
# 数据加载器,使用 DistributedSampler 分配数据
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])
dataset = datasets.MNIST(root='./data', train=True, transform=transform, download=True)
sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank, shuffle=True)
dataloader = DataLoader(dataset, batch_size=64, sampler=sampler)
# 训练循环
for epoch in range(10):
sampler.set_epoch(epoch)
for batch_idx, (data, target) in enumerate(dataloader):
data, target = data.to(local_rank), target.to(local_rank)
# 前向传播
optimizer.zero_grad()
output = model(data)
loss = criterion(output, target)
# 反向传播和优化
loss.backward()
optimizer.step()
if batch_idx % 100 == 0 and rank == 0: # 仅主进程打印日志
print(f"Epoch [{epoch+1}/10], Step [{batch_idx}/{len(dataloader)}], Loss: {loss.item():.4f}")
# 清理进程组
dist.destroy_process_group()
if __name__ == "__main__":
train()
在使用 torchrun
启动分布式训练时,torchrun
会自动计算和设置 RANK
和 WORLD_SIZE
环境变量,这样每个进程都能知道其在分布式系统中的全局位置。这是 torchrun
的一个重要特性,帮助简化分布式训练的设置流程。
解释
torchrun
使用以下命令行参数来推断和设置这些环境变量:
--nnodes
:指定分布式系统中的总节点(机器)数量。--nproc_per_node
:指定每台机器上启动的进程数量(通常等于 GPU 数量)。--node_rank
:指定当前机器的节点编号(从0
开始),用于确定当前机器的RANK
起始位置。--master_addr
和--master_port
:指定主节点的 IP 地址和端口,确保所有节点可以连接到主节点进行通信。
torchrun
根据这些参数计算并设置 WORLD_SIZE
和 RANK
,每个进程启动时就能自动获得这些环境变量。
WORLD_SIZE
:全局进程数,即分布式系统中的总进程数。RANK
:每个进程的全局唯一编号,用来标识当前进程在整个分布式系统中的位置。- 每台机器上的
RANK
是按顺序递增的,从主节点开始。 - 例如,如果有两台机器,每台机器启动 4 个进程,那么
RANK
在主节点(node_rank=0
)上为0-3
,在第二台机器(node_rank=1
)上为4-7
。
- 每台机器上的
假设有两台机器,每台机器上分别有 4 个 GPU,配置如下:
-
- 主节点(机器 1)启动命令:
torchrun --nnodes=2 --nproc_per_node=4 --node_rank=0 --master_addr="<主节点IP地址>" --master_port="12355" script.py
-
WORLD_SIZE
自动设置为2 * 4 = 8
。 -
RANK
自动为0-3
,对应机器 1 上的 4 个进程。 -
- 第二台机器(机器 2)启动命令:
torchrun --nnodes=2 --nproc_per_node=4 --node_rank=1 --master_addr="<主节点IP地址>" --master_port="12355" script.py
WORLD_SIZE
仍为8
。RANK
自动为4-7
,对应机器 2 上的 4 个进程。
注意
torchrun 不支持不同的 nproc_per_node,我们需要手动启动每个机器的进程,确保每个进程的 RANK 和 WORLD_SIZE 环境变量设置正确。
假设我们有两台机器,且每台机器上 GPU 数量不同:机器 1有 4 个 GPU,机器 2有 2 个 GPU。
主节点(机器 1)启动命令:
# 启动第一个进程
RANK=0 torchrun --nproc_per_node=1 --nnodes=2 --node_rank=0 --master_addr="<主节点IP地址>" --master_port="12355" --world_size=6 script.py &
# 启动第二个进程
RANK=1 torchrun --nproc_per_node=1 --nnodes=2 --node_rank=0 --master_addr="<主节点IP地址>" --master_port="12355" --world_size=6 script.py &
# 启动第三个进程
RANK=2 torchrun --nproc_per_node=1 --nnodes=2 --node_rank=0 --master_addr="<主节点IP地址>" --master_port="12355" --world_size=6 script.py &
# 启动第四个进程
RANK=3 torchrun --nproc_per_node=1 --nnodes=2 --node_rank=0 --master_addr="<主节点IP地址>" --master_port="12355" --world_size=6 script.py &
第二台机器(机器 2)启动命令:
export WORLD_SIZE=6
# 启动第一个进程
RANK=4 torchrun --nproc_per_node=1 --nnodes=2 --node_rank=1 --master_addr="<主节点IP地址>" --master_port="12355" --world_size=6 script.py &
# 启动第二个进程
RANK=5 torchrun --nproc_per_node=1 --nnodes=2 --node_rank=1 --master_addr="<主节点IP地址>" --master_port="12355" --world_size=6 script.py &
2. 分布式模型并行(Distributed Model Parallelism)
分布式模型并行将模型拆分到不同节点的 GPU 上,通常用于单节点 GPU 无法容纳完整模型的超大规模模型(如 GPT-3)。
- 参数分配:模型参数按功能模块或层次划分到不同的 GPU。
- 前向与反向传播:每个 GPU 只计算自己负责的部分,然后将输出传递给下一层的 GPU。
- 适用场景:适用于超大模型的训练(如 Transformer、大规模语言模型等),尤其在内存受限时。
3. 混合并行(Hybrid Parallelism)
混合并行结合数据并行和模型并行的优势,用于训练规模极大的模型。可以同时在不同 GPU 上进行数据并行训练,并在每个 GPU 内部按模型并行分配参数。
- 实现框架:
- Megatron-LM:NVIDIA 提供的框架,支持混合并行训练超大规模 Transformer 模型。
- DeepSpeed:微软开发的分布式训练优化库,支持数据并行、模型并行、流水线并行以及零冗余优化。
- 适用场景:超大规模模型的训练,适合在分布式多 GPU 集群中实现。
四、TPU(张量处理单元)
TPU 是 Google 专门为深度学习设计的硬件加速器,尤其针对矩阵运算进行了优化。
TPU 的特点
- 高度优化的矩阵运算:TPU 专为深度学习而设计,适合卷积神经网络(CNN)和循环神经网络(RNN)等高密度计算任务。
- 高内存带宽:TPU 拥有很高的内存带宽,可以加速训练过程中数据的加载和计算。
使用 TPU 的方法
- Google Colab 和 Google Cloud Platform(GCP):可以在 Google Colab 免费使用 TPU,GCP 提供的 TPU 资源适合更大规模任务。
- TensorFlow 高效兼容:TPU 最适合 TensorFlow 框架的使用,Keras 和 TensorFlow 可以在单行代码中切换 TPU 支持。
TPU 使用场景
- Transformer 模型:如 BERT、GPT 等 NLP 模型。
- 大规模深度神经网络:适合需要更大数据处理能力的深度神经网络。