RWKV_v4neo_Trainer_STK.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "private_outputs": true,
      "provenance": [],
      "authorship_tag": "ABX9TyMwD5p1C1TxilKKDAQ2k1Kd",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU",
    "gpuClass": "standard"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/mrsteyk/RWKV-LM-deepspeed/blob/master/RWKV_v4neo_Trainer_STK.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# **RWKV-v4neo finetuner**\n",
        "\n",
        "**THIS IS A COLAB NOTEBOOK USE LOCAL SCRIPTS IF YOU RUN LOCALLY**\n",
        "\n",
        "Certified `Your session crashed after using all available RAM.` moment\n",
        "\n",
        "This colab is a port of [https://github.com/mrsteyk/RWKV-LM-deepspeed](https://github.com/mrsteyk/RWKV-LM-deepspeed)"
      ],
      "metadata": {
        "id": "NWQZtHqEguTx"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pVvKh-nNgn-R",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title Prereqs\n",
        "%cd /content\n",
        "!git clone --depth 1 https://github.com/mrsteyk/RWKV-LM-deepspeed.git\n",
        "%cd RWKV-LM-deepspeed\n",
        "!git pull\n",
        "%cd RWKV-v4neo\n",
        "\n",
        "!pip install deepspeed pytorch_lightning transformers psutil"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Tokenize your stuff\n",
        "import numpy as np\n",
        "\n",
        "from transformers import GPTNeoXTokenizerFast\n",
        "tokenizer = GPTNeoXTokenizerFast.from_pretrained(\"EleutherAI/gpt-neox-20b\")\n",
        "\n",
        "input_file = 'train.txt'\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)"
      ],
      "metadata": {
        "id": "XND3fCZdhqXU",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!wget https://huggingface.co/BlinkDL/rwkv-4-pile-430m/resolve/main/RWKV-4-Pile-430M-20220808-8066.pth"
      ],
      "metadata": {
        "id": "hMDI5uE1oc6h"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Initial Imports\n",
        "import types\n",
        "\n",
        "import deepspeed\n",
        "import os\n",
        "import torch\n",
        "import torch.utils.data\n",
        "from torch.utils.data import DataLoader\n",
        "from pytorch_lightning import Trainer\n",
        "from pytorch_lightning.utilities import rank_zero_info\n",
        "from pytorch_lightning.callbacks import DeviceStatsMonitor, ModelCheckpoint\n",
        "\n",
        "from deepspeed.utils.zero_to_fp32 import load_state_dict_from_zero_checkpoint\n",
        "\n",
        "import dataset\n",
        "import lr_warmup\n",
        "from soft_embedding_hotswap import SoftEmbedding"
      ],
      "metadata": {
        "id": "LukqgBpchiBW",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Arguments\n",
        "\n",
        "args = types.SimpleNamespace()\n",
        "\n",
        "args.vocab_size_delta = 1\n",
        "args.allgather_bucket_size = 200\n",
        "args.reduce_bucket_size = 200\n",
        "\n",
        "args.data_file = \"train.npy\"\n",
        "\n",
        "args.batch_size = 2\n",
        "\n",
        "args.load_model_cont = '' # ZeRO checkpoint\n",
        "\n",
        "args.soft_emb_tune = False\n",
        "args.soft_emb_tokens = 50\n",
        "\n",
        "args.load_model_init = './RWKV-4-Pile-430M-20220808-8066.pth' # Initialise weights with this\n",
        "args.layerwise_lr = True\n",
        "args.ctx_len = 1024\n",
        "\n",
        "# For Soft Embeddings try using larger lr and epsilon (0.1 and 1e-6)\n",
        "args.lr_init = 1e-5 # 6e-4 for L12-D768, 4e-4 for L24-D1024, 3e-4 for L24-D2048\n",
        "args.lr_final = 1e-5\n",
        "args.warmup_steps = 0 # try 50 if you load a model\n",
        "args.beta1 = 0.9\n",
        "args.beta2 = 0.999 # def: 0.99; use 0.999 when your model is close to convergence\n",
        "args.adam_eps = 1e-8\n",
        "\n",
        "\n",
        "# Hyperparameters\n",
        "args.vocab_size = 50277\n",
        "args.n_layer = 24\n",
        "args.n_embd = 1024\n",
        "args.pre_ffn = False\n",
        "args.head_qk = 0\n",
        "\n",
        "args.tiny_att_dim = 0\n",
        "args.tiny_att_layer = -999\n",
        "\n",
        "# Trainer stuff\n",
        "args.accelerator = \"gpu\"\n",
        "args.devices = 1\n",
        "args.precision = 16 #\"bf16\" # Do T4's support bf16?\n",
        "args.strategy = 'single_device'"
      ],
      "metadata": {
        "id": "MUEASp15hj0d"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Trainer prologue (if you change context len you need to restart)\n",
        "args.betas = (args.beta1, args.beta2)\n",
        "rank_zero_info(args)\n",
        "\n",
        "assert args.precision in [32, 16, \"bf16\"]\n",
        "if args.precision == 16:\n",
        "    os.environ[\"RWKV_FLOAT_MODE\"] = \"fp16\"\n",
        "else:\n",
        "    os.environ[\"RWKV_FLOAT_MODE\"] = str(args.precision)\n",
        "os.environ[\"RWKV_T_MAX\"] = str(args.ctx_len + (args.soft_emb_tokens if args.soft_emb_tune else 0))\n",
        "\n",
        "# Now we can import the model after setting that stupid T max envvar\n",
        "import model as M\n",
        "model = M.RWKV(args)\n",
        "# model = None\n",
        "\n",
        "if args.load_model_cont != '' and not args.soft_emb_tune:\n",
        "    # load_state_dict_from_zero_checkpoint(model, args.load_model_cont)\n",
        "    pass\n",
        "elif args.load_model_init != '':\n",
        "    if os.path.isdir(args.load_model_init):\n",
        "        load_state_dict_from_zero_checkpoint(model, args.load_model_init)\n",
        "        model.cpu()\n",
        "        if args.precision == 16:\n",
        "            model.half()\n",
        "        elif args.precision == \"bf16\":\n",
        "            model.bfloat16()\n",
        "    else:\n",
        "        d = torch.load(args.load_model_init, map_location='cpu')\n",
        "        if list(d.keys())[0].startswith(\"_forward_module.\"):\n",
        "            d = {n[len(\"_forward_module.\"):]: d[n] for n in d.keys()}\n",
        "        model.load_state_dict(d)\n",
        "    # model = M.RWKV(args).load_from_checkpoint(args.load_model_init)\n",
        "else:\n",
        "    # TODO?\n",
        "    # model = M.RWKV(args)\n",
        "    model.generate_init_weight()\n",
        "\n",
        "if args.vocab_size_delta > 0:\n",
        "    new_vocab_size = args.vocab_size + args.vocab_size_delta\n",
        "    model.resize_emb(new_vocab_size)\n",
        "    args.vocab_size = new_vocab_size\n",
        "\n",
        "if args.soft_emb_tune:\n",
        "    # meme hard, die young\n",
        "    print(\"### буду погибать молодым/малоДЫМ(а)\")\n",
        "    args.layerwise_lr = False\n",
        "    for p in model.parameters():\n",
        "        p.requires_grad = False\n",
        "    model.emb_hotswap = True\n",
        "    assert args.soft_emb_tokens < args.vocab_size, \"Soft Embedding can't eat more than the `emb`\"\n",
        "    model.emb = SoftEmbedding(model.emb, n_tokens=args.soft_emb_tokens, initialize_from_vocab=True)\n",
        "\n",
        "lr_meme = lr_warmup.LearningWarmUpCallback(args)\n",
        "device_stats = DeviceStatsMonitor(cpu_stats=True)\n",
        "val_loss_checkpointing = ModelCheckpoint(\n",
        "    filename=\"epoch-{epoch:02d}-val_loss-{val_loss:.2f}\",\n",
        "    # save_on_train_epoch_end=True,\n",
        "    # save_weights_only=True,\n",
        "    save_top_k=3,\n",
        "    mode='min',\n",
        "    monitor=\"val_loss\",\n",
        "    auto_insert_metric_name=False,\n",
        ")\n",
        "epoch_checkpointing = ModelCheckpoint(\n",
        "    filename=\"epoch-{epoch:02d}\",\n",
        "    save_on_train_epoch_end=True,\n",
        "    save_top_k=1,\n",
        "    mode='max',\n",
        "    monitor=\"epoch\",\n",
        "    auto_insert_metric_name=False,\n",
        ")\n",
        "\n",
        "trainer = Trainer.from_argparse_args(\n",
        "    args,\n",
        "    callbacks=[lr_meme, device_stats, val_loss_checkpointing, epoch_checkpointing],\n",
        ")\n",
        "if \"deepspeed\" in args.strategy:\n",
        "    trainer.strategy.config[\"zero_optimization\"][\"allgather_bucket_size\"] = args.allgather_bucket_size * 1e6\n",
        "    trainer.strategy.config[\"zero_optimization\"][\"reduce_bucket_size\"] = args.reduce_bucket_size * 1e6\n",
        "    rank_zero_info(trainer.strategy.config)\n",
        "\n",
        "if \"single_device\" == args.strategy:\n",
        "    trainer.strategy._root_device = torch.device('cuda:0')\n",
        "\n",
        "train_data = dataset.MyDataSet(args)\n",
        "\n",
        "# TODO(mrsteyk): Allow different validation files\n",
        "# use 20% of training data for validation\n",
        "train_set_size = int(len(train_data) * 0.8)\n",
        "valid_set_size = len(train_data) - train_set_size\n",
        "\n",
        "# split the train set into two\n",
        "seed = torch.Generator().manual_seed(42)\n",
        "train_data, valid_data = torch.utils.data.random_split(train_data, [train_set_size, valid_set_size], generator=seed)\n",
        "\n",
        "# data_loader = DataLoader(train_data, shuffle=False, pin_memory=True, batch_size=args.batch_size, num_workers=1, persistent_workers=False, drop_last=True)\n",
        "train_loader = DataLoader(train_data, shuffle=True, pin_memory=True, batch_size=args.batch_size)\n",
        "valid_loader = DataLoader(valid_data, shuffle=False, pin_memory=True, batch_size=args.batch_size)"
      ],
      "metadata": {
        "id": "jr7cv9OTkstX"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Train\n",
        "\n",
        "%load_ext tensorboard\n",
        "%tensorboard --logdir /content/RWKV-LM-deepspeed/RWKV-v4neo/lightning_logs\n",
        "\n",
        "model.cuda()\n",
        "trainer.fit(model, train_loader, valid_loader, ckpt_path=args.load_model_cont if args.load_model_cont != ''  else None)"
      ],
      "metadata": {
        "id": "mmc2Yw7Zlo12"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}