tocmo0nlord/axolotl
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

experiment with nvcr pytorch image for torch 2.5.1

Browse Source
This commit is contained in:
Wing Lian
2024-12-06 11:07:27 -05:00
parent 5bef19064b
commit 684b543aa1
3 changed files with 15 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
docker/Dockerfile
View File

@@ -1,5 +1,6 @@
ARG BASE_IMAGE=axolotlai/axolotl-base
ARG BASE_TAG=main-base
FROM axolotlai/axolotl-base:$BASE_TAG
FROM $BASE_IMAGE:$BASE_TAG
ARG TORCH_CUDA_ARCH_LIST="7.0 7.5 8.0 8.6+PTX"
ARG AXOLOTL_EXTRAS=""

4
docker/Dockerfile-base
View File

@@ -3,7 +3,9 @@ ARG CUDNN_VERSION="8"
ARG UBUNTU_VERSION="22.04"
ARG MAX_JOBS=4
FROM nvidia/cuda:$CUDA_VERSION-cudnn$CUDNN_VERSION-devel-ubuntu$UBUNTU_VERSION AS base-builder
ARG BASE_IMAGE=nvidia/cuda
ARG BASE_TAG=""
FROM $BASE_IMAGE:${BASE_TAG:-$CUDA_VERSION-cudnn$CUDNN_VERSION-devel-ubuntu$UBUNTU_VERSION} AS base-builder
ENV PATH="/root/miniconda3/bin:${PATH}"