From 7019509daacd51f33c66d81aec9f4804409be90d Mon Sep 17 00:00:00 2001 From: Morgan McGuire Date: Sat, 12 Aug 2023 17:17:11 +0100 Subject: [PATCH] Add wandb_entity to wandb options, update example configs, update README (#361) * Update wandb_entity and add wandb descriptions * add wandb to config section * remove trailing whitespace for pre-commit hook * remove trailing whitespace for pre-commit hook --------- Co-authored-by: Morgan McGuire Co-authored-by: Wing Lian --- README.md | 21 +++++++++++++++++---- examples/cerebras/qlora.yml | 1 + examples/falcon/config-7b-lora.yml | 1 + examples/falcon/config-7b-qlora.yml | 1 + examples/falcon/config-7b.yml | 1 + examples/gptj/qlora.yml | 1 + examples/gptq-lora-7b/config.yml | 1 + examples/jeopardy-bot/config.yml | 1 + examples/llama-2/lora.yml | 1 + examples/llama-2/qlora.yml | 1 + examples/mpt-7b/config.yml | 1 + examples/openllama-3b/config.yml | 1 + examples/openllama-3b/lora.yml | 1 + examples/openllama-3b/qlora.yml | 1 + examples/pythia-12b/config.yml | 1 + examples/pythia/lora.yml | 1 + examples/redpajama/config-3b.yml | 1 + examples/replit-3b/config-lora.yml | 1 + examples/xgen-7b/xgen-7b-8k-qlora.yml | 1 + src/axolotl/utils/wandb.py | 2 ++ 20 files changed, 37 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index 011b13903..e9f0aaefb 100644 --- a/README.md +++ b/README.md @@ -401,11 +401,12 @@ lora_out_dir: lora_fan_in_fan_out: false # wandb configuration if you're using it -wandb_mode: -wandb_project: +wandb_mode: # "offline" to save run metadata locally and not sync to the server, "disabled" to turn off wandb +wandb_project: # your wandb project name +wandb_entity: # a wandb Team name if using a Team wandb_watch: -wandb_run_id: -wandb_log_model: # 'checkpoint' +wandb_run_id: # set the name of your wandb run +wandb_log_model: # "checkpoint" to log model to wandb Artifacts every `save_steps` or "end" to log only at the end of training # where to save the finished model to output_dir: ./completed-model @@ -553,6 +554,18 @@ fsdp_config: - llama Deepspeed: append `ACCELERATE_USE_DEEPSPEED=true` in front of finetune command +##### Weights & Biases Logging + +- wandb options +```yaml +wandb_mode: +wandb_project: +wandb_entity: +wandb_watch: +wandb_run_id: +wandb_log_model: +``` + ### Inference Pass the appropriate flag to the train command: diff --git a/examples/cerebras/qlora.yml b/examples/cerebras/qlora.yml index 2f3700249..eec16193d 100644 --- a/examples/cerebras/qlora.yml +++ b/examples/cerebras/qlora.yml @@ -23,6 +23,7 @@ lora_target_modules: lora_target_linear: lora_fan_in_fan_out: wandb_project: +wandb_entity: wandb_watch: wandb_run_id: wandb_log_model: diff --git a/examples/falcon/config-7b-lora.yml b/examples/falcon/config-7b-lora.yml index 8aa585851..a5cbdc00d 100644 --- a/examples/falcon/config-7b-lora.yml +++ b/examples/falcon/config-7b-lora.yml @@ -24,6 +24,7 @@ lora_target_modules: lora_target_linear: true lora_fan_in_fan_out: wandb_project: +wandb_entity: wandb_watch: wandb_run_id: wandb_log_model: diff --git a/examples/falcon/config-7b-qlora.yml b/examples/falcon/config-7b-qlora.yml index 6168ff2d5..72b09b87d 100644 --- a/examples/falcon/config-7b-qlora.yml +++ b/examples/falcon/config-7b-qlora.yml @@ -38,6 +38,7 @@ lora_target_linear: true lora_fan_in_fan_out: wandb_project: +wandb_entity: wandb_watch: wandb_run_id: wandb_log_model: diff --git a/examples/falcon/config-7b.yml b/examples/falcon/config-7b.yml index b267566ce..46f4caff1 100644 --- a/examples/falcon/config-7b.yml +++ b/examples/falcon/config-7b.yml @@ -24,6 +24,7 @@ lora_target_modules: lora_target_linear: true lora_fan_in_fan_out: wandb_project: +wandb_entity: wandb_watch: wandb_run_id: wandb_log_model: diff --git a/examples/gptj/qlora.yml b/examples/gptj/qlora.yml index f2427f4d4..b41c946fc 100644 --- a/examples/gptj/qlora.yml +++ b/examples/gptj/qlora.yml @@ -20,6 +20,7 @@ lora_target_modules: lora_target_linear: true lora_fan_in_fan_out: wandb_project: +wandb_entity: wandb_watch: wandb_run_id: wandb_log_model: diff --git a/examples/gptq-lora-7b/config.yml b/examples/gptq-lora-7b/config.yml index 351a2bf97..d5fbe3f13 100644 --- a/examples/gptq-lora-7b/config.yml +++ b/examples/gptq-lora-7b/config.yml @@ -22,6 +22,7 @@ lora_target_modules: - v_proj lora_fan_in_fan_out: false wandb_project: llama-7b-lora-int4 +wandb_entity: wandb_watch: wandb_run_id: wandb_log_model: diff --git a/examples/jeopardy-bot/config.yml b/examples/jeopardy-bot/config.yml index b803c6074..1aaa1b557 100644 --- a/examples/jeopardy-bot/config.yml +++ b/examples/jeopardy-bot/config.yml @@ -18,6 +18,7 @@ lora_dropout: lora_target_modules: lora_fan_in_fan_out: false wandb_project: +wandb_entity: wandb_watch: wandb_run_id: wandb_log_model: diff --git a/examples/llama-2/lora.yml b/examples/llama-2/lora.yml index 494538fff..b7f985cbf 100644 --- a/examples/llama-2/lora.yml +++ b/examples/llama-2/lora.yml @@ -26,6 +26,7 @@ lora_target_linear: true lora_fan_in_fan_out: wandb_project: +wandb_entity: wandb_watch: wandb_run_id: wandb_log_model: diff --git a/examples/llama-2/qlora.yml b/examples/llama-2/qlora.yml index 64728ac3d..b6a1f69f2 100644 --- a/examples/llama-2/qlora.yml +++ b/examples/llama-2/qlora.yml @@ -27,6 +27,7 @@ lora_target_linear: true lora_fan_in_fan_out: wandb_project: +wandb_entity: wandb_watch: wandb_run_id: wandb_log_model: diff --git a/examples/mpt-7b/config.yml b/examples/mpt-7b/config.yml index f21ce7022..0efd28cee 100644 --- a/examples/mpt-7b/config.yml +++ b/examples/mpt-7b/config.yml @@ -20,6 +20,7 @@ lora_target_modules: - v_proj lora_fan_in_fan_out: false wandb_project: mpt-alpaca-7b +wandb_entity: wandb_watch: wandb_run_id: wandb_log_model: diff --git a/examples/openllama-3b/config.yml b/examples/openllama-3b/config.yml index 4372876eb..0d8144d6b 100644 --- a/examples/openllama-3b/config.yml +++ b/examples/openllama-3b/config.yml @@ -22,6 +22,7 @@ lora_target_modules: lora_target_linear: lora_fan_in_fan_out: wandb_project: +wandb_entity: wandb_watch: wandb_run_id: wandb_log_model: diff --git a/examples/openllama-3b/lora.yml b/examples/openllama-3b/lora.yml index d1f252455..acf0826c9 100644 --- a/examples/openllama-3b/lora.yml +++ b/examples/openllama-3b/lora.yml @@ -28,6 +28,7 @@ lora_target_modules: - o_proj lora_fan_in_fan_out: wandb_project: +wandb_entity: wandb_watch: wandb_run_id: wandb_log_model: diff --git a/examples/openllama-3b/qlora.yml b/examples/openllama-3b/qlora.yml index dcad9bbcd..d8c43df82 100644 --- a/examples/openllama-3b/qlora.yml +++ b/examples/openllama-3b/qlora.yml @@ -22,6 +22,7 @@ lora_target_modules: lora_target_linear: true lora_fan_in_fan_out: wandb_project: +wandb_entity: wandb_watch: wandb_run_id: wandb_log_model: diff --git a/examples/pythia-12b/config.yml b/examples/pythia-12b/config.yml index 3b3d91630..535e5cd37 100644 --- a/examples/pythia-12b/config.yml +++ b/examples/pythia-12b/config.yml @@ -23,6 +23,7 @@ lora_target_modules: lora_target_linear: true lora_fan_in_fan_out: true # pythia/GPTNeoX lora specific wandb_project: +wandb_entity: wandb_watch: wandb_run_id: wandb_log_model: diff --git a/examples/pythia/lora.yml b/examples/pythia/lora.yml index e2b28f218..dfe7ef846 100644 --- a/examples/pythia/lora.yml +++ b/examples/pythia/lora.yml @@ -17,6 +17,7 @@ lora_target_modules: lora_target_linear: lora_fan_in_fan_out: true # pythia/GPTNeoX lora specific wandb_project: +wandb_entity: wandb_watch: wandb_run_id: wandb_log_model: diff --git a/examples/redpajama/config-3b.yml b/examples/redpajama/config-3b.yml index 869c0883e..10297d860 100644 --- a/examples/redpajama/config-3b.yml +++ b/examples/redpajama/config-3b.yml @@ -21,6 +21,7 @@ lora_target_modules: - v_proj lora_fan_in_fan_out: false wandb_project: redpajama-alpaca-3b +wandb_entity: wandb_watch: wandb_run_id: wandb_log_model: diff --git a/examples/replit-3b/config-lora.yml b/examples/replit-3b/config-lora.yml index c757e720d..71a80d17d 100644 --- a/examples/replit-3b/config-lora.yml +++ b/examples/replit-3b/config-lora.yml @@ -20,6 +20,7 @@ lora_target_modules: - mlp_down lora_fan_in_fan_out: wandb_project: lora-replit +wandb_entity: wandb_watch: wandb_run_id: wandb_log_model: diff --git a/examples/xgen-7b/xgen-7b-8k-qlora.yml b/examples/xgen-7b/xgen-7b-8k-qlora.yml index e74651eb6..0f215afb9 100644 --- a/examples/xgen-7b/xgen-7b-8k-qlora.yml +++ b/examples/xgen-7b/xgen-7b-8k-qlora.yml @@ -37,6 +37,7 @@ lora_target_linear: true lora_fan_in_fan_out: wandb_project: +wandb_entity: wandb_watch: wandb_run_id: wandb_log_model: diff --git a/src/axolotl/utils/wandb.py b/src/axolotl/utils/wandb.py index d22b932cb..6c3af3177 100644 --- a/src/axolotl/utils/wandb.py +++ b/src/axolotl/utils/wandb.py @@ -9,6 +9,8 @@ def setup_wandb_env_vars(cfg): elif cfg.wandb_project and len(cfg.wandb_project) > 0: os.environ["WANDB_PROJECT"] = cfg.wandb_project cfg.use_wandb = True + if cfg.wandb_entity and len(cfg.wandb_entity) > 0: + os.environ["WANDB_ENTITY"] = cfg.wandb_entity if cfg.wandb_watch and len(cfg.wandb_watch) > 0: os.environ["WANDB_WATCH"] = cfg.wandb_watch if cfg.wandb_log_model and len(cfg.wandb_log_model) > 0: