Upload GPTRefactForCausalLM

config.json

@@ -20,7 +20,7 @@
   "n_layer": 32,
   "n_positions": 4096,
   "scale_attention_softmax_in_fp32": true,
-  "torch_dtype": "float16",
+  "torch_dtype": "bfloat16",
   "transformers_version": "4.31.0",
   "use_cache": true,
   "vocab_size": 49216

modeling_gpt_refact.py

@@ -101,7 +101,6 @@ def get_alibi_biases(
     # Multiply them pair-wise to get the AliBi bias matrix
     biases = distance[:, :, None] * m[None, None, :]
     biases = biases.permute(2, 0, 1)[None, :, :T, :T]
-    biases = biases.repeat(B, 1, 1, 1)
     return biases.contiguous()
 
 
@@ -132,8 +131,7 @@ class Attention(nn.Module):
         self.attention_bias_in_fp32 = config.attention_bias_in_fp32
 
         self.q = nn.Linear(self.embed_dim, self.embed_dim, bias=False)
-        self.k = nn.Linear(self.embed_dim, self.head_dim, bias=False)
-        self.v = nn.Linear(self.embed_dim, self.head_dim, bias=False)
+        self.kv = nn.Linear(self.embed_dim, self.head_dim * 2, bias=False)
         self.c_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=False)
 
     def _get_mask_value(self, device, dtype):
@@ -200,8 +198,8 @@ class Attention(nn.Module):
         Tuple[torch.Tensor, Optional[torch.Tensor], Tuple[torch.Tensor, ...]],
     ]:
         query = self.q(hidden_states)
-        key = self.k(hidden_states)
-        value = self.v(hidden_states)
+        kv = self.kv(hidden_states)
+        key, value = kv.split(self.head_dim, dim=-1)
 
         if layer_past is not None:
             past_key, past_value = layer_past
@@ -231,15 +229,14 @@ class MLP(nn.Module):
         embed_dim = config.hidden_size
         hidden_dim = intermediate_size
         hidden_dim = int(2 * hidden_dim / 3)
-        hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)
-        self.linear_1 = nn.Linear(embed_dim, hidden_dim, bias=False)
-        self.linear_3 = nn.Linear(embed_dim, hidden_dim, bias=False)
-        self.c_proj = nn.Linear(hidden_dim, embed_dim, bias=False)
+        self.hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)
+        self.gate_up_proj = nn.Linear(embed_dim, self.hidden_dim * 2, bias=False)
+        self.c_proj = nn.Linear(self.hidden_dim, embed_dim, bias=False)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        x1 = F.silu(self.linear_1(x))
-        x2 = self.linear_3(x)
-        x = self.c_proj(x1 * x2)
+        up_proj = self.gate_up_proj(x)
+        x1, x2 = torch.split(up_proj, self.hidden_dim, dim=-1)
+        x = self.c_proj(F.silu(x1) * x2)
         return x
 
 
@@ -264,7 +261,6 @@ class GPTRefactBlock(nn.Module):
         self.ln_1 = LayerNormNoBias(hidden_size, eps=config.layer_norm_epsilon)
         self.attn = Attention(config, layer_idx=layer_idx)
         self.ln_2 = LayerNormNoBias(hidden_size, eps=config.layer_norm_epsilon)
-
         self.mlp = MLP(self.inner_dim, config)
 
     def forward(

pytorch_model.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:1092c5efe56fe5b04360ba0d4ac231e8b03f9d1d0b8633b8ed678f73bdcb021a
-size 3171776281
+oid sha256:6bf4dc20907069119671fdaf9f7b79d0260cd36ab94626f4af4fdd5a157d0205
+size 3171755929