PyTorch在CentOS上的并行计算能力概览
在CentOS上，PyTorch可通过单机多卡、多机多卡与CPU多核协同，覆盖数据并行、模型并行、流水线并行与分布式训练等主流方式。实际性能主要受GPU型号与数量、NCCL/网络、数据管道效率与混合精度等因素影响。一般建议优先采用DistributedDataParallel（DDP）进行多GPU训练，并在数据侧使用多进程异步加载与pin_memory来喂数与减少GPU空转。
并行方式与适用场景
| 并行方式 | 适用场景 | 关键要点 |
|---|---|---|
| DataParallel（DP） | 快速上手、单机多卡 | 单进程多线程，易用；在多卡下性能通常不如DDP。 |
| DistributedDataParallel（DDP） | 单机多卡/多机多卡 | 每卡一个进程，通信后端推荐NCCL，扩展性与稳定性更好。 |
| 模型并行 | 单卡显存不足、超大模型 | 将模型不同层切分到不同GPU，需精细显存与通信权衡。 |
| 流水线并行 | 超大模型（如LLM） | 将模型分阶段并行，需处理气泡与边界开销。 |
| CPU多线程/数据加载并行 | CPU密集型预处理、数据增强 | DataLoader(num_workers>0, pin_memory=True) 提升吞吐。 |
| 混合精度训练（AMP） | 加速训练、降低显存 | 使用torch.cuda.amp.autocast/GradScaler，在保持精度的同时提速。 |
快速上手示例
# train.py
import torch, torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, DistributedSampler
import torch.multiprocessing as mp
def main(rank, world_size):
dist.init_process_group(backend='nccl', init_method='env://', world_size=world_size, rank=rank)
torch.cuda.set_device(rank)
model = torch.nn.Linear(10, 1).cuda(rank)
ddp_model = DDP(model, device_ids=[rank])
dataset = torch.utils.data.TensorDataset(torch.randn(1024, 10), torch.randn(1024, 1))
sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank, shuffle=True)
loader = DataLoader(dataset, batch_size=32, sampler=sampler, num_workers=4, pin_memory=True)
for epoch in range(3):
sampler.set_epoch(epoch)
for x, y in loader:
x, y = x.cuda(rank, non_blocking=True), y.cuda(rank, non_blocking=True)
loss = ddp_model(x).sum()
loss.backward()
# 简化：直接step（实际请加梯度裁剪/同步等）
dist.destroy_process_group()
if __name__ == '__main__':
world_size = 2
mp.spawn(main, args=(world_size,), nprocs=world_size, join=True)
要点：使用NCCL后端、每进程绑定单卡、用DistributedSampler与pin_memory提升数据吞吐。
性能优化要点
集群与运维实践