9f942c95f693bd4e3f53fac4ddf2b42f01533af8281ff8513ca0df539b8cf746

experiment/rwkv-x-exp/multi-size-train/v5-L6-D2048-part1.ipynb

@@ -0,0 +1,685 @@
+{
+ "cells": [
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "id": "f5ebb021",
+   "metadata": {
+    "papermill": {
+     "duration": 0.002819,
+     "end_time": "2023-09-29T04:50:11.669839",
+     "exception": false,
+     "start_time": "2023-09-29T04:50:11.667020",
+     "status": "completed"
+    },
+    "tags": []
+   },
+   "source": [
+    "# RWKV v5 multi-size training experiment\n",
+    "\n",
+    "**Note:** This project assumes you have the rwkv-infctx conda env setup"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "id": "6e6abc3f",
+   "metadata": {
+    "papermill": {
+     "duration": 0.00214,
+     "end_time": "2023-09-29T04:50:11.676239",
+     "exception": false,
+     "start_time": "2023-09-29T04:50:11.674099",
+     "status": "completed"
+    },
+    "tags": []
+   },
+   "source": [
+    "# Basic Setup"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "3d2405bd",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2023-09-29T04:50:11.682830Z",
+     "iopub.status.busy": "2023-09-29T04:50:11.682139Z",
+     "iopub.status.idle": "2023-09-29T04:50:12.432460Z",
+     "shell.execute_reply": "2023-09-29T04:50:12.431486Z"
+    },
+    "papermill": {
+     "duration": 0.756299,
+     "end_time": "2023-09-29T04:50:12.434815",
+     "exception": false,
+     "start_time": "2023-09-29T04:50:11.678516",
+     "status": "completed"
+    },
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "# First lets setup the various directories, and init the model\n",
+    "!mkdir -p ../../../../model/\n",
+    "!mkdir -p ../../../../datapath/\n",
+    "!mkdir -p ../../../../checkpoint/"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "66fd5201",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2023-09-29T04:50:12.446546Z",
+     "iopub.status.busy": "2023-09-29T04:50:12.446098Z",
+     "iopub.status.idle": "2023-09-29T04:50:12.454394Z",
+     "shell.execute_reply": "2023-09-29T04:50:12.453644Z"
+    },
+    "papermill": {
+     "duration": 0.018125,
+     "end_time": "2023-09-29T04:50:12.456177",
+     "exception": false,
+     "start_time": "2023-09-29T04:50:12.438052",
+     "status": "completed"
+    },
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "DEEPSPEED_STRAT: deepspeed_stage_1\n",
+      "ENABLE_WANDB: True\n",
+      "GPU_DEVICES: auto\n",
+      "NOTEBOOK_DIR: /actions-runner/_work/RWKV-infctx-trainer/RWKV-infctx-trainer/notebook/experiment/rwkv-x-exp/multi-size-train\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",
+    "EMBED_SCALE=0.01\n",
+    "EMBED_SCALE_LABEL=str(EMBED_SCALE).replace(\".\", \"_\")\n",
+    "\n",
+    "LAYER_COUNT=6\n",
+    "EMBED_SIZE=2048\n",
+    "\n",
+    "WANDB_PREFIX=f\"[Multi-size] v5-L{LAYER_COUNT}-D{EMBED_SIZE}-E{EMBED_SCALE}\"\n",
+    "FILENAME_PREFIX=f\"v5-L{LAYER_COUNT}-D{EMBED_SIZE}-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",
+    "PROJECT_DIR=os.path.abspath(os.path.join(NOTEBOOK_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",
+    "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": 3,
+   "id": "e0b56789",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2023-09-29T04:50:12.464627Z",
+     "iopub.status.busy": "2023-09-29T04:50:12.464037Z",
+     "iopub.status.idle": "2023-09-29T04:50:42.488005Z",
+     "shell.execute_reply": "2023-09-29T04:50:42.486665Z"
+    },
+    "papermill": {
+     "duration": 30.031629,
+     "end_time": "2023-09-29T04:50:42.490859",
+     "exception": false,
+     "start_time": "2023-09-29T04:50:12.459230",
+     "status": "completed"
+    },
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[2023-09-29 04:50:16,856] [INFO] [real_accelerator.py:133:get_accelerator] Setting ds_accelerator to cuda (auto detect)\r\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[RWKV.model] Running RWKV model using 'torch-jit' with torch '2.0.1+cu118'\r\n",
+      "---- Initializing model ----\r\n",
+      "No of layers: 6\r\n",
+      "Embedding size: 2048\r\n",
+      "Output model path: ../model/v5-L6-D2048-E0_01-neox-v5base-init.pth\r\n",
+      "Vocab size: 50277\r\n",
+      "Emb scale: 0.01\r\n",
+      "Note: this process takes a significant time (and ram) for large models\r\n",
+      "---- ----- ----\r\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "50277 2048  -0.01 emb.weight\r\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "2048  2048  1.0  blocks.0.att.gate.weight\r\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "2048  2048  1.0  blocks.0.att.receptance.weight\r\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "2048  2048  1.0  blocks.0.att.key.weight\r\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "2048  2048  1.0  blocks.0.att.value.weight\r\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "2048  2048  0    blocks.0.att.output.weight\r\n",
+      "7168  2048  1.0  blocks.0.ffn.key.weight\r\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "2048  2048  0    blocks.0.ffn.receptance.weight\r\n",
+      "2048  7168  0    blocks.0.ffn.value.weight\r\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "2048  2048  1.0  blocks.1.att.gate.weight\r\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "2048  2048  1.0  blocks.1.att.receptance.weight\r\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "2048  2048  1.0  blocks.1.att.key.weight\r\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "2048  2048  1.0  blocks.1.att.value.weight\r\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "2048  2048  0    blocks.1.att.output.weight\r\n",
+      "7168  2048  1.0  blocks.1.ffn.key.weight\r\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "2048  2048  0    blocks.1.ffn.receptance.weight\r\n",
+      "2048  7168  0    blocks.1.ffn.value.weight\r\n",
+      "2048  2048  1.0  blocks.2.att.gate.weight\r\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "2048  2048  1.0  blocks.2.att.receptance.weight\r\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "2048  2048  1.0  blocks.2.att.key.weight\r\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "2048  2048  1.0  blocks.2.att.value.weight\r\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "2048  2048  0    blocks.2.att.output.weight\r\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "7168  2048  1.0  blocks.2.ffn.key.weight\r\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "2048  2048  0    blocks.2.ffn.receptance.weight\r\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "2048  7168  0    blocks.2.ffn.value.weight\r\n",
+      "2048  2048  1.0  blocks.3.att.gate.weight\r\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "2048  2048  1.0  blocks.3.att.receptance.weight\r\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "2048  2048  1.0  blocks.3.att.key.weight\r\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "2048  2048  1.0  blocks.3.att.value.weight\r\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "2048  2048  0    blocks.3.att.output.weight\r\n",
+      "7168  2048  1.0  blocks.3.ffn.key.weight\r\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "2048  2048  0    blocks.3.ffn.receptance.weight\r\n",
+      "2048  7168  0    blocks.3.ffn.value.weight\r\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "2048  2048  1.0  blocks.4.att.gate.weight\r\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "2048  2048  1.0  blocks.4.att.receptance.weight\r\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "2048  2048  1.0  blocks.4.att.key.weight\r\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "2048  2048  1.0  blocks.4.att.value.weight\r\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "2048  2048  0    blocks.4.att.output.weight\r\n",
+      "7168  2048  1.0  blocks.4.ffn.key.weight\r\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "2048  2048  0    blocks.4.ffn.receptance.weight\r\n",
+      "2048  7168  0    blocks.4.ffn.value.weight\r\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "2048  2048  1.0  blocks.5.att.gate.weight\r\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "2048  2048  1.0  blocks.5.att.receptance.weight\r\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "2048  2048  1.0  blocks.5.att.key.weight\r\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "2048  2048  1.0  blocks.5.att.value.weight\r\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "2048  2048  0    blocks.5.att.output.weight\r\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "7168  2048  1.0  blocks.5.ffn.key.weight\r\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "2048  2048  0    blocks.5.ffn.receptance.weight\r\n",
+      "2048  7168  0    blocks.5.ffn.value.weight\r\n",
+      "50277 2048  0.5  head.weight\r\n"
+     ]
+    }
+   ],
+   "source": [
+    "# Init the model\n",
+    "!cd \"{TRAINER_DIR}\" && \\\n",
+    "    python3 ./init_model.py \\\n",
+    "        --n_layer {LAYER_COUNT} --n_embd {EMBED_SIZE} \\\n",
+    "        --emb-scale \"{EMBED_SCALE}\" \\\n",
+    "        --vocab_size neox --skip-if-exists \\\n",
+    "        \"../model/{FILENAME_PREFIX}-neox-v5base-init.pth\""
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "e3057f8b",
+   "metadata": {
+    "papermill": {
+     "duration": 0.006306,
+     "end_time": "2023-09-29T04:50:42.503924",
+     "exception": false,
+     "start_time": "2023-09-29T04:50:42.497618",
+     "status": "completed"
+    },
+    "tags": []
+   },
+   "source": [
+    "## Enwiki Stage 1 : Foundation 4k model training"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "c06c6ad2",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2023-09-29T04:50:42.519776Z",
+     "iopub.status.busy": "2023-09-29T04:50:42.518734Z",
+     "iopub.status.idle": "2023-09-29T04:50:47.537394Z",
+     "shell.execute_reply": "2023-09-29T04:50:47.535673Z"
+    },
+    "papermill": {
+     "duration": 5.029265,
+     "end_time": "2023-09-29T04:50:47.539698",
+     "exception": false,
+     "start_time": "2023-09-29T04:50:42.510433",
+     "status": "completed"
+    },
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Traceback (most recent call last):\r\n",
+      "  File \"/actions-runner/_work/RWKV-infctx-trainer/RWKV-infctx-trainer/RWKV-v5/preload_datapath.py\", line 20, in <module>\r\n",
+      "    assert os.path.exists(config_file), \"Config file does not exist\"\r\n",
+      "AssertionError: Config file does not exist\r\n"
+     ]
+    }
+   ],
+   "source": [
+    "# Lets preload the requried dataset \n",
+    "!cd \"{TRAINER_DIR}\" && \\\n",
+    "    python3 preload_datapath.py \"{NOTEBOOK_DIR}/v5base-enwiki-4k-part1.yaml\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "4cc7e34f",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2023-09-29T04:50:47.552392Z",
+     "iopub.status.busy": "2023-09-29T04:50:47.551853Z",
+     "iopub.status.idle": "2023-09-29T04:50:47.806392Z",
+     "shell.execute_reply": "2023-09-29T04:50:47.805379Z"
+    },
+    "papermill": {
+     "duration": 0.264553,
+     "end_time": "2023-09-29T04:50:47.809133",
+     "exception": false,
+     "start_time": "2023-09-29T04:50:47.544580",
+     "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 \"{NOTEBOOK_DIR}/v5base-enwiki-4k.yaml\" \\\n",
+    "        --trainer.logger.init_args.name=\"{WANDB_PREFIX} - Enwiki-4k Part 1 (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-p1/\" \\\n",
+    "        --model.load_model=\"../model/{FILENAME_PREFIX}-neox-v5base-init.pth\" \\\n",
+    "        --model.ctx_len=4096 \\\n",
+    "        --model.bptt_learning_range=1"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "0b3b8134",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2023-09-29T04:50:47.825099Z",
+     "iopub.status.busy": "2023-09-29T04:50:47.824495Z",
+     "iopub.status.idle": "2023-09-29T04:50:48.327589Z",
+     "shell.execute_reply": "2023-09-29T04:50:48.326466Z"
+    },
+    "papermill": {
+     "duration": 0.514109,
+     "end_time": "2023-09-29T04:50:48.330177",
+     "exception": false,
+     "start_time": "2023-09-29T04:50:47.816068",
+     "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/v5-L6-D2048-E0_01-enwiki-4k-p1.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-p1/last.ckpt\" \"../model/{FILENAME_PREFIX}-enwiki-4k-p1.pth\" \"bf16\"\n",
+    "!cd \"{TRAINER_DIR}\" && ls -alh \"../model/{FILENAME_PREFIX}-enwiki-4k-p1.pth\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "92869fb9",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2023-09-29T04:50:48.346924Z",
+     "iopub.status.busy": "2023-09-29T04:50:48.346311Z",
+     "iopub.status.idle": "2023-09-29T04:50:48.600443Z",
+     "shell.execute_reply": "2023-09-29T04:50:48.599423Z"
+    },
+    "papermill": {
+     "duration": 0.26565,
+     "end_time": "2023-09-29T04:50:48.603118",
+     "exception": false,
+     "start_time": "2023-09-29T04:50:48.337468",
+     "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-p1.pth\" \"cuda fp32\""
+   ]
+  }
+ ],
+ "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": 38.643439,
+   "end_time": "2023-09-29T04:50:49.032467",
+   "environment_variables": {},
+   "exception": null,
+   "input_path": "/actions-runner/_work/RWKV-infctx-trainer/RWKV-infctx-trainer/notebook/experiment/rwkv-x-exp/multi-size-train/v5-L6-D2048-part1.ipynb",
+   "output_path": "/actions-runner/_work/RWKV-infctx-trainer/RWKV-infctx-trainer/output/experiment/rwkv-x-exp/multi-size-train/v5-L6-D2048-part1.ipynb",
+   "parameters": {},
+   "start_time": "2023-09-29T04:50:10.389028",
+   "version": "2.4.0"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}