From 5dba5c82a80279331b96b65b15e8707234bbcd98 Mon Sep 17 00:00:00 2001 From: Wing Lian Date: Fri, 25 Apr 2025 21:10:54 -0400 Subject: [PATCH] fix support for wandb run_name for rl trainers (#2566) [skip ci] * fix support for wandb run_name for rl trainers * prefer to use wandb random names for run_name --- src/axolotl/core/trainer_builder.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/src/axolotl/core/trainer_builder.py b/src/axolotl/core/trainer_builder.py index 44f8c5d2b..970b02075 100755 --- a/src/axolotl/core/trainer_builder.py +++ b/src/axolotl/core/trainer_builder.py @@ -1048,6 +1048,9 @@ class HFRLTrainerBuilder(TrainerBuilderBase): if self.cfg.rpo_alpha is not None: training_args_kwargs["rpo_alpha"] = self.cfg.rpo_alpha + if self.cfg.use_wandb: + training_args_kwargs["run_name"] = self.cfg.wandb_name + training_args_cls = None blocklist_args_kwargs = [] if self.cfg.rl == "simpo": @@ -1118,6 +1121,12 @@ class HFRLTrainerBuilder(TrainerBuilderBase): **training_args_kwargs, ) + # unset run_name so wandb sets up experiment names + if self.cfg.use_wandb and training_args.run_name == training_args.output_dir: + training_args.run_name = ( # pylint: disable=attribute-defined-outside-init + None + ) + return training_args def build(self, total_num_steps):