IREE项目中Llama 8B FP16模型编译失败问题分析

问题背景

在IREE深度学习编译器项目中，开发团队发现了一个关于Llama 8B FP16模型编译失败的问题。该问题出现在PR#20320合并后，原本可以正常编译的模型开始出现共享内存超限的错误。

错误现象

编译过程中出现的具体错误信息表明，某个函数操作使用了286720字节的共享内存，超过了65536字节的限制。错误发生在torch.aten._scaled_dot_product_flash_attention_for_cpu操作处，涉及维度为[4,32,?,128]的张量计算。

技术分析

问题根源

经过开发团队深入分析，发现问题出在IREE的维度折叠优化阶段。具体表现为：

1. 在注意力机制计算过程中，维度折叠优化产生了不合理的中间展开操作
2. 这种展开操作导致了共享内存使用量激增
3. 问题主要出现在注意力机制中掩码和注意力计算之间的维度转换

复现案例

开发团队提供了两个简化测试用例来复现这个问题：

1. 第一个测试用例(test1)在优化过程中产生了不必要的中间展开操作
2. 第二个测试用例(test2)则没有产生这种问题

这两个测试用例的主要区别在于输入张量的维度排列方式，这影响了优化器的决策过程。

优化过程分析

通过分析优化器的处理流程，可以观察到：

1. 原始计算涉及高维张量操作(如4x8x4x?x32x128)
2. 优化器尝试折叠某些维度以减少计算复杂度
3. 在某些情况下，优化器会插入不必要的展开操作
4. 这些展开操作导致内存使用量超出硬件限制

解决方案

开发团队提出了以下解决方案：

1. 暂时禁用注意力机制中的rope融合优化
2. 修复维度折叠优化中的过度激进策略
3. 确保优化后的计算图不会产生不必要的高内存消耗操作

技术影响

这个问题对IREE项目的影响包括：

1. 影响Llama等大型语言模型的编译和优化
2. 揭示了维度折叠优化中的潜在问题
3. 为未来优化器改进提供了重要参考案例

总结

IREE项目中出现的Llama 8B FP16编译问题，反映了深度学习编译器在复杂模型优化过程中面临的挑战。通过深入分析维度折叠优化的行为，开发团队不仅解决了当前问题，也为未来类似问题的诊断和解决积累了宝贵经验。这类问题的解决有助于提升IREE对大型语言模型的支持能力，推动深度学习编译器技术的进步。