Phase 3, Section 3: GSPMD Propagation

Phase 3: Sharding (the distribution axis). Section 3 of 4: GSPMD Propagation.

1. The Inference Problem

In a distributed tensor program, manually annotating the sharding layout of every intermediate tensor is unergonomic and error-prone. The compiler should infer optimal sharding layouts for intermediate operations given only sparse constraints (e.g., the input weights are replicated, but the data batches are sharded).

We adapt the GSPMD (General and Scalable Parallelization for ML Computation Graphs) propagation rules into a formalized bidirectional type-inference pass.

2. Propagation Rules

The inference engine runs over the typed IR and propagates sharding annotations forwards and backwards through the dataflow graph.

2.1 Elementwise Operations (Add, Mul, etc.)

Elementwise operations require both operands and the result to share the exact same sharding layout.

Rule:

• If A is @sharded(axis=0) and B is unannotated, B is inferred as @sharded(axis=0).
• If A is @sharded(axis=0) and B is @replicated, this is a conflict, resolved by coercions (Section 4).

2.2 Contractions (Matmul)

Matrix multiplication (C = A @ B) imposes specific structural constraints depending on the contracted and batch axes.

Rules:

1. Batch Parallelism: If A is sharded on a batch (non-contracting) axis, C is sharded on that same axis. B must be replicated or sharded on a compatible batch axis.
2. Weight Parallelism: If B is sharded on its non-contracting output axis, C is sharded on that same axis. A must be replicated.
3. Inner Product (Contracting Axis) Parallelism: If A and B are sharded on their shared contracting axis, they generate partial sums. The output C is technically a "partial" tensor that requires an all_reduce to become @replicated.

2.3 Reductions

For an operation reduce(sum, axis=d):

• If the input is sharded on an axis other than d, the output inherits that sharding.
• If the input is sharded on d, the reduction is partial, requiring an all_reduce collective to yield a @replicated output.

3. Conflict Resolution

When the type-checker encounters a node where forward and backward propagation derive incompatible sharding layouts, it inserts a Type-Changing Coercion (see Section 4) to bridge the gap.

The inference pass assigns a cost to each coercion (e.g., all_gather is more expensive than a no-op). The goal of the propagation engine is to minimize the total network communication cost of the inserted coercions.