Splitting the model: five parallelisms and the traffic they make切分模型：五種平行化與它們產生的流量

3.1Memory needed to train a model訓練模型的記憶體需求

Chapter 1 explained why a large language model needs thousands of GPUs to finish training in a reasonable time. Besides speed, memory capacity is an important reason why a model must be spread across many GPUs. During training a GPU stores not only the model's parameters but also the gradients, the values the optimizer keeps and the intermediate results that the backward pass needs. This chapter first estimates the memory these data require, then introduces the basic operations GPUs use to exchange data and the five common ways of splitting a training job, each called a parallelism. Each way creates a different communication demand, which lands on a different kind of network link from Chapter 2.

第 1 章說明了大型語言模型為什麼需要數千顆 GPU 才能在合理時間內完成訓練。除了運算速度，記憶體容量也是模型必須分散到多顆 GPU 的重要原因。訓練時，GPU 除了儲存模型參數，還需要保留梯度、優化器狀態及反向傳播所需的中間結果。本章先估算這些資料的記憶體需求，再介紹 GPU 之間的基本通訊操作，以及五種常見的平行化方式。每種方式都會產生不同的通訊需求，對應到第 2 章介紹的不同網路鏈路。

All the models in this book use the Transformer architecture, a stack of many layers with a similar structure. Each layer consists mainly of two parts, Attention and an MLP. Attention lets each token update its representation from the information in the other tokens of the sequence, while the MLP processes each token on its own, and most of this computation is matrix multiplication. The model represents each token as a vector of length $h$, and $h$ is called the hidden dimension, which is 16,384 for Llama 3 405B.

本書討論的模型都採用 Transformer 架構，由多個結構相似的層堆疊而成。每層主要包含 Attention 和 MLP 兩個部分。Attention 讓每個 token 根據序列中其他 token 的資訊更新表示，而 MLP 則獨立處理每個 token。這些計算大多由矩陣乘法組成。模型使用長度為 $h$ 的向量表示每個 token，其中 $h$ 稱為隱藏維度（hidden dimension）。例如，Llama 3 405B 的隱藏維度為 16,384。

Memory for the model state模型狀態的記憶體需求

Training must store far more than the model weights. In the common recipe of BF16 mixed-precision training, trained with Adam, the forward and backward passes compute with 16-bit floating-point numbers, but the optimizer still keeps 32-bit master weights so that tiny parameter updates are not rounded away for lack of precision.

訓練時需要儲存的資料遠多於模型權重本身。以常見的 BF16 混合精度訓練搭配 Adam 優化器為例，前向與反向傳播使用 16-bit 浮點數進行計算，但優化器仍保留 32-bit 的主權重，以避免微小的參數更新因精度不足而被捨去。

In this setup each parameter needs 2 bytes for its BF16 weight, 2 bytes for its BF16 gradient and three FP32 values, namely the master weight (4 bytes), Adam's first moment $m$ (4 bytes) and its second moment $v$ (4 bytes). The last three together are called the optimizer state, so each parameter needs in total

在這種配置下，每個參數需要 2 bytes 的 BF16 權重、2 bytes 的 BF16 梯度，以及三組 FP32 資料，分別是主權重（4 bytes）、Adam 的一階動量 $m$（4 bytes）和二階動量 $v$（4 bytes）。後三者合稱優化器狀態（optimizer states），因此每個參數總共需要：

The optimizer state accounts for 12 of these bytes, the largest part of the model state and the first data worth spreading across GPUs (Section 3.3).

其中優化器狀態占了 12 bytes，是模型狀態中最大的一部分，也是最值得優先分散到多顆 GPU 的資料（第 3.3 節）。

This memory grows in direct proportion to the number of parameters. With 80 GB of HBM on one H100 and 16 bytes per parameter, models of different sizes need the following:

這項記憶體需求隨參數量直接增加。假設一顆 H100 GPU 配備 80 GB HBM，且每個參數需要 16 bytes，則不同規模模型的記憶體需求如下：

The model state of Llama 3 405B alone takes about 6.48 TB, the entire memory of 81 H100s, and even a 7B model already exceeds the capacity of one H100. These numbers do not yet include the activations and other temporary data, so real training needs more memory. Other optimizers or lower-precision optimizer state can reduce the requirement, but large models still have to spread their data across many GPUs.

Llama 3 405B 光是模型狀態就需要約 6.48 TB，相當於 81 顆 H100 的全部記憶體容量。即使是 7B 模型，也已經超過單顆 H100 的容量。這些數字尚未包含激活值與其他暫存資料，因此實際訓練需要更多記憶體。使用其他優化器或較低精度的優化器狀態可以減少需求，但大型模型仍然需要將資料分散到多顆 GPU。

For a mixture-of-experts (MoE) model, the memory requirement depends on the total number of parameters. Each token activates only a few experts, but the parameters of every expert must be stored in the system, so the size of the model cannot be estimated from the parameters active at any one time (Section 3.8).

對混合專家模型（Mixture of Experts, MoE）而言，記憶體需求取決於總參數量。雖然每個 token 只會啟用少數專家，但所有專家的參數都必須儲存在系統中，因此不能只根據每次啟用的參數量估算模型大小（第 3.8 節）。

Memory for activations激活值的記憶體需求

Besides the model state, training must also store the activations produced by the forward pass. These intermediate results are used in the backward pass to compute the gradients, so they must be kept until the corresponding backward computation has finished, and their memory grows with the sequence length, the batch size and the model's dimensions.

除了模型狀態，訓練還必須儲存前向傳播產生的激活值（activations）。這些中間結果會在反向傳播時用來計算梯度，因此需要保留到對應的反向運算完成。激活值的記憶體需求會隨序列長度、批次大小及模型維度增加。

To limit the memory of a single pass, training usually splits a large batch into several micro-batches and runs the forward and backward passes on them one after another. The gradients of the micro-batches are accumulated, and the parameters are updated once after all of them have been processed, which keeps the training batch size the same while reducing the activations that must be kept at any one time.

為了控制單次運算的記憶體需求，訓練通常將一個較大的 batch 拆成數個 micro-batches，依序執行前向與反向傳播。各個 micro-batch 的梯度會累積起來，等全部處理完畢後，再進行一次參數更新。如此便能在維持相同訓練 batch 大小的同時，減少單次需要保留的激活值。

The rest of this chapter describes the size of the work with one set of symbols. A sequence has $s$ tokens and a micro-batch holds $b$ sequences, so each micro-batch carries $bs$ tokens. Each copy of the model processes $G$ micro-batches before every parameter update, and $D$ copies run side by side under data parallelism (Section 3.3), so one optimizer step processes $G\,b\,D$ sequences, the global batch. The degrees of the other splits are $t$ for tensor parallelism and $p$ for the stages of pipeline parallelism, and the model has $L$ layers of hidden dimension $h$. Unless a number is marked per micro-batch, the byte counts of this chapter are per GPU per optimizer step, which is the unit in which Section 3.7 adds them up.

本章其餘部分用同一組符號描述工作量。一條序列有 $s$ 個 token，一個 micro-batch 包含 $b$ 條序列，因此每個 micro-batch 有 $bs$ 個 token。每份模型副本在每次參數更新前處理 $G$ 個 micro-batch，資料平行（第 3.3 節）讓 $D$ 份副本同時運作，所以一次參數更新處理 $G\,b\,D$ 條序列，也就是全域 batch。其他切分方式的度數記為：張量平行 $t$、管線平行的階段數 $p$；模型共有 $L$ 層，隱藏維度為 $h$。除非特別標明是每個 micro-batch，本章的位元組數都以每顆 GPU、每次參數更新計算，第 3.7 節也以這個單位把它們加總。

The values in the last column are those of Llama 3's 16,384-GPU run with 8K-token sequences (Llama 3 Table 4, second row), the layout that Section 3.7 works through in full.

最後一欄是 Llama 3 以 8K token 序列、16,384 顆 GPU 訓練時的數值（Llama 3 Table 4 第二列），第 3.7 節會完整推算這個配置。

A Transformer layer produces many kinds of intermediate results, including the projections of Attention, the hidden representations of the MLP and the data needed by operations such as normalization. According to Megatron's study of activation recomputation, for a conventional Transformer layer with sequence length $s$, micro-batch size $b$ and hidden dimension $h$, using 16-bit floating-point numbers, the activation memory is approximately

一個 Transformer 層會產生多種中間結果，包括 Attention 的投影、MLP 的隱藏表示，以及正規化等運算所需的資料。根據 Megatron 的激活值重算研究，對序列長度為 $s$、micro-batch 大小為 $b$、隱藏維度為 $h$ 的傳統 Transformer 層，使用 16-bit 浮點數時，激活值記憶體可以近似為：

The full formula also has a $5as^2b$ term for the intermediate matrices of Attention, where $a$ is the number of attention heads. Modern fused attention kernels avoid storing the full attention matrix, so this term is left out here. The coefficient 34 comes from Megatron's analysis of the conventional Transformer, and the actual value changes with the model's structure and with how the kernels are implemented.

完整公式還包含 Attention 中間矩陣的 $5abs^2$ 項，其中 $a$ 是 Attention head 數量。現代融合式 Attention kernel 可以避免儲存完整的 Attention 矩陣，因此這裡先省略該項。係數 34 來自 Megatron 對傳統 Transformer 的分析，實際數值會隨模型結構與 kernel 實作而改變。

Using the dimensions of Llama 3 405B, with sequence length $s=8192$, one sequence per micro-batch ($b=1$) and hidden dimension $h=16384$, each layer needs about

以 Llama 3 405B 的維度估算，假設序列長度 $s=8192$，每個 micro-batch 包含一條序列（$b=1$），隱藏維度 $h=16384$，則每層約需要：

Llama 3 405B has 126 layers, so keeping the activations of every layer would take about 575 GB by this formula, more than the memory of seven H100s. This shows why activations can still take a large share of memory even after the model's parameters have been spread across many GPUs.

Llama 3 405B 共有 126 層，若完整保留所有層的激活值，依此公式估算約需 575 GB，相當於超過七顆 H100 的記憶體。這也說明了為什麼即使模型參數已經分散到多顆 GPU，激活值仍然可能占用大量記憶體。

There are two main ways to reduce the memory of activations. The first is model parallelism, which spreads different layers or intermediate tensors across several GPUs so that each GPU keeps only part of them (Section 3.7).

降低激活值記憶體需求有兩種主要方法。第一種是透過模型平行化，將不同層或中間張量分散到多顆 GPU，使每顆 GPU 只需保存其中一部分（第 3.7 節）。

The second is activation checkpointing, also called activation recomputation. During the forward pass it keeps only some of the intermediate results, such as the input of each layer, and recomputes the other activations when the backward pass needs them. If each layer keeps only its input, the checkpointed data take about $2sbh$ bytes, far less than the $34sbh$ above. The price is that part of the forward computation must be run again during the backward pass, and recomputing every layer in full costs about one extra forward pass.

第二種是激活值檢查點（activation checkpointing），也稱為 activation recomputation。其做法是在前向傳播時只保留部分中間結果，例如每層的輸入，等反向傳播需要其他激活值時再重新計算。若每層只保存輸入，檢查點資料量約為 $2sbh$ bytes，相較於前面的 $34sbh$ 大幅降低。代價是反向傳播期間必須重新執行部分前向運算，若完整重算每層，額外計算量約相當於再進行一次前向傳播。

Distributed training must therefore account for the memory of both the model state and the activations. The first is set mainly by the number of parameters, while the second grows with the number of tokens and the model's dimensions. The different ways of splitting a job reduce the memory burden on each GPU by partitioning these data or by rearranging the computation, but each brings its own communication cost.

因此，分散式訓練必須同時考慮模型狀態與激活值的記憶體需求。前者主要由參數量決定，後者則隨 token 數量及模型維度增加。不同的平行化方式透過切分這些資料或重新安排計算，降低單顆 GPU 的記憶體負擔，但也會帶來不同的通訊成本。

3.2How GPUs exchange data: collectivesGPU 如何交換資料：集體操作

Once a model or its training data is spread across several GPUs, the GPUs need to exchange data. The simplest way is point-to-point communication, in which one GPU sends and another receives. Distributed training, however, usually needs a whole group of GPUs to exchange data together, for example to sum the gradients that each one computed or to gather model parameters spread over different GPUs. Communication in which a group of GPUs takes part together is called a collective operation.

當模型或訓練資料分散到多顆 GPU 後，GPU 之間就需要交換資料。最簡單的方式是點對點傳輸（point-to-point communication），由一顆 GPU 發送（send），另一顆 GPU 接收（receive）。然而，分散式訓練通常需要一整組 GPU 共同交換資料，例如將各自計算的梯度加總，或收集分散在不同 GPU 上的模型參數。這類由多顆 GPU 共同參與的通訊稱為集體操作（collective communication）。

Distributed training mainly uses four collectives. The table below assumes that $N$ GPUs take part and that $S$ is the size of the complete data being exchanged, and the traffic it lists is the number of bytes each GPU must send.

分散式訓練主要使用四種集體操作。下表假設有 $N$ 顆 GPU 參與，$S$ 表示要交換的完整資料大小。表中的通訊量是每顆 GPU 需要送出的位元組數。

The pictograms below show what each collective does to the buffers of four GPUs. A letter is a position in the buffer, the color says whose data a chunk holds, and Σ marks a chunk summed over all four GPUs.

下方的圖示顯示每種集體操作對四顆 GPU 的緩衝區做了什麼。字母是緩衝區中的位置，顏色表示該塊是誰的資料，Σ 表示對四顆 GPU 加總後的塊。

Of these, an all-reduce can be split into one reduce-scatter followed by one all-gather. First the GPUs sum their data and each takes a different part of the result, and then they exchange these parts so that every GPU ends with the complete sum. This is also how the ring all-reduce works.

其中，All-reduce 可以拆成一次 Reduce-scatter 加上一次 All-gather。首先，各 GPU 將資料加總，並分別取得結果的不同部分；接著交換這些部分，使每顆 GPU 都得到完整的加總結果。這也是環形 All-reduce的基本運作方式。

Another important property of collectives is that the traffic of each GPU usually does not shrink noticeably as more GPUs take part. For an all-reduce with the ring algorithm, for example, each GPU sends

集體操作的另一個重要特性是，每顆 GPU 的通訊量通常不會隨參與 GPU 數量增加而明顯減少。以 All-reduce 為例，使用環形演算法時，每顆 GPU 送出的資料量為：

When $N$ is large this approaches $2S$, so each GPU still exchanges roughly the same amount of data however many GPUs are added. The all-to-all follows a different pattern, because each GPU sends different data to every other GPU and the traffic becomes many-to-many. Chapter 4 describes the algorithms that carry out these operations and the time they take.

當 $N$ 很大時，這個值趨近 $2S$。因此，即使增加更多 GPU，每顆 GPU 仍須交換大致相同的資料量。All-to-all 的通訊模式則有所不同。每顆 GPU 都要向其他 GPU 傳送不同的資料，形成多對多的通訊需求。第 4 章將進一步介紹這些操作的執行演算法與通訊時間。

3.3Data parallelism and ZeRO資料平行與 ZeRO

Data parallelism (DP) is the most direct way to distribute training. Each GPU keeps a complete copy of the model, a replica, but processes different data of the training batch. After the GPUs finish their forward and backward passes independently, they hold different gradients, so to make every replica apply the same parameter update the GPUs must first sum and average the gradients with an all-reduce and only then run the optimizer.

資料平行（Data Parallelism, DP）是最直接的分散式訓練方式。每顆 GPU 保留一份完整的模型副本（replica），但處理訓練批次中不同的資料。各 GPU 獨立完成前向與反向傳播後，會得到不同的梯度。為了確保所有模型副本使用相同的參數更新，GPU 必須先透過 All-reduce 將梯度加總並取平均，再執行優化器更新。

If the model has $P$ parameters and the gradients are stored in BF16, the gradient on each GPU is $S=2P$ bytes. With a DP group of $N$ GPUs and a ring all-reduce, each GPU sends per step

假設模型有 $P$ 個參數，梯度以 BF16 儲存，則每顆 GPU 的梯度大小為 $S=2P$ bytes。若 DP 群組包含 $N$ 顆 GPU，使用環形 All-reduce 時，每顆 GPU 每步需要發送：

The BF16 gradients of a 70B model, for example, take about 140 GB, and a full all-reduce makes each GPU send about 280 GB, which shows why DP produces so much network traffic. In practice the gradients that are already computed can start their all-reduce during the backward pass, overlapping with the gradient computation of other layers and reducing the time the GPUs spend waiting for communication.

例如，一個 70B 模型的 BF16 梯度大小約為 140 GB，完整 All-reduce 的每 GPU 發送量約為 280 GB。這也說明了為什麼 DP 會產生大量網路流量。實際訓練可以在反向傳播時，讓已完成計算的梯度先開始 All-reduce，與其他層的梯度計算重疊，減少 GPU 等待通訊的時間。

The gradient size and the actual traffic need to be kept apart here. For a model with $P$ parameters the BF16 gradient is $2P$ bytes, but in one ring all-reduce each GPU sends about $4P$ bytes. Some papers report the former, such as HPN's 5.5 GB gradient message, while others count the bytes actually carried by the network.

這裡需要區分梯度大小與實際通訊量。對具有 $P$ 個參數的模型，BF16 梯度大小是 $2P$ bytes，但一次環形 All-reduce 中，每顆 GPU 送出的資料約為 $4P$ bytes。有些論文報告的是前者，例如 HPN 的 5.5 GB 梯度訊息；另一些則計算實際在網路上傳輸的位元組數。

ZeRO: spreading the model stateZeRO：分散模型狀態

Plain DP also wastes memory. Every GPU keeps the complete model weights, gradients and optimizer state and applies the same parameter update once the averaged gradient is available, so with the mixed-precision Adam training of Section 3.1, which needs 16 bytes per parameter, every DP replica stores the same data again.

一般 DP 還有記憶體浪費的問題。每顆 GPU 都保存完整的模型權重、梯度和優化器狀態，並在取得平均梯度後執行相同的參數更新。以第 3.1 節的混合精度 Adam 訓練為例，每個參數需要 16 bytes，因此所有 DP 副本都要重複儲存這些資料。

ZeRO (Zero Redundancy Optimizer) removes these duplicate copies step by step by spreading the model state across the GPUs of the DP group. ZeRO has three stages, and each stage shards one more kind of data:

ZeRO（Zero Redundancy Optimizer）透過將模型狀態分散到 DP 群組中的不同 GPU，逐步消除這些重複副本。ZeRO 分成三個階段，每個階段進一步切分一種資料（分片）：

• ZeRO-1 shards the optimizer state. Each GPU keeps only $1/N$ of the FP32 master weights and Adam moments and updates the corresponding parameters, while the BF16 weights and gradients are still fully replicated.
• ZeRO-2 also shards the gradients. A GPU obtains the gradients it is responsible for with a reduce-scatter, updates the corresponding parameters and finally synchronizes the updated weights with an all-gather.
• ZeRO-3 shards the model weights as well. Each GPU normally keeps only $1/N$ of the model state and gathers the complete weights of a layer with an all-gather only when that layer runs, after which they can be sharded again. PyTorch's Fully Sharded Data Parallel (FSDP) works this way.

• ZeRO-1 將優化器狀態分片。每顆 GPU 只保留 $1/N$ 的 FP32 主權重與 Adam 動量，並負責更新對應的參數。BF16 權重與梯度仍保留完整副本。
• ZeRO-2 進一步將梯度分片。GPU 透過 Reduce-scatter 取得自己負責的梯度，再更新對應的參數，最後利用 All-gather 同步更新後的權重。
• ZeRO-3 連模型權重也一起分片。每顆 GPU 平時只保存 $1/N$ 的模型狀態，執行某層時才透過 All-gather 收集該層的完整權重，計算完成後即可重新分片。PyTorch 的 Fully Sharded Data Parallel（FSDP）採用這類方式。

With a DP group of $N$ GPUs, the memory each stage keeps resident per parameter is:

假設 DP 群組有 $N$ 顆 GPU，各階段每個參數所需的常駐記憶體如下：

As Figure 3.3 shows, each ZeRO stage shards one more part of the model state, and at $N=4$ ZeRO-3 leaves each GPU with only a quarter of the model state resident. Running a layer, however, still needs extra temporary space for the weights collected by the all-gather.

如圖 3.3 所示，ZeRO 每增加一個階段，就多分片一部分模型狀態。當 $N=4$ 時，ZeRO-3 讓每顆 GPU 只需常駐保存四分之一的模型狀態。不過，執行某層時仍需要額外的暫存空間，放置 All-gather 收集到的權重。

What the two collectives carry兩次集體操作搬的是什麼

The two collectives carry different data in the two methods. In plain DP, the all-gather that completes the all-reduce spreads the averaged gradients, and every GPU then repeats the whole Adam update with its full optimizer state. In ZeRO-2, each GPU uses the reduce-scattered average to update only the shard it is responsible for, and the all-gather spreads the updated weights instead, which yields the same model because Adam still sees the averaged gradient. ZeRO-1 follows the same sequence but keeps its full local gradient buffer, while ZeRO-3 moves its all-gather of the weights to just before each layer runs (Panel 3.1 steps through all four methods).

兩種方法中，這兩次集體操作搬運的資料並不相同。一般 DP 中，完成 All-reduce 的 All-gather 散播的是平均梯度，之後每顆 GPU 再用完整的優化器狀態重複整個 Adam 更新。ZeRO-2 則讓每顆 GPU 用 Reduce-scatter 得到的平均梯度只更新自己負責的那一片，All-gather 散播的改為更新後的權重，由於 Adam 用的仍是平均梯度，得到的模型完全相同。ZeRO-1 的順序相同，但保留完整的本地梯度緩衝區，ZeRO-3 則把權重的 All-gather 移到每一層執行之前（面板 3.1 逐步呈現這四種方法）。

Animation 3.1 uses four GPUs to show how ZeRO-3 first gathers the weights of a layer with an all-gather, runs the computation and then distributes the gradients with a reduce-scatter.

動畫 3.1 以四顆 GPU 為例，展示 ZeRO-3 如何先透過 All-gather 收集一層的權重，執行計算後再以 Reduce-scatter 分配梯度。

Gradient accumulation梯度累加

One training step usually contains several micro-batches. The GPUs accumulate the gradients computed for each micro-batch and synchronize the gradients and update the model only after all of them are done, so plain DP needs only one gradient all-reduce per step, and its traffic is set mainly by the size of the model and does not depend on the number of micro-batches. How often ZeRO-3 gathers the weights also depends on whether the full weights are kept from one micro-batch to the next.

一次訓練步驟通常包含多個 micro-batch。GPU 會先累積各個 micro-batch 計算出的梯度，等全部完成後才同步梯度並更新模型。因此，一般 DP 每步只需要執行一次梯度 All-reduce，其通訊量主要由模型大小決定，與 micro-batch 數量無關。ZeRO-3 的權重收集次數則還取決於是否在不同 micro-batch 之間保留完整權重。

As the batch size grows, the same cost of gradient synchronization is spread over more tokens, which lowers the average DP traffic per token. This is why optimizers that support large-batch training, such as LAMB, also help reduce communication cost besides improving training efficiency (MegaScale §3.1). Llama 3 further accumulates and reduce-scatters the gradients in FP32 for numerical stability, which also raises the gradient traffic from 2 to 4 bytes per parameter.

當 batch size 增加時，同樣的梯度同步成本可以分攤到更多 token，使每個 token 的平均 DP 通訊成本下降。這也是 LAMB 等支援大型 batch 訓練的優化器除了改善訓練效率，也有助於降低通訊成本的原因（MegaScale §3.1）。Llama 3 進一步使用 FP32 進行梯度累積與 Reduce-scatter，以提高數值穩定性，但也讓梯度通訊量從每參數 2 bytes 增加至 4 bytes。

The analysis of this section assumes one gradient synchronization per optimizer step. A real step contains many micro-batches and, with pipeline parallelism, a pipeline schedule, and both change how often ZeRO must communicate, so Section 3.5 returns to ZeRO once pipelines have been introduced.

本節的分析假設每次參數更新只同步一次梯度。實際的一步包含許多 micro-batch，搭配管線平行時還有管線排程，兩者都會改變 ZeRO 的通訊次數，因此第 3.5 節在介紹管線之後會再回到 ZeRO。

3.4Tensor parallelism: communication inside every layer張量平行：每層內部的通訊

Tensor parallelism (TP) spreads the weight matrices of every layer of the model across several GPUs, so that each GPU stores and computes only part of them and the GPUs then combine their results through communication. This lowers the memory each GPU needs and also lets several GPUs carry out together a matrix computation that one GPU would otherwise do alone.

張量平行（Tensor Parallelism, TP）將模型每一層的權重矩陣分散到多顆 GPU，讓每顆 GPU 只需儲存並計算其中一部分，再透過通訊合併結果。這樣既能降低單顆 GPU 的記憶體需求，也能讓多顆 GPU 同時執行原本由一顆 GPU 負責的矩陣運算。

The communication of TP happens inside every layer, and both the forward and the backward passes need to exchange data. These exchanges carry mainly activations, and the computation that follows usually has to wait until the exchange is complete, so TP communication sits on the critical path of training and is especially sensitive to communication latency and bandwidth. To lower this cost, real systems usually place a TP group inside one server, where it communicates at high speed over NVLink and NVSwitch, and a common TP degree is 8, matching a server with eight GPUs.

TP 的通訊發生在每一層內部，前向與反向傳播都需要交換資料。這些交換主要傳輸激活值，且後續計算往往必須等待交換完成才能繼續。因此，TP 通訊位於訓練的關鍵路徑（critical path）上，對通訊延遲與頻寬特別敏感。為了降低這項成本，實際系統通常將 TP 群組安排在同一台伺服器內，利用 NVLink 和 NVSwitch 進行高速通訊。常見的 TP 度數為 8，對應一台具有八顆 GPU 的伺服器。

Megatron's matrix splitMegatron 的矩陣切分

Megatron-LM combines two ways of splitting a matrix to reduce the communication each layer needs. Take the MLP of a Transformer as an example, whose computation consists of two matrix multiplications:

Megatron-LM 透過搭配兩種矩陣切分方式，減少每層需要執行的通訊操作。以 Transformer 的 MLP 為例，其計算包含兩次矩陣乘法：

Here $X$ is the input activation of size $h\times bs$, where $h$ is the hidden dimension and $bs$ is the number of tokens in the micro-batch. The first weight matrix $A$ expands the hidden dimension from $h$ to $4h$, and the second matrix $B$ shrinks it back to $h$.

其中，$X$ 是大小為 $h\times bs$ 的輸入激活值，$h$ 是隱藏維度，$bs$ 是 micro-batch 中的 token 數量。第一個權重矩陣 $A$ 將隱藏維度從 $h$ 擴展至 $4h$，第二個矩陣 $B$ 再將其縮回 $h$。

With a TP group of $t$ GPUs, Megatron splits $A$ into $t$ parts along its output dimension, so each GPU holds one $A_k$ and computes on its own

假設 TP 群組包含 $t$ 顆 GPU，Megatron 將 $A$ 沿輸出維度分成 $t$ 份，每顆 GPU 持有一個 $A_k$，並獨立計算：

Because GeLU acts on each element separately, each GPU can process its own part directly without exchanging intermediate results.

由於 GeLU 是逐元素運算，各 GPU 可以直接處理自己負責的部分，無須交換中間結果。

Megatron then splits $B$ into the matching $t$ parts along its input dimension. Each GPU computes a partial output from its own $Y_k$ and $B_k$, and an all-reduce finally adds these results together:

接著，Megatron 將 $B$ 沿輸入維度分成對應的 $t$ 份。每顆 GPU 使用自己的 $Y_k$ 與 $B_k$ 計算部分輸出，最後透過 All-reduce 將這些結果加總：

In this way the two matrix multiplications of the MLP need only one all-reduce at the end (Figure 3.4). Attention follows the same approach, assigning the attention heads to different GPUs that compute independently and then adding up the partial results of the output projection. The backward pass needs the corresponding sum across GPUs when it computes the gradient of the input.

這樣一來，MLP 的兩次矩陣乘法只需要在最後執行一次 All-reduce（圖 3.4）。Attention 也採用類似方式，先將 Attention heads 分配給不同 GPU 獨立計算，再將輸出投影的部分結果加總。反向傳播則需要在計算輸入梯度時進行對應的跨 GPU 加總。

Standard Megatron TP therefore needs four all-reduces per Transformer layer and per micro-batch, two in the forward pass and two in the backward pass.

因此，在標準 Megatron TP 中，每個 Transformer 層、每個 micro-batch 需要四次 All-reduce，前向兩次、反向兩次。

Figures 3.5 and 3.6 show how Llama 3 405B applies this split with TP = 8. In attention, $W_Q$, $W_K$ and $W_V$ are split by heads, that is along their output dimension, so each GPU computes 16 query heads with one KV head entirely on its own, and $W_O$ is split along its input dimension, which leaves each GPU with a partial sum of the output. In the MLP, the gate and up projections of SwiGLU are split along the FFN dimension into slices of 6,656, and the down projection along its input dimension. Each block therefore has two points of communication. At its output, marked g, the forward pass all-reduces the partial sums, and at its input, marked f, the backward pass all-reduces the gradient of $X$, because every GPU used the full $X$ and the gradients that the slices contribute to it must be added. The column and row splits stay the same in both passes, and only the point that carries the all-reduce moves from g in the forward pass to f in the backward pass.

圖 3.5 與 圖 3.6 顯示 Llama 3 405B 在 TP = 8 下如何套用這種切分。在注意力中，$W_Q$、$W_K$ 與 $W_V$ 依注意力頭切分，也就是沿輸出維度切分，因此每顆 GPU 可以完全獨立地計算 16 個 query 頭與 1 個 KV 頭，而 $W_O$ 沿輸入維度切分，使每顆 GPU 得到輸出的部分和。在 MLP 中，SwiGLU 的 gate 與 up 投影沿 FFN 維度切成每片 6,656，down 投影則沿輸入維度切分。因此每個區塊有兩個通訊點。在區塊輸出端的 g，前向傳播對部分和做 All-reduce；在區塊輸入端的 f，反向傳播對 $X$ 的梯度做 All-reduce，因為每顆 GPU 都用了完整的 $X$，各切片對它的梯度必須加總。前向與反向使用同一組行列切分，改變的只是 All-reduce 的位置，前向在 g，反向在 f。

Megatron later added sequence parallelism (SP), which spreads operations such as LayerNorm and Dropout across GPUs along the sequence dimension and so removes duplicated copies of activations. With SP, each all-reduce is carried out instead by a reduce-scatter together with an all-gather. Under the ideal ring model the total bytes transferred stay the same, while the memory of some activations falls to $1/t$ of what it was.

Megatron 後來進一步加入序列平行（Sequence Parallelism, SP），將 LayerNorm、Dropout 等運算沿序列維度分散至不同 GPU，減少重複儲存的激活值。加入 SP 後，原本的 All-reduce 會改由 Reduce-scatter 與 All-gather 配合完成。在理想的環形通訊模型下，總傳輸位元組數維持相同，但部分激活值的記憶體需求可降低至原本的 $1/t$。

Why TP usually stays inside the server為什麼 TP 通常留在伺服器內？

The communication cost of TP comes mainly from two factors. First, every layer must exchange activations of full size, a single message can reach tens of megabytes, and the traffic of a whole layer can add up to hundreds of megabytes. Second, these operations usually sit between adjacent matrix multiplications, so the GPU must obtain the combined result before it can continue with the next computation. The gradient synchronization of DP, by contrast, can proceed layer by layer during the backward pass and is easier to overlap with the computation of other layers.

TP 的通訊成本主要來自兩個因素。首先，每層都需要交換完整大小的激活值，單次訊息可能達到數十 MB，整層累積的傳輸量可達數百 MB。其次，這些通訊操作通常位於相鄰矩陣乘法之間，GPU 必須取得合併後的結果，才能繼續後面的計算。相較之下，DP 的梯度同步可以隨著反向傳播逐層進行，較容易與其他層的計算重疊。

This makes TP particularly dependent on fast GPU interconnects. An H100 has up to 900 GB/s of NVLink bandwidth counted in both directions, about 450 GB/s in each direction, while a 400 Gb/s NIC offers only 50 GB/s. Placing a TP group inside one eight-GPU server lets its collectives use the high bandwidth of NVSwitch and avoids frequent transfers of activations over the slower network between servers.

這使得 TP 特別依賴高速的 GPU 互連。以 H100 為例，每顆 GPU 的 NVLink 雙向總頻寬最高為 900 GB/s，約相當於每方向 450 GB/s，而一張 400 Gb/s NIC 的頻寬只有 50 GB/s。若將 TP 群組放在同一台八 GPU 伺服器內，便可以利用 NVSwitch 提供的高頻寬進行集體操作，避免頻繁透過較慢的跨機網路傳輸激活值。

Take the dimensions of GPT-3 175B as an example, with hidden dimension $h=12{,}288$, 2,048 tokens per micro-batch, BF16 activations and TP degree $t=8$. The message of each all-reduce is

以 GPT-3 175B 的模型維度為例，假設隱藏維度 $h=12{,}288$，每個 micro-batch 包含 2,048 個 token，激活值使用 BF16，且 TP 度數 $t=8$。每次 All-reduce 的訊息大小為：

With a ring all-reduce each GPU sends $2(t-1)S/t$ bytes per all-reduce, so the four all-reduces of one layer's forward and backward passes send in total

使用環形 All-reduce 時，每顆 GPU 每次需要傳送 $2(t-1)S/t$ bytes。因此，一層前向與反向傳播的四次 All-reduce，合計需要傳送：

To compare communication with computation, assume an effective GPU rate of 400 TFLOP/s. The matrix multiplications of a Transformer layer's forward and backward passes take about $72bsh^2/t$ FLOPs, ignoring the attention scores, which with the parameters above gives about 7.0 ms of computation per GPU. Without any overlap of communication and computation, the result is:

為了比較通訊與運算時間，假設 GPU 的有效計算速率為 400 TFLOP/s，Transformer 層的前向與反向矩陣乘法合計約需 $72bsh^2/t$ FLOPs，忽略 Attention 分數計算。代入上述參數，每顆 GPU 的計算時間約為 7.0 ms。若不考慮通訊與計算重疊，結果如下：

The example shows that with the same matrix split, TP communication over the NICs between servers can take nearly as long as the computation and leave the GPUs waiting for data much of the time, while NVLink cuts this cost sharply.

這個例子顯示，同樣的矩陣切分方式，若透過跨機 NIC 進行 TP 通訊，通訊時間可能接近計算時間，使 GPU 有大量時間等待資料。使用 NVLink 則能大幅降低這項成本。

The 400G column is a simplified model in which every byte of the collective crosses one NIC. A TP group that spans servers would run its collectives over both NVLink and the NICs, with hierarchical algorithms that keep part of the traffic inside each server, so the real cost depends on the topology and the algorithm (Chapter 4).

400G 這一欄是簡化模型，假設集體操作的每個位元組都經過一張 NIC。跨伺服器的 TP 群組實際上會同時使用 NVLink 與 NIC，並以階層式演算法讓部分流量留在伺服器內，因此實際成本取決於拓樸與演算法（第 4 章）。

The difference also shows in the ratio of communication time to computation time. With GPU compute rate $F$ and communication bandwidth $B$, the simplified model above gives

這個差異也可以從通訊時間與計算時間的比例看出。設 GPU 計算速率為 $F$、通訊頻寬為 $B$，在上述簡化模型中：

The ratio does not depend on the number of tokens in the micro-batch and is set mainly by the TP degree, the model's hidden dimension and the ratio of the GPU's compute rate to the communication bandwidth, so when GPU compute grows without a matching rise in communication bandwidth, the communication bottleneck of TP becomes more pronounced.

可以看到，比例與 micro-batch 的 token 數無關，主要取決於 TP 度數、模型隱藏維度，以及 GPU 計算能力和通訊頻寬之間的比例。當 GPU 計算能力提高，卻沒有相應增加通訊頻寬時，TP 的通訊瓶頸就會更加明顯。

Real systems also try to overlap TP communication with computation. MegaScale cuts the matrix multiplications of the MLP into smaller blocks so that the all-gather or reduce-scatter of some data runs at the same time as the matrix computation of other blocks, and in its ablation study this TP overlap raised MFU by 2.2 percentage points (MegaScale §3.2, §6.1).

實際系統也會設法重疊 TP 的通訊與計算。MegaScale 將 MLP 的矩陣乘法切成較小的區塊，讓部分資料的 All-gather 或 Reduce-scatter 與其他區塊的矩陣運算同時進行。在其消融實驗中，TP 通訊重疊使 MFU 提升了 2.2 個百分點（MegaScale §3.2、§6.1）。

Large-model training therefore usually keeps TP inside one server, for example with a TP group of eight GPUs, and scales to more servers with other kinds of parallelism. If the TP degree rises to 16, the group must span two eight-GPU servers, and the bandwidth and latency of the NICs between them become important factors in performance.

因此，大型模型訓練通常優先將 TP 限制在單台伺服器內，例如使用八顆 GPU 組成一個 TP 群組，再透過其他平行化方式擴展到更多伺服器。若 TP 度數增加至 16，群組就必須跨越兩台八 GPU 伺服器，跨機 NIC 的頻寬與延遲也會成為影響效能的重要因素。

Panel 3.2 compares the time of the matrix multiplications and of the collectives for different TP degrees, hidden dimensions and numbers of tokens per micro-batch. Its values use the simplified model without overlap to show how communication bandwidth affects TP performance, while real training time also depends on the efficiency of the collectives, the GPU utilization and how far communication and computation can overlap.

面板 3.2 比較不同 TP 度數、隱藏維度與 micro-batch token 數下，矩陣乘法和集體操作所需的時間。數值使用不重疊的簡化模型，反映通訊頻寬對 TP 效能的影響；實際訓練時間則還取決於集體操作效率、GPU 利用率，以及通訊與計算能夠重疊的程度。

3.5Pipeline parallelism: stages, micro-batches and bubbles管線平行：階段、micro-batch 與氣泡

Pipeline parallelism (PP) assigns consecutive layers of the model to different groups of GPUs, and each group is called a pipeline stage. In the forward pass the stages compute in turn and pass their output activations to the next stage, and in the backward pass the gradients of those activations travel back in the opposite direction. Stages exchange data with point-to-point send and receive, and because the traffic is usually far smaller than that of DP and TP, PP suits deployments that span racks, a property Meta uses to place different pipeline stages far apart in its network (Chapter 6).

管線平行（Pipeline Parallelism, PP）將模型的連續層分配給不同的 GPU 群組，每個群組稱為一個管線階段（stage）。前向傳播時，各階段依序計算，並將輸出的激活值傳給下一階段；反向傳播則沿相反方向傳回激活值的梯度。階段間透過點對點的 send 與 receive 傳遞資料，通訊量通常遠小於 DP 與 TP，因此適合跨機櫃部署。Meta 便利用這個特性，將不同管線階段安排在網路上相距較遠的位置（第 6 章）。

If the whole batch went through the pipeline at once, the stages could only work in turn and most GPUs would be waiting. To let the stages work at the same time, PP cuts the batch into many micro-batches that flow through the pipeline one after another. When the second stage starts on the first micro-batch, the first stage can move on to the next one, so the stages work like a conveyor belt, each handling a different micro-batch at the same time.

若整個 batch 一次通過管線，各階段便只能輪流工作，大部分 GPU 都處於等待狀態。為了讓不同階段同時執行，PP 將 batch 切成多個 micro-batch，讓它們依序流過管線。當第二階段開始處理第一個 micro-batch，第一階段便能接著處理下一個。如此一來，多個階段便能像輸送帶一樣，同時處理不同的 micro-batch。

The bubble氣泡

Some stages still have to wait at the start and at the end of the pipeline. At first, the later stages wait for activations from the earlier ones, and near the end of the backward pass the earlier stages finish their work one after another and sit idle. These waiting times are called the pipeline bubble.

管線開始與結束時，仍有部分階段必須等待。最初，後面的階段要等前面的激活值傳來；到了反向傳播的尾聲，前面的階段則會陸續完成工作並閒置。這些等待時間稱為管線氣泡（pipeline bubble）。

In the GPipe schedule, all micro-batches first finish their forward passes and then run their backward passes together. With $p$ stages, $G$ micro-batches per step and equal compute time in every stage, the bubble fraction is

在 GPipe 排程中，所有 micro-batch 先完成正向傳播，再統一執行反向傳播。假設管線有 $p$ 個階段、每步處理 $G$ 個 micro-batch，且各階段計算時間相同，氣泡比例為：

With $p=8$ and $G=8$, for example, 7 of the 15 time units are idle, a bubble of 47%, and raising $G$ to 64 brings it down to about 10%. With the number of stages and the schedule fixed, more micro-batches per step is therefore the most direct way to shrink the bubble.

例如，當 $p=8$、$G=8$ 時，總共 15 個時間單位中有 7 個閒置，氣泡佔 47%。若將 $G$ 增加到 64，氣泡便降至約 10%。因此，在階段數與排程固定時，增加每步的 micro-batch 數量，是降低氣泡最直接的方法。

This ties the efficiency of the pipeline to the training batch size. MegaScale uses the LAMB optimizer, which allows a batch four times larger, and reports 87.5% fewer pipeline bubbles (MegaScale §3.1). By the formula above, four times as many micro-batches of a fixed size cut the bubble relative to compute time by 75%, and the further improvement that the paper reports also involves differences in scheduling.

這也讓管線效率與訓練的 batch size 產生關聯。MegaScale 採用 LAMB 最佳化器，允許使用四倍大的 batch，並報告管線氣泡減少 87.5%（MegaScale §3.1）。依上述公式，若 micro-batch 大小固定，數量增加四倍會使氣泡相對計算時間減少 75%；論文報告的額外改善也涉及排程上的差異。

Another common schedule is 1F1B (one forward, one backward). After a warm-up, it alternates the forward pass of a new micro-batch with the backward pass of an earlier one. Its bubble fraction equals that of GPipe, but the first stage needs to keep the activations of only about $p$ micro-batches instead of all $G$, which lowers memory considerably, and Section 3.7 assumes this schedule.

另一種常見排程是1F1B（One Forward, One Backward）。它在暖機後交替執行一個新 micro-batch 的正向傳播，以及一個較早 micro-batch 的反向傳播。1F1B 的氣泡比例與 GPipe 相同，但第一階段最多只需保留約 $p$ 個 micro-batch 的激活值，而 GPipe 必須保留全部 $G$ 個，因此能顯著降低記憶體需求。第 3.7 節採用的便是這種假設。

Animation 3.2 steps through the warm-up, steady state and drain of the pipeline one slot at a time and counts the idle time of every stage. By default it uses $p=4$ and $G=8$, with a forward computation taking one slot and a backward computation two, and every operation starts as soon as its stage is free and its input has arrived, with transfers taking no time. In this setting a step takes $3(G+p-1)=33$ slots, of which each stage is idle for $3(p-1)=9$, a bubble of about 27%. Switching between GPipe and 1F1B shows that the two have the same bubble while the first stage holds a different number of activations in flight, and Panel 3.3 compares more stages and interleaved layouts.

動畫 3.2 逐格顯示管線的暖機、穩態與收尾，並統計每個階段的閒置時間。預設使用 $p=4$、$G=8$，每次正向計算佔一格，反向計算佔兩格。所有操作都在階段空閒且輸入已到達時開始，傳輸時間視為零。在此設定下，一步需要 $3(G+p-1)=33$ 格，每個階段有 $3(p-1)=9$ 格閒置，氣泡比例約為 27%。切換 GPipe 與 1F1B，可以觀察兩者相同的氣泡比例，以及第一階段在途激活值數量的差異。更多階段數與交錯配置可在面板 3.3中比較。

Interleaving交錯

An interleaved pipeline further splits each stage into $v$ smaller chunks of layers and assigns them to the same group of GPUs, so each GPU holds several non-adjacent runs of layers and one micro-batch passes through the same group of GPUs several times (Figure 3.8).

交錯管線（Interleaved Pipeline）進一步將每個階段拆成 $v$ 個較小的 layer chunk，並分配給同一組 GPU。各 GPU 因此持有多段不相鄰的層，一個 micro-batch 也會多次經過同一組 GPU（圖 3.8）。

Because each chunk takes only $1/v$ of the original compute time, the overhead of warm-up and drain falls with it, and the ratio of bubble to compute time becomes (Llama 3 §3.3.2)

由於每個 chunk 的計算時間縮短為原本的 $1/v$，暖機與收尾的開銷也隨之降低。其氣泡與計算時間的比例為（Llama 3 §3.3.2）：

Take a 16-layer model with $p=4$. A conventional pipeline ($v=1$) gives each GPU four consecutive layers, and each micro-batch crosses three stage boundaries. With $v=2$, each GPU instead holds two non-adjacent runs of layers, GPU 1 for example layers 1–2 and 9–10, and a micro-batch then crosses seven boundaries.

以 16 層模型、$p=4$ 為例，傳統管線（$v=1$）讓每顆 GPU 持有連續四層，每個 micro-batch 依序跨越三個階段邊界。當 $v=2$ 時，每顆 GPU 改為持有兩段不相鄰的層，例如 GPU 1 持有第 1–2 層與第 9–10 層。此時 micro-batch 需要跨越七個邊界。

The price of interleaving is more communication. The stage boundaries each micro-batch crosses rise from $p-1$ to $vp-1$, so the number of point-to-point messages grows about $v$ times while the size of each message stays the same.

交錯的代價是更多通訊。每個 micro-batch 經過的階段邊界由 $p-1$ 增加到 $vp-1$，點對點訊息數量約增加 $v$ 倍，而每則訊息的大小維持不變。

Llama 3 makes three further adjustments to the pipeline load. It makes the number of micro-batches in flight a tunable parameter, so that memory use can be traded against the bubble. The first stage, which besides its Transformer layers stores the embedding table and more activations from the warm-up, is under the most memory pressure, while the last stage computes the output projection and the loss and so carries the heaviest compute, so Llama 3 removes one Transformer layer from each end to balance the load of the stages (Llama 3 §3.3.2).

Llama 3 另外針對管線負載做了三項調整。首先，它將在途 micro-batch 數量設為可調參數，允許在記憶體使用量與氣泡之間取捨。其次，第一階段除了 Transformer 層，還須儲存 embedding 表與較多暖機期間的激活值，因此記憶體壓力最大；最後階段則負責輸出投影與損失計算，計算負擔較重。Llama 3 從兩端各移除一層 Transformer，以平衡各階段的負載（Llama 3 §3.3.2）。

MegaScale instead improves the scheduling of transfers between stages. During the warm-up it decouples send from receive, so that a stage can keep computing on inputs it has already received while data is sent in the background, and by overlapping this communication with computation MegaScale reports an MFU gain of 2.5 percentage points (MegaScale §3.2).

MegaScale 則改善階段間的傳輸排程。它在暖機期間將 send 與 receive 解耦，讓階段能在背景傳送資料的同時，繼續處理已收到的輸入。透過通訊與計算重疊，MegaScale 報告 MFU 提高了 2.5 個百分點（MegaScale §3.2）。

Panel 3.3 lets you set the number of stages $p$, the micro-batches $G$ and the interleaved chunks $v$, and switch between GPipe, 1F1B and interleaved schedules, while it shows the bubble fraction, the activation peak and the number of point-to-point messages. With the default $p=8$, $G=8$ and $v=1$, the bubble is 87.5% of the compute time and 46.7% of the total, assuming that a backward computation takes twice as long as a forward one, that each stage runs its operations in order and that it starts as soon as its input arrives. Try first raising $G$ above $p$ to compare the activation peaks of GPipe and 1F1B, then set $p=16$ and $G=16$ and raise $v$ step by step to watch the bubble fall as $(p-1)/(vG)$ while the traffic grows. The interleaved schedule follows Megatron's order of operations and needs $G$ to be a multiple of $p$, and if it is not, the panel rounds up and marks the adjustment.

面板 3.3 可調整階段數 $p$、micro-batch 數 $G$、交錯 chunk 數 $v$，並切換 GPipe、1F1B 與交錯排程。它同時呈現氣泡比例、激活值峰值與點對點訊息數量。預設 $p=8$、$G=8$、$v=1$ 時，氣泡相對計算時間為 87.5%，佔總時間的 46.7%。圖中假設反向計算花費正向的兩倍時間，各階段依序執行操作，並在輸入到達後立即開始計算。可以先將 $G$ 增加到超過 $p$，比較 GPipe 與 1F1B 的激活值峰值；再設定 $p=16$、$G=16$，逐步增加 $v$，觀察氣泡隨 $(p-1)/(vG)$ 下降，以及通訊量增加的情況。交錯排程使用 Megatron 的操作順序，要求 $G$ 為 $p$ 的倍數；若設定不符，面板會自動向上取整並標示調整結果。

Bytes per boundary每個邊界的位元組

What a stage sends to the next is the activation tensor that one micro-batch produces at the stage boundary, and in BF16 each message is

階段間傳送的資料，是一個 micro-batch 在階段邊界產生的激活值張量。若使用 BF16，每則訊息的大小為：

where $h$ is the hidden size. Megatron splits this tensor across the $t$ TP GPUs of the stage, so each GPU sends only $1/t$ of it to the corresponding GPU of the next stage (Figure 3.9), and the backward pass returns the gradients in the same way. In each step every GPU sends $G$ messages in each direction, multiplied by $v$ under an interleaved schedule.

其中 $h$ 為 hidden size。Megatron 會將這個張量切分到階段內的 $t$ 顆 TP GPU，因此每顆 GPU 只需將其中的 $1/t$ 傳給下一階段對應的 GPU（圖 3.9）。反向傳播則以相同方式傳回梯度。每步中，每顆 GPU 在每個方向傳送 $G$ 則訊息；若使用交錯排程，訊息數還需乘上 $v$。

For GPT-3 175B, with one 2,048-token sequence per micro-batch and $G=16$:

以 GPT-3 175B 為例，假設每個 micro-batch 包含一條 2,048-token 序列，且 $G=16$：

A single message of about 6 MB matches HPN's Table 3, and the roughly 200 MB per step is three orders of magnitude below the traffic of TP.

單則訊息約為 6 MB，與 HPN Table 3 相符。每步的資料量約為 200 MB，比 TP 通訊低三個數量級。

PP messages are small, each GPU talks only to fixed partners in the neighbouring stages, and an activation can be sent asynchronously as soon as it is ready, so PP tolerates longer network paths. This is why Alibaba places pipeline traffic on its slower links and Meta deploys different stages in different racks (HPN §7, Llama 3 §3.3.2; Chapter 6).

PP 的訊息量較小，而且每顆 GPU 只需與相鄰階段的固定夥伴通訊，激活值準備好後即可非同步傳送，因此適合經過較長的網路路徑。這也是阿里巴巴將管線流量安排在較慢的鏈路上，以及 Meta 將不同階段部署於不同機櫃的原因（HPN §7、Llama 3 §3.3.2；第 6 章）。

How gradient accumulation and the pipeline schedule change ZeRO梯度累加與管線排程如何影響 ZeRO

In real training a micro-batch usually holds one to a few sequences of 2,048 to 8,192 tokens, and large models keep it small because of memory and reach a large batch by accumulating many micro-batches. With the symbols of Section 3.1, each copy of the model accumulates $G$ micro-batches before an update, so one optimizer step processes

實際訓練中，一個 micro-batch 通常包含一到數條長 2,048 到 8,192 個 token 的序列。大型模型受記憶體限制而使用較小的 micro-batch，再透過累積許多 micro-batch 達到很大的 batch。沿用第 3.1 節的符號，每份模型副本在更新前累積 $G$ 個 micro-batch，因此一次參數更新處理的 token 數為：

TP, PP and CP split one replica of the model across several GPUs, so only the DP degree $D$ multiplies the batch. The benchmark configurations that NVIDIA Megatron Bridge publishes show how large $G$ is in practice:

TP、PP 與 CP 把一份模型副本切到多顆 GPU 上，所以只有 DP 度數 $D$ 會放大 batch。NVIDIA Megatron Bridge 公開的效能測試配置顯示了 $G$ 在實務上有多大：

Because $G$ is tens to hundreds, what matters is how each method handles the gradients of many micro-batches. The ZeRO paper counts one gradient reduction per training step and does not discuss gradient accumulation (ZeRO §7.2), but plain DP and ZeRO-1 keep a full gradient buffer, accumulate locally and reduce once, whereas ZeRO-2 keeps only $1/D$ of the gradients and so reduce-scatters after every micro-batch by default in DeepSpeed, and ZeRO-3 also gathers the weights again for every micro-batch. The table below counts, for each method and schedule, the model state a GPU holds during the step and the bytes it sends per optimizer step in units of $fS$, where $f=(D-1)/D$. Black rows are the schedules used in practice, and grey rows are possible but rarely chosen.

因為 $G$ 是數十到數百，關鍵在於各方法如何處理許多 micro-batch 的梯度。ZeRO 論文以每個訓練步驟歸約一次梯度計算，並未討論梯度累加（ZeRO §7.2）。一般 DP 與 ZeRO-1 保留完整的梯度緩衝區，可以在本地累積、只歸約一次；ZeRO-2 只保留 $1/D$ 的梯度，因此在 DeepSpeed 的預設排程中每個 micro-batch 之後都要做 Reduce-scatter，而 ZeRO-3 還要為每個 micro-batch 重新收集權重。下表列出各方法與排程在整步期間每顆 GPU 占用的模型狀態，以及每次參數更新的發送量，以 $fS$ 為單位，其中 $f=(D-1)/D$。黑色列是實務上採用的排程，灰色列可行但很少選用。