Update modeling_internlm2.py

modeling_internlm2.py

@@ -311,9 +311,9 @@ class InternLM2Attention(nn.Module):
         self.attncut = True
         self.headcut = True
         self.layercut = True
-        self.layercut_idx = 36
-        self.offset = 41
-        head_num=24
+        self.layercut_idx = 36 # number of layer kept 
+        self.offset = 41  # system prompt token length
+        head_num=24 # num of heads kept on average
         self.mask = torch.load("headcut_mask/internvl2.0_26B/mask_"+str(head_num)+".pth")
     def _init_rope(self):
         if self.config.rope_scaling is None: