Upgrade verions

huggingface/pytorch/training/docker/2.6/py3/cu126/Dockerfile.gpu

@@ -12,7 +12,7 @@ LABEL dlc_major_version="2"
 
 # version args
 ARG TRANSFORMERS_VERSION=4.51.3
-ARG DATASETS_VERSION=3.5.0
+ARG DATASETS_VERSION=2.18.0
 ARG HUGGINGFACE_HUB_VERSION=0.30.0
 ARG DIFFUSERS_VERSION=0.33.1
 ARG EVALUATE_VERSION=0.4.3
@@ -22,6 +22,8 @@ ARG PEFT_VERSION=0.14.0
 ARG FLASH_ATTN_VERSION=2.7.4.post1
 ARG NINJA_VERSION=1.11.1.3
 ARG PYTHON=python3
+ARG MULTIPROCESS_VERSION=0.70.16
+ARG DILL_VERSION=0.3.8
 
 # TODO: Remove when the base image is updated
 RUN pip install --upgrade pip \
@@ -47,8 +49,10 @@ RUN pip install --no-cache-dir \
     ninja==${NINJA_VERSION} \
     trl==${TRL_VERSION} \
     peft==${PEFT_VERSION} \
-    flash-attn==${FLASH_ATTN_VERSION}
-
+    flash-attn==${FLASH_ATTN_VERSION} \
+    dill==${DILL_VERSION} \
+    multiprocess==${MULTIPROCESS_VERSION} \
+    "pathos<0.3.3"
 
 # hf_transfer will be a built-in feature, remove the env variable then
 ENV HF_HUB_ENABLE_HF_TRANSFER="1"