init v0.23.0

Signed-off-by: Sun Ruoxi <sunruoxi@4paradigm.com>

examples/offline_external_launcher.py

@@ -63,17 +63,47 @@ from multiprocessing import Process
 from time import sleep
 
 import torch
+from safetensors.torch import load_file
 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
+    destroy_distributed_environment,
+    destroy_model_parallel,
+    get_tp_group,
+)
+from vllm.utils.mem_constants import GiB_bytes
+from vllm.utils.network_utils import get_open_port
 
 os.environ["VLLM_USE_MODELSCOPE"] = "True"
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 
 
-def parse_args():
+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")
     parser.add_argument(
         "--model",
@@ -81,55 +111,41 @@ def parse_args():
         default="Qwen/Qwen3-0.6B",
         help="Model name or path",
     )
-    parser.add_argument("--tp-size",
-                        type=int,
-                        default=1,
-                        help="Tensor parallel size")
-    parser.add_argument("--node-size",
-                        type=int,
-                        default=1,
-                        help="Total number of nodes")
-    parser.add_argument("--node-rank",
-                        type=int,
-                        default=0,
-                        help="Rank of the current node")
-    parser.add_argument("--proc-per-node",
-                        type=int,
-                        default=1,
-                        help="Number of processes per node")
-    parser.add_argument("--master-addr",
-                        type=str,
-                        default="",
-                        help="Master node IP address")
-    parser.add_argument("--master-port",
-                        type=int,
-                        default=0,
-                        help="Master node port")
-    parser.add_argument("--enforce-eager",
-                        action="store_true",
-                        help="Enforce eager mode execution.")
-    parser.add_argument("--trust-remote-code",
-                        action="store_true",
-                        help="Trust remote code.")
-    parser.add_argument("--enable-expert-parallel",
-                        action="store_true",
-                        help="Enable expert parallel, used in MOE models.")
-    parser.add_argument("--enable-sleep-mode",
-                        action="store_true",
-                        help="Enable sleep mode for the engine.")
-    parser.add_argument("--temperature",
-                        type=float,
-                        default=0.8,
-                        help="Float that controls the randomness of the sampling.")
-    parser.add_argument("--model-weight-gib",
-                        type=float,
-                        default=None,
-                        help="Model weight memory usage in GiB (e.g., 1.0 for 0.5B model).")
+    parser.add_argument("--tp-size", type=int, default=1, help="Tensor parallel size")
+    parser.add_argument("--node-size", type=int, default=1, help="Total number of nodes")
+    parser.add_argument("--node-rank", type=int, default=0, help="Rank of the current node")
+    parser.add_argument("--proc-per-node", type=int, default=1, help="Number of processes per node")
+    parser.add_argument("--master-addr", type=str, default="", help="Master node IP address")
+    parser.add_argument("--master-port", type=int, default=0, help="Master node port")
+    parser.add_argument("--enforce-eager", action="store_true", help="Enforce eager mode execution.")
+    parser.add_argument("--trust-remote-code", action="store_true", help="Trust remote code.")
+    parser.add_argument(
+        "--enable-expert-parallel", action="store_true", help="Enable expert parallel, used in MOE models."
+    )
+    parser.add_argument("--enable-sleep-mode", action="store_true", help="Enable sleep mode for the engine.")
+    parser.add_argument(
+        "--temperature", type=float, default=0.8, help="Float that controls the randomness of the sampling."
+    )
+    parser.add_argument(
+        "--model-weight-gib",
+        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:
         if args.model_weight_gib is None or args.temperature != 0:
-            parser.error("model-weight-gib must be provided, and temperature must be zero when enable-sleep-mode is set.")
+            parser.error(
+                "model-weight-gib must be provided, and temperature must be zero when enable-sleep-mode is set."
+            )
         if args.model_weight_gib <= 0:
             parser.error("model-weight-gib must be greater than 0 when enable-sleep-mode is set.")
         if args.model == parser.get_default("model") and args.model_weight_gib is None:
@@ -152,6 +168,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)
@@ -186,22 +203,31 @@ def main(
         enable_sleep_mode=enable_sleep_mode,
     )
     tp_ranks = get_tp_group().ranks
-    print(f'TP RANKS: {tp_ranks}')
+    print(f"TP RANKS: {tp_ranks}")
 
     outputs = llm.generate(prompts, sampling_params)
 
     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
-            print(f"Freed memory: {freed_bytes / 1024 ** 3:.2f} GiB")
+            print(f"Freed memory: {freed_bytes / 1024**3:.2f} GiB")
             # 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
@@ -214,8 +240,7 @@ def main(
             break
         prompt = output.prompt
         generated_text = output.outputs[0].text
-        print(f"Global rank: {rank}, Prompt: {prompt!r}, "
-              f"Generated text: {generated_text!r}")
+        print(f"Global rank: {rank}, Prompt: {prompt!r}, Generated text: {generated_text!r}")
 
     # Give engines time to pause their processing loops before exiting.
     sleep(5)
@@ -251,24 +276,26 @@ if __name__ == "__main__":
     world_size = node_size * proc_per_node
 
     procs = []
-    for local_rank, rank in enumerate(
-            range(proc_per_node * node_rank, proc_per_node * (node_rank + 1))):
-        proc = Process(target=main,
-                       args=(
-                           local_rank,
-                           rank,
-                           master_addr,
-                           master_port,
-                           args.model_weight_gib,
-                           args.model,
-                           world_size,
-                           tp_size,
-                           args.enable_expert_parallel,
-                           args.enforce_eager,
-                           args.trust_remote_code,
-                           args.enable_sleep_mode,
-                           args.temperature,
-                       ))
+    for local_rank, rank in enumerate(range(proc_per_node * node_rank, proc_per_node * (node_rank + 1))):
+        proc = Process(
+            target=main,
+            args=(
+                local_rank,
+                rank,
+                master_addr,
+                master_port,
+                args.model_weight_gib,
+                args.model,
+                world_size,
+                tp_size,
+                args.enable_expert_parallel,
+                args.enforce_eager,
+                args.trust_remote_code,
+                args.enable_sleep_mode,
+                args.temperature,
+                args.sleep_mode_level,
+            ),
+        )
 
         proc.start()
         procs.append(proc)
@@ -276,9 +303,7 @@ if __name__ == "__main__":
     for proc in procs:
         proc.join(timeout=600)
         if proc.exitcode is None:
-            print(
-                f"Killing process {proc.pid} that didn't stop within 30 minutes."
-            )
+            print(f"Killing process {proc.pid} that didn't stop within 30 minutes.")
             proc.kill()
             exit_code = 1
         elif proc.exitcode: