add wandb support

.gitignore

@@ -3,4 +3,5 @@ __pycache__/
 /venv
 .vscode
 
+/wandb
 /data

README.md

@@ -60,13 +60,14 @@ file_mounts:
 setup: |
   git clone https://github.com/zetavg/LLaMA-LoRA-Tuner.git llama_lora_tuner
   cd llama_lora_tuner && pip install -r requirements.lock.txt
+  pip install wandb
   cd ..
   echo 'Dependencies installed.'
 
 # Start the app.
 run: |
   echo 'Starting...'
-  python llama_lora_tuner/app.py --data_dir='/data' --base_model='decapoda-research/llama-7b-hf' --share
+  python llama_lora_tuner/app.py --data_dir='/data' --wandb_api_key "$([ -f /data/secrets/wandb_api_key ] && cat /data/secrets/wandb_api_key | tr -d '\n')" --base_model='decapoda-research/llama-7b-hf' --share
 ```
 
 Then launch a cluster to run the task:

app.py

@@ -5,21 +5,37 @@ import fire
 import gradio as gr
 
 from llama_lora.globals import Global
+from llama_lora.models import prepare_base_model
 from llama_lora.ui.main_page import main_page, get_page_title, main_page_custom_css
 from llama_lora.utils.data import init_data_dir
 
 
+
 def main(
-    load_8bit: bool = False,
     base_model: str = "",
     data_dir: str = "",
     # Allows to listen on all interfaces by providing '0.0.0.0'.
     server_name: str = "127.0.0.1",
     share: bool = False,
     skip_loading_base_model: bool = False,
+    load_8bit: bool = False,
     ui_show_sys_info: bool = True,
     ui_dev_mode: bool = False,
+    wandb_api_key: str = "",
+    wandb_project: str = "",
 ):
+    '''
+    Start the LLaMA-LoRA Tuner UI.
+
+    :param base_model: (required) The name of the default base model to use.
+    :param data_dir: (required) The path to the directory to store data.
+    :param server_name: Allows to listen on all interfaces by providing '0.0.0.0'.
+    :param share: Create a public Gradio URL.
+
+    :param wandb_api_key: The API key for Weights & Biases. Setting either this or `wandb_project` will enable Weights & Biases.
+    :param wandb_project: The default project name for Weights & Biases. Setting either this or `wandb_api_key` will enable Weights & Biases.
+    '''
+
     base_model = base_model or os.environ.get("LLAMA_LORA_BASE_MODEL", "")
     data_dir = data_dir or os.environ.get("LLAMA_LORA_DATA_DIR", "")
     assert (
@@ -34,12 +50,22 @@ def main(
     Global.data_dir = os.path.abspath(data_dir)
     Global.load_8bit = load_8bit
 
+    if len(wandb_api_key) > 0:
+        Global.enable_wandb = True
+        Global.wandb_api_key = wandb_api_key
+    if len(wandb_project) > 0:
+        Global.enable_wandb = True
+        Global.wandb_project = wandb_project
+
     Global.ui_dev_mode = ui_dev_mode
     Global.ui_show_sys_info = ui_show_sys_info
 
     os.makedirs(data_dir, exist_ok=True)
     init_data_dir()
 
+    if (not skip_loading_base_model) and (not ui_dev_mode):
+        prepare_base_model(base_model)
+
     with gr.Blocks(title=get_page_title(), css=main_page_custom_css()) as demo:
         main_page()
 

llama_lora/globals.py

@@ -40,6 +40,11 @@ class Global:
     gpu_total_cores = None  # GPU total cores
     gpu_total_memory = None
 
+    # WandB
+    enable_wandb = False
+    wandb_api_key = None
+    default_wandb_project = "llama-lora-tuner"
+
     # UI related
     ui_title: str = "LLaMA-LoRA Tuner"
     ui_emoji: str = "🦙🎛️"

llama_lora/lib/finetune.py

@@ -50,8 +50,32 @@ def train(
     save_total_limit: int = 3,
     logging_steps: int = 10,
     # logging
-    callbacks: List[Any] = []
+    callbacks: List[Any] = [],
+    # wandb params
+    wandb_api_key = None,
+    wandb_project: str = "",
+    wandb_run_name: str = "",
+    wandb_watch: str = "false",  # options: false | gradients | all
+    wandb_log_model: str = "true",  # options: false | true
 ):
+    if wandb_api_key:
+        os.environ["WANDB_API_KEY"] = wandb_api_key
+    if wandb_project:
+        os.environ["WANDB_PROJECT"] = wandb_project
+    if wandb_run_name:
+        os.environ["WANDB_RUN_NAME"] = wandb_run_name
+    if wandb_watch:
+        os.environ["WANDB_WATCH"] = wandb_watch
+    if wandb_log_model:
+        os.environ["WANDB_LOG_MODEL"] = wandb_log_model
+    use_wandb = (wandb_project and len(wandb_project) > 0) or (
+            "WANDB_PROJECT" in os.environ and len(os.environ["WANDB_PROJECT"]) > 0
+        )
+    if use_wandb:
+        os.environ['WANDB_MODE'] = "online"
+    else:
+        os.environ['WANDB_MODE'] = "disabled"
+
     if os.path.exists(output_dir):
         if (not os.path.isdir(output_dir)) or os.path.exists(os.path.join(output_dir, 'adapter_config.json')):
             raise ValueError(
@@ -204,8 +228,8 @@ def train(
             load_best_model_at_end=True if val_set_size > 0 else False,
             ddp_find_unused_parameters=False if ddp else None,
             group_by_length=group_by_length,
-            # report_to="wandb" if use_wandb else None,
-            # run_name=wandb_run_name if use_wandb else None,
+            report_to="wandb" if use_wandb else None,
+            run_name=wandb_run_name if use_wandb else None,
         ),
         data_collator=transformers.DataCollatorForSeq2Seq(
             tokenizer, pad_to_multiple_of=8, return_tensors="pt", padding=True

llama_lora/ui/finetune_ui.py

@@ -491,7 +491,10 @@ Train data (first 10):
             save_steps,  # save_steps
             save_total_limit,  # save_total_limit
             logging_steps,  # logging_steps
-            training_callbacks  # callbacks
+            training_callbacks,  # callbacks
+            Global.wandb_api_key,  # wandb_api_key
+            Global.default_wandb_project if Global.enable_wandb else None,  # wandb_project
+            model_name  # wandb_run_name
         )
 
         logs_str = "\n".join([json.dumps(log)