🐰大模型分布式训练篇——从零实现 Tensor Para (English)

Generated: 2026-06-22 11:52:13

---

The Pitfalls of Distributed Training: It Took Me Two Whole Weeks to Finally Understand TP, SP, and Communication Overlap

Hey, friend! Let me tell you about a topic I've got a love-hate relationship with—distributed training. Two years ago, when I first wrote Tensor Parallel code, I almost smashed my keyboard!

Guess what? I was testing on a 6B model with TP size=2, thinking, "The memory usage should definitely be cut in half!" But when I ran it... it only dropped by about 40%! I stared at the screen for ten minutes, then started frantically digging through my code. Finally, I found it—I had chained the communication and computation together. The GPU was just sitting there waiting for data, wasting a ton of time.

Speaking of which, have you ever run into something similar? Don't worry—today I'm going to dump all the lessons I've learned from my mistakes and all the hard-earned knowledge, straight to you!

---

1. What Exactly Does TP Slice? Don't Rush to Write Code—Let Me Give You an Analogy First

Imagine you have a giant cake—this is your weight matrix W, shaped [h, h']. Computing on a single card is like one person eating the whole cake: Y = X @ W. Simple, right?

But the cake is too big for one person. What do you do? Slice it!

Column Parallel: It's like cutting the cake vertically down the middle, each card gets half. W1 takes shape [h, h'/2], W2 takes [h, h'/2]. Each card computes its own Y1 = X @ W1, Y2 = X @ W2. Finally, you need to stitch the results together—that requires one AllGather communication, just like two people combining their halves back into a whole cake.

Row Parallel: This time, cut horizontally! W1 takes shape [h/2, h'], W2 the same. Each card gets half of the input X1 and X2, computes Y1 = X1 @ W1, Y2 = X2 @ W2, and then sums the results—requiring one AllReduce, like two people adding up their individual scores to get the total.

Sounds easy, right? I thought so too.

But there's a trap that nearly drove me crazy—I naively assumed that if I sliced things in the forward pass, the backward pass gradients would be handled automatically. What happened? The model's loss curve didn't match the single-card run at all! It took me three days of debugging to realize I had written the communication in the backward pass incorrectly. Annoying, isn't it!

---

2. Sequence Parallel—Another Trap That Made My Memory Explode

While running TP, I hit my second major pitfall.

When the sequence length got long (say, 2048), memory would blow up. I was baffled: "Didn't I already slice the weights? Why is the memory still so huge?"

Then it hit me—TP only slices the weights and computation, but it doesn't touch the activations!

What are activations? They're the intermediate results computed by each layer, stored for use during backpropagation. For a sequence, every token has its own activation values. The longer the sequence, the larger the activations. It's like cooking while also writing down the recipe—the more dishes you make, the more notes you accumulate, and you eventually run out of paper!

That's when Sequence Parallel (SP) comes to the rescue! In a nutshell, it slices along the sequence length dimension. Originally, each card processed the entire sequence s; now, each card only handles a s/n length sequence. The activations drop to 1/n directly! Awesome, right?

But the trade-off is increased communication. Because Attention needs global Key and Value, SP requires an AllGather before Attention to collect all KVs, compute, and then each card proceeds on its own. Like writing an essay where you need to reference the notes from everyone in the class—you have to borrow them first, read them, then return them.

I tested this on an 8-card machine: without SP, at sequence length 4096 the memory overflowed; with SP, I could run up to 8192 with plenty of headroom! But the training time increased by about 15%. This is a classic case of "trading time for space"—worth it or not? Depends on whether you have enough memory.

---

3. Computation-Communication Overlap—So Worth It! No More Waiting Idly for the GPU

So far, I've been describing a serial pattern: compute, then communicate, then compute again. It's like cooking: you chop vegetables, then wait for water to boil, then put them in the pot. What are you doing while waiting for the water to boil? Standing around!

Think about it—when your GPU is communicating, what's it doing? It's waiting for data to come from other cards. During that time, the compute units are idle. What a waste!

Can we overlap computation and communication? In theory, yes! It's like boiling water while chopping vegetables at the same time—no delays.

How exactly?

Step 1: Compute a portion
Step 2: Start communication (async, launch it)
Step 3: Meanwhile, continue computing the rest
Step 4: When communication finishes, merge the results

The key is to split your computation into two independent parts—one that depends on the communication result, and one that doesn't. Like cooking: seasoning requires waiting for the water to boil (dependent), but chopping side ingredients doesn't (independent).

Applied to TP implementation: For Column Parallel, after computing the local result Y1 in the forward pass, you need an AllGather to merge. But you can kick off the AllGather first, and while waiting, process subsequent Layer Norm or other computations that don't depend on the complete Y.

In my actual tests, communication overlap reduced TP's communication overhead by about 40%! Exactly how much depends on your network bandwidth and compute ratio. If your network is 10 Gbps and your GPU is an A100, the overlap effect is quite noticeable; with NVLink and high-end GPUs, the benefit might be smaller.

---

4. Details in Code Implementation—Pitfalls I've Already Stepped Into for You

A reader once asked me: "Why is there allreduce instead of reducescatter in the backward of your ColumnLinear?"

Good question! It depends on the specific scenario.

In the code I posted in that article, the forward uses identity (no communication), and the backward does all_reduce. Why?

Because in the forward pass, each card only computes its local Y, passes it on without communication, and the full gradients are only needed when computing the loss later. In the backward pass, each card's gradient is local (only for its own weights), and an all_reduce sum is required to get the correct global gradient. It's like each person only calculates the score for the part they're responsible for, and at the end you need to sum them up to know the total.

But there's a variation: If you already performed communication in the forward pass (e.g., an AllGather to merge outputs), then you can use ReduceScatter in the backward for better efficiency. In short, the choice of communication operation depends on when you split and when you merge.

Later, I implemented both approaches in my open-source code, ran a comparison, and found that the direct identity + backward all_reduce is friendlier for small TP sizes and simpler to code. However, when TP size is 4 or larger, the forward AllGather + backward ReduceScatter combination yields higher throughput.

Remember this rule: Use simple solutions for small scales, and optimized solutions for large scales. Don't start getting fancy too early!

---

5. Actual Results—Guess What? Enabling TP Actually Made Training Faster!

After all that theory, let's look at what happened when I actually ran it.

I tested on a small Transformer with 12 layers, hidden=512, MLP=2048, and compared three configurations:

Configuration	Training Time	Memory Usage	Final Loss

No TP	67.9 min	100%	2.625

TP size=2	50.5 min	55%	2.625

TP size=4	39.4 min