convert_neo_pytorch_model_to_jax.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "GPT-Neo pytorch_model.bin to JAX conversion tool.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "hwNbvZdI5pbk",
        "outputId": "15809bbd-45d1-434a-9d9d-66c00fa62c67"
      },
      "source": [
        "#@title Load checkpoint\n",
        "path_to_checkpoint = \"/content/drive/MyDrive/pytorch_model.bin\"  #@param {type:\"string\"}\n",
        "checkpoint_type = \"GPT-Neo-2.7B\"  #@param [\"GPT-Neo-1.3B\", \"GPT-Neo-2.7B\"]\n",
        "import os\n",
        "from termcolor import colored\n",
        "from IPython.display import clear_output\n",
        "\n",
        "from google.colab import drive\n",
        "drive.mount('/content/drive/')\n",
        "import torch\n",
        "print(\"Reading from checkpoint...\")\n",
        "torch_checkpoint = torch.load(path_to_checkpoint, map_location='cpu')\n",
        "print(\"Done.\")\n",
        "\n",
        "if checkpoint_type == \"GPT-Neo-1.3B\":\n",
        "    total_shards = 8\n",
        "    d_model = 2048\n",
        "    layers = 24\n",
        "else:\n",
        "    total_shards = 4\n",
        "    d_model = 2560\n",
        "    layers = 32\n",
        "for i in range(total_shards):\n",
        "    os.makedirs(f\"jax_checkpoint/shard_{i}\")\n",
        "pieces = 16"
      ],
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Mounted at /content/drive/\n",
            "Reading from checkpoint...\n",
            "Done.\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TmEE1tcf6ADW",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "53420d0e-5554-4012-e29b-3d98041e33df"
      },
      "source": [
        "#@title Convert checkpoint to be JAX-compatible { display-mode: \"form\" }\n",
        "from termcolor import colored\n",
        "from IPython.display import clear_output\n",
        "import torch\n",
        "import numpy as np\n",
        "import jax.numpy as jnp\n",
        "\n",
        "def reshard_reverse(x, old_shape, is_shard_bias=False):\n",
        "    if len(x.shape) == 1:\n",
        "        assert False\n",
        "        out = x[0:1]\n",
        "\n",
        "    elif len(x.shape) == 2:\n",
        "        #print(f\"LN/bias\")\n",
        "        if old_shape[1] == x.shape[1]:\n",
        "            #print(\"LN\")\n",
        "            if not is_shard_bias:\n",
        "                out = np.tile(x[0:1], (total_shards, 1))\n",
        "            else:\n",
        "                #print(\"shard bias\")\n",
        "                out = np.tile(x[0:1], (total_shards, 1)) / total_shards\n",
        "        else:\n",
        "            #print(\"bias\")\n",
        "            out = x.reshape(old_shape)\n",
        "\n",
        "    elif len(x.shape) == 3:\n",
        "        if x.shape[0] * x.shape[2] == old_shape[2]:\n",
        "            #print(\"case 1\")\n",
        "            out = x.reshape(old_shape)\n",
        "        elif x.shape[0] * x.shape[1] == old_shape[1]:\n",
        "            #print(\"case 2\")\n",
        "            out = jnp.transpose(x.reshape((old_shape[1], old_shape[0], old_shape[2])), (1, 0, 2))\n",
        "        else:\n",
        "            raise Exception(f\"unimplemented, {x.shape}, {old_shape}\")\n",
        "    else:\n",
        "        raise Exception(f\"unimplemented, {x}\")\n",
        "    #flattened, structure = jax.tree_flatten(out)\n",
        "    #return flattened\n",
        "    return out\n",
        "\n",
        "def get_old_shape(t, dim=2):\n",
        "    if len(t.shape) == 2:\n",
        "        shard_shape = t.shape\n",
        "        if dim == 1:\n",
        "            assert shard_shape[0] % total_shards == 0\n",
        "            return (shard_shape[0] // total_shards, shard_shape[1])\n",
        "        elif dim == 2:\n",
        "            assert shard_shape[1] % total_shards == 0\n",
        "            return (shard_shape[0], shard_shape[1] // total_shards)\n",
        "        else:\n",
        "            raise ValueError(f\"unsupported dim {dim}\")\n",
        "    if len(t.shape) == 1:\n",
        "        assert t.shape[0] % total_shards == 0\n",
        "        return (t.shape[0] // total_shards,)\n",
        "    else:\n",
        "        raise ValueError(f\"unsupported shape {t.shape}\")\n",
        "\n",
        "\n",
        "def split(a, n):\n",
        "    k, m = divmod(len(a), n)\n",
        "    return (a[i * k + min(i, m):(i + 1) * k + min(i + 1, m)] for i in range(n))\n",
        "\n",
        "def save(cpu_flattened):\n",
        "    for i in range(total_shards):\n",
        "        cpu_flattened_chunked = split(cpu_flattened, pieces)\n",
        "        for j, chunk in enumerate(cpu_flattened_chunked):\n",
        "            with open(f\"jax_checkpoint/shard_{i}/{j}.npz\", \"wb\") as f:\n",
        "                np.savez(f, *map(lambda c: c[i], chunk))\n",
        "\n",
        "\n",
        "transforms = [\n",
        "    (\"transformer.wpe.weight\", False, 2),\n",
        "    (\"transformer.wte.weight\", False, 1)\n",
        "]\n",
        "\n",
        "checkpoint = []\n",
        "\n",
        "layer_names = sorted(map(str, range(layers)))\n",
        "for layer in layer_names:\n",
        "    transforms.extend([\n",
        "        (f\"transformer.h.{layer}.attn.attention.q_proj.weight\", False, 2),\n",
        "        (f\"transformer.h.{layer}.attn.attention.v_proj.weight\", False, 2),\n",
        "        (f\"transformer.h.{layer}.attn.attention.k_proj.weight\", False, 2),\n",
        "        (f\"transformer.h.{layer}.attn.attention.out_proj.bias\", True, None),\n",
        "        (f\"transformer.h.{layer}.attn.attention.out_proj.weight\", False, 1),\n",
        "        (f\"transformer.h.{layer}.mlp.c_fc.bias\", False, 1),\n",
        "        (f\"transformer.h.{layer}.mlp.c_fc.weight\", False, 2),\n",
        "        (f\"transformer.h.{layer}.mlp.c_proj.bias\", True, None),\n",
        "        (f\"transformer.h.{layer}.mlp.c_proj.weight\", False, 1),\n",
        "        (f\"transformer.h.{layer}.ln_1.bias\", False, None),\n",
        "        (f\"transformer.h.{layer}.ln_1.weight\", False, None),\n",
        "        (f\"transformer.h.{layer}.ln_2.bias\", False, None),\n",
        "        (f\"transformer.h.{layer}.ln_2.weight\", False, None),\n",
        "    ])\n",
        "transforms.extend([\n",
        "    (\"transformer.ln_f.bias\", False, None),\n",
        "    (\"transformer.ln_f.weight\", False, None),\n",
        "])\n",
        "\n",
        "for i in range(len(transforms)):\n",
        "    transform = transforms.pop(0)\n",
        "\n",
        "    params = torch_checkpoint[transform[0]]\n",
        "\n",
        "    # Pad input and output embeddings with 0 at the bottom to have 50400 rows\n",
        "    # instead of 50257 rows (the padding value doesn't have to be 0, it doesn't\n",
        "    # even have to be a constant value; the only thing the padding affects is\n",
        "    # it adds junk logits to the end of the logits array the transformer returns\n",
        "    # without affecting the other logits)\n",
        "    if transform[0] in (\"transformer.wte.weight\", \"lm_head.weight\"):\n",
        "        params = torch.cat((params, torch.zeros(143, params.shape[1])))\n",
        "    \n",
        "    # torch.nn.Linear uses a transposed version of the equivalent tensor that\n",
        "    # haiku.Linear uses, so we have to un-transpose the tensor first\n",
        "    if not any(s in transform[0] for s in (\"wte\", \"wpe\")):\n",
        "        params = params.T\n",
        "\n",
        "    if transform[2] is not None:\n",
        "        old_shape = (total_shards,) + get_old_shape(params, transform[2])\n",
        "    else:\n",
        "        old_shape = (total_shards, params.shape[0],)\n",
        "    print(f\"< [{transform[0]}] {params.shape} to {old_shape}\")\n",
        "\n",
        "    params = np.asarray(params[None], dtype=jnp.bfloat16)\n",
        "    params = reshard_reverse(params, old_shape, is_shard_bias=transform[1])\n",
        "\n",
        "    if np.isnan(params).any() or np.isinf(params).any():\n",
        "        raise ValueError(f\"bfloat16 overflow/underflow\")\n",
        "\n",
        "    print(f\"> [{transform[0]}] {params.shape}\")\n",
        "    assert params.shape == old_shape\n",
        "    checkpoint.append(params)\n",
        "\n",
        "# Append the checkpoint step number (can be set to an arbitrary value, in this\n",
        "# case 0, as long as we're only using inference and not training the model)\n",
        "checkpoint.append(np.zeros(total_shards, dtype=np.int32))\n",
        "\n",
        "print(\"saving\")\n",
        "save(checkpoint)\n",
        "del checkpoint\n",
        "print(colored(\"DONE! The JAX checkpoint is now stored at /content/jax_checkpoint\", \"green\"))"
      ],
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "< [transformer.wpe.weight] torch.Size([2048, 2560]) to (4, 2048, 640)\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "> [transformer.wpe.weight] (4, 2048, 640)\n",
            "< [transformer.wte.weight] torch.Size([50400, 2560]) to (4, 12600, 2560)\n",
            "> [transformer.wte.weight] (4, 12600, 2560)\n",
            "< [transformer.h.0.attn.attention.q_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.0.attn.attention.q_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.0.attn.attention.v_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.0.attn.attention.v_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.0.attn.attention.k_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.0.attn.attention.k_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.0.attn.attention.out_proj.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.0.attn.attention.out_proj.bias] (4, 2560)\n",
            "< [transformer.h.0.attn.attention.out_proj.weight] torch.Size([2560, 2560]) to (4, 640, 2560)\n",
            "> [transformer.h.0.attn.attention.out_proj.weight] (4, 640, 2560)\n",
            "< [transformer.h.0.mlp.c_fc.bias] torch.Size([10240]) to (4, 2560)\n",
            "> [transformer.h.0.mlp.c_fc.bias] (4, 2560)\n",
            "< [transformer.h.0.mlp.c_fc.weight] torch.Size([2560, 10240]) to (4, 2560, 2560)\n",
            "> [transformer.h.0.mlp.c_fc.weight] (4, 2560, 2560)\n",
            "< [transformer.h.0.mlp.c_proj.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.0.mlp.c_proj.bias] (4, 2560)\n",
            "< [transformer.h.0.mlp.c_proj.weight] torch.Size([10240, 2560]) to (4, 2560, 2560)\n",
            "> [transformer.h.0.mlp.c_proj.weight] (4, 2560, 2560)\n",
            "< [transformer.h.0.ln_1.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.0.ln_1.bias] (4, 2560)\n",
            "< [transformer.h.0.ln_1.weight] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.0.ln_1.weight] (4, 2560)\n",
            "< [transformer.h.0.ln_2.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.0.ln_2.bias] (4, 2560)\n",
            "< [transformer.h.0.ln_2.weight] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.0.ln_2.weight] (4, 2560)\n",
            "< [transformer.h.1.attn.attention.q_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.1.attn.attention.q_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.1.attn.attention.v_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.1.attn.attention.v_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.1.attn.attention.k_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.1.attn.attention.k_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.1.attn.attention.out_proj.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.1.attn.attention.out_proj.bias] (4, 2560)\n",
            "< [transformer.h.1.attn.attention.out_proj.weight] torch.Size([2560, 2560]) to (4, 640, 2560)\n",
            "> [transformer.h.1.attn.attention.out_proj.weight] (4, 640, 2560)\n",
            "< [transformer.h.1.mlp.c_fc.bias] torch.Size([10240]) to (4, 2560)\n",
            "> [transformer.h.1.mlp.c_fc.bias] (4, 2560)\n",
            "< [transformer.h.1.mlp.c_fc.weight] torch.Size([2560, 10240]) to (4, 2560, 2560)\n",
            "> [transformer.h.1.mlp.c_fc.weight] (4, 2560, 2560)\n",
            "< [transformer.h.1.mlp.c_proj.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.1.mlp.c_proj.bias] (4, 2560)\n",
            "< [transformer.h.1.mlp.c_proj.weight] torch.Size([10240, 2560]) to (4, 2560, 2560)\n",
            "> [transformer.h.1.mlp.c_proj.weight] (4, 2560, 2560)\n",
            "< [transformer.h.1.ln_1.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.1.ln_1.bias] (4, 2560)\n",
            "< [transformer.h.1.ln_1.weight] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.1.ln_1.weight] (4, 2560)\n",
            "< [transformer.h.1.ln_2.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.1.ln_2.bias] (4, 2560)\n",
            "< [transformer.h.1.ln_2.weight] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.1.ln_2.weight] (4, 2560)\n",
            "< [transformer.h.10.attn.attention.q_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.10.attn.attention.q_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.10.attn.attention.v_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.10.attn.attention.v_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.10.attn.attention.k_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.10.attn.attention.k_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.10.attn.attention.out_proj.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.10.attn.attention.out_proj.bias] (4, 2560)\n",
            "< [transformer.h.10.attn.attention.out_proj.weight] torch.Size([2560, 2560]) to (4, 640, 2560)\n",
            "> [transformer.h.10.attn.attention.out_proj.weight] (4, 640, 2560)\n",
            "< [transformer.h.10.mlp.c_fc.bias] torch.Size([10240]) to (4, 2560)\n",
            "> [transformer.h.10.mlp.c_fc.bias] (4, 2560)\n",
            "< [transformer.h.10.mlp.c_fc.weight] torch.Size([2560, 10240]) to (4, 2560, 2560)\n",
            "> [transformer.h.10.mlp.c_fc.weight] (4, 2560, 2560)\n",
            "< [transformer.h.10.mlp.c_proj.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.10.mlp.c_proj.bias] (4, 2560)\n",
            "< [transformer.h.10.mlp.c_proj.weight] torch.Size([10240, 2560]) to (4, 2560, 2560)\n",
            "> [transformer.h.10.mlp.c_proj.weight] (4, 2560, 2560)\n",
            "< [transformer.h.10.ln_1.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.10.ln_1.bias] (4, 2560)\n",
            "< [transformer.h.10.ln_1.weight] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.10.ln_1.weight] (4, 2560)\n",
            "< [transformer.h.10.ln_2.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.10.ln_2.bias] (4, 2560)\n",
            "< [transformer.h.10.ln_2.weight] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.10.ln_2.weight] (4, 2560)\n",
            "< [transformer.h.11.attn.attention.q_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.11.attn.attention.q_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.11.attn.attention.v_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.11.attn.attention.v_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.11.attn.attention.k_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.11.attn.attention.k_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.11.attn.attention.out_proj.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.11.attn.attention.out_proj.bias] (4, 2560)\n",
            "< [transformer.h.11.attn.attention.out_proj.weight] torch.Size([2560, 2560]) to (4, 640, 2560)\n",
            "> [transformer.h.11.attn.attention.out_proj.weight] (4, 640, 2560)\n",
            "< [transformer.h.11.mlp.c_fc.bias] torch.Size([10240]) to (4, 2560)\n",
            "> [transformer.h.11.mlp.c_fc.bias] (4, 2560)\n",
            "< [transformer.h.11.mlp.c_fc.weight] torch.Size([2560, 10240]) to (4, 2560, 2560)\n",
            "> [transformer.h.11.mlp.c_fc.weight] (4, 2560, 2560)\n",
            "< [transformer.h.11.mlp.c_proj.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.11.mlp.c_proj.bias] (4, 2560)\n",
            "< [transformer.h.11.mlp.c_proj.weight] torch.Size([10240, 2560]) to (4, 2560, 2560)\n",
            "> [transformer.h.11.mlp.c_proj.weight] (4, 2560, 2560)\n",
            "< [transformer.h.11.ln_1.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.11.ln_1.bias] (4, 2560)\n",
            "< [transformer.h.11.ln_1.weight] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.11.ln_1.weight] (4, 2560)\n",
            "< [transformer.h.11.ln_2.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.11.ln_2.bias] (4, 2560)\n",
            "< [transformer.h.11.ln_2.weight] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.11.ln_2.weight] (4, 2560)\n",
            "< [transformer.h.12.attn.attention.q_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.12.attn.attention.q_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.12.attn.attention.v_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.12.attn.attention.v_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.12.attn.attention.k_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.12.attn.attention.k_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.12.attn.attention.out_proj.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.12.attn.attention.out_proj.bias] (4, 2560)\n",
            "< [transformer.h.12.attn.attention.out_proj.weight] torch.Size([2560, 2560]) to (4, 640, 2560)\n",
            "> [transformer.h.12.attn.attention.out_proj.weight] (4, 640, 2560)\n",
            "< [transformer.h.12.mlp.c_fc.bias] torch.Size([10240]) to (4, 2560)\n",
            "> [transformer.h.12.mlp.c_fc.bias] (4, 2560)\n",
            "< [transformer.h.12.mlp.c_fc.weight] torch.Size([2560, 10240]) to (4, 2560, 2560)\n",
            "> [transformer.h.12.mlp.c_fc.weight] (4, 2560, 2560)\n",
            "< [transformer.h.12.mlp.c_proj.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.12.mlp.c_proj.bias] (4, 2560)\n",
            "< [transformer.h.12.mlp.c_proj.weight] torch.Size([10240, 2560]) to (4, 2560, 2560)\n",
            "> [transformer.h.12.mlp.c_proj.weight] (4, 2560, 2560)\n",
            "< [transformer.h.12.ln_1.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.12.ln_1.bias] (4, 2560)\n",
            "< [transformer.h.12.ln_1.weight] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.12.ln_1.weight] (4, 2560)\n",
            "< [transformer.h.12.ln_2.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.12.ln_2.bias] (4, 2560)\n",
            "< [transformer.h.12.ln_2.weight] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.12.ln_2.weight] (4, 2560)\n",
            "< [transformer.h.13.attn.attention.q_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.13.attn.attention.q_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.13.attn.attention.v_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.13.attn.attention.v_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.13.attn.attention.k_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.13.attn.attention.k_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.13.attn.attention.out_proj.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.13.attn.attention.out_proj.bias] (4, 2560)\n",
            "< [transformer.h.13.attn.attention.out_proj.weight] torch.Size([2560, 2560]) to (4, 640, 2560)\n",
            "> [transformer.h.13.attn.attention.out_proj.weight] (4, 640, 2560)\n",
            "< [transformer.h.13.mlp.c_fc.bias] torch.Size([10240]) to (4, 2560)\n",
            "> [transformer.h.13.mlp.c_fc.bias] (4, 2560)\n",
            "< [transformer.h.13.mlp.c_fc.weight] torch.Size([2560, 10240]) to (4, 2560, 2560)\n",
            "> [transformer.h.13.mlp.c_fc.weight] (4, 2560, 2560)\n",
            "< [transformer.h.13.mlp.c_proj.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.13.mlp.c_proj.bias] (4, 2560)\n",
            "< [transformer.h.13.mlp.c_proj.weight] torch.Size([10240, 2560]) to (4, 2560, 2560)\n",
            "> [transformer.h.13.mlp.c_proj.weight] (4, 2560, 2560)\n",
            "< [transformer.h.13.ln_1.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.13.ln_1.bias] (4, 2560)\n",
            "< [transformer.h.13.ln_1.weight] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.13.ln_1.weight] (4, 2560)\n",
            "< [transformer.h.13.ln_2.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.13.ln_2.bias] (4, 2560)\n",
            "< [transformer.h.13.ln_2.weight] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.13.ln_2.weight] (4, 2560)\n",
            "< [transformer.h.14.attn.attention.q_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.14.attn.attention.q_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.14.attn.attention.v_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.14.attn.attention.v_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.14.attn.attention.k_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.14.attn.attention.k_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.14.attn.attention.out_proj.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.14.attn.attention.out_proj.bias] (4, 2560)\n",
            "< [transformer.h.14.attn.attention.out_proj.weight] torch.Size([2560, 2560]) to (4, 640, 2560)\n",
            "> [transformer.h.14.attn.attention.out_proj.weight] (4, 640, 2560)\n",
            "< [transformer.h.14.mlp.c_fc.bias] torch.Size([10240]) to (4, 2560)\n",
            "> [transformer.h.14.mlp.c_fc.bias] (4, 2560)\n",
            "< [transformer.h.14.mlp.c_fc.weight] torch.Size([2560, 10240]) to (4, 2560, 2560)\n",
            "> [transformer.h.14.mlp.c_fc.weight] (4, 2560, 2560)\n",
            "< [transformer.h.14.mlp.c_proj.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.14.mlp.c_proj.bias] (4, 2560)\n",
            "< [transformer.h.14.mlp.c_proj.weight] torch.Size([10240, 2560]) to (4, 2560, 2560)\n",
            "> [transformer.h.14.mlp.c_proj.weight] (4, 2560, 2560)\n",
            "< [transformer.h.14.ln_1.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.14.ln_1.bias] (4, 2560)\n",
            "< [transformer.h.14.ln_1.weight] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.14.ln_1.weight] (4, 2560)\n",
            "< [transformer.h.14.ln_2.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.14.ln_2.bias] (4, 2560)\n",
            "< [transformer.h.14.ln_2.weight] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.14.ln_2.weight] (4, 2560)\n",
            "< [transformer.h.15.attn.attention.q_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.15.attn.attention.q_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.15.attn.attention.v_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.15.attn.attention.v_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.15.attn.attention.k_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.15.attn.attention.k_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.15.attn.attention.out_proj.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.15.attn.attention.out_proj.bias] (4, 2560)\n",
            "< [transformer.h.15.attn.attention.out_proj.weight] torch.Size([2560, 2560]) to (4, 640, 2560)\n",
            "> [transformer.h.15.attn.attention.out_proj.weight] (4, 640, 2560)\n",
            "< [transformer.h.15.mlp.c_fc.bias] torch.Size([10240]) to (4, 2560)\n",
            "> [transformer.h.15.mlp.c_fc.bias] (4, 2560)\n",
            "< [transformer.h.15.mlp.c_fc.weight] torch.Size([2560, 10240]) to (4, 2560, 2560)\n",
            "> [transformer.h.15.mlp.c_fc.weight] (4, 2560, 2560)\n",
            "< [transformer.h.15.mlp.c_proj.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.15.mlp.c_proj.bias] (4, 2560)\n",
            "< [transformer.h.15.mlp.c_proj.weight] torch.Size([10240, 2560]) to (4, 2560, 2560)\n",
            "> [transformer.h.15.mlp.c_proj.weight] (4, 2560, 2560)\n",
            "< [transformer.h.15.ln_1.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.15.ln_1.bias] (4, 2560)\n",
            "< [transformer.h.15.ln_1.weight] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.15.ln_1.weight] (4, 2560)\n",
            "< [transformer.h.15.ln_2.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.15.ln_2.bias] (4, 2560)\n",
            "< [transformer.h.15.ln_2.weight] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.15.ln_2.weight] (4, 2560)\n",
            "< [transformer.h.16.attn.attention.q_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.16.attn.attention.q_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.16.attn.attention.v_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.16.attn.attention.v_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.16.attn.attention.k_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.16.attn.attention.k_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.16.attn.attention.out_proj.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.16.attn.attention.out_proj.bias] (4, 2560)\n",
            "< [transformer.h.16.attn.attention.out_proj.weight] torch.Size([2560, 2560]) to (4, 640, 2560)\n",
            "> [transformer.h.16.attn.attention.out_proj.weight] (4, 640, 2560)\n",
            "< [transformer.h.16.mlp.c_fc.bias] torch.Size([10240]) to (4, 2560)\n",
            "> [transformer.h.16.mlp.c_fc.bias] (4, 2560)\n",
            "< [transformer.h.16.mlp.c_fc.weight] torch.Size([2560, 10240]) to (4, 2560, 2560)\n",
            "> [transformer.h.16.mlp.c_fc.weight] (4, 2560, 2560)\n",
            "< [transformer.h.16.mlp.c_proj.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.16.mlp.c_proj.bias] (4, 2560)\n",
            "< [transformer.h.16.mlp.c_proj.weight] torch.Size([10240, 2560]) to (4, 2560, 2560)\n",
            "> [transformer.h.16.mlp.c_proj.weight] (4, 2560, 2560)\n",
            "< [transformer.h.16.ln_1.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.16.ln_1.bias] (4, 2560)\n",
            "< [transformer.h.16.ln_1.weight] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.16.ln_1.weight] (4, 2560)\n",
            "< [transformer.h.16.ln_2.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.16.ln_2.bias] (4, 2560)\n",
            "< [transformer.h.16.ln_2.weight] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.16.ln_2.weight] (4, 2560)\n",
            "< [transformer.h.17.attn.attention.q_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.17.attn.attention.q_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.17.attn.attention.v_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.17.attn.attention.v_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.17.attn.attention.k_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.17.attn.attention.k_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.17.attn.attention.out_proj.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.17.attn.attention.out_proj.bias] (4, 2560)\n",
            "< [transformer.h.17.attn.attention.out_proj.weight] torch.Size([2560, 2560]) to (4, 640, 2560)\n",
            "> [transformer.h.17.attn.attention.out_proj.weight] (4, 640, 2560)\n",
            "< [transformer.h.17.mlp.c_fc.bias] torch.Size([10240]) to (4, 2560)\n",
            "> [transformer.h.17.mlp.c_fc.bias] (4, 2560)\n",
            "< [transformer.h.17.mlp.c_fc.weight] torch.Size([2560, 10240]) to (4, 2560, 2560)\n",
            "> [transformer.h.17.mlp.c_fc.weight] (4, 2560, 2560)\n",
            "< [transformer.h.17.mlp.c_proj.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.17.mlp.c_proj.bias] (4, 2560)\n",
            "< [transformer.h.17.mlp.c_proj.weight] torch.Size([10240, 2560]) to (4, 2560, 2560)\n",
            "> [transformer.h.17.mlp.c_proj.weight] (4, 2560, 2560)\n",
            "< [transformer.h.17.ln_1.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.17.ln_1.bias] (4, 2560)\n",
            "< [transformer.h.17.ln_1.weight] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.17.ln_1.weight] (4, 2560)\n",
            "< [transformer.h.17.ln_2.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.17.ln_2.bias] (4, 2560)\n",
            "< [transformer.h.17.ln_2.weight] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.17.ln_2.weight] (4, 2560)\n",
            "< [transformer.h.18.attn.attention.q_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.18.attn.attention.q_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.18.attn.attention.v_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.18.attn.attention.v_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.18.attn.attention.k_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.18.attn.attention.k_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.18.attn.attention.out_proj.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.18.attn.attention.out_proj.bias] (4, 2560)\n",
            "< [transformer.h.18.attn.attention.out_proj.weight] torch.Size([2560, 2560]) to (4, 640, 2560)\n",
            "> [transformer.h.18.attn.attention.out_proj.weight] (4, 640, 2560)\n",
            "< [transformer.h.18.mlp.c_fc.bias] torch.Size([10240]) to (4, 2560)\n",
            "> [transformer.h.18.mlp.c_fc.bias] (4, 2560)\n",
            "< [transformer.h.18.mlp.c_fc.weight] torch.Size([2560, 10240]) to (4, 2560, 2560)\n",
            "> [transformer.h.18.mlp.c_fc.weight] (4, 2560, 2560)\n",
            "< [transformer.h.18.mlp.c_proj.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.18.mlp.c_proj.bias] (4, 2560)\n",
            "< [transformer.h.18.mlp.c_proj.weight] torch.Size([10240, 2560]) to (4, 2560, 2560)\n",
            "> [transformer.h.18.mlp.c_proj.weight] (4, 2560, 2560)\n",
            "< [transformer.h.18.ln_1.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.18.ln_1.bias] (4, 2560)\n",
            "< [transformer.h.18.ln_1.weight] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.18.ln_1.weight] (4, 2560)\n",
            "< [transformer.h.18.ln_2.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.18.ln_2.bias] (4, 2560)\n",
            "< [transformer.h.18.ln_2.weight] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.18.ln_2.weight] (4, 2560)\n",
            "< [transformer.h.19.attn.attention.q_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.19.attn.attention.q_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.19.attn.attention.v_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.19.attn.attention.v_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.19.attn.attention.k_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.19.attn.attention.k_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.19.attn.attention.out_proj.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.19.attn.attention.out_proj.bias] (4, 2560)\n",
            "< [transformer.h.19.attn.attention.out_proj.weight] torch.Size([2560, 2560]) to (4, 640, 2560)\n",
            "> [transformer.h.19.attn.attention.out_proj.weight] (4, 640, 2560)\n",
            "< [transformer.h.19.mlp.c_fc.bias] torch.Size([10240]) to (4, 2560)\n",
            "> [transformer.h.19.mlp.c_fc.bias] (4, 2560)\n",
            "< [transformer.h.19.mlp.c_fc.weight] torch.Size([2560, 10240]) to (4, 2560, 2560)\n",
            "> [transformer.h.19.mlp.c_fc.weight] (4, 2560, 2560)\n",
            "< [transformer.h.19.mlp.c_proj.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.19.mlp.c_proj.bias] (4, 2560)\n",
            "< [transformer.h.19.mlp.c_proj.weight] torch.Size([10240, 2560]) to (4, 2560, 2560)\n",
            "> [transformer.h.19.mlp.c_proj.weight] (4, 2560, 2560)\n",
            "< [transformer.h.19.ln_1.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.19.ln_1.bias] (4, 2560)\n",
            "< [transformer.h.19.ln_1.weight] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.19.ln_1.weight] (4, 2560)\n",
            "< [transformer.h.19.ln_2.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.19.ln_2.bias] (4, 2560)\n",
            "< [transformer.h.19.ln_2.weight] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.19.ln_2.weight] (4, 2560)\n",
            "< [transformer.h.2.attn.attention.q_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.2.attn.attention.q_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.2.attn.attention.v_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.2.attn.attention.v_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.2.attn.attention.k_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.2.attn.attention.k_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.2.attn.attention.out_proj.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.2.attn.attention.out_proj.bias] (4, 2560)\n",
            "< [transformer.h.2.attn.attention.out_proj.weight] torch.Size([2560, 2560]) to (4, 640, 2560)\n",
            "> [transformer.h.2.attn.attention.out_proj.weight] (4, 640, 2560)\n",
            "< [transformer.h.2.mlp.c_fc.bias] torch.Size([10240]) to (4, 2560)\n",
            "> [transformer.h.2.mlp.c_fc.bias] (4, 2560)\n",
            "< [transformer.h.2.mlp.c_fc.weight] torch.Size([2560, 10240]) to (4, 2560, 2560)\n",
            "> [transformer.h.2.mlp.c_fc.weight] (4, 2560, 2560)\n",
            "< [transformer.h.2.mlp.c_proj.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.2.mlp.c_proj.bias] (4, 2560)\n",
            "< [transformer.h.2.mlp.c_proj.weight] torch.Size([10240, 2560]) to (4, 2560, 2560)\n",
            "> [transformer.h.2.mlp.c_proj.weight] (4, 2560, 2560)\n",
            "< [transformer.h.2.ln_1.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.2.ln_1.bias] (4, 2560)\n",
            "< [transformer.h.2.ln_1.weight] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.2.ln_1.weight] (4, 2560)\n",
            "< [transformer.h.2.ln_2.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.2.ln_2.bias] (4, 2560)\n",
            "< [transformer.h.2.ln_2.weight] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.2.ln_2.weight] (4, 2560)\n",
            "< [transformer.h.20.attn.attention.q_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.20.attn.attention.q_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.20.attn.attention.v_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.20.attn.attention.v_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.20.attn.attention.k_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.20.attn.attention.k_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.20.attn.attention.out_proj.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.20.attn.attention.out_proj.bias] (4, 2560)\n",
            "< [transformer.h.20.attn.attention.out_proj.weight] torch.Size([2560, 2560]) to (4, 640, 2560)\n",
            "> [transformer.h.20.attn.attention.out_proj.weight] (4, 640, 2560)\n",
            "< [transformer.h.20.mlp.c_fc.bias] torch.Size([10240]) to (4, 2560)\n",
            "> [transformer.h.20.mlp.c_fc.bias] (4, 2560)\n",
            "< [transformer.h.20.mlp.c_fc.weight] torch.Size([2560, 10240]) to (4, 2560, 2560)\n",
            "> [transformer.h.20.mlp.c_fc.weight] (4, 2560, 2560)\n",
            "< [transformer.h.20.mlp.c_proj.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.20.mlp.c_proj.bias] (4, 2560)\n",
            "< [transformer.h.20.mlp.c_proj.weight] torch.Size([10240, 2560]) to (4, 2560, 2560)\n",
            "> [transformer.h.20.mlp.c_proj.weight] (4, 2560, 2560)\n",
            "< [transformer.h.20.ln_1.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.20.ln_1.bias] (4, 2560)\n",
            "< [transformer.h.20.ln_1.weight] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.20.ln_1.weight] (4, 2560)\n",
            "< [transformer.h.20.ln_2.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.20.ln_2.bias] (4, 2560)\n",
            "< [transformer.h.20.ln_2.weight] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.20.ln_2.weight] (4, 2560)\n",
            "< [transformer.h.21.attn.attention.q_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.21.attn.attention.q_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.21.attn.attention.v_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.21.attn.attention.v_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.21.attn.attention.k_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.21.attn.attention.k_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.21.attn.attention.out_proj.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.21.attn.attention.out_proj.bias] (4, 2560)\n",
            "< [transformer.h.21.attn.attention.out_proj.weight] torch.Size([2560, 2560]) to (4, 640, 2560)\n",
            "> [transformer.h.21.attn.attention.out_proj.weight] (4, 640, 2560)\n",
            "< [transformer.h.21.mlp.c_fc.bias] torch.Size([10240]) to (4, 2560)\n",
            "> [transformer.h.21.mlp.c_fc.bias] (4, 2560)\n",
            "< [transformer.h.21.mlp.c_fc.weight] torch.Size([2560, 10240]) to (4, 2560, 2560)\n",
            "> [transformer.h.21.mlp.c_fc.weight] (4, 2560, 2560)\n",
            "< [transformer.h.21.mlp.c_proj.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.21.mlp.c_proj.bias] (4, 2560)\n",
            "< [transformer.h.21.mlp.c_proj.weight] torch.Size([10240, 2560]) to (4, 2560, 2560)\n",
            "> [transformer.h.21.mlp.c_proj.weight] (4, 2560, 2560)\n",
            "< [transformer.h.21.ln_1.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.21.ln_1.bias] (4, 2560)\n",
            "< [transformer.h.21.ln_1.weight] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.21.ln_1.weight] (4, 2560)\n",
            "< [transformer.h.21.ln_2.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.21.ln_2.bias] (4, 2560)\n",
            "< [transformer.h.21.ln_2.weight] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.21.ln_2.weight] (4, 2560)\n",
            "< [transformer.h.22.attn.attention.q_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.22.attn.attention.q_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.22.attn.attention.v_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.22.attn.attention.v_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.22.attn.attention.k_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.22.attn.attention.k_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.22.attn.attention.out_proj.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.22.attn.attention.out_proj.bias] (4, 2560)\n",
            "< [transformer.h.22.attn.attention.out_proj.weight] torch.Size([2560, 2560]) to (4, 640, 2560)\n",
            "> [transformer.h.22.attn.attention.out_proj.weight] (4, 640, 2560)\n",
            "< [transformer.h.22.mlp.c_fc.bias] torch.Size([10240]) to (4, 2560)\n",
            "> [transformer.h.22.mlp.c_fc.bias] (4, 2560)\n",
            "< [transformer.h.22.mlp.c_fc.weight] torch.Size([2560, 10240]) to (4, 2560, 2560)\n",
            "> [transformer.h.22.mlp.c_fc.weight] (4, 2560, 2560)\n",
            "< [transformer.h.22.mlp.c_proj.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.22.mlp.c_proj.bias] (4, 2560)\n",
            "< [transformer.h.22.mlp.c_proj.weight] torch.Size([10240, 2560]) to (4, 2560, 2560)\n",
            "> [transformer.h.22.mlp.c_proj.weight] (4, 2560, 2560)\n",
            "< [transformer.h.22.ln_1.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.22.ln_1.bias] (4, 2560)\n",
            "< [transformer.h.22.ln_1.weight] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.22.ln_1.weight] (4, 2560)\n",
            "< [transformer.h.22.ln_2.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.22.ln_2.bias] (4, 2560)\n",
            "< [transformer.h.22.ln_2.weight] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.22.ln_2.weight] (4, 2560)\n",
            "< [transformer.h.23.attn.attention.q_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.23.attn.attention.q_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.23.attn.attention.v_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.23.attn.attention.v_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.23.attn.attention.k_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.23.attn.attention.k_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.23.attn.attention.out_proj.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.23.attn.attention.out_proj.bias] (4, 2560)\n",
            "< [transformer.h.23.attn.attention.out_proj.weight] torch.Size([2560, 2560]) to (4, 640, 2560)\n",
            "> [transformer.h.23.attn.attention.out_proj.weight] (4, 640, 2560)\n",
            "< [transformer.h.23.mlp.c_fc.bias] torch.Size([10240]) to (4, 2560)\n",
            "> [transformer.h.23.mlp.c_fc.bias] (4, 2560)\n",
            "< [transformer.h.23.mlp.c_fc.weight] torch.Size([2560, 10240]) to (4, 2560, 2560)\n",
            "> [transformer.h.23.mlp.c_fc.weight] (4, 2560, 2560)\n",
            "< [transformer.h.23.mlp.c_proj.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.23.mlp.c_proj.bias] (4, 2560)\n",
            "< [transformer.h.23.mlp.c_proj.weight] torch.Size([10240, 2560]) to (4, 2560, 2560)\n",
            "> [transformer.h.23.mlp.c_proj.weight] (4, 2560, 2560)\n",
            "< [transformer.h.23.ln_1.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.23.ln_1.bias] (4, 2560)\n",
            "< [transformer.h.23.ln_1.weight] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.23.ln_1.weight] (4, 2560)\n",
            "< [transformer.h.23.ln_2.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.23.ln_2.bias] (4, 2560)\n",
            "< [transformer.h.23.ln_2.weight] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.23.ln_2.weight] (4, 2560)\n",
            "< [transformer.h.24.attn.attention.q_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.24.attn.attention.q_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.24.attn.attention.v_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.24.attn.attention.v_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.24.attn.attention.k_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.24.attn.attention.k_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.24.attn.attention.out_proj.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.24.attn.attention.out_proj.bias] (4, 2560)\n",
            "< [transformer.h.24.attn.attention.out_proj.weight] torch.Size([2560, 2560]) to (4, 640, 2560)\n",
            "> [transformer.h.24.attn.attention.out_proj.weight] (4, 640, 2560)\n",
            "< [transformer.h.24.mlp.c_fc.bias] torch.Size([10240]) to (4, 2560)\n",
            "> [transformer.h.24.mlp.c_fc.bias] (4, 2560)\n",
            "< [transformer.h.24.mlp.c_fc.weight] torch.Size([2560, 10240]) to (4, 2560, 2560)\n",
            "> [transformer.h.24.mlp.c_fc.weight] (4, 2560, 2560)\n",
            "< [transformer.h.24.mlp.c_proj.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.24.mlp.c_proj.bias] (4, 2560)\n",
            "< [transformer.h.24.mlp.c_proj.weight] torch.Size([10240, 2560]) to (4, 2560, 2560)\n",
            "> [transformer.h.24.mlp.c_proj.weight] (4, 2560, 2560)\n",
            "< [transformer.h.24.ln_1.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.24.ln_1.bias] (4, 2560)\n",
            "< [transformer.h.24.ln_1.weight] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.24.ln_1.weight] (4, 2560)\n",
            "< [transformer.h.24.ln_2.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.24.ln_2.bias] (4, 2560)\n",
            "< [transformer.h.24.ln_2.weight] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.24.ln_2.weight] (4, 2560)\n",
            "< [transformer.h.25.attn.attention.q_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.25.attn.attention.q_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.25.attn.attention.v_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.25.attn.attention.v_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.25.attn.attention.k_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.25.attn.attention.k_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.25.attn.attention.out_proj.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.25.attn.attention.out_proj.bias] (4, 2560)\n",
            "< [transformer.h.25.attn.attention.out_proj.weight] torch.Size([2560, 2560]) to (4, 640, 2560)\n",
            "> [transformer.h.25.attn.attention.out_proj.weight] (4, 640, 2560)\n",
            "< [transformer.h.25.mlp.c_fc.bias] torch.Size([10240]) to (4, 2560)\n",
            "> [transformer.h.25.mlp.c_fc.bias] (4, 2560)\n",
            "< [transformer.h.25.mlp.c_fc.weight] torch.Size([2560, 10240]) to (4, 2560, 2560)\n",
            "> [transformer.h.25.mlp.c_fc.weight] (4, 2560, 2560)\n",
            "< [transformer.h.25.mlp.c_proj.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.25.mlp.c_proj.bias] (4, 2560)\n",
            "< [transformer.h.25.mlp.c_proj.weight] torch.Size([10240, 2560]) to (4, 2560, 2560)\n",
            "> [transformer.h.25.mlp.c_proj.weight] (4, 2560, 2560)\n",
            "< [transformer.h.25.ln_1.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.25.ln_1.bias] (4, 2560)\n",
            "< [transformer.h.25.ln_1.weight] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.25.ln_1.weight] (4, 2560)\n",
            "< [transformer.h.25.ln_2.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.25.ln_2.bias] (4, 2560)\n",
            "< [transformer.h.25.ln_2.weight] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.25.ln_2.weight] (4, 2560)\n",
            "< [transformer.h.26.attn.attention.q_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.26.attn.attention.q_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.26.attn.attention.v_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.26.attn.attention.v_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.26.attn.attention.k_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.26.attn.attention.k_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.26.attn.attention.out_proj.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.26.attn.attention.out_proj.bias] (4, 2560)\n",
            "< [transformer.h.26.attn.attention.out_proj.weight] torch.Size([2560, 2560]) to (4, 640, 2560)\n",
            "> [transformer.h.26.attn.attention.out_proj.weight] (4, 640, 2560)\n",
            "< [transformer.h.26.mlp.c_fc.bias] torch.Size([10240]) to (4, 2560)\n",
            "> [transformer.h.26.mlp.c_fc.bias] (4, 2560)\n",
            "< [transformer.h.26.mlp.c_fc.weight] torch.Size([2560, 10240]) to (4, 2560, 2560)\n",
            "> [transformer.h.26.mlp.c_fc.weight] (4, 2560, 2560)\n",
            "< [transformer.h.26.mlp.c_proj.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.26.mlp.c_proj.bias] (4, 2560)\n",
            "< [transformer.h.26.mlp.c_proj.weight] torch.Size([10240, 2560]) to (4, 2560, 2560)\n",
            "> [transformer.h.26.mlp.c_proj.weight] (4, 2560, 2560)\n",
            "< [transformer.h.26.ln_1.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.26.ln_1.bias] (4, 2560)\n",
            "< [transformer.h.26.ln_1.weight] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.26.ln_1.weight] (4, 2560)\n",
            "< [transformer.h.26.ln_2.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.26.ln_2.bias] (4, 2560)\n",
            "< [transformer.h.26.ln_2.weight] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.26.ln_2.weight] (4, 2560)\n",
            "< [transformer.h.27.attn.attention.q_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.27.attn.attention.q_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.27.attn.attention.v_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.27.attn.attention.v_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.27.attn.attention.k_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.27.attn.attention.k_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.27.attn.attention.out_proj.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.27.attn.attention.out_proj.bias] (4, 2560)\n",
            "< [transformer.h.27.attn.attention.out_proj.weight] torch.Size([2560, 2560]) to (4, 640, 2560)\n",
            "> [transformer.h.27.attn.attention.out_proj.weight] (4, 640, 2560)\n",
            "< [transformer.h.27.mlp.c_fc.bias] torch.Size([10240]) to (4, 2560)\n",
            "> [transformer.h.27.mlp.c_fc.bias] (4, 2560)\n",
            "< [transformer.h.27.mlp.c_fc.weight] torch.Size([2560, 10240]) to (4, 2560, 2560)\n",
            "> [transformer.h.27.mlp.c_fc.weight] (4, 2560, 2560)\n",
            "< [transformer.h.27.mlp.c_proj.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.27.mlp.c_proj.bias] (4, 2560)\n",
            "< [transformer.h.27.mlp.c_proj.weight] torch.Size([10240, 2560]) to (4, 2560, 2560)\n",
            "> [transformer.h.27.mlp.c_proj.weight] (4, 2560, 2560)\n",
            "< [transformer.h.27.ln_1.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.27.ln_1.bias] (4, 2560)\n",
            "< [transformer.h.27.ln_1.weight] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.27.ln_1.weight] (4, 2560)\n",
            "< [transformer.h.27.ln_2.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.27.ln_2.bias] (4, 2560)\n",
            "< [transformer.h.27.ln_2.weight] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.27.ln_2.weight] (4, 2560)\n",
            "< [transformer.h.28.attn.attention.q_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.28.attn.attention.q_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.28.attn.attention.v_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.28.attn.attention.v_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.28.attn.attention.k_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.28.attn.attention.k_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.28.attn.attention.out_proj.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.28.attn.attention.out_proj.bias] (4, 2560)\n",
            "< [transformer.h.28.attn.attention.out_proj.weight] torch.Size([2560, 2560]) to (4, 640, 2560)\n",
            "> [transformer.h.28.attn.attention.out_proj.weight] (4, 640, 2560)\n",
            "< [transformer.h.28.mlp.c_fc.bias] torch.Size([10240]) to (4, 2560)\n",
            "> [transformer.h.28.mlp.c_fc.bias] (4, 2560)\n",
            "< [transformer.h.28.mlp.c_fc.weight] torch.Size([2560, 10240]) to (4, 2560, 2560)\n",
            "> [transformer.h.28.mlp.c_fc.weight] (4, 2560, 2560)\n",
            "< [transformer.h.28.mlp.c_proj.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.28.mlp.c_proj.bias] (4, 2560)\n",
            "< [transformer.h.28.mlp.c_proj.weight] torch.Size([10240, 2560]) to (4, 2560, 2560)\n",
            "> [transformer.h.28.mlp.c_proj.weight] (4, 2560, 2560)\n",
            "< [transformer.h.28.ln_1.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.28.ln_1.bias] (4, 2560)\n",
            "< [transformer.h.28.ln_1.weight] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.28.ln_1.weight] (4, 2560)\n",
            "< [transformer.h.28.ln_2.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.28.ln_2.bias] (4, 2560)\n",
            "< [transformer.h.28.ln_2.weight] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.28.ln_2.weight] (4, 2560)\n",
            "< [transformer.h.29.attn.attention.q_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.29.attn.attention.q_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.29.attn.attention.v_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.29.attn.attention.v_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.29.attn.attention.k_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.29.attn.attention.k_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.29.attn.attention.out_proj.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.29.attn.attention.out_proj.bias] (4, 2560)\n",
            "< [transformer.h.29.attn.attention.out_proj.weight] torch.Size([2560, 2560]) to (4, 640, 2560)\n",
            "> [transformer.h.29.attn.attention.out_proj.weight] (4, 640, 2560)\n",
            "< [transformer.h.29.mlp.c_fc.bias] torch.Size([10240]) to (4, 2560)\n",
            "> [transformer.h.29.mlp.c_fc.bias] (4, 2560)\n",
            "< [transformer.h.29.mlp.c_fc.weight] torch.Size([2560, 10240]) to (4, 2560, 2560)\n",
            "> [transformer.h.29.mlp.c_fc.weight] (4, 2560, 2560)\n",
            "< [transformer.h.29.mlp.c_proj.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.29.mlp.c_proj.bias] (4, 2560)\n",
            "< [transformer.h.29.mlp.c_proj.weight] torch.Size([10240, 2560]) to (4, 2560, 2560)\n",
            "> [transformer.h.29.mlp.c_proj.weight] (4, 2560, 2560)\n",
            "< [transformer.h.29.ln_1.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.29.ln_1.bias] (4, 2560)\n",
            "< [transformer.h.29.ln_1.weight] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.29.ln_1.weight] (4, 2560)\n",
            "< [transformer.h.29.ln_2.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.29.ln_2.bias] (4, 2560)\n",
            "< [transformer.h.29.ln_2.weight] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.29.ln_2.weight] (4, 2560)\n",
            "< [transformer.h.3.attn.attention.q_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.3.attn.attention.q_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.3.attn.attention.v_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.3.attn.attention.v_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.3.attn.attention.k_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.3.attn.attention.k_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.3.attn.attention.out_proj.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.3.attn.attention.out_proj.bias] (4, 2560)\n",
            "< [transformer.h.3.attn.attention.out_proj.weight] torch.Size([2560, 2560]) to (4, 640, 2560)\n",
            "> [transformer.h.3.attn.attention.out_proj.weight] (4, 640, 2560)\n",
            "< [transformer.h.3.mlp.c_fc.bias] torch.Size([10240]) to (4, 2560)\n",
            "> [transformer.h.3.mlp.c_fc.bias] (4, 2560)\n",
            "< [transformer.h.3.mlp.c_fc.weight] torch.Size([2560, 10240]) to (4, 2560, 2560)\n",
            "> [transformer.h.3.mlp.c_fc.weight] (4, 2560, 2560)\n",
            "< [transformer.h.3.mlp.c_proj.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.3.mlp.c_proj.bias] (4, 2560)\n",
            "< [transformer.h.3.mlp.c_proj.weight] torch.Size([10240, 2560]) to (4, 2560, 2560)\n",
            "> [transformer.h.3.mlp.c_proj.weight] (4, 2560, 2560)\n",
            "< [transformer.h.3.ln_1.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.3.ln_1.bias] (4, 2560)\n",
            "< [transformer.h.3.ln_1.weight] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.3.ln_1.weight] (4, 2560)\n",
            "< [transformer.h.3.ln_2.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.3.ln_2.bias] (4, 2560)\n",
            "< [transformer.h.3.ln_2.weight] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.3.ln_2.weight] (4, 2560)\n",
            "< [transformer.h.30.attn.attention.q_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.30.attn.attention.q_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.30.attn.attention.v_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.30.attn.attention.v_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.30.attn.attention.k_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.30.attn.attention.k_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.30.attn.attention.out_proj.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.30.attn.attention.out_proj.bias] (4, 2560)\n",
            "< [transformer.h.30.attn.attention.out_proj.weight] torch.Size([2560, 2560]) to (4, 640, 2560)\n",
            "> [transformer.h.30.attn.attention.out_proj.weight] (4, 640, 2560)\n",
            "< [transformer.h.30.mlp.c_fc.bias] torch.Size([10240]) to (4, 2560)\n",
            "> [transformer.h.30.mlp.c_fc.bias] (4, 2560)\n",
            "< [transformer.h.30.mlp.c_fc.weight] torch.Size([2560, 10240]) to (4, 2560, 2560)\n",
            "> [transformer.h.30.mlp.c_fc.weight] (4, 2560, 2560)\n",
            "< [transformer.h.30.mlp.c_proj.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.30.mlp.c_proj.bias] (4, 2560)\n",
            "< [transformer.h.30.mlp.c_proj.weight] torch.Size([10240, 2560]) to (4, 2560, 2560)\n",
            "> [transformer.h.30.mlp.c_proj.weight] (4, 2560, 2560)\n",
            "< [transformer.h.30.ln_1.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.30.ln_1.bias] (4, 2560)\n",
            "< [transformer.h.30.ln_1.weight] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.30.ln_1.weight] (4, 2560)\n",
            "< [transformer.h.30.ln_2.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.30.ln_2.bias] (4, 2560)\n",
            "< [transformer.h.30.ln_2.weight] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.30.ln_2.weight] (4, 2560)\n",
            "< [transformer.h.31.attn.attention.q_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.31.attn.attention.q_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.31.attn.attention.v_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.31.attn.attention.v_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.31.attn.attention.k_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.31.attn.attention.k_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.31.attn.attention.out_proj.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.31.attn.attention.out_proj.bias] (4, 2560)\n",
            "< [transformer.h.31.attn.attention.out_proj.weight] torch.Size([2560, 2560]) to (4, 640, 2560)\n",
            "> [transformer.h.31.attn.attention.out_proj.weight] (4, 640, 2560)\n",
            "< [transformer.h.31.mlp.c_fc.bias] torch.Size([10240]) to (4, 2560)\n",
            "> [transformer.h.31.mlp.c_fc.bias] (4, 2560)\n",
            "< [transformer.h.31.mlp.c_fc.weight] torch.Size([2560, 10240]) to (4, 2560, 2560)\n",
            "> [transformer.h.31.mlp.c_fc.weight] (4, 2560, 2560)\n",
            "< [transformer.h.31.mlp.c_proj.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.31.mlp.c_proj.bias] (4, 2560)\n",
            "< [transformer.h.31.mlp.c_proj.weight] torch.Size([10240, 2560]) to (4, 2560, 2560)\n",
            "> [transformer.h.31.mlp.c_proj.weight] (4, 2560, 2560)\n",
            "< [transformer.h.31.ln_1.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.31.ln_1.bias] (4, 2560)\n",
            "< [transformer.h.31.ln_1.weight] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.31.ln_1.weight] (4, 2560)\n",
            "< [transformer.h.31.ln_2.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.31.ln_2.bias] (4, 2560)\n",
            "< [transformer.h.31.ln_2.weight] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.31.ln_2.weight] (4, 2560)\n",
            "< [transformer.h.4.attn.attention.q_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.4.attn.attention.q_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.4.attn.attention.v_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.4.attn.attention.v_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.4.attn.attention.k_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.4.attn.attention.k_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.4.attn.attention.out_proj.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.4.attn.attention.out_proj.bias] (4, 2560)\n",
            "< [transformer.h.4.attn.attention.out_proj.weight] torch.Size([2560, 2560]) to (4, 640, 2560)\n",
            "> [transformer.h.4.attn.attention.out_proj.weight] (4, 640, 2560)\n",
            "< [transformer.h.4.mlp.c_fc.bias] torch.Size([10240]) to (4, 2560)\n",
            "> [transformer.h.4.mlp.c_fc.bias] (4, 2560)\n",
            "< [transformer.h.4.mlp.c_fc.weight] torch.Size([2560, 10240]) to (4, 2560, 2560)\n",
            "> [transformer.h.4.mlp.c_fc.weight] (4, 2560, 2560)\n",
            "< [transformer.h.4.mlp.c_proj.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.4.mlp.c_proj.bias] (4, 2560)\n",
            "< [transformer.h.4.mlp.c_proj.weight] torch.Size([10240, 2560]) to (4, 2560, 2560)\n",
            "> [transformer.h.4.mlp.c_proj.weight] (4, 2560, 2560)\n",
            "< [transformer.h.4.ln_1.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.4.ln_1.bias] (4, 2560)\n",
            "< [transformer.h.4.ln_1.weight] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.4.ln_1.weight] (4, 2560)\n",
            "< [transformer.h.4.ln_2.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.4.ln_2.bias] (4, 2560)\n",
            "< [transformer.h.4.ln_2.weight] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.4.ln_2.weight] (4, 2560)\n",
            "< [transformer.h.5.attn.attention.q_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.5.attn.attention.q_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.5.attn.attention.v_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.5.attn.attention.v_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.5.attn.attention.k_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.5.attn.attention.k_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.5.attn.attention.out_proj.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.5.attn.attention.out_proj.bias] (4, 2560)\n",
            "< [transformer.h.5.attn.attention.out_proj.weight] torch.Size([2560, 2560]) to (4, 640, 2560)\n",
            "> [transformer.h.5.attn.attention.out_proj.weight] (4, 640, 2560)\n",
            "< [transformer.h.5.mlp.c_fc.bias] torch.Size([10240]) to (4, 2560)\n",
            "> [transformer.h.5.mlp.c_fc.bias] (4, 2560)\n",
            "< [transformer.h.5.mlp.c_fc.weight] torch.Size([2560, 10240]) to (4, 2560, 2560)\n",
            "> [transformer.h.5.mlp.c_fc.weight] (4, 2560, 2560)\n",
            "< [transformer.h.5.mlp.c_proj.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.5.mlp.c_proj.bias] (4, 2560)\n",
            "< [transformer.h.5.mlp.c_proj.weight] torch.Size([10240, 2560]) to (4, 2560, 2560)\n",
            "> [transformer.h.5.mlp.c_proj.weight] (4, 2560, 2560)\n",
            "< [transformer.h.5.ln_1.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.5.ln_1.bias] (4, 2560)\n",
            "< [transformer.h.5.ln_1.weight] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.5.ln_1.weight] (4, 2560)\n",
            "< [transformer.h.5.ln_2.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.5.ln_2.bias] (4, 2560)\n",
            "< [transformer.h.5.ln_2.weight] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.5.ln_2.weight] (4, 2560)\n",
            "< [transformer.h.6.attn.attention.q_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.6.attn.attention.q_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.6.attn.attention.v_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.6.attn.attention.v_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.6.attn.attention.k_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.6.attn.attention.k_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.6.attn.attention.out_proj.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.6.attn.attention.out_proj.bias] (4, 2560)\n",
            "< [transformer.h.6.attn.attention.out_proj.weight] torch.Size([2560, 2560]) to (4, 640, 2560)\n",
            "> [transformer.h.6.attn.attention.out_proj.weight] (4, 640, 2560)\n",
            "< [transformer.h.6.mlp.c_fc.bias] torch.Size([10240]) to (4, 2560)\n",
            "> [transformer.h.6.mlp.c_fc.bias] (4, 2560)\n",
            "< [transformer.h.6.mlp.c_fc.weight] torch.Size([2560, 10240]) to (4, 2560, 2560)\n",
            "> [transformer.h.6.mlp.c_fc.weight] (4, 2560, 2560)\n",
            "< [transformer.h.6.mlp.c_proj.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.6.mlp.c_proj.bias] (4, 2560)\n",
            "< [transformer.h.6.mlp.c_proj.weight] torch.Size([10240, 2560]) to (4, 2560, 2560)\n",
            "> [transformer.h.6.mlp.c_proj.weight] (4, 2560, 2560)\n",
            "< [transformer.h.6.ln_1.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.6.ln_1.bias] (4, 2560)\n",
            "< [transformer.h.6.ln_1.weight] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.6.ln_1.weight] (4, 2560)\n",
            "< [transformer.h.6.ln_2.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.6.ln_2.bias] (4, 2560)\n",
            "< [transformer.h.6.ln_2.weight] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.6.ln_2.weight] (4, 2560)\n",
            "< [transformer.h.7.attn.attention.q_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.7.attn.attention.q_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.7.attn.attention.v_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.7.attn.attention.v_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.7.attn.attention.k_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.7.attn.attention.k_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.7.attn.attention.out_proj.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.7.attn.attention.out_proj.bias] (4, 2560)\n",
            "< [transformer.h.7.attn.attention.out_proj.weight] torch.Size([2560, 2560]) to (4, 640, 2560)\n",
            "> [transformer.h.7.attn.attention.out_proj.weight] (4, 640, 2560)\n",
            "< [transformer.h.7.mlp.c_fc.bias] torch.Size([10240]) to (4, 2560)\n",
            "> [transformer.h.7.mlp.c_fc.bias] (4, 2560)\n",
            "< [transformer.h.7.mlp.c_fc.weight] torch.Size([2560, 10240]) to (4, 2560, 2560)\n",
            "> [transformer.h.7.mlp.c_fc.weight] (4, 2560, 2560)\n",
            "< [transformer.h.7.mlp.c_proj.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.7.mlp.c_proj.bias] (4, 2560)\n",
            "< [transformer.h.7.mlp.c_proj.weight] torch.Size([10240, 2560]) to (4, 2560, 2560)\n",
            "> [transformer.h.7.mlp.c_proj.weight] (4, 2560, 2560)\n",
            "< [transformer.h.7.ln_1.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.7.ln_1.bias] (4, 2560)\n",
            "< [transformer.h.7.ln_1.weight] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.7.ln_1.weight] (4, 2560)\n",
            "< [transformer.h.7.ln_2.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.7.ln_2.bias] (4, 2560)\n",
            "< [transformer.h.7.ln_2.weight] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.7.ln_2.weight] (4, 2560)\n",
            "< [transformer.h.8.attn.attention.q_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.8.attn.attention.q_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.8.attn.attention.v_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.8.attn.attention.v_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.8.attn.attention.k_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.8.attn.attention.k_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.8.attn.attention.out_proj.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.8.attn.attention.out_proj.bias] (4, 2560)\n",
            "< [transformer.h.8.attn.attention.out_proj.weight] torch.Size([2560, 2560]) to (4, 640, 2560)\n",
            "> [transformer.h.8.attn.attention.out_proj.weight] (4, 640, 2560)\n",
            "< [transformer.h.8.mlp.c_fc.bias] torch.Size([10240]) to (4, 2560)\n",
            "> [transformer.h.8.mlp.c_fc.bias] (4, 2560)\n",
            "< [transformer.h.8.mlp.c_fc.weight] torch.Size([2560, 10240]) to (4, 2560, 2560)\n",
            "> [transformer.h.8.mlp.c_fc.weight] (4, 2560, 2560)\n",
            "< [transformer.h.8.mlp.c_proj.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.8.mlp.c_proj.bias] (4, 2560)\n",
            "< [transformer.h.8.mlp.c_proj.weight] torch.Size([10240, 2560]) to (4, 2560, 2560)\n",
            "> [transformer.h.8.mlp.c_proj.weight] (4, 2560, 2560)\n",
            "< [transformer.h.8.ln_1.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.8.ln_1.bias] (4, 2560)\n",
            "< [transformer.h.8.ln_1.weight] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.8.ln_1.weight] (4, 2560)\n",
            "< [transformer.h.8.ln_2.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.8.ln_2.bias] (4, 2560)\n",
            "< [transformer.h.8.ln_2.weight] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.8.ln_2.weight] (4, 2560)\n",
            "< [transformer.h.9.attn.attention.q_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.9.attn.attention.q_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.9.attn.attention.v_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.9.attn.attention.v_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.9.attn.attention.k_proj.weight] torch.Size([2560, 2560]) to (4, 2560, 640)\n",
            "> [transformer.h.9.attn.attention.k_proj.weight] (4, 2560, 640)\n",
            "< [transformer.h.9.attn.attention.out_proj.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.9.attn.attention.out_proj.bias] (4, 2560)\n",
            "< [transformer.h.9.attn.attention.out_proj.weight] torch.Size([2560, 2560]) to (4, 640, 2560)\n",
            "> [transformer.h.9.attn.attention.out_proj.weight] (4, 640, 2560)\n",
            "< [transformer.h.9.mlp.c_fc.bias] torch.Size([10240]) to (4, 2560)\n",
            "> [transformer.h.9.mlp.c_fc.bias] (4, 2560)\n",
            "< [transformer.h.9.mlp.c_fc.weight] torch.Size([2560, 10240]) to (4, 2560, 2560)\n",
            "> [transformer.h.9.mlp.c_fc.weight] (4, 2560, 2560)\n",
            "< [transformer.h.9.mlp.c_proj.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.9.mlp.c_proj.bias] (4, 2560)\n",
            "< [transformer.h.9.mlp.c_proj.weight] torch.Size([10240, 2560]) to (4, 2560, 2560)\n",
            "> [transformer.h.9.mlp.c_proj.weight] (4, 2560, 2560)\n",
            "< [transformer.h.9.ln_1.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.9.ln_1.bias] (4, 2560)\n",
            "< [transformer.h.9.ln_1.weight] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.9.ln_1.weight] (4, 2560)\n",
            "< [transformer.h.9.ln_2.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.9.ln_2.bias] (4, 2560)\n",
            "< [transformer.h.9.ln_2.weight] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.h.9.ln_2.weight] (4, 2560)\n",
            "< [transformer.ln_f.bias] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.ln_f.bias] (4, 2560)\n",
            "< [transformer.ln_f.weight] torch.Size([2560]) to (4, 2560)\n",
            "> [transformer.ln_f.weight] (4, 2560)\n",
            "saving\n",
            "\u001b[32mDONE! The JAX checkpoint is now stored at /content/jax_checkpoint\u001b[0m\n"
          ]
        }
      ]
    }
  ]
}