Skip to main content

parallel_state管理

parallel_state​

初始化_TP 和 _PP通信组​

该函数在model_runner.py中调用初始化TP和PP组，底层通信就是torch调用的nccl + gloo。

sglang的通信大致分为三种：

  • torch下的nccl + gloo，这里的tp、pp计算传递tensor
  • zmq，tokennizer、sheduder、detokennizer之间传递req、kv_indices等
  • mooncake等传递kv_cache，比如在pd分离时的mooncake实现，底层优化了网络传输
def initialize_model_parallel(
    tensor_model_parallel_size: int = 1,
    pipeline_model_parallel_size: int = 1,
    backend: Optional[str] = None,
) -> None:

比如有8个卡，tp = 2(分成了4个组)， pp=4（将8卡分成了两个组）：

4 tensor model-parallel groups:
    [g0, g1], [g2, g3], [g4, g5], [g6, g7]
2 pipeline model-parallel groups:
    [g0, g2, g4, g6], [g1, g3, g5, g7]

这里关于pytorch分布式的用法可以参考文章：Sglang 源码学习笔记（三）- 分布式和并行（以deepseek 为例）（WIP）

init_process_group 必须首先调用：所有分布式操作（包括new_group）必须在默认进程组初始化后进行 全局唯一：每个进程只需调用一次，定义整个分布式训练的基础环境。 new_group 依赖默认进程组：必须在init_process_group之后调用，用于创建子组。 可多次调用：可根据需要创建多个子组，每个子组包含不同的进程集合。

Scatter_mode​

在下面这个enable-dp-attention的配置下，[卡1、卡2、卡3、卡4]，卡1、卡2属于dp1的tp_group， 卡3、卡4属于dp2的tp_group。

# workspace/sglang/python/sglang/srt/layers/communicator.py
class ScatterMode(Enum):
    """
    Suppose we have TP=4, DP=2, enable-dp-attention, and the system handles seq a,b,c,d
    Model input/output: [ab, ab, cd, cd] for four ranks respectively
    SCATTERED: [a, b, c, d]
    TP_ATTN_FULL: [ab, ab, cd, cd], i.e. all ranks inside a TP attn group have full data of the group
    FULL: [abcd, abcd, abcd, abcd]
    """

    SCATTERED = auto()
    TP_ATTN_FULL = auto()
    FULL = auto()

    @staticmethod
    def model_input_output():
        """The scatter mode for model forward pass input and output data"""
        return ScatterMode.TP_ATTN_FULL