并行推理指南

本文档详细介绍 TeleFuser 的分布式并行推理架构，包括原理介绍、配置方法和使用示例。 Tensor 数据流、同步协议、transport 职责和性能约束见通信架构。

概述

TeleFuser 提供多维度并行推理能力，支持以下并行策略：

并行类型	描述	适用场景
数据并行 (DP)	复制模型到多个 GPU，并行处理不同数据	吞吐量优化
CFG 并行	并行计算 positive/negative prompt	CFG 加速
序列并行 (SP)	将长序列分割到多个 GPU	长视频生成
流水线并行 (PP)	将模型层分割到多个 GPU	大模型推理
张量并行 (TP)	将张量维度分割到多个 GPU	大模型推理

架构设计

Device Mesh 布局

TeleFuser 使用 PyTorch DeviceMesh 管理分布式并行，维度顺序为：

DP -> CFG -> SP (ring, ulysses) -> PP -> TP

from telefuser.distributed import create_device_mesh_from_config
from telefuser.core.config import ParallelConfig

config = ParallelConfig(
    device_ids=[0, 1, 2, 3],
    dp_degree=1,
    cfg_degree=2,
    sp_ulysses_degree=2,
    sp_ring_degree=1,
    pp_degree=1,
    tp_degree=1,
)

device_mesh = create_device_mesh_from_config(config)

核心模块

telefuser/distributed/
├── device_mesh.py      # DeviceMesh 创建和进程组管理
├── collectives.py      # 共享的连续 gather 与 reduction 原语
├── pp_comm.py          # 流水线并行 P2P 通信
├── ulysses_comm.py     # Ulysses All-to-All 通信原语
├── ring.py             # Ring Attention P2P 通信
├── parallel_shard.py   # 序列并行张量分片/反分片
├── vae_spatial.py      # 按高度分片的 VAE halo 交换
├── fsdp.py             # FSDP 数据并行
└── tp_parallelize.py   # 张量并行工具

模型代码只负责张量布局和重建语义；可复用的 collective buffer 分配与 reduction 提交统一放在 collectives.py，策略专用协议仍保留在各自模块中。

序列并行

序列并行用于处理超长序列（如长视频），将序列维度分割到多个 GPU。

Ulysses Attention

基于 All-to-All 通信的序列并行：

原理： 1. 每个 GPU 持有序列的一部分 2. 通过 All-to-All 将 heads 重新分配 3. 每个 GPU 拥有完整序列的部分 heads 4. 本地计算注意力后，再通过 All-to-All 恢复

数据流：

输入: (B, S_LOCAL, H_GLOBAL, D)
  -> All-to-All QKV -> (B, S_GLOBAL, H_LOCAL, D)
  -> 本地注意力计算
  -> All-to-All O -> (B, S_LOCAL, H_GLOBAL, D)

特点： - 通信开销：2次 All-to-All（QKV + Output） - 适合中等长度序列 - 需要头数能被 GPU 数整除

在 PyTorch 2.11 及以上版本中，TeleFuser 优先使用 PyTorch Symmetric Memory 分配并 rendezvous 成组 Q/K/V scatter 的对端 target buffer，再由源码编译的 tf-kernel Copy Engine 算子通过一条高优先级 stream 直接 写入这些由 PyTorch 管理的映射。Q、K、V 保持独立提交以保留 projection overlap，同时共享一次 grouped handshake。communicator 及其缓存 buffer 由模型持有，因此模型 offload 或销毁时即可释放资源，无需 worker 在 process-group teardown 阶段执行额外 collective。

主包仍支持 torch>=2.6.0：能力检测采用延迟方式，较旧 PyTorch 会优先回退到源码编译的 CUDA IPC backend， 不可用时再回退到 PyTorch/NCCL。单次 collective 和 output gather 继续使用实测更快的 PyTorch/NCCL 路径。 设计、执行时间线、H100 测量、限制和 Related Work 见 CUDA IPC Ulysses 技术文章。

Ring Attention

基于 P2P 通信的序列并行：

原理： 1. 每个 GPU 持有 Q 的一部分和 K/V 的一部分 2. K/V 在 GPU 环中轮转 3. 每个 GPU 依次看到所有 K/V 块 4. 使用在线 softmax 合并注意力输出

算法流程：

for step in range(world_size):
    # 1. 计算当前 KV 块的注意力
    out, lse = attention(q, k, v)

    # 2. 发送当前 KV 到下一个 GPU
    # 3. 从上一个 GPU 接收新的 KV
    next_k, next_v = send_recv_kv(k, v)

    # 4. 使用在线 softmax 合并结果
    out, lse = merge_attn_states(prev_out, prev_lse, out, lse)

    # 5. 更新 KV
    k, v = next_k, next_v

特点： - 支持任意长度序列 - 通信与计算可重叠 - 需要支持 log-sum-exp 的注意力实现

USP (Unified Sequence Parallelism)

Ulysses + Ring 组合策略，支持更大规模并行：

原理： 1. Ring 维度：序列分割 2. Ulysses 维度：heads 分割 3. 两种策略互补，支持更多 GPU

配置示例：

# 4 GPU: ring=2, ulysses=2
config = ParallelConfig(
    device_ids=[0, 1, 2, 3],
    sp_ring_degree=2,
    sp_ulysses_degree=2,
)

异步 Ulysses (async_usp_forward)

异步 All-to-All 实现，重叠计算和通信：

# 发起异步 All-to-All
q_wait = ulysses_scatter_heads(q, group, tag="q", barrier=False)
k_wait = ulysses_scatter_heads(k, group, tag="k", barrier=False)
v_wait = ulysses_scatter_heads(v, group, tag="v")

# 等待完成
q = q_wait()
k = k_wait()
v = v_wait()

# 计算注意力
x = attention(q, k, v)

# 异步 All-to-All 输出
out_wait = ulysses_gather_heads(x, group, num_heads=num_heads)
out = out_wait()

流水线并行 (PP)

将模型层分割到多个 GPU，实现大模型推理。

跨 Worker Tensor 通道

相邻 stage 如果属于不同 ParallelWorker group，可以使用 WorkerTensorChannel 连接。CPU tensor 使用 multiprocessing shared memory；CUDA tensor 对每个稳定 tensor profile 使用两个有界的 producer-owned IPC slot，每个 IPC allocation 只打开一次并持续复用。Pool handle 与各 rank 的 offset 作为私有 WorkerTensorRef 元数据沿现有控制路径传输。双向复用的跨进程 CUDA event 既保证 producer staging 完成后 consumer 才 copy，也保证所有 consumer copy 完成后 producer 才覆盖 slot。Consumer 先记录 completion event 再发布 generation ACK，producer staging stream 等待这些 event；整个过程不需要 device-wide synchronization。

from telefuser.worker import ParallelWorker, WorkerTensorChannel

latent_channel = WorkerTensorChannel(consumer_world_size=vae_parallel_config.world_size)
denoise_worker = ParallelWorker(
    denoise_stage,
    tensor_output_channel=latent_channel,
    tensor_output_methods=("denoise",),
)
vae_worker = ParallelWorker(vae_stage, tensor_input_channels=(latent_channel,))

Channel 不会自动创建。拥有 stage 拓扑的 pipeline 必须创建并持有每条 channel，把它传给两端的 ParallelWorker，并在 consumer 和 producer worker 停止后关闭。Actor 或 scheduler edge 只携带 WorkerTensorRef metadata；支持取消的 actor 流程必须通过 consumer worker 释放末端 reference。完整的跨 stage wiring 与生命周期 contract 参见 通信架构。

当 consumer rank 只处理互不重叠的 tensor slice 时可设置 shard_dim。Producer 只 staging 一次，各 rank 仅 copy 自己的 slice。LingBot 空间并行 VAE 使用 shard_dim=-2，因此跨卡 peer-copy 聚合通信量保持为一个 logical latent，而不会随 VAE world size 放大；计入 producer 本地 staging 后，设备搬运预算为两个 logical latent。每个 channel 最多池化八种稳定 CUDA tensor profile；更多动态 profile 会回退到 PyTorch CUDA IPC， 避免 retained HBM 无界增长。

该路径是单 producer、单 consumer group 的点对点 FIFO。只有 tensor 的完整 consumer 集合就是所连接的 worker group 时才能启用。需要在主进程读取 tensor 的调用可以传入 _tensor_transport=False。Consumer 必须保持 producer 顺序；关闭时先停止 consumer、再停止 producer，最后关闭 channel。Scheduler 取消末端 artifact 时，必须按 producer 顺序调用 consumer worker 的 discard_tensor_refs(ref, sync=True)。该操作不会在父进程物化 tensor，并会 立即释放 CPU shared memory 或 CUDA IPC storage；正常 receive 仍会兜底丢弃更早的已取消 FIFO entry。

常规 worker 派发同样只向各 rank 发送 shared-memory 或 CUDA IPC handle，最终 device placement 由接收 rank 完成；主进程不再为每张目标 GPU 分配临时副本。

可使用本地 SGLang checkout 在相同 GPU 和 tensor shape 上运行端到端延迟门禁：

python tools/validation/benchmark_tensor_channel_vs_sglang.py

该比较同时计入两种实现的 producer staging、元数据传输、target copy、target 同步与 slot ACK。默认门禁使用 200 个样本：p50 不得比 SGLang 高 5% 以上；p95 上限取 10% 与 0.05 ms 中较宽者，以覆盖亚毫秒级 multiprocessing 调度抖动。

原理

Stage 0 (GPU 0):  Embedding + Layers [0:N/4]
       ↓ send hidden states
Stage 1 (GPU 1):  Layers [N/4:N/2]
       ↓ send hidden states
Stage 2 (GPU 2):  Layers [N/2:3N/4]
       ↓ send hidden states
Stage 3 (GPU 3):  Layers [3N/4:N] + Head
       ↓ output

P2P 通信

from telefuser.distributed import PipelineP2PComm, get_pp_group

pp_group = get_pp_group(device_mesh)
comm = PipelineP2PComm(pp_group)

# 发送隐藏状态到下一阶段
if not comm.is_last_stage:
    comm.send_latent(hidden_states)

# 从上一阶段接收隐藏状态
if not comm.is_first_stage:
    hidden_states = comm.recv_latent(shape=latent_shape)

层分配

# 自动分配层到各阶段
start_idx, end_idx = comm.get_stage_indices(num_layers)

# 示例：40 层，4 阶段
# Stage 0: [0:10]
# Stage 1: [10:20]
# Stage 2: [20:30]
# Stage 3: [30:40]

PP Forward 实现

def pp_forward(self, x, timestep, context, latent_shape, **kwargs):
    # 第一阶段：Embedding + 首批层
    if self.is_pp_first_stage:
        x = self.patch_embedding(x)
        x, grid_size = self.patchify(x)
        x = self.forward_blocks_pp(x, timestep, context, **kwargs)
        self.pp_comm.send_latent(x)
        return None

    # 中间阶段：接收 + 处理 + 发送
    elif not self.is_pp_last_stage:
        x = self.pp_comm.recv_latent(shape=latent_shape)
        x = self.forward_blocks_pp(x, timestep, context, **kwargs)
        self.pp_comm.send_latent(x)
        return None

    # 最后阶段：接收 + 处理 + 输出
    else:
        x = self.pp_comm.recv_latent(shape=latent_shape)
        x = self.forward_blocks_pp(x, timestep, context, **kwargs)
        x = self.head(x)
        return x

CFG 并行

将 Classifier-Free Guidance 的 positive/negative prompt 并行计算：

原理

标准 CFG 计算：

noise_pred = noise_neg + cfg_scale * (noise_pos - noise_neg)

CFG 并行将 positive 和 negative 分配到不同 GPU：

GPU 0: 计算 noise_pos (positive prompt)
GPU 1: 计算 noise_neg (negative prompt)
然后：All-Gather 合并结果

使用方法

# 2 GPU CFG 并行
config = ParallelConfig(
    device_ids=[0, 1],
    cfg_degree=2,
)

# 在模型中启用
dit.enable_cfgp()

实现

# 分片
cfg_parallel_shard(device_mesh, [x, timestep, context, ...])

# 根据 CFG rank 选择 positive/negative
cond_flag = False if get_cfg_rank(device_mesh) else True

# 计算
output = model(x, context, cond_flag=cond_flag)

# 合并
output = cfg_parallel_unshard(device_mesh, [output])[0]

数据并行 (DP)

使用 FSDP 进行数据并行训练/推理：

FSDP1

from telefuser.distributed.fsdp import shard_model

model = shard_model(
    model,
    device_id=device_id,
    sharding_strategy=ShardingStrategy.FULL_SHARD,
    wrap_module_names=["blocks"],
    param_dtype=torch.bfloat16,
)

FSDP2

from telefuser.distributed.fsdp import shard_model_fsdp2

model = shard_model_fsdp2(
    model,
    wrap_module_names=["blocks"],
    param_dtype=torch.bfloat16,
)

张量并行 (TP)

将张量维度分割到多个 GPU：

使用方法

from telefuser.distributed.tp_parallelize import parallelize_module
from torch.distributed.tensor.parallel import ColwiseParallel, RowwiseParallel

tp_plan = {
    "self_attn.q": ColwiseParallel(),
    "self_attn.k": ColwiseParallel(),
    "self_attn.v": ColwiseParallel(),
    "self_attn.o": RowwiseParallel(),
    "ffn.0": ColwiseParallel(),
    "ffn.2": RowwiseParallel(),
}

model = parallelize_module(model, device_mesh, tp_plan)

注意事项

• 在所选 pipeline 支持时，SP 和 TP 可以作为独立的 device mesh 维度同时启用
• 需要确保头数能被 TP 度数整除

Worker 实现

ParallelWorker

多进程并行 Worker，使用 multiprocessing.spawn：

from telefuser.worker import ParallelWorker

worker = ParallelWorker(stage)

# 调用方法
result = worker.process(latents, ...)

# 关闭
del worker

特点： - 每个 GPU 一个进程 - 自动初始化进程组 - 支持同步/异步调用

RayWorker

Ray 集群分布式 Worker：

from telefuser.worker import create_ray_worker

worker = create_ray_worker(stage, enable_parallel=True)
result = worker.process.remote(latents, ...)

配置说明

ParallelConfig

@dataclass
class ParallelConfig:
    device_ids: list | None = None        # GPU ID 列表
    dp_degree: int = 1                    # 数据并行度
    cfg_degree: int = 1                   # CFG 并行度
    sp_ulysses_degree: int = 1            # Ulysses 序列并行度
    sp_ring_degree: int = 1               # Ring 序列并行度
    pp_degree: int = 1                    # 流水线并行度
    tp_degree: int = 1                    # 张量并行度
    enable_fsdp: bool = False             # 启用 FSDP
    timeout: int = 600                    # 超时时间（秒）
    queue_with_cpu: bool = False          # 使用 CPU 队列
    worker_intra_op_threads: int = 1      # 每个 worker 的 CPU intra-op 线程数

默认值 1 与 torchrun 一致，可避免多个 GPU worker 的 CPU 线程过度竞争。只有分布式 stage 包含大量 CPU 计算时才需要提高该值；它不会修改父进程的线程池。

验证规则

# 设备数必须等于各并行度乘积
world_size = dp * cfg * sp_ring * sp_ulysses * pp * tp

# SP 和 TP 可以作为独立的 mesh 维度组合。
# 所选 pipeline 必须实现并校验请求的组合。

使用示例

单 GPU 推理

from telefuser.core.config import ParallelConfig

config = ParallelConfig(device_ids=[0])

2 GPU Ulysses 序列并行

config = ParallelConfig(
    device_ids=[0, 1],
    sp_ulysses_degree=2,
)
pipe_config.dit_config.parallel_config = config
pipe_config.enable_denoising_parallel = True

4 GPU CFG + Ulysses

config = ParallelConfig(
    device_ids=[0, 1, 2, 3],
    cfg_degree=2,
    sp_ulysses_degree=2,
)

4 GPU USP

config = ParallelConfig(
    device_ids=[0, 1, 2, 3],
    sp_ring_degree=2,
    sp_ulysses_degree=2,
)

4 GPU 流水线并行

config = ParallelConfig(
    device_ids=[0, 1, 2, 3],
    pp_degree=4,
)

8 GPU 混合并行

# DP=2, CFG=2, Ulysses=2
config = ParallelConfig(
    device_ids=[0, 1, 2, 3, 4, 5, 6, 7],
    dp_degree=2,
    cfg_degree=2,
    sp_ulysses_degree=2,
)

Wan Video 示例

from telefuser.pipelines.wan_video.wan21_video import (
    Wan21VideoPipeline,
    Wan21VideoPipelineConfig,
)
from telefuser.core.config import AttentionConfig, AttnImplType

# 创建 Pipeline
pipe = Wan21VideoPipeline(device="cuda", torch_dtype=torch.bfloat16)
pipe_config = Wan21VideoPipelineConfig()

# 配置并行
if gpu_num > 1:
    pipe_config.dit_config.parallel_config.device_ids = list(range(gpu_num))
    pipe_config.dit_config.parallel_config.sp_ulysses_degree = 2
    pipe_config.enable_denoising_parallel = True

# 配置注意力
pipe_config.dit_config.attention_config = AttentionConfig.dense_attention(
    AttnImplType.FLASH_ATTN_2
)

# 初始化
pipe.init(module_manager, pipe_config)

# 推理
video = pipe(
    prompt="A stylish girl playing with her dog",
    num_inference_steps=40,
    num_frames=81,
    cfg_scale=6.0,
)

Device Mesh 工具函数

from telefuser.distributed import (
    # 进程组
    get_dp_group, get_dp_rank, get_dp_world_size,
    get_cfg_group, get_cfg_rank, get_cfg_world_size,
    get_ulysses_group, get_ulysses_rank, get_ulysses_world_size,
    get_ring_group, get_ring_rank, get_ring_world_size,
    get_pp_group, get_pp_rank, get_pp_world_size,
    get_tp_group, get_tp_rank, get_tp_world_size,

    # PP 辅助
    is_pipeline_first_stage,
    is_pipeline_last_stage,

    # 策略检测
    get_attention_strategy,  # 返回 "local", "ulysses", "ring", "usp"

    # 通信
    ulysses_scatter_heads,
    ulysses_gather_heads,
    RingP2PComm,
    PipelineP2PComm,
    merge_attn_states,
    ring_attention_forward,
)

# 获取当前注意力策略
strategy = get_attention_strategy(device_mesh)
# "local": 无序列并行
# "ulysses": 仅 Ulysses
# "ring": 仅 Ring
# "usp": Ulysses + Ring 组合

性能优化建议

选择并行策略

场景	推荐策略	说明
短视频 (81帧)	单 GPU 或 CFG=2	通信开销小
中等视频 (161帧)	Ulysses=2	All-to-All 效率高
长视频 (241+帧)	Ring 或 USP	支持任意长度
大模型 (14B)	PP 或 FSDP	模型分割

FSDP vs TP 选择

FSDP 和 TP 以不同方式切分模型权重，推理通信的扩展规律也不同。FSDP 会 all-gather 参数单元，TP 会在 张量并行算子后归约激活张量；不能只根据模型大小或 sequence 长度判断哪一种更快。

条件	候选策略	原因
最大 FSDP wrapping unit 加激活可放入单卡	FSDP 可行	FSDP 在 all-gather 时必须完整 materialize 一个 wrapped unit。
单卡无法放下一个 wrapped unit	需要 TP 或 PP	仅切分 FSDP 常驻 shard 并不能消除完整单元的 all-gather 峰值。
参数 all-gather 主导暴露通信时间	TP	TP 让参数 shard 常驻，转而通信激活。
激活归约主导暴露通信时间，尤其 B * S * H 较大时	FSDP 或 TP + SP	FSDP 参数流量大多与 sequence 长度无关；SP 会降低 TP 的局部 sequence，但自身也会引入通信。
TP group 位于高速 NVLink/NVSwitch 域内	TP 值得优先评估	大量对延迟敏感的激活 collective 可受益于低延迟、高带宽互联。

第一轮可以按每张卡、每个 denoising step 估算通信字节数。设 FSDP degree 为 f，每个 wrapped unit 的完整参数字节数为 P_u，每个请求的 all-gather 次数为 r_u：

V_fsdp ~= sum_u(r_u * P_u * (f - 1) / f)

设 TP degree 为 t、SP degree 为 s、激活元素字节数为 e_a，一次 TP collective j 处理的激活 形状近似为 [B, S / s, H]：

V_tp ~= sum_j(2 * (t - 1) / t * B * (S / s) * H * e_a)

该求和通常包括每个 Transformer block 中的 row-parallel reduction。Ulysses 或 Ring 流量必须单独加入； SP 虽然降低 TP 的局部激活大小，但并不是免费的。应在目标拓扑上把字节数转换成实际决策：

T_comm ~= N_collectives * latency + V_wire / effective_bandwidth
T_step ~= T_compute + exposed(T_comm after overlap)

在实际模型、sequence 长度、dtype、denoising step 数和 mesh degree 下，采集 p50/p95 step time、原始和 暴露通信时间、collective 次数及每卡峰值 HBM。FSDP 可与 PP、SP 组合；所选 pipeline 已实现并校验时，TP 和 SP 也可以组合。

通信优化

1. 使用异步通信：async_usp_forward 重叠计算和通信
2. 批量通信：batch_isend_irecv 减少通信次数
3. FP8 量化：减少通信数据量

内存优化

1. 序列并行：减少每个 GPU 的序列长度
2. 流水线并行：减少每个 GPU 的层数
3. CPU Offload：将权重卸载到 CPU

故障排除

设备数不匹配

RuntimeError: device num 4 and world size 2 not match

解决方案：确保 len(device_ids) == dp * cfg * sp_ring * sp_ulysses * pp * tp

Ring Attention 需要 LSE

RuntimeError: Ring attention requires log-sum-exp from attention implementation

解决方案：使用支持 LSE 的注意力实现（Flash Attention ⅔/4）。

通信超时

RuntimeError: ParallelWorker timeout

解决方案：增加 timeout 参数值，或检查网络连接。

参考资料

• Ulysses: Sequence Parallelism
• Ring Attention
• GPipe: Pipeline Parallelism
• PyTorch Distributed