Nonuniform-Tensor-Parallelism: Mitigating GPU failure impact for Scaled-up LLM Training

TL;DR

Nonuniform Tensor Parallellism (NTP) addresses the problem of amplified throughput losses in large-scale LLM training caused by GPU failures within high-degree tensor-parallel scale-up domains. By allowing DP replicas to operate at reduced TP degrees and coupling this with a power-boosting rack design, NTP achieves throughput losses close to the fraction of failed GPUs, and near-zero loss when paired with power boosts (NTP-PW). Implemented on top of Megatron with gradient resharding and evaluated via a high-fidelity simulator and prototype experiments, NTP reduces loss from up to $12\%$ to about $3\%$ (and <$1\%$ for power-boosted variants) under varied failure scenarios. This approach enables robust, efficient LLM pretraining on future NVLink-scale clusters without large reserves of spare GPUs.

Abstract

LLM training is scaled up to 10Ks of GPUs by a mix of data-(DP) and model-parallel (MP) execution. Critical to achieving efficiency is tensor-parallel (TP; a form of MP) execution within tightly-coupled subsets of GPUs, referred to as a scale-up domain, and the larger the scale-up domain the better the performance. New datacenter architectures are emerging with more GPUs able to be tightly-coupled in a scale-up domain, such as moving from 8 GPUs to 72 GPUs connected via NVLink. Unfortunately, larger scale-up domains increase the blast-radius of failures, with a failure of single GPU potentially impacting TP execution on the full scale-up domain, which can degrade overall LLM training throughput dramatically. With as few as 0.1% of GPUs being in a failed state, a high TP-degree job can experience nearly 10% reduction in LLM training throughput. We propose nonuniform-tensor-parallelism (NTP) to mitigate this amplified impact of GPU failures. In NTP, a DP replica that experiences GPU failures operates at a reduced TP degree, contributing throughput equal to the percentage of still-functional GPUs. We also propose a rack-design with improved electrical and thermal capabilities in order to sustain power-boosting of scale-up domains that have experienced failures; combined with NTP, this can allow the DP replica with the reduced TP degree (i.e., with failed GPUs) to keep up with the others, thereby achieving near-zero throughput loss for large-scale LLM training.

Figures (14)

• Figure 1: Hybrid parallelism: here we show one DP replica of a 2-layer transformer partitioned into two pipeline stages (left), we zoom in on one MLP layer to show how it is partitioned onto 4 GPUs using TP (right). Not shown are the TP partitions for Attention layers (which are partitioned across the same 4 GPUs as the MLP layer within a pipeline stage), see section \ref{['sec:design']} for a rigorous formulation of TP Attention and MLP.
• Figure 2: Effects of NVL domain size and TP degree on per-GPU throughput when scaling training. Throughputs are normalized to NVL32 at 16K GPUs. [daiyaan: y-axis start at 0, fix the colors, make y-axis lost-throughput, make it clear that NVL16 is using unlimited TP ]
• Figure 3: Larger TP/scale-up domains experience higher failure-amplification; the same number of failed GPUs impacts a larger fraction of the cluster. Solid lines are median lost and shaded are maximum.
• Figure 4: Observed and predicted failure rates and realistic failure recovery times result in high failure fractions.
• Figure 5: Illustration of NTP (top) with a PyTorch trace of our prototype (bottom). Presync-reshard (dashed) runs in the NVL comm stream and is overlapped with computation (which is completely unaffected and is identical to healthy operation), postsync-reshard (dotted) runs in the NVL comm stream and is overlapped with gradient allreduce (solid) which incurs slightly increased volume. The trace shows 5 different gradients being resharded as soon as they are computed, all 5 gradients being synchronized together in one allreduce operation, and then being resharded post-sync one after another.