stats.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "f3876fb2-fec7-43b1-a65a-531fb1e6cedf",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from IPython.display import display\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.ticker import FormatStrFormatter\n",
    "pd.options.display.max_columns = None\n",
    "pd.options.display.max_rows = None\n",
    "from modules import *\n",
    "from execution import *\n",
    "import json\n",
    "import pprint\n",
    "import models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "203561ac-b9f6-4fd7-a4a3-34fae1e87d4d",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "model is {'l': 2048, 'e': 25600, 'h': 160, 'depth': 128, 'f': 102400}\n",
      "num parameters = 1006.665728B\n",
      "total flops = 12.569785006994787PFLOPs\n",
      "flop ratio = 1.9074197861849107\n"
     ]
    }
   ],
   "source": [
    "''' choose the transformer architecture hyperparameters\n",
    "    get some high-level stats like flops and mem  '''\n",
    "\n",
    "model_str = 'gpt3_1T'  # what model? choose from models.py (gpt3_1T, vit_era5) \n",
    "model = models.models[model_str]\n",
    "\n",
    "# alternatively define your  model here with sequence length (l), embed (e), heads (h), depth (d)\n",
    "# model = {'l': 64800*0.5, 'e': 6144, 'h': 32, 'depth': 48}\n",
    "\n",
    "# set model hyperparams\n",
    "l = model['l']\n",
    "e = model['e']\n",
    "f = 4 * e\n",
    "model['f'] = f\n",
    "h = model['h']\n",
    "depth = model['depth']\n",
    "print('model is {}'.format(model))\n",
    "\n",
    "# which system? \n",
    "with open('systems/config-H200.json', 'r') as file:\n",
    "    system = json.load(file)\n",
    "\n",
    "# get some overall stats for the model \n",
    "df_mlp = mlp_1d(1, l, e, f, parallelism={'m': 1}, topology={'t': 1}, system=system)\n",
    "df_sa = sa_1d(1, l, e, h, parallelism={'m': 1}, topology={'t': 1}, flash_attention=True, system=system)\n",
    "flops = (df_mlp['flops_fwd'].sum() + df_mlp['flops_bwd'].sum() + df_sa['flops_fwd'].sum() + df_sa['flops_bwd'].sum()) * depth\n",
    "flop_ratio = (df_mlp['flops_fwd'].sum() + df_mlp['flops_bwd'].sum()) / (df_sa['flops_fwd'].sum() + df_sa['flops_bwd'].sum())\n",
    "param_count = ((df_mlp['weights_mem'].sum() + df_sa['weights_mem'].sum()) * depth) / system['element_size']\n",
    "print('num parameters = {}B'.format(param_count/1E9))\n",
    "print('total flops = {}PFLOPs'.format(flops/1E3))\n",
    "print('flop ratio = {}'.format(flop_ratio))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "1731f175-74bb-40ce-9a02-ae930b3d2217",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "num gpus = 2048, nvs domain size = 4, #possible candidates = 1810\n",
      "[(102.463805377367,\n",
      "  124.64239104,\n",
      "  {'dp': 16, 'mbs': 1, 'pp': 32, 'tp': 4},\n",
      "  {'acts_mem': 87.28346624000001,\n",
      "   'bubble_frac': 0.1060093337200795,\n",
      "   'comp_frac': 0.8093073814289401,\n",
      "   'dp_comm_frac': 0.018558749107651158,\n",
      "   'flops_per_gpu': 25139.570013989574,\n",
      "   'mem': 124.64239104,\n",
      "   'mem_frac': 0.0037681621009692263,\n",
      "   'nv_dp': 1,\n",
      "   'nv_pp': 1,\n",
      "   'nv_tp': 4,\n",
      "   'pp_comm_frac': 0.009656994091519603,\n",
      "   't': 39.97509154490915,\n",
      "   't_bubble': 4.237732820075002,\n",
      "   't_comp': 32.352136660592585,\n",
      "   't_dp_comm': 0.741887694537356,\n",
      "   't_mem': 0.150632624942302,\n",
      "   't_pp_comm': 0.38603922285714287,\n",
      "   't_tp_comm': 2.1066625219047626,\n",
      "   'tp_comm_frac': 0.05269937955084051,\n",
      "   'wts_grad_mem': 15.7300736,\n",
      "   'wts_mem': 15.7300736,\n",
      "   'wts_optimizer_states_mem': 5.898777600000001})]\n"
     ]
    }
   ],
   "source": [
    "''' what is the optimal configuration? '''\n",
    "# set your inputs\n",
    "system['nvlink_size'] = 4              # change the nvlink size if needed\n",
    "parallel_strat = '1D'                  # 1D, 2D: summa, 2D-seqp: context parallel\n",
    "total_gpus = 2048                      # total number of GPUs\n",
    "global_batch_size = 4096               # global batch size\n",
    "\n",
    "# note that 2D and 2D-seqp have a much larger design space and can take several minutes to run\n",
    "# to run a specific config, see the next cell\n",
    "if parallel_strat == '1D':\n",
    "    configs = execute_1d(model, [total_gpus], global_batch_size=global_batch_size, system=system, verbose=False, nlargest=100)\n",
    "elif parallel_strat == '2D-seqp': # context parallel 2D TP\n",
    "    configs = execute_seqp(model, [total_gpus], global_batch_size=global_batch_size, system=system, verbose=False, nlargest=100)\n",
    "elif parallel_strat == '2D': # SUMMA 2D TP\n",
    "    configs = execute_2d(model, [total_gpus], global_batch_size=global_batch_size, system=system, verbose=False, nlargest=100)\n",
    "else:\n",
    "    assert False, 'parallel strategy not valid!'\n",
    "    \n",
    "top_configs_to_print = 1 # how many configs to print? max 100 but dont print all \n",
    "pprint.pprint(configs[0][0:top_configs_to_print]) \n",
    "\n",
    "### what info do the configs give you?\n",
    "# here's an example config:\n",
    "# [(102.463805377367,                                # throughput in samples/sec\n",
    "#   124.64239104,                                    # memory consumed on the GPU in GB\n",
    "#   {'dp': 16, 'mbs': 1, 'pp': 32, 'tp': 4},         # optimal parallel config: dp (data parallel) mbs (microbatchsize) pp (pipeline) tp (tensor parallel; can be two nums) \n",
    "#   {'acts_mem': 87.28346624000001,                  # activation mem in GB\n",
    "#    'bubble_frac': 0.1060093337200795,              # fraction of time spent in pipeline bubbles\n",
    "#    'tp_comm_frac': 0.05269937955084051,            # fraction of time spent in TP comms\n",
    "#    'dp_comm_frac': 0.018558749107651158,           # fraction of time spent in DP comms\n",
    "#    'comp_frac': 0.8093073814289401,                # fraction of time spent in compute\n",
    "#    'flops_per_gpu': 25139.570013989574,            # flops per GPU\n",
    "#    'mem': 124.64239104,                            # memory consumed on the GPU in GB\n",
    "#    'mem_frac': 0.0037681621009692263,              # fraction of time spent in memory accesses from HBM\n",
    "#    'nv_dp': 1,                                     # number of GPUs in a node belonging to data parallel grp on NVLINK\n",
    "#    'nv_pp': 1,                                     # number of GPUs in a node belonging to pipeline parallel grp on NVLINK\n",
    "#    'nv_tp': 4,                                     # number of GPUs in a node belonging to tensor parallel grp on NVLINK\n",
    "#    'pp_comm_frac': 0.009656994091519603,           # fraction of time spent in PP comms\n",
    "#    't': 39.97509154490915,                         # total time per iteration in s\n",
    "#    't_bubble': 4.237732820075002,                  # time spent in bubbles\n",
    "#    't_comm': 2.1066625219047626,                   # time spent in tp comms\n",
    "#    't_comp': 32.352136660592585,                   # time spent in compute\n",
    "#    't_dp_comm': 0.741887694537356,                 # time spent in dp comms\n",
    "#    't_mem': 0.150632624942302,                     # time spent in memory accesses from HBM\n",
    "#    't_pp_comm': 0.38603922285714287,               # time spent in pp comms\n",
    "#    'wts_mem': 15.7300736,                          # weights mem in GB\n",
    "#    'wts_grad_mem': 15.7300736,                     # weights gradients mem in GB\n",
    "#    'wts_optimizer_states_mem': 5.898777600000001}) # optimizer stats mem in GB\n",
    "\n",
    "# if empty, then no feasible config (model is too big to fit on the GPU memory: reduce batch size etc..)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "b0bc5f4e-508a-4408-96e9-220dfde94091",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tp = 4, pp = 8, dp = 64, total = 2048, #microbatches = 64\n",
      "\n",
      "############## Single microbatch stats ##############\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>name</th>\n",
       "      <th>weights_mem</th>\n",
       "      <th>weights_grad_mem</th>\n",
       "      <th>flops_fwd</th>\n",
       "      <th>activation_buffer</th>\n",
       "      <th>comm_fwd</th>\n",
       "      <th>comm_fwd_type</th>\n",
       "      <th>flops_bwd</th>\n",
       "      <th>comm_bwd</th>\n",
       "      <th>comm_bwd_type</th>\n",
       "      <th>t_fwd</th>\n",
       "      <th>t_fwd_comm</th>\n",
       "      <th>t_fwd_comp</th>\n",
       "      <th>t_fwd_mem</th>\n",
       "      <th>intensity_fwd</th>\n",
       "      <th>t_bwd</th>\n",
       "      <th>t_bwd_comm</th>\n",
       "      <th>t_bwd_comp</th>\n",
       "      <th>t_bwd_mem</th>\n",
       "      <th>intensity_bwd</th>\n",
       "      <th>t</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>fc1</td>\n",
       "      <td>1.310720</td>\n",
       "      <td>1.310720</td>\n",
       "      <td>2.684302</td>\n",
       "      <td>0.104858</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>reducescatter</td>\n",
       "      <td>5.368001</td>\n",
       "      <td>0.104858</td>\n",
       "      <td>reducescatter</td>\n",
       "      <td>0.003411</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.003411</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>10.769524</td>\n",
       "      <td>0.007059</td>\n",
       "      <td>0.000257</td>\n",
       "      <td>0.006802</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>10.736752</td>\n",
       "      <td>0.010470</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>fc1-bias</td>\n",
       "      <td>0.000051</td>\n",
       "      <td>0.000051</td>\n",
       "      <td>0.000052</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>none</td>\n",
       "      <td>0.000052</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>none</td>\n",
       "      <td>0.000022</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000020</td>\n",
       "      <td>0.000001</td>\n",
       "      <td>0.937491</td>\n",
       "      <td>0.000022</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000020</td>\n",
       "      <td>0.000001</td>\n",
       "      <td>0.937491</td>\n",
       "      <td>0.000044</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>act1</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000419</td>\n",
       "      <td>0.104858</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>none</td>\n",
       "      <td>0.000682</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>none</td>\n",
       "      <td>0.000044</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000024</td>\n",
       "      <td>0.000020</td>\n",
       "      <td>0.547450</td>\n",
       "      <td>0.000066</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000026</td>\n",
       "      <td>0.000039</td>\n",
       "      <td>0.402336</td>\n",
       "      <td>0.000109</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>dpr1</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000052</td>\n",
       "      <td>0.052429</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>none</td>\n",
       "      <td>0.000052</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>none</td>\n",
       "      <td>0.000055</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000020</td>\n",
       "      <td>0.000034</td>\n",
       "      <td>0.375180</td>\n",
       "      <td>0.000055</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000020</td>\n",
       "      <td>0.000034</td>\n",
       "      <td>0.375180</td>\n",
       "      <td>0.000109</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>fc2</td>\n",
       "      <td>1.310720</td>\n",
       "      <td>1.310720</td>\n",
       "      <td>2.684302</td>\n",
       "      <td>0.104858</td>\n",
       "      <td>0.104858</td>\n",
       "      <td>reducescatter</td>\n",
       "      <td>5.368001</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>reducescatter</td>\n",
       "      <td>0.003668</td>\n",
       "      <td>0.000257</td>\n",
       "      <td>0.003411</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>10.769524</td>\n",
       "      <td>0.006802</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.006802</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>10.736752</td>\n",
       "      <td>0.010470</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>fc2-bias</td>\n",
       "      <td>0.000051</td>\n",
       "      <td>0.000051</td>\n",
       "      <td>0.000052</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>none</td>\n",
       "      <td>0.000052</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>none</td>\n",
       "      <td>0.000022</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000020</td>\n",
       "      <td>0.000001</td>\n",
       "      <td>0.937491</td>\n",
       "      <td>0.000022</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000020</td>\n",
       "      <td>0.000001</td>\n",
       "      <td>0.937491</td>\n",
       "      <td>0.000044</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>dpr2</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000013</td>\n",
       "      <td>0.013107</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>none</td>\n",
       "      <td>0.000013</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>none</td>\n",
       "      <td>0.000020</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000020</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>1.473812</td>\n",
       "      <td>0.000020</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000020</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>1.473812</td>\n",
       "      <td>0.000040</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>ln1</td>\n",
       "      <td>0.000102</td>\n",
       "      <td>0.000102</td>\n",
       "      <td>0.000118</td>\n",
       "      <td>0.026214</td>\n",
       "      <td>0.104858</td>\n",
       "      <td>allgather</td>\n",
       "      <td>0.000157</td>\n",
       "      <td>0.104858</td>\n",
       "      <td>reducescatter</td>\n",
       "      <td>0.000278</td>\n",
       "      <td>0.000257</td>\n",
       "      <td>0.000021</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>1.928110</td>\n",
       "      <td>0.000284</td>\n",
       "      <td>0.000257</td>\n",
       "      <td>0.000021</td>\n",
       "      <td>0.000006</td>\n",
       "      <td>0.785619</td>\n",
       "      <td>0.000563</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "       name  weights_mem  weights_grad_mem  flops_fwd  activation_buffer  \\\n",
       "0       fc1     1.310720          1.310720   2.684302           0.104858   \n",
       "1  fc1-bias     0.000051          0.000051   0.000052           0.000000   \n",
       "2      act1     0.000000          0.000000   0.000419           0.104858   \n",
       "3      dpr1     0.000000          0.000000   0.000052           0.052429   \n",
       "4       fc2     1.310720          1.310720   2.684302           0.104858   \n",
       "5  fc2-bias     0.000051          0.000051   0.000052           0.000000   \n",
       "6      dpr2     0.000000          0.000000   0.000013           0.013107   \n",
       "7       ln1     0.000102          0.000102   0.000118           0.026214   \n",
       "\n",
       "   comm_fwd  comm_fwd_type  flops_bwd  comm_bwd  comm_bwd_type     t_fwd  \\\n",
       "0  0.000000  reducescatter   5.368001  0.104858  reducescatter  0.003411   \n",
       "1  0.000000           none   0.000052  0.000000           none  0.000022   \n",
       "2  0.000000           none   0.000682  0.000000           none  0.000044   \n",
       "3  0.000000           none   0.000052  0.000000           none  0.000055   \n",
       "4  0.104858  reducescatter   5.368001  0.000000  reducescatter  0.003668   \n",
       "5  0.000000           none   0.000052  0.000000           none  0.000022   \n",
       "6  0.000000           none   0.000013  0.000000           none  0.000020   \n",
       "7  0.104858      allgather   0.000157  0.104858  reducescatter  0.000278   \n",
       "\n",
       "   t_fwd_comm  t_fwd_comp  t_fwd_mem  intensity_fwd     t_bwd  t_bwd_comm  \\\n",
       "0    0.000000    0.003411   0.000000      10.769524  0.007059    0.000257   \n",
       "1    0.000000    0.000020   0.000001       0.937491  0.000022    0.000000   \n",
       "2    0.000000    0.000024   0.000020       0.547450  0.000066    0.000000   \n",
       "3    0.000000    0.000020   0.000034       0.375180  0.000055    0.000000   \n",
       "4    0.000257    0.003411   0.000000      10.769524  0.006802    0.000000   \n",
       "5    0.000000    0.000020   0.000001       0.937491  0.000022    0.000000   \n",
       "6    0.000000    0.000020   0.000000       1.473812  0.000020    0.000000   \n",
       "7    0.000257    0.000021   0.000000       1.928110  0.000284    0.000257   \n",
       "\n",
       "   t_bwd_comp  t_bwd_mem  intensity_bwd         t  \n",
       "0    0.006802   0.000000      10.736752  0.010470  \n",
       "1    0.000020   0.000001       0.937491  0.000044  \n",
       "2    0.000026   0.000039       0.402336  0.000109  \n",
       "3    0.000020   0.000034       0.375180  0.000109  \n",
       "4    0.006802   0.000000      10.736752  0.010470  \n",
       "5    0.000020   0.000001       0.937491  0.000044  \n",
       "6    0.000020   0.000000       1.473812  0.000040  \n",
       "7    0.000021   0.000006       0.785619  0.000563  "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>name</th>\n",
       "      <th>weights_mem</th>\n",
       "      <th>weights_grad_mem</th>\n",
       "      <th>flops_fwd</th>\n",
       "      <th>activation_buffer</th>\n",
       "      <th>comm_fwd</th>\n",
       "      <th>comm_fwd_type</th>\n",
       "      <th>flops_bwd</th>\n",
       "      <th>comm_bwd</th>\n",
       "      <th>comm_bwd_type</th>\n",
       "      <th>t_fwd</th>\n",
       "      <th>t_fwd_comm</th>\n",
       "      <th>t_fwd_comp</th>\n",
       "      <th>t_fwd_mem</th>\n",
       "      <th>intensity_fwd</th>\n",
       "      <th>t_bwd</th>\n",
       "      <th>t_bwd_comm</th>\n",
       "      <th>t_bwd_comp</th>\n",
       "      <th>t_bwd_mem</th>\n",
       "      <th>intensity_bwd</th>\n",
       "      <th>t</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>qkv</td>\n",
       "      <td>0.983040</td>\n",
       "      <td>0.983040</td>\n",
       "      <td>2.013227</td>\n",
       "      <td>0.104858</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>reducescatter</td>\n",
       "      <td>4.025988</td>\n",
       "      <td>0.104858</td>\n",
       "      <td>reducescatter</td>\n",
       "      <td>0.002563</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.002563</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>10.548086</td>\n",
       "      <td>0.005364</td>\n",
       "      <td>0.000257</td>\n",
       "      <td>0.005106</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>10.505729</td>\n",
       "      <td>0.007927</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>fusedla</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.112155</td>\n",
       "      <td>0.105185</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>none</td>\n",
       "      <td>0.279225</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>none</td>\n",
       "      <td>0.000162</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000162</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>7.390304</td>\n",
       "      <td>0.000373</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000373</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>8.518728</td>\n",
       "      <td>0.000534</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>vproj</td>\n",
       "      <td>0.327680</td>\n",
       "      <td>0.327680</td>\n",
       "      <td>0.671036</td>\n",
       "      <td>0.026214</td>\n",
       "      <td>0.104858</td>\n",
       "      <td>reducescatter</td>\n",
       "      <td>1.342000</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>reducescatter</td>\n",
       "      <td>0.001125</td>\n",
       "      <td>0.000257</td>\n",
       "      <td>0.000868</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>9.079748</td>\n",
       "      <td>0.001715</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.001715</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>8.974639</td>\n",
       "      <td>0.002840</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>vproj-bias</td>\n",
       "      <td>0.000051</td>\n",
       "      <td>0.000051</td>\n",
       "      <td>0.000052</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>none</td>\n",
       "      <td>0.000052</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>none</td>\n",
       "      <td>0.000022</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000020</td>\n",
       "      <td>0.000001</td>\n",
       "      <td>0.937491</td>\n",
       "      <td>0.000022</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000020</td>\n",
       "      <td>0.000001</td>\n",
       "      <td>0.937491</td>\n",
       "      <td>0.000044</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>dpr_v</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000013</td>\n",
       "      <td>0.013107</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>none</td>\n",
       "      <td>0.000013</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>none</td>\n",
       "      <td>0.000020</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000020</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>1.473812</td>\n",
       "      <td>0.000020</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000020</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>1.473812</td>\n",
       "      <td>0.000040</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>ln2</td>\n",
       "      <td>0.000102</td>\n",
       "      <td>0.000102</td>\n",
       "      <td>0.000118</td>\n",
       "      <td>0.026214</td>\n",
       "      <td>0.104858</td>\n",
       "      <td>allgather</td>\n",
       "      <td>0.000157</td>\n",
       "      <td>0.104858</td>\n",
       "      <td>reducescatter</td>\n",
       "      <td>0.000278</td>\n",
       "      <td>0.000257</td>\n",
       "      <td>0.000021</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>1.928110</td>\n",
       "      <td>0.000284</td>\n",
       "      <td>0.000257</td>\n",
       "      <td>0.000021</td>\n",
       "      <td>0.000006</td>\n",
       "      <td>0.785619</td>\n",
       "      <td>0.000563</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "         name  weights_mem  weights_grad_mem  flops_fwd  activation_buffer  \\\n",
       "0         qkv     0.983040          0.983040   2.013227           0.104858   \n",
       "1     fusedla     0.000000          0.000000   0.112155           0.105185   \n",
       "2       vproj     0.327680          0.327680   0.671036           0.026214   \n",
       "3  vproj-bias     0.000051          0.000051   0.000052           0.000000   \n",
       "4       dpr_v     0.000000          0.000000   0.000013           0.013107   \n",
       "5         ln2     0.000102          0.000102   0.000118           0.026214   \n",
       "\n",
       "   comm_fwd  comm_fwd_type  flops_bwd  comm_bwd  comm_bwd_type     t_fwd  \\\n",
       "0  0.000000  reducescatter   4.025988  0.104858  reducescatter  0.002563   \n",
       "1  0.000000           none   0.279225  0.000000           none  0.000162   \n",
       "2  0.104858  reducescatter   1.342000  0.000000  reducescatter  0.001125   \n",
       "3  0.000000           none   0.000052  0.000000           none  0.000022   \n",
       "4  0.000000           none   0.000013  0.000000           none  0.000020   \n",
       "5  0.104858      allgather   0.000157  0.104858  reducescatter  0.000278   \n",
       "\n",
       "   t_fwd_comm  t_fwd_comp  t_fwd_mem  intensity_fwd     t_bwd  t_bwd_comm  \\\n",
       "0    0.000000    0.002563   0.000000      10.548086  0.005364    0.000257   \n",
       "1    0.000000    0.000162   0.000000       7.390304  0.000373    0.000000   \n",
       "2    0.000257    0.000868   0.000000       9.079748  0.001715    0.000000   \n",
       "3    0.000000    0.000020   0.000001       0.937491  0.000022    0.000000   \n",
       "4    0.000000    0.000020   0.000000       1.473812  0.000020    0.000000   \n",
       "5    0.000257    0.000021   0.000000       1.928110  0.000284    0.000257   \n",
       "\n",
       "   t_bwd_comp  t_bwd_mem  intensity_bwd         t  \n",
       "0    0.005106   0.000000      10.505729  0.007927  \n",
       "1    0.000373   0.000000       8.518728  0.000534  \n",
       "2    0.001715   0.000000       8.974639  0.002840  \n",
       "3    0.000020   0.000001       0.937491  0.000044  \n",
       "4    0.000020   0.000000       1.473812  0.000040  \n",
       "5    0.000021   0.000006       0.785619  0.000563  "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "############## Final stats ##############\n",
      "{'acts_mem': 87.28346624000001,\n",
      " 'bubble_frac': 0.09137148818975546,\n",
      " 'comp_frac': 0.7787387200882211,\n",
      " 'dp_comm_frac': 0.07323204836105164,\n",
      " 'flops_per_gpu': 25139.570013989574,\n",
      " 'mem': 219.02283264,\n",
      " 'mem_frac': 0.0036258333964687417,\n",
      " 'pp_comm_frac': 0.0023230590104872887,\n",
      " 't': 41.544276438350856,\n",
      " 't_bubble': 3.7959623639387114,\n",
      " 't_comp': 32.352136660592585,\n",
      " 't_dp_comm': 3.0423724612582084,\n",
      " 't_mem': 0.150632624942302,\n",
      " 't_pp_comm': 0.09650980571428572,\n",
      " 't_tp_comm': 2.1066625219047626,\n",
      " 'tp_comm_frac': 0.05070885095401577,\n",
      " 'wts_grad_mem': 62.9202944,\n",
      " 'wts_mem': 62.9202944,\n",
      " 'wts_optimizer_states_mem': 5.898777600000001}\n"
     ]
    }
   ],
   "source": [
    "''' play around with the configurations and \n",
    "   show each layer of the transformer and \n",
    "   associated states (flops, mem, comms, \n",
    "   intensities, etc.) and final times \n",
    "   for single iteration '''\n",
    "\n",
    "# set your configurations\n",
    "mbs = 1                                # microbatch size  \n",
    "tp1 = 4                                # tensor parallelism (dim1)\n",
    "tp2 = 1                                # tensor parallelism (dim2): set to 1 if 1D\n",
    "pp = 8                                 # pipeline parallelism\n",
    "t1 = 4                                 # how many TP GPUs on nvlink\n",
    "t2 = 1                                 # how many TP GPUs on nvlink for dim2 \n",
    "t_dp = 1                               # how many DP GPUs on nvlink\n",
    "t_pp = 1                               # how many PP GPUs on nvlink\n",
    "tp = tp1 * tp2                         # total tensor parallelism (set automatically)\n",
    "dp = total_gpus // (tp * pp)           # data parallelism (set automatically)\n",
    "nm = (global_batch_size // dp) // mbs  # number of microbatches (set automatically)\n",
    "\n",
    "assert t1 * t2 * t_dp * t_pp <= system['nvlink_size'], 'allocated more GPUs to nvlink than available!'\n",
    "print('tp = {}, pp = {}, dp = {}, total = {}, #microbatches = {}'.format(tp,pp,dp,tp*pp*dp,nm))\n",
    "\n",
    "\n",
    "if parallel_strat == '1D':\n",
    "    m1 = tp # quick variable change for consistency\n",
    "    # get all layer stats\n",
    "    df_mlp = mlp_1d(mbs, l, e, f, parallelism={'m': m1}, topology={'t': t1}, system=system)\n",
    "    df_sa = sa_1d(mbs, l, e, h, parallelism={'m': m1}, topology={'t': t1}, flash_attention=True, system=system)\n",
    "    # get the DP and PP comms\n",
    "    df_dp = dataparallel(modules=[df_mlp, df_sa], depth=(depth//pp), dp=dp, t_dp=t_dp, overlap=True, system=system)\n",
    "    p2p_comm_vol = float(df_mlp.loc[df_mlp['name'] == 'ln1']['activation_buffer']) # activation maps that are P2P sent btw GPUs\n",
    "    df_pp = pipelineparallel(modules=[df_mlp, df_sa], number_micro_batches=nm, comm_vol=p2p_comm_vol, pp=pp, t_pp=t_pp, overlap=False, system=system)\n",
    "    \n",
    "elif parallel_strat == '2D-seqp': # context parallel 2D TP\n",
    "    m1 = tp1 # quick variable change for consistency\n",
    "    m2 = tp2\n",
    "    df_mlp = mlp_seqp(mbs, l, e, f, parallelism={'m1': m1, 'm2': m2}, topology={'t1': t1, 't2': t2}, system=system)\n",
    "    df_sa = sa_seqp(mbs, l, e, h, parallelism={'m1': m1, 'm2': m2}, topology={'t1': t1, 't2': t2}, flash_attention=True, system=system)\n",
    "    # DP has some context parallel comms as well\n",
    "    df_dp = dataparallel(modules=[df_mlp, df_sa], depth=(depth//pp), dp=dp*tp2, t_dp=t_dp*t2, overlap=True, system=system)\n",
    "    p2p_comm_vol = float(df_mlp.loc[df_mlp['name'] == 'ln1']['activation_buffer'])\n",
    "    df_pp = pipelineparallel(modules=[df_mlp, df_sa], number_micro_batches=nm, comm_vol=p2p_comm_vol, pp=pp, t_pp=t_pp, overlap=False, system=system)\n",
    "    \n",
    "elif parallel_strat == '2D': # SUMMA 2D TP\n",
    "    m1 = tp1 # quick variable change for consistency\n",
    "    m2 = tp2\n",
    "    df_mlp = mlp_2d(mbs, l, e, f, parallelism={'m1': m1, 'm2': m2}, topology={'t1': t1, 't2': t2}, system=system)\n",
    "    df_sa = sa_2d_seqp(mbs, l, e, h, parallelism={'m1': m1, 'm2': m2}, topology={'t1': t1, 't2': t2}, flash_attention=True, system=system)\n",
    "    df_dp = dataparallel(modules=[df_mlp, df_sa], depth=(depth//pp), dp=dp, t_dp=t_dp, overlap=True, system=system)\n",
    "    p2p_comm_vol = float(df_mlp.loc[df_mlp['name'] == 'ln1']['activation_buffer'])\n",
    "    df_pp = pipelineparallel(modules=[df_mlp, df_sa], number_micro_batches=nm, comm_vol=p2p_comm_vol, pp=pp, t_pp=t_pp, overlap=False, system=system)\n",
    "    \n",
    "else:\n",
    "    assert False, 'parallel strategy not valid!'\n",
    "\n",
    "# print some single mbs layer stats before pipeline and zero sharding corrections\n",
    "print('\\n############## Single microbatch stats ##############')\n",
    "print_df(df_mlp, df_sa)\n",
    "\n",
    "# correct for pipeline bubbles and zero1 sharding\n",
    "(t, mem), stats = totals(df_mlp, df_sa, df_dp, df_pp, depth, pp=pp, dp=dp, number_micro_batches=nm)\n",
    "\n",
    "print('\\n\\n############## Final stats ##############')\n",
    "pprint.pprint(stats)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28e29638-1eb4-4fd8-9f29-8231cc1b9b24",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pytorch-2.0.1",
   "language": "python",
   "name": "pytorch-2.0.1"
  },
  "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.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}