Fix deepspeed version

ai-safety-foundation · Jan 16, 2024 · 5792e55 · 5792e55
1 parent 0da6696
commit 5792e55
Show file tree

Hide file tree

Showing 2 changed files with 133 additions and 7 deletions.
diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -7,17 +7,20 @@
     readme="README.md"
     version="0.0.0"
 
+    # Note: Zstandard is required for downloading datasets such as The Pile
     [tool.poetry.dependencies]
         datasets=">=2.15.0"
+        deepspeed={version=">=0.12.6", extras=["deepspeed"], optional=false}
         einops=">=0.6"
+        mpi4py={version=">=3.1.5", extras=["deepspeed"], optional=true}
         pydantic=">=2.5.2"
         python=">=3.10, <3.12"
         strenum=">=0.4.15"
         tokenizers=">=0.15.0"
         torch=">=2.1.1"
         transformers=">=4.35.2"
         wandb=">=0.16.1"
-        zstandard=">=0.22.0"    # Required for downloading datasets such as The Pile
+        zstandard=">=0.22.0"
 
     [tool.poetry.group]
         [tool.poetry.group.dev.dependencies]