..

mosaicml · Oct 27, 2023 · d162855 · d162855
1 parent bfc092e
commit d162855
Show file tree

Hide file tree

Showing 5 changed files with 7 additions and 7 deletions.
diff --git a/llmfoundry/__init__.py b/llmfoundry/__init__.py
@@ -4,15 +4,15 @@
 import torch
 
 try:
+    from llmfoundry import optim, utils
+
     # Before importing any transformers models, we need to disable transformers flash attention if
     # we are in an environment with flash attention version <2. Transformers hard errors on a not properly
     # gated import otherwise.
-    import transformers
-
-    from llmfoundry.models.utils.flash_attn_checker import is_flash_v1_installed
+    from llmfoundry.utils.flash_attn_checker import is_flash_v1_installed
     if is_flash_v1_installed():
+        import transformers
         transformers.utils.is_flash_attn_available = lambda: False
-    from llmfoundry import optim, utils
     from llmfoundry.data import (ConcatTokensDataset,
                                  MixtureOfDenoisersCollator, NoConcatDataset,
                                  Seq2SeqFinetuningCollator,

diff --git a/llmfoundry/models/hf/hf_causal_lm.py b/llmfoundry/models/hf/hf_causal_lm.py
@@ -25,7 +25,7 @@
 from llmfoundry.models.hf.hf_fsdp import hf_get_init_device
 from llmfoundry.models.hf.model_wrapper import HuggingFaceModelWithZLoss
 from llmfoundry.models.utils import init_empty_weights
-from llmfoundry.models.utils.flash_attn_checker import is_flash_v2_installed
+from llmfoundry.utils.flash_attn_checker import is_flash_v2_installed
 
 try:
     from peft.peft_model import PeftModel

diff --git a/llmfoundry/models/layers/attention.py b/llmfoundry/models/layers/attention.py
@@ -15,7 +15,7 @@
 
 from llmfoundry.models.layers.fc import FC_CLASS_REGISTRY
 from llmfoundry.models.layers.norm import NORM_CLASS_REGISTRY
-from llmfoundry.models.utils.flash_attn_checker import (is_flash_v1_installed,
+from llmfoundry.utils.flash_attn_checker import (is_flash_v1_installed,
                                                         is_flash_v2_installed)
 
 

diff --git a/...oundry/models/utils/flash_attn_checker.py → llmfoundry/utils/flash_attn_checker.py b/...oundry/models/utils/flash_attn_checker.py → llmfoundry/utils/flash_attn_checker.py
diff --git a/tests/test_huggingface_flash.py b/tests/test_huggingface_flash.py
@@ -14,7 +14,7 @@
 
 from llmfoundry import COMPOSER_MODEL_REGISTRY
 from llmfoundry.models.hf.hf_fsdp import rgetattr
-from llmfoundry.models.utils.flash_attn_checker import (is_flash_v1_installed,
+from llmfoundry.utils.flash_attn_checker import (is_flash_v1_installed,
                                                         is_flash_v2_installed)
 from llmfoundry.utils.builders import build_tokenizer