Chapter 3: Sharded Matrices and How to Multiply Them

Scaling Book Exercises – Chapter 3

Chapter: 3 – Sharded Matrices and How to Multiply Them
Source: handwritten Onyx Boox A4 PDF
Note: Transcribed from handwritten solutions; book markdown used only for exercise statements and notation.

Exercise 1 – replicated sharding

Exercise statement

An array is sharded $𝐴 [𝐼_{𝑋}, 𝐽, 𝐾, \dots]$ (i.e., only sharded across $𝑋$ ), with a mesh Mesh({'X': 4, 'Y': 8, 'Z': 2}). What is the ratio of the total number of bytes taken up by $𝐴$ across all chips to the size of one copy of the array?

Solution

Scan pages: 1-2.

Let $𝐴 [𝐼_{𝑋}, 𝐽, 𝐾, \dots]$ and $Mesh = {𝑋 : 4, 𝑌 : 8, 𝑍 : 2}$ .

At each device tuple $(𝑖, 𝑗, 𝑘)$ , size of array shard is:

Size of shard = \frac{| 𝐼 |}{𝑋} \cdot 𝑃 \cdot BytesForPrecision where 𝑃 = \prod_{𝑑 \in dimensions - {𝐼}} | 𝑑 |

Assuming “one copy of array” means unsharded array, ratio is then:

\frac{𝑋 \cdot 𝑌 \cdot 𝑍 \cdot | 𝐼 | \cdot 𝑋^{- 1} \cdot 𝑃 \cdot BytesForPrecision}{| 𝐼 | \cdot 𝑃 \cdot BytesForPrecision} = 16

If copy means one shard, then ratio is:

\frac{𝑋 \cdot 𝑌 \cdot 𝑍 \cdot | 𝐼 | \cdot 𝑋^{- 1} \cdot 𝑃 \cdot BytesForPrecision}{| 𝐼 | \cdot 𝑋^{- 1} \cdot 𝑃 \cdot BytesForPrecision} = 64

Exercise 2 – AllGather latency

Exercise statement

How long should ${AllGather}_{𝑋} ([𝐵_{𝑋}, 𝐷_{𝑌}])$ take on a TPU v4p 4x4x4 slice with mesh Mesh({'X': 4, 'Y': 4, 'Z': 4}) if $𝐵 = 1024$ and $𝐷 = 4096$ in bfloat16? How about ${AllGather}_{𝑋 𝑌} ([𝐵_{𝑋}, 𝐷_{𝑌}])$ ? How about ${AllReduce}_{𝑍} ([𝐵_{𝑋}, 𝐷_{𝑌}] {𝑈_{𝑍}})$ ?

Solution

Scan pages: 2-5.

$v4p :$

HBM capacity = $32 GB$
HBM BW = $1.2 𝑒 12 Bytes/s$
FLOPs/s (bf16) = $2.75 𝑒 14$
ICI Bidi = $9.0 𝑒 10$

Slice of $4 \times 4 \times 4$ $v4p$ . $Mesh = {𝑋 : 4, 𝑌 : 4, 𝑍 : 4}$ . Let $𝐵 = 1024$ and $𝐷 = 4096$ in bfloat16. Assuming $𝐴 [𝐵_{𝑋}, 𝐷_{𝑌}]$ , each device tuple in the mesh maps to:

(𝑖, 𝑗, 𝑘) \mapsto 𝐴 [𝑖 \cdot \frac{| 𝐵 |}{𝑋} : (𝑖 + 1) \cdot \frac{| 𝐵 |}{𝑋}, 𝑗 \cdot \frac{| 𝐷 |}{𝑌} : (𝑗 + 1) \cdot \frac{| 𝐷 |}{𝑌}]

Therefore size per shard is:

\frac{| 𝐵 |}{𝑋} \cdot \frac{| 𝐷 |}{𝑌} \cdot 2 = \frac{𝐵^{2}}{2} bytes

since $𝐷 = 4 \cdot 𝐵$ .

Since v4p slice $4 \times 4 \times 4$ has wrap-around links, we get that:

𝑇 ({AllGather}_{𝑋} ([𝐵_{𝑋}, 𝐷_{𝑌}])) = \frac{𝑋}{2} \cdot \frac{𝐵^{2}}{2 \cdot 𝑊_{uni}} = \frac{𝐵^{2}}{4.5 𝑒 10} \approx 2 \cdot 10^{- 5} s = 0.02 ms.

For allgather on 2 axis we don’t have an specific algorithm but we can derive a lower bound:

𝑇 ({AllGather}_{𝑋 𝑌} ([𝐵_{𝑋}, 𝐷_{𝑌}])) \geq \frac{2 \cdot 𝐵 \cdot 𝐷}{2 \cdot 𝑊_{ICI}} = \frac{2 𝐵 \cdot 4 𝐵}{2 𝑊_{ICI}} = \frac{4 𝐵^{2}}{9 𝑒 10} \approx 0.04 ms.

$𝑇 ({AllReduce}_{𝑍} ([𝐵_{𝑋}, 𝐷_{𝑌}] {𝑈_{𝑍}})) = \frac{𝑍}{2} \cdot (\frac{𝐵^{2}}{2 \cdot 𝑍 \cdot 𝑊_{uni}}) + \frac{𝑍}{2} \cdot (\frac{𝐵^{2}}{2 \cdot 𝑍 \cdot 𝑊_{uni}}) = \frac{𝐵^{2}}{4 \cdot 𝑊_{uni}} + \frac{𝐵^{2}}{4 \cdot 𝑊_{uni}} = \frac{2 𝐵^{2}}{4 𝑊_{uni}} = \frac{𝐵^{2}}{2 𝑊_{uni}} = \frac{𝐵^{2}}{𝑊_{ICI}} = \frac{1024^{2}}{9 𝑒 10} = 1 \cdot 10^{- 4} = 0.01 ms.$

Exercise 3 – latency-bound AllGather

Exercise statement

Let’s say we’re performing an ${AllGather}_{𝑋 ([𝐵_{𝑋}])}$ but $𝐵$ is very small, say 128. How long should this take on a TPU v4p 4x4x4 slice with mesh Mesh({'X': 4, 'Y': 4, 'Z': 4}) in bfloat16? Hint: you’re probably latency bound.

Solution

Scan pages: 6-7.

TPU v4p.

Let $Mesh = {𝑋 : 4, 𝑌 : 4, 𝑍 : 4}$ Let $𝐵 = 128 and 𝐴 = bfloat16 [𝐵_{𝑋}]$

Assuming communication model $𝑇_{step} = 𝛼 + \frac{𝐷}{BW}$ where $𝛼 = 1 𝜇 s$ , and wraparound links, $𝑇 ({AllGather}_{𝑋} ([𝐵_{𝑋}])) = \frac{𝑋}{2} \cdot (𝛼 + \frac{𝐵 \cdot 2}{𝑋 \cdot 𝑊_{uni}}) = 2 \cdot 𝛼 + \frac{2 \cdot 𝐵}{𝑊_{ICI}} = 2 𝜇 s + \frac{256}{9 𝑒 10} \approx 2 𝜇 s + 28.5 \cdot 10^{- 10} = 2 𝜇 s + 28.5 \cdot 10^{- 3} \cdot 10^{- 6} = 2 𝜇 s + 0.0285 𝜇 s$

So yeah I’m bound by latency.

Exercise 4 – matmul strategies

Exercise statement

To perform $𝑋 [𝐵, 𝐷] \cdot_{𝐷} 𝑌 [𝐷_{𝑋}, 𝐹] \to 𝑍 [𝐵, 𝐹]$ , in this section we tell you to perform ${AllGather}_{𝑋} (𝑌 [𝐷_{𝑋}, 𝐹])$ and multiply the fully replicated matrices (Case 2, Strategy 1). Instead, you could multiply the local shards like $𝑋 [𝐵, 𝐷_{𝑋}] \cdot_{𝐷} 𝑌 [𝐷_{𝑋}, 𝐹] \to 𝑍 [𝐵, 𝐹] {𝑈_{𝑋}}$ (Case 3, Strategy 2), and then ${AllReduce}_{𝑋} (𝑍 [𝐵, 𝐹] {𝑈_{𝑋}})$ . How many FLOPs and comms does each of these perform? Which is better and why?

Solution

Scan pages: 7-16.

𝑋 [𝐵, 𝐷] \cdot_{𝐷} 𝑌 [𝐷_{𝑋}, 𝐹] \to 𝑍 [𝐵, 𝐹]

Strategy 1:

${AllGather}_{𝑋} (𝑌 [𝐷_{𝑋}, 𝐹]) \to 𝑌 [𝐷, 𝐹]$
$𝑋 [𝐵, 𝐷] \cdot_{𝐷} 𝑌 [𝐷, 𝐹] \to 𝑍 [𝐵, 𝐹]$

Note that $2 \cdot 𝑋 =$ number of directed links. Each link participates in $\frac{𝑋}{2}$ steps of the bidirectional ring algo and each payload is $\frac{𝐷 \cdot 𝐹 \cdot 2}{𝑋}$ , assuming bfloat16:

AllGather [𝐷_{𝑋}, 𝐹] Comms = 2 \cdot (\frac{𝑋}{2} \cdot \frac{𝐷 \cdot 𝐹 \cdot 2}{𝑋}) = 2 𝐷 𝐹

Comms for MatMul $= 2 𝐵 𝐷 + 2 𝐷 𝐹 + 2 𝐵 𝐹$ $\Rightarrow Total comms strategy 1 = 2 𝐷 𝐹 + 2 𝐵 𝐷 + 2 𝐷 𝐹 + 2 𝐵 𝐹 = 4 𝐷 𝐹 + 2 𝐵 𝐷 + 2 𝐵 𝐹 .$

$FLOPs for strategy 1 = 2 \cdot 𝐵 \cdot 𝐷 \cdot 𝐹$

Strategy 2:

$𝑋 [𝐵, 𝐷_{𝑋}] \cdot_{𝐷} 𝑌 [𝐷_{𝑋}, 𝐹] \to 𝑍 [𝐵, 𝐹] {𝑈_{𝑋}}$
${AllReduce}_{𝑋} (𝑍 [𝐵, 𝐹] {𝑈_{𝑋}})$

$Comms MatMul = 2 \cdot 𝐵 \cdot \frac{𝐷}{𝑋} + 2 \cdot \frac{𝐷}{𝑋} \cdot 𝐹 + 2 \cdot 𝐵 \cdot 𝐹 .$

$Comms {AllReduce}_{𝑋} (𝑍 [𝐵, 𝐹] {𝑈_{𝑋}}) = 2 \cdot (2 (\frac{𝑋}{2} \cdot \frac{𝐵 \cdot 𝐹 \cdot 2}{𝑋})) = 4 𝐵 𝐹 .$

$Total comms strategy 2 = 4 𝐵 𝐹 + 2 \cdot 𝐵 \cdot \frac{𝐷}{𝑋} + 2 \cdot \frac{𝐷}{𝑋} \cdot 𝐹 + 2 𝐵 𝐹 = 6 𝐵 𝐹 + 2 \cdot 𝐵 \cdot \frac{𝐷}{𝑋} + 2 \cdot \frac{𝐷}{𝑋} \cdot 𝐹 .$

$FLOPs strategy 2 = 2 \cdot 𝐵 \cdot \frac{𝐷}{𝑋} \cdot 𝐹 .$

From the book they say you can overlap collective communications with the matrix multiplication itself. So then we can work only with the communications of the collectives to see which strategy is better.

So then assume:

\begin{aligned} 𝑇_{strategy1} & = max (\frac{2 𝐷 𝐹}{𝑊_{ICI}}, \frac{2 𝐵 𝐷 𝐹}{FLOPs/s}) \\ 𝑇_{strategy2} & = max (\frac{4 𝐵 𝐹}{𝑊_{ICI}}, \frac{2 𝐵 𝐷 𝐹}{𝑋 \cdot FLOPs/s}) \\ 𝐴 𝐼_{1} & = \frac{2 𝐵 𝐷 𝐹}{2 𝐷 𝐹} = 𝐵 FLOPs/Byte \\ 𝐴 𝐼_{2} & = \frac{2 𝐵 𝐷 𝐹 \cdot 𝑋^{- 1}}{4 𝐵 𝐹} = \frac{𝐷}{2 𝑋} FLOPs/Byte \end{aligned}

Let’s fix $𝑋 \in {4, 8, 16}$ valid in v4p slice, and they will still have a wraparound link. $𝐴 𝐼_{1} (𝐵), 𝐴 𝐼_{2} (𝐷; 𝑋) for fixed values of 𝑋 .$

Achievable FLOPs:

\begin{aligned} Achievable FLOPs (𝐵) & = min (𝐵 \cdot 9 𝑒 10, 2.75 𝑒 14) \\ Achievable FLOPs (𝐷; 𝑋) & = min (\frac{𝐷}{2 𝑋} \cdot 9 𝑒 10, 2.75 𝑒 14) \end{aligned}

Algo 1 branches by $𝐵 > 3055$ .Algo 2 branches by $𝐷 > 24 𝐾$ . So in total 4 combinations:

$𝐵 > 3055$ and $𝐷 > 24 𝐾$ $\Rightarrow$ Both algos are compute bound.

\begin{aligned} 𝑇_{math} (Strat 1) < 𝑇_{math} (Strat 2) \\ \Leftrightarrow 2 𝐵 𝐷 𝐹 < \frac{2 𝐵 𝐷 𝐹}{𝑋} for 𝑋 = 4 \\ \Leftrightarrow 2 < \frac{1}{2} Contradiction. \end{aligned}

Therefore $𝑇_{math} (Strat 2) \geq 𝑇_{math} (Strat 1)$ and in fact it holds for every $𝑋 > 1$ . Just the condition for when Algo 2 is compute bound changes, but assuming both compute bound strategy 2 is better.

$𝐵 > 3055$ and $𝐷 < 24 𝐾$ $\Rightarrow$ Algo 1 compute bound and Algo 2 communication bound.

\begin{aligned} 𝑇_{math} (Strat 1) < 𝑇_{comms} (Strat 2) \\ \Leftrightarrow \frac{2 𝐵 𝐷 𝐹}{FLOPs/s} < \frac{4 𝐵 𝐹}{𝑊_{ICI}} \\ \Leftrightarrow \frac{2 𝐷}{4} < \frac{FLOPs}{𝑊_{ICI}} = 3055 \\ \Leftrightarrow \frac{𝐷}{2} < 3055 \Rightarrow 𝐷 < 6110 \end{aligned}

So strategy 1 in this case wins for $𝐷 < 6110$ otherwise strategy 2 is better.

$𝐵 < 3055$ and $𝐷 > 24 𝐾$ $\Rightarrow$ Algo 1 communication bound and Algo 2 compute bound.

\begin{aligned} 𝑇_{comms} (Strat 1) < 𝑇_{math} (Strat 2) \\ \Leftrightarrow \frac{2 𝐷 𝐹}{𝑊_{ICI}} < \frac{2 𝐵 𝐷 𝐹}{𝑋 \cdot FLOPs/s}) for 𝑋 = 4 \\ \Leftrightarrow \frac{2}{𝑊_{ICI}} < \frac{2 𝐵}{4 \cdot FLOPs/s} \\ \Leftrightarrow 𝐵 > 4 \cdot \frac{FLOPs/s}{𝑊_{ICI}} = 4 \cdot 3055 \approx 12 𝐾 \end{aligned}

which is a contradiction, so then strategy 2 in this setting is always better, and in fact it will hold for all $𝑋 > 2$ . It just changes the condition when algo 2 is compute bound.

$𝐵 \leq 3055$ and $𝐷 \leq 24 𝐾$ $\Rightarrow$ Both algos are communication bound.

\begin{aligned} 𝑇_{comms} (Strat 1) < 𝑇_{comms} (Strat 2) \\ \Leftrightarrow 2 𝐷 𝐹 < 4 𝐵 𝐹 \\ \Leftrightarrow 2 𝐷 < 4 𝐵 \Rightarrow 𝐷 < 2 \cdot 𝐵 \end{aligned}

So if $𝐷 < 2 \cdot 𝐵$ , strategy 1 wins, otherwise strategy 2 wins. In the book they often assume $𝐷 ≫ 𝐵$ , so then $𝐷 > 2 \cdot 𝐵$ and strategy 2 is better.

Exercise 5 – minimum latency

Exercise statement

Let’s say I want to do a matmul $𝐴 [𝐼, 𝐽] \cdot_{𝐽} 𝐵 [𝐽, 𝐾] \to 𝐶 [𝐼, 𝐾]$ on a TPU v4p 4x4x4 with the lowest possible latency. Assume the inputs can be sharded arbitrarily but the result should be fully replicated. How should my inputs be sharded? What is the total FLOPs and comms time?

Solution

Scan pages: 17-25.

Assume we have a $4 \times 4 \times 4$ slice of v4p devices with wraparound links. Assume the intended operation is:

𝐴 [𝐼, 𝐽] \cdot_{𝐽} 𝐵 [𝐽, 𝐾] \to 𝐶 [𝐼, 𝐾]

with bfloat16 precision.

Note that there is only one operation where we do FLOPs, the matrix multiplication (ignoring AllReduce). This allows to divide the space of possible sharding strategies in three groups:

Shardings resulting in collective communications before doing matmul.
Shardings resulting in collective communications only after doing matmul.
Shardings resulting in collective communications before and after matmul.

In all groups, note that if we are going to do a collective operation, we should choose shardings for which the lower bound in $𝑇_{comms}$ is the smallest, since an optimal algo can get closer to this lower bound. Note that for either $𝑂 \in {𝐴, 𝐵, 𝐶}$ we have:

𝑇_{comms} \geq \frac{size (𝑂)}{𝑁_{axes} \cdot 𝑊_{ICI}} .

So the smallest lower bound is achieved using shardings that use all the axis of our $4 \times 4 \times 4$ v4p slice. This argument wipes out our space of possible shardings by a lot.

For group 1:

1.1 $𝐴 [𝐼, 𝐽_{𝑋 𝑌 𝑍}] \cdot 𝐵 [𝐽, 𝐾] \to 𝐶 [𝐼, 𝐾]$

1.2 $𝐴 [𝐼, 𝐽] \cdot 𝐵 [𝐽_{𝑋 𝑌 𝑍}, 𝐾] \to 𝐶 [𝐼, 𝐾]$

For group 2:

2.1 $𝐴 [𝐼_{𝑋 𝑌 𝑍}, 𝐽] \cdot 𝐵 [𝐽, 𝐾] \to 𝐶 [𝐼, 𝐾]$

2.2 $𝐴 [𝐼, 𝐽] \cdot 𝐵 [𝐽, 𝐾_{𝑋 𝑌 𝑍}] \to 𝐶 [𝐼, 𝐾]$

2.3 $𝐴 [𝐼, 𝐽_{𝑋 𝑌 𝑍}] \cdot 𝐵 [𝐽_{𝑋 𝑌 𝑍}, 𝐾] \to 𝐶 [𝐼, 𝐾]$

For group 3:

3.1 $𝐴 [𝐼_{𝑋 𝑌 𝑍}, 𝐽] \cdot 𝐵 [𝐽_{𝑋 𝑌 𝑍}, 𝐾] \to 𝐶 [𝐼, 𝐾]$

3.2 $𝐴 [𝐼, 𝐽_{𝑋 𝑌 𝑍}] \cdot 𝐵 [𝐽, 𝐾_{𝑋 𝑌 𝑍}] \to 𝐶 [𝐼, 𝐾]$

3.3 $𝐴 [𝐼_{𝑋 𝑌 𝑍}, 𝐽] \cdot 𝐽 [𝐽, 𝐾_{𝑋 𝑌 𝑍}] \to 𝐶 [𝐼, 𝐾]$

Note that in group 3, depending in which axis you allgather first, you end back to a sharding that belongs to group 1 or 2.

Assume $𝐽 ≫ 𝐼$ and $𝐾 \approx 4 \cdot 𝐽$ .

Then for group 1:

1.1 $𝑇_{comms} = 2 \cdot 𝐼 \cdot 𝐽, 𝑇_{math} = 2 \cdot 𝐼 \cdot 𝐽 \cdot 𝐾$ . If $\frac{2 𝐼 𝐽 𝐾}{2 𝐼 𝐽} = 𝐾 > 3055 \Rightarrow compute bound$

1.2 $𝑇_{comms} = 2 \cdot 𝐽 \cdot 𝐾, 𝑇_{math} = 2 \cdot 𝐼 \cdot 𝐽 \cdot 𝐾$ . If $𝐼 > 3055 \Rightarrow compute bound$ .

Then for group 2:

2.1 $𝑇_{comms} = 2 \cdot 𝐼 \cdot 𝐾, 𝑇_{math} = \frac{2 𝐼 𝐽 𝐾}{𝑋 \cdot 𝑌 \cdot 𝑍}$ . If $\frac{𝐽}{𝑋 \cdot 𝑌 \cdot 𝑍} > 3055 \Rightarrow compute bound$ i.e $𝐽 > 64 \cdot 3055 = 196 𝐾$ not realistic as $𝐵 [𝐽, 𝐾]$ will not fit in HBM.

2.2 $𝑇_{comms} = 2 \cdot 𝐼 \cdot 𝐾, 𝑇_{math} = \frac{2 𝐼 𝐽 𝐾}{𝑋 \cdot 𝑌 \cdot 𝑍}$ . If $\frac{𝐽}{𝑋 \cdot 𝑌 \cdot 𝑍} > 3055 \Rightarrow compute bound$ i.e. $𝐽 > 196 𝐾$ . So again not realistic.

2.3 $𝑇_{comms} = 4 \cdot 𝐼 \cdot 𝐾, 𝑇_{math} = \frac{2 𝐼 𝐽 𝐾}{𝑋 \cdot 𝑌 \cdot 𝑍}$ . If $\frac{𝐽}{2 𝑋 \cdot 𝑌 \cdot 𝑍} > 3055 \Rightarrow compute bound$ i.e $𝐽 > 391 𝐾$ . Not realistic.

Hence:

\begin{aligned} 𝑇 (Algo 1.1) & = \frac{2 𝐼 𝐽 𝐾}{FLOPs/s} \approx \frac{8 𝐽^{2}}{FLOPs/s} \\ 𝑇 (Algo 1.2) & = \frac{2 𝐼 𝐽 𝐾}{FLOPs/s} \approx \frac{8 𝐽^{2}}{FLOPs/s} \\ or \\ = \frac{2 𝐽 𝐾}{𝑊_{ICI}} = \frac{8 𝐽^{2}}{𝑊_{ICI}} \\ 𝑇 (Algo 2.2) & = \frac{2 𝐼 𝐾}{𝑊_{ICI}} \approx \frac{8 𝐽}{𝑊_{ICI}} \end{aligned}

So then Algo 2.2 and 2.1 are better than Algo 1.2 and Algo 1.1 when all are communication bound. And if Algo 1.2 and Algo 1.1 are compute bound, Algo 2.1 and Algo 2.2 are better when:

\frac{8 𝐽}{𝑊_{ICI}} < \frac{8 𝐽^{2}}{FLOPs/s} \Rightarrow 𝐽 > 3055

In that case, between Algo 2.1 and Algo 2.2, Algo 2.2 is better since $𝐾 ≫ 𝐼$ and:

comms Algo 2.1 = \frac{2 𝐼 𝐽}{𝑋 𝑌 𝑍} + 2 𝐽 𝐾 + 2 𝐼 𝐾

is dominated by $2 𝐽 𝐾$ .

comms Algo 2.2 = 2 𝐼 𝐾 + \frac{2 𝐽 𝐾}{𝑋 𝑌 𝑍} + 2 𝐼 𝐾

is dominated by $2 𝐽 \frac{𝐾}{𝑋 𝑌 𝑍}$ . So if both memory bound, Algo 2.2 is better. And we enter compute bound with smaller values of $𝐼$ with Algo 2.2, since if $𝐼 = 𝐴 𝐼 (Algo 2.2) > 230$ we are compute bound and for Algo 2.1 we have:

𝐴 𝐼 = \frac{2 𝐼 𝐽 𝐾}{𝑋 𝑌 𝑍} \cdot \frac{1}{2 𝐽 𝐾} = \frac{2}{64} \cdot 𝐼 > 230 \Rightarrow 𝐼 > 7360

Exercise 6 – sharded matmul cases on TPU v5e 4x4

Exercise statement

Let’s say we want to perform $𝐴 [𝐼_{𝑋}, 𝐽_{𝑌}] \cdot_{𝐽} 𝐵 [𝐽_{𝑌}, 𝐾] \to 𝐶 [𝐼_{𝑋}, 𝐾]$ on TPU v5e 4x4. What communication do we perform? How much time is spent on communication vs. computation?

What about $𝐴 [𝐼_{𝑋}, 𝐽] \cdot_{𝐽} 𝐵 [𝐽_{𝑋}, 𝐾_{𝑌}] \to 𝐶 [𝐼_{𝑋}, 𝐾_{𝑌}]$ ? This is the most standard setting for training where we combine data, tensor, and ZeRO sharding.

What about $𝐴 [𝐼_{𝑋}, 𝐽] \cdot_{𝐽} 𝐵 [𝐽, 𝐾_{𝑌}] \to 𝐶 [𝐼_{𝑋}, 𝐾_{𝑌}]$ ? This is standard for inference, where we do pure tensor parallelism plus data.

Solution

Scan pages: 25-28.

Assume we have a $4 \times 4$ slice of TPU v5e. Since $4 \times 4$ is not a full pod in this setting we don’t have wraparound links.

Let:

v5e ICI BW unidirectional = $4.5 𝑒 10$
v5e FLOPs/s (bf16) = $1.97 𝑒 14$

Assuming BFLOAT16.

For $𝐴 [𝐼_{𝑋}, 𝐽_{𝑌}] \cdot_{𝐽} 𝐵 [𝐽_{𝑌}, 𝐾] \to 𝐶 [𝐼_{𝑋}, 𝐾]$ :

$𝑂 [𝐼_{𝑋}, 𝐾] {𝑈_{𝑌}} = 𝐴 [𝐼_{𝑋}, 𝐽_{𝑌}] \cdot 𝐽_{local} [𝐽_{𝑌}, 𝐾]$
$𝐶 [𝐼_{𝑋}, 𝐾] = {AllReduce}_{𝑌} (𝑂 [𝐼_{𝑋}, 𝐾] {𝑈_{𝑌}})$

We cannot use a bidirectional ring algorithm for the ReduceScatter + AllGather, since no wraparound links in a $4 \times 4$ slice. Hence:

\begin{matrix} Communication Load = 2 \cdot (| 𝑌 | - 1) \cdot \frac{𝐼 \cdot 𝐾 \cdot 2}{| 𝑌 | | 𝑋 |} \approx \frac{2 \cdot 𝐼 \cdot 𝐾 \cdot 2}{| 𝑋 |} \Rightarrow 𝑇_{comms} = \frac{4 𝐼 𝐾}{| 𝑋 | \cdot BW} \\ 𝑇_{math} = \frac{2 𝐼 𝐽 𝐾}{| 𝑋 | | 𝑌 | \cdot FLOPs/s} . \end{matrix}

Therefore:

\frac{𝑇_{comms}}{𝑇_{math}} = \frac{4 𝐼 𝐾}{| 𝑋 | \cdot BW} \cdot \frac{| 𝑋 | | 𝑌 | \cdot FLOPs/s}{2 𝐼 𝐽 𝐾} = \frac{4 | 𝑌 |}{𝐽} \cdot \frac{FLOPs/s}{BW} = \frac{16}{𝐽} \cdot 4380 \approx \frac{70080}{𝐽} .

For $𝐴 [𝐼_{𝑋}, 𝐽] \cdot_{𝐽} 𝐵 [𝐽_{𝑋}, 𝐾_{𝑌}] \to 𝐶 [𝐼_{𝑋}, 𝐾_{𝑌}]$ :

$𝐵 [𝐽, 𝐾_{𝑌}] = {AllGather}_{𝑋} (𝐵 [𝐽_{𝑋}, 𝐾_{𝑌}])$
$𝐶 [𝐼_{𝑋}, 𝐾_{𝑌}] = 𝐴 [𝐼_{𝑋}, 𝐽] \cdot 𝐵 [𝐽, 𝐾_{𝑌}]$

\begin{matrix} Communication load = (| 𝑋 | - 1) \cdot (\frac{2 𝐽 \cdot 𝐾}{| 𝑋 | | 𝑌 |}) \approx 2 \cdot 𝐽 \cdot 𝐾 \cdot {| 𝑌 |}^{- 1} \Rightarrow 𝑇_{comms} = \frac{2 𝐽 𝐾}{| 𝑌 | \cdot BW} \\ 𝑇_{math} = \frac{2 𝐼 𝐽 𝐾}{| 𝑋 | | 𝑌 | \cdot FLOPs/s} . \end{matrix}

Therefore:

\frac{𝑇_{comms}}{𝑇_{math}} = \frac{2 𝐽 𝐾}{| 𝑌 |} \cdot \frac{| 𝑋 | | 𝑌 | \cdot FLOPs/s}{2 𝐼 𝐽 𝐾 \cdot BW} = \frac{| 𝑋 |}{𝐼} \cdot 4380 = \frac{17520}{𝐼} .

For $𝐴 [𝐼_{𝑋}, 𝐽] \cdot_{𝐽} 𝐵 [𝐽, 𝐾_{𝑌}] \to 𝐶 [𝐼_{𝑋}, 𝐾_{𝑌}]$ :

No communication needs to be done.

𝑇_{math} = \frac{2 𝐼 𝐽 𝐾}{| 𝑋 | | 𝑌 | \cdot FLOPs/s}

All the time is in computation.

Exercise 7 – Transformer block sharding with memory limit

Exercise statement

A typical Transformer block has two matrices $𝑊_{in} [𝐷, 𝐹]$ and $𝑊_{out} [𝐹, 𝐷]$ where $𝐹 ≫ 𝐷$ . Say we have a batch size B. Then the full block is $In [𝐵, 𝐷] \cdot 𝑊_{in} [𝐷, 𝐹] \cdot 𝑊_{out} [𝐹, 𝐷]$ . Let’s pick $𝐷 = 8192$ , $𝐹 = 32768$ , and $𝐵 = 128$ and assume everything is in bfloat16. Assume we’re running on a TPU v5e 2x2 slice but let’s pretend each TPU only has 300MB of free memory. How should In, $𝑊_{in}$ , $𝑊_{out}$ , and Out be sharded to stay below the memory limit while minimizing overall time? How much time is spent on comms and FLOPs? Hint: the final output doesn’t need to be fully replicated, but it should be sharded the same as the input so the layer can be repeated.

Solution

Scan pages: 29-32.

Assume bfloat16. Let $𝐷 = 8192$ , $𝐹 = 32768$ , $𝐵 = 128$ . The intended operation is:

In [𝐵, 𝐷] \cdot 𝑊_{in} [𝐷, 𝐹] \cdot 𝑊_{out} [𝐹, 𝐷]

$Size (𝑊_{in}) = 8192^{2} \cdot 4 \cdot 2 \approx 536 MB$ .

$Size (𝑊_{out}) = Size (𝑊_{in})$ .

$Size (In) \approx 2 MB$ .

$Size (Aux [𝐵, 𝐹]) \approx 8 MB$ .

Since we only have left 300MB, then we have to shard $𝑊_{in}$ and $𝑊_{out}$ so that both sharded $𝑊_{in}$ and $𝑊_{out}$ can be found in each device. Sharding just one dimension along one axis does not work since then both $𝑊_{in}$ and $𝑊_{out}$ do not fit in a single device:

\frac{Size (𝑊_{in})}{2} + \frac{Size (𝑊_{out})}{2} > 300 MB

Therefore we have to either shard each dimension with a single axis i.e:

𝑊_{in} [𝐷_{𝑋}, 𝐹_{𝑌}] \forall \in {𝑋, 𝑌}

where $𝑋, 𝑌$ are the axis of the device mesh, or shard one dimension along both axis of the mesh i.e:

𝑊_{in} [𝐷, 𝐹_{𝑋 𝑌}] or 𝑊_{in} [𝐷_{𝑋 𝑌}, 𝐹]

In both cases we shrink both $Size (𝑊_{in})$ and $Size (𝑊_{out})$ by 4 and:

\frac{Size (𝑊_{in})}{4} + \frac{Size (𝑊_{out})}{4} \leq 300 MB

The way to shard to minimize overall time is:

In [𝐵, 𝐷] @ 𝑊_{in} [𝐷, 𝐹_{𝑋 𝑌}] @ 𝑊_{out} [𝐹_{𝑋 𝑌}, 𝐷]

where we construct a 1-D ring in $2 \times 2$ slice of v5e TPUs.

For our chosen sharding we only do one collective communication at the end, specifically an AllReduce = ReduceScatter + AllGather. Every other sharding strategy will require:

𝑇_{comms} (Alternative sharding) \geq \frac{4 𝐵 𝐷}{𝑊_{ICI}}

You could also do:

In [𝐵, 𝐷_{𝑋 𝑌}] \cdot 𝑊_{in} [𝐷, 𝐹_{𝑋 𝑌}] \cdot 𝑊_{out} [𝐹_{𝑋 𝑌}, 𝐷]

with an AllGather at the beginning and a ReduceScatter at the end and get same:

𝑇_{comms} = \frac{2 𝐵 𝐷}{𝑊_{ICI}} + \frac{2 𝐵 𝐷}{𝑊_{ICI}} = \frac{4 𝐵 𝐷}{𝑊_{ICI}}

So for our chosen sharding strategy:

$Aux [𝐵, 𝐹_{𝑋 𝑌}] = In [𝐵, 𝐷] \cdot 𝑊_{in} [𝐷, 𝐹_{𝑋 𝑌}]$
$𝑂 [𝐵, 𝐷] {𝑈_{𝑋 𝑌}} = Aux [𝐵, 𝐹_{𝑋 𝑌}] \cdot 𝑊_{out} [𝐹_{𝑋 𝑌}, 𝐷]$
$𝑂 [𝐵, 𝐷] = {AllReduce}_{𝑋 𝑌} (𝑂 [𝐵, 𝐷] {𝑈_{𝑋 𝑌}})$

Therefore:

\begin{aligned} 𝑇_{math} & = \frac{4 𝐵 𝐷 𝐹}{| 𝑋 | | 𝑌 | \cdot 1.97 𝑒 14} = \frac{128 \cdot 8192 \cdot 32768}{1.97 𝑒 14} \approx 0.0002 s = 0.2 ms \\ 𝑇_{comms} & = \frac{4 𝐵 𝐷}{9.0 𝑒 10} \approx 0.00004 s = 4 \cdot 10^{- 5} = 4 \cdot 10^{- 2} \cdot 10^{- 3} s = 0.04 ms \end{aligned}

Exercise 9 – another strategy for sharded matmuls?

Exercise statement

Above, we claimed that when only one input to a matmul is sharded along its contracting dimension, we should AllGather the sharded matrix and perform the resulting contraction locally. Another strategy you might think of is to perform the sharded matmul and then AllReduce the result, as if both inputs were sharded along the contracting dimension, i.e. $𝐴 [𝐼, 𝐽_{𝑋}] *_{𝐽} 𝐵 [𝐽, 𝐾] \to 𝐶 [𝐼, 𝐾]$ by way of:

$𝐶 [𝐼, 𝐾] {𝑈_{𝑋}} = 𝐴 [𝐼, 𝐽_{𝑋}] \cdot 𝐵 [𝐽_{𝑋}, 𝐾]$
$𝐶 [𝐼, 𝐾] = AllReduce (𝐶 [𝐼, 𝐾] {𝑈_{𝑋}})$

Answer the following:

Explicitly write out this algorithm for matrices $𝐴 [𝑁, 𝑀]$ and $𝐵 [𝑀, 𝐾]$ , using indices to show exactly what computation is done on what device. Assume $𝐴$ is sharded as $𝐴 [𝐼, 𝐽_{𝑋}]$ across ND devices, and you want your output to be replicated across all devices.
Now suppose you are ok with the final result not being replicated on each device, but instead sharded, across either the N or K dimension. How would the algorithm above change?
Looking purely at the communication cost of the strategy above, in part 2 not 1, how does this communication cost compare to the communication cost of the algorithm in which we first AllGather A and then do the matmul?

Solution

Scan pages: 33-34.

9.1

Let $𝐴 [𝑁, 𝑀]$ be sharded as $𝐴 [𝐼, 𝐽_{𝑋}]$ where $| 𝑋 | = 𝐷$ . Let $𝐵 [𝑀, 𝐾]$ be sharded as $𝐵 [𝐽_{𝑋}, 𝐾]$ . Assume $\frac{𝑀}{𝐷} \in ℕ$ .

Define:

𝐴_{𝑖} ≔ 𝐴 [:, 𝑖 \cdot \frac{𝑀}{𝐷} : (𝑖 + 1) \cdot \frac{𝑀}{𝐷}]

for $0 \leq 𝑖 \leq 𝐷 - 1$ , where we use Python slicing notation and semantics.

Define:

𝐵_{𝑖} ≔ 𝐵 [𝑖 \cdot \frac{𝑀}{𝐷} : (𝑖 + 1) \cdot \frac{𝑀}{𝐷}, :]

for $0 \leq 𝑖 \leq 𝐷 - 1$ .

The computation at each device $𝑖$ is then $𝐴_{𝑖} \cdot 𝐵_{𝑖} \in ℝ^{𝑁 \times 𝐾}$ . After that we do $𝐶 [𝐼, 𝐾] = \sum_{𝑖 = 0}^{𝐷 - 1} 𝐴_{𝑖} \cdot 𝐵_{𝑖}$ with an AllReduce.

9.2

$𝐶 [𝐼, 𝐾] {𝑈_{𝑋}} = 𝐴 [𝐼, 𝐽_{𝑋}] \cdot 𝐵 [𝐽_{𝑋}, 𝐾]$
${ReduceScatter}_{𝑋, 𝐼 or 𝐾} (𝐶 [𝐼, 𝐾] {𝑈_{𝑋}})$

9.3

Assume bfloat16.

$Communication Load Part 2 = 2 \cdot 𝑁 \cdot 𝐾$ .

$Communication load AllGather first = 2 \cdot 𝑀 \cdot 𝐾$ .

$Ratio = \frac{2 \cdot 𝑁 \cdot 𝐾}{2 \cdot 𝑀 \cdot 𝐾} = \frac{𝑁}{𝑀}$ .

Exercise 10 – Fun with AllToAll

Exercise statement

In the table above, it was noted that the time to perform an AllToAll is a factor of 4 lower than the time to perform an AllGather or ReduceScatter, in the regime where we are throughput-bound. In this problem we will see where that factor of 4 comes from, and also see how this factor would change if we only had single-direction ICI links, rather than bidirectional ICI links.

Let’s start with the single-direction case first. Imagine we have D devices in a ring topology and want to do either an AllGather or a ReduceScatter on an N x N matrix $𝐴 [𝐼_{𝑋}, 𝐽]$ , say $𝐷$ divides $𝑁$ for simplicity. Describe the comms involved in these two collectives, and calculate the total number of scalars, floats or ints, which are transferred across a single ICI link during the entirety of this algorithm.
Now let’s think about an AllToAll, still in the single-directional ICI case. How is the algorithm different in this case than the all-gather case? Calculate the number of scalars that are transferred across a single ICI link in this algorithm.
You should have found that the ratio between your answers to part (a) and part (b) is a nice number. Explain where this factor comes from in simple terms.
Now let’s add bidirectional communication. How does this affect the total time needed in the all-gather case?
How does adding bidirectional communication affect the total time needed in the AllToAll case?
Now simply explain the ratio between AllGather time and AllToAll time in a bidirectional ring.

Solution

Scan pages: 35-40.

10.1

Let $𝐴 [𝑁, 𝑁]$ be sharded as $𝐴 [𝐼_{𝑋}, 𝐽]$ where $| 𝑋 | = 𝐷$ and $\frac{𝑁}{𝐷} \in ℕ$ .

With only one direction, the ring algo requires $𝐷 - 1$ steps to finish. Define a link $𝑖$ as a directed edge $(𝑖, 𝑖 + 1 mod 𝐷)$ for $0 \leq 𝑖 \leq 𝐷 - 1$ . Note that at one step of the algo link $𝑖$ transfers $\frac{𝑁^{2}}{𝐷}$ scalars. Since there are $𝐷 - 1$ steps, the scalars being transferred by a single link during the entirety of the algorithm is:

(𝐷 - 1) \cdot \frac{𝑁^{2}}{𝐷}

for both AllGather and ReduceScatter.

10.2

At step $𝑘$ of the AllToAll algorithm, where $1 \leq 𝑘 \leq 𝐷 - 1$ , the directed link $(𝑖, 𝑖 + 1 mod 𝐷)$ transfers:

\frac{𝑁^{2}}{𝐷} - 𝑘 \cdot \frac{𝑁^{2}}{𝐷^{2}}

Therefore the scalars over the entirety of the algorithm in a single link is:

\sum_{𝑘 = 1}^{𝐷 - 1} (\frac{𝑁^{2}}{𝐷} - 𝑘 \cdot \frac{𝑁^{2}}{𝐷^{2}}) = (𝐷 - 1) \frac{𝑁^{2}}{𝐷} - \frac{(𝐷 - 1) 𝐷 \cdot 𝑁^{2}}{2 \cdot 𝐷^{2}} = \frac{𝐷 - 1}{𝐷} \cdot (𝑁^{2} - \frac{𝑁^{2}}{2})

10.3

\begin{aligned} Scalars (AllGather) & = (𝐷 - 1) \cdot \frac{𝑁^{2}}{𝐷} \\ Scalars (AllToAll) & = \frac{𝐷 - 1}{𝐷} \cdot (𝑁^{2} - \frac{𝑁^{2}}{2}) \\ \frac{Scalars (AllToAll)}{Scalars (AllGather)} & = \frac{1}{2} \end{aligned}

It comes from the fact that the payload per link is smaller at each step.

10.4

With bidirectional communication the algo for AllGather takes $\frac{𝐷}{2}$ to finish, so then:

\frac{𝐷}{2} \cdot (\frac{𝑁^{2}}{𝐷}) = \frac{𝑁^{2}}{2} scalars

being transferred by a single link $(𝑖, 𝑖 + 1 mod 𝐷)$ or a link $(𝑖 + 1 mod 𝐷, 𝑖)$ .

\frac{Scalars (AllGather Bidi)}{Scalars (AllGather)} = \frac{\frac{𝑁^{2}}{2}}{(𝐷 - 1) \cdot \frac{𝑁^{2}}{𝐷}}

= \frac{1}{2} \cdot \frac{𝐷}{𝐷 - 1} \approx \frac{1}{2} as 𝐷 \to \infty

It takes half the time.

10.5

With bidirectional communication the algorithm for AllToAll takes $\frac{𝐷}{2}$ to finish, so an upper bound is:

\sum_{𝑘 = 1}^{\frac{𝐷}{2}} (\frac{𝑁^{2}}{𝐷 \cdot 2} - 𝑘 \frac{𝑁^{2}}{𝐷^{2}}) \approx \frac{𝐷}{2} \cdot \frac{𝑁^{2}}{𝐷 \cdot 2} - \frac{𝑁^{2}}{𝐷^{2}} \cdot \frac{𝐷^{2}}{8} = \frac{𝑁^{2}}{4} - \frac{𝑁^{2}}{8} = \frac{1}{4} \cdot (𝑁^{2} - \frac{𝑁^{2}}{2})

scalars being transferred by a single link $(𝑖, 𝑖 + 1 mod 𝐷)$ or $(𝑖 + 1 mod 𝐷, 𝑖)$ . So it takes $\frac{1}{4}$ of the time of the single direction AllToAll.

10.6

\frac{Scalars (Bidi AllToAll)}{Scalars (Bidi AllGather)} = \frac{𝑁^{2}}{8} \cdot \frac{2}{𝑁^{2}} = \frac{2}{8} = \frac{1}{4}