Merge pull request #90 from RajeshThallam/main

Fix DeepSeek 671B recipe to read from quantized checkpoints

Author: vipannalla

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
 ```

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" ]
+ENTRYPOINT [ "/bin/bash" ]