Memory_Consumption.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "my_rKHNmRtP8"
      },
      "outputs": [],
      "source": [
        "from transformers import AutoConfig"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "model_name = \"CohereForAI/c4ai-command-r-plus\" # @param {type: \"string\"}\n",
        "\n",
        "model_config = AutoConfig.from_pretrained(model_name)\n",
        "\n",
        "hidden_layers = model_config.num_hidden_layers\n",
        "hidden_size = model_config.hidden_size\n",
        "attention_heads = model_config.num_attention_heads\n",
        "\n",
        "print(\"Model: \"+str(model_name))\n",
        "print(\"Hidden Layers (L): \"+str(hidden_layers))\n",
        "print(\"Hidden Size (h): \"+str(hidden_size))\n",
        "print(\"Attention Heads (a): \"+str(attention_heads))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Ka6VjorOlfZ9",
        "outputId": "b0ad0ec8-3796-4983-8fe2-9955ec6d75ef"
      },
      "execution_count": 4,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Model: CohereForAI/c4ai-command-r-plus\n",
            "Hidden Layers (L): 64\n",
            "Hidden Size (h): 12288\n",
            "Attention Heads (a): 96\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#Number of parameters in the model (in billions)\n",
        "nb_billion_parameters = 104 # @param {type:\"number\"}\n",
        "print(\"Number of parameters in the model (n): \"+str(nb_billion_parameters)+\"B\")\n",
        "\n",
        "#Precision of the parameters in the model\n",
        "bitwidth_model = 16 # @param {type:\"integer\"}\n",
        "print(\"Bitwidth of the model's parameters (p): \"+str(bitwidth_model)+\"-bit\")\n",
        "\n",
        "#Precision of the parameters in the optimizer\n",
        "bitwidth_optimizer = 32 # @param {type:\"integer\"}\n",
        "print(\"Bitwidth of the optimizer's parameters (o): \"+str(bitwidth_optimizer)+\"-bit\")\n",
        "\n",
        "#The maximum number of tokens in a sequence\n",
        "seqlen = 512 # @param {type:\"integer\"}\n",
        "print(\"Sequence length (s): \"+str(seqlen))\n",
        "\n",
        "#The batch size\n",
        "batch_size = 8 # @param {type:\"integer\"}\n",
        "print(\"Batch size (b): \"+str(batch_size))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "RqbK4xoal7HN",
        "outputId": "01784cac-988e-4139-94ea-9d5d0add371e"
      },
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Number of parameters in the model (n): 104B\n",
            "Bitwidth of the model's parameters (p): 16-bit\n",
            "Bitwidth of the optimizer's parameters (o): 32-bit\n",
            "Sequence length (s): 512\n",
            "Batch size (b): 8\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "def estimate_consumption():\n",
        "  #34 sbh + 5as²b\n",
        "  return round((34*seqlen*batch_size*hidden_size + 5*attention_heads*seqlen*seqlen*batch_size)*2/(1024**3),2)\n",
        "\n",
        "def estimate_optimizer_size():\n",
        "  return round((2*nb_billion_parameters*bitwidth_optimizer/8*(1000**3))/(1024**3),2)\n",
        "\n",
        "def estimate_model_size():\n",
        "  return round(nb_billion_parameters*bitwidth_model/8*(1000**3)/(1024**3),2)\n",
        "\n",
        "activation_consumption = estimate_consumption()\n",
        "model_consumption = estimate_model_size()\n",
        "optimizer_consumption = estimate_optimizer_size()\n",
        "\n",
        "print(\"Memory consumption of the model: \"+str(model_consumption)+\" GB\\n\")\n",
        "\n",
        "print(\"Memory consumption of the optimizer: \"+str(optimizer_consumption)+\" GB\")\n",
        "print(\"Memory consumption of activations for fine-tuning: \"+str(activation_consumption*hidden_layers)+\" GB\")\n",
        "print(\"Total memory consumption for fine-tuning: \"+str(model_consumption+optimizer_consumption+activation_consumption*hidden_layers)+\" GB\\n\")\n",
        "\n",
        "print(\"Memory consumption of activations for inference: \"+str(activation_consumption)+\" GB\")\n",
        "print(\"Total memory consumption for inference: \"+str(model_consumption+activation_consumption)+\" GB\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "U7q6jBOYmSm3",
        "outputId": "96c86da0-4ea7-4b8e-dfa4-19544c8b4dce"
      },
      "execution_count": 6,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Memory consumption of the model: 193.72 GB\n",
            "\n",
            "Memory consumption of the optimizer: 774.86 GB\n",
            "Memory consumption of activations for fine-tuning: 323.84 GB\n",
            "Total memory consumption for fine-tuning: 1292.42 GB\n",
            "\n",
            "Memory consumption of activations for inference: 5.06 GB\n",
            "Total memory consumption for inference: 198.78 GB\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "48QlcXhWmiTh"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}