0
代码
代码
Issues
Pull Requests
讨论
Wiki
分析
0
ascend-robotascend-robotdocs:阶段大模型检查
74429164创建于 6 天前历史提交

SGLang精度数据采集（SGLang版本<0.5.11）

简介

msProbe工具通过在SGLang框架中负责模型前向传播执行的核心类ModelRunner中添加PrecisionDebugger接口并启动推理的方式， 采集模型在运行过程中的精度数据。

dump "statistics"模式的性能膨胀大小与"tensor"模式采集的数据量大小，可以参考dump基线。

注意：

  • 本文档仅适用于SGLang版本<0.5.11的数据采集，需通过侵入式修改SGLang源码的方式开启msProbe工具能力。若当前SGLang版本>=0.5.11，此版本及更高版本已原生内置msProbe工具，可直接指定参数--msprobe-dump-config进行精度数据采集，具体操作请参考文档《SGLang精度数据采集（SGLang版本>=0.5.11）》。

  • 采集数据前，需要指定SGLang框架的--disable-cuda-graph参数关闭图模式。

  • 使用SGLang框架的在线模式采集数据，需要指定SGLang框架的--skip-server-warmup参数关闭warmup，避免采集到warmup阶段的数据。

  • 如果遇到dynamo相关报错，可设置环境变量export TORCHDYNAMO_DISABLE=1全局关闭dynamo。

  • 使用SGLang框架的PD分离模式采集数据，Router启动时会发送/health请求，/health会触发模型forward，需要设置环境变量export SGLANG_ENABLE_HEALTH_ENDPOINT_GENERATION=0，这样/health请求只返回200，不会触发模型forward，避免msProbe采集不需要的/health请求阶段的数据。

  • 本工具提供固定的API支持列表，若需要删除或增加dump的API，可以在support_wrap_ops.yaml文件内手动修改，如下示例：

    functional:  # functional为算子类别，找到对应的类别，在该类别下按照下列格式删除或添加API
  - conv1d
  - conv2d
  - conv3d

    删除API的场景：部分模型代码逻辑会存在API原生类型校验，工具执行dump操作时，对模型的API封装可能与模型的原生API类型不一致，此时可能引发校验失败，详见《FAQ》中“异常情况”的第10条。

使用前准备

环境准备

安装msProbe工具，详情请参见《msProbe安装指南》。

约束

仅支持采集基于PyTorch框架实现的模型，暂不支持PyTorch版本>=2.7的dynamo场景。

快速入门

以下通过一个简单的示例，展示如何在SGLang框架中使用msProbe工具进行精度数据采集。

  1. 配置文件创建

    在当前目录下创建config.json文件，用于配置dump参数。内容示例如下：

      {
    "task": "statistics",
    "dump_path": "/home/data_dump",
    "rank": [],
    "step": [],
    "level": "mix",
    "async_dump": false,
    "statistics": {
      "scope": [],
      "list": [],
      "data_mode": [
        "all"
      ],
      "summary_mode": "statistics"
    }
  }

    config.json文件详细介绍请参见《配置文件介绍》。

  2. SGLang框架中使能msProbe工具

    找到SGLang框架ModelRunner类所属文件：sglang/srt/model_executor/model_runner.py

    • ModelRunner类的__init__方法中添加PrecisionDebugger接口，传入config.json文件路径。

      class ModelRunner(ModelRunnerKVCacheMixin):
    """ModelRunner runs the forward passes of the models."""
     
    def __init__(
        self,
        model_config: ModelConfig,
        mem_fraction_static: float,
        gpu_id: int,
        tp_rank: int,
        ...
    ):
        
        ################################ msprobe ################################
        from msprobe.pytorch import PrecisionDebugger, seed_all
        seed_all(mode=True)
        self.debugger = PrecisionDebugger(config_path="./config.json")
        ################################ msprobe ################################
        # Parse args
        self.mem_fraction_static = mem_fraction_static
        self.device = server_args.device
        self.gpu_id = gpu_id
        self.tp_rank = tp_rank
        self.tp_size = tp_size
      ...

    • ModelRunner类的forward方法中添加start、stop和step接口。

      SGLang框架默认场景:

      class ModelRunner(ModelRunnerKVCacheMixin):
    """ModelRunner runs the forward passes of the models."""
    
    def __init__(
        self,
        model_config: ModelConfig,
        mem_fraction_static: float,
        gpu_id: int,
        tp_rank: int,
        ...
    ):
        
    def forward(
        self,
        forward_batch: ForwardBatch,
        skip_attn_backend_init: bool = False,
        pp_proxy_tensors: Optional[PPProxyTensors] = None,
        reinit_attn_backend: bool = False,
        split_forward_count: int = 1,
    ) -> ModelRunnerOutput:
        self.forward_pass_id += 1
        ################################ msprobe ################################
        if hasattr(self, 'debugger'):
            self.debugger.start(model=self.model)
        ################################ msprobe ################################
        
        ...
      
        ################################ msprobe ################################ 
        if hasattr(self, 'debugger'):
            self.debugger.stop()
            self.debugger.step()
        ################################ msprobe ################################
    
      return output

    SGLang框架启用DP场景（--dp-size>1），需要在start接口中配置rank_id参数：

     if hasattr(self, 'debugger'):
    self.debugger.start(model=self.model, rank_id=self.gpu_id)

  3. 启动SGLang框架模型推理，开始采集数据

    • 在线模式

      1. 启动服务

        #!/bin/bash
export TORCHDYNAMO_DISABLE=1

python3 -m sglang.launch_server \
 --model-path Qwen/Qwen2.5-0.5B-Instruct \
 --host 127.0.0.1 \
 --port 1027 \
 --disable-cuda-graph \
 --skip-server-warmup

      2. 发送请求，自动开始dump

         curl -H "Content-type: application/json" \
  -X POST \
  -d '{
      "model": "Qwen/Qwen2.5-0.5B-Instruct",
      "messages": [
          {
              "role": "user",
              "content": "Hello, my name is"
          }
      ],
      "max_tokens": 10
  }' \
  http://127.0.0.1:1027/v1/chat/completions

    • 离线模式

      离线脚本示例如下，运行将自动开始dump：

      import os
import asyncio
  
import sglang as sgl
import sglang.test.doc_patch
from sglang.utils import async_stream_and_merge, stream_and_merge
  
def main():
    llm = sgl.Engine(model_path="Qwen/Qwen2.5-0.5B-Instruct", disable_cuda_graph=True)
  
    prompts = [
        "Hello, my name is",
        "The president of the United States is",
        "The capital of France is",
        "The future of AI is"
    ]
  
    sampling_params = {"temperature": 0.8, "top_p": 0.95}
  
    outputs = llm.generate(prompts, sampling_params)
    for prompt, output in zip(prompts, outputs):
        print("===============================")
        print(f"Prompt: {prompt}\nGenerated text: {output['text']}")
  
  
if __name__ == '__main__':
    main()

PD分离场景数据采集

以下通过一个简单的示例，展示如何在SGLang框架PD分离场景中使用msProbe工具进行精度数据采集。

  1. 配置文件创建

    在当前目录下创建config文件，用于配置dump参数。

    • 采集prefill和decode阶段的数据，则创建config_prefill.json和config_decode.json。其中两个json配置文件的"dump_path"要不同，避免dump写入冲突。
    • 仅采集prefill阶段的数据，则创建config_prefill.json。
    • 仅采集decode阶段的数据，则创建config_decode.json。

    内容示例如下：

      {
    "task": "statistics",
    "dump_path": "/home/data_dump",
    "rank": [],
    "step": [],
    "level": "mix",
    "async_dump": false,
    "statistics": {
      "scope": [],
      "list": [],
      "data_mode": [
        "all"
      ],
      "summary_mode": "statistics"
    }
  }

    config文件详细介绍请参见《配置文件介绍》。

  2. SGLang框架中使能msProbe工具

    1. 找到SGLang框架ModelRunner类所属文件：sglang/srt/model_executor/model_runner.py

      ModelRunner类的__init__方法中添加PrecisionDebugger接口，传入config_prefill.json或config_decode.json文件路径。

      基于传入的config文件，可选择采集prefill和decode阶段、仅prefill阶段或仅decode阶段。

      class ModelRunner(ModelRunnerKVCacheMixin):
    """ModelRunner runs the forward passes of the models."""
    
    def __init__(
        self,
        model_config: ModelConfig,
        mem_fraction_static: float,
        gpu_id: int,
        tp_rank: int,
        ...
    ):
        
        ################################ msprobe ################################
        from msprobe.pytorch import PrecisionDebugger, seed_all
        seed_all(mode=True)
        config_path = ""
        disagg_mode = server_args.disaggregation_mode
        if disagg_mode == "prefill":
            config_path = "./config_prefill.json" # 不采集prefill阶段则修改为config_path = ""
        elif disagg_mode == "decode":
            config_path = "./config_decode.json" # 不采集decode阶段则修改为config_path = ""
        if config_path:
            self.debugger = PrecisionDebugger(config_path=config_path)
        ################################ msprobe ################################
        # Parse args
        self.mem_fraction_static = mem_fraction_static
        self.device = server_args.device
        self.gpu_id = gpu_id
        self.tp_rank = tp_rank
        self.tp_size = tp_size
      ...

      ModelRunner类的forward方法中添加start、stop和step接口，请参考快速入门中的添加方式。

    2. SGLang框架启用DP场景（--dp-size>1），需要固定bootstrap_room值。

      • 背景描述

        SGLang开启PD分离部署后，Prefill阶段会强制采用follow_bootstrap_room调度规则，依靠bootstrap_room数值分配请求到不同显卡， 调度规则为目标dp_rank = bootstrap_room % dp_size。SGLang默认Router服务会自动随机生成bootstrap_room数值， 该值目前无法通过启动参数配置。当dp-size>1、一次性批量处理多条请求时，随机数值会导致请求乱序分配在不同显卡上，该问题会造成多次采集数据不一致、实验无法复现等问题。

        因此，需要手动修改bootstrap_room为固定数值，锁定请求的分配显卡，统一调度规则，保障数据采集固定、运行结果可复现。

      • 操作方式

        找到SGLang框架源码文件：sglang/srt/managers/io_struct.py，重新设置bootstrap_room值（考虑到源码使用random.randint(0, 2**63 - 1)生成19位整数，这里也用固定的19位整数）。

        @dataclass
class GenerateReqInput(BaseReq):
    ...

    def _normalize_bootstrap_params(self, num):
        """Normalize bootstrap parameters for batch processing."""

        ...

        # Normalize bootstrap_room
        ################################ msprobe ################################
        self.bootstrap_room = 6347036608774465186
        ################################ msprobe ################################
        if self.bootstrap_room is None:
            self.bootstrap_room = [None] * num
        elif not isinstance(self.bootstrap_room, list):
            self.bootstrap_room = [self.bootstrap_room + i for i in range(num)]
        elif isinstance(self.bootstrap_room, list):
            self.bootstrap_room = self.bootstrap_room * self.parallel_sample_num

        ...

  3. 启动SGLang框架PD分离场景模型推理，开始采集数据。以下示例基于NPU设备。

    1. 启动Prefill服务

      使用msProbe工具dump时，需要添加export SGLANG_ENABLE_HEALTH_ENDPOINT_GENERATION=0，--disable-cuda-graph和--skip-server-warmup。示例如下：

      # Enabling CPU Affinity
export SGLANG_SET_CPU_AFFINITY=1

# 使用msProbe工具dump时，避免Router发送健康检查请求触发模型forward
export SGLANG_ENABLE_HEALTH_ENDPOINT_GENERATION=0

# PIP: recommended to config first Prefill Server IP
# PORT: one free port
# all sglang servers need to be config the same PIP and PORT,
export ASCEND_MF_STORE_URL="tcp://PIP:PORT"
# if you are Atlas 800I A2 hardware and use rdma for kv cache transfer, add this parameter
export ASCEND_MF_TRANSFER_PROTOCOL="device_rdma"
python3 -m sglang.launch_server \
    --model-path /home/models/Qwen2.5-0.5B-Instruct \
    --disaggregation-mode prefill \
    --disaggregation-transfer-backend ascend \
    --disaggregation-bootstrap-port 8995 \
    --attention-backend ascend \
    --device npu \
    --base-gpu-id 0 \
    --tp-size 1 \
    --host 127.0.0.1 \
    --port 8000 \
    --disable-cuda-graph \
    --skip-server-warmup

    2. 启动Decode服务

      使用msProbe工具dump时，需要添加export SGLANG_ENABLE_HEALTH_ENDPOINT_GENERATION=0，--disable-cuda-graph和--skip-server-warmup。示例如下：

      # 使用msProbe工具dump时，避免Router发送健康检查请求触发模型forward
export SGLANG_ENABLE_HEALTH_ENDPOINT_GENERATION=0

# PIP: recommended to config first Prefill Server IP
# PORT: one free port
# all sglang servers need to be config the same PIP and PORT,
export ASCEND_MF_STORE_URL="tcp://PIP:PORT"
# if you are Atlas 800I A2 hardware and use rdma for kv cache transfer, add this parameter
export ASCEND_MF_TRANSFER_PROTOCOL="device_rdma"
python3 -m sglang.launch_server \
    --model-path /home/models/Qwen2.5-0.5B-Instruct \
    --disaggregation-mode decode \
    --disaggregation-transfer-backend ascend \
    --attention-backend ascend \
    --device npu \
    --base-gpu-id 1 \
    --tp-size 1 \
    --host 127.0.0.1 \
    --port 8001 \
    --disable-cuda-graph \
    --skip-server-warmup

    3. 启动路由

      示例如下：

      python3 -m sglang_router.launch_router \
    --pd-disaggregation \
    --policy cache_aware \
    --prefill http://127.0.0.1:8000 8995 \
    --decode http://127.0.0.1:8001 \
    --host 127.0.0.1 \
    --port 6688

    4. 发送请求，自动开始dump

      示例如下：

      curl -H "Content-type: application/json" \
-X POST \
-d '{
    "model": "Qwen/Qwen2.5-0.5B-Instruct",
    "messages": [
        {
            "role": "user",
            "content": "Hello, my name is"
        }
    ],
    "max_tokens": 10
}' \
http://127.0.0.1:6688/v1/chat/completions

数据采集功能介绍

SGLang精度数据采集详细功能以及采集的dump数据结构与PyTorch场景一致，具体请参见《PyTorch场景精度数据采集》。