add megablocks, move things

deploy/Dockerfile

@@ -1,6 +1,4 @@
-FROM --platform=amd64 nvcr.io/nvidia/cuda:11.8.0-devel-ubuntu22.04 as base
-
-ARG MAX_JOBS
+FROM --platform=amd64 nvcr.io/nvidia/cuda:12.1.0-devel-ubuntu22.04 as base
 
 WORKDIR /workspace
 
@@ -13,21 +11,21 @@ RUN apt update && \
 # https://developer.nvidia.com/cuda-gpus
 ENV TORCH_CUDA_ARCH_LIST "7.0;7.2;7.5;8.0;8.6;8.9;9.0"
 
-# We have to manually install Torch otherwise apex & xformers won't build
-RUN pip3 install "torch>=2.0.0"
-# To enable H100 PCIe support, install PyTorch >=2.2.0 by uncommenting the following line
-# RUN pip3 install "torch==2.2.0.dev20231018+cu118" --index-url https://download.pytorch.org/whl/nightly/cu118
+RUN pip3 install "torch==2.1.1"
 
 # This build is slow but NVIDIA does not provide binaries. Increase MAX_JOBS as needed.
+RUN pip3 install "git+https://github.com/stanford-futuredata/megablocks.git"
+RUN pip3 install "git+https://github.com/vllm-project/vllm.git"
+RUN pip3 install "xformers==0.0.23" "transformers==4.36.0" "fschat[model_worker]==0.2.34"
+
 RUN git clone https://github.com/NVIDIA/apex && \
     cd apex && git checkout 2386a912164b0c5cfcd8be7a2b890fbac5607c82 && \
     sed -i '/check_cuda_torch_binary_vs_bare_metal(CUDA_HOME)/d' setup.py && \
     python3 setup.py install --cpp_ext --cuda_ext
 
-RUN pip3 install "xformers==0.0.22" "transformers==4.34.0" "vllm==0.2.0" "fschat[model_worker]==0.2.30"
 
 COPY entrypoint.sh .
 
 RUN chmod +x /workspace/entrypoint.sh
 
-ENTRYPOINT ["/workspace/entrypoint.sh"]
+ENTRYPOINT ["/workspace/entrypoint.sh"]