diff --git a/python/sglang/srt/debug_utils/dumper.py b/python/sglang/srt/debug_utils/dumper.py index 6c7d3f6e5..63af64221 100644 --- a/python/sglang/srt/debug_utils/dumper.py +++ b/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()] diff --git a/test/registered/debug_utils/test_dumper.py b/test/registered/debug_utils/test_dumper.py index a76f0f7ff..d0c568c04 100644 --- a/test/registered/debug_utils/test_dumper.py +++ b/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"