A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 16:00 UTC
[Bugfix][V1] Fix bug from putting llm_engine.model_executor in a background process (#15367)
Signed-off-by: wwl2755 <wangwenlong2755@gmail.com>
wwl2755 committed
Apr 3, 2025 at 02:32 UTC
463bbb1835b8bbb1a80e6286c6396f6f3182c4f7
5 files changed
+167
-4
examples/offline_inference/load_sharded_state.py
new
+93
@@ -0,0 +1,93 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+"""
3
+Validates the loading of a model saved with the sharded_state format.
4
+This script demonstrates how to load a model that was previously saved
5
+using save_sharded_state.py and validates it by running inference.
6
+Example usage:
7
+(First need to save a sharded_state mode)
8
+
9
+python save_sharded_state.py \
10
+ --model /path/to/load \
11
+ --quantization deepspeedfp \
12
+ --tensor-parallel-size 8 \
13
+ --output /path/to/save/sharded/modele
14
+
15
+python load_sharded_state.py \
16
+ --model /path/to/saved/sharded/model \
17
+ --load-format sharded_state \
18
+ --quantization deepspeedfp \
19
+ --tensor-parallel-size 8 \
20
+ --prompt "Hello, my name is" \
21
+ --max-tokens 50
22
+"""
23
+
24
+import dataclasses
25
+
26
+from vllm import LLM, EngineArgs, SamplingParams
27
+from vllm.utils import FlexibleArgumentParser
28
+
29
+
30
+def parse_args():
31
+ parser = FlexibleArgumentParser()
32
+ # Add engine arguments
33
+ EngineArgs.add_cli_args(parser)
34
+
35
+ # Override default load_format for clarity
36
+ parser.set_defaults(load_format="sharded_state")
37
+
38
+ # Add validation arguments
39
+ parser.add_argument("--prompt",
40
+ type=str,
41
+ default="Hello, world!",
42
+ help="Prompt for validation")
43
+ parser.add_argument("--max-tokens",
44
+ type=int,
45
+ default=100,
46
+ help="Maximum number of tokens to generate")
47
+ parser.add_argument("--temperature",
48
+ type=float,
49
+ default=0.7,
50
+ help="Sampling temperature")
51
+ parser.add_argument("--top-p",
52
+ type=float,
53
+ default=1.0,
54
+ help="Top-p sampling parameter")
55
+
56
+ return parser.parse_args()
57
+
58
+
59
+def main():
60
+ args = parse_args()
61
+ engine_args = EngineArgs.from_cli_args(args)
62
+
63
+ print(f"Loading model from {engine_args.model} "
64
+ f"using format {engine_args.load_format}")
65
+ print(f"Tensor parallel size: {engine_args.tensor_parallel_size}")
66
+
67
+ # Load the model using engine args
68
+ llm = LLM(**dataclasses.asdict(engine_args))
69
+
70
+ # Prepare sampling parameters
71
+ sampling_params = SamplingParams(
72
+ temperature=args.temperature,
73
+ top_p=args.top_p,
74
+ max_tokens=args.max_tokens,
75
+ )
76
+
77
+ print("\nRunning inference:")
78
+ print(f"Prompt: {args.prompt}")
79
+
80
+ # Generate completion
81
+ outputs = llm.generate(args.prompt, sampling_params)
82
+
83
+ # Display generated text
84
+ print("\nGenerated outputs:")
85
+ for output in outputs:
86
+ generated_text = output.outputs[0].text
87
+ print("-" * 50)
88
+ print(f"Full output: {args.prompt}{generated_text}")
89
+ print("-" * 50)
90
+
91
+
92
+if __name__ == "__main__":
93
+ main()
\ No newline at end of file
examples/offline_inference/save_sharded_state.py
+19
-4
@@ -57,10 +57,25 @@ def main(args):
57
# Prepare output directory
58
Path(args.output).mkdir(exist_ok=True)
59
# Dump worker states to output directory
60
- model_executor = llm.llm_engine.model_executor
61
- model_executor.save_sharded_state(path=args.output,
62
- pattern=args.file_pattern,
63
- max_size=args.max_file_size)
60
+
61
+ # Check which engine version is being used
62
+ is_v1_engine = hasattr(llm.llm_engine, "engine_core")
63
+
64
+ if is_v1_engine:
65
+ # For V1 engine, we need to use engine_core.save_sharded_state
66
+ print("Using V1 engine save path")
67
+ llm.llm_engine.engine_core.save_sharded_state(
68
+ path=args.output,
69
+ pattern=args.file_pattern,
70
+ max_size=args.max_file_size)
71
+ else:
72
+ # For V0 engine
73
+ print("Using V0 engine save path")
74
+ model_executor = llm.llm_engine.model_executor
75
+ model_executor.save_sharded_state(path=args.output,
76
+ pattern=args.file_pattern,
77
+ max_size=args.max_file_size)
78
+
79
# Copy metadata files to output directory
80
for file in os.listdir(model_path):
81
if os.path.splitext(file)[1] not in (".bin", ".pt", ".safetensors"):
vllm/v1/engine/core.py
+10
@@ -285,6 +285,16 @@ class EngineCore:
285
def pin_lora(self, lora_id: int) -> bool:
286
return self.model_executor.pin_lora(lora_id)
287
288
+ def save_sharded_state(
289
+ self,
290
+ path: str,
291
+ pattern: Optional[str] = None,
292
+ max_size: Optional[int] = None,
293
+ ) -> None:
294
+ self.model_executor.save_sharded_state(path=path,
295
+ pattern=pattern,
296
+ max_size=max_size)
297
+
298
def collective_rpc(self,
299
method: Union[str, Callable[..., _R]],
300
timeout: Optional[float] = None,
vllm/v1/engine/core_client.py
+31
@@ -119,6 +119,12 @@ class EngineCoreClient(ABC):
119
def pin_lora(self, lora_id: int) -> bool:
120
raise NotImplementedError
121
122
+ def save_sharded_state(self,
123
+ path: str,
124
+ pattern: Optional[str] = None,
125
+ max_size: Optional[int] = None) -> None:
126
+ raise NotImplementedError
127
+
128
def collective_rpc(self,
129
method: Union[str, Callable[..., _R]],
130
timeout: Optional[float] = None,
@@ -162,6 +168,12 @@ class EngineCoreClient(ABC):
168
async def pin_lora_async(self, lora_id: int) -> bool:
169
raise NotImplementedError
170
171
+ async def save_sharded_state_async(self,
172
+ path: str,
173
+ pattern: Optional[str] = None,
174
+ max_size: Optional[int] = None) -> None:
175
+ raise NotImplementedError
176
+
177
async def collective_rpc_async(
178
self,
179
method: Union[str, Callable[..., _R]],
@@ -227,6 +239,12 @@ class InprocClient(EngineCoreClient):
239
def pin_lora(self, lora_id: int) -> bool:
240
return self.engine_core.pin_lora(lora_id)
241
242
+ def save_sharded_state(self,
243
+ path: str,
244
+ pattern: Optional[str] = None,
245
+ max_size: Optional[int] = None) -> None:
246
+ self.engine_core.save_sharded_state(path, pattern, max_size)
247
+
248
def collective_rpc(self,
249
method: Union[str, Callable[..., _R]],
250
timeout: Optional[float] = None,
@@ -537,6 +555,12 @@ class SyncMPClient(MPClient):
555
return self.call_utility("collective_rpc", method, timeout, args,
556
kwargs)
557
558
+ def save_sharded_state(self,
559
+ path: str,
560
+ pattern: Optional[str] = None,
561
+ max_size: Optional[int] = None) -> None:
562
+ self.call_utility("save_sharded_state", path, pattern, max_size)
563
+
564
565
class AsyncMPClient(MPClient):
566
"""Asyncio-compatible client for multi-proc EngineCore."""
@@ -668,6 +692,13 @@ class AsyncMPClient(MPClient):
692
async def pin_lora_async(self, lora_id: int) -> bool:
693
return await self.call_utility_async("pin_lora", lora_id)
694
695
+ async def save_sharded_state_async(self,
696
+ path: str,
697
+ pattern: Optional[str] = None,
698
+ max_size: Optional[int] = None) -> None:
699
+ await self.call_utility_async("save_sharded_state", path, pattern,
700
+ max_size)
701
+
702
async def collective_rpc_async(
703
self,
704
method: Union[str, Callable[..., _R]],
vllm/v1/worker/gpu_worker.py
+14
@@ -269,6 +269,20 @@ class Worker(WorkerBase):
269
# worker will always be healthy as long as it's running.
270
return
271
272
+ def save_sharded_state(
273
+ self,
274
+ path: str,
275
+ pattern: Optional[str] = None,
276
+ max_size: Optional[int] = None,
277
+ ) -> None:
278
+ from vllm.model_executor.model_loader.loader import ShardedStateLoader
279
+ ShardedStateLoader.save_model(
280
+ self.model_runner.model,
281
+ path,
282
+ pattern=pattern,
283
+ max_size=max_size,
284
+ )
285
+
286
287
def init_worker_distributed_environment(
288
parallel_config: ParallelConfig,