{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "d16d19a7",
   "metadata": {
    "papermill": {
     "duration": 0.004251,
     "end_time": "2023-09-06T01:53:47.272181",
     "exception": false,
     "start_time": "2023-09-06T01:53:47.267930",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# RWKV v5\n",
    "\n",
    "Simple memory training for a small model\n",
    "\n",
    "**Note:** This project assumes you have the rwkv-infctx conda env setup"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "e7eaa6b5",
   "metadata": {
    "papermill": {
     "duration": 0.002635,
     "end_time": "2023-09-06T01:53:47.277736",
     "exception": false,
     "start_time": "2023-09-06T01:53:47.275101",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Basic Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "43c42c12",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-09-06T01:53:47.281975Z",
     "iopub.status.busy": "2023-09-06T01:53:47.281689Z",
     "iopub.status.idle": "2023-09-06T01:53:48.160757Z",
     "shell.execute_reply": "2023-09-06T01:53:48.159854Z"
    },
    "papermill": {
     "duration": 0.883492,
     "end_time": "2023-09-06T01:53:48.162633",
     "exception": false,
     "start_time": "2023-09-06T01:53:47.279141",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CITATION.cff  RWKV-v4wavenet\t RWKV-v5headsize32  checkpoint\tnotebook\r\n",
      "LICENSE       RWKV-v5\t\t RWKV-v5r2\t    datapath\toutput\r\n",
      "README.md     RWKV-v5altwavenet  RWKV-v5rstack\t    docker\r\n",
      "RWKV-v4neo    RWKV-v5headsize2x  RWKV-v5wavenet     model\r\n"
     ]
    }
   ],
   "source": [
    "# First lets setup the various directories, and init the model\n",
    "!ls ../../../../../\n",
    "!mkdir -p ../../../../../model/\n",
    "!mkdir -p ../../../../../datapath/\n",
    "!mkdir -p ../../../../../checkpoint/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "4f145f3c",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-09-06T01:53:48.169937Z",
     "iopub.status.busy": "2023-09-06T01:53:48.169610Z",
     "iopub.status.idle": "2023-09-06T01:53:50.281603Z",
     "shell.execute_reply": "2023-09-06T01:53:50.280695Z"
    },
    "papermill": {
     "duration": 2.117825,
     "end_time": "2023-09-06T01:53:50.283651",
     "exception": false,
     "start_time": "2023-09-06T01:53:48.165826",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\r\n",
      "\u001b[0m"
     ]
    }
   ],
   "source": [
    "# Additional dependencies for eval stuff\n",
    "!pip install -q aiocsv aiofiles"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "e478c7e4",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-09-06T01:53:50.291201Z",
     "iopub.status.busy": "2023-09-06T01:53:50.290951Z",
     "iopub.status.idle": "2023-09-06T01:53:50.300137Z",
     "shell.execute_reply": "2023-09-06T01:53:50.299422Z"
    },
    "papermill": {
     "duration": 0.014857,
     "end_time": "2023-09-06T01:53:50.301805",
     "exception": false,
     "start_time": "2023-09-06T01:53:50.286948",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "DEEPSPEED_STRAT: deepspeed_stage_1\n",
      "ENABLE_WANDB: True\n",
      "GPU_DEVICES: auto\n",
      "DIR_NAME: L12-D2048-E1e-1-ctx4k\n",
      "NOTEBOOK_DIR: /actions-runner/_work/RWKV-infctx-trainer/RWKV-infctx-trainer/notebook/experiment/rwkv-x-exp/v5-r3-memory/L12-D2048-E1e-1-ctx4k\n",
      "INFERENCE_DIR: /actions-runner/_work/RWKV-infctx-trainer/RWKV-infctx-trainer/RWKV-v5\n",
      "TRAINER_DIR: /actions-runner/_work/RWKV-infctx-trainer/RWKV-infctx-trainer/RWKV-v5\n",
      "PROJECT_DIR: /actions-runner/_work/RWKV-infctx-trainer/RWKV-infctx-trainer\n"
     ]
    }
   ],
   "source": [
    "DEEPSPEED_STRAT=\"deepspeed_stage_1\"\n",
    "GPU_DEVICES=\"auto\"\n",
    "ENABLE_WANDB=True\n",
    "\n",
    "# Layer count and embed dim to start with\n",
    "LAYER_COUNT=12\n",
    "EMBED_DIM=2048\n",
    "\n",
    "EMBED_SCALE=0.1\n",
    "EMBED_SCALE_LABEL=str(EMBED_SCALE).replace(\".\", \"_\")\n",
    "\n",
    "WANDB_PREFIX=f\"v5r3-L{LAYER_COUNT}-D{EMBED_DIM}-E{EMBED_SCALE}\"\n",
    "FILENAME_PREFIX=f\"v5r3-L{LAYER_COUNT}-D{EMBED_DIM}-E{EMBED_SCALE_LABEL}\"\n",
    "\n",
    "print(\"DEEPSPEED_STRAT:\", DEEPSPEED_STRAT)\n",
    "print(\"ENABLE_WANDB:\", ENABLE_WANDB)\n",
    "print(\"GPU_DEVICES:\", GPU_DEVICES)\n",
    "\n",
    "if ENABLE_WANDB:\n",
    "    WANDB_MODE=\"online\"\n",
    "else:\n",
    "    WANDB_MODE=\"disabled\"\n",
    "\n",
    "# Computing the notebook, and various paths\n",
    "import os\n",
    "NOTEBOOK_DIR=os.path.dirname(os.path.abspath(\"__file__\"))\n",
    "CONFIG_DIR=os.path.abspath(os.path.join(NOTEBOOK_DIR, \"../\"))\n",
    "PROJECT_DIR=os.path.abspath(os.path.join(CONFIG_DIR, \"../../../../\"))\n",
    "TRAINER_DIR=os.path.abspath(os.path.join(PROJECT_DIR, \"./RWKV-v5/\"))\n",
    "INFERENCE_DIR=os.path.abspath(os.path.join(PROJECT_DIR, \"./RWKV-v5/\"))\n",
    "\n",
    "# Get the notebook dir name\n",
    "DIR_NAME=os.path.basename(NOTEBOOK_DIR)\n",
    "\n",
    "# Log names and dir\n",
    "print(\"DIR_NAME:\", DIR_NAME)\n",
    "print(\"NOTEBOOK_DIR:\", NOTEBOOK_DIR)\n",
    "print(\"INFERENCE_DIR:\", INFERENCE_DIR)\n",
    "print(\"TRAINER_DIR:\", TRAINER_DIR)\n",
    "print(\"PROJECT_DIR:\", PROJECT_DIR)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "c0d1081e",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-09-06T01:53:50.309388Z",
     "iopub.status.busy": "2023-09-06T01:53:50.309130Z",
     "iopub.status.idle": "2023-09-06T01:53:50.530630Z",
     "shell.execute_reply": "2023-09-06T01:53:50.529804Z"
    },
    "papermill": {
     "duration": 0.227251,
     "end_time": "2023-09-06T01:53:50.532377",
     "exception": false,
     "start_time": "2023-09-06T01:53:50.305126",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/usr/bin/sh: 1: cd: can't cd to {TRAINER_DIR}\r\n"
     ]
    }
   ],
   "source": [
    "# Init the model\n",
    "!cd \"{TRAINER_DIR}\" && \\\n",
    "    export RWKV_WAVENET_LAYERS=\"{RWKV_WAVENET_LAYERS}\" && \\\n",
    "    python3 ./init_model.py \\\n",
    "        --n_layer \"{LAYER_COUNT}\" --n_embd \"{EMBED_DIM}\" \\\n",
    "        --emb-scale \"{EMBED_SCALE}\" \\\n",
    "        --vocab_size neox --skip-if-exists \\\n",
    "        \"../model/{FILENAME_PREFIX}-neox-init.pth\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ddc1086b",
   "metadata": {
    "papermill": {
     "duration": 0.003123,
     "end_time": "2023-09-06T01:53:50.538962",
     "exception": false,
     "start_time": "2023-09-06T01:53:50.535839",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Enwiki Stage 1 : Foundation 4k model training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "636df5aa",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-09-06T01:53:50.544775Z",
     "iopub.status.busy": "2023-09-06T01:53:50.544524Z",
     "iopub.status.idle": "2023-09-06T01:53:59.821234Z",
     "shell.execute_reply": "2023-09-06T01:53:59.820428Z"
    },
    "papermill": {
     "duration": 9.281152,
     "end_time": "2023-09-06T01:53:59.823158",
     "exception": false,
     "start_time": "2023-09-06T01:53:50.542006",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\r",
      "Saving the dataset (0/5 shards):   0%|         | 0/81505 [00:00<?, ? examples/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\r",
      "Saving the dataset (0/5 shards):   6%| | 5000/81505 [00:00<00:01, 40020.30 examp"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\r",
      "Saving the dataset (0/5 shards):  13%|▏| 11000/81505 [00:00<00:01, 43937.85 exam"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\r",
      "Saving the dataset (0/5 shards):  20%|▏| 16301/81505 [00:00<00:01, 46422.21 exam\r",
      "Saving the dataset (1/5 shards):  20%|▏| 16301/81505 [00:00<00:01, 46422.21 exam"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\r",
      "Saving the dataset (1/5 shards):  27%|▎| 22301/81505 [00:00<00:01, 48661.83 exam"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\r",
      "Saving the dataset (1/5 shards):  35%|▎| 28301/81505 [00:00<00:01, 50444.99 exam"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\r",
      "Saving the dataset (2/5 shards):  40%|▍| 32602/81505 [00:00<00:00, 50444.99 exam"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\r",
      "Saving the dataset (2/5 shards):  44%|▍| 35602/81505 [00:00<00:00, 51589.83 exam"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\r",
      "Saving the dataset (2/5 shards):  51%|▌| 41602/81505 [00:00<00:00, 45740.30 exam"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\r",
      "Saving the dataset (2/5 shards):  58%|▌| 47602/81505 [00:01<00:00, 46874.25 exam\r",
      "Saving the dataset (3/5 shards):  60%|▌| 48903/81505 [00:01<00:00, 46874.25 exam"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\r",
      "Saving the dataset (3/5 shards):  67%|▋| 54903/81505 [00:01<00:00, 48825.20 exam"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\r",
      "Saving the dataset (3/5 shards):  75%|▋| 60903/81505 [00:01<00:00, 50738.43 exam"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\r",
      "Saving the dataset (4/5 shards):  80%|▊| 65204/81505 [00:01<00:00, 50738.43 exam"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\r",
      "Saving the dataset (4/5 shards):  84%|▊| 68204/81505 [00:01<00:00, 51451.93 exam"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\r",
      "Saving the dataset (4/5 shards):  92%|▉| 75204/81505 [00:01<00:00, 44264.66 exam"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\r",
      "Saving the dataset (4/5 shards): 100%|▉| 81204/81505 [00:01<00:00, 45739.71 exam\r",
      "Saving the dataset (5/5 shards): 100%|█| 81505/81505 [00:01<00:00, 45739.71 exam\r",
      "Saving the dataset (5/5 shards): 100%|█| 81505/81505 [00:01<00:00, 47312.07 exam\r\n",
      "\r",
      "Saving the dataset (0/1 shards):   0%|           | 0/410 [00:00<?, ? examples/s]\r",
      "Saving the dataset (1/1 shards): 100%|█| 410/410 [00:00<00:00, 40194.11 examples\r",
      "Saving the dataset (1/1 shards): 100%|█| 410/410 [00:00<00:00, 38725.95 examples\r\n"
     ]
    }
   ],
   "source": [
    "# Lets preload the requried dataset \n",
    "!cd \"{TRAINER_DIR}\" && \\\n",
    "    python3 preload_datapath.py \"{CONFIG_DIR}/config-enwiki-4k.yaml\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "14045bce",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-09-06T01:53:59.833777Z",
     "iopub.status.busy": "2023-09-06T01:53:59.833515Z",
     "iopub.status.idle": "2023-09-06T01:54:00.054878Z",
     "shell.execute_reply": "2023-09-06T01:54:00.054138Z"
    },
    "papermill": {
     "duration": 0.228539,
     "end_time": "2023-09-06T01:54:00.056582",
     "exception": false,
     "start_time": "2023-09-06T01:53:59.828043",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/usr/bin/sh: 1: cd: can't cd to {TRAINER_DIR}\r\n"
     ]
    }
   ],
   "source": [
    "# Start the foundation model training\n",
    "!cd \"{TRAINER_DIR}\" && \\\n",
    "    export RWKV_WAVENET_LAYERS=\"{RWKV_WAVENET_LAYERS}\" && \\\n",
    "    export WANDB_MODE=\"{WANDB_MODE}\" && \\\n",
    "    python lightning_trainer.py fit \\\n",
    "        -c \"{CONFIG_DIR}/config-enwiki-4k.yaml\" \\\n",
    "        --trainer.logger.init_args.name=\"{WANDB_PREFIX} - Enwiki-4k Foundation (train-ctx=4k, {DEEPSPEED_STRAT})\" \\\n",
    "        --trainer.strategy=\"{DEEPSPEED_STRAT}\" \\\n",
    "        --trainer.devices=\"{GPU_DEVICES}\" \\\n",
    "        --trainer.callbacks.init_args.dirpath=\"../checkpoint/{FILENAME_PREFIX}-enwiki-4k/\" \\\n",
    "        --model.load_model=\"../model/{FILENAME_PREFIX}-neox-init.pth\" \\\n",
    "        --model.ctx_len=4096 \\\n",
    "        --model.bptt_learning_range=1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "498791c4",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-09-06T01:54:00.067029Z",
     "iopub.status.busy": "2023-09-06T01:54:00.066776Z",
     "iopub.status.idle": "2023-09-06T01:54:00.505182Z",
     "shell.execute_reply": "2023-09-06T01:54:00.504440Z"
    },
    "papermill": {
     "duration": 0.445543,
     "end_time": "2023-09-06T01:54:00.506877",
     "exception": false,
     "start_time": "2023-09-06T01:54:00.061334",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/usr/bin/sh: 1: python: not found\r\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ls: cannot access '../model/v5r3-L12-D2048-E0_1-enwiki-4k.pth': No such file or directory\r\n"
     ]
    }
   ],
   "source": [
    "# Lets export the model from the checkpoint\n",
    "!cd \"{TRAINER_DIR}\" && \\\n",
    "    python export_checkpoint.py \"../checkpoint/{FILENAME_PREFIX}-enwiki-4k/last.ckpt\" \"../model/{FILENAME_PREFIX}-enwiki-4k.pth\" \"bf16\"\n",
    "!cd \"{TRAINER_DIR}\" && ls -alh \"../model/{FILENAME_PREFIX}-enwiki-4k.pth\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "eeae81c5",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-09-06T01:54:00.517808Z",
     "iopub.status.busy": "2023-09-06T01:54:00.517561Z",
     "iopub.status.idle": "2023-09-06T01:54:00.734217Z",
     "shell.execute_reply": "2023-09-06T01:54:00.733477Z"
    },
    "papermill": {
     "duration": 0.223994,
     "end_time": "2023-09-06T01:54:00.735863",
     "exception": false,
     "start_time": "2023-09-06T01:54:00.511869",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/usr/bin/sh: 1: cd: can't cd to {INFERENCE_DIR}\r\n"
     ]
    }
   ],
   "source": [
    "# # Lets do a quick dragon prompt validation\n",
    "!cd \"{INFERENCE_DIR}\" && \\\n",
    "    export RWKV_WAVENET_LAYERS=\"{RWKV_WAVENET_LAYERS}\" && \\\n",
    "    python3 dragon_test.py \"../model/{FILENAME_PREFIX}-enwiki-4k.pth\" \"cuda fp32\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "e7fcd1dc",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-09-06T01:54:00.747330Z",
     "iopub.status.busy": "2023-09-06T01:54:00.746699Z",
     "iopub.status.idle": "2023-09-06T01:54:00.976416Z",
     "shell.execute_reply": "2023-09-06T01:54:00.975671Z"
    },
    "papermill": {
     "duration": 0.237171,
     "end_time": "2023-09-06T01:54:00.978059",
     "exception": false,
     "start_time": "2023-09-06T01:54:00.740888",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "python3: can't open file '/actions-runner/_work/RWKV-infctx-trainer/RWKV-infctx-trainer/notebook/experiment/rwkv-x-exp/v5-r3-memory/L12-D2048-E1e-1-ctx4k/../memory_script/eval_v5_memory_guided.py': [Errno 2] No such file or directory\r\n"
     ]
    }
   ],
   "source": [
    "# Lets do a quick memory test\n",
    "!export RWKV_WAVENET_LAYERS=\"{RWKV_WAVENET_LAYERS}\" && \\\n",
    "        python3 ../memory_script/eval_v5_memory_guided.py \"{PROJECT_DIR}/model/{FILENAME_PREFIX}-enwiki-4k.pth\""
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "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.12"
  },
  "papermill": {
   "default_parameters": {},
   "duration": 14.829906,
   "end_time": "2023-09-06T01:54:01.098909",
   "environment_variables": {},
   "exception": null,
   "input_path": "/actions-runner/_work/RWKV-infctx-trainer/RWKV-infctx-trainer/notebook/experiment/rwkv-x-exp/v5-r3-memory/L12-D2048-E1e-1-ctx4k/part1.ipynb",
   "output_path": "/actions-runner/_work/RWKV-infctx-trainer/RWKV-infctx-trainer/output/experiment/rwkv-x-exp/v5-r3-memory/L12-D2048-E1e-1-ctx4k/part1.ipynb",
   "parameters": {},
   "start_time": "2023-09-06T01:53:46.269003",
   "version": "2.4.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}