DualPipe from Scratch: Implementing DeepSeek's 5D Parallelism in PyTorch - Dev Jadhav, ING Bank

This talk focuses on the Dual Pipe implementation in PyTorch, presented by Theo Jadav, a tech lead machine learning engineer. He begins by addressing the inefficiencies in training large models across multiple GPUs, specifically highlighting that traditional scheduling methods can lead to significant idle time and wasted resources. The Dual Pipe scheduling technique aims to optimize this process by allowing two microbatches to be processed simultaneously in opposite directions. Jadav outlines the foundational concepts of memory requirements, scheduling algorithms, and model architecture, explaining how they relate to efficient parallel processing. He also dives into five critical discoveries made during implementation, which are not covered in the original DeepSeek V3 paper. The session concludes with the practical implications of these findings, emphasizing the importance of effective scheduling in high-parameter models and the applicability of Dual Pipe to various architectures.

Hi. Welcome everyone. Quick question for the room. If you're training a model across eight or more GPUs in pipeline, what percentage of the compute time do you think your GPUs are actually doing useful work? Any guesses? I think most people will guess 70, maybe 80%. The answer with the standard G pipe schedule is closer to 57%. Nearly half of your GPU time and your budget are wasted

on idle time. That's the problem I spent last year working on. I am Theo Jadav, a tech lead machine learning engineer in ING Bank, but today I am here in my personal capacity to talk about my open source work. The DeepSeek V3 paper describes 5D parallelism and dual pipe at high level, but leaves critical implementation details undocumented. Today, I shall present my open source PyTorch reference implementation

that fills those gaps. I have three parts for you. First, what is dual pipe and why does it matter? Then, how we implemented it, the five discoveries that paper doesn't tell you. Finally, why this matters beyond DeepSeek, where it generalizes, what research it builds Uh prerequisite for this talk is to have a basic understanding of torch.distributed. Everything is open source and Apache 2.0 license. Uh it is

uh the link for the GitHub repository is on the screen. Before diving into the dual pipe, let me frame three key bottlenecks in any large model training system. First, memory. 671 billion parameters nearly rough roughly uh needs 6 terabytes of memory, including gradients, optimizer states, and activations. One GPU has around 80 GB. You physically cannot fit this mo uh model without splitting it across dozens of devices.

Second, the scheduler. Once you split the model into pipeline stages, how you schedule the computation determines how much GPU time is productive versus wasted. This is where the dual pipe lives. It's a scheduling algorithm for pipeline parallelism. Third, model architecture. DeepSpeed V3 uses 256 mixture of expert multi-head latent attention and five dimensions of parallelism. The architectural choices dictates which scheduling strategies are even possible. Today, we are

focusing on second bottleneck, uh specifically our PyTorch implementation of dual pipe. The details that paper doesn't uh tell Here's the problem dual pipe was invented to solve. This is a G-pipe schedule with four pipeline stages and four microbatches. Look at the stage zero. It processes all four forward microbatches, then sits idle while backward passes propagates. That time uh that that time is the idle and it is

the bubble. The formula is straightforward. The bubble overhead equals P minus 1 divided by M, where P is the pipeline stages and M is the microbatches. With DeepSpeed configuration, P equals to 8 and M equals to 16 microbatches. That's 7 divided by 16, which is 43.75%. Nearly half of your GPUs time is wasted. At 2048 GPUs running for months, that's millions of dollars in idle compute. DeepSpeed

needed to get this below 5%. That's why they invented dual pipe. Dual pipe's core insight: run two microbatches simultaneously, but in opposite direction. Stream A sends microbatches from stage zero to stage P-1. That's standard left-to-right pipeline. Stream B sends microbatches in reverse from stage P-1 to zero. Uh what happens when stream A's forward passes occupy the first pipeline stages, stream B's forward passes fill the last stages.

The bubble gets filled from both the ends. The number 3% bubble overhead instead of 47%. And 55% throughput gain over Gpipe. And hierarchical all-to-all architecture gives you four times reduction in mixture of expert dispatch overhead. Now, let me show you how we build this in PyTorch and the five critical implementation details that paper doesn't give you. All five dimensions are implemented with torch.distributed process groups over here.

Tensor parallelism splits weight matrices within the node. Tensor parallelism equal to four matching the NVLink topology. Pipeline layers into the stages. PP equals to eight. And this is what dual pipe schedules. Data parallelism replicates the model using distributed data parallel. DP equals to 64. Expert parallelism distributes 256 mixture of experts. EP Sequence parallelism splits activations along the sequence dimension to reduce the memory. They compose multiplicatively. 4

* 8 * 64, which is equals to 2048 GPUs. Each dimensions gets its own process groups created with dist.new_group. The group compose via rank mapping. Each GPU knows exactly which process group it belongs for each dimensions. This formula is the single most important implementation detail in dual pipe. And it is not in the Deep Seek paper. We derive it by analyzing the scheduling pattern. Warm-up steps equals

to P floor divided by two plus M minus one times two. P floor divided by two is the number of steps for one stream to reach the midpoint of the pipeline. Since the two streams travels in opposite direction, they meet in the middle after P over two steps. M minus one times two is the extra step each additional micro batch needs. That factor of two comes from

the bidirectional nature. With P equals to eight and M equals to the warm-up equals to 4 plus 30, which is equal to 34 steps before steady state. Getting this wrong is catastrophic. Too few warm-up steps, you start backward propagation passes before activation exist. Training crashes with NaN. Even if you off by one, the uh the pattern the overlap pattern breaks and torch.distributed.barrier hangs forever, a deadlock. And

with the warm-up formula, you get exactly 3% bubble overhead. And both streams overlap perfectly. Here's the actual PyTorch implementation from our repo, dual_pipe.py file. The schedule step enum captures four phases, warm-up forward, steady state forward, steady state backward, and cool down. This is one of the three teachable abstraction from our submission also. Making the schedule step enum not a string, not an integer, let you unit test

the entire scheduling logic without the GPUs. Run by test test/torch/test_dual_pipe.py repository. You don't need a CUDA to run this one. We got three scheduling bugs this way before they even hit distributed code. Line 12 is the warm-up formula we just discussed. The generate schedule method walks through the total steps and classifies each one. In steady state, both streams are active active simultaneously. That's where the bidirectional overlap

happens. The second critical implementation detail. Communication overlap using torch.distributed async primitives. The pattern has four steps. Step one, start async send of backward gradients using dist.isend. This returns a work handle and the GPU continues immediately. Step two, while the network interface shifts the data, compute the forward pass for the next micro batch. Step three, call handle.wait to ensure the backward send completed. Step four, start async send

of a forward output. The key distinction is asynchronous. It returns immediately and the GPU keeps computing. dist.send is a synchronous. It blocks the GPU until the data is fully delivered. Using the blocking version serializes everything and destroys the entire overlap benefit. The async API is not optional over here. It's a architecturally essential for dual pipe. Discovery number one, directly from our CFP, YK_PE is shared across heads

in decoupled rotational position encoding. The paper says they decouple the key vector. K equals the concatenation of K underscore C and K underscore PE. But, it's ambiguous about whether K underscore PE is computed per head or shared across all the heads. From our MLA.py implementation, we found it it must be shared. Compute K underscore PE once with ROPE applied to the projection, then unsqueeze to add the

head dimension, and expand to broadcast Why does this matter? If you implement K underscore PE as per head, the model still trains. No error, no NaN, the gradients are technically non-zero, but each head learns redundant position information instead of sharing it. The model converges to a measurably worse solution. And you won't know unless you are carefully comparing against the reference output. Discovery number two. The critical timing

of bias update in auxiliary loss pre-load balancing and the exponential moving average initialization bug. DeepSeek V3 uses a bias term instead of auxiliary loss to balance expert loads, but the initialization isn't specified in the paper. We found that initializing the exponential moving average with torch.zeros creates a terrible routing for the first 100 steps. Some expert receives 10 times more tokens than others. Three experts were handling 80%

of all the tokens. This took us two weeks to track it down because there's no error. The model trains just wastefully. The fix, initialize with torch.full at one over number of experts, makes uniform distribution. And critically, update the bias before making routing decisions, not after. The order is update EMA, then compute bias, then add to scores, and then route tokens. Reversing step three and four causes oscillating

expert loads. Discovery three from our CFP, how sigmoid routing separates selection scores from gate values. With f.softmax, all 256 expert scores must sum to one. Changing one expert score automatically change every other expert's probability. This creates 256 way gradient competition on every token. With torch.sigmoid, each expert score is computed independently. The selection step picks top K by raw sigmoid score. Then and only then the selected gates

are normalized by dividing by their sum. The decoupling is critical at 256 experts. Selection ask which experts using independent probabilities, getting asked how much weight, and normalizes only after the selection is made. Two separate decisions, not one couple decision. Tensor parallelism in PyTorch. Over here, we use column-to-row composition trick. Column parallel, each rank calls torch.mm it's with its column chunk of the weight matrix. Partial output column,

no communication needed. Row parallel, each rank computes torch.mm with its row chunk. Then dist.all_reduce sums the partials. The key pattern compose column parallel followed by row parallel. The split output from the first layer fits directly to the split input to the second. This eliminates one per transformer block. At hundreds of blocks, that's a thousands of save communication rounds. Hierarchical all-to-all, our CFP topic on two levels. communication

reducing mixture of expert dispatch overhead by four times. Two process groups created with this.new_group. One for intra-node rank using NVLink at 900 GB/s communication. One for inter-node leaders using InfiniBand at 50 GB/s. per second 18 times slower. The dispatch become torch. all to all calls. Level one Local dispatch within the node group. Sort tokens, so local experts are served first. Level two Only ship tokens that absolutely

must cross the node boundaries. Result four times fewer messages on the slow inter-node network. Standard torch.distributed primitives, no custom NCCL kernel needed. This bug from our CFP cause causal mask position offset is the most insidious because it fails completely silently. With sequence parallelism, each GPU sees a chunk of this full sequence. GPU zero has tokens zero to 2047. GPU one has a token 2048 to 4095 and

so on. Each GPU's local indices start at zero. The bug using torch.tril with local indices to create a causal mask. GPU one's mask act if its chunk starts at position zero, allowing tokens to attend the future tokens in other chunks. The fix one-line offset equals this.get_rank of the ESP process group times chunk size. Pass that offset to the mask creation function. Simple arithmetic. Catastrophic if miss. No

error, no NaN, no crash. The model trains with subtle information leakage and produces slightly worse perplexity. We caught it by visualizing attention patterns against a non-parallel reference. A capacity matrix data class, one of our three teachable abstraction from CFP. It tracks per expert capacity, load ratio, and overflow count. When an expert exceeds capacity, we use torch.argsort on gate scores. Highest score stays, lowest gets dropped. The critical

detail, dropped tokens are redistributed to underloaded expert, not discarded. Every token contributes to at least one expert's gradient, maintaining training signal quality. Over here. The three teachable abstraction from our CFP, all are pure PyTorch. The schedule makes pipeline schedule composable and testable. Run the pytest test/torch/test_dualpipe.py. No CUDA We tested dual pipe, one forward, one backward, and Gpipe with the same test harness. Capacity matrix data class decouples

monitoring from routing logic. Drop it into any torch mixture of expert, TensorBoard compatible, log expert utilization per step. Expert specialization tracker monitors which expert learn which token patterns using EMA smooth statistics. Catches expert collapse before loss degrades. Weight and bias integration ready. These patterns make complex distributed PyTorch code maintainable. Without them, you have a research prototype. With them, you have something others can learn and build on.

Part three, why this matters beyond our implementation and Deep Seek. The benchmark result. Gpipe has 47% One forward, one backward has 30% interleave. One forward, one backward has 15% and dual pipe has 3% bubble overhead. 55% throughput gain for the same hardware. And why deep six specifically? Because 671 billion parameters with 256 experts creates extreme parallelism demand. You need pipeline parallelism equals eight stages. That's where the

bubble is catastrophic. And dual pipe's bidirectional nature shines. For small smaller models with PP equals to two the bubble is already small and giving you a diminishing returns. But any model using eight or more pipeline stages, like Lama 3405 billion parameter model, Mistral, the current Grok models uh those all uses dual pipes uh and it can be directly applied to them. The scheduler is model agnostic. It's

pure torch.distributed. Dual pipe extend the 1F1B scheduling work from Narayanan and is compatible with the zero from Raj Bhandari and composes with mega Megatron-LM tensor parallelism pattern. Four takeaways. First, the scheduler is the key leverage point in pipeline parallelism training. Dual pipe cuts the bubble from 47% to 3%. That's the difference between wasting nearly half your compute and wasting almost none. Second, the warm-up formula that makes

the dual pipe work is not in the paper. Getting it wrong by one step means torch.distributed deadlock. All GPUs hanging forever. five silent bugs that don't crash but degrades quality. KPE sharing, bias update timings, sigmoid routing separation, causal mask offset, capacity dropping priority. Every one of them fails silently. Fourth, three teachable abstractions. Schedule step, capacity metric, expert that turn the complex distributed PyTorch code into something testable,

monitorable, and extensible. The code is on the GitHub. That's it. That's all. Thank you for attending. If you have any questions, I'm here.