NeoX_20b.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 2,
  "metadata": {
    "colab": {
      "name": "NeoX_20b",
      "private_outputs": true,
      "provenance": [],
      "machine_shape": "hm",
      "background_execution": "on",
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3.10.4 64-bit"
    },
    "interpreter": {
      "hash": "97cc609b13305c559618ec78a438abc56230b9381f827f22d070313b9a1f3777"
    },
    "language_info": {
      "name": "python",
      "version": "3.10.4",
      "mimetype": "text/x-python",
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "pygments_lexer": "ipython3",
      "nbconvert_exporter": "python",
      "file_extension": ".py"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# installs"
      ],
      "metadata": {
        "id": "_7yXuY9uUtym"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "source": [
        "!pip install tokenizers==0.12.0\n",
        "!pip install transformers"
      ],
      "outputs": [],
      "metadata": {
        "id": "7c-D1-JyUtPm"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "source": [
        "# some code"
      ],
      "outputs": [],
      "metadata": {
        "id": "Nzho67_KUvrL"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "source": [
        "!ls"
      ],
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            " 20B_checkpoints   20b.ipynb  'NeoX_20b (1).ipynb'   nohup.out\t wget-log\n"
          ]
        }
      ],
      "metadata": {}
    },
    {
      "cell_type": "code",
      "execution_count": 23,
      "source": [
        "class Args20b:\n",
        "    vocab_size = 50432\n",
        "    hidden_size = 6144\n",
        "    num_attention_heads = 64\n",
        "    rotary_pct = 0.25\n",
        "    rotary_emb_base = 10000\n",
        "    layernorm_epsilon = 1e-5\n",
        "    num_layers = 44\n"
      ],
      "outputs": [],
      "metadata": {
        "id": "VEu7gwNFkO6A"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 24,
      "source": [
        "import torch\n",
        "\n",
        "\n",
        "class RotaryEmbedding(torch.nn.Module):\n",
        "\n",
        "    def __init__(self, dim, base=10000, device=None):\n",
        "        super().__init__()\n",
        "        inv_freq = 1. / (base ** (torch.arange(0, dim, 2).float().to(device) / dim))\n",
        "        self.register_buffer('inv_freq', inv_freq)\n",
        "        self.max_seq_len_cached = None\n",
        "        self.cos_cached = None\n",
        "        self.sin_cached = None\n",
        "\n",
        "    def forward(self, x, seq_dim=1, seq_len=None):\n",
        "        if seq_len is None:\n",
        "            seq_len = x.shape[seq_dim]\n",
        "        if self.max_seq_len_cached is None or (seq_len > self.max_seq_len_cached):\n",
        "            self.max_seq_len_cached = seq_len\n",
        "            t = torch.arange(self.max_seq_len_cached, device=x.device, dtype=self.inv_freq.dtype)\n",
        "            freqs = torch.einsum('i,j->ij', t, self.inv_freq)\n",
        "            # Different from paper, but it uses a different permutation in order to obtain the same calculation\n",
        "            emb = torch.cat((freqs, freqs), dim=-1).to(x.device)\n",
        "            # [sx, 1 (b * np), hn]\n",
        "            self.cos_cached = emb.cos()[:, None, None, :]\n",
        "            self.sin_cached = emb.sin()[:, None, None, :]\n",
        "        return self.cos_cached[:seq_len, ...], self.sin_cached[:seq_len, ...]\n",
        "\n",
        "\n",
        "def rotate_half(x):\n",
        "    x1, x2 = x[..., :x.shape[-1] // 2], x[..., x.shape[-1] // 2:]\n",
        "    return torch.cat((-x2, x1), dim=x1.ndim - 1)  # dim=-1 triggers a bug in earlier torch versions\n",
        "\n",
        "\n",
        "# @torch.jit.script\n",
        "def apply_rotary_pos_emb(q, k, cos, sin, offset: int = 0):\n",
        "    cos, sin = cos[offset:q.shape[0] + offset, ...], sin[offset:q.shape[0] + offset, ...]\n",
        "    return (q * cos) + (rotate_half(q) * sin), (k * cos) + (rotate_half(k) * sin)"
      ],
      "outputs": [],
      "metadata": {
        "id": "5BDC8mK9Ugjf"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 25,
      "source": [
        "import torch.nn as nn\n",
        "import torch\n",
        "import math\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "class NeoX20BModel(nn.Module):\n",
        "    def __init__(self, args, use_cache=False, device=None):\n",
        "        super().__init__()\n",
        "        self.use_cache = use_cache\n",
        "        self.embed_in = nn.Embedding(args.vocab_size, args.hidden_size, device=device)\n",
        "        self.layer_list = nn.ModuleList([])\n",
        "        for layer_i in range(args.num_layers):\n",
        "            self.layer_list.append(TransformerLayer(args, use_cache, device=device))\n",
        "        self.final_layer_norm = nn.LayerNorm(\n",
        "            args.hidden_size,\n",
        "            eps=args.layernorm_epsilon,\n",
        "            device=device,\n",
        "        )\n",
        "        self.logits_out = nn.Linear(\n",
        "            args.hidden_size,\n",
        "            args.vocab_size,\n",
        "            bias=False,\n",
        "            device=device,\n",
        "        )\n",
        "\n",
        "    def forward(self, x, attention_mask=None, layer_past=None):\n",
        "        if attention_mask is None:\n",
        "            attention_mask = generate_mask(x.shape[1]).to(x.device)\n",
        "        if self.use_cache:\n",
        "            if layer_past is None:\n",
        "                kv_length = x.shape[1]\n",
        "            else:\n",
        "                kv_length = layer_past[0].shape[1] + 1\n",
        "            attention_mask = attention_mask[..., :x.shape[1], :kv_length]\n",
        "\n",
        "        if layer_past is None:\n",
        "            layer_past = [None] * len(self.layer_list)\n",
        "        kv_cache_list = []\n",
        "        hidden_states = self.embed_in(x)\n",
        "        hidden_states = self.pre_transformer_transpose(hidden_states)\n",
        "\n",
        "        for layer_i, layer in enumerate(self.layer_list):\n",
        "            hidden_states, kv_cache = layer(\n",
        "                x=hidden_states,\n",
        "                attention_mask=attention_mask,\n",
        "                layer_past=layer_past[layer_i],\n",
        "            )\n",
        "            kv_cache_list.append(kv_cache)\n",
        "        hidden_states = self.post_transformer_transpose(hidden_states)\n",
        "        hidden_states = self.final_layer_norm(hidden_states)\n",
        "        logits = self.logits_out(hidden_states)\n",
        "        if self.use_cache:\n",
        "            return logits, kv_cache_list\n",
        "        else:\n",
        "            return logits\n",
        "\n",
        "    @classmethod\n",
        "    def pre_transformer_transpose(cls, x):\n",
        "        return x.transpose(0, 1).contiguous()\n",
        "\n",
        "    @classmethod\n",
        "    def post_transformer_transpose(cls, x):\n",
        "        return x.transpose(0, 1).contiguous()\n",
        "\n",
        "\n",
        "class TransformerLayer(nn.Module):\n",
        "    def __init__(self, args, use_cache, device=None):\n",
        "        super().__init__()\n",
        "        self.use_cache = use_cache\n",
        "        self.input_layernorm = nn.LayerNorm(\n",
        "            args.hidden_size,\n",
        "            eps=args.layernorm_epsilon,\n",
        "            device=device,\n",
        "        )\n",
        "        self.post_attention_layernorm = nn.LayerNorm(\n",
        "            args.hidden_size,\n",
        "            eps=args.layernorm_epsilon,\n",
        "            device=device,\n",
        "        )\n",
        "        self.attention = SelfAttention(args, self.use_cache, device=device)\n",
        "        self.mlp = MLP(args)\n",
        "\n",
        "    def forward(self, x, attention_mask, layer_past=None):\n",
        "        residual = x\n",
        "        ln_output = self.input_layernorm(x)\n",
        "        attention_output, kv_cache = self.attention(\n",
        "            ln_output,\n",
        "            attention_mask,\n",
        "            layer_past=layer_past,\n",
        "        )\n",
        "        post_attn_ln = self.post_attention_layernorm(x)\n",
        "        mlp_output = self.mlp(hidden_states=post_attn_ln)\n",
        "        output = residual + mlp_output + attention_output\n",
        "        return output, kv_cache\n",
        "\n",
        "\n",
        "class SelfAttention(nn.Module):\n",
        "    def __init__(self, args, use_cache=False, device=None):\n",
        "        super().__init__()\n",
        "        self.hidden_size = args.hidden_size\n",
        "        self.use_cache = use_cache\n",
        "        self.num_attention_heads = args.num_attention_heads\n",
        "        self.hidden_size_per_attention_head = args.hidden_size // args.num_attention_heads\n",
        "        self.rotary_ndims = int(self.hidden_size_per_attention_head * args.rotary_pct)\n",
        "        self.rotary_emb = RotaryEmbedding(\n",
        "            self.rotary_ndims,\n",
        "            base=args.rotary_emb_base,\n",
        "            device=device,\n",
        "        )\n",
        "        self.query_key_value = nn.Linear(\n",
        "            args.hidden_size,\n",
        "            3 * args.hidden_size,\n",
        "            device=device,\n",
        "        )\n",
        "        self.norm_factor = math.sqrt(self.hidden_size_per_attention_head)\n",
        "        self.dense = nn.Linear(\n",
        "            args.hidden_size,\n",
        "            args.hidden_size,\n",
        "            device=device,\n",
        "        )\n",
        "\n",
        "    def forward(self, hidden_states, attention_mask, layer_past=None):\n",
        "        has_layer_past = layer_past is not None and layer_past.numel() > 0\n",
        "\n",
        "        # Compute QKV\n",
        "        # Attention heads [sq, b, h] --> [sq, b, (np * 3 * hn)]\n",
        "        qkv = self.query_key_value(hidden_states)\n",
        "\n",
        "        # [sq, b, (np * 3 * hn)] --> [sq, b, np, 3 * hn]\n",
        "        new_qkv_shape = qkv.size()[:-1] + (\n",
        "            self.num_attention_heads,\n",
        "            3 * self.hidden_size_per_attention_head,\n",
        "        )\n",
        "        qkv = qkv.view(*new_qkv_shape)\n",
        "\n",
        "        # [sq, b, np, 3 * hn] --> 3 [sq, b, np, hn]\n",
        "        query_layer = qkv[..., :self.hidden_size_per_attention_head]\n",
        "        key_layer = qkv[..., self.hidden_size_per_attention_head: 2 * self.hidden_size_per_attention_head]\n",
        "        value_layer = qkv[..., 2 * self.hidden_size_per_attention_head:]\n",
        "\n",
        "        # Compute rotary embeddings\n",
        "        query_rot, query_pass = (\n",
        "            query_layer[..., : self.rotary_ndims],\n",
        "            query_layer[..., self.rotary_ndims:],\n",
        "        )\n",
        "        key_rot, key_pass = (\n",
        "            key_layer[..., : self.rotary_ndims],\n",
        "            key_layer[..., self.rotary_ndims:],\n",
        "        )\n",
        "        seq_len = key_layer.shape[0]\n",
        "        offset = 0\n",
        "        if has_layer_past:\n",
        "            offset = layer_past[0].shape[0]\n",
        "            seq_len += offset\n",
        "        cos, sin = self.rotary_emb(value_layer, seq_len=seq_len)\n",
        "        query_layer, key_layer = apply_rotary_pos_emb(\n",
        "            query_rot, key_rot, cos, sin, offset=offset,\n",
        "        )\n",
        "        query_layer = torch.cat((query_layer, query_pass), dim=-1)\n",
        "        key_layer = torch.cat((key_layer, key_pass), dim=-1)\n",
        "\n",
        "        # Cache QKV values\n",
        "        if has_layer_past:\n",
        "            past_key, past_value = layer_past\n",
        "            key_layer = torch.cat((past_key.type_as(key_layer), key_layer), dim=0)\n",
        "            value_layer = torch.cat((past_value.type_as(value_layer), value_layer), dim=0)\n",
        "        if self.use_cache:\n",
        "            kv_cache = torch.stack((key_layer, value_layer))\n",
        "        else:\n",
        "            kv_cache = None\n",
        "\n",
        "        # Compute attention\n",
        "        # noinspection PyTypeChecker\n",
        "        context_layer = self.attention(\n",
        "            query_layer, key_layer, value_layer, attention_mask\n",
        "        )\n",
        "\n",
        "        # Reshape outputs\n",
        "        # [b, np, sq, hn] --> [sq, b, np, hn]\n",
        "        context_layer = context_layer.permute(2, 0, 1, 3).contiguous()\n",
        "\n",
        "        # [sq, b, np, hn] --> [sq, b, hp]\n",
        "        new_context_layer_shape = context_layer.size()[:-2] + (\n",
        "            self.hidden_size,\n",
        "        )\n",
        "        context_layer = context_layer.view(*new_context_layer_shape)\n",
        "\n",
        "        # =================\n",
        "        # Output. [sq, b, h]\n",
        "        # =================\n",
        "        output = self.dense(context_layer)\n",
        "\n",
        "        return output, kv_cache\n",
        "\n",
        "    def attention(self, query_layer, key_layer, value_layer, attention_mask):\n",
        "        # ===================================\n",
        "        # Raw attention scores. [b, np, s, s]\n",
        "        # ===================================\n",
        "\n",
        "        # [b, np, sq, sk]\n",
        "        output_size = (\n",
        "            query_layer.size(1),\n",
        "            query_layer.size(2),\n",
        "            query_layer.size(0),\n",
        "            key_layer.size(0),\n",
        "        )\n",
        "\n",
        "        # [sq, b, np, hn] -> [sq, b * np, hn]\n",
        "        query_layer = query_layer.view(\n",
        "            output_size[2], output_size[0] * output_size[1], -1\n",
        "        )\n",
        "        key_layer = key_layer.view(output_size[3], output_size[0] * output_size[1], -1)\n",
        "\n",
        "        # preallocating result tensor: [b * np, sq, sk]\n",
        "        matmul_result = torch.empty(\n",
        "            output_size[0] * output_size[1],\n",
        "            output_size[2],\n",
        "            output_size[3],\n",
        "            dtype=query_layer.dtype,\n",
        "            device=query_layer.device,\n",
        "        )\n",
        "\n",
        "        # Raw attention scores. [b * np, sq, sk]\n",
        "        matmul_result = torch.baddbmm(\n",
        "            matmul_result,\n",
        "            query_layer.transpose(0, 1),  # [b * np, sq, hn]\n",
        "            key_layer.transpose(0, 1).transpose(1, 2),  # [b * np, hn, sk]\n",
        "            beta=0.0,\n",
        "            alpha=(1.0 / self.norm_factor),\n",
        "        )\n",
        "\n",
        "        # change view to [b, np, sq, sk]\n",
        "        attention_scores = matmul_result.view(*output_size)\n",
        "\n",
        "        # ==================================================\n",
        "        # Update attention mask for inference. [b, np, sq, sk]\n",
        "        # ==================================================\n",
        "\n",
        "        # ===========================\n",
        "        # Attention probs and dropout\n",
        "        # ===========================\n",
        "\n",
        "        # attention scores and attention mask [b, np, sq, sk]\n",
        "        masked_scores = attention_mask_func(attention_scores, attention_mask) \\\n",
        "            if attention_mask is not None else attention_scores\n",
        "        attention_probs = torch.nn.Softmax(dim=-1)(masked_scores)\n",
        "\n",
        "        #         # This is actually dropping out entire tokens to attend to, which might\n",
        "        #         # seem a bit unusual, but is taken from the original Transformer paper.\n",
        "        #         attention_probs = self.attention_dropout(attention_probs)\n",
        "\n",
        "        # =========================\n",
        "        # Context layer. [sq, b, hp]\n",
        "        # =========================\n",
        "\n",
        "        # value_layer -> context layer.\n",
        "        # [sk, b, np, hn] --> [b, np, sq, hn]\n",
        "\n",
        "        # context layer shape: [b, np, sq, hn]\n",
        "        output_size = (\n",
        "            value_layer.size(1),\n",
        "            value_layer.size(2),\n",
        "            query_layer.size(0),\n",
        "            value_layer.size(3),\n",
        "        )\n",
        "\n",
        "        # change view [sk, b * np, hn]\n",
        "        value_layer = value_layer.view(\n",
        "            value_layer.size(0), output_size[0] * output_size[1], -1\n",
        "        )\n",
        "\n",
        "        # change view [b * np, sq, sk]\n",
        "        attention_probs = attention_probs.view(\n",
        "            output_size[0] * output_size[1], output_size[2], -1\n",
        "        )\n",
        "\n",
        "        # matmul: [b * np, sq, hn]\n",
        "        context_layer = torch.bmm(attention_probs, value_layer.transpose(0, 1))\n",
        "\n",
        "        # change view [b, np, sq, hn]\n",
        "        context_layer = context_layer.view(*output_size)\n",
        "        return context_layer\n",
        "\n",
        "\n",
        "class MLP(nn.Module):\n",
        "    def __init__(self, args, device=None):\n",
        "        super().__init__()\n",
        "        ff_dim = 4 * args.hidden_size\n",
        "        self.dense_h_to_4h = nn.Linear(args.hidden_size, ff_dim, device=device)\n",
        "        self.dense_4h_to_h = nn.Linear(ff_dim, args.hidden_size, device=device)\n",
        "\n",
        "    def forward(self, hidden_states):\n",
        "        intermediate_parallel = self.dense_h_to_4h(hidden_states)\n",
        "        intermediate_parallel = bias_gelu_impl(intermediate_parallel)\n",
        "        output = self.dense_4h_to_h(intermediate_parallel)\n",
        "        return output\n",
        "\n",
        "\n",
        "# noinspection PyAbstractClass\n",
        "class GeLUFunction(torch.autograd.Function):\n",
        "    # noinspection PyMethodOverriding\n",
        "    @staticmethod\n",
        "    # bias is an optional argument\n",
        "    def forward(ctx, inputs):\n",
        "        ctx.save_for_backward(inputs)\n",
        "        return gelu(inputs)\n",
        "\n",
        "    # noinspection PyMethodOverriding\n",
        "    @staticmethod\n",
        "    def backward(ctx, grad_output):\n",
        "        inputs = ctx.saved_tensors\n",
        "        tmp = gelu_back(grad_output, inputs)\n",
        "        return tmp, tmp\n",
        "\n",
        "\n",
        "bias_gelu_impl = GeLUFunction.apply\n",
        "\n",
        "\n",
        "def generate_mask(seq_len):\n",
        "    return torch.tril(torch.ones((1, 1, seq_len, seq_len), dtype=torch.bool))\n",
        "\n",
        "\n",
        "def attention_mask_func(attention_scores, ltor_mask):\n",
        "    \"\"\"Assign -10000.0 to False cells in ltor_mask\"\"\"\n",
        "    attention_scores.masked_fill_(~ltor_mask, -10000.0)\n",
        "    return attention_scores\n",
        "\n",
        "\n",
        "@torch.jit.script\n",
        "def gelu(x):\n",
        "    return x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))\n",
        "\n",
        "\n",
        "# gradient of tanh approximation of gelu\n",
        "# gradient of actual gelu is:\n",
        "# 0.5 * (1. + torch.erf(x * 0.70710678)) + 0.3989423 * x * torch.exp(-0.5 * x * x)\n",
        "@torch.jit.script\n",
        "def gelu_back(g, x):\n",
        "    tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x))\n",
        "    # sqrt(2/pi) * 3 * 0.044715 -> 0.1070322243\n",
        "    ff = 0.5 * x * (\n",
        "            (1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)\n",
        "    ) + 0.5 * (1 + tanh_out)\n",
        "    return ff * g"
      ],
      "outputs": [],
      "metadata": {
        "id": "dBFKeAEcUdSm"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 26,
      "source": [
        "import os\n",
        "from tqdm import auto as tqdm_lib\n",
        "\n",
        "import torch\n",
        "import tokenizers\n",
        "\n",
        "\n",
        "\n",
        "def create_model(checkpoint_path, use_cache=False, device=torch.device(\"cuda:0\")):\n",
        "    \"\"\"\n",
        "    To prevent allocation memory on CPU, we initialize on 'meta' and individually\n",
        "    port each module over to 'device' as we load each state dict.\n",
        "    :param checkpoint_path: Path to the checkpoint folder\n",
        "    :param use_cache: whether to use cache (i.e. for efficient generation)\n",
        "    :param device: device that you want the model to end up on\n",
        "    :return: model\n",
        "    \"\"\"\n",
        "    # Instantiate model\n",
        "    pbar = tqdm_lib.tqdm(total=48)\n",
        "    pbar.set_description(\"Instantiating model (~1 min)\")\n",
        "    model = NeoX20BModel(Args20b, use_cache=use_cache, device=\"meta\")\n",
        "    if 'cuda' in device:\n",
        "        model = model.half().to_empty(device=device)\n",
        "    else:\n",
        "        model = model.to_empty(device=device)\n",
        "    pbar.update(1)\n",
        "\n",
        "    # Load transformer layers\n",
        "    for layer_i in range(Args20b.num_layers):\n",
        "        pbar.set_description(f\"Loading layer {layer_i}\")\n",
        "        filename_tp1 = f\"layer_{layer_i + 2:02d}-model_00-model_states.pt\"\n",
        "        filename_tp2 = f\"layer_{layer_i + 2:02d}-model_01-model_states.pt\"\n",
        "        loaded_tp1 = torch.load(os.path.join(checkpoint_path, filename_tp1))\n",
        "        loaded_tp2 = torch.load(os.path.join(checkpoint_path, filename_tp2))\n",
        "        state_dict = {}\n",
        "        # Good\n",
        "        # Keys where we concatenate on the second dim\n",
        "        for key in [\n",
        "            \"attention.dense.weight\",\n",
        "            \"mlp.dense_4h_to_h.weight\",\n",
        "        ]:\n",
        "            state_dict[key] = torch.cat([loaded_tp1[key], loaded_tp2[key]], dim=1)\n",
        "        # Mapping individual split weights to custom split implementations\n",
        "        # Layer Norms\n",
        "        # Choose 1\n",
        "        state_dict[\"input_layernorm.weight\"] = (\n",
        "            loaded_tp1[\"input_layernorm.weight\"] + loaded_tp2[\"input_layernorm.weight\"]) / 2\n",
        "        state_dict[\"input_layernorm.bias\"] = (\n",
        "            loaded_tp1[\"input_layernorm.bias\"] + loaded_tp2[\"input_layernorm.bias\"]) / 2\n",
        "        state_dict[\"post_attention_layernorm.weight\"] = (\n",
        "            loaded_tp1[\"post_attention_layernorm.weight\"] + loaded_tp2[\"post_attention_layernorm.weight\"]) / 2\n",
        "        state_dict[\"post_attention_layernorm.bias\"] = (\n",
        "            loaded_tp1[\"post_attention_layernorm.bias\"] + loaded_tp2[\"post_attention_layernorm.bias\"]) / 2\n",
        "        # LinearWithTPMerge\n",
        "        state_dict[\"mlp.dense_h_to_4h.weight\"] = torch.cat([\n",
        "            loaded_tp1[\"mlp.dense_h_to_4h.weight\"],\n",
        "            loaded_tp2[\"mlp.dense_h_to_4h.weight\"],\n",
        "        ], dim=0)\n",
        "        state_dict[\"mlp.dense_h_to_4h.bias\"] = torch.cat([\n",
        "            loaded_tp1[\"mlp.dense_h_to_4h.bias\"],\n",
        "            loaded_tp2[\"mlp.dense_h_to_4h.bias\"],\n",
        "        ], dim=0)\n",
        "        state_dict[\"attention.query_key_value.weight\"] = torch.cat([\n",
        "            loaded_tp1[\"attention.query_key_value.weight\"],\n",
        "            loaded_tp2[\"attention.query_key_value.weight\"],\n",
        "        ], dim=0)\n",
        "        state_dict[\"attention.query_key_value.bias\"] = torch.cat([\n",
        "            loaded_tp1[\"attention.query_key_value.bias\"],\n",
        "            loaded_tp2[\"attention.query_key_value.bias\"],\n",
        "        ], dim=0)\n",
        "        # LinearWithTPSplitBias\n",
        "        state_dict[\"mlp.dense_4h_to_h.bias\"] = (\n",
        "            loaded_tp1[\"mlp.dense_4h_to_h.bias\"]\n",
        "            + loaded_tp2[\"mlp.dense_4h_to_h.bias\"]\n",
        "        )\n",
        "        state_dict[\"attention.dense.bias\"] = (\n",
        "            loaded_tp1[\"attention.dense.bias\"]\n",
        "            + loaded_tp2[\"attention.dense.bias\"]\n",
        "        )\n",
        "        # Just take one\n",
        "        state_dict[\"attention.rotary_emb.inv_freq\"] = loaded_tp1[\"attention.rotary_emb.inv_freq\"]\n",
        "        model.layer_list[layer_i].load_state_dict(state_dict)\n",
        "        del loaded_tp1\n",
        "        del loaded_tp2\n",
        "        pbar.update(1)\n",
        "\n",
        "    # Load input embedding\n",
        "    pbar.set_description(f\"Loading input embedding\")\n",
        "    loaded_tp1 = torch.load(os.path.join(checkpoint_path, \"layer_00-model_00-model_states.pt\"))\n",
        "    loaded_tp2 = torch.load(os.path.join(checkpoint_path, \"layer_00-model_01-model_states.pt\"))\n",
        "    model.embed_in.load_state_dict({\"weight\": torch.cat([\n",
        "        loaded_tp1[\"word_embeddings.weight\"],\n",
        "        loaded_tp2[\"word_embeddings.weight\"],\n",
        "    ], dim=0)})\n",
        "    del loaded_tp1\n",
        "    del loaded_tp2\n",
        "    pbar.update(1)\n",
        "\n",
        "    # Load final layer norm\n",
        "    pbar.set_description(f\"Loading final layer norm\")\n",
        "    loaded_tp1 = torch.load(os.path.join(checkpoint_path, \"layer_47-model_00-model_states.pt\"))\n",
        "    loaded_tp2 = torch.load(os.path.join(checkpoint_path, \"layer_47-model_01-model_states.pt\"))\n",
        "    model.final_layer_norm.load_state_dict({\n",
        "        \"weight\": (loaded_tp1[\"norm.weight\"] + loaded_tp2[\"norm.weight\"])/2,\n",
        "        \"bias\": (loaded_tp1[\"norm.bias\"] + loaded_tp2[\"norm.bias\"])/2,\n",
        "    })\n",
        "    del loaded_tp1\n",
        "    del loaded_tp2\n",
        "    pbar.update(1)\n",
        "\n",
        "    # Load output embedding\n",
        "    pbar.set_description(f\"Loading output embedding\")\n",
        "    loaded_tp1 = torch.load(os.path.join(checkpoint_path, \"layer_48-model_00-model_states.pt\"))\n",
        "    loaded_tp2 = torch.load(os.path.join(checkpoint_path, \"layer_48-model_01-model_states.pt\"))\n",
        "    model.logits_out.load_state_dict({\n",
        "        \"weight\": torch.cat([\n",
        "            loaded_tp1[\"final_linear.weight\"],\n",
        "            loaded_tp2[\"final_linear.weight\"],\n",
        "        ], dim=0),\n",
        "    })\n",
        "    del loaded_tp1\n",
        "    del loaded_tp2\n",
        "    pbar.update(1)\n",
        "    pbar.set_description(\"Done.\")\n",
        "\n",
        "    return model\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "def create_tokenizer(tokenizer_path):\n",
        "    return tokenizers.Tokenizer.from_file(tokenizer_path)"
      ],
      "outputs": [],
      "metadata": {
        "id": "NjeoQnoRUYhK"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 27,
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "from tqdm import auto as tqdm_lib\n",
        "\n",
        "\n",
        "def greedy_generate(model: nn.Module, input_ids: torch.Tensor, max_seq_len: int,\n",
        "                    verbose=True):\n",
        "    \"\"\"Generate greedily from 20B.\n",
        "    :param model: NeoX20BModel\n",
        "    :param input_ids: token IDs [batch_size, seq_len]\n",
        "    :param max_seq_len: max sequence length to generate up to (includes input_ids)\n",
        "    :param verbose: whether to print progress\n",
        "    :return: List of token IDs\n",
        "    \"\"\"\n",
        "    initial_input_length = input_ids.shape[1]\n",
        "    current_input_ids = input_ids\n",
        "    layer_past = None\n",
        "    layer_past_length = 0\n",
        "    all_token_ids = input_ids.tolist()\n",
        "    batch_size = len(all_token_ids)\n",
        "\n",
        "    if verbose:\n",
        "        trange = tqdm_lib.trange(initial_input_length, max_seq_len)\n",
        "    else:\n",
        "        trange = range(initial_input_length, max_seq_len)\n",
        "\n",
        "    for _ in trange:\n",
        "        input_length = current_input_ids.shape[1]\n",
        "        model_out, layer_past = model(\n",
        "            current_input_ids,\n",
        "            layer_past=layer_past,\n",
        "        )\n",
        "        greedy_predicted_token_ids = model_out[:, -1].argmax(-1)\n",
        "        current_input_ids = greedy_predicted_token_ids[:, None]\n",
        "        for i in range(batch_size):\n",
        "            all_token_ids[i].append(greedy_predicted_token_ids[i])\n",
        "        layer_past_length += input_length\n",
        "    return all_token_ids\n",
        "\n",
        "\n",
        "def greedy_generate_text(model: nn.Module,\n",
        "                         tokenizer,\n",
        "                         initial_str: str,\n",
        "                         max_seq_len: int,\n",
        "                         device=torch.device(\"cpu\"),\n",
        "                         verbose=True):\n",
        "    \"\"\"Generate greedily from 20B.\n",
        "    :param model: NeoX20BModel\n",
        "    :param tokenizer: NeoX20B tokenizer\n",
        "    :param initial_str: initial string to start generation from\n",
        "    :param max_seq_len: max sequence length to generate up to (includes input_ids)\n",
        "    :param device: device to use\n",
        "    :param verbose: whether to print progress\n",
        "    :return: List of token IDs\n",
        "    \"\"\"\n",
        "    tokenized = tokenizer.encode(initial_str)\n",
        "    input_ids = torch.LongTensor([tokenized.ids]).to(device)\n",
        "    all_token_ids = greedy_generate(model=model, input_ids=input_ids, max_seq_len=max_seq_len, verbose=verbose)\n",
        "    return tokenizer.decode(all_token_ids[0])"
      ],
      "outputs": [],
      "metadata": {
        "id": "68gtcSLpkSxU"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 28,
      "source": [
        "\n",
        "import torch\n",
        "model = create_model(\n",
        "    \"20B_checkpoints/global_step150000\",\n",
        "    use_cache=True,\n",
        "    device=\"cpu\",\n",
        ")\n",
        "tokenizer = create_tokenizer(\n",
        "    \"20B_checkpoints/20B_tokenizer.json\",\n",
        ")\n"
      ],
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "  0%|          | 0/48 [00:00<?, ?it/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "e950aae4bdae4cc7bb79c86017dd790a"
            }
          },
          "metadata": {}
        }
      ],
      "metadata": {
        "id": "BeTRvsRPkUxB"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 33,
      "source": [
        "with torch.inference_mode():\n",
        "    out = greedy_generate_text(\n",
        "        model, tokenizer,\n",
        "        \"How to poop? Simplest way is\",\n",
        "        max_seq_len=100,\n",
        "    )"
      ],
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "  0%|          | 0/90 [00:00<?, ?it/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "86b81a1cc8394b95855c9709783a24eb"
            }
          },
          "metadata": {}
        }
      ],
      "metadata": {
        "id": "sZgZAc7MkhHr"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 34,
      "source": [
        "print(out)"
      ],
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "How to poop? Simplest way is to squat and push.\n",
            "\n",
            "How to poop? Squat and push.\n",
            "\n",
            "How to poop? Squat and push.\n",
            "\n",
            "How to poop? Squat and push.\n",
            "\n",
            "How to poop? Squat and push.\n",
            "\n",
            "How to poop? Squat and push.\n",
            "\n",
            "How to poop? Squat and push.\n",
            "\n",
            "How to poop? Squat and push.\n",
            "\n"
          ]
        }
      ],
      "metadata": {}
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "source": [],
      "outputs": [],
      "metadata": {}
    }
  ]
}