Core 分布式

Process group 初始化、上下文并行切分聚合、模型并行 group、rank 日志、EMA 广播与 pipeline 调度。

最后更新: 2026年7月22日

本页内容

Distributed Core 把 process-group 拓扑和 tensor 移动分开。初始化负责创建全局与模型并行 group；rank/group 查询负责暴露拓扑；collective helper 再完成切分、聚合、广播或同步，模型接入不需要重复处理这些边界情况。

导入这个 package 需要分布式模型 runtime 依赖，包括 Torch 和 Loguru；dist_init 还必须有 CUDA。下面展示的 cp_group=None 是接口契约示例，并不表示 distributed package 属于最小 CLI 环境。

单进程行为是明确契约

上下文并行数据 helper 把 None 解释为“并行已关闭”。因此同一份模型代码可以在一个进程和多个进程中使用，不需要在每个调用点手写分支。

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 x

使用真实 CP group 时，split_inputs_cp 要求序列维长度能被 group size 整除。cat_outputs_cp 按 rank 顺序聚合等形状本地 tensor。cat_outputs_cp_with_grad 会在聚合后恢复本 rank 的 autograd 引用，只有梯度必须穿过本地 shard 时才使用。

典型上下文并行顺序

Rank 0 或 group 最小 rank 可以只加载一次输入；broadcast_split_tensor 先广播完整形状和数据，再给每个 rank 返回一个 shard；模型本地计算；最后 cat_outputs_cp 恢复全局序列。切分和聚合必须使用相同 seq_dim 与 process group。

broadcast 同时支持 tensor 和 Python object，但 object collective 会序列化数据，不适合高吞吐模型 activation。find_split 为 Megatron 兼容 runtime 规划时间与空间 CP 维度，并会更新兼容 parallel state；它不是普通 tensor chunk helper。

初始化边界

dist_init 是面向 CUDA 推理 runtime 的高层初始化器，配置需要提供 backend、timeout、CP size 和 PP size。它读取 RANK 与 WORLD_SIZE，绑定本地 CUDA device，验证 cp_size * pp_size，初始化模型并行 group，并在需要时创建 pipeline scheduler。

不要在初始化之前调用模型并行 rank/group accessor。在测试或会重建拓扑的常驻 worker 中，应使用 destroy_model_parallel 对称清理。print_rank_0 在非分布式执行中也安全，适合只应出现一次的日志。

完整参考

以下为该类别的生成签名。可用本页符号索引跳转；源码链接指向各惰性导出背后的具体实现。

20 个公开符号

`broadcast`

func

def broadcast(item: torch.Tensor | str | None,process_group: ProcessGroup | None = None) -> torch.Tensor | str | None

worldfoundry.core.distributed.broadcastfrom worldfoundry.core.distributed import broadcast

源码

简介

broadcast — Broadcast a tensor or object from the minimum rank in `process_group. 属于 Core 分布式辅助（集合通信与 context-parallel 切分/聚合）。标注返回类型：torch.Tensor | str | None`。

参数

itemtorch.Tensor | str | None
process_groupProcessGroup | None: 默认值: None

返回值: torch.Tensor | str | None

`broadcast_dtensor_model_states`

func

def broadcast_dtensor_model_states(model: torch.nn.Module, mesh: DeviceMesh) -> None

worldfoundry.core.broadcast_dtensor_model_statesfrom worldfoundry.core import broadcast_dtensor_model_states

源码

简介

broadcast_dtensor_model_states — Broadcast model parameters and buffers from the first rank in the replicate mesh. 属于 Core 分布式辅助（集合通信与 context-parallel 切分/聚合）。标注返回类型：None。

参数

modeltorch.nn.Module
meshDeviceMesh

返回值: None

`broadcast_split_tensor`

func

def broadcast_split_tensor(tensor: torch.Tensor | None,seq_dim: int,process_group: ProcessGroup | None = None) -> torch.Tensor | None

worldfoundry.core.distributed.broadcast_split_tensorfrom worldfoundry.core.distributed import broadcast_split_tensor

源码

简介

broadcast_split_tensor — Broadcast a tensor from the minimum CP rank, then return this rank's shard. 属于 Core 分布式辅助（集合通信与 context-parallel 切分/聚合）。标注返回类型：torch.Tensor | None。

参数

tensortorch.Tensor | None
seq_dimint
process_groupProcessGroup | None: 默认值: None

返回值: torch.Tensor | None

`cat_outputs_cp`

func

def cat_outputs_cp(x: Tensor,seq_dim: int,cp_group: ProcessGroup | None = None) -> Tensor

worldfoundry.core.distributed.cat_outputs_cpfrom worldfoundry.core.distributed import cat_outputs_cp

源码

简介

cat_outputs_cp — Gather and concatenate per-rank tensors along `seq_dim. 属于 Core 分布式辅助（集合通信与 context-parallel 切分/聚合）。标注返回类型：Tensor`。

参数

xTensor: This rank's local tensor.
seq_dimint: Concatenation dimension.
cp_groupProcessGroup | None: CP process group; `None returns x` unchanged.默认值: None

异常

RuntimeError: `all_gather` failed.

返回值: Tensor — Tensor with the gathered shards concatenated along `seq_dim`.

`cat_outputs_cp_with_grad`

func

def cat_outputs_cp_with_grad(x: Tensor,seq_dim: int,cp_group: ProcessGroup | None = None) -> Tensor

worldfoundry.core.distributed.cat_outputs_cp_with_gradfrom worldfoundry.core.distributed import cat_outputs_cp_with_grad

源码

简介

cat_outputs_cp_with_grad — Gather CP shards while preserving the local rank's autograd graph. 属于 Core 分布式辅助（集合通信与 context-parallel 切分/聚合）。标注返回类型：Tensor。

参数

xTensor
seq_dimint
cp_groupProcessGroup | None: 默认值: None

返回值: Tensor

`destroy_model_parallel`

func

def destroy_model_parallel()

worldfoundry.core.destroy_model_parallelfrom worldfoundry.core import destroy_model_parallel

源码

简介

清空已保存的进程组句柄，拆除 model-parallel 分组。

`dist_init`

func

def dist_init(config)

worldfoundry.core.dist_initfrom worldfoundry.core import dist_init

源码

简介

dist_init — Initialize torch.distributed plus WorldFoundry CP/PP groups. 属于 Core 分布式辅助（集合通信与 context-parallel 切分/聚合）。

参数

config: Runtime config whose `engine_config provides distributed_backend, distributed_timeout_minutes, cp_size, and pp_size`.

说明

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.

`DTensorFastEmaModelUpdater`

cls

class DTensorFastEmaModelUpdater()

worldfoundry.core.DTensorFastEmaModelUpdaterfrom worldfoundry.core import DTensorFastEmaModelUpdater

源码

简介

DTensorFastEmaModelUpdater — Foreach-based EMA updater that operates on local DTensor shards. 属于 Core 分布式辅助（集合通信与 context-parallel 切分/聚合）。

方法

methcopy_to(src_model: torch.nn.Module, tgt_model: torch.nn.Module) -> None源码

简介

该类型上的公开 method。

参数

src_modeltorch.nn.Module
tgt_modeltorch.nn.Module

返回值: None

methupdate_average(src_model: torch.nn.Module,tgt_model: torch.nn.Module,beta: float = 0.9999) -> None源码

简介

该类型上的公开 method。

参数

src_modeltorch.nn.Module
tgt_modeltorch.nn.Module
betafloat: 默认值: 0.9999

返回值: None

methcache(parameters: Any, is_cpu: bool = False) -> None源码

简介

该类型上的公开 method。

参数

parametersAny
is_cpubool: 默认值: False

返回值: None

methrestore(parameters: Any) -> None源码

简介

该类型上的公开 method。

参数

parametersAny

返回值: None

`find_split`

func

def find_split(shape_tensor: torch.Size,cp_size: int,patch_values: tuple[int, int, int] = (1, 2, 2),view_factor: int = 1) -> torch.Size

worldfoundry.core.distributed.find_splitfrom worldfoundry.core.distributed import find_split

源码

简介

find_split — Find the post-context-parallel temporal/spatial split shape. 属于 Core 分布式辅助（集合通信与 context-parallel 切分/聚合）。标注返回类型：torch.Size。

参数

shape_tensortorch.Size
cp_sizeint
patch_valuestuple[int, int, int]: 默认值: (1, 2, 2)
view_factorint: 默认值: 1

返回值: torch.Size

`get_local_tensor_if_dtensor`

func

def get_local_tensor_if_dtensor(tensor)

worldfoundry.core.get_local_tensor_if_dtensorfrom worldfoundry.core import get_local_tensor_if_dtensor

源码

简介

get_local_tensor_if_dtensor — Return the local shard for DTensor inputs; leave regular tensors unchanged. 属于 Core 分布式辅助（集合通信与 context-parallel 切分/聚合）。

参数

tensor

`init_pp_scheduler`

func

def init_pp_scheduler()

worldfoundry.core.init_pp_schedulerfrom worldfoundry.core import init_pp_scheduler

源码

简介

初始化 pipeline-parallel 运行使用的 PPScheduler 单例。

异常

AssertionError: If the PPScheduler is already initialized.

`initialize_model_parallel`

func

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') -> None

worldfoundry.core.initialize_model_parallelfrom worldfoundry.core import initialize_model_parallel

源码

简介

initialize_model_parallel — Initialize model data parallel groups. 属于 Core 分布式辅助（集合通信与 context-parallel 切分/聚合）。标注返回类型：None。

参数

tp_sizeint: The number of GPUs to split individual tensors across.默认值: 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.默认值: 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.默认值: 1
nccl_communicator_config_pathOptional[str]: Path to the yaml file of NCCL communicator configurations. min_ctas, max_ctas, and cga_cluster_size can be set for each communicator.默认值: None
distributed_timeout_minutesint: Timeout, in minutes,for operations executed against distributed process groups. See PyTorch documentation at默认值: 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.默认值: 'tp-cp-pp-dp'

返回值: None

`is_last_rank`

func

def is_last_rank()

worldfoundry.core.is_last_rankfrom worldfoundry.core import is_last_rank

源码

简介

is_last_rank — Return whether this worker is the final rank in the global process group. 属于 Core 分布式辅助（集合通信与 context-parallel 切分/聚合）。

`is_last_tp_cp_rank`

func

def is_last_tp_cp_rank()

worldfoundry.core.is_last_tp_cp_rankfrom worldfoundry.core import is_last_tp_cp_rank

源码

简介

is_last_tp_cp_rank — Return whether this worker is last in the combined tensor/context group. 属于 Core 分布式辅助（集合通信与 context-parallel 切分/聚合）。

`model_parallel_is_initialized`

func

def model_parallel_is_initialized()

worldfoundry.core.model_parallel_is_initializedfrom worldfoundry.core import model_parallel_is_initialized

源码

简介

model_parallel_is_initialized — Check if model and data parallel groups are initialized. 属于 Core 分布式辅助（集合通信与 context-parallel 切分/聚合）。

`pp_scheduler`

func

def pp_scheduler() -> PPScheduler

worldfoundry.core.pp_schedulerfrom worldfoundry.core import pp_scheduler

源码

简介

获取当前 pipeline-parallel 调度器实例。

异常

AssertionError: If the PPScheduler has not been initialized.

返回值: PPScheduler — PPScheduler: The current PPScheduler instance.

`PPScheduler`

cls

class PPScheduler()

worldfoundry.core.PPSchedulerfrom worldfoundry.core import PPScheduler

源码

简介

PPScheduler — Minimal point-to-point scheduler for adjacent pipeline stages. 属于 Core 分布式辅助（集合通信与 context-parallel 切分/聚合）。

源码 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`.

方法

methisend_next(tensor: torch.Tensor) -> torch.distributed.Work源码

简介

isend_next — Asynchronously send a tensor to the next pipeline and return the send handle.

参数

tensortorch.Tensor: The tensor to be sent.

返回值: torch.distributed.Work — torch.distributed.Work: The handle for the send operation.

methirecv_prev(buffer: torch.Tensor) -> torch.distributed.Work源码

简介

irecv_prev — Asynchronously receive a tensor from the previous pipeline and return the receive handle.

参数

buffertorch.Tensor: The buffer tensor for receiving data.

返回值: torch.distributed.Work — torch.distributed.Work: The handle for the receive operation.

methrecv_prev_data(shape: torch.Size, dtype: torch.dtype) -> torch.Tensor源码

简介

recv_prev_data — Receive data from the previous pipeline and return the received tensor.

参数

shapetorch.Size: The shape of the tensor to receive.
dtypetorch.dtype: The data type of the tensor to receive.

返回值: torch.Tensor — torch.Tensor: The received tensor.

methqueue_irecv_prev(shape: torch.Size, dtype: torch.dtype) -> None源码

简介

queue_irecv_prev — Put the asynchronously received tensor and handle into the receive queue.

参数

shapetorch.Size: The shape of the tensor to receive.
dtypetorch.dtype: The data type of the tensor to receive.

返回值: None

methqueue_irecv_prev_data() -> torch.Tensor源码

简介

queue_irecv_prev_data — Get a tensor from the receive queue and wait for the receive operation to complete.

返回值: torch.Tensor — torch.Tensor: The received tensor obtained from the queue.

`print_per_rank`

func

def print_per_rank(message: object) -> None

worldfoundry.core.print_per_rankfrom worldfoundry.core import print_per_rank

源码

简介

print_per_rank — Emit one informational log record from every calling rank. 属于 Core 分布式辅助（集合通信与 context-parallel 切分/聚合）。标注返回类型：None。

参数

messageobject

返回值: None

`print_rank_0`

func

def print_rank_0(message: object) -> None

worldfoundry.core.print_rank_0from worldfoundry.core import print_rank_0

源码

简介

print_rank_0 — Emit one informational record on rank zero, or in a non-distributed process. 属于 Core 分布式辅助（集合通信与 context-parallel 切分/聚合）。标注返回类型：None。

参数

messageobject

返回值: None

`split_inputs_cp`

func

def split_inputs_cp(x: Tensor,seq_dim: int,cp_group: ProcessGroup | None = None) -> Tensor

worldfoundry.core.distributed.split_inputs_cpfrom worldfoundry.core.distributed import split_inputs_cp

源码

简介

split_inputs_cp — Slice a tensor along `seq_dim to this rank's CP shard. 属于 Core 分布式辅助（集合通信与 context-parallel 切分/聚合）。标注返回类型：Tensor`。

参数

xTensor: Input tensor.
seq_dimint: Dimension to split along (negative indexing supported).
cp_groupProcessGroup | None: CP process group; `None returns x` unchanged.默认值: None

异常

AssertionError: `seq_dim` is not divisible by the CP size.

返回值: Tensor — Contiguous slice of length `x.shape[seq_dim] // cp_size`.