Hugging Face's logo

TRL documentation

MergeModelCallback

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v1.4.0).
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

MergeModelCallback

class trl.experimental.merge_model_callback.MergeModelCallback

< >

( merge_config: MergeConfig | None = None merge_at_every_checkpoint: bool = False push_to_hub: bool = False )

Parameters

  • merge_config (experimental.merge_model_callback.MergeConfig, optional) — Configuration used for the merging process. If not provided, the default MergeConfig is used.
  • merge_at_every_checkpoint (bool, optional, defaults to False) — Whether to merge the model at every checkpoint.
  • push_to_hub (bool, optional, defaults to False) — Whether to push the merged model to the Hub after merging.

A TrainerCallback that merges the policy model (the model being trained) with another model based on a merge configuration.

Example:

from trl.experimental.merge_model_callback import MergeConfig, MergeModelCallback

config = MergeConfig()
merge_callback = MergeModelCallback(config)
trainer = DPOTrainer(..., callbacks=[merge_callback])
Update on GitHub

←KTO MiniLLM→