From 4d84f886e7acd6504233002076d711acd5c7eb9c Mon Sep 17 00:00:00 2001 From: Yuhong Guo Date: Wed, 5 Nov 2025 19:30:01 +0800 Subject: [PATCH] Refactor `--debug-tensor-dump-layers` to list (#12691) --- .../debug_utils/tensor_dump_forward_hook.py | 23 +++++++++++++++---- python/sglang/srt/server_args.py | 8 +++---- .../test_tensor_dump_forward_hook.py | 4 +++- 3 files changed, 26 insertions(+), 9 deletions(-) diff --git a/python/sglang/srt/debug_utils/tensor_dump_forward_hook.py b/python/sglang/srt/debug_utils/tensor_dump_forward_hook.py index 2a9a1d4c6..66a266075 100644 --- a/python/sglang/srt/debug_utils/tensor_dump_forward_hook.py +++ b/python/sglang/srt/debug_utils/tensor_dump_forward_hook.py @@ -13,6 +13,7 @@ The file contains a series of key-value pairs, where the keys correspond to oper import logging import os from pathlib import Path +from typing import List, Optional import torch @@ -24,7 +25,12 @@ logger = logging.getLogger(__name__) class TensorDumper: def __init__( - self, dump_dir: str, dump_layers: int, tp_size: int, tp_rank: int, pp_rank: int + self, + dump_dir: str, + dump_layers: Optional[List[int]], + tp_size: int, + tp_rank: int, + pp_rank: int, ): self._dump_layers = dump_layers self._forward_pass_id = 0 @@ -94,11 +100,15 @@ class TensorDumper: top_level_model = True else: cur_name = prefix + "." + name - if self._dump_layers > 0 and name.isdigit() and prefix == layers_prefix: + if ( + self._dump_layers is not None + and name.isdigit() + and prefix == layers_prefix + ): # If we only need n layers, skip the reset layers. # Most models' layout is like model.layers.0. cur_layer = int(name) - if cur_layer >= self._dump_layers: + if cur_layer not in self._dump_layers: continue if module is not None: _, sub_count = self._add_hook_recursive( @@ -129,7 +139,12 @@ class TensorDumper: def register_forward_hook_for_model( - model, dump_dir: str, dump_layers: int, tp_size: int, tp_rank: int, pp_rank: int + model, + dump_dir: str, + dump_layers: Optional[List[int]], + tp_size: int, + tp_rank: int, + pp_rank: int, ): tensor_dumper = TensorDumper(dump_dir, dump_layers, tp_size, tp_rank, pp_rank) # Most models have the layerout like: diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 156be80ac..7eebb34f8 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -519,8 +519,8 @@ class ServerArgs: # Debug tensor dumps debug_tensor_dump_output_folder: Optional[str] = None - # -1 mean dump all layers. - debug_tensor_dump_layers: int = -1 + # None means dump all layers. + debug_tensor_dump_layers: Optional[List[int]] = None # TODO(guoyuhong): clean the old dumper code. debug_tensor_dump_input_file: Optional[str] = None debug_tensor_dump_inject: bool = False @@ -3424,8 +3424,8 @@ class ServerArgs: parser.add_argument( "--debug-tensor-dump-layers", type=int, - default=-1, - help="The layer number for dumping tensors.", + nargs="+", + help="The layer ids to dump. Dump all layers if not specified.", ) parser.add_argument( "--debug-tensor-dump-input-file", diff --git a/test/srt/debug_utils/test_tensor_dump_forward_hook.py b/test/srt/debug_utils/test_tensor_dump_forward_hook.py index 6300b7726..c691f0f15 100644 --- a/test/srt/debug_utils/test_tensor_dump_forward_hook.py +++ b/test/srt/debug_utils/test_tensor_dump_forward_hook.py @@ -13,6 +13,7 @@ from sglang.srt.distributed.parallel_state import ( from sglang.srt.layers.layernorm import RMSNorm from sglang.srt.layers.linear import LinearBase from sglang.srt.models.qwen2 import Qwen2MLP +from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler from sglang.srt.utils import add_prefix TEST_HIDDEN_SIZE = 32 @@ -63,6 +64,7 @@ def init_weights(module): def test_model_forward_dump(tmp_path): + set_global_server_args_for_scheduler(ServerArgs(model_path="dummy")) init_distributed_environment( backend="nccl", world_size=1, @@ -75,7 +77,7 @@ def test_model_forward_dump(tmp_path): model.apply(init_weights) model = model.cuda().bfloat16() dumper = register_forward_hook_for_model( - model, tmp_path / "sglang_dump", -1, 0, 0, 0 + model, tmp_path / "sglang_dump", [0], 0, 0, 0 ) dir_path = dumper.get_dump_dir()