add optional dependencies in pyproject.toml

Azure · Jan 30, 2024 · df2725e · df2725e
1 parent 5d2f9b1
commit df2725e
Showing 1 changed file with 8 additions and 3 deletions.
diff --git a/pyproject.toml b/pyproject.toml
@@ -38,11 +38,8 @@ classifiers=[
 ]
 dependencies = [
     "torch",
-    "transformer-engine@git+https://github.com/NVIDIA/TransformerEngine.git@v0.11#egg=transformer-engine",
-    "flash-attn==1.0.9",
     "colorlog>=6.7.0",
     "deepspeed==0.9.2",
-    "mpi4py",
 ]
 dynamic = ["version"]
 
@@ -56,6 +53,14 @@ test = [
     "pytest>=7.2.0",
     "yapf>=0.32.0",
 ]
+nvidia = [
+    "transformer-engine@git+https://github.com/NVIDIA/TransformerEngine.git@v0.11#egg=transformer-engine",
+    "flash-attn==1.0.9",
+    "mpi4py",
+]
+amd = [
+    # TransformerEngine
+]
 
 [project.urls]
 homepage = "https://github.com/Azure/MS-AMP"