rlightning.utils.distributed

Distributed training utilities module.

This module provides utilities for distributed training with PyTorch, including process group initialization, communication contexts, and collective operations.

rlightning.utils.distributed.initialize_distributed_env(rank: int, world_size: int, backend: str = 'nccl', dist_url: str = 'env://', timeout: int = 1800) → None[source]

Initialize the PyTorch distributed environment.

This function wraps CommContext.init_distributed_env to provide a convenient interface for initializing distributed training.

Parameters:

  • rank – Global rank of this process.

  • world_size – Total number of processes in the distributed group.

  • backend – Communication backend (‘nccl’ or ‘gloo’). Defaults to ‘nccl’.

  • dist_url – URL for distributed initialization. Defaults to ‘env://’.

  • timeout – Timeout in seconds for initialization. Defaults to 1800.

Subpackages

Submodules