Skip to main content

Python class

Allreduce

Allreduce

class max.nn.Allreduce(num_accelerators)

source

Bases: Module

Layer to perform allreduce operation with automatic implementation selection.

Automatically chooses between peer-to-peer optimized allreduce and naive device-to-device transfer based on accelerator connectivity.

Parameters:

num_accelerators (int) – Number of accelerators participating in the allreduce operation

Initialize the Allreduce layer with a specified number of accelerators.

Parameters:

num_accelerators (int) – Number of accelerators to use for allreduce

Raises:

ValueError – If num_accelerators is less than 1

devices

devices: list[Accelerator]

source

List of accelerators involved in the allreduce operation.