notebooks/hf_adapted_inference.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "632a649d",
   "metadata": {},
   "source": [
    "# Simple HuggingFace inference with Huggingface Adapted FMS models"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3d1500f8",
   "metadata": {},
   "source": [
    "*Note: This notebook is using Torch 2.1.0 and Transformers 4.35.0.dev0*\n",
    "\n",
    "If you would like to run a similar pipeline using a script, please view the following file: `scripts/hf_compile_example.py`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "36289ab1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from transformers import pipeline, AutoTokenizer\n",
    "from fms.models import get_model\n",
    "from fms.models.hf import to_hf_api"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ee7b5344",
   "metadata": {},
   "source": [
    "## load Huggingface Adapted FMS model\n",
    "\n",
    "Simply get the Huggingface model and convert it to an equivalent HF adapted FMS model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "7b02ec29-289e-425a-960e-a66d6521730a",
   "metadata": {},
   "outputs": [],
   "source": [
    "architecture = \"llama\"\n",
    "variant = \"13b\"\n",
    "model_path = \"/path/to/hf_model\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7fa0364d",
   "metadata": {},
   "source": [
    "If you intend to use half tensors, you must set the default device to cuda and default dtype to half tensors prior to loading the model to save space in memory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "51c92721",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.set_default_device(\"cuda\")\n",
    "torch.set_default_dtype(torch.half)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2f05d812",
   "metadata": {},
   "source": [
    "get the model and wrap in huggingface adapter api"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "63af00ae-b041-4e2f-b882-c99389d967cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = get_model(\n",
    "    architecture,\n",
    "    variant,\n",
    "    model_path=model_path,\n",
    "    source=\"hf\",\n",
    "    device_type=\"cuda\",\n",
    "    norm_eps=1e-6,\n",
    ")\n",
    "model = to_hf_api(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "480c816c",
   "metadata": {},
   "source": [
    "## Simple inference with Huggingface pipelines"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "01bf9ad6",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained(model_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "149e9ae7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[{'generated_text': 'I believe the meaning of life is to find your purpose and to fulfill it.\\n\\nI believe that everyone has a unique purpose in life, and that'}]\n",
      "1.14 s ± 0 ns per loop (mean ± std. dev. of 1 run, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit -r 1 -n 1\n",
    "pipe = pipeline(\n",
    "    task=\"text-generation\",\n",
    "    model=model,\n",
    "    max_new_tokens=25,\n",
    "    tokenizer=tokenizer,\n",
    "    device=\"cuda\",\n",
    ")\n",
    "prompt = \"\"\"I believe the meaning of life is\"\"\"\n",
    "result = pipe(prompt)\n",
    "print(result)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eb08e175",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "2cb14521",
   "metadata": {},
   "source": [
    "## Compilation\n",
    "\n",
    "All fms models support torch compile for faster inference, therefore Huggingface Adapted FMS models also support this feature. \n",
    "\n",
    "*Note: `generate` calls the underlying decoder and not the model itself, which requires compiling the underlying decoder.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "a24655fb",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.decoder = torch.compile(model.decoder)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dbdfd3ce",
   "metadata": {},
   "source": [
    "Because compile is lazy, we first just do a single generation pipeline to compile the graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "fa3c43f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "pipe = pipeline(\n",
    "    task=\"text-generation\",\n",
    "    model=model,\n",
    "    max_new_tokens=25,\n",
    "    tokenizer=tokenizer,\n",
    "    device=\"cuda\",\n",
    ")\n",
    "prompt = \"\"\"I believe the meaning of life is\"\"\"\n",
    "result = pipe(prompt)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "55206d51",
   "metadata": {},
   "source": [
    "At this point, the graph should be compiled and we can get proper performance numbers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "7b22dc6b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[{'generated_text': 'I believe the meaning of life is to find your purpose and to fulfill it.\\n\\nI believe that everyone has a unique purpose in life, and that'}]\n",
      "648 ms ± 0 ns per loop (mean ± std. dev. of 1 run, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit -r 1 -n 1\n",
    "pipe = pipeline(\n",
    "    task=\"text-generation\",\n",
    "    model=model,\n",
    "    max_new_tokens=25,\n",
    "    tokenizer=tokenizer,\n",
    "    device=\"cuda\",\n",
    ")\n",
    "prompt = \"\"\"I believe the meaning of life is\"\"\"\n",
    "result = pipe(prompt)\n",
    "print(result)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63f44374",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}