rlightning.utils.distributed.group_initializer

Process group initializers for distributed training.

This module provides classes for initializing PyTorch distributed process groups with various parallel modes including data parallel, tensor parallel, sequence parallel, and weight transfer groups.

class rlightning.utils.distributed.group_initializer.CommMode[source]

Bases: object

Communication mode identifiers for process groups.

GLOBAL

Global communication group spanning all processes.

Type:

str

INTRA_NODE

Communication group within a single node.

Type:

str

INTER_NODE

Communication group across nodes (one process per node).

Type:

str

GLOBAL: str = 'GLOBAL'
INTER_NODE: str = 'INTER_NODE'
INTRA_NODE: str = 'INTRA_NODE'
class rlightning.utils.distributed.group_initializer.DPGroupInitializer(rank: int, world_size: int, data_parallel_size: int)[source]

Bases: ProcessGroupInitializer

Data parallel group initializer.

Creates data parallel groups based on the specified parallelism size. Groups are formed by selecting processes at regular intervals.

data_parallel_size

Number of processes in each data parallel group.

process_num_between_dp_rank

Stride between data parallel ranks.

data_parallel_size: int
init_dist_group(backend: str = 'nccl', use_cpu: bool = False) → Tuple[int | None, int | None, MockModule('torch.distributed.ProcessGroup') | None, List[int] | None, str][source]

Initialize data parallel groups.

Creates multiple groups where each group contains processes separated by process_num_between_dp_rank stride.

Parameters:

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

  • use_cpu – If True, use CPU tensors.

Returns:

Tuple of (local_rank, group_world_size, process_group, ranks_in_group, mode).

process_num_between_dp_rank: int
class rlightning.utils.distributed.group_initializer.InterNodeGroupInitializer(rank: int, world_size: int, local_world_size: int)[source]

Bases: ProcessGroupInitializer

Inter-node process group initializer.

Creates process groups for communication across nodes, with one representative process per node in each group.

local_world_size

Number of processes per node.

node_count

Total number of nodes in the cluster.

init_dist_group(backend: str = 'nccl', use_cpu: bool = False) → Tuple[int | None, int | None, MockModule('torch.distributed.ProcessGroup') | None, List[int] | None, str][source]

Initialize inter-node communication groups.

Creates groups where each group contains processes with the same local rank across different nodes.

Parameters:

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

  • use_cpu – If True, use CPU tensors.

Returns:

Tuple of (local_rank, group_world_size, process_group, ranks_in_group, mode).

local_world_size: int
node_count: int
class rlightning.utils.distributed.group_initializer.IntraNodeGroupInitializer(rank: int, world_size: int, ranks: List[int])[source]

Bases: ProcessGroupInitializer

Intra-node process group initializer.

Creates a process group for communication within a single node.

ranks

List of global ranks within the node.

init_dist_group(backend: str = 'nccl', use_cpu: bool = False) → Tuple[int | None, int | None, MockModule('torch.distributed.ProcessGroup') | None, List[int] | None, str][source]

Initialize the intra-node communication group.

Parameters:

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

  • use_cpu – If True, use CPU tensors.

Returns:

Tuple of (local_rank, group_world_size, process_group, ranks_in_group, mode).

ranks: List[int]
class rlightning.utils.distributed.group_initializer.ParallelMode[source]

Bases: CommMode

Parallel mode identifiers extending CommMode.

TRAIN_DATA_PARALLEL

Data parallel group for training.

Type:

str

DATA_PARALLEL

General data parallel group.

Type:

str

TENSOR_PARALLEL

Tensor parallel group for model parallelism.

Type:

str

SEQUENCE_PARALLEL

Sequence parallel group for long sequences.

Type:

str

WEIGHT_TRANSFER

Group for transferring model weights.

Type:

str

DATA_PARALLEL: str = 'DATA_PARALLEL'
SEQUENCE_PARALLEL: str = 'SEQUENCE_PARALLEL'
TENSOR_PARALLEL: str = 'TENSOR_PARALLEL'
TRAIN_DATA_PARALLEL: str = 'TRAIN_DATA_PARALLEL'
WEIGHT_TRANSFER: str = 'WEIGHT_TRANSFER'
class rlightning.utils.distributed.group_initializer.ProcessGroupInitializer(rank: int, world_size: int)[source]

Bases: ABC

Abstract base class for process group initialization.

This class provides the foundation for creating PyTorch distributed process groups with various configurations.

rank

Global rank of this process.

world_size

Total number of processes.

local_rank

Local rank within the initialized group.

ranks_in_group

List of global ranks in the group.

process_group

The initialized ProcessGroup instance.

group_world_size

Size of the initialized group.

mode

Communication mode identifier for this group.

group_world_size: int | None
abstractmethod init_dist_group(backend: str = 'nccl', use_cpu: bool = False) → Tuple[int | None, int | None, MockModule('torch.distributed.ProcessGroup') | None, List[int] | None, str | None][source]

Initialize the distributed group.

Parameters:

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

  • use_cpu – If True, use CPU tensors (forces ‘gloo’ backend).

Returns:

Tuple of (local_rank, group_world_size, process_group, ranks_in_group, mode).

local_rank: int | None
mode: str | None
process_group: MockModule('torch.distributed.ProcessGroup') | None
rank: int
ranks_in_group: List[int] | None
world_size: int
class rlightning.utils.distributed.group_initializer.SPGroupInitializer(rank: int, world_size: int, sequence_parallel_size: int)[source]

Bases: ProcessGroupInitializer

Sequence parallel group initializer.

Creates sequence parallel groups for parallelizing long sequence processing across multiple processes.

sequence_parallel_size

Number of processes in each sequence parallel group.

sequence_parallel_group_num

Total number of sequence parallel groups.

init_dist_group(backend: str = 'nccl', use_cpu: bool = False) → Tuple[int | None, int | None, MockModule('torch.distributed.ProcessGroup') | None, List[int] | None, str][source]

Initialize sequence parallel groups.

Creates groups of consecutive processes for sequence parallelism.

Parameters:

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

  • use_cpu – If True, use CPU tensors.

Returns:

Tuple of (local_rank, group_world_size, process_group, ranks_in_group, mode).

sequence_parallel_group_num: int
sequence_parallel_size: int
class rlightning.utils.distributed.group_initializer.TPGroupInitializer(rank: int, world_size: int, tensor_parallel_size: int)[source]

Bases: ProcessGroupInitializer

Tensor parallel group initializer.

Creates tensor parallel groups for model parallelism where model parameters are sharded across processes.

tensor_parallel_size

Number of processes in each tensor parallel group.

tensor_parallel_group_num

Total number of tensor parallel groups.

init_dist_group(backend: str = 'nccl', use_cpu: bool = False) → Tuple[int | None, int | None, MockModule('torch.distributed.ProcessGroup') | None, List[int] | None, str][source]

Initialize tensor parallel groups.

Creates groups of consecutive processes for tensor parallelism.

Parameters:

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

  • use_cpu – If True, use CPU tensors.

Returns:

Tuple of (local_rank, group_world_size, process_group, ranks_in_group, mode).

tensor_parallel_group_num: int
tensor_parallel_size: int
class rlightning.utils.distributed.group_initializer.TrainDPGroupInitializer(rank: int, world_size: int, ranks: List[int])[source]

Bases: ProcessGroupInitializer

Data parallel group initializer for training.

Creates a process group containing specified ranks for data parallel training synchronization.

ranks

List of global ranks to include in this group.

init_dist_group(backend: str = 'nccl', use_cpu: bool = False) → Tuple[int | None, int | None, MockModule('torch.distributed.ProcessGroup') | None, List[int] | None, str][source]

Initialize the training data parallel group.

Parameters:

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

  • use_cpu – If True, use CPU tensors.

Returns:

Tuple of (local_rank, group_world_size, process_group, ranks_in_group, mode).

ranks: List[int]
class rlightning.utils.distributed.group_initializer.WeightTransferGroupInitializer(rank: int, world_size: int, ranks: List[int])[source]

Bases: ProcessGroupInitializer

Weight transfer process group initializer.

Creates a process group for transferring model weights between processes.

ranks

List of global ranks participating in weight transfer.

init_dist_group(backend: str = 'nccl', use_cpu: bool = False) → Tuple[int | None, int | None, MockModule('torch.distributed.ProcessGroup') | None, List[int] | None, str][source]

Initialize the weight transfer communication group.

Parameters:

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

  • use_cpu – If True, use CPU tensors.

Returns:

Tuple of (local_rank, group_world_size, process_group, ranks_in_group, mode).

ranks: List[int]