关于VMM在Torch中的应用与优化效果的疑问

AI

关于VMM的一些看法与实践经验分享：1. VMM的概念其实很早就出现了，在我看来，它本质上是一种基础的API能力，类似于malloc。然而，要真正用好它、管理好它，尤其是在特定场景下发挥其优势，则需要大量额外的工作。这就好比jemalloc、tcmalloc这样的高性能内存分配器，或者像Redis这样的数据结构存储系统，甚至包括我之前从事的EMC文件系统领域——无论哪家存储系统的本质操作无非是open、read和write，但如何优化这些操作以满足不同需求才是关键所在。因此，我认为VMM本身并不是问题的核心，而是解决问题的一个重要切入点。2. 在我们的实际开发过程中，PyTorch并没有提供我们所需的抽象类，所以我们做了一些定制化的修改。此外，随着项目推进，所解决的问题也在不断演变。大约两年前，我们基于VMM实现了一些针对推理场景（主要是ONNXRuntime）的优化，并不是单纯为了解决碎片化问题（相关工作还在整理中）。后来尝试将其应用于训练时的碎片优化，效果并不显著，因为当时的碎片问题尚不严重。直到后来复现了ColossalAI的相关工作后，才意识到大模型中的碎片问题确实非常突出，于是我们针对这一场景进行了专门的定制优化，最终取得了不错的成效。 关于为何PyTorch没有更早引入类似机制，我认为这是由具体需求驱动的。如果问题尚未达到足够严重的程度，就没有必要进行大规模改动。当时我们团队并未参考后来的2.1版本（各团队独立开发），所以建议大家先仔细对比代码后再发表意见，这样会更加负责任。实际上，双方的实现差异较大，核心设计也不尽相同。不过，我们最近已经注意到这一点，并且正在进行一些对比测试，初步结果显示我们的方案效果更好（GitHub上已有反馈）。当然，这两种方案可以集成在一起，这也是我们当前正在努力的方向。3. 除了上述内容外，我们还实现了一些其他功能，计划近期开源，这些功能在某些特定场景下表现尤为出色。4. 当前，我们的研究重点又回到了其他场景。总体思路是能否通过统一管理和调度虚拟地址空间（V）与物理地址空间（P），构建一个集中的管理层，避免每次遇到问题时都针对具体对象单独优化（这种做法容易导致小池子现象并加剧碎片化）。此外，这也意味着内核层面可能需要做出相应调整。当然，这一观点存在争议，最终还是要看实际效果说话。作为系统层开发者，我希望通过多承担一些底层复杂性，让上层应用更加简洁高效。目前，这部分工作仍在整理完善中。需要注意的是，部分媒体报道可能存在一定偏差，对此我们已作出修正。下周我们将举办一次线上技术交流活动，欢迎大家参与讨论并提出宝贵意见。同时，我们也创建了一个微信群（群号已更新至GitHub页面），欢迎感兴趣的朋友加入交流。