ChunkFlow优化扩散模型推理,让数据预取更聪明地避开通信拥堵。
ChunkFlow: Communication-Aware Chunked Prefetching for Layerwise Offloading in Distributed Diffusion Transformer Inference
- 根据通信与计算的动态关系,自适应调整数据预取粒度。
- 在双H100+PCIe环境下,最多提升1.28倍推理速度,内存峰值降低49%。
- 适合小负载场景下的高效扩散模型推理,尤其对通信敏感系统有用。
层间卸载通过从主机内存预取后续层来降低大型扩散变换器(DiT)推理的GPU内存占用,但其有效性依赖于将预取延迟隐藏在每层计算之后。当每卡计算量较小时,该假设失效。此外,在仅支持PCIe的节点上,预取与GPU间集体通信(如all-reduce、all-to-all)竞争共享PCIe链路,导致即使计算足够也无法隐藏预取延迟。本文将层间卸载重新建模为预取与通信的协同调度问题,基于首个一阶分析模型预测预取是否可被计算掩盖。在此基础上,设计了ChunkFlow:一种通信感知、分块粒度的卸载运行时,能自适应让步于集体通信,并平滑地在GPU内存与预取量之间权衡。在三种代表性扩散变换器上,使用两块H100 GPU与Ulysses序列并行,ChunkFlow相比SGLang现有层间卸载最多提升1.28倍步骤时间速度,相较于无卸载基线峰值内存减少49%,且在工作负载足够大时步时延基本不变;同时在小负载场景下展现出可调的内存-延迟权衡,几乎消除步骤时间开销。
原文摘要 · Abstract (English)
Layerwise offloading reduces the GPU memory footprint of large diffusion transformer (DiT) inference by prefetching upcoming layers from host memory, but its effectiveness hinges on hiding prefetch latency behind per-layer computation. This assumption breaks down when the per-GPU compute workload is small. Moreover, on PCIe-only nodes, prefetch and inter-GPU collective communications such as all-reduce and all-to-all contend on the shared PCIe path, exposing prefetch latency even when compute would otherwise hide it. We revisit layerwise offloading as a co-scheduling problem between prefetch and communication, guided by a first-order analytical model that predicts when prefetch can be hidden by computation. Building on this model, we design ChunkFlow, a communication-aware, chunk-granular offloading runtime that adaptively yields to collective communication and smoothly trades GPU memory for prefetch volume. On three representative diffusion transformers running on two H100 GPUs over PCIe with Ulysses sequence parallelism, ChunkFlow delivers up to 1.28x step-time speedup over SGLang's existing layerwise offloading, reduces peak GPU memory by up to 49% over the no-offload baseline at near-identical step time once the workload is large enough, and exposes a tunable memory-latency tradeoff that recovers near-zero step-time overhead in the small-workload regime.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。