Fix llama4 test with hf access token

Author: shuningjin

tests/end_to_end/tpu/llama4/2_test_llama4.sh

@@ -6,7 +6,7 @@
 
 # The flow of this file is to take the MaxText (unscanned Orbax) checkpoint and run inference on a TPU VM.
 
-# Example Usage: export BASE_OUTPUT_PATH=/path/to/GCS/bucket; export MODEL_VARIATION=llama4-17b-[16e/128e]; bash tests/end_to_end/tpu/llama4/2_test_llama4.sh
+# Example Usage: export HF_TOKEN=<huggingface_access_token>; export BASE_OUTPUT_PATH=/path/to/GCS/bucket; export MODEL_VARIATION=llama4-17b-[16e/128e]; bash tests/end_to_end/tpu/llama4/2_test_llama4.sh
 # Use the same BASE_OUTPUT_PATH and MODEL_VARIATION for both 1_test_llama4.sh & 1_test_llama4.sh.
 
 # In order to generate the Llama4 golden logits, please see this script: tests/assets/logits_generation/golden_llama4_17b_16e_128e_export.ipynb
@@ -36,4 +36,4 @@ echo using BASE_OUTPUT_PATH = ${BASE_OUTPUT_PATH}
 export UNSCANNED_CKPT_PATH=${BASE_OUTPUT_PATH}/unscanned/0/items
 
 # Step 2: run logit checking
-python3 -m tests.utils.forward_pass_logit_checker "${MAXTEXT_CONFIGS_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/maxtext/configs}"//base.yml tokenizer_path=${TOKENIZER_PATH} load_parameters_path=${UNSCANNED_CKPT_PATH} run_name=forward_pass_test_${MODEL_VARIATION} attention=dot_product per_device_batch_size=1 model_name=${MODEL_VARIATION} max_prefill_predict_length=4 max_target_length=4 scan_layers=false --atol=0.01 --rtol=0.01 async_checkpointing=false sparse_matmul=false weight_dtype=float32 dtype=float32 activations_in_float32=true matmul_precision=float32 float32_logits=true float32_qk_product=true ici_expert_parallelism=16
+python3 -m tests.utils.forward_pass_logit_checker "${MAXTEXT_CONFIGS_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/maxtext/configs}"//base.yml tokenizer_path=${TOKENIZER_PATH} load_parameters_path=${UNSCANNED_CKPT_PATH} run_name=forward_pass_test_${MODEL_VARIATION} attention=dot_product per_device_batch_size=1 model_name=${MODEL_VARIATION} max_prefill_predict_length=4 max_target_length=4 scan_layers=false --atol=0.01 --rtol=0.01 async_checkpointing=false sparse_matmul=false weight_dtype=float32 dtype=float32 activations_in_float32=true matmul_precision=float32 float32_logits=true float32_qk_product=true ici_expert_parallelism=16 hf_access_token=${HF_TOKEN}

tests/utils/forward_pass_logit_checker.py

@@ -33,6 +33,8 @@
 # Utilities and Colab notebooks used to generate the golden logits are available in tests/assets/logits_generation.
 # For example:
 #   tests/assets/logits_generation/golden_llama2-7b_export.ipynb
+#
+# NOTE: Be careful of updates. This file affects scripts under `tests/end_to_end/tpu` for nightly tests.
 
 """Check if the logits generated by a model's src/maxtext/HF implementation matches golden logits for the same inputs.