remove

mosaicml · Feb 29, 2024 · e202aac · e202aac
1 parent 5abbca0
commit e202aac
Show file tree

Hide file tree

Showing 3 changed files with 56 additions and 86 deletions.
diff --git a/llmfoundry/__init__.py b/llmfoundry/__init__.py
@@ -1,67 +1,49 @@
 # Copyright 2022 MosaicML LLM Foundry authors
 # SPDX-License-Identifier: Apache-2.0
 
-import torch
+import warnings
 
-try:
-    import warnings
+# bitsandbytes is a very noisy library. A lot of it is print statements that we can't easily suppress,
+# but we can at least suppress a bunch of spurious warnings.
+warnings.filterwarnings('ignore', category=UserWarning, module='bitsandbytes')
 
-    # bitsandbytes is a very noisy library. A lot of it is print statements that we can't easily suppress,
-    # but we can at least suppress a bunch of spurious warnings.
-    warnings.filterwarnings('ignore',
-                            category=UserWarning,
-                            module='bitsandbytes')
+import logging
 
-    import logging
+from llmfoundry.utils.logging_utils import SpecificWarningFilter
 
-    from llmfoundry.utils.logging_utils import SpecificWarningFilter
+# Filter out Hugging Face warning for not using a pinned revision of the model
+hf_dynamic_modules_logger = logging.getLogger(
+    'transformers.dynamic_module_utils')
+new_files_warning_filter = SpecificWarningFilter(
+    'A new version of the following files was downloaded from')
 
-    # Filter out Hugging Face warning for not using a pinned revision of the model
-    hf_dynamic_modules_logger = logging.getLogger(
-        'transformers.dynamic_module_utils')
-    new_files_warning_filter = SpecificWarningFilter(
-        'A new version of the following files was downloaded from')
+hf_dynamic_modules_logger.addFilter(new_files_warning_filter)
 
-    hf_dynamic_modules_logger.addFilter(new_files_warning_filter)
+# 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
 
-    # 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 import optim, utils
+from llmfoundry.data import (ConcatTokensDataset, MixtureOfDenoisersCollator,
+                             NoConcatDataset, Seq2SeqFinetuningCollator,
+                             build_finetuning_dataloader,
+                             build_text_denoising_dataloader)
+from llmfoundry.models.hf import (ComposerHFCausalLM, ComposerHFPrefixLM,
+                                  ComposerHFT5)
+from llmfoundry.models.layers.attention import (
+    MultiheadAttention, attn_bias_shape, build_alibi_bias, build_attn_bias,
+    flash_attn_fn, is_flash_v1_installed,
+    scaled_multihead_dot_product_attention, triton_flash_attn_fn)
+from llmfoundry.models.layers.blocks import MPTBlock
+from llmfoundry.models.layers.ffn import FFN_CLASS_REGISTRY, MPTMLP, build_ffn
+from llmfoundry.models.model_registry import COMPOSER_MODEL_REGISTRY
+from llmfoundry.models.mpt import (ComposerMPTCausalLM, MPTConfig,
+                                   MPTForCausalLM, MPTModel, MPTPreTrainedModel)
+from llmfoundry.tokenizers import TiktokenTokenizerWrapper
 
-    from llmfoundry import optim, utils
-    from llmfoundry.data import (ConcatTokensDataset,
-                                 MixtureOfDenoisersCollator, NoConcatDataset,
-                                 Seq2SeqFinetuningCollator,
-                                 build_finetuning_dataloader,
-                                 build_text_denoising_dataloader)
-    from llmfoundry.models.hf import (ComposerHFCausalLM, ComposerHFPrefixLM,
-                                      ComposerHFT5)
-    from llmfoundry.models.layers.attention import (
-        MultiheadAttention, attn_bias_shape, build_alibi_bias, build_attn_bias,
-        flash_attn_fn, is_flash_v1_installed,
-        scaled_multihead_dot_product_attention, triton_flash_attn_fn)
-    from llmfoundry.models.layers.blocks import MPTBlock
-    from llmfoundry.models.layers.ffn import (FFN_CLASS_REGISTRY, MPTMLP,
-                                              build_ffn)
-    from llmfoundry.models.model_registry import COMPOSER_MODEL_REGISTRY
-    from llmfoundry.models.mpt import (ComposerMPTCausalLM, MPTConfig,
-                                       MPTForCausalLM, MPTModel,
-                                       MPTPreTrainedModel)
-    from llmfoundry.tokenizers import TiktokenTokenizerWrapper
-    if is_flash_v1_installed():
-        transformers.utils.is_flash_attn_available = lambda: False
-
-except ImportError as e:
-    try:
-        is_cuda_available = torch.cuda.is_available()
-    except:
-        is_cuda_available = False
-
-    extras = '.[gpu]' if is_cuda_available else '.'
-    raise ImportError(
-        f'Please make sure to pip install {extras} to get the requirements for the LLM example.'
-    ) from e
+if is_flash_v1_installed():
+    transformers.utils.is_flash_attn_available = lambda: False
 
 __all__ = [
     'build_text_denoising_dataloader',

diff --git a/llmfoundry/callbacks/__init__.py b/llmfoundry/callbacks/__init__.py
@@ -1,23 +1,16 @@
 # Copyright 2022 MosaicML LLM Foundry authors
 # SPDX-License-Identifier: Apache-2.0
 
-try:
-    from llmfoundry.callbacks.async_eval_callback import AsyncEval
-    from llmfoundry.callbacks.curriculum_learning_callback import \
-        CurriculumLearning
-    from llmfoundry.callbacks.eval_gauntlet_callback import EvalGauntlet
-    from llmfoundry.callbacks.fdiff_callback import FDiffMetrics
-    from llmfoundry.callbacks.hf_checkpointer import HuggingFaceCheckpointer
-    from llmfoundry.callbacks.monolithic_ckpt_callback import \
-        MonolithicCheckpointSaver
-    from llmfoundry.callbacks.resumption_callbacks import (GlobalLRScaling,
-                                                           LayerFreezing)
-    from llmfoundry.callbacks.scheduled_gc_callback import \
-        ScheduledGarbageCollector
-except ImportError as e:
-    raise ImportError(
-        'Please make sure to pip install . to get requirements for llm-foundry.'
-    ) from e
+from llmfoundry.callbacks.async_eval_callback import AsyncEval
+from llmfoundry.callbacks.curriculum_learning_callback import CurriculumLearning
+from llmfoundry.callbacks.eval_gauntlet_callback import EvalGauntlet
+from llmfoundry.callbacks.fdiff_callback import FDiffMetrics
+from llmfoundry.callbacks.hf_checkpointer import HuggingFaceCheckpointer
+from llmfoundry.callbacks.monolithic_ckpt_callback import \
+    MonolithicCheckpointSaver
+from llmfoundry.callbacks.resumption_callbacks import (GlobalLRScaling,
+                                                       LayerFreezing)
+from llmfoundry.callbacks.scheduled_gc_callback import ScheduledGarbageCollector
 
 __all__ = [
     'FDiffMetrics',

diff --git a/llmfoundry/utils/__init__.py b/llmfoundry/utils/__init__.py
@@ -1,22 +1,17 @@
 # Copyright 2022 MosaicML LLM Foundry authors
 # SPDX-License-Identifier: Apache-2.0
 
-try:
-    from llmfoundry.utils.builders import (build_algorithm, build_callback,
-                                           build_icl_evaluators, build_logger,
-                                           build_optimizer, build_scheduler,
-                                           build_tokenizer)
-    from llmfoundry.utils.checkpoint_conversion_helpers import (
-        convert_and_save_ft_weights, get_hf_tokenizer_from_composer_state_dict)
-    from llmfoundry.utils.config_utils import (calculate_batch_size_info,
-                                               log_config, pop_config,
-                                               update_batch_size_info)
-    from llmfoundry.utils.model_download_utils import (
-        download_from_hf_hub, download_from_http_fileserver)
-except ImportError as e:
-    raise ImportError(
-        'Please make sure to pip install . to get requirements for llm-foundry.'
-    ) from e
+from llmfoundry.utils.builders import (build_algorithm, build_callback,
+                                       build_icl_evaluators, build_logger,
+                                       build_optimizer, build_scheduler,
+                                       build_tokenizer)
+from llmfoundry.utils.checkpoint_conversion_helpers import (
+    convert_and_save_ft_weights, get_hf_tokenizer_from_composer_state_dict)
+from llmfoundry.utils.config_utils import (calculate_batch_size_info,
+                                           log_config, pop_config,
+                                           update_batch_size_info)
+from llmfoundry.utils.model_download_utils import (
+    download_from_hf_hub, download_from_http_fileserver)
 
 __all__ = [
     'build_callback',