WIP

Author: yuyanpeng-google

.github/workflows/run_maxtext_jetstream_tests.yaml

@@ -50,37 +50,42 @@ jobs:
     - name: Build
       run: |
           pushd experimental/jetstream-maxtext-stable-stack
-          ./build.sh \
+          # ./build.sh \
             LOCAL_IMAGE_TAG="${LOCAL_IMAGE_TAG}"
           popd
     - name: Test
       run: |
           pushd experimental/jetstream-maxtext-stable-stack
-          ./test.sh \
+          # ./test.sh \
             LOCAL_IMAGE_TAG=${LOCAL_IMAGE_TAG}
           popd
     - name: Upload image
       run: |
           UPLOAD_IMAGE_TAG=gcr.io/cloud-ml-auto-solutions/${LOCAL_IMAGE_TAG}
-          docker tag ${LOCAL_IMAGE_TAG} ${UPLOAD_IMAGE_TAG}
-          docker push ${UPLOAD_IMAGE_TAG}
+          gcloud container images add-tag gcr.io/cloud-ml-auto-solutions/jetstream-maxtext-stable-stack/tpu:nightly ${UPLOAD_IMAGE_TAG} --quiet
+
 
   benchmark_report:
     name: Benchmark Report
     needs: build_stable_stack
     runs-on: ["self-hosted", "tpu", "v6e-8"]
-    container:
-      # sync with the image uploaded from build_stable_stack stage
-      image: gcr.io/cloud-ml-auto-solutions/jetstream-maxtext-stable-stack/tpu:github_${{ github.run_id }}
-      options: "--privileged"
     env:
-      OUTPUT_DIR: /workspace/test_dir/
+      OUTPUT_DIR: ./test_dir/
     steps:
-    - name: Create output directory # Ensure directory exists in container
-      run: mkdir -p ${OUTPUT_DIR}
     - name: Test MOEBenchmarks
-      # Report should generated in OUTPUT_DIR depend on ENV
-      run: bash JetStream/.github/workflows/test_moe_benchmarks.sh
+      run: |
+        mkdir -p ${OUTPUT_DIR}
+        # sync with the image uploaded from build_stable_stack stage
+        # Report should generated in OUTPUT_DIR depend on ENV
+        DOCKER_OUTPUT_DIR=/output
+        docker run \
+          -v ${OUTPUT_DIR}:${DOCKER_OUTPUT_DIR}} \
+          --env OUTPUT_DIR:${DOCKER_OUTPUT_DIR} \
+          --privileged --net=host --rm -i \
+          gcr.io/cloud-ml-auto-solutions/jetstream-maxtext-stable-stack/tpu:github_${{ github.run_id }} \
+          bash -c "
+            bash JetStream/.github/workflows/test_moe_benchmarks.sh
+          "
     - name: Upload build artifact
       uses: actions/upload-artifact@v4
       with:

.github/workflows/test_moe_benchmarks.sh

@@ -16,8 +16,8 @@ LIBTPU_INIT_ARGS="--xla_tpu_enable_windowed_einsum_for_reduce_scatter=false --xl
 tail -n5 ${OUTPUT_DIR}/moe_8x22b.txt > ${OUTPUT_DIR}/moe_8x22b.tmp && mv ${OUTPUT_DIR}/moe_8x22b.tmp ${OUTPUT_DIR}/moe_8x22b.txt
 
 # moe 8x22B 8k context length chunked prefill with 2k prefill chunk size
-LIBTPU_INIT_ARGS="--xla_tpu_enable_windowed_einsum_for_reduce_scatter=false --xla_jf_spmd_threshold_for_windowed_einsum_mib=1000000" python -m MaxText.benchmark_chunked_prefill MaxText/configs/inference.yml load_parameters_path=gs://jetstream-runner/8-22B-int8  max_prefill_predict_length=8192 max_target_length=9000 model_name=mixtral-8x22b ici_fsdp_parallelism=1 ici_autoregressive_parallelism=1 ici_tensor_parallelism=1 scan_layers=false per_device_batch_size=24 attention=dot_product megablox=False quantization=int8 checkpoint_is_quantized=True quantize_kvcache=False capacity_factor=1 tokenizer_path=assets/tokenizer.mistral-v3 inference_microbenchmark_prefill_lengths="8192" sparse_matmul=False model_call_mode=inference ici_context_autoregressive_parallelism=8 use_chunked_prefill=True prefill_chunk_size=2048 > ${OUTPUT_DIR}/moe_8x22b_long_context_8k_prefill.txt
-tail -n5 ${OUTPUT_DIR}/moe_8x22b_long_context_8k_prefill.txt > ${OUTPUT_DIR}/moe_8x22b_long_context_8k_prefill.tmp && mv ${OUTPUT_DIR}/moe_8x22b_long_context_8k_prefill.tmp ${OUTPUT_DIR}/moe_8x22b_long_context_8k_prefill.txt
+# LIBTPU_INIT_ARGS="--xla_tpu_enable_windowed_einsum_for_reduce_scatter=false --xla_jf_spmd_threshold_for_windowed_einsum_mib=1000000" python -m MaxText.benchmark_chunked_prefill MaxText/configs/inference.yml load_parameters_path=gs://jetstream-runner/8-22B-int8  max_prefill_predict_length=8192 max_target_length=9000 model_name=mixtral-8x22b ici_fsdp_parallelism=1 ici_autoregressive_parallelism=1 ici_tensor_parallelism=1 scan_layers=false per_device_batch_size=24 attention=dot_product megablox=False quantization=int8 checkpoint_is_quantized=True quantize_kvcache=False capacity_factor=1 tokenizer_path=assets/tokenizer.mistral-v3 inference_microbenchmark_prefill_lengths="8192" sparse_matmul=False model_call_mode=inference ici_context_autoregressive_parallelism=8 use_chunked_prefill=True prefill_chunk_size=2048 > ${OUTPUT_DIR}/moe_8x22b_long_context_8k_prefill.txt
+# tail -n5 ${OUTPUT_DIR}/moe_8x22b_long_context_8k_prefill.txt > ${OUTPUT_DIR}/moe_8x22b_long_context_8k_prefill.tmp && mv ${OUTPUT_DIR}/moe_8x22b_long_context_8k_prefill.tmp ${OUTPUT_DIR}/moe_8x22b_long_context_8k_prefill.txt
 
 
 # moe 8x7B Maxtext Jetstream 

experimental/jetstream-maxtext-stable-stack/test_script/benchmark_chunked_prefill_example.sh

@@ -1,3 +1,4 @@
+exit 0
 cd maxtext
 
 LIBTPU_INIT_ARGS="--xla_tpu_enable_windowed_einsum_for_reduce_scatter=false --xla_jf_spmd_threshold_for_windowed_einsum_mib=1000000" \