feat: add ORPO support

Author: eghuzefa

tests/sft/dpo/dpo_trainer_test.py

@@ -122,11 +122,11 @@ def test_dpo_trainer(
       rejected_mask,
       use_ref_model,
   ):
-    model = tc.ToyTransformer(rngs=nnx.Rngs(0))
+    model = tc.ToyTransformer(config=tc.ModelConfig(), rngs=nnx.Rngs(0))
     original_variables = jax.tree.map(jnp.copy, nnx.state(model, nnx.Param))
     ref_model = None
     if use_ref_model:
-      ref_model = tc.ToyTransformer(rngs=nnx.Rngs(0))
+      ref_model = tc.ToyTransformer(config=tc.ModelConfig(), rngs=nnx.Rngs(0))
     dpo_config = dpo_lib.DPOTrainingConfig(
         eval_every_n_steps=5,
         max_steps=10,
@@ -169,11 +169,15 @@ def test_dpo_trainer(
         "log_probs/rejected",
     ]:
       self.assertLen(
-          dpo_trainer.metrics_logger.get_metric_history(metric_name, "train"),
+          dpo_trainer.metrics_logger.get_metric_history(
+              "", metric_name, "train"
+          ),
           dpo_trainer._train_steps,
       )
       self.assertLen(
-          dpo_trainer.metrics_logger.get_metric_history(metric_name, "eval"),
+          dpo_trainer.metrics_logger.get_metric_history(
+              "", metric_name, "eval"
+          ),
           3,
       )
 
@@ -201,11 +205,13 @@ def test_dpo_trainer(
   def test_dpo_trainer_with_string_inputs(self, train_ds):
     tokenizer = tc.MockVocab()
     model = tc.ToyTransformer(
-        rngs=nnx.Rngs(0), vocab_size=tokenizer.GetPieceSize()
+        config=tc.ModelConfig(vocab_size=tokenizer.GetPieceSize()),
+        rngs=nnx.Rngs(0),
     )
     original_variables = jax.tree.map(jnp.copy, nnx.state(model, nnx.Param))
     ref_model = tc.ToyTransformer(
-        rngs=nnx.Rngs(0), vocab_size=tokenizer.GetPieceSize()
+        config=tc.ModelConfig(vocab_size=tokenizer.GetPieceSize()),
+        rngs=nnx.Rngs(0),
     )
     original_ref_variables = jax.tree.map(
         jnp.copy, nnx.state(ref_model, nnx.Param)
@@ -240,13 +246,15 @@ def test_dpo_trainer_with_string_inputs(self, train_ds):
         "rewards/accuracy",
     ]:
       self.assertLen(
-          dpo_trainer.metrics_logger.get_metric_history(metric_name, "train"),
+          dpo_trainer.metrics_logger.get_metric_history(
+              "", metric_name, "train"
+          ),
           dpo_trainer._train_steps,
       )
 
   def test_dpo_loss_fn(self):
     np.random.seed(0)
-    model = tc.ToyTransformer(rngs=nnx.Rngs(0))
+    model = tc.ToyTransformer(config=tc.ModelConfig(), rngs=nnx.Rngs(0))
     per_token_logps = np.random.normal(0, 5, size=(8, 4))
     ref_per_token_logps = np.random.normal(0, 5, size=(8, 4)).sum(axis=-1)
     train_example = dpo_lib.TrainExample(
@@ -262,20 +270,26 @@ def test_dpo_loss_fn(self):
     with mock.patch.object(
         common, "get_per_token_logps", return_value=jnp.array(per_token_logps)
     ):
-      loss, _ = dpo_lib.dpo_loss_fn(model, train_example, 0.1, 0)
+      loss, _ = dpo_lib.dpo_loss_fn(
+          model, train_example, beta=0.1, label_smoothing=0
+      )
       np.testing.assert_allclose(loss, 0.753059, atol=1e-5)
 
-      loss, _ = dpo_lib.dpo_loss_fn(model, train_example, 0.1, 0.3)
+      loss, _ = dpo_lib.dpo_loss_fn(
+          model, train_example, beta=0.1, label_smoothing=0.3
+      )
       np.testing.assert_allclose(loss, 0.925447, atol=1e-5)
 
   def test_dpo_prepare_inputs_for_strings(self):
     tokenizer = tc.MockVocab()
 
     model = tc.ToyTransformer(
-        rngs=nnx.Rngs(0), vocab_size=tokenizer.GetPieceSize()
+        config=tc.ModelConfig(vocab_size=tokenizer.GetPieceSize()),
+        rngs=nnx.Rngs(0),
     )
     ref_model = tc.ToyTransformer(
-        rngs=nnx.Rngs(0), vocab_size=tokenizer.GetPieceSize()
+        config=tc.ModelConfig(vocab_size=tokenizer.GetPieceSize()),
+        rngs=nnx.Rngs(0),
     )
     dpo_trainer = dpo_lib.DPOTrainer(
         model=model,
@@ -328,8 +342,8 @@ def test_dpo_prepare_inputs_for_strings(self):
     self.assertEqual(out.logits_to_keep, 3)
 
   def test_dpo_prepare_inputs(self):
-    model = tc.ToyTransformer(rngs=nnx.Rngs(0))
-    ref_model = tc.ToyTransformer(rngs=nnx.Rngs(0))
+    model = tc.ToyTransformer(config=tc.ModelConfig(), rngs=nnx.Rngs(0))
+    ref_model = tc.ToyTransformer(config=tc.ModelConfig(), rngs=nnx.Rngs(0))
     dpo_trainer = dpo_lib.DPOTrainer(
         model=model,
         ref_model=ref_model,