Design Space 04: Distributed Programming Model

Series: AI-Native Programming Language Research — Design Space Memos Document 4 of 6 (planned: shapes, AD-through-effects, sparsity, distribution, verification, hardware abstraction)

1. Problem statement

Modern foundation-model training runs across thousands of accelerators. Inference for frontier models uses pipelines that span dozens of devices. Even fine-tuning a 70B model is a multi-GPU exercise. Distribution is no longer an optimization — it is the default execution mode for non-trivial ML.

Yet the programming model is essentially HPC code from 1995 with a thin Python wrapper. The user thinks in terms of:

• Process groups and ranks (MPI heritage).
• Hand-coded all_reduce / all_gather / all_to_all calls.
• Layer-by-layer parallelism strategies (Megatron-LM, DeepSpeed) chosen by hand.
• Hybrid 3D / 5D parallelism configurations baked into framework setup files.
• Manual placement of optimizer states, activations, and gradients across devices.

The single most successful piece of compiler-driven distribution is GSPMD (Xu et al., 2021): the user annotates a few key tensors with sharding specifications, and the compiler propagates sharding through the rest of the graph and inserts collectives automatically. JAX pjit / shard_map, PyTorch's DTensor, and Alpa all build on this idea. It works. And it works because, in essence, GSPMD treats sharding as a type and propagation as type inference — without saying so explicitly.

The thesis of this memo: sharding is the fourth refinement axis on the tensor type, alongside shape (memo 1), aliasing (memo 2), and sparsity (memo 3). Treating it that way makes communication a type-changing coercion, makes auto-parallelization a constraint-solving problem over the type, and unifies distribution with the rest of the language design. GSPMD has the implementation; what's missing is the formal type system and the surface language that exposes it.

2. Taxonomy of parallelism

A distributed ML program is decomposed along several orthogonal axes. The main ones:

Real frontier models use all of these simultaneously. PaLM used 12-way TP × 64-way DP × 2-way PP. Frontier MoE models layer EP on top. Long-context models add SP and CP. The combinatorics are real, and choosing the right configuration is an optimization problem nobody actually solves — they pick a configuration that worked for a similar model and call it done.

The taxonomic observation: all of these are statements about how a tensor's axes map to a device mesh's axes. DP is "batch axis sharded over the data mesh axis." TP is "hidden axis sharded over the model mesh axis." SP is "sequence axis sharded over the data mesh axis." Different names for the same kind of object: a sharding specification.

This is exactly what GSPMD made explicit. It just hasn't propagated into language design.

3. Mesh as a type-level construct

A device mesh is a logical grid of devices with named axes. A 64-GPU cluster might be configured as a [8, 8] mesh with axes ("data", "model"), or as a [4, 4, 4] mesh with axes ("data", "model", "pipeline").

In current systems the mesh is a runtime configuration object (jax.sharding.Mesh, torch.distributed.DeviceMesh). In a type-aware language it should be a type-level construct: a kind, a phantom type, or a parameter on the tensor type. Concretely:

Tensor f16 [B, S, D] on mesh M with shard {B: data, S: seq, D: model}

The mesh M is a type-level value with named axes; the sharding specification maps tensor axes to mesh axes. This is exactly what jax.sharding.NamedSharding represents at runtime — the proposal is to lift it to compile time so the type checker can reason about it.

What you get from this lift:

1. Mesh-axis arithmetic in the type system. A [B, S, D] tensor with B → data and another [D, V] tensor with D → model produce a result whose sharding can be computed from the operation, just as the result shape can. Most of GSPMD's propagation rules are typeable this way.

2. Sharding mismatches as type errors. If you try to add two tensors with incompatible shardings, the compiler tells you, just as it would for incompatible shapes. Today this is a runtime error or a silent-resharding-with-bad-performance.

3. Resharding as an explicit (and visible) coercion. When the compiler inserts an all-gather to reshard a tensor, it shows up in the IR (and ideally in profiler tooling) as an explicit operation, not as a hidden cost.

4. Mesh polymorphism for portability. A function written generically over a mesh axis variable can run on different mesh sizes without modification. This is the "scale agnostic" property GSPMD argues for, made first-class in the type system.

The closest existing work: PartIR (Google), shard_map (JAX), and PyTorch DTensor all expose explicit sharding types. None of them lift the mesh itself to compile time.

4. Communication as type-changing operations

The most compelling consequence of typing sharding: collectives become type-changing operations whose insertion the compiler can infer.

Concrete examples:

• all_gather(x, axis="model"): changes the sharding of x from {D: model} to {D: replicated}. The result type is Tensor f16 [B, S, D] with no model-axis sharding.
• reduce_scatter(x, axis="data"): changes {} to {B: data} and reduces across the data axis. Used in ZeRO-3.
• all_to_all(x, src_axis, dst_axis): rearranges shardings; central to expert parallelism.
• all_reduce(x): type-preserving in shape, but the post-condition is "this tensor is consistent across the group."

In a typed setting, the compiler runs an inference pass that walks the graph, propagates sharding types through operations, and inserts the cheapest collective when shardings don't match at an operation boundary. This is what GSPMD does, dressed up as type inference.

Cost-aware insertion. The cheap-collective-insertion problem is non-trivial because there are usually multiple correct ways to make shardings agree. all_gather then operate vs. operate-then-all_reduce produce the same result with different costs. Alpa's contribution (Zheng et al., OSDI 2022) is essentially a cost model for this choice. A typed language should expose the cost model as part of the compiler's type-directed search.

Effect interpretation. From the memo 2 framing: communication is an effect. Comm is a effect with operations all_reduce, all_gather, etc. A handler interprets these on a particular substrate (NCCL, Gloo, MPI, simulator). The typed sharding determines which communication is needed; the handler determines how it's executed. The compositional handler story extends naturally.

5. Comparison across systems

The pattern visible in this table: the more compiler-driven a system is, the closer it gets to type-system semantics — but none of them complete the journey. GSPMD has the propagation; Alpa has the cost model; JAX has the surface API. The unified design has all three plus integration with the rest of the type system.

The bottom three rows (composing with shape, sparsity, AD) are the integration points where current systems break down. Sharding × sparsity is the production pain point in MoE deployment; sharding × dynamic shape is the LLM-serving pain point; sharding × AD just barely works because everyone treats AD as a separate compiler pass that re-runs sharding propagation on the backward.

6. The auto-parallelization problem

Given a model and a mesh, what is the optimal sharding strategy? This is the open question.

Alpa's framing (Zheng et al., OSDI 2022): split the problem into intra-operator (how to shard one matmul across devices) and inter-operator (how to assign layers to pipeline stages) parallelism. Solve intra-op with integer linear programming; solve inter-op with dynamic programming. This produces sharding plans competitive with hand-tuned Megatron-LM on standard benchmarks.

What Alpa doesn't solve:

• The cost model is approximate, especially for cross-host communication.
• Dynamic-shape workloads are not directly supported.
• Sparsity (MoE, structured sparse) is not in the cost model.
• The search is offline; it can't adapt during training.
• Heterogeneous hardware (CPU/GPU/TPU mixes, different GPU generations) is not modeled.

What a language-level approach could add:

1. Constraint-based search, where the user gives partial annotations (e.g., "the embedding axis is replicated, the batch axis is data-parallel") and the compiler infers the rest. This is like Hindley-Milner type inference where the user can annotate and the rest is inferred.

2. Compositional cost models, where the cost of a complex pipeline is computed from the costs of its parts. The handler-based effect framing helps: each Comm handler has a cost model.

3. Online repartitioning, where the runtime can change the sharding strategy based on observed load (e.g., MoE expert imbalance). The type system records valid shardings; the runtime picks among them.

4. Verification of sharding correctness, in the sense of "the program with this sharding produces the same result as the unsharded program." This is again where memo 5 (verification) becomes relevant — sharding transformations are program transformations whose semantic preservation is non-trivial.

7. The hard composition problems

Where the unified-type-system thesis gets stress-tested.

7.1 Sharding × dynamic shape

If the sequence length is symbolic, can it be sharded along the SP axis? Yes — but only if the shape is a multiple of the mesh axis size, which is a constraint on the symbolic shape variable. This wants to be a refinement-typed property: Tensor f16 [B, S where S % data_axis_size == 0, D].

LLM serving with continuous batching is the canonical hard case. The batch dimension grows and shrinks dynamically; the sharding along it must be valid for the current batch size. Most current systems either pad to a fixed multiple (wasted compute) or skip data parallelism entirely on the dynamic axis.

7.2 Sharding × sparsity

Sharding a sparse tensor across devices: how? Two approaches:

• Shard by index: divide the index space evenly across devices. Simple but produces severe load imbalance if nonzeros are clustered.
• Shard by nonzero count: divide the nonzeros evenly. Better load balance but breaks spatial locality and complicates indexing.

For dynamic sparsity (MoE), neither is satisfactory. Expert parallelism shards experts (which are dense) across devices; the sparsity is in the routing, not in the weights. The all-to-all that dispatches tokens to experts is a giant communication operation that bottlenecks scaling. Nobody has solved load-balanced expert parallelism with dynamic routing.

7.3 Sharding × AD

The backward pass produces gradients with shardings determined by the forward shardings. In simple cases this is automatic (GSPMD handles it). The hard cases:

• ZeRO-3 / FSDP: parameters are gathered just-in-time during forward, the gathered version produces forward activations, the backward needs the gathered version to be available again, and the gradient is reduce-scattered back. The shardings change during execution. This is currently implemented via hooks into the autograd engine.
• Pipeline parallelism with activation checkpointing: the forward pass discards activations, the backward pass recomputes them, and the recomputation must use the same sharding. Mostly works in current systems but is fragile.
• Gradient accumulation across micro-batches with different shardings: requires per-micro-batch reduce-scatter and final all-reduce, with careful synchronization.

The unified-type-system answer: the AD transformation propagates sharding types through the backward pass, just as it propagates shapes and sparsity. The collectives in the backward pass are inferred the same way as in the forward.

7.4 Heterogeneous meshes

Modern training clusters are not uniform. Different GPU generations, CPU offload, optical fabric vs. PCIe, etc. The communication cost between devices is non-uniform. Pure GSPMD-style propagation assumes a uniform mesh; production systems handle heterogeneity with bespoke code.

The cleanest framing: a mesh has not just axes but also cost annotations on its edges. Sharding choice considers both correctness and the cost graph. This is closer to a heterogeneous parallel computing model (cf. Sequoia, Legion) than to BSP.

8. Open research problems

1. A fully typed sharding calculus. Formal type system where sharding is a type, mesh is a kind/parameter, collectives are coercions, and resharding is a coercion. Prove subject reduction (well-typed sharded programs reduce to well-typed sharded programs) and semantic preservation (sharded program = unsharded program).

2. Sharding inference with partial annotations. Like Hindley-Milner, but for sharding. The user annotates a few key tensors; the compiler infers the rest. Currently in JAX you must annotate at every pjit boundary; better inference is possible.

3. Cost-aware sharding inference. Given multiple correct shardings, pick the cheapest. Alpa solves a version of this; doing it as part of the type system is open.

4. Sharding for dynamic shape. What's the right abstraction for "this axis is sharded if it's a multiple of the mesh size, replicated otherwise"? Refinement types, conditional shardings, or something new?

5. Sharding for dynamic sparsity (MoE). Load-balanced expert parallelism is the production-scale open problem. Solutions in MegaBlocks address the kernel level; the distributed MoE story is mostly hand-tuned.

6. Verified sharding transformations. The CompCert-style question: can we prove that a sharded program produces the same result (up to numerical reordering) as the unsharded one? Connects to memo 5.

7. Heterogeneous and hierarchical meshes. Real clusters are NUMA, multi-tier, multi-generation. Sharding should account for this; current systems mostly don't.

8. Online resharding. Training dynamics can favor different shardings at different points (warm-up vs. main training, different MoE load). A language that makes resharding cheap and the type system enforces correctness across resharding could enable adaptive parallelism.

9. Interaction with effects (memo 2). Communication-as-effect composes with other effects (Random, State) in non-trivial ways. The order of with reverseMode handle and with allReduce handle matters. Working this out compositionally is open.

9. Recommendation for an AI-native language design

Concrete positions, in priority order:

• Mesh is a kind / type-level construct. Programs are parameterized over meshes the way generic functions are parameterized over types. f : forall M : Mesh. Tensor f16 [B, D] on M with {B: M.data} -> ....

• Sharding is a refinement on the tensor type, alongside shape and sparsity. The single tensor type carries all four refinement axes. This is the unified-type-system thesis.

• Collectives are type-changing coercions inferred by the compiler. Users write the algorithm; the compiler inserts all_gather, reduce_scatter, all_to_all. Alpa's cost model becomes part of the inference.

• Resharding is explicit in the IR but inferable from context. The user can also write reshard(x, new_spec) to force a particular layout when they know better than the compiler.

• Communication is an effect (Comm), with handlers for NCCL, Gloo, MPI, and simulators. This composes with the AD effect (Diff), the random effect (Random), and the implicit effect (Implicit) from memo 2.

• Pipeline parallelism via continuations. Pipeline schedules (1F1B, interleaved 1F1B, etc.) are control-flow patterns that fit naturally into a CPS-friendly language with effect handlers.

• Cost-model-driven sharding inference, with the user able to annotate partial shardings. The inference is a constraint solve over the type system, not a separate optimizer pass.

• No process-rank concept in the surface language. The mesh is the only explicit notion of distribution; per-device behavior is derived from sharding.

The thesis: the GSPMD compiler is a type inferencer in disguise, and the auto-parallelization problem is constraint solving over the type. A language that puts sharding in the type system from day one absorbs GSPMD, Alpa, and ZeRO into one framework rather than three. The mesh becomes a kind; collectives become coercions; pipeline schedules become continuations; auto-parallelization becomes inference.

10. References

Foundational distributed systems / HPC

• Valiant, L. (1990). A Bridging Model for Parallel Computation. CACM. — BSP, the foundational parallel cost model.
• Culler, D. et al. (1993). LogP: Towards a Realistic Model of Parallel Computation. PPoPP. — More refined cost model.
• Gabriel, E. et al. (2004). Open MPI: Goals, Concept, and Design of a Next Generation MPI Implementation. — The standard message-passing baseline.

Compiler-driven sharding

• Xu, Y. et al. (2021). GSPMD: General and Scalable Parallelization for ML Computation Graphs. arXiv:2105.04663. — Single most important reference for this memo.
• Zheng, L. et al. (2022). Alpa: Automating Inter- and Intra-Operator Parallelism for Distributed Deep Learning. OSDI 2022. — Auto-parallelization with cost models.
• Lepikhin, D. et al. (2020). GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding. ICLR 2021. — GSPMD's predecessor with MoE focus.
• Barham, P. et al. (2022). Pathways: Asynchronous Distributed Dataflow for ML. MLSys 2022. — Google's distributed runtime.

Hand-coded parallelism (the strategies the compiler should subsume)

• Shoeybi, M. et al. (2019). Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism. arXiv:1909.08053. — TP for transformers.
• Huang, Y. et al. (2019). GPipe: Efficient Training of Giant Neural Networks Using Pipeline Parallelism. NeurIPS 2019.
• Rajbhandari, S. et al. (2020). ZeRO: Memory Optimizations Toward Training Trillion Parameter Models. SC 2020.
• Narayanan, D. et al. (2021). Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM. SC 2021. — 3D parallelism in production.
• Korthikanti, V. et al. (2022). Reducing Activation Recomputation in Large Transformer Models. MLSys 2023. — Sequence parallelism.
• Liu, H. et al. (2023). Ring Attention with Blockwise Transformers for Near-Infinite Context. arXiv:2310.01889. — Context parallelism.
• Rasley, J. et al. (2020). DeepSpeed: System Optimizations Enable Training Deep Learning Models with Over 100 Billion Parameters. KDD 2020.

Distributed tensor abstractions

• Bradbury, J. et al. JAX: composable transformations of Python+NumPy programs. — pjit and shard_map.
• PyTorch team. DTensor: PyTorch Distributed Tensor. — RFCs and documentation.
• Zhao, Y. et al. (2023). PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel. VLDB 2023.

Parallel programming languages (research lineage)

• Chamberlain, B., Callahan, D. & Zima, H. (2007). Parallel Programmability and the Chapel Language. IJHPCA. — Place-based parallelism.
• Charles, P. et al. (2005). X10: An Object-Oriented Approach to Non-Uniform Cluster Computing. OOPSLA 2005.
• Fatahalian, K. et al. (2006). Sequoia: Programming the Memory Hierarchy. SC 2006. — Hierarchical memory model.
• Bauer, M. et al. (2012). Legion: Expressing Locality and Independence with Logical Regions. SC 2012. — Region-based programming.
• Frigo, M., Leiserson, C. & Randall, K. (1998). The Implementation of the Cilk-5 Multithreaded Language. PLDI 1998. — Work stealing.

Polyhedral / scheduling foundations (relevant for tiling and distribution)

• Bondhugula, U. et al. (2008). A Practical Automatic Polyhedral Parallelizer and Locality Optimizer. PLDI 2008. — Pluto.
• Verdoolaege, S. (2010). isl: An Integer Set Library for the Polyhedral Model. — Tooling.

Production-scale training reports (concrete configurations)

• Smith, S. et al. (2022). Using DeepSpeed and Megatron to Train Megatron-Turing NLG 530B. arXiv:2201.11990.
• Chowdhery, A. et al. (2022). PaLM: Scaling Language Modeling with Pathways. arXiv:2204.02311. — 6144 TPU v4 chip configuration.
• Touvron, H. et al. (2023). LLaMA / LLaMA 2 papers. — Training infrastructure descriptions.

11. Connections to the rest of the series

• Memo 1 (shapes): dynamic-shape sharding (sequence-length growing during decoding) is the production blocker. Refinement-typed shapes carrying mesh-divisibility constraints are the natural answer.
• Memo 2 (AD-through-effects): communication is an effect; the AD transformation must propagate sharding through the backward. The handler composition story extends naturally.
• Memo 3 (sparsity): sharding × sparsity is the MoE-at-scale problem and is open even with hand-coding. Expert parallelism is a hack on dense expert weights; the underlying compositional question (how does dynamic sparsity interact with sharding?) is unsolved.
• Memo 5 (verification): verified sharding transformations — proving that a sharded program is equivalent to the unsharded one — is a credible CompCert-style research target.
• Memo 6 (hardware abstraction): heterogeneous meshes, NUMA-aware sharding, and tile-level kernel placement are where this memo and memo 6 meet.

12. Cross-memo synthesis

Four memos in, the unified-type-system thesis can be stated more precisely:

The "tensor" of an AI-native language is a type with several orthogonal refinement axes:

1. Shape — dimension sizes, possibly symbolic, possibly with constraints (memo 1).
2. Effects — purity, mutation, randomness, divergence, the things AD needs to know (memo 2).
3. Structure — sparsity patterns, low-rank factorizations, quantization, hardware layouts (memo 3).
4. Sharding — distribution across a typed mesh (this memo).

Operations on tensors compose along all four axes. The compiler infers the result's refinements from the operands. Mismatches at operation boundaries are either type errors or compiler-inserted coercions (broadcast for shape, resharding for sharding, format conversion for sparsity, AD-handler insertion for effects).

The user writes algorithms; the compiler manages the four refinement axes. This is the program model that current systems implement piecemeal.

Memos 5 (verification) and 6 (hardware abstraction) will test whether this framework holds up under the verification load and the hardware-portability load respectively. If it does, the synthesis is the design document for an actual language.

13. Next memo

Design Space 05: Verification. ML compilers do enormous program transformations: AD, sharding propagation, kernel fusion, layout selection, quantization. Each transformation must preserve semantics, and most do not have proofs. A CompCert-for-ML, or a verified-AD, or a verified-sharding-pass, is a real research target — and the most concretely achievable contribution in this whole design space, given existing verification infrastructure (Coq, Lean 4, F*, Frama-C). This is the memo most directly aligned with formal-methods backgrounds.