From e82524016146ec8672093320189315cf463adb3a Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Mon, 24 Aug 2026 12:45:08 -0700 Subject: [PATCH 1/5] Add reusable TPU v5e setup script --- .../nanogpt_one_head/scripts/setup_tpu_v5e.sh | 314 ++++++++++++++++++ 1 file changed, 314 insertions(+) create mode 100644 baseline/nanogpt_one_head/scripts/setup_tpu_v5e.sh diff --git a/baseline/nanogpt_one_head/scripts/setup_tpu_v5e.sh b/baseline/nanogpt_one_head/scripts/setup_tpu_v5e.sh new file mode 100644 index 0000000..295f937 --- /dev/null +++ b/baseline/nanogpt_one_head/scripts/setup_tpu_v5e.sh @@ -0,0 +1,314 @@ +#!/usr/bin/env bash +# Reproducible bootstrap for the one-head nanoGPT baseline on a Cloud TPU v5e VM. +# +# Run this *inside* an active TPU VM after cloning rg_optimizers. + +set -Eeuo pipefail +IFS=$'\n\t' + +SCRIPT_NAME="$(basename "$0")" +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +PROJECT_DIR="$(cd "${SCRIPT_DIR}/.." && pwd)" +PYTHON_BIN="${PYTHON_BIN:-python3}" +TORCH_VERSION="${TORCH_VERSION:-2.6.0}" +ACCELERATOR_TYPE="${TPU_ACCELERATOR_TYPE:-v5litepod-4}" +STORAGE_MODE="auto" +PERSISTENT_ROOT="" +RUN_QUICK_SMOKE=0 +SMOKE_OPTIMIZER="adamw" +FORCE_FRAMEWORK_REINSTALL=0 + +log() { + printf '[tpu-setup] %s\n' "$*" +} + +warn() { + printf '[tpu-setup] WARNING: %s\n' "$*" >&2 +} + +fail() { + printf '[tpu-setup] ERROR: %s\n' "$*" >&2 + exit 1 +} + +on_error() { + local exit_code=$? + printf '[tpu-setup] ERROR: command failed at line %s (exit=%s)\n' "${BASH_LINENO[0]}" "$exit_code" >&2 + exit "$exit_code" +} +trap on_error ERR + +usage() { + cat <= 2)) || fail "--persistent-root requires a path" + STORAGE_MODE="persistent" + PERSISTENT_ROOT="$2" + shift 2 + ;; + --accelerator-type) + (($# >= 2)) || fail "--accelerator-type requires a value" + ACCELERATOR_TYPE="$2" + shift 2 + ;; + --torch-version) + (($# >= 2)) || fail "--torch-version requires a value" + TORCH_VERSION="$2" + shift 2 + ;; + --force-framework-reinstall) + FORCE_FRAMEWORK_REINSTALL=1 + shift + ;; + --run-quick-smoke) + RUN_QUICK_SMOKE=1 + shift + ;; + --optimizer) + (($# >= 2)) || fail "--optimizer requires a value" + SMOKE_OPTIMIZER="$2" + shift 2 + ;; + -h|--help) + usage + exit 0 + ;; + *) + fail "unknown option: $1" + ;; + esac +done + +[[ -f "${PROJECT_DIR}/pyproject.toml" ]] || \ + fail "could not find ${PROJECT_DIR}/pyproject.toml; run this script from the cloned repository" +command -v "$PYTHON_BIN" >/dev/null 2>&1 || fail "Python executable not found: ${PYTHON_BIN}" + +if [[ "$STORAGE_MODE" == "auto" ]]; then + if [[ -d /mnt/disks/rg-data && -w /mnt/disks/rg-data ]]; then + STORAGE_MODE="persistent" + PERSISTENT_ROOT="/mnt/disks/rg-data" + else + STORAGE_MODE="ephemeral" + fi +fi + +if [[ "$STORAGE_MODE" == "persistent" ]]; then + [[ -n "$PERSISTENT_ROOT" ]] || fail "persistent storage selected without a root" + [[ -d "$PERSISTENT_ROOT" ]] || fail "persistent root does not exist: ${PERSISTENT_ROOT}" + [[ -w "$PERSISTENT_ROOT" ]] || fail "persistent root is not writable: ${PERSISTENT_ROOT}" +else + warn "using ephemeral /tmp storage; checkpoints and data disappear when the TPU VM is deleted" + warn "do not use --ephemeral for a long or scientific run" +fi + +log "project directory: ${PROJECT_DIR}" +log "python: $($PYTHON_BIN --version 2>&1)" +log "accelerator type: ${ACCELERATOR_TYPE}" +log "storage mode: ${STORAGE_MODE}" + +log "upgrading pip, setuptools, and wheel (required to avoid UNKNOWN-0.0.0 installs)" +"$PYTHON_BIN" -m pip install --user --upgrade pip setuptools wheel + +# Remove the bad artifact produced by the old TPU VM packaging toolchain, if present. +if "$PYTHON_BIN" -m pip show UNKNOWN >/dev/null 2>&1; then + log "removing stale UNKNOWN-0.0.0 package" + "$PYTHON_BIN" -m pip uninstall -y UNKNOWN +fi + +framework_matches() { + PJRT_DEVICE=TPU "$PYTHON_BIN" - "$TORCH_VERSION" <<'PY' +import re +import sys + +expected = sys.argv[1] + +def major_minor(value: str) -> tuple[int, int]: + match = re.match(r"^\s*(\d+)\.(\d+)", value) + if not match: + raise RuntimeError(f"cannot parse version: {value!r}") + return int(match.group(1)), int(match.group(2)) + +import torch +import torch_xla + +if major_minor(torch.__version__) != major_minor(expected): + raise SystemExit(1) +if major_minor(torch_xla.__version__) != major_minor(expected): + raise SystemExit(1) +PY +} + +if ((FORCE_FRAMEWORK_REINSTALL)) || ! framework_matches; then + log "installing matching torch=${TORCH_VERSION} and torch_xla=${TORCH_VERSION}" + "$PYTHON_BIN" -m pip uninstall -y torch torch_xla torchvision || true + "$PYTHON_BIN" -m pip install --user \ + "torch==${TORCH_VERSION}" \ + "torch_xla[tpu]==${TORCH_VERSION}" \ + -f https://storage.googleapis.com/libtpu-releases/index.html \ + -f https://storage.googleapis.com/libtpu-wheels/index.html +else + log "matching torch/torch_xla ${TORCH_VERSION%.*}.x stack is already installed" +fi + +log "installing the nanoGPT package without editable mode" +"$PYTHON_BIN" -m pip install --user --upgrade "$PROJECT_DIR" + +ENV_DIR="${HOME}/.config/rg_optimizers" +ENV_FILE="${ENV_DIR}/tpu_env.sh" +mkdir -p "$ENV_DIR" +{ + printf '# Generated by %s. Safe to source from ~/.bashrc.\n' "$SCRIPT_NAME" + printf 'export PATH="$HOME/.local/bin:$PATH"\n' + printf 'export PJRT_DEVICE=TPU\n' + printf 'export TPU_ACCELERATOR_TYPE=%q\n' "$ACCELERATOR_TYPE" + printf 'unset XLA_USE_BF16\n' + printf 'unset XLA_DOWNCAST_BF16\n' + if [[ "$STORAGE_MODE" == "persistent" ]]; then + printf 'unset RG_NANOGPT_ALLOW_EPHEMERAL_TPU_STORAGE\n' + printf 'export RG_TPU_PERSISTENT_ROOT=%q\n' "$PERSISTENT_ROOT" + else + printf 'unset RG_TPU_PERSISTENT_ROOT\n' + printf 'export RG_NANOGPT_ALLOW_EPHEMERAL_TPU_STORAGE=1\n' + fi +} > "$ENV_FILE" +chmod 600 "$ENV_FILE" + +BASHRC_LINE='[ -f "$HOME/.config/rg_optimizers/tpu_env.sh" ] && source "$HOME/.config/rg_optimizers/tpu_env.sh"' +if ! grep -Fqx "$BASHRC_LINE" "${HOME}/.bashrc" 2>/dev/null; then + printf '\n%s\n' "$BASHRC_LINE" >> "${HOME}/.bashrc" +fi + +# shellcheck disable=SC1090 +source "$ENV_FILE" + +log "verifying the PyTorch/XLA runtime" +"$PYTHON_BIN" - "$TORCH_VERSION" <<'PY' +import json +import re +import sys + +expected = sys.argv[1] + +def major_minor(value: str) -> tuple[int, int]: + match = re.match(r"^\s*(\d+)\.(\d+)", value) + if not match: + raise RuntimeError(f"cannot parse version: {value!r}") + return int(match.group(1)), int(match.group(2)) + +import torch +import torch_xla +import torch_xla.core.xla_model as xm +import torch_xla.runtime as xr + +if major_minor(torch.__version__) != major_minor(expected): + raise RuntimeError(f"torch version mismatch: {torch.__version__} vs {expected}") +if major_minor(torch_xla.__version__) != major_minor(expected): + raise RuntimeError(f"torch_xla version mismatch: {torch_xla.__version__} vs {expected}") +if str(xr.device_type()).upper() != "TPU": + raise RuntimeError(f"PJRT device is not TPU: {xr.device_type()!r}") + +devices = xm.get_xla_supported_devices("TPU") +if not devices: + raise RuntimeError("PyTorch/XLA found no TPU devices") + +print(json.dumps({ + "torch": torch.__version__, + "torch_xla": torch_xla.__version__, + "pjrt_device": xr.device_type(), + "tpu_devices": devices, +}, indent=2)) +PY + +RG_ENV="$(command -v rg-onehead-env || true)" +[[ -n "$RG_ENV" ]] || RG_ENV="${HOME}/.local/bin/rg-onehead-env" +[[ -x "$RG_ENV" ]] || fail "rg-onehead-env was not installed" + +log "verifying nanoGPT runtime and storage resolution" +"$RG_ENV" --device auto + +if ((RUN_QUICK_SMOKE)); then + case "$SMOKE_OPTIMIZER" in + adamw|muon|sgd_momentum) ;; + *) fail "unsupported quick-smoke optimizer: ${SMOKE_OPTIMIZER}" ;; + esac + + QUICK_CONFIG="${TMPDIR:-/tmp}/rg_nanogpt_tpu_quick.yaml" + log "writing non-scientific quick-smoke config: ${QUICK_CONFIG}" + "$PYTHON_BIN" - "${PROJECT_DIR}/configs/tpu_smoke.yaml" "$QUICK_CONFIG" <<'PY' +import sys +from pathlib import Path + +import yaml + +source = Path(sys.argv[1]) +target = Path(sys.argv[2]) +with source.open(encoding="utf-8") as handle: + cfg = yaml.safe_load(handle) + +cfg["protocol"]["name"] = "rg_nanogpt_one_head_tpu_quick_smoke" +cfg["protocol"]["description"] = ( + "Non-scientific TPU/XLA compatibility and throughput smoke test." +) +cfg["dataset"]["train_tokens"] = 4_000_000 +cfg["dataset"]["val_tokens"] = 100_000 +cfg["dataset"]["test_tokens"] = 100_000 +cfg["training"]["target_epochs"] = 0.10 +cfg["training"]["epoch_interval"] = 0.10 +cfg["training"]["eval_interval_steps"] = 25 +cfg["training"]["checkpoint_interval_steps"] = 25 +cfg["weightwatcher"]["enabled"] = False + +with target.open("w", encoding="utf-8") as handle: + yaml.safe_dump(cfg, handle, sort_keys=False) +PY + + RG_PREPARE="$(command -v rg-onehead-prepare || true)" + RG_TRAIN="$(command -v rg-onehead-train || true)" + [[ -n "$RG_PREPARE" ]] || RG_PREPARE="${HOME}/.local/bin/rg-onehead-prepare" + [[ -n "$RG_TRAIN" ]] || RG_TRAIN="${HOME}/.local/bin/rg-onehead-train" + + log "preparing the small pinned FineWeb-Edu cache" + "$RG_PREPARE" --config "$QUICK_CONFIG" --force + + log "running the quick TPU smoke test with optimizer=${SMOKE_OPTIMIZER}" + "$RG_TRAIN" \ + --config "$QUICK_CONFIG" \ + --optimizer "$SMOKE_OPTIMIZER" \ + --device auto \ + --no-resume +fi + +log "setup complete" +log "future SSH shells load ${ENV_FILE} automatically" +log "for this already-open shell, run: source ${ENV_FILE}" +if ((RUN_QUICK_SMOKE == 0)); then + log "optional quick test: ${SCRIPT_NAME} --ephemeral --run-quick-smoke" +fi From 734b04959722e10c1a0d4abfdd9465a0aa1f1961 Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Mon, 24 Aug 2026 12:46:07 -0700 Subject: [PATCH 2/5] Document reproducible TPU v5e bootstrap workflow --- baseline/nanogpt_one_head/TPU.md | 157 ++++++++++++++++++++++++++++--- 1 file changed, 143 insertions(+), 14 deletions(-) diff --git a/baseline/nanogpt_one_head/TPU.md b/baseline/nanogpt_one_head/TPU.md index 8e86a0e..d34b5f8 100644 --- a/baseline/nanogpt_one_head/TPU.md +++ b/baseline/nanogpt_one_head/TPU.md @@ -16,20 +16,146 @@ size, data sampling, checkpoint ownership, and WeightWatcher ownership; the current runner refuses a multi-process XLA launch rather than silently changing the experiment. -## Install on a TPU VM +## TPU Builders v5e Flex-Start quick path -Use a Python/PyTorch combination supported by the installed PyTorch/XLA release, -then install the TPU extra from this directory: +The following is a known-good path for a one-hour disposable v5e session. Run +the provisioning commands from Google Cloud Shell, not from inside a TPU VM. + +### 1. Request the TPU + +```bash +gcloud alpha compute tpus queued-resources create tpu-v5e-request \ + --zone=us-west4-a \ + --accelerator-type=v5litepod-4 \ + --runtime-version=v2-alpha-tpuv5-lite \ + --node-id=tpu-v5e-node \ + --provisioning-model=flex-start \ + --max-run-duration=1h \ + --valid-until-duration=30m \ + --labels=purpose=flex-start +``` + +Check the request until its state is `ACTIVE`: + +```bash +gcloud alpha compute tpus queued-resources describe tpu-v5e-request \ + --zone=us-west4-a \ + --format='value(state.state)' +``` + +`--valid-until-duration=30m` is the capacity-acquisition window. +`--max-run-duration=1h` starts after provisioning and automatically terminates +the TPU after at most one hour. + +SSH into the active TPU VM: + +```bash +gcloud compute tpus tpu-vm ssh tpu-v5e-node \ + --project=YOUR_PROJECT_ID \ + --zone=us-west4-a +``` + +To clean up early from Cloud Shell: + +```bash +gcloud alpha compute tpus queued-resources delete tpu-v5e-request \ + --zone=us-west4-a \ + --force \ + --quiet +``` + +### 2. Clone and bootstrap the TPU VM + +Inside the TPU VM: + +```bash +cd /tmp +git clone https://github.com/CalculatedContent/rg_optimizers.git +cd rg_optimizers/baseline/nanogpt_one_head +bash scripts/setup_tpu_v5e.sh --ephemeral +``` + +The script records the required environment in: + +```text +~/.config/rg_optimizers/tpu_env.sh +``` + +It also adds an idempotent source line to `~/.bashrc`, so future SSH shells load +the TPU environment automatically. For the shell that launched the setup +script, load it explicitly after the script returns: + +```bash +source ~/.config/rg_optimizers/tpu_env.sh +``` + +For a disposable setup plus a small, non-scientific AdamW throughput test: + +```bash +bash scripts/setup_tpu_v5e.sh --ephemeral --run-quick-smoke +``` + +That optional quick smoke creates a reduced 4M/100k/100k-token corpus, disables +WeightWatcher, and runs approximately 49 optimizer steps. It is only a TPU/XLA +compatibility and throughput check; it must not appear in scientific result +tables. + +For a real run, mount durable storage and use: + +```bash +bash scripts/setup_tpu_v5e.sh \ + --persistent-root /mnt/disks/rg-data +``` + +### What the setup script fixes and verifies + +The stock TPU VM may contain an old packaging toolchain and an incompatible +PyTorch/PyTorch-XLA pair. The reusable script therefore: + +1. upgrades user-level `pip`, `setuptools`, and `wheel` before installing this + package, preventing the erroneous `UNKNOWN-0.0.0` build observed with the + stock toolchain; +2. removes a stale `UNKNOWN` package if one exists; +3. installs matching `torch==2.6.0` and `torch_xla[tpu]==2.6.0` binaries from + the TPU wheel indexes when the installed major/minor versions do not match; +4. installs this package without editable mode; +5. exports `PJRT_DEVICE=TPU` and the TPU provenance label + `TPU_ACCELERATOR_TYPE=v5litepod-4`; +6. unsets `XLA_USE_BF16` and `XLA_DOWNCAST_BF16`, preserving the float32 + reference protocol; +7. configures either persistent or explicitly ephemeral storage; and +8. verifies both the raw XLA devices and `rg-onehead-env --device auto`. + +The expected direct XLA check on a `v5litepod-4` is: + +```text +torch: 2.6.0+cu124 +torch_xla: 2.6.0 +TPU devices: ['xla:0', 'xla:1', 'xla:2', 'xla:3'] +``` + +The baseline still uses only `xla:0` in its current single-process protocol. + +## Manual installation fallback + +The setup script is preferred. The equivalent core installation sequence is: ```bash cd baseline/nanogpt_one_head -python -m pip install -e '.[tpu]' + +python3 -m pip install --user --upgrade pip setuptools wheel +python3 -m pip uninstall -y torch torch_xla torchvision +python3 -m pip install --user \ + 'torch==2.6.0' \ + 'torch_xla[tpu]==2.6.0' \ + -f https://storage.googleapis.com/libtpu-releases/index.html \ + -f https://storage.googleapis.com/libtpu-wheels/index.html +python3 -m pip install --user . ``` -PyTorch and PyTorch/XLA must have matching major/minor versions. The runner -checks this before training and reports a direct error if they do not match. -The reference protocol is float32; `XLA_USE_BF16` and `XLA_DOWNCAST_BF16` must -not be enabled. +PyTorch and PyTorch/XLA must have matching major/minor versions. A mismatch can +surface as an `_XLAC` import failure with an undefined PyTorch symbol. The +runner also checks the versions before training and reports a direct error. ## Persistent TPU storage @@ -99,20 +225,23 @@ rg-onehead-env --device auto ``` It prints JSON containing the selected accelerator, XLA runtime information, -and resolved data/results/plot roots. On a TPU VM, the output must show: +and resolved data/results/plot roots. On a v5e TPU VM, the output should include: ```text accelerator: tpu +device: xla:0 pjrt_device: TPU +tpu_accelerator_type: v5litepod-4 +xla_addressable_device_count: 4 xla_process_count: 1 ``` -and the roots must point to the attached durable volume. +For a scientific run, the roots must point to the attached durable volume. -## Smoke test +## Full integration smoke test -After the pinned corpus is available on the persistent volume, run the short -non-scientific integration test: +After the pinned corpus is available on the persistent volume, run the committed +integration test: ```bash rg-onehead-train \ @@ -122,7 +251,7 @@ rg-onehead-train \ --no-resume ``` -The smoke test exercises model transfer, Muon, auxiliary AdamW, XLA step +This smoke test exercises model transfer, Muon, auxiliary AdamW, XLA step boundaries, evaluation, CPU BLEU, CPU WeightWatcher, portable checkpoints, and persistent path resolution. It is not an optimizer result and should never be included in the scientific tables. From 93e213c8840d55e37a79ee14a789efdcce3d983d Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Mon, 24 Aug 2026 12:51:08 -0700 Subject: [PATCH 3/5] Relocate TPU setup utility to experiment root --- baseline/nanogpt_one_head/setup_tpu_v5e.sh | 317 +++++++++++++++++++++ 1 file changed, 317 insertions(+) create mode 100644 baseline/nanogpt_one_head/setup_tpu_v5e.sh diff --git a/baseline/nanogpt_one_head/setup_tpu_v5e.sh b/baseline/nanogpt_one_head/setup_tpu_v5e.sh new file mode 100644 index 0000000..5f43c8a --- /dev/null +++ b/baseline/nanogpt_one_head/setup_tpu_v5e.sh @@ -0,0 +1,317 @@ +#!/usr/bin/env bash +# Reproducible bootstrap for the one-head nanoGPT baseline on a Cloud TPU v5e VM. +# +# Run this *inside* an active TPU VM after cloning rg_optimizers. + +set -Eeuo pipefail +IFS=$'\n\t' + +SCRIPT_NAME="$(basename "$0")" +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +PROJECT_DIR="${SCRIPT_DIR}" +PYTHON_BIN="${PYTHON_BIN:-python3}" +USER_ROOT="$(getent passwd "$(id -u)" | cut -d: -f6)" +TORCH_VERSION="${TORCH_VERSION:-2.6.0}" +ACCELERATOR_TYPE="${TPU_ACCELERATOR_TYPE:-v5litepod-4}" +STORAGE_MODE="auto" +PERSISTENT_ROOT="" +RUN_QUICK_SMOKE=0 +SMOKE_OPTIMIZER="adamw" +FORCE_FRAMEWORK_REINSTALL=0 + +log() { + printf '[tpu-setup] %s\n' "$*" +} + +warn() { + printf '[tpu-setup] WARNING: %s\n' "$*" >&2 +} + +fail() { + printf '[tpu-setup] ERROR: %s\n' "$*" >&2 + exit 1 +} + +on_error() { + local exit_code=$? + printf '[tpu-setup] ERROR: command failed at line %s (exit=%s)\n' "${BASH_LINENO[0]}" "$exit_code" >&2 + exit "$exit_code" +} +trap on_error ERR + +usage() { + cat <= 2)) || fail "--persistent-root requires a path" + STORAGE_MODE="persistent" + PERSISTENT_ROOT="$2" + shift 2 + ;; + --accelerator-type) + (($# >= 2)) || fail "--accelerator-type requires a value" + ACCELERATOR_TYPE="$2" + shift 2 + ;; + --torch-version) + (($# >= 2)) || fail "--torch-version requires a value" + TORCH_VERSION="$2" + shift 2 + ;; + --force-framework-reinstall) + FORCE_FRAMEWORK_REINSTALL=1 + shift + ;; + --run-quick-smoke) + RUN_QUICK_SMOKE=1 + shift + ;; + --optimizer) + (($# >= 2)) || fail "--optimizer requires a value" + SMOKE_OPTIMIZER="$2" + shift 2 + ;; + -h|--help) + usage + exit 0 + ;; + *) + fail "unknown option: $1" + ;; + esac +done + +[[ -f "${PROJECT_DIR}/pyproject.toml" ]] || \ + fail "could not find ${PROJECT_DIR}/pyproject.toml; run this script from the cloned repository" +command -v "$PYTHON_BIN" >/dev/null 2>&1 || fail "Python executable not found: ${PYTHON_BIN}" +[[ -n "$USER_ROOT" && -d "$USER_ROOT" ]] || fail "could not resolve the current user directory" + +if [[ "$STORAGE_MODE" == "auto" ]]; then + if [[ -d /mnt/disks/rg-data && -w /mnt/disks/rg-data ]]; then + STORAGE_MODE="persistent" + PERSISTENT_ROOT="/mnt/disks/rg-data" + else + STORAGE_MODE="ephemeral" + fi +fi + +if [[ "$STORAGE_MODE" == "persistent" ]]; then + [[ -n "$PERSISTENT_ROOT" ]] || fail "persistent storage selected without a root" + [[ -d "$PERSISTENT_ROOT" ]] || fail "persistent root does not exist: ${PERSISTENT_ROOT}" + [[ -w "$PERSISTENT_ROOT" ]] || fail "persistent root is not writable: ${PERSISTENT_ROOT}" +else + warn "using ephemeral /tmp storage; checkpoints and data disappear when the TPU VM is deleted" + warn "do not use --ephemeral for a long or scientific run" +fi + +log "project directory: ${PROJECT_DIR}" +log "python: $($PYTHON_BIN --version 2>&1)" +log "accelerator type: ${ACCELERATOR_TYPE}" +log "storage mode: ${STORAGE_MODE}" + +log "upgrading pip, setuptools, and wheel (required to avoid UNKNOWN-0.0.0 installs)" +"$PYTHON_BIN" -m pip install --user --upgrade pip setuptools wheel + +# Remove the bad artifact produced by the old TPU VM packaging toolchain, if present. +if "$PYTHON_BIN" -m pip show UNKNOWN >/dev/null 2>&1; then + log "removing stale UNKNOWN-0.0.0 package" + "$PYTHON_BIN" -m pip uninstall -y UNKNOWN +fi + +framework_matches() { + PJRT_DEVICE=TPU "$PYTHON_BIN" - "$TORCH_VERSION" <<'PY' +import re +import sys + +expected = sys.argv[1] + +def major_minor(value: str) -> tuple[int, int]: + match = re.match(r"^\s*(\d+)\.(\d+)", value) + if not match: + raise RuntimeError(f"cannot parse version: {value!r}") + return int(match.group(1)), int(match.group(2)) + +import torch +import torch_xla + +if major_minor(torch.__version__) != major_minor(expected): + raise SystemExit(1) +if major_minor(torch_xla.__version__) != major_minor(expected): + raise SystemExit(1) +PY +} + +if ((FORCE_FRAMEWORK_REINSTALL)) || ! framework_matches >/dev/null 2>&1; then + log "installing matching torch=${TORCH_VERSION} and torch_xla=${TORCH_VERSION}" + "$PYTHON_BIN" -m pip uninstall -y torch torch_xla torchvision || true + "$PYTHON_BIN" -m pip install --user \ + "torch==${TORCH_VERSION}" \ + "torch_xla[tpu]==${TORCH_VERSION}" \ + -f https://storage.googleapis.com/libtpu-releases/index.html \ + -f https://storage.googleapis.com/libtpu-wheels/index.html +else + log "matching torch/torch_xla ${TORCH_VERSION%.*}.x stack is already installed" +fi + +log "installing the nanoGPT package without editable mode" +"$PYTHON_BIN" -m pip install --user --upgrade "$PROJECT_DIR" + +ENV_DIR="${USER_ROOT}/.config/rg_optimizers" +ENV_FILE="${ENV_DIR}/tpu_env.sh" +mkdir -p "$ENV_DIR" +{ + printf '# Generated by %s. Safe to source from the user shell profile.\n' "$SCRIPT_NAME" + printf 'export PATH=%q:$PATH\n' "${USER_ROOT}/.local/bin" + printf 'export PJRT_DEVICE=TPU\n' + printf 'export TPU_ACCELERATOR_TYPE=%q\n' "$ACCELERATOR_TYPE" + printf 'unset XLA_USE_BF16\n' + printf 'unset XLA_DOWNCAST_BF16\n' + if [[ "$STORAGE_MODE" == "persistent" ]]; then + printf 'unset RG_NANOGPT_ALLOW_EPHEMERAL_TPU_STORAGE\n' + printf 'export RG_TPU_PERSISTENT_ROOT=%q\n' "$PERSISTENT_ROOT" + else + printf 'unset RG_TPU_PERSISTENT_ROOT\n' + printf 'export RG_NANOGPT_ALLOW_EPHEMERAL_TPU_STORAGE=1\n' + fi +} > "$ENV_FILE" +chmod 600 "$ENV_FILE" + +BASHRC_FILE="${USER_ROOT}/.bashrc" +printf -v BASHRC_LINE '[ -f %q ] && source %q' "$ENV_FILE" "$ENV_FILE" +if ! grep -Fqx "$BASHRC_LINE" "$BASHRC_FILE" 2>/dev/null; then + printf '\n%s\n' "$BASHRC_LINE" >> "$BASHRC_FILE" +fi + +# shellcheck disable=SC1090 +source "$ENV_FILE" + +log "verifying the PyTorch/XLA runtime" +"$PYTHON_BIN" - "$TORCH_VERSION" <<'PY' +import json +import re +import sys + +expected = sys.argv[1] + +def major_minor(value: str) -> tuple[int, int]: + match = re.match(r"^\s*(\d+)\.(\d+)", value) + if not match: + raise RuntimeError(f"cannot parse version: {value!r}") + return int(match.group(1)), int(match.group(2)) + +import torch +import torch_xla +import torch_xla.core.xla_model as xm +import torch_xla.runtime as xr + +if major_minor(torch.__version__) != major_minor(expected): + raise RuntimeError(f"torch version mismatch: {torch.__version__} vs {expected}") +if major_minor(torch_xla.__version__) != major_minor(expected): + raise RuntimeError(f"torch_xla version mismatch: {torch_xla.__version__} vs {expected}") +if str(xr.device_type()).upper() != "TPU": + raise RuntimeError(f"PJRT device is not TPU: {xr.device_type()!r}") + +devices = xm.get_xla_supported_devices("TPU") +if not devices: + raise RuntimeError("PyTorch/XLA found no TPU devices") + +print(json.dumps({ + "torch": torch.__version__, + "torch_xla": torch_xla.__version__, + "pjrt_device": xr.device_type(), + "tpu_devices": devices, +}, indent=2)) +PY + +RG_ENV="$(command -v rg-onehead-env || true)" +[[ -n "$RG_ENV" ]] || RG_ENV="${USER_ROOT}/.local/bin/rg-onehead-env" +[[ -x "$RG_ENV" ]] || fail "rg-onehead-env was not installed" + +log "verifying nanoGPT runtime and storage resolution" +"$RG_ENV" --device auto + +if ((RUN_QUICK_SMOKE)); then + case "$SMOKE_OPTIMIZER" in + adamw|muon|sgd_momentum) ;; + *) fail "unsupported quick-smoke optimizer: ${SMOKE_OPTIMIZER}" ;; + esac + + QUICK_CONFIG="${TMPDIR:-/tmp}/rg_nanogpt_tpu_quick.yaml" + log "writing non-scientific quick-smoke config: ${QUICK_CONFIG}" + "$PYTHON_BIN" - "${PROJECT_DIR}/configs/tpu_smoke.yaml" "$QUICK_CONFIG" <<'PY' +import sys +from pathlib import Path + +import yaml + +source = Path(sys.argv[1]) +target = Path(sys.argv[2]) +with source.open(encoding="utf-8") as handle: + cfg = yaml.safe_load(handle) + +cfg["protocol"]["name"] = "rg_nanogpt_one_head_tpu_quick_smoke" +cfg["protocol"]["description"] = ( + "Non-scientific TPU/XLA compatibility and throughput smoke test." +) +cfg["dataset"]["train_tokens"] = 4_000_000 +cfg["dataset"]["val_tokens"] = 100_000 +cfg["dataset"]["test_tokens"] = 100_000 +cfg["training"]["target_epochs"] = 0.10 +cfg["training"]["epoch_interval"] = 0.10 +cfg["training"]["eval_interval_steps"] = 25 +cfg["training"]["checkpoint_interval_steps"] = 25 +cfg["weightwatcher"]["enabled"] = False + +with target.open("w", encoding="utf-8") as handle: + yaml.safe_dump(cfg, handle, sort_keys=False) +PY + + RG_PREPARE="$(command -v rg-onehead-prepare || true)" + RG_TRAIN="$(command -v rg-onehead-train || true)" + [[ -n "$RG_PREPARE" ]] || RG_PREPARE="${USER_ROOT}/.local/bin/rg-onehead-prepare" + [[ -n "$RG_TRAIN" ]] || RG_TRAIN="${USER_ROOT}/.local/bin/rg-onehead-train" + + log "preparing the small pinned FineWeb-Edu cache" + "$RG_PREPARE" --config "$QUICK_CONFIG" --force + + log "running the quick TPU smoke test with optimizer=${SMOKE_OPTIMIZER}" + "$RG_TRAIN" \ + --config "$QUICK_CONFIG" \ + --optimizer "$SMOKE_OPTIMIZER" \ + --device auto \ + --no-resume +fi + +log "setup complete" +log "future SSH shells load ${ENV_FILE} automatically" +log "for this already-open shell, run: source ${ENV_FILE}" +if ((RUN_QUICK_SMOKE == 0)); then + log "optional quick test: ${SCRIPT_NAME} --ephemeral --run-quick-smoke" +fi From 1f3fab0aacd4657c578b8d1978c4a59548b0ffdd Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Mon, 24 Aug 2026 12:51:13 -0700 Subject: [PATCH 4/5] Remove forbidden nanoGPT scripts directory --- .../nanogpt_one_head/scripts/setup_tpu_v5e.sh | 314 ------------------ 1 file changed, 314 deletions(-) delete mode 100644 baseline/nanogpt_one_head/scripts/setup_tpu_v5e.sh diff --git a/baseline/nanogpt_one_head/scripts/setup_tpu_v5e.sh b/baseline/nanogpt_one_head/scripts/setup_tpu_v5e.sh deleted file mode 100644 index 295f937..0000000 --- a/baseline/nanogpt_one_head/scripts/setup_tpu_v5e.sh +++ /dev/null @@ -1,314 +0,0 @@ -#!/usr/bin/env bash -# Reproducible bootstrap for the one-head nanoGPT baseline on a Cloud TPU v5e VM. -# -# Run this *inside* an active TPU VM after cloning rg_optimizers. - -set -Eeuo pipefail -IFS=$'\n\t' - -SCRIPT_NAME="$(basename "$0")" -SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -PROJECT_DIR="$(cd "${SCRIPT_DIR}/.." && pwd)" -PYTHON_BIN="${PYTHON_BIN:-python3}" -TORCH_VERSION="${TORCH_VERSION:-2.6.0}" -ACCELERATOR_TYPE="${TPU_ACCELERATOR_TYPE:-v5litepod-4}" -STORAGE_MODE="auto" -PERSISTENT_ROOT="" -RUN_QUICK_SMOKE=0 -SMOKE_OPTIMIZER="adamw" -FORCE_FRAMEWORK_REINSTALL=0 - -log() { - printf '[tpu-setup] %s\n' "$*" -} - -warn() { - printf '[tpu-setup] WARNING: %s\n' "$*" >&2 -} - -fail() { - printf '[tpu-setup] ERROR: %s\n' "$*" >&2 - exit 1 -} - -on_error() { - local exit_code=$? - printf '[tpu-setup] ERROR: command failed at line %s (exit=%s)\n' "${BASH_LINENO[0]}" "$exit_code" >&2 - exit "$exit_code" -} -trap on_error ERR - -usage() { - cat <= 2)) || fail "--persistent-root requires a path" - STORAGE_MODE="persistent" - PERSISTENT_ROOT="$2" - shift 2 - ;; - --accelerator-type) - (($# >= 2)) || fail "--accelerator-type requires a value" - ACCELERATOR_TYPE="$2" - shift 2 - ;; - --torch-version) - (($# >= 2)) || fail "--torch-version requires a value" - TORCH_VERSION="$2" - shift 2 - ;; - --force-framework-reinstall) - FORCE_FRAMEWORK_REINSTALL=1 - shift - ;; - --run-quick-smoke) - RUN_QUICK_SMOKE=1 - shift - ;; - --optimizer) - (($# >= 2)) || fail "--optimizer requires a value" - SMOKE_OPTIMIZER="$2" - shift 2 - ;; - -h|--help) - usage - exit 0 - ;; - *) - fail "unknown option: $1" - ;; - esac -done - -[[ -f "${PROJECT_DIR}/pyproject.toml" ]] || \ - fail "could not find ${PROJECT_DIR}/pyproject.toml; run this script from the cloned repository" -command -v "$PYTHON_BIN" >/dev/null 2>&1 || fail "Python executable not found: ${PYTHON_BIN}" - -if [[ "$STORAGE_MODE" == "auto" ]]; then - if [[ -d /mnt/disks/rg-data && -w /mnt/disks/rg-data ]]; then - STORAGE_MODE="persistent" - PERSISTENT_ROOT="/mnt/disks/rg-data" - else - STORAGE_MODE="ephemeral" - fi -fi - -if [[ "$STORAGE_MODE" == "persistent" ]]; then - [[ -n "$PERSISTENT_ROOT" ]] || fail "persistent storage selected without a root" - [[ -d "$PERSISTENT_ROOT" ]] || fail "persistent root does not exist: ${PERSISTENT_ROOT}" - [[ -w "$PERSISTENT_ROOT" ]] || fail "persistent root is not writable: ${PERSISTENT_ROOT}" -else - warn "using ephemeral /tmp storage; checkpoints and data disappear when the TPU VM is deleted" - warn "do not use --ephemeral for a long or scientific run" -fi - -log "project directory: ${PROJECT_DIR}" -log "python: $($PYTHON_BIN --version 2>&1)" -log "accelerator type: ${ACCELERATOR_TYPE}" -log "storage mode: ${STORAGE_MODE}" - -log "upgrading pip, setuptools, and wheel (required to avoid UNKNOWN-0.0.0 installs)" -"$PYTHON_BIN" -m pip install --user --upgrade pip setuptools wheel - -# Remove the bad artifact produced by the old TPU VM packaging toolchain, if present. -if "$PYTHON_BIN" -m pip show UNKNOWN >/dev/null 2>&1; then - log "removing stale UNKNOWN-0.0.0 package" - "$PYTHON_BIN" -m pip uninstall -y UNKNOWN -fi - -framework_matches() { - PJRT_DEVICE=TPU "$PYTHON_BIN" - "$TORCH_VERSION" <<'PY' -import re -import sys - -expected = sys.argv[1] - -def major_minor(value: str) -> tuple[int, int]: - match = re.match(r"^\s*(\d+)\.(\d+)", value) - if not match: - raise RuntimeError(f"cannot parse version: {value!r}") - return int(match.group(1)), int(match.group(2)) - -import torch -import torch_xla - -if major_minor(torch.__version__) != major_minor(expected): - raise SystemExit(1) -if major_minor(torch_xla.__version__) != major_minor(expected): - raise SystemExit(1) -PY -} - -if ((FORCE_FRAMEWORK_REINSTALL)) || ! framework_matches; then - log "installing matching torch=${TORCH_VERSION} and torch_xla=${TORCH_VERSION}" - "$PYTHON_BIN" -m pip uninstall -y torch torch_xla torchvision || true - "$PYTHON_BIN" -m pip install --user \ - "torch==${TORCH_VERSION}" \ - "torch_xla[tpu]==${TORCH_VERSION}" \ - -f https://storage.googleapis.com/libtpu-releases/index.html \ - -f https://storage.googleapis.com/libtpu-wheels/index.html -else - log "matching torch/torch_xla ${TORCH_VERSION%.*}.x stack is already installed" -fi - -log "installing the nanoGPT package without editable mode" -"$PYTHON_BIN" -m pip install --user --upgrade "$PROJECT_DIR" - -ENV_DIR="${HOME}/.config/rg_optimizers" -ENV_FILE="${ENV_DIR}/tpu_env.sh" -mkdir -p "$ENV_DIR" -{ - printf '# Generated by %s. Safe to source from ~/.bashrc.\n' "$SCRIPT_NAME" - printf 'export PATH="$HOME/.local/bin:$PATH"\n' - printf 'export PJRT_DEVICE=TPU\n' - printf 'export TPU_ACCELERATOR_TYPE=%q\n' "$ACCELERATOR_TYPE" - printf 'unset XLA_USE_BF16\n' - printf 'unset XLA_DOWNCAST_BF16\n' - if [[ "$STORAGE_MODE" == "persistent" ]]; then - printf 'unset RG_NANOGPT_ALLOW_EPHEMERAL_TPU_STORAGE\n' - printf 'export RG_TPU_PERSISTENT_ROOT=%q\n' "$PERSISTENT_ROOT" - else - printf 'unset RG_TPU_PERSISTENT_ROOT\n' - printf 'export RG_NANOGPT_ALLOW_EPHEMERAL_TPU_STORAGE=1\n' - fi -} > "$ENV_FILE" -chmod 600 "$ENV_FILE" - -BASHRC_LINE='[ -f "$HOME/.config/rg_optimizers/tpu_env.sh" ] && source "$HOME/.config/rg_optimizers/tpu_env.sh"' -if ! grep -Fqx "$BASHRC_LINE" "${HOME}/.bashrc" 2>/dev/null; then - printf '\n%s\n' "$BASHRC_LINE" >> "${HOME}/.bashrc" -fi - -# shellcheck disable=SC1090 -source "$ENV_FILE" - -log "verifying the PyTorch/XLA runtime" -"$PYTHON_BIN" - "$TORCH_VERSION" <<'PY' -import json -import re -import sys - -expected = sys.argv[1] - -def major_minor(value: str) -> tuple[int, int]: - match = re.match(r"^\s*(\d+)\.(\d+)", value) - if not match: - raise RuntimeError(f"cannot parse version: {value!r}") - return int(match.group(1)), int(match.group(2)) - -import torch -import torch_xla -import torch_xla.core.xla_model as xm -import torch_xla.runtime as xr - -if major_minor(torch.__version__) != major_minor(expected): - raise RuntimeError(f"torch version mismatch: {torch.__version__} vs {expected}") -if major_minor(torch_xla.__version__) != major_minor(expected): - raise RuntimeError(f"torch_xla version mismatch: {torch_xla.__version__} vs {expected}") -if str(xr.device_type()).upper() != "TPU": - raise RuntimeError(f"PJRT device is not TPU: {xr.device_type()!r}") - -devices = xm.get_xla_supported_devices("TPU") -if not devices: - raise RuntimeError("PyTorch/XLA found no TPU devices") - -print(json.dumps({ - "torch": torch.__version__, - "torch_xla": torch_xla.__version__, - "pjrt_device": xr.device_type(), - "tpu_devices": devices, -}, indent=2)) -PY - -RG_ENV="$(command -v rg-onehead-env || true)" -[[ -n "$RG_ENV" ]] || RG_ENV="${HOME}/.local/bin/rg-onehead-env" -[[ -x "$RG_ENV" ]] || fail "rg-onehead-env was not installed" - -log "verifying nanoGPT runtime and storage resolution" -"$RG_ENV" --device auto - -if ((RUN_QUICK_SMOKE)); then - case "$SMOKE_OPTIMIZER" in - adamw|muon|sgd_momentum) ;; - *) fail "unsupported quick-smoke optimizer: ${SMOKE_OPTIMIZER}" ;; - esac - - QUICK_CONFIG="${TMPDIR:-/tmp}/rg_nanogpt_tpu_quick.yaml" - log "writing non-scientific quick-smoke config: ${QUICK_CONFIG}" - "$PYTHON_BIN" - "${PROJECT_DIR}/configs/tpu_smoke.yaml" "$QUICK_CONFIG" <<'PY' -import sys -from pathlib import Path - -import yaml - -source = Path(sys.argv[1]) -target = Path(sys.argv[2]) -with source.open(encoding="utf-8") as handle: - cfg = yaml.safe_load(handle) - -cfg["protocol"]["name"] = "rg_nanogpt_one_head_tpu_quick_smoke" -cfg["protocol"]["description"] = ( - "Non-scientific TPU/XLA compatibility and throughput smoke test." -) -cfg["dataset"]["train_tokens"] = 4_000_000 -cfg["dataset"]["val_tokens"] = 100_000 -cfg["dataset"]["test_tokens"] = 100_000 -cfg["training"]["target_epochs"] = 0.10 -cfg["training"]["epoch_interval"] = 0.10 -cfg["training"]["eval_interval_steps"] = 25 -cfg["training"]["checkpoint_interval_steps"] = 25 -cfg["weightwatcher"]["enabled"] = False - -with target.open("w", encoding="utf-8") as handle: - yaml.safe_dump(cfg, handle, sort_keys=False) -PY - - RG_PREPARE="$(command -v rg-onehead-prepare || true)" - RG_TRAIN="$(command -v rg-onehead-train || true)" - [[ -n "$RG_PREPARE" ]] || RG_PREPARE="${HOME}/.local/bin/rg-onehead-prepare" - [[ -n "$RG_TRAIN" ]] || RG_TRAIN="${HOME}/.local/bin/rg-onehead-train" - - log "preparing the small pinned FineWeb-Edu cache" - "$RG_PREPARE" --config "$QUICK_CONFIG" --force - - log "running the quick TPU smoke test with optimizer=${SMOKE_OPTIMIZER}" - "$RG_TRAIN" \ - --config "$QUICK_CONFIG" \ - --optimizer "$SMOKE_OPTIMIZER" \ - --device auto \ - --no-resume -fi - -log "setup complete" -log "future SSH shells load ${ENV_FILE} automatically" -log "for this already-open shell, run: source ${ENV_FILE}" -if ((RUN_QUICK_SMOKE == 0)); then - log "optional quick test: ${SCRIPT_NAME} --ephemeral --run-quick-smoke" -fi From 830672b395f8d7e34f94eb9fd8b4fe461b297061 Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Mon, 24 Aug 2026 12:51:52 -0700 Subject: [PATCH 5/5] Align TPU runbook with repository path invariants --- baseline/nanogpt_one_head/TPU.md | 24 ++++++++++++++---------- 1 file changed, 14 insertions(+), 10 deletions(-) diff --git a/baseline/nanogpt_one_head/TPU.md b/baseline/nanogpt_one_head/TPU.md index d34b5f8..c289bd9 100644 --- a/baseline/nanogpt_one_head/TPU.md +++ b/baseline/nanogpt_one_head/TPU.md @@ -72,27 +72,31 @@ Inside the TPU VM: cd /tmp git clone https://github.com/CalculatedContent/rg_optimizers.git cd rg_optimizers/baseline/nanogpt_one_head -bash scripts/setup_tpu_v5e.sh --ephemeral +bash setup_tpu_v5e.sh --ephemeral ``` -The script records the required environment in: +The script resolves the current user's login directory dynamically and records +the required environment below that directory. To inspect the exact path: -```text -~/.config/rg_optimizers/tpu_env.sh +```bash +USER_ROOT="$(getent passwd "$(id -u)" | cut -d: -f6)" +ENV_FILE="${USER_ROOT}/.config/rg_optimizers/tpu_env.sh" +printf '%s\n' "$ENV_FILE" ``` -It also adds an idempotent source line to `~/.bashrc`, so future SSH shells load -the TPU environment automatically. For the shell that launched the setup -script, load it explicitly after the script returns: +It also adds an idempotent source line to the current user's `.bashrc`, so +future SSH shells load the TPU environment automatically. For the shell that +launched the setup script, load it explicitly after the script returns: ```bash -source ~/.config/rg_optimizers/tpu_env.sh +USER_ROOT="$(getent passwd "$(id -u)" | cut -d: -f6)" +source "${USER_ROOT}/.config/rg_optimizers/tpu_env.sh" ``` For a disposable setup plus a small, non-scientific AdamW throughput test: ```bash -bash scripts/setup_tpu_v5e.sh --ephemeral --run-quick-smoke +bash setup_tpu_v5e.sh --ephemeral --run-quick-smoke ``` That optional quick smoke creates a reduced 4M/100k/100k-token corpus, disables @@ -103,7 +107,7 @@ tables. For a real run, mount durable storage and use: ```bash -bash scripts/setup_tpu_v5e.sh \ +bash setup_tpu_v5e.sh \ --persistent-root /mnt/disks/rg-data ```