From 5be8b555a0c14428ba224fd8c3c78689582cc87b Mon Sep 17 00:00:00 2001
From: NanoCode012 <kevinvong@rocketmail.com>
Date: Tue, 27 Feb 2024 15:46:44 +0900
Subject: [PATCH] fix: checkpoint saving with deepspeed (#1321)

---
 requirements.txt | 2 +-
 setup.py         | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/requirements.txt b/requirements.txt
index 8dce6daa7..722a9c644 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -5,7 +5,7 @@ transformers @ git+https://github.com/huggingface/transformers.git@ae49b218c3d71
 tokenizers==0.15.0
 bitsandbytes>=0.41.1
 accelerate==0.26.1
-deepspeed>=0.13.1
+deepspeed==0.13.1
 pydantic>=2.5.3
 addict
 fire
diff --git a/setup.py b/setup.py
index 85d9eae36..4a949a60d 100644
--- a/setup.py
+++ b/setup.py
@@ -74,7 +74,7 @@ setup(
             "fused-dense-lib  @ git+https://github.com/Dao-AILab/flash-attention@v2.3.3#subdirectory=csrc/fused_dense_lib",
         ],
         "deepspeed": [
-            "deepspeed>=0.13.1",
+            "deepspeed==0.13.1",
             "deepspeed-kernels",
         ],
         "mamba-ssm": [