Add NVIDIA TransformerEngine

recipes/transformer-engine/001-cmake.patch

@@ -0,0 +1,11 @@
+--- a/build_tools/build_ext.py
++++ b/build_tools/build_ext.py
+@@ -56,7 +56,7 @@ class CMakeExtension(setuptools.Extension):
+             cmake_path,
+             "-B",
+             build_dir,
+-            f"-DPython_EXECUTABLE={sys.executable}",
++            "-DPython_EXECUTABLE=$PYTHON",
+             f"-DPython_INCLUDE_DIR={sysconfig.get_path('include')}",
+             f"-DCMAKE_BUILD_TYPE={build_type}",
+             f"-DCMAKE_INSTALL_PREFIX={install_dir}",

recipes/transformer-engine/activate.sh

@@ -0,0 +1,5 @@
+if [[ -n "${CUDA_HOME:-}" ]]
+then
+    export CUDA_HOME_BACKUP="${CUDA_HOME}"
+fi
+export CUDA_HOME="${CONDA_PREFIX}"

recipes/transformer-engine/build_te.sh

@@ -0,0 +1,30 @@
+#!/bin/bash
+set -euxo pipefail
+
+cat > $RECIPE_DIR/gcc_shim <<"EOF"
+#!/bin/sh
+exec $GCC -I$PREFIX/include "$@"
+EOF
+
+chmod +x $RECIPE_DIR/gcc_shim
+export CC="$RECIPE_DIR/gcc_shim"
+
+# Re-arrange build files to match the expected layout.
+ln -s $PREFIX/nvvm $PREFIX/targets/x86_64-linux/nvvm
+cp $BUILD_PREFIX/targets/x86_64-linux/include/fatbinary_section.h $PREFIX/targets/x86_64-linux/include
+cp $PREFIX/include/cudnn*.h $PREFIX/targets/x86_64-linux/include
+
+echo "Installing transformer-engine"
+NVTE_NO_LOCAL_VERSION=1 ${PYTHON} -m pip install .
+
+# Remove re-arranged files.
+rm -rf $PREFIX/targets/x86_64-linux/nvvm
+rm -rf $PREFIX/targets/x86_64-linux/include/fatbinary_section.h
+rm -rf $PREFIX/targets/x86_64-linux/include/cudnn*.h
+
+
+mkdir -p $PREFIX/etc/conda/activate.d
+cp $RECIPE_DIR/activate.sh $PREFIX/etc/conda/activate.d/transformer-engine-activate.sh
+
+mkdir -p $PREFIX/etc/conda/deactivate.d
+cp $RECIPE_DIR/deactivate.sh $PREFIX/etc/conda/deactivate.d/transformer-engine-deactivate.sh

recipes/transformer-engine/deactivate.sh

@@ -0,0 +1,5 @@
+if [[ -n "${CUDA_HOME_BACKUP:-}" ]]
+then
+    export CUDA_HOME="${CUDA_HOME_BACKUP}"
+    unset CUDA_HOME_BACKUP
+fi

recipes/transformer-engine/meta.yaml

@@ -0,0 +1,106 @@
+{% set version = "1.13" %}
+{% set torch_proc_type = "cuda" if cuda_compiler_version != "None" else "cpu" %}
+
+{% if cuda_compiler_version in (None, "None", True, False) %}
+{% set cuda_major = 0 %}
+{% else %}
+{% set cuda_major = environ.get("cuda_compiler_version", "11.8").split(".")[0] | int %}
+{% endif %}
+
+package:
+  name: transformer-engine
+  version: {{ version }}
+
+source:
+  - url: https://github.com/NVIDIA/TransformerEngine/archive/refs/tags/v{{ version }}.tar.gz
+    sha256: 90dd8f694cba29707c502a4f4d12871cb9d4c1f157d329f22cb3d382de8468ae
+    patches:
+      - 001-cmake.patch
+
+  - url: https://github.com/NVIDIA/cudnn-frontend/archive/refs/tags/v1.8.0.tar.gz
+    sha256: 5907bf8baa8f3b74ad430d54027aadd14a0a61166ae4b00c2daece0df384409e
+    folder: 3rdparty/cudnn-frontend
+
+  - url: https://github.com/google/googletest/archive/refs/tags/v1.14.0.tar.gz
+    sha256: 8ad598c73ad796e0d8280b082cebd82a630d73e73cd3c70057938a6501bba5d7
+    folder: 3rdparty/googletest
+
+build:
+  number: 0
+
+outputs:
+  - name: transformer-engine-torch
+    version: {{ version }}
+    script: build_te.sh
+
+    build:
+      skip: true  # [not (cuda_compiler_version or "").startswith("12")]
+      skip: true  # [not (linux64 or aarch64)]
+      script_env:
+        - CUDA_HOME={{ BUILD_PREFIX }}
+        - NVTE_FRAMEWORK=pytorch
+      run_exports:
+        - {{ pin_subpackage('transformer-engine-torch', max_pin='x.x') }}
+
+    requirements:
+      build:
+        - python                                 # [build_platform != target_platform]
+        - cross-python_{{ target_platform }}     # [build_platform != target_platform]
+        - pytorch                                # [build_platform != target_platform]
+        - pytorch =*={{ torch_proc_type }}*      # [build_platform != target_platform]
+        - {{ stdlib('c') }}
+        - {{ compiler('c') }}
+        - {{ compiler('cxx') }}
+        - {{ compiler('cuda') }}  # [cuda_compiler_version not in (undefined, 'None')]
+        - cuda-nvvm
+        - cmake >=3.21
+        - ninja
+
+      host:
+        - python
+        - pip
+        - pybind11
+        - pytorch
+        - pytorch =*={{ torch_proc_type }}*
+        - setuptools
+        - pydantic
+        - importlib-metadata >=1.0
+        - packaging
+        - cuda-driver-dev
+        - cuda-cudart-dev
+        - cuda-nvrtc-dev
+        - cuda-nvtx-dev
+        - cuda-nvml-dev
+        - libcublas-dev
+        - libcusparse-dev
+        - libcusolver-dev
+        - cudnn
+        - cuda-profiler-api
+
+      run:
+        - python
+        - pydantic
+        - importlib-metadata >=1.0
+
+      run_constrained:
+        # additional run constraint to the one from the (version-only) run_export;
+        # constraining the CPU builds to CPU pytorch isn't 100% necessary, but cleaner
+        - pytorch =*={{ torch_proc_type }}*
+
+    test:
+      imports:
+        - transformer_engine.pytorch
+        - transformer_engine.common
+        - transformer_engine.common.recipe
+
+about:
+  home: https://github.com/NVIDIA/TransformerEngine
+  summary: A library for accelerating Transformer models on NVIDIA GPUs.
+  license: Apache-2.0
+  license_file: LICENSE
+  dev_url: https://github.com/NVIDIA/TransformerEngine
+
+extra:
+  feedstock-name: transformer-engine
+  recipe-maintainers:
+    - pstjohn