diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index a843aa87d..b10002c91 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -268,6 +268,7 @@ 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) # TBO SGLANG_TBO_DEBUG = EnvBool(False) diff --git a/python/sglang/srt/eplb/expert_distribution.py b/python/sglang/srt/eplb/expert_distribution.py index d40439409..3fa9fcbce 100644 --- a/python/sglang/srt/eplb/expert_distribution.py +++ b/python/sglang/srt/eplb/expert_distribution.py @@ -20,6 +20,7 @@ import time from abc import ABC 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 @@ -43,6 +44,14 @@ logger = logging.getLogger(__name__) _OutputMode = Literal["file", "object"] +@dataclass +class ExpertDistributionMetrics: + eplb_balancedness: torch.Tensor + + def copy_to_cpu(self): + self.eplb_balancedness = self.eplb_balancedness.to("cpu", non_blocking=True) + + class ExpertDistributionRecorder(ABC): """Global expert distribution recording""" @@ -78,7 +87,7 @@ class ExpertDistributionRecorder(ABC): @contextmanager def with_forward_pass(self, forward_pass_id: int, forward_batch: ForwardBatch): - yield + yield {} def on_select_experts(self, topk_ids: torch.Tensor): pass @@ -157,12 +166,13 @@ class _ExpertDistributionRecorderReal(ExpertDistributionRecorder): @contextmanager def with_forward_pass(self, forward_pass_id: int, forward_batch: ForwardBatch): + outputs = {} with self._current_forward_pass_id.with_value(forward_pass_id): self._on_forward_pass_start(forward_batch) try: - yield + yield outputs finally: - self._on_forward_pass_end(forward_pass_id) + self._on_forward_pass_end(forward_pass_id, outputs) @contextmanager def disable_this_region(self): @@ -181,12 +191,14 @@ class _ExpertDistributionRecorderReal(ExpertDistributionRecorder): gatherer.reset() gatherer.on_forward_pass_start(forward_batch) - def _on_forward_pass_end(self, forward_pass_id: int): + def _on_forward_pass_end(self, forward_pass_id: int, outputs: Dict[str, Any]): if not self._recording: return for gatherer_key, gatherer in self._single_pass_gatherers.items(): single_pass_data = gatherer.collect() - self._accumulator.append(forward_pass_id, gatherer_key, single_pass_data) + self._accumulator.append( + forward_pass_id, gatherer_key, single_pass_data, outputs + ) def on_select_experts(self, topk_ids: torch.Tensor): self._on_hook("on_select_experts", topk_ids=topk_ids) @@ -636,6 +648,7 @@ class _Accumulator(ABC): forward_pass_id: int, gatherer_key: str, single_pass_data: Dict, + outputs: Dict[str, Any], ): pass @@ -659,18 +672,19 @@ class _UtilizationRateAccumulatorMixin(_Accumulator): self._expert_dispatch_collector = ExpertDispatchCollector( self._expert_location_metadata.ep_size ) - self._collection_counter = 0 + self._metric_heatmap_collection_counter = 0 def append( self, forward_pass_id: int, gatherer_key: str, single_pass_data: Dict, + outputs: Dict[str, Any], ): - super().append(forward_pass_id, gatherer_key, single_pass_data) + super().append(forward_pass_id, gatherer_key, single_pass_data, outputs) if self._enable: - self._append_utilization_rate( - forward_pass_id, single_pass_data["global_physical_count"] + return self._append_utilization_rate( + forward_pass_id, single_pass_data["global_physical_count"], outputs ) def reset(self): @@ -679,7 +693,10 @@ class _UtilizationRateAccumulatorMixin(_Accumulator): self._history.clear() def _append_utilization_rate( - self, forward_pass_id: int, single_pass_global_physical_count: torch.Tensor + self, + forward_pass_id: int, + single_pass_global_physical_count: torch.Tensor, + outputs: Dict[str, Any], ): gpu_physical_count = compute_gpu_physical_count( single_pass_global_physical_count, @@ -691,27 +708,37 @@ class _UtilizationRateAccumulatorMixin(_Accumulator): ) if self._rank == 0: - self._collect_metrics_if_needed(gpu_physical_count) + self._handle_metric_eplb_heatmap(gpu_physical_count) - utilization_rate_tensor = compute_utilization_rate(gpu_physical_count) - utilization_rate = torch.mean(utilization_rate_tensor).item() - self._history.append(utilization_rate) - - 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:.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()]}" + utilization_rate_gpu = torch.mean( + compute_utilization_rate(gpu_physical_count) ) + if envs.SGLANG_ENABLE_EPLB_BALANCEDNESS_METRIC.get(): + print(f"hi {self._rank=} {utilization_rate_gpu=}") + outputs["metrics"] = ExpertDistributionMetrics( + eplb_balancedness=utilization_rate_gpu, + ) + 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) - def _collect_metrics_if_needed(self, gpu_physical_count: torch.Tensor): + 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): # sglang:eplb_gpu_physical_count metric is disabled if SGLANG_EPLB_HEATMAP_COLLECTION_INTERVAL <= 0 interval = get_int_env_var("SGLANG_EPLB_HEATMAP_COLLECTION_INTERVAL", 0) - if interval > 0 and self._collection_counter % interval == 0: + if interval > 0 and self._metric_heatmap_collection_counter % interval == 0: for layer_idx in range(self._expert_location_metadata.num_layers): count_of_layer = ( self._expert_dispatch_collector.eplb_gpu_physical_count.labels( @@ -728,7 +755,7 @@ class _UtilizationRateAccumulatorMixin(_Accumulator): if count > 0: count_of_layer._sum.inc(count * gpu_rank) count_of_layer._buckets[gpu_rank].inc(count) - self._collection_counter += 1 + self._metric_heatmap_collection_counter += 1 class _DequeCollection: @@ -767,8 +794,9 @@ class _DetailAccumulator(_UtilizationRateAccumulatorMixin): forward_pass_id: int, gatherer_key: str, single_pass_data: Dict, + outputs: Dict[str, Any], ): - super().append(forward_pass_id, gatherer_key, single_pass_data) + super().append(forward_pass_id, gatherer_key, single_pass_data, outputs) def _process_object(obj): if isinstance(obj, torch.Tensor): @@ -824,8 +852,9 @@ class _StatAccumulator(_UtilizationRateAccumulatorMixin): forward_pass_id: int, gatherer_key: str, single_pass_data: Dict, + outputs: Dict[str, Any], ): - super().append(forward_pass_id, gatherer_key, single_pass_data) + super().append(forward_pass_id, gatherer_key, single_pass_data, outputs) # Can optimize if overhead here is large self._global_physical_count_of_buffered_step.append( single_pass_data["global_physical_count"] diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 983093d00..8e62c3f7a 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -2221,6 +2221,7 @@ class Scheduler( if result.copy_done is not None: result.copy_done.synchronize() + self.log_batch_result_stats(batch, result) self.maybe_send_health_check_signal() def maybe_send_health_check_signal(self): diff --git a/python/sglang/srt/managers/scheduler_metrics_mixin.py b/python/sglang/srt/managers/scheduler_metrics_mixin.py index a4d884ab5..57cb3f0c5 100644 --- a/python/sglang/srt/managers/scheduler_metrics_mixin.py +++ b/python/sglang/srt/managers/scheduler_metrics_mixin.py @@ -4,7 +4,7 @@ import logging import time from collections import defaultdict from contextlib import contextmanager -from typing import TYPE_CHECKING, List, Optional +from typing import TYPE_CHECKING, List, Optional, Union from sglang.srt.disaggregation.kv_events import EventPublisherFactory, KVEventBatch from sglang.srt.disaggregation.utils import DisaggregationMode @@ -12,12 +12,13 @@ from sglang.srt.environ import envs from sglang.srt.managers.io_struct import GetLoadReqInput, GetLoadReqOutput from sglang.srt.managers.schedule_policy import PrefillAdder from sglang.srt.managers.scheduler import Req, ScheduleBatch +from sglang.srt.managers.utils import GenerationBatchResult from sglang.srt.metrics.collector import SchedulerMetricsCollector, SchedulerStats from sglang.srt.utils import get_bool_env_var from sglang.srt.utils.device_timer import DeviceTimer if TYPE_CHECKING: - from sglang.srt.managers.scheduler import Scheduler + from sglang.srt.managers.scheduler import EmbeddingBatchResult, Scheduler logger = logging.getLogger(__name__) @@ -395,6 +396,22 @@ class SchedulerMetricsMixin: self._emit_kv_metrics() self._publish_kv_events() + def log_batch_result_stats( + self: Scheduler, + 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(), + ) + def _emit_kv_metrics(self: Scheduler): if not self.enable_kv_cache_events: return diff --git a/python/sglang/srt/managers/tp_worker.py b/python/sglang/srt/managers/tp_worker.py index 0a132c709..dff873edd 100644 --- a/python/sglang/srt/managers/tp_worker.py +++ b/python/sglang/srt/managers/tp_worker.py @@ -406,6 +406,7 @@ class TpModelWorker(BaseTpWorker): batch_result = GenerationBatchResult( logits_output=logits_output, can_run_cuda_graph=can_run_cuda_graph, + expert_distribution_metrics=out.expert_distribution_metrics, ) if is_verify: @@ -460,6 +461,7 @@ class TpModelWorker(BaseTpWorker): return GenerationBatchResult( pp_hidden_states_proxy_tensors=pp_proxy_tensors, can_run_cuda_graph=can_run_cuda_graph, + expert_distribution_metrics=out.expert_distribution_metrics, ) def forward_batch_split_prefill(self, batch: ScheduleBatch): @@ -482,6 +484,7 @@ class TpModelWorker(BaseTpWorker): batch_result = GenerationBatchResult( logits_output=logits_output, can_run_cuda_graph=can_run_cuda_graph, + expert_distribution_metrics=out.expert_distribution_metrics, ) batch_result.next_token_ids = next_token_ids return batch_result diff --git a/python/sglang/srt/managers/utils.py b/python/sglang/srt/managers/utils.py index 06469a061..cdbb9dec8 100644 --- a/python/sglang/srt/managers/utils.py +++ b/python/sglang/srt/managers/utils.py @@ -6,6 +6,7 @@ from typing import TYPE_CHECKING, List, Optional import torch +from sglang.srt.eplb.expert_distribution import ExpertDistributionMetrics from sglang.srt.layers.logits_processor import LogitsProcessorOutput from sglang.srt.managers.overlap_utils import FutureIndices from sglang.srt.managers.schedule_batch import Req @@ -44,6 +45,9 @@ class GenerationBatchResult: # relay path: forward stream -> next step forward next_draft_input: Optional[EagleDraftInput] = None + # metrics + expert_distribution_metrics: Optional[ExpertDistributionMetrics] = None + def copy_to_cpu(self, return_logprob: bool): """Copy tensors to CPU in overlap scheduling. Only the tensors which are needed for processing results are copied, @@ -67,6 +71,9 @@ class GenerationBatchResult: if self.accept_lens is not None: self.accept_lens = self.accept_lens.to("cpu", non_blocking=True) + if (x := self.expert_distribution_metrics) is not None: + x.copy_to_cpu() + self.copy_done.record() @classmethod diff --git a/python/sglang/srt/metrics/collector.py b/python/sglang/srt/metrics/collector.py index 65c856041..e824de88d 100644 --- a/python/sglang/srt/metrics/collector.py +++ b/python/sglang/srt/metrics/collector.py @@ -19,6 +19,7 @@ from dataclasses import dataclass, field from typing import Dict, List, Optional, Union from sglang.srt.disaggregation.utils import DisaggregationMode +from sglang.srt.environ import envs from sglang.srt.metrics.utils import exponential_buckets, generate_buckets from sglang.srt.server_args import ServerArgs from sglang.srt.utils import get_bool_env_var @@ -241,7 +242,7 @@ class SchedulerMetricsCollector: labels: Dict[str, str], ) -> None: # We need to import prometheus_client after setting the env variable `PROMETHEUS_MULTIPROC_DIR` - from prometheus_client import Counter, Gauge, Histogram + from prometheus_client import Counter, Gauge, Histogram, Summary self.labels = labels self.last_log_time = time.perf_counter() @@ -641,6 +642,15 @@ class SchedulerMetricsCollector: labelnames=list(labels.keys()) + ["mode"], ) + if ( + labels["moe_ep_rank"] == 0 + ) and envs.SGLANG_ENABLE_EPLB_BALANCEDNESS_METRIC.get(): + self.eplb_balancedness = Summary( + name="sglang:eplb_balancedness", + documentation="Balancedness of MoE in expert parallelism.", + labelnames=list(labels.keys()) + ["forward_mode"], + ) + self.new_token_ratio = Gauge( name="sglang:new_token_ratio", documentation="The new token ratio.", @@ -698,6 +708,13 @@ class SchedulerMetricsCollector: mode = "decode_cuda_graph" if value else "decode_none" self.cuda_graph_passes_total.labels(**self.labels, mode=mode).inc(1) + def increment_eplb_balancedness( + self, forward_mode: str, balancedness: float + ) -> None: + self.eplb_balancedness.labels(**self.labels, forward_mode=forward_mode).observe( + balancedness + ) + def increment_realtime_tokens( self, prefill_compute_tokens=0, prefill_cache_tokens=0, decode_tokens=0 ): diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 4ab0240dd..479723d74 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -68,6 +68,7 @@ from sglang.srt.elastic_ep.elastic_ep import ElasticEPStateManager from sglang.srt.environ import envs from sglang.srt.eplb.eplb_manager import EPLBManager from sglang.srt.eplb.expert_distribution import ( + ExpertDistributionMetrics, ExpertDistributionRecorder, get_global_expert_distribution_recorder, set_global_expert_distribution_recorder, @@ -272,6 +273,7 @@ class RankZeroFilter(logging.Filter): class ModelRunnerOutput: logits_output: Union[LogitsProcessorOutput, PPProxyTensors] can_run_graph: bool + expert_distribution_metrics: Optional[ExpertDistributionMetrics] = None class ModelRunner: @@ -2738,7 +2740,7 @@ class ModelRunner: with get_global_expert_distribution_recorder().with_forward_pass( self.forward_pass_id, forward_batch, - ): + ) as recorder_outputs: output = self._forward_raw( forward_batch, skip_attn_backend_init, @@ -2746,6 +2748,7 @@ class ModelRunner: reinit_attn_backend, split_forward_count, ) + output.expert_distribution_metrics = recorder_outputs.get("metrics") if self.eplb_manager is not None: self.eplb_manager.on_forward_pass_end()