Mesh

substrax.mesh builds device meshes and describes how arrays are laid out on them.

from substrax.mesh import DeviceMeshManager, FSDPStrategy, ShardingConfig, fsdp_rules

mesh = DeviceMeshManager.create_device_mesh({"data": 2, "model": 4})
rules = fsdp_rules(data_axis="data", model_axis="model")
config = ShardingConfig(data_parallel_size=2, tensor_parallel_size=4, fsdp_enabled=True)
strategy = FSDPStrategy("model", mesh_axis=1, min_weight_size=config.fsdp_min_weight_size)

Three layers, from low to high:

Layer	Names	Role
Meshes	DeviceMeshManager	jax.make_mesh over the visible devices, plus mesh introspection
Rules	MeshRules, data_parallel_rules, fsdp_rules, create_named_sharding, partition_spec_for_names	Map logical axis names to mesh axes and build NamedShardings
Strategies	ShardingConfig, ParallelismConfig, ShardingStrategy and its data-parallel, FSDP, tensor-parallel, pipeline-parallel and multi-dimensional implementations	Decide the partition spec of every parameter in a model

substrax.mesh

Device meshes, mesh rules, partition-spec helpers and sharding strategies.

DeviceMeshManager

Manager for creating and configuring JAX device meshes.

This class provides utilities for creating device meshes for different distributed training configurations, including data-parallel, model-parallel, and hybrid approaches.

create_device_mesh(mesh_shape, devices=None, *, axis_types=None) staticmethod

Create a JAX device mesh with the specified shape.

Every axis is AxisType.Auto unless axis_types says otherwise: the data-parallel helpers in :mod:substrax.spmd leave sharding inference to XLA (the gradient all-reduce of a batch-sharded step is inferred from the input sharding). jax.make_mesh defaults to Explicit axes from jax 0.11, under which a contraction over a sharded batch axis, the backward pass of any layer, raises unless every such op names its out_sharding.

Parameters:

Name	Type	Description	Default
mesh_shape	dict[str, int] | list[tuple[str, int]]	The shape of the mesh, specified either as a dictionary mapping axis names to sizes, or as a list of (name, size) tuples.	required
devices	Sequence[DeviceLike] | None	The devices to use, such as a slice of jax.devices(). If None, uses all available devices.	None
axis_types	Sequence[AxisType] | None	One AxisType per axis; Auto for every axis when None.	None

Returns:

Type	Description
Mesh	A JAX device mesh.

Raises:

Type	Description
ValueError	If the mesh shape is incompatible with the number of devices, or axis_types does not name one type per axis.

create_data_parallel_mesh(num_devices=None) staticmethod

Create a data-parallel device mesh.

Parameters:

Name	Type	Description	Default
num_devices	int | None	Optional number of devices to use. If None, uses all available devices.	None

Returns:

Type	Description
Mesh	A JAX device mesh configured for data-parallel training.

create_model_parallel_mesh(num_devices) staticmethod

Create a model-parallel device mesh.

Parameters:

Name	Type	Description	Default
num_devices	int	Number of devices to use for model parallelism.	required

Returns:

Type	Description
Mesh	A JAX device mesh configured for model-parallel training.

Raises:

Type	Description
ValueError	If there aren't enough devices available.

create_hybrid_mesh(data_parallel_size, model_parallel_size) staticmethod

Create a hybrid data-parallel and model-parallel device mesh.

Parameters:

Name	Type	Description	Default
data_parallel_size	int	Number of devices to use for data parallelism.	required
model_parallel_size	int	Number of devices to use for model parallelism.	required

Returns:

Type	Description
Mesh	A JAX device mesh configured for hybrid parallel training.

Raises:

Type	Description
ValueError	If there aren't enough devices available.

get_mesh_info(mesh) staticmethod

Get information about a device mesh.

Parameters:

Name	Type	Description	Default
mesh	Mesh	The device mesh to inspect.	required

Returns:

Type	Description
dict[str, int | dict[str, int]]	A dictionary containing information about the mesh, including the
dict[str, int | dict[str, int]]	number of devices and the size of each axis.

MeshRules(*, data=None, embed=None, mlp=None, heads=None) dataclass

Logical-to-physical mesh axis mapping for SPMD sharding.

Maps logical dimension names to physical mesh axis names. Following the pattern from Flax NNX FSDP examples.

Attributes:

Name	Type	Description
data	str | None	Mesh axis for the batch/data dimension.
embed	str | None	Mesh axis for embedding dimensions.
mlp	str | None	Mesh axis for MLP hidden dimensions.
heads	str | None	Mesh axis for attention head dimensions.

__call__(*keys)

Look up mesh axes for the given logical dimension names.

Parameters:

Name	Type	Description	Default
*keys	str	Logical dimension names (must be attributes of MeshRules).	()

Returns:

Type	Description
tuple[str | None, ...]	Tuple of mesh axis names (or None for unmapped dimensions).

DataParallelStrategy(axis_name, mesh_axis)

Bases: ShardingStrategy

Data parallel sharding strategy.

Shards the batch dimension across devices while replicating model parameters and computation.

get_partition_spec(tensor_shape)

Get partition spec for data parallel sharding.

Only shards the batch dimension, leaves others replicated.

apply_sharding(array, mesh)

Apply data parallel sharding to array.

FSDPStrategy(axis_name, mesh_axis, min_weight_size=1024)

Bases: ShardingStrategy

Fully Sharded Data Parallel strategy.

Shards model parameters across devices to reduce memory usage while maintaining training efficiency.

Initialize FSDP strategy.

Parameters:

Name	Type	Description	Default
axis_name	str	Name of the mesh axis	required
mesh_axis	int	Index of the mesh axis	required
min_weight_size	int	Minimum first dimension size to enable sharding	1024

should_shard_weight(weight)

Determine if a weight should be sharded based on its size.

Parameters:

Name	Type	Description	Default
weight	Array	Weight array to check	required

Returns:

Type	Description
bool	True if weight should be sharded, False otherwise

get_partition_spec(tensor_shape)

Get partition spec for FSDP sharding.

Shards along the first dimension of weight tensors.

get_gradient_partition_spec(tensor_shape)

Get partition spec for gradient sharding (same as weights).

apply_sharding(array, mesh)

Apply FSDP sharding to array.

MultiDimensionalStrategy(strategies, config)

Multi-dimensional parallelism strategy combining multiple approaches.

Combines different sharding strategies (data, tensor, FSDP, pipeline) to achieve optimal performance for large-scale training.

Initialize multi-dimensional strategy.

Parameters:

Name	Type	Description	Default
strategies	Mapping[str, ShardingStrategy]	Dictionary mapping strategy names to strategy instances	required
config	ParallelismConfig	Sharding configuration for the multi-dimensional strategy	required

get_combined_partition_spec(tensor_name, tensor_shape)

Get combined partition spec from all strategies.

Parameters:

Name	Type	Description	Default
tensor_name	str	Name/type of the tensor	required
tensor_shape	tuple[str, ...]	Shape dimension names of the tensor	required

Returns:

Type	Description
PartitionSpec	Combined PartitionSpec

resolve_sharding_conflicts(tensor_name, proposed_specs)

Resolve conflicts between multiple proposed partition specs.

Parameters:

Name	Type	Description	Default
tensor_name	str	Name of the tensor	required
proposed_specs	dict[str, PartitionSpec]	Dictionary of strategy names to proposed specs	required

Returns:

Type	Description
PartitionSpec	Resolved PartitionSpec

ParallelismConfig(mesh_shape, mesh_axis_names, sharding_config) dataclass

Complete parallelism configuration including mesh topology.

Combines sharding configuration with device mesh setup.

is_valid()

Validate that mesh shape matches sharding configuration.

from_sharding_config(config) classmethod

Create parallelism config from sharding configuration.

PipelineParallelStrategy(axis_name, mesh_axis, num_stages)

Bases: ShardingStrategy

Pipeline parallel sharding strategy.

Distributes model layers across devices to enable pipeline parallelism for very large models that don't fit on single devices.

Initialize pipeline parallel strategy.

Parameters:

Name	Type	Description	Default
axis_name	str	Name of the mesh axis	required
mesh_axis	int	Index of the mesh axis	required
num_stages	int	Number of pipeline stages	required

assign_layers_to_stages(num_layers)

Assign layers to pipeline stages.

Parameters:

Name	Type	Description	Default
num_layers	int	Total number of layers in the model	required

Returns:

Type	Description
list[int]	list of layer counts per stage

get_partition_spec(tensor_shape)

Get partition spec for pipeline parallel sharding.

Pipeline parallelism doesn't shard individual tensors, but rather assigns entire layers to different devices.

get_forward_communication_pattern()

Get communication pattern for forward pass.

Returns:

Type	Description
list[tuple[int, int]]	list of (source_stage, dest_stage) pairs

get_backward_communication_pattern()

Get communication pattern for backward pass.

Returns:

Type	Description
list[tuple[int, int]]	list of (source_stage, dest_stage) pairs

apply_sharding(array, mesh)

Apply pipeline parallel sharding to array.

Pipeline parallelism handles layer assignment rather than tensor sharding.

ShardingConfig(data_parallel_size=1, tensor_parallel_size=1, pipeline_parallel_size=1, fsdp_enabled=False, fsdp_min_weight_size=1024) dataclass

Configuration for multi-dimensional parallelism setup.

Defines the parallelism dimensions and FSDP settings for a model.

get_total_device_count()

Calculate total devices needed for this configuration.

from_device_count(device_count) classmethod

Create optimal sharding config for given device count.

Uses heuristics to balance different parallelism dimensions.

ShardingStrategy(axis_name, mesh_axis)

Bases: ABC

Abstract base class for sharding strategies.

Defines the interface that all sharding strategies must implement for consistent handling of different parallelism types.

Initialize sharding strategy.

Parameters:

Name	Type	Description	Default
axis_name	str	Name of the mesh axis for this strategy	required
mesh_axis	int	Index of the mesh axis	required

get_partition_spec(tensor_shape) abstractmethod

Get partition specification for a tensor with given shape names.

Parameters:

Name	Type	Description	Default
tensor_shape	tuple[str, ...]	Tuple of dimension names for the tensor	required

Returns:

Type	Description
PartitionSpec	PartitionSpec defining how to shard the tensor

apply_sharding(array, mesh) abstractmethod

Apply sharding to an array using the given mesh.

Parameters:

Name	Type	Description	Default
array	Array	JAX array to shard	required
mesh	Mesh	Device mesh for sharding	required

Returns:

Type	Description
Array	Sharded array

TensorParallelStrategy(axis_name, mesh_axis, shard_dimension=None)

Bases: ShardingStrategy

Tensor parallel sharding strategy.

Shards model computation across devices by splitting tensors along specific dimensions (typically features).

Initialize tensor parallel strategy.

Parameters:

Name	Type	Description	Default
axis_name	str	Name of the mesh axis	required
mesh_axis	int	Index of the mesh axis	required
shard_dimension	str | None	Preferred dimension to shard ('in_features' or 'out_features')	None

get_partition_spec(tensor_shape)

Get partition spec for tensor parallel sharding.

get_linear_weight_spec()

Get partition spec for linear layer weights.

get_attention_qkv_spec()

Get partition spec for attention QKV projections.

get_attention_output_spec()

Get partition spec for attention output projection.

apply_sharding(array, mesh)

Apply tensor parallel sharding to array.

create_named_sharding(mesh, *axis_names)

Create a NamedSharding from a mesh and axis names.

Parameters:

Name	Type	Description	Default
mesh	Mesh	The device mesh.	required
*axis_names	str | None	Axis names for each dimension. Use None for replicated dimensions. If no axes are provided, creates a fully replicated sharding.	()

Returns:

Type	Description
NamedSharding	A JAX NamedSharding object.

data_parallel_rules(data_axis='data')

Create MeshRules for pure data-parallel training.

Shards only the batch dimension across devices.

Parameters:

Name	Type	Description	Default
data_axis	str	Name of the mesh axis for data parallelism.	'data'

Returns:

Type	Description
MeshRules	MeshRules with only the data axis mapped.

fsdp_rules(data_axis='data', model_axis='model')

Create MeshRules for Fully Sharded Data Parallel (FSDP) training.

Shards the batch dimension across the data axis and model weight dimensions across the model axis.

Parameters:

Name	Type	Description	Default
data_axis	str	Name of the mesh axis for data parallelism.	'data'
model_axis	str	Name of the mesh axis for model parallelism.	'model'

Returns:

Type	Description
MeshRules	MeshRules with data and model axes mapped.

partition_spec_for_names(rules, *logical_names)

Apply MeshRules to create a PartitionSpec for a tensor.

Maps logical dimension names through the rules to produce physical mesh axis assignments.

Parameters:

Name	Type	Description	Default
rules	MeshRules	The MeshRules mapping to apply.	required
*logical_names	str	Logical dimension names for each tensor axis.	()

Returns:

Type	Description
PartitionSpec	A PartitionSpec with physical mesh axis assignments.