numbers/bfloat16-vs-float16-study.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "black-beijing",
   "metadata": {},
   "source": [
    "# float16 vs bfloat16 numerical properties comparison\n",
    "\n",
    "This a short notebook to help understand `fp16` vs `bfloat16` in particular when converting a model trained\n",
    "in `bfloat16` to mixed precision - it should be possible to look at the numbers to know which ranges\n",
    "are safe and which need to be scaled/avoided.\n",
    "\n",
    "I needed to do that in the context of trying to understand why bfloat16 t5/mt5 models that were pretrained in bfloat16 had a lot of `nan`/`inf` problems when finetuned in mixed precision."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "eastern-variation",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "adult-daughter",
   "metadata": {},
   "source": [
    "This is the main function, that tries to do very simply increments in `bfloat16` and then converting the result to `float16` and showing the discrepancies."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "resistant-chile",
   "metadata": {},
   "outputs": [],
   "source": [
    "def find_mismatch(start, incr):\n",
    "    bf16 = torch.tensor(start, dtype=torch.bfloat16)\n",
    "    print(f\"\\nfp32 start={start:.2e} using increment={incr}\")\n",
    "    print(f\"{'bfloat16':>18} {'float16':>18} {'diff':>8}\")\n",
    "    c = 0\n",
    "    tries = 0\n",
    "    while c < 8:\n",
    "        fp16 = bf16.to(torch.float16)\n",
    "        if not (fp16 == bf16):\n",
    "            print(f\"{bf16:.16f} {fp16:.16f} {torch.sub(fp16.to(dtype=torch.float32), bf16):+.2e}\")\n",
    "            c += 1\n",
    "        bf16 += incr\n",
    "        tries += 1\n",
    "        if tries >= 1e5:\n",
    "            print(f\"gave up finding mismatch after {tries} steps\")\n",
    "            return"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "applied-damages",
   "metadata": {},
   "source": [
    "## Underflow for fp16\n",
    "\n",
    "when numbers become 0.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "cooperative-government",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "fp32 start=1.00e-08 using increment=1e-09\n",
      "          bfloat16            float16     diff\n",
      "0.0000000100117177 0.0000000000000000 -1.00e-08\n",
      "0.0000000110012479 0.0000000000000000 -1.10e-08\n",
      "0.0000000119907781 0.0000000000000000 -1.20e-08\n",
      "0.0000000129803084 0.0000000000000000 -1.30e-08\n",
      "0.0000000139698386 0.0000000000000000 -1.40e-08\n",
      "0.0000000150175765 0.0000000000000000 -1.50e-08\n",
      "0.0000000160653144 0.0000000000000000 -1.61e-08\n",
      "0.0000000171130523 0.0000000000000000 -1.71e-08\n"
     ]
    }
   ],
   "source": [
    "find_mismatch(1e-08, 1e-09)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "decimal-fraction",
   "metadata": {},
   "source": [
    "## Subnormal range for fp16\n",
    "\n",
    "starting from 5.96e-8 \n",
    "\n",
    "usually expensive and very low precision"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "statutory-procurement",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "fp32 start=1.00e-07 using increment=1e-08\n",
      "          bfloat16            float16     diff\n",
      "0.0000001001171768 0.0000001192092896 +1.91e-08\n",
      "0.0000001098960638 0.0000001192092896 +9.31e-09\n",
      "0.0000001201406121 0.0000001192092896 -9.31e-10\n",
      "0.0000001303851604 0.0000001192092896 -1.12e-08\n",
      "0.0000001406297088 0.0000001192092896 -2.14e-08\n",
      "0.0000001508742571 0.0000001788139343 +2.79e-08\n",
      "0.0000001611188054 0.0000001788139343 +1.77e-08\n",
      "0.0000001713633537 0.0000001788139343 +7.45e-09\n"
     ]
    }
   ],
   "source": [
    "# very limited range for fp16\n",
    "find_mismatch(1e-07, 1e-08)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "distributed-puppy",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "fp32 start=1.00e-06 using increment=1e-07\n",
      "          bfloat16            float16     diff\n",
      "0.0000009983778000 0.0000010132789612 +1.49e-08\n",
      "0.0000010952353477 0.0000010728836060 -2.24e-08\n",
      "0.0000012889504433 0.0000013113021851 +2.24e-08\n",
      "0.0000013858079910 0.0000013709068298 -1.49e-08\n",
      "0.0000014826655388 0.0000014901161194 +7.45e-09\n",
      "0.0000015795230865 0.0000015497207642 -2.98e-08\n",
      "0.0000016763806343 0.0000016689300537 -7.45e-09\n",
      "0.0000017732381821 0.0000017881393433 +1.49e-08\n"
     ]
    }
   ],
   "source": [
    "# things starting to improve slightly for fp16\n",
    "find_mismatch(1e-06, 1e-07)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "expired-drinking",
   "metadata": {},
   "source": [
    "## Normal numbers\n",
    "\n",
    "Min positive normal fp16: 6.104e-05 (`np.finfo(np.float16).tiny`)\n",
    "\n",
    "These ranges match much better and thus will not easily find a mismatch if at all"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "seven-caution",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "fp32 start=1.00e-05 using increment=1e-06\n",
      "          bfloat16            float16     diff\n",
      "gave up finding mismatch after 100000 steps\n",
      "\n",
      "fp32 start=1.00e-04 using increment=1e-06\n",
      "          bfloat16            float16     diff\n",
      "gave up finding mismatch after 100000 steps\n",
      "\n",
      "fp32 start=1.00e-03 using increment=0.0001\n",
      "          bfloat16            float16     diff\n",
      "gave up finding mismatch after 100000 steps\n",
      "\n",
      "fp32 start=1.00e-02 using increment=0.001\n",
      "          bfloat16            float16     diff\n",
      "gave up finding mismatch after 100000 steps\n",
      "\n",
      "fp32 start=1.00e-01 using increment=0.01\n",
      "          bfloat16            float16     diff\n",
      "gave up finding mismatch after 100000 steps\n",
      "\n",
      "fp32 start=1.00e+01 using increment=1e-06\n",
      "          bfloat16            float16     diff\n",
      "gave up finding mismatch after 100000 steps\n",
      "\n",
      "fp32 start=1.00e+01 using increment=10.0\n",
      "          bfloat16            float16     diff\n",
      "gave up finding mismatch after 100000 steps\n",
      "\n",
      "fp32 start=1.00e+04 using increment=1\n",
      "          bfloat16            float16     diff\n",
      "gave up finding mismatch after 100000 steps\n"
     ]
    }
   ],
   "source": [
    "find_mismatch(1e-05, 1e-06)\n",
    "find_mismatch(1e-04, 1e-06)\n",
    "find_mismatch(1e-03, 1e-04)\n",
    "find_mismatch(1e-02, 1e-03)\n",
    "find_mismatch(1e-01, 1e-02)\n",
    "find_mismatch(1e1, 1e-06)\n",
    "find_mismatch(1e1, 1e1)\n",
    "find_mismatch(1e4, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "mighty-injection",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "fp32 start=5.00e+04 using increment=1000.0\n",
      "          bfloat16            float16     diff\n",
      "66048.0000000000000000 inf +inf\n",
      "67072.0000000000000000 inf +inf\n",
      "68096.0000000000000000 inf +inf\n",
      "69120.0000000000000000 inf +inf\n",
      "70144.0000000000000000 inf +inf\n",
      "71168.0000000000000000 inf +inf\n",
      "72192.0000000000000000 inf +inf\n",
      "73216.0000000000000000 inf +inf\n"
     ]
    }
   ],
   "source": [
    "# hitting max range for fp16\n",
    "find_mismatch(5e4, 1e3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "alleged-lobby",
   "metadata": {},
   "outputs": [],
   "source": [
    "# --- roundoff ---\n",
    "# fp16 4.88e-4\n",
    "# bf16 3.91e-3"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cheap-natural",
   "metadata": {},
   "source": [
    "## Big numbers\n",
    "\n",
    "`bfloat16` numbers have a terrible range for numbers `> 1` but `fp16` matches those exactly\n",
    "e.g. one can't represent 283 in bf16\n",
    "\n",
    "```\n",
    "python -c \"import torch; print( torch.tensor(283, dtype=torch.bfloat16) )\"\n",
    "tensor(284., dtype=torch.bfloat16)\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "integrated-individual",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "282.00\n",
      "284.00\n",
      "286.00\n"
     ]
    }
   ],
   "source": [
    "start = 280\n",
    "fp32 = torch.tensor(start, dtype=torch.float32)\n",
    "for i in range(3):\n",
    "    bf16 = fp32.to(torch.bfloat16)\n",
    "    bf16d = bf16\n",
    "    while bf16 == bf16d:\n",
    "        fp32 += 1\n",
    "        bf16d = fp32.to(torch.bfloat16)\n",
    "    print(f\"{bf16d:.2f}\")\n",
    "# 282\n",
    "# 284\n",
    "# 286"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae153e44",
   "metadata": {},
   "source": [
    "## How many positions between 2 numbers\n",
    "\n",
    "Let's see how many `fp16` numbers can fit between `bf16` numbers - which should help to understand how converting a model trained in `fp16` to `bf16` in a way quantizes the model - since there are less `bf16` numbers in the same range."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "278e6249",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "9 fp16s: [0.10009765625, 0.10015869140625, 0.1002197265625, 0.10028076171875, 0.100341796875, 0.10040283203125, 0.1004638671875, 0.10052490234375, 0.1005859375]\n",
      "2 bf16s: [0.10009765625, 0.1005859375]\n"
     ]
    }
   ],
   "source": [
    "fp16 = torch.tensor(0.1001, dtype=torch.float16)\n",
    "bf16 = torch.tensor(0.1001, dtype=torch.bfloat16)\n",
    "fp16s = [fp16]\n",
    "bf16s = [bf16]\n",
    "\n",
    "delta = 0.00001\n",
    "for i in range(100):\n",
    "    fp16_new = fp16 + delta*i\n",
    "    bf16_new = bf16 + delta*i\n",
    "    if fp16s[-1] != fp16_new:\n",
    "        fp16s.append(fp16_new)\n",
    "    if bf16s[-1] != bf16_new:\n",
    "        bf16s.append(bf16_new)\n",
    "    if len(bf16s) > 1 and bf16s[-1] == fp16s[-1]:\n",
    "        break\n",
    "        \n",
    "print(f\"{len(fp16s)} fp16s: {[x.item() for x in fp16s]}\")\n",
    "print(f\"{len(bf16s)} bf16s: {[x.item() for x in bf16s]}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fa0ae862",
   "metadata": {},
   "source": [
    "So it can be seen that in this particular range of numbers every 8 \"positions\" in `fp16` get remapped to a single \"position\" in `bf16`. As `exponent(fp16) = 10` and `exponent(bf16) = 7` - so we have `2**3=8` different positions between 2 representations."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "revolutionary-force",
   "metadata": {},
   "source": [
    "# Math"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "systematic-latex",
   "metadata": {},
   "source": [
    "## Summation\n",
    "\n",
    "A very narrow dynamic range means that for largish numbers NN trained in `bfloat16` **expects** bad\n",
    "precision and when the precision is suddenly higher unexpected outcomes happen:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "unlike-raising",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(284., dtype=torch.bfloat16)\n",
      "tensor(283., dtype=torch.float16)\n"
     ]
    }
   ],
   "source": [
    "# small sum\n",
    "print(torch.tensor(282, dtype=torch.bfloat16)+1) # 284\n",
    "print(torch.tensor(282, dtype=torch.float16)+1)  # 283"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "competitive-average",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(2848., dtype=torch.bfloat16)\n",
      "tensor(2830., dtype=torch.float16)\n"
     ]
    }
   ],
   "source": [
    "# sum several of these\n",
    "print(torch.tensor(283, dtype=torch.bfloat16)*10) # 2848\n",
    "print(torch.tensor(283, dtype=torch.float16)*10)  # 2830"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "choice-enemy",
   "metadata": {},
   "source": [
    "As you can see numbers start to diverge quickly!\n",
    "\n",
    "Now in practice we typically add up thousands of numbers.\n",
    "\n",
    "The solution is to always do this kind of operations in double precision of the operands and then if needed casting back to the original. i.e. the accumulate of `sum(fp16_tensor)` should be at least a `float32` tensor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "liquid-purple",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(inf, dtype=torch.float16)"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "tensor(250394.1875)"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = torch.rand((10000)).half()*50\n",
    "\n",
    "# this overflows\n",
    "x.sum()\n",
    "# this succeeds\n",
    "x.sum(dtype=torch.float32)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dated-scope",
   "metadata": {},
   "source": [
    "## Getting overflows\n",
    "\n",
    "Full numbers range: ``float16: ±65,504``\n",
    "\n",
    "So fp16 overflows easily in say variance calculation when you try to just square a number bigger than `256` - as it'd overflow, i.e. you get `inf`! so `256**2+1` will be `inf`\n",
    "\n",
    "You can't even do `pow(2)` for fp16 in pytorch, the following will give an error: that it doesn't suppor power for fp16.\n",
    "\n",
    "`torch.tensor(256, dtype=torch.float16).pow(2)`\n",
    "\n",
    "You have to cast to `float32` first:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "quick-local",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(65024., dtype=torch.float16)"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = torch.tensor(255, dtype=torch.float16)\n",
    "x_squared = x.float().pow(2)\n",
    "x_squared.to(dtype=torch.float16)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "nonprofit-charleston",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(inf, dtype=torch.float16)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# let's cross into the overflow\n",
    "x += 1\n",
    "x_squared = x.float().pow(2)\n",
    "x_squared.to(dtype=torch.float16)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dominican-retro",
   "metadata": {},
   "source": [
    "And that's how `inf` comes about.\n",
    "\n",
    "Or if you need to create one, you can just do:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "returning-certificate",
   "metadata": {},
   "outputs": [],
   "source": [
    "t_inf = torch.tensor(float('inf'))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eastern-orleans",
   "metadata": {},
   "source": [
    "If you need to compare if a tensor has `inf` elements:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "ethical-wayne",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(True)"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.isinf(t_inf).any()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "derived-canberra",
   "metadata": {},
   "source": [
    "## Getting NaNs \n",
    "\n",
    "While there are many ways to get `NaN` during calculations, the most common for machine learning are:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "formed-certificate",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(nan)"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 0/0\n",
    "t_zero = torch.tensor(0)\n",
    "t_zero/t_zero"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "authorized-window",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(nan)"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# inf/inf\n",
    "t_inf = torch.tensor(float('inf'))\n",
    "t_inf/t_inf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "elegant-grant",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(nan)"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 0*inf\n",
    "t_zero * t_inf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "several-council",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(nan)"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# inf - inf\n",
    "t_inf - t_inf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "local-change",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(nan)"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# to get one explicitly\n",
    "t_nan = torch.tensor(float('nan'))\n",
    "t_nan"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "hawaiian-evaluation",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(True)"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# comparison\n",
    "torch.isnan(t_nan).any()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "heated-equivalent",
   "metadata": {},
   "source": [
    "# Debugging process"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "inside-investor",
   "metadata": {},
   "source": [
    "As you can see, since ML is mostly matrix multiplications, which is sums and multiplications, it's enough to get one `inf` or `nan`, and the whole training goes down the rails.\n",
    "\n",
    "Here is a helper that you can run after suspect functions to see if the output gets any `inf` or `nan`s and also if you want to get an indication on whether you have some large numbers that are likely to overflow - remember in fp16 65K is the biggest number one can have."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "alpine-wrist",
   "metadata": {},
   "outputs": [],
   "source": [
    "def detect_overflow(var, ctx):\n",
    "    \"\"\"\n",
    "    Report the count of ``nan`` and ``inf`` entries in the tensor.\n",
    "\n",
    "    This is useful for detecting overflows/underflows and best to call right after the function that did some math that\n",
    "    modified the variable in question.\n",
    "\n",
    "    Args:\n",
    "        var: tensor variable to check\n",
    "        ctx: the message to print as a context\n",
    "    \"\"\"\n",
    "    if torch.isnan(var).any().item():\n",
    "        logger.warning(f\"{ctx} has nans\")\n",
    "    if torch.isinf(var).any().item():\n",
    "        logger.warning(f\"{ctx} has inf\")\n",
    "\n",
    "    # if needed to monitor large elements can enable the following\n",
    "    if 0:\n",
    "        n100 = var[torch.ge(var.abs(), 100)]\n",
    "        if n100.numel() > 0:\n",
    "            logger.warning(f\"{ctx}:  n100={n100.numel()}\")\n",
    "        n1000 = var[torch.ge(var.abs(), 1000)]\n",
    "        if n1000.numel() > 0:\n",
    "            logger.warning(f\"{ctx}: n1000={n1000.numel()}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "vocal-passing",
   "metadata": {},
   "source": [
    "So, if you training gives you say a loss of `nan`, you can go to the layers of your model and inject this function, in one or more places, e.g.:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "exempt-environment",
   "metadata": {},
   "outputs": [],
   "source": [
    "def forward(x):\n",
    "    detect_overflow(x, \"x / enter\")\n",
    "    y = self.ff(x)\n",
    "    detect_overflow(x, \"y / after ff\")              "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "final-tracy",
   "metadata": {},
   "source": [
    "or you use an advanced debugger you can assign watches that will immediately tell you if a tensor just got some `inf`s, by having a dynamically evaluated watch expression: `torch.isinf(x).any().item()` - in this example we watch the tensor `x`. So as you step through the code you can visually immediately see if it went from `False` to `True`. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "caroline-group",
   "metadata": {},
   "source": [
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "imported-empty",
   "metadata": {},
   "source": [
    "# Disabling subnormal numbers in pytorch\n",
    "\n",
    "In some systems subnormal number calculation can be suboptimial as it's often done in software, so if your network deals a lot with subnormal numbers you might want to disable those and scale your numbers to a normal range instead.\n",
    "\n",
    "The following demonstrates how it works in pytorch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "fuzzy-pound",
   "metadata": {
    "run_control": {
     "marked": false
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([0.])"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "tensor([1.0000e-39])"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "_ = torch.set_flush_denormal(True)\n",
    "torch.tensor([1e-39], dtype=torch.float32)\n",
    "_ = torch.set_flush_denormal(False)\n",
    "torch.tensor([1e-39], dtype=torch.float32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "surrounded-upper",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([1.0133e-06], dtype=torch.float16)"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "tensor([1.0133e-06], dtype=torch.float16)"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# broken for fp16\n",
    "_ = torch.set_flush_denormal(True)\n",
    "torch.tensor([1e-6], dtype=torch.float16)\n",
    "_ = torch.set_flush_denormal(False)\n",
    "torch.tensor([1e-6], dtype=torch.float16)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "attempted-deficit",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([0.], dtype=torch.bfloat16)"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "tensor([1.0102e-39], dtype=torch.bfloat16)"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "_ = torch.set_flush_denormal(True)\n",
    "torch.tensor([1e-39], dtype=torch.bfloat16)\n",
    "_ = torch.set_flush_denormal(False)\n",
    "torch.tensor([1e-39], dtype=torch.bfloat16)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "hollywood-karma",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "indie-message",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%javascript # prevent committing an unsaved notebook\n",
    "IPython.notebook.save_notebook()"
   ]
  }
 ],
 "metadata": {
  "hide_input": false,
  "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.8.15"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {
    "height": "calc(100% - 180px)",
    "left": "10px",
    "top": "150px",
    "width": "256.391px"
   },
   "toc_section_display": true,
   "toc_window_display": true
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}