Add checkpointing support for Tinker SkyRL backend #992
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Add full checkpoint save/load functionality to SkyRLTrainBackend: - save_checkpoint(): Saves model + optimizer + scheduler state as uncompressed tar - load_checkpoint(): Restores full training state from tar checkpoint - save_sampler_checkpoint(): Exports model weights in HuggingFace format for inference Implementation wraps WorkerDispatch checkpoint methods and handles tar packaging. Uses uncompressed tar to avoid 5-10 minute gzip bottleneck on 6-7GB FSDP checkpoints. Co-Authored-By: Claude Sonnet 4.5 <[email protected]>
Co-Authored-By: Claude Sonnet 4.5 <[email protected]>
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request introduces checkpointing capabilities to the SkyRLTrainBackend, enabling training state to be saved and restored via tar archives. However, a high-severity path traversal vulnerability was identified in the checkpoint loading logic due to the unsafe use of tarfile.extractall(). Additionally, the worker initialization logic uses trust_remote_code=True when loading model configurations, which poses a significant security risk if untrusted model identifiers are processed. Beyond security, there's also a minor race condition and some code duplication that could benefit from refactoring for improved maintainability.
Co-Authored-By: Claude Sonnet 4.5 <[email protected]>
Address PR review feedback: 1. Security: Add filter='data' to tarfile.extractall() to prevent path traversal (TarSlip) attacks where malicious archives could write outside the temp directory 2. Refactor: Extract duplicate validation logic into _validate_model_state() helper method (used by all 3 checkpoint methods) 3. Remove redundant os.path.exists() check that creates TOCTOU race condition - tarfile.open() already raises FileNotFoundError 4. Refactor: Extract common tar creation logic into _create_tar_from_directory() helper method to reduce duplication Co-Authored-By: Claude Sonnet 4.5 <[email protected]>
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Summary
Implements checkpointing for the SkyRL-Train backend in Tinker, enabling periodic checkpoint saves and training resumption.
Changes
Adds three checkpoint methods to
SkyRLTrainBackend:save_checkpoint()- Saves full training state (model + optimizer + scheduler) as tar archiveload_checkpoint()- Restores full training state from checkpointsave_sampler_checkpoint()- Exports model-only checkpoint in HuggingFace formatImplementation
WorkerDispatch.save_checkpoint()andload_checkpoint()methods