feat: update to use wandb_name and allow id

examples/qwen/lora.yml

@@ -31,7 +31,7 @@ lora_fan_in_fan_out:
 wandb_project:
 wandb_entity:
 wandb_watch:
-wandb_run_id:
+wandb_name:
 wandb_log_model:
 
 gradient_accumulation_steps: 4

examples/qwen/qlora.yml

@@ -31,7 +31,7 @@ lora_fan_in_fan_out:
 wandb_project:
 wandb_entity:
 wandb_watch:
-wandb_run_id:
+wandb_name:
 wandb_log_model:
 
 gradient_accumulation_steps: 4

src/axolotl/core/trainer_builder.py

@@ -643,7 +643,7 @@ def build(self, total_num_steps):
         training_arguments_kwargs["group_by_length"] = self.cfg.group_by_length
         training_arguments_kwargs["report_to"] = "wandb" if self.cfg.use_wandb else None
         training_arguments_kwargs["run_name"] = (
-            self.cfg.wandb_run_id if self.cfg.use_wandb else None
+            self.cfg.wandb_name if self.cfg.use_wandb else None
         )
         training_arguments_kwargs["optim"] = (
             self.cfg.optimizer if self.cfg.optimizer else "adamw_hf"

src/axolotl/utils/config.py

@@ -397,12 +397,11 @@ def validate_config(cfg):
             "Gradient checkpointing is broken for Qwen models for transformers>=4.35.0, except main branch."
         )
 
-    if cfg.wandb_run_id and len(cfg.wandb_run_id) > 0:
+    if cfg.wandb_run_id and not cfg.wandb_name:
         cfg.wandb_name = cfg.wandb_run_id
-        cfg.wandb_run_id = None
 
         LOG.warning(
-            "wandb_run_id is not recommended anymore. Please use wandb_name instead."
+            "wandb_run_id sets the ID of the run. If you would like to set the name, please use wandb_name instead."
         )
 
     # TODO

tests/test_validation.py

@@ -682,7 +682,13 @@ def test_warmup_step_no_conflict(self):
 
         validate_config(cfg)
 
-    def test_wandb_rename_run_id_to_name(self):
+
+class ValidationWandbTest(ValidationTest):
+    """
+    Validation test for wandb
+    """
+
+    def test_wandb_set_run_id_to_name(self):
         cfg = DictDefault(
             {
                 "wandb_run_id": "foo",
@@ -692,12 +698,12 @@ def test_wandb_rename_run_id_to_name(self):
         with self._caplog.at_level(logging.WARNING):
             validate_config(cfg)
             assert any(
-                "wandb_run_id is not recommended anymore. Please use wandb_name instead."
+                "wandb_run_id sets the ID of the run. If you would like to set the name, please use wandb_name instead."
                 in record.message
                 for record in self._caplog.records
             )
 
-            assert cfg.wandb_name == "foo" and cfg.wandb_run_id is None
+            assert cfg.wandb_name == "foo" and cfg.wandb_run_id == "foo"
 
         cfg = DictDefault(
             {
@@ -707,11 +713,14 @@ def test_wandb_rename_run_id_to_name(self):
 
         validate_config(cfg)
 
+        assert cfg.wandb_name == "foo" and cfg.wandb_run_id is None
+
     def test_wandb_sets_env(self):
         cfg = DictDefault(
             {
                 "wandb_project": "foo",
                 "wandb_name": "bar",
+                "wandb_run_id": "bat",
                 "wandb_entity": "baz",
                 "wandb_mode": "online",
                 "wandb_watch": "false",
@@ -725,6 +734,7 @@ def test_wandb_sets_env(self):
 
         assert os.environ.get("WANDB_PROJECT", "") == "foo"
         assert os.environ.get("WANDB_NAME", "") == "bar"
+        assert os.environ.get("WANDB_RUN_ID", "") == "bat"
         assert os.environ.get("WANDB_ENTITY", "") == "baz"
         assert os.environ.get("WANDB_MODE", "") == "online"
         assert os.environ.get("WANDB_WATCH", "") == "false"