RWKV_v4neo_Fine_Tuning.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/resloved/RWKV-notebooks/blob/master/RWKV_v4neo_Fine_Tuning.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Vx7KFfeieD7z"
      },
      "source": [
        "# RWKV-v4neo Fine-Tuning\n",
        "\n",
        "[RWKV](https://github.com/BlinkDL/RWKV-LM) is an RNN with transformer-level performance\n",
        "\n",
        "\n",
        "This notebook aims to streamline fine-tuning RWKV-v4 models"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7JFIiAsrfvJy"
      },
      "source": [
        "\n",
        "## Setup"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "g_qFjgYmtSfK"
      },
      "outputs": [],
      "source": [
        "#@title Google Drive Options { display-mode: \"form\" }\n",
        "save_models_to_drive = True #@param {type:\"boolean\"}\n",
        "drive_mount = '/content/drive' #@param {type:\"string\"}\n",
        "output_dir = 'rwkv-v4neo-rnn-pile-tuning' #@param {type:\"string\"}\n",
        "tuned_model_name = 'tuned' #@param {type:\"string\"}\n",
        "\n",
        "import os\n",
        "from google.colab import drive\n",
        "if save_models_to_drive:\n",
        "    from google.colab import drive\n",
        "    drive.mount(drive_mount)\n",
        "    \n",
        "output_path = f\"{drive_mount}/MyDrive/{output_dir}\" if save_models_to_drive else f\"/content/{output_dir}\"\n",
        "os.makedirs(f\"{output_path}/{tuned_model_name}\", exist_ok=True)\n",
        "os.makedirs(f\"{output_path}/base_models/\", exist_ok=True)\n",
        "\n",
        "print(f\"Saving models to {output_path}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "eivKJ6FP1_9z"
      },
      "outputs": [],
      "source": [
        "!nvidia-smi"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "R4lt0FTegJw9"
      },
      "outputs": [],
      "source": [
        "!git clone https://github.com/blinkdl/RWKV-LM\n",
        "repo_dir = \"/content/RWKV-LM/RWKV-v4neo\"\n",
        "%cd $repo_dir"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RDavUrBsgKIV"
      },
      "outputs": [],
      "source": [
        "!pip install transformers pytorch-lightning==1.9 deepspeed wandb ninja"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Wt7y7vR6e6U3"
      },
      "source": [
        "## Load Base Model\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KIgagN-Se3wi",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title Base Model Options\n",
        "#@markdown Using any of the listed options will download the checkpoint from huggingface\n",
        "\n",
        "base_model_name = \"RWKV-4-Pile-169M\" #@param [\"RWKV-4-Pile-1B5\", \"RWKV-4-Pile-430M\", \"RWKV-4-Pile-169M\"]\n",
        "base_model_url = f\"https://huggingface.co/BlinkDL/{base_model_name.lower()}\"\n",
        "\n",
        "if base_model_name == \"RWKV-4-Pile-169M\":\n",
        "    n_layer = 12\n",
        "    n_embd = 768\n",
        "elif base_model_name == \"RWKV-4-Pile-430M\":\n",
        "    n_layer = 24\n",
        "    n_embd = 1024\n",
        "elif base_model_name == \"RWKV-4-Pile-1B5\":\n",
        "    n_layer = 24\n",
        "    n_embd = 2048\n",
        "\n",
        "!git lfs clone $base_model_url\n",
        "\n",
        "from glob import glob\n",
        "base_model_path = glob(f\"{base_model_name.lower()}/{base_model_name}*.pth\")[0]\n",
        "\n",
        "print(f\"Using {base_model_path} as base\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hCOPnLelfJgP"
      },
      "source": [
        "## Generate Training Data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wW5OmlXmvaIU",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title Training Data Options\n",
        "#@markdown `input_file` should be the path to a single file that contains the text you want to fine-tune with.\n",
        "#@markdown Either upload a file to this notebook instance or reference a file in your Google drive.\n",
        "\n",
        "import numpy as np\n",
        "from transformers import PreTrainedTokenizerFast\n",
        "\n",
        "tokenizer = PreTrainedTokenizerFast(tokenizer_file=f'{repo_dir}/20B_tokenizer.json')\n",
        "\n",
        "input_file = \"/content/drive/MyDrive/training.txt\" #@param {type:\"string\"}\n",
        "output_file = 'train.npy'\n",
        "\n",
        "print(f'Tokenizing {input_file} (VERY slow. please wait)')\n",
        "\n",
        "data_raw = open(input_file, encoding=\"utf-8\").read()\n",
        "print(f'Raw length = {len(data_raw)}')\n",
        "\n",
        "data_code = tokenizer.encode(data_raw)\n",
        "print(f'Tokenized length = {len(data_code)}')\n",
        "\n",
        "out = np.array(data_code, dtype='uint16')\n",
        "np.save(output_file, out, allow_pickle=False)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "I4lz-3maeIwY"
      },
      "source": [
        "## Training"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fuCw5_ASwMud"
      },
      "outputs": [],
      "source": [
        "#@title Begin Training with these Options { display-mode: \"form\" }\n",
        "n_epoch = 100 #@param {type:\"integer\"}\n",
        "epoch_save_frequency = 25 #@param {type:\"integer\"}\n",
        "batch_size =  11 #@param {type:\"integer\"} \n",
        "ctx_len = 384 #@param {type:\"integer\"}\n",
        "precision = 'fp16' #@param ['fp16', 'bf16', 'bf32'] {type:\"string\"}\n",
        "\n",
        "epoch_save_path = f\"{output_path}/{tuned_model_name}\"\n",
        "\n",
        "\n",
        "!python train.py \\\n",
        "--load_model $base_model_path \\\n",
        "--wandb \"\" \\\n",
        "--proj_dir $output_dir \\\n",
        "--data_file  \"train.npy\" \\\n",
        "--data_type \"numpy\" \\\n",
        "--vocab_size 50277 \\\n",
        "--ctx_len $ctx_len \\\n",
        "--epoch_steps 1000 \\\n",
        "--epoch_count $n_epoch \\\n",
        "--epoch_begin 0 \\\n",
        "--epoch_save $epoch_save_frequency \\\n",
        "--micro_bsz 8 \\\n",
        "--n_layer $n_layer \\\n",
        "--n_embd $n_embd \\\n",
        "--pre_ffn 0 \\\n",
        "--head_qk 0 \\\n",
        "--lr_init 1e-5 \\\n",
        "--lr_final 1e-5 \\\n",
        "--warmup_steps 0 \\\n",
        "--beta1 0.9 \\\n",
        "--beta2 0.999 \\\n",
        "--adam_eps 1e-8 \\\n",
        "--accelerator gpu \\\n",
        "--devices 1 \\\n",
        "--precision $precision \\\n",
        "--strategy deepspeed_stage_2 \\\n",
        "--grad_cp 0"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "provenance": [],
      "mount_file_id": "1ITfSRrRpVO9Xz-6MYDWVHHWykp02L4qi",
      "authorship_tag": "ABX9TyM/egA8GK0VL1hdCHnm+AQd",
      "include_colab_link": true
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}