ShardedDataParallelStrategy

class braintools.trainer.ShardedDataParallelStrategy(mesh=None, data_axis='data')

Sharded data parallel strategy using jax.sharding.

Uses modern JAX sharding APIs for more flexible data distribution.

Parameters:

• mesh (Mesh | None) – Device mesh for sharding. Default creates a 1D mesh.

• data_axis (str) – Name of the data parallel axis in the mesh.

Examples

>>> strategy = ShardedDataParallelStrategy()
>>> trainer = Trainer(strategy=strategy)

property devices: List[Any]

List of devices used by this strategy.

property mesh: Mesh

The device mesh.

property name: str

Strategy name.

property num_devices: int

Number of devices used by this strategy.

setup(model, optimizer)

Set up sharded training.

Return type:

Tuple[Any, Any]

training_step(model, optimizer, batch, loss_fn, param_states)

Sharded training step.

Return type:

Tuple[Any, PyTree]