update reference model with policy within the spin and dpo trainers t…

…hemselves, to ready for arbitrary ordering of fine tuning steps

README.md

@@ -98,9 +98,11 @@ sft_dataset = create_mock_dataset(100, lambda: (torch.randint(0, 256, (256,)), t
 spin_trainer = SPINTrainer(
     transformer,
     max_seq_len = 16,
-    sft_dataset = sft_dataset,
-    spin_λ = 0.1,
-    checkpoint_every = 100
+    train_sft_dataset = sft_dataset,
+    checkpoint_every = 100,
+    spin_kwargs = dict(
+        λ = 0.1,
+    ),
 )
 
 spin_trainer()

self_rewarding_lm_pytorch/dpo.py

@@ -434,7 +434,10 @@ def forward(
         self,
         train_self_reward_dataset: Optional[Dataset] = None
     ):
-        self.early_stopper.clear_early_checkpoint_folder()
+        self.model.update_reference_model_with_policy()
+
+        if exists(self.early_stopper):
+            self.early_stopper.clear_early_checkpoint_folder()
 
         train_dataloader = self.train_dataloader
 

self_rewarding_lm_pytorch/self_rewarding_lm_pytorch.py

@@ -803,8 +803,6 @@ def forward(
 
             spin_trainer()
 
-            self.spin.update_reference_model_with_policy()
-
             self.save(f'spin.{spin_cycle}.ckpt.pt', overwrite = overwrite_checkpoints)
 
 
@@ -816,8 +814,6 @@ def forward(
 
             dpo_trainer(dpo_dataset_from_self_reward)
 
-            self.dpo.update_reference_model_with_policy()
-
             self.save(f'self-reward.{iterate_num}.ckpt.pt', overwrite = overwrite_checkpoints)
 
         self.print(f'self-reward training done')

self_rewarding_lm_pytorch/spin.py

@@ -292,6 +292,8 @@ def forward(self, overwrite_checkpoints: bool = True):
         Algorithm 1 - https://arxiv.org/abs/2401.01335v1
         """
 
+        self.model.update_reference_model_with_policy()
+
         self.steps = 0
         self.model.train()
 

setup.py

@@ -3,7 +3,7 @@
 setup(
   name = 'self-rewarding-lm-pytorch',
   packages = find_packages(exclude=[]),
-  version = '0.0.38',
+  version = '0.0.39',
   license='MIT',
   description = 'Self Rewarding LM - Pytorch',
   author = 'Phil Wang',