Update modeling_internlm2.py

modeling_internlm2.py

@@ -310,9 +310,9 @@ class InternLM2Attention(nn.Module):
         self.attncut = True
         self.headcut = True
         self.layercut = True
-        self.layercut_idx = 24
-        self.offset = 41
-        head_num=24
+        self.layercut_idx = 24 # 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_8B/mask_"+str(head_num)+".pth")
         
     def _init_rope(self):