awni commited on
Commit
a9e00be
1 Parent(s): a0166e9

292214023e4a9d271867930743d46930a457d312c750c50a47266210e375baad

Browse files
README.md ADDED
@@ -0,0 +1,38 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ license_link: https://huggingface.co/tencent/Tencent-Hunyuan-Large/blob/main/LICENSE.txt
3
+ language:
4
+ - en
5
+ pipeline_tag: text-generation
6
+ library_name: transformers
7
+ tags:
8
+ - mlx
9
+ base_model: tencent-community/Hunyuan-A52B-Instruct
10
+ ---
11
+
12
+ # mlx-community/Hunyuan-A52B-Instruct-3bit
13
+
14
+ The Model [mlx-community/Hunyuan-A52B-Instruct-3bit](https://huggingface.co/mlx-community/Hunyuan-A52B-Instruct-3bit) was
15
+ converted to MLX format from [tencent-community/Hunyuan-A52B-Instruct](https://huggingface.co/tencent-community/Hunyuan-A52B-Instruct)
16
+ using mlx-lm version **0.19.3**.
17
+
18
+ ## Use with mlx
19
+
20
+ ```bash
21
+ pip install mlx-lm
22
+ ```
23
+
24
+ ```python
25
+ from mlx_lm import load, generate
26
+
27
+ model, tokenizer = load("mlx-community/Hunyuan-A52B-Instruct-3bit")
28
+
29
+ prompt="hello"
30
+
31
+ if hasattr(tokenizer, "apply_chat_template") and tokenizer.chat_template is not None:
32
+ messages = [{"role": "user", "content": prompt}]
33
+ prompt = tokenizer.apply_chat_template(
34
+ messages, tokenize=False, add_generation_prompt=True
35
+ )
36
+
37
+ response = generate(model, tokenizer, prompt=prompt, verbose=True)
38
+ ```
config.json ADDED
@@ -0,0 +1,55 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "architectures": [
3
+ "HunYuanForCausalLM"
4
+ ],
5
+ "attention_bias": false,
6
+ "attention_dropout": 0.0,
7
+ "auto_map": {
8
+ "AutoConfig": "configuration_hunyuan.HunYuanConfig",
9
+ "AutoModel": "modeling_hunyuan.HunyuanModel",
10
+ "AutoModelForCausalLM": "modeling_hunyuan.HunYuanForCausalLM"
11
+ },
12
+ "bos_token_id": 1,
13
+ "capacity_factor": 1.0,
14
+ "cla_share_factor": 2,
15
+ "eos_token_id": 2,
16
+ "hidden_act": "silu",
17
+ "hidden_size": 6400,
18
+ "initializer_range": 0.02,
19
+ "intermediate_size": 18304,
20
+ "max_position_embeddings": 131072,
21
+ "model_type": "hunyuan",
22
+ "moe_drop_tokens": false,
23
+ "moe_random_routing_dropped_token": false,
24
+ "moe_topk": 1,
25
+ "num_attention_heads": 80,
26
+ "num_experts": 16,
27
+ "num_hidden_layers": 64,
28
+ "num_key_value_heads": 8,
29
+ "num_shared_expert": 1,
30
+ "pad_token_id": 0,
31
+ "pretraining_tp": 1,
32
+ "quantization": {
33
+ "group_size": 32,
34
+ "bits": 3
35
+ },
36
+ "quantization_config": {
37
+ "group_size": 32,
38
+ "bits": 3
39
+ },
40
+ "rms_norm_eps": 1e-05,
41
+ "rope_scaling": {
42
+ "alpha": 1000.0,
43
+ "factor": 1.0,
44
+ "type": "dynamic"
45
+ },
46
+ "rope_theta": 10000.0,
47
+ "tie_word_embeddings": true,
48
+ "torch_dtype": "bfloat16",
49
+ "transformers_version": "4.45.0.dev0",
50
+ "use_cache": true,
51
+ "use_cla": true,
52
+ "use_mixed_mlp_moe": true,
53
+ "use_qk_norm": true,
54
+ "vocab_size": 129024
55
+ }
configuration_hunyuan.py ADDED
@@ -0,0 +1,206 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (C) 2024 THL A29 Limited, a Tencent company. All rights reserved.
2
+ #
3
+ # Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # https://github.com/Tencent/Tencent-Hunyuan-Large/blob/main/License.docx
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+ """ HunYuan model configuration"""
15
+
16
+ from transformers.configuration_utils import PretrainedConfig
17
+ from transformers.utils import logging
18
+
19
+
20
+ logger = logging.get_logger(__name__)
21
+
22
+
23
+ class HunYuanConfig(PretrainedConfig):
24
+ r"""
25
+ This is the configuration class to store the configuration of a [`HunYuanModel`]. It is used to instantiate an
26
+ HunYuan model according to the specified arguments, defining the model architecture. Instantiating a configuration
27
+ with the defaults will yield a similar configuration to that of the HunYuan-7B.
28
+
29
+ Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
30
+ documentation from [`PretrainedConfig`] for more information.
31
+
32
+
33
+ Args:
34
+ vocab_size (`int`, *optional*, defaults to 32000):
35
+ Vocabulary size of the HunYuan model. Defines the number of different tokens that can be represented by the
36
+ `inputs_ids` passed when calling [`HunYuanModel`]
37
+ hidden_size (`int`, *optional*, defaults to 4096):
38
+ Dimension of the hidden representations.
39
+ intermediate_size (`int`, *optional*, defaults to 11008):
40
+ Dimension of the MLP representations.
41
+ num_hidden_layers (`int`, *optional*, defaults to 32):
42
+ Number of hidden layers in the Transformer decoder.
43
+ num_attention_heads (`int`, *optional*, defaults to 32):
44
+ Number of attention heads for each attention layer in the Transformer decoder.
45
+ num_key_value_heads (`int`, *optional*):
46
+ This is the number of key_value heads that should be used to implement Grouped Query Attention. If
47
+ `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
48
+ `num_key_value_heads=1 the model will use Multi Query Attention (MQA) otherwise GQA is used. When
49
+ converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed
50
+ by meanpooling all the original heads within that group. For more details checkout [this
51
+ paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to
52
+ `num_attention_heads`.
53
+ hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
54
+ The non-linear activation function (function or string) in the decoder.
55
+ max_position_embeddings (`int`, *optional*, defaults to 2048):
56
+ The maximum sequence length that this model might ever be used with.
57
+ initializer_range (`float`, *optional*, defaults to 0.02):
58
+ The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
59
+ rms_norm_eps (`float`, *optional*, defaults to 1e-06):
60
+ The epsilon used by the rms normalization layers.
61
+ use_cache (`bool`, *optional*, defaults to `True`):
62
+ Whether or not the model should return the last key/values attentions (not used by all models). Only
63
+ relevant if `config.is_decoder=True`.
64
+ pad_token_id (`int`, *optional*):
65
+ Padding token id.
66
+ bos_token_id (`int`, *optional*, defaults to 1):
67
+ Beginning of stream token id.
68
+ eos_token_id (`int`, *optional*, defaults to 2):
69
+ End of stream token id.
70
+ pretraining_tp (`int`, *optional*, defaults to 1):
71
+ Experimental feature. Tensor parallelism rank used during pretraining. Please refer to [this
72
+ document](https://huggingface.co/docs/transformers/parallelism) to understand more about it. This value is
73
+ necessary to ensure exact reproducibility of the pretraining results. Please refer to [this
74
+ issue](https://github.com/pytorch/pytorch/issues/76232).
75
+ tie_word_embeddings (`bool`, *optional*, defaults to `False`):
76
+ Whether to tie weight embeddings
77
+ rope_theta (`float`, *optional*, defaults to 10000.0):
78
+ The base period of the RoPE embeddings.
79
+ rope_scaling (`Dict`, *optional*):
80
+ Dictionary containing the scaling configuration for the RoPE embeddings. Currently supports two scaling
81
+ strategies: linear and dynamic. Their scaling factor must be a float greater than 1. The expected format is
82
+ `{"type": strategy name, "factor": scaling factor}`. When using this flag, don't update
83
+ `max_position_embeddings` to the expected new maximum. See the following thread for more information on how
84
+ these scaling strategies behave:
85
+ https://www.reddit.com/r/LocalLLaMA/comments/14mrgpr/dynamically_scaled_rope_further_increases/. This is an
86
+ experimental feature, subject to breaking API changes in future versions.
87
+ attention_bias (`bool`, defaults to `False`, *optional*, defaults to `False`):
88
+ Whether to use a bias in the query, key, value and output projection layers during self-attention.
89
+ attention_dropout (`float`, *optional*, defaults to 0.0):
90
+ The dropout ratio for the attention probabilities.
91
+ use_qk_norm (`bool`, *optional*, defaults to `False`):
92
+ Whether query and key in attention use norm
93
+ use_cla (`bool`, *optional*, defaults to `False`):
94
+ Whether to use CLA in attention
95
+ cla_share_factor (`int`, *optional*, defaults to 1):
96
+ The share factor of CLA
97
+ """
98
+
99
+ model_type = "hunyuan"
100
+ keys_to_ignore_at_inference = ["past_key_values"]
101
+
102
+ def __init__(
103
+ self,
104
+ vocab_size=290943,
105
+ hidden_size=4096,
106
+ intermediate_size=11008,
107
+ num_hidden_layers=32,
108
+ num_attention_heads=32,
109
+ num_key_value_heads=None,
110
+ hidden_act="silu",
111
+ max_position_embeddings=2048,
112
+ initializer_range=0.02,
113
+ rms_norm_eps=1e-5,
114
+ use_cache=True,
115
+ pad_token_id=0,
116
+ bos_token_id=1,
117
+ eos_token_id=2,
118
+ pretraining_tp=1,
119
+ tie_word_embeddings=False,
120
+ rope_theta=10000.0,
121
+ rope_scaling=None,
122
+ attention_bias=False,
123
+ attention_dropout=0.0,
124
+ use_qk_norm=False,
125
+ use_cla=False,
126
+ cla_share_factor=1,
127
+ num_experts=1,
128
+ use_mixed_mlp_moe=False,
129
+ num_shared_expert=1,
130
+ moe_topk=1,
131
+ capacity_factor=1.0,
132
+ moe_drop_tokens=False,
133
+ moe_random_routing_dropped_token=False,
134
+ **kwargs,
135
+ ):
136
+ self.vocab_size = vocab_size
137
+ self.max_position_embeddings = max_position_embeddings
138
+ self.hidden_size = hidden_size
139
+ self.intermediate_size = intermediate_size
140
+ self.num_hidden_layers = num_hidden_layers
141
+ self.num_attention_heads = num_attention_heads
142
+ self.num_experts = num_experts
143
+ self.use_mixed_mlp_moe = use_mixed_mlp_moe
144
+ self.num_shared_expert = num_shared_expert
145
+ self.moe_topk = moe_topk
146
+ self.capacity_factor = capacity_factor
147
+ self.moe_drop_tokens = moe_drop_tokens
148
+ self.moe_random_routing_dropped_token = moe_random_routing_dropped_token
149
+
150
+ # for backward compatibility
151
+ if num_key_value_heads is None:
152
+ num_key_value_heads = num_attention_heads
153
+
154
+ self.num_key_value_heads = num_key_value_heads
155
+ self.hidden_act = hidden_act
156
+ self.initializer_range = initializer_range
157
+ self.rms_norm_eps = rms_norm_eps
158
+ self.pretraining_tp = pretraining_tp
159
+ self.use_cache = use_cache
160
+ self.rope_theta = rope_theta
161
+ self.rope_scaling = rope_scaling
162
+ # self._rope_scaling_validation() # TODO: Need validation?
163
+ self.attention_bias = attention_bias
164
+ self.attention_dropout = attention_dropout
165
+ self.use_qk_norm = use_qk_norm
166
+ self.use_cla = use_cla
167
+ self.cla_share_factor = cla_share_factor
168
+
169
+ super().__init__(
170
+ pad_token_id=pad_token_id,
171
+ bos_token_id=bos_token_id,
172
+ eos_token_id=eos_token_id,
173
+ tie_word_embeddings=tie_word_embeddings,
174
+ **kwargs,
175
+ )
176
+
177
+ def _rope_scaling_validation(self):
178
+ """
179
+ Validate the `rope_scaling` configuration.
180
+ """
181
+ if self.rope_scaling is None:
182
+ return
183
+
184
+ if not isinstance(self.rope_scaling, dict) or len(self.rope_scaling) != 2:
185
+ raise ValueError(
186
+ "`rope_scaling` must be a dictionary with with two fields, `type` and `factor` or `type` and `alpha`, "
187
+ f"got {self.rope_scaling}"
188
+ )
189
+ rope_scaling_type = self.rope_scaling.get("type", None)
190
+ rope_scaling_factor = self.rope_scaling.get("factor", None)
191
+ rope_scaling_alpha = self.rope_scaling.get("alpha", None)
192
+ if rope_scaling_type is None or rope_scaling_type not in ["linear", "dynamic"]:
193
+ raise ValueError(
194
+ f"`rope_scaling`'s type field must be one of ['linear', 'dynamic'], got {rope_scaling_type}"
195
+ )
196
+ if rope_scaling_factor is None and rope_scaling_alpha is None:
197
+ raise ValueError(f"`rope_scaling`'s factor or alpha field must be have one, got both of none")
198
+ if rope_scaling_factor is not None:
199
+ if not isinstance(rope_scaling_factor, float) or rope_scaling_factor <= 1.0:
200
+ raise ValueError(f"`rope_scaling`'s factor field must be a float > 1.0, got {rope_scaling_factor}")
201
+ if rope_scaling_alpha is not None:
202
+ if not isinstance(rope_scaling_alpha, float) or rope_scaling_alpha <= 1.0:
203
+ raise ValueError(f"`rope_scaling`'s alpha field must be a float > 1.0, got {rope_scaling_alpha}")
204
+
205
+
206
+
hy.tiktoken ADDED
The diff for this file is too large to render. See raw diff
 
model-00039-of-00039.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:6d3927cc41cf8727a3d7d0feb4ed7f603dabd8a0688ea3b1cb4e10b3a5c1fd8b
3
+ size 1874369087
model.safetensors.index.json ADDED
The diff for this file is too large to render. See raw diff
 
modeling_hunyuan.py ADDED
@@ -0,0 +1,1726 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (C) 2024 THL A29 Limited, a Tencent company. All rights reserved.
2
+ #
3
+ # Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # https://github.com/Tencent/Tencent-Hunyuan-Large/blob/main/License.docx
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+ #
15
+ """ PyTorch HunYuan model."""
16
+
17
+ import math
18
+ import warnings
19
+ from typing import List, Optional, Tuple, Union
20
+
21
+ import torch
22
+ from torch import Tensor
23
+ import torch.nn.functional as F
24
+ import torch.utils.checkpoint
25
+ from torch import nn
26
+ from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
27
+
28
+ from transformers.activations import ACT2FN
29
+ from transformers.cache_utils import Cache, DynamicCache
30
+ from transformers.modeling_attn_mask_utils import (
31
+ AttentionMaskConverter,
32
+ _prepare_4d_attention_mask,
33
+ _prepare_4d_causal_attention_mask,
34
+ _prepare_4d_causal_attention_mask_for_sdpa,
35
+ )
36
+ from transformers.modeling_outputs import (
37
+ BaseModelOutputWithPast,
38
+ CausalLMOutputWithPast,
39
+ SequenceClassifierOutputWithPast
40
+ )
41
+ from transformers.modeling_utils import PreTrainedModel
42
+ from transformers.pytorch_utils import ALL_LAYERNORM_LAYERS, is_torch_greater_or_equal_than_1_13
43
+ from transformers.utils import (
44
+ add_start_docstrings,
45
+ add_start_docstrings_to_model_forward,
46
+ is_flash_attn_2_available,
47
+ is_flash_attn_greater_or_equal_2_10,
48
+ logging,
49
+ replace_return_docstrings,
50
+ )
51
+ from transformers.utils.import_utils import is_torch_fx_available
52
+ from .configuration_hunyuan import HunYuanConfig
53
+
54
+
55
+ if is_flash_attn_2_available():
56
+ from flash_attn import flash_attn_func, flash_attn_varlen_func
57
+ from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input # noqa
58
+
59
+
60
+ # This makes `_prepare_4d_causal_attention_mask` a leaf function in the FX graph.
61
+ # It means that the function will not be traced through and simply appear as a node in the graph.
62
+ if is_torch_fx_available():
63
+ if not is_torch_greater_or_equal_than_1_13:
64
+ import torch.fx
65
+
66
+ _prepare_4d_causal_attention_mask = torch.fx.wrap(_prepare_4d_causal_attention_mask)
67
+
68
+
69
+ logger = logging.get_logger(__name__)
70
+
71
+ _CONFIG_FOR_DOC = "HunYuanConfig"
72
+
73
+
74
+ def topkgating(logits: Tensor, topk: int):
75
+ logits = logits.float()
76
+ gates = F.softmax(logits, dim=1)
77
+ expert_capacity = topk * gates.shape[0]
78
+ num_experts = int(gates.shape[1])
79
+ # Top-k router probability and corresponding expert indices for each token.
80
+ # Shape: [tokens_per_group, num_selected_experts].
81
+ expert_gate, expert_index = torch.topk(gates, topk)
82
+ expert_mask = F.one_hot(expert_index, num_experts)
83
+ # For a given token, determine if it was routed to a given expert.
84
+ # Shape: [tokens_per_group, num_experts]
85
+ expert_mask_aux = expert_mask.max(dim=-2)[0]
86
+ tokens_per_group_and_expert = torch.mean(expert_mask_aux.float(), dim=-2)
87
+ router_prob_per_group_and_expert = torch.mean(gates.float(), dim=-2)
88
+ l_aux = num_experts**2 * torch.mean(tokens_per_group_and_expert * router_prob_per_group_and_expert)
89
+
90
+ gates_s = torch.clamp(
91
+ torch.matmul(expert_mask.float(), gates.unsqueeze(-1)).sum(dim=1), min=torch.finfo(gates.dtype).eps
92
+ )
93
+ router_probs = gates / gates_s
94
+ # Make num_selected_experts the leading axis to ensure that top-1 choices
95
+ # have priority over top-2 choices, which have priority over top-3 choices,
96
+ # etc.
97
+ expert_index = torch.transpose(expert_index, 0, 1)
98
+ # Shape: [num_selected_experts * tokens_per_group]
99
+ expert_index = expert_index.reshape(-1)
100
+
101
+ # Create mask out of indices.
102
+ # Shape: [tokens_per_group * num_selected_experts, num_experts].
103
+ expert_mask = F.one_hot(expert_index, num_experts).to(torch.int32)
104
+ exp_counts = torch.sum(expert_mask, dim=0).detach()
105
+
106
+ # Experts have a fixed capacity that we cannot exceed. A token's priority
107
+ # within the expert's buffer is given by the masked, cumulative capacity of
108
+ # its target expert.
109
+ # Shape: [tokens_per_group * num_selected_experts, num_experts].
110
+ token_priority = torch.cumsum(expert_mask, dim=0) * expert_mask - 1
111
+ # Shape: [num_selected_experts, tokens_per_group, num_experts].
112
+ token_priority = token_priority.reshape((topk, -1, num_experts))
113
+ # Shape: [tokens_per_group, num_selected_experts, num_experts].
114
+ token_priority = torch.transpose(token_priority, 0, 1)
115
+ # For each token, across all selected experts, select the only non-negative
116
+ # (unmasked) priority. Now, for group G routing to expert E, token T has
117
+ # non-negative priority (i.e. token_priority[G,T,E] >= 0) if and only if E
118
+ # is its targeted expert.
119
+ # Shape: [tokens_per_group, num_experts].
120
+ token_priority = torch.max(token_priority, dim=1)[0]
121
+
122
+ # Token T can only be routed to expert E if its priority is positive and
123
+ # less than the expert capacity. One-hot matrix will ignore indices outside
124
+ # the range [0, expert_capacity).
125
+ # Shape: [tokens_per_group, num_experts, expert_capacity].
126
+ valid_mask = torch.logical_and(token_priority >= 0, token_priority < expert_capacity)
127
+ token_priority = torch.masked_fill(token_priority, ~valid_mask, 0)
128
+ dispatch_mask = F.one_hot(token_priority, expert_capacity).to(torch.bool)
129
+ valid_mask = valid_mask.unsqueeze(-1).expand(-1, -1, expert_capacity)
130
+ dispatch_mask = torch.masked_fill(dispatch_mask, ~valid_mask, 0)
131
+
132
+ # The combine array will be used for combining expert outputs, scaled by the
133
+ # router probabilities. Shape: [num_groups, tokens_per_group, num_experts,
134
+ # expert_capacity].
135
+ combine_weights = torch.einsum("...te,...tec->...tec", router_probs, dispatch_mask)
136
+ exp_counts_capacity = torch.sum(dispatch_mask)
137
+ exp_capacity_rate = exp_counts_capacity / (logits.shape[0]*topk)
138
+
139
+ return [l_aux, exp_capacity_rate], combine_weights, dispatch_mask, exp_counts
140
+
141
+
142
+ def top1gating(logits: Tensor, random_routing_dropped_token: bool = False):
143
+ """Implements Top1Gating on logits."""
144
+ # everything is in fp32 in this function
145
+ logits = logits.float()
146
+ gates = F.softmax(logits, dim=1)
147
+ capacity = gates.shape[0]
148
+
149
+ # Create a mask for 1st's expert per token
150
+ # noisy gating
151
+ indices1_s = torch.argmax(gates, dim=1)
152
+ num_experts = int(gates.shape[1])
153
+ mask1 = F.one_hot(indices1_s, num_classes=num_experts)
154
+
155
+ # gating decisions
156
+ # exp_counts = torch.sum(mask1, dim=0).detach().to('cpu')
157
+ exp_counts = torch.sum(mask1, dim=0).detach()
158
+
159
+ # Compute l_aux
160
+ me = torch.mean(gates, dim=0)
161
+ ce = torch.mean(mask1.float(), dim=0)
162
+ l_aux = torch.sum(me * ce) * num_experts
163
+ mask1_rand = mask1
164
+
165
+ top_idx = torch.topk(mask1_rand, k=capacity, dim=0)[1]
166
+
167
+ new_mask1 = mask1 * torch.zeros_like(mask1).scatter_(0, top_idx, 1)
168
+ mask1 = new_mask1
169
+ mask1_bk = mask1
170
+ if random_routing_dropped_token:
171
+ not_full = capacity - new_mask1.sum(dim=0)
172
+ sorted_notfull, indices_notfull = torch.sort(not_full, descending=True)
173
+ sorted_notfull = sorted_notfull.to(torch.int64)
174
+ not_full_experts_ids = torch.repeat_interleave(indices_notfull, sorted_notfull)
175
+ shuffle_not_full_ids = torch.randperm(not_full_experts_ids.shape[0])
176
+ not_full_experts_ids = not_full_experts_ids[shuffle_not_full_ids]
177
+ indices1_s_after_drop = torch.argmax(new_mask1, dim=1)
178
+ # get drop idx
179
+ drop_mask = 1 - new_mask1.sum(dim=1)
180
+ drop_mask = drop_mask.bool()
181
+ drop_idx = drop_mask.nonzero().view(-1)
182
+ drop_num = drop_mask.sum().to(torch.int64)
183
+ indices1_s_after_drop.scatter_(0, drop_idx, not_full_experts_ids[:drop_num])
184
+ nodrop_mask1 = F.one_hot(indices1_s_after_drop, num_classes=num_experts)
185
+ mask1 = nodrop_mask1
186
+
187
+ # Compute locations in capacity buffer
188
+ locations1 = torch.cumsum(mask1, dim=0) - 1
189
+
190
+ # Store the capacity location for each token
191
+ locations1_s = torch.sum(locations1 * mask1, dim=1)
192
+
193
+ # Normalize gate probabilities
194
+ mask1_float = mask1.float()
195
+ gates = gates * mask1_float
196
+
197
+ locations1_sc = F.one_hot(locations1_s, num_classes=capacity).float() # one hot to float
198
+ combine_weights = torch.einsum("se,sc->sec", gates, locations1_sc)
199
+
200
+ dispatch_mask = combine_weights.bool()
201
+
202
+ exp_counts_capacity = torch.sum(mask1_bk)
203
+ exp_capacity_rate = exp_counts_capacity / (logits.shape[0])
204
+ return [l_aux, exp_capacity_rate], combine_weights, dispatch_mask, exp_counts
205
+
206
+
207
+ def _get_unpad_data(attention_mask):
208
+ seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
209
+ indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
210
+ max_seqlen_in_batch = seqlens_in_batch.max().item()
211
+ cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.torch.int32), (1, 0))
212
+ return (
213
+ indices,
214
+ cu_seqlens,
215
+ max_seqlen_in_batch,
216
+ )
217
+
218
+
219
+ def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None):
220
+ warnings.warn(
221
+ "Calling `transformers.models.llama.modeling_llama._prepare_4d_attention_mask` is deprecated and will be "
222
+ "removed in v4.37. Use `transformers.modeling_attn_mask_utils._prepare_4d_attention_mask"
223
+ )
224
+ return _prepare_4d_attention_mask(mask=mask, dtype=dtype, tgt_len=tgt_len)
225
+
226
+
227
+ def _make_causal_mask(
228
+ input_ids_shape: torch.Size, dtype: torch.dtype, device: torch.device, past_key_values_length: int = 0
229
+ ):
230
+ warnings.warn(
231
+ "Calling `transformers.models.llama.modeling_llama._make_causal_mask` is deprecated and will be removed in "
232
+ "v4.37. Use `transformers.models.llama.modeling_llama.AttentionMaskConverter._make_causal_mask"
233
+ )
234
+ return AttentionMaskConverter._make_causal_mask(
235
+ input_ids_shape=input_ids_shape, dtype=dtype, device=device, past_key_values_length=past_key_values_length
236
+ )
237
+
238
+
239
+ class HunYuanRMSNorm(nn.Module):
240
+ def __init__(self, hidden_size, eps=1e-6):
241
+ """
242
+ HunYuanRMSNorm is equivalent to T5LayerNorm
243
+ """
244
+ super().__init__()
245
+ self.weight = nn.Parameter(torch.ones(hidden_size))
246
+ self.variance_epsilon = eps
247
+
248
+ def forward(self, hidden_states):
249
+ input_dtype = hidden_states.dtype
250
+ hidden_states = hidden_states.to(torch.float32)
251
+ variance = hidden_states.pow(2).mean(-1, keepdim=True)
252
+ hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
253
+ return self.weight * hidden_states.to(input_dtype)
254
+
255
+
256
+ ALL_LAYERNORM_LAYERS.append(HunYuanRMSNorm)
257
+
258
+
259
+ class HunYuanRotaryEmbedding(nn.Module):
260
+ def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):
261
+ super().__init__()
262
+
263
+ self.dim = dim
264
+ self.max_position_embeddings = max_position_embeddings
265
+ self.base = base
266
+ inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))
267
+ inv_freq = inv_freq.bfloat16()
268
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
269
+
270
+ # Build here to make `torch.jit.trace` work.
271
+ self._set_cos_sin_cache(
272
+ seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype()
273
+ )
274
+
275
+ def _set_cos_sin_cache(self, seq_len, device, dtype):
276
+ self.max_seq_len_cached = seq_len
277
+ t = torch.arange(self.max_seq_len_cached, device=device, dtype=torch.float32)
278
+
279
+ freqs = torch.outer(t, self.inv_freq)
280
+ # Different from paper, but it uses a different permutation in order to obtain the same calculation
281
+ emb = torch.cat((freqs, freqs), dim=-1).float()
282
+ self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
283
+ self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
284
+
285
+ def forward(self, x, seq_len=None):
286
+ # x: [bs, num_attention_heads, seq_len, head_size]
287
+ if seq_len > self.max_seq_len_cached:
288
+ self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype)
289
+
290
+ return (
291
+ self.cos_cached[:seq_len].to(dtype=x.dtype),
292
+ self.sin_cached[:seq_len].to(dtype=x.dtype),
293
+ )
294
+
295
+
296
+ class HunYuanLinearScalingRotaryEmbedding(HunYuanRotaryEmbedding):
297
+ """HunYuanRotaryEmbedding extended with linear scaling. Credits to the Reddit user /u/kaiokendev"""
298
+
299
+ def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0):
300
+ self.scaling_factor = scaling_factor
301
+ super().__init__(dim, max_position_embeddings, base, device)
302
+
303
+ def _set_cos_sin_cache(self, seq_len, device, dtype):
304
+ self.max_seq_len_cached = seq_len
305
+ t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)
306
+ t = t / self.scaling_factor
307
+
308
+ freqs = torch.outer(t, self.inv_freq)
309
+ # Different from paper, but it uses a different permutation in order to obtain the same calculation
310
+ emb = torch.cat((freqs, freqs), dim=-1)
311
+ self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
312
+ self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
313
+
314
+
315
+ class HunYuanDynamicNTKScalingRotaryEmbedding(HunYuanRotaryEmbedding):
316
+ """
317
+ HunYuanRotaryEmbedding extended with Dynamic NTK scaling.
318
+ Credits to the Reddit users /u/bloc97 and /u/emozilla
319
+ """
320
+
321
+ def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0):
322
+ self.scaling_factor = scaling_factor
323
+ super().__init__(dim, max_position_embeddings, base, device)
324
+
325
+ def _set_cos_sin_cache(self, seq_len, device, dtype):
326
+ self.max_seq_len_cached = seq_len
327
+
328
+ if seq_len > self.max_position_embeddings:
329
+ base = self.base * (
330
+ (self.scaling_factor * seq_len / self.max_position_embeddings) - (self.scaling_factor - 1)
331
+ ) ** (self.dim / (self.dim - 2))
332
+ inv_freq = 1.0 / (base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))
333
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
334
+
335
+ t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)
336
+
337
+ freqs = torch.outer(t, self.inv_freq)
338
+ # Different from paper, but it uses a different permutation in order to obtain the same calculation
339
+ emb = torch.cat((freqs, freqs), dim=-1)
340
+ self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
341
+ self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
342
+
343
+
344
+ class HunYuanDynamicNTKAlphaRotaryEmbedding(HunYuanRotaryEmbedding):
345
+ """
346
+ HunYuanRotaryEmbedding extended with Dynamic NTK scaling.
347
+ Credits to the Reddit users /u/bloc97 and /u/emozilla
348
+ """
349
+
350
+ def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_alpha=1.0):
351
+ self.scaling_alpha = scaling_alpha
352
+ super().__init__(dim, max_position_embeddings, base, device)
353
+
354
+ def _set_cos_sin_cache(self, seq_len, device, dtype):
355
+ self.max_seq_len_cached = seq_len
356
+ base = self.base * self.scaling_alpha ** (self.dim / (self.dim-2))
357
+ inv_freq = 1.0 / (base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))
358
+
359
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
360
+
361
+ t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)
362
+
363
+ freqs = torch.outer(t, self.inv_freq)
364
+ # Different from paper, but it uses a different permutation in order to obtain the same calculation
365
+ emb = torch.cat((freqs, freqs), dim=-1)
366
+ self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
367
+ self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
368
+
369
+
370
+ def rotate_half(x):
371
+ """Rotates half the hidden dims of the input."""
372
+ x1 = x[..., : x.shape[-1] // 2]
373
+ x2 = x[..., x.shape[-1] // 2:]
374
+ return torch.cat((-x2, x1), dim=-1)
375
+
376
+
377
+ def apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1):
378
+ """Applies Rotary Position Embedding to the query and key tensors.
379
+
380
+ Args:
381
+ q (`torch.Tensor`): The query tensor.
382
+ k (`torch.Tensor`): The key tensor.
383
+ cos (`torch.Tensor`): The cosine part of the rotary embedding.
384
+ sin (`torch.Tensor`): The sine part of the rotary embedding.
385
+ position_ids (`torch.Tensor`):
386
+ The position indices of the tokens corresponding to the query and key tensors. For example, this can be
387
+ used to pass offsetted position ids when working with a KV-cache.
388
+ unsqueeze_dim (`int`, *optional*, defaults to 1):
389
+ The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
390
+ sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
391
+ that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
392
+ k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
393
+ cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
394
+ the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
395
+ Returns:
396
+ `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
397
+ """
398
+ cos = cos[position_ids].unsqueeze(unsqueeze_dim)
399
+ sin = sin[position_ids].unsqueeze(unsqueeze_dim)
400
+ q_embed = (q * cos) + (rotate_half(q) * sin)
401
+ k_embed = (k * cos) + (rotate_half(k) * sin)
402
+ return q_embed, k_embed
403
+
404
+
405
+ class HunYuanMLP(nn.Module):
406
+ def __init__(self, config: HunYuanConfig, layer_idx=None, is_shared_mlp=False):
407
+ super().__init__()
408
+ self.config = config
409
+ self.layer_idx = layer_idx
410
+ self.hidden_size = config.hidden_size
411
+ if is_shared_mlp:
412
+ self.intermediate_size = config.intermediate_size * config.num_shared_expert
413
+ else:
414
+ self.intermediate_size = config.intermediate_size
415
+ self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
416
+ self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
417
+ self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
418
+ self.act_fn = ACT2FN[config.hidden_act]
419
+
420
+ def forward(self, x):
421
+ if self.config.pretraining_tp > 1:
422
+ slice = self.intermediate_size // self.config.pretraining_tp
423
+ gate_proj_slices = self.gate_proj.weight.split(slice, dim=0)
424
+ up_proj_slices = self.up_proj.weight.split(slice, dim=0)
425
+ down_proj_slices = self.down_proj.weight.split(slice, dim=1)
426
+
427
+ gate_proj = torch.cat(
428
+ [F.linear(x, gate_proj_slices[i]) for i in range(self.config.pretraining_tp)], dim=-1
429
+ )
430
+ up_proj = torch.cat([F.linear(x, up_proj_slices[i]) for i in range(self.config.pretraining_tp)], dim=-1)
431
+
432
+ intermediate_states = (self.act_fn(gate_proj) * up_proj).split(slice, dim=2)
433
+ down_proj = [
434
+ F.linear(intermediate_states[i], down_proj_slices[i]) for i in range(self.config.pretraining_tp)
435
+ ]
436
+ down_proj = sum(down_proj)
437
+ else:
438
+ down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
439
+
440
+ return down_proj
441
+
442
+
443
+ class HunYuanTopKGate(nn.Module):
444
+ def __init__(self, config: HunYuanConfig, layer_idx: Optional[int] = None):
445
+ super().__init__()
446
+ self.config = config
447
+ self.layer_idx = layer_idx
448
+ self.moe_topk = config.moe_topk
449
+ self.drop_tokens = config.moe_drop_tokens
450
+ self.min_capacity = 8
451
+ self.random_routing_dropped_token = config.moe_random_routing_dropped_token
452
+ self.wg = nn.Linear(config.hidden_size, config.num_experts, bias=False, dtype=torch.float32)
453
+
454
+ def forward(self, hidden_states):
455
+ bsz, seq_len, hidden_size = hidden_states.shape
456
+ hidden_states = hidden_states.reshape(-1, hidden_size)
457
+ if self.wg.weight.dtype == torch.float32:
458
+ hidden_states = hidden_states.float()
459
+ logits = self.wg(hidden_states)
460
+ if self.moe_topk == 1:
461
+ gate_output = top1gating(logits, random_routing_dropped_token=self.random_routing_dropped_token)
462
+ else:
463
+ gate_output = topkgating(logits, self.moe_topk)
464
+
465
+ return gate_output
466
+
467
+
468
+ class HunYuanMoE(nn.Module):
469
+ def __init__(self, config: HunYuanConfig, layer_idx: Optional[int] = None):
470
+ super().__init__()
471
+ self.config = config
472
+ self.layer_idx = layer_idx
473
+ self.moe_topk = config.moe_topk
474
+ self.num_experts = config.num_experts
475
+ if config.use_mixed_mlp_moe:
476
+ self.shared_mlp = HunYuanMLP(config, layer_idx=layer_idx, is_shared_mlp=True)
477
+ self.gate = HunYuanTopKGate(config, layer_idx=layer_idx)
478
+ self.experts = nn.ModuleList(
479
+ [HunYuanMLP(config, layer_idx=layer_idx, is_shared_mlp=False) for _ in range(config.num_experts)]
480
+ )
481
+
482
+ def forward(self, hidden_states):
483
+ bsz, seq_len, hidden_size = hidden_states.shape
484
+
485
+ if self.config.use_mixed_mlp_moe:
486
+ hidden_states_mlp = self.shared_mlp(hidden_states)
487
+
488
+ l_moe, combine_weights, dispatch_mask, exp_counts = self.gate(hidden_states)
489
+
490
+ reshaped_input = hidden_states.reshape(-1, hidden_size)
491
+
492
+ dispatched_input = torch.einsum("sec,sm->ecm", dispatch_mask.type_as(hidden_states), reshaped_input)
493
+
494
+ chunks = dispatched_input.chunk(self.num_experts, dim=0)
495
+ expert_outputs = []
496
+ for chunk, expert in zip(chunks, self.experts):
497
+ expert_outputs.append(expert(chunk))
498
+
499
+ expert_output = torch.cat(expert_outputs, dim=0)
500
+ combined_output = torch.einsum("sec,ecm->sm", combine_weights.type_as(hidden_states), expert_output)
501
+ combined_output = combined_output.reshape(bsz, seq_len, hidden_size)
502
+
503
+ if self.config.use_mixed_mlp_moe:
504
+ output = hidden_states_mlp + combined_output
505
+ else:
506
+ output = combined_output
507
+
508
+ return output
509
+
510
+
511
+ def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
512
+ """
513
+ This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
514
+ num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
515
+ """
516
+ batch, num_key_value_heads, slen, head_dim = hidden_states.shape
517
+ if n_rep == 1:
518
+ return hidden_states
519
+ hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
520
+ return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
521
+
522
+
523
+ class HunYuanAttention(nn.Module):
524
+ """Multi-headed attention from 'Attention Is All You Need' paper"""
525
+
526
+ def __init__(self, config: HunYuanConfig, layer_idx: Optional[int] = None):
527
+ super().__init__()
528
+ self.config = config
529
+ self.layer_idx = layer_idx
530
+ # layer_idx 从 0 开始
531
+ self.attention_type = 'cross' if config.use_cla and layer_idx % config.cla_share_factor != 0 else 'self'
532
+ if layer_idx is None:
533
+ logger.warning_once(
534
+ f"Instantiating {self.__class__.__name__} without passing `layer_idx` is not recommended and will "
535
+ "to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` "
536
+ "when creating this class."
537
+ )
538
+
539
+ self.attention_dropout = config.attention_dropout
540
+ self.hidden_size = config.hidden_size
541
+ self.num_heads = config.num_attention_heads
542
+ self.head_dim = self.hidden_size // self.num_heads
543
+ self.num_key_value_heads = config.num_key_value_heads
544
+ self.num_key_value_groups = self.num_heads // self.num_key_value_heads
545
+ self.max_position_embeddings = config.max_position_embeddings
546
+ self.rope_theta = config.rope_theta
547
+ self.is_causal = True
548
+ self.use_qk_norm = config.use_qk_norm
549
+
550
+ if (self.head_dim * self.num_heads) != self.hidden_size:
551
+ raise ValueError(
552
+ f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"
553
+ f" and `num_heads`: {self.num_heads})."
554
+ )
555
+
556
+ self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=config.attention_bias)
557
+ if self.attention_type == 'self':
558
+ self.k_proj = nn.Linear(
559
+ self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias
560
+ )
561
+ self.v_proj = nn.Linear(
562
+ self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias
563
+ )
564
+ self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=config.attention_bias)
565
+ if self.use_qk_norm:
566
+ self.query_layernorm = HunYuanRMSNorm(self.head_dim, eps=config.rms_norm_eps)
567
+ self.key_layernorm = HunYuanRMSNorm(self.head_dim, eps=config.rms_norm_eps)
568
+ self._init_rope()
569
+
570
+ def _init_rope(self):
571
+ if self.config.rope_scaling is None:
572
+ self.rotary_emb = HunYuanRotaryEmbedding(
573
+ self.head_dim,
574
+ max_position_embeddings=self.max_position_embeddings,
575
+ base=self.rope_theta,
576
+ )
577
+ else:
578
+ scaling_type = self.config.rope_scaling["type"]
579
+ scaling_factor = self.config.rope_scaling["factor"]
580
+ scaling_alpha = self.config.rope_scaling["alpha"]
581
+ if scaling_type == "linear":
582
+ self.rotary_emb = HunYuanLinearScalingRotaryEmbedding(
583
+ self.head_dim,
584
+ max_position_embeddings=self.max_position_embeddings,
585
+ scaling_factor=scaling_factor,
586
+ base=self.rope_theta,
587
+ )
588
+ elif scaling_type == "dynamic":
589
+ if scaling_alpha:
590
+ self.rotary_emb = HunYuanDynamicNTKAlphaRotaryEmbedding(
591
+ self.head_dim,
592
+ max_position_embeddings=self.max_position_embeddings,
593
+ scaling_alpha=scaling_alpha,
594
+ base=self.rope_theta,
595
+ )
596
+ else:
597
+ self.rotary_emb = HunYuanDynamicNTKScalingRotaryEmbedding(
598
+ self.head_dim,
599
+ max_position_embeddings=self.max_position_embeddings,
600
+ scaling_factor=scaling_factor,
601
+ base=self.rope_theta,
602
+ )
603
+ else:
604
+ raise ValueError(f"Unknown RoPE scaling type {scaling_type}")
605
+
606
+ def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
607
+ return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
608
+
609
+ def forward(
610
+ self,
611
+ hidden_states: torch.Tensor,
612
+ attention_mask: Optional[torch.Tensor] = None,
613
+ position_ids: Optional[torch.LongTensor] = None,
614
+ past_key_value: Optional[Cache] = None,
615
+ output_attentions: bool = False,
616
+ use_cache: bool = False,
617
+ kv_states: torch.Tensor = None,
618
+ **kwargs,
619
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
620
+ if "padding_mask" in kwargs:
621
+ warnings.warn(
622
+ "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use "
623
+ "`attention_mask` instead.`"
624
+ )
625
+
626
+ bsz, q_len, _ = hidden_states.size()
627
+
628
+ if self.config.pretraining_tp > 1:
629
+ query_slices = self.q_proj.weight.split(
630
+ (self.num_heads * self.head_dim) // self.config.pretraining_tp, dim=0
631
+ )
632
+ query_states = [F.linear(hidden_states, query_slices[i]) for i in range(self.config.pretraining_tp)]
633
+ query_states = torch.cat(query_states, dim=-1)
634
+
635
+ if self.attention_type == "cross" and kv_states is not None and isinstance(kv_states, tuple):
636
+ orig_key_states, orig_value_states = kv_states
637
+ key_states, value_states = kv_states
638
+ else:
639
+ key_value_slicing = (self.num_key_value_heads * self.head_dim) // self.config.pretraining_tp
640
+ key_slices = self.k_proj.weight.split(key_value_slicing, dim=0)
641
+ value_slices = self.v_proj.weight.split(key_value_slicing, dim=0)
642
+
643
+ key_states = [F.linear(hidden_states, key_slices[i]) for i in range(self.config.pretraining_tp)]
644
+ key_states = torch.cat(key_states, dim=-1)
645
+
646
+ value_states = [F.linear(hidden_states, value_slices[i]) for i in range(self.config.pretraining_tp)]
647
+ value_states = torch.cat(value_states, dim=-1)
648
+ orig_key_states, orig_value_states = key_states, value_states
649
+
650
+ else:
651
+ query_states = self.q_proj(hidden_states)
652
+ if self.attention_type == "cross" and kv_states is not None and isinstance(kv_states, tuple):
653
+ orig_key_states, orig_value_states = kv_states
654
+ key_states, value_states = kv_states
655
+ else:
656
+ key_states = self.k_proj(hidden_states)
657
+ value_states = self.v_proj(hidden_states)
658
+ orig_key_states, orig_value_states = key_states, value_states
659
+
660
+ query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
661
+ key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
662
+ value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
663
+
664
+ kv_seq_len = key_states.shape[-2]
665
+ if past_key_value is not None:
666
+ if self.layer_idx is None:
667
+ raise ValueError(
668
+ f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} "
669
+ "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class "
670
+ "with a layer index."
671
+ )
672
+ kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
673
+ cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
674
+ query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
675
+
676
+ if self.use_qk_norm:
677
+ query_states = self.query_layernorm(query_states)
678
+ key_states = self.key_layernorm(key_states)
679
+
680
+ if past_key_value is not None:
681
+ cache_kwargs = {"sin": sin, "cos": cos} # Specific to RoPE models
682
+ key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
683
+
684
+ key_states = repeat_kv(key_states, self.num_key_value_groups)
685
+ value_states = repeat_kv(value_states, self.num_key_value_groups)
686
+
687
+ attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
688
+
689
+ if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len):
690
+ raise ValueError(
691
+ f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is"
692
+ f" {attn_weights.size()}"
693
+ )
694
+
695
+ if attention_mask is not None:
696
+ if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
697
+ raise ValueError(
698
+ f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
699
+ )
700
+ attn_weights = attn_weights + attention_mask
701
+
702
+ # upcast attention to fp32
703
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
704
+ attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training)
705
+ attn_output = torch.matmul(attn_weights, value_states)
706
+
707
+ if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):
708
+ raise ValueError(
709
+ f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"
710
+ f" {attn_output.size()}"
711
+ )
712
+
713
+ attn_output = attn_output.transpose(1, 2).contiguous()
714
+
715
+ attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
716
+
717
+ if self.config.pretraining_tp > 1:
718
+ attn_output = attn_output.split(self.hidden_size // self.config.pretraining_tp, dim=2)
719
+ o_proj_slices = self.o_proj.weight.split(self.hidden_size // self.config.pretraining_tp, dim=1)
720
+ attn_output = sum([F.linear(attn_output[i], o_proj_slices[i]) for i in range(self.config.pretraining_tp)])
721
+ else:
722
+ attn_output = self.o_proj(attn_output)
723
+
724
+ if not output_attentions:
725
+ attn_weights = None
726
+
727
+ return attn_output, attn_weights, past_key_value, (orig_key_states, orig_value_states)
728
+
729
+
730
+ class HunYuanFlashAttention2(HunYuanAttention):
731
+ """
732
+ HunYuan flash attention module. This module inherits from `HunYuanAttention` as the weights of the module stays
733
+ untouched. The only required change would be on the forward pass where it needs to correctly call the public API of
734
+ flash attention and deal with padding tokens in case the input contains any of them.
735
+ """
736
+
737
+ def __init__(self, *args, **kwargs):
738
+ super().__init__(*args, **kwargs)
739
+
740
+ self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10()
741
+
742
+ def forward(
743
+ self,
744
+ hidden_states: torch.Tensor,
745
+ attention_mask: Optional[torch.LongTensor] = None,
746
+ position_ids: Optional[torch.LongTensor] = None,
747
+ past_key_value: Optional[Cache] = None,
748
+ output_attentions: bool = False,
749
+ use_cache: bool = False,
750
+ kv_states: torch.Tensor = None,
751
+ **kwargs,
752
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
753
+ # HunYuanFlashAttention2 attention does not support output_attentions
754
+ if "padding_mask" in kwargs:
755
+ warnings.warn(
756
+ "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use "
757
+ "`attention_mask` instead.`"
758
+ )
759
+
760
+ # overwrite attention_mask with padding_mask
761
+ attention_mask = kwargs.pop("padding_mask")
762
+
763
+ bsz, q_len, _ = hidden_states.size()
764
+
765
+ query_states = self.q_proj(hidden_states)
766
+ if self.attention_type == "cross" and kv_states is not None and isinstance(kv_states, tuple):
767
+ orig_key_states, orig_value_states = kv_states
768
+ key_states, value_states = kv_states
769
+ else:
770
+ key_states = self.k_proj(hidden_states)
771
+ value_states = self.v_proj(hidden_states)
772
+ orig_key_states, orig_value_states = key_states, value_states
773
+
774
+ # Flash attention requires the input to have the shape
775
+ # batch_size x seq_length x head_dim x hidden_dim
776
+ # therefore we just need to keep the original shape
777
+ query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
778
+ key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
779
+ value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
780
+
781
+ kv_seq_len = key_states.shape[-2]
782
+ if past_key_value is not None:
783
+ kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
784
+ cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
785
+ query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
786
+
787
+ if self.use_qk_norm:
788
+ query_states = self.query_layernorm(query_states)
789
+ key_states = self.key_layernorm(key_states)
790
+
791
+ if past_key_value is not None:
792
+ cache_kwargs = {"sin": sin, "cos": cos} # Specific to RoPE models
793
+ key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
794
+
795
+ query_states = query_states.transpose(1, 2)
796
+ key_states = key_states.transpose(1, 2)
797
+ value_states = value_states.transpose(1, 2)
798
+
799
+ dropout_rate = self.attention_dropout if self.training else 0.0
800
+
801
+ # In PEFT, usually we cast the layer norms in float32 for training stability reasons
802
+ # therefore the input hidden states gets silently casted in float32. Hence, we need
803
+ # cast them back in the correct dtype just to be sure everything works as expected.
804
+ # This might slowdown training & inference so it is recommended to not cast the LayerNorms
805
+ # in fp32. (HunYuanRMSNorm handles it correctly)
806
+
807
+ input_dtype = query_states.dtype
808
+ if input_dtype == torch.float32:
809
+ # Handle the case where the model is quantized
810
+ if hasattr(self.config, "_pre_quantization_dtype"):
811
+ target_dtype = self.config._pre_quantization_dtype
812
+ else:
813
+ target_dtype = self.q_proj.weight.dtype
814
+
815
+ logger.warning_once(
816
+ f"The input hidden states seems to be silently casted in float32, this might be related to"
817
+ f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in"
818
+ f" {target_dtype}."
819
+ )
820
+
821
+ query_states = query_states.to(target_dtype)
822
+ key_states = key_states.to(target_dtype)
823
+ value_states = value_states.to(target_dtype)
824
+
825
+ attn_output = self._flash_attention_forward(
826
+ query_states, key_states, value_states, attention_mask, q_len, dropout=dropout_rate
827
+ )
828
+
829
+ attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous()
830
+ attn_output = self.o_proj(attn_output)
831
+
832
+ return attn_output, None, past_key_value, (orig_key_states, orig_value_states)
833
+
834
+ def _flash_attention_forward(
835
+ self, query_states, key_states, value_states, attention_mask, query_length, dropout=0.0, softmax_scale=None
836
+ ):
837
+ """
838
+ Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token
839
+ first unpad the input, then computes the attention scores and pad the final attention scores.
840
+
841
+ Args:
842
+ query_states (`torch.Tensor`):
843
+ Input query states to be passed to Flash Attention API
844
+ key_states (`torch.Tensor`):
845
+ Input key states to be passed to Flash Attention API
846
+ value_states (`torch.Tensor`):
847
+ Input value states to be passed to Flash Attention API
848
+ attention_mask (`torch.Tensor`):
849
+ The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the
850
+ position of padding tokens and 1 for the position of non-padding tokens.
851
+ dropout (`int`, *optional*):
852
+ Attention dropout
853
+ softmax_scale (`float`, *optional*):
854
+ The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim)
855
+ """
856
+ if not self._flash_attn_uses_top_left_mask:
857
+ causal = self.is_causal
858
+ else:
859
+ causal = self.is_causal and query_length != 1
860
+
861
+ # Contains at least one padding token in the sequence
862
+ if attention_mask is not None:
863
+ batch_size = query_states.shape[0]
864
+ query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._upad_input(
865
+ query_states, key_states, value_states, attention_mask, query_length
866
+ )
867
+
868
+ cu_seqlens_q, cu_seqlens_k = cu_seq_lens
869
+ max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens
870
+
871
+ attn_output_unpad = flash_attn_varlen_func(
872
+ query_states,
873
+ key_states,
874
+ value_states,
875
+ cu_seqlens_q=cu_seqlens_q,
876
+ cu_seqlens_k=cu_seqlens_k,
877
+ max_seqlen_q=max_seqlen_in_batch_q,
878
+ max_seqlen_k=max_seqlen_in_batch_k,
879
+ dropout_p=dropout,
880
+ softmax_scale=softmax_scale,
881
+ causal=causal,
882
+ )
883
+
884
+ attn_output = pad_input(attn_output_unpad, indices_q, batch_size, query_length)
885
+ else:
886
+ attn_output = flash_attn_func(
887
+ query_states, key_states, value_states, dropout, softmax_scale=softmax_scale, causal=causal
888
+ )
889
+
890
+ return attn_output
891
+
892
+ def _upad_input(self, query_layer, key_layer, value_layer, attention_mask, query_length):
893
+ indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)
894
+ batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape
895
+
896
+ key_layer = index_first_axis(
897
+ key_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k
898
+ )
899
+ value_layer = index_first_axis(
900
+ value_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k
901
+ )
902
+ if query_length == kv_seq_len:
903
+ query_layer = index_first_axis(
904
+ query_layer.reshape(batch_size * kv_seq_len, self.num_heads, head_dim), indices_k
905
+ )
906
+ cu_seqlens_q = cu_seqlens_k
907
+ max_seqlen_in_batch_q = max_seqlen_in_batch_k
908
+ indices_q = indices_k
909
+ elif query_length == 1:
910
+ max_seqlen_in_batch_q = 1
911
+ cu_seqlens_q = torch.arange(
912
+ batch_size + 1, dtype=torch.int32, device=query_layer.device
913
+ ) # There is a memcpy here, that is very bad.
914
+ indices_q = cu_seqlens_q[:-1]
915
+ query_layer = query_layer.squeeze(1)
916
+ else:
917
+ # The -q_len: slice assumes left padding.
918
+ attention_mask = attention_mask[:, -query_length:]
919
+ query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(query_layer, attention_mask)
920
+
921
+ return (
922
+ query_layer,
923
+ key_layer,
924
+ value_layer,
925
+ indices_q,
926
+ (cu_seqlens_q, cu_seqlens_k),
927
+ (max_seqlen_in_batch_q, max_seqlen_in_batch_k),
928
+ )
929
+
930
+
931
+ class HunYuanSdpaAttention(HunYuanAttention):
932
+ """
933
+ HunYuan attention module using torch.nn.functional.scaled_dot_product_attention. This module inherits from
934
+ `HunYuanAttention` as the weights of the module stays untouched. The only changes are on the forward pass to adapt
935
+ to SDPA API.
936
+ """
937
+
938
+ # Adapted from HunYuanAttention.forward
939
+ def forward(
940
+ self,
941
+ hidden_states: torch.Tensor,
942
+ attention_mask: Optional[torch.Tensor] = None,
943
+ position_ids: Optional[torch.LongTensor] = None,
944
+ past_key_value: Optional[Cache] = None,
945
+ output_attentions: bool = False,
946
+ use_cache: bool = False,
947
+ kv_states: torch.Tensor = None,
948
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
949
+ if output_attentions:
950
+ logger.warning_once(
951
+ 'HunYuanModel is using HunYuanSdpaAttention,'
952
+ 'but `torch.nn.functional.scaled_dot_product_attention`'
953
+ 'does not support `output_attentions=True`. Falling back to the manual attention implementation, '
954
+ 'but specifying the manual implementation will be required from Transformers version v5.0.0 onwards. '
955
+ 'This warning can be removed using the argument `attn_implementation="eager"` when loading the model.'
956
+ )
957
+ return super().forward(
958
+ hidden_states=hidden_states,
959
+ attention_mask=attention_mask,
960
+ position_ids=position_ids,
961
+ past_key_value=past_key_value,
962
+ output_attentions=output_attentions,
963
+ use_cache=use_cache,
964
+ )
965
+
966
+ bsz, q_len, _ = hidden_states.size()
967
+
968
+ query_states = self.q_proj(hidden_states)
969
+ if self.attention_type == "cross" and kv_states is not None and isinstance(kv_states, tuple):
970
+ orig_key_states, orig_value_states = kv_states
971
+ key_states, value_states = kv_states
972
+ else:
973
+ key_states = self.k_proj(hidden_states)
974
+ value_states = self.v_proj(hidden_states)
975
+ orig_key_states, orig_value_states = key_states, value_states
976
+
977
+ query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
978
+ key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
979
+ value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
980
+
981
+ kv_seq_len = key_states.shape[-2]
982
+ if past_key_value is not None:
983
+ kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
984
+ cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
985
+
986
+ query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
987
+
988
+ if self.use_qk_norm:
989
+ query_states = self.query_layernorm(query_states)
990
+ key_states = self.key_layernorm(key_states)
991
+
992
+ if past_key_value is not None:
993
+ cache_kwargs = {"sin": sin, "cos": cos} # Specific to RoPE models
994
+ key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
995
+
996
+ key_states = repeat_kv(key_states, self.num_key_value_groups)
997
+ value_states = repeat_kv(value_states, self.num_key_value_groups)
998
+
999
+ if attention_mask is not None:
1000
+ if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
1001
+ raise ValueError(
1002
+ f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
1003
+ )
1004
+
1005
+ # SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with
1006
+ # custom attn_mask,
1007
+ # Reference: https://github.com/pytorch/pytorch/issues/112577.
1008
+ if query_states.device.type == "cuda" and attention_mask is not None:
1009
+ query_states = query_states.contiguous()
1010
+ key_states = key_states.contiguous()
1011
+ value_states = value_states.contiguous()
1012
+
1013
+ attn_output = torch.nn.functional.scaled_dot_product_attention(
1014
+ query_states,
1015
+ key_states,
1016
+ value_states,
1017
+ attn_mask=attention_mask,
1018
+ dropout_p=self.attention_dropout if self.training else 0.0,
1019
+ # The q_len > 1 is necessary to match with AttentionMaskConverter.to_causal_4d that does not create a
1020
+ # causal mask in case q_len == 1.
1021
+ is_causal=self.is_causal and attention_mask is None and q_len > 1,
1022
+ )
1023
+
1024
+ attn_output = attn_output.transpose(1, 2).contiguous()
1025
+ attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
1026
+
1027
+ attn_output = self.o_proj(attn_output)
1028
+
1029
+ return attn_output, None, past_key_value, (orig_key_states, orig_value_states)
1030
+
1031
+
1032
+ HUNYUAN_ATTENTION_CLASSES = {
1033
+ "eager": HunYuanAttention,
1034
+ "flash_attention_2": HunYuanFlashAttention2,
1035
+ "sdpa": HunYuanSdpaAttention,
1036
+ }
1037
+
1038
+
1039
+ class HunYuanDecoderLayer(nn.Module):
1040
+ def __init__(self, config: HunYuanConfig, layer_idx: int):
1041
+ super().__init__()
1042
+ self.hidden_size = config.hidden_size
1043
+ self.layer_idx = layer_idx
1044
+
1045
+ self.self_attn = HUNYUAN_ATTENTION_CLASSES[config._attn_implementation](config=config, layer_idx=layer_idx)
1046
+
1047
+ if config.num_experts > 1:
1048
+ self.mlp = HunYuanMoE(config, layer_idx=layer_idx)
1049
+ else:
1050
+ self.mlp = HunYuanMLP(config, layer_idx=layer_idx, is_shared_mlp=False)
1051
+ self.input_layernorm = HunYuanRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
1052
+ self.post_attention_layernorm = HunYuanRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
1053
+
1054
+ def forward(
1055
+ self,
1056
+ hidden_states: torch.Tensor,
1057
+ attention_mask: Optional[torch.Tensor] = None,
1058
+ position_ids: Optional[torch.LongTensor] = None,
1059
+ past_key_value: Optional[Tuple[torch.Tensor]] = None,
1060
+ output_attentions: Optional[bool] = False,
1061
+ use_cache: Optional[bool] = False,
1062
+ kv_states: Optional[Tuple[torch.Tensor]] = None,
1063
+ **kwargs,
1064
+ ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
1065
+ """
1066
+ Args:
1067
+ hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
1068
+ attention_mask (`torch.FloatTensor`, *optional*):
1069
+ attention mask of size `(batch_size, sequence_length)` if flash attention is used or `(batch_size, 1,
1070
+ query_sequence_length, key_sequence_length)` if default attention is used.
1071
+ output_attentions (`bool`, *optional*):
1072
+ Whether or not to return the attentions tensors of all attention layers. See `attentions` under
1073
+ returned tensors for more detail.
1074
+ use_cache (`bool`, *optional*):
1075
+ If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
1076
+ (see `past_key_values`).
1077
+ past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
1078
+ kv_states (`Tuple(torch.FloatTensor)`, *optional*): Used when CLA is enabled,
1079
+ key and value states from past attention blocks
1080
+ """
1081
+ if "padding_mask" in kwargs:
1082
+ warnings.warn(
1083
+ "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use "
1084
+ "`attention_mask` instead.`"
1085
+ )
1086
+
1087
+ residual = hidden_states
1088
+
1089
+ hidden_states = self.input_layernorm(hidden_states)
1090
+
1091
+ # Self Attention
1092
+ hidden_states, self_attn_weights, present_key_value, kv_states = self.self_attn(
1093
+ hidden_states=hidden_states,
1094
+ attention_mask=attention_mask,
1095
+ position_ids=position_ids,
1096
+ past_key_value=past_key_value,
1097
+ output_attentions=output_attentions,
1098
+ use_cache=use_cache,
1099
+ kv_states=kv_states,
1100
+ **kwargs,
1101
+ )
1102
+ hidden_states = residual + hidden_states
1103
+
1104
+ # Fully Connected
1105
+ residual = hidden_states
1106
+ hidden_states = self.post_attention_layernorm(hidden_states)
1107
+ hidden_states = self.mlp(hidden_states)
1108
+ hidden_states = residual + hidden_states
1109
+
1110
+ outputs = (hidden_states,)
1111
+
1112
+ if output_attentions:
1113
+ outputs += (self_attn_weights,)
1114
+
1115
+ if use_cache:
1116
+ outputs += (present_key_value,)
1117
+
1118
+ outputs += (kv_states,)
1119
+
1120
+ return outputs
1121
+
1122
+
1123
+ HUNYUAN_START_DOCSTRING = r"""
1124
+ This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
1125
+ library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
1126
+ etc.)
1127
+
1128
+ This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
1129
+ Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
1130
+ and behavior.
1131
+
1132
+ Parameters:
1133
+ config ([`HunYuanConfig`]):
1134
+ Model configuration class with all the parameters of the model. Initializing with a config file does not
1135
+ load the weights associated with the model, only the configuration. Check out the
1136
+ [`~PreTrainedModel.from_pretrained`] method to load the model weights.
1137
+ """
1138
+
1139
+
1140
+ @add_start_docstrings(
1141
+ "The bare HunYuan Model outputting raw hidden-states without any specific head on top.",
1142
+ HUNYUAN_START_DOCSTRING,
1143
+ )
1144
+ class HunYuanPreTrainedModel(PreTrainedModel):
1145
+ config_class = HunYuanConfig
1146
+ base_model_prefix = "model"
1147
+ supports_gradient_checkpointing = True
1148
+ _no_split_modules = ["HunYuanDecoderLayer"]
1149
+ _skip_keys_device_placement = "past_key_values"
1150
+ _supports_flash_attn_2 = True
1151
+ _supports_sdpa = True
1152
+ _supports_cache_class = True
1153
+
1154
+ def _init_weights(self, module):
1155
+ std = self.config.initializer_range
1156
+ if isinstance(module, nn.Linear):
1157
+ module.weight.data.normal_(mean=0.0, std=std)
1158
+ if module.bias is not None:
1159
+ module.bias.data.zero_()
1160
+ elif isinstance(module, nn.Embedding):
1161
+ module.weight.data.normal_(mean=0.0, std=std)
1162
+ if module.padding_idx is not None:
1163
+ module.weight.data[module.padding_idx].zero_()
1164
+
1165
+
1166
+ HUNYUAN_INPUTS_DOCSTRING = r"""
1167
+ Args:
1168
+ input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
1169
+ Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide
1170
+ it.
1171
+
1172
+ Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
1173
+ [`PreTrainedTokenizer.__call__`] for details.
1174
+
1175
+ [What are input IDs?](../glossary#input-ids)
1176
+ attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
1177
+ Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
1178
+
1179
+ - 1 for tokens that are **not masked**,
1180
+ - 0 for tokens that are **masked**.
1181
+
1182
+ [What are attention masks?](../glossary#attention-mask)
1183
+
1184
+ Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
1185
+ [`PreTrainedTokenizer.__call__`] for details.
1186
+
1187
+ If `past_key_values` is used, optionally only the last `input_ids` have to be input (see
1188
+ `past_key_values`).
1189
+
1190
+ If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`]
1191
+ and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more
1192
+ information on the default strategy.
1193
+
1194
+ - 1 indicates the head is **not masked**,
1195
+ - 0 indicates the head is **masked**.
1196
+ position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
1197
+ Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
1198
+ config.n_positions - 1]`.
1199
+
1200
+ [What are position IDs?](../glossary#position-ids)
1201
+ past_key_values (`Cache` or `tuple(tuple(torch.FloatTensor))`, *optional*):
1202
+ Pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention
1203
+ blocks) that can be used to speed up sequential decoding. This typically consists in the `past_key_values`
1204
+ returned by the model at a previous stage of decoding, when `use_cache=True` or `config.use_cache=True`.
1205
+
1206
+ Two formats are allowed:
1207
+ - a [`~cache_utils.Cache`] instance;
1208
+ - Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of
1209
+ shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`). This is also known as the legacy
1210
+ cache format.
1211
+
1212
+ The model will output the same cache format that is fed as input. If no `past_key_values` are passed, the
1213
+ legacy cache format will be returned.
1214
+
1215
+ If `past_key_values` are used, the user can optionally input only the last `input_ids` (those that don't
1216
+ have their past key value states given to this model) of shape `(batch_size, 1)` instead of all `input_ids`
1217
+ of shape `(batch_size, sequence_length)`.
1218
+ inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
1219
+ Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
1220
+ is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
1221
+ model's internal embedding lookup matrix.
1222
+ use_cache (`bool`, *optional*):
1223
+ If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
1224
+ `past_key_values`).
1225
+ output_attentions (`bool`, *optional*):
1226
+ Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
1227
+ tensors for more detail.
1228
+ output_hidden_states (`bool`, *optional*):
1229
+ Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
1230
+ more detail.
1231
+ return_dict (`bool`, *optional*):
1232
+ Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
1233
+ """
1234
+
1235
+
1236
+ @add_start_docstrings(
1237
+ "The bare HunYuan Model outputting raw hidden-states without any specific head on top.",
1238
+ HUNYUAN_START_DOCSTRING,
1239
+ )
1240
+ class HunYuanModel(HunYuanPreTrainedModel):
1241
+ """
1242
+ Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`HunYuanDecoderLayer`]
1243
+
1244
+ Args:
1245
+ config: HunYuanConfig
1246
+ """
1247
+
1248
+ def __init__(self, config: HunYuanConfig):
1249
+ super().__init__(config)
1250
+ self.padding_idx = config.pad_token_id
1251
+ self.vocab_size = config.vocab_size
1252
+
1253
+ self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
1254
+ self.layers = nn.ModuleList(
1255
+ [HunYuanDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
1256
+ )
1257
+ self._use_sdpa = config._attn_implementation == "sdpa"
1258
+ self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2"
1259
+ self.norm = HunYuanRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
1260
+
1261
+ self.cla = config.use_cla
1262
+ self.cla_share_factor = config.cla_share_factor
1263
+
1264
+ self.gradient_checkpointing = False
1265
+ # Initialize weights and apply final processing
1266
+ self.post_init()
1267
+
1268
+ def get_input_embeddings(self):
1269
+ return self.embed_tokens
1270
+
1271
+ def set_input_embeddings(self, value):
1272
+ self.embed_tokens = value
1273
+
1274
+ @add_start_docstrings_to_model_forward(HUNYUAN_INPUTS_DOCSTRING)
1275
+ def forward(
1276
+ self,
1277
+ input_ids: torch.LongTensor = None,
1278
+ attention_mask: Optional[torch.Tensor] = None,
1279
+ position_ids: Optional[torch.LongTensor] = None,
1280
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
1281
+ inputs_embeds: Optional[torch.FloatTensor] = None,
1282
+ use_cache: Optional[bool] = None,
1283
+ output_attentions: Optional[bool] = None,
1284
+ output_hidden_states: Optional[bool] = None,
1285
+ return_dict: Optional[bool] = None,
1286
+ ) -> Union[Tuple, BaseModelOutputWithPast]:
1287
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
1288
+ output_hidden_states = (
1289
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
1290
+ )
1291
+ use_cache = use_cache if use_cache is not None else self.config.use_cache
1292
+
1293
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
1294
+
1295
+ # retrieve input_ids and inputs_embeds
1296
+ if input_ids is not None and inputs_embeds is not None:
1297
+ raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
1298
+ elif input_ids is not None:
1299
+ batch_size, seq_length = input_ids.shape[:2]
1300
+ elif inputs_embeds is not None:
1301
+ batch_size, seq_length = inputs_embeds.shape[:2]
1302
+ else:
1303
+ raise ValueError("You have to specify either input_ids or inputs_embeds")
1304
+
1305
+ if self.gradient_checkpointing and self.training:
1306
+ if use_cache:
1307
+ logger.warning_once(
1308
+ "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
1309
+ )
1310
+ use_cache = False
1311
+
1312
+ past_key_values_length = 0
1313
+ if use_cache:
1314
+ use_legacy_cache = not isinstance(past_key_values, Cache)
1315
+ if use_legacy_cache:
1316
+ past_key_values = DynamicCache.from_legacy_cache(past_key_values)
1317
+ past_key_values_length = past_key_values.get_usable_length(seq_length)
1318
+
1319
+ if position_ids is None:
1320
+ device = input_ids.device if input_ids is not None else inputs_embeds.device
1321
+ position_ids = torch.arange(
1322
+ past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=device
1323
+ )
1324
+ position_ids = position_ids.unsqueeze(0)
1325
+
1326
+ if inputs_embeds is None:
1327
+ inputs_embeds = self.embed_tokens(input_ids)
1328
+
1329
+ # Fix lora with gradient checkpointing training
1330
+ if self.training and inputs_embeds.is_leaf:
1331
+ inputs_embeds.requires_grad = True
1332
+
1333
+ if self._use_flash_attention_2:
1334
+ # 2d mask is passed through the layers
1335
+ attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) else None
1336
+ elif self._use_sdpa and not output_attentions:
1337
+ # output_attentions=True can not be supported when using SDPA, and we fall back on
1338
+ # the manual implementation that requires a 4D causal mask in all cases.
1339
+ attention_mask = _prepare_4d_causal_attention_mask_for_sdpa(
1340
+ attention_mask,
1341
+ (batch_size, seq_length),
1342
+ inputs_embeds,
1343
+ past_key_values_length,
1344
+ )
1345
+ else:
1346
+ # 4d mask is passed through the layers
1347
+ attention_mask = _prepare_4d_causal_attention_mask(
1348
+ attention_mask, (batch_size, seq_length), inputs_embeds, past_key_values_length
1349
+ )
1350
+
1351
+ # embed positions
1352
+ hidden_states = inputs_embeds
1353
+
1354
+ # decoder layers
1355
+ all_hidden_states = () if output_hidden_states else None
1356
+ all_self_attns = () if output_attentions else None
1357
+ next_decoder_cache = None
1358
+
1359
+ prev_kv_states = None
1360
+ for layer_idx, decoder_layer in enumerate(self.layers):
1361
+ if output_hidden_states:
1362
+ all_hidden_states += (hidden_states,)
1363
+
1364
+ if self.gradient_checkpointing and self.training:
1365
+ layer_outputs = self._gradient_checkpointing_func(
1366
+ decoder_layer.__call__,
1367
+ hidden_states,
1368
+ attention_mask,
1369
+ position_ids,
1370
+ past_key_values,
1371
+ output_attentions,
1372
+ use_cache,
1373
+ prev_kv_states,
1374
+ )
1375
+ else:
1376
+ layer_outputs = decoder_layer(
1377
+ hidden_states,
1378
+ attention_mask=attention_mask,
1379
+ position_ids=position_ids,
1380
+ past_key_value=past_key_values,
1381
+ output_attentions=output_attentions,
1382
+ use_cache=use_cache,
1383
+ kv_states=prev_kv_states
1384
+ )
1385
+
1386
+ hidden_states = layer_outputs[0]
1387
+
1388
+ if use_cache:
1389
+ next_decoder_cache = layer_outputs[2 if output_attentions else 1]
1390
+
1391
+ if output_attentions:
1392
+ all_self_attns += (layer_outputs[1],)
1393
+
1394
+ kv_states = layer_outputs[-1]
1395
+
1396
+ if self.cla and layer_idx % self.cla_share_factor == 0:
1397
+ prev_kv_states = kv_states
1398
+
1399
+ hidden_states = self.norm(hidden_states)
1400
+
1401
+ # add hidden states from the last decoder layer
1402
+ if output_hidden_states:
1403
+ all_hidden_states += (hidden_states,)
1404
+
1405
+ next_cache = None
1406
+ if use_cache:
1407
+ next_cache = next_decoder_cache.to_legacy_cache() if use_legacy_cache else next_decoder_cache
1408
+ if not return_dict:
1409
+ return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)
1410
+ return BaseModelOutputWithPast(
1411
+ last_hidden_state=hidden_states,
1412
+ past_key_values=next_cache,
1413
+ hidden_states=all_hidden_states,
1414
+ attentions=all_self_attns,
1415
+ )
1416
+
1417
+
1418
+ class HunYuanForCausalLM(HunYuanPreTrainedModel):
1419
+ _tied_weights_keys = ["lm_head.weight"]
1420
+
1421
+ def __init__(self, config: HunYuanConfig):
1422
+ super().__init__(config)
1423
+ self.model = HunYuanModel(config)
1424
+ self.vocab_size = config.vocab_size
1425
+ self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
1426
+
1427
+ # Initialize weights and apply final processing
1428
+ self.post_init()
1429
+
1430
+ def get_input_embeddings(self):
1431
+ return self.model.embed_tokens
1432
+
1433
+ def set_input_embeddings(self, value):
1434
+ self.model.embed_tokens = value
1435
+
1436
+ def get_output_embeddings(self):
1437
+ return self.lm_head
1438
+
1439
+ def set_output_embeddings(self, new_embeddings):
1440
+ self.lm_head = new_embeddings
1441
+
1442
+ def set_decoder(self, decoder):
1443
+ self.model = decoder
1444
+
1445
+ def get_decoder(self):
1446
+ return self.model
1447
+
1448
+ @add_start_docstrings_to_model_forward(HUNYUAN_INPUTS_DOCSTRING)
1449
+ @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)
1450
+ def forward(
1451
+ self,
1452
+ input_ids: torch.LongTensor = None,
1453
+ attention_mask: Optional[torch.Tensor] = None,
1454
+ position_ids: Optional[torch.LongTensor] = None,
1455
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
1456
+ inputs_embeds: Optional[torch.FloatTensor] = None,
1457
+ labels: Optional[torch.LongTensor] = None,
1458
+ use_cache: Optional[bool] = None,
1459
+ output_attentions: Optional[bool] = None,
1460
+ output_hidden_states: Optional[bool] = None,
1461
+ return_dict: Optional[bool] = None,
1462
+ ) -> Union[Tuple, CausalLMOutputWithPast]:
1463
+ r"""
1464
+ Args:
1465
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
1466
+ Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
1467
+ config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
1468
+ (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
1469
+
1470
+ Returns:
1471
+
1472
+ Example:
1473
+
1474
+ ```python
1475
+ >>> from transformers import AutoTokenizer, AutoModelForCausalLM
1476
+
1477
+ >>> model = AutoModelForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS)
1478
+ >>> tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER)
1479
+
1480
+ >>> prompt = "Hey, are you conscious? Can you talk to me?"
1481
+ >>> inputs = tokenizer(prompt, return_tensors="pt")
1482
+
1483
+ >>> # Generate
1484
+ >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
1485
+ >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
1486
+ "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
1487
+ ```"""
1488
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
1489
+ output_hidden_states = (
1490
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
1491
+ )
1492
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
1493
+
1494
+ # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
1495
+ outputs = self.model(
1496
+ input_ids=input_ids,
1497
+ attention_mask=attention_mask,
1498
+ position_ids=position_ids,
1499
+ past_key_values=past_key_values,
1500
+ inputs_embeds=inputs_embeds,
1501
+ use_cache=use_cache,
1502
+ output_attentions=output_attentions,
1503
+ output_hidden_states=output_hidden_states,
1504
+ return_dict=return_dict,
1505
+ )
1506
+
1507
+ hidden_states = outputs[0]
1508
+ if self.config.pretraining_tp > 1:
1509
+ lm_head_slices = self.lm_head.weight.split(self.vocab_size // self.config.pretraining_tp, dim=0)
1510
+ logits = [F.linear(hidden_states, lm_head_slices[i]) for i in range(self.config.pretraining_tp)]
1511
+ logits = torch.cat(logits, dim=-1)
1512
+ else:
1513
+ logits = self.lm_head(hidden_states)
1514
+ logits = logits.float()
1515
+
1516
+ loss = None
1517
+ if labels is not None:
1518
+ # Shift so that tokens < n predict n
1519
+ shift_logits = logits[..., :-1, :].contiguous()
1520
+ shift_labels = labels[..., 1:].contiguous()
1521
+ # Flatten the tokens
1522
+ loss_fct = CrossEntropyLoss()
1523
+ shift_logits = shift_logits.view(-1, self.config.vocab_size)
1524
+ shift_labels = shift_labels.view(-1)
1525
+ # Enable model parallelism
1526
+ shift_labels = shift_labels.to(shift_logits.device)
1527
+ loss = loss_fct(shift_logits, shift_labels)
1528
+
1529
+ if not return_dict:
1530
+ output = (logits,) + outputs[1:]
1531
+ return (loss,) + output if loss is not None else output
1532
+
1533
+ return CausalLMOutputWithPast(
1534
+ loss=loss,
1535
+ logits=logits,
1536
+ past_key_values=outputs.past_key_values,
1537
+ hidden_states=outputs.hidden_states,
1538
+ attentions=outputs.attentions,
1539
+ )
1540
+
1541
+ def prepare_inputs_for_generation(
1542
+ self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, **kwargs
1543
+ ):
1544
+ if past_key_values is not None:
1545
+ if isinstance(past_key_values, Cache):
1546
+ cache_length = past_key_values.get_seq_length()
1547
+ past_length = past_key_values.seen_tokens
1548
+ max_cache_length = past_key_values.get_max_length()
1549
+ else:
1550
+ cache_length = past_length = past_key_values[0][0].shape[2]
1551
+ max_cache_length = None
1552
+
1553
+ # Keep only the unprocessed tokens:
1554
+ # 1 - If the length of the attention_mask exceeds the length of input_ids, then we are in a setting where
1555
+ # some of the inputs are exclusivelly passed as part of the cache (e.g. when passing input_embeds as
1556
+ # input)
1557
+ if attention_mask is not None and attention_mask.shape[1] > input_ids.shape[1]:
1558
+ input_ids = input_ids[:, -(attention_mask.shape[1] - past_length):]
1559
+ # 2 - If the past_length is smaller than input_ids', then input_ids holds all input tokens. We can discard
1560
+ # input_ids based on the past_length.
1561
+ elif past_length < input_ids.shape[1]:
1562
+ input_ids = input_ids[:, past_length:]
1563
+ # 3 - Otherwise (past_length >= input_ids.shape[1]), let's assume input_ids only has unprocessed tokens.
1564
+
1565
+ # If we are about to go beyond the maximum cache length, we need to crop the input attention mask.
1566
+ if (
1567
+ max_cache_length is not None
1568
+ and attention_mask is not None
1569
+ and cache_length + input_ids.shape[1] > max_cache_length
1570
+ ):
1571
+ attention_mask = attention_mask[:, -max_cache_length:]
1572
+
1573
+ position_ids = kwargs.get("position_ids", None)
1574
+ if attention_mask is not None and position_ids is None:
1575
+ # create position_ids on the fly for batch generation
1576
+ position_ids = attention_mask.long().cumsum(-1) - 1
1577
+ position_ids.masked_fill_(attention_mask == 0, 1)
1578
+ if past_key_values:
1579
+ position_ids = position_ids[:, -input_ids.shape[1]:]
1580
+
1581
+ # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
1582
+ if inputs_embeds is not None and past_key_values is None:
1583
+ model_inputs = {"inputs_embeds": inputs_embeds}
1584
+ else:
1585
+ model_inputs = {"input_ids": input_ids}
1586
+
1587
+ model_inputs.update(
1588
+ {
1589
+ "position_ids": position_ids,
1590
+ "past_key_values": past_key_values,
1591
+ "use_cache": kwargs.get("use_cache"),
1592
+ "attention_mask": attention_mask,
1593
+ }
1594
+ )
1595
+ return model_inputs
1596
+
1597
+ @staticmethod
1598
+ def _reorder_cache(past_key_values, beam_idx):
1599
+ reordered_past = ()
1600
+ for layer_past in past_key_values:
1601
+ reordered_past += (
1602
+ tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past),
1603
+ )
1604
+ return reordered_past
1605
+
1606
+
1607
+ @add_start_docstrings(
1608
+ """
1609
+ The HunYuan Model transformer with a sequence classification head on top (linear layer).
1610
+
1611
+ [`HunYuanForSequenceClassification`] uses the last token in order to do the classification, as other causal models
1612
+ (e.g. GPT-2) do.
1613
+
1614
+ Since it does classification on the last token, it requires to know the position of the last token. If a
1615
+ `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If
1616
+ no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the
1617
+ padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in
1618
+ each row of the batch).
1619
+ """,
1620
+ HUNYUAN_START_DOCSTRING,
1621
+ )
1622
+ class HunYuanForSequenceClassification(HunYuanPreTrainedModel):
1623
+ def __init__(self, config):
1624
+ super().__init__(config)
1625
+ self.num_labels = config.num_labels
1626
+ self.model = HunYuanModel(config)
1627
+ self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False)
1628
+
1629
+ # Initialize weights and apply final processing
1630
+ self.post_init()
1631
+
1632
+ def get_input_embeddings(self):
1633
+ return self.model.embed_tokens
1634
+
1635
+ def set_input_embeddings(self, value):
1636
+ self.model.embed_tokens = value
1637
+
1638
+ @add_start_docstrings_to_model_forward(HUNYUAN_INPUTS_DOCSTRING)
1639
+ def forward(
1640
+ self,
1641
+ input_ids: torch.LongTensor = None,
1642
+ attention_mask: Optional[torch.Tensor] = None,
1643
+ position_ids: Optional[torch.LongTensor] = None,
1644
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
1645
+ inputs_embeds: Optional[torch.FloatTensor] = None,
1646
+ labels: Optional[torch.LongTensor] = None,
1647
+ use_cache: Optional[bool] = None,
1648
+ output_attentions: Optional[bool] = None,
1649
+ output_hidden_states: Optional[bool] = None,
1650
+ return_dict: Optional[bool] = None,
1651
+ ) -> Union[Tuple, SequenceClassifierOutputWithPast]:
1652
+ r"""
1653
+ labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
1654
+ Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
1655
+ config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
1656
+ `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
1657
+ """
1658
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
1659
+
1660
+ transformer_outputs = self.model(
1661
+ input_ids,
1662
+ attention_mask=attention_mask,
1663
+ position_ids=position_ids,
1664
+ past_key_values=past_key_values,
1665
+ inputs_embeds=inputs_embeds,
1666
+ use_cache=use_cache,
1667
+ output_attentions=output_attentions,
1668
+ output_hidden_states=output_hidden_states,
1669
+ return_dict=return_dict,
1670
+ )
1671
+ hidden_states = transformer_outputs[0]
1672
+ logits = self.score(hidden_states)
1673
+
1674
+ if input_ids is not None:
1675
+ batch_size = input_ids.shape[0]
1676
+ else:
1677
+ batch_size = inputs_embeds.shape[0]
1678
+
1679
+ if self.config.pad_token_id is None and batch_size != 1:
1680
+ raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.")
1681
+ if self.config.pad_token_id is None:
1682
+ sequence_lengths = -1
1683
+ else:
1684
+ if input_ids is not None:
1685
+ sequence_lengths = (torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1).to(
1686
+ logits.device
1687
+ )
1688
+ else:
1689
+ sequence_lengths = -1
1690
+
1691
+ pooled_logits = logits[torch.arange(batch_size, device=logits.device), sequence_lengths]
1692
+
1693
+ loss = None
1694
+ if labels is not None:
1695
+ labels = labels.to(logits.device)
1696
+ if self.config.problem_type is None:
1697
+ if self.num_labels == 1:
1698
+ self.config.problem_type = "regression"
1699
+ elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
1700
+ self.config.problem_type = "single_label_classification"
1701
+ else:
1702
+ self.config.problem_type = "multi_label_classification"
1703
+
1704
+ if self.config.problem_type == "regression":
1705
+ loss_fct = MSELoss()
1706
+ if self.num_labels == 1:
1707
+ loss = loss_fct(pooled_logits.squeeze(), labels.squeeze())
1708
+ else:
1709
+ loss = loss_fct(pooled_logits, labels)
1710
+ elif self.config.problem_type == "single_label_classification":
1711
+ loss_fct = CrossEntropyLoss()
1712
+ loss = loss_fct(pooled_logits.view(-1, self.num_labels), labels.view(-1))
1713
+ elif self.config.problem_type == "multi_label_classification":
1714
+ loss_fct = BCEWithLogitsLoss()
1715
+ loss = loss_fct(pooled_logits, labels)
1716
+ if not return_dict:
1717
+ output = (pooled_logits,) + transformer_outputs[1:]
1718
+ return ((loss,) + output) if loss is not None else output
1719
+
1720
+ return SequenceClassifierOutputWithPast(
1721
+ loss=loss,
1722
+ logits=pooled_logits,
1723
+ past_key_values=transformer_outputs.past_key_values,
1724
+ hidden_states=transformer_outputs.hidden_states,
1725
+ attentions=transformer_outputs.attentions,
1726
+ )
special_tokens_map.json ADDED
@@ -0,0 +1,12 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "additional_special_tokens": [
3
+ "<|startoftext|>",
4
+ "<|extra_0|>",
5
+ "<|extra_4|>",
6
+ "<|extra_5|>",
7
+ "<|eos|>"
8
+ ],
9
+ "bos_token": "<|startoftext|>",
10
+ "eos_token": "<|endoftext|>",
11
+ "pad_token": "<|pad|>"
12
+ }
test.py ADDED
@@ -0,0 +1,49 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (C) 2024 THL A29 Limited, a Tencent company. All rights reserved.
2
+ #
3
+ # Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # https://github.com/Tencent/Tencent-Hunyuan-Large/blob/main/License.docx
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+
15
+ from tokenizers import ByteLevelBPETokenizer
16
+ from transformers import AutoTokenizer
17
+
18
+ # Step 1: Initialize ByteLevelBPETokenizer
19
+ #tokenizer = ByteLevelBPETokenizer(
20
+ # "vocab.json",
21
+ # "merges.txt"
22
+ #)
23
+
24
+ # Step 2: Save the tokenizer configuration
25
+ #tokenizer.save_model("auto_model")
26
+
27
+ # Step 3: Load the tokenizer using AutoTokenizer
28
+ auto_tokenizer = AutoTokenizer.from_pretrained("./", use_fast=False, trust_remote_code=True)
29
+
30
+ # Test the tokenizer
31
+ text = "Hello, world!"
32
+ encoded = auto_tokenizer.encode(text)
33
+ decoded = auto_tokenizer.decode(encoded)
34
+
35
+ print("Encoded:", encoded)
36
+ print("Decoded:", decoded)
37
+
38
+ messages = [
39
+ {"role": "system", "content": "You are a helpful assistant."},
40
+ {"role": "user", "content": "Hello, how are you?"},
41
+ {"role": "assistant", "content": "I'm good, thank you! How can I help you today?"},
42
+ {"role": "user", "content": "Nothing"},
43
+ ]
44
+
45
+ print('messages:', messages)
46
+ ids = auto_tokenizer.apply_chat_template(messages)
47
+ print(f"input_ids:\t{ids}")
48
+ text = auto_tokenizer.decode(ids)
49
+ print(f"input_text:\t[{text}]")
test4consistent.py ADDED
@@ -0,0 +1,31 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (C) 2024 THL A29 Limited, a Tencent company. All rights reserved.
2
+ #
3
+ # Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # https://github.com/Tencent/Tencent-Hunyuan-Large/blob/main/License.docx
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+ # test tokenizer encode & decode consistency
15
+ from transformers import AutoTokenizer
16
+ tokenizer = AutoTokenizer.from_pretrained('/tokenizer_exp/other_tokenizer_vocab/hy', local_files_only=True, trust_remote_code=True)
17
+
18
+ test_data = [line.strip() for line in open('/tokenizer_exp/data/test.txt', 'r').readlines()]
19
+
20
+ num_origi_len = 0
21
+ num_token_len = 0
22
+
23
+ for d in test_data:
24
+ a = tokenizer.encode(d)
25
+ num_origi_len += len(d)
26
+ num_token_len += len(a)
27
+ b = tokenizer.decode(a)
28
+ assert b == d, f"encode & decode not consistent: {d} vs {b}"
29
+
30
+ print(f" original length: {num_origi_len}")
31
+ print(f" token length: {num_token_len}")
tokenization_hy.py ADDED
@@ -0,0 +1,354 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (C) 2024 THL A29 Limited, a Tencent company. All rights reserved.
2
+ #
3
+ # Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # https://github.com/Tencent/Tencent-Hunyuan-Large/blob/main/License.docx
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+
15
+ import os
16
+ import base64
17
+ import logging
18
+ import tiktoken
19
+ import unicodedata
20
+ from transformers import PreTrainedTokenizer, AddedToken
21
+ from typing import Collection, Dict, List, Set, Tuple, Union
22
+
23
+
24
+ logger = logging.getLogger(__name__)
25
+
26
+
27
+ VOCAB_FILES_NAMES = {"vocab_file": "hy.tiktoken"}
28
+ PAT_STR = r"""(?i:'s|'t|'re|'ve|'m|'ll|'d)|""" \
29
+ r"""[^\r\n\p{L}\p{N}]?\p{L}+|""" \
30
+ r"""\p{N}|""" \
31
+ r""" ?[^\s\p{L}\p{N}]+[\r\n]*|""" \
32
+ r"""\s*[\r\n]+|""" \
33
+ r"""\s+(?!\S)|""" \
34
+ r"""\s+"""
35
+ # default eod_token and bod_token of our base model
36
+ ENDOFTEXT = "<|endoftext|>"
37
+ STARTOFTEXT = "<|startoftext|>"
38
+
39
+ # extra flag token for other training
40
+ BOSTOKEN = "<|bos|>"
41
+ EOSTOKEN = "<|eos|>"
42
+
43
+ PADTOKEN = "<|pad|>"
44
+
45
+ # extra special tokens for the tokenizer
46
+ # as the default behavior is changed to allow special tokens in
47
+ # regular texts, the surface forms of special tokens need to be
48
+ # as different as possible to minimize the impact
49
+ EXTRAS = tuple((f"<|extra_{i}|>" for i in range(204)))
50
+
51
+ SPECIAL_START_ID = 127957
52
+
53
+
54
+ def _load_tiktoken_bpe(tiktoken_bpe_file: str) -> Dict[bytes, int]:
55
+ dic = {}
56
+ rank = 0
57
+ for i, line in enumerate(open(tiktoken_bpe_file, "rb")):
58
+ if line:
59
+ token, _ = line.split()
60
+ # skip duplicated tokens, this should not happen
61
+ if base64.b64decode(token) in dic:
62
+ raise ValueError(f"!ERROR: duplicated token {token} in your vocab file")
63
+ dic[base64.b64decode(token)] = int(rank)
64
+ rank += 1
65
+ return dic
66
+
67
+
68
+ # special tokens for pretrain and finetune models
69
+ SPECIAL_TOKENS = tuple(
70
+ enumerate(
71
+ (
72
+ (
73
+ ENDOFTEXT,
74
+ STARTOFTEXT,
75
+ BOSTOKEN,
76
+ EOSTOKEN,
77
+ PADTOKEN,
78
+ )
79
+ + EXTRAS
80
+ ),
81
+ start=SPECIAL_START_ID,
82
+ )
83
+ )
84
+
85
+ SPECIAL_TOKENS_SET = set(t for i, t in SPECIAL_TOKENS)
86
+
87
+
88
+ class HYTokenizer(PreTrainedTokenizer):
89
+ """
90
+ HunYuan Tokenizer Initialization. We extend `tiktoken` vocab and
91
+ the default EOD & BOD special tokens are used for base model.
92
+
93
+ Args:
94
+ vocab_file (`str`):
95
+ Path to the vocabulary file.
96
+
97
+ errors (`str`):
98
+ How to handle errors in decoding UTF-8 byte sequences.
99
+ use ignore if you are in streaming inference
100
+
101
+ bod_token (`str` or `tokenizers.AddedToken`, *optional*, defaults to `""<|startoftext|>""`):
102
+ The beginning of document token that was used for training. can be modified by your task.
103
+ default to be `<|startoftext|>` for released base model.
104
+
105
+ eod_token (`str` or `tokenizers.AddedToken`, *optional*, defaults to `""<|endoftext|>""`):
106
+ The end of document token that was used for training. can be modified by your task.
107
+ default to be `<|endoftext|>` for released base model.
108
+
109
+ bos_token (`str` or `tokenizers.AddedToken`, *optional*, defaults to `None`):
110
+ The start or sep special token that was used for some training tasks.
111
+ default to be `<|startoftext|>` for released base model.
112
+ It can be set to `<|bos|>` when you training for some other task
113
+
114
+ eos_token (`str` or `tokenizers.AddedToken`, *optional*, defaults to `None`):
115
+ The end or sep special token that was used for some training tasks.
116
+ default to be `<|endoftext|>` for released base model.
117
+ It can be set to `<|eos|>` when you training for some other task
118
+
119
+ pad_token (`str` or `tokenizers.AddedToken`, *optional*):
120
+ A special token used to make arrays of tokens the same size for batching purpose. Will then be ignored by
121
+ attention mechanisms or loss computation.
122
+
123
+ special_vocab_file (str, *optional*):
124
+ Customed special extra vocab file, same format with hy.tiktoken.
125
+ **Be careful** to use the extra special vocab, it will may cause the model loading collapse.
126
+ The data line be like:
127
+ `PHxhYmN8Pg== 0`
128
+ the id followed `base64.encode(str)` is unused, we will reset them in case of collision
129
+
130
+ add_bod_token (`bool`, *optional*, defaults to `True`):
131
+ Whether or not to add an `bos_token` at the start of documents.
132
+ This will effect `build_inputs_with_special_tokens` method
133
+
134
+ add_eod_token (`bool`, *optional*, defaults to `False`):
135
+ Whether or not to add an `eos_token` at the end of documents.
136
+ This will effect `build_inputs_with_special_tokens` method
137
+
138
+ """
139
+ vocab_files_names = VOCAB_FILES_NAMES
140
+
141
+ def __init__(
142
+ self,
143
+ vocab_file,
144
+ errors="replace",
145
+ bod_token="<|startoftext|>",
146
+ eod_token="<|endoftext|>",
147
+ bos_token="<|startoftext|>",
148
+ eos_token="<|endoftext|>",
149
+ pad_token="<|pad|>",
150
+ add_bod_token=True,
151
+ add_eod_token=True,
152
+ **kwargs,
153
+ ):
154
+ super().__init__(**kwargs)
155
+
156
+ self.errors = errors
157
+ self.mergeable_ranks = _load_tiktoken_bpe(vocab_file) # type: Dict[bytes, int]
158
+ self.special_tokens = {
159
+ token: index
160
+ for index, token in SPECIAL_TOKENS
161
+ }
162
+
163
+ enc = tiktoken.Encoding(
164
+ "HunYuan",
165
+ pat_str=PAT_STR,
166
+ mergeable_ranks=self.mergeable_ranks,
167
+ special_tokens=self.special_tokens,
168
+ )
169
+ assert (
170
+ len(self.mergeable_ranks) + len(self.special_tokens) == enc.n_vocab
171
+ ), f"{len(self.mergeable_ranks)} + {len(self.special_tokens)} != {enc.n_vocab} in encoding"
172
+
173
+ self.decoder = {
174
+ v: k for k, v in self.mergeable_ranks.items()
175
+ } # type: dict[int, bytes|str]
176
+ self.decoder.update({v: k for k, v in self.special_tokens.items()})
177
+
178
+ self.tokenizer = enc
179
+
180
+ self.bod_token, self.bod_id = bod_token, self.special_tokens[bod_token]
181
+ self.eod_token, self.eod_id = eod_token, self.special_tokens[eod_token]
182
+ self.bos_token, self.bos_id = bos_token, self.special_tokens[bos_token]
183
+ self.eos_token, self.eos_id = eos_token, self.special_tokens[eos_token]
184
+ self.pad_token, self.pad_id = pad_token, self.special_tokens[pad_token]
185
+
186
+ self._num_special_token = len(self.special_tokens)
187
+
188
+ self.add_bod_token = add_bod_token
189
+ self.add_eod_token = add_eod_token
190
+
191
+ def __getstate__(self):
192
+ state = self.__dict__.copy()
193
+ del state["tokenizer"]
194
+ return state
195
+
196
+ def __setstate__(self, state):
197
+ self.__dict__.update(state)
198
+ enc = tiktoken.Encoding(
199
+ "HunYuan",
200
+ pat_str=PAT_STR,
201
+ mergeable_ranks=self.mergeable_ranks,
202
+ special_tokens=self.special_tokens,
203
+ )
204
+ self.tokenizer = enc
205
+
206
+ def __len__(self) -> int:
207
+ return self.tokenizer.n_vocab
208
+
209
+ def get_vocab(self) -> Dict[bytes, int]:
210
+ """Return the vocabulary as a dictionary, without special tokens."""
211
+ return self.mergeable_ranks
212
+
213
+ def convert_tokens_to_ids(
214
+ self, tokens: Union[bytes, str, List[Union[bytes, str]]]
215
+ ) -> List[int]:
216
+ ids = []
217
+ if isinstance(tokens, (str, bytes)):
218
+ if tokens in self.special_tokens:
219
+ return self.special_tokens[tokens]
220
+ else:
221
+ return self.mergeable_ranks.get(tokens)
222
+ for token in tokens:
223
+ if token in self.special_tokens:
224
+ ids.append(self.special_tokens[token])
225
+ else:
226
+ ids.append(self.mergeable_ranks.get(token))
227
+ return ids
228
+
229
+ def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None):
230
+ bod_token_id = [self.bod_id] if self.add_bod_token else []
231
+ eod_token_id = [self.eod_id] if self.add_eod_token else []
232
+ output = bod_token_id + token_ids_0 + eod_token_id
233
+ if token_ids_1 is not None:
234
+ output = output + bod_token_id + token_ids_1 + eod_token_id
235
+ return output
236
+
237
+ def _add_tokens(
238
+ self,
239
+ new_tokens: Union[List[str], List[AddedToken]],
240
+ special_tokens: bool = False,
241
+ ) -> List[Tuple[int, str]]:
242
+ """do not support adding tokens"""
243
+ if not special_tokens and new_tokens:
244
+ raise ValueError("Adding regular tokens is not supported")
245
+ for token in new_tokens:
246
+ surface_form = token.content if isinstance(token, AddedToken) else token
247
+ if surface_form not in SPECIAL_TOKENS_SET:
248
+ raise ValueError("Adding unknown special tokens is not supported")
249
+ return 0
250
+
251
+ def save_vocabulary(self, save_directory: str, **kwargs) -> Tuple[str]:
252
+ """
253
+ Save only the vocabulary of the tokenizer (vocabulary).
254
+ Returns:
255
+ `Tuple(str)`: Paths to the files saved.
256
+ """
257
+ file_path = os.path.join(save_directory, "hy.tiktoken")
258
+ with open(file_path, "w", encoding="utf8") as w:
259
+ for k, v in self.mergeable_ranks.items():
260
+ line = base64.b64encode(k).decode("utf8") + " " + str(v) + "\n"
261
+ w.write(line)
262
+ return (file_path,)
263
+
264
+ def tokenize(
265
+ self,
266
+ text: str,
267
+ allowed_special: Union[Set, str] = "all",
268
+ disallowed_special: Union[Collection, str] = (),
269
+ **kwargs,
270
+ ) -> List[Union[bytes, str]]:
271
+ """
272
+ Converts a string in a sequence of tokens.
273
+ Args:
274
+ text (`str`):
275
+ The sequence to be encoded.
276
+ allowed_special (`Literal["all"]` or `set`):
277
+ The surface forms of the tokens to be encoded as special tokens in regular texts.
278
+ Default to "all".
279
+ disallowed_special (`Literal["all"]` or `Collection`):
280
+ The surface forms of the tokens that should not be in regular texts and trigger errors.
281
+ Default to an empty tuple.
282
+ kwargs (additional keyword arguments, *optional*):
283
+ Will be passed to the underlying model specific encode method.
284
+ Returns:
285
+ `List[bytes|str]`: The list of tokens.
286
+ """
287
+ tokens = []
288
+ text = unicodedata.normalize("NFC", text)
289
+
290
+ # this implementation takes a detour: text -> token id -> token surface forms
291
+ for t in self.tokenizer.encode(
292
+ text, allowed_special=allowed_special, disallowed_special=disallowed_special
293
+ ):
294
+ tokens.append(self.decoder[t])
295
+ return tokens
296
+
297
+ def convert_tokens_to_string(self, tokens: List[Union[bytes, str]]) -> str:
298
+ """
299
+ Converts a sequence of tokens in a single string.
300
+ """
301
+ text = ""
302
+ temp = b""
303
+ for t in tokens:
304
+ if isinstance(t, str):
305
+ if temp:
306
+ text += temp.decode("utf-8", errors=self.errors)
307
+ temp = b""
308
+ text += t
309
+ elif isinstance(t, bytes):
310
+ temp += t
311
+ else:
312
+ raise TypeError("token should only be of type types or str")
313
+ if temp:
314
+ text += temp.decode("utf-8", errors=self.errors)
315
+ return text
316
+
317
+ @property
318
+ def vocab_size(self):
319
+ return self.tokenizer.n_vocab
320
+
321
+ def _convert_id_to_token(self, index: int) -> Union[bytes, str]:
322
+ """Converts an id to a token, special tokens included"""
323
+ if index in self.decoder:
324
+ return self.decoder[index]
325
+ raise ValueError("unknown ids")
326
+
327
+ def _convert_token_to_id(self, token: Union[bytes, str]) -> int:
328
+ """Converts a token to an id using the vocab, special tokens included"""
329
+ if token in self.special_tokens:
330
+ return self.special_tokens[token]
331
+ if token in self.mergeable_ranks:
332
+ return self.mergeable_ranks[token]
333
+ raise ValueError("unknown token")
334
+
335
+ def _tokenize(self, text: str, **kwargs):
336
+ """
337
+ Converts a string in a sequence of tokens (string), using the tokenizer. Split in words for word-based
338
+ vocabulary or sub-words for sub-word-based vocabularies (BPE/SentencePieces/WordPieces).
339
+ Do NOT take care of added tokens.
340
+ """
341
+ raise NotImplementedError
342
+
343
+ def _decode(
344
+ self,
345
+ token_ids: Union[int, List[int]],
346
+ skip_special_tokens: bool = False,
347
+ errors: str = None,
348
+ **kwargs,
349
+ ) -> str:
350
+ if isinstance(token_ids, int):
351
+ token_ids = [token_ids]
352
+ if skip_special_tokens:
353
+ token_ids = [i for i in token_ids if i < self.eod_id]
354
+ return self.tokenizer.decode(token_ids, errors=errors or self.errors)
tokenizer_config.json ADDED
@@ -0,0 +1,23 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "added_tokens_decoder": {},
3
+ "additional_special_tokens": [
4
+ "<|startoftext|>",
5
+ "<|extra_0|>",
6
+ "<|extra_4|>",
7
+ "<|extra_5|>",
8
+ "<|eos|>"
9
+ ],
10
+ "auto_map": {
11
+ "AutoTokenizer": [
12
+ "tokenization_hy.HYTokenizer",
13
+ null
14
+ ]
15
+ },
16
+ "bos_token": "<|startoftext|>",
17
+ "chat_template": "{% set context = {'has_head': true} %}{% set loop_messages = messages %}{% for message in loop_messages %}{% set content = message['content'] %}{% if loop.index0 == 0 %}{% if content == '' %}{% set _ = context.update({'has_head': false}) %}{% else %}{% set content = '<|startoftext|>' + content + '<|extra_4|>' %}{% endif %}{% endif %}{% if message['role'] == 'user' %}{% if loop.index0 == 1 and not context.has_head %}{% set content = '<|startoftext|>' + content %}{% endif %}{% if loop.index0 == 1 and context.has_head %}{% set content = content + '<|extra_0|>' %}{% else %}{% set content = '<|startoftext|>' + content + '<|extra_0|>' %}{% endif %}{% elif message['role'] == 'assistant' %}{% set content = content + '<|eos|>' %}{% endif %}{{ content }}{% endfor %}",
18
+ "clean_up_tokenization_spaces": false,
19
+ "eos_token": "<|endoftext|>",
20
+ "model_max_length": 1048576,
21
+ "pad_token": "<|pad|>",
22
+ "tokenizer_class": "HYTokenizer"
23
+ }