ColossalAI中混合并行优化器的参数分片内存问题分析

问题背景

在ColossalAI项目的混合并行训练场景中，当使用张量并行(Tensor Parallelism)技术时，模型参数会被分片到不同的GPU上，以减少内存消耗并实现并行计算。然而，开发者发现优化器仍然保留着未分片的完整模型参数，导致内存无法被有效释放，造成了额外的内存开销。

问题现象

通过以下典型代码示例可以观察到问题现象：

colossalai.launch_from_torch(config={})
plugin = HybridParallelPlugin(tp_size=4, pp_size=1)
optimizer = Adam(model.parameters())
model, optimizer, *_ = booster.booster(model, optimizer, ...)

检查发现：

• 模型参数已被正确分片（如形状从[50257, 768]变为[12565, 768]）
• 但优化器中的参数仍保持未分片状态（仍为[50257, 768]）
• 混合精度优化器的参数映射关系也出现了不匹配

问题根源

深入分析后发现，问题主要出在模型预处理阶段的词表嵌入层调整逻辑上：

1. 当词表大小不能被TP维度整除时，会调用HuggingFace的resize_token_embeddings方法
2. 该方法会创建一个全新的nn.Embedding层，而非原地修改
3. 导致模型参数ID发生变化，与优化器中保留的原始参数失去关联

具体表现为：

• GPT2和BERT等模型的默认词表大小不能被常见TP维度(如8)整除
• 这些模型会触发重建嵌入层的逻辑
• 而OPT、Falcon等模型的词表大小则通常可以整除

解决方案

提出了一个原地调整词表嵌入层的改进方案：

def resize_token_embedding_inplace(num_new_tokens: int, embedding: nn.Embedding):
    # 原地调整词表嵌入层
    embedding.num_embeddings = new_num_tokens
    embedding.weight.data = nn.functional.pad(
        embedding.weight.data,
        (0, 0, 0, new_num_tokens - embedding.weight.size(0)),
        "constant",
        0,
    )

该方案的优势在于：

1. 使用nn.functional.pad直接扩展张量，避免创建新层
2. 保持参数ID不变，确保优化器能正确跟踪参数
3. 内存效率更高，不会产生冗余参数副本

技术影响

这个问题对训练过程有多方面影响：

1. 内存消耗：未释放的完整参数副本增加了约1/TP_size倍的内存占用
2. 训练正确性：混合精度优化器可能跳过部分参数的更新
3. 性能表现：额外的内存压力可能影响整体训练效率

最佳实践建议

对于使用ColossalAI进行混合并行训练的用户，建议：

1. 检查模型词表大小与TP维度的整除性
2. 对于需要调整词表的情况，优先使用原地调整方案
3. 监控训练过程中的内存使用情况
4. 验证参数更新是否覆盖了所有分片

总结

ColossalAI作为领先的大模型训练框架，其混合并行功能需要精细的内存管理。这个参数分片问题的发现和解决，体现了深度学习系统优化中参数生命周期管理的重要性。通过采用原地调整策略，既保证了训练的正确性，又提升了内存使用效率，为大规模模型训练提供了更可靠的解决方案。