17 lines
428 B
Docker
17 lines
428 B
Docker
ARG BASE_TAG=main
|
|
FROM winglian/axolotl-base:$BASE_TAG
|
|
|
|
ARG TORCH_CUDA_ARCH_LIST="7.0 7.5 8.0 8.6+PTX"
|
|
|
|
RUN apt-get update && \
|
|
apt-get install -y vim
|
|
|
|
WORKDIR /workspace
|
|
|
|
# The base image ships with `pydantic==1.8.2` which is not working
|
|
RUN python3 -m pip install -U --no-cache-dir pydantic
|
|
|
|
ARG REF=main
|
|
RUN git clone https://github.com/winglian/axolotl && cd axolotl && git checkout $REF && \
|
|
pip install -e .[int4]
|