naming fix

attn.py

@@ -27,13 +27,13 @@ class Attention(nn.Module):
 
         self.device = torch.device("cuda")
         self.bsz = config.bsz
-        self.c_attn = nn.Linear(
+        self.attn = nn.Linear(
             config.n_embd, 3 * config.n_embd, bias=config.bias, dtype=torch_dtype
         )
-        self.c_proj = nn.Linear(
+        self.o_proj = nn.Linear(
             config.n_embd, config.n_embd, bias=config.bias, dtype=torch_dtype
         )
-        self.c_proj.SCALE_INIT = 1
+        self.o_proj.SCALE_INIT = 1
         self.dropout = config.dropout
         self.resid_dropout = nn.Dropout(self.dropout)
         self.alibi_slopes = self._get_alibi_slopes(self.n_heads)
@@ -65,7 +65,7 @@ class Attention(nn.Module):
     def forward(self, x):
         bsz, seq_len, d_in = x.size()
 
-        qkv = self.c_attn(x)
+        qkv = self.attn(x)
         q, k, v = torch.chunk(qkv, 3, dim=2)
 
         q = q.view(bsz, seq_len, self.n_heads, d_in // self.n_heads)
@@ -82,7 +82,7 @@ class Attention(nn.Module):
             softcap=self.softcap,  # https://arxiv.org/pdf/2408.00118
         )
         y = y.contiguous().view(bsz, seq_len, d_in)
-        y = self.resid_dropout(self.c_proj(y))
+        y = self.resid_dropout(self.o_proj(y))
         return y
 
 class AttentionSDPA(nn.Module):
@@ -98,15 +98,15 @@ class AttentionSDPA(nn.Module):
 
         self.device = torch.device("cuda") # Technically don't need CUDA for SDPA
         self.bsz = config.bsz
-        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias, dtype=torch_dtype)
-        self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias, dtype=torch_dtype)
+        self.attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias, dtype=torch_dtype)
+        self.o_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias, dtype=torch_dtype)
         self.dropout = config.dropout
         self.resid_dropout = nn.Dropout(self.dropout)
 
     def forward(self, x):
         bsz, seq_len, d_in = x.size()
 
-        qkv = self.c_attn(x)
+        qkv = self.attn(x)
         q, k, v = torch.chunk(qkv, 3, dim=2)
 
         q = q.view(bsz, seq_len, self.n_heads, d_in // self.n_heads).transpose(1, 2)
@@ -121,5 +121,5 @@ class AttentionSDPA(nn.Module):
 
         y = y.transpose(1, 2).contiguous().view(bsz, seq_len, d_in)
 
-        y = self.resid_dropout(self.c_proj(y))
+        y = self.resid_dropout(self.o_proj(y))
         return y

configuration_minitransformer.py

@@ -8,7 +8,7 @@ class MiniTransformerConfig(PretrainedConfig):
         self,
         bsz: int = 1,
         n_embd: int = 768,
-        n_heads: int = 24,
+        n_heads: int = 12,
         n_layers: int = 27,
         seq_len: int = 8192,
         window_size: int = 8192,

modeling_minitransformer.py

@@ -121,12 +121,12 @@ class MiniTransformer(PreTrainedModel):
         elif isinstance(module, nn.Embedding):
             torch.nn.init.normal_(module.weight, mean=0.0, std=self.std)
         elif isinstance(module, Attention):
-            torch.nn.init.xavier_normal_(module.c_attn.weight)
-            torch.nn.init.xavier_normal_(module.c_proj.weight)
-            if module.c_attn.bias is not None:
-                torch.nn.init.zeros_(module.c_attn.bias)
-            if module.c_proj.bias is not None:
-                torch.nn.init.zeros_(module.c_proj.bias)
+            torch.nn.init.xavier_normal_(module.attn.weight)
+            torch.nn.init.xavier_normal_(module.o_proj.weight)
+            if module.attn.bias is not None:
+                torch.nn.init.zeros_(module.attn.bias)
+            if module.o_proj.bias is not None:
+                torch.nn.init.zeros_(module.o_proj.bias)
     @staticmethod
     def top_k_top_p_filtering(
         logits: torch.Tensor,