Skip to content

Commit cb15ef8

Browse files
yuyanpeng-googleyuyanpeng-google
committed
WIP
1 parent 19b2086 commit cb15ef8

3 files changed

Lines changed: 34 additions & 27 deletions

File tree

‎.github/workflows/run_maxtext_jetstream_tests.yaml‎

Lines changed: 31 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -50,13 +50,13 @@ jobs:
5050
- name: Build
5151
run: |
5252
pushd experimental/jetstream-maxtext-stable-stack
53-
./build.sh \
53+
# ./build.sh \
5454
LOCAL_IMAGE_TAG="${LOCAL_IMAGE_TAG}"
5555
popd
5656
- name: Test
5757
run: |
5858
pushd experimental/jetstream-maxtext-stable-stack
59-
./test.sh \
59+
# ./test.sh \
6060
LOCAL_IMAGE_TAG=${LOCAL_IMAGE_TAG}
6161
popd
6262
- name: Upload image
@@ -65,42 +65,48 @@ jobs:
6565
docker tag ${LOCAL_IMAGE_TAG} ${UPLOAD_IMAGE_TAG}
6666
docker push ${UPLOAD_IMAGE_TAG}
6767
68+
6869
benchmark_report:
6970
name: Benchmark Report
7071
needs: build_stable_stack
7172
runs-on: ["self-hosted", "tpu", "v6e-8"]
72-
container:
73-
# sync with the image uploaded from build_stable_stack stage
74-
image: gcr.io/cloud-ml-auto-solutions/jetstream-maxtext-stable-stack/tpu:github_${{ github.run_id }}
75-
options: "--privileged"
7673
env:
77-
OUTPUT_DIR: /workspace/test_dir/
74+
OUTPUT_DIR: ./test_dir/
7875
steps:
79-
- name: Create output directory # Ensure directory exists in container
80-
run: mkdir -p ${OUTPUT_DIR}
8176
- name: Test MOEBenchmarks
82-
# Report should generated in OUTPUT_DIR depend on ENV
83-
run: bash JetStream/.github/workflows/test_moe_benchmarks.sh
77+
run: |
78+
mkdir -p ${OUTPUT_DIR}
79+
# sync with the image uploaded from build_stable_stack stage
80+
# Report should generated in OUTPUT_DIR depend on ENV
81+
DOCKER_OUTPUT_DIR=/output
82+
docker run \
83+
-v ${OUTPUT_DIR}:${DOCKER_OUTPUT_DIR}} \
84+
--env OUTPUT_DIR:${DOCKER_OUTPUT_DIR} \
85+
--privileged --net=host --rm -i \
86+
gcr.io/cloud-ml-auto-solutions/jetstream-maxtext-stable-stack/tpu:github_${{ github.run_id }} \
87+
bash -c "
88+
bash JetStream/.github/workflows/test_moe_benchmarks.sh
89+
"
8490
- name: Upload build artifact
8591
uses: actions/upload-artifact@v4
8692
with:
8793
name: benchmark_report
8894
path: ${{ env.OUTPUT_DIR }}
8995

90-
clean_up_on_fail:
91-
if: ${{ failure() }}
92-
needs: [build_stable_stack, benchmark_report]
93-
name: "Clean up"
94-
runs-on: ["self-hosted"]
95-
permissions:
96-
contents: read
97-
issues: write # for failed-build-issue
98-
steps:
99-
- name: Authenticate gcloud
100-
run: gcloud auth configure-docker gcr.io --quiet
101-
- name: Delete TPU image
102-
# sync with the image uploaded from build_stable_stack stage
103-
run: gcloud container images delete gcr.io/cloud-ml-auto-solutions/jetstream-maxtext-stable-stack/tpu:github_${{ github.run_id }} --force-delete-tags --quiet
96+
# clean_up_on_fail:
97+
# if: ${{ failure() }}
98+
# needs: [build_stable_stack, benchmark_report]
99+
# name: "Clean up"
100+
# runs-on: ["self-hosted"]
101+
# permissions:
102+
# contents: read
103+
# issues: write # for failed-build-issue
104+
# steps:
105+
# - name: Authenticate gcloud
106+
# run: gcloud auth configure-docker gcr.io --quiet
107+
# - name: Delete TPU image
108+
# # sync with the image uploaded from build_stable_stack stage
109+
# run: gcloud container images delete gcr.io/cloud-ml-auto-solutions/jetstream-maxtext-stable-stack/tpu:github_${{ github.run_id }} --force-delete-tags --quiet
104110

105111
tag_night_image:
106112
needs: [build_stable_stack, benchmark_report]

‎.github/workflows/test_moe_benchmarks.sh‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,8 +16,8 @@ LIBTPU_INIT_ARGS="--xla_tpu_enable_windowed_einsum_for_reduce_scatter=false --xl
1616
tail -n5 ${OUTPUT_DIR}/moe_8x22b.txt > ${OUTPUT_DIR}/moe_8x22b.tmp && mv ${OUTPUT_DIR}/moe_8x22b.tmp ${OUTPUT_DIR}/moe_8x22b.txt
1717

1818
# moe 8x22B 8k context length chunked prefill with 2k prefill chunk size
19-
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
20-
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
19+
# 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
20+
# 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
2121

2222

2323
# moe 8x7B Maxtext Jetstream

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
exit 0
12
cd maxtext
23

34
LIBTPU_INIT_ARGS="--xla_tpu_enable_windowed_einsum_for_reduce_scatter=false --xla_jf_spmd_threshold_for_windowed_einsum_mib=1000000" \

0 commit comments

Comments
 (0)