From 3ebd85bf1c9d9e259778354f9e3dc9a80867ba9a Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Mon, 2 Mar 2026 18:46:03 +0800 Subject: [PATCH] Enhance sglang engine dumping tests in dump comparator (#19681) --- .../debug_utils/comparator/conftest.py | 23 ++ .../test_engine_dumper_comparator_e2e.py | 330 ++++++++++++------ 2 files changed, 249 insertions(+), 104 deletions(-) diff --git a/test/registered/debug_utils/comparator/conftest.py b/test/registered/debug_utils/comparator/conftest.py index f356008e3..4b032ede5 100644 --- a/test/registered/debug_utils/comparator/conftest.py +++ b/test/registered/debug_utils/comparator/conftest.py @@ -1,7 +1,30 @@ +import warnings + +warnings.filterwarnings( + "ignore", message="builtin type Swig.*", category=DeprecationWarning +) + import pytest from sglang.srt.debug_utils.comparator.output_types import report_sink +collect_ignore_glob: list[str] = [] + + +def pytest_configure(config: pytest.Config) -> None: + config.addinivalue_line( + "filterwarnings", + "ignore:Unknown config option. asyncio_mode:pytest.PytestConfigWarning", + ) + config.addinivalue_line( + "filterwarnings", + "ignore:builtin type Swig.*:DeprecationWarning", + ) + config.addinivalue_line( + "filterwarnings", + "ignore:Named tensors and all their associated APIs:UserWarning", + ) + @pytest.fixture(autouse=True) def _reset_report_sink() -> None: diff --git a/test/registered/debug_utils/test_engine_dumper_comparator_e2e.py b/test/registered/debug_utils/test_engine_dumper_comparator_e2e.py index 05d46761d..49c301ff0 100644 --- a/test/registered/debug_utils/test_engine_dumper_comparator_e2e.py +++ b/test/registered/debug_utils/test_engine_dumper_comparator_e2e.py @@ -1,9 +1,13 @@ """E2E test: source patcher + dumper + comparator on SGLang server. -Patches Qwen3DecoderLayer.forward to insert dumper.dump() calls, -launches 1-GPU baseline and 2-GPU TP=2 target servers, runs inference, -verifies patched dump fields exist, then runs comparator to verify -numerical consistency. +Patches Qwen3MoeDecoderLayer.forward (and related methods) to insert +dumper.dump() calls at 7 points, launches servers with Qwen3-30B-A3B +(MOE model), runs inference, verifies patched dump fields exist, then +runs comparator to verify numerical consistency. + +Test cases: +- test_patch_dump_and_compare: TP=2 baseline vs TP=4 target +- test_dp_attention: TP=2 baseline vs TP=2+DP=2+dp-attention target The dumper.apply_source_patches() auto-injects ``from ... import dumper`` so the YAML only needs ``dumper.dump(...)`` calls. @@ -13,15 +17,15 @@ import os import subprocess import tempfile from pathlib import Path +from typing import Optional import pytest import requests -from sglang.srt.debug_utils.comparator.output_types import ( - AnyRecord, - SummaryRecord, - parse_record_json, +pytestmark = pytest.mark.filterwarnings( + "ignore:Unknown config option. asyncio_mode:pytest.PytestConfigWarning", ) + from sglang.srt.utils import kill_process_tree from sglang.test.ci.ci_register import register_cuda_ci from sglang.test.test_utils import ( @@ -30,115 +34,233 @@ from sglang.test.test_utils import ( popen_launch_server, ) -register_cuda_ci(est_time=120, suite="nightly-2-gpu", nightly=True) +register_cuda_ci(est_time=300, suite="nightly-4-gpu", nightly=True) -MODEL = "Qwen/Qwen3-0.6B" +MODEL = "Qwen/Qwen3-30B-A3B" +BASELINE_TP = 2 +TARGET_TP = 4 EXP_NAME = "e2e_source_patcher" DUMPER_FILTER = "layer_id in [0, 1, 2]" +_FIELDS_TO_VERIFY: list[str] = [ + # decoder layer level (aligned with miles) + "layer_input", + "attn_output", + "pre_mlp_residual", + "mlp_output", + # attention internals + "attn_pre_o_proj", + # moe internals + "moe_router_logits", + "moe_expert_output", +] + PATCH_CONFIG_YAML: str = """\ patches: - - target: sglang.srt.models.qwen3.Qwen3DecoderLayer.forward + # --- decoder layer level (aligned with miles test) --- + - target: sglang.srt.models.qwen3_moe.Qwen3MoeDecoderLayer.forward edits: - - match: "hidden_states = self.mlp(hidden_states)" - prepend: "dumper.dump('patched_attn_output', hidden_states, dims='t h')" - - match: "return hidden_states, residual" - prepend: "dumper.dump('patched_mlp_output', hidden_states, dims='t h')" + - match: | + hidden_states, residual = ( + self.layer_communicator.prepare_attn_and_capture_last_layer_outputs( + hidden_states, + residual, + forward_batch, + captured_last_layer_outputs=captured_last_layer_outputs, + **kwargs, + ) + ) + append: "dumper.dump('layer_input', hidden_states, dims='t h # tp:replicated')" + - match: | + hidden_states = self.self_attn( + positions=positions, + hidden_states=hidden_states, + forward_batch=forward_batch, + ) + append: "dumper.dump('attn_output', hidden_states, dims='t h[tp:partial]')" + - match: | + hidden_states, residual = self.layer_communicator.prepare_mlp( + hidden_states, residual, forward_batch + ) + append: "dumper.dump('pre_mlp_residual', hidden_states, dims='t h # tp:replicated')" + - match: | + hidden_states = self.mlp( + hidden_states, forward_batch, should_allreduce_fusion, use_reduce_scatter + ) + append: "dumper.dump('mlp_output', hidden_states, dims='t h[tp:partial]')" + + # --- attention internals --- + - target: sglang.srt.models.qwen3_moe.Qwen3MoeAttention.forward_core + edits: + - match: "output, _ = self.o_proj(attn_output)" + prepend: "dumper.dump('attn_pre_o_proj', attn_output, dims='t attn_h[tp]')" + + # --- moe internals --- + - target: sglang.srt.models.qwen3_moe.Qwen3MoeSparseMoeBlock.forward_normal + edits: + - match: "router_logits, _ = self.gate(hidden_states)" + append: "dumper.dump('moe_router_logits', router_logits, dims='t num_experts # tp:replicated')" + - match: "final_hidden_states = self.experts(hidden_states, topk_output)" + append: "dumper.dump('moe_expert_output', final_hidden_states, dims='t h[tp:partial]')" +""" + +PATCH_CONFIG_DP_ATTENTION_YAML: str = """\ +patches: + # --- decoder layer level (aligned with miles test) --- + # In dp-attention mode: attn tensors are NOT TP-sharded (attn_tp_size=1), + # and mlp_output is already all-reduced inside forward_normal(). + # layer_input is dumped after prepare_attn which DP-distributes tokens, + # so it needs dp:=attn_dp to filter to the non-empty DP rank. + - target: sglang.srt.models.qwen3_moe.Qwen3MoeDecoderLayer.forward + edits: + - match: | + hidden_states, residual = ( + self.layer_communicator.prepare_attn_and_capture_last_layer_outputs( + hidden_states, + residual, + forward_batch, + captured_last_layer_outputs=captured_last_layer_outputs, + **kwargs, + ) + ) + append: "dumper.dump('layer_input', hidden_states, dims='t h # tp:replicated dp:=attn_dp')" + - match: | + hidden_states = self.self_attn( + positions=positions, + hidden_states=hidden_states, + forward_batch=forward_batch, + ) + append: "dumper.dump('attn_output', hidden_states, dims='t h # tp:replicated')" + - match: | + hidden_states, residual = self.layer_communicator.prepare_mlp( + hidden_states, residual, forward_batch + ) + append: "dumper.dump('pre_mlp_residual', hidden_states, dims='t h # tp:replicated')" + - match: | + hidden_states = self.mlp( + hidden_states, forward_batch, should_allreduce_fusion, use_reduce_scatter + ) + append: "dumper.dump('mlp_output', hidden_states, dims='t h # tp:replicated')" + + # --- attention internals --- + - target: sglang.srt.models.qwen3_moe.Qwen3MoeAttention.forward_core + edits: + - match: "output, _ = self.o_proj(attn_output)" + prepend: "dumper.dump('attn_pre_o_proj', attn_output, dims='t attn_h # tp:replicated')" + + # --- moe internals --- + - target: sglang.srt.models.qwen3_moe.Qwen3MoeSparseMoeBlock.forward_normal + edits: + - match: "router_logits, _ = self.gate(hidden_states)" + append: "dumper.dump('moe_router_logits', router_logits, dims='t num_experts # tp:replicated')" + - match: "final_hidden_states = self.experts(hidden_states, topk_output)" + append: "dumper.dump('moe_expert_output', final_hidden_states, dims='t h[tp:partial]')" """ class TestSourcePatcherE2ESGLang: - """E2E: patch Qwen3 forward -> dump -> compare 1gpu vs 2gpu-tp2.""" + """E2E: patch Qwen3Moe forward -> dump -> compare.""" - @pytest.mark.timeout(300) def test_patch_dump_and_compare(self, tmp_path: Path) -> None: - patched_fields: list[str] = ["patched_attn_output", "patched_mlp_output"] - base_url: str = DEFAULT_URL_FOR_TEST - - config_path: Path = tmp_path / "patch_config.yaml" - config_path.write_text(PATCH_CONFIG_YAML) - - # Run 1: baseline (1 GPU) - baseline_dir: Path = tmp_path / "baseline" - _run_server_and_generate( - dump_dir=baseline_dir, - config_path=config_path, - tp=1, - base_url=base_url, - ) - _verify_patched_fields(dump_dir=baseline_dir, field_names=patched_fields) - - # Run 2: target (2 GPU TP=2) - target_dir: Path = tmp_path / "target" - _run_server_and_generate( - dump_dir=target_dir, - config_path=config_path, - tp=2, - base_url=base_url, - ) - _verify_patched_fields(dump_dir=target_dir, field_names=patched_fields) - - # Compare baseline vs target - baseline_exp: Path = baseline_dir / EXP_NAME - target_exp: Path = target_dir / EXP_NAME - - result: subprocess.CompletedProcess[str] = subprocess.run( - [ - "python", - "-m", - "sglang.srt.debug_utils.comparator", - "--baseline-path", - str(baseline_exp), - "--target-path", - str(target_exp), - "--output-format", - "json", - "--grouping", - "logical", - ], - capture_output=True, - text=True, + """TP=2 baseline vs TP=4 target.""" + _run_e2e_scenario( + tmp_path=tmp_path, + target_tp=TARGET_TP, ) - debug_file: Path = _save_comparator_output( - stdout=result.stdout, stderr=result.stderr - ) - print(f"Comparator debug output: {debug_file}") + def test_dp_attention(self, tmp_path: Path) -> None: + """TP=2 baseline vs TP=2+DP=2+dp-attention target. - assert result.returncode == 0, ( - f"Comparator failed (rc={result.returncode}). " - f"Debug output: {debug_file}" - ) - - records: list[AnyRecord] = [ - parse_record_json(line) - for line in result.stdout.strip().splitlines() - if line.strip() - ] - assert ( - len(records) > 0 - ), f"Comparator produced no output records. Debug: {debug_file}" - - summary: SummaryRecord = _find_summary(records=records, debug_file=debug_file) - assert ( - summary.passed > 0 - ), f"No comparisons passed (total={summary.total}). Debug: {debug_file}" - assert summary.failed == 0, ( - f"{summary.failed} comparisons failed " - f"(passed={summary.passed}, skipped={summary.skipped}). " - f"Debug: {debug_file}" + In dp-attention mode (attn_tp_size=1, attn_dp_size=2), attention + tensors are NOT TP-sharded and mlp_output is already all-reduced. + A separate patch config with corrected dims is used for the target. + """ + _run_e2e_scenario( + tmp_path=tmp_path, + target_tp=BASELINE_TP, + extra_target_server_args=["--dp", "2", "--enable-dp-attention"], + target_patch_config_yaml=PATCH_CONFIG_DP_ATTENTION_YAML, ) # --------------------------------- helpers --------------------------------- +def _run_e2e_scenario( + *, + tmp_path: Path, + target_tp: int, + extra_target_server_args: Optional[list[str]] = None, + target_patch_config_yaml: Optional[str] = None, +) -> None: + """Full e2e: write patch config -> baseline run -> target run -> compare.""" + base_url: str = DEFAULT_URL_FOR_TEST + + baseline_config_path: Path = tmp_path / "patch_config.yaml" + baseline_config_path.write_text(PATCH_CONFIG_YAML) + + target_config_path: Path = tmp_path / "patch_config_target.yaml" + target_config_path.write_text(target_patch_config_yaml or PATCH_CONFIG_YAML) + + baseline_dir: Path = tmp_path / "baseline" + _run_server_and_generate( + dump_dir=baseline_dir, + config_path=baseline_config_path, + tp=BASELINE_TP, + base_url=base_url, + ) + _verify_patched_fields(dump_dir=baseline_dir, field_names=_FIELDS_TO_VERIFY) + + target_dir: Path = tmp_path / "target" + _run_server_and_generate( + dump_dir=target_dir, + config_path=target_config_path, + tp=target_tp, + base_url=base_url, + extra_server_args=extra_target_server_args, + ) + _verify_patched_fields(dump_dir=target_dir, field_names=_FIELDS_TO_VERIFY) + + baseline_exp: Path = baseline_dir / EXP_NAME + target_exp: Path = target_dir / EXP_NAME + + cmd: list[str] = [ + "python", + "-m", + "sglang.srt.debug_utils.comparator", + "--baseline-path", + str(baseline_exp), + "--target-path", + str(target_exp), + "--output-format", + "json", + "--allow-skipped-pattern", + "input_ids|positions", + ] + + result: subprocess.CompletedProcess[str] = subprocess.run( + cmd, + capture_output=True, + text=True, + ) + + debug_file: Path = _save_comparator_output( + stdout=result.stdout, stderr=result.stderr + ) + print(f"Comparator debug output: {debug_file}") + + assert result.returncode == 0, ( + f"Comparator failed (rc={result.returncode}). " f"Debug output: {debug_file}" + ) + + def _run_server_and_generate( *, dump_dir: Path, config_path: Path, tp: int, base_url: str, + extra_server_args: Optional[list[str]] = None, ) -> None: """Launch SGLang server with source patcher + dumper, send a generate request.""" env: dict[str, str] = { @@ -149,11 +271,24 @@ def _run_server_and_generate( "DUMPER_SERVER_PORT": "reuse", } + server_args: list[str] = [ + "--tp", + str(tp), + "--max-total-tokens", + "128", + "--mem-fraction-static", + "0.5", + "--disable-cuda-graph", + "--disable-radix-cache", + ] + if extra_server_args: + server_args.extend(extra_server_args) + proc = popen_launch_server( MODEL, base_url, timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, - other_args=["--tp", str(tp), "--max-total-tokens", "128"], + other_args=server_args, env=env, ) try: @@ -170,7 +305,7 @@ def _run_server_and_generate( f"{base_url}/generate", json={ "text": "The capital of France is", - "sampling_params": {"max_new_tokens": 8}, + "sampling_params": {"max_new_tokens": 1, "temperature": 0}, }, ) assert resp.status_code == 200, f"Generate failed: {resp.text}" @@ -188,19 +323,6 @@ def _verify_patched_fields(*, dump_dir: Path, field_names: list[str]) -> None: ) -def _find_summary(*, records: list[AnyRecord], debug_file: Path) -> SummaryRecord: - """Extract the SummaryRecord from comparator output.""" - summaries: list[SummaryRecord] = [ - r for r in records if isinstance(r, SummaryRecord) - ] - assert len(summaries) == 1, ( - f"Expected 1 summary record, got {len(summaries)}. " - f"Record types: {[type(r).__name__ for r in records]}. " - f"Debug: {debug_file}" - ) - return summaries[0] - - def _save_comparator_output(*, stdout: str, stderr: str) -> Path: """Save comparator stdout+stderr to a temp file that persists for debugging.""" fd, path_str = tempfile.mkstemp(prefix="comparator_e2e_", suffix=".log", dir="/tmp")