异构计算训练：混合不同型号 GPU 与 CPU 的训练

简介

异构计算训练指在同一训练任务中同时使用不同类型或不同型号的计算设备（如不同代次的GPU、不同厂商的GPU、甚至CPU），以充分利用集群中多样化的硬件资源。这种策略能够降低硬件采购成本、绕过单设备显存/内存瓶颈，并加速模型收敛。本教程将带你从零搭建一个支持混合GPU与CPU的训练环境，并给出可落地的PyTorch代码示例。

为什么需要异构训练

• 硬件复用：企业常有余量硬件，如旧款V100与新款A100并存，异构训练让它们协作而非闲置。
• 显存突破：大模型单卡放不下时，可将不同层切分到不同设备，用更小显存卡训练更大模型。
• 成本控制：CPU内存容量大、成本低，可将嵌入表、预处理等放置于CPU，GPU专注密集计算。
• 弹性扩展：动态加入不同性能设备，避免“木桶效应”。

异构训练的技术挑战

• 计算负载不均：快设备等慢设备，导致整体吞吐下降。
• 通信开销：跨设备、跨主机传输数据时延高，需最小化传输量。
• 框架支持差异：原生DataParallel只支持同构GPU，需使用分布式API或模型并行策略。
• 调试复杂：设备放置错误可能导致性能骤降或静默错误。

核心概念：设备放置与通信后端

设备放置（Device Placement）

PyTorch中通过 torch.device 指定张量位置，模型的不同部分可显式发送到不同设备：

# 将模型前三层放在GPU1，后面放在GPU2
model.layer1.to('cuda:0')
model.layer2.to('cuda:1')

框架会自动追踪跨设备张量的依赖，插入必要的拷贝操作，但这可能产生隐式同步，需谨慎设计。

通信后端

跨设备通信依赖集合通信库：

• NCCL：NVIDIA出品，专为GPU间通信优化，支持多机多卡，推荐用于GPU。
• Gloo：适用于CPU和GPU，跨平台性好，但GPU间吞吐不如NCCL。
• MPI：传统HPC，可配合CPU集群使用。

PyTorch分布式包用字符串指定后端：dist.init_process_group(backend='nccl')。

环境准备与驱动要求

• GPU驱动版本需支持CUDA Toolkit，建议所有GPU使用统一CUDA版本（通常驱动向后兼容）。
• 安装相同版本的PyTorch，推荐预编译CUDA版本与硬件匹配。
• 若混合不同型号GPU（如Tesla T4与A10），确保NCCL版本兼容，可通过设置 NCCL_SOCKET_IFNAME 环境变量控制通信网卡。

# 检查CUDA可用设备
python -c "import torch; print(torch.cuda.device_count())"

实战：PyTorch混合不同型号GPU训练

我们使用 Distributed Data Parallel (DDP) 与模型并行结合，将模型切成两段，分别分配到不同GPU，并用流水线方式减少等待。

场景设定

• GPU 0：NVIDIA RTX 3080 (10GB显存)
• GPU 1：NVIDIA RTX 3060 (12GB显存)
• 模型：简易Transformer，嵌入层+12层编码器+输出层。

由于3080算力强但显存小，我们将嵌入层和最后两层放在3060，中间计算密集型层放在3080，平衡负载。

模型切分与Pipeline包装

import torch
import torch.nn as nn
import torch.distributed as dist
from torch.distributed.pipeline.sync import Pipe

# 假设Transformer模型返回阶段列表
class Stage0(nn.Module):
    """嵌入 + 前4个编码器层"""
    def __init__(self, vocab_size, d_model):
        super().__init__()
        self.embed = nn.Embedding(vocab_size, d_model)
        self.encoder_layers = nn.TransformerEncoderLayer(d_model, nhead=8, batch_first=True)
        self.encoders = nn.TransformerEncoder(self.encoder_layers, num_layers=4)
    def forward(self, x):
        return self.encoders(self.embed(x))

class Stage1(nn.Module):
    """中间5-8编码器层"""
    def __init__(self, d_model):
        super().__init__()
        layer = nn.TransformerEncoderLayer(d_model, nhead=8, batch_first=True)
        self.encoders = nn.TransformerEncoder(layer, num_layers=4)
    def forward(self, x):
        return self.encoders(x)

class Stage2(nn.Module):
    """最后4层 + 输出层"""
    def __init__(self, d_model, vocab_size):
        super().__init__()
        layer = nn.TransformerEncoderLayer(d_model, nhead=8, batch_first=True)
        self.encoders = nn.TransformerEncoder(layer, num_layers=4)
        self.out = nn.Linear(d_model, vocab_size)
    def forward(self, x):
        return self.out(self.encoders(x))

# 包装为Pipe需要按阶段顺序构造module列表
model = nn.Sequential(
    Stage0(vocab_size=30000, d_model=512),
    Stage1(d_model=512),
    Stage2(d_model=512, vocab_size=30000)
)

# 配置设备映射：阶段0 -> cuda:0 (3060), 阶段1 -> cuda:1 (3080), 阶段2 -> cuda:0
devices = ['cuda:0', 'cuda:1', 'cuda:0']
model = Pipe(model, chunks=8, devices=devices)  # chunks控制微批次数量

Pipe 会自动将每个阶段放到指定设备，并通过微批次流水线隐藏通信延迟。

数据加载与训练循环

from torch.utils.data import DataLoader, Dataset

# 模拟数据集
class DummyDataset(Dataset):
    def __len__(self): return 1000
    def __getitem__(self, idx):
        return torch.randint(0, 30000, (128,)), torch.randint(0, 30000, (128,))

dataset = DummyDataset()
loader = DataLoader(dataset, batch_size=32, shuffle=True)

# 初始化分布式（即使单机多卡也建议用DDP环境）
dist.init_process_group(backend='nccl', init_method='tcp://127.0.0.1:29500', rank=0, world_size=1)
torch.cuda.set_device(0)  # 主设备

optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)

model.train()
for epoch in range(5):
    for src, tgt in loader:
        # Pipe要求输入和目标在同一设备（阶段0所在设备）
        src, tgt = src.to('cuda:0'), tgt.to('cuda:0')
        output = model(src)
        loss = nn.functional.cross_entropy(output.view(-1, 30000), tgt.view(-1))
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    print(f'Epoch {epoch} loss: {loss.item():.4f}')

负载均衡调优

• 通过 chunks 参数控制微批次数，较大的chunks可提高GPU利用率，但增加显存。
• 可统计每个阶段的执行时间，用 torch.cuda.synchronize() 测量，手动调整各阶段层数或设备分配。
• 若设备算力差异大，可将更厚的一层切分到更快GPU，例如3080上放置6层，3060上2层。

混合CPU参与训练

当某些操作（如嵌入查找、数据预处理、自定义损失）在CPU上更经济时，可把模型子图放在CPU，并通过NCCL（仅GPU）或Gloo后端与GPU通信。

示例：大型嵌入表离线存放在CPU

class HybridModel(nn.Module):
    def __init__(self, num_embeddings, embedding_dim):
        super().__init__()
        # 嵌入表常驻CPU，可节省数十GB显存
        self.emb = nn.Embedding(num_embeddings, embedding_dim, sparse=True).to('cpu')
        self.lstm = nn.LSTM(embedding_dim, 256, batch_first=True).to('cuda')
        self.fc = nn.Linear(256, 10).to('cuda')

    def forward(self, x):
        # x在CPU上查询嵌入
        emb = self.emb(x.cpu())
        # 只将必要的嵌入向量移至GPU
        out, _ = self.lstm(emb.to('cuda'))
        return self.fc(out)

model = HybridModel(100000, 128)
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

for data, target in dataloader:
    data, target = data.to('cpu'), target.to('cuda')  # 标签在GPU
    output = model(data)          # 前向自动跨设备
    loss = nn.functional.cross_entropy(output, target)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

注意：CPU上的参数更新通过优化器自动处理，但CPU->GPU的每一次前向都会触发数据传输，需确保计算量远大于传输开销。

性能优化要点

减少跨设备同步

• 避免在训练循环中调用 torch.cuda.synchronize() 或 .item()。
• 使用非阻塞拷贝：tensor.to(device, non_blocking=True)。

流水线micro-batch大小

增大chunks降低流水线空泡，但会增加显存占用。可通过实验找到最佳值。

梯度累加

小显存设备可能无法承载常规batch size，用梯度累加模拟更大批量：

accumulation_steps = 4
for i, (data, target) in enumerate(loader):
    output = model(data)
    loss = criterion(output, target) / accumulation_steps
    loss.backward()
    if (i+1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

带宽感知的任务分配

使用 torch.cuda.get_device_properties() 获取设备算力和显存带宽，将带宽敏感操作（如大矩阵乘法）放在高带宽设备。

常见问题与排查

Q：混合GPU训练时速度反而不如单卡？ A：检查是否因小批次导致流水线空泡严重，或设备间通信成为瓶颈。尝试调大微批次，并确保使用NCCL后端。

Q：RuntimeError: Expected all tensors to be on the same device A：检查损失函数和目标张量的设备，确保它们在同一设备，或使用 loss.to(device)。

Q：CPU参与训练时梯度更新异常缓慢？ A：稀疏嵌入（sparse=True）可加速，或使用Adam等自适应优化器，或将CPU参数单独设置更大学习率。

Q：不同型号GPU显存不对等导致OOM？ A：将大显存卡放置显存密集层（如注意力），或将嵌入表切分到多个设备（模型并行）。

总结

异构计算训练不是妥协，而是一种充分利用现有资源、降低成本的先进训练模式。通过PyTorch的Pipe、设备放置与DDP，你可以在混合GPU/CPU集群上训练大型模型。核心在于量化每个设备的算力与显存，合理切分计算图，并最大化通信与计算的重叠。动手尝试本教程的代码，你就能让异构硬件协同工作，释放闲置算力。