More control over DPO training process

CHANGELOG.md

@@ -39,7 +39,8 @@ to communicate reliably.
 - Agent-level hyper-parameter to enable quantisation for self-hosted models.
 - Using Liger kernel in DPO training for increased speed and lower memory usage.
 - Now logging prompt and completion token lengths for self-hosted models.
-- Now attempting several times to start the vLLM server. 
+- Now attempting several times to start the vLLM server.
+- Allowed setting DPO number of epochs.
 
 
 ### Changed
@@ -52,6 +53,7 @@ to communicate reliably.
   run`), which has allowed displaying more log messages.
 - Sanitising model name a little more when creating training job name, which helps
   Hugging Face identify the base model more easily.
+- Allowed setting DPO training logging period, and set it to 1 by default.
 
 
 ### Fixed

nip/code_validation/agents.py

@@ -1752,12 +1752,17 @@ async def _make_fine_tune_api_call_openai(
         file_id = uploaded_file.id
 
         method_key = {"type": method}
+        shared_hyperparameters = {"n_epochs": self.num_epochs}
         if method == "dpo":
             if self.shared_agent_params.dpo_beta is None:
                 beta = "auto"
             else:
                 beta = self.shared_agent_params.dpo_beta
-            method_key["dpo"] = {"hyperparameters": {"beta": beta}}
+            method_key["dpo"] = {
+                "hyperparameters": {"beta": beta, **shared_hyperparameters}
+            }
+        elif method == "supervised":
+            method_key["supervised"] = {"hyperparameters": shared_hyperparameters}
 
         # Create the fine-tune job
         while True:
@@ -1825,6 +1830,7 @@ async def _make_fine_tune_api_call_lm_server(
             dpo_config=LmDpoTrainingConfig(
                 beta=self.shared_agent_params.dpo_beta,
                 learning_rate=self.rl_learning_rate,
+                num_train_epochs=self.num_epochs,
             ),
             training_lora_config=LmLoraAdapterConfig(
                 r=self.shared_agent_params.lora_rank,

nip/language_model_server/trainers/dpo.py

@@ -219,15 +219,19 @@ def train(config: LmTrainingConfig, dataset: Dataset, job_id: str, new_model_nam
         gradient_checkpointing=config.gradient_checkpointing,
         per_device_train_batch_size=config.per_device_train_batch_size,
         use_liger_kernel=config.use_liger_kernel,
+        logging_steps=config.logging_steps,
         seed=config.seed,
+        log_level="info",
     )
 
     ignore_training_lora_config = False
 
     if not is_model_peft(config.model_name):
+        logger.info(f"Loading model {config.model_name!r}...")
         model = AutoModelForCausalLM.from_pretrained(config.model_name)
 
     else:
+        logger.info(f"Model {config.model_name!r} is a LoRA model. Loading...")
         model_lora_config = LoraConfig.from_pretrained(config.model_name)
 
         # When reusing the LoRA adapter, make sure the model's LoRA configuration is
@@ -275,6 +279,8 @@ def train(config: LmTrainingConfig, dataset: Dataset, job_id: str, new_model_nam
             # already LoRA-adapted and the trainer will train the existing adapter.
             ignore_training_lora_config = True
 
+    logger.info(f"Model {config.model_name!r} loaded successfully.")
+
     if ignore_training_lora_config or config.training_lora_config is None:
         training_lora_config = None
     else:
@@ -337,9 +343,12 @@ def main():
 
 
 if __name__ == "__main__":
-    logging.basicConfig(
-        level=logging.INFO,
-        format="[%(asctime)s %(levelname)s] %(message)s",
-        datefmt="%x %X",
+    root_logger = logging.getLogger()
+    root_logger.setLevel(logging.INFO)
+    handler = logging.StreamHandler()
+    formatter = logging.Formatter(
+        "[%(asctime)s %(levelname)s] %(message)s", datefmt="%x %X"
     )
+    handler.setFormatter(formatter)
+    root_logger.addHandler(handler)
     main()

nip/language_model_server/types.py

@@ -135,6 +135,9 @@ class LmDpoTrainingConfig(BaseModel):
     learning_rate: float
     """The learning rate for the DPO training."""
 
+    num_train_epochs: int
+    """The number of epochs to train the model for."""
+
     max_prompt_length: int | None = None
     """The maximum length of the prompt sequence."""
 
@@ -219,6 +222,9 @@ class LmTrainingConfig(BaseModel):
     This can improve training speed and lower memory usage.
     """
 
+    logging_steps: int = 1
+    """The period (in steps) at which to log training metrics."""
+
 
 class CreateTrainingJobRequest(BaseModel):
     """A request to create a new training job."""

nip/parameters/agents.py

@@ -504,6 +504,9 @@ class PureTextAgentParameters(AgentParameters):
         The quantization method to use for model inference. This is only relevant when
         using a self-hosted model. It controls how the model weights are quantized to
         reduce memory usage, at the cost of some accuracy.
+    num_epochs : int | None
+        The number of epochs to train the model for. If ``None``, we use the global
+        ``rl.num_epochs`` parameter, which is set in the RL trainer parameters.
     dpo_beta : float | None
         The beta parameter for to use when training the model with DPO. This is a float
         between 0 and 2, which controls how strictly the new model will adhere to its
@@ -577,6 +580,8 @@ class PureTextAgentParameters(AgentParameters):
 
     quantization: VllmQuantization = "none"
 
+    num_epochs: Optional[int] = None
+
     dpo_beta: Optional[float] = None
 
     use_lora: bool = True

nip/scenario_base/agents.py

@@ -505,6 +505,8 @@ class SharedAgentParams:
 
         fine_tune_from_scratch: bool
 
+        num_epochs: Optional[int]
+
         dpo_beta: Optional[float]
 
         use_lora: bool
@@ -544,6 +546,18 @@ def rl_learning_rate(self) -> float:
 
         return learning_rate
 
+    @property
+    def num_epochs(self) -> int:
+        """The number of epochs to train the model for.
+
+        We first look at the ``num_epochs`` shared agent parameter. If this is not set,
+        we use the global hyperparameter ``rl.num_epochs``.
+        """
+        if self.shared_agent_params.num_epochs is not None:
+            return self.shared_agent_params.num_epochs
+        else:
+            return self.hyper_params.rl.num_epochs
+
     @property
     def lora_alpha(self) -> float:
         """The computed LoRA alpha value for the group.

scripts/config/cv_experiment/single_experiment.json5

@@ -20,6 +20,7 @@
             "vllm_server_port": 8000,
             "temperature": null,
             "dpo_beta": 0.1,
+            "num_epochs": 3,
             "lora_rank": 32,
             "lora_alpha_scale": 1.0,
             "lora_dropout": 0.05,
@@ -41,6 +42,7 @@
             "vllm_server_port": 8000,
             "temperature": 1.0,
             "dpo_beta": 0.1,
+            "num_epochs": 3,
             "lora_rank": 32,
             "lora_alpha_scale": 1.0,
             "lora_dropout": 0.05,

scripts/cv_experiment.py

@@ -87,6 +87,7 @@ def _construct_params(combo: dict, cmd_args: Namespace) -> HyperParameters:
             lora_alpha_scale=combo["verifier.lora_alpha_scale"],
             lora_dropout=combo["verifier.lora_dropout"],
             quantization=combo["verifier.quantization"],
+            num_epochs=combo["verifier.num_epochs"],
         ),
     )
 
@@ -110,6 +111,7 @@ def _construct_params(combo: dict, cmd_args: Namespace) -> HyperParameters:
         lora_alpha_scale=combo["prover.lora_alpha_scale"],
         lora_dropout=combo["prover.lora_dropout"],
         quantization=combo["prover.quantization"],
+        num_epochs=combo["prover.num_epochs"],
     )
 
     if combo["provers_share_model"]: