@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 178 lines 5.12 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 from typing import Any
4
5 from transformers.configuration_utils import PreTrainedConfig
6
7
8 class Step3VisionEncoderConfig(PreTrainedConfig):
9 model_type = "step3_vision_encoder"
10
11 def __init__(
12 self,
13 hidden_size=1792,
14 intermediate_size=3072,
15 output_hidden_size=4096,
16 num_hidden_layers=63,
17 num_attention_heads=16,
18 num_channels=3,
19 image_size=728,
20 patch_size=14,
21 hidden_act="quick_gelu",
22 layer_norm_eps=1e-5,
23 **kwargs,
24 ):
25 self.hidden_size = hidden_size
26 self.intermediate_size = intermediate_size
27 self.output_hidden_size = output_hidden_size
28 self.num_hidden_layers = num_hidden_layers
29 self.num_attention_heads = num_attention_heads
30 self.num_channels = num_channels
31 self.patch_size = patch_size
32 self.image_size = image_size
33 self.layer_norm_eps = layer_norm_eps
34 self.hidden_act = hidden_act
35 super().__init__(**kwargs)
36
37
38 class Step3TextConfig(PreTrainedConfig):
39 model_type = "step3_text"
40 architectures = ["Step3TextForCausalLM"]
41
42 def __init__(
43 self,
44 hidden_size: int = 7168,
45 intermediate_size: int = 18432,
46 num_attention_heads: int = 64,
47 num_attention_groups: int = 1,
48 num_hidden_layers: int = 61,
49 max_seq_len: int = 65536,
50 vocab_size: int = 128815,
51 rms_norm_eps: float = 1e-5,
52 moe_intermediate_size: int = 5120,
53 moe_num_experts: int = 48,
54 moe_top_k: int = 3,
55 rope_parameters: dict[str, Any] | None = None,
56 max_position_embedding: int = 65536,
57 share_expert_dim: int = 5120,
58 share_q_dim: int = 2048,
59 head_dim: int = 256,
60 norm_expert_weight: bool = False,
61 moe_layers_enum: tuple[int, ...] = (
62 4,
63 5,
64 6,
65 7,
66 8,
67 9,
68 10,
69 11,
70 12,
71 13,
72 14,
73 15,
74 16,
75 17,
76 18,
77 19,
78 20,
79 21,
80 22,
81 23,
82 24,
83 25,
84 26,
85 27,
86 28,
87 29,
88 30,
89 31,
90 32,
91 33,
92 34,
93 35,
94 36,
95 37,
96 38,
97 39,
98 40,
99 41,
100 42,
101 43,
102 44,
103 45,
104 46,
105 47,
106 48,
107 49,
108 50,
109 51,
110 52,
111 53,
112 54,
113 55,
114 56,
115 57,
116 58,
117 59,
118 ),
119 **kwargs,
120 ) -> None:
121 self.hidden_size = hidden_size
122 self.intermediate_size = intermediate_size
123 self.num_attention_heads = num_attention_heads
124 self.num_attention_groups = num_attention_groups
125 self.num_hidden_layers = num_hidden_layers
126 self.max_seq_len = max_seq_len
127 self.vocab_size = vocab_size
128 self.rms_norm_eps = rms_norm_eps
129 self.moe_intermediate_size = moe_intermediate_size
130 self.moe_num_experts = moe_num_experts
131 self.moe_top_k = moe_top_k
132 # Try to set `rope_scaling` if available, otherwise use `rope_parameters`
133 rope_scaling = kwargs.pop("rope_scaling", None)
134 rope_parameters = rope_scaling or rope_parameters or {"rope_type": "default"}
135 rope_theta = kwargs.pop("rope_theta", 500000.0)
136 if "rope_theta" not in rope_parameters:
137 rope_parameters["rope_theta"] = rope_theta
138 self.rope_parameters = rope_parameters
139 self.max_position_embedding = max_position_embedding
140 self.share_expert_dim = share_expert_dim
141 self.share_q_dim = share_q_dim
142 self.head_dim = head_dim
143 self.norm_expert_weight = norm_expert_weight
144 self.moe_layers_enum = moe_layers_enum
145
146 super().__init__(**kwargs)
147
148
149 class Step3VLConfig(PreTrainedConfig):
150 model_type = "step3_vl"
151
152 def __init__(
153 self,
154 vision_config: dict | Step3VisionEncoderConfig | None = None,
155 text_config: dict | Step3TextConfig | None = None,
156 understand_projector_stride: int = 1,
157 projector_bias: bool = True,
158 image_token_id: int = 128001,
159 **kwargs,
160 ) -> None:
161 if vision_config is None:
162 vision_config = Step3VisionEncoderConfig()
163 elif isinstance(vision_config, dict):
164 vision_config = Step3VisionEncoderConfig(**vision_config)
165 self.vision_config = vision_config
166
167 if text_config is None:
168 text_config = Step3TextConfig()
169 elif isinstance(text_config, dict):
170 text_config = Step3TextConfig(**text_config)
171 self.text_config = text_config
172
173 self.understand_projector_stride = understand_projector_stride
174 self.projector_bias = projector_bias
175 self.hidden_size = text_config.hidden_size
176 self.image_token_id = image_token_id
177
178 super().__init__(**kwargs)