SGLang 混合并行

SGLang 从 0.4.10 开始支持 Hybrid TP + EP (混合并行），也就是 EP 内再做 TP ，参考：[Feature] Hybrid EP and TP by ch-wan · Pull Request #8590 · sgl-project/sglang

1. 背景知识简述

在大规模语言模型（LLM）训练/推理中，处于一些考量，比如单个 GPU 无法容纳整个模型或者为了减少冗余参数对显存的占用，通常会采用以下并行策略：

Tensor Parallelism (TP)：将模型的权重（如矩阵乘法）在多个 GPU 上切分，每个 GPU 只计算一部分。
Expert Parallelism (EP)：在 MoE 架构中，不同的"专家"（expert subnetworks）被分配到不同 GPU 上，每个 token 只激活部分专家。
混合并行：例如 MoE 内可能同时使用 TP（在每个 expert 内部做张量并行）和 EP（跨 expert 分布）。

为了实现这些并行策略，框架会创建多个 通信组（process groups），每个组内的 GPU 可以高效通信（如 AllReduce、AllGather 等）。

2. 混合并行参数

SGLang 从 0.4.10 开始有 6 个与分布式并行计算有关的参数和对应的获取接口。

2.1. 并行参数获取

python 复制代码

def get_tensor_model_parallel_world_size():
    """Return world size for the tensor model parallel group."""
    return get_tp_group().world_size


def get_tensor_model_parallel_rank():
    """Return my rank for the tensor model parallel group."""
    return get_tp_group().rank_in_group


def get_moe_expert_parallel_world_size():
    """Return world size for the moe expert parallel group."""
    return get_moe_ep_group().world_size


def get_moe_expert_parallel_rank():
    """Return my rank for the moe expert parallel group."""
    return get_moe_ep_group().rank_in_group


def get_moe_tensor_parallel_world_size():
    """Return world size for the moe tensor parallel group."""
    return get_moe_tp_group().world_size


def get_moe_tensor_parallel_rank():
    """Return my rank for the moe tensor parallel group."""
    return get_moe_tp_group().rank_in_group

2.2. 代码逐行解释

假设目前单机内有 8 张 GPU，某个 MoE 模型有 16 个专家。

2.2.1. 张量模型并行（Tensor Model Parallelism）

python 复制代码

def get_tensor_model_parallel_world_size():
    """Return world size for the tensor model parallel group."""
    return get_tp_group().world_size

返回 张量并行组 中的 GPU 数量（即 world size）。
get_tp_group() 返回一个表示张量并行通信组的对象，.world_size 是该组中进程总数。

python 复制代码

def get_tensor_model_parallel_rank():
    """Return my rank for the tensor model parallel group."""
    return get_tp_group().rank_in_group

返回当前进程在张量并行组中的 本地 rank（编号），从 0 开始。
用于判断"我是这个 TP 组中的第几个 GPU"。

例如：如果 TP=8，那么这 8 个 GPU 组成一个 TP 组，每个 GPU 的 rank_in_group 是 0~7。

2.2.2. MoE 专家并行（Expert Parallelism）

python 复制代码

def get_moe_expert_parallel_world_size():
    return get_moe_ep_group().world_size

返回 MoE 专家并行组 的大小（即有多少 GPU 共享或分发 experts）。

python 复制代码

def get_moe_expert_parallel_rank():
    return get_moe_ep_group().rank_in_group

返回当前进程在 MoE 专家并行组中的 rank。
用于决定"我负责哪些 experts"。

举例：如果有 16 个 experts，EP=4，那么每个 EP 负责 4 个 experts。

2.2.3. MoE 张量并行（Tensor Parallel inside MoE）

python 复制代码

def get_moe_tensor_parallel_world_size():
    return get_moe_tp_group().world_size
def get_moe_tensor_parallel_rank():
    return get_moe_tp_group().rank_in_group

这些函数用于 MoE 内部的张量并行。
在 MoE 架构中，每个 expert 本身可能很大，因此 每个 expert 内部也可以做张量并行。
moe_tp_group 是专为 MoE 中的张量并行创建的通信组，可能与主模型的 TP 组不同（例如，主模型 TP=8，EP=4，moe_tp_size = TP/EP = 2）。

注意：有些框架会将 MoE 的 TP 和主模型 TP 合并，但 SGLang 这里显式区分了 tp_group 和 moe_tp_group，说明支持独立配置。

2.3. 参数之间的限制

SGLang 内部对 tp_size、ep_size 等参数有一些限制

tp_size >= ep_size
tp_size % ep_size == 0
num_physical_experts % ep_size == 0
moe_ep_rank = tp_rank // (server_args.tp_size // server_args.ep_size)

3. 总结

函数	用途
`get_tensor_model_parallel_world_size/rank`	主模型的张量并行组信息
`get_moe_expert_parallel_world_size/rank`	MoE 专家并行组信息
`get_moe_tensor_parallel_world_size/rank`	MoE 内部每个 expert 的TP张量并行组信息

SGLang 这些函数为分布式计算提供了 上下文感知：每个 GPU 可以知道自己在不同并行维度中的角色，从而正确地进行通信、切分权重、路由 tokens（MoE）等操作。