don't double eplb prefix

Signed-off-by: Patryk Saffer <[email protected]>

Author: Patryk999

tests/v1/e2e/test_eplb_offline.py

@@ -48,7 +48,7 @@ def test_eplb_model(
         )
 
         # Save EPLB statistics to disk
-        eplb_config_save = EPLBConfig(eplb_save_dir="/tmp")
+        eplb_config_save = EPLBConfig(save_dir="/tmp")
         llm = LLM(eplb_config=eplb_config_save, **llm_args)
         llm.generate(test_prompts, sampling_config)
         del llm
@@ -57,7 +57,7 @@ def test_eplb_model(
 
         # Load EPLB statistics from disk
         eplb_config_load = EPLBConfig(
-            eplb_load_path="/tmp/global_expert_load_window_i0.safetensors"
+            load_path="/tmp/global_expert_load_window_i0.safetensors"
         )
         llm = LLM(eplb_config=eplb_config_load, **llm_args)
         llm.generate(test_prompts, sampling_config)

vllm/config/parallel.py

@@ -62,15 +62,15 @@ class EPLBConfig:
     This is turned off by default since it will cause communication overhead.
     """
 
-    eplb_save_dir: Path | None = None
+    save_dir: Path | None = None
     """Directory to save expert load balance metrics."""
-    eplb_load_path: Path | None = None
+    load_path: Path | None = None
     """Path to load expert load balance metrics."""
 
     @property
-    def eplb_record_metrics(self) -> bool:
-        return self.eplb_save_dir is not None or (
-            self.eplb_save_dir is None and self.eplb_load_path is None
+    def record_metrics(self) -> bool:
+        return self.save_dir is not None or (
+            self.save_dir is None and self.load_path is None
         )
 
 

vllm/distributed/eplb/eplb_state.py

@@ -151,11 +151,11 @@ class EplbState:
     See:
     https://github.com/vllm-project/vllm/pull/22167#pullrequestreview-3086143856
     """
-    eplb_load_path: Path | None = None
+    load_path: Path | None = None
     """
     Path for loading eplb initial state.
     """
-    eplb_save_dir: Path | None = None
+    save_dir: Path | None = None
     """
     Path where eplb states will be saved.
     """
@@ -298,7 +298,7 @@ def build(
         )
 
         eplb_load_path = parallel_config.eplb_config.eplb_load_path
-        eplb_save_dir = parallel_config.eplb_config.eplb_save_dir
+        eplb_save_dir = parallel_config.eplb_config.save_dir
 
         eplb_step_interval = parallel_config.eplb_config.step_interval
         if eplb_load_path is not None or eplb_save_dir is not None:
@@ -498,8 +498,8 @@ def rearrange(
             time_start = time.perf_counter()
             logger.info("Rearranging experts %s...", "(profile)" if is_profile else "")
 
-        if self.eplb_load_path is not None and self.expert_rearrangement_step == 0:
-            global_expert_load_window = load_eplb_state(self.eplb_load_path).to(
+        if self.load_path is not None and self.expert_rearrangement_step == 0:
+            global_expert_load_window = load_eplb_state(self.load_path).to(
                 self.physical_to_logical_map.device
             )
         elif global_expert_load is None:

vllm/engine/arg_utils.py

@@ -400,8 +400,8 @@ class EngineArgs:
     eplb_window_size: int = EPLBConfig.window_size
     eplb_step_interval: int = EPLBConfig.step_interval
     eplb_log_balancedness: bool = EPLBConfig.log_balancedness
-    eplb_save_dir: Path | None = EPLBConfig.eplb_save_dir
-    eplb_load_path: Path | None = EPLBConfig.eplb_load_path
+    eplb_save_dir: Path | None = EPLBConfig.save_dir
+    eplb_load_path: Path | None = EPLBConfig.load_path
     max_parallel_loading_workers: int | None = (
         ParallelConfig.max_parallel_loading_workers
     )

vllm/model_executor/models/transformers/moe.py

@@ -236,7 +236,7 @@ def forward(self, *args, **kwargs):
         # Expert parallel load balancing kwargs
         enable_eplb = self.parallel_config.enable_eplb
         num_redundant_experts = self.parallel_config.eplb_config.num_redundant_experts
-        eplb_record_metrics = self.parallel_config.eplb_config.eplb_record_metrics
+        eplb_record_metrics = self.parallel_config.eplb_config.record_metrics
 
         # MixtureOfExperts mixin settings
         ep_size = get_ep_group().world_size