[Test] enable external launcher and add e2e test for sleep mode in level2 (#3344)

### What this PR does / why we need it?
1. Enable tests/e2e/multicard/test_external_launcher.py
2. Add e2e test for  sleep mode in level2

### Does this PR introduce _any_ user-facing change?
not involved

### How was this patch tested?
CI passed with existing test.

- vLLM version: v0.11.0rc3
- vLLM main: https://github.com/vllm-project/vllm/commit/v0.11.0

Signed-off-by: huangxialu <huangxialu1@huawei.com>
Co-authored-by: Shangwei-Li <lishangwei2@huawei.com>
This commit is contained in:
huangxialu
2025-10-11 17:29:38 +08:00
committed by GitHub
parent ecb1713dfc
commit e8c871ed0a
3 changed files with 101 additions and 5 deletions

View File

@@ -67,11 +67,38 @@ from vllm import LLM, SamplingParams
from vllm.distributed.parallel_state import ( # noqa E402
destroy_distributed_environment, destroy_model_parallel, get_tp_group)
from vllm.utils import get_open_port, GiB_bytes
from safetensors.torch import load_file
os.environ["VLLM_USE_MODELSCOPE"] = "True"
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
def patch_vllm_moe_model_weight_loader(model):
model = getattr(model, "model", None) or getattr(model, "language_model", None)
if model is None:
raise ValueError("The provided model does not have a valid 'model' or 'language_model' attribute.")
for layer in model.layers:
mlp_attr = "mlp"
mlp = getattr(layer, mlp_attr)
param_dict = dict(mlp.named_parameters())
for name, param in param_dict.items():
if "w13_weight" in name or "w2_weight" in name:
param.weight_loader = mlp.experts.weight_loader
def load_and_merge_safetensors(directory):
if not os.path.isdir(directory):
raise ValueError(f"The provided directory does not exist: {directory}")
merged_dict = {}
for filename in os.listdir(directory):
if filename.endswith(".safetensors"):
file_path = os.path.join(directory, filename)
print(f"loading file: {file_path}")
f = load_file(file_path)
merged_dict.update(f)
return merged_dict
def parse_args():
parser = argparse.ArgumentParser(description="External launcher Inference")
@@ -125,6 +152,11 @@ def parse_args():
type=float,
default=None,
help="Model weight memory usage in GiB (e.g., 1.0 for 0.5B model).")
parser.add_argument("--sleep-mode-level",
type=int,
choices=[1, 2],
default=1,
help="Sleep mode level: 1 or 2. This example of level 2 is only supported for dense model.")
args = parser.parse_args()
if args.enable_sleep_mode:
@@ -152,6 +184,7 @@ def main(
trust_remote_code: bool = True,
enable_sleep_mode: bool = False,
temperature: float = 0.8,
sleep_mode_level: int = 1,
):
os.environ["MASTER_ADDR"] = master_addr
os.environ["MASTER_PORT"] = str(master_port)
@@ -193,7 +226,7 @@ def main(
if enable_sleep_mode:
if rank == 0:
free_bytes_before_sleep, total = torch.npu.mem_get_info()
llm.sleep(level=1)
llm.sleep(level=sleep_mode_level)
if rank == 0:
free_bytes_after_sleep, total = torch.npu.mem_get_info()
freed_bytes = free_bytes_after_sleep - free_bytes_before_sleep
@@ -201,7 +234,16 @@ def main(
# now the freed memory should be larger than the model weights
assert freed_bytes >= model_weight_gib / tensor_parallel_size * GiB_bytes
llm.wake_up()
if sleep_mode_level == 1:
llm.wake_up()
else:
llm.wake_up(tags=["weights"])
run_model = llm.llm_engine.model_executor.driver_worker.worker.model_runner.model
patch_vllm_moe_model_weight_loader(run_model)
sd = load_and_merge_safetensors(model)
run_model.load_weights(sd.items())
llm.wake_up(tags=["kv_cache"])
outputs_after_wakeup = llm.generate(prompts, sampling_params)
if rank == 0:
# cmp output
@@ -268,6 +310,7 @@ if __name__ == "__main__":
args.trust_remote_code,
args.enable_sleep_mode,
args.temperature,
args.sleep_mode_level,
))
proc.start()