mixle.experimental.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 local Q/K/V and streams preceding K/V blocks through an online-softmax accumulator. The in-process reference holds all simulated ranks, but no rank’s attention calculation concatenates a full-sequence K/V tensor; the returned receipt records that memory invariant.

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.