Merge branch 'main' into activation-ops

Author: timmoon10

.github/workflows/docs.yml

@@ -17,8 +17,8 @@ jobs:
         uses: actions/checkout@v3
       - name: 'Install dependencies'
         run: |
-          pip install sphinx==5.1.1 sphinx_rtd_theme==1.0.0 nbsphinx==0.8.10 IPython ipython_genutils==0.2.0 ipywidgets==8.0.2 astroid==2.15.7
-          pip install breathe==4.34.0 sphinx-autoapi==2.0.1
+          pip install sphinx==8.1.3 sphinx_rtd_theme==3.0.1 nbsphinx==0.9.5 IPython ipython_genutils==0.2.0 ipywidgets==8.0.2 astroid==3.3.2
+          pip install breathe==4.35.0 sphinx-autoapi==3.3.2
           sudo apt-get install -y pandoc graphviz doxygen
           export GIT_SHA=$(git show-ref --hash HEAD)
       - name: 'Build docs'

.gitignore

@@ -22,9 +22,7 @@ __pycache__
 .hypothesis
 .devcontainer.json
 tests/cpp/build/
-docs/_build
 .ipynb_checkpoints
-docs/doxygen
 *.log
 CMakeFiles/CMakeSystem.cmake
 sdist/
@@ -40,3 +38,4 @@ dist/
 downloads/
 .pytest_cache/
 compile_commands.json
+.nfs

3rdparty/cudnn-frontend

@@ -11,7 +11,6 @@
 from .utils import (
     all_files_in_dir,
     cuda_archs,
-    cuda_path,
     cuda_version,
 )
 
@@ -29,9 +28,6 @@ def setup_pytorch_extension(
     sources = [
         csrc_source_files / "common.cu",
         csrc_source_files / "ts_fp8_op.cpp",
-        csrc_source_files / "userbuffers" / "ipcsocket.cc",
-        csrc_source_files / "userbuffers" / "userbuffers.cu",
-        csrc_source_files / "userbuffers" / "userbuffers-host.cpp",
     ] + all_files_in_dir(extensions_dir)
 
     # Header files
@@ -85,19 +81,14 @@ def setup_pytorch_extension(
                 continue  # Already handled
             nvcc_flags.extend(["-gencode", f"arch=compute_{arch},code=sm_{arch}"])
 
-    # Libraries
-    library_dirs = []
-    libraries = []
-    if bool(int(os.getenv("NVTE_UB_WITH_MPI", 0))):
+    if bool(int(os.getenv("NVTE_UB_WITH_MPI", "0"))):
         assert (
             os.getenv("MPI_HOME") is not None
-        ), "MPI_HOME must be set when compiling with NVTE_UB_WITH_MPI=1"
-        mpi_home = Path(os.getenv("MPI_HOME"))
-        include_dirs.append(mpi_home / "include")
+        ), "MPI_HOME=/path/to/mpi must be set when compiling with NVTE_UB_WITH_MPI=1!"
+        mpi_path = Path(os.getenv("MPI_HOME"))
+        include_dirs.append(mpi_path / "include")
         cxx_flags.append("-DNVTE_UB_WITH_MPI")
         nvcc_flags.append("-DNVTE_UB_WITH_MPI")
-        library_dirs.append(mpi_home / "lib")
-        libraries.append("mpi")
 
     # Construct PyTorch CUDA extension
     sources = [str(path) for path in sources]
@@ -112,6 +103,4 @@ def setup_pytorch_extension(
             "cxx": cxx_flags,
             "nvcc": nvcc_flags,
         },
-        libraries=[str(lib) for lib in libraries],
-        library_dirs=[str(lib_dir) for lib_dir in library_dirs],
     )

build_tools/pytorch.py

@@ -0,0 +1,3 @@
+_build
+doxygen
+sphinx_rtd_theme

docs/.gitignore

@@ -16,5 +16,10 @@ help:
 
 # Catch-all target: route all unknown targets to Sphinx using the new
 # "make mode" option.  $(O) is meant as a shortcut for $(SPHINXOPTS).
-%: Makefile
-	@$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
+%: Makefile sphinx_rtd_theme
+	PYTHONPATH=sphinx_rtd_theme:$(PYTHONPATH) $(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
+
+# Patch Sphinx RTD theme 3.0.1 to add version selector in sidebar
+sphinx_rtd_theme:
+	git clone --depth=1 -b 3.0.1 --single-branch https://github.com/readthedocs/sphinx_rtd_theme.git
+	bash -c "cd sphinx_rtd_theme; git apply ../version_select.patch"

docs/Makefile

@@ -51,3 +51,7 @@ pyTorch
 .. autoapifunction:: transformer_engine.pytorch.moe_permute
 
 .. autoapifunction:: transformer_engine.pytorch.moe_unpermute
+
+.. autoapifunction:: transformer_engine.pytorch.initialize_ub
+
+.. autoapifunction:: transformer_engine.pytorch.destroy_ub

docs/api/pytorch.rst

@@ -2,38 +2,30 @@
 #
 # See LICENSE for license information.
 
+import datetime
 import os
-import sys
-import sphinx_rtd_theme
-from sphinx.ext.autodoc.mock import mock
-from sphinx.ext.autodoc import between, ClassDocumenter, AttributeDocumenter
-from sphinx.util import inspect
-from builtins import str
-from enum import Enum
-import re
+import pathlib
 import subprocess
-from pathlib import Path
-from datetime import date
-
-te_path = os.path.dirname(os.path.realpath(__file__))
+from builtins import str
 
-with open(te_path + "/../build_tools/VERSION.txt", "r") as f:
-    te_version = f.readline().strip()
+# Basic project info
+project = "Transformer Engine"
+author = "NVIDIA CORPORATION & AFFILIATES"
 
+# Copyright statement
 release_year = 2022
-
-current_year = date.today().year
+current_year = datetime.date.today().year
 if current_year == release_year:
     copyright_year = release_year
 else:
     copyright_year = str(release_year) + "-" + str(current_year)
+copyright = f"{copyright_year}, NVIDIA CORPORATION & AFFILIATES. All rights reserved."
 
-project = "Transformer Engine"
-copyright = "{}, NVIDIA CORPORATION & AFFILIATES. All rights reserved.".format(copyright_year)
-author = "NVIDIA CORPORATION & AFFILIATES"
+# Transformer Engine root directory
+root_path = pathlib.Path(__file__).resolve().parent.parent
 
+# Git hash
 git_sha = os.getenv("GIT_SHA")
-
 if not git_sha:
     try:
         git_sha = (
@@ -44,31 +36,16 @@
         )
     except:
         git_sha = "0000000"
-
 git_sha = git_sha[:7] if len(git_sha) > 7 else git_sha
 
-if "dev" in te_version:
-    version = str(te_version + "-" + git_sha)
+# Version
+with open(root_path / "build_tools" / "VERSION.txt", "r") as f:
+    _raw_version = f.readline().strip()
+if "dev" in _raw_version:
+    version = str(_raw_version + "-" + git_sha)
 else:
-    version = str(te_version)
-release = te_version
-
-# hack: version is used for html creation, so put the version picker
-# link here as well:
-option_on = " selected"
-option_off = ""
-release_opt = option_on
-option_nr = 0
-version = (
-    version
-    + """<br/>
-Version select: <select onChange="window.location.href = this.value" onFocus="this.selectedIndex = {0}">
-    <option value="https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/index.html"{1}>Current release</option>
-    <option value="https://docs.nvidia.com/deeplearning/transformer-engine/documentation-archive.html">Older releases</option>
-</select>""".format(
-        option_nr, release_opt
-    )
-)
+    version = str(_raw_version)
+release = _raw_version
 
 # -- General configuration ---------------------------------------------------
 # https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration
@@ -92,12 +69,10 @@
 
 pygments_style = "sphinx"
 
-
 # -- Options for HTML output -------------------------------------------------
 # https://www.sphinx-doc.org/en/master/usage/configuration.html#options-for-html-output
 
 html_theme = "sphinx_rtd_theme"
-html_theme_path = [sphinx_rtd_theme.get_html_theme_path()]
 html_static_path = ["_static"]
 html_show_sphinx = False
 
@@ -106,7 +81,12 @@
     "css/nvidia_footer.css",
 ]
 
-html_theme_options = {"display_version": True, "collapse_navigation": False, "logo_only": False}
+html_theme_options = {
+    "collapse_navigation": False,
+    "logo_only": False,
+    "version_selector": False,
+    "language_selector": False,
+}
 
 napoleon_custom_sections = [
     ("Parallelism parameters", "params_style"),
@@ -116,8 +96,8 @@
     ("FP8-related parameters", "params_style"),
 ]
 
-breathe_projects = {"TransformerEngine": os.path.abspath("doxygen/xml/")}
+breathe_projects = {"TransformerEngine": root_path / "docs" / "doxygen" / "xml"}
 breathe_default_project = "TransformerEngine"
 
 autoapi_generate_api_docs = False
-autoapi_dirs = ["../transformer_engine"]
+autoapi_dirs = [root_path / "transformer_engine"]

docs/conf.py

@@ -0,0 +1,21 @@
+diff --git a/sphinx_rtd_theme/layout.html b/sphinx_rtd_theme/layout.html
+index e6a38b1..579eaec 100644
+--- a/sphinx_rtd_theme/layout.html
++++ b/sphinx_rtd_theme/layout.html
+@@ -124,6 +124,16 @@
+             {%- endif %}
+           </a>
+
++          {# Show TE version and version selector #}
++          <div class="version">
++            {{ version }}
++            <br>
++            Version select: <select onChange="window.location.href = this.value" onFocus="this.selectedIndex = {0}">
++              <option value="https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/index.html"{1}>Current release</option>
++              <option value="https://docs.nvidia.com/deeplearning/transformer-engine/documentation-archive.html">Older releases</option>
++            </select>
++          </div>
++
+           {%- if READTHEDOCS or DEBUG %}
+             {%- if theme_version_selector or theme_language_selector %}
+               <div class="switch-menus">

docs/version_select.patch

@@ -0,0 +1,14 @@
+# Copyright (c) 2022-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+"""Shared functions for the encoder tests"""
+from functools import lru_cache
+
+from transformer_engine.transformer_engine_jax import get_device_compute_capability
+
+
+@lru_cache
+def is_bf16_supported():
+    """Return if BF16 has hardware supported"""
+    gpu_arch = get_device_compute_capability(0)
+    return gpu_arch >= 80