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 | Nonebroadcast — Broadcast a tensor or object from the minimum rank in `process_group. 属于 Core 分布式辅助(集合通信与 context-parallel 切分/聚合)。 标注返回类型:torch.Tensor | str | None`。
参数
itemtorch.Tensor | str | Noneprocess_groupProcessGroup | None默认值 None返回值
def broadcast_dtensor_model_states(model: torch.nn.Module,mesh: DeviceMesh) -> Nonebroadcast_dtensor_model_states — Broadcast model parameters and buffers from the first rank in the replicate mesh. 属于 Core 分布式辅助(集合通信与 context-parallel 切分/聚合)。 标注返回类型:None。
参数
modeltorch.nn.ModulemeshDeviceMesh返回值
def broadcast_split_tensor(tensor: torch.Tensor | None,seq_dim: int,process_group: ProcessGroup | None = None) -> torch.Tensor | Nonebroadcast_split_tensor — Broadcast a tensor from the minimum CP rank, then return this rank's shard. 属于 Core 分布式辅助(集合通信与 context-parallel 切分/聚合)。 标注返回类型:torch.Tensor | None。
参数
tensortorch.Tensor | Noneseq_dimintprocess_groupProcessGroup | None默认值 None返回值
def cat_outputs_cp(x: Tensor,seq_dim: int,cp_group: ProcessGroup | None = None) -> Tensorcat_outputs_cp — Gather and concatenate per-rank tensors along `seq_dim. 属于 Core 分布式辅助(集合通信与 context-parallel 切分/聚合)。 标注返回类型:Tensor`。
参数
xTensorseq_dimintcp_groupProcessGroup | None默认值 NoneNone returns x` unchanged.返回值
seq_dim`.异常
RuntimeError- `
all_gather` failed.
def cat_outputs_cp_with_grad(x: Tensor,seq_dim: int,cp_group: ProcessGroup | None = None) -> Tensorcat_outputs_cp_with_grad — Gather CP shards while preserving the local rank's autograd graph. 属于 Core 分布式辅助(集合通信与 context-parallel 切分/聚合)。 标注返回类型:Tensor。
参数
xTensorseq_dimintcp_groupProcessGroup | None默认值 None返回值
def destroy_model_parallel(destroy_process_groups: bool = True) -> None简介
清空已保存的进程组句柄,拆除 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返回值
def dist_init(config)dist_init — Initialize torch.distributed plus WorldFoundry CP/PP groups. 属于 Core 分布式辅助(集合通信与 context-parallel 切分/聚合)。
参数
configengine_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.
class DTensorFastEmaModelUpdater()DTensorFastEmaModelUpdater — Foreach-based EMA updater that operates on local DTensor shards. 属于 Core 分布式辅助(集合通信与 context-parallel 切分/聚合)。
方法
copy_to — Overwrite target local shards with source local shards (no broadcast).
参数
src_modeltorch.nn.Moduletgt_modeltorch.nn.Module返回值
update_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.Moduletgt_modeltorch.nn.Modulebetafloat默认值 0.9999返回值
简介
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.
参数
parametersAnyis_cpubool默认值 False返回值
restore — Write cached local shards back; `strict=False` zip matches FSDP holes.
参数
parametersAny返回值
def find_split(shape_tensor: torch.Size,cp_size: int,patch_values: tuple[int, int, int] = (1, 2, 2),view_factor: int = 1) -> torch.Sizefind_split — Find the post-context-parallel temporal/spatial split shape. 属于 Core 分布式辅助(集合通信与 context-parallel 切分/聚合)。 标注返回类型:torch.Size。
参数
shape_tensortorch.Sizecp_sizeintpatch_valuestuple[int, int, int]默认值 (1, 2, 2)view_factorint默认值 1返回值
def get_local_tensor_if_dtensor(tensor)get_local_tensor_if_dtensor — Return the local shard for DTensor inputs; leave regular tensors unchanged. 属于 Core 分布式辅助(集合通信与 context-parallel 切分/聚合)。
参数
tensordef 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') -> Noneinitialize_model_parallel — Initialize model data parallel groups. 属于 Core 分布式辅助(集合通信与 context-parallel 切分/聚合)。 标注返回类型:None。
参数
tp_sizeint默认值 1pp_sizeint默认值 1cp_sizeint默认值 1nccl_communicator_config_pathOptional[str]默认值 Nonemin_ctas, max_ctas, and cga_cluster_size can be set for each communicator.distributed_timeout_minutesint默认值 30orderstr默认值 'tp-cp-pp-dp'返回值
def is_last_rank()is_last_rank — Return whether this worker is the final rank in the global process group. 属于 Core 分布式辅助(集合通信与 context-parallel 切分/聚合)。
def 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 — Check if model and data parallel groups are initialized. 属于 Core 分布式辅助(集合通信与 context-parallel 切分/聚合)。
def pp_scheduler() -> PPScheduler获取当前 pipeline-parallel 调度器实例。
返回值
异常
AssertionError- If the PPScheduler has not been initialized.
class 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`.
方法
isend_next — Asynchronously send a tensor to the next pipeline and return the send handle.
参数
tensortorch.Tensor返回值
irecv_prev — Asynchronously receive a tensor from the previous pipeline and return the receive handle.
参数
buffertorch.Tensor返回值
recv_prev_data — Receive data from the previous pipeline and return the received tensor.
参数
shapetorch.Sizedtypetorch.dtype返回值
queue_irecv_prev — Put the asynchronously received tensor and handle into the receive queue.
参数
shapetorch.Sizedtypetorch.dtype返回值
queue_irecv_prev_data — Get a tensor from the receive queue and wait for the receive operation to complete.
返回值
def print_per_rank(message: object) -> Noneprint_per_rank — Emit one informational log record from every calling rank. 属于 Core 分布式辅助(集合通信与 context-parallel 切分/聚合)。 标注返回类型:None。
参数
messageobject返回值
def print_rank_0(message: object) -> Noneprint_rank_0 — Emit one informational record on rank zero, or in a non-distributed process. 属于 Core 分布式辅助(集合通信与 context-parallel 切分/聚合)。 标注返回类型:None。
参数
messageobject返回值
def split_inputs_cp(x: Tensor,seq_dim: int,cp_group: ProcessGroup | None = None) -> Tensorsplit_inputs_cp — Slice a tensor along `seq_dim to this rank's CP shard. 属于 Core 分布式辅助(集合通信与 context-parallel 切分/聚合)。 标注返回类型:Tensor`。
参数
xTensorseq_dimintcp_groupProcessGroup | None默认值 NoneNone returns x` unchanged.返回值
x.shape[seq_dim] // cp_size`.异常
AssertionError- `
seq_dim` is not divisible by the CP size.