[Usage]: Does vLLM compile draft model?

[ngsitrong26]

I'm not sure vLLM compiles draft model. Does vLLM compile both the model base and draft model following compilation configuration?

[youkaichao]

Right now it's only compiling the model base, not the draft model.

[BKM1804]

INFO 02-18 04:51:36 model_runner.py:1116] Loading model weights took 1.4123 GB
INFO 02-18 04:51:36 spec_decode_worker.py:335] [Speculative Decoding] Use batch expansion for scoring proposals.
INFO 02-18 04:51:41 backends.py:577] Using cache directory: /root/.cache/vllm/torch_compile_cache/800ea12410/rank_0 for vLLM's torch.compile
INFO 02-18 04:51:41 backends.py:585] Dynamo bytecode transform time: 5.63 s
INFO 02-18 04:51:45 backends.py:207] Directly lookup the graph for shape None from the cache
INFO 02-18 04:51:45 backends.py:321] Compiling a graph for general shape takes 0.34 s
INFO 02-18 04:51:47 monitor.py:31] torch.compile takes 5.96 s in total
INFO 02-18 04:51:48 worker.py:266] Memory profiling takes 11.94 seconds
INFO 02-18 04:51:48 worker.py:266] the current vLLM instance can use total_gpu_memory (23.55GiB) x gpu_memory_utilization (0.96) = 22.61GiB
INFO 02-18 04:51:48 worker.py:266] model weights take 6.02GiB; non_torch_memory takes 0.10GiB; PyTorch activation peak memory takes 1.96GiB; the rest of the memory reserved for KV Cache is 14.54GiB.
WARNING 02-18 04:51:48 config.py:991] compressed-tensors MLA support requires fp8 activations and weights in group 'group_0', but got activations type 'float' and weights type 'float'.
WARNING 02-18 04:51:48 config.py:991]  Full config: {'config_groups': {'group_0': {'input_activations': {'block_structure': None, 'dynamic': False, 'group_size': None, 'num_bits': 8, 'observer': 'minmax', 'observer_kwargs': {}, 'strategy': 'tensor', 'symmetric': True, 'type': 'float'}, 'output_activations': None, 'targets': ['Linear'], 'weights': {'block_structure': None, 'dynamic': False, 'group_size': None, 'num_bits': 8, 'observer': 'minmax', 'observer_kwargs': {}, 'strategy': 'tensor', 'symmetric': True, 'type': 'float'}}}, 'format': 'float-quantized', 'global_compression_ratio': 1.4516128438127633, 'ignore': ['lm_head'], 'kv_cache_scheme': None, 'quant_method': 'compressed-tensors', 'quantization_status': 'frozen'}
WARNING 02-18 04:51:48 config.py:991] compressed-tensors MLA support requires fp8 activations and weights in group 'group_0', but got activations type 'float' and weights type 'float'.
WARNING 02-18 04:51:48 config.py:991]  Full config: {'config_groups': {'group_0': {'input_activations': {'block_structure': None, 'dynamic': False, 'group_size': None, 'num_bits': 8, 'observer': 'minmax', 'observer_kwargs': {}, 'strategy': 'tensor', 'symmetric': True, 'type': 'float'}, 'output_activations': None, 'targets': ['Linear'], 'weights': {'block_structure': None, 'dynamic': False, 'group_size': None, 'num_bits': 8, 'observer': 'minmax', 'observer_kwargs': {}, 'strategy': 'tensor', 'symmetric': True, 'type': 'float'}}}, 'format': 'float-quantized', 'global_compression_ratio': 1.4516128438127633, 'ignore': ['lm_head'], 'kv_cache_scheme': None, 'quant_method': 'compressed-tensors', 'quantization_status': 'frozen'}
INFO 02-18 04:51:48 executor_base.py:108] # CUDA blocks: 6615, # CPU blocks: 2340
INFO 02-18 04:51:48 executor_base.py:113] Maximum concurrency for 20000 tokens per request: 5.29x
INFO 02-18 04:51:53 model_runner.py:1435] Starting to compile model unsloth/Llama-3.2-3B-Instruct... ###############################################################################
INFO 02-18 04:51:53 model_runner.py:1436] Capturing cudagraphs for decoding. This may lead to unexpected consequences if the model is not static. To run the model in eager mode, set 'enforce_eager=True' or use '--enforce-eager' in the CLI. If out-of-memory error occurs during cudagraph capture, consider decreasing `gpu_memory_utilization` or switching to eager mode. You can also reduce the `max_num_seqs` as needed to decrease memory usage.
Capturing CUDA graph shapes: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 35/35 [00:16<00:00,  2.16it/s]
INFO 02-18 04:52:09 model_runner.py:1564] Graph capturing finished in 16 secs, took 0.22 GiB
WARNING 02-18 04:52:09 config.py:991] compressed-tensors MLA support requires fp8 activations and weights in group 'group_0', but got activations type 'float' and weights type 'float'.
WARNING 02-18 04:52:09 config.py:991]  Full config: {'config_groups': {'group_0': {'input_activations': {'block_structure': None, 'dynamic': False, 'group_size': None, 'num_bits': 8, 'observer': 'minmax', 'observer_kwargs': {}, 'strategy': 'tensor', 'symmetric': True, 'type': 'float'}, 'output_activations': None, 'targets': ['Linear'], 'weights': {'block_structure': None, 'dynamic': False, 'group_size': None, 'num_bits': 8, 'observer': 'minmax', 'observer_kwargs': {}, 'strategy': 'tensor', 'symmetric': True, 'type': 'float'}}}, 'format': 'float-quantized', 'global_compression_ratio': 1.4516128438127633, 'ignore': ['lm_head'], 'kv_cache_scheme': None, 'quant_method': 'compressed-tensors', 'quantization_status': 'frozen'}
WARNING 02-18 04:52:09 config.py:991] compressed-tensors MLA support requires fp8 activations and weights in group 'group_0', but got activations type 'float' and weights type 'float'.
WARNING 02-18 04:52:09 config.py:991]  Full config: {'config_groups': {'group_0': {'input_activations': {'block_structure': None, 'dynamic': False, 'group_size': None, 'num_bits': 8, 'observer': 'minmax', 'observer_kwargs': {}, 'strategy': 'tensor', 'symmetric': True, 'type': 'float'}, 'output_activations': None, 'targets': ['Linear'], 'weights': {'block_structure': None, 'dynamic': False, 'group_size': None, 'num_bits': 8, 'observer': 'minmax', 'observer_kwargs': {}, 'strategy': 'tensor', 'symmetric': True, 'type': 'float'}}}, 'format': 'float-quantized', 'global_compression_ratio': 1.4516128438127633, 'ignore': ['lm_head'], 'kv_cache_scheme': None, 'quant_method': 'compressed-tensors', 'quantization_status': 'frozen'}
INFO 02-18 04:52:11 model_runner.py:1435] Starting to compile model neuralmagic/Llama-3.2-1B-Instruct-FP8... ###############################################################################
INFO 02-18 04:52:11 model_runner.py:1436] Capturing cudagraphs for decoding. This may lead to unexpected consequences if the model is not static. To run the model in eager mode, set 'enforce_eager=True' or use '--enforce-eager' in the CLI. If out-of-memory error occurs during cudagraph capture, consider decreasing `gpu_memory_utilization` or switching to eager mode. You can also reduce the `max_num_seqs` as needed to decrease memory usage.
Capturing CUDA graph shapes:   0%|                                                                                                                                                                                                                                                 | 0/35 [00:00<?, ?it/s]INFO 02-18 04:52:15 backends.py:577] Using cache directory: /root/.cache/vllm/torch_compile_cache/a66abe35c3/rank_0 for vLLM's torch.compile
INFO 02-18 04:52:15 backends.py:585] Dynamo bytecode transform time: 3.77 s
INFO 02-18 04:52:17 backends.py:207] Directly lookup the graph for shape None from the cache
INFO 02-18 04:52:17 backends.py:321] Compiling a graph for general shape takes 0.23 s
INFO 02-18 04:52:18 monitor.py:31] torch.compile takes 4.00 s in total
Capturing CUDA graph shapes: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 35/35 [00:25<00:00,  1.35it/s]
INFO 02-18 04:52:37 model_runner.py:1564] Graph capturing finished in 26 secs, took 0.21 GiB

I try show log by add

logger.info("Starting to compile model %s... ###############################################################################", self.model_config.model)

in the capture_model function. I think it's compiling the draft model. @youkaichao

[ngsitrong26]

INFO 02-18 07:08:34 model_runner.py:1111] Starting to load model unsloth/Llama-3.2-3B-Instruct...
INFO 02-18 07:08:39 config.py:2959] Using VllmBackend with config: model='unsloth/Llama-3.2-3B-Instruct', speculative_config=SpeculativeConfig(draft_model='neuralmagic/Llama-3.2-1B-Instruct-FP8', num_spec_tokens=7), tokenizer='tau-vision/llama-tokenizer-fix', skip_tokenizer_init=False, tokenizer_mode=auto, revision=None, override_neuron_config=None, tokenizer_revision=None, trust_remote_code=False, dtype=torch.bfloat16, max_seq_len=20000, download_dir=None, load_format=auto, tensor_parallel_size=1, pipeline_parallel_size=1, disable_custom_all_reduce=False, quantization=None, enforce_eager=False, kv_cache_dtype=auto,  device_config=cuda, decoding_config=DecodingConfig(guided_decoding_backend='xgrammar'), observability_config=ObservabilityConfig(otlp_traces_endpoint=None, collect_model_forward_time=False, collect_model_execute_time=False), seed=0, served_model_name=vllm-api, num_scheduler_steps=1, multi_step_stream_outputs=True, enable_prefix_caching=True, chunked_prefill_enabled=False, use_async_output_proc=False, disable_mm_preprocessor_cache=False, mm_processor_kwargs=None, pooler_config=None, compilation_config={"level":3,"splitting_ops":[],"use_inductor":true,"compile_sizes":[],"use_cudagraph":false,"cudagraph_capture_sizes":[256,248,240,232,224,216,208,200,192,184,176,168,160,152,144,136,128,120,112,104,96,88,80,72,64,56,48,40,32,24,16,8,4,2,1],"max_capture_size":256}
INFO 02-18 07:08:40 weight_utils.py:251] Using model weights format ['*.safetensors']
Loading safetensors checkpoint shards:   0% Completed | 0/2 [00:00<?, ?it/s]
Loading safetensors checkpoint shards:  50% Completed | 1/2 [00:02<00:02,  2.03s/it]
Loading safetensors checkpoint shards: 100% Completed | 2/2 [00:08<00:00,  4.49s/it]
Loading safetensors checkpoint shards: 100% Completed | 2/2 [00:08<00:00,  4.13s/it]

INFO 02-18 07:08:52 model_runner.py:1116] Loading model weights took 6.0160 GB
INFO 02-18 07:08:52 model_runner.py:1111] Starting to load model neuralmagic/Llama-3.2-1B-Instruct-FP8...
INFO 02-18 07:08:53 config.py:2959] Using VllmBackend with config: model='neuralmagic/Llama-3.2-1B-Instruct-FP8', speculative_config=SpeculativeConfig(draft_model='neuralmagic/Llama-3.2-1B-Instruct-FP8', num_spec_tokens=7), tokenizer='tau-vision/llama-tokenizer-fix', skip_tokenizer_init=False, tokenizer_mode=auto, revision=None, override_neuron_config=None, tokenizer_revision=None, trust_remote_code=False, dtype=torch.bfloat16, max_seq_len=20000, download_dir=None, load_format=auto, tensor_parallel_size=1, pipeline_parallel_size=1, disable_custom_all_reduce=False, quantization=compressed-tensors, enforce_eager=False, kv_cache_dtype=auto,  device_config=cuda, decoding_config=DecodingConfig(guided_decoding_backend='xgrammar'), observability_config=ObservabilityConfig(otlp_traces_endpoint=None, collect_model_forward_time=False, collect_model_execute_time=False), seed=0, served_model_name=neuralmagic/Llama-3.2-1B-Instruct-FP8, num_scheduler_steps=1, multi_step_stream_outputs=True, enable_prefix_caching=True, chunked_prefill_enabled=False, use_async_output_proc=True, disable_mm_preprocessor_cache=False, mm_processor_kwargs=None, pooler_config=None, compilation_config={"level":3,"splitting_ops":[],"use_inductor":true,"compile_sizes":[],"use_cudagraph":false,"cudagraph_capture_sizes":[256,248,240,232,224,216,208,200,192,184,176,168,160,152,144,136,128,120,112,104,96,88,80,72,64,56,48,40,32,24,16,8,4,2,1],"max_capture_size":256}

As above, I added a logger and saw VllmBackend appear twice
@youkaichao

[simon-mo]

Good feature request! I would imagine draft model to be especially useful to be compiled.

[github-actions]

This issue has been automatically marked as stale because it has not had any activity within 90 days. It will be automatically closed if no further activity occurs within 30 days. Leave a comment if you feel this issue should remain open. Thank you!

[github-actions]

This issue has been automatically closed due to inactivity. Please feel free to reopen if you feel it is still relevant. Thank you!