Megatron-DeepSpeed项目中TP=1时训练LLaMA2模型失败的排查与解决

问题背景

在使用Megatron-DeepSpeed框架进行LLaMA2模型训练时，开发人员遇到了一个特殊现象：当设置张量并行度(TP)为1时，使用8块H800 GPU进行训练会失败，而TP≥2时却能正常运行。这个问题涉及到深度学习分布式训练中的关键配置参数，值得深入分析。

错误现象分析

训练过程中出现的核心错误信息是CUDA设备端断言失败，具体表现为：

RuntimeError: CUDA error: device-side assert triggered
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.

错误发生在执行torch.embedding()操作时，提示srcIndex < srcSelectDimSize断言失败。这表明在词嵌入层操作中，输入的索引值超出了词表大小的范围。

初步排查

开发人员首先尝试了独立的测试代码来验证PyTorch的embedding函数：

import torch

weight = torch.nn.Parameter(torch.empty(32000,4096,dtype=torch.float16,device='cuda'))
torch.nn.init.uniform_(weight, -1, 1)
input = torch.randint(0, 32000, (1, 4096), dtype=torch.int64, device='cuda')
embedding_output = torch.nn.functional.embedding(input,weight,None,None,2,False,False)

这段测试代码能够正常运行，说明PyTorch的embedding函数本身没有问题。这表明问题可能出在训练过程中的数据预处理环节，或者是分布式训练环境下的特殊配置。

深入分析

在Megatron-DeepSpeed框架中，当TP=1时，模型参数和数据分布与TP≥2时有显著不同：

1. 参数分布差异：TP=1时，所有参数都完整存在于每个GPU上；TP≥2时，参数被切分到不同GPU上
2. 数据流差异：TP=1时数据并行度更高，可能需要更大的显存处理完整批次
3. 同步机制差异：不同TP配置下梯度同步和参数更新的方式不同

问题根源

经过深入排查，发现问题确实出在数据预处理环节。当TP=1时，数据加载和处理流程与TP≥2时有所不同，导致某些输入token的索引值超出了预设词表大小的范围。这种情况在TP≥2时由于数据切分方式不同而没有触发。

解决方案

最终解决方案是检查并修正数据预处理流程：

1. 确保训练数据中的token都映射到有效的词表范围内
2. 验证数据加载器在不同TP配置下的行为一致性
3. 检查tokenizer是否正确配置并与词表大小匹配

经验总结

这个案例提供了几个有价值的经验：

1. 分布式训练配置敏感性：不同的并行策略可能导致不同的错误表现，不能仅凭一种配置的工作状态判断问题
2. 数据预处理的重要性：即使模型架构正确，数据问题也可能导致难以诊断的错误
3. 系统性排查方法：从简单测试到复杂场景逐步验证，是定位深度学习问题的有效方法

对于使用Megatron-DeepSpeed框架的开发者，建议在遇到类似问题时：

1. 首先验证数据预处理流程
2. 尝试不同的并行配置组合
3. 使用小规模数据和简单配置进行问题复现
4. 逐步增加复杂度直到问题重现

这种系统性的调试方法可以显著提高解决分布式训练问题的效率。