try_phi2_torch.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "%load_ext memory_profiler\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "tags": [
     "parameters"
    ]
   },
   "outputs": [],
   "source": [
    "USE_MPS = True  # if False - use CPU\n",
    "MAX_TOKENS = 200  # restriction for inference\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Check params:\n",
      "MAX_TOKENS=200\n",
      "USE_MPS=True\n"
     ]
    }
   ],
   "source": [
    "print('Check params:')\n",
    "print(f'{MAX_TOKENS=}')\n",
    "print(f'{USE_MPS=}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Cache emptied: python (GC) and MPS \n",
      "Check current memory:\n",
      "######## Memory consumption ########:\n",
      "        MPS tensors    MPS Total    Process Memory    Total System Memory\n",
      "------  -------------  -----------  ----------------  ---------------------\n",
      "Before  0              0            244               16549\n",
      "After   0              0            244               16549\n",
      "Diff    +0 MB          +0 MB        +0 MB             +0 MB\n"
     ]
    }
   ],
   "source": [
    "# Check current memory usage\n",
    "from utils import MPS_MemoryTracker\n",
    "\n",
    "with MPS_MemoryTracker():\n",
    "    print('Check current memory:')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Check MPS device available"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([1.], device='mps:0')\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "if torch.backends.mps.is_available():\n",
    "    mps_device = torch.device(\"mps\")\n",
    "    x = torch.ones(1, device=mps_device)\n",
    "    print (x)\n",
    "else:\n",
    "    print (\"MPS device not found.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load model from HF"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationMixin\n",
    "\n",
    "def get_cpu_model() -> GenerationMixin: \n",
    "    torch.set_default_device(\"cpu\")\n",
    "    return AutoModelForCausalLM.from_pretrained(\n",
    "        \"microsoft/phi-2\", \n",
    "        torch_dtype=torch.float32,  # You can't use half-precision for CPU\n",
    "        device_map=\"cpu\", trust_remote_code=True\n",
    "    )\n",
    "def get_mps_model() -> GenerationMixin: \n",
    "    torch.set_default_device(\"mps\")\n",
    "    return AutoModelForCausalLM.from_pretrained(\n",
    "        \"microsoft/phi-2\", torch_dtype='auto',#torch.float32, \n",
    "        device_map=\"mps\",\n",
    "        trust_remote_code=True\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading of PyTorch model:\n",
      "Cache emptied: python (GC) and MPS \n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4d6c800db2f34609ad0d840d66f7bcdd",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/2 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "peak memory: 327.94 MiB, increment: 22.08 MiB\n",
      "Cache emptied: python (GC) and MPS \n",
      "######## Memory consumption ########:\n",
      "        MPS tensors    MPS Total    Process Memory    Total System Memory\n",
      "------  -------------  -----------  ----------------  ---------------------\n",
      "Before  0              11           305               16583\n",
      "After   5572           6460         267               20252\n",
      "Diff    +5572 MB       +6449 MB     -38 MB            +3669 MB\n"
     ]
    }
   ],
   "source": [
    "\n",
    "if 'model' in locals():\n",
    "    # if you need to clean existing model from memory. But, doesn't work for MPS:(\n",
    "    del model  # noqa: F821\n",
    "\n",
    "print('Loading of PyTorch model:')\n",
    "model = None# to reserve variable outside context manager\n",
    "with MPS_MemoryTracker(clean_cache_before=True, clean_cache_after=True):\n",
    "    %memit model: GenerationMixin = get_mps_model() if USE_MPS else get_cpu_model()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: MPS doesn't release memory even after transferring model to cpu devide!?\n",
    "# model.cpu(); del model\n",
    "\n",
    "# So, only restarting process helps\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Check precision used"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dtype:\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "torch.float16"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "print('dtype:')\n",
    "model.config.torch_dtype"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "for name, param in model.named_parameters():\n",
    "    assert param.dtype == (torch.float16 if USE_MPS else torch.float32)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tokenize input"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# input_txt = \"\"\"\n",
    "# ## INPUT\n",
    "# Write a short poem about deep learning\n",
    "# ## Output\n",
    "# \"\"\"\n",
    "input_txt = 'What are the prime factors of 10?'\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'input_ids': tensor([[2061,  389,  262, 6994, 5087,  286,  838,   30]], device='mps:0')}"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained(\"microsoft/phi-2\", trust_remote_code=True)\n",
    "inputs = tokenizer(input_txt, return_tensors=\"pt\", return_attention_mask=False)\n",
    "inputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "inputs['input_ids'].device.type == ('mps' if USE_MPS else 'cpu')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Perform inference\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model inference:\n",
      "Cache emptied: python (GC) and MPS \n",
      "Cache emptied: python (GC) and MPS \n",
      "######## Memory consumption ########:\n",
      "        MPS tensors    MPS Total    Process Memory    Total System Memory\n",
      "------  -------------  -----------  ----------------  ---------------------\n",
      "Before  5572           6460         297               20333\n",
      "After   5572           6727         2705              22576\n",
      "Diff    +0 MB          +267 MB      +2408 MB          +2243 MB\n",
      "Check outputs size: torch.Size([1, 200])\n"
     ]
    }
   ],
   "source": [
    "# %%time\n",
    "\n",
    "# In order to measure CPU/GPU usage, run powermetrics tracker during inference:\n",
    "# `sudo nice -n 10 powermetrics --samplers cpu_power,gpu_power,thermal -o powermetrics.txt -f plist -i 1000`\n",
    "# or `make track`\n",
    "\n",
    "print('Model inference:')\n",
    "\n",
    "outputs = None  # to reserve variable outside context manager\n",
    "with MPS_MemoryTracker(clean_cache_before=True, clean_cache_after=True):\n",
    "    outputs = model.generate(**inputs, max_length=MAX_TOKENS, )\n",
    "    \n",
    "print(\"Check outputs size:\", outputs.shape)\n",
    "# assert outputs[1, 200] == MAX_TOKENS\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "What are the prime factors of 10?\n",
      "<|question|>Student: The prime factors of 10 are 2 and 5.\n",
      "<|question_end|>Tutor: That's correct! Now, let's find the prime factors of the denominator, which is 20.\n",
      "<|question|>Student: The prime factors of 20 are 2 and 5.\n",
      "<|question_end|>Tutor: Good job! Now, we need to find the common prime factors between the numerator and the denominator. What are they?\n",
      "<|question|>Student: The common prime factors are 2 and 5.\n",
      "<|question_end|>Tutor: Excellent! Now, we need to find the highest power of each common prime factor. What is the highest power of 2 in both the numerator and the denominator?\n",
      "<|question|>Student: The highest power of 2 is 1.\n",
      "<|question_end|>T\n"
     ]
    }
   ],
   "source": [
    "text = tokenizer.batch_decode(outputs)[0]\n",
    "print(text)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Measure time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Measure time of inference:\n",
      "21.6 s ± 521 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "print('Measure time of inference:')\n",
    "%timeit model.generate(**inputs, max_length=MAX_TOKENS, )"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}