llama-multidoc-multipack-flex.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "f6a37387",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = '1'\n",
    "os.environ['TF_FORCE_GPU_ALLOW_GROWTH'] = 'true'\n",
    "os.environ[\"HF_HUB_ENABLE_HF_TRANSFER\"] = \"1\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "54a6b493",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer\n",
    "from peft import LoraConfig, get_peft_model, AutoPeftModelForCausalLM, PeftModel\n",
    "import torch\n",
    "import transformers\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "f10e8dda",
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip3.10 uninstall transformers -y\n",
    "# !pip3.10 install -e ."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "a79188b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained('HuggingFaceTB/SmolLM2-135M-Instruct')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "eda02a81",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = AutoModelForCausalLM.from_pretrained(\n",
    "    'HuggingFaceTB/SmolLM2-135M-Instruct', attn_implementation = 'flex_attention'\n",
    ").cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "4a45f4e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_sdpa = AutoModelForCausalLM.from_pretrained(\n",
    "    'HuggingFaceTB/SmolLM2-135M-Instruct', attn_implementation = 'sdpa'\n",
    ").cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "de7f4f4f",
   "metadata": {},
   "outputs": [],
   "source": [
    "texts = [\n",
    "    'how to solve world hunger',\n",
    "    '1+1',\n",
    "]\n",
    "input_ids, position_ids, lengths = [], [], [0]\n",
    "for t in texts:\n",
    "    d = [\n",
    "        {'role': 'user', 'content': t}\n",
    "    ]\n",
    "    d = tokenizer.apply_chat_template(d)\n",
    "    input_ids.extend(d)\n",
    "    position_ids.extend(list(range(len(d))))\n",
    "    lengths.append(len(d))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "afaff11c",
   "metadata": {},
   "outputs": [],
   "source": [
    "lengths = torch.tensor(np.cumsum(lengths)).cuda()\n",
    "input_ids = torch.tensor(input_ids).cuda()\n",
    "position_ids = torch.tensor(position_ids, dtype = torch.int32).cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "dcc63e84",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(5.5380, device='cuda:0', grad_fn=<NllLossBackward0>)"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model(input_ids = input_ids[None],\n",
    "      labels = input_ids[None],\n",
    "      position_ids = position_ids[None],\n",
    "      attention_mask = [lengths]).loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "5f855e7b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def block_diagonal_concat_inverted(*masks, dtype=torch.bfloat16):\n",
    "    total_size = sum(mask.size(0) for mask in masks)\n",
    "    combined_mask = torch.zeros(total_size, total_size, dtype=dtype)\n",
    "\n",
    "    current_pos = 0\n",
    "\n",
    "    for mask in masks:\n",
    "        size = mask.size(0)\n",
    "        combined_mask[current_pos:current_pos + size, current_pos:current_pos + size] = mask\n",
    "        current_pos += size\n",
    "\n",
    "    min_value = torch.finfo(dtype).min if dtype.is_floating_point else torch.iinfo(dtype).min\n",
    "    inverted_mask = torch.where(combined_mask == 1, torch.tensor(0, dtype=dtype), min_value)\n",
    "    return inverted_mask.unsqueeze(0)\n",
    "\n",
    "masks = []\n",
    "for f in [lengths]:\n",
    "    masks_ = []\n",
    "    masking = torch.diff(f)\n",
    "    for m in masking:\n",
    "        masks_.append(torch.tril(torch.ones(m, m)))\n",
    "    \n",
    "    masks.append(block_diagonal_concat_inverted(*masks_, dtype = model.dtype))\n",
    "    \n",
    "masks = torch.stack(masks, 0).to('cuda')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "49c30ba6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(5.5380, device='cuda:0', grad_fn=<NllLossBackward0>)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model_sdpa(input_ids = input_ids[None],\n",
    "      labels = input_ids[None],\n",
    "      position_ids = position_ids[None],\n",
    "      attention_mask = masks).loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f7ed0196",
   "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
}