Strategy

class braintools.trainer.Strategy

Abstract base class for distributed training strategies.

Strategies define how model parameters and computations are distributed across devices.

barrier()

Synchronization barrier across devices.

broadcast(tensor, src=0)

Broadcast tensor from source to all devices.

Parameters:

• tensor (Any) – Tensor to broadcast.

• src (int) – Source device index.

Returns:

Broadcasted tensor.

Return type:

Any

abstract property devices: List[Any]

List of devices used by this strategy.

property is_distributed: bool

Whether this is a distributed strategy.

abstract property name: str

Strategy name.

abstract property num_devices: int

Number of devices used by this strategy.

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

Reduce tensor across devices.

Parameters:

• tensor (Any) – Tensor to reduce.

• op (str) – Reduction operation (‘mean’, ‘sum’, ‘min’, ‘max’).

• axis_name (str) – Name of the axis for reduction.

Returns:

Reduced tensor.

Return type:

Any

abstractmethod setup(model, optimizer)

Set up the model and optimizer for distributed training.

Parameters:

• model (Any) – The model to distribute.

• optimizer (Any) – The optimizer to distribute.

Returns:

The distributed model and optimizer.

Return type:

Tuple[Any, Any]

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

Execute a single training step.

Parameters:

• model (Any) – The model.

• optimizer (Any) – The optimizer.

• batch (Any) – The input batch.

• loss_fn (Callable) – Loss function to compute.

• param_states (PyTree) – Parameter states for gradient computation.

Returns:

Loss value and metrics.

Return type:

Tuple[Any, PyTree]