From f2d1b7c4ad8533f8dac3c375d12cfc16db92c899 Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Fri, 27 Feb 2026 08:11:15 +0800 Subject: [PATCH] Support sequence parallel and trivial dims in dump comparator (#19460) --- .../comparator/aligner/unsharder/planner.py | 8 +++++- python/sglang/srt/debug_utils/dumper.py | 12 +++++++++ .../aligner/unsharder/test_parallel_info.py | 16 ++++++++++++ .../aligner/unsharder/test_planner.py | 25 +++++++++++++++++-- 4 files changed, 58 insertions(+), 3 deletions(-) diff --git a/python/sglang/srt/debug_utils/comparator/aligner/unsharder/planner.py b/python/sglang/srt/debug_utils/comparator/aligner/unsharder/planner.py index e51888012..2f713e212 100644 --- a/python/sglang/srt/debug_utils/comparator/aligner/unsharder/planner.py +++ b/python/sglang/srt/debug_utils/comparator/aligner/unsharder/planner.py @@ -31,9 +31,15 @@ def compute_unsharder_plan( sharded_axis_infos: dict[ParallelAxis, DimSpec] = { spec.parallel: spec for spec in dim_specs if spec.parallel is not None } - sharded_axes: set[ParallelAxis] = set(sharded_axis_infos) + sharded_axes_raw: set[ParallelAxis] = set(sharded_axis_infos) all_axes: set[ParallelAxis] = {axis for info in parallel_infos for axis in info} + + # axis annotated in dims but absent from all parallel_infos -> axis_size=1, skip + sharded_axes: set[ParallelAxis] = sharded_axes_raw & all_axes + sharded_axis_infos = { + k: v for k, v in sharded_axis_infos.items() if k in sharded_axes + } replicated_axes: set[ParallelAxis] = all_axes - sharded_axes if not sharded_axes and not replicated_axes: diff --git a/python/sglang/srt/debug_utils/dumper.py b/python/sglang/srt/debug_utils/dumper.py index 1cafcdede..3b7563f5f 100644 --- a/python/sglang/srt/debug_utils/dumper.py +++ b/python/sglang/srt/debug_utils/dumper.py @@ -1238,6 +1238,18 @@ class _MegatronPlugin(_FrameworkPlugin): except (AttributeError, AssertionError): info["megatron_error"] = True + # Megatron sequence parallel reuses the TP group (no dedicated parallel state API). + # When sequence_parallel=True, inject sp_rank/sp_size for the comparator unsharder. + try: + from megatron.training.global_vars import get_args + + args = get_args() + if getattr(args, "sequence_parallel", False) and "tp_rank" in info: + info["sp_rank"] = info["tp_rank"] + info["sp_size"] = info["tp_size"] + except (ImportError, AssertionError, AttributeError): + pass + return info def convert_value( diff --git a/test/registered/debug_utils/comparator/aligner/unsharder/test_parallel_info.py b/test/registered/debug_utils/comparator/aligner/unsharder/test_parallel_info.py index 85060c5b0..3c42117a5 100644 --- a/test/registered/debug_utils/comparator/aligner/unsharder/test_parallel_info.py +++ b/test/registered/debug_utils/comparator/aligner/unsharder/test_parallel_info.py @@ -54,6 +54,22 @@ class TestNormalizeParallelInfo: with pytest.raises(ValueError, match="multiple parallel_info"): normalize_parallel_info(meta) + def test_megatron_with_sp(self) -> None: + """Megatron SP reuses TP group: sp_rank==tp_rank, sp_size==tp_size.""" + meta = { + "megatron_parallel_info": { + "tp_rank": 1, + "tp_size": 4, + "sp_rank": 1, + "sp_size": 4, + } + } + result = normalize_parallel_info(meta) + assert result == { + ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=4), + ParallelAxis.SP: AxisInfo(axis_rank=1, axis_size=4), + } + def test_size_1_filtered(self) -> None: meta = { "sglang_parallel_info": { diff --git a/test/registered/debug_utils/comparator/aligner/unsharder/test_planner.py b/test/registered/debug_utils/comparator/aligner/unsharder/test_planner.py index beaeaa98d..a80ca5738 100644 --- a/test/registered/debug_utils/comparator/aligner/unsharder/test_planner.py +++ b/test/registered/debug_utils/comparator/aligner/unsharder/test_planner.py @@ -38,10 +38,13 @@ class TestComputeUnsharderPlan: with pytest.raises(ValueError, match="Inconsistent axis_size"): compute_unsharder_plan(dim_specs, parallel_infos) - def test_missing_axis_in_parallel_info_raises(self) -> None: + def test_missing_axis_in_all_parallel_infos_skipped(self) -> None: + """Axis in dims but absent from all parallel_infos -> axis_size=1, auto-skip.""" dim_specs = parse_dims("h(tp)") parallel_infos = [{ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2)}] - with pytest.raises(ValueError, match="missing parallel_info"): + # TP not in any parallel_info → skipped; CP is replicated but only 1 rank + # with size=2 → incomplete coverage + with pytest.raises(ValueError, match="axis_rank coverage"): compute_unsharder_plan(dim_specs, parallel_infos) def test_empty_parallel_infos_raises(self) -> None: @@ -232,6 +235,24 @@ class TestComputeUnsharderPlan: assert len(plans[2].groups) == 1 assert len(plans[2].groups[0]) == 2 + def test_sp_in_dims_but_not_in_parallel_info(self) -> None: + """s(sp) in dims but SP absent from parallel_info (SP disabled), should auto-skip.""" + dim_specs = parse_dims("s(sp) b h(tp)") + parallel_infos = [ + {ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2)}, + {ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2)}, + ] + plans = compute_unsharder_plan(dim_specs, parallel_infos) + assert len(plans) == 1 + assert plans[0].axis == ParallelAxis.TP + + def test_all_dims_sharded_but_single_gpu(self) -> None: + """Single GPU (TP=1, CP=1), dims has s(cp) h(tp) but parallel_info is empty.""" + dim_specs = parse_dims("b s(cp) h(tp) d") + parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [{}] + plans = compute_unsharder_plan(dim_specs, parallel_infos) + assert plans == [] + def test_sharded_axis_missing_from_rank_raises(self) -> None: """A world_rank missing a sharded axis raises ValueError.""" dim_specs = parse_dims("s(cp) h(tp)")