feat(trainer): Just-in-time (JIT) asynchronous checkpointing using SIGTERM signals
Author: efazal
Pull request: yes

# What does this PR do?

This PR adds Just-In-Time (JIT) asynchronous checkpointing to the Trainer API, enabling graceful handling of training interruptions in preemptible workloads. When enabled via `enable_jit_checkpoint=True`, the trainer
  automatically saves progress upon receiving a SIGTERM signal, preventing loss of training progress in shared cluster environments (e.g., Kueue).

## Motivation

Modern ML training increasingly relies on orchestration platforms like Kubeflow and workload managers like Kueue in shared Kubernetes clusters. These platforms require the ability to:

1. **Pause and Resume Jobs**: Dynamically reallocate resources by pausing long-running training jobs and resuming them later without losing progress
2. **Kueue Preemption**: Support preemptible workloads where lower-priority jobs can be terminated (via SIGTERM) to make room for higher-priority workloads
3. 
4. **Dynamic Scaling**: Enable training jobs to scale down/up based on cluster resource availability

Currently, if a training job is terminated before the next periodic checkpoint, all progress since the last checkpoint is lost. This wastes significant compute resources and training time.

## Implementation

  **New Components:**
    - `trainer_jit_checkpoint.py`: Core implementation with `CheckpointManager` and `JITCheckpointCallback`
    - Registers SIGTERM signal handler to intercept termination signals
    - Implements 3-second kill-wait period to distinguish SIGTERM from SIGKILL
    - Thread-based checkpoint execution to avoid blocking training loop

  **Modified Components:**
  - `trainer.py`: Automatically adds JITCheckpointCallback when `enable_jit_checkpoint=True`
  - `training_args.py`: New `enable_jit_checkpoint` boolean flag (configurable via `ENABLE_JIT_CHECKPOINT` environment variable)

**Environment Variable Support Rationale:**
 The new environment variables (`ENABLE_JIT_CHECKPOINT`, `SAVE_STRATEGY`, `SAVE_STEPS`, `SAVE_TOTAL_LIMIT`, `OUTPUT_DIR`, `RESUME_FROM_CHECKPOINT`) allow orchestration platforms like Kubeflow to automatically configure checkpointing behavior
  without requiring users to modify their training scripts. This enables:
  - Centralized checkpoint configuration through cluster operators
  - Dynamic checkpoint policies based on workload type
  - Seamless integration with existing training code

  **Key Features:**
  - Graceful training termination after checkpoint completion
  - Environment variable support for easier configuration in cluster environments
  - Compatible with existing checkpoint management (save_steps, save_strategy, etc.)

  ## Testing

  Comprehensive unit tests added in `tests/trainer/test_trainer_jit_checkpoint.py` covering:
  - CheckpointManager initialization and configuration
  - SIGTERM signal handler flow and duplicate signal handling
  - Callback integration with Trainer lifecycle hooks

Fixes #38961

## Who can review?

@SunMarc - This PR adds emergency checkpointing functionality to the Trainer class

## Before submitting
- [x] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [x] Was this discussed/approved via a Github issue or the [issue](https://github.com/huggingface/transformers/issues/38961)? Please add a link
      to it if that's the case.
- [x] Updated Trainer args doc.
- [x] Did you write any new necessary tests?


Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

<!-- Your PR will be replied to more quickly if you can figure out the right person to tag with @

 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- text models: @ArthurZucker @Cyrilvallez
- vision models: @yonigozlan @molbap
- audio models: @eustlb @ebezzam @vasqu
- multimodal models: @zucchini-nlp
- graph models: @clefourrier

Library:

- generate: @zucchini-nlp (visual-language models) or @gante (all others)
- continuous batching: @remi-or @ArthurZucker @McPatate
- pipelines: @Rocketknight1
- tokenizers: @ArthurZucker and @itazap
- trainer: @SunMarc
- attention: @vasqu @ArthurZucker @CyrilVallez
- model loading (from pretrained, etc): @CyrilVallez
- distributed: @3outeille @ArthurZucker
- CIs: @ydshieh

Integrations:

- ray/raytune: @richardliaw, @amogkam
- Big Model Inference: @SunMarc
- quantization: @SunMarc @MekkCyber
- kernels: @MekkCyber @drbh
- peft: @BenjaminBossan @githubnemo

Devices/Backends:

- AMD ROCm: @ivarflakstad
- Intel XPU: @IlyasMoutawwakil
- Ascend NPU: @ivarflakstad 

Documentation: @stevhliu

Research projects are not maintained and should be taken as is.

 -->


Comments:
@SunMarc @Rocketknight1 could you please review this PR. Thank you. 

@SunMarc this is definitely a pure or nearly-pure code agent PR, so leaving it up to you to decide if you want it in `Trainer`; feel free to close if you don't

@SunMarc  AI was indeed used to to write tests, for this PR description, and so on. I'm wondering if there is a policy against using AI whatsover in the transformers project. I've manually tested code with multiple different models to verify JIT checkpointing works on sigterm signals for over a month now. I'm willing take reviews and fix them. I've tried implement this without modifying the trainer code too much by keeping it separate and it's flag enabled. I would appreciate it if you could review it please. Thank you. 



@sgugger @stas00 could you please review this PR. Thank you. 

@efazal, Sylvain and I are no longer with HF.

In general the PR looks great and +1 (for whatever it counts) to add this functionality. Let's discuss the implementation.

Question: I had similar functionality implemented in other frameworks I worked on, but I was using a check between trainer steps.

Otherwise how do you know that when you trigger the weight saving you're not in the middle of a weight update and you could end up with inconsistent weights?

The way I did it is via a flag - sigterm -> raise a flag -> trainer checks the flag as soon as step is finished -> goes into checkpoint saving -> exit - no threads are needed and the weights are consistent.

I see no benefit for doing async saving here also because the trainer can't continue anyway, so the last step's results will be lost anyway if sigterm caught it in the middle.

Thanks for your review @sfc-gh-sbekman 🙏, and apologies I didn't realize you were not with HF anymore.

You made a great point about when a checkpoint gets triggered during weight update (optimizer pass). To avoid checkpointing during optimizer pass, I'm creating a new cuda stream, and I wait for the main cuda stream operations to complete, and after weight updates completed, checkpointing happens async in a new cuda stream. 

Again, you rightly said that async checkpointing is unnecessary given that the training is going to stop anyway. I just wanted to bring this approach forward to possibly extend async checkpointing to existing periodic checkpointing, since its currently synchronous and blocking GPU while periodic checkpoint takes place making trainings roughly 30% more expensive (in periodic checkpoint enabled trainings). But I think this approach is no longer needed since Pytorch addressed this problem recently and added support to [asynchronously checkpoint](https://pytorch.org/blog/reducing-checkpointing-times/) and also added support to save model in [safetensors](https://pytorch.org/blog/huggingface-safetensors-support-in-pytorch-distributed-checkpointing/?ajs_aid=4a55dd7f-1ce3-4743-9079-3833e1790fdd) format. So I will update the PR to remove async checkpointing.

I like the approach of doing this at the end of step as you suggested. In the current implementation, the `_sigterm_handler()` immediately triggers checkpointing. Instead, I will mark the flag `checkpoint_requested` to `True` and as part of the callbacks `on_step_end()`, `on_step_begin`, and `on_pre_optimizer_step()` I will trigger checkpoint based on the flag. I'm planning to trigger in 3 different callbacks, since we are targeting graceful shutdowns, and in Kubernetes the default is 30 seconds graceful shutdown period so we need to checkpoint asap and the time it takes to checkpoint varies based on the model being trained/fine-tuned.  The 30 second default can we overridden, but I doubt most users will consciously do this. I hope this makes sense. 

I think your async checkpointing as a replacement for blocking normal operation checkpointing would be a fantastic addition - perhaps in a separate PR?

re: k8s - there should be a config setting in your job definition which can tell the scheduler how much grace time your job needs to shutdown. So the bigger the model the more time you usually need, and 30 secs will not be enough in most cases since you need to have enough time for 1 full step + saving. I recommend documenting how this can be done. 

SLURM jobs definitely have this feature - it's covered in detail here: https://github.com/stas00/ml-engineering/tree/11111d63c92edf3c10c25515851083de75cac629/training/fault-tolerance#dealing-with-forced-job-preemption Let's add the k8s graceful setup at that section? We can work together on it if you PR in.

Additionally, when using such panic event-based saving techniques it's important to also ensure the checkpoint was fully saved as it could get the carpet yanked under 1sec too soon and the user won't be the wiser why their can't resume. Typically I use some sort of flag file which I write when the checkpoint has finished writing, so when resuming it'd first check that the flag file is there. or this can be done in reverse, touching "checkpoint-is-incomplete.txt" file and then deleting it when the checkpoint writing has finished, which perhaps would integrate better, be more obvious to the user inspecting its files.


Yes, absolutely. I would contribute non blocking checkpointing separately after this PR goes through. 

I think that makes sense. For this feature to work consistently with larger models, the default 30 sec time period wouldn't suffice. And yes, there is a config which can propagate the value of graceful shutdown period to the pods `terminationGracePeriodSeconds` field. In most cloud-native based trainers like Kubeflow for example allows to do this done via CRD job pod `template.spec`. I shall document this in the trainer arguments file for the arg `enable_jit_checkpoint`. Might take inspiration from your ml-engineering repo, to let the user know how to decide the amount of graceful shutdown time period. 

I'm interested in raising a PR to your ml-engineering repo with regards to k8s graceful shown. I shall ping you when I do so.

Also, another valid problem you addressed. I did in fact run into incomplete checkpoint files when trying to resume training  when I tested with larger models. I addressed this problem by performing a check within my training function to check if the lastest checkpoint contains the files: model.safetensors, optimizer.pt, scheduler.pt, and rng.states.pth. To confirm if the checkpoint is valid, otherwise go back to a earlier checkpoint file to resume or start over. But I find your idea brilliant to do it in reverse and I shall do that. And perhaps later, as part of a separate PR we could leverage this file to determine if the training can be resumed with latest checkpoint, or defer to earlier checkpoint if available. 

Thanks a mil for the great ideas. I shall implement them soon. 




Glad to hear you found my suggestions useful, Esa

> by performing a check within my training function to check if the lastest checkpoint contains the files: model.safetensors, optimizer.pt, scheduler.pt, and rng.states.pth.

this is insufficient since the file could exist but be incomplete if the writing got aborted mid-way. But as you said you will switch to a sentinel control file approach.

You might like the other solution I came up with when training BLOOM-176B - the kill and save switches: 
- https://github.com/stas00/ml-engineering/tree/11111d63c92edf3c10c25515851083de75cac629/training/fault-tolerance#kill-switch
- https://github.com/stas00/ml-engineering/tree/11111d63c92edf3c10c25515851083de75cac629/training/fault-tolerance#save-switch

if you have those in place, you don't even need a flag on sigterm - now all you need is to drop a save switch file, but either way it works. The variable flag is likely to be a tad faster.

Looking forward to your PRs wrt k8s recipes, I'm a newbie there so need more time to gain expertise before I can share recipes with others. I like SLURM so much better. The k8s' level of complexity is insane compared to SLURM. You need one sysadmin to manage a SLURM cluster, you need a whole team to deal with k8s.

> this is insufficient since the file could exist but be incomplete if the writing got aborted mid-way. But as you said you will switch to a sentinel control file approach

That's true, definitely not a sustainable approach. I've implemented it as you suggested and it's a great approach.

Yes, I did read up the fault-tolerance doc, it was certainly a very interesting read. Yeah, I suppose we are trying to achieve the kill switch approach via a flag and yes, it's a bit quicker since we can trigger checkpoint in many callbacks. 

I've updated the PR now based on your comments so far. I still need to manually verify this which I will do next and confirm. 

Again, thanks for all the info you've provided thus far @sfc-gh-sbekman . Please have a look at the PR and let me know if you have any concerns in particular. I've updated the docs to mention about grace time period in both in K8s and slurp orchestrators.  

should these be more explicit and have the trainer prefix? `SAVE_STRATEGY, SAVE_STEPS, SAVE_TOTAL_LIMIT` 

In the context of the code it's obvious what `save_` implies, but as an env var this is too ambiguous. Perhaps `SAVE_CHECKPOINT_*`?

Oddly I don't see the trainer using some consistent prefix for its env vars, may be I missed something - ideally those should start with `HF_TRAINER_` or some such not to pollute the env, but ask the current maintainers to which one to use. e.g. HF Accelerate env vars all start with `ACCELERATE_`.

Good catch on the env prefix concern. Upon looking around in the repo, I could see `HF_` prefix used for `HF_HOME` and `HF_TOKEN`, and I figured since there are other trainers like TRL which is also based on HF Transformer trainer, I updated the checkpoint variables also likewise i.e.  `HF_ENABLE_JIT_CHECKPOINT`, so even TRL users can benefit from this feature and it works across HF ecosystem I suppose. I hope it's a good idea. Maybe @SunMarc can correct me here perhaps. 

I did a bunch of manual testing by finetuning a llama3.2 8B instruct model on multinode setting. Works pretty well as before, expect now its synchronous. I also tested JIT checkpointing by coinciding it with periodic checkpoint to see how it behaves if JIT checkpoint was triggered before and during periodic. Following were the result:

- JIT checkpoint triggered right **before** periodic checkpoint: Since we set the `control.should_save = False` when JIT checkpoint is triggered, so it doesn't periodic checkpoint, and job is terminated which is what we expect.
- JIT checkpoint triggered **during** periodic checkpoint: The periodic check completes and the job is terminated. But I suspect that if the graceful period was exessive, the training loop continues after the periodic checkpoint and the callback would trigger JIT checkpoint as well which may possibly lead to an incomplete checkpoint.  WDYT? @sfc-gh-sbekman 

@SunMarc could you please add other reviewers to this PR please. Would also appreciate your review. Thank you! 

@ArthurZucker could you please review this PR and add anybody else who might be required please. Thank you!

@SunMarc @ArthurZucker gentle ping here for a review please. 🙏

Sorry for the delay, I was off for quite some time. I will review it shortly 

Thank you @SunMarc. I look forward to your reviews! 

cc @qgallouedec if you can have a look also. might be interesting for trl 

Hi @SunMarc @qgallouedec. I was wondering if you were able to review this PR. A gentle ping here as a reminder. Thank you!  🙏

> Thanks for the PR and sorry for the long wait ! Left a couple of comments

Thanks for the review comments @SunMarc. I will address them right away. 

@bot /style 

Style bot fixed some files and pushed the changes.

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_41723). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

Thanks a mil for your review @SunMarc. I will definitely add to the docs, do you want me to add to this PR? I suppose you mean a separate PR for it. 

Yeah we can merge this first !