This document is relevant for: Trn2, Trn3

Component-Level DP Sharding#

Overview#

Each model component (Attention, Embedding, MLP, LM Head) can independently shard its weights across TP * component_dp_size devices instead of just TP. This reduces per-device weight memory at the cost of additional collectives during decode.

Decode-only. Prefill uses standard TP sharding and is unaffected by these flags.

DI (Disaggregated Inference) required. Since prefill and decode use different effective TP sizes, weights must be loaded separately per graph. This is naturally supported in DI setups where prefill and decode run on separate compiled graphs.

Configuration#

NeuronConfig(
    attention_dp_size=4,  # existing — Q/O (and optionally K/V) sharding
    embedding_dp_size=4,            # new — Embedding vocab sharding
    mlp_dp_size=4,                  # new — MLP intermediate sharding
    lm_head_dp_size=4,              # new — LM Head vocab sharding
)

Each parameter controls one component independently:

Parameter

Component

What’s sharded

Default

attention_dp_size

Attention Q/O/KV

Q heads, O projection, optionally KV heads

1

embedding_dp_size

Embedding

Vocabulary dimension

1

mlp_dp_size

Dense MLP

Intermediate dimension (gate/up/down projections)

1

lm_head_dp_size

LM Head

Vocabulary dimension (output projection)

1

Constraints:

  • Each *_dp_size must be a positive integer that divides the overall dp_size

  • embedding_dp_size and lm_head_dp_size are independent (can differ even with tied embeddings — weights are loaded separately)

  • mlp_dp_size > 1 is incompatible with ep_degree > 1 (MoE models use EP for expert sharding instead)

Process Groups#

Two types of groups are created per component:

Component TP group (for weight sharding and TP collectives)#

A rectangular block of TP * component_dp_size ranks. Passed as tp_group to the module (VocabDimShardedEmbedding, ColumnParallelLinear, or used for MLP all-reduce). The module sees this as its TP group — weight sizing, weight loading rank, and collectives all use this group.

DP Column Group (for batch all-gather and slice)#

A column of component_dp_size ranks at the same TP position across DP groups. Used by the decoder layer’s _dp_transition function to gather/scatter batches between modules with different DP sizes.

Example: TP=8, DP=4, all dp_sizes=4#

         TP rank:  0   1   2   3   4   5   6   7
         ─────────────────────────────────────────
  DP0:           [ 0   1   2   3   4   5   6   7 ]
  DP1:           [ 8   9  10  11  12  13  14  15 ]
  DP2:           [16  17  18  19  20  21  22  23 ]
  DP3:           [24  25  26  27  28  29  30  31 ]

Component TP group (size 32): [0,1,...,31]  — all ranks
  Used by: embedding, MLP, LM head, attention for weight sharding + all-reduce

DP Column (size 4): [0,8,16,24], [1,9,17,25], ..., [7,15,23,31]
  Used by: _dp_transition for batch all-gather/slice between modules

Group Aliasing#

When multiple components share the same dp_size, their groups are aliased (same GroupCoordinator object, created once). A cache in _create_neuron_groups ensures no duplicate groups.

Groups Always Created#

All Component TP groups and DP column groups are created unconditionally, even when dp_size=1. With dp_size=1, the Component TP group equals the regular TP group, and the DP column group is a size-1 self-group. This ensures _dp_transition works without conditional logic.

Batch State Machine (Decode)#

During decode, the batch dimension may be “gathered” across multiple DP ranks. The batch gathered size tracks how many DP ranks’ batches are present in the tensor. The _dp_transition function handles transitions between different gathered states.

_dp_transition#

def _dp_transition(x, current_group, target_group, dim=0):
    current_dp = current_group.world_size
    target_dp = target_group.world_size
    if current_dp == target_dp:
        return x                          # no-op
    if current_dp > target_dp:
        # Down: slice to keep this rank's target sub-group batches
        per_dp = x.shape[dim] // current_dp
        start = (current_group.rank_in_group // target_dp) * target_dp * per_dp
        return x.narrow(dim, start, target_dp * per_dp)
    # Up: go through local to avoid duplicates, then all-gather to target
    per_dp = x.shape[dim] // current_dp
    x = x.narrow(dim, current_group.rank_in_group * per_dp, per_dp)
    return target_group.all_gather(x, dim=dim)

Decoder Layer Flow#

The decoder layer has fixed internal transitions (no state threading needed):

Input arrives at mlp_dp (from previous layer or embedding transition)

  _dp_transition(mlp_dp → attn_dp)
  ┌─ Self Attention ──────────────────────────────────────┐
  │  Input: gathered at attn_dp                           │
  │  Megakernel: QKV + attention + O projection           │
  │  attn_tp_group.all_reduce: sums TP+DP weight partials │
  │  Output: gathered at attn_dp                          │
  └───────────────────────────────────────────────────────┘
  residual add (both at attn_dp)
  _dp_transition(attn_dp → mlp_dp)

  ┌─ MLP ────────────────────────────────────────────────┐
  │  Input: gathered at mlp_dp                            │
  │  gate/up/down projections                             │
  │  mlp_tp_group.all_reduce: sums TP+DP weight partials  │
  │  Output: gathered at mlp_dp                           │
  └──────────────────────────────────────────────────────┘
  residual add (both at mlp_dp)

Output at mlp_dp

Full Model Flow (Decode)#

Embedding:
  all-gather input_ids across emb_dp column        [B_local] → [B_local * emb_dp]
  VocabDimShardedEmbedding (uses emb Component TP group)
  _dp_transition(emb_dp → mlp_dp)                 transition once before layer loop

Layer 0..N (each layer):
  _dp_transition(mlp_dp → attn_dp)                no-op if equal
  Attention (uses attn Component TP group all-reduce)
  _dp_transition(attn_dp → mlp_dp)                no-op if equal
  MLP (uses mlp Component TP group all-reduce)

Norm (elementwise, state-agnostic)

LM Head:
  Slice to local batch                             [B_local * mlp_dp] → [B_local]
  index_select(sampling_positions)                 [B_local] → [B_selected]
  all-gather across lm_head_dp column              [B_selected] → [B_selected * lm_head_dp]
  ColumnParallelLinear (uses lm_head Component TP group)
  Slice to local batch                             [B_selected * lm_head_dp] → [B_selected]

All dp_sizes Equal#

When attn_dp == mlp_dp == emb_dp == lm_head_dp:

  • All _dp_transition calls are no-ops (same group sizes)

  • Embedding gathers once, data stays gathered through all layers

  • Only the LM head slices back to local at the end

Collectives per layer: 3 (attn Component TP group all-reduce + MLP Component TP group all-reduce + residual add) vs current without optimization: 7 (AG + RS + TP-AR + AG + slice + TP-super-AR per layer)

Mixed Case#

When dp_sizes differ, _dp_transition inserts the minimal collective at each boundary. Each transition goes through local (slice → all-gather) to avoid duplicated data from overlapping gathered states.

Weight Sharding and Loading#

Embedding and LM Head#

The Component TP group is passed directly as tp_group to VocabDimShardedEmbedding and ColumnParallelLinear. These modules automatically size their weights based on the group’s world_size and use rank_in_group for sharding.

Since load_weights passes tp_rank (not the component TP group rank) to the checkpoint loader, with_rank_override wraps the weight loader to substitute the component TP group rank:

emb_loader = sharding_weight_loader(shard_dim=0, shard_size=..., num_shards=emb_tp_size)
emb_loader = with_rank_override(emb_loader, rank=emb_tp_group.rank_in_group)
set_weight_loader(embed_tokens.weight, emb_loader)

MLP#

Same pattern as Embedding/LM Head — with_rank_override using the MLP Component TP group rank:

gate_up_loader = sharding_weight_loader(shard_dim=1, shard_size=..., num_shards=mlp_tp_size)
gate_up_loader = with_rank_override(gate_up_loader, rank=mlp_tp_group.rank_in_group)

Attention#

Attention weight sharding is unchanged from the existing attention_dp_size implementation (interleaved effective rank for Q/O, separate handling for K/V).

Attention Changes#

The attention decode path was simplified as part of this work:

Before: Internal all-gather + reduce-scatter + TP all-reduce (3 separate collectives) After: Caller provides gathered input via _dp_transition, attention uses single attn_tp_group.all_reduce (1 collective)

The attention Component TP group (_NEURON_ATTENTION_TP, size TP * attn_dp) replaces the previous two-step collective. The all_reduce across the component TP group sums both TP partials and DP weight-shard partials in one operation.

Internal all-to-all for Q/KV head swapping is unchanged (handled by the megakernel).

Implementation Files#

File

Change

model/neuron_config.py

embedding_dp_size, lm_head_dp_size, mlp_dp_size config fields

parallel/neuron_parallel_state.py

Component TP groups (_NEURON_*_TP) and DP column groups (_NEURON_*_DP) for all components; _build_component_tp_group_ranks() helper; _NEURON_ATTENTION_TP component TP group

model/llama3/model.py

_dp_transition() helper; LlamaDecoderLayer batch state transitions; LlamaMLP uses MLP Component TP group; LlamaModel embedding DP bracket; LlamaForCausalLM LM head DP bracket; LlamaAttention uses attention Component TP group all-reduce

utils/executor.py

embedding_dp_size, lm_head_dp_size, mlp_dp_size params threaded through

vllm/worker/neuron_worker.py

Reads new config fields and passes to init_neuron_distributed_environment

This document is relevant for: Trn2, Trn3