AutoStrategy

class braintools.trainer.AutoStrategy(prefer_fsdp=False)

Automatically select the best strategy based on available devices.

This strategy inspects the available hardware and selects an appropriate distribution strategy.

Parameters:

prefer_fsdp (bool) – Whether to prefer FSDP over simple data parallelism when multiple devices are available.

Examples

>>> strategy = AutoStrategy()
>>> trainer = Trainer(strategy=strategy)

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

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='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

property selected_strategy: Strategy

The automatically selected strategy.

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]

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]