DataParallelStrategy

class braintools.trainer.DataParallelStrategy(devices=None, axis_name='batch')

Data parallel strategy using pmap.

Replicates the model across all devices and splits data batches. Gradients are synchronized using all-reduce.

Parameters:

• devices (List[Any] | None) – List of devices to use. Default is all available devices.

• axis_name (str) – Name of the batch axis for collective operations.

Examples

>>> strategy = DataParallelStrategy()
>>> trainer = Trainer(strategy=strategy)

broadcast(tensor, src=0)

Broadcast tensor from the source replica to all replicas.

Implemented via all_gather + index rather than ppermute: a broadcast is one-source-to-many-destinations, which ppermute cannot express (it requires unique sources and destinations). (T-D)

Return type:

Any

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.

reduce(tensor, op='mean', axis_name=None)

Reduce tensor across devices.

Return type:

Any

setup(model, optimizer)

Replicate model parameters across devices.

Return type:

Tuple[Any, Any]

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

Parallel training step.

The batch should already be sharded with shape (num_devices, batch_size, …).

Return type:

Tuple[Any, PyTree]