diff --git a/docker/Dockerfile b/docker/Dockerfile
index c3f4dee907..9a18e3ccc9 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -11,8 +11,8 @@
 ARG CUDA_VERSION=11.3.1
 
 # Calculate the base image based on CUDA_VERSION
-ARG BASE_IMAGE=${CUDA_VERSION:+"nvidia/cuda:${CUDA_VERSION}-cudnn8-devel-ubuntu20.04"}
-ARG BASE_IMAGE=${BASE_IMAGE:-"ubuntu:20.04"}
+ARG BASE_IMAGE=${CUDA_VERSION:+"nvidia/cuda:${CUDA_VERSION}-cudnn8-devel-ubuntu22.04"}
+ARG BASE_IMAGE=${BASE_IMAGE:-"ubuntu:22.04"}
 
 # The Python version to install
 ARG PYTHON_VERSION=3.10
@@ -251,7 +251,7 @@ ARG MOFED_VERSION
 
 RUN if [ -n "$MOFED_VERSION" ] ; then \
         wget -qO - http://www.mellanox.com/downloads/ofed/RPM-GPG-KEY-Mellanox | sudo apt-key add - && \
-        wget -P /etc/apt/sources.list.d/ http://linux.mellanox.com/public/repo/mlnx_ofed/$MOFED_VERSION/ubuntu20.04/mellanox_mlnx_ofed.list && \
+        wget -P /etc/apt/sources.list.d/ http://linux.mellanox.com/public/repo/mlnx_ofed/$MOFED_VERSION/ubuntu22.04/mellanox_mlnx_ofed.list && \
         apt-get update && \
         apt-get install -y mlnx-ofed-dpdk-upstream-libs-user-only  ; \
     fi
diff --git a/docker/README.md b/docker/README.md
index d4af7cc4ad..65fdf796f2 100644
--- a/docker/README.md
+++ b/docker/README.md
@@ -30,15 +30,15 @@ To install composer, once inside the image, run `pip install mosaicml`.
 <!-- BEGIN_PYTORCH_BUILD_MATRIX -->
 | Linux Distro   | Flavor   | PyTorch Version   | CUDA Version        | Python Version   | Docker Tags                                                                                                                                                                                                          |
 |----------------|----------|-------------------|---------------------|------------------|----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
-| Ubuntu 20.04   | Base     | 2.5.1             | 12.4.1 (Infiniband) | 3.11             | `mosaicml/pytorch:latest`, `mosaicml/pytorch:2.5.1_cu124-python3.11-ubuntu20.04`                 |
-| Ubuntu 20.04   | Base     | 2.5.1             | 12.4.1 (EFA)        | 3.11             | `mosaicml/pytorch:latest-aws`, `mosaicml/pytorch:2.5.1_cu124-python3.11-ubuntu20.04-aws` |
-| Ubuntu 20.04   | Base     | 2.5.1             | cpu                 | 3.11             | `mosaicml/pytorch:latest_cpu`, `mosaicml/pytorch:2.5.1_cpu-python3.11-ubuntu20.04`             |
-| Ubuntu 20.04   | Base     | 2.4.1             | 12.4.1 (Infiniband) | 3.11             | `mosaicml/pytorch:2.4.1_cu124-python3.11-ubuntu20.04`                                                                                        |
-| Ubuntu 20.04   | Base     | 2.4.1             | 12.4.1 (EFA)        | 3.11             | `mosaicml/pytorch:2.4.1_cu124-python3.11-ubuntu20.04-aws`                                                                                |
-| Ubuntu 20.04   | Base     | 2.4.1             | cpu                 | 3.11             | `mosaicml/pytorch:2.4.1_cpu-python3.11-ubuntu20.04`                                                                                            |
-| Ubuntu 20.04   | Base     | 2.3.1             | 12.1.1 (Infiniband) | 3.11             | `mosaicml/pytorch:2.3.1_cu121-python3.11-ubuntu20.04`                                                                                        |
-| Ubuntu 20.04   | Base     | 2.3.1             | 12.1.1 (EFA)        | 3.11             | `mosaicml/pytorch:2.3.1_cu121-python3.11-ubuntu20.04-aws`                                                                                |
-| Ubuntu 20.04   | Base     | 2.3.1             | cpu                 | 3.11             | `mosaicml/pytorch:2.3.1_cpu-python3.11-ubuntu20.04`                                                                                            |
+| Ubuntu 20.04   | Base     | 2.5.1             | 12.4.1 (Infiniband) | 3.11             | `mosaicml/pytorch:latest`, `mosaicml/pytorch:2.5.1_cu124-python3.11-ubuntu22.04`                 |
+| Ubuntu 20.04   | Base     | 2.5.1             | 12.4.1 (EFA)        | 3.11             | `mosaicml/pytorch:latest-aws`, `mosaicml/pytorch:2.5.1_cu124-python3.11-ubuntu22.04-aws` |
+| Ubuntu 20.04   | Base     | 2.5.1             | cpu                 | 3.11             | `mosaicml/pytorch:latest_cpu`, `mosaicml/pytorch:2.5.1_cpu-python3.11-ubuntu22.04`             |
+| Ubuntu 20.04   | Base     | 2.4.1             | 12.4.1 (Infiniband) | 3.11             | `mosaicml/pytorch:2.4.1_cu124-python3.11-ubuntu22.04`                                                                                        |
+| Ubuntu 20.04   | Base     | 2.4.1             | 12.4.1 (EFA)        | 3.11             | `mosaicml/pytorch:2.4.1_cu124-python3.11-ubuntu22.04-aws`                                                                                |
+| Ubuntu 20.04   | Base     | 2.4.1             | cpu                 | 3.11             | `mosaicml/pytorch:2.4.1_cpu-python3.11-ubuntu22.04`                                                                                            |
+| Ubuntu 20.04   | Base     | 2.3.1             | 12.1.1 (Infiniband) | 3.11             | `mosaicml/pytorch:2.3.1_cu121-python3.11-ubuntu22.04`                                                                                        |
+| Ubuntu 20.04   | Base     | 2.3.1             | 12.1.1 (EFA)        | 3.11             | `mosaicml/pytorch:2.3.1_cu121-python3.11-ubuntu22.04-aws`                                                                                |
+| Ubuntu 20.04   | Base     | 2.3.1             | cpu                 | 3.11             | `mosaicml/pytorch:2.3.1_cpu-python3.11-ubuntu22.04`                                                                                            |
 <!-- END_PYTORCH_BUILD_MATRIX -->
 
 **Note**: The `mosaicml/pytorch:latest`, `mosaicml/pytorch:latest_cpu`, and `mosaicml/pytorch:latest-aws`
diff --git a/docker/build_matrix.yaml b/docker/build_matrix.yaml
index 9c3d9eed36..cf24bd0f6a 100644
--- a/docker/build_matrix.yaml
+++ b/docker/build_matrix.yaml
@@ -1,6 +1,6 @@
 # This file is automatically generated by generate_build_matrix.py. DO NOT EDIT!
 - AWS_OFI_NCCL_VERSION: ''
-  BASE_IMAGE: nvidia/cuda:12.4.1-cudnn-devel-ubuntu20.04
+  BASE_IMAGE: nvidia/cuda:12.4.1-cudnn-devel-ubuntu22.04
   CUDA_VERSION: 12.4.1
   IMAGE_NAME: torch-2-5-1-cu124
   MOFED_VERSION: latest-23.10
@@ -10,14 +10,14 @@
   PYTORCH_NIGHTLY_VERSION: ''
   PYTORCH_VERSION: 2.5.1
   TAGS:
-  - mosaicml/pytorch:2.5.1_cu124-python3.11-ubuntu20.04
-  - ghcr.io/databricks-mosaic/pytorch:2.5.1_cu124-python3.11-ubuntu20.04
+  - mosaicml/pytorch:2.5.1_cu124-python3.11-ubuntu22.04
+  - ghcr.io/databricks-mosaic/pytorch:2.5.1_cu124-python3.11-ubuntu22.04
   - mosaicml/pytorch:latest
   - ghcr.io/databricks-mosaic/pytorch:latest
   TARGET: pytorch_stage
   TORCHVISION_VERSION: 0.20.1
 - AWS_OFI_NCCL_VERSION: v1.11.0-aws
-  BASE_IMAGE: nvidia/cuda:12.4.1-cudnn-devel-ubuntu20.04
+  BASE_IMAGE: nvidia/cuda:12.4.1-cudnn-devel-ubuntu22.04
   CUDA_VERSION: 12.4.1
   IMAGE_NAME: torch-2-5-1-cu124-aws
   MOFED_VERSION: ''
@@ -27,8 +27,8 @@
   PYTORCH_NIGHTLY_VERSION: ''
   PYTORCH_VERSION: 2.5.1
   TAGS:
-  - mosaicml/pytorch:2.5.1_cu124-python3.11-ubuntu20.04-aws
-  - ghcr.io/databricks-mosaic/pytorch:2.5.1_cu124-python3.11-ubuntu20.04-aws
+  - mosaicml/pytorch:2.5.1_cu124-python3.11-ubuntu22.04-aws
+  - ghcr.io/databricks-mosaic/pytorch:2.5.1_cu124-python3.11-ubuntu22.04-aws
   - mosaicml/pytorch:latest-aws
   - ghcr.io/databricks-mosaic/pytorch:latest-aws
   TARGET: pytorch_stage
@@ -44,14 +44,14 @@
   PYTORCH_NIGHTLY_VERSION: ''
   PYTORCH_VERSION: 2.5.1
   TAGS:
-  - mosaicml/pytorch:2.5.1_cpu-python3.11-ubuntu20.04
-  - ghcr.io/databricks-mosaic/pytorch:2.5.1_cpu-python3.11-ubuntu20.04
+  - mosaicml/pytorch:2.5.1_cpu-python3.11-ubuntu22.04
+  - ghcr.io/databricks-mosaic/pytorch:2.5.1_cpu-python3.11-ubuntu22.04
   - mosaicml/pytorch:latest_cpu
   - ghcr.io/databricks-mosaic/pytorch:latest_cpu
   TARGET: pytorch_stage
   TORCHVISION_VERSION: 0.20.1
 - AWS_OFI_NCCL_VERSION: ''
-  BASE_IMAGE: nvidia/cuda:12.4.1-cudnn-devel-ubuntu20.04
+  BASE_IMAGE: nvidia/cuda:12.4.1-cudnn-devel-ubuntu22.04
   CUDA_VERSION: 12.4.1
   IMAGE_NAME: torch-2-4-1-cu124
   MOFED_VERSION: latest-23.10
@@ -61,12 +61,12 @@
   PYTORCH_NIGHTLY_VERSION: ''
   PYTORCH_VERSION: 2.4.1
   TAGS:
-  - mosaicml/pytorch:2.4.1_cu124-python3.11-ubuntu20.04
-  - ghcr.io/databricks-mosaic/pytorch:2.4.1_cu124-python3.11-ubuntu20.04
+  - mosaicml/pytorch:2.4.1_cu124-python3.11-ubuntu22.04
+  - ghcr.io/databricks-mosaic/pytorch:2.4.1_cu124-python3.11-ubuntu22.04
   TARGET: pytorch_stage
   TORCHVISION_VERSION: 0.19.1
 - AWS_OFI_NCCL_VERSION: v1.11.0-aws
-  BASE_IMAGE: nvidia/cuda:12.4.1-cudnn-devel-ubuntu20.04
+  BASE_IMAGE: nvidia/cuda:12.4.1-cudnn-devel-ubuntu22.04
   CUDA_VERSION: 12.4.1
   IMAGE_NAME: torch-2-4-1-cu124-aws
   MOFED_VERSION: ''
@@ -76,8 +76,8 @@
   PYTORCH_NIGHTLY_VERSION: ''
   PYTORCH_VERSION: 2.4.1
   TAGS:
-  - mosaicml/pytorch:2.4.1_cu124-python3.11-ubuntu20.04-aws
-  - ghcr.io/databricks-mosaic/pytorch:2.4.1_cu124-python3.11-ubuntu20.04-aws
+  - mosaicml/pytorch:2.4.1_cu124-python3.11-ubuntu22.04-aws
+  - ghcr.io/databricks-mosaic/pytorch:2.4.1_cu124-python3.11-ubuntu22.04-aws
   TARGET: pytorch_stage
   TORCHVISION_VERSION: 0.19.1
 - AWS_OFI_NCCL_VERSION: ''
@@ -91,12 +91,12 @@
   PYTORCH_NIGHTLY_VERSION: ''
   PYTORCH_VERSION: 2.4.1
   TAGS:
-  - mosaicml/pytorch:2.4.1_cpu-python3.11-ubuntu20.04
-  - ghcr.io/databricks-mosaic/pytorch:2.4.1_cpu-python3.11-ubuntu20.04
+  - mosaicml/pytorch:2.4.1_cpu-python3.11-ubuntu22.04
+  - ghcr.io/databricks-mosaic/pytorch:2.4.1_cpu-python3.11-ubuntu22.04
   TARGET: pytorch_stage
   TORCHVISION_VERSION: 0.19.1
 - AWS_OFI_NCCL_VERSION: ''
-  BASE_IMAGE: nvidia/cuda:12.1.1-cudnn8-devel-ubuntu20.04
+  BASE_IMAGE: nvidia/cuda:12.1.1-cudnn8-devel-ubuntu22.04
   CUDA_VERSION: 12.1.1
   IMAGE_NAME: torch-2-3-1-cu121
   MOFED_VERSION: latest-23.10
@@ -119,12 +119,12 @@
   PYTORCH_NIGHTLY_VERSION: ''
   PYTORCH_VERSION: 2.3.1
   TAGS:
-  - mosaicml/pytorch:2.3.1_cu121-python3.11-ubuntu20.04
-  - ghcr.io/databricks-mosaic/pytorch:2.3.1_cu121-python3.11-ubuntu20.04
+  - mosaicml/pytorch:2.3.1_cu121-python3.11-ubuntu22.04
+  - ghcr.io/databricks-mosaic/pytorch:2.3.1_cu121-python3.11-ubuntu22.04
   TARGET: pytorch_stage
   TORCHVISION_VERSION: 0.18.1
 - AWS_OFI_NCCL_VERSION: v1.11.0-aws
-  BASE_IMAGE: nvidia/cuda:12.1.1-cudnn8-devel-ubuntu20.04
+  BASE_IMAGE: nvidia/cuda:12.1.1-cudnn8-devel-ubuntu22.04
   CUDA_VERSION: 12.1.1
   IMAGE_NAME: torch-2-3-1-cu121-aws
   MOFED_VERSION: ''
@@ -147,8 +147,8 @@
   PYTORCH_NIGHTLY_VERSION: ''
   PYTORCH_VERSION: 2.3.1
   TAGS:
-  - mosaicml/pytorch:2.3.1_cu121-python3.11-ubuntu20.04-aws
-  - ghcr.io/databricks-mosaic/pytorch:2.3.1_cu121-python3.11-ubuntu20.04-aws
+  - mosaicml/pytorch:2.3.1_cu121-python3.11-ubuntu22.04-aws
+  - ghcr.io/databricks-mosaic/pytorch:2.3.1_cu121-python3.11-ubuntu22.04-aws
   TARGET: pytorch_stage
   TORCHVISION_VERSION: 0.18.1
 - AWS_OFI_NCCL_VERSION: ''
@@ -162,12 +162,12 @@
   PYTORCH_NIGHTLY_VERSION: ''
   PYTORCH_VERSION: 2.3.1
   TAGS:
-  - mosaicml/pytorch:2.3.1_cpu-python3.11-ubuntu20.04
-  - ghcr.io/databricks-mosaic/pytorch:2.3.1_cpu-python3.11-ubuntu20.04
+  - mosaicml/pytorch:2.3.1_cpu-python3.11-ubuntu22.04
+  - ghcr.io/databricks-mosaic/pytorch:2.3.1_cpu-python3.11-ubuntu22.04
   TARGET: pytorch_stage
   TORCHVISION_VERSION: 0.18.1
 - AWS_OFI_NCCL_VERSION: ''
-  BASE_IMAGE: nvidia/cuda:12.4.1-cudnn-devel-ubuntu20.04
+  BASE_IMAGE: nvidia/cuda:12.4.1-cudnn-devel-ubuntu22.04
   COMPOSER_INSTALL_COMMAND: mosaicml[all]==0.27.0
   CUDA_VERSION: 12.4.1
   IMAGE_NAME: composer-0-27-0
diff --git a/docker/generate_build_matrix.py b/docker/generate_build_matrix.py
index 2d6c713a4a..1fb476b351 100644
--- a/docker/generate_build_matrix.py
+++ b/docker/generate_build_matrix.py
@@ -35,10 +35,10 @@ def _get_torchvision_version(pytorch_version: str):
 
 def _get_base_image(cuda_version: str):
     if not cuda_version:
-        return 'ubuntu:20.04'
+        return 'ubuntu:22.04'
     if cuda_version == '12.4.1':
-        return f'nvidia/cuda:12.4.1-cudnn-devel-ubuntu20.04'
-    return f'nvidia/cuda:{cuda_version}-cudnn8-devel-ubuntu20.04'
+        return f'nvidia/cuda:12.4.1-cudnn-devel-ubuntu22.04'
+    return f'nvidia/cuda:{cuda_version}-cudnn8-devel-ubuntu22.04'
 
 
 def _get_cuda_version(pytorch_version: str, use_cuda: bool):
@@ -112,8 +112,8 @@ def _get_pytorch_tags(python_version: str, pytorch_version: str, cuda_version: s
     tags = []
     cuda_version_tag = _get_cuda_version_tag(cuda_version)
     tags += [
-        f'{base_image_name}:{pytorch_version}_{cuda_version_tag}-python{python_version}-ubuntu20.04',
-        f'{ghcr_base_image_name}:{pytorch_version}_{cuda_version_tag}-python{python_version}-ubuntu20.04',
+        f'{base_image_name}:{pytorch_version}_{cuda_version_tag}-python{python_version}-ubuntu22.04',
+        f'{ghcr_base_image_name}:{pytorch_version}_{cuda_version_tag}-python{python_version}-ubuntu22.04',
     ]
 
     if python_version == PRODUCTION_PYTHON_VERSION and pytorch_version == PRODUCTION_PYTORCH_VERSION:
@@ -294,7 +294,7 @@ def _main():
                 interconnect = 'EFA'
         cuda_version = f"{entry['CUDA_VERSION']} ({interconnect})" if entry['CUDA_VERSION'] else 'cpu'
         table.append([
-            'Ubuntu 20.04',  # Linux distro
+            'Ubuntu 22.04',  # Linux distro
             'Base',  # Flavor
             entry['PYTORCH_VERSION'],  # Pytorch version
             cuda_version,  # Cuda version