diff --git a/inference/trillium/JetStream-Maxtext/DeepSeek-R1-671B/README.md b/inference/trillium/JetStream-Maxtext/DeepSeek-R1-671B/README.md index 3e80a53..d69baea 100644 --- a/inference/trillium/JetStream-Maxtext/DeepSeek-R1-671B/README.md +++ b/inference/trillium/JetStream-Maxtext/DeepSeek-R1-671B/README.md @@ -363,7 +363,7 @@ helm install -f values.yaml \ --set clusterName=$CLUSTER_NAME \ --set job.jax_tpu_image.repository=${ARTIFACT_REGISTRY}/${JETSTREAM_MAXTEXT_IMAGE} \ --set job.jax_tpu_image.tag=${JETSTREAM_MAXTEXT_VERSION} \ ---set maxtext_config.load_parameters_path=${GCS_CKPT_PATH_UNSCANNED}/0/items \ +--set maxtext_config.load_parameters_path=${GCS_CKPT_PATH_QUANTIZED} \ jetstream-pathways \ $RECIPE_ROOT/serve-model ``` diff --git a/inference/trillium/JetStream-Maxtext/DeepSeek-R1-671B/docker/Dockerfile b/inference/trillium/JetStream-Maxtext/DeepSeek-R1-671B/docker/Dockerfile index 9efcb66..17903b8 100644 --- a/inference/trillium/JetStream-Maxtext/DeepSeek-R1-671B/docker/Dockerfile +++ b/inference/trillium/JetStream-Maxtext/DeepSeek-R1-671B/docker/Dockerfile @@ -44,8 +44,11 @@ ENV XCLOUD_ENVIRONMENT=GCP # Install JetStream and MaxText RUN git clone https://github.com/AI-Hypercomputer/JetStream.git && \ -git clone https://github.com/AI-Hypercomputer/maxtext.git && \ -git clone https://github.com/google/aqt.git + cd JetStream && git checkout 29329e8 && cd .. && \ + git clone https://github.com/AI-Hypercomputer/maxtext.git && \ + cd maxtext && git checkout 4f63ce61 && cd .. && \ + git clone https://github.com/google/aqt.git && \ + cd aqt && git checkout 28504f5 && cd .. RUN cd /maxtext && bash setup.sh && pip install torch --index-url https://download.pytorch.org/whl/cpu @@ -56,4 +59,4 @@ RUN cd /JetStream && pip install -e . RUN apt -y update && apt-get -y install python3-dev && apt-get -y install build-essential RUN cp -r /aqt/aqt/* /usr/local/lib/python3.10/dist-packages/aqt/ -ENTRYPOINT [ "/bin/bash" ] \ No newline at end of file +ENTRYPOINT [ "/bin/bash" ]