Add pytorch dependencies to CI Dockerfile and enable test (#503)

alpa-projects · Jun 12, 2022 · ef6df24 · ef6df24
1 parent b47da7b
commit ef6df24
Show file tree

Hide file tree

Showing 6 changed files with 13 additions and 1 deletion.
diff --git a/docker/unittest.Dockerfile b/docker/unittest.Dockerfile
@@ -23,6 +23,18 @@ RUN source python3.9-env/bin/activate && pip install --upgrade pip \
   tqdm scipy numba pulp tensorstore prospector yapf coverage cmake  \
     pybind11 ray[default] matplotlib
 
+# Install PyTorch dependencies
+RUN git clone https://github.com/pytorch/functorch /functorch
+RUN source python3.7-env/bin/activate \
+  && pip install torch torchdistx --pre --extra-index-url https://download.pytorch.org/whl/nightly/cpu \
+  && pushd /functorch && python setup.py install && popd
+RUN source python3.8-env/bin/activate \
+  && pip install torch torchdistx --pre --extra-index-url https://download.pytorch.org/whl/nightly/cpu \
+  && pushd /functorch && python setup.py install && popd
+RUN source python3.9-env/bin/activate \
+  && pip install torch torchdistx --pre --extra-index-url https://download.pytorch.org/whl/nightly/cpu \
+  && pushd /functorch && python setup.py install && popd
+
 # We determine the CUDA version at `docker build ...` phase
 ARG JAX_CUDA_VERSION=11.1
 COPY scripts/install_cuda.sh /install_cuda.sh

diff --git a/docs/install.rst b/docs/install.rst
@@ -144,7 +144,7 @@ To enable Alpa for PyTorch, install the following dependencies:
 
     # Install nightly version of torch and torchdistx
     pip3 uninstall -y torch torchdistx
-    pip install torchdistx --pre --extra-index-url https://download.pytorch.org/whl/nightly/cpu
+    pip install torch torchdistx --pre --extra-index-url https://download.pytorch.org/whl/nightly/cpu
 
     # Build functorch from source
     git clone https://github.com/pytorch/functorch

diff --git a/tests/torch/test_torch_dict_input.py → tests/test_torch_dict_input.py b/tests/torch/test_torch_dict_input.py → tests/test_torch_dict_input.py
diff --git a/tests/torch/test_torch_reshape.py → tests/test_torch_reshape.py b/tests/torch/test_torch_reshape.py → tests/test_torch_reshape.py
diff --git a/tests/torch/test_torch_simple.py → tests/test_torch_simple.py b/tests/torch/test_torch_simple.py → tests/test_torch_simple.py
diff --git a/tests/torch/test_torch_zhen.py → tests/test_torch_zhen.py b/tests/torch/test_torch_zhen.py → tests/test_torch_zhen.py