How to Parallelize a Transformer for Training | How To Scale Your Model

[jacobaustin123]

Training parallelism, discussed!

[roee89871324]

Nit: "ZeRO-{1,2,3} are used to refer to sharding the weights, gradients and optimizer states in this way, respectively"
This is actually in reverse order in the paper. Optimiser state is 1, gradients is 2 and weights is 3 (since weights have some overhead to normal dp they kept it "last")

[jacobaustin123]

Thanks for catching this, I reversed the ordering. In writing this, we basically found that as long as you're careful about AllGathering the weights ahead of time, ZeRO-3 should always be OK from a roofline standpoint even thought it adds extra comms. I think that wasn't so clear to me from the ZeRO paper.

[chillmang]

How do Mixture of Experts (MoE) and GShard fit into this context? Thank you.

[voyagerdd]

For pure data parallelism, does num_params=HBM per device/10 need to take into account gradient storage space besides parameters and optimizer states?

[jacobaustin123]

This would be with like batch size 1 or something, not enough to be useful. This is just the absolute bare minimum. I'll add a note.

[cnberry]

nit: Seems like you are splitting the difference when approximately max model size in TPU v5p. Takeaway value and text should be consistent.

"TPUv5p pod with 96GB of HBM and pure data parallelism this is about 10B parameters."
"For TPU v5p this is roughly 9B parameters."

[jacobaustin123]

Thanks, I've updated this to be 9B everywhere.

[emergenz]

nit:
Increasing parallelism or reducing batch size both tend to make us more communication-bound because they reduce the amount of compute performed per chip.

[jacobaustin123]

Fixed, thanks!

[bonpyt]

While pure FSDP parallelism dominates for very large batch sizes, in the regime where batch size over number of chips is between roughly 400 and 850, a mixed FSDP + MP strategy is required in order to be comms-bound.

Shouldn't it be to be compute-bound?

[jacobaustin123]

Fixed!

[bonpyt]

Question 1 / Attention parameters: missing number of layers L.

[jacobaustin123]

This should be fixed now, thank you!

[manavgarg]

I think the diagrams for "Tensor Parallelism" and "Mixed FSDP and Tensor Parallelism" sections needs to be switched. The diagram for "Tensor Parallelism" shows local shape as [B // N, D // M] whereas the one for "Mixed FSDP and Tensor Parallelism" show as "[B, D // M]" and it should be inverse?

[jacobaustin123]

Good catch, I think these were swapped. I think these should now be fixed.

[tengyifei]

Just a question about DCN vs FSDP bandwidth requirements. The book mentions "Because DCN has lower bandwidth, it’s typically too slow to do much useful FSDP". By doing that math, it would appear that Data Parallel and FSDP has the same amount of communication (one all-reduces the gradients, another all-gathers the weights and reduce-scatters the gradients), could you explain why is FSDP much more bandwidth intensive than DDP (and hence unsuitable for DCN)?

[awgu]

ZeRO-{1,2,3} are used to refer to sharding the optimizer states, gradients, and weights in this way, respectively. Since all have the same communication cost, we can basically always do ZeRO-3 sharding, which shards the parameters, gradients, and optimizer states across a set of devices.

I do not think that it is true that all three have the same communication cost. Across all variations of implementations I have personally seen, ZeRO-3 includes the extra parameter all-gather in backward, while ZeRO-2 does not. It may be that in your case, you often find that the backward parameter all-gather can be fully overlapped, but I would still not say that the communication cost is the same (since there are certainly many real cases where it cannot be overlapped).

[jacobaustin123]

First of all, I haven't seen many non-TPU implementations so there might be a difference there.

WRT your point about the extra parameter all-gather, the key here is that DP -> FSDP turns a backward pass AllReduce into an AllGather and a ReduceScatter, which has the same overall cost. You're right that it also adds some communication during the forward pass that's absent in the pure DP version, but the backward pass comms volume ought to be the same.

[jacobaustin123]

I added a new footnote to clarify this, let me know if it seems good to you.

[kerrywang]

For data parallelism the compute $T_{math}$ is computed as:
$$T_\text{math} = \frac{2 \cdot 2 \cdot 2 \cdot B \cdot L \cdot D \cdot F} {X \cdot C}$$

I am trying to understand what the three two each means:

1. matmul (multiply and addition)
2. there are 2 matrix multiplication

the final2 is a bit confusing to me, if you are counting both forward and backward pass, shouldn't it be 3?
Thanks!

[jacobaustin123]

I'll add a note about this. We're just looking at the backward pass here since the forward pass as no comms. The last 2 just comes from the 2x matmuls in the backward pass. We can't include the forward pass because we can't overlap the gradient all reduce with the forward pass, since we don't have gradients yet.

[iankur]

@jacobaustin123 T_math expression here includes the number of layers L also in the numerator whereas it should not depend on L as used in other places?

[jacobaustin123]

Fixed!

[boyrealmadred]

In Figure: FSDP shard, shouldn't it be Wout instead Win since the logical shape is [D,F] ?

[jacobaustin123]

I don't think so, in general W_in has shape [D, F] (where F >> D) and W_out has shape [F, D].

[jesse7chen]

I'm curious at what point we are able to overlap comms and compute for tensor parallelism? It doesn't seem possible at first glance since the scattered activations from the previous layer need to be immediately re-gathered for the next layer's computation, with no compute in between. I feel like it is implied that we can overlap it somehow since we still derive the conditions when tensor parallelism becomes comms bound, but based on the sequence of events, it seems like it'll be comms-bound no matter what.

[jacobaustin123]

They can generally be overlapped. We need to do [B, D_Y] @ [D, F_Y], and we can do this by starting to do the local matmul over D, then shuffle the shards around as we keep doing the matmul. See the JAX section for an example.

[tengyifei]

Is this overlapping technique same as the "Overlap Communication with Dependent Computation via Decomposition in Large Deep Learning Models" paper aka collective matmul?

[jacobaustin123]

Yes

[gitnicos]

In Pure Parallelism section Communication time it is worth mentioning where the factor 2 comes = it counts for redundancy of bidirectional communication in a ring to cut down number of hops that was talked in Section 3 (not previous section as mentioned in the text).

Generally, I think this kind of brief repetitions do help to stay focused on main subject of discussion rather than scrambling to recall exactly where some factors do come from.

[gitnicos]

In the FSDP diagram above "Note that the activations (left) are not sharded along the contracting dimension, which is what forces us to gather". It is important reminding reasoning behind what necessitates gathering - that true activation function is non-linear and it can't be sharded.

Generally, I see activation is mistakenly used instead of pre-activation term. Pre-activation can be sharded because it is linear function, activation can not - it is non-linear. But this is a problem of many publications that abuse the term, in my opinion.

[manishucsd]

Equation (2) above has the text "Thus, for a single layer in the backward pass, we have" followed by "L" in in the FLOP/s equation (2).

[jacobaustin123]

Good point. Deleted.

[manishucsd]

From the FSDP section, please help me understand what you mean here

"FSDP shards the contracting dimension of the MLP weights along the data dimension."

While this is true that the contracting dimension is shared for the W_in, but the contracting dimension F for W_out is not sharded? It is still D which is not the contracting dimension.

[jacobaustin123]

Yes this is correct. I've fixed this.

[sshkhr]

Is jax.lax.psum_scatter same as reduce scatter where reduction is the sum operation? If so, why was a specific instance of reduce scatter implement in the JAX API instead of a more generic reduce scatter?

[jacobaustin123]

In general JAX/XLA primitives generally assume all reductions are sums. psum_scatter is a ReduceScatter with a sum, just as psum is an AllReduce with a sum. There are pmin/pmean/pmax operations for AllReduces.

[fidlej]

There is a minor typo in the Question 1 answer.
The Vocabulary parameters are correctly 0.3e9. But the total adds instead 0.39e9 (an extra 0.09e9).

Thanks the book.

[nishanthdikkala]

Thanks for the amazing book! :)

However, in practice, as our batch size increases, we tend to see diminishing returns in training since our gradients become almost noise-free. We also sometimes see training instability.

Are there any references (papers or blog posts) documenting this behavior?

[jacobaustin123]

https://arxiv.org/pdf/1812.06162 and https://arxiv.org/pdf/2001.08361 are two good references although it's a well-known phenomenon, generally the "optimal batch size" scales with the gradient "noise level". In general this is just a question of the goodness of our gradient estimators, which scale like 1 / sqrt(N).

[nishanthdikkala]

Thanks!

[yejingxin]

[jacobaustin123]

They can be overlapped with something called a collective matmul (link). The JAX section later in the book has an example of how this can be implemented.

[yejingxin]

For question 3, given F is small, and seqlen=32k is large enough (close to 8D=40k), so that attention flops is not neglectable, T_math will be T/8D=80% longer than the analytic model in this content, FSDP per device batch size threshold can be approximated as 848/(1+80%)=471, therefore, pure FSDP may be best.

[jacobaustin123]

This is an interesting and good point. Once we add in attention, a bunch of things change a bit, but after giving it some thought I think it's not necessarily super substantial. During training, we often do sequence sharding/ring attention which adds its own ICI bandwidth requirements. That means whatever extra computation is done also involves additional communication, roughly at a 1:1 ratio (moving Qs or Ks around between devices). If we aren't doing that and this truly requires no additional communication, then it may be that we want to do more FSDP than this book prescribes, since we can get away with it.

[eitanporat]

The black curve is the amount of time spent on model FLOPs, meaning any batch size where this is lower than all comms costs is strictly comms bound. You’ll notice the black curve intersects the green curve at about 1.6e10, as predicted.

Why 1.6e10? Is this a typo?

[jacobaustin123]

If you plug in the formula above, you get 4096 * 4 * 2550^2 / (2 * 8192 * 4) = 1.6e10.

[eitanporat]

irrespective of the computation itself the diagram shows the intersection at O(1.0e6)

[eitanporat]

moreover the figure has a label

We note that all strategies become comms‑bound below batch size 🔴 1.5e6,
which is in line with our expected 4096 * 2 * 2550² / (8192 * 4) = 🟢 1.6e10

[jacobaustin123]

Ah that should be 1.6e6.

[tushar00jain]

Is $$B_X$$ in data parallel correct? If it's a 2D mesh, this would indicate that the batch is also being copied. Should this be $$B_XY$$ instead so that each device gets a separate batch?

[tushar00jain]

btw $*_D$ notation is also not defined anywhere?

I can only see $*_{LOCAL}$

[jacobaustin123]

Where does pipelining use Z? I was intending for it only to use Z for clarity.

[jacobaustin123]

*_D means reduction along the D dimension, like an einsum.

[tushar00jain]

Where does pipelining use Z? I was intending for it only to use Z for clarity.

There's a $D_Y$ in the $Out$

$*_D$ means reduction along the D dimension, like an einsum.

Is this needed? Since In/W/Out etc. are matrices. If you're using things in the brackets as indices, then the summation would be implicit anyway?

[jacobaustin123]

The Y in pipelining was a typo. Yikes. Fixed now.

is this needed?

No, it's just nice for clarity in case the contracting dimensions aren't always at the beginning/end. It's defined in the first section I think.

[mayun8172]

T_comm and T_math is calculated for one layer in dp section. Are later sections calculated for one block instead of one layer?

[mayun8172]

oh，sorry. I know now

[zhiyisun]

Hello, the tensor parallelism figure shows two groups (TPU 0/1 and TPU 2/3) with replicated B and sharded D of In. For 1D model parallelism, wouldn’t a single sequence of 2 TPU along the Y-axis suffice (either group of TPU0/1 or TPU2/3)? The division into two identical groups seems redundant and may suggest an X-axis split, but this isn’t explained. Is this intentional (e.g., 2D TP) or a diagram oversight? Thanks for the valuable resource!

[AakashKumarNain]

@jacobaustin123 There is one thing that I feel can be added to the (FSDP + TP) strategy. The algorithm presented here is clean and clear, but if someone goes through the shard_map documentation and read this example, pretty sure it can cause confusion. May be we can add those examples from the docs for each strategy here?

[damek]

Could you explicitly define the effect of the number of mesh axes used by an algorithm, i.e., $M_X$ for FSDP (and symmetrically $M_Y$ for TP), on communication time, in one place?

Right now this idea appears across sections but I can't find it stated in a single sentence. In "Combining FSDP and Tensor Parallelism," the communication models effectively imply linear scaling with the number of axes:

$$ T_{\text{FSDP comms}} \approx \frac{4,D,F}{Y,W_{\mathrm{ici}},M_X}, \qquad T_{\text{TP comms}} \approx \frac{4,B,D}{X,W_{\mathrm{ici}},M_Y}. $$

This reads as: each additional independent ICI axis provides another “pipe” of bandwidth, so communication time scales as $1/M$ in the throughput-bound regime. Is that the way to think about it?

Separately, the "Sharding" chapter states the same intuition for multi-axis collectives (e.g., AllGather over multiple mesh axes), noting that operating over multiple axes increases the available bandwidth by a factor of the number of axes:

“What happens when we AllGather over multiple axes? When we gather over multiple axes, we have multiple dimensions of ICI over which to perform the gather. … This increases the available bandwidth by a factor of …”
(https://jax-ml.github.io/scaling-book/sharding/)

Perhaps after introducing $M_X$ and $M_Y$ in the Training chapter, you could add a one-liner like:

Definition. $M_X$ (resp. $M_Y$) is the number of hardware mesh axes that FSDP (resp. TP) collectives span.
Effect (bandwidth-bound): Using $M$ axes provides (\approx M\times) aggregate link bandwidth, so collective time scales $\propto 1/M_X$ (FSDP) and $\propto 1/M_Y$ (TP).

Probably there are caveats to this. It would be good to understand those as well.

Thanks for this!

[jacobaustin123]

Hi Damek. Yes I agree that this isn't introduced very cleanly. Yes, it provides another pipe that increases the total bandwidth proportionally. I'll add a comment to the training section.

[jacobaustin123]

I've added the note you suggested.

[colefranks]

You mentioned that layer pipelining is less essential in TPUs due to their faster comms. Could it still lead to gains over TP on TPU?