SpectraX — True MPMD Pipeline Parallelism for JAX

TL;DR

SpectraX is a JAX-native neural-network library built around true MPMD pipeline parallelism — each physical rank compiles and runs its own distinct XLA program, with no shared shard_map HLO and no SPMD-same-shape constraint. It exposes an eager-Module API (modules are also JAX pytrees), nine pipeline schedules (GPipe, Std1F1B, Eager1F1B, ZeroBubbleH1, InterleavedH1/GPipe/1F1B+1, DualPipeV, KimiK2), and a unified spx.run(mesh) entry point that dispatches to SPMD (pjit) or MPMD (sxcall / sxjit) based on whether the mesh has an mpmd_axis. SpectraX is the NN core under EasyDeL (Llama, DeepSeek-V3, Qwen3, Gemma, GPT-OSS, Mamba, Whisper — 77 model families), so the same primitives are running production training and serving today.

Key claims

True MPMD compilation: sxjit traces a function, splits its jaxpr at sxstage_iter markers, and emits one XLA executable per pipeline rank — stages can differ in class, parameter shapes, and I/O shapes, so heterogeneous pipelines (embed → blocks → head) are first-class rather than padded to a common shape [README “MPMD Runtime”].
sxstage_region lets each branch or tower (vision V0→V3, text T0→T3) define its own local stage sequence and schedule, instead of pretending a multimodal stack is one long linear pipeline [README “Stage regions — multimodal pipelines without fake stages”].
Nine pipeline schedules ship behind a single Schedule interface (microbatches=, virtual_stages=), each exposing bubble_ratio, peak_activations, total_steps for analytical comparison before launch [README “Supported schedules”].
Forward, backward, and schedule cells all execute through the same scheduled MPMD dispatcher — legacy host-side schedule walking (fuse_1f1b, fuse_zb, chunks) is rejected on the true MPMD path rather than silently falling back to shard-map [README “Choosing & using a schedule”].
Symbolic-axis sharding (BATCH, EMBED, TP, FSDP, SP, EP) is resolved at runtime against the active mesh and runtime mode (train vs autoregressive decode); layer code never names PartitionSpecs directly. MPMD stage meshes drop the pipeline axis before inner SPMD layout to keep FSDP/TP/EP constraints stage-local [README “Sharding: the symbolic-axis system”].
One selector DSL powers grad(wrt=...), jit(mutable=...), Optimizer(wrt=...), freeze(...), iter_variables(select=...) — by collection name, by Variable subclass, by path glob, with set algebra [README “Selector DSL”].
Module-aware transforms cover spx.jit / grad / value_and_grad / vmap / scan / remat_scan / remat / cond / switch / while_loop / fori_loop / eval_shape / jvp / vjp — all respect the mutable= collection contract [README “Module-aware JAX transforms”].
Production claim: EasyDeL — covering 77 model families including Llama 3, DeepSeek-V3, Qwen3, Gemma, GPT-OSS, Mamba, Whisper — uses SpectraX as the NN core, with shared sharding, KV-cache, and pipeline plumbing [README “Used by”].
Built-in FP8 training with delayed scaling + rolling amax history through a fp8_meta mutable collection, and a LoRA path via nn.wrap_lora(base, rank, alpha) + spx.grad(wrt="lora") selector [README “Features”].
Dispatch overhead microbenchmark: ~150 µs vs. “300–2000 µs” for other JAX frameworks on a tiny 2-layer CPU transformer; reported 1.83× over flax.nnx at d=64 batch-2, 2.0× at d=48; on TPU-8B compute-bound workloads the Python gap shrinks but stays positive [README “Why SpectraX?” + “Benchmark”].

Method

The MPMD runtime is the engineering centerpiece. sxjit accepts a function written against an eager spx.Module graph, traces it once, then partitions the resulting jaxpr at sxstage_iter markers. Each cluster becomes a distinct XLA program compiled for its own rank; activation transport across boundaries honors the PartitionSpec declared on the marker. sxstage_region does the same trick at branch scope, so vision and text towers each get their own stage indices and schedules under an outer MPMD spx.jit.

Schedules are first-class objects rather than driver-loop conventions. GPipe, Std1F1B, Eager1F1B, ZeroBubbleH1 are “flat” (one stage per rank); InterleavedH1, InterleavedGPipe, Interleaved1F1BPlusOne, DualPipeV, KimiK2 use virtual_stages>1 to shrink the pipeline bubble at the cost of extra cross-rank transport. The scheduler emits fused forward / backward / weight-grad cells directly into the traced jaxpr (treduce is the binding primitive), so there’s no host-side stepping loop that diverges from the compiled schedule.

spx.run(model, inputs, targets, mesh, mode, loss_fn, microbatches, schedule) is the user-facing entry: if mesh.is_mpmd it routes to sxcall / sxjit, otherwise to pjit. For homogeneous transformer stacks it auto-detects a ModuleList named blocks and slices it evenly across pipeline stages. The sharding layer is symbolic — layer code uses tokens (BATCH, EMBED, TP, FSDP, SP, EP) via PartitionManager/apply_logical_sharding, and a PartitionAxis config maps tokens to mesh axes per mode. Pre-baked shapes ship for HiddenStateSharding, AttnQSharding, AttnKVSharding, RowWise, ColumnWise, Replicated, plus an Expert* family for MoE.

State lives in Variable cells (Parameter, Buffer, user subclasses) tagged with a collection name (parameters, buffers, lora, fp8_meta); the Selector DSL composes set-algebra over these. The jit cache is keyed on the input modules’ GraphDef snapshot — value mutation doesn’t invalidate, structural changes (attribute add/remove/replace) do.

Results

This is a library README, not an evaluation paper, so “results” are micro-benchmark numbers and the production deployment claim:

1.83× wall-clock vs. flax.nnx on a 2-layer / d=64 / batch-2 CPU transformer; 2.0× at d=48; “positive but smaller” on TPU-8B compute-bound runs (no exact number quoted for that case).
Schedule analytics (bubble_ratio, peak_activations, total_steps) are exposed so users can pick Std1F1B for O(n_stages) peak memory, ZeroBubbleH1 to fill the 1F1B bubble with weight-grad work, InterleavedH1(v=2) to shrink the bubble by v at the cost of more transport, or DualPipeV / KimiK2 for DeepSeek-V3 / Moonshot-K2-style bidirectional / interleaved schedules.
Production deployment via EasyDeL (77 model families: Llama, DeepSeek-V3, Qwen3, Gemma, GPT-OSS, Mamba, Whisper) is the strongest “does it work at scale” claim; no published throughput numbers vs. competing JAX stacks at frontier scale.
License is AGPL-3.0-or-later; requires Python 3.11+ and JAX ≥ 0.9.2; pip install spectrax-lib with CUDA/TPU extras.

Why it’s interesting

This is the JAX-side counterpart to the PyTorch infrastructure papers the wiki has been tracking. veScale-FSDP: Flexible and High-Performance FSDP at Scale argued that the missing FSDP primitive (RaggedShard) is what blocks modern optimizers (Muon, Shampoo) and block-wise FP8 from running without intrusive model surgery — that’s the data-parallel / ZeRO axis on PyTorch. SpectraX picks up the pipeline-parallel axis on JAX and makes the same kind of argument: SPMD-with-shard-map can’t express heterogeneous stages or multimodal branch pipelines without padding, and a properly per-rank compiled program is the right primitive. The two systems also share the “one symbolic-sharding contract through which everything routes” philosophy (RaggedShard placements vs. PartitionAxis tokens).

It also complements the productized side. Scaling and Optimizing Frontier Model Training (Fireworks AI) catalogs four parallelism dimensions (FSDP / Pipeline / Context / Expert) composed per training shape on PyTorch + GB200, and describes its own “streaming pipeline schedule redesigned for RL’s asynchronous data arrival.” SpectraX is the open, library-level analogue of that pipeline scheduling layer — with DualPipeV (DeepSeek-style) and KimiK2 schedules shipped as named objects rather than hidden in a proprietary trainer. For the wiki’s io-aware-kernel-design thread (FA4, Newton-Schulz on Blackwell), SpectraX is the upstream of where those kernels would land in a JAX training stack; it’s the substrate, not the kernel. Lastly: the JAX ML ecosystem currently does not have a stable concept page on the wiki — most JAX-relevant filings live in narrower threads (Enabling Up to 41% Faster Pre-training: MXFP8 and DeepEP for DeepSeek-V3 on B200 with TorchTitan is PyTorch-side, Gram Newton-Schulz: A Fast, Hardware-Aware Newton-Schulz Algorithm for Muon is kernel-side). SpectraX is the first system-side filing keyed on JAX as the host runtime.