From 2812d10244a4ad29bf80d51e12e67a6daee18ecb Mon Sep 17 00:00:00 2001
From: Jun Gong <jungong@anyscale.com>
Date: Fri, 19 May 2023 14:52:25 -0700
Subject: [PATCH] [Doc] Pin deepspeed version to fix fine_tuning example
 (#35477) (#35563)

Signed-off-by: Jun Gong <jungong@anyscale.com>
---
 .../examples/gptj_deepspeed_fine_tuning.ipynb | 20 +++++++++++++++----
 .../gptj_deepspeed_env.yaml                   |  6 +++---
 2 files changed, 19 insertions(+), 7 deletions(-)

diff --git a/doc/source/ray-air/examples/gptj_deepspeed_fine_tuning.ipynb b/doc/source/ray-air/examples/gptj_deepspeed_fine_tuning.ipynb
index fa2c7cf04ab5..552d904180b0 100644
--- a/doc/source/ray-air/examples/gptj_deepspeed_fine_tuning.ipynb
+++ b/doc/source/ray-air/examples/gptj_deepspeed_fine_tuning.ipynb
@@ -30,6 +30,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [
@@ -42,7 +43,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "#! pip install \"datasets\" \"evaluate\" \"accelerate>=0.16.0\" \"transformers>=4.26.0\" \"torch>=1.12.0\" \"deepspeed\""
+    "#! pip install \"datasets\" \"evaluate\" \"accelerate==0.18.0\" \"transformers>=4.26.0\" \"torch>=1.12.0\" \"deepspeed==0.8.3\""
    ]
   },
   {
@@ -57,6 +58,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [
@@ -78,6 +80,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [
@@ -153,10 +156,14 @@
     "        \"pip\": [\n",
     "            \"datasets\",\n",
     "            \"evaluate\",\n",
-    "            \"accelerate>=0.16.0\",\n",
-    "            \"transformers>=4.26.0\",\n",
+    "            # Latest combination of accelerate==0.19.0 and transformers==4.29.0\n",
+    "            # seems to have issues with DeepSpeed process group initialization,\n",
+    "            # and will result in a batch_size validation problem.\n",
+    "            # TODO(jungong) : get rid of the pins once the issue is fixed.\n",
+    "            \"accelerate==0.16.0\",\n",
+    "            \"transformers==4.26.0\",\n",
     "            \"torch>=1.12.0\",\n",
-    "            \"deepspeed\",\n",
+    "            \"deepspeed==0.9.2\",\n",
     "        ]\n",
     "    }\n",
     ")"
@@ -234,6 +241,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [
@@ -308,6 +316,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [
@@ -340,6 +349,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [
@@ -969,6 +979,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [
@@ -997,6 +1008,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [
diff --git a/release/air_examples/gptj_deepspeed_finetuning/gptj_deepspeed_env.yaml b/release/air_examples/gptj_deepspeed_finetuning/gptj_deepspeed_env.yaml
index cbecd450c849..1a36e9c03b31 100644
--- a/release/air_examples/gptj_deepspeed_finetuning/gptj_deepspeed_env.yaml
+++ b/release/air_examples/gptj_deepspeed_finetuning/gptj_deepspeed_env.yaml
@@ -7,10 +7,10 @@ python:
   pip_packages:
     - "datasets"
     - "evaluate"
-    - "accelerate>=0.16.0"
-    - "transformers>=4.26.0"
+    - "accelerate==0.16.0"
+    - "transformers==4.26.0"
     - "torch>=1.12.0"
-    - "deepspeed"
+    - "deepspeed==0.9.2"
     - myst-parser==0.15.2
     - myst-nb==0.13.1
     - jupytext==1.13.6