SGLang 混合并行

SGLang 从 0.4.10 开始支持 Hybrid TP + EP (混合并行),也就是 EP 内再做 TP ,参考:[Feature Hybrid EP and TP by ch-wan · Pull Request #8590 · sgl-project/sglang](https://link.juejin.cn?target=https%3A%2F%2Fgithub.com%2Fsgl-project%2Fsglang%2Fpull%2F8590 "https://github.com/sgl-project/sglang/pull/8590")

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_groupmoe_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)等操作。

相关推荐
vx-Biye_Design2 小时前
SSM伴侣动物伴护星小程序06330-计算机课程设计、毕业设计
spring boot·后端·elasticsearch·小程序·架构·课程设计·idea
2分钟速写快排4 小时前
什么是 RAG?如何用 RAG 实现一个用户记忆?
前端·后端·ai编程
面向Google编程4 小时前
Apache Iceberg Variant 类型:v3 半结构化数据不再「二选一」
大数据·后端
jason成都5 小时前
Spring WebFlux 适配达梦新方案|dm‑r2dbc:Netty 异步传输的实验性 R2DBC 驱动
java·后端·spring
北岛贰6 小时前
迷茫焦虑期,我做了一个带支付带官网的 AI 聊天虚拟恋人 App
前端·人工智能·后端
IT_陈寒9 小时前
React子组件莫名其妙重渲染?你可能漏了这个Hook
前端·人工智能·后端
小兵金林9 小时前
AI 生成代码的思考与实践
后端·ai编程
码事漫谈9 小时前
DeepSeek 明天又降价(涵历史价格对比)
后端
前端精髓9 小时前
NestJS 是什么(对着 Spring Boot 一起学习)
spring boot·后端·学习
Flynt9 小时前
"只输出 JSON"根本不够:LLM 结构化输出的坑,我实测了 60 次调用
llm·json