Add parallel-rank dump filenames and pipeline-global layer remapping to dumper (#26850)

python/sglang/srt/debug_utils/dumper.py

@@ -155,6 +155,9 @@ class DumperConfig(_BaseConfig):
     # Fully-qualified Python path "pkg.subpkg.module.fn_name"
     # None -> use the default identity-by-rank fallback in _Grafter._default_transform.
     grafter_transform_path: Optional[str] = None
+    # When True, append parallel-rank tags (pp_rank/tp_rank/...) to dump filenames so
+    # tensors from different ranks do not collide when dumped into a shared directory.
+    include_parallel_rank_in_filename: bool = False
 
     @classmethod
     def _env_prefix(cls) -> str:
@@ -312,6 +315,10 @@ class _Dumper:
         **kwargs,
     ) -> None:
         for param_name, param in model.named_parameters():
+            for plugin in _plugins:
+                param_name = (
+                    plugin.transform_model_param_name(model, param_name) or param_name
+                )
             self._dump_inner(
                 name=f"{name_prefix}__{param_name}",
                 value=param,
@@ -567,6 +574,8 @@ class _Dumper:
             dump_index=self._state.dump_index,
             **tags,
         )
+        if self._config.include_parallel_rank_in_filename:
+            full_kwargs.update(_collect_parallel_rank_tags())
         full_filename = _format_tags(full_kwargs) + ".pt"
         path = Path(self._config.dir) / self._config.exp_name / full_filename
 
@@ -1248,6 +1257,26 @@ def _materialize_value(value):
     return value
 
 
+_PARALLEL_RANK_KEYS = ("pp_rank", "tp_rank", "cp_rank", "ep_rank", "etp_rank")
+
+
+def _collect_parallel_rank_tags() -> dict[str, int]:
+    """Collect parallel-rank tags from framework plugins for use in dump filenames.
+
+    Merges the ``_PARALLEL_RANK_KEYS`` reported by each plugin's
+    ``collect_parallel_info()``; the first plugin to report a given key wins.
+    """
+    result: dict[str, int] = {}
+    for plugin in _plugins:
+        info = plugin.collect_parallel_info()
+        if not info:
+            continue
+        for key in _PARALLEL_RANK_KEYS:
+            if key in info and key not in result:
+                result[key] = info[key]
+    return result
+
+
 def _format_tags(kwargs: dict) -> str:
     return "___".join(f"{k}={v}" for k, v in kwargs.items())
 
@@ -1645,6 +1674,16 @@ class _FrameworkPlugin(ABC):
     def detect_recompute_status(self) -> _RecomputeStatus:
         return _RecomputeStatus.DISABLED
 
+    def transform_model_param_name(
+        self, model: "torch.nn.Module", param_name: str
+    ) -> Optional[str]:
+        """Return a rewritten parameter name, or None to keep the original.
+
+        Used by ``dump_model`` to canonicalize parameter names across parallel
+        layouts (e.g. mapping pipeline-local layer indices to global ones).
+        """
+        return None
+
 
 class _SGLangPlugin(_FrameworkPlugin):
     _available = True
@@ -1848,6 +1887,65 @@ class _MegatronPlugin(_FrameworkPlugin):
         except (ImportError, AttributeError):
             return _RecomputeStatus.DISABLED
 
+    def transform_model_param_name(
+        self, model: "torch.nn.Module", param_name: str
+    ) -> Optional[str]:
+        """Rewrite pipeline-local layer indices to global ones in a param name.
+
+        With pipeline parallelism, ``model.named_parameters()`` reports layer
+        indices local to the current PP stage (e.g. ``layers.0`` on every stage).
+        Adding the stage's ``get_transformer_layer_offset`` makes the dumped
+        names globally unique and comparable across stages. Returns None (keep the
+        original name) when not applicable.
+        """
+        if not self._available:
+            return None
+
+        try:
+            pp_size = self._mpu.get_pipeline_model_parallel_world_size()
+        except (AttributeError, AssertionError):
+            return None
+        if pp_size <= 1:
+            return None
+
+        config = self._get_model_config(model)
+        if config is None:
+            return None
+
+        offset = self._get_transformer_layer_offset(config)
+        if not offset:
+            return None
+
+        def _add_offset(match: "re.Match") -> str:
+            return f"layers.{int(match.group(1)) + offset}"
+
+        return re.sub(r"layers\.(\d+)", _add_offset, param_name)
+
+    @staticmethod
+    def _get_transformer_layer_offset(config) -> int:
+        """Return the PP-stage layer offset for ``config``, or 0 if unavailable."""
+        try:
+            from megatron.core.transformer.transformer_layer import (
+                get_transformer_layer_offset,
+            )
+
+            return get_transformer_layer_offset(config)
+        except (ImportError, AttributeError, AssertionError):
+            return 0
+
+    @staticmethod
+    def _get_model_config(model: "torch.nn.Module"):
+        """Unwrap nested ``.module`` wrappers to reach the Megatron model config."""
+        inner = model
+        for _ in range(10):
+            if hasattr(inner, "config"):
+                return inner.config
+            if hasattr(inner, "module"):
+                inner = inner.module
+            else:
+                break
+        return None
+
 
 _plugins: list[_FrameworkPlugin] = [_SGLangPlugin(), _MegatronPlugin()]
 

test/registered/debug_utils/test_dumper.py

@@ -16,6 +16,7 @@ import torch.distributed as dist
 
 from sglang.srt.debug_utils.dumper import (
     DumperConfig,
+    _collect_parallel_rank_tags,
     _collective_with_timeout,
     _compare_tensors_quick,
     _deepcopy_or_clone,
@@ -1126,6 +1127,210 @@ class TestDumpModel:
         assert all("grad" in f for f in filenames)
 
 
+class TestParallelRankInFilename:
+    def test_config_default_false(self):
+        """include_parallel_rank_in_filename defaults to False."""
+        assert DumperConfig().include_parallel_rank_in_filename is False
+
+    def test_config_from_kv_pairs(self):
+        """include_parallel_rank_in_filename is parsed as a bool from kv pairs."""
+        cfg = DumperConfig.from_kv_pairs(["include_parallel_rank_in_filename=true"])
+        assert cfg.include_parallel_rank_in_filename is True
+
+    def test_collect_tags_merges_keys_across_plugins(self, monkeypatch):
+        """_collect_parallel_rank_tags keeps only rank keys, merging across plugins."""
+        plugin_a = type(
+            "PluginA",
+            (),
+            {"collect_parallel_info": lambda self: {"pp_rank": 1, "ignored": 9}},
+        )()
+        plugin_b = type(
+            "PluginB",
+            (),
+            {"collect_parallel_info": lambda self: {"tp_rank": 2, "cp_rank": 3}},
+        )()
+        monkeypatch.setattr(
+            "sglang.srt.debug_utils.dumper._plugins", [plugin_a, plugin_b]
+        )
+
+        tags = _collect_parallel_rank_tags()
+        assert tags == {"pp_rank": 1, "tp_rank": 2, "cp_rank": 3}
+
+    def test_collect_tags_first_plugin_wins_on_conflict(self, monkeypatch):
+        """When two plugins report the same rank key, the first plugin wins."""
+        plugin_a = type(
+            "PluginA", (), {"collect_parallel_info": lambda self: {"pp_rank": 1}}
+        )()
+        plugin_b = type(
+            "PluginB", (), {"collect_parallel_info": lambda self: {"pp_rank": 7}}
+        )()
+        monkeypatch.setattr(
+            "sglang.srt.debug_utils.dumper._plugins", [plugin_a, plugin_b]
+        )
+
+        assert _collect_parallel_rank_tags() == {"pp_rank": 1}
+
+    def test_collect_tags_skips_empty_plugin_info(self, monkeypatch):
+        """Plugins that report no parallel info are skipped without error."""
+        plugin_empty = type(
+            "PluginEmpty", (), {"collect_parallel_info": lambda self: {}}
+        )()
+        plugin_real = type(
+            "PluginReal", (), {"collect_parallel_info": lambda self: {"tp_rank": 4}}
+        )()
+        monkeypatch.setattr(
+            "sglang.srt.debug_utils.dumper._plugins", [plugin_empty, plugin_real]
+        )
+
+        assert _collect_parallel_rank_tags() == {"tp_rank": 4}
+
+    def test_disabled_filename_has_no_rank_tags(self, tmp_path, monkeypatch):
+        """When disabled, dump filenames do not include parallel-rank tags."""
+        monkeypatch.setattr(
+            _SGLangPlugin,
+            "collect_parallel_info",
+            lambda self: {"pp_rank": 2, "tp_rank": 3},
+        )
+
+        d = _make_test_dumper(tmp_path, include_parallel_rank_in_filename=False)
+        d.dump("hidden", torch.randn(3))
+
+        filenames = _get_filenames(tmp_path)
+        assert filenames
+        assert all("pp_rank=" not in f and "tp_rank=" not in f for f in filenames)
+
+    def test_enabled_filename_includes_rank_tags(self, tmp_path, monkeypatch):
+        """When enabled, dump filenames include the collected parallel-rank tags."""
+        monkeypatch.setattr(
+            _SGLangPlugin,
+            "collect_parallel_info",
+            lambda self: {"pp_rank": 2, "tp_rank": 3},
+        )
+
+        d = _make_test_dumper(tmp_path, include_parallel_rank_in_filename=True)
+        d.dump("hidden", torch.randn(3))
+
+        filenames = _get_filenames(tmp_path)
+        assert any("pp_rank=2" in f and "tp_rank=3" in f for f in filenames)
+
+
+class TestTransformModelParamName:
+    def test_base_plugin_returns_none(self):
+        """The default plugin hook keeps the original name (returns None)."""
+        plugin = _SGLangPlugin()
+        assert (
+            plugin.transform_model_param_name(torch.nn.Linear(2, 2), "layers.0.weight")
+            is None
+        )
+
+    def test_dump_model_keeps_name_without_transform(self, tmp_path):
+        """With no plugin rewriting names, dump_model uses the original param name."""
+        model = torch.nn.Module()
+        model.layers = torch.nn.ModuleList([torch.nn.Linear(2, 2, bias=False)])
+        d = _make_test_dumper(
+            tmp_path, enable_model_value=True, enable_model_grad=False
+        )
+
+        d.dump_model(model, name_prefix="m")
+
+        _assert_files(_get_filenames(tmp_path), exist=["m__layers.0.weight"])
+
+    def test_dump_model_applies_plugin_transform(self, tmp_path, monkeypatch):
+        """dump_model rewrites param names through the plugin transform hook."""
+
+        def _shift_layers(self, model, param_name):
+            return re.sub(
+                r"layers\.(\d+)", lambda m: f"layers.{int(m.group(1)) + 4}", param_name
+            )
+
+        monkeypatch.setattr(_SGLangPlugin, "transform_model_param_name", _shift_layers)
+
+        model = torch.nn.Module()
+        model.layers = torch.nn.ModuleList([torch.nn.Linear(2, 2, bias=False)])
+        d = _make_test_dumper(
+            tmp_path, enable_model_value=True, enable_model_grad=False
+        )
+
+        d.dump_model(model, name_prefix="m")
+
+        _assert_files(
+            _get_filenames(tmp_path),
+            exist=["m__layers.4.weight"],
+            not_exist=["m__layers.0.weight"],
+        )
+
+    def test_get_model_config_unwraps_module_chain(self):
+        """_get_model_config peels nested .module wrappers to find .config."""
+        config = object()
+        leaf = type("Leaf", (), {"config": config})()
+        wrapped = type("W", (), {"module": type("W2", (), {"module": leaf})()})()
+        assert _MegatronPlugin._get_model_config(wrapped) is config
+
+    def test_get_model_config_returns_none_when_absent(self):
+        """_get_model_config returns None when no .config is reachable."""
+        assert _MegatronPlugin._get_model_config(object()) is None
+
+
+class _FakeMpu:
+    _pp_size = 2
+
+    @classmethod
+    def get_pipeline_model_parallel_world_size(cls):
+        return cls._pp_size
+
+
+class TestMegatronTransformModelParamName:
+    @pytest.fixture(autouse=True)
+    def _patch_megatron(self, monkeypatch):
+        monkeypatch.setattr(_MegatronPlugin, "_available", True)
+        monkeypatch.setattr(_MegatronPlugin, "_mpu", _FakeMpu, raising=False)
+        monkeypatch.setattr(
+            _MegatronPlugin,
+            "_get_model_config",
+            staticmethod(lambda model: object()),
+        )
+        monkeypatch.setattr(
+            _MegatronPlugin,
+            "_get_transformer_layer_offset",
+            staticmethod(lambda config: 4),
+        )
+
+    def test_remaps_local_layer_index_to_global(self):
+        """layers.N is shifted by the PP-stage offset; other tokens untouched."""
+        plugin = _MegatronPlugin()
+        result = plugin.transform_model_param_name(object(), "decoder.layers.0.weight")
+        assert result == "decoder.layers.4.weight"
+
+    def test_remaps_all_layer_occurrences(self):
+        """Every ``layers.N`` occurrence in the name is shifted."""
+        plugin = _MegatronPlugin()
+        result = plugin.transform_model_param_name(object(), "layers.1.x.layers.2.y")
+        assert result == "layers.5.x.layers.6.y"
+
+    def test_returns_none_when_not_available(self, monkeypatch):
+        """No transform when megatron is unavailable."""
+        monkeypatch.setattr(_MegatronPlugin, "_available", False)
+        plugin = _MegatronPlugin()
+        assert plugin.transform_model_param_name(object(), "layers.0.weight") is None
+
+    def test_returns_none_when_pp_size_one(self, monkeypatch):
+        """No transform without pipeline parallelism (pp_size == 1)."""
+        monkeypatch.setattr(_FakeMpu, "_pp_size", 1)
+        plugin = _MegatronPlugin()
+        assert plugin.transform_model_param_name(object(), "layers.0.weight") is None
+        monkeypatch.setattr(_FakeMpu, "_pp_size", 2)
+
+    def test_returns_none_when_offset_zero(self, monkeypatch):
+        """A zero offset (e.g. first PP stage) leaves the name unchanged (None)."""
+        monkeypatch.setattr(
+            _MegatronPlugin,
+            "_get_transformer_layer_offset",
+            staticmethod(lambda config: 0),
+        )
+        plugin = _MegatronPlugin()
+        assert plugin.transform_model_param_name(object(), "layers.0.weight") is None
+
+
 class TestCleanup:
     def test_cleanup_removes_old_dumps(self, tmp_path):
         old_dir = tmp_path / "dump_old"