Clean up the publicly exported API

llmfoundry/__init__.py

@@ -19,47 +19,39 @@
 
 hf_dynamic_modules_logger.addFilter(new_files_warning_filter)
 
-from llmfoundry import algorithms, callbacks, loggers, optim, registry, utils
-from llmfoundry.data import (ConcatTokensDataset, NoConcatDataset,
-                             Seq2SeqFinetuningCollator,
-                             build_finetuning_dataloader)
-from llmfoundry.models.hf import ComposerHFCausalLM, ComposerHFT5
-from llmfoundry.models.layers.attention import (
-    MultiheadAttention, attn_bias_shape, build_alibi_bias, build_attn_bias,
-    flash_attn_fn, scaled_multihead_dot_product_attention)
-from llmfoundry.models.layers.blocks import MPTBlock
-from llmfoundry.models.layers.ffn import MPTMLP
+from llmfoundry import (algorithms, callbacks, cli, data, eval, interfaces,
+                        loggers, metrics, models, optim, tokenizers, utils)
+from llmfoundry.data import StreamingFinetuningDataset, StreamingTextDataset
+from llmfoundry.eval import InContextLearningDataset, InContextLearningMetric
+from llmfoundry.models.hf import ComposerHFCausalLM
 from llmfoundry.models.mpt import (ComposerMPTCausalLM, MPTConfig,
                                    MPTForCausalLM, MPTModel, MPTPreTrainedModel)
-from llmfoundry.tokenizers import TiktokenTokenizerWrapper
+from llmfoundry.optim import DecoupledLionW
 
 __all__ = [
-    'build_finetuning_dataloader',
-    'Seq2SeqFinetuningCollator',
-    'MPTBlock',
-    'MPTMLP',
+    'StreamingFinetuningDataset',
+    'StreamingTextDataset',
+    'InContextLearningDataset',
+    'InContextLearningMetric',
+    'ComposerHFCausalLM',
     'MPTConfig',
     'MPTPreTrainedModel',
     'MPTModel',
     'MPTForCausalLM',
     'ComposerMPTCausalLM',
-    'ComposerHFCausalLM',
-    'ComposerHFT5',
-    'scaled_multihead_dot_product_attention',
-    'flash_attn_fn',
-    'MultiheadAttention',
-    'NoConcatDataset',
-    'ConcatTokensDataset',
-    'attn_bias_shape',
-    'build_attn_bias',
-    'build_alibi_bias',
-    'optim',
-    'utils',
-    'loggers',
+    'DecoupledLionW',
     'algorithms',
     'callbacks',
-    'TiktokenTokenizerWrapper',
-    'registry',
+    'cli',
+    'data',
+    'eval',
+    'interfaces',
+    'loggers',
+    'metrics',
+    'models',
+    'optim',
+    'tokenizers',
+    'utils',
 ]
 
 __version__ = '0.8.0.dev0'

llmfoundry/callbacks/async_eval_callback.py

@@ -27,6 +27,8 @@
 
 log = logging.getLogger(__name__)
 
+__all__ = ['AsyncEval']
+
 REQUIRED_PARAMS_FOR_EVAL = {
     'device_eval_batch_size',
     'icl_tasks',  # only required for eval, may not be specified in pure training

llmfoundry/callbacks/curriculum_learning_callback.py

@@ -20,6 +20,8 @@
 
 log = logging.getLogger(__name__)
 
+__all__ = ['CurriculumLearning']
+
 
 @experimental_class('CurriculumLearning callback')
 class CurriculumLearning(CallbackWithConfig):

llmfoundry/callbacks/fdiff_callback.py

@@ -8,6 +8,8 @@
 from composer.core import Callback, State
 from composer.loggers import Logger
 
+__all__ = ['FDiffMetrics']
+
 
 class FDiffMetrics(Callback):
     """Rate of change of metrics.

llmfoundry/callbacks/hf_checkpointer.py

@@ -35,6 +35,8 @@
 
 log = logging.getLogger(__name__)
 
+__all__ = ['HuggingFaceCheckpointer']
+
 _LICENSE_FILE_PATTERN = re.compile(r'license(\.[a-z]+|$)', re.IGNORECASE)
 
 

llmfoundry/callbacks/log_mbmoe_tok_per_expert_callback.py

@@ -9,6 +9,8 @@
 from composer.loggers import Logger
 from composer.utils import dist
 
+__all__ = ['MegaBlocksMoE_TokPerExpert']
+
 
 class MegaBlocksMoE_TokPerExpert(Callback):
     """Log tokens per expert for MegaBlocks MoE.
@@ -44,7 +46,7 @@ class MegaBlocksMoE_TokPerExpert(Callback):
 
     Args:
         log_interval (int, optional): The interval on which to log (Default: 10).
-        log_every_layer (bool, optional): Enable logging ever layer's statisictics (True) or log
+        log_every_layer (bool, optional): Enable logging ever layer's statistics (True) or log
             only aggregate statistics (Default: False).
         all_reduce_stats (bool, optional): Enable aggregating statistics across gpus (True) or log
             statistics for GPU 0 (Default: False).

llmfoundry/callbacks/monolithic_ckpt_callback.py

@@ -15,6 +15,8 @@
 from composer.utils import (dist, format_name_with_dist_and_time, parse_uri,
                             reproducibility)
 
+__all__ = ['MonolithicCheckpointSaver']
+
 
 class MonolithicCheckpointSaver(Callback):
     """Save a monolithic checkpoint every N batches.

llmfoundry/callbacks/resumption_callbacks.py

@@ -7,6 +7,8 @@
 from composer.core import Callback, State
 from composer.loggers import Logger
 
+from llmfoundry.utils.warnings import experimental_class
+
 __all__ = [
     'GlobalLRScaling',
     'LayerFreezing',
@@ -15,6 +17,7 @@
 log = logging.getLogger(__name__)
 
 
+@experimental_class('GlobalLRScaling')
 class GlobalLRScaling(Callback):
     """GlobalLRScaling.
 
@@ -52,6 +55,7 @@ def fit_start(self, state: State, logger: Logger) -> None:
             ]
 
 
+@experimental_class('LayerFreezing')
 class LayerFreezing(Callback):
     """LayerFreezing.
 

llmfoundry/callbacks/scheduled_gc_callback.py

@@ -8,6 +8,8 @@
 from composer.core import Callback, State
 from composer.loggers import Logger
 
+__all__ = ['ScheduledGarbageCollector']
+
 
 def gc_cuda():
     """Garbage collect Torch (CUDA) memory."""

llmfoundry/data/__init__.py

@@ -4,9 +4,14 @@
 from llmfoundry.data.data import ConcatTokensDataset, NoConcatDataset
 from llmfoundry.data.dataloader import build_dataloader
 from llmfoundry.data.finetuning import (Seq2SeqFinetuningCollator,
+                                        StreamingFinetuningDataset,
                                         build_finetuning_dataloader)
-from llmfoundry.data.text_data import (StreamingTextDataset,
-                                       build_text_dataloader)
+from llmfoundry.data.packing import (BinPackCollator, auto_packing_ratio,
+                                     profile_packing)
+from llmfoundry.data.text_data import (ConcatenatedSequenceCollatorWrapper,
+                                       StreamingTextDataset,
+                                       build_text_dataloader,
+                                       get_tokens_per_batch_func)
 from llmfoundry.registry import dataloaders
 
 dataloaders.register('text', func=build_text_dataloader)
@@ -15,9 +20,15 @@
 __all__ = [
     'Seq2SeqFinetuningCollator',
     'build_finetuning_dataloader',
+    'StreamingFinetuningDataset',
     'StreamingTextDataset',
     'build_text_dataloader',
     'NoConcatDataset',
     'ConcatTokensDataset',
     'build_dataloader',
+    'BinPackCollator',
+    'auto_packing_ratio',
+    'profile_packing',
+    'ConcatenatedSequenceCollatorWrapper',
+    'get_tokens_per_batch_func',
 ]

llmfoundry/data/data.py

@@ -11,6 +11,11 @@
 from torch.utils.data import IterableDataset
 from transformers import PreTrainedTokenizerBase
 
+__all__ = [
+    'ConcatTokensDataset',
+    'NoConcatDataset',
+]
+
 
 class NoConcatDataset(IterableDataset):
     """An IterableDataset that returns text samples for MDSWriter.

llmfoundry/data/dataloader.py

@@ -10,6 +10,10 @@
 from llmfoundry import registry
 from llmfoundry.utils.registry_utils import construct_from_registry
 
+__all__ = [
+    'build_dataloader',
+]
+
 
 def build_dataloader(cfg: DictConfig, tokenizer: PreTrainedTokenizerBase,
                      device_batch_size: int) -> DataSpec:

llmfoundry/data/finetuning/__init__.py

@@ -3,5 +3,16 @@
 
 from llmfoundry.data.finetuning.collator import Seq2SeqFinetuningCollator
 from llmfoundry.data.finetuning.dataloader import build_finetuning_dataloader
+from llmfoundry.data.finetuning.tasks import (StreamingFinetuningDataset,
+                                              dataset_constructor,
+                                              is_valid_ift_example,
+                                              tokenize_formatted_example)
 
-__all__ = ['Seq2SeqFinetuningCollator', 'build_finetuning_dataloader']
+__all__ = [
+    'Seq2SeqFinetuningCollator',
+    'build_finetuning_dataloader',
+    'dataset_constructor',
+    'tokenize_formatted_example',
+    'is_valid_ift_example',
+    'StreamingFinetuningDataset',
+]

llmfoundry/data/finetuning/collator.py

@@ -10,6 +10,10 @@
 
 log = logging.getLogger(__name__)
 
+__all__ = [
+    'Seq2SeqFinetuningCollator',
+]
+
 # HuggingFace hardcodes the ignore index to -100
 _HF_IGNORE_INDEX = -100
 

llmfoundry/data/finetuning/dataloader.py

@@ -23,6 +23,10 @@
 
 log = logging.getLogger(__name__)
 
+__all__ = [
+    'build_finetuning_dataloader',
+]
+
 # HuggingFace hardcodes the ignore index to -100
 _HF_IGNORE_INDEX = -100
 

llmfoundry/data/finetuning/tasks.py

@@ -72,7 +72,12 @@ def preprocessing_fn(example: Dict) -> Dict[str, str]:
 
 log = logging.getLogger(__name__)
 
-__all__ = ['dataset_constructor']
+__all__ = [
+    'dataset_constructor',
+    'tokenize_formatted_example',
+    'is_valid_ift_example',
+    'StreamingFinetuningDataset',
+]
 
 _ALLOWED_RESPONSE_KEYS = {'response', 'completion'}
 _ALLOWED_PROMPT_KEYS = {'prompt'}

llmfoundry/data/packing.py

@@ -13,6 +13,12 @@
 
 log = logging.getLogger(__name__)
 
+__all__ = [
+    'BinPackCollator',
+    'auto_packing_ratio',
+    'profile_packing',
+]
+
 
 class BinPackCollator:
     """Utility collator for packing to reduce padding."""

llmfoundry/data/text_data.py

@@ -22,6 +22,13 @@
 
 log = logging.getLogger(__name__)
 
+__all__ = [
+    'StreamingTextDataset',
+    'build_text_dataloader',
+    'ConcatenatedSequenceCollatorWrapper',
+    'get_tokens_per_batch_func',
+]
+
 
 class StreamingTextDataset(StreamingDataset):
     """Generic text dataset using MosaicML's StreamingDataset.

llmfoundry/eval/__init__.py

@@ -1,2 +1,31 @@
 # Copyright 2024 MosaicML LLM Foundry authors
 # SPDX-License-Identifier: Apache-2.0
+
+from llmfoundry.eval.datasets.in_context_learning_evaluation import (
+    InContextLearningCodeEvalDataset, InContextLearningDataset,
+    InContextLearningGenerationTaskWithAnswersDataset,
+    InContextLearningLMTaskDataset, InContextLearningMultipleChoiceTaskDataset,
+    InContextLearningSchemaTaskDataset, get_icl_task_dataloader)
+from llmfoundry.eval.metrics.nlp import (
+    InContextLearningCodeEvalAccuracy,
+    InContextLearningGenerationExactMatchAccuracy, InContextLearningLMAccuracy,
+    InContextLearningLMExpectedCalibrationError,
+    InContextLearningMCExpectedCalibrationError, InContextLearningMetric,
+    InContextLearningMultipleChoiceAccuracy)
+
+__all__ = [
+    'InContextLearningDataset',
+    'InContextLearningLMTaskDataset',
+    'InContextLearningMultipleChoiceTaskDataset',
+    'InContextLearningSchemaTaskDataset',
+    'InContextLearningCodeEvalDataset',
+    'InContextLearningGenerationTaskWithAnswersDataset',
+    'get_icl_task_dataloader',
+    'InContextLearningMetric',
+    'InContextLearningLMAccuracy',
+    'InContextLearningMultipleChoiceAccuracy',
+    'InContextLearningGenerationExactMatchAccuracy',
+    'InContextLearningCodeEvalAccuracy',
+    'InContextLearningLMExpectedCalibrationError',
+    'InContextLearningMCExpectedCalibrationError',
+]

llmfoundry/eval/datasets/__init__.py

@@ -8,11 +8,13 @@
     InContextLearningGenerationTaskWithAnswersDataset,
     InContextLearningLMTaskDataset, InContextLearningMultipleChoiceTaskDataset,
     InContextLearningSchemaTaskDataset, get_icl_task_dataloader)
-from llmfoundry.eval.datasets.utils import (get_continuation_span,
-                                            get_fewshot_sample_idxs,
-                                            make_padded_input, strip_data,
-                                            tokenizer_needs_prefix_space,
-                                            trim_context)
+
+# isort: off
+from llmfoundry.eval.datasets.utils import (
+    MultiTokenEOSCriteria, convert_tokens_to_tensors, get_continuation_span,
+    get_fewshot_sample_idxs, make_padded_input, stop_sequences_criteria,
+    strip_data, tokenizer_needs_prefix_space, trim_context)
+# isort: on
 
 __all__ = [
     'InContextLearningDataset',
@@ -22,10 +24,13 @@
     'InContextLearningMultipleChoiceTaskDataset',
     'InContextLearningSchemaTaskDataset',
     'get_icl_task_dataloader',
+    'MultiTokenEOSCriteria',
     'strip_data',
     'tokenizer_needs_prefix_space',
     'trim_context',
     'get_continuation_span',
-    'get_fewshot_sample_idxs',
     'make_padded_input',
+    'convert_tokens_to_tensors',
+    'get_fewshot_sample_idxs',
+    'stop_sequences_criteria',
 ]