@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 161 lines 6.01 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from transformers.configuration_utils import PreTrainedConfig
5
6
7 class KimiLinearConfig(PreTrainedConfig):
8 model_type = "kimi_linear"
9 keys_to_ignore_at_inference = ["past_key_values"]
10
11 def __init__(
12 self,
13 model_type="kimi_linear",
14 vocab_size=163840,
15 hidden_size=4096,
16 head_dim=None,
17 intermediate_size=11008,
18 num_hidden_layers=32,
19 num_attention_heads=32,
20 num_key_value_heads=None,
21 hidden_act="silu",
22 initializer_range=0.02,
23 rms_norm_eps=1e-6,
24 use_cache=True,
25 pad_token_id=0,
26 bos_token_id=1,
27 eos_token_id=2,
28 rope_parameters=None,
29 tie_word_embeddings=False,
30 moe_intermediate_size: int | None = None,
31 moe_renormalize: bool = True,
32 moe_router_activation_func: str = "sigmoid",
33 num_experts: int | None = None,
34 num_experts_per_token: int | None = None,
35 num_shared_experts: int = 0,
36 routed_scaling_factor: float = 1.0,
37 first_k_dense_replace: int = 0,
38 moe_layer_freq: int = 1,
39 use_grouped_topk: bool = True,
40 num_expert_group: int = 1,
41 topk_group: int = 1,
42 q_lora_rank: int | None = None,
43 kv_lora_rank: int | None = None,
44 qk_nope_head_dim: int | None = None,
45 qk_rope_head_dim: int | None = None,
46 v_head_dim: int | None = None,
47 mla_use_nope: bool | None = False,
48 mla_use_output_gate: bool | None = False,
49 num_nextn_predict_layers: int = 0,
50 linear_attn_config: dict | None = None,
51 attn_res_block_size: int | None = None,
52 latent_moe_use_norm: bool = False,
53 activation_situ_beta: float | None = None,
54 activation_situ_linear_beta: float | None = None,
55 max_position_embeddings: int = 4096,
56 routed_expert_hidden_size: int | None = None,
57 topk_method: str = "noaux_tc",
58 **kwargs,
59 ):
60 self.model_type = model_type
61 self.vocab_size = vocab_size
62 self.hidden_size = hidden_size
63 self.head_dim = (
64 head_dim if head_dim is not None else hidden_size // num_attention_heads
65 )
66 self.intermediate_size = intermediate_size
67 self.num_hidden_layers = num_hidden_layers
68 self.num_attention_heads = num_attention_heads
69
70 # for backward compatibility
71 if num_key_value_heads is None:
72 num_key_value_heads = num_attention_heads
73
74 self.num_key_value_heads = num_key_value_heads
75 self.hidden_act = hidden_act
76 self.initializer_range = initializer_range
77 self.rms_norm_eps = rms_norm_eps
78 self.use_cache = use_cache
79 # Try to set `rope_scaling` if available, otherwise use `rope_parameters`
80 rope_scaling = kwargs.pop("rope_scaling", None)
81 rope_parameters = rope_scaling or rope_parameters or {"rope_type": "default"}
82 rope_theta = kwargs.pop("rope_theta", 10000.0)
83 if "rope_theta" not in rope_parameters:
84 rope_parameters["rope_theta"] = rope_theta
85 self.rope_parameters = rope_parameters
86
87 self.q_lora_rank = q_lora_rank
88 self.kv_lora_rank = kv_lora_rank
89 self.qk_nope_head_dim = qk_nope_head_dim
90 self.qk_rope_head_dim = qk_rope_head_dim
91 self.v_head_dim = v_head_dim
92 self.mla_use_nope = mla_use_nope
93 self.mla_use_output_gate = mla_use_output_gate
94 # moe config
95 self.num_experts = num_experts
96 self.num_experts_per_token = num_experts_per_token
97 self.moe_renormalize = moe_renormalize
98 self.num_shared_experts = num_shared_experts
99 self.routed_scaling_factor = routed_scaling_factor
100 self.moe_router_activation_func = moe_router_activation_func
101 assert self.moe_router_activation_func in ("softmax", "sigmoid")
102 self.moe_intermediate_size = moe_intermediate_size
103 self.first_k_dense_replace = first_k_dense_replace
104 self.moe_layer_freq = moe_layer_freq
105 self.use_grouped_topk = use_grouped_topk
106 self.num_expert_group = num_expert_group
107 self.topk_group = topk_group
108 self.num_nextn_predict_layers = num_nextn_predict_layers
109
110 self.attn_res_block_size = attn_res_block_size
111 self.latent_moe_use_norm = latent_moe_use_norm
112 self.activation_situ_beta = activation_situ_beta
113 self.activation_situ_linear_beta = activation_situ_linear_beta
114 self.max_position_embeddings = max_position_embeddings
115 self.routed_expert_hidden_size = routed_expert_hidden_size
116 self.topk_method = topk_method
117
118 if linear_attn_config is not None:
119 assert linear_attn_config["kda_layers"] is not None
120 assert linear_attn_config["full_attn_layers"] is not None
121 self.linear_attn_config = linear_attn_config
122
123 super().__init__(
124 pad_token_id=pad_token_id,
125 bos_token_id=bos_token_id,
126 eos_token_id=eos_token_id,
127 tie_word_embeddings=tie_word_embeddings,
128 **kwargs,
129 )
130
131 @property
132 def is_mla(self):
133 return (
134 self.q_lora_rank is not None
135 or self.kv_lora_rank is not None
136 or self.qk_nope_head_dim is not None
137 or self.qk_rope_head_dim is not None
138 or self.v_head_dim is not None
139 or self.mla_use_nope is True
140 )
141
142 @property
143 def is_moe(self):
144 return self.num_experts is not None
145
146 @property
147 def is_linear_attn(self) -> bool:
148 return not (
149 self.linear_attn_config is None
150 or (
151 isinstance(self.linear_attn_config, dict)
152 and self.linear_attn_config["kda_layers"] is not None
153 and len(self.linear_attn_config["kda_layers"]) == 0
154 )
155 )
156
157 def is_kda_layer(self, layer_idx: int):
158 return (
159 self.linear_attn_config is not None
160 and (layer_idx + 1) in self.linear_attn_config["kda_layers"]
161 )