pure_jax.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "COMBINE_BYTES = 130*1024*1024\n",
    "os.environ[\"XLA_FLAGS\"] = f\"--xla_gpu_simplify_all_fp_conversions --xla_gpu_all_reduce_combine_threshold_bytes={COMBINE_BYTES}\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "dict_keys(['blocks.0.att.key.weight', 'blocks.0.att.output.weight', 'blocks.0.att.receptance.weight', 'blocks.0.att.time_decay', 'blocks.0.att.time_first', 'blocks.0.att.time_mix_k', 'blocks.0.att.time_mix_r', 'blocks.0.att.time_mix_v', 'blocks.0.att.value.weight', 'blocks.0.ffn.key.weight', 'blocks.0.ffn.receptance.weight', 'blocks.0.ffn.time_mix_k', 'blocks.0.ffn.time_mix_r', 'blocks.0.ffn.value.weight', 'blocks.0.ln0.bias', 'blocks.0.ln0.weight', 'blocks.0.ln1.bias', 'blocks.0.ln1.weight', 'blocks.0.ln2.bias', 'blocks.0.ln2.weight', 'blocks.1.att.key.weight', 'blocks.1.att.output.weight', 'blocks.1.att.receptance.weight', 'blocks.1.att.time_decay', 'blocks.1.att.time_first', 'blocks.1.att.time_mix_k', 'blocks.1.att.time_mix_r', 'blocks.1.att.time_mix_v', 'blocks.1.att.value.weight', 'blocks.1.ffn.key.weight', 'blocks.1.ffn.receptance.weight', 'blocks.1.ffn.time_mix_k', 'blocks.1.ffn.time_mix_r', 'blocks.1.ffn.value.weight', 'blocks.1.ln1.bias', 'blocks.1.ln1.weight', 'blocks.1.ln2.bias', 'blocks.1.ln2.weight', 'blocks.10.att.key.weight', 'blocks.10.att.output.weight', 'blocks.10.att.receptance.weight', 'blocks.10.att.time_decay', 'blocks.10.att.time_first', 'blocks.10.att.time_mix_k', 'blocks.10.att.time_mix_r', 'blocks.10.att.time_mix_v', 'blocks.10.att.value.weight', 'blocks.10.ffn.key.weight', 'blocks.10.ffn.receptance.weight', 'blocks.10.ffn.time_mix_k', 'blocks.10.ffn.time_mix_r', 'blocks.10.ffn.value.weight', 'blocks.10.ln1.bias', 'blocks.10.ln1.weight', 'blocks.10.ln2.bias', 'blocks.10.ln2.weight', 'blocks.11.att.key.weight', 'blocks.11.att.output.weight', 'blocks.11.att.receptance.weight', 'blocks.11.att.time_decay', 'blocks.11.att.time_first', 'blocks.11.att.time_mix_k', 'blocks.11.att.time_mix_r', 'blocks.11.att.time_mix_v', 'blocks.11.att.value.weight', 'blocks.11.ffn.key.weight', 'blocks.11.ffn.receptance.weight', 'blocks.11.ffn.time_mix_k', 'blocks.11.ffn.time_mix_r', 'blocks.11.ffn.value.weight', 'blocks.11.ln1.bias', 'blocks.11.ln1.weight', 'blocks.11.ln2.bias', 'blocks.11.ln2.weight', 'blocks.12.att.key.weight', 'blocks.12.att.output.weight', 'blocks.12.att.receptance.weight', 'blocks.12.att.time_decay', 'blocks.12.att.time_first', 'blocks.12.att.time_mix_k', 'blocks.12.att.time_mix_r', 'blocks.12.att.time_mix_v', 'blocks.12.att.value.weight', 'blocks.12.ffn.key.weight', 'blocks.12.ffn.receptance.weight', 'blocks.12.ffn.time_mix_k', 'blocks.12.ffn.time_mix_r', 'blocks.12.ffn.value.weight', 'blocks.12.ln1.bias', 'blocks.12.ln1.weight', 'blocks.12.ln2.bias', 'blocks.12.ln2.weight', 'blocks.13.att.key.weight', 'blocks.13.att.output.weight', 'blocks.13.att.receptance.weight', 'blocks.13.att.time_decay', 'blocks.13.att.time_first', 'blocks.13.att.time_mix_k', 'blocks.13.att.time_mix_r', 'blocks.13.att.time_mix_v', 'blocks.13.att.value.weight', 'blocks.13.ffn.key.weight', 'blocks.13.ffn.receptance.weight', 'blocks.13.ffn.time_mix_k', 'blocks.13.ffn.time_mix_r', 'blocks.13.ffn.value.weight', 'blocks.13.ln1.bias', 'blocks.13.ln1.weight', 'blocks.13.ln2.bias', 'blocks.13.ln2.weight', 'blocks.14.att.key.weight', 'blocks.14.att.output.weight', 'blocks.14.att.receptance.weight', 'blocks.14.att.time_decay', 'blocks.14.att.time_first', 'blocks.14.att.time_mix_k', 'blocks.14.att.time_mix_r', 'blocks.14.att.time_mix_v', 'blocks.14.att.value.weight', 'blocks.14.ffn.key.weight', 'blocks.14.ffn.receptance.weight', 'blocks.14.ffn.time_mix_k', 'blocks.14.ffn.time_mix_r', 'blocks.14.ffn.value.weight', 'blocks.14.ln1.bias', 'blocks.14.ln1.weight', 'blocks.14.ln2.bias', 'blocks.14.ln2.weight', 'blocks.15.att.key.weight', 'blocks.15.att.output.weight', 'blocks.15.att.receptance.weight', 'blocks.15.att.time_decay', 'blocks.15.att.time_first', 'blocks.15.att.time_mix_k', 'blocks.15.att.time_mix_r', 'blocks.15.att.time_mix_v', 'blocks.15.att.value.weight', 'blocks.15.ffn.key.weight', 'blocks.15.ffn.receptance.weight', 'blocks.15.ffn.time_mix_k', 'blocks.15.ffn.time_mix_r', 'blocks.15.ffn.value.weight', 'blocks.15.ln1.bias', 'blocks.15.ln1.weight', 'blocks.15.ln2.bias', 'blocks.15.ln2.weight', 'blocks.16.att.key.weight', 'blocks.16.att.output.weight', 'blocks.16.att.receptance.weight', 'blocks.16.att.time_decay', 'blocks.16.att.time_first', 'blocks.16.att.time_mix_k', 'blocks.16.att.time_mix_r', 'blocks.16.att.time_mix_v', 'blocks.16.att.value.weight', 'blocks.16.ffn.key.weight', 'blocks.16.ffn.receptance.weight', 'blocks.16.ffn.time_mix_k', 'blocks.16.ffn.time_mix_r', 'blocks.16.ffn.value.weight', 'blocks.16.ln1.bias', 'blocks.16.ln1.weight', 'blocks.16.ln2.bias', 'blocks.16.ln2.weight', 'blocks.17.att.key.weight', 'blocks.17.att.output.weight', 'blocks.17.att.receptance.weight', 'blocks.17.att.time_decay', 'blocks.17.att.time_first', 'blocks.17.att.time_mix_k', 'blocks.17.att.time_mix_r', 'blocks.17.att.time_mix_v', 'blocks.17.att.value.weight', 'blocks.17.ffn.key.weight', 'blocks.17.ffn.receptance.weight', 'blocks.17.ffn.time_mix_k', 'blocks.17.ffn.time_mix_r', 'blocks.17.ffn.value.weight', 'blocks.17.ln1.bias', 'blocks.17.ln1.weight', 'blocks.17.ln2.bias', 'blocks.17.ln2.weight', 'blocks.18.att.key.weight', 'blocks.18.att.output.weight', 'blocks.18.att.receptance.weight', 'blocks.18.att.time_decay', 'blocks.18.att.time_first', 'blocks.18.att.time_mix_k', 'blocks.18.att.time_mix_r', 'blocks.18.att.time_mix_v', 'blocks.18.att.value.weight', 'blocks.18.ffn.key.weight', 'blocks.18.ffn.receptance.weight', 'blocks.18.ffn.time_mix_k', 'blocks.18.ffn.time_mix_r', 'blocks.18.ffn.value.weight', 'blocks.18.ln1.bias', 'blocks.18.ln1.weight', 'blocks.18.ln2.bias', 'blocks.18.ln2.weight', 'blocks.19.att.key.weight', 'blocks.19.att.output.weight', 'blocks.19.att.receptance.weight', 'blocks.19.att.time_decay', 'blocks.19.att.time_first', 'blocks.19.att.time_mix_k', 'blocks.19.att.time_mix_r', 'blocks.19.att.time_mix_v', 'blocks.19.att.value.weight', 'blocks.19.ffn.key.weight', 'blocks.19.ffn.receptance.weight', 'blocks.19.ffn.time_mix_k', 'blocks.19.ffn.time_mix_r', 'blocks.19.ffn.value.weight', 'blocks.19.ln1.bias', 'blocks.19.ln1.weight', 'blocks.19.ln2.bias', 'blocks.19.ln2.weight', 'blocks.2.att.key.weight', 'blocks.2.att.output.weight', 'blocks.2.att.receptance.weight', 'blocks.2.att.time_decay', 'blocks.2.att.time_first', 'blocks.2.att.time_mix_k', 'blocks.2.att.time_mix_r', 'blocks.2.att.time_mix_v', 'blocks.2.att.value.weight', 'blocks.2.ffn.key.weight', 'blocks.2.ffn.receptance.weight', 'blocks.2.ffn.time_mix_k', 'blocks.2.ffn.time_mix_r', 'blocks.2.ffn.value.weight', 'blocks.2.ln1.bias', 'blocks.2.ln1.weight', 'blocks.2.ln2.bias', 'blocks.2.ln2.weight', 'blocks.20.att.key.weight', 'blocks.20.att.output.weight', 'blocks.20.att.receptance.weight', 'blocks.20.att.time_decay', 'blocks.20.att.time_first', 'blocks.20.att.time_mix_k', 'blocks.20.att.time_mix_r', 'blocks.20.att.time_mix_v', 'blocks.20.att.value.weight', 'blocks.20.ffn.key.weight', 'blocks.20.ffn.receptance.weight', 'blocks.20.ffn.time_mix_k', 'blocks.20.ffn.time_mix_r', 'blocks.20.ffn.value.weight', 'blocks.20.ln1.bias', 'blocks.20.ln1.weight', 'blocks.20.ln2.bias', 'blocks.20.ln2.weight', 'blocks.21.att.key.weight', 'blocks.21.att.output.weight', 'blocks.21.att.receptance.weight', 'blocks.21.att.time_decay', 'blocks.21.att.time_first', 'blocks.21.att.time_mix_k', 'blocks.21.att.time_mix_r', 'blocks.21.att.time_mix_v', 'blocks.21.att.value.weight', 'blocks.21.ffn.key.weight', 'blocks.21.ffn.receptance.weight', 'blocks.21.ffn.time_mix_k', 'blocks.21.ffn.time_mix_r', 'blocks.21.ffn.value.weight', 'blocks.21.ln1.bias', 'blocks.21.ln1.weight', 'blocks.21.ln2.bias', 'blocks.21.ln2.weight', 'blocks.22.att.key.weight', 'blocks.22.att.output.weight', 'blocks.22.att.receptance.weight', 'blocks.22.att.time_decay', 'blocks.22.att.time_first', 'blocks.22.att.time_mix_k', 'blocks.22.att.time_mix_r', 'blocks.22.att.time_mix_v', 'blocks.22.att.value.weight', 'blocks.22.ffn.key.weight', 'blocks.22.ffn.receptance.weight', 'blocks.22.ffn.time_mix_k', 'blocks.22.ffn.time_mix_r', 'blocks.22.ffn.value.weight', 'blocks.22.ln1.bias', 'blocks.22.ln1.weight', 'blocks.22.ln2.bias', 'blocks.22.ln2.weight', 'blocks.23.att.key.weight', 'blocks.23.att.output.weight', 'blocks.23.att.receptance.weight', 'blocks.23.att.time_decay', 'blocks.23.att.time_first', 'blocks.23.att.time_mix_k', 'blocks.23.att.time_mix_r', 'blocks.23.att.time_mix_v', 'blocks.23.att.value.weight', 'blocks.23.ffn.key.weight', 'blocks.23.ffn.receptance.weight', 'blocks.23.ffn.time_mix_k', 'blocks.23.ffn.time_mix_r', 'blocks.23.ffn.value.weight', 'blocks.23.ln1.bias', 'blocks.23.ln1.weight', 'blocks.23.ln2.bias', 'blocks.23.ln2.weight', 'blocks.3.att.key.weight', 'blocks.3.att.output.weight', 'blocks.3.att.receptance.weight', 'blocks.3.att.time_decay', 'blocks.3.att.time_first', 'blocks.3.att.time_mix_k', 'blocks.3.att.time_mix_r', 'blocks.3.att.time_mix_v', 'blocks.3.att.value.weight', 'blocks.3.ffn.key.weight', 'blocks.3.ffn.receptance.weight', 'blocks.3.ffn.time_mix_k', 'blocks.3.ffn.time_mix_r', 'blocks.3.ffn.value.weight', 'blocks.3.ln1.bias', 'blocks.3.ln1.weight', 'blocks.3.ln2.bias', 'blocks.3.ln2.weight', 'blocks.4.att.key.weight', 'blocks.4.att.output.weight', 'blocks.4.att.receptance.weight', 'blocks.4.att.time_decay', 'blocks.4.att.time_first', 'blocks.4.att.time_mix_k', 'blocks.4.att.time_mix_r', 'blocks.4.att.time_mix_v', 'blocks.4.att.value.weight', 'blocks.4.ffn.key.weight', 'blocks.4.ffn.receptance.weight', 'blocks.4.ffn.time_mix_k', 'blocks.4.ffn.time_mix_r', 'blocks.4.ffn.value.weight', 'blocks.4.ln1.bias', 'blocks.4.ln1.weight', 'blocks.4.ln2.bias', 'blocks.4.ln2.weight', 'blocks.5.att.key.weight', 'blocks.5.att.output.weight', 'blocks.5.att.receptance.weight', 'blocks.5.att.time_decay', 'blocks.5.att.time_first', 'blocks.5.att.time_mix_k', 'blocks.5.att.time_mix_r', 'blocks.5.att.time_mix_v', 'blocks.5.att.value.weight', 'blocks.5.ffn.key.weight', 'blocks.5.ffn.receptance.weight', 'blocks.5.ffn.time_mix_k', 'blocks.5.ffn.time_mix_r', 'blocks.5.ffn.value.weight', 'blocks.5.ln1.bias', 'blocks.5.ln1.weight', 'blocks.5.ln2.bias', 'blocks.5.ln2.weight', 'blocks.6.att.key.weight', 'blocks.6.att.output.weight', 'blocks.6.att.receptance.weight', 'blocks.6.att.time_decay', 'blocks.6.att.time_first', 'blocks.6.att.time_mix_k', 'blocks.6.att.time_mix_r', 'blocks.6.att.time_mix_v', 'blocks.6.att.value.weight', 'blocks.6.ffn.key.weight', 'blocks.6.ffn.receptance.weight', 'blocks.6.ffn.time_mix_k', 'blocks.6.ffn.time_mix_r', 'blocks.6.ffn.value.weight', 'blocks.6.ln1.bias', 'blocks.6.ln1.weight', 'blocks.6.ln2.bias', 'blocks.6.ln2.weight', 'blocks.7.att.key.weight', 'blocks.7.att.output.weight', 'blocks.7.att.receptance.weight', 'blocks.7.att.time_decay', 'blocks.7.att.time_first', 'blocks.7.att.time_mix_k', 'blocks.7.att.time_mix_r', 'blocks.7.att.time_mix_v', 'blocks.7.att.value.weight', 'blocks.7.ffn.key.weight', 'blocks.7.ffn.receptance.weight', 'blocks.7.ffn.time_mix_k', 'blocks.7.ffn.time_mix_r', 'blocks.7.ffn.value.weight', 'blocks.7.ln1.bias', 'blocks.7.ln1.weight', 'blocks.7.ln2.bias', 'blocks.7.ln2.weight', 'blocks.8.att.key.weight', 'blocks.8.att.output.weight', 'blocks.8.att.receptance.weight', 'blocks.8.att.time_decay', 'blocks.8.att.time_first', 'blocks.8.att.time_mix_k', 'blocks.8.att.time_mix_r', 'blocks.8.att.time_mix_v', 'blocks.8.att.value.weight', 'blocks.8.ffn.key.weight', 'blocks.8.ffn.receptance.weight', 'blocks.8.ffn.time_mix_k', 'blocks.8.ffn.time_mix_r', 'blocks.8.ffn.value.weight', 'blocks.8.ln1.bias', 'blocks.8.ln1.weight', 'blocks.8.ln2.bias', 'blocks.8.ln2.weight', 'blocks.9.att.key.weight', 'blocks.9.att.output.weight', 'blocks.9.att.receptance.weight', 'blocks.9.att.time_decay', 'blocks.9.att.time_first', 'blocks.9.att.time_mix_k', 'blocks.9.att.time_mix_r', 'blocks.9.att.time_mix_v', 'blocks.9.att.value.weight', 'blocks.9.ffn.key.weight', 'blocks.9.ffn.receptance.weight', 'blocks.9.ffn.time_mix_k', 'blocks.9.ffn.time_mix_r', 'blocks.9.ffn.value.weight', 'blocks.9.ln1.bias', 'blocks.9.ln1.weight', 'blocks.9.ln2.bias', 'blocks.9.ln2.weight', 'emb.weight', 'head.weight', 'ln_out.bias', 'ln_out.weight'])"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import jax\n",
    "import jax.numpy as jnp\n",
    "from safetensors.numpy import load_file\n",
    "\n",
    "rwkv = load_file(\"../../RWKV-LM-deepspeed/RWKV-v4neo/RWKV-4-Pile-430M-20220808-8066.safetensors\")\n",
    "\n",
    "rwkv.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(1024, 1024) (1024,)\n",
      "24\n"
     ]
    }
   ],
   "source": [
    "print(rwkv[\"blocks.0.att.receptance.weight\"].shape, rwkv[\"blocks.0.att.time_decay\"].shape)\n",
    "\n",
    "BLK = \"blocks.\"\n",
    "blks = 1\n",
    "for k in rwkv.keys():\n",
    "    if k.startswith(BLK):\n",
    "        blks = max(int(k.split('.')[1]) + 1, blks)\n",
    "print(blks)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "@jax.jit\n",
    "def AT(x, state, i: int, time_mix_k, time_mix_v, time_mix_r, key, output, receptance, time_decay, time_first, value):\n",
    "    xk = x * time_mix_k + state[5*i+1].astype(time_mix_k.dtype) * (1 - time_mix_k)\n",
    "    xv = x * time_mix_v + state[5*i+1].astype(time_mix_v.dtype) * (1 - time_mix_v)\n",
    "    xr = x * time_mix_r + state[5*i+1].astype(time_mix_r.dtype) * (1 - time_mix_r)\n",
    "    state = state.at[5*i+1].set(x)\n",
    "\n",
    "    r = jax.nn.sigmoid(xr @ receptance)\n",
    "    k = (xk @ key).astype(jnp.float32)\n",
    "    v = (xv @ value).astype(jnp.float32)\n",
    "\n",
    "    aa = state[5*i+2]\n",
    "    bb = state[5*i+3]\n",
    "    pp = state[5*i+4]\n",
    "    ww = time_first + k\n",
    "    p = jnp.maximum(pp, ww)\n",
    "    e1 = jnp.exp(pp - p)\n",
    "    e2 = jnp.exp(ww - p)\n",
    "    a = e1 * aa + e2 * v\n",
    "    b = e1 * bb + e2\n",
    "    ww = pp + time_decay\n",
    "    p = jnp.maximum(ww, k)\n",
    "    e1 = jnp.exp(ww - p)\n",
    "    e2 = jnp.exp(k - p)\n",
    "    state = state.at[5*i+2].set(e1 * aa + e2 * v)\n",
    "    state = state.at[5*i+3].set(e1 * bb + e2)\n",
    "    state = state.at[5*i+4].set(p)\n",
    "\n",
    "    wkv = (a / b).astype(key.dtype)\n",
    "    return (r * wkv) @ output, state\n",
    "\n",
    "class Attention():\n",
    "    def __init__(self, sd, i, dtype=jnp.float32):\n",
    "        self.dtype = dtype\n",
    "        self.key = jax.device_put(jnp.transpose(sd[f\"blocks.{i}.att.key.weight\"].astype(dtype)))\n",
    "        self.output = jax.device_put(jnp.transpose(sd[f\"blocks.{i}.att.output.weight\"].astype(dtype) / (2 ** (i // 6))))\n",
    "        self.receptance = jax.device_put(jnp.transpose(sd[f\"blocks.{i}.att.receptance.weight\"].astype(dtype)))\n",
    "        # f32 only\n",
    "        self.time_decay = -jnp.exp(jax.device_put(sd[f\"blocks.{i}.att.time_decay\"].astype(jnp.float32)))\n",
    "        self.time_first = jax.device_put(sd[f\"blocks.{i}.att.time_first\"].astype(jnp.float32))\n",
    "        # ---\n",
    "        self.time_mix_k = jax.device_put(sd[f\"blocks.{i}.att.time_mix_k\"].astype(dtype).squeeze())\n",
    "        self.time_mix_v = jax.device_put(sd[f\"blocks.{i}.att.time_mix_v\"].astype(dtype).squeeze())\n",
    "        self.time_mix_r = jax.device_put(sd[f\"blocks.{i}.att.time_mix_r\"].astype(dtype).squeeze())\n",
    "        self.value = jax.device_put(jnp.transpose(sd[f\"blocks.{i}.att.value.weight\"].astype(dtype)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "att = Attention(rwkv, 0)\n",
    "key = jax.random.PRNGKey(0)\n",
    "x = jax.random.normal(key, (1024,), dtype=att.dtype)\n",
    "state = jnp.array([0.0]*5*1024, dtype=jnp.float32).reshape((5,1024))\n",
    "# AT(x, state, 0, att.time_mix_k, att.time_mix_v, att.time_mix_r, att.key, att.output, att.receptance, att.time_decay, att.time_first, att.value)\n",
    "%timeit AT(x, state, 0, att.time_mix_k, att.time_mix_v, att.time_mix_r, att.key, att.output, att.receptance, att.time_decay, att.time_first, att.value)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "@jax.jit\n",
    "def FFN(x, state, i: int, key, receptance, time_mix_k, time_mix_r, value):\n",
    "    xk = x * time_mix_k + state[5*i+0] * (1 - time_mix_k)\n",
    "    xr = x * time_mix_r + state[5*i+0] * (1 - time_mix_r)\n",
    "    state = state.at[5*i+0].set(x)\n",
    "\n",
    "    r = jax.nn.sigmoid(xr @ receptance)\n",
    "    k = jnp.square(jax.nn.relu(xk @ key))\n",
    "    # print(i, k, value)\n",
    "    kv = k @ value\n",
    "\n",
    "    return r * kv, state\n",
    "\n",
    "class Ffn():\n",
    "    def __init__(self, sd, i, dtype=jnp.float32):\n",
    "        self.dtype = dtype\n",
    "        self.key = jax.device_put(jnp.transpose(sd[f\"blocks.{i}.ffn.key.weight\"].astype(dtype)))\n",
    "        self.receptance = jax.device_put(jnp.transpose(sd[f\"blocks.{i}.ffn.receptance.weight\"].astype(dtype)))\n",
    "        self.time_mix_k = jax.device_put(sd[f\"blocks.{i}.ffn.time_mix_k\"].astype(dtype).squeeze())\n",
    "        self.time_mix_r = jax.device_put(sd[f\"blocks.{i}.ffn.time_mix_r\"].astype(dtype).squeeze())\n",
    "        self.value = jax.device_put(jnp.transpose(sd[f\"blocks.{i}.ffn.value.weight\"].astype(dtype)) / (2 ** (i // 6)))\n",
    "        # print(i, self.value)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ffn = Ffn(rwkv, 0)\n",
    "key = jax.random.PRNGKey(0)\n",
    "x = jax.random.normal(key, (1024,), dtype=att.dtype)\n",
    "state = jnp.array([0.0]*5*1024, dtype=jnp.float32).reshape((5,1024))\n",
    "# FFN(x, state, 0, ffn.key, ffn.receptance, ffn.time_mix_k, ffn.time_mix_r, value)\n",
    "%timeit FFN(x, state, 0, ffn.key, ffn.receptance, ffn.time_mix_k, ffn.time_mix_r, ffn.value)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "@jax.jit\n",
    "def LN(x, w, b):\n",
    "    mean = jnp.mean(x)\n",
    "    v = jnp.var(x)\n",
    "    o = x - mean\n",
    "    i = w * jax.lax.rsqrt(v + 1e-5)\n",
    "    return o * i + b\n",
    "\n",
    "class Ln():\n",
    "    def __init__(self, sd, n, dtype=jnp.float32):\n",
    "        self.dtype = dtype\n",
    "        self.weight = jax.device_put(sd[f\"{n}.weight\"].astype(dtype))\n",
    "        self.bias = jax.device_put(sd[f\"{n}.bias\"].astype(dtype))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Block():\n",
    "    def __init__(self, sd, i, dtype=jnp.float32):\n",
    "        self.dtype = dtype\n",
    "        self.ln1 = Ln(sd, f\"blocks.{i}.ln1\", dtype=dtype)\n",
    "        self.att = Attention(sd, i, dtype=dtype)\n",
    "        self.ln2 = Ln(sd, f\"blocks.{i}.ln2\", dtype=dtype)\n",
    "        self.ffn = Ffn(sd, i, dtype=dtype)\n",
    "\n",
    "def BLOCK(x, state, i, block: Block):\n",
    "    xx = LN(x, block.ln1.weight, block.ln1.bias)\n",
    "    xx, state = AT(xx, state, i, block.att.time_mix_k, block.att.time_mix_v, block.att.time_mix_r, block.att.key, block.att.output, block.att.receptance, block.att.time_decay, block.att.time_first, block.att.value)\n",
    "    x = x + xx\n",
    "    xx = LN(x, block.ln2.weight, block.ln2.bias)\n",
    "    xx, state = FFN(xx, state, i, block.ffn.key, block.ffn.receptance, block.ffn.time_mix_k, block.ffn.time_mix_r, block.ffn.value)\n",
    "    return x + xx, state\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "block = Block(rwkv, 0)\n",
    "key = jax.random.PRNGKey(0)\n",
    "x = jax.random.normal(key, (1024,), dtype=block.dtype)\n",
    "state = jnp.array([0.0]*5*1024, dtype=jnp.float32).reshape((5,1024))\n",
    "print(BLOCK(x, state, 0, block))\n",
    "# %timeit BLOCK(x, state, 0, block)\n",
    "del block, x, state, key"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "class RWKV():\n",
    "    def __init__(self, sd, blks, dtype=jnp.float32):\n",
    "        emb = jax.device_put(sd[f\"emb.weight\"].astype(dtype))\n",
    "        ln0 = Ln(sd, \"blocks.0.ln0\", dtype=dtype)\n",
    "        ln = jax.vmap(LN, (0, None, None), 0)\n",
    "        self.emb = ln(emb, ln0.weight, ln0.bias)\n",
    "        self.blocks = [Block(sd, i, dtype=dtype) for i in range(blks)]\n",
    "        self.ln_out = Ln(sd, \"ln_out\", dtype=dtype)\n",
    "        self.head = jax.device_put(sd[f\"head.weight\"].astype(dtype))\n",
    "    \n",
    "    def __call__(self, x, state):\n",
    "        x = self.emb[x]\n",
    "        for i,b in enumerate(self.blocks):\n",
    "            x, state = BLOCK(x, state, i, b)\n",
    "        x = LN(x, self.ln_out.weight, self.ln_out.bias)\n",
    "        return self.head @ x, state\n",
    "    \n",
    "    def token(self, x):\n",
    "        return self.emb[x]\n",
    "    \n",
    "    def preproc(self, x, state):\n",
    "        x = self.emb[x]\n",
    "        for i,b in enumerate(self.blocks):\n",
    "            x, state = BLOCK(x, state, i, b)\n",
    "        return state"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "from jax.config import config\n",
    "# config.update(\"jax_debug_nans\", True)\n",
    "# config.update(\"jax_debug_infs\", True)\n",
    "# config.update(\"jax_disable_jit\", True)\n",
    "\n",
    "rwkv_ = RWKV(rwkv, blks, dtype=jnp.float16)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "32.7 ms ± 239 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)\n"
     ]
    }
   ],
   "source": [
    "key = jax.random.PRNGKey(0)\n",
    "state = jnp.array(([0.0]*1024*4 + [-1e30]*1024)*blks, dtype=jnp.float32).reshape((5*blks,1024))\n",
    "# print(rwkv_.emb.shape, state[4+5])\n",
    "# rwkv_(147, state)\n",
    "%timeit rwkv_(0, state)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tokenizers\n",
    "\n",
    "tokenizer = tokenizers.Tokenizer.from_file(\"../../RWKV-LM-deepspeed/20B_tokenizer_openchatgpt.json\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Finished preproc!\n",
      "--- --- ---\n",
      "Hi! I am new in the JAX community. I am not sure if I am able to get any more information about your work. I am not sure if you are able to get any more information about your work. I am not sure if you are able to get any more information about your work. I am not sure if you are able to get any more information about your work. I am not sure if you are able to get any more information about your work. I am not sure if you are able to get any more information about your work. I am not sure if you are able to get any more information about your work. I am not sure if you are able to get any more information about your work. I am not sure if you are able to get any more information about your work. I am not sure if you are able to get any more information about your work. I am not sure if you are able to get any more information about your work. I am not sure if you are able to get any more information about your work. I am not sure if you are able to get any more information about your work. I am not sure if you are able to get any more information about your work. I am not sure if you are able to get any more information about your work. I am not sure if you are able to get any more information about your work. I am not sure if you are able to get any more information about your work. I am not sure if you are able to get any more information about your work. I am not sure if you are able to get any more information about your work. I am not sure if you are able to get any more information about your work. I am not sure if you are able to get any more information about your work. I am not sure if you are able to get any more information about your work. I am not sure if you are able to get any more information about your work. I am not sure if you are able to get any more information about your work. I am not sure if you are able to get any more information about your work. I am not sure if you are able to get any more information about your work. I am not sure if you are able to get any more information about your work. I am not sure if you are able to get any more information about your work. I am not sure if you are able to get any more information about your work. I am not sure if you are able to get any more information about your work. I am not sure if you are able to get any more information about your work. I am not sure if you are able to get any more information about your work. I am not sure if you are able to get any more information about your work. I am not sure if you are able to get any more information about your work. I am not sure if you are able to get any more information about your work. I am not sure if you are able to get any more information about your work. I am not sure if you are able to get any more information about your work. I am not sure if you are able to get any more information about your work. I am not sure if you are able to get any more information about your work. I am not sure if you are able to get any more information about your work. I am not sure if you are able to get any more information about your work. I am not sure if you are able to get any more information about your work. I am not sure if you are able to get any more information about your work. I am not sure if you are able to get any more information about your work. I am not sure if you are able to get any more information about your work. I"
     ]
    }
   ],
   "source": [
    "state = jnp.array(([0.0]*1024*4 + [-1e30]*1024)*blks, dtype=jnp.float32).reshape((5*blks,1024))\n",
    "\n",
    "context = \"\"\"Hi! I am new in the JAX community\"\"\"\n",
    "tokens = tokenizer.encode(context).ids\n",
    "\n",
    "for i in tokens[:-1]:\n",
    "    state = rwkv_.preproc(i, state)\n",
    "\n",
    "logits, state = rwkv_(tokens[-1], state)\n",
    "print(\"Finished preproc!\\n--- --- ---\")\n",
    "\n",
    "print(context, end='')\n",
    "for i in range(767):\n",
    "    token = jnp.argmax(logits)\n",
    "\n",
    "    print(tokenizer.decode([token]), end='', flush=True)\n",
    "    if token == 0:\n",
    "        break\n",
    "\n",
    "    logits, state = rwkv_(token, state)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv",
   "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.9"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "108a77ccb50483248237e4437e3c4c648b04e6f2ff2fc5096b4500deb5d70dcc"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}