[论文解读] Marius: Learning Massive Graph Embeddings on a Single Machine

Jason Mohoney, Roger Waleffe+3 more|arXiv (Cornell University)|Jan 20, 2021

Advanced Graph Neural Networks参考文献 33被引用 10

一句话总结

Marius 提出了一种单机框架，通过优化数据移动过程中的分区缓存和缓冲区感知数据排序（BETA），实现大规模图嵌入的训练，相比当前最先进系统，训练速度最高提升10倍，同时在仅使用16 GB GPU和64 GB CPU内存的情况下，可扩展至包含超过十亿条边、总参数量达550 GB的图模型。

ABSTRACT

We propose a new framework for computing the embeddings of large-scale graphs on a single machine. A graph embedding is a fixed length vector representation for each node (and/or edge-type) in a graph and has emerged as the de-facto approach to apply modern machine learning on graphs. We identify that current systems for learning the embeddings of large-scale graphs are bottlenecked by data movement, which results in poor resource utilization and inefficient training. These limitations require state-of-the-art systems to distribute training across multiple machines. We propose Marius, a system for efficient training of graph embeddings that leverages partition caching and buffer-aware data orderings to minimize disk access and interleaves data movement with computation to maximize utilization. We compare Marius against two state-of-the-art industrial systems on a diverse array of benchmarks. We demonstrate that Marius achieves the same level of accuracy but is up to one order of magnitude faster. We also show that Marius can scale training to datasets an order of magnitude beyond a single machine's GPU and CPU memory capacity, enabling training of configurations with more than a billion edges and 550 GB of total parameters on a single machine with 16 GB of GPU memory and 64 GB of CPU memory. Marius is open-sourced at www.marius-project.org.

研究动机与目标

• 解决单机大规模图嵌入训练中数据移动的瓶颈问题。
• 克服因图数据超出GPU显存导致的不规则内存访问模式所引发的GPU利用率低下问题。
• 通过优化I/O与计算重叠，实现超越单个GPU和CPU内存容量限制的图嵌入训练。
• 设计一种系统，在显著提升训练吞吐量的同时，保持与现有工业系统相当的高精度。
• 开发一种新型数据排序方案（BETA）和流水线化架构，以最小化磁盘I/O并最大化硬件利用率。

提出的方法

• 引入流水线化训练架构，将数据移动、数据传输和计算操作交错执行，以隐藏I/O延迟。
• 实现分区缓存机制，通过重用频繁访问的节点嵌入分区，减少冗余磁盘读取。
• 提出BETA排序——一种缓冲区感知的数据排序方案，通过将邻近节点分区的边分组，最小化I/O操作。
• 对节点嵌入采用异步训练并引入有界延迟机制，对边嵌入采用同步训练，以支持异构边类型。
• 采用小批量随机梯度下降（SGD）并结合高效的点积运算，降低计算开销。
• 通过缓冲区大小、分区数量和延迟边界等系统配置参数的调优，实现吞吐量与嵌入质量之间的平衡。

实验结果

研究问题

• RQ1 是否可以在不使用分布式系统的情况下，实现在单机上对十亿条边图进行嵌入训练？
• RQ2 如何最小化数据移动开销，以实现在超出GPU显存的训练中实现高GPU利用率？
• RQ3 何种数据排序策略可在保持大规模图嵌入训练效率的同时，降低I/O成本？
• RQ4 系统级优化（如分区缓存和流水线化）在不牺牲模型精度的前提下，能在多大程度上提升训练吞吐量？
• RQ5 超参数（如批量大小、缓冲区大小和延迟边界）对训练性能和嵌入质量的影响如何？

主要发现

• Marius 实现了高达70%的平均GPU利用率，远超 DGL-KE 的10%和 PBG 的30%以下，显著提升了硬件利用率。
• 在 Freebase86m 数据集上，Marius 相比 DGL-KE 和 PBG 最快提升10倍，同时保持相同的模型精度。
• Marius 在仅使用16 GB GPU和64 GB CPU内存的情况下，可扩展至包含超过十亿条边、总参数量达550 GB的图模型。
• BETA 排序通过最小化节点嵌入的磁盘访问，显著降低I/O成本，尤其在高分区数场景下效果显著。
• 在大规模图上，使用较大的批量大小（约10,000）可提升吞吐量且不损失精度，但超过某一阈值后收益递减。
• 缓冲区容量的四次方便比显著减少磁盘交换次数，表明缓冲区大小配置对性能至关重要。