[论文解读] Flashlight: Enabling Innovation in Tools for Machine Learning

Jacob Kahn, Vineel Pratap+12 more|arXiv (Cornell University)|Jan 29, 2022

Explainable Artificial Intelligence (XAI)被引用 14

一句话总结

Flashlight 是一个轻量级、模块化的机器学习库，旨在通过支持快速原型设计新型张量计算、内存管理和分布式训练技术，加速系统研究。该库以极低开销构建，完全控制内部机制，使研究人员能够用自定义实现替换核心操作（如张量加法），在性能上与 PyTorch 和 TensorFlow 竞争的同时，支持向主流框架的上游贡献。

ABSTRACT

As the computational requirements for machine learning systems and the size and complexity of machine learning frameworks increases, essential framework innovation has become challenging. While computational needs have driven recent compiler, networking, and hardware advancements, utilization of those advancements by machine learning tools is occurring at a slower pace. This is in part due to the difficulties involved in prototyping new computational paradigms with existing frameworks. Large frameworks prioritize machine learning researchers and practitioners as end users and pay comparatively little attention to systems researchers who can push frameworks forward -- we argue that both are equally important stakeholders. We introduce Flashlight, an open-source library built to spur innovation in machine learning tools and systems by prioritizing open, modular, customizable internals and state-of-the-art, research-ready models and training setups across a variety of domains. Flashlight allows systems researchers to rapidly prototype and experiment with novel ideas in machine learning computation and has low overhead, competing with and often outperforming other popular machine learning frameworks. We see Flashlight as a tool enabling research that can benefit widely used libraries downstream and bring machine learning and systems researchers closer together. Flashlight is available at https://github.com/flashlight/flashlight .

研究动机与目标

• 解决由于机器学习框架日益庞大和复杂，导致创新日益困难的问题。
• 通过提供针对系统创新优化的框架，弥合机器学习研究人员与系统研究人员之间的鸿沟。
• 支持新型计算范式的快速原型设计，包括自定义内存管理器和分布式计算模型。
• 提供一个与 PyTorch 和 TensorFlow 等 monolithic 框架相比更轻量、更模块化的替代方案，支持对新内部组件的端到端实验。
• 通过在生产就绪、高性能环境中测试和验证新技术，促进向主流框架的上游贡献。

提出的方法

• 将 Flashlight 设计为一个轻量级、模块化的 C++ 库，提供简洁、可组合的 API，用于张量操作和计算图。
• 在多个领域（如自然语言处理、计算机视觉）实现完整、生产就绪的参考模型，作为系统研究的测试平台。
• 通过可插拔的内存管理器和自定义分配器后端，暴露对内存分配和张量布局的底层控制。
• 支持任意设备映射和计算调度，支持新型分布式训练策略的实验。
• 通过继承或组合，允许无缝替换核心操作（如张量加法），而无需修改现有操作符或基准测试代码。
• 集成自定义遥测和性能分析工具，实现实验过程中对操作级别内存行为和性能的监控。

实验结果

研究问题

• RQ1 一个轻量级、模块化的机器学习库是否能在支持系统级创新方面超越 monolithic 框架？
• RQ2 研究人员在完全控制分配和碎片化的情况下，能否高效地原型设计并基准测试新型内存管理技术？
• RQ3 在不破坏现有代码或无需大量重构的前提下，自定义张量操作（如优化的加法）在多大程度上可以被集成到框架中？
• RQ4 是否可以在一个灵活、可扩展的框架中实现并评估通用的分布式计算与优化器状态分片方法？
• RQ5 如何在不依赖庞大复杂代码库的情况下，对新型计算原语进行大规模端到端测试？

主要发现

• 与 PyTorch 和 TensorFlow 相比，Flashlight 的代码库规模和二进制文件大小均小几个数量级，显著提升了迭代速度并简化了维护。
• 研究人员能够快速原型设计并测试一种限制大缓存块拆分的内存管理器，在多个模型上将内部碎片率降低超过 20%。
• 自定义张量操作（如新型张量加法实现）可无缝替换所有操作符，无需修改现有代码，证明了框架的完全互操作性。
• 该框架的模块化设计支持开发通用的分布式训练技术，包括缓冲区分片和动态设备调度，且无硬编码假设。
• 性能基准测试表明，Flashlight 的参考实现通常在端到端训练工作负载中表现优于或等同于 PyTorch 和 TensorFlow。
• Flashlight 研究成果已成功上游贡献至 PyTorch，包括对其缓存分配器的改进，验证了该框架在真实世界中的影响。