From 099cfc527aea6e33df707c00192808d4910193fc Mon Sep 17 00:00:00 2001 From: Rajesh Thallam Date: Wed, 17 Sep 2025 10:23:17 -0700 Subject: [PATCH 1/2] Fix DeepSeek 671B recipe to read from quantized checkpoints --- inference/trillium/JetStream-Maxtext/DeepSeek-R1-671B/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 ``` From e21e564804c4aea773490425177d15c02886a940 Mon Sep 17 00:00:00 2001 From: Rajesh Thallam Date: Fri, 19 Sep 2025 15:12:21 -0700 Subject: [PATCH 2/2] For DeepSeek inference pin JAX/MaxText to specific commits --- .../JetStream-Maxtext/DeepSeek-R1-671B/docker/Dockerfile | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) 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" ]