Skip to main content

Python class

Signals

Signals

class max.nn.Signals(devices)

source

Bases: object

Signal buffers used for peer-to-peer communication in allreduce.

Device code uses these buffers by enabling peer-to-peer access. Then thread blocks use the buffers to implement barriers for synchronization, and to hold intermediate communication results.

Args: num_gpus: Number of GPUs involved in the allreduce.

Parameters:

devices (list[DeviceRef])

NUM_BYTES

NUM_BYTES = 537919488

source

The size of the signal buffers used for communication in allreduce.

buffers()

buffers()

source

Allocates and returns buffers used for communication in allreduce.

Enables peer-to-peer access between all GPUs (idempotent) and synchronizes so that buffers are ready for use when this method returns.

Return type:

list[Buffer]

devices

devices: list[DeviceRef]

source

List of graph devices that these signals communicate between.

input_types()

input_types()

source

Gets graph input types corresponding to these signal buffers.

Return type:

list[BufferType]