Support sequence parallel and trivial dims in dump comparator (#19460)

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:

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(

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": {

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)")