xDiT项目多GPU环境下Flux模型运行问题分析与解决

问题背景

在使用xDiT项目的Flux模型进行多GPU并行计算时，用户遇到了CUDA非法内存访问的错误。该问题主要出现在使用2个或4个H100 GPU的情况下，而单GPU运行则表现正常。错误发生在调用pipe.prepare_run(input_config, steps=1)方法时，系统抛出RuntimeError: CUDA error: an illegal memory access was encountered异常。

错误现象分析

在多GPU环境下运行Flux模型时，系统报告了以下关键错误信息：

1. CUDA非法内存访问错误，通常表明GPU尝试访问了未分配或已释放的内存区域
2. 错误发生在T5模型的相对位置编码计算过程中
3. 当使用torchrun在SLURM环境下运行时，还出现了NCCL通信相关的错误

环境配置

问题出现的环境配置如下：

• Python 3.12
• PyTorch 2.5
• CUDA 12.4
• H100 GPU（2-4卡）
• SLURM作业调度系统

问题排查过程

初步排查

1. 确认单GPU环境工作正常，问题仅出现在多GPU场景
2. 检查了数据并行度(pipefusion_parallel_degree)设置，尝试了2-4的不同值
3. 验证了基本的PyTorch多GPU通信功能正常

深入分析

通过错误堆栈发现，问题出现在T5文本编码器的相对位置编码计算阶段。具体是在计算相对位置桶(relative position bucket)时发生的CUDA内存访问违规。这表明：

1. 可能是模型并行实现中的内存分配问题
2. 或者是多GPU间数据同步时的内存访问冲突
3. 也可能是某些组件的版本兼容性问题

解决方案

经过多次尝试和验证，最终发现问题与apex库的版本有关。具体解决步骤如下：

1. 更新apex库到最新版本
2. 确保所有相关组件的版本兼容性
3. 重新测试SD3和Flux模型的多GPU运行

更新后，不仅Flux模型的多GPU运行问题得到解决，SD3模型的相关错误也一并消除。SD3之前报告的"RuntimeError: expected scalar type Float but found Half"错误同样源于apex库的版本问题。

经验总结

1. 在多GPU环境下运行大型模型时，组件版本兼容性至关重要
2. CUDA非法内存访问错误可能有多种根源，需要系统性地排查
3. 类似问题可先验证基础通信功能，再逐步深入模型内部
4. 保持关键组件(如apex)的更新可以避免许多潜在问题

最佳实践建议

对于在xDiT项目中使用多GPU运行Flux模型的用户，建议：

1. 定期更新apex等关键库
2. 在多GPU环境中先运行简单的测试程序验证基础功能
3. 注意PyTorch、CUDA和各组件间的版本匹配
4. 在SLURM环境中确保正确的资源分配和进程启动方式

通过系统性的问题排查和版本管理，可以有效避免类似的多GPU运行问题，充分发挥硬件加速性能。