diff --git a/llmfoundry/__init__.py b/llmfoundry/__init__.py index b7e69ff31d..f65338c987 100644 --- 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 index 8ba9ae988d..7dbed6af28 100644 --- 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 index 83af2153a2..54dc9bd5a8 100644 --- 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',