mixle.utils.parallel.tensor_pipeline_context_parallel module

TP/PP/CP for CausalLM, atop the existing FSDP2 support (F1).

mixle/models/transformer.py names the destination directly: “At frontier scale the same module is what a vendored TorchTitan/Megatron trainer shards (FSDP2/TP/PP).” torch_neural.py already gives the data-parallel dimension (DDP on CPU, FSDP2/ZeRO-3 on CUDA). This module adds the three ORTHOGONAL sharding dimensions a frontier trainer composes with FSDP2 – “N-D parallelism”: FSDP2 shards params/optimizer state across the data-parallel group while TP/PP/CP further shard the MODEL and the SEQUENCE across independent device groups:

• TP (ColumnParallelLinear / RowParallelLinear, tp_shard_causal_lm()) – splits CausalAttention’s qkv/proj and the MLP’s two Linear layers across tp_size ranks (Megatron-style: column-parallel then row-parallel, so exactly one all-reduce per sublayer), by HEAD for attention (each rank owns whole heads, never a fraction of one) and by hidden-unit block for the MLP.

• PP (pp_partition_causal_lm(), pipeline_forward()) – splits model.blocks into pp_size contiguous stages (stage 0 also owns the embeddings, the last stage also owns ln/head), and runs a GPipe-style microbatched pipeline: stages are threads connected by queues, so microbatches genuinely overlap in flight across “devices” (this repo’s existing thread-based distributed-simulation pattern – see multiprocessing.py / mpi.py).

• CP (cp_shard_sequence(), cp_forward_causal_lm()) – splits the SEQUENCE into cp_size contiguous chunks. Token/position embeddings, LayerNorm, the MLP, and the LM head are all per-position and need no communication; only attention needs the other chunks’ K/V, so each rank computes its local K/V, all-gathers everyone else’s (one collective per block), and runs LOCAL causal attention with an explicit offset mask (its query positions against the FULL key sequence). This is the “simpler chunked approach” the roadmap calls out as an acceptable scope cut vs. incremental ring-attention: it reconstructs bit-for-bit-equivalent output (same collective volume as ring attention, just gathered up front instead of streamed rank-to-rank) and is exact and testable without incremental overlap.

None of this touches real multi-GPU: there are no 512 A100s in this environment (or in CI), so the roadmap’s “70B-config across >=512 GPUs at published-comparable MFU” acceptance number is NOT measured here and cannot honestly be claimed from a laptop/CI run – see the test module’s docstring for what IS verified (exact-match correctness of the TP/PP/CP mechanism at small scale). What’s built here is the real sharding/reconstruction MATH, following the structure a TorchTitan integration would slot into (tp_size/pp_size/cp_size device-mesh axes orthogonal to FSDP2’s data-parallel axis); a full TorchTitan integration would additionally need: real multi-GPU process groups per axis (NCCL, not the in-process simulation here), overlap of TP’s all-reduce with compute, 1F1B (not GPipe fill-drain) pipeline scheduling, and incremental ring-attention communication for CP’s memory profile at long context.

validate_tp_pp_cp_plan(model, tp_size=1, pp_size=1, cp_size=1)

Validate a (tp_size, pp_size, cp_size) plan against a real CausalLM’s dimensions.

Raises ValueError with an actionable message if the plan does not divide the model cleanly – the same checks tp_shard_causal_lm() / pp_partition_causal_lm() / cp_shard_sequence() enforce structurally, surfaced up front so lm.fit(distributed=True, ...) fails fast on a bad plan instead of partway through a run. This is the plan-construction half of the tp_size/ pp_size/cp_size knobs on fit(); wiring the validated plan into per-axis NCCL process groups (real SPMD TP/PP/CP execution, composed with the existing FSDP2 data-parallel group) is the multi-GPU piece this environment cannot exercise – see the module docstring and torch_neural.py’s FSDP2 CUDA branch, which carries the identical caveat (“correct per the API, only exercised on multi-GPU”).

Parameters:

• model (Module)

• tp_size (int)

• pp_size (int)

• cp_size (int)

Return type:

None

class ColumnParallelLinear(weight, bias)

Bases: object

A Linear’s OUTPUT dimension split across ranks; reconstruction is a concat (all-gather).

weight[r] is a contiguous row-block of the dense weight (out_features split into n_ranks chunks); bias[r] the matching bias slice (or None). Each rank’s local matmul x @ weight[r].T + bias[r] is exactly the corresponding output slice of the dense layer, so concatenating the ranks’ outputs along the last dim reconstructs the dense output.

Parameters:

• weight (list[Any])

• bias (list[Any] | None)

forward(x)

Reference/non-distributed reconstruction: run every shard and all-gather (concat).

Parameters:

x (Any)

Return type:

Any

class RowParallelLinear(weight, bias)

Bases: object

A Linear’s INPUT dimension split across ranks; reconstruction is a sum (all-reduce).

weight[r] is a contiguous column-block of the dense weight (in_features split into n_ranks chunks). Each rank’s local matmul against ITS input slice sums, across ranks, to the dense output; the bias is carried by rank 0 only (added once) so the sum stays exact.

Parameters:

• weight (list[Any])

• bias (Any | None)

forward(x_shards)

Reference/non-distributed reconstruction: sum every shard’s partial output (all-reduce).

Parameters:

x_shards (list[Any])

Return type:

Any

class TPAttentionShard(n_head_local, qkv_weight, qkv_bias, proj_weight, proj_bias)

Bases: object

One rank’s shard of a CausalAttention: whole heads of qkv (column) + matching proj rows (row).

Parameters:

• n_head_local (int)

• qkv_weight (Any)

• qkv_bias (Any)

• proj_weight (Any)

• proj_bias (Any | None)

class TPBlockShard(attn: 'list[TPAttentionShard]', mlp_fc1: 'ColumnParallelLinear', mlp_fc2: 'RowParallelLinear')

Bases: object

Parameters:

• attn (list[TPAttentionShard])

• mlp_fc1 (ColumnParallelLinear)

• mlp_fc2 (RowParallelLinear)

class TPCausalLMShard(blocks: 'list[TPBlockShard]', tp_size: 'int')

Bases: object

Parameters:

• blocks (list[TPBlockShard])

• tp_size (int)

tp_shard_attention(attn, tp_size)

Shard a CausalAttention into tp_size head-parallel ranks.

Parameters:

• attn (Module)

• tp_size (int)

Return type:

list[TPAttentionShard]

tp_attention_forward(x, shards)

Run head-parallel attention across the (simulated) ranks and reconstruct the dense output.

Each rank: local qkv projection (its whole heads only) -> local causal attention -> partial (b, t, head_dim * n_head_local) activation. All-gather (concat, in rank order == head order) reconstructs the o the dense CausalAttention would compute; the row-parallel proj then sums the ranks’ partial output projections (all-reduce) plus rank 0’s bias – exactly the dense proj(o).

Parameters:

• x (Any)

• shards (list[TPAttentionShard])

Return type:

Any

tp_shard_causal_lm(model, tp_size)

Shard every block of a CausalLM for tp_size-way TP.

Token/position embeddings and the final ln/head are NOT sharded here (they are cheap relative to attention/MLP and, per Megatron, are typically the sequence-/vocab-parallel dimension rather than TP proper) – this covers the attention+MLP sharding the spec calls out explicitly.

Parameters:

• model (Module)

• tp_size (int)

Return type:

TPCausalLMShard

tp_forward_causal_lm(model, x, tp_shard)

Forward an input through the TP-sharded blocks (attention/MLP), embeddings/head run dense.

Parameters:

• model (Module)

• x (Any)

• tp_shard (TPCausalLMShard)

Return type:

Any

class PPStage(blocks, *, tok=None, pos=None, ln=None, head=None)

Bases: Module

One pipeline stage: a contiguous slice of model.blocks, optionally with embeddings and/or the final ln/head (stage 0 embeds, the last stage projects to logits).

Parameters:

• blocks (list[nn.Module])

• tok (Any)

• pos (Any)

• ln (Any)

• head (Any)

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Parameters:

x (Any)

Return type:

Any

pp_partition_causal_lm(model, pp_size)

Split model.blocks into pp_size contiguous stages (GPipe-style layer partition).

Stage 0 additionally owns the token/position embeddings; the LAST stage additionally owns the final ln/head – so stage 0 takes raw token ids and the last stage emits logits, and every intermediate stage is a pure activation-in/activation-out block group (what gets pipelined).

Parameters:

• model (Module)

• pp_size (int)

Return type:

list[PPStage]

pipeline_forward(stages, x, n_microbatches)

GPipe-style microbatched pipeline: split x’s batch dim, run stages as threads-with-queues.

Each stage is a thread reading its input queue and writing to the next stage’s; the driver feeds microbatches into stage 0’s queue back-to-back (no waiting for one to finish before starting the next), so microbatches genuinely overlap in flight across stages – the “devices” this repo’s existing thread-based distributed-simulation tests stand in for real ranks with (see multiprocessing.py). Since every op here (LayerNorm, attention, MLP, embeddings) is batch-independent, splitting the batch into microbatches and reassembling in order is exactly equivalent to running the whole batch through the un-partitioned model.

Parameters:

• stages (list[PPStage])

• x (Any)

• n_microbatches (int)

Return type:

Any

cp_shard_sequence(x, cp_size)

Split a (batch, seq) (or (batch, seq, ...)) tensor into cp_size contiguous sequence chunks along dim 1 – each rank keeps one chunk resident (never materializes the full sequence).

Parameters:

• x (Any)

• cp_size (int)

Return type:

list[Any]

cp_attention_forward(attn, chunks)

Context-parallel attention: each rank computes local Q/K/V, all-gathers K/V (one collective), then runs LOCAL causal attention of its Q chunk against the FULL (gathered) K/V with an explicit offset causal mask. Returns the per-rank output chunks (concat along seq to reconstruct the dense CausalAttention output) – this is the “simpler chunked” CP scope noted in the module docstring: same total K/V communication volume as ring attention, gathered eagerly instead of streamed incrementally rank-to-rank (that overlap is the piece a real ring-attention CP would add).

Parameters:

• attn (Module)

• chunks (list[Any])

Return type:

list[Any]

cp_forward_causal_lm(model, x, cp_size)

Full CP forward: per-block, only attention needs the K/V all-gather – embeddings, LayerNorm, MLP, and the LM head are all per-position (no communication) and run locally on each chunk. Returns per-position logits for the WHOLE sequence ((batch, seq, vocab)), reconstructed by concatenating the ranks’ chunks – so CP correctness is checked at every position, not just last.

Parameters:

• model (Module)

• x (Any)

• cp_size (int)

Return type:

Any