GPT-J-6B-8bit-HeadlineGeneration.ipynb

{"nbformat":4,"nbformat_minor":0,"metadata":{"colab":{"provenance":[],"machine_shape":"hm","toc_visible":true,"authorship_tag":"ABX9TyNFTnT/mHoKW1iX/wN38Ayf"},"kernelspec":{"name":"python3","display_name":"Python 3"},"language_info":{"name":"python"},"accelerator":"GPU","gpuClass":"premium"},"cells":[{"cell_type":"markdown","source":["# Training GPT-J-6B with 8-bit weights for Headline Generation\n","This notebook contains the code to train a GPT Model using Transfer Learning on the task of Headline Generation from News Articles.\n","\n","\n","* GPT-J-6B: A 6 billion parameter, autoregressive text generation model by EleutherAI trained on [The Pile](https://pile.eleuther.ai/) using [Mesh Transformer JAX](https://github.com/kingoflolz/mesh-transformer-jax/) (Ben Wang and Aran Komatsuzaki).\n","* GPT-J-6B-8bit: A quantized GPT-J-6B with 8-bit weights for scalable and cost-efficient fine-tuning by Hivemind with [LoRA](https://arxiv.org/pdf/2106.09685.pdf) and [8-bit Adam](https://arxiv.org/abs/2110.02861).\n","* LoRA Adapter Implementation is created by [Denis Mazur](https://github.com/deniskamazur): [Notebook](https://colab.research.google.com/drive/1ft6wQU0BhqG5PRlwgaZJv2VukKKjU4Es)\n","\n"],"metadata":{"id":"MdSeD1gxv-G0"}},{"cell_type":"markdown","source":["## Installs & Imports\n","Freeze library versions for reproduction"],"metadata":{"id":"aoOr_j2owTUt"}},{"cell_type":"code","execution_count":null,"metadata":{"id":"nBspToWRSbyN"},"outputs":[],"source":["!pip install transformers==4.25.1\n","!pip install bitsandbytes==0.35.4\n","!pip install datasets==2.7.1\n","!pip install accelerate==0.15.0"]},{"cell_type":"code","source":["import transformers\n","import pandas as pd\n","\n","# Torch Imports\n","import torch\n","import torch.nn.functional as F\n","from torch import nn\n","from torch.cuda.amp import custom_fwd, custom_bwd\n","from torch.utils.data import DataLoader\n","\n","import accelerate\n","\n","from bitsandbytes.functional import quantize_blockwise, dequantize_blockwise\n","from bitsandbytes.optim import Adam8bit\n","\n","from datasets import load_dataset\n","\n","from google.colab import drive"],"metadata":{"id":"uKqoV5LDTflt"},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":["## Convert Model to 8 bits\n","We convert EleutherAI's GPT-J-6B model to 8 bits using facebook's [bitsandbytes](https://github.com/facebookresearch/bitsandbytes) library. This reduces the model's size from 20Gb down to 6Gb.\n","* large weight tensors are quantized using dynamic 8-bit quantization and de-quantized just-in-time for multiplication\n","* using gradient checkpoints to store one only activation per layer: using dramatically less memory at the cost of 30% slower training"],"metadata":{"id":"FV7F4q2YwYSo"}},{"cell_type":"code","source":["class FrozenBNBLinear(nn.Module):\n","    def __init__(self, weight, absmax, code, bias=None):\n","        assert isinstance(bias, nn.Parameter) or bias is None\n","        super().__init__()\n","        self.out_features, self.in_features = weight.shape\n","        self.register_buffer(\"weight\", weight.requires_grad_(False))\n","        self.register_buffer(\"absmax\", absmax.requires_grad_(False))\n","        self.register_buffer(\"code\", code.requires_grad_(False))\n","        self.adapter = None\n","        self.bias = bias\n"," \n","    def forward(self, input):\n","        output = DequantizeAndLinear.apply(input, self.weight, self.absmax, self.code, self.bias)\n","        if self.adapter:\n","            output += self.adapter(input)\n","        return output\n"," \n","    @classmethod\n","    def from_linear(cls, linear: nn.Linear) -> \"FrozenBNBLinear\":\n","        weights_int8, state = quantize_blockise_lowmemory(linear.weight)\n","        return cls(weights_int8, *state, linear.bias)\n"," \n","    def __repr__(self):\n","        return f\"{self.__class__.__name__}({self.in_features}, {self.out_features})\"\n"," \n"," \n","class DequantizeAndLinear(torch.autograd.Function): \n","    @staticmethod\n","    @custom_fwd\n","    def forward(ctx, input: torch.Tensor, weights_quantized: torch.ByteTensor,\n","                absmax: torch.FloatTensor, code: torch.FloatTensor, bias: torch.FloatTensor):\n","        weights_deq = dequantize_blockwise(weights_quantized, absmax=absmax, code=code)\n","        ctx.save_for_backward(input, weights_quantized, absmax, code)\n","        ctx._has_bias = bias is not None\n","        return F.linear(input, weights_deq, bias).clone()\n"," \n","    @staticmethod\n","    @custom_bwd\n","    def backward(ctx, grad_output: torch.Tensor):\n","        assert not ctx.needs_input_grad[1] and not ctx.needs_input_grad[2] and not ctx.needs_input_grad[3]\n","        input, weights_quantized, absmax, code = ctx.saved_tensors\n","        # grad_output: [*batch, out_features]\n","        weights_deq = dequantize_blockwise(weights_quantized, absmax=absmax, code=code)\n","        grad_input = grad_output @ weights_deq\n","        grad_bias = grad_output.flatten(0, -2).sum(dim=0) if ctx._has_bias else None\n","        return grad_input, None, None, None, grad_bias\n"," \n"," \n","class FrozenBNBEmbedding(nn.Module):\n","    def __init__(self, weight, absmax, code):\n","        super().__init__()\n","        self.num_embeddings, self.embedding_dim = weight.shape\n","        self.register_buffer(\"weight\", weight.requires_grad_(False))\n","        self.register_buffer(\"absmax\", absmax.requires_grad_(False))\n","        self.register_buffer(\"code\", code.requires_grad_(False))\n","        self.adapter = None\n"," \n","    def forward(self, input, **kwargs):\n","        with torch.no_grad():\n","            # note: both quantized weights and input indices are not differentiable\n","            weight_deq = dequantize_blockwise(self.weight, absmax=self.absmax, code=self.code)\n","            output = F.embedding(input, weight_deq, **kwargs)\n","        if self.adapter:\n","            output += self.adapter(input)\n","        return output \n"," \n","    @classmethod\n","    def from_embedding(cls, embedding: nn.Embedding) -> \"FrozenBNBEmbedding\":\n","        weights_int8, state = quantize_blockise_lowmemory(embedding.weight)\n","        return cls(weights_int8, *state)\n"," \n","    def __repr__(self):\n","        return f\"{self.__class__.__name__}({self.num_embeddings}, {self.embedding_dim})\"\n"," \n"," \n","def quantize_blockise_lowmemory(matrix: torch.Tensor, chunk_size: int = 2 ** 20):\n","    assert chunk_size % 4096 == 0\n","    code = None\n","    chunks = []\n","    absmaxes = []\n","    flat_tensor = matrix.view(-1)\n","    for i in range((matrix.numel() - 1) // chunk_size + 1):\n","        input_chunk = flat_tensor[i * chunk_size: (i + 1) * chunk_size].clone()\n","        quantized_chunk, (absmax_chunk, code) = quantize_blockwise(input_chunk, code=code)\n","        chunks.append(quantized_chunk)\n","        absmaxes.append(absmax_chunk)\n"," \n","    matrix_i8 = torch.cat(chunks).reshape_as(matrix)\n","    absmax = torch.cat(absmaxes)\n","    return matrix_i8, (absmax, code)\n"," \n"," \n","def convert_to_int8(model):\n","    \"\"\"Convert linear and embedding modules to 8-bit with optional adapters\"\"\"\n","    for module in list(model.modules()):\n","        for name, child in module.named_children():\n","            if isinstance(child, nn.Linear):\n","                print(name, child)\n","                setattr(\n","                    module,\n","                    name,\n","                    FrozenBNBLinear(\n","                        weight=torch.zeros(child.out_features, child.in_features, dtype=torch.uint8),\n","                        absmax=torch.zeros((child.weight.numel() - 1) // 4096 + 1),\n","                        code=torch.zeros(256),\n","                        bias=child.bias,\n","                    ),\n","                )\n","            elif isinstance(child, nn.Embedding):\n","                setattr(\n","                    module,\n","                    name,\n","                    FrozenBNBEmbedding(\n","                        weight=torch.zeros(child.num_embeddings, child.embedding_dim, dtype=torch.uint8),\n","                        absmax=torch.zeros((child.weight.numel() - 1) // 4096 + 1),\n","                        code=torch.zeros(256),\n","                    )\n","                )\n","\n","class GPTJBlock(transformers.models.gptj.modeling_gptj.GPTJBlock):\n","    def __init__(self, config):\n","        super().__init__(config)\n","\n","        convert_to_int8(self.attn)\n","        convert_to_int8(self.mlp)\n","\n","\n","class GPTJModel(transformers.models.gptj.modeling_gptj.GPTJModel):\n","    def __init__(self, config):\n","        super().__init__(config)\n","        convert_to_int8(self)\n","        \n","\n","class GPTJForCausalLM(transformers.models.gptj.modeling_gptj.GPTJForCausalLM):\n","    def __init__(self, config):\n","        super().__init__(config)\n","        convert_to_int8(self)\n","\n","\n","transformers.models.gptj.modeling_gptj.GPTJBlock = GPTJBlock"],"metadata":{"id":"lzw5FWEHT25v"},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":["We use the model configuration and the tokenizer of GPT-J-6B and set `pad_token` as `eos_token`.\n"],"metadata":{"id":"bPMcBs3tY_Pb"}},{"cell_type":"code","source":["config = transformers.GPTJConfig.from_pretrained(\"EleutherAI/gpt-j-6B\")\n","config.pad_token_id = config.eos_token_id\n","\n","tokenizer = transformers.AutoTokenizer.from_pretrained(\"EleutherAI/gpt-j-6B\")\n","tokenizer.pad_token = config.pad_token_id"],"metadata":{"id":"RYUBLX9ZUFZ5"},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":["## Load pretrained model\n","We load the pre-trained gpt-j-6b with 8-bit weights from [huggingface](https://huggingface.co/hivemind/gpt-j-6B-8bit). To reduce the peak RAM usage, we add the argument, `low_cpu_mem_usage=True` to `from_pretrained`."],"metadata":{"id":"U2Ze3RFeqo38"}},{"cell_type":"code","source":["gpt = GPTJForCausalLM.from_pretrained(\"hivemind/gpt-j-6B-8bit\", low_cpu_mem_usage=True)"],"metadata":{"id":"dCE0aTo2UJy0"},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":["## Add LoRA Adapter\n","Low-Rank Adaptation, or LoRA, which freezes the pretrained model weights and injects trainable rank decomposition matrices into each layer of the Transformer architecture, greatly reducing the number of trainable parameters for downstream tasks.\n","* We set `adapter_dim` from 16 to 4\n","* We set the Dropout `p` from 0 to 0.1"],"metadata":{"id":"0rf_Esgkx0AP"}},{"cell_type":"code","source":["def add_adapters(model, adapter_dim=4, p = 0.1):\n","    assert adapter_dim > 0\n","\n","    for name, module in model.named_modules():\n","      if isinstance(module, FrozenBNBLinear):\n","          if \"attn\" in name or \"mlp\" in name or \"head\" in name:\n","              print(\"Adding adapter to\", name)\n","              module.adapter = nn.Sequential(\n","                nn.Linear(module.in_features, adapter_dim, bias=False),\n","                nn.Dropout(p=p),\n","                nn.Linear(adapter_dim, module.out_features, bias=False),\n","            )\n","              print(\"Initializing\", name)\n","              nn.init.zeros_(module.adapter[2].weight)\n","\n","          else:\n","              print(\"Not adding adapter to\", name)\n","      elif isinstance(module, FrozenBNBEmbedding):\n","          print(\"Adding adapter to\", name)\n","          module.adapter = nn.Sequential(\n","                nn.Embedding(module.num_embeddings, adapter_dim),\n","                nn.Dropout(p=p),\n","                nn.Linear(adapter_dim, module.embedding_dim, bias=False),\n","            )\n","          print(\"Initializing\", name)\n","          nn.init.zeros_(module.adapter[2].weight)"],"metadata":{"id":"1DxucseIXrrv"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["add_adapters(gpt)\n","device = 'cuda' if torch.cuda.is_available() else 'cpu'\n","gpt.to(device)"],"metadata":{"id":"6xP3hmT5ax7b"},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":["## Load Dataset\n","We load our dataset as csv files from Google Drive. Data Structure:\n","\n","| pair                                              |   |   |   |   |\n","|---------------------------------------------------|---|---|---|---|\n","| [Text:] Lorem Ipsum<br>[Title:] Title of Lorem    |   |   |   |   |\n","| [Text:] Dolor Sit Amet<br>[Title:] Title of Dolor |   |   |   |   |\n","|                                                   |   |   |   |   |"],"metadata":{"id":"X-rVQTXbyKWi"}},{"cell_type":"code","source":["drive.mount('/content/drive')"],"metadata":{"id":"hsE9DrhXTpvg"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["dataset = load_dataset('csv', data_files={'train': '/content/drive/MyDrive/train.csv', 'test': '/content/drive/MyDrive/test.csv'})"],"metadata":{"id":"GPP2NOmsUnkU"},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":["## Tokenize Data\n","We tokenize our dataset and set `max_length=2048` for our task. Since GPTJ models operate with a maximum total token count of 2048 tokens, one Text-Title-Pair must not exceed this token count. You can use [this tool](https://beta.openai.com/tokenizer) to understand how a piece of text would be tokenized and the total count of tokens in that piece of text."],"metadata":{"id":"a5OFHXVUrW3x"}},{"cell_type":"code","source":["def tokenize_function(examples):\n","    return tokenizer(examples[\"pair\"], padding=True, truncation=True, max_length=2048)\n","\n","tokenized_datasets = dataset.map(tokenize_function, batched=True)\n","tokenized_datasets = tokenized_datasets.remove_columns([\"pair\"])\n","\n","#Convert to Torch Format\n","tokenized_datasets.set_format(\"torch\")"],"metadata":{"id":"5goVmhEqW8WX"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["full_train_dataset = tokenized_datasets[\"train\"]\n","train_dataloader = DataLoader(full_train_dataset, shuffle=False, batch_size=8)"],"metadata":{"id":"hgU13nR9W_vK"},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":["## Training Params\n","We define the training parameters and initialize the Optimizer, Schedluer and Scaler."],"metadata":{"id":"q4myrq39x-Mn"}},{"cell_type":"code","source":["num_epochs = 5\n","num_training_steps = num_epochs * len(train_dataloader)\n","num_warmup_steps = int(num_training_steps*0.1)\n","learning_rate = 1e-5 # Initial Learning Rate of 0.00001"],"metadata":{"id":"2OmGfuvViVk5"},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":["We activate gradient checkpointing for the model to reduce memory load."],"metadata":{"id":"SZWZJwppinOd"}},{"cell_type":"code","source":["gpt.gradient_checkpointing_enable()"],"metadata":{"id":"1XVEOYhVXCe8"},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":["We set a savepath for the model"],"metadata":{"id":"v3CQ2d0Vwwvn"}},{"cell_type":"code","source":["filepath = '/content/drive/MyDrive/model.pt'"],"metadata":{"id":"xbs_a7ocXFNM"},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":["We initialize the Optimizer, Schedluer and Scaler for training"],"metadata":{"id":"ToLy8Dzsi69p"}},{"cell_type":"code","source":["optimizer = Adam8bit(gpt.parameters(), lr=learning_rate, weight_decay=0.01)\n","lr_scheduler = transformers.get_linear_schedule_with_warmup(optimizer, num_warmup_steps, num_training_steps)\n","scaler = torch.cuda.amp.GradScaler()"],"metadata":{"id":"9-t6k6yTitJU"},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":["## Training Loop\n","We train the model and save it 👏"],"metadata":{"id":"SZ0SUgIiw08b"}},{"cell_type":"code","source":["progress_bar = tqdm(range(num_training_steps))\n","gpt.train()\n","k = 0\n","\n","# Iterate through epochs\n","for epoch in range(num_epochs):\n","    # Iterate through our training data in batches\n","    for batch in train_dataloader:\n","        k = k + 1\n","        if k % 500 == 0:\n","            print(k)\n","\n","            # Define a custom model state dict\n","            state = {'k': k, 'epoch': num_epochs, 'lr_scheduler': lr_scheduler.state_dict(\n","            ), 'state_dict': gpt.state_dict(), 'optimizer': optimizer.state_dict()}\n","\n","            # Save Model with Torch\n","            torch.save(state, filepath)\n","\n","        # Define batch to train on\n","        batch = {k: v.to(device) for k, v in batch.items()}\n","\n","        # Zeroing out the gradients\n","        # Explanation: https://stackoverflow.com/a/48009142\n","        optimizer.zero_grad()\n","\n","        # Runs the forward pass with autocasting\n","        with torch.cuda.amp.autocast():\n","\n","            # Feed batch in custom forward function\n","            out = gpt.forward(**batch,)\n","            \n","            # Custom loss function\n","            loss = F.cross_entropy(out.logits[:, :-1, :].flatten(0, -2),\n","                                   batch['input_ids'][:, 1:].flatten(),\n","                                   reduction='mean', \n","                                   label_smoothing=0.1)\n","\n","        print(loss)\n","\n","        # Scales loss\n","        # Calls backward() on scaled loss to create scaled gradients\n","        scaler.scale(loss).backward()\n","\n","        # Unscales gradients held by optimizer’s assigned parameters\n","        scaler.unscale_(optimizer)\n","\n","        # clips the norm of the overall gradient by concatenating all parameters passed to the function\n","        # documentation: https://pytorch.org/docs/stable/_modules/torch/nn/utils/clip_grad.html\n","        torch.nn.utils.clip_grad_norm_(gpt.parameters(), 1.0)\n","\n","        # Updates the scale for next iteration.\n","        scaler.step(optimizer)\n","        scaler.update()\n","\n","        # Updates initial learning rate\n","        lr_scheduler.step()\n","\n","        progress_bar.update(1)"],"metadata":{"id":"h-fah4lpXIK0"},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":["## Evaluate Model\n","We evaluate the model and generate a headline."],"metadata":{"id":"yLTZql-pw4F4"}},{"cell_type":"code","source":["gpt.eval()\n","\n","input_text=\"[Text:]Lorem Ipsum\\n\\n[Titel:]\"\n","\n","with torch.no_grad():\n","  prompt = tokenizer(cleaned_text, truncation=True, padding=True, return_tensors='pt')\n","  prompt = {key: value.to(device) for key, value in prompt.items()}\n","  out = gpt.generate(**prompt, min_length=5, max_new_tokens=256, top_p=0.7, temperature=1.0, do_sample=True)\n","  print(tokenizer.decode(out[0]))"],"metadata":{"id":"WpaFYVxHt9r4"},"execution_count":null,"outputs":[]}]}