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
Fix Arcee model weight loading: Add custom load_weights (#21725)
Signed-off-by: alyosha-swamy <raghav@arcee.ai>
Raghav Ravishankar committed
Aug 4, 2025 at 16:39 UTC
a5fff3bd49a5ea888cf0dbdfe7ecf140455fa8d4
2 files changed
+80
-6
tests/models/registry.py
+1
-2
@@ -139,8 +139,7 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
139
trust_remote_code=True),
140
"AquilaForCausalLM": _HfExamplesInfo("BAAI/AquilaChat2-7B",
141
trust_remote_code=True),
142
- "ArceeForCausalLM": _HfExamplesInfo("arcee-ai/AFM-4.5B-Base",
143
- is_available_online=False),
142
+ "ArceeForCausalLM": _HfExamplesInfo("arcee-ai/AFM-4.5B-Base"),
143
"ArcticForCausalLM": _HfExamplesInfo("Snowflake/snowflake-arctic-instruct",
144
trust_remote_code=True),
145
"BaiChuanForCausalLM": _HfExamplesInfo("baichuan-inc/Baichuan-7B",
vllm/model_executor/models/arcee.py
+79
-4
@@ -24,10 +24,12 @@ from vllm.model_executor.layers.linear import (ColumnParallelLinear,
24
from vllm.model_executor.layers.logits_processor import LogitsProcessor
25
from vllm.model_executor.layers.vocab_parallel_embedding import (
26
DEFAULT_VOCAB_PADDING_SIZE, ParallelLMHead, VocabParallelEmbedding)
27
+from vllm.model_executor.model_loader.weight_utils import (
28
+ default_weight_loader, maybe_remap_kv_scale_name)
29
from vllm.sequence import IntermediateTensors
30
31
from .interfaces import SupportsLoRA, SupportsPP
30
-from .utils import (AutoWeightsLoader, PPMissingLayer,
32
+from .utils import (AutoWeightsLoader, PPMissingLayer, is_pp_missing_parameter,
33
make_empty_intermediate_tensors_factory, make_layers)
34
35
@@ -260,6 +262,81 @@ class ArceeModel(nn.Module):
262
return hidden_states, aux_hidden_states
263
return hidden_states
264
265
+ def load_weights(self, weights: Iterable[tuple[str,
266
+ torch.Tensor]]) -> set[str]:
267
+ """Load weights, mapping q/k/v projections to fused qkv_proj."""
268
+ stacked_params_mapping = [
269
+ (".qkv_proj", ".q_proj", "q"),
270
+ (".qkv_proj", ".k_proj", "k"),
271
+ (".qkv_proj", ".v_proj", "v"),
272
+ ]
273
+
274
+ params_dict = dict(self.named_parameters())
275
+ loaded_params: set[str] = set()
276
+
277
+ for name, loaded_weight in weights:
278
+ if "rotary_emb.inv_freq" in name:
279
+ continue
280
+ if ("rotary_emb.cos_cached" in name
281
+ or "rotary_emb.sin_cached" in name):
282
+ continue
283
+
284
+ if (self.quant_config is not None and
285
+ (scale_name := self.quant_config.get_cache_scale(name))):
286
+ param = params_dict[scale_name]
287
+ weight_loader = getattr(param, "weight_loader",
288
+ default_weight_loader)
289
+ loaded_weight = (loaded_weight if loaded_weight.dim() == 0 else
290
+ loaded_weight[0])
291
+ weight_loader(param, loaded_weight)
292
+ loaded_params.add(scale_name)
293
+ continue
294
+
295
+ if "scale" in name:
296
+ remapped_name = maybe_remap_kv_scale_name(name, params_dict)
297
+ if remapped_name is None:
298
+ continue
299
+ name = remapped_name
300
+
301
+ mapped = False
302
+ for param_name, weight_name, shard_id in stacked_params_mapping:
303
+ if weight_name not in name:
304
+ continue
305
+
306
+ name = name.replace(weight_name, param_name)
307
+
308
+ if name.endswith(".bias") and name not in params_dict:
309
+ mapped = True
310
+ break
311
+
312
+ if is_pp_missing_parameter(name, self):
313
+ mapped = True
314
+ break
315
+
316
+ param = params_dict[name]
317
+ weight_loader = param.weight_loader # type: ignore[attr-defined]
318
+ weight_loader(param, loaded_weight, shard_id)
319
+ loaded_params.add(name)
320
+ mapped = True
321
+ break
322
+
323
+ if mapped:
324
+ continue
325
+
326
+ if name.endswith(".bias") and name not in params_dict:
327
+ continue
328
+
329
+ if is_pp_missing_parameter(name, self):
330
+ continue
331
+
332
+ param = params_dict[name]
333
+ weight_loader = getattr(param, "weight_loader",
334
+ default_weight_loader)
335
+ weight_loader(param, loaded_weight)
336
+ loaded_params.add(name)
337
+
338
+ return loaded_params
339
+
340
341
class ArceeForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
342
"""Arcee Model for causal language modeling, integrated with vLLM
@@ -304,8 +381,7 @@ class ArceeForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
381
else:
382
# Placeholder for lm_head on non-last ranks
383
self.lm_head = PPMissingLayer()
307
- # Provide a reference to the model's method for generating empty
308
- # tensors (used in pipeline parallel schedule)
384
+
385
self.make_empty_intermediate_tensors = (
386
self.model.make_empty_intermediate_tensors)
387
@@ -316,7 +392,6 @@ class ArceeForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
392
intermediate_tensors: Optional[IntermediateTensors] = None,
393
inputs_embeds: Optional[torch.Tensor] = None
394
) -> Union[torch.Tensor, IntermediateTensors]:
319
- # Forward pass through the Arcee model backbone
395
model_output = self.model(input_ids=input_ids,
396
positions=positions,
397
intermediate_tensors=intermediate_tensors,