Build ec2 with MAX_JOBS=18 (#5216)

* build ec2 with MAX_JOBS=18

* revert toml

* Revert "build ec2 with MAX_JOBS=18"

This reverts commit cb92482.

* remove unpin packages

---------

Co-authored-by: Yadan Wei <[email protected]>

Author: Yadan-Wei

pytorch/training/docker/2.7/py3/Dockerfile.ec2.cpu.core_packages.json

@@ -21,14 +21,6 @@
   "tornado": {
     "version_specifier": ">=6.5.1",
     "skip": "True"
-  },
-  "s3torchconnector": {
-    "version_specifier": "==1.4.3",
-     "skip": "True"
-  },
-  "accelerate": {
-    "version_specifier": "==1.10.1",
-     "skip": "True"
   },
    "thinc": {
     "version_specifier": "==8.3.4"

pytorch/training/docker/2.7/py3/Dockerfile.sagemaker.cpu.core_packages.json

@@ -18,14 +18,6 @@
     "version_specifier": "==0.2.4",
     "skip": "True"
   },
-  "s3torchconnector": {
-    "version_specifier": "==1.4.3",
-    "skip": "True"
-  },
-  "accelerate": {
-    "version_specifier": "==1.10.1",
-    "skip": "True"
-  },
   "thinc": {
     "version_specifier": "==8.3.4"
   },

pytorch/training/docker/2.7/py3/cu128/Dockerfile.ec2.gpu.core_packages.json

@@ -22,14 +22,6 @@
     "version_specifier": "==2.3",
     "skip": "True"
   },
-  "s3torchconnector": {
-    "version_specifier": "==1.4.3",
-    "skip": "True"
-  },
-  "accelerate": {
-    "version_specifier": "==1.10.1",
-    "skip": "True"
-  },
   "thinc": {
     "version_specifier": "==8.3.4"
   },

pytorch/training/docker/2.7/py3/cu128/Dockerfile.gpu

@@ -89,7 +89,7 @@ RUN cd /tmp \
 ENV NVTE_FRAMEWORK=pytorch
 # Install flash-attn using instructions from https://github.com/Dao-AILab/flash-attention#installation-and-features
 # Set MAX_JOBS=4 to avoid OOM issues in installation process
-RUN MAX_JOBS=4 pip install --no-cache-dir flash-attn==${FLASH_ATTN_VERSION} --no-build-isolation
+RUN MAX_JOBS=18 pip install --no-cache-dir flash-attn==${FLASH_ATTN_VERSION} --no-build-isolation
 # Install TE using instructions from https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/installation.html
 RUN pip install --no-cache-dir git+https://github.com/NVIDIA/TransformerEngine.git@release_v${TE_VERSION} --no-build-isolation
 

pytorch/training/docker/2.7/py3/cu128/Dockerfile.sagemaker.gpu.core_packages.json

@@ -22,14 +22,6 @@
     "version_specifier": "==2.3",
     "skip": "True"
   },
-  "s3torchconnector": {
-    "version_specifier": "==1.4.3",
-    "skip": "True"
-  },
-  "accelerate": {
-    "version_specifier": "==1.10.1",
-    "skip": "True"
-  },
   "thinc": {
     "version_specifier": "==8.3.4"
   },