fix(mha): passed dropout_p as 0 in mha eval

DeepLink-org · Jan 29, 2024 · af74d9f · af74d9f
1 parent 1229e2d
commit af74d9f
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/deeplink_ext/internlm_ops/mha/mha.py b/deeplink_ext/internlm_ops/mha/mha.py
@@ -48,7 +48,7 @@ def forward(self, qkv, causal=None, cu_seqlens=None, max_seqlen=None):
             # padded
             return DeepLinkMultiHeadAttentionQKVPackedFunc.apply(
                 qkv,
-                self.dropout_p,
+                self.dropout_p if self.training else 0.0,
                 self.softmax_scale,
                 causal if causal is not None else self.causal,
                 False,
@@ -59,7 +59,7 @@ def forward(self, qkv, causal=None, cu_seqlens=None, max_seqlen=None):
                 qkv,
                 cu_seqlens,
                 max_seqlen,
-                self.dropout_p,
+                self.dropout_p if self.training else 0.0,
                 self.softmax_scale,
                 causal if causal is not None else self.causal,
                 False,