Device strategies

Accelerator and DeviceSelection describe the values accepted by Trainer and create_strategy().

SingleDeviceStrategy places complete state and batches on one selected CPU, GPU, or TPU device. DataParallelStrategy replicates state and shards non-scalar batch leaves along their leading dimension. Batch sizes must be divisible by the number of local shards.

Experimental distributed execution

DataParallelStrategy provides experimental synchronous data parallelism. It does not provide full Lightning DDP parity. Call initialize_distributed() before any JAX backend operation for multi-process use. Every process must run the same compiled steps and collectives in the same order. Logging and checkpoint writes remain restricted to global rank zero.

Ordinary CI covers single-process CPU behavior. Validate the intended multi-device or multi-host environment before a long run.

Strategy [source]

Strategy(devices: Sequence[JaxDevice])

Define device placement operations used by the host-side trainer.

Attributes:

Name	Type	Description
process_index		Global JAX process index.
is_global_zero		Whether this process owns rank-zero logging side effects.

Initialize a strategy over explicit JAX devices.

Parameters:

Name	Type	Description	Default
devices	Sequence[JaxDevice]	Non-empty ordered device collection.	required

Raises:

Type	Description
ValueError	If devices is empty.

root_device property

root_device: JaxDevice

Return the first strategy device addressable by this process.

The root device is used to prepare persistent sampler state and explicit PRNG keys. Final batches still pass through :meth:place_batch, which may shard them over several devices.

Returns:

Type	Description
JaxDevice	First process-local strategy device.

Raises:

Type	Description
RuntimeError	If the strategy contains no device addressable by this process.

place_state [source]

place_state(state: Any) -> Any

Place a functional training-state PyTree.

Parameters:

Name	Type	Description	Default
state	Any	Training-state PyTree.	required

Returns:

Type	Description
Any	Device-placed training state.

place_batch [source]

place_batch(batch: Any) -> Any

Place a training-batch PyTree.

Parameters:

Name	Type	Description	Default
batch	Any	Training-batch PyTree.	required

Returns:

Type	Description
Any	Device-placed training batch.

SingleDeviceStrategy [source]

SingleDeviceStrategy(device: JaxDevice)

Bases: Strategy

Place state and batches on one explicit CPU, GPU, or TPU device.

Initialize single-device placement.

Parameters:

Name	Type	Description	Default
device	JaxDevice	JAX device used by the compiled training step.	required

place_state [source]

place_state(state: Any) -> Any

Place the training state on the configured device.

Parameters:

Name	Type	Description	Default
state	Any	Training-state PyTree.	required

Returns:

Type	Description
Any	State placed on device.

place_batch [source]

place_batch(batch: Any) -> Any

Place the batch on the configured device.

Parameters:

Name	Type	Description	Default
batch	Any	Training-batch PyTree.	required

Returns:

Type	Description
Any	Batch placed on device.

DataParallelStrategy [source]

DataParallelStrategy(
    devices: Sequence[JaxDevice], axis_name: str = "data"
)

Bases: Strategy

Replicate state and shard batch-leading axes over a named data mesh.

The mesh may span several processes after :func:initialize_distributed has been called. Every process must run the same fit loop and encounter compiled collectives in the same order.

Attributes:

Name	Type	Description
mesh		One-dimensional global device mesh named data by default.

Initialize data-parallel placement.

Parameters:

Name	Type	Description	Default
devices	Sequence[JaxDevice]	Global devices forming the data-parallel mesh.	required
axis_name	str	Logical mesh-axis name.	'data'

Raises:

Type	Description
ValueError	If fewer than two devices are supplied.

place_state [source]

place_state(state: Any) -> Any

Replicate every training-state array over the mesh.

Parameters:

Name	Type	Description	Default
state	Any	Training-state PyTree.	required

Returns:

Type	Description
Any	Replicated global training state.

place_batch [source]

place_batch(batch: Any) -> Any

Shard each non-scalar batch leaf along its leading dimension.

Parameters:

Name	Type	Description	Default
batch	Any	Training-batch PyTree whose non-scalar leading dimensions are divisible by the mesh size.	required

Returns:

Type	Description
Any	Global batch with leading axes sharded over the data mesh.

Raises:

Type	Description
ValueError	If a non-scalar leading dimension is not divisible by the device count.

create_strategy [source]

create_strategy(
    accelerator: Accelerator = "auto",
    devices: DeviceSelection = 1,
) -> Strategy

Create a single-device or data-parallel strategy from visible JAX devices.

Parameters:

Name	Type	Description	Default
accelerator	Accelerator	One of auto, cpu, gpu, or tpu.	'auto'
devices	DeviceSelection	Device count, explicit backend-device indices, or auto for every matching device.	1

Returns:

Type	Description
Strategy	Strategy over the selected devices.

Raises:

Type	Description
ValueError	If the accelerator, count, or device indices are invalid.
RuntimeError	If the requested accelerator has no visible devices.

initialize_distributed [source]

initialize_distributed(
    coordinator_address: str | None = None,
    num_processes: int | None = None,
    process_id: int | None = None,
    local_device_ids: int | Sequence[int] | None = None,
) -> None

Initialize JAX multi-process execution before any backend operation.

Parameters:

Name	Type	Description	Default
coordinator_address	str | None	Coordinator host:port, or None for supported environment auto-detection.	None
num_processes	int | None	Total number of participating JAX processes.	None
process_id	int | None	Unique zero-based identifier for this process.	None
local_device_ids	int | Sequence[int] | None	Process-local device indices made visible to this process.	None

Raises:

Type	Description
RuntimeError	If JAX distributed execution has already been initialized.