feat(trainer): emergency checkpointing on crashes & SIGTERM/SIGINT#39140
feat(trainer): emergency checkpointing on crashes & SIGTERM/SIGINT#39140AyushSharma173 wants to merge 4 commits into
Conversation
* new flag * save model/optimizer/scheduler/RNG on crash or signal * tests:
|
|
||
| try: | ||
| if args.push_to_hub: | ||
| try: |
There was a problem hiding this comment.
Wrapping the whole training loop into a try/expect would make the use of debugger quite challenging, no?
There was a problem hiding this comment.
Good point. I’ve adjusted the implementation so the try except now only surrounds the call to inner_training_loop (via a small _run_inner_training_loop_with_emergency_ckpt helper).
Everything that happens earlier in train is executed outside a try so debuggers will still break exactly at the offending line. The emergency checkpoint logic is therefore triggered only for failures that occur after training has actually started, which is where we have state worth saving.
There was a problem hiding this comment.
Thanks for narrowing the scope of the try/epxect, although I think that the remark remains for everything inside the training loop.
Not saying that's it's a major issue for this feature, but I feel like it could make the development work a lot more annoying
…ethod for emergency checkpointing logic.
|
cc @SunMarc as well |
|
CC: @SunMarc for your attn here |
|
hey @AyushSharma173 I don't think saving it in |
|
@SunMarc A gentle ping here for review. Also @AyushSharma173 LMK if you don't have the bandwidth to continue as I have some time now and can it pick up. |
|
@yaswanth19 Hey sure. You can pickup this issue and try to address the concerns raised so far. I dont have time for next week, so feel free. |
|
Hi @AyushSharma173, This is great work! I've been working on a similar feature and only now came across your PR. I took a proactive approach as opposed to your reactive approach to solve this and named it Just-in-Time (JIT) checkpointing. And I also wanted to ensure the trainer code was not disrupted a lot. For that reason I took a modular approach. I'm also leveraging SIGTERM signals for this implementation, and the key difference would we that I'm doing the checkpoint asynchronously. Please have a look at my approach implemented in my fork and let me know your thoughts as I would love to contribute these changes as part of this PR or even a fresh PR if you think that would be wise: |
|
what about just: try:
trainer.train()
except:
trainer.save_model() |
|
Hi @AyushSharma173 @qgallouedec @SunMarc, I was wondering if I could pickup this issue and continue working on the feature. In my previous comment I've added a link to my fork with the approach I took to implement emergency checkpoint. It addresses all the review comments here so far. Please have a look and let me know if I can go ahead and submit a new revision to this PR. Thank you. |
What does this PR do?
Adds failure-safe training to the 🤗
Trainer.enable_emergency_checkpointinTrainingArgumentsFalse(opt-in)Trainer.trainSIGTERM/SIGINT(e.g. job pre-emption, Ctrl-C)<output_dir>/checkpoint-emergencyand contains model weights, optimizer, scheduler, scaler, RNG state
plus a minimal
trainer_state.json.TrainingArguments(..., resume_from_checkpoint=".../checkpoint-emergency").Motivation
Issue #38961 requests a robust way to avoid losing progress when a run is killed mid-epoch or by OOM/infra errors.
Current work-arounds (very frequent
save_stepsor a user-leveltry/except) are slow or brittle.This PR makes the feature one-line (
enable_emergency_checkpoint=True) and cost-free when disabled.Implementation details
training_args.py– new dataclass fieldenable_emergency_checkpoint.trainer.pyatexithandlers when the flag is ontrain()loop intry/exceptand calls_common_emergency_save_common_emergency_saveis idempotent & rank-safe (_emergency_save_running / _completedflags).tests/trainer/test_emergency_ckpt.pyFixes # (issue)
Closes #38961.
Before submitting
make fixup && make qualitylocally.Who can review?
Trainer reviewers: @zach-huggingface, @SunMarc, @qgallouedec