XiaoduoAILab commited on
Commit
e44944c
1 Parent(s): 47da563

Upload 2 files

Browse files
Files changed (2) hide show
  1. configuration_xmodel.py +88 -0
  2. modeling_xmodel.py +739 -0
configuration_xmodel.py ADDED
@@ -0,0 +1,88 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (c) 2023 XiaoDuo AI. All rights reserved.
2
+
3
+ from transformers.configuration_utils import PretrainedConfig
4
+ from transformers.utils import logging
5
+ from typing_extensions import Self
6
+
7
+ logger = logging.get_logger(__name__)
8
+
9
+
10
+ class XModelConfig(PretrainedConfig):
11
+ model_type = "xmodel_32000"
12
+ keys_to_ignore_at_inference = ["past_key_values"]
13
+
14
+ def __init__(
15
+ self,
16
+ vocab_size=32000,
17
+ hidden_size=4096,
18
+ intermediate_size=None,
19
+ num_hidden_layers=32,
20
+ num_attention_heads=32,
21
+ num_key_value_heads=32,
22
+ hidden_act="silu",
23
+ max_position_embeddings=32768,
24
+ initializer_range=0.02,
25
+ rms_norm_eps=1e-5,
26
+ use_cache=True,
27
+ pad_token_id=0,
28
+ bos_token_id=1,
29
+ eos_token_id=2,
30
+ pretraining_tp=1,
31
+ tie_word_embeddings=False,
32
+ **kwargs,
33
+ ):
34
+ self.vocab_size = vocab_size
35
+ self.max_position_embeddings = max_position_embeddings
36
+ self.hidden_size = hidden_size
37
+ # self.intermediate_size = intermediate_size
38
+ if intermediate_size is None:
39
+ self.intermediate_size = find_multiple(int(8 * hidden_size / 3), 256)
40
+ else:
41
+ self.intermediate_size = intermediate_size
42
+ self.num_hidden_layers = num_hidden_layers
43
+ self.num_attention_heads = num_attention_heads
44
+ self.num_key_value_heads = num_key_value_heads
45
+ self.hidden_act = hidden_act
46
+ self.initializer_range = initializer_range
47
+ self.rms_norm_eps = rms_norm_eps
48
+ self.pretraining_tp = pretraining_tp
49
+ self.use_cache = use_cache
50
+ self.auto_map = {
51
+ "AutoConfig": "configuration_xmodel.XModelConfig",
52
+ "AutoModelForCausalLM": "modeling_xmodel.XModelForCausalLM"
53
+ }
54
+
55
+ super().__init__(
56
+ pad_token_id=pad_token_id,
57
+ bos_token_id=bos_token_id,
58
+ eos_token_id=eos_token_id,
59
+ tie_word_embeddings=tie_word_embeddings,
60
+ **kwargs,
61
+ )
62
+
63
+ @classmethod
64
+ def from_name(cls, name: str) -> Self:
65
+ return cls(**xmodel_configs[name])
66
+
67
+
68
+ xmodel_configs = {
69
+ "nano": dict(num_hidden_layers=6, num_attention_heads=6, num_key_value_heads=1, hidden_size=192),
70
+ "micro": dict(num_hidden_layers=6, num_attention_heads=6, num_key_value_heads=1, hidden_size=384),
71
+ "tiny": dict(num_hidden_layers=8, num_attention_heads=8, num_key_value_heads=2, hidden_size=512),
72
+ "small": dict(num_hidden_layers=12, num_attention_heads=12, num_key_value_heads=3, hidden_size=768),
73
+ # GPT-1 & Bert-Base
74
+ "medium": dict(num_hidden_layers=24, num_attention_heads=16, num_key_value_heads=4, hidden_size=1024), # Bert-Large
75
+ "large": dict(num_hidden_layers=24, num_attention_heads=16, num_key_value_heads=4, hidden_size=1536),
76
+ "xl": dict(num_hidden_layers=24, num_attention_heads=32, num_key_value_heads=4, hidden_size=2048), # GPT-2
77
+ "3B": dict(num_hidden_layers=32, num_attention_heads=32, num_key_value_heads=4, hidden_size=2560),
78
+ "7B": dict(num_hidden_layers=32, num_attention_heads=32, num_key_value_heads=32, hidden_size=4096),
79
+ "13B": dict(num_hidden_layers=40, num_attention_heads=40, num_key_value_heads=40, hidden_size=5120),
80
+ "34B": dict(num_hidden_layers=48, num_attention_heads=64, num_key_value_heads=8, hidden_size=8192),
81
+ "70B": dict(num_hidden_layers=80, num_attention_heads=64, num_key_value_heads=8, hidden_size=8192), # Llama
82
+ }
83
+
84
+
85
+ def find_multiple(n: int, k: int) -> int:
86
+ if n % k == 0:
87
+ return n
88
+ return n + k - (n % k)
modeling_xmodel.py ADDED
@@ -0,0 +1,739 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (c) 2023 XiaoDuo AI. All rights reserved.
2
+
3
+ import math
4
+ from typing import List, Optional, Tuple, Union
5
+
6
+ import torch
7
+ import torch.utils.checkpoint
8
+ import transformers
9
+ from torch import nn
10
+ from torch.nn import CrossEntropyLoss
11
+ from torch.nn import functional as F
12
+ from transformers.activations import ACT2FN
13
+ from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
14
+ from transformers.utils import logging
15
+
16
+ from .configuration_xmodel import XModelConfig
17
+
18
+ logger = logging.get_logger(__name__)
19
+ torch2 = torch.__version__.split('.')[0] == '2'
20
+
21
+
22
+ # Copied from transformers.models.bart.modeling_bart._make_causal_mask
23
+ def _make_causal_mask(
24
+ input_ids_shape: torch.Size, dtype: torch.dtype, device: torch.device, past_key_values_length: int = 0
25
+ ):
26
+ """
27
+ Make causal mask used for bi-directional self-attention.
28
+ """
29
+ bsz, tgt_len = input_ids_shape
30
+ mask = torch.full((tgt_len, tgt_len), torch.tensor(torch.finfo(dtype).min, device=device), device=device)
31
+ mask_cond = torch.arange(mask.size(-1), device=device)
32
+ mask.masked_fill_(mask_cond < (mask_cond + 1).view(mask.size(-1), 1), 0)
33
+ mask = mask.to(dtype)
34
+
35
+ if past_key_values_length > 0:
36
+ mask = torch.cat([torch.zeros(tgt_len, past_key_values_length, dtype=dtype, device=device), mask], dim=-1)
37
+ return mask[None, None, :, :].expand(bsz, 1, tgt_len, tgt_len + past_key_values_length)
38
+
39
+
40
+ # Copied from transformers.models.bart.modeling_bart._expand_mask
41
+ def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None):
42
+ """
43
+ Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`.
44
+ """
45
+ bsz, src_len = mask.size()
46
+ tgt_len = tgt_len if tgt_len is not None else src_len
47
+
48
+ expanded_mask = mask[:, None, None, :].expand(bsz, 1, tgt_len, src_len).to(dtype)
49
+
50
+ inverted_mask = 1.0 - expanded_mask
51
+
52
+ return inverted_mask.masked_fill(inverted_mask.to(torch.bool), torch.finfo(dtype).min)
53
+
54
+
55
+ class RMSNorm(nn.Module):
56
+ def __init__(self, hidden_size, eps=1e-6):
57
+ super().__init__()
58
+ self.weight = nn.Parameter(torch.ones(hidden_size))
59
+ self.variance_epsilon = eps
60
+
61
+ def forward(self, hidden_states):
62
+ variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True)
63
+ hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
64
+
65
+ # convert into half-precision if necessary
66
+ if self.weight.dtype in [torch.float16, torch.bfloat16]:
67
+ hidden_states = hidden_states.to(self.weight.dtype)
68
+
69
+ return self.weight * hidden_states
70
+
71
+
72
+ class RotaryEmbedding(torch.nn.Module):
73
+ def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):
74
+ super().__init__()
75
+ inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float().to(device) / dim))
76
+ self.register_buffer("inv_freq", inv_freq)
77
+
78
+ # Build here to make `torch.jit.trace` work.
79
+ self.max_seq_len_cached = max_position_embeddings
80
+ t = torch.arange(self.max_seq_len_cached, device=self.inv_freq.device, dtype=self.inv_freq.dtype)
81
+ freqs = torch.einsum("i,j->ij", t, self.inv_freq)
82
+ # Different from paper, but it uses a different permutation in order to obtain the same calculation
83
+ emb = torch.cat((freqs, freqs), dim=-1)
84
+ self.register_buffer("cos_cached", emb.cos()[None, None, :, :], persistent=False)
85
+ self.register_buffer("sin_cached", emb.sin()[None, None, :, :], persistent=False)
86
+
87
+ def forward(self, x, seq_len=None):
88
+ # x: [bs, num_attention_heads, seq_len, head_size]
89
+ # This `if` block is unlikely to be run after we build sin/cos in `__init__`. Keep the logic here just in case.
90
+ if seq_len > self.max_seq_len_cached:
91
+ self.max_seq_len_cached = seq_len
92
+ t = torch.arange(self.max_seq_len_cached, device=x.device, dtype=self.inv_freq.dtype)
93
+ freqs = torch.einsum("i,j->ij", t, self.inv_freq)
94
+ # Different from paper, but it uses a different permutation in order to obtain the same calculation
95
+ emb = torch.cat((freqs, freqs), dim=-1).to(x.device)
96
+ self.register_buffer("cos_cached", emb.cos()[None, None, :, :], persistent=False)
97
+ self.register_buffer("sin_cached", emb.sin()[None, None, :, :], persistent=False)
98
+ return (
99
+ self.cos_cached[:, :, :seq_len, ...].to(dtype=x.dtype),
100
+ self.sin_cached[:, :, :seq_len, ...].to(dtype=x.dtype),
101
+ )
102
+
103
+
104
+ def rotate_half(x):
105
+ """Rotates half the hidden dims of the input."""
106
+ x1 = x[..., : x.shape[-1] // 2]
107
+ x2 = x[..., x.shape[-1] // 2:]
108
+ return torch.cat((-x2, x1), dim=-1)
109
+
110
+
111
+ def apply_rotary_pos_emb(q, k, cos, sin, position_ids):
112
+ # The first two dimensions of cos and sin are always 1, so we can `squeeze` them.
113
+ cos = cos.squeeze(1).squeeze(0) # [seq_len, dim]
114
+ sin = sin.squeeze(1).squeeze(0) # [seq_len, dim]
115
+ cos = cos[position_ids].unsqueeze(1) # [bs, 1, seq_len, dim]
116
+ sin = sin[position_ids].unsqueeze(1) # [bs, 1, seq_len, dim]
117
+ q_embed = (q * cos) + (rotate_half(q) * sin)
118
+ k_embed = (k * cos) + (rotate_half(k) * sin)
119
+ return q_embed, k_embed
120
+
121
+
122
+ class MLP(nn.Module):
123
+ def __init__(
124
+ self,
125
+ hidden_size: int,
126
+ intermediate_size: int,
127
+ hidden_act: str,
128
+ ):
129
+ super().__init__()
130
+ self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
131
+ self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)
132
+ self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
133
+ self.act_fn = ACT2FN[hidden_act]
134
+
135
+ def forward(self, x):
136
+ out = self.gate_proj(x)
137
+ out = self.act_fn(out)
138
+ out = out * self.up_proj(x)
139
+ out = self.down_proj(out)
140
+ return out
141
+
142
+
143
+ def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
144
+ """
145
+ This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
146
+ num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
147
+ """
148
+ batch, num_key_value_heads, slen, head_dim = hidden_states.shape
149
+ if n_rep == 1:
150
+ return hidden_states
151
+ hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
152
+ return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
153
+
154
+
155
+ class Attention(nn.Module):
156
+ """Multi-headed attention from 'Attention Is All You Need' paper"""
157
+
158
+ def __init__(self, config: XModelConfig):
159
+ super().__init__()
160
+ self.config = config
161
+ self.hidden_size = config.hidden_size
162
+ self.num_heads = config.num_attention_heads
163
+ self.head_dim = self.hidden_size // self.num_heads
164
+ self.num_key_value_heads = config.num_key_value_heads
165
+ self.num_key_value_groups = self.num_heads // self.num_key_value_heads
166
+ self.max_position_embeddings = config.max_position_embeddings
167
+
168
+ if (self.head_dim * self.num_heads) != self.hidden_size:
169
+ raise ValueError(
170
+ f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"
171
+ f" and `num_heads`: {self.num_heads})."
172
+ )
173
+ self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)
174
+ self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
175
+ self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
176
+ self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)
177
+ self.rotary_emb = RotaryEmbedding(self.head_dim, max_position_embeddings=self.max_position_embeddings)
178
+
179
+ def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
180
+ return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
181
+
182
+ def forward(
183
+ self,
184
+ hidden_states: torch.Tensor,
185
+ attention_mask: Optional[torch.Tensor] = None,
186
+ position_ids: Optional[torch.LongTensor] = None,
187
+ past_key_value: Optional[Tuple[torch.Tensor]] = None,
188
+ output_attentions: bool = False,
189
+ use_cache: bool = False,
190
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
191
+ bsz, q_len, _ = hidden_states.size()
192
+
193
+ query_states = self.q_proj(hidden_states).view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
194
+ key_states = self.k_proj(hidden_states).view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1,
195
+ 2)
196
+ value_states = self.v_proj(hidden_states).view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1,
197
+ 2)
198
+
199
+ kv_seq_len = key_states.shape[-2]
200
+ if past_key_value is not None:
201
+ kv_seq_len += past_key_value[0].shape[-2]
202
+ cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
203
+ query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
204
+
205
+ if past_key_value is not None:
206
+ # reuse k, v, self_attention
207
+ key_states = torch.cat([past_key_value[0], key_states], dim=2)
208
+ value_states = torch.cat([past_key_value[1], value_states], dim=2)
209
+
210
+ past_key_value = (key_states, value_states) if use_cache else None
211
+
212
+ key_states = repeat_kv(key_states, self.num_key_value_groups)
213
+ value_states = repeat_kv(value_states, self.num_key_value_groups)
214
+
215
+ if torch2:
216
+ with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=True, enable_mem_efficient=True):
217
+ attn_output = F.scaled_dot_product_attention(query_states, key_states, value_states,
218
+ attn_mask=attention_mask)
219
+ else:
220
+ attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
221
+
222
+ if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len):
223
+ raise ValueError(
224
+ f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is"
225
+ f" {attn_weights.size()}"
226
+ )
227
+
228
+ if attention_mask is not None:
229
+ if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
230
+ raise ValueError(
231
+ f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
232
+ )
233
+ attn_weights = attn_weights + attention_mask
234
+
235
+ # upcast attention to fp32
236
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
237
+ # self.attention_dropout
238
+ attn_weights = nn.functional.dropout(attn_weights, training=self.training)
239
+ attn_output = torch.matmul(attn_weights, value_states)
240
+
241
+ if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):
242
+ raise ValueError(
243
+ f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"
244
+ f" {attn_output.size()}"
245
+ )
246
+
247
+ attn_output = attn_output.transpose(1, 2)
248
+
249
+ attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
250
+
251
+ attn_output = self.o_proj(attn_output)
252
+
253
+ if not output_attentions:
254
+ attn_weights = None
255
+
256
+ return attn_output, attn_weights, past_key_value
257
+
258
+
259
+ class DecoderLayer(nn.Module):
260
+ def __init__(self, config: XModelConfig):
261
+ super().__init__()
262
+ self.hidden_size = config.hidden_size
263
+ self.self_attn = Attention(config=config)
264
+ self.mlp = MLP(
265
+ hidden_size=self.hidden_size,
266
+ intermediate_size=config.intermediate_size,
267
+ hidden_act=config.hidden_act,
268
+ )
269
+ self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
270
+ self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
271
+
272
+ def forward(
273
+ self,
274
+ hidden_states: torch.Tensor,
275
+ attention_mask: Optional[torch.Tensor] = None,
276
+ position_ids: Optional[torch.LongTensor] = None,
277
+ past_key_value: Optional[Tuple[torch.Tensor]] = None,
278
+ output_attentions: Optional[bool] = False,
279
+ use_cache: Optional[bool] = False,
280
+ ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
281
+ """
282
+ Args:
283
+ hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
284
+ attention_mask (`torch.FloatTensor`, *optional*): attention mask of size
285
+ `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
286
+ output_attentions (`bool`, *optional*):
287
+ Whether or not to return the attentions tensors of all attention layers. See `attentions` under
288
+ returned tensors for more detail.
289
+ use_cache (`bool`, *optional*):
290
+ If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
291
+ (see `past_key_values`).
292
+ past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
293
+ """
294
+
295
+ residual = hidden_states
296
+
297
+ hidden_states = self.input_layernorm(hidden_states)
298
+
299
+ # Self Attention
300
+ hidden_states, self_attn_weights, present_key_value = self.self_attn(
301
+ hidden_states=hidden_states,
302
+ attention_mask=attention_mask,
303
+ position_ids=position_ids,
304
+ past_key_value=past_key_value,
305
+ output_attentions=output_attentions,
306
+ use_cache=use_cache,
307
+ )
308
+ hidden_states = residual + hidden_states
309
+
310
+ # Fully Connected
311
+ residual = hidden_states
312
+ hidden_states = self.post_attention_layernorm(hidden_states)
313
+ hidden_states = self.mlp(hidden_states)
314
+ hidden_states = residual + hidden_states
315
+
316
+ outputs = (hidden_states,)
317
+
318
+ if output_attentions:
319
+ outputs += (self_attn_weights,)
320
+
321
+ if use_cache:
322
+ outputs += (present_key_value,)
323
+
324
+ return outputs
325
+
326
+
327
+ class PreTrainedModel(transformers.PreTrainedModel):
328
+ config_class = XModelConfig
329
+ base_model_prefix = "model"
330
+ supports_gradient_checkpointing = True
331
+ _no_split_modules = ["DecoderLayer"]
332
+ _keys_to_ignore_on_load_unexpected = [r"decoder\.version"]
333
+
334
+ def _init_weights(self, module):
335
+ std = self.config.initializer_range
336
+ if isinstance(module, nn.Linear):
337
+ module.weight.data.normal_(mean=0.0, std=std)
338
+ if module.bias is not None:
339
+ module.bias.data.zero_()
340
+ elif isinstance(module, nn.Embedding):
341
+ module.weight.data.normal_(mean=0.0, std=std)
342
+ if module.padding_idx is not None:
343
+ module.weight.data[module.padding_idx].zero_()
344
+
345
+ def _set_gradient_checkpointing(self, module, value=False):
346
+ if isinstance(module, Model):
347
+ module.gradient_checkpointing = value
348
+
349
+
350
+ class Model(PreTrainedModel):
351
+ """
352
+ Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`DecoderLayer`]
353
+
354
+ Args:
355
+ config: XModelConfig
356
+ """
357
+
358
+ def __init__(self, config: XModelConfig):
359
+ super().__init__(config)
360
+ self.padding_idx = config.pad_token_id
361
+ self.vocab_size = config.vocab_size
362
+
363
+ self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
364
+ self.layers = nn.ModuleList([DecoderLayer(config) for _ in range(config.num_hidden_layers)])
365
+ self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
366
+
367
+ self.gradient_checkpointing = False
368
+ # Initialize weights and apply final processing
369
+ self.post_init()
370
+
371
+ def get_input_embeddings(self):
372
+ return self.embed_tokens
373
+
374
+ def set_input_embeddings(self, value):
375
+ self.embed_tokens = value
376
+
377
+ # Copied from transformers.models.bart.modeling_bart.BartDecoder._prepare_decoder_attention_mask
378
+ def _prepare_decoder_attention_mask(self, attention_mask, input_shape, inputs_embeds, past_key_values_length):
379
+ # create causal mask
380
+ # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
381
+ combined_attention_mask = None
382
+ if input_shape[-1] > 1:
383
+ combined_attention_mask = _make_causal_mask(
384
+ input_shape,
385
+ inputs_embeds.dtype,
386
+ device=inputs_embeds.device,
387
+ past_key_values_length=past_key_values_length,
388
+ )
389
+
390
+ if attention_mask is not None:
391
+ # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
392
+ expanded_attn_mask = _expand_mask(attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1]).to(
393
+ inputs_embeds.device
394
+ )
395
+ combined_attention_mask = (
396
+ expanded_attn_mask if combined_attention_mask is None else expanded_attn_mask + combined_attention_mask
397
+ )
398
+
399
+ return combined_attention_mask
400
+
401
+ def forward(
402
+ self,
403
+ input_ids: torch.LongTensor = None,
404
+ attention_mask: Optional[torch.Tensor] = None,
405
+ position_ids: Optional[torch.LongTensor] = None,
406
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
407
+ inputs_embeds: Optional[torch.FloatTensor] = None,
408
+ use_cache: Optional[bool] = None,
409
+ output_attentions: Optional[bool] = None,
410
+ output_hidden_states: Optional[bool] = None,
411
+ return_dict: Optional[bool] = None,
412
+ ) -> Union[Tuple, BaseModelOutputWithPast]:
413
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
414
+ output_hidden_states = (
415
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
416
+ )
417
+ use_cache = use_cache if use_cache is not None else self.config.use_cache
418
+
419
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
420
+
421
+ # retrieve input_ids and inputs_embeds
422
+ if input_ids is not None and inputs_embeds is not None:
423
+ raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time")
424
+ elif input_ids is not None:
425
+ batch_size, seq_length = input_ids.shape
426
+ elif inputs_embeds is not None:
427
+ batch_size, seq_length, _ = inputs_embeds.shape
428
+ else:
429
+ raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds")
430
+
431
+ seq_length_with_past = seq_length
432
+ past_key_values_length = 0
433
+
434
+ if past_key_values is not None:
435
+ past_key_values_length = past_key_values[0][0].shape[2]
436
+ seq_length_with_past = seq_length_with_past + past_key_values_length
437
+
438
+ if position_ids is None:
439
+ device = input_ids.device if input_ids is not None else inputs_embeds.device
440
+ position_ids = torch.arange(
441
+ past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=device
442
+ )
443
+ position_ids = position_ids.unsqueeze(0).view(-1, seq_length)
444
+ else:
445
+ position_ids = position_ids.view(-1, seq_length).long()
446
+
447
+ if inputs_embeds is None:
448
+ inputs_embeds = self.embed_tokens(input_ids)
449
+ # embed positions
450
+ if attention_mask is None:
451
+ attention_mask = torch.ones(
452
+ (batch_size, seq_length_with_past), dtype=torch.bool, device=inputs_embeds.device
453
+ )
454
+ attention_mask = self._prepare_decoder_attention_mask(
455
+ attention_mask, (batch_size, seq_length), inputs_embeds, past_key_values_length
456
+ )
457
+
458
+ hidden_states = inputs_embeds
459
+
460
+ if self.gradient_checkpointing and self.training:
461
+ if use_cache:
462
+ logger.warning_once(
463
+ "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
464
+ )
465
+ use_cache = False
466
+
467
+ # decoder layers
468
+ all_hidden_states = () if output_hidden_states else None
469
+ all_self_attns = () if output_attentions else None
470
+ next_decoder_cache = () if use_cache else None
471
+
472
+ for idx, decoder_layer in enumerate(self.layers):
473
+ if output_hidden_states:
474
+ all_hidden_states += (hidden_states,)
475
+
476
+ past_key_value = past_key_values[idx] if past_key_values is not None else None
477
+
478
+ if self.gradient_checkpointing and self.training:
479
+
480
+ def create_custom_forward(module):
481
+ def custom_forward(*inputs):
482
+ # None for past_key_value
483
+ return module(*inputs, output_attentions, None)
484
+
485
+ return custom_forward
486
+
487
+ layer_outputs = torch.utils.checkpoint.checkpoint(
488
+ create_custom_forward(decoder_layer),
489
+ hidden_states,
490
+ attention_mask,
491
+ position_ids,
492
+ None,
493
+ )
494
+ else:
495
+ layer_outputs = decoder_layer(
496
+ hidden_states,
497
+ attention_mask=attention_mask,
498
+ position_ids=position_ids,
499
+ past_key_value=past_key_value,
500
+ output_attentions=output_attentions,
501
+ use_cache=use_cache,
502
+ )
503
+ # print('debug_attention_mask', type(attention_mask),attention_mask.dtype)
504
+ # print('debug_position_ids', type(position_ids),position_ids.dtype)
505
+ hidden_states = layer_outputs[0]
506
+
507
+ if use_cache:
508
+ next_decoder_cache += (layer_outputs[2 if output_attentions else 1],)
509
+
510
+ if output_attentions:
511
+ all_self_attns += (layer_outputs[1],)
512
+
513
+ hidden_states = self.norm(hidden_states)
514
+
515
+ # add hidden states from the last decoder layer
516
+ if output_hidden_states:
517
+ all_hidden_states += (hidden_states,)
518
+
519
+ next_cache = next_decoder_cache if use_cache else None
520
+ if not return_dict:
521
+ return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)
522
+ return BaseModelOutputWithPast(
523
+ last_hidden_state=hidden_states,
524
+ past_key_values=next_cache,
525
+ hidden_states=all_hidden_states,
526
+ attentions=all_self_attns,
527
+ )
528
+
529
+
530
+ class XModelForCausalLM(PreTrainedModel):
531
+ def __init__(self, config):
532
+ super().__init__(config)
533
+ self.model = Model(config)
534
+
535
+ self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
536
+
537
+ # Initialize weights and apply final processing
538
+ self.post_init()
539
+
540
+ def get_input_embeddings(self):
541
+ return self.model.embed_tokens
542
+
543
+ def set_input_embeddings(self, value):
544
+ self.model.embed_tokens = value
545
+
546
+ def get_output_embeddings(self):
547
+ return self.lm_head
548
+
549
+ def set_output_embeddings(self, new_embeddings):
550
+ self.lm_head = new_embeddings
551
+
552
+ def set_decoder(self, decoder):
553
+ self.model = decoder
554
+
555
+ def get_decoder(self):
556
+ return self.model
557
+
558
+ def forward(
559
+ self,
560
+ input_ids: torch.LongTensor = None,
561
+ attention_mask: Optional[torch.Tensor] = None,
562
+ position_ids: Optional[torch.LongTensor] = None,
563
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
564
+ inputs_embeds: Optional[torch.FloatTensor] = None,
565
+ labels: Optional[torch.LongTensor] = None,
566
+ use_cache: Optional[bool] = None,
567
+ output_attentions: Optional[bool] = None,
568
+ output_hidden_states: Optional[bool] = None,
569
+ return_dict: Optional[bool] = None,
570
+ ) -> Union[Tuple, CausalLMOutputWithPast]:
571
+ r"""
572
+ Args:
573
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
574
+ Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
575
+ config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
576
+ (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
577
+
578
+ Returns:
579
+
580
+ Example:
581
+
582
+ ```python
583
+ >>> from transformers import AutoTokenizer, AutoModelForCausalLM
584
+
585
+ >>> model = AutoModelForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS)
586
+ >>> tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER)
587
+
588
+ >>> prompt = "Hey, are you consciours? Can you talk to me?"
589
+ >>> inputs = tokenizer(prompt, return_tensors="pt")
590
+
591
+ >>> # Generate
592
+ >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
593
+ >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
594
+ "Hey, are you consciours? Can you talk to me?\nI'm not consciours, but I can talk to you."
595
+ ```"""
596
+
597
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
598
+ output_hidden_states = (
599
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
600
+ )
601
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
602
+
603
+ # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
604
+ outputs = self.model(
605
+ input_ids=input_ids,
606
+ attention_mask=attention_mask,
607
+ position_ids=position_ids,
608
+ past_key_values=past_key_values,
609
+ inputs_embeds=inputs_embeds,
610
+ use_cache=use_cache,
611
+ output_attentions=output_attentions,
612
+ output_hidden_states=output_hidden_states,
613
+ return_dict=return_dict,
614
+ )
615
+
616
+ hidden_states = outputs[0]
617
+ logits = self.lm_head(hidden_states)
618
+
619
+ loss = None
620
+ if labels is not None:
621
+ # Shift so that tokens < n predict n
622
+ shift_logits = logits[..., :-1, :].contiguous()
623
+ shift_labels = labels[..., 1:].contiguous()
624
+ # Flatten the tokens
625
+ loss_fct = CrossEntropyLoss()
626
+ shift_logits = shift_logits.view(-1, self.config.vocab_size)
627
+ shift_labels = shift_labels.view(-1)
628
+ # Enable model parallelism
629
+ shift_labels = shift_labels.to(shift_logits.device)
630
+ loss = loss_fct(shift_logits, shift_labels)
631
+
632
+ if not return_dict:
633
+ output = (logits,) + outputs[1:]
634
+ return (loss,) + output if loss is not None else output
635
+
636
+ return CausalLMOutputWithPast(
637
+ loss=loss,
638
+ logits=logits,
639
+ past_key_values=outputs.past_key_values,
640
+ hidden_states=outputs.hidden_states,
641
+ attentions=outputs.attentions,
642
+ )
643
+
644
+ def prepare_inputs_for_generation(
645
+ self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, **kwargs
646
+ ):
647
+ if past_key_values:
648
+ input_ids = input_ids[:, -1:]
649
+
650
+ position_ids = kwargs.get("position_ids", None)
651
+ if attention_mask is not None and position_ids is None:
652
+ # create position_ids on the fly for batch generation
653
+ position_ids = attention_mask.long().cumsum(-1) - 1
654
+ position_ids.masked_fill_(attention_mask == 0, 1)
655
+ if past_key_values:
656
+ position_ids = position_ids[:, -1].unsqueeze(-1)
657
+
658
+ # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
659
+ if inputs_embeds is not None and past_key_values is None:
660
+ model_inputs = {"inputs_embeds": inputs_embeds}
661
+ else:
662
+ model_inputs = {"input_ids": input_ids}
663
+
664
+ model_inputs.update(
665
+ {
666
+ "position_ids": position_ids,
667
+ "past_key_values": past_key_values,
668
+ "use_cache": kwargs.get("use_cache"),
669
+ "attention_mask": attention_mask,
670
+ }
671
+ )
672
+ return model_inputs
673
+
674
+ @staticmethod
675
+ def _reorder_cache(past_key_values, beam_idx):
676
+ reordered_past = ()
677
+ for layer_past in past_key_values:
678
+ reordered_past += (tuple(past_state.index_select(0, beam_idx) for past_state in layer_past),)
679
+ return reordered_past
680
+
681
+ def get_num_params(self, non_embedding=True):
682
+ """
683
+ Return the number of parameters in the model.
684
+ For non-embedding count (default), the position embeddings get subtracted.
685
+ The token embeddings would too, except due to the parameter sharing these
686
+ params are actually used as weights in the final layer, so we include them.
687
+ """
688
+ n_params = sum(p.numel() for p in self.parameters())
689
+ # if non_embedding:
690
+ # n_params -= self.transformer.wte.weight.numel()
691
+ return n_params
692
+
693
+ def estimate_mfu(self, fwdbwd_per_iter, dt, max_length=None, device_model='A100', dtype='float32'):
694
+ """ estimate model flops utilization (MFU) in units of A100 bfloat16 peak FLOPS """
695
+ # first estimate the number of flops we do per iteration.
696
+ # see PaLM paper Appendix B as ref: https://arxiv.org/abs/2204.02311
697
+ N = self.get_num_params()
698
+ n_layer = self.config.num_hidden_layers
699
+ n_head = self.config.num_attention_heads
700
+ n_embd = self.config.hidden_size
701
+
702
+ if max_length is None:
703
+ max_length = self.config.max_position_embeddings
704
+
705
+ L, H, Q, T = n_layer, n_head, n_embd // n_head, max_length
706
+ flops_per_token = 6 * N + 12 * L * H * Q * T
707
+ flops_per_fwdbwd = flops_per_token * T
708
+ flops_per_iter = flops_per_fwdbwd * fwdbwd_per_iter
709
+ # express our flops throughput as ratio of A100 bfloat16 peak flops
710
+ flops_achieved = flops_per_iter * (1.0 / dt) # per second
711
+
712
+ if device_model is None:
713
+ device_model = torch.cuda.get_device_name(0)
714
+
715
+ flops_promised = 312e12 # A100 GPU bfloat16 peak flops is 312 TFLOPS
716
+ if device_model == 'DCU' and dtype == 'float16':
717
+ flops_promised = 23.6e12
718
+ elif device_model == 'DCU' and dtype == 'float32':
719
+ flops_promised = 11.8e12
720
+ elif device_model == 'NVIDIA V100' or 'V100' in device_model:
721
+ flops_promised = 28e12
722
+ elif device_model == 'NVIDIA H100' or 'H100' in device_model or device_model == 'NVIDIA H800' or 'H800' in device_model:
723
+ flops_promised = 1513e12
724
+
725
+ mfu = flops_achieved / flops_promised
726
+ return flops_achieved, mfu
727
+
728
+ def flops_per_token(self, max_length=None, non_embedding=False):
729
+ N = self.get_num_params()
730
+ if non_embedding:
731
+ N -= self.config.vocab_size * self.config.hidden_size * 2
732
+ n_layer = self.config.num_hidden_layers
733
+ n_head = self.config.num_attention_heads
734
+ n_embd = self.config.hidden_size
735
+ if max_length is None:
736
+ max_length = self.config.max_position_embeddings
737
+ L, H, Q, T = n_layer, n_head, n_embd // n_head, max_length
738
+ flops_per_token = 6 * N + 12 * L * H * Q * T
739
+ return flops_per_token