Making tokenizers optional in the building of LLMs (#1781)

Author: ethantang-db

Diff for: llmfoundry/command_utils/train.py

@@ -361,10 +361,13 @@ def train(cfg: DictConfig) -> Trainer:
             )
 
     # Build tokenizer
-    log.info('Building tokenizer...')
-    tokenizer_name = train_cfg.tokenizer['name']
-    tokenizer_kwargs = train_cfg.tokenizer.get('kwargs', {})
-    tokenizer = build_tokenizer(tokenizer_name, tokenizer_kwargs)
+    tokenizer = None
+    tokenizer_name = None
+    if train_cfg.tokenizer:
+        log.info('Building tokenizer...')
+        tokenizer_name = train_cfg.tokenizer['name']
+        tokenizer_kwargs = train_cfg.tokenizer.get('kwargs', {})
+        tokenizer = build_tokenizer(tokenizer_name, tokenizer_kwargs)
 
     # Scheduler
     scheduler_name: str = train_cfg.scheduler.pop('name')

Diff for: llmfoundry/data/contrastive_pairs/dataloader.py

@@ -185,7 +185,7 @@ def _tokenize(
 
 def build_pairs_dataloader(
     dataset: dict[str, Any],
-    tokenizer: PreTrainedTokenizerBase,
+    tokenizer: Optional[PreTrainedTokenizerBase],
     device_batch_size: int,
     drop_last: bool,
     num_workers: int,
@@ -195,6 +195,11 @@ def build_pairs_dataloader(
     timeout: int = 0,
     max_hard_negatives: Optional[int] = None,
 ) -> DataSpec:
+    if tokenizer is None:
+        raise ValueError(
+            'Tokenizer is required for contrastive pairs dataloader',
+        )
+
     dataset_cfg = dataset
     streams_dict = dataset.pop('streams', None)
     eos_token_id = dataset.pop('eos_token_id', None)

Diff for: llmfoundry/data/dataloader.py

@@ -3,7 +3,7 @@
 
 """Dataloader builder utilities."""
 
-from typing import Any, Union
+from typing import Any, Optional, Union
 
 from composer import DataSpec
 from transformers import PreTrainedTokenizerBase
@@ -18,14 +18,14 @@
 
 def build_dataloader(
     cfg: dict[str, Any],
-    tokenizer: PreTrainedTokenizerBase,
+    tokenizer: Optional[PreTrainedTokenizerBase],
     device_batch_size: Union[int, float],
 ) -> DataSpec:
     """Builds a dataloader from a config.
 
     Args:
         cfg (DictConfig): An omegaconf dictionary used to configure the loader.
-        tokenizer (PreTrainedTokenizerBase): The tokenizer that the model will use.
+        tokenizer (Optional[PreTrainedTokenizerBase]): The tokenizer that the model will use.
         device_batch_size (int): The size of the batches (number of examples)
             that the dataloader will produce.
     """

Diff for: llmfoundry/data/finetuning/collator.py

@@ -313,7 +313,6 @@ def __call__(self,
         else:
             batch = self._process_and_batch_encoder_decoder(examples)
 
-        # Add any batch_metadata
         batch_size = batch['input_ids'].shape[0]
         batch.update({
             k: torch.tensor([v] * batch_size)

Diff for: llmfoundry/data/finetuning/dataloader.py

@@ -54,7 +54,7 @@
 
 
 def build_finetuning_dataloader(
-    tokenizer: PreTrainedTokenizerBase,
+    tokenizer: Optional[PreTrainedTokenizerBase],
     device_batch_size: Union[int, float],
     dataset: dict[str, Any],
     num_workers: int,
@@ -179,6 +179,9 @@ def build_finetuning_dataloader(
         padding/waste rates for different `cfg.dataset.packing_ratio` choices,
         given a starting workload YAML.
     """
+    if tokenizer is None:
+        raise ValueError('Tokenizer is required for finetuning dataloader')
+
     dataset_cfg = dataset
     is_streaming = (
         dataset_cfg.get('remote') is not None or

Diff for: llmfoundry/data/text_data.py

@@ -301,7 +301,7 @@ def build_streams(streams: Optional[dict[str, Any]] = None,):
 
 
 def build_text_dataloader(
-    tokenizer: PreTrainedTokenizerBase,
+    tokenizer: Optional[PreTrainedTokenizerBase],
     device_batch_size: Union[int, float],
     dataset: dict[str, Any],
     drop_last: bool,
@@ -311,6 +311,8 @@ def build_text_dataloader(
     persistent_workers: bool = True,
     timeout: int = 0,
 ) -> DataSpec:
+    if tokenizer is None:
+        raise ValueError('Tokenizer is required for text dataloader')
 
     dataset_cfg = dataset
 

Diff for: llmfoundry/utils/builders.py

@@ -58,7 +58,7 @@ def build_evaluators(
     icl_tasks_config: Optional[Union[str, list[dict[str, Any]]]],
     eval_gauntlet_config: Optional[Union[str, dict[str, Any]]],
     *,
-    tokenizer: PreTrainedTokenizerBase,
+    tokenizer: Optional[PreTrainedTokenizerBase],
     device_eval_batch_size: Union[int, float],
     icl_seq_len: int,
     icl_subset_num_batches: Optional[int],
@@ -75,10 +75,13 @@ def build_evaluators(
     logger_keys = []
     eval_gauntlet_callback = None
     if icl_tasks_config is not None:
+        if tokenizer is None:
+            raise ValueError('Tokenizer is required for icl tasks')
         if not isinstance(device_eval_batch_size, int):
             raise ValueError(
                 'device_eval_batch_size should be an int for icl tasks.',
             )
+
         icl_evaluators, logger_keys, eval_gauntlet_callback = build_icl_data_and_gauntlet(
             icl_tasks_config,
             eval_gauntlet_config,
@@ -94,7 +97,7 @@ def build_evaluators(
 
 def build_eval_loaders(
     eval_loader_config: Union[dict[str, Any], list[dict[str, Any]]],
-    tokenizer: PreTrainedTokenizerBase,
+    tokenizer: Optional[PreTrainedTokenizerBase],
     device_eval_batch_size: Union[int, float],
 ) -> list[Evaluator]:
     evaluators: list[Evaluator] = []
@@ -225,7 +228,7 @@ def build_save_planner(name: str, **kwargs: Any) -> SavePlanner:
 def build_composer_model(
     name: str,
     cfg: dict[str, Any],
-    tokenizer: PreTrainedTokenizerBase,
+    tokenizer: Optional[PreTrainedTokenizerBase],
     init_context: Optional[ContextManager] = None,
     master_weights_dtype: Optional[str] = None,
 ) -> ComposerModel:
@@ -234,7 +237,7 @@ def build_composer_model(
     Args:
         name (str): Name of the model to build.
         cfg (DictConfig): Configuration for the model.
-        tokenizer (PreTrainedTokenizerBase): Tokenizer to use.
+        tokenizer (Optional[PreTrainedTokenizerBase]): Tokenizer to use.
         init_context (Optional[ContextManager], optional): Context manager to use for initialization. Defaults to None.
         master_weights_dtype (Optional[str], optional): Master weights dtype. Defaults to None.
 

Diff for: llmfoundry/utils/mosaicml_logger_utils.py

@@ -87,19 +87,21 @@ def log_train_analytics(
     train_loader_config: dict[str, Any],
     eval_loader_config: Optional[Union[dict[str, Any], list[dict[str, Any]]]],
     callback_configs: Optional[dict[str, Any]],
-    tokenizer_name: str,
+    tokenizer_name: Optional[str],
     load_path: Optional[str],
     icl_tasks_config: Optional[Union[list[dict[str, Any]], str]],
     eval_gauntlet: Optional[Union[dict[str, Any], str]],
 ):
     """Logs analytics for runs using the `train.py` script."""
     train_loader_dataset = train_loader_config.get('dataset', {})
     metrics: dict[str, Any] = {
-        'llmfoundry/tokenizer_name': tokenizer_name,
         'llmfoundry/script': 'train',
         'llmfoundry/train_loader_name': train_loader_config.get('name'),
     }
 
+    if tokenizer_name is not None:
+        metrics['llmfoundry/tokenizer_name'] = tokenizer_name
+
     if callback_configs is not None:
         metrics['llmfoundry/callbacks'] = [
             name for name, _ in callback_configs.items()

Diff for: tests/data/test_dataloader.py

@@ -26,7 +26,11 @@
 from llmfoundry.command_utils import convert_dataset_hf
 from llmfoundry.command_utils.data_prep.convert_finetuning_dataset import \
     get_columns_and_format
-from llmfoundry.data import build_dataloader, build_finetuning_dataloader
+from llmfoundry.data import (
+    build_dataloader,
+    build_finetuning_dataloader,
+    build_pairs_dataloader,
+)
 from llmfoundry.data.finetuning.collator import (
     validate_target_settings,
 )
@@ -1557,3 +1561,18 @@ def test_text_dataloader_with_extra_keys(tiny_gpt2_tokenizer: PreTrainedTokenize
                 tokenizer=tokenizer,
                 device_batch_size=device_batch_size,
             ).dataloader
+
+
+@pytest.mark.parametrize(
+        'build_fn',
+        [build_finetuning_dataloader, build_text_dataloader, build_pairs_dataloader])
+def test_tokenizer_none(build_fn: Callable):
+    params = {
+        'device_batch_size': 2,
+        'dataset': {},
+        'num_workers': 0,
+        'drop_last': False,
+    }
+
+    with pytest.raises(ValueError, match='Tokenizer is required'):
+        _ = build_fn(tokenizer=None, **params)

Diff for: tests/eval/test_in_context_learning_datasets.py

@@ -38,7 +38,7 @@
     InContextLearningLMAccuracy,
     InContextLearningMultipleChoiceAccuracy,
 )
-from llmfoundry.utils.builders import build_icl_evaluators
+from llmfoundry.utils.builders import build_evaluators, build_icl_evaluators
 
 
 def test_strip_data():
@@ -2652,3 +2652,16 @@ def test_bc_question_prelimiter(
     assert len(evaluators) == 1
     evaluator = evaluators[0]
     assert evaluator.dataloader.dataloader.dataset.prelimiter == 'This is a question: '  # type: ignore
+
+
+def test_icl_no_tokenizer():
+    with pytest.raises(ValueError, match='Tokenizer is required for icl tasks'):
+        _ = build_evaluators(
+            eval_loader_config=None,
+            icl_tasks_config=[],
+            eval_gauntlet_config=None,
+            tokenizer=None,
+            device_eval_batch_size=2,
+            icl_seq_len=128,
+            icl_subset_num_batches=2,
+        )