Fix triton import guards (NVIDIA#8552)

* Fix triton import guards Signed-off-by: Michal Futrega <[email protected]> * Update attention.py Signed-off-by: Michal Futrega <[email protected]> --------- Signed-off-by: Michal Futrega <[email protected]> Signed-off-by: duongvdo <[email protected]>
va290 · Mar 10, 2024 · e78c1a3 · e78c1a3
1 parent 14ee191
commit e78c1a3
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/nemo/collections/nlp/modules/common/megatron/attention.py b/nemo/collections/nlp/modules/common/megatron/attention.py
@@ -73,7 +73,7 @@
     # pinned triton version for flash-attention triton https://github.com/HazyResearch/flash-attention/blob/main/flash_attn/flash_attn_triton.py#L3
     assert pkg_resources.get_distribution("triton").version == '2.0.0.dev20221202'
 
-except (ImportError, ModuleNotFoundError, AssertionError):
+except (ImportError, ModuleNotFoundError, AssertionError, pkg_resources.DistributionNotFound):
 
     flash_attn_func_triton = None