本文目录导读：
我来详细介绍如何使用Horovod进行Python脚本的分布式训练。
基础安装
# 安装Horovod pip install horovod # 或根据GPU支持安装 HOROVOD_GPU_CUDA_HOME=/usr/local/cuda pip install horovod
基本使用模板（TensorFlow示例）
import torch
import horovod.torch as hvd
import tensorflow as tf
import horovod.tensorflow as hvd
# 初始化Horovod
hvd.init()
# 获取当前进程信息
rank = hvd.rank() # 当前进程ID
local_rank = hvd.local_rank() # 本地进程ID
size = hvd.size() # 总进程数
# 设置GPU
if torch.cuda.is_available():
torch.cuda.set_device(hvd.local_rank())
# 或 TensorFlow
gpus = tf.config.experimental.list_physical_devices('GPU')
for gpu in gpus:
tf.config.experimental.set_memory_growth(gpu, True)
if gpus:
tf.config.experimental.set_visible_devices(gpus[hvd.local_rank()], 'GPU')
PyTorch完整示例
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
import horovod.torch as hvd
def main():
# 1. 初始化
hvd.init()
# 2. 设置设备
device = torch.device(f'cuda:{hvd.local_rank()}' if torch.cuda.is_available() else 'cpu')
# 3. 创建模型
model = torchvision.models.resnet18()
model = model.to(device)
# 4. 包装模型（广播初始参数）
model = hvd.BroadcastGlobalVariablesWrapper(0, args=(model,))
# 5. 创建优化器并包装
optimizer = optim.SGD(model.parameters(), lr=0.01 * hvd.size())
optimizer = hvd.DistributedOptimizer(
optimizer,
named_parameters=model.named_parameters()
)
# 6. 数据加载（每个进程有不同数据）
train_dataset = torchvision.datasets.CIFAR10(
root='./data',
train=True,
transform=transforms.ToTensor(),
download=True
)
# 分布式数据采样器
train_sampler = torch.utils.data.distributed.DistributedSampler(
train_dataset,
num_replicas=hvd.size(),
rank=hvd.rank()
)
train_loader = torch.utils.data.DataLoader(
train_dataset,
batch_size=32,
sampler=train_sampler,
num_workers=2
)
# 7. 训练循环
for epoch in range(10):
train_sampler.set_epoch(epoch)
for batch_idx, (data, target) in enumerate(train_loader):
data, target = data.to(device), target.to(device)
optimizer.zero_grad()
output = model(data)
loss = nn.CrossEntropyLoss()(output, target)
loss.backward()
optimizer.step()
if batch_idx % 100 == 0 and hvd.rank() == 0:
print(f'Epoch: {epoch} Batch: {batch_idx} Loss: {loss.item():.4f}')
if __name__ == '__main__':
main()
运行命令
# 单机多GPU
horovodrun -np 4 -H localhost:4 python train.py
# 多机训练（假设4台机器，每台4个GPU）
horovodrun -np 16 -H host1:4,host2:4,host3:4,host4:4 python train.py
# 使用mpirun
mpirun -np 4 -H localhost:4 python train.py
# 使用OpenMPI
mpirun --allow-run-as-root -np 4 -bind-to none -map-by slot \
-x NCCL_DEBUG=INFO \
-x LD_LIBRARY_PATH \
-x PATH \
python train.py
分布式优化器配置
# 不同优化器类型
import horovod.torch as hvd
# TensorFlow
opt = tf.optimizers.Adam(0.001)
opt = hvd.DistributedOptimizer(opt)
# PyTorch
optimizer = optim.SGD(model.parameters(), lr=0.01)
optimizer = hvd.DistributedOptimizer(
optimizer,
named_parameters=model.named_parameters(),
compression=hvd.Compression.none, # 可选: fpure, fp16等
backward_passes_per_step=1,
op=hvd.Average # 可选: Sum, Average, Adasum
)
高级配置
import horovod.torch as hvd
# 1. 梯度压缩
optimizer = hvd.DistributedOptimizer(
optimizer,
compression=hvd.Compression.fp16 # 使用FP16压缩
)
# 2. 异步训练
hvd.Async()
# 3. 梯度累积
optimizer = hvd.DistributedOptimizer(
optimizer,
backward_passes_per_step=4 # 累积4步梯度再更新
)
# 4. 自定义广播
hvd.broadcast_parameters(model.state_dict(), root_rank=0)
hvd.broadcast_optimizer_state(optimizer, root_rank=0)
# 5. 进度监控
if hvd.rank() == 0:
# 只在主进程打印
print(f"Epoch {epoch}, Loss: {loss:.4f}")
# 6. 模型保存
if hvd.rank() == 0:
torch.save(model.state_dict(), 'model.pth')
性能优化建议
# 1. 设置NCCL相关环境变量
import os
os.environ['NCCL_SOCKET_IFNAME'] = 'eth0'
os.environ['NCCL_IB_DISABLE'] = '1'
os.environ['NCCL_DEBUG'] = 'INFO'
# 2. 调整批大小（总批大小不变）
batch_size_per_gpu = TOTAL_BATCH_SIZE // hvd.size()
# 3. 使用梯度累积代替减小批大小
accumulation_steps = 4
optimizer.zero_grad()
for i, (data, target) in enumerate(train_loader):
output = model(data)
loss = criterion(output, target) / accumulation_steps
loss.backward()
if (i + 1) % accumulation_steps == 0:
optimizer.step()
optimizer.zero_grad()
# 4. 数据加载优化
train_loader = DataLoader(
dataset,
batch_size=batch_size,
num_workers=4, # 增加worker数
pin_memory=True, # 固定内存加速
prefetch_factor=2 # 预取数据
)
常见问题排查
# 1. 网络问题
export NCCL_DEBUG=INFO
export NCCL_SOCKET_IFNAME=eth0
# 2. 显存不足
export PYTORCH_NO_CUDA_MEMORY_CACHING=1
# 3. 设置超时
export HOROVOD_TIMEOUT=300 # 秒
# 4. 测试horovod是否正常
horovodrun -np 2 python -c "import horovod.torch as hvd; hvd.init(); print(f'Rank {hvd.rank()} of {hvd.size()}');"
监控和调试
# 状态检查
if hvd.rank() == 0:
print(f"World size: {hvd.size()}")
print(f"Local rank: {hvd.local_rank()}")
# 同步检查点
if hvd.rank() == 0:
torch.save(model.state_dict(), 'checkpoint.pth')
# 所有进程等待
hvd.allreduce(torch.tensor(1.0))
# 集合操作
tensor = torch.tensor([1.0, 2.0]).cuda()
avg = hvd.allreduce(tensor, op=hvd.Average)
完整的生产环境脚本示例
# train_horovod.py
import argparse
import os
import horovod.torch as hvd
import torch
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--batch-size', type=int, default=32)
parser.add_argument('--epochs', type=int, default=10)
parser.add_argument('--lr', type=float, default=0.01)
args = parser.parse_args()
# 初始化
hvd.init()
# 根据总进程数调整学习率
lr_scaled = args.lr * hvd.size()
# 数据并行
batch_size_per_rank = args.batch_size // hvd.size()
# 训练逻辑
print(f"Training with {hvd.size()} processes, "
f"batch size per GPU: {batch_size_per_rank}, "
f"learning rate: {lr_scaled}")
if __name__ == '__main__':
main()
使用时根据您的具体框架（TensorFlow/PyTorch/MXNet）和需求选择相应的配置即可。
