CentOS上PyTorch并行计算能力全景
在CentOS上,PyTorch可通过数据并行、分布式数据并行(DDP)、模型并行与流水线并行充分利用多核CPU与多GPU资源;结合NCCL、CUDA/cuDNN与Intel MKL等底层库,可实现单机多卡与多机多卡的高效扩展。对于单机多卡优先选用DDP,多机训练需正确配置通信后端与进程启动方式;CPU侧可通过MKL与多线程提升数据预处理与算子效率。
并行策略与适用场景
| 并行策略 | 适用场景 | 关键要点 |
|---|---|---|
| DataParallel(DP) | 快速上手、单机多卡 | 单进程多线程,易用但效率低于DDP,扩展性受限 |
| DistributedDataParallel(DDP) | 单机多卡/多机多卡 | 每卡一个进程,通信高效,推荐默认方案 |
| 模型并行(MP) | 单卡显存不足的大模型 | 手动切分模型层到不同GPU,需显式传输张量 |
| 流水线并行(PP) | 超大模型(如LLM) | 将模型分阶段并行,需精细切分与调度 |
| CPU多线程/异步 | CPU密集型预处理/推理 | 结合MKL与DataLoader多进程,提高吞吐 |
以上策略在CentOS上的实现与取舍要点可参考PyTorch官方用法与社区实践。
环境准备与安装要点
快速上手示例
# 假设单机4卡
torchrun --nproc_per_node=4 train.py
# train.py 关键片段
import torch, torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data.distributed import DistributedSampler
def main():
dist.init_process_group(backend="nccl")
rank = dist.get_rank()
torch.cuda.set_device(rank)
model = YourModel().cuda(rank)
model = DDP(model, device_ids=[rank])
dataset = YourDataset()
sampler = DistributedSampler(dataset, shuffle=True)
loader = DataLoader(dataset, batch_size=..., sampler=sampler, num_workers=4, pin_memory=True)
for epoch in range(epochs):
sampler.set_epoch(epoch)
for x, y in loader:
x, y = x.cuda(rank, non_blocking=True), y.cuda(rank, non_blocking=True)
# 前向、损失、反向、优化
要点:使用NCCL后端、每进程绑定单卡、DistributedSampler保证数据不重复、DataLoader开启多进程+pin_memory提升吞吐。
from torch.cuda.amp import autocast, GradScaler
scaler = GradScaler()
for x, y in loader:
x, y = x.cuda(), y.cuda()
with autocast():
out = model(x)
loss = criterion(out, y)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
AMP可在保持精度的同时降低显存占用并加速训练。
性能优化与排障要点