OpenDiT项目中CogVideoX序列并行实现的问题分析与解决方案

背景介绍

在OpenDiT项目的视频生成模型CogVideoX中，序列并行(Sequence Parallel)是实现大规模模型训练的重要技术手段。该技术通过将输入序列沿序列维度分割到不同GPU上进行并行计算，从而突破单卡显存限制，实现更长序列的训练。

问题发现

在CogVideoX的注意力机制实现中，编码器隐藏状态(encoder_hidden_states)和解码器隐藏状态(hidden_states)会沿序列维度进行拼接。然而，当前视频系统(videosys)中的序列并行实现存在一个关键问题：

1. 系统仅对hidden_states沿序列维度进行分割
2. 却将完整的encoder_hidden_states与分割后的hidden_states进行拼接
3. 这种实现导致经过all-to-all通信后的注意力计算语义与原始实现不一致

问题影响

虽然生成的视频帧看起来自然，但对比实验显示：

• 单GPU(无并行)生成结果
• 2GPU(cp_size=2)生成结果
• 4GPU(cp_size=2,sp_size=2)生成结果

三者之间存在明显差异，例如画面中狗左侧落叶的位置和形态不一致。这种差异在理论上不应该出现，表明并行实现确实改变了模型的计算逻辑。

技术分析

正确的序列并行实现应该：

1. 同时对encoder_hidden_states和hidden_states沿序列维度进行分割
2. 保持两部分分割方式的一致性
3. 确保拼接操作在分割后的张量上进行

当前实现的问题在于仅分割了部分输入，破坏了注意力计算的完整性，导致模型行为出现偏差。

解决方案

项目维护者已确认该问题，并在相关提交中修复了此bug。修复方案主要包括：

1. 统一处理encoder_hidden_states和hidden_states的分割
2. 确保注意力计算在分割后的张量上保持语义一致性
3. 验证不同并行配置下的生成结果一致性

经验总结

在实现序列并行时，需要特别注意：

1. 所有需要沿序列维度操作的张量必须采用相同的分割策略
2. 拼接、注意力等操作需要在分割后的张量上保持计算语义不变
3. 需要通过严格的对比测试验证并行实现的正确性

这个问题提醒我们，在分布式训练实现中，任何细小的不一致都可能导致模型行为的改变，需要格外谨慎处理各部分的并行策略一致性。