helios.core.distributed

Classes

DistributedInfo

Bundle information regarding distributed state.

Functions

is_using_torchrun(→ bool)	Check if the current process was launched from torchrun.
init_dist(→ None)	Initialize the distributed process group.
shutdown_dist(→ None)	Shutdown the distributed process group.
is_dist_active(→ bool)	Check if torch.distributed is active.
get_dist_info(→ DistributedInfo)	Get the distributed state of the current run.
get_local_rank(→ int)	Get the local rank of the device the process is running on.
get_global_rank(→ int)	Get the global rank of the device the process is running on.
gather_into_tensor(→ torch.Tensor)	Gathers the tensors across all processes and merges them into a single tensor.
all_reduce_tensors(→ torch.Tensor)	Reduces tensors across all processes so all have the same value.
global_print(→ None)	Print wrapper that only prints on the specified global rank.
local_print(→ None)	Print wrapper that only prints on the specified local rank.
safe_barrier(→ None)	Safe wrapper for torch.distributed.barrier.

Module Contents

helios.core.distributed.is_using_torchrun() → bool

Check if the current process was launched from torchrun.

This will check to see if the environment variables that are set by torchrun exist. The list of variables is taken directly from the documentation and can be seen here.

Returns:

True if the run was started from torchrun, false otherwise.

helios.core.distributed.init_dist(backend: str = 'nccl', rank: int | None = None, world_size: int | None = None) → None

Initialize the distributed process group.

The optional values for rank and world_size must be omitted if distributed training is handled through torchrun. If distributed training is started manually, then both arguments must be provided.

Parameters:

• backend – the backend to use. Defaults to “nccl”.

• rank – the (optional) rank of the current GPU.

• world_size – the (optional) number of GPUs in the system.

Raises:

ValueError – if either of rank or world_size are None (but not both).

helios.core.distributed.shutdown_dist() → None

Shutdown the distributed process group.

helios.core.distributed.is_dist_active() → bool

Check if torch.distributed is active.

class helios.core.distributed.DistributedInfo

Bundle information regarding distributed state.

To understand what these mean, consider a run being performed over two nodes, each with 2 GPUs. Suppose we have a single process per GPU. Then the following values are assigned:

• local_rank: 0 or 1 for both nodes. 0 for the first GPU, 1 for the second.

• rank: 0, 1, 2, 3 for each of the GPUs over the two nodes.

• local_world_size: 2 for both nodes.

• world_size: 4 (2 nodes with 2 workers per node).

Parameters:

• local_rank – the local rank.

• rank – the global rank.

• local_world_size – the local world size.

• world_size – the global world size.

local_rank: int = 0

rank: int = 0

local_world_size: int = 1

world_size: int = 1

helios.core.distributed.get_dist_info() → DistributedInfo

Get the distributed state of the current run.

If distributed training is not used, then both ranks are set to 0 and both world sizes are set to 1.

Returns:

The information of the current distributed run.

helios.core.distributed.get_local_rank() → int

Get the local rank of the device the process is running on.

If distributed training is not used, 0 is returned.

Returns:

The local rank of the current device.

helios.core.distributed.get_global_rank() → int

Get the global rank of the device the process is running on.

If distributed training is not used, 0 is returned.

Returns:

The global rank of the current device.

helios.core.distributed.gather_into_tensor(tensor: torch.Tensor, size: tuple[int]) → torch.Tensor

Gathers the tensors across all processes and merges them into a single tensor.

Parameters:

• tensor – the tensor to merge

• size – the dimensions of the output tensor.

Returns:

The resulting tensor containing all gathered tensors.

Raises:

RuntimeError – if distributed training hasn’t been initialised.

helios.core.distributed.all_reduce_tensors(tensor: torch.Tensor | list[torch.Tensor], **kwargs) → torch.Tensor

Reduces tensors across all processes so all have the same value.

Parameters:

• tensor – the input tensor(s) to reduce. If the input is a list of tensors, they will be concatenated into a single tensor.

• kwargs – additional options for torch.distributed.all_reduce

Returns:

The reduced tensor.

Raises:

RuntimeError – if distributed training has not been initialised.

helios.core.distributed.global_print(*args: Any, global_rank: int = 0, **kwargs: Any) → None

Print wrapper that only prints on the specified global rank.

Parameters:

• *args – positional arguments to pass in to Python’s print.

• global_rank – the global rank the print should happen on. Defaults to 0.

• **kwargs – keyword arguments to pass in to Python’s print.

helios.core.distributed.local_print(*args: Any, local_rank: int = 0, **kwargs: Any) → None

Print wrapper that only prints on the specified local rank.

Parameters:

• *args – positional arguments to pass in to Python’s print.

• local_rank – the local rank the print should happen on. Defaults to 0.

• **kwargs – keyword arguments to pass in to Python’s print.

helios.core.distributed.safe_barrier(**kwargs: Any) → None

Safe wrapper for torch.distributed.barrier.

The wrapper is “safe” in the sense that it is valid to call this function regardless of whether the code is currently using distributed training or not.

Parameters:

**kwargs – keyword arguments to torch.distributed.barrier.