ComfyUI中Lumina模型在MPS后端上的复数张量操作问题解析

问题背景

在ComfyUI项目中，用户在使用Lumina 2.0模型时遇到了一个技术障碍。当在配备M1-M3芯片的Mac设备上运行时，模型渲染过程会失败，错误信息显示MPS后端不支持对复数张量执行repeat操作。

技术细节分析

该问题核心在于PyTorch的MPS后端对复数张量操作的支持不完善。具体表现为：

1. 错误触发点：在Lumina模型的rope_embedder模块中，尝试对复数张量执行repeat操作时抛出异常
2. 错误信息：RuntimeError: repeat(): Not supported for complex yet!
3. 影响范围：主要影响使用Apple Silicon芯片(M1-M3)且使用MPS后端的用户

解决方案演进

初始临时方案

用户最初提出的解决方案是将张量临时转移到CPU执行操作：

def __call__(self, ids: torch.Tensor):
    self.freqs_cis = [freqs_cis.to("cpu") for freqs_cis in self.freqs_cis]
    # ...执行操作...
    return torch.cat(result, dim=-1).to(ids.device)

这种方案虽然可行，但存在性能损耗和潜在精度问题。

改进方案

更优的解决方案是将复数张量分解为实部和虚部分别处理：

def __call__(self, ids: torch.Tensor):
    # 移动数据到相同设备
    self.freqs_cis = [freqs_cis.to(ids.device) for freqs_cis in self.freqs_cis]
    
    result = []
    for i in range(len(self.axes_dims)):
        # 分离实部和虚部
        freqs_cis_real = self.freqs_cis[i].real
        freqs_cis_imag = self.freqs_cis[i].imag
        
        # 执行独立操作
        gathered_real = torch.gather(freqs_cis_real.unsqueeze(0).repeat(...), ...)
        gathered_imag = torch.gather(freqs_cis_imag.unsqueeze(0).repeat(...), ...)
        
        # 重新组合复数
        result.append(torch.complex(gathered_real, gathered_imag))
    
    return torch.cat(result, dim=-1)

官方修复

项目维护者确认该问题已得到修复，建议用户更新到最新版本。

技术原理深入

复数在PyTorch中以两个浮点数(实部和虚部)的形式存储。MPS后端对复数操作的支持相对滞后是因为：

1. 硬件限制：Apple Silicon的GPU架构对复数运算的原生支持有限
2. 优先级问题：复数运算在深度学习中的使用频率相对较低
3. 实现复杂性：复数运算需要特殊的处理逻辑和优化

最佳实践建议

1. 版本更新：优先使用官方修复后的版本
2. 回退方案：如必须使用旧版，可采用分解实部虚部的方案
3. 性能考量：避免在关键路径上频繁进行设备间数据传输
4. 兼容性测试：针对不同硬件后端进行充分测试

总结

这个问题展示了深度学习框架在跨平台支持中面临的挑战。随着Apple Silicon的普及，MPS后端的重要性日益提升，开发者需要关注这类平台特定的兼容性问题。理解复数运算的实现原理和跨平台差异，对于开发稳健的深度学习应用至关重要。