gemma-q.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "import penzai\n",
    "from penzai import pz\n",
    "pz.ts.register_as_default()\n",
    "pz.ts.register_autovisualize_magic()\n",
    "pz.enable_interactive_context()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from micrlhf.llama import LlamaTransformer\n",
    "from micrlhf.scan import sequential_to_scan\n",
    "llama = LlamaTransformer.from_pretrained(\"models/gemma-2-2b-it-q4_k_s.gguf\", device_map=\"tpu:0\",\n",
    "                                         from_type=\"gemma2\",\n",
    "                                         load_on_cpu=True,\n",
    "                                         )\n",
    "llama = sequential_to_scan(llama)\n",
    "llama = llama.to_tpu()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"NousResearch/gemma-2b-it-tokenizer\")\n",
    "prompt = tokenizer.apply_chat_template([\n",
    "    {\"role\": \"user\", \"content\": \"Who are you?\"},\n",
    "    {\"role\": \"assistant\", \"content\": \"Hello, I am a language model that exists and stuff. How can I help you today?\"},\n",
    "] * 1_000, tokenize=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from penzai.toolshed import jit_wrapper\n",
    "from micrlhf.flash import flashify\n",
    "tokens = pz.nx.wrap([tokenizer.encode(prompt)[:128]] * 128, \"batch\", \"seq\")\n",
    "inputs = llama.inputs.from_basic_segments(tokens)\n",
    "llama_jitted = jit_wrapper.Jitted(llama)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "ename": "TypeError",
     "evalue": "shard_map in_specs argument must be a pytree of `jax.sharding.PartitionSpec` instances, but:\n\n  in_specs[3] is mp of type str, \n\n  in_specs[6] is mp of type str, \n\n  in_specs[9] is mp of type str, \n\n  in_specs[12] is mp of type str, \n\nCheck the in_specs values passed to shard_map.",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[5], line 8\u001b[0m\n\u001b[1;32m      6\u001b[0m     loss \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m-\u001b[39mpz\u001b[38;5;241m.\u001b[39mnx\u001b[38;5;241m.\u001b[39mnmap(\u001b[38;5;28;01mlambda\u001b[39;00m l, t: jnp\u001b[38;5;241m.\u001b[39mtake_along_axis(jax\u001b[38;5;241m.\u001b[39mnn\u001b[38;5;241m.\u001b[39mlog_softmax(l[:\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m], \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m), t[\u001b[38;5;241m1\u001b[39m:, \u001b[38;5;28;01mNone\u001b[39;00m], \u001b[38;5;241m1\u001b[39m)\u001b[38;5;241m.\u001b[39mmean())(logits\u001b[38;5;241m.\u001b[39muntag(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mseq\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mvocabulary\u001b[39m\u001b[38;5;124m\"\u001b[39m), tokens\u001b[38;5;241m.\u001b[39muntag(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mseq\u001b[39m\u001b[38;5;124m\"\u001b[39m))\n\u001b[1;32m      7\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m loss\n\u001b[0;32m----> 8\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[43mlfn\u001b[49m\u001b[43m(\u001b[49m\u001b[43mllama_jitted\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43minputs\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241m.\u001b[39munwrap(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mbatch\u001b[39m\u001b[38;5;124m\"\u001b[39m))\n",
      "    \u001b[0;31m[... skipping hidden 11 frame]\u001b[0m\n",
      "Cell \u001b[0;32mIn[5], line 5\u001b[0m, in \u001b[0;36mlfn\u001b[0;34m(llama_jitted, inputs)\u001b[0m\n\u001b[1;32m      3\u001b[0m \u001b[38;5;129m@jax\u001b[39m\u001b[38;5;241m.\u001b[39mjit\n\u001b[1;32m      4\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mlfn\u001b[39m(llama_jitted, inputs):\n\u001b[0;32m----> 5\u001b[0m     logits \u001b[38;5;241m=\u001b[39m \u001b[43mllama_jitted\u001b[49m\u001b[43m(\u001b[49m\u001b[43minputs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m      6\u001b[0m     loss \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m-\u001b[39mpz\u001b[38;5;241m.\u001b[39mnx\u001b[38;5;241m.\u001b[39mnmap(\u001b[38;5;28;01mlambda\u001b[39;00m l, t: jnp\u001b[38;5;241m.\u001b[39mtake_along_axis(jax\u001b[38;5;241m.\u001b[39mnn\u001b[38;5;241m.\u001b[39mlog_softmax(l[:\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m], \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m), t[\u001b[38;5;241m1\u001b[39m:, \u001b[38;5;28;01mNone\u001b[39;00m], \u001b[38;5;241m1\u001b[39m)\u001b[38;5;241m.\u001b[39mmean())(logits\u001b[38;5;241m.\u001b[39muntag(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mseq\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mvocabulary\u001b[39m\u001b[38;5;124m\"\u001b[39m), tokens\u001b[38;5;241m.\u001b[39muntag(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mseq\u001b[39m\u001b[38;5;124m\"\u001b[39m))\n\u001b[1;32m      7\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m loss\n",
      "File \u001b[0;32m~/.cache/pypoetry/virtualenvs/micrlhf-progress-a058ydGG-py3.10/lib/python3.10/site-packages/penzai/toolshed/jit_wrapper.py:62\u001b[0m, in \u001b[0;36mJitted.__call__\u001b[0;34m(self, argument)\u001b[0m\n\u001b[1;32m     60\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m__call__\u001b[39m(\u001b[38;5;28mself\u001b[39m, argument: Any, \u001b[38;5;241m/\u001b[39m) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Any:\n\u001b[1;32m     61\u001b[0m   \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m---> 62\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43m_flat_jit_call_layer\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbody\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43margument\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     63\u001b[0m   \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mTypeError\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m exc:\n\u001b[1;32m     64\u001b[0m     impls \u001b[38;5;241m=\u001b[39m pz\u001b[38;5;241m.\u001b[39mselect(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mbody)\u001b[38;5;241m.\u001b[39mat_instances_of(pz\u001b[38;5;241m.\u001b[39mde\u001b[38;5;241m.\u001b[39mEffectRuntimeImpl)\n",
      "    \u001b[0;31m[... skipping hidden 11 frame]\u001b[0m\n",
      "File \u001b[0;32m~/.cache/pypoetry/virtualenvs/micrlhf-progress-a058ydGG-py3.10/lib/python3.10/site-packages/penzai/toolshed/jit_wrapper.py:37\u001b[0m, in \u001b[0;36m_flat_jit_call_layer\u001b[0;34m(layer, arg)\u001b[0m\n\u001b[1;32m     34\u001b[0m \u001b[38;5;129m@jax\u001b[39m\u001b[38;5;241m.\u001b[39mjit\n\u001b[1;32m     35\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m_flat_jit_call_layer\u001b[39m(layer, arg):\n\u001b[1;32m     36\u001b[0m \u001b[38;5;250m  \u001b[39m\u001b[38;5;124;03m\"\"\"Helper to call a callable pytree with an argument.\"\"\"\u001b[39;00m\n\u001b[0;32m---> 37\u001b[0m   \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mlayer\u001b[49m\u001b[43m(\u001b[49m\u001b[43marg\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/.cache/pypoetry/virtualenvs/micrlhf-progress-a058ydGG-py3.10/lib/python3.10/site-packages/penzai/core/layer.py:82\u001b[0m, in \u001b[0;36mchecked_layer_call.<locals>.wrapper\u001b[0;34m(self, argument)\u001b[0m\n\u001b[1;32m     80\u001b[0m output_structure \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39moutput_structure()\n\u001b[1;32m     81\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m jax\u001b[38;5;241m.\u001b[39mnamed_scope(\u001b[38;5;28mtype\u001b[39m(\u001b[38;5;28mself\u001b[39m)\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__name__\u001b[39m):\n\u001b[0;32m---> 82\u001b[0m   result \u001b[38;5;241m=\u001b[39m \u001b[43mfunc\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43margument\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     83\u001b[0m shapecheck\u001b[38;5;241m.\u001b[39mcheck_structure(\n\u001b[1;32m     84\u001b[0m     result,\n\u001b[1;32m     85\u001b[0m     pattern\u001b[38;5;241m=\u001b[39moutput_structure,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     90\u001b[0m     ),\n\u001b[1;32m     91\u001b[0m )\n\u001b[1;32m     92\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m result\n",
      "File \u001b[0;32m~/micrlhf-progress/micrlhf/llama.py:286\u001b[0m, in \u001b[0;36mLlamaTransformer.__call__\u001b[0;34m(self, inputs)\u001b[0m\n\u001b[1;32m    284\u001b[0m \u001b[38;5;129m@pz\u001b[39m\u001b[38;5;241m.\u001b[39mchecked_layer_call\n\u001b[1;32m    285\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m__call__\u001b[39m(\u001b[38;5;28mself\u001b[39m, inputs: LlamaInputs) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m pz\u001b[38;5;241m.\u001b[39mnx\u001b[38;5;241m.\u001b[39mNamedArray:\n\u001b[0;32m--> 286\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbody\u001b[49m\u001b[43m(\u001b[49m\u001b[43m(\u001b[49m\u001b[43minputs\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtokens\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43minputs\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mpositions\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43minputs\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mattention_mask\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/.cache/pypoetry/virtualenvs/micrlhf-progress-a058ydGG-py3.10/lib/python3.10/site-packages/penzai/core/layer.py:82\u001b[0m, in \u001b[0;36mchecked_layer_call.<locals>.wrapper\u001b[0;34m(self, argument)\u001b[0m\n\u001b[1;32m     80\u001b[0m output_structure \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39moutput_structure()\n\u001b[1;32m     81\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m jax\u001b[38;5;241m.\u001b[39mnamed_scope(\u001b[38;5;28mtype\u001b[39m(\u001b[38;5;28mself\u001b[39m)\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__name__\u001b[39m):\n\u001b[0;32m---> 82\u001b[0m   result \u001b[38;5;241m=\u001b[39m \u001b[43mfunc\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43margument\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     83\u001b[0m shapecheck\u001b[38;5;241m.\u001b[39mcheck_structure(\n\u001b[1;32m     84\u001b[0m     result,\n\u001b[1;32m     85\u001b[0m     pattern\u001b[38;5;241m=\u001b[39moutput_structure,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     90\u001b[0m     ),\n\u001b[1;32m     91\u001b[0m )\n\u001b[1;32m     92\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m result\n",
      "File \u001b[0;32m~/.cache/pypoetry/virtualenvs/micrlhf-progress-a058ydGG-py3.10/lib/python3.10/site-packages/penzai/data_effects/side_input.py:146\u001b[0m, in \u001b[0;36mWithSideInputsFromInputTuple.__call__\u001b[0;34m(self, argument)\u001b[0m\n\u001b[1;32m    136\u001b[0m impls \u001b[38;5;241m=\u001b[39m {\n\u001b[1;32m    137\u001b[0m     tag: SideInputEffectImpl(_value\u001b[38;5;241m=\u001b[39mval, _handler_id\u001b[38;5;241m=\u001b[39m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mhandler_id)\n\u001b[1;32m    138\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m tag, val \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mzip\u001b[39m(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mside_input_tags, side_inputs)\n\u001b[1;32m    139\u001b[0m }\n\u001b[1;32m    140\u001b[0m handled_body \u001b[38;5;241m=\u001b[39m (\n\u001b[1;32m    141\u001b[0m     selectors\u001b[38;5;241m.\u001b[39mselect(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mbody)\n\u001b[1;32m    142\u001b[0m     \u001b[38;5;241m.\u001b[39mat_instances_of(HandledSideInputRef)\n\u001b[1;32m    143\u001b[0m     \u001b[38;5;241m.\u001b[39mwhere(\u001b[38;5;28;01mlambda\u001b[39;00m ref: ref\u001b[38;5;241m.\u001b[39mhandler_id \u001b[38;5;241m==\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mhandler_id)\n\u001b[1;32m    144\u001b[0m     \u001b[38;5;241m.\u001b[39mapply(\u001b[38;5;28;01mlambda\u001b[39;00m ref: impls[ref\u001b[38;5;241m.\u001b[39mtag])\n\u001b[1;32m    145\u001b[0m )\n\u001b[0;32m--> 146\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mhandled_body\u001b[49m\u001b[43m(\u001b[49m\u001b[43minner_arg\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/.cache/pypoetry/virtualenvs/micrlhf-progress-a058ydGG-py3.10/lib/python3.10/site-packages/penzai/nn/grouping.py:79\u001b[0m, in \u001b[0;36mSequential.__call__\u001b[0;34m(self, value)\u001b[0m\n\u001b[1;32m     77\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m i, layer \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28menumerate\u001b[39m(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39msublayers):\n\u001b[1;32m     78\u001b[0m   \u001b[38;5;28;01mwith\u001b[39;00m jax\u001b[38;5;241m.\u001b[39mnamed_scope(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mi\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m):\n\u001b[0;32m---> 79\u001b[0m     value \u001b[38;5;241m=\u001b[39m \u001b[43mlayer\u001b[49m\u001b[43m(\u001b[49m\u001b[43mvalue\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     80\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m value\n",
      "File \u001b[0;32m~/micrlhf-progress/micrlhf/scan.py:35\u001b[0m, in \u001b[0;36mScanSequential.__call__\u001b[0;34m(self, inputs)\u001b[0m\n\u001b[1;32m     33\u001b[0m     h_ \u001b[38;5;241m=\u001b[39m pz\u001b[38;5;241m.\u001b[39mnx\u001b[38;5;241m.\u001b[39mwrap(h_\u001b[38;5;241m.\u001b[39munwrap(\u001b[38;5;241m*\u001b[39mh\u001b[38;5;241m.\u001b[39mnamed_axes\u001b[38;5;241m.\u001b[39mkeys()), \u001b[38;5;241m*\u001b[39mh\u001b[38;5;241m.\u001b[39mnamed_axes\u001b[38;5;241m.\u001b[39mkeys())\n\u001b[1;32m     34\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m h_, \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[0;32m---> 35\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mjax\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mlax\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mscan\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m     36\u001b[0m \u001b[43m    \u001b[49m\u001b[43mscanner\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     37\u001b[0m \u001b[43m    \u001b[49m\u001b[43minputs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     38\u001b[0m \u001b[43m    \u001b[49m\u001b[43mlayer_nx\u001b[49m\u001b[43m)\u001b[49m[\u001b[38;5;241m0\u001b[39m]\n",
      "    \u001b[0;31m[... skipping hidden 9 frame]\u001b[0m\n",
      "File \u001b[0;32m~/micrlhf-progress/micrlhf/scan.py:32\u001b[0m, in \u001b[0;36mScanSequential.__call__.<locals>.scanner\u001b[0;34m(h, l)\u001b[0m\n\u001b[1;32m     30\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mscanner\u001b[39m(h, l):\n\u001b[1;32m     31\u001b[0m     l_ \u001b[38;5;241m=\u001b[39m eqx\u001b[38;5;241m.\u001b[39mcombine(l, layer_base, is_leaf\u001b[38;5;241m=\u001b[39mis_nx)\n\u001b[0;32m---> 32\u001b[0m     h_ \u001b[38;5;241m=\u001b[39m \u001b[43ml_\u001b[49m\u001b[43m(\u001b[49m\u001b[43mh\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     33\u001b[0m     h_ \u001b[38;5;241m=\u001b[39m pz\u001b[38;5;241m.\u001b[39mnx\u001b[38;5;241m.\u001b[39mwrap(h_\u001b[38;5;241m.\u001b[39munwrap(\u001b[38;5;241m*\u001b[39mh\u001b[38;5;241m.\u001b[39mnamed_axes\u001b[38;5;241m.\u001b[39mkeys()), \u001b[38;5;241m*\u001b[39mh\u001b[38;5;241m.\u001b[39mnamed_axes\u001b[38;5;241m.\u001b[39mkeys())\n\u001b[1;32m     34\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m h_, \u001b[38;5;28;01mNone\u001b[39;00m\n",
      "File \u001b[0;32m~/.cache/pypoetry/virtualenvs/micrlhf-progress-a058ydGG-py3.10/lib/python3.10/site-packages/penzai/nn/grouping.py:79\u001b[0m, in \u001b[0;36mSequential.__call__\u001b[0;34m(self, value)\u001b[0m\n\u001b[1;32m     77\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m i, layer \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28menumerate\u001b[39m(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39msublayers):\n\u001b[1;32m     78\u001b[0m   \u001b[38;5;28;01mwith\u001b[39;00m jax\u001b[38;5;241m.\u001b[39mnamed_scope(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mi\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m):\n\u001b[0;32m---> 79\u001b[0m     value \u001b[38;5;241m=\u001b[39m \u001b[43mlayer\u001b[49m\u001b[43m(\u001b[49m\u001b[43mvalue\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     80\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m value\n",
      "File \u001b[0;32m~/.cache/pypoetry/virtualenvs/micrlhf-progress-a058ydGG-py3.10/lib/python3.10/site-packages/penzai/nn/combinators.py:59\u001b[0m, in \u001b[0;36mResidual.__call__\u001b[0;34m(self, value)\u001b[0m\n\u001b[1;32m     50\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m__call__\u001b[39m(\u001b[38;5;28mself\u001b[39m, value: Any) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Any:\n\u001b[1;32m     51\u001b[0m \u001b[38;5;250m  \u001b[39m\u001b[38;5;124;03m\"\"\"Runs each of the sublayers in sequence, then adds back the original input.\u001b[39;00m\n\u001b[1;32m     52\u001b[0m \n\u001b[1;32m     53\u001b[0m \u001b[38;5;124;03m  Args:\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     57\u001b[0m \u001b[38;5;124;03m    The sum of the input to the residual block and the output of the child.\u001b[39;00m\n\u001b[1;32m     58\u001b[0m \u001b[38;5;124;03m  \"\"\"\u001b[39;00m\n\u001b[0;32m---> 59\u001b[0m   delta_value \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mdelta\u001b[49m\u001b[43m(\u001b[49m\u001b[43mvalue\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     60\u001b[0m   \u001b[38;5;28;01mreturn\u001b[39;00m delta_value \u001b[38;5;241m+\u001b[39m value\n",
      "File \u001b[0;32m~/.cache/pypoetry/virtualenvs/micrlhf-progress-a058ydGG-py3.10/lib/python3.10/site-packages/penzai/nn/grouping.py:79\u001b[0m, in \u001b[0;36mSequential.__call__\u001b[0;34m(self, value)\u001b[0m\n\u001b[1;32m     77\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m i, layer \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28menumerate\u001b[39m(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39msublayers):\n\u001b[1;32m     78\u001b[0m   \u001b[38;5;28;01mwith\u001b[39;00m jax\u001b[38;5;241m.\u001b[39mnamed_scope(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mi\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m):\n\u001b[0;32m---> 79\u001b[0m     value \u001b[38;5;241m=\u001b[39m \u001b[43mlayer\u001b[49m\u001b[43m(\u001b[49m\u001b[43mvalue\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     80\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m value\n",
      "File \u001b[0;32m~/.cache/pypoetry/virtualenvs/micrlhf-progress-a058ydGG-py3.10/lib/python3.10/site-packages/penzai/nn/attention.py:129\u001b[0m, in \u001b[0;36mAttention.__call__\u001b[0;34m(self, x)\u001b[0m\n\u001b[1;32m    119\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m__call__\u001b[39m(\u001b[38;5;28mself\u001b[39m, x: named_axes\u001b[38;5;241m.\u001b[39mNamedArray) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m named_axes\u001b[38;5;241m.\u001b[39mNamedArray:\n\u001b[1;32m    120\u001b[0m \u001b[38;5;250m  \u001b[39m\u001b[38;5;124;03m\"\"\"Runs the attention computation.\u001b[39;00m\n\u001b[1;32m    121\u001b[0m \n\u001b[1;32m    122\u001b[0m \u001b[38;5;124;03m  Args:\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    127\u001b[0m \u001b[38;5;124;03m    The final output of the ``attn_value_to_output`` sublayer.\u001b[39;00m\n\u001b[1;32m    128\u001b[0m \u001b[38;5;124;03m  \"\"\"\u001b[39;00m\n\u001b[0;32m--> 129\u001b[0m   query \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43minput_to_query\u001b[49m\u001b[43m(\u001b[49m\u001b[43mx\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    130\u001b[0m   key \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39minput_to_key(x)\n\u001b[1;32m    131\u001b[0m   value \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39minput_to_value(x)\n",
      "File \u001b[0;32m~/.cache/pypoetry/virtualenvs/micrlhf-progress-a058ydGG-py3.10/lib/python3.10/site-packages/penzai/nn/grouping.py:79\u001b[0m, in \u001b[0;36mSequential.__call__\u001b[0;34m(self, value)\u001b[0m\n\u001b[1;32m     77\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m i, layer \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28menumerate\u001b[39m(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39msublayers):\n\u001b[1;32m     78\u001b[0m   \u001b[38;5;28;01mwith\u001b[39;00m jax\u001b[38;5;241m.\u001b[39mnamed_scope(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mi\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m):\n\u001b[0;32m---> 79\u001b[0m     value \u001b[38;5;241m=\u001b[39m \u001b[43mlayer\u001b[49m\u001b[43m(\u001b[49m\u001b[43mvalue\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     80\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m value\n",
      "File \u001b[0;32m~/.cache/pypoetry/virtualenvs/micrlhf-progress-a058ydGG-py3.10/lib/python3.10/site-packages/penzai/core/layer.py:82\u001b[0m, in \u001b[0;36mchecked_layer_call.<locals>.wrapper\u001b[0;34m(self, argument)\u001b[0m\n\u001b[1;32m     80\u001b[0m output_structure \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39moutput_structure()\n\u001b[1;32m     81\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m jax\u001b[38;5;241m.\u001b[39mnamed_scope(\u001b[38;5;28mtype\u001b[39m(\u001b[38;5;28mself\u001b[39m)\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__name__\u001b[39m):\n\u001b[0;32m---> 82\u001b[0m   result \u001b[38;5;241m=\u001b[39m \u001b[43mfunc\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43margument\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     83\u001b[0m shapecheck\u001b[38;5;241m.\u001b[39mcheck_structure(\n\u001b[1;32m     84\u001b[0m     result,\n\u001b[1;32m     85\u001b[0m     pattern\u001b[38;5;241m=\u001b[39moutput_structure,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     90\u001b[0m     ),\n\u001b[1;32m     91\u001b[0m )\n\u001b[1;32m     92\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m result\n",
      "File \u001b[0;32m~/micrlhf-progress/micrlhf/quantizers.py:103\u001b[0m, in \u001b[0;36mQuantizedLinear.__call__\u001b[0;34m(self, inputs)\u001b[0m\n\u001b[1;32m    101\u001b[0m inputs \u001b[38;5;241m=\u001b[39m pz\u001b[38;5;241m.\u001b[39mnx\u001b[38;5;241m.\u001b[39mnmap(jnp\u001b[38;5;241m.\u001b[39mravel)(inputs\u001b[38;5;241m.\u001b[39muntag(\u001b[38;5;241m*\u001b[39m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39min_features\u001b[38;5;241m.\u001b[39mkeys()))\u001b[38;5;241m.\u001b[39mtag(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124min_features\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m    102\u001b[0m inputs \u001b[38;5;241m=\u001b[39m pz\u001b[38;5;241m.\u001b[39mnx\u001b[38;5;241m.\u001b[39mnmap(jnp\u001b[38;5;241m.\u001b[39mravel)(inputs\u001b[38;5;241m.\u001b[39muntag(\u001b[38;5;241m*\u001b[39mbatch_shape\u001b[38;5;241m.\u001b[39mkeys()))\u001b[38;5;241m.\u001b[39mtag(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mbatch\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m--> 103\u001b[0m outputs \u001b[38;5;241m=\u001b[39m pz\u001b[38;5;241m.\u001b[39mnx\u001b[38;5;241m.\u001b[39mwrap(\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mquant_linear\u001b[49m\u001b[43m(\u001b[49m\u001b[43minputs\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43munwrap\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mbatch\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43min_features\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mbatch\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mout_features\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m    104\u001b[0m outputs \u001b[38;5;241m=\u001b[39m pz\u001b[38;5;241m.\u001b[39mnx\u001b[38;5;241m.\u001b[39mnmap(\u001b[38;5;28;01mlambda\u001b[39;00m x: x\u001b[38;5;241m.\u001b[39mreshape(\u001b[38;5;241m*\u001b[39mbatch_shape\u001b[38;5;241m.\u001b[39mvalues()))(outputs\u001b[38;5;241m.\u001b[39muntag(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mbatch\u001b[39m\u001b[38;5;124m\"\u001b[39m))\u001b[38;5;241m.\u001b[39mtag(\u001b[38;5;241m*\u001b[39mbatch_shape\u001b[38;5;241m.\u001b[39mkeys())\n\u001b[1;32m    105\u001b[0m outputs \u001b[38;5;241m=\u001b[39m pz\u001b[38;5;241m.\u001b[39mnx\u001b[38;5;241m.\u001b[39mnmap(\u001b[38;5;28;01mlambda\u001b[39;00m x: x\u001b[38;5;241m.\u001b[39mreshape(\u001b[38;5;241m*\u001b[39m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mout_features\u001b[38;5;241m.\u001b[39mvalues()))(outputs\u001b[38;5;241m.\u001b[39muntag(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mout_features\u001b[39m\u001b[38;5;124m\"\u001b[39m))\u001b[38;5;241m.\u001b[39mtag(\u001b[38;5;241m*\u001b[39m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mout_features\u001b[38;5;241m.\u001b[39mkeys())\n",
      "File \u001b[0;32m~/micrlhf-progress/micrlhf/quantizers.py:373\u001b[0m, in \u001b[0;36mLinearQuantizedTranspose.quant_linear\u001b[0;34m(self, inputs)\u001b[0m\n\u001b[1;32m    371\u001b[0m params \u001b[38;5;241m=\u001b[39m (t\u001b[38;5;241m.\u001b[39munwrap(\u001b[38;5;241m*\u001b[39m((ia, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mquant_group\u001b[39m\u001b[38;5;124m\"\u001b[39m, oa) \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mis_transposed \u001b[38;5;28;01melse\u001b[39;00m (\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mquant_group\u001b[39m\u001b[38;5;124m\"\u001b[39m, ia, oa))) \u001b[38;5;28;01mfor\u001b[39;00m t \u001b[38;5;129;01min\u001b[39;00m params)\n\u001b[1;32m    372\u001b[0m mesh \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mmesh\n\u001b[0;32m--> 373\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mmatmul_fast\u001b[49m\u001b[43m(\u001b[49m\u001b[43minputs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mparams\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mmesh\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mmesh\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mkernel\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mkernel\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    374\u001b[0m \u001b[43m                   \u001b[49m\u001b[43mbatch_axis\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbatch_axis\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43min_axis\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43min_axis\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mout_axis\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mout_axis\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    375\u001b[0m \u001b[43m                   \u001b[49m\u001b[43mis_transpose\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mis_transposed\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/micrlhf-progress/micrlhf/quantizers.py:313\u001b[0m, in \u001b[0;36mmatmul_fast\u001b[0;34m(inputs, kernel, mesh, batch_axis, in_axis, out_axis, is_transpose, *tensors)\u001b[0m\n\u001b[1;32m    310\u001b[0m         outputs \u001b[38;5;241m=\u001b[39m jax\u001b[38;5;241m.\u001b[39mlax\u001b[38;5;241m.\u001b[39mpsum(outputs, axis_name\u001b[38;5;241m=\u001b[39min_axis)\n\u001b[1;32m    311\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m outputs\n\u001b[0;32m--> 313\u001b[0m result \u001b[38;5;241m=\u001b[39m \u001b[43mjax\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mexperimental\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mshard_map\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mshard_map\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    314\u001b[0m \u001b[43m    \u001b[49m\u001b[43mkernel_call\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    315\u001b[0m \u001b[43m    \u001b[49m\u001b[43mmesh\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mmesh\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    316\u001b[0m \u001b[43m    \u001b[49m\u001b[43min_specs\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m(\u001b[49m\n\u001b[1;32m    317\u001b[0m \u001b[43m        \u001b[49m\u001b[43mP\u001b[49m\u001b[43m(\u001b[49m\u001b[43mbatch_axis\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43min_axis\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    318\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m+\u001b[39;49m\u001b[43m \u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    319\u001b[0m \u001b[43m        \u001b[49m\u001b[43mP\u001b[49m\u001b[43m(\u001b[49m\u001b[43min_axis\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43;01mNone\u001b[39;49;00m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mout_axis\u001b[49m\u001b[43m)\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43;01mif\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[38;5;129;43;01mnot\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mis_transpose\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43;01melse\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mP\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43;01mNone\u001b[39;49;00m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43min_axis\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mout_axis\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    320\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43m \u001b[49m\u001b[38;5;28;43mlen\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43mtensors\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    321\u001b[0m \u001b[43m    \u001b[49m\u001b[43mout_specs\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mP\u001b[49m\u001b[43m(\u001b[49m\u001b[43mbatch_axis\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mout_axis\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    322\u001b[0m \u001b[43m    \u001b[49m\u001b[43mcheck_rep\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mFalse\u001b[39;49;00m\n\u001b[1;32m    323\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m(inputs, \u001b[38;5;241m*\u001b[39mtensors)\n\u001b[1;32m    324\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m x_pad \u001b[38;5;129;01mor\u001b[39;00m y_pad:\n\u001b[1;32m    325\u001b[0m     result \u001b[38;5;241m=\u001b[39m result\u001b[38;5;241m.\u001b[39mreshape(batch_mesh, inputs\u001b[38;5;241m.\u001b[39mshape[\u001b[38;5;241m0\u001b[39m] \u001b[38;5;241m/\u001b[39m\u001b[38;5;241m/\u001b[39m batch_mesh, out_mesh, \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m\n\u001b[1;32m    326\u001b[0m                         )[:, :per_block_size, :, :per_mp_output_size]\n",
      "    \u001b[0;31m[... skipping hidden 3 frame]\u001b[0m\n",
      "File \u001b[0;32m~/.cache/pypoetry/virtualenvs/micrlhf-progress-a058ydGG-py3.10/lib/python3.10/site-packages/jax/experimental/shard_map.py:255\u001b[0m, in \u001b[0;36m_check_specs\u001b[0;34m(error_type, specs, auto)\u001b[0m\n\u001b[1;32m    249\u001b[0m           msgs\u001b[38;5;241m.\u001b[39mappend(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m  \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mprefix\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m_specs\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mkeystr(key)\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m refers to \u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;28mrepr\u001b[39m(name)\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m    250\u001b[0m   \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(\n\u001b[1;32m    251\u001b[0m       \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mshard_map \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mprefix\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m_specs argument cannot refer to an axis \u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    252\u001b[0m       \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mmarked auto (\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mauto\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m), but:\u001b[39m\u001b[38;5;130;01m\\n\u001b[39;00m\u001b[38;5;130;01m\\n\u001b[39;00m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    253\u001b[0m       \u001b[38;5;241m+\u001b[39m \u001b[38;5;124m'\u001b[39m\u001b[38;5;130;01m\\n\u001b[39;00m\u001b[38;5;130;01m\\n\u001b[39;00m\u001b[38;5;124m'\u001b[39m\u001b[38;5;241m.\u001b[39mjoin(msgs) \u001b[38;5;241m+\u001b[39m \u001b[38;5;124m'\u001b[39m\u001b[38;5;130;01m\\n\u001b[39;00m\u001b[38;5;130;01m\\n\u001b[39;00m\u001b[38;5;124m'\u001b[39m\n\u001b[1;32m    254\u001b[0m       \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mCheck the \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mprefix\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m_specs values passed to shard_map.\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m--> 255\u001b[0m \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mTypeError\u001b[39;00m(\n\u001b[1;32m    256\u001b[0m     \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mshard_map \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mprefix\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m_specs argument must be a pytree of \u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    257\u001b[0m     \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m`jax.sharding.PartitionSpec` instances, but:\u001b[39m\u001b[38;5;130;01m\\n\u001b[39;00m\u001b[38;5;130;01m\\n\u001b[39;00m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    258\u001b[0m     \u001b[38;5;241m+\u001b[39m \u001b[38;5;124m'\u001b[39m\u001b[38;5;130;01m\\n\u001b[39;00m\u001b[38;5;130;01m\\n\u001b[39;00m\u001b[38;5;124m'\u001b[39m\u001b[38;5;241m.\u001b[39mjoin(msgs) \u001b[38;5;241m+\u001b[39m \u001b[38;5;124m'\u001b[39m\u001b[38;5;130;01m\\n\u001b[39;00m\u001b[38;5;130;01m\\n\u001b[39;00m\u001b[38;5;124m'\u001b[39m\n\u001b[1;32m    259\u001b[0m     \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mCheck the \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mprefix\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m_specs values passed to shard_map.\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n",
      "\u001b[0;31mTypeError\u001b[0m: shard_map in_specs argument must be a pytree of `jax.sharding.PartitionSpec` instances, but:\n\n  in_specs[3] is mp of type str, \n\n  in_specs[6] is mp of type str, \n\n  in_specs[9] is mp of type str, \n\n  in_specs[12] is mp of type str, \n\nCheck the in_specs values passed to shard_map."
     ]
    }
   ],
   "source": [
    "import jax\n",
    "import jax.numpy as jnp\n",
    "@jax.jit\n",
    "def lfn(llama_jitted, inputs):\n",
    "    logits = llama_jitted(inputs)\n",
    "    loss = -pz.nx.nmap(lambda l, t: jnp.take_along_axis(jax.nn.log_softmax(l[:-1], -1), t[1:, None], 1).mean())(logits.untag(\"seq\", \"vocabulary\"), tokens.untag(\"seq\"))\n",
    "    return loss\n",
    "print(lfn(llama_jitted, inputs).unwrap(\"batch\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import jax.numpy as jnp\n",
    "import numpy as np\n",
    "import jax\n",
    "from functools import partial\n",
    "from jax.experimental import pallas as pl\n",
    "from jax.experimental.pallas import tpu as pltpu\n",
    "from jax.sharding import PartitionSpec as P\n",
    "\n",
    "\n",
    "compute = lambda f: (f(jnp.arange(256, dtype=jnp.uint8)), (f(jnp.arange(256, dtype=jnp.uint8).astype(jnp.int8))).astype(jnp.uint8))\n",
    "\n",
    "# https://github.com/99991/pygguf/blob/829886d0726c89c6f6c0d8c39b0d507ec1604077/gguf.py#L209\n",
    "def matmul_4bit(inputs, scale_factors, scale_offsets, qs1, qs2, based=False):\n",
    "    num_blocks = scale_factors.shape[0]\n",
    "    # assert num_blocks == scale_offsets.shape[0] == qs1.shape[0] == qs2.shape[0] == 8\n",
    "\n",
    "    sr = jax.lax.shift_right_logical\n",
    "\n",
    "    def switch(x):\n",
    "        x = x.astype(jnp.float16)\n",
    "        # x = x.astype(jnp.float32) * 10\n",
    "        x = x.astype(jnp.float32)\n",
    "        return x\n",
    "    scale_factors = switch(scale_factors)\n",
    "    scale_offsets = switch(scale_offsets)\n",
    "    if based:\n",
    "        scale_factors = scale_factors.transpose(1, 0, 2).reshape(1, 1, num_blocks, -1).astype(jnp.float32)\n",
    "        scale_offsets = scale_offsets.transpose(1, 0, 2).reshape(1, 1, num_blocks, -1).astype(jnp.float32)\n",
    "        qs1 = qs1.transpose(1, 0, 2).reshape(12, 1, num_blocks, -1)\n",
    "        qs2 = qs2.transpose(1, 0, 2).reshape(4, 32, num_blocks, -1)\n",
    "    else:\n",
    "        scale_factors = scale_factors.transpose(0, 2, 1).reshape(-1, 1, 1).astype(jnp.float32)\n",
    "        scale_offsets = scale_offsets.transpose(0, 2, 1).reshape(-1, 1, 1).astype(jnp.float32)\n",
    "        qs1 = qs1.transpose(0, 2, 1).reshape(-1, 12, 1)\n",
    "        qs2 = qs2.transpose(0, 2, 1).reshape(-1, 4, 32)\n",
    "\n",
    "    qs1 = qs1.astype(jnp.int32)\n",
    "    qs2 = qs2.astype(jnp.int32)\n",
    "    i8tou8 = lambda x: jnp.where(x < 0, 256 + x, x)\n",
    "    qs1 = i8tou8(qs1)\n",
    "    qs2 = i8tou8(qs2)\n",
    "\n",
    "\n",
    "    # max 63\n",
    "    if based:\n",
    "        chunk1 = qs1[0:4]\n",
    "        chunk2 = qs1[4:8]\n",
    "        chunk3 = qs1[8:]\n",
    "        factor_scale = jnp.concatenate([chunk1 & 0b111111, (chunk3 & 15) | (sr(chunk1, 6) << 4)], axis=0)\n",
    "        offset_scale = jnp.concatenate([chunk2 & 0b111111, (sr(chunk3, 4) % 16) | (sr(chunk2, 6) << 4)], axis=0)\n",
    "    else:\n",
    "        chunk1 = qs1[:, 0:4]\n",
    "        chunk2 = qs1[:, 4:8]\n",
    "        chunk3 = qs1[:, 8:]\n",
    "        factor_scale = jnp.concatenate([chunk1 & 0b111111, (chunk3 & 15) | (sr(chunk1, 6) << 4)], axis=1)\n",
    "        offset_scale = jnp.concatenate([chunk2 & 0b111111, (sr(chunk3, 4) % 16) | (sr(chunk2, 6) << 4)], axis=1)\n",
    "\n",
    "    basify = lambda x: x  # x.astype(jnp.int8) if not based else x\n",
    "    factors = scale_factors * basify(factor_scale)\n",
    "    offsets = scale_offsets * basify(offset_scale)\n",
    "\n",
    "    # max 15\n",
    "    if based:\n",
    "        # qs2 = jnp.stack([qs2 & 0xf, sr(qs2, 4)], axis=1).reshape(8, 32, num_blocks, -1)\n",
    "        qs2 = jnp.concatenate([qs2 & 0xf, sr(qs2, 4)], axis=1).reshape(8, 32, num_blocks, -1)\n",
    "    else:\n",
    "        qs2 = jnp.stack([qs2 & 0xf, sr(qs2, 4)], axis=2).reshape(-1, 8, 32)\n",
    "\n",
    "    matrix = factors * basify(qs2) - offsets\n",
    "    # matrix = basify(qs2).astype(jnp.float32)\n",
    "    if not based:\n",
    "        matrix = matrix.reshape(num_blocks, -1, 256).transpose(0, 2, 1)\n",
    "        return inputs @ matrix.reshape(inputs.shape[-1], -1)\n",
    "    else:\n",
    "        matrix = matrix.reshape(256, num_blocks, -1)\n",
    "        inputs = inputs.reshape(inputs.shape[0], num_blocks, 256)\n",
    "        result = jax.lax.dot_general(inputs, matrix, (((2, 1), (0, 1)), ((), ())))\n",
    "        return result\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'matmul_4bit' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[7], line 11\u001b[0m\n\u001b[1;32m      8\u001b[0m qs2 \u001b[38;5;241m=\u001b[39m jax\u001b[38;5;241m.\u001b[39mrandom\u001b[38;5;241m.\u001b[39mrandint(key_qs2, (a \u001b[38;5;241m/\u001b[39m\u001b[38;5;241m/\u001b[39m bs, \u001b[38;5;241m128\u001b[39m, b), \u001b[38;5;241m0\u001b[39m, \u001b[38;5;241m255\u001b[39m, dtype\u001b[38;5;241m=\u001b[39mjnp\u001b[38;5;241m.\u001b[39muint8)\u001b[38;5;241m.\u001b[39mview(jnp\u001b[38;5;241m.\u001b[39mint8)\n\u001b[1;32m      9\u001b[0m inputs \u001b[38;5;241m=\u001b[39m jax\u001b[38;5;241m.\u001b[39mrandom\u001b[38;5;241m.\u001b[39mnormal(jax\u001b[38;5;241m.\u001b[39mrandom\u001b[38;5;241m.\u001b[39mPRNGKey(\u001b[38;5;241m2\u001b[39m), (c, a), dtype\u001b[38;5;241m=\u001b[39mjnp\u001b[38;5;241m.\u001b[39mbfloat16) \u001b[38;5;241m/\u001b[39m (a \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39m \u001b[38;5;241m0.5\u001b[39m)\n\u001b[0;32m---> 11\u001b[0m d \u001b[38;5;241m=\u001b[39m \u001b[43mmatmul_4bit\u001b[49m(inputs, scale_factors, scale_offsets, qs1, qs2)\n\u001b[1;32m     12\u001b[0m \u001b[38;5;28mprint\u001b[39m(d)\n\u001b[1;32m     13\u001b[0m e \u001b[38;5;241m=\u001b[39m matmul_4bit(inputs, scale_factors, scale_offsets, qs1, qs2, based\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m)\n",
      "\u001b[0;31mNameError\u001b[0m: name 'matmul_4bit' is not defined"
     ]
    }
   ],
   "source": [
    "a, b, c = 256 * 32, 1024, 1024\n",
    "bs = 256\n",
    "key = jax.random.key(0)\n",
    "key_scale_factors, key_scale_offsets, key_qs1, key_qs2 = jax.random.split(key, 4)\n",
    "scale_factors = jax.random.normal(key_scale_factors, (a // bs, 1, b), dtype=jnp.bfloat16)\n",
    "scale_offsets = jax.random.normal(key_scale_offsets, (a // bs, 1, b), dtype=jnp.bfloat16)\n",
    "qs1 = jax.random.randint(key_qs1, (a // bs, 12, b), 0, 255, dtype=jnp.uint8).view(jnp.int8)\n",
    "qs2 = jax.random.randint(key_qs2, (a // bs, 128, b), 0, 255, dtype=jnp.uint8).view(jnp.int8)\n",
    "inputs = jax.random.normal(jax.random.PRNGKey(2), (c, a), dtype=jnp.bfloat16) / (a ** 0.5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[ 239.15594    743.40454   -211.24771   ...   56.04226    170.22482\n",
      "  -286.69595  ]\n",
      " [  -2.4591599  551.1879    -373.605     ...  -27.097576   149.38379\n",
      "   -57.147476 ]\n",
      " [-220.97462   -212.58714    357.2976    ...  342.2035    -220.43375\n",
      "   257.03244  ]\n",
      " ...\n",
      " [-314.5901      95.02302   -653.9875    ...  233.09988   -152.33174\n",
      "   305.01868  ]\n",
      " [  64.707375   -47.08348   -289.0493    ... -278.11423    165.84293\n",
      "  -159.15613  ]\n",
      " [  -0.9532623  356.84503   -123.46204   ...  -80.8894     394.97888\n",
      "  -458.64697  ]]\n",
      "[[ 0.0000000e+00  1.2207031e-04  0.0000000e+00 ...  1.9073486e-05\n",
      "  -1.0681152e-04 -9.1552734e-05]\n",
      " [ 2.2888184e-05  1.2207031e-04  1.8310547e-04 ...  1.1444092e-04\n",
      "   1.5258789e-05  4.1961670e-05]\n",
      " [-4.5776367e-05  9.1552734e-05 -1.2207031e-04 ... -3.0517578e-05\n",
      "   7.6293945e-05  0.0000000e+00]\n",
      " ...\n",
      " [-1.2207031e-04 -5.3405762e-05 -6.1035156e-05 ...  0.0000000e+00\n",
      "   0.0000000e+00  6.1035156e-05]\n",
      " [ 5.3405762e-05 -3.8146973e-06  9.1552734e-05 ...  1.5258789e-04\n",
      "  -1.5258789e-05 -1.5258789e-05]\n",
      " [ 3.4332275e-05  0.0000000e+00  7.6293945e-06 ... -5.3405762e-05\n",
      "  -6.1035156e-05 -6.1035156e-05]]\n",
      "0.0008544922\n"
     ]
    }
   ],
   "source": [
    "d = matmul_4bit(inputs, scale_factors, scale_offsets, qs1, qs2)\n",
    "print(d)\n",
    "e = matmul_4bit(inputs, scale_factors, scale_offsets, qs1, qs2, based=True)\n",
    "print(d - e)\n",
    "print(jnp.abs(d - e).max())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from micrlhf.quantizers import matmul_4bit_kernel, matmul_fast\n",
    "mesh = jax.sharding.Mesh(np.array(jax.devices(\"tpu\")[:1]).reshape((1, 1)), (\"dp\", \"mp\"))\n",
    "f = matmul_fast(\n",
    "    inputs,#.reshape(-1, 8, 256).swapaxes(1, 2).reshape(inputs.shape),\n",
    "    # inputs,\n",
    "    scale_factors.swapaxes(0, 1).astype(jnp.float16), scale_offsets.swapaxes(0, 1).astype(jnp.float16),\n",
    "    qs1.swapaxes(0, 1), qs2.swapaxes(0, 1),\n",
    "    kernel=matmul_4bit_kernel, mesh=mesh, in_axis=None, out_axis=None, is_transpose=True).block_until_ready()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<script> (()=>{ if (customElements.get('treescope-container') === undefined) { class TreescopeContainer extends HTMLElement { constructor() { super(); this.attachShadow({mode: \"open\"}); this.defns = {}; this.state = {}; } } customElements.define(\"treescope-container\", TreescopeContainer); } if (customElements.get('treescope-run-here') === undefined) { class RunHere extends HTMLElement { constructor() { super() } connectedCallback() { const run = child => { const fn = new Function(child.textContent); child.textContent = \"\"; fn.call(this); this.remove(); }; const child = this.querySelector(\"script\"); if (child) { run(child); } else { new MutationObserver(()=>{ run(this.querySelector(\"script\")); }).observe(this, {childList: true}); } } } customElements.define(\"treescope-run-here\", RunHere); } })(); </script> <treescope-container class=\"treescope_out_74a9149034d44dbb9cffabcfa29e8dfa\" ></treescope-container> <treescope-run-here><script type=\"application/octet-stream\"> const root = ( Array.from(document.getElementsByClassName( \"treescope_out_74a9149034d44dbb9cffabcfa29e8dfa\")) .filter((elt) => !elt.dataset.setup) )[0]; root.dataset.setup = 1; const msg = document.createElement(\"span\"); msg.style = \"color: #aaaaaa; font-family: monospace; transition: opacity 0.2s; opacity: 0.0;\"; msg.textContent = \"(Loading...)\"; root.state.loadingMsg = msg; root.shadowRoot.appendChild(msg); root.state.chain = new Promise((resolve, reject) => { const observer = new IntersectionObserver((entries) => { for (const entry of entries) { if (entry.isIntersecting) { resolve(); observer.disconnect(); return; } } }, {rootMargin: \"1000px\"}); window.setTimeout(() => { if (root.loadingMsg) { root.loadingMsg.style.opacity = \"1.0\"; } observer.observe(root); }, 0); }); root.state.deferring = false; const _insertNode = (node) => { for (let oldScript of node.querySelectorAll(\"script\")) { let newScript = document.createElement(\"script\"); newScript.type = oldScript.type; newScript.textContent = oldScript.textContent; oldScript.parentNode.replaceChild(newScript, oldScript); } if (root.state.loadingMsg) { root.state.loadingMsg.remove(); root.state.loadingMsg = null; } root.shadowRoot.appendChild(node); }; root.defns.insertContent = ((contentNode, compressed) => { if (compressed) { root.state.deferring = true; } if (root.state.deferring) { root.state.chain = (async () => { await root.state.chain; if (compressed) { const encoded = contentNode.textContent; const blob = new Blob([ Uint8Array.from(atob(encoded), (m) => m.codePointAt(0)) ]); const reader = blob.stream().pipeThrough( new DecompressionStream(\"deflate\") ).pipeThrough( new TextDecoderStream(\"utf-8\") ).getReader(); const parts = []; while (true) { const step = await reader.read(); if (step.done) { break; } parts.push(step.value); } const tpl = document.createElement('template'); tpl.innerHTML = parts.join(\"\"); _insertNode(tpl.content); } else { _insertNode(contentNode.content); } })(); } else { _insertNode(contentNode.content); } }); </script></treescope-run-here><div style=\"display:none\"> <script type=\"application/octet-stream\" >eNrdWI1T4kgW/1d6slUrnEPkG0GxLiBfzqCjOKPj7hbXSTpJS9IdOx0Qr/zf73UCCIruzO3c7t1BFSHd77vfe/3rPozk3CdHuhSERBYPyVhwLtE/UcgjKilnDSSIjyWdkgPkcCZzDg6oP2+ggDMehdiC8ZlHJcklLw0UChjxaSRzieicnIcwyjiDYRNbE1fwmNk5i/tcNFLWA7R4M30gAHnUll4DOVQCGZOEyQMUYtumzM35xJENVLQ8pYSRnEeo68FIQa8oMUxiCjav2BZ/clMaUZP6VILlOJZ8RZujTArKImrlIvpA0tmFuY+He2l4DlfhyYmYgU4BY5ElaCiR8q+5g8PQpxZWEdvjliTKe0FwsHOUyWSbRxBQ0BdJZBOHRaiJpEcj3SXyAqJ9ym2Syeoej6SezINrRKJxSJhy2bCUVMX0y2/bZvqY2T6BaRb7/kGqQQczR5wzGM3MuJhk0boN/AqG1NTGsKSWGgyJcLgIMLOIzvgsk03WFxRkXsygXMp0iErFLMihDso8s1r3CXOlh5pNlFckb5ouiIwFg7gj4kfkyTAvZsqy56IjjzpS2ZcQqD+P8H1FQwayitl8pgtyF5NIGowGyXJ1BQ5IJo1JVsk4eKEojCMvDePBFh+XKpqpG294+e02KCvShZTcdf20KsdJ5UC2hkqWGiG+fI/IFBJ8sZLKuuRdn5C5CromNGXQgli3fBxFH6E4F3Iz2krmOIA01JbKH7MQT0j/JMePDve2FYBNpygR2NQ224eGJDbBU3Lf1PIa4gyMAbcZ0L2V9tt9zSiepZcalF3asICDCEHscQg9g3jct4lIMiZpJD/hmvouWlbC0kBUYqhR8E13gByboIqBDa80u0e0STZOrQNqm0agdL5sas8J0RHysUn8RsMkUC9kzSor+Rxs1Zc2tlxBdbZFA8wfPOmiLGl3ps9Vh3xVp8enRLzUbGMxiQh2YZnYS+4G4zLzbMjDUeYokXm0NQ4JT8PyiDUhdjaL/pZ9skGxbmda0m9YmHToBtr5tVgxrZ2/0rxNpleNrP4JRqp1VIpjEakFDDnsU0Rs0UujH6c2KYVEUS6pt+i1HP8xWp/ck+RevtSi02jsUBHJMWdjlf5bSuutUtKLFVVNW5cK/WHz0xV/bqLyKsDCBWSRmpEU9OMf1KZDb52bsZSwpW9rQE/T25JWQ9ozKggkwLXtxL+SQtnWnqG9nSGGrKDYR6N5YHI/QmexVP7aqJ1ywjOcQ2HkZsScAHJLO28AvdyDnRBwFZPATnFE7BXe+4nk1ffgZZqn3Akgy+t1Ejz3Mq2PLV5sb3dPnPoMR2ML9gEI7IofO3Jj91j26bd0PuPZVLkeejTFIpPL2VjiHGawsMmWn10fVkoUjhGYLbM5EYsKESIQMQCpOR7L73NlZQEsDCX2u01LEpXoHQ1CLiRmL2Sbgk8IG6uRp2b0+9FdY1uL53KZH3WFCsAwe2wBprQFYQtTN1E7yNwk/EHdda10FhvpRqla2LcycLQAXFsI7xMopEcSK/6Vvf8xS0wuAMIsLLG5BN+VFevBm3lJaH0fh1BFv49Evr/NvK4hxR+JjoSI3EO62a/S/Ag7tqn4gQczhZXRO0MIPNcdwQPA31YcKOysmk+kT7EPID2TzeoRB3SetCSFstVTT/uiQtjf2Bm1HZRF2dWZJvIIHOTgxEBmqD0ajZQ3IzWmjjHJJBwSEmg7mjMr84+/L7qxpbrcv9mZ17EwUwc5fzE2W5yhy3lIs0hYDRQLP6PaREPN78244xQPTGhE1fJ7O1/vDV2jZSSfwblh8ORf62IGv/2uYXSMtz6twDDcCf9gDzqt9uyrYVx+bZ8Yw0GrbXTd+0H/oyej1pASt9Q9vi5+HFS/TkdhTD8NK5eFk+vBxZfh9Gr4ID/Nu9327pU7uaSt47xHj8/jk47du833zT1nOrDDuw9V7+6K0vN4yHpe3/ksjc/V1qkoG90Bm3Sq1uc4ZrsXlTsrmsymTtffu7t3O3zfNU9mvf1C39hjxkXloxAnhYtd9yF/YeeNE6fgntbas95t0c3zeXxRqwWdQnXWv66fuW5ILifzMhmYDxXLFGc9iQ33fHA6O8bRPDqPB4Prq053Znw6Dwdf7c97e7tu7bJ2XZJ558OnO2NaAZkfjdOaMZwZgftwMdqNb0akc31fdKrWw2n5oj+vxC3jw0PrNuyGJdo/b3fyN/Gn8qjGnNbHTr87DAy6uz/tFD1W8Gq75pfZ9e2sL6bHvc9tdut0Oq7cPbNufL9WqbdPZq19r14eDnujUu/GcINB5bZ1XpeXPdKvd1qtQa907JYv9r5ac9PowZp++bBnnPewQYZt3+g/dM7cG+lWW5/cs7PBcWtCzyuk27put7oWzYee4CGD3AhvOseFh8Jk5LQd6c0/sL6Nu1HfyZ8Gvc5ptWUbd1++hFhGo5vAtjGtF52Hevkzvb2rhoGonvGv7REVvWB60iuNrkalbqdotc6dy92+z8NeuRvNKti9q+7TGzI69cMr1uoPiD0UJL6667WDwlVXTEaj+0qxenUVzQywKIuSSwyZ2UnSeke1+X/Az6r6sc1DaL5PJZlcvei6/gbF+7RmfwNZbx9xveQuINkf060bZEN6MAtl0h1086YGSvCSq/IFssUOq8YiaA9KhIIMap/FM0wlYnhKXSy5gGM+DU2Oha3PBJXkEtB15kkWOLuQ9XQdgG07o63hCXURAFouaUAAeGSWN0Uv+AQJABy8YH18j4r5fD7ZQqH5wpaaSZDxdr1roEF7Mk6dCZYdTN2daOgn1MXUh8YmOVLE75LOBns4i7EP3ZhCzAi2FebZXY/d4lLjd64zFH5a3mdsHq6fY3zt6DDdzA4pC+PFTqMl25fJ77WtQhY7HUymuxwYkTBv6vVAkgAVk7EEYEKklriRm5A51IjX1LSjDHgBDMC+MGGd/QW60t6aX05Su6mtrlNMXKqY+3atUDKtcr1QwSWnVrDhWS/nrUrd1P66QG3zeNs1kHb0sy8PbvG9nmzuyPE5lqUi5DAUMfrZlQerGK4/3m8I3kRC2hFaEr8gWsKUzcjsPIHiHcRZW9VGc+c7m0Ny95XdQSvE3tQyPg5MGyc3eo3k95f8b1kNJft6U1tD9wANthOrQ9rz4wgQJ6Wi0gz+L3w52hqobZH4hsyqlIplsk/282bNLDtlgvetfZKvwlu+Xqnul/60zPrjGWUmKVWo/t+mVOF7Uqrw16UUtCizWLadklMtF0p5s27nyzap2tCpaqZT+R9KqW9sUm/lyvv/3nwqfk8+FX9cPqWP7OrVptOjfwFqdh7u</script> <treescope-run-here><script type=\"application/octet-stream\"> const root = ( Array.from(document.getElementsByClassName( \"treescope_out_74a9149034d44dbb9cffabcfa29e8dfa\")) .filter((elt) => !elt.dataset['step0']) )[0]; root.dataset['step0'] = 1; root.defns.insertContent( this.parentNode.querySelector('script[type=\"application/octet-stream\"]'), true ); this.parentNode.remove(); </script></treescope-run-here> </div>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div style=\"display:none\"> <script type=\"application/octet-stream\" >eNrtWk9z2jgU/yoaZ7aBNjgYbLAhMLOb7s700kP3sIdOxyNLAnsjLFcWadJOvvs+yTZgIISk6W67dQ7Yst57ev/1k+MLmlyjXN1yNrFokmcc345QKlJmoYROrJmQIWUzJiWjYYT7XuTTodOPiBs4Hu7Phg6Fa+B2iRdE1vQiz3AKv1re1I6EUmLBk5SNUqFa9kxwiiPOwlRQNopx3ppyHDE+rc+ESsznvOAZkZiRK0bbbfSyjb6gDFOapPOOEtkIdW3HY4sx0it0YpbMYzVCjm2efe4kKWU3MB6jTOSJSkQ6QpJxrJJrNkZ3NhEcjMNRJNk1SDbDEToZ+D3c7a4JiFgsWKrCTzFLQ3YD9lFGn82ealVs/tarpstFxOQGQTfwHW+wJsiVBD9sSug7nuOtCCBc2i1rgohjcjVGEfzOpVimtFNOXGPZ6nQiLshV8agNVEJSJgsvO9kNygVP6CHKItQPEd8hO77NmATNrkLJFkKxMBbXYOi20rtqiitO4pbv/qLD7nrI6QaFxE84D3MiBefAroQ9S2SuQpWQq6OlBq7taLndADn9gT3cK9jE7+By7Z0FlcRplXvbiyMnRwznDBK1I5bqDM1EqjqfTBpvzemoxiQOoboICzmbqY1S0ENdCx6J63TSSFoTyqJAVpQfl5ineMHCTLJZcgOURoM8+cw0la/L6M7OFdYJG5I44VSy1GRUqnACJukbqA3jrW26lSuKDEHTLZJnK6FV1zK5NkYLLOfgtsItBHPS6pEYdcAil8TtY3T9ZqqVxVKoVpSJ0WpvsdRanbEDIsdj9BK1HDBnk6WUmwMVa7fbKx+sWTvH81brViVtOFciq6ed8vG/6M11UzIe7EhMk2VeGuiwBVh3v2FG069T4sEU/1Lrm0V8dUM8om8eJr6zQeI32UfXQf6ardTsjaa1wa6UbSa+RhLG9UoylhORQVsSQtkr4nABmqIDEpJU61Q0ZE0E1jxIeMxi98opVb44L1CMgTSIcJznE6uuZol3KvhUixrsyDTyXQ+P9+UjOHxs1UXXQoeSPCw2FpGavQSIixBfJGm2VEjBNjqxTFAjcWPtFVLGH1jPDVN9vfVGrKCymbIQxQp3rththlU8sbbUq21sZTai+p5kHeOrcnJr77GmGUs/42QdOaitDDTJASHMMOd687QhJAZeFpe3Qr2rpGfa6FZ9/uK89Njmuju7Wb2ErQPUO6QCNtQkxRyQTCYnj3KuXePd5+oS4FnTFyc3veH4VynxbatvB8HQCXruGaImA2ZcYNXvtQuiuv1nu9EAyRxnOaPWFFXEO0QVxq1rdarjES2hW6SnSKSXHADP5FTFSW6Dfe+gzN5CyrXadixyZcORIc3tGORALmrOkGiGlqZvnxbO0I8nVovjRUQx0oU6Mr9137StdXmt2Ebo9EG+0zEqOXAKtY6LrnV6CexIBwHuV+eV8wOXJ+QDHJyemA0J3ZsLxWnAmgauN/B6cPRyfMiC/1u4E/qUYAPXfxhqU4VPDLbmPbQJPH/PP3CitQ6mxsne/Nlp4n/jm6pJm4Z1qBMfWl3jDThZrNv0wXlr+nwB2HwpAL2Xq7GZQEWD1fYZw4rhi7kaH0ih76y8tLVPKTDD91wldmz5bfnurPHvo/x7L3Jbv5JDWy8KrGn7cdGpAegHy2hVOgdAzG45VReaXE8vHvGm0uv3XOYzvxsNI3fmMuwTn3UHMOoG3sDvr99UNuC+AfffB7jveV5VEZEpCWfQ4PqfB9f3/eGg6ziu7ze4vsH1Da5vcH2D6xv//kC4fi9+eUZIT9wg6rl01p8NXKffjQLadSkbUOIFw2jmNZC+gfTfH6R3bX/oes37+p8a1/vDrtcLBg2ub3B9g+sbXN/g+sa/PxSuvw/EPADuV9iwI5dpJ2ZSf0xDYPGqNeIMYCkxvjsXRDHVARDF8MKaEpHmQGX6wRuaowl6bx3/GfIZso7/T0CN+KEzxodxoRkVRCcF6LWbo2NUag/ml0Qluf1xyeTtn4wzooRsWVvfKVkr1pnEc93yC9alvrUJeEax1+Xwj5Jid7lLDucnYKxk2OBlltJL3XlbFZFNNJXRV8ml+TwODluoxZlCCTB3x3C5WAfA5iydqxievnqlPx/TdMXk5kow/J0zffvb7RtYrOJ+n3yABTRPLpaSsNdw6rvXKSf61GehV2iLvRjagBU51MBfiYpbK2nFR7jl2sZSYFjP6i+Nr01k7lZuqgmq+a6tk9qkKSTynhz+B97Nv08=</script> <treescope-run-here><script type=\"application/octet-stream\"> const root = ( Array.from(document.getElementsByClassName( \"treescope_out_74a9149034d44dbb9cffabcfa29e8dfa\")) .filter((elt) => !elt.dataset['step1']) )[0]; root.dataset['step1'] = 1; root.defns.insertContent( this.parentNode.querySelector('script[type=\"application/octet-stream\"]'), true ); this.parentNode.remove(); </script></treescope-run-here> </div>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<treescope-run-here><script type=\"application/octet-stream\"> const root = ( Array.from(document.getElementsByClassName( \"treescope_out_74a9149034d44dbb9cffabcfa29e8dfa\")) .filter((elt) => !elt.dataset.stolen) )[0]; root.dataset.stolen = 1; this.parentNode.replaceChild(root, this); </script></treescope-run-here>"
      ],
      "text/plain": [
       "(<jax.Array(3.9971924, dtype=float32)>,\n",
       " <jax.Array(255, dtype=bfloat16)>,\n",
       " <jax.Array(254.87454, dtype=float32)>)"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jnp.abs(f - e).max(), jnp.abs(f).mean(), jnp.abs(e).mean()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "jax:    0.4.31\n",
      "jaxlib: 0.4.31\n",
      "numpy:  1.26.4\n",
      "python: 3.10.12 (main, Jul 29 2024, 16:56:48) [GCC 11.4.0]\n",
      "jax.devices (4 total, 4 local): [TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0) TpuDevice(id=1, process_index=0, coords=(1,0,0), core_on_chip=0) TpuDevice(id=2, process_index=0, coords=(0,1,0), core_on_chip=0) TpuDevice(id=3, process_index=0, coords=(1,1,0), core_on_chip=0)]\n",
      "process_count: 1\n",
      "platform: uname_result(system='Linux', node='t1v-n-43c1cc56-w-0', release='5.19.0-1022-gcp', version='#24~22.04.1-Ubuntu SMP Sun Apr 23 09:51:08 UTC 2023', machine='x86_64')\n",
      "\n"
     ]
    }
   ],
   "source": [
    " import jax; jax.print_environment_info()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "micrlhf-progress-a058ydGG-py3.12",
   "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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}