Fix torch 2.10 version pinning and increase CPU image size baseline

- Split pip install into separate commands to prevent dependency resolver
  from downgrading torch 2.10.0 to 2.9.1
- Add torch version constraint when installing fastai/accelerate/spacy
- Increase CPU image_size_baseline from 7200 to 12000 in buildspec files

Author: Bhanu Teja Goshikonda

pytorch/training/buildspec-2-10-ec2.yml

@@ -44,7 +44,7 @@ images:
   BuildEC2CPUPTTrainPy3DockerImage:
     <<: *TRAINING_REPOSITORY
     build: &PYTORCH_CPU_TRAINING_PY3 false
-    image_size_baseline: 7200
+    image_size_baseline: 12000
     device_type: &DEVICE_TYPE cpu
     python_version: &DOCKER_PYTHON_VERSION py3
     tag_python_version: &TAG_PYTHON_VERSION py313

pytorch/training/buildspec-2-10-sm.yml

@@ -44,7 +44,7 @@ images:
   BuildSageMakerCPUPTTrainPy3DockerImage:
     <<: *TRAINING_REPOSITORY
     build: &PYTORCH_CPU_TRAINING_PY3 false
-    image_size_baseline: 7200
+    image_size_baseline: 12000
     device_type: &DEVICE_TYPE cpu
     python_version: &DOCKER_PYTHON_VERSION py3
     tag_python_version: &TAG_PYTHON_VERSION py313

pytorch/training/docker/2.10/py3/Dockerfile.cpu

@@ -197,20 +197,26 @@ RUN pip install --no-cache-dir \
     pytz \
     tzdata
 
-# Install PyTorch
+# Install PyTorch - split into separate commands to prevent dependency resolver from downgrading torch
 RUN pip install --no-cache-dir -U torch==${PYTORCH_VERSION} \
     torchvision==${TORCHVISION_VERSION} \
     torchaudio==${TORCHAUDIO_VERSION} \
-    --index-url https://download.pytorch.org/whl/cpu \
-    && pip install --no-cache-dir -U torchtnt==${TORCHTNT_VERSION} \
+    --index-url https://download.pytorch.org/whl/cpu
+
+# Install torch ecosystem packages
+RUN pip install --no-cache-dir -U torchtnt==${TORCHTNT_VERSION} \
     torchdata==${TORCHDATA_VERSION} \
-    s3torchconnector \
+    s3torchconnector
+
+# Install ML packages with torch version constraint to prevent downgrade
+# pin numpy requirement for fastai dependency
+# requires explicit declaration of spacy, thinc, blis
+# pin thinc due to incompatibility with numpy 1.26.4 (sagemaker doesn't support latest numpy)
+RUN pip install --no-cache-dir -U \
+    "torch==${PYTORCH_VERSION}" \
     fastai \
     accelerate \
-    # pin numpy requirement for fastai dependency
-    # requires explicit declaration of spacy, thic, blis
     spacy \
-    # pin thinc due to incompatibility with numpy 1.26.4 (sagemaker doesn't support latest numpy)
     thinc \
     blis \
     numpy \

pytorch/training/docker/2.10/py3/cu130/Dockerfile.gpu

@@ -116,21 +116,27 @@ RUN pip install --no-cache-dir \
     pytz \
     tzdata
 
-# Install PyTorch
+# Install PyTorch - split into separate commands to prevent dependency resolver from downgrading torch
 RUN pip install --no-cache-dir -U torch==${PYTORCH_VERSION} \
     torchvision==${TORCHVISION_VERSION} \
     torchaudio==${TORCHAUDIO_VERSION} \
-    --index-url https://download.pytorch.org/whl/cu130 \
-    && pip install --no-cache-dir -U torchtnt==${TORCHTNT_VERSION} \
+    --index-url https://download.pytorch.org/whl/cu130
+
+# Install torch ecosystem packages
+RUN pip install --no-cache-dir -U torchtnt==${TORCHTNT_VERSION} \
     torchdata==${TORCHDATA_VERSION} \
     triton \
-    s3torchconnector \
+    s3torchconnector
+
+# Install ML packages with torch version constraint to prevent downgrade
+# pin numpy requirement for fastai dependency
+# requires explicit declaration of spacy, thinc, blis
+# pin thinc due to incompatibility with numpy 1.26.4 (sagemaker doesn't support latest numpy)
+RUN pip install --no-cache-dir -U \
+    "torch==${PYTORCH_VERSION}" \
     fastai \
     accelerate \
-    # pin numpy requirement for fastai dependency
-    # requires explicit declaration of spacy, thic, blis
     spacy \
-    # pin thinc due to incompatibility with numpy 1.26.4 (sagemaker doesn't support latest numpy)
     thinc \
     blis \
     numpy \