LLaMA-Factory 多机多卡训练中的 Ray 分布式训练问题分析与解决方案

问题背景

在使用 LLaMA-Factory 进行大规模语言模型微调时，研究人员经常需要利用多台设备的计算资源来加速训练过程。Ray 作为一个分布式计算框架，能够有效地协调多台设备上的 GPU 资源进行并行训练。然而，在实际部署过程中，用户可能会遇到一些技术挑战。

典型错误现象

当用户尝试使用 Ray 在 3 台设备（每台配备 1 张 GPU）上进行模型微调时，虽然训练过程看似完成，但系统最终会抛出 AssertionError 错误。错误日志显示，问题出现在 torch.distributed.destroy_process_group() 方法调用时，系统断言进程组(pg)不为 None 的条件失败。

问题根源分析

经过深入调查，发现这个问题源于 Ray 框架内部的一个设计缺陷。具体表现为：

1. 重复销毁进程组：系统在训练过程中会两次尝试销毁分布式进程组，第一次是在 LLaMA-Factory 的训练函数内部，第二次是在 Ray 的 on_shutdown 回调中。

2. 条件判断不足：Ray 原始的代码仅通过一个布尔标志(destroy_process_group)来决定是否销毁进程组，而没有检查进程组是否已经被销毁或是否仍然存在。

解决方案

针对这个问题，有两种可行的解决路径：

方案一：修改 Ray 源码

在所有节点上修改 Ray 的配置文件：

/root/anaconda3/envs/train/lib/python3.10/site-packages/ray/train/torch/config.py

将原来的条件判断：

if destroy_process_group:

修改为更健壮的检查方式：

if dist.is_initialized():

这种方法直接解决了框架层面的问题，确保只有在分布式环境确实初始化的情况下才会尝试销毁进程组。

方案二：修改 LLaMA-Factory 训练逻辑

另一种思路是调整 LLaMA-Factory 自身的训练流程，避免在训练函数内部过早销毁进程组，将这个职责完全交给 Ray 的 shutdown 回调处理。

多机训练中的检查点管理

在多机多卡训练场景下，检查点的保存位置也是一个需要注意的技术点。目前 Ray 会随机选择一个节点保存训练结果，这可能给后续的模型管理和部署带来不便。

推荐的解决方案包括：

1. 使用共享存储：设置一个 NFS 服务器作为主节点，其他工作节点挂载共享文件夹，确保所有检查点都存储在统一位置。

2. 定期同步机制：实现一个定期将检查点从各工作节点同步到主节点的机制，保证数据的集中管理。

最佳实践建议

1. 环境一致性：确保所有参与训练的节点具有相同的软件环境和依赖版本，避免因环境差异导致的问题。

2. 日志监控：密切关注训练过程中的日志输出，及时发现并解决潜在问题。

3. 资源预留：在训练大规模模型时，为系统操作预留一定的计算资源，避免因资源耗尽导致训练中断。

4. 定期验证：在长时间训练过程中，定期验证模型检查点的完整性和可用性。

总结

分布式训练是加速大规模语言模型微调的有效手段，但在实际部署过程中可能会遇到各种技术挑战。通过深入理解框架的工作原理和仔细排查错误原因，可以有效解决这些问题。本文介绍的 Ray 分布式训练问题及其解决方案，为使用 LLaMA-Factory 进行多机多卡训练的研究人员提供了实用的技术参考。