RWKV_v2_RNN_Pile_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_v2_RNN_Pile_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-v2-RNN-Pile 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-v2-RNN-Pile](https://github.com/BlinkDL/RWKV-v2-RNN-Pile) as detailed [here](https://github.com/BlinkDL/RWKV-v2-RNN-Pile)"
      ]
    },
    {
      "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-v2-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",
        "\n",
        "drive.mount(drive_mount, force_remount=True)\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",
        "repo_dir = \"/content/RWKV-v2-RNN-Pile\"\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": "S6_RM1m292ZW"
      },
      "outputs": [],
      "source": [
        "if save_models_to_drive:\n",
        "    from google.colab import drive\n",
        "    drive.mount('/content/drive')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "R4lt0FTegJw9"
      },
      "outputs": [],
      "source": [
        "!git clone https://github.com/blinkdl/RWKV-v2-RNN-Pile"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RDavUrBsgKIV"
      },
      "outputs": [],
      "source": [
        "!pip install transformers 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 github\n",
        "\n",
        "base_model_name = \"v2-430M-332B\" #@param [\"v2-430M-332B\", \"v2-430M-235B\", \"v2-430M-123B\", \"v2-430M-57B\", \"v2-430M-50B\"]\n",
        "check_for_existing_model = True #@param {type: \"boolean\"}\n",
        "\n",
        "class BaseModel():\n",
        "    def __init__(self, name, tag, k_eps=1e-9):\n",
        "        self.name = name\n",
        "        self.tag = tag\n",
        "        self.k_eps = k_eps\n",
        "\n",
        "model_options = {\n",
        "    \"v2-430M-332B\": BaseModel(\"v2-430M-332B\", \"20220615-10803\", k_eps=1e-8),\n",
        "    \"v2-430M-235B\": BaseModel(\"v2-430M-235B\", \"20220605-7663\", k_eps=1e-8),\n",
        "    \"v2-430M-123B\": BaseModel(\"v2-430M-123B\", \"20220524-4006\"),\n",
        "    \"v2-430M-57B\":  BaseModel(\"v2-430M-57B\", \"20220515-1853\"),\n",
        "    \"v2-430M-50B\":  BaseModel(\"v2-430M-50B\", \"20220501-6548\"),\n",
        "}\n",
        "\n",
        "#@markdown Alternatively if you want to use a custom base model (for example a previous checkpoint) simply input the path to the model.\n",
        "#@markdown This can be either a `.zip` archive that contains the model or the `.pth` file.\n",
        "#@markdown You should still set the `base_model_name` to whichever model your custom one is based on, as it's used to determine some training variables.\n",
        "\n",
        "use_custom_model = False #@param {type: \"boolean\"}\n",
        "custom_model_path = \"\" #@param {type: \"string\"}\n",
        "\n",
        "import requests, shutil\n",
        "from zipfile import ZipFile\n",
        "\n",
        "base_model = model_options[base_model_name]\n",
        "\n",
        "release_api = \"https://api.github.com/repos/blinkdl/RWKV-v2-RNN-Pile/releases/tags\"\n",
        "\n",
        "base_model_path = \"\"\n",
        "\n",
        "saved_model_dir = f\"{output_path}/base_models/{base_model.name}\"\n",
        "saved_model_path = f\"{saved_model_dir}/{base_model.name}.zip\"\n",
        "\n",
        "# Get model\n",
        "if use_custom_model:\n",
        "    base_model_path = custom_model_path\n",
        "else:\n",
        "    if check_for_existing_model and os.path.exists(saved_model_path):\n",
        "        base_model_path = saved_model_path\n",
        "    if not base_model_path:\n",
        "        os.makedirs(saved_model_dir, exist_ok=True)\n",
        "        r = requests.get(f\"{release_api}/{base_model.tag}\")\n",
        "        url = r.json()[\"assets\"][0][\"browser_download_url\"]\n",
        "        print(f\"Downloading {base_model.name} from {url} this may take a while...\")\n",
        "        try:\n",
        "            with requests.get(url, stream=True) as r:\n",
        "                with open(saved_model_path, 'wb') as f:\n",
        "                    shutil.copyfileobj(r.raw, f)\n",
        "            base_model_path = saved_model_path\n",
        "        except:\n",
        "            os.remove(saved_model_path)\n",
        "            raise RuntimeError(f\"Download from {url} had an unexpected failure\")\n",
        "\n",
        "# Prepare model\n",
        "if base_model_path.endswith(\".zip\"):\n",
        "    print(f\"Unzipping {base_model_path}\")\n",
        "    with ZipFile(base_model_path, 'r') as z:\n",
        "        for file in z.filelist:\n",
        "            if file.filename.endswith(\".pth\"):\n",
        "                z.extract(file)\n",
        "                base_model_path = file.filename\n",
        "                break\n",
        "\n",
        "if not base_model_path.endswith(\".pth\"):\n",
        "    raise RuntimeError(\"base_model must be either a .zip file that contains the model or a .pth file\")\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": {
        "cellView": "form",
        "id": "wW5OmlXmvaIU"
      },
      "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}/RWKV-v2/20B_tokenizer.json')\n",
        "\n",
        "input_file = \"/content/drive/MyDrive/train.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": [
        "## Fine-tune"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fuCw5_ASwMud"
      },
      "outputs": [],
      "source": [
        "#@title Fine-tuning Options { display-mode: \"form\" }\n",
        "\n",
        "#@markdown By default the fine tuning is handled in GPT mode as it trains much faster,\n",
        "#@markdown however it uses much more VRAM.\n",
        "#@markdown\n",
        "#@markdown The suggested settings for training are normally a `ctx_len` of 768, a `batch_size` of 8, and `B_GROUP_FORWARD` being 8.\n",
        "#@markdown However with the limited VRAM you get with a P100 I've found `336`/`4`/`4` doable. \n",
        "#@markdown\n",
        "#@markdown As always your mileage may vary, fiddle with the numbers yourself.\n",
        "#@markdown\n",
        "#@markdown ---\n",
        "#@markdown \n",
        "\n",
        "#@markdown Enable `use_wandb` if you want to track your training run via [Weights & Biases](https://wandb.ai)\n",
        "use_wandb = True #@param {type:\"boolean\"}\n",
        "#@markdown Epochs in this context are really \"mini-epochs\" that are quite short. They have a fixed length of `ctx_len * epoch_length_fixed` tokens\n",
        "n_epoch = 100#@param {type:\"integer\"}\n",
        "epoch_save_frequency = 10 #@param {type:\"integer\"}\n",
        "ctx_len = 336 #@param {type:\"integer\"}\n",
        "#@markdown If finetuning OOMs, consider lowering the batch size before lowering `ctx_len`\n",
        "batch_size =  4#@param {type:\"integer\"} \n",
        "#@markdown `batch_size` must be divisible by both `B_GROUP_FORWARD` and `B_GROUP_BACKWARD`\n",
        "B_GROUP_FORWARD =  4#@param {type:\"integer\"}\n",
        "B_GROUP_BACKWARD =  2#@param {type:\"integer\"}\n",
        "\n",
        "import logging\n",
        "import datetime\n",
        "import json\n",
        "import torch\n",
        "import numpy as np\n",
        "import math\n",
        "import sys\n",
        "\n",
        "import wandb\n",
        "\n",
        "import torch.nn as nn\n",
        "import torch.optim as optim\n",
        "from torch.nn import functional as F\n",
        "from torch.utils.data import Dataset\n",
        "from torch.utils.cpp_extension import load\n",
        "from torch.utils.data.dataloader import DataLoader\n",
        "from torch.optim.lr_scheduler import LambdaLR\n",
        "\n",
        "from tqdm.auto import tqdm\n",
        "\n",
        "torch.backends.cudnn.benchmark = True\n",
        "torch.backends.cudnn.allow_tf32 = True\n",
        "torch.backends.cuda.matmul.allow_tf32 = True\n",
        "torch.cuda.empty_cache()\n",
        "\n",
        "n_layer = 24\n",
        "n_embd = 1024\n",
        "vocab_size = 50277\n",
        "\n",
        "model_type = 'RWKV'\n",
        "datafile = 'train.npy'\n",
        "\n",
        "#@markdown If your training data uses something similar to what is already in [The Pile](https://pile.eleuther.ai/)\n",
        "#@markdown consider setting `lr_init` to 1e-5\n",
        "lr_init = 2e-5#@param {type:\"number\"}\n",
        "lr_final = 1e-5#@param {type:\"number\"}\n",
        "\n",
        "epoch_length_fixed = 10000\n",
        "\n",
        "epoch_save_path = f\"{output_path}/\"\n",
        "\n",
        "grad_norm_clip = 1.0\n",
        "warmup_tokens = 0\n",
        "\n",
        "betas = (0.9, 0.99)\n",
        "eps = 1e-8\n",
        "\n",
        "num_workers = 0\n",
        "\n",
        "logger = logging.getLogger(__name__)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xvV3C1wEgwJJ"
      },
      "outputs": [],
      "source": [
        "#@title model_train.py { display-mode: \"form\" }\n",
        "########################################################################################################\n",
        "# CUDA Kernel\n",
        "########################################################################################################\n",
        "\n",
        "T_MAX = 1024 # increase this if your ctx_len > 1024\n",
        "\n",
        "timex_cuda = load(name=\"timex\", sources=[f\"{repo_dir}/cuda/timex_op.cpp\", f\"{repo_dir}/cuda/timex_cuda.cu\"],\n",
        "                  verbose=True, extra_cuda_cflags=['--use_fast_math', '--extra-device-vectorization', f'-DTmax={T_MAX}', f'-DBF={B_GROUP_FORWARD}', f'-DBB={B_GROUP_BACKWARD}'])\n",
        "\n",
        "\n",
        "class TimeX(torch.autograd.Function):\n",
        "    @staticmethod\n",
        "    def forward(ctx, w, k, B, C, T, eps):\n",
        "        ctx.B = B\n",
        "        ctx.C = C\n",
        "        ctx.T = T\n",
        "        assert ctx.T % 4 == 0 and ctx.T <= T_MAX and ctx.B % B_GROUP_FORWARD == 0 and ctx.B % B_GROUP_BACKWARD == 0\n",
        "        w = w.contiguous()\n",
        "        k = k.contiguous()\n",
        "        ctx.save_for_backward(w, k)\n",
        "        wk = torch.empty((B, C, T), device='cuda',\n",
        "                         memory_format=torch.contiguous_format)\n",
        "        timex_cuda.forward(w, k, wk, eps, B, C, T)\n",
        "        return wk\n",
        "\n",
        "    @staticmethod\n",
        "    def backward(ctx, gwk):\n",
        "        assert ctx.T % 4 == 0 and ctx.T <= T_MAX and ctx.B % B_GROUP_FORWARD == 0 and ctx.B % B_GROUP_BACKWARD == 0\n",
        "        w, k = ctx.saved_tensors\n",
        "        gw = torch.empty((ctx.B, ctx.C, ctx.T), device='cuda',\n",
        "                         memory_format=torch.contiguous_format)\n",
        "        gk = torch.empty((ctx.B, ctx.C, ctx.T), device='cuda',\n",
        "                         memory_format=torch.contiguous_format)\n",
        "        timex_cuda.backward(w, k, gwk.contiguous(), gw,\n",
        "                            gk, ctx.B, ctx.C, ctx.T)\n",
        "        return (gw.sum(dim=0), gk, None, None, None, None)\n",
        "\n",
        "########################################################################################################\n",
        "# RWKV: RWKV Time-mix + RWKV Channel-mix\n",
        "########################################################################################################\n",
        "\n",
        "RWKV_K_CLAMP = 60  # e^60 = 1e26\n",
        "RWKV_K_EPS = base_model.k_eps\n",
        "RWKV_HEAD_QK_DIM = 256\n",
        "\n",
        "class RWKV_TimeMix(nn.Module):\n",
        "    def __init__(self, config, layer_id):\n",
        "        super().__init__()\n",
        "        self.layer_id = layer_id\n",
        "        self.ctx_len = config.ctx_len\n",
        "        self.n_embd = config.n_embd\n",
        "\n",
        "        attn_sz = config.n_embd\n",
        "\n",
        "        ############# fancy init of time_w curves ###################################\n",
        "        f1_begin = 3.0\n",
        "        f1_end = 1.2\n",
        "        f2_begin = 0.65\n",
        "        f2_end = 0.4\n",
        "        with torch.no_grad():  # initial time_w curves for better convergence\n",
        "            decay_speed = torch.ones(attn_sz, 1)\n",
        "            first_sa_layer_id = 1\n",
        "            for h in range(attn_sz):\n",
        "                f1 = f1_begin + (layer_id-first_sa_layer_id) / \\\n",
        "                    (config.n_layer-1-first_sa_layer_id) * (f1_end - f1_begin)\n",
        "                f2 = f2_begin + (layer_id-first_sa_layer_id) / \\\n",
        "                    (config.n_layer-1-first_sa_layer_id) * (f2_end - f2_begin)\n",
        "                if layer_id == first_sa_layer_id:\n",
        "                    f1 += 0.5\n",
        "                if layer_id == config.n_layer-2:\n",
        "                    f2 = 0.4\n",
        "                if layer_id == config.n_layer-1:\n",
        "                    f2 = 0.37\n",
        "                decay_speed[h][0] = math.pow(f2, h / (attn_sz-1) * 7) * f1\n",
        "        self.time_decay = nn.Parameter(torch.log(decay_speed)) # will use exp(self.time_decay) to ensure time_decay > 0\n",
        "        self.time_curve = torch.tensor(\n",
        "            [-(config.ctx_len - 2 - i) for i in range(config.ctx_len-1)]).unsqueeze(0)\n",
        "        self.time_curve = self.time_curve.to('cuda')\n",
        "        self.time_first = nn.Parameter(torch.ones(attn_sz, 1) * math.log(0.3))\n",
        "        #############################################################################\n",
        "\n",
        "        self.time_shift = nn.ZeroPad2d((0, 0, 1, -1))\n",
        "        with torch.no_grad():  # init to \"shift half of the channels\"\n",
        "            ww = torch.ones(1, 1, config.n_embd)\n",
        "            for i in range(config.n_embd // 2):\n",
        "                ww[0, 0, i] = 0\n",
        "        self.time_mix = nn.Parameter(ww)\n",
        "\n",
        "        self.key = nn.Linear(config.n_embd, attn_sz, bias=False)\n",
        "        self.value = nn.Linear(config.n_embd, attn_sz, bias=False)\n",
        "        self.receptance = nn.Linear(config.n_embd, attn_sz, bias=False)\n",
        "\n",
        "        self.output = nn.Linear(attn_sz, config.n_embd, bias=False)\n",
        "\n",
        "        self.key.scale_init = 0\n",
        "        self.receptance.scale_init = 0\n",
        "        self.output.scale_init = 0\n",
        "\n",
        "    def forward(self, x):\n",
        "        B, T, C = x.size()\n",
        "        assert T == self.ctx_len\n",
        "\n",
        "        x = x * self.time_mix + self.time_shift(x) * (1 - self.time_mix)\n",
        "\n",
        "        k = self.key(x).transpose(-1, -2)\n",
        "        v = self.value(x).transpose(-1, -2)\n",
        "        r = self.receptance(x)\n",
        "\n",
        "        # RWKV_K_CLAMP can be removed if the CUDA kernel substracts the correct k_max for each k (I will do this later)\n",
        "        k = torch.clamp(k, max=RWKV_K_CLAMP)\n",
        "        k = torch.exp(k)\n",
        "        kv = k * v\n",
        "\n",
        "        self.time_w = torch.cat(\n",
        "            [torch.exp(self.time_decay) * self.time_curve, self.time_first], dim=-1)\n",
        "        w = torch.exp(self.time_w)\n",
        "\n",
        "        wkv = TimeX.apply(w, kv, B, C, T, 0)\n",
        "        # RWKV_K_EPS can be removed if the CUDA kernel sets 0/0 = 0 (I will do this later)\n",
        "        wk = TimeX.apply(w, k, B, C, T, RWKV_K_EPS)\n",
        "\n",
        "        rwkv = torch.sigmoid(r) * (wkv / wk).transpose(-1, -2)\n",
        "        rwkv = self.output(rwkv)\n",
        "        return rwkv\n",
        "\n",
        "\n",
        "class RWKV_ChannelMix(nn.Module):\n",
        "    def __init__(self, config, layer_id):\n",
        "        super().__init__()\n",
        "        self.layer_id = layer_id\n",
        "\n",
        "        self.time_shift = nn.ZeroPad2d((0, 0, 1, -1))\n",
        "\n",
        "        with torch.no_grad():  # init to \"shift half of the channels\"\n",
        "            x = torch.ones(1, 1, config.n_embd)\n",
        "            for i in range(config.n_embd // 2):\n",
        "                x[0, 0, i] = 0\n",
        "        self.time_mix = nn.Parameter(x)\n",
        "\n",
        "        hidden_sz = 4 * config.n_embd\n",
        "        self.key = nn.Linear(config.n_embd, hidden_sz, bias=False)\n",
        "        self.receptance = nn.Linear(config.n_embd, config.n_embd, bias=False)\n",
        "        self.value = nn.Linear(hidden_sz, config.n_embd, bias=False)\n",
        "\n",
        "        self.value.scale_init = 0\n",
        "        self.receptance.scale_init = 0\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = x * self.time_mix + self.time_shift(x) * (1 - self.time_mix)\n",
        "\n",
        "        k = self.key(x)\n",
        "        k = torch.square(torch.relu(k))\n",
        "        kv = self.value(k)\n",
        "\n",
        "        rkv = torch.sigmoid(self.receptance(x)) * kv\n",
        "        return rkv\n",
        "\n",
        "########################################################################################################\n",
        "# The GPT Model with our blocks\n",
        "########################################################################################################\n",
        "\n",
        "\n",
        "class GPTConfig:\n",
        "    def __init__(self, vocab_size, ctx_len, **kwargs):\n",
        "        self.vocab_size = vocab_size\n",
        "        self.ctx_len = ctx_len\n",
        "        for k, v in kwargs.items():\n",
        "            setattr(self, k, v)\n",
        "\n",
        "\n",
        "class Block(nn.Module):\n",
        "    def __init__(self, config, layer_id):\n",
        "        super().__init__()\n",
        "        self.config = config\n",
        "        self.layer_id = layer_id\n",
        "\n",
        "        self.ln1 = nn.LayerNorm(config.n_embd)\n",
        "        self.ln2 = nn.LayerNorm(config.n_embd)\n",
        "\n",
        "        if self.layer_id == 0 and self.config.model_type == 'RWKV-ffnPre':\n",
        "            self.ffnPre = RWKV_ChannelMix(config, layer_id+1000)\n",
        "        else:\n",
        "            self.att = RWKV_TimeMix(config, layer_id)\n",
        "\n",
        "        self.ffn = RWKV_ChannelMix(config, layer_id)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = self.ln1(x)\n",
        "        if self.layer_id == 0 and self.config.model_type == 'RWKV-ffnPre':\n",
        "            x = x + self.ffnPre(x)  # better in some cases\n",
        "        else:\n",
        "            x = x + self.att(x)\n",
        "        x = self.ln2(x)\n",
        "        x = x + self.ffn(x)\n",
        "        return x\n",
        "\n",
        "\n",
        "class GPT(nn.Module):\n",
        "    def __init__(self, config):\n",
        "        super().__init__()\n",
        "        self.step = 0\n",
        "        self.config = config\n",
        "\n",
        "        self.emb = nn.Embedding(config.vocab_size, config.n_embd)\n",
        "\n",
        "        self.blocks = nn.Sequential(*[Block(config, i)\n",
        "                                    for i in range(config.n_layer)])\n",
        "\n",
        "        self.ln_out = nn.LayerNorm(config.n_embd)\n",
        "        self.head = nn.Linear(config.n_embd, config.vocab_size, bias=False)\n",
        "\n",
        "        # self.head_q = nn.Linear(config.n_embd, RWKV_HEAD_QK_DIM, bias=False)\n",
        "        # self.head_q.scale_init = 0\n",
        "        # self.head_k = nn.Linear(config.n_embd, RWKV_HEAD_QK_DIM, bias=False)\n",
        "        # self.head_k.scale_init = 0.1\n",
        "        # self.register_buffer(\"copy_mask\", torch.tril(\n",
        "        #     torch.ones(config.ctx_len, config.ctx_len)))\n",
        "\n",
        "        self.ctx_len = config.ctx_len\n",
        "\n",
        "        # RWKV_Init(self, config)\n",
        "\n",
        "        logger.info(\"number of parameters: %e\", sum(p.numel()\n",
        "                    for p in self.parameters()))\n",
        "\n",
        "    def get_ctx_len(self):\n",
        "        return self.ctx_len\n",
        "\n",
        "    def _init_weights(self, module):\n",
        "        if isinstance(module, (nn.Linear)):\n",
        "            module.weight.data.normal_(mean=0.0, std=0.01)\n",
        "        if isinstance(module, (nn.Embedding)):\n",
        "            module.weight.data.normal_(mean=0.0, std=1e-5)\n",
        "        if isinstance(module, nn.Linear) and module.bias is not None:\n",
        "            module.bias.data.zero_()\n",
        "\n",
        "    def configure_optimizers(self, train_config):\n",
        "        # separate out all parameters to those that will and won't experience regularizing weight decay\n",
        "        decay = set()\n",
        "        no_decay = set()\n",
        "\n",
        "        for mn, m in self.named_modules():  # here we disable weight_decay\n",
        "            for pn, p in m.named_parameters():\n",
        "                fpn = '%s.%s' % (mn, pn) if mn else pn  # full param name\n",
        "                no_decay.add(fpn)\n",
        "\n",
        "        param_dict = {pn: p for pn, p in self.named_parameters()}\n",
        "        inter_params = decay & no_decay\n",
        "        union_params = decay | no_decay\n",
        "        assert len(\n",
        "            inter_params) == 0, \"parameters %s made it into both decay/no_decay sets!\" % (str(inter_params), )\n",
        "        assert len(param_dict.keys() - union_params) == 0, \"parameters %s were not separated into either decay/no_decay set!\" \\\n",
        "            % (str(param_dict.keys() - union_params), )\n",
        "\n",
        "        optim_groups = [\n",
        "            {\"params\": [param_dict[pn]\n",
        "                        for pn in sorted(list(no_decay))], \"weight_decay\": 0.0},\n",
        "        ]\n",
        "\n",
        "        optimizer = torch.optim.Adam(\n",
        "            optim_groups, lr=train_config.learning_rate, betas=train_config.betas, eps=train_config.eps)\n",
        "\n",
        "        return optimizer\n",
        "\n",
        "    def forward(self, idx, targets=None):\n",
        "        self.step += 1\n",
        "        B, T = idx.size()\n",
        "        assert T <= self.ctx_len, \"Cannot forward, because len(input) > model ctx_len.\"\n",
        "        x = self.emb(idx)\n",
        "        x = self.blocks(x)\n",
        "        x = self.ln_out(x)\n",
        "        x = self.head(x)\n",
        "        # q = self.head_q(x)[:, :T, :]\n",
        "        # k = self.head_k(x)[:, :T, :]\n",
        "        # c = (q @ k.transpose(-2, -1)) * (1.0 / RWKV_HEAD_QK_DIM)\n",
        "        # c = c.masked_fill(self.copy_mask[:T, :T] == 0, 0)\n",
        "\n",
        "        # c = c @ F.one_hot(idx, num_classes=self.config.vocab_size).float()\n",
        "        # x = self.head(x) + c\n",
        "\n",
        "        loss = None\n",
        "        if targets is not None:\n",
        "            loss = F.cross_entropy(x.view(-1, x.size(-1)), targets.view(-1))\n",
        "\n",
        "        return x, loss"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "EHIJp8vmmWX8"
      },
      "outputs": [],
      "source": [
        "#@title trainer.py { display-mode: \"form\" }\n",
        "\n",
        "log_file = open(\"mylog.txt\", \"a\")\n",
        "\n",
        "class TrainerConfig:\n",
        "    max_epochs = 10\n",
        "    batch_size = 64\n",
        "    learning_rate = 4e-4\n",
        "    betas = (0.9, 0.99)\n",
        "    eps = 1e-8\n",
        "    grad_norm_clip = 1.0\n",
        "    lr_decay = True  # linear warmup followed by cosine decay\n",
        "    warmup_tokens = 0\n",
        "    final_tokens = 0\n",
        "    epoch_save_frequency = 0\n",
        "    epoch_save_path = 'trained-'\n",
        "    num_workers = 0  # for DataLoader\n",
        "\n",
        "    def __init__(self, **kwargs):\n",
        "        for k, v in kwargs.items():\n",
        "            setattr(self, k, v)\n",
        "\n",
        "\n",
        "class Trainer:\n",
        "\n",
        "    def __init__(self, model, train_dataset, test_dataset, config):\n",
        "        self.model = model\n",
        "        self.train_dataset = train_dataset\n",
        "        self.test_dataset = test_dataset\n",
        "        self.config = config\n",
        "        self.avg_loss = -1\n",
        "        self.steps = 0\n",
        "\n",
        "        if use_wandb and 'wandb' in sys.modules:\n",
        "            cfg = model.config\n",
        "            for k in config.__dict__:\n",
        "                setattr(cfg, k, config.__dict__[k])  # combine cfg\n",
        "            wandb.init(project=\"RWKV-LM\", name=f\"{tuned_model_name}-{self.get_run_name()}\", config=cfg, save_code=False)\n",
        "\n",
        "        self.device = 'cpu'\n",
        "        if torch.cuda.is_available():  # take over whatever gpus are on the system\n",
        "            self.device = torch.cuda.current_device()\n",
        "\n",
        "    def get_run_name(self):\n",
        "        raw_model = self.model.module if hasattr(\n",
        "            self.model, \"module\") else self.model\n",
        "        cfg = raw_model.config\n",
        "        run_name = str(cfg.vocab_size) + '-' + str(cfg.ctx_len) + '-' + str(cfg.n_layer) + '-' + str(cfg.n_embd)\n",
        "        return run_name\n",
        "\n",
        "    def train(self):\n",
        "        model, config = self.model, self.config\n",
        "        raw_model = model.module if hasattr(self.model, \"module\") else model\n",
        "        optimizer = raw_model.configure_optimizers(config)\n",
        "\n",
        "        def run_epoch(split):\n",
        "            is_train = split == 'train'\n",
        "            model.train(is_train)\n",
        "            data = self.train_dataset if is_train else self.test_dataset\n",
        "\n",
        "            if config.num_workers > 0:\n",
        "                loader = DataLoader(data, shuffle=False, pin_memory=True,\n",
        "                                    batch_size=config.batch_size,\n",
        "                                    num_workers=config.num_workers)\n",
        "            else:\n",
        "                loader = DataLoader(data, shuffle=False,\n",
        "                                    batch_size=config.batch_size,\n",
        "                                    num_workers=config.num_workers)\n",
        "\n",
        "            pbar = tqdm(enumerate(loader), total=len(\n",
        "                loader), bar_format='{l_bar}{bar:10}{r_bar}{bar:-10b}') if is_train else enumerate(loader)\n",
        "\n",
        "            for it, (x, y) in pbar:\n",
        "                x = x.to(self.device)  # place data on the correct device\n",
        "                y = y.to(self.device)\n",
        "\n",
        "                with torch.set_grad_enabled(is_train):\n",
        "                    _, loss = model(x, y)  # forward the model\n",
        "\n",
        "                if is_train:  # backprop and update the parameters\n",
        "                    model.zero_grad()\n",
        "                    loss.backward()\n",
        "\n",
        "                    if config.grad_norm_clip > 0:\n",
        "                        torch.nn.utils.clip_grad_norm_(\n",
        "                            model.parameters(), config.grad_norm_clip)\n",
        "\n",
        "                    optimizer.step()\n",
        "\n",
        "                    if config.lr_decay:  # decay the learning rate based on our progress\n",
        "                        # number of tokens processed this step (i.e. label is not -100)\n",
        "                        self.tokens += (y >= 0).sum()\n",
        "                        lr_final_factor = config.lr_final / config.learning_rate\n",
        "                        if self.tokens < config.warmup_tokens:\n",
        "                            # linear warmup\n",
        "                            lr_mult = lr_final_factor + \\\n",
        "                                (1 - lr_final_factor) * float(self.tokens) / \\\n",
        "                                float(config.warmup_tokens)\n",
        "                            progress = 0\n",
        "                        else:\n",
        "                            # cosine learning rate decay\n",
        "                            progress = float(self.tokens - config.warmup_tokens) / float(\n",
        "                                max(1, config.final_tokens - config.warmup_tokens))\n",
        "                            lr_mult = (0.5 + lr_final_factor / 2) + (0.5 - lr_final_factor /\n",
        "                                                                     2) * math.cos(math.pi * progress)  # better 1.0 ~ 0.1\n",
        "                        lr = config.learning_rate * lr_mult\n",
        "                        for param_group in optimizer.param_groups:\n",
        "                            param_group['lr'] = lr\n",
        "                    else:\n",
        "                        lr = config.learning_rate\n",
        "\n",
        "                    now_loss = loss.item()  # report progress\n",
        "                    self.lr = lr\n",
        "\n",
        "                    if use_wandb and 'wandb' in sys.modules:\n",
        "                        wandb.log({\"loss\": now_loss},\n",
        "                                  step=self.steps * self.config.batch_size)\n",
        "                    self.steps += 1\n",
        "\n",
        "                    if self.avg_loss < 0:\n",
        "                        self.avg_loss = now_loss\n",
        "                    else:\n",
        "                        factor = 1 / (it + 1)\n",
        "                        self.avg_loss = self.avg_loss * \\\n",
        "                            (1.0 - factor) + now_loss * factor\n",
        "                    pbar.set_description(\n",
        "                        f\"mini-epoch {epoch+1} prog {progress*100.0:.2f}% iter {it}: ppl {math.exp(self.avg_loss):.2f} loss {self.avg_loss:.4f} lr {lr:e}\")\n",
        "\n",
        "        self.tokens = 0  # counter used for learning rate decay\n",
        "        for epoch in range(config.max_epochs):\n",
        "\n",
        "            run_epoch('train')\n",
        "\n",
        "            log_file.write(\n",
        "                f'{epoch+1} {self.avg_loss:.6f} {math.exp(self.avg_loss):.4f} {self.lr:.8f} {datetime.datetime.now()} \\n')\n",
        "            log_file.flush()\n",
        "\n",
        "            if (self.config.epoch_save_frequency > 0 and epoch % self.config.epoch_save_frequency == 0) or (epoch == config.max_epochs - 1):\n",
        "                # DataParallel wrappers keep raw model object in .module\n",
        "                raw_model = self.model.module if hasattr(\n",
        "                    self.model, \"module\") else self.model\n",
        "                torch.save(raw_model.state_dict(), f\"{output_path}/{tuned_model_name}/{tuned_model_name}-{base_model.name}-{trainer.get_run_name()}-{epoch}.pth\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "uyCW0Wp4gk4p"
      },
      "outputs": [],
      "source": [
        "#@title Load training data { display-mode: \"form\" }\n",
        "\n",
        "class Dataset(Dataset):\n",
        "    def __init__(self, data, vocab_size, ctx_len, epoch_length_fixed):\n",
        "        data_size, vocab_size = len(data), vocab_size\n",
        "        print('data has %d tokens, %d unique.' % (data_size, vocab_size))\n",
        "        self.ctx_len = ctx_len\n",
        "        self.epoch_length_fixed = epoch_length_fixed\n",
        "        self.vocab_size = vocab_size\n",
        "        self.data = data\n",
        "\n",
        "    def __len__(self):\n",
        "        return self.epoch_length_fixed\n",
        "\n",
        "    def __getitem__(self, idx):\n",
        "        i = np.random.randint(0, len(self.data) - (self.ctx_len + 1))\n",
        "        dix = self.data[i:i+self.ctx_len+1]\n",
        "        x = torch.tensor(dix[:-1], dtype=torch.long,\n",
        "                         device=torch.device('cuda'))\n",
        "        y = torch.tensor(dix[1:], dtype=torch.long,\n",
        "                         device=torch.device('cuda'))\n",
        "        return x, y\n",
        "\n",
        "print('loading data... ' + datafile)\n",
        "train_dataset = Dataset(np.load(datafile).astype('int'), vocab_size, ctx_len, epoch_length_fixed)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "c5nfvk9lkHo3"
      },
      "outputs": [],
      "source": [
        "#@title Start fine-tuning { display-mode: \"form\" }\n",
        "\n",
        "np.set_printoptions(precision=4, suppress=True, linewidth=200)\n",
        "logging.basicConfig(format=\"%(asctime)s - %(levelname)s - %(name)s - %(message)s\",\n",
        "                    datefmt=\"%Y-%m-%d %H:%M:%S\", level=logging.INFO,)\n",
        "\n",
        "model = GPT(GPTConfig(train_dataset.vocab_size, train_dataset.ctx_len, model_type=model_type,\n",
        "                      n_layer=n_layer, n_embd=n_embd)).cuda()\n",
        "\n",
        "print('loading ' + base_model_path)\n",
        "m2 = torch.load(base_model_path)\n",
        "model.load_state_dict(m2)\n",
        "\n",
        "print('model', model_type, 'epoch', n_epoch, 'batchsz', batch_size, 'betas',\n",
        "      betas, 'eps', eps, 'ctx', ctx_len, 'layer', n_layer, 'embd', n_embd, )\n",
        "tconf = TrainerConfig(model_type=model_type, max_epochs=n_epoch, batch_size=batch_size,\n",
        "                      learning_rate=lr_init, lr_decay=True, lr_final=lr_final, betas=betas, eps=eps, grad_norm_clip=grad_norm_clip,\n",
        "                      warmup_tokens=warmup_tokens, final_tokens=n_epoch*len(train_dataset)*ctx_len, num_workers=num_workers, epoch_save_frequency=epoch_save_frequency, epoch_save_path=epoch_save_path)\n",
        "trainer = Trainer(model, train_dataset, None, tconf)\n",
        "\n",
        "trainer.train()\n",
        "\n",
        "tuned_model_path = f\"{output_path}/{tuned_model_name}/{tuned_model_name}-{base_model.name}-{trainer.get_run_name()}.pth\"\n",
        "torch.save(model.state_dict(), tuned_model_path)"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Test fine-tuned model (Inference)"
      ],
      "metadata": {
        "id": "JemAl28e9Ah6"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Load GPT model\n",
        "model_gpt = GPT(GPTConfig(vocab_size, ctx_len, n_embd=n_embd, model_type='RWKV', n_layer=n_layer)).cuda()\n",
        "loaded = torch.load(tuned_model_path, map_location='cuda')\n",
        "model_gpt.load_state_dict(loaded)"
      ],
      "metadata": {
        "cellView": "form",
        "id": "MIcbnFSpYo4f"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Run Inference\n",
        "#@markdown Infrerence continues until it reaches your models contex length\n",
        "#@markdown (this takes into account the lenght of your initial prompt).\n",
        "#@markdown Simply rerun this cell for new output, no need to run anything else.\n",
        "#@markdown\n",
        "#@markdown If you've already fine-tuned a model or simply want to run inference on one of the base models. You can use this inference [notebook](https://colab.research.google.com/drive/1TMq2bHfRw2TyCY582yisk8yK9B8zZKhL#scrollTo=VXYK73bc9FcR).\n",
        "prompt = \"\" #@param {type:\"string\"}\n",
        "temperature = 1.0 #@param {type:\"number\"}\n",
        "top_p = 0.7 #@param {type:\"number\"}\n",
        "\n",
        "ctx = tokenizer.encode(prompt)\n",
        "ids = torch.zeros(1, ctx_len, dtype=int, device='cuda')\n",
        "ids[:,:len(ctx)] = torch.tensor(ctx)\n",
        "print(prompt, end='')\n",
        "\n",
        "def sample_logits(logits, temperature=1.0, top_p=0.7):\n",
        "    probs = F.softmax(logits, dim=-1)\n",
        "    sorted_probs, _ = torch.sort(probs, descending=True)\n",
        "    cumulative_probs = torch.cumsum(sorted_probs, dim=-1)\n",
        "    cutoff = float(torch.argmax(sorted_probs[cumulative_probs > top_p]))\n",
        "    probs[probs < cutoff] = 0\n",
        "\n",
        "    if temperature != 1.0:\n",
        "        probs = probs.pow(1.0 / temperature)\n",
        "\n",
        "    return torch.multinomial(probs, num_samples=1)[0]\n",
        "\n",
        "for i in range(len(ctx), ctx_len):\n",
        "    with torch.no_grad():\n",
        "        logits = model_gpt(ids.expand(B_GROUP_FORWARD, ctx_len))[0][0,i-1,:]\n",
        "    token = sample_logits(logits, temperature, top_p)\n",
        "    print(tokenizer.decode(token), end='')\n",
        "    ids[:,i] = token"
      ],
      "metadata": {
        "cellView": "form",
        "id": "VXYK73bc9FcR"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "RWKV-v2-RNN-Pile Fine-Tuning",
      "provenance": [],
      "private_outputs": true,
      "mount_file_id": "1ITfSRrRpVO9Xz-6MYDWVHHWykp02L4qi",
      "authorship_tag": "ABX9TyNdQrqKPap4hSCjQhahxKHp",
      "include_colab_link": true
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}