ttnn.all_reduce

ttnn.all_reduce(input_tensor: ttnn.Tensor, *, cluster_axis: int = None, subdevice_id: ttnn.SubDeviceId | None, memory_config: ttnn.MemoryConfig | None, num_links: int = None, topology: ttnn.Topology = None) → ttnn.Tensor

All-reduce operation across devices with Sum reduction. If cluster axis is specified, the all-reduce is performed on tensor shards along the cluster axis, resulting in identical tensor shards across all devices along the cluster axis. If it is not specified, then we reduce across all devices in the mesh. All-reduce is a collective operation that reduces data from all devices using the Sum operation and returns the result to all devices.

Parameters:

input_tensor (ttnn.Tensor) – Input tensor to be reduced.

Keyword Arguments:

cluster_axis (int, optional) – The axis on the mesh device to reduce across. Defaults to None.
subdevice_id (ttnn.SubDeviceId, optional) – Subdevice id for worker cores.
memory_config (ttnn.MemoryConfig, optional) – Output memory configuration.
num_links (int, optional) – Number of links to use for the all_reduce operation. Defaults to None.
topology (ttnn.Topology, optional) – Fabric topology. Defaults to None.

Returns:

ttnn.Tensor – The reduced tensor with the same shape as the input tensor. The output tensor is identical across all devices along the cluster axis if specified, otherwise it is identical across all devices in the mesh.

Supported dtypes and layouts:

Dtypes

Layouts

BFLOAT16, BFLOAT8_B, FLOAT32

TILE, ROW_MAJOR

The reduction is performed in BFLOAT16; BFLOAT8_B inputs are up-cast to BFLOAT16 for the reduction and cast back, and FLOAT32 takes a dedicated reduction path. Input must be rank 2 or greater. The output preserves the input dtype and shape.

Memory Support:

Interleaved: DRAM and L1
Sharded: supported in DRAM or L1. A sharded input is converted to interleaved internally (its buffer type is preserved), and the result is written back according to the requested output memory_config.

Example

torch_input = torch.randn([1, 1, 32, 256], dtype=torch.bfloat16)
tt_input = ttnn.from_torch(
    torch_input,
    dtype=ttnn.bfloat16,
    layout=ttnn.TILE_LAYOUT,
    device=mesh_device,
    mesh_mapper=ttnn.ReplicateTensorToMesh(mesh_device),
)

# Sum-reduce across devices; result is replicated and shape is preserved.
output = ttnn.all_reduce(tt_input, cluster_axis=1)
logger.info(output.shape)  # [1, 1, 32, 256]