c6b636f9fb · tduncan/Nvidia-Optimization-Stack-VLLM

tests/v1/spec_decode/test_eagle.py

+6 -52

@@ -117,34 +117,13 @@ def test_prepare_inputs():

117 ])

118 @mock.patch('vllm.v1.spec_decode.eagle.get_pp_group')

119 @mock.patch('vllm.v1.spec_decode.eagle.get_layers_from_vllm_config')

120 -@mock.patch('vllm.v1.spec_decode.eagle.ModelRegistry')

121 -@mock.patch('vllm.v1.spec_decode.eagle.get_model_loader')

122 -@mock.patch('vllm.v1.spec_decode.eagle.set_default_torch_dtype')

123 -@mock.patch('vllm.v1.spec_decode.eagle.set_current_vllm_config')

124 -def test_load_model(mock_set_config, mock_set_dtype, mock_get_loader,

125 - mock_registry, mock_get_layers, mock_get_pp_group, method,

120 +@mock.patch('vllm.v1.spec_decode.eagle.get_model')

121 +def test_load_model(mock_get_model, mock_get_layers, mock_get_pp_group, method,

122 proposer_helper, draft_model_dir, target_attribute_path):

123

128 - # Setup mock for model class

129 - mock_model_cls = mock.MagicMock()

130 - mock_registry.resolve_model_cls.return_value = (mock_model_cls,

131 - "test_arch")

132 -

133 - # Create a real context manager for mocks

134 - class MockContextManager:

135 -

136 - def __init__(self):

137 - pass

138 -

139 - def __enter__(self):

140 - return None

141 -

142 - def __exit__(self, exc_type, exc_val, exc_tb):

143 - return False

144 -

145 - # Make the mocks return actual context manager objects

146 - mock_set_dtype.return_value = MockContextManager()

147 - mock_set_config.return_value = MockContextManager()

124 + # Setup model mock

125 + mock_model = mock.MagicMock()

126 + mock_get_model.return_value = mock_model

127

128 # Setup mocks for attention layers

129 target_attn_layers = {

@@ -164,25 +143,6 @@ def test_load_model(mock_set_config, mock_set_dtype, mock_get_loader,

143 mock_pp_group.world_size = 2 if method == "eagle" else 1

144 mock_get_pp_group.return_value = mock_pp_group

145

167 - # Setup model loader mock

168 - mock_loader = mock.MagicMock()

169 - mock_get_loader.return_value = mock_loader

170 -

171 - # Setup model mock

172 - mock_model = mock.MagicMock()

173 - mock_model_cls.return_value = mock_model

174 - mock_model.to.return_value = mock_model

175 -

176 - # Configure mock to test the attribute sharing path

177 - if method == "eagle":

178 - # For eagle, test the lm_head path

179 - mock_model.load_weights.return_value = {

180 - "model.embed_tokens.weight": torch.zeros(1)

181 - }

182 - else:

183 - # For eagle3, test the embed_tokens path

184 - mock_model.load_weights.return_value = {}

185 -

146 # Setup target model with the appropriate attributes

147 target_model = mock.MagicMock()

148

@@ -204,13 +164,7 @@ def test_load_model(mock_set_config, mock_set_dtype, mock_get_loader,

164 proposer.load_model(target_model)

165

166 # Verify common interactions

207 - mock_get_loader.assert_called_once()

208 - mock_model_cls.assert_called_once()

209 - mock_model.to.assert_called_once()

210 - mock_model.load_weights.assert_called_once()

211 -

212 - # Verify the loader was called with the right config

213 - mock_get_loader.assert_called_once_with(proposer.vllm_config.load_config)

167 + mock_get_model.assert_called_once()

168

169 # Verify the specific attribute sharing based on the method

170 if method == "eagle":

vllm/model_executor/model_loader/__init__.py

+10 -3

@@ -1,8 +1,10 @@

1 # SPDX-License-Identifier: Apache-2.0

2

3 +from typing import Optional

4 +

5 from torch import nn

6

5 -from vllm.config import LoadConfig, LoadFormat, VllmConfig

7 +from vllm.config import LoadConfig, LoadFormat, ModelConfig, VllmConfig

8 from vllm.model_executor.model_loader.base_loader import BaseModelLoader

9 from vllm.model_executor.model_loader.bitsandbytes_loader import (

10 BitsAndBytesModelLoader)

@@ -47,9 +49,14 @@ def get_model_loader(load_config: LoadConfig) -> BaseModelLoader:

49 return DefaultModelLoader(load_config)

50

51

50 -def get_model(*, vllm_config: VllmConfig) -> nn.Module:

52 +def get_model(*,

53 + vllm_config: VllmConfig,

54 + model_config: Optional[ModelConfig] = None) -> nn.Module:

55 loader = get_model_loader(vllm_config.load_config)

52 - return loader.load_model(vllm_config=vllm_config)

56 + if model_config is None:

57 + model_config = vllm_config.model_config

58 + return loader.load_model(vllm_config=vllm_config,

59 + model_config=model_config)

60

61

62 __all__ = [

vllm/model_executor/model_loader/base_loader.py

+2 -1

@@ -18,6 +18,7 @@ class BaseModelLoader(ABC):

18 raise NotImplementedError

19

20 @abstractmethod

21 - def load_model(self, *, vllm_config: VllmConfig) -> nn.Module:

21 + def load_model(self, *, vllm_config: VllmConfig,

22 + model_config: ModelConfig) -> nn.Module:

23 """Load a model with the given configurations."""

24 raise NotImplementedError

vllm/model_executor/model_loader/bitsandbytes_loader.py

+2 -3

@@ -569,10 +569,9 @@ class BitsAndBytesModelLoader(BaseModelLoader):

569 def download_model(self, model_config: ModelConfig) -> None:

570 self._prepare_weights(model_config.model, model_config.revision)

571

572 - def load_model(self, vllm_config: VllmConfig) -> nn.Module:

572 + def load_model(self, vllm_config: VllmConfig,

573 + model_config: ModelConfig) -> nn.Module:

574 device_config = vllm_config.device_config

574 - model_config = vllm_config.model_config

575 -

575 with set_default_torch_dtype(model_config.dtype):

576 with torch.device(device_config.device):

577

vllm/model_executor/model_loader/default_loader.py

+4 -3

@@ -264,13 +264,14 @@ class DefaultModelLoader(BaseModelLoader):

264 fall_back_to_pt=True,

265 allow_patterns_overrides=None)

266

267 - def load_model(self, vllm_config: VllmConfig) -> nn.Module:

267 + def load_model(self, vllm_config: VllmConfig,

268 + model_config: ModelConfig) -> nn.Module:

269 device_config = vllm_config.device_config

269 - model_config = vllm_config.model_config

270 target_device = torch.device(device_config.device)

271 with set_default_torch_dtype(model_config.dtype):

272 with target_device:

273 - model = initialize_model(vllm_config=vllm_config)

273 + model = initialize_model(vllm_config=vllm_config,

274 + model_config=model_config)

275

276 weights_to_load = {name for name, _ in model.named_parameters()}

277 loaded_weights = model.load_weights(

vllm/model_executor/model_loader/dummy_loader.py

+2 -2

@@ -22,9 +22,9 @@ class DummyModelLoader(BaseModelLoader):

22 def download_model(self, model_config: ModelConfig) -> None:

23 pass # Nothing to download

24

25 - def load_model(self, vllm_config: VllmConfig) -> nn.Module:

25 + def load_model(self, vllm_config: VllmConfig,

26 + model_config: ModelConfig) -> nn.Module:

27 device_config = vllm_config.device_config

27 - model_config = vllm_config.model_config

28 target_device = torch.device(device_config.device)

29 with set_default_torch_dtype(model_config.dtype):

30 with target_device:

vllm/model_executor/model_loader/gguf_loader.py

+2 -2

@@ -92,9 +92,9 @@ class GGUFModelLoader(BaseModelLoader):

92 def download_model(self, model_config: ModelConfig) -> None:

93 self._prepare_weights(model_config.model)

94

95 - def load_model(self, vllm_config: VllmConfig) -> nn.Module:

95 + def load_model(self, vllm_config: VllmConfig,

96 + model_config: ModelConfig) -> nn.Module:

97 device_config = vllm_config.device_config

97 - model_config = vllm_config.model_config

98 local_model_path = self._prepare_weights(model_config.model)

99 gguf_weights_map = self._get_gguf_weights_map(model_config)

100 # we can only know if tie word embeddings after mapping weights

vllm/model_executor/model_loader/runai_streamer_loader.py

+2 -3

@@ -100,11 +100,10 @@ class RunaiModelStreamerLoader(BaseModelLoader):

100 """Download model if necessary"""

101 self._prepare_weights(model_config.model, model_config.revision)

102

103 - def load_model(self, vllm_config: VllmConfig) -> nn.Module:

103 + def load_model(self, vllm_config: VllmConfig,

104 + model_config: ModelConfig) -> nn.Module:

105 """Perform streaming of the model to destination"""

106 device_config = vllm_config.device_config

106 - model_config = vllm_config.model_config

107 -

107 target_device = torch.device(device_config.device)

108 with set_default_torch_dtype(model_config.dtype):

109 with target_device:

vllm/model_executor/model_loader/sharded_state_loader.py

+2 -2

@@ -100,9 +100,9 @@ class ShardedStateLoader(BaseModelLoader):

100 def download_model(self, model_config: ModelConfig) -> None:

101 self._prepare_weights(model_config.model, model_config.revision)

102

103 - def load_model(self, vllm_config: VllmConfig) -> nn.Module:

103 + def load_model(self, vllm_config: VllmConfig,

104 + model_config: ModelConfig) -> nn.Module:

105 device_config = vllm_config.device_config

105 - model_config = vllm_config.model_config

106 target_device = torch.device(device_config.device)

107

108 from vllm.distributed import get_tensor_model_parallel_rank

vllm/model_executor/model_loader/tensorizer_loader.py

+2 -2

@@ -93,8 +93,8 @@ class TensorizerLoader(BaseModelLoader):

93 with self.tensorizer_config.open_stream():

94 pass

95

96 - def load_model(self, vllm_config: VllmConfig) -> nn.Module:

97 - model_config = vllm_config.model_config

96 + def load_model(self, vllm_config: VllmConfig,

97 + model_config: ModelConfig) -> nn.Module:

98 parallel_config = vllm_config.parallel_config

99 self._verify_config(model_config, parallel_config)

100

vllm/model_executor/model_loader/utils.py

+3 -1

@@ -42,9 +42,11 @@ def initialize_model(

42 *,

43 prefix: str = "",

44 model_class: Optional[type[nn.Module]] = None,

45 + model_config: Optional[ModelConfig] = None,

46 ) -> nn.Module:

47 """Initialize a model with the given configurations."""

47 - model_config = vllm_config.model_config

48 + if model_config is None:

49 + model_config = vllm_config.model_config

50 if model_class is None:

51 model_class, _ = get_model_architecture(model_config)

52

vllm/model_executor/models/llama_eagle.py

+4 -2

@@ -130,13 +130,15 @@ class LlamaModel(nn.Module):

130

131 class EagleLlamaForCausalLM(LlamaForCausalLM):

132

133 - def __init__(self, *, vllm_config: VllmConfig, start_layer_id: int = 0):

133 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

134 nn.Module.__init__(self)

135 self.config = vllm_config. \

136 speculative_config.draft_model_config.hf_config

137 + target_layer_num = vllm_config.model_config.get_num_layers(

138 + vllm_config.parallel_config)

139 self.model = LlamaModel(vllm_config=vllm_config,

140 prefix="model",

139 - start_layer_id=start_layer_id)

141 + start_layer_id=target_layer_num)

142

143 logit_scale = getattr(self.config, "logit_scale", 1.0)

144 self.logits_processor = LogitsProcessor(self.config.vocab_size,

vllm/model_executor/models/llama_eagle3.py

+6 -5

@@ -175,13 +175,15 @@ class LlamaModel(nn.Module):

175

176 class Eagle3LlamaForCausalLM(LlamaForCausalLM):

177

178 - def __init__(self, *, vllm_config: VllmConfig, start_layer_id: int = 0):

178 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

179 nn.Module.__init__(self)

180 self.config = vllm_config. \

181 speculative_config.draft_model_config.hf_config

182 + target_layer_num = vllm_config.model_config.get_num_layers(

183 + vllm_config.parallel_config)

184 self.model = LlamaModel(vllm_config=vllm_config,

183 - start_layer_id=start_layer_id,

184 - prefix="model")

185 + prefix="model",

186 + start_layer_id=target_layer_num)

187

188 logit_scale = getattr(self.config, "logit_scale", 1.0)

189 self.lm_head = ParallelLMHead(

@@ -193,8 +195,7 @@ class Eagle3LlamaForCausalLM(LlamaForCausalLM):

195 self.logits_processor = LogitsProcessor(self.config.draft_vocab_size,

196 scale=logit_scale)

197 self.draft_id_to_target_id = nn.Parameter(

196 - torch.zeros((self.config.draft_vocab_size),

197 - dtype=torch.long).type(torch.LongTensor),

198 + torch.zeros(self.config.draft_vocab_size, dtype=torch.long),

199 requires_grad=False,

200 )

201

vllm/model_executor/models/medusa.py

+1 -4

@@ -51,10 +51,7 @@ class Medusa(nn.Module):

51 needs to have truncated_vocab_size (=k) as an attribute."""

52

53 def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:

54 - if hasattr(vllm_config, 'draft_model_config'):

55 - config = vllm_config.draft_model_config.hf_config

56 - else:

57 - config = vllm_config.model_config.hf_config

54 + config = vllm_config.speculative_config.draft_model_config.hf_config

55 super().__init__()

56 self.config = config

57 self.blocks = nn.ModuleList([

vllm/v1/spec_decode/eagle.py

+6 -32

@@ -4,14 +4,11 @@ import torch.nn as nn

4

5 from vllm.attention.layer import Attention

6 from vllm.config import (CompilationLevel, VllmConfig,

7 - get_layers_from_vllm_config, set_current_vllm_config)

7 + get_layers_from_vllm_config)

8 from vllm.distributed.parallel_state import get_pp_group

9 from vllm.forward_context import set_forward_context

10 from vllm.logger import init_logger

11 -from vllm.model_executor.model_loader import get_model_loader

12 -from vllm.model_executor.model_loader.utils import (

13 - process_weights_after_loading, set_default_torch_dtype)

14 -from vllm.model_executor.models import ModelRegistry

11 +from vllm.model_executor.model_loader import get_model

12 from vllm.model_executor.models.llama_eagle3 import Eagle3LlamaForCausalLM

13 from vllm.triton_utils import tl, triton

14 from vllm.v1.attention.backends.flash_attn import FlashAttentionMetadata

@@ -280,51 +277,28 @@ class EagleProposer:

277 return cu_num_tokens, token_indices

278

279 def load_model(self, target_model: nn.Module) -> None:

283 - loader = get_model_loader(self.vllm_config.load_config)

284 - target_layer_num = self.vllm_config.model_config.get_num_layers(

285 - self.vllm_config.parallel_config)

280 + draft_model_config = \

281 + self.vllm_config.speculative_config.draft_model_config

282 target_attn_layer_names = set(

283 get_layers_from_vllm_config(self.vllm_config, Attention).keys())

284

289 - draft_model_config = \

290 - self.vllm_config.speculative_config.draft_model_config

291 - # FIXME(lily): This does not handle with distributed inference.

292 - target_device = self.vllm_config.device_config.device

293 - # We need to set the vllm_config here to register attention

294 - # layers in the forward context.

295 - with set_default_torch_dtype(

296 - draft_model_config.dtype), set_current_vllm_config(

297 - self.vllm_config):

298 - draft_model_cls, arch = ModelRegistry.resolve_model_cls(

299 - draft_model_config.architectures)

300 - self.model = draft_model_cls(

301 - vllm_config=self.vllm_config,

302 - start_layer_id=target_layer_num).to(target_device)

285 + self.model = get_model(vllm_config=self.vllm_config,

286 + model_config=draft_model_config)

287

288 draft_attn_layer_names = (

289 get_layers_from_vllm_config(self.vllm_config, Attention).keys() -

290 target_attn_layer_names)

291 assert len(draft_attn_layer_names) == 1

292 self.attn_layer_name = next(iter(draft_attn_layer_names))

309 - loaded_weights = self.model.load_weights(

310 - loader.get_all_weights(draft_model_config, self.model))

311 -

312 - process_weights_after_loading(self.model, draft_model_config,

313 - target_device)

293

294 # share embed_tokens with the target model if needed

295 if get_pp_group().world_size == 1:

317 - assert "model.embed_tokens.weight" not in loaded_weights, \

318 - "For PP = 1, Eagle draft should share embed with target model"

296 logger.info(

297 "The EAGLE head shares the same vocab embedding" \

298 " with the target model."

299 )

300 self.model.model.embed_tokens = target_model.model.embed_tokens

301 else:

325 - assert "model.embed_tokens.weight" in loaded_weights, \

326 - "For PP > 1, Eagle draft checkpoint should its own copy of "

327 - " the model.embed_tokens.weight"

302 logger.info(

303 "Since PP > 1, the EAGLE head loaded its own vocab embedding" \

304 " weights instead of sharing them with the target model."

vllm/v1/spec_decode/medusa.py

+5 -18

@@ -3,12 +3,10 @@

3 import torch

4 import torch.nn as nn

5

6 -from vllm.config import VllmConfig, set_current_vllm_config

6 +from vllm.config import VllmConfig

7 from vllm.forward_context import set_forward_context

8 from vllm.logger import init_logger

9 -from vllm.model_executor.model_loader import get_model_loader

10 -from vllm.model_executor.model_loader.utils import set_default_torch_dtype

11 -from vllm.model_executor.models.medusa import Medusa

9 +from vllm.model_executor.model_loader import get_model

10 from vllm.v1.sample.metadata import SamplingMetadata

11

12 # Initialize logger

@@ -49,20 +47,9 @@ class MedusaProposer:

47 return [list(row) for row in zip(*draft_tokens)]

48

49 def load_model(self, target_model: nn.Module) -> None:

52 - # Get model loader and config

53 - loader = get_model_loader(self.vllm_config.load_config)

54 - draft_config = self.vllm_config.speculative_config.draft_model_config

55 -

56 - # Load model with proper dtype and config

57 - with set_default_torch_dtype(draft_config.dtype), \

58 - set_current_vllm_config(self.vllm_config):

59 - self.model = Medusa(

60 - vllm_config=self.vllm_config.speculative_config).to(

61 - self.device)

62 -

63 - # Load model weights

64 - weights = loader.get_all_weights(draft_config, self.model)

65 - self.model.load_weights(weights)

50 + self.model = get_model(vllm_config=self.vllm_config,

51 + model_config=self.vllm_config.

52 + speculative_config.draft_model_config)

53

54 @torch.inference_mode()

55 def dummy_run(self, num_tokens: int) -> None:

[V1][Spec Decoding] Use model_loader.get_model() to load models (#18273)