ZenFlow: Enabling Stall-Free Offloading Training via Asynchronous Updates

TL;DR

ZenFlow is a novel offloading framework for large language model training that prioritizes important parameters, enabling asynchronous updates and significantly reducing GPU stalls and PCIe traffic while maintaining model accuracy.

Contribution

ZenFlow introduces a parameter prioritization and asynchronous update mechanism that decouples GPU and CPU updates, improving training efficiency for large models.

Findings

Up to 5x end-to-end training speedup

2x reduction in PCIe traffic

Over 85% reduction in GPU stalls

Abstract

Fine-tuning large language models (LLMs) often exceeds GPU memory limits, prompting systems to offload model states to CPU memory. However, existing offloaded training frameworks like ZeRO-Offload treat all parameters equally and update the full model on the CPU, causing severe GPU stalls, where fast, expensive GPUs sit idle waiting for slow CPU updates and limited-bandwidth PCIe transfers. We present ZenFlow, a new offloading framework that prioritizes important parameters and decouples updates between GPU and CPU. ZenFlow performs in-place updates of important gradients on GPU, while asynchronously offloading and accumulating less important ones on CPU, fully overlapping CPU work with GPU computation. To scale across GPUs, ZenFlow introduces a lightweight gradient selection method that exploits a novel spatial and temporal locality property of important gradients, avoiding costly global synchronization. ZenFlow achieves up to 5x end-to-end speedup, 2x lower PCIe traffic, and reduces GPU stalls by over 85 percent, all while preserving accuracy.