trl
trl copied to clipboard
MergeModelCallBack
What does this PR do?
Fixes #2241
Since the focus was on replicating the checkpoint merging methods from the paper, I have covered only Linear, TIES, SLERP, and DARE-TIES merging methods. These were the ones used in the paper and were primarily tested with the DPOTrainer.
Please provide feedback on any issues and suggest improvements regarding the structure of the files in this PR, as well as the documentation, to enhance clarity.
Here's an example of usage:
pip install trl[mergekit]
since it's an optional dependency
from trl.mergekit_utils import MergeConfig
Trainer = DPOTrainer(...)
config = MergeConfig("ties")
config.target_model_path = "path_to_target_model" #if none is provided, the reference model will be used
config.policy_model_weight = 0.7 #optional
config.target_model_weight = 0.3 #optional
merge_callback = MergeModelCallBack(config,push_to_hub=True,merge_at_every_checkpoint=True)
Trainer.add_callback(callback)
Before submitting
- [ ] 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, Pull Request section?
- [x] Was this discussed/approved via a GitHub issue? Please add a link to it if that's the case.
- [x] Did you make sure to update the documentation with your changes? Here are the documentation guidelines.
- [ ] Did you write any new necessary tests?
Who can review?
@lewtun @qgallouedec