---
id: spmd-types
version: "0.2.3"
license: BSD 3-Clause License (c) Meta Platforms, Inc. and affiliates. Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are… (full text in the JSON record)
license_treatment: permissive
maintenance: active
---

# spmd-types — A type system for distributed (SPMD) tensor computations in PyTorch

License: permissive · Maintenance: active · Downloads: 111.7K/mo

## What it is and what it does

spmd_types is a type system for writing distributed tensor computations with explicit, verifiable semantics. It provides two complementary abstractions: local SPMD types track whether gradients are pending reduction or fully synchronized, enabling safe use of Megatron-style differentiable collectives; global SPMD types offer a DTensor-like interface where code has identical semantics whether run on a single device or distributed, but with explicit communication operations so redistributes are never implicit.

The package lets you verify correctness of distributed training logic—checking gradient computation and parallelization equivalence—without running full end-to-end distributed training. It includes runtime type checking via a checker module, allowing you to annotate tensors with their layout and synchronization state and assert those invariants hold at each step.

Use it for:

- Porting training frameworks to verify gradients are correctly reduced across ranks without E2E training runs.
- Typechecking distributed code to catch mismatches between expected and actual tensor layouts before deployment.
- Writing distributed logic that is semantically identical whether run on a single device or across multiple ranks.
- Verifying that collective operations are applied to tensors in the correct synchronization state.
- Debugging gradient flow by asserting tensor state at each computation step.

## Worth the install?

AI-flagged interpretation of the facts above — verify before relying.

Provides type annotations and runtime checking for distributed tensor computations, tracking data layout and synchronization state across SPMD parallelism patterns.

Yes, if you are building or maintaining distributed training code. The package is actively maintained, has no external dependencies, and offers a concrete way to catch synchronization bugs before expensive distributed runs. It is in beta (global SPMD types are under construction), so expect the API to evolve; use it for new projects where you can adapt to changes.

## Install

pip install spmd-types
uv add spmd-types
poetry add spmd-types

## Installing spmd-types

Before you install: Installs with no runtime dependencies and low friction. The package is actively maintained with a recent release 8 days old; the repository shows steady activity.

License in practice: BSD 3-Clause License permits commercial and private use with attribution and liability disclaimer. No restrictions on modification or redistribution.

Quickstart:

pip install spmd_types

import spmd_types as spmd
import spmd_types.checker

with spmd.set_current_mesh(mesh), spmd.checker.typecheck():
    spmd.assert_type(x, {dp: spmd.R, tp: spmd.P})
    y = spmd.all_reduce(x, tp, src=spmd.P, dst=spmd.R)
    spmd.assert_type(y, {dp: spmd.R, tp: spmd.R})

Requires Python 3.10 or later; integration with distributed tensor frameworks assumed in usage context.

Verify before relying:

- Whether the type system catches all realistic gradient-reduction bugs in production training code.
- Performance overhead of runtime type checking during training.
- Maturity and completeness of global SPMD types (noted as 'actively under construction').
- Compatibility with frameworks beyond those mentioned in documentation.

## Package facts

- License: BSD 3-Clause License (c) Meta Platforms, Inc. and affiliates. Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are… (full text in the JSON record) (permissive)
- Python support: supports_current
- Install friction: low
- Maintenance: active
- Downloads: 111.7K/month (top 15,000 on PyPI)
- Known vulnerabilities: none known

## Tags

distributed tensor types, SPMD type checking, collective operations typing, distributed training type safety, gradient reduction tracking, parallel tensor layout verification, megatron-style collectives, distributed-training, type-checking

[View on SkillFed](https://skillfed.io/packages/spmd-types) · [View on PyPI](https://pypi.org/project/spmd-types/)