Cephalo: Harnessing Heterogeneous GPU Clusters for Training Transformer Models

TL;DR

Cephalo is presented, a system that holistically balances both compute and memory usage by decoupling compute distribution from training state assignment and uses an optimizer to efficiently distribute the compute workload and storage of training state to account for GPU heterogeneity in the cluster.

Abstract

Training transformer models requires substantial GPU compute and memory resources. In homogeneous clusters, distributed strategies allocate resources evenly, but this approach is inefficient for heterogeneous clusters, where GPUs differ in power and memory. As high-end GPUs are costly and limited in availability, heterogeneous clusters with diverse GPU types are becoming more common. Existing methods attempt to balance compute across GPUs based on capacity but often underutilize compute due to memory constraints. We present Cephalo, a system that optimizes compute and memory usage by decoupling compute distribution from training state assignment. Cephalo outperforms state-of-the-art methods by achieving significantly higher training throughput while supporting larger models and batch sizes.

Figures (12)

• Figure 1: Hourly AWS GPU availability over 12-hour period.
• Figure 2: GPU TFlops (FP32) vs. Memory Capacity.
• Figure 3: Architecture of Cephalo.
• Figure 4: Gradient accumulation in FSDP (top) vs Cephalo (bottom). The diagram illustrates gradient accumulation over $2$ microbatches on a model consisting of $3$ FSDP units. $F_{ij}$ and $B_{ij}$ are the forwards and backwards passes of the $i$th FSDP unit on the $j$th microbatch. $AG_i$ and $RS_i$ are the AllGather and ReduceScatter collectives for the $i$th FSDP unit.
• Figure 5: Training latency and memory allocated for compute as the microbatch size increases for Bert-Large.