main
py 61 lines 2.46 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import regex
5
6 from vllm.model_executor.models.qwen3_vl import Qwen3VLForConditionalGeneration
7 from vllm.model_executor.models.utils import WeightsMapper
8
9
10 class Cosmos3ForConditionalGeneration(Qwen3VLForConditionalGeneration):
11 # Cosmos3 unified checkpoints store a Qwen3-VL understanding tower
12 # alongside a generation tower in a flat key layout. This mapper drops
13 # the generation tower weights and rewrites the understanding tower keys
14 # into the nested form expected by Qwen3VLForConditionalGeneration.
15 hf_to_vllm_mapper = WeightsMapper(
16 orig_to_new_regex={
17 regex.compile(
18 r"^(layers\.|embed_tokens\.|norm\.)(.+)$"
19 ): r"language_model.model.\1\2",
20 regex.compile(
21 r"^(blocks\.|merger\.|patch_embed\.|pos_embed\.|deepstack_merger_list\.)"
22 ): r"visual.\1",
23 regex.compile(r"^audio_modality_embed(?:\..*)?$"): None,
24 regex.compile(r"^action_modality_embed(?:\..*)?$"): None,
25 },
26 orig_to_new_substr={
27 "_moe_gen": None,
28 ".add_q_proj.": None,
29 ".add_k_proj.": None,
30 ".add_v_proj.": None,
31 ".to_add_out.": None,
32 ".norm_added_q.": None,
33 ".norm_added_k.": None,
34 ".to_q.": ".q_proj.",
35 ".to_k.": ".k_proj.",
36 ".to_v.": ".v_proj.",
37 ".to_out.": ".o_proj.",
38 ".norm_q.": ".q_norm.",
39 ".norm_k.": ".k_norm.",
40 # ModelOpt-native dialect (diffusers/transformers read these; vLLM reads
41 # weight_scale/input_scale instead), drop so AutoWeightsLoader passes
42 ".input_quantizer.": None,
43 ".weight_quantizer.": None,
44 ".output_quantizer.": None,
45 },
46 orig_to_new_prefix={
47 "proj_in.": None,
48 "proj_out.": None,
49 "time_embedder.": None,
50 "audio_proj_in.": None,
51 "audio_proj_out.": None,
52 "action_proj_in.": None,
53 "action_proj_out.": None,
54 "lm_head.": "language_model.lm_head.",
55 },
56 )
57
58 # Cosmos3 unified Diffusers checkpoints store reasoner weights across
59 # transformer/ and vision_encoder/. Match both while excluding VAE and
60 # sound_tokenizer weights; the mapper drops generation-only tensors.
61 allow_patterns_overrides = ["[tv]*er/*.safetensors"]