minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Add explicit EPLB balancedness reporting modes (#34998)

Browse Source
This commit is contained in:
Lianmin Zheng
2026-08-16 15:31:11 -07:00
committed by GitHub
parent 77cadf6b98
commit f61f584347
10 changed files with 133 additions and 50 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/docs/advanced_features/server_arguments.mdx
+4 -4
View File
@@ -1885,10 +1885,10 @@ Please consult the documentation below and [server_args.py](https://github.com/s
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Type: int</td>
</tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--enable-expert-distribution-metrics`</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Enable logging metrics for expert balancedness</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`False`</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>bool flag (set to enable)</td>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--expert-balancedness-report-mode`</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Where to report expert balancedness.</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>off</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><code>off</code>, <code>server_log</code>, <code>prometheus</code>, <code>both</code></td>
</tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--deepep-config`</td>
docs/docs/hardware-platforms/ascend-npus/reference/support_features.mdx
+3 -3
View File
@@ -1709,9 +1709,9 @@ If the value is int8, you must also set the environment variable:DEEP_NORMAL_MOD
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>A2, A3</td>
</tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--enable-expert-distribution-metrics`</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>`False`</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>bool flag (set to enable)</td>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--expert-balancedness-report-mode`</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><code>off</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>off</code>, <code>server_log</code>, <code>prometheus</code>, <code>both</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>A2, A3</td>
</tr>
<tr>
docs/docs/references/environment_variables.mdx
+1 -1
View File
@@ -1745,7 +1745,7 @@ SGLang supports various environment variables that can be used to configure its
</tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_ENABLE_EPLB_BALANCEDNESS_METRIC</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Emit an EPLB balancedness metric.</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Removed. Use <code>--expert-balancedness-report-mode=prometheus</code> to emit the EPLB balancedness metric.</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>false</code></td>
</tr>
<tr>
python/sglang/srt/arg_groups/argparse_actions.py
+4 -1
View File
@@ -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."
)
python/sglang/srt/environ.py
-1
View File
@@ -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(
python/sglang/srt/eplb/expert_distribution.py
+38 -24
View File
@@ -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,
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
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()
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()]}"
)
# 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):
python/sglang/srt/managers/scheduler_components/metrics_reporter.py
+39 -4
View File
@@ -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,15 +957,44 @@ 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:
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=m.eplb_balancedness.item(),
balancedness=balancedness,
)
def _emit_forward_pass_metrics(
python/sglang/srt/observability/metrics_collector.py
+3 -3
View File
@@ -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.",
python/sglang/srt/server_args.py
+35 -4
View File
@@ -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
test/registered/npu/basic_function/parallel_strategy/expert_parallelism/test_npu_eplb_min_rebalancing_utilization_threshold.py
+2 -1
View File
@@ -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",
]