diff --git a/requirements.txt b/requirements.txt index 3f17e5d329..ab1ca76ec2 100644 --- a/requirements.txt +++ b/requirements.txt @@ -6,7 +6,7 @@ tokenizers>=0.19.1 bitsandbytes==0.44.0 accelerate==0.34.2 datasets==2.21.0 -deepspeed==0.14.4 +deepspeed==0.15.1 pydantic==2.6.3 addict fire diff --git a/setup.py b/setup.py index 1b64fadaef..6f7d86f360 100644 --- a/setup.py +++ b/setup.py @@ -86,7 +86,7 @@ def parse_requirements(): "fused-dense-lib @ git+https://github.com/Dao-AILab/flash-attention@v2.6.2#subdirectory=csrc/fused_dense_lib", ], "deepspeed": [ - "deepspeed==0.14.4", + "deepspeed==0.15.1", "deepspeed-kernels", ], "mamba-ssm": [