🔀 Add MergeModelCallBack

docs/source/callbacks.mdx

@@ -14,4 +14,8 @@
 
 ## LogCompletionsCallback
 
-[[autodoc]] LogCompletionsCallback
+[[autodoc]] LogCompletionsCallback
+
+## MergeModelCallback
+
+[[autodoc]] MergeModelCallback

setup.py

@@ -85,9 +85,10 @@
     # Windows support is partially supported with DeepSpeed https://github.com/microsoft/DeepSpeed/tree/master#windows
     "deepspeed": ["deepspeed>=0.14.4; sys_platform != 'win32'"],
     "diffusers": ["diffusers>=0.18.0"],
+    "judges": ["openai>=1.23.2", "llm-blender>=0.0.2"],
     # liger-kernel depends on triton, which is only available on Linux https://github.com/triton-lang/triton#compatibility
     "liger": ["liger-kernel>=0.4.0; sys_platform != 'win32'"],
-    "judges": ["openai>=1.23.2", "llm-blender>=0.0.2"],
+    "mergekit": ["mergekit>=0.0.5.1"],
     "peft": ["peft>=0.8.0"],
     "quantization": ["bitsandbytes"],
     "scikit": ["scikit-learn"],

tests/test_callbacks.py

@@ -20,9 +20,12 @@
 from datasets import load_dataset
 from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig, Trainer, TrainingArguments
 from transformers.testing_utils import require_peft, require_wandb
+from transformers.trainer_utils import get_last_checkpoint
 from transformers.utils import is_peft_available
 
-from trl import BasePairwiseJudge, LogCompletionsCallback, WinRateCallback
+from tests.testing_utils import require_mergekit
+from trl import BasePairwiseJudge, DPOConfig, DPOTrainer, LogCompletionsCallback, MergeModelCallback, WinRateCallback
+from trl.mergekit_utils import MergeConfig
 
 
 if is_peft_available():
@@ -266,3 +269,64 @@ def test_basic(self):
 
             # Check that the prompt is in the log
             self.assertIn(self.dataset["test"][0]["prompt"], completions["data"][0])
+
+
+@require_mergekit
+class MergeModelCallbackTester(unittest.TestCase):
+    def setUp(self):
+        self.model = AutoModelForCausalLM.from_pretrained("trl-internal-testing/tiny-random-LlamaForCausalLM")
+        self.tokenizer = AutoTokenizer.from_pretrained("trl-internal-testing/tiny-random-LlamaForCausalLM")
+        self.dataset = load_dataset("trl-internal-testing/zen", "standard_preference", split="train")
+
+    def test_callback(self):
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            training_args = DPOConfig(
+                output_dir=tmp_dir,
+                num_train_epochs=1,
+                report_to="none",
+                save_strategy="steps",
+                save_steps=1,
+            )
+            config = MergeConfig()
+            merge_callback = MergeModelCallback(config)
+            trainer = DPOTrainer(
+                model=self.model,
+                args=training_args,
+                train_dataset=self.dataset,
+                tokenizer=self.tokenizer,
+                callbacks=[merge_callback],
+            )
+            trainer.train()
+            last_checkpoint = get_last_checkpoint(tmp_dir)
+            merged_path = os.path.join(last_checkpoint, "merged")
+            self.assertTrue(os.path.isdir(merged_path), "Merged folder does not exist in the last checkpoint.")
+
+    def test_every_checkpoint(self):
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            training_args = DPOConfig(
+                output_dir=tmp_dir,
+                num_train_epochs=1,
+                report_to="none",
+                save_strategy="steps",
+                save_steps=1,
+            )
+            config = MergeConfig()
+            merge_callback = MergeModelCallback(config, merge_at_every_checkpoint=True)
+            trainer = DPOTrainer(
+                model=self.model,
+                args=training_args,
+                train_dataset=self.dataset,
+                tokenizer=self.tokenizer,
+                callbacks=[merge_callback],
+            )
+            trainer.train()
+
+            checkpoints = sorted(
+                [os.path.join(tmp_dir, cp) for cp in os.listdir(tmp_dir) if cp.startswith("checkpoint-")]
+            )
+
+            for checkpoint in checkpoints:
+                merged_path = os.path.join(checkpoint, "merged")
+                self.assertTrue(
+                    os.path.isdir(merged_path), f"Merged folder does not exist in checkpoint {checkpoint}."
+                )

tests/testing_utils.py

@@ -17,6 +17,7 @@
 from transformers import is_bitsandbytes_available, is_sklearn_available, is_wandb_available
 
 from trl import BaseBinaryJudge, BasePairwiseJudge, is_diffusers_available, is_llm_blender_available
+from trl.import_utils import is_mergekit_available
 
 
 # transformers.testing_utils contains a require_bitsandbytes function, but relies on pytest markers which we don't use
@@ -35,6 +36,20 @@ def require_diffusers(test_case):
     return unittest.skipUnless(is_diffusers_available(), "test requires diffusers")(test_case)
 
 
+def require_llm_blender(test_case):
+    """
+    Decorator marking a test that requires llm-blender. Skips the test if llm-blender is not available.
+    """
+    return unittest.skipUnless(is_llm_blender_available(), "test requires llm-blender")(test_case)
+
+
+def require_mergekit(test_case):
+    """
+    Decorator marking a test that requires mergekit. Skips the test if mergekit is not available.
+    """
+    return unittest.skipUnless(is_mergekit_available(), "test requires mergekit")(test_case)
+
+
 def require_no_wandb(test_case):
     """
     Decorator marking a test that requires no wandb. Skips the test if wandb is available.
@@ -49,13 +64,6 @@ def require_sklearn(test_case):
     return unittest.skipUnless(is_sklearn_available(), "test requires sklearn")(test_case)
 
 
-def require_llm_blender(test_case):
-    """
-    Decorator marking a test that requires llm-blender. Skips the test if llm-blender is not available.
-    """
-    return unittest.skipUnless(is_llm_blender_available(), "test requires llm-blender")(test_case)
-
-
 class RandomBinaryJudge(BaseBinaryJudge):
     """
     Random binary judge, for testing purposes.

trl/__init__.py

@@ -70,6 +70,7 @@
         "KTOConfig",
         "KTOTrainer",
         "LogCompletionsCallback",
+        "MergeModelCallback",
         "ModelConfig",
         "NashMDConfig",
         "NashMDTrainer",
@@ -91,7 +92,7 @@
         "XPOConfig",
         "XPOTrainer",
     ],
-    "trainer.callbacks": ["RichProgressCallback", "SyncRefModelCallback"],
+    "trainer.callbacks": ["MergeModelCallback", "RichProgressCallback", "SyncRefModelCallback"],
     "trainer.utils": ["get_kbit_device_map", "get_peft_config", "get_quantization_config"],
     "utils": ["ScriptArguments"],
 }
@@ -159,6 +160,7 @@
         KTOConfig,
         KTOTrainer,
         LogCompletionsCallback,
+        MergeModelCallback,
         ModelConfig,
         NashMDConfig,
         NashMDTrainer,

trl/import_utils.py

@@ -24,6 +24,7 @@
 _deepspeed_available = _is_package_available("deepspeed")
 _diffusers_available = _is_package_available("diffusers")
 _llm_blender_available = _is_package_available("llm_blender")
+_mergekit_available = _is_package_available("mergekit")
 _rich_available = _is_package_available("rich")
 _unsloth_available = _is_package_available("unsloth")
 
@@ -40,6 +41,10 @@ def is_llm_blender_available() -> bool:
     return _llm_blender_available
 
 
+def is_mergekit_available() -> bool:
+    return _mergekit_available
+
+
 def is_rich_available() -> bool:
     return _rich_available