FHElium

Appearance

fhelium.distributed._state ​

Process-group initialization and rank-local device helpers.

local_device ​

function View source

python
def local_device() -> torch.device: ...

Return rank-local device metadata selected for the current process.

init selects the CUDA device identified by LOCAL_RANK. The helper reads the current CUDA device when CUDA is available and otherwise returns CPU; it is available before process-group initialization.

Returns

The current CUDA device when CUDA is available, otherwise the CPU device. World size does not affect the result.

init ​

function View source

python
def init(*, backend: str | None=None, init_method: str | None=None, timeout: timedelta | None=None, world_size: int=-1, rank: int=-1, store: Any | None=None, pg_options: Any | None=None, device_id: torch.device | int | None=None) -> None: ...

Initialize the PyTorch default process group and local CUDA device.

The helper uses PyTorch's standard process group and the torchrun rank environment. A direct world-size-one program uses a local HashStore while retaining ordinary PyTorch collective and asynchronous Work behavior. rank is the global rank in the newly created default process group; initialization completes synchronously and returns None.

Parameters

  • backend: Process-group backend. Defaults to "nccl" when CUDA is available and "gloo" otherwise.
  • init_method: Optional PyTorch rendezvous URL. When omitted for a direct world-size-one launch, a local HashStore is used unless store is supplied.
  • timeout: Optional process-group operation timeout forwarded to torch.distributed.init_process_group.
  • world_size: Number of processes in the default group. -1 resolves WORLD_SIZE from the environment.
  • rank: Global rank in the default group. -1 resolves RANK from the environment.
  • store: Optional rendezvous key-value store forwarded to PyTorch.
  • pg_options: Optional backend-specific process-group options forwarded to PyTorch.
  • device_id: Rank-local device. When omitted with CUDA available, cuda:LOCAL_RANK is selected. A CUDA device is required when CUDA is available; only a CPU device is accepted otherwise.

Returns

None. If a process group is already initialized, the selected CUDA device is still applied and process-group initialization is skipped.

Raises

  • ValueError: If device_id is incompatible with CUDA availability.
  • RuntimeError: If torch.distributed is unavailable or PyTorch cannot initialize the requested process group.

shutdown ​

function View source

python
def shutdown(group: torch.distributed.ProcessGroup | None=None) -> None: ...

Destroy an initialized process group.

Parameters

  • group: Process group to destroy. None selects the PyTorch default process group; this is a group object, not a rank identifier.

Returns

None. An unavailable or uninitialized distributed runtime, including a direct process before init, is a no-op. A world-size-one group is destroyed in the same way as any other initialized group. No asynchronous torch.distributed.Work handle is returned.