通过垂直调度与优化步重叠，提升大模型在固态盘上训练的效率。
GreedySnake: Accelerating SSD-Offloaded LLM Training with Efficient Scheduling and Optimizer Step Overlapping
- 采用垂直调度：完成一层所有微批次后再进入下一层。
- 在单卡和多卡上分别实现最高2.53倍的训练吞吐提升。
- 适合大规模语言模型训练，尤其适用于资源受限场景。
SSD离线训练为降低大模型训练成本提供了实用且有前景的方案。本文基于微批次梯度累积，提出GreedySnake系统，采用垂直调度策略——在进入下一层前完成当前层所有微批次的执行。相比传统的水平调度（即按顺序执行微批次），GreedySnake在较小批量下实现了更高训练吞吐，使系统更接近屋顶线模型预测的理想性能。为进一步缓解I/O瓶颈，GreedySnake将部分优化器步骤与下一迭代的前向计算重叠进行。在A100 GPU上的实验表明，与ZeRO-Infinity相比，GreedySnake在1张卡上对GPT-65B实现1.96倍、4张卡上实现1.93倍的吞吐提升；对GPT-175B则在单卡上达到2.53倍的提升。
原文摘要 · Abstract (English)
SSD-offloaded training offers a practical and promising approach to making LLM training cost-effective. Building on gradient accumulation with micro-batches, this paper introduces GreedySnake, a new SSD-offloaded training system that employs vertical scheduling, which executes all microbatches of a layer before proceeding to the next. Compared to existing systems that use horizontal scheduling (i.e., executing micro-batches sequentially), GreedySnake achieves higher training throughput with smaller batch sizes, bringing the system much closer to the ideal scenario predicted by the roofline model. To further mitigate the I/O bottleneck, GreedySnake overlaps part of the optimization step with the forward pass of the next iteration. Experimental results on A100 GPUs show that GreedySnake achieves saturated training throughput improvements over ZeRO-Infinity: 1.96x on 1 GPU and 1.93x on 4 GPUs for GPT-65B, and 2.53x on 1 GPU for GPT-175B.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。