diff --git a/CHANGELOG.md b/CHANGELOG.md index 94dc693..778a510 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/nip/code_validation/agents.py b/nip/code_validation/agents.py index 5c50c58..d18a5cc 100644 --- a/nip/code_validation/agents.py +++ b/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, diff --git a/nip/language_model_server/trainers/dpo.py b/nip/language_model_server/trainers/dpo.py index 4d84839..dc283a5 100644 --- a/nip/language_model_server/trainers/dpo.py +++ b/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() diff --git a/nip/language_model_server/types.py b/nip/language_model_server/types.py index d2e2da8..5de5a50 100644 --- a/nip/language_model_server/types.py +++ b/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.""" diff --git a/nip/parameters/agents.py b/nip/parameters/agents.py index 9edf7f0..4df6a73 100644 --- a/nip/parameters/agents.py +++ b/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 diff --git a/nip/scenario_base/agents.py b/nip/scenario_base/agents.py index d76a9cc..ab518dd 100644 --- a/nip/scenario_base/agents.py +++ b/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. diff --git a/scripts/config/cv_experiment/single_experiment.json5 b/scripts/config/cv_experiment/single_experiment.json5 index 8c4a2b3..e198cd0 100644 --- a/scripts/config/cv_experiment/single_experiment.json5 +++ b/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, diff --git a/scripts/cv_experiment.py b/scripts/cv_experiment.py index f1995ac..fc81cf3 100644 --- a/scripts/cv_experiment.py +++ b/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"]: