Intel Extension for PyTorch 中的 CPU 内存泄漏问题分析与解决方案

问题背景

在使用 Intel Extension for PyTorch (IPEX) 进行深度学习模型训练时，用户报告了一个严重的 CPU 内存泄漏问题。当在 Data Center GPU Max 1550 等 Intel GPU 上训练模型时，CPU 内存会线性增长，最终导致内存不足(OOM)错误。这个问题在 IPEX 2.0.120 和 OneAPI 2024.0 环境下尤为明显。

问题现象

用户提供了一个基于 Transformer 架构的最小可复现示例。在训练过程中，可以观察到以下现象：

1. 在 Intel GPU (XPU) 上运行时，CPU 内存以约 10MB/s 的速度持续增长
2. 同样的代码在 NVIDIA Ampere 架构 GPU 上运行时，CPU 内存保持稳定
3. 内存泄漏会导致训练过程最终因内存耗尽而失败

技术分析

问题根源

经过 Intel 开发团队的调查，这个内存泄漏问题源于 IPEX 内部的内存管理机制。具体来说，是在某些特定操作后未能正确释放 CPU 内存。这个问题在 IPEX 2.1.30 及更高版本中得到了修复。

影响范围

该问题主要影响以下配置环境：

• IPEX 版本低于 2.1.2
• OneAPI 2024.0 版本
• 使用 Intel XPU 进行模型训练的场景

解决方案

推荐方案

升级到最新版本的 IPEX 和 OneAPI 可以完全解决这个问题：

1. 升级到 IPEX 2.1.30 或更高版本
2. 使用 OneAPI 2024.1 或更高版本

测试表明，在升级后的环境中，CPU 内存使用量在训练过程中保持稳定，不再出现线性增长的情况。

验证结果

升级后，用户观察到的内存使用情况如下：

Number of Model Parameters: 190,168,040
Epoch 1, Memory Usage: 1.44 GB
Epoch 2, Memory Usage: 1.89 GB
Epoch 3, Memory Usage: 1.89 GB
...
Epoch 10, Memory Usage: 1.89 GB

相比之下，在问题版本中的内存使用会持续增长：

Number of Model Parameters: 190,168,040
Epoch 1, Memory Usage: 1.68 GB
Epoch 2, Memory Usage: 2.48 GB
...
Epoch 10, Memory Usage: 8.53 GB

最佳实践建议

1. 保持软件更新：定期检查并升级 IPEX 和 OneAPI 到最新版本，以获得最佳性能和稳定性
2. 监控内存使用：在训练过程中监控 CPU 和 GPU 内存使用情况，及时发现潜在问题
3. 版本兼容性：确保 IPEX 版本与 OneAPI 版本兼容，避免因版本不匹配导致的问题

结论

Intel Extension for PyTorch 作为优化 Intel 硬件上 PyTorch 性能的重要工具，其开发团队持续改进和修复已知问题。对于遇到类似 CPU 内存泄漏问题的用户，升级到最新版本是最有效的解决方案。这也提醒我们，在使用深度学习框架和扩展时，保持软件栈的更新是确保稳定运行的重要前提。