Core distributed
Process-group setup, context-parallel split/gather, model-parallel groups, rank-aware logging, EMA broadcast, and pipeline scheduling.
On this page
Distributed Core separates process-group topology from tensor movement. Initialization creates global and model-parallel groups. Rank/group queries expose that topology. Collective helpers then split, gather, broadcast, or synchronize tensors without forcing each model integration to reproduce edge cases.
Importing this package requires the distributed model runtime dependencies, including Torch and Loguru. dist_init additionally requires CUDA; the cp_group=None behavior shown below is a contract example, not a claim that the distributed package belongs to the minimal CLI environment.
Single-process behavior is intentional
The context-parallel data helpers accept None as “parallelism disabled.” This makes the same model code usable in one process and many processes without branching at every call site.
import torch
from worldfoundry.core.distributed import cat_outputs_cp, split_inputs_cp
x = torch.arange(24).reshape(2, 3, 4)
local = split_inputs_cp(x, seq_dim=1, cp_group=None)
restored = cat_outputs_cp(local, seq_dim=1, cp_group=None)
assert local is x
assert restored is xWith a real CP group, split_inputs_cp requires the sequence dimension to be divisible by the group size. cat_outputs_cp gathers equal-shaped local tensors in rank order. cat_outputs_cp_with_grad restores the local rank's autograd reference after gathering; use it only when gradients must flow through the local shard.
Typical context-parallel sequence
Rank zero or the minimum rank can load an input once, broadcast_split_tensor broadcasts its full shape/data and returns one shard per rank, the model computes locally, and cat_outputs_cp restores the global sequence. The split and gather must use the same seq_dim and process group.
broadcast supports tensors and Python objects, but object collectives serialize data and are inappropriate for high-volume model activations. find_split plans temporal/spatial CP dimensions for Megatron-compatible runtimes and mutates the compatibility parallel state; it is not a generic tensor chunker.
Initialization boundary
dist_init is the high-level CUDA inference initializer for runtimes whose config supplies backend, timeout, CP size, and PP size. It reads RANK and WORLD_SIZE, binds the local CUDA device, verifies cp_size * pp_size, initializes model-parallel groups, and creates a pipeline scheduler when needed.
Do not call model-parallel rank/group accessors before initialization. Pair lifecycle teardown with destroy_model_parallel in tests or long-lived workers that rebuild topology. print_rank_0 is safe outside distributed execution and is preferable for messages that should appear once.
Complete reference
The blocks below are the generated signatures for this category. Use the on-page symbol index to jump; source links open the defining implementation behind each lazy export.
20 public symbols
broadcast
funcdef broadcast(item: torch.Tensor | str | None,process_group: ProcessGroup | None = None) -> torch.Tensor | str | Noneworldfoundry.core.distributed.broadcastfrom worldfoundry.core.distributed import broadcastOverview
Broadcast a tensor or object from the minimum rank in `process_group. Belongs to Core distributed helpers (collectives and context-parallel split/gather). Annotated return type: torch.Tensor | str | None`.
Parameters
itemtorch.Tensor | str | Noneprocess_groupProcessGroup | None- default:
None
Returns: torch.Tensor | str | None
def broadcast_dtensor_model_states(model: torch.nn.Module, mesh: DeviceMesh) -> Noneworldfoundry.core.broadcast_dtensor_model_statesfrom worldfoundry.core import broadcast_dtensor_model_statesOverview
Broadcast model parameters and buffers from the first rank in the replicate mesh. Belongs to Core distributed helpers (collectives and context-parallel split/gather). Annotated return type: None.
Parameters
modeltorch.nn.ModulemeshDeviceMesh
Returns: None
def broadcast_split_tensor(tensor: torch.Tensor | None,seq_dim: int,process_group: ProcessGroup | None = None) -> torch.Tensor | Noneworldfoundry.core.distributed.broadcast_split_tensorfrom worldfoundry.core.distributed import broadcast_split_tensorOverview
Broadcast a tensor from the minimum CP rank, then return this rank's shard. Belongs to Core distributed helpers (collectives and context-parallel split/gather). Annotated return type: torch.Tensor | None.
Parameters
tensortorch.Tensor | Noneseq_dimintprocess_groupProcessGroup | None- default:
None
Returns: torch.Tensor | None
cat_outputs_cp
funcdef cat_outputs_cp(x: Tensor,seq_dim: int,cp_group: ProcessGroup | None = None) -> Tensorworldfoundry.core.distributed.cat_outputs_cpfrom worldfoundry.core.distributed import cat_outputs_cpOverview
Gather and concatenate per-rank tensors along `seq_dim. Belongs to Core distributed helpers (collectives and context-parallel split/gather). Annotated return type: Tensor`.
Parameters
xTensor- This rank's local tensor.
seq_dimint- Concatenation dimension.
cp_groupProcessGroup | None- CP process group; `
Nonereturnsx` unchanged.default:None
Raises
RuntimeError- `
all_gather` failed.
Returns: Tensor — Tensor with the gathered shards concatenated along `seq_dim`.
def cat_outputs_cp_with_grad(x: Tensor,seq_dim: int,cp_group: ProcessGroup | None = None) -> Tensorworldfoundry.core.distributed.cat_outputs_cp_with_gradfrom worldfoundry.core.distributed import cat_outputs_cp_with_gradOverview
Gather CP shards while preserving the local rank's autograd graph. Belongs to Core distributed helpers (collectives and context-parallel split/gather). Annotated return type: Tensor.
Parameters
xTensorseq_dimintcp_groupProcessGroup | None- default:
None
Returns: Tensor
def destroy_model_parallel()worldfoundry.core.destroy_model_parallelfrom worldfoundry.core import destroy_model_parallelOverview
Tear down model-parallel process groups by clearing the stored group handles.
dist_init
funcOverview
Initialize torch.distributed plus WorldFoundry CP/PP groups. Belongs to Core distributed helpers (collectives and context-parallel split/gather).
Parameters
config- Runtime config whose `
engine_configprovidesdistributed_backend,distributed_timeout_minutes,cp_size, andpp_size`.
Notes
Rank and world size come from `RANK and WORLD_SIZE. The function requires CUDA, binds each rank to a local device, verifies cp_size * pp_size == world_size`, and initializes the pipeline scheduler when pipeline parallelism is active.
class DTensorFastEmaModelUpdater()worldfoundry.core.DTensorFastEmaModelUpdaterfrom worldfoundry.core import DTensorFastEmaModelUpdaterOverview
Foreach-based EMA updater that operates on local DTensor shards. Belongs to Core distributed helpers (collectives and context-parallel split/gather).
Methods
Overview
Public method on this type.
Parameters
src_modeltorch.nn.Moduletgt_modeltorch.nn.Module
Returns: None
update_average(src_model: torch.nn.Module,tgt_model: torch.nn.Module,beta: float = 0.9999) -> NonesourceOverview
Public method on this type.
Parameters
src_modeltorch.nn.Moduletgt_modeltorch.nn.Modulebetafloat- default:
0.9999
Returns: None
Overview
Public method on this type.
Parameters
parametersAnyis_cpubool- default:
False
Returns: None
Overview
Public method on this type.
Parameters
parametersAny
Returns: None
find_split
funcdef find_split(shape_tensor: torch.Size,cp_size: int,patch_values: tuple[int, int, int] = (1, 2, 2),view_factor: int = 1) -> torch.Sizeworldfoundry.core.distributed.find_splitfrom worldfoundry.core.distributed import find_splitOverview
Find the post-context-parallel temporal/spatial split shape. Belongs to Core distributed helpers (collectives and context-parallel split/gather). Annotated return type: torch.Size.
Parameters
shape_tensortorch.Sizecp_sizeintpatch_valuestuple[int, int, int]- default:
(1, 2, 2) view_factorint- default:
1
Returns: torch.Size
def get_local_tensor_if_dtensor(tensor)worldfoundry.core.get_local_tensor_if_dtensorfrom worldfoundry.core import get_local_tensor_if_dtensorOverview
Return the local shard for DTensor inputs; leave regular tensors unchanged. Belongs to Core distributed helpers (collectives and context-parallel split/gather).
Parameters
tensor
def init_pp_scheduler()worldfoundry.core.init_pp_schedulerfrom worldfoundry.core import init_pp_schedulerOverview
Initialize the pipeline-parallel scheduler singleton used by PP runs.
Raises
AssertionError- If the PPScheduler is already initialized.
def initialize_model_parallel(tp_size: int = 1,pp_size: int = 1,cp_size: int = 1,nccl_communicator_config_path: Optional[str] = None,distributed_timeout_minutes: int = 30,order: str = 'tp-cp-pp-dp') -> Noneworldfoundry.core.initialize_model_parallelfrom worldfoundry.core import initialize_model_parallelOverview
Initialize model data parallel groups. Belongs to Core distributed helpers (collectives and context-parallel split/gather). Annotated return type: None.
Parameters
tp_sizeint- The number of GPUs to split individual tensors across.default:
1 pp_sizeint- The number of tensor parallel GPU groups to split the Transformer layers across. For example, if tp_size is 4 and pp_size is 2, the model will be split into 2 groups of 4 GPUs.default:
1 cp_sizeint- The number of tensor parallel GPU groups to split the network input sequence length across. Compute of attention module requires tokens of full sequence length, so GPUs in a context parallel group need to communicate with each other to exchange information of other sequence chunks. Each GPU and its counterparts in other tensor parallel groups compose a context parallel group. For example, assume we have 8 GPUs, if tensor model parallel size is 4 and context parallel size is 2, the network input will be split into two sequence chunks, which are processed by 2 different groups of 4 GPUs. One chunk is processed by GPU0-3, the other chunk is processed by GPU4-7. Four groups are build to do context parallel communications: [GPU0, GPU4], [GPU1, GPU5], [GPU2, GPU6], and [GPU3, GPU7]. Context parallelism partitions sequence length, so it has no impact on weights, which means weights are duplicated among GPUs in a context parallel group. Hence, weight gradients all-reduce is required in backward. For simplicity, we piggyback GPUs of context parallelism on data parallel group for weight gradient all-reduce.default:
1 nccl_communicator_config_pathOptional[str]- Path to the yaml file of NCCL communicator configurations.
min_ctas,max_ctas, andcga_cluster_sizecan be set for each communicator.default:None distributed_timeout_minutesint- Timeout, in minutes,for operations executed against distributed process groups. See PyTorch documentation atdefault:
30 orderstr- The rank initialization order of parallelism. Now we support tp-dp-pp and tp-pp-dp orders. Let's say we have a total of 16 GPUs denoted by g0 ... g15 and we use 2 GPUs to parallelize the model tensor, and 4 GPUs to parallelize the model pipeline. The present function will create 8 tensor model-parallel groups, 4 pipeline model-parallel groups and 8 data-parallel groups as: 8 data_parallel groups: [g0, g2], [g1, g3], [g4, g6], [g5, g7], [g8, g10], [g9, g11], [g12, g14], [g13, g15] 8 tensor model-parallel groups: [g0, g1], [g2, g3], [g4, g5], [g6, g7], [g8, g9], [g10, g11], [g12, g13], [g14, g15] 4 pipeline model-parallel groups: [g0, g4, g8, g12], [g1, g5, g9, g13], [g2, g6, g10, g14], [g3, g7, g11, g15] Note that for efficiency, the caller should make sure adjacent ranks are on the same DGX box. For example if we are using 2 DGX-1 boxes with a total of 16 GPUs, rank 0 to 7 belong to the first box and ranks 8 to 15 belong to the second box.default:
'tp-cp-pp-dp'
Returns: None
is_last_rank
funcdef is_last_rank()worldfoundry.core.is_last_rankfrom worldfoundry.core import is_last_rankOverview
Return whether this worker is the final rank in the global process group. Belongs to Core distributed helpers (collectives and context-parallel split/gather).
def is_last_tp_cp_rank()worldfoundry.core.is_last_tp_cp_rankfrom worldfoundry.core import is_last_tp_cp_rankOverview
Return whether this worker is last in the combined tensor/context group. Belongs to Core distributed helpers (collectives and context-parallel split/gather).
def model_parallel_is_initialized()worldfoundry.core.model_parallel_is_initializedfrom worldfoundry.core import model_parallel_is_initializedOverview
Check if model and data parallel groups are initialized. Belongs to Core distributed helpers (collectives and context-parallel split/gather).
pp_scheduler
funcdef pp_scheduler() -> PPSchedulerworldfoundry.core.pp_schedulerfrom worldfoundry.core import pp_schedulerOverview
Return the current pipeline-parallel scheduler instance.
Raises
AssertionError- If the PPScheduler has not been initialized.
Returns: PPScheduler — PPScheduler: The current PPScheduler instance.
PPScheduler
clsclass PPScheduler()worldfoundry.core.PPSchedulerfrom worldfoundry.core import PPSchedulerOverview
Minimal point-to-point scheduler for adjacent pipeline stages. Belongs to Core distributed helpers (collectives and context-parallel split/gather).
Source docstring
Minimal point-to-point scheduler for adjacent pipeline stages.
The scheduler owns a CUDA receive buffer queue and exposes synchronous and asynchronous helpers around the model-parallel pipeline group. Initialize the process-wide instance with `init_pp_scheduler and retrieve it with pp_scheduler`.
Methods
Overview
Asynchronously send a tensor to the next pipeline and return the send handle.
Parameters
tensortorch.Tensor- The tensor to be sent.
Returns: torch.distributed.Work — torch.distributed.Work: The handle for the send operation.
Overview
Asynchronously receive a tensor from the previous pipeline and return the receive handle.
Parameters
buffertorch.Tensor- The buffer tensor for receiving data.
Returns: torch.distributed.Work — torch.distributed.Work: The handle for the receive operation.
Overview
Receive data from the previous pipeline and return the received tensor.
Parameters
shapetorch.Size- The shape of the tensor to receive.
dtypetorch.dtype- The data type of the tensor to receive.
Returns: torch.Tensor — torch.Tensor: The received tensor.
Overview
Put the asynchronously received tensor and handle into the receive queue.
Parameters
shapetorch.Size- The shape of the tensor to receive.
dtypetorch.dtype- The data type of the tensor to receive.
Returns: None
Overview
Get a tensor from the receive queue and wait for the receive operation to complete.
Returns: torch.Tensor — torch.Tensor: The received tensor obtained from the queue.
print_per_rank
funcdef print_per_rank(message: object) -> Noneworldfoundry.core.print_per_rankfrom worldfoundry.core import print_per_rankOverview
Emit one informational log record from every calling rank. Belongs to Core distributed helpers (collectives and context-parallel split/gather). Annotated return type: None.
Parameters
messageobject
Returns: None
print_rank_0
funcdef print_rank_0(message: object) -> Noneworldfoundry.core.print_rank_0from worldfoundry.core import print_rank_0Overview
Emit one informational record on rank zero, or in a non-distributed process. Belongs to Core distributed helpers (collectives and context-parallel split/gather). Annotated return type: None.
Parameters
messageobject
Returns: None
split_inputs_cp
funcdef split_inputs_cp(x: Tensor,seq_dim: int,cp_group: ProcessGroup | None = None) -> Tensorworldfoundry.core.distributed.split_inputs_cpfrom worldfoundry.core.distributed import split_inputs_cpOverview
Slice a tensor along `seq_dim to this rank's CP shard. Belongs to Core distributed helpers (collectives and context-parallel split/gather). Annotated return type: Tensor`.
Parameters
xTensor- Input tensor.
seq_dimint- Dimension to split along (negative indexing supported).
cp_groupProcessGroup | None- CP process group; `
Nonereturnsx` unchanged.default:None
Raises
AssertionError- `
seq_dim` is not divisible by the CP size.
Returns: Tensor — Contiguous slice of length `x.shape[seq_dim] // cp_size`.