Meta open-sources spmd_types: a type system to verify distributed PyTorch training correctness

austinvhuang · x · 2026-09-19

meta-pytorch/spmdtypes, inspired by JAX's sharding-in-types and adapted for PyTorch, offers two type systems: Local SPMD types for safely using Megatron-style differentiable collectives by tracking pending gradient reductions, and Global SPMD types, a DTensor-like abstraction with explicit communication ops so code behaves identically on one device or many. It lets you statically check gradient correctness or parallelization equivalence without running a full E2E distributed training job.

Original post →

More from Infra

Infra channel →