rwkv5_world_run.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "c2d7d23f-3a68-4436-a70c-88aa822de290",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "RWKV_JIT_ON 1 RWKV_CUDA_ON 1 RESCALE_LAYER 6\n",
      "\n",
      "Loading ./model/0.1-1/rwkv-10.pth ...\n",
      "Model detected: v5.2\n",
      "Strategy: (total 12+1=13 layers)\n",
      "* cuda [float16, float16], store 13 layers\n",
      "0-cuda-float16-float16 1-cuda-float16-float16 2-cuda-float16-float16 3-cuda-float16-float16 4-cuda-float16-float16 5-cuda-float16-float16 6-cuda-float16-float16 7-cuda-float16-float16 8-cuda-float16-float16 9-cuda-float16-float16 10-cuda-float16-float16 11-cuda-float16-float16 12-cuda-float16-float16 \n",
      "emb.weight                        f16      cpu  65536   768 \n",
      "blocks.0.ln1.weight               f16   cuda:0    768       \n",
      "blocks.0.ln1.bias                 f16   cuda:0    768       \n",
      "blocks.0.ln2.weight               f16   cuda:0    768       \n",
      "blocks.0.ln2.bias                 f16   cuda:0    768       \n",
      "blocks.0.att.time_mix_k           f16   cuda:0    768       \n",
      "blocks.0.att.time_mix_v           f16   cuda:0    768       \n",
      "blocks.0.att.time_mix_r           f16   cuda:0    768       \n",
      "blocks.0.att.time_mix_g           f16   cuda:0    768       \n",
      "blocks.0.att.time_decay           f32   cuda:0     12    64 \n",
      "blocks.0.att.time_first           f32   cuda:0     12    64 \n",
      "blocks.0.att.receptance.weight    f16   cuda:0    768   768 \n",
      "blocks.0.att.key.weight           f16   cuda:0    768   768 \n",
      "blocks.0.att.value.weight         f16   cuda:0    768   768 \n",
      "blocks.0.att.output.weight        f16   cuda:0    768   768 \n",
      "blocks.0.att.gate.weight          f16   cuda:0    768   768 \n",
      "blocks.0.att.ln_x.weight          f32   cuda:0    768       \n",
      "blocks.0.att.ln_x.bias            f32   cuda:0    768       \n",
      "blocks.0.ffn.time_mix_k           f16   cuda:0    768       \n",
      "blocks.0.ffn.time_mix_r           f16   cuda:0    768       \n",
      "blocks.0.ffn.key.weight           f16   cuda:0    768  2688 \n",
      "blocks.0.ffn.receptance.weight    f16   cuda:0    768   768 \n",
      "blocks.0.ffn.value.weight         f16   cuda:0   2688   768 \n",
      "............................................................................................................................................................................................................................\n",
      "blocks.11.ln1.weight              f16   cuda:0    768       \n",
      "blocks.11.ln1.bias                f16   cuda:0    768       \n",
      "blocks.11.ln2.weight              f16   cuda:0    768       \n",
      "blocks.11.ln2.bias                f16   cuda:0    768       \n",
      "blocks.11.att.time_mix_k          f16   cuda:0    768       \n",
      "blocks.11.att.time_mix_v          f16   cuda:0    768       \n",
      "blocks.11.att.time_mix_r          f16   cuda:0    768       \n",
      "blocks.11.att.time_mix_g          f16   cuda:0    768       \n",
      "blocks.11.att.time_decay          f32   cuda:0     12    64 \n",
      "blocks.11.att.time_first          f32   cuda:0     12    64 \n",
      "blocks.11.att.receptance.weight   f16   cuda:0    768   768 \n",
      "blocks.11.att.key.weight          f16   cuda:0    768   768 \n",
      "blocks.11.att.value.weight        f16   cuda:0    768   768 \n",
      "blocks.11.att.output.weight       f16   cuda:0    768   768 \n",
      "blocks.11.att.gate.weight         f16   cuda:0    768   768 \n",
      "blocks.11.att.ln_x.weight         f32   cuda:0    768       \n",
      "blocks.11.att.ln_x.bias           f32   cuda:0    768       \n",
      "blocks.11.ffn.time_mix_k          f16   cuda:0    768       \n",
      "blocks.11.ffn.time_mix_r          f16   cuda:0    768       \n",
      "blocks.11.ffn.key.weight          f16   cuda:0    768  2688 \n",
      "blocks.11.ffn.receptance.weight   f16   cuda:0    768   768 \n",
      "blocks.11.ffn.value.weight        f16   cuda:0   2688   768 \n",
      "ln_out.weight                     f16   cuda:0    768       \n",
      "ln_out.bias                       f16   cuda:0    768       \n",
      "head.weight                       f16   cuda:0    768 65536 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using /root/.cache/torch_extensions/py38_cu117 as PyTorch extensions root...\n",
      "No modifications detected for re-loaded extension module rwkv5, skipping build step...\n",
      "Loading extension module rwkv5...\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "기아 타이거즈 헬멧: 플레이어는 에너지와 접촉을 하며 인공 지능을 사용할 수 있습니다.\n",
      "야구팀: 마이크로소프트가 컨퍼런스에 진출한 후에는 펑크와 리조트에 가입하여 특정 요구를 구현하는 것이 중요합니다.\n",
      "\n",
      "프로그래밍: 윈도우 7과 함께 출시되었으며, 마이크로소프트가 리눅스 엔진에서 선보일 수 있도록 도와줍니다.\n",
      "\n",
      "요약하자면, 기술이 개선되면 새로운 펑크와 리눅스 펑크가 필요하다는 것을 의미하세요. 예를 들어, AWS의 모든 제품은 매출과 기타 프로세스에 대한 적절한\n",
      "\n",
      "[ 2.7539062 -5.3242188 -5.4179688 ... -5.3671875 -5.3320312 -5.3320312]\n",
      "[ 2.7519531 -5.3203125 -5.4140625 ... -5.3632812 -5.328125  -5.328125 ]\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "\n",
    "# set these before import RWKV\n",
    "os.environ['RWKV_JIT_ON'] = '1'\n",
    "os.environ[\"RWKV_CUDA_ON\"] = '1' # '1' to compile CUDA kernel (10x faster), requires c++ compiler & cuda libraries\n",
    "\n",
    "########################################################################################################\n",
    "#\n",
    "# Use '/' in model path, instead of '\\'. Use ctx4096 models if you need long ctx.\n",
    "#\n",
    "# fp16 = good for GPU (!!! DOES NOT support CPU !!!)\n",
    "# fp32 = good for CPU\n",
    "# bf16 = worse accuracy, supports CPU\n",
    "# xxxi8 (example: fp16i8, fp32i8) = xxx with int8 quantization to save 50% VRAM/RAM, slower, slightly less accuracy\n",
    "#\n",
    "# We consider [ln_out+head] to be an extra layer, so L12-D768 (169M) has \"13\" layers, L24-D2048 (1.5B) has \"25\" layers, etc.\n",
    "# Strategy Examples: (device = cpu/cuda/cuda:0/cuda:1/...)\n",
    "# 'cpu fp32' = all layers cpu fp32\n",
    "# 'cuda fp16' = all layers cuda fp16\n",
    "# 'cuda fp16i8' = all layers cuda fp16 with int8 quantization\n",
    "# 'cuda fp16i8 *10 -> cpu fp32' = first 10 layers cuda fp16i8, then cpu fp32 (increase 10 for better speed)\n",
    "# 'cuda:0 fp16 *10 -> cuda:1 fp16 *8 -> cpu fp32' = first 10 layers cuda:0 fp16, then 8 layers cuda:1 fp16, then cpu fp32\n",
    "#\n",
    "# Basic Strategy Guide: (fp16i8 works for any GPU)\n",
    "# 100% VRAM = 'cuda fp16'                   # all layers cuda fp16\n",
    "#  98% VRAM = 'cuda fp16i8 *1 -> cuda fp16' # first 1 layer  cuda fp16i8, then cuda fp16\n",
    "#  96% VRAM = 'cuda fp16i8 *2 -> cuda fp16' # first 2 layers cuda fp16i8, then cuda fp16\n",
    "#  94% VRAM = 'cuda fp16i8 *3 -> cuda fp16' # first 3 layers cuda fp16i8, then cuda fp16\n",
    "#  ...\n",
    "#  50% VRAM = 'cuda fp16i8'                 # all layers cuda fp16i8\n",
    "#  48% VRAM = 'cuda fp16i8 -> cpu fp32 *1'  # most layers cuda fp16i8, last 1 layer  cpu fp32\n",
    "#  46% VRAM = 'cuda fp16i8 -> cpu fp32 *2'  # most layers cuda fp16i8, last 2 layers cpu fp32\n",
    "#  44% VRAM = 'cuda fp16i8 -> cpu fp32 *3'  # most layers cuda fp16i8, last 3 layers cpu fp32\n",
    "#  ...\n",
    "#   0% VRAM = 'cpu fp32'                    # all layers cpu fp32\n",
    "#\n",
    "# Use '+' for STREAM mode, which can save VRAM too, and it is sometimes faster\n",
    "# 'cuda fp16i8 *10+' = first 10 layers cuda fp16i8, then fp16i8 stream the rest to it (increase 10 for better speed)\n",
    "#\n",
    "# Extreme STREAM: 3G VRAM is enough to run RWKV 14B (slow. will be faster in future)\n",
    "# 'cuda fp16i8 *0+ -> cpu fp32 *1' = stream all layers cuda fp16i8, last 1 layer [ln_out+head] cpu fp32\n",
    "#\n",
    "# ########################################################################################################\n",
    "\n",
    "from rwkv.model import RWKV\n",
    "from rwkv.utils import PIPELINE, PIPELINE_ARGS\n",
    "\n",
    "# download models: https://huggingface.co/BlinkDL\n",
    "model = RWKV(model='./model/0.1-1/rwkv-10', strategy='cuda fp16')\n",
    "pipeline = PIPELINE(model, \"rwkv_vocab_v20230424\") # 20B_tokenizer.json is in https://github.com/BlinkDL/ChatRWKV\n",
    "\n",
    "ctx = \"\\n기아 타이거즈 \"\n",
    "print(ctx, end='')\n",
    "\n",
    "def my_print(s):\n",
    "    print(s, end='', flush=True)\n",
    "\n",
    "# For alpha_frequency and alpha_presence, see \"Frequency and presence penalties\":\n",
    "# https://platform.openai.com/docs/api-reference/parameter-details\n",
    "\n",
    "args = PIPELINE_ARGS(temperature = 1.0, top_p = 0.7, top_k = 100, # top_k = 0 then ignore\n",
    "                     alpha_frequency = 0.25,\n",
    "                     alpha_presence = 0.25,\n",
    "                     alpha_decay = 0.996, # gradually decay the penalty\n",
    "                     token_ban = [0], # ban the generation of some tokens\n",
    "                     token_stop = [], # stop generation whenever you see any token here\n",
    "                     chunk_len = 256) # split input into chunks to save VRAM (shorter -> slower)\n",
    "\n",
    "pipeline.generate(ctx, token_count=200, args=args, callback=my_print)\n",
    "print('\\n')\n",
    "\n",
    "out, state = model.forward([187, 510, 1563, 310, 247], None)\n",
    "print(out.detach().cpu().numpy())                   # get logits\n",
    "out, state = model.forward([187, 510], None)\n",
    "out, state = model.forward([1563], state)           # RNN has state (use deepcopy to clone states)\n",
    "out, state = model.forward([310, 247], state)\n",
    "print(out.detach().cpu().numpy())                   # same result as above\n",
    "print('\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c2375a08-c984-449d-9a02-f4190212a11a",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.8.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}