Merge branch 'dev-2'

LLaMA_LoRA.ipynb

@@ -6,7 +6,7 @@
       "provenance": [],
       "private_outputs": true,
       "toc_visible": true,
-      "authorship_tag": "ABX9TyOQclbOikH+aaXWlbaTkEqc",
+      "authorship_tag": "ABX9TyMHMc4PwWLbRlhFol+WRzoT",
       "include_colab_link": true
     },
     "kernelspec": {
@@ -34,7 +34,8 @@
       "cell_type": "markdown",
       "source": [
         "# 🦙🎛️ LLaMA-LoRA\n",
-        "\n",        "TL;DR: **Runtime > Run All** (`⌘/Ctrl+F9`). Takes about 5 minutes to start. You will be promped to authorize Google Drive access."
+        "\n",
+        "TL;DR: **Runtime > Run All** (`⌘/Ctrl+F9`). Takes about 5 minutes to start. You will be promped to authorize Google Drive access."
       ],
       "metadata": {
         "id": "bb4nzBvLfZUj"
@@ -309,41 +310,6 @@
       },
       "execution_count": null,
       "outputs": []
-    },
-    {
-      "cell_type": "markdown",
-      "source": [
-        "# Reset"
-      ],
-      "metadata": {
-        "id": "RW09SrCZpqpa"
-      }
-    },
-    {
-      "cell_type": "code",
-      "source": [
-        "# @title Kill Session { display-mode: \"form\" }\n",
-        "# @markdown If you ran out of runtime resources, you can **check the following \n",
-        "# @markdown checkbox and run this code cell to kill the runtime session** while\n",
-        "# @markdown preserving your downloaded data.\n",
-        "do_kill_session = False # @param {type:\"boolean\"}\n",
-        "# @markdown You will need to re-run this notebook from start after doing this.\n",
-        "#\n",
-        "# @markdown All data that are saved to disk, including Python dependencies, base\n",
-        "# @markdown models will all be preserved, so the second run will be much faster.\n",
-        "\n",
-        "import os\n",
-        "def kill_session():\n",
-        "    os.kill(os.getpid(), 9)\n",
-        "\n",
-        "if do_kill_session:\n",
-        "    kill_session()"
-      ],
-      "metadata": {
-        "id": "bM4sY2tVps8U"
-      },
-      "execution_count": null,
-      "outputs": []
     }
   ]
 }

README.md

@@ -15,6 +15,100 @@ Making evaluating and fine-tuning LLaMA models with low-rank adaptation (LoRA) e
   * Supports Stanford Alpaca [seed_tasks](https://github.com/tatsu-lab/stanford_alpaca/blob/main/seed_tasks.jsonl), [alpaca_data](https://github.com/tatsu-lab/stanford_alpaca/blob/main/alpaca_data.json) and [OpenAI "prompt"-"completion"](https://platform.openai.com/docs/guides/fine-tuning/data-formatting) format.
 
 
+## How to Start
+
+There are various ways to run this app:
+
+* **[Run on Google Colab](#run-on-google-colab)**: The simplest way to get started, all you need is a Google account. Standard (free) GPU runtime is sufficient to run generation and training with micro batch size of 8. However, the text generation and training is much slower than on other cloud services, and Colab might terminate the execution in inactivity while running long tasks.
+* **[Run on a cloud service via SkyPilot](#run-on-a-cloud-service-via-skypilot)**: If you have a cloud service (Lambda Labs, GCP, AWS, or Azure) account, you can use SkyPilot to run the app on a cloud service. A cloud bucket can be mounted to preserve your data.
+* **[Run locally](#run-locally)**: Depends on the hardware you have.
+
+### Run On Google Colab
+
+Open [this Colab Notebook](https://colab.research.google.com/github/zetavg/LLaMA-LoRA/blob/main/LLaMA_LoRA.ipynb) and select **Runtime > Run All** (`⌘/Ctrl+F9`).
+
+You will be prompted to authorize Google Drive access, as Google Drive will be used to store your data. See the "Config"/"Google Drive" section for settings and more info.
+
+After approximately 5 minutes of running, you will see the public URL in the output of the "Launch"/"Start Gradio UI 🚀" section (like `Running on public URL: https://xxxx.gradio.live`). Open the URL in your browser to use the app.
+
+### Run on a cloud service via SkyPilot
+
+After following the [installation guide of SkyPilot](https://skypilot.readthedocs.io/en/latest/getting-started/installation.html), create a `.yaml` to define a task for running the app:
+
+```yaml
+# llama-lora-multitool.yaml
+
+resources:
+  accelerators: A10:1  # 1x NVIDIA A10 GPU
+  cloud: lambda  # Optional; if left out, SkyPilot will automatically pick the cheapest cloud.
+
+file_mounts:
+  # Mount a presisted cloud storage that will be used as the data directory.
+  # (to store train datasets trained models)
+  # See https://skypilot.readthedocs.io/en/latest/reference/storage.html for details.
+  /data:
+    name: llama-lora-multitool-data  # Make sure this name is unique or you own this bucket. If it does not exists, SkyPilot will try to create a bucket with this name.
+    store: gcs  # Could be either of [s3, gcs]
+    mode: MOUNT
+
+# Clone the LLaMA-LoRA repo and install its dependencies.
+setup: |
+  git clone https://github.com/zetavg/LLaMA-LoRA.git llama_lora
+  cd llama_lora && pip install -r requirements.txt
+  cd ..
+  echo 'Dependencies installed.'
+
+# Start the app.
+run: |
+  echo 'Starting...'
+  python llama_lora/app.py --data_dir='/data' --base_model='decapoda-research/llama-7b-hf' --share
+```
+
+Then launch a cluster to run the task:
+
+```
+sky launch -c llama-lora-multitool llama-lora-multitool.yaml
+```
+
+`-c ...` is an optional flag to specify a cluster name. If not specified, SkyPilot will automatically generate one.
+
+You will see the public URL of the app in the terminal. Open the URL in your browser to use the app.
+
+Note that exiting `sky launch` will only exit log streaming and will not stop the task. You can use `sky queue --skip-finished` to see the status of running or pending tasks, `sky logs <cluster_name> <job_id>` connect back to log streaming, and `sky cancel <cluster_name> <job_id>` to stop a task.
+
+When you are done, run `sky stop <cluster_name>` to stop the cluster. To terminate a cluster instead, run `sky down <cluster_name>`.
+
+### Run locally
+
+<details>
+  <summary>Prepare environment with conda</summary>
+
+  ```bash
+  conda create -y -n llama-lora-multitool python=3.8
+  conda activate llama-lora-multitool
+  ```
+</details>
+
+```bash
+pip install -r requirements.txt
+python app.py --data_dir='./data' --base_model='decapoda-research/llama-7b-hf' --share
+```
+
+You will see the local and public URLs of the app in the terminal. Open the URL in your browser to use the app.
+
+For more options, see `python app.py --help`.
+
+<details>
+  <summary>UI development mode</summary>
+
+  To test the UI without loading the language model, use the `--ui_dev_mode` flag:
+
+  ```bash
+  python app.py --data_dir='./data' --base_model='decapoda-research/llama-7b-hf' --share --ui_dev_mode
+  ```
+</details>
+
+
 ## Acknowledgements
 
 * https://github.com/tloen/alpaca-lora

llama_lora/globals.py

@@ -3,6 +3,9 @@ import subprocess
 
 from typing import Any, Dict, List, Optional, Tuple, Union
 
+from numba import cuda
+import nvidia_smi
+
 from .lib.finetune import train
 
 
@@ -25,6 +28,12 @@ class Global:
     # Model related
     model_has_been_used = False
 
+    # GPU Info
+    gpu_cc = None  # GPU compute capability
+    gpu_sms = None  # GPU total number of SMs
+    gpu_total_cores = None  # GPU total cores
+    gpu_total_memory = None
+
     # UI related
     ui_title: str = "LLaMA-LoRA"
     ui_emoji: str = "🦙🎛️"
@@ -60,3 +69,57 @@ commit_hash = get_git_commit_hash()
 
 if commit_hash:
     Global.version = commit_hash[:8]
+
+
+def load_gpu_info():
+    try:
+        cc_cores_per_SM_dict = {
+            (2, 0): 32,
+            (2, 1): 48,
+            (3, 0): 192,
+            (3, 5): 192,
+            (3, 7): 192,
+            (5, 0): 128,
+            (5, 2): 128,
+            (6, 0): 64,
+            (6, 1): 128,
+            (7, 0): 64,
+            (7, 5): 64,
+            (8, 0): 64,
+            (8, 6): 128,
+            (8, 9): 128,
+            (9, 0): 128
+        }
+        # the above dictionary should result in a value of "None" if a cc match
+        # is not found.  The dictionary needs to be extended as new devices become
+        # available, and currently does not account for all Jetson devices
+        device = cuda.get_current_device()
+        device_sms = getattr(device, 'MULTIPROCESSOR_COUNT')
+        device_cc = device.compute_capability
+        cores_per_sm = cc_cores_per_SM_dict.get(device_cc)
+        total_cores = cores_per_sm*device_sms
+        print("GPU compute capability: ", device_cc)
+        print("GPU total number of SMs: ", device_sms)
+        print("GPU total cores: ", total_cores)
+        Global.gpu_cc = device_cc
+        Global.gpu_sms = device_sms
+        Global.gpu_total_cores = total_cores
+
+        nvidia_smi.nvmlInit()
+        handle = nvidia_smi.nvmlDeviceGetHandleByIndex(0)
+        info = nvidia_smi.nvmlDeviceGetMemoryInfo(handle)
+        total_memory = info.total
+
+        total_memory_mb = total_memory / (1024 ** 2)
+        total_memory_gb = total_memory / (1024 ** 3)
+
+        # Print the memory size
+        print(
+            f"GPU total memory: {total_memory} bytes ({total_memory_mb:.2f} MB) ({total_memory_gb:.2f} GB)")
+        Global.gpu_total_memory = total_memory
+
+    except Exception as e:
+        print(f"Notice: cannot get GPU info: {e}")
+
+
+load_gpu_info()

llama_lora/models.py

@@ -102,6 +102,14 @@ def load_base_model():
             )
 
 
+def clear_cache():
+    gc.collect()
+
+    # if not shared.args.cpu: # will not be running on CPUs anyway
+    with torch.no_grad():
+        torch.cuda.empty_cache()
+
+
 def unload_models():
     del Global.loaded_base_model
     Global.loaded_base_model = None
@@ -109,11 +117,7 @@ def unload_models():
     del Global.loaded_tokenizer
     Global.loaded_tokenizer = None
 
-    gc.collect()
-
-    # if not shared.args.cpu: # will not be running on CPUs anyway
-    with torch.no_grad():
-        torch.cuda.empty_cache()
+    clear_cache()
 
     Global.model_has_been_used = False
 

llama_lora/ui/finetune_ui.py

@@ -9,7 +9,9 @@ from random_word import RandomWords
 from transformers import TrainerCallback
 
 from ..globals import Global
-from ..models import get_base_model, get_tokenizer, unload_models_if_already_used
+from ..models import (
+    get_base_model, get_tokenizer,
+    clear_cache, unload_models_if_already_used)
 from ..utils.data import (
     get_available_template_names,
     get_available_dataset_names,
@@ -238,6 +240,12 @@ def parse_plain_text_input(
     return result
 
 
+should_training_progress_track_tqdm = True
+
+if Global.gpu_total_cores is not None and Global.gpu_total_cores > 2560:
+    should_training_progress_track_tqdm = False
+
+
 def do_train(
     # Dataset
     template,
@@ -258,9 +266,15 @@ def do_train(
     lora_alpha,
     lora_dropout,
     model_name,
-    progress=gr.Progress(track_tqdm=True),
+    progress=gr.Progress(track_tqdm=should_training_progress_track_tqdm),
 ):
     try:
+        clear_cache()
+        # If model has been used in inference, we need to unload it first.
+        # Otherwise, we'll get a 'Function MmBackward0 returned an invalid
+        # gradient at index 1 - expected device meta but got cuda:0' error.
+        unload_models_if_already_used()
+
         prompter = Prompter(template)
         variable_names = prompter.get_variable_names()
 
@@ -312,7 +326,32 @@ def do_train(
                 'completion': d['output']}
             for d in data]
 
+        def get_progress_text(epoch, epochs, last_loss):
+            progress_detail = f"Epoch {math.ceil(epoch)}/{epochs}"
+            if last_loss is not None:
+                progress_detail += f", Loss: {last_loss:.4f}"
+            return f"Training... ({progress_detail})"
+
         if Global.ui_dev_mode:
+            Global.should_stop_training = False
+
+            for i in range(300):
+                if (Global.should_stop_training):
+                    return
+                epochs = 3
+                epoch = i / 100
+                last_loss = None
+                if (i > 20):
+                    last_loss = 3 + (i - 0) * (0.5 - 3) / (300 - 0)
+
+                progress(
+                    (i, 300),
+                    desc="(Simulate) " +
+                    get_progress_text(epoch, epochs, last_loss)
+                )
+
+                time.sleep(0.1)
+
             message = f"""Currently in UI dev mode, not doing the actual training.
 
 Train options: {json.dumps({
@@ -368,16 +407,21 @@ Train data (first 10):
 
         training_callbacks = [UiTrainerCallback]
 
-        # If model has been used in inference, we need to unload it first.
-        # Otherwise, we'll get a 'Function MmBackward0 returned an invalid
-        # gradient at index 1 - expected device meta but got cuda:0' error.
+        Global.should_stop_training = False
+
+        # Do this again right before training to make sure the model is not used in inference.
         unload_models_if_already_used()
+        clear_cache()
 
-        Global.should_stop_training = False
+        base_model = get_base_model()
+        tokenizer = get_tokenizer()
+
+        # Do not let other tqdm iterations interfere the progress reporting after training starts.
+        # progress.track_tqdm = False  # setting this dynamically is not working, determining if track_tqdm should be enabled based on GPU cores at start instead.
 
         results = Global.train_fn(
-            get_base_model(),  # base_model
-            get_tokenizer(),  # tokenizer
+            base_model,  # base_model
+            tokenizer,  # tokenizer
             os.path.join(Global.data_dir, "lora_models",
                          model_name),  # output_dir
             train_data,
@@ -398,7 +442,8 @@ Train data (first 10):
             training_callbacks  # callbacks
         )
 
-        logs_str = "\n".join([json.dumps(log) for log in log_history]) or "None"
+        logs_str = "\n".join([json.dumps(log)
+                             for log in log_history]) or "None"
 
         result_message = f"Training ended:\n{str(results)}\n\nLogs:\n{logs_str}"
         print(result_message)
@@ -557,9 +602,18 @@ def finetune_ui():
             )
 
         with gr.Row():
+            micro_batch_size_default_value = 1
+
+            if Global.gpu_total_cores is not None and Global.gpu_total_memory is not None:
+                memory_per_core = Global.gpu_total_memory / Global.gpu_total_cores
+                if memory_per_core >= 6291456:
+                    micro_batch_size_default_value = 8
+                elif memory_per_core >= 4000000:  # ?
+                    micro_batch_size_default_value = 4
+
             with gr.Column():
                 micro_batch_size = gr.Slider(
-                    minimum=1, maximum=100, step=1, value=8,
+                    minimum=1, maximum=100, step=1, value=micro_batch_size_default_value,
                     label="Micro Batch Size",
                     info="The number of examples in each mini-batch for gradient computation. A smaller micro_batch_size reduces memory usage but may increase training time."
                 )

llama_lora/ui/inference_ui.py

@@ -1,5 +1,6 @@
 import gradio as gr
 import time
+import json
 
 import torch
 import transformers
@@ -47,10 +48,30 @@ def do_inference(
                 lora_model_name = path_of_available_lora_model
 
         if Global.ui_dev_mode:
-            message = f"Currently in UI dev mode, not running actual inference.\n\nLoRA model: {lora_model_name}\n\nYour prompt is:\n\n{prompt}"
+            message = f"Hi, I’m currently in UI-development mode and do not have access to resources to process your request. However, this behavior is similar to what will actually happen, so you can try and see how it will work!\n\nBase model: {Global.base_model}\nLoRA model: {lora_model_name}\n\nThe following text is your prompt:\n\n{prompt}"
             print(message)
+
+            if stream_output:
+                def word_generator(sentence):
+                    lines = message.split('\n')
+                    out = ""
+                    for line in lines:
+                        words = line.split(' ')
+                        for i in range(len(words)):
+                            if out:
+                                out += ' '
+                            out += words[i]
+                            yield out
+                        out += "\n"
+                        yield out
+
+                for partial_sentence in word_generator(message):
+                    yield partial_sentence, json.dumps(list(range(len(partial_sentence.split()))), indent=2)
+                    time.sleep(0.05)
+
+                return
             time.sleep(1)
-            yield message, '[0]'
+            yield message, json.dumps(list(range(len(message.split()))), indent=2)
             return
 
         model = get_base_model()
@@ -224,7 +245,7 @@ def inference_ui():
                         preview_prompt = gr.Textbox(
                             show_label=False, interactive=False, elem_id="inference_preview_prompt")
                         update_prompt_preview_btn = gr.Button(
-                            "↻", elem_id="inference_update_prompt_preview_btn", full_width=False)
+                            "↻", elem_id="inference_update_prompt_preview_btn")
                         update_prompt_preview_btn.style(size="sm")
 
                 # with gr.Column():

requirements.txt

@@ -7,6 +7,8 @@ datasets
 fire
 git+https://github.com/huggingface/peft.git
 git+https://github.com/huggingface/transformers.git
+numba
+nvidia-ml-py3
 gradio
 loralib
 sentencepiece