Forget about the Dropout here in the figure, we will check it over the SP section.

so let's start from a simple example of FFN (dense, no MoE), which is more easy to understand:

Without TP, it is:

With TP=4, we split along $H_{\mathrm{ffn}}$:

where

So on GPU i:

Then the final output is $Z = \sum_{i=1}^{4} Z_i$. The derivation is as follows:

Equivalently, $Y = [Y_1, Y_2, Y_3, Y_4]$, and

so

then let's move to Attention, which is a little bit more complex. But the basic idea is still the same.

Let's consider the vanilla Self-Attention:

For head h, where $d_{\text{head}} = H / N_h$,

Concatenate all heads:

Output projection:

With tensor parallelism

We split heads across GPUs, since heads are decoupled in the attention mechanism. (say, for context, you get this qi,kj problem)

Let TP degree be p=4. Split heads across GPUs, so each GPU holds $N_h/p$ heads.

where

and split $W_O$ row-wise:

On GPU i:

with

GPU i computes attention only for its own head subset:

Then local output projection:

Each O_i is a partial sum. Final output:

equivalently,

Why this works

Since

and

we have

Takeaway: matrix multiply with a block-partitioned inner dimension is a sum of independent sub-products. That's why AllReduce (sum) gives the exact answer.

1. X (full) × W_first_i (sharded)  →  intermediate_i (sharded) | no comm
2. compute on intermediate_i independently                     | no comm
3. intermediate_i × W_last_i (sharded)  →  partial_sum_i       | no comm
4. AllReduce(partial_sums)  →  correct output                  | 1 AllReduce

Takeaway: VPP is the best you can do in upstream Megatron. The research frontier has moved on to zero-bubble variants, and DeepSeek uses DualPipe in production — but adopting either requires patching a fork.

PP cuts the model's layers into sequential chunks and places each chunk on a different GPU. A forward pass becomes a relay — GPU 0 runs its layers, hands the activation to GPU 1, and so on. Backward flows in reverse.

The communication between adjacent stages is small (one activation tensor per hand-off), and only between neighbors — which is why PP is the natural choice for crossing slow inter-node links (InfiniBand) where TP's all-reduce bandwidth would be too expensive.

PP=4, M=1; F/B each take one slot; communication time omitted.
.. = idle

Time:   00 01 02 03 04 05 06 07
GPU 0:  F0 .. .. .. .. .. .. B0
GPU 1:  .. F0 .. .. .. .. B0 ..
GPU 2:  .. .. F0 .. .. B0 .. ..
GPU 3:  .. .. .. F0 B0 .. .. ..

PP=4, M=4; F/B each take one slot; communication time omitted.
.. = idle

Time:   00 01 02 03 04 05 06 07 08 09 10 11 12 13
GPU 0:  F0 F1 F2 F3 .. .. .. .. .. .. B3 B2 B1 B0
GPU 1:  .. F0 F1 F2 F3 .. .. .. .. B3 B2 B1 B0 ..
GPU 2:  .. .. F0 F1 F2 F3 .. .. B3 B2 B1 B0 .. ..
GPU 3:  .. .. .. F0 F1 F2 F3 B3 B2 B1 B0 .. .. ..

PP=4, M=7; F/B each take one slot; communication time omitted.
.. = idle

Time:   00 01 02 03 04 05 06 07 08 09 10 11 12 13 14 15 16 17 18 19
GPU 0:  F0 F1 F2 F3 .. .. .. B0 F4 B1 F5 B2 F6 B3 .. B4 .. B5 .. B6
GPU 1:  .. F0 F1 F2 .. .. B0 F3 B1 F4 B2 F5 B3 F6 B4 .. B5 .. B6 ..
GPU 2:  .. .. F0 F1 .. B0 F2 B1 F3 B2 F4 B3 F5 B4 F6 B5 .. B6 .. ..
GPU 3:  .. .. .. F0 B0 F1 B1 F2 B2 F3 B3 F4 B4 F5 B5 F6 B6 .. .. ..

PP=4, VPP=2, 32 layers:
  GPU 0: chunks {0, 4} -> layers [0-3]   AND [16-19]
  GPU 1: chunks {1, 5} -> layers [4-7]   AND [20-23]
  GPU 2: chunks {2, 6} -> layers [8-11]  AND [24-27]
  GPU 3: chunks {3, 7} -> layers [12-15] AND [28-31]

cN = model chunk N; arrows show forward order

GPU0/c0 -> GPU1/c1 -> GPU2/c2 -> GPU3/c3
                                    |
                                    v
GPU3/c7 <- GPU2/c6 <- GPU1/c5 <- GPU0/c4

Standard 1F1B:
  [B_input + B_weight] -> send dInput to the previous stage

Zero Bubble:
  [B_input]            -> send dInput to the previous stage
  [B_weight]           -> ready before the optimizer step
       ^
       can run later in a free slot on the same GPU

Dependencies only; horizontal positions are not timestamps.

Same physical GPU order in every row:

Standard 1F1B:
  Forward:   GPU 0 -> GPU 1 -> GPU 2 -> GPU 3
  Backward:  GPU 0 <- GPU 1 <- GPU 2 <- GPU 3

DualPipe, forward directions:
  Stream A:  GPU 0 -> GPU 1 -> GPU 2 -> GPU 3
  Stream B:  GPU 0 <- GPU 1 <- GPU 2 <- GPU 3
  Each stream's backward runs opposite to its forward.

A possible overlap on GPU i (schematic, not to scale):
  A forward:  [EP dispatch        ]
  B backward: [independent compute]

In an MoE (Mixture-of-Experts) model, each transformer block contains a router that sends tokens to a subset of many parallel FFN "experts." The whole point of MoE is to scale total parameters without scaling per-token compute — each token only activates its top-k experts. But those expert weights still live in memory. At scale this becomes the dominant cost: DeepSeek-V3, for example, has 256 routed experts per MoE layer, each a SwiGLU FFN with intermediate dim 2,048 on a hidden dim of 7,168 — about 44M params (~88 MB in bf16) per expert. One MoE layer's experts alone take ~22 GB; across all 58 MoE layers, total expert params exceed 1 TB. No single GPU can hold that. EP distributes experts across GPUs: with EP=8, each GPU holds 256/8 = 32 experts per layer, and tokens are routed to the right GPU via all-to-all communication.

SP is 'doing a favor' after TP. We set it on by default

Problem: In a Transformer layer with TP, the weight matrices (attention, FFN) are already split across GPUs. LayerNorm has small scale/bias parameters; Dropout has none. LayerNorm normalizes each token across its hidden dimension; Dropout operates element-wise. In vanilla TP, every GPU holds the full [S, H] activation for these ops, which is redundant and wastes memory.

Insight: Vanilla TP uses AllReduce after each block (attention, FFN). An AllReduce is mathematically equivalent to ReduceScatter + AllGather. SP simply pulls these two steps apart and inserts LayerNorm/Dropout in between:

Vanilla TP (per block):

compute on sharded weights → AllReduce → LayerNorm(full [S, H])
                                ↑ equivalent to ReduceScatter + AllGather back-to-back

SP (per block):

compute on sharded weights → ReduceScatter → LayerNorm(partial [S/TP, H]) → AllGather
                                   ↑ each GPU gets a correct slice of S

After ReduceScatter, each GPU holds the correct, fully-reduced result for S/TP tokens with the full hidden dimension H. Since LayerNorm normalizes each token independently across H, computing it on [S/TP, H] gives the same result as on [S, H]. The AllGather then reconstructs the full sequence before the next TP-parallelized op.

Why the split is S/TP: ReduceScatter across a group of N GPUs always produces N chunks — one per GPU. SP reuses the TP group, which has TP GPUs, so the chunk count is TP. This is a constraint of the group size, not a design choice.

Takeaways:

• Same communication volume as vanilla TP (ReduceScatter + AllGather = AllReduce)
• Activation memory per GPU reduced from [S, H] to [S/TP, H] for LayerNorm/Dropout
• No new GPU groups, no new communication — just reordering what TP already pays for
• Mathematically identical output; always enable when TP > 1

Mode 1: Ulysses (a2a) — swap layout, not data movement during attention

Ulysses is an algorithm from the DeepSpeed paper⁶, not a single implementation. Two open-source versions exist: DeepSpeed's original Python implementation, modifiable step-by-step; and NVIDIA Transformer Engine's fused CUDA implementation, faster but exposed as one opaque forward() call. Upstream Megatron-LM with cp_comm_type='a2a' dispatches to TE. In fact, [a2a, p2p, allgather, hybrid] are all implemented in TE.

The idea: instead of moving KV during attention, rearrange the data once before attention so each rank ends up with the full sequence but only a subset of heads. Then run a vanilla local attention. Reverse the rearrangement after.

This picture shows clearly how Ulysses works, the most important thing is to understand these two alltoall communications. Without loss of generality, we explain the first alltoall communication, that turns local [N/P, d] tensor to local [N, d/P], since the last alltoall communication is symmetric.

To be short, it is split into P of [N/P, d/P] on each GPU, and then send to the corresponding GPU to get [N, d/P].

Initially GPU p holds $X^{(p)}$, one block of rows.

Then each local tensor $X^{(p)}$ is further split into P column-blocks:

During all-to-all, GPU p sends $X^{(p)}_{q}$ to GPU q and GPU q receives $X^{(0)}_{q}, X^{(1)}_{q}, \dots, X^{(P-1)}_{q}$. Then GPU q concatenates these along the sequence dimension:

Mode 2: Ring (p2p) — pass the KV around

Before Step in, let's have a fast recap of the online softmax in flash attn

Keep each block's running max m, denominator l, and output out; when two blocks merge, both get rescaled to the new max before being added — the code below is exactly that.

m_new = max(m_old, m_current)
l_new = exp(m_old - m_new) * l_old + exp(m_current - m_new) * l_current
out_new = (exp(m_old - m_new) * l_old * out_old + exp(m_current - m_new) * l_current * out_current) / l_new

This tells us that softmax can be computed block by block (in flash attention), and ring attn is doing the same thing, but in a $S/CP$ way.

The figure is quite intuitive, since ffn is token-wise, so we only need to consider attn. At a high level, the pattern is:

1. each GPU starts with its own local KV block,
2. it computes attention between its local queries and the KV block currently in hand,
3. meanwhile, it sends that KV block to the next GPU and receives another KV block from the previous GPU,
4. it repeats this process round by round until it has seen all KV blocks it needs.

So Ring Attention turns the global attention problem into a streaming computation: compute on one KV block, pass it on, receive the next one, and keep accumulating.

The communication pattern

Each round, every GPU does exactly two communication operations: • isend its current KV package to the right neighbor • irecv the next KV package from the left neighbor

These operations are non-blocking. That means they are only launched first; the GPU does not stop and wait immediately. The actual transfer proceeds on a communication stream, while the main compute stream can continue running the attention kernel on the KV block that is already available.

This is the whole point: communication and computation can overlap.

Round    Compute (default)   | Transfer (cp_stream)
0        KV-own (buffer 0)   | send own; recv KV-A -> buffer 1
1        KV-A   (buffer 1)   | send A;   recv KV-B -> buffer 0
2        KV-B   (buffer 0)   | send B;   recv KV-C -> buffer 1
3        KV-C   (buffer 1)   | no further transfer

Send right; receive from the left.
Reuse a buffer only after its compute and transfer finish.

Compute is still flash-attn.

  Full attention (S × S)
    ├── Ring splits it into CP² blocks, each (S/CP × S/CP)     ← across GPUs
    │     └── Flash Attention splits each block into tiny tiles  ← within one GPU's SRAM
    │           └── each tile: ~128 × 128, fits in SRAM
    │
    └── Online softmax merges at BOTH levels
          - inner: tile-to-tile within Flash Attention
          - outer: chunk-to-chunk across ring steps
          - same (m, l, out) formula both times

For causal we have the zigzag assignment, to ensure the load balance.

Mode 3: AllGather — just bring all KV to everyone

AllGather CP is the brute-force mode. Dead simple.

The idea: before attention, every GPU just gathers the full K and V from all CP ranks. Now every GPU has full KV in memory. Then each GPU runs plain, standard
attention — its local Q (S/CP queries) against the full K, V (all S keys/values). No ring, no partial merges, no online softmax juggling. Just normal attention
after one big collective.

Not recommended, since it is not efficient.

Hybrid CP (Hierarchical Context Parallel)

Hybrid CP splits one CP group into two nested communication levels.

Suppose you set CP=16 and cp-comm-type=a2a+p2p. Megatron then has to decide:

• which ranks should do all-to-all locally
• which groups should talk through the outer ring

That is exactly what --hierarchical-context-parallel-sizes specifies. You pass two integers [inner, outer] such that:

Example: 2 nodes x 8 GPUs = 16-way CP

Use [8, 2]:

Within each node: 8-rank a2a over NVLink.

Node 0: [R0  R1  R2  R3  R4  R5  R6  R7 ]
          |   |   |   |   |   |   |   |
Node 1: [R8  R9  R10 R11 R12 R13 R14 R15]

Each vertical link is a 2-rank p2p group over IB:
{R0,R8}, {R1,R9}, ..., {R7,R15}

• inner = 8: the 8 GPUs inside each node do Ulysses-style all-to-all over NVLink.
• outer = 2: the two nodes communicate with each other through a ring over InfiniBand.

Why this layout works:

• all-to-all is bandwidth-heavy, so it should stay on the fast intra-node fabric
• ring is more latency-tolerant because its communication can overlap with compute

Why [2, 8] is a bad idea:

• each 2-GPU pair would do tiny all-to-alls
• the outer ring would span 8 groups across nodes
• you would push much more traffic onto InfiniBand and waste NVLink bandwidth

Mental shortcut:

• inner_size = GPUs per node, or more generally GPUs inside one fast-link domain
• outer_size = number of such domains participating in the CP group

One constraint: in Transformer Engine v2.18's a2a path, inner_size must divide the local Q and KV head counts after TP, not just the model's total num_attention_heads, because Ulysses splits by heads^**.

Dynamic CP for variable-length packing (--hybrid-context-parallel)

Despite the name, this has nothing to do with hierarchical a2a+p2p. It solves a different problem: when training with packed sequences (multiple variable-length samples concatenated into one batch), not every sub-sample needs the same CP degree. A 128K sub-sample needs many CP ranks; a 2K sub-sample fits on one GPU and gains nothing from splitting.

--max-seqlen-per-dp-cp-rank tells the scheduler the memory capacity of a single GPU, measured in tokens: "one GPU can hold at most this many tokens without OOM." The scheduler then decides, for each sub-sample, how many CP ranks to spread it across:

For example, with max-seqlen-per-dp-cp-rank = 16384:

Long samples get spread across many ranks, short samples stay local, and no GPU wastes memory on a tiny sample split across an unnecessarily large CP group.

For a simplified 7B example with bf16 params/grads and fp32 Adam moments, excluding fp32 master parameters, each DP rank holds:

That's already over an 80 GB H100 — and 7B is small. With DP=8, every replica wastes 7× the memory it needs to. The DP algorithms below are progressively more aggressive ways to share that wasted memory across the DP group.

Default behavior. No DP flags set.

The baseline. Every DP rank holds the full state. After backward:

Backward done on every rank → each rank has its local 14 GB grad tensor
                            → AllReduce across DP group  (1 collective)
                            → every rank now has the same averaged grad
                            → optimizer step runs locally on every rank (redundantly)

Comm: 1 AllReduce per step (just gradients). Memory: nothing saved. Every rank holds full params + grads + optimizer state.

This is what you get if you set none of the DP flags. It's the simplest, but the most memory-hungry.

Enabled by --use-distributed-optimizer. Megatron's own implementation, conceptually identical to DeepSpeed's ZeRO-1.

The cheapest, biggest memory win. Optimizer state is the largest single component (Adam keeps two fp32 tensors per parameter — momentum and variance, ~56 GB for our 7B model). Sharding it across DP ranks recovers most of the wasted memory immediately.

Each rank now holds:
  - Full params      (still 14 GB)
  - Full grads       (still 14 GB)
  - 1/DP optim state (e.g., 7 GB instead of 56 GB at DP=8)

The communication pattern flips:

Backward → ReduceScatter grads (instead of AllReduce)
           Each rank now holds 1/DP of the grad — the slice corresponding to
           its 1/DP of the optimizer state.

Optimizer → Local optimizer step on the 1/DP slice.

Forward (next step) → AllGather params, because each rank only updated its slice.

Comm: 1 ReduceScatter (grads) + 1 AllGather (params) per step. Same total volume as DDP's AllReduce, just split into two halves. Memory: optimizer state shrinks by DP×. For our 7B example, recovers ~50 GB per rank.

This is the highest-leverage move you can make. Almost every production training job uses at least this.

Enabled by --use-megatron-fsdp (which requires --use-distributed-optimizer). The depth is controlled by --data-parallel-sharding-strategy, which defaults to 'optim_grads_params' (full ZeRO-3) when FSDP is enabled.

Going further. The --data-parallel-sharding-strategy flag exposes four levels of sharding aggressiveness:

The 'optim' and 'optim_grads' strategies save more memory at modest extra comm cost. Going to full ZeRO-3 introduces a qualitatively new behavior worth understanding.

The runtime cost of sharding parameters. With ZeRO-3, each GPU only holds 1/DP of every layer's parameters. But to actually run a layer's forward, you need the whole weight matrix for that layer. So Megatron-FSDP issues a just-in-time AllGather before each layer's forward, materializes the full weight for one matmul, then releases it:

Layer 1 forward:  AllGather params  →  matmul  →  release
Layer 2 forward:  AllGather params  →  matmul  →  release
...

One AllGather per layer per forward. Same on backward. This is much more comm volume than ZeRO-1, but the reward is that the param memory itself is sharded — the biggest possible memory saving.

Megatron-FSDP enforces a hard requirement: it only works when --use-distributed-optimizer is also set. The two flags must travel together.

For the full-shard setup below, use --use-megatron-fsdp with 'optim_grads_params' and --num-distributed-optimizer-instances > 1. Companion flag --outer-dp-sharding-strategy controls outer-group behavior.

Full ZeRO-3 has a problem at large scale. When DP is large enough to cross node boundaries, the per-layer AllGather goes over InfiniBand (~50 GB/s) instead of NVLink (~600 GB/s). At 60 layers × 1 forward + 1 backward, you pay 120 cross-node AllGathers per micro-batch. The comm dominates wall-clock.

HSDP fixes this by splitting the DP dimension into two nested levels:

HSDP layout (DP=64, num_instances=8):

  Inner level (within each node, 8 GPUs):    Full ZeRO-3 sharding
  Outer level (across the 8 nodes):          Replication

  Node 0           Node 1           ...    Node 7
  +-------------+  +-------------+         +-------------+
  | GPU 0: 1/8  |  | GPU 0: 1/8  |         | GPU 0: 1/8  |
  | GPU 1: 1/8  |  | GPU 1: 1/8  |   ...   | GPU 1: 1/8  |
  | ...         |  | ...         |         | ...         |
  | GPU 7: 1/8  |  | GPU 7: 1/8  |         | GPU 7: 1/8  |
  |             |  |             |         |             |
  | Full copy   |  | Full copy   |         | Full copy   |
  | of model    |  | of model    |         | of model    |
  +-------------+  +-------------+         +-------------+
         ^                ^                       ^
     replicas of each other (the "outer" axis = 8 replicas)

What changes:

• The per-layer AllGather (params) → only across the inner group of 8 GPUs → NVLink only, fast
• The gradient sync (once per step) → across the outer group of 8 nodes → cross-node IB, but it had to happen anyway in any DP scheme
• Memory cost: each node now holds a full copy of the model. So HSDP uses outer× more memory than full ZeRO-3 — in the diagram, 8 replicas instead of 1. Worth it because cross-node param AllGathers are gone.

The companion --outer-dp-sharding-strategy controls the outer group:

• 'no_shard' (default): outer is pure replication — the picture above
• 'optim': outer also shards optimizer state across replicas. Saves more memory, costs more cross-node comm. Only valid when inner is 'optim_grads_params'.

In Megatron, set --num-distributed-optimizer-instances=N where N is the desired outer degree (number of replicas). Inner = DP / N.

Enabled by --use-torch-fsdp2. Mutually exclusive with --use-megatron-fsdp.

Megatron also exposes PyTorch's native FSDP2 as an alternative implementation of the same ZeRO-3 algorithm. It's largely experimental — the upstream code lists six incompatibilities:

• Not compatible with PP (pipeline parallelism)
• Not compatible with EP (expert parallelism)
• Not compatible with MCore distributed optimizer
• Not compatible with fp16
• Not compatible with gradient accumulation fusion
• Requires --ckpt-format torch_dist or torch_dcp, plus --untie-embeddings-and-output-weights

For production large-scale training (which usually involves PP and/or EP), FSDP2 isn't a realistic choice. It exists for small experiments and for users committed to the PyTorch-native stack.

One ZeRO-2-like toggle worth noting: --torch-fsdp2-no-reshard-after-forward keeps full params in memory between forward and backward, avoiding the backward AllGather. Parameters are still resharded after backward by default; this flag alone does not keep them permanently replicated^††.

The same overlap recipe we've seen for PP, CP, and EP applies to DP. ZeRO's ReduceScatter and AllGather are in the critical path; both can be hidden:

Two PP×DP interaction flags:

• --no-align-grad-reduce — disable the synchronized launch of grad reduces across PP stages (alignment defaults ON)
• --no-align-param-gather — same, for param gathers (alignment defaults ON with interleaved PP; otherwise disabled)^§

When alignment is on, the pipeline schedule aligns the launch points of DP collectives, rather than enforcing an exact wall-clock instant^§. Predictable, but can over-subscribe NCCL. Disabling alignment lets each PP stage fire when it's actually ready — better concurrency, more complex scheduling.

One more optimization: --create-all-gather-group creates a separate process group for AllGather so it can run concurrently with ReduceScatter (instead of sharing one group's resources).

Let $P$, $G$, $O$ be a model partition's parameter, gradient, and optimizer-state bytes, and $d$ the sharding-group size (in the CP=1 examples here, $d = DP$; for HSDP with $N$ replicas, let $h = d/N$ be the inner-group size).

For the 7B example above ($P=14$, $G=14$, $O=56$ GB) at $d=8$, ZeRO-1/2/3 give 35 / 22.75 / 10.5 GB — 41.7% / 27.1% / 12.5% of the 84 GB baseline; HSDP's fraction is $1/h$ (12.5% at $h=8$), not a fixed 50%. These are idealized model-state formulas, not peak-memory guarantees.

Decision tree:

• Model fits in DDP memory? → leave everything off. You get DDP.
• Need more memory? → --use-distributed-optimizer (ZeRO-1). Recovers most of the optimizer state. Almost always worth it.
• Still not enough? → add --use-megatron-fsdp and pick 'optim_grads' or 'optim_grads_params' via --data-parallel-sharding-strategy.
• Full ZeRO-3 with DP crossing nodes? → add --num-distributed-optimizer-instances=N to switch to HSDP. Inner = DP / N.

expert_DP is computed separately because the expert view of world_size factors differently: $\text{expert\_DP} = \text{world\_size} / (\text{expert\_TP} \times PP \times EP)$. The same DP/HSDP/FSDP machinery applies to expert parameters within their own DP group. The "Relation between Parallelism" section below works through how the dense and expert views coexist on the same physical GPUs.

GPU (device) A single accelerator chip. In modern training setups, a typical machine has 8 GPUs, all visible via nvidia-smi as devices cuda:0 through cuda:7.

Node One physical machine — usually one server with 8 GPUs, its own host CPU, RAM, and NICs. GPUs within a node talk to each other over fast intra-node links (NVLink / NVSwitch). GPUs in different nodes talk over the slower inter-node fabric (InfiniBand / RoCE).

Process One operating-system process running the training script. In Megatron-LM (and almost every modern framework), there is one process per GPU. The process owns its GPU, holds its slice of the model, and runs forward/backward.

So an 8-GPU node runs 8 processes in parallel. A 64-GPU job (8 nodes × 8 GPUs) runs 64 processes total. They communicate through NCCL (or Gloo) collectives, not shared memory.

Every process has two integer IDs:

local_rank The process's index within its own node, ranging from 0 to gpus_per_node - 1. On an 8-GPU node, local ranks are 0..7. This is the ID you use to pick which GPU to bind to:

torch.cuda.set_device(local_rank)

global_rank (often just called rank) The process's index across the entire job, ranging from 0 to world_size - 1. On a 64-GPU job, global ranks are 0..63. This is the ID NCCL uses to address messages between processes.

world_size The total number of processes in the job. For 8 nodes × 8 GPUs, world_size = 64.

The relationship is straightforward when nodes are uniform:

global_rank = node_id * gpus_per_node + local_rank

A process group is a subset of ranks that participate in the same collective communication call (AllReduce, AllGather, etc.). Megatron-LM creates several overlapping process groups, one per parallelism dimension:

Each rank belongs to all of these groups simultaneously — its position is one coordinate in a multi-dimensional grid. For example, with TP=4, PP=2, CP=8, DP=1 on 64 GPUs, a single rank has a (tp_rank, pp_rank, cp_rank, dp_rank) tuple identifying it within each group.

Takeaway: a "rank" is an integer ID for a process, not a role. The same process is tp_rank=2 in the TP group, pp_rank=0 in the PP group, cp_rank=5 in the CP group, and so on — all at the same time. Process groups are the primitive that makes this possible.