local_inference/fp8-405B.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Run Llama-3.1-405B-FP8-Instruct "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Note that running the FP8 model requires GPUs with compute capability > 9. A potential working setup would be 8*H100**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's first install the required libraries:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! pip install transformers accelerate"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that for torch and [fbgemm-gpu](https://huggingface.co/docs/transformers/main/en/quantization/fbgemm_fp8) libraries, you might need to download the nighly version. Just follow the instruction here :\n",
    "https://pytorch.org/FBGEMM/fbgemm_gpu-development/InstallationInstructions.html\n",
    "\n",
    "Change with your version of cuda. In this example, we are installing the nighlty version with cuda 12.1\n",
    "```\n",
    "pip install torch --index-url https://download.pytorch.org/whl/cu121/\n",
    "pip install fbgemm-gpu --index-url https://download.pytorch.org/whl/cu121/\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! pip install torch fbgemm-gpu"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We import the required libraries : "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's load the model. The model has already been quantized with fbgemm_fp8 as specified in the model's [config.json](https://huggingface.co/meta-llama/Llama-3.1-405B-Instruct-FP8/blob/main/config.json), so we don't need to specify a `quantization_config` and can load the quantized model as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_name = \"meta-llama/Meta-Llama-3.1-405B-Instruct-FP8\"\n",
    "\n",
    "quantized_model = AutoModelForCausalLM.from_pretrained(\n",
    "\tmodel_name, device_map=\"auto\", torch_dtype=torch.bfloat16)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then, we need to prepare the inputs: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
    "input_text = \"What are we having for dinner?\"\n",
    "input_ids = tokenizer(input_text, return_tensors=\"pt\").to(\"cuda\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we can generate the output ! "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output = quantized_model.generate(**input_ids, max_new_tokens=10)\n",
    "print(tokenizer.decode(output[0], skip_special_tokens=True))"
   ]
  }
 ],
 "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}