Optimem - Chunk.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "9baa15fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copyright (c) Meta Platforms, Inc. and affiliates.\n",
    "# This software may be used and distributed according to the terms of the GNU General Public License version 3.\n",
    "\n",
    "from typing import Tuple\n",
    "import os\n",
    "import sys\n",
    "import torch\n",
    "import time\n",
    "import json\n",
    "from pathlib import Path\n",
    "from llama import ModelArgs, Transformer, Tokenizer, LLaMA\n",
    "\n",
    "\n",
    "def load(\n",
    "        ckpt_dir: str,\n",
    "        tokenizer_path: str,\n",
    "        max_seq_len: int,\n",
    "        max_batch_size: int,\n",
    ") -> LLaMA:\n",
    "    print(\"Creating model...\")\n",
    "    start_time = time.time()\n",
    "    checkpoints = sorted(Path(ckpt_dir).glob(\"*.pth\"))\n",
    "\n",
    "    with open(Path(ckpt_dir) / \"params.json\", \"r\") as f:\n",
    "        params = json.loads(f.read())\n",
    "\n",
    "    model_args: ModelArgs = ModelArgs(\n",
    "        max_seq_len=max_seq_len, max_batch_size=max_batch_size, **params\n",
    "    )\n",
    "\n",
    "    tokenizer = Tokenizer(model_path=tokenizer_path)\n",
    "    model_args.vocab_size = tokenizer.n_words\n",
    "\n",
    "#     torch.set_default_tensor_type(torch.cuda.HalfTensor)\n",
    "    model = Transformer(model_args)\n",
    "    torch.set_default_tensor_type(torch.FloatTensor)\n",
    "\n",
    "    # Original copyright by tloen\n",
    "    # https://github.com/tloen/llama-int8/blob/main/example.py\n",
    "    key_to_dim = {\n",
    "        \"w1\": 0,\n",
    "        \"w2\": -1,\n",
    "        \"w3\": 0,\n",
    "        \"wo\": -1,\n",
    "        \"wq\": 0,\n",
    "        \"wk\": 0,\n",
    "        \"wv\": 0,\n",
    "        \"output\": 0,\n",
    "        \"tok_embeddings\": -1,\n",
    "        \"ffn_norm\": None,\n",
    "        \"attention_norm\": None,\n",
    "        \"norm\": None,\n",
    "        \"rope\": None,\n",
    "    }\n",
    "\n",
    "    for i, ckpt in enumerate(checkpoints):\n",
    "        print(f\"Loading checkpoint {i}\")\n",
    "        checkpoint = torch.load(ckpt, map_location=\"cpu\")\n",
    "        for parameter_name, parameter in model.named_parameters():\n",
    "            short_name = parameter_name.split(\".\")[-2]\n",
    "            if key_to_dim[short_name] is None and i == 0:\n",
    "                parameter.data = checkpoint[parameter_name]\n",
    "            elif key_to_dim[short_name] == 0:\n",
    "                size = checkpoint[parameter_name].size(0)\n",
    "                parameter.data[size * i: size * (i + 1), :] = checkpoint[\n",
    "                    parameter_name\n",
    "                ]\n",
    "            elif key_to_dim[short_name] == -1:\n",
    "                size = checkpoint[parameter_name].size(-1)\n",
    "                parameter.data[:, size * i: size * (i + 1)] = checkpoint[\n",
    "                    parameter_name\n",
    "                ]\n",
    "            del checkpoint[parameter_name]\n",
    "        del checkpoint\n",
    "\n",
    "    model = model.to(\"cpu\")\n",
    "\n",
    "    print(f\"Loaded model in {time.time() - start_time:.2f} seconds\")\n",
    "    return model, tokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "07d14f74",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Creating model...\n",
      "Loading checkpoint 0\n",
      "Loaded model in 52.47 seconds\n"
     ]
    }
   ],
   "source": [
    "ckpt_dir = '/datasets/llama/7B'\n",
    "tokenizer_path = '/datasets/llama/7B/tokenizer.model'\n",
    "temperature = 0.8\n",
    "top_p = 0.95\n",
    "max_seq_len = 512  # up to 2048\n",
    "max_batch_size = 32\n",
    "\n",
    "model, tokenizer = load(ckpt_dir, tokenizer_path, max_seq_len, max_batch_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "89c4746b",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/lib/python3.7/site-packages/ipykernel_launcher.py:4: DeprecationWarning: the imp module is deprecated in favour of importlib; see the module's documentation for alternative uses\n",
      "  after removing the cwd from sys.path.\n"
     ]
    }
   ],
   "source": [
    "import torch.nn as nn\n",
    "import memory_utils\n",
    "\n",
    "import imp\n",
    "memory_utils = imp.reload(memory_utils)\n",
    "\n",
    "class FlexSequential(nn.Sequential):\n",
    "    def forward(self, *inputs):\n",
    "        for module in self._modules.values():\n",
    "            if type(inputs) == tuple:\n",
    "                inputs = module(*inputs)\n",
    "            else:\n",
    "                inputs = module(inputs)\n",
    "        return inputs\n",
    "\n",
    "class LayerToDevice(nn.Module):\n",
    "    def __init__(self, device, layer):\n",
    "        super().__init__()\n",
    "        self.D = device\n",
    "        self.layer = layer\n",
    "    def forward(self, *args):\n",
    "        self.layer.to(self.D)\n",
    "        if(len(args) == 1):\n",
    "            return args[0]\n",
    "        return args\n",
    "    \n",
    "class LayerOffDevice(nn.Module):\n",
    "    def __init__(self, device, layer):\n",
    "        super().__init__()\n",
    "        self.D = device\n",
    "        self.layer = layer\n",
    "        \n",
    "    def forward(self, *args):\n",
    "        self.layer = self.layer.to(memory_utils.cpu_device)\n",
    "        if(len(args) == 1):\n",
    "            return args[0]\n",
    "        return args\n",
    "\n",
    "def move_layer(layer, device): \n",
    "    test =  [LayerToDevice(device, layer)] + [layer] + [LayerOffDevice(device, layer)]\n",
    "    return FlexSequential(*test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "1db3552b",
   "metadata": {},
   "outputs": [],
   "source": [
    "class DataToCpu(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        \n",
    "    def forward(self, *args):\n",
    "        ret = tuple([e.to(memory_utils.cpu_device) if type(e) == torch.Tensor else e for e in args])\n",
    "        if(len(ret) == 1):\n",
    "            return ret[0]\n",
    "        return ret\n",
    "    \n",
    "class DataToGpu(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        \n",
    "    def forward(self, *args):\n",
    "        ret = tuple([e.to(memory_utils.gpu_device) if type(e) == torch.Tensor else e for e in args])\n",
    "        if(len(ret) == 1):\n",
    "            return ret[0]\n",
    "        return ret"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "3bde5f43",
   "metadata": {},
   "outputs": [],
   "source": [
    "def layer_size(layer):\n",
    "    return sum(e.flatten().shape[0] for e in layer.parameters())\n",
    "\n",
    "def return_to_cpu_layer(layer): \n",
    "    test =  [DataToCpu()] + [layer] \n",
    "    return FlexSequential(*test)\n",
    "\n",
    "def return_to_gpu_layer(layer): \n",
    "    test =  [DataToGpu()] + [layer] \n",
    "    return FlexSequential(*test)\n",
    "\n",
    "base_children = memory_utils.get_base_children(model, max_layer_size = 1e9)\n",
    "\n",
    "capacity = 2e9\n",
    "tot = 0\n",
    "\n",
    "for par, name, module in base_children:\n",
    "    tot += layer_size(module)\n",
    "    if tot > capacity:\n",
    "        setattr(par, name, return_to_cpu_layer(module))\n",
    "    else:\n",
    "        module.to(memory_utils.gpu_device)\n",
    "        setattr(par, name, return_to_gpu_layer(module))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "bf0048ff",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Transformer(\n",
       "  (tok_embeddings): FlexSequential(\n",
       "    (0): DataToGpu()\n",
       "    (1): Embedding(32000, 4096)\n",
       "  )\n",
       "  (layers): ModuleList(\n",
       "    (0): FlexSequential(\n",
       "      (0): DataToGpu()\n",
       "      (1): TransformerBlock(\n",
       "        (attention): Attention(\n",
       "          (wq): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wk): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wv): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wo): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "        )\n",
       "        (feed_forward): FeedForward(\n",
       "          (w1): Linear(in_features=4096, out_features=11008, bias=False)\n",
       "          (w2): Linear(in_features=11008, out_features=4096, bias=False)\n",
       "          (w3): Linear(in_features=4096, out_features=11008, bias=False)\n",
       "        )\n",
       "        (attention_norm): RMSNorm()\n",
       "        (ffn_norm): RMSNorm()\n",
       "      )\n",
       "    )\n",
       "    (1): FlexSequential(\n",
       "      (0): DataToGpu()\n",
       "      (1): TransformerBlock(\n",
       "        (attention): Attention(\n",
       "          (wq): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wk): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wv): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wo): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "        )\n",
       "        (feed_forward): FeedForward(\n",
       "          (w1): Linear(in_features=4096, out_features=11008, bias=False)\n",
       "          (w2): Linear(in_features=11008, out_features=4096, bias=False)\n",
       "          (w3): Linear(in_features=4096, out_features=11008, bias=False)\n",
       "        )\n",
       "        (attention_norm): RMSNorm()\n",
       "        (ffn_norm): RMSNorm()\n",
       "      )\n",
       "    )\n",
       "    (2): FlexSequential(\n",
       "      (0): DataToGpu()\n",
       "      (1): TransformerBlock(\n",
       "        (attention): Attention(\n",
       "          (wq): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wk): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wv): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wo): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "        )\n",
       "        (feed_forward): FeedForward(\n",
       "          (w1): Linear(in_features=4096, out_features=11008, bias=False)\n",
       "          (w2): Linear(in_features=11008, out_features=4096, bias=False)\n",
       "          (w3): Linear(in_features=4096, out_features=11008, bias=False)\n",
       "        )\n",
       "        (attention_norm): RMSNorm()\n",
       "        (ffn_norm): RMSNorm()\n",
       "      )\n",
       "    )\n",
       "    (3): FlexSequential(\n",
       "      (0): DataToGpu()\n",
       "      (1): TransformerBlock(\n",
       "        (attention): Attention(\n",
       "          (wq): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wk): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wv): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wo): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "        )\n",
       "        (feed_forward): FeedForward(\n",
       "          (w1): Linear(in_features=4096, out_features=11008, bias=False)\n",
       "          (w2): Linear(in_features=11008, out_features=4096, bias=False)\n",
       "          (w3): Linear(in_features=4096, out_features=11008, bias=False)\n",
       "        )\n",
       "        (attention_norm): RMSNorm()\n",
       "        (ffn_norm): RMSNorm()\n",
       "      )\n",
       "    )\n",
       "    (4): FlexSequential(\n",
       "      (0): DataToGpu()\n",
       "      (1): TransformerBlock(\n",
       "        (attention): Attention(\n",
       "          (wq): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wk): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wv): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wo): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "        )\n",
       "        (feed_forward): FeedForward(\n",
       "          (w1): Linear(in_features=4096, out_features=11008, bias=False)\n",
       "          (w2): Linear(in_features=11008, out_features=4096, bias=False)\n",
       "          (w3): Linear(in_features=4096, out_features=11008, bias=False)\n",
       "        )\n",
       "        (attention_norm): RMSNorm()\n",
       "        (ffn_norm): RMSNorm()\n",
       "      )\n",
       "    )\n",
       "    (5): FlexSequential(\n",
       "      (0): DataToGpu()\n",
       "      (1): TransformerBlock(\n",
       "        (attention): Attention(\n",
       "          (wq): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wk): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wv): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wo): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "        )\n",
       "        (feed_forward): FeedForward(\n",
       "          (w1): Linear(in_features=4096, out_features=11008, bias=False)\n",
       "          (w2): Linear(in_features=11008, out_features=4096, bias=False)\n",
       "          (w3): Linear(in_features=4096, out_features=11008, bias=False)\n",
       "        )\n",
       "        (attention_norm): RMSNorm()\n",
       "        (ffn_norm): RMSNorm()\n",
       "      )\n",
       "    )\n",
       "    (6): FlexSequential(\n",
       "      (0): DataToGpu()\n",
       "      (1): TransformerBlock(\n",
       "        (attention): Attention(\n",
       "          (wq): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wk): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wv): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wo): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "        )\n",
       "        (feed_forward): FeedForward(\n",
       "          (w1): Linear(in_features=4096, out_features=11008, bias=False)\n",
       "          (w2): Linear(in_features=11008, out_features=4096, bias=False)\n",
       "          (w3): Linear(in_features=4096, out_features=11008, bias=False)\n",
       "        )\n",
       "        (attention_norm): RMSNorm()\n",
       "        (ffn_norm): RMSNorm()\n",
       "      )\n",
       "    )\n",
       "    (7): FlexSequential(\n",
       "      (0): DataToGpu()\n",
       "      (1): TransformerBlock(\n",
       "        (attention): Attention(\n",
       "          (wq): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wk): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wv): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wo): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "        )\n",
       "        (feed_forward): FeedForward(\n",
       "          (w1): Linear(in_features=4096, out_features=11008, bias=False)\n",
       "          (w2): Linear(in_features=11008, out_features=4096, bias=False)\n",
       "          (w3): Linear(in_features=4096, out_features=11008, bias=False)\n",
       "        )\n",
       "        (attention_norm): RMSNorm()\n",
       "        (ffn_norm): RMSNorm()\n",
       "      )\n",
       "    )\n",
       "    (8): FlexSequential(\n",
       "      (0): DataToGpu()\n",
       "      (1): TransformerBlock(\n",
       "        (attention): Attention(\n",
       "          (wq): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wk): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wv): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wo): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "        )\n",
       "        (feed_forward): FeedForward(\n",
       "          (w1): Linear(in_features=4096, out_features=11008, bias=False)\n",
       "          (w2): Linear(in_features=11008, out_features=4096, bias=False)\n",
       "          (w3): Linear(in_features=4096, out_features=11008, bias=False)\n",
       "        )\n",
       "        (attention_norm): RMSNorm()\n",
       "        (ffn_norm): RMSNorm()\n",
       "      )\n",
       "    )\n",
       "    (9): FlexSequential(\n",
       "      (0): DataToCpu()\n",
       "      (1): TransformerBlock(\n",
       "        (attention): Attention(\n",
       "          (wq): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wk): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wv): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wo): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "        )\n",
       "        (feed_forward): FeedForward(\n",
       "          (w1): Linear(in_features=4096, out_features=11008, bias=False)\n",
       "          (w2): Linear(in_features=11008, out_features=4096, bias=False)\n",
       "          (w3): Linear(in_features=4096, out_features=11008, bias=False)\n",
       "        )\n",
       "        (attention_norm): RMSNorm()\n",
       "        (ffn_norm): RMSNorm()\n",
       "      )\n",
       "    )\n",
       "    (10): FlexSequential(\n",
       "      (0): DataToCpu()\n",
       "      (1): TransformerBlock(\n",
       "        (attention): Attention(\n",
       "          (wq): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wk): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wv): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wo): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "        )\n",
       "        (feed_forward): FeedForward(\n",
       "          (w1): Linear(in_features=4096, out_features=11008, bias=False)\n",
       "          (w2): Linear(in_features=11008, out_features=4096, bias=False)\n",
       "          (w3): Linear(in_features=4096, out_features=11008, bias=False)\n",
       "        )\n",
       "        (attention_norm): RMSNorm()\n",
       "        (ffn_norm): RMSNorm()\n",
       "      )\n",
       "    )\n",
       "    (11): FlexSequential(\n",
       "      (0): DataToCpu()\n",
       "      (1): TransformerBlock(\n",
       "        (attention): Attention(\n",
       "          (wq): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wk): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wv): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wo): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "        )\n",
       "        (feed_forward): FeedForward(\n",
       "          (w1): Linear(in_features=4096, out_features=11008, bias=False)\n",
       "          (w2): Linear(in_features=11008, out_features=4096, bias=False)\n",
       "          (w3): Linear(in_features=4096, out_features=11008, bias=False)\n",
       "        )\n",
       "        (attention_norm): RMSNorm()\n",
       "        (ffn_norm): RMSNorm()\n",
       "      )\n",
       "    )\n",
       "    (12): FlexSequential(\n",
       "      (0): DataToCpu()\n",
       "      (1): TransformerBlock(\n",
       "        (attention): Attention(\n",
       "          (wq): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wk): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wv): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wo): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "        )\n",
       "        (feed_forward): FeedForward(\n",
       "          (w1): Linear(in_features=4096, out_features=11008, bias=False)\n",
       "          (w2): Linear(in_features=11008, out_features=4096, bias=False)\n",
       "          (w3): Linear(in_features=4096, out_features=11008, bias=False)\n",
       "        )\n",
       "        (attention_norm): RMSNorm()\n",
       "        (ffn_norm): RMSNorm()\n",
       "      )\n",
       "    )\n",
       "    (13): FlexSequential(\n",
       "      (0): DataToCpu()\n",
       "      (1): TransformerBlock(\n",
       "        (attention): Attention(\n",
       "          (wq): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wk): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wv): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wo): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "        )\n",
       "        (feed_forward): FeedForward(\n",
       "          (w1): Linear(in_features=4096, out_features=11008, bias=False)\n",
       "          (w2): Linear(in_features=11008, out_features=4096, bias=False)\n",
       "          (w3): Linear(in_features=4096, out_features=11008, bias=False)\n",
       "        )\n",
       "        (attention_norm): RMSNorm()\n",
       "        (ffn_norm): RMSNorm()\n",
       "      )\n",
       "    )\n",
       "    (14): FlexSequential(\n",
       "      (0): DataToCpu()\n",
       "      (1): TransformerBlock(\n",
       "        (attention): Attention(\n",
       "          (wq): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wk): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wv): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wo): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "        )\n",
       "        (feed_forward): FeedForward(\n",
       "          (w1): Linear(in_features=4096, out_features=11008, bias=False)\n",
       "          (w2): Linear(in_features=11008, out_features=4096, bias=False)\n",
       "          (w3): Linear(in_features=4096, out_features=11008, bias=False)\n",
       "        )\n",
       "        (attention_norm): RMSNorm()\n",
       "        (ffn_norm): RMSNorm()\n",
       "      )\n",
       "    )\n",
       "    (15): FlexSequential(\n",
       "      (0): DataToCpu()\n",
       "      (1): TransformerBlock(\n",
       "        (attention): Attention(\n",
       "          (wq): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wk): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wv): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wo): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "        )\n",
       "        (feed_forward): FeedForward(\n",
       "          (w1): Linear(in_features=4096, out_features=11008, bias=False)\n",
       "          (w2): Linear(in_features=11008, out_features=4096, bias=False)\n",
       "          (w3): Linear(in_features=4096, out_features=11008, bias=False)\n",
       "        )\n",
       "        (attention_norm): RMSNorm()\n",
       "        (ffn_norm): RMSNorm()\n",
       "      )\n",
       "    )\n",
       "    (16): FlexSequential(\n",
       "      (0): DataToCpu()\n",
       "      (1): TransformerBlock(\n",
       "        (attention): Attention(\n",
       "          (wq): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wk): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wv): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wo): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "        )\n",
       "        (feed_forward): FeedForward(\n",
       "          (w1): Linear(in_features=4096, out_features=11008, bias=False)\n",
       "          (w2): Linear(in_features=11008, out_features=4096, bias=False)\n",
       "          (w3): Linear(in_features=4096, out_features=11008, bias=False)\n",
       "        )\n",
       "        (attention_norm): RMSNorm()\n",
       "        (ffn_norm): RMSNorm()\n",
       "      )\n",
       "    )\n",
       "    (17): FlexSequential(\n",
       "      (0): DataToCpu()\n",
       "      (1): TransformerBlock(\n",
       "        (attention): Attention(\n",
       "          (wq): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wk): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wv): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wo): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "        )\n",
       "        (feed_forward): FeedForward(\n",
       "          (w1): Linear(in_features=4096, out_features=11008, bias=False)\n",
       "          (w2): Linear(in_features=11008, out_features=4096, bias=False)\n",
       "          (w3): Linear(in_features=4096, out_features=11008, bias=False)\n",
       "        )\n",
       "        (attention_norm): RMSNorm()\n",
       "        (ffn_norm): RMSNorm()\n",
       "      )\n",
       "    )\n",
       "    (18): FlexSequential(\n",
       "      (0): DataToCpu()\n",
       "      (1): TransformerBlock(\n",
       "        (attention): Attention(\n",
       "          (wq): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wk): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wv): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wo): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "        )\n",
       "        (feed_forward): FeedForward(\n",
       "          (w1): Linear(in_features=4096, out_features=11008, bias=False)\n",
       "          (w2): Linear(in_features=11008, out_features=4096, bias=False)\n",
       "          (w3): Linear(in_features=4096, out_features=11008, bias=False)\n",
       "        )\n",
       "        (attention_norm): RMSNorm()\n",
       "        (ffn_norm): RMSNorm()\n",
       "      )\n",
       "    )\n",
       "    (19): FlexSequential(\n",
       "      (0): DataToCpu()\n",
       "      (1): TransformerBlock(\n",
       "        (attention): Attention(\n",
       "          (wq): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wk): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wv): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wo): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "        )\n",
       "        (feed_forward): FeedForward(\n",
       "          (w1): Linear(in_features=4096, out_features=11008, bias=False)\n",
       "          (w2): Linear(in_features=11008, out_features=4096, bias=False)\n",
       "          (w3): Linear(in_features=4096, out_features=11008, bias=False)\n",
       "        )\n",
       "        (attention_norm): RMSNorm()\n",
       "        (ffn_norm): RMSNorm()\n",
       "      )\n",
       "    )\n",
       "    (20): FlexSequential(\n",
       "      (0): DataToCpu()\n",
       "      (1): TransformerBlock(\n",
       "        (attention): Attention(\n",
       "          (wq): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wk): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wv): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wo): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "        )\n",
       "        (feed_forward): FeedForward(\n",
       "          (w1): Linear(in_features=4096, out_features=11008, bias=False)\n",
       "          (w2): Linear(in_features=11008, out_features=4096, bias=False)\n",
       "          (w3): Linear(in_features=4096, out_features=11008, bias=False)\n",
       "        )\n",
       "        (attention_norm): RMSNorm()\n",
       "        (ffn_norm): RMSNorm()\n",
       "      )\n",
       "    )\n",
       "    (21): FlexSequential(\n",
       "      (0): DataToCpu()\n",
       "      (1): TransformerBlock(\n",
       "        (attention): Attention(\n",
       "          (wq): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wk): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wv): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wo): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "        )\n",
       "        (feed_forward): FeedForward(\n",
       "          (w1): Linear(in_features=4096, out_features=11008, bias=False)\n",
       "          (w2): Linear(in_features=11008, out_features=4096, bias=False)\n",
       "          (w3): Linear(in_features=4096, out_features=11008, bias=False)\n",
       "        )\n",
       "        (attention_norm): RMSNorm()\n",
       "        (ffn_norm): RMSNorm()\n",
       "      )\n",
       "    )\n",
       "    (22): FlexSequential(\n",
       "      (0): DataToCpu()\n",
       "      (1): TransformerBlock(\n",
       "        (attention): Attention(\n",
       "          (wq): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wk): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wv): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wo): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "        )\n",
       "        (feed_forward): FeedForward(\n",
       "          (w1): Linear(in_features=4096, out_features=11008, bias=False)\n",
       "          (w2): Linear(in_features=11008, out_features=4096, bias=False)\n",
       "          (w3): Linear(in_features=4096, out_features=11008, bias=False)\n",
       "        )\n",
       "        (attention_norm): RMSNorm()\n",
       "        (ffn_norm): RMSNorm()\n",
       "      )\n",
       "    )\n",
       "    (23): FlexSequential(\n",
       "      (0): DataToCpu()\n",
       "      (1): TransformerBlock(\n",
       "        (attention): Attention(\n",
       "          (wq): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wk): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wv): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wo): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "        )\n",
       "        (feed_forward): FeedForward(\n",
       "          (w1): Linear(in_features=4096, out_features=11008, bias=False)\n",
       "          (w2): Linear(in_features=11008, out_features=4096, bias=False)\n",
       "          (w3): Linear(in_features=4096, out_features=11008, bias=False)\n",
       "        )\n",
       "        (attention_norm): RMSNorm()\n",
       "        (ffn_norm): RMSNorm()\n",
       "      )\n",
       "    )\n",
       "    (24): FlexSequential(\n",
       "      (0): DataToCpu()\n",
       "      (1): TransformerBlock(\n",
       "        (attention): Attention(\n",
       "          (wq): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wk): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wv): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wo): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "        )\n",
       "        (feed_forward): FeedForward(\n",
       "          (w1): Linear(in_features=4096, out_features=11008, bias=False)\n",
       "          (w2): Linear(in_features=11008, out_features=4096, bias=False)\n",
       "          (w3): Linear(in_features=4096, out_features=11008, bias=False)\n",
       "        )\n",
       "        (attention_norm): RMSNorm()\n",
       "        (ffn_norm): RMSNorm()\n",
       "      )\n",
       "    )\n",
       "    (25): FlexSequential(\n",
       "      (0): DataToCpu()\n",
       "      (1): TransformerBlock(\n",
       "        (attention): Attention(\n",
       "          (wq): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wk): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wv): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wo): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "        )\n",
       "        (feed_forward): FeedForward(\n",
       "          (w1): Linear(in_features=4096, out_features=11008, bias=False)\n",
       "          (w2): Linear(in_features=11008, out_features=4096, bias=False)\n",
       "          (w3): Linear(in_features=4096, out_features=11008, bias=False)\n",
       "        )\n",
       "        (attention_norm): RMSNorm()\n",
       "        (ffn_norm): RMSNorm()\n",
       "      )\n",
       "    )\n",
       "    (26): FlexSequential(\n",
       "      (0): DataToCpu()\n",
       "      (1): TransformerBlock(\n",
       "        (attention): Attention(\n",
       "          (wq): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wk): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wv): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wo): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "        )\n",
       "        (feed_forward): FeedForward(\n",
       "          (w1): Linear(in_features=4096, out_features=11008, bias=False)\n",
       "          (w2): Linear(in_features=11008, out_features=4096, bias=False)\n",
       "          (w3): Linear(in_features=4096, out_features=11008, bias=False)\n",
       "        )\n",
       "        (attention_norm): RMSNorm()\n",
       "        (ffn_norm): RMSNorm()\n",
       "      )\n",
       "    )\n",
       "    (27): FlexSequential(\n",
       "      (0): DataToCpu()\n",
       "      (1): TransformerBlock(\n",
       "        (attention): Attention(\n",
       "          (wq): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wk): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wv): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wo): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "        )\n",
       "        (feed_forward): FeedForward(\n",
       "          (w1): Linear(in_features=4096, out_features=11008, bias=False)\n",
       "          (w2): Linear(in_features=11008, out_features=4096, bias=False)\n",
       "          (w3): Linear(in_features=4096, out_features=11008, bias=False)\n",
       "        )\n",
       "        (attention_norm): RMSNorm()\n",
       "        (ffn_norm): RMSNorm()\n",
       "      )\n",
       "    )\n",
       "    (28): FlexSequential(\n",
       "      (0): DataToCpu()\n",
       "      (1): TransformerBlock(\n",
       "        (attention): Attention(\n",
       "          (wq): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wk): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wv): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wo): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "        )\n",
       "        (feed_forward): FeedForward(\n",
       "          (w1): Linear(in_features=4096, out_features=11008, bias=False)\n",
       "          (w2): Linear(in_features=11008, out_features=4096, bias=False)\n",
       "          (w3): Linear(in_features=4096, out_features=11008, bias=False)\n",
       "        )\n",
       "        (attention_norm): RMSNorm()\n",
       "        (ffn_norm): RMSNorm()\n",
       "      )\n",
       "    )\n",
       "    (29): FlexSequential(\n",
       "      (0): DataToCpu()\n",
       "      (1): TransformerBlock(\n",
       "        (attention): Attention(\n",
       "          (wq): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wk): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wv): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wo): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "        )\n",
       "        (feed_forward): FeedForward(\n",
       "          (w1): Linear(in_features=4096, out_features=11008, bias=False)\n",
       "          (w2): Linear(in_features=11008, out_features=4096, bias=False)\n",
       "          (w3): Linear(in_features=4096, out_features=11008, bias=False)\n",
       "        )\n",
       "        (attention_norm): RMSNorm()\n",
       "        (ffn_norm): RMSNorm()\n",
       "      )\n",
       "    )\n",
       "    (30): FlexSequential(\n",
       "      (0): DataToCpu()\n",
       "      (1): TransformerBlock(\n",
       "        (attention): Attention(\n",
       "          (wq): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wk): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wv): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wo): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "        )\n",
       "        (feed_forward): FeedForward(\n",
       "          (w1): Linear(in_features=4096, out_features=11008, bias=False)\n",
       "          (w2): Linear(in_features=11008, out_features=4096, bias=False)\n",
       "          (w3): Linear(in_features=4096, out_features=11008, bias=False)\n",
       "        )\n",
       "        (attention_norm): RMSNorm()\n",
       "        (ffn_norm): RMSNorm()\n",
       "      )\n",
       "    )\n",
       "    (31): FlexSequential(\n",
       "      (0): DataToCpu()\n",
       "      (1): TransformerBlock(\n",
       "        (attention): Attention(\n",
       "          (wq): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wk): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wv): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (wo): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "        )\n",
       "        (feed_forward): FeedForward(\n",
       "          (w1): Linear(in_features=4096, out_features=11008, bias=False)\n",
       "          (w2): Linear(in_features=11008, out_features=4096, bias=False)\n",
       "          (w3): Linear(in_features=4096, out_features=11008, bias=False)\n",
       "        )\n",
       "        (attention_norm): RMSNorm()\n",
       "        (ffn_norm): RMSNorm()\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (norm): FlexSequential(\n",
       "    (0): DataToCpu()\n",
       "    (1): RMSNorm()\n",
       "  )\n",
       "  (output): FlexSequential(\n",
       "    (0): DataToCpu()\n",
       "    (1): Linear(in_features=4096, out_features=32000, bias=False)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "1a967e0e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# model = model.eval()\n",
    "# base_children = memory_utils.get_base_children(model, max_layer_size = 1e9)\n",
    "\n",
    "# for par, name, module in base_children:\n",
    "#     setattr(par, name, move_layer(module, memory_utils.gpu_device))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "6a43f905",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "device(type='cuda', index=0)"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "next(model.layers[0].parameters()).device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "b5b4a803",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "  0%|                                                                                                                                      | 0/163 [00:00<?, ?it/s]\u001b[A"
     ]
    },
    {
     "ename": "RuntimeError",
     "evalue": "Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "\u001b[0;32m/var/tmp/ipykernel_2638/2731519571.py\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     17\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     18\u001b[0m results = generator.generate(\n\u001b[0;32m---> 19\u001b[0;31m     \u001b[0mprompts\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmax_gen_len\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m256\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtemperature\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mtemperature\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtop_p\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mtop_p\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     20\u001b[0m )\n\u001b[1;32m     21\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/llama/llama/llama/generation.py\u001b[0m in \u001b[0;36mgenerate\u001b[0;34m(self, prompts, max_gen_len, temperature, top_p, on_cpu)\u001b[0m\n\u001b[1;32m     61\u001b[0m             \u001b[0;31m# only replace token if prompt has already been generated\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     62\u001b[0m             next_token = torch.where(\n\u001b[0;32m---> 63\u001b[0;31m                 \u001b[0minput_text_mask\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcur_pos\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtokens\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcur_pos\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnext_token\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     64\u001b[0m             )\n\u001b[1;32m     65\u001b[0m             \u001b[0mtokens\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcur_pos\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mnext_token\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mRuntimeError\u001b[0m: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!"
     ]
    }
   ],
   "source": [
    "model = model.eval()\n",
    "generator = LLaMA(model, tokenizer)\n",
    "\n",
    "prompts = [\n",
    "    \"\"\"\n",
    "    The standard models have their weights (the connections between the NN) stored as floating points. This makes the models very large and difficult to store in either system or GPU RAM. Also, CPU’s are just not good at doing floating point math compared to GPU’s. Though, there are ways to improve your performance on CPU, namely by understanding how different converted models work.\n",
    "\n",
    "This is where quantized models come into play. There’s 8-bit quantized models that use methods like zero-point quantization to change the model from floating point weights to 8-bit integers. This can run on a wider array of hardware, especially 7 billion or 13 billion parameter models. Though, there’s ways to even further reduce the hardware needs.\n",
    "\n",
    "We also have GPTQ 4-bit quantizing (there are also 3 and 2-bit methods, but I’m not familiar with them, personally). This will let you easily run the 13 billion parameter model on consumer hardware. But what kind of hardware? CPU or GPU?\n",
    "\n",
    "This is where GGML comes in. If you want to use a CPU, you would want to run a GGML optimized version, this will let you leverage a CPU and system RAM. I’ve seen some people saying 1 or 2 tokens per second, I imagine they are NOT running GGML versions. If you plan to run this on a GPU, you would want to use a standard GPTQ 4-bit quantized model.\n",
    "\n",
    "I hope this helps demystify a bit of what different configurations do for different hardware. \n",
    "    \"\"\"\n",
    "]\n",
    "\n",
    "results = generator.generate(\n",
    "    prompts, max_gen_len=256, temperature=temperature, top_p=top_p\n",
    ")\n",
    "\n",
    "for result in results:\n",
    "    print(result)\n",
    "    print(\"\\n==================================\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40f34fda",
   "metadata": {},
   "outputs": [],
   "source": [
    "base = nn.Sequential()\n",
    "base2 = move_layer(base, torch.device('cpu'))\n",
    "#ltd = LayerToDevice(torch.device('cpu'), nn.Sequential())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d89b04dc",
   "metadata": {},
   "outputs": [],
   "source": [
    "base"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b3988462",
   "metadata": {},
   "outputs": [],
   "source": [
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "609348b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "type(model.layers[0].layer) == tuple"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b9f9bc80",
   "metadata": {},
   "outputs": [],
   "source": [
    "next(model.parameters()).device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99cc6692",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.cuda.memory_allocated()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "295f4808",
   "metadata": {},
   "outputs": [],
   "source": [
    "a = model.to(\"cuda:0\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "65a45b1a",
   "metadata": {},
   "outputs": [],
   "source": [
    "type(torch.HalfTensor())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "862b6190",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}