diff --git a/python/sglang/srt/debug_utils/dumper.py b/python/sglang/srt/debug_utils/dumper.py index 3d4c368bc..d4cafb1fc 100644 --- a/python/sglang/srt/debug_utils/dumper.py +++ b/python/sglang/srt/debug_utils/dumper.py @@ -508,12 +508,10 @@ class _NonIntrusiveDumper: cls, module_name: str, module: "torch.nn.Module" ) -> Optional[dict]: if cls._LAYER_NAME_RE.fullmatch(module_name): - # Megatron - if hasattr(module, "layer_number"): - return {"layer_id": module.layer_number - 1} - # SGLang - if hasattr(module, "layer_id"): - return {"layer_id": module.layer_id} + for plugin in _plugins: + layer_id = plugin.detect_layer_id(module) + if layer_id is not None: + return {"layer_id": layer_id} return None def _register_ctx_hooks(self, module: "torch.nn.Module", *, ctx: dict) -> None: @@ -560,31 +558,10 @@ class _NonIntrusiveDumper: return {"": tensors[0]} return {str(i): t for i, t in enumerate(tensors)} - # SGLang specific - try: - from sglang.srt.layers.logits_processor import LogitsProcessorOutput - from sglang.srt.model_executor.forward_batch_info import ( - ForwardBatch, - PPProxyTensors, - ) - - if isinstance(value, LogitsProcessorOutput): - return {"next_token_logits": value.next_token_logits} - if isinstance(value, ForwardBatch): - if skip_forward_batch: - return {} - return { - "input_ids": value.input_ids, - "seq_lens": value.seq_lens, - "positions": value.positions, - } - if isinstance(value, PPProxyTensors): - return {k: v for k, v in value.tensors.items()} - except ImportError: - pass - - # Megatron specific - # TODO + for plugin in _plugins: + result = plugin.convert_value(value, skip_forward_batch=skip_forward_batch) + if result is not None: + return result return {} @@ -710,99 +687,13 @@ def _compute_static_meta(): "world_size": _get_world_size(), } - if x := _collect_sglang_parallel_info(): - result["sglang_parallel_info"] = x - if x := _collect_megatron_parallel_info(): - result["megatron_parallel_info"] = x + for plugin in _plugins: + if info := plugin.collect_parallel_info(): + result[f"{plugin.name}_parallel_info"] = info return result -def _collect_sglang_parallel_info(): - info = {} - - try: - from sglang.srt.distributed import ( - get_moe_expert_parallel_rank, - get_moe_expert_parallel_world_size, - get_moe_tensor_parallel_rank, - get_moe_tensor_parallel_world_size, - get_pipeline_model_parallel_rank, - get_pipeline_model_parallel_world_size, - get_tensor_model_parallel_rank, - get_tensor_model_parallel_world_size, - ) - - info["tp_rank"] = get_tensor_model_parallel_rank() - info["tp_size"] = get_tensor_model_parallel_world_size() - info["pp_rank"] = get_pipeline_model_parallel_rank() - info["pp_size"] = get_pipeline_model_parallel_world_size() - info["moe_ep_rank"] = get_moe_expert_parallel_rank() - info["moe_ep_size"] = get_moe_expert_parallel_world_size() - info["moe_tp_rank"] = get_moe_tensor_parallel_rank() - info["moe_tp_size"] = get_moe_tensor_parallel_world_size() - except (ImportError, AttributeError, AssertionError): - info["distributed_error"] = True - - try: - from sglang.srt.layers.dp_attention import ( - get_attention_dp_rank, - get_attention_dp_size, - get_attention_tp_rank, - get_attention_tp_size, - get_local_attention_dp_rank, - get_local_attention_dp_size, - is_dp_attention_enabled, - ) - - info["enable_dp_attention"] = is_dp_attention_enabled() - info["attn_tp_rank"] = get_attention_tp_rank() - info["attn_tp_size"] = get_attention_tp_size() - info["attn_dp_rank"] = get_attention_dp_rank() - info["attn_dp_size"] = get_attention_dp_size() - info["local_attn_dp_rank"] = get_local_attention_dp_rank() - info["local_attn_dp_size"] = get_local_attention_dp_size() - except (ImportError, AttributeError, AssertionError): - info["dp_attention_error"] = True - - return info - - -def _collect_megatron_parallel_info(): - info = {} - - try: - from megatron.core import parallel_state as mpu - - info["tp_rank"] = mpu.get_tensor_model_parallel_rank() - info["tp_size"] = mpu.get_tensor_model_parallel_world_size() - info["pp_rank"] = mpu.get_pipeline_model_parallel_rank() - info["pp_size"] = mpu.get_pipeline_model_parallel_world_size() - info["dp_rank"] = mpu.get_data_parallel_rank() - info["dp_size"] = mpu.get_data_parallel_world_size() - info["cp_rank"] = mpu.get_context_parallel_rank() - info["cp_size"] = mpu.get_context_parallel_world_size() - info["vpp_rank"] = mpu.get_virtual_pipeline_model_parallel_rank() - info["vpp_size"] = mpu.get_virtual_pipeline_model_parallel_world_size() - info["ep_rank"] = mpu.get_expert_model_parallel_rank() - info["ep_size"] = mpu.get_expert_model_parallel_world_size() - info["etp_rank"] = mpu.get_expert_tensor_parallel_rank() - info["etp_size"] = mpu.get_expert_tensor_parallel_world_size() - info["edp_rank"] = mpu.get_expert_data_parallel_rank() - info["edp_size"] = mpu.get_expert_data_parallel_world_size() - info["tcp_rank"] = mpu.get_tensor_and_context_parallel_rank() - info["tcp_size"] = mpu.get_tensor_and_context_parallel_world_size() - info["etmp_rank"] = mpu.get_expert_tensor_and_model_parallel_rank() - info["etmp_size"] = mpu.get_expert_tensor_and_model_parallel_world_size() - info["tp_src_rank"] = mpu.get_tensor_model_parallel_src_rank() - info["mp_src_rank"] = mpu.get_model_parallel_src_rank() - info["dp_src_rank"] = mpu.get_data_parallel_src_rank() - except (ImportError, AttributeError, AssertionError): - info["megatron_error"] = True - - return info - - # -------------------------------------- http control server ------------------------------------------ @@ -1009,6 +900,168 @@ def _get_local_ip_by_remote() -> Optional[str]: return None +# -------------------------------------- framework plugins ------------------------------------------ + + +class _FrameworkPlugin(ABC): + @property + @abstractmethod + def name(self) -> str: ... + + @abstractmethod + def collect_parallel_info(self) -> dict: ... + + @abstractmethod + def convert_value( + self, value: Any, *, skip_forward_batch: bool + ) -> Optional[dict[str, "torch.Tensor"]]: + """Return converted tensors dict, or None if this plugin doesn't handle the value.""" + ... + + @abstractmethod + def detect_layer_id(self, module: "torch.nn.Module") -> Optional[int]: + """Return 0-indexed layer_id, or None if not detectable.""" + ... + + +class _SGLangPlugin(_FrameworkPlugin): + _available = True + try: + from sglang.srt import distributed as _dist + from sglang.srt.layers import dp_attention as _dp_attn + from sglang.srt.layers.logits_processor import LogitsProcessorOutput + from sglang.srt.model_executor.forward_batch_info import ( + ForwardBatch, + PPProxyTensors, + ) + except ImportError: + _available = False + + @property + def name(self) -> str: + return "sglang" + + def collect_parallel_info(self) -> dict: + if not self._available: + return {} + + info = {} + + try: + info["tp_rank"] = self._dist.get_tensor_model_parallel_rank() + info["tp_size"] = self._dist.get_tensor_model_parallel_world_size() + info["pp_rank"] = self._dist.get_pipeline_model_parallel_rank() + info["pp_size"] = self._dist.get_pipeline_model_parallel_world_size() + info["moe_ep_rank"] = self._dist.get_moe_expert_parallel_rank() + info["moe_ep_size"] = self._dist.get_moe_expert_parallel_world_size() + info["moe_tp_rank"] = self._dist.get_moe_tensor_parallel_rank() + info["moe_tp_size"] = self._dist.get_moe_tensor_parallel_world_size() + except (AttributeError, AssertionError): + info["distributed_error"] = True + + try: + info["enable_dp_attention"] = self._dp_attn.is_dp_attention_enabled() + info["attn_tp_rank"] = self._dp_attn.get_attention_tp_rank() + info["attn_tp_size"] = self._dp_attn.get_attention_tp_size() + info["attn_dp_rank"] = self._dp_attn.get_attention_dp_rank() + info["attn_dp_size"] = self._dp_attn.get_attention_dp_size() + info["local_attn_dp_rank"] = self._dp_attn.get_local_attention_dp_rank() + info["local_attn_dp_size"] = self._dp_attn.get_local_attention_dp_size() + except (AttributeError, AssertionError): + info["dp_attention_error"] = True + + return info + + def convert_value( + self, value: Any, *, skip_forward_batch: bool + ) -> Optional[dict[str, "torch.Tensor"]]: + if not self._available: + return None + + if isinstance(value, self.LogitsProcessorOutput): + return {"next_token_logits": value.next_token_logits} + if isinstance(value, self.ForwardBatch): + if skip_forward_batch: + return {} + return { + "input_ids": value.input_ids, + "seq_lens": value.seq_lens, + "positions": value.positions, + } + if isinstance(value, self.PPProxyTensors): + return {k: v for k, v in value.tensors.items()} + + return None + + def detect_layer_id(self, module: "torch.nn.Module") -> Optional[int]: + if hasattr(module, "layer_id"): + return module.layer_id + return None + + +class _MegatronPlugin(_FrameworkPlugin): + _available = True + try: + from megatron.core import parallel_state as _mpu + except ImportError: + _available = False + + @property + def name(self) -> str: + return "megatron" + + def collect_parallel_info(self) -> dict: + if not self._available: + return {} + + info = {} + try: + info["tp_rank"] = self._mpu.get_tensor_model_parallel_rank() + info["tp_size"] = self._mpu.get_tensor_model_parallel_world_size() + info["pp_rank"] = self._mpu.get_pipeline_model_parallel_rank() + info["pp_size"] = self._mpu.get_pipeline_model_parallel_world_size() + info["dp_rank"] = self._mpu.get_data_parallel_rank() + info["dp_size"] = self._mpu.get_data_parallel_world_size() + info["cp_rank"] = self._mpu.get_context_parallel_rank() + info["cp_size"] = self._mpu.get_context_parallel_world_size() + info["vpp_rank"] = self._mpu.get_virtual_pipeline_model_parallel_rank() + info["vpp_size"] = ( + self._mpu.get_virtual_pipeline_model_parallel_world_size() + ) + info["ep_rank"] = self._mpu.get_expert_model_parallel_rank() + info["ep_size"] = self._mpu.get_expert_model_parallel_world_size() + info["etp_rank"] = self._mpu.get_expert_tensor_parallel_rank() + info["etp_size"] = self._mpu.get_expert_tensor_parallel_world_size() + info["edp_rank"] = self._mpu.get_expert_data_parallel_rank() + info["edp_size"] = self._mpu.get_expert_data_parallel_world_size() + info["tcp_rank"] = self._mpu.get_tensor_and_context_parallel_rank() + info["tcp_size"] = self._mpu.get_tensor_and_context_parallel_world_size() + info["etmp_rank"] = self._mpu.get_expert_tensor_and_model_parallel_rank() + info["etmp_size"] = ( + self._mpu.get_expert_tensor_and_model_parallel_world_size() + ) + info["tp_src_rank"] = self._mpu.get_tensor_model_parallel_src_rank() + info["mp_src_rank"] = self._mpu.get_model_parallel_src_rank() + info["dp_src_rank"] = self._mpu.get_data_parallel_src_rank() + except (AttributeError, AssertionError): + info["megatron_error"] = True + + return info + + def convert_value( + self, value: Any, *, skip_forward_batch: bool + ) -> Optional[dict[str, "torch.Tensor"]]: + return None + + def detect_layer_id(self, module: "torch.nn.Module") -> Optional[int]: + if hasattr(module, "layer_number"): + return module.layer_number - 1 + return None + + +_plugins: list[_FrameworkPlugin] = [_SGLangPlugin(), _MegatronPlugin()] + + # -------------------------------------- singleton ------------------------------------------ diff --git a/test/registered/debug_utils/test_dumper.py b/test/registered/debug_utils/test_dumper.py index 0c143f95f..b96ece72a 100644 --- a/test/registered/debug_utils/test_dumper.py +++ b/test/registered/debug_utils/test_dumper.py @@ -13,14 +13,14 @@ import torch import torch.distributed as dist from sglang.srt.debug_utils.dumper import ( - _collect_megatron_parallel_info, - _collect_sglang_parallel_info, _collective_with_timeout, _Dumper, _DumperConfig, _format_tags, _materialize_value, + _MegatronPlugin, _obj_to_dict, + _SGLangPlugin, _torch_save, dumper, get_tensor_info, @@ -529,10 +529,10 @@ class TestStaticMetadata: assert meta1 is meta2 def test_parallel_info_graceful_fallback(self): - sglang_info = _collect_sglang_parallel_info() + sglang_info = _SGLangPlugin().collect_parallel_info() assert isinstance(sglang_info, dict) - megatron_info = _collect_megatron_parallel_info() + megatron_info = _MegatronPlugin().collect_parallel_info() assert isinstance(megatron_info, dict) def test_dump_includes_static_meta(self, tmp_path):