rlightning.utils.distributed.collective

Collective communication operations for distributed training.

This module provides wrappers around PyTorch distributed collective operations that work with the CommContext singleton for group management. Includes scatter, broadcast, gather, all_reduce, all_gather, and sequence parallel communication utilities.

rlightning.utils.distributed.collective.all_gather(output_tensor: MockModule('torch.Tensor'), input_tensor: MockModule('torch.Tensor'), comm_mode: CommMode, async_op: bool = False) → Any | None

Single tensor all gather. Gathers a single tensor from all ranks, and puts them in a single output tensor.

Parameters:

• output_tensor (Tensor) – Output tensor. It should contain correctly-sized tensors to be used for output of the collective.

• input_tensor (Tensor) – Tensor to be broadcast from current process.

• comm_mode (CommMode) – Communication mode registered in CommContext.

• async_op (bool, optional) – Whether this op should be an async op

Returns:

Async work handle, if async_op is set to True. None, if not async_op or if not part of the group

rlightning.utils.distributed.collective.all_reduce(tensor: MockModule('torch.Tensor'), comm_mode: CommMode, op: MockModule('torch.distributed.ReduceOp') = MockModule('torch.distributed.ReduceOp.SUM'), async_op: bool = False) → Any | None

Reduces the tensor data across all machines in such a way that all get the final result.

After the call tensor is going to be bitwise identical in all processes.

Complex tensors are supported.

Parameters:

• tensor (Tensor) – Input and output of the collective. The function operates in-place.

• comm_mode (CommMode) – Communication mode registered in CommContext.

• op (optional) – One of the values from torch.distributed.ReduceOp enum. Specifies an operation used for element-wise reductions.

• async_op (bool, optional) – Whether this op should be an async op.

Returns:

Async work handle, if async_op is set to True. None, if not async_op or if not part of the group

rlightning.utils.distributed.collective.all_reduce_dict(dictionary: dict, comm_mode: CommMode, op=MockModule('torch.distributed.ReduceOp.SUM'), dtype=MockModule('torch.float32'))

Reduces the dictionary data across all machines in such a way that all get the final result.

rlightning.utils.distributed.collective.all_to_all(input_: MockModule('torch.Tensor'), gather_dim: int, scatter_dim: int, group: Any)

Perform all-to-all communication with autograd support.

Redistributes tensor data across processes by scattering along one dimension and gathering along another.

Parameters:

• input – Input tensor to redistribute.

• gather_dim – Dimension to gather along.

• scatter_dim – Dimension to scatter along.

• group – Process group for communication.

Returns:

Redistributed tensor.

rlightning.utils.distributed.collective.broadcast(tensor: MockModule('torch.Tensor'), comm_mode: CommMode, src: int = 0, async_op: bool = False) → Any | None

Broadcasts the tensor to the whole group.

tensor must have the same number of elements in all processes participating in the collective.

Parameters:

• tensor (Tensor) – Data to be sent if src is the rank of current process, and tensor to be used to save received data otherwise.

• comm_mode (CommMode) – Communication mode registered in CommContext.

• src (int) – Source rank.

• async_op (bool, optional) – Whether this op should be an async op.

Returns:

Async work handle, if async_op is set to True. None, if not async_op or if not part of the group

rlightning.utils.distributed.collective.broadcast_object_list(object_list: List[Any], comm_mode: CommMode, src: int = 0) → None

Broadcasts python objects based on torch.distributed.broadcast_object_list

Parameters:

• object_list (List[Any]) – List of input objects to broadcast. Each object must be picklable. Only objects on the src rank will be broadcast, but each rank must provide lists of equal sizes.

• src (int) – Source rank from which to broadcast object_list.

• comm_mode (CommMode) – Communication mode registered in CommContext.

Returns:

None

rlightning.utils.distributed.collective.gather(tensor: MockModule('torch.Tensor'), comm_mode: CommMode, gather_list: List[MockModule('torch.Tensor')] | None = None, dst: int = 0, async_op: bool = False) → Any | None

Gathers a list of tensors in a single process.

Parameters:

• tensor (Tensor) – Input tensor.

• comm_mode (CommMode) – Communication mode registered in CommContext.

• gather_list (list[Tensor], optional) – List of appropriately-sized tensors to use for gathered data (default is None, must be specified on the destination rank)

• dst (int, optional) – Destination rank (default is 0)

• async_op (bool, optional) – Whether this op should be an async op

Returns:

Async work handle, if async_op is set to True. None, if not async_op or if not part of the group

rlightning.utils.distributed.collective.gather_from_sequence_parallel_region(input_: MockModule('torch.Tensor'), rank0_only: bool = True)

Gather tensor from sequence parallel region.

Parameters:

• input – Local tensor chunk.

• rank0_only – If True, indicates the original tensor was only on rank 0.

Returns:

Concatenated tensor from all ranks.

rlightning.utils.distributed.collective.scatter(tensor: MockModule('torch.Tensor'), comm_mode: CommMode, scatter_list: List[MockModule('torch.Tensor')] | None = None, src: int = 0, async_op: bool = False) → Any | None

custom scatter operation.

Parameters:

• tensor (Tensor) – Output tensor.

• comm_mode (CommMode) – Communication mode registered in CommContext.

• scatter_list (list[Tensor]) – List of tensors to scatter (default is None, must be specified on the source rank).

• src (int) – Src rank.

• async_op (bool) – Whether this op should be an async op.

Returns:

Async work handle, if async_op is set to True. None, if not async_op or if not part of the group

rlightning.utils.distributed.collective.scatter_to_sequence_parallel_region(input_: MockModule('torch.Tensor'), rank0_only: bool = True)

Scatter tensor to sequence parallel region.

Parameters:

• input – Input tensor to scatter.

• rank0_only – If True, scatter from rank 0; otherwise, split locally.

Returns:

Scattered tensor chunk for this rank.