fix(mha.py): fix evaluation argu key err (#370)

InternLM · Nov 18, 2024 · aee457c · aee457c
1 parent 330c03a
commit aee457c
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/internlm/model/modules/mha.py b/internlm/model/modules/mha.py
@@ -211,7 +211,7 @@ def _training(self, x, **kwargs):
 
         # self attention
         kwargs = _convert_cu_seqlens_for_qksplited(kwargs)
-        if gpc.config.data.use_packed_dataset is False:
+        if gpc.config.data.use_packed_dataset is False or self.training is False:
             kwargs.pop("max_seqlen_q", None)
             kwargs.pop("max_seqlen_k", None)
         context = self.inner_attn(q, k, v, **kwargs)
@@ -529,7 +529,7 @@ def _training(self, x, **kwargs):
 
         kv = torch.concat([k.unsqueeze(2), v.unsqueeze(2)], dim=2)
 
-        if gpc.config.data.use_packed_dataset is False:
+        if gpc.config.data.use_packed_dataset is False or self.training is False:
             kwargs.pop("max_seqlen_q", None)
             kwargs.pop("max_seqlen_k", None)