抽取self.downsample_v与self.downsample_q的共同部分，并使用可分离卷积降低参数量

model/model.py

@@ -174,12 +174,12 @@ class CrossAttention(nn.Module):
         super().__init__()
         self.config = config
         self.num_heads = 8  
-        self.head_dim = 768 // self.num_heads  
-        self.to_q = nn.Linear(768, 768, bias=False)
-        self.to_k = nn.Linear(768, 768, bias=False)
-        self.to_v = nn.Linear(768, 768, bias=False)
+        self.head_dim = self.config.dim // self.num_heads  
+        self.to_q = nn.Linear(self.config.dim, self.config.dim, bias=False)
+        self.to_k = nn.Linear(self.config.dim, self.config.dim, bias=False)
+        self.to_v = nn.Linear(self.config.dim, self.config.dim, bias=False)
         
-        self.to_out = nn.Linear(768, 768, bias=False)
+        self.to_out = nn.Linear(self.config.dim, self.config.dim, bias=False)
     
     def forward(self, x, db, context_mask=None, pos_emb=None):
         batch_size = x.size(0)
@@ -205,7 +205,7 @@ class CrossAttention(nn.Module):
 
         context = torch.matmul(attn_weights, v)
         
-        context = context.transpose(1, 2).contiguous().view(batch_size, -1, 768)
+        context = context.transpose(1, 2).contiguous().view(batch_size, -1, self.config.dim)
 
         context = self.to_out(context)
         
@@ -523,19 +523,36 @@ class MiniMindLM(PreTrainedModel):
         self.norm = RMSNorm(params.dim, eps=params.norm_eps)
         self.output = nn.Linear(params.dim, params.vocab_size, bias=False)
         self.tok_embeddings.weight = self.output.weight
-        self.downsample_v = nn.Sequential(
-            nn.Conv1d(511*8,128*8,kernel_size=1,padding='same'),
-            nn.Conv1d(128*8,128,kernel_size=1,padding='same'),
-            nn.Conv1d(128,8,kernel_size=1,padding='same')
+        
+        # Calculate input dimension
+        input_dim = (self.params.max_seq_len-1)*self.params.n_layers
+        # Use a bottleneck architecture to reduce parameters
+        bottleneck_dim = 256  # Significantly smaller bottleneck dimension
+        
+        # Factorized shared downsampling using two smaller convolutions
+        self.shared_downsample = nn.Sequential(
+            # First reduce input dimension to bottleneck
+            nn.Conv1d(input_dim, bottleneck_dim, kernel_size=1, padding='same'),
+            nn.ReLU(),  # Non-linearity to improve representation capacity
+            # Then expand to target dimension
+            nn.Conv1d(bottleneck_dim, 128*8, kernel_size=1, padding='same')
         )
-        self.downsample_q = nn.Sequential(
-            nn.Conv1d(511*8,128*8,kernel_size=1,padding='same'),
-            nn.Conv1d(128*8,512,kernel_size=1,padding='same')
+        
+        # Specific layers for v path
+        self.downsample_v_specific = nn.Sequential(
+            nn.Conv1d(128*8, 128, kernel_size=1, padding='same'),
+            nn.Conv1d(128, 8, kernel_size=1, padding='same')
+        )
+        
+        # Specific layers for q path
+        self.downsample_q_specific = nn.Sequential(
+            nn.Conv1d(128*8, 512, kernel_size=1, padding='same')
         )
         self.register_buffer("pos_cis",
                              precompute_pos_cis(dim=params.dim // params.n_heads, theta=params.rope_theta),
                              persistent=False)
         self.OUT = CausalLMOutputWithPast()
+        self.params = params
 
     def forward(self,
                 input_ids: Optional[torch.Tensor] = None,
@@ -565,12 +582,14 @@ class MiniMindLM(PreTrainedModel):
         # 使用detach()分离计算图，避免多次反向传播
         h_tensor = torch.cat(h_list,dim=0).permute(1,0,2,3)
         h_tensor_detached = h_tensor.detach()
-        h_tensor_detached = h_tensor_detached.reshape(h_tensor_detached.shape[0],-1,768)
+        h_tensor_detached = h_tensor_detached.reshape(h_tensor_detached.shape[0],-1,self.params.dim)
         
         # 数据库更新逻辑与主计算图分离
         with torch.no_grad():
-            z_v = self.downsample_v(h_tensor_detached)
-            z_q = self.downsample_q(h_tensor_detached)
+            # Compute shared downsampling layer once
+            shared_features = self.shared_downsample(h_tensor_detached)
+            z_v = self.downsample_v_specific(shared_features)
+            z_q = self.downsample_q_specific(shared_features)
             z_k = self.extract_db.q_to_k(z_q)
             self.extract_db.updata_value(z_k,z_v)
         

train_pretrain.py

@@ -179,7 +179,7 @@ if __name__ == "__main__":
     parser.add_argument("--out_dir", type=str, default="out")
     # 若要以最快速度实现zero则epochs设置为1轮；否则应当利用有限的数据训练2~6个epochs。
     parser.add_argument("--epochs", type=int, default=3)
-    parser.add_argument("--batch_size", type=int, default=32)
+    parser.add_argument("--batch_size", type=int, default=1)
     parser.add_argument("--learning_rate", type=float, default=5e-4)
     parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu") #如果GPU可用，则使用GPU，否则使用CPU。
     parser.add_argument("--dtype", type=str, default="bfloat16")
@@ -193,8 +193,8 @@ if __name__ == "__main__":
     parser.add_argument("--log_interval", type=int, default=100) #日志打印间隔，用于控制日志打印的频率。
     parser.add_argument("--save_interval", type=int, default=100) #模型保存间隔，用于控制模型保存的频率。
     parser.add_argument('--local_rank', type=int, default=-1) #本地进程编号，用于分布式训练。
-    parser.add_argument('--dim', default=768, type=int) #模型维度，用于控制模型的大小。
-    parser.add_argument('--n_layers', default=8, type=int) #层数，用于控制模型层数。
+    parser.add_argument('--dim', default=592, type=int) #模型维度，用于控制模型的大小。
+    parser.add_argument('--n_layers', default=4, type=int) #层数，用于控制模型层数。
     parser.add_argument('--max_seq_len', default=512, type=int) #最大序列长度，用于控制输入序列的最大长度。
     parser.add_argument('--use_moe', default=False, type=bool) #是否使用MOE，用于控制是否使用MOE。
     parser.add_argument("--data_path", type=str, default="./dataset/pretrain_hq.jsonl") #数据路径，用于控制数据集的路径。