update unfused

Signed-off-by: Ananth Subramaniam <[email protected]>
ananthsub · Dec 16, 2024 · 8d1c589 · 8d1c589
1 parent 6d5a201
commit 8d1c589
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/nemo/collections/llm/gpt/model/gemma.py b/nemo/collections/llm/gpt/model/gemma.py
@@ -75,6 +75,9 @@ class GemmaConfig7B(GemmaConfig):
     num_query_groups: int = 16
     ffn_hidden_size: int = 24576
 
+    # Disable cuDNN attention since TE 1.8 does not support head dim > 128
+    attention_backend: AttnBackend = AttnBackend.flash
+
 
 class CodeGemmaConfig2B(GemmaConfig2B):
     pass