diff --git a/Dockerfile.ci.dev b/Dockerfile.ci.dev index cd879b1bbc..e6073c1713 100644 --- a/Dockerfile.ci.dev +++ b/Dockerfile.ci.dev @@ -23,6 +23,14 @@ RUN apt-get update && \ wget https://github.com/mikefarah/yq/releases/download/v4.44.1/yq_linux_amd64 -O /usr/local/bin/yq && \ chmod a+x /usr/local/bin/yq +COPY --from=build_causal_conv1d /opt/causal_conv1d-*.whl ./ +COPY --from=build_grouped_gemm /opt/grouped_gemm-*.whl ./ +COPY --from=build_mamba_ssm /opt/mamba_ssm-*.whl ./ + +RUN pip install causal_conv1d-*.whl \ + mamba_ssm-*.whl \ + grouped_gemm-*.whl + # Since megatron does not have any dependencies (and isn't a dependency to any other package), we can install it separately to make everything a bit quicker ARG MCORE_REPO ARG MCORE_REF @@ -47,14 +55,6 @@ git checkout $MCORE_BACKWARDS_REF rm -rf megatron; cp -a /opt/megatron-lm/megatron ./ EOF -COPY --from=build_causal_conv1d /opt/causal_conv1d-*.whl ./ -COPY --from=build_grouped_gemm /opt/grouped_gemm-*.whl ./ -COPY --from=build_mamba_ssm /opt/mamba_ssm-*.whl ./ - -RUN pip install causal_conv1d-*.whl \ - mamba_ssm-*.whl \ - grouped_gemm-*.whl - RUN PY_ENV=pytorch:24.07 pip install -e /opt/megatron-lm ENV PYTHONPATH="/opt/megatron-lm:$PYTHONPATH" ENV NVTE_FLASH_ATTN=0 diff --git a/Dockerfile.ci.lts b/Dockerfile.ci.lts index efc9ba470e..af4698dae5 100644 --- a/Dockerfile.ci.lts +++ b/Dockerfile.ci.lts @@ -24,6 +24,14 @@ RUN apt-get update && \ wget https://github.com/mikefarah/yq/releases/download/v4.44.1/yq_linux_amd64 -O /usr/local/bin/yq && \ chmod a+x /usr/local/bin/yq +COPY --from=build_causal_conv1d /opt/causal_conv1d-*.whl ./ +COPY --from=build_grouped_gemm /opt/grouped_gemm-*.whl ./ +COPY --from=build_mamba_ssm /opt/mamba_ssm-*.whl ./ + +RUN pip install causal_conv1d-*.whl \ + mamba_ssm-*.whl \ + grouped_gemm-*.whl + # Since megatron does not have any dependencies (and isn't a dependency to any other package), we can install it separately to make everything a bit quicker ARG MCORE_REPO ARG MCORE_REF @@ -48,14 +56,6 @@ git checkout $MCORE_BACKWARDS_REF rm -rf megatron; cp -a /opt/megatron-lm/megatron ./ EOF -COPY --from=build_causal_conv1d /opt/causal_conv1d-*.whl ./ -COPY --from=build_grouped_gemm /opt/grouped_gemm-*.whl ./ -COPY --from=build_mamba_ssm /opt/mamba_ssm-*.whl ./ - -RUN pip install causal_conv1d-*.whl \ - mamba_ssm-*.whl \ - grouped_gemm-*.whl - RUN PY_ENV=pytorch:24.01 \ CAUSAL_CONV1D_FORCE_BUILD=TRUE \ MAMBA_FORCE_BUILD=TRUE \