# 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 the per-device weight memory. The cost is an additional collective on each path that uses the enlarged group. **The MLP shards on both paths.** The MLP shards over `TP * mlp_dp_size` on the prefill graph and on the decode graph. The prefill path gathers the rows and reduces the results over the same enlarged group. So one process can serve the prefill shape and the decode shape from one graph. **The target and the Eagle3 draft shard their MLP separately.** `mlp_dp_size` widens the group for the target MLP. `draft_mlp_dp_size` widens it for the Eagle3 draft MLP. The two degrees are independent, because a MoE target must keep `mlp_dp_size` at 1 (see the constraints below) while its dense draft MLP still gains from the wider group. **Decode-only for the other components.** Prefill uses standard TP sharding for the attention, the embedding, and the LM head. `attention_dp_size`, `embedding_dp_size`, and `lm_head_dp_size` do not change the prefill graph. **DI (Disaggregated Inference) recommended.** A component that uses a different effective TP size on prefill and on decode must load its weights separately per graph. DI supports this, because prefill and decode run on separate compiled graphs. `mlp_dp_size` does not have this constraint. It uses the same group on both paths. ## Configuration ```python 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 — target MLP intermediate sharding draft_mlp_dp_size=4, # new — Eagle3 draft 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 (target) | Intermediate dimension (gate/up/down projections) | 1 | | `draft_mlp_dp_size` | Dense MLP (Eagle3 draft) | 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 expert parallelism (`ep_degree > 1` or `--enable-expert-parallel`) **on the static cross-DP-EP dispatch**. A MoE model shards its experts with EP, so `mlp_dp_size` would only change the batch layout of the target: the decoder layer all-gathers the rows over the MLP DP column and the dispatch all-gathers them again over the same ranks, which makes `mlp_dp_size` identical copies of one row block. The expert kernel computes every copy and the combine slices the copies back off. The platform validation rejects that combination when a server asks for it. The check lives there rather than on `NeuronConfig`, because which dispatch the MoE takes depends on the vLLM parallel config and the resolved all2all backend, and a `NeuronConfig` built on its own sees neither. Use `draft_mlp_dp_size` to widen the draft MLP group on a MoE target. - The all2all-v dispatch (`all2all_backend="neuron"`) is exempt. It keeps the rows sequence parallel, so the decoder layer skips the MLP DP transition and the dispatch never gathers over the DP group. There is no second gather, so `mlp_dp_size > 1` is allowed there. - `draft_mlp_dp_size` is compatible with expert parallelism, because it changes the draft MLP only. It does not change the row layout that the target hands to the draft. ## 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 ```text 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 ```python 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): ```text 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) ```text 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: ```python 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: ```python 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`, `draft_mlp_dp_size` config fields; rejects `mlp_dp_size > 1` with EP | | `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`, `draft_mlp_dp_size` params threaded through | | `vllm/worker/neuron_worker.py` | Reads new config fields and passes to `init_neuron_distributed_environment` |