ColossalAI模型并行训练中词汇表大小与张量并行数不匹配问题分析

问题背景

在使用ColossalAI进行大规模语言模型训练时，当词汇表大小(vocab_size)不能被张量并行数(tp_size)整除时，会出现模型检查点加载失败的问题。这是一个典型的模型并行训练中的边界条件问题，值得深入探讨其成因和解决方案。

现象描述

在ColossalAI项目中，当使用Llama模型进行训练时，如果设置vocab_size=65535（不能被常见的并行数如2整除），在保存检查点后尝试重新加载模型时，会遇到如下错误：

RuntimeError: Error(s) in loading state_dict for LlamaForCausalLM:
    size mismatch for model.embed_tokens.weight: copying a param with shape torch.Size([32768, 1024]) from checkpoint, the shape in current model is torch.Size([65535, 1024]).
    size mismatch for lm_head.weight: copying a param with shape torch.Size([32768, 1024]) from checkpoint, the shape in current model is torch.Size([65535, 1024]).

技术分析

张量并行与词汇表切分

在模型并行训练中，词汇表通常会被切分到不同的设备上。理想情况下，词汇表大小应该能被张量并行数整除，这样每个设备可以获得相同大小的词汇表分片。例如：

• 当vocab_size=65536且tp_size=2时，每个设备获得32768个词向量
• 当vocab_size=65535且tp_size=2时，理论上应该获得32767.5个词向量，这在实践中无法实现

填充机制

ColossalAI采用了填充(padding)机制来处理这种不整除的情况：

1. 将原始词汇表大小向上填充到最近的能被tp_size整除的数值
2. 例如65535会被填充到65536
3. 在计算时使用掩码(mask)来忽略填充部分的影响

问题根源

通过调试信息可以发现，在保存检查点时：

1. 参数首先被标记为填充张量(ptensor=True)
2. 然后进行解填充操作(unpad)
3. 最后进行张量收集(gather)

然而，当前的实现顺序存在问题：

• 应该在收集完整张量后再进行解填充操作
• 当前顺序导致只收集了填充后的分片，丢失了原始词汇表大小的信息

解决方案建议

短期修复

调整参数保存的顺序：

1. 先收集分布式张量
2. 然后进行解填充操作
3. 最后保存完整的、未填充的参数

长期改进

1. 在模型初始化时增加vocab_size检查，当检测到不整除情况时：
   • 发出明确警告
   • 提供自动填充选项
2. 改进检查点格式，显式存储原始词汇表大小信息
3. 增强错误处理，提供更友好的错误信息

影响范围

此问题影响所有使用ColossalAI进行模型并行训练且词汇表大小不满足整除条件的场景，特别是：

• 自定义词汇表大小的模型
• 使用非2的幂次作为张量并行数的情况
• 多阶段训练中需要加载检查点的场景

最佳实践

为避免此类问题，建议：

1. 设计模型时尽量使vocab_size能被常见的并行数(如2、4、8等)整除
2. 如果必须使用特定词汇表大小，考虑手动填充到最近的合适数值
3. 在保存检查点前验证参数形状是否符合预期

总结

ColossalAI中的这一边界条件问题揭示了模型并行训练中资源分配的重要性。通过深入理解张量切分和填充机制，开发者可以更好地设计模型架构和训练配置，避免潜在问题。该问题的修复将提升框架在非理想条件下的鲁棒性，为更灵活的模型训练提供支持。