From 990b5896bc8abe11b6cce0eee0edf2629f85d9b6 Mon Sep 17 00:00:00 2001
From: NanoCode012 <nano@axolotl.ai>
Date: Tue, 1 Apr 2025 23:25:05 +0700
Subject: [PATCH] fix: downgrade deepspeed to fix grad checkpoint oom (#2465)
 [skip ci]

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

diff --git a/requirements.txt b/requirements.txt
index 567b446dd..dde64f392 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -16,7 +16,7 @@ transformers==4.50.3
 tokenizers>=0.21.1
 accelerate==1.5.2
 datasets==3.5.0
-deepspeed==0.16.4
+deepspeed==0.15.4
 trl==0.16.0
 
 optimum==1.16.2
diff --git a/setup.py b/setup.py
index aef8182af..29719b1f3 100644
--- a/setup.py
+++ b/setup.py
@@ -112,7 +112,7 @@ extras_require = {
         "yunchang==0.6.0",
     ],
     "deepspeed": [
-        "deepspeed==0.16.4",
+        "deepspeed==0.15.4",
         "deepspeed-kernels",
     ],
     "mamba-ssm": [