Update transformers to 4.51 (#1790)

Author: dakinggg

llmfoundry/callbacks/curriculum_learning_callback.py

@@ -237,7 +237,7 @@ def _build_train_loader(
         try:
             return build_dataloader(
                 train_loader_config,
-                self._tokenizer,
+                self._tokenizer,  # type: ignore
                 self._device_train_batch_size,
             )
         except BaseContextualError as e:

llmfoundry/callbacks/hf_checkpointer.py

@@ -13,7 +13,7 @@
 import warnings
 from multiprocessing.context import SpawnProcess
 from pathlib import Path
-from typing import Any, Optional, Sequence, Union
+from typing import TYPE_CHECKING, Any, Optional, Sequence, Union
 
 import numpy as np
 import torch
@@ -48,6 +48,9 @@
 from llmfoundry.utils.huggingface_hub_utils import \
     edit_files_for_hf_compatibility
 
+if TYPE_CHECKING:
+    from peft import PeftModel
+
 try:
     import transformer_engine.pytorch as te
     is_te_imported = True
@@ -487,9 +490,9 @@ def _any_register_processes_error(self, device: Device) -> bool:
 
     def transform_model_and_tokenizer(
         self,
-        model: PreTrainedModel,
+        model: Union[PreTrainedModel, 'PeftModel'],
         tokenizer: PreTrainedTokenizerBase,
-    ) -> tuple[PreTrainedModel, PreTrainedTokenizerBase]:
+    ) -> tuple[Union[PreTrainedModel, 'PeftModel'], PreTrainedTokenizerBase]:
         """Transform the model and tokenizer before saving.
 
         This allows a subclass to modify the model and tokenizer before saving. The base class implementation will
@@ -537,8 +540,8 @@ def pre_register_edit(self, local_save_path: str):
 
     def transform_model_pre_registration(
         self,
-        model: PreTrainedModel,
-    ) -> PreTrainedModel:
+        model: Union[PreTrainedModel, 'PeftModel'],
+    ) -> Union[PreTrainedModel, 'PeftModel']:
         """Transform the model before registering with MLflow.
 
         This allows a subclass to modify the model before registering with MLflow. The base class implementation will
@@ -565,17 +568,23 @@ def _get_hf_model(self, state: State):
         log.debug('Gathering state dict')
 
         if state.is_model_ddp:
-            original_model: PreTrainedModel = state.model.module.model  # type: ignore
+            original_model: Union[
+                PreTrainedModel,
+                'PeftModel'] = state.model.module.model  # type: ignore
             state_dict_model = state.model.module.model  # type: ignore
-            original_tokenizer = state.model.module.tokenizer  # type: ignore
+            original_tokenizer: PreTrainedTokenizerBase = state.model.module.tokenizer  # type: ignore
         elif isinstance(state.model.model, FSDP):
-            original_model: PreTrainedModel = state.model.model.module  # type: ignore
+            original_model: Union[
+                PreTrainedModel,
+                'PeftModel'] = state.model.model.module  # type: ignore
             state_dict_model = state.model.model  # type: ignore
-            original_tokenizer = state.model.tokenizer  # type: ignore
+            original_tokenizer: PreTrainedTokenizerBase = state.model.tokenizer  # type: ignore
         else:
-            original_model: PreTrainedModel = state.model.model  # type: ignore
+            original_model: Union[
+                PreTrainedModel,
+                'PeftModel'] = state.model.model  # type: ignore
             state_dict_model = state.model.model  # type: ignore
-            original_tokenizer = state.model.tokenizer  # type: ignore
+            original_tokenizer: PreTrainedTokenizerBase = state.model.tokenizer  # type: ignore
 
         cpu_offload = True
 
@@ -631,7 +640,7 @@ def tensor_hook(
 
             # Transform HF config before building 2nd model copy
             new_config = self.transform_config(
-                original_config=original_model.config,
+                original_config=original_model.config,  # type: ignore
             )
 
             log.debug(f'Creating new model instance')
@@ -640,25 +649,33 @@ def tensor_hook(
             # initialization cost.
             with init_empty_weights():
                 if self.using_peft:
-                    active_adapter = original_model.active_adapter
-                    base_model = original_model.get_base_model()
+                    from peft import PeftModel
+                    assert isinstance(original_model, PeftModel)
+                    active_adapter = original_model.active_adapter  # type: ignore
+                    base_model: PreTrainedModel = original_model.get_base_model(  # type: ignore
+                    )
                     new_base_model_instance = type(base_model)(new_config)
 
                     new_model_instance = type(original_model)(
-                        new_base_model_instance,
-                        original_model.peft_config[active_adapter],
+                        new_base_model_instance,  # type: ignore
+                        original_model.
+                        peft_config[active_adapter],  # type: ignore
                     )
                     del new_base_model_instance
                 else:
+                    assert isinstance(original_model, PreTrainedModel)
                     new_model_instance = type(original_model)(new_config)
                     if new_model_instance.generation_config is not None:
+                        assert original_model.generation_config is not None
                         new_model_instance.generation_config.update(
                             **original_model.generation_config.to_dict(),
                         )
 
             # Then load the state dict in with "assign" so that the state dict
             # is loaded properly even though the model is initially on meta device.
-            new_model_instance.load_state_dict(state_dict, assign=True)
+            new_model_instance.load_state_dict(  # type: ignore
+                state_dict, assign=True,
+            )
             del state_dict
 
             # Transform the model and tokenizer before saving
@@ -671,11 +688,14 @@ def tensor_hook(
             if self.pretrained_model_name is not None:
                 new_model_instance.name_or_path = self.pretrained_model_name
                 if self.using_peft:
+                    from peft import PeftModel
+                    assert isinstance(new_model_instance, PeftModel)
                     new_model_instance.base_model.name_or_path = self.pretrained_model_name
-                    for k in new_model_instance.peft_config.keys():
-                        new_model_instance.peft_config[
+                    for k in new_model_instance.peft_config.keys(  # type: ignore
+                    ):
+                        new_model_instance.peft_config[  # type: ignore
                             k
-                        ].base_model_name_or_path = self.pretrained_model_name
+                        ].base_model_name_or_path = self.pretrained_model_name  # type: ignore
 
             log.debug('Saving Hugging Face checkpoint to disk')
 
@@ -686,7 +706,7 @@ def _register_hf_model(
         temp_save_dir: str,
         original_tokenizer: PreTrainedTokenizerBase,
         use_temp_dir: bool,
-        new_model_instance: PreTrainedModel,
+        new_model_instance: Union[PreTrainedModel, 'PeftModel'],
     ):
         assert new_model_instance is not None
         new_model_instance = self.transform_model_pre_registration(
@@ -802,7 +822,7 @@ def _save_checkpoint(
                     )
 
                 # Only need to edit files for MPT because it has custom code
-                if new_model_instance.config.model_type == 'mpt':
+                if new_model_instance.config.model_type == 'mpt':  # type: ignore
                     log.debug('Editing MPT files for HuggingFace compatibility')
                     edit_files_for_hf_compatibility(
                         temp_save_dir,
@@ -837,6 +857,12 @@ def _save_checkpoint(
                     None,
                 )
                 if model_name is not None:
+                    from peft import PeftModel
+                    assert isinstance(new_model_instance, PeftModel)
+                    assert isinstance(
+                        new_model_instance.model,
+                        PreTrainedModel,
+                    )
                     new_model_instance.name_or_path = model_name
                     new_model_instance.model.name_or_path = model_name
                     new_model_instance.base_model.name_or_path = model_name

llmfoundry/command_utils/data_prep/convert_dataset_hf.py

@@ -169,7 +169,7 @@ def build_hf_dataset(
     bos_text: str = '',
     eos_text: str = '',
     no_wrap: bool = False,
-    tokenizer: PreTrainedTokenizerBase = None,
+    tokenizer: Optional[PreTrainedTokenizerBase] = None,
     data_subset: Union[str, None] = None,
 ) -> IterableDataset:
     """Build an IterableDataset over the HF C4 or pile source data.
@@ -206,9 +206,10 @@ def build_hf_dataset(
             raise ValueError(f'max_length must be set.')
         if bos_text + eos_text == '':
             test_tokens = tokenizer('test')
-            if test_tokens['input_ids'][
-                0] != tokenizer.bos_token_id and test_tokens['input_ids'][
-                    -1] != tokenizer.eos_token_id:
+            if test_tokens['input_ids'][  # type: ignore
+                0] != tokenizer.bos_token_id and test_tokens[
+                    'input_ids'][  # type: ignore
+                        -1] != tokenizer.eos_token_id:
                 tok_error_msg = 'This tokenizer does not insert an EOS nor BOS token. '
                 tok_error_msg += 'Concatenating with this tokenizer will result in sequences being '
                 tok_error_msg += 'attached without a separating token. Please use another tokenizer, '

llmfoundry/command_utils/data_prep/convert_dataset_json.py

@@ -29,7 +29,7 @@ def build_hf_dataset(
     bos_text: str = '',
     eos_text: str = '',
     no_wrap: bool = False,
-    tokenizer: PreTrainedTokenizerBase = None,
+    tokenizer: Optional[PreTrainedTokenizerBase] = None,
 ) -> IterableDataset:
     """Build an IterableDataset over the HF C4 or pile source data.
 
@@ -70,9 +70,10 @@ def build_hf_dataset(
             raise ValueError(f'max_length must be set.')
         if bos_text + eos_text == '':
             test_tokens = tokenizer('test')
-            if test_tokens['input_ids'][
-                0] != tokenizer.bos_token_id and test_tokens['input_ids'][
-                    -1] != tokenizer.eos_token_id:
+            if test_tokens['input_ids'][  # type: ignore
+                0] != tokenizer.bos_token_id and test_tokens[
+                    'input_ids'][  # type: ignore
+                        -1] != tokenizer.eos_token_id:
                 tok_error_msg = 'This tokenizer does not insert an EOS nor BOS token. '
                 tok_error_msg += 'Concatenating with this tokenizer will result in sequences being '
                 tok_error_msg += 'attached without a separating token. Please use another tokenizer, '
@@ -118,6 +119,7 @@ def convert_dataset_json(
     """
     if concat_tokens is not None:
         mode = ConcatMode.CONCAT_TOKENS
+        assert tokenizer is not None
         built_tokenizer = AutoTokenizer.from_pretrained(tokenizer)
         # we will enforce length, so suppress warnings about sequences too long for the model
         built_tokenizer.model_max_length = int(1e30)

llmfoundry/command_utils/data_prep/convert_finetuning_dataset.py

@@ -167,8 +167,9 @@ def convert_finetuning_dataset(
 
     tokenizer_kwargs = tokenizer_kwargs
     tokenizer_kwargs.update({'model_max_length': max_seq_len})
+    built_tokenizer = None
     if tokenizer:
-        tokenizer = build_tokenizer(tokenizer, tokenizer_kwargs)
+        built_tokenizer = build_tokenizer(tokenizer, tokenizer_kwargs)
 
     for i, split_name in enumerate(splits):
         data_file = None
@@ -184,7 +185,7 @@ def convert_finetuning_dataset(
         # Determine the output columns
         columns, example_type = get_columns_and_format(
             dataset=loaded_dataset,
-            tokenizing=tokenizer is not None,
+            tokenizing=built_tokenizer is not None,
             preprocessing_fn=preprocessing_fn,
         )
         # Prepare the iterables
@@ -226,10 +227,10 @@ def convert_finetuning_dataset(
                         'Encountered an error when checking example for proper formatting. ' +\
                         f'example={formatted_sample}',
                     ) from e
-                if tokenizer is not None:
+                if built_tokenizer is not None:
                     sample = tokenize_formatted_example(
                         formatted_sample,
-                        tokenizer=tokenizer,
+                        tokenizer=built_tokenizer,
                     )
                     if not is_valid_ift_example(
                         max_seq_len,
@@ -259,7 +260,7 @@ def convert_finetuning_dataset(
                     else:
                         out.write(formatted_sample)
 
-        if tokenizer is not None and examples_removed > 0:
+        if built_tokenizer is not None and examples_removed > 0:
             warnings.warn(
                 f'Dropped {examples_removed} examples where the prompt was longer than {max_seq_len}, '
                 +

llmfoundry/command_utils/data_prep/convert_text_to_mds.py

@@ -9,7 +9,7 @@
 from concurrent.futures import ProcessPoolExecutor
 from functools import partial
 from glob import glob
-from typing import Iterable, Optional, cast
+from typing import Any, Iterable, Optional, cast
 
 import numpy as np
 from composer.utils import (
@@ -81,7 +81,7 @@ def __iter__(self) -> Iterable[dict[str, NDArray]]:
                             truncation=False,
                             padding=False,
                         )
-                        iids = encoded['input_ids']
+                        iids = cast(Any, encoded['input_ids'])
 
                         # If this is not the first chunk, remove the BOS token
                         if not first_chunk:

llmfoundry/command_utils/eval.py

@@ -80,13 +80,13 @@ def evaluate_model(
     tokenizer_cfg = tokenizer
     tokenizer_name = tokenizer_cfg['name']
     tokenizer_kwargs = tokenizer_cfg.get('kwargs', {})
-    tokenizer = build_tokenizer(tokenizer_name, tokenizer_kwargs)
+    built_tokenizer = build_tokenizer(tokenizer_name, tokenizer_kwargs)
 
     evaluators, logger_keys, eval_gauntlet_callback = build_evaluators(
         eval_loader_config,
         icl_tasks,
         eval_gauntlet_config,
-        tokenizer=tokenizer,
+        tokenizer=built_tokenizer,
         device_eval_batch_size=device_eval_batch_size,
         icl_seq_len=max_seq_len,
         icl_subset_num_batches=icl_subset_num_batches,
@@ -124,7 +124,7 @@ def evaluate_model(
     name = model.pop('name')
     composer_model = build_composer_model(
         name=name,
-        tokenizer=tokenizer,
+        tokenizer=built_tokenizer,
         init_context=init_context,
         cfg=model,
     )

llmfoundry/data/contrastive_pairs/dataloader.py

@@ -175,8 +175,8 @@ def _tokenize(
         text_samples_negatives = text_samples['negative']
         assert isinstance(text_samples_negatives, list)  # pyright type check
         text_samples_list.extend(text_samples_negatives)
-        return self.tokenizer(
-            text_samples_list,
+        return self.tokenizer(  # type: ignore
+            text_samples_list,  # type: ignore
             truncation=True,
             padding='max_length',
             max_length=self.max_seq_len,

llmfoundry/data/data.py

@@ -72,24 +72,29 @@ def __init__(
         self.eos_text = eos_text
         self.should_wrap = not no_wrap
 
-        self.bos_tokens = self.tokenizer(
+        bos_ids = self.tokenizer(
             self.bos_text,
             truncation=False,
             padding=False,
             add_special_tokens=False,
         )['input_ids']
+        assert isinstance(bos_ids, list)
+
+        self.bos_tokens: list[int] = bos_ids
         if len(self.bos_tokens) > 1:
             warnings.warn(
                 f'You specified --concat_tokens with --bos_text, but your BOS text is not tokenizing to one token\
                 , instead we got {self.bos_tokens}. Quit if this was in error.',
             )
 
-        self.eos_tokens = self.tokenizer(
+        eos_ids = self.tokenizer(
             self.eos_text,
             truncation=False,
             padding=False,
             add_special_tokens=False,
         )['input_ids']
+        assert isinstance(eos_ids, list)
+        self.eos_tokens: list[int] = eos_ids
         if len(self.eos_tokens) > 1:
             warnings.warn(
                 f'You specified --concat_tokens with --eos_text, but your EOS text is not tokenizing to one token\
@@ -99,8 +104,10 @@ def __init__(
         eos_text_provided = self.eos_text != ''
         bos_text_provided = self.bos_text != ''
         test_text = self.tokenizer('')
+        test_text_iids = test_text['input_ids']
+        assert isinstance(test_text_iids, list)
         if len(
-            test_text['input_ids'],
+            test_text_iids,
         ) > 0 and (eos_text_provided or bos_text_provided):
             message = 'both eos and bos' if eos_text_provided and bos_text_provided else (
                 'eos_text' if eos_text_provided else 'bos_text'
@@ -155,11 +162,12 @@ def __iter__(self) -> Iterable[dict[str, NDArray]]:
         buffer = []
         for sample in self.hf_dataset:
             encoded = self.tokenizer(
-                sample['text'],
+                sample['text'],  # type: ignore
                 truncation=False,
                 padding=False,
             )
             iids = encoded['input_ids']
+            assert isinstance(iids, list)
             buffer = buffer + self.bos_tokens + iids + self.eos_tokens
             while len(buffer) >= self.max_length:
                 concat_sample = buffer[:self.max_length]