diff --git a/requirements.txt b/requirements.txt
index 75c8dd4..03916f0 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -9,4 +9,4 @@ transformers>=4.22.0
 torch>=1.10.0
 tqdm>=4.63.0
 cpm_kernels
-deepspeed
+deepspeed>0.6.1
\ No newline at end of file
diff --git a/setup.py b/setup.py
index af193f7..3260672 100644
--- a/setup.py
+++ b/setup.py
@@ -19,6 +19,8 @@ setup(
         "tokenizers>=0.11.0,<=0.11.4",
         "torch>=1.10.0",
         "tqdm>=4.63.0",
+        "cpm_kernels",
+        "deepspeed>0.6.1",
     ],
     entry_points={}
 )