Core 分布式

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

本页内容

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 个公开符号

def broadcast(item: torch.Tensor | str | None,process_group: ProcessGroup | None = None) -> torch.Tensor | str | None
worldfoundry.core.distributedfrom 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
def broadcast_dtensor_model_states(model: torch.nn.Module,mesh: DeviceMesh) -> None
worldfoundry.corefrom 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
def broadcast_split_tensor(tensor: torch.Tensor | None,seq_dim: int,process_group: ProcessGroup | None = None) -> torch.Tensor | None
worldfoundry.core.distributedfrom 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
def cat_outputs_cp(x: Tensor,seq_dim: int,cp_group: ProcessGroup | None = None) -> Tensor
worldfoundry.core.distributedfrom 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默认值 None
CP process group; `None returns x` unchanged.

返回值

TensorTensor with the gathered shards concatenated along `seq_dim`.

异常

RuntimeError
`all_gather` failed.
def cat_outputs_cp_with_grad(x: Tensor,seq_dim: int,cp_group: ProcessGroup | None = None) -> Tensor
worldfoundry.core.distributedfrom 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
def destroy_model_parallel(destroy_process_groups: bool = True) -> None
worldfoundry.corefrom worldfoundry.core import destroy_model_parallel

简介

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

源码 docstring

Destroy owned subgroups and clear all model-parallel state.

The default process group is deliberately left alive; its lifecycle is owned by the distributed launcher. Set `destroy_process_groups=False` only when an external framework has already destroyed the subgroups.

参数

destroy_process_groupsbool默认值 True

返回值

None
def dist_init(config)
worldfoundry.corefrom 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. CC-17 WONTFIX: not a wrapper of init_torch_distributed. Backend and timeout come from config, device binding is rank % device_count (not LOCAL_RANK`), and this function also creates CP/PP groups plus the PP scheduler.

worldfoundry.corefrom 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源码

copy_to — Overwrite target local shards with source local shards (no broadcast).

参数

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

简介

update_average — In-place EMA on local FP32 shards: `tgt = beta * tgt + (1 - beta) * src`.

源码 docstring

In-place EMA on local FP32 shards: `tgt = beta * tgt + (1 - beta) * src`.

EMA in a lower dtype silently underflows the `1 - beta` term, so the target must already be FP32.

参数

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

返回值

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

简介

cache — Snapshot local shards before a temporary overwrite (eval / checkpoint).

源码 docstring

Snapshot local shards before a temporary overwrite (eval / checkpoint).

Nested :meth:cache without :meth:restore is refused so a later restore cannot apply the wrong generation of weights.

参数

parametersAny
is_cpubool默认值 False

返回值

None
methrestore(parameters: Any) -> None源码

restore — Write cached local shards back; `strict=False` zip matches FSDP holes.

参数

parametersAny

返回值

None
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.distributedfrom 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
worldfoundry.corefrom 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
worldfoundry.corefrom worldfoundry.core import init_pp_scheduler

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

异常

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') -> None
worldfoundry.corefrom worldfoundry.core import initialize_model_parallel

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

参数

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

返回值

None
worldfoundry.corefrom 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 切分/聚合)。

worldfoundry.corefrom 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 切分/聚合)。

worldfoundry.corefrom worldfoundry.core import model_parallel_is_initialized

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

worldfoundry.corefrom worldfoundry.core import pp_scheduler

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

返回值

PPSchedulerPPScheduler: The current PPScheduler instance.

异常

AssertionError
If the PPScheduler has not been initialized.
class PPScheduler()
worldfoundry.corefrom 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.Worktorch.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.Worktorch.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.Tensortorch.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.Tensortorch.Tensor: The received tensor obtained from the queue.
def print_per_rank(message: object) -> None
worldfoundry.corefrom worldfoundry.core import print_per_rank

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

参数

messageobject

返回值

None
def print_rank_0(message: object) -> None
worldfoundry.corefrom 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
def split_inputs_cp(x: Tensor,seq_dim: int,cp_group: ProcessGroup | None = None) -> Tensor
worldfoundry.core.distributedfrom 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默认值 None
CP process group; `None returns x` unchanged.

返回值

TensorContiguous slice of length `x.shape[seq_dim] // cp_size`.

异常

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