GPT-J-6B-Low-Vram-UI.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "54d46ac1",
   "metadata": {},
   "source": [
    "# Import"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ca1af70",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "from transformers import GPTNeoForCausalLM, AutoTokenizer\n",
    "import torch\n",
    "import transformers\n",
    "import copy\n",
    "import gc\n",
    "import pickle\n",
    "import torch.cuda.comm\n",
    "import time\n",
    "\n",
    "from IPython.display import HTML, display\n",
    "import ipywidgets as widgets"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "72d1eac8",
   "metadata": {},
   "source": [
    "# Import model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ecf4653",
   "metadata": {},
   "outputs": [],
   "source": [
    "low_ram_method = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a8d9ceea",
   "metadata": {},
   "outputs": [],
   "source": [
    "if low_ram_method:\n",
    "    directory = \"parts/\"\n",
    "    with open(directory + \"emptymodel\" + \".pkl\", 'rb') as f:\n",
    "        model = pickle.load(f)\n",
    "    count = 0\n",
    "    for i in model.parameters():\n",
    "        count += 1\n",
    "        with open(directory + str(count) + \".pkl\", 'rb') as f:\n",
    "            current_layer = pickle.load(f)\n",
    "        i.data = current_layer.data\n",
    "print(count)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "553c439a",
   "metadata": {},
   "source": [
    "# Tokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "806812ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained(\"gpt2\")\n",
    "tokenizer.pad_token = tokenizer.eos_token\n",
    "tokenizer.padding_side = 'left'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9b7da8fe",
   "metadata": {},
   "source": [
    "# Set vram, shared/pinned memory and regular ram usage\n",
    "\n",
    "increase ram_blocks if cuda runs out of memory, decrease max_shared_ram_blocks if it raises an error in the code. Current settings are for 8gb vram and 8gb shared memory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "380c54f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "breakmodel = True\n",
    "ram_blocks = 23\n",
    "max_shared_ram_blocks = 18\n",
    "if ram_blocks > len(model.transformer.h):\n",
    "    ram_blocks = len(model.transformer.h)\n",
    "if ram_blocks < 2:\n",
    "    ram_blocks = 2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c19f48ec",
   "metadata": {},
   "source": [
    "# Modify the forward pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "525ee8aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import GPTNeoForCausalLM,GPTNeoModel\n",
    "from transformers.modeling_outputs import BaseModelOutputWithPast\n",
    "from transformers.models.gpt_neo.modeling_gpt_neo import GPTNeoAttentionMixin"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9b002a98",
   "metadata": {},
   "outputs": [],
   "source": [
    "def new_forward(\n",
    "        self,\n",
    "        input_ids=None,\n",
    "        past_key_values=None,\n",
    "        attention_mask=None,\n",
    "        token_type_ids=None,\n",
    "        position_ids=None,\n",
    "        head_mask=None,\n",
    "        inputs_embeds=None,\n",
    "        use_cache=None,\n",
    "        output_attentions=None,\n",
    "        output_hidden_states=None,\n",
    "        return_dict=None,\n",
    "    ):\n",
    "    global breakmodel\n",
    "\n",
    "    if breakmodel:\n",
    "        global ram_blocks\n",
    "        global max_shared_ram_blocks\n",
    "        \n",
    "        if not hasattr(self, 'extrastorage'):\n",
    "            import copy\n",
    "            setattr(self,\"extrastorage\",{})\n",
    "            torch.cuda.empty_cache()\n",
    "            \n",
    "            for i in range(ram_blocks,len(self.h)):\n",
    "                self.h[i].to(\"cuda\")\n",
    "            \n",
    "            for i in range(ram_blocks):\n",
    "                self.h[i].to(\"cpu\")\n",
    "                self.extrastorage[i] = copy.deepcopy(self.h[i])\n",
    "                smalltensor = torch.tensor(0).to(\"cuda\")\n",
    "                for param1 in self.h[i].parameters():\n",
    "                    param1.data = smalltensor\n",
    "                self.h[i].to(\"cuda\")\n",
    "\n",
    "            for i in range(len(self.h)):\n",
    "                for param in self.h[i].parameters():\n",
    "                    param.requires_grad = False\n",
    "                    param.data = param.data.detach()\n",
    "                    gc.collect()\n",
    "                    torch.cuda.empty_cache()\n",
    "                    \n",
    "            for i in range(ram_blocks):\n",
    "                for param in self.extrastorage[i].parameters():\n",
    "                    param.requires_grad = False\n",
    "                    if i < max_shared_ram_blocks:\n",
    "                        try:\n",
    "                            param.data = param.data.detach().pin_memory()\n",
    "                        except:\n",
    "                            raise ValueError('max_shared_ram_blocks is set too high, please set it to '+str(i))\n",
    "                    else:\n",
    "                        param.data = param.data.detach()\n",
    "                    gc.collect()\n",
    "                    torch.cuda.empty_cache()\n",
    "\n",
    "            for param1,param2 in zip(self.h[0].parameters(),self.extrastorage[0].parameters()):\n",
    "                param1.data = param2.data.to(\"cuda\", non_blocking=False).detach()\n",
    "\n",
    "            for param1,param2 in zip(self.h[ram_blocks-1].parameters(),self.extrastorage[ram_blocks-1].parameters()):\n",
    "                param1.data = param2.data.to(\"cuda\", non_blocking=False).detach()\n",
    "    #END MODEL BREAK EDITS\n",
    "    \n",
    "        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions\n",
    "        output_hidden_states = (\n",
    "            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states\n",
    "        )\n",
    "        use_cache = use_cache if use_cache is not None else self.config.use_cache\n",
    "        return_dict = return_dict if return_dict is not None else self.config.use_return_dict\n",
    "\n",
    "        if input_ids is not None and inputs_embeds is not None:\n",
    "            raise ValueError(\"You cannot specify both input_ids and inputs_embeds at the same time\")\n",
    "        elif input_ids is not None:\n",
    "            input_shape = input_ids.size()\n",
    "            input_ids = input_ids.view(-1, input_shape[-1])\n",
    "            batch_size = input_ids.shape[0]\n",
    "        elif inputs_embeds is not None:\n",
    "            input_shape = inputs_embeds.size()[:-1]\n",
    "            batch_size = inputs_embeds.shape[0]\n",
    "        else:\n",
    "            raise ValueError(\"You have to specify either input_ids or inputs_embeds\")\n",
    "\n",
    "        device = input_ids.device if input_ids is not None else inputs_embeds.device\n",
    "\n",
    "        if token_type_ids is not None:\n",
    "            token_type_ids = token_type_ids.view(-1, input_shape[-1])\n",
    "        if position_ids is not None:\n",
    "            position_ids = position_ids.view(-1, input_shape[-1])\n",
    "\n",
    "        if past_key_values is None:\n",
    "            past_length = 0\n",
    "            past_key_values = tuple([None] * len(self.h))\n",
    "        else:\n",
    "            past_length = past_key_values[0][0].size(-2)\n",
    "\n",
    "        device = input_ids.device if input_ids is not None else inputs_embeds.device\n",
    "        if position_ids is None:\n",
    "            position_ids = torch.arange(past_length, input_shape[-1] + past_length, dtype=torch.long, device=device)\n",
    "            position_ids = position_ids.unsqueeze(0).view(-1, input_shape[-1])\n",
    "\n",
    "        # Attention mask.\n",
    "        if attention_mask is not None:\n",
    "            assert batch_size > 0, \"batch_size has to be defined and > 0\"\n",
    "            global_attention_mask = attention_mask.view(batch_size, -1)\n",
    "            # We create a 3D attention mask from a 2D tensor mask.\n",
    "            # Sizes are [batch_size, 1, 1, to_seq_length]\n",
    "            # So we can broadcast to [batch_size, num_heads, from_seq_length, to_seq_length]\n",
    "            # this attention mask is more simple than the triangular masking of causal attention\n",
    "            # used in OpenAI GPT, we just need to prepare the broadcast dimension here.\n",
    "            global_attention_mask = global_attention_mask[:, None, None, :]\n",
    "\n",
    "            # Since global_attention_mask is 1.0 for positions we want to attend and 0.0 for\n",
    "            # masked positions, this operation will create a tensor which is 0.0 for\n",
    "            # positions we want to attend and -10000.0 for masked positions.\n",
    "            # Since we are adding it to the raw scores before the softmax, this is\n",
    "            # effectively the same as removing these entirely.\n",
    "            global_attention_mask = global_attention_mask.to(dtype=self.dtype)  # fp16 compatibility\n",
    "            global_attention_mask = (1.0 - global_attention_mask) * -10000.0\n",
    "        else:\n",
    "            global_attention_mask = None\n",
    "\n",
    "        # Local causal attention mask\n",
    "        batch_size, seq_length = input_shape\n",
    "        full_seq_length = seq_length + past_length\n",
    "\n",
    "        # Prepare head mask if needed\n",
    "        # 1.0 in head_mask indicate we keep the head\n",
    "        # attention_probs has shape bsz x num_heads x N x N\n",
    "        # head_mask has shape n_layer x batch x num_heads x N x N\n",
    "        head_mask = self.get_head_mask(head_mask, self.config.num_layers)\n",
    "\n",
    "        if inputs_embeds is None:\n",
    "            inputs_embeds = self.wte(input_ids)\n",
    "\n",
    "        if self.rotary:\n",
    "            hidden_states = inputs_embeds\n",
    "        else:\n",
    "            position_embeds = self.wpe(position_ids)\n",
    "            hidden_states = inputs_embeds + position_embeds\n",
    "\n",
    "        if token_type_ids is not None:\n",
    "            token_type_embeds = self.wte(token_type_ids)\n",
    "            hidden_states = hidden_states + token_type_embeds\n",
    "\n",
    "        hidden_states = self.drop(hidden_states)\n",
    "\n",
    "        output_shape = input_shape + (hidden_states.size(-1),)\n",
    "\n",
    "        presents = () if use_cache else None\n",
    "        all_self_attentions = () if output_attentions else None\n",
    "        all_hidden_states = () if output_hidden_states else None\n",
    "        \n",
    "        \n",
    "        if breakmodel :\n",
    "            copystream = torch.cuda.Stream(device=0,priority = -1)\n",
    "        \n",
    "        for i, (block, layer_past) in enumerate(zip(self.h, past_key_values)):\n",
    "            \n",
    "            if breakmodel :\n",
    "                if i in range(ram_blocks):\n",
    "                    index1 = (i+1)%ram_blocks\n",
    "                    for param1,param2 in zip(self.h[index1].parameters(),self.h[(i-1)%ram_blocks].parameters()):\n",
    "                        param1.data = param2.data\n",
    "                    for param1,param2 in zip(self.h[index1].parameters(),self.extrastorage[index1].parameters()):\n",
    "                        with torch.cuda.stream(copystream):\n",
    "                            torch.cuda.comm.broadcast(param2.data,out = [param1.data])\n",
    "            \n",
    "            \n",
    "            attn_type = self.config.attention_layers[i]\n",
    "            attn_mask = global_attention_mask\n",
    "\n",
    "            if output_hidden_states:\n",
    "                all_hidden_states = all_hidden_states + (hidden_states,)\n",
    "\n",
    "            if getattr(self.config, \"gradient_checkpointing\", False) and self.training:\n",
    "\n",
    "                if use_cache:\n",
    "                    logger.warning(\n",
    "                        \"`use_cache=True` is incompatible with `config.gradient_checkpointing=True`. Setting \"\n",
    "                        \"`use_cache=False`...\"\n",
    "                    )\n",
    "                    use_cache = False\n",
    "\n",
    "                def create_custom_forward(module):\n",
    "                    def custom_forward(*inputs):\n",
    "                        # None for past_key_value\n",
    "                        return module(*inputs, use_cache, output_attentions)\n",
    "\n",
    "                    return custom_forward\n",
    "\n",
    "                outputs = torch.utils.checkpoint.checkpoint(\n",
    "                    create_custom_forward(block),\n",
    "                    hidden_states,\n",
    "                    None,\n",
    "                    attn_mask,\n",
    "                    head_mask[i],\n",
    "                )\n",
    "            else:\n",
    "                outputs = block(\n",
    "                    hidden_states,\n",
    "                    layer_past=layer_past,\n",
    "                    attention_mask=attn_mask,\n",
    "                    head_mask=head_mask[i],\n",
    "                    use_cache=use_cache,\n",
    "                    output_attentions=output_attentions,\n",
    "                )\n",
    "\n",
    "            hidden_states = outputs[0]\n",
    "            if use_cache is True:\n",
    "                presents = presents + (outputs[1],)\n",
    "\n",
    "            if output_attentions:\n",
    "                all_self_attentions = all_self_attentions + (outputs[2 if use_cache else 1],)\n",
    "\n",
    "                \n",
    "            if breakmodel:\n",
    "                if i in range(ram_blocks):\n",
    "                    torch.cuda.synchronize()\n",
    "                    torch.cuda.empty_cache()   \n",
    "\n",
    "        if breakmodel:\n",
    "            del copystream\n",
    "\n",
    "        torch.cuda.empty_cache()   \n",
    "        \n",
    "        \n",
    "        hidden_states = self.ln_f(hidden_states)\n",
    "\n",
    "        hidden_states = hidden_states.view(*output_shape)\n",
    "        # Add last hidden state\n",
    "        if output_hidden_states:\n",
    "            all_hidden_states = all_hidden_states + (hidden_states,)\n",
    "\n",
    "        if not return_dict:\n",
    "            return tuple(v for v in [hidden_states, presents, all_hidden_states, all_self_attentions] if v is not None)\n",
    "\n",
    "        return BaseModelOutputWithPast(\n",
    "            last_hidden_state=hidden_states,\n",
    "            past_key_values=presents,\n",
    "            hidden_states=all_hidden_states,\n",
    "            attentions=all_self_attentions,\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2989d10a",
   "metadata": {},
   "outputs": [],
   "source": [
    "if breakmodel:\n",
    "    model.eval().half().to(\"cpu\")\n",
    "    \n",
    "    model.lm_head.to(\"cuda\")\n",
    "    model.transformer.wte.to(\"cuda\")\n",
    "    model.transformer.ln_f.to(\"cuda\")\n",
    "    \n",
    "    gc.collect()\n",
    "    print(GPTNeoModel.forward)\n",
    "    print(new_forward)\n",
    "    GPTNeoModel.forward = new_forward\n",
    "    print(GPTNeoModel.forward)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1fecb64e",
   "metadata": {},
   "source": [
    "# Generation Settings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aabc112a",
   "metadata": {},
   "outputs": [],
   "source": [
    "tail_free_sampling = 0.95 #@param {type:\"number\"}\n",
    "top_k = 80 #@param {type:\"number\"}\n",
    "top_p = 0.8 #@param {type:\"number\"}\n",
    "temperature =  0.7#@param {type:\"number\"}\n",
    "number_generated_tokens =  25#@param {type:\"integer\"}\n",
    "repetition_penalty = 1.1 #@param {type:\"number\"}\n",
    "repetition_penalty_range = 512 #@param {type:\"number\"}\n",
    "repetition_penalty_slope = 3.33 #@param {type:\"number\"}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3d248759",
   "metadata": {},
   "source": [
    "# Few runs for testing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2ae01d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "basic_prompt = \"test \" * 100000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "01c7d3a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs = tokenizer(basic_prompt, return_tensors=\"pt\",truncation=True,max_length=2000).to(\"cuda\")\n",
    "with torch.no_grad():\n",
    "    outputs = model(**inputs)\n",
    "del outputs\n",
    "torch.cuda.empty_cache()\n",
    "start_time = time.time()\n",
    "with torch.no_grad():\n",
    "    for i in range(1):\n",
    "        outputs = model(**inputs)\n",
    "print(time.time()  - start_time)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "282b304e",
   "metadata": {},
   "outputs": [],
   "source": [
    "del inputs,outputs\n",
    "torch.cuda.empty_cache()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5fbc3585",
   "metadata": {},
   "source": [
    "# Functions for ui"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f03dbe1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def more_text(inputtext):\n",
    "    #return \"Epictest\"\n",
    "    with torch.no_grad():\n",
    "        with torch.cuda.amp.autocast(enabled=True):\n",
    "            \n",
    "            context = 2000\n",
    "            overhead = 50\n",
    "\n",
    "            currpoint = len(inputtext)\n",
    "            inputs = tokenizer(inputtext[-currpoint:], return_tensors=\"pt\",truncation=True,max_length=context+overhead)\n",
    "            if inputs.input_ids[0].size()[0] == context+overhead:            \n",
    "\n",
    "                low = 0\n",
    "                high = len(inputtext)\n",
    "                currpoint = 0\n",
    "\n",
    "                #BINARY SEARCH FOR A POINT WHERE TOKENIZER RETURNS BETWEEN CONTEXT AND CONTEXT + OVERHEAD TOKENS\n",
    "                while low <= high:\n",
    "\n",
    "                    currpoint = (high + low) // 2\n",
    "\n",
    "                    # If x is greater, ignore left half\n",
    "                    inputs = tokenizer(inputtext[-currpoint:], return_tensors=\"pt\",truncation=True,max_length=context+overhead)\n",
    "\n",
    "                    if inputs.input_ids[0].size()[0] < context:\n",
    "                        low = currpoint + 1\n",
    "\n",
    "                    # If x is smaller, ignore right half\n",
    "                    elif inputs.input_ids[0].size()[0] == context + overhead :\n",
    "                        high = currpoint - 1\n",
    "\n",
    "                    # means x is present at mid\n",
    "                    else:\n",
    "                        break\n",
    "\n",
    "                ids = tokenizer(inputtext[-currpoint:], return_tensors=\"pt\",truncation=True,max_length=context+overhead).input_ids\n",
    "            else:\n",
    "                ids = tokenizer(inputtext[-currpoint:], return_tensors=\"pt\",truncation=True,max_length=context+overhead,padding = 'max_length').input_ids\n",
    "\n",
    "            ids = ids[:,-context:]        \n",
    "            n_ids = ids.shape[1]\n",
    "            if n_ids < 1:\n",
    "                n_ids = 1\n",
    "                ids = torch.tensor([[tokenizer.eos_token_id]])\n",
    "            max_length = n_ids + number_generated_tokens\n",
    "\n",
    "            gc.collect()\n",
    "\n",
    "            basic_output = model.generate(\n",
    "                ids.long().to(\"cuda\"),\n",
    "                do_sample=True,\n",
    "                num_beams=1,\n",
    "                min_length=max_length,\n",
    "                max_length=max_length,\n",
    "                temperature=temperature,\n",
    "                top_k = top_k,\n",
    "                top_p = top_p,\n",
    "                repetition_penalty = repetition_penalty,\n",
    "                repetition_penalty_range = repetition_penalty_range,\n",
    "                repetition_penalty_slope = repetition_penalty_slope,\n",
    "                use_cache=True,\n",
    "                pad_token_id=tokenizer.eos_token_id,\n",
    "                num_return_sequences = 1\n",
    "            ).long()\n",
    "\n",
    "            gc.collect()\n",
    "            torch.cuda.empty_cache()\n",
    "\n",
    "            return tokenizer.decode(basic_output[0][-number_generated_tokens:])\n",
    "\n",
    "    #print(time.time()  - start_time)\n",
    "    #print(number_generated_tokens)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22049455",
   "metadata": {},
   "outputs": [],
   "source": [
    "def change_text(_):\n",
    "    global box\n",
    "    temp_message = \"\\nCOMPUTE IN PROGRESS\"\n",
    "    box.value += temp_message\n",
    "    data = more_text(box.value[:-len(temp_message)])\n",
    "    box.value = box.value[:-len(temp_message)]\n",
    "    box.value += data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a8385566",
   "metadata": {},
   "source": [
    "# UI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "db0ea22b",
   "metadata": {},
   "outputs": [],
   "source": [
    "box = widgets.Textarea(value='', placeholder='', description='', disabled=False, rows=20, layout={\"width\": \"950px\"})\n",
    "button = widgets.Button(description='continue')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1718324",
   "metadata": {},
   "outputs": [],
   "source": [
    "display(box)\n",
    "display(button)\n",
    "button.on_click(change_text)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d2dbcb3",
   "metadata": {},
   "source": [
    "\n",
    "# Timing\n",
    "Use the last 2 cells to time the notebook <br>\n",
    "1. Put some text in the box<br>\n",
    "2. run the cell 1 to get the start time<br>\n",
    "3. Press the continue button<br>\n",
    "4. Run the cell 2.<br>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "49223d95",
   "metadata": {},
   "outputs": [],
   "source": [
    "# cell 1\n",
    "time.sleep(10)\n",
    "start_time = time.time()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8297c863",
   "metadata": {},
   "outputs": [],
   "source": [
    "# cell 2\n",
    "print(time.time() - start_time )"
   ]
  }
 ],
 "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}