Nicer error message for undefined symbol

llmfoundry/__init__.py

@@ -9,6 +9,23 @@
 
 import logging
 
+try:
+    from flash_attn import flash_attn_func
+    del flash_attn_func
+except ImportError as e:
+    if 'undefined symbol' in str(e):
+        raise ImportError(
+            'The flash_attn package is not installed correctly. Usually this means that your runtime version'
+            +
+            ' of PyTorch is different from the version that flash_attn was installed with, which can occur when your'
+            +
+            ' workflow has resulted in PyTorch being reinstalled. This probably happened because you are using an old Docker image'
+            +
+            ' with the latest version of LLM Foundry. Check that the PyTorch version in your Docker image matches the PyTorch version'
+            +
+            ' in LLM Foundry setup.py and update accordingly. The latest Docker image can be found in the README.',
+        ) from e
+
 from llmfoundry.utils.logging_utils import SpecificWarningFilter
 
 # Filter out Hugging Face warning for not using a pinned revision of the model