pharaouk Crystalcareai commited on
Commit
b0e4fff
0 Parent(s):

Duplicate from Crystalcareai/Quiet-Star-Custom

Browse files

Co-authored-by: Lucas Atkins <Crystalcareai@users.noreply.huggingface.co>

.gitattributes ADDED
@@ -0,0 +1,35 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ *.7z filter=lfs diff=lfs merge=lfs -text
2
+ *.arrow filter=lfs diff=lfs merge=lfs -text
3
+ *.bin filter=lfs diff=lfs merge=lfs -text
4
+ *.bz2 filter=lfs diff=lfs merge=lfs -text
5
+ *.ckpt filter=lfs diff=lfs merge=lfs -text
6
+ *.ftz filter=lfs diff=lfs merge=lfs -text
7
+ *.gz filter=lfs diff=lfs merge=lfs -text
8
+ *.h5 filter=lfs diff=lfs merge=lfs -text
9
+ *.joblib filter=lfs diff=lfs merge=lfs -text
10
+ *.lfs.* filter=lfs diff=lfs merge=lfs -text
11
+ *.mlmodel filter=lfs diff=lfs merge=lfs -text
12
+ *.model filter=lfs diff=lfs merge=lfs -text
13
+ *.msgpack filter=lfs diff=lfs merge=lfs -text
14
+ *.npy filter=lfs diff=lfs merge=lfs -text
15
+ *.npz filter=lfs diff=lfs merge=lfs -text
16
+ *.onnx filter=lfs diff=lfs merge=lfs -text
17
+ *.ot filter=lfs diff=lfs merge=lfs -text
18
+ *.parquet filter=lfs diff=lfs merge=lfs -text
19
+ *.pb filter=lfs diff=lfs merge=lfs -text
20
+ *.pickle filter=lfs diff=lfs merge=lfs -text
21
+ *.pkl filter=lfs diff=lfs merge=lfs -text
22
+ *.pt filter=lfs diff=lfs merge=lfs -text
23
+ *.pth filter=lfs diff=lfs merge=lfs -text
24
+ *.rar filter=lfs diff=lfs merge=lfs -text
25
+ *.safetensors filter=lfs diff=lfs merge=lfs -text
26
+ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
27
+ *.tar.* filter=lfs diff=lfs merge=lfs -text
28
+ *.tar filter=lfs diff=lfs merge=lfs -text
29
+ *.tflite filter=lfs diff=lfs merge=lfs -text
30
+ *.tgz filter=lfs diff=lfs merge=lfs -text
31
+ *.wasm filter=lfs diff=lfs merge=lfs -text
32
+ *.xz filter=lfs diff=lfs merge=lfs -text
33
+ *.zip filter=lfs diff=lfs merge=lfs -text
34
+ *.zst filter=lfs diff=lfs merge=lfs -text
35
+ *tfevents* filter=lfs diff=lfs merge=lfs -text
README.md ADDED
@@ -0,0 +1,6 @@
 
 
 
 
 
 
 
1
+ ---
2
+ datasets:
3
+ - open-web-math/open-web-math
4
+ ---
5
+
6
+ Mistral-7b with continued pretraining using Quiet-STaR (https://arxiv.org/abs/2403.09629) for generating 8 thought tokens before each output token.
added_tokens.json ADDED
@@ -0,0 +1,4 @@
 
 
 
 
 
1
+ {
2
+ "<|endthought|>": 32000,
3
+ "<|startthought|>": 32001
4
+ }
config.json ADDED
@@ -0,0 +1,41 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "_name_or_path": "Crystalcareai/Quiet-Star-Custom",
3
+ "architectures": [
4
+ "QuietForCausalLM"
5
+ ],
6
+ "attention_dropout": 0.0,
7
+ "auto_map": {
8
+ "AutoConfig": "Crystalcareai/Quiet-Star-Custom--configuration_quiet.QuietConfig",
9
+ "AutoModel": "Crystalcareai/Quiet-Star-Custom--modeling_quiet.QuietModel",
10
+ "AutoModelForCausalLM": "Crystalcareai/Quiet-Star-Custom--modeling_quiet.QuietForCausalLM"
11
+ },
12
+ "bos_token_id": 1,
13
+ "eos_token_id": 2,
14
+ "hidden_act": "silu",
15
+ "hidden_size": 4096,
16
+ "initializer_range": 0.02,
17
+ "intermediate_size": 14336,
18
+ "max_position_embeddings": 32768,
19
+ "max_thoughts": 10,
20
+ "merged_lm_and_talk_heads": false,
21
+ "merged_lm_and_think_heads": true,
22
+ "merged_talk_heads": true,
23
+ "model_type": "quiet",
24
+ "num_attention_heads": 32,
25
+ "num_hidden_layers": 32,
26
+ "num_key_value_heads": 8,
27
+ "rms_norm_eps": 1e-05,
28
+ "rope_theta": 10000.0,
29
+ "sliding_window": 4096,
30
+ "tie_word_embeddings": false,
31
+ "torch_dtype": "bfloat16",
32
+ "transformers_version": "4.37.0.dev0",
33
+ "use_cache": true,
34
+ "use_complex_talk_head": true,
35
+ "use_complex_think_head": false,
36
+ "use_concat_talk_head": true,
37
+ "use_shallow_talk": false,
38
+ "use_shallow_think": true,
39
+ "use_weighted_talk_head": true,
40
+ "vocab_size": 32002
41
+ }
configuration_quiet.py ADDED
@@ -0,0 +1,172 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # coding=utf-8
2
+ # Copyright 2023 Quiet AI and the HuggingFace Inc. team. All rights reserved.
3
+ #
4
+ # Licensed under the Apache License, Version 2.0 (the "License");
5
+ # you may not use this file except in compliance with the License.
6
+ # You may obtain a copy of the License at
7
+ #
8
+ # http://www.apache.org/licenses/LICENSE-2.0
9
+ #
10
+ # Unless required by applicable law or agreed to in writing, software
11
+ # distributed under the License is distributed on an "AS IS" BASIS,
12
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
+ # See the License for the specific language governing permissions and
14
+ # limitations under the License.
15
+ """ Quiet model configuration"""
16
+
17
+ from transformers.configuration_utils import PretrainedConfig
18
+ from transformers.utils import logging
19
+
20
+
21
+ logger = logging.get_logger(__name__)
22
+
23
+ QUIET_PRETRAINED_CONFIG_ARCHIVE_MAP = {
24
+ "quietai/Quiet-7B-v0.1": "https://huggingface.co/quietai/Quiet-7B-v0.1/resolve/main/config.json",
25
+ "quietai/Quiet-7B-Instruct-v0.1": "https://huggingface.co/quietai/Quiet-7B-Instruct-v0.1/resolve/main/config.json",
26
+ }
27
+
28
+
29
+ class QuietConfig(PretrainedConfig):
30
+ r"""
31
+ This is the configuration class to store the configuration of a [`QuietModel`]. It is used to instantiate an
32
+ Quiet model according to the specified arguments, defining the model architecture. Instantiating a configuration
33
+ with the defaults will yield a similar configuration to that of the Quiet-7B-v0.1 or Quiet-7B-Instruct-v0.1.
34
+ [quietai/Quiet-7B-v0.1](https://huggingface.co/quietai/Quiet-7B-v0.1)
35
+ [quietai/Quiet-7B-Instruct-v0.1](https://huggingface.co/quietai/Quiet-7B-Instruct-v0.1)
36
+ Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
37
+ documentation from [`PretrainedConfig`] for more information.
38
+ Args:
39
+ vocab_size (`int`, *optional*, defaults to 32000):
40
+ Vocabulary size of the Quiet model. Defines the number of different tokens that can be represented by the
41
+ `inputs_ids` passed when calling [`QuietModel`]
42
+ hidden_size (`int`, *optional*, defaults to 4096):
43
+ Dimension of the hidden representations.
44
+ intermediate_size (`int`, *optional*, defaults to 14336):
45
+ Dimension of the MLP representations.
46
+ num_hidden_layers (`int`, *optional*, defaults to 32):
47
+ Number of hidden layers in the Transformer encoder.
48
+ num_attention_heads (`int`, *optional*, defaults to 32):
49
+ Number of attention heads for each attention layer in the Transformer encoder.
50
+ num_key_value_heads (`int`, *optional*, defaults to 8):
51
+ This is the number of key_value heads that should be used to implement Grouped Query Attention. If
52
+ `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
53
+ `num_key_value_heads=1 the model will use Multi Query Attention (MQA) otherwise GQA is used. When
54
+ converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed
55
+ by meanpooling all the original heads within that group. For more details checkout [this
56
+ paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to `8`.
57
+ hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
58
+ The non-linear activation function (function or string) in the decoder.
59
+ max_position_embeddings (`int`, *optional*, defaults to `4096*32`):
60
+ The maximum sequence length that this model might ever be used with. Quiet's sliding window attention
61
+ allows sequence of up to 4096*32 tokens.
62
+ initializer_range (`float`, *optional*, defaults to 0.02):
63
+ The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
64
+ rms_norm_eps (`float`, *optional*, defaults to 1e-06):
65
+ The epsilon used by the rms normalization layers.
66
+ use_cache (`bool`, *optional*, defaults to `True`):
67
+ Whether or not the model should return the last key/values attentions (not used by all models). Only
68
+ relevant if `config.is_decoder=True`.
69
+ pad_token_id (`int`, *optional*):
70
+ The id of the padding token.
71
+ bos_token_id (`int`, *optional*, defaults to 1):
72
+ The id of the "beginning-of-sequence" token.
73
+ eos_token_id (`int`, *optional*, defaults to 2):
74
+ The id of the "end-of-sequence" token.
75
+ tie_word_embeddings (`bool`, *optional*, defaults to `False`):
76
+ Whether the model's input and output word embeddings should be tied.
77
+ rope_theta (`float`, *optional*, defaults to 10000.0):
78
+ The base period of the RoPE embeddings.
79
+ sliding_window (`int`, *optional*, defaults to 4096):
80
+ Sliding window attention window size. If not specified, will default to `4096`.
81
+ attention_dropout (`float`, *optional*, defaults to 0.0):
82
+ The dropout ratio for the attention probabilities.
83
+ ```python
84
+ >>> from transformers import QuietModel, QuietConfig
85
+ >>> # Initializing a Quiet 7B style configuration
86
+ >>> configuration = QuietConfig()
87
+ >>> # Initializing a model from the Quiet 7B style configuration
88
+ >>> model = QuietModel(configuration)
89
+ >>> # Accessing the model configuration
90
+ >>> configuration = model.config
91
+ ```"""
92
+
93
+ model_type = "quiet"
94
+ keys_to_ignore_at_inference = ["past_key_values"]
95
+
96
+ def __init__(
97
+ self,
98
+ vocab_size=32000,
99
+ hidden_size=4096,
100
+ intermediate_size=14336,
101
+ num_hidden_layers=32,
102
+ num_attention_heads=32,
103
+ num_key_value_heads=8,
104
+ hidden_act="silu",
105
+ max_position_embeddings=4096 * 32,
106
+ initializer_range=0.02,
107
+ rms_norm_eps=1e-6,
108
+ use_cache=True,
109
+ pad_token_id=None,
110
+ bos_token_id=1,
111
+ eos_token_id=2,
112
+ tie_word_embeddings=False,
113
+ rope_theta=10000.0,
114
+ complexity_factor = 0.5,
115
+ sliding_window=4096,
116
+ attention_dropout=0.0,
117
+ max_thoughts=16,
118
+ max_time=None,
119
+ max_temperature=10,
120
+ merged_talk_heads=True,
121
+ merged_lm_and_talk_heads=False,
122
+ merged_lm_and_think_heads=True,
123
+ use_concat_talk_head=True,
124
+ use_shallow_think=True,
125
+ use_shallow_talk=False,
126
+ use_complex_think_head=False,
127
+ use_complex_talk_head=True,
128
+ use_weighted_talk_head=True,
129
+ hidden_dropout_prob=0.0,
130
+ **kwargs,
131
+ ):
132
+ self.vocab_size = vocab_size
133
+ self.max_position_embeddings = max_position_embeddings
134
+ self.hidden_size = hidden_size
135
+ self.intermediate_size = intermediate_size
136
+ self.num_hidden_layers = num_hidden_layers
137
+ self.num_attention_heads = num_attention_heads
138
+ self.sliding_window = sliding_window
139
+
140
+ # for backward compatibility
141
+ if num_key_value_heads is None:
142
+ num_key_value_heads = num_attention_heads
143
+
144
+ self.num_key_value_heads = num_key_value_heads
145
+ self.hidden_act = hidden_act
146
+ self.initializer_range = initializer_range
147
+ self.rms_norm_eps = rms_norm_eps
148
+ self.use_cache = use_cache
149
+ self.rope_theta = rope_theta
150
+ self.attention_dropout = attention_dropout
151
+ self.max_thoughts = max_thoughts
152
+ self.max_time = max_time
153
+ self.complexity_factor = complexity_factor
154
+ self.max_temperature = max_temperature
155
+ self.merged_talk_heads = merged_talk_heads
156
+ self.merged_lm_and_talk_heads = merged_lm_and_talk_heads
157
+ self.merged_lm_and_think_heads = merged_lm_and_think_heads
158
+ self.use_concat_talk_head = use_concat_talk_head
159
+ self.use_shallow_think = use_shallow_think
160
+ self.use_shallow_talk = use_shallow_talk
161
+ self.use_complex_think_head = use_complex_think_head
162
+ self.use_complex_talk_head = use_complex_talk_head
163
+ self.use_weighted_talk_head = use_weighted_talk_head
164
+ self.hidden_dropout_prob = hidden_dropout_prob
165
+
166
+ super().__init__(
167
+ pad_token_id=pad_token_id,
168
+ bos_token_id=bos_token_id,
169
+ eos_token_id=eos_token_id,
170
+ tie_word_embeddings=tie_word_embeddings,
171
+ **kwargs,
172
+ )
generate.py ADDED
@@ -0,0 +1,210 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ from transformers.generation.utils import (
3
+ GenerationMixin,
4
+ validate_stopping_criteria,
5
+ StoppingCriteriaList,
6
+ )
7
+ from transformers import TextStreamer
8
+
9
+
10
+ def custom_generate(
11
+ self,
12
+ input_ids,
13
+ attention_mask=None,
14
+ max_new_tokens=None,
15
+ min_length=None,
16
+ do_sample=None,
17
+ early_stopping=None,
18
+ num_beams=None,
19
+ temperature=None,
20
+ top_k=None,
21
+ top_p=None,
22
+ repetition_penalty=None,
23
+ bad_words_ids=None,
24
+ bos_token_id=None,
25
+ pad_token_id=None,
26
+ eos_token_id=None,
27
+ streamer=None,
28
+ length_penalty=None,
29
+ no_repeat_ngram_size=None,
30
+ num_return_sequences=None,
31
+ decoder_start_token_id=None,
32
+ use_cache=None,
33
+ num_beam_groups=None,
34
+ diversity_penalty=None,
35
+ prefix_allowed_tokens_fn=None,
36
+ output_attentions=None,
37
+ output_hidden_states=None,
38
+ output_scores=None,
39
+ return_dict_in_generate=None,
40
+ forced_bos_token_id=None,
41
+ forced_eos_token_id=None,
42
+ remove_invalid_values=None,
43
+ synced_gpus=None,
44
+ **kwargs,
45
+ ):
46
+ device = input_ids.device
47
+ with torch.no_grad():
48
+ finished_generating = torch.zeros(len(input_ids), dtype=torch.bool, device=device)
49
+
50
+ if max_new_tokens is None:
51
+ max_new_tokens = 50 # Default value if not specified
52
+ for cur_token_idx in range(max_new_tokens):
53
+ # Sample the next token
54
+ new_ids = self(
55
+ input_ids[~finished_generating],
56
+ attention_mask=attention_mask[~finished_generating] if attention_mask is not None else None,
57
+ **kwargs
58
+ )['logits']
59
+
60
+ # Mask out the start and end thought tokens so we don't accidentally sample them
61
+ new_ids[:, :, self.tokenizer.vocab_size:] = -float("inf")
62
+
63
+ for list_idx, answer_idx in enumerate((~finished_generating).nonzero(as_tuple=True)[0]):
64
+ # Find the index of the last token that is not padding
65
+ base_answer_ids = input_ids[answer_idx]
66
+ new_answer_ids = new_ids[list_idx]
67
+ last_token_idx = (base_answer_ids != self.tokenizer.pad_token_id).nonzero(as_tuple=True)[0].max()
68
+
69
+ new_ids_sampled = torch.multinomial(
70
+ torch.nn.functional.softmax(new_answer_ids[last_token_idx] / temperature, dim=-1), 1)
71
+
72
+ # Assign the new id to the last token
73
+ if last_token_idx + 1 >= len(base_answer_ids):
74
+ # Add padding everywhere
75
+ new_padding = torch.full((len(input_ids), 1), self.tokenizer.pad_token_id, dtype=torch.long,
76
+ device=device)
77
+ input_ids = torch.cat([input_ids, new_padding], dim=-1)
78
+ if attention_mask is not None:
79
+ attention_mask = torch.cat([attention_mask, torch.zeros_like(new_padding)], dim=-1)
80
+
81
+ if attention_mask is not None:
82
+ attention_mask[answer_idx, last_token_idx + 1] = 1
83
+ input_ids[answer_idx, last_token_idx + 1] = new_ids_sampled
84
+
85
+ if new_ids_sampled == self.tokenizer.eos_token_id or new_ids_sampled == self.tokenizer.bos_token_id or new_ids_sampled == self.tokenizer.pad_token_id:
86
+ finished_generating[answer_idx] = 1
87
+
88
+ # Check if the end token is generated
89
+ if new_ids_sampled == self.tokenizer.convert_tokens_to_ids("</s>"):
90
+ finished_generating[answer_idx] = 1
91
+
92
+ if finished_generating.all():
93
+ break
94
+
95
+ if streamer is not None:
96
+ streamer.put(new_ids_sampled)
97
+
98
+ generated_token_ids = input_ids.tolist()
99
+ return generated_token_ids, attention_mask
100
+
101
+
102
+ def generate(
103
+ self,
104
+ input_ids,
105
+ attention_mask=None,
106
+ max_new_tokens=None,
107
+ min_length=None,
108
+ do_sample=None,
109
+ early_stopping=None,
110
+ num_beams=None,
111
+ temperature=1.1,
112
+ streamer=None,
113
+ top_k=None,
114
+ top_p=None,
115
+ repetition_penalty=None,
116
+ bad_words_ids=None,
117
+ bos_token_id=None,
118
+ pad_token_id=None,
119
+ eos_token_id=None,
120
+ length_penalty=None,
121
+ no_repeat_ngram_size=None,
122
+ num_return_sequences=None,
123
+ decoder_start_token_id=None,
124
+ use_cache=None,
125
+ num_beam_groups=None,
126
+ diversity_penalty=None,
127
+ prefix_allowed_tokens_fn=None,
128
+ output_attentions=None,
129
+ output_hidden_states=None,
130
+ output_scores=None,
131
+ return_dict_in_generate=None,
132
+ forced_bos_token_id=None,
133
+ forced_eos_token_id=None,
134
+ remove_invalid_values=None,
135
+ synced_gpus=None,
136
+ n_ahead=12,
137
+ n_ahead_talk=4,
138
+ merged_talk_heads=True,
139
+ merged_lm_and_talk_heads=False,
140
+ merged_lm_and_think_heads=True,
141
+ use_concat_talk_head=True,
142
+ use_shallow_think=True,
143
+ use_shallow_talk=False,
144
+ use_complex_think_head=False,
145
+ use_complex_talk_head=True,
146
+ use_weighted_talk_head=True,
147
+ trust_remote_code=True,
148
+ torch_dtype=torch.bfloat16,
149
+ **model_kwargs,
150
+ ):
151
+ # Set model attributes
152
+ self.max_thoughts = n_ahead + n_ahead_talk + 1
153
+ self.merged_talk_heads = merged_talk_heads
154
+ self.merged_lm_and_talk_heads = merged_lm_and_talk_heads
155
+ self.merged_lm_and_think_heads = merged_lm_and_think_heads
156
+ self.use_concat_talk_head = use_concat_talk_head
157
+ self.use_shallow_think = use_shallow_think
158
+ self.use_shallow_talk = use_shallow_talk
159
+ self.use_complex_think_head = use_complex_think_head
160
+ self.use_complex_talk_head = use_complex_talk_head
161
+ self.use_weighted_talk_head = use_weighted_talk_head
162
+
163
+ # Set model properties
164
+ self.use_end_thought_token = True
165
+ self.use_start_thought_token = True
166
+ self.n_ahead = n_ahead
167
+ self.n_passes = 1
168
+ self.eval_mode = True
169
+ self.first_run = False
170
+ self.rm_initialized = True
171
+ self.original_mode = False
172
+
173
+ generated_token_ids, attention_mask = custom_generate(
174
+ self,
175
+ input_ids=input_ids,
176
+ attention_mask=attention_mask,
177
+ max_new_tokens=max_new_tokens,
178
+ min_length=min_length,
179
+ do_sample=do_sample,
180
+ early_stopping=early_stopping,
181
+ num_beams=num_beams,
182
+ temperature=temperature,
183
+ top_k=top_k,
184
+ top_p=top_p,
185
+ repetition_penalty=repetition_penalty,
186
+ bad_words_ids=bad_words_ids,
187
+ bos_token_id=bos_token_id,
188
+ pad_token_id=pad_token_id,
189
+ eos_token_id=eos_token_id,
190
+ length_penalty=length_penalty,
191
+ no_repeat_ngram_size=no_repeat_ngram_size,
192
+ num_return_sequences=num_return_sequences,
193
+ decoder_start_token_id=decoder_start_token_id,
194
+ use_cache=use_cache,
195
+ num_beam_groups=num_beam_groups,
196
+ diversity_penalty=diversity_penalty,
197
+ prefix_allowed_tokens_fn=prefix_allowed_tokens_fn,
198
+ output_attentions=output_attentions,
199
+ output_hidden_states=output_hidden_states,
200
+ output_scores=output_scores,
201
+ return_dict_in_generate=return_dict_in_generate,
202
+ forced_bos_token_id=forced_bos_token_id,
203
+ forced_eos_token_id=forced_eos_token_id,
204
+ remove_invalid_values=remove_invalid_values,
205
+ synced_gpus=synced_gpus,
206
+ streamer=streamer,
207
+ **model_kwargs,
208
+ )
209
+
210
+ return generated_token_ids, attention_mask
generation_config.json ADDED
@@ -0,0 +1,6 @@
 
 
 
 
 
 
 
1
+ {
2
+ "_from_model_config": true,
3
+ "bos_token_id": 1,
4
+ "eos_token_id": 2,
5
+ "transformers_version": "4.37.0.dev0"
6
+ }
inference.py ADDED
@@ -0,0 +1,121 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer
3
+
4
+ def compute_memory_used_pct(device):
5
+ memory_used = torch.cuda.max_memory_allocated(device) / (1024**3)
6
+ memory_pct = (
7
+ memory_used
8
+ / (torch.cuda.get_device_properties(device).total_memory / (1024**3))
9
+ * 100
10
+ )
11
+ return memory_pct
12
+
13
+ model_path = "./out"
14
+
15
+ n_ahead = 8
16
+ n_ahead_talk = 4
17
+ merged_talk_heads = True
18
+
19
+ # Load the model
20
+ model = AutoModelForCausalLM.from_pretrained(
21
+ model_path,
22
+ max_thoughts=n_ahead + n_ahead_talk + 1,
23
+ merged_talk_heads=merged_talk_heads,
24
+ merged_lm_and_talk_heads=False,
25
+ merged_lm_and_think_heads=True,
26
+ use_concat_talk_head=True,
27
+ use_shallow_think=True,
28
+ use_shallow_talk=False,
29
+ use_complex_think_head=False,
30
+ use_complex_talk_head=True,
31
+ use_weighted_talk_head=True,
32
+ trust_remote_code=True,
33
+ torch_dtype=torch.bfloat16,
34
+ device_map="auto",
35
+ )
36
+
37
+ # Load the tokenizer and assign it to the model instance for compatibility
38
+ tokenizer = AutoTokenizer.from_pretrained(model_path)
39
+ model.tokenizer = tokenizer
40
+
41
+ model.use_end_thought_token = True
42
+ model.use_start_thought_token = True
43
+ model.wandb_enabled = True
44
+ model.n_ahead = n_ahead
45
+ model.n_passes = 2
46
+ model.eval_mode = True
47
+ model.first_run = False
48
+ model.kill_after = 100
49
+ model.rm_initialized = True
50
+ model.original_mode = False
51
+
52
+ # Custom generate function
53
+ def custom_generate(model, input_ids, attention_mask, max_new_tokens, streamer, **kwargs):
54
+ with torch.no_grad():
55
+ finished_generating = torch.zeros(len(input_ids), dtype=torch.bool, device=input_ids.device)
56
+ for cur_token_idx in range(max_new_tokens):
57
+ # Sample the next token
58
+ new_ids = model(
59
+ input_ids[~finished_generating],
60
+ attention_mask=attention_mask[~finished_generating]
61
+ )['logits']
62
+ # Mask out the start and end thought tokens so we don't accidentally sample them
63
+ new_ids[:, :, model.tokenizer.vocab_size:] = -float("inf")
64
+ for list_idx, answer_idx in enumerate((~finished_generating).nonzero(as_tuple=True)[0]):
65
+ # Find the index of the last token that is not padding
66
+ base_answer_ids = input_ids[answer_idx]
67
+ new_answer_ids = new_ids[list_idx]
68
+ last_token_idx = (base_answer_ids != model.tokenizer.pad_token_id).nonzero(as_tuple=True)[0].max()
69
+
70
+ new_ids_sampled = torch.multinomial(
71
+ torch.nn.functional.softmax(new_answer_ids[last_token_idx] / kwargs.get("temperature", 1.0), dim=-1), 1)
72
+ # Assign the new id to the last token
73
+ if last_token_idx + 1 >= len(base_answer_ids):
74
+ # Add padding everywhere
75
+ new_padding = torch.full((len(input_ids), 1), model.tokenizer.pad_token_id, dtype=torch.long,
76
+ device=input_ids.device)
77
+ input_ids = torch.cat([input_ids, new_padding], dim=-1)
78
+ attention_mask = torch.cat([attention_mask, torch.zeros_like(new_padding)], dim=-1)
79
+ attention_mask[answer_idx, last_token_idx + 1] = 1
80
+ input_ids[answer_idx, last_token_idx + 1] = new_ids_sampled
81
+ if new_ids_sampled == model.tokenizer.eos_token_id or new_ids_sampled == model.tokenizer.bos_token_id or new_ids_sampled == model.tokenizer.pad_token_id:
82
+ finished_generating[answer_idx] = 1
83
+ # Check if the end token is generated
84
+ if new_ids_sampled == model.tokenizer.convert_tokens_to_ids("<|/assistant|>"):
85
+ finished_generating[answer_idx] = 1
86
+ if finished_generating.all():
87
+ break
88
+ streamer.put(new_ids_sampled)
89
+ return input_ids, attention_mask
90
+
91
+ # Formulate your prompt
92
+ prompt_template = "[INST] {prompt} [/INST]"
93
+
94
+ prompt = "You're standing on the surface of the Earth. "\
95
+ "You walk one mile south, one mile west and one mile north. "\
96
+ "You end up exactly where you started. Where are you?"
97
+
98
+ # Convert prompt to tokens
99
+ tokens = tokenizer(prompt_template.format(prompt=prompt), return_tensors='pt').input_ids.to(model.device)
100
+
101
+ # Generate an attention mask
102
+ attention_mask = torch.where(tokens != tokenizer.pad_token_id, torch.ones_like(tokens), torch.zeros_like(tokens)).to(model.device)
103
+
104
+ streamer = TextStreamer(tokenizer, skip_prompt=False, skip_special_tokens=True)
105
+
106
+ # Generate output using the custom generate function
107
+ output_ids, _ = custom_generate(
108
+ model,
109
+ input_ids=tokens,
110
+ attention_mask=attention_mask,
111
+ max_new_tokens=512,
112
+ streamer=streamer,
113
+ temperature=0.9,
114
+ )
115
+
116
+ generated_text = ""
117
+
118
+ print() # Print a newline after streaming is complete
119
+
120
+ # Cleanup if necessary
121
+ torch.cuda.empty_cache()
model-00001-of-00003.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:1d0220963e92ec181cd2b99c0421eb03ce9b5dfa938b66df39af1122ebb7ba76
3
+ size 4943211664
model-00002-of-00003.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:871f3846be90a6c4515931c237d4efb109b37b2f4061670e10c744db0c659461
3
+ size 4999819336
model-00003-of-00003.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:690edbd65da074f63493d8baf4c9c93a9ee93ad8bec46c8e360595573fa444e6
3
+ size 4641221080
model.safetensors.index.json ADDED
@@ -0,0 +1,305 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "metadata": {
3
+ "total_size": 14584217600
4
+ },
5
+ "weight_map": {
6
+ "end_embedding": "model-00001-of-00003.safetensors",
7
+ "lm_head.weight": "model-00003-of-00003.safetensors",
8
+ "model.embed_tokens.weight": "model-00001-of-00003.safetensors",
9
+ "model.layers.0.input_layernorm.weight": "model-00001-of-00003.safetensors",
10
+ "model.layers.0.mlp.down_proj.weight": "model-00001-of-00003.safetensors",
11
+ "model.layers.0.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
12
+ "model.layers.0.mlp.up_proj.weight": "model-00001-of-00003.safetensors",
13
+ "model.layers.0.post_attention_layernorm.weight": "model-00001-of-00003.safetensors",
14
+ "model.layers.0.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
15
+ "model.layers.0.self_attn.o_proj.weight": "model-00001-of-00003.safetensors",
16
+ "model.layers.0.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
17
+ "model.layers.0.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
18
+ "model.layers.1.input_layernorm.weight": "model-00001-of-00003.safetensors",
19
+ "model.layers.1.mlp.down_proj.weight": "model-00001-of-00003.safetensors",
20
+ "model.layers.1.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
21
+ "model.layers.1.mlp.up_proj.weight": "model-00001-of-00003.safetensors",
22
+ "model.layers.1.post_attention_layernorm.weight": "model-00001-of-00003.safetensors",
23
+ "model.layers.1.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
24
+ "model.layers.1.self_attn.o_proj.weight": "model-00001-of-00003.safetensors",
25
+ "model.layers.1.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
26
+ "model.layers.1.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
27
+ "model.layers.10.input_layernorm.weight": "model-00002-of-00003.safetensors",
28
+ "model.layers.10.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
29
+ "model.layers.10.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
30
+ "model.layers.10.mlp.up_proj.weight": "model-00001-of-00003.safetensors",
31
+ "model.layers.10.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
32
+ "model.layers.10.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
33
+ "model.layers.10.self_attn.o_proj.weight": "model-00001-of-00003.safetensors",
34
+ "model.layers.10.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
35
+ "model.layers.10.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
36
+ "model.layers.11.input_layernorm.weight": "model-00002-of-00003.safetensors",
37
+ "model.layers.11.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
38
+ "model.layers.11.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
39
+ "model.layers.11.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
40
+ "model.layers.11.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
41
+ "model.layers.11.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
42
+ "model.layers.11.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
43
+ "model.layers.11.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
44
+ "model.layers.11.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
45
+ "model.layers.12.input_layernorm.weight": "model-00002-of-00003.safetensors",
46
+ "model.layers.12.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
47
+ "model.layers.12.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
48
+ "model.layers.12.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
49
+ "model.layers.12.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
50
+ "model.layers.12.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
51
+ "model.layers.12.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
52
+ "model.layers.12.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
53
+ "model.layers.12.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
54
+ "model.layers.13.input_layernorm.weight": "model-00002-of-00003.safetensors",
55
+ "model.layers.13.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
56
+ "model.layers.13.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
57
+ "model.layers.13.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
58
+ "model.layers.13.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
59
+ "model.layers.13.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
60
+ "model.layers.13.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
61
+ "model.layers.13.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
62
+ "model.layers.13.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
63
+ "model.layers.14.input_layernorm.weight": "model-00002-of-00003.safetensors",
64
+ "model.layers.14.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
65
+ "model.layers.14.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
66
+ "model.layers.14.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
67
+ "model.layers.14.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
68
+ "model.layers.14.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
69
+ "model.layers.14.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
70
+ "model.layers.14.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
71
+ "model.layers.14.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
72
+ "model.layers.15.input_layernorm.weight": "model-00002-of-00003.safetensors",
73
+ "model.layers.15.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
74
+ "model.layers.15.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
75
+ "model.layers.15.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
76
+ "model.layers.15.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
77
+ "model.layers.15.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
78
+ "model.layers.15.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
79
+ "model.layers.15.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
80
+ "model.layers.15.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
81
+ "model.layers.16.input_layernorm.weight": "model-00002-of-00003.safetensors",
82
+ "model.layers.16.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
83
+ "model.layers.16.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
84
+ "model.layers.16.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
85
+ "model.layers.16.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
86
+ "model.layers.16.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
87
+ "model.layers.16.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
88
+ "model.layers.16.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
89
+ "model.layers.16.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
90
+ "model.layers.17.input_layernorm.weight": "model-00002-of-00003.safetensors",
91
+ "model.layers.17.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
92
+ "model.layers.17.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
93
+ "model.layers.17.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
94
+ "model.layers.17.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
95
+ "model.layers.17.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
96
+ "model.layers.17.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
97
+ "model.layers.17.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
98
+ "model.layers.17.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
99
+ "model.layers.18.input_layernorm.weight": "model-00002-of-00003.safetensors",
100
+ "model.layers.18.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
101
+ "model.layers.18.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
102
+ "model.layers.18.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
103
+ "model.layers.18.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
104
+ "model.layers.18.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
105
+ "model.layers.18.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
106
+ "model.layers.18.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
107
+ "model.layers.18.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
108
+ "model.layers.19.input_layernorm.weight": "model-00002-of-00003.safetensors",
109
+ "model.layers.19.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
110
+ "model.layers.19.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
111
+ "model.layers.19.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
112
+ "model.layers.19.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
113
+ "model.layers.19.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
114
+ "model.layers.19.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
115
+ "model.layers.19.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
116
+ "model.layers.19.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
117
+ "model.layers.2.input_layernorm.weight": "model-00001-of-00003.safetensors",
118
+ "model.layers.2.mlp.down_proj.weight": "model-00001-of-00003.safetensors",
119
+ "model.layers.2.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
120
+ "model.layers.2.mlp.up_proj.weight": "model-00001-of-00003.safetensors",
121
+ "model.layers.2.post_attention_layernorm.weight": "model-00001-of-00003.safetensors",
122
+ "model.layers.2.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
123
+ "model.layers.2.self_attn.o_proj.weight": "model-00001-of-00003.safetensors",
124
+ "model.layers.2.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
125
+ "model.layers.2.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
126
+ "model.layers.20.input_layernorm.weight": "model-00002-of-00003.safetensors",
127
+ "model.layers.20.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
128
+ "model.layers.20.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
129
+ "model.layers.20.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
130
+ "model.layers.20.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
131
+ "model.layers.20.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
132
+ "model.layers.20.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
133
+ "model.layers.20.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
134
+ "model.layers.20.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
135
+ "model.layers.21.input_layernorm.weight": "model-00002-of-00003.safetensors",
136
+ "model.layers.21.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
137
+ "model.layers.21.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
138
+ "model.layers.21.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
139
+ "model.layers.21.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
140
+ "model.layers.21.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
141
+ "model.layers.21.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
142
+ "model.layers.21.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
143
+ "model.layers.21.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
144
+ "model.layers.22.input_layernorm.weight": "model-00003-of-00003.safetensors",
145
+ "model.layers.22.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
146
+ "model.layers.22.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
147
+ "model.layers.22.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
148
+ "model.layers.22.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
149
+ "model.layers.22.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
150
+ "model.layers.22.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
151
+ "model.layers.22.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
152
+ "model.layers.22.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
153
+ "model.layers.23.input_layernorm.weight": "model-00003-of-00003.safetensors",
154
+ "model.layers.23.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
155
+ "model.layers.23.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
156
+ "model.layers.23.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
157
+ "model.layers.23.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
158
+ "model.layers.23.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
159
+ "model.layers.23.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
160
+ "model.layers.23.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
161
+ "model.layers.23.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
162
+ "model.layers.24.input_layernorm.weight": "model-00003-of-00003.safetensors",
163
+ "model.layers.24.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
164
+ "model.layers.24.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
165
+ "model.layers.24.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
166
+ "model.layers.24.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
167
+ "model.layers.24.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
168
+ "model.layers.24.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
169
+ "model.layers.24.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
170
+ "model.layers.24.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
171
+ "model.layers.25.input_layernorm.weight": "model-00003-of-00003.safetensors",
172
+ "model.layers.25.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
173
+ "model.layers.25.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
174
+ "model.layers.25.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
175
+ "model.layers.25.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
176
+ "model.layers.25.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
177
+ "model.layers.25.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
178
+ "model.layers.25.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
179
+ "model.layers.25.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
180
+ "model.layers.26.input_layernorm.weight": "model-00003-of-00003.safetensors",
181
+ "model.layers.26.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
182
+ "model.layers.26.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
183
+ "model.layers.26.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
184
+ "model.layers.26.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
185
+ "model.layers.26.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
186
+ "model.layers.26.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
187
+ "model.layers.26.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
188
+ "model.layers.26.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
189
+ "model.layers.27.input_layernorm.weight": "model-00003-of-00003.safetensors",
190
+ "model.layers.27.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
191
+ "model.layers.27.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
192
+ "model.layers.27.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
193
+ "model.layers.27.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
194
+ "model.layers.27.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
195
+ "model.layers.27.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
196
+ "model.layers.27.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
197
+ "model.layers.27.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
198
+ "model.layers.28.input_layernorm.weight": "model-00003-of-00003.safetensors",
199
+ "model.layers.28.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
200
+ "model.layers.28.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
201
+ "model.layers.28.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
202
+ "model.layers.28.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
203
+ "model.layers.28.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
204
+ "model.layers.28.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
205
+ "model.layers.28.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
206
+ "model.layers.28.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
207
+ "model.layers.29.input_layernorm.weight": "model-00003-of-00003.safetensors",
208
+ "model.layers.29.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
209
+ "model.layers.29.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
210
+ "model.layers.29.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
211
+ "model.layers.29.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
212
+ "model.layers.29.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
213
+ "model.layers.29.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
214
+ "model.layers.29.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
215
+ "model.layers.29.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
216
+ "model.layers.3.input_layernorm.weight": "model-00001-of-00003.safetensors",
217
+ "model.layers.3.mlp.down_proj.weight": "model-00001-of-00003.safetensors",
218
+ "model.layers.3.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
219
+ "model.layers.3.mlp.up_proj.weight": "model-00001-of-00003.safetensors",
220
+ "model.layers.3.post_attention_layernorm.weight": "model-00001-of-00003.safetensors",
221
+ "model.layers.3.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
222
+ "model.layers.3.self_attn.o_proj.weight": "model-00001-of-00003.safetensors",
223
+ "model.layers.3.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
224
+ "model.layers.3.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
225
+ "model.layers.30.input_layernorm.weight": "model-00003-of-00003.safetensors",
226
+ "model.layers.30.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
227
+ "model.layers.30.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
228
+ "model.layers.30.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
229
+ "model.layers.30.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
230
+ "model.layers.30.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
231
+ "model.layers.30.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
232
+ "model.layers.30.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
233
+ "model.layers.30.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
234
+ "model.layers.31.input_layernorm.weight": "model-00003-of-00003.safetensors",
235
+ "model.layers.31.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
236
+ "model.layers.31.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
237
+ "model.layers.31.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
238
+ "model.layers.31.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
239
+ "model.layers.31.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
240
+ "model.layers.31.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
241
+ "model.layers.31.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
242
+ "model.layers.31.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
243
+ "model.layers.4.input_layernorm.weight": "model-00001-of-00003.safetensors",
244
+ "model.layers.4.mlp.down_proj.weight": "model-00001-of-00003.safetensors",
245
+ "model.layers.4.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
246
+ "model.layers.4.mlp.up_proj.weight": "model-00001-of-00003.safetensors",
247
+ "model.layers.4.post_attention_layernorm.weight": "model-00001-of-00003.safetensors",
248
+ "model.layers.4.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
249
+ "model.layers.4.self_attn.o_proj.weight": "model-00001-of-00003.safetensors",
250
+ "model.layers.4.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
251
+ "model.layers.4.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
252
+ "model.layers.5.input_layernorm.weight": "model-00001-of-00003.safetensors",
253
+ "model.layers.5.mlp.down_proj.weight": "model-00001-of-00003.safetensors",
254
+ "model.layers.5.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
255
+ "model.layers.5.mlp.up_proj.weight": "model-00001-of-00003.safetensors",
256
+ "model.layers.5.post_attention_layernorm.weight": "model-00001-of-00003.safetensors",
257
+ "model.layers.5.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
258
+ "model.layers.5.self_attn.o_proj.weight": "model-00001-of-00003.safetensors",
259
+ "model.layers.5.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
260
+ "model.layers.5.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
261
+ "model.layers.6.input_layernorm.weight": "model-00001-of-00003.safetensors",
262
+ "model.layers.6.mlp.down_proj.weight": "model-00001-of-00003.safetensors",
263
+ "model.layers.6.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
264
+ "model.layers.6.mlp.up_proj.weight": "model-00001-of-00003.safetensors",
265
+ "model.layers.6.post_attention_layernorm.weight": "model-00001-of-00003.safetensors",
266
+ "model.layers.6.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
267
+ "model.layers.6.self_attn.o_proj.weight": "model-00001-of-00003.safetensors",
268
+ "model.layers.6.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
269
+ "model.layers.6.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
270
+ "model.layers.7.input_layernorm.weight": "model-00001-of-00003.safetensors",
271
+ "model.layers.7.mlp.down_proj.weight": "model-00001-of-00003.safetensors",
272
+ "model.layers.7.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
273
+ "model.layers.7.mlp.up_proj.weight": "model-00001-of-00003.safetensors",
274
+ "model.layers.7.post_attention_layernorm.weight": "model-00001-of-00003.safetensors",
275
+ "model.layers.7.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
276
+ "model.layers.7.self_attn.o_proj.weight": "model-00001-of-00003.safetensors",
277
+ "model.layers.7.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
278
+ "model.layers.7.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
279
+ "model.layers.8.input_layernorm.weight": "model-00001-of-00003.safetensors",
280
+ "model.layers.8.mlp.down_proj.weight": "model-00001-of-00003.safetensors",
281
+ "model.layers.8.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
282
+ "model.layers.8.mlp.up_proj.weight": "model-00001-of-00003.safetensors",
283
+ "model.layers.8.post_attention_layernorm.weight": "model-00001-of-00003.safetensors",
284
+ "model.layers.8.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
285
+ "model.layers.8.self_attn.o_proj.weight": "model-00001-of-00003.safetensors",
286
+ "model.layers.8.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
287
+ "model.layers.8.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
288
+ "model.layers.9.input_layernorm.weight": "model-00001-of-00003.safetensors",
289
+ "model.layers.9.mlp.down_proj.weight": "model-00001-of-00003.safetensors",
290
+ "model.layers.9.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
291
+ "model.layers.9.mlp.up_proj.weight": "model-00001-of-00003.safetensors",
292
+ "model.layers.9.post_attention_layernorm.weight": "model-00001-of-00003.safetensors",
293
+ "model.layers.9.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
294
+ "model.layers.9.self_attn.o_proj.weight": "model-00001-of-00003.safetensors",
295
+ "model.layers.9.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
296
+ "model.layers.9.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
297
+ "model.norm.weight": "model-00003-of-00003.safetensors",
298
+ "start_embedding": "model-00001-of-00003.safetensors",
299
+ "talk_head.0.0.bias": "model-00003-of-00003.safetensors",
300
+ "talk_head.0.0.weight": "model-00003-of-00003.safetensors",
301
+ "talk_head.0.2.bias": "model-00003-of-00003.safetensors",
302
+ "talk_head.0.2.weight": "model-00003-of-00003.safetensors",
303
+ "talk_head.0.4.weight": "model-00003-of-00003.safetensors"
304
+ }
305
+ }
modeling_quiet.py ADDED
@@ -0,0 +1,2335 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # coding=utf-8
2
+ # Copyright 2023 Quiet AI and the HuggingFace Inc. team. All rights reserved.
3
+ #
4
+ # This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
5
+ # and OPT implementations in this library. It has been modified from its
6
+ # original forms to accommodate minor architectural differences compared
7
+ # to GPT-NeoX and OPT used by the Meta AI team that trained the model.
8
+ #
9
+ # Licensed under the Apache License, Version 2.0 (the "License");
10
+ # you may not use this file except in compliance with the License.
11
+ # You may obtain a copy of the License at
12
+ #
13
+ # http://www.apache.org/licenses/LICENSE-2.0
14
+ #
15
+ # Unless required by applicable law or agreed to in writing, software
16
+ # distributed under the License is distributed on an "AS IS" BASIS,
17
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
18
+ # See the License for the specific language governing permissions and
19
+ # limitations under the License.
20
+ """ PyTorch Quiet model."""
21
+ import inspect
22
+ import math
23
+ import pdb
24
+ import warnings
25
+ from collections import defaultdict
26
+ from typing import List, Optional, Tuple, Union
27
+
28
+ import torch
29
+ import torch.nn.functional as F
30
+ import torch.utils.checkpoint
31
+ from torch import nn
32
+ from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
33
+ from transformers.generation.utils import GenerationMixin
34
+ from transformers.generation.stopping_criteria import StoppingCriteriaList, validate_stopping_criteria
35
+ from transformers import TextStreamer, AutoTokenizer
36
+
37
+ from transformers.activations import ACT2FN
38
+ from transformers.cache_utils import Cache, DynamicCache
39
+ from transformers.modeling_attn_mask_utils import _prepare_4d_causal_attention_mask
40
+ from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast, SequenceClassifierOutputWithPast
41
+ from transformers.modeling_utils import PreTrainedModel
42
+ from transformers.utils import (
43
+ add_start_docstrings,
44
+ add_start_docstrings_to_model_forward,
45
+ is_flash_attn_2_available,
46
+ is_flash_attn_greater_or_equal_2_10,
47
+ logging,
48
+ replace_return_docstrings,
49
+ )
50
+ from .configuration_quiet import QuietConfig
51
+
52
+ import time
53
+ from typing import Optional, List
54
+
55
+
56
+ if is_flash_attn_2_available():
57
+ from flash_attn import flash_attn_func, flash_attn_varlen_func
58
+ from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input # noqa
59
+
60
+ _flash_supports_window_size = "window_size" in list(inspect.signature(flash_attn_func).parameters)
61
+
62
+
63
+ logger = logging.get_logger(__name__)
64
+
65
+ _CONFIG_FOR_DOC = "QuietConfig"
66
+
67
+
68
+ def _prepare_4d_causal_attention_mask_for_sdpa(attention_mask, input_shape, inputs_embeds, past_key_values_length):
69
+ # Compute the attention mask correctly
70
+ bsz, tgt_len = input_shape
71
+
72
+ # Create a 4D attention mask from a 2D tensor mask.
73
+ # The shape of the output attention mask is (batch_size, 1, tgt_len, src_len)
74
+ # The values are either 0 or 1, where 0 means padding and 1 means non-padding.
75
+ combined_attention_mask = None
76
+ if attention_mask is not None:
77
+ # What if attention_mask is not None and has a shape of (batch_size, 1, tgt_len, src_len)
78
+ # In this case, we can just use it directly.
79
+ if attention_mask.dim() == 4:
80
+ combined_attention_mask = attention_mask
81
+ # What if attention_mask is not None and has a shape of (batch_size, 1, tgt_len)
82
+ # In this case, we need to expand it to (batch_size, 1, tgt_len, src_len)
83
+ elif attention_mask.dim() == 3:
84
+ expanded_attn_mask = attention_mask[:, None, :, :]
85
+ combined_attention_mask = expanded_attn_mask
86
+ # What if attention_mask is not None and has a shape of (batch_size, tgt_len)
87
+ # In this case, we need to expand it to (batch_size, 1, tgt_len, src_len)
88
+ elif attention_mask.dim() == 2:
89
+ # Provided a padding mask of dimensions [batch_size, seq_length]
90
+ # - if the model is a decoder, apply a causal mask in addition to the padding mask
91
+ # - if the model is an encoder, make the mask broadcastable to [batch_size, num_heads, seq_length, seq_length]
92
+ if past_key_values_length > 0:
93
+ attention_mask = attention_mask.to(dtype=torch.long)
94
+ attention_mask = attention_mask[:, past_key_values_length:]
95
+ expanded_attn_mask = attention_mask[:, None, None, :]
96
+ combined_attention_mask = expanded_attn_mask
97
+ else:
98
+ raise ValueError(
99
+ "Wrong shape for input_ids (shape {}) or attention_mask (shape {})".format(
100
+ input_shape, attention_mask.shape
101
+ )
102
+ )
103
+
104
+ # Since attention_mask is 1.0 for positions we want to attend and 0.0 for
105
+ # masked positions, this operation will create a tensor which is 0.0 for
106
+ # positions we want to attend and -10000.0 for masked positions.
107
+ # Since we are adding it to the raw scores before the softmax, this is
108
+ # effectively the same as removing these entirely.
109
+ if combined_attention_mask is not None:
110
+ # Ensure the attention mask values are within a reasonable range
111
+ combined_attention_mask = combined_attention_mask.clamp(min=0, max=1)
112
+
113
+ # Convert the attention mask to bfloat16
114
+ combined_attention_mask = combined_attention_mask.to(torch.bfloat16)
115
+
116
+ # Normalize the attention mask values to be between 0 and 1
117
+ combined_attention_mask = (1.0 - combined_attention_mask) * -10000.0
118
+ else:
119
+ combined_attention_mask = torch.zeros(
120
+ (bsz, 1, tgt_len, tgt_len), dtype=torch.bfloat16, device=inputs_embeds.device
121
+ )
122
+
123
+ return combined_attention_mask
124
+
125
+
126
+ # Copied from transformers.models.llama.modeling_llama._get_unpad_data
127
+ def _get_unpad_data(attention_mask):
128
+ seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
129
+ indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
130
+ max_seqlen_in_batch = seqlens_in_batch.max().item()
131
+ cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.torch.int32), (1, 0))
132
+ return (
133
+ indices,
134
+ cu_seqlens,
135
+ max_seqlen_in_batch,
136
+ )
137
+
138
+
139
+ # Copied from transformers.models.llama.modeling_llama.LlamaRMSNorm with Llama->Quiet
140
+ class QuietRMSNorm(nn.Module):
141
+ def __init__(self, hidden_size, eps=1e-6):
142
+ super().__init__()
143
+ self.weight = nn.Parameter(torch.ones(hidden_size))
144
+ self.variance_epsilon = eps
145
+
146
+
147
+ def forward(self, hidden_states):
148
+ input_dtype = hidden_states.dtype
149
+ hidden_states = hidden_states.to(torch.float32)
150
+ variance = hidden_states.pow(2).mean(-1, keepdim=True)
151
+ hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
152
+ return hidden_states.to(input_dtype) * self.weight.to(hidden_states.device)
153
+
154
+
155
+ # Copied from transformers.models.llama.modeling_llama.LlamaRotaryEmbedding with Llama->Quiet
156
+ class QuietRotaryEmbedding(nn.Module):
157
+ def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):
158
+ super().__init__()
159
+
160
+ self.dim = dim
161
+ self.max_position_embeddings = max_position_embeddings
162
+ self.base = base
163
+ inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))
164
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
165
+
166
+ # Build here to make `torch.jit.trace` work.
167
+ self._set_cos_sin_cache(
168
+ seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype()
169
+ )
170
+
171
+ def _set_cos_sin_cache(self, seq_len, device, dtype):
172
+ self.max_seq_len_cached = seq_len
173
+ t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)
174
+
175
+ freqs = torch.outer(t, self.inv_freq)
176
+ # Different from paper, but it uses a different permutation in order to obtain the same calculation
177
+ emb = torch.cat((freqs, freqs), dim=-1)
178
+ self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
179
+ self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
180
+
181
+ def forward(self, x, seq_len=None):
182
+ # x: [bs, num_attention_heads, seq_len, head_size]
183
+ if seq_len > self.max_seq_len_cached:
184
+ self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype)
185
+
186
+ return (
187
+ self.cos_cached[:seq_len].to(dtype=x.dtype),
188
+ self.sin_cached[:seq_len].to(dtype=x.dtype),
189
+ )
190
+
191
+
192
+ # Copied from transformers.models.llama.modeling_llama.rotate_half
193
+ def rotate_half(x):
194
+ """Rotates half the hidden dims of the input."""
195
+ x1 = x[..., : x.shape[-1] // 2]
196
+ x2 = x[..., x.shape[-1] // 2 :]
197
+ return torch.cat((-x2, x1), dim=-1)
198
+
199
+
200
+ # Copied from transformers.models.llama.modeling_llama.apply_rotary_pos_emb
201
+ def apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1):
202
+ """Applies Rotary Position Embedding to the query and key tensors.
203
+ Args:
204
+ q (`torch.Tensor`): The query tensor.
205
+ k (`torch.Tensor`): The key tensor.
206
+ cos (`torch.Tensor`): The cosine part of the rotary embedding.
207
+ sin (`torch.Tensor`): The sine part of the rotary embedding.
208
+ position_ids (`torch.Tensor`):
209
+ The position indices of the tokens corresponding to the query and key tensors. For example, this can be
210
+ used to pass offsetted position ids when working with a KV-cache.
211
+ unsqueeze_dim (`int`, *optional*, defaults to 1):
212
+ The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
213
+ sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
214
+ that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
215
+ k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
216
+ cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
217
+ the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
218
+ Returns:
219
+ `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
220
+ """
221
+ cos = cos[position_ids].unsqueeze(unsqueeze_dim)
222
+ sin = sin[position_ids].unsqueeze(unsqueeze_dim)
223
+ q_embed = (q * cos) + (rotate_half(q) * sin)
224
+ k_embed = (k * cos) + (rotate_half(k) * sin)
225
+ return q_embed, k_embed
226
+
227
+
228
+ class QuietMLP(nn.Module):
229
+ def __init__(self, config):
230
+ super().__init__()
231
+ self.config = config
232
+ self.hidden_size = config.hidden_size
233
+ self.intermediate_size = config.intermediate_size
234
+ self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
235
+ self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
236
+ self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
237
+ self.act_fn = ACT2FN[config.hidden_act]
238
+
239
+ def forward(self, x):
240
+ return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
241
+
242
+
243
+ # Copied from transformers.models.llama.modeling_llama.repeat_kv
244
+ def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
245
+ """
246
+ This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
247
+ num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
248
+ """
249
+
250
+ # pdb.set_trace()
251
+ batch, num_key_value_heads, slen, head_dim = hidden_states.shape
252
+ if n_rep == 1:
253
+ return hidden_states
254
+ hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
255
+ return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
256
+
257
+
258
+ class QuietAttention(nn.Module):
259
+ """
260
+ Multi-headed attention from 'Attention Is All You Need' paper. Modified to use sliding window attention: Longformer
261
+ and "Generating Long Sequences with Sparse Transformers".
262
+ """
263
+
264
+ def __init__(self, config: QuietConfig, layer_idx: Optional[int] = None):
265
+ super().__init__()
266
+ self.config = config
267
+ self.layer_idx = layer_idx
268
+ if layer_idx is None:
269
+ logger.warning_once(
270
+ f"Instantiating {self.__class__.__name__} without passing `layer_idx` is not recommended and will "
271
+ "to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` "
272
+ "when creating this class."
273
+ )
274
+
275
+ self.hidden_size = config.hidden_size
276
+ self.num_heads = config.num_attention_heads
277
+ self.head_dim = self.hidden_size // self.num_heads
278
+ self.num_key_value_heads = config.num_key_value_heads
279
+ self.num_key_value_groups = self.num_heads // self.num_key_value_heads
280
+ self.max_position_embeddings = config.max_position_embeddings
281
+ self.rope_theta = config.rope_theta
282
+ self.is_causal = True
283
+ self.attention_dropout = config.attention_dropout
284
+ self._attn_implementation = config._attn_implementation
285
+
286
+ if (self.head_dim * self.num_heads) != self.hidden_size:
287
+ raise ValueError(
288
+ f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"
289
+ f" and `num_heads`: {self.num_heads})."
290
+ )
291
+ self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)
292
+ self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
293
+ self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
294
+ self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)
295
+
296
+ self.rotary_emb = QuietRotaryEmbedding(
297
+ self.head_dim,
298
+ max_position_embeddings=self.max_position_embeddings,
299
+ base=self.rope_theta,
300
+ )
301
+
302
+ def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
303
+ return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
304
+
305
+ def forward(
306
+ self,
307
+ hidden_states: torch.Tensor,
308
+ attention_mask: Optional[torch.Tensor] = None,
309
+ position_ids: Optional[torch.LongTensor] = None,
310
+ past_key_value: Optional[Cache] = None,
311
+ output_attentions: bool = False,
312
+ use_cache: bool = False,
313
+ **kwargs,
314
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
315
+ if "padding_mask" in kwargs:
316
+ warnings.warn(
317
+ "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"
318
+ )
319
+ bsz, q_len, _ = hidden_states.size()
320
+
321
+ query_states = self.q_proj(hidden_states)
322
+ key_states = self.k_proj(hidden_states)
323
+ value_states = self.v_proj(hidden_states)
324
+
325
+ query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
326
+ key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
327
+ value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
328
+
329
+ kv_seq_len = key_states.shape[-2]
330
+ if past_key_value is not None:
331
+ if self.layer_idx is None:
332
+ raise ValueError(
333
+ f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} "
334
+ "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class "
335
+ "with a layer index."
336
+ )
337
+ kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
338
+ cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
339
+ query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
340
+
341
+ if past_key_value is not None:
342
+ cache_kwargs = {"sin": sin, "cos": cos} # Specific to RoPE models
343
+ key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
344
+
345
+ # repeat k/v heads if n_kv_heads < n_heads
346
+ key_states = repeat_kv(key_states, self.num_key_value_groups)
347
+ value_states = repeat_kv(value_states, self.num_key_value_groups)
348
+
349
+ attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
350
+
351
+ if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len):
352
+ raise ValueError(
353
+ f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is"
354
+ f" {attn_weights.size()}"
355
+ )
356
+ if self._attn_implementation == "flash_attention_2":
357
+ # Prepare attention mask for flash-attn
358
+ attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) else None
359
+ elif self._attn_implementation == "sdpa":
360
+ # Prepare attention mask for SDPA
361
+ if attention_mask is None or attention_mask.dim() == 2:
362
+ attention_mask = _prepare_4d_causal_attention_mask(
363
+ attention_mask,
364
+ (batch_size, seq_length),
365
+ inputs_embeds,
366
+ past_key_values_length,
367
+ sliding_window=self.config.sliding_window,
368
+ )
369
+ else:
370
+ # Prepare attention mask for other implementations
371
+ if attention_mask is None or attention_mask.dim() == 2:
372
+ attention_mask = _prepare_4d_causal_attention_mask(
373
+ attention_mask,
374
+ (batch_size, seq_length),
375
+ inputs_embeds,
376
+ past_key_values_length,
377
+ sliding_window=self.config.sliding_window,
378
+ )
379
+
380
+ if attention_mask is not None:
381
+ if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
382
+ raise ValueError(
383
+ f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
384
+ )
385
+
386
+ attn_weights = attn_weights + attention_mask
387
+
388
+ # upcast attention to fp32
389
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
390
+ attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training)
391
+ attn_output = torch.matmul(attn_weights, value_states)
392
+
393
+ if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):
394
+ raise ValueError(
395
+ f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"
396
+ f" {attn_output.size()}"
397
+ )
398
+
399
+ attn_output = attn_output.transpose(1, 2).contiguous()
400
+ attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
401
+
402
+ attn_output = self.o_proj(attn_output)
403
+
404
+ if not output_attentions:
405
+ attn_weights = None
406
+
407
+ return attn_output, attn_weights, past_key_value
408
+
409
+
410
+ class QuietFlashAttention2(QuietAttention):
411
+ """
412
+ Quiet flash attention module. This module inherits from `QuietAttention` as the weights of the module stays
413
+ untouched. The only required change would be on the forward pass where it needs to correctly call the public API of
414
+ flash attention and deal with padding tokens in case the input contains any of them.
415
+ """
416
+
417
+ # Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2.__init__
418
+ def __init__(self, *args, **kwargs):
419
+ super().__init__(*args, **kwargs)
420
+
421
+ # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1.
422
+ # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignement, that was made default for flash_attn>=2.1. This attribute is used to handle this difference. Reference: https://github.com/Dao-AILab/flash-attention/releases/tag/v2.1.0.
423
+ # Beware that with flash_attn<2.1, using q_seqlen != k_seqlen (except for the case q_seqlen == 1) produces a wrong mask (top-left).
424
+ self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10()
425
+
426
+ def forward(
427
+ self,
428
+ hidden_states: torch.Tensor,
429
+ attention_mask: Optional[torch.Tensor] = None,
430
+ position_ids: Optional[torch.LongTensor] = None,
431
+ past_key_value: Optional[Cache] = None,
432
+ output_attentions: bool = False,
433
+ use_cache: bool = False,
434
+ **kwargs,
435
+ ):
436
+ if "padding_mask" in kwargs:
437
+ warnings.warn(
438
+ "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"
439
+ )
440
+
441
+ # overwrite attention_mask with padding_mask
442
+ attention_mask = kwargs.pop("padding_mask")
443
+ bsz, q_len, _ = hidden_states.size()
444
+
445
+ query_states = self.q_proj(hidden_states)
446
+ key_states = self.k_proj(hidden_states)
447
+ value_states = self.v_proj(hidden_states)
448
+
449
+ query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
450
+ key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
451
+ value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
452
+
453
+ kv_seq_len = key_states.shape[-2]
454
+ if past_key_value is not None:
455
+ if self.layer_idx is None:
456
+ raise ValueError(
457
+ f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} "
458
+ "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class "
459
+ "with a layer index."
460
+ )
461
+ kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
462
+
463
+ # Because the input can be padded, the absolute sequence length depends on the max position id.
464
+ rotary_seq_len = max(kv_seq_len, position_ids[:, -1].max().item()) + 1
465
+ cos, sin = self.rotary_emb(value_states, seq_len=rotary_seq_len)
466
+
467
+ query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
468
+
469
+ use_sliding_windows = (
470
+ _flash_supports_window_size
471
+ and getattr(self.config, "sliding_window", None) is not None
472
+ and kv_seq_len > self.config.sliding_window
473
+ )
474
+
475
+ if not _flash_supports_window_size:
476
+ logger.warning_once(
477
+ "The current flash attention version does not support sliding window attention, for a more memory efficient implementation"
478
+ " make sure to upgrade flash-attn library."
479
+ )
480
+
481
+ if past_key_value is not None:
482
+ # Activate slicing cache only if the config has a value `sliding_windows` attribute
483
+ cache_has_contents = past_key_value.get_seq_length(self.layer_idx) > 0
484
+ if (
485
+ getattr(self.config, "sliding_window", None) is not None
486
+ and kv_seq_len > self.config.sliding_window
487
+ and cache_has_contents
488
+ ):
489
+ slicing_tokens = 1 - self.config.sliding_window
490
+
491
+ past_key = past_key_value[self.layer_idx][0]
492
+ past_value = past_key_value[self.layer_idx][1]
493
+
494
+ past_key = past_key[:, :, slicing_tokens:, :].contiguous()
495
+ past_value = past_value[:, :, slicing_tokens:, :].contiguous()
496
+
497
+ if past_key.shape[-2] != self.config.sliding_window - 1:
498
+ raise ValueError(
499
+ f"past key must have a shape of (`batch_size, num_heads, self.config.sliding_window-1, head_dim`), got"
500
+ f" {past_key.shape}"
501
+ )
502
+
503
+ if attention_mask is not None:
504
+ attention_mask = attention_mask[:, slicing_tokens:]
505
+ attention_mask = torch.cat([attention_mask, torch.ones_like(attention_mask[:, -1:])], dim=-1)
506
+
507
+ cache_kwargs = {"sin": sin, "cos": cos} # Specific to RoPE models
508
+ key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
509
+
510
+ # repeat k/v heads if n_kv_heads < n_heads
511
+ key_states = repeat_kv(key_states, self.num_key_value_groups)
512
+ value_states = repeat_kv(value_states, self.num_key_value_groups)
513
+ dropout_rate = 0.0 if not self.training else self.attention_dropout
514
+
515
+ # In PEFT, usually we cast the layer norms in float32 for training stability reasons
516
+ # therefore the input hidden states gets silently casted in float32. Hence, we need
517
+ # cast them back in float16 just to be sure everything works as expected.
518
+ input_dtype = query_states.dtype
519
+ if input_dtype == torch.float32:
520
+ if torch.is_autocast_enabled():
521
+ target_dtype = torch.get_autocast_gpu_dtype()
522
+ # Handle the case where the model is quantized
523
+ elif hasattr(self.config, "_pre_quantization_dtype"):
524
+ target_dtype = self.config._pre_quantization_dtype
525
+ else:
526
+ target_dtype = self.q_proj.weight.dtype
527
+
528
+ logger.warning_once(
529
+ f"The input hidden states seems to be silently casted in float32, this might be related to"
530
+ f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in"
531
+ f" {target_dtype}."
532
+ )
533
+
534
+ query_states = query_states.to(target_dtype)
535
+ key_states = key_states.to(target_dtype)
536
+ value_states = value_states.to(target_dtype)
537
+
538
+ # Reashape to the expected shape for Flash Attention
539
+ query_states = query_states.transpose(1, 2)
540
+ key_states = key_states.transpose(1, 2)
541
+ value_states = value_states.transpose(1, 2)
542
+
543
+ attn_output = self._flash_attention_forward(
544
+ query_states,
545
+ key_states,
546
+ value_states,
547
+ attention_mask,
548
+ q_len,
549
+ dropout=dropout_rate,
550
+ use_sliding_windows=use_sliding_windows,
551
+ )
552
+
553
+ attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous()
554
+ attn_output = self.o_proj(attn_output)
555
+
556
+ if not output_attentions:
557
+ attn_weights = None
558
+
559
+ return attn_output, attn_weights, past_key_value
560
+
561
+ def _flash_attention_forward(
562
+ self,
563
+ query_states,
564
+ key_states,
565
+ value_states,
566
+ attention_mask,
567
+ query_length,
568
+ dropout=0.0,
569
+ softmax_scale=None,
570
+ use_sliding_windows=False,
571
+ ):
572
+ """
573
+ Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token
574
+ first unpad the input, then computes the attention scores and pad the final attention scores.
575
+ Args:
576
+ query_states (`torch.Tensor`):
577
+ Input query states to be passed to Flash Attention API
578
+ key_states (`torch.Tensor`):
579
+ Input key states to be passed to Flash Attention API
580
+ value_states (`torch.Tensor`):
581
+ Input value states to be passed to Flash Attention API
582
+ attention_mask (`torch.Tensor`):
583
+ The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the
584
+ position of padding tokens and 1 for the position of non-padding tokens.
585
+ dropout (`int`, *optional*):
586
+ Attention dropout
587
+ softmax_scale (`float`, *optional*):
588
+ The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim)
589
+ use_sliding_windows (`bool`, *optional*):
590
+ Whether to activate sliding window attention.
591
+ """
592
+ if not self._flash_attn_uses_top_left_mask:
593
+ causal = self.is_causal
594
+ else:
595
+ # TODO: Remove the `query_length != 1` check once Flash Attention for RoCm is bumped to 2.1. For details, please see the comment in LlamaFlashAttention2 __init__.
596
+ causal = self.is_causal and query_length != 1
597
+
598
+ # Ensure attention_mask has the correct shape and values
599
+ if attention_mask is not None:
600
+ if attention_mask.dim() == 4:
601
+ # Convert 4D attention mask to 2D
602
+ attention_mask = attention_mask.squeeze(1).squeeze(1)
603
+ elif attention_mask.dim() != 2:
604
+ raise ValueError(
605
+ f"Invalid attention mask dimension: {attention_mask.dim()}. Expected 2D or 4D mask."
606
+ )
607
+
608
+ # Ensure attention_mask has values of 0 and 1
609
+ attention_mask = attention_mask.to(torch.bool).to(torch.int32)
610
+
611
+ # Contains at least one padding token in the sequence
612
+ if attention_mask is not None:
613
+ batch_size = query_states.shape[0]
614
+ query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._upad_input(
615
+ query_states, key_states, value_states, attention_mask, query_length
616
+ )
617
+
618
+ cu_seqlens_q, cu_seqlens_k = cu_seq_lens
619
+ max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens
620
+
621
+ if not use_sliding_windows:
622
+ attn_output_unpad = flash_attn_varlen_func(
623
+ query_states,
624
+ key_states,
625
+ value_states,
626
+ cu_seqlens_q=cu_seqlens_q,
627
+ cu_seqlens_k=cu_seqlens_k,
628
+ max_seqlen_q=max_seqlen_in_batch_q,
629
+ max_seqlen_k=max_seqlen_in_batch_k,
630
+ dropout_p=dropout,
631
+ softmax_scale=softmax_scale,
632
+ causal=causal,
633
+ )
634
+ else:
635
+ attn_output_unpad = flash_attn_varlen_func(
636
+ query_states,
637
+ key_states,
638
+ value_states,
639
+ cu_seqlens_q=cu_seqlens_q,
640
+ cu_seqlens_k=cu_seqlens_k,
641
+ max_seqlen_q=max_seqlen_in_batch_q,
642
+ max_seqlen_k=max_seqlen_in_batch_k,
643
+ dropout_p=dropout,
644
+ softmax_scale=softmax_scale,
645
+ causal=causal,
646
+ window_size=(self.config.sliding_window, self.config.sliding_window),
647
+ )
648
+
649
+ attn_output = pad_input(attn_output_unpad, indices_q, batch_size, query_length)
650
+ else:
651
+ if not use_sliding_windows:
652
+ attn_output = flash_attn_func(
653
+ query_states,
654
+ key_states,
655
+ value_states,
656
+ dropout,
657
+ softmax_scale=softmax_scale,
658
+ causal=causal,
659
+ )
660
+ else:
661
+ attn_output = flash_attn_func(
662
+ query_states,
663
+ key_states,
664
+ value_states,
665
+ dropout,
666
+ softmax_scale=softmax_scale,
667
+ causal=causal,
668
+ window_size=(self.config.sliding_window, self.config.sliding_window),
669
+ )
670
+
671
+ return attn_output
672
+
673
+ def _upad_input(self, query_layer, key_layer, value_layer, attention_mask, query_length):
674
+ batch_size, kv_seq_len, num_heads, head_dim = key_layer.shape
675
+
676
+ # On the first iteration we need to properly re-create the padding mask
677
+ # by slicing it on the proper place
678
+ if kv_seq_len != attention_mask.shape[-1]:
679
+ attention_mask_num_tokens = attention_mask.shape[-1]
680
+ attention_mask = attention_mask[:, attention_mask_num_tokens - kv_seq_len :]
681
+
682
+ indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)
683
+
684
+ key_layer = index_first_axis(key_layer.reshape(batch_size * kv_seq_len, num_heads, head_dim), indices_k)
685
+ value_layer = index_first_axis(value_layer.reshape(batch_size * kv_seq_len, num_heads, head_dim), indices_k)
686
+
687
+ if query_length == kv_seq_len:
688
+ query_layer = index_first_axis(
689
+ query_layer.reshape(batch_size * kv_seq_len, num_heads, head_dim), indices_k
690
+ )
691
+ cu_seqlens_q = cu_seqlens_k
692
+ max_seqlen_in_batch_q = max_seqlen_in_batch_k
693
+ indices_q = indices_k
694
+ elif query_length == 1:
695
+ max_seqlen_in_batch_q = 1
696
+ cu_seqlens_q = torch.arange(
697
+ batch_size + 1, dtype=torch.int32, device=query_layer.device
698
+ ) # There is a memcpy here, that is very bad.
699
+ indices_q = cu_seqlens_q[:-1]
700
+ query_layer = query_layer.squeeze(1)
701
+ else:
702
+ # The -q_len: slice assumes left padding.
703
+ attention_mask = attention_mask[:, -query_length:]
704
+ query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(query_layer, attention_mask)
705
+
706
+ return (
707
+ query_layer,
708
+ key_layer,
709
+ value_layer,
710
+ indices_q,
711
+ (cu_seqlens_q, cu_seqlens_k),
712
+ (max_seqlen_in_batch_q, max_seqlen_in_batch_k),
713
+ )
714
+
715
+
716
+ # Copied from transformers.models.llama.modeling_llama.LlamaSdpaAttention with Llama->Quiet
717
+ class QuietSdpaAttention(QuietAttention):
718
+ """
719
+ Quiet attention module using torch.nn.functional.scaled_dot_product_attention. This module inherits from
720
+ `QuietAttention` as the weights of the module stays untouched. The only changes are on the forward pass to adapt to
721
+ SDPA API.
722
+ """
723
+
724
+ # Adapted from QuietAttention.forward
725
+ def forward(
726
+ self,
727
+ hidden_states: torch.Tensor,
728
+ attention_mask: Optional[torch.Tensor] = None,
729
+ position_ids: Optional[torch.LongTensor] = None,
730
+ past_key_value: Optional[Cache] = None,
731
+ output_attentions: bool = False,
732
+ use_cache: bool = False,
733
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
734
+ if output_attentions:
735
+ # TODO: Improve this warning with e.g. `model.config.attn_implementation = "manual"` once this is implemented.
736
+ logger.warning_once(
737
+ "QuietModel is using QuietSdpaAttention, but `torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to the manual attention implementation, "
738
+ 'but specifying the manual implementation will be required from Transformers version v5.0.0 onwards. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.'
739
+ )
740
+ return super().forward(
741
+ hidden_states=hidden_states,
742
+ attention_mask=attention_mask,
743
+ position_ids=position_ids,
744
+ past_key_value=past_key_value,
745
+ output_attentions=output_attentions,
746
+ use_cache=use_cache,
747
+ )
748
+ bsz, q_len, _ = hidden_states.size()
749
+
750
+ query_states = self.q_proj(hidden_states)
751
+ key_states = self.k_proj(hidden_states)
752
+ value_states = self.v_proj(hidden_states)
753
+
754
+ query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
755
+ key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
756
+ value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
757
+
758
+ kv_seq_len = key_states.shape[-2]
759
+ if past_key_value is not None:
760
+ kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
761
+ cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
762
+
763
+ query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
764
+
765
+ if past_key_value is not None:
766
+ cache_kwargs = {"sin": sin, "cos": cos} # Specific to RoPE models
767
+ key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
768
+
769
+ key_states = repeat_kv(key_states, self.num_key_value_groups)
770
+ value_states = repeat_kv(value_states, self.num_key_value_groups)
771
+
772
+ if attention_mask is not None:
773
+ if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
774
+ raise ValueError(
775
+ f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
776
+ )
777
+ attention_mask = attention_mask.to(query_states.dtype)
778
+ # SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with custom attn_mask,
779
+ # Reference: https://github.com/pytorch/pytorch/issues/112577.
780
+ if query_states.device.type == "cuda" and attention_mask is not None:
781
+ query_states = query_states.contiguous()
782
+ key_states = key_states.contiguous()
783
+ value_states = value_states.contiguous()
784
+
785
+ attn_output = torch.nn.functional.scaled_dot_product_attention(
786
+ query_states,
787
+ key_states,
788
+ value_states,
789
+ attn_mask=attention_mask.to(query_states.device) if attention_mask is not None else None,
790
+ dropout_p=self.attention_dropout if self.training else 0.0,
791
+ # The q_len > 1 is necessary to match with AttentionMaskConverter.to_causal_4d that does not create a causal mask in case q_len == 1.
792
+ is_causal=self.is_causal and attention_mask is None and q_len > 1,
793
+ )
794
+
795
+ attn_output = attn_output.transpose(1, 2).contiguous()
796
+ attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
797
+
798
+ attn_output = self.o_proj(attn_output)
799
+
800
+ return attn_output, None, past_key_value
801
+
802
+
803
+ QUIET_ATTENTION_CLASSES = {
804
+ "eager": QuietAttention,
805
+ "flash_attention_2": QuietFlashAttention2,
806
+ "sdpa": QuietSdpaAttention,
807
+ }
808
+
809
+
810
+ class QuietDecoderLayer(nn.Module):
811
+ def __init__(self, config: QuietConfig, layer_idx: int):
812
+ super().__init__()
813
+ self.hidden_size = config.hidden_size
814
+
815
+ self.self_attn = QUIET_ATTENTION_CLASSES[config._attn_implementation](config, layer_idx)
816
+
817
+ self.mlp = QuietMLP(config)
818
+ self.input_layernorm = QuietRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
819
+ self.post_attention_layernorm = QuietRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
820
+
821
+ def forward(
822
+ self,
823
+ hidden_states: torch.Tensor,
824
+ attention_mask: Optional[torch.Tensor] = None,
825
+ position_ids: Optional[torch.LongTensor] = None,
826
+ past_key_value: Optional[Tuple[torch.Tensor]] = None,
827
+ output_attentions: Optional[bool] = False,
828
+ use_cache: Optional[bool] = False,
829
+ **kwargs,
830
+ ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
831
+ if "padding_mask" in kwargs:
832
+ warnings.warn(
833
+ "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"
834
+ )
835
+ """
836
+ Args:
837
+ hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
838
+ attention_mask (`torch.FloatTensor`, *optional*): attention mask of size
839
+ `(batch, sequence_length)` where padding elements are indicated by 0.
840
+ output_attentions (`bool`, *optional*):
841
+ Whether or not to return the attentions tensors of all attention layers. See `attentions` under
842
+ returned tensors for more detail.
843
+ use_cache (`bool`, *optional*):
844
+ If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
845
+ (see `past_key_values`).
846
+ past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
847
+ """
848
+
849
+ residual = hidden_states
850
+
851
+ hidden_states = self.input_layernorm(hidden_states)
852
+
853
+ # Self Attention
854
+ hidden_states, self_attn_weights, present_key_value = self.self_attn(
855
+ hidden_states=hidden_states,
856
+ attention_mask=attention_mask,
857
+ position_ids=position_ids,
858
+ past_key_value=past_key_value,
859
+ output_attentions=output_attentions,
860
+ use_cache=use_cache,
861
+ )
862
+ hidden_states = residual.to(hidden_states.device) + hidden_states
863
+
864
+ # Fully Connected
865
+ residual = hidden_states
866
+ hidden_states = self.post_attention_layernorm(hidden_states)
867
+ hidden_states = self.mlp(hidden_states)
868
+ hidden_states = residual + hidden_states
869
+
870
+ outputs = (hidden_states,)
871
+
872
+ if output_attentions:
873
+ outputs += (self_attn_weights,)
874
+
875
+ if use_cache:
876
+ outputs += (present_key_value,)
877
+
878
+ return outputs
879
+
880
+
881
+ QUIET_START_DOCSTRING = r"""
882
+ This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
883
+ library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
884
+ etc.)
885
+ This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
886
+ Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
887
+ and behavior.
888
+ Parameters:
889
+ config ([`QuietConfig`]):
890
+ Model configuration class with all the parameters of the model. Initializing with a config file does not
891
+ load the weights associated with the model, only the configuration. Check out the
892
+ [`~PreTrainedModel.from_pretrained`] method to load the model weights.
893
+ """
894
+
895
+
896
+ @add_start_docstrings(
897
+ "The bare Quiet Model outputting raw hidden-states without any specific head on top.",
898
+ QUIET_START_DOCSTRING,
899
+ )
900
+ class QuietPreTrainedModel(PreTrainedModel):
901
+ config_class = QuietConfig
902
+ base_model_prefix = "model"
903
+ supports_gradient_checkpointing = True
904
+ _no_split_modules = ["QuietDecoderLayer"]
905
+ _skip_keys_device_placement = "past_key_values"
906
+ _supports_flash_attn_2 = True
907
+ _supports_sdpa = True
908
+ _supports_cache_class = True
909
+
910
+ def _init_weights(self, module):
911
+ std = self.config.initializer_range
912
+ if isinstance(module, nn.Linear):
913
+ module.weight.data.normal_(mean=0.0, std=std)
914
+ if module.bias is not None:
915
+ module.bias.data.zero_()
916
+ elif isinstance(module, nn.Embedding):
917
+ module.weight.data.normal_(mean=0.0, std=std)
918
+ if module.padding_idx is not None:
919
+ module.weight.data[module.padding_idx].zero_()
920
+
921
+
922
+ QUIET_INPUTS_DOCSTRING = r"""
923
+ Args:
924
+ input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
925
+ Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide
926
+ it.
927
+ Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
928
+ [`PreTrainedTokenizer.__call__`] for details.
929
+ [What are input IDs?](../glossary#input-ids)
930
+ attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
931
+ Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
932
+ - 1 for tokens that are **not masked**,
933
+ - 0 for tokens that are **masked**.
934
+ [What are attention masks?](../glossary#attention-mask)
935
+ Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
936
+ [`PreTrainedTokenizer.__call__`] for details.
937
+ If `past_key_values` is used, optionally only the last `decoder_input_ids` have to be input (see
938
+ `past_key_values`).
939
+ If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`]
940
+ and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more
941
+ information on the default strategy.
942
+ - 1 indicates the head is **not masked**,
943
+ - 0 indicates the head is **masked**.
944
+ position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
945
+ Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
946
+ config.n_positions - 1]`.
947
+ [What are position IDs?](../glossary#position-ids)
948
+ past_key_values (`Cache` or `tuple(tuple(torch.FloatTensor))`, *optional*):
949
+ Pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention
950
+ blocks) that can be used to speed up sequential decoding. This typically consists in the `past_key_values`
951
+ returned by the model at a previous stage of decoding, when `use_cache=True` or `config.use_cache=True`.
952
+ Two formats are allowed:
953
+ - a [`~cache_utils.Cache`] instance;
954
+ - Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of
955
+ shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`). This is also known as the legacy
956
+ cache format.
957
+ The model will output the same cache format that is fed as input. If no `past_key_values` are passed, the
958
+ legacy cache format will be returned.
959
+ If `past_key_values` are used, the user can optionally input only the last `input_ids` (those that don't
960
+ have their past key value states given to this model) of shape `(batch_size, 1)` instead of all `input_ids`
961
+ of shape `(batch_size, sequence_length)`.
962
+ inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
963
+ Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
964
+ is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
965
+ model's internal embedding lookup matrix.
966
+ use_cache (`bool`, *optional*):
967
+ If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
968
+ `past_key_values`).
969
+ output_attentions (`bool`, *optional*):
970
+ Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
971
+ tensors for more detail.
972
+ output_hidden_states (`bool`, *optional*):
973
+ Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
974
+ more detail.
975
+ return_dict (`bool`, *optional*):
976
+ Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
977
+ """
978
+
979
+
980
+ @add_start_docstrings(
981
+ "The bare Quiet Model outputting raw hidden-states without any specific head on top.",
982
+ QUIET_START_DOCSTRING,
983
+ )
984
+ class QuietModel(QuietPreTrainedModel):
985
+ """
986
+ Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`QuietDecoderLayer`]
987
+ Args:
988
+ config: QuietConfig
989
+ """
990
+
991
+ def __init__(self, config: QuietConfig):
992
+ super().__init__(config)
993
+ self.padding_idx = config.pad_token_id
994
+ self.vocab_size = config.vocab_size
995
+
996
+ self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
997
+ self.layers = nn.ModuleList(
998
+ [QuietDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
999
+ )
1000
+ self._attn_implementation = config._attn_implementation
1001
+ self.norm = QuietRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
1002
+
1003
+ self.gradient_checkpointing = False
1004
+ # Initialize weights and apply final processing
1005
+ self.post_init()
1006
+
1007
+ def get_input_embeddings(self):
1008
+ return self.embed_tokens
1009
+
1010
+ def set_input_embeddings(self, value):
1011
+ self.embed_tokens = value
1012
+
1013
+ @add_start_docstrings_to_model_forward(QUIET_INPUTS_DOCSTRING)
1014
+ def forward(
1015
+ self,
1016
+ input_ids: torch.LongTensor = None,
1017
+ attention_mask: Optional[torch.Tensor] = None,
1018
+ position_ids: Optional[torch.LongTensor] = None,
1019
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
1020
+ inputs_embeds: Optional[torch.FloatTensor] = None,
1021
+ use_cache: Optional[bool] = None,
1022
+ output_attentions: Optional[bool] = None,
1023
+ output_hidden_states: Optional[bool] = None,
1024
+ return_dict: Optional[bool] = None,
1025
+ ) -> Union[Tuple, BaseModelOutputWithPast]:
1026
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
1027
+ output_hidden_states = (
1028
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
1029
+ )
1030
+ use_cache = use_cache if use_cache is not None else self.config.use_cache
1031
+
1032
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
1033
+
1034
+ # retrieve input_ids and inputs_embeds
1035
+ if input_ids is not None and inputs_embeds is not None:
1036
+ raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time")
1037
+ elif input_ids is not None:
1038
+ batch_size, seq_length = input_ids.shape
1039
+ elif inputs_embeds is not None:
1040
+ batch_size, seq_length, _ = inputs_embeds.shape
1041
+ else:
1042
+ raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds")
1043
+
1044
+ if self.gradient_checkpointing and self.training:
1045
+ if use_cache:
1046
+ logger.warning_once(
1047
+ "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
1048
+ )
1049
+ use_cache = False
1050
+
1051
+ past_key_values_length = 0
1052
+
1053
+ if use_cache:
1054
+ use_legacy_cache = not isinstance(past_key_values, Cache)
1055
+ if use_legacy_cache:
1056
+ past_key_values = DynamicCache.from_legacy_cache(past_key_values)
1057
+ past_key_values_length = past_key_values.get_usable_length(seq_length)
1058
+
1059
+ if position_ids is None:
1060
+ device = input_ids.device if input_ids is not None else inputs_embeds.device
1061
+ position_ids = torch.arange(
1062
+ past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=device
1063
+ )
1064
+ position_ids = position_ids.unsqueeze(0).view(-1, seq_length)
1065
+ else:
1066
+ position_ids = position_ids.view(-1, seq_length).long()
1067
+
1068
+ if inputs_embeds is None:
1069
+ inputs_embeds = self.embed_tokens(input_ids)
1070
+
1071
+ if self._attn_implementation == "flash_attention_2":
1072
+ # 2d mask is passed through the layers
1073
+ attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) else None
1074
+ elif self._attn_implementation == "sdpa" and not output_attentions and attention_mask.dim() == 2 and False:
1075
+ # output_attentions=True can not be supported when using SDPA, and we fall back on
1076
+ # the manual implementation that requires a 4D causal mask in all cases.
1077
+ attention_mask = _prepare_4d_causal_attention_mask_for_sdpa(
1078
+ attention_mask,
1079
+ (batch_size, seq_length),
1080
+ inputs_embeds,
1081
+ past_key_values_length,
1082
+ )
1083
+ elif attention_mask is None or attention_mask.dim() == 2:
1084
+ # 4d mask is passed through the layers
1085
+ attention_mask = _prepare_4d_causal_attention_mask(
1086
+ attention_mask,
1087
+ (batch_size, seq_length),
1088
+ inputs_embeds,
1089
+ past_key_values_length,
1090
+ sliding_window=self.config.sliding_window,
1091
+ )
1092
+
1093
+ hidden_states = inputs_embeds
1094
+
1095
+ # decoder layers
1096
+ all_hidden_states = () if output_hidden_states else None
1097
+ all_self_attns = () if output_attentions else None
1098
+ next_decoder_cache = None
1099
+
1100
+ for decoder_layer in self.layers:
1101
+ if output_hidden_states:
1102
+ all_hidden_states += (hidden_states,)
1103
+
1104
+ if self.gradient_checkpointing and self.training:
1105
+ layer_outputs = self._gradient_checkpointing_func(
1106
+ decoder_layer.__call__,
1107
+ hidden_states,
1108
+ attention_mask,
1109
+ position_ids,
1110
+ past_key_values,
1111
+ output_attentions,
1112
+ use_cache,
1113
+ )
1114
+ else:
1115
+ layer_outputs = decoder_layer(
1116
+ hidden_states,
1117
+ attention_mask=attention_mask,
1118
+ position_ids=position_ids,
1119
+ past_key_value=past_key_values,
1120
+ output_attentions=output_attentions,
1121
+ use_cache=use_cache,
1122
+ )
1123
+
1124
+ hidden_states = layer_outputs[0]
1125
+
1126
+ if use_cache:
1127
+ next_decoder_cache = layer_outputs[2 if output_attentions else 1]
1128
+
1129
+ if output_attentions:
1130
+ all_self_attns += (layer_outputs[1],)
1131
+
1132
+ hidden_states = self.norm(hidden_states)
1133
+
1134
+ # add hidden states from the last decoder layer
1135
+ if output_hidden_states:
1136
+ all_hidden_states += (hidden_states,)
1137
+
1138
+ next_cache = None
1139
+ if use_cache:
1140
+ next_cache = next_decoder_cache.to_legacy_cache() if use_legacy_cache else next_decoder_cache
1141
+
1142
+ if not return_dict:
1143
+ return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)
1144
+ return BaseModelOutputWithPast(
1145
+ last_hidden_state=hidden_states,
1146
+ past_key_values=next_cache,
1147
+ hidden_states=all_hidden_states,
1148
+ attentions=all_self_attns,
1149
+ )
1150
+
1151
+ def nonzero_mean(x, axis=None):
1152
+ if axis is not None:
1153
+ return x.sum(axis) / (x != 0).sum(axis)
1154
+ return x.sum() / (x != 0).sum()
1155
+
1156
+ def loss_mean(x):
1157
+ return x.sum() / (x != 0).sum()
1158
+
1159
+ class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
1160
+ _tied_weights_keys = ["lm_head.weight"]
1161
+
1162
+ def __init__(self, config):
1163
+ super().__init__(config)
1164
+ self.model = QuietModel(config)
1165
+ self.vocab_size = config.vocab_size
1166
+ self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
1167
+ # self.router_aux_loss_coef = config.router_aux_loss_coef
1168
+ # self.num_experts = config.num_experts
1169
+ # self.num_experts_per_tok = config.num_experts_per_tok
1170
+ self.max_thoughts = config.max_thoughts
1171
+ self.merged_lm_and_talk_heads = config.merged_lm_and_talk_heads
1172
+ self.use_concat_talk_head = config.use_concat_talk_head
1173
+ self.use_shallow_talk = config.use_shallow_talk
1174
+ self.use_complex_talk_head = config.use_complex_talk_head
1175
+ self.use_weighted_talk_head = config.use_weighted_talk_head
1176
+ # the weighted head will output a single value, so it can't be passed to the lm head
1177
+ assert not (self.use_weighted_talk_head and self.use_shallow_talk)
1178
+
1179
+ self.n_ahead = 1
1180
+ self.n_ahead_talk = 1
1181
+ self.n_passes = 1
1182
+ self.n_tokens_print = 1
1183
+ self.gradient_accumulation_steps = 1
1184
+ self.training_steps = 0
1185
+ self.tokenizer = AutoTokenizer.from_pretrained("Crystalcareai/Quiet-Star-Custom")
1186
+ self.start_token_id = None
1187
+ self.end_token_id = None
1188
+ self.rm_initialized = False
1189
+ self.residual_talk_head = True
1190
+ self.thought_init_std_scale = 1e-2
1191
+
1192
+ self.final_only_mode = False
1193
+ self.first_and_last_mode = True
1194
+ self.first_only = False
1195
+ self.original_loss_weight = 0.5
1196
+
1197
+ self.cumulative_residual = False
1198
+ self.clever_residual = False
1199
+ self.skip_residual = False
1200
+ self.no_residual = True
1201
+
1202
+ self.optimize_lm_head_only_at_start = False
1203
+ self.optimize_model_only_at_start = False
1204
+
1205
+ if self.optimize_model_only_at_start:
1206
+ raise NotImplementedError
1207
+ self.train_only_thinking_embedding = False
1208
+ self.weighted_embeddings = False
1209
+ self.use_start_thought_token = True
1210
+ self.use_end_thought_token = True
1211
+ self.initialize_thought_embedding_to_normal = False
1212
+ self.initial_start_token = "---"
1213
+ self.initial_end_token = "---"
1214
+ self.output_logits_at_the_end = True
1215
+
1216
+ self.wandb_enabled = False
1217
+ self.gumbel_temperature = 0.001
1218
+
1219
+ self.use_policy_loss = True
1220
+ self.include_policy_loss = True
1221
+ self.trice_mode = True
1222
+ self.remove_negative_rewards = True
1223
+ self.use_policy_loss_for_end_thought = True
1224
+
1225
+ self.base_original_mode = False
1226
+ self.original_mode = False
1227
+
1228
+ self.thought_prefix = "(Let's think step by step"
1229
+ self.tokenized_thought_prefix = None
1230
+ self.log_dict = defaultdict(int)
1231
+ self.eval_log_dict = defaultdict(int)
1232
+ self.loss_mean = loss_mean
1233
+
1234
+ self.start_embedding = nn.Parameter(torch.zeros(2, self.model.config.hidden_size))
1235
+ self.end_embedding = nn.Parameter(torch.zeros(2, self.model.config.hidden_size))
1236
+
1237
+ self.policy_loss_beta = 1e6
1238
+ self.embedding_scale = 1e2
1239
+ self.temperature = nn.Parameter(torch.ones(1))
1240
+ self.max_temperature = config.max_temperature
1241
+ self.reinforce_temperature = 3
1242
+ self.base_loss_beta = 1
1243
+ self.thinking_usefulness_head = nn.Linear(self.model.config.hidden_size, 1)
1244
+ self.thinking_threshold = 0.5
1245
+ self.thinking_usefulness_loss_weight = 1e-2
1246
+
1247
+ # Not used in the paper:
1248
+ self.use_thought_prefix = False
1249
+ self.use_reparam_for_thought_embeddings = False
1250
+ self.use_upper_triangular = False
1251
+ self.subtract_mean_reward = False
1252
+ self.comparison_mode = False
1253
+ self.gumbel_detach = False
1254
+
1255
+ # For visualization
1256
+ self.eval_mode = False
1257
+
1258
+ num_talk = 1
1259
+ talk_input_dim = config.hidden_size if not self.use_concat_talk_head else config.hidden_size * 2
1260
+ if self.use_weighted_talk_head:
1261
+ talk_output_dim = 1
1262
+ else:
1263
+ talk_output_dim = config.hidden_size if self.use_shallow_talk else config.vocab_size
1264
+
1265
+ if not self.merged_lm_and_talk_heads:
1266
+ if self.use_complex_talk_head:
1267
+ self.talk_head = nn.ModuleList([nn.Sequential(
1268
+ nn.Linear(talk_input_dim, config.hidden_size),
1269
+ nn.ReLU(),
1270
+ nn.Linear(config.hidden_size, config.hidden_size),
1271
+ nn.ReLU(),
1272
+ nn.Linear(config.hidden_size, talk_output_dim, bias=False)
1273
+ )])
1274
+ else:
1275
+ self.talk_head = nn.ModuleList([nn.Sequential(
1276
+ nn.Linear(talk_input_dim, talk_output_dim, bias=False)
1277
+ )])
1278
+
1279
+ self.apply(self._init_weights)
1280
+
1281
+ # Add dropout regularization
1282
+ self.dropout = nn.Dropout(config.hidden_dropout_prob)
1283
+
1284
+ # Initialize weights and apply final processing
1285
+ self.post_init()
1286
+
1287
+ def get_input_embeddings(self):
1288
+ return self.model.embed_tokens
1289
+
1290
+ def set_input_embeddings(self, value):
1291
+ self.model.embed_tokens = value
1292
+
1293
+ def get_output_embeddings(self):
1294
+ return self.lm_head
1295
+
1296
+ def set_output_embeddings(self, new_embeddings):
1297
+ self.lm_head = new_embeddings
1298
+
1299
+ def set_decoder(self, decoder):
1300
+ self.model = decoder
1301
+
1302
+ def get_decoder(self):
1303
+ return self.model
1304
+
1305
+ def _init_weights(self, module):
1306
+ if isinstance(module, nn.Linear):
1307
+ nn.init.xavier_uniform_(module.weight)
1308
+ if module.bias is not None:
1309
+ nn.init.constant_(module.bias, 0)
1310
+ elif isinstance(module, nn.Embedding):
1311
+ nn.init.xavier_uniform_(module.weight)
1312
+
1313
+ @torch.no_grad()
1314
+ def infer(
1315
+ self,
1316
+ input_ids: torch.LongTensor,
1317
+ attention_mask: Optional[torch.Tensor] = None,
1318
+ position_ids: Optional[torch.LongTensor] = None,
1319
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
1320
+ inputs_embeds: Optional[torch.FloatTensor] = None,
1321
+ use_cache: Optional[bool] = None,
1322
+ output_attentions: Optional[bool] = None,
1323
+ output_hidden_states: Optional[bool] = None,
1324
+ return_dict: Optional[bool] = None,
1325
+ ):
1326
+ batch_size, seq_len = input_ids.shape
1327
+
1328
+ # Save the original input_ids and attention_mask for later use
1329
+ original_input_ids = input_ids.clone()
1330
+ original_attention_mask = attention_mask.clone() if attention_mask is not None else None
1331
+
1332
+ # Append the start thought token to the input sequence
1333
+ start_thought_token_id = self.tokenizer.convert_tokens_to_ids("<|startthought|>")
1334
+ input_ids = torch.cat([input_ids, torch.tensor([[start_thought_token_id]] * batch_size).to(input_ids.device)], dim=-1)
1335
+ seq_len += 1
1336
+
1337
+ # Update the attention mask
1338
+ if attention_mask is not None:
1339
+ attention_mask = torch.cat([attention_mask, torch.ones((batch_size, 1)).to(attention_mask.device)], dim=-1)
1340
+
1341
+ # Generate the continuation
1342
+ continuation_length = self.n_ahead - 2
1343
+ new_key_values = past_key_values
1344
+
1345
+ # Initialize next_token_id with a default value
1346
+ next_token_id = torch.zeros(batch_size, dtype=torch.long).to(input_ids.device)
1347
+
1348
+ start_time = time.time()
1349
+ for continuation_idx in range(continuation_length):
1350
+ outputs = self.model(
1351
+ input_ids=input_ids if continuation_idx == 0 else next_token_id.unsqueeze(-1).to(input_ids.device),
1352
+ attention_mask=attention_mask,
1353
+ position_ids=position_ids,
1354
+ past_key_values=new_key_values,
1355
+ inputs_embeds=inputs_embeds,
1356
+ use_cache=True,
1357
+ output_attentions=output_attentions,
1358
+ output_hidden_states=output_hidden_states,
1359
+ return_dict=return_dict,
1360
+ )
1361
+ new_key_values = outputs.past_key_values
1362
+
1363
+ hidden_states = outputs[0]
1364
+
1365
+ logits = self.lm_head(hidden_states)
1366
+ logits = logits[:, -1, :] # Only consider the last token
1367
+
1368
+ # Apply Gumbel-Softmax to the logits
1369
+ next_token_logits = F.gumbel_softmax(logits, tau=self.gumbel_temperature, hard=True, dim=-1)
1370
+ next_token_id = torch.argmax(next_token_logits, dim=-1)
1371
+
1372
+ # Append the generated token to the input sequence
1373
+ # input_ids = torch.cat([input_ids, next_token_id.unsqueeze(-1).to(input_ids.device)], dim=-1)
1374
+ seq_len += 1
1375
+
1376
+ # Update the attention mask
1377
+ if attention_mask is not None:
1378
+ attention_mask = torch.cat([attention_mask, torch.ones((batch_size, 1)).to(attention_mask.device)], dim=-1)
1379
+
1380
+ # Append the end thought token to the input sequence
1381
+ end_thought_token_id = self.tokenizer.convert_tokens_to_ids("<|endthought|>")
1382
+ input_ids = torch.cat([input_ids, torch.tensor([[end_thought_token_id]] * batch_size).to(input_ids.device)], dim=-1)
1383
+ seq_len += 1
1384
+
1385
+ # Update the attention mask
1386
+ if attention_mask is not None:
1387
+ attention_mask = torch.cat([attention_mask, torch.ones((batch_size, 1)).to(attention_mask.device)], dim=-1)
1388
+
1389
+ # Get the hidden states before and after the thought
1390
+ outputs_before = self.model(
1391
+ input_ids=original_input_ids,
1392
+ attention_mask=original_attention_mask,
1393
+ position_ids=position_ids,
1394
+ past_key_values=past_key_values,
1395
+ inputs_embeds=inputs_embeds,
1396
+ use_cache=use_cache,
1397
+ output_attentions=output_attentions,
1398
+ output_hidden_states=output_hidden_states,
1399
+ return_dict=return_dict,
1400
+ )
1401
+ hidden_states_before = outputs_before[0][:, -1:, :]
1402
+
1403
+ # two new tokens: last continuation token and end thought token
1404
+ outputs_after = self.model(
1405
+ input_ids=torch.cat([next_token_id.unsqueeze(-1).to(input_ids.device), torch.tensor([[end_thought_token_id]] * batch_size).to(input_ids.device)], dim=-1),
1406
+ attention_mask=torch.cat([attention_mask[:, -1:], torch.ones((batch_size, 1)).to(attention_mask.device)], dim=-1),
1407
+ position_ids=position_ids,
1408
+ past_key_values=new_key_values,
1409
+ inputs_embeds=inputs_embeds,
1410
+ use_cache=use_cache,
1411
+ output_attentions=output_attentions,
1412
+ output_hidden_states=output_hidden_states,
1413
+ return_dict=return_dict,
1414
+ )
1415
+ hidden_states_after = outputs_after[0][:, -1:, :]
1416
+
1417
+ # Apply the talk head to get the mixing weight
1418
+ mixing_weight = self.talk_head[0](torch.cat([hidden_states_before, hidden_states_after], dim=-1))
1419
+
1420
+ # Apply the mixing weight to the hidden states
1421
+ mixed_hidden_states = (1 - mixing_weight) * hidden_states_before + mixing_weight * hidden_states_after
1422
+
1423
+ # Apply the language model head to get the final logits
1424
+ logits = self.lm_head(mixed_hidden_states)
1425
+ return logits
1426
+
1427
+ @torch.no_grad()
1428
+ def generate(
1429
+ self,
1430
+ input_ids: torch.LongTensor,
1431
+ attention_mask: Optional[torch.Tensor] = None,
1432
+ max_new_tokens: Optional[int] = None,
1433
+ temperature: float = 1.1,
1434
+ **kwargs,
1435
+ ):
1436
+ if attention_mask is None:
1437
+ # Create a default attention mask if not provided
1438
+ attention_mask = torch.ones_like(input_ids)
1439
+
1440
+ from .generate import generate
1441
+ return generate(self, input_ids=input_ids, attention_mask=attention_mask, max_new_tokens=max_new_tokens, temperature=temperature, **kwargs)
1442
+
1443
+ @add_start_docstrings_to_model_forward(QUIET_INPUTS_DOCSTRING)
1444
+ @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)
1445
+ def forward(
1446
+ self,
1447
+ input_ids: torch.LongTensor = None,
1448
+ attention_mask: Optional[torch.Tensor] = None,
1449
+ position_ids: Optional[torch.LongTensor] = None,
1450
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
1451
+ inputs_embeds: Optional[torch.FloatTensor] = None,
1452
+ labels: Optional[torch.LongTensor] = None,
1453
+ use_cache: Optional[bool] = None,
1454
+ # output_router_logits: Optional[bool] = None,
1455
+ output_attentions: Optional[bool] = None,
1456
+ output_hidden_states: Optional[bool] = None,
1457
+ return_dict: Optional[bool] = None,
1458
+ ) -> Union[Tuple, CausalLMOutputWithPast]:
1459
+ r"""
1460
+ Args:
1461
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
1462
+ Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
1463
+ config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
1464
+ (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
1465
+ Returns:
1466
+ Example:
1467
+ ```python
1468
+ >>> from transformers import AutoTokenizer, QuietForCausalLM
1469
+ >>> model = QuietForCausalLM.from_pretrained("quietai/Quiet-7B-v0.1")
1470
+ >>> tokenizer = AutoTokenizer.from_pretrained("quietai/Quiet-7B-v0.1")
1471
+ >>> prompt = "Hey, are you conscious? Can you talk to me?"
1472
+ >>> inputs = tokenizer(prompt, return_tensors="pt")
1473
+ >>> # Generate
1474
+ >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
1475
+ >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
1476
+ "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
1477
+ ```"""
1478
+
1479
+ if not self.training:
1480
+ n_ahead_talk_to_restore = self.n_ahead_talk
1481
+ n_passes_to_restore = self.n_passes
1482
+ self.n_ahead_talk = 1
1483
+ self.n_passes = 1
1484
+
1485
+ # aux_loss = None
1486
+ # output_router_logits = output_router_logits if output_router_logits is not None else self.config.output_router_logits
1487
+ # if output_router_logits:
1488
+ # router_logits = outputs.router_logits if return_dict else outputs[-1]
1489
+ # if router_logits is not None:
1490
+ # aux_loss = load_balancing_loss_func(
1491
+ # router_logits,
1492
+ # self.num_experts,
1493
+ # self.num_experts_per_tok,
1494
+ # attention_mask,
1495
+ # )
1496
+ # if labels is not None:
1497
+ # loss += self.router_aux_loss_coef * aux_loss.to(loss.device)
1498
+
1499
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
1500
+ output_hidden_states = (
1501
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
1502
+ )
1503
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
1504
+
1505
+ assert self.cumulative_residual or self.clever_residual or self.skip_residual or self.no_residual
1506
+ assert not (self.skip_residual and self.use_policy_loss)
1507
+
1508
+ if self.tokenized_thought_prefix is None and self.use_thought_prefix:
1509
+ self.tokenized_thought_prefix = self.tokenizer(self.thought_prefix, return_tensors="pt", add_special_tokens=False)["input_ids"]
1510
+
1511
+ def apply_head(head, states, detach=False):
1512
+ if detach:
1513
+ head_weight = head.weight.detach()
1514
+ else:
1515
+ head_weight = head.weight
1516
+ head_weight = head_weight.to(states.device)
1517
+ return (head_weight @ states.transpose(-1, -2)).transpose(-1, -2).contiguous()
1518
+
1519
+ def idx_if_sequential(head, idx=0):
1520
+ if isinstance(head, nn.Sequential) or isinstance(head, nn.ModuleList):
1521
+ return idx_if_sequential(head[idx], idx=idx)
1522
+ return head
1523
+
1524
+ def none_repeat_interleave(x, n):
1525
+ if x is None:
1526
+ return x
1527
+ return x.repeat_interleave(n, dim=0)
1528
+
1529
+ if self.n_passes > 1:
1530
+ input_ids = none_repeat_interleave(input_ids, self.n_passes)
1531
+ attention_mask = none_repeat_interleave(attention_mask, self.n_passes)
1532
+ position_ids = none_repeat_interleave(position_ids, self.n_passes)
1533
+ inputs_embeds = none_repeat_interleave(inputs_embeds, self.n_passes)
1534
+ labels = none_repeat_interleave(labels, self.n_passes)
1535
+ if past_key_values is not None:
1536
+ past_key_values = [none_repeat_interleave(p, self.n_passes) for p in past_key_values]
1537
+ cur_token_indices = torch.arange(input_ids.shape[1], device=input_ids.device)
1538
+
1539
+ self.tokenizer_has_start_thought_token = True
1540
+ self.tokenizer_has_end_thought_token = True
1541
+ if self.start_token_id is None:
1542
+ self.start_token_id = self.tokenizer.convert_tokens_to_ids("<|startthought|>")
1543
+ if self.start_token_id == 0:
1544
+ self.start_token_id = self.tokenizer.bos_token_id
1545
+ self.tokenizer_has_start_thought_token = False
1546
+ elif self.use_start_thought_token:
1547
+ # base_start_id = self.tokenizer.convert_tokens_to_ids(self.initial_start_token)
1548
+ base_start_id = self.tokenizer.encode(self.initial_start_token, add_special_tokens=False)[0]
1549
+ if self.initialize_thought_embedding_to_normal:
1550
+ self.start_embedding.data = torch.zeros_like(self.start_embedding.data)
1551
+ else:
1552
+ self.start_embedding.data[0] = self.model.embed_tokens.weight.data[base_start_id].clone().detach() / self.embedding_scale
1553
+ self.start_embedding.data[1] = torch.log(self.model.embed_tokens.weight.data.std(dim=0) * self.thought_init_std_scale / self.embedding_scale)
1554
+ if self.end_token_id is None:
1555
+ self.end_token_id = self.tokenizer.convert_tokens_to_ids("<|endthought|>")
1556
+ if self.end_token_id == 0:
1557
+ self.end_token_id = self.tokenizer.eos_token_id
1558
+ self.tokenizer_has_end_thought_token = False
1559
+ elif self.use_end_thought_token:
1560
+ # base_end_id = self.tokenizer.convert_tokens_to_ids(self.initial_end_token)
1561
+ base_end_id = self.tokenizer.encode(self.initial_end_token, add_special_tokens=False)[0]
1562
+ if self.initialize_thought_embedding_to_normal:
1563
+ self.end_embedding.data = torch.zeros_like(self.end_embedding.data)
1564
+ else:
1565
+ self.end_embedding.data[0] = self.model.embed_tokens.weight.data[base_end_id].clone().detach() / self.embedding_scale
1566
+ self.end_embedding.data[1] = torch.log(self.model.embed_tokens.weight.data.std(dim=0) * self.thought_init_std_scale / self.embedding_scale)
1567
+
1568
+ if not self.rm_initialized and (self.n_ahead > 1 or not self.base_original_mode):
1569
+ self.rm_initialized = True
1570
+ if not self.use_shallow_talk:
1571
+ head = self.talk_head[0]
1572
+ cur_head = head[-1] if isinstance(head, nn.Sequential) else head
1573
+ talk_input_dim = cur_head.weight.data.shape[1]
1574
+ talk_output_dim = 1 if self.use_weighted_talk_head else self.lm_head.weight.data.shape[0]
1575
+ cur_head.weight.data = torch.zeros(talk_output_dim, talk_input_dim, device=cur_head.weight.device, dtype=cur_head.weight.dtype)
1576
+ else:
1577
+ # convert to identity transform
1578
+ def lambda_transform(cur_head):
1579
+ # pdb.set_trace()
1580
+ if cur_head.weight.data.shape[0] != cur_head.weight.data.shape[1]:
1581
+ return torch.cat([
1582
+ torch.eye(
1583
+ cur_head.weight.data.shape[0],
1584
+ device=cur_head.weight.device,
1585
+ dtype=cur_head.weight.dtype
1586
+ ),
1587
+ torch.zeros(
1588
+ cur_head.weight.data.shape[0],
1589
+ cur_head.weight.data.shape[1] - cur_head.weight.data.shape[0],
1590
+ device=cur_head.weight.device,
1591
+ dtype=cur_head.weight.dtype
1592
+ )], dim=1)
1593
+ return torch.eye(
1594
+ cur_head.weight.data.shape[0],
1595
+ device=cur_head.weight.device,
1596
+ dtype=cur_head.weight.dtype
1597
+ )
1598
+ if isinstance(self.talk_head[0], nn.Sequential):
1599
+ for cur_head in self.talk_head[0]:
1600
+ # if it has weights
1601
+ if hasattr(cur_head, "weight"):
1602
+ cur_head.weight.data = lambda_transform(cur_head)
1603
+ else:
1604
+ self.talk_head[-1].weight.data = lambda_transform(self.talk_head[0])
1605
+
1606
+ loss = None
1607
+ prev_rm_tokens = None
1608
+ cur_rm_tokens = None
1609
+ prev_rm_logits = None
1610
+ prev_sample_probs = None
1611
+ did_skip_sampling = None
1612
+ skip_sampling = None
1613
+ sample_probs = None
1614
+ hidden_states = None
1615
+ logits = None
1616
+ talk_kl_penalty = None
1617
+ rm_logits = None
1618
+ residual_logits = None
1619
+ probabilities_2d = None
1620
+ prev_probabilities_2d = None
1621
+ policy_reward = None
1622
+ logits_to_output = None
1623
+ batch_size, seq_len = input_ids.shape
1624
+ base_input_ids = input_ids.clone()
1625
+ loss_list = []
1626
+ dqn_loss_list = []
1627
+ sampled_token_history = []
1628
+ sample_probs_history = []
1629
+ action_loglikelihoods_list = []
1630
+
1631
+ temperature = self.temperature
1632
+
1633
+ if self.use_end_thought_token or self.use_start_thought_token:
1634
+ if not self.use_reparam_for_thought_embeddings:
1635
+ start_embedding = self.start_embedding[0].unsqueeze(0) * self.embedding_scale * temperature
1636
+ end_embedding = self.end_embedding[0].unsqueeze(0) * self.embedding_scale * temperature
1637
+ else:
1638
+ start_embedding = self.start_embedding * self.embedding_scale * temperature
1639
+ end_embedding = self.end_embedding * self.embedding_scale * temperature
1640
+ base_embeddings = self.model.embed_tokens.weight
1641
+ if self.train_only_thinking_embedding:
1642
+ base_embeddings = base_embeddings.detach()
1643
+
1644
+ # # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
1645
+ fwd_iters = 1 if self.original_mode else self.n_ahead + self.n_ahead_talk - 1
1646
+ for ahead_idx in range(fwd_iters):
1647
+ past_key_values_length = 0
1648
+ if past_key_values is not None:
1649
+ use_legacy_cache = not isinstance(past_key_values, Cache)
1650
+ if use_legacy_cache:
1651
+ past_key_values = DynamicCache.from_legacy_cache(past_key_values)
1652
+ past_key_values_length = past_key_values.get_usable_length(seq_len)
1653
+
1654
+ if position_ids is None:
1655
+ device = input_ids.device if input_ids is not None else inputs_embeds.device
1656
+ position_ids = torch.arange(
1657
+ past_key_values_length, seq_len + past_key_values_length, dtype=torch.long, device=device
1658
+ )
1659
+ position_ids = position_ids.unsqueeze(0).view(-1, seq_len)
1660
+ else:
1661
+ position_ids = position_ids.view(-1, seq_len).long()
1662
+
1663
+ if inputs_embeds is None:
1664
+ contains_start = self.use_start_thought_token and (input_ids == self.start_token_id).any()
1665
+ contains_end = self.use_end_thought_token and (input_ids == self.end_token_id).any()
1666
+ contains_thought = contains_start or contains_end
1667
+ if contains_thought:
1668
+ thought_id = self.start_token_id if contains_start else self.end_token_id
1669
+ cur_thought_embedding = start_embedding if contains_start else end_embedding
1670
+ if self.use_reparam_for_thought_embeddings:
1671
+ inputs_embeds = torch.randn(batch_size, seq_len, self.model.config.hidden_size, device=input_ids.device, dtype=cur_thought_embedding.dtype)
1672
+ inputs_embeds = inputs_embeds.detach() * torch.exp(cur_thought_embedding[1]) + cur_thought_embedding[0]
1673
+ if contains_start:
1674
+ sampled_start = inputs_embeds.clone().detach()
1675
+ if contains_end:
1676
+ sampled_end = inputs_embeds.clone().detach()
1677
+ else:
1678
+ inputs_embeds = cur_thought_embedding.unsqueeze(0).repeat(batch_size, seq_len, 1)
1679
+ else:
1680
+ with torch.set_grad_enabled(not self.train_only_thinking_embedding):
1681
+ inputs_embeds = self.model.embed_tokens(input_ids)
1682
+
1683
+ if self.n_ahead != 1 or self.n_ahead_talk != 1 or self.comparison_mode:
1684
+ if attention_mask is None:
1685
+ base_attention_mask = torch.triu(torch.ones(seq_len, seq_len), diagonal=0).to(input_ids.device)
1686
+ base_attention_mask = base_attention_mask.view(1, 1, seq_len, seq_len)
1687
+ base_attention_mask = base_attention_mask.repeat(input_ids.shape[0], 1, 1, 1)
1688
+ attention_mask = base_attention_mask
1689
+ # breakpoint()
1690
+ elif attention_mask.dim() == 2:
1691
+ if seq_len + past_key_values_length != attention_mask.shape[-1]:
1692
+ # breakpoint()
1693
+ attention_mask = torch.cat(
1694
+ [torch.ones((attention_mask.shape[0], past_key_values_length), dtype=attention_mask.dtype, device=attention_mask.device), attention_mask],
1695
+ dim=-1
1696
+ )
1697
+ # # if the attention mask
1698
+ attention_mask = _prepare_4d_causal_attention_mask(
1699
+ attention_mask,
1700
+ (batch_size, seq_len),
1701
+ inputs_embeds,
1702
+ past_key_values_length,
1703
+ sliding_window=self.config.sliding_window,
1704
+ )
1705
+
1706
+ outputs = self.model(
1707
+ # input_ids=input_ids,
1708
+ attention_mask=attention_mask,
1709
+ position_ids=position_ids,
1710
+ past_key_values=past_key_values,
1711
+ inputs_embeds=inputs_embeds,
1712
+ use_cache=use_cache,
1713
+ output_attentions=output_attentions,
1714
+ output_hidden_states=output_hidden_states,
1715
+ # output_router_logits=output_router_logits,
1716
+ return_dict=return_dict,
1717
+ )
1718
+
1719
+ prev_hidden_states = hidden_states
1720
+ hidden_states = outputs[0]
1721
+ prev_rm_logits = rm_logits # for policy gradient
1722
+ prev_rm_tokens = cur_rm_tokens # for policy gradient
1723
+
1724
+ if ahead_idx == 0:
1725
+ hidden_states_lm = hidden_states
1726
+ logits = self.lm_head(hidden_states_lm)
1727
+ base_hidden_states = hidden_states.clone()
1728
+ initial_loss_logits = logits.clone()
1729
+ if self.optimize_lm_head_only_at_start or self.optimize_model_only_at_start:
1730
+ logits = logits.detach()
1731
+ base_hidden_states = base_hidden_states.detach()
1732
+ if self.optimize_model_only_at_start:
1733
+ hidden_states = hidden_states.detach()
1734
+ base_logits = logits.clone()
1735
+ else:
1736
+ talk_hidden_states = hidden_states
1737
+ if self.merged_lm_and_talk_heads:
1738
+ assert self.no_residual
1739
+ residual_logits = self.lm_head(hidden_states)
1740
+ talk_hidden_states = hidden_states
1741
+ else:
1742
+ if ahead_idx > self.n_ahead - 1:
1743
+ cur_base_hidden = torch.cat([
1744
+ base_hidden_states[..., ahead_idx - self.n_ahead + 1:, :],
1745
+ base_hidden_states[..., :ahead_idx - self.n_ahead + 1, :]
1746
+ ], dim=-2)
1747
+ else:
1748
+ cur_base_hidden = base_hidden_states
1749
+
1750
+ if self.use_concat_talk_head:
1751
+ # concatenate the hidden states with the original hidden states
1752
+ head_input_hidden_states = torch.cat([cur_base_hidden, talk_hidden_states], dim=-1)
1753
+ else:
1754
+ head_input_hidden_states = talk_hidden_states
1755
+
1756
+ residual_logits = self.talk_head[0](head_input_hidden_states)
1757
+ if self.use_shallow_talk:
1758
+ residual_logits = apply_head(self.lm_head, residual_logits, detach=self.optimize_lm_head_only_at_start)
1759
+ residual_logits = residual_logits.to(logits.device)
1760
+ if self.use_weighted_talk_head:
1761
+ # combine the cur_base_hidden with the talk_hidden_states according to the weighted head
1762
+ residual_logits = cur_base_hidden * (1 - residual_logits) + talk_hidden_states * residual_logits
1763
+ residual_logits = apply_head(self.lm_head, residual_logits, detach=self.optimize_lm_head_only_at_start)
1764
+
1765
+ assert sum([self.cumulative_residual, self.clever_residual, self.skip_residual, self.no_residual]) == 1
1766
+ if self.clever_residual:
1767
+ if ahead_idx >= self.n_ahead - 1:
1768
+ # get the logits shifted according to the current talk ahead
1769
+ cur_base_logits = torch.cat([
1770
+ base_logits[..., ahead_idx - self.n_ahead + 1:, :],
1771
+ base_logits[..., :ahead_idx - self.n_ahead + 1, :]
1772
+ ], dim=-2)
1773
+ if self.optimize_lm_head_only_at_start:
1774
+ cur_base_logits = cur_base_logits.detach()
1775
+ logits = cur_base_logits + residual_logits
1776
+ else:
1777
+ logits += residual_logits / self.n_ahead
1778
+ elif self.cumulative_residual:
1779
+ if self.residual_talk_head:
1780
+ if ahead_idx < self.n_ahead:
1781
+ logits += residual_logits
1782
+ else:
1783
+ # get the logits shifted according to the current talk ahead
1784
+ cur_base_logits = torch.cat([
1785
+ base_logits[..., ahead_idx - self.n_ahead + 1:, :],
1786
+ base_logits[..., :ahead_idx - self.n_ahead + 1, :]
1787
+ ], dim=-2)
1788
+ if self.optimize_lm_head_only_at_start:
1789
+ cur_base_logits = cur_base_logits.detach()
1790
+ logits = cur_base_logits + residual_logits
1791
+ else:
1792
+ if ahead_idx < self.n_ahead:
1793
+ logits += residual_logits
1794
+ else:
1795
+ logits = residual_logits
1796
+ elif self.skip_residual:
1797
+ if ahead_idx >= self.n_ahead:
1798
+ # get the logits shifted according to the current talk ahead
1799
+ cur_base_logits = torch.cat([
1800
+ base_logits[..., ahead_idx - self.n_ahead + 1:, :],
1801
+ base_logits[..., :ahead_idx - self.n_ahead + 1, :]
1802
+ ], dim=-2)
1803
+ if self.optimize_lm_head_only_at_start:
1804
+ cur_base_logits = cur_base_logits.detach()
1805
+ logits = cur_base_logits
1806
+ elif self.no_residual:
1807
+ logits = residual_logits
1808
+ else:
1809
+ logits = base_logits + residual_logits
1810
+
1811
+ attempted = False
1812
+ talk_loss_list = []
1813
+ if self.original_mode or (self.n_ahead == 1) or (self.comparison_mode and ahead_idx == 0):# or (self.optimize_lm_head_only_at_start and ahead_idx == 0):
1814
+ loss = None
1815
+ attempted = True
1816
+
1817
+ if labels is not None:
1818
+ for shift_amount in range(self.n_ahead_talk):
1819
+ # Shift so that tokens < n predict n
1820
+ # ab[cde]f
1821
+ # abc[def]
1822
+ if ahead_idx == 0 and self.optimize_lm_head_only_at_start:
1823
+ loss_logits = initial_loss_logits
1824
+ else:
1825
+ loss_logits = logits
1826
+ shift_logits = loss_logits[..., shift_amount:-1, :].contiguous()
1827
+ shift_labels = labels[..., 1 + shift_amount:].contiguous()
1828
+ # Flatten the tokens
1829
+ loss_fct = CrossEntropyLoss(reduction="none")
1830
+ # print("Shift logits before:", shift_logits)
1831
+ shift_logits = shift_logits.view(-1, self.config.vocab_size)
1832
+ shift_labels = shift_labels.view(-1).clone()
1833
+ # print("shift logits after:", shift_logits)
1834
+ # Enable model parallelism
1835
+ shift_labels[shift_labels == self.tokenizer.pad_token_id] = -100
1836
+ shift_labels = shift_labels.to(shift_logits.device)
1837
+ loss = loss_fct(shift_logits, shift_labels)
1838
+ if not self.comparison_mode and not (self.optimize_lm_head_only_at_start and (self.n_ahead + self.n_ahead_talk > 2)) or self.original_mode:
1839
+ loss_list.append(loss)
1840
+ talk_loss_list.append(nonzero_mean(loss).detach())
1841
+
1842
+ if not attempted or self.comparison_mode:
1843
+ rm_hidden_states = hidden_states
1844
+ # print("Magnitude of RM hidden states before RM head", rm_hidden_states.norm())
1845
+ rm_logits = apply_head(self.lm_head, rm_hidden_states, detach=self.optimize_lm_head_only_at_start)
1846
+
1847
+ # don't allow it to predict the thinking token
1848
+ if self.tokenizer_has_start_thought_token:
1849
+ rm_logits[..., self.start_token_id] = -1e10
1850
+ if self.tokenizer_has_end_thought_token:
1851
+ rm_logits[..., self.end_token_id] = -1e10
1852
+ probabilities = rm_logits
1853
+ if probabilities_2d is not None:
1854
+ prev_probabilities_2d = probabilities_2d.clone()
1855
+ probabilities_2d = probabilities.view(-1, probabilities.size(-1))
1856
+
1857
+ did_skip_sampling = skip_sampling
1858
+ skip_sampling = False
1859
+ if ahead_idx == 0 and self.use_start_thought_token:
1860
+ override_token = self.start_token_id
1861
+ elif self.use_thought_prefix and ahead_idx < self.tokenized_thought_prefix.shape[-1]:
1862
+ override_token = self.tokenized_thought_prefix[..., ahead_idx]
1863
+ elif ahead_idx == self.n_ahead - 2 and self.use_end_thought_token:
1864
+ override_token = self.end_token_id
1865
+ else:
1866
+ override_token = None
1867
+ if override_token is not None and self.n_ahead > 1:
1868
+ # always start with the start token
1869
+ probabilities_2d = torch.zeros_like(probabilities_2d)
1870
+ probabilities_2d[:, override_token] = 1.0
1871
+ skip_sampling = True
1872
+ elif ahead_idx >= self.n_ahead - 1:
1873
+ if labels is not None: # we're in the talk phase
1874
+ cur_talk_n = ahead_idx - (self.n_ahead - 1) + 1
1875
+ # print("Setting rm to labels", cur_talk_n, "during", ahead_idx)
1876
+ shift_labels = labels[..., cur_talk_n:].contiguous().to(probabilities_2d.device)
1877
+ padding = torch.full_like(
1878
+ labels[..., :cur_talk_n],
1879
+ self.tokenizer.pad_token_id,
1880
+ dtype=torch.long,
1881
+ device=shift_labels.device
1882
+ )
1883
+ new_rm_tokens = torch.cat(
1884
+ [shift_labels, padding],
1885
+ dim=-1
1886
+ )
1887
+
1888
+ # print((new_rm_tokens > self.vocab_size - 1).any().item())
1889
+ new_rm_tokens = torch.clamp(new_rm_tokens, 0, self.vocab_size - 1)
1890
+
1891
+ # Now safely convert rm tokens to one-hot
1892
+ probabilities_2d = F.one_hot(new_rm_tokens, num_classes=self.vocab_size).reshape(-1, self.vocab_size).to(probabilities_2d.dtype)
1893
+ else:
1894
+ continue
1895
+ temperature = self.gumbel_temperature if self.training else 0.001
1896
+ prev_sample_probs = sample_probs
1897
+ sample_probs = probabilities_2d
1898
+ if ahead_idx < self.n_ahead - 1 and not skip_sampling:
1899
+ probabilities_2d = F.gumbel_softmax(sample_probs, tau=temperature, hard=True, dim=-1)
1900
+ if self.gumbel_detach:
1901
+ probabilities_2d = probabilities_2d.detach()
1902
+ sampled_token_history.append(probabilities_2d.argmax(dim=-1).detach().cpu())
1903
+ # convert rm logits directly to embeddings
1904
+ contains_start = self.use_start_thought_token and (probabilities_2d[..., self.start_token_id].sum() > 0)
1905
+ contains_end = self.use_end_thought_token and (probabilities_2d[..., self.end_token_id].sum() > 0)
1906
+ contains_thought = contains_start or contains_end
1907
+
1908
+
1909
+ if not contains_thought:
1910
+ with torch.set_grad_enabled(not self.train_only_thinking_embedding):
1911
+ inputs_embeds = probabilities_2d @ (self.model.embed_tokens.weight.to(probabilities.device).to(probabilities.dtype) * temperature)
1912
+ else:
1913
+ thought_id = self.start_token_id if contains_start else self.end_token_id
1914
+ cur_thought_embedding = start_embedding if contains_start else end_embedding
1915
+ if self.use_reparam_for_thought_embeddings:
1916
+ inputs_embeds = torch.randn(batch_size, seq_len, self.model.config.hidden_size, device=input_ids.device, dtype=cur_thought_embedding.dtype)
1917
+ inputs_embeds = inputs_embeds * torch.exp(cur_thought_embedding[1]) + cur_thought_embedding[0]
1918
+ if contains_start:
1919
+ sampled_start = inputs_embeds.clone().detach()
1920
+ else:
1921
+ sampled_end = inputs_embeds.clone().detach()
1922
+ else:
1923
+ inputs_embeds = cur_thought_embedding.unsqueeze(0).repeat(batch_size, seq_len, 1)
1924
+ inputs_embeds = inputs_embeds.view(probabilities.size(0), probabilities.size(1), -1).to(self.model.embed_tokens.weight.dtype)
1925
+ inputs_embeds = inputs_embeds.view(probabilities.size(0), probabilities.size(1), -1).to(self.model.embed_tokens.weight.dtype)
1926
+
1927
+ # Predict the usefulness of thinking at each token position
1928
+ thinking_usefulness = self.thinking_usefulness_head(hidden_states).squeeze(-1)
1929
+
1930
+ # Apply a threshold to decide where to generate thoughts
1931
+ generate_thought_mask = thinking_usefulness > self.thinking_threshold
1932
+
1933
+ # Compute the regularization loss for thinking usefulness prediction
1934
+ thinking_usefulness_loss = torch.mean(thinking_usefulness * (1 - generate_thought_mask.float()))
1935
+
1936
+ # Add the regularization loss to the total loss
1937
+ if loss is not None:
1938
+ loss = loss + self.thinking_usefulness_loss_weight * thinking_usefulness_loss
1939
+ else:
1940
+ loss = self.thinking_usefulness_loss_weight * thinking_usefulness_loss
1941
+
1942
+
1943
+ if len(attention_mask.shape) == 2:
1944
+ breakpoint()
1945
+ else:
1946
+ original_attention = attention_mask[..., :attention_mask.shape[-2]]
1947
+ if self.use_upper_triangular:
1948
+ new_attention = original_attention
1949
+ else:
1950
+ original_attention = original_attention == attention_mask.max()
1951
+ # because eye isn't implemented for BF16, we need to handle the case
1952
+ if not attention_mask.dtype == torch.bfloat16:
1953
+ new_attention = torch.eye(
1954
+ seq_len, dtype=attention_mask.dtype, device=attention_mask.device
1955
+ )
1956
+ else:
1957
+ new_attention = torch.eye(
1958
+ seq_len, dtype=torch.float32, device=attention_mask.device
1959
+ ).to(attention_mask.dtype)
1960
+
1961
+ new_attention = new_attention.view(1, 1, seq_len, seq_len).repeat(input_ids.shape[0], 1, 1, 1)
1962
+ new_attention = new_attention * original_attention
1963
+ new_attention[new_attention == 0] = attention_mask.min()
1964
+ new_attention[new_attention == 1] = attention_mask.max()
1965
+ attention_mask = torch.cat([attention_mask, new_attention], dim=-1)
1966
+ past_key_values = outputs.past_key_values
1967
+ position_ids = position_ids + 1
1968
+
1969
+ if labels is not None and (self.n_ahead > 1 or not self.base_original_mode):
1970
+ # Shift so that tokens < n predict n
1971
+ # logits: abcdef -> bcdef? -> cdef??
1972
+ # labels: abcdef -> ?bcdef -> ??cdef
1973
+ if ahead_idx == 0 and self.optimize_lm_head_only_at_start:
1974
+ loss_logits = initial_loss_logits
1975
+ else:
1976
+ loss_logits = logits
1977
+ shift_idx = 1 + max(0, ahead_idx - (self.n_ahead - 1))
1978
+ shift_logits = loss_logits[..., :-shift_idx, :].contiguous()
1979
+ shift_labels = labels[..., shift_idx:].contiguous()
1980
+ # Flatten the tokens
1981
+ loss_fct = CrossEntropyLoss(reduction="none")
1982
+ shift_logits = shift_logits.view(-1, self.config.vocab_size)
1983
+ shift_labels = shift_labels.view(-1)
1984
+ # Enable model parallelism
1985
+ shift_labels = shift_labels.to(shift_logits.device)
1986
+ # if shift_labels.min() == self.tokenizer.pad_token_id:
1987
+ shift_labels = torch.where(shift_labels == self.tokenizer.pad_token_id, -100, shift_labels)
1988
+ unreduced_loss = loss_fct(shift_logits, shift_labels)
1989
+ # print("Loss:", unreduced_loss.item()) # Print the loss before checking for NaN values
1990
+ if torch.any(unreduced_loss != unreduced_loss):
1991
+ # pdb.set_trace()
1992
+ raise ValueError("NaN loss")
1993
+ unreduced_loss = unreduced_loss.reshape(logits.shape[0], -1)
1994
+ loss_list.append(unreduced_loss)
1995
+
1996
+
1997
+ if self.use_policy_loss and ahead_idx > 0 and (ahead_idx > 1 or not self.use_start_thought_token):
1998
+ # we treat the change in loss as the reward
1999
+ previous_loss = loss_list[-2]
2000
+ # for example, suppose n_ahead = 3 and n_ahead_talk = 2
2001
+ # note that we end at self.n_ahead + self.n_ahead_talk - 2
2002
+ # in this case, 5 - 2 = 3, so we end at ahead_idx = 3
2003
+ # we also predict the next token at ahead_idx = 2
2004
+ # when we get to ahead_idx = 2, we predict ahead
2005
+ # so we shift by 1
2006
+ # note that this is ahead_idx = n_ahead - 1
2007
+ # when we get to ahead_idx = 3, we predict ahead
2008
+ # so we shift by 2
2009
+ # note that this is ahead_idx = n_ahead
2010
+ if ahead_idx < self.n_ahead - 1:
2011
+ shift_amount = 0
2012
+ reward_scale = 1.0
2013
+ original_dqn_reward = torch.sign(previous_loss - unreduced_loss).detach() * reward_scale
2014
+ if self.first_and_last_mode:
2015
+ original_dqn_reward = original_dqn_reward * 0.0
2016
+ else:
2017
+ # logits vs cur_policy_shift_logits
2018
+ # let's look at rm_logits and prev_rm_logits
2019
+ shift_amount = max(0, ahead_idx - (self.n_ahead - 1))
2020
+ # let's say shift_amount = 2
2021
+ # abcdefg -> bcdefg? -> cdefg??
2022
+ # logits = [a b]c d e f[g]
2023
+ # labels = [a b c]d e f g
2024
+ cur_policy_shift_logits = initial_loss_logits[..., shift_amount:-1, :].contiguous().detach()
2025
+ cur_policy_shift_labels = labels[..., 1 + shift_amount:].contiguous()
2026
+ # Flatten the tokens
2027
+ cur_policy_loss_fct = CrossEntropyLoss(reduction="none")
2028
+ cur_policy_shift_logits = cur_policy_shift_logits.view(-1, self.config.vocab_size)
2029
+ cur_policy_shift_labels = cur_policy_shift_labels.view(-1).clone()
2030
+ # Enable model parallelism
2031
+ cur_policy_shift_labels[cur_policy_shift_labels == self.tokenizer.pad_token_id] = -100
2032
+ cur_policy_shift_labels = cur_policy_shift_labels.to(cur_policy_shift_labels.device)
2033
+ cur_policy_reward_base_loss = loss_fct(
2034
+ cur_policy_shift_logits, cur_policy_shift_labels.to(cur_policy_shift_logits.device)
2035
+ ).reshape(logits.shape[0], -1)
2036
+ original_dqn_reward = cur_policy_reward_base_loss.detach() - unreduced_loss
2037
+
2038
+ if not did_skip_sampling:
2039
+ nonzero_indices = prev_probabilities_2d.nonzero()
2040
+ action_loglikelihoods = F.log_softmax(prev_sample_probs / self.reinforce_temperature, dim=-1)[nonzero_indices[:, 0], nonzero_indices[:, 1]]
2041
+ action_loglikelihoods_2d = action_loglikelihoods.reshape(batch_size, -1)[:, :-1 - shift_amount]
2042
+ action_loglikelihoods_list.append(action_loglikelihoods_2d)
2043
+ if policy_reward is None:
2044
+ policy_reward = original_dqn_reward[:, :-(self.n_ahead_talk - shift_amount)]
2045
+ else:
2046
+ if self.n_ahead_talk > shift_amount:
2047
+ added_reward = original_dqn_reward[:, :-(self.n_ahead_talk - shift_amount)]
2048
+ else:
2049
+ added_reward = original_dqn_reward
2050
+ policy_reward += added_reward
2051
+
2052
+ for action_loglikelihoods_2d in action_loglikelihoods_list:
2053
+ train_policy_reward = policy_reward
2054
+
2055
+ # discard rewards below the mean
2056
+ if self.trice_mode and self.n_passes > 1:
2057
+ batched_policy_reward = train_policy_reward.reshape(-1, self.n_passes, train_policy_reward.shape[-1])
2058
+ # average over the passes
2059
+ train_policy_reward = batched_policy_reward - batched_policy_reward.mean(dim=1, keepdim=True)
2060
+ train_policy_reward = train_policy_reward.reshape(-1, train_policy_reward.shape[-1])
2061
+
2062
+ if self.subtract_mean_reward:
2063
+ train_policy_reward = train_policy_reward - train_policy_reward.mean()
2064
+ if self.remove_negative_rewards:
2065
+ fixed_policy_reward = train_policy_reward.detach().clamp(min=0)
2066
+ else:
2067
+ fixed_policy_reward = train_policy_reward.detach()
2068
+
2069
+ # Normalize rewards
2070
+ fixed_policy_reward = (fixed_policy_reward - fixed_policy_reward.mean()) / (fixed_policy_reward.std() + 1e-8)
2071
+ actor_loss = -fixed_policy_reward * action_loglikelihoods_2d[:, :policy_reward.shape[-1]].to(policy_reward.device)
2072
+ if action_loglikelihoods_2d.mean() < -1e4 and not self.use_policy_loss_just_for_thoughts:
2073
+ # This will only happen when we force the next token to be the end of thought token
2074
+ break
2075
+ dqn_loss_list.append(actor_loss.mean())
2076
+
2077
+ if loss_list:
2078
+ if self.first_and_last_mode:
2079
+ loss = sum(
2080
+ self.loss_mean(loss_list[-(i + 1)]) for i in range(self.n_ahead_talk)
2081
+ ) * (1 - self.original_loss_weight) / self.n_ahead_talk
2082
+ loss = loss + self.loss_mean(loss_list[0]) * self.original_loss_weight
2083
+ # Let's NaN out the others
2084
+ # e.g. if n_ahead_talk = 2 and the list is 5 long, we want to NaN out 1, 2 but keep 0, 3, 4
2085
+ for i in range(1, len(loss_list) - self.n_ahead_talk):
2086
+ loss_list[i] = loss_list[i] * math.nan
2087
+ elif self.first_only:
2088
+ loss = self.loss_mean(loss_list[0])
2089
+ elif self.final_only_mode:
2090
+ loss = sum(
2091
+ self.loss_mean(loss_list[-i]) for i in range(1, self.n_ahead_talk + 1)
2092
+ ) / self.n_ahead_talk
2093
+ else:
2094
+ loss = None
2095
+ for i in range(len(loss_list)):
2096
+ cur_loss = self.loss_mean(loss_list[i])
2097
+ if loss is not None:
2098
+ loss = loss + cur_loss.to(loss.device)
2099
+ else:
2100
+ loss = cur_loss
2101
+ loss = loss / len(loss_list)
2102
+ loss = loss + thinking_usefulness_loss
2103
+
2104
+ base_loss_scale = 0.6
2105
+ policy_loss_scale = 0.03
2106
+
2107
+ loss = loss * base_loss_scale
2108
+
2109
+ if dqn_loss_list:
2110
+ dqn_loss = sum(dqn_loss_list) / len(dqn_loss_list)
2111
+ if self.include_policy_loss:
2112
+ if loss is not None:
2113
+ loss += dqn_loss * policy_loss_scale
2114
+ else:
2115
+ loss = dqn_loss * self.policy_loss_beta
2116
+
2117
+ if not return_dict:
2118
+ output = (logits,) + outputs[1:]
2119
+ return (loss,) + output if loss is not None else output
2120
+
2121
+ base_log_dict = {
2122
+ f"loss_{i}": nonzero_mean(loss_list[i]) for i in range(len(loss_list))
2123
+ }
2124
+
2125
+ if loss is not None:
2126
+ base_log_dict["loss_train"] = loss.item()
2127
+
2128
+ if not self.training:
2129
+ self.n_ahead_talk = n_ahead_talk_to_restore
2130
+ self.n_passes = n_passes_to_restore
2131
+
2132
+ del start_embedding
2133
+ del end_embedding
2134
+ torch.cuda.empty_cache()
2135
+
2136
+
2137
+ return CausalLMOutputWithPast(
2138
+ loss=loss if loss is not None else None,
2139
+ logits=(rm_logits if self.n_ahead > 1 else logits) if not self.output_logits_at_the_end else logits,
2140
+ past_key_values=outputs.past_key_values,
2141
+ hidden_states=outputs.hidden_states,
2142
+ attentions=outputs.attentions,
2143
+ )
2144
+
2145
+
2146
+
2147
+ def prepare_inputs_for_generation(
2148
+ self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, **kwargs
2149
+ ):
2150
+ # Omit tokens covered by past_key_values
2151
+ if past_key_values is not None:
2152
+ if isinstance(past_key_values, Cache):
2153
+ cache_length = past_key_values.get_seq_length()
2154
+ past_length = past_key_values.seen_tokens
2155
+ max_cache_length = past_key_values.get_max_length()
2156
+ else:
2157
+ cache_length = past_length = past_key_values[0][0].shape[2]
2158
+ max_cache_length = None
2159
+
2160
+ # Keep only the unprocessed tokens:
2161
+ # 1 - If the length of the attention_mask exceeds the length of input_ids, then we are in a setting where
2162
+ # some of the inputs are exclusively passed as part of the cache (e.g. when passing inputs_embeds as
2163
+ # input)
2164
+ if attention_mask is not None and attention_mask.shape[1] > input_ids.shape[1]:
2165
+ input_ids = input_ids[:, -(attention_mask.shape[1] - past_length) :]
2166
+ # 2 - If the past_length is smaller than input_ids', then input_ids holds all input tokens. We can discard
2167
+ # input_ids based on the past_length.
2168
+ elif past_length < input_ids.shape[1]:
2169
+ input_ids = input_ids[:, past_length:]
2170
+ # 3 - Otherwise (past_length >= input_ids.shape[1]), let's assume input_ids only has unprocessed tokens.
2171
+
2172
+ # If we are about to go beyond the maximum cache length, we need to crop the input attention mask.
2173
+ if (
2174
+ max_cache_length is not None
2175
+ and attention_mask is not None
2176
+ and cache_length + input_ids.shape[1] > max_cache_length
2177
+ ):
2178
+ attention_mask = attention_mask[:, -max_cache_length:]
2179
+
2180
+ position_ids = kwargs.get("position_ids", None)
2181
+ if attention_mask is not None and position_ids is None:
2182
+ # create position_ids on the fly for batch generation
2183
+ position_ids = attention_mask.long().cumsum(-1) - 1
2184
+ position_ids.masked_fill_(attention_mask == 0, 1)
2185
+ if past_key_values:
2186
+ position_ids = position_ids[:, -input_ids.shape[1] :]
2187
+
2188
+ # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
2189
+ if inputs_embeds is not None and past_key_values is None:
2190
+ model_inputs = {"inputs_embeds": inputs_embeds}
2191
+ else:
2192
+ model_inputs = {"input_ids": input_ids}
2193
+
2194
+ model_inputs.update(
2195
+ {
2196
+ "position_ids": position_ids,
2197
+ "past_key_values": past_key_values,
2198
+ "use_cache": kwargs.get("use_cache"),
2199
+ "attention_mask": attention_mask,
2200
+ }
2201
+ )
2202
+ return model_inputs
2203
+
2204
+ @staticmethod
2205
+ def _reorder_cache(past_key_values, beam_idx):
2206
+ reordered_past = ()
2207
+ for layer_past in past_key_values:
2208
+ reordered_past += (
2209
+ tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past),
2210
+ )
2211
+ return reordered_past
2212
+
2213
+
2214
+
2215
+
2216
+ @add_start_docstrings(
2217
+ """
2218
+ The Quiet Model transformer with a sequence classification head on top (linear layer).
2219
+ [`QuietForSequenceClassification`] uses the last token in order to do the classification, as other causal models
2220
+ (e.g. GPT-2) do.
2221
+ Since it does classification on the last token, it requires to know the position of the last token. If a
2222
+ `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If
2223
+ no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the
2224
+ padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in
2225
+ each row of the batch).
2226
+ """,
2227
+ QUIET_START_DOCSTRING,
2228
+ )
2229
+ # Copied from transformers.models.llama.modeling_llama.LlamaForSequenceClassification with Llama->Quiet, LLAMA->QUIET
2230
+ class QuietForSequenceClassification(QuietPreTrainedModel):
2231
+ def __init__(self, config):
2232
+ super().__init__(config)
2233
+ self.num_labels = config.num_labels
2234
+ self.model = QuietModel(config)
2235
+ self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False)
2236
+
2237
+ # Initialize weights and apply final processing
2238
+ self.post_init()
2239
+
2240
+ def get_input_embeddings(self):
2241
+ return self.model.embed_tokens
2242
+
2243
+ def set_input_embeddings(self, value):
2244
+ self.model.embed_tokens = value
2245
+
2246
+ @add_start_docstrings_to_model_forward(QUIET_INPUTS_DOCSTRING)
2247
+ def forward(
2248
+ self,
2249
+ input_ids: torch.LongTensor = None,
2250
+ attention_mask: Optional[torch.Tensor] = None,
2251
+ position_ids: Optional[torch.LongTensor] = None,
2252
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
2253
+ inputs_embeds: Optional[torch.FloatTensor] = None,
2254
+ labels: Optional[torch.LongTensor] = None,
2255
+ use_cache: Optional[bool] = None,
2256
+ output_attentions: Optional[bool] = None,
2257
+ output_hidden_states: Optional[bool] = None,
2258
+ return_dict: Optional[bool] = None,
2259
+ ) -> Union[Tuple, SequenceClassifierOutputWithPast]:
2260
+ r"""
2261
+ labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
2262
+ Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
2263
+ config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
2264
+ `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
2265
+ """
2266
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
2267
+
2268
+ transformer_outputs = self.model(
2269
+ input_ids,
2270
+ attention_mask=attention_mask,
2271
+ position_ids=position_ids,
2272
+ past_key_values=past_key_values,
2273
+ inputs_embeds=inputs_embeds,
2274
+ use_cache=use_cache,
2275
+ output_attentions=output_attentions,
2276
+ output_hidden_states=output_hidden_states,
2277
+ return_dict=return_dict,
2278
+ )
2279
+ hidden_states = transformer_outputs[0]
2280
+ logits = self.score(hidden_states)
2281
+
2282
+ if input_ids is not None:
2283
+ batch_size = input_ids.shape[0]
2284
+ else:
2285
+ batch_size = inputs_embeds.shape[0]
2286
+
2287
+ if self.config.pad_token_id is None and batch_size != 1:
2288
+ raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.")
2289
+ if self.config.pad_token_id is None:
2290
+ sequence_lengths = -1
2291
+ else:
2292
+ if input_ids is not None:
2293
+ # if no pad token found, use modulo instead of reverse indexing for ONNX compatibility
2294
+ sequence_lengths = torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1
2295
+ sequence_lengths = sequence_lengths % input_ids.shape[-1]
2296
+ sequence_lengths = sequence_lengths.to(logits.device)
2297
+ else:
2298
+ sequence_lengths = -1
2299
+
2300
+ pooled_logits = logits[torch.arange(batch_size, device=logits.device), sequence_lengths]
2301
+
2302
+ loss = None
2303
+ if labels is not None:
2304
+ labels = labels.to(logits.device)
2305
+ if self.config.problem_type is None:
2306
+ if self.num_labels == 1:
2307
+ self.config.problem_type = "regression"
2308
+ elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
2309
+ self.config.problem_type = "single_label_classification"
2310
+ else:
2311
+ self.config.problem_type = "multi_label_classification"
2312
+
2313
+ if self.config.problem_type == "regression":
2314
+ loss_fct = MSELoss()
2315
+ if self.num_labels == 1:
2316
+ loss = loss_fct(pooled_logits.squeeze(), labels.squeeze())
2317
+ else:
2318
+ loss = loss_fct(pooled_logits, labels)
2319
+ elif self.config.problem_type == "single_label_classification":
2320
+ loss_fct = CrossEntropyLoss()
2321
+ loss = loss_fct(pooled_logits.view(-1, self.num_labels), labels.view(-1))
2322
+ elif self.config.problem_type == "multi_label_classification":
2323
+ loss_fct = BCEWithLogitsLoss()
2324
+ loss = loss_fct(pooled_logits, labels)
2325
+ if not return_dict:
2326
+ output = (pooled_logits,) + transformer_outputs[1:]
2327
+ return ((loss,) + output) if loss is not None else output
2328
+
2329
+ return SequenceClassifierOutputWithPast(
2330
+ loss=loss,
2331
+ logits=pooled_logits,
2332
+ past_key_values=transformer_outputs.past_key_values,
2333
+ hidden_states=transformer_outputs.hidden_states,
2334
+ attentions=transformer_outputs.attentions,
2335
+ )
optuna.py ADDED
@@ -0,0 +1,153 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import optuna
2
+ import torch
3
+ import random
4
+ from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments
5
+ from datasets import load_dataset
6
+ from trl import SFTTrainer
7
+ import time
8
+
9
+ # Set random seed for reproducibility
10
+ random_seed = 42
11
+ torch.manual_seed(random_seed)
12
+ random.seed(random_seed)
13
+
14
+ # Load dataset
15
+ dataset = load_dataset("tatsu-lab/alpaca", split="train")
16
+
17
+
18
+ def chatml_format(example):
19
+ """Format the dataset for training, accounting for empty columns."""
20
+ return {
21
+ "instruction": example['instruction'] if 'instruction' in example else " \n",
22
+ "input": example['input'] if 'input' in example else " \n",
23
+ "system": example['system'] if 'system' in example else " \n",
24
+ "output": example['output'] if 'output' in example else " \n",
25
+ }
26
+
27
+
28
+ # Format dataset
29
+ dataset = dataset.map(chatml_format, remove_columns=dataset.column_names)
30
+
31
+ # Define the model initialization function
32
+ def model_init(trial=None):
33
+ original = False
34
+ params = {}
35
+ if trial is not None:
36
+ n_ahead = 1
37
+ n_ahead_talk = 1
38
+ n_passes = 1
39
+ gumbel_temperature = 1
40
+ use_start_thought_token = True
41
+ use_end_thought_token = True
42
+ include_policy_loss = True
43
+ gumbel_detach = True
44
+ merged_talk_heads = True
45
+ residual_think_head = False
46
+ optimize_lm_head_only_at_start = False
47
+
48
+ model_id = "Crystalcareai/Quiet-Star-Custom"
49
+ tokenizer_id = model_id
50
+
51
+ model = AutoModelForCausalLM.from_pretrained(
52
+ model_id,
53
+ torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
54
+ max_thoughts=n_ahead + n_ahead_talk + 1,
55
+ merged_talk_heads=merged_talk_heads,
56
+ merged_lm_and_talk_heads=False,
57
+ merged_lm_and_think_heads=True,
58
+ use_concat_talk_head=True,
59
+ use_shallow_think=True,
60
+ use_shallow_talk=False,
61
+ use_complex_think_head=False,
62
+ use_complex_talk_head=True,
63
+ use_weighted_talk_head=True,
64
+ trust_remote_code=True,
65
+ device_map="auto",
66
+ )
67
+
68
+ tokenizer = AutoTokenizer.from_pretrained(tokenizer_id, truncation=True, padding="left")
69
+ tokenizer.pad_token_id = tokenizer.eos_token_id
70
+
71
+ special_tokens_to_add = []
72
+ if model.use_start_thought_token:
73
+ special_tokens_to_add.append("<|startthought|>")
74
+ if model.use_end_thought_token:
75
+ special_tokens_to_add.append("<|endthought|>")
76
+ if special_tokens_to_add:
77
+ tokenizer.add_special_tokens({"additional_special_tokens": special_tokens_to_add})
78
+ model.resize_token_embeddings(len(tokenizer))
79
+ model.tokenizer = tokenizer
80
+ for name, module in model.named_modules():
81
+ if "embed" in name:
82
+ print(module, flush=True)
83
+
84
+ model.gumbel_detach = gumbel_detach
85
+ model.include_policy_loss = include_policy_loss
86
+ model.use_end_thought_token = use_end_thought_token
87
+ model.use_start_thought_token = use_start_thought_token
88
+ model.n_ahead = n_ahead
89
+ model.n_ahead_talk = n_ahead_talk
90
+ model.n_passes = n_passes
91
+ model.residual_think_head = residual_think_head
92
+ model.gumbel_temperature = gumbel_temperature
93
+ model.original_mode = original
94
+ model.config_params = params
95
+ model.run_start = int(time.time())
96
+ model.train()
97
+ return model
98
+
99
+ # Define the objective function for Optuna
100
+ # Define the objective function for Optuna
101
+ def objective(trial):
102
+ # Hyperparameters to be optimized
103
+ learning_rate = trial.suggest_float("learning_rate", 1e-07, 1e-06, log=True)
104
+ max_grad_norm = trial.suggest_float("max_grad_norm", 0.3, 1.0)
105
+ warmup_steps = trial.suggest_int("warmup_steps", 0, 20)
106
+ gradient_accumulation_steps = trial.suggest_int("gradient_accumulation_steps", 4, 8)
107
+
108
+ model = model_init(trial)
109
+
110
+ training_args = TrainingArguments(
111
+ output_dir="./out",
112
+ num_train_epochs=3,
113
+ max_steps=30,
114
+ per_device_train_batch_size=1,
115
+ logging_steps=1,
116
+ optim="lion_32bit",
117
+ save_strategy="steps",
118
+ save_steps=3000,
119
+ gradient_accumulation_steps=gradient_accumulation_steps,
120
+ learning_rate=learning_rate,
121
+ max_grad_norm=max_grad_norm,
122
+ warmup_steps=warmup_steps,
123
+ lr_scheduler_type="cosine",
124
+ report_to="none" # Disable reporting to avoid errors related to WandB in this context
125
+ )
126
+
127
+ trainer = SFTTrainer(
128
+ args=training_args,
129
+ train_dataset=dataset,
130
+ model=model,
131
+ tokenizer=model.tokenizer,
132
+ max_seq_length=1024,
133
+ dataset_text_field="output",
134
+ )
135
+
136
+ # Train the model and get the training loss
137
+ train_result = trainer.train()
138
+ loss = train_result.training_loss
139
+
140
+ return loss
141
+
142
+
143
+ # Create a study and optimize
144
+ study = optuna.create_study(storage="sqlite:///db.sqlite3")
145
+ study.optimize(objective, n_trials=100)
146
+
147
+ # Print the best trial
148
+ print("Best trial:")
149
+ trial = study.best_trial
150
+ print(f" Loss: {trial.value}")
151
+ print(" Params: ")
152
+ for key, value in trial.params.items():
153
+ print(f" {key}: {value}")
sft-dora-alpaca.py ADDED
@@ -0,0 +1,162 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ torch.backends.cuda.matmul.allow_tf32 = True
3
+ import random
4
+ from transformers import AutoTokenizer, AutoModelForCausalLM, TextGenerationPipeline, AutoConfig
5
+ from datasets import load_dataset
6
+ from transformers import TrainingArguments
7
+ from trl import SFTTrainer
8
+ from peft import LoraConfig
9
+ # from accelerate import infer_auto_device_map, init_empty_weights, dispatch_model
10
+ from torch.nn import CrossEntropyLoss
11
+
12
+ import time
13
+ random_seed = 42
14
+ torch.manual_seed(random_seed)
15
+ random.seed(random_seed)
16
+
17
+ dataset = load_dataset("Vezora/Tested-22k-Python-Alpaca", split="train")
18
+
19
+ def chatml_format(example):
20
+ """Format the dataset for training, accounting for empty columns."""
21
+ return {
22
+ "instruction": example['instruction'] if 'instruction' in example else " \n",
23
+ "input": example['input'] if 'input' in example else " \n",
24
+ "system": example['system'] if 'system' in example else " \n",
25
+ "output": example['output'] if 'output' in example else " \n",
26
+ }
27
+
28
+ # Format dataset
29
+ dataset = dataset.map(chatml_format, remove_columns=dataset.column_names)
30
+
31
+ n_ahead_talk_global = 4
32
+ n_passes_global = 2
33
+ n_ahead_global = 8
34
+ n_examples = 0
35
+
36
+ def model_init(params):
37
+ original = False
38
+ if params is None:
39
+ params = {}
40
+ else:
41
+ params = params.params
42
+ # save params to file
43
+ n_ahead = params.get("n_ahead", n_ahead_global if not original else 1)
44
+ n_ahead_talk = params.get("n_ahead_talk", n_ahead_talk_global if not original else 1)
45
+ n_passes = params.get("n_passes", n_passes_global if not original else 1)
46
+ gumbel_temperature = params.get("gumbel_temperature", 1)
47
+ use_start_thought_token = params.get("use_start_thought_token", True)
48
+ use_end_thought_token = params.get("use_end_thought_token", True)
49
+ include_policy_loss = params.get("include_policy_loss", True)
50
+ gumbel_detach = params.get("gumbel_detach", True)
51
+ merged_talk_heads = params.get("merged_talk_heads", True)
52
+ residual_think_head = params.get("residual_think_head", False)
53
+ optimize_lm_head_only_at_start = params.get("optimize_lm_head_only_at_start", False)
54
+
55
+ model_id = "Crystalcareai/Quiet-Star-Custom"
56
+ tokenizer_id = model_id
57
+ print("Loading model")
58
+
59
+ model = AutoModelForCausalLM.from_pretrained(
60
+ model_id,
61
+ torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
62
+ max_thoughts=n_ahead + n_ahead_talk + 1,
63
+ merged_talk_heads=merged_talk_heads,
64
+ merged_lm_and_talk_heads=False,
65
+ merged_lm_and_think_heads=True,
66
+ use_concat_talk_head=True,
67
+ use_shallow_think=True,
68
+ use_shallow_talk=False,
69
+ use_complex_think_head=False,
70
+ use_complex_talk_head=True,
71
+ use_weighted_talk_head=True,
72
+ trust_remote_code=True,
73
+ device_map="auto",
74
+ )
75
+ print("Loaded model")
76
+
77
+ tokenizer = AutoTokenizer.from_pretrained(tokenizer_id, truncation=True, padding_side="right")
78
+ tokenizer.pad_token_id = tokenizer.eos_token_id
79
+
80
+ special_tokens_to_add = []
81
+ if model.use_start_thought_token:
82
+ special_tokens_to_add.append("<|startthought|>")
83
+ if model.use_end_thought_token:
84
+ special_tokens_to_add.append("<|endthought|>")
85
+ if special_tokens_to_add:
86
+ tokenizer.add_special_tokens({"additional_special_tokens": special_tokens_to_add})
87
+ model.resize_token_embeddings(len(tokenizer))
88
+ model.tokenizer = tokenizer
89
+ for name, module in model.named_modules():
90
+ if "embed" in name:
91
+ print(module, flush=True)
92
+
93
+ model.gumbel_detach = gumbel_detach
94
+ model.include_policy_loss = include_policy_loss
95
+ model.use_end_thought_token = use_end_thought_token
96
+ model.use_start_thought_token = use_start_thought_token
97
+ model.n_ahead = n_ahead
98
+ model.n_ahead_talk = n_ahead_talk
99
+ model.n_passes = n_passes
100
+ model.residual_think_head = residual_think_head
101
+ model.optimize_lm_head_only_at_start = optimize_lm_head_only_at_start
102
+ model.gumbel_temperature = gumbel_temperature
103
+ model.original_mode = original
104
+ model.config_params = params
105
+ model.run_start = int(time.time())
106
+ model.train()
107
+ return model
108
+
109
+ max_seq_length = 1024
110
+ run_id = int(time.time())
111
+ training_args = TrainingArguments(
112
+ output_dir="./out",
113
+ num_train_epochs=3,
114
+ per_device_train_batch_size=1,
115
+ gradient_checkpointing=False,
116
+ gradient_accumulation_steps=8,
117
+ optim="lion_32bit",
118
+ logging_steps=1,
119
+ save_strategy="steps",
120
+ save_steps=300,
121
+ max_steps=1000,
122
+ bf16=True,
123
+ tf32=False,
124
+ learning_rate=6e-05,
125
+ max_grad_norm=0.3,
126
+ warmup_ratio=0.06,
127
+ lr_scheduler_type="cosine",
128
+ push_to_hub=False,
129
+ report_to="wandb"
130
+ )
131
+
132
+ peft_config = LoraConfig(
133
+ r = 16, # Choose any number > 0 ! Suggested 8, 16, 32, 64, 128
134
+ target_modules = ["q_proj", "k_proj"],
135
+ lora_alpha = 16,
136
+ lora_dropout = 0, # Supports any, but = 0 is optimized
137
+ bias = "none",
138
+ use_dora=True,
139
+ )
140
+
141
+
142
+ torch.autograd.set_detect_anomaly(True)
143
+
144
+ # Set the device for each process
145
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
146
+ # torch.cuda.set_device(device)
147
+
148
+ model = model_init(None) # Initialize the model
149
+
150
+ tokenizer = model.tokenizer
151
+
152
+ trainer = SFTTrainer(
153
+ args=training_args,
154
+ train_dataset=dataset,
155
+ model=model,
156
+ tokenizer=tokenizer,
157
+ max_seq_length=max_seq_length,
158
+ dataset_text_field="output",
159
+ peft_config=peft_config,
160
+ )
161
+
162
+ trainer.train()
special_tokens_map.json ADDED
@@ -0,0 +1,40 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "additional_special_tokens": [
3
+ {
4
+ "content": "<|endthought|>",
5
+ "lstrip": false,
6
+ "normalized": false,
7
+ "rstrip": false,
8
+ "single_word": false
9
+ },
10
+ {
11
+ "content": "<|startthought|>",
12
+ "lstrip": false,
13
+ "normalized": false,
14
+ "rstrip": false,
15
+ "single_word": false
16
+ }
17
+ ],
18
+ "bos_token": {
19
+ "content": "<s>",
20
+ "lstrip": false,
21
+ "normalized": false,
22
+ "rstrip": false,
23
+ "single_word": false
24
+ },
25
+ "eos_token": {
26
+ "content": "</s>",
27
+ "lstrip": false,
28
+ "normalized": false,
29
+ "rstrip": false,
30
+ "single_word": false
31
+ },
32
+ "pad_token": "</s>",
33
+ "unk_token": {
34
+ "content": "<unk>",
35
+ "lstrip": false,
36
+ "normalized": false,
37
+ "rstrip": false,
38
+ "single_word": false
39
+ }
40
+ }
tokenizer.json ADDED
The diff for this file is too large to render. See raw diff
 
tokenizer.model ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:dadfd56d766715c61d2ef780a525ab43b8e6da4de6865bda3d95fdef5e134055
3
+ size 493443
tokenizer_config.json ADDED
@@ -0,0 +1,61 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "add_bos_token": true,
3
+ "add_eos_token": false,
4
+ "added_tokens_decoder": {
5
+ "0": {
6
+ "content": "<unk>",
7
+ "lstrip": false,
8
+ "normalized": false,
9
+ "rstrip": false,
10
+ "single_word": false,
11
+ "special": true
12
+ },
13
+ "1": {
14
+ "content": "<s>",
15
+ "lstrip": false,
16
+ "normalized": false,
17
+ "rstrip": false,
18
+ "single_word": false,
19
+ "special": true
20
+ },
21
+ "2": {
22
+ "content": "</s>",
23
+ "lstrip": false,
24
+ "normalized": false,
25
+ "rstrip": false,
26
+ "single_word": false,
27
+ "special": true
28
+ },
29
+ "32000": {
30
+ "content": "<|endthought|>",
31
+ "lstrip": false,
32
+ "normalized": false,
33
+ "rstrip": false,
34
+ "single_word": false,
35
+ "special": true
36
+ },
37
+ "32001": {
38
+ "content": "<|startthought|>",
39
+ "lstrip": false,
40
+ "normalized": false,
41
+ "rstrip": false,
42
+ "single_word": false,
43
+ "special": true
44
+ }
45
+ },
46
+ "additional_special_tokens": [
47
+ "<|endthought|>",
48
+ "<|startthought|>"
49
+ ],
50
+ "bos_token": "<s>",
51
+ "clean_up_tokenization_spaces": false,
52
+ "eos_token": "</s>",
53
+ "legacy": true,
54
+ "model_max_length": 1000000000000000019884624838656,
55
+ "pad_token": "</s>",
56
+ "sp_model_kwargs": {},
57
+ "spaces_between_special_tokens": false,
58
+ "tokenizer_class": "LlamaTokenizer",
59
+ "unk_token": "<unk>",
60
+ "use_default_system_prompt": false
61
+ }
train-h100-sharegpt-sft.py ADDED
@@ -0,0 +1,193 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ torch.backends.cuda.matmul.allow_tf32 = True
3
+ import random
4
+ from transformers import AutoTokenizer, AutoModelForCausalLM, TextGenerationPipeline, AutoConfig, BitsAndBytesConfig
5
+ from datasets import load_dataset
6
+ from transformers import TrainingArguments
7
+ from accelerate import infer_auto_device_map, init_empty_weights, dispatch_model
8
+ from trl import SFTTrainer
9
+ from peft import LoraConfig
10
+ from torch.nn import CrossEntropyLoss
11
+ import time
12
+ import gc
13
+
14
+ random_seed = 42
15
+ torch.manual_seed(random_seed)
16
+ random.seed(random_seed)
17
+
18
+ dataset = load_dataset("HuggingFaceH4/orca-math-word-problems-200k", split="train_sft").select(range(1000))
19
+
20
+
21
+ n_ahead_talk_global = 4
22
+ n_passes_global = 1
23
+ n_ahead_global = 4
24
+ # n_examples = 1000
25
+ # full_batch_size = 8
26
+
27
+ def model_init(params):
28
+ original = False
29
+ if params is None:
30
+ params = {}
31
+ else:
32
+ params = params.params
33
+ # save params to file
34
+ n_ahead = params.get("n_ahead", n_ahead_global if not original else 1)
35
+ n_ahead_talk = params.get("n_ahead_talk", n_ahead_talk_global if not original else 1)
36
+ n_passes = params.get("n_passes", n_passes_global if not original else 1)
37
+ gumbel_temperature = params.get("gumbel_temperature", 1)
38
+ use_start_thought_token = params.get("use_start_thought_token", True)
39
+ use_end_thought_token = params.get("use_end_thought_token", True)
40
+ include_policy_loss = params.get("include_policy_loss", True)
41
+ gumbel_detach = params.get("gumbel_detach", True)
42
+ merged_talk_heads = params.get("merged_talk_heads", True)
43
+ residual_think_head = params.get("residual_think_head", False)
44
+ optimize_lm_head_only_at_start = params.get("optimize_lm_head_only_at_start", False)
45
+
46
+ model_id = "Crystalcareai/Quiet-Star-Custom"
47
+ tokenizer_id = model_id
48
+ print("Loading model")
49
+
50
+ model = AutoModelForCausalLM.from_pretrained(
51
+ model_id,
52
+ torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
53
+ max_thoughts=n_ahead + n_ahead_talk + 1,
54
+ merged_talk_heads=merged_talk_heads,
55
+ merged_lm_and_talk_heads=False,
56
+ merged_lm_and_think_heads=True,
57
+ use_concat_talk_head=True,
58
+ use_shallow_think=True,
59
+ use_shallow_talk=False,
60
+ use_complex_think_head=False,
61
+ use_complex_talk_head=True,
62
+ use_weighted_talk_head=True,
63
+ trust_remote_code=True,
64
+ device_map="auto",
65
+ )
66
+ print("Loaded model")
67
+
68
+ tokenizer = AutoTokenizer.from_pretrained(tokenizer_id, truncation=True, padding_side="right")
69
+ tokenizer.pad_token_id = tokenizer.eos_token_id
70
+
71
+ special_tokens_to_add = []
72
+ if model.use_start_thought_token:
73
+ special_tokens_to_add.append("<|startthought|>")
74
+ if model.use_end_thought_token:
75
+ special_tokens_to_add.append("<|endthought|>")
76
+ if special_tokens_to_add:
77
+ tokenizer.add_special_tokens({"additional_special_tokens": special_tokens_to_add})
78
+ model.tokenizer = tokenizer
79
+ for name, module in model.named_modules():
80
+ if "embed" in name:
81
+ print(module, flush=True)
82
+
83
+ model.gumbel_detach = gumbel_detach
84
+ model.include_policy_loss = include_policy_loss
85
+ model.use_end_thought_token = use_end_thought_token
86
+ model.use_start_thought_token = use_start_thought_token
87
+ model.n_ahead = n_ahead
88
+ model.n_ahead_talk = n_ahead_talk
89
+ model.n_passes = n_passes
90
+ model.residual_think_head = residual_think_head
91
+ model.optimize_lm_head_only_at_start = optimize_lm_head_only_at_start
92
+ model.gumbel_temperature = gumbel_temperature
93
+ model.original_mode = original
94
+ model.config_params = params
95
+ model.run_start = int(time.time())
96
+ model.train()
97
+ return model
98
+
99
+ max_seq_length = 1024
100
+ run_id = int(time.time())
101
+ training_args = TrainingArguments(
102
+ output_dir="./out",
103
+ num_train_epochs=1,
104
+ per_device_train_batch_size=1,
105
+ gradient_checkpointing=False,
106
+ gradient_accumulation_steps=8,
107
+ optim="adamw_torch_fused",
108
+ logging_steps=1,
109
+ save_strategy="steps",
110
+ save_steps=100,
111
+ max_steps=-1,
112
+ # auto_find_batch_size=True,
113
+ weight_decay=0.001,
114
+ bf16=True,
115
+
116
+ tf32=True,
117
+ learning_rate=2e-10,
118
+ max_grad_norm=0,
119
+ warmup_steps=20,
120
+ lr_scheduler_type="cosine",
121
+ push_to_hub=False,
122
+ report_to="wandb"
123
+ )
124
+
125
+ peft_config = LoraConfig(
126
+ r = 8, # Choose any number > 0 ! Suggested 8, 16, 32, 64, 128
127
+ target_modules =["q_proj", "v_proj"],
128
+ lora_alpha = 32,
129
+ lora_dropout = 0, # Supports any, but = 0 is optimized
130
+ bias = "none",
131
+ use_dora=True,
132
+ task_type="CAUSAL_LM"
133
+ )
134
+
135
+ torch.autograd.set_detect_anomaly(True)
136
+
137
+ # class CustomSFTTrainer(SFTTrainer):
138
+ # def __init__(self, *args, **kwargs):
139
+ # super().__init__(*args, **kwargs)
140
+ # self.beta = 0.9 # momentum factor
141
+ # self.clip_factor = 1.0 # clipping factor
142
+ # self.moving_avg = 0.0
143
+
144
+ # def training_step(self, model, inputs):
145
+ # model.train()
146
+ # inputs = self._prepare_inputs(inputs)
147
+
148
+ # outputs = model(**inputs)
149
+ # loss = outputs.loss if isinstance(outputs, dict) else outputs[0]
150
+
151
+ # if self.args.gradient_accumulation_steps > 1:
152
+ # loss = loss / self.args.gradient_accumulation_steps
153
+
154
+ # loss.backward()
155
+
156
+ # # Compute gradients and their norm
157
+ # grad_norm = torch.sqrt(sum(p.grad.data.norm().to(model.device)**2 for p in model.parameters() if p.grad is not None))
158
+
159
+ # # Update moving average and apply gradient clipping
160
+ # if self.state.global_step == 0:
161
+ # self.moving_avg = grad_norm
162
+ # else:
163
+ # self.moving_avg = self.beta * self.moving_avg + (1 - self.beta) * grad_norm
164
+
165
+ # if grad_norm > self.clip_factor * self.moving_avg:
166
+ # clip_coef = (self.clip_factor * self.moving_avg / grad_norm).item()
167
+ # for param in model.parameters():
168
+ # if param.grad is not None:
169
+ # param.grad.data.mul_(clip_coef)
170
+
171
+ # if (self.state.global_step + 1) % self.args.gradient_accumulation_steps == 0:
172
+ # self.optimizer.step()
173
+ # self.lr_scheduler.step()
174
+ # model.zero_grad()
175
+ # self.state.global_step += 1
176
+
177
+ # # Return the loss as a Tensor
178
+ # return loss
179
+
180
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
181
+
182
+ model = model_init(None)
183
+
184
+ trainer = SFTTrainer(
185
+ model=model,
186
+ args=training_args,
187
+ train_dataset=dataset,
188
+ tokenizer=model.tokenizer,
189
+ max_seq_length=max_seq_length,
190
+ peft_config=peft_config,
191
+ )
192
+
193
+ trainer.train()