Distributed Operations

About

The d9d.core.dist_ops package wraps torch.distributed collective operations and allocates their output buffers for you. With plain PyTorch, you must pre-allocate the outputs yourself, for example a list of empty tensors for all_gather. The package also has operations for variadic shapes. They let ranks exchange tensors without knowing the shapes of the incoming tensors in advance.

Usage

Gathering Tensors

Gather tensors of the same shape from all ranks.

import torch
from d9d.core.dist_context import DistributedContext, FLAT_DOMAIN
from d9d.core.dist_ops import all_gather

# Setup.
ctx: DistributedContext = ...
mesh = ctx.mesh_for(FLAT_DOMAIN)
group = mesh.get_group()
rank = mesh.get_local_rank()

# Each rank has a tensor of the same shape, but with different values.
local_tensor = torch.ones((2, 2), device="cuda") * rank

# Gather.
gathered_tensors = all_gather(local_tensor, group=group)

for i, t in enumerate(gathered_tensors):
    print(f"From rank {i}: {t}")

Gathering Tensors with Variadic Shapes

Gather tensors whose shapes differ across ranks.

import torch
from d9d.core.dist_context import DistributedContext, FLAT_DOMAIN
from d9d.core.dist_ops import all_gather_variadic_shape

# Setup.
ctx: DistributedContext = ...
mesh = ctx.mesh_for(FLAT_DOMAIN)
group = mesh.get_group()
rank = mesh.get_local_rank()

# Rank 0 has shape (1,), rank 1 has shape (2,), ...
local_tensor = torch.randn((rank + 1,), device="cuda")

# Gather: the shapes are exchanged first.
gathered_tensors = all_gather_variadic_shape(local_tensor, group=group)

for i, t in enumerate(gathered_tensors):
    print(f"Rank {i} sent shape: {t.shape}")

Gathering Objects

Gather Python objects from all ranks. The objects must be picklable.

from d9d.core.dist_context import DistributedContext, FLAT_DOMAIN
from d9d.core.dist_ops import all_gather_object

# Setup.
ctx: DistributedContext = ...
mesh = ctx.mesh_for(FLAT_DOMAIN)
group = mesh.get_group()
rank = mesh.get_local_rank()

# Local data.
my_metadata = {
    "rank": rank,
    "status": "ready"
}

# Gather.
results = all_gather_object(my_metadata, group=group)

for data in results:
    print(f"Rank {data['rank']} sent {data}")

API Reference

`d9d.core.dist_ops`

Wrappers around torch.distributed collectives that allocate the output buffers.

`all_gather(tensor, group, async_op=False)`

Gathers tensors from the whole process group to all ranks.

Tensors on all ranks must have the same shape and dtype as the local tensor. The output list is allocated on all ranks.

Parameters:

Name	Type	Description	Default
`tensor`	`Tensor`	The local tensor to send.	required
`group`	`ProcessGroup`	The process group to work on.	required
`async_op`	`bool`	Whether the operation should be asynchronous.	`False`

Returns:

Type	Description
`list[Tensor] \| tuple[list[Tensor], Work]`	If `async_op` is `False`, a list of gathered tensors. If `async_op` is `True`, a tuple of this
`list[Tensor] \| tuple[list[Tensor], Work]`	list and the work handle.

`all_gather_object(obj, group)`

Gathers picklable objects from the whole process group to all ranks.

Wraps torch.distributed.all_gather_object and allocates the output list on all ranks.

Parameters:

Name	Type	Description	Default
`obj`	`T`	The local object to send. Must be picklable.	required
`group`	`ProcessGroup`	The process group to work on.	required

Returns:

Type	Description
`list[T]`	A list of objects from all ranks.

`all_gather_variadic_shape(tensor, group, async_op=False)`

Gathers tensors of different shapes from the whole process group to all ranks.

Unlike all_gather, this function first exchanges the tensor shapes of all ranks.

Parameters:

Name	Type	Description	Default
`tensor`	`Tensor`	The local tensor to send.	required
`group`	`ProcessGroup`	The process group to work on.	required
`async_op`	`bool`	Whether the data gathering should be asynchronous. The shape exchange is always synchronous.	`False`

Returns:

Type	Description
`list[Tensor] \| tuple[list[Tensor], Work]`	If `async_op` is `False`, a list of gathered tensors of varying shapes. If `async_op` is `True`,
`list[Tensor] \| tuple[list[Tensor], Work]`	a tuple of this list and the work handle.

`gather(tensor, group, group_dst, async_op=False)`

Gathers tensors from the process group to a specific destination rank.

Tensors on all ranks must have the same shape and dtype as the local tensor. The output list is allocated on the destination rank.

Parameters:

Name	Type	Description	Default
`tensor`	`Tensor`	The local tensor to send.	required
`group`	`ProcessGroup`	The process group to work on.	required
`group_dst`	`int`	The rank within the group that will receive the tensors.	required
`async_op`	`bool`	Whether the operation should be asynchronous.	`False`

Returns:

Type	Description
`list[Tensor] \| tuple[list[Tensor] \| None, Work] \| None`	If `async_op` is `False`, a list of tensors on the destination rank and `None` on other ranks.
`list[Tensor] \| tuple[list[Tensor] \| None, Work] \| None`	If `async_op` is `True`, a tuple of this value and the work handle.

`gather_object(obj, group, group_dst)`

Gathers picklable objects from the whole process group to a specific destination rank.

Wraps torch.distributed.gather_object and allocates the output list on the destination rank.

Parameters:

Name	Type	Description	Default
`obj`	`T`	The local object to send. Must be picklable.	required
`group`	`ProcessGroup`	The process group to work on.	required
`group_dst`	`int`	The rank within the group that will receive the objects.	required

Returns:

Type	Description
`list[T] \| None`	A list of objects from all ranks on the destination rank, `None` on other ranks.

`gather_variadic_shape(tensor, group, group_dst)`

Gathers tensors of different shapes from the process group to a specific rank.

This function first exchanges the tensor shapes of all ranks. Then it gathers the tensors with point-to-point communication (isend and irecv). It does not support async_op.

Parameters:

Name	Type	Description	Default
`tensor`	`Tensor`	The local tensor to send.	required
`group`	`ProcessGroup`	The process group to work on.	required
`group_dst`	`int`	The rank within the group that will receive the tensors.	required

Returns:

Type	Description
`list[Tensor] \| None`	A list of tensors of varying shapes on the destination rank, `None` on other ranks.