assert to make sure flash-attn>=2.4.3

volcengine · Feb 5, 2025 · cc68a04 · cc68a04
1 parent ddfc04e
commit cc68a04
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/verl/utils/torch_functional.py b/verl/utils/torch_functional.py
@@ -63,8 +63,9 @@ def logprobs_from_logits(logits, labels):
 
 
 def logprobs_from_logits_flash_attn(logits, labels):
-    output = -cross_entropy_loss(logits, labels)[0]
-    return output
+    output = cross_entropy_loss(logits, labels)
+    assert isinstance(output, tuple), "please make sure flash-attn>=2.4.3 where cross_entropy_loss returns Tuple[losses, z_losses]."
+    return -output[0]
 
 
 def logprobs_from_logits_naive(logits, labels):