transformers 5.0.0rc0__py3-none-any.whl → 5.0.0rc2__py3-none-any.whl

transformers/tokenization_utils_tokenizers.py

@@ -20,18 +20,19 @@ import copy
 import json
 import os
 from collections import defaultdict
+from collections.abc import Iterable
 from shutil import copyfile
 from typing import Any, Optional, Union
 
 import tokenizers.pre_tokenizers as pre_tokenizers_fast
+from huggingface_hub import is_offline_mode
 from tokenizers import AddedToken, processors
 from tokenizers import Encoding as EncodingFast
 from tokenizers import Tokenizer as TokenizerFast
-from tokenizers import normalizers as tokenizers_normalizers
 from tokenizers.decoders import Decoder as DecoderFast
+from tokenizers.models import BPE, Unigram
 from tokenizers.trainers import BpeTrainer, UnigramTrainer, WordLevelTrainer, WordPieceTrainer
 
-from .convert_slow_tokenizer import convert_slow_tokenizer
 from .integrations.ggml import convert_gguf_tokenizer
 from .modeling_gguf_pytorch_utils import load_gguf_checkpoint
 from .tokenization_utils_base import (
@@ -41,8 +42,9 @@ from .tokenization_utils_base import (
     PreTrainedTokenizerBase,
     TextInput,
     TruncationStrategy,
+    generate_merges,
 )
-from .utils import PaddingStrategy, add_end_docstrings, is_offline_mode, logging
+from .utils import PaddingStrategy, add_end_docstrings, logging
 
 
 logger = logging.get_logger(__name__)
@@ -90,26 +92,162 @@ class TokenizersBackend(PreTrainedTokenizerBase):
     """
 
     vocab_files_names = VOCAB_FILES_NAMES
+    model = None
+    _tokenizer = None
+
+    @classmethod
+    def convert_to_native_format(cls, trust_remote_code=False, **kwargs):
+        """s
+        Build a `tokenizers.Tokenizer` backend from the available serialization files (tokenizer.json, sentencepiece
+        models, tekken.json, vocab/merges).
+        """
+        # Preserve kwargs for possible downstream use
+        local_kwargs = dict(kwargs)
+        fast_tokenizer_file = local_kwargs.pop("tokenizer_file", None)
+
+        if (
+            fast_tokenizer_file is not None
+            and os.path.isfile(fast_tokenizer_file)
+            and (cls is TokenizersBackend or "__init__" not in cls.__dict__ or trust_remote_code)
+        ):
+            local_kwargs["tokenizer_object"] = TokenizerFast.from_file(fast_tokenizer_file)
+            return local_kwargs
+        elif fast_tokenizer_file is not None and os.path.isfile(fast_tokenizer_file):
+            # we extract vocab / merges from the tokenizer file to pass them to __init__
+            processor = TokenizerFast.from_file(fast_tokenizer_file).post_processor
+            with open(fast_tokenizer_file, encoding="utf-8") as tokenizer_handle:
+                tokenizer_json = json.load(tokenizer_handle)
+            vocab = tokenizer_json.get("model", {}).get("vocab", None)
+            if cls.model is None:
+                if isinstance(vocab, list):
+                    vocab = list(map(tuple, vocab))  # TODO just for now
+            elif cls.model.__name__ == "Unigram":
+                if vocab and isinstance(vocab[0], (list, tuple)):
+                    vocab = [tuple(item) for item in vocab]
+            elif cls.model.__name__ == "WordLevel":
+                vocab = {token: i for i, token in enumerate(vocab)}
+            elif cls.model.__name__ == "BPE" or cls.model.__name__ == "WordPiece":
+                if isinstance(vocab, list):
+                    vocab = {token[0] if isinstance(token, list) else token: i for i, token in enumerate(vocab)}
+            local_kwargs["vocab"] = vocab
+
+            model_type = getattr(cls, "model", None)
+            if "merges" in tokenizer_json.get("model", {}) and (model_type and model_type.__name__ == "BPE"):
+                merges = tokenizer_json["model"]["merges"]
+                merges = [tuple(merge.split(" ")) if isinstance(merge, str) else tuple(merge) for merge in merges]
+                local_kwargs["merges"] = merges
+
+            if processor is not None:
+                local_kwargs["post_processor"] = processor
+            return local_kwargs
+
+        vocab_file = local_kwargs.get("vocab_file")
+        merges_file = local_kwargs.get("merges_file")
+        vocab = local_kwargs.get("vocab")
+        merges = local_kwargs.get("merges")
+
+        # Tekken converter (Mistral)
+        if isinstance(vocab_file, str) and vocab_file.endswith("tekken.json") and os.path.isfile(vocab_file):
+            from .convert_slow_tokenizer import MistralConverter
+
+            local_kwargs["vocab"], local_kwargs["merges"] = MistralConverter(
+                vocab_file=vocab_file
+            ).extract_vocab_merges_from_model(vocab_file)
+            return local_kwargs
+
+        # SentencePiece model (with TikToken fallback)
+        if isinstance(vocab_file, str) and os.path.isfile(vocab_file) and vocab_file.endswith(".model"):
+            try:
+                from .convert_slow_tokenizer import SentencePieceExtractor
+
+                local_kwargs = SentencePieceExtractor(vocab_file).extract(cls.model, **local_kwargs)
+                try:
+                    from .convert_slow_tokenizer import SLOW_TO_FAST_CONVERTERS
+
+                    converter_class = SLOW_TO_FAST_CONVERTERS.get(cls.__name__)
+                    if converter_class is not None and hasattr(converter_class, "convert_from_spm"):
+                        local_kwargs = converter_class.convert_from_spm(**local_kwargs)
+                except Exception as e:
+                    logger.warning(
+                        f"Could not reorder vocab using converter for {cls.__name__} due to {e}. Falling back to raw SentencePiece extraction."
+                    )
+                # what used to be in `convert_slow`
+                if hasattr(cls, "convert_from_spm_model"):
+                    local_kwargs = cls.convert_from_spm_model(**local_kwargs)
+            except Exception as e:  # TODO only catch deserialization error here!
+                logger.warning(
+                    f"Could not extract SentencePiece model from {vocab_file} using sentencepiece library due to {e}. "
+                    "Falling back to TikToken extractor."
+                )
+                from .convert_slow_tokenizer import TikTokenConverter
+
+                local_kwargs["vocab"], local_kwargs["merges"] = TikTokenConverter(
+                    vocab_file=vocab_file, extra_special_tokens=local_kwargs.get("extra_special_tokens")
+                ).extract_vocab_merges_from_model(vocab_file)
+
+            return local_kwargs
+
+        # Fallback to standard vocab/merges files if they existed!
+        if vocab is None and isinstance(vocab_file, str) and os.path.isfile(vocab_file):
+            local_kwargs["vocab"] = vocab_file
+            vocab = local_kwargs["vocab"]
+        if merges is None and isinstance(merges_file, str) and os.path.isfile(merges_file):
+            local_kwargs["merges"] = merges_file
+            merges = local_kwargs["merges"]
+
+        # Generate merges automatically when not provided for BPE tokenizers
+        if merges is None and cls.model is not None and cls.model.__name__ == "BPE" and isinstance(vocab, dict):
+            # Gather special tokens from kwargs to skip in merge generation
+            def _iter_special_tokens(values: Iterable[Any]) -> list[str]:
+                collected: list[str] = []
+                for val in values:
+                    if val is None:
+                        continue
+                    if isinstance(val, (list, tuple)):
+                        collected.extend(_iter_special_tokens(val))
+                    else:
+                        collected.append(str(val))
+                return collected
+
+            special_tokens_keys = [
+                "pad_token",
+                "unk_token",
+                "bos_token",
+                "eos_token",
+                "sep_token",
+                "cls_token",
+                "mask_token",
+                "additional_special_tokens",
+                "extra_special_tokens",
+            ]
+            skip_tokens: set[str] = set()
+            for key in special_tokens_keys:
+                if key in local_kwargs:
+                    skip_tokens.update(_iter_special_tokens([local_kwargs[key]]))
+
+            merges = generate_merges(vocab, skip_tokens=skip_tokens)
+            local_kwargs["merges"] = merges
+        return local_kwargs
 
     def __init__(self, *args, **kwargs):
         tokenizer_object = kwargs.pop("tokenizer_object", None)
-        slow_tokenizer = kwargs.pop("__slow_tokenizer", None)
         gguf_file = kwargs.pop("gguf_file", None)
         fast_tokenizer_file = kwargs.pop("tokenizer_file", None)
-        from_slow = kwargs.pop("from_slow", False)
         # Note: added_tokens_decoder is NOT popped - it's passed to super().__init__() for processing
         added_tokens_decoder = kwargs.get("added_tokens_decoder", {})
         # Store add_prefix_space before super().__init__() to ensure it's not overridden
         add_prefix_space = kwargs.get("add_prefix_space", False)
+        vocab_file = kwargs.get("vocab_file")
+
+        vocab = kwargs.get("vocab")
+        merges = kwargs.get("merges")
 
+        fast_tokenizer = None
         if tokenizer_object is not None:
             fast_tokenizer = copy.deepcopy(tokenizer_object)
-        elif fast_tokenizer_file is not None and not from_slow:
+        elif fast_tokenizer_file is not None and os.path.isfile(fast_tokenizer_file):
             # We have a serialization from tokenizers which let us directly build the backend
             fast_tokenizer = TokenizerFast.from_file(fast_tokenizer_file)
-        elif slow_tokenizer:
-            # We need to convert a slow tokenizer to build the backend
-            fast_tokenizer = convert_slow_tokenizer(slow_tokenizer)
         elif gguf_file is not None:
             # We need to convert a slow tokenizer to build the backend
             gguf_param = load_gguf_checkpoint(kwargs.get("vocab_file"))
@@ -120,18 +258,16 @@ class TokenizersBackend(PreTrainedTokenizerBase):
             kwargs.update(tokenizer_config)
             if len(additional_kwargs) > 0:
                 kwargs.update(additional_kwargs)
-        elif self.slow_tokenizer_class is not None and slow_tokenizer is not False:
-            # We need to create and convert a slow tokenizer to build the backend
-            slow_tokenizer = self.slow_tokenizer_class(*args, **kwargs)
-            fast_tokenizer = convert_slow_tokenizer(slow_tokenizer)
-        elif not slow_tokenizer:
-            # We tried loading a slow_tokenizer with spm and failed, try to load with tiktoken
-            self.vocab_file = kwargs.get("vocab_file")
-            # V5: Set _extra_special_tokens directly for converter
-            self._extra_special_tokens = kwargs.get("extra_special_tokens", [])
-            fast_tokenizer = convert_slow_tokenizer(self, from_tiktoken=True)
-            slow_tokenizer = None
-        else:
+        elif self._tokenizer is None and vocab is not None:
+            # Build from vocab/merges extracted by convert_to_native_format
+            if merges is not None:
+                vocab_dict = vocab if isinstance(vocab, dict) else {w: i for i, (w, _) in enumerate(vocab)}
+                fast_tokenizer = TokenizerFast(BPE(vocab=vocab_dict, merges=merges, fuse_unk=True, dropout=None))
+            elif isinstance(vocab, dict):
+                fast_tokenizer = TokenizerFast(BPE(vocab=vocab, merges=[], fuse_unk=True, dropout=None))
+            elif isinstance(vocab, list) and vocab and isinstance(vocab[0], (tuple, list)):
+                fast_tokenizer = TokenizerFast(Unigram(vocab=vocab, unk_id=kwargs.get("unk_id", 0)))
+        elif self._tokenizer is None:
             raise ValueError(
                 "Couldn't instantiate the backend tokenizer from one of: \n"
                 "(1) a `tokenizers` library serialization file, \n"
@@ -139,11 +275,16 @@ class TokenizersBackend(PreTrainedTokenizerBase):
                 "(3) an equivalent slow tokenizer class to instantiate and convert. \n"
                 "You need to have sentencepiece or tiktoken installed to convert a slow tokenizer to a fast one."
             )
+        # Only set defaults when creating TokenizersBackend from scratch
+        if fast_tokenizer_file is None and tokenizer_object is None and self._tokenizer is None:
+            kwargs.setdefault("bos_token", "<s>")
+            kwargs.setdefault("eos_token", "</s>")
 
-        self._tokenizer = fast_tokenizer
+        if fast_tokenizer is not None:
+            self._tokenizer = fast_tokenizer
 
-        if slow_tokenizer is not None:
-            kwargs.update(slow_tokenizer.init_kwargs)
+        if self._tokenizer is None:
+            raise ValueError("The backend tokenizer is not correctly initialized.")
 
         _truncation = self._tokenizer.truncation
 
@@ -169,8 +310,17 @@ class TokenizersBackend(PreTrainedTokenizerBase):
         if "backend" not in kwargs:
             kwargs["backend"] = "tokenizers"
 
+        explicit_bos_eos_in_kwargs = "add_bos_token" in kwargs or "add_eos_token" in kwargs
+        self._add_bos_token = kwargs.get("add_bos_token", False)
+        self._add_eos_token = kwargs.get("add_eos_token", False)
+        if post_processor := kwargs.pop("post_processor", None):  # most reliable way to get the post-processor
+            self._tokenizer.post_processor = post_processor
+        self._should_update_post_processor = explicit_bos_eos_in_kwargs or self._tokenizer.post_processor is None
         # We call this after having initialized the backend tokenizer because we update it.
         super().__init__(**kwargs)
+
+        if vocab_file is not None:
+            self.vocab_file = vocab_file
         # Ensure add_prefix_space is set correctly after parent init
         self.add_prefix_space = add_prefix_space
         self._tokenizer.encode_special_tokens = self.split_special_tokens
@@ -210,7 +360,7 @@ class TokenizersBackend(PreTrainedTokenizerBase):
                 tokens.append(token)
             if tokens:
                 # These tokens are from the special tokens map
-                self.add_tokens(tokens, special_tokens=True)
+                self.add_tokens(tokens)
 
         try:
             vocab_size = self._tokenizer.get_vocab_size()
@@ -228,6 +378,12 @@ class TokenizersBackend(PreTrainedTokenizerBase):
                 **kwargs,
             )
 
+        self._should_update_post_processor = (
+            self._should_update_post_processor or self._tokenizer.post_processor is None
+        )
+        if self._should_update_post_processor:
+            self.update_post_processor()
+
     @property
     def is_fast(self) -> bool:
         return True
@@ -273,7 +429,7 @@ class TokenizersBackend(PreTrainedTokenizerBase):
         # If eos_token is None and add_eos_token is True, silently disable add_eos_token
         # This allows tokenizers to set add_eos_token even if eos_token is not configured
         if eos is None and self.add_eos_token:
-            self._add_eos_token = False
+            self.add_eos_token = False
             return
 
         single = f"{(bos + ':0 ') if self.add_bos_token else ''}$A:0{(' ' + eos + ':0') if self.add_eos_token else ''}"
@@ -320,98 +476,24 @@ class TokenizersBackend(PreTrainedTokenizerBase):
             if token_value is None:
                 continue
             if isinstance(token_value, AddedToken):
-                if self._tokenizer.token_to_id(str(token_value)) is None:
-                    tokens_to_add.append(token_value)
+                tokens_to_add.append(token_value)
             elif isinstance(token_value, str):
-                if self._tokenizer.token_to_id(token_value) is None:
-                    tokens_to_add.append(AddedToken(token_value, special=True, normalized=False))
+                tokens_to_add.append(AddedToken(token_value, special=True, normalized=False))
 
         # V5: Check extra special tokens
         for token in self._extra_special_tokens:
             if isinstance(token, AddedToken):
-                if self._tokenizer.token_to_id(str(token)) is None:
-                    tokens_to_add.append(token)
+                tokens_to_add.append(token)
             elif isinstance(token, str):
-                if self._tokenizer.token_to_id(token) is None:
-                    tokens_to_add.append(AddedToken(token, special=True, normalized=False))
+                tokens_to_add.append(AddedToken(token, special=True, normalized=False))
 
         if tokens_to_add:
             # Ensure special tokens are added as such to the backend
             self.add_tokens(tokens_to_add, special_tokens=True)
 
-        if hasattr(self, "_add_bos_token") or hasattr(self, "_add_eos_token"):
+        if getattr(self, "_should_update_post_processor", True) or self._tokenizer.post_processor is None:
             self.update_post_processor()
 
-        # Update add_prefix_space in the pre_tokenizer if needed
-        if hasattr(self, "add_prefix_space"):
-            try:
-                tokenizer_json = json.loads(self.backend_tokenizer.to_str())
-                pre_tok = tokenizer_json.get("pre_tokenizer", {})
-
-                # Recursively update add_prefix_space in pretokenizers
-                def update_add_prefix_space(pretok_dict, value):
-                    updated = False
-                    if pretok_dict.get("type") == "Sequence":
-                        for nested in pretok_dict.get("pretokenizers", []):
-                            updated |= update_add_prefix_space(nested, value)
-                    elif "add_prefix_space" in pretok_dict and pretok_dict["add_prefix_space"] != value:
-                        pretok_dict["add_prefix_space"] = value
-                        updated = True
-                    return updated
-
-                if update_add_prefix_space(pre_tok, self.add_prefix_space):
-                    self._tokenizer = TokenizerFast.from_str(json.dumps(tokenizer_json))
-            except Exception:
-                pass
-
-        # Ensure normalizer flags (lowercase/accents/chinese chars) reflect tokenizer attributes
-        try:
-            normalizer = self.backend_tokenizer.normalizer
-            if normalizer is not None:
-                norm_state = json.loads(normalizer.__getstate__())
-                norm_type = norm_state.get("type")
-
-                desired_lowercase = getattr(self, "do_lower_case", None)
-                desired_strip_accents = getattr(self, "strip_accents", None)
-                # Some tokenizers expose keep_accents instead of strip_accents
-                if desired_strip_accents is None and hasattr(self, "keep_accents") and "strip_accents" in norm_state:
-                    keep_accents_value = getattr(self, "keep_accents")
-                    if keep_accents_value is not None:
-                        desired_strip_accents = not keep_accents_value
-                desired_handle_chinese = getattr(self, "tokenize_chinese_chars", None)
-
-                updated = False
-                if (
-                    desired_lowercase is not None
-                    and "lowercase" in norm_state
-                    and norm_state["lowercase"] != desired_lowercase
-                ):
-                    norm_state["lowercase"] = desired_lowercase
-                    updated = True
-                if (
-                    desired_strip_accents is not None
-                    and "strip_accents" in norm_state
-                    and norm_state["strip_accents"] != desired_strip_accents
-                ):
-                    norm_state["strip_accents"] = desired_strip_accents
-                    updated = True
-                if (
-                    desired_handle_chinese is not None
-                    and "handle_chinese_chars" in norm_state
-                    and norm_state["handle_chinese_chars"] != desired_handle_chinese
-                ):
-                    norm_state["handle_chinese_chars"] = desired_handle_chinese
-                    updated = True
-
-                if updated and norm_type is not None:
-                    norm_class = getattr(tokenizers_normalizers, norm_type, None)
-                    if norm_class is not None:
-                        norm_state.pop("type", None)
-                        self.backend_tokenizer.normalizer = norm_class(**norm_state)
-        except Exception:
-            # Best-effort: do not block initialization on normalizer reconciliation
-            pass
-
     @property
     def vocab_size(self) -> int:
         """
@@ -839,6 +921,8 @@ class TokenizersBackend(PreTrainedTokenizerBase):
 
         if isinstance(token_ids, int):
             token_ids = [token_ids]
+        if isinstance(token_ids, dict):
+            token_ids = token_ids["input_ids"]
         return self._tokenizer.decode(token_ids, skip_special_tokens=skip_special_tokens)
 
     def _save_pretrained(
@@ -1132,7 +1216,7 @@ class TokenizersBackend(PreTrainedTokenizerBase):
                         ]
                     ):
                         return tokenizer
-                elif transformers_version and version.parse(transformers_version) >= version.parse("4.57.3"):
+                elif transformers_version and version.parse(transformers_version) > version.parse("4.57.3"):
                     return tokenizer
 
                 mistral_config_detected = True

transformers/trainer.py

@@ -642,6 +642,16 @@ class Trainer:
                 "You should subclass `Trainer` and override the `create_optimizer_and_scheduler` method."
             )
         default_callbacks = DEFAULT_CALLBACKS + get_reporting_integration_callbacks(self.args.report_to)
+
+        # Add JIT checkpoint callback if enabled
+        if self.args.enable_jit_checkpoint:
+            from .trainer_jit_checkpoint import JITCheckpointCallback
+
+            jit_callback = JITCheckpointCallback()
+            default_callbacks = default_callbacks + [jit_callback]
+            # Set trainer reference for JIT callback after initialization
+            jit_callback.set_trainer(self)
+
         callbacks = default_callbacks if callbacks is None else default_callbacks + callbacks
         self.callback_handler = CallbackHandler(
             callbacks, self.model, self.processing_class, self.optimizer, self.lr_scheduler
@@ -1661,6 +1671,12 @@ class Trainer:
                 optimizer_cls = AdamW8bit
             else:
                 raise ValueError("Invalid optimizer")
+            optimizer_kwargs.update(
+                {
+                    "block_size": optim_args.get("block_size", 256),
+                    "bf16_stochastic_round": strtobool(optim_args.get("bf16_stochastic_round", "False")),
+                }
+            )
             optimizer_kwargs.update(adam_kwargs)
         elif args.optim in [
             OptimizerNames.SCHEDULE_FREE_RADAM,
@@ -2338,6 +2354,9 @@ class Trainer:
 
         if self.is_fsdp_enabled:
             self.model = self.model_wrapped = model
+            # Fix `got mixed torch.Tensor and DTensor` error in model.generate() for FSDP2 with LoRA
+            if hasattr(self.model, "generate"):
+                dist.fsdp.register_fsdp_forward_method(self.model, "generate")
 
         # for the rest of this function `model` is the outside model, whether it was wrapped or not
         if model is not self.model:
@@ -2428,8 +2447,6 @@ class Trainer:
 
         for epoch in range(epochs_trained, num_train_epochs):
             epoch_dataloader = train_dataloader
-            if hasattr(epoch_dataloader, "set_epoch"):
-                epoch_dataloader.set_epoch(epoch)
 
             steps_in_epoch = (
                 len(epoch_dataloader)
@@ -2450,6 +2467,9 @@ class Trainer:
                 elif steps_trained_in_current_epoch == 0:
                     self._load_rng_state(resume_from_checkpoint)
 
+            if hasattr(epoch_dataloader, "set_epoch"):
+                epoch_dataloader.set_epoch(epoch)
+
             epoch_iterator = iter(epoch_dataloader)
             # We chunkify the epoch iterator into gradient accumulation steps `n` batches
             remainder = steps_in_epoch % args.gradient_accumulation_steps
@@ -2788,7 +2808,7 @@ class Trainer:
                 )
             else:
                 # We load the model state dict on the CPU to avoid an OOM error.
-                if self.args.save_safetensors and os.path.isfile(safe_weights_file):
+                if os.path.isfile(safe_weights_file):
                     state_dict = safetensors.torch.load_file(safe_weights_file, device="cpu")
                 else:
                     check_torch_load_is_safe()
@@ -2828,9 +2848,7 @@ class Trainer:
                 logger.warning(f"Could not load adapter model, make sure to have PEFT >= {MIN_PEFT_VERSION} installed")
         else:
             # We load the sharded checkpoint
-            load_result = load_sharded_checkpoint(
-                model, resume_from_checkpoint, strict=is_sagemaker_mp_enabled(), prefer_safe=self.args.save_safetensors
-            )
+            load_result = load_sharded_checkpoint(model, resume_from_checkpoint, strict=is_sagemaker_mp_enabled())
             if not is_sagemaker_mp_enabled():
                 self._issue_warnings_after_load(load_result)
 
@@ -2913,7 +2931,7 @@ class Trainer:
                         has_been_loaded = False
                 else:
                     # We load the model state dict on the CPU to avoid an OOM error.
-                    if self.args.save_safetensors and os.path.isfile(best_safe_model_path):
+                    if os.path.isfile(best_safe_model_path):
                         state_dict = safetensors.torch.load_file(best_safe_model_path, device="cpu")
                     else:
                         check_torch_load_is_safe()
@@ -3932,6 +3950,9 @@ class Trainer:
         # Both standard transformer models and Liger-patched models handle shift_labels correctly,
         # so we can directly use the computed loss from the model output.
         # See: https://huggingface.co/docs/accelerate/en/concept_guides/sequence_parallelism
+        if "labels" not in inputs and "shift_labels" in inputs:
+            # DeepSpeed SP Dataloader removes "labels" but we need it, otherwise, we won't compute the loss.
+            inputs["labels"] = inputs["shift_labels"]
         outputs = model(**inputs)
         loss = outputs.loss
 
@@ -4007,7 +4028,16 @@ class Trainer:
                     self._save(output_dir, state_dict=state_dict)
         elif self.is_deepspeed_enabled:
             try:
-                state_dict = self.accelerator.get_state_dict(self.deepspeed)
+                accept_exclude_frozen_parameters = "exclude_frozen_parameters" in set(
+                    inspect.signature(self.model_wrapped.save_checkpoint).parameters.keys()
+                )
+                zero3_sharding = self.deepspeed.config.get("zero_optimization", {}).get("stage", None) == 3
+                if accept_exclude_frozen_parameters and _is_peft_model(self.model) and zero3_sharding:
+                    # When using PEFT with DeepSpeed ZeRO Stage 3,
+                    # we do not need to load the frozen parameters
+                    state_dict = self.deepspeed._zero3_consolidated_16bit_state_dict(exclude_frozen_parameters=True)
+                else:
+                    state_dict = self.accelerator.get_state_dict(self.deepspeed)
                 if self.args.should_save:
                     self._save(output_dir, state_dict=state_dict)
             except ValueError:
@@ -4067,12 +4097,7 @@ class Trainer:
                 model = model.module.module
                 unwrapped_model = self.accelerator.unwrap_model(model)
                 if isinstance(unwrapped_model, supported_classes):
-                    unwrapped_model.save_pretrained(
-                        output_dir,
-                        state_dict=full_state_dict,
-                        save_function=xm.save,
-                        safe_serialization=self.args.save_safetensors,
-                    )
+                    unwrapped_model.save_pretrained(output_dir, state_dict=full_state_dict)
                 else:
                     logger.info("Trainer.model is not a `PreTrainedModel`, only saving its state dict.")
                     xm.save(full_state_dict, os.path.join(output_dir, WEIGHTS_NAME))
@@ -4082,8 +4107,6 @@ class Trainer:
                     output_dir,
                     is_main_process=self.args.should_save,
                     state_dict=xm._maybe_convert_to_cpu(model.state_dict()),
-                    save_function=xm.save,
-                    safe_serialization=self.args.save_safetensors,
                 )
             else:
                 logger.info("Trainer.model is not a `PreTrainedModel`, only saving its state dict.")
@@ -4093,8 +4116,6 @@ class Trainer:
             model.save_pretrained(
                 output_dir,
                 is_main_process=self.args.should_save,
-                save_function=xm.save,
-                safe_serialization=self.args.save_safetensors,
                 state_dict=xm._maybe_convert_to_cpu(model.state_dict()),
             )
         if self.processing_class is not None and self.args.should_save:
@@ -4115,20 +4136,15 @@ class Trainer:
 
             if isinstance(self.accelerator.unwrap_model(self.model, keep_torch_compile=False), supported_classes):
                 self.accelerator.unwrap_model(self.model, keep_torch_compile=False).save_pretrained(
-                    output_dir, state_dict=state_dict, safe_serialization=self.args.save_safetensors
+                    output_dir, state_dict=state_dict
                 )
             else:
                 logger.info("Trainer.model is not a `PreTrainedModel`, only saving its state dict.")
-                if self.args.save_safetensors:
-                    safetensors.torch.save_file(
-                        state_dict, os.path.join(output_dir, SAFE_WEIGHTS_NAME), metadata={"format": "pt"}
-                    )
-                else:
-                    torch.save(state_dict, os.path.join(output_dir, WEIGHTS_NAME))
+                safetensors.torch.save_file(
+                    state_dict, os.path.join(output_dir, SAFE_WEIGHTS_NAME), metadata={"format": "pt"}
+                )
         else:
-            self.model.save_pretrained(
-                output_dir, state_dict=state_dict, safe_serialization=self.args.save_safetensors
-            )
+            self.model.save_pretrained(output_dir, state_dict=state_dict)
 
         if self.processing_class is not None:
             self.processing_class.save_pretrained(output_dir)
@@ -4827,6 +4843,7 @@ class Trainer:
         if not self.args.hub_always_push and self.push_in_progress is not None and not self.push_in_progress.is_done():
             return
 
+        self.callback_handler.on_push_begin(self.args, self.state, self.control)
         output_dir = self.args.output_dir
         # To avoid a new synchronization of all model weights, we just copy the file from the checkpoint folder
         modeling_files = [CONFIG_NAME, GENERATION_CONFIG_NAME, WEIGHTS_NAME, SAFE_WEIGHTS_NAME]
@@ -4921,6 +4938,8 @@ class Trainer:
             The URL of the repository where the model was pushed if `blocking=False`, or a `Future` object tracking the
             progress of the commit if `blocking=True`.
         """
+        self.callback_handler.on_push_begin(self.args, self.state, self.control)
+
         model_name = kwargs.pop("model_name", None)
         if model_name is None and self.args.should_save:
             if self.args.hub_model_id is None:
@@ -5074,14 +5093,14 @@ class Trainer:
         self.is_tp_enabled = False
         if getattr(self.model, "tp_size", None) is not None and self.model.tp_size > 1:
             self.is_tp_enabled = True
-            if self.args.parallelism_config is not None:
-                if is_accelerate_available("1.10.1"):
-                    if self.args.parallelism_config is not None:
+            if self.args.parallelism_config is None:
+                if is_accelerate_available("1.12.0"):
+                    if self.args.parallelism_config is None:
                         from accelerate import ParallelismConfig
 
                         args["parallelism_config"] = ParallelismConfig(tp_size=self.model.tp_size)
                 else:
-                    raise ValueError("Requires accelerate>1.10.1 to use Tensor Parallelism.")
+                    raise ValueError("Requires accelerate>1.12.0 to use Tensor Parallelism.")
 
         if is_accelerate_available("1.2.0"):
             # it we don't have the correct version, we will rely on env var instead that were set in TrainingArguments

transformers/trainer_callback.py

@@ -420,6 +420,11 @@ class TrainerCallback:
         Event called after a prediction step.
         """
 
+    def on_push_begin(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
+        """
+        Event called before pushing the model to the hub, at the beginning of Trainer.push_to_hub and Trainer._push_from_checkpoint.
+        """
+
 
 class CallbackHandler(TrainerCallback):
     """Internal class that just calls the list of callbacks in order."""
@@ -532,6 +537,9 @@ class CallbackHandler(TrainerCallback):
     def on_prediction_step(self, args: TrainingArguments, state: TrainerState, control: TrainerControl):
         return self.call_event("on_prediction_step", args, state, control)
 
+    def on_push_begin(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
+        return self.call_event("on_push_begin", args, state, control, **kwargs)
+
     def call_event(self, event, args, state, control, **kwargs):
         for callback in self.callbacks:
             result = getattr(callback, event)(