2022-01-07 06:54:04 +00:00
|
|
|
FROM nvcr.io/nvidia/pytorch:21.07-py3
|
|
|
|
|
|
|
|
|
|
# install dependencies
|
2022-01-18 05:35:18 +00:00
|
|
|
RUN pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple \
|
|
|
|
|
&& pip install -U pip setuptools wheel \
|
2022-01-07 06:54:04 +00:00
|
|
|
&& pip install pytest tensorboard deepspeed apex
|
|
|
|
|
|
|
|
|
|
# install colossalai
|
|
|
|
|
RUN git clone https://github.com/hpcaitech/ColossalAI.git \
|
|
|
|
|
&& cd ./ColossalAI \
|
2022-02-14 09:09:30 +00:00
|
|
|
&& pip install -v --no-cache-dir .
|
