torch-compile-llama-3.2-3b.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "3280345f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = '2'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "440f2e85",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer\n",
    "from transformers.models.llama.modeling_llama import LlamaForCausalLM\n",
    "from peft import LoraConfig, get_peft_model, AutoPeftModelForCausalLM, PeftModel\n",
    "from cut_cross_entropy.transformers import cce_patch\n",
    "from cut_cross_entropy import linear_cross_entropy\n",
    "import torch\n",
    "import transformers\n",
    "import numpy as np\n",
    "import random\n",
    "from torchviz import make_dot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "da209bcd",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained('unsloth/Llama-3.2-3B-Instruct')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "9913eb2b",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Llama(LlamaForCausalLM):\n",
    "    def __init__(self, config):\n",
    "        super().__init__(config)\n",
    "        \n",
    "    def forward(self, **kwargs):\n",
    "        labels = kwargs.pop('labels')\n",
    "        kwargs.pop('output_hidden_states', None)\n",
    "        super_out = super().forward(**kwargs, output_hidden_states = True)\n",
    "        if labels is not None:\n",
    "            embeddings = super_out.hidden_states[-1]\n",
    "            auto_shift_loss = linear_cross_entropy(embeddings, self.lm_head, labels, shift=True, impl = 'torch_compile')\n",
    "            return {'loss': auto_shift_loss}\n",
    "        return super_out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "269d2494",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = Llama.from_pretrained(\n",
    "    'unsloth/Llama-3.2-3B-Instruct',\n",
    "    torch_dtype = torch.bfloat16\n",
    ").cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "270b500a",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_auto = AutoModelForCausalLM.from_pretrained(\n",
    "    'unsloth/Llama-3.2-3B-Instruct',\n",
    "    torch_dtype = torch.bfloat16\n",
    ").cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "3ca9fdd9",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "rank = 256\n",
    "peft_config = LoraConfig(\n",
    "        lora_alpha=rank * 2,\n",
    "        lora_dropout=0.0,\n",
    "        r=rank,\n",
    "        bias=\"none\",\n",
    "        task_type=\"CAUSAL_LM\",\n",
    "        target_modules=[\"embed_tokens\", \"lm_head\"],\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "4c568df4",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/husein/.local/lib/python3.10/site-packages/peft/tuners/tuners_utils.py:543: UserWarning: Model with `tie_word_embeddings=True` and the tied_target_modules=['lm_head'] are part of the adapter. This can lead to complications, for example when merging the adapter or converting your model to formats other than safetensors. See for example https://github.com/huggingface/peft/issues/2018.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "model = get_peft_model(model, peft_config)\n",
    "model_auto = get_peft_model(model_auto, peft_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "19b18e39",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[128000, 128006,   9125, 128007,    271,  38766,   1303,  33025,   2696,\n",
       "             25,   6790,    220,   2366,     18,    198,  15724,   2696,     25,\n",
       "            220,   2437,   4448,    220,   2366,     20,    271, 128009, 128006,\n",
       "            882, 128007,    271,  13347,      0, 128009]], device='cuda:0')"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "input_ids = tokenizer.apply_chat_template([\n",
    "    {'role': 'user', 'content': 'Hi!'}\n",
    "], return_tensors = 'pt').cuda()\n",
    "input_ids"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "e860d57c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'loss': tensor(6.8273, device='cuda:0', grad_fn=<CompiledFunctionBackward>)}"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "o = model(input_ids = input_ids, labels = input_ids)\n",
    "o"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "08e59548",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(6.8273, device='cuda:0', grad_fn=<NllLossBackward0>)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "o = model_auto(input_ids = input_ids, labels = input_ids)\n",
    "o.loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "546c24bd",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "python3.10",
   "language": "python",
   "name": "python3.10"
  },
  "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.10.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}