15_OptimizeAIML/15_OptimizeAIML.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "4b665981-c662-466b-b242-12ef9e1487ff",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": []
   },
   "source": [
    "![NCAR UCAR Logo](../NCAR_CISL_NSF_banner.jpeg)\n",
    "# Optimizing AI/ML Workflows in Python for GPUs\n",
    "\n",
    "By: Daniel Howard [dhoward@ucar.edu](mailto:dhoward@ucar.edu), Consulting Services Group, CISL & NCAR \n",
    "\n",
    "Date: August 25th, 2022"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e4c92c10-7da4-4ffa-b5d4-586bfd7e366b",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    },
    "tags": []
   },
   "source": [
    "In this notebook we analyse the overall workflow of typical machine learning/deep learning projects, emphasizing how to work towards optimal performance on GPUs. We will NOT cover theory of or how to implement AI based projects. We will cover:\n",
    "\n",
    "1. Background on machine learning research in Earth sciences\n",
    "2. Setting up Python virtual `conda` environments\n",
    "    * The RAPIDS AI software suite\n",
    "    * GPU enabled TensorFlow and PyTorch\n",
    "3. Enabling tuning and profiling with TensorFlow and PyTorch\n",
    "4. Profiling with DLProf/TensorBoard and performance optimizations for NVIDIA Tensor Cores"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3c8efae1-8e3d-44f0-9f11-6e44bf9f3663",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    },
    "tags": []
   },
   "source": [
    "## Workshop Etiquette\n",
    "* Please mute yourself and turn off video during the session.\n",
    "* Questions may be submitted in the chat and will be answered when appropriate. You may also raise your hand, unmute, and ask questions during Q&A at the end of the presentation.\n",
    "* By participating, you are agreeing to [UCAR’s Code of Conduct](https://www.ucar.edu/who-we-are/ethics-integrity/codes-conduct/participants)\n",
    "* Recordings & other material will be archived & shared publicly.\n",
    "* Feel free to follow up with the GPU workshop team via Slack or submit support requests to [rchelp.ucar.edu](https://support.ucar.edu)\n",
    "    * Office Hours: Asynchronous support via [Slack](https://ncargpuusers.slack.com) or schedule a time with an organizer"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f0f05b32-1a5e-43de-9c5f-c13cf031b22a",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    },
    "tags": []
   },
   "source": [
    "## Start a JupyterHub Session\n",
    "\n",
    "Head to the [NCAR JupyterHub portal](https://jupyterhub.hpc.ucar.edu/stable) and __start a JupyterHub session on Casper PBS Login Node__ and open the notebook at `15_OptimizeAIML/15_OptimizeAIML.ipynb`. Be sure to clone (if needed) and update/pull the NCAR GPU_workshop directory. You are welcome to use an interactive GPU node for the final few cells of this notebook\n",
    "\n",
    "```shell\n",
    "# Use the JupyterHub GitHub GUI on the left panel or the below shell commands\n",
    "git clone git@github.com:NCAR/GPU_workshop.git\n",
    "git pull\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "56bdd4b9-fc2f-4add-bb4a-18876443c44c",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    },
    "tags": []
   },
   "source": [
    "## Notebook Setup\n",
    "__The`GPU_TYPE=gp100` nodes do not have tensor cores!__ Thus, the `gpuworkshop` queue is not as useful for this session. Saying as much, please set `GPU_TYPE=v100` and use the `gpudev` or `casper` queue both during the workshop and for independent work. See [Casper queue documentation](https://arc.ucar.edu/knowledge_base/72581396#StartingCasperjobswithPBS-Concurrentresourcelimits) for more info.  "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "35c0e7e1-15f1-4b29-b507-cfad49e3a18d",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": []
   },
   "source": [
    "# Machine Learning and Deep Learning?\n",
    "\n",
    "ML and DL are statistical models that are designed to learn and predict behavior from a large amount of input training data.\n",
    "\n",
    "<img src=\"img/Venn_AI.jpg\" alt=\"Venn diagram of AI\" style=\"width:600px;\"/>\n",
    "\n",
    "The BAMS article \"[Outlook for Exploiting Artificial Intelligence in the Earth and Environmental Sciences](https://journals.ametsoc.org/view/journals/bams/102/5/BAMS-D-20-0031.1.xml)\" by Boukabara, et al highlights additional applications of AI in the Earth Sciences."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7a17e44c-7814-4551-a09d-d81172ebcad6",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    },
    "tags": []
   },
   "source": [
    "## Overiew of an Earth Science AI Workflow - Remote Sensing\n",
    "\n",
    "Multiple steps are needed to enable AI for Earth Science. __GPUs are critical in the most expensive step, model building and training__, since they perform well with matrix algebra, foundational to ML methods.\n",
    "\n",
    "<img src=\"img/ML_RemoteSensing.png\" alt=\"Conceptual overview of the derivation of digital entities from remotely sensed data using deep-learning techniques\" style=\"width:400px;\"/>\n",
    "\n",
    "Image: [Object Detection and Image Segmentation with Deep Learning on Earth Observation Data: A Review—Part II: Applications](https://www.mdpi.com/2072-4292/12/18/3053/htm) by Hoeser, et al"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2daba0b0-8eaf-4c0d-a88a-58c0b8678ba4",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    },
    "tags": []
   },
   "source": [
    "## Why Use AI for Earth Science?\n",
    "Earth Science is largely built on physics based theories and dynamical interactions with the biosphere. Today, these models have scaled to enormous sizes, consuming significant computational resources and data storage.\n",
    "\n",
    "![Left, E3SM on ANL-Theta, 2022; Right, ECMWF on ORNL-Summit, 2021](img/ECMWF_and_Sat.png)\n",
    "\n",
    "4km global runs of [E3SM](https://e3sm.org/) (left) over 100 forecast years uses 120M core-hours and 250 GB/forecast day, or 12 PB. 1km ECMWF runs (right), as [in this article](https://climatemodeling.science.energy.gov/news/nils-wedi-1-km-resolution-ecmwf-esm-simulation) and by Nils Wedi [keynote at ESMD 2020](https://acme-climate.atlassian.net/wiki/download/attachments/1929707601/WEDI_Nils_D1S1_2020-1026.pdf?api=v2).\n",
    "\n",
    "__AI offers an opportunity to reduce computational resources required__. Feel free to consult [A Review of Earth Artificial Intelligence](https://www.sciencedirect.com/science/article/pii/S0098300422000036) for current \"Grand Challenges\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b29d8f54-f229-43f8-8292-e66b40f75a14",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    },
    "tags": []
   },
   "source": [
    "## Surrogate Models\n",
    "\n",
    "Novel ways can be explored to use Earth Science data to reduce required computational resources. A __surrogate model__ in machine learning is a __statistical model__ designed to more efficiently approximate the output of a physics based model.\n",
    "\n",
    "<img src=\"img/Surrogate_Modeling.png\" alt=\"Surrogate Modeling, Guo\" style=\"width:650px;\"/>\n",
    "\n",
    "Image: [Introduction to Surrogate Modeling](https://towardsdatascience.com/an-introduction-to-surrogate-modeling-part-i-fundamentals-84697ce4d241), Shuai Guo. See \"[Learning Nonlinear Dynamical Systems from Data Using Scientific Machine Learning](https://anl.app.box.com/s/hvrk3t8qpg2u218ynggjn87jw1864h12)\" by Maulik, ANL."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "792b68a6-1fb9-4fd1-8982-d62103564cc5",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    },
    "tags": []
   },
   "source": [
    "## Neural Ordinary Differential Equations\n",
    "\n",
    "For example, a stabilized neural ODE can be designed to accurately simulate shocks and chaotic dynamics.\n",
    "\n",
    "![Neural ODE with Neural Networks](img/Neural_ODE.png)\n",
    "\n",
    "See paper by Linot, et al \"[Stabilized Neural Ordinary Differential Equations for Long-Time Forecasting of Dynamical Systems](https://arxiv.org/abs/2203.15706)\"."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d8df2a2c-9a65-4f22-a050-38a40fde35dc",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    },
    "tags": []
   },
   "source": [
    "## Physics Informed Neural Networks (PINNs)\n",
    "Other applications to consider are __Physics Informed Neural Networks__. PINNs attempt to embed known physics relationships into the design of a machine learning model. This may include defining the Navier-Stokes conservation laws as conditions to minimize in a ML model's loss function.\n",
    "\n",
    "<img src=\"img/PINN_NS.png\" alt=\"Physics Informed Neural Networks with Navier Stokes\" style=\"width:700px;\"/>\n",
    "\n",
    "Image: [Wikipedia - Physics Informed Nearal Networks](https://en.wikipedia.org/wiki/Physics-informed_neural_networks)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d4da771-b506-4e8d-afcb-373f666569bd",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    },
    "tags": []
   },
   "source": [
    "## Resources for Engaging and Learning AI in Earth Sciences\n",
    "\n",
    "Feel free to reach out to [rchelp@ucar.edu](mailto:rchelp@ucar.edu) if you want assistance recreating environments for any below code examples.\n",
    "\n",
    "1. OLCF AI 4 Science Fluid Flow Tutorial ([GitHub](https://github.com/muralikrishnangm/tutorial-ai4science-fluidflow)) - Uses [MiniWeatherML](https://github.com/mrnorman/miniWeatherML)\n",
    "2. OpenHackathons GPU Bootcamp ([GitHub](https://github.com/openhackathons-org/gpubootcamp/)) - [HPC AI Examples](https://github.com/openhackathons-org/gpubootcamp/tree/master/hpc_ai) for PINNs, CFD, and Climate\n",
    "3. NSF AI Institute for Research on Trustworthy AI in Weather, Climate, and Coastal Oceanography ([AI2ES.org](https://www.ai2es.org/)) - [Education Materials](https://www.ai2es.org/products/education/) and [2022 Trust-a-thon GitHub](https://github.com/ai2es/tai4es-trustathon-2022)\n",
    "4. Argonne ALCF\n",
    "   * [2021 Simulation, Data, and Learning Workshop for AI](https://www.alcf.anl.gov/events/2021-alcf-simulation-data-and-learning-workshop) ([GitHub](https://github.com/argonne-lcf/sdl_ai_workshop))- Detailed [DL profiling tutorial notebooks](https://github.com/argonne-lcf/sdl_ai_workshop/tree/master/04_profilingDeepLearning) & [video](https://www.youtube.com/watch?v=cdLIlOUnRCU)\n",
    "   * [2022 Introduction to AI-driven Science on Supercomputers](https://www.alcf.anl.gov/alcf-ai-science-training-series) ([GitHub](https://github.com/argonne-lcf/ai-science-training-series))\n",
    "5. Data Driven Atmospheric and Water Dynamics Beucler Lab (U. of Lausanne - Switzerland)\n",
    "   * [Getting Started with Machine Learning](https://wp.unil.ch/dawn/getting-started-with-machine-learning/) curated resource list\n",
    "6. [NOAA Workshop on Leveraging Artificial Intelligence in Environmental Sciences](https://www.noaa.gov/ai/events/4th-noaa-ai-workshop) - 4th Workshop free to register, virtual Sept 6-9 2022\n",
    "7. National Academies - 2022 workshop [Machine Learning and Artificial Intelligence to Advance Earth System Science: Opportunities and Challenges](https://www.nationalacademies.org/event/02-07-2022/machine-learning-and-artificial-intelligence-to-advance-earth-system-science-opportunities-and-challenges-a-workshop)\n",
    "8. [Climate Informatics](http://www.climateinformatics.org/) community - [Conferences](http://www.climateinformatics.org/conferences/) and [Hackathons](http://www.climateinformatics.org/hackathons/)\n",
    "9. Book - [Deep learning for the Earth Sciences -- A comprehensive approach to remote sensing, climate science and geosciences](https://github.com/DL4ES/DL4ES)\n",
    "10. [climatechange.ai](https://www.climatechange.ai/) - Global initiative to catalyze impactful work at the intersection of climate change and machine learning."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "926bf02f-a906-486a-9db9-cc2e1b027ffc",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": []
   },
   "source": [
    "# How to Manage Python Software for ML and DL Models\n",
    "The Python ecosystem already provides many robust pre-built software packages and libraries which are continually maintained. __Learning about and employing the Python ecosystem well can simplify the process of using machine learning tools__. \n",
    "\n",
    "The kernel `GPU_Workshop` already has many useful packages plus others (notably [Horovod](https://horovod.ai/) for distributed deep learning) which you are welcome to explore on your own beyond this workshop. \n",
    "\n",
    "Run the below cell to get a listing of all packages installed in the `GPU_Workshop` conda environment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "75d124be-a2c0-469e-90c4-4b4c072b1dd3",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "!mamba list -p /glade/work/dhoward/conda/envs/GPU_Workshop/"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bc93f803-db47-43f0-9145-29d5d2efbffd",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    },
    "tags": []
   },
   "source": [
    "## Setting Up Conda Environments\n",
    "\n",
    "Since ensuring compatibility and reproducibility is difficult across python package environments, __you are encouraged to maintain your own personalized `conda` virtual environments__. Nonetheless, NCAR provides a base set of commonly used Python packages via the [NCAR Package Library (NPL)](https://arc.ucar.edu/knowledge_base/83853599). NPL does include the faster package management tool `mamba` which uses the same command syntax as `conda`.\n",
    "\n",
    "If you prefer to install your own and not use `module load conda`, we encourage [Mambaforge](https://github.com/conda-forge/miniforge). In general, `mamba` is safe to use compared to `conda`. To update all non-pinned packages in an environment, you can use `mamba update --all`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f5d988aa-bc70-41a1-a6dc-a2506b287f1c",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    },
    "tags": []
   },
   "source": [
    "## Choosing Conda Channels\n",
    "\n",
    "To source packages, the channel `conda-forge` is recommended and set as priority on Casper but other channels you may consider are `ncar`, `nvidia`, `rapidsai`, `intel`, `pytorch`, and `anaconda` among others. \n",
    "\n",
    "* Learn to manage channels [here](https://docs.conda.io/projects/conda/en/latest/user-guide/tasks/manage-channels.html) using your `$HOME/.condarc` file\n",
    "* Define pinned packages, ie packages that should stay at a specific version or use a specific build type, via the `/path/to/env/conda-meta/pinned` file"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "504c6a68-a7a7-49b0-a672-86900bb55320",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": []
   },
   "source": [
    "## RAPIDS AI Environment\n",
    "\n",
    "`rapidsai` channel provides [RAPIDS](https://rapids.ai/about.html), an open source, NVIDIA maintained suite for end-to-end data science and analytics pipelines on GPUs. Feel free to explore RAPIDS [Getting Started Notebooks](https://docs.rapids.ai/start).\n",
    "\n",
    "<img src=\"img/ScaleUpOut_RAPIDS.png\" alt=\"Scale Up and Out in Python\" style=\"width:700px;\"/>\n",
    "\n",
    "__Scale Up with RAPIDS__ tools and __Scale Out with Dask/UCX or Horovod__ tools."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0305b861-4390-4436-b7f8-9b5500739254",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    },
    "tags": []
   },
   "source": [
    "### Python Packages and RAPIDS Equivalents\n",
    "\n",
    "![Python Packages and RAPIDS Equivalents](img/PythonPackages_CPUtoGPU.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4927cb7b-6b9e-43b4-a952-5e4dac6887fe",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    },
    "tags": []
   },
   "source": [
    "### Install RAPIDS environment\n",
    "\n",
    "Setting flexible channel priority via `conda config --set channel_priority flexible` or in `~/.condarc`, follow install directions [here](https://rapids.ai/start.html) or by running:\n",
    "\n",
    "```shell\n",
    "conda create -n rapids-22.08 -c rapidsai -c nvidia -c conda-forge  \\\n",
    "rapids=22.08 python=3.9 cudatoolkit=11.5\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a39c763-8e48-4e51-b48a-d406ce16d066",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    },
    "tags": []
   },
   "source": [
    "## Installing Customized Python Packages\n",
    "\n",
    "For more personalized environments, an example process to setup a `conda` environment on Casper is below: \n",
    "```shell\n",
    "module load conda\n",
    "# Creates environment in /glade/work/$USER/conda-envs/my-env-name or a fully specified path\n",
    "mamba create -n my-env-name\n",
    "mamba activate my-env-name\n",
    "\n",
    "# The Python version installed here will automatically be pinned\n",
    "# Recommend to not use the latest Python version (3.10+) given compatibility issues\n",
    "mamba install python=3.9*\n",
    "\n",
    "# Ensures we get MKL optimized packages to run on Casper's Intel CPUs\n",
    "mamba install numpy scipy pandas scikit-learn xarray \"libblas=*=*mkl\"\n",
    "# Ensures common packages provide MPI support (typically defaults to OpenMPI). \n",
    "# Useful to pin packages in `/path/to/env/conda-meta/pinned` file.\n",
    "mamba install mpi4py fftw=*=mpi* h5py=*=mpi* netcdf4=*=mpi* \n",
    "```\n",
    "\n",
    "To highlight, adding `<package-name>=<version>=<build-type>` is important to ensure you install the most relevant and performant version for your needs.\n",
    "\n",
    "For example, `libblas=*=*mkl` guarantees you get the Intel MKL optimized versions of packages that utilize the BLAS library. The `*` is a wildcard for the latest version or other build specifications/hashes."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1725a184-51c4-44ed-9ade-911f2da7f254",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    },
    "tags": []
   },
   "source": [
    "## GPU Enabled Python Packages and Tools\n",
    "ML libraries [pytorch](https://pytorch.org/get-started/locally/) and [tensorflow](https://www.tensorflow.org/install/pip) require additional steps to ensure they are installed with GPU support.\n",
    "\n",
    "```shell\n",
    "mamba install cudatoolkit cudnn cupy nvtx\n",
    "# Make sure package wheel ID includes *cuda* to verify GPU support\n",
    "mamba install pytorch=1.12.1=cuda112*\n",
    "# Don’t use tensorflow-gpu package as package solver is inconsistent in condo-forge channel\n",
    "# TF recommends pip install for latest official version but conda-forge versions also work\n",
    "mamba install tensorflow=2.9.1=cuda112*\n",
    "\n",
    "# Enables added profiling capabilities, only available via pip and PyPI or NVIDIA's package index\n",
    "pip install nvidia-pyindex\n",
    "pip install nvidia-dlprof nvidia-dlprof-pytorch-nvtx\n",
    "pip install tensorboard_plugin_profile\n",
    "```\n",
    "\n",
    "Each library's documentation linked above has more info about installation options. As of this workshop, TensorFlow guarantees support up to CUDA v11.2 and PyTorch up to CUDA v11.6 so we specified builds with `=cuda112*`. Run `mamba search <package>` to view all available packages given available channels.\n",
    "\n",
    "TensorFlow recommends installation via `pip` for their offical versions but the community does tend to maintain similar quality releases via `conda-forge`. Combining `pip` with `conda/mamba` installs should be avoided if possible due to greater difficulty in maintaining environments."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "131f0165-88ac-4078-8e76-a990a9064bb8",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    },
    "tags": []
   },
   "source": [
    "### Horovod for Distributed Deep Learning\n",
    "\n",
    "For distributed deep learning with [Horovod](https://horovod.ai/) instead of Dask, see below or [Horovod installation documentation](https://horovod.readthedocs.io/en/stable/install_include.html) for how to use pip to install Horovod from PyPI on Casper.\n",
    "```shell\n",
    "module load cuda/11.7 gnu/10.1.0\n",
    "mamba install pip gxx_linux-64 cmake nccl\n",
    "export HOROVOD_NCCL_HOME=$CONDA_PREFIX\n",
    "export HOROVOD_CUDA_HOME=$CUDA_HOME\n",
    "HOROVOD_GPU_OPERATIONS=NCCL pip install horovod[tensorflow,keras,pytorch]\n",
    "horovodrun --check-build\n",
    "```\n",
    "\n",
    "Note the specification of `HOROVOD_GPU_OPERATIONS=NCCL` to use NVIDIA's Collective Communication Library. An `MPI` option is also selectable for CUDA-aware MPI libraries. Find more details about Horovod's GPU tensor operations and [GPU install options here](https://github.com/horovod/horovod/blob/master/docs/gpus.rst).\n",
    "\n",
    "A useful tutorial for Horovod was given as part of the [Argonne Training Program on Extreme-Scale Computing](https://extremecomputingtraining.anl.gov/agenda-2022/) (ATPESC) - [Data Parallel Deep Learning](https://anl.app.box.com/s/ujkvbb8glmq7n6gzjhxza7vx7n3wa86u)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0c65eeb2-ff7a-40a0-a23d-b7e87291c21d",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    },
    "tags": []
   },
   "source": [
    "## Sharing Package Environments\n",
    "Once your environment is setup, you can share or give access to your Python virtual environments, which is vitally important to consider towards __enabling reproducible science__.\n",
    "\n",
    "1. On a shared cluster, share a path to your environment, see `mamba env list`. Make sure you provide read access plus write access if you want others to be able to modify the environment. Then run `mamba activate /path/to/env`\n",
    "2. Others may instead clone a readable environment with `mamba create --name cloned_env --clone /path/to/original_env`\n",
    "3. To distribute your environment, run `mamba env export > my-env.yml`. Others can then install this environment with `mamba env create -f /path/to/yaml-file`"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1fb6913d-045d-4b38-b5e1-39aca3d18df5",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": []
   },
   "source": [
    "# Running a Profiler on TensorFlow and PyTorch Models\n",
    "Both `tensorflow` and `pytorch` have built in tools and `tensorboard` GUI interface for DL profiling, which typically run profiles during the training portion of a deep learning model. Base guides for using these built-in tools follow:\n",
    "\n",
    "* PyTorch\n",
    "    * [Profiler Tutorial](https://pytorch.org/tutorials/recipes/recipes/profiler_recipe.html)\n",
    "    * [Building a Benchmark Tutorial](https://pytorch.org/tutorials/recipes/recipes/benchmark.html)\n",
    "    * [PyTorch Profiler with TensorBoard Tutorial](https://pytorch.org/tutorials/intermediate/tensorboard_profiler_tutorial.html)\n",
    "* TensorFlow\n",
    "    * [TensorFlow Profiler Guide](https://www.tensorflow.org/guide/profiler) \n",
    "    * [TensorBoard Profiler Analysis Guide](https://www.tensorflow.org/tensorboard/tensorboard_profiling_keras)\n",
    "* TensorBoard - [Callbacks API Class](https://keras.io/api/callbacks/tensorboard/)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d6e7348-8f48-4a07-ba24-6350bf137c06",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    },
    "tags": []
   },
   "source": [
    "## Easy Ways to Implement TensorFlow and PyTorch Profilers\n",
    "__PyTorch__\n",
    "\n",
    "```python\n",
    "model = models.resnet18().cuda()\n",
    "inputs = torch.randn(5, 3, 224, 224).cuda()\n",
    "\n",
    "with profile(activities=[\n",
    "        ProfilerActivity.CPU, ProfilerActivity.CUDA], record_shapes=True) as prof:\n",
    "    with record_function(\"model_inference\"):\n",
    "        model(inputs)\n",
    "\n",
    "print(prof.key_averages().table(sort_by=\"cuda_time_total\", row_limit=10))\n",
    "```\n",
    "\n",
    "The `record_shapes` parameter ensures the profiler collects data on the data pipeline types, notably tensor shapes.\n",
    "\n",
    "__TensorFlow__ - See [API](https://www.tensorflow.org/api_docs/python/tf/profiler/experimental/start) for additional options\n",
    "\n",
    "```python\n",
    "import tensorflow as tf\n",
    "\n",
    "tf.profiler.experimental.start('/path/to/log/output/')\n",
    "\n",
    "# ... training loop ...\n",
    "\n",
    "tf.profiler.experimental.stop()\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ea9c2692-db71-4902-9b6a-0bcf4968aee5",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    },
    "tags": []
   },
   "source": [
    "## Using NVIDIA Tools for Profiling DL Models\n",
    "\n",
    "The tools `nsys` and `ncu` are similarly adaptable to run against DL Python codes. The [`dlprof` tool](https://docs.nvidia.com/deeplearning/frameworks/dlprof-user-guide/) was previously developed to run `nsys` on DL models then output a TensorBoard interface. However, `dlprof` is no longer being developed in favor of the previous built in profiling methods.\n",
    "\n",
    "* PyTorch\n",
    "    * DNN Layer annotations are disabled by default\n",
    "    * Use `with torch.autograd.profiler.emit_nvtx():`\n",
    "    * Manually with `torch.cuda.nvtx.range_(push/pop)`\n",
    "    * TensorRT backend is already annotated\n",
    "* TensorFlow\n",
    "    * Annotated by default with NVTX, _only in [NVIDIA NGC containers](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/tensorflow) or `nvidia-pyindex` TF 1.X containers_\n",
    "        * `export TF_DISABLE_NVTX_RANGES=1` to disable for production\n",
    "    * For TensorFlow 2.X, must manually inline NVTX ranges or use `dlprof --mode=tensorflow2 ...`\n",
    "    \n",
    "NVIDIA provides their own guides, such as [NVIDIA Deep Learning Performance](https://docs.nvidia.com/deeplearning/performance/index.html). A small example using the `nsys`/`ncu` tools and `dlprof` with DL models can be found [here](https://github.com/argonne-lcf/sdl_ai_workshop/tree/master/04_profilingDeepLearning/NvidiaProfiler). `dlprof` can still work well in NVIDIA [NGC Containers](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/tensorflow) but compatibility elsewhere is not well supported."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8c11d95c-76a3-4d3a-b9a2-7422c49d81c3",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    },
    "tags": []
   },
   "source": [
    "## Common Performance Considerations\n",
    "1. I/O\n",
    "    * Use designated TF/PT data loaders\n",
    "        * TensorFlow - [Better Performance with the `tf.data` API](https://www.tensorflow.org/guide/data_performance)\n",
    "        * PyTorch - [Datasets & Dataloaders](https://pytorch.org/tutorials/beginner/basics/data_tutorial.html)\n",
    "    * Multithreading, eg [Multi-Worker Training with Keras](https://www.tensorflow.org/tutorials/distribute/multi_worker_with_keras)\n",
    "2. CPU to/from GPU data copies\n",
    "    * Rewrite code with TF/PT tensors or use CuPy, etc\n",
    "    * Overlap copy and computation\n",
    "3. Batch size - Increase batch size up to GPU is saturated\n",
    "4. Precision (Background: See Theo Mary's [Mixed Precision Artithmetic](https://www.youtube.com/watch?v=9ZnwfPvAlHM) talk at London Math Society)\n",
    "    * Consider mixed precision, [NVIDIA Mixed Precision Training Guide](https://docs.nvidia.com/deeplearning/performance/mixed-precision-training/index.html)\n",
    "    * Automatic Mixed Precision (AMP) settings\n",
    "        * [PT Guide](https://pytorch.org/docs/stable/notes/amp_examples.html): `scaler = torch.cuda.amp.GradScaler()`\n",
    "        * [TF Guide](https://www.tensorflow.org/guide/mixed_precision): `policy = mixed_precision.Policy('mixed_float16'); mixed_precision.set_global_policy(policy)`\n",
    "    * Ensure usage of Tensor Cores with Mixed Precision\n",
    "    \n",
    "TensorFlow provides a comprehensive guide, [Optimize TensorFlow GPU performance with the TensorFlow Profiler](https://www.tensorflow.org/guide/gpu_performance_analysis)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a28c1611-bac6-4bd7-a01d-10e0aa01a055",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    },
    "tags": []
   },
   "source": [
    "### Performance Improvements with Tensor Cores\n",
    "\n",
    "![Tensor Core Arithmetic and Estimated Throughput](img/Tensorcore.png)\n",
    "\n",
    "Per NVIDIA's recommendation on [Optimizing for Tensor Cores](https://docs.nvidia.com/deeplearning/performance/mixed-precision-training/index.html#opt-tensor-cores), setting parameters such as matrix dimension sizes, batch sizes, convolution layer channel counts, etc. as __multiples of 8__ is optimal due to tensor core shape constraints.\n",
    "\n",
    "Utilizing mixed precision and tensor cores effectively can lead to [theoretical throughput performance](https://3s81si1s5ygj3mzby34dq6qf-wpengine.netdna-ssl.com/wp-content/uploads/2020/05/nvidia-pascal-volta-ampere-comparison-table.jpg) of __9.70 TeraFLOPS__ for FP64 arithmetic up to __78.0 TeraFLOPS__ for FP16 arithmetic on A100 GPUs."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1a3aa1f0-0b5d-4252-b968-c270ffde2bd3",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": []
   },
   "source": [
    "# Profiler Runs of a Geomagnetic Field LSTM Model\n",
    "This Long Short-Term Memory (LSTM) example comes courtesy of the [Trustworthy AI for Environmental Science Trust-a-thon](https://github.com/ai2es/tai4es-trustathon-2022/tree/main/space). You can follow the original example, with data preparation and explanation of how the LSTM model is implemented in the [source notebook](https://github.com/ai2es/tai4es-trustathon-2022/blob/main/space/magnet_lstm_tutorial.ipynb).\n",
    "\n",
    "To begin, let's first download data to use for training and validation of our LSTM model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ded53087-fc96-4f8e-ac7e-72f83e724c97",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%capture captured_io\n",
    "%%bash\n",
    "\n",
    "# Download data we need. If a directory \"data/\" already exists, we'll assume the data are already downloaded.\n",
    "#      The above \"magic\" statements are used to capture shell in/out and to run the following Bash commands.\n",
    "if [ ! -d \"data\" ]; then\n",
    "  wget --verbose https://ngdc.noaa.gov/geomag/data/geomag/magnet/public.zip\n",
    "  wget --verbose https://ngdc.noaa.gov/geomag/data/geomag/magnet/private.zip\n",
    "  unzip public.zip\n",
    "  unzip private.zip\n",
    "  mkdir -v data\n",
    "  mv -v public private data/\n",
    "  mv -v public.zip private.zip data/\n",
    "fi\n",
    "# Uncomment for debugging if you have trouble downloading:\n",
    "#print(captured_io)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3cf02223-47ee-4765-b0e0-e83f3fe93ae3",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    },
    "tags": []
   },
   "source": [
    "## Profile the `magnet_lstm_tutorial.py` Python Script\n",
    "The full Geomagnetic Field LSTM model is condensed into the Python file [magnet_lstm_tutorial.py](magnet_lstm_tutorial.py). Recall that profiling does not require analyzing the full runtime of most models. In DL, most operations are highliy repetitive so the profiler only needs to sample a small portion of the runtime. __Minimizing the time for profile runs can speed up the iterative development process__.\n",
    "\n",
    "* __TODO Line 237__: Adjust parameter `n_epochs=1` in order to minimize profiling time.\n",
    "* __TODO Line 295/301__: Add TensorBoard callbacks as defined below.\n",
    "\n",
    "```python\n",
    "tboard_callback = keras.callbacks.TensorBoard(\n",
    "    log_dir = \"profile_results\", histogram_freq = 1, profile_batch = '500,520')\n",
    "\n",
    "...\n",
    "\n",
    "model.fit(\n",
    "    ...,\n",
    "    callbacks = [tboard_callback]\n",
    ")\n",
    "```\n",
    "\n",
    "__Question__: How else could you minimize runtime of a \"profile run\" but still maintain model configuration parameters equivalent to production runs?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5eab299f-febe-4e5c-886c-d06e9d6f4660",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%bash\n",
    "qsub pbs_job.sh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "643e5ee4-6342-41cd-be99-475d7969c2b8",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%bash\n",
    "# Run this cell if in an interactive GPU node job\n",
    "#module load cuda/11.7 &> /dev/null\n",
    "#python magnet_lstm_tutorial.py"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b39ef77a-18de-40ca-9b20-36ed20884bd5",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    },
    "tags": []
   },
   "source": [
    "## Open the Profile Report in TensorBoard\n",
    "Typically you'll run this on the login node of Casper and will need to do some `ssh` port forwarding to access the server. You can generally follow these steps:\n",
    "\n",
    "1. `ssh -L$PORTA:localhost:$PORTB $USER@casper.ucar.edu`\n",
    "2. `module load conda`\n",
    "3. `mamba activate /glade/work/dhoward/conda/env/GPU_Workshop`\n",
    "3. `cd /path/to/log/output`\n",
    "4. `tensorboard --port $PORTB --bind_all --logdir </path/to/log/output/>`  and wait for message that says server has started\n",
    "5. Open browser on your laptop and go to `localhost:$PORTA`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5cf61e0e-264b-4350-9c24-58a5d770d368",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# If on a local machine with GPU, use these commands to open the profile.\n",
    "# Otherwise, port forwarding is needed on Casper\n",
    "# %load_ext tensorboard\n",
    "# %tensorboard --logdir=profile_results"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4fb12bec-2418-4752-b16b-10dab298afd0",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": []
   },
   "source": [
    "## Analyizing Profiles in TensorBoard\n",
    "\n",
    "<img src=\"img/TensorBoard_PerformanceSummary.png\" alt=\"TensorBoard Performance Summary\" style=\"width:900px;\"/>\n",
    "\n",
    "Performance improvement heuristics are often provided with links to more detailed information."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "73a677dd-ecde-41fd-ae29-a57f7a29b196",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    },
    "tags": []
   },
   "source": [
    "![TensorBoard Kernel Stats](img/TensorBoard_KernelStats.png)\n",
    "\n",
    "Important to emphasize __where Tensor Core use is eligible__ in your model and determine if it's appropriate for employing reduced precision."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1153c6b3-f896-4aea-9a55-82e7f6869dc3",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    },
    "tags": []
   },
   "source": [
    "![TensorBoard Trace Timeline](img/TensorBoard_Trace.png)\n",
    "\n",
    "Use the `trace_viewer` to get an overall timeline of your deep learning model."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "GPU Workshop",
   "language": "python",
   "name": "pygpu"
  },
  "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}