Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

MergeModelCallBack #2282

Draft
wants to merge 4 commits into
base: main
Choose a base branch
from

Conversation

August-murr
Copy link
Contributor

@August-murr August-murr commented Oct 25, 2024

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).
  • Did you read the contributor guideline,
    Pull Request section?
  • Was this discussed/approved via a GitHub issue? Please add a link
    to it if that's the case.
  • 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

@@ -103,6 +103,7 @@
"deepspeed": ["deepspeed>=0.14.4"],
"quantization": ["bitsandbytes<=0.41.1"],
"llm_judge": ["openai>=1.23.2", "llm-blender>=0.0.2"],
"mergekit": ["mergekit @ git+https://github.com/arcee-ai/mergekit.git"],
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is not ideal and very messy, but it was the only way I could install mergekit from source since "pip" installed an older version

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

Add model merging callback
1 participant