forked from unslothai/notebooks
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathDockerfile_DGX_Spark
More file actions
31 lines (26 loc) · 1.03 KB
/
Dockerfile_DGX_Spark
File metadata and controls
31 lines (26 loc) · 1.03 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
FROM nvcr.io/nvidia/pytorch:25.11-py3
# Set CUDA environment variables
ENV CUDA_HOME=/usr/local/cuda-13.0
ENV CUDA_PATH=$CUDA_HOME
ENV PATH=$CUDA_HOME/bin:$PATH
ENV LD_LIBRARY_PATH=$CUDA_HOME/lib64:$LD_LIBRARY_PATH
ENV C_INCLUDE_PATH=$CUDA_HOME/include:$C_INCLUDE_PATH
ENV CPLUS_INCLUDE_PATH=$CUDA_HOME/include:$CPLUS_INCLUDE_PATH
# Install triton from source for latest blackwell support
RUN git clone https://github.com/triton-lang/triton.git && \
cd triton && \
git checkout c5d671f91d90f40900027382f98b17a3e04045f6 && \
pip install -r python/requirements.txt && \
pip install . && \
cd ..
# Install xformers from source for blackwell support
RUN git clone --depth=1 https://github.com/facebookresearch/xformers --recursive && \
cd xformers && \
export TORCH_CUDA_ARCH_LIST="12.1" && \
python setup.py install && \
cd ..
# Install unsloth and other dependencies
RUN pip install --no-deps bitsandbytes==0.48.0 transformers==4.56.2 trl==0.22.2
RUN pip install unsloth unsloth_zoo
# Launch the shell
CMD ["/bin/bash"]