Shaltiel commited on
Commit
f4d185c
1 Parent(s): 6824dd0

Upload 11 files

Browse files
config.json ADDED
@@ -0,0 +1,37 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "architectures": [
3
+ "MegatronGPTForCausalLM"
4
+ ],
5
+ "attention_dropout": 0.0,
6
+ "auto_map": {
7
+ "AutoConfig": "configuration_megatron_gpt.MegatronGPTConfig",
8
+ "AutoModelForCausalLM": "modeling_megatron_gpt.MegatronGPTForCausalLM"
9
+ },
10
+ "bias": true,
11
+ "bos_token_id": 2,
12
+ "classifier_dropout": 0.0,
13
+ "eos_token_id": 2,
14
+ "hidden_act": "gelu",
15
+ "hidden_dropout": 0.0,
16
+ "hidden_size": 4096,
17
+ "initializer_range": 0.01,
18
+ "intermediate_size": 10880,
19
+ "layer_norm_eps": 1e-05,
20
+ "max_position_embeddings": 2048,
21
+ "model_type": "megatron_gpt",
22
+ "normalization": "layernorm1p",
23
+ "normalize_attention_scores": true,
24
+ "num_attention_heads": 32,
25
+ "num_hidden_layers": 32,
26
+ "pad_token_id": 0,
27
+ "rope_scaling": null,
28
+ "rotary_emb_base": 10000,
29
+ "rotary_pct": 0.5,
30
+ "self_attention_relative_position_bias": true,
31
+ "tie_word_embeddings": false,
32
+ "tokenizer_class": "GPT2Tokenizer",
33
+ "torch_dtype": "float16",
34
+ "transformers_version": "4.31.0",
35
+ "use_cache": true,
36
+ "vocab_size": 56064
37
+ }
configuration_megatron_gpt.py ADDED
@@ -0,0 +1,176 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # coding=utf-8
2
+ # Copyright 2022 EleutherAI 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 implementation 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 Nemo Framework
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
+ """ MegatronGPT model configuration"""
21
+
22
+ from transformers.configuration_utils import PretrainedConfig
23
+ from transformers.utils import logging
24
+
25
+
26
+ logger = logging.get_logger(__name__)
27
+
28
+ class MegatronGPTConfig(PretrainedConfig):
29
+ r"""
30
+ This is the configuration class to store the configuration of a [`MegatronGPTModel`]. It is used to instantiate an
31
+ MegatronGPT model according to the specified arguments, defining the model architecture. Instantiating a configuration
32
+ with the defaults will yield a similar configuration to that of the MegatronGPT 1.4B parameter architecture.
33
+
34
+ Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
35
+ documentation from [`PretrainedConfig`] for more information.
36
+
37
+ Args:
38
+ vocab_size (`int`, *optional*, defaults to 50432):
39
+ Vocabulary size of the MegatronGPT model. Defines the number of different tokens that can be represented by the
40
+ `inputs_ids` passed when calling [`MegatronGPTModel`].
41
+ hidden_size (`int`, *optional*, defaults to 6144):
42
+ Dimension of the encoder layers and the pooler layer.
43
+ num_hidden_layers (`int`, *optional*, defaults to 44):
44
+ Number of hidden layers in the Transformer encoder.
45
+ num_attention_heads (`int`, *optional*, defaults to 64):
46
+ Number of attention heads for each attention layer in the Transformer encoder.
47
+ intermediate_size (`int`, *optional*, defaults to 24576):
48
+ Dimension of the "intermediate" (i.e., feed-forward) layer in the Transformer encoder.
49
+ hidden_act (`str` or `function`, *optional*, defaults to `"gelu"`):
50
+ The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
51
+ `"relu"`, `"selu"` and `"gelu_new"` are supported.
52
+ bias (`bool`, *optional*, defaults to True)
53
+ Whether or not to include a bias in every linear layer
54
+ rotary_pct (`float`, *optional*, defaults to 0.25):
55
+ percentage of hidden dimensions to allocate to rotary embeddings
56
+ rotary_emb_base (`int`, *optional*, defaults to 10000)
57
+ base for computing rotary embeddings frequency
58
+ attention_dropout (`float`, *optional*, defaults to 0.0):
59
+ The dropout ratio probability of the attention score.
60
+ hidden_dropout (`float`, *optional*, defaults to 0.0):
61
+ The dropout ratio of (1) the word embeddings, (2) the post-attention hidden states, and (3) the post-mlp
62
+ hidden states.
63
+ classifier_dropout (`float`, *optional*, defaults to 0.1):
64
+ Argument used when doing token classification, used in the model [`MegatronGPTForTokenClassification`].
65
+ The dropout ratio for the hidden layer.
66
+ max_position_embeddings (`int`, *optional*, defaults to 2048):
67
+ The maximum sequence length that this model might ever be used with. Typically set this to something large
68
+ just in case (e.g., 512 or 1024 or 2048).
69
+ normalize_attention_scores (`bool`, *optional*, defaults to `True`)
70
+ Whether to scale the output Q * K^T by 1 / sqrt(hidden_size_per_head).
71
+ initializer_range (`float`, *optional*, defaults to 1e-5):
72
+ The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
73
+ layer_norm_eps (`float`, *optional*, defaults to 1e-12):
74
+ The epsilon used by the layer normalization layers.
75
+ normalization (`string`, *optional*, defaults to `layernorm1p`)
76
+ The type of normalization to use for the LayerNorm layers, either `layernorm` or `layernorm1p`
77
+ use_cache (`bool`, *optional*, defaults to `True`):
78
+ Whether or not the model should return the last key/values attentions (not used by all models). Only
79
+ relevant if `config.is_decoder=True`.
80
+ self_attention_relative_position_bias (`bool`, *optional*, defaults to `True`):
81
+ Whether to calculate and apply the relative position bias within the attention function.
82
+ If this is False, then model.generate will require you to calculate the triangular attention
83
+ mask and pass it through in the attention mask.
84
+ skip_flash_attention (`bool`, *optional*, defaults to `False`):
85
+ When calculating attention, whether to attempt to use flash attention if it's installed, or to always skip and use the regular method.
86
+ rope_scaling (`Dict`, *optional*):
87
+ Dictionary containing the scaling configuration for the RoPE embeddings. Currently supports three scaling
88
+ strategies: linear and dynamic. Their scaling factor must be an float greater than 1. The expected format
89
+ is `{"type": strategy name, "factor": scaling factor}`. When using this flag, don't update
90
+ `max_position_embeddings` to the expected new maximum. See the following thread for more information on how
91
+ these scaling strategies behave:
92
+ https://www.reddit.com/r/LocalLLaMA/comments/14mrgpr/dynamically_scaled_rope_further_increases/. This is an
93
+ experimental feature, subject to breaking API changes in future versions.
94
+
95
+ """
96
+ model_type = "megatron_gpt"
97
+
98
+ def __init__(
99
+ self,
100
+ vocab_size=56064,
101
+ hidden_size=2048,
102
+ num_hidden_layers=24,
103
+ num_attention_heads=16,
104
+ intermediate_size=5440,
105
+ hidden_act="fast-swiglu",
106
+ bias=True,
107
+ rotary_pct=0.5,
108
+ rotary_emb_base=10000,
109
+ attention_dropout=0.0,
110
+ hidden_dropout=0.0,
111
+ classifier_dropout=0.0,
112
+ normalize_attention_scores=True,
113
+ max_position_embeddings=2048,
114
+ initializer_range=0.01,
115
+ layer_norm_eps=1e-5,
116
+ normalization='layernorm1p',
117
+ use_cache=True,
118
+ self_attention_relative_position_bias=True,
119
+ bos_token_id=0,
120
+ eos_token_id=2,
121
+ tie_word_embeddings=False,
122
+ rope_scaling=None,
123
+ skip_flash_attention=False,
124
+ **kwargs,
125
+ ):
126
+ super().__init__(bos_token_id=bos_token_id, eos_token_id=eos_token_id, **kwargs)
127
+ self.vocab_size = vocab_size
128
+ self.max_position_embeddings = max_position_embeddings
129
+ self.hidden_size = hidden_size
130
+ self.num_hidden_layers = num_hidden_layers
131
+ self.num_attention_heads = num_attention_heads
132
+ self.intermediate_size = intermediate_size
133
+ self.hidden_act = hidden_act
134
+ self.bias = bias
135
+ self.rotary_pct = rotary_pct
136
+ self.rotary_emb_base = rotary_emb_base
137
+ self.attention_dropout = attention_dropout
138
+ self.hidden_dropout = hidden_dropout
139
+ self.classifier_dropout = classifier_dropout
140
+ self.normalize_attention_scores = normalize_attention_scores
141
+ self.initializer_range = initializer_range
142
+ self.layer_norm_eps = layer_norm_eps
143
+ self.normalization = normalization
144
+ self.use_cache = use_cache
145
+ self.self_attention_relative_position_bias = self_attention_relative_position_bias
146
+ self.tie_word_embeddings = tie_word_embeddings
147
+ self.skip_flash_attention = skip_flash_attention
148
+ self.rope_scaling = rope_scaling
149
+ self._rope_scaling_validation()
150
+
151
+ if self.hidden_size % self.num_attention_heads != 0:
152
+ raise ValueError(
153
+ "The hidden size is not divisble by the number of attention heads! Make sure to update them!"
154
+ )
155
+
156
+ # Copied from transformers.models.llama.configuration_llama.LlamaConfig._rope_scaling_validation
157
+ def _rope_scaling_validation(self):
158
+ """
159
+ Validate the `rope_scaling` configuration.
160
+ """
161
+ if self.rope_scaling is None:
162
+ return
163
+
164
+ if not isinstance(self.rope_scaling, dict) or len(self.rope_scaling) != 2:
165
+ raise ValueError(
166
+ "`rope_scaling` must be a dictionary with with two fields, `name` and `factor`, "
167
+ f"got {self.rope_scaling}"
168
+ )
169
+ rope_scaling_type = self.rope_scaling.get("type", None)
170
+ rope_scaling_factor = self.rope_scaling.get("factor", None)
171
+ if rope_scaling_type is None or rope_scaling_type not in ["linear", "dynamic"]:
172
+ raise ValueError(
173
+ f"`rope_scaling`'s name field must be one of ['linear', 'dynamic'], got {rope_scaling_type}"
174
+ )
175
+ if rope_scaling_factor is None or not isinstance(rope_scaling_factor, float) or rope_scaling_factor <= 1.0:
176
+ raise ValueError(f"`rope_scaling`'s factor field must be an float > 1, got {rope_scaling_factor}")
generation_config.json ADDED
@@ -0,0 +1,7 @@
 
 
 
 
 
 
 
 
1
+ {
2
+ "_from_model_config": true,
3
+ "bos_token_id": 2,
4
+ "eos_token_id": 2,
5
+ "pad_token_id": 0,
6
+ "transformers_version": "4.31.0"
7
+ }
merges.txt ADDED
The diff for this file is too large to render. See raw diff
 
modeling_megatron_gpt.py ADDED
@@ -0,0 +1,1206 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # coding=utf-8
2
+ # Copyright 2022 EleutherAI 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 implementation 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 Nemo Framework
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
+
21
+ """ PyTorch MegatronGPT model."""
22
+
23
+ from dataclasses import dataclass
24
+ import math
25
+ from typing import Optional, Tuple, Union
26
+
27
+ import torch
28
+ import torch.utils.checkpoint
29
+ from torch import nn
30
+ from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
31
+
32
+ from transformers.activations import ACT2FN
33
+ from transformers.file_utils import (
34
+ add_start_docstrings,
35
+ add_start_docstrings_to_model_forward,
36
+ replace_return_docstrings,
37
+ )
38
+ from transformers.modeling_outputs import (
39
+ BaseModelOutputWithPast,
40
+ CausalLMOutputWithPast,
41
+ QuestionAnsweringModelOutput,
42
+ SequenceClassifierOutputWithPast,
43
+ TokenClassifierOutput,
44
+ )
45
+ from transformers.modeling_utils import PreTrainedModel
46
+ from transformers.utils import logging
47
+ # try to load using a relative path, but if it fails try loading it directly
48
+ from .configuration_megatron_gpt import MegatronGPTConfig
49
+
50
+ try:
51
+ from flash_attn.bert_padding import unpad_input, pad_input
52
+ from flash_attn import flash_attn_varlen_func as flash_attn_func
53
+ HAS_FLASH = True
54
+ except:
55
+ try:
56
+ from flash_attn.flash_attn_interface import flash_attn_unpadded_func as flash_attn_func
57
+ HAS_FLASH = True
58
+ except:
59
+ HAS_FLASH = False
60
+
61
+
62
+ def get_activation(act):
63
+ if act in ["gelu", "geglu", "fast-geglu"]:
64
+ act = 'gelu'
65
+ elif act in ["reglu", "fast-reglu"]:
66
+ act = 'relu'
67
+ elif act in ["swiglu", "fast-swiglu"]:
68
+ act = 'silu'
69
+ return ACT2FN[act]
70
+
71
+ logger = logging.get_logger(__name__)
72
+
73
+ _CONFIG_FOR_DOC = "MegatronGPTConfig"
74
+
75
+ @dataclass
76
+ class CausalLMOutputWithPastAndEncoding(CausalLMOutputWithPast):
77
+ encoding_states: Optional[torch.FloatTensor] = None
78
+
79
+ class MegatronGPTPreTrainedModel(PreTrainedModel):
80
+ """
81
+ An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
82
+ models.
83
+ """
84
+
85
+ config_class = MegatronGPTConfig
86
+ base_model_prefix = "megatron_gpt"
87
+ supports_gradient_checkpointing = True
88
+ _no_split_modules = ["MegatronGPTLayer"]
89
+ _skip_keys_device_placement = "past_key_values"
90
+
91
+ def _init_weights(self, module):
92
+ """Initialize the weights"""
93
+ if isinstance(module, nn.Linear):
94
+ module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
95
+ if module.bias is not None:
96
+ module.bias.data.zero_()
97
+ elif isinstance(module, nn.Embedding):
98
+ module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
99
+ if module.padding_idx is not None:
100
+ module.weight.data[module.padding_idx].zero_()
101
+ elif isinstance(module, nn.LayerNorm):
102
+ module.bias.data.zero_()
103
+ module.weight.data.fill_(1.0)
104
+
105
+ def _set_gradient_checkpointing(self, module, value=False):
106
+ if isinstance(module, MegatronGPTModel):
107
+ module.gradient_checkpointing = value
108
+
109
+
110
+ class MegatronGPTAttention(nn.Module):
111
+ def __init__(self, config):
112
+ super().__init__()
113
+ self.config = config
114
+ self.self_attention_relative_position_bias = config.self_attention_relative_position_bias
115
+ self.num_attention_heads = config.num_attention_heads
116
+ self.hidden_size = config.hidden_size
117
+ if self.hidden_size % self.num_attention_heads != 0:
118
+ raise ValueError(
119
+ "The hidden size is not divisble by the number of attention heads! Make sure to update them"
120
+ )
121
+ self.head_size = self.hidden_size // self.num_attention_heads
122
+ self.rotary_ndims = int(self.head_size * config.rotary_pct)
123
+ self._init_bias(config.max_position_embeddings)
124
+
125
+ self.register_buffer("masked_bias", torch.tensor(-1e9), persistent=False)
126
+ self._init_rope()
127
+
128
+ self.norm_factor_float = math.sqrt(self.head_size if config.normalize_attention_scores else 1.0)
129
+ self.register_buffer(
130
+ "norm_factor",
131
+ torch.tensor(self.norm_factor_float, dtype=torch.float32).to(torch.get_default_dtype()),
132
+ persistent=False,
133
+ )
134
+ self.query_key_value = nn.Linear(config.hidden_size, 3 * config.hidden_size, bias=config.bias)
135
+ self.dense = nn.Linear(config.hidden_size, config.hidden_size, bias=config.bias)
136
+ self.attention_dropout = nn.Dropout(config.attention_dropout)
137
+
138
+ def _init_bias(self, max_positions, device=None):
139
+ self.register_buffer(
140
+ "bias",
141
+ torch.tril(torch.ones((max_positions, max_positions), dtype=torch.bool)).view(
142
+ 1, 1, max_positions, max_positions
143
+ ),
144
+ persistent=False,
145
+ )
146
+ if device is not None:
147
+ self.bias = self.bias.to(device)
148
+
149
+ def _init_rope(self):
150
+ if self.config.rope_scaling is None:
151
+ self.rotary_emb = MegatronGPTRotaryEmbedding(
152
+ self.rotary_ndims, self.config.max_position_embeddings, base=self.config.rotary_emb_base
153
+ )
154
+ else:
155
+ scaling_type = self.config.rope_scaling["type"]
156
+ scaling_factor = self.config.rope_scaling["factor"]
157
+ if scaling_type == "linear":
158
+ self.rotary_emb = MegatronGPTLinearScalingRotaryEmbedding(
159
+ self.rotary_ndims,
160
+ self.config.max_position_embeddings,
161
+ base=self.config.rotary_emb_base,
162
+ scaling_factor=scaling_factor,
163
+ )
164
+ elif scaling_type == "dynamic":
165
+ self.rotary_emb = MegatronGPTDynamicNTKScalingRotaryEmbedding(
166
+ self.rotary_ndims,
167
+ self.config.max_position_embeddings,
168
+ base=self.config.rotary_emb_base,
169
+ scaling_factor=scaling_factor,
170
+ )
171
+ else:
172
+ raise ValueError(f"Unknown RoPE scaling type {scaling_type}")
173
+
174
+ def forward(
175
+ self,
176
+ hidden_states: torch.FloatTensor,
177
+ attention_mask: torch.FloatTensor,
178
+ position_ids: torch.LongTensor,
179
+ head_mask: Optional[torch.FloatTensor] = None,
180
+ layer_past: Optional[Tuple[torch.Tensor]] = None,
181
+ use_cache: Optional[bool] = False,
182
+ output_attentions: Optional[bool] = False,
183
+ ):
184
+ has_layer_past = layer_past is not None
185
+
186
+ # Compute QKV
187
+ # Attention heads [batch, seq_len, hidden_size]
188
+ # --> [batch, seq_len, (np * 3 * head_size)]
189
+ qkv = self.query_key_value(hidden_states)
190
+
191
+ # [batch, seq_len, (num_heads * 3 * head_size)]
192
+ # --> [batch, seq_len, num_heads, 3 * head_size]
193
+ new_qkv_shape = qkv.size()[:-1] + (self.num_attention_heads, 3 * self.head_size)
194
+ qkv = qkv.view(*new_qkv_shape)
195
+
196
+ # [batch, seq_len, num_attention_heads, 3 * head_size] --> 3 [batch, num_attention_heads, seq_len, head_size]
197
+ query = qkv[..., : self.head_size].permute(0, 2, 1, 3)
198
+ key = qkv[..., self.head_size : 2 * self.head_size].permute(0, 2, 1, 3)
199
+ value = qkv[..., 2 * self.head_size :].permute(0, 2, 1, 3)
200
+
201
+ # Compute rotary embeddings on rotary_ndims
202
+ query_rot = query[..., : self.rotary_ndims]
203
+ query_pass = query[..., self.rotary_ndims :]
204
+ key_rot = key[..., : self.rotary_ndims]
205
+ key_pass = key[..., self.rotary_ndims :]
206
+
207
+ # Compute token offset for rotary embeddings (when decoding)
208
+ seq_len = key.shape[-2]
209
+ if has_layer_past:
210
+ seq_len = seq_len + layer_past[0].shape[-2]
211
+ cos, sin = self.rotary_emb(value, seq_len=seq_len)
212
+ query, key = apply_rotary_pos_emb(query_rot, key_rot, cos, sin, position_ids)
213
+ query = torch.cat((query, query_pass), dim=-1)
214
+ key = torch.cat((key, key_pass), dim=-1)
215
+
216
+ # Cache QKV values
217
+ if has_layer_past:
218
+ past_key = layer_past[0]
219
+ past_value = layer_past[1]
220
+ key = torch.cat((past_key, key), dim=-2)
221
+ value = torch.cat((past_value, value), dim=-2)
222
+ present = (key, value) if use_cache else None
223
+
224
+ # Compute attention
225
+ if not HAS_FLASH or output_attentions or head_mask is not None or self.config.skip_flash_attention:
226
+ attn_output, attn_weights = self._attn(query, key, value, attention_mask, head_mask)
227
+ else:
228
+ attn_output = self._flash_attn(query, key, value, attention_mask)
229
+
230
+ # Reshape outputs
231
+ attn_output = self._merge_heads(attn_output, self.num_attention_heads, self.head_size)
232
+ attn_output = self.dense(attn_output)
233
+
234
+ outputs = (attn_output, present)
235
+ if output_attentions:
236
+ outputs += (attn_weights,)
237
+
238
+ return outputs
239
+
240
+ @classmethod
241
+ def _split_heads(cls, tensor, num_attention_heads, attn_head_size):
242
+ """
243
+ Splits hidden dim into attn_head_size and num_attention_heads
244
+ """
245
+ # tensor: [bs, seq_len, hidden_size]
246
+ new_shape = tensor.size()[:-1] + (num_attention_heads, attn_head_size)
247
+ # -> [bs, seq_len, num_attention_heads, attn_head_size]
248
+ tensor = tensor.view(new_shape)
249
+ # -> [bs, num_attention_heads, seq_len, attn_head_size]
250
+ tensor = tensor.permute(0, 2, 1, 3)
251
+ return tensor
252
+
253
+ @classmethod
254
+ def _merge_heads(cls, tensor, num_attention_heads, attn_head_size):
255
+ """
256
+ Merges attn_head_size dim and num_attn_heads dim into hidden dim
257
+ """
258
+ # tensor [bs, num_attention_heads, seq_len, attn_head_size]
259
+ tensor = tensor.permute(0, 2, 1, 3).contiguous()
260
+ # -> [bs, seq_len, num_attention_heads, attn_head_size]
261
+ tensor = tensor.view(tensor.size(0), tensor.size(1), num_attention_heads * attn_head_size)
262
+ # -> [bs, seq_len, hidden_size]
263
+ return tensor
264
+
265
+ def _flash_attn(self, query, key, value, attention_mask=None):
266
+ # q, k, v: [bs, num_attention_heads, seq_len, attn_head_size]
267
+ # compute causal mask from causal mask buffer
268
+ batch_size, num_attention_heads, query_seq_length, attn_head_size = query.size()
269
+
270
+ # transpose_for_scores_flash returns b s h d
271
+ query_layer = query.transpose(1, 2).half()
272
+ key_layer = key.transpose(1, 2).half()
273
+ value_layer = value.transpose(1, 2).half()
274
+
275
+ # fix the mask
276
+ attention_mask = (attention_mask == 0).int().squeeze(1).squeeze(1)
277
+ query_layer, query_indicies, cu_seqlens_q, max_seqlen_q = unpad_input(query_layer, attention_mask[:, -query_seq_length:])
278
+ key_layer, _, cu_seqlens_k, max_seqlen_k = unpad_input(key_layer, attention_mask)
279
+ value_layer, _, _, _ = unpad_input(value_layer, attention_mask)
280
+
281
+ # returns [batch * seq, nheads, headdim]
282
+ context_layer = flash_attn_func(query_layer, key_layer, value_layer,
283
+ cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k,
284
+ dropout_p=self.config.attention_dropout, softmax_scale=1 / self.norm_factor_float, causal=self.self_attention_relative_position_bias if max_seqlen_q > 1 else False)
285
+
286
+ # fix the shape to be [bs, num_attention_heads, seq_len, attn_head_size]
287
+ context_layer = pad_input(context_layer, query_indicies, batch_size, query_seq_length)
288
+ context_layer = context_layer.view(batch_size, query_seq_length, num_attention_heads, attn_head_size) \
289
+ .transpose(1, 2)
290
+
291
+ return context_layer.to(value.dtype)
292
+
293
+ def _attn(self, query, key, value, attention_mask=None, head_mask=None):
294
+ # q, k, v: [bs, num_attention_heads, seq_len, attn_head_size]
295
+ # compute causal mask from causal mask buffer
296
+ batch_size, num_attention_heads, query_length, attn_head_size = query.size()
297
+ key_length = key.size(-2)
298
+
299
+ # dynamically increase the causal mask with the key length, if needed.
300
+ if key_length > self.bias.shape[-1]:
301
+ self._init_bias(key_length, device=key.device)
302
+ causal_mask = self.bias[:, :, key_length - query_length : key_length, :key_length]
303
+
304
+ query = query.view(batch_size * num_attention_heads, query_length, attn_head_size)
305
+ key = key.view(batch_size * num_attention_heads, key_length, attn_head_size)
306
+ attn_scores = torch.zeros(
307
+ batch_size * num_attention_heads,
308
+ query_length,
309
+ key_length,
310
+ dtype=query.dtype,
311
+ device=key.device,
312
+ )
313
+ attn_scores = torch.baddbmm(
314
+ attn_scores,
315
+ query,
316
+ key.transpose(1, 2),
317
+ beta=0.0,
318
+ alpha=(torch.tensor(1.0, dtype=self.norm_factor.dtype, device=self.norm_factor.device) / self.norm_factor),
319
+ )
320
+ attn_scores = attn_scores.view(batch_size, num_attention_heads, query_length, key_length)
321
+
322
+ if self.self_attention_relative_position_bias:
323
+ mask_value = torch.finfo(attn_scores.dtype).min
324
+ # Need to be a tensor, otherwise we get error: `RuntimeError: expected scalar type float but found double`.
325
+ # Need to be on the same device, otherwise `RuntimeError: ..., x and y to be on the same device`
326
+ mask_value = torch.tensor(mask_value, dtype=attn_scores.dtype).to(attn_scores.device)
327
+ attn_scores = torch.where(causal_mask, attn_scores, mask_value)
328
+
329
+ if attention_mask is not None:
330
+ # Apply the attention mask
331
+ attn_scores = attn_scores + attention_mask
332
+
333
+ attn_weights = nn.functional.softmax(attn_scores, dim=-1)
334
+ attn_weights = attn_weights.to(value.dtype)
335
+
336
+ # Mask heads if we want to
337
+ if head_mask is not None:
338
+ attn_weights = attn_weights * head_mask
339
+
340
+ attn_weights = self.attention_dropout(attn_weights)
341
+
342
+ attn_output = torch.matmul(attn_weights, value)
343
+ return attn_output, attn_weights
344
+
345
+
346
+ def attention_mask_func(attention_scores, ltor_mask):
347
+ attention_scores.masked_fill_(~ltor_mask, torch.finfo(attention_scores.dtype).min)
348
+ return attention_scores
349
+
350
+
351
+ class MegatronGPTRotaryEmbedding(torch.nn.Module):
352
+ def __init__(self, dim, max_position_embeddings, base=10000, device=None):
353
+ super().__init__()
354
+
355
+ self.dim = dim
356
+ self.max_position_embeddings = max_position_embeddings
357
+ self.base = base
358
+ inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))
359
+ self.register_buffer("inv_freq", inv_freq)
360
+
361
+ # Build here to make `torch.jit.trace` work.
362
+ self._set_cos_sin_cache(seq_len=max_position_embeddings, device=self.inv_freq.device)
363
+
364
+ def _set_cos_sin_cache(self, seq_len, device):
365
+ self.max_seq_len_cached = seq_len
366
+ t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)
367
+
368
+ freqs = torch.einsum("i,j->ij", t, self.inv_freq)
369
+ # Different from paper, but it uses a different permutation in order to obtain the same calculation
370
+ emb = torch.cat((freqs, freqs), dim=-1)
371
+ self.cos_cached = emb.cos()[None, None, :, :]
372
+ self.sin_cached = emb.sin()[None, None, :, :]
373
+
374
+ def forward(self, x, seq_len=None):
375
+ # x: [bs, num_attention_heads, seq_len, head_size]
376
+ if seq_len > self.max_seq_len_cached:
377
+ self._set_cos_sin_cache(seq_len=seq_len, device=x.device)
378
+ return self.cos_cached[:seq_len, ...].to(x.device), self.sin_cached[:seq_len, ...].to(x.device)
379
+
380
+
381
+ class MegatronGPTLinearScalingRotaryEmbedding(MegatronGPTRotaryEmbedding):
382
+ """MegatronGPTRotaryEmbedding extended with linear scaling. Credits to the Reddit user /u/kaiokendev"""
383
+
384
+ def __init__(self, dim, max_position_embeddings, base=10000, device=None, scaling_factor=1.0):
385
+ self.scaling_factor = scaling_factor
386
+ super().__init__(dim, max_position_embeddings, base, device)
387
+
388
+ def _set_cos_sin_cache(self, seq_len, device):
389
+ self.max_seq_len_cached = seq_len
390
+ t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)
391
+ t = t / self.scaling_factor
392
+
393
+ freqs = torch.einsum("i,j->ij", t, self.inv_freq)
394
+ # Different from paper, but it uses a different permutation in order to obtain the same calculation
395
+ emb = torch.cat((freqs, freqs), dim=-1)
396
+ self.cos_cached = emb.cos()[None, None, :, :]
397
+ self.sin_cached = emb.sin()[None, None, :, :]
398
+
399
+
400
+ class MegatronGPTDynamicNTKScalingRotaryEmbedding(MegatronGPTRotaryEmbedding):
401
+ """MegatronGPTRotaryEmbedding extended with Dynamic NTK scaling. Credits to the Reddit users /u/bloc97 and /u/emozilla"""
402
+
403
+ def __init__(self, dim, max_position_embeddings, base=10000, device=None, scaling_factor=1.0):
404
+ self.scaling_factor = scaling_factor
405
+ super().__init__(dim, max_position_embeddings, base, device)
406
+
407
+ def _set_cos_sin_cache(self, seq_len, device):
408
+ self.max_seq_len_cached = seq_len
409
+
410
+ if seq_len > self.max_position_embeddings:
411
+ base = self.base * (
412
+ (self.scaling_factor * seq_len / self.max_position_embeddings) - (self.scaling_factor - 1)
413
+ ) ** (self.dim / (self.dim - 2))
414
+ inv_freq = 1.0 / (base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))
415
+ self.register_buffer("inv_freq", inv_freq)
416
+
417
+ t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)
418
+
419
+ freqs = torch.einsum("i,j->ij", t, self.inv_freq)
420
+ # Different from paper, but it uses a different permutation in order to obtain the same calculation
421
+ emb = torch.cat((freqs, freqs), dim=-1)
422
+ self.cos_cached = emb.cos()[None, None, :, :]
423
+ self.sin_cached = emb.sin()[None, None, :, :]
424
+
425
+
426
+ def rotate_half(x):
427
+ """Rotates half the hidden dims of the input."""
428
+ x1 = x[..., : x.shape[-1] // 2]
429
+ x2 = x[..., x.shape[-1] // 2 :]
430
+ return torch.cat((-x2, x1), dim=-1)
431
+
432
+
433
+ def apply_rotary_pos_emb(q, k, cos, sin, position_ids):
434
+ gather_indices = position_ids[:, None, :, None] # [bs, 1, seq_len, 1]
435
+ gather_indices = gather_indices.repeat(1, cos.shape[1], 1, cos.shape[3])
436
+ cos = torch.gather(cos.repeat(gather_indices.shape[0], 1, 1, 1), 2, gather_indices)
437
+ sin = torch.gather(sin.repeat(gather_indices.shape[0], 1, 1, 1), 2, gather_indices)
438
+ q_embed = (q * cos) + (rotate_half(q) * sin)
439
+ k_embed = (k * cos) + (rotate_half(k) * sin)
440
+ return q_embed, k_embed
441
+
442
+
443
+ class MegatronGPTMLP(nn.Module):
444
+ def __init__(self, config):
445
+ super().__init__()
446
+ self.fast_glu_activation = config.hidden_act in ['fast-geglu', 'fast-swiglu', 'fast-reglu']
447
+ self.glu_activation_family = self.fast_glu_activation or config.hidden_act in ['geglu','reglu','swiglu']
448
+
449
+ self.dense_h_to_4h = nn.Linear(config.hidden_size, config.intermediate_size * (2 if self.fast_glu_activation else 1), bias=config.bias)
450
+ if config.hidden_act in ['geglu', 'reglu', 'swiglu']:
451
+ self.dense_h_to_4h_2 = nn.Linear(config.hidden_size, config.intermediate_size, bias=config.bias)
452
+
453
+ self.dense_4h_to_h = nn.Linear(config.intermediate_size, config.hidden_size, bias=config.bias)
454
+ self.act = get_activation(config.hidden_act)
455
+
456
+ def forward(self, hidden_states):
457
+ intermediate_states = self.dense_h_to_4h(hidden_states)
458
+ if self.glu_activation_family:
459
+ if self.fast_glu_activation:
460
+ intermediate_states, intermediate_states_2 = torch.chunk(intermediate_states, 2, dim=-1)
461
+ else:
462
+ intermediate_states_2 = self.dense_h_to_4h_2(hidden_states)
463
+
464
+ hidden_states = self.act(intermediate_states) * intermediate_states_2
465
+ else:
466
+ hidden_states = self.act(intermediate_states)
467
+ hidden_states = self.dense_4h_to_h(hidden_states)
468
+ return hidden_states
469
+
470
+
471
+ class MegatronGPTLayer(nn.Module):
472
+ def __init__(self, config, layer_idx):
473
+ super().__init__()
474
+ self.input_layernorm = MegatronGPTLayerNorm(config.normalization, config.hidden_size, eps=config.layer_norm_eps)
475
+ self.post_attention_layernorm = MegatronGPTLayerNorm(config.normalization, config.hidden_size, eps=config.layer_norm_eps)
476
+ self.post_attention_dropout = nn.Dropout(config.hidden_dropout)
477
+ self.post_mlp_dropout = nn.Dropout(config.hidden_dropout)
478
+ self.self_attention = MegatronGPTAttention(config)
479
+ self.mlp = MegatronGPTMLP(config)
480
+ self.layer_idx = layer_idx
481
+
482
+ def forward(
483
+ self,
484
+ hidden_states: Optional[torch.FloatTensor],
485
+ attention_mask: Optional[torch.FloatTensor] = None,
486
+ position_ids: Optional[torch.LongTensor] = None,
487
+ head_mask: Optional[torch.FloatTensor] = None,
488
+ use_cache: Optional[bool] = False,
489
+ layer_past: Optional[Tuple[torch.Tensor]] = None,
490
+ output_attentions: Optional[bool] = False,
491
+ ):
492
+ attention_layer_outputs = self.self_attention(
493
+ self.input_layernorm(hidden_states),
494
+ attention_mask=attention_mask,
495
+ position_ids=position_ids,
496
+ layer_past=layer_past,
497
+ head_mask=head_mask,
498
+ use_cache=use_cache,
499
+ output_attentions=output_attentions,
500
+ )
501
+ attn_output = attention_layer_outputs[0] # output_attn: attn_output, present, (attn_weights)
502
+ attn_output = self.post_attention_dropout(attn_output)
503
+
504
+ # pseudocode:
505
+ # x = x + attn(ln1(x))
506
+ # x = x + mlp(ln2(x))
507
+ attn_output = attn_output + hidden_states
508
+ mlp_output = self.mlp(self.post_attention_layernorm(attn_output))
509
+ mlp_output = self.post_mlp_dropout(mlp_output)
510
+ hidden_states = mlp_output + attn_output
511
+
512
+ outputs = attention_layer_outputs[1:]
513
+ if use_cache:
514
+ outputs = (hidden_states,) + outputs # hidden_states, present, (attn_weights)
515
+ else:
516
+ outputs = (hidden_states,) + outputs[1:] # hidden_states, (attn_weights)
517
+
518
+ return outputs
519
+
520
+ class MegatronGPTLayerNorm(torch.nn.LayerNorm):
521
+ def __init__(self, normalization, normalized_shape, eps=1e-05, elementwise_affine=True, device=None, dtype=None):
522
+ normalization = normalization.lower()
523
+ assert normalization in ['layernorm', 'layernorm1p', 'rmsnorm']
524
+ if normalization == 'rmsnorm':
525
+ torch.nn.Module.__init__(self)
526
+ self.weight = nn.Parameter(torch.ones(normalized_shape))
527
+ self.variance_epsilon = eps
528
+ else:
529
+ super().__init__(
530
+ normalized_shape=normalized_shape,
531
+ eps=eps,
532
+ elementwise_affine=elementwise_affine,
533
+ device=device,
534
+ dtype=dtype,
535
+ )
536
+ self.normalization = normalization
537
+
538
+ def forward(self, x):
539
+ if self.normalization == 'rmsnorm':
540
+ input_dtype = x.dtype
541
+ x = x.to(torch.float32)
542
+ variance = x.pow(2).mean(-1, keepdim=True)
543
+ x = x * torch.rsqrt(variance + self.variance_epsilon)
544
+ return self.weight * x.to(input_dtype)
545
+ else:
546
+ weight_bias = 1 if self.normalization == 'layernorm1p' else 0
547
+ return torch.nn.functional.layer_norm(
548
+ x, self.normalized_shape, self.weight + weight_bias, self.bias, self.eps
549
+ )
550
+
551
+
552
+
553
+ MEGATRON_GPT_START_DOCSTRING = r"""
554
+ This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) sub-class. Use
555
+ it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and
556
+ behavior.
557
+
558
+ Parameters:
559
+ config ([`~MegatronGPTConfig`]): Model configuration class with all the parameters of the model.
560
+ Initializing with a config file does not load the weights associated with the model, only the
561
+ configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
562
+ """
563
+
564
+ MEGATRON_GPT_INPUTS_DOCSTRING = r"""
565
+ Args:
566
+ input_ids (`torch.LongTensor` of shape `({0})`):
567
+ Indices of input sequence tokens in the vocabulary.
568
+
569
+ Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
570
+ [`PreTrainedTokenizer.__call__`] for details.
571
+
572
+ [What are input IDs?](../glossary#input-ids)
573
+ attention_mask (`torch.FloatTensor` of shape `({0})`, *optional*):
574
+ Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
575
+
576
+ - 1 for tokens that are **not masked**,
577
+ - 0 for tokens that are **masked**.
578
+
579
+ [What are attention masks?](../glossary#attention-mask)
580
+ position_ids (`torch.LongTensor` of shape `({0})`, *optional*):
581
+ Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
582
+ config.n_positions - 1]`.
583
+
584
+ [What are position IDs?](../glossary#position-ids)
585
+ head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*):
586
+ Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`:
587
+
588
+ - 1 indicates the head is **not masked**,
589
+ - 0 indicates the head is **masked**.
590
+
591
+ inputs_embeds (`torch.FloatTensor` of shape `({0}, hidden_size)`, *optional*):
592
+ Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
593
+ is useful if you want more control over how to convert *input_ids* indices into associated vectors than the
594
+ model's internal embedding lookup matrix.
595
+ output_attentions (`bool`, *optional*):
596
+ Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
597
+ tensors for more detail.
598
+ output_hidden_states (`bool`, *optional*):
599
+ Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
600
+ more detail.
601
+ return_dict (`bool`, *optional*):
602
+ Whether or not to return a [`~file_utils.ModelOutput`] instead of a plain tuple.
603
+ """
604
+
605
+
606
+ @add_start_docstrings(
607
+ "The bare MegatronGPT Model transformer outputting raw hidden-states without any specific head on top.",
608
+ MEGATRON_GPT_START_DOCSTRING,
609
+ )
610
+ class MegatronGPTModel(MegatronGPTPreTrainedModel):
611
+ def __init__(self, config):
612
+ super().__init__(config)
613
+ self.config = config
614
+
615
+ self.embed_in = nn.Embedding(config.vocab_size, config.hidden_size)
616
+ self.emb_dropout = nn.Dropout(config.hidden_dropout)
617
+ self.layers = nn.ModuleList([MegatronGPTLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)])
618
+ self.final_layernorm = MegatronGPTLayerNorm(config.normalization, config.hidden_size, eps=config.layer_norm_eps)
619
+
620
+ self.gradient_checkpointing = False
621
+
622
+ # Initialize weights and apply final processing
623
+ self.post_init()
624
+
625
+ def get_input_embeddings(self):
626
+ return self.embed_in
627
+
628
+ def set_input_embeddings(self, value):
629
+ self.embed_in = value
630
+
631
+ @add_start_docstrings_to_model_forward(MEGATRON_GPT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
632
+ def forward(
633
+ self,
634
+ input_ids: Optional[torch.LongTensor] = None,
635
+ attention_mask: Optional[torch.FloatTensor] = None,
636
+ position_ids: Optional[torch.LongTensor] = None,
637
+ head_mask: Optional[torch.FloatTensor] = None,
638
+ inputs_embeds: Optional[torch.FloatTensor] = None,
639
+ past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
640
+ use_cache: Optional[bool] = None,
641
+ output_attentions: Optional[bool] = None,
642
+ output_hidden_states: Optional[bool] = None,
643
+ return_dict: Optional[bool] = None,
644
+ ) -> Union[Tuple, BaseModelOutputWithPast]:
645
+ r"""
646
+ past_key_values (`tuple(tuple(torch.FloatTensor))` of length `config.n_layers` with each tuple having 4 tensors of shape `(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`):
647
+ Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding.
648
+ If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that
649
+ don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all
650
+ `decoder_input_ids` of shape `(batch_size, sequence_length)`.
651
+ use_cache (`bool`, *optional*):
652
+ If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
653
+ `past_key_values`).
654
+ """
655
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
656
+ output_hidden_states = (
657
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
658
+ )
659
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
660
+ use_cache = use_cache if use_cache is not None else self.config.use_cache
661
+
662
+ if input_ids is not None and inputs_embeds is not None:
663
+ raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
664
+ elif input_ids is not None:
665
+ input_shape = input_ids.size()
666
+ elif inputs_embeds is not None:
667
+ input_shape = inputs_embeds.size()[:-1]
668
+ else:
669
+ raise ValueError("You have to specify either input_ids or inputs_embeds")
670
+
671
+ batch_size, seq_length = input_shape
672
+
673
+ if past_key_values is None:
674
+ past_length = 0
675
+ past_key_values = tuple([None] * self.config.num_hidden_layers)
676
+ else:
677
+ past_length = past_key_values[0][0].size(-2)
678
+
679
+ if position_ids is None:
680
+ device = input_ids.device if input_ids is not None else inputs_embeds.device
681
+ position_ids = torch.arange(past_length, seq_length + past_length, dtype=torch.long, device=device)
682
+ position_ids = position_ids.unsqueeze(0).view(-1, seq_length)
683
+ else:
684
+ position_ids = position_ids.view(-1, seq_length).long()
685
+
686
+ # Attention mask.
687
+ if attention_mask is not None:
688
+ assert batch_size > 0, "batch_size has to be defined and > 0"
689
+ attention_mask = attention_mask.view(batch_size, -1)
690
+ # We create a 3D attention mask from a 2D tensor mask.
691
+ # Sizes are [batch_size, 1, 1, to_seq_length]
692
+ # So we can broadcast to [batch_size, num_heads, from_seq_length, to_seq_length]
693
+ # this attention mask is more simple than the triangular masking of causal attention
694
+ # used in OpenAI GPT, we just need to prepare the broadcast dimension here.
695
+ attention_mask = attention_mask[:, None, None, :]
696
+
697
+ # Since attention_mask is 1.0 for positions we want to attend and 0.0 for
698
+ # masked positions, this operation will create a tensor which is 0.0 for
699
+ # positions we want to attend and the dtype's smallest value for masked positions.
700
+ # Since we are adding it to the raw scores before the softmax, this is
701
+ # effectively the same as removing these entirely.
702
+ attention_mask = attention_mask.to(dtype=self.dtype) # fp16 compatibility
703
+ attention_mask = (1.0 - attention_mask) * torch.finfo(self.dtype).min
704
+
705
+ # Prepare head mask if needed
706
+ # 1.0 in head_mask indicate we keep the head
707
+ # attention_probs has shape bsz x n_heads x N x N
708
+ # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
709
+ # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
710
+ head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)
711
+
712
+ if inputs_embeds is None:
713
+ inputs_embeds = self.embed_in(input_ids)
714
+
715
+ hidden_states = self.emb_dropout(inputs_embeds)
716
+
717
+ if self.gradient_checkpointing and self.training:
718
+ if use_cache:
719
+ logger.warning(
720
+ "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
721
+ )
722
+ use_cache = False
723
+
724
+ presents = () if use_cache else None
725
+ all_attentions = () if output_attentions else None
726
+ all_hidden_states = () if output_hidden_states else None
727
+ for i, (layer, layer_past) in enumerate(zip(self.layers, past_key_values)):
728
+ if output_hidden_states:
729
+ all_hidden_states = all_hidden_states + (hidden_states,)
730
+
731
+ if self.gradient_checkpointing and self.training:
732
+
733
+ def create_custom_forward(module):
734
+ def custom_forward(*inputs):
735
+ # None for layer_past
736
+ return module(*inputs, use_cache, None, output_attentions)
737
+
738
+ return custom_forward
739
+
740
+ outputs = torch.utils.checkpoint.checkpoint(
741
+ create_custom_forward(layer),
742
+ hidden_states,
743
+ attention_mask,
744
+ position_ids,
745
+ head_mask[i],
746
+ )
747
+ else:
748
+ outputs = layer(
749
+ hidden_states,
750
+ attention_mask=attention_mask,
751
+ position_ids=position_ids,
752
+ head_mask=head_mask[i],
753
+ layer_past=layer_past,
754
+ use_cache=use_cache,
755
+ output_attentions=output_attentions,
756
+ )
757
+ hidden_states = outputs[0]
758
+ if use_cache is True:
759
+ presents = presents + (outputs[1],)
760
+ if output_attentions:
761
+ all_attentions = all_attentions + (outputs[2 if use_cache else 1],)
762
+
763
+ hidden_states = self.final_layernorm(hidden_states)
764
+ # Add last hidden state
765
+ if output_hidden_states:
766
+ all_hidden_states = all_hidden_states + (hidden_states,)
767
+
768
+ if not return_dict:
769
+ return tuple(v for v in [hidden_states, presents, all_hidden_states, all_attentions] if v is not None)
770
+
771
+ return BaseModelOutputWithPast(
772
+ last_hidden_state=hidden_states,
773
+ past_key_values=presents,
774
+ hidden_states=all_hidden_states,
775
+ attentions=all_attentions,
776
+ )
777
+
778
+
779
+ @add_start_docstrings(
780
+ """MegatronGPT Model with a `language modeling` head on top for CLM fine-tuning.""", MEGATRON_GPT_START_DOCSTRING
781
+ )
782
+ class MegatronGPTForCausalLM(MegatronGPTPreTrainedModel):
783
+ _tied_weights_keys = ["embed_out.weight"]
784
+
785
+ def __init__(self, config):
786
+ super().__init__(config)
787
+
788
+ self.megatron_gpt = MegatronGPTModel(config)
789
+ self.embed_out = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
790
+
791
+ # Initialize weights and apply final processing
792
+ self.post_init()
793
+
794
+ def get_output_embeddings(self):
795
+ return self.embed_out
796
+
797
+ def set_output_embeddings(self, new_embeddings):
798
+ self.embed_out = new_embeddings
799
+
800
+ @add_start_docstrings_to_model_forward(MEGATRON_GPT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
801
+ @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)
802
+ def forward(
803
+ self,
804
+ input_ids: Optional[torch.LongTensor] = None,
805
+ attention_mask: Optional[torch.FloatTensor] = None,
806
+ position_ids: Optional[torch.LongTensor] = None,
807
+ inputs_embeds: Optional[torch.FloatTensor] = None,
808
+ head_mask: Optional[torch.FloatTensor] = None,
809
+ past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
810
+ labels: Optional[torch.LongTensor] = None,
811
+ use_cache: Optional[bool] = None,
812
+ output_attentions: Optional[bool] = None,
813
+ output_hidden_states: Optional[bool] = None,
814
+ return_dict: Optional[bool] = None,
815
+ ) -> Union[Tuple, CausalLMOutputWithPastAndEncoding]:
816
+ r"""
817
+ past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
818
+ Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape
819
+ `(batch_size, num_heads, sequence_length, embed_size_per_head)`) and 2 additional tensors of shape
820
+ `(batch_size, num_heads, encoder_sequence_length, embed_size_per_head)`. The two additional tensors are
821
+ only required when the model is used as a decoder in a Sequence to Sequence model.
822
+
823
+ Contains pre-computed hidden-states (key and values in the self-attention blocks that can be used (see
824
+ `past_key_values` input) to speed up sequential decoding.
825
+
826
+ If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that
827
+ don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all
828
+ `decoder_input_ids` of shape `(batch_size, sequence_length)`.
829
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
830
+ Labels for computing the left-to-right language modeling loss (next word prediction). Indices should be in
831
+ `[-100, 0, ..., config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are
832
+ ignored (masked), the loss is only computed for the tokens with labels n `[0, ..., config.vocab_size]`.
833
+ use_cache (`bool`, *optional*):
834
+ If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
835
+ `past_key_values`).
836
+
837
+ Returns:
838
+ """
839
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
840
+
841
+ outputs = self.megatron_gpt(
842
+ input_ids,
843
+ attention_mask=attention_mask,
844
+ position_ids=position_ids,
845
+ head_mask=head_mask,
846
+ inputs_embeds=inputs_embeds,
847
+ past_key_values=past_key_values,
848
+ use_cache=use_cache,
849
+ output_attentions=output_attentions,
850
+ output_hidden_states=output_hidden_states,
851
+ return_dict=return_dict,
852
+ )
853
+
854
+ hidden_states = outputs[0]
855
+ lm_logits = self.embed_out(hidden_states)
856
+
857
+ lm_loss = None
858
+ if labels is not None:
859
+ # move labels to correct device to enable model parallelism
860
+ labels = labels.to(lm_logits.device)
861
+ # we are doing next-token prediction; shift prediction scores and input ids by one
862
+ shift_logits = lm_logits[:, :-1, :].contiguous()
863
+ labels = labels[:, 1:].contiguous()
864
+ loss_fct = CrossEntropyLoss()
865
+ lm_loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), labels.view(-1))
866
+
867
+ if not return_dict:
868
+ output = (lm_logits,) + outputs[1:]
869
+ return ((lm_loss,) + output) if lm_loss is not None else output
870
+
871
+ return CausalLMOutputWithPastAndEncoding(
872
+ loss=lm_loss,
873
+ logits=lm_logits,
874
+ past_key_values=outputs.past_key_values,
875
+ hidden_states=outputs.hidden_states,
876
+ attentions=outputs.attentions,
877
+ encoding_states=hidden_states
878
+ )
879
+
880
+ def prepare_inputs_for_generation(
881
+ self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, **kwargs
882
+ ):
883
+ input_shape = input_ids.shape
884
+
885
+ # cut decoder_input_ids if past is used
886
+ if past_key_values and past_key_values[0] is not None:
887
+ input_ids = input_ids[:, -1:]
888
+
889
+ position_ids = kwargs.get("position_ids", None)
890
+ if attention_mask is not None and position_ids is None:
891
+ # create position_ids on the fly for batch generation
892
+ position_ids = attention_mask.long().cumsum(-1) - 1
893
+ position_ids.masked_fill_(attention_mask == 0, 1)
894
+ if past_key_values:
895
+ position_ids = position_ids[:, -1].unsqueeze(-1)
896
+
897
+ # if model is used as a decoder in encoder-decoder model, the decoder attention mask is created on the fly
898
+ if attention_mask is None:
899
+ attention_mask = input_ids.new_ones(input_shape)
900
+
901
+ # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
902
+ if inputs_embeds is not None and past_key_values is None:
903
+ model_inputs = {"inputs_embeds": inputs_embeds}
904
+ else:
905
+ model_inputs = {"input_ids": input_ids}
906
+
907
+ model_inputs.update(
908
+ {
909
+ "attention_mask": attention_mask,
910
+ "past_key_values": past_key_values,
911
+ "position_ids": position_ids,
912
+ }
913
+ )
914
+
915
+ return model_inputs
916
+
917
+ def _reorder_cache(self, past_key_values, beam_idx):
918
+ reordered_past = ()
919
+ for layer_past in past_key_values:
920
+ reordered_past += (
921
+ tuple(past_state.index_select(0, beam_idx) for past_state in layer_past[:2]) + layer_past[2:],
922
+ )
923
+ return reordered_past
924
+
925
+
926
+ @add_start_docstrings(
927
+ """
928
+ The MegatronGPT Model transformer with a sequence classification head on top (linear layer).
929
+
930
+ [`MegatronGPTForSequenceClassification`] uses the last token in order to do the classification, as other causal models
931
+ (e.g. GPT-1) do.
932
+
933
+ Since it does classification on the last token, it requires to know the position of the last token. If a
934
+ `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If
935
+ no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the
936
+ padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in
937
+ each row of the batch).
938
+ """,
939
+ MEGATRON_GPT_START_DOCSTRING,
940
+ )
941
+ class MegatronGPTForSequenceClassification(MegatronGPTPreTrainedModel):
942
+ def __init__(self, config):
943
+ super().__init__(config)
944
+ self.num_labels = config.num_labels
945
+ self.megatron_gpt = MegatronGPTModel(config)
946
+ self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False)
947
+
948
+ # Initialize weights and apply final processing
949
+ self.post_init()
950
+
951
+ @add_start_docstrings_to_model_forward(MEGATRON_GPT_INPUTS_DOCSTRING)
952
+ def forward(
953
+ self,
954
+ input_ids: Optional[torch.LongTensor] = None,
955
+ attention_mask: Optional[torch.FloatTensor] = None,
956
+ position_ids: Optional[torch.LongTensor] = None,
957
+ inputs_embeds: Optional[torch.FloatTensor] = None,
958
+ head_mask: Optional[torch.FloatTensor] = None,
959
+ past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
960
+ labels: Optional[torch.LongTensor] = None,
961
+ use_cache: Optional[bool] = None,
962
+ output_attentions: Optional[bool] = None,
963
+ output_hidden_states: Optional[bool] = None,
964
+ return_dict: Optional[bool] = None,
965
+ ) -> Union[Tuple[torch.Tensor], SequenceClassifierOutputWithPast]:
966
+ r"""
967
+ labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
968
+ Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
969
+ config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
970
+ `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
971
+ """
972
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
973
+
974
+ outputs = self.megatron_gpt(
975
+ input_ids,
976
+ attention_mask=attention_mask,
977
+ position_ids=position_ids,
978
+ head_mask=head_mask,
979
+ inputs_embeds=inputs_embeds,
980
+ past_key_values=past_key_values,
981
+ use_cache=use_cache,
982
+ output_attentions=output_attentions,
983
+ output_hidden_states=output_hidden_states,
984
+ return_dict=return_dict,
985
+ )
986
+ hidden_states = outputs[0]
987
+ logits = self.score(hidden_states)
988
+
989
+ if input_ids is not None:
990
+ batch_size, sequence_length = input_ids.shape[:2]
991
+ else:
992
+ batch_size, sequence_length = inputs_embeds.shape[:2]
993
+
994
+ if self.config.pad_token_id is None and batch_size != 1:
995
+ raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.")
996
+ if self.config.pad_token_id is None:
997
+ sequence_lengths = -1
998
+ else:
999
+ if input_ids is not None:
1000
+ sequence_lengths = (torch.ne(input_ids, self.config.pad_token_id).sum(-1) - 1).to(logits.device)
1001
+ else:
1002
+ sequence_lengths = -1
1003
+ logger.warning(
1004
+ f"{self.__class__.__name__} will not detect padding tokens in `inputs_embeds`. Results may be "
1005
+ "unexpected if using padding tokens in conjunction with `inputs_embeds.`"
1006
+ )
1007
+
1008
+ pooled_logits = logits[torch.arange(batch_size, device=logits.device), sequence_lengths]
1009
+
1010
+ loss = None
1011
+ if labels is not None:
1012
+ labels = labels.to(logits.device)
1013
+ if self.config.problem_type is None:
1014
+ if self.num_labels == 1:
1015
+ self.config.problem_type = "regression"
1016
+ elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
1017
+ self.config.problem_type = "single_label_classification"
1018
+ else:
1019
+ self.config.problem_type = "multi_label_classification"
1020
+
1021
+ if self.config.problem_type == "regression":
1022
+ loss_fct = MSELoss()
1023
+ if self.num_labels == 1:
1024
+ loss = loss_fct(pooled_logits.squeeze(), labels.squeeze())
1025
+ else:
1026
+ loss = loss_fct(pooled_logits, labels)
1027
+ elif self.config.problem_type == "single_label_classification":
1028
+ loss_fct = CrossEntropyLoss()
1029
+ loss = loss_fct(pooled_logits.view(-1, self.num_labels), labels.view(-1))
1030
+ elif self.config.problem_type == "multi_label_classification":
1031
+ loss_fct = BCEWithLogitsLoss()
1032
+ loss = loss_fct(pooled_logits, labels)
1033
+ if not return_dict:
1034
+ output = (pooled_logits,) + outputs[1:]
1035
+ return ((loss,) + output) if loss is not None else output
1036
+
1037
+ return SequenceClassifierOutputWithPast(
1038
+ loss=loss,
1039
+ logits=pooled_logits,
1040
+ past_key_values=outputs.past_key_values,
1041
+ hidden_states=outputs.hidden_states,
1042
+ attentions=outputs.attentions,
1043
+ )
1044
+
1045
+
1046
+ class MegatronGPTForTokenClassification(MegatronGPTPreTrainedModel):
1047
+ def __init__(self, config):
1048
+ super().__init__(config)
1049
+ self.num_labels = config.num_labels
1050
+
1051
+ self.megatron_gpt = MegatronGPTModel(config)
1052
+ self.dropout = nn.Dropout(config.classifier_dropout)
1053
+ self.classifier = nn.Linear(config.hidden_size, config.num_labels)
1054
+
1055
+ # Initialize weights and apply final processing
1056
+ self.post_init()
1057
+
1058
+ @add_start_docstrings_to_model_forward(MEGATRON_GPT_INPUTS_DOCSTRING)
1059
+ def forward(
1060
+ self,
1061
+ input_ids: Optional[torch.LongTensor] = None,
1062
+ past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None,
1063
+ attention_mask: Optional[torch.FloatTensor] = None,
1064
+ token_type_ids: Optional[torch.LongTensor] = None,
1065
+ position_ids: Optional[torch.LongTensor] = None,
1066
+ head_mask: Optional[torch.FloatTensor] = None,
1067
+ inputs_embeds: Optional[torch.FloatTensor] = None,
1068
+ labels: Optional[torch.LongTensor] = None,
1069
+ use_cache: Optional[bool] = None,
1070
+ output_attentions: Optional[bool] = None,
1071
+ output_hidden_states: Optional[bool] = None,
1072
+ return_dict: Optional[bool] = None,
1073
+ ) -> Union[Tuple, TokenClassifierOutput]:
1074
+ r"""
1075
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
1076
+ Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
1077
+ config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
1078
+ `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
1079
+ """
1080
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
1081
+
1082
+ outputs = self.megatron_gpt(
1083
+ input_ids,
1084
+ past_key_values=past_key_values,
1085
+ attention_mask=attention_mask,
1086
+ position_ids=position_ids,
1087
+ head_mask=head_mask,
1088
+ inputs_embeds=inputs_embeds,
1089
+ use_cache=use_cache,
1090
+ output_attentions=output_attentions,
1091
+ output_hidden_states=output_hidden_states,
1092
+ return_dict=return_dict,
1093
+ )
1094
+
1095
+ hidden_states = outputs[0]
1096
+ hidden_states = self.dropout(hidden_states)
1097
+ logits = self.classifier(hidden_states)
1098
+
1099
+ loss = None
1100
+ if labels is not None:
1101
+ labels = labels.to(logits.device)
1102
+ loss_fct = CrossEntropyLoss()
1103
+ loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
1104
+
1105
+ if not return_dict:
1106
+ output = (logits,) + outputs[2:]
1107
+ return ((loss,) + output) if loss is not None else output
1108
+
1109
+ return TokenClassifierOutput(
1110
+ loss=loss,
1111
+ logits=logits,
1112
+ hidden_states=outputs.hidden_states,
1113
+ attentions=outputs.attentions,
1114
+ )
1115
+
1116
+
1117
+ @add_start_docstrings(
1118
+ """
1119
+ The Megatron-GPT Model transformer with a span classification head on top for extractive question-answering tasks like
1120
+ SQuAD (a linear layer on top of the hidden-states output to compute `span start logits` and `span end logits`).
1121
+ """,
1122
+ MEGATRON_GPT_START_DOCSTRING,
1123
+ )
1124
+ class MegatronGPTForQuestionAnswering(MegatronGPTPreTrainedModel):
1125
+ def __init__(self, config):
1126
+ super().__init__(config)
1127
+ self.num_labels = config.num_labels
1128
+ self.megatron_gpt = MegatronGPTModel(config)
1129
+ self.qa_outputs = nn.Linear(config.hidden_size, 2)
1130
+
1131
+ # Initialize weights and apply final processing
1132
+ self.post_init()
1133
+
1134
+ @add_start_docstrings_to_model_forward(MEGATRON_GPT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
1135
+ def forward(
1136
+ self,
1137
+ input_ids: Optional[torch.LongTensor] = None,
1138
+ attention_mask: Optional[torch.FloatTensor] = None,
1139
+ token_type_ids: Optional[torch.LongTensor] = None,
1140
+ position_ids: Optional[torch.LongTensor] = None,
1141
+ head_mask: Optional[torch.FloatTensor] = None,
1142
+ inputs_embeds: Optional[torch.FloatTensor] = None,
1143
+ start_positions: Optional[torch.LongTensor] = None,
1144
+ end_positions: Optional[torch.LongTensor] = None,
1145
+ output_attentions: Optional[bool] = None,
1146
+ output_hidden_states: Optional[bool] = None,
1147
+ return_dict: Optional[bool] = None,
1148
+ ) -> Union[Tuple, QuestionAnsweringModelOutput]:
1149
+ r"""
1150
+ start_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
1151
+ Labels for position (index) of the start of the labelled span for computing the token classification loss.
1152
+ Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
1153
+ are not taken into account for computing the loss.
1154
+ end_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
1155
+ Labels for position (index) of the end of the labelled span for computing the token classification loss.
1156
+ Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
1157
+ are not taken into account for computing the loss.
1158
+ """
1159
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
1160
+
1161
+ outputs = self.megatron_gpt(
1162
+ input_ids,
1163
+ attention_mask=attention_mask,
1164
+ position_ids=position_ids,
1165
+ head_mask=head_mask,
1166
+ inputs_embeds=inputs_embeds,
1167
+ output_attentions=output_attentions,
1168
+ output_hidden_states=output_hidden_states,
1169
+ return_dict=return_dict,
1170
+ )
1171
+
1172
+ sequence_output = outputs[0]
1173
+
1174
+ logits = self.qa_outputs(sequence_output)
1175
+ start_logits, end_logits = logits.split(1, dim=-1)
1176
+ start_logits = start_logits.squeeze(-1).contiguous()
1177
+ end_logits = end_logits.squeeze(-1).contiguous()
1178
+
1179
+ total_loss = None
1180
+ if start_positions is not None and end_positions is not None:
1181
+ # If we are on multi-GPU, split add a dimension
1182
+ if len(start_positions.size()) > 1:
1183
+ start_positions = start_positions.squeeze(-1).to(start_logits.device)
1184
+ if len(end_positions.size()) > 1:
1185
+ end_positions = end_positions.squeeze(-1).to(end_logits.device)
1186
+ # sometimes the start/end positions are outside our model inputs, we ignore these terms
1187
+ ignored_index = start_logits.size(1)
1188
+ start_positions = start_positions.clamp(0, ignored_index)
1189
+ end_positions = end_positions.clamp(0, ignored_index)
1190
+
1191
+ loss_fct = CrossEntropyLoss(ignore_index=ignored_index)
1192
+ start_loss = loss_fct(start_logits, start_positions)
1193
+ end_loss = loss_fct(end_logits, end_positions)
1194
+ total_loss = (start_loss + end_loss) / 2
1195
+
1196
+ if not return_dict:
1197
+ output = (start_logits, end_logits) + outputs[2:]
1198
+ return ((total_loss,) + output) if total_loss is not None else output
1199
+
1200
+ return QuestionAnsweringModelOutput(
1201
+ loss=total_loss,
1202
+ start_logits=start_logits,
1203
+ end_logits=end_logits,
1204
+ hidden_states=outputs.hidden_states,
1205
+ attentions=outputs.attentions,
1206
+ )
pytorch_model-00001-of-00002.bin ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:cc0f52eb712f34385164c02b0f6d45456574fd7b02c6aebebfd25f7213f1c071
3
+ size 9970836963
pytorch_model-00002-of-00002.bin ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:fc9549a9ce647fd7d3f798987224afd2a27d970d84379bd435ad6e8adb863b92
3
+ size 950158711
pytorch_model.bin.index.json ADDED
@@ -0,0 +1,427 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "metadata": {
3
+ "total_size": 10920847360
4
+ },
5
+ "weight_map": {
6
+ "embed_out.weight": "pytorch_model-00002-of-00002.bin",
7
+ "megatron_gpt.embed_in.weight": "pytorch_model-00001-of-00002.bin",
8
+ "megatron_gpt.final_layernorm.bias": "pytorch_model-00002-of-00002.bin",
9
+ "megatron_gpt.final_layernorm.weight": "pytorch_model-00002-of-00002.bin",
10
+ "megatron_gpt.layers.0.input_layernorm.bias": "pytorch_model-00001-of-00002.bin",
11
+ "megatron_gpt.layers.0.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
12
+ "megatron_gpt.layers.0.mlp.dense_4h_to_h.bias": "pytorch_model-00001-of-00002.bin",
13
+ "megatron_gpt.layers.0.mlp.dense_4h_to_h.weight": "pytorch_model-00001-of-00002.bin",
14
+ "megatron_gpt.layers.0.mlp.dense_h_to_4h.bias": "pytorch_model-00001-of-00002.bin",
15
+ "megatron_gpt.layers.0.mlp.dense_h_to_4h.weight": "pytorch_model-00001-of-00002.bin",
16
+ "megatron_gpt.layers.0.post_attention_layernorm.bias": "pytorch_model-00001-of-00002.bin",
17
+ "megatron_gpt.layers.0.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
18
+ "megatron_gpt.layers.0.self_attention.dense.bias": "pytorch_model-00001-of-00002.bin",
19
+ "megatron_gpt.layers.0.self_attention.dense.weight": "pytorch_model-00001-of-00002.bin",
20
+ "megatron_gpt.layers.0.self_attention.query_key_value.bias": "pytorch_model-00001-of-00002.bin",
21
+ "megatron_gpt.layers.0.self_attention.query_key_value.weight": "pytorch_model-00001-of-00002.bin",
22
+ "megatron_gpt.layers.0.self_attention.rotary_emb.inv_freq": "pytorch_model-00001-of-00002.bin",
23
+ "megatron_gpt.layers.1.input_layernorm.bias": "pytorch_model-00001-of-00002.bin",
24
+ "megatron_gpt.layers.1.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
25
+ "megatron_gpt.layers.1.mlp.dense_4h_to_h.bias": "pytorch_model-00001-of-00002.bin",
26
+ "megatron_gpt.layers.1.mlp.dense_4h_to_h.weight": "pytorch_model-00001-of-00002.bin",
27
+ "megatron_gpt.layers.1.mlp.dense_h_to_4h.bias": "pytorch_model-00001-of-00002.bin",
28
+ "megatron_gpt.layers.1.mlp.dense_h_to_4h.weight": "pytorch_model-00001-of-00002.bin",
29
+ "megatron_gpt.layers.1.post_attention_layernorm.bias": "pytorch_model-00001-of-00002.bin",
30
+ "megatron_gpt.layers.1.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
31
+ "megatron_gpt.layers.1.self_attention.dense.bias": "pytorch_model-00001-of-00002.bin",
32
+ "megatron_gpt.layers.1.self_attention.dense.weight": "pytorch_model-00001-of-00002.bin",
33
+ "megatron_gpt.layers.1.self_attention.query_key_value.bias": "pytorch_model-00001-of-00002.bin",
34
+ "megatron_gpt.layers.1.self_attention.query_key_value.weight": "pytorch_model-00001-of-00002.bin",
35
+ "megatron_gpt.layers.1.self_attention.rotary_emb.inv_freq": "pytorch_model-00001-of-00002.bin",
36
+ "megatron_gpt.layers.10.input_layernorm.bias": "pytorch_model-00001-of-00002.bin",
37
+ "megatron_gpt.layers.10.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
38
+ "megatron_gpt.layers.10.mlp.dense_4h_to_h.bias": "pytorch_model-00001-of-00002.bin",
39
+ "megatron_gpt.layers.10.mlp.dense_4h_to_h.weight": "pytorch_model-00001-of-00002.bin",
40
+ "megatron_gpt.layers.10.mlp.dense_h_to_4h.bias": "pytorch_model-00001-of-00002.bin",
41
+ "megatron_gpt.layers.10.mlp.dense_h_to_4h.weight": "pytorch_model-00001-of-00002.bin",
42
+ "megatron_gpt.layers.10.post_attention_layernorm.bias": "pytorch_model-00001-of-00002.bin",
43
+ "megatron_gpt.layers.10.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
44
+ "megatron_gpt.layers.10.self_attention.dense.bias": "pytorch_model-00001-of-00002.bin",
45
+ "megatron_gpt.layers.10.self_attention.dense.weight": "pytorch_model-00001-of-00002.bin",
46
+ "megatron_gpt.layers.10.self_attention.query_key_value.bias": "pytorch_model-00001-of-00002.bin",
47
+ "megatron_gpt.layers.10.self_attention.query_key_value.weight": "pytorch_model-00001-of-00002.bin",
48
+ "megatron_gpt.layers.10.self_attention.rotary_emb.inv_freq": "pytorch_model-00001-of-00002.bin",
49
+ "megatron_gpt.layers.11.input_layernorm.bias": "pytorch_model-00001-of-00002.bin",
50
+ "megatron_gpt.layers.11.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
51
+ "megatron_gpt.layers.11.mlp.dense_4h_to_h.bias": "pytorch_model-00001-of-00002.bin",
52
+ "megatron_gpt.layers.11.mlp.dense_4h_to_h.weight": "pytorch_model-00001-of-00002.bin",
53
+ "megatron_gpt.layers.11.mlp.dense_h_to_4h.bias": "pytorch_model-00001-of-00002.bin",
54
+ "megatron_gpt.layers.11.mlp.dense_h_to_4h.weight": "pytorch_model-00001-of-00002.bin",
55
+ "megatron_gpt.layers.11.post_attention_layernorm.bias": "pytorch_model-00001-of-00002.bin",
56
+ "megatron_gpt.layers.11.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
57
+ "megatron_gpt.layers.11.self_attention.dense.bias": "pytorch_model-00001-of-00002.bin",
58
+ "megatron_gpt.layers.11.self_attention.dense.weight": "pytorch_model-00001-of-00002.bin",
59
+ "megatron_gpt.layers.11.self_attention.query_key_value.bias": "pytorch_model-00001-of-00002.bin",
60
+ "megatron_gpt.layers.11.self_attention.query_key_value.weight": "pytorch_model-00001-of-00002.bin",
61
+ "megatron_gpt.layers.11.self_attention.rotary_emb.inv_freq": "pytorch_model-00001-of-00002.bin",
62
+ "megatron_gpt.layers.12.input_layernorm.bias": "pytorch_model-00001-of-00002.bin",
63
+ "megatron_gpt.layers.12.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
64
+ "megatron_gpt.layers.12.mlp.dense_4h_to_h.bias": "pytorch_model-00001-of-00002.bin",
65
+ "megatron_gpt.layers.12.mlp.dense_4h_to_h.weight": "pytorch_model-00001-of-00002.bin",
66
+ "megatron_gpt.layers.12.mlp.dense_h_to_4h.bias": "pytorch_model-00001-of-00002.bin",
67
+ "megatron_gpt.layers.12.mlp.dense_h_to_4h.weight": "pytorch_model-00001-of-00002.bin",
68
+ "megatron_gpt.layers.12.post_attention_layernorm.bias": "pytorch_model-00001-of-00002.bin",
69
+ "megatron_gpt.layers.12.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
70
+ "megatron_gpt.layers.12.self_attention.dense.bias": "pytorch_model-00001-of-00002.bin",
71
+ "megatron_gpt.layers.12.self_attention.dense.weight": "pytorch_model-00001-of-00002.bin",
72
+ "megatron_gpt.layers.12.self_attention.query_key_value.bias": "pytorch_model-00001-of-00002.bin",
73
+ "megatron_gpt.layers.12.self_attention.query_key_value.weight": "pytorch_model-00001-of-00002.bin",
74
+ "megatron_gpt.layers.12.self_attention.rotary_emb.inv_freq": "pytorch_model-00001-of-00002.bin",
75
+ "megatron_gpt.layers.13.input_layernorm.bias": "pytorch_model-00001-of-00002.bin",
76
+ "megatron_gpt.layers.13.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
77
+ "megatron_gpt.layers.13.mlp.dense_4h_to_h.bias": "pytorch_model-00001-of-00002.bin",
78
+ "megatron_gpt.layers.13.mlp.dense_4h_to_h.weight": "pytorch_model-00001-of-00002.bin",
79
+ "megatron_gpt.layers.13.mlp.dense_h_to_4h.bias": "pytorch_model-00001-of-00002.bin",
80
+ "megatron_gpt.layers.13.mlp.dense_h_to_4h.weight": "pytorch_model-00001-of-00002.bin",
81
+ "megatron_gpt.layers.13.post_attention_layernorm.bias": "pytorch_model-00001-of-00002.bin",
82
+ "megatron_gpt.layers.13.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
83
+ "megatron_gpt.layers.13.self_attention.dense.bias": "pytorch_model-00001-of-00002.bin",
84
+ "megatron_gpt.layers.13.self_attention.dense.weight": "pytorch_model-00001-of-00002.bin",
85
+ "megatron_gpt.layers.13.self_attention.query_key_value.bias": "pytorch_model-00001-of-00002.bin",
86
+ "megatron_gpt.layers.13.self_attention.query_key_value.weight": "pytorch_model-00001-of-00002.bin",
87
+ "megatron_gpt.layers.13.self_attention.rotary_emb.inv_freq": "pytorch_model-00001-of-00002.bin",
88
+ "megatron_gpt.layers.14.input_layernorm.bias": "pytorch_model-00001-of-00002.bin",
89
+ "megatron_gpt.layers.14.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
90
+ "megatron_gpt.layers.14.mlp.dense_4h_to_h.bias": "pytorch_model-00001-of-00002.bin",
91
+ "megatron_gpt.layers.14.mlp.dense_4h_to_h.weight": "pytorch_model-00001-of-00002.bin",
92
+ "megatron_gpt.layers.14.mlp.dense_h_to_4h.bias": "pytorch_model-00001-of-00002.bin",
93
+ "megatron_gpt.layers.14.mlp.dense_h_to_4h.weight": "pytorch_model-00001-of-00002.bin",
94
+ "megatron_gpt.layers.14.post_attention_layernorm.bias": "pytorch_model-00001-of-00002.bin",
95
+ "megatron_gpt.layers.14.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
96
+ "megatron_gpt.layers.14.self_attention.dense.bias": "pytorch_model-00001-of-00002.bin",
97
+ "megatron_gpt.layers.14.self_attention.dense.weight": "pytorch_model-00001-of-00002.bin",
98
+ "megatron_gpt.layers.14.self_attention.query_key_value.bias": "pytorch_model-00001-of-00002.bin",
99
+ "megatron_gpt.layers.14.self_attention.query_key_value.weight": "pytorch_model-00001-of-00002.bin",
100
+ "megatron_gpt.layers.14.self_attention.rotary_emb.inv_freq": "pytorch_model-00001-of-00002.bin",
101
+ "megatron_gpt.layers.15.input_layernorm.bias": "pytorch_model-00001-of-00002.bin",
102
+ "megatron_gpt.layers.15.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
103
+ "megatron_gpt.layers.15.mlp.dense_4h_to_h.bias": "pytorch_model-00001-of-00002.bin",
104
+ "megatron_gpt.layers.15.mlp.dense_4h_to_h.weight": "pytorch_model-00001-of-00002.bin",
105
+ "megatron_gpt.layers.15.mlp.dense_h_to_4h.bias": "pytorch_model-00001-of-00002.bin",
106
+ "megatron_gpt.layers.15.mlp.dense_h_to_4h.weight": "pytorch_model-00001-of-00002.bin",
107
+ "megatron_gpt.layers.15.post_attention_layernorm.bias": "pytorch_model-00001-of-00002.bin",
108
+ "megatron_gpt.layers.15.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
109
+ "megatron_gpt.layers.15.self_attention.dense.bias": "pytorch_model-00001-of-00002.bin",
110
+ "megatron_gpt.layers.15.self_attention.dense.weight": "pytorch_model-00001-of-00002.bin",
111
+ "megatron_gpt.layers.15.self_attention.query_key_value.bias": "pytorch_model-00001-of-00002.bin",
112
+ "megatron_gpt.layers.15.self_attention.query_key_value.weight": "pytorch_model-00001-of-00002.bin",
113
+ "megatron_gpt.layers.15.self_attention.rotary_emb.inv_freq": "pytorch_model-00001-of-00002.bin",
114
+ "megatron_gpt.layers.16.input_layernorm.bias": "pytorch_model-00001-of-00002.bin",
115
+ "megatron_gpt.layers.16.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
116
+ "megatron_gpt.layers.16.mlp.dense_4h_to_h.bias": "pytorch_model-00001-of-00002.bin",
117
+ "megatron_gpt.layers.16.mlp.dense_4h_to_h.weight": "pytorch_model-00001-of-00002.bin",
118
+ "megatron_gpt.layers.16.mlp.dense_h_to_4h.bias": "pytorch_model-00001-of-00002.bin",
119
+ "megatron_gpt.layers.16.mlp.dense_h_to_4h.weight": "pytorch_model-00001-of-00002.bin",
120
+ "megatron_gpt.layers.16.post_attention_layernorm.bias": "pytorch_model-00001-of-00002.bin",
121
+ "megatron_gpt.layers.16.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
122
+ "megatron_gpt.layers.16.self_attention.dense.bias": "pytorch_model-00001-of-00002.bin",
123
+ "megatron_gpt.layers.16.self_attention.dense.weight": "pytorch_model-00001-of-00002.bin",
124
+ "megatron_gpt.layers.16.self_attention.query_key_value.bias": "pytorch_model-00001-of-00002.bin",
125
+ "megatron_gpt.layers.16.self_attention.query_key_value.weight": "pytorch_model-00001-of-00002.bin",
126
+ "megatron_gpt.layers.16.self_attention.rotary_emb.inv_freq": "pytorch_model-00001-of-00002.bin",
127
+ "megatron_gpt.layers.17.input_layernorm.bias": "pytorch_model-00001-of-00002.bin",
128
+ "megatron_gpt.layers.17.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
129
+ "megatron_gpt.layers.17.mlp.dense_4h_to_h.bias": "pytorch_model-00001-of-00002.bin",
130
+ "megatron_gpt.layers.17.mlp.dense_4h_to_h.weight": "pytorch_model-00001-of-00002.bin",
131
+ "megatron_gpt.layers.17.mlp.dense_h_to_4h.bias": "pytorch_model-00001-of-00002.bin",
132
+ "megatron_gpt.layers.17.mlp.dense_h_to_4h.weight": "pytorch_model-00001-of-00002.bin",
133
+ "megatron_gpt.layers.17.post_attention_layernorm.bias": "pytorch_model-00001-of-00002.bin",
134
+ "megatron_gpt.layers.17.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
135
+ "megatron_gpt.layers.17.self_attention.dense.bias": "pytorch_model-00001-of-00002.bin",
136
+ "megatron_gpt.layers.17.self_attention.dense.weight": "pytorch_model-00001-of-00002.bin",
137
+ "megatron_gpt.layers.17.self_attention.query_key_value.bias": "pytorch_model-00001-of-00002.bin",
138
+ "megatron_gpt.layers.17.self_attention.query_key_value.weight": "pytorch_model-00001-of-00002.bin",
139
+ "megatron_gpt.layers.17.self_attention.rotary_emb.inv_freq": "pytorch_model-00001-of-00002.bin",
140
+ "megatron_gpt.layers.18.input_layernorm.bias": "pytorch_model-00001-of-00002.bin",
141
+ "megatron_gpt.layers.18.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
142
+ "megatron_gpt.layers.18.mlp.dense_4h_to_h.bias": "pytorch_model-00001-of-00002.bin",
143
+ "megatron_gpt.layers.18.mlp.dense_4h_to_h.weight": "pytorch_model-00001-of-00002.bin",
144
+ "megatron_gpt.layers.18.mlp.dense_h_to_4h.bias": "pytorch_model-00001-of-00002.bin",
145
+ "megatron_gpt.layers.18.mlp.dense_h_to_4h.weight": "pytorch_model-00001-of-00002.bin",
146
+ "megatron_gpt.layers.18.post_attention_layernorm.bias": "pytorch_model-00001-of-00002.bin",
147
+ "megatron_gpt.layers.18.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
148
+ "megatron_gpt.layers.18.self_attention.dense.bias": "pytorch_model-00001-of-00002.bin",
149
+ "megatron_gpt.layers.18.self_attention.dense.weight": "pytorch_model-00001-of-00002.bin",
150
+ "megatron_gpt.layers.18.self_attention.query_key_value.bias": "pytorch_model-00001-of-00002.bin",
151
+ "megatron_gpt.layers.18.self_attention.query_key_value.weight": "pytorch_model-00001-of-00002.bin",
152
+ "megatron_gpt.layers.18.self_attention.rotary_emb.inv_freq": "pytorch_model-00001-of-00002.bin",
153
+ "megatron_gpt.layers.19.input_layernorm.bias": "pytorch_model-00001-of-00002.bin",
154
+ "megatron_gpt.layers.19.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
155
+ "megatron_gpt.layers.19.mlp.dense_4h_to_h.bias": "pytorch_model-00001-of-00002.bin",
156
+ "megatron_gpt.layers.19.mlp.dense_4h_to_h.weight": "pytorch_model-00001-of-00002.bin",
157
+ "megatron_gpt.layers.19.mlp.dense_h_to_4h.bias": "pytorch_model-00001-of-00002.bin",
158
+ "megatron_gpt.layers.19.mlp.dense_h_to_4h.weight": "pytorch_model-00001-of-00002.bin",
159
+ "megatron_gpt.layers.19.post_attention_layernorm.bias": "pytorch_model-00001-of-00002.bin",
160
+ "megatron_gpt.layers.19.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
161
+ "megatron_gpt.layers.19.self_attention.dense.bias": "pytorch_model-00001-of-00002.bin",
162
+ "megatron_gpt.layers.19.self_attention.dense.weight": "pytorch_model-00001-of-00002.bin",
163
+ "megatron_gpt.layers.19.self_attention.query_key_value.bias": "pytorch_model-00001-of-00002.bin",
164
+ "megatron_gpt.layers.19.self_attention.query_key_value.weight": "pytorch_model-00001-of-00002.bin",
165
+ "megatron_gpt.layers.19.self_attention.rotary_emb.inv_freq": "pytorch_model-00001-of-00002.bin",
166
+ "megatron_gpt.layers.2.input_layernorm.bias": "pytorch_model-00001-of-00002.bin",
167
+ "megatron_gpt.layers.2.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
168
+ "megatron_gpt.layers.2.mlp.dense_4h_to_h.bias": "pytorch_model-00001-of-00002.bin",
169
+ "megatron_gpt.layers.2.mlp.dense_4h_to_h.weight": "pytorch_model-00001-of-00002.bin",
170
+ "megatron_gpt.layers.2.mlp.dense_h_to_4h.bias": "pytorch_model-00001-of-00002.bin",
171
+ "megatron_gpt.layers.2.mlp.dense_h_to_4h.weight": "pytorch_model-00001-of-00002.bin",
172
+ "megatron_gpt.layers.2.post_attention_layernorm.bias": "pytorch_model-00001-of-00002.bin",
173
+ "megatron_gpt.layers.2.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
174
+ "megatron_gpt.layers.2.self_attention.dense.bias": "pytorch_model-00001-of-00002.bin",
175
+ "megatron_gpt.layers.2.self_attention.dense.weight": "pytorch_model-00001-of-00002.bin",
176
+ "megatron_gpt.layers.2.self_attention.query_key_value.bias": "pytorch_model-00001-of-00002.bin",
177
+ "megatron_gpt.layers.2.self_attention.query_key_value.weight": "pytorch_model-00001-of-00002.bin",
178
+ "megatron_gpt.layers.2.self_attention.rotary_emb.inv_freq": "pytorch_model-00001-of-00002.bin",
179
+ "megatron_gpt.layers.20.input_layernorm.bias": "pytorch_model-00001-of-00002.bin",
180
+ "megatron_gpt.layers.20.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
181
+ "megatron_gpt.layers.20.mlp.dense_4h_to_h.bias": "pytorch_model-00001-of-00002.bin",
182
+ "megatron_gpt.layers.20.mlp.dense_4h_to_h.weight": "pytorch_model-00001-of-00002.bin",
183
+ "megatron_gpt.layers.20.mlp.dense_h_to_4h.bias": "pytorch_model-00001-of-00002.bin",
184
+ "megatron_gpt.layers.20.mlp.dense_h_to_4h.weight": "pytorch_model-00001-of-00002.bin",
185
+ "megatron_gpt.layers.20.post_attention_layernorm.bias": "pytorch_model-00001-of-00002.bin",
186
+ "megatron_gpt.layers.20.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
187
+ "megatron_gpt.layers.20.self_attention.dense.bias": "pytorch_model-00001-of-00002.bin",
188
+ "megatron_gpt.layers.20.self_attention.dense.weight": "pytorch_model-00001-of-00002.bin",
189
+ "megatron_gpt.layers.20.self_attention.query_key_value.bias": "pytorch_model-00001-of-00002.bin",
190
+ "megatron_gpt.layers.20.self_attention.query_key_value.weight": "pytorch_model-00001-of-00002.bin",
191
+ "megatron_gpt.layers.20.self_attention.rotary_emb.inv_freq": "pytorch_model-00001-of-00002.bin",
192
+ "megatron_gpt.layers.21.input_layernorm.bias": "pytorch_model-00001-of-00002.bin",
193
+ "megatron_gpt.layers.21.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
194
+ "megatron_gpt.layers.21.mlp.dense_4h_to_h.bias": "pytorch_model-00001-of-00002.bin",
195
+ "megatron_gpt.layers.21.mlp.dense_4h_to_h.weight": "pytorch_model-00001-of-00002.bin",
196
+ "megatron_gpt.layers.21.mlp.dense_h_to_4h.bias": "pytorch_model-00001-of-00002.bin",
197
+ "megatron_gpt.layers.21.mlp.dense_h_to_4h.weight": "pytorch_model-00001-of-00002.bin",
198
+ "megatron_gpt.layers.21.post_attention_layernorm.bias": "pytorch_model-00001-of-00002.bin",
199
+ "megatron_gpt.layers.21.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
200
+ "megatron_gpt.layers.21.self_attention.dense.bias": "pytorch_model-00001-of-00002.bin",
201
+ "megatron_gpt.layers.21.self_attention.dense.weight": "pytorch_model-00001-of-00002.bin",
202
+ "megatron_gpt.layers.21.self_attention.query_key_value.bias": "pytorch_model-00001-of-00002.bin",
203
+ "megatron_gpt.layers.21.self_attention.query_key_value.weight": "pytorch_model-00001-of-00002.bin",
204
+ "megatron_gpt.layers.21.self_attention.rotary_emb.inv_freq": "pytorch_model-00001-of-00002.bin",
205
+ "megatron_gpt.layers.22.input_layernorm.bias": "pytorch_model-00001-of-00002.bin",
206
+ "megatron_gpt.layers.22.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
207
+ "megatron_gpt.layers.22.mlp.dense_4h_to_h.bias": "pytorch_model-00001-of-00002.bin",
208
+ "megatron_gpt.layers.22.mlp.dense_4h_to_h.weight": "pytorch_model-00001-of-00002.bin",
209
+ "megatron_gpt.layers.22.mlp.dense_h_to_4h.bias": "pytorch_model-00001-of-00002.bin",
210
+ "megatron_gpt.layers.22.mlp.dense_h_to_4h.weight": "pytorch_model-00001-of-00002.bin",
211
+ "megatron_gpt.layers.22.post_attention_layernorm.bias": "pytorch_model-00001-of-00002.bin",
212
+ "megatron_gpt.layers.22.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
213
+ "megatron_gpt.layers.22.self_attention.dense.bias": "pytorch_model-00001-of-00002.bin",
214
+ "megatron_gpt.layers.22.self_attention.dense.weight": "pytorch_model-00001-of-00002.bin",
215
+ "megatron_gpt.layers.22.self_attention.query_key_value.bias": "pytorch_model-00001-of-00002.bin",
216
+ "megatron_gpt.layers.22.self_attention.query_key_value.weight": "pytorch_model-00001-of-00002.bin",
217
+ "megatron_gpt.layers.22.self_attention.rotary_emb.inv_freq": "pytorch_model-00001-of-00002.bin",
218
+ "megatron_gpt.layers.23.input_layernorm.bias": "pytorch_model-00001-of-00002.bin",
219
+ "megatron_gpt.layers.23.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
220
+ "megatron_gpt.layers.23.mlp.dense_4h_to_h.bias": "pytorch_model-00001-of-00002.bin",
221
+ "megatron_gpt.layers.23.mlp.dense_4h_to_h.weight": "pytorch_model-00001-of-00002.bin",
222
+ "megatron_gpt.layers.23.mlp.dense_h_to_4h.bias": "pytorch_model-00001-of-00002.bin",
223
+ "megatron_gpt.layers.23.mlp.dense_h_to_4h.weight": "pytorch_model-00001-of-00002.bin",
224
+ "megatron_gpt.layers.23.post_attention_layernorm.bias": "pytorch_model-00001-of-00002.bin",
225
+ "megatron_gpt.layers.23.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
226
+ "megatron_gpt.layers.23.self_attention.dense.bias": "pytorch_model-00001-of-00002.bin",
227
+ "megatron_gpt.layers.23.self_attention.dense.weight": "pytorch_model-00001-of-00002.bin",
228
+ "megatron_gpt.layers.23.self_attention.query_key_value.bias": "pytorch_model-00001-of-00002.bin",
229
+ "megatron_gpt.layers.23.self_attention.query_key_value.weight": "pytorch_model-00001-of-00002.bin",
230
+ "megatron_gpt.layers.23.self_attention.rotary_emb.inv_freq": "pytorch_model-00001-of-00002.bin",
231
+ "megatron_gpt.layers.24.input_layernorm.bias": "pytorch_model-00001-of-00002.bin",
232
+ "megatron_gpt.layers.24.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
233
+ "megatron_gpt.layers.24.mlp.dense_4h_to_h.bias": "pytorch_model-00001-of-00002.bin",
234
+ "megatron_gpt.layers.24.mlp.dense_4h_to_h.weight": "pytorch_model-00001-of-00002.bin",
235
+ "megatron_gpt.layers.24.mlp.dense_h_to_4h.bias": "pytorch_model-00001-of-00002.bin",
236
+ "megatron_gpt.layers.24.mlp.dense_h_to_4h.weight": "pytorch_model-00001-of-00002.bin",
237
+ "megatron_gpt.layers.24.post_attention_layernorm.bias": "pytorch_model-00001-of-00002.bin",
238
+ "megatron_gpt.layers.24.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
239
+ "megatron_gpt.layers.24.self_attention.dense.bias": "pytorch_model-00001-of-00002.bin",
240
+ "megatron_gpt.layers.24.self_attention.dense.weight": "pytorch_model-00001-of-00002.bin",
241
+ "megatron_gpt.layers.24.self_attention.query_key_value.bias": "pytorch_model-00001-of-00002.bin",
242
+ "megatron_gpt.layers.24.self_attention.query_key_value.weight": "pytorch_model-00001-of-00002.bin",
243
+ "megatron_gpt.layers.24.self_attention.rotary_emb.inv_freq": "pytorch_model-00001-of-00002.bin",
244
+ "megatron_gpt.layers.25.input_layernorm.bias": "pytorch_model-00001-of-00002.bin",
245
+ "megatron_gpt.layers.25.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
246
+ "megatron_gpt.layers.25.mlp.dense_4h_to_h.bias": "pytorch_model-00001-of-00002.bin",
247
+ "megatron_gpt.layers.25.mlp.dense_4h_to_h.weight": "pytorch_model-00001-of-00002.bin",
248
+ "megatron_gpt.layers.25.mlp.dense_h_to_4h.bias": "pytorch_model-00001-of-00002.bin",
249
+ "megatron_gpt.layers.25.mlp.dense_h_to_4h.weight": "pytorch_model-00001-of-00002.bin",
250
+ "megatron_gpt.layers.25.post_attention_layernorm.bias": "pytorch_model-00001-of-00002.bin",
251
+ "megatron_gpt.layers.25.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
252
+ "megatron_gpt.layers.25.self_attention.dense.bias": "pytorch_model-00001-of-00002.bin",
253
+ "megatron_gpt.layers.25.self_attention.dense.weight": "pytorch_model-00001-of-00002.bin",
254
+ "megatron_gpt.layers.25.self_attention.query_key_value.bias": "pytorch_model-00001-of-00002.bin",
255
+ "megatron_gpt.layers.25.self_attention.query_key_value.weight": "pytorch_model-00001-of-00002.bin",
256
+ "megatron_gpt.layers.25.self_attention.rotary_emb.inv_freq": "pytorch_model-00001-of-00002.bin",
257
+ "megatron_gpt.layers.26.input_layernorm.bias": "pytorch_model-00001-of-00002.bin",
258
+ "megatron_gpt.layers.26.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
259
+ "megatron_gpt.layers.26.mlp.dense_4h_to_h.bias": "pytorch_model-00001-of-00002.bin",
260
+ "megatron_gpt.layers.26.mlp.dense_4h_to_h.weight": "pytorch_model-00001-of-00002.bin",
261
+ "megatron_gpt.layers.26.mlp.dense_h_to_4h.bias": "pytorch_model-00001-of-00002.bin",
262
+ "megatron_gpt.layers.26.mlp.dense_h_to_4h.weight": "pytorch_model-00001-of-00002.bin",
263
+ "megatron_gpt.layers.26.post_attention_layernorm.bias": "pytorch_model-00001-of-00002.bin",
264
+ "megatron_gpt.layers.26.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
265
+ "megatron_gpt.layers.26.self_attention.dense.bias": "pytorch_model-00001-of-00002.bin",
266
+ "megatron_gpt.layers.26.self_attention.dense.weight": "pytorch_model-00001-of-00002.bin",
267
+ "megatron_gpt.layers.26.self_attention.query_key_value.bias": "pytorch_model-00001-of-00002.bin",
268
+ "megatron_gpt.layers.26.self_attention.query_key_value.weight": "pytorch_model-00001-of-00002.bin",
269
+ "megatron_gpt.layers.26.self_attention.rotary_emb.inv_freq": "pytorch_model-00001-of-00002.bin",
270
+ "megatron_gpt.layers.27.input_layernorm.bias": "pytorch_model-00001-of-00002.bin",
271
+ "megatron_gpt.layers.27.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
272
+ "megatron_gpt.layers.27.mlp.dense_4h_to_h.bias": "pytorch_model-00001-of-00002.bin",
273
+ "megatron_gpt.layers.27.mlp.dense_4h_to_h.weight": "pytorch_model-00001-of-00002.bin",
274
+ "megatron_gpt.layers.27.mlp.dense_h_to_4h.bias": "pytorch_model-00001-of-00002.bin",
275
+ "megatron_gpt.layers.27.mlp.dense_h_to_4h.weight": "pytorch_model-00001-of-00002.bin",
276
+ "megatron_gpt.layers.27.post_attention_layernorm.bias": "pytorch_model-00001-of-00002.bin",
277
+ "megatron_gpt.layers.27.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
278
+ "megatron_gpt.layers.27.self_attention.dense.bias": "pytorch_model-00001-of-00002.bin",
279
+ "megatron_gpt.layers.27.self_attention.dense.weight": "pytorch_model-00001-of-00002.bin",
280
+ "megatron_gpt.layers.27.self_attention.query_key_value.bias": "pytorch_model-00001-of-00002.bin",
281
+ "megatron_gpt.layers.27.self_attention.query_key_value.weight": "pytorch_model-00001-of-00002.bin",
282
+ "megatron_gpt.layers.27.self_attention.rotary_emb.inv_freq": "pytorch_model-00001-of-00002.bin",
283
+ "megatron_gpt.layers.28.input_layernorm.bias": "pytorch_model-00001-of-00002.bin",
284
+ "megatron_gpt.layers.28.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
285
+ "megatron_gpt.layers.28.mlp.dense_4h_to_h.bias": "pytorch_model-00001-of-00002.bin",
286
+ "megatron_gpt.layers.28.mlp.dense_4h_to_h.weight": "pytorch_model-00001-of-00002.bin",
287
+ "megatron_gpt.layers.28.mlp.dense_h_to_4h.bias": "pytorch_model-00001-of-00002.bin",
288
+ "megatron_gpt.layers.28.mlp.dense_h_to_4h.weight": "pytorch_model-00001-of-00002.bin",
289
+ "megatron_gpt.layers.28.post_attention_layernorm.bias": "pytorch_model-00001-of-00002.bin",
290
+ "megatron_gpt.layers.28.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
291
+ "megatron_gpt.layers.28.self_attention.dense.bias": "pytorch_model-00001-of-00002.bin",
292
+ "megatron_gpt.layers.28.self_attention.dense.weight": "pytorch_model-00001-of-00002.bin",
293
+ "megatron_gpt.layers.28.self_attention.query_key_value.bias": "pytorch_model-00001-of-00002.bin",
294
+ "megatron_gpt.layers.28.self_attention.query_key_value.weight": "pytorch_model-00001-of-00002.bin",
295
+ "megatron_gpt.layers.28.self_attention.rotary_emb.inv_freq": "pytorch_model-00001-of-00002.bin",
296
+ "megatron_gpt.layers.29.input_layernorm.bias": "pytorch_model-00001-of-00002.bin",
297
+ "megatron_gpt.layers.29.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
298
+ "megatron_gpt.layers.29.mlp.dense_4h_to_h.bias": "pytorch_model-00001-of-00002.bin",
299
+ "megatron_gpt.layers.29.mlp.dense_4h_to_h.weight": "pytorch_model-00001-of-00002.bin",
300
+ "megatron_gpt.layers.29.mlp.dense_h_to_4h.bias": "pytorch_model-00001-of-00002.bin",
301
+ "megatron_gpt.layers.29.mlp.dense_h_to_4h.weight": "pytorch_model-00001-of-00002.bin",
302
+ "megatron_gpt.layers.29.post_attention_layernorm.bias": "pytorch_model-00001-of-00002.bin",
303
+ "megatron_gpt.layers.29.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
304
+ "megatron_gpt.layers.29.self_attention.dense.bias": "pytorch_model-00001-of-00002.bin",
305
+ "megatron_gpt.layers.29.self_attention.dense.weight": "pytorch_model-00001-of-00002.bin",
306
+ "megatron_gpt.layers.29.self_attention.query_key_value.bias": "pytorch_model-00001-of-00002.bin",
307
+ "megatron_gpt.layers.29.self_attention.query_key_value.weight": "pytorch_model-00001-of-00002.bin",
308
+ "megatron_gpt.layers.29.self_attention.rotary_emb.inv_freq": "pytorch_model-00001-of-00002.bin",
309
+ "megatron_gpt.layers.3.input_layernorm.bias": "pytorch_model-00001-of-00002.bin",
310
+ "megatron_gpt.layers.3.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
311
+ "megatron_gpt.layers.3.mlp.dense_4h_to_h.bias": "pytorch_model-00001-of-00002.bin",
312
+ "megatron_gpt.layers.3.mlp.dense_4h_to_h.weight": "pytorch_model-00001-of-00002.bin",
313
+ "megatron_gpt.layers.3.mlp.dense_h_to_4h.bias": "pytorch_model-00001-of-00002.bin",
314
+ "megatron_gpt.layers.3.mlp.dense_h_to_4h.weight": "pytorch_model-00001-of-00002.bin",
315
+ "megatron_gpt.layers.3.post_attention_layernorm.bias": "pytorch_model-00001-of-00002.bin",
316
+ "megatron_gpt.layers.3.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
317
+ "megatron_gpt.layers.3.self_attention.dense.bias": "pytorch_model-00001-of-00002.bin",
318
+ "megatron_gpt.layers.3.self_attention.dense.weight": "pytorch_model-00001-of-00002.bin",
319
+ "megatron_gpt.layers.3.self_attention.query_key_value.bias": "pytorch_model-00001-of-00002.bin",
320
+ "megatron_gpt.layers.3.self_attention.query_key_value.weight": "pytorch_model-00001-of-00002.bin",
321
+ "megatron_gpt.layers.3.self_attention.rotary_emb.inv_freq": "pytorch_model-00001-of-00002.bin",
322
+ "megatron_gpt.layers.30.input_layernorm.bias": "pytorch_model-00001-of-00002.bin",
323
+ "megatron_gpt.layers.30.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
324
+ "megatron_gpt.layers.30.mlp.dense_4h_to_h.bias": "pytorch_model-00002-of-00002.bin",
325
+ "megatron_gpt.layers.30.mlp.dense_4h_to_h.weight": "pytorch_model-00002-of-00002.bin",
326
+ "megatron_gpt.layers.30.mlp.dense_h_to_4h.bias": "pytorch_model-00002-of-00002.bin",
327
+ "megatron_gpt.layers.30.mlp.dense_h_to_4h.weight": "pytorch_model-00002-of-00002.bin",
328
+ "megatron_gpt.layers.30.post_attention_layernorm.bias": "pytorch_model-00001-of-00002.bin",
329
+ "megatron_gpt.layers.30.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
330
+ "megatron_gpt.layers.30.self_attention.dense.bias": "pytorch_model-00001-of-00002.bin",
331
+ "megatron_gpt.layers.30.self_attention.dense.weight": "pytorch_model-00001-of-00002.bin",
332
+ "megatron_gpt.layers.30.self_attention.query_key_value.bias": "pytorch_model-00001-of-00002.bin",
333
+ "megatron_gpt.layers.30.self_attention.query_key_value.weight": "pytorch_model-00001-of-00002.bin",
334
+ "megatron_gpt.layers.30.self_attention.rotary_emb.inv_freq": "pytorch_model-00001-of-00002.bin",
335
+ "megatron_gpt.layers.31.input_layernorm.bias": "pytorch_model-00002-of-00002.bin",
336
+ "megatron_gpt.layers.31.input_layernorm.weight": "pytorch_model-00002-of-00002.bin",
337
+ "megatron_gpt.layers.31.mlp.dense_4h_to_h.bias": "pytorch_model-00002-of-00002.bin",
338
+ "megatron_gpt.layers.31.mlp.dense_4h_to_h.weight": "pytorch_model-00002-of-00002.bin",
339
+ "megatron_gpt.layers.31.mlp.dense_h_to_4h.bias": "pytorch_model-00002-of-00002.bin",
340
+ "megatron_gpt.layers.31.mlp.dense_h_to_4h.weight": "pytorch_model-00002-of-00002.bin",
341
+ "megatron_gpt.layers.31.post_attention_layernorm.bias": "pytorch_model-00002-of-00002.bin",
342
+ "megatron_gpt.layers.31.post_attention_layernorm.weight": "pytorch_model-00002-of-00002.bin",
343
+ "megatron_gpt.layers.31.self_attention.dense.bias": "pytorch_model-00002-of-00002.bin",
344
+ "megatron_gpt.layers.31.self_attention.dense.weight": "pytorch_model-00002-of-00002.bin",
345
+ "megatron_gpt.layers.31.self_attention.query_key_value.bias": "pytorch_model-00002-of-00002.bin",
346
+ "megatron_gpt.layers.31.self_attention.query_key_value.weight": "pytorch_model-00002-of-00002.bin",
347
+ "megatron_gpt.layers.31.self_attention.rotary_emb.inv_freq": "pytorch_model-00002-of-00002.bin",
348
+ "megatron_gpt.layers.4.input_layernorm.bias": "pytorch_model-00001-of-00002.bin",
349
+ "megatron_gpt.layers.4.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
350
+ "megatron_gpt.layers.4.mlp.dense_4h_to_h.bias": "pytorch_model-00001-of-00002.bin",
351
+ "megatron_gpt.layers.4.mlp.dense_4h_to_h.weight": "pytorch_model-00001-of-00002.bin",
352
+ "megatron_gpt.layers.4.mlp.dense_h_to_4h.bias": "pytorch_model-00001-of-00002.bin",
353
+ "megatron_gpt.layers.4.mlp.dense_h_to_4h.weight": "pytorch_model-00001-of-00002.bin",
354
+ "megatron_gpt.layers.4.post_attention_layernorm.bias": "pytorch_model-00001-of-00002.bin",
355
+ "megatron_gpt.layers.4.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
356
+ "megatron_gpt.layers.4.self_attention.dense.bias": "pytorch_model-00001-of-00002.bin",
357
+ "megatron_gpt.layers.4.self_attention.dense.weight": "pytorch_model-00001-of-00002.bin",
358
+ "megatron_gpt.layers.4.self_attention.query_key_value.bias": "pytorch_model-00001-of-00002.bin",
359
+ "megatron_gpt.layers.4.self_attention.query_key_value.weight": "pytorch_model-00001-of-00002.bin",
360
+ "megatron_gpt.layers.4.self_attention.rotary_emb.inv_freq": "pytorch_model-00001-of-00002.bin",
361
+ "megatron_gpt.layers.5.input_layernorm.bias": "pytorch_model-00001-of-00002.bin",
362
+ "megatron_gpt.layers.5.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
363
+ "megatron_gpt.layers.5.mlp.dense_4h_to_h.bias": "pytorch_model-00001-of-00002.bin",
364
+ "megatron_gpt.layers.5.mlp.dense_4h_to_h.weight": "pytorch_model-00001-of-00002.bin",
365
+ "megatron_gpt.layers.5.mlp.dense_h_to_4h.bias": "pytorch_model-00001-of-00002.bin",
366
+ "megatron_gpt.layers.5.mlp.dense_h_to_4h.weight": "pytorch_model-00001-of-00002.bin",
367
+ "megatron_gpt.layers.5.post_attention_layernorm.bias": "pytorch_model-00001-of-00002.bin",
368
+ "megatron_gpt.layers.5.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
369
+ "megatron_gpt.layers.5.self_attention.dense.bias": "pytorch_model-00001-of-00002.bin",
370
+ "megatron_gpt.layers.5.self_attention.dense.weight": "pytorch_model-00001-of-00002.bin",
371
+ "megatron_gpt.layers.5.self_attention.query_key_value.bias": "pytorch_model-00001-of-00002.bin",
372
+ "megatron_gpt.layers.5.self_attention.query_key_value.weight": "pytorch_model-00001-of-00002.bin",
373
+ "megatron_gpt.layers.5.self_attention.rotary_emb.inv_freq": "pytorch_model-00001-of-00002.bin",
374
+ "megatron_gpt.layers.6.input_layernorm.bias": "pytorch_model-00001-of-00002.bin",
375
+ "megatron_gpt.layers.6.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
376
+ "megatron_gpt.layers.6.mlp.dense_4h_to_h.bias": "pytorch_model-00001-of-00002.bin",
377
+ "megatron_gpt.layers.6.mlp.dense_4h_to_h.weight": "pytorch_model-00001-of-00002.bin",
378
+ "megatron_gpt.layers.6.mlp.dense_h_to_4h.bias": "pytorch_model-00001-of-00002.bin",
379
+ "megatron_gpt.layers.6.mlp.dense_h_to_4h.weight": "pytorch_model-00001-of-00002.bin",
380
+ "megatron_gpt.layers.6.post_attention_layernorm.bias": "pytorch_model-00001-of-00002.bin",
381
+ "megatron_gpt.layers.6.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
382
+ "megatron_gpt.layers.6.self_attention.dense.bias": "pytorch_model-00001-of-00002.bin",
383
+ "megatron_gpt.layers.6.self_attention.dense.weight": "pytorch_model-00001-of-00002.bin",
384
+ "megatron_gpt.layers.6.self_attention.query_key_value.bias": "pytorch_model-00001-of-00002.bin",
385
+ "megatron_gpt.layers.6.self_attention.query_key_value.weight": "pytorch_model-00001-of-00002.bin",
386
+ "megatron_gpt.layers.6.self_attention.rotary_emb.inv_freq": "pytorch_model-00001-of-00002.bin",
387
+ "megatron_gpt.layers.7.input_layernorm.bias": "pytorch_model-00001-of-00002.bin",
388
+ "megatron_gpt.layers.7.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
389
+ "megatron_gpt.layers.7.mlp.dense_4h_to_h.bias": "pytorch_model-00001-of-00002.bin",
390
+ "megatron_gpt.layers.7.mlp.dense_4h_to_h.weight": "pytorch_model-00001-of-00002.bin",
391
+ "megatron_gpt.layers.7.mlp.dense_h_to_4h.bias": "pytorch_model-00001-of-00002.bin",
392
+ "megatron_gpt.layers.7.mlp.dense_h_to_4h.weight": "pytorch_model-00001-of-00002.bin",
393
+ "megatron_gpt.layers.7.post_attention_layernorm.bias": "pytorch_model-00001-of-00002.bin",
394
+ "megatron_gpt.layers.7.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
395
+ "megatron_gpt.layers.7.self_attention.dense.bias": "pytorch_model-00001-of-00002.bin",
396
+ "megatron_gpt.layers.7.self_attention.dense.weight": "pytorch_model-00001-of-00002.bin",
397
+ "megatron_gpt.layers.7.self_attention.query_key_value.bias": "pytorch_model-00001-of-00002.bin",
398
+ "megatron_gpt.layers.7.self_attention.query_key_value.weight": "pytorch_model-00001-of-00002.bin",
399
+ "megatron_gpt.layers.7.self_attention.rotary_emb.inv_freq": "pytorch_model-00001-of-00002.bin",
400
+ "megatron_gpt.layers.8.input_layernorm.bias": "pytorch_model-00001-of-00002.bin",
401
+ "megatron_gpt.layers.8.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
402
+ "megatron_gpt.layers.8.mlp.dense_4h_to_h.bias": "pytorch_model-00001-of-00002.bin",
403
+ "megatron_gpt.layers.8.mlp.dense_4h_to_h.weight": "pytorch_model-00001-of-00002.bin",
404
+ "megatron_gpt.layers.8.mlp.dense_h_to_4h.bias": "pytorch_model-00001-of-00002.bin",
405
+ "megatron_gpt.layers.8.mlp.dense_h_to_4h.weight": "pytorch_model-00001-of-00002.bin",
406
+ "megatron_gpt.layers.8.post_attention_layernorm.bias": "pytorch_model-00001-of-00002.bin",
407
+ "megatron_gpt.layers.8.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
408
+ "megatron_gpt.layers.8.self_attention.dense.bias": "pytorch_model-00001-of-00002.bin",
409
+ "megatron_gpt.layers.8.self_attention.dense.weight": "pytorch_model-00001-of-00002.bin",
410
+ "megatron_gpt.layers.8.self_attention.query_key_value.bias": "pytorch_model-00001-of-00002.bin",
411
+ "megatron_gpt.layers.8.self_attention.query_key_value.weight": "pytorch_model-00001-of-00002.bin",
412
+ "megatron_gpt.layers.8.self_attention.rotary_emb.inv_freq": "pytorch_model-00001-of-00002.bin",
413
+ "megatron_gpt.layers.9.input_layernorm.bias": "pytorch_model-00001-of-00002.bin",
414
+ "megatron_gpt.layers.9.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
415
+ "megatron_gpt.layers.9.mlp.dense_4h_to_h.bias": "pytorch_model-00001-of-00002.bin",
416
+ "megatron_gpt.layers.9.mlp.dense_4h_to_h.weight": "pytorch_model-00001-of-00002.bin",
417
+ "megatron_gpt.layers.9.mlp.dense_h_to_4h.bias": "pytorch_model-00001-of-00002.bin",
418
+ "megatron_gpt.layers.9.mlp.dense_h_to_4h.weight": "pytorch_model-00001-of-00002.bin",
419
+ "megatron_gpt.layers.9.post_attention_layernorm.bias": "pytorch_model-00001-of-00002.bin",
420
+ "megatron_gpt.layers.9.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
421
+ "megatron_gpt.layers.9.self_attention.dense.bias": "pytorch_model-00001-of-00002.bin",
422
+ "megatron_gpt.layers.9.self_attention.dense.weight": "pytorch_model-00001-of-00002.bin",
423
+ "megatron_gpt.layers.9.self_attention.query_key_value.bias": "pytorch_model-00001-of-00002.bin",
424
+ "megatron_gpt.layers.9.self_attention.query_key_value.weight": "pytorch_model-00001-of-00002.bin",
425
+ "megatron_gpt.layers.9.self_attention.rotary_emb.inv_freq": "pytorch_model-00001-of-00002.bin"
426
+ }
427
+ }
special_tokens_map.json ADDED
@@ -0,0 +1,30 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "bos_token": {
3
+ "content": "<|endoftext|>",
4
+ "lstrip": false,
5
+ "normalized": true,
6
+ "rstrip": false,
7
+ "single_word": false
8
+ },
9
+ "eos_token": {
10
+ "content": "<|endoftext|>",
11
+ "lstrip": false,
12
+ "normalized": true,
13
+ "rstrip": false,
14
+ "single_word": false
15
+ },
16
+ "pad_token": {
17
+ "content": "<pad>",
18
+ "lstrip": false,
19
+ "normalized": true,
20
+ "rstrip": false,
21
+ "single_word": false
22
+ },
23
+ "unk_token": {
24
+ "content": "<unk>",
25
+ "lstrip": false,
26
+ "normalized": true,
27
+ "rstrip": false,
28
+ "single_word": false
29
+ }
30
+ }
tokenizer_config.json ADDED
@@ -0,0 +1,40 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "add_bos_token": false,
3
+ "add_prefix_space": false,
4
+ "bos_token": {
5
+ "__type": "AddedToken",
6
+ "content": "<|endoftext|>",
7
+ "lstrip": false,
8
+ "normalized": true,
9
+ "rstrip": false,
10
+ "single_word": false
11
+ },
12
+ "clean_up_tokenization_spaces": true,
13
+ "eos_token": {
14
+ "__type": "AddedToken",
15
+ "content": "<|endoftext|>",
16
+ "lstrip": false,
17
+ "normalized": true,
18
+ "rstrip": false,
19
+ "single_word": false
20
+ },
21
+ "errors": "replace",
22
+ "model_max_length": 1000000000000000019884624838656,
23
+ "pad_token": {
24
+ "__type": "AddedToken",
25
+ "content": "<pad>",
26
+ "lstrip": false,
27
+ "normalized": true,
28
+ "rstrip": false,
29
+ "single_word": false
30
+ },
31
+ "tokenizer_class": "GPT2Tokenizer",
32
+ "unk_token": {
33
+ "__type": "AddedToken",
34
+ "content": "<unk>",
35
+ "lstrip": false,
36
+ "normalized": true,
37
+ "rstrip": false,
38
+ "single_word": false
39
+ }
40
+ }
vocab.json ADDED
The diff for this file is too large to render. See raw diff