Add augment to version string and allow install path overriding with NVTE_INSTALL_PATH

Author: cbcase

setup.py

@@ -27,6 +27,9 @@ def te_version() -> str:
     Includes Git commit as local version, unless suppressed with
     NVTE_NO_LOCAL_VERSION environment variable.
 
+    [augment] We replace the git hash with the string "augment"
+    to identify this as a custom augment release.
+
     """
     with open(root_path / "VERSION", "r") as f:
         version = f.readline().strip()
@@ -43,7 +46,7 @@ def te_version() -> str:
             pass
         else:
             commit = output.stdout.strip()
-            version += f"+{commit}"
+            version += "+augment"
     return version
 
 @lru_cache(maxsize=1)
@@ -290,6 +293,7 @@ def add_unique(l: List[str], vals: Union[str, List[str]]) -> None:
 
     # Framework-specific requirements
     if "pytorch" in frameworks():
+        # [augment] We remove the versioning requirement on flash-attn
         add_unique(install_reqs, ["torch", "flash-attn"])
         add_unique(test_reqs, ["numpy", "onnxruntime", "torchvision"])
     if "jax" in frameworks():

transformer_engine/common/__init__.py

@@ -13,6 +13,11 @@
 def get_te_path():
     """Find Transformer Engine install path using pip"""
 
+    # [augment] Allow overriding install path with an envvar. Should be something like:
+    # > NVTE_INSTALL_PATH=/opt/conda/lib/python3.9/site-packages
+    if "NVTE_INSTALL_PATH" in os.environ:
+        return os.environ["NVTE_INSTALL_PATH"]
+
     command = [sys.executable, "-m", "pip", "show", "transformer_engine"]
     result = subprocess.run(command, capture_output=True, check=True, text=True)
     result = result.stdout.replace("\n", ":").split(":")