Pyramid-Flow项目多GPU推理问题分析与解决方案

问题背景

在Pyramid-Flow项目中使用flux模型进行视频生成任务时，用户遇到了RuntimeError错误。该错误主要发生在尝试使用多GPU进行推理时，具体表现为张量尺寸不匹配问题。

错误现象分析

当用户尝试使用3个GPU运行Pyramid-Flow的i2v(图像到视频)或t2v(文本到视频)任务时，系统会抛出以下关键错误：

RuntimeError: The size of tensor a (248) must match the size of tensor b (249) at non-singleton dimension 1

深入分析错误堆栈，问题出现在flux模块的旋转位置编码(ROPE)实现部分。具体来说，在计算旋转位置编码时，两个张量的第一维度尺寸不匹配(248 vs 249)，导致无法完成广播操作。

多GPU配置影响

进一步测试发现，GPU数量配置直接影响该问题的出现：

1. 2个GPU：运行正常
2. 3个GPU：出现上述尺寸不匹配错误
3. 4个GPU：出现新的矩阵乘法维度不匹配错误

解决方案

根据项目维护者的建议和实际测试结果，目前推荐的解决方案是：

1. 对于pyramid_flux模型，使用2个GPU进行推理
2. 对于pyramid_mmdit模型，可以使用2个或4个GPU进行推理

技术原理分析

该问题本质上与模型并行计算时的张量切分策略有关。Pyramid-Flow中的flux模块采用了特定的注意力机制实现，其中包含旋转位置编码(ROPE)的计算。在多GPU环境下，当GPU数量不是2的幂次时，可能导致：

1. 张量切分不均匀
2. 各GPU间计算负载不均衡
3. 广播操作维度不匹配

最佳实践建议

1. 始终使用2的幂次数量的GPU(如2,4,8等)
2. 确保输入图像的宽高设置与模型变体匹配
3. 对于384p变体，保持默认的640×384分辨率
4. 监控GPU内存使用情况，合理选择batch size

结论

Pyramid-Flow项目在多GPU推理时对GPU数量配置较为敏感。通过合理配置GPU数量(特别是使用2的幂次)，可以避免大部分张量维度不匹配问题，确保视频生成任务的稳定执行。未来随着项目的更新，这一问题有望得到更完善的解决。