fix num_labels= 1 test fail

src/axolotl/core/builders/causal.py

@@ -421,6 +421,13 @@ def build(self, total_num_steps):
             trainer_kwargs["dataset_tags"] = [
                 d["path"] for d in self.cfg.datasets if not Path(d["path"]).is_dir()
             ]
+        # TRL's RewardTrainer validates num_labels=1 on pre-loaded models; ensure the
+        # config reflects this regardless of how the model was instantiated.
+        if (
+            self.cfg.reward_model
+            and getattr(self.model.config, "num_labels", None) != 1
+        ):
+            self.model.config.num_labels = 1
         trainer = trainer_cls(
             model=self.model,
             train_dataset=self.train_dataset,

src/axolotl/utils/schemas/validation.py

@@ -254,6 +254,23 @@ def hint_reward_model_pad(cls, data):
                 data["pad_to_sequence_len"] = True
         return data
 
+    @model_validator(mode="before")
+    @classmethod
+    def set_reward_model_defaults(cls, data):
+        if data.get("reward_model"):
+            if data.get("num_labels") is None:
+                data["num_labels"] = 1
+            if not (data.get("type_of_model") or data.get("model_type")):
+                data["model_type"] = "AutoModelForSequenceClassification"
+
+        if data.get("process_reward_model"):
+            if data.get("num_labels") is None:
+                data["num_labels"] = 2
+            if not (data.get("type_of_model") or data.get("model_type")):
+                data["model_type"] = "AutoModelForTokenClassification"
+
+        return data
+
     @model_validator(mode="before")
     @classmethod
     def check_gas_bsz(cls, data):

tests/core/test_builders.py

@@ -536,7 +536,7 @@ def test_training_arguments(self, sft_cfg, model, tokenizer):
         "cfg_string",
         [
             "sft_cfg",
-            # "rm_cfg",  # TODO fix for num_labels = 2 vs 1
+            "rm_cfg",
             "prm_cfg",
         ],
     )

tests/patched/test_validation.py

@@ -277,6 +277,34 @@ def test_model_type_remap(self, minimal_cfg):
         new_cfg = validate_config(cfg)
         assert new_cfg.type_of_model == "AutoModelForCausalLM"
 
+    def test_reward_model_defaults(self, minimal_cfg):
+        cfg = (
+            DictDefault(
+                {
+                    "reward_model": True,
+                }
+            )
+            | minimal_cfg
+        )
+
+        new_cfg = validate_config(cfg)
+        assert new_cfg.num_labels == 1
+        assert new_cfg.type_of_model == "AutoModelForSequenceClassification"
+
+    def test_process_reward_model_defaults(self, minimal_cfg):
+        cfg = (
+            DictDefault(
+                {
+                    "process_reward_model": True,
+                }
+            )
+            | minimal_cfg
+        )
+
+        new_cfg = validate_config(cfg)
+        assert new_cfg.num_labels == 2
+        assert new_cfg.type_of_model == "AutoModelForTokenClassification"
+
     def test_model_revision_remap(self, minimal_cfg):
         cfg = (
             DictDefault(