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

transformers/tokenization_utils_tokenizers.py

@@ -20,18 +20,18 @@ 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.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 +41,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 +91,157 @@ 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":
+                vocab = list(map(tuple, 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")
 
+        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 +252,7 @@ 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:
             raise ValueError(
                 "Couldn't instantiate the backend tokenizer from one of: \n"
                 "(1) a `tokenizers` library serialization file, \n"
@@ -139,11 +260,11 @@ 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."
             )
+        if fast_tokenizer is not None:
+            self._tokenizer = fast_tokenizer
 
-        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
 
@@ -168,9 +289,17 @@ class TokenizersBackend(PreTrainedTokenizerBase):
         # Set backend to "tokenizers" if not already set
         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
@@ -228,6 +357,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 +408,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 +455,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:
         """
@@ -1132,7 +1193,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
@@ -2338,6 +2348,8 @@ 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
+            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 +2440,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 +2460,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 +2801,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 +2841,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 +2924,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()
@@ -4067,12 +4078,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 +4088,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 +4097,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 +4117,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)
@@ -5074,14 +5071,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_jit_checkpoint.py

@@ -0,0 +1,126 @@
+import os
+import signal
+import threading
+from typing import Optional
+
+from .trainer_callback import TrainerCallback
+from .trainer_utils import PREFIX_CHECKPOINT_DIR
+from .utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+
+class CheckpointManager:
+    def __init__(self, trainer, kill_wait: int = 3):
+        """
+        Initialize the CheckpointManager for Just-In-Time checkpoint handling.
+
+        Args:
+            trainer: The Trainer instance that will be used to save checkpoints when SIGTERM is received.
+            kill_wait (`int`, *optional*, defaults to 3): Grace period to distinguish between SIGTERM and SIGKILL.
+        """
+        self.trainer = trainer
+        self.is_checkpoint_requested = False
+        self._original_sigterm_handler = None
+        self.kill_wait = kill_wait
+
+    def setup_signal_handler(self):
+        self._original_sigterm_handler = signal.signal(signal.SIGTERM, self._sigterm_handler)
+        logger.info("JIT checkpoint signal handler registered for SIGTERM")
+
+    def _sigterm_handler(self, signum, frame):
+        if self.is_checkpoint_requested:
+            return
+
+        logger.info(f"SIGTERM received, will request JIT checkpoint after {self.kill_wait}s")
+        threading.Timer(self.kill_wait, self._enable_checkpoint).start()
+
+    def _enable_checkpoint(self):
+        logger.info("Kill wait period elapsed, requesting checkpoint")
+        self.is_checkpoint_requested = True
+
+    def execute_jit_checkpoint(self):
+        try:
+            # Set checkpoint flag to False to avoid multiple checkpoints getting triggered by other callbacks
+            self.is_checkpoint_requested = False
+
+            logger.info("Starting JIT checkpointing...")
+            current_step = self.trainer.state.global_step
+            logger.info(f"Saving JIT checkpoint at step {current_step}")
+
+            output_dir = self.trainer._get_output_dir(trial=None)
+            checkpoint_folder = f"{PREFIX_CHECKPOINT_DIR}-{current_step}"
+            checkpoint_path = os.path.join(output_dir, checkpoint_folder)
+
+            # Create checkpoint directory
+            os.makedirs(checkpoint_path, exist_ok=True)
+
+            # Create a sentinel file to indicate checkpointing is in progress
+            sentinel_file = os.path.join(output_dir, checkpoint_folder, "checkpoint-is-incomplete.txt")
+            with open(sentinel_file, "w") as f:
+                f.write(f"Checkpoint started at step {current_step} and in progress...")
+            logger.info(f"Created checkpoint progress sentinel marker file: {sentinel_file}")
+
+            # Invoke the trainer's checkpoint method directly
+            self.trainer._save_checkpoint(self.trainer.model, trial=None)
+
+            # Remove sentinel file upon successful checkpointing
+            if os.path.exists(sentinel_file):
+                os.remove(sentinel_file)
+                logger.info("Sentinel marker file removed")
+
+            logger.info("Immediate JIT checkpoint completed successfully")
+
+        except Exception as e:
+            logger.error(f"Failed to save JIT checkpoint: {e}")
+            raise
+
+
+class JITCheckpointCallback(TrainerCallback):
+    """
+    Callback for Just-In-Time checkpointing on SIGTERM signals.
+
+    When SIGTERM is received, the checkpoint manager sets `is_checkpoint_requested=True`.
+    The callbacks detect this flag and set `control.should_training_stop=True`, which signals
+    the Trainer's training loop to exit gracefully after saving the checkpoint.
+    """
+
+    def __init__(self):
+        self.trainer = None
+        self.jit_manager: Optional[CheckpointManager] = None
+
+    def set_trainer(self, trainer):
+        self.trainer = trainer
+        if trainer.args.enable_jit_checkpoint:
+            self.jit_manager = CheckpointManager(trainer=trainer)
+            self.jit_manager.setup_signal_handler()
+            logger.info("JIT checkpointing enabled")
+
+    def on_pre_optimizer_step(self, args, state, control, **kwargs):
+        if self.jit_manager and self.jit_manager.is_checkpoint_requested:
+            control.should_training_stop = True
+            self.jit_manager.execute_jit_checkpoint()
+
+    def on_step_begin(self, args, state, control, **kwargs):
+        if self.jit_manager and self.jit_manager.is_checkpoint_requested:
+            control.should_training_stop = True
+            self.jit_manager.execute_jit_checkpoint()
+
+    def on_step_end(self, args, state, control, **kwargs):
+        if self.jit_manager and self.jit_manager.is_checkpoint_requested:
+            control.should_save = False
+            control.should_training_stop = True
+            self.jit_manager.execute_jit_checkpoint()
+
+    def on_epoch_end(self, args, state, control, **kwargs):
+        if self.jit_manager and self.jit_manager.is_checkpoint_requested:
+            control.should_save = False
+            control.should_training_stop = True
+            self.jit_manager.execute_jit_checkpoint()
+
+    def on_train_end(self, args, state, control, **kwargs):
+        #  Restore original SIGTERM handler
+        if self.jit_manager and self.jit_manager._original_sigterm_handler is not None:
+            signal.signal(signal.SIGTERM, self.jit_manager._original_sigterm_handler)
+            logger.info("Restored original SIGTERM handler after training completion")

transformers/trainer_utils.py

@@ -924,7 +924,7 @@ def load_sharded_checkpoint(model, folder, strict=True, prefer_safe=True):
     shard_files = list(set(index["weight_map"].values()))
 
     # If strict=True, error before loading any of the state dicts.
-    # TODO: Here, update the weigth map with the config.dynamic_weight_conversion
+    # TODO: Here, update the weight map with the config.dynamic_weight_conversion
     loaded_keys = index["weight_map"].keys()
     model_keys = model.state_dict().keys()
     missing_keys = [key for key in model_keys if key not in loaded_keys]