tocmo0nlord/axolotl
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

Add MLFlow run name option in config (#1961)

Browse Source 
Co-authored-by: Adam Hazell <adam.hazell@mindfoundry.ai>
This commit is contained in:
Adam Hazell
2024-10-11 18:33:06 +01:00
committed by GitHub
parent e73b8dff8d
commit 922db77521
3 changed files with 8 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
docs/config.qmd
View File

@@ -266,6 +266,7 @@ wandb_log_model: # "checkpoint" to log model to wandb Artifacts every `save_step
# mlflow configuration if you're using it # mlflow configuration if you're using it
mlflow_tracking_uri: # URI to mlflow mlflow_tracking_uri: # URI to mlflow
mlflow_experiment_name: # Your experiment name mlflow_experiment_name: # Your experiment name
mlflow_run_name: # Your run name
hf_mlflow_log_artifacts: # set to true to copy each saved checkpoint on each save to mlflow artifact registry hf_mlflow_log_artifacts: # set to true to copy each saved checkpoint on each save to mlflow artifact registry
# Comet configuration if you're using it # Comet configuration if you're using it

9
src/axolotl/core/trainer_builder.py
View File

@@ -1445,9 +1445,12 @@ class HFCausalTrainerBuilder(TrainerBuilderBase):
report_to.append("comet_ml") report_to.append("comet_ml")
training_arguments_kwargs["report_to"] = report_to training_arguments_kwargs["report_to"] = report_to
training_arguments_kwargs["run_name"] = ( if self.cfg.use_wandb:
self.cfg.wandb_name if self.cfg.use_wandb else None training_arguments_kwargs["run_name"] = self.cfg.wandb_name
) elif self.cfg.use_mlflow:
training_arguments_kwargs["run_name"] = self.cfg.mlflow_run_name
else:
training_arguments_kwargs["run_name"] = None
training_arguments_kwargs["optim"] = ( training_arguments_kwargs["optim"] = (
self.cfg.optimizer if self.cfg.optimizer else "adamw_hf" self.cfg.optimizer if self.cfg.optimizer else "adamw_hf"
) )

1
src/axolotl/utils/config/models/input/v0_4_1/__init__.py
View File

@@ -447,6 +447,7 @@ class MLFlowConfig(BaseModel):
use_mlflow: Optional[bool] = None use_mlflow: Optional[bool] = None
mlflow_tracking_uri: Optional[str] = None mlflow_tracking_uri: Optional[str] = None
mlflow_experiment_name: Optional[str] = None mlflow_experiment_name: Optional[str] = None
mlflow_run_name: Optional[str] = None
hf_mlflow_log_artifacts: Optional[bool] = None hf_mlflow_log_artifacts: Optional[bool] = None