keyfan commited on
Commit
8d9db92
1 Parent(s): fa044d9

Upload folder using huggingface_hub

Browse files
config.json ADDED
@@ -0,0 +1,37 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "_name_or_path": "hf/",
3
+ "architectures": [
4
+ "GrokForCausalLM"
5
+ ],
6
+ "attention_dropout": 0.0,
7
+ "attn_output_multiplier": 0.08838834764831845,
8
+ "auto_map": {
9
+ "AutoConfig": "configuration_grok.GrokConfig",
10
+ "AutoModelForCausalLM": "modeling_grok.GrokForCausalLM"
11
+ },
12
+ "bos_token_id": 1,
13
+ "embedding_multiplier_scale": 78.38367176906169,
14
+ "eos_token_id": 2,
15
+ "hidden_act": "gelu",
16
+ "hidden_size": 6144,
17
+ "initializer_range": 0.02,
18
+ "intermediate_size": 32768,
19
+ "max_position_embeddings": 8192,
20
+ "model_type": "grok",
21
+ "num_attention_heads": 48,
22
+ "num_experts_per_tok": 2,
23
+ "num_hidden_layers": 64,
24
+ "num_key_value_heads": 8,
25
+ "num_local_experts": 8,
26
+ "output_multiplier_scale": 0.5773502691896257,
27
+ "output_router_logits": false,
28
+ "pad_token_id": 0,
29
+ "rms_norm_eps": 1e-05,
30
+ "rope_theta": 100000.0,
31
+ "router_aux_loss_coef": 0.02,
32
+ "sliding_window": null,
33
+ "torch_dtype": "bfloat16",
34
+ "transformers_version": "4.38.2",
35
+ "use_cache": true,
36
+ "vocab_size": 131072
37
+ }
configuration_grok.py ADDED
@@ -0,0 +1,151 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # coding=utf-8
2
+ # Copyright 2023 Mixtral 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
+ """ Grok 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
+
24
+ class GrokConfig(PretrainedConfig):
25
+ r"""
26
+ This is the configuration class to store the configuration of a [`GrokModel`].
27
+
28
+ Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
29
+ documentation from [`PretrainedConfig`] for more information.
30
+
31
+
32
+ Args:
33
+ vocab_size (`int`, *optional*, defaults to 32000):
34
+ Vocabulary size of the Mixtral model. Defines the number of different tokens that can be represented by the
35
+ `inputs_ids` passed when calling [`MixtralModel`]
36
+ hidden_size (`int`, *optional*, defaults to 4096):
37
+ Dimension of the hidden representations.
38
+ intermediate_size (`int`, *optional*, defaults to 14336):
39
+ Dimension of the MLP representations.
40
+ num_hidden_layers (`int`, *optional*, defaults to 32):
41
+ Number of hidden layers in the Transformer encoder.
42
+ num_attention_heads (`int`, *optional*, defaults to 32):
43
+ Number of attention heads for each attention layer in the Transformer encoder.
44
+ num_key_value_heads (`int`, *optional*, defaults to 8):
45
+ This is the number of key_value heads that should be used to implement Grouped Query Attention. If
46
+ `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
47
+ `num_key_value_heads=1 the model will use Multi Query Attention (MQA) otherwise GQA is used. When
48
+ converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed
49
+ by meanpooling all the original heads within that group. For more details checkout [this
50
+ paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to `8`.
51
+ hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
52
+ The non-linear activation function (function or string) in the decoder.
53
+ max_position_embeddings (`int`, *optional*, defaults to `4096*32`):
54
+ The maximum sequence length that this model might ever be used with. Mixtral's sliding window attention
55
+ allows sequence of up to 4096*32 tokens.
56
+ initializer_range (`float`, *optional*, defaults to 0.02):
57
+ The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
58
+ rms_norm_eps (`float`, *optional*, defaults to 1e-05):
59
+ The epsilon used by the rms normalization layers.
60
+ use_cache (`bool`, *optional*, defaults to `True`):
61
+ Whether or not the model should return the last key/values attentions (not used by all models). Only
62
+ relevant if `config.is_decoder=True`.
63
+ pad_token_id (`int`, *optional*):
64
+ The id of the padding token.
65
+ bos_token_id (`int`, *optional*, defaults to 1):
66
+ The id of the "beginning-of-sequence" token.
67
+ eos_token_id (`int`, *optional*, defaults to 2):
68
+ The id of the "end-of-sequence" token.
69
+ tie_word_embeddings (`bool`, *optional*, defaults to `True`):
70
+ Whether the model's input and output word embeddings should be tied.
71
+ rope_theta (`float`, *optional*, defaults to 100000.0):
72
+ The base period of the RoPE embeddings.
73
+ attention_dropout (`float`, *optional*, defaults to 0.0):
74
+ The dropout ratio for the attention probabilities.
75
+ num_experts_per_tok (`int`, *optional*, defaults to 2):
76
+ The number of experts to root per-token, can be also interpreted as the `top-p` routing
77
+ parameter
78
+ num_local_experts (`int`, *optional*, defaults to 8):
79
+ Number of experts per Sparse MLP layer.
80
+ output_router_logits (`bool`, *optional*, defaults to `False`):
81
+ Whether or not the router logits should be returned by the model. Enabeling this will also
82
+ allow the model to output the auxiliary loss. See [here]() for more details
83
+ router_aux_loss_coef (`float`, *optional*, defaults to 0.001):
84
+ The aux loss factor for the total loss.
85
+
86
+ """
87
+
88
+ model_type = "grok"
89
+ keys_to_ignore_at_inference = ["past_key_values"]
90
+
91
+ def __init__(
92
+ self,
93
+ vocab_size=131072,
94
+ hidden_size=6144,
95
+ intermediate_size=32768,
96
+ num_hidden_layers=64,
97
+ num_attention_heads=48,
98
+ num_key_value_heads=8,
99
+ hidden_act="silu",
100
+ max_position_embeddings=4096,
101
+ initializer_range=0.02,
102
+ rms_norm_eps=1e-5,
103
+ use_cache=True,
104
+ pad_token_id=0,
105
+ bos_token_id=1,
106
+ eos_token_id=2,
107
+ tie_word_embeddings=True,
108
+ rope_theta=1e5,
109
+ attention_dropout=0.0,
110
+ num_experts_per_tok=2,
111
+ num_local_experts=8,
112
+ output_router_logits=False,
113
+ router_aux_loss_coef=0.001,
114
+ output_multiplier_scale=0.5773502691896257,
115
+ embedding_multiplier_scale=78.38367176906169,
116
+ attn_output_multiplier=0.08838834764831845,
117
+ **kwargs,
118
+ ):
119
+ self.vocab_size = vocab_size
120
+ self.max_position_embeddings = max_position_embeddings
121
+ self.hidden_size = hidden_size
122
+ self.intermediate_size = intermediate_size
123
+ self.num_hidden_layers = num_hidden_layers
124
+ self.num_attention_heads = num_attention_heads
125
+
126
+ # for backward compatibility
127
+ if num_key_value_heads is None:
128
+ num_key_value_heads = num_attention_heads
129
+
130
+ self.num_key_value_heads = num_key_value_heads
131
+ self.hidden_act = hidden_act
132
+ self.initializer_range = initializer_range
133
+ self.rms_norm_eps = rms_norm_eps
134
+ self.use_cache = use_cache
135
+ self.rope_theta = rope_theta
136
+ self.attention_dropout = attention_dropout
137
+
138
+ self.num_experts_per_tok = num_experts_per_tok
139
+ self.num_local_experts = num_local_experts
140
+ self.output_router_logits = output_router_logits
141
+ self.router_aux_loss_coef = router_aux_loss_coef
142
+ self.output_multiplier_scale = output_multiplier_scale
143
+ self.embedding_multiplier_scale = embedding_multiplier_scale
144
+ self.attn_output_multiplier = attn_output_multiplier
145
+ super().__init__(
146
+ pad_token_id=pad_token_id,
147
+ bos_token_id=bos_token_id,
148
+ eos_token_id=eos_token_id,
149
+ tie_word_embeddings=tie_word_embeddings,
150
+ **kwargs,
151
+ )
generation_config.json ADDED
@@ -0,0 +1,7 @@
 
 
 
 
 
 
 
 
1
+ {
2
+ "_from_model_config": true,
3
+ "bos_token_id": 1,
4
+ "eos_token_id": 2,
5
+ "pad_token_id": 0,
6
+ "transformers_version": "4.38.2"
7
+ }
modeling_grok.py ADDED
@@ -0,0 +1,838 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # coding=utf-8
2
+ # Modified from https://raw.githubusercontent.com/huggingface/transformers/v4.38.2/src/transformers/models/mixtral/modeling_mixtral.py
3
+ # Copyright 2023 Mistral AI and the HuggingFace Inc. team. All rights reserved.
4
+ #
5
+ # This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
6
+ # and OPT implementations in this library. It has been modified from its
7
+ # original forms to accommodate minor architectural differences compared
8
+ # to GPT-NeoX and OPT used by the Meta AI team that trained the model.
9
+ #
10
+ # Licensed under the Apache License, Version 2.0 (the "License");
11
+ # you may not use this file except in compliance with the License.
12
+ # You may obtain a copy of the License at
13
+ #
14
+ # http://www.apache.org/licenses/LICENSE-2.0
15
+ #
16
+ # Unless required by applicable law or agreed to in writing, software
17
+ # distributed under the License is distributed on an "AS IS" BASIS,
18
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
19
+ # See the License for the specific language governing permissions and
20
+ # limitations under the License.
21
+ """ PyTorch Grok-1 model."""
22
+ import inspect
23
+ import math
24
+ import warnings
25
+ from typing import List, Optional, Tuple, Union
26
+
27
+ import torch
28
+ import torch.nn.functional as F
29
+ import torch.utils.checkpoint
30
+ from torch import nn
31
+ from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
32
+
33
+ from transformers.activations import ACT2FN
34
+ from transformers.cache_utils import Cache, DynamicCache
35
+ from transformers.modeling_attn_mask_utils import (
36
+ _prepare_4d_causal_attention_mask,
37
+ )
38
+ from transformers.modeling_outputs import (
39
+ MoeCausalLMOutputWithPast,
40
+ MoeModelOutputWithPast,
41
+ SequenceClassifierOutputWithPast,
42
+ )
43
+ from transformers.modeling_utils import PreTrainedModel
44
+ from transformers.pytorch_utils import is_torch_greater_or_equal_than_1_13
45
+ from transformers.utils import (
46
+ add_start_docstrings,
47
+ add_start_docstrings_to_model_forward,
48
+ logging,
49
+ replace_return_docstrings,
50
+ )
51
+ from transformers.utils.import_utils import is_torch_fx_available
52
+ from .configuration_grok import GrokConfig
53
+
54
+
55
+ # This makes `_prepare_4d_causal_attention_mask` a leaf function in the FX graph.
56
+ # It means that the function will not be traced through and simply appear as a node in the graph.
57
+ if is_torch_fx_available():
58
+ if not is_torch_greater_or_equal_than_1_13:
59
+ import torch.fx
60
+
61
+ _prepare_4d_causal_attention_mask = torch.fx.wrap(_prepare_4d_causal_attention_mask)
62
+
63
+
64
+ logger = logging.get_logger(__name__)
65
+
66
+
67
+ # Copied from transformers.models.llama.modeling_llama._get_unpad_data
68
+ def _get_unpad_data(attention_mask):
69
+ seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
70
+ indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
71
+ max_seqlen_in_batch = seqlens_in_batch.max().item()
72
+ cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0))
73
+ return (
74
+ indices,
75
+ cu_seqlens,
76
+ max_seqlen_in_batch,
77
+ )
78
+
79
+
80
+ # Copied from transformers.models.llama.modeling_llama.LlamaRMSNorm with Llama->Grok
81
+ class GrokRMSNorm(nn.Module):
82
+ def __init__(self, hidden_size, eps=1e-6):
83
+ """
84
+ GrokRMSNorm is equivalent to T5LayerNorm
85
+ """
86
+ super().__init__()
87
+ self.weight = nn.Parameter(torch.ones(hidden_size))
88
+ self.variance_epsilon = eps
89
+
90
+ def forward(self, hidden_states):
91
+ input_dtype = hidden_states.dtype
92
+ hidden_states = hidden_states.to(torch.float32)
93
+ variance = hidden_states.pow(2).mean(-1, keepdim=True)
94
+ hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
95
+ return self.weight * hidden_states.to(input_dtype)
96
+
97
+
98
+ # Copied from transformers.models.mistral.modeling_mistral.MistralRotaryEmbedding with Mistral->Grok
99
+ class GrokRotaryEmbedding(nn.Module):
100
+ def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):
101
+ super().__init__()
102
+
103
+ self.dim = dim
104
+ self.max_position_embeddings = max_position_embeddings
105
+ self.base = base
106
+ inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64).float().to(device) / self.dim))
107
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
108
+
109
+ # Build here to make `torch.jit.trace` work.
110
+ self._set_cos_sin_cache(
111
+ seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype()
112
+ )
113
+
114
+ def _set_cos_sin_cache(self, seq_len, device, dtype):
115
+ self.max_seq_len_cached = seq_len
116
+ t = torch.arange(self.max_seq_len_cached, device=device, dtype=torch.int64).type_as(self.inv_freq)
117
+
118
+ freqs = torch.outer(t, self.inv_freq)
119
+ # Different from paper, but it uses a different permutation in order to obtain the same calculation
120
+ emb = torch.cat((freqs, freqs), dim=-1)
121
+ self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
122
+ self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
123
+
124
+ def forward(self, x, seq_len=None):
125
+ # x: [bs, num_attention_heads, seq_len, head_size]
126
+ if seq_len > self.max_seq_len_cached:
127
+ self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype)
128
+
129
+ return (
130
+ self.cos_cached[:seq_len].to(dtype=x.dtype),
131
+ self.sin_cached[:seq_len].to(dtype=x.dtype),
132
+ )
133
+
134
+
135
+ # Copied from transformers.models.llama.modeling_llama.rotate_half
136
+ def rotate_half(x):
137
+ """Rotates half the hidden dims of the input."""
138
+ x1 = x[..., : x.shape[-1] // 2]
139
+ x2 = x[..., x.shape[-1] // 2 :]
140
+ return torch.cat((-x2, x1), dim=-1)
141
+
142
+
143
+ # Copied from transformers.models.mistral.modeling_mistral.apply_rotary_pos_emb
144
+ def apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1):
145
+ """Applies Rotary Position Embedding to the query and key tensors.
146
+
147
+ Args:
148
+ q (`torch.Tensor`): The query tensor.
149
+ k (`torch.Tensor`): The key tensor.
150
+ cos (`torch.Tensor`): The cosine part of the rotary embedding.
151
+ sin (`torch.Tensor`): The sine part of the rotary embedding.
152
+ position_ids (`torch.Tensor`):
153
+ The position indices of the tokens corresponding to the query and key tensors. For example, this can be
154
+ used to pass offsetted position ids when working with a KV-cache.
155
+ unsqueeze_dim (`int`, *optional*, defaults to 1):
156
+ The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
157
+ sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
158
+ that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
159
+ k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
160
+ cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
161
+ the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
162
+ Returns:
163
+ `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
164
+ """
165
+ cos = cos[position_ids].unsqueeze(unsqueeze_dim)
166
+ sin = sin[position_ids].unsqueeze(unsqueeze_dim)
167
+ q_embed = (q * cos) + (rotate_half(q) * sin)
168
+ k_embed = (k * cos) + (rotate_half(k) * sin)
169
+ return q_embed, k_embed
170
+
171
+
172
+ # Copied from transformers.models.llama.modeling_llama.repeat_kv
173
+ def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
174
+ """
175
+ This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
176
+ num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
177
+ """
178
+ batch, num_key_value_heads, slen, head_dim = hidden_states.shape
179
+ if n_rep == 1:
180
+ return hidden_states
181
+ hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
182
+ return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
183
+
184
+
185
+ class GrokAttention(nn.Module):
186
+ """
187
+ Multi-headed attention from 'Attention Is All You Need' paper.
188
+ """
189
+
190
+ def __init__(self, config: GrokConfig, layer_idx: Optional[int] = None):
191
+ super().__init__()
192
+ self.config = config
193
+ self.layer_idx = layer_idx
194
+ if layer_idx is None:
195
+ logger.warning_once(
196
+ f"Instantiating {self.__class__.__name__} without passing a `layer_idx` is not recommended and will "
197
+ "lead to errors during the forward call if caching is used. Please make sure to provide a `layer_idx` "
198
+ "when creating this class."
199
+ )
200
+
201
+ self.hidden_size = config.hidden_size
202
+ self.num_heads = config.num_attention_heads
203
+ self.head_dim = self.hidden_size // self.num_heads
204
+ self.num_key_value_heads = config.num_key_value_heads
205
+ self.num_key_value_groups = self.num_heads // self.num_key_value_heads
206
+ self.max_position_embeddings = config.max_position_embeddings
207
+ self.rope_theta = config.rope_theta
208
+ self.attn_output_multiplier = config.attn_output_multiplier
209
+ self.is_causal = True
210
+ self.attention_dropout = config.attention_dropout
211
+
212
+ if (self.head_dim * self.num_heads) != self.hidden_size:
213
+ raise ValueError(
214
+ f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"
215
+ f" and `num_heads`: {self.num_heads})."
216
+ )
217
+ self.query = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)
218
+ self.key = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
219
+ self.value = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
220
+ self.linear = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)
221
+
222
+ self.rotary_emb = GrokRotaryEmbedding(
223
+ self.head_dim,
224
+ max_position_embeddings=self.max_position_embeddings,
225
+ base=self.rope_theta,
226
+ )
227
+
228
+ def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
229
+ return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
230
+
231
+ def forward(
232
+ self,
233
+ hidden_states: torch.Tensor,
234
+ attention_mask: Optional[torch.Tensor] = None,
235
+ position_ids: Optional[torch.LongTensor] = None,
236
+ past_key_value: Optional[Cache] = None,
237
+ output_attentions: bool = False,
238
+ use_cache: bool = False,
239
+ **kwargs,
240
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
241
+ if "padding_mask" in kwargs:
242
+ warnings.warn(
243
+ "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"
244
+ )
245
+ bsz, q_len, _ = hidden_states.size()
246
+
247
+ query_states = self.query(hidden_states)
248
+ key_states = self.key(hidden_states)
249
+ value_states = self.value(hidden_states)
250
+
251
+ query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
252
+ key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
253
+ value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
254
+
255
+ kv_seq_len = key_states.shape[-2]
256
+ if past_key_value is not None:
257
+ if self.layer_idx is None:
258
+ raise ValueError(
259
+ f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} "
260
+ "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class "
261
+ "with a layer index."
262
+ )
263
+ kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
264
+ cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
265
+ query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
266
+
267
+ if past_key_value is not None:
268
+ cache_kwargs = {"sin": sin, "cos": cos} # Specific to RoPE models
269
+ key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
270
+
271
+ # repeat k/v heads if n_kv_heads < n_heads
272
+ key_states = repeat_kv(key_states, self.num_key_value_groups)
273
+ value_states = repeat_kv(value_states, self.num_key_value_groups)
274
+
275
+ attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) * self.attn_output_multiplier
276
+ attn_logits = 30 * torch.tanh(attn_weights / 30)
277
+
278
+ if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len):
279
+ raise ValueError(
280
+ f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is"
281
+ f" {attn_weights.size()}"
282
+ )
283
+
284
+ if attention_mask is not None:
285
+ if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
286
+ raise ValueError(
287
+ f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
288
+ )
289
+
290
+ attn_weights = attn_weights + attention_mask
291
+
292
+ # upcast attention to fp32
293
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
294
+ attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training)
295
+ attn_output = torch.matmul(attn_weights, value_states)
296
+
297
+ if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):
298
+ raise ValueError(
299
+ f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"
300
+ f" {attn_output.size()}"
301
+ )
302
+
303
+ attn_output = attn_output.transpose(1, 2).contiguous()
304
+ attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
305
+
306
+ attn_output = self.linear(attn_output)
307
+
308
+ if not output_attentions:
309
+ attn_weights = None
310
+
311
+ return attn_output, attn_weights, past_key_value
312
+
313
+
314
+ class GrokBlockSparseTop2MLP(nn.Module):
315
+ def __init__(self, config: GrokConfig):
316
+ super().__init__()
317
+ self.ffn_dim = config.intermediate_size
318
+ self.hidden_dim = config.hidden_size
319
+
320
+ self.linear_v = nn.Linear(self.hidden_dim, self.ffn_dim, bias=False)
321
+ self.linear_1 = nn.Linear(self.ffn_dim, self.hidden_dim, bias=False)
322
+ self.linear = nn.Linear(self.hidden_dim, self.ffn_dim, bias=False)
323
+
324
+ self.act_fn = ACT2FN[config.hidden_act]
325
+
326
+ def forward(self, hidden_states):
327
+ current_hidden_states = self.act_fn(self.linear(hidden_states)) * self.linear_v(hidden_states)
328
+ current_hidden_states = self.linear_1(current_hidden_states)
329
+ return current_hidden_states
330
+
331
+
332
+ class GrokDecoderLayer(nn.Module):
333
+ def __init__(self, config: GrokConfig, layer_idx: int):
334
+ super().__init__()
335
+ self.hidden_size = config.hidden_size
336
+ self.ffn_dim = config.intermediate_size
337
+ self.num_experts = config.num_local_experts
338
+ self.top_k = config.num_experts_per_tok
339
+
340
+ self.multi_head_attention = GrokAttention(config, layer_idx)
341
+ self.router = nn.Linear(self.hidden_size, self.num_experts, bias=False)
342
+ self.moe = nn.ModuleList([GrokBlockSparseTop2MLP(config) for _ in range(self.num_experts)])
343
+
344
+ self.rms_norm = GrokRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
345
+ self.rms_norm_1 = GrokRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
346
+ self.rms_norm_2 = GrokRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
347
+ self.rms_norm_3 = GrokRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
348
+
349
+ def forward(
350
+ self,
351
+ hidden_states: torch.Tensor,
352
+ attention_mask: Optional[torch.Tensor] = None,
353
+ position_ids: Optional[torch.LongTensor] = None,
354
+ past_key_value: Optional[Tuple[torch.Tensor]] = None,
355
+ output_attentions: Optional[bool] = False,
356
+ output_router_logits: Optional[bool] = False,
357
+ use_cache: Optional[bool] = False,
358
+ **kwargs,
359
+ ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
360
+ if "padding_mask" in kwargs:
361
+ warnings.warn(
362
+ "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"
363
+ )
364
+ """
365
+ Args:
366
+ hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
367
+ attention_mask (`torch.FloatTensor`, *optional*): attention mask of size
368
+ `(batch, sequence_length)` where padding elements are indicated by 0.
369
+ past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
370
+ output_attentions (`bool`, *optional*):
371
+ Whether or not to return the attentions tensors of all attention layers. See `attentions` under
372
+ returned tensors for more detail.
373
+ output_router_logits (`bool`, *optional*):
374
+ Whether or not to return the logits of all the routers. They are useful for computing the router loss, and
375
+ should not be returned during inference.
376
+ use_cache (`bool`, *optional*):
377
+ If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
378
+ (see `past_key_values`).
379
+ """
380
+
381
+ residual = hidden_states
382
+
383
+ hidden_states = self.rms_norm(hidden_states)
384
+
385
+ # Self Attention
386
+ hidden_states, self_attn_weights, present_key_value = self.multi_head_attention(
387
+ hidden_states=hidden_states,
388
+ attention_mask=attention_mask,
389
+ position_ids=position_ids,
390
+ past_key_value=past_key_value,
391
+ output_attentions=output_attentions,
392
+ use_cache=use_cache,
393
+ )
394
+ hidden_states = residual + self.rms_norm_1(hidden_states)
395
+
396
+ # Fully Connected
397
+ residual = hidden_states
398
+ hidden_states = self.rms_norm_2(hidden_states)
399
+
400
+ batch_size, sequence_length, hidden_dim = hidden_states.shape
401
+ hidden_states = hidden_states.view(-1, hidden_dim)
402
+ # router_logits: (batch * sequence_length, n_experts)
403
+ router_logits = self.router(hidden_states)
404
+
405
+ routing_weights = F.softmax(router_logits, dim=1, dtype=torch.float)
406
+ routing_weights, selected_experts = torch.topk(routing_weights, self.top_k, dim=-1)
407
+ # we cast back to the input dtype
408
+ routing_weights = routing_weights.to(hidden_states.dtype)
409
+
410
+ final_hidden_states = torch.zeros(
411
+ (batch_size * sequence_length, hidden_dim), dtype=hidden_states.dtype, device=hidden_states.device
412
+ )
413
+
414
+ # One hot encode the selected experts to create an expert mask
415
+ # this will be used to easily index which expert is going to be sollicitated
416
+ expert_mask = torch.nn.functional.one_hot(selected_experts, num_classes=self.num_experts).permute(2, 1, 0)
417
+
418
+ # Loop over all available experts in the model and perform the computation on each expert
419
+ for expert_idx in range(self.num_experts):
420
+ expert_layer = self.moe[expert_idx]
421
+ idx, top_x = torch.where(expert_mask[expert_idx])
422
+
423
+ if top_x.shape[0] == 0:
424
+ continue
425
+
426
+ # in torch it is faster to index using lists than torch tensors
427
+ top_x_list = top_x.tolist()
428
+ idx_list = idx.tolist()
429
+
430
+ # Index the correct hidden states and compute the expert hidden state for
431
+ # the current expert. We need to make sure to multiply the output hidden
432
+ # states by `routing_weights` on the corresponding tokens (top-1 and top-2)
433
+ current_state = hidden_states[None, top_x_list].reshape(-1, hidden_dim)
434
+ current_hidden_states = expert_layer(current_state) * routing_weights[top_x_list, idx_list, None]
435
+
436
+ # However `index_add_` only support torch tensors for indexing so we'll use
437
+ # the `top_x` tensor here.
438
+ final_hidden_states.index_add_(0, top_x, current_hidden_states.to(hidden_states.dtype))
439
+ hidden_states = final_hidden_states.reshape(batch_size, sequence_length, hidden_dim)
440
+
441
+ hidden_states = residual + self.rms_norm_3(hidden_states)
442
+
443
+ outputs = (hidden_states,)
444
+
445
+ if output_attentions:
446
+ outputs += (self_attn_weights,)
447
+
448
+ if use_cache:
449
+ outputs += (present_key_value,)
450
+
451
+ if output_router_logits:
452
+ outputs += (router_logits,)
453
+
454
+ return outputs
455
+
456
+
457
+ # Copied from transformers.models.mistral.modeling_mistral.MistralPreTrainedModel with Mistral->Grok
458
+ class GrokPreTrainedModel(PreTrainedModel):
459
+ config_class = GrokConfig
460
+ base_model_prefix = "transformer"
461
+ supports_gradient_checkpointing = True
462
+ _no_split_modules = ["GrokDecoderLayer"]
463
+ _skip_keys_device_placement = "past_key_values"
464
+ _keys_to_ignore_on_load_missing = [r"lm_head.*."]
465
+ _supports_flash_attn_2 = False
466
+ _supports_sdpa = False
467
+
468
+ def _init_weights(self, module):
469
+ pass
470
+
471
+
472
+ # Copied from transformers.models.mistral.modeling_mistral.MistralModel with Mistral->Grok
473
+ class GrokModel(GrokPreTrainedModel):
474
+ """
475
+ Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`GrokDecoderLayer`]
476
+
477
+ Args:
478
+ config: GrokConfig
479
+ """
480
+
481
+ def __init__(self, config: GrokConfig):
482
+ super().__init__(config)
483
+ self.padding_idx = config.pad_token_id
484
+ self.vocab_size = config.vocab_size
485
+ self.embedding_multiplier_scale = config.embedding_multiplier_scale
486
+
487
+ self.in_out_embed = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
488
+ self.decoder_layer = nn.ModuleList(
489
+ [GrokDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
490
+ )
491
+ self.rms_norm = GrokRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
492
+
493
+ self.gradient_checkpointing = False
494
+ # Initialize weights and apply final processing
495
+ self.post_init()
496
+
497
+ def get_input_embeddings(self):
498
+ return self.in_out_embed
499
+
500
+ def set_input_embeddings(self, value):
501
+ self.in_out_embed = value
502
+
503
+ # Ignore copy
504
+ def forward(
505
+ self,
506
+ input_ids: torch.LongTensor = None,
507
+ attention_mask: Optional[torch.Tensor] = None,
508
+ position_ids: Optional[torch.LongTensor] = None,
509
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
510
+ inputs_embeds: Optional[torch.FloatTensor] = None,
511
+ use_cache: Optional[bool] = None,
512
+ output_attentions: Optional[bool] = None,
513
+ output_hidden_states: Optional[bool] = None,
514
+ output_router_logits: Optional[bool] = None,
515
+ return_dict: Optional[bool] = None,
516
+ ) -> Union[Tuple, MoeModelOutputWithPast]:
517
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
518
+ output_router_logits = (
519
+ output_router_logits if output_router_logits is not None else self.config.output_router_logits
520
+ )
521
+ output_hidden_states = (
522
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
523
+ )
524
+ use_cache = use_cache if use_cache is not None else self.config.use_cache
525
+
526
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
527
+
528
+ # retrieve input_ids and inputs_embeds
529
+ if input_ids is not None and inputs_embeds is not None:
530
+ raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time")
531
+ elif input_ids is not None:
532
+ batch_size, seq_length = input_ids.shape
533
+ elif inputs_embeds is not None:
534
+ batch_size, seq_length, _ = inputs_embeds.shape
535
+ else:
536
+ raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds")
537
+
538
+ past_key_values_length = 0
539
+
540
+ if self.gradient_checkpointing and self.training:
541
+ if use_cache:
542
+ logger.warning_once(
543
+ "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
544
+ )
545
+ use_cache = False
546
+
547
+ if use_cache:
548
+ use_legacy_cache = not isinstance(past_key_values, Cache)
549
+ if use_legacy_cache:
550
+ past_key_values = DynamicCache.from_legacy_cache(past_key_values)
551
+ past_key_values_length = past_key_values.get_usable_length(seq_length)
552
+
553
+ if position_ids is None:
554
+ device = input_ids.device if input_ids is not None else inputs_embeds.device
555
+ position_ids = torch.arange(
556
+ past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=device
557
+ )
558
+ position_ids = position_ids.unsqueeze(0).view(-1, seq_length)
559
+ else:
560
+ position_ids = position_ids.view(-1, seq_length).long()
561
+
562
+ if inputs_embeds is None:
563
+ inputs_embeds = self.in_out_embed(input_ids)
564
+
565
+ attention_mask = _prepare_4d_causal_attention_mask(
566
+ attention_mask,
567
+ (batch_size, seq_length),
568
+ inputs_embeds,
569
+ past_key_values_length,
570
+ )
571
+
572
+ hidden_states = inputs_embeds
573
+ hidden_states *= self.embedding_multiplier_scale
574
+
575
+ # decoder layers
576
+ all_hidden_states = () if output_hidden_states else None
577
+ all_self_attns = () if output_attentions else None
578
+ all_router_logits = () if output_router_logits else None
579
+ next_decoder_cache = None
580
+
581
+ for decoder_layer in self.decoder_layer:
582
+ if output_hidden_states:
583
+ all_hidden_states += (hidden_states,)
584
+
585
+ if self.gradient_checkpointing and self.training:
586
+ layer_outputs = self._gradient_checkpointing_func(
587
+ decoder_layer.__call__,
588
+ hidden_states,
589
+ attention_mask,
590
+ position_ids,
591
+ past_key_values,
592
+ output_attentions,
593
+ output_router_logits,
594
+ use_cache,
595
+ )
596
+ else:
597
+ layer_outputs = decoder_layer(
598
+ hidden_states,
599
+ attention_mask=attention_mask,
600
+ position_ids=position_ids,
601
+ past_key_value=past_key_values,
602
+ output_attentions=output_attentions,
603
+ output_router_logits=output_router_logits,
604
+ use_cache=use_cache,
605
+ )
606
+
607
+ hidden_states = layer_outputs[0]
608
+
609
+ if use_cache:
610
+ next_decoder_cache = layer_outputs[2 if output_attentions else 1]
611
+
612
+ if output_attentions:
613
+ all_self_attns += (layer_outputs[1],)
614
+
615
+ if output_router_logits:
616
+ all_router_logits += (layer_outputs[-1],)
617
+
618
+ hidden_states = self.rms_norm(hidden_states)
619
+
620
+ # add hidden states from the last decoder layer
621
+ if output_hidden_states:
622
+ all_hidden_states += (hidden_states,)
623
+
624
+ next_cache = None
625
+ if use_cache:
626
+ next_cache = next_decoder_cache.to_legacy_cache() if use_legacy_cache else next_decoder_cache
627
+
628
+ if not return_dict:
629
+ return tuple(
630
+ v
631
+ for v in [hidden_states, next_cache, all_hidden_states, all_self_attns, all_router_logits]
632
+ if v is not None
633
+ )
634
+ return MoeModelOutputWithPast(
635
+ last_hidden_state=hidden_states,
636
+ past_key_values=next_cache,
637
+ hidden_states=all_hidden_states,
638
+ attentions=all_self_attns,
639
+ router_logits=all_router_logits,
640
+ )
641
+
642
+
643
+ class GrokForCausalLM(GrokPreTrainedModel):
644
+ _tied_weights_keys = ["lm_head.weight"]
645
+
646
+ def __init__(self, config):
647
+ super().__init__(config)
648
+ self.transformer = GrokModel(config)
649
+ self.vocab_size = config.vocab_size
650
+ self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
651
+ self.router_aux_loss_coef = config.router_aux_loss_coef
652
+ self.num_experts = config.num_local_experts
653
+ self.num_experts_per_tok = config.num_experts_per_tok
654
+ self.output_multiplier_scale = config.output_multiplier_scale
655
+ # Initialize weights and apply final processing
656
+ self.post_init()
657
+
658
+ def get_input_embeddings(self):
659
+ return self.transformer.in_out_embed
660
+
661
+ def set_input_embeddings(self, value):
662
+ self.transformer.in_out_embed = value
663
+
664
+ def get_output_embeddings(self):
665
+ return self.lm_head
666
+
667
+ def set_output_embeddings(self, new_embeddings):
668
+ self.lm_head = new_embeddings
669
+
670
+ def set_decoder(self, decoder):
671
+ self.transformer = decoder
672
+
673
+ def get_decoder(self):
674
+ return self.transformer
675
+
676
+ def _tie_weights(self):
677
+ self._tie_or_clone_weights(self.lm_head, self.get_input_embeddings())
678
+
679
+ # Ignore copy
680
+ def forward(
681
+ self,
682
+ input_ids: torch.LongTensor = None,
683
+ attention_mask: Optional[torch.Tensor] = None,
684
+ position_ids: Optional[torch.LongTensor] = None,
685
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
686
+ inputs_embeds: Optional[torch.FloatTensor] = None,
687
+ labels: Optional[torch.LongTensor] = None,
688
+ use_cache: Optional[bool] = None,
689
+ output_attentions: Optional[bool] = None,
690
+ output_hidden_states: Optional[bool] = None,
691
+ output_router_logits: Optional[bool] = None,
692
+ return_dict: Optional[bool] = None,
693
+ ) -> Union[Tuple, MoeCausalLMOutputWithPast]:
694
+ r"""
695
+ Args:
696
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
697
+ Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
698
+ config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
699
+ (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
700
+
701
+ Returns:
702
+
703
+ """
704
+
705
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
706
+ output_router_logits = (
707
+ output_router_logits if output_router_logits is not None else self.config.output_router_logits
708
+ )
709
+
710
+ output_hidden_states = (
711
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
712
+ )
713
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
714
+
715
+ # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
716
+ outputs = self.transformer(
717
+ input_ids=input_ids,
718
+ attention_mask=attention_mask,
719
+ position_ids=position_ids,
720
+ past_key_values=past_key_values,
721
+ inputs_embeds=inputs_embeds,
722
+ use_cache=use_cache,
723
+ output_attentions=output_attentions,
724
+ output_hidden_states=output_hidden_states,
725
+ output_router_logits=output_router_logits,
726
+ return_dict=return_dict,
727
+ )
728
+
729
+ hidden_states = outputs[0]
730
+ logits = self.lm_head(hidden_states)
731
+ logits = logits * self.output_multiplier_scale
732
+ logits = logits.float()
733
+
734
+ loss = None
735
+ if labels is not None:
736
+ # Shift so that tokens < n predict n
737
+ shift_logits = logits[..., :-1, :].contiguous()
738
+ shift_labels = labels[..., 1:].contiguous()
739
+ # Flatten the tokens
740
+ loss_fct = CrossEntropyLoss()
741
+ shift_logits = shift_logits.view(-1, self.config.vocab_size)
742
+ shift_labels = shift_labels.view(-1)
743
+ # Enable model parallelism
744
+ shift_labels = shift_labels.to(shift_logits.device)
745
+ loss = loss_fct(shift_logits, shift_labels)
746
+
747
+ aux_loss = None
748
+ if output_router_logits:
749
+ aux_loss = load_balancing_loss_func(
750
+ outputs.router_logits if return_dict else outputs[-1],
751
+ self.num_experts,
752
+ self.num_experts_per_tok,
753
+ attention_mask,
754
+ )
755
+ if labels is not None:
756
+ loss += self.router_aux_loss_coef * aux_loss.to(loss.device) # make sure to reside in the same device
757
+
758
+ if not return_dict:
759
+ output = (logits,) + outputs[1:]
760
+ if output_router_logits:
761
+ output = (aux_loss,) + output
762
+ return (loss,) + output if loss is not None else output
763
+
764
+ return MoeCausalLMOutputWithPast(
765
+ loss=loss,
766
+ aux_loss=aux_loss,
767
+ logits=logits,
768
+ past_key_values=outputs.past_key_values,
769
+ hidden_states=outputs.hidden_states,
770
+ attentions=outputs.attentions,
771
+ router_logits=outputs.router_logits,
772
+ )
773
+
774
+ def prepare_inputs_for_generation(
775
+ self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, **kwargs
776
+ ):
777
+ # Omit tokens covered by past_key_values
778
+ if past_key_values is not None:
779
+ if isinstance(past_key_values, Cache):
780
+ cache_length = past_key_values.get_seq_length()
781
+ past_length = past_key_values.seen_tokens
782
+ max_cache_length = past_key_values.get_max_length()
783
+ else:
784
+ cache_length = past_length = past_key_values[0][0].shape[2]
785
+ max_cache_length = None
786
+
787
+ # Keep only the unprocessed tokens:
788
+ # 1 - If the length of the attention_mask exceeds the length of input_ids, then we are in a setting where
789
+ # some of the inputs are exclusively passed as part of the cache (e.g. when passing input_embeds as
790
+ # input)
791
+ if attention_mask is not None and attention_mask.shape[1] > input_ids.shape[1]:
792
+ input_ids = input_ids[:, -(attention_mask.shape[1] - past_length) :]
793
+ # 2 - If the past_length is smaller than input_ids', then input_ids holds all input tokens. We can discard
794
+ # input_ids based on the past_length.
795
+ elif past_length < input_ids.shape[1]:
796
+ input_ids = input_ids[:, past_length:]
797
+ # 3 - Otherwise (past_length >= input_ids.shape[1]), let's assume input_ids only has unprocessed tokens.
798
+
799
+ # If we are about to go beyond the maximum cache length, we need to crop the input attention mask.
800
+ if (
801
+ max_cache_length is not None
802
+ and attention_mask is not None
803
+ and cache_length + input_ids.shape[1] > max_cache_length
804
+ ):
805
+ attention_mask = attention_mask[:, -max_cache_length:]
806
+
807
+ position_ids = kwargs.get("position_ids", None)
808
+ if attention_mask is not None and position_ids is None:
809
+ # create position_ids on the fly for batch generation
810
+ position_ids = attention_mask.long().cumsum(-1) - 1
811
+ position_ids.masked_fill_(attention_mask == 0, 1)
812
+ if past_key_values:
813
+ position_ids = position_ids[:, -input_ids.shape[1] :]
814
+
815
+ # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
816
+ if inputs_embeds is not None and past_key_values is None:
817
+ model_inputs = {"inputs_embeds": inputs_embeds}
818
+ else:
819
+ model_inputs = {"input_ids": input_ids}
820
+
821
+ model_inputs.update(
822
+ {
823
+ "position_ids": position_ids,
824
+ "past_key_values": past_key_values,
825
+ "use_cache": kwargs.get("use_cache"),
826
+ "attention_mask": attention_mask,
827
+ }
828
+ )
829
+ return model_inputs
830
+
831
+ @staticmethod
832
+ def _reorder_cache(past_key_values, beam_idx):
833
+ reordered_past = ()
834
+ for layer_past in past_key_values:
835
+ reordered_past += (
836
+ tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past),
837
+ )
838
+ return reordered_past
pytorch_model-00001-of-00019.bin ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:49bf50ef3d70afd8e7d199223f8472bc7296f2407ec2f43809ca62276ebaa9a9
3
+ size 24688084712
pytorch_model-00002-of-00019.bin ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:cd2e170b63a6532af4f1e7cde4fb51c34bd810aaa0ebbf33345f36f1d08fe9ac
3
+ size 35962479774
pytorch_model-00003-of-00019.bin ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:abb2c6f340f4d0f122c92479f0c2e9324d8f415af1bed4a645e6a374dd7693c3
3
+ size 32917463282
pytorch_model-00004-of-00019.bin ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:b9fdff65e71dc305f23181d18ea28b35c5b740f3d624a6d7e92495ccb7cb7a07
3
+ size 35962479902
pytorch_model-00005-of-00019.bin ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:88ec608b37942f333a7274017d4d7ef47e8c420ec3d1c251a0002e08da59a852
3
+ size 32917463410
pytorch_model-00006-of-00019.bin ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:52ad50ece14037cc7d7165d720b02c19619e45d354ff1f46b6b568975106e1a2
3
+ size 35962479902
pytorch_model-00007-of-00019.bin ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:15b5b209f1c94aed54d67b5bccd67be3696c042ce8af4d04317a834aa70c9feb
3
+ size 32917463410
pytorch_model-00008-of-00019.bin ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:1a78713c58d353f9835ac0913f80e45c57aa34a93813bba937e39ac37ed09c46
3
+ size 35962479902
pytorch_model-00009-of-00019.bin ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:ccaa848a8f0f1245019b69e8bf883f51d3a36ba9f251813cd48c28fca597b1c1
3
+ size 32917463410
pytorch_model-00010-of-00019.bin ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:0e11fc345ffdf8cb429deeca6ad79ebbf77a4684e845e6857a4003235c6e074e
3
+ size 35962479902
pytorch_model-00011-of-00019.bin ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:8d180e294abda6722fe27332a1d7311f6256c490db667ce8041345fd912f06e2
3
+ size 32917463410
pytorch_model-00012-of-00019.bin ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:aaaeeece29933f8ed8e3bd2ba431f80ee692f525b0a33ff9fbc9e4e319c354d4
3
+ size 35962479902
pytorch_model-00013-of-00019.bin ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:1dd78618c0245d6388b13a88eba6577b958a10cd88342a060c6ca9ad559c0feb
3
+ size 32917463410
pytorch_model-00014-of-00019.bin ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:0a0c17526baf9136881fd3438bf347e527bcdbb6c014bc25e2fea8d621ee85a6
3
+ size 35962479902
pytorch_model-00015-of-00019.bin ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:af1b44d316af3cb0d66e4e2f3703e0956a70b77b0f1754cea611aea499ec25a9
3
+ size 32917463410
pytorch_model-00016-of-00019.bin ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:3082fac8e8f9a2d0271b1a88ce1812b6ef41ff40cdfca7e4d9ce16d0241e36b1
3
+ size 35962479902
pytorch_model-00017-of-00019.bin ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:931d80d3cc5157b219ccaaa3c5b6ab5c3996300162cc6f6cf27780f2f0de7262
3
+ size 32917463410
pytorch_model-00018-of-00019.bin ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:74d3dfe9067c0185982e8a6fcf91612da6e118782eb92849452289af1bc75d06
3
+ size 35962479902
pytorch_model-00019-of-00019.bin ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:adf12c09aae1a83384d945c440524cd8e2e3ce4dffb83bad3f96857a8a55a770
3
+ size 19679996468
pytorch_model.bin.index.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:5c69d7cbad192fa2c9d14e2a77fbfdc11597c68907b043ddb0260e3d28eddd7f
3
+ size 2229219
tokenizer_config.json ADDED
@@ -0,0 +1,42 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "add_bos_token": false,
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
+ },
30
+ "additional_special_tokens": [],
31
+ "bos_token": "<s>",
32
+ "clean_up_tokenization_spaces": false,
33
+ "eos_token": "</s>",
34
+ "legacy": true,
35
+ "model_max_length": 1000000000000000019884624838656,
36
+ "pad_token": null,
37
+ "sp_model_kwargs": {},
38
+ "spaces_between_special_tokens": false,
39
+ "tokenizer_class": "LlamaTokenizer",
40
+ "unk_token": "<unk>",
41
+ "use_default_system_prompt": false
42
+ }