Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Bump torch 2.5 #1665

Draft
wants to merge 2 commits into
base: main
Choose a base branch
from
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 4 additions & 4 deletions .github/workflows/docker.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -17,11 +17,11 @@ jobs:
strategy:
matrix:
include:
- name: "2.4.0_cu124"
base_image: mosaicml/pytorch:2.4.0_cu124-python3.11-ubuntu20.04
- name: "2.5.1_cu124"
base_image: mosaicml/pytorch:2.5.1_cu124-python3.11-ubuntu20.04
dep_groups: "[all]"
- name: "2.4.0_cu124_aws"
base_image: mosaicml/pytorch:2.4.0_cu124-python3.11-ubuntu20.04-aws
- name: "2.5.1_cu124_aws"
base_image: mosaicml/pytorch:2.5.1_cu124-python3.11-ubuntu20.04-aws
dep_groups: "[all]"
steps:

Expand Down
4 changes: 2 additions & 2 deletions .github/workflows/pr-cpu.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -21,9 +21,9 @@ jobs:
strategy:
matrix:
include:
- name: "cpu-2.4.0"
- name: "cpu-2.5.1"
pip_deps: "[all-cpu]"
container: mosaicml/pytorch:2.4.0_cpu-python3.11-ubuntu20.04
container: mosaicml/pytorch:2.5.1_cpu-python3.11-ubuntu20.04
markers: "not gpu"
pytest_command: "coverage run -m pytest"
steps:
Expand Down
12 changes: 6 additions & 6 deletions .github/workflows/pr-gpu.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -22,8 +22,8 @@ jobs:
fail-fast: false
matrix:
include:
- name: "gpu-2.4.0-1"
container: mosaicml/llm-foundry:2.4.0_cu124-latest
- name: "gpu-2.5.1-1"
container: mosaicml/llm-foundry:2.5.1_cu124-latest
markers: "gpu"
pip_deps: "[all]"
pytest_command: "coverage run -m pytest"
Expand Down Expand Up @@ -51,8 +51,8 @@ jobs:
fail-fast: false
matrix:
include:
- name: "gpu-2.4.0-2"
container: mosaicml/llm-foundry:2.4.0_cu124-latest
- name: "gpu-2.5.1-2"
container: mosaicml/llm-foundry:2.5.1_cu124-latest
markers: "gpu"
pip_deps: "[all]"
pytest_command: "coverage run -m pytest"
Expand Down Expand Up @@ -80,8 +80,8 @@ jobs:
fail-fast: false
matrix:
include:
- name: "gpu-2.4.0-4"
container: mosaicml/llm-foundry:2.4.0_cu124-latest
- name: "gpu-2.5.1-4"
container: mosaicml/llm-foundry:2.5.1_cu124-latest
markers: "gpu"
pip_deps: "[all]"
pytest_command: "coverage run -m pytest"
Expand Down
4 changes: 2 additions & 2 deletions .github/workflows/release.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -93,7 +93,7 @@ jobs:
${{ env.AWS_DOCKER_TAG }}
${{ env.AWS_LATEST_TAG }}
build-args: |
BASE_IMAGE=mosaicml/pytorch:2.4.0_cu124-python3.11-ubuntu20.04-aws
BASE_IMAGE=mosaicml/pytorch:2.5.1_cu124-python3.11-ubuntu20.04-aws
BRANCH_NAME=${{ env.BRANCH_NAME }}
DEP_GROUPS=[all]
KEEP_FOUNDRY=true
Expand All @@ -108,7 +108,7 @@ jobs:
${{ env.DOCKER_TAG }}
${{ env.LATEST_TAG }}
build-args: |
BASE_IMAGE=mosaicml/pytorch:2.4.0_cu124-python3.11-ubuntu20.04
BASE_IMAGE=mosaicml/pytorch:2.5.1_cu124-python3.11-ubuntu20.04
BRANCH_NAME=${{ env.BRANCH_NAME }}
DEP_GROUPS=[all]
KEEP_FOUNDRY=true
12 changes: 6 additions & 6 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -113,24 +113,24 @@ If you have success/failure using LLM Foundry on other systems, please let us kn

| Device | Torch Version | Cuda Version | Status |
| -------------- | ------------- | ------------ | ---------------------------- |
| A100-40GB/80GB | 2.4.0 | 12.4 | :white_check_mark: Supported |
| H100-80GB | 2.4.0 | 12.4 | :white_check_mark: Supported |
| A100-40GB/80GB | 2.5.1 | 12.4 | :white_check_mark: Supported |
| H100-80GB | 2.5.1 | 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.4.0_cu124-36ab1ba` or take the latest one using `mosaicml/llm-foundry:2.4.0_cu124-latest`.
You can select a specific commit hash such as `mosaicml/llm-foundry:2.5.1_cu124-36ab1ba` or take the latest one using `mosaicml/llm-foundry:2.5.1_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.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 |
| `mosaicml/pytorch:2.5.1_cu124-python3.11-ubuntu20.04` | 2.5.1 | 12.4 (Infiniband) | No |
| `mosaicml/llm-foundry:2.5.1_cu124-latest` | 2.5.1 | 12.4 (Infiniband) | Yes |
| `mosaicml/llm-foundry:2.5.1_cu124_aws-latest` | 2.5.1 | 12.4 (EFA) | Yes |


# Installation
Expand Down
10 changes: 5 additions & 5 deletions setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -52,12 +52,12 @@
]

install_requires = [
'mosaicml[libcloud,wandb,oci,gcs,mlflow]>=0.26.0,<0.27',
'mosaicml[libcloud,wandb,oci,gcs,mlflow]>=0.27.0,<0.28',
'mlflow>=2.14.1,<2.18',
'accelerate>=0.25,<1.2', # for HF inference `device_map`
'transformers>=4.43.2,<4.47',
'mosaicml-streaming>=0.9.0,<0.10',
'torch>=2.4.0,<2.4.1',
'torch>=2.5.1,<2.5.2',
'datasets>=2.20.0,<2.21',
'fsspec==2023.6.0', # newer version results in a bug in datasets that duplicates data
'sentencepiece==0.2.0',
Expand Down Expand Up @@ -91,15 +91,15 @@
]

extra_deps['databricks'] = [
'mosaicml[databricks]>=0.26.0,<0.27',
'mosaicml[databricks]>=0.27.0,<0.28',
'numpy<2',
'databricks-sql-connector>=3,<4',
'databricks-connect==14.1.0',
'lz4>=4,<5',
]

extra_deps['tensorboard'] = [
'mosaicml[tensorboard]>=0.26.0,<0.27',
'mosaicml[tensorboard]>=0.27.0,<0.28',
]

# Flash 2 group kept for backwards compatibility
Expand All @@ -110,7 +110,7 @@
extra_deps['gpu'] = copy.deepcopy(extra_deps['gpu-flash2'])

extra_deps['peft'] = [
'mosaicml[peft]>=0.26.0,<0.27',
'mosaicml[peft]>=0.27.0,<0.28',
]

extra_deps['openai'] = [
Expand Down
Loading