Skip to content

fix(gpt-oss-20b): adjust max_kl_div threshold to 1.2 and support custom BASE_OUTPUT_DIRECTORY #11

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
hsuan-lun-chiang wants to merge 4 commits into
base: main
Choose a base branch
from
Draft

fix(gpt-oss-20b): adjust max_kl_div threshold to 1.2 and support custom BASE_OUTPUT_DIRECTORY #11

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
a89705d
fix(gpt-oss-20b): adjust max_kl_div threshold to 1.2 for checkpoint c…
hsuan-lun-chiang Oct 13, 2026
4b28692
fix(gpt-oss-20b): support custom BASE_OUTPUT_DIRECTORY and adjust max…
hsuan-lun-chiang Oct 13, 2026
2826bcd
fix(gpt-oss-20b): copy tests into runner dockerfile and support custo…
hsuan-lun-chiang Oct 13, 2026
49c9046
fix(gcs_utils): fix mkdir_and_check_permissions for newly created GCS…
hsuan-lun-chiang Oct 13, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions src/dependencies/dockerfiles/maxtext_runner.Dockerfile
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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/
5 changes: 3 additions & 2 deletions src/maxtext/utils/gcs_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand All @@ -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.
Expand Down
4 changes: 2 additions & 2 deletions tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss_to_mt.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Loading