FullyShardedDataParallelStrategy

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

Fully Sharded Data Parallel (FSDP) strategy.

Shards both model parameters and gradients across devices, similar to PyTorch’s FSDP.

Parameters:

• mesh (Mesh | None) – Device mesh for sharding.

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

• model_axis (str | None) – Name of the model parallel axis (for tensor parallelism).

Examples

>>> strategy = FullyShardedDataParallelStrategy()
>>> trainer = Trainer(strategy=strategy)

property devices: List[Any]

List of devices used by this strategy.

property name: str

Strategy name.

property num_devices: int

Number of devices used by this strategy.

setup(model, optimizer)

Set up FSDP by sharding model parameters.

Return type:

Tuple[Any, Any]

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

FSDP training step with parameter gathering and gradient scattering.

Return type:

Tuple[Any, PyTree]