bump_torch

.github/workflows/docker.yaml

@@ -17,14 +17,6 @@ jobs:
     strategy:
       matrix:
         include:
-        - name: "2.3.1_cu121"
-          base_image: mosaicml/pytorch:2.3.1_cu121-python3.11-ubuntu20.04
-          dep_groups: "[all]"
-          te_commit: b5a7c9f
-        - name: "2.3.1_cu121_aws"
-          base_image: mosaicml/pytorch:2.3.1_cu121-python3.11-ubuntu20.04-aws
-          dep_groups: "[all]"
-          te_commit: b5a7c9f
         - name: "2.4.0_cu124"
           base_image: mosaicml/pytorch:2.4.0_cu124-python3.11-ubuntu20.04
           dep_groups: "[all]"

.github/workflows/pr-cpu.yaml

@@ -20,9 +20,9 @@ jobs:
     strategy:
       matrix:
         include:
-        - name: "cpu-2.3.1"
+        - name: "cpu-2.4.0"
           pip_deps: "[all-cpu]"
-          container: mosaicml/pytorch:2.3.1_cpu-python3.11-ubuntu20.04
+          container: mosaicml/pytorch:2.4.0_cpu-python3.11-ubuntu20.04
           markers: "not gpu"
           pytest_command: "coverage run -m pytest"
     steps:

.github/workflows/pr-gpu.yaml

@@ -22,8 +22,8 @@ jobs:
       fail-fast: false
       matrix:
         include:
-        - name: "gpu-2.3.1-1"
-          container: mosaicml/llm-foundry:2.3.1_cu121-latest
+        - name: "gpu-2.4.0-1"
+          container: mosaicml/llm-foundry:2.4.0_cu124-latest
           markers: "gpu"
           pip_deps: "[all]"
           pytest_command: "coverage run -m pytest"
@@ -51,8 +51,8 @@ jobs:
       fail-fast: false
       matrix:
         include:
-        - name: "gpu-2.3.1-2"
-          container: mosaicml/llm-foundry:2.3.1_cu121-latest
+        - name: "gpu-2.4.0-2"
+          container: mosaicml/llm-foundry:2.4.0_cu124-latest
           markers: "gpu"
           pip_deps: "[all]"
           pytest_command: "coverage run -m pytest"
@@ -80,8 +80,8 @@ jobs:
       fail-fast: false
       matrix:
         include:
-        - name: "gpu-2.3.1-4"
-          container: mosaicml/llm-foundry:2.3.1_cu121-latest
+        - name: "gpu-2.4.0-4"
+          container: mosaicml/llm-foundry:2.4.0_cu124-latest
           markers: "gpu"
           pip_deps: "[all]"
           pytest_command: "coverage run -m pytest"

README.md

@@ -107,30 +107,30 @@ Something missing? Contribute with a PR!
 
 
 # Hardware and Software Requirements
-This codebase has been tested with PyTorch 2.2 with NVIDIA A100s and H100s.
+This codebase has been tested with PyTorch 2.4 with NVIDIA A100s and H100s.
 This codebase may also work on systems with other devices, such as consumer NVIDIA cards and AMD cards, but we are not actively testing these systems.
 If you have success/failure using LLM Foundry on other systems, please let us know in a Github issue and we will update the support matrix!
 
 | Device         | Torch Version | Cuda Version | Status                       |
 | -------------- | ------------- | ------------ | ---------------------------- |
-| A100-40GB/80GB | 2.3.1         | 12.1         | :white_check_mark: Supported |
-| H100-80GB      | 2.3.1         | 12.1         | :white_check_mark: Supported |
+| A100-40GB/80GB | 2.4.0         | 12.4         | :white_check_mark: Supported |
+| H100-80GB      | 2.4.0         | 12.4         | :white_check_mark: Supported |
 
 ## MosaicML Docker Images
 We highly recommend using our prebuilt Docker images. You can find them here: https://hub.docker.com/orgs/mosaicml/repositories.
 
 The `mosaicml/pytorch` images are pinned to specific PyTorch and CUDA versions, and are stable and rarely updated.
 
 The `mosaicml/llm-foundry` images are built with new tags upon every commit to the `main` branch.
-You can select a specific commit hash such as `mosaicml/llm-foundry:2.3.1_cu121-36ab1ba` or take the latest one using `mosaicml/llm-foundry:2.3.1_cu121-latest`.
+You can select a specific commit hash such as `mosaicml/llm-foundry:2.4.0_cu124-36ab1ba` or take the latest one using `mosaicml/llm-foundry:2.4.0_cu124-latest`.
 
 **Please Note:** The `mosaicml/llm-foundry` images do not come with the `llm-foundry` package preinstalled, just the dependencies. You will still need to `pip install llm-foundry` either from PyPi or from source.
 
 | Docker Image                                           | Torch Version | Cuda Version      | LLM Foundry dependencies installed? |
 | ------------------------------------------------------ | ------------- | ----------------- | ----------------------------------- |
-| `mosaicml/pytorch:2.3.1_cu121-python3.11-ubuntu20.04`  | 2.3.1         | 12.1 (Infiniband) | No                                  |
-| `mosaicml/llm-foundry:2.3.1_cu121-latest`              | 2.3.1         | 12.1 (Infiniband) | Yes                                 |
-| `mosaicml/llm-foundry:2.3.1_cu121_aws-latest`          | 2.3.1         | 12.1 (EFA)        | Yes                                 |
+| `mosaicml/pytorch:2.4.0_cu124-python3.11-ubuntu20.04`  | 2.4.0         | 12.4 (Infiniband) | No                                  |
+| `mosaicml/llm-foundry:2.4.0_cu124-latest`              | 2.4.0         | 12.4 (Infiniband) | Yes                                 |
+| `mosaicml/llm-foundry:2.4.0_cu124_aws-latest`          | 2.4.0         | 12.4 (EFA)        | Yes                                 |
 
 
 # Installation

mcli/mcli-1b-eval.yaml

@@ -9,7 +9,7 @@ integrations:
 command: |
   cd llm-foundry/scripts/
   composer eval/eval.py /mnt/config/parameters.yaml
-image: mosaicml/llm-foundry:2.3.1_cu121-latest
+image: mosaicml/llm-foundry:2.4.0_cu124-latest
 name: mpt-1b-eval
 
 compute:

mcli/mcli-1b-max-seq-len-8k.yaml

@@ -17,7 +17,7 @@ command: |
     --out_root ./my-copy-c4 --splits train_small val_small \
     --concat_tokens 8192 --tokenizer EleutherAI/gpt-neox-20b --eos_text '<|endoftext|>'
   composer train/train.py /mnt/config/parameters.yaml
-image: mosaicml/llm-foundry:2.3.1_cu121-latest
+image: mosaicml/llm-foundry:2.4.0_cu124-latest
 name: mpt-1b-ctx-8k-gpus-8
 
 compute:

mcli/mcli-1b.yaml

@@ -21,7 +21,7 @@ command: |
     eval_loader.dataset.split=val_small \
     max_duration=100ba \
     eval_interval=0
-image: mosaicml/llm-foundry:2.3.1_cu121-latest
+image: mosaicml/llm-foundry:2.4.0_cu124-latest
 name: mpt-1b-gpus-8
 
 compute:

mcli/mcli-benchmark-mpt.yaml

@@ -6,7 +6,7 @@ compute:
   # cluster: TODO # Name of the cluster to use for this run
   # gpu_type: a100_80gb # Type of GPU to use. We use a100_80gb in our experiments
 
-image: mosaicml/llm-foundry:2.3.1_cu121-latest
+image: mosaicml/llm-foundry:2.4.0_cu124-latest
 
 integrations:
 - integration_type: git_repo

mcli/mcli-convert-composer-to-hf.yaml

@@ -13,7 +13,7 @@ command: |
     --hf_output_path s3://bucket/folder/hf/ \
     --output_precision bf16 \
 
-image: mosaicml/llm-foundry:2.3.1_cu121-latest
+image: mosaicml/llm-foundry:2.4.0_cu124-latest
 name: convert-composer-hf
 
 compute:

mcli/mcli-hf-eval.yaml

@@ -16,7 +16,7 @@ gpu_num: 8
 # gpu_type:
 # cluster:  # replace with your cluster here!
 
-image: mosaicml/llm-foundry:2.3.1_cu121-latest
+image: mosaicml/llm-foundry:2.4.0_cu124-latest
 
 # The below is injected as a YAML file: /mnt/config/parameters.yaml
 parameters:

mcli/mcli-hf-generate.yaml

@@ -35,7 +35,7 @@ command: |
       "Here's a quick recipe for baking chocolate chip cookies: Start by" \
       "The best 5 cities to visit in Europe are"
 
-image: mosaicml/llm-foundry:2.3.1_cu121-latest
+image: mosaicml/llm-foundry:2.4.0_cu124-latest
 name: hf-generate
 
 compute:

mcli/mcli-llama2-finetune.yaml

@@ -9,7 +9,7 @@ integrations:
 command: |
   cd llm-foundry/scripts
   composer train/train.py /mnt/config/parameters.yaml
-image: mosaicml/llm-foundry:2.3.1_cu121-latest
+image: mosaicml/llm-foundry:2.4.0_cu124-latest
 name: llama2-finetune
 
 compute:

mcli/mcli-openai-eval.yaml

@@ -16,7 +16,7 @@ gpu_num:  #
 gpu_type:  #
 cluster:  # replace with your cluster here!
 
-image: mosaicml/llm-foundry:2.3.1_cu121-latest
+image: mosaicml/llm-foundry:2.4.0_cu124-latest
 
 # The below is injected as a YAML file: /mnt/config/parameters.yaml
 parameters:

mcli/mcli-pretokenize-oci-upload.yaml

@@ -1,5 +1,5 @@
 name: c4-2k-pre-tokenized
-image: mosaicml/llm-foundry:2.3.1_cu121-latest
+image: mosaicml/llm-foundry:2.4.0_cu124-latest
 compute:
   gpus: 8  # Number of GPUs to use
 

setup.py

@@ -57,7 +57,7 @@
     'accelerate>=0.25,<0.34',  # for HF inference `device_map`
     'transformers>=4.43.2,<4.44',
     'mosaicml-streaming>=0.8.1,<0.9',
-    'torch>=2.3.0,<2.4.1',
+    'torch>=2.4.0,<2.4.1',
     'datasets>=2.19,<2.20',
     'fsspec==2023.6.0',  # newer version results in a bug in datasets that duplicates data
     'sentencepiece==0.2.0',
@@ -118,8 +118,8 @@
 ]
 
 extra_deps['megablocks'] = [
-    'megablocks==0.5.1',
-    'grouped-gemm==0.1.4',
+    'megablocks==0.6.1',
+    'grouped-gemm==0.1.6',
 ]
 
 extra_deps['databricks-serverless'] = {