Torchtitan项目在H200 GPU集群上的性能优化实践

问题背景

在使用Torchtitan项目进行Llama3-70B模型的多节点训练时，用户报告在4节点32块H200 GPU环境下仅获得约200 tokens/s的处理速度。这一性能远低于预期，特别是在单节点8B模型训练表现正常的情况下。

性能瓶颈分析

通过对用户提供的性能追踪文件分析，发现主要存在以下问题：

1. 通信瓶颈：FSDP全收集操作耗时超过前向计算的两倍以上，表明跨节点通信成为主要性能瓶颈
2. 配置不当：环境变量CUDA_LAUNCH_BLOCKING被错误设置为1，导致CUDA内核同步执行
3. 网络适配：集群使用InfiniBand而非EFA，但未针对此进行优化配置

解决方案

网络配置优化

对于使用InfiniBand的集群环境，建议配置以下关键环境变量：

export NCCL_SOCKET_IFNAME=eth0
export NCCL_IB_HCA=ibp
export NCCL_MIN_CTAS=32
export UCX_NET_DEVICES=ibp0:1,ibp1:1,ibp2:1,ibp3:1,ibp4:1,ibp5:1,ibp6:1,ibp7:1
export SHARP_COLL_ENABLE_PCI_RELAXED_ORDERING=1
export NCCL_COLLNET_ENABLE=0
export NCCL_DEBUG=INFO
export NCCL_ALGO=NVLSTREE

训练策略调整

1. 采用HSDP策略：建议使用8路分片保持FSDP全收集/归约分散操作在节点内部完成
2. 禁用同步执行：确保设置CUDA_LAUNCH_BLOCKING=0以避免不必要的同步

性能验证

1. 基准测试：首先运行NCCL测试验证集群带宽性能
2. 小规模验证：使用Llama3-8B模型进行单节点FSDP训练验证基本性能
3. 逐步扩展：从单节点扩展到多节点，观察性能变化

实际效果

经验证，在正确配置的H200集群上，Torchtitan项目可以实现：

• HSDP2策略下达到1600 tokens/s的处理速度
• 多节点扩展性能接近线性增长
• 通信开销显著降低

总结

H200 GPU集群上的性能问题主要是由于网络配置不当和训练策略选择造成的。通过优化网络环境变量、调整分片策略以及正确设置CUDA执行模式，可以显著提升Torchtitan项目在大模型训练中的性能表现。对于使用InfiniBand的集群环境，特别注意网络适配器的正确配置是获得最佳性能的关键。