Big_science_Inference_accelerate.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyPOZbI9DfpyFAzGeoDTSfFC",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/abhilash1910/Framework-Optimization/blob/master/Big_science_Inference_accelerate.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8L344W-SfHVP"
      },
      "outputs": [],
      "source": [
        "from argparse import Namespace\n",
        "\n",
        "import torch\n",
        "from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer\n",
        "import accelerate\n",
        "from utils import Model, get_downloaded_model_path, print_rank_n\n",
        "\n",
        "\n",
        "class HFAccelerateModel(Model):\n",
        "    def __init__(self, args: Namespace) -> None:\n",
        "        print_rank_n(\"Loading model...\")\n",
        "\n",
        "        downloaded_model_path = get_downloaded_model_path(args.model_name)\n",
        "\n",
        "        self.tokenizer = AutoTokenizer.from_pretrained(downloaded_model_path)\n",
        "        self.pad = self.tokenizer.pad_token_id\n",
        "        print(accelerate.modeling.get_max_memory())\n",
        "        kwargs = {\n",
        "            \"pretrained_model_name_or_path\": downloaded_model_path,\n",
        "            \"device_map\": \"auto\",\n",
        "            \"max_memory\": get_max_memory_per_gpu_dict(\n",
        "                args.dtype,\n",
        "                args.model_name\n",
        "            )\n",
        "        }\n",
        "        if (args.dtype == torch.int8):\n",
        "            kwargs[\"load_in_8bit\"] = True\n",
        "        else:\n",
        "            kwargs[\"torch_dtype\"] = args.dtype\n",
        "\n",
        "        self.model = AutoModelForCausalLM.from_pretrained(**kwargs)\n",
        "\n",
        "        self.model.requires_grad_(False)\n",
        "        self.model.eval()\n",
        "        self.input_device = \"cuda:0\"\n",
        "\n",
        "        print_rank_n(\"Model loaded\")\n",
        "\n",
        "\n",
        "def get_max_memory_per_gpu_dict(dtype, model_name):\n",
        "    \"\"\" try to generate the memory map based on what we know about the model and the available hardware \"\"\"\n",
        "\n",
        "    # figure out the memory map - the minimum per gpu required to load the model\n",
        "    n_gpus = torch.cuda.device_count()\n",
        "\n",
        "    if model_name == \"bigscience/bloom\" and n_gpus == 8 and torch.cuda.get_device_properties(0).total_memory > 79*2**30:\n",
        "        # hand crafted optimized memory map for 8x80 setup over BLOOM\n",
        "        # this works with bs=40\n",
        "        if (dtype in [torch.bfloat16, torch.float16]):\n",
        "            max_memory_per_gpu = {0: '0GIB', 1: '51GIB', 2: '51GIB', 3: '51GIB',\n",
        "                                  4: '51GIB', 5: '51GIB', 6: '51GIB', 7: '51GIB'}\n",
        "        elif (dtype == torch.int8):\n",
        "            max_memory_per_gpu = {0: '0GIB', 1: '26GIB', 2: '26GIB', 3: '26GIB',\n",
        "                                  4: '26GIB', 5: '26GIB', 6: '26GIB', 7: '26GIB'}\n",
        "        print_rank_n(\"Max memory per gpu:\", max_memory_per_gpu)\n",
        "        return max_memory_per_gpu\n",
        "    try:\n",
        "        # model_params calculation, as we don't have a model yet to do:\n",
        "        #model_params = sum(dict((p.data_ptr(), p.numel()) for p in model.parameters()).values())\n",
        "\n",
        "        config = AutoConfig.from_pretrained(model_name)\n",
        "        h = config.hidden_size\n",
        "        l = config.n_layer\n",
        "        v = config.vocab_size\n",
        "        # from https://github.com/bigscience-workshop/bigscience/tree/6917a3b5fefcf439d3485ca184b4d9f6ab605150/math#model-sizing\n",
        "        model_params = l*(12*h**2 + 13*h) + v*h + 4*h\n",
        "    except:\n",
        "        print_rank_n(\n",
        "            f\"The model {model_name} has a broken config file. Please notify the owner\")\n",
        "        raise\n",
        "\n",
        "    if (dtype == torch.int8):\n",
        "        bytes = 1\n",
        "    else:\n",
        "        bytes = torch.finfo(dtype).bits / 8\n",
        "    param_memory_total_in_bytes = model_params * bytes\n",
        "    # add 5% since weight sizes aren't the same and some GPU may need more memory\n",
        "    param_memory_per_gpu_in_bytes = int(\n",
        "        param_memory_total_in_bytes / n_gpus * 1.10)\n",
        "    print_rank_n(\n",
        "        f\"Estimating {param_memory_per_gpu_in_bytes/2**30:0.2f}GB per gpu for weights\")\n",
        "\n",
        "    # check the real available memory\n",
        "    # load cuda kernels first and only measure the real free memory after loading (shorter by ~2GB)\n",
        "    torch.ones(1).cuda()\n",
        "    max_memory_per_gpu_in_bytes = torch.cuda.mem_get_info(0)[0]\n",
        "    if max_memory_per_gpu_in_bytes < param_memory_per_gpu_in_bytes:\n",
        "        raise ValueError(\n",
        "            f\"Unable to generate the memory map automatically as the needed estimated memory per gpu ({param_memory_per_gpu_in_bytes/2**30:0.2f}GB) is bigger than the available per gpu memory ({max_memory_per_gpu_in_bytes/2**30:0.2f}GB)\")\n",
        "\n",
        "    max_memory_per_gpu = {\n",
        "        i: param_memory_per_gpu_in_bytes for i in range(torch.cuda.device_count())}\n",
        "    print(\"Max memory per gpu:\", max_memory_per_gpu)\n",
        "    return max_memory_per_gpu"
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "6jg11OHIfO8e"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}