mirror of https://github.com/hpcaitech/ColossalAI
You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
11 lines
323 B
11 lines
323 B
3 years ago
|
FROM nvcr.io/nvidia/pytorch:21.07-py3
|
||
|
|
|
||
|
|
# install dependencies
|
||
|
|
RUN pip install -U pip setuptools wheel \
|
||
|
|
&& pip install pytest tensorboard deepspeed apex
|
||
|
|
|
||
|
|
# install colossalai
|
||
|
|
RUN git clone https://github.com/hpcaitech/ColossalAI.git \
|
||
|
|
&& cd ./ColossalAI \
|
||
|
|
&& pip install -v --no-cache-dir --global-option="--cuda_ext" .
|