Explicitly pin PyTorch versions (mosaicml#1688)

* Bound torch to < 1.13
* Updated Docker flow to pin torchtext version

Author: bandish-shah

docker/Dockerfile

@@ -25,6 +25,11 @@ ARG PYTORCH_VERSION=1.11.0
 # version that corresponds to the PyTorch version
 ARG TORCHVISION_VERSION=0.12.0
 
+# The Torchtext version to install.
+# Reference https://github.com/pytorch/text#installation to determine the Torchtext
+# version that corresponds to the PyTorch version
+ARG TORCHTEXT_VERSION=0.12.0
+
 # In the Dockerimage, Pillow-SIMD is installed instead of Pillow. To trick pip into thinking that
 # Pillow is also installed (so it won't override it with a future pip install), a Pillow stub is included
 # PILLOW_PSEUDOVERSION is the Pillow version that pip thinks is installed
@@ -174,11 +179,16 @@ RUN pip${PYTHON_VERSION} install --no-cache-dir --upgrade /tmp/pillow_stub && \
 #################
 ARG PYTORCH_VERSION
 ARG TORCHVISION_VERSION
+ARG TORCHTEXT_VERSION
+
+# Set so supporting PyTorch packages such as Torchvision, Torchaudio, Torchtext pin PyTorch version
+ENV PYTORCH_VERSION=${PYTORCH_VERSION}
 
 RUN CUDA_VERSION_TAG=$(python${PYTHON_VERSION} -c "print('cu' + ''.join('${CUDA_VERSION}'.split('.')[:2]) if '${CUDA_VERSION}' else 'cpu')") && \
     pip${PYTHON_VERSION} install --no-cache-dir --find-links https://download.pytorch.org/whl/torch_stable.html \
         torch==${PYTORCH_VERSION}+${CUDA_VERSION_TAG} \
-        torchvision==${TORCHVISION_VERSION}+${CUDA_VERSION_TAG}
+        torchvision==${TORCHVISION_VERSION}+${CUDA_VERSION_TAG} \
+        torchtext==${TORCHTEXT_VERSION}
 
 ###################################
 # Mellanox OFED driver installation

docker/build_matrix.yaml

@@ -6,6 +6,7 @@
   TAGS:
   - mosaicml/pytorch:1.10.2_cu113-python3.7-ubuntu20.04
   TARGET: pytorch_stage
+  TORCHTEXT_VERSION: 0.11.2
   TORCHVISION_VERSION: 0.11.3
 - BASE_IMAGE: ubuntu:20.04
   CUDA_VERSION: ''
@@ -15,6 +16,7 @@
   TAGS:
   - mosaicml/pytorch:1.10.2_cpu-python3.7-ubuntu20.04
   TARGET: pytorch_stage
+  TORCHTEXT_VERSION: 0.11.2
   TORCHVISION_VERSION: 0.11.3
 - BASE_IMAGE: nvidia/cuda:11.5.2-cudnn8-devel-ubuntu20.04
   CUDA_VERSION: 11.5.2
@@ -23,6 +25,7 @@
   TAGS:
   - mosaicml/pytorch:1.11.0_cu115-python3.8-ubuntu20.04
   TARGET: pytorch_stage
+  TORCHTEXT_VERSION: 0.12.0
   TORCHVISION_VERSION: 0.12.0
 - BASE_IMAGE: ubuntu:20.04
   CUDA_VERSION: ''
@@ -32,6 +35,7 @@
   TAGS:
   - mosaicml/pytorch:1.11.0_cpu-python3.8-ubuntu20.04
   TARGET: pytorch_stage
+  TORCHTEXT_VERSION: 0.12.0
   TORCHVISION_VERSION: 0.12.0
 - BASE_IMAGE: nvidia/cuda:11.6.2-cudnn8-devel-ubuntu20.04
   CUDA_VERSION: 11.6.2
@@ -41,6 +45,7 @@
   - mosaicml/pytorch:1.12.1_cu116-python3.9-ubuntu20.04
   - mosaicml/pytorch:latest
   TARGET: pytorch_stage
+  TORCHTEXT_VERSION: 0.13.1
   TORCHVISION_VERSION: 0.13.1
 - BASE_IMAGE: nvidia/cuda:11.6.2-cudnn8-devel-ubuntu20.04
   CUDA_VERSION: 11.6.2
@@ -50,6 +55,7 @@
   - mosaicml/pytorch_vision:1.12.1_cu116-python3.9-ubuntu20.04
   - mosaicml/pytorch_vision:latest
   TARGET: vision_stage
+  TORCHTEXT_VERSION: 0.13.1
   TORCHVISION_VERSION: 0.13.1
 - BASE_IMAGE: ubuntu:20.04
   CUDA_VERSION: ''
@@ -60,6 +66,7 @@
   - mosaicml/pytorch:1.12.1_cpu-python3.9-ubuntu20.04
   - mosaicml/pytorch:latest_cpu
   TARGET: pytorch_stage
+  TORCHTEXT_VERSION: 0.13.1
   TORCHVISION_VERSION: 0.13.1
 - BASE_IMAGE: ubuntu:20.04
   CUDA_VERSION: ''
@@ -70,6 +77,7 @@
   - mosaicml/pytorch_vision:1.12.1_cpu-python3.9-ubuntu20.04
   - mosaicml/pytorch_vision:latest_cpu
   TARGET: vision_stage
+  TORCHTEXT_VERSION: 0.13.1
   TORCHVISION_VERSION: 0.13.1
 - BASE_IMAGE: nvidia/cuda:11.6.2-cudnn8-devel-ubuntu20.04
   COMPOSER_INSTALL_COMMAND: mosaicml[all]
@@ -79,6 +87,7 @@
   TAGS:
   - mosaicml/composer:latest
   TARGET: composer_stage
+  TORCHTEXT_VERSION: 0.13.1
   TORCHVISION_VERSION: 0.13.1
 - BASE_IMAGE: ubuntu:20.04
   COMPOSER_INSTALL_COMMAND: mosaicml[all]
@@ -88,6 +97,7 @@
   TAGS:
   - mosaicml/composer:latest_cpu
   TARGET: composer_stage
+  TORCHTEXT_VERSION: 0.13.1
   TORCHVISION_VERSION: 0.13.1
 - BASE_IMAGE: nvidia/cuda:11.6.2-cudnn8-devel-ubuntu20.04
   COMPOSER_INSTALL_COMMAND: mosaicml[all]==0.11.0
@@ -97,6 +107,7 @@
   TAGS:
   - mosaicml/composer:0.11.0
   TARGET: composer_stage
+  TORCHTEXT_VERSION: 0.13.1
   TORCHVISION_VERSION: 0.13.1
 - BASE_IMAGE: ubuntu:20.04
   COMPOSER_INSTALL_COMMAND: mosaicml[all]==0.11.0
@@ -106,6 +117,7 @@
   TAGS:
   - mosaicml/composer:0.11.0_cpu
   TARGET: composer_stage
+  TORCHTEXT_VERSION: 0.13.1
   TORCHVISION_VERSION: 0.13.1
 - BASE_IMAGE: nvidia/cuda:11.6.2-cudnn8-devel-ubuntu20.04
   COMPOSER_INSTALL_COMMAND: mosaicml[all]GIT_COMMIT
@@ -114,6 +126,7 @@
   PYTORCH_VERSION: 1.12.1
   TAGS: []
   TARGET: composer_stage
+  TORCHTEXT_VERSION: 0.13.1
   TORCHVISION_VERSION: 0.13.1
 - BASE_IMAGE: ubuntu:20.04
   COMPOSER_INSTALL_COMMAND: mosaicml[all]GIT_COMMIT
@@ -122,4 +135,5 @@
   PYTORCH_VERSION: 1.12.1
   TAGS: []
   TARGET: composer_stage
+  TORCHTEXT_VERSION: 0.13.1
   TORCHVISION_VERSION: 0.13.1

docker/generate_build_matrix.py

@@ -39,6 +39,16 @@ def _get_torchvision_version(pytorch_version: str):
     raise ValueError(f'Invalid pytorch_version: {pytorch_version}')
 
 
+def _get_torchtext_version(pytorch_version: str):
+    if pytorch_version == '1.12.1':
+        return '0.13.1'
+    if pytorch_version == '1.11.0':
+        return '0.12.0'
+    if pytorch_version == '1.10.2':
+        return '0.11.2'
+    raise ValueError(f'Invalid pytorch_version: {pytorch_version}')
+
+
 def _get_base_image(cuda_version: str):
     if not cuda_version:
         return 'ubuntu:20.04'
@@ -136,6 +146,8 @@ def _main():
                 stage,
             'TORCHVISION_VERSION':
                 _get_torchvision_version(pytorch_version),
+            'TORCHTEXT_VERSION':
+                _get_torchtext_version(pytorch_version),
             'TAGS':
                 _get_pytorch_tags(
                     python_version=python_version,
@@ -175,6 +187,7 @@ def _main():
             'PYTORCH_VERSION': pytorch_version,
             'TARGET': 'composer_stage',
             'TORCHVISION_VERSION': _get_torchvision_version(pytorch_version),
+            'TORCHTEXT_VERSION': _get_torchtext_version(pytorch_version),
             'COMPOSER_INSTALL_COMMAND': f'mosaicml[all]{composer_version}',
             'TAGS': _get_composer_tags(
                 composer_version=composer_version,

setup.py

@@ -78,7 +78,7 @@ def package_files(prefix: str, directory: str, extension: str):
     'torchmetrics>=0.7.0,<0.10.0',
     'torch_optimizer>=0.3.0,<0.4',
     'torchvision>=0.10.0',  # torchvision has strict pytorch requirements
-    'torch>=1.10,<2',
+    'torch>=1.10,<1.13',
     'yahp>=0.1.3,<0.2',
     'requests>=2.26.0,<3',
     'numpy>=1.21.5,<1.23.0',