diff --git a/docs/docs/advanced_features/server_arguments.mdx b/docs/docs/advanced_features/server_arguments.mdx index 5ee20bead..ac4291753 100644 --- a/docs/docs/advanced_features/server_arguments.mdx +++ b/docs/docs/advanced_features/server_arguments.mdx @@ -1885,10 +1885,10 @@ Please consult the documentation below and [server_args.py](https://github.com/s Type: int - `--enable-expert-distribution-metrics` - Enable logging metrics for expert balancedness - `False` - bool flag (set to enable) + `--expert-balancedness-report-mode` + Where to report expert balancedness. + off + off, server_log, prometheus, both `--deepep-config` diff --git a/docs/docs/hardware-platforms/ascend-npus/reference/support_features.mdx b/docs/docs/hardware-platforms/ascend-npus/reference/support_features.mdx index 734f0d3ff..4c0fe9429 100644 --- a/docs/docs/hardware-platforms/ascend-npus/reference/support_features.mdx +++ b/docs/docs/hardware-platforms/ascend-npus/reference/support_features.mdx @@ -1709,9 +1709,9 @@ If the value is int8, you must also set the environment variable:DEEP_NORMAL_MOD A2, A3 - `--enable-expert-distribution-metrics` - `False` - bool flag (set to enable) + `--expert-balancedness-report-mode` + off + off, server_log, prometheus, both A2, A3 diff --git a/docs/docs/references/environment_variables.mdx b/docs/docs/references/environment_variables.mdx index 7bca0227e..a0248b71e 100644 --- a/docs/docs/references/environment_variables.mdx +++ b/docs/docs/references/environment_variables.mdx @@ -1745,7 +1745,7 @@ SGLang supports various environment variables that can be used to configure its SGLANG_ENABLE_EPLB_BALANCEDNESS_METRIC - Emit an EPLB balancedness metric. + Removed. Use --expert-balancedness-report-mode=prometheus to emit the EPLB balancedness metric. false diff --git a/python/sglang/srt/arg_groups/argparse_actions.py b/python/sglang/srt/arg_groups/argparse_actions.py index 3fdd39e97..06b98ceb4 100644 --- a/python/sglang/srt/arg_groups/argparse_actions.py +++ b/python/sglang/srt/arg_groups/argparse_actions.py @@ -30,12 +30,15 @@ def print_deprecated_warning(message: str): class DeprecatedAction(argparse.Action): - def __init__(self, option_strings, dest, nargs=0, **kwargs): + def __init__(self, option_strings, dest, error_message=None, nargs=0, **kwargs): + self.error_message = error_message super(DeprecatedAction, self).__init__( option_strings, dest, nargs=nargs, **kwargs ) def __call__(self, parser, namespace, values, option_string=None): + if self.error_message is not None: + parser.error(self.error_message) print_deprecated_warning( f"The command line argument '{option_string}' is deprecated and will be removed in future versions." ) diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 5e8799ce4..6e0a8a62b 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -939,7 +939,6 @@ class Envs: SGLANG_LOG_EXPERT_LOCATION_METADATA = EnvBool(False) SGLANG_EXPERT_DISTRIBUTION_RECORDER_DIR = EnvStr("/tmp") SGLANG_EPLB_HEATMAP_COLLECTION_INTERVAL = EnvInt(0) - SGLANG_ENABLE_EPLB_BALANCEDNESS_METRIC = EnvBool(False) # Chunk size for the rebalance expert-weight P2P exchange; set # >= num_physical_experts to submit a single batch_isend_irecv. SGLANG_EPLB_P2P_BATCH_CHUNK_SIZE = EnvIntWithAlias( diff --git a/python/sglang/srt/eplb/expert_distribution.py b/python/sglang/srt/eplb/expert_distribution.py index e28f583ee..2ea7cb7d8 100644 --- a/python/sglang/srt/eplb/expert_distribution.py +++ b/python/sglang/srt/eplb/expert_distribution.py @@ -22,7 +22,17 @@ from collections import deque from contextlib import contextmanager from dataclasses import dataclass from pathlib import Path -from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Tuple, Type +from typing import ( + TYPE_CHECKING, + Any, + Dict, + List, + Literal, + Optional, + Sequence, + Tuple, + Type, +) import einops import torch @@ -47,16 +57,22 @@ logger = logging.getLogger(__name__) # --------------------------------------- Entrypoint ----------------------------------------- _OutputMode = Literal["file", "object"] +EPLB_BALANCEDNESS_WINDOW_SIZES = (10, 100, 1000) @dataclass class ExpertDistributionMetrics: + forward_pass_id: int eplb_balancedness: torch.Tensor + gpu_physical_count_sum: Optional[torch.Tensor] + reset_server_log_history: bool def map_device_tensors(self, fn): # Device-tensor fields only; caller injects the copy+safety primitive # (see GenerationBatchResult.copy_to_cpu). self.eplb_balancedness = fn(self.eplb_balancedness) + if self.gpu_physical_count_sum is not None: + self.gpu_physical_count_sum = fn(self.gpu_physical_count_sum) class ExpertDistributionRecorder(ABC): @@ -159,9 +175,10 @@ class _ExpertDistributionRecorderReal(ExpertDistributionRecorder): for k in self._accumulator.get_single_pass_gatherer_keys() } - if server_args.enable_expert_distribution_metrics: + if server_args.should_report_expert_balancedness(): logger.info( - "ExpertDistributionRecorder auto start record since enable_expert_distribution_metrics" + "ExpertDistributionRecorder auto start record since " + f"expert_balancedness_report_mode={server_args.expert_balancedness_report_mode}" ) self.start_record() @@ -699,11 +716,12 @@ class _UtilizationRateAccumulatorMixin(_Accumulator): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) - self._enable = self._server_args.enable_expert_distribution_metrics + self._enable = self._server_args.should_report_expert_balancedness() if self._enable: - self.window_sizes = [10, 100, 1000] + self.window_sizes = EPLB_BALANCEDNESS_WINDOW_SIZES self._history = _DequeCollection(maxlens=self.window_sizes) + self._reset_server_log_history = True self._rank = torch.distributed.get_rank() expert_dispatch_cls = resolve_collector_class( self._server_args, @@ -732,6 +750,7 @@ class _UtilizationRateAccumulatorMixin(_Accumulator): super().reset() if self._enable: self._history.clear() + self._reset_server_log_history = True def _append_utilization_rate( self, @@ -757,27 +776,22 @@ class _UtilizationRateAccumulatorMixin(_Accumulator): should_track_history = not math.isclose( self._server_args.eplb_min_rebalancing_utilization_threshold, 1.0 ) - if envs.SGLANG_ENABLE_EPLB_BALANCEDNESS_METRIC.get(): - outputs["metrics"] = ExpertDistributionMetrics( - eplb_balancedness=utilization_rate_gpu, - ) - if should_track_history: - self._history.append(utilization_rate_gpu.item()) - else: - # TODO maybe refactor this part to also avoid a `.item()` gpu->cpu sync - utilization_rate_cpu = utilization_rate_gpu.item() - self._history.append(utilization_rate_cpu) - gpu_physical_count_sum = gpu_physical_count.sum().item() + should_log = ( + self._server_args.should_log_expert_balancedness_to_server_log() + ) + outputs["metrics"] = ExpertDistributionMetrics( + forward_pass_id=forward_pass_id, + eplb_balancedness=utilization_rate_gpu, + gpu_physical_count_sum=( + gpu_physical_count.sum() if should_log else None + ), + reset_server_log_history=self._reset_server_log_history, + ) + self._reset_server_log_history = False - logger.info( - f"[Expert Balancedness] " - f"forward_pass_id={forward_pass_id} " - f"current_pass_balancedness={utilization_rate_cpu:.03f} " - f"{''.join(f'last_{size}_average_balancedness={value:.03f} ' for size, value in self._history.mean().items())} " - f"gpu_physical_count_sum={gpu_physical_count_sum}" - # f"current_pass_per_layer={[round(x, 2) for x in utilization_rate_tensor.cpu().tolist()]}" - ) + if should_track_history: + self._history.append(utilization_rate_gpu.item()) # TODO refactor def _handle_metric_eplb_heatmap(self, gpu_physical_count: torch.Tensor): @@ -804,7 +818,7 @@ class _UtilizationRateAccumulatorMixin(_Accumulator): class _DequeCollection: - def __init__(self, maxlens: List[int]): + def __init__(self, maxlens: Sequence[int]): self._dequeues = [deque(maxlen=maxlen) for maxlen in maxlens] def append(self, value): diff --git a/python/sglang/srt/managers/scheduler_components/metrics_reporter.py b/python/sglang/srt/managers/scheduler_components/metrics_reporter.py index 73a49a66d..60bb04953 100644 --- a/python/sglang/srt/managers/scheduler_components/metrics_reporter.py +++ b/python/sglang/srt/managers/scheduler_components/metrics_reporter.py @@ -5,12 +5,13 @@ import logging import math import tempfile import time -from collections import defaultdict +from collections import defaultdict, deque from dataclasses import dataclass from typing import TYPE_CHECKING, List, Optional, Tuple, Union from sglang.srt.disaggregation.utils import DisaggregationMode from sglang.srt.environ import envs +from sglang.srt.eplb.expert_distribution import EPLB_BALANCEDNESS_WINDOW_SIZES from sglang.srt.managers.schedule_batch import ScheduleBatch from sglang.srt.managers.utils import GenerationBatchResult from sglang.srt.observability.metrics_collector import ( @@ -112,6 +113,11 @@ class SchedulerMetricsReporter: ) self._init_metrics(self.tp_rank, self.pp_rank, self.dp_rank) self._install_device_timer_on_runners() + # Keep log history after the existing async result copy so reporting does + # not synchronize the model stream once per generated token. + self._eplb_balancedness_history = [ + deque(maxlen=window_size) for window_size in EPLB_BALANCEDNESS_WINDOW_SIZES + ] def _init_metrics( self, @@ -951,16 +957,45 @@ class SchedulerMetricsReporter: batch: ScheduleBatch, result: Union[GenerationBatchResult, EmbeddingBatchResult], ): - if not self.enable_metrics: - return if not isinstance(result, GenerationBatchResult): return if (m := result.expert_distribution_metrics) is not None: - self.metrics_collector.increment_eplb_balancedness( - forward_mode=batch.forward_mode.name.lower(), - balancedness=m.eplb_balancedness.item(), - ) + balancedness = m.eplb_balancedness.item() + + if ( + self.scheduler.server_args.should_log_expert_balancedness_to_server_log() + ): + if m.reset_server_log_history: + for history in self._eplb_balancedness_history: + history.clear() + for history in self._eplb_balancedness_history: + history.append(balancedness) + balancedness_history_means = { + history.maxlen: sum(history) / len(history) + for history in self._eplb_balancedness_history + if len(history) > 0 + } + assert m.gpu_physical_count_sum is not None + gpu_physical_count_sum = m.gpu_physical_count_sum.item() + + logger.info( + f"[Expert Balancedness] " + f"forward_pass_id={m.forward_pass_id} " + f"current_pass_balancedness={balancedness:.03f} " + f"{''.join(f'last_{size}_average_balancedness={value:.03f} ' for size, value in balancedness_history_means.items())} " + f"gpu_physical_count_sum={gpu_physical_count_sum}" + ) + + if ( + self.enable_metrics + and self.scheduler.server_args.should_export_expert_balancedness_to_prometheus() + ): + assert self.metrics_collector is not None + self.metrics_collector.increment_eplb_balancedness( + forward_mode=batch.forward_mode.name.lower(), + balancedness=balancedness, + ) def _emit_forward_pass_metrics( self, diff --git a/python/sglang/srt/observability/metrics_collector.py b/python/sglang/srt/observability/metrics_collector.py index a5976c120..adc15d97a 100644 --- a/python/sglang/srt/observability/metrics_collector.py +++ b/python/sglang/srt/observability/metrics_collector.py @@ -24,7 +24,6 @@ from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, Dict, List, Mapping, Optional, Set, Union from sglang.srt.disaggregation.utils import DisaggregationMode -from sglang.srt.environ import envs from sglang.srt.model_executor.forward_batch_info import ForwardMode from sglang.srt.observability.utils import exponential_buckets, generate_buckets from sglang.srt.server_args import ServerArgs @@ -240,10 +239,10 @@ class SchedulerMetricsCollector(_StatLoggerDIMixin): def __init__( self, labels: Dict[str, str], + server_args: ServerArgs, enable_lora: bool = False, enable_hierarchical_cache: bool = False, enable_streaming_session: bool = False, - server_args: Optional[ServerArgs] = None, ) -> None: # We need to import prometheus_client after setting the env variable `PROMETHEUS_MULTIPROC_DIR` from prometheus_client import Counter as _PromCounter @@ -874,7 +873,8 @@ class SchedulerMetricsCollector(_StatLoggerDIMixin): # ================================================================= if ( labels["moe_ep_rank"] == 0 - ) and envs.SGLANG_ENABLE_EPLB_BALANCEDNESS_METRIC.get(): + and server_args.should_export_expert_balancedness_to_prometheus() + ): self.eplb_balancedness = Summary( name="sglang:eplb_balancedness", documentation="Balancedness of MoE in expert parallelism.", diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 4f23cc0d8..bebbf3194 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -2409,9 +2409,11 @@ class ServerArgs: "Circular buffer size of expert distribution recorder. Set to -1 to denote infinite buffer.", NS("exec.moe"), ] = None - enable_expert_distribution_metrics: A[ - bool, "Enable logging metrics for expert balancedness", NS("exec.moe") - ] = False + expert_balancedness_report_mode: A[ + Literal["off", "server_log", "prometheus", "both"], + "Where to report expert balancedness. Options: off, server_log, prometheus, both.", + NS("exec.moe"), + ] = "off" deepep_config: A[ Optional[str], "Tuned DeepEP config suitable for your own cluster. It can be either a string with JSON content or a file path.", @@ -7198,7 +7200,14 @@ class ServerArgs: # ===== END TO BE REFACTORED ==== def _handle_expert_distribution_metrics(self): - if self.enable_expert_distribution_metrics and ( + if "SGLANG_ENABLE_EPLB_BALANCEDNESS_METRIC" in os.environ: + raise ValueError( + "SGLANG_ENABLE_EPLB_BALANCEDNESS_METRIC is no longer supported. Use " + "--expert-balancedness-report-mode with one of: off, server_log, " + "prometheus, both." + ) + + if self.should_report_expert_balancedness() and ( self.expert_distribution_recorder_mode is None ): self.expert_distribution_recorder_mode = "stat" @@ -8620,6 +8629,19 @@ class ServerArgs: ) # --- Deprecated argument registrations --- + parser.add_argument( + "--enable-expert-distribution-metrics", + action=DeprecatedAction, + error_message=( + "--enable-expert-distribution-metrics is no longer supported. Use " + "--expert-balancedness-report-mode with one of: off, server_log, " + "prometheus, both." + ), + help=( + "Removed. Use --expert-balancedness-report-mode with one of: " + "off, server_log, prometheus, both." + ), + ) parser.add_argument( "--stream-output", action=DeprecatedStoreTrueAction, @@ -9593,6 +9615,15 @@ class ServerArgs: "dp_size": self.dp_size, } + def should_report_expert_balancedness(self) -> bool: + return self.expert_balancedness_report_mode != "off" + + def should_log_expert_balancedness_to_server_log(self) -> bool: + return self.expert_balancedness_report_mode in ("server_log", "both") + + def should_export_expert_balancedness_to_prometheus(self) -> bool: + return self.expert_balancedness_report_mode in ("prometheus", "both") + def m3_fp8_attn_gemm_enabled(args) -> bool: """Whether MiniMax-M3 attention GEMMs run in fp8 (no opt-in flag; active diff --git a/test/registered/npu/basic_function/parallel_strategy/expert_parallelism/test_npu_eplb_min_rebalancing_utilization_threshold.py b/test/registered/npu/basic_function/parallel_strategy/expert_parallelism/test_npu_eplb_min_rebalancing_utilization_threshold.py index 1c23e7df1..dcba511fe 100644 --- a/test/registered/npu/basic_function/parallel_strategy/expert_parallelism/test_npu_eplb_min_rebalancing_utilization_threshold.py +++ b/test/registered/npu/basic_function/parallel_strategy/expert_parallelism/test_npu_eplb_min_rebalancing_utilization_threshold.py @@ -51,7 +51,8 @@ class TestEplbMinRebalancingUtilizationThresholdBase(CustomTestCase): 50, "--expert-distribution-recorder-buffer-size", 50, - "--enable-expert-distribution-metrics", + "--expert-balancedness-report-mode", + "server_log", "--eplb-rebalance-layers-per-chunk", "1", ]