prepare_load_from_hf_pytorch.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "import jax.numpy as jnp\n",
    "import flax.nnx as nn\n",
    "import torch\n",
    "from IPython.core.interactiveshell import InteractiveShell\n",
    "InteractiveShell.ast_node_interactivity = \"all\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import GPT2LMHeadModel\n",
    "\n",
    "model_type = 'gpt2'\n",
    "model_hf = GPT2LMHeadModel.from_pretrained(model_type, cache_dir=\"/var/local/ML/TRAIN/STAGE\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !cd /var/local/ML/TRAIN/pico_shakespeare && cat $(find ../STAGE/*gpt2* -name 'config.json')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "GPT2Config {\n",
       "  \"_name_or_path\": \"gpt2\",\n",
       "  \"activation_function\": \"gelu_new\",\n",
       "  \"architectures\": [\n",
       "    \"GPT2LMHeadModel\"\n",
       "  ],\n",
       "  \"attn_pdrop\": 0.1,\n",
       "  \"bos_token_id\": 50256,\n",
       "  \"embd_pdrop\": 0.1,\n",
       "  \"eos_token_id\": 50256,\n",
       "  \"initializer_range\": 0.02,\n",
       "  \"layer_norm_epsilon\": 1e-05,\n",
       "  \"model_type\": \"gpt2\",\n",
       "  \"n_ctx\": 1024,\n",
       "  \"n_embd\": 768,\n",
       "  \"n_head\": 12,\n",
       "  \"n_inner\": null,\n",
       "  \"n_layer\": 12,\n",
       "  \"n_positions\": 1024,\n",
       "  \"reorder_and_upcast_attn\": false,\n",
       "  \"resid_pdrop\": 0.1,\n",
       "  \"scale_attn_by_inverse_layer_idx\": false,\n",
       "  \"scale_attn_weights\": true,\n",
       "  \"summary_activation\": null,\n",
       "  \"summary_first_dropout\": 0.1,\n",
       "  \"summary_proj_to_labels\": true,\n",
       "  \"summary_type\": \"cls_index\",\n",
       "  \"summary_use_proj\": true,\n",
       "  \"task_specific_params\": {\n",
       "    \"text-generation\": {\n",
       "      \"do_sample\": true,\n",
       "      \"max_length\": 50\n",
       "    }\n",
       "  },\n",
       "  \"transformers_version\": \"4.43.4\",\n",
       "  \"use_cache\": true,\n",
       "  \"vocab_size\": 50257\n",
       "}"
      ]
     },
     "execution_count": 42,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model_hf.config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "odict_keys(['transformer.wte.weight', 'transformer.wpe.weight', 'transformer.h.0.ln_1.weight', 'transformer.h.0.ln_1.bias', 'transformer.h.0.attn.c_attn.weight', 'transformer.h.0.attn.c_attn.bias', 'transformer.h.0.attn.c_proj.weight', 'transformer.h.0.attn.c_proj.bias', 'transformer.h.0.ln_2.weight', 'transformer.h.0.ln_2.bias', 'transformer.h.0.mlp.c_fc.weight', 'transformer.h.0.mlp.c_fc.bias', 'transformer.h.0.mlp.c_proj.weight', 'transformer.h.0.mlp.c_proj.bias', 'transformer.h.1.ln_1.weight', 'transformer.h.1.ln_1.bias', 'transformer.h.1.attn.c_attn.weight', 'transformer.h.1.attn.c_attn.bias', 'transformer.h.1.attn.c_proj.weight', 'transformer.h.1.attn.c_proj.bias', 'transformer.h.1.ln_2.weight', 'transformer.h.1.ln_2.bias', 'transformer.h.1.mlp.c_fc.weight', 'transformer.h.1.mlp.c_fc.bias', 'transformer.h.1.mlp.c_proj.weight', 'transformer.h.1.mlp.c_proj.bias', 'transformer.h.2.ln_1.weight', 'transformer.h.2.ln_1.bias', 'transformer.h.2.attn.c_attn.weight', 'transformer.h.2.attn.c_attn.bias', 'transformer.h.2.attn.c_proj.weight', 'transformer.h.2.attn.c_proj.bias', 'transformer.h.2.ln_2.weight', 'transformer.h.2.ln_2.bias', 'transformer.h.2.mlp.c_fc.weight', 'transformer.h.2.mlp.c_fc.bias', 'transformer.h.2.mlp.c_proj.weight', 'transformer.h.2.mlp.c_proj.bias', 'transformer.h.3.ln_1.weight', 'transformer.h.3.ln_1.bias', 'transformer.h.3.attn.c_attn.weight', 'transformer.h.3.attn.c_attn.bias', 'transformer.h.3.attn.c_proj.weight', 'transformer.h.3.attn.c_proj.bias', 'transformer.h.3.ln_2.weight', 'transformer.h.3.ln_2.bias', 'transformer.h.3.mlp.c_fc.weight', 'transformer.h.3.mlp.c_fc.bias', 'transformer.h.3.mlp.c_proj.weight', 'transformer.h.3.mlp.c_proj.bias', 'transformer.h.4.ln_1.weight', 'transformer.h.4.ln_1.bias', 'transformer.h.4.attn.c_attn.weight', 'transformer.h.4.attn.c_attn.bias', 'transformer.h.4.attn.c_proj.weight', 'transformer.h.4.attn.c_proj.bias', 'transformer.h.4.ln_2.weight', 'transformer.h.4.ln_2.bias', 'transformer.h.4.mlp.c_fc.weight', 'transformer.h.4.mlp.c_fc.bias', 'transformer.h.4.mlp.c_proj.weight', 'transformer.h.4.mlp.c_proj.bias', 'transformer.h.5.ln_1.weight', 'transformer.h.5.ln_1.bias', 'transformer.h.5.attn.c_attn.weight', 'transformer.h.5.attn.c_attn.bias', 'transformer.h.5.attn.c_proj.weight', 'transformer.h.5.attn.c_proj.bias', 'transformer.h.5.ln_2.weight', 'transformer.h.5.ln_2.bias', 'transformer.h.5.mlp.c_fc.weight', 'transformer.h.5.mlp.c_fc.bias', 'transformer.h.5.mlp.c_proj.weight', 'transformer.h.5.mlp.c_proj.bias', 'transformer.h.6.ln_1.weight', 'transformer.h.6.ln_1.bias', 'transformer.h.6.attn.c_attn.weight', 'transformer.h.6.attn.c_attn.bias', 'transformer.h.6.attn.c_proj.weight', 'transformer.h.6.attn.c_proj.bias', 'transformer.h.6.ln_2.weight', 'transformer.h.6.ln_2.bias', 'transformer.h.6.mlp.c_fc.weight', 'transformer.h.6.mlp.c_fc.bias', 'transformer.h.6.mlp.c_proj.weight', 'transformer.h.6.mlp.c_proj.bias', 'transformer.h.7.ln_1.weight', 'transformer.h.7.ln_1.bias', 'transformer.h.7.attn.c_attn.weight', 'transformer.h.7.attn.c_attn.bias', 'transformer.h.7.attn.c_proj.weight', 'transformer.h.7.attn.c_proj.bias', 'transformer.h.7.ln_2.weight', 'transformer.h.7.ln_2.bias', 'transformer.h.7.mlp.c_fc.weight', 'transformer.h.7.mlp.c_fc.bias', 'transformer.h.7.mlp.c_proj.weight', 'transformer.h.7.mlp.c_proj.bias', 'transformer.h.8.ln_1.weight', 'transformer.h.8.ln_1.bias', 'transformer.h.8.attn.c_attn.weight', 'transformer.h.8.attn.c_attn.bias', 'transformer.h.8.attn.c_proj.weight', 'transformer.h.8.attn.c_proj.bias', 'transformer.h.8.ln_2.weight', 'transformer.h.8.ln_2.bias', 'transformer.h.8.mlp.c_fc.weight', 'transformer.h.8.mlp.c_fc.bias', 'transformer.h.8.mlp.c_proj.weight', 'transformer.h.8.mlp.c_proj.bias', 'transformer.h.9.ln_1.weight', 'transformer.h.9.ln_1.bias', 'transformer.h.9.attn.c_attn.weight', 'transformer.h.9.attn.c_attn.bias', 'transformer.h.9.attn.c_proj.weight', 'transformer.h.9.attn.c_proj.bias', 'transformer.h.9.ln_2.weight', 'transformer.h.9.ln_2.bias', 'transformer.h.9.mlp.c_fc.weight', 'transformer.h.9.mlp.c_fc.bias', 'transformer.h.9.mlp.c_proj.weight', 'transformer.h.9.mlp.c_proj.bias', 'transformer.h.10.ln_1.weight', 'transformer.h.10.ln_1.bias', 'transformer.h.10.attn.c_attn.weight', 'transformer.h.10.attn.c_attn.bias', 'transformer.h.10.attn.c_proj.weight', 'transformer.h.10.attn.c_proj.bias', 'transformer.h.10.ln_2.weight', 'transformer.h.10.ln_2.bias', 'transformer.h.10.mlp.c_fc.weight', 'transformer.h.10.mlp.c_fc.bias', 'transformer.h.10.mlp.c_proj.weight', 'transformer.h.10.mlp.c_proj.bias', 'transformer.h.11.ln_1.weight', 'transformer.h.11.ln_1.bias', 'transformer.h.11.attn.c_attn.weight', 'transformer.h.11.attn.c_attn.bias', 'transformer.h.11.attn.c_proj.weight', 'transformer.h.11.attn.c_proj.bias', 'transformer.h.11.ln_2.weight', 'transformer.h.11.ln_2.bias', 'transformer.h.11.mlp.c_fc.weight', 'transformer.h.11.mlp.c_fc.bias', 'transformer.h.11.mlp.c_proj.weight', 'transformer.h.11.mlp.c_proj.bias', 'transformer.ln_f.weight', 'transformer.ln_f.bias', 'lm_head.weight'])"
      ]
     },
     "execution_count": 43,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sd_hf = model_hf.state_dict()\n",
    "sd_hf.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Tensor"
      ]
     },
     "execution_count": 44,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "torch.Size([50257, 768])"
      ]
     },
     "execution_count": 44,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "type(sd_hf['transformer.wte.weight'])\n",
    "sd_hf['transformer.wte.weight'].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Embed(\n",
       "  embedding=Param(\n",
       "    value=Array(shape=(50257, 768), dtype=float32)\n",
       "  ),\n",
       "  num_embeddings=50257,\n",
       "  features=768,\n",
       "  dtype=dtype('float32'),\n",
       "  param_dtype=<class 'jax.numpy.float32'>,\n",
       "  embedding_init=<function variance_scaling.<locals>.init at 0x7f88b8b31310>\n",
       ")"
      ]
     },
     "execution_count": 45,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vocab_size, embedding_size = sd_hf['transformer.wte.weight'].shape\n",
    "wte = nn.Embed(num_embeddings=vocab_size, features=embedding_size, rngs=nn.Rngs(0))\n",
    "wte"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "State({\n",
       "  'embedding': VariableState(\n",
       "    type=Param,\n",
       "    value=Array([[ 1.1953074e-02, -4.3702841e-02,  1.5583087e-02, ...,\n",
       "             2.9827623e-02,  2.3185154e-02, -5.0074987e-02],\n",
       "           [-7.3833148e-06, -2.3771707e-02,  3.2675751e-02, ...,\n",
       "             1.2370082e-02, -1.8245960e-02, -5.5854514e-02],\n",
       "           [-6.4068988e-02,  1.0926131e-02, -9.7181993e-03, ...,\n",
       "            -1.6451136e-03, -1.8916763e-02, -7.8727528e-02],\n",
       "           ...,\n",
       "           [-4.5188973e-03, -8.1994740e-04,  1.7434264e-02, ...,\n",
       "             1.5338360e-02,  2.8312072e-02,  2.1429532e-04],\n",
       "           [-1.0427453e-03,  1.4039346e-02,  4.0459871e-02, ...,\n",
       "            -3.7717942e-02, -1.7851518e-02, -4.7507521e-02],\n",
       "           [ 2.5526977e-03, -1.7003939e-02,  2.0834690e-02, ...,\n",
       "            -3.3392377e-02, -8.9475606e-04, -4.4884817e-03]], dtype=float32)\n",
       "  )\n",
       "})"
      ]
     },
     "execution_count": 46,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "np.float32(-0.11010301)"
      ]
     },
     "execution_count": 46,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "Array(0.01195307, dtype=float32)"
      ]
     },
     "execution_count": 46,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "jaxlib.xla_extension.ArrayImpl"
      ]
     },
     "execution_count": 46,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "jaxlib.xla_extension.ArrayImpl"
      ]
     },
     "execution_count": 46,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "graphdef, state = nn.split(wte)\n",
    "state\n",
    "sd_hf['transformer.wte.weight'].cpu().numpy()[0][0] \n",
    "state[\"embedding\"].value[0][0]\n",
    "type(state[\"embedding\"].value)\n",
    "state[\"embedding\"].value = jnp.array(sd_hf['transformer.wte.weight'].cpu().numpy())\n",
    "type(state[\"embedding\"].value)\n",
    "nn.update(wte, state)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[('transformer.h.0.ln_1.weight', torch.Size([768])),\n",
       " ('transformer.h.0.ln_1.bias', torch.Size([768])),\n",
       " ('transformer.h.0.attn.c_attn.weight', torch.Size([768, 2304])),\n",
       " ('transformer.h.0.attn.c_attn.bias', torch.Size([2304])),\n",
       " ('transformer.h.0.attn.c_proj.weight', torch.Size([768, 768])),\n",
       " ('transformer.h.0.attn.c_proj.bias', torch.Size([768])),\n",
       " ('transformer.h.0.ln_2.weight', torch.Size([768])),\n",
       " ('transformer.h.0.ln_2.bias', torch.Size([768])),\n",
       " ('transformer.h.0.mlp.c_fc.weight', torch.Size([768, 3072])),\n",
       " ('transformer.h.0.mlp.c_fc.bias', torch.Size([3072])),\n",
       " ('transformer.h.0.mlp.c_proj.weight', torch.Size([3072, 768])),\n",
       " ('transformer.h.0.mlp.c_proj.bias', torch.Size([768]))]"
      ]
     },
     "execution_count": 47,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "[(x,model_hf.state_dict()[x].shape) for x in model_hf.state_dict().keys() if \"h.0\" in x]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "LayerNorm(\n",
       "  scale=Param(\n",
       "    value=Array(shape=(768,), dtype=float32)\n",
       "  ),\n",
       "  bias=Param(\n",
       "    value=Array(shape=(768,), dtype=float32)\n",
       "  ),\n",
       "  num_features=768,\n",
       "  epsilon=1e-06,\n",
       "  dtype=None,\n",
       "  param_dtype=<class 'jax.numpy.float32'>,\n",
       "  use_bias=True,\n",
       "  use_scale=True,\n",
       "  bias_init=<function zeros at 0x7f88b93d3160>,\n",
       "  scale_init=<function ones at 0x7f88b93d3310>,\n",
       "  reduction_axes=-1,\n",
       "  feature_axes=-1,\n",
       "  axis_name=None,\n",
       "  axis_index_groups=None,\n",
       "  use_fast_variance=True\n",
       ")"
      ]
     },
     "execution_count": 48,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ln_features, = sd_hf['transformer.h.0.ln_1.weight'].shape\n",
    "ln = nn.LayerNorm(num_features=ln_features, rngs=nn.Rngs(0))\n",
    "ln"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "LayerNorm(\n",
       "  scale=Param(\n",
       "    value=Array(shape=(768,), dtype=float32)\n",
       "  ),\n",
       "  bias=Param(\n",
       "    value=Array(shape=(768,), dtype=float32)\n",
       "  ),\n",
       "  num_features=768,\n",
       "  epsilon=1e-06,\n",
       "  dtype=None,\n",
       "  param_dtype=<class 'jax.numpy.float32'>,\n",
       "  use_bias=True,\n",
       "  use_scale=True,\n",
       "  bias_init=<function zeros at 0x7f88b93d3160>,\n",
       "  scale_init=<function ones at 0x7f88b93d3310>,\n",
       "  reduction_axes=-1,\n",
       "  feature_axes=-1,\n",
       "  axis_name=None,\n",
       "  axis_index_groups=None,\n",
       "  use_fast_variance=True\n",
       ")"
      ]
     },
     "execution_count": 49,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ln.scale.value = jnp.array(sd_hf['transformer.h.0.ln_1.weight'].cpu().numpy())\n",
    "ln.bias.value = jnp.array(sd_hf['transformer.h.0.ln_1.bias'].cpu().numpy())\n",
    "ln"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "MultiHeadAttention(\n",
       "  num_heads=12,\n",
       "  in_features=768,\n",
       "  qkv_features=768,\n",
       "  out_features=768,\n",
       "  dtype=None,\n",
       "  param_dtype=<class 'jax.numpy.float32'>,\n",
       "  broadcast_dropout=True,\n",
       "  dropout_rate=0.0,\n",
       "  deterministic=None,\n",
       "  precision=None,\n",
       "  kernel_init=<function variance_scaling.<locals>.init at 0x7f88b8b31040>,\n",
       "  out_kernel_init=None,\n",
       "  bias_init=<function zeros at 0x7f88b93d3160>,\n",
       "  out_bias_init=None,\n",
       "  use_bias=True,\n",
       "  attention_fn=<function dot_product_attention at 0x7f88b8b31e50>,\n",
       "  decode=False,\n",
       "  normalize_qk=False,\n",
       "  qkv_dot_general=None,\n",
       "  out_dot_general=None,\n",
       "  qkv_dot_general_cls=None,\n",
       "  out_dot_general_cls=None,\n",
       "  head_dim=64,\n",
       "  query=LinearGeneral(\n",
       "    in_features=(768,),\n",
       "    out_features=(12, 64),\n",
       "    axis=(-1,),\n",
       "    batch_axis=FrozenDict({}),\n",
       "    use_bias=True,\n",
       "    dtype=None,\n",
       "    param_dtype=<class 'jax.numpy.float32'>,\n",
       "    kernel_init=<function variance_scaling.<locals>.init at 0x7f88b8b31040>,\n",
       "    bias_init=<function zeros at 0x7f88b93d3160>,\n",
       "    precision=None,\n",
       "    dot_general=None,\n",
       "    dot_general_cls=None,\n",
       "    kernel=Param(\n",
       "      value=Array(shape=(768, 12, 64), dtype=float32)\n",
       "    ),\n",
       "    bias=Param(\n",
       "      value=Array(shape=(12, 64), dtype=float32)\n",
       "    )\n",
       "  ),\n",
       "  key=LinearGeneral(\n",
       "    in_features=(768,),\n",
       "    out_features=(12, 64),\n",
       "    axis=(-1,),\n",
       "    batch_axis=FrozenDict({}),\n",
       "    use_bias=True,\n",
       "    dtype=None,\n",
       "    param_dtype=<class 'jax.numpy.float32'>,\n",
       "    kernel_init=<function variance_scaling.<locals>.init at 0x7f88b8b31040>,\n",
       "    bias_init=<function zeros at 0x7f88b93d3160>,\n",
       "    precision=None,\n",
       "    dot_general=None,\n",
       "    dot_general_cls=None,\n",
       "    kernel=Param(\n",
       "      value=Array(shape=(768, 12, 64), dtype=float32)\n",
       "    ),\n",
       "    bias=Param(\n",
       "      value=Array(shape=(12, 64), dtype=float32)\n",
       "    )\n",
       "  ),\n",
       "  value=LinearGeneral(\n",
       "    in_features=(768,),\n",
       "    out_features=(12, 64),\n",
       "    axis=(-1,),\n",
       "    batch_axis=FrozenDict({}),\n",
       "    use_bias=True,\n",
       "    dtype=None,\n",
       "    param_dtype=<class 'jax.numpy.float32'>,\n",
       "    kernel_init=<function variance_scaling.<locals>.init at 0x7f88b8b31040>,\n",
       "    bias_init=<function zeros at 0x7f88b93d3160>,\n",
       "    precision=None,\n",
       "    dot_general=None,\n",
       "    dot_general_cls=None,\n",
       "    kernel=Param(\n",
       "      value=Array(shape=(768, 12, 64), dtype=float32)\n",
       "    ),\n",
       "    bias=Param(\n",
       "      value=Array(shape=(12, 64), dtype=float32)\n",
       "    )\n",
       "  ),\n",
       "  query_ln=None,\n",
       "  key_ln=None,\n",
       "  out=LinearGeneral(\n",
       "    in_features=(12, 64),\n",
       "    out_features=(768,),\n",
       "    axis=(-2, -1),\n",
       "    batch_axis=FrozenDict({}),\n",
       "    use_bias=True,\n",
       "    dtype=None,\n",
       "    param_dtype=<class 'jax.numpy.float32'>,\n",
       "    kernel_init=<function variance_scaling.<locals>.init at 0x7f88b8b31040>,\n",
       "    bias_init=<function zeros at 0x7f88b93d3160>,\n",
       "    precision=None,\n",
       "    dot_general=None,\n",
       "    dot_general_cls=None,\n",
       "    kernel=Param(\n",
       "      value=Array(shape=(12, 64, 768), dtype=float32)\n",
       "    ),\n",
       "    bias=Param(\n",
       "      value=Array(shape=(768,), dtype=float32)\n",
       "    )\n",
       "  ),\n",
       "  rngs=None,\n",
       "  cached_key=None,\n",
       "  cached_value=None,\n",
       "  cache_index=None\n",
       ")"
      ]
     },
     "execution_count": 50,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mha = nn.MultiHeadAttention(\n",
    "num_heads=12, \n",
    "in_features=768, \n",
    "qkv_features=768,\n",
    "out_features=768,\n",
    "decode=False, \n",
    "rngs=nn.Rngs(0)\n",
    ")\n",
    "mha"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Prep"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[(('h', 0, 'attn', 'c_attn'), 'Linear'),\n",
       " (('h', 0, 'attn', 'c_proj'), 'Linear'),\n",
       " (('h', 0, 'attn'), 'CausalSelfAttention'),\n",
       " (('h', 0, 'ln_1'), 'LayerNorm'),\n",
       " (('h', 0, 'ln_2'), 'LayerNorm'),\n",
       " (('h', 0, 'mlp', 'c_fc'), 'Linear'),\n",
       " (('h', 0, 'mlp', 'c_proj'), 'Linear'),\n",
       " (('h', 0, 'mlp'), 'MLP'),\n",
       " (('h', 0), 'Block'),\n",
       " (('h', 1, 'attn', 'c_attn'), 'Linear'),\n",
       " (('h', 1, 'attn', 'c_proj'), 'Linear'),\n",
       " (('h', 1, 'attn'), 'CausalSelfAttention'),\n",
       " (('h', 1, 'ln_1'), 'LayerNorm'),\n",
       " (('h', 1, 'ln_2'), 'LayerNorm'),\n",
       " (('h', 1, 'mlp', 'c_fc'), 'Linear'),\n",
       " (('h', 1, 'mlp', 'c_proj'), 'Linear'),\n",
       " (('h', 1, 'mlp'), 'MLP'),\n",
       " (('h', 1), 'Block'),\n",
       " (('h', 2, 'attn', 'c_attn'), 'Linear'),\n",
       " (('h', 2, 'attn', 'c_proj'), 'Linear'),\n",
       " (('h', 2, 'attn'), 'CausalSelfAttention'),\n",
       " (('h', 2, 'ln_1'), 'LayerNorm'),\n",
       " (('h', 2, 'ln_2'), 'LayerNorm'),\n",
       " (('h', 2, 'mlp', 'c_fc'), 'Linear'),\n",
       " (('h', 2, 'mlp', 'c_proj'), 'Linear'),\n",
       " (('h', 2, 'mlp'), 'MLP'),\n",
       " (('h', 2), 'Block'),\n",
       " (('h', 3, 'attn', 'c_attn'), 'Linear'),\n",
       " (('h', 3, 'attn', 'c_proj'), 'Linear'),\n",
       " (('h', 3, 'attn'), 'CausalSelfAttention'),\n",
       " (('h', 3, 'ln_1'), 'LayerNorm'),\n",
       " (('h', 3, 'ln_2'), 'LayerNorm'),\n",
       " (('h', 3, 'mlp', 'c_fc'), 'Linear'),\n",
       " (('h', 3, 'mlp', 'c_proj'), 'Linear'),\n",
       " (('h', 3, 'mlp'), 'MLP'),\n",
       " (('h', 3), 'Block'),\n",
       " (('h', 4, 'attn', 'c_attn'), 'Linear'),\n",
       " (('h', 4, 'attn', 'c_proj'), 'Linear'),\n",
       " (('h', 4, 'attn'), 'CausalSelfAttention'),\n",
       " (('h', 4, 'ln_1'), 'LayerNorm'),\n",
       " (('h', 4, 'ln_2'), 'LayerNorm'),\n",
       " (('h', 4, 'mlp', 'c_fc'), 'Linear'),\n",
       " (('h', 4, 'mlp', 'c_proj'), 'Linear'),\n",
       " (('h', 4, 'mlp'), 'MLP'),\n",
       " (('h', 4), 'Block'),\n",
       " (('h', 5, 'attn', 'c_attn'), 'Linear'),\n",
       " (('h', 5, 'attn', 'c_proj'), 'Linear'),\n",
       " (('h', 5, 'attn'), 'CausalSelfAttention'),\n",
       " (('h', 5, 'ln_1'), 'LayerNorm'),\n",
       " (('h', 5, 'ln_2'), 'LayerNorm'),\n",
       " (('h', 5, 'mlp', 'c_fc'), 'Linear'),\n",
       " (('h', 5, 'mlp', 'c_proj'), 'Linear'),\n",
       " (('h', 5, 'mlp'), 'MLP'),\n",
       " (('h', 5), 'Block'),\n",
       " (('h', 6, 'attn', 'c_attn'), 'Linear'),\n",
       " (('h', 6, 'attn', 'c_proj'), 'Linear'),\n",
       " (('h', 6, 'attn'), 'CausalSelfAttention'),\n",
       " (('h', 6, 'ln_1'), 'LayerNorm'),\n",
       " (('h', 6, 'ln_2'), 'LayerNorm'),\n",
       " (('h', 6, 'mlp', 'c_fc'), 'Linear'),\n",
       " (('h', 6, 'mlp', 'c_proj'), 'Linear'),\n",
       " (('h', 6, 'mlp'), 'MLP'),\n",
       " (('h', 6), 'Block'),\n",
       " (('h', 7, 'attn', 'c_attn'), 'Linear'),\n",
       " (('h', 7, 'attn', 'c_proj'), 'Linear'),\n",
       " (('h', 7, 'attn'), 'CausalSelfAttention'),\n",
       " (('h', 7, 'ln_1'), 'LayerNorm'),\n",
       " (('h', 7, 'ln_2'), 'LayerNorm'),\n",
       " (('h', 7, 'mlp', 'c_fc'), 'Linear'),\n",
       " (('h', 7, 'mlp', 'c_proj'), 'Linear'),\n",
       " (('h', 7, 'mlp'), 'MLP'),\n",
       " (('h', 7), 'Block'),\n",
       " (('h', 8, 'attn', 'c_attn'), 'Linear'),\n",
       " (('h', 8, 'attn', 'c_proj'), 'Linear'),\n",
       " (('h', 8, 'attn'), 'CausalSelfAttention'),\n",
       " (('h', 8, 'ln_1'), 'LayerNorm'),\n",
       " (('h', 8, 'ln_2'), 'LayerNorm'),\n",
       " (('h', 8, 'mlp', 'c_fc'), 'Linear'),\n",
       " (('h', 8, 'mlp', 'c_proj'), 'Linear'),\n",
       " (('h', 8, 'mlp'), 'MLP'),\n",
       " (('h', 8), 'Block'),\n",
       " (('h', 9, 'attn', 'c_attn'), 'Linear'),\n",
       " (('h', 9, 'attn', 'c_proj'), 'Linear'),\n",
       " (('h', 9, 'attn'), 'CausalSelfAttention'),\n",
       " (('h', 9, 'ln_1'), 'LayerNorm'),\n",
       " (('h', 9, 'ln_2'), 'LayerNorm'),\n",
       " (('h', 9, 'mlp', 'c_fc'), 'Linear'),\n",
       " (('h', 9, 'mlp', 'c_proj'), 'Linear'),\n",
       " (('h', 9, 'mlp'), 'MLP'),\n",
       " (('h', 9), 'Block'),\n",
       " (('h', 10, 'attn', 'c_attn'), 'Linear'),\n",
       " (('h', 10, 'attn', 'c_proj'), 'Linear'),\n",
       " (('h', 10, 'attn'), 'CausalSelfAttention'),\n",
       " (('h', 10, 'ln_1'), 'LayerNorm'),\n",
       " (('h', 10, 'ln_2'), 'LayerNorm'),\n",
       " (('h', 10, 'mlp', 'c_fc'), 'Linear'),\n",
       " (('h', 10, 'mlp', 'c_proj'), 'Linear'),\n",
       " (('h', 10, 'mlp'), 'MLP'),\n",
       " (('h', 10), 'Block'),\n",
       " (('h', 11, 'attn', 'c_attn'), 'Linear'),\n",
       " (('h', 11, 'attn', 'c_proj'), 'Linear'),\n",
       " (('h', 11, 'attn'), 'CausalSelfAttention'),\n",
       " (('h', 11, 'ln_1'), 'LayerNorm'),\n",
       " (('h', 11, 'ln_2'), 'LayerNorm'),\n",
       " (('h', 11, 'mlp', 'c_fc'), 'Linear'),\n",
       " (('h', 11, 'mlp', 'c_proj'), 'Linear'),\n",
       " (('h', 11, 'mlp'), 'MLP'),\n",
       " (('h', 11), 'Block'),\n",
       " (('lm_head',), 'Linear'),\n",
       " (('ln_f',), 'LayerNorm'),\n",
       " (('wpe',), 'Embed'),\n",
       " (('wte',), 'Embed'),\n",
       " ((), 'GPT')]"
      ]
     },
     "execution_count": 51,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from jax_gpt2 import GPT, GPTConfig\n",
    "import flax.nnx as nn\n",
    "config = GPTConfig()\n",
    "model = GPT(config, nn.Rngs(0))\n",
    "[(x[0], type(x[1]).__name__) for x in model.iter_modules()]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'Block', 'CausalSelfAttention', 'Embed', 'GPT', 'LayerNorm', 'Linear', 'MLP'}"
      ]
     },
     "execution_count": 52,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "set([type(x[1]).__name__ for x in model.iter_modules()])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "76"
      ]
     },
     "execution_count": 53,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Transposing  lm_head\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "149"
      ]
     },
     "execution_count": 53,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from transformers import GPT2LMHeadModel\n",
    "import jax\n",
    "\n",
    "model_type = 'gpt2'\n",
    "model_hf = GPT2LMHeadModel.from_pretrained(model_type, cache_dir=\"/var/local/ML/TRAIN/STAGE\")\n",
    "\n",
    "jax_modules_dict = {}\n",
    "for module_pair in model.iter_modules():\n",
    "    if type(module_pair[1]).__name__  in ['Block', 'CausalSelfAttention', 'GPT', 'MLP']:\n",
    "        continue\n",
    "    module_path = '.'.join([str(x) for x in module_pair[0]])\n",
    "    module = module_pair[1]\n",
    "    jax_modules_dict[module_path] = module\n",
    "\n",
    "len(jax_modules_dict.keys())\n",
    "\n",
    "equivalent_jax_modules = []\n",
    "hf_sd = model_hf.state_dict()\n",
    "\n",
    "for param in hf_sd:\n",
    "    if 'transformer' in param:\n",
    "        key = '.'.join(param.split(\".\")[1:-1])\n",
    "    else:\n",
    "        key = '.'.join(param.split(\".\")[:-1])\n",
    "    equivalent_jax_module = jax_modules_dict[key]\n",
    "    equivalent_jax_modules.append(type(equivalent_jax_module).__name__)\n",
    "\n",
    "    if type(equivalent_jax_module).__name__ == 'Embed':\n",
    "        inner = equivalent_jax_module.embedding\n",
    "    elif type(equivalent_jax_module).__name__ == 'Linear':\n",
    "        if 'weight' in param:\n",
    "            inner = equivalent_jax_module.kernel\n",
    "        else:\n",
    "            inner = equivalent_jax_module.bias\n",
    "    elif type(equivalent_jax_module).__name__ == 'LayerNorm':\n",
    "        if 'weight' in param:\n",
    "            inner = equivalent_jax_module.scale\n",
    "        else:\n",
    "            inner = equivalent_jax_module.bias\n",
    "    \n",
    "    if inner.value.shape == tuple(hf_sd[param].shape):\n",
    "        inner.value = jnp.array(hf_sd[param].cpu().numpy())\n",
    "    elif inner.value.shape == tuple(hf_sd[param].shape)[::-1]:\n",
    "        print(\"Transposing \", key)\n",
    "        inner.value = jnp.array(hf_sd[param].cpu().numpy().T)\n",
    "\n",
    "assert len(equivalent_jax_modules) == len(model_hf.state_dict())\n",
    "\n",
    "len(equivalent_jax_modules)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "jax",
   "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.9.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}