From 1c83a1a02081834eb84afa7292eef479133541fa Mon Sep 17 00:00:00 2001 From: NanoCode012 Date: Thu, 22 May 2025 19:18:27 +0700 Subject: [PATCH] feat(doc): clarify minimum pytorch and cuda to use blackwell (#2704) [skip ci] --- docs/docker.qmd | 4 ++++ docs/installation.qmd | 8 ++++++++ 2 files changed, 12 insertions(+) diff --git a/docs/docker.qmd b/docs/docker.qmd index e208d3222..d665eaf5b 100644 --- a/docs/docker.qmd +++ b/docs/docker.qmd @@ -8,6 +8,10 @@ format: This section describes the different Docker images that are released by AxolotlAI at [Docker Hub](https://hub.docker.com/u/axolotlai). +::: {.callout-important} +For Blackwell GPUs, please use the tags with Pytorch 2.7.0 and CUDA 12.8. +::: + ## Base The base image is the most minimal image that can install Axolotl. It is based on the `nvidia/cuda` image. It includes python, torch, git, git-lfs, awscli, pydantic, and more. diff --git a/docs/installation.qmd b/docs/installation.qmd index 0cf5ffceb..b429992b6 100644 --- a/docs/installation.qmd +++ b/docs/installation.qmd @@ -25,6 +25,10 @@ Please make sure to have Pytorch installed before installing Axolotl in your loc Follow the instructions at: [https://pytorch.org/get-started/locally/](https://pytorch.org/get-started/locally/) ::: +::: {.callout-important} +For Blackwell GPUs, please use Pytorch 2.7.0 and CUDA 12.8. +::: + ### PyPI Installation (Recommended) {#sec-pypi} ```{.bash} @@ -72,6 +76,10 @@ docker run --privileged --gpus '"all"' --shm-size 10g --rm -it \ ``` ::: +::: {.callout-important} +For Blackwell GPUs, please use `axolotlai/axolotl:main-py3.11-cu128-2.7.0` or the cloud variant `axolotlai/axolotl-cloud:main-py3.11-cu128-2.7.0`. +::: + Please refer to the [Docker documentation](docker.qmd) for more information on the different Docker images that are available. ## Cloud Environments {#sec-cloud}