Use latest Ray custom images in the LLM fine-tuning with Ray example

Author: astefanutti

examples/ray-finetune-llm-deepspeed/ray_finetune_llm_deepspeed.ipynb

@@ -67,12 +67,12 @@
     "    worker_memory_requests=128,\n",
     "    worker_memory_limits=256,\n",
     "    head_memory=128,\n",
-    "    # Use the following parameters with NVIDIA GPUs \n",
-    "    image=\"quay.io/rhoai/ray:2.23.0-py39-cu121\",\n",
+    "    # Use the following parameters with NVIDIA GPUs\n",
+    "    image=\"quay.io/rhoai/ray:2.35.0-py39-cu121-torch24-fa26\",\n",
     "    head_extended_resource_requests={'nvidia.com/gpu':1},\n",
     "    worker_extended_resource_requests={'nvidia.com/gpu':1},\n",
     "    # Or replace them with these parameters for AMD GPUs\n",
-    "    # image=\"quay.io/rhoai/ray:2.35.0-py39-rocm61-torch24\",\n",
+    "    # image=\"quay.io/rhoai/ray:2.35.0-py39-rocm61-torch24-fa26\",\n",
     "    # head_extended_resource_requests={'amd.com/gpu':1},\n",
     "    # worker_extended_resource_requests={'amd.com/gpu':1},\n",
     "))"

examples/ray-finetune-llm-deepspeed/requirements.txt

@@ -1,11 +1,5 @@
 accelerate==0.31.0
 awscliv2==2.3.0
 datasets==2.19.2
-deepspeed==0.14.4
-# Flash Attention 2 requires PyTorch to be installed first
-# See https://github.com/Dao-AILab/flash-attention/issues/453
-https://github.com/Dao-AILab/flash-attention/releases/download/v2.6.3/flash_attn-2.6.3+cu123torch2.3cxx11abiFALSE-cp39-cp39-linux_x86_64.whl
 peft==0.11.1
-ray[train]==2.23.0
-torch==2.3.1
 transformers==4.44.0