fine-tune-code-llama.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_5rfduRD1CSF"
      },
      "source": [
        "# Fine tune your own private Copilot\n",
        "\n",
        "The integration between GitHub and Colab has been annoyingly difficult. While it's possible to open a notebook from a GitHub link in Colab, unfortunately, none of the rest of the repository content is brought into the Colab runtime. This makes it cumbersome to make use of other materials saved in your repo, that includes your dataset preprocessing scripts, structured training code, and maybe even the dataset itself. People have compromised and resorted to alternative solutions to complete a fine tuning lifecycle:\n",
        "\n",
        "1. First create some dataset and put it in GDrive or a Hugging Face dataset repo.\n",
        "2. Put up some code in notebook and run it in Colab, loading models from a Hugging Face model repo.\n",
        "3. Save the fine tuned model back into a Hugging Face model repo.\n",
        "4. Evaluate the fine tuned model. And if it's not ideal, go back to step 1.\n",
        "\n",
        "This breaks one project into three pieces stored in different places: a dataset repo, a source code (notebook) repo, and a model repo, and there's no good way to cross reference between their individual versions. For example, if one fine tuning lifecycle deteriorates, one has to manually search back into three parallel history, letting alone the difficulty to revert to a good base.\n",
        "\n",
        "In this guide we demonstrate that one can\n",
        "1. Version **all** three pieces together in one GitHub repo managed by [XetData](https://github.com/apps/XetData) GitHub app.\n",
        "2. Clones **only** what you need in the training to Colab runtime using [Lazy clone](https://xethub.com/assets/docs/large-repos/lazy-clone) feature.\n",
        "\n",
        "\n",
        "This fine tuning example uses a Lora approach on top of [Code Llama](https://ai.meta.com/blog/code-llama-large-language-model-coding/), quantizing the base model to int 8, freezing its weights and only training an adapter. Please accept their License at https://ai.meta.com/resources/models-and-libraries/llama-downloads/. Much of the code is refactored from [[1]](https://github.com/tloen/alpaca-lora), [[2]](https://github.com/samlhuillier/code-llama-fine-tune-notebook/tree/main), [[3]](https://github.com/pacman100/DHS-LLM-Workshop/tree/main/personal_copilot).\n",
        "\n",
        "\n",
        "*Avoid running this on V100 GPUs as [BF16 is not supported on V100](https://github.com/facebookresearch/llama-recipes/issues/284) and will otherwise throw out errors."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0mu9JczX1CSH"
      },
      "source": [
        "### Set up environment\n",
        "This installs necessary training libraries and [git-xet](https://xethub.com/assets/docs/getting-started/install) that adds natural support for managing large files to Git, also sets up authorization to access your github repo. To get started, first create a GitHub personal access token as mentioned [here](https://docs.github.com/en/authentication/keeping-your-account-and-data-secure/managing-your-personal-access-tokens)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "tags": [],
        "id": "6lxQKE5MW310"
      },
      "outputs": [],
      "source": [
        "# Install python dependencies\n",
        "!pip install tqdm nbformat\n",
        "!pip install git+https://github.com/huggingface/transformers.git@main bitsandbytes\n",
        "!pip install git+https://github.com/huggingface/peft.git@main\n",
        "!pip install datasets\n",
        "import locale # colab workaround\n",
        "locale.getpreferredencoding = lambda x=False:\"UTF-8\" # colab workaround"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Download and install git-xet\n",
        "!curl -fsSLO https://github.com/xetdata/xet-tools/releases/latest/download/xet-linux-x86_64.tar.gz\n",
        "!tar -xvf xet-linux-x86_64.tar.gz && rm xet-linux-x86_64.tar.gz\n",
        "!mv git-xet /usr/local/bin\n",
        "!git xet install"
      ],
      "metadata": {
        "id": "hJ-nvsLiXqU3"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0A5O9EwlW311"
      },
      "outputs": [],
      "source": [
        "# Set up authorization to access your repo where models, source code, etc. are versioned.\n",
        "from IPython.display import clear_output\n",
        "user = input(\"GitHub user name?\")\n",
        "%env GH_USER=$user\n",
        "email = input(\"GitHub user email?\")\n",
        "%env GH_USER_EMAIL=$email\n",
        "token = input(\"GitHub token?\")\n",
        "%env GH_TOKEN=$token\n",
        "%env XET_LOG_PATH=log.txt\n",
        "clear_output()"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# The repo that contains the model and fine tuning code.\n",
        "model_repo = \"LLM_fine_tuning\" # change to your own model repo\n",
        "model_repo_url = f\"https://{user}:{token}@github.com/{user}/{model_repo}.git\""
      ],
      "metadata": {
        "id": "bUgFwNrB_RuT"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Configure git for later commit author info\n",
        "!git config --global user.name $GH_USER\n",
        "!git config --global user.email $GH_USER_EMAIL\n",
        "\n",
        "# Clone in lazy mode so as to materialize files on need basis\n",
        "!git xet clone --lazy {model_repo_url} {model_repo}"
      ],
      "metadata": {
        "id": "Vs7NWkwmXKXE"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oTeYW8z51CSH",
        "tags": []
      },
      "outputs": [],
      "source": [
        "from datetime import datetime\n",
        "import os\n",
        "import sys\n",
        "\n",
        "import torch\n",
        "from peft import (\n",
        "    LoraConfig,\n",
        "    get_peft_model,\n",
        "    get_peft_model_state_dict,\n",
        "    prepare_model_for_int8_training,\n",
        "    set_peft_model_state_dict,\n",
        ")\n",
        "from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments, Trainer, DataCollatorForSeq2Seq\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6ig7NvWN1CSJ"
      },
      "source": [
        "### Load model\n",
        "Load the model from the cloned model repo. This is the base Code Llama model or your fine tuned model saved from previous runs. You can drop other LLMs into this repo, resting assured that XetData [supports per-repository limit of over 100TB and no per file or number of file limits](https://xethub.com/assets/docs/)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rMnU93bY1CSJ",
        "tags": []
      },
      "outputs": [],
      "source": [
        "model_name = \"CodeLlama-7b-hf\" # the model that you want to fine tune on\n",
        "!cd {model_repo} && git xet materialize {model_name} # brings the model files to local\n",
        "\n",
        "base_model = f\"./{model_repo}/{model_name}\"\n",
        "model = AutoModelForCausalLM.from_pretrained(\n",
        "    base_model,\n",
        "    load_in_8bit=True,\n",
        "    torch_dtype=torch.float16,\n",
        "    device_map=\"auto\",\n",
        ")\n",
        "tokenizer = AutoTokenizer.from_pretrained(base_model)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_2VXqJJe1CSJ"
      },
      "source": [
        "### Check base model\n",
        "As a baseline, let's first check how does the existing model behave.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "uiyAff1a1CSJ"
      },
      "outputs": [],
      "source": [
        "eval_prompt = \"\"\"\n",
        "def parse_url(url, force_domain='xethub.com', partial_remote=False):\n",
        "    '''\n",
        "    Parses a Xet URL of the form\n",
        "     - xet://user/repo/branch/[path]\n",
        "     - /user/repo/branch/[path]\n",
        "\n",
        "    Into a XetPathInfo which forms it as remote=https://[domain]/user/repo\n",
        "    branch=[branch] and path=[path].\n",
        "\n",
        "    branches with '/' are not supported.\n",
        "\n",
        "    If partial_remote==True, allows [repo] to be optional. i.e. it will\n",
        "    parse /user or xet://user\n",
        "    '''\n",
        "\n",
        "    <FILL_ME>\n",
        "\"\"\"\n",
        "\n",
        "model_input = tokenizer(eval_prompt, return_tensors=\"pt\").to(\"cuda\")\n",
        "\n",
        "model.eval()\n",
        "with torch.no_grad():\n",
        "    print(tokenizer.decode(model.generate(**model_input, max_new_tokens=300)[0], skip_special_tokens=True))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "544rmsdU1CSJ"
      },
      "source": [
        "I get the below output, which is erroneous and contains repetitive code.\n",
        "```\n",
        "if url.startswith('xet://'):\n",
        "        url = url[len('xet://'):]\n",
        "\n",
        "    if url.startswith('/'):\n",
        "        url = url[1:]\n",
        "\n",
        "    if '/' in url:\n",
        "        repo, branch, path = url.split('/', 2)\n",
        "    else:\n",
        "        repo, branch = url.split('/', 1)\n",
        "        path = ''\n",
        "\n",
        "    if not repo:\n",
        "        raise ValueError('No repo specified')\n",
        "\n",
        "    if not branch:\n",
        "        raise ValueError('No branch specified')\n",
        "\n",
        "    if not path:\n",
        "        path = ''\n",
        "\n",
        "    if not partial_remote:\n",
        "        if not repo.startswith('xet-'):\n",
        "            raise ValueError('Invalid repo name')\n",
        "\n",
        "        if not repo.endswith('.git'):\n",
        "            raise ValueError('Invalid repo name')\n",
        "\n",
        "        repo = repo[4:-4]\n",
        "\n",
        "    if not repo:\n",
        "        raise ValueError('No repo specified')\n",
        "\n",
        "    if not branch:\n",
        "        raise ValueError('No branch specified')\n",
        "\n",
        "    if not path:\n",
        "        path = ''\n",
        "\n",
        "    if not path.startswith('/'):\n",
        "        path = '/' + path\n",
        "\n",
        "    if not path.endswith('/'):\n",
        "        path = path + '/'\n",
        "\n",
        "    return XetPathInfo(\n",
        "        remote='https://'\n",
        "```"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4M9KyT0S1CSH"
      },
      "source": [
        "### Load dataset\n",
        "This example fine tunes using the [pyxet](https://github.com/xetdata/pyxet) project source code as the dataset. We use the scripts stored together in the model repo to clone pyxet and extract source code into a pandas DataFrame of format `['repo_id', 'file_path', 'content']`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "tags": [],
        "id": "KXXSZVwKW311"
      },
      "outputs": [],
      "source": [
        "import importlib\n",
        "myscripts=importlib.import_module(f\"{model_repo}.scripts\")\n",
        "import pandas as pd\n",
        "from datasets import Dataset\n",
        "\n",
        "# Clones a source code repository as fine tuning data\n",
        "username='xetdata'\n",
        "repository='pyxet'\n",
        "parquet_file = myscripts.create_dataset_from_git_repo(username,repository)\n",
        "# Optionally you can save this dataset back to the model repo\n",
        "df = pd.read_parquet(parquet_file)\n",
        "dataset = Dataset.from_pandas(df, split=\"train\")\n",
        "train_dataset = dataset.train_test_split(test_size=0.1)[\"train\"]\n",
        "eval_dataset = dataset.train_test_split(test_size=0.1)[\"test\"]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "puuOXL2R1CSJ"
      },
      "source": [
        "### Tokenization\n",
        "Source code files come with drastically different length, feeding them directly into tensors requires padding and/or truncation which either amplifies memory usage or discards information. We split file content into constant length (512 tokens) chunks and tokenize each chunk."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SA2BqSzW1CSK"
      },
      "outputs": [],
      "source": [
        "tokenized_train_dataset, tokenized_val_dataset = myscripts.constant_length_token_seq_from(tokenizer, train_dataset, eval_dataset, seq_length=512)\n",
        "tokenized_train_dataset = [s for s in tokenized_train_dataset]\n",
        "tokenized_val_dataset = [s for s in tokenized_val_dataset]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Z5ByhitV1CSK"
      },
      "source": [
        "### 5. Setup Lora"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9q_26pz71CSK"
      },
      "outputs": [],
      "source": [
        "model.train() # put model back into training mode\n",
        "model = prepare_model_for_int8_training(model)\n",
        "\n",
        "config = LoraConfig(\n",
        "    r=16,\n",
        "    lora_alpha=16,\n",
        "    target_modules=[\n",
        "    \"q_proj\",\n",
        "    \"k_proj\",\n",
        "    \"v_proj\",\n",
        "    \"o_proj\",\n",
        "],\n",
        "    lora_dropout=0.05,\n",
        "    bias=\"none\",\n",
        "    task_type=\"CAUSAL_LM\",\n",
        ")\n",
        "model = get_peft_model(model, config)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pYSsnciQ1CSK"
      },
      "source": [
        "### 6. Training arguments\n",
        "If you run out of GPU memory, change per_device_train_batch_size. The gradient_accumulation_steps variable should ensure this doesn't affect batch dynamics during the training run. All the other variables are standard stuff that I wouldn't recommend messing with:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5iv4txxu1CSK"
      },
      "outputs": [],
      "source": [
        "if torch.cuda.device_count() > 1:\n",
        "    # keeps Trainer from trying its own DataParallelism when more than 1 gpu is available\n",
        "    model.is_parallelizable = True\n",
        "    model.model_parallel = True"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "uvEi1kP21CSK"
      },
      "outputs": [],
      "source": [
        "batch_size = 128\n",
        "per_device_train_batch_size = 4\n",
        "gradient_accumulation_steps = batch_size // per_device_train_batch_size\n",
        "output_dir = \"code-llama\" # to write checkpoints\n",
        "\n",
        "training_args = TrainingArguments(\n",
        "        per_device_train_batch_size=per_device_train_batch_size,\n",
        "        gradient_accumulation_steps=gradient_accumulation_steps,\n",
        "        warmup_steps=100,\n",
        "        max_steps=400,\n",
        "        #num_train_epochs=3,\n",
        "        learning_rate=3e-4,\n",
        "        fp16=True,\n",
        "        logging_steps=10,\n",
        "        optim=\"adamw_torch\",\n",
        "        evaluation_strategy=\"steps\", # if val_set_size > 0 else \"no\",\n",
        "        save_strategy=\"steps\",\n",
        "        eval_steps=20,\n",
        "        save_steps=20,\n",
        "        output_dir=output_dir,\n",
        "        # save_total_limit=3,\n",
        "        load_best_model_at_end=False,\n",
        "        # ddp_find_unused_parameters=False if ddp else None,\n",
        "        report_to=\"none\", # if use_wandb else \"none\",\n",
        "        run_name=f\"codellama-{datetime.now().strftime('%Y-%m-%d-%H-%M')}\", # if use_wandb else None,\n",
        "    )\n",
        "\n",
        "trainer = Trainer(\n",
        "    model=model,\n",
        "    train_dataset=tokenized_train_dataset,\n",
        "    eval_dataset=tokenized_val_dataset,\n",
        "    args=training_args,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1aJp6Jxl1CSK"
      },
      "source": [
        "Then we do some pytorch-related optimizations (which just make training faster but don't affect accuracy):"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6ycCHZZl1CSK"
      },
      "outputs": [],
      "source": [
        "model.config.use_cache = False\n",
        "\n",
        "if torch.__version__ >= \"2\" and sys.platform != \"win32\":\n",
        "    print(\"compiling the model\")\n",
        "    model = torch.compile(model)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bF5oWKxK1CSL"
      },
      "outputs": [],
      "source": [
        "trainer.train()"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "On a T4 GPU, I'm getting the below training speed:\n",
        " - 21 steps or ~3.70 epochs for 2:04:20\n",
        " - 109 steps or ~20.56 epochs for 9:29:53"
      ],
      "metadata": {
        "id": "FyLkS5vqOLVq"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "model.save_pretrained(output_dir)"
      ],
      "metadata": {
        "id": "Fhlh6w4AxLoG"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "j1dRQLMT1CSU"
      },
      "source": [
        "### Load the final checkpoint\n",
        "\n",
        "If loading from a certain checkpoint, load the base model and the adapters separately, the checkpoint directory should contain an adapter_config.json and an adapter_model.safetensors:\n",
        "```\n",
        "from peft import PeftModel\n",
        "model = AutoModelForCausalLM.from_pretrained(\n",
        "    base_model,\n",
        "    load_in_8bit=True,\n",
        "    torch_dtype=torch.float16,\n",
        "    device_map=\"auto\",\n",
        ")\n",
        "model = PeftModel.from_pretrained(model, [checkpoint directory])\n",
        "```"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "roqy_WRi1CSU"
      },
      "source": [
        "Otherwise, directly try the same prompt as before:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GrCxouNp1CSU"
      },
      "outputs": [],
      "source": [
        "eval_prompt = \"\"\"\n",
        "def parse_url(url, force_domain='xethub.com', partial_remote=False):\n",
        "    '''\n",
        "    Parses a Xet URL of the form\n",
        "     - xet://user/repo/branch/[path]\n",
        "     - /user/repo/branch/[path]\n",
        "\n",
        "    Into a XetPathInfo which forms it as remote=https://[domain]/user/repo\n",
        "    branch=[branch] and path=[path].\n",
        "\n",
        "    branches with '/' are not supported.\n",
        "\n",
        "    If partial_remote==True, allows [repo] to be optional. i.e. it will\n",
        "    parse /user or xet://user\n",
        "    '''\n",
        "\n",
        "    <FILL_ME>\n",
        "\"\"\"\n",
        "\n",
        "model_input = tokenizer(eval_prompt, return_tensors=\"pt\").to(\"cuda\")\n",
        "\n",
        "model.eval()\n",
        "with torch.no_grad():\n",
        "    print(tokenizer.decode(model.generate(**model_input, max_new_tokens=512)[0], skip_special_tokens=True))\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "I6IcTOCq1CSU"
      },
      "source": [
        "And the model outputs the below code after training for 100 steps. This is much better than before the fine tuning.\n",
        "```\n",
        "import re\n",
        "    from .xet_path import XetPathInfo\n",
        "\n",
        "    if url.startswith('xet://'):\n",
        "        url = url[len('xet://'):]\n",
        "\n",
        "    if url.startswith('/'):\n",
        "        url = url[1:]\n",
        "\n",
        "    if len(url) == 0:\n",
        "        raise ValueError('Invalid Xet URL')\n",
        "\n",
        "    parts = url.split('/')\n",
        "    if len(parts) < 3:\n",
        "        raise ValueError('Invalid Xet URL')\n",
        "\n",
        "    if len(parts) == 3:\n",
        "        branch = parts[2]\n",
        "        path = ''\n",
        "    else:\n",
        "        branch = parts[2]\n",
        "        path = '/'.join(parts[3:])\n",
        "\n",
        "    if len(branch) == 0:\n",
        "        raise ValueError('Invalid Xet URL')\n",
        "\n",
        "    if partial_remote:\n",
        "        if len(parts) == 2:\n",
        "            return XetPathInfo(f'https://{force_domain}/{parts[0]}', branch, path)\n",
        "        else:\n",
        "            return XetPathInfo(f'https://{force_domain}/{parts[0]}/{parts[1]}', branch, path)\n",
        "    else:\n",
        "        return XetPathInfo(f'https://{force_domain}/{parts[0]}/{parts[1]}', branch, path)\n",
        "```\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Finally merge the adapter into the model and save the model back to the repo.\n",
        "model = model.merge_and_unload()\n",
        "model.save_pretrained(base_model)\n",
        "commit_id=!cd {repository} && git rev-parse --short HEAD\n",
        "commit_id=commit_id[0]\n",
        "!cd {model_repo} && git add {model_name} && git commit -m \"Fine tuned model trained on {username}/{repository}@{commit_id}\" && git push"
      ],
      "metadata": {
        "id": "yHEgbSFyMJTN"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "###At last, save this notebook back to your GitHub repository if there are changes."
      ],
      "metadata": {
        "id": "_UNd_1G1c7IJ"
      }
    }
  ],
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4"
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.10.10"
    },
    "accelerator": "GPU"
  },
  "nbformat": 4,
  "nbformat_minor": 0
}