update README w deepspeed info (#605)

2023-09-22 00:15:52 -04:00
parent d5f8589021
commit c25ba7939b
6 changed files with 114 additions and 88 deletions
--- a/setup.py
+++ b/setup.py
@@ -13,7 +13,12 @@ def parse_requirements():
                # Handle custom index URLs
                _, url = line.split()
                _dependency_links.append(url)
-            elif "flash-attn" not in line and line and line[0] != "#":
+            elif (
+                "flash-attn" not in line
+                and "deepspeed" not in line
+                and line
+                and line[0] != "#"
+            ):
                # Handle standard packages
                _install_requires.append(line)
    return _install_requires, _dependency_links
@@ -35,7 +40,7 @@ setup(
        "flash-attn": [
            "flash-attn>=2.2.1",
        ],
-        "extras": [
+        "deepspeed": [
            "deepspeed",
        ],
    },