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 14, 05:30 UTC
[Core] Enable sharded state loader for V1 engine and enhance test coverage (#25308)
Signed-off-by: pengdrumli <pengdrumli@tencent.com>
lirong committed
Sep 20, 2025 at 21:15 UTC
d88918e4c2d189eb25724a6cff9a9028c67daa07
2 files changed
+12
-14
tests/test_sharded_state_loader.py
+12
-8
@@ -57,10 +57,19 @@ def llama_3p2_1b_files():
57
58
def _run_writer(input_dir, output_dir, weights_patterns, **kwargs):
59
llm_sharded_writer = LLM(model=input_dir, **kwargs)
60
-
60
+ # Check which engine version is being used
61
+ is_v1_engine = hasattr(llm_sharded_writer.llm_engine, "engine_core")
62
# Dump worker states to output directory
62
- llm_sharded_writer.llm_engine.model_executor.save_sharded_state(
63
- path=output_dir)
63
+ if is_v1_engine:
64
+ # For V1 engine, we need to use engine_core.save_sharded_state
65
+ print("Using V1 engine save path")
66
+ llm_sharded_writer.llm_engine.engine_core.save_sharded_state(
67
+ path=output_dir)
68
+ else:
69
+ # For V0 engine
70
+ print("Using V0 engine save path")
71
+ model_executor = llm_sharded_writer.llm_engine.model_executor
72
+ model_executor.save_sharded_state(path=output_dir)
73
74
# Copy metadata files to output directory
75
for file in os.listdir(input_dir):
@@ -91,8 +100,6 @@ def test_sharded_state_loader(enable_lora, tp_size, num_gpus_available,
100
gpu_memory_utilization = 0.8
101
input_dir = llama_3p2_1b_files
102
ctx = mp.get_context("spawn")
94
- # The interface in v1 engine has changed, run in v1 engine will hang.
95
- monkeypatch.setenv("VLLM_USE_V1", "0")
103
104
# Run in separate processes for memory & CUDA isolation
105
with TemporaryDirectory() as output_dir:
@@ -100,7 +107,6 @@ def test_sharded_state_loader(enable_lora, tp_size, num_gpus_available,
107
args=(input_dir, output_dir, weights_patterns),
108
kwargs=dict(
109
tensor_parallel_size=tp_size,
103
- distributed_executor_backend="mp",
110
gpu_memory_utilization=gpu_memory_utilization,
111
enforce_eager=True,
112
))
@@ -112,7 +118,6 @@ def test_sharded_state_loader(enable_lora, tp_size, num_gpus_available,
118
p = ctx.Process(target=_run_generate,
119
args=(input_dir, queue),
120
kwargs=dict(
115
- distributed_executor_backend="mp",
121
enable_lora=enable_lora,
122
gpu_memory_utilization=gpu_memory_utilization,
123
tensor_parallel_size=tp_size,
@@ -133,7 +138,6 @@ def test_sharded_state_loader(enable_lora, tp_size, num_gpus_available,
138
p = ctx.Process(target=_run_generate,
139
args=(output_dir, queue),
140
kwargs=dict(
136
- distributed_executor_backend="mp",
141
enable_lora=enable_lora,
142
gpu_memory_utilization=gpu_memory_utilization,
143
tensor_parallel_size=tp_size,
vllm/engine/arg_utils.py
-6
@@ -1486,12 +1486,6 @@ class EngineArgs:
1486
#############################################################
1487
# Unsupported Feature Flags on V1.
1488
1489
- if self.load_format == "sharded_state":
1490
- _raise_or_fallback(
1491
- feature_name=f"--load_format {self.load_format}",
1492
- recommend_to_remove=False)
1493
- return False
1494
-
1489
if (self.logits_processor_pattern
1490
!= EngineArgs.logits_processor_pattern):
1491
_raise_or_fallback(feature_name="--logits-processor-pattern",