From a89705d57f22e675b1fe703048b7d76678ffeb1e Mon Sep 17 00:00:00 2001 From: hsuan-lun-chiang Date: Tue, 13 Oct 2026 04:12:05 +0000 Subject: [PATCH 1/4] fix(gpt-oss-20b): adjust max_kl_div threshold to 1.2 for checkpoint conversion test --- tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss_to_mt.sh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss_to_mt.sh b/tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss_to_mt.sh index e9667eaa91..47637ad09f 100644 --- a/tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss_to_mt.sh +++ b/tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss_to_mt.sh @@ -74,6 +74,6 @@ python3 -m tests.utils.forward_pass_logit_checker \ dtype=float32 \ attention=dot_product \ --golden_logits_path=/tmp/golden_data_gpt-oss-20b.jsonl \ - --max_kl_div=1.0 \ + --max_kl_div=1.2 \ hardware=cpu \ skip_jax_distributed_system=True \ No newline at end of file From 4b286921751fa42d13373ee463d0dc98f93dcfa5 Mon Sep 17 00:00:00 2001 From: hsuan-lun-chiang Date: Tue, 13 Oct 2026 04:56:10 +0000 Subject: [PATCH 2/4] fix(gpt-oss-20b): support custom BASE_OUTPUT_DIRECTORY and adjust max_kl_div threshold to 1.2 --- tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss_to_mt.sh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss_to_mt.sh b/tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss_to_mt.sh index 47637ad09f..654af108d4 100644 --- a/tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss_to_mt.sh +++ b/tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss_to_mt.sh @@ -20,7 +20,7 @@ run_id=${1:-$(date +%Y-%m-%d-%H-%M-%S)} MODEL_NAME='gpt-oss-20b' # Non-Googlers please remember to point BASE_OUTPUT_DIRECTORY to the GCS paths where you want to store scanned and unscanned checkpoints -BASE_OUTPUT_DIRECTORY=gs://runner-maxtext-logs/${MODEL_NAME}/to_maxtext +BASE_OUTPUT_DIRECTORY=${BASE_OUTPUT_DIRECTORY:-gs://runner-maxtext-logs/${MODEL_NAME}/to_maxtext} # Step 1: Install torch python3 -m pip install torch --index-url https://download.pytorch.org/whl/cpu From 2826bcdbbbdc55df3d4df22fa21287a4e360e155 Mon Sep 17 00:00:00 2001 From: hsuan-lun-chiang Date: Tue, 13 Oct 2026 05:11:58 +0000 Subject: [PATCH 3/4] fix(gpt-oss-20b): copy tests into runner dockerfile and support custom BASE_OUTPUT_DIRECTORY --- src/dependencies/dockerfiles/maxtext_runner.Dockerfile | 1 + 1 file changed, 1 insertion(+) diff --git a/src/dependencies/dockerfiles/maxtext_runner.Dockerfile b/src/dependencies/dockerfiles/maxtext_runner.Dockerfile index d85511c848..b45eda3a72 100644 --- a/src/dependencies/dockerfiles/maxtext_runner.Dockerfile +++ b/src/dependencies/dockerfiles/maxtext_runner.Dockerfile @@ -19,3 +19,4 @@ COPY ${PACKAGE_DIR}/maxtext/assets/ "${MAXTEXT_ASSETS_ROOT}" # Copy all files except assets from local workspace into docker container COPY --exclude=${PACKAGE_DIR}/maxtext/assets/ ${PACKAGE_DIR}/maxtext/ src/maxtext/ +COPY tests/ tests/ From 49c90461fe9456ffd2c3c3d794744146a1775b7e Mon Sep 17 00:00:00 2001 From: hsuan-lun-chiang Date: Tue, 13 Oct 2026 05:27:27 +0000 Subject: [PATCH 4/4] fix(gcs_utils): fix mkdir_and_check_permissions for newly created GCS paths --- src/maxtext/utils/gcs_utils.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/maxtext/utils/gcs_utils.py b/src/maxtext/utils/gcs_utils.py index ade7fb79ad..2287937ac8 100644 --- a/src/maxtext/utils/gcs_utils.py +++ b/src/maxtext/utils/gcs_utils.py @@ -291,7 +291,8 @@ def mkdir_and_check_permissions(path: str | epath.Path) -> epath.Path: if isinstance(path, str): path = epath.Path(path) - if path.as_posix().startswith("gs://"): + is_gcs = path.as_posix().startswith("gs://") + if is_gcs: if len(path.parts) < 3: raise ValueError(f"Invalid GCS path (missing bucket name): '{path}'") bucket_name = path.parts[2] @@ -301,7 +302,7 @@ def mkdir_and_check_permissions(path: str | epath.Path) -> epath.Path: except Exception as e: raise FileNotFoundError(f"GCS bucket 'gs://{bucket_name}' not found or accessible.") from e path.mkdir(exist_ok=True, parents=True) - if not path.exists(): + if not is_gcs and not path.exists(): raise PermissionError(f"Failed to create the directory '{path}'. Please check that you have write access.") # Verify write permissions by creating and deleting a temporary file.