From 7f63d336ef143a1737c8973394cc4553043c34df Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Wed, 30 Sep 2026 17:23:17 -0700 Subject: [PATCH 01/36] Add single-host multi-chip XLA SPMD MuonClip training Add four-chip data parallelism, global QK clipping, full-state resume checks, long-run and smoke configs, and a GitHub-based TPU runbook. Preserve tied embedding weights through XLA device conversion. --- baseline/nanogpt_one_head/TPU.md | 15 +- baseline/nanogpt_one_head/TPU_SPMD.md | 194 ++++++++++++++++++ .../configs/muonclip_tpu_spmd_long.yaml | 126 ++++++++++++ .../configs/muonclip_tpu_spmd_smoke.yaml | 123 +++++++++++ .../src/rg_nanogpt_one_head/engine.py | 4 + .../src/rg_nanogpt_one_head/evaluation.py | 9 +- .../src/rg_nanogpt_one_head/model.py | 9 + .../src/rg_nanogpt_one_head/muonclip.py | 5 + .../rg_nanogpt_one_head/muonclip_resilient.py | 16 +- .../src/rg_nanogpt_one_head/run_utils.py | 5 + .../src/rg_nanogpt_one_head/runtime.py | 3 + .../src/rg_nanogpt_one_head/tpu_spmd.py | 102 +++++++++ .../src/rg_nanogpt_one_head/tpu_spmd_check.py | 185 +++++++++++++++++ .../src/rg_nanogpt_one_head/train_loop.py | 6 +- .../src/rg_nanogpt_one_head/training.py | 4 + .../nanogpt_one_head/tests/test_tpu_spmd.py | 98 +++++++++ 16 files changed, 883 insertions(+), 21 deletions(-) create mode 100644 baseline/nanogpt_one_head/TPU_SPMD.md create mode 100644 baseline/nanogpt_one_head/configs/muonclip_tpu_spmd_long.yaml create mode 100644 baseline/nanogpt_one_head/configs/muonclip_tpu_spmd_smoke.yaml create mode 100644 baseline/nanogpt_one_head/src/rg_nanogpt_one_head/tpu_spmd.py create mode 100644 baseline/nanogpt_one_head/src/rg_nanogpt_one_head/tpu_spmd_check.py create mode 100644 baseline/nanogpt_one_head/tests/test_tpu_spmd.py diff --git a/baseline/nanogpt_one_head/TPU.md b/baseline/nanogpt_one_head/TPU.md index c289bd9e..0156393d 100644 --- a/baseline/nanogpt_one_head/TPU.md +++ b/baseline/nanogpt_one_head/TPU.md @@ -8,13 +8,11 @@ this order: TPU/XLA -> CUDA -> Apple MPS -> CPU ``` -The implementation is intentionally **single-process**. On a multi-chip TPU -slice it uses one XLA device. This preserves the reference batch size, gradient -accumulation, optimizer-step count, and learning-rate schedule. A future -multi-device protocol must explicitly define gradient reduction, global batch -size, data sampling, checkpoint ownership, and WeightWatcher ownership; the -current runner refuses a multi-process XLA launch rather than silently changing -the experiment. +Historical configs use one XLA device. For **one MuonClip model across all four +chips**, use the opt-in SPMD configs and [multi-chip runbook](TPU_SPMD.md). +SPMD retains one Python process, one global batch/RNG stream, and one checkpoint +writer while XLA partitions the batch across the chips. Multi-process and +multi-host launches remain unsupported by this trainer. ## TPU Builders v5e Flex-Start quick path @@ -138,7 +136,8 @@ 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. +Historical configs use one chip. With `runtime.tpu_spmd: true`, `xla:0` is a +logical device spanning the configured physical chips; see [TPU_SPMD.md](TPU_SPMD.md). ## Manual installation fallback diff --git a/baseline/nanogpt_one_head/TPU_SPMD.md b/baseline/nanogpt_one_head/TPU_SPMD.md new file mode 100644 index 00000000..73d292a7 --- /dev/null +++ b/baseline/nanogpt_one_head/TPU_SPMD.md @@ -0,0 +1,194 @@ +# One MuonClip run across four TPU chips + +This is single-host XLA SPMD data parallelism for the v5e-4, with one Python +process and one checkpoint/WeightWatcher owner. It is not a four-run sweep. +Multi-host TPU slices are rejected. The old single-chip configs remain opt-out. + +## Provision from Cloud Shell + +If you already submitted the request from chat, reuse its variables and skip +creation. Otherwise: + +```bash +export RG_PROJECT=YOUR_PROJECT_ID +export RG_ZONE=us-west4-a +export RG_REQUEST="ww-long-$(date -u +%Y%m%d-%H%M%S)" +export RG_NODE="${RG_REQUEST}-node" +declare -p RG_PROJECT RG_ZONE RG_REQUEST RG_NODE > "$HOME/ww-long-session.env" +gcloud alpha compute tpus queued-resources create "$RG_REQUEST" \ + --project="$RG_PROJECT" --zone="$RG_ZONE" --node-id="$RG_NODE" \ + --accelerator-type=v5litepod-4 --runtime-version=v2-alpha-tpuv5-lite \ + --provisioning-model=flex-start --max-run-duration=72h \ + --valid-until-duration=1h --labels=purpose=muonclip-longrun + +gcloud alpha compute tpus queued-resources describe "$RG_REQUEST" \ + --project="$RG_PROJECT" --zone="$RG_ZONE" --format='yaml(state)' +``` + +Wait for `ACTIVE`. The acquisition window is one hour; the run limit is 72 +hours after provisioning. Flex-start supports up to seven days. To restore +variables in another Cloud Shell session, source `~/ww-long-session.env`. + +## Durable storage + +List existing disks first; an existing data disk can be reused only when it is +available for attachment. Do not detach a disk from an active experiment. + +```bash +gcloud compute disks list --project="$RG_PROJECT" \ + --filter="zone:($RG_ZONE)" --format='table(name,sizeGb,type.basename(),users)' +``` + +For a NEW dedicated disk (skip creation if deliberately reusing an existing +one, and set `RG_DISK` to that disk's name): + +```bash +export RG_DISK="${RG_REQUEST}-data" +gcloud compute disks create "$RG_DISK" --project="$RG_PROJECT" \ + --zone="$RG_ZONE" --size=100GB --type=pd-balanced + +gcloud alpha compute tpus tpu-vm attach-disk "$RG_NODE" \ + --project="$RG_PROJECT" --zone="$RG_ZONE" --disk="$RG_DISK" --mode=read-write + +declare -p RG_PROJECT RG_ZONE RG_REQUEST RG_NODE RG_DISK > "$HOME/ww-long-session.env" +gcloud compute tpus tpu-vm ssh "$RG_NODE" --project="$RG_PROJECT" --zone="$RG_ZONE" +``` + +Inside the VM, inspect `lsblk -f` and `ls -l /dev/disk/by-id/`. Identify the +attached data filesystem by its size and UUID. The guest device alias may be +`google-persistent-disk-1`; it need not match the Cloud disk resource name. +Mount the existing filesystem using the UUID shown by `lsblk -f`: + +```bash +lsblk -f +ls -l /dev/disk/by-id/ +export RG_DATA_UUID=YOUR_EXISTING_DATA_FILESYSTEM_UUID +sudo mkdir -p /mnt/disks/rg-data +sudo mount "UUID=$RG_DATA_UUID" /mnt/disks/rg-data +``` + +For a **new blank disk only**, set `RG_DEVICE` to its verified device path and +format it with `sudo mkfs.ext4 -m 0 "$RG_DEVICE"`, then read its UUID with +`lsblk -f` and mount as above. Never format a reused data disk. If it has a +partition, use the UUID of the filesystem-bearing partition. Then: + +```bash +sudo chown "$(id -u):$(id -g)" /mnt/disks/rg-data +findmnt /mnt/disks/rg-data +``` + +The data disk survives TPU deletion and continues to incur storage charges. +The script below deliberately requires this mount for the long experiment. + +## Install the branch inside the TPU VM + +```bash +mountpoint -q /mnt/disks/rg-data +cd /mnt/disks/rg-data +git clone --branch codex/tpu-spmd-muonclip \ + https://github.com/CalculatedContent/rg_optimizers.git rg_optimizers_spmd +cd rg_optimizers_spmd/baseline/nanogpt_one_head +unset TPU_VISIBLE_CHIPS TPU_PROCESS_BOUNDS TPU_CHIPS_PER_PROCESS_BOUNDS +unset XLA_USE_SPMD XLA_USE_BF16 XLA_DOWNCAST_BF16 +bash setup_tpu_v5e.sh --persistent-root /mnt/disks/rg-data +source "$HOME/.config/rg_optimizers/tpu_env.sh" +``` + +The config enables SPMD before XLA device creation. Do not use `torchrun`, +`xmp.spawn`, or one process per chip. Do not run the older per-chip sweep on +these same chips at the same time. + +## Acceptance check and measured speed + +This check downloads no data. It uses synthetic tokens only to test numerical +correctness, global QK clipping, evaluation and full-state checkpoint resume. +The optional throughput benchmark uses the real model shape and global batch. + +```bash +python3 -m rg_nanogpt_one_head.tpu_spmd_check --backend tpu --chips 4 \ + --benchmark-config configs/muonclip_tpu_spmd_long.yaml --benchmark-steps 30 \ + --output /mnt/disks/rg-data/spmd-four-chip-check.json +``` + +Compare one chip using the SAME global batch and model shape, in a fresh process: + +```bash +TPU_VISIBLE_CHIPS=0 TPU_PROCESS_BOUNDS=1,1,1 TPU_CHIPS_PER_PROCESS_BOUNDS=1,1,1 \ +python3 -m rg_nanogpt_one_head.tpu_spmd_check --backend tpu --chips 1 \ + --benchmark-config configs/muonclip_tpu_spmd_long.yaml --benchmark-steps 30 \ + --output /mnt/disks/rg-data/spmd-one-chip-check.json +``` + +A correctness check must pass before continuing. Benchmark timing excludes the +first five warm-up updates and excludes WeightWatcher/checkpoint/evaluation I/O. +Do not interpret CPU-XLA throughput as TPU performance. Four chips need not be +faster for this small model. + +## FineWeb smoke, then long run + +Set `RG_DATA` to an existing compatible 80M/1M/1M-token cache or let the trainer +prepare it at a new path. The example reuses the existing one-head cache. +Keep smoke and long results separate. + +```bash +export RG_DATA=/mnt/disks/rg-data/rg-nanogpt-one-head/data +python3 -m rg_nanogpt_one_head.muonclip --device tpu --optimizer muon_clip \ + --config configs/muonclip_tpu_spmd_smoke.yaml --seeds 1337 \ + --data-root "$RG_DATA" --results-root /mnt/disks/rg-data/muonclip-spmd-smoke +``` + +For the long run, use `tmux new -s muonclip-long` (install tmux if absent), then: + +```bash +python3 -u -m rg_nanogpt_one_head.muonclip_resilient --device tpu \ + --config configs/muonclip_tpu_spmd_long.yaml --seed 1337 \ + --data-root "$RG_DATA" \ + --results-root /mnt/disks/rg-data/muonclip-spmd-long \ + --max-no-progress-failures 3 +``` + +Detach with Ctrl-b then d. Reattach with `tmux attach -t muonclip-long`. +Run the same command to resume after replacing a VM and mounting the disk. +Keep the source revision, config, dependencies, and four-chip topology fixed; +resume checks reject a changed protocol. The supervisor handles worker process +failures, not provisioning/replacing an expired VM. + +## Protocol and monitoring + +- Global microbatch: 32 sequences, context 256, accumulation 1 = 8,192 tokens + per update. On four chips each contributes eight sequences / 2,048 tokens. +- Replicated weights and gradients; global mean-loss gradients before clipping, + momentum, and Muon Newton–Schulz. Never average independent Muon updates. +- QK maxima reduce over the global batch and every accumulation microbatch; + replicated per-head maxima drive identical clipping. +- Explicit FP32 baseline; BF16 is not enabled or claimed validated. +- 2,150,000 updates, peak LR 2e-4, 2,000 warm-up updates, full-horizon cosine to + 2e-5. This is a starting protocol, not an assertion of optimality. +- Full optimizer/model/RNG/sampler checkpoint every 500 updates; existing + finite-state validation and atomic replacement preserve the last good state. +- Train/validation loss, perplexity and accuracy every 1,000 updates; held-out + test report at completion. CPU WeightWatcher and permanent model snapshots + every 10,000 updates plus endpoints. Raw and clipped alpha stay separate. +- WeightWatcher remains synchronous and runs once per snapshot. This commit + provides multi-chip support; adaptive per-layer LR/backtracking is not enabled. +- The model remains one block/head, width 128, GPT-2 tokenizer, clean FineWeb-Edu. + Synthetic acceptance tokens are not experimental memorization data. + +Live outputs are under `muonclip-spmd-long/muon_clip/seed_1337/`: `metrics.csv`, +`muonclip_qk.csv`, spectral outputs and `checkpoint_latest.pt`. + +## Stop the allocation early (Cloud Shell) + +After training has stopped and its durable checkpoint is verified: + +```bash +source "$HOME/ww-long-session.env" +gcloud alpha compute tpus queued-resources delete "$RG_REQUEST" \ + --project="$RG_PROJECT" --zone="$RG_ZONE" --force +``` + +## References + +- https://docs.pytorch.org/xla/release/r2.6/perf/spmd_basic.html +- https://docs.cloud.google.com/tpu/docs/request-using-flex-start +- https://docs.cloud.google.com/tpu/docs/attach-durable-block-storage diff --git a/baseline/nanogpt_one_head/configs/muonclip_tpu_spmd_long.yaml b/baseline/nanogpt_one_head/configs/muonclip_tpu_spmd_long.yaml new file mode 100644 index 00000000..67514614 --- /dev/null +++ b/baseline/nanogpt_one_head/configs/muonclip_tpu_spmd_long.yaml @@ -0,0 +1,126 @@ +# batch_size is GLOBAL: 32 sequences / 4 chips = 8 sequences per chip. +protocol: + name: rg_nanogpt_muonclip_tpu_spmd_long + version: 16 + description: 'Single-host four-chip SPMD MuonClip run: 2,150,000 updates, 8,192 + global tokens/update. Global batches are sharded, weights and optimizer updates + are replicated. Full-horizon cosine with 2,000-step warmup; periodic evaluation, + CPU WeightWatcher and resumable checkpoints. No adaptive alpha controller is enabled.' +dataset: + name: HuggingFaceFW/fineweb-edu + config: sample-10BT + split: train + revision: 593b3a867298afb8ce42625a270ef20ddcad28f9 + tokenizer: gpt2 + train_tokens: 80000000 + val_tokens: 1000000 + test_tokens: 1000000 +model: + vocab_size: 50257 + block_size: 256 + n_layer: 1 + n_head: 1 + n_embd: 128 + dropout: 0.0 + bias: false + tie_weights: true +training: + seeds: + - 1337 + batch_size: 32 + grad_accum_steps: 1 + target_epochs: 220.16 + epoch_interval: 1.024 + eval_interval_steps: 1000 + eval_batches: 8 + checkpoint_interval_steps: 500 + grad_clip: 1.0 +optimizer_profiles: + sgd_momentum: + display_name: SGD + Nesterov momentum + family: sgd + learning_rate: 0.05 + min_learning_rate: 0.005 + warmup_fraction: 0.1 + lr_schedule_epochs: 1.0 + schedule: warmup_cosine + momentum: 0.9 + dampening: 0.0 + nesterov: true + weight_decay: 0.01 + adamw: + display_name: AdamW + family: adamw + learning_rate: 0.0006 + min_learning_rate: 6.0e-05 + warmup_fraction: 0.01 + lr_schedule_epochs: 1.0 + schedule: warmup_cosine + beta1: 0.9 + beta2: 0.95 + epsilon: 1.0e-08 + weight_decay: 0.1 + muon: + display_name: Muon + auxiliary AdamW + family: muon + matrix_learning_rate: 0.02 + matrix_min_learning_rate: 0.002 + aux_learning_rate: 0.0003 + aux_min_learning_rate: 3.0e-05 + warmup_fraction: 0.05 + lr_schedule_epochs: 1.0 + schedule: warmup_cosine + momentum: 0.95 + nesterov: true + newton_schulz_steps: 5 + muon_epsilon: 1.0e-07 + matrix_weight_decay: 0.01 + beta1: 0.9 + beta2: 0.95 + epsilon: 1.0e-08 + aux_weight_decay: 0.01 + muon_clip: + display_name: MuonClip + RMS-matched updates + auxiliary AdamW + family: muon_clip + learning_rate: 0.0002 + min_learning_rate: 2.0e-05 + warmup_fraction: 0.0009302325581395349 + lr_schedule_epochs: 220.16 + schedule: warmup_cosine + momentum: 0.95 + nesterov: false + newton_schulz_steps: 5 + muon_epsilon: 1.0e-07 + weight_decay: 0.1 + update_rms_scale: 0.2 + qk_clip_threshold: 100.0 + qk_clip_balance: 0.5 + qk_diagnostics_interval: 1000 + beta1: 0.9 + beta2: 0.95 + epsilon: 1.0e-08 +evaluation: + train_probe_seed: 21001 + validation_probe_seed: 22001 + test_probe_seed: 23001 + bleu_probe_seed: 24001 + bleu_examples: 64 + bleu_prompt_tokens: 64 + bleu_continuation_tokens: 32 + bleu_batch_size: 4 +weightwatcher: + enabled: true + ERG: true + randomize: true + strict: true + min_evals: 20 + fix_fingers: clip_xmax + max_fingers: 10 + require_raw_alpha: true +runtime: + matmul_precision: high + mps_fallback: true + deterministic_algorithms: false + empty_mps_cache_after_weightwatcher: true + tpu_spmd: true + tpu_expected_chips: 4 diff --git a/baseline/nanogpt_one_head/configs/muonclip_tpu_spmd_smoke.yaml b/baseline/nanogpt_one_head/configs/muonclip_tpu_spmd_smoke.yaml new file mode 100644 index 00000000..0b377523 --- /dev/null +++ b/baseline/nanogpt_one_head/configs/muonclip_tpu_spmd_smoke.yaml @@ -0,0 +1,123 @@ +protocol: + name: rg_nanogpt_muonclip_tpu_spmd_smoke + version: 16 + description: Short real-FineWeb multi-chip trainer check; separate results from + the long run. +dataset: + name: HuggingFaceFW/fineweb-edu + config: sample-10BT + split: train + revision: 593b3a867298afb8ce42625a270ef20ddcad28f9 + tokenizer: gpt2 + train_tokens: 80000000 + val_tokens: 1000000 + test_tokens: 1000000 +model: + vocab_size: 50257 + block_size: 256 + n_layer: 1 + n_head: 1 + n_embd: 128 + dropout: 0.0 + bias: false + tie_weights: true +training: + seeds: + - 1337 + batch_size: 32 + grad_accum_steps: 1 + target_epochs: 0.002048 + epoch_interval: 0.002048 + eval_interval_steps: 10 + eval_batches: 1 + checkpoint_interval_steps: 5 + grad_clip: 1.0 +optimizer_profiles: + sgd_momentum: + display_name: SGD + Nesterov momentum + family: sgd + learning_rate: 0.05 + min_learning_rate: 0.005 + warmup_fraction: 0.1 + lr_schedule_epochs: 0.002048 + schedule: warmup_cosine + momentum: 0.9 + dampening: 0.0 + nesterov: true + weight_decay: 0.01 + adamw: + display_name: AdamW + family: adamw + learning_rate: 0.0006 + min_learning_rate: 6.0e-05 + warmup_fraction: 0.01 + lr_schedule_epochs: 0.002048 + schedule: warmup_cosine + beta1: 0.9 + beta2: 0.95 + epsilon: 1.0e-08 + weight_decay: 0.1 + muon: + display_name: Muon + auxiliary AdamW + family: muon + matrix_learning_rate: 0.02 + matrix_min_learning_rate: 0.002 + aux_learning_rate: 0.0003 + aux_min_learning_rate: 3.0e-05 + warmup_fraction: 0.05 + lr_schedule_epochs: 0.002048 + schedule: warmup_cosine + momentum: 0.95 + nesterov: true + newton_schulz_steps: 5 + muon_epsilon: 1.0e-07 + matrix_weight_decay: 0.01 + beta1: 0.9 + beta2: 0.95 + epsilon: 1.0e-08 + aux_weight_decay: 0.01 + muon_clip: + display_name: MuonClip + RMS-matched updates + auxiliary AdamW + family: muon_clip + learning_rate: 0.0002 + min_learning_rate: 2.0e-05 + warmup_fraction: 0.1 + lr_schedule_epochs: 0.002048 + schedule: warmup_cosine + momentum: 0.95 + nesterov: false + newton_schulz_steps: 5 + muon_epsilon: 1.0e-07 + weight_decay: 0.1 + update_rms_scale: 0.2 + qk_clip_threshold: 100.0 + qk_clip_balance: 0.5 + qk_diagnostics_interval: 5 + beta1: 0.9 + beta2: 0.95 + epsilon: 1.0e-08 +evaluation: + train_probe_seed: 21001 + validation_probe_seed: 22001 + test_probe_seed: 23001 + bleu_probe_seed: 24001 + bleu_examples: 4 + bleu_prompt_tokens: 64 + bleu_continuation_tokens: 32 + bleu_batch_size: 4 +weightwatcher: + enabled: true + ERG: true + randomize: true + strict: true + min_evals: 20 + fix_fingers: clip_xmax + max_fingers: 10 + require_raw_alpha: true +runtime: + matmul_precision: high + mps_fallback: true + deterministic_algorithms: false + empty_mps_cache_after_weightwatcher: true + tpu_spmd: true + tpu_expected_chips: 4 diff --git a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/engine.py b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/engine.py index 51d57f51..63d7654c 100644 --- a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/engine.py +++ b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/engine.py @@ -47,6 +47,7 @@ synchronize, ) from .train_loop import execute_training_loop +from .tpu_spmd import initialize as initialize_spmd, replicate_model def run_one( @@ -67,6 +68,7 @@ def run_one( if resume and overwrite: raise ValueError("resume and overwrite are mutually exclusive") + initialize_spmd(cfg, device) data_root = Path(data_root) results_root = Path(results_root) run_dir = run_directory(results_root, optimizer_name, int(seed)) @@ -128,6 +130,7 @@ def run_one( ) model = GPT(GPTConfig(**cfg["model"])).to(resolved_device) + replicate_model(model) handles = make_optimizer_handles(model, profile) train_generator = torch.Generator(device="cpu").manual_seed( int(seed) + 11 @@ -209,6 +212,7 @@ def run_one( ) resumed_from_checkpoint = True model.to(resolved_device) + replicate_model(model) synchronize(resolved_device) truncate_spectral_after(run_dir, start_step) truncate_muonclip_qk_after(run_dir, start_step) diff --git a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/evaluation.py b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/evaluation.py index 71aeb63d..997fa97a 100644 --- a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/evaluation.py +++ b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/evaluation.py @@ -9,6 +9,7 @@ import torch.nn as nn from .runtime import is_xla_device, mark_step, synchronize, tree_to_cpu +from .tpu_spmd import batch_to_device @dataclass(frozen=True) @@ -145,8 +146,8 @@ def evaluate_probe( top5_correct = 0 total = 0 for x_cpu, y_cpu in probe: - x = x_cpu.to(device) - y = y_cpu.to(device) + x = batch_to_device(x_cpu, device) + y = batch_to_device(y_cpu, device) logits, loss = model(x, y) if loss is None: raise RuntimeError( @@ -184,8 +185,8 @@ def evaluate_probe( total = 0 batches = 0 for x_cpu, y_cpu in probe: - x = x_cpu.to(device) - y = y_cpu.to(device) + x = batch_to_device(x_cpu, device) + y = batch_to_device(y_cpu, device) logits, loss = model(x, y) if loss is None: raise RuntimeError("evaluation forward pass did not return loss") diff --git a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/model.py b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/model.py index 029b4267..caf606f2 100644 --- a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/model.py +++ b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/model.py @@ -186,6 +186,15 @@ def __init__(self, cfg: GPTConfig) -> None: std=residual_std, ) + def _apply(self, fn, recurse=True): + # CPU -> XLA may replace each Parameter object independently, breaking + # the embedding/head alias established in __init__. Retie BEFORE an + # optimizer is built, including for resumed models and CPU WW copies. + result = super()._apply(fn, recurse=recurse) + if self.cfg.tie_weights: + self.lm_head.weight = self.token_embedding.weight + return result + @staticmethod def _init_module(module: nn.Module) -> None: if isinstance(module, (nn.Linear, nn.Embedding)): diff --git a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/muonclip.py b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/muonclip.py index 83278195..df9b44bb 100644 --- a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/muonclip.py +++ b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/muonclip.py @@ -18,6 +18,8 @@ import torch import torch.nn.functional as F +from .tpu_spmd import replicate + _INSTALLED = False _CURRENT_RUN_DIR: Path | None = None @@ -409,6 +411,9 @@ def _record_qk_logits(attention, scores: torch.Tensor) -> None: .amax(dim=0) .float() ) + # The amax includes the GLOBAL batch axis. Replication forces the + # compiler to complete the cross-chip max before per-head QK clipping. + replicate(value) previous = getattr(attention, "_muonclip_max_logits", None) setattr( attention, diff --git a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/muonclip_resilient.py b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/muonclip_resilient.py index c353ff9d..32fe07b5 100644 --- a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/muonclip_resilient.py +++ b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/muonclip_resilient.py @@ -1,12 +1,13 @@ from __future__ import annotations -"""Progress-aware fresh-process recovery for long MuonClip MPS runs. +"""Progress-aware fresh-process recovery for long MuonClip accelerator runs. -This launcher runs exactly one MuonClip worker process at a time. If Metal/MPS -terminates a worker, the next worker resumes from ``checkpoint_latest.pt``. +This launcher runs exactly one MuonClip worker process at a time. If the +accelerator runtime terminates a worker, the next worker resumes from +``checkpoint_latest.pt``. The retry budget counts only consecutive failures that do not advance the -verified checkpoint. Therefore intermittent Metal failures may be recovered +verified checkpoint. Intermittent accelerator failures may be recovered throughout a long run, while deterministic failures at one checkpoint stop after a small bounded number of attempts. """ @@ -98,7 +99,7 @@ def _worker_command(args: argparse.Namespace) -> list[str]: "--results-root", str(Path(args.results_root).expanduser().resolve()), "--device", - "mps", + getattr(args, "device", "mps"), "--mps-worker", "--mps-retries", "0", @@ -131,7 +132,7 @@ def run_resilient(args: argparse.Namespace) -> int: attempt += 1 (run_dir / "run_failed.json").unlink(missing_ok=True) print( - "[one-head-resilient] starting fresh MPS worker " + "[one-head-resilient] starting fresh accelerator worker " f"attempt={attempt} checkpoint_step=" f"{last_verified_step if last_verified_step is not None else 'none'} " f"no_progress_failures={no_progress_failures}/" @@ -283,10 +284,11 @@ def run_resilient(args: argparse.Namespace) -> int: def main() -> None: parser = argparse.ArgumentParser( description=( - "Run a long MuonClip MPS experiment with progress-aware " + "Run a long MuonClip experiment with progress-aware " "fresh-process checkpoint recovery" ) ) + parser.add_argument("--device", choices=("auto", "tpu", "xla", "mps", "cuda", "cpu"), default="mps") parser.add_argument("--config", required=True) parser.add_argument("--seed", type=int, required=True) parser.add_argument("--data-root", required=True) diff --git a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/run_utils.py b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/run_utils.py index 79905bba..ae4cec56 100644 --- a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/run_utils.py +++ b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/run_utils.py @@ -129,6 +129,11 @@ def runtime_identity_payload(metadata: dict) -> dict: *_COMMON_RUNTIME_IDENTITY_FIELDS, *_ACCELERATOR_RUNTIME_IDENTITY_FIELDS.get(accelerator, ()), ) + if accelerator == "tpu": + # Old single-chip manifests remain readable. New SPMD runs record + # topology so a masked device cannot silently change a resumed run. + metadata = {"xla_spmd": False, "xla_spmd_chips": 1, **metadata} + fields = (*fields, "xla_spmd", "xla_spmd_chips") missing = [field for field in fields if field not in metadata] if missing: raise RuntimeError( diff --git a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/runtime.py b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/runtime.py index 2cea820e..115b368c 100644 --- a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/runtime.py +++ b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/runtime.py @@ -552,6 +552,9 @@ def runtime_metadata(device: torch.device) -> dict[str, Any]: ), } ) + if device.type == "xla": + from .tpu_spmd import metadata as spmd_metadata + metadata.update(spmd_metadata()) block_id, block_source = _hardware_block_identity(metadata) metadata["hardware_block_id"] = block_id metadata["hardware_block_id_source"] = block_source diff --git a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/tpu_spmd.py b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/tpu_spmd.py new file mode 100644 index 00000000..690a27b5 --- /dev/null +++ b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/tpu_spmd.py @@ -0,0 +1,102 @@ +from __future__ import annotations + +"""Single-host XLA SPMD data parallelism; batch sizes are always GLOBAL. + +There is one Python process, RNG stream, optimizer and checkpoint writer. XLA +partitions the batch and inserts collectives for global reductions. Never use +xm.optimizer_step or manually divide the loss by the chip count in this mode. +""" + +import os +from typing import Any + +import numpy as np +import torch + +_MESH: Any = None +_XS: Any = None +_CHIPS = 0 + + +def initialize(cfg: dict, requested: str | torch.device) -> None: + """Must run before choose_device creates any XLA tensors.""" + enabled = cfg.get("runtime", {}).get("tpu_spmd", False) + if not isinstance(enabled, bool): + raise ValueError("runtime.tpu_spmd must be a boolean") + if not enabled: + if _MESH is not None or os.environ.get("XLA_USE_SPMD") == "1": + raise RuntimeError("SPMD is already enabled; use a tpu_spmd config in a fresh process") + return + if str(requested) not in {"auto", "tpu", "xla", "xla:0"}: + raise ValueError("runtime.tpu_spmd requires --device tpu (or auto)") + expected = int(cfg["runtime"].get("tpu_expected_chips", 4)) + batch = int(cfg["training"]["batch_size"]) + if expected < 1 or batch % expected: + raise ValueError("global training.batch_size must be divisible by tpu_expected_chips") + os.environ.setdefault("PJRT_DEVICE", "TPU") + from .runtime import _load_xla + + modules = _load_xla(required=True) + assert modules is not None + _, xr, _ = modules + if xr.device_type() != "TPU": + raise RuntimeError("TPU SPMD requested but PJRT_DEVICE is not TPU") + _initialize_mesh(xr, expected) + + +def _initialize_mesh(xr, expected: int) -> None: + """Kept separate so the same sharding code can be checked on XLA CPU.""" + global _MESH, _XS, _CHIPS + if _MESH is not None: + if _CHIPS != expected: + raise RuntimeError("Cannot change the SPMD mesh in a running process") + return + xr.use_spmd() + # This implementation deliberately uses one host, as on the v5e-4. + # Multi-host SPMD requires a separate input/checkpoint ownership protocol. + count = int(xr.global_runtime_device_count()) + local = int(xr.addressable_runtime_device_count()) + if int(xr.process_count()) != 1 or local != count: + raise RuntimeError("TPU SPMD trainer currently requires one host/process") + if count != expected: + raise RuntimeError( + f"Expected {expected} TPU chips, found {count}; check TPU_VISIBLE_CHIPS " + "and remove per-chip sweep environment settings" + ) + import torch_xla.distributed.spmd as xs + + _XS = xs + _MESH = xs.Mesh(np.arange(count), (count,), ("data",)) + _CHIPS = count + + +def metadata() -> dict: + return {"xla_spmd": _MESH is not None, "xla_spmd_chips": _CHIPS or 1} + + +def replicate(tensor: torch.Tensor) -> torch.Tensor: + if _MESH is not None and tensor.device.type == "xla": + _XS.mark_sharding(tensor, _MESH, (None,) * tensor.ndim) + return tensor + + +def replicate_model(model) -> None: + for tensor in (*model.parameters(), *model.buffers()): + replicate(tensor) + + +def replicate_gradients(model) -> None: + # Replicated gradients MUST precede gradient clipping, momentum and NS. + # These are constraints on GLOBAL tensors, not a second gradient average. + for parameter in model.parameters(): + if parameter.grad is not None: + replicate(parameter.grad) + + +def batch_to_device(tensor: torch.Tensor, device: torch.device) -> torch.Tensor: + value = tensor.to(device) + if _MESH is not None and value.device.type == "xla": + if value.ndim < 1 or value.shape[0] % _CHIPS: + raise ValueError("SPMD batch must divide evenly across the data mesh") + _XS.mark_sharding(value, _MESH, ("data",) + (None,) * (value.ndim - 1)) + return value diff --git a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/tpu_spmd_check.py b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/tpu_spmd_check.py new file mode 100644 index 00000000..618f96d2 --- /dev/null +++ b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/tpu_spmd_check.py @@ -0,0 +1,185 @@ +from __future__ import annotations + +"""Opt-in numerical, resume and throughput check on actual XLA devices. + +Run with --backend tpu on the VM; --backend cpu is a local XLA compiler test, +not evidence of TPU hardware performance. This never downloads a corpus. +""" + +import argparse +from copy import deepcopy +import json +import os +from pathlib import Path +import tempfile +import time + +import torch + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--backend", choices=("tpu", "cpu"), default="tpu") + parser.add_argument("--chips", type=int, default=4) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--benchmark-config", type=Path) + parser.add_argument("--benchmark-steps", type=int, default=30) + args = parser.parse_args() + if args.chips < 1 or args.benchmark_steps < 1: + parser.error("chips and benchmark-steps must be positive") + os.environ["PJRT_DEVICE"] = args.backend.upper() + if args.backend == "cpu": + os.environ["CPU_NUM_DEVICES"] = str(args.chips) + torch.set_num_threads(1) + + import torch_xla.core.xla_model as xm + import torch_xla.runtime as xr + import torch_xla + from . import tpu_spmd as spmd + from .muonclip import install_muonclip_extension + + install_muonclip_extension() + from .checkpoints import save_training_checkpoint, load_training_checkpoint_for_resume + from .config import load_config, optimizer_profile + from .evaluation import evaluate_probe + from .model import GPT, GPTConfig + from .optimizers import make_optimizer_handles, zero_grad, optimizer_step + from .runtime import mark_step, synchronize, tree_to_cpu, parameter_snapshot + + spmd._initialize_mesh(xr, args.chips) + device = xm.xla_device() + cfg_path = Path.cwd() / "configs" / "muonclip_reference.yaml" + cfg = load_config(cfg_path) + cfg["model"].update(vocab_size=64, block_size=8, n_embd=16, n_head=2, n_layer=1) + cfg["training"].update(batch_size=2 * args.chips, grad_accum_steps=2) + cfg["runtime"].update(tpu_spmd=True, tpu_expected_chips=args.chips) + profile = optimizer_profile(cfg, "muon_clip") + # Force clipping to activate so a wrong cross-chip max cannot pass quietly. + profile.update(qk_clip_threshold=0.0001, qk_diagnostics_interval=100) + torch.manual_seed(314) + reference = GPT(GPTConfig(**cfg["model"])) + model = deepcopy(reference).to(device) + spmd.replicate_model(model) + assert model.token_embedding.weight is model.lm_head.weight, "XLA transfer broke tied weights" + cpu_handles = make_optimizer_handles(reference, profile) + handles = make_optimizer_handles(model, profile) + generator = torch.Generator().manual_seed(2718) + batches = [ + (torch.randint(64, (2 * args.chips, 8), generator=generator), + torch.randint(64, (2 * args.chips, 8), generator=generator)) + for _ in range(2) + ] + + def update(net, opts, target): + zero_grad(opts) + for x, y in batches: + _, loss = net(spmd.batch_to_device(x, target), spmd.batch_to_device(y, target)) + (loss / len(batches)).backward() + spmd.replicate_gradients(net) + gradients = {n: p.grad.detach().cpu().clone() for n, p in net.named_parameters()} + maxima = [b.attn._muonclip_max_logits.detach().cpu().clone() for b in net.blocks] + torch.nn.utils.clip_grad_norm_(net.parameters(), 1.0, foreach=False) + optimizer_step(opts) + mark_step(target) + synchronize(target) + return gradients, maxima + + cpu_grad, cpu_max = update(reference, cpu_handles, torch.device("cpu")) + xla_grad, xla_max = update(model, handles, device) + for name in cpu_grad: + torch.testing.assert_close(xla_grad[name], cpu_grad[name], atol=3e-5, rtol=3e-3) + for actual, expected in zip(xla_max, cpu_max): + torch.testing.assert_close(actual, expected, atol=3e-5, rtol=3e-3) + for name, expected in reference.state_dict().items(): + torch.testing.assert_close(model.state_dict()[name].cpu(), expected, atol=3e-5, rtol=3e-3) + for block_max in xla_max: + assert bool((block_max > profile["qk_clip_threshold"]).any()), "QK clip was not exercised" + + probe_cpu = evaluate_probe(reference, batches, torch.device("cpu")) + probe_xla = evaluate_probe(model, batches, device) + for metric in ("loss", "accuracy", "top5_accuracy"): + torch.testing.assert_close(torch.tensor(probe_xla[metric]), torch.tensor(probe_cpu[metric]), atol=3e-5, rtol=3e-3) + + with tempfile.TemporaryDirectory() as temporary: + path = Path(temporary) / "checkpoint.pt" + save_training_checkpoint( + path, model=model, handles=handles, step=1, + best_validation_loss=probe_xla["loss"], best_validation_step=1, + elapsed_seconds=1.0, fingerprint="spmd-numerical-check", cfg=cfg, + optimizer_name="muon_clip", seed=314, train_generator=generator, + resume_diagnostics={"previous_eval_snapshot": parameter_snapshot(model), + "last_grad_pre": 1.0, "last_grad_post": 1.0, "last_clipped": False}, + ) + payload = torch.load(path, map_location="cpu", weights_only=False) + assert all(t.device.type == "cpu" for t in payload["model"].values()) + resumed = GPT(GPTConfig(**cfg["model"])).to(device) + spmd.replicate_model(resumed) + resumed_handles = make_optimizer_handles(resumed, profile) + restored_generator = torch.Generator() + state = load_training_checkpoint_for_resume( + path, model=resumed, handles=resumed_handles, + expected_fingerprint="spmd-numerical-check", train_generator=restored_generator, + ) + spmd.replicate_model(resumed) + assert state[0] == 1 + assert torch.equal(restored_generator.get_state(), generator.get_state()) + update(model, handles, device) + update(resumed, resumed_handles, device) + for name, expected in tree_to_cpu(model.state_dict()).items(): + torch.testing.assert_close(resumed.state_dict()[name].cpu(), expected, atol=1e-6, rtol=1e-5) + # Check the next update as well: a missing momentum/Adam state can hide + # behind an apparently successful model-only checkpoint load. + update(model, handles, device) + update(resumed, resumed_handles, device) + for name, expected in tree_to_cpu(model.state_dict()).items(): + torch.testing.assert_close(resumed.state_dict()[name].cpu(), expected, atol=1e-6, rtol=1e-5) + + report = {"passed": True, "backend": xr.device_type(), **spmd.metadata(), + "torch": torch.__version__, "torch_xla": torch_xla.__version__, + "checks": ["global gradients", "global per-head QK maxima", "clipped update", + "train/eval metrics", "CPU checkpoint", "optimizer and sampler resume"], + "max_gradient_error": max(float((cpu_grad[n] - xla_grad[n]).abs().max()) for n in cpu_grad)} + + if args.benchmark_config: + benchmark_cfg = load_config(args.benchmark_config) + batch = int(benchmark_cfg["training"]["batch_size"]) + accum = int(benchmark_cfg["training"]["grad_accum_steps"]) + context = int(benchmark_cfg["model"]["block_size"]) + if batch % args.chips: + raise ValueError("Benchmark global batch must divide evenly across chips") + net = GPT(GPTConfig(**benchmark_cfg["model"])).to(device) + spmd.replicate_model(net) + opts = make_optimizer_handles(net, optimizer_profile(benchmark_cfg, "muon_clip")) + x = spmd.batch_to_device(torch.randint(net.cfg.vocab_size, (batch, context)), device) + y = spmd.batch_to_device(torch.randint(net.cfg.vocab_size, (batch, context)), device) + # Includes the same sampler copy/sharding boundary in each timed step. + # Uses synthetic tokens; excludes WW, evaluation and checkpoint I/O. + x_cpu, y_cpu = x.cpu(), y.cpu() + def benchmark_step(): + zero_grad(opts) + for _ in range(accum): + _, loss = net(spmd.batch_to_device(x_cpu, device), spmd.batch_to_device(y_cpu, device)) + (loss / accum).backward() + spmd.replicate_gradients(net) + torch.nn.utils.clip_grad_norm_(net.parameters(), float(benchmark_cfg["training"]["grad_clip"]), foreach=False) + optimizer_step(opts) + mark_step(device) + for _ in range(5): + benchmark_step() + synchronize(device) + start = time.perf_counter() + for _ in range(args.benchmark_steps): + benchmark_step() + synchronize(device) + seconds = time.perf_counter() - start + report["benchmark"] = {"steps": args.benchmark_steps, "seconds": seconds, + "global_tokens_per_update": batch * accum * context, + "tokens_per_second": args.benchmark_steps * batch * accum * context / seconds, + "includes_monitoring": False} + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(report, indent=2) + "\n") + print(json.dumps(report, indent=2), flush=True) + + +if __name__ == "__main__": + main() diff --git a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/train_loop.py b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/train_loop.py index ef6ebaeb..83abc122 100644 --- a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/train_loop.py +++ b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/train_loop.py @@ -25,6 +25,7 @@ ) from .spectral import run_weightwatcher from .random_canaries import RandomCanaryExperiment +from .tpu_spmd import batch_to_device, replicate_gradients def _require_finite_metrics( @@ -519,8 +520,8 @@ def execute_training_loop( completed_step=completed_steps, micro_index=micro_index, ) - x = x_cpu.to(device) - y = y_cpu.to(device) + x = batch_to_device(x_cpu, device) + y = batch_to_device(y_cpu, device) _, loss = model(x, y) if loss is None: raise RuntimeError( @@ -528,6 +529,7 @@ def execute_training_loop( ) (loss / grad_accum).backward() + replicate_gradients(model) grad_pre_tensor = gradient_norm(model.parameters()) clip = float(cfg["training"]["grad_clip"]) if clip > 0: diff --git a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/training.py b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/training.py index 85a78ae9..2d66a50c 100644 --- a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/training.py +++ b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/training.py @@ -25,6 +25,7 @@ from .engine import run_one from .run_utils import run_directory, run_is_complete from .runtime import choose_device +from .tpu_spmd import initialize as initialize_spmd def _resolve_roots( @@ -110,6 +111,7 @@ def run_optimizer_replicates( progress: bool = True, ) -> list[Path]: del config_path + initialize_spmd(cfg, device) data_path, results_path, resolved_device = _resolve_roots( data_root=data_root, results_root=results_root, @@ -153,6 +155,7 @@ def run_all_replicates( overwrite: bool = False, progress: bool = True, ) -> list[Path]: + initialize_spmd(cfg, device) data_path, results_path, resolved_device = _resolve_roots( data_root=data_root, results_root=results_root, @@ -548,6 +551,7 @@ def main() -> None: if not seeds: parser.error("at least one seed is required") + initialize_spmd(cfg, args.device) resolved_device = choose_device(args.device) if ( resolved_device.type == "mps" diff --git a/baseline/nanogpt_one_head/tests/test_tpu_spmd.py b/baseline/nanogpt_one_head/tests/test_tpu_spmd.py new file mode 100644 index 00000000..451e99ff --- /dev/null +++ b/baseline/nanogpt_one_head/tests/test_tpu_spmd.py @@ -0,0 +1,98 @@ +import os +from pathlib import Path +import subprocess +import sys +from types import SimpleNamespace + +import pytest +import torch + +from rg_nanogpt_one_head import tpu_spmd as spmd +from rg_nanogpt_one_head.model import GPT, GPTConfig +from rg_nanogpt_one_head.muonclip_resilient import _worker_command + + +def test_cpu_transfer_preserves_tied_weights_even_when_parameters_are_replaced(): + model = GPT(GPTConfig(vocab_size=32, block_size=8, n_embd=8)) + previous = torch.__future__.get_overwrite_module_params_on_conversion() + try: + torch.__future__.set_overwrite_module_params_on_conversion(True) + model.to(dtype=torch.float64) + assert model.token_embedding.weight is model.lm_head.weight + finally: + torch.__future__.set_overwrite_module_params_on_conversion(previous) + + +def test_untied_model_stays_untied(): + model = GPT(GPTConfig(vocab_size=32, block_size=8, n_embd=8, tie_weights=False)) + model.to(dtype=torch.float64) + assert model.token_embedding.weight is not model.lm_head.weight + + +@pytest.mark.parametrize("device", ["cpu", "cuda", "mps"]) +def test_spmd_rejects_wrong_accelerator(device): + cfg = {"runtime": {"tpu_spmd": True}, "training": {"batch_size": 8}} + with pytest.raises(ValueError, match="requires --device tpu"): + spmd.initialize(cfg, device) + + +def test_spmd_rejects_indivisible_global_batch(): + cfg = {"runtime": {"tpu_spmd": True, "tpu_expected_chips": 4}, "training": {"batch_size": 2}} + with pytest.raises(ValueError, match="divisible"): + spmd.initialize(cfg, "tpu") + + +def test_spmd_rejects_silent_mode_switch(monkeypatch): + monkeypatch.setenv("XLA_USE_SPMD", "1") + with pytest.raises(RuntimeError, match="fresh process"): + spmd.initialize({"runtime": {}}, "tpu") + + +@pytest.mark.parametrize("global_count,local_count,process_count,message", [ + (1, 1, 1, "Expected 4"), (8, 4, 2, "one host/process"), +]) +def test_spmd_rejects_wrong_topology(global_count, local_count, process_count, message): + xr = SimpleNamespace(use_spmd=lambda: None, + global_runtime_device_count=lambda: global_count, + addressable_runtime_device_count=lambda: local_count, + process_count=lambda: process_count) + with pytest.raises(RuntimeError, match=message): + spmd._initialize_mesh(xr, 4) + + +def test_supervisor_passes_tpu_device(tmp_path): + args = SimpleNamespace(config=tmp_path / "c.yaml", seed=1337, + data_root=tmp_path / "data", results_root=tmp_path / "results", device="tpu") + cmd = _worker_command(args) + assert cmd[cmd.index("--device") + 1] == "tpu" + + +def test_longrun_global_batch_and_schedule(): + from rg_nanogpt_one_head.muonclip import install_muonclip_extension + install_muonclip_extension() + from rg_nanogpt_one_head.config import load_config, tokens_per_step, max_steps, lr_schedule_steps, warmup_steps + root = Path(__file__).resolve().parents[1] + cfg = load_config(root / "configs/muonclip_tpu_spmd_long.yaml") + assert cfg["weightwatcher"]["fix_fingers"] == "clip_xmax" + assert cfg["weightwatcher"]["require_raw_alpha"] is True + assert tokens_per_step(cfg) == 8192 + assert cfg["training"]["batch_size"] // cfg["runtime"]["tpu_expected_chips"] == 8 + assert max_steps(cfg) == 2150000 + profile = cfg["optimizer_profiles"]["muon_clip"] + assert lr_schedule_steps(cfg, profile) == 2150000 + assert warmup_steps(profile, 2150000) == 2000 + smoke = load_config(root / "configs/muonclip_tpu_spmd_smoke.yaml") + assert max_steps(smoke) == 20 + + +@pytest.mark.skipif(os.environ.get("RG_TEST_XLA_SPMD") != "1", reason="opt-in real XLA CPU test") +def test_real_four_device_xla_gradients_clipping_and_resume(tmp_path): + root = Path(__file__).resolve().parents[1] + environment = {**os.environ, "PJRT_DEVICE": "CPU", "CPU_NUM_DEVICES": "4", "OMP_NUM_THREADS": "1"} + environment.pop("XLA_USE_SPMD", None) + result = subprocess.run([ + sys.executable, "-m", "rg_nanogpt_one_head.tpu_spmd_check", "--backend", "cpu", + "--chips", "4", "--output", str(tmp_path / "report.json"), + ], cwd=root, env=environment, capture_output=True, text=True, timeout=180) + assert result.returncode == 0, result.stdout + result.stderr + assert '"passed": true' in result.stdout From 763844f15f2cd995287c9e712b7100ef55d54532 Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Wed, 30 Sep 2026 17:58:57 -0700 Subject: [PATCH 02/36] Set explicit XLA matmul precision for TPU training and checks Apply the native XLA 2.6 precision control, use highest precision in both SPMD configurations, and validate HLO precision before gradient comparisons. Keep existing numerical tolerances. Record precision in runtime identity. Validation: 79 tests passed, including the real four-device CPU-XLA check. Physical TPU rerun remains pending. --- baseline/nanogpt_one_head/TPU_SPMD.md | 6 ++++ .../configs/muonclip_tpu_spmd_long.yaml | 2 +- .../configs/muonclip_tpu_spmd_smoke.yaml | 2 +- .../src/rg_nanogpt_one_head/run_utils.py | 4 ++- .../src/rg_nanogpt_one_head/runtime.py | 28 ++++++++++++++-- .../src/rg_nanogpt_one_head/tpu_spmd_check.py | 33 +++++++++++++++---- .../tests/test_runtime_tpu.py | 21 ++++++++++++ .../nanogpt_one_head/tests/test_tpu_spmd.py | 2 ++ 8 files changed, 86 insertions(+), 12 deletions(-) diff --git a/baseline/nanogpt_one_head/TPU_SPMD.md b/baseline/nanogpt_one_head/TPU_SPMD.md index 73d292a7..aea31c00 100644 --- a/baseline/nanogpt_one_head/TPU_SPMD.md +++ b/baseline/nanogpt_one_head/TPU_SPMD.md @@ -162,6 +162,11 @@ failures, not provisioning/replacing an expired VM. - QK maxima reduce over the global batch and every accumulation microbatch; replicated per-head maxima drive identical clipping. - Explicit FP32 baseline; BF16 is not enabled or claimed validated. +- Both SPMD configs request `matmul_precision: highest`. The runtime sets XLA's + separate native precision control as well as PyTorch's setting. PyTorch/XLA + 2.6 does not inherit this setting from `torch.set_float32_matmul_precision`. + The check verifies the emitted HLO precision and a precision-sensitive matrix + product before comparing gradients. Numerical tolerances are unchanged. - 2,150,000 updates, peak LR 2e-4, 2,000 warm-up updates, full-horizon cosine to 2e-5. This is a starting protocol, not an assertion of optimality. - Full optimizer/model/RNG/sampler checkpoint every 500 updates; existing @@ -190,5 +195,6 @@ gcloud alpha compute tpus queued-resources delete "$RG_REQUEST" \ ## References - https://docs.pytorch.org/xla/release/r2.6/perf/spmd_basic.html +- https://docs.pytorch.org/xla/master/tutorials/precision_tutorial.html - https://docs.cloud.google.com/tpu/docs/request-using-flex-start - https://docs.cloud.google.com/tpu/docs/attach-durable-block-storage diff --git a/baseline/nanogpt_one_head/configs/muonclip_tpu_spmd_long.yaml b/baseline/nanogpt_one_head/configs/muonclip_tpu_spmd_long.yaml index 67514614..ba452fee 100644 --- a/baseline/nanogpt_one_head/configs/muonclip_tpu_spmd_long.yaml +++ b/baseline/nanogpt_one_head/configs/muonclip_tpu_spmd_long.yaml @@ -118,7 +118,7 @@ weightwatcher: max_fingers: 10 require_raw_alpha: true runtime: - matmul_precision: high + matmul_precision: highest mps_fallback: true deterministic_algorithms: false empty_mps_cache_after_weightwatcher: true diff --git a/baseline/nanogpt_one_head/configs/muonclip_tpu_spmd_smoke.yaml b/baseline/nanogpt_one_head/configs/muonclip_tpu_spmd_smoke.yaml index 0b377523..a00e2487 100644 --- a/baseline/nanogpt_one_head/configs/muonclip_tpu_spmd_smoke.yaml +++ b/baseline/nanogpt_one_head/configs/muonclip_tpu_spmd_smoke.yaml @@ -115,7 +115,7 @@ weightwatcher: max_fingers: 10 require_raw_alpha: true runtime: - matmul_precision: high + matmul_precision: highest mps_fallback: true deterministic_algorithms: false empty_mps_cache_after_weightwatcher: true diff --git a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/run_utils.py b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/run_utils.py index ae4cec56..9f367dba 100644 --- a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/run_utils.py +++ b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/run_utils.py @@ -101,6 +101,7 @@ def _package_versions() -> dict[str, str]: "tpu": ( "torch_xla_version", "pjrt_device", + "xla_matmul_precision", "tpu_accelerator_type", "xla_process_count", "xla_process_index", @@ -132,7 +133,8 @@ def runtime_identity_payload(metadata: dict) -> dict: if accelerator == "tpu": # Old single-chip manifests remain readable. New SPMD runs record # topology so a masked device cannot silently change a resumed run. - metadata = {"xla_spmd": False, "xla_spmd_chips": 1, **metadata} + metadata = {"xla_spmd": False, "xla_spmd_chips": 1, + "xla_matmul_precision": "unset", **metadata} fields = (*fields, "xla_spmd", "xla_spmd_chips") missing = [field for field in fields if field not in metadata] if missing: diff --git a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/runtime.py b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/runtime.py index 115b368c..18a16c47 100644 --- a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/runtime.py +++ b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/runtime.py @@ -27,6 +27,7 @@ "TPU_CHIPS_PER_HOST_BOUNDS", "TPU_HOST_BOUNDS", ) +_XLA_MATMUL_PRECISION: str | None = None def _major_minor(version: str) -> tuple[int, int] | None: @@ -201,10 +202,29 @@ def _xla_process_index(xr: Any) -> int: return 0 +def configure_matmul_precision(device: torch.device, precision: str) -> None: + """Configure both ATen and XLA; XLA 2.6 has a separate precision control.""" + global _XLA_MATMUL_PRECISION + mapping = {"highest": "HIGHEST", "high": "HIGH", "medium": "DEFAULT"} + if precision not in mapping: + raise ValueError(f"Unsupported matmul precision: {precision!r}") + torch.set_float32_matmul_precision(precision) + if device.type == "xla": + modules = _load_xla(required=True) + assert modules is not None + torch_xla, _, _ = modules + # The public torch_xla.backends wrapper is newer than our pinned 2.6 + # stack. Its native binding is present in 2.6 and controls HLO dot ops. + setter = getattr(getattr(torch_xla, "_XLAC", None), "_xla_set_mat_mul_precision", None) + if not callable(setter): + raise RuntimeError("PyTorch/XLA cannot set explicit matrix multiplication precision") + setter(mapping[precision]) + _XLA_MATMUL_PRECISION = mapping[precision].lower() + + def configure_runtime(device: torch.device, cfg: dict) -> None: - torch.set_float32_matmul_precision( - str(cfg["runtime"].get("matmul_precision", "high")) - ) + precision = str(cfg["runtime"].get("matmul_precision", "high")) + configure_matmul_precision(torch.device("cpu"), precision) if device.type == "cuda": allow_tf32 = bool(cfg["runtime"].get("allow_tf32", False)) torch.backends.cuda.matmul.allow_tf32 = allow_tf32 @@ -253,6 +273,7 @@ def configure_runtime(device: torch.device, cfg: dict) -> None: f"Detected process_count={count}. Run the ordinary single-process " "launcher or implement an explicitly distributed protocol." ) + configure_matmul_precision(device, precision) deterministic = bool( cfg["runtime"].get("deterministic_algorithms", False) ) @@ -542,6 +563,7 @@ def runtime_metadata(device: torch.device) -> dict[str, Any]: torch_xla, "__version__", "unknown" ), "pjrt_device": _xla_device_type(xr), + "xla_matmul_precision": _XLA_MATMUL_PRECISION or "unset", "tpu_accelerator_type": str( os.environ.get("TPU_ACCELERATOR_TYPE", "unknown") ), diff --git a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/tpu_spmd_check.py b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/tpu_spmd_check.py index 618f96d2..904fa501 100644 --- a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/tpu_spmd_check.py +++ b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/tpu_spmd_check.py @@ -44,15 +44,31 @@ def main() -> None: from .evaluation import evaluate_probe from .model import GPT, GPTConfig from .optimizers import make_optimizer_handles, zero_grad, optimizer_step - from .runtime import mark_step, synchronize, tree_to_cpu, parameter_snapshot + from .runtime import configure_runtime, mark_step, synchronize, tree_to_cpu, parameter_snapshot spmd._initialize_mesh(xr, args.chips) device = xm.xla_device() cfg_path = Path.cwd() / "configs" / "muonclip_reference.yaml" cfg = load_config(cfg_path) + benchmark_cfg = load_config(args.benchmark_config) if args.benchmark_config else None + precision = str(benchmark_cfg["runtime"].get("matmul_precision", "highest")) if benchmark_cfg else "highest" cfg["model"].update(vocab_size=64, block_size=8, n_embd=16, n_head=2, n_layer=1) cfg["training"].update(batch_size=2 * args.chips, grad_accum_steps=2) - cfg["runtime"].update(tpu_spmd=True, tpu_expected_chips=args.chips) + cfg["runtime"].update(tpu_spmd=True, tpu_expected_chips=args.chips, matmul_precision=precision) + # torch.set_float32_matmul_precision alone does not configure XLA 2.6. + # Apply the same explicit runtime policy used by the actual trainer. + configure_runtime(device, cfg) + print(f"[spmd-check] backend={xr.device_type()} chips={args.chips} matmul_precision={precision}", flush=True) + # Inspect emitted HLO as well as metadata: CPU-XLA accepts reduced-precision + # HLO but can execute it with FP32, hiding a missing TPU precision setting. + probe_input = torch.full((128, 128), 1.0 + 2.0 ** -8).to(device) + probe_product = probe_input @ probe_input + probe_hlo = torch_xla._XLAC._get_xla_tensors_hlo([probe_product]) + expected_precision = {"highest": "highest", "high": "high"}.get(precision) + if expected_precision: + assert f"operand_precision={{{expected_precision},{expected_precision}}}" in probe_hlo, "XLA HLO is missing requested matmul precision" + if precision == "highest": + torch.testing.assert_close(probe_product.cpu(), torch.full((128, 128), 128 * (1.0 + 2.0 ** -8) ** 2), atol=1e-6, rtol=1e-6) profile = optimizer_profile(cfg, "muon_clip") # Force clipping to activate so a wrong cross-chip max cannot pass quietly. profile.update(qk_clip_threshold=0.0001, qk_diagnostics_interval=100) @@ -86,8 +102,10 @@ def update(net, opts, target): cpu_grad, cpu_max = update(reference, cpu_handles, torch.device("cpu")) xla_grad, xla_max = update(model, handles, device) + print("[spmd-check] comparing global gradients, QK maxima and clipped update", flush=True) for name in cpu_grad: - torch.testing.assert_close(xla_grad[name], cpu_grad[name], atol=3e-5, rtol=3e-3) + torch.testing.assert_close(xla_grad[name], cpu_grad[name], atol=3e-5, rtol=3e-3, + msg=lambda message, name=name: f"Gradient {name}: {message}") for actual, expected in zip(xla_max, cpu_max): torch.testing.assert_close(actual, expected, atol=3e-5, rtol=3e-3) for name, expected in reference.state_dict().items(): @@ -100,6 +118,7 @@ def update(net, opts, target): for metric in ("loss", "accuracy", "top5_accuracy"): torch.testing.assert_close(torch.tensor(probe_xla[metric]), torch.tensor(probe_cpu[metric]), atol=3e-5, rtol=3e-3) + print("[spmd-check] checking full-state checkpoint and two resumed updates", flush=True) with tempfile.TemporaryDirectory() as temporary: path = Path(temporary) / "checkpoint.pt" save_training_checkpoint( @@ -136,12 +155,14 @@ def update(net, opts, target): report = {"passed": True, "backend": xr.device_type(), **spmd.metadata(), "torch": torch.__version__, "torch_xla": torch_xla.__version__, - "checks": ["global gradients", "global per-head QK maxima", "clipped update", + "xla_matmul_precision": "default" if precision == "medium" else precision, + "checks": ["explicit HLO matmul precision", "global gradients", "global per-head QK maxima", "clipped update", "train/eval metrics", "CPU checkpoint", "optimizer and sampler resume"], "max_gradient_error": max(float((cpu_grad[n] - xla_grad[n]).abs().max()) for n in cpu_grad)} - if args.benchmark_config: - benchmark_cfg = load_config(args.benchmark_config) + print("[spmd-check] numerical and resume checks passed", flush=True) + if benchmark_cfg is not None: + print("[spmd-check] benchmarking training shape (5 warm-up updates)", flush=True) batch = int(benchmark_cfg["training"]["batch_size"]) accum = int(benchmark_cfg["training"]["grad_accum_steps"]) context = int(benchmark_cfg["model"]["block_size"]) diff --git a/baseline/nanogpt_one_head/tests/test_runtime_tpu.py b/baseline/nanogpt_one_head/tests/test_runtime_tpu.py index b541c843..402c11a2 100644 --- a/baseline/nanogpt_one_head/tests/test_runtime_tpu.py +++ b/baseline/nanogpt_one_head/tests/test_runtime_tpu.py @@ -2,6 +2,7 @@ from pathlib import Path import sys +from types import SimpleNamespace import pytest import torch @@ -143,6 +144,26 @@ def test_configure_runtime_rejects_unregistered_bf16( runtime.configure_runtime(torch.device("xla"), cfg) +@pytest.mark.parametrize("requested,expected", [("highest", "HIGHEST"), ("high", "HIGH"), ("medium", "DEFAULT")]) +def test_xla_runtime_sets_native_matmul_precision(fake_xla, monkeypatch, requested, expected): + calls = [] + monkeypatch.setenv("TPU_ACCELERATOR_TYPE", "v5litepod-4") + monkeypatch.setattr(FakeTorchXLA, "_XLAC", SimpleNamespace(_xla_set_mat_mul_precision=calls.append), raising=False) + monkeypatch.setattr(runtime, "_XLA_MATMUL_PRECISION", None) + previous = torch.get_float32_matmul_precision() + try: + runtime.configure_runtime(torch.device("xla"), {"runtime": {"matmul_precision": requested}}) + assert calls == [expected] + assert runtime.runtime_metadata(torch.device("xla"))["xla_matmul_precision"] == expected.lower() + finally: + torch.set_float32_matmul_precision(previous) + + +def test_xla_precision_cannot_silently_fall_back(fake_xla): + with pytest.raises(RuntimeError, match="cannot set explicit"): + runtime.configure_matmul_precision(torch.device("xla"), "highest") + + def test_cpu_tree_conversion_detaches_nested_tensors(): value = { "a": torch.tensor([1.0], requires_grad=True), diff --git a/baseline/nanogpt_one_head/tests/test_tpu_spmd.py b/baseline/nanogpt_one_head/tests/test_tpu_spmd.py index 451e99ff..26ee54bd 100644 --- a/baseline/nanogpt_one_head/tests/test_tpu_spmd.py +++ b/baseline/nanogpt_one_head/tests/test_tpu_spmd.py @@ -73,6 +73,7 @@ def test_longrun_global_batch_and_schedule(): from rg_nanogpt_one_head.config import load_config, tokens_per_step, max_steps, lr_schedule_steps, warmup_steps root = Path(__file__).resolve().parents[1] cfg = load_config(root / "configs/muonclip_tpu_spmd_long.yaml") + assert cfg["runtime"]["matmul_precision"] == "highest" assert cfg["weightwatcher"]["fix_fingers"] == "clip_xmax" assert cfg["weightwatcher"]["require_raw_alpha"] is True assert tokens_per_step(cfg) == 8192 @@ -82,6 +83,7 @@ def test_longrun_global_batch_and_schedule(): assert lr_schedule_steps(cfg, profile) == 2150000 assert warmup_steps(profile, 2150000) == 2000 smoke = load_config(root / "configs/muonclip_tpu_spmd_smoke.yaml") + assert smoke["runtime"]["matmul_precision"] == "highest" assert max_steps(smoke) == 20 From 0416ad9858daf992b3156d923717aa5eb96e2f11 Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Thu, 1 Oct 2026 10:27:10 -0700 Subject: [PATCH 03/36] Support open-ended MuonClip continuation with periodic test accuracy Preserve full optimizer, RNG and sampler state across bounded training segments. Track cumulative steps and test probes; add pause, recovery and checkpoint retention. Include TPU runbook and CPU/XLA regression coverage. --- baseline/nanogpt_one_head/README.md | 6 + baseline/nanogpt_one_head/TPU_CONTINUATION.md | 191 +++++++++++++ baseline/nanogpt_one_head/TPU_SPMD.md | 3 + baseline/nanogpt_one_head/pyproject.toml | 1 + .../src/rg_nanogpt_one_head/checkpoints.py | 1 + .../src/rg_nanogpt_one_head/completion.py | 51 +++- .../src/rg_nanogpt_one_head/config.py | 11 + .../src/rg_nanogpt_one_head/continuation.py | 144 ++++++++++ .../src/rg_nanogpt_one_head/engine.py | 24 ++ .../src/rg_nanogpt_one_head/monitor.py | 61 ++++- .../src/rg_nanogpt_one_head/muonclip.py | 7 + .../rg_nanogpt_one_head/muonclip_continue.py | 253 ++++++++++++++++++ .../rg_nanogpt_one_head/muonclip_resilient.py | 24 +- .../src/rg_nanogpt_one_head/run_utils.py | 8 + .../src/rg_nanogpt_one_head/train_loop.py | 33 ++- .../tests/continuation_scenario.py | 134 ++++++++++ .../tests/test_continuation.py | 92 +++++++ 17 files changed, 1031 insertions(+), 13 deletions(-) create mode 100644 baseline/nanogpt_one_head/TPU_CONTINUATION.md create mode 100644 baseline/nanogpt_one_head/src/rg_nanogpt_one_head/continuation.py create mode 100644 baseline/nanogpt_one_head/src/rg_nanogpt_one_head/muonclip_continue.py create mode 100644 baseline/nanogpt_one_head/tests/continuation_scenario.py create mode 100644 baseline/nanogpt_one_head/tests/test_continuation.py diff --git a/baseline/nanogpt_one_head/README.md b/baseline/nanogpt_one_head/README.md index eb760141..169602a9 100644 --- a/baseline/nanogpt_one_head/README.md +++ b/baseline/nanogpt_one_head/README.md @@ -367,3 +367,9 @@ update paths, common probe identity, exact split-writing and cache corruption detection, checkpoint round-tripping, LR logging semantics, Student-t intervals, direct `ERG_gap`/`num_traps` handling, tiny CPU training, and notebook structure. The same tests run in the repository's baseline CI. + +## Extended TPU training + +See [TPU_CONTINUATION.md](TPU_CONTINUATION.md) for full-state continuation beyond +the original training horizon, periodic test accuracy, retained metric history, +and safe pause/resume with bounded checkpoint storage. diff --git a/baseline/nanogpt_one_head/TPU_CONTINUATION.md b/baseline/nanogpt_one_head/TPU_CONTINUATION.md new file mode 100644 index 00000000..14875353 --- /dev/null +++ b/baseline/nanogpt_one_head/TPU_CONTINUATION.md @@ -0,0 +1,191 @@ +# Continuing MuonClip for an extended period + +`rg_nanogpt_one_head.muonclip_continue` can run an explicit number of additional +updates or keep launching successive segments until stopped. It imports a full +training checkpoint into a new run identity. The original experiment remains +intact. Do not edit the original run's duration and try to bypass its fingerprint. + +## Behavior + +- Preserves model weights, Muon momentum, Adam moments and per-parameter update + counters, RNG state, and the training sampler. The new segment resets only + logging intervals and its own best-validation selection. Its step zero is the + imported trained model, not a new random initialization. +- Defaults to the parent's minimum learning rate, held constant: **2e-5** for + `muonclip_tpu_spmd_long.yaml`. There is no repeated LR ramp. `--learning-rate` + can specify a different positive constant for an explicitly new series. + This is a conservative continuation setting, not a claim of optimality. +- Defaults to one million additional updates per segment, with no overall cap + under `--forever`. The global counter includes all parent updates; the first + extension of the 2,150,000-update baseline starts at global step 2,150,000. +- Inherits four-chip SPMD, precision, batch size, gradient clipping, optimizer + algorithm, dataset, and probe seeds. Single-host/four-chip support is retained; + this change does not add multi-host training or change the model architecture. +- Inherits full rolling checkpoints every 500 updates, train/validation + evaluation every 1,000, and WeightWatcher/epoch snapshots every 10,000. + Segment endpoints are measured too. Raw and clipped alpha remain separate. +- Adds fixed test-probe loss, perplexity, bits/token, top-1 and top-5 token + accuracy every 10,000 updates by default (`--test-interval-steps`). For the + current configuration, each probe contains 8 x 32 x 256 = **65,536 token + predictions**, using the same test windows across all segments. This is not + evaluation of every token in the million-token test split. +- Validation loss alone selects each segment's best checkpoint. Test outcomes + do not drive automatic checkpoint selection or an LR controller. Once humans + monitor test results to choose further experiments, the test set is no longer + an untouched final holdout. BLEU remains an endpoint diagnostic. +- Keeps all CSV histories, spectral outputs, manifests, configs, test reports, + and lineage. By default, keeps full and epoch checkpoints for the latest + **three completed segments plus the active segment**. Older segment checkpoint + files are removed only after completion validation; their archival marker is + `checkpoints_pruned.json`. The original parent run is never pruned. +- Pauses at a saved checkpoint on a stop request or below **5 GiB free disk**. + Metrics and logs still grow over time, so storage is not literally unlimited. + Exit code 75 means a deliberate pause; it does not consume the retry budget. +- Uses the existing progress-aware worker recovery and a writer lock inherited + by workers. A second driver cannot write the same series while either the + first driver or its worker holds the lock. Status records after host loss may + be stale; resuming acquires the lock and checks the durable state. + +## Install after the current baseline finishes + +From **Cloud Shell**, enter the current TPU VM: + +```bash +gcloud compute tpus tpu-vm ssh ww-long-20260930-232752-node \ + --project=tpu-builders-504820 --zone=us-west4-a +``` + +Run the following **inside the TPU VM**. Use the existing pinned dependencies; +this code adds no external packages. Keep the current source checkout available +for the original experiment. Do not reinstall or update dependencies mid-run. + +```bash +set -e +source "$HOME/.config/rg_optimizers/tpu_env.sh" +mountpoint -q /mnt/disks/rg-data +export RG_PARENT=/mnt/disks/rg-data/muonclip-spmd-long/muon_clip/seed_1337 +test -s "$RG_PARENT/run_complete.json" +test -s "$RG_PARENT/checkpoint_final.pt" +cd /mnt/disks/rg-data +git clone --branch codex/tpu-long-continuation \ + https://github.com/CalculatedContent/rg_optimizers.git rg_optimizers_continuation +cd rg_optimizers_continuation/baseline/nanogpt_one_head +export PYTHONPATH="$PWD/src" +git rev-parse HEAD > /mnt/disks/rg-data/continuation-source-commit.txt +python3 -m pip freeze > /mnt/disks/rg-data/continuation-environment.txt +export RG_DATA=/mnt/disks/rg-data/rg-nanogpt-one-head/data +export RG_SERIES=/mnt/disks/rg-data/muonclip-extended +``` + +Confirm the previous supervisor/worker exited before using the chips. The +completion marker is written near the end; `pgrep -af rg_nanogpt_one_head` can +show whether a worker still owns the TPU. Separate series roots have separate +locks; they do not arbitrate TPU ownership across independent experiments. + +First rerun the four-chip numerical/resume acceptance check, then a short +continuation from the actual trained checkpoint. These commands use the TPU +and should run sequentially after the previous experiment has exited: + +```bash +python3 -m rg_nanogpt_one_head.tpu_spmd_check --backend tpu --chips 4 \ + --output /mnt/disks/rg-data/continuation-spmd-check.json + +python3 -u -m rg_nanogpt_one_head.muonclip_continue start \ + --series-root /mnt/disks/rg-data/muonclip-continuation-smoke \ + --from-checkpoint "$RG_PARENT/checkpoint_final.pt" \ + --data-root "$RG_DATA" --device tpu \ + --additional-steps 20 --segment-steps 20 --test-interval-steps 10 +``` + +The smoke command returns zero only after its full completion audit succeeds. +Its worker log is in `segments/segment_000001/launch.log` under the smoke root. +If it fails, inspect that log before launching the long series. The original +checkpoint is unchanged by the smoke test. + +## Start, monitor, pause, and resume + +Start the extended series from the original final checkpoint: + +```bash +python3 -u -m rg_nanogpt_one_head.muonclip_continue start \ + --series-root "$RG_SERIES" \ + --from-checkpoint "$RG_PARENT/checkpoint_final.pt" \ + --data-root "$RG_DATA" --device tpu --forever \ + --segment-steps 1000000 --test-interval-steps 10000 \ + --keep-segments 3 --min-free-disk-gb 5 --background +``` + +`--background` detaches the driver from SSH. The command prints its PID and +driver-log path (`/mnt/disks/rg-data/muonclip-extended.driver.log`). The child +performs startup validation; check the log and status after launching. Detailed +training logs are in each segment's `launch.log`. To run a finite extension, +replace `--forever` with `--additional-steps 1000000`. + +```bash +python3 -m rg_nanogpt_one_head.muonclip_continue status --series-root "$RG_SERIES" +python3 -m rg_nanogpt_one_head.monitor --series-root "$RG_SERIES" \ + --interval 60 --no-clear +``` + +The monitor joins the baseline and all extensions using cumulative steps and +epochs. It reports latest validation accuracy, latest measured test accuracy, +recent test history, and layer alphas. Blank test entries at intermediate +validation steps mean no test probe ran at that step. In the recent history +table, accuracy is a fraction; the headline uses percent. + +Export the full retained trajectory for later plotting: + +```bash +python3 -m rg_nanogpt_one_head.monitor --series-root "$RG_SERIES" --once \ + --export-metrics /mnt/disks/rg-data/muonclip-extended-history.csv +``` + +Request a deliberate pause, then wait for `status` to report `paused` and the +worker to exit before detaching the disk or replacing the TPU: + +```bash +python3 -m rg_nanogpt_one_head.muonclip_continue stop --series-root "$RG_SERIES" +python3 -m rg_nanogpt_one_head.muonclip_continue status --series-root "$RG_SERIES" +``` + +Resume later using the **same checkout/commit, dependencies, mount path and +four-chip configuration**. Ordinary in-segment resumes retain strict runtime +and protocol checks. The resume command removes the stop request and continues +the active segment from its full checkpoint: + +```bash +source "$HOME/.config/rg_optimizers/tpu_env.sh" +cd /mnt/disks/rg-data/rg_optimizers_continuation/baseline/nanogpt_one_head +export PYTHONPATH="$PWD/src" +export RG_SERIES=/mnt/disks/rg-data/muonclip-extended +python3 -u -m rg_nanogpt_one_head.muonclip_continue resume \ + --series-root "$RG_SERIES" --background +``` + +Resume uses the options recorded in `series.json`; it does not reschedule a +running series from new CLI hyperparameters. For a deliberate LR change, pause +and start a **new series root** from the saved full checkpoint, specifying +`--learning-rate`. Preserve both roots. Model-only epoch snapshots cannot +restart optimization; full `checkpoint_latest.pt` and `checkpoint_final.pt` can. + +## TPU allocation lifetime + +The software's `--forever` option does not extend the TPU allocation. The +existing runbook requested 72 hours. Check its actual termination timestamp +from **Cloud Shell**: + +```bash +gcloud alpha compute tpus queued-resources describe ww-long-20260930-232752 \ + --project=tpu-builders-504820 --zone=us-west4-a --format=yaml +``` + +[Google's Flex-start documentation](https://docs.cloud.google.com/tpu/docs/request-using-flex-start) +states that VMs are deleted at the requested duration and that requests can be +up to seven days. A longer experiment must span allocations. Preserve the data +disk, mount it at the same path on the replacement VM, reproduce the pinned +environment, then use `resume`. The supervisor does not provision cloud VMs or +reattach disks. An abrupt termination can lose updates after the last durable +checkpoint; it does not require restarting the experiment from initialization. + +No alpha-targeting feedback controller is enabled. Continued training may +change accuracy and spectra, but cannot guarantee every fitted alpha reaches 2. diff --git a/baseline/nanogpt_one_head/TPU_SPMD.md b/baseline/nanogpt_one_head/TPU_SPMD.md index aea31c00..0374a6ca 100644 --- a/baseline/nanogpt_one_head/TPU_SPMD.md +++ b/baseline/nanogpt_one_head/TPU_SPMD.md @@ -1,5 +1,8 @@ # One MuonClip run across four TPU chips +For extending a trained checkpoint, periodic test accuracy, bounded checkpoint +retention, and open-ended training, see [TPU_CONTINUATION.md](TPU_CONTINUATION.md). + This is single-host XLA SPMD data parallelism for the v5e-4, with one Python process and one checkpoint/WeightWatcher owner. It is not a four-run sweep. Multi-host TPU slices are rejected. The old single-chip configs remain opt-out. diff --git a/baseline/nanogpt_one_head/pyproject.toml b/baseline/nanogpt_one_head/pyproject.toml index 2dae8847..e52480ef 100644 --- a/baseline/nanogpt_one_head/pyproject.toml +++ b/baseline/nanogpt_one_head/pyproject.toml @@ -36,6 +36,7 @@ rg-onehead-env = "rg_nanogpt_one_head.runtime:main" rg-onehead-monitor = "rg_nanogpt_one_head.monitor:main" rg-onehead-validation = "rg_nanogpt_one_head.validation_report:main" rg-onehead-muonclip = "rg_nanogpt_one_head.muonclip:main" +rg-onehead-muonclip-continue = "rg_nanogpt_one_head.muonclip_continue:main" rg-onehead-muonclip-walk = "rg_nanogpt_one_head.muonclip_capture:main" rg-onehead-muonclip-movie = "rg_nanogpt_one_head.muonclip_esd_movie:main" rg-onehead-muon-svd = "rg_nanogpt_one_head.muon_svd_runner:main" diff --git a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/checkpoints.py b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/checkpoints.py index 8818f251..f423e494 100644 --- a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/checkpoints.py +++ b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/checkpoints.py @@ -272,6 +272,7 @@ def save_training_checkpoint( "model_state_sha256": model_state_sha256(model_state), "optimizer_state_sha256": optimizer_state_sha256(optimizer_states), "step": int(step), + "global_step": int(cfg.get("continuation", {}).get("global_step_offset", 0)) + int(step), "best_validation_loss": float(best_validation_loss), "best_validation_step": int(best_validation_step), "elapsed_seconds": float(elapsed_seconds), diff --git a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/completion.py b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/completion.py index 247ab679..7be3a68e 100644 --- a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/completion.py +++ b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/completion.py @@ -60,6 +60,33 @@ class CompletedRunValidationError(RuntimeError): """A nominally completed run is missing, stale, or inconsistent.""" +def _validate_test_monitoring(frame, label, interval, total_steps): + steps = pd.to_numeric(frame["step"], errors="raise") + due = steps.mod(interval).eq(0) | steps.eq(total_steps) + if label == "metrics.csv": + expected = {*range(0, total_steps + 1, interval), total_steps} + if not expected.issubset(set(steps)): + _fail("metrics.csv is missing periodic test measurements") + columns = ( + "test_loss", "test_perplexity", "test_bits_per_token", + "test_accuracy", "test_top5_accuracy", "test_generalization_gap", + ) + for column in columns: + values = pd.to_numeric(frame[column], errors="coerce") + if not values.loc[~due].isna().all() or not all(math.isfinite(v) for v in values.loc[due]): + _fail(f"{label} {column} violates periodic test measurement schedule") + observed = frame.loc[due] + for _, row in observed.iterrows(): + if not 0 <= row["test_accuracy"] <= row["test_top5_accuracy"] <= 1: + _fail(f"{label} invalid test accuracy range") + if row["test_loss"] < 0 or not math.isclose(math.exp(row["test_loss"]), row["test_perplexity"], rel_tol=1e-9): + _fail(f"{label} inconsistent test loss/perplexity") + if not math.isclose(row["test_loss"] / math.log(2), row["test_bits_per_token"], rel_tol=1e-9): + _fail(f"{label} inconsistent test bits per token") + if not math.isclose(row["test_loss"] - row["train_loss"], row["test_generalization_gap"], abs_tol=1e-9): + _fail(f"{label} inconsistent test generalization gap") + + def _fail(message: str) -> NoReturn: raise CompletedRunValidationError( "completed one-head nanoGPT run is stale or inconsistent: " @@ -399,10 +426,16 @@ def validate_completed_run( final_test = test_results.get("final") selected_test = test_results.get("validation_selected") test_policy = str(test_results.get("policy", "")).lower() + test_interval = int(manifest.get("evaluation", {}).get("test_interval_steps", 0)) + if test_interval > 0: + if "monitoring" not in test_policy or "validation" not in test_policy: + _fail("test_results.json does not declare the monitoring test policy") if ( + test_interval == 0 and ( "held out" not in test_policy or "validation" not in test_policy or "never" not in test_policy + ) ): _fail("test_results.json does not declare the held-out test policy") if not isinstance(final_test, dict) or not isinstance(selected_test, dict): @@ -494,6 +527,13 @@ def validate_completed_run( metric_steps = _step_tuple(metrics, "metrics.csv") epoch_steps = _step_tuple(epoch_metrics, "epoch_metrics.csv") summary_steps = _step_tuple(summary, "spectral/summary.csv") + lineage = manifest.get("continuation") + if lineage: + offset = _as_int(lineage.get("global_step_offset"), "global step offset") + _expect(_as_int(completion.get("global_step"), "completion global step"), offset + total_steps, "completion global step") + for frame in (metrics, epoch_metrics): + if "global_step" not in frame or not pd.to_numeric(frame["global_step"], errors="coerce").eq(frame["step"] + offset).all(): + _fail("continuation global steps do not match local steps plus parent offset") weightwatcher = manifest.get("weightwatcher", {}) if not isinstance(weightwatcher, dict): _fail("manifest weightwatcher configuration is not a mapping") @@ -518,9 +558,16 @@ def validate_completed_run( f"{label} is missing held-out placeholder columns " + ", ".join(sorted(missing_held_out)) ) + monitored_columns = { + "test_loss", "test_perplexity", "test_bits_per_token", + "test_accuracy", "test_top5_accuracy", "test_generalization_gap", + } if test_interval > 0 else set() + if monitored_columns: + _validate_test_monitoring(frame, label, test_interval, total_steps) leaked = [ column for column in _HELD_OUT_CURVE_COLUMNS + if column not in monitored_columns if not frame[column].isna().all() ] if leaked: @@ -537,7 +584,7 @@ def validate_completed_run( if "test_held_out" not in epoch_metrics.columns: _fail("epoch_metrics.csv has no test_held_out column") held_out = pd.to_numeric(epoch_metrics["test_held_out"], errors="coerce") - if held_out.isna().any() or not held_out.astype(int).eq(1).all(): + if held_out.isna().any() or not held_out.eq(int(test_interval == 0)).all(): _fail("epoch_metrics.csv does not mark every test curve as held out") required_layer_columns = { @@ -836,6 +883,8 @@ def validate_completed_run( schema_version=5, ) computed_model_hash = model_state_sha256(payload["model"]) + if lineage: + _expect(_as_int(payload.get("global_step"), f"{filename} global step"), offset + expected_step, f"{filename} global step") if str(payload.get("model_state_sha256", "")) != computed_model_hash: _fail(f"{filename} model-state SHA-256 does not match") optimizer_states = payload.get("optimizers") diff --git a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/config.py b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/config.py index 4359ded2..da57b8c5 100644 --- a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/config.py +++ b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/config.py @@ -236,6 +236,10 @@ def validate_config(cfg: dict[str, Any]) -> None: raise ValueError(f"dataset.{key} is too small for the context length") training = cfg["training"] + if "max_steps" in training: + value = training["max_steps"] + if isinstance(value, bool) or not isinstance(value, int) or value < 1: + raise ValueError("training.max_steps must be a positive integer") for key in ( "batch_size", "grad_accum_steps", "target_epochs", "eval_interval_steps", "eval_batches", "checkpoint_interval_steps", @@ -270,6 +274,9 @@ def validate_config(cfg: dict[str, Any]) -> None: ) evaluation = cfg["evaluation"] + interval = evaluation.get("test_interval_steps", 0) + if isinstance(interval, bool) or not isinstance(interval, int) or interval < 0: + raise ValueError("evaluation.test_interval_steps must be a nonnegative integer") for key in ( "bleu_examples", "bleu_prompt_tokens", "bleu_continuation_tokens", "bleu_batch_size", @@ -398,6 +405,8 @@ def _steps_for_epochs(cfg: dict[str, Any], epochs: float, train_tokens: int) -> def max_steps(cfg: dict[str, Any], train_tokens: int | None = None) -> int: + if "max_steps" in cfg["training"]: + return int(cfg["training"]["max_steps"]) train_tokens = int(train_tokens or cfg["dataset"]["train_tokens"]) return _steps_for_epochs(cfg, float(cfg["training"]["target_epochs"]), train_tokens) @@ -483,5 +492,7 @@ def protocol_fingerprint( "seed": int(seed), "data_metadata": data_metadata, } + if "continuation" in cfg: + payload["continuation"] = cfg["continuation"] canonical = json.dumps(payload, sort_keys=True, separators=(",", ":"), default=str) return hashlib.sha256(canonical.encode("utf-8")).hexdigest() diff --git a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/continuation.py b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/continuation.py new file mode 100644 index 00000000..3a5d0ef1 --- /dev/null +++ b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/continuation.py @@ -0,0 +1,144 @@ +"""Explicit imports of full training state into separately identified extensions.""" +from __future__ import annotations + +from copy import deepcopy +import hashlib +import json +import math +from pathlib import Path +import shutil + +import torch + +from .checkpoints import load_training_checkpoint_for_resume +from .config import tokens_per_step, validate_config + + +class TrainingPaused(SystemExit): + """Exit without making a deliberate pause consume the recovery budget.""" + + def __init__(self, reason: str): + print(f"[one-head-pause] {reason}; restart checkpoint retained", flush=True) + super().__init__(75) + + +def file_sha256(path: Path) -> str: + digest = hashlib.sha256() + with Path(path).open("rb") as handle: + for block in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(block) + return digest.hexdigest() + + +def pause_reason(cfg: dict, run_dir: Path) -> str | None: + training = cfg["training"] + stop = training.get("stop_file") + if stop and Path(stop).exists(): + return f"stop requested by {stop}" + reserve = float(training.get("min_free_disk_gb", 0)) + if reserve > 0 and shutil.disk_usage(run_dir).free < reserve * 1024**3: + return f"free disk space fell below {reserve:g} GiB" + return None + + +def make_continuation_config( + parent_path: Path, *, steps: int, learning_rate: float | None, + test_interval: int, stop_file: Path, min_free_disk_gb: float, +) -> dict: + if steps < 1 or test_interval < 1 or min_free_disk_gb < 0: + raise ValueError("steps/test interval must be positive; disk reserve nonnegative") + parent_path = Path(parent_path).resolve() + before_hash = file_sha256(parent_path) + parent = torch.load(parent_path, map_location="cpu", weights_only=False) + if before_hash != file_sha256(parent_path): + raise RuntimeError("parent checkpoint changed during import; use an immutable final checkpoint") + if parent.get("optimizer_name") != "muon_clip" or not parent.get("optimizers"): + raise ValueError("continuation requires a full MuonClip training checkpoint") + if int(parent.get("step", 0)) <= 0 or parent.get("resume_diagnostics") is None: + raise ValueError("parent must contain trained state and deterministic resume diagnostics") + cfg = deepcopy(parent["config"]) + if cfg.get("memorization", {}).get("enabled"): + raise ValueError("continuation currently supports clean-corpus training only") + profile = cfg["optimizer_profiles"]["muon_clip"] + rate = float(profile["min_learning_rate"] if learning_rate is None else learning_rate) + if not math.isfinite(rate) or rate <= 0: + raise ValueError("continuation learning rate must be finite and positive") + epochs = steps * tokens_per_step(cfg) / int(cfg["dataset"]["train_tokens"]) + cfg["protocol"] = { + "name": "rg_nanogpt_muonclip_continuation", "version": 1, + "description": "Full-state continuation with constant LR and fixed test monitoring", + } + cfg["training"].update( + max_steps=int(steps), target_epochs=epochs, + stop_file=str(stop_file.resolve()), min_free_disk_gb=float(min_free_disk_gb), + ) + for item in cfg["optimizer_profiles"].values(): + item.pop("lr_schedule_epochs", None) + profile.update(learning_rate=rate, min_learning_rate=rate, warmup_fraction=0.0) + cfg["evaluation"]["test_interval_steps"] = int(test_interval) + cfg["continuation"] = { + "parent_checkpoint": str(parent_path), "parent_file_sha256": before_hash, + "parent_fingerprint": parent["fingerprint"], + "global_step_offset": int(parent.get("global_step", parent["step"])), + "parent_local_step": int(parent["step"]), + "seed": int(parent["seed"]), + "policy": "preserve weights, momentum, Adam steps, RNG and sampler; reset only phase diagnostics and phase best-validation selection", + } + validate_config(cfg) + return cfg + + +def import_parent_state(cfg, *, model, handles, train_generator, data_metadata, seed, current_runtime): + """Changing run identity is allowed only through this recorded import path.""" + lineage = cfg["continuation"] + path = Path(lineage["parent_checkpoint"]) + if file_sha256(path) != lineage["parent_file_sha256"]: + raise RuntimeError("parent checkpoint file SHA-256 does not match") + parent = torch.load(path, map_location="cpu", weights_only=False) + if parent["optimizer_name"] != "muon_clip" or parent["seed"] != seed: + raise RuntimeError("continuation parent optimizer/seed mismatch") + if int(parent.get("global_step", parent["step"])) != lineage["global_step_offset"]: + raise RuntimeError("continuation global step does not match parent") + old = parent["config"] + for key in ("model", "dataset", "runtime", "weightwatcher"): + if old[key] != cfg[key]: + raise RuntimeError(f"continuation cannot change {key}") + for key in ("batch_size", "grad_accum_steps", "grad_clip", "eval_batches"): + if old["training"][key] != cfg["training"][key]: + raise RuntimeError(f"continuation cannot change training.{key}") + ignored = {"learning_rate", "min_learning_rate", "warmup_fraction", "lr_schedule_epochs"} + for key in set(old["optimizer_profiles"]["muon_clip"]) | set(cfg["optimizer_profiles"]["muon_clip"]): + if key not in ignored and old["optimizer_profiles"]["muon_clip"].get(key) != cfg["optimizer_profiles"]["muon_clip"].get(key): + raise RuntimeError(f"continuation cannot change optimizer {key}") + for key, value in old["evaluation"].items(): + if key != "test_interval_steps" and cfg["evaluation"].get(key) != value: + raise RuntimeError(f"continuation cannot change evaluation.{key}") + manifest = json.loads((path.parent / "manifest.json").read_text()) + if manifest["protocol_fingerprint"] != lineage["parent_fingerprint"]: + raise RuntimeError("parent manifest/checkpoint fingerprint mismatch") + if manifest["data_metadata"] != data_metadata: + raise RuntimeError("continuation data inventory differs from parent") + numerical_fields = ( + "accelerator", "torch_version", "torch_xla_version", "float32_matmul_precision", + "xla_matmul_precision", "xla_spmd", "xla_spmd_chips", "tpu_accelerator_type", + ) + for key in numerical_fields: + if manifest["runtime_environment"].get(key) != current_runtime.get(key): + raise RuntimeError(f"continuation numerical runtime changed: {key}") + loaded = load_training_checkpoint_for_resume( + path, model=model, handles=handles, + expected_fingerprint=lineage["parent_fingerprint"], train_generator=train_generator, + ) + if loaded[0] != lineage["parent_local_step"] or loaded[4] is None: + raise RuntimeError("parent restart state is incomplete") + if file_sha256(path) != lineage["parent_file_sha256"]: + raise RuntimeError("parent checkpoint changed while restoring state") + # This counter controls ONLY QK CSV interval boundaries. Momentum buffers + # and auxiliary Adam's per-parameter update counters remain untouched. + for handle in handles: + # The CLI may load MuonClip as __main__, so identity checks against a + # second imported module would silently miss the active optimizer. + reset = getattr(handle.optimizer, "reset_phase_diagnostics", None) + if callable(reset): + reset() + return loaded[4] diff --git a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/engine.py b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/engine.py index 63d7654c..c83ad8f0 100644 --- a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/engine.py +++ b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/engine.py @@ -48,6 +48,7 @@ ) from .train_loop import execute_training_loop from .tpu_spmd import initialize as initialize_spmd, replicate_model +from .continuation import import_parent_state, pause_reason, TrainingPaused def run_one( @@ -241,6 +242,15 @@ def run_one( f"cannot resume {run_dir}: checkpoint_latest.pt is missing" ) + if cfg.get("continuation") and resume_checkpoint is None: + resume_diagnostics = import_parent_state( + cfg, model=model, handles=handles, train_generator=train_generator, + data_metadata=data_metadata, seed=int(seed), current_runtime=current_runtime, + ) + model.to(resolved_device) + replicate_model(model) + synchronize(resolved_device) + write_manifest( run_dir, cfg=cfg, @@ -275,6 +285,7 @@ def run_one( optimizer_name=optimizer_name, seed=int(seed), train_generator=train_generator, + resume_diagnostics=resume_diagnostics, ) if progress: print( @@ -290,6 +301,10 @@ def run_one( flush=True, ) + reason = pause_reason(cfg, run_dir) + if reason: + raise TrainingPaused(reason) + metrics_path = run_dir / "metrics.csv" epoch_metrics_path = run_dir / "epoch_metrics.csv" prepare_csv( @@ -327,6 +342,7 @@ def run_one( arrays=arrays, train_probe=train_probe, val_probe=val_probe, + test_probe=test_probe, device=resolved_device, optimizer_name=optimizer_name, seed=int(seed), @@ -418,6 +434,13 @@ def run_one( "final": final_test, "validation_selected": best_test, } + if int(eval_cfg.get("test_interval_steps", 0)) > 0: + test_results["policy"] = ( + "fixed test probe used for monitoring; validation selects checkpoints; " + "test never selects checkpoints automatically; not an untouched held-out audit" + ) + test_results["test_interval_steps"] = int(eval_cfg["test_interval_steps"]) + test_results["probe_tokens"] = eval_batches * batch_size * block_size (run_dir / "test_results.json").write_text( json.dumps( test_results, @@ -432,6 +455,7 @@ def run_one( "optimizer": optimizer_name, "seed": int(seed), "optimizer_steps": int(total_steps), + "global_step": int(cfg.get("continuation", {}).get("global_step_offset", 0)) + int(total_steps), "train_epochs": float( total_steps * tokens_per_step(cfg) / train_tokens ), diff --git a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/monitor.py b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/monitor.py index 79ef9641..98c5aae5 100644 --- a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/monitor.py +++ b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/monitor.py @@ -1,6 +1,7 @@ from __future__ import annotations import argparse +import json from pathlib import Path import time @@ -72,6 +73,45 @@ def load_monitor_frames( return metrics, layers +def load_series_frames(root: str | Path) -> tuple[pd.DataFrame, pd.DataFrame]: + """Join retained CSV history even after older model files are pruned.""" + root = Path(root) + state = json.loads((root / "series.json").read_text()) + dirs = [Path(state["origin_checkpoint"]).parent] + dirs.extend(root / record["directory"] / "muon_clip" / f"seed_{state['seed']}" + for record in state["completed_segments"]) + if state.get("active_segment"): + dirs.append(root / state["active_segment"] / "muon_clip" / f"seed_{state['seed']}") + collected = [[], []] + for run_dir in dirs: + manifest_path = run_dir / "manifest.json" + if not manifest_path.is_file(): + continue + manifest = json.loads(manifest_path.read_text()) + offset = int((manifest.get("continuation") or {}).get("global_step_offset", 0)) + tokens = int(manifest["tokens_per_step"]) + train_tokens = int(manifest["data_metadata"]["splits"]["train"]) + for index, frame in enumerate(load_monitor_frames(run_dir)): + if frame.empty: + continue + frame = frame.copy() + frame["step"] = pd.to_numeric(frame["step"], errors="coerce") + offset + frame["epoch"] = frame["step"] * tokens / train_tokens + frame["segment_run_dir"] = str(run_dir) + if run_dir == dirs[0]: + frame = frame.loc[frame["step"] <= state["origin_global_step"]] + collected[index].append(frame) + frames = [] + for index, parts in enumerate(collected): + if not parts: + frames.append(pd.DataFrame()) + continue + frame = pd.concat(parts, ignore_index=True) + keys = ["step", "matrix_name"] if index else ["step"] + frames.append(frame.drop_duplicates(keys, keep="last").sort_values(keys)) + return tuple(frames) + + def _finite_summary(values: pd.Series) -> str: array = pd.to_numeric(values, errors="coerce").to_numpy(dtype=float) finite = array[np.isfinite(array)] @@ -130,6 +170,16 @@ def format_monitor_snapshot( ] ) + if not metrics.empty and "test_accuracy" in metrics.columns: + measured = metrics.loc[pd.to_numeric(metrics["test_accuracy"], errors="coerce").notna()] + if not measured.empty: + row = measured.iloc[-1] + lines.extend(["", f"LATEST FIXED TEST PROBE: step={int(row['step'])} " + f"test_loss={float(row['test_loss']):.4f} " + f"test_acc={100 * float(row['test_accuracy']):.2f}%", + "RECENT TEST PROBE HISTORY", + _format_table(measured[["step", "epoch", "test_loss", "test_accuracy", "val_accuracy"]].tail(recent))]) + if layers.empty: lines.extend(["", "Waiting for spectral/layers.csv..."]) return "\n".join(lines) @@ -240,6 +290,8 @@ def main() -> None: ) ) parser.add_argument("--run-dir") + parser.add_argument("--series-root", help="join all continuation segments on cumulative steps") + parser.add_argument("--export-metrics", type=Path, help="write the displayed metrics history as CSV") parser.add_argument("--results-root") parser.add_argument("--optimizer", default="muon") parser.add_argument("--seed", type=int, default=1337) @@ -259,7 +311,7 @@ def main() -> None: if args.recent < 1: parser.error("--recent must be positive") - run_dir = resolve_run_dir( + run_dir = Path(args.series_root) if args.series_root else resolve_run_dir( run_dir=args.run_dir, results_root=args.results_root, optimizer=args.optimizer, @@ -269,7 +321,12 @@ def main() -> None: try: while True: - metrics, layers = load_monitor_frames(run_dir) + metrics, layers = (load_series_frames(run_dir) if args.series_root else load_monitor_frames(run_dir)) + if args.export_metrics: + args.export_metrics.parent.mkdir(parents=True, exist_ok=True) + temporary = args.export_metrics.with_suffix(".csv.tmp") + metrics.to_csv(temporary, index=False) + temporary.replace(args.export_metrics) if not args.no_clear and not args.once: print("\033[2J\033[H", end="") print( diff --git a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/muonclip.py b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/muonclip.py index df9b44bb..c124ce61 100644 --- a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/muonclip.py +++ b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/muonclip.py @@ -372,6 +372,13 @@ def step(self, closure=None): self._flush_diagnostics() return loss + def reset_phase_diagnostics(self) -> None: + """Start a new logging interval without resetting optimization history.""" + self.step_index = 0 + self.last_diagnostics = {} + self._diagnostic_interval_state = None + self.reset_qk_tracking() + def state_dict(self) -> dict[str, Any]: payload = super().state_dict() payload["muonclip_global_state"] = { diff --git a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/muonclip_continue.py b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/muonclip_continue.py new file mode 100644 index 00000000..ef995da0 --- /dev/null +++ b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/muonclip_continue.py @@ -0,0 +1,253 @@ +"""Checkpoint-based finite or open-ended MuonClip training series. + +The supervisor stays on CPU. Every segment uses a fresh accelerator worker; +only full-state checkpoints connect segments. No cloud resources are created. +""" +from __future__ import annotations + +import argparse +from contextlib import contextmanager +import fcntl +import json +import math +import os +from pathlib import Path +import signal +import subprocess +import sys + +import yaml + +from .completion import validate_completed_run +from .continuation import make_continuation_config, pause_reason +from .muonclip_resilient import _atomic_json, _utc_now, run_resilient +from .run_utils import run_directory +from .provenance import source_fingerprint_payload, scientific_dependency_versions + + +def series_environment() -> dict: + return {"source": source_fingerprint_payload(), "dependencies": scientific_dependency_versions()} + + +@contextmanager +def series_lock(root: Path): + with (root / "driver.lock").open("a") as handle: + try: + fcntl.flock(handle, fcntl.LOCK_EX | fcntl.LOCK_NB) + except BlockingIOError as exc: + raise RuntimeError(f"a continuation driver or worker already holds {root}/driver.lock") from exc + # Keep the lock alive in the accelerator worker even if this supervisor + # dies, preventing a second launcher from writing the same checkpoint. + yield handle.fileno() + + +def prune_completed_segments(root: Path, state: dict) -> None: + """Delete only named checkpoint files in older, validated, owned segments.""" + keep = int(state["keep_segments"]) + if keep < 2: + raise ValueError("retain at least two completed segments") + for record in state["completed_segments"][:-keep]: + segment = (root / record["directory"]).resolve() + if not segment.is_relative_to(root.resolve()) or segment.parent != (root / "segments").resolve() or not segment.name.startswith("segment_"): + raise RuntimeError("refusing to prune outside owned segment directories") + run_dir = run_directory(segment, "muon_clip", state["seed"]) + if run_dir.resolve() != segment / "muon_clip" / f"seed_{state['seed']}": + raise RuntimeError("refusing to prune a symlinked run") + complete = json.loads((run_dir / "run_complete.json").read_text()) + if complete["fingerprint"] != record["fingerprint"]: + raise RuntimeError("refusing to prune a segment with changed completion identity") + paths = [run_dir / f"checkpoint_{name}.pt" for name in ("initial", "latest", "best", "final")] + paths.extend((run_dir / "epoch_checkpoints").glob("model_epoch_*.pt")) + deleted = [] + for path in paths: + if path.is_symlink() or not path.resolve().is_relative_to(run_dir): + raise RuntimeError("refusing to prune a checkpoint symlink") + if path.is_file(): + path.unlink() + deleted.append(str(path.relative_to(run_dir))) + if deleted: + _atomic_json(run_dir / "checkpoints_pruned.json", { + "pruned_at_utc": _utc_now(), "files": deleted, + "policy": "metrics, spectra, config, lineage and test results retained; full checkpoint validation occurred before archival", + }) + + +def _save_state(root: Path, state: dict) -> None: + state["updated_at_utc"] = _utc_now() + _atomic_json(root / "series.json", state) + + +def drive_series(root: Path, state: dict, lock_fd: int) -> int: + if state["environment"] != series_environment(): + raise RuntimeError("series source/dependencies changed; restore the pinned environment or explicitly start a new series") + stop_file = root / "STOP" + def request_stop(signum, frame): + del signum, frame + stop_file.touch() + print("[one-head-series] stop requested; waiting for the next saved checkpoint", flush=True) + + old_handlers = {s: signal.signal(s, request_stop) for s in (signal.SIGTERM, signal.SIGINT)} + try: + state["pid"] = os.getpid() + state["status"] = "running" + _save_state(root, state) + while state["additional_steps"] is None or state["completed_steps"] < state["additional_steps"]: + if state["environment"] != series_environment(): + raise RuntimeError("series source/dependencies changed between segments") + if stop_file.exists(): + state["status"] = "paused" + _save_state(root, state) + return 75 + index = len(state["completed_segments"]) + 1 + relative = f"segments/segment_{index:06d}" + segment = root / relative + config_path = segment / "config.yaml" + if not state.get("active_segment"): + steps = state["segment_steps"] + if state["additional_steps"] is not None: + steps = min(steps, state["additional_steps"] - state["completed_steps"]) + parent = Path(state["latest_checkpoint"]) + cfg = make_continuation_config( + parent, steps=steps, learning_rate=state["learning_rate"], + test_interval=state["test_interval_steps"], stop_file=stop_file, + min_free_disk_gb=state["min_free_disk_gb"], + ) + segment.mkdir(parents=True, exist_ok=True) + temporary = config_path.with_suffix(".yaml.tmp") + temporary.write_text(yaml.safe_dump(cfg, sort_keys=False)) + temporary.replace(config_path) + state["active_segment"] = relative + _save_state(root, state) + elif state["active_segment"] != relative: + raise RuntimeError("active segment does not match series history") + cfg = yaml.safe_load(config_path.read_text()) + reason = pause_reason(cfg, root) + if reason: + print(f"[one-head-series] {reason}", flush=True) + state["status"] = "paused" + state["pause_reason"] = reason + _save_state(root, state) + return 75 + print(f"[one-head-series] segment={index} global_start={cfg['continuation']['global_step_offset']} steps={cfg['training']['max_steps']}", flush=True) + args = argparse.Namespace( + config=str(config_path), seed=state["seed"], data_root=state["data_root"], + results_root=str(segment), device=state["device"], + max_no_progress_failures=state["max_no_progress_failures"], + retry_delay_seconds=5.0, lock_fd=lock_fd, + worker_log=str(segment / "launch.log"), + ) + code = run_resilient(args) + if code: + state["status"] = "paused" if code == 75 else "failed" + state["last_exit_code"] = code + _save_state(root, state) + return code + run_dir = run_directory(segment, "muon_clip", state["seed"]) + # Independently audit artifacts before using a segment as a parent + # or considering any older checkpoint for retention pruning. + validate_completed_run(run_dir, verify_checkpoints=True) + complete = json.loads((run_dir / "run_complete.json").read_text()) + state["completed_segments"].append({ + "directory": relative, "fingerprint": complete["fingerprint"], + "steps": complete["optimizer_steps"], "global_step": complete["global_step"], + }) + state["completed_steps"] += complete["optimizer_steps"] + state["latest_checkpoint"] = str(run_dir / "checkpoint_final.pt") + state["active_segment"] = None + _save_state(root, state) + prune_completed_segments(root, state) + state["status"] = "completed" + _save_state(root, state) + return 0 + except Exception as exc: + state["status"] = "failed" + state["error"] = f"{type(exc).__name__}: {exc}" + _save_state(root, state) + raise + finally: + for sig, handler in old_handlers.items(): + signal.signal(sig, handler) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("action", choices=("start", "resume", "stop", "status")) + parser.add_argument("--series-root", type=Path, required=True) + parser.add_argument("--from-checkpoint", type=Path) + parser.add_argument("--data-root", type=Path) + parser.add_argument("--device", choices=("tpu", "xla", "cpu", "cuda", "mps"), default="tpu") + duration = parser.add_mutually_exclusive_group() + duration.add_argument("--forever", action="store_true") + duration.add_argument("--additional-steps", type=int) + parser.add_argument("--segment-steps", type=int, default=1000000) + parser.add_argument("--learning-rate", type=float) + parser.add_argument("--test-interval-steps", type=int, default=10000) + parser.add_argument("--keep-segments", type=int, default=3) + parser.add_argument("--min-free-disk-gb", type=float, default=5.0) + parser.add_argument("--max-no-progress-failures", type=int, default=3) + parser.add_argument("--background", action="store_true") + args = parser.parse_args() + root = args.series_root.expanduser().resolve() + if args.action in ("stop", "status", "resume") and not (root / "series.json").is_file(): + parser.error("series.json is missing; choose an existing series") + if args.action == "stop": + (root / "STOP").touch() + print("Stop requested. The worker will pause after its next atomic checkpoint.") + return + if args.action == "status": + print((root / "series.json").read_text()) + print("Recorded status may be stale after host loss; inspect the checkpoint and worker log before resuming.") + return + if args.action == "start": + if not args.from_checkpoint or not args.data_root or not (args.forever or args.additional_steps): + parser.error("start requires --from-checkpoint, --data-root and --forever or --additional-steps") + if args.segment_steps < 1 or args.test_interval_steps < 1 or args.keep_segments < 2 or args.max_no_progress_failures < 1: + parser.error("steps/interval/retry budget must be positive; keep at least two segments") + if args.additional_steps is not None and args.additional_steps < 1: + parser.error("--additional-steps must be positive") + if not math.isfinite(args.min_free_disk_gb) or args.min_free_disk_gb < 0: + parser.error("--min-free-disk-gb must be finite and nonnegative") + if root.exists() and any(path.name != "driver.lock" for path in root.iterdir()): + parser.error("start requires an empty series directory; use resume for an existing series") + root.mkdir(parents=True, exist_ok=True) + if args.background: + # Launch the same pinned interpreter and arguments; only the foreground + # child creates state and acquires the writer lock. + command = [sys.executable, "-u", "-m", "rg_nanogpt_one_head.muonclip_continue", + *[x for x in sys.argv[1:] if x != "--background"]] + # Keep this log outside a new root, which must remain empty for start. + log = root.parent / f"{root.name}.driver.log" + with log.open("a") as output: + child = subprocess.Popen(command, stdin=subprocess.DEVNULL, stdout=output, stderr=subprocess.STDOUT, start_new_session=True) + print(f"Started driver pid={child.pid}; log={log}") + return + from .muonclip import install_muonclip_extension + install_muonclip_extension() + with series_lock(root) as lock_fd: + if args.action == "start": + cfg = make_continuation_config( + args.from_checkpoint, steps=args.segment_steps, learning_rate=args.learning_rate, + test_interval=args.test_interval_steps, stop_file=root / "STOP", min_free_disk_gb=args.min_free_disk_gb, + ) + state = { + "schema_version": 1, "created_at_utc": _utc_now(), "seed": cfg["continuation"]["seed"], + "origin_checkpoint": str(args.from_checkpoint.expanduser().resolve()), + "origin_global_step": cfg["continuation"]["global_step_offset"], + "latest_checkpoint": str(args.from_checkpoint.expanduser().resolve()), + "data_root": str(args.data_root.expanduser().resolve()), "device": args.device, + "additional_steps": args.additional_steps, "segment_steps": args.segment_steps, + "learning_rate": cfg["optimizer_profiles"]["muon_clip"]["learning_rate"], + "test_interval_steps": args.test_interval_steps, "keep_segments": args.keep_segments, + "min_free_disk_gb": args.min_free_disk_gb, "max_no_progress_failures": args.max_no_progress_failures, + "completed_steps": 0, "completed_segments": [], "active_segment": None, + "environment": series_environment(), + } + else: + state = json.loads((root / "series.json").read_text()) + (root / "STOP").unlink(missing_ok=True) + code = drive_series(root, state, lock_fd) + raise SystemExit(code) + + +if __name__ == "__main__": + main() diff --git a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/muonclip_resilient.py b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/muonclip_resilient.py index 32fe07b5..db2b5cbf 100644 --- a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/muonclip_resilient.py +++ b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/muonclip_resilient.py @@ -13,6 +13,7 @@ """ import argparse +from contextlib import nullcontext from datetime import datetime, timezone import json import os @@ -159,13 +160,26 @@ def run_resilient(args: argparse.Namespace) -> int: }, ) - result = subprocess.run( - command, - env=environment, - check=False, - ) + options = {} + if getattr(args, "lock_fd", None) is not None: + options.update(pass_fds=(args.lock_fd,), start_new_session=True) + log_path = getattr(args, "worker_log", None) + with (Path(log_path).open("a") if log_path else nullcontext()) as output: + if output is not None: + options.update(stdout=output, stderr=subprocess.STDOUT) + result = subprocess.run(command, env=environment, check=False, **options) return_code = int(result.returncode) + if return_code == 75: + _atomic_json(status_path, { + "completed": False, "running": False, "paused": True, + "last_exit_code": 75, "run_dir": str(run_dir), + "checkpoint_path": str(latest_checkpoint), + "last_verified_checkpoint_step": _checkpoint_step(latest_checkpoint), + "updated_at_utc": _utc_now(), + }) + return 75 + if return_code == 0: final_step = _checkpoint_step(latest_checkpoint) _atomic_json( diff --git a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/run_utils.py b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/run_utils.py index 9f367dba..09337bea 100644 --- a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/run_utils.py +++ b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/run_utils.py @@ -27,6 +27,7 @@ from .runtime import runtime_metadata METRIC_FIELDS = [ + "global_step", "global_epoch", "step", "tokens_seen", "epoch", "elapsed_sec", "tokens_per_sec", "primary_lr", "auxiliary_lr", "train_loss", "train_perplexity", "train_bits_per_token", "train_accuracy", "train_top5_accuracy", @@ -400,6 +401,7 @@ def write_manifest( "data_metadata": data_metadata, "training": cfg["training"], "evaluation": cfg["evaluation"], + "continuation": cfg.get("continuation"), "weightwatcher": cfg["weightwatcher"], "tokens_per_step": tokens_per_step(cfg), "max_steps": int(total_steps), @@ -418,6 +420,12 @@ def write_manifest( "not translation BLEU" ), } + if int(cfg["evaluation"].get("test_interval_steps", 0)) > 0: + payload["test_policy"] = ( + "fixed test probe used for monitoring; validation selects the best " + "checkpoint; test never selects checkpoints automatically; this is " + "not an untouched held-out evaluation after human monitoring" + ) temporary = run_dir / "manifest.json.tmp" temporary.write_text( json.dumps( diff --git a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/train_loop.py b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/train_loop.py index 83abc122..44271e44 100644 --- a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/train_loop.py +++ b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/train_loop.py @@ -26,6 +26,12 @@ from .spectral import run_weightwatcher from .random_canaries import RandomCanaryExperiment from .tpu_spmd import batch_to_device, replicate_gradients +from .continuation import pause_reason, TrainingPaused + + +def _test_due(step: int, cfg: dict, total_steps: int) -> bool: + interval = int(cfg["evaluation"].get("test_interval_steps", 0)) + return interval > 0 and (step % interval == 0 or step == total_steps) def _require_finite_metrics( @@ -96,6 +102,7 @@ def _evaluation_due( step % int(cfg["training"]["eval_interval_steps"]) == 0 or step in epoch_steps or step == total_steps + or _test_due(step, cfg, total_steps) ) @@ -182,6 +189,7 @@ def execute_training_loop( latest_checkpoint: Path, best_checkpoint: Path, progress: bool, + test_probe=None, ) -> tuple[float, int, float, dict]: batch_size = int(cfg["training"]["batch_size"]) grad_accum = int(cfg["training"]["grad_accum_steps"]) @@ -229,7 +237,7 @@ def execute_training_loop( else float("nan") ), } - if start_step > 0: + if start_step > 0 or cfg.get("continuation"): for handle in handles: last_update_lrs[handle.role] = float(handle.lr) @@ -303,9 +311,14 @@ def execute_training_loop( "continuation_token_accuracy": float("nan"), "continuation_exact_match": float("nan"), } - # Keep the test split genuinely held out during optimization. - # Final and validation-selected checkpoints are evaluated once, - # after training, by engine.checkpoint_eval. + # Periodic test monitoring is opt-in. Historical baseline configs + # keep these placeholders empty until the post-training audit. + if _test_due(completed_steps, cfg, total_steps): + if test_probe is None: + raise RuntimeError("periodic test monitoring requires a fixed test probe") + test_metrics = evaluate_probe(model, test_probe, device) + if not all(math.isfinite(float(value)) for value in test_metrics.values()): + raise FloatingPointError("non-finite test probe metrics") tokens_seen = int(completed_steps * step_tokens) actual_epoch = tokens_seen / max(1, train_tokens) @@ -318,6 +331,8 @@ def execute_training_loop( weight_norm = model_weight_norm(model) current_mps, driver_mps = mps_memory_megabytes(device) row = { + "global_step": int(cfg.get("continuation", {}).get("global_step_offset", 0)) + completed_steps, + "global_epoch": (int(cfg.get("continuation", {}).get("global_step_offset", 0)) + completed_steps) * step_tokens / train_tokens, "step": int(completed_steps), "tokens_seen": tokens_seen, "epoch": float(actual_epoch), @@ -438,7 +453,7 @@ def execute_training_loop( "nominal_epoch": nominal_epoch, "checkpoint_path": str(checkpoint_path), "test_monitoring_only": 1, - "test_held_out": 1, + "test_held_out": int(not cfg["evaluation"].get("test_interval_steps", 0)), } ) epoch_handle.flush() @@ -485,6 +500,10 @@ def execute_training_loop( if not math.isfinite(eta) else f"{eta / 60:.1f}m" ) + test_text = ( + f"test_acc={100 * test_metrics['accuracy']:.2f}% " + if math.isfinite(test_metrics["accuracy"]) else "" + ) print( "[one-head-train] " f"optimizer={optimizer_name} seed={seed} " @@ -496,6 +515,7 @@ def execute_training_loop( f"val_loss={val_metrics['loss']:.4f} " f"val_ppl={val_metrics['perplexity']:.2f} " f"val_acc={100 * val_metrics['accuracy']:.2f}% " + f"{test_text}" f"eta={eta_text}", flush=True, ) @@ -597,6 +617,9 @@ def execute_training_loop( last_clipped=last_clipped, ), ) + reason = pause_reason(cfg, run_dir) + if reason: + raise TrainingPaused(reason) synchronize(device) if final_resume_diagnostics is None: diff --git a/baseline/nanogpt_one_head/tests/continuation_scenario.py b/baseline/nanogpt_one_head/tests/continuation_scenario.py new file mode 100644 index 00000000..35763ecb --- /dev/null +++ b/baseline/nanogpt_one_head/tests/continuation_scenario.py @@ -0,0 +1,134 @@ +"""Real CPU training and fresh-process continuation exercise, invoked by pytest.""" +from copy import deepcopy +import hashlib +import json +from pathlib import Path +import subprocess +import sys + +import numpy as np +import pandas as pd +import torch + +from rg_nanogpt_one_head.muonclip import install_muonclip_extension +install_muonclip_extension() +from rg_nanogpt_one_head.config import load_config +from rg_nanogpt_one_head.continuation import make_continuation_config, TrainingPaused, file_sha256 +from rg_nanogpt_one_head.checkpoints import optimizer_state_sha256 +from rg_nanogpt_one_head.completion import validate_completed_run +from rg_nanogpt_one_head.monitor import load_series_frames +from rg_nanogpt_one_head.training import run_one + +root = Path(sys.argv[1]) +cfg = deepcopy(load_config("configs/muonclip_reference.yaml")) +cfg["dataset"].update(name="unit/fineweb", config="unit", revision="unit-revision", train_tokens=2048, val_tokens=512, test_tokens=512) +cfg["model"].update(vocab_size=64, block_size=8, n_layer=1, n_head=1, n_embd=32, dropout=0.0) +cfg["training"].update(seeds=[13], batch_size=2, grad_accum_steps=2, target_epochs=0.125, max_steps=8, epoch_interval=1.0, eval_interval_steps=2, eval_batches=1, checkpoint_interval_steps=2) +cfg["evaluation"].update(bleu_examples=2, bleu_prompt_tokens=3, bleu_continuation_tokens=2, bleu_batch_size=2, test_interval_steps=2) +cfg["weightwatcher"].update(min_evals=5, fix_fingers=False) +for p in cfg["optimizer_profiles"].values(): + p.pop("lr_schedule_epochs", None) +cfg["optimizer_profiles"]["muon_clip"].update(learning_rate=2e-5, min_learning_rate=2e-5, warmup_fraction=0.0, qk_diagnostics_interval=2) +data = root / "data" +data.mkdir(parents=True) +rng = np.random.default_rng(7) +splits = {k: cfg["dataset"][k + "_tokens"] for k in ("train", "val", "test")} +files = {} +for split, size in splits.items(): + path = data / f"{split}.bin" + rng.integers(0, 64, size=size, dtype=np.uint16).tofile(path) + files[split] = {"path": path.name, "sha256": file_sha256(path), "bytes": path.stat().st_size} +(data / "meta.json").write_text(json.dumps({ + "schema_version": 2, "tokenizer": "gpt2", "vocab_size": 64, "dtype": "uint16", "splits": splits, + "document_disjoint_splits": True, "dataset_name": "unit/fineweb", "dataset_config": "unit", + "dataset_split": "train", "dataset_revision": "unit-revision", "eot_token": 0, "files": files, +})) + +def run(configuration, name): + return run_one(cfg=configuration, data_root=data, results_root=root / name, optimizer_name="muon_clip", seed=13, device="cpu", progress=False) + +def load(path): + return torch.load(path, map_location="cpu", weights_only=False) + +def same_training_state(a, b): + assert a["model_state_sha256"] == b["model_state_sha256"] + assert torch.equal(a["train_generator_state"], b["train_generator_state"]) + assert torch.equal(a["torch_random_state"], b["torch_random_state"]) + aa, bb = deepcopy(a["optimizers"]), deepcopy(b["optimizers"]) + # Phase-local QK statistics have no effect on the update rule. + aa[0].pop("muonclip_global_state") + bb[0].pop("muonclip_global_state") + assert optimizer_state_sha256(aa) == optimizer_state_sha256(bb) + +whole = run(cfg, "whole") +parent_cfg = deepcopy(cfg) +parent_cfg["training"].update(max_steps=4, target_epochs=0.0625) +parent = run(parent_cfg, "parent") +parent_path = parent / "checkpoint_final.pt" +parent_hash = file_sha256(parent_path) +extension = make_continuation_config(parent_path, steps=4, learning_rate=None, test_interval=2, stop_file=root / "STOP", min_free_disk_gb=0) +child = run(extension, "child") +validate_completed_run(child) +same_training_state(load(whole / "checkpoint_final.pt"), load(child / "checkpoint_final.pt")) +same_training_state(load(parent_path), load(child / "checkpoint_initial.pt")) +assert load(child / "checkpoint_final.pt")["global_step"] == 8 +frame = pd.read_csv(child / "metrics.csv") +assert frame.global_step.tolist() == [4, 6, 8] +assert frame.test_accuracy.notna().all() +assert frame.loc[frame.step > 0, "primary_lr"].eq(2e-5).all() + +# Pause after a saved update; resume must reproduce exactly the uninterrupted state. +import rg_nanogpt_one_head.train_loop as loop +original_save = loop.save_training_checkpoint +def pause_after_checkpoint(path, **kwargs): + result = original_save(path, **kwargs) + if Path(path).name == "checkpoint_latest.pt" and kwargs["step"] == 2: + (root / "STOP").touch() + return result +loop.save_training_checkpoint = pause_after_checkpoint +try: + run(extension, "paused") + raise AssertionError("pause was not honored") +except TrainingPaused as exc: + assert exc.code == 75 +finally: + loop.save_training_checkpoint = original_save +(root / "STOP").unlink() +resumed = run(extension, "paused") +validate_completed_run(resumed) +same_training_state(load(child / "checkpoint_final.pt"), load(resumed / "checkpoint_final.pt")) + +# Mutating a protected parent-derived field must be rejected before child training. +bad = deepcopy(extension) +bad["optimizer_profiles"]["muon_clip"]["momentum"] = 0.5 +try: + run(bad, "bad") + raise AssertionError("optimizer mutation accepted") +except RuntimeError as exc: + assert "cannot change optimizer" in str(exc) +assert file_sha256(parent_path) == parent_hash + +# Exercise actual CLI -> CPU supervisor -> fresh worker across THREE segments, +# strict completion audit, bounded checkpoint retention, and joined monitoring. +series = root / "series" +command = [sys.executable, "-m", "rg_nanogpt_one_head.muonclip_continue", "start", + "--series-root", str(series), "--from-checkpoint", str(parent_path), "--data-root", str(data), + "--device", "cpu", "--additional-steps", "6", "--segment-steps", "2", "--test-interval-steps", "2", + "--keep-segments", "2", "--min-free-disk-gb", "0"] +subprocess.run(command, check=True) +state = json.loads((series / "series.json").read_text()) +assert state["status"] == "completed" and state["completed_steps"] == 6 +assert len(state["completed_segments"]) == 3 +first = series / "segments/segment_000001/muon_clip/seed_13" +assert not (first / "checkpoint_final.pt").exists() +assert (first / "checkpoints_pruned.json").is_file() +assert (first / "metrics.csv").is_file() +assert Path(state["latest_checkpoint"]).is_file() +assert load(Path(state["latest_checkpoint"]))["global_step"] == 10 +metrics, layers = load_series_frames(series) +assert metrics.step.tolist() == [0, 2, 4, 6, 8, 10] +assert layers.step.max() == 10 +assert file_sha256(parent_path) == parent_hash +# A completed finite series resumes idempotently even with pruned older phases. +subprocess.run([sys.executable, "-m", "rg_nanogpt_one_head.muonclip_continue", "resume", "--series-root", str(series)], check=True) +print("CONTINUATION_SCENARIO_PASSED") diff --git a/baseline/nanogpt_one_head/tests/test_continuation.py b/baseline/nanogpt_one_head/tests/test_continuation.py new file mode 100644 index 00000000..fcdeb54a --- /dev/null +++ b/baseline/nanogpt_one_head/tests/test_continuation.py @@ -0,0 +1,92 @@ +from pathlib import Path +import subprocess +import sys +from types import SimpleNamespace + +import pandas as pd +import pytest + +from rg_nanogpt_one_head.continuation import pause_reason +from rg_nanogpt_one_head.completion import _validate_test_monitoring, CompletedRunValidationError +from rg_nanogpt_one_head.muonclip_continue import series_lock, prune_completed_segments +import rg_nanogpt_one_head.muonclip_continue as series +import rg_nanogpt_one_head.muonclip_resilient as resilient + + +def test_real_continuation_resume_and_segment_retention(tmp_path): + result = subprocess.run([sys.executable, "tests/continuation_scenario.py", str(tmp_path)], + cwd=Path(__file__).resolve().parents[1], capture_output=True, text=True, timeout=240) + assert result.returncode == 0, result.stdout + "\n" + result.stderr + assert "CONTINUATION_SCENARIO_PASSED" in result.stdout + + +def test_supervisor_honors_pause_without_retry(tmp_path, monkeypatch): + args = SimpleNamespace(config=tmp_path / "cfg.yaml", seed=13, data_root=tmp_path / "data", + results_root=tmp_path / "results", device="cpu", max_no_progress_failures=3, retry_delay_seconds=0) + calls = [] + def pause(*args, **kwargs): + calls.append(1) + return SimpleNamespace(returncode=75) + monkeypatch.setattr(resilient.subprocess, "run", pause) + assert resilient.run_resilient(args) == 75 + assert len(calls) == 1 + + +def test_duplicate_writer_lock(tmp_path): + with series_lock(tmp_path): + with pytest.raises(RuntimeError, match="already holds"): + with series_lock(tmp_path): + pass + + +def test_low_disk_pause(tmp_path, monkeypatch): + monkeypatch.setattr("rg_nanogpt_one_head.continuation.shutil.disk_usage", lambda _: SimpleNamespace(free=100)) + assert "free disk" in pause_reason({"training": {"min_free_disk_gb": 1}}, tmp_path) + + +def test_retention_rejects_outside_series(tmp_path): + state = {"seed": 13, "keep_segments": 2, "completed_segments": [{"directory": "../outside"}] * 3} + with pytest.raises(RuntimeError, match="outside owned"): + prune_completed_segments(tmp_path, state) + + +def test_periodic_probe_validator_rejects_missing_measurement(): + frame = pd.DataFrame({"step": [0, 4]}) + with pytest.raises(CompletedRunValidationError, match="missing periodic"): + _validate_test_monitoring(frame, "metrics.csv", 2, 4) + + +def test_open_ended_series_stops_at_saved_boundary(tmp_path, monkeypatch): + import json + monkeypatch.setattr(series, "series_environment", lambda: {"source": "fixed"}) + monkeypatch.setattr(series, "make_continuation_config", lambda *args, **kwargs: { + "continuation": {"global_step_offset": 123}, "training": {"max_steps": 5}, + }) + monkeypatch.setattr(series, "pause_reason", lambda *args: None) + monkeypatch.setattr(series, "validate_completed_run", lambda *args, **kwargs: None) + monkeypatch.setattr(series, "prune_completed_segments", lambda *args: None) + calls = [] + def worker(args): + calls.append(1) + run_dir = Path(args.results_root) / "muon_clip/seed_13" + run_dir.mkdir(parents=True) + (run_dir / "run_complete.json").write_text(json.dumps({ + "fingerprint": "verified", "optimizer_steps": 5, "global_step": 128, + })) + (tmp_path / "STOP").touch() + return 0 + monkeypatch.setattr(series, "run_resilient", worker) + state = {"environment": {"source": "fixed"}, "additional_steps": None, + "completed_steps": 0, "completed_segments": [], "segment_steps": 5, + "latest_checkpoint": "parent.pt", "learning_rate": 2e-5, "test_interval_steps": 2, + "min_free_disk_gb": 0, "seed": 13, "data_root": str(tmp_path), "device": "cpu", + "max_no_progress_failures": 3} + assert series.drive_series(tmp_path, state, 1) == 75 + assert len(calls) == 1 and state["completed_steps"] == 5 + assert state["status"] == "paused" and state["active_segment"] is None + + +def test_series_environment_change_is_rejected(tmp_path, monkeypatch): + monkeypatch.setattr(series, "series_environment", lambda: {"source": "changed"}) + with pytest.raises(RuntimeError, match="source/dependencies changed"): + series.drive_series(tmp_path, {"environment": {"source": "original"}}, 1) From 7e811f17ee9b41b174bf9b8938646cffedfd0903 Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Thu, 1 Oct 2026 10:53:35 -0700 Subject: [PATCH 04/36] Fix short clip-Xmax continuation completion validation Require the exact configured snapshot grid for continuation segments while preserving the historical campaign minimum. Exercise real raw/clipped-alpha training, pause/resume and fresh-worker segments in the integration test. Validation: 53 continuation, completion and clip-Xmax campaign tests passed. --- baseline/nanogpt_one_head/TPU_CONTINUATION.md | 11 ++++++ .../src/rg_nanogpt_one_head/completion.py | 25 +++++++++++-- .../tests/continuation_scenario.py | 8 +++-- .../tests/test_continuation.py | 36 ++++++++++++++++++- 4 files changed, 75 insertions(+), 5 deletions(-) diff --git a/baseline/nanogpt_one_head/TPU_CONTINUATION.md b/baseline/nanogpt_one_head/TPU_CONTINUATION.md index 14875353..a1e269e3 100644 --- a/baseline/nanogpt_one_head/TPU_CONTINUATION.md +++ b/baseline/nanogpt_one_head/TPU_CONTINUATION.md @@ -102,6 +102,17 @@ Its worker log is in `segments/segment_000001/launch.log` under the smoke root. If it fails, inspect that log before launching the long series. The original checkpoint is unchanged by the smoke test. +Short continuation segments are audited against their exact configured snapshot +grid (initial and final for this smoke), not the historical comparison campaign's +ten-snapshot minimum. Missing scheduled snapshots still fail validation. + +If the first release rejected this smoke with `fewer than ten permanent states`, +update the continuation checkout after the worker exits, then rerun the smoke +using a **new** root such as `/mnt/disks/rg-data/muonclip-continuation-smoke-v2`. +Keep the old smoke as evidence. Its source identity is pinned, so do not edit +its manifest or use `resume` to bypass the source check. Both smokes start from +the original `muonclip-spmd-long` final checkpoint; those weights are unchanged. + ## Start, monitor, pause, and resume Start the extended series from the original final checkpoint: diff --git a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/completion.py b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/completion.py index 7be3a68e..37244e3b 100644 --- a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/completion.py +++ b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/completion.py @@ -17,6 +17,7 @@ optimizer_state_sha256, require_finite_checkpoint_state, ) +from .config import epoch_step_map _REQUIRED_FILES = ( "run_complete.json", @@ -60,6 +61,27 @@ class CompletedRunValidationError(RuntimeError): """A nominally completed run is missing, stale, or inconsistent.""" +def _validate_clip_snapshot_schedule(manifest: dict, observed_steps: tuple[int, ...]) -> None: + if not manifest.get("continuation"): + # Historical comparison campaigns require this minimum sampling density. + if len(observed_steps) < 10: + _fail("clip-Xmax campaign has fewer than ten permanent states") + return + # A continuation may be a short acceptance run or a final partial segment. + # Require its entire configured grid instead of a campaign-wide minimum. + try: + train_tokens = int(manifest["data_metadata"]["splits"]["train"]) + cfg = { + "training": manifest["training"], "model": manifest["model"], + "dataset": {"train_tokens": train_tokens}, + } + expected_steps = tuple(epoch_step_map(cfg, train_tokens)) + except (KeyError, TypeError, ValueError, OverflowError) as exc: + _fail(f"continuation snapshot schedule cannot be reconstructed: {exc}") + if observed_steps != expected_steps: + _fail("continuation permanent states do not cover the exact configured snapshot grid") + + def _validate_test_monitoring(frame, label, interval, total_steps): steps = pd.to_numeric(frame["step"], errors="raise") due = steps.mod(interval).eq(0) | steps.eq(total_steps) @@ -777,8 +799,7 @@ def validate_completed_run( "clip_xmax" ).all(): _fail("spectral rows do not declare clipped alpha as primary") - if len(epoch_steps) < 10: - _fail("clip-Xmax campaign has fewer than ten permanent states") + _validate_clip_snapshot_schedule(manifest, epoch_steps) for column in ( "alpha_raw_n", "alpha_raw_median", diff --git a/baseline/nanogpt_one_head/tests/continuation_scenario.py b/baseline/nanogpt_one_head/tests/continuation_scenario.py index 35763ecb..01a9b9f3 100644 --- a/baseline/nanogpt_one_head/tests/continuation_scenario.py +++ b/baseline/nanogpt_one_head/tests/continuation_scenario.py @@ -25,10 +25,10 @@ cfg["model"].update(vocab_size=64, block_size=8, n_layer=1, n_head=1, n_embd=32, dropout=0.0) cfg["training"].update(seeds=[13], batch_size=2, grad_accum_steps=2, target_epochs=0.125, max_steps=8, epoch_interval=1.0, eval_interval_steps=2, eval_batches=1, checkpoint_interval_steps=2) cfg["evaluation"].update(bleu_examples=2, bleu_prompt_tokens=3, bleu_continuation_tokens=2, bleu_batch_size=2, test_interval_steps=2) -cfg["weightwatcher"].update(min_evals=5, fix_fingers=False) +cfg["weightwatcher"].update(min_evals=5, fix_fingers="clip_xmax", max_fingers=10, require_raw_alpha=True) for p in cfg["optimizer_profiles"].values(): p.pop("lr_schedule_epochs", None) -cfg["optimizer_profiles"]["muon_clip"].update(learning_rate=2e-5, min_learning_rate=2e-5, warmup_fraction=0.0, qk_diagnostics_interval=2) +cfg["optimizer_profiles"]["muon_clip"].update(learning_rate=2e-5, min_learning_rate=2e-5, warmup_fraction=0.0, qk_diagnostics_interval=1000) data = root / "data" data.mkdir(parents=True) rng = np.random.default_rng(7) @@ -69,6 +69,10 @@ def same_training_state(a, b): extension = make_continuation_config(parent_path, steps=4, learning_rate=None, test_interval=2, stop_file=root / "STOP", min_free_disk_gb=0) child = run(extension, "child") validate_completed_run(child) +assert len(pd.read_csv(child / "epoch_metrics.csv")) == 2 +spectra = pd.read_csv(child / "spectral/layers.csv") +assert spectra["finger_policy"].eq("fix_fingers=clip_xmax").all() +assert np.isfinite(spectra[["alpha_raw", "alpha_clip_xmax"]].to_numpy()).all() same_training_state(load(whole / "checkpoint_final.pt"), load(child / "checkpoint_final.pt")) same_training_state(load(parent_path), load(child / "checkpoint_initial.pt")) assert load(child / "checkpoint_final.pt")["global_step"] == 8 diff --git a/baseline/nanogpt_one_head/tests/test_continuation.py b/baseline/nanogpt_one_head/tests/test_continuation.py index fcdeb54a..296ab531 100644 --- a/baseline/nanogpt_one_head/tests/test_continuation.py +++ b/baseline/nanogpt_one_head/tests/test_continuation.py @@ -7,7 +7,9 @@ import pytest from rg_nanogpt_one_head.continuation import pause_reason -from rg_nanogpt_one_head.completion import _validate_test_monitoring, CompletedRunValidationError +from rg_nanogpt_one_head.completion import ( + _validate_test_monitoring, _validate_clip_snapshot_schedule, CompletedRunValidationError, +) from rg_nanogpt_one_head.muonclip_continue import series_lock, prune_completed_segments import rg_nanogpt_one_head.muonclip_continue as series import rg_nanogpt_one_head.muonclip_resilient as resilient @@ -90,3 +92,35 @@ def test_series_environment_change_is_rejected(tmp_path, monkeypatch): monkeypatch.setattr(series, "series_environment", lambda: {"source": "changed"}) with pytest.raises(RuntimeError, match="source/dependencies changed"): series.drive_series(tmp_path, {"environment": {"source": "original"}}, 1) + + +def _snapshot_manifest(steps): + # The production TPU shape and cadence, including a short final segment. + return { + "continuation": {"global_step_offset": 2150000}, + "data_metadata": {"splits": {"train": 80000000}}, + "model": {"block_size": 256}, + "training": {"batch_size": 32, "grad_accum_steps": 1, "max_steps": steps, + "target_epochs": steps * 8192 / 80000000, "epoch_interval": 1.024}, + } + + +@pytest.mark.parametrize("steps,expected", [ + (20, (0, 20)), + (25000, (0, 10000, 20000, 25000)), + (1000000, tuple(range(0, 1000001, 10000))), +]) +def test_clip_continuation_requires_its_planned_grid(steps, expected): + _validate_clip_snapshot_schedule(_snapshot_manifest(steps), expected) + + +@pytest.mark.parametrize("observed", [(0, 25000), (0, 10000, 25000), (0, 10000, 20000)]) +def test_clip_continuation_missing_snapshots_are_rejected(observed): + with pytest.raises(CompletedRunValidationError, match="exact configured snapshot grid"): + _validate_clip_snapshot_schedule(_snapshot_manifest(25000), observed) + + +def test_historical_clip_campaign_still_requires_ten_snapshots(): + with pytest.raises(CompletedRunValidationError, match="fewer than ten"): + _validate_clip_snapshot_schedule({}, (0, 20)) + _validate_clip_snapshot_schedule({}, tuple(range(10))) From f53bf3e8d5e66391656891cf63dacef06a28a1f3 Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Fri, 2 Oct 2026 12:46:01 -0700 Subject: [PATCH 05/36] Add continuous eight-chip MuonClip experiment with fixed probes and GCS checkpoints --- baseline/nanogpt_one_head/README.md | 7 + .../configs/muonclip_continuous8.yaml | 130 +++++++++ .../nanogpt_one_head/continuous8/README.md | 113 ++++++++ .../continuous8/cloudshell.py | 154 +++++++++++ .../continuous8/requirements.txt | 14 + .../nanogpt_one_head/continuous8/startup.sh | 55 ++++ .../nanogpt_one_head/continuous8/worker.sh | 66 +++++ .../src/rg_nanogpt_one_head/checkpoints.py | 2 + .../src/rg_nanogpt_one_head/config.py | 5 + .../src/rg_nanogpt_one_head/continuous_run.py | 71 +++++ .../rg_nanogpt_one_head/continuous_support.py | 246 ++++++++++++++++++ .../src/rg_nanogpt_one_head/engine.py | 6 + .../src/rg_nanogpt_one_head/train_loop.py | 4 + .../tests/test_continuous8.py | 194 ++++++++++++++ 14 files changed, 1067 insertions(+) create mode 100644 baseline/nanogpt_one_head/configs/muonclip_continuous8.yaml create mode 100644 baseline/nanogpt_one_head/continuous8/README.md create mode 100644 baseline/nanogpt_one_head/continuous8/cloudshell.py create mode 100644 baseline/nanogpt_one_head/continuous8/requirements.txt create mode 100644 baseline/nanogpt_one_head/continuous8/startup.sh create mode 100644 baseline/nanogpt_one_head/continuous8/worker.sh create mode 100644 baseline/nanogpt_one_head/src/rg_nanogpt_one_head/continuous_run.py create mode 100644 baseline/nanogpt_one_head/src/rg_nanogpt_one_head/continuous_support.py create mode 100644 baseline/nanogpt_one_head/tests/test_continuous8.py diff --git a/baseline/nanogpt_one_head/README.md b/baseline/nanogpt_one_head/README.md index 169602a9..cd5caf19 100644 --- a/baseline/nanogpt_one_head/README.md +++ b/baseline/nanogpt_one_head/README.md @@ -373,3 +373,10 @@ notebook structure. The same tests run in the repository's baseline CI. See [TPU_CONTINUATION.md](TPU_CONTINUATION.md) for full-state continuation beyond the original training horizon, periodic test accuracy, retained metric history, and safe pause/resume with bounded checkpoint storage. + +## Continuous single-host eight-chip run + +For the fresh 124M-parameter MuonClip experiment with 5B training tokens, fixed +document token-error probes, raw/clipped alpha and synchronous Cloud Storage +checkpoints, see [continuous8/README.md](continuous8/README.md). This launcher +disables automatic restarts and does not use the segmented continuation runner. diff --git a/baseline/nanogpt_one_head/configs/muonclip_continuous8.yaml b/baseline/nanogpt_one_head/configs/muonclip_continuous8.yaml new file mode 100644 index 00000000..5693d613 --- /dev/null +++ b/baseline/nanogpt_one_head/configs/muonclip_continuous8.yaml @@ -0,0 +1,130 @@ +protocol: + name: muonclip_continuous8_124m_5bt + version: 1 + description: Fresh continuous single-host eight-chip run; no segments, no automatic + restart; fixed document probes; durable GCS checkpoints. +dataset: + name: HuggingFaceFW/fineweb-edu + config: sample-10BT + split: train + revision: 593b3a867298afb8ce42625a270ef20ddcad28f9 + tokenizer: gpt2 + train_tokens: 5000000000 + val_tokens: 10000000 + test_tokens: 10000000 +model: + vocab_size: 50257 + block_size: 256 + n_layer: 12 + n_head: 12 + n_embd: 768 + dropout: 0.0 + bias: false + tie_weights: true +training: + seeds: + - 1337 + batch_size: 32 + grad_accum_steps: 4 + target_epochs: 6.5536 + epoch_interval: 0.0131072 + eval_interval_steps: 500 + eval_batches: 16 + checkpoint_interval_steps: 2000 + grad_clip: 1.0 + max_steps: 1000000 + min_free_disk_gb: 20 +optimizer_profiles: + sgd_momentum: + display_name: SGD + Nesterov momentum + family: sgd + learning_rate: 0.05 + min_learning_rate: 0.005 + warmup_fraction: 0.1 + schedule: warmup_cosine + momentum: 0.9 + dampening: 0.0 + nesterov: true + weight_decay: 0.01 + adamw: + display_name: AdamW + family: adamw + learning_rate: 0.0006 + min_learning_rate: 6.0e-05 + warmup_fraction: 0.01 + schedule: warmup_cosine + beta1: 0.9 + beta2: 0.95 + epsilon: 1.0e-08 + weight_decay: 0.1 + muon: + display_name: Muon + auxiliary AdamW + family: muon + matrix_learning_rate: 0.02 + matrix_min_learning_rate: 0.002 + aux_learning_rate: 0.0003 + aux_min_learning_rate: 3.0e-05 + warmup_fraction: 0.05 + schedule: warmup_cosine + momentum: 0.95 + nesterov: true + newton_schulz_steps: 5 + muon_epsilon: 1.0e-07 + matrix_weight_decay: 0.01 + beta1: 0.9 + beta2: 0.95 + epsilon: 1.0e-08 + aux_weight_decay: 0.01 + muon_clip: + display_name: MuonClip + RMS-matched updates + auxiliary AdamW + family: muon_clip + learning_rate: 0.0002 + min_learning_rate: 2.0e-05 + warmup_fraction: 0.02 + schedule: warmup_cosine + momentum: 0.95 + nesterov: false + newton_schulz_steps: 5 + muon_epsilon: 1.0e-07 + weight_decay: 0.1 + update_rms_scale: 0.2 + qk_clip_threshold: 100.0 + qk_clip_balance: 0.5 + qk_diagnostics_interval: 500 + beta1: 0.9 + beta2: 0.95 + epsilon: 1.0e-08 + lr_schedule_steps: 100000 +evaluation: + train_probe_seed: 21001 + validation_probe_seed: 22001 + test_probe_seed: 23001 + bleu_probe_seed: 24001 + bleu_examples: 64 + bleu_prompt_tokens: 64 + bleu_continuation_tokens: 32 + bleu_batch_size: 4 + test_interval_steps: 500 + document_probe: true + probe_documents: 512 +weightwatcher: + enabled: true + ERG: true + randomize: true + strict: true + min_evals: 20 + fix_fingers: clip_xmax + max_fingers: 10 + require_raw_alpha: true +runtime: + matmul_precision: highest + mps_fallback: true + deterministic_algorithms: false + empty_mps_cache_after_weightwatcher: true + tpu_spmd: true + tpu_expected_chips: 8 +continuous: + enabled: true + cloud_required: true + max_wall_hours: 144 + auto_resume: false diff --git a/baseline/nanogpt_one_head/continuous8/README.md b/baseline/nanogpt_one_head/continuous8/README.md new file mode 100644 index 00000000..f961f8f1 --- /dev/null +++ b/baseline/nanogpt_one_head/continuous8/README.md @@ -0,0 +1,113 @@ +# Continuous eight-chip MuonClip experiment + +This is a fresh single-host v5e-8 experiment, not a continuation of the earlier +one-layer run. The older extension changed the learning-rate phase and reset +phase diagnostics. That alone does not establish a broken checkpoint restore; +this experiment removes deliberate process/phase restarts from the scientific run. + +## Protocol + +- One seed (1337), one Python training process, no automatic retry or resume. +- NanoGPT: 12 layers, 12 heads, width 768, context 256, tied GPT-2 embeddings, + no dropout, approximately 124M parameters. MuonClip with the existing QK + clipping implementation; no alpha-based adaptive learning-rate controller. +- FineWeb-Edu `sample-10BT`, revision pinned in the YAML. Exactly 5B train, + 10M validation and 10M test tokens, split at document boundaries. +- Global microbatch 32 sequences, four accumulation steps: 32,768 tokens/update. + Weights are replicated; batch is sharded across eight chips on ONE host. +- Fixed maximum 1,000,000 updates (32.768B token presentations, 6.5536 corpus + equivalents). Learning rate warms up for 2,000 updates to 2e-4, follows its + preregistered cosine to 2e-5 by update 100,000, then stays at that floor. + No schedule is rebuilt or extended after launch. This is a new model/data + regime, not a controlled replication of the old model. +- Stop request after 144 training wall-hours, honored at the next full + checkpoint. Actual achievable update count is unknown until benchmarked. + Cloud allocation is capped at seven days, including setup. Hardware can fail; + no promise of uninterrupted infrastructure or a particular correlation is made. + +## Fixed measurements + +512 distinct eligible documents per split, one 257-token window per document, +256 teacher-forced next-token predictions per window. Seeds, document IDs, +absolute offsets, window SHA256 and corpus hashes are recorded once and checked. +The probe consumes no training RNG. This uses the same *definition* of token +error as the earlier document audit, but a larger, new probe/corpus; its numerical +values must not be presented as an exact replication of yesterday's probe. + +`token_error.csv` records top-1 token error percentages every 500 updates. +`alpha_token_error.csv` pairs errors and WeightWatcher measurements at exactly +one model state every 2,000 updates, including a model tensor hash and probe hash. +All 72 Q/K/V/O/MLP_IN/MLP_OUT matrices are retained in `spectral/layers.csv`. +Raw and clipped alpha remain separate. An aggregate is NaN if the expected +matrix count is not present; changing subsets must not create a spurious trend. +`alpha_token_error.png` plots raw mean/min alpha against token error with ordinary +regression lines, step coloring, no detrending, and no step-zero point. Repeated +checkpoints are dependent observations: Pearson r is descriptive, not causal. +One seed does not provide across-seed error bars. + +The test split is repeatedly monitored and is not used for checkpoint selection +or an adaptive training controller. Validation loss selects the best checkpoint. +It is a monitored test set, not an untouched final confirmation set. + +## Storage and failure behavior + +A dedicated 500 GB persistent disk is attached to the single host and mounted at +`/mnt/disks/rg-data`. GCS is the durable experiment archive: +`gs://tpu-builders-504820-ww-continuous8/runs/ww-continuous8-20261002/`. +The pinned dataset is uploaded before scientific training begins. Code commit, +resolved dependency versions, preflight report, fixed probes, metrics, spectra, +plots and logs are recorded. Full checkpoints include optimizer buffers/counters, +RNG/sampler state, learning-rate/config identity and monitoring state. + +Checkpoint files upload synchronously with CRC32C checking. A receipt containing +SHA256, generation, size and identity is written only after successful upload. +`LATEST_RESUMABLE.json` only points to a full-state checkpoint with the required +resume diagnostics (or initialization). All checkpoint artifacts are retained; +large archives incur storage costs. Backup failure stops the run after retaining +the local checkpoint, rather than silently running without durable protection. +Abrupt hardware loss may lose work since the last uploaded checkpoint. + +No systemd restart, continuation supervisor, reboot startup replay, or implicit +scientific directory reuse is permitted. A later manual recovery would be a +separately identified resumed run and is not exposed by this launcher. Recoverable +checkpoints do not constitute a guarantee that any future environment reproduces +TPU updates bit-for-bit. + +Before the scientific process, the existing TPU preflight checks global gradients, +QK clipping, CPU/TPU evaluation agreement and restored optimizer/sampler updates, +then benchmarks the proposed large model shape. That separate disposable test +may restore a checkpoint; the scientific run always starts from step zero. +CPU tests additionally check exact continued-vs-restored weights, optimizer hashes, +sampler state, LR scheduling and metrics. Actual TPU validation must pass on the +allocated hardware before the launch script starts scientific training. + +## Launch and monitor + +From a clean checkout of `codex/continuous-muonclip-8`, on Cloud Shell: + +```bash +python3 baseline/nanogpt_one_head/continuous8/cloudshell.py launch --delete-old-experiments +``` + +The explicit cleanup flag deletes earlier `ww-long-`, `ww-mem-`, `ww-mem2-` and +`ww-v6e16-` queues/nodes in us-west4-a and us-east5-a, their discovered attached +data disks plus `ww-full-data-20260929`, and snapshots of those disks. The exact +inventory is saved to `~/continuous8-cleanup.json`. Other project resources, +unrelated buckets and local/downloaded files are not deleted. Inventory or +permission failures stop the operation. Re-running launch does not duplicate an +existing continuous8 queue or restart its training. + +```bash +python3 baseline/nanogpt_one_head/continuous8/cloudshell.py status +``` + +The queue may wait for capacity or fail quota validation. A startup script fetches +the exact launch commit and starts setup via a systemd service with `Restart=no`. +A persistent claim blocks reruns after reboot. The TPU uses a dedicated service +account with object-admin permission on this experiment bucket, not TPU-admin. +Consequently it cannot delete itself: early completion/failure leaves the TPU +allocated until explicit deletion or the seven-day expiry. Review status promptly. +At the published v5e Flex-start rate of $0.60/chip-hour, eight chips cost $4.80/hour, +or $806.40 for the seven-day cap, plus disk/bucket/network charges. Credits and +remaining balance must be checked in the project's billing account; this script +does not assert that sufficient credits remain. diff --git a/baseline/nanogpt_one_head/continuous8/cloudshell.py b/baseline/nanogpt_one_head/continuous8/cloudshell.py new file mode 100644 index 00000000..cf10f55c --- /dev/null +++ b/baseline/nanogpt_one_head/continuous8/cloudshell.py @@ -0,0 +1,154 @@ +#!/usr/bin/env python3 +"""Run in Cloud Shell: scoped old-experiment cleanup and one v5e-8 request.""" +import argparse +import json +from pathlib import Path +import subprocess +import tempfile + +PROJECT = 'tpu-builders-504820' +ZONE = 'us-west4-a' +QUEUE = 'ww-continuous8-20261002' +NODE = QUEUE + '-node' +DISK = QUEUE + '-data' +BUCKET = PROJECT + '-ww-continuous8' +SA_ID = 'rg-continuous-tpu' +SA = SA_ID + '@' + PROJECT + '.iam.gserviceaccount.com' +OLD_PREFIXES = ('ww-long-', 'ww-mem2-', 'ww-mem-', 'ww-v6e16-') +OLD_DISKS = {'ww-full-data-20260929'} + + +def gc(*args, capture=False): + return subprocess.run(['gcloud', *args, '--project='+PROJECT], check=True, + text=True, stdout=subprocess.PIPE if capture else None).stdout + + +def inventory(*args): + return json.loads(gc(*args, '--format=json', capture=True)) + + +def short(resource): + return resource['name'].rsplit('/', 1)[-1] + + +def old(name): + return name.startswith(OLD_PREFIXES) + + +def cleanup(): + # Snapshot inventory before deleting resources; record exact targets. + plan = {'queues': [], 'nodes': [], 'disks': [], 'snapshots': []} + for zone in ('us-west4-a', 'us-east5-a'): + for q in inventory('alpha', 'compute', 'tpus', 'queued-resources', 'list', '--zone='+zone): + if old(short(q)): + detail = inventory('alpha', 'compute', 'tpus', 'queued-resources', 'describe', short(q), '--zone='+zone) + specs = detail.get('tpu', {}).get('nodeSpec', []) + if any(not old(s.get('nodeId', '')) for s in specs): + raise RuntimeError('Old queue includes unrecognized nodes; refusing deletion') + plan['queues'].append((zone, short(q))) + for n in inventory('compute', 'tpus', 'tpu-vm', 'list', '--zone='+zone): + if old(short(n)): + plan['nodes'].append((zone, short(n))) + for d in n.get('dataDisks', []): + source = d['sourceDisk'] + if f'projects/{PROJECT}/zones/{zone}/disks/' not in source: + raise RuntimeError('Unexpected old disk project/zone') + OLD_DISKS.add(source.rsplit('/', 1)[-1]) + for d in inventory('compute', 'disks', 'list'): + if short(d) in OLD_DISKS: + plan['disks'].append((d['zone'].rsplit('/',1)[-1], short(d))) + for s in inventory('compute', 'snapshots', 'list'): + if s.get('sourceDisk','').rsplit('/',1)[-1] in OLD_DISKS or short(s).startswith('ww-long-saved-'): + plan['snapshots'].append(short(s)) + audit = Path.home() / 'continuous8-cleanup.json' + audit.write_text(json.dumps(plan, indent=2)) + print('Deleting these earlier experiment resources:', json.dumps(plan, indent=2), flush=True) + for zone, name in plan['queues']: + gc('alpha','compute','tpus','queued-resources','delete', name, '--zone='+zone, '--force', '--quiet') + # Refresh; queue deletion usually already removed its node. + for zone in ('us-west4-a', 'us-east5-a'): + for n in inventory('compute','tpus','tpu-vm','list','--zone='+zone): + if old(short(n)): + gc('compute','tpus','tpu-vm','delete',short(n),'--zone='+zone,'--quiet') + for zone, name in plan['disks']: + gc('compute','disks','delete',name,'--zone='+zone,'--quiet') + for name in plan['snapshots']: + gc('compute','snapshots','delete',name,'--quiet') + # Successful inventories are required; access errors cannot masquerade as absence. + for zone in ('us-west4-a', 'us-east5-a'): + for resource in ('queued-resources','tpu-vm'): + remaining = inventory('alpha','compute','tpus',resource,'list','--zone='+zone) + if any(old(short(x)) for x in remaining): + raise RuntimeError('Earlier experiment TPU/queue still exists') + print('Earlier targeted TPU experiments and their discovered disks/snapshots removed.', flush=True) + + +def status(): + detail = inventory('alpha','compute','tpus','queued-resources','describe',QUEUE,'--zone='+ZONE) + print(json.dumps(detail, indent=2)) + print(f'Cloud results: gs://{BUCKET}/runs/{QUEUE}/') + if detail.get('state',{}).get('state') == 'ACTIVE': + gc('compute','tpus','tpu-vm','ssh',NODE,'--zone='+ZONE,'--worker=0', + '--command=systemctl --no-pager status rg-continuous8.service; tail -n 35 /mnt/disks/rg-data/continuous8/run.log') + + +def provision(): + root = Path(__file__).resolve().parents[3] + commit = subprocess.check_output(['git','-C',str(root),'rev-parse','HEAD'],text=True).strip() + dirty = subprocess.check_output(['git','-C',str(root),'status','--porcelain'],text=True).strip() + if dirty: + raise RuntimeError('Provision from a clean checked-in source tree') + requests = inventory('alpha','compute','tpus','queued-resources','list','--zone='+ZONE) + if any(short(q)==QUEUE for q in requests): + print('Existing continuous-run request retained; no second run or restart.') + status() + return + gc('services','enable','tpu.googleapis.com','compute.googleapis.com','storage.googleapis.com','iam.googleapis.com') + buckets = inventory('storage','buckets','list') + if not any(x.get('name','').removeprefix('gs://').rstrip('/')==BUCKET for x in buckets): + gc('storage','buckets','create','gs://'+BUCKET,'--location=us-west4','--uniform-bucket-level-access') + accounts = inventory('iam','service-accounts','list') + if not any(x.get('email')==SA for x in accounts): + gc('iam','service-accounts','create',SA_ID,'--display-name=Continuous MuonClip TPU storage') + gc('storage','buckets','add-iam-policy-binding','gs://'+BUCKET, + '--member=serviceAccount:'+SA,'--role=roles/storage.objectAdmin') + disks = inventory('compute','disks','list') + if any(short(d)==DISK for d in disks): + raise RuntimeError('Dedicated run disk exists without its queue. Keep it intact and inspect the previous attempt before retrying.') + gc('compute','disks','create',DISK,'--zone='+ZONE,'--size=500GB','--type=pd-balanced', + '--labels=experiment=continuous8') + uri = f'gs://{BUCKET}/runs/{QUEUE}' + template = Path(__file__).with_name('startup.sh').read_text() + startup = template.replace('__COMMIT__',commit).replace('__GCS_URI__',uri) + state = dict(project=PROJECT,zone=ZONE,queue=QUEUE,node=NODE,disk=DISK,gcs_uri=uri,commit=commit) + (Path.home()/'continuous8-resources.json').write_text(json.dumps(state,indent=2)) + with tempfile.NamedTemporaryFile(mode='w',suffix='.sh') as f: + f.write(startup); f.flush() + gc('alpha','compute','tpus','queued-resources','create',QUEUE, + '--zone='+ZONE,'--node-id='+NODE,'--accelerator-type=v5litepod-8', + '--runtime-version=v2-alpha-tpuv5-lite','--provisioning-model=flex-start', + '--max-run-duration=7d','--valid-until-duration=24h', + '--service-account='+SA,'--scopes=https://www.googleapis.com/auth/cloud-platform', + '--data-disk=source=projects/'+PROJECT+'/zones/'+ZONE+'/disks/'+DISK+',mode=read-write', + '--metadata-from-file=startup-script='+f.name,'--labels=experiment=continuous8', + '--quiet','--async') + print('Submitted one continuous8 request. It starts setup and preflight automatically when allocated.') + print('Inspect with: python3 baseline/nanogpt_one_head/continuous8/cloudshell.py status') + print('Run settings: '+str(Path.home()/'continuous8-resources.json')) + print('Cloud results: '+uri) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument('action', choices=['launch','status']) + parser.add_argument('--delete-old-experiments', action='store_true') + args = parser.parse_args() + if args.action == 'status': + status() + else: + if args.delete_old_experiments: + cleanup() + provision() + +if __name__ == '__main__': + main() diff --git a/baseline/nanogpt_one_head/continuous8/requirements.txt b/baseline/nanogpt_one_head/continuous8/requirements.txt new file mode 100644 index 00000000..44965d1b --- /dev/null +++ b/baseline/nanogpt_one_head/continuous8/requirements.txt @@ -0,0 +1,14 @@ +# Framework installed separately from Google's TPU wheel repository. +numpy==1.26.4 +pandas==2.2.3 +scipy==1.14.1 +matplotlib==3.9.4 +PyYAML==6.0.2 +datasets==3.3.2 +tiktoken==0.9.0 +sacrebleu==2.5.1 +weightwatcher==0.7.7 +powerlaw==2.0.0 +packaging==24.2 +google-cloud-storage==2.19.0 +pytest==8.3.5 diff --git a/baseline/nanogpt_one_head/continuous8/startup.sh b/baseline/nanogpt_one_head/continuous8/startup.sh new file mode 100644 index 00000000..37911212 --- /dev/null +++ b/baseline/nanogpt_one_head/continuous8/startup.sh @@ -0,0 +1,55 @@ +#!/usr/bin/env bash +# Provisioner substitutes the pinned Git SHA and the dedicated GCS run URI. +set -Eeuo pipefail +DEVICE=/dev/disk/by-id/google-persistent-disk-1 +for attempt in $(seq 1 60); do + [ -b "$DEVICE" ] && break + sleep 2 +done +test -b "$DEVICE" +TYPE="$(blkid -s TYPE -o value "$DEVICE" || true)" +if [ -z "$TYPE" ]; then + # The provisioner creates this new dedicated disk; never format another device. + test "$(lsblk -dn -o TYPE "$DEVICE")" = disk + mkfs.ext4 -F "$DEVICE" +elif [ "$TYPE" != ext4 ]; then + echo 'Unexpected filesystem; refusing to format.' >&2 + exit 1 +fi +mkdir -p /mnt/disks/rg-data +mountpoint -q /mnt/disks/rg-data || mount "$DEVICE" /mnt/disks/rg-data +BASE=/mnt/disks/rg-data/continuous8 +mkdir -p "$BASE" +# Persistent guard prevents either setup or training from restarting after a reboot. +if ! mkdir "$BASE/STARTED_ONCE"; then + echo 'Already started once. Automatic restart is forbidden.' + exit 0 +fi +exec > >(tee -a "$BASE/startup.log") 2>&1 +apt-get update +DEBIAN_FRONTEND=noninteractive apt-get install -y python3-venv git +mkdir -p "$BASE/repo" +git -C "$BASE/repo" init +git -C "$BASE/repo" remote add origin https://github.com/CalculatedContent/rg_optimizers.git +git -C "$BASE/repo" fetch --depth 1 origin __COMMIT__ +git -C "$BASE/repo" checkout --detach FETCH_HEAD +cat > /etc/systemd/system/rg-continuous8.service < "$BASE/environment.lock.txt" +git rev-parse HEAD > "$BASE/source_commit.txt" +# Upload credentials test before corpus preparation or training. +"$PY" - <<'PY' +import os +from pathlib import Path +from rg_nanogpt_one_head.continuous_support import CloudPublisher +p=CloudPublisher(os.environ['RG_CONTINUOUS_GCS_URI']) +base=Path('/mnt/disks/rg-data/continuous8') +p.claim({'commit':(base/'source_commit.txt').read_text().strip(), 'automatic_restart':False}) +for name in ('environment.lock.txt','source_commit.txt'): + p.file(base/name, name) +p.json({'status':'preflight'}, 'SETUP_STATUS.json') +PY +# Tiny, independent test: global gradients, clipping, metrics, optimizer + RNG restore. +# Also measures the full proposed model shape before the long run. +"$PY" -m rg_nanogpt_one_head.tpu_spmd_check --backend tpu --chips 8 \ + --benchmark-config configs/muonclip_continuous8.yaml --benchmark-steps 10 \ + --output "$BASE/preflight.json" +"$PY" - <<'PY' +import os +from rg_nanogpt_one_head.continuous_support import CloudPublisher +CloudPublisher(os.environ['RG_CONTINUOUS_GCS_URI']).file('/mnt/disks/rg-data/continuous8/preflight.json', 'preflight.json') +PY +# Pinned dataset revision; exact document-disjoint split sizes and content hashes. +"$PY" -m rg_nanogpt_one_head.data --config configs/muonclip_continuous8.yaml \ + --output-dir "$BASE/data" +"$PY" - <<'PY' +import os +from pathlib import Path +from rg_nanogpt_one_head.continuous_support import CloudPublisher, sha_file +p=CloudPublisher(os.environ['RG_CONTINUOUS_GCS_URI']) +base=Path('/mnt/disks/rg-data/continuous8/data') +receipts={} +for path in sorted(base.iterdir()): + if path.is_file(): + receipts[path.name]=p.file(path, 'data/'+path.name) + receipts[path.name]['sha256']=sha_file(path) +p.json(receipts, 'data/COMPLETE.json') +p.json({'status':'training'}, 'SETUP_STATUS.json') +PY +# No resilient supervisor, --resume, or automatic retry anywhere in this path. +"$PY" -u -m rg_nanogpt_one_head.continuous_run \ + --config configs/muonclip_continuous8.yaml --data-root "$BASE/data" \ + --results-root "$BASE/results" --device tpu diff --git a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/checkpoints.py b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/checkpoints.py index f423e494..a079a157 100644 --- a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/checkpoints.py +++ b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/checkpoints.py @@ -177,6 +177,8 @@ def _atomic_torch_save(payload: dict[str, Any], path: Path) -> Path: temporary = path.with_suffix(path.suffix + ".tmp") torch.save(payload, temporary) temporary.replace(path) + from .continuous_support import publish_checkpoint + publish_checkpoint(path, payload) return path diff --git a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/config.py b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/config.py index da57b8c5..012e0c3e 100644 --- a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/config.py +++ b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/config.py @@ -418,6 +418,11 @@ def lr_schedule_steps( ) -> int: """Return the LR horizon, which may be shorter than training.""" train_tokens = int(train_tokens or cfg["dataset"]["train_tokens"]) + if "lr_schedule_steps" in profile: + steps = profile["lr_schedule_steps"] + if isinstance(steps, bool) or not isinstance(steps, int) or not 1 <= steps <= max_steps(cfg, train_tokens): + raise ValueError("lr_schedule_steps must be an integer within the training horizon") + return steps training_epochs = float(cfg["training"]["target_epochs"]) schedule_epochs = float(profile.get("lr_schedule_epochs", training_epochs)) if not 0 < schedule_epochs <= training_epochs: diff --git a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/continuous_run.py b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/continuous_run.py new file mode 100644 index 00000000..ec4e3cb3 --- /dev/null +++ b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/continuous_run.py @@ -0,0 +1,71 @@ +"""One fresh scientific process. No retry loop, continuation or resume option.""" +from __future__ import annotations +import argparse +import os +from pathlib import Path +import signal +import threading +import time + +from .continuous_support import atomic_json, publisher, publish_metadata + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument('--config', required=True) + parser.add_argument('--data-root', required=True) + parser.add_argument('--results-root', required=True) + parser.add_argument('--device', default='tpu', choices=['tpu', 'cpu']) + args = parser.parse_args() + from .muonclip import install_muonclip_extension + install_muonclip_extension() + from .config import load_config + from .training import run_optimizer_replicates + cfg = load_config(args.config) + if not cfg.get('continuous', {}).get('enabled') or cfg.get('continuation'): + raise ValueError('A fresh continuous-run config is required') + root = Path(args.results_root) + root.mkdir(parents=True, exist_ok=True) + # An exclusive claim survives VM reboots and remains even after failure. + with (root/'CONTINUOUS_STARTED.json').open('x') as f: + import json + json.dump(dict(pid=os.getpid(), started_unix=time.time(), start_step=0, + automatic_restart=False), f) + stop = root/'STOP' + cfg['training']['stop_file'] = str(stop.resolve()) + sink = publisher(cfg) + if sink: + sink.file(args.config, 'config.yaml') + sink.file(root/'CONTINUOUS_STARTED.json', 'CONTINUOUS_STARTED.json') + def request_stop(*unused): + stop.touch() + for sig in (signal.SIGTERM, signal.SIGINT): + signal.signal(sig, request_stop) + timer = threading.Timer(float(cfg['continuous']['max_wall_hours'])*3600, request_stop) + timer.daemon = True + timer.start() + status = {'status':'running', 'pid':os.getpid(), 'start_step':0, 'restarts':0} + try: + run_optimizer_replicates(cfg=cfg, config_path=args.config, + optimizer_name='muon_clip', seeds=(1337,), data_root=args.data_root, + results_root=root, device=args.device, resume=False, overwrite=False) + status['status'] = 'completed' + except SystemExit as exc: + status.update(status='stopped_at_checkpoint' if exc.code == 75 else 'failed', exit_code=exc.code) + raise + except BaseException as exc: + status.update(status='failed', error=f'{type(exc).__name__}: {exc}') + raise + finally: + timer.cancel() + status['ended_unix'] = time.time() + atomic_json(root/'CONTINUOUS_STATUS.json', status) + if sink: + sink.json(status, 'CONTINUOUS_STATUS.json') + run_dir = root/'muon_clip'/'seed_1337' + if run_dir.exists(): + publish_metadata(cfg, run_dir) + + +if __name__ == '__main__': + main() diff --git a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/continuous_support.py b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/continuous_support.py new file mode 100644 index 00000000..29fe7842 --- /dev/null +++ b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/continuous_support.py @@ -0,0 +1,246 @@ +"""Opt-in continuous-run probes, paired measurements and durable publication. + +No training RNG is consumed here. GCS writes are synchronous: a failed backup +raises instead of allowing an apparently healthy but unprotected long run. +""" +from __future__ import annotations +import csv +import hashlib +import json +import math +import os +from pathlib import Path +import time + +import numpy as np +import torch + + +def atomic_json(path, value): + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + tmp = path.with_suffix(path.suffix + '.tmp') + tmp.write_text(json.dumps(value, indent=2, allow_nan=False) + '\n') + tmp.replace(path) + + +def sha_file(path): + digest = hashlib.sha256() + with Path(path).open('rb') as f: + for block in iter(lambda: f.read(8 * 1024 * 1024), b''): + digest.update(block) + return digest.hexdigest() + + +def document_windows(data, count, width, seed, eot=50256): + # Scan in bounded chunks: do not allocate a 5 GB boolean corpus mask. + ends = [] + chunk = 4_000_000 + for offset in range(0, len(data), chunk): + ends.extend((np.flatnonzero(data[offset:offset+chunk] == eot) + offset).tolist()) + starts = np.asarray([0, *[e + 1 for e in ends]], dtype=np.int64) + ends = np.asarray([*ends, len(data)], dtype=np.int64) + eligible = np.flatnonzero(ends - starts >= width) + if len(eligible) < count: + raise ValueError(f'Need {count} eligible documents, found {len(eligible)}') + rng = np.random.default_rng(seed) + docs = rng.choice(eligible, count, replace=False) + offsets = np.asarray([rng.integers(starts[i], ends[i]-width+1) for i in docs]) + windows = np.stack([np.asarray(data[o:o+width], dtype=np.int64) for o in offsets]) + if np.any(windows == eot): + raise RuntimeError('Probe crossed a document boundary') + return windows, docs, offsets + + +def build_document_probes(cfg, arrays, run_dir, metadata): + ev = cfg['evaluation'] + count = int(ev['probe_documents']) + batch = int(cfg['training']['batch_size']) + width = int(cfg['model']['block_size']) + 1 + if count % batch: + raise ValueError('probe_documents must divide into full evaluation batches') + records, probes = {}, [] + for split, key in [('train','train_probe_seed'), ('val','validation_probe_seed'), ('test','test_probe_seed')]: + windows, docs, offsets = document_windows(arrays[split], count, width, int(ev[key]), int(metadata['eot_token'])) + records[split] = dict(document_ids=docs.tolist(), token_offsets=offsets.tolist(), + windows_sha256=hashlib.sha256(windows.tobytes()).hexdigest(), + data_sha256=metadata['files'][split]['sha256']) + probes.append([(torch.from_numpy(windows[i:i+batch,:-1].copy()), + torch.from_numpy(windows[i:i+batch,1:].copy())) + for i in range(0, count, batch)]) + record = dict(schema_version=1, documents=count, context=width-1, splits=records, + metric='100 * incorrect argmax next-token predictions / scored tokens; teacher forced', + tokens_per_split=count*(width-1)) + path = Path(run_dir)/'fixed_document_probe.json' + if path.exists() and json.loads(path.read_text()) != record: + raise RuntimeError('Fixed document probe changed') + atomic_json(path, record) + return probes + + +class CloudPublisher: + def __init__(self, uri): + if not uri.startswith('gs://') or '/' not in uri[5:]: + raise ValueError('Expected gs://bucket/run-prefix') + from google.cloud import storage + bucket, self.prefix = uri[5:].split('/', 1) + self.prefix = self.prefix.strip('/') + if not self.prefix: + raise ValueError('An experiment-specific bucket prefix is required') + self.bucket = storage.Client().bucket(bucket) + self.metadata_cache = {} + + def file(self, path, relative): + blob = self.bucket.blob(self.prefix + '/' + relative) + blob.upload_from_filename(str(path), checksum='crc32c', timeout=600) + blob.reload() + if int(blob.size) != Path(path).stat().st_size: + raise RuntimeError('Cloud object size differs from local artifact') + return {'object': blob.name, 'generation': str(blob.generation), + 'bytes': int(blob.size), 'crc32c': blob.crc32c} + + def claim(self, value): + # A deleted/recreated VM must never silently overwrite an earlier archive. + self.bucket.blob(self.prefix + '/RUN_CLAIM.json').upload_from_string( + json.dumps(value), content_type='application/json', checksum='crc32c', + if_generation_match=0, timeout=120) + + def snapshot_text_file(self, path, relative): + # Runtime libraries may append to the log during upload. Upload a fixed + # prefix, never compare its length to a subsequently growing live file. + with Path(path).open('rb') as f: + content = f.read(os.fstat(f.fileno()).st_size) + self.bucket.blob(self.prefix + '/' + relative).upload_from_string( + content, content_type='text/plain', checksum='crc32c', timeout=120) + + def json(self, value, relative): + self.bucket.blob(self.prefix + '/' + relative).upload_from_string( + json.dumps(value, indent=2, allow_nan=False), content_type='application/json', + checksum='crc32c', timeout=120) + + +_PUBLISHERS = {} +def publisher(cfg): + if not cfg.get('continuous', {}).get('enabled'): + return None + uri = os.environ.get('RG_CONTINUOUS_GCS_URI', '') + if not uri: + if cfg['continuous'].get('cloud_required', True): + raise RuntimeError('RG_CONTINUOUS_GCS_URI is required') + return None + if uri not in _PUBLISHERS: + _PUBLISHERS[uri] = CloudPublisher(uri) + return _PUBLISHERS[uri] + + +def publish_metadata(cfg, run_dir): + sink = publisher(cfg) + if sink is None: + return + root = Path(run_dir) + # Called on the training thread, after writers flush, with no concurrent writer. + for path in sorted(root.rglob('*')): + if path.is_file() and path.suffix in {'.json', '.csv', '.yaml', '.png', '.pdf'}: + stat = path.stat() + identity = (stat.st_mtime_ns, stat.st_size) + if sink.metadata_cache.get(str(path)) != identity: + sink.file(path, 'results/' + path.relative_to(root).as_posix()) + sink.metadata_cache[str(path)] = identity + log = os.environ.get('RG_CONTINUOUS_RUN_LOG') + if log and Path(log).is_file(): + sink.snapshot_text_file(log, 'run.log') + + +def publish_checkpoint(path, payload): + cfg = payload.get('config', {}) + sink = publisher(cfg) + if sink is None: + return + path = Path(path) + step = int(payload['step']) + relative = f'checkpoints/step_{step:09d}/{path.name}' + receipt = sink.file(path, relative) + receipt.update(step=step, sha256=sha_file(path), fingerprint=payload['fingerprint'], + model_state_sha256=payload['model_state_sha256'], + resumable=bool(payload.get('optimizers')) and + (step == 0 or payload.get('resume_diagnostics') is not None)) + # A completion marker is published only AFTER the complete verified upload. + sink.json(receipt, relative + '.complete.json') + if receipt['resumable'] and path.name in {'checkpoint_latest.pt', 'checkpoint_final.pt', 'checkpoint_initial.pt'}: + sink.json(receipt, 'LATEST_RESUMABLE.json') + root = path.parent.parent if path.parent.name == 'epoch_checkpoints' else path.parent + publish_metadata(cfg, root) + print(f'[continuous-backup] saved step={step} file={path.name}', flush=True) + + +def record_pair(cfg, run_dir, row, summary): + if not cfg.get('continuous', {}).get('enabled'): + return + count = 6 * int(cfg['model']['n_layer']) + pair = {k: row[k] for k in ('step', 'tokens_seen', 'elapsed_sec', 'primary_lr', 'train_loss', 'val_loss', 'test_loss')} + pair.update(train_token_error_pct=100*(1-row['train_accuracy']), + test_token_error_pct=100*(1-row['test_accuracy']), + model_state_sha256=summary['model_state_sha256'], + probe_sha256=sha_file(Path(run_dir)/'fixed_document_probe.json'), + expected_matrices=count) + for kind in ('raw', 'clip_xmax'): + n = int(summary.get(f'alpha_{kind}_n', 0)) + pair[f'alpha_{kind}_n'] = n + for stat in ('mean', 'min'): + # Never silently average a changing subset of layers. + pair[f'alpha_{kind}_{stat}'] = summary.get(f'alpha_{kind}_{stat}', float('nan')) if n == count else float('nan') + path = Path(run_dir)/'alpha_token_error.csv' + exists = path.exists() + with path.open('a', newline='') as f: + writer = csv.DictWriter(f, fieldnames=list(pair)) + if not exists: + writer.writeheader() + writer.writerow(pair) + f.flush() + os.fsync(f.fileno()) + plot_pairs(path) + publish_metadata(cfg, run_dir) + + +def plot_pairs(path): + # No pandas CSV parser: the previous environment segfaulted in that path. + import matplotlib + matplotlib.use('Agg') + import matplotlib.pyplot as plt + with Path(path).open() as f: + rows = [r for r in csv.DictReader(f) if int(r['step']) > 0] + fig, axes = plt.subplots(1, 2, figsize=(10, 4), constrained_layout=True) + for ax, key in zip(axes, ('alpha_raw_mean', 'alpha_raw_min')): + pairs = [(float(r[key]), float(r['test_token_error_pct']), int(r['step'])) for r in rows] + pairs = np.asarray([p for p in pairs if all(math.isfinite(v) for v in p)], dtype=float).reshape(-1, 3) + if len(pairs): + x, y, steps = pairs.T + dots = ax.scatter(x, y, c=steps, s=18, cmap='viridis') + fig.colorbar(dots, ax=ax, label='Training step') + if len(pairs) >= 3 and np.ptp(x) > 0 and np.ptp(y) > 0: + slope, intercept = np.polyfit(x, y, 1) + line = np.array([x.min(), x.max()]) + ax.plot(line, slope*line+intercept, color='black', lw=1) + ax.set_title(f"Pearson r={np.corrcoef(x,y)[0,1]:.3f}; n={len(x)}") + ax.set_xlabel(key.replace('_', ' ')) + ax.set_ylabel('Fixed-test token error (%)') + fig.suptitle('Continuous MuonClip — no detrending; step zero excluded') + fig.savefig(Path(path).with_suffix('.png'), dpi=160) + plt.close(fig) + + +def record_token_errors(cfg, run_dir, row): + if not cfg.get('continuous', {}).get('enabled'): + return + result = {'step': row['step'], 'tokens_seen': row['tokens_seen'], + 'train_token_error_pct': 100*(1-row['train_accuracy']), + 'val_token_error_pct': 100*(1-row['val_accuracy']), + 'test_token_error_pct': 100*(1-row['test_accuracy']), + 'probe_sha256': sha_file(Path(run_dir)/'fixed_document_probe.json')} + path = Path(run_dir)/'token_error.csv' + exists = path.exists() + with path.open('a', newline='') as f: + w = csv.DictWriter(f, fieldnames=list(result)) + if not exists: w.writeheader() + w.writerow(result) + publish_metadata(cfg, run_dir) diff --git a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/engine.py b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/engine.py index c83ad8f0..1ba067d3 100644 --- a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/engine.py +++ b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/engine.py @@ -69,6 +69,8 @@ def run_one( if resume and overwrite: raise ValueError("resume and overwrite are mutually exclusive") + if cfg.get("continuous", {}).get("enabled") and (resume or overwrite or cfg.get("continuation")): + raise ValueError("Continuous experiment must start fresh: no resume, overwrite, or segments") initialize_spmd(cfg, device) data_root = Path(data_root) results_root = Path(results_root) @@ -268,6 +270,10 @@ def run_one( model=model, ) + if cfg["evaluation"].get("document_probe", False): + from .continuous_support import build_document_probes + train_probe, val_probe, test_probe = build_document_probes(cfg, arrays, run_dir, data_metadata) + # Persist an immutable step-zero checkpoint before evaluation, WeightWatcher, # or the first optimizer update. This makes initial-versus-final angular # analysis use two actual checkpoint files rather than seed reconstruction. diff --git a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/train_loop.py b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/train_loop.py index 44271e44..eae7e7a0 100644 --- a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/train_loop.py +++ b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/train_loop.py @@ -420,6 +420,8 @@ def execute_training_loop( } metrics_writer.writerow(row) metrics_handle.flush() + from .continuous_support import record_token_errors + record_token_errors(cfg, run_dir, row) if canaries is not None: canary_summary = canaries.evaluate(model, device=device, step=completed_steps, epoch=actual_epoch) if progress: @@ -468,6 +470,8 @@ def execute_training_loop( seed=int(seed), fingerprint=fingerprint, ) + from .continuous_support import record_pair + record_pair(cfg, run_dir, row, ww_summary) if progress: print( "[one-head-ww] " diff --git a/baseline/nanogpt_one_head/tests/test_continuous8.py b/baseline/nanogpt_one_head/tests/test_continuous8.py new file mode 100644 index 00000000..b74e65e7 --- /dev/null +++ b/baseline/nanogpt_one_head/tests/test_continuous8.py @@ -0,0 +1,194 @@ +import csv +from pathlib import Path +import numpy as np +import pytest +import torch +from rg_nanogpt_one_head import continuous_support as cs + + +def test_probe_fixed_document_disjoint_and_rng_independent(tmp_path): + data=np.tile(np.r_[np.arange(20),63],40).astype(np.uint16) + a=cs.document_windows(data,16,9,123,eot=63) + b=cs.document_windows(data,16,9,123,eot=63) + for x,y in zip(a,b):np.testing.assert_array_equal(x,y) + assert len(set(a[1]))==16 and not np.any(a[0]==63) + cfg={'training':{'batch_size':4},'model':{'block_size':8},'evaluation':{ + 'probe_documents':16,'train_probe_seed':1,'validation_probe_seed':2,'test_probe_seed':3}} + meta={'eot_token':63,'files':{s:{'sha256':'a'} for s in ('train','val','test')}} + arrays={s:data for s in ('train','val','test')} + state=torch.get_rng_state().clone() + probes=cs.build_document_probes(cfg,arrays,tmp_path,meta) + assert torch.equal(state,torch.get_rng_state()) + assert sum(y.numel() for x,y in probes[2])==128 + cfg['evaluation']['test_probe_seed']=99 + with pytest.raises(RuntimeError,match='probe changed'): + cs.build_document_probes(cfg,arrays,tmp_path,meta) + + +def test_failed_upload_never_publishes_completion_marker(tmp_path,monkeypatch): + events=[] + class Sink: + def file(self,*a):events.append('file');raise OSError('upload failed') + def json(self,*a):events.append('marker') + monkeypatch.setattr(cs,'publisher',lambda c:Sink()) + p=tmp_path/'checkpoint_latest.pt';p.write_bytes(b'checkpoint') + with pytest.raises(OSError,match='upload failed'): + cs.publish_checkpoint(p,dict(config={},step=4)) + assert events==['file'] and p.exists() + + +def test_latest_pointer_requires_resume_state(tmp_path,monkeypatch): + calls=[] + class Sink: + def file(self,*a):return {'generation':'1','bytes':10} + def json(self,value,path):calls.append(path) + monkeypatch.setattr(cs,'publisher',lambda c:Sink()) + monkeypatch.setattr(cs,'publish_metadata',lambda *a:None) + p=tmp_path/'checkpoint_latest.pt';p.write_bytes(b'checkpoint') + payload=dict(config={},step=4,fingerprint='x',model_state_sha256='m',optimizers=[{}]) + cs.publish_checkpoint(p,payload) + assert 'LATEST_RESUMABLE.json' not in calls + payload['resume_diagnostics']={'valid':True} + cs.publish_checkpoint(p,payload) + assert calls[-1]=='LATEST_RESUMABLE.json' + + +def test_pair_refuses_changing_layer_subset(tmp_path,monkeypatch): + monkeypatch.setattr(cs,'plot_pairs',lambda *a:None) + monkeypatch.setattr(cs,'publish_metadata',lambda *a:None) + (tmp_path/'fixed_document_probe.json').write_text('{}') + row=dict(step=2000,tokens_seen=10,elapsed_sec=1,primary_lr=.1,train_loss=2, + val_loss=3,test_loss=4,train_accuracy=.3,test_accuracy=.2) + summary={'model_state_sha256':'m','alpha_raw_n':5,'alpha_raw_mean':2.1, + 'alpha_clip_xmax_n':6,'alpha_clip_xmax_mean':2.,'alpha_clip_xmax_min':1.9} + cs.record_pair({'continuous':{'enabled':True},'model':{'n_layer':1}},tmp_path,row,summary) + with (tmp_path/'alpha_token_error.csv').open() as f:r=next(csv.DictReader(f)) + assert r['alpha_raw_mean']=='nan' and float(r['test_token_error_pct'])==80 + assert r['alpha_clip_xmax_mean']=='2.0' + + +def test_guard_rejects_resume_before_tpu_initialization(tmp_path): + from rg_nanogpt_one_head.muonclip import install_muonclip_extension + install_muonclip_extension() + from rg_nanogpt_one_head.engine import run_one + with pytest.raises(ValueError,match='must start fresh'): + run_one(cfg={'continuous':{'enabled':True}},data_root=tmp_path, + results_root=tmp_path,optimizer_name='muon_clip',seed=1,device='tpu',resume=True) + + +def test_fixed_config_alignment_and_schedule(): + from rg_nanogpt_one_head.muonclip import install_muonclip_extension + install_muonclip_extension() + from rg_nanogpt_one_head.config import load_config,epoch_step_map,lr_schedule_steps,optimizer_profile,warmup_steps + root=Path(__file__).resolve().parents[1] + cfg=load_config(root/'configs/muonclip_continuous8.yaml') + assert cfg['model']['n_layer']==12 and cfg['runtime']['tpu_expected_chips']==8 + assert cfg['dataset']['train_tokens']==5_000_000_000 and 'continuation' not in cfg + assert list(epoch_step_map(cfg))==list(range(0,1_000_001,2000)) + p=optimizer_profile(cfg,'muon_clip') + assert lr_schedule_steps(cfg,p)==100_000 and warmup_steps(p,100_000)==2000 + + +def test_muonclip_checkpoint_preserves_next_updates_lr_optimizer_and_sampler(tmp_path): + from copy import deepcopy + from rg_nanogpt_one_head.muonclip import install_muonclip_extension + install_muonclip_extension() + from rg_nanogpt_one_head.model import GPT,GPTConfig + from rg_nanogpt_one_head.config import load_config,optimizer_profile + from rg_nanogpt_one_head.optimizers import make_optimizer_handles,zero_grad,optimizer_step,set_learning_rates,optimizer_state_dict + from rg_nanogpt_one_head.checkpoints import save_training_checkpoint,load_training_checkpoint_for_resume,optimizer_state_sha256 + from rg_nanogpt_one_head.evaluation import random_batch,evaluate_probe,fixed_probe + from rg_nanogpt_one_head.runtime import parameter_snapshot + torch.set_num_threads(1) + cfg=deepcopy(load_config(Path(__file__).resolve().parents[1]/'configs/muonclip_reference.yaml')) + cfg['model'].update(vocab_size=64,block_size=8,n_embd=16,n_head=2,n_layer=2) + cfg['optimizer_profiles']['muon_clip']['qk_clip_threshold']=0.01 + torch.manual_seed(123) + net=GPT(GPTConfig(**cfg['model'])) + handles=make_optimizer_handles(net,optimizer_profile(cfg,'muon_clip')) + gen=torch.Generator().manual_seed(44) + data=np.tile(np.arange(64,dtype=np.uint16),20) + def advance(model,opts,generator,start,end): + for step in range(start,end): + set_learning_rates(opts,update_index=step,total_steps=10,warmup_steps=2) + x,y=random_batch(data,batch_size=4,block_size=8,generator=generator) + zero_grad(opts) + model(x,y)[1].backward() + torch.nn.utils.clip_grad_norm_(model.parameters(),1.,foreach=False) + optimizer_step(opts) + advance(net,handles,gen,0,3) + path=tmp_path/'checkpoint.pt' + save_training_checkpoint(path,model=net,handles=handles,step=3,best_validation_loss=4., + best_validation_step=2,elapsed_seconds=1.,fingerprint='test',cfg=cfg, + optimizer_name='muon_clip',seed=123,train_generator=gen, + resume_diagnostics={'previous_eval_snapshot':parameter_snapshot(net), + 'last_grad_pre':1.,'last_grad_post':1.,'last_clipped':False}) + advance(net,handles,gen,3,7) + resumed=GPT(GPTConfig(**cfg['model'])) + rh=make_optimizer_handles(resumed,optimizer_profile(cfg,'muon_clip')) + rg=torch.Generator() + loaded=load_training_checkpoint_for_resume(path,model=resumed,handles=rh, + expected_fingerprint='test',train_generator=rg) + advance(resumed,rh,rg,loaded[0],7) + for name,value in net.state_dict().items(): + assert torch.equal(value,resumed.state_dict()[name]),name + assert optimizer_state_sha256(optimizer_state_dict(handles))==optimizer_state_sha256(optimizer_state_dict(rh)) + assert torch.equal(gen.get_state(),rg.get_state()) + probe=fixed_probe(data,batch_size=4,block_size=8,n_batches=2,seed=99) + assert evaluate_probe(net,probe,torch.device('cpu'))==evaluate_probe(resumed,probe,torch.device('cpu')) + + +def test_fresh_continuous_engine_writes_fixed_probe_and_pairs(tmp_path,monkeypatch): + from copy import deepcopy + from rg_nanogpt_one_head.muonclip import install_muonclip_extension + install_muonclip_extension() + from rg_nanogpt_one_head.config import load_config + from rg_nanogpt_one_head.data import write_token_splits + from rg_nanogpt_one_head.checkpoints import model_state_sha256 + import rg_nanogpt_one_head.train_loop as loop + import rg_nanogpt_one_head.run_utils as utils + from rg_nanogpt_one_head.training import run_one + cfg=deepcopy(load_config(Path(__file__).resolve().parents[1]/'configs/muonclip_continuous8.yaml')) + cfg['model'].update(vocab_size=64,block_size=8,n_embd=16,n_head=2,n_layer=2) + cfg['runtime']['tpu_spmd']=False + cfg['dataset'].update(train_tokens=256,val_tokens=128,test_tokens=128) + cfg['training'].update(batch_size=2,grad_accum_steps=1,max_steps=4,target_epochs=.25, + epoch_interval=.125,eval_interval_steps=2,eval_batches=2, + checkpoint_interval_steps=2,min_free_disk_gb=0) + cfg['continuous']['cloud_required']=False + cfg['evaluation'].update(probe_documents=4,test_interval_steps=2,bleu_examples=2, + bleu_prompt_tokens=3,bleu_continuation_tokens=2,bleu_batch_size=2) + cfg['optimizer_profiles']['muon_clip']['lr_schedule_steps']=4 + for p in cfg['optimizer_profiles'].values():p.pop('lr_schedule_epochs',None) + class Encoder: + n_vocab=64 + eot_token=63 + def encode_ordinary(self,text):return list(range(16)) + data=tmp_path/'data' + write_token_splits(['doc']*100,Encoder(),data,train_tokens=256,val_tokens=128,test_tokens=128, + dataset_metadata={'dataset_name':cfg['dataset']['name'],'dataset_config':cfg['dataset']['config'], + 'dataset_split':'train','dataset_revision':cfg['dataset']['revision'],'tokenizer':'gpt2'}) + def spectrum(model,*a,**kw): + return {'model_state_sha256':model_state_sha256(model.state_dict()), + 'alpha_raw_n':12,'alpha_raw_mean':3.,'alpha_raw_min':2.5, + 'alpha_clip_xmax_n':12,'alpha_clip_xmax_mean':2.9,'alpha_clip_xmax_min':2.4} + monkeypatch.setattr(loop,'run_weightwatcher',spectrum) + monkeypatch.setattr(utils,'evaluate_bleu',lambda *a,**kw:{'bleu':0.}) + # Stop at the full step-4 checkpoint: this tests the actual monitoring loop + # without post-run spectral completion checks against our synthetic spectrum. + original=loop.save_training_checkpoint + def save(path,**kw): + result=original(path,**kw) + if kw['step']==4:raise RuntimeError('test stopped after saved step 4') + return result + monkeypatch.setattr(loop,'save_training_checkpoint',save) + with pytest.raises(RuntimeError,match='test stopped'): + run_one(cfg=cfg,data_root=data,results_root=tmp_path/'results',optimizer_name='muon_clip', + seed=1337,device='cpu',resume=False,progress=False) + run=tmp_path/'results'/'muon_clip'/'seed_1337' + with (run/'alpha_token_error.csv').open() as f:rows=list(csv.DictReader(f)) + assert [int(r['step']) for r in rows]==[0,2] + assert len({r['probe_sha256'] for r in rows})==1 + assert all(r['model_state_sha256'] for r in rows) + assert (run/'checkpoint_latest.pt').is_file() + assert (run/'alpha_token_error.png').is_file() From ab585394598e132f245aeea4922f43a0d507ef2e Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Fri, 2 Oct 2026 13:19:04 -0700 Subject: [PATCH 06/36] Bound MuonClip TPU pilots to one six-hour or two four-hour allocations Prepare shared 5B-token corpus on Cloud Shell before requesting TPUs. Preserve fixed probes and continuous MuonClip training; force a full checkpoint at the wall-time stop, record measured throughput projections, and cap allocation with Flex-start expiry. Add independent seed support and budget-ordering tests. Validation: 24 CPU tests passed; one hardware-only check skipped. TPU preflight must pass on allocation before scientific training starts. --- baseline/nanogpt_one_head/README.md | 4 +- .../configs/muonclip_continuous8.yaml | 11 +- .../nanogpt_one_head/continuous8/README.md | 63 +++++++--- .../continuous8/cloudshell.py | 108 ++++++++++++------ .../continuous8/prepare_cloud_data.py | 59 ++++++++++ .../nanogpt_one_head/continuous8/startup.sh | 7 ++ .../nanogpt_one_head/continuous8/worker.sh | 64 +++++++---- .../src/rg_nanogpt_one_head/continuous_run.py | 19 ++- .../rg_nanogpt_one_head/continuous_support.py | 11 ++ .../src/rg_nanogpt_one_head/data.py | 4 +- .../src/rg_nanogpt_one_head/train_loop.py | 9 +- .../tests/test_continuous8.py | 26 ++++- .../tests/test_continuous8_launcher.py | 59 ++++++++++ 13 files changed, 355 insertions(+), 89 deletions(-) create mode 100644 baseline/nanogpt_one_head/continuous8/prepare_cloud_data.py create mode 100644 baseline/nanogpt_one_head/tests/test_continuous8_launcher.py diff --git a/baseline/nanogpt_one_head/README.md b/baseline/nanogpt_one_head/README.md index cd5caf19..ad35a446 100644 --- a/baseline/nanogpt_one_head/README.md +++ b/baseline/nanogpt_one_head/README.md @@ -379,4 +379,6 @@ and safe pause/resume with bounded checkpoint storage. For the fresh 124M-parameter MuonClip experiment with 5B training tokens, fixed document token-error probes, raw/clipped alpha and synchronous Cloud Storage checkpoints, see [continuous8/README.md](continuous8/README.md). This launcher -disables automatic restarts and does not use the segmented continuation runner. +defaults to one six-hour v5e-8 allocation, with a two-machine/four-hour option. +Data preparation happens on Cloud Shell before TPU allocation; automatic +restarts and the segmented continuation runner are disabled. diff --git a/baseline/nanogpt_one_head/configs/muonclip_continuous8.yaml b/baseline/nanogpt_one_head/configs/muonclip_continuous8.yaml index 5693d613..c9346ce3 100644 --- a/baseline/nanogpt_one_head/configs/muonclip_continuous8.yaml +++ b/baseline/nanogpt_one_head/configs/muonclip_continuous8.yaml @@ -1,6 +1,6 @@ protocol: - name: muonclip_continuous8_124m_5bt - version: 1 + name: muonclip_continuous8_124m_5bt_pilot + version: 2 description: Fresh continuous single-host eight-chip run; no segments, no automatic restart; fixed document probes; durable GCS checkpoints. dataset: @@ -24,13 +24,14 @@ model: training: seeds: - 1337 + - 2027 batch_size: 32 grad_accum_steps: 4 target_epochs: 6.5536 - epoch_interval: 0.0131072 + epoch_interval: 0.0065536 eval_interval_steps: 500 eval_batches: 16 - checkpoint_interval_steps: 2000 + checkpoint_interval_steps: 500 grad_clip: 1.0 max_steps: 1000000 min_free_disk_gb: 20 @@ -126,5 +127,5 @@ runtime: continuous: enabled: true cloud_required: true - max_wall_hours: 144 + max_wall_hours: 5.5 auto_resume: false diff --git a/baseline/nanogpt_one_head/continuous8/README.md b/baseline/nanogpt_one_head/continuous8/README.md index f961f8f1..84d18576 100644 --- a/baseline/nanogpt_one_head/continuous8/README.md +++ b/baseline/nanogpt_one_head/continuous8/README.md @@ -7,7 +7,12 @@ this experiment removes deliberate process/phase restarts from the scientific ru ## Protocol -- One seed (1337), one Python training process, no automatic retry or resume. +The default budget is **one machine for six hours**. The alternative is two +independent machines for four hours each; they are different seeds, not multi-host +training of a single model. + +- Seed 1337 by default; the two-machine option adds seed 2027. Each machine has + one Python training process, no automatic retry or resume. - NanoGPT: 12 layers, 12 heads, width 768, context 256, tied GPT-2 embeddings, no dropout, approximately 124M parameters. MuonClip with the existing QK clipping implementation; no alpha-based adaptive learning-rate controller. @@ -20,10 +25,12 @@ this experiment removes deliberate process/phase restarts from the scientific ru preregistered cosine to 2e-5 by update 100,000, then stays at that floor. No schedule is rebuilt or extended after launch. This is a new model/data regime, not a controlled replication of the old model. -- Stop request after 144 training wall-hours, honored at the next full - checkpoint. Actual achievable update count is unknown until benchmarked. - Cloud allocation is capped at seven days, including setup. Hardware can fail; - no promise of uninterrupted infrastructure or a particular correlation is made. +- Google Cloud enforces a six-hour allocation limit (four hours per machine for + the two-machine option). Setup, downloads, compilation and diagnostics consume + allocation time. A stop is requested 30 minutes before that limit, measured + conservatively from VM boot; it saves after the next optimizer update. + Checkpoints also save every 500 updates. Actual throughput is measured on + hardware. A failure may interrupt training; automatic restart remains disabled. ## Fixed measurements @@ -36,14 +43,15 @@ values must not be presented as an exact replication of yesterday's probe. `token_error.csv` records top-1 token error percentages every 500 updates. `alpha_token_error.csv` pairs errors and WeightWatcher measurements at exactly -one model state every 2,000 updates, including a model tensor hash and probe hash. +one model state every 1,000 updates, including a model tensor hash and probe hash. All 72 Q/K/V/O/MLP_IN/MLP_OUT matrices are retained in `spectral/layers.csv`. Raw and clipped alpha remain separate. An aggregate is NaN if the expected matrix count is not present; changing subsets must not create a spurious trend. `alpha_token_error.png` plots raw mean/min alpha against token error with ordinary regression lines, step coloring, no detrending, and no step-zero point. Repeated checkpoints are dependent observations: Pearson r is descriptive, not causal. -One seed does not provide across-seed error bars. +The single-seed option does not provide across-seed error bars. Two seeds give +only a limited estimate of seed variability; compare matched training steps. The test split is repeatedly monitored and is not used for checkpoint selection or an adaptive training controller. Validation loss selects the best checkpoint. @@ -51,10 +59,14 @@ It is a monitored test set, not an untouched final confirmation set. ## Storage and failure behavior -A dedicated 500 GB persistent disk is attached to the single host and mounted at +A dedicated 200 GB persistent disk is attached to each host and mounted at `/mnt/disks/rg-data`. GCS is the durable experiment archive: -`gs://tpu-builders-504820-ww-continuous8/runs/ww-continuous8-20261002/`. -The pinned dataset is uploaded before scientific training begins. Code commit, +`gs://tpu-builders-504820-ww-continuous8/runs/ww-continuous8-pilot-20261002-sSEED/`. +The dataset is prepared once on Cloud Shell CPU and uploaded **before any TPU +is requested**. This can take hours and needs about 13 GB free under `/tmp`; +it is outside the TPU allocation budget. Both machines download the same +SHA256-verified files. Keep Cloud Shell open during preparation. No additional +CPU VM is provisioned by this launcher. Code commit, resolved dependency versions, preflight report, fixed probes, metrics, spectra, plots and logs are recorded. Full checkpoints include optimizer buffers/counters, RNG/sampler state, learning-rate/config identity and monitoring state. @@ -86,16 +98,21 @@ allocated hardware before the launch script starts scientific training. From a clean checkout of `codex/continuous-muonclip-8`, on Cloud Shell: ```bash -python3 baseline/nanogpt_one_head/continuous8/cloudshell.py launch --delete-old-experiments +python3 baseline/nanogpt_one_head/continuous8/cloudshell.py launch --machines 1 --hours 6 ``` -The explicit cleanup flag deletes earlier `ww-long-`, `ww-mem-`, `ww-mem2-` and -`ww-v6e16-` queues/nodes in us-west4-a and us-east5-a, their discovered attached +For two independent four-hour seeds, use `--machines 2 --hours 4`. + +The optional `--delete-old-experiments` flag deletes earlier `ww-long-`, `ww-mem-`, `ww-mem2-` and +`ww-v6e16-` and the earlier seven-day continuous-run queues/nodes in us-west4-a and us-east5-a, their discovered attached data disks plus `ww-full-data-20260929`, and snapshots of those disks. The exact inventory is saved to `~/continuous8-cleanup.json`. Other project resources, unrelated buckets and local/downloaded files are not deleted. Inventory or permission failures stop the operation. Re-running launch does not duplicate an -existing continuous8 queue or restart its training. +existing pilot queue or restart its training. If only one of two queue submissions +succeeds, the successful request is retained and status must be checked; there is +no automatic second attempt. Disk creation or quota errors can leave a dedicated +disk that must be inspected before retrying. ```bash python3 baseline/nanogpt_one_head/continuous8/cloudshell.py status @@ -106,8 +123,22 @@ the exact launch commit and starts setup via a systemd service with `Restart=no` A persistent claim blocks reruns after reboot. The TPU uses a dedicated service account with object-admin permission on this experiment bucket, not TPU-admin. Consequently it cannot delete itself: early completion/failure leaves the TPU -allocated until explicit deletion or the seven-day expiry. Review status promptly. +allocated until explicit deletion or the requested four/six-hour expiry. Review status promptly. At the published v5e Flex-start rate of $0.60/chip-hour, eight chips cost $4.80/hour, -or $806.40 for the seven-day cap, plus disk/bucket/network charges. Credits and +or **$28.80 for one six-hour machine**, **$38.40 for two four-hour machines**, +plus disk/bucket/network charges. Storage is retained after TPU expiry and +continues to incur charges until explicitly deleted. Credits and remaining balance must be checked in the project's billing account; this script does not assert that sufficient credits remain. + +## How far will the run get? + +`BENCHMARK_PROJECTION.json` reports measured training-only throughput and an +optimistic upper estimate based on the remaining allocation. It excludes +monitoring and uploads. `results/progress.json`, updated at evaluation, projects +from observed throughput including training-time evaluation/diagnostics/I/O. +Neither is a guarantee. With 32,768 token presentations/update, divide processed +tokens by 32,768 for update count, or by 5B for corpus-equivalent passes. The +sampler draws random windows; token presentations are not a count of unique +tokens visited. A larger corpus reduces repeated sampling but does not guarantee +better test accuracy or that alpha will fall below two within six hours. diff --git a/baseline/nanogpt_one_head/continuous8/cloudshell.py b/baseline/nanogpt_one_head/continuous8/cloudshell.py index cf10f55c..864d0f05 100644 --- a/baseline/nanogpt_one_head/continuous8/cloudshell.py +++ b/baseline/nanogpt_one_head/continuous8/cloudshell.py @@ -1,21 +1,20 @@ #!/usr/bin/env python3 -"""Run in Cloud Shell: scoped old-experiment cleanup and one v5e-8 request.""" +"""Cloud Shell: bounded v5e-8 experiments with data prepared before allocation.""" import argparse import json from pathlib import Path import subprocess +import sys import tempfile PROJECT = 'tpu-builders-504820' ZONE = 'us-west4-a' -QUEUE = 'ww-continuous8-20261002' -NODE = QUEUE + '-node' -DISK = QUEUE + '-data' +RUN_PREFIX = 'ww-continuous8-pilot-20261002' BUCKET = PROJECT + '-ww-continuous8' SA_ID = 'rg-continuous-tpu' SA = SA_ID + '@' + PROJECT + '.iam.gserviceaccount.com' -OLD_PREFIXES = ('ww-long-', 'ww-mem2-', 'ww-mem-', 'ww-v6e16-') -OLD_DISKS = {'ww-full-data-20260929'} +OLD_PREFIXES = ('ww-long-', 'ww-mem2-', 'ww-mem-', 'ww-v6e16-', 'ww-continuous8-20261002') +OLD_DISKS = {'ww-full-data-20260929', 'ww-continuous8-20261002-data'} def gc(*args, capture=False): @@ -55,7 +54,8 @@ def cleanup(): raise RuntimeError('Unexpected old disk project/zone') OLD_DISKS.add(source.rsplit('/', 1)[-1]) for d in inventory('compute', 'disks', 'list'): - if short(d) in OLD_DISKS: + if short(d) in OLD_DISKS or old(short(d)): + OLD_DISKS.add(short(d)) plan['disks'].append((d['zone'].rsplit('/',1)[-1], short(d))) for s in inventory('compute', 'snapshots', 'list'): if s.get('sourceDisk','').rsplit('/',1)[-1] in OLD_DISKS or short(s).startswith('ww-long-saved-'): @@ -84,23 +84,28 @@ def cleanup(): def status(): - detail = inventory('alpha','compute','tpus','queued-resources','describe',QUEUE,'--zone='+ZONE) - print(json.dumps(detail, indent=2)) - print(f'Cloud results: gs://{BUCKET}/runs/{QUEUE}/') - if detail.get('state',{}).get('state') == 'ACTIVE': - gc('compute','tpus','tpu-vm','ssh',NODE,'--zone='+ZONE,'--worker=0', - '--command=systemctl --no-pager status rg-continuous8.service; tail -n 35 /mnt/disks/rg-data/continuous8/run.log') - - -def provision(): + requests = inventory('alpha','compute','tpus','queued-resources','list','--zone='+ZONE) + for item in requests: + if not short(item).startswith(RUN_PREFIX): + continue + name = short(item) + detail = inventory('alpha','compute','tpus','queued-resources','describe',name,'--zone='+ZONE) + print(json.dumps(detail, indent=2)) + print(f'Cloud results: gs://{BUCKET}/runs/{name}/') + if detail.get('state',{}).get('state') == 'ACTIVE': + gc('compute','tpus','tpu-vm','ssh',name+'-node','--zone='+ZONE,'--worker=0', + '--command=systemctl --no-pager status rg-continuous8.service; tail -n 35 /mnt/disks/rg-data/continuous8/run.log') + + +def provision(machines, hours): root = Path(__file__).resolve().parents[3] commit = subprocess.check_output(['git','-C',str(root),'rev-parse','HEAD'],text=True).strip() dirty = subprocess.check_output(['git','-C',str(root),'status','--porcelain'],text=True).strip() if dirty: raise RuntimeError('Provision from a clean checked-in source tree') requests = inventory('alpha','compute','tpus','queued-resources','list','--zone='+ZONE) - if any(short(q)==QUEUE for q in requests): - print('Existing continuous-run request retained; no second run or restart.') + if any(short(q).startswith(RUN_PREFIX) for q in requests): + print('Existing pilot request retained; no additional allocation or restart.') status() return gc('services','enable','tpu.googleapis.com','compute.googleapis.com','storage.googleapis.com','iam.googleapis.com') @@ -112,43 +117,72 @@ def provision(): gc('iam','service-accounts','create',SA_ID,'--display-name=Continuous MuonClip TPU storage') gc('storage','buckets','add-iam-policy-binding','gs://'+BUCKET, '--member=serviceAccount:'+SA,'--role=roles/storage.objectAdmin') + seeds = (1337, 2027)[:machines] + runs = [(seed, f'{RUN_PREFIX}-s{seed}') for seed in seeds] disks = inventory('compute','disks','list') - if any(short(d)==DISK for d in disks): + if any(short(d).startswith(RUN_PREFIX) for d in disks): raise RuntimeError('Dedicated run disk exists without its queue. Keep it intact and inspect the previous attempt before retrying.') - gc('compute','disks','create',DISK,'--zone='+ZONE,'--size=500GB','--type=pd-balanced', - '--labels=experiment=continuous8') - uri = f'gs://{BUCKET}/runs/{QUEUE}' + + # Tokenization happens once on Cloud Shell CPU, outside the TPU budget. + temp = Path(tempfile.gettempdir())/'rg-continuous8-cloudshell' + temp.mkdir(exist_ok=True) + py = temp/'venv/bin/python' + if not py.exists(): + subprocess.run([sys.executable,'-m','venv',str(temp/'venv')], check=True) + subprocess.run([str(py),'-m','pip','install','--disable-pip-version-check','--no-cache-dir', + 'numpy==1.26.4','datasets==3.3.2','tiktoken==0.9.0','PyYAML==6.0.2'], check=True) + exp = Path(__file__).resolve().parents[1] + data_uri = f'gs://{BUCKET}/corpora/{RUN_PREFIX}' + print(f'Plan: {machines} machine(s), {hours}h maximum each; TPU compute ${machines*hours*8*0.60:.2f} plus storage.',flush=True) + subprocess.run([str(py),str(Path(__file__).with_name('prepare_cloud_data.py')), + '--config',str(exp/'configs/muonclip_continuous8.yaml'), + '--output-dir',str(temp/'data'),'--gcs-uri',data_uri],check=True) + template = Path(__file__).with_name('startup.sh').read_text() - startup = template.replace('__COMMIT__',commit).replace('__GCS_URI__',uri) - state = dict(project=PROJECT,zone=ZONE,queue=QUEUE,node=NODE,disk=DISK,gcs_uri=uri,commit=commit) + state = dict(project=PROJECT,zone=ZONE,commit=commit,hours=hours,machines=machines, + data_uri=data_uri,runs=[]) (Path.home()/'continuous8-resources.json').write_text(json.dumps(state,indent=2)) - with tempfile.NamedTemporaryFile(mode='w',suffix='.sh') as f: - f.write(startup); f.flush() - gc('alpha','compute','tpus','queued-resources','create',QUEUE, - '--zone='+ZONE,'--node-id='+NODE,'--accelerator-type=v5litepod-8', - '--runtime-version=v2-alpha-tpuv5-lite','--provisioning-model=flex-start', - '--max-run-duration=7d','--valid-until-duration=24h', - '--service-account='+SA,'--scopes=https://www.googleapis.com/auth/cloud-platform', - '--data-disk=source=projects/'+PROJECT+'/zones/'+ZONE+'/disks/'+DISK+',mode=read-write', - '--metadata-from-file=startup-script='+f.name,'--labels=experiment=continuous8', - '--quiet','--async') - print('Submitted one continuous8 request. It starts setup and preflight automatically when allocated.') + for seed, queue in runs: + node, disk = queue+'-node', queue+'-data' + uri = f'gs://{BUCKET}/runs/{queue}' + gc('compute','disks','create',disk,'--zone='+ZONE,'--size=200GB','--type=pd-balanced', + '--labels=experiment=continuous8') + startup = (template.replace('__COMMIT__',commit).replace('__GCS_URI__',uri) + .replace('__DATA_URI__',data_uri).replace('__SEED__',str(seed)) + .replace('__STOP_HOURS__',str(hours-0.5))) + state['runs'].append(dict(seed=seed,queue=queue,node=node,disk=disk,gcs_uri=uri)) + (Path.home()/'continuous8-resources.json').write_text(json.dumps(state,indent=2)) + with tempfile.NamedTemporaryFile(mode='w',suffix='.sh') as f: + f.write(startup); f.flush() + gc('alpha','compute','tpus','queued-resources','create',queue, + '--zone='+ZONE,'--node-id='+node,'--accelerator-type=v5litepod-8', + '--runtime-version=v2-alpha-tpuv5-lite','--provisioning-model=flex-start', + f'--max-run-duration={hours}h','--valid-until-duration=24h', + '--service-account='+SA,'--scopes=https://www.googleapis.com/auth/cloud-platform', + '--data-disk=source=projects/'+PROJECT+'/zones/'+ZONE+'/disks/'+disk+',mode=read-write', + '--metadata-from-file=startup-script='+f.name,'--labels=experiment=continuous8', + '--quiet','--async') + print('Submitted '+queue+'; cloud results: '+uri,flush=True) + print('Setup and preflight start automatically when each machine is allocated.') print('Inspect with: python3 baseline/nanogpt_one_head/continuous8/cloudshell.py status') print('Run settings: '+str(Path.home()/'continuous8-resources.json')) - print('Cloud results: '+uri) def main(): parser = argparse.ArgumentParser(description=__doc__) parser.add_argument('action', choices=['launch','status']) parser.add_argument('--delete-old-experiments', action='store_true') + parser.add_argument('--machines', type=int, choices=[1,2], default=1) + parser.add_argument('--hours', type=int, choices=[4,6], default=6) args = parser.parse_args() if args.action == 'status': status() else: + if (args.machines, args.hours) not in ((1,6),(2,4)): + parser.error('Supported budgets: --machines 1 --hours 6 or --machines 2 --hours 4') if args.delete_old_experiments: cleanup() - provision() + provision(args.machines, args.hours) if __name__ == '__main__': main() diff --git a/baseline/nanogpt_one_head/continuous8/prepare_cloud_data.py b/baseline/nanogpt_one_head/continuous8/prepare_cloud_data.py new file mode 100644 index 00000000..c2ce2fdd --- /dev/null +++ b/baseline/nanogpt_one_head/continuous8/prepare_cloud_data.py @@ -0,0 +1,59 @@ +#!/usr/bin/env python3 +"""Prepare the shared corpus on Cloud Shell before any TPU is requested.""" +import argparse +import hashlib +import importlib.util +import json +import os +from pathlib import Path +import shutil +import subprocess + +import yaml + +PROJECT = 'tpu-builders-504820' + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument('--config', required=True) + parser.add_argument('--output-dir', required=True) + parser.add_argument('--gcs-uri', required=True) + args = parser.parse_args() + cfg = yaml.safe_load(Path(args.config).read_text()) + output = Path(args.output_dir) + output.mkdir(parents=True, exist_ok=True) + os.environ['HF_HOME'] = str(output.parent/'hf-cache') + # Use the very same split writer and validator as training, without importing + # the package's torch/plotting exports into the small Cloud Shell environment. + source = Path(__file__).resolve().parents[1]/'src/rg_nanogpt_one_head/data.py' + spec = importlib.util.spec_from_file_location('corpus_writer', source) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + try: + module.validate_prepared_data(output, cfg) + except FileNotFoundError: + needed = 2*sum(int(cfg['dataset'][s+'_tokens']) for s in ('train','val','test')) + if shutil.disk_usage(output).free < needed + 3*1024**3: + raise RuntimeError('Insufficient Cloud Shell temporary disk space for corpus and download buffers') + print('Preparing shared corpus on Cloud Shell CPU. No TPUs have been requested.', flush=True) + module.prepare_fineweb_edu(cfg, output) + receipts = {} + for name in ('train.bin', 'val.bin', 'test.bin', 'meta.json'): + path = output/name + subprocess.run(['gcloud','storage','cp',str(path),args.gcs_uri+'/'+name, + '--project='+PROJECT], check=True) + receipts[name] = {'bytes':path.stat().st_size, 'sha256':module._sha256(path)} + manifest = {'files':receipts, 'dataset':cfg['dataset'], + 'dataset_config_sha256':hashlib.sha256(json.dumps(cfg['dataset'],sort_keys=True).encode()).hexdigest()} + complete = output/'COMPLETE.json' + complete.write_text(json.dumps(manifest, indent=2)+'\n') + # Gcloud verifies transfers; publish the completion marker only after all + # files succeeded. TPU workers also verify SHA256 after downloading. + subprocess.run(['gcloud','storage','cp',str(complete),args.gcs_uri+'/COMPLETE.json', + '--project='+PROJECT], check=True) + print('Shared cloud corpus complete. TPU provisioning can now begin.', flush=True) + + +if __name__ == '__main__': + main() diff --git a/baseline/nanogpt_one_head/continuous8/startup.sh b/baseline/nanogpt_one_head/continuous8/startup.sh index 37911212..ea77a354 100644 --- a/baseline/nanogpt_one_head/continuous8/startup.sh +++ b/baseline/nanogpt_one_head/continuous8/startup.sh @@ -1,6 +1,9 @@ #!/usr/bin/env bash # Provisioner substitutes the pinned Git SHA and the dedicated GCS run URI. set -Eeuo pipefail +# Stop training 30+ minutes before the server-enforced allocation limit. +# VM uptime includes setup; a slow setup cannot extend the requested budget. +DEADLINE=$(python3 -c 'import time; print(time.time()-float(open("/proc/uptime").read().split()[0])+__STOP_HOURS__*3600)') DEVICE=/dev/disk/by-id/google-persistent-disk-1 for attempt in $(seq 1 60); do [ -b "$DEVICE" ] && break @@ -41,8 +44,12 @@ Wants=network-online.target [Service] Type=simple Environment=RG_CONTINUOUS_GCS_URI=__GCS_URI__ +Environment=RG_CONTINUOUS_DATA_URI=__DATA_URI__ +Environment=RG_CONTINUOUS_SEED=__SEED__ +Environment=RG_CONTINUOUS_DEADLINE_UNIX=$DEADLINE ExecStart=/bin/bash $BASE/repo/baseline/nanogpt_one_head/continuous8/worker.sh Restart=no +SuccessExitStatus=75 KillSignal=SIGTERM TimeoutStopSec=1800 StandardOutput=append:$BASE/run.log diff --git a/baseline/nanogpt_one_head/continuous8/worker.sh b/baseline/nanogpt_one_head/continuous8/worker.sh index e610a6db..752efaa6 100644 --- a/baseline/nanogpt_one_head/continuous8/worker.sh +++ b/baseline/nanogpt_one_head/continuous8/worker.sh @@ -21,46 +21,70 @@ PY="$BASE/venv/bin/python" "$PY" -m pip install --no-deps -e . "$PY" -m pip freeze > "$BASE/environment.lock.txt" git rev-parse HEAD > "$BASE/source_commit.txt" -# Upload credentials test before corpus preparation or training. +# Upload credentials test before data download or training. "$PY" - <<'PY' import os from pathlib import Path from rg_nanogpt_one_head.continuous_support import CloudPublisher p=CloudPublisher(os.environ['RG_CONTINUOUS_GCS_URI']) base=Path('/mnt/disks/rg-data/continuous8') -p.claim({'commit':(base/'source_commit.txt').read_text().strip(), 'automatic_restart':False}) +p.claim({'commit':(base/'source_commit.txt').read_text().strip(), + 'seed':int(os.environ['RG_CONTINUOUS_SEED']), 'automatic_restart':False}) for name in ('environment.lock.txt','source_commit.txt'): p.file(base/name, name) p.json({'status':'preflight'}, 'SETUP_STATUS.json') PY +# Download the one shared CPU-prepared corpus and verify every file hash. +"$PY" - <<'PY' +import json, os +from pathlib import Path +import yaml +from rg_nanogpt_one_head.continuous_support import CloudPublisher, sha_file +from rg_nanogpt_one_head.data import validate_prepared_data +source=CloudPublisher(os.environ['RG_CONTINUOUS_DATA_URI']) +manifest=json.loads(source.bucket.blob(source.prefix+'/COMPLETE.json').download_as_text()) +cfg=yaml.safe_load(Path('configs/muonclip_continuous8.yaml').read_text()) +if manifest['dataset'] != cfg['dataset']: + raise RuntimeError('Shared cloud corpus does not match experiment configuration') +base=Path('/mnt/disks/rg-data/continuous8/data') +base.mkdir(exist_ok=True) +for name in ('train.bin','val.bin','test.bin','meta.json'): + path=base/name + source.bucket.blob(source.prefix+'/'+name).download_to_filename(str(path), checksum='crc32c', timeout=600) + entry=manifest['files'][name] + if path.stat().st_size != entry['bytes'] or sha_file(path) != entry['sha256']: + raise RuntimeError('Shared data checksum failed: '+name) +validate_prepared_data(base,cfg) +p=CloudPublisher(os.environ['RG_CONTINUOUS_GCS_URI']) +p.json({'uri':os.environ['RG_CONTINUOUS_DATA_URI'], 'manifest':manifest}, 'DATA_SOURCE.json') +print('Shared data downloaded and SHA256 verified.', flush=True) +PY # Tiny, independent test: global gradients, clipping, metrics, optimizer + RNG restore. # Also measures the full proposed model shape before the long run. "$PY" -m rg_nanogpt_one_head.tpu_spmd_check --backend tpu --chips 8 \ --benchmark-config configs/muonclip_continuous8.yaml --benchmark-steps 10 \ --output "$BASE/preflight.json" "$PY" - <<'PY' -import os -from rg_nanogpt_one_head.continuous_support import CloudPublisher -CloudPublisher(os.environ['RG_CONTINUOUS_GCS_URI']).file('/mnt/disks/rg-data/continuous8/preflight.json', 'preflight.json') -PY -# Pinned dataset revision; exact document-disjoint split sizes and content hashes. -"$PY" -m rg_nanogpt_one_head.data --config configs/muonclip_continuous8.yaml \ - --output-dir "$BASE/data" -"$PY" - <<'PY' -import os +import json, os, time from pathlib import Path -from rg_nanogpt_one_head.continuous_support import CloudPublisher, sha_file +from rg_nanogpt_one_head.continuous_support import CloudPublisher p=CloudPublisher(os.environ['RG_CONTINUOUS_GCS_URI']) -base=Path('/mnt/disks/rg-data/continuous8/data') -receipts={} -for path in sorted(base.iterdir()): - if path.is_file(): - receipts[path.name]=p.file(path, 'data/'+path.name) - receipts[path.name]['sha256']=sha_file(path) -p.json(receipts, 'data/COMPLETE.json') +path=Path('/mnt/disks/rg-data/continuous8/preflight.json') +p.file(path, 'preflight.json') +report=json.loads(path.read_text()) +remaining=max(0,float(os.environ['RG_CONTINUOUS_DEADLINE_UNIX'])-time.time()) +rate=report['benchmark']['tokens_per_second'] +projection={'remaining_training_window_seconds':remaining, + 'benchmark_tokens_per_second':rate, + 'optimistic_token_presentations':int(rate*remaining), + 'optimistic_updates':int(rate*remaining/report['benchmark']['global_tokens_per_update']), + 'caveat':'Synthetic training throughput; excludes evaluation, spectra, initialization and checkpoint upload overhead.'} +p.json(projection, 'BENCHMARK_PROJECTION.json') +print('BENCHMARK PROJECTION:', json.dumps(projection), flush=True) p.json({'status':'training'}, 'SETUP_STATUS.json') PY # No resilient supervisor, --resume, or automatic retry anywhere in this path. "$PY" -u -m rg_nanogpt_one_head.continuous_run \ --config configs/muonclip_continuous8.yaml --data-root "$BASE/data" \ - --results-root "$BASE/results" --device tpu + --results-root "$BASE/results" --device tpu --seed "$RG_CONTINUOUS_SEED" \ + --deadline-unix "$RG_CONTINUOUS_DEADLINE_UNIX" diff --git a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/continuous_run.py b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/continuous_run.py index ec4e3cb3..88ca77e8 100644 --- a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/continuous_run.py +++ b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/continuous_run.py @@ -16,6 +16,8 @@ def main(): parser.add_argument('--data-root', required=True) parser.add_argument('--results-root', required=True) parser.add_argument('--device', default='tpu', choices=['tpu', 'cpu']) + parser.add_argument('--seed', type=int, default=1337) + parser.add_argument('--deadline-unix', type=float) args = parser.parse_args() from .muonclip import install_muonclip_extension install_muonclip_extension() @@ -24,12 +26,20 @@ def main(): cfg = load_config(args.config) if not cfg.get('continuous', {}).get('enabled') or cfg.get('continuation'): raise ValueError('A fresh continuous-run config is required') + if args.seed not in cfg['training']['seeds']: + raise ValueError('Seed must be one of the registered training seeds') + seconds = float(cfg['continuous']['max_wall_hours']) * 3600 + if args.deadline_unix is not None: + seconds = min(seconds, args.deadline_unix - time.time()) + if seconds <= 0: + raise RuntimeError('Allocation training deadline passed during setup; refusing to start') root = Path(args.results_root) root.mkdir(parents=True, exist_ok=True) # An exclusive claim survives VM reboots and remains even after failure. with (root/'CONTINUOUS_STARTED.json').open('x') as f: import json json.dump(dict(pid=os.getpid(), started_unix=time.time(), start_step=0, + seed=args.seed, deadline_unix=args.deadline_unix, automatic_restart=False), f) stop = root/'STOP' cfg['training']['stop_file'] = str(stop.resolve()) @@ -41,13 +51,14 @@ def request_stop(*unused): stop.touch() for sig in (signal.SIGTERM, signal.SIGINT): signal.signal(sig, request_stop) - timer = threading.Timer(float(cfg['continuous']['max_wall_hours'])*3600, request_stop) + timer = threading.Timer(seconds, request_stop) timer.daemon = True timer.start() - status = {'status':'running', 'pid':os.getpid(), 'start_step':0, 'restarts':0} + status = {'status':'running', 'pid':os.getpid(), 'start_step':0, 'restarts':0, + 'seed':args.seed, 'deadline_unix':args.deadline_unix} try: run_optimizer_replicates(cfg=cfg, config_path=args.config, - optimizer_name='muon_clip', seeds=(1337,), data_root=args.data_root, + optimizer_name='muon_clip', seeds=(args.seed,), data_root=args.data_root, results_root=root, device=args.device, resume=False, overwrite=False) status['status'] = 'completed' except SystemExit as exc: @@ -62,7 +73,7 @@ def request_stop(*unused): atomic_json(root/'CONTINUOUS_STATUS.json', status) if sink: sink.json(status, 'CONTINUOUS_STATUS.json') - run_dir = root/'muon_clip'/'seed_1337' + run_dir = root/'muon_clip'/f'seed_{args.seed}' if run_dir.exists(): publish_metadata(cfg, run_dir) diff --git a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/continuous_support.py b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/continuous_support.py index 29fe7842..52c9c47a 100644 --- a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/continuous_support.py +++ b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/continuous_support.py @@ -243,4 +243,15 @@ def record_token_errors(cfg, run_dir, row): w = csv.DictWriter(f, fieldnames=list(result)) if not exists: w.writeheader() w.writerow(result) + if row['step'] > 0 and row['elapsed_sec'] > 0: + remaining = max(0, float(os.environ.get('RG_CONTINUOUS_DEADLINE_UNIX', time.time()))-time.time()) + rate = row['tokens_seen']/row['elapsed_sec'] + progress = dict(step=row['step'], tokens_seen=row['tokens_seen'], + measured_tokens_per_second=rate, + remaining_training_window_seconds=remaining, + projected_total_token_presentations=int(row['tokens_seen']+remaining*rate), + caveat='Projection assumes measured throughput including monitoring stays constant.') + atomic_json(Path(run_dir)/'progress.json', progress) + print(f"[continuous-progress] step={row['step']} measured={rate:,.0f} tokens/s; " + f"projected total={progress['projected_total_token_presentations']:,} token presentations", flush=True) publish_metadata(cfg, run_dir) diff --git a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/data.py b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/data.py index 8c7a6b68..beb18d83 100644 --- a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/data.py +++ b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/data.py @@ -11,8 +11,6 @@ import numpy as np -from .config import load_config, roots - TOKEN_DTYPE = np.dtype(np.uint16) SPLIT_NAMES = ("train", "val", "test") @@ -327,6 +325,8 @@ def load_memmaps( def main() -> None: + from .config import load_config, roots + parser = argparse.ArgumentParser( description="Prepare the pinned FineWeb-Edu one-head baseline corpus" ) diff --git a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/train_loop.py b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/train_loop.py index eae7e7a0..54de48d8 100644 --- a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/train_loop.py +++ b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/train_loop.py @@ -568,7 +568,12 @@ def execute_training_loop( last_update_lrs = dict(next_update_lrs) new_step = completed_steps + 1 - if _resume_diagnostics_due( + continuous_stop = bool( + cfg.get('continuous', {}).get('enabled') + and cfg['training'].get('stop_file') + and Path(cfg['training']['stop_file']).exists() + ) + if continuous_stop or _resume_diagnostics_due( new_step, cfg=cfg, epoch_steps=epoch_steps, @@ -588,7 +593,7 @@ def execute_training_loop( else False ) - checkpoint_due = _checkpoint_due( + checkpoint_due = continuous_stop or _checkpoint_due( new_step, cfg=cfg, epoch_steps=epoch_steps, diff --git a/baseline/nanogpt_one_head/tests/test_continuous8.py b/baseline/nanogpt_one_head/tests/test_continuous8.py index b74e65e7..c25b124f 100644 --- a/baseline/nanogpt_one_head/tests/test_continuous8.py +++ b/baseline/nanogpt_one_head/tests/test_continuous8.py @@ -1,4 +1,5 @@ import csv +import math from pathlib import Path import numpy as np import pytest @@ -84,7 +85,7 @@ def test_fixed_config_alignment_and_schedule(): cfg=load_config(root/'configs/muonclip_continuous8.yaml') assert cfg['model']['n_layer']==12 and cfg['runtime']['tpu_expected_chips']==8 assert cfg['dataset']['train_tokens']==5_000_000_000 and 'continuation' not in cfg - assert list(epoch_step_map(cfg))==list(range(0,1_000_001,2000)) + assert list(epoch_step_map(cfg))==list(range(0,1_000_001,1000)) p=optimizer_profile(cfg,'muon_clip') assert lr_schedule_steps(cfg,p)==100_000 and warmup_steps(p,100_000)==2000 @@ -138,7 +139,8 @@ def advance(model,opts,generator,start,end): assert evaluate_probe(net,probe,torch.device('cpu'))==evaluate_probe(resumed,probe,torch.device('cpu')) -def test_fresh_continuous_engine_writes_fixed_probe_and_pairs(tmp_path,monkeypatch): +@pytest.mark.parametrize('stop_early',[False,True]) +def test_fresh_continuous_engine_writes_fixed_probe_and_pairs(tmp_path,monkeypatch,stop_early): from copy import deepcopy from rg_nanogpt_one_head.muonclip import install_muonclip_extension install_muonclip_extension() @@ -174,6 +176,26 @@ def spectrum(model,*a,**kw): 'alpha_clip_xmax_n':12,'alpha_clip_xmax_mean':2.9,'alpha_clip_xmax_min':2.4} monkeypatch.setattr(loop,'run_weightwatcher',spectrum) monkeypatch.setattr(utils,'evaluate_bleu',lambda *a,**kw:{'bleu':0.}) + if stop_early: + # A deadline between regular checkpoints must save after the next update, + # with current finite diagnostics, rather than wait 500 more steps. + stop=tmp_path/'STOP' + cfg['training']['stop_file']=str(stop) + update=loop.optimizer_step + def request_after_update(*a,**kw): + result=update(*a,**kw) + stop.touch() + return result + monkeypatch.setattr(loop,'optimizer_step',request_after_update) + from rg_nanogpt_one_head.continuation import TrainingPaused + with pytest.raises(TrainingPaused): + run_one(cfg=cfg,data_root=data,results_root=tmp_path/'results',optimizer_name='muon_clip', + seed=2027,device='cpu',resume=False,progress=False) + checkpoint=torch.load(tmp_path/'results/muon_clip/seed_2027/checkpoint_latest.pt',weights_only=False) + assert checkpoint['step']==1 + assert checkpoint['seed']==2027 + assert math.isfinite(checkpoint['resume_diagnostics']['last_grad_pre']) + return # Stop at the full step-4 checkpoint: this tests the actual monitoring loop # without post-run spectral completion checks against our synthetic spectrum. original=loop.save_training_checkpoint diff --git a/baseline/nanogpt_one_head/tests/test_continuous8_launcher.py b/baseline/nanogpt_one_head/tests/test_continuous8_launcher.py new file mode 100644 index 00000000..b36314f9 --- /dev/null +++ b/baseline/nanogpt_one_head/tests/test_continuous8_launcher.py @@ -0,0 +1,59 @@ +"""Exercise provisioning without network access or real cloud mutations.""" +import importlib.util +from pathlib import Path +import subprocess + +import pytest + + +@pytest.mark.parametrize('machines,hours', [(1,6),(2,4)]) +def test_cloud_budget_and_data_before_allocation(tmp_path, monkeypatch, machines, hours): + file = Path(__file__).resolve().parents[1]/'continuous8/cloudshell.py' + spec = importlib.util.spec_from_file_location('pilot_launcher', file) + m = importlib.util.module_from_spec(spec) + spec.loader.exec_module(m) + monkeypatch.setattr(Path, 'home', classmethod(lambda cls: tmp_path)) + monkeypatch.setattr(m.tempfile,'gettempdir',lambda:str(tmp_path)) + monkeypatch.setattr(m.subprocess,'check_output', + lambda cmd,**kw: '\n' if 'status' in cmd else 'abcdef123\n') + def inventory(*args): + if args[:3] == ('storage','buckets','list'): return [{'name':m.BUCKET}] + if args[:3] == ('iam','service-accounts','list'): return [{'email':m.SA}] + return [] + monkeypatch.setattr(m,'inventory',inventory) + prepared = False + calls = [] + def run(cmd,**kw): + nonlocal prepared + if any(str(x).endswith('prepare_cloud_data.py') for x in cmd): prepared = True + return subprocess.CompletedProcess(cmd,0) + monkeypatch.setattr(m.subprocess,'run',run) + def gc(*args,**kw): + calls.append(args) + if args[:5] == ('alpha','compute','tpus','queued-resources','create'): + assert prepared, 'No TPU allocation before CPU preparation succeeds' + assert f'--max-run-duration={hours}h' in args + source = next(x.split('=',2)[-1] for x in args if x.startswith('--metadata-from-file=')) + startup = Path(source).read_text() + assert '__' not in startup + assert f'+{hours-0.5}*3600' in startup + assert 'Restart=no' in startup + assert 'RG_CONTINUOUS_DATA_URI=gs://' in startup + monkeypatch.setattr(m,'gc',gc) + m.provision(machines,hours) + create = [c for c in calls if c[:5] == ('alpha','compute','tpus','queued-resources','create')] + assert len(create) == machines + assert len({c[5] for c in create}) == machines + assert all('--accelerator-type=v5litepod-8' in c for c in create) + + +def test_dataset_writer_import_requires_no_torch(): + import sys + source = Path(__file__).resolve().parents[1]/'src/rg_nanogpt_one_head/data.py' + code = ( + 'import importlib.util,sys; ' + f's=importlib.util.spec_from_file_location("writer",{str(source)!r}); ' + 'm=importlib.util.module_from_spec(s); s.loader.exec_module(m); ' + 'assert "torch" not in sys.modules' + ) + subprocess.run([sys.executable,'-c',code],check=True) From 3393b14e990f2119ab94da11a64aebd9e764b9c9 Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Fri, 2 Oct 2026 16:33:01 -0700 Subject: [PATCH 07/36] Run large-corpus preparation on TPU VM and make launch failures visible Submit the bounded TPU request from Cloud Shell without a long /tmp preparation job. Use the VM CPU and persistent disk for ordered parallel tokenization, validate and archive the corpus, then train in one continuous process. Preserve the six-hour allocation cap and no-restart guard. Save launch phases/errors in Cloud Shell HOME and a transcript via run.sh; report empty status explicitly. Validation: 27 CPU tests passed, one hardware check skipped. Parallel and serial split files match byte for byte; tested budget flags and failure records. --- baseline/nanogpt_one_head/README.md | 5 +- .../configs/muonclip_continuous8.yaml | 2 + .../nanogpt_one_head/continuous8/README.md | 36 ++++- .../continuous8/cloudshell.py | 137 +++++++++++++++--- .../continuous8/prepare_tpu_data.py | 51 +++++++ baseline/nanogpt_one_head/continuous8/run.sh | 7 + .../nanogpt_one_head/continuous8/worker.sh | 52 +++---- .../src/rg_nanogpt_one_head/data.py | 24 ++- .../tests/test_continuous8_launcher.py | 60 +++++++- 9 files changed, 314 insertions(+), 60 deletions(-) create mode 100644 baseline/nanogpt_one_head/continuous8/prepare_tpu_data.py create mode 100644 baseline/nanogpt_one_head/continuous8/run.sh diff --git a/baseline/nanogpt_one_head/README.md b/baseline/nanogpt_one_head/README.md index ad35a446..2b7f0d64 100644 --- a/baseline/nanogpt_one_head/README.md +++ b/baseline/nanogpt_one_head/README.md @@ -380,5 +380,6 @@ For the fresh 124M-parameter MuonClip experiment with 5B training tokens, fixed document token-error probes, raw/clipped alpha and synchronous Cloud Storage checkpoints, see [continuous8/README.md](continuous8/README.md). This launcher defaults to one six-hour v5e-8 allocation, with a two-machine/four-hour option. -Data preparation happens on Cloud Shell before TPU allocation; automatic -restarts and the segmented continuation runner are disabled. +Data preparation and training run on the TPU VM and its persistent disk, +independently of Cloud Shell. Automatic restarts and the segmented continuation +runner are disabled; the allocation cap includes setup and tokenization. diff --git a/baseline/nanogpt_one_head/configs/muonclip_continuous8.yaml b/baseline/nanogpt_one_head/configs/muonclip_continuous8.yaml index c9346ce3..c4bcf102 100644 --- a/baseline/nanogpt_one_head/configs/muonclip_continuous8.yaml +++ b/baseline/nanogpt_one_head/configs/muonclip_continuous8.yaml @@ -9,6 +9,8 @@ dataset: split: train revision: 593b3a867298afb8ce42625a270ef20ddcad28f9 tokenizer: gpt2 + encoding_workers: 16 + encoding_batch_size: 256 train_tokens: 5000000000 val_tokens: 10000000 test_tokens: 10000000 diff --git a/baseline/nanogpt_one_head/continuous8/README.md b/baseline/nanogpt_one_head/continuous8/README.md index 84d18576..b33ee84b 100644 --- a/baseline/nanogpt_one_head/continuous8/README.md +++ b/baseline/nanogpt_one_head/continuous8/README.md @@ -62,11 +62,14 @@ It is a monitored test set, not an untouched final confirmation set. A dedicated 200 GB persistent disk is attached to each host and mounted at `/mnt/disks/rg-data`. GCS is the durable experiment archive: `gs://tpu-builders-504820-ww-continuous8/runs/ww-continuous8-pilot-20261002-sSEED/`. -The dataset is prepared once on Cloud Shell CPU and uploaded **before any TPU -is requested**. This can take hours and needs about 13 GB free under `/tmp`; -it is outside the TPU allocation budget. Both machines download the same -SHA256-verified files. Keep Cloud Shell open during preparation. No additional -CPU VM is provisioned by this launcher. Code commit, +The launcher requests the TPU immediately after resource checks. Data preparation +runs on the TPU VM CPU, using 16 encoding threads and the attached persistent +disk; Cloud Shell is only the submission client. Ordered parallel encoding +preserves the serial writer's token bytes and document-disjoint splits. The corpus +is uploaded and SHA256-verified before scientific training starts. This work +counts against the four/six-hour allocation. No additional CPU VM is created. +Each seed archives its corpus under a separate prefix; corpus hashes must match +before comparing independent seeds. Code commit, resolved dependency versions, preflight report, fixed probes, metrics, spectra, plots and logs are recorded. Full checkpoints include optimizer buffers/counters, RNG/sampler state, learning-rate/config identity and monitoring state. @@ -98,7 +101,7 @@ allocated hardware before the launch script starts scientific training. From a clean checkout of `codex/continuous-muonclip-8`, on Cloud Shell: ```bash -python3 baseline/nanogpt_one_head/continuous8/cloudshell.py launch --machines 1 --hours 6 +bash baseline/nanogpt_one_head/continuous8/run.sh launch --machines 1 --hours 6 ``` For two independent four-hour seeds, use `--machines 2 --hours 4`. @@ -142,3 +145,24 @@ tokens by 32,768 for update count, or by 5B for corpus-equivalent passes. The sampler draws random windows; token presentations are not a count of unique tokens visited. A larger corpus reduces repeated sampling but does not guarantee better test accuracy or that alpha will fall below two within six hours. + +## Launch visibility and failure diagnosis + +The launcher records each phase and any exception in `~/continuous8-launch.json`. +The `run.sh` wrapper also saves the terminal transcript in +`~/continuous8-launch.log`. These small files live in Cloud Shell HOME; the large +corpus and training job do not. After the queued request is successfully submitted, +Cloud Shell may disconnect without stopping setup or training on the TPU VM. + +`cloudshell.py status` explicitly reports `NO TPU REQUEST` when none exists. +`cloudshell.py check` performs only read-only local/cloud diagnostics. VM setup +phase is recorded in `SETUP_STATUS.json`; `WORKER_STATUS.json` records worker +exit, including failures before training; `CONTINUOUS_STATUS.json` records the +scientific process outcome. If dependency installation fails before cloud upload +libraries are available, inspect the persistent disk's startup/run log over SSH. + +The earlier launcher did long preparation in Cloud Shell `/tmp`. No request/data +was found after one attempt; the initial error was not retained, so its exact +cause is unknown. Cloud Shell VM disposal can lose `/tmp`, and the old empty +status output did not distinguish preparation from failure. This version moves +the long work off Cloud Shell and makes those states explicit. diff --git a/baseline/nanogpt_one_head/continuous8/cloudshell.py b/baseline/nanogpt_one_head/continuous8/cloudshell.py index 864d0f05..d592eea6 100644 --- a/baseline/nanogpt_one_head/continuous8/cloudshell.py +++ b/baseline/nanogpt_one_head/continuous8/cloudshell.py @@ -1,11 +1,14 @@ #!/usr/bin/env python3 -"""Cloud Shell: bounded v5e-8 experiments with data prepared before allocation.""" +"""Submit bounded v5e-8 experiments; prepare data and train on the TPU VM.""" import argparse +from datetime import datetime, timezone import json +import os from pathlib import Path import subprocess import sys import tempfile +import traceback PROJECT = 'tpu-builders-504820' ZONE = 'us-west4-a' @@ -15,6 +18,85 @@ SA = SA_ID + '@' + PROJECT + '.iam.gserviceaccount.com' OLD_PREFIXES = ('ww-long-', 'ww-mem2-', 'ww-mem-', 'ww-v6e16-', 'ww-continuous8-20261002') OLD_DISKS = {'ww-full-data-20260929', 'ww-continuous8-20261002-data'} +_LAUNCH_RECORD = None + + +def boot_id(): + path = Path('/proc/sys/kernel/random/boot_id') + return path.read_text().strip() if path.exists() else 'unknown' + + +def launch_record(**changes): + """Small durable progress record in HOME; never put the corpus in HOME.""" + global _LAUNCH_RECORD + if _LAUNCH_RECORD is None: + _LAUNCH_RECORD = dict(pid=os.getpid(), boot_id=boot_id(), + started_utc=datetime.now(timezone.utc).isoformat()) + _LAUNCH_RECORD.update(changes, updated_utc=datetime.now(timezone.utc).isoformat()) + path = Path.home()/'continuous8-launch.json' + temp = path.with_suffix('.json.tmp') + temp.write_text(json.dumps(_LAUNCH_RECORD, indent=2)+'\n') + temp.replace(path) + if 'phase' in changes: + print('[launch] '+changes['phase'], flush=True) + + +def show_local_launch(): + path = Path.home()/'continuous8-launch.json' + if not path.exists(): + print('No persistent launch record. Earlier launcher versions did not save one.') + return + record = json.loads(path.read_text()) + alive = False + if record.get('boot_id') == boot_id() and isinstance(record.get('pid'),int): + try: + os.kill(record['pid'], 0) + cmdline = Path(f"/proc/{record['pid']}/cmdline").read_bytes() + alive = b'cloudshell.py' in cmdline and b'launch' in cmdline + except (ProcessLookupError, FileNotFoundError): + pass + if record.get('status') == 'running' and not alive: + record['observed_status'] = 'INTERRUPTED: recorded launcher process is no longer present' + print(json.dumps(record, indent=2)) + + +def check_environment(): + """Read-only diagnosis: no installs, cleanup, resource creation, or launch.""" + print('READ-ONLY CHECK. No TPUs or other resources will be created/deleted.', flush=True) + show_local_launch() + root = Path(__file__).resolve().parents[3] + print('Python:', sys.version.split()[0]) + print('Source:', subprocess.check_output(['git','-C',str(root),'rev-parse','HEAD'],text=True).strip()) + dirty = subprocess.check_output(['git','-C',str(root),'status','--porcelain'],text=True).strip() + print('Checkout:', 'DIRTY (launch will refuse)' if dirty else 'clean') + print('Data preparation runs on the TPU VM and its persistent disk, not Cloud Shell.') + data = Path(tempfile.gettempdir())/'rg-continuous8-cloudshell/data' + print('Local preparation:', 'absent' if not data.exists() else str(data)) + if data.exists(): + for p in sorted(data.iterdir()): + if p.is_file(): print(f' {p.name}: {p.stat().st_size:,} bytes') + errors = [] + checks = [ + ('TPU requests', ('alpha','compute','tpus','queued-resources','list','--zone='+ZONE)), + ('TPU VMs', ('compute','tpus','tpu-vm','list','--zone='+ZONE)), + ('Experiment bucket', ('storage','buckets','describe','gs://'+BUCKET)), + ('Experiment service account', ('iam','service-accounts','describe',SA)), + ] + for label,args in checks: + print(label+':',flush=True) + try: + value = inventory(*args) + if isinstance(value,list): + print(json.dumps([{'name':short(x),'state':x.get('state')} for x in value],indent=2)) + else: + print(json.dumps({k:value[k] for k in ('name','email','location','disabled') if k in value})) + except subprocess.CalledProcessError: + errors.append(label) + print('CHECK FAILED: '+label+' (see gcloud error above)',flush=True) + print('Persistent log:', Path.home()/'continuous8-launch.log') + print('Read-only checks finished; creation permissions and data preparation are not validated.') + if errors: + print('Failed checks:', ', '.join(errors)) def gc(*args, capture=False): @@ -84,7 +166,13 @@ def cleanup(): def status(): + show_local_launch() requests = inventory('alpha','compute','tpus','queued-resources','list','--zone='+ZONE) + if not any(short(item).startswith(RUN_PREFIX) for item in requests): + print('NO TPU REQUEST for this pilot. It is not training or waiting for TPU capacity.') + print('This launcher prepares data on the TPU VM: no request means the new pipeline has not started.') + print('Log: '+str(Path.home()/'continuous8-launch.log')) + return for item in requests: if not short(item).startswith(RUN_PREFIX): continue @@ -98,6 +186,7 @@ def status(): def provision(machines, hours): + launch_record(status='running',phase='checking source and existing requests', machines=machines, hours=hours) root = Path(__file__).resolve().parents[3] commit = subprocess.check_output(['git','-C',str(root),'rev-parse','HEAD'],text=True).strip() dirty = subprocess.check_output(['git','-C',str(root),'status','--porcelain'],text=True).strip() @@ -105,38 +194,35 @@ def provision(machines, hours): raise RuntimeError('Provision from a clean checked-in source tree') requests = inventory('alpha','compute','tpus','queued-resources','list','--zone='+ZONE) if any(short(q).startswith(RUN_PREFIX) for q in requests): + launch_record(status='existing_request',phase='existing TPU request; no new allocation') print('Existing pilot request retained; no additional allocation or restart.') status() return + launch_record(phase='enabling required APIs') gc('services','enable','tpu.googleapis.com','compute.googleapis.com','storage.googleapis.com','iam.googleapis.com') + launch_record(phase='checking or creating experiment bucket') buckets = inventory('storage','buckets','list') if not any(x.get('name','').removeprefix('gs://').rstrip('/')==BUCKET for x in buckets): gc('storage','buckets','create','gs://'+BUCKET,'--location=us-west4','--uniform-bucket-level-access') + launch_record(phase='checking or creating experiment service account') accounts = inventory('iam','service-accounts','list') if not any(x.get('email')==SA for x in accounts): gc('iam','service-accounts','create',SA_ID,'--display-name=Continuous MuonClip TPU storage') + launch_record(phase='granting experiment bucket access') gc('storage','buckets','add-iam-policy-binding','gs://'+BUCKET, '--member=serviceAccount:'+SA,'--role=roles/storage.objectAdmin') seeds = (1337, 2027)[:machines] runs = [(seed, f'{RUN_PREFIX}-s{seed}') for seed in seeds] + launch_record(phase='checking for previous experiment disks') disks = inventory('compute','disks','list') if any(short(d).startswith(RUN_PREFIX) for d in disks): raise RuntimeError('Dedicated run disk exists without its queue. Keep it intact and inspect the previous attempt before retrying.') - # Tokenization happens once on Cloud Shell CPU, outside the TPU budget. - temp = Path(tempfile.gettempdir())/'rg-continuous8-cloudshell' - temp.mkdir(exist_ok=True) - py = temp/'venv/bin/python' - if not py.exists(): - subprocess.run([sys.executable,'-m','venv',str(temp/'venv')], check=True) - subprocess.run([str(py),'-m','pip','install','--disable-pip-version-check','--no-cache-dir', - 'numpy==1.26.4','datasets==3.3.2','tiktoken==0.9.0','PyYAML==6.0.2'], check=True) - exp = Path(__file__).resolve().parents[1] + # Only submit from Cloud Shell. All long work is performed by the TPU VM's + # systemd service on its attached persistent disk, independent of this shell. data_uri = f'gs://{BUCKET}/corpora/{RUN_PREFIX}' print(f'Plan: {machines} machine(s), {hours}h maximum each; TPU compute ${machines*hours*8*0.60:.2f} plus storage.',flush=True) - subprocess.run([str(py),str(Path(__file__).with_name('prepare_cloud_data.py')), - '--config',str(exp/'configs/muonclip_continuous8.yaml'), - '--output-dir',str(temp/'data'),'--gcs-uri',data_uri],check=True) + print('Allocation time includes software setup, on-VM data preparation, preflight and training.', flush=True) template = Path(__file__).with_name('startup.sh').read_text() state = dict(project=PROJECT,zone=ZONE,commit=commit,hours=hours,machines=machines, @@ -145,15 +231,17 @@ def provision(machines, hours): for seed, queue in runs: node, disk = queue+'-node', queue+'-data' uri = f'gs://{BUCKET}/runs/{queue}' + launch_record(phase='creating run disk for seed '+str(seed)) gc('compute','disks','create',disk,'--zone='+ZONE,'--size=200GB','--type=pd-balanced', '--labels=experiment=continuous8') startup = (template.replace('__COMMIT__',commit).replace('__GCS_URI__',uri) - .replace('__DATA_URI__',data_uri).replace('__SEED__',str(seed)) + .replace('__DATA_URI__',data_uri+'-s'+str(seed)).replace('__SEED__',str(seed)) .replace('__STOP_HOURS__',str(hours-0.5))) state['runs'].append(dict(seed=seed,queue=queue,node=node,disk=disk,gcs_uri=uri)) (Path.home()/'continuous8-resources.json').write_text(json.dumps(state,indent=2)) with tempfile.NamedTemporaryFile(mode='w',suffix='.sh') as f: f.write(startup); f.flush() + launch_record(phase='submitting TPU request '+queue) gc('alpha','compute','tpus','queued-resources','create',queue, '--zone='+ZONE,'--node-id='+node,'--accelerator-type=v5litepod-8', '--runtime-version=v2-alpha-tpuv5-lite','--provisioning-model=flex-start', @@ -163,6 +251,7 @@ def provision(machines, hours): '--metadata-from-file=startup-script='+f.name,'--labels=experiment=continuous8', '--quiet','--async') print('Submitted '+queue+'; cloud results: '+uri,flush=True) + launch_record(status='submitted',phase='TPU requests submitted; waiting for capacity/setup') print('Setup and preflight start automatically when each machine is allocated.') print('Inspect with: python3 baseline/nanogpt_one_head/continuous8/cloudshell.py status') print('Run settings: '+str(Path.home()/'continuous8-resources.json')) @@ -170,19 +259,29 @@ def provision(machines, hours): def main(): parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument('action', choices=['launch','status']) + parser.add_argument('action', choices=['launch','status','check']) parser.add_argument('--delete-old-experiments', action='store_true') parser.add_argument('--machines', type=int, choices=[1,2], default=1) parser.add_argument('--hours', type=int, choices=[4,6], default=6) args = parser.parse_args() - if args.action == 'status': + if args.action == 'check': + check_environment() + elif args.action == 'status': status() else: if (args.machines, args.hours) not in ((1,6),(2,4)): parser.error('Supported budgets: --machines 1 --hours 6 or --machines 2 --hours 4') - if args.delete_old_experiments: - cleanup() - provision(args.machines, args.hours) + try: + launch_record(status='running',phase='starting launcher',machines=args.machines,hours=args.hours) + if args.delete_old_experiments: + launch_record(phase='removing previously authorized old experiment resources') + cleanup() + provision(args.machines, args.hours) + except BaseException as exc: + launch_record(status='failed',error=f'{type(exc).__name__}: {exc}') + traceback.print_exc() + print('LAUNCH FAILED. See '+str(Path.home()/'continuous8-launch.json'),file=sys.stderr,flush=True) + raise SystemExit(1) if __name__ == '__main__': main() diff --git a/baseline/nanogpt_one_head/continuous8/prepare_tpu_data.py b/baseline/nanogpt_one_head/continuous8/prepare_tpu_data.py new file mode 100644 index 00000000..3bb1bec2 --- /dev/null +++ b/baseline/nanogpt_one_head/continuous8/prepare_tpu_data.py @@ -0,0 +1,51 @@ +#!/usr/bin/env python3 +"""Prepare or recover the pinned corpus on the TPU VM's durable data disk.""" +import json +import os +from pathlib import Path + +from google.api_core.exceptions import NotFound +import yaml + +from rg_nanogpt_one_head.continuous_support import CloudPublisher, sha_file +from rg_nanogpt_one_head.data import prepare_fineweb_edu, validate_prepared_data + + +def main(): + cfg = yaml.safe_load(Path('configs/muonclip_continuous8.yaml').read_text()) + data = Path('/mnt/disks/rg-data/continuous8/data') + data.mkdir(exist_ok=True) + archive = CloudPublisher(os.environ['RG_CONTINUOUS_DATA_URI']) + run = CloudPublisher(os.environ['RG_CONTINUOUS_GCS_URI']) + run.json({'status':'preparing_data','location':'TPU VM CPU and persistent disk'},'SETUP_STATUS.json') + try: + manifest = json.loads(archive.bucket.blob(archive.prefix+'/COMPLETE.json').download_as_text()) + except NotFound: + manifest = None + if manifest is not None: + if manifest['dataset'] != cfg['dataset']: + raise RuntimeError('Archived corpus does not match the experiment configuration') + for name in ('train.bin','val.bin','test.bin','meta.json'): + path = data/name + archive.bucket.blob(archive.prefix+'/'+name).download_to_filename(str(path),checksum='crc32c',timeout=600) + receipt = manifest['files'][name] + if path.stat().st_size != receipt['bytes'] or sha_file(path) != receipt['sha256']: + raise RuntimeError('Archived data checksum failed: '+name) + validate_prepared_data(data,cfg) + else: + print('Preparing 5B-token corpus on the TPU VM CPU; Cloud Shell can disconnect.',flush=True) + prepare_fineweb_edu(cfg,data) + receipts = {} + run.json({'status':'uploading_data'},'SETUP_STATUS.json') + for name in ('train.bin','val.bin','test.bin','meta.json'): + receipts[name] = archive.file(data/name,name) + receipts[name]['sha256'] = sha_file(data/name) + manifest = {'dataset':cfg['dataset'],'files':receipts} + archive.json(manifest,'COMPLETE.json') + run.json({'uri':os.environ['RG_CONTINUOUS_DATA_URI'],'manifest':manifest},'DATA_SOURCE.json') + run.snapshot_text_file('/mnt/disks/rg-data/continuous8/run.log','run.log') + print('Corpus complete: document-disjoint splits and SHA256 verified; cloud copy saved.',flush=True) + + +if __name__ == '__main__': + main() diff --git a/baseline/nanogpt_one_head/continuous8/run.sh b/baseline/nanogpt_one_head/continuous8/run.sh new file mode 100644 index 00000000..d18cea19 --- /dev/null +++ b/baseline/nanogpt_one_head/continuous8/run.sh @@ -0,0 +1,7 @@ +#!/usr/bin/env bash +# Persistent terminal transcript. This does not make Cloud Shell a durable VM. +set -Eeuo pipefail +export PYTHONUNBUFFERED=1 +SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd) +echo "Log: $HOME/continuous8-launch.log" +python3 -u "$SCRIPT_DIR/cloudshell.py" "$@" 2>&1 | tee -a "$HOME/continuous8-launch.log" diff --git a/baseline/nanogpt_one_head/continuous8/worker.sh b/baseline/nanogpt_one_head/continuous8/worker.sh index 752efaa6..15eaa9e6 100644 --- a/baseline/nanogpt_one_head/continuous8/worker.sh +++ b/baseline/nanogpt_one_head/continuous8/worker.sh @@ -13,6 +13,26 @@ unset XLA_USE_BF16 XLA_DOWNCAST_BF16 TPU_VISIBLE_CHIPS cd "$EXP" python3 -m venv "$BASE/venv" PY="$BASE/venv/bin/python" +finish_worker() { + local result=$? + trap - EXIT + set +e + echo "Worker exit code: $result" + "$PY" - "$result" <<'PYCODE' +import json, os, sys, time +from pathlib import Path +code=int(sys.argv[1]) +status={'status':'finished' if code in (0,75) else 'failed','exit_code':code,'ended_unix':time.time()} +base=Path('/mnt/disks/rg-data/continuous8') +(base/'WORKER_STATUS.json').write_text(json.dumps(status)+'\n') +from rg_nanogpt_one_head.continuous_support import CloudPublisher +p=CloudPublisher(os.environ['RG_CONTINUOUS_GCS_URI']) +p.json(status,'WORKER_STATUS.json') +p.snapshot_text_file(base/'run.log','run.log') +PYCODE + exit "$result" +} +trap finish_worker EXIT "$PY" -m pip install --upgrade 'pip==25.0.1' 'setuptools==75.8.2' 'wheel==0.45.1' "$PY" -m pip install 'torch==2.6.0' 'torch_xla[tpu]==2.6.0' \ -f https://storage.googleapis.com/libtpu-releases/index.html \ @@ -34,36 +54,18 @@ for name in ('environment.lock.txt','source_commit.txt'): p.file(base/name, name) p.json({'status':'preflight'}, 'SETUP_STATUS.json') PY -# Download the one shared CPU-prepared corpus and verify every file hash. -"$PY" - <<'PY' -import json, os -from pathlib import Path -import yaml -from rg_nanogpt_one_head.continuous_support import CloudPublisher, sha_file -from rg_nanogpt_one_head.data import validate_prepared_data -source=CloudPublisher(os.environ['RG_CONTINUOUS_DATA_URI']) -manifest=json.loads(source.bucket.blob(source.prefix+'/COMPLETE.json').download_as_text()) -cfg=yaml.safe_load(Path('configs/muonclip_continuous8.yaml').read_text()) -if manifest['dataset'] != cfg['dataset']: - raise RuntimeError('Shared cloud corpus does not match experiment configuration') -base=Path('/mnt/disks/rg-data/continuous8/data') -base.mkdir(exist_ok=True) -for name in ('train.bin','val.bin','test.bin','meta.json'): - path=base/name - source.bucket.blob(source.prefix+'/'+name).download_to_filename(str(path), checksum='crc32c', timeout=600) - entry=manifest['files'][name] - if path.stat().st_size != entry['bytes'] or sha_file(path) != entry['sha256']: - raise RuntimeError('Shared data checksum failed: '+name) -validate_prepared_data(base,cfg) -p=CloudPublisher(os.environ['RG_CONTINUOUS_GCS_URI']) -p.json({'uri':os.environ['RG_CONTINUOUS_DATA_URI'], 'manifest':manifest}, 'DATA_SOURCE.json') -print('Shared data downloaded and SHA256 verified.', flush=True) -PY # Tiny, independent test: global gradients, clipping, metrics, optimizer + RNG restore. # Also measures the full proposed model shape before the long run. "$PY" -m rg_nanogpt_one_head.tpu_spmd_check --backend tpu --chips 8 \ --benchmark-config configs/muonclip_continuous8.yaml --benchmark-steps 10 \ --output "$BASE/preflight.json" +"$PY" - <<'PYCODE' +import os +from rg_nanogpt_one_head.continuous_support import CloudPublisher +CloudPublisher(os.environ['RG_CONTINUOUS_GCS_URI']).file('/mnt/disks/rg-data/continuous8/preflight.json','preflight.json') +PYCODE +# Dataset preparation runs on the TPU VM CPU and durable disk, not Cloud Shell. +"$PY" -u continuous8/prepare_tpu_data.py "$PY" - <<'PY' import json, os, time from pathlib import Path diff --git a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/data.py b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/data.py index beb18d83..3e9ff283 100644 --- a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/data.py +++ b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/data.py @@ -1,8 +1,10 @@ from __future__ import annotations import argparse +from concurrent.futures import ThreadPoolExecutor import hashlib import json +from itertools import islice import os from pathlib import Path import sys @@ -38,6 +40,21 @@ def _encode_document(text: str, encoder: Encoder) -> np.ndarray: return np.asarray(tokens, dtype=TOKEN_DTYPE) +def _encoded_documents(texts, encoder, workers=1, batch_size=256): + """Bounded parallel tokenization; preserve exactly the original document order.""" + if workers < 1 or batch_size < 1: + raise ValueError('encoding workers and batch size must be positive') + if workers == 1: + for text in texts: + yield _encode_document(str(text), encoder) + return + iterator = iter(texts) + with ThreadPoolExecutor(max_workers=workers) as pool: + while batch := list(islice(iterator, batch_size)): + # executor.map yields in input order, independent of worker timing. + yield from pool.map(lambda text: _encode_document(str(text), encoder), batch) + + def write_token_splits( texts: Iterable[str], encoder: Encoder, @@ -48,6 +65,8 @@ def write_token_splits( test_tokens: int, dataset_metadata: dict[str, object] | None = None, progress_every_documents: int = 2_000, + encoding_workers: int = 1, + encoding_batch_size: int = 256, ) -> dict[str, object]: """Write exact, document-disjoint splits without loading the corpus in RAM.""" @@ -75,11 +94,10 @@ def write_token_splits( started = time.monotonic() try: - for text in texts: + for encoded in _encoded_documents(texts, encoder, encoding_workers, encoding_batch_size): if split_index >= len(split_names): break documents += 1 - encoded = _encode_document(str(text), encoder) split = split_names[split_index] remaining = targets[split] - written[split] take = min(remaining, len(encoded)) @@ -295,6 +313,8 @@ def prepare_fineweb_edu( train_tokens=int(dataset_cfg["train_tokens"]), val_tokens=int(dataset_cfg["val_tokens"]), test_tokens=int(dataset_cfg["test_tokens"]), + encoding_workers=int(dataset_cfg.get('encoding_workers',1)), + encoding_batch_size=int(dataset_cfg.get('encoding_batch_size',256)), dataset_metadata={ "dataset_name": str(dataset_cfg["name"]), "dataset_config": str(dataset_cfg["config"]), diff --git a/baseline/nanogpt_one_head/tests/test_continuous8_launcher.py b/baseline/nanogpt_one_head/tests/test_continuous8_launcher.py index b36314f9..e57cb5b7 100644 --- a/baseline/nanogpt_one_head/tests/test_continuous8_launcher.py +++ b/baseline/nanogpt_one_head/tests/test_continuous8_launcher.py @@ -7,7 +7,7 @@ @pytest.mark.parametrize('machines,hours', [(1,6),(2,4)]) -def test_cloud_budget_and_data_before_allocation(tmp_path, monkeypatch, machines, hours): +def test_cloud_budget_without_cloudshell_preparation(tmp_path, monkeypatch, machines, hours): file = Path(__file__).resolve().parents[1]/'continuous8/cloudshell.py' spec = importlib.util.spec_from_file_location('pilot_launcher', file) m = importlib.util.module_from_spec(spec) @@ -21,17 +21,13 @@ def inventory(*args): if args[:3] == ('iam','service-accounts','list'): return [{'email':m.SA}] return [] monkeypatch.setattr(m,'inventory',inventory) - prepared = False calls = [] def run(cmd,**kw): - nonlocal prepared - if any(str(x).endswith('prepare_cloud_data.py') for x in cmd): prepared = True - return subprocess.CompletedProcess(cmd,0) + raise AssertionError('Launcher must not install dependencies or prepare data in Cloud Shell: '+str(cmd)) monkeypatch.setattr(m.subprocess,'run',run) def gc(*args,**kw): calls.append(args) if args[:5] == ('alpha','compute','tpus','queued-resources','create'): - assert prepared, 'No TPU allocation before CPU preparation succeeds' assert f'--max-run-duration={hours}h' in args source = next(x.split('=',2)[-1] for x in args if x.startswith('--metadata-from-file=')) startup = Path(source).read_text() @@ -45,6 +41,9 @@ def gc(*args,**kw): assert len(create) == machines assert len({c[5] for c in create}) == machines assert all('--accelerator-type=v5litepod-8' in c for c in create) + import json + record=json.loads((tmp_path/'continuous8-launch.json').read_text()) + assert record['status']=='submitted' def test_dataset_writer_import_requires_no_torch(): @@ -57,3 +56,52 @@ def test_dataset_writer_import_requires_no_torch(): 'assert "torch" not in sys.modules' ) subprocess.run([sys.executable,'-c',code],check=True) + + +def test_empty_status_is_explicit(tmp_path,monkeypatch,capsys): + file=Path(__file__).resolve().parents[1]/'continuous8/cloudshell.py' + spec=importlib.util.spec_from_file_location('status_launcher',file) + m=importlib.util.module_from_spec(spec) + spec.loader.exec_module(m) + monkeypatch.setattr(Path,'home',classmethod(lambda cls:tmp_path)) + monkeypatch.setattr(m,'inventory',lambda *args:[]) + m.status() + assert 'NO TPU REQUEST' in capsys.readouterr().out + + +def test_launch_failure_persists_exact_phase(tmp_path,monkeypatch): + import json + import sys + file=Path(__file__).resolve().parents[1]/'continuous8/cloudshell.py' + spec=importlib.util.spec_from_file_location('failed_launcher',file) + m=importlib.util.module_from_spec(spec) + spec.loader.exec_module(m) + monkeypatch.setattr(Path,'home',classmethod(lambda cls:tmp_path)) + monkeypatch.setattr(sys,'argv',[str(file),'launch']) + def fail(*args): + m.launch_record(phase='granting experiment bucket access') + raise RuntimeError('permission denied in test') + monkeypatch.setattr(m,'provision',fail) + with pytest.raises(SystemExit) as caught: + m.main() + assert caught.value.code==1 + record=json.loads((tmp_path/'continuous8-launch.json').read_text()) + assert record['status']=='failed' + assert record['phase']=='granting experiment bucket access' + assert 'permission denied' in record['error'] + + +def test_parallel_data_matches_serial(tmp_path): + from rg_nanogpt_one_head.data import write_token_splits + class Encoder: + n_vocab=64 + eot_token=63 + def encode_ordinary(self,text):return [int(x) for x in text.split()] + texts=['1 2 3','4 5','6 7 8 9','10','11 12 13 14','15 16']*10 + common=dict(train_tokens=25,val_tokens=13,test_tokens=19,progress_every_documents=0) + serial=write_token_splits(texts,Encoder(),tmp_path/'serial',**common) + parallel=write_token_splits(iter(texts),Encoder(),tmp_path/'parallel',encoding_workers=4, + encoding_batch_size=3,**common) + assert serial==parallel + for name in ('train.bin','val.bin','test.bin','meta.json'): + assert (tmp_path/'serial'/name).read_bytes()==(tmp_path/'parallel'/name).read_bytes() From 22afb4a6198014f0a691cffeaf25caee88b0d56a Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Fri, 2 Oct 2026 17:33:11 -0700 Subject: [PATCH 08/36] Recover interrupted dependency downloads before continuous TPU training Use a persistent pip cache, 300-second timeouts and bounded install retries. Install pinned CPU PyTorch with XLA to avoid unnecessary CUDA downloads. Preserve the allocation deadline; refuse concurrent workers or any scientific run restart before changing dependencies. Keep local failure status when the cloud reporter is unavailable during installation. Validation: 18 local tests passed, including interrupted downloads, exhausted retries and expired allocation deadlines. Shell syntax and diff checks passed. TPU execution remains subject to the existing on-machine preflight. --- .../nanogpt_one_head/continuous8/README.md | 11 ++++++ .../continuous8/install_dependencies.sh | 36 +++++++++++++++++ .../nanogpt_one_head/continuous8/worker.sh | 33 ++++++++++------ .../tests/test_continuous8_launcher.py | 39 +++++++++++++++++++ 4 files changed, 107 insertions(+), 12 deletions(-) create mode 100644 baseline/nanogpt_one_head/continuous8/install_dependencies.sh diff --git a/baseline/nanogpt_one_head/continuous8/README.md b/baseline/nanogpt_one_head/continuous8/README.md index b33ee84b..711b5c03 100644 --- a/baseline/nanogpt_one_head/continuous8/README.md +++ b/baseline/nanogpt_one_head/continuous8/README.md @@ -161,6 +161,17 @@ exit, including failures before training; `CONTINUOUS_STATUS.json` records the scientific process outcome. If dependency installation fails before cloud upload libraries are available, inspect the persistent disk's startup/run log over SSH. +Package downloads use a 300-second socket timeout, eight connection retries, +and up to three attempts per install command, with a cache on the data disk. +The pinned PyTorch 2.6 CPU wheel plus XLA 2.6 supplies TPU support without CUDA +downloads. Installation retries retain the original allocation deadline. +If installation exhausts its retries, a manual service start after updating the +worker is allowed only before scientific training has started. The worker locks +out concurrent setup and checks `results/CONTINUOUS_STARTED.json` before changing +dependencies. Existing cloud run claims are never deleted or overwritten, so a +failure after claiming the run still stops on a duplicate claim. Training itself +has no retry or resume loop; no deadline or restart guard is reset for this repair. + The earlier launcher did long preparation in Cloud Shell `/tmp`. No request/data was found after one attempt; the initial error was not retained, so its exact cause is unknown. Cloud Shell VM disposal can lose `/tmp`, and the old empty diff --git a/baseline/nanogpt_one_head/continuous8/install_dependencies.sh b/baseline/nanogpt_one_head/continuous8/install_dependencies.sh new file mode 100644 index 00000000..ae577aac --- /dev/null +++ b/baseline/nanogpt_one_head/continuous8/install_dependencies.sh @@ -0,0 +1,36 @@ +#!/usr/bin/env bash +# Sourced by worker.sh; PY, BASE and the working directory are already set. +# pip 25.0.1 connection retries do not recover a mid-download read timeout, +# so retry the install command as well. Successful downloads remain cached. +export PIP_CACHE_DIR="$BASE/pip-cache" +export PIP_DEFAULT_TIMEOUT=300 PIP_RETRIES=8 PIP_PROGRESS_BAR=off +export PIP_DISABLE_PIP_VERSION_CHECK=1 +pip_install() { + local attempt result=1 + for attempt in 1 2 3; do + if [ -n "${RG_CONTINUOUS_DEADLINE_UNIX:-}" ] && + [ "$(date +%s)" -ge "${RG_CONTINUOUS_DEADLINE_UNIX%%.*}" ]; then + echo 'Original allocation deadline passed; refusing further installation.' >&2 + return 1 + fi + echo "Dependency install attempt $attempt/3: $*" + if "$PY" -m pip install "$@"; then + return 0 + else + result=$? + fi + if [ "$attempt" -lt 3 ]; then + echo "Dependency install exited $result; retrying in 5 seconds." >&2 + sleep 5 + fi + done + return "$result" +} +pip_install --upgrade 'pip==25.0.1' 'setuptools==75.8.2' 'wheel==0.45.1' +# XLA supplies the TPU backend. Avoid the unnecessary CUDA dependency downloads. +pip_install 'torch==2.6.0+cpu' --index-url https://download.pytorch.org/whl/cpu +pip_install 'torch_xla[tpu]==2.6.0' \ + -f https://storage.googleapis.com/libtpu-releases/index.html \ + -f https://storage.googleapis.com/libtpu-wheels/index.html +pip_install -r continuous8/requirements.txt +pip_install --no-deps --no-build-isolation -e . diff --git a/baseline/nanogpt_one_head/continuous8/worker.sh b/baseline/nanogpt_one_head/continuous8/worker.sh index 15eaa9e6..963750f2 100644 --- a/baseline/nanogpt_one_head/continuous8/worker.sh +++ b/baseline/nanogpt_one_head/continuous8/worker.sh @@ -1,9 +1,17 @@ #!/usr/bin/env bash -# Invoked once by systemd, with Restart=no. Scientific process never resumes. +# Invoked by systemd, with Restart=no. A failed dependency install may be retried +# manually before training starts. The scientific process never resumes. set -Eeuo pipefail BASE=/mnt/disks/rg-data/continuous8 REPO="$BASE/repo" EXP="$REPO/baseline/nanogpt_one_head" +# Serialize manual setup retries and reject all scientific-run reuse before pip. +exec 9>"$BASE/worker.lock" +flock -n 9 || { echo 'Another worker is running; refusing duplicate setup.'; exit 1; } +if [ -e "$BASE/results/CONTINUOUS_STARTED.json" ]; then + echo 'Scientific training already started; refusing setup or training restart.' + exit 1 +fi export OMP_NUM_THREADS=4 OPENBLAS_NUM_THREADS=4 MKL_NUM_THREADS=4 export PJRT_DEVICE=TPU TPU_ACCELERATOR_TYPE=v5litepod-8 export RG_TPU_PERSISTENT_ROOT=/mnt/disks/rg-data @@ -25,20 +33,21 @@ code=int(sys.argv[1]) status={'status':'finished' if code in (0,75) else 'failed','exit_code':code,'ended_unix':time.time()} base=Path('/mnt/disks/rg-data/continuous8') (base/'WORKER_STATUS.json').write_text(json.dumps(status)+'\n') -from rg_nanogpt_one_head.continuous_support import CloudPublisher -p=CloudPublisher(os.environ['RG_CONTINUOUS_GCS_URI']) -p.json(status,'WORKER_STATUS.json') -p.snapshot_text_file(base/'run.log','run.log') +try: + from rg_nanogpt_one_head.continuous_support import CloudPublisher + p=CloudPublisher(os.environ['RG_CONTINUOUS_GCS_URI']) + p.json(status,'WORKER_STATUS.json') + p.snapshot_text_file(base/'run.log','run.log') +except Exception as exc: + print(f'Cloud exit report unavailable: {type(exc).__name__}: {exc}', flush=True) + print(f'Local status and log retained in {base}', flush=True) PYCODE exit "$result" } trap finish_worker EXIT -"$PY" -m pip install --upgrade 'pip==25.0.1' 'setuptools==75.8.2' 'wheel==0.45.1' -"$PY" -m pip install '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 -"$PY" -m pip install -r continuous8/requirements.txt -"$PY" -m pip install --no-deps -e . +# Only package installation is retried. Preflight, data preparation and scientific +# training remain single-attempt, with the original allocation deadline. +source "$EXP/continuous8/install_dependencies.sh" "$PY" -m pip freeze > "$BASE/environment.lock.txt" git rev-parse HEAD > "$BASE/source_commit.txt" # Upload credentials test before data download or training. @@ -85,7 +94,7 @@ p.json(projection, 'BENCHMARK_PROJECTION.json') print('BENCHMARK PROJECTION:', json.dumps(projection), flush=True) p.json({'status':'training'}, 'SETUP_STATUS.json') PY -# No resilient supervisor, --resume, or automatic retry anywhere in this path. +# No resilient supervisor, --resume, or automatic retry of scientific training. "$PY" -u -m rg_nanogpt_one_head.continuous_run \ --config configs/muonclip_continuous8.yaml --data-root "$BASE/data" \ --results-root "$BASE/results" --device tpu --seed "$RG_CONTINUOUS_SEED" \ diff --git a/baseline/nanogpt_one_head/tests/test_continuous8_launcher.py b/baseline/nanogpt_one_head/tests/test_continuous8_launcher.py index e57cb5b7..a9e901e4 100644 --- a/baseline/nanogpt_one_head/tests/test_continuous8_launcher.py +++ b/baseline/nanogpt_one_head/tests/test_continuous8_launcher.py @@ -105,3 +105,42 @@ def encode_ordinary(self,text):return [int(x) for x in text.split()] assert serial==parallel for name in ('train.bin','val.bin','test.bin','meta.json'): assert (tmp_path/'serial'/name).read_bytes()==(tmp_path/'parallel'/name).read_bytes() + + +@pytest.mark.parametrize('failures,deadline,expected_calls,success', [ + (1, '4102444800.5', 6, True), # Recover a interrupted wheel download. + (99, '4102444800.5', 3, False), # Stop before later install stages. + (0, '1.5', 0, False), # A retry cannot extend the original allocation. +]) +def test_dependency_download_recovery(tmp_path, failures, deadline, expected_calls, success): + import json + import os + import shlex + import sys + installer = Path(__file__).resolve().parents[1]/'continuous8/install_dependencies.sh' + fake = tmp_path/'fake_python' + fake.write_text(f'#!{sys.executable}\n' + ''' +import json, os, pathlib, sys +p = pathlib.Path(os.environ['BASE'])/'pip_calls.jsonl' +count = len(p.read_text().splitlines()) if p.exists() else 0 +with p.open('a') as f: + f.write(json.dumps({'args':sys.argv[1:], + 'timeout':os.environ['PIP_DEFAULT_TIMEOUT'], + 'cache':os.environ['PIP_CACHE_DIR']})+'\\n') +sys.exit(2 if count < int(os.environ['FAKE_FAILURES']) else 0) +''') + fake.chmod(0o755) + env = dict(os.environ, BASE=str(tmp_path), PY=str(fake), + FAKE_FAILURES=str(failures), RG_CONTINUOUS_DEADLINE_UNIX=deadline) + result = subprocess.run(['bash','-ec', + 'sleep() { :; }; source '+shlex.quote(str(installer))], env=env, + text=True, capture_output=True) + assert (result.returncode == 0) == success, result.stdout+result.stderr + log = tmp_path/'pip_calls.jsonl' + calls = [json.loads(x) for x in log.read_text().splitlines()] if log.exists() else [] + assert len(calls) == expected_calls + assert all(x['timeout'] == '300' and x['cache'] == str(tmp_path/'pip-cache') for x in calls) + if success: + assert calls[0]['args'] == calls[1]['args'] + assert any('torch==2.6.0+cpu' in x['args'] for x in calls) + assert calls[-1]['args'][-1] == '.' From 641ad6ade10808e3f74f825c4e9ce02dec839dcd Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Fri, 2 Oct 2026 18:30:15 -0700 Subject: [PATCH 09/36] Replace pilot with one 24-hour TPU and double paired alpha monitoring Add a Cloud Shell replacement command that stops the old worker, removes TPU requests and VMs in the two used zones, verifies deletion, and reuses the existing data disk for one fresh v5e-8 request. Preserve old results, share installed packages and data caches, and block duplicate allocation on repeated invocation. Use a separate 24-hour configuration with paired token-error/alpha measurements every 500 updates and a 23.5-hour stop deadline from VM boot. Archive all checkpoints in GCS; retain three verified epoch files locally plus rolling checkpoints. Add benchmark stage messages and periodic waiting heartbeats. Validation: 35 CPU tests passed, one hardware-only test skipped. Tests cover reallocation ordering/failure/duplicate guards, preserved disk, the 24-hour budget, unchanged optimizer/probes, aligned measurements and verified pruning. TPU numerical and shape preflight must pass again on the new allocation. --- .../configs/muonclip_continuous8_24h.yaml | 134 +++++++++++++++++ .../nanogpt_one_head/continuous8/README.md | 42 ++++++ .../continuous8/install_dependencies.sh | 2 +- .../continuous8/prepare_tpu_data.py | 8 +- .../continuous8/reallocate24.py | 138 ++++++++++++++++++ .../nanogpt_one_head/continuous8/worker.sh | 27 ++-- .../rg_nanogpt_one_head/continuous_support.py | 15 ++ .../src/rg_nanogpt_one_head/tpu_spmd_check.py | 20 ++- .../tests/test_continuous8.py | 50 +++++++ .../tests/test_reallocate24.py | 70 +++++++++ 10 files changed, 489 insertions(+), 17 deletions(-) create mode 100644 baseline/nanogpt_one_head/configs/muonclip_continuous8_24h.yaml create mode 100644 baseline/nanogpt_one_head/continuous8/reallocate24.py create mode 100644 baseline/nanogpt_one_head/tests/test_reallocate24.py diff --git a/baseline/nanogpt_one_head/configs/muonclip_continuous8_24h.yaml b/baseline/nanogpt_one_head/configs/muonclip_continuous8_24h.yaml new file mode 100644 index 00000000..45ddd793 --- /dev/null +++ b/baseline/nanogpt_one_head/configs/muonclip_continuous8_24h.yaml @@ -0,0 +1,134 @@ +protocol: + name: muonclip_continuous8_124m_5bt_24h + version: 2 + description: Fresh continuous single-host eight-chip run; no segments, no automatic + restart; fixed document probes; durable GCS checkpoints. +dataset: + name: HuggingFaceFW/fineweb-edu + config: sample-10BT + split: train + revision: 593b3a867298afb8ce42625a270ef20ddcad28f9 + tokenizer: gpt2 + encoding_workers: 16 + encoding_batch_size: 256 + train_tokens: 5000000000 + val_tokens: 10000000 + test_tokens: 10000000 +model: + vocab_size: 50257 + block_size: 256 + n_layer: 12 + n_head: 12 + n_embd: 768 + dropout: 0.0 + bias: false + tie_weights: true +training: + seeds: + - 1337 + - 2027 + batch_size: 32 + grad_accum_steps: 4 + target_epochs: 6.5536 + epoch_interval: 0.0032768 + eval_interval_steps: 500 + eval_batches: 16 + checkpoint_interval_steps: 500 + grad_clip: 1.0 + max_steps: 1000000 + min_free_disk_gb: 20 +optimizer_profiles: + sgd_momentum: + display_name: SGD + Nesterov momentum + family: sgd + learning_rate: 0.05 + min_learning_rate: 0.005 + warmup_fraction: 0.1 + schedule: warmup_cosine + momentum: 0.9 + dampening: 0.0 + nesterov: true + weight_decay: 0.01 + adamw: + display_name: AdamW + family: adamw + learning_rate: 0.0006 + min_learning_rate: 6.0e-05 + warmup_fraction: 0.01 + schedule: warmup_cosine + beta1: 0.9 + beta2: 0.95 + epsilon: 1.0e-08 + weight_decay: 0.1 + muon: + display_name: Muon + auxiliary AdamW + family: muon + matrix_learning_rate: 0.02 + matrix_min_learning_rate: 0.002 + aux_learning_rate: 0.0003 + aux_min_learning_rate: 3.0e-05 + warmup_fraction: 0.05 + schedule: warmup_cosine + momentum: 0.95 + nesterov: true + newton_schulz_steps: 5 + muon_epsilon: 1.0e-07 + matrix_weight_decay: 0.01 + beta1: 0.9 + beta2: 0.95 + epsilon: 1.0e-08 + aux_weight_decay: 0.01 + muon_clip: + display_name: MuonClip + RMS-matched updates + auxiliary AdamW + family: muon_clip + learning_rate: 0.0002 + min_learning_rate: 2.0e-05 + warmup_fraction: 0.02 + schedule: warmup_cosine + momentum: 0.95 + nesterov: false + newton_schulz_steps: 5 + muon_epsilon: 1.0e-07 + weight_decay: 0.1 + update_rms_scale: 0.2 + qk_clip_threshold: 100.0 + qk_clip_balance: 0.5 + qk_diagnostics_interval: 500 + beta1: 0.9 + beta2: 0.95 + epsilon: 1.0e-08 + lr_schedule_steps: 100000 +evaluation: + train_probe_seed: 21001 + validation_probe_seed: 22001 + test_probe_seed: 23001 + bleu_probe_seed: 24001 + bleu_examples: 64 + bleu_prompt_tokens: 64 + bleu_continuation_tokens: 32 + bleu_batch_size: 4 + test_interval_steps: 500 + document_probe: true + probe_documents: 512 +weightwatcher: + enabled: true + ERG: true + randomize: true + strict: true + min_evals: 20 + fix_fingers: clip_xmax + max_fingers: 10 + require_raw_alpha: true +runtime: + matmul_precision: highest + mps_fallback: true + deterministic_algorithms: false + empty_mps_cache_after_weightwatcher: true + tpu_spmd: true + tpu_expected_chips: 8 +continuous: + enabled: true + cloud_required: true + max_wall_hours: 23.5 + auto_resume: false + keep_local_epoch_checkpoints: 3 diff --git a/baseline/nanogpt_one_head/continuous8/README.md b/baseline/nanogpt_one_head/continuous8/README.md index 711b5c03..97cb9fa2 100644 --- a/baseline/nanogpt_one_head/continuous8/README.md +++ b/baseline/nanogpt_one_head/continuous8/README.md @@ -98,6 +98,48 @@ allocated hardware before the launch script starts scientific training. ## Launch and monitor +### Replace the pilot with a 24-hour allocation + +For the explicitly authorized 24-hour experiment, run from a clean checkout: + +```bash +python3 -u baseline/nanogpt_one_head/continuous8/reallocate24.py +``` + +This stops the pilot worker, deletes TPU queued requests and VMs in **us-west4-a +and us-east5-a**, verifies that both zones are empty, then submits exactly one +`v5litepod-8` Flex-start request with a 24-hour maximum. Other zones are not +inventoried. All data disks and cloud buckets are retained. The existing 200 GB +pilot disk must exist and is reused without formatting. Failed deletion prevents +new allocation. Repeating the command while the new request exists does nothing +destructive and creates no duplicate. + +The new request is `ww-continuous8-24h-20261003-s1337`; its fresh run directory is +`/mnt/disks/rg-data/continuous8-24h-20261003-s1337`, with a separate GCS run prefix. +It reuses the original `continuous8/venv`, `pip-cache`, `hf-cache` and `data` paths. +Setup performs all TPU correctness/benchmark checks again on the replacement VM. +Scientific training starts from step zero; no previous scientific checkpoint is +loaded. Benchmark messages identify each warm-up/timed update and emit a waiting +heartbeat every minute; a heartbeat is not evidence of completed device work. + +The 24-hour config keeps model, dataset, optimizer, learning-rate schedule and +fixed probes identical to the six-hour pilot configuration, but pairs alpha and +token error **every 500 updates** instead of every 1,000. Full checkpoints remain +every 500 updates. All checkpoint uploads remain archived in GCS, while the last +three uploaded epoch checkpoint files are retained locally alongside latest, +best and initialization/final checkpoints. Local epoch files are pruned only +after a completion receipt and a matching local SHA256; unverified files remain. + +Training is requested to stop at **23.5 hours from VM boot**, leaving 30 minutes +before the 24-hour server expiry. Setup, data preparation, compilation, monitoring +and upload all consume allocation time: this is not 24 guaranteed training hours. +The new allocation can be deleted earlier after analysis. Compute cap at the +published v5e Flex-start price is **$115.20**, plus storage/network charges and +the already incurred cost of the old allocation. The Cloud Shell command records +phases, errors and resource IDs in `~/continuous8-24h-resources.json`. + +### Original shorter pilot + From a clean checkout of `codex/continuous-muonclip-8`, on Cloud Shell: ```bash diff --git a/baseline/nanogpt_one_head/continuous8/install_dependencies.sh b/baseline/nanogpt_one_head/continuous8/install_dependencies.sh index ae577aac..d078b2e7 100644 --- a/baseline/nanogpt_one_head/continuous8/install_dependencies.sh +++ b/baseline/nanogpt_one_head/continuous8/install_dependencies.sh @@ -2,7 +2,7 @@ # Sourced by worker.sh; PY, BASE and the working directory are already set. # pip 25.0.1 connection retries do not recover a mid-download read timeout, # so retry the install command as well. Successful downloads remain cached. -export PIP_CACHE_DIR="$BASE/pip-cache" +export PIP_CACHE_DIR="${PIP_CACHE_DIR:-$BASE/pip-cache}" export PIP_DEFAULT_TIMEOUT=300 PIP_RETRIES=8 PIP_PROGRESS_BAR=off export PIP_DISABLE_PIP_VERSION_CHECK=1 pip_install() { diff --git a/baseline/nanogpt_one_head/continuous8/prepare_tpu_data.py b/baseline/nanogpt_one_head/continuous8/prepare_tpu_data.py index 3bb1bec2..8bf64cec 100644 --- a/baseline/nanogpt_one_head/continuous8/prepare_tpu_data.py +++ b/baseline/nanogpt_one_head/continuous8/prepare_tpu_data.py @@ -12,9 +12,9 @@ def main(): - cfg = yaml.safe_load(Path('configs/muonclip_continuous8.yaml').read_text()) - data = Path('/mnt/disks/rg-data/continuous8/data') - data.mkdir(exist_ok=True) + cfg = yaml.safe_load(Path(os.environ.get('RG_CONTINUOUS_CONFIG','configs/muonclip_continuous8.yaml')).read_text()) + data = Path(os.environ.get('RG_CONTINUOUS_DATA_ROOT','/mnt/disks/rg-data/continuous8/data')) + data.mkdir(parents=True,exist_ok=True) archive = CloudPublisher(os.environ['RG_CONTINUOUS_DATA_URI']) run = CloudPublisher(os.environ['RG_CONTINUOUS_GCS_URI']) run.json({'status':'preparing_data','location':'TPU VM CPU and persistent disk'},'SETUP_STATUS.json') @@ -43,7 +43,7 @@ def main(): manifest = {'dataset':cfg['dataset'],'files':receipts} archive.json(manifest,'COMPLETE.json') run.json({'uri':os.environ['RG_CONTINUOUS_DATA_URI'],'manifest':manifest},'DATA_SOURCE.json') - run.snapshot_text_file('/mnt/disks/rg-data/continuous8/run.log','run.log') + run.snapshot_text_file(os.environ.get('RG_CONTINUOUS_RUN_LOG','/mnt/disks/rg-data/continuous8/run.log'),'run.log') print('Corpus complete: document-disjoint splits and SHA256 verified; cloud copy saved.',flush=True) diff --git a/baseline/nanogpt_one_head/continuous8/reallocate24.py b/baseline/nanogpt_one_head/continuous8/reallocate24.py new file mode 100644 index 00000000..2b08f57a --- /dev/null +++ b/baseline/nanogpt_one_head/continuous8/reallocate24.py @@ -0,0 +1,138 @@ +#!/usr/bin/env python3 +"""Replace experiment TPUs in the two used zones with one 24-hour v5e-8. + +Retain every data disk and bucket; reuse the existing pilot data disk. This is a +fresh scientific run in a new directory/archive, never a checkpoint continuation. +""" +from datetime import datetime, timezone +import json +from pathlib import Path +import subprocess +import sys +import time + +from cloudshell import PROJECT, ZONE, BUCKET, SA, gc, inventory, short + +ZONES = ('us-west4-a', 'us-east5-a') +QUEUE = 'ww-continuous8-24h-20261003-s1337' +NODE = QUEUE+'-node' +DISK = 'ww-continuous8-pilot-20261002-s1337-data' +OLD_NODE = 'ww-continuous8-pilot-20261002-s1337-node' +BASE = '/mnt/disks/rg-data/continuous8-24h-20261003-s1337' +SHARED = '/mnt/disks/rg-data/continuous8' +URI = f'gs://{BUCKET}/runs/{QUEUE}' +DATA_URI = f'gs://{BUCKET}/corpora/ww-continuous8-pilot-20261002-s1337' + + +def make_startup(commit): + source = Path(__file__).with_name('startup.sh').read_text() + source = source.replace('BASE=/mnt/disks/rg-data/continuous8', 'BASE='+BASE) + # Reusing a known disk must never format a missing/incorrect filesystem. + start = source.index('if [ -z "$TYPE" ]; then') + end = source.index('mkdir -p /mnt/disks/rg-data', start) + source = source[:start] + '''if [ "$TYPE" != ext4 ]; then + echo 'Expected the existing ext4 data disk; refusing to format.' >&2 + exit 1 +fi +''' + source[end:] + source = source.replace('Environment=RG_CONTINUOUS_SEED=__SEED__', + 'Environment=RG_CONTINUOUS_SEED=__SEED__\n' + 'Environment=RG_CONTINUOUS_BASE=$BASE\n' + 'Environment=RG_CONTINUOUS_SHARED_BASE='+SHARED+'\n' + 'Environment=RG_CONTINUOUS_CONFIG=$BASE/repo/baseline/nanogpt_one_head/configs/muonclip_continuous8_24h.yaml') + return (source.replace('__COMMIT__', commit).replace('__GCS_URI__', URI) + .replace('__DATA_URI__', DATA_URI).replace('__SEED__', '1337') + .replace('__STOP_HOURS__', '23.5')) + + +def save_record(record, **changes): + record.update(changes, updated_utc=datetime.now(timezone.utc).isoformat()) + path = Path.home()/'continuous8-24h-resources.json' + tmp = path.with_suffix('.tmp') + tmp.write_text(json.dumps(record, indent=2)+'\n') + tmp.replace(path) + if 'phase' in changes: + print('[24h] '+changes['phase'], flush=True) + + +def main(): + record = dict(project=PROJECT, zone=ZONE, queue=QUEUE, node=NODE, disk=DISK, + base=BASE, gcs_uri=URI, hours=24, machines=1, seed=1337) + try: + # Complete all read-only checks and build the startup script first. + root = Path(__file__).resolve().parents[3] + commit = subprocess.check_output(['git','-C',str(root),'rev-parse','HEAD'],text=True).strip() + if subprocess.check_output(['git','-C',str(root),'status','--porcelain'],text=True).strip(): + raise RuntimeError('Use a clean checkout so the allocated VM receives the same code.') + save_record(record, commit=commit, phase='checking existing resources', status='preparing') + queues = {z:inventory('alpha','compute','tpus','queued-resources','list','--zone='+z) for z in ZONES} + if any(short(q)==QUEUE for q in queues[ZONE]): + save_record(record, status='existing', phase='24-hour request already exists; no deletion or duplicate launch') + print(json.dumps(inventory('alpha','compute','tpus','queued-resources','describe',QUEUE, + '--zone='+ZONE),indent=2)) + return + nodes = {z:inventory('compute','tpus','tpu-vm','list','--zone='+z) for z in ZONES} + disk = inventory('compute','disks','describe',DISK,'--zone='+ZONE) + inventory('storage','buckets','describe','gs://'+BUCKET) + inventory('iam','service-accounts','describe',SA) + if disk.get('zone','').rsplit('/',1)[-1] != ZONE: + raise RuntimeError('Existing data disk is in the wrong zone.') + save_record(record, old_queues=queues, old_nodes=nodes, preserved_disk=disk) + startup = make_startup(commit) + subprocess.run(['bash','-n'],input=startup,text=True,check=True) + startup_path = Path.home()/'continuous8-24h-startup.sh' + startup_path.write_text(startup) + print('One v5e-8, 24-hour allocation; compute cap $115.20 plus storage.',flush=True) + print('Stop requested at 23.5h from boot; setup is included. Measurements and checkpoints every 500 updates.',flush=True) + print('Removing TPU queues/VMs in us-west4-a and us-east5-a; retaining all data disks and buckets.',flush=True) + if any(short(n)==OLD_NODE for n in nodes[ZONE]): + save_record(record, phase='stopping the old worker and flushing the persistent disk') + command = '''sudo bash -ec ' +test -d /mnt/disks/rg-data/continuous8 +mountpoint -q /mnt/disks/rg-data +if ! timeout 120 systemctl stop rg-continuous8.service; then + systemctl kill --kill-who=all --signal=SIGKILL rg-continuous8.service +fi +sync +' ''' + gc('compute','tpus','tpu-vm','ssh',OLD_NODE,'--zone='+ZONE,'--worker=0','--command='+command) + save_record(record, phase='deleting old TPU requests and their nodes', status='replacing') + for zone, items in queues.items(): + for item in items: + gc('alpha','compute','tpus','queued-resources','delete',short(item), + '--zone='+zone,'--force','--quiet') + for zone in ZONES: + for item in inventory('compute','tpus','tpu-vm','list','--zone='+zone): + gc('compute','tpus','tpu-vm','delete',short(item),'--zone='+zone,'--quiet') + for zone in ZONES: + for kind in ('queued-resources','tpu-vm'): + if inventory('alpha','compute','tpus',kind,'list','--zone='+zone): + raise RuntimeError(f'{zone}: {kind} still present; no new allocation submitted.') + save_record(record, phase='checking that the preserved disk is detached') + for attempt in range(13): + disk = inventory('compute','disks','describe',DISK,'--zone='+ZONE) + if not disk.get('users'): + break + if attempt == 12: + raise RuntimeError('Data disk is still attached; preserved, but no new request submitted.') + time.sleep(5) + save_record(record, phase='submitting the new 24-hour request') + gc('alpha','compute','tpus','queued-resources','create',QUEUE, + '--zone='+ZONE,'--node-id='+NODE,'--accelerator-type=v5litepod-8', + '--runtime-version=v2-alpha-tpuv5-lite','--provisioning-model=flex-start', + '--max-run-duration=24h','--valid-until-duration=24h', + '--service-account='+SA,'--scopes=https://www.googleapis.com/auth/cloud-platform', + '--data-disk=source=projects/'+PROJECT+'/zones/'+ZONE+'/disks/'+DISK+',mode=read-write', + '--metadata-from-file=startup-script='+str(startup_path), + '--labels=experiment=continuous8-24h','--quiet','--async') + save_record(record, status='submitted', phase='request submitted; setup starts automatically when allocated') + print('Node: '+NODE+'\nLog: '+BASE+'/run.log\nCloud results: '+URI,flush=True) + print('Record: '+str(Path.home()/'continuous8-24h-resources.json'),flush=True) + except BaseException as exc: + save_record(record, status='failed', error=f'{type(exc).__name__}: {exc}') + print('STOPPED: '+str(exc)+'\nData disks and cloud archives were retained.',file=sys.stderr,flush=True) + raise + + +if __name__ == '__main__': + main() diff --git a/baseline/nanogpt_one_head/continuous8/worker.sh b/baseline/nanogpt_one_head/continuous8/worker.sh index 963750f2..9a7ae3c4 100644 --- a/baseline/nanogpt_one_head/continuous8/worker.sh +++ b/baseline/nanogpt_one_head/continuous8/worker.sh @@ -2,9 +2,14 @@ # Invoked by systemd, with Restart=no. A failed dependency install may be retried # manually before training starts. The scientific process never resumes. set -Eeuo pipefail -BASE=/mnt/disks/rg-data/continuous8 +BASE="${RG_CONTINUOUS_BASE:-/mnt/disks/rg-data/continuous8}" +SHARED="${RG_CONTINUOUS_SHARED_BASE:-$BASE}" +export RG_CONTINUOUS_BASE="$BASE" +export RG_CONTINUOUS_DATA_ROOT="$SHARED/data" REPO="$BASE/repo" EXP="$REPO/baseline/nanogpt_one_head" +CONFIG="${RG_CONTINUOUS_CONFIG:-$EXP/configs/muonclip_continuous8.yaml}" +export RG_CONTINUOUS_CONFIG="$CONFIG" # Serialize manual setup retries and reject all scientific-run reuse before pip. exec 9>"$BASE/worker.lock" flock -n 9 || { echo 'Another worker is running; refusing duplicate setup.'; exit 1; } @@ -16,11 +21,12 @@ export OMP_NUM_THREADS=4 OPENBLAS_NUM_THREADS=4 MKL_NUM_THREADS=4 export PJRT_DEVICE=TPU TPU_ACCELERATOR_TYPE=v5litepod-8 export RG_TPU_PERSISTENT_ROOT=/mnt/disks/rg-data export RG_CONTINUOUS_RUN_LOG="$BASE/run.log" -export HF_HOME="$BASE/hf-cache" +export HF_HOME="$SHARED/hf-cache" +export PIP_CACHE_DIR="$SHARED/pip-cache" unset XLA_USE_BF16 XLA_DOWNCAST_BF16 TPU_VISIBLE_CHIPS cd "$EXP" -python3 -m venv "$BASE/venv" -PY="$BASE/venv/bin/python" +python3 -m venv "$SHARED/venv" +PY="$SHARED/venv/bin/python" finish_worker() { local result=$? trap - EXIT @@ -31,7 +37,7 @@ import json, os, sys, time from pathlib import Path code=int(sys.argv[1]) status={'status':'finished' if code in (0,75) else 'failed','exit_code':code,'ended_unix':time.time()} -base=Path('/mnt/disks/rg-data/continuous8') +base=Path(os.environ['RG_CONTINUOUS_BASE']) (base/'WORKER_STATUS.json').write_text(json.dumps(status)+'\n') try: from rg_nanogpt_one_head.continuous_support import CloudPublisher @@ -56,7 +62,7 @@ import os from pathlib import Path from rg_nanogpt_one_head.continuous_support import CloudPublisher p=CloudPublisher(os.environ['RG_CONTINUOUS_GCS_URI']) -base=Path('/mnt/disks/rg-data/continuous8') +base=Path(os.environ['RG_CONTINUOUS_BASE']) p.claim({'commit':(base/'source_commit.txt').read_text().strip(), 'seed':int(os.environ['RG_CONTINUOUS_SEED']), 'automatic_restart':False}) for name in ('environment.lock.txt','source_commit.txt'): @@ -66,12 +72,13 @@ PY # Tiny, independent test: global gradients, clipping, metrics, optimizer + RNG restore. # Also measures the full proposed model shape before the long run. "$PY" -m rg_nanogpt_one_head.tpu_spmd_check --backend tpu --chips 8 \ - --benchmark-config configs/muonclip_continuous8.yaml --benchmark-steps 10 \ + --benchmark-config "$CONFIG" --benchmark-steps 10 \ --output "$BASE/preflight.json" "$PY" - <<'PYCODE' import os +from pathlib import Path from rg_nanogpt_one_head.continuous_support import CloudPublisher -CloudPublisher(os.environ['RG_CONTINUOUS_GCS_URI']).file('/mnt/disks/rg-data/continuous8/preflight.json','preflight.json') +CloudPublisher(os.environ['RG_CONTINUOUS_GCS_URI']).file(Path(os.environ['RG_CONTINUOUS_BASE'])/'preflight.json','preflight.json') PYCODE # Dataset preparation runs on the TPU VM CPU and durable disk, not Cloud Shell. "$PY" -u continuous8/prepare_tpu_data.py @@ -80,7 +87,7 @@ import json, os, time from pathlib import Path from rg_nanogpt_one_head.continuous_support import CloudPublisher p=CloudPublisher(os.environ['RG_CONTINUOUS_GCS_URI']) -path=Path('/mnt/disks/rg-data/continuous8/preflight.json') +path=Path(os.environ['RG_CONTINUOUS_BASE'])/'preflight.json' p.file(path, 'preflight.json') report=json.loads(path.read_text()) remaining=max(0,float(os.environ['RG_CONTINUOUS_DEADLINE_UNIX'])-time.time()) @@ -96,6 +103,6 @@ p.json({'status':'training'}, 'SETUP_STATUS.json') PY # No resilient supervisor, --resume, or automatic retry of scientific training. "$PY" -u -m rg_nanogpt_one_head.continuous_run \ - --config configs/muonclip_continuous8.yaml --data-root "$BASE/data" \ + --config "$CONFIG" --data-root "$RG_CONTINUOUS_DATA_ROOT" \ --results-root "$BASE/results" --device tpu --seed "$RG_CONTINUOUS_SEED" \ --deadline-unix "$RG_CONTINUOUS_DEADLINE_UNIX" diff --git a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/continuous_support.py b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/continuous_support.py index 52c9c47a..098697e7 100644 --- a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/continuous_support.py +++ b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/continuous_support.py @@ -169,6 +169,21 @@ def publish_checkpoint(path, payload): if receipt['resumable'] and path.name in {'checkpoint_latest.pt', 'checkpoint_final.pt', 'checkpoint_initial.pt'}: sink.json(receipt, 'LATEST_RESUMABLE.json') root = path.parent.parent if path.parent.name == 'epoch_checkpoints' else path.parent + if path.parent.name == 'epoch_checkpoints': + keep = int(cfg.get('continuous', {}).get('keep_local_epoch_checkpoints', 0)) + if keep > 0: + receipts = root/'checkpoint_receipts' + atomic_json(receipts/(path.name+'.json'), receipt) + verified = [] + for record in receipts.glob('*.pt.json'): + value = json.loads(record.read_text()) + candidate = path.parent/record.name.removesuffix('.json') + if candidate.is_file(): + verified.append((int(value['step']), candidate, value)) + for _, candidate, value in sorted(verified, reverse=True)[keep:]: + # Never prune an unverified or subsequently changed local file. + if candidate.stat().st_size == value['bytes'] and sha_file(candidate) == value['sha256']: + candidate.unlink() publish_metadata(cfg, root) print(f'[continuous-backup] saved step={step} file={path.name}', flush=True) diff --git a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/tpu_spmd_check.py b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/tpu_spmd_check.py index 904fa501..1cb7ed2f 100644 --- a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/tpu_spmd_check.py +++ b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/tpu_spmd_check.py @@ -12,6 +12,7 @@ import os from pathlib import Path import tempfile +import threading import time import torch @@ -163,6 +164,13 @@ def update(net, opts, target): print("[spmd-check] numerical and resume checks passed", flush=True) if benchmark_cfg is not None: print("[spmd-check] benchmarking training shape (5 warm-up updates)", flush=True) + bench_started = time.monotonic() + stage = ["model initialization"] + finished = threading.Event() + def heartbeat(): + while not finished.wait(60): + print(f"[spmd-check] waiting in {stage[0]}; benchmark elapsed={time.monotonic()-bench_started:.0f}s (no completion yet)", flush=True) + threading.Thread(target=heartbeat, daemon=True).start() batch = int(benchmark_cfg["training"]["batch_size"]) accum = int(benchmark_cfg["training"]["grad_accum_steps"]) context = int(benchmark_cfg["model"]["block_size"]) @@ -185,14 +193,22 @@ def benchmark_step(): torch.nn.utils.clip_grad_norm_(net.parameters(), float(benchmark_cfg["training"]["grad_clip"]), foreach=False) optimizer_step(opts) mark_step(device) - for _ in range(5): + for index in range(5): + stage[0] = f"warm-up update {index+1}/5" + print(f"[spmd-check] starting {stage[0]}", flush=True) benchmark_step() + stage[0] = "warm-up device synchronization" synchronize(device) + print("[spmd-check] warm-up completed", flush=True) start = time.perf_counter() - for _ in range(args.benchmark_steps): + for index in range(args.benchmark_steps): + stage[0] = f"timed update {index+1}/{args.benchmark_steps}" + print(f"[spmd-check] submitting {stage[0]}", flush=True) benchmark_step() + stage[0] = "timed device synchronization" synchronize(device) seconds = time.perf_counter() - start + finished.set() report["benchmark"] = {"steps": args.benchmark_steps, "seconds": seconds, "global_tokens_per_update": batch * accum * context, "tokens_per_second": args.benchmark_steps * batch * accum * context / seconds, diff --git a/baseline/nanogpt_one_head/tests/test_continuous8.py b/baseline/nanogpt_one_head/tests/test_continuous8.py index c25b124f..3050c069 100644 --- a/baseline/nanogpt_one_head/tests/test_continuous8.py +++ b/baseline/nanogpt_one_head/tests/test_continuous8.py @@ -90,6 +90,56 @@ def test_fixed_config_alignment_and_schedule(): assert lr_schedule_steps(cfg,p)==100_000 and warmup_steps(p,100_000)==2000 +def test_24h_config_has_aligned_pairs_without_changing_training_dynamics(): + from rg_nanogpt_one_head.muonclip import install_muonclip_extension + install_muonclip_extension() + from rg_nanogpt_one_head.config import load_config, epoch_step_map + root = Path(__file__).resolve().parents[1]/'configs' + before = load_config(root/'muonclip_continuous8.yaml') + after = load_config(root/'muonclip_continuous8_24h.yaml') + for key in ('dataset','model','optimizer_profiles','evaluation','runtime'): + assert before[key] == after[key] + assert list(epoch_step_map(after)) == list(range(0,1_000_001,500)) + assert after['training']['checkpoint_interval_steps'] == 500 + assert after['training']['eval_interval_steps'] == 500 + assert after['continuous']['max_wall_hours'] == 23.5 + assert after['continuous']['keep_local_epoch_checkpoints'] == 3 + + +def test_only_uploaded_unchanged_epoch_checkpoints_are_pruned(tmp_path, monkeypatch): + class Sink: + fail = False + def file(self, path, relative): + if self.fail: raise OSError('upload failed') + return {'generation':'1', 'bytes':path.stat().st_size} + def json(self, *args): pass + sink = Sink() + monkeypatch.setattr(cs, 'publisher', lambda cfg:sink) + monkeypatch.setattr(cs, 'publish_metadata', lambda *args:None) + directory = tmp_path/'epoch_checkpoints' + directory.mkdir() + unknown = directory/'unknown.pt' + unknown.write_bytes(b'not uploaded') + config = {'continuous':{'keep_local_epoch_checkpoints':2}} + def publish(step): + path = directory/f'epoch_{step:06d}.pt' + path.write_bytes(str(step).encode()) + cs.publish_checkpoint(path, dict(config=config,step=step,fingerprint='f',model_state_sha256='m')) + return path + first = publish(1) + second = publish(2) + first.write_bytes(b'x') # Same size, changed contents must survive pruning. + third = publish(3) + assert first.exists() and unknown.exists() + first.write_bytes(b'1') + fourth = publish(4) + assert not first.exists() and not second.exists() + assert third.exists() and fourth.exists() and unknown.exists() + sink.fail = True + with pytest.raises(OSError, match='upload failed'): publish(5) + assert third.exists() and fourth.exists() and (directory/'epoch_000005.pt').exists() + + def test_muonclip_checkpoint_preserves_next_updates_lr_optimizer_and_sampler(tmp_path): from copy import deepcopy from rg_nanogpt_one_head.muonclip import install_muonclip_extension diff --git a/baseline/nanogpt_one_head/tests/test_reallocate24.py b/baseline/nanogpt_one_head/tests/test_reallocate24.py new file mode 100644 index 00000000..795a85a6 --- /dev/null +++ b/baseline/nanogpt_one_head/tests/test_reallocate24.py @@ -0,0 +1,70 @@ +"""Simulate the destructive provisioning sequence without contacting GCP.""" +import importlib.util +from pathlib import Path +import subprocess +import pytest + + +def launcher(monkeypatch): + folder = Path(__file__).resolve().parents[1]/'continuous8' + monkeypatch.syspath_prepend(str(folder)) + spec = importlib.util.spec_from_file_location('reallocation',folder/'reallocate24.py') + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def test_startup_reuses_disk_without_format_and_has_full_deadline(monkeypatch): + m = launcher(monkeypatch) + source = m.make_startup('a'*40) + subprocess.run(['bash','-n'],input=source,text=True,check=True) + assert 'mkfs' not in source and '__' not in source + assert '+23.5*3600' in source + assert 'Restart=no' in source + assert 'BASE='+m.BASE in source + assert 'RG_CONTINUOUS_SHARED_BASE='+m.SHARED in source + assert 'muonclip_continuous8_24h.yaml' in source + + +@pytest.mark.parametrize('existing,fail_delete',[(False,False),(True,False),(False,True)]) +def test_delete_before_create_preserve_disk_and_avoid_duplicate(tmp_path,monkeypatch,existing,fail_delete): + m = launcher(monkeypatch) + monkeypatch.setattr(Path,'home',classmethod(lambda cls:tmp_path)) + monkeypatch.setattr(m.subprocess,'check_output',lambda args,**kw:'' if 'status' in args else 'a'*40+'\n') + # Keep real bash syntax validation; only cloud calls are simulated. + queues = {z:[{'name':f'old-{z}'}] for z in m.ZONES} + if existing: queues[m.ZONE].append({'name':m.QUEUE}) + nodes = {z:[] for z in m.ZONES} + nodes[m.ZONE] = [{'name':m.OLD_NODE}] + nodes['us-east5-a'] = [{'name':'standalone-old'}] + calls = [] + def inventory(*args): + zone = next((x.split('=',1)[1] for x in args if x.startswith('--zone=')),m.ZONE) + if 'queued-resources' in args: return queues[zone] + if 'tpu-vm' in args: return nodes[zone] + if args[:3] == ('compute','disks','describe'): return {'name':m.DISK,'zone':m.ZONE,'users':[]} + return {} + def gc(*args,**kw): + calls.append(args) + zone = next((x.split('=',1)[1] for x in args if x.startswith('--zone=')),m.ZONE) + if 'delete' in args: + if fail_delete: raise RuntimeError('deletion failed') + if 'queued-resources' in args: + queues[zone] = [] + if zone == m.ZONE: nodes[zone] = [] + else: nodes[zone] = [] + if 'create' in args: + assert all(not values for values in queues.values()) + assert all(not values for values in nodes.values()) + assert '--max-run-duration=24h' in args + assert '--accelerator-type=v5litepod-8' in args + assert any(m.DISK in x for x in args if x.startswith('--data-disk=')) + monkeypatch.setattr(m,'inventory',inventory) + monkeypatch.setattr(m,'gc',gc) + if fail_delete: + with pytest.raises(RuntimeError,match='deletion failed'):m.main() + else:m.main() + creates = [c for c in calls if 'create' in c] + assert len(creates) == (0 if existing or fail_delete else 1) + assert not any('disks' in c and 'delete' in c for c in calls) + if existing: assert calls == [] From 3508a96c8eed61eaf063a4ffffe1b223e0643d1a Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Sat, 3 Oct 2026 17:31:46 -0700 Subject: [PATCH 10/36] Add GPT-2 Small FineWeb MuonClip WeightWatcher experiment --- baseline/nanogpt_one_head/README.md | 6 + .../configs/gpt2_small_cpu_smoke.yaml | 56 ++++ .../gpt2_small_fineweb_adamw_baseline.yaml | 56 ++++ .../gpt2_small_fineweb_muonclip_baseline.yaml | 68 +++++ .../gpt2_small_fineweb_muonclip_long_ww.yaml | 69 +++++ .../gpt2small/LOCAL_VALIDATION.json | 23 ++ baseline/nanogpt_one_head/gpt2small/README.md | 119 ++++++++ .../nanogpt_one_head/gpt2small/analyze.py | 44 +++ .../nanogpt_one_head/gpt2small/cloudshell.py | 27 ++ baseline/nanogpt_one_head/gpt2small/on_tpu.sh | 37 +++ .../gpt2small/prepare_existing.py | 68 +++++ .../nanogpt_one_head/gpt2small/validate.py | 77 +++++ .../rg_nanogpt_one_head/gpt2_experiment.py | 289 ++++++++++++++++++ .../tests/test_gpt2_experiment.py | 100 ++++++ 14 files changed, 1039 insertions(+) create mode 100644 baseline/nanogpt_one_head/configs/gpt2_small_cpu_smoke.yaml create mode 100644 baseline/nanogpt_one_head/configs/gpt2_small_fineweb_adamw_baseline.yaml create mode 100644 baseline/nanogpt_one_head/configs/gpt2_small_fineweb_muonclip_baseline.yaml create mode 100644 baseline/nanogpt_one_head/configs/gpt2_small_fineweb_muonclip_long_ww.yaml create mode 100644 baseline/nanogpt_one_head/gpt2small/LOCAL_VALIDATION.json create mode 100644 baseline/nanogpt_one_head/gpt2small/README.md create mode 100644 baseline/nanogpt_one_head/gpt2small/analyze.py create mode 100644 baseline/nanogpt_one_head/gpt2small/cloudshell.py create mode 100644 baseline/nanogpt_one_head/gpt2small/on_tpu.sh create mode 100644 baseline/nanogpt_one_head/gpt2small/prepare_existing.py create mode 100644 baseline/nanogpt_one_head/gpt2small/validate.py create mode 100644 baseline/nanogpt_one_head/src/rg_nanogpt_one_head/gpt2_experiment.py create mode 100644 baseline/nanogpt_one_head/tests/test_gpt2_experiment.py diff --git a/baseline/nanogpt_one_head/README.md b/baseline/nanogpt_one_head/README.md index 2b7f0d64..369fca48 100644 --- a/baseline/nanogpt_one_head/README.md +++ b/baseline/nanogpt_one_head/README.md @@ -383,3 +383,9 @@ defaults to one six-hour v5e-8 allocation, with a two-machine/four-hour option. Data preparation and training run on the TPU VM and its persistent disk, independently of Cloud Shell. Automatic restarts and the segmented continuation runner are disabled; the allocation cap includes setup and tokenization. + +## GPT-2 Small validation + +See [gpt2small/README.md](gpt2small/README.md) for the isolated context-1024 +AdamW/MuonClip validation, corpus reuse, bounded checkpoints and resume tests. +The long experiment requires a separate explicit launch. diff --git a/baseline/nanogpt_one_head/configs/gpt2_small_cpu_smoke.yaml b/baseline/nanogpt_one_head/configs/gpt2_small_cpu_smoke.yaml new file mode 100644 index 00000000..1936b785 --- /dev/null +++ b/baseline/nanogpt_one_head/configs/gpt2_small_cpu_smoke.yaml @@ -0,0 +1,56 @@ +run_id: tiny_cpu_smoke +seed: 1337 +model: + vocab_size: 64 + block_size: 8 + n_layer: 2 + n_head: 2 + n_embd: 32 + dropout: 0.0 + bias: true + tie_weights: true +dataset: + name: HuggingFaceFW/fineweb-edu + config: sample-10BT + split: train + revision: 593b3a867298afb8ce42625a270ef20ddcad28f9 + tokenizer: gpt2 + encoding_workers: 16 + encoding_batch_size: 256 + train_tokens: 512 + val_tokens: 128 + test_tokens: 128 +runtime: + matmul_precision: highest + mps_fallback: true + deterministic_algorithms: false + empty_mps_cache_after_weightwatcher: true + tpu_spmd: false + tpu_expected_chips: 8 +training: + batch_size: 2 + grad_accum_steps: 1 + max_steps: 4 + max_tokens: 64 + warmup_steps: 1 + schedule_steps: 4 + grad_clip: 1.0 +eval_batches: 1 +metrics_interval: 2 +ww: + enabled: false + interval: 0 + steps: [] + logarithmic: false + min_evals: 20 +optimizer: + display_name: AdamW + family: adamw + learning_rate: 0.0006 + min_learning_rate: 6.0e-05 + warmup_fraction: 0.01 + schedule: warmup_cosine + beta1: 0.9 + beta2: 0.95 + epsilon: 1.0e-08 + weight_decay: 0.1 diff --git a/baseline/nanogpt_one_head/configs/gpt2_small_fineweb_adamw_baseline.yaml b/baseline/nanogpt_one_head/configs/gpt2_small_fineweb_adamw_baseline.yaml new file mode 100644 index 00000000..9bc415c9 --- /dev/null +++ b/baseline/nanogpt_one_head/configs/gpt2_small_fineweb_adamw_baseline.yaml @@ -0,0 +1,56 @@ +run_id: gpt2_small_fineweb_adamw_baseline_s1337 +seed: 1337 +model: + vocab_size: 50257 + block_size: 1024 + n_layer: 12 + n_head: 12 + n_embd: 768 + dropout: 0.0 + bias: true + tie_weights: true +dataset: + name: HuggingFaceFW/fineweb-edu + config: sample-10BT + split: train + revision: 593b3a867298afb8ce42625a270ef20ddcad28f9 + tokenizer: gpt2 + encoding_workers: 16 + encoding_batch_size: 256 + train_tokens: 5000000000 + val_tokens: 10000000 + test_tokens: 10000000 +runtime: + matmul_precision: highest + mps_fallback: true + deterministic_algorithms: false + empty_mps_cache_after_weightwatcher: true + tpu_spmd: true + tpu_expected_chips: 8 +training: + batch_size: 8 + grad_accum_steps: 4 + max_steps: 1000 + max_tokens: 32768000 + warmup_steps: 20 + schedule_steps: 1000 + grad_clip: 1.0 +eval_batches: 2 +metrics_interval: 25 +ww: + enabled: false + interval: 0 + steps: [] + logarithmic: false + min_evals: 20 +optimizer: + display_name: AdamW + family: adamw + learning_rate: 0.0006 + min_learning_rate: 6.0e-05 + warmup_fraction: 0.01 + schedule: warmup_cosine + beta1: 0.9 + beta2: 0.95 + epsilon: 1.0e-08 + weight_decay: 0.1 diff --git a/baseline/nanogpt_one_head/configs/gpt2_small_fineweb_muonclip_baseline.yaml b/baseline/nanogpt_one_head/configs/gpt2_small_fineweb_muonclip_baseline.yaml new file mode 100644 index 00000000..1472b7b9 --- /dev/null +++ b/baseline/nanogpt_one_head/configs/gpt2_small_fineweb_muonclip_baseline.yaml @@ -0,0 +1,68 @@ +run_id: gpt2_small_fineweb_muonclip_baseline_s1337 +seed: 1337 +model: + vocab_size: 50257 + block_size: 1024 + n_layer: 12 + n_head: 12 + n_embd: 768 + dropout: 0.0 + bias: true + tie_weights: true +dataset: + name: HuggingFaceFW/fineweb-edu + config: sample-10BT + split: train + revision: 593b3a867298afb8ce42625a270ef20ddcad28f9 + tokenizer: gpt2 + encoding_workers: 16 + encoding_batch_size: 256 + train_tokens: 5000000000 + val_tokens: 10000000 + test_tokens: 10000000 +runtime: + matmul_precision: highest + mps_fallback: true + deterministic_algorithms: false + empty_mps_cache_after_weightwatcher: true + tpu_spmd: true + tpu_expected_chips: 8 +training: + batch_size: 8 + grad_accum_steps: 4 + max_steps: 1000 + max_tokens: 32768000 + warmup_steps: 20 + schedule_steps: 1000 + grad_clip: 1.0 +eval_batches: 2 +metrics_interval: 25 +ww: + enabled: true + interval: 0 + steps: + - 25 + - 1000 + logarithmic: false + min_evals: 20 +optimizer: + display_name: MuonClip + RMS-matched updates + auxiliary AdamW + family: muon_clip + learning_rate: 0.02 + min_learning_rate: 0.002 + warmup_fraction: 0.02 + schedule: warmup_cosine + momentum: 0.95 + nesterov: true + newton_schulz_steps: 5 + muon_epsilon: 1.0e-07 + weight_decay: 0.1 + update_rms_scale: 0.2 + qk_clip_threshold: 100.0 + qk_clip_balance: 0.5 + qk_diagnostics_interval: 25 + beta1: 0.9 + beta2: 0.95 + epsilon: 1.0e-08 + aux_learning_rate: 0.0006 + aux_min_learning_rate: 6.0e-05 diff --git a/baseline/nanogpt_one_head/configs/gpt2_small_fineweb_muonclip_long_ww.yaml b/baseline/nanogpt_one_head/configs/gpt2_small_fineweb_muonclip_long_ww.yaml new file mode 100644 index 00000000..24ee84ff --- /dev/null +++ b/baseline/nanogpt_one_head/configs/gpt2_small_fineweb_muonclip_long_ww.yaml @@ -0,0 +1,69 @@ +run_id: gpt2_small_fineweb_muonclip_long_ww_s1337 +seed: 1337 +model: + vocab_size: 50257 + block_size: 1024 + n_layer: 12 + n_head: 12 + n_embd: 768 + dropout: 0.0 + bias: true + tie_weights: true +dataset: + name: HuggingFaceFW/fineweb-edu + config: sample-10BT + split: train + revision: 593b3a867298afb8ce42625a270ef20ddcad28f9 + tokenizer: gpt2 + encoding_workers: 16 + encoding_batch_size: 256 + train_tokens: 5000000000 + val_tokens: 10000000 + test_tokens: 10000000 +runtime: + matmul_precision: highest + mps_fallback: true + deterministic_algorithms: false + empty_mps_cache_after_weightwatcher: true + tpu_spmd: true + tpu_expected_chips: 8 +training: + batch_size: 8 + grad_accum_steps: 4 + max_tokens: 100000000000 + warmup_steps: 2000 + schedule_steps: 3051758 + grad_clip: 1.0 +eval_batches: 2 +metrics_interval: 500 +ww: + enabled: true + interval: 10000 + steps: + - 100 + - 500 + - 1000 + logarithmic: false + min_evals: 20 +optimizer: + display_name: MuonClip + RMS-matched updates + auxiliary AdamW + family: muon_clip + learning_rate: 0.02 + min_learning_rate: 0.002 + warmup_fraction: 0.02 + schedule: warmup_cosine + momentum: 0.95 + nesterov: true + newton_schulz_steps: 5 + muon_epsilon: 1.0e-07 + weight_decay: 0.1 + update_rms_scale: 0.2 + qk_clip_threshold: 100.0 + qk_clip_balance: 0.5 + qk_diagnostics_interval: 25 + beta1: 0.9 + beta2: 0.95 + epsilon: 1.0e-08 + aux_learning_rate: 0.0006 + aux_min_learning_rate: 6.0e-05 +long_run: true diff --git a/baseline/nanogpt_one_head/gpt2small/LOCAL_VALIDATION.json b/baseline/nanogpt_one_head/gpt2small/LOCAL_VALIDATION.json new file mode 100644 index 00000000..ec66a161 --- /dev/null +++ b/baseline/nanogpt_one_head/gpt2small/LOCAL_VALIDATION.json @@ -0,0 +1,23 @@ +{ + "environment": "CPU, torch 2.6.0, Python 3.12; no TPU access", + "model_parameters": 124439808, + "real_weightwatcher": { + "records": 72, + "seconds": 109.84400583000024, + "raw_success": 72, + "clipped_success": 72, + "randomized_success": 72 + }, + "tpu_validation": "NOT RUN", + "long_run": "NOT STARTED", + "tests": { + "new_integration_tests": "7 passed", + "focused_training_optimizer_data_spmd_tests": "43 passed, 1 skipped (before adding one transaction-recovery test)", + "repository_suite": "221 passed, 1 skipped, 3 failed", + "existing_failures": [ + "test_single_matrix_wrapper_preserves_loader_metadata", + "test_notebook_displays_the_actual_results_variable", + "test_repo_shell_scripts_never_reference_home" + ] + } +} diff --git a/baseline/nanogpt_one_head/gpt2small/README.md b/baseline/nanogpt_one_head/gpt2small/README.md new file mode 100644 index 00000000..21568297 --- /dev/null +++ b/baseline/nanogpt_one_head/gpt2small/README.md @@ -0,0 +1,119 @@ +# GPT-2 Small / FineWeb-Edu validation + +This is a new, isolated experiment. Existing tiny and continuous8 workflows are unchanged. +It reuses the repository GPT, advanced MuonClip, SPMD, corpus validation and WeightWatcher adapter. + +## Model and data + +124,439,808 trainable parameters: 12 independent blocks, 12 heads, width 768, +head dimension 64, MLP width 3072, context 1024, vocabulary 50257. GELU, LayerNorm, +causal attention and tied embedding/output weights; biases enabled. The six separate +projection matrices per block yield exactly 72 named trajectories. No QKV fusion or +layer sharing is introduced. + +The existing corpus is `/mnt/disks/rg-data/continuous8/data`: 5 billion training +GPT-2 tokens and 10 million tokens each for validation/test, stored as three uint16 +files with document indexes. Every launch validates metadata, sizes, SHA256 identities +and the document-disjoint split contract. This workflow never prepares/downloads data. +Evaluation uses deterministic windows within each split; train/validation/test probes +are fixed across optimizers and resumes. They are not the speedrun benchmark protocol. + +## Configurations + +- `configs/gpt2_small_fineweb_adamw_baseline.yaml`: AdamW 6e-4, betas .9/.95, + weight decay .1, gradient clipping 1.0. +- `configs/gpt2_small_fineweb_muonclip_baseline.yaml`: existing advanced MuonClip, + matrix LR .02, momentum .95, Nesterov, five NS iterations, QK threshold 100; + auxiliary AdamW LR 6e-4 independently configurable. +- `configs/gpt2_small_fineweb_muonclip_long_ww.yaml`: 100B token presentations; + refuses CLI launch without `--allow-long-run`. No loss-based stopping. +- `configs/gpt2_small_cpu_smoke.yaml`: tiny synthetic-corpus CPU integration tests. + +**LR caveat:** advanced MuonClip scales the orthogonal update by +`update_rms_scale * sqrt(max(matrix.shape))`, here .2 times that square root. +The requested .02 is retained explicitly, but is not equivalent to the plain Muon +parameterization. Its stability is NOT established by this configuration. Do not +call this a reproduced reference benchmark or launch a long run before validating it. + +Global microbatch 8, accumulation 4, context 1024 = 32,768 tokens/update, split over +eight chips. Batch size, accumulation, warmup, cosine horizon, floor, max_steps and +max_tokens are configurable. Stop-after is an interruption point, not an LR-horizon +change. max_tokens is rounded up to the next complete optimizer update. + +## Run only the short validation on the existing TPU + +From a clean Cloud Shell checkout of this commit: + +```bash +python3 baseline/nanogpt_one_head/gpt2small/cloudshell.py +``` + +This verifies the existing active queue/node, gracefully stops the old trainer, +archives compact scientific results and verifies the archive, then deletes only +recognized model checkpoint files in the known old run directory. It does not +format/delete disks, remove the corpus/caches/indexes, or allocate a replacement TPU. +Old scalar records are retained as well as archived; ambiguous output is not removed. + +The exact commit is checked out under a fresh persistent-disk validation directory. +The sequence is AdamW through step 4, MuonClip through step 4, then fresh-process +resume of each through step 25. The baseline horizon stays fixed at 1000 steps. +MuonClip WW runs at steps 4 and 25. All earlier metric/WW file hashes must survive +resume unchanged. Short validation stops on nonfinite metrics, increasing training +probe NLL, incomplete matrix inventory or null controls. Failed alpha fits are retained +as NaN with status/reason, never replaced with clipped alpha. + +The cutoff is read from the previous allocation record; five minutes are reserved +for checkpoint/backup. Setup/compilation/WW may consume the remaining window: an +incomplete validation is reported honestly, and no new machine is requested. +Outputs are copied to `gs://tpu-builders-504820-ww-continuous8/gpt2small/` on exit; +a failed cloud copy is an error and leaves persistent-disk outputs intact. + +25 updates prove plumbing and initial learning direction only. They do not establish +convergence or reproduce CE 3.28. A longer matched benchmark is needed before claiming +normal GPT-2 training quality, and before the long scientific experiment. + +## Records, checkpoints and timing + +Per-step immutable JSON scalar and WW records are written incrementally. NLL, +perplexity, top-1 accuracy, error (fraction), steps, token presentations, wall time, +LR and pre-clipping gradient norm are included. All WW library columns are retained, +including raw/clipped fits, randomized distance, bounds, KS/D, spectral measures and +fit status. `xmin/xmax/D` are library-returned clipped-fit fields, not invented raw-fit +bounds. Raw failure is explicit; alpha is never constrained toward 2. + +Three rolling full checkpoints contain model, both optimizer states where applicable, +step-derived cosine scheduler state/config, token count, data sampler RNG, all training +RNG states and run identity. Checkpoint writes are atomic; old checkpoints are pruned +only after publishing the new pointer. A checkpoint includes pending metrics/WW records: +resume completes missing writes, appends later records, and refuses histories ahead of +the checkpoint rather than erasing data. Config/data/software version mismatch fails +closed. Optional milestone copies are independent of the rolling set. A single writer +lock prevents concurrent mutation. CPU exact-resume tests compare model and optimizer +states bit for bit. Numerical equivalence on TPU replacement still requires TPU testing. + +The first two update durations include compilation and execution and are reported +separately; they are **not pure compiler timing**. Native XLA CompileTime/ExecuteTime +metrics are also saved in `logs/xla_compile_metrics_after_step_*.txt`. Later synchronized update durations +measure training throughput, excluding evaluation/WW/checkpoint overhead. End-to-end +throughput is reported separately. The validation runner intentionally synchronizes +once per update to measure completed TPU work; the long configuration only synchronizes +at measurement boundaries after the first two updates. This is not a reproduced speedrun. + +WW supports fixed `interval`, explicit `steps`, and `logarithmic` 1/2/5-per-decade +schedules. A full pass reports seconds and recommends at least 9x that duration of +training between passes for at most 10% WW-only overhead. Include other overhead when +selecting the eventual long-run schedule. Scalar checkpoints are bounded in count; +no thousands of full checkpoint copies are retained. + +## Analysis and tests + +```bash +python gpt2small/analyze.py /mnt/disks/rg-data/gpt2small/VALIDATION/muonclip +PYTHONPATH=src python -m pytest tests/test_gpt2_experiment.py -q +``` + +Analysis creates CSVs and loss/perplexity/error-vs-token plots; mean/min raw and clipped +alpha trajectories; raw-alpha/error regression plots; all-layer raw/clipped plots +by matrix type; and per-matrix initial/latest/min/delta/recent slope/recent variance. +Correlations along one trajectory do not establish causation or independent-sample +significance. Tokens are presentations, not necessarily unique training tokens. diff --git a/baseline/nanogpt_one_head/gpt2small/analyze.py b/baseline/nanogpt_one_head/gpt2small/analyze.py new file mode 100644 index 00000000..35c5faf8 --- /dev/null +++ b/baseline/nanogpt_one_head/gpt2small/analyze.py @@ -0,0 +1,44 @@ +"""python analyze.py RUN_DIRECTORY: immutable records -> CSV and scientific plots.""" +import json +from pathlib import Path +import sys +import numpy as np +import pandas as pd +import matplotlib +matplotlib.use('Agg') +import matplotlib.pyplot as plt +root = Path(sys.argv[1]); out = root / 'summaries'; out.mkdir(exist_ok=True) +m = pd.DataFrame([json.loads(p.read_text()) for p in sorted((root/'metrics').glob('*.json'))]) +if m.empty: raise SystemExit('No measurements yet') +m.to_csv(out/'metrics.csv',index=False) +for metric in ('nll','perplexity','token_error'): + fig, ax=plt.subplots() + for split in ('train','val','test'): ax.plot(m.tokens_seen,m[f'{split}_{metric}'],label=split) + ax.set(xlabel='Token presentations',ylabel=metric); ax.legend(); fig.tight_layout(); fig.savefig(out/f'{metric}_vs_tokens.png'); plt.close(fig) +rows=[r for p in sorted((root/'ww_metrics').glob('*.json')) for r in json.loads(p.read_text())['records']] +if not rows: raise SystemExit('Scalar plots saved; no WW snapshots yet') +w=pd.DataFrame(rows); w.to_csv(out/'ww_metrics.csv',index=False) +summary=[] +for variant in ('alpha_raw','alpha_clip_xmax'): + grouped=w.groupby('tokens_seen')[variant].agg(['mean','min','std']).reset_index() + for stat in ('mean','min'): + fig,ax=plt.subplots(); ax.plot(grouped.tokens_seen,grouped[stat],'o-'); ax.set(xlabel='Token presentations',ylabel=f'{stat} {variant}'); fig.tight_layout(); fig.savefig(out/f'{stat}_{variant}_vs_tokens.png'); plt.close(fig) + if variant=='alpha_raw': + paired=grouped.merge(m[['tokens_seen','test_token_error']],on='tokens_seen',validate='one_to_one').dropna() + fig,ax=plt.subplots(); ax.scatter(paired[stat],paired.test_token_error) + if len(paired)>2 and paired[stat].nunique()>1: + slope,intercept=np.polyfit(paired[stat],paired.test_token_error,1); x=np.sort(paired[stat]); ax.plot(x,intercept+slope*x) + ax.set(xlabel=f'{stat} raw alpha',ylabel='Test token error'); fig.tight_layout(); fig.savefig(out/f'test_error_vs_{stat}_raw.png'); plt.close(fig) + fig,axes=plt.subplots(3,2,figsize=(13,12)) + for ax,kind in zip(axes.flat,('Q','K','V','O','MLP_IN','MLP_OUT')): + for name,g in w[w.matrix_type==kind].groupby('matrix_name'): + g=g.sort_values('tokens_seen'); ax.plot(g.tokens_seen,g[variant],label=name) + ax.axhline(2,color='gray',ls=':'); ax.set(title=kind,xlabel='Token presentations',ylabel=variant); ax.legend(fontsize=6,ncol=3) + fig.tight_layout(); fig.savefig(out/f'all_layers_{variant}.png'); plt.close(fig) + for name,g in w.groupby('matrix_name'): + g=g.sort_values('tokens_seen').dropna(subset=[variant]); recent=g.tail(5) + if g.empty: continue + slope=np.polyfit(recent.tokens_seen,recent[variant],1)[0] if len(recent)>1 else np.nan + summary.append({'matrix_name':name,'variant':variant,'initial':g[variant].iloc[0],'latest':g[variant].iloc[-1], 'minimum':g[variant].min(),'delta':g[variant].iloc[-1]-g[variant].iloc[0], 'recent_slope_per_token':slope,'recent_variance':recent[variant].var()}) +pd.DataFrame(summary).to_csv(out/'matrix_summary.csv',index=False) +print(out) diff --git a/baseline/nanogpt_one_head/gpt2small/cloudshell.py b/baseline/nanogpt_one_head/gpt2small/cloudshell.py new file mode 100644 index 00000000..71967237 --- /dev/null +++ b/baseline/nanogpt_one_head/gpt2small/cloudshell.py @@ -0,0 +1,27 @@ +"""Use existing TPU only. No allocation, formatting, disk deletion or corpus download.""" +from pathlib import Path +import datetime +import json +import shlex +import subprocess + +PROJECT='tpu-builders-504820'; ZONE='us-west4-a' +QUEUE='ww-continuous8-24h-20261003-s1337'; NODE=QUEUE+'-node' +HERE=Path(__file__).resolve().parent; REPO=HERE.parents[2] +def run(*args,capture=False): + return subprocess.run(args,check=True,text=True,stdout=subprocess.PIPE if capture else None).stdout +commit=run('git','-C',str(REPO),'rev-parse','HEAD',capture=True).strip() +if run('git','-C',str(REPO),'status','--porcelain',capture=True).strip(): raise SystemExit('Use a clean checkout of the pushed commit') +args=['--project='+PROJECT,'--zone='+ZONE] +queue=json.loads(run('gcloud','alpha','compute','tpus','queued-resources','describe',QUEUE,*args,'--format=json',capture=True)) +if queue.get('state',{}).get('state')!='ACTIVE': raise SystemExit('Existing TPU is not ACTIVE. No replacement will be allocated.') +print(json.dumps(queue,indent=2),flush=True) +node=json.loads(run('gcloud','compute','tpus','tpu-vm','describe',NODE,*args,'--format=json',capture=True)) +print(json.dumps(node,indent=2),flush=True) +root='/mnt/disks/rg-data/gpt2small/validation-'+datetime.datetime.now(datetime.timezone.utc).strftime('%Y%m%d-%H%M%S') +remote='/tmp/rg-gpt2-validation-'+commit[:12] +run('gcloud','compute','tpus','tpu-vm','ssh',NODE,*args,'--worker=0','--command=mkdir -p '+shlex.quote(remote)) +run('gcloud','compute','tpus','tpu-vm','scp',str(HERE/'on_tpu.sh'),str(HERE/'prepare_existing.py'),NODE+':'+remote+'/',*args,'--worker=0') +command=shlex.join(['bash',remote+'/on_tpu.sh',commit,root]) +run('gcloud','compute','tpus','tpu-vm','ssh',NODE,*args,'--worker=0','--command='+command) +print('Saved validation:',root) diff --git a/baseline/nanogpt_one_head/gpt2small/on_tpu.sh b/baseline/nanogpt_one_head/gpt2small/on_tpu.sh new file mode 100644 index 00000000..f98f2040 --- /dev/null +++ b/baseline/nanogpt_one_head/gpt2small/on_tpu.sh @@ -0,0 +1,37 @@ +#!/usr/bin/env bash +set -euo pipefail +commit=$1 +root=$2 +mountpoint -q /mnt/disks/rg-data +shared=/mnt/disks/rg-data/continuous8 +python="$shared/venv/bin/python" +test -x "$python" +# This script and prepare_existing.py are transferred from the pushed commit by Cloud Shell. +sudo "$python" "$(dirname "$0")/prepare_existing.py" "$root" +sudo chown -R "$(id -u):$(id -g)" "$root" +repo="$root/repo" +git clone --no-checkout https://github.com/CalculatedContent/rg_optimizers.git "$repo" +git -C "$repo" checkout --detach "$commit" +test "$(git -C "$repo" rev-parse HEAD)" = "$commit" +git -C "$repo" rev-parse HEAD +export PYTHONPATH="$repo/baseline/nanogpt_one_head/src" +export OMP_NUM_THREADS=4 OPENBLAS_NUM_THREADS=4 MKL_NUM_THREADS=4 +export PJRT_DEVICE=TPU +cd "$repo/baseline/nanogpt_one_head" +"$python" -c 'import torch, torch_xla, weightwatcher, yaml; print("Installed dependencies loaded")' +deadline=$("$python" -c 'import json,sys; print(json.load(open(sys.argv[1]))["deadline_unix"])' "$root/old_allocation.json") +bucket=gs://tpu-builders-504820-ww-continuous8/gpt2small +backup() { + rc=$? + trap - EXIT + sync + # Keep backups separate from prior pilot archives. No remote deletion. + if ! gcloud storage rsync "$root" "$bucket/$(basename "$root")" --recursive --exclude='repo/.*' --project=tpu-builders-504820; then + echo "Cloud backup FAILED; all outputs remain on persistent disk: $root" >&2 + exit 1 + fi + exit "$rc" +} +trap backup EXIT +"$python" -u gpt2small/validate.py --root "$root" --data "$shared/data" --deadline "$deadline" 2>&1 | tee "$root/validation.log" +echo 'Validation finished. No long experiment launched.' diff --git a/baseline/nanogpt_one_head/gpt2small/prepare_existing.py b/baseline/nanogpt_one_head/gpt2small/prepare_existing.py new file mode 100644 index 00000000..2080f13e --- /dev/null +++ b/baseline/nanogpt_one_head/gpt2small/prepare_existing.py @@ -0,0 +1,68 @@ +"""Run as root on the known VM: graceful stop, compact archive, narrow checkpoint cleanup.""" +import hashlib +import json +import os +from pathlib import Path +import subprocess +import sys +import tarfile +import time +MOUNT=Path('/mnt/disks/rg-data') +OLD=MOUNT/'continuous8-24h-20261003-s1337' +DATA=MOUNT/'continuous8/data' +ROOT=Path(sys.argv[1]) +if ROOT.resolve() == MOUNT/'gpt2small' or not ROOT.resolve().is_relative_to(MOUNT/'gpt2small'): + raise SystemExit('Unexpected new output root; no cleanup') +if OLD.resolve() != OLD: + raise SystemExit('Unexpected old output symlink; no cleanup') +if not os.path.ismount(MOUNT): raise SystemExit('Persistent disk is not mounted; no cleanup') +if DATA.is_symlink() or not DATA.resolve().is_relative_to(MOUNT): raise SystemExit('Unexpected corpus path') +meta=json.loads((DATA/'meta.json').read_text()) +if meta.get('tokenizer')!='gpt2' or not meta.get('document_disjoint_splits'): raise SystemExit('Corpus identity/isolation check failed') +for split in ('train','val','test'): + f=DATA/f'{split}.bin' + if f.stat().st_size!=2*meta['splits'][split]: raise SystemExit('Corpus file size mismatch') +ROOT.mkdir(parents=True,exist_ok=True) +subprocess.run(['df','-h',str(MOUNT)],check=True) +subprocess.run(['du','-sh',*[str(p) for p in MOUNT.iterdir()]],check=True) +subprocess.run(['systemctl','--no-pager','--full','status','rg-continuous8.service'],check=False) +claim=json.loads((OLD/'results/CONTINUOUS_STARTED.json').read_text()) +(ROOT/'old_allocation.json').write_text(json.dumps(claim,indent=2)) +(ROOT/'data_inventory.json').write_text(json.dumps({'path':str(DATA),'bytes':sum(p.stat().st_size for p in DATA.rglob('*') if p.is_file()),'token_shards':3,'metadata':meta},indent=2)) +(OLD/'results/STOP').touch() +# Let the old trainer finish its current update and final checkpoint; never SIGKILL it here. +end=time.monotonic()+600 +while True: + running=[] + for p in Path('/proc').glob('[0-9]*/cmdline'): + try: args=p.read_bytes().split(b'\0') + except (OSError,ProcessLookupError): continue + if b'rg_nanogpt_one_head.continuous_run' in args: running.append(p.parent.name) + if not running: break + if time.monotonic()>end: raise SystemExit('Graceful stop timed out; no files deleted') + time.sleep(5) +subprocess.run(['systemctl','stop','rg-continuous8.service'],check=True) +os.sync() +run=OLD/'results/muon_clip/seed_1337' +if run.resolve() != run or not (run/'manifest.json').is_file(): raise SystemExit('Cannot identify old run outputs; no cleanup') +archive=ROOT/'old_scientific_results.tgz' +if archive.exists(): raise SystemExit('Archive already exists; refusing repeated cleanup') +files=[p for p in OLD.rglob('*') if p.is_file() and not p.is_symlink() + and 'repo' not in p.relative_to(OLD).parts + and p.suffix in ('.csv','.json','.yaml','.log','.txt')] +with tarfile.open(archive,'w:gz') as tar: + for p in files: tar.add(p,arcname=str(p.relative_to(OLD)),recursive=False) +# Verify every archived file before deleting any model outputs. +with tarfile.open(archive) as tar: + for p in files: + member=tar.extractfile(str(p.relative_to(OLD))) + if member is None or hashlib.sha256(member.read()).digest()!=hashlib.sha256(p.read_bytes()).digest(): + raise SystemExit('Archive verification failed; no deletion') +removed=[] +for p in run.glob('*.pt'): + if p.is_symlink(): raise SystemExit('Unexpected checkpoint symlink') + if p.name.startswith(('checkpoint_','model_epoch_')): + removed.append({'path':str(p),'bytes':p.stat().st_size}); p.unlink() +(ROOT/'cleanup.json').write_text(json.dumps({'archive':str(archive),'removed':removed,'preserved_data':str(DATA)},indent=2)) +os.sync(); subprocess.run(['df','-h',str(MOUNT)],check=True) +print('Old training stopped; corpus protected; compact archive verified; old model checkpoint files removed.',flush=True) diff --git a/baseline/nanogpt_one_head/gpt2small/validate.py b/baseline/nanogpt_one_head/gpt2small/validate.py new file mode 100644 index 00000000..257e2971 --- /dev/null +++ b/baseline/nanogpt_one_head/gpt2small/validate.py @@ -0,0 +1,77 @@ +"""Bounded validation only. Called on the existing TPU after clean shutdown.""" +import argparse +import hashlib +import json +import os +from pathlib import Path +import subprocess +import sys +import time +import yaml + +p=argparse.ArgumentParser() +p.add_argument('--root',required=True); p.add_argument('--data',required=True) +p.add_argument('--deadline',type=float,required=True) +a=p.parse_args(); root=Path(a.root); base=Path(__file__).resolve().parents[1] +root.mkdir(parents=True,exist_ok=True) +for folder in ('logs','metrics','ww_metrics','checkpoints','summaries','configs'): (root/folder).mkdir(exist_ok=True) +configs={} +for opt in ('adamw','muonclip'): + cfg=yaml.safe_load((base/'configs'/f'gpt2_small_fineweb_{opt}_baseline.yaml').read_text()) + cfg['run_id']=root.name+'_'+opt + # Only control stop points here. Keep the planned LR/token horizon immutable across resume. + cfg['metrics_interval']=4 + cfg['benchmark_sync_every_step']=True + if opt=='muonclip': cfg['ww']['steps']=[4,25] + path=root/'configs'/f'{opt}.yaml'; path.write_text(yaml.safe_dump(cfg,sort_keys=False)); configs[opt]=path +report={'status':'running','deadline_unix':a.deadline,'long_run_started':False,'phases':[]} + +def persist(): + (root/'summaries/validation_report.json').write_text(json.dumps(report,indent=2)) + +def history(run): + return {str(x.relative_to(run)):hashlib.sha256(x.read_bytes()).hexdigest() + for folder in ('metrics','ww_metrics') for x in (run/folder).glob('*.json')} + +try: + for stop in (4,25): + for opt in ('adamw','muonclip'): + # Reserve five minutes for final checkpoint and backup; never start after the cutoff. + if time.time()>a.deadline-300: raise RuntimeError('Insufficient allocation time for next phase; no new allocation requested') + run=root/opt; old=history(run) + cmd=[sys.executable,'-u','-m','rg_nanogpt_one_head.gpt2_experiment', + '--config',str(configs[opt]),'--data-root',a.data,'--output',str(run), + '--device','tpu','--stop-after',str(stop),'--deadline-unix',str(a.deadline-300)] + if stop>4: cmd.append('--resume') + print('START',opt,'through step',stop,flush=True) + with (root/'logs'/f'{opt}_{stop}.log').open('x') as log: + subprocess.run(cmd,stdout=log,stderr=subprocess.STDOUT,check=True) + state=json.loads((run/'status.json').read_text()) + if state['step']!=stop: raise RuntimeError(f'{opt} stopped before requested step {stop}') + now=history(run) + if any(now.get(k)!=v for k,v in old.items()): raise RuntimeError('Resume changed previous scientific history') + rows=[json.loads(x.read_text()) for x in sorted((run/'metrics').glob('*.json'))] + initial,latest=rows[0],rows[-1] + if latest['train_nll']>=initial['train_nll']: raise RuntimeError(f'{opt} training probe NLL did not decrease; inspect before proceeding') + if latest['val_nll']>initial['val_nll']+1: raise RuntimeError(f'{opt} validation NLL increased substantially') + ww=[] + for x in (run/'ww_metrics').glob('*.json'): + measured=json.loads(x.read_text()); rec=measured['records'] + if len(rec)!=72 or len({r['matrix_name'] for r in rec})!=72: raise RuntimeError('WW matrix inventory failure') + if any(sum(r['matrix_type']==kind for r in rec)!=12 for kind in ('Q','K','V','O','MLP_IN','MLP_OUT')): raise RuntimeError('WW type counts incorrect') + ww.append({'file':x.name,'seconds':measured['seconds'], + 'raw_success':sum(r['raw_fit_status']=='success' for r in rec), + 'clipped_success':sum(r['clipped_fit_status']=='success' for r in rec), + 'null_success':sum(r['null_status']=='success' for r in rec)}) + if opt=='muonclip' and (not ww or any(x['null_success']!=72 for x in ww)): + raise RuntimeError('Randomized WW control incomplete; see retained records') + report['phases'].append({'optimizer':opt,'step':stop,'metrics':latest,'ww':ww, + 'resume_history_unchanged':bool(old)}) + persist(); print(json.dumps(report['phases'][-1]),flush=True) + subprocess.run([sys.executable,str(base/'gpt2small/analyze.py'),str(run)],check=True) + report['status']='short_validation_completed' + report['interpretation']='25 updates demonstrate functionality and initial direction only; not reproduction of a NanoGPT speedrun benchmark or proof of long-run stability.' +except Exception as exc: + report['status']='failed_or_incomplete'; report['error']=str(exc); persist(); raise +finally: + persist() diff --git a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/gpt2_experiment.py b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/gpt2_experiment.py new file mode 100644 index 00000000..7ec36c4d --- /dev/null +++ b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/gpt2_experiment.py @@ -0,0 +1,289 @@ +"""Isolated, token-budgeted GPT-2 validation; never downloads data or allocates TPUs.""" +from __future__ import annotations +import argparse +import copy +import fcntl +import hashlib +import json +import math +import os +from pathlib import Path +import random +import time + +import numpy as np +import torch +import yaml +from .model import GPT, GPTConfig, transformer_matrix_items +from . import runtime as rt, tpu_spmd as spmd +from .data import load_memmaps +from .muonclip import install_muonclip_extension +from . import optimizers +from .spectral import WeightMatrixHolder, _attach_matrix_metadata + + +def atomic_json(path, value): + path = Path(path); path.parent.mkdir(parents=True, exist_ok=True) + tmp = path.with_suffix(path.suffix + '.tmp') + with tmp.open('w') as f: + json.dump(value, f, indent=2, default=str, allow_nan=True) + f.flush(); os.fsync(f.fileno()) + os.replace(tmp, path) + + +def append_record(path, row): + """One immutable JSON per step: no truncation, silent overwrite, or duplicate.""" + path = Path(path); path.parent.mkdir(parents=True, exist_ok=True) + if path.exists(): + raise FileExistsError(f'Already recorded: {path}') + atomic_json(path, row) + + +def due(step, spec): + return (step in spec.get('steps', []) or + bool(spec.get('interval', 0) and step % spec['interval'] == 0) or + bool(spec.get('logarithmic', False) and step > 0 and + step in {int(k * 10**p) for p in range(12) for k in (1, 2, 5)})) + + +def batch(array, generator, size, context, device): + if len(array) <= context: + raise ValueError('Split is shorter than one context') + starts = torch.randint(len(array) - context, (size,), generator=generator).tolist() + x = torch.from_numpy(np.stack([np.array(array[i:i+context], dtype=np.int64) for i in starts])) + y = torch.from_numpy(np.stack([np.array(array[i+1:i+context+1], dtype=np.int64) for i in starts])) + return spmd.batch_to_device(x, device), spmd.batch_to_device(y, device) + + +@torch.no_grad() +def evaluate(model, arrays, cfg, device): + model.eval(); result = {} + for j, split in enumerate(('train', 'val', 'test')): + gen = torch.Generator().manual_seed(cfg['seed'] + 20000 + j) + losses, correct = [], [] + for _ in range(cfg['eval_batches']): + x, y = batch(arrays[split], gen, cfg['training']['batch_size'], cfg['model']['block_size'], device) + logits, loss = model(x, y) + losses.append(loss.detach()); correct.append((logits.argmax(-1) == y).float().mean()) + rt.mark_step(device) + nll, accuracy = torch.stack(losses).mean().item(), torch.stack(correct).mean().item() + if not math.isfinite(nll): + raise RuntimeError(f'Nonfinite {split} NLL') + result.update({f'{split}_nll': nll, f'{split}_perplexity': math.exp(min(nll, 700)), + f'{split}_accuracy': accuracy, f'{split}_token_error': 1 - accuracy}) + model.train(); return result + + +def make_handles(model, cfg): + install_muonclip_extension() + profile = copy.deepcopy(cfg['optimizer']) + handles = optimizers.make_optimizer_handles(model, profile) + if profile['family'] == 'muon_clip': + aux = next(h for h in handles if h.role == 'auxiliary') + aux.peak_lr = profile['aux_learning_rate']; aux.min_lr = profile['aux_min_learning_rate'] + for group in aux.optimizer.param_groups: + group['lr'] = aux.peak_lr + return handles + + +def measure_ww(model, cfg, identity, metrics): + """Existing clip_xmax + randomized null; never substitute clipped for raw.""" + import weightwatcher as ww + started = time.monotonic() + device = rt.model_device(model) + states = (random.getstate(), np.random.get_state(), torch.get_rng_state(), rt.capture_accelerator_rng_state(device)) + try: + seed = cfg['seed'] + 1000003 + identity['step'] + random.seed(seed); np.random.seed(seed % (2**32-1)); torch.manual_seed(seed) + holder = WeightMatrixHolder(model) + frame = ww.WeightWatcher(model=holder).analyze( + ERG=True, randomize=True, plot=False, fix_fingers='clip_xmax', + max_fingers=10, min_evals=cfg['ww'].get('min_evals', 20)) + frame = _attach_matrix_metadata(frame, holder.matrix_metadata) + expected = {m['matrix_name'] for m in holder.matrix_metadata} + if len(frame) != len(expected) or set(frame.matrix_name) != expected: + raise RuntimeError('Incomplete/duplicated WeightWatcher matrix inventory') + records = [] + for raw in frame.to_dict('records'): + def number(key): + try: return float(raw.get(key, float('nan'))) + except (ValueError, TypeError): return float('nan') + a, c, null = number('raw_alpha'), number('alpha'), number('rand_distance') + fit_ok = str(raw.get('status', 'success')) == 'success' + valid = math.isfinite(a) and a > 0 and fit_ok + clipped_valid = math.isfinite(c) and c > 0 and fit_ok + records.append({**raw, **identity, **metrics, + 'layer': f"L{int(raw['block']):02d}", + 'matrix_type': str(raw['matrix_type']).removeprefix('W_'), + 'alpha_raw': a if valid else float('nan'), + 'alpha_clip_xmax': c if clipped_valid else float('nan'), + 'raw_fit_status': 'success' if valid else 'failed', + 'raw_fit_reason': '' if valid else str(raw.get('warning', 'raw_alpha missing/nonfinite/nonpositive')), + 'clipped_fit_status': 'success' if clipped_valid else 'failed', + 'null_status': 'success' if math.isfinite(null) else 'failed', + 'randomized_distance': null}) + return {'records': records, 'seconds': time.monotonic() - started} + finally: + random.setstate(states[0]); np.random.set_state(states[1]); torch.set_rng_state(states[2]) + rt.restore_accelerator_rng_state(states[3], device) + + +def save_checkpoint(root, payload, keep=3, milestone=False): + root = Path(root); root.mkdir(parents=True, exist_ok=True) + path = root / f"step_{payload['step']:09d}.pt" + tmp = path.with_suffix('.tmp') + with tmp.open('wb') as f: + torch.save(rt.tree_to_cpu(payload), f); f.flush(); os.fsync(f.fileno()) + os.replace(tmp, path) + atomic_json(root / 'latest.json', {'file': path.name, 'step': payload['step']}) + if milestone: + import shutil + (root / 'milestones').mkdir(exist_ok=True) + shutil.copy2(path, root / 'milestones' / path.name) + for old in sorted(root.glob('step_*.pt'))[:-max(2, keep)]: + old.unlink() + return path + + +def architecture(model): + c = model.cfg + return {'parameters': model.parameter_count(), 'layers': c.n_layer, 'heads': c.n_head, + 'd_model': c.n_embd, 'head_dim': c.n_embd // c.n_head, 'context': c.block_size, + 'vocab_size': c.vocab_size, 'tied': model.lm_head.weight is model.token_embedding.weight, + 'logical_matrices': len(transformer_matrix_items(model)), **spmd.metadata()} + + +def train(cfg, data_root, output, *, device='cpu', resume=False, stop_after=None, deadline=None): + output = Path(output); output.mkdir(parents=True, exist_ok=True) + with (output / 'writer.lock').open('a') as lock: + fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB) + return _train(cfg, data_root, output, device, resume, stop_after, deadline) + + +def _train(cfg, data_root, output, device, resume, stop_after, deadline): + cfg = copy.deepcopy(cfg) + t = cfg['training']; context = cfg['model']['block_size'] + step_tokens = t['batch_size'] * t['grad_accum_steps'] * context + total = min(t.get('max_steps', 10**12), math.ceil(t['max_tokens'] / step_tokens)) + if not 0 <= t['warmup_steps'] < t['schedule_steps'] or total < 1: + raise ValueError('Invalid training horizon/warmup') + spmd.initialize(cfg, device); dev = rt.choose_device(device); rt.configure_runtime(dev, cfg) + rt.seed_everything(cfg['seed'], dev) + metadata, arrays = load_memmaps(data_root, cfg) # hashes and document isolation; NO download + versions = {'torch': torch.__version__, 'numpy': np.__version__} + fingerprint = hashlib.sha256(json.dumps({'config': cfg, 'data': metadata, 'versions': versions}, sort_keys=True).encode()).hexdigest() + model = GPT(GPTConfig(**cfg['model'])).to(dev); spmd.replicate_model(model) + handles = make_handles(model, cfg) + gen = torch.Generator().manual_seed(cfg['seed'] + 11) + step = 0; elapsed = 0.; last_grad = None; compile_seconds = 0.; steady_seconds = 0.; steady_steps = 0 + latest = output / 'checkpoints/latest.json' + if resume: + pointer = json.loads(latest.read_text()) + state = torch.load(latest.parent / pointer['file'], map_location='cpu', weights_only=False) + if state['fingerprint'] != fingerprint or state['run_id'] != cfg['run_id']: + raise RuntimeError('Resume config/data/run ID mismatch') + model.load_state_dict(state['model']); spmd.replicate_model(model) + optimizers.load_optimizer_state_dict(handles, state['optimizers']) + gen.set_state(state['data_rng']); torch.set_rng_state(state['torch_rng']) + random.setstate(state['python_rng']); np.random.set_state(state['numpy_rng']) + rt.restore_accelerator_rng_state(state['accelerator_rng'], dev) + step = state['step']; elapsed = state['wall_time']; last_grad = state['gradient_norm'] + if state['tokens_seen'] != step * step_tokens or state['scheduler_step'] != step: + raise RuntimeError('Checkpoint step/token/scheduler mismatch') + for folder in ('metrics', 'ww_metrics'): + if any(int(p.stem) > step for p in (output / folder).glob('*.json')): + raise RuntimeError('History ahead of checkpoint; refusing to erase scientific records') + # Complete the checkpoint-first measurement transaction after interrupted writes. + for folder, row in state['pending_records'].items(): + path = output / folder / f'{step:09d}.json' + if not path.exists(): append_record(path, row) + atomic_json(output / 'resume_verified.json', {'step': step, 'tokens_seen': step * step_tokens, + 'optimizer_states': [len(h.optimizer.state) for h in handles], 'fingerprint': fingerprint}) + elif latest.exists() or (output / 'manifest.json').exists(): + raise FileExistsError('Use a fresh output directory or explicit --resume') + atomic_json(output / 'manifest.json', {'config': cfg, 'data': metadata, 'data_path': str(Path(data_root).resolve()), + 'fingerprint': fingerprint, 'architecture': architecture(model), 'runtime': rt.runtime_metadata(dev), + 'effective_batch_tokens': step_tokens}) + print(json.dumps({'architecture': architecture(model), 'batch_tokens': step_tokens}), flush=True) + started = time.monotonic(); initial_step = step + + def record(step, final=False): + nonlocal elapsed + metrics = evaluate(model, arrays, cfg, dev) + wall = elapsed + time.monotonic() - started + identity = {'run_id': cfg['run_id'], 'optimizer': cfg['optimizer']['family'], 'seed': cfg['seed'], + 'step': step, 'tokens_seen': step * step_tokens, 'wall_time': wall, + 'learning_rate': handles[0].optimizer.param_groups[0]['lr'] if step else 0., 'gradient_norm': last_grad} + row = {**identity, **metrics, 'first_two_updates_seconds_including_compile': compile_seconds, + 'steady_training_tokens_per_second': steady_steps * step_tokens / max(steady_seconds, 1e-9), + 'end_to_end_tokens_per_second': (step-initial_step) * step_tokens / max(time.monotonic()-started, 1e-9)} + pending = {'metrics': row} + if cfg['ww']['enabled'] and due(step, cfg['ww']): + measured = measure_ww(model, cfg, identity, metrics) + measured['recommended_interval_seconds_for_10pct'] = 9 * measured['seconds'] + pending['ww_metrics'] = measured + # Checkpoint FIRST includes the pending scalar/WW transaction. On resume, finish missing rows. + save_checkpoint(output / 'checkpoints', { + 'run_id': cfg['run_id'], 'config': cfg, 'fingerprint': fingerprint, 'model': model.state_dict(), + 'optimizers': optimizers.optimizer_state_dict(handles), 'step': step, 'tokens_seen': step*step_tokens, + 'scheduler_step': step, 'data_rng': gen.get_state(), 'torch_rng': torch.get_rng_state(), + 'python_rng': random.getstate(), 'numpy_rng': np.random.get_state(), + 'accelerator_rng': rt.capture_accelerator_rng_state(dev), 'wall_time': elapsed + time.monotonic()-started, + 'gradient_norm': last_grad, 'pending_records': pending}, + milestone=step in cfg.get('milestones', [])) + for folder, value in pending.items(): append_record(output / folder / f'{step:09d}.json', value) + print(json.dumps(row), flush=True) + + if not resume: record(0) + training_window = time.monotonic(); last_timed_step = step + while step < total: + if (stop_after is not None and step >= stop_after) or (deadline and time.time() >= deadline) or (output / 'STOP').exists(): + break + optimizers.zero_grad(handles) + for handle in handles: + lr = optimizers.cosine_learning_rate(step, total_steps=t['schedule_steps'], warmup_steps=t['warmup_steps'], + peak_lr=handle.peak_lr, min_lr=handle.min_lr) + handle.set_lr(lr) + for _ in range(t['grad_accum_steps']): + x, y = batch(arrays['train'], gen, t['batch_size'], context, dev) + _, loss = model(x, y); (loss / t['grad_accum_steps']).backward() + spmd.replicate_gradients(model) + norm = torch.nn.utils.clip_grad_norm_(model.parameters(), t['grad_clip']) + optimizers.optimizer_step(handles); rt.mark_step(dev) + step += 1 + measurement_due = (step % cfg['metrics_interval'] == 0 or due(step, cfg['ww']) + or step == total or step == stop_after) + if step <= initial_step + 2 or measurement_due or cfg.get('benchmark_sync_every_step', False): + rt.synchronize(dev) + seconds = time.monotonic() - training_window + if step <= initial_step + 2: compile_seconds += seconds + else: steady_seconds += seconds; steady_steps += step - last_timed_step + last_timed_step = step + if step == initial_step + 2 and dev.type == 'xla': + from torch_xla.debug import metrics as xla_metrics + log = output / 'logs' / f'xla_compile_metrics_after_step_{step}.txt' + log.parent.mkdir(exist_ok=True) + log.write_text(xla_metrics.metrics_report()) + if measurement_due: + last_grad = norm.item() + if not math.isfinite(last_grad): raise RuntimeError('Nonfinite gradient norm') + record(step) + training_window = time.monotonic() + if not (output / 'metrics' / f'{step:09d}.json').exists(): + last_grad = norm.item(); record(step, final=True) + atomic_json(output / 'status.json', {'step': step, 'tokens_seen': step*step_tokens, + 'completed': step >= total, 'stopped': step < total, 'long_run_launched': cfg.get('long_run', False)}) + return output + + +def main(): + p = argparse.ArgumentParser(description=__doc__) + p.add_argument('--config', required=True); p.add_argument('--data-root', required=True) + p.add_argument('--output', required=True); p.add_argument('--device', default='tpu') + p.add_argument('--resume', action='store_true'); p.add_argument('--stop-after', type=int) + p.add_argument('--deadline-unix', type=float); p.add_argument('--allow-long-run', action='store_true') + a = p.parse_args(); cfg = yaml.safe_load(Path(a.config).read_text()) + if cfg.get('long_run') and not a.allow_long_run: p.error('Long experiment requires explicit --allow-long-run') + train(cfg, a.data_root, a.output, device=a.device, resume=a.resume, stop_after=a.stop_after, deadline=a.deadline_unix) + +if __name__ == '__main__': main() diff --git a/baseline/nanogpt_one_head/tests/test_gpt2_experiment.py b/baseline/nanogpt_one_head/tests/test_gpt2_experiment.py new file mode 100644 index 00000000..77debd14 --- /dev/null +++ b/baseline/nanogpt_one_head/tests/test_gpt2_experiment.py @@ -0,0 +1,100 @@ +import copy +import json +from pathlib import Path +import sys +import types +import numpy as np +import pandas as pd +import pytest +import torch +import yaml +from rg_nanogpt_one_head.model import GPT, GPTConfig, transformer_matrix_items +from rg_nanogpt_one_head.data import write_token_splits +from rg_nanogpt_one_head import gpt2_experiment as g +BASE=Path(__file__).resolve().parents[1] + +def config(name='cpu_smoke'): + return yaml.safe_load((BASE/'configs'/f'gpt2_small_{name}.yaml').read_text()) + +def data(tmp,c): + class Encoder: + n_vocab=64; eot_token=63 + def encode_ordinary(self,text): return list(range(16)) + root=tmp/'data' + write_token_splits(['document']*100,Encoder(),root,train_tokens=512,val_tokens=128,test_tokens=128, + dataset_metadata={'dataset_name':c['dataset']['name'],'dataset_config':c['dataset']['config'], + 'dataset_split':'train','dataset_revision':c['dataset']['revision'],'tokenizer':'gpt2'}) + return root + +def test_full_architecture(): + c=config('fineweb_adamw_baseline') + with torch.device('meta'): model=GPT(GPTConfig(**c['model'])) + a=g.architecture(model) + assert (a['layers'],a['heads'],a['d_model'],a['context'],a['logical_matrices'])==(12,12,768,1024,72) + assert a['parameters']==124439808 and a['tied'] + items=transformer_matrix_items(model) + assert len({id(w) for _,_,_,w in items})==72 + assert all(tuple(w.shape)==(768,768) for _,kind,_,w in items if kind in ('W_Q','W_K','W_V','W_O')) + +@pytest.mark.parametrize('family',['adamw','muon_clip']) +def test_exact_resume_and_append(tmp_path,family): + torch.set_num_threads(1) + c=config(); c['optimizer']=config('fineweb_'+('adamw' if family=='adamw' else 'muonclip')+'_baseline')['optimizer'] + d=data(tmp_path,c) + g.train(c,d,tmp_path/'full') + g.train(c,d,tmp_path/'split',stop_after=2) + before=(tmp_path/'split/metrics/000000002.json').read_bytes() + g.train(c,d,tmp_path/'split',resume=True) + def saved(root): + return torch.load(root/'checkpoints/step_000000004.pt',weights_only=False) + a,b=saved(tmp_path/'full'),saved(tmp_path/'split') + for name in a['model']: torch.testing.assert_close(a['model'][name],b['model'][name],rtol=0,atol=0) + assert torch.equal(a['data_rng'],b['data_rng']) and a['tokens_seen']==b['tokens_seen']==64 + from rg_nanogpt_one_head.checkpoints import optimizer_state_sha256 + assert optimizer_state_sha256(a['optimizers'])==optimizer_state_sha256(b['optimizers']) + assert (tmp_path/'split/metrics/000000002.json').read_bytes()==before + assert len(list((tmp_path/'split/metrics').glob('*.json')))==3 + with pytest.raises(FileExistsError): g.append_record(tmp_path/'split/metrics/000000002.json',{}) + row=json.loads((tmp_path/'split/metrics/000000004.json').read_text()) + assert row['test_accuracy']+row['test_token_error']==1 + assert row['test_perplexity']==pytest.approx(np.exp(row['test_nll'])) + assert row['train_nll'] Date: Sat, 3 Oct 2026 17:47:11 -0700 Subject: [PATCH 11/36] Fix GPT-2 TPU identity and log permissions; retry without cleanup --- baseline/nanogpt_one_head/gpt2small/on_tpu.sh | 3 +-- .../gpt2small/retry_existing.sh | 25 +++++++++++++++++++ .../gpt2small/tpu_environment.sh | 14 +++++++++++ .../nanogpt_one_head/gpt2small/validate.py | 14 ++++++++++- 4 files changed, 53 insertions(+), 3 deletions(-) create mode 100644 baseline/nanogpt_one_head/gpt2small/retry_existing.sh create mode 100644 baseline/nanogpt_one_head/gpt2small/tpu_environment.sh diff --git a/baseline/nanogpt_one_head/gpt2small/on_tpu.sh b/baseline/nanogpt_one_head/gpt2small/on_tpu.sh index f98f2040..1f288c4f 100644 --- a/baseline/nanogpt_one_head/gpt2small/on_tpu.sh +++ b/baseline/nanogpt_one_head/gpt2small/on_tpu.sh @@ -15,8 +15,7 @@ git -C "$repo" checkout --detach "$commit" test "$(git -C "$repo" rev-parse HEAD)" = "$commit" git -C "$repo" rev-parse HEAD export PYTHONPATH="$repo/baseline/nanogpt_one_head/src" -export OMP_NUM_THREADS=4 OPENBLAS_NUM_THREADS=4 MKL_NUM_THREADS=4 -export PJRT_DEVICE=TPU +source "$repo/baseline/nanogpt_one_head/gpt2small/tpu_environment.sh" cd "$repo/baseline/nanogpt_one_head" "$python" -c 'import torch, torch_xla, weightwatcher, yaml; print("Installed dependencies loaded")' deadline=$("$python" -c 'import json,sys; print(json.load(open(sys.argv[1]))["deadline_unix"])' "$root/old_allocation.json") diff --git a/baseline/nanogpt_one_head/gpt2small/retry_existing.sh b/baseline/nanogpt_one_head/gpt2small/retry_existing.sh new file mode 100644 index 00000000..6646a82b --- /dev/null +++ b/baseline/nanogpt_one_head/gpt2small/retry_existing.sh @@ -0,0 +1,25 @@ +#!/usr/bin/env bash +# Retry validation ONLY. Never invokes old-run cleanup or downloads data. +set -euo pipefail +old=$1 +base=$(cd "$(dirname "$0")/.." && pwd) +source "$base/gpt2small/tpu_environment.sh" +mountpoint -q /mnt/disks/rg-data +test -f "$old/old_allocation.json" +test -f /mnt/disks/rg-data/continuous8/data/meta.json +root="${old}-retry-$(date -u +%Y%m%d-%H%M%S)" +mkdir "$root" +cp "$old/old_allocation.json" "$root/old_allocation.json" +python=/mnt/disks/rg-data/continuous8/venv/bin/python +export PYTHONPATH="$base/src" +git -C "$base" rev-parse HEAD | tee "$root/commit.txt" +deadline=$("$python" -c 'import json,sys; print(json.load(open(sys.argv[1]))["deadline_unix"])' "$root/old_allocation.json") +backup() { + rc=$? + trap - EXIT + sync + gcloud storage rsync "$root" "gs://tpu-builders-504820-ww-continuous8/gpt2small/$(basename "$root")" --recursive --project=tpu-builders-504820 || exit 1 + exit "$rc" +} +trap backup EXIT +"$python" -u "$base/gpt2small/validate.py" --root "$root" --data /mnt/disks/rg-data/continuous8/data --deadline "$deadline" 2>&1 | tee "$root/validation.log" diff --git a/baseline/nanogpt_one_head/gpt2small/tpu_environment.sh b/baseline/nanogpt_one_head/gpt2small/tpu_environment.sh new file mode 100644 index 00000000..f7cd77f8 --- /dev/null +++ b/baseline/nanogpt_one_head/gpt2small/tpu_environment.sh @@ -0,0 +1,14 @@ +# Source before importing torch_xla. Matches the verified v5litepod-8 allocation. +export PJRT_DEVICE=TPU TPU_ACCELERATOR_TYPE=v5litepod-8 +export OMP_NUM_THREADS=4 OPENBLAS_NUM_THREADS=4 MKL_NUM_THREADS=4 +if systemctl is-active --quiet rg-continuous8.service; then + echo 'Old training service is still active; refusing concurrent TPU use.' >&2 + exit 1 +fi +# The previous service ran as root; libtpu writes here even for a user-run job. +if [ -L /tmp/tpu_logs ]; then + echo 'Unexpected TPU log directory symlink.' >&2 + exit 1 +fi +sudo install -d -m 0755 -o "$(id -u)" -g "$(id -g)" /tmp/tpu_logs +test -w /tmp/tpu_logs diff --git a/baseline/nanogpt_one_head/gpt2small/validate.py b/baseline/nanogpt_one_head/gpt2small/validate.py index 257e2971..dcf3c952 100644 --- a/baseline/nanogpt_one_head/gpt2small/validate.py +++ b/baseline/nanogpt_one_head/gpt2small/validate.py @@ -7,6 +7,7 @@ import subprocess import sys import time +import threading import yaml p=argparse.ArgumentParser() @@ -45,7 +46,18 @@ def history(run): if stop>4: cmd.append('--resume') print('START',opt,'through step',stop,flush=True) with (root/'logs'/f'{opt}_{stop}.log').open('x') as log: - subprocess.run(cmd,stdout=log,stderr=subprocess.STDOUT,check=True) + child=subprocess.Popen(cmd,stdout=subprocess.PIPE,stderr=subprocess.STDOUT,text=True,bufsize=1) + def stream(): + for line in child.stdout: + log.write(line); log.flush(); print(line,end='',flush=True) + reader=threading.Thread(target=stream,daemon=True); reader.start() + while True: + try: + rc=child.wait(timeout=30); break + except subprocess.TimeoutExpired: + print(f'WAIT {opt} through step {stop}: process {child.pid} still running',flush=True) + reader.join() + if rc: raise subprocess.CalledProcessError(rc,cmd) state=json.loads((run/'status.json').read_text()) if state['step']!=stop: raise RuntimeError(f'{opt} stopped before requested step {stop}') now=history(run) From 6491a90f64a1f2574608048d107978cb312a29c9 Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Sat, 3 Oct 2026 17:48:49 -0700 Subject: [PATCH 12/36] Create GPT-2 retry output with invoking user ownership --- baseline/nanogpt_one_head/gpt2small/retry_existing.sh | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/baseline/nanogpt_one_head/gpt2small/retry_existing.sh b/baseline/nanogpt_one_head/gpt2small/retry_existing.sh index 6646a82b..2c3eaf98 100644 --- a/baseline/nanogpt_one_head/gpt2small/retry_existing.sh +++ b/baseline/nanogpt_one_head/gpt2small/retry_existing.sh @@ -8,7 +8,8 @@ mountpoint -q /mnt/disks/rg-data test -f "$old/old_allocation.json" test -f /mnt/disks/rg-data/continuous8/data/meta.json root="${old}-retry-$(date -u +%Y%m%d-%H%M%S)" -mkdir "$root" +sudo mkdir "$root" +sudo chown "$(id -u):$(id -g)" "$root" cp "$old/old_allocation.json" "$root/old_allocation.json" python=/mnt/disks/rg-data/continuous8/venv/bin/python export PYTHONPATH="$base/src" From d8cb1efafd8c2f1cce0f2b529afe7479d2817367 Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Sat, 3 Oct 2026 18:06:58 -0700 Subject: [PATCH 13/36] Diagnose invalid gradients before TPU updates and back up using object permissions --- baseline/nanogpt_one_head/gpt2small/backup.py | 19 ++++++++++++ .../gpt2small/diagnose_adamw.sh | 30 +++++++++++++++++++ baseline/nanogpt_one_head/gpt2small/on_tpu.sh | 2 +- .../gpt2small/retry_existing.sh | 2 +- .../nanogpt_one_head/gpt2small/validate.py | 1 + .../rg_nanogpt_one_head/gpt2_experiment.py | 30 +++++++++++++++++-- .../tests/test_gpt2_experiment.py | 12 ++++++++ 7 files changed, 92 insertions(+), 4 deletions(-) create mode 100644 baseline/nanogpt_one_head/gpt2small/backup.py create mode 100644 baseline/nanogpt_one_head/gpt2small/diagnose_adamw.sh diff --git a/baseline/nanogpt_one_head/gpt2small/backup.py b/baseline/nanogpt_one_head/gpt2small/backup.py new file mode 100644 index 00000000..8c6851b9 --- /dev/null +++ b/baseline/nanogpt_one_head/gpt2small/backup.py @@ -0,0 +1,19 @@ +"""Upload validation artifacts using object permissions; no bucket-metadata read.""" +import argparse +import json +from pathlib import Path +from rg_nanogpt_one_head.continuous_support import CloudPublisher +p=argparse.ArgumentParser(); p.add_argument('root'); a=p.parse_args() +root=Path(a.root).resolve() +if not root.is_dir(): raise SystemExit('Output directory missing') +publisher=CloudPublisher('gs://tpu-builders-504820-ww-continuous8/gpt2small/'+root.name) +receipts=[] +for path in sorted(root.rglob('*')): + relative=path.relative_to(root) + if path.is_symlink() or not path.is_file() or 'repo' in relative.parts or path.suffix=='.tmp': continue + print('Uploading',relative,flush=True) + receipts.append(publisher.file(path,relative.as_posix())) +receipt=root/'CLOUD_BACKUP_VERIFIED.json' +receipt.write_text(json.dumps({'files':receipts,'method':'object upload with CRC32C and object-size verification'},indent=2)) +publisher.file(receipt,receipt.name) +print('Cloud backup verified:',root.name,flush=True) diff --git a/baseline/nanogpt_one_head/gpt2small/diagnose_adamw.sh b/baseline/nanogpt_one_head/gpt2small/diagnose_adamw.sh new file mode 100644 index 00000000..4ae91cbd --- /dev/null +++ b/baseline/nanogpt_one_head/gpt2small/diagnose_adamw.sh @@ -0,0 +1,30 @@ +#!/usr/bin/env bash +# Four guarded updates only; caller supplies verified allocation expiry minus backup reserve. +set -euo pipefail +old=$1 +deadline=$2 +base=$(cd "$(dirname "$0")/.." && pwd) +source "$base/gpt2small/tpu_environment.sh" +python=/mnt/disks/rg-data/continuous8/venv/bin/python +export PYTHONPATH="$base/src" +"$python" -c 'import sys,time; assert float(sys.argv[1])>time.time()+60,"Too little allocation time; diagnostic not started"' "$deadline" +root="$old/diagnostic-$(date -u +%Y%m%d-%H%M%S)" +mkdir "$root" +"$python" - "$old" "$root" <<'PY' +import sys,yaml +from pathlib import Path +old,root=map(Path,sys.argv[1:]) +cfg=yaml.safe_load((old/'configs/adamw.yaml').read_text()) +cfg.update(run_id=root.name,validation_gradient_checks=True,metrics_interval=1,benchmark_sync_every_step=True) +cfg['ww']['enabled']=False +(root/'config.yaml').write_text(yaml.safe_dump(cfg,sort_keys=False)) +PY +backup() { + rc=$? + trap - EXIT + sync + "$python" "$base/gpt2small/backup.py" "$root" || exit 1 + exit "$rc" +} +trap backup EXIT +"$python" -u -m rg_nanogpt_one_head.gpt2_experiment --config "$root/config.yaml" --data-root /mnt/disks/rg-data/continuous8/data --output "$root/adamw" --device tpu --stop-after 4 --deadline-unix "$deadline" 2>&1 | tee "$root/diagnostic.log" diff --git a/baseline/nanogpt_one_head/gpt2small/on_tpu.sh b/baseline/nanogpt_one_head/gpt2small/on_tpu.sh index 1f288c4f..e851c1c6 100644 --- a/baseline/nanogpt_one_head/gpt2small/on_tpu.sh +++ b/baseline/nanogpt_one_head/gpt2small/on_tpu.sh @@ -25,7 +25,7 @@ backup() { trap - EXIT sync # Keep backups separate from prior pilot archives. No remote deletion. - if ! gcloud storage rsync "$root" "$bucket/$(basename "$root")" --recursive --exclude='repo/.*' --project=tpu-builders-504820; then + if ! "$python" "$repo/baseline/nanogpt_one_head/gpt2small/backup.py" "$root"; then echo "Cloud backup FAILED; all outputs remain on persistent disk: $root" >&2 exit 1 fi diff --git a/baseline/nanogpt_one_head/gpt2small/retry_existing.sh b/baseline/nanogpt_one_head/gpt2small/retry_existing.sh index 2c3eaf98..71ed2993 100644 --- a/baseline/nanogpt_one_head/gpt2small/retry_existing.sh +++ b/baseline/nanogpt_one_head/gpt2small/retry_existing.sh @@ -19,7 +19,7 @@ backup() { rc=$? trap - EXIT sync - gcloud storage rsync "$root" "gs://tpu-builders-504820-ww-continuous8/gpt2small/$(basename "$root")" --recursive --project=tpu-builders-504820 || exit 1 + "$python" "$base/gpt2small/backup.py" "$root" || exit 1 exit "$rc" } trap backup EXIT diff --git a/baseline/nanogpt_one_head/gpt2small/validate.py b/baseline/nanogpt_one_head/gpt2small/validate.py index dcf3c952..22946dbb 100644 --- a/baseline/nanogpt_one_head/gpt2small/validate.py +++ b/baseline/nanogpt_one_head/gpt2small/validate.py @@ -23,6 +23,7 @@ # Only control stop points here. Keep the planned LR/token horizon immutable across resume. cfg['metrics_interval']=4 cfg['benchmark_sync_every_step']=True + cfg['validation_gradient_checks']=True if opt=='muonclip': cfg['ww']['steps']=[4,25] path=root/'configs'/f'{opt}.yaml'; path.write_text(yaml.safe_dump(cfg,sort_keys=False)); configs[opt]=path report={'status':'running','deadline_unix':a.deadline,'long_run_started':False,'phases':[]} diff --git a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/gpt2_experiment.py b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/gpt2_experiment.py index 7ec36c4d..dcdca0c8 100644 --- a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/gpt2_experiment.py +++ b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/gpt2_experiment.py @@ -128,6 +128,27 @@ def number(key): rt.restore_accelerator_rng_state(states[3], device) +def require_finite_update(model, norm, losses, output, step): + """Diagnostic pre-update barrier for validation. Never apply known-invalid gradients.""" + value=float(norm.detach().cpu()) + loss_values=[float(x.detach().cpu()) for x in losses] + if math.isfinite(value) and all(math.isfinite(x) for x in loss_values): + print(json.dumps({'before_update':step,'microbatch_losses':loss_values,'gradient_norm':value}),flush=True) + return + gradients=[] + for name,p in model.named_parameters(): + if p.grad is None: continue + grad=p.grad.detach().float().cpu() + finite=torch.isfinite(grad) + gradients.append({'parameter':name,'shape':list(grad.shape), + 'nonfinite_elements':int((~finite).sum()), + 'max_finite_abs':float(grad[finite].abs().max()) if finite.any() else None}) + report={'before_update':step,'gradient_norm':value,'microbatch_losses':loss_values,'gradients':gradients} + atomic_json(Path(output)/'nonfinite_diagnostics.json',report) + bad=[r['parameter'] for r in gradients if r['nonfinite_elements']] + raise RuntimeError(f'Invalid loss/gradient BEFORE update {step}; bad matrices={bad}; see nonfinite_diagnostics.json') + + def save_checkpoint(root, payload, keep=3, milestone=False): root = Path(root); root.mkdir(parents=True, exist_ok=True) path = root / f"step_{payload['step']:09d}.pt" @@ -244,11 +265,16 @@ def record(step, final=False): lr = optimizers.cosine_learning_rate(step, total_steps=t['schedule_steps'], warmup_steps=t['warmup_steps'], peak_lr=handle.peak_lr, min_lr=handle.min_lr) handle.set_lr(lr) + losses = [] for _ in range(t['grad_accum_steps']): x, y = batch(arrays['train'], gen, t['batch_size'], context, dev) - _, loss = model(x, y); (loss / t['grad_accum_steps']).backward() + _, loss = model(x, y); losses.append(loss.detach()) + (loss / t['grad_accum_steps']).backward() spmd.replicate_gradients(model) - norm = torch.nn.utils.clip_grad_norm_(model.parameters(), t['grad_clip']) + norm = rt.gradient_norm(model.parameters()) + if cfg.get('validation_gradient_checks', False): + require_finite_update(model, norm, losses, output, step + 1) + torch.nn.utils.clip_grad_norm_(model.parameters(), t['grad_clip'], foreach=False) optimizers.optimizer_step(handles); rt.mark_step(dev) step += 1 measurement_due = (step % cfg['metrics_interval'] == 0 or due(step, cfg['ww']) diff --git a/baseline/nanogpt_one_head/tests/test_gpt2_experiment.py b/baseline/nanogpt_one_head/tests/test_gpt2_experiment.py index 77debd14..28784bfd 100644 --- a/baseline/nanogpt_one_head/tests/test_gpt2_experiment.py +++ b/baseline/nanogpt_one_head/tests/test_gpt2_experiment.py @@ -98,3 +98,15 @@ def test_interrupted_record_transaction_recovers(tmp_path): g.train(c,d,out,resume=True,stop_after=2) assert missing.read_bytes()==expected assert len(list((out/'metrics').glob('*.json')))==2 + + +def test_invalid_gradient_diagnostic_precedes_update(tmp_path): + model=GPT(GPTConfig(**config()['model'])) + for p in model.parameters(): p.grad=torch.ones_like(p) + name,p=next(iter(model.named_parameters())); p.grad.view(-1)[0]=float('nan') + before=p.detach().clone() + with pytest.raises(RuntimeError,match='BEFORE update 1'): + g.require_finite_update(model,torch.tensor(float('nan')),[torch.tensor(11.)],tmp_path,1) + report=json.loads((tmp_path/'nonfinite_diagnostics.json').read_text()) + assert [r['parameter'] for r in report['gradients'] if r['nonfinite_elements']]==[name] + assert torch.equal(before,p.detach()) From 105c8f684eb6eefc9092b4bf060366209406a226 Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Sat, 3 Oct 2026 18:31:19 -0700 Subject: [PATCH 14/36] Replace expiring TPU for bounded GPT-2 validation while preserving FineWeb disk --- baseline/nanogpt_one_head/gpt2small/README.md | 33 ++++ .../gpt2small/reallocate_validation.py | 167 ++++++++++++++++++ .../gpt2small/replacement_startup.sh | 59 +++++++ .../gpt2small/replacement_worker.sh | 38 ++++ .../tests/test_gpt2_reallocation.py | 76 ++++++++ 5 files changed, 373 insertions(+) create mode 100644 baseline/nanogpt_one_head/gpt2small/reallocate_validation.py create mode 100644 baseline/nanogpt_one_head/gpt2small/replacement_startup.sh create mode 100644 baseline/nanogpt_one_head/gpt2small/replacement_worker.sh create mode 100644 baseline/nanogpt_one_head/tests/test_gpt2_reallocation.py diff --git a/baseline/nanogpt_one_head/gpt2small/README.md b/baseline/nanogpt_one_head/gpt2small/README.md index 21568297..04fed3ca 100644 --- a/baseline/nanogpt_one_head/gpt2small/README.md +++ b/baseline/nanogpt_one_head/gpt2small/README.md @@ -72,6 +72,39 @@ a failed cloud copy is an error and leaves persistent-disk outputs intact. convergence or reproduce CE 3.28. A longer matched benchmark is needed before claiming normal GPT-2 training quality, and before the long scientific experiment. +## Replace an expired allocation while retaining FineWeb + +From the updated, clean Cloud Shell checkout: + +```bash +python3 baseline/nanogpt_one_head/gpt2small/reallocate_validation.py launch +python3 baseline/nanogpt_one_head/gpt2small/reallocate_validation.py status +``` + +The launcher replaces only `ww-continuous8-24h-20261003-s1337` and its node. +It retains the existing `ww-continuous8-pilot-20261002-s1337-data` disk and all +cloud objects. It waits for disk detachment and attaches that same disk to one +new v5litepod-8 with a server-enforced four-hour allocation limit, including setup. +The startup script mounts the existing ext4 filesystem; it never formats a disk. +The corpus and Python environment are reused without downloading or reinstalling. +The fixed replacement queue name prevents duplicate launches on repeated commands. + +Short validation runs independently of Cloud Shell in `rg-gpt2-validation.service`. +The new root is `/mnt/disks/rg-data/gpt2small/ww-gpt2-validation-20261004-s1337`. +Cloud uploads are tested before training using object permissions, CRC32C and +size verification. Exit backup uses this same uploader rather than bucket-metadata +operations. The complete logs and outputs remain on the persistent disk on failure. + +This is a fresh validation from initialization, not continuation of the failed run. +Before each optimizer update, the runner checks loss and gradient norm. A failure +writes `nonfinite_diagnostics.json` with parameter names and nonfinite element +counts, then stops before applying the invalid update. The earlier nonfinite +gradient's cause is still unconfirmed; these checks do not claim to fix it. +Only if the short AdamW checks pass does MuonClip validation proceed, followed by +the existing resume checks. No long run starts automatically. The service does +not restart automatically after failures or reboot. A finished service does not +delete its TPU: the allocation limit remains four hours unless stopped earlier. + ## Records, checkpoints and timing Per-step immutable JSON scalar and WW records are written incrementally. NLL, diff --git a/baseline/nanogpt_one_head/gpt2small/reallocate_validation.py b/baseline/nanogpt_one_head/gpt2small/reallocate_validation.py new file mode 100644 index 00000000..0e8993ab --- /dev/null +++ b/baseline/nanogpt_one_head/gpt2small/reallocate_validation.py @@ -0,0 +1,167 @@ +#!/usr/bin/env python3 +"""Replace the expiring TPU, retain its corpus disk, run bounded validation only.""" +import argparse +from datetime import datetime, timezone +import json +from pathlib import Path +import subprocess +import time + +PROJECT = 'tpu-builders-504820' +ZONE = 'us-west4-a' +BUCKET = PROJECT + '-ww-continuous8' +SA = 'rg-continuous-tpu@' + PROJECT + '.iam.gserviceaccount.com' +OLD_QUEUE = 'ww-continuous8-24h-20261003-s1337' +OLD_NODE = OLD_QUEUE + '-node' +QUEUE = 'ww-gpt2-validation-20261004-s1337' +NODE = QUEUE + '-node' +DISK = 'ww-continuous8-pilot-20261002-s1337-data' +DISK_PATH = f'projects/{PROJECT}/zones/{ZONE}/disks/{DISK}' +ROOT = '/mnt/disks/rg-data/gpt2small/' + QUEUE +HOURS = 4 + + +def gc(*args, capture=False): + return subprocess.run(['gcloud', *args, '--project=' + PROJECT], check=True, + text=True, stdout=subprocess.PIPE if capture else None).stdout + + +def inventory(*args): + return json.loads(gc(*args, '--format=json', capture=True)) + + +def named(items, name): + return next((x for x in items if x['name'].rsplit('/', 1)[-1] == name), None) + + +def queues(): + return inventory('alpha', 'compute', 'tpus', 'queued-resources', 'list', '--zone=' + ZONE) + + +def nodes(): + return inventory('compute', 'tpus', 'tpu-vm', 'list', '--zone=' + ZONE) + + +def save(**changes): + path = Path.home() / 'gpt2-validation-resources.json' + record = json.loads(path.read_text()) if path.exists() else {} + record.update(project=PROJECT, zone=ZONE, queue=QUEUE, node=NODE, disk=DISK, + root=ROOT, hours=HOURS, updated_utc=datetime.now(timezone.utc).isoformat(), **changes) + temp = path.with_suffix('.tmp'); temp.write_text(json.dumps(record, indent=2) + '\n'); temp.replace(path) + if 'phase' in changes: + print('[validation] ' + changes['phase'], flush=True) + + +def make_startup(commit): + text = Path(__file__).with_name('replacement_startup.sh').read_text() + return text.replace('__COMMIT__', commit).replace('__ROOT__', ROOT).replace('__HOURS__', str(HOURS)) + + +def status(): + request = named(queues(), QUEUE) + if request is None: + print('No replacement request exists.'); return + detail = inventory('alpha', 'compute', 'tpus', 'queued-resources', 'describe', QUEUE, '--zone=' + ZONE) + state = detail.get('state', {}).get('state') + print('TPU request:', state, flush=True) + print('Cloud outputs: gs://' + BUCKET + '/gpt2small/' + QUEUE, flush=True) + if state == 'ACTIVE': + command = ('sudo systemctl --no-pager --full status rg-gpt2-validation.service || true; ' + f'sudo tail -n 35 {ROOT}/startup.log {ROOT}/run.log; ' + f'if sudo test -f {ROOT}/summaries/validation_report.json; then ' + f'sudo cat {ROOT}/summaries/validation_report.json; fi') + gc('compute', 'tpus', 'tpu-vm', 'ssh', NODE, '--zone=' + ZONE, '--worker=0', '--command=' + command) + + +def launch(): + # Check for our fixed replacement name before any destructive action. + current_queues = queues(); current_nodes = nodes() + if named(current_queues, QUEUE) or named(current_nodes, NODE): + print('Replacement already exists; no deletion, duplicate allocation or validation restart.'); return + repo = Path(__file__).resolve().parents[3] + if subprocess.check_output(['git', '-C', str(repo), 'status', '--porcelain'], text=True).strip(): + raise RuntimeError('Launch from a clean checkout.') + commit = subprocess.check_output(['git', '-C', str(repo), 'rev-parse', 'HEAD'], text=True).strip() + disk = inventory('compute', 'disks', 'describe', DISK, '--zone=' + ZONE) + if disk['zone'].rsplit('/', 1)[-1] != ZONE: + raise RuntimeError('Unexpected disk zone.') + old_queue = named(current_queues, OLD_QUEUE) + if old_queue: + detail = inventory('alpha', 'compute', 'tpus', 'queued-resources', 'describe', OLD_QUEUE, '--zone=' + ZONE) + specs = detail.get('tpu', {}).get('nodeSpec', []) + if len(specs) != 1 or specs[0].get('nodeId') != OLD_NODE: + raise RuntimeError('Old request has unexpected nodes; refusing deletion.') + old_node = named(current_nodes, OLD_NODE) + if old_node: + detail = inventory('compute', 'tpus', 'tpu-vm', 'describe', OLD_NODE, '--zone=' + ZONE) + sources = [d.get('sourceDisk', '').removeprefix('https://www.googleapis.com/compute/v1/') + for d in detail.get('dataDisks', [])] + if sources != [DISK_PATH]: + raise RuntimeError('Old TPU does not have exactly the expected preserved data disk.') + elif disk.get('users') and not old_queue: + raise RuntimeError('Data disk is in use by another resource; no changes made.') + source = make_startup(commit) + subprocess.run(['bash', '-n'], input=source, text=True, check=True) + startup = Path.home() / 'gpt2-validation-startup.sh'; startup.write_text(source) + save(commit=commit, preserved_disk=disk, phase='replacing only the expired/expiring experiment TPU') + print('One v5e-8, at most 4 hours including setup. Estimated compute $19.20 plus storage.', flush=True) + print('Keeping the existing disk, FineWeb, environments and cloud archives. Short validation only.', flush=True) + if old_node and detail.get('state') == 'READY': + # Do not terminate a newly started diagnostic or any unrelated workload. + command = """sudo bash -se <<'CHECK' +if pgrep -af '[p]ython.*(gpt2_experiment|gpt2small/validate.py|continuous_run)'; then + echo 'A trainer is still running; replacement aborted.' >&2 + exit 1 +fi +sync +CHECK""" + try: + gc('compute', 'tpus', 'tpu-vm', 'ssh', OLD_NODE, '--zone=' + ZONE, '--worker=0', '--command=' + command) + except subprocess.CalledProcessError: + # Expiry can race this read-only check. Proceed only if no READY old VM remains. + remaining = named(nodes(), OLD_NODE) + if remaining and remaining.get('state') == 'READY': + raise + if old_queue: + gc('alpha', 'compute', 'tpus', 'queued-resources', 'delete', OLD_QUEUE, + '--zone=' + ZONE, '--force', '--quiet') + if named(nodes(), OLD_NODE): + gc('compute', 'tpus', 'tpu-vm', 'delete', OLD_NODE, '--zone=' + ZONE, '--quiet') + if named(queues(), OLD_QUEUE) or named(nodes(), OLD_NODE): + raise RuntimeError('Old TPU still present; no replacement submitted.') + save(phase='waiting for the preserved data disk to detach') + for attempt in range(25): + if not inventory('compute', 'disks', 'describe', DISK, '--zone=' + ZONE).get('users'): + break + if attempt == 24: + raise RuntimeError('Disk still attached; preserved, no replacement submitted. Rerun later.') + time.sleep(5) + save(phase='submitting four-hour validation allocation') + gc('alpha', 'compute', 'tpus', 'queued-resources', 'create', QUEUE, '--zone=' + ZONE, + '--node-id=' + NODE, '--accelerator-type=v5litepod-8', '--runtime-version=v2-alpha-tpuv5-lite', + '--provisioning-model=flex-start', '--max-run-duration=4h', '--valid-until-duration=4h', + '--service-account=' + SA, '--scopes=https://www.googleapis.com/auth/cloud-platform', + '--data-disk=source=' + DISK_PATH + ',mode=read-write', + '--metadata-from-file=startup-script=' + str(startup), + '--labels=experiment=gpt2-validation', '--quiet', '--async') + save(phase='submitted; validation starts automatically when capacity is allocated') + print('Node: ' + NODE + '\nLog: ' + ROOT + '/run.log', flush=True) + print('Check: python3 baseline/nanogpt_one_head/gpt2small/reallocate_validation.py status', flush=True) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument('action', choices=['launch', 'status'], nargs='?', default='launch') + args = parser.parse_args() + if args.action == 'status': + status() + else: + try: + launch() + except BaseException as exc: + save(phase='stopped; disk and corpus retained', error=str(exc)) + raise + + +if __name__ == '__main__': + main() diff --git a/baseline/nanogpt_one_head/gpt2small/replacement_startup.sh b/baseline/nanogpt_one_head/gpt2small/replacement_startup.sh new file mode 100644 index 00000000..ee72749a --- /dev/null +++ b/baseline/nanogpt_one_head/gpt2small/replacement_startup.sh @@ -0,0 +1,59 @@ +#!/usr/bin/env bash +# Runs as root on a replacement VM. This disk already contains the verified corpus. +set -Eeuo pipefail +DEVICE=/dev/disk/by-id/google-persistent-disk-1 +for attempt in $(seq 1 60); do + [ -b "$DEVICE" ] && break + sleep 2 +done +test -b "$DEVICE" +test "$(blkid -s TYPE -o value "$DEVICE")" = ext4 +mkdir -p /mnt/disks/rg-data +mountpoint -q /mnt/disks/rg-data || mount "$DEVICE" /mnt/disks/rg-data +test "$(readlink -f "$(findmnt -n -o SOURCE --target /mnt/disks/rg-data)")" = "$(readlink -f "$DEVICE")" +ROOT=__ROOT__ +mkdir -p "$ROOT" +exec > >(tee -a "$ROOT/startup.log") 2>&1 +test -s /mnt/disks/rg-data/continuous8/data/meta.json +test -s /mnt/disks/rg-data/continuous8/data/train.bin +# No automatic restart after reboot, including a reboot during initialization. +if ! mkdir "$ROOT/STARTED_ONCE"; then + echo 'Already started once. Retaining all previous outputs; no automatic restart.' + exit 0 +fi +DEADLINE=$(python3 -c 'import time; print(time.time()-float(open("/proc/uptime").read().split()[0])+__HOURS__*3600-600)') +python3 - "$ROOT" "$DEADLINE" <<'PY' +import json,sys +from pathlib import Path +Path(sys.argv[1], 'allocation.json').write_text(json.dumps({ + 'validation_deadline_unix':float(sys.argv[2]), 'max_run_hours':__HOURS__, + 'backup_reserve_seconds':600, 'fresh_validation':True, 'long_run_started':False},indent=2)) +PY +if ! command -v git >/dev/null; then + apt-get update + DEBIAN_FRONTEND=noninteractive apt-get install -y git +fi +mkdir "$ROOT/repo" +git -C "$ROOT/repo" init +git -C "$ROOT/repo" remote add origin https://github.com/CalculatedContent/rg_optimizers.git +git -C "$ROOT/repo" fetch --depth 1 origin __COMMIT__ +git -C "$ROOT/repo" checkout --detach FETCH_HEAD +test "$(git -C "$ROOT/repo" rev-parse HEAD)" = __COMMIT__ +echo __COMMIT__ > "$ROOT/commit.txt" +cat > /etc/systemd/system/rg-gpt2-validation.service <&2 + exit 1 + fi + exit "$rc" +} +trap finish EXIT +source "$base/gpt2small/tpu_environment.sh" +mountpoint -q /mnt/disks/rg-data +test -x "$python" +# Reuse installed dependencies and the verified corpus. Do not reinstall or redownload. +"$python" -c 'import sys,torch,torch_xla,weightwatcher,yaml; print("Dependencies loaded:",sys.version,torch.__version__,torch_xla.__version__,flush=True)' +"$python" - "$root" <<'PY' +import json,sys +from pathlib import Path +from rg_nanogpt_one_head.continuous_support import CloudPublisher +root=Path(sys.argv[1]) +probe=root/'backup_probe.json'; probe.write_text(json.dumps({'run':root.name,'phase':'before_training'})) +sink=CloudPublisher('gs://tpu-builders-504820-ww-continuous8/gpt2small/'+root.name) +receipt=sink.file(probe,probe.name) +(root/'backup_probe_receipt.json').write_text(json.dumps(receipt,indent=2)) +print('Cloud upload permission and checksum verification passed before training.',flush=True) +PY +"$python" -u "$base/gpt2small/validate.py" --root "$root" \ + --data /mnt/disks/rg-data/continuous8/data --deadline "$deadline" 2>&1 | tee "$root/validation.log" +echo 'Short validation completed. No long experiment launched.' diff --git a/baseline/nanogpt_one_head/tests/test_gpt2_reallocation.py b/baseline/nanogpt_one_head/tests/test_gpt2_reallocation.py new file mode 100644 index 00000000..fa5fbcd0 --- /dev/null +++ b/baseline/nanogpt_one_head/tests/test_gpt2_reallocation.py @@ -0,0 +1,76 @@ +"""Exercise allocation ordering and preservation gates without contacting GCP.""" +import importlib.util +from pathlib import Path +import subprocess + +import pytest + + +def load(): + path = Path(__file__).resolve().parents[1] / 'gpt2small/reallocate_validation.py' + spec = importlib.util.spec_from_file_location('gpt2_reallocation', path) + module = importlib.util.module_from_spec(spec); spec.loader.exec_module(module) + return module + + +def test_startup_mounts_only_existing_disk_and_never_restarts(): + m = load(); script = m.make_startup('a' * 40) + subprocess.run(['bash', '-n'], input=script, text=True, check=True) + assert '__' not in script and 'mkfs' not in script + assert 'readlink -f "$DEVICE"' in script + assert '+4*3600-600' in script and 'Restart=no' in script + assert 'STARTED_ONCE' in script and 'replacement_worker.sh' in script + worker = Path(m.__file__).with_name('replacement_worker.sh').read_text() + subprocess.run(['bash', '-n'], input=worker, text=True, check=True) + assert worker.index('receipt=sink.file') < worker.index('gpt2small/validate.py') + assert 'prepare_tpu_data' not in worker and '--allow-long-run' not in worker + + +@pytest.mark.parametrize('case', ['normal', 'duplicate', 'wrong_disk', 'delete_failure', 'other_user']) +def test_replacement_sequence(tmp_path, monkeypatch, case): + m = load(); calls = [] + old_present = case != 'other_user' + monkeypatch.setattr(Path, 'home', classmethod(lambda cls: tmp_path)) + monkeypatch.setattr(m.subprocess, 'check_output', + lambda args, **kw: '' if 'status' in args else 'a'*40 + '\n') + + def inventory(*args): + if 'queued-resources' in args: + if 'describe' in args: + return {'tpu': {'nodeSpec': [{'nodeId': m.OLD_NODE}]}} + return ([{'name': m.QUEUE}] if case == 'duplicate' else + [{'name': m.OLD_QUEUE}] if old_present else []) + if 'tpu-vm' in args: + if 'describe' in args: + return {'state': 'READY', 'dataDisks': [{'sourceDisk': + 'wrong' if case == 'wrong_disk' else m.DISK_PATH}]} + return [{'name': m.OLD_NODE, 'state': 'READY'}] if old_present else [] + if args[:3] == ('compute', 'disks', 'describe'): + return {'zone': m.ZONE, 'users': ['an-attached-vm'] if old_present or case == 'other_user' else []} + raise AssertionError(args) + + def gc(*args, **kwargs): + nonlocal old_present + calls.append(args) + if 'delete' in args: + if case == 'delete_failure': + raise RuntimeError('delete failed') + assert 'queued-resources' in args and m.OLD_QUEUE in args + old_present = False + if 'create' in args: + assert not old_present + assert '--max-run-duration=4h' in args + assert '--accelerator-type=v5litepod-8' in args + assert '--data-disk=source=' + m.DISK_PATH + ',mode=read-write' in args + + monkeypatch.setattr(m, 'inventory', inventory) + monkeypatch.setattr(m, 'gc', gc) + if case in ('wrong_disk', 'delete_failure', 'other_user'): + with pytest.raises(RuntimeError): + m.launch() + else: + m.launch() + assert sum('create' in call for call in calls) == (case == 'normal') + assert not any('disks' in call and 'delete' in call for call in calls) + if case in ('duplicate', 'wrong_disk', 'other_user'): + assert calls == [] From 4631a3dac08fc81976c117fe7d2250ba52f5b16e Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Sat, 3 Oct 2026 18:43:23 -0700 Subject: [PATCH 15/36] Replace four-hour validation TPU with an optional 48-hour allocation --- baseline/nanogpt_one_head/gpt2small/README.md | 23 ++++++++++++ .../gpt2small/reallocate_validation.py | 36 +++++++++++++++---- .../tests/test_gpt2_reallocation.py | 22 ++++++++---- 3 files changed, 68 insertions(+), 13 deletions(-) diff --git a/baseline/nanogpt_one_head/gpt2small/README.md b/baseline/nanogpt_one_head/gpt2small/README.md index 04fed3ca..57cd1fa9 100644 --- a/baseline/nanogpt_one_head/gpt2small/README.md +++ b/baseline/nanogpt_one_head/gpt2small/README.md @@ -105,6 +105,29 @@ the existing resume checks. No long run starts automatically. The service does not restart automatically after failures or reboot. A finished service does not delete its TPU: the allocation limit remains four hours unless stopped earlier. +### Allocate 48 hours instead + +```bash +python3 baseline/nanogpt_one_head/gpt2small/reallocate_validation.py launch --hours 48 +python3 baseline/nanogpt_one_head/gpt2small/reallocate_validation.py status --hours 48 +``` + +The queued-resource API exposes no update operation for extending the requested +lifetime. This replaces only the four-hour validation request with +`ww-gpt2-validation-48h-20261004-s1337`. If its validation service is already +running, it is stopped before deletion; all existing files remain on the same +data disk. The new VM mounts that disk, reuses FineWeb and the environment, and +runs fresh short validation in a separate directory. Repeating this command +does not replace an existing 48-hour request or create a second machine. + +Both the server-enforced allocation limit and the startup deadline use 48 hours; +the worker deadline reserves ten minutes, and validation reserves another five. +The queue can wait up to four hours for capacity, independently of the 48-hour +allocation lifetime. At the published $0.60/chip-hour Flex-start rate, eight +chips for 48 hours cost $230.40 before storage and any earlier allocation usage. +Short validation still stops after its checks; it does not automatically start +the long experiment. The allocation remains available until deletion or expiry. + ## Records, checkpoints and timing Per-step immutable JSON scalar and WW records are written incrementally. NLL, diff --git a/baseline/nanogpt_one_head/gpt2small/reallocate_validation.py b/baseline/nanogpt_one_head/gpt2small/reallocate_validation.py index 0e8993ab..be9cfee5 100644 --- a/baseline/nanogpt_one_head/gpt2small/reallocate_validation.py +++ b/baseline/nanogpt_one_head/gpt2small/reallocate_validation.py @@ -19,6 +19,23 @@ DISK_PATH = f'projects/{PROJECT}/zones/{ZONE}/disks/{DISK}' ROOT = '/mnt/disks/rg-data/gpt2small/' + QUEUE HOURS = 4 +OLD_SERVICE = None + + +def configure(hours): + """The 48-hour request replaces only our previous four-hour validation VM.""" + global HOURS, OLD_QUEUE, OLD_NODE, QUEUE, NODE, ROOT, OLD_SERVICE + if hours not in (4, 48): + raise ValueError('Supported allocation lengths are 4 and 48 hours.') + HOURS = hours + OLD_QUEUE = ('ww-continuous8-24h-20261003-s1337' if hours == 4 + else 'ww-gpt2-validation-20261004-s1337') + OLD_NODE = OLD_QUEUE + '-node' + QUEUE = ('ww-gpt2-validation-20261004-s1337' if hours == 4 + else 'ww-gpt2-validation-48h-20261004-s1337') + NODE = QUEUE + '-node' + ROOT = '/mnt/disks/rg-data/gpt2small/' + QUEUE + OLD_SERVICE = 'rg-gpt2-validation.service' if hours == 48 else None def gc(*args, capture=False): @@ -103,12 +120,15 @@ def launch(): source = make_startup(commit) subprocess.run(['bash', '-n'], input=source, text=True, check=True) startup = Path.home() / 'gpt2-validation-startup.sh'; startup.write_text(source) - save(commit=commit, preserved_disk=disk, phase='replacing only the expired/expiring experiment TPU') - print('One v5e-8, at most 4 hours including setup. Estimated compute $19.20 plus storage.', flush=True) + save(commit=commit, preserved_disk=disk, phase='replacing only ' + OLD_QUEUE) + print(f'One v5e-8, at most {HOURS} hours including setup. Estimated compute ${HOURS*8*.60:.2f} plus storage.', flush=True) print('Keeping the existing disk, FineWeb, environments and cloud archives. Short validation only.', flush=True) if old_node and detail.get('state') == 'READY': - # Do not terminate a newly started diagnostic or any unrelated workload. - command = """sudo bash -se <<'CHECK' + # A requested duration change may interrupt the known validation service. + # Other trainer processes must still prevent deletion of this VM. + stop = (f'if systemctl cat {OLD_SERVICE} >/dev/null 2>&1; then\n' + f' systemctl stop {OLD_SERVICE}\nfi\n') if OLD_SERVICE else '' + command = "sudo bash -se <<'CHECK'\n" + stop + """ if pgrep -af '[p]ython.*(gpt2_experiment|gpt2small/validate.py|continuous_run)'; then echo 'A trainer is still running; replacement aborted.' >&2 exit 1 @@ -136,23 +156,25 @@ def launch(): if attempt == 24: raise RuntimeError('Disk still attached; preserved, no replacement submitted. Rerun later.') time.sleep(5) - save(phase='submitting four-hour validation allocation') + save(phase=f'submitting {HOURS}-hour allocation with short validation') gc('alpha', 'compute', 'tpus', 'queued-resources', 'create', QUEUE, '--zone=' + ZONE, '--node-id=' + NODE, '--accelerator-type=v5litepod-8', '--runtime-version=v2-alpha-tpuv5-lite', - '--provisioning-model=flex-start', '--max-run-duration=4h', '--valid-until-duration=4h', + '--provisioning-model=flex-start', f'--max-run-duration={HOURS}h', '--valid-until-duration=4h', '--service-account=' + SA, '--scopes=https://www.googleapis.com/auth/cloud-platform', '--data-disk=source=' + DISK_PATH + ',mode=read-write', '--metadata-from-file=startup-script=' + str(startup), '--labels=experiment=gpt2-validation', '--quiet', '--async') save(phase='submitted; validation starts automatically when capacity is allocated') print('Node: ' + NODE + '\nLog: ' + ROOT + '/run.log', flush=True) - print('Check: python3 baseline/nanogpt_one_head/gpt2small/reallocate_validation.py status', flush=True) + print(f'Check: python3 baseline/nanogpt_one_head/gpt2small/reallocate_validation.py status --hours {HOURS}', flush=True) def main(): parser = argparse.ArgumentParser(description=__doc__) parser.add_argument('action', choices=['launch', 'status'], nargs='?', default='launch') + parser.add_argument('--hours', type=int, choices=[4, 48], default=4) args = parser.parse_args() + configure(args.hours) if args.action == 'status': status() else: diff --git a/baseline/nanogpt_one_head/tests/test_gpt2_reallocation.py b/baseline/nanogpt_one_head/tests/test_gpt2_reallocation.py index fa5fbcd0..52bd8ecd 100644 --- a/baseline/nanogpt_one_head/tests/test_gpt2_reallocation.py +++ b/baseline/nanogpt_one_head/tests/test_gpt2_reallocation.py @@ -13,12 +13,13 @@ def load(): return module -def test_startup_mounts_only_existing_disk_and_never_restarts(): - m = load(); script = m.make_startup('a' * 40) +@pytest.mark.parametrize('hours', [4, 48]) +def test_startup_mounts_only_existing_disk_and_never_restarts(hours): + m = load(); m.configure(hours); script = m.make_startup('a' * 40) subprocess.run(['bash', '-n'], input=script, text=True, check=True) assert '__' not in script and 'mkfs' not in script assert 'readlink -f "$DEVICE"' in script - assert '+4*3600-600' in script and 'Restart=no' in script + assert f'+{hours}*3600-600' in script and 'Restart=no' in script assert 'STARTED_ONCE' in script and 'replacement_worker.sh' in script worker = Path(m.__file__).with_name('replacement_worker.sh').read_text() subprocess.run(['bash', '-n'], input=worker, text=True, check=True) @@ -27,8 +28,9 @@ def test_startup_mounts_only_existing_disk_and_never_restarts(): @pytest.mark.parametrize('case', ['normal', 'duplicate', 'wrong_disk', 'delete_failure', 'other_user']) -def test_replacement_sequence(tmp_path, monkeypatch, case): - m = load(); calls = [] +@pytest.mark.parametrize('hours', [4, 48]) +def test_replacement_sequence(tmp_path, monkeypatch, case, hours): + m = load(); m.configure(hours); calls = [] old_present = case != 'other_user' monkeypatch.setattr(Path, 'home', classmethod(lambda cls: tmp_path)) monkeypatch.setattr(m.subprocess, 'check_output', @@ -59,7 +61,7 @@ def gc(*args, **kwargs): old_present = False if 'create' in args: assert not old_present - assert '--max-run-duration=4h' in args + assert f'--max-run-duration={hours}h' in args assert '--accelerator-type=v5litepod-8' in args assert '--data-disk=source=' + m.DISK_PATH + ',mode=read-write' in args @@ -74,3 +76,11 @@ def gc(*args, **kwargs): assert not any('disks' in call and 'delete' in call for call in calls) if case in ('duplicate', 'wrong_disk', 'other_user'): assert calls == [] + if case == 'normal': + ssh = next(c for c in calls if 'ssh' in c) + command = next(a for a in ssh if a.startswith('--command=')) + assert ('systemctl stop rg-gpt2-validation.service' in command) == (hours == 48) + if hours == 48: + assert m.OLD_NODE == 'ww-gpt2-validation-20261004-s1337-node' + assert m.NODE == 'ww-gpt2-validation-48h-20261004-s1337-node' + assert command.index('systemctl stop') < command.index('pgrep') From 02c3d6a2e6fa21a6f47de31b7195b3f1c5edd107 Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Sat, 3 Oct 2026 21:18:28 -0700 Subject: [PATCH 16/36] Capture stalled GPT-2 validation stacks before stopping only its service --- baseline/nanogpt_one_head/gpt2small/README.md | 25 +++ .../gpt2small/capture_stall.py | 183 ++++++++++++++++++ .../tests/test_gpt2_stall_capture.py | 72 +++++++ 3 files changed, 280 insertions(+) create mode 100644 baseline/nanogpt_one_head/gpt2small/capture_stall.py create mode 100644 baseline/nanogpt_one_head/tests/test_gpt2_stall_capture.py diff --git a/baseline/nanogpt_one_head/gpt2small/README.md b/baseline/nanogpt_one_head/gpt2small/README.md index 57cd1fa9..1d8ec1d2 100644 --- a/baseline/nanogpt_one_head/gpt2small/README.md +++ b/baseline/nanogpt_one_head/gpt2small/README.md @@ -130,6 +130,31 @@ the long experiment. The allocation remains available until deletion or expiry. ## Records, checkpoints and timing +### Capture a stalled validation + +From Cloud Shell on the updated branch: + +```bash +python3 baseline/nanogpt_one_head/gpt2small/capture_stall.py --stop +``` + +This targets the 48-hour validation run above. It captures process CPU/memory, +logs, saved XLA metrics and live Python/native stacks before stopping only +`rg-gpt2-validation.service`. A bounded, optional py-spy installation goes in a +temporary directory, leaving the training environment and checkout unchanged. +If profiling is unavailable, process/log diagnostics are still retained. +Reports remain under the run's `diagnostics/stall-*` directory on the persistent +disk; a compact summary is printed for sharing. This command does not upload them. +Omit `--stop` for a read-only capture. A blocked update may be lost on stop; no new +checkpoint is promised. Existing checkpoints, FineWeb and the TPU allocation +remain. The allocation continues to incur compute usage until deletion/expiry. + +The validator's `WAIT` lines only indicate a live process, not a completed update. +Likewise `before_update: 2` confirms the gradient check before update 2, not its +completion. An extended wait requires inspection, not an assumed compilation ETA. + +### Scientific records + Per-step immutable JSON scalar and WW records are written incrementally. NLL, perplexity, top-1 accuracy, error (fraction), steps, token presentations, wall time, LR and pre-clipping gradient norm are included. All WW library columns are retained, diff --git a/baseline/nanogpt_one_head/gpt2small/capture_stall.py b/baseline/nanogpt_one_head/gpt2small/capture_stall.py new file mode 100644 index 00000000..f76c53e9 --- /dev/null +++ b/baseline/nanogpt_one_head/gpt2small/capture_stall.py @@ -0,0 +1,183 @@ +"""Capture a live validation stall, optionally stop its service; keep TPU/data. + +Run from Cloud Shell. The script sends itself over SSH without changing the +running checkout or Python environment. No GCP allocation/storage mutations. +""" +import argparse +import datetime as dt +import json +import os +from pathlib import Path +import shlex +import subprocess +import sys +import tempfile +import time + +PROJECT = 'tpu-builders-504820' +ZONE = 'us-west4-a' +RUN = 'ww-gpt2-validation-48h-20261004-s1337' +ROOT = Path('/mnt/disks/rg-data/gpt2small') / RUN +SERVICE = 'rg-gpt2-validation.service' +PYTHON = '/mnt/disks/rg-data/continuous8/venv/bin/python' + + +def service_info(): + result = subprocess.run( + ['systemctl', 'show', SERVICE, '--property=ExecStart,ActiveState,SubState,MainPID,ControlGroup'], + capture_output=True, text=True, check=True, timeout=15) + return dict(line.split('=', 1) for line in result.stdout.splitlines() if '=' in line) + + +def verify_target(info, root): + expected = str(root / 'repo/baseline/nanogpt_one_head/gpt2small/replacement_worker.sh') + if expected not in info.get('ExecStart', ''): + raise RuntimeError('Service belongs to another run; refusing to stop it.') + + +def training_pids(root, proc=Path('/proc')): + found = [] + for folder in proc.iterdir(): + if not folder.name.isdigit(): + continue + try: + args = (folder / 'cmdline').read_bytes().decode().strip('\0').split('\0') + if 'rg_nanogpt_one_head.gpt2_experiment' not in args or '--output' not in args: + continue + output = Path(args[args.index('--output') + 1]) + if output.parent == root and output.name in ('adamw', 'muonclip'): + found.append(int(folder.name)) + except (OSError, UnicodeError, IndexError): + continue + return sorted(found) + + +def capture_command(destination, args, timeout=20): + with destination.open('w') as out: + try: + result = subprocess.run(args, stdout=out, stderr=subprocess.STDOUT, timeout=timeout) + out.write(f'\n[exit code {result.returncode}]\n') + return result.returncode == 0 + except (OSError, subprocess.TimeoutExpired) as exc: + out.write(f'\n[capture unavailable: {exc}]\n') + return False + + +def stop_service(root, report): + verify_target(service_info(), root) + # Queue the stop so the client's timeout cannot leave an unobserved stop job. + subprocess.run(['systemctl', 'stop', '--no-block', SERVICE], check=True, timeout=15) + for attempt in range(30): + state = service_info() + if state.get('ActiveState') in ('inactive', 'failed') and state.get('MainPID') == '0': + report['service_stopped'] = True + return + if attempt in (0, 10, 20): + print('Waiting for validation service to stop...', flush=True) + time.sleep(2) + # The blocked update may ignore TERM. Only this verified service is targeted. + verify_target(service_info(), root) + report['forced_service_kill'] = True + subprocess.run(['systemctl', 'kill', '--kill-who=all', '--signal=KILL', SERVICE], + check=True, timeout=15) + for _ in range(10): + state = service_info() + if state.get('ActiveState') in ('inactive', 'failed') and state.get('MainPID') == '0': + report['service_stopped'] = True + return + time.sleep(1) + raise RuntimeError('Stop is not confirmed; inspect service-after.txt. No TPU was deleted.') + + +def on_tpu(stop=False): + if os.geteuid() != 0: + raise RuntimeError('Run the remote capture as root.') + if not os.path.ismount('/mnt/disks/rg-data') or not ROOT.is_dir(): + raise RuntimeError('Expected mounted disk/run directory is missing.') + info = service_info() + verify_target(info, ROOT) + stamp = dt.datetime.now(dt.timezone.utc).strftime('%Y%m%d-%H%M%S-%f') + dest = ROOT / 'diagnostics' / ('stall-' + stamp) + dest.mkdir(parents=True) + pids = training_pids(ROOT) + report = {'captured_at_utc': stamp, 'run': str(ROOT), 'training_pids': pids, + 'service_before': info, 'stop_requested': stop, 'service_stopped': False} + print('Saving diagnostics:', dest, flush=True) + try: + capture_command(dest / 'memory.txt', ['free', '-h']) + capture_command(dest / 'disk.txt', ['df', '-h', '/mnt/disks/rg-data']) + capture_command(dest / 'journal.txt', ['journalctl', '-u', SERVICE, '-n', '80', '--no-pager']) + for source in (ROOT / 'run.log', ROOT / 'validation.log', ROOT / 'logs/adamw_4.log', + ROOT / 'adamw/status.json', ROOT / 'adamw/checkpoints/latest.json', + ROOT / 'adamw/logs/xla_compile_metrics_after_step_2.txt'): + name = str(source.relative_to(ROOT)).replace('/', '_') + capture_command(dest / name, ['tail', '-n', '150', str(source)]) + for pid in pids: + for name in ('status', 'stat', 'io', 'wchan', 'stack', 'cgroup'): + capture_command(dest / f'pid-{pid}-{name}.txt', ['cat', f'/proc/{pid}/{name}']) + capture_command(dest / f'pid-{pid}-cpu-before.txt', + ['ps', '-p', str(pid), '-o', 'pid,stat,etime,time,pcpu,rss,wchan:32']) + if pids: + # Install only the profiler binary in a temporary directory; no training dependencies change. + with tempfile.TemporaryDirectory(prefix='rg-stall-profiler-') as tools: + installed = capture_command(dest / 'profiler-install.txt', [PYTHON, '-m', 'pip', 'install', + '--target', tools, '--no-deps', '--no-cache-dir', '--only-binary=:all:', + '--disable-pip-version-check', '--progress-bar', 'off', '--retries', '1', + '--timeout', '15', 'py-spy==0.4.1'], timeout=60) + if installed: + for pid in pids: + print('Capturing Python and native stacks for PID', pid, flush=True) + spy = str(Path(tools) / 'bin/py-spy') + capture_command(dest / f'pid-{pid}-python-stack.txt', + [spy, 'dump', '--pid', str(pid)], timeout=15) + capture_command(dest / f'pid-{pid}-native-stack.txt', + [spy, 'dump', '--native', '--pid', str(pid)], timeout=20) + else: + print('Profiler unavailable; process/log diagnostics retained.', flush=True) + time.sleep(3) + for pid in pids: + capture_command(dest / f'pid-{pid}-cpu-after.txt', + ['ps', '-p', str(pid), '-o', 'pid,stat,etime,time,pcpu,rss,wchan:32']) + capture_command(dest / f'pid-{pid}-threads.txt', + ['ps', '-L', '-p', str(pid), '-o', 'pid,tid,stat,time,pcpu,wchan:32,comm']) + if stop: + print('Stopping validation; the unfinished update may be lost.', flush=True) + stop_service(ROOT, report) + finally: + capture_command(dest / 'service-after.txt', ['systemctl', 'status', SERVICE, '--no-pager', '--full']) + (dest / 'capture.json').write_text(json.dumps(report, indent=2)) + os.sync() + for filename in ['adamw_checkpoints_latest.json'] + [ + f'pid-{pid}-{part}.txt' for pid in pids for part in ('cpu-before', 'cpu-after', 'python-stack')]: + print('\n' + filename, flush=True) + print('\n'.join((dest / filename).read_text().splitlines()[:65]) + if (dest / filename).is_file() else 'Not captured.', flush=True) + print('\nFull diagnostics saved on persistent disk:', dest, flush=True) + if report['service_stopped']: + print('Validation stopped. TPU allocation, FineWeb and existing checkpoints retained.', flush=True) + elif not stop: + print('Read-only capture complete; validation was not stopped.', flush=True) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument('--on-tpu', action='store_true', help=argparse.SUPPRESS) + parser.add_argument('--stop', action='store_true', help='Stop only the validation service after capture') + args = parser.parse_args() + if args.on_tpu: + on_tpu(args.stop) + return 0 + remote = ['sudo', 'python3', '-c', Path(__file__).read_text(), '--on-tpu'] + if args.stop: + remote.append('--stop') + # Do not alter the running checkout or ask it to fetch/checkout another commit. + return subprocess.run(['gcloud', 'compute', 'tpus', 'tpu-vm', 'ssh', RUN + '-node', + '--project=' + PROJECT, '--zone=' + ZONE, '--worker=0', '--command=' + shlex.join(remote)]).returncode + + +if __name__ == '__main__': + try: + sys.exit(main()) + except Exception as exc: + print('Diagnostic failed:', exc, file=sys.stderr) + sys.exit(1) diff --git a/baseline/nanogpt_one_head/tests/test_gpt2_stall_capture.py b/baseline/nanogpt_one_head/tests/test_gpt2_stall_capture.py new file mode 100644 index 00000000..e99ad0a7 --- /dev/null +++ b/baseline/nanogpt_one_head/tests/test_gpt2_stall_capture.py @@ -0,0 +1,72 @@ +import importlib.util +from pathlib import Path + +import pytest + + +spec = importlib.util.spec_from_file_location( + 'stall_capture', Path(__file__).parents[1] / 'gpt2small/capture_stall.py') +capture = importlib.util.module_from_spec(spec) +spec.loader.exec_module(capture) + + +def service(root, active='active', pid='6800'): + return {'ExecStart': f'/bin/bash {root}/repo/baseline/nanogpt_one_head/gpt2small/replacement_worker.sh {root}', + 'ActiveState': active, 'MainPID': pid} + + +def test_process_selection_excludes_other_training_runs(tmp_path): + root = Path('/mnt/disks/rg-data/gpt2small/current') + for pid, args in { + 11: ['python', '-m', 'rg_nanogpt_one_head.gpt2_experiment', '--output', str(root / 'adamw')], + 12: ['python', '-m', 'rg_nanogpt_one_head.gpt2_experiment', '--output', str(root / 'muonclip')], + 13: ['python', '-m', 'rg_nanogpt_one_head.gpt2_experiment', '--output', str(root.parent / 'other/adamw')], + 14: ['python', 'validate.py', '--root', str(root)], + }.items(): + folder = tmp_path / str(pid) + folder.mkdir() + (folder / 'cmdline').write_bytes(('\0'.join(args) + '\0').encode()) + assert capture.training_pids(root, tmp_path) == [11, 12] + + +def test_stop_refuses_other_service_root(monkeypatch): + monkeypatch.setattr(capture, 'service_info', lambda: service('/some/other/run')) + commands = [] + monkeypatch.setattr(capture.subprocess, 'run', lambda args, **kwargs: commands.append(args)) + with pytest.raises(RuntimeError, match='another run'): + capture.stop_service(capture.ROOT, {}) + assert commands == [] + + +def test_clean_service_stop_does_not_force_kill(monkeypatch): + states = iter([service(capture.ROOT), service(capture.ROOT, 'inactive', '0')]) + monkeypatch.setattr(capture, 'service_info', lambda: next(states)) + commands = [] + monkeypatch.setattr(capture.subprocess, 'run', lambda args, **kwargs: commands.append(args)) + report = {} + capture.stop_service(capture.ROOT, report) + assert report == {'service_stopped': True} + assert commands == [['systemctl', 'stop', '--no-block', capture.SERVICE]] + + +def test_stuck_service_kill_is_scoped_and_verified(monkeypatch): + states = iter([service(capture.ROOT)] * 32 + [service(capture.ROOT, 'failed', '0')]) + monkeypatch.setattr(capture, 'service_info', lambda: next(states)) + monkeypatch.setattr(capture.time, 'sleep', lambda seconds: None) + commands = [] + monkeypatch.setattr(capture.subprocess, 'run', lambda args, **kwargs: commands.append(args)) + report = {} + capture.stop_service(capture.ROOT, report) + assert report == {'service_stopped': True, 'forced_service_kill': True} + assert commands[-1] == ['systemctl', 'kill', '--kill-who=all', '--signal=KILL', capture.SERVICE] + + +def test_changed_service_is_not_force_killed(monkeypatch): + states = iter([service(capture.ROOT)] * 31 + [service('/different/run')]) + monkeypatch.setattr(capture, 'service_info', lambda: next(states)) + monkeypatch.setattr(capture.time, 'sleep', lambda seconds: None) + commands = [] + monkeypatch.setattr(capture.subprocess, 'run', lambda args, **kwargs: commands.append(args)) + with pytest.raises(RuntimeError, match='another run'): + capture.stop_service(capture.ROOT, {}) + assert len(commands) == 1 From f887593793f3bad6a3319630b3f9688c670ac379 Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Sat, 3 Oct 2026 21:29:36 -0700 Subject: [PATCH 17/36] Move GPT-2 Small into its own package and fail promptly on invalid validation gradients --- baseline/README.md | 5 ++ .../gpt2small => gpt2_small}/README.md | 51 ++++++++++++++----- .../configs/gpt2_small_cpu_smoke.yaml | 0 .../gpt2_small_fineweb_adamw_baseline.yaml | 0 .../gpt2_small_fineweb_muonclip_baseline.yaml | 0 .../gpt2_small_fineweb_muonclip_long_ww.yaml | 0 baseline/gpt2_small/pyproject.toml | 16 ++++++ .../scripts}/LOCAL_VALIDATION.json | 0 .../scripts}/analyze.py | 0 .../scripts}/backup.py | 0 .../scripts}/capture_stall.py | 38 ++++++++++++-- .../scripts}/cloudshell.py | 0 .../scripts}/diagnose_adamw.sh | 8 +-- .../scripts}/on_tpu.sh | 10 ++-- .../scripts}/prepare_existing.py | 0 .../scripts}/reallocate_validation.py | 4 +- .../scripts}/replacement_startup.sh | 2 +- .../scripts}/replacement_worker.sh | 8 +-- .../scripts}/retry_existing.sh | 8 +-- .../scripts}/tpu_environment.sh | 0 .../scripts}/validate.py | 19 ++++--- .../gpt2_small/scripts/validation_watchdog.py | 21 ++++++++ .../gpt2_small/src/rg_gpt2_small/__init__.py | 1 + .../src/rg_gpt2_small/experiment.py} | 44 +++++++++------- .../tests/test_gpt2_experiment.py | 29 +++++++++-- .../tests/test_gpt2_reallocation.py | 4 +- .../tests/test_gpt2_stall_capture.py | 10 ++-- .../tests/test_validation_watchdog.py | 37 ++++++++++++++ baseline/nanogpt_one_head/README.md | 2 +- 29 files changed, 244 insertions(+), 73 deletions(-) rename baseline/{nanogpt_one_head/gpt2small => gpt2_small}/README.md (79%) rename baseline/{nanogpt_one_head => gpt2_small}/configs/gpt2_small_cpu_smoke.yaml (100%) rename baseline/{nanogpt_one_head => gpt2_small}/configs/gpt2_small_fineweb_adamw_baseline.yaml (100%) rename baseline/{nanogpt_one_head => gpt2_small}/configs/gpt2_small_fineweb_muonclip_baseline.yaml (100%) rename baseline/{nanogpt_one_head => gpt2_small}/configs/gpt2_small_fineweb_muonclip_long_ww.yaml (100%) create mode 100644 baseline/gpt2_small/pyproject.toml rename baseline/{nanogpt_one_head/gpt2small => gpt2_small/scripts}/LOCAL_VALIDATION.json (100%) rename baseline/{nanogpt_one_head/gpt2small => gpt2_small/scripts}/analyze.py (100%) rename baseline/{nanogpt_one_head/gpt2small => gpt2_small/scripts}/backup.py (100%) rename baseline/{nanogpt_one_head/gpt2small => gpt2_small/scripts}/capture_stall.py (81%) rename baseline/{nanogpt_one_head/gpt2small => gpt2_small/scripts}/cloudshell.py (100%) rename baseline/{nanogpt_one_head/gpt2small => gpt2_small/scripts}/diagnose_adamw.sh (69%) rename baseline/{nanogpt_one_head/gpt2small => gpt2_small/scripts}/on_tpu.sh (75%) rename baseline/{nanogpt_one_head/gpt2small => gpt2_small/scripts}/prepare_existing.py (100%) rename baseline/{nanogpt_one_head/gpt2small => gpt2_small/scripts}/reallocate_validation.py (97%) rename baseline/{nanogpt_one_head/gpt2small => gpt2_small/scripts}/replacement_startup.sh (95%) rename baseline/{nanogpt_one_head/gpt2small => gpt2_small/scripts}/replacement_worker.sh (86%) rename baseline/{nanogpt_one_head/gpt2small => gpt2_small/scripts}/retry_existing.sh (70%) rename baseline/{nanogpt_one_head/gpt2small => gpt2_small/scripts}/tpu_environment.sh (100%) rename baseline/{nanogpt_one_head/gpt2small => gpt2_small/scripts}/validate.py (88%) create mode 100644 baseline/gpt2_small/scripts/validation_watchdog.py create mode 100644 baseline/gpt2_small/src/rg_gpt2_small/__init__.py rename baseline/{nanogpt_one_head/src/rg_nanogpt_one_head/gpt2_experiment.py => gpt2_small/src/rg_gpt2_small/experiment.py} (91%) rename baseline/{nanogpt_one_head => gpt2_small}/tests/test_gpt2_experiment.py (80%) rename baseline/{nanogpt_one_head => gpt2_small}/tests/test_gpt2_reallocation.py (95%) rename baseline/{nanogpt_one_head => gpt2_small}/tests/test_gpt2_stall_capture.py (84%) create mode 100644 baseline/gpt2_small/tests/test_validation_watchdog.py diff --git a/baseline/README.md b/baseline/README.md index 0dcfec5c..64bb39a1 100644 --- a/baseline/README.md +++ b/baseline/README.md @@ -382,3 +382,8 @@ GitHub Actions additionally compiles all Python sources, parses notebook code cells, runs a real pinned WeightWatcher integration, and executes a pinned nanochat CPU model/optimizer preflight. These bounded checks do not replace the full three-seed long-horizon campaigns or the required target-MPS preflight. + +## GPT-2 Small + +See [gpt2_small/README.md](gpt2_small/README.md) for the 124M-parameter, +12-layer, 12-head FineWeb TPU validation workflow. diff --git a/baseline/nanogpt_one_head/gpt2small/README.md b/baseline/gpt2_small/README.md similarity index 79% rename from baseline/nanogpt_one_head/gpt2small/README.md rename to baseline/gpt2_small/README.md index 1d8ec1d2..1aea7d57 100644 --- a/baseline/nanogpt_one_head/gpt2small/README.md +++ b/baseline/gpt2_small/README.md @@ -1,7 +1,31 @@ # GPT-2 Small / FineWeb-Edu validation -This is a new, isolated experiment. Existing tiny and continuous8 workflows are unchanged. -It reuses the repository GPT, advanced MuonClip, SPMD, corpus validation and WeightWatcher adapter. +This is the GPT-2 Small experiment, under `baseline/gpt2_small`. Its training module +is `rg_gpt2_small.experiment`; configurations and launchers live here. It reuses the +shared GPT implementation, optimizers, SPMD, corpus validation and WeightWatcher +adapter from the older `rg_nanogpt_one_head` infrastructure package. That dependency +does not constrain the head count. Existing one-head and continuous8 workflows stay +in their original directory. Existing disk and bucket paths are retained for data reuse. + +The TPU launch scripts expose both source packages through `PYTHONPATH`. For local +development, from the repository root, install both with +`pip install -e baseline/nanogpt_one_head -e baseline/gpt2_small`. + +## Current validation failure + +The saved 2026-10-04 04:19 UTC stack from commit `4631a3d` is inside the +nonfinite-gradient diagnostic's per-parameter CPU copy. That branch is reached only +after detecting a nonfinite loss or aggregate gradient norm. The latest checkpoint +pointer was step 0. This is a failed numerical validation, not evidence of healthy +training. The origin of the nonfinite result is still unconfirmed. + +The failure handler now writes the scalar failure and raises immediately, avoiding +the full gradient-copy loop. Validation synchronizes XLA before reading diagnostic +scalars, prints its pre-update stage, and emits Python stacks every five minutes. +Each validation phase has a 30-minute limit (also bounded by allocation time); +on expiry its child is terminated and the report says validation is incomplete. +These changes improve failure reporting and bound waits; TPU numerical stability +has not yet been demonstrated. They do not change the model, data or learning rates. ## Model and data @@ -45,7 +69,7 @@ change. max_tokens is rounded up to the next complete optimizer update. From a clean Cloud Shell checkout of this commit: ```bash -python3 baseline/nanogpt_one_head/gpt2small/cloudshell.py +python3 baseline/gpt2_small/scripts/cloudshell.py ``` This verifies the existing active queue/node, gracefully stops the old trainer, @@ -77,8 +101,8 @@ normal GPT-2 training quality, and before the long scientific experiment. From the updated, clean Cloud Shell checkout: ```bash -python3 baseline/nanogpt_one_head/gpt2small/reallocate_validation.py launch -python3 baseline/nanogpt_one_head/gpt2small/reallocate_validation.py status +python3 baseline/gpt2_small/scripts/reallocate_validation.py launch +python3 baseline/gpt2_small/scripts/reallocate_validation.py status ``` The launcher replaces only `ww-continuous8-24h-20261003-s1337` and its node. @@ -97,8 +121,8 @@ operations. The complete logs and outputs remain on the persistent disk on failu This is a fresh validation from initialization, not continuation of the failed run. Before each optimizer update, the runner checks loss and gradient norm. A failure -writes `nonfinite_diagnostics.json` with parameter names and nonfinite element -counts, then stops before applying the invalid update. The earlier nonfinite +writes `nonfinite_diagnostics.json` with the step, loss values and aggregate gradient +norm, then stops before applying the invalid update. The earlier nonfinite gradient's cause is still unconfirmed; these checks do not claim to fix it. Only if the short AdamW checks pass does MuonClip validation proceed, followed by the existing resume checks. No long run starts automatically. The service does @@ -108,8 +132,8 @@ delete its TPU: the allocation limit remains four hours unless stopped earlier. ### Allocate 48 hours instead ```bash -python3 baseline/nanogpt_one_head/gpt2small/reallocate_validation.py launch --hours 48 -python3 baseline/nanogpt_one_head/gpt2small/reallocate_validation.py status --hours 48 +python3 baseline/gpt2_small/scripts/reallocate_validation.py launch --hours 48 +python3 baseline/gpt2_small/scripts/reallocate_validation.py status --hours 48 ``` The queued-resource API exposes no update operation for extending the requested @@ -135,7 +159,7 @@ the long experiment. The allocation remains available until deletion or expiry. From Cloud Shell on the updated branch: ```bash -python3 baseline/nanogpt_one_head/gpt2small/capture_stall.py --stop +python3 baseline/gpt2_small/scripts/capture_stall.py --stop ``` This targets the 48-hour validation run above. It captures process CPU/memory, @@ -148,6 +172,9 @@ disk; a compact summary is printed for sharing. This command does not upload the Omit `--stop` for a read-only capture. A blocked update may be lost on stop; no new checkpoint is promised. Existing checkpoints, FineWeb and the TPU allocation remain. The allocation continues to incur compute usage until deletion/expiry. +Use `--show-last` to print the saved capture, native stack and current service state +without another profiler installation or stop. Both old and new package paths are +recognized when diagnosing the already allocated machine. The validator's `WAIT` lines only indicate a live process, not a completed update. Likewise `before_update: 2` confirms the gradient check before update 2, not its @@ -189,8 +216,8 @@ no thousands of full checkpoint copies are retained. ## Analysis and tests ```bash -python gpt2small/analyze.py /mnt/disks/rg-data/gpt2small/VALIDATION/muonclip -PYTHONPATH=src python -m pytest tests/test_gpt2_experiment.py -q +python baseline/gpt2_small/scripts/analyze.py /mnt/disks/rg-data/gpt2small/VALIDATION/muonclip +PYTHONPATH=baseline/gpt2_small/src:baseline/nanogpt_one_head/src python -m pytest baseline/gpt2_small/tests -q ``` Analysis creates CSVs and loss/perplexity/error-vs-token plots; mean/min raw and clipped diff --git a/baseline/nanogpt_one_head/configs/gpt2_small_cpu_smoke.yaml b/baseline/gpt2_small/configs/gpt2_small_cpu_smoke.yaml similarity index 100% rename from baseline/nanogpt_one_head/configs/gpt2_small_cpu_smoke.yaml rename to baseline/gpt2_small/configs/gpt2_small_cpu_smoke.yaml diff --git a/baseline/nanogpt_one_head/configs/gpt2_small_fineweb_adamw_baseline.yaml b/baseline/gpt2_small/configs/gpt2_small_fineweb_adamw_baseline.yaml similarity index 100% rename from baseline/nanogpt_one_head/configs/gpt2_small_fineweb_adamw_baseline.yaml rename to baseline/gpt2_small/configs/gpt2_small_fineweb_adamw_baseline.yaml diff --git a/baseline/nanogpt_one_head/configs/gpt2_small_fineweb_muonclip_baseline.yaml b/baseline/gpt2_small/configs/gpt2_small_fineweb_muonclip_baseline.yaml similarity index 100% rename from baseline/nanogpt_one_head/configs/gpt2_small_fineweb_muonclip_baseline.yaml rename to baseline/gpt2_small/configs/gpt2_small_fineweb_muonclip_baseline.yaml diff --git a/baseline/nanogpt_one_head/configs/gpt2_small_fineweb_muonclip_long_ww.yaml b/baseline/gpt2_small/configs/gpt2_small_fineweb_muonclip_long_ww.yaml similarity index 100% rename from baseline/nanogpt_one_head/configs/gpt2_small_fineweb_muonclip_long_ww.yaml rename to baseline/gpt2_small/configs/gpt2_small_fineweb_muonclip_long_ww.yaml diff --git a/baseline/gpt2_small/pyproject.toml b/baseline/gpt2_small/pyproject.toml new file mode 100644 index 00000000..18f17d3b --- /dev/null +++ b/baseline/gpt2_small/pyproject.toml @@ -0,0 +1,16 @@ +[build-system] +requires = ["setuptools>=68", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "rg-gpt2-small" +version = "0.1.0" +description = "GPT-2 Small: 12-layer, 12-head FineWeb TPU experiments" +requires-python = ">=3.10" +dependencies = ["rg-nanogpt-one-head==0.5.1"] + +[project.scripts] +rg-gpt2-train = "rg_gpt2_small.experiment:main" + +[tool.setuptools.packages.find] +where = ["src"] diff --git a/baseline/nanogpt_one_head/gpt2small/LOCAL_VALIDATION.json b/baseline/gpt2_small/scripts/LOCAL_VALIDATION.json similarity index 100% rename from baseline/nanogpt_one_head/gpt2small/LOCAL_VALIDATION.json rename to baseline/gpt2_small/scripts/LOCAL_VALIDATION.json diff --git a/baseline/nanogpt_one_head/gpt2small/analyze.py b/baseline/gpt2_small/scripts/analyze.py similarity index 100% rename from baseline/nanogpt_one_head/gpt2small/analyze.py rename to baseline/gpt2_small/scripts/analyze.py diff --git a/baseline/nanogpt_one_head/gpt2small/backup.py b/baseline/gpt2_small/scripts/backup.py similarity index 100% rename from baseline/nanogpt_one_head/gpt2small/backup.py rename to baseline/gpt2_small/scripts/backup.py diff --git a/baseline/nanogpt_one_head/gpt2small/capture_stall.py b/baseline/gpt2_small/scripts/capture_stall.py similarity index 81% rename from baseline/nanogpt_one_head/gpt2small/capture_stall.py rename to baseline/gpt2_small/scripts/capture_stall.py index f76c53e9..9917df2d 100644 --- a/baseline/nanogpt_one_head/gpt2small/capture_stall.py +++ b/baseline/gpt2_small/scripts/capture_stall.py @@ -30,8 +30,9 @@ def service_info(): def verify_target(info, root): - expected = str(root / 'repo/baseline/nanogpt_one_head/gpt2small/replacement_worker.sh') - if expected not in info.get('ExecStart', ''): + expected = [str(root / ('repo/baseline/' + path)) for path in ( + 'gpt2_small/scripts/replacement_worker.sh', 'nanogpt_one_head/gpt2small/replacement_worker.sh')] + if not any(path in info.get('ExecStart', '') for path in expected): raise RuntimeError('Service belongs to another run; refusing to stop it.') @@ -42,7 +43,7 @@ def training_pids(root, proc=Path('/proc')): continue try: args = (folder / 'cmdline').read_bytes().decode().strip('\0').split('\0') - if 'rg_nanogpt_one_head.gpt2_experiment' not in args or '--output' not in args: + if not {'rg_gpt2_small.experiment', 'rg_nanogpt_one_head.gpt2_experiment'}.intersection(args) or '--output' not in args: continue output = Path(args[args.index('--output') + 1]) if output.parent == root and output.name in ('adamw', 'muonclip'): @@ -159,17 +160,46 @@ def on_tpu(stop=False): print('Read-only capture complete; validation was not stopped.', flush=True) +def show_last(): + """Read the saved capture without touching the trainer, allocation or data.""" + captures = sorted((ROOT / 'diagnostics').glob('stall-*')) + if not captures: + raise RuntimeError('No saved stall capture found.') + dest = captures[-1] + print('Saved diagnostics:', dest, flush=True) + for path in [dest/'capture.json', dest/'service-after.txt', + dest/'adamw_logs_xla_compile_metrics_after_step_2.txt', + *sorted(dest.glob('pid-*-native-stack.txt'))]: + print('\n===', path.name, '===', flush=True) + if not path.is_file(): + print('Not captured.', flush=True) + continue + lines = path.read_text(errors='replace').splitlines() + print('\n'.join(lines[:120]), flush=True) + if 'native-stack' in path.name: + keywords = ('Compile', 'compile', 'Execute', 'execute', 'Hlo', 'hlo', 'llvm', 'xla::') + matches = [line for line in lines[120:] if any(word in line for word in keywords)] + print('\nAdditional compiler/execution frames:\n' + '\n'.join(matches[:100]), flush=True) + print('\n=== Current validation service ===', flush=True) + print(json.dumps(service_info(), indent=2), flush=True) + + def main(): parser = argparse.ArgumentParser(description=__doc__) parser.add_argument('--on-tpu', action='store_true', help=argparse.SUPPRESS) parser.add_argument('--stop', action='store_true', help='Stop only the validation service after capture') + parser.add_argument('--show-last', action='store_true', help='Read the saved diagnostic; do not capture or stop') args = parser.parse_args() + if args.stop and args.show_last: + parser.error('--stop and --show-last cannot be combined') if args.on_tpu: - on_tpu(args.stop) + show_last() if args.show_last else on_tpu(args.stop) return 0 remote = ['sudo', 'python3', '-c', Path(__file__).read_text(), '--on-tpu'] if args.stop: remote.append('--stop') + if args.show_last: + remote.append('--show-last') # Do not alter the running checkout or ask it to fetch/checkout another commit. return subprocess.run(['gcloud', 'compute', 'tpus', 'tpu-vm', 'ssh', RUN + '-node', '--project=' + PROJECT, '--zone=' + ZONE, '--worker=0', '--command=' + shlex.join(remote)]).returncode diff --git a/baseline/nanogpt_one_head/gpt2small/cloudshell.py b/baseline/gpt2_small/scripts/cloudshell.py similarity index 100% rename from baseline/nanogpt_one_head/gpt2small/cloudshell.py rename to baseline/gpt2_small/scripts/cloudshell.py diff --git a/baseline/nanogpt_one_head/gpt2small/diagnose_adamw.sh b/baseline/gpt2_small/scripts/diagnose_adamw.sh similarity index 69% rename from baseline/nanogpt_one_head/gpt2small/diagnose_adamw.sh rename to baseline/gpt2_small/scripts/diagnose_adamw.sh index 4ae91cbd..923bdb9f 100644 --- a/baseline/nanogpt_one_head/gpt2small/diagnose_adamw.sh +++ b/baseline/gpt2_small/scripts/diagnose_adamw.sh @@ -4,9 +4,9 @@ set -euo pipefail old=$1 deadline=$2 base=$(cd "$(dirname "$0")/.." && pwd) -source "$base/gpt2small/tpu_environment.sh" +source "$base/scripts/tpu_environment.sh" python=/mnt/disks/rg-data/continuous8/venv/bin/python -export PYTHONPATH="$base/src" +export PYTHONPATH="$base/src:$base/../nanogpt_one_head/src" "$python" -c 'import sys,time; assert float(sys.argv[1])>time.time()+60,"Too little allocation time; diagnostic not started"' "$deadline" root="$old/diagnostic-$(date -u +%Y%m%d-%H%M%S)" mkdir "$root" @@ -23,8 +23,8 @@ backup() { rc=$? trap - EXIT sync - "$python" "$base/gpt2small/backup.py" "$root" || exit 1 + "$python" "$base/scripts/backup.py" "$root" || exit 1 exit "$rc" } trap backup EXIT -"$python" -u -m rg_nanogpt_one_head.gpt2_experiment --config "$root/config.yaml" --data-root /mnt/disks/rg-data/continuous8/data --output "$root/adamw" --device tpu --stop-after 4 --deadline-unix "$deadline" 2>&1 | tee "$root/diagnostic.log" +"$python" -u -m rg_gpt2_small.experiment --config "$root/config.yaml" --data-root /mnt/disks/rg-data/continuous8/data --output "$root/adamw" --device tpu --stop-after 4 --deadline-unix "$deadline" 2>&1 | tee "$root/diagnostic.log" diff --git a/baseline/nanogpt_one_head/gpt2small/on_tpu.sh b/baseline/gpt2_small/scripts/on_tpu.sh similarity index 75% rename from baseline/nanogpt_one_head/gpt2small/on_tpu.sh rename to baseline/gpt2_small/scripts/on_tpu.sh index e851c1c6..813442f9 100644 --- a/baseline/nanogpt_one_head/gpt2small/on_tpu.sh +++ b/baseline/gpt2_small/scripts/on_tpu.sh @@ -14,9 +14,9 @@ git clone --no-checkout https://github.com/CalculatedContent/rg_optimizers.git " git -C "$repo" checkout --detach "$commit" test "$(git -C "$repo" rev-parse HEAD)" = "$commit" git -C "$repo" rev-parse HEAD -export PYTHONPATH="$repo/baseline/nanogpt_one_head/src" -source "$repo/baseline/nanogpt_one_head/gpt2small/tpu_environment.sh" -cd "$repo/baseline/nanogpt_one_head" +export PYTHONPATH="$repo/baseline/gpt2_small/src:$repo/baseline/nanogpt_one_head/src" +source "$repo/baseline/gpt2_small/scripts/tpu_environment.sh" +cd "$repo/baseline/gpt2_small" "$python" -c 'import torch, torch_xla, weightwatcher, yaml; print("Installed dependencies loaded")' deadline=$("$python" -c 'import json,sys; print(json.load(open(sys.argv[1]))["deadline_unix"])' "$root/old_allocation.json") bucket=gs://tpu-builders-504820-ww-continuous8/gpt2small @@ -25,12 +25,12 @@ backup() { trap - EXIT sync # Keep backups separate from prior pilot archives. No remote deletion. - if ! "$python" "$repo/baseline/nanogpt_one_head/gpt2small/backup.py" "$root"; then + if ! "$python" "$repo/baseline/gpt2_small/scripts/backup.py" "$root"; then echo "Cloud backup FAILED; all outputs remain on persistent disk: $root" >&2 exit 1 fi exit "$rc" } trap backup EXIT -"$python" -u gpt2small/validate.py --root "$root" --data "$shared/data" --deadline "$deadline" 2>&1 | tee "$root/validation.log" +"$python" -u scripts/validate.py --root "$root" --data "$shared/data" --deadline "$deadline" 2>&1 | tee "$root/validation.log" echo 'Validation finished. No long experiment launched.' diff --git a/baseline/nanogpt_one_head/gpt2small/prepare_existing.py b/baseline/gpt2_small/scripts/prepare_existing.py similarity index 100% rename from baseline/nanogpt_one_head/gpt2small/prepare_existing.py rename to baseline/gpt2_small/scripts/prepare_existing.py diff --git a/baseline/nanogpt_one_head/gpt2small/reallocate_validation.py b/baseline/gpt2_small/scripts/reallocate_validation.py similarity index 97% rename from baseline/nanogpt_one_head/gpt2small/reallocate_validation.py rename to baseline/gpt2_small/scripts/reallocate_validation.py index be9cfee5..dc68d212 100644 --- a/baseline/nanogpt_one_head/gpt2small/reallocate_validation.py +++ b/baseline/gpt2_small/scripts/reallocate_validation.py @@ -129,7 +129,7 @@ def launch(): stop = (f'if systemctl cat {OLD_SERVICE} >/dev/null 2>&1; then\n' f' systemctl stop {OLD_SERVICE}\nfi\n') if OLD_SERVICE else '' command = "sudo bash -se <<'CHECK'\n" + stop + """ -if pgrep -af '[p]ython.*(gpt2_experiment|gpt2small/validate.py|continuous_run)'; then +if pgrep -af '[p]ython.*(gpt2_experiment|rg_gpt2_small.experiment|gpt2small/validate.py|gpt2_small/scripts/validate.py|continuous_run)'; then echo 'A trainer is still running; replacement aborted.' >&2 exit 1 fi @@ -166,7 +166,7 @@ def launch(): '--labels=experiment=gpt2-validation', '--quiet', '--async') save(phase='submitted; validation starts automatically when capacity is allocated') print('Node: ' + NODE + '\nLog: ' + ROOT + '/run.log', flush=True) - print(f'Check: python3 baseline/nanogpt_one_head/gpt2small/reallocate_validation.py status --hours {HOURS}', flush=True) + print(f'Check: python3 baseline/gpt2_small/scripts/reallocate_validation.py status --hours {HOURS}', flush=True) def main(): diff --git a/baseline/nanogpt_one_head/gpt2small/replacement_startup.sh b/baseline/gpt2_small/scripts/replacement_startup.sh similarity index 95% rename from baseline/nanogpt_one_head/gpt2small/replacement_startup.sh rename to baseline/gpt2_small/scripts/replacement_startup.sh index ee72749a..c365d16c 100644 --- a/baseline/nanogpt_one_head/gpt2small/replacement_startup.sh +++ b/baseline/gpt2_small/scripts/replacement_startup.sh @@ -47,7 +47,7 @@ After=network-online.target Wants=network-online.target [Service] Type=simple -ExecStart=/bin/bash $ROOT/repo/baseline/nanogpt_one_head/gpt2small/replacement_worker.sh $ROOT $DEADLINE +ExecStart=/bin/bash $ROOT/repo/baseline/gpt2_small/scripts/replacement_worker.sh $ROOT $DEADLINE Restart=no KillSignal=SIGTERM TimeoutStopSec=600 diff --git a/baseline/nanogpt_one_head/gpt2small/replacement_worker.sh b/baseline/gpt2_small/scripts/replacement_worker.sh similarity index 86% rename from baseline/nanogpt_one_head/gpt2small/replacement_worker.sh rename to baseline/gpt2_small/scripts/replacement_worker.sh index c42bb7b0..cb9a1b08 100644 --- a/baseline/nanogpt_one_head/gpt2small/replacement_worker.sh +++ b/baseline/gpt2_small/scripts/replacement_worker.sh @@ -4,20 +4,20 @@ root=$1 deadline=$2 base=$(cd "$(dirname "$0")/.." && pwd) python=/mnt/disks/rg-data/continuous8/venv/bin/python -export PYTHONPATH="$base/src" +export PYTHONPATH="$base/src:$base/../nanogpt_one_head/src" finish() { rc=$? trap - EXIT echo "Validation worker exit code: $rc" sync - if ! "$python" "$base/gpt2small/backup.py" "$root"; then + if ! "$python" "$base/scripts/backup.py" "$root"; then echo "Cloud backup FAILED; files remain on the preserved disk: $root" >&2 exit 1 fi exit "$rc" } trap finish EXIT -source "$base/gpt2small/tpu_environment.sh" +source "$base/scripts/tpu_environment.sh" mountpoint -q /mnt/disks/rg-data test -x "$python" # Reuse installed dependencies and the verified corpus. Do not reinstall or redownload. @@ -33,6 +33,6 @@ receipt=sink.file(probe,probe.name) (root/'backup_probe_receipt.json').write_text(json.dumps(receipt,indent=2)) print('Cloud upload permission and checksum verification passed before training.',flush=True) PY -"$python" -u "$base/gpt2small/validate.py" --root "$root" \ +"$python" -u "$base/scripts/validate.py" --root "$root" \ --data /mnt/disks/rg-data/continuous8/data --deadline "$deadline" 2>&1 | tee "$root/validation.log" echo 'Short validation completed. No long experiment launched.' diff --git a/baseline/nanogpt_one_head/gpt2small/retry_existing.sh b/baseline/gpt2_small/scripts/retry_existing.sh similarity index 70% rename from baseline/nanogpt_one_head/gpt2small/retry_existing.sh rename to baseline/gpt2_small/scripts/retry_existing.sh index 71ed2993..3ab25b96 100644 --- a/baseline/nanogpt_one_head/gpt2small/retry_existing.sh +++ b/baseline/gpt2_small/scripts/retry_existing.sh @@ -3,7 +3,7 @@ set -euo pipefail old=$1 base=$(cd "$(dirname "$0")/.." && pwd) -source "$base/gpt2small/tpu_environment.sh" +source "$base/scripts/tpu_environment.sh" mountpoint -q /mnt/disks/rg-data test -f "$old/old_allocation.json" test -f /mnt/disks/rg-data/continuous8/data/meta.json @@ -12,15 +12,15 @@ sudo mkdir "$root" sudo chown "$(id -u):$(id -g)" "$root" cp "$old/old_allocation.json" "$root/old_allocation.json" python=/mnt/disks/rg-data/continuous8/venv/bin/python -export PYTHONPATH="$base/src" +export PYTHONPATH="$base/src:$base/../nanogpt_one_head/src" git -C "$base" rev-parse HEAD | tee "$root/commit.txt" deadline=$("$python" -c 'import json,sys; print(json.load(open(sys.argv[1]))["deadline_unix"])' "$root/old_allocation.json") backup() { rc=$? trap - EXIT sync - "$python" "$base/gpt2small/backup.py" "$root" || exit 1 + "$python" "$base/scripts/backup.py" "$root" || exit 1 exit "$rc" } trap backup EXIT -"$python" -u "$base/gpt2small/validate.py" --root "$root" --data /mnt/disks/rg-data/continuous8/data --deadline "$deadline" 2>&1 | tee "$root/validation.log" +"$python" -u "$base/scripts/validate.py" --root "$root" --data /mnt/disks/rg-data/continuous8/data --deadline "$deadline" 2>&1 | tee "$root/validation.log" diff --git a/baseline/nanogpt_one_head/gpt2small/tpu_environment.sh b/baseline/gpt2_small/scripts/tpu_environment.sh similarity index 100% rename from baseline/nanogpt_one_head/gpt2small/tpu_environment.sh rename to baseline/gpt2_small/scripts/tpu_environment.sh diff --git a/baseline/nanogpt_one_head/gpt2small/validate.py b/baseline/gpt2_small/scripts/validate.py similarity index 88% rename from baseline/nanogpt_one_head/gpt2small/validate.py rename to baseline/gpt2_small/scripts/validate.py index 22946dbb..8af55150 100644 --- a/baseline/nanogpt_one_head/gpt2small/validate.py +++ b/baseline/gpt2_small/scripts/validate.py @@ -9,11 +9,15 @@ import time import threading import yaml +from validation_watchdog import wait_for_child p=argparse.ArgumentParser() p.add_argument('--root',required=True); p.add_argument('--data',required=True) p.add_argument('--deadline',type=float,required=True) +p.add_argument('--phase-timeout-seconds',type=float,default=1800) a=p.parse_args(); root=Path(a.root); base=Path(__file__).resolve().parents[1] +if not 0 < a.phase_timeout_seconds <= 3600: + p.error('--phase-timeout-seconds must be in (0, 3600]') root.mkdir(parents=True,exist_ok=True) for folder in ('logs','metrics','ww_metrics','checkpoints','summaries','configs'): (root/folder).mkdir(exist_ok=True) configs={} @@ -41,7 +45,7 @@ def history(run): # Reserve five minutes for final checkpoint and backup; never start after the cutoff. if time.time()>a.deadline-300: raise RuntimeError('Insufficient allocation time for next phase; no new allocation requested') run=root/opt; old=history(run) - cmd=[sys.executable,'-u','-m','rg_nanogpt_one_head.gpt2_experiment', + cmd=[sys.executable,'-u','-m','rg_gpt2_small.experiment', '--config',str(configs[opt]),'--data-root',a.data,'--output',str(run), '--device','tpu','--stop-after',str(stop),'--deadline-unix',str(a.deadline-300)] if stop>4: cmd.append('--resume') @@ -52,12 +56,11 @@ def stream(): for line in child.stdout: log.write(line); log.flush(); print(line,end='',flush=True) reader=threading.Thread(target=stream,daemon=True); reader.start() - while True: - try: - rc=child.wait(timeout=30); break - except subprocess.TimeoutExpired: - print(f'WAIT {opt} through step {stop}: process {child.pid} still running',flush=True) - reader.join() + try: + rc=wait_for_child(child,f'{opt} through step {stop}', + min(time.time()+a.phase_timeout_seconds,a.deadline-300)) + finally: + reader.join(timeout=30) if rc: raise subprocess.CalledProcessError(rc,cmd) state=json.loads((run/'status.json').read_text()) if state['step']!=stop: raise RuntimeError(f'{opt} stopped before requested step {stop}') @@ -81,7 +84,7 @@ def stream(): report['phases'].append({'optimizer':opt,'step':stop,'metrics':latest,'ww':ww, 'resume_history_unchanged':bool(old)}) persist(); print(json.dumps(report['phases'][-1]),flush=True) - subprocess.run([sys.executable,str(base/'gpt2small/analyze.py'),str(run)],check=True) + subprocess.run([sys.executable,str(base/'scripts/analyze.py'),str(run)],check=True) report['status']='short_validation_completed' report['interpretation']='25 updates demonstrate functionality and initial direction only; not reproduction of a NanoGPT speedrun benchmark or proof of long-run stability.' except Exception as exc: diff --git a/baseline/gpt2_small/scripts/validation_watchdog.py b/baseline/gpt2_small/scripts/validation_watchdog.py new file mode 100644 index 00000000..42c39406 --- /dev/null +++ b/baseline/gpt2_small/scripts/validation_watchdog.py @@ -0,0 +1,21 @@ +"""Bound validation phases even if an XLA operation never returns.""" +import subprocess +import time + + +def wait_for_child(child, label, deadline): + while True: + remaining = deadline - time.time() + if remaining <= 0: + child.terminate() + try: + child.wait(timeout=20) + except subprocess.TimeoutExpired: + child.kill() + child.wait(timeout=20) + raise RuntimeError(f'{label}: phase time limit exceeded; child stopped. Validation incomplete.') + try: + return child.wait(timeout=min(30, remaining)) + except subprocess.TimeoutExpired: + print(f'WAIT {label}: process {child.pid} alive; completion not confirmed; ' + f'{max(0, deadline-time.time()):.0f}s to phase cutoff', flush=True) diff --git a/baseline/gpt2_small/src/rg_gpt2_small/__init__.py b/baseline/gpt2_small/src/rg_gpt2_small/__init__.py new file mode 100644 index 00000000..1b311d27 --- /dev/null +++ b/baseline/gpt2_small/src/rg_gpt2_small/__init__.py @@ -0,0 +1 @@ +"""GPT-2 Small (12 layers, 12 heads) validation and training.""" diff --git a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/gpt2_experiment.py b/baseline/gpt2_small/src/rg_gpt2_small/experiment.py similarity index 91% rename from baseline/nanogpt_one_head/src/rg_nanogpt_one_head/gpt2_experiment.py rename to baseline/gpt2_small/src/rg_gpt2_small/experiment.py index dcdca0c8..afa405cc 100644 --- a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/gpt2_experiment.py +++ b/baseline/gpt2_small/src/rg_gpt2_small/experiment.py @@ -3,6 +3,7 @@ import argparse import copy import fcntl +import faulthandler import hashlib import json import math @@ -14,12 +15,12 @@ import numpy as np import torch import yaml -from .model import GPT, GPTConfig, transformer_matrix_items -from . import runtime as rt, tpu_spmd as spmd -from .data import load_memmaps -from .muonclip import install_muonclip_extension -from . import optimizers -from .spectral import WeightMatrixHolder, _attach_matrix_metadata +from rg_nanogpt_one_head.model import GPT, GPTConfig, transformer_matrix_items +from rg_nanogpt_one_head import runtime as rt, tpu_spmd as spmd +from rg_nanogpt_one_head.data import load_memmaps +from rg_nanogpt_one_head.muonclip import install_muonclip_extension +from rg_nanogpt_one_head import optimizers +from rg_nanogpt_one_head.spectral import WeightMatrixHolder, _attach_matrix_metadata def atomic_json(path, value): @@ -130,23 +131,23 @@ def number(key): def require_finite_update(model, norm, losses, output, step): """Diagnostic pre-update barrier for validation. Never apply known-invalid gradients.""" + print(json.dumps({'checking_before_update':step}),flush=True) + # Materialize the pending XLA graph once before any individual host reads. + rt.synchronize(norm.device) value=float(norm.detach().cpu()) loss_values=[float(x.detach().cpu()) for x in losses] if math.isfinite(value) and all(math.isfinite(x) for x in loss_values): print(json.dumps({'before_update':step,'microbatch_losses':loss_values,'gradient_norm':value}),flush=True) return - gradients=[] - for name,p in model.named_parameters(): - if p.grad is None: continue - grad=p.grad.detach().float().cpu() - finite=torch.isfinite(grad) - gradients.append({'parameter':name,'shape':list(grad.shape), - 'nonfinite_elements':int((~finite).sum()), - 'max_finite_abs':float(grad[finite].abs().max()) if finite.any() else None}) - report={'before_update':step,'gradient_norm':value,'microbatch_losses':loss_values,'gradients':gradients} + # Persist the known failure immediately. Copying every gradient to the host + # previously stalled here and hid the nonfinite result for over two hours. + report={'status':'invalid_before_update','before_update':step, + 'gradient_norm':value,'microbatch_losses':loss_values, + 'gradient_details_collected':False, + 'reason':'Nonfinite loss or aggregate gradient norm; invalid optimizer update was not applied.'} atomic_json(Path(output)/'nonfinite_diagnostics.json',report) - bad=[r['parameter'] for r in gradients if r['nonfinite_elements']] - raise RuntimeError(f'Invalid loss/gradient BEFORE update {step}; bad matrices={bad}; see nonfinite_diagnostics.json') + print(json.dumps(report),flush=True) + raise RuntimeError(f'Invalid loss/gradient BEFORE update {step}; see nonfinite_diagnostics.json') def save_checkpoint(root, payload, keep=3, milestone=False): @@ -178,7 +179,14 @@ def train(cfg, data_root, output, *, device='cpu', resume=False, stop_after=None output = Path(output); output.mkdir(parents=True, exist_ok=True) with (output / 'writer.lock').open('a') as lock: fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB) - return _train(cfg, data_root, output, device, resume, stop_after, deadline) + diagnostic = cfg.get('validation_gradient_checks', False) + if diagnostic: + faulthandler.dump_traceback_later(300, repeat=True) + try: + return _train(cfg, data_root, output, device, resume, stop_after, deadline) + finally: + if diagnostic: + faulthandler.cancel_dump_traceback_later() def _train(cfg, data_root, output, device, resume, stop_after, deadline): diff --git a/baseline/nanogpt_one_head/tests/test_gpt2_experiment.py b/baseline/gpt2_small/tests/test_gpt2_experiment.py similarity index 80% rename from baseline/nanogpt_one_head/tests/test_gpt2_experiment.py rename to baseline/gpt2_small/tests/test_gpt2_experiment.py index 28784bfd..fc1b4295 100644 --- a/baseline/nanogpt_one_head/tests/test_gpt2_experiment.py +++ b/baseline/gpt2_small/tests/test_gpt2_experiment.py @@ -10,7 +10,7 @@ import yaml from rg_nanogpt_one_head.model import GPT, GPTConfig, transformer_matrix_items from rg_nanogpt_one_head.data import write_token_splits -from rg_nanogpt_one_head import gpt2_experiment as g +from rg_gpt2_small import experiment as g BASE=Path(__file__).resolve().parents[1] def config(name='cpu_smoke'): @@ -100,13 +100,36 @@ def test_interrupted_record_transaction_recovers(tmp_path): assert len(list((out/'metrics').glob('*.json')))==2 -def test_invalid_gradient_diagnostic_precedes_update(tmp_path): +def test_invalid_gradient_diagnostic_precedes_update(tmp_path, monkeypatch): model=GPT(GPTConfig(**config()['model'])) for p in model.parameters(): p.grad=torch.ones_like(p) name,p=next(iter(model.named_parameters())); p.grad.view(-1)[0]=float('nan') before=p.detach().clone() + def no_gradient_walk(*args, **kwargs): + raise AssertionError('Failure diagnostics must not walk/copy every gradient') + monkeypatch.setattr(model,'named_parameters',no_gradient_walk) with pytest.raises(RuntimeError,match='BEFORE update 1'): g.require_finite_update(model,torch.tensor(float('nan')),[torch.tensor(11.)],tmp_path,1) report=json.loads((tmp_path/'nonfinite_diagnostics.json').read_text()) - assert [r['parameter'] for r in report['gradients'] if r['nonfinite_elements']]==[name] + assert report['status']=='invalid_before_update' and report['gradient_details_collected'] is False + assert np.isnan(report['gradient_norm']) and report['microbatch_losses']==[11.] assert torch.equal(before,p.detach()) + + +def test_validation_synchronizes_before_scalar_transfer(tmp_path,monkeypatch): + calls=[] + class Scalar: + device=torch.device('cpu') + def detach(self): return self + def cpu(self): + assert calls==['sync'] + return torch.tensor(1.) + monkeypatch.setattr(g.rt,'synchronize',lambda device:calls.append('sync')) + g.require_finite_update(None,Scalar(),[Scalar()],tmp_path,1) + assert not (tmp_path/'nonfinite_diagnostics.json').exists() + + +def test_nonfinite_loss_fails_even_with_finite_norm(tmp_path): + with pytest.raises(RuntimeError,match='BEFORE update 3'): + g.require_finite_update(None,torch.tensor(1.),[torch.tensor(float('inf'))],tmp_path,3) + assert json.loads((tmp_path/'nonfinite_diagnostics.json').read_text())['before_update']==3 diff --git a/baseline/nanogpt_one_head/tests/test_gpt2_reallocation.py b/baseline/gpt2_small/tests/test_gpt2_reallocation.py similarity index 95% rename from baseline/nanogpt_one_head/tests/test_gpt2_reallocation.py rename to baseline/gpt2_small/tests/test_gpt2_reallocation.py index 52bd8ecd..cef913e9 100644 --- a/baseline/nanogpt_one_head/tests/test_gpt2_reallocation.py +++ b/baseline/gpt2_small/tests/test_gpt2_reallocation.py @@ -7,7 +7,7 @@ def load(): - path = Path(__file__).resolve().parents[1] / 'gpt2small/reallocate_validation.py' + path = Path(__file__).resolve().parents[1] / 'scripts/reallocate_validation.py' spec = importlib.util.spec_from_file_location('gpt2_reallocation', path) module = importlib.util.module_from_spec(spec); spec.loader.exec_module(module) return module @@ -23,7 +23,7 @@ def test_startup_mounts_only_existing_disk_and_never_restarts(hours): assert 'STARTED_ONCE' in script and 'replacement_worker.sh' in script worker = Path(m.__file__).with_name('replacement_worker.sh').read_text() subprocess.run(['bash', '-n'], input=worker, text=True, check=True) - assert worker.index('receipt=sink.file') < worker.index('gpt2small/validate.py') + assert worker.index('receipt=sink.file') < worker.index('scripts/validate.py') assert 'prepare_tpu_data' not in worker and '--allow-long-run' not in worker diff --git a/baseline/nanogpt_one_head/tests/test_gpt2_stall_capture.py b/baseline/gpt2_small/tests/test_gpt2_stall_capture.py similarity index 84% rename from baseline/nanogpt_one_head/tests/test_gpt2_stall_capture.py rename to baseline/gpt2_small/tests/test_gpt2_stall_capture.py index e99ad0a7..dc0a8458 100644 --- a/baseline/nanogpt_one_head/tests/test_gpt2_stall_capture.py +++ b/baseline/gpt2_small/tests/test_gpt2_stall_capture.py @@ -5,22 +5,22 @@ spec = importlib.util.spec_from_file_location( - 'stall_capture', Path(__file__).parents[1] / 'gpt2small/capture_stall.py') + 'stall_capture', Path(__file__).parents[1] / 'scripts/capture_stall.py') capture = importlib.util.module_from_spec(spec) spec.loader.exec_module(capture) def service(root, active='active', pid='6800'): - return {'ExecStart': f'/bin/bash {root}/repo/baseline/nanogpt_one_head/gpt2small/replacement_worker.sh {root}', + return {'ExecStart': f'/bin/bash {root}/repo/baseline/gpt2_small/scripts/replacement_worker.sh {root}', 'ActiveState': active, 'MainPID': pid} def test_process_selection_excludes_other_training_runs(tmp_path): root = Path('/mnt/disks/rg-data/gpt2small/current') for pid, args in { - 11: ['python', '-m', 'rg_nanogpt_one_head.gpt2_experiment', '--output', str(root / 'adamw')], - 12: ['python', '-m', 'rg_nanogpt_one_head.gpt2_experiment', '--output', str(root / 'muonclip')], - 13: ['python', '-m', 'rg_nanogpt_one_head.gpt2_experiment', '--output', str(root.parent / 'other/adamw')], + 11: ['python', '-m', 'rg_gpt2_small.experiment', '--output', str(root / 'adamw')], + 12: ['python', '-m', 'rg_gpt2_small.experiment', '--output', str(root / 'muonclip')], + 13: ['python', '-m', 'rg_gpt2_small.experiment', '--output', str(root.parent / 'other/adamw')], 14: ['python', 'validate.py', '--root', str(root)], }.items(): folder = tmp_path / str(pid) diff --git a/baseline/gpt2_small/tests/test_validation_watchdog.py b/baseline/gpt2_small/tests/test_validation_watchdog.py new file mode 100644 index 00000000..102bf08c --- /dev/null +++ b/baseline/gpt2_small/tests/test_validation_watchdog.py @@ -0,0 +1,37 @@ +import importlib.util +from pathlib import Path +import subprocess + +import pytest + +spec=importlib.util.spec_from_file_location('watchdog',Path(__file__).parents[1]/'scripts/validation_watchdog.py') +watchdog=importlib.util.module_from_spec(spec) +spec.loader.exec_module(watchdog) + + +class Child: + pid=123 + def __init__(self,ignores_term=False): + self.ignores_term=ignores_term; self.calls=[] + def terminate(self): self.calls.append('terminate') + def kill(self): self.calls.append('kill') + def wait(self,timeout): + if self.ignores_term and self.calls==['terminate']: + raise subprocess.TimeoutExpired('test-child',timeout) + return 0 + + +@pytest.mark.parametrize('ignores_term',[False,True]) +def test_phase_cutoff_terminates_child(ignores_term,monkeypatch): + monkeypatch.setattr(watchdog.time,'time',lambda:100) + child=Child(ignores_term) + with pytest.raises(RuntimeError,match='Validation incomplete'): + watchdog.wait_for_child(child,'adamw through step 4',99) + assert child.calls==(['terminate','kill'] if ignores_term else ['terminate']) + + +def test_successful_child_is_not_stopped(monkeypatch): + monkeypatch.setattr(watchdog.time,'time',lambda:100) + child=Child() + assert watchdog.wait_for_child(child,'adamw',200)==0 + assert child.calls==[] diff --git a/baseline/nanogpt_one_head/README.md b/baseline/nanogpt_one_head/README.md index 369fca48..5c17bc9a 100644 --- a/baseline/nanogpt_one_head/README.md +++ b/baseline/nanogpt_one_head/README.md @@ -386,6 +386,6 @@ runner are disabled; the allocation cap includes setup and tokenization. ## GPT-2 Small validation -See [gpt2small/README.md](gpt2small/README.md) for the isolated context-1024 +See [GPT-2 Small](../gpt2_small/README.md) for the isolated context-1024 AdamW/MuonClip validation, corpus reuse, bounded checkpoints and resume tests. The long experiment requires a separate explicit launch. From 19e2bb081773c1ca1c13879f8d88db23c115f3fa Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Sat, 3 Oct 2026 21:44:42 -0700 Subject: [PATCH 18/36] Log TPU port failures and add a bounded four-update numerical diagnostic --- baseline/gpt2_small/README.md | 21 +++ baseline/gpt2_small/TPU_PORT_BUGS.md | 66 +++++++++ baseline/gpt2_small/scripts/diagnose_adamw.sh | 12 +- baseline/gpt2_small/scripts/retry_adamw.py | 127 ++++++++++++++++++ .../gpt2_small/scripts/run_adamw_probe.py | 54 ++++++++ .../src/rg_gpt2_small/experiment.py | 38 +++++- .../src/rg_gpt2_small/port_debug.py | 110 +++++++++++++++ baseline/gpt2_small/tests/test_port_debug.py | 101 ++++++++++++++ 8 files changed, 523 insertions(+), 6 deletions(-) create mode 100644 baseline/gpt2_small/TPU_PORT_BUGS.md create mode 100644 baseline/gpt2_small/scripts/retry_adamw.py create mode 100644 baseline/gpt2_small/scripts/run_adamw_probe.py create mode 100644 baseline/gpt2_small/src/rg_gpt2_small/port_debug.py create mode 100644 baseline/gpt2_small/tests/test_port_debug.py diff --git a/baseline/gpt2_small/README.md b/baseline/gpt2_small/README.md index 1aea7d57..f82cda24 100644 --- a/baseline/gpt2_small/README.md +++ b/baseline/gpt2_small/README.md @@ -13,6 +13,27 @@ development, from the repository root, install both with ## Current validation failure +Observed failures and their attribution are tracked in [TPU_PORT_BUGS.md](TPU_PORT_BUGS.md). +Diagnosing bugs in the PyTorch/TPU port is an explicit experiment objective. + +From an updated Cloud Shell checkout, use the existing 48-hour TPU for a bounded +four-update AdamW diagnostic (no allocation or data download): + +```bash +python3 baseline/gpt2_small/scripts/retry_adamw.py start +python3 baseline/gpt2_small/scripts/retry_adamw.py status +``` + +It runs independently of Cloud Shell under a new service and fresh output directory, +pins the checked-in source, and refuses concurrent trainers. Training is limited to +20 minutes with up to 10 more for backup. It records per-layer gradients, clipped +gradients, weights and optimizer moments as finite flags/extrema; only small reduced +tables move to CPU. Full checkpoints are saved at each completed update. Source and +runtime versions, input offsets, XLA metrics, stage times and failures accompany the +run. `TPU_PORT_FAILURE.json` records observed failure without assuming upstream fault; +`PROBE_STATUS.json` reports completion or failure. Existing cloud upload verification +runs on exit. This is diagnostic instrumentation, not a throughput measurement. + The saved 2026-10-04 04:19 UTC stack from commit `4631a3d` is inside the nonfinite-gradient diagnostic's per-parameter CPU copy. That branch is reached only after detecting a nonfinite loss or aggregate gradient norm. The latest checkpoint diff --git a/baseline/gpt2_small/TPU_PORT_BUGS.md b/baseline/gpt2_small/TPU_PORT_BUGS.md new file mode 100644 index 00000000..5d468f34 --- /dev/null +++ b/baseline/gpt2_small/TPU_PORT_BUGS.md @@ -0,0 +1,66 @@ +# TPU port bug log + +Finding and isolating correctness/performance defects in the PyTorch-to-TPU port is +an explicit objective of this project. Preserve failures as evidence. A passing +CPU smoke test does not establish TPU correctness. Do not silently lower learning +rates, change data, or skip nonfinite checks to make a run appear successful. + +## 2026-10-04: GPT-2 AdamW nonfinite result, followed by stalled failure reporting + +Status: numerical cause open; blocking diagnostic implementation replaced. +Upstream attribution: **unconfirmed**. No upstream issue has been submitted. + +- Run: `ww-gpt2-validation-48h-20261004-s1337`, original commit `4631a3d`. +- Machine: one v5litepod-8, eight chips, SPMD; project `tpu-builders-504820`, + zone `us-west4-a`. Reported runtime: Python 3.10.12, PyTorch 2.6.0+cpu, + PyTorch/XLA 2.6.0. The library's `+cpu` build string does not identify where + the model computations executed; the run explicitly selected XLA/TPU. +- Model: GPT-2 Small, 124,439,808 parameters; 12 layers, 12 heads, width 768, + context 1024, tied embeddings. FineWeb reused from the persistent disk. +- Before updates 1 and 2, reported losses were finite (approximately 11.01 and + 10.24) and aggregate gradient norms were 16.177856 and 7.746079. +- After approximately 2.5 hours, the four-update check had not completed. + Latest checkpoint pointer: step 0. Process 6802 had roughly 212 GiB RSS. +- Saved Python and native stacks identify `require_finite_update`, line 141, + at `p.grad.detach().float().cpu()`. This branch executes only after detecting + a nonfinite loss or aggregate gradient norm. It does not reveal which scalar + failed, the first affected matrix, or the numerical root cause. +- Native frames include `THPVariable_cpu` and tensor conversion. They confirm + waiting in the host transfer path; they do not prove a hardware failure, + compiler bug, deadlock, or that all training ran on CPU. +- Service subsequently confirmed `MainPID=0`, `ActiveState=inactive`, + `SubState=dead`. TPU allocation and data remain available. +- Evidence on disk: + `/mnt/disks/rg-data/gpt2small/ww-gpt2-validation-48h-20261004-s1337/diagnostics/stall-20261004-041944-714308`. + +### Reporting defect and correction + +The old failure handler copied full gradients to CPU, serially, before writing +the failure report. That code could stall and hide the already detected failure. +Commit `f887593` removed these copies, added an execution barrier before host +scalar reads, saved the scalar failure immediately, and bounded validation phases. +These changes have local test coverage; they have not established numerical +correctness on TPU. GPT-2 now has its own `baseline/gpt2_small` package. + +### Next diagnostic and attribution criteria + +Run four fresh AdamW updates on the same allocation, with the original model, +corpus, seed and optimizer hyperparameters. Save each completed update. Record: + +- Source commit, Python/torch/torch_xla/libtpu versions, relevant XLA settings. +- Exact input-window offsets, corpus identities, initialization/rolling full states. +- Per-tensor finite flags and extrema before clipping, after clipping, and after + the optimizer update; parameter names plus Adam moment names. +- XLA compilation/execution counters, fallback counters, stage timestamps, + Python tracebacks and a structured failure report. + +Checks reduce tensors on device and transfer a small summary table, never full +gradients. Additional synchronization is recorded as instrumentation: it can change +fusion/compilation behavior. A pass under instrumentation does not by itself clear +the original execution path. The diagnostic stops after four updates or 20 minutes; +up to 10 additional minutes are reserved for verified cloud backup. Disk evidence +remains if upload fails. No long experiment starts automatically. + +To attribute an upstream bug, isolate the first failing operation and compare a +matched CPU/TPU replay with the same inputs, weights and optimizer state. Preserve +both results and a minimal reproducer before claiming a PyTorch/XLA defect. diff --git a/baseline/gpt2_small/scripts/diagnose_adamw.sh b/baseline/gpt2_small/scripts/diagnose_adamw.sh index 923bdb9f..33929159 100644 --- a/baseline/gpt2_small/scripts/diagnose_adamw.sh +++ b/baseline/gpt2_small/scripts/diagnose_adamw.sh @@ -8,14 +8,18 @@ source "$base/scripts/tpu_environment.sh" python=/mnt/disks/rg-data/continuous8/venv/bin/python export PYTHONPATH="$base/src:$base/../nanogpt_one_head/src" "$python" -c 'import sys,time; assert float(sys.argv[1])>time.time()+60,"Too little allocation time; diagnostic not started"' "$deadline" -root="$old/diagnostic-$(date -u +%Y%m%d-%H%M%S)" -mkdir "$root" +root="${3:-$old/diagnostic-$(date -u +%Y%m%d-%H%M%S)}" +mkdir -p "$root" +test ! -e "$root/config.yaml" +export RG_GPT2_SOURCE_COMMIT=$(git -C "$base" rev-parse HEAD) +export PT_XLA_DEBUG_LEVEL=2 "$python" - "$old" "$root" <<'PY' import sys,yaml from pathlib import Path old,root=map(Path,sys.argv[1:]) cfg=yaml.safe_load((old/'configs/adamw.yaml').read_text()) -cfg.update(run_id=root.name,validation_gradient_checks=True,metrics_interval=1,benchmark_sync_every_step=True) +cfg.update(run_id=root.name,validation_gradient_checks=True,validation_tensor_checks=True, + metrics_interval=1,benchmark_sync_every_step=True) cfg['ww']['enabled']=False (root/'config.yaml').write_text(yaml.safe_dump(cfg,sort_keys=False)) PY @@ -27,4 +31,4 @@ backup() { exit "$rc" } trap backup EXIT -"$python" -u -m rg_gpt2_small.experiment --config "$root/config.yaml" --data-root /mnt/disks/rg-data/continuous8/data --output "$root/adamw" --device tpu --stop-after 4 --deadline-unix "$deadline" 2>&1 | tee "$root/diagnostic.log" +"$python" -u "$base/scripts/run_adamw_probe.py" "$root" "$deadline" 2>&1 | tee "$root/diagnostic.log" diff --git a/baseline/gpt2_small/scripts/retry_adamw.py b/baseline/gpt2_small/scripts/retry_adamw.py new file mode 100644 index 00000000..921f0d4b --- /dev/null +++ b/baseline/gpt2_small/scripts/retry_adamw.py @@ -0,0 +1,127 @@ +"""Launch/status a bounded diagnostic on the EXISTING 48-hour TPU only.""" +import argparse +import datetime as dt +import fcntl +import json +import os +from pathlib import Path +import re +import shlex +import subprocess +import sys +import time + +PROJECT='tpu-builders-504820'; ZONE='us-west4-a' +QUEUE='ww-gpt2-validation-48h-20261004-s1337'; NODE=QUEUE+'-node' +BASE=Path('/mnt/disks/rg-data/gpt2small'); OLD=BASE/QUEUE +LATEST=BASE/'PORT_CHECK_LATEST.json' + + +def run(args,**kwargs): + return subprocess.run(args,check=True,text=True,**kwargs) + + +def active(unit): + result=subprocess.run(['systemctl','show',unit,'--property=ActiveState','--value'], + text=True,capture_output=True,timeout=10) + return result.stdout.strip() in ('active','activating','deactivating','reloading') + + +def assert_idle(): + if active('rg-gpt2-validation.service') or active('rg-continuous8.service'): + raise RuntimeError('An existing training service is active; no diagnostic launched.') + if LATEST.exists() and active(json.loads(LATEST.read_text())['unit']): + raise RuntimeError('Port diagnostic already active. Use status; no second run launched.') + modules={'rg_gpt2_small.experiment','rg_nanogpt_one_head.gpt2_experiment', + 'rg_nanogpt_one_head.continuous_run','rg_nanogpt_one_head.tpu_spmd_check'} + for path in Path('/proc').glob('[0-9]*/cmdline'): + try: args=set(path.read_bytes().decode().split('\0')) + except (OSError,UnicodeError): continue + if args.intersection(modules): + raise RuntimeError(f'Trainer process {path.parent.name} is active; no concurrent diagnostic launched.') + + +def launch_remote(commit): + if os.geteuid()!=0 or not os.path.ismount('/mnt/disks/rg-data'): + raise RuntimeError('Requires root and the existing mounted data disk.') + if not re.fullmatch('[0-9a-f]{40}',commit): raise ValueError('Expected pinned commit SHA') + with (BASE/'port-check-launch.lock').open('a') as lock: + fcntl.flock(lock,fcntl.LOCK_EX|fcntl.LOCK_NB) + assert_idle() + allocation=json.loads((OLD/'allocation.json').read_text()) + allocation_deadline=float(allocation['validation_deadline_unix']) + if allocation_deadline-time.time()<900: + raise RuntimeError('Less than 15 minutes remain; diagnostic not launched.') + if not Path('/mnt/disks/rg-data/continuous8/data/train.bin').is_file(): + raise RuntimeError('Preserved FineWeb is missing; no download will be started.') + stamp=dt.datetime.now(dt.timezone.utc).strftime('%Y%m%d-%H%M%S') + root=BASE/('port-check-'+stamp); root.mkdir() + repo=root/'repo'; repo.mkdir() + run(['git','-C',str(repo),'init','-q']) + run(['git','-C',str(repo),'remote','add','origin','https://github.com/CalculatedContent/rg_optimizers.git']) + run(['git','-C',str(repo),'fetch','--depth','1','origin',commit],timeout=180) + run(['git','-C',str(repo),'checkout','--detach',commit]) + (root/'commit.txt').write_text(commit+'\n') + deadline=min(time.time()+1200,allocation_deadline-300) + if deadline-time.time()<300: raise RuntimeError('Too little time remains after source checkout.') + unit='rg-gpt2-port-check-'+stamp+'.service' + record={'root':str(root),'unit':unit,'commit':commit,'node':NODE, + 'training_deadline_unix':deadline,'maximum_training_seconds':1200, + 'purpose':'four-update AdamW numerical diagnostic','long_run_started':False} + (root/'launch.json').write_text(json.dumps(record,indent=2)) + command=['systemd-run','--unit='+unit,'--property=Type=exec','--property=Restart=no', + '--property=RuntimeMaxSec='+str(int(min(1800,allocation_deadline-time.time()))), + '--property=TimeoutStopSec=30','--property=KillMode=control-group', + '--property=StandardOutput=append:'+str(root/'run.log'), + '--property=StandardError=append:'+str(root/'run.log'), + '/bin/bash',str(repo/'baseline/gpt2_small/scripts/diagnose_adamw.sh'), + str(OLD),str(deadline),str(root)] + run(command) + temp=LATEST.with_suffix('.tmp'); temp.write_text(json.dumps(record,indent=2)); temp.replace(LATEST) + print('Diagnostic started:',unit,flush=True) + print('Log:',root/'run.log',flush=True) + print('Four AdamW updates only; training capped at 20 minutes, plus up to 10 minutes for backup.',flush=True) + print('Existing TPU allocation, FineWeb and prior outputs retained.',flush=True) + + +def status_remote(): + if not LATEST.exists(): + print('No new port diagnostic has been launched.'); return + record=json.loads(LATEST.read_text()); root=Path(record['root']) + print(json.dumps(record,indent=2),flush=True) + subprocess.run(['systemctl','--no-pager','--full','status',record['unit']],check=False) + for file in ('PROBE_STATUS.json','adamw/TPU_PORT_FAILURE.json','adamw/diagnostics/current_stage.json'): + path=root/file + if path.is_file(): print('\n'+file+'\n'+path.read_text(),flush=True) + subprocess.run(['tail','-n','70',str(root/'run.log')],check=False) + + +def main(): + parser=argparse.ArgumentParser(description=__doc__) + parser.add_argument('action',choices=('start','status')) + parser.add_argument('--on-tpu',action='store_true',help=argparse.SUPPRESS) + parser.add_argument('--commit',help=argparse.SUPPRESS) + args=parser.parse_args() + if args.on_tpu: + launch_remote(args.commit) if args.action=='start' else status_remote() + return 0 + commit=None + if args.action=='start': + repo=Path(__file__).resolve().parents[3] + dirty=run(['git','-C',str(repo),'status','--porcelain'],capture_output=True).stdout + if dirty.strip(): raise RuntimeError('Use a clean checkout of the pushed commit.') + commit=run(['git','-C',str(repo),'rev-parse','HEAD'],capture_output=True).stdout.strip() + queue=json.loads(run(['gcloud','alpha','compute','tpus','queued-resources','describe',QUEUE, + '--project='+PROJECT,'--zone='+ZONE,'--format=json'],capture_output=True).stdout) + if queue.get('state',{}).get('state')!='ACTIVE': + raise RuntimeError('The existing TPU is not ACTIVE; no allocation requested.') + remote=['sudo','python3','-c',Path(__file__).read_text(),args.action,'--on-tpu'] + if commit: remote+=['--commit',commit] + return subprocess.run(['gcloud','compute','tpus','tpu-vm','ssh',NODE, + '--project='+PROJECT,'--zone='+ZONE,'--worker=0','--command='+shlex.join(remote)]).returncode + + +if __name__=='__main__': + try: sys.exit(main()) + except Exception as exc: + print('Port diagnostic:',exc,file=sys.stderr); sys.exit(1) diff --git a/baseline/gpt2_small/scripts/run_adamw_probe.py b/baseline/gpt2_small/scripts/run_adamw_probe.py new file mode 100644 index 00000000..b4bd54f2 --- /dev/null +++ b/baseline/gpt2_small/scripts/run_adamw_probe.py @@ -0,0 +1,54 @@ +"""Four updates only, with an external deadline even if XLA never returns.""" +import argparse +import json +from pathlib import Path +import subprocess +import sys +import time + +from validation_watchdog import wait_for_child + + +def main(): + parser=argparse.ArgumentParser() + parser.add_argument('root',type=Path); parser.add_argument('deadline',type=float) + args=parser.parse_args(); root=args.root + report={'status':'running','long_run_started':False,'training_deadline_unix':args.deadline, + 'purpose':'TPU-port numerical debugging; not a performance benchmark'} + status=root/'PROBE_STATUS.json' + def persist(): + tmp=status.with_suffix('.tmp'); tmp.write_text(json.dumps(report,indent=2)); tmp.replace(status) + persist() + try: + if args.deadline<=time.time(): raise RuntimeError('Diagnostic training window expired before launch') + cmd=[sys.executable,'-X','faulthandler','-u','-m','rg_gpt2_small.experiment','--config',str(root/'config.yaml'), + '--data-root','/mnt/disks/rg-data/continuous8/data','--output',str(root/'adamw'), + '--device','tpu','--stop-after','4','--deadline-unix',str(args.deadline)] + child=subprocess.Popen(cmd) + rc=wait_for_child(child,'AdamW port diagnostic, four updates',args.deadline) + report['child_exit_code']=rc + if rc: raise RuntimeError(f'AdamW diagnostic exited with code {rc}; inspect run.log and TPU_PORT_FAILURE.json') + state=json.loads((root/'adamw/status.json').read_text()) + if state['step']!=4: raise RuntimeError('Stopped before completing four updates') + report.update(status='four_updates_completed',completed_updates=4, + interpretation='Diagnostic checks passed only; full validation and long-run stability remain unproven.') + print(json.dumps(report),flush=True) + return 0 + except Exception as exc: + report.update(status='failed_or_incomplete',error=str(exc)) + failure=root/'adamw/TPU_PORT_FAILURE.json' + if not failure.exists(): + stage=root/'adamw/diagnostics/current_stage.json' + failure.parent.mkdir(parents=True,exist_ok=True) + failure.write_text(json.dumps({ + 'status':'process_failed_or_timed_out','attribution':'unconfirmed', + 'exception':str(exc),'child_exit_code':report.get('child_exit_code'), + 'last_stage':json.loads(stage.read_text()) if stage.exists() else None, + 'evidence':['../run.log','diagnostics/','../config.yaml','../commit.txt']},indent=2)) + print(json.dumps(report),flush=True) + return 1 + finally: + persist() + + +if __name__=='__main__': sys.exit(main()) diff --git a/baseline/gpt2_small/src/rg_gpt2_small/experiment.py b/baseline/gpt2_small/src/rg_gpt2_small/experiment.py index afa405cc..b54a90f8 100644 --- a/baseline/gpt2_small/src/rg_gpt2_small/experiment.py +++ b/baseline/gpt2_small/src/rg_gpt2_small/experiment.py @@ -21,6 +21,7 @@ from rg_nanogpt_one_head.muonclip import install_muonclip_extension from rg_nanogpt_one_head import optimizers from rg_nanogpt_one_head.spectral import WeightMatrixHolder, _attach_matrix_metadata +from . import port_debug def atomic_json(path, value): @@ -47,10 +48,12 @@ def due(step, spec): step in {int(k * 10**p) for p in range(12) for k in (1, 2, 5)})) -def batch(array, generator, size, context, device): +def batch(array, generator, size, context, device, trace=None): if len(array) <= context: raise ValueError('Split is shorter than one context') starts = torch.randint(len(array) - context, (size,), generator=generator).tolist() + if trace is not None: + trace.append(starts) x = torch.from_numpy(np.stack([np.array(array[i:i+context], dtype=np.int64) for i in starts])) y = torch.from_numpy(np.stack([np.array(array[i+1:i+context+1], dtype=np.int64) for i in starts])) return spmd.batch_to_device(x, device), spmd.batch_to_device(y, device) @@ -180,10 +183,17 @@ def train(cfg, data_root, output, *, device='cpu', resume=False, stop_after=None with (output / 'writer.lock').open('a') as lock: fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB) diagnostic = cfg.get('validation_gradient_checks', False) + port_checks = cfg.get('validation_tensor_checks', False) + if port_checks: + port_debug.environment(output) if diagnostic: faulthandler.dump_traceback_later(300, repeat=True) try: return _train(cfg, data_root, output, device, resume, stop_after, deadline) + except Exception as exc: + if port_checks: + port_debug.failure(output,exc) + raise finally: if diagnostic: faulthandler.cancel_dump_traceback_later() @@ -192,6 +202,9 @@ def train(cfg, data_root, output, *, device='cpu', resume=False, stop_after=None def _train(cfg, data_root, output, device, resume, stop_after, deadline): cfg = copy.deepcopy(cfg) t = cfg['training']; context = cfg['model']['block_size'] + port_checks = cfg.get('validation_tensor_checks', False) + if port_checks: + port_debug.stage(output,'initializing_runtime_and_data',0) step_tokens = t['batch_size'] * t['grad_accum_steps'] * context total = min(t.get('max_steps', 10**12), math.ceil(t['max_tokens'] / step_tokens)) if not 0 <= t['warmup_steps'] < t['schedule_steps'] or total < 1: @@ -274,16 +287,37 @@ def record(step, final=False): peak_lr=handle.peak_lr, min_lr=handle.min_lr) handle.set_lr(lr) losses = [] + input_windows = [] if port_checks else None + if port_checks: + port_debug.stage(output,'building_forward_backward_graph',step+1) for _ in range(t['grad_accum_steps']): - x, y = batch(arrays['train'], gen, t['batch_size'], context, dev) + x, y = batch(arrays['train'], gen, t['batch_size'], context, dev,trace=input_windows) _, loss = model(x, y); losses.append(loss.detach()) (loss / t['grad_accum_steps']).backward() + if port_checks: + atomic_json(output/'diagnostics'/f'{step+1:06d}-input-windows.json', + {'update':step+1,'split':'train','context':context,'microbatch_offsets':input_windows, + 'corpus_identity':'manifest.json/data','config_fingerprint':fingerprint}) spmd.replicate_gradients(model) norm = rt.gradient_norm(model.parameters()) + if port_checks: + port_debug.check_tensors( + [('gradient/'+name,p.grad) for name,p in model.named_parameters() if p.grad is not None] + + [('microbatch_loss/'+str(i),loss) for i,loss in enumerate(losses)] + + [('global_gradient_norm',norm)],output,'before_clipping',step+1,dev) if cfg.get('validation_gradient_checks', False): require_finite_update(model, norm, losses, output, step + 1) torch.nn.utils.clip_grad_norm_(model.parameters(), t['grad_clip'], foreach=False) + if port_checks: + port_debug.check_tensors( + [('gradient/'+name,p.grad) for name,p in model.named_parameters() if p.grad is not None], + output,'after_clipping',step+1,dev) + port_debug.stage(output,'optimizer_update_started',step+1) optimizers.optimizer_step(handles); rt.mark_step(dev) + if port_checks: + port_debug.check_tensors(port_debug.optimizer_tensors(model,handles), + output,'after_optimizer',step+1,dev) + port_debug.stage(output,'optimizer_update_completed',step+1) step += 1 measurement_due = (step % cfg['metrics_interval'] == 0 or due(step, cfg['ww']) or step == total or step == stop_after) diff --git a/baseline/gpt2_small/src/rg_gpt2_small/port_debug.py b/baseline/gpt2_small/src/rg_gpt2_small/port_debug.py new file mode 100644 index 00000000..59ee9bcb --- /dev/null +++ b/baseline/gpt2_small/src/rg_gpt2_small/port_debug.py @@ -0,0 +1,110 @@ +"""Small, device-reduced numerical reports for diagnosing the TPU port. + +Only scalar summaries cross to the CPU. Full gradients are never copied here. +Instrumentation changes execution boundaries, so record it as a diagnostic run. +""" +import importlib.metadata +import json +import math +import os +from pathlib import Path +import platform +import time +import traceback + +import torch +from rg_nanogpt_one_head import runtime as rt + + +def write(path, value): + path=Path(path); path.parent.mkdir(parents=True,exist_ok=True) + tmp=path.with_suffix(path.suffix+'.tmp') + with tmp.open('w') as f: + json.dump(value,f,indent=2,allow_nan=False) + f.flush(); os.fsync(f.fileno()) + tmp.replace(path) + + +def scalar(value): + return value if math.isfinite(value) else str(value) + + +def environment(output): + versions={} + for package in ('torch','torch-xla','libtpu','libtpu-nightly','numpy','PyYAML'): + try: versions[package]=importlib.metadata.version(package) + except importlib.metadata.PackageNotFoundError: versions[package]=None + variables=('PJRT_DEVICE','TPU_ACCELERATOR_TYPE','XLA_USE_BF16','XLA_DOWNCAST_BF16', + 'XLA_FLAGS','XLA_MATMUL_PRECISION','PT_XLA_DEBUG_LEVEL','OMP_NUM_THREADS') + write(Path(output)/'diagnostics/environment.json',{ + 'python':platform.python_version(),'platform':platform.platform(),'versions':versions, + 'source_commit':os.environ.get('RG_GPT2_SOURCE_COMMIT'), + 'environment':{key:os.environ.get(key) for key in variables}, + 'mode':'diagnostic; additional synchronization and per-tensor reductions enabled', + 'attribution':'unconfirmed: model/training code, port, runtime, and hardware not yet isolated'}) + + +def stage(output,name,step): + row={'stage':name,'update':step,'unix_time':time.time()} + write(Path(output)/'diagnostics/current_stage.json',row) + print(json.dumps(row),flush=True) + + +def xla_metrics(output,label,device): + if torch.device(device).type!='xla': return + from torch_xla.debug import metrics + path=Path(output)/'diagnostics'/f'xla-{label}.txt' + path.parent.mkdir(parents=True,exist_ok=True) + path.write_text(metrics.metrics_report()) + + +def check_tensors(named,output,label,step,device): + named=list(named) + if not named: return + stage(output,label+'_started',step) + rows=[] + for name,tensor in named: + value=tensor.detach().float() + rows.append(torch.stack((torch.isfinite(value).all().float(),value.amin(),value.amax()))) + table=torch.stack(rows) + rt.synchronize(torch.device(device)) + # A single small transfer, following an explicit execution barrier. + values=table.cpu().tolist() + records=[] + for (name,tensor),(finite,minimum,maximum) in zip(named,values): + negative_variance=name.endswith('/exp_avg_sq') and minimum<0 + records.append({'tensor':name,'shape':list(tensor.shape),'dtype':str(tensor.dtype), + 'all_finite':bool(finite),'min':scalar(minimum),'max':scalar(maximum), + 'negative_second_moment':negative_variance}) + bad=[r for r in records if not r['all_finite'] or r['negative_second_moment']] + report={'stage':label,'update':step,'records':records,'invalid_tensors':bad} + write(Path(output)/'diagnostics'/f'{step:06d}-{label}.json',report) + xla_metrics(output,f'{step:06d}-{label}',device) + if bad: + write(Path(output)/'diagnostics/first_invalid_tensors.json',report) + print(json.dumps({'invalid_stage':label,'update':step,'invalid_tensors':bad}),flush=True) + raise RuntimeError(f'Invalid tensor at {label}, update {step}; see diagnostics/first_invalid_tensors.json') + stage(output,label+'_passed',step) + + +def optimizer_tensors(model,handles): + names={id(p):name for name,p in model.named_parameters()} + yield from (('weight/'+name,p) for name,p in model.named_parameters()) + for handle in handles: + for parameter,state in handle.optimizer.state.items(): + for key,value in state.items(): + # Adam's CPU step counter is metadata, not a TPU moment tensor. + if torch.is_tensor(value) and value.device==parameter.device and value.numel()>1: + yield f'{handle.role}/{names[id(parameter)]}/{key}',value + + +def failure(output,error): + current=Path(output)/'diagnostics/current_stage.json' + write(Path(output)/'TPU_PORT_FAILURE.json',{ + 'status':'failed','attribution':'unconfirmed; not yet demonstrated to be an upstream TPU/XLA bug', + 'source_commit':os.environ.get('RG_GPT2_SOURCE_COMMIT'), + 'last_stage':json.loads(current.read_text()) if current.exists() else None, + 'exception_type':type(error).__name__,'exception':str(error), + 'traceback':traceback.format_exc(), + 'evidence':['manifest.json','diagnostics/environment.json','diagnostics/'], + 'next_comparison':'Replay the failing batch/state on CPU and TPU before upstream attribution.'}) diff --git a/baseline/gpt2_small/tests/test_port_debug.py b/baseline/gpt2_small/tests/test_port_debug.py new file mode 100644 index 00000000..ca80f216 --- /dev/null +++ b/baseline/gpt2_small/tests/test_port_debug.py @@ -0,0 +1,101 @@ +import importlib.util +import json +from pathlib import Path +import types + +import pytest +import torch + +from rg_gpt2_small import port_debug as debug +from rg_gpt2_small import experiment as g +from test_gpt2_experiment import config,data + + +def test_nonfinite_layer_report_uses_one_small_transfer(tmp_path,monkeypatch): + shapes=[]; original=torch.Tensor.cpu + def cpu(tensor,*args,**kwargs): + shapes.append(tuple(tensor.shape)); return original(tensor,*args,**kwargs) + monkeypatch.setattr(torch.Tensor,'cpu',cpu) + with pytest.raises(RuntimeError,match='before_clipping'): + debug.check_tensors([('gradient/good',torch.ones(128,128)), + ('gradient/bad',torch.tensor([float('nan'),1.]))], + tmp_path,'before_clipping',3,'cpu') + row=json.loads((tmp_path/'diagnostics/first_invalid_tensors.json').read_text()) + assert [x['tensor'] for x in row['invalid_tensors']]==['gradient/bad'] + assert shapes==[(2,3)] + assert row['invalid_tensors'][0]['min']=='nan' + + +def test_negative_second_moment_detected(tmp_path): + with pytest.raises(RuntimeError,match='after_optimizer'): + debug.check_tensors([('primary/layer/exp_avg_sq',torch.tensor([-.1,1.]))], + tmp_path,'after_optimizer',2,'cpu') + row=json.loads((tmp_path/'diagnostics/first_invalid_tensors.json').read_text()) + assert row['invalid_tensors'][0]['negative_second_moment'] is True + + +def test_cpu_diagnostic_records_stages_inputs_and_states(tmp_path): + torch.set_num_threads(1) + c=config(); c.update(validation_tensor_checks=True,validation_gradient_checks=True,metrics_interval=1) + out=tmp_path/'run'; g.train(c,data(tmp_path,c),out,stop_after=2) + assert json.loads((out/'status.json').read_text())['step']==2 + for step in (1,2): + for stage in ('before_clipping','after_clipping','after_optimizer'): + row=json.loads((out/f'diagnostics/{step:06d}-{stage}.json').read_text()) + assert not row['invalid_tensors'] + windows=json.loads((out/f'diagnostics/{step:06d}-input-windows.json').read_text()) + assert len(windows['microbatch_offsets'])==c['training']['grad_accum_steps'] + assert (out/'diagnostics/environment.json').is_file() + + +def test_training_failure_keeps_initial_checkpoint_and_layer_evidence(tmp_path,monkeypatch): + torch.set_num_threads(1) + c=config(); c.update(validation_tensor_checks=True,validation_gradient_checks=True) + original=g.GPT + def broken_model(*args,**kwargs): + model=original(*args,**kwargs) + model.token_embedding.weight.register_hook(lambda gradient:gradient*float('nan')) + return model + monkeypatch.setattr(g,'GPT',broken_model) + out=tmp_path/'failed' + with pytest.raises(RuntimeError,match='before_clipping'): + g.train(c,data(tmp_path,c),out) + report=json.loads((out/'TPU_PORT_FAILURE.json').read_text()) + assert report['attribution'].startswith('unconfirmed') + assert json.loads((out/'checkpoints/latest.json').read_text())['step']==0 + bad=json.loads((out/'diagnostics/first_invalid_tensors.json').read_text())['invalid_tensors'] + assert 'gradient/token_embedding.weight' in [r['tensor'] for r in bad] + + +def load_retry(): + spec=importlib.util.spec_from_file_location('retry',Path(__file__).parents[1]/'scripts/retry_adamw.py') + module=importlib.util.module_from_spec(spec); spec.loader.exec_module(module); return module + + +def test_retry_refuses_active_training(monkeypatch): + module=load_retry() + monkeypatch.setattr(module,'active',lambda unit:unit=='rg-gpt2-validation.service') + with pytest.raises(RuntimeError,match='existing training service'): + module.assert_idle() + + +def test_retry_is_bounded_and_never_allocates(tmp_path,monkeypatch): + module=load_retry(); old=tmp_path/'old'; old.mkdir() + (old/'allocation.json').write_text(json.dumps({'validation_deadline_unix':100000.})) + monkeypatch.setattr(module,'BASE',tmp_path); monkeypatch.setattr(module,'OLD',old) + monkeypatch.setattr(module,'LATEST',tmp_path/'latest.json') + monkeypatch.setattr(module.os,'geteuid',lambda:0) + monkeypatch.setattr(module.os.path,'ismount',lambda path:True) + monkeypatch.setattr(module,'assert_idle',lambda:None) + original=Path.is_file + monkeypatch.setattr(Path,'is_file',lambda path:True if str(path).endswith('/data/train.bin') else original(path)) + monkeypatch.setattr(module.time,'time',lambda:1000.) + calls=[] + monkeypatch.setattr(module,'run',lambda args,**kwargs:calls.append(args)) + module.launch_remote('a'*40) + record=json.loads(module.LATEST.read_text()) + assert record['training_deadline_unix']==2200. + systemd=next(c for c in calls if c[0]=='systemd-run') + assert '--property=RuntimeMaxSec=1800' in systemd + assert '--property=Restart=no' in systemd + assert not any(c[0]=='gcloud' for c in calls) From 649b197d814c3b2b6a1efd67be89f9b32ce78f93 Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Sat, 3 Oct 2026 22:07:41 -0700 Subject: [PATCH 19/36] Record native SIGABRT in the first instrumented TPU gradient check --- baseline/gpt2_small/TPU_PORT_BUGS.md | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/baseline/gpt2_small/TPU_PORT_BUGS.md b/baseline/gpt2_small/TPU_PORT_BUGS.md index 5d468f34..7b6d0d51 100644 --- a/baseline/gpt2_small/TPU_PORT_BUGS.md +++ b/baseline/gpt2_small/TPU_PORT_BUGS.md @@ -64,3 +64,23 @@ remains if upload fails. No long experiment starts automatically. To attribute an upstream bug, isolate the first failing operation and compare a matched CPU/TPU replay with the same inputs, weights and optimizer state. Preserve both results and a minimal reproducer before claiming a PyTorch/XLA defect. + +## 2026-10-04: instrumented diagnostic aborts at its first gradient check + +Status: open; fatal native message still required for diagnosis. + +- Run: `port-check-20261004-045633`, commit `19e2bb0`. +- Supervisor report: child exit code `-6` (SIGABRT), with last recorded stage + `before_clipping_started`, update 1, Unix time `1791089838.8617651`. +- No per-tensor results from this check were reported. This abort does not by + itself establish a nonfinite gradient, a particular failing operator, or an + upstream runtime/hardware bug. It is a separate observed failure from the + earlier numerical check and stalled host copy. +- Cloud backup was explicitly verified for this run. Logs, initialization, + exact input-window offsets and environment metadata remain on disk and in + its cloud prefix. The next evidence to inspect is the fatal native message + immediately preceding the abort in `run.log`. +- A later launch was blocked by an untracked repository-root `FETCH_HEAD` + file in Cloud Shell. This local checkout issue is separate from the TPU + abort; moving that file outside the checkout preserves it and clears this + particular cleanliness check. The real Git metadata is under `.git`. From 08f10cf6f864cdee1c100a2b811049bd8e3ec505 Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Sat, 3 Oct 2026 22:10:16 -0700 Subject: [PATCH 20/36] Localize diagnostic SIGABRT to XLA stack shape inference --- baseline/gpt2_small/TPU_PORT_BUGS.md | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/baseline/gpt2_small/TPU_PORT_BUGS.md b/baseline/gpt2_small/TPU_PORT_BUGS.md index 7b6d0d51..7a968aef 100644 --- a/baseline/gpt2_small/TPU_PORT_BUGS.md +++ b/baseline/gpt2_small/TPU_PORT_BUGS.md @@ -69,6 +69,15 @@ both results and a minimal reproducer before claiming a PyTorch/XLA defect. Status: open; fatal native message still required for diagnosis. +Further stack evidence localizes this abort to `port_debug.py:68`: the new +diagnostic's `torch.stack((isfinite(value).all().float(), value.amin(), value.amax()))`. +Native frames include `torch_xla::Stack::Stack`, `XlaNode::GetOpShape`, and +`XLANativeFunctions::stack`. The abort occurs while assembling diagnostic summaries, +before the first optimizer update. It therefore does not reproduce or explain the +earlier nonfinite-result failure. CPU tests passed this operation, but TPU behavior +has not passed validation. The preceding native assertion/status text is still +needed; the Python abort trace alone is insufficient to identify its cause. + - Run: `port-check-20261004-045633`, commit `19e2bb0`. - Supervisor report: child exit code `-6` (SIGABRT), with last recorded stage `before_clipping_started`, update 1, Unix time `1791089838.8617651`. From 6f8d59214f51be00212a9956182a0c697225f086 Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Sat, 3 Oct 2026 22:29:27 -0700 Subject: [PATCH 21/36] Run continuous MuonClip with per-tensor diagnostic disabled and verified rolling backups Preserve the existing TPU, FineWeb corpus and all crash evidence. Launch one fresh process with scalar finite guards, paired token error/alpha monitoring, a stage stall watchdog, and allocation cutoff. Retain all metrics with bounded local/cloud checkpoints. Validation: 43 CPU tests passed; TPU stability remains unverified. --- baseline/gpt2_small/README.md | 45 ++++++ baseline/gpt2_small/TPU_PORT_BUGS.md | 16 +++ .../gpt2_small/scripts/muonclip_worker.sh | 25 ++++ baseline/gpt2_small/scripts/retry_adamw.py | 3 + baseline/gpt2_small/scripts/run_muonclip.py | 132 ++++++++++++++++++ .../gpt2_small/scripts/supervise_muonclip.py | 106 ++++++++++++++ .../src/rg_gpt2_small/experiment.py | 44 ++++-- .../src/rg_gpt2_small/run_backup.py | 40 ++++++ .../gpt2_small/tests/test_muonclip_night.py | 130 +++++++++++++++++ 9 files changed, 531 insertions(+), 10 deletions(-) create mode 100644 baseline/gpt2_small/scripts/muonclip_worker.sh create mode 100644 baseline/gpt2_small/scripts/run_muonclip.py create mode 100644 baseline/gpt2_small/scripts/supervise_muonclip.py create mode 100644 baseline/gpt2_small/src/rg_gpt2_small/run_backup.py create mode 100644 baseline/gpt2_small/tests/test_muonclip_night.py diff --git a/baseline/gpt2_small/README.md b/baseline/gpt2_small/README.md index f82cda24..be6394b7 100644 --- a/baseline/gpt2_small/README.md +++ b/baseline/gpt2_small/README.md @@ -11,6 +11,51 @@ The TPU launch scripts expose both source packages through `PYTHONPATH`. For loc development, from the repository root, install both with `pip install -e baseline/nanogpt_one_head -e baseline/gpt2_small`. +## Continue with MuonClip on the existing allocation + +The user authorized a fresh continuous MuonClip run after the diagnostic crash. +Use the current 48-hour TPU, installed environment and preserved FineWeb: + +```bash +python3 baseline/gpt2_small/scripts/run_muonclip.py start +python3 baseline/gpt2_small/scripts/run_muonclip.py status +``` + +This starts one fresh process under `rg-gpt2-muonclip-.service` with +`Restart=no`, in `/mnt/disks/rg-data/gpt2small/muonclip-night-`. +It never invokes the cleanup/reallocation scripts, changes the installed packages, +downloads the corpus, or writes into previous experiment directories/cloud prefixes. +The saved `port-check-20261004-045633` crash evidence remains intact. A shared launch +lock and checks for existing services/trainers prevent simultaneous TPU jobs. + +Only the crashing per-tensor gradient diagnostic is disabled. A scalar finite +loss/gradient-norm guard remains before each update. The 124M GPT-2 model, data, +MuonClip/auxiliary AdamW settings, batch size and long-run LR schedule are retained. +This bypass does not fix or explain the earlier numerical failure or establish TPU +stability. It is an explicitly requested experimental run, not a passed validation. + +Training stays in one process from initialization until the existing allocation +cutoff (20 minutes before its recorded expiry), manual STOP, token budget or error. +There is no AdamW gate, stop/resume transition or automatic restart. A watchdog stops +a phase with no progress for 30 minutes and saves a failure report. First completed +updates are printed explicitly; a watch heartbeat alone is not completion evidence. + +Token error, NLL and full checkpoints are recorded at steps 0, 1, 2, 4 and every 25 +updates. All 72 matrix raw/clipped alphas are paired with the same-step token metrics +at step 25 and every 100 updates. Evaluation windows stay fixed. Synchronization, +evaluation, spectra and checkpoint upload overhead contribute to elapsed runtime. + +Three rolling local full checkpoints and an initialization milestone limit disk +usage. During training, verified uploads keep three rotating cloud slots plus +initialization; every scalar/spectral JSON record is retained. The cloud prefix is +`gs://tpu-builders-504820-ww-continuous8/gpt2small/`. +`muonclip/checkpoints/LATEST_VERIFIED.json` is published only after checkpoint and +metrics uploads, with object generation and CRC32C. It names a cloud slot, not a +local filename; validate generation/checksum when downloading for recovery. A failed +upload stops the trainer, retains local evidence, and does not advance this pointer. +Exit backup additionally archives the remaining local files using object permissions. +Old-run checkpoints and cloud objects are never deleted by this workflow. + ## Current validation failure Observed failures and their attribution are tracked in [TPU_PORT_BUGS.md](TPU_PORT_BUGS.md). diff --git a/baseline/gpt2_small/TPU_PORT_BUGS.md b/baseline/gpt2_small/TPU_PORT_BUGS.md index 7a968aef..551ec606 100644 --- a/baseline/gpt2_small/TPU_PORT_BUGS.md +++ b/baseline/gpt2_small/TPU_PORT_BUGS.md @@ -93,3 +93,19 @@ needed; the Python abort trace alone is insufficient to identify its cause. file in Cloud Shell. This local checkout issue is separate from the TPU abort; moving that file outside the checkout preserves it and clears this particular cleanliness check. The real Git metadata is under `.git`. + +### User-authorized MuonClip bypass, 2026-10-04 + +`run_muonclip.py` starts a separate continuous MuonClip experiment on the remaining +48-hour allocation. Its config sets `validation_tensor_checks=false` and +`validation_gradient_checks=false`, bypassing the crashing per-tensor stack and its +verbose validation path. `finite_update_guard=true` still checks scalar losses and +aggregate norm before the optimizer update; the existing numerical guard is not +removed. Model/data/optimizer and long-run LR settings are unchanged. Additional +synchronization, reporting and denser checkpoint/spectral measurements are explicit. + +The workaround has CPU integration coverage; TPU success is not claimed. It is not +a repair or root-cause diagnosis for either SIGABRT or the earlier nonfinite result. +The earlier source checkouts, diagnostics, initialization checkpoints and verified +cloud archives are retained. New failures produce separate evidence in the new run. +No automatic restart, cleanup, disk formatting or TPU reallocation is performed. diff --git a/baseline/gpt2_small/scripts/muonclip_worker.sh b/baseline/gpt2_small/scripts/muonclip_worker.sh new file mode 100644 index 00000000..e6f11a19 --- /dev/null +++ b/baseline/gpt2_small/scripts/muonclip_worker.sh @@ -0,0 +1,25 @@ +#!/usr/bin/env bash +# One fresh process, on the current allocation and preserved corpus. +set -Eeuo pipefail +root=$1 +deadline=$2 +service_deadline=$3 +base=$(cd "$(dirname "$0")/.." && pwd) +python=/mnt/disks/rg-data/continuous8/venv/bin/python +export PYTHONPATH="$base/src:$base/../nanogpt_one_head/src" +export RG_GPT2_SOURCE_COMMIT=$(git -C "$base" rev-parse HEAD) +export RG_GPT2_GCS_URI="gs://tpu-builders-504820-ww-continuous8/gpt2small/$(basename "$root")" +export RG_GPT2_RUN_LOG="$root/run.log" +source "$base/scripts/tpu_environment.sh" +mountpoint -q /mnt/disks/rg-data +test -x "$python" +finish() { + rc=$? + trap - EXIT + echo "MuonClip worker exit code: $rc; retaining all persistent-disk evidence." + sync + "$python" "$base/scripts/backup.py" "$root" || exit 1 + exit "$rc" +} +trap finish EXIT +"$python" -u "$base/scripts/supervise_muonclip.py" "$root" "$deadline" "$service_deadline" diff --git a/baseline/gpt2_small/scripts/retry_adamw.py b/baseline/gpt2_small/scripts/retry_adamw.py index 921f0d4b..ac566c6a 100644 --- a/baseline/gpt2_small/scripts/retry_adamw.py +++ b/baseline/gpt2_small/scripts/retry_adamw.py @@ -32,6 +32,9 @@ def assert_idle(): raise RuntimeError('An existing training service is active; no diagnostic launched.') if LATEST.exists() and active(json.loads(LATEST.read_text())['unit']): raise RuntimeError('Port diagnostic already active. Use status; no second run launched.') + night=BASE/'MUONCLIP_LATEST.json' + if night.exists() and active(json.loads(night.read_text())['unit']): + raise RuntimeError('MuonClip service is active; no concurrent diagnostic launched.') modules={'rg_gpt2_small.experiment','rg_nanogpt_one_head.gpt2_experiment', 'rg_nanogpt_one_head.continuous_run','rg_nanogpt_one_head.tpu_spmd_check'} for path in Path('/proc').glob('[0-9]*/cmdline'): diff --git a/baseline/gpt2_small/scripts/run_muonclip.py b/baseline/gpt2_small/scripts/run_muonclip.py new file mode 100644 index 00000000..90205d03 --- /dev/null +++ b/baseline/gpt2_small/scripts/run_muonclip.py @@ -0,0 +1,132 @@ +"""Launch a fresh continuous MuonClip run on the EXISTING 48-hour TPU.""" +import argparse +import datetime as dt +import fcntl +import json +import os +from pathlib import Path +import re +import shlex +import subprocess +import sys +import time + +PROJECT='tpu-builders-504820'; ZONE='us-west4-a' +QUEUE='ww-gpt2-validation-48h-20261004-s1337'; NODE=QUEUE+'-node' +BASE=Path('/mnt/disks/rg-data/gpt2small'); OLD=BASE/QUEUE +LATEST=BASE/'MUONCLIP_LATEST.json' + + +def run(args,**kwargs): + return subprocess.run(args,check=True,text=True,**kwargs) + + +def active(unit): + result=subprocess.run(['systemctl','show',unit,'--property=ActiveState','--value'], + text=True,capture_output=True,timeout=10) + return result.stdout.strip() in ('active','activating','deactivating','reloading') + + +def assert_idle(): + if active('rg-gpt2-validation.service') or active('rg-continuous8.service'): + raise RuntimeError('An existing training service is active; no training launched.') + if LATEST.exists() and active(json.loads(LATEST.read_text())['unit']): + raise RuntimeError('MuonClip is already active. Use status; no second run launched.') + night=BASE/'PORT_CHECK_LATEST.json' + if night.exists() and active(json.loads(night.read_text())['unit']): + raise RuntimeError('Port diagnostic is active; no concurrent training launched.') + modules={'rg_gpt2_small.experiment','rg_nanogpt_one_head.gpt2_experiment', + 'rg_nanogpt_one_head.continuous_run','rg_nanogpt_one_head.tpu_spmd_check'} + for path in Path('/proc').glob('[0-9]*/cmdline'): + try: args=set(path.read_bytes().decode().split('\0')) + except (OSError,UnicodeError): continue + if args.intersection(modules): + raise RuntimeError(f'Trainer process {path.parent.name} is active; no concurrent diagnostic launched.') + + +def launch_remote(commit): + if os.geteuid()!=0 or not os.path.ismount('/mnt/disks/rg-data'): + raise RuntimeError('Requires root and the existing mounted data disk.') + if not re.fullmatch('[0-9a-f]{40}',commit): raise ValueError('Expected pinned commit SHA') + with (BASE/'port-check-launch.lock').open('a') as lock: + fcntl.flock(lock,fcntl.LOCK_EX|fcntl.LOCK_NB) + assert_idle() + allocation=json.loads((OLD/'allocation.json').read_text()) + allocation_deadline=float(allocation['validation_deadline_unix']) + if allocation_deadline-time.time()<900: + raise RuntimeError('Less than 15 minutes remain; training not launched.') + if not Path('/mnt/disks/rg-data/continuous8/data/train.bin').is_file(): + raise RuntimeError('Preserved FineWeb is missing; no download will be started.') + stamp=dt.datetime.now(dt.timezone.utc).strftime('%Y%m%d-%H%M%S') + root=BASE/('muonclip-night-'+stamp); root.mkdir() + repo=root/'repo'; repo.mkdir() + run(['git','-C',str(repo),'init','-q']) + run(['git','-C',str(repo),'remote','add','origin','https://github.com/CalculatedContent/rg_optimizers.git']) + run(['git','-C',str(repo),'fetch','--depth','1','origin',commit],timeout=180) + run(['git','-C',str(repo),'checkout','--detach',commit]) + (root/'commit.txt').write_text(commit+'\n') + deadline=allocation_deadline-600 + if deadline-time.time()<300: raise RuntimeError('Too little time remains after source checkout.') + unit='rg-gpt2-muonclip-'+stamp+'.service' + record={'root':str(root),'unit':unit,'commit':commit,'node':NODE, + 'training_deadline_unix':deadline,'service_deadline_unix':allocation_deadline, + 'purpose':'continuous MuonClip with scalar finite guard; per-tensor diagnostic disabled', + 'cloud_uri':'gs://tpu-builders-504820-ww-continuous8/gpt2small/'+root.name} + (root/'launch.json').write_text(json.dumps(record,indent=2)) + command=['systemd-run','--unit='+unit,'--property=Type=exec','--property=Restart=no', + '--property=RuntimeMaxSec='+str(int(allocation_deadline-time.time())), + '--property=TimeoutStopSec=30','--property=KillMode=control-group', + '--property=StandardOutput=append:'+str(root/'run.log'), + '--property=StandardError=append:'+str(root/'run.log'), + '/bin/bash',str(repo/'baseline/gpt2_small/scripts/muonclip_worker.sh'), + str(root),str(deadline),str(allocation_deadline)] + run(command) + temp=LATEST.with_suffix('.tmp'); temp.write_text(json.dumps(record,indent=2)); temp.replace(LATEST) + print('MuonClip service started:',unit,flush=True) + print('Log:',root/'run.log',flush=True) + print('Training cutoff UTC:',dt.datetime.fromtimestamp(deadline,dt.timezone.utc).isoformat(),flush=True) + print('Checkpoints/token error every 25 updates; raw/clipped alpha every 100; no automatic restart.',flush=True) + print('Existing TPU allocation, FineWeb and prior outputs retained.',flush=True) + + +def status_remote(): + if not LATEST.exists(): + print('No overnight MuonClip run has been launched.'); return + record=json.loads(LATEST.read_text()); root=Path(record['root']) + print(json.dumps(record,indent=2),flush=True) + subprocess.run(['systemctl','--no-pager','--full','status',record['unit']],check=False) + for file in ('RUN_STATUS.json','muonclip/TPU_PORT_FAILURE.json','muonclip/progress.json','muonclip/cloud_checkpoint.json'): + path=root/file + if path.is_file(): print('\n'+file+'\n'+path.read_text(),flush=True) + subprocess.run(['tail','-n','70',str(root/'run.log')],check=False) + + +def main(): + parser=argparse.ArgumentParser(description=__doc__) + parser.add_argument('action',choices=('start','status')) + parser.add_argument('--on-tpu',action='store_true',help=argparse.SUPPRESS) + parser.add_argument('--commit',help=argparse.SUPPRESS) + args=parser.parse_args() + if args.on_tpu: + launch_remote(args.commit) if args.action=='start' else status_remote() + return 0 + commit=None + if args.action=='start': + repo=Path(__file__).resolve().parents[3] + dirty=run(['git','-C',str(repo),'status','--porcelain'],capture_output=True).stdout + if dirty.strip(): raise RuntimeError('Use a clean checkout of the pushed commit.') + commit=run(['git','-C',str(repo),'rev-parse','HEAD'],capture_output=True).stdout.strip() + queue=json.loads(run(['gcloud','alpha','compute','tpus','queued-resources','describe',QUEUE, + '--project='+PROJECT,'--zone='+ZONE,'--format=json'],capture_output=True).stdout) + if queue.get('state',{}).get('state')!='ACTIVE': + raise RuntimeError('The existing TPU is not ACTIVE; no allocation requested.') + remote=['sudo','python3','-c',Path(__file__).read_text(),args.action,'--on-tpu'] + if commit: remote+=['--commit',commit] + return subprocess.run(['gcloud','compute','tpus','tpu-vm','ssh',NODE, + '--project='+PROJECT,'--zone='+ZONE,'--worker=0','--command='+shlex.join(remote)]).returncode + + +if __name__=='__main__': + try: sys.exit(main()) + except Exception as exc: + print('MuonClip launch:',exc,file=sys.stderr); sys.exit(1) diff --git a/baseline/gpt2_small/scripts/supervise_muonclip.py b/baseline/gpt2_small/scripts/supervise_muonclip.py new file mode 100644 index 00000000..96c54dd6 --- /dev/null +++ b/baseline/gpt2_small/scripts/supervise_muonclip.py @@ -0,0 +1,106 @@ +"""Supervise one continuous trainer; stop on failure/stall, never restart.""" +import argparse +import json +import os +from pathlib import Path +import subprocess +import sys +import time + +import yaml + + +def write(path, value): + tmp=path.with_suffix('.tmp'); tmp.write_text(json.dumps(value,indent=2)); tmp.replace(path) + + +def prepare_config(template, run_id): + cfg=yaml.safe_load(Path(template).read_text()) + cfg.update(run_id=run_id, validation_tensor_checks=False, validation_gradient_checks=False, + finite_update_guard=True, progress_reporting=True, benchmark_sync_every_step=True, + cloud_checkpoints=True, metrics_interval=25, metrics_steps=[1,2,4], milestones=[0]) + cfg['ww'].update(enabled=True, interval=100, steps=[25], logarithmic=False) + return cfg + + +def stop(child): + if child.poll() is not None: return + child.terminate() + try: child.wait(timeout=20) + except subprocess.TimeoutExpired: + child.kill(); child.wait(timeout=20) + + +def watch(child, output, deadline, grace_deadline, stall_seconds=1800): + last_change=time.monotonic(); previous=None; requested=False + while True: + rc=child.poll() + if rc is not None: return rc + now=time.time() + if now>=deadline and not requested: + (output/'STOP').touch(); requested=True + print('Allocation cutoff reached; requesting final save.',flush=True) + if now>=grace_deadline: + raise RuntimeError('Trainer did not exit within allocation cutoff grace period.') + path=output/'progress.json' + if path.exists(): + current=path.read_text() # atomic file replacement by trainer + if current!=previous: + previous=current; last_change=time.monotonic() + if time.monotonic()-last_change>=stall_seconds: + raise RuntimeError('No training stage completed/changed for 30 minutes; stopping stalled run.') + try: return child.wait(timeout=30) + except subprocess.TimeoutExpired: + stage=json.loads(previous) if previous else {'stage':'starting','completed_step':0} + print('[muonclip-watch] '+json.dumps(stage),flush=True) + + +def main(): + parser=argparse.ArgumentParser(description=__doc__) + parser.add_argument('root',type=Path); parser.add_argument('deadline',type=float) + parser.add_argument('service_deadline',type=float) + args=parser.parse_args(); root=args.root; output=root/'muonclip' + base=Path(__file__).resolve().parents[1] + child=None + report={'status':'starting','long_run_started':False,'training_deadline_unix':args.deadline, + 'per_tensor_diagnostic':False,'scalar_finite_guard':True,'automatic_restart':False} + try: + if args.deadline<=time.time()+60: raise RuntimeError('Too little allocation time remaining.') + output.mkdir() # Refuse reuse/overwrite of any previous run. + cfg=prepare_config(base/'configs/gpt2_small_fineweb_muonclip_long_ww.yaml',root.name) + with (root/'config.yaml').open('x') as f: yaml.safe_dump(cfg,f,sort_keys=False) + from rg_nanogpt_one_head.continuous_support import CloudPublisher + sink=CloudPublisher(os.environ['RG_GPT2_GCS_URI']) + sink.claim({'run_id':root.name,'source_commit':os.environ['RG_GPT2_SOURCE_COMMIT']}) + for name in ('config.yaml','commit.txt','launch.json'): + sink.file(root/name,name) + print('Cloud upload/checksum verification passed; starting one fresh MuonClip process.',flush=True) + cmd=[sys.executable,'-X','faulthandler','-u','-m','rg_gpt2_small.experiment', + '--config',str(root/'config.yaml'),'--data-root','/mnt/disks/rg-data/continuous8/data', + '--output',str(output),'--device','tpu','--allow-long-run', + '--deadline-unix',str(args.deadline)] + child=subprocess.Popen(cmd) + report.update(status='running',long_run_started=True,pid=child.pid) + write(root/'RUN_STATUS.json',report) + rc=watch(child,output,args.deadline,min(args.deadline+300,args.service_deadline-300)) + report['child_exit_code']=rc + if rc: raise RuntimeError(f'MuonClip exited with code {rc}; inspect run.log.') + state=json.loads((output/'status.json').read_text()) + report.update(status='completed' if state['completed'] else 'stopped',training=state) + return 0 + except Exception as exc: + if child is not None: stop(child) + report.update(status='failed_or_incomplete',error=str(exc)) + failure=output/'TPU_PORT_FAILURE.json' + if output.exists() and not failure.exists(): + path=output/'progress.json' + write(failure,{'attribution':'unconfirmed','exception':str(exc), + 'child_exit_code':child.returncode if child is not None else None, + 'last_progress':json.loads(path.read_text()) if path.exists() else None}) + print(json.dumps(report),flush=True) + return 1 + finally: + write(root/'RUN_STATUS.json',report) + + +if __name__=='__main__': sys.exit(main()) diff --git a/baseline/gpt2_small/src/rg_gpt2_small/experiment.py b/baseline/gpt2_small/src/rg_gpt2_small/experiment.py index b54a90f8..d3b62105 100644 --- a/baseline/gpt2_small/src/rg_gpt2_small/experiment.py +++ b/baseline/gpt2_small/src/rg_gpt2_small/experiment.py @@ -132,15 +132,15 @@ def number(key): rt.restore_accelerator_rng_state(states[3], device) -def require_finite_update(model, norm, losses, output, step): +def require_finite_update(model, norm, losses, output, step, verbose=True): """Diagnostic pre-update barrier for validation. Never apply known-invalid gradients.""" - print(json.dumps({'checking_before_update':step}),flush=True) + if verbose: print(json.dumps({'checking_before_update':step}),flush=True) # Materialize the pending XLA graph once before any individual host reads. rt.synchronize(norm.device) value=float(norm.detach().cpu()) loss_values=[float(x.detach().cpu()) for x in losses] if math.isfinite(value) and all(math.isfinite(x) for x in loss_values): - print(json.dumps({'before_update':step,'microbatch_losses':loss_values,'gradient_norm':value}),flush=True) + if verbose: print(json.dumps({'before_update':step,'microbatch_losses':loss_values,'gradient_norm':value}),flush=True) return # Persist the known failure immediately. Copying every gradient to the host # previously stalled here and hid the nonfinite result for over two hours. @@ -182,16 +182,17 @@ def train(cfg, data_root, output, *, device='cpu', resume=False, stop_after=None output = Path(output); output.mkdir(parents=True, exist_ok=True) with (output / 'writer.lock').open('a') as lock: fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB) - diagnostic = cfg.get('validation_gradient_checks', False) + reporting = cfg.get('progress_reporting', False) + diagnostic = cfg.get('validation_gradient_checks', False) or reporting port_checks = cfg.get('validation_tensor_checks', False) - if port_checks: + if port_checks or reporting: port_debug.environment(output) if diagnostic: faulthandler.dump_traceback_later(300, repeat=True) try: return _train(cfg, data_root, output, device, resume, stop_after, deadline) except Exception as exc: - if port_checks: + if port_checks or reporting: port_debug.failure(output,exc) raise finally: @@ -203,6 +204,15 @@ def _train(cfg, data_root, output, device, resume, stop_after, deadline): cfg = copy.deepcopy(cfg) t = cfg['training']; context = cfg['model']['block_size'] port_checks = cfg.get('validation_tensor_checks', False) + def progress(stage, completed_step): + if cfg.get('progress_reporting', False): + atomic_json(output/'progress.json', {'stage':stage, 'completed_step':completed_step, + 'unix_time':time.time()}) + progress('initializing_runtime_and_data', 0) + cloud = None + if cfg.get('cloud_checkpoints', False): + from .run_backup import RunBackup + cloud = RunBackup(output, os.environ['RG_GPT2_GCS_URI']) if port_checks: port_debug.stage(output,'initializing_runtime_and_data',0) step_tokens = t['batch_size'] * t['grad_accum_steps'] * context @@ -251,6 +261,7 @@ def _train(cfg, data_root, output, device, resume, stop_after, deadline): def record(step, final=False): nonlocal elapsed + progress('evaluating', step) metrics = evaluate(model, arrays, cfg, dev) wall = elapsed + time.monotonic() - started identity = {'run_id': cfg['run_id'], 'optimizer': cfg['optimizer']['family'], 'seed': cfg['seed'], @@ -261,11 +272,13 @@ def record(step, final=False): 'end_to_end_tokens_per_second': (step-initial_step) * step_tokens / max(time.monotonic()-started, 1e-9)} pending = {'metrics': row} if cfg['ww']['enabled'] and due(step, cfg['ww']): + progress('weightwatcher', step) measured = measure_ww(model, cfg, identity, metrics) measured['recommended_interval_seconds_for_10pct'] = 9 * measured['seconds'] pending['ww_metrics'] = measured # Checkpoint FIRST includes the pending scalar/WW transaction. On resume, finish missing rows. - save_checkpoint(output / 'checkpoints', { + progress('saving_checkpoint', step) + checkpoint = save_checkpoint(output / 'checkpoints', { 'run_id': cfg['run_id'], 'config': cfg, 'fingerprint': fingerprint, 'model': model.state_dict(), 'optimizers': optimizers.optimizer_state_dict(handles), 'step': step, 'tokens_seen': step*step_tokens, 'scheduler_step': step, 'data_rng': gen.get_state(), 'torch_rng': torch.get_rng_state(), @@ -275,6 +288,10 @@ def record(step, final=False): milestone=step in cfg.get('milestones', [])) for folder, value in pending.items(): append_record(output / folder / f'{step:09d}.json', value) print(json.dumps(row), flush=True) + if cloud: + progress('cloud_backup', step) + cloud.publish(checkpoint, step) + progress('measurement_completed', step) if not resume: record(0) training_window = time.monotonic(); last_timed_step = step @@ -282,6 +299,7 @@ def record(step, final=False): if (stop_after is not None and step >= stop_after) or (deadline and time.time() >= deadline) or (output / 'STOP').exists(): break optimizers.zero_grad(handles) + progress('update_started', step) for handle in handles: lr = optimizers.cosine_learning_rate(step, total_steps=t['schedule_steps'], warmup_steps=t['warmup_steps'], peak_lr=handle.peak_lr, min_lr=handle.min_lr) @@ -305,8 +323,9 @@ def record(step, final=False): [('gradient/'+name,p.grad) for name,p in model.named_parameters() if p.grad is not None] + [('microbatch_loss/'+str(i),loss) for i,loss in enumerate(losses)] + [('global_gradient_norm',norm)],output,'before_clipping',step+1,dev) - if cfg.get('validation_gradient_checks', False): - require_finite_update(model, norm, losses, output, step + 1) + if cfg.get('validation_gradient_checks', False) or cfg.get('finite_update_guard', False): + require_finite_update(model, norm, losses, output, step + 1, + verbose=cfg.get('validation_gradient_checks', False) or step < 4) torch.nn.utils.clip_grad_norm_(model.parameters(), t['grad_clip'], foreach=False) if port_checks: port_debug.check_tensors( @@ -319,10 +338,14 @@ def record(step, final=False): output,'after_optimizer',step+1,dev) port_debug.stage(output,'optimizer_update_completed',step+1) step += 1 + progress('update_submitted', step-1) measurement_due = (step % cfg['metrics_interval'] == 0 or due(step, cfg['ww']) - or step == total or step == stop_after) + or step in cfg.get('metrics_steps', []) or step == total or step == stop_after) if step <= initial_step + 2 or measurement_due or cfg.get('benchmark_sync_every_step', False): rt.synchronize(dev) + progress('update_completed', step) + if cfg.get('progress_reporting', False) and (step <= 4 or step % 25 == 0): + print(json.dumps({'completed_update':step,'tokens_seen':step*step_tokens}),flush=True) seconds = time.monotonic() - training_window if step <= initial_step + 2: compile_seconds += seconds else: steady_seconds += seconds; steady_steps += step - last_timed_step @@ -341,6 +364,7 @@ def record(step, final=False): last_grad = norm.item(); record(step, final=True) atomic_json(output / 'status.json', {'step': step, 'tokens_seen': step*step_tokens, 'completed': step >= total, 'stopped': step < total, 'long_run_launched': cfg.get('long_run', False)}) + progress('stopped' if step < total else 'completed', step) return output diff --git a/baseline/gpt2_small/src/rg_gpt2_small/run_backup.py b/baseline/gpt2_small/src/rg_gpt2_small/run_backup.py new file mode 100644 index 00000000..66098808 --- /dev/null +++ b/baseline/gpt2_small/src/rg_gpt2_small/run_backup.py @@ -0,0 +1,40 @@ +"""Verified rolling cloud checkpoints and all scalar/spectral records for a fresh run.""" +import json +import os +from pathlib import Path + +from rg_nanogpt_one_head.continuous_support import CloudPublisher + + +class RunBackup: + def __init__(self, output, uri, sink=None): + self.output = Path(output) + self.sink = sink or CloudPublisher(uri) + saved = self.output/'cloud_checkpoint.json' + self.sequence = json.loads(saved.read_text())['next_slot_sequence'] if saved.exists() else 0 + self.sent = set() + + def publish(self, checkpoint, step): + # Three slots bound storage without deleting any prior experiment objects. + # Never overwrite the slot referenced by the last successful pointer. + name = f'muonclip/checkpoints/slot_{self.sequence % 3}.pt' + receipt = self.sink.file(checkpoint, name) + pointer = {**receipt, 'step': step, 'file': name, + 'local_file': Path(checkpoint).name, + 'next_slot_sequence': self.sequence+1, + 'note': 'Validate object generation and CRC32C when recovering this rolling checkpoint.'} + if step == 0: + self.sink.file(checkpoint, 'muonclip/checkpoints/initial.pt') + for folder in ('metrics', 'ww_metrics', 'diagnostics'): + for path in sorted((self.output/folder).glob('*.json')): + if path not in self.sent: + self.sink.file(path, 'muonclip/'+path.relative_to(self.output).as_posix()) + self.sent.add(path) + self.sink.file(self.output/'manifest.json', 'muonclip/manifest.json') + self.sink.json(pointer, 'muonclip/checkpoints/LATEST_VERIFIED.json') + tmp = self.output/'cloud_checkpoint.tmp' + tmp.write_text(json.dumps(pointer, indent=2)); tmp.replace(self.output/'cloud_checkpoint.json') + self.sequence += 1 + log = os.environ.get('RG_GPT2_RUN_LOG') + if log and Path(log).is_file(): self.sink.snapshot_text_file(log, 'run.log') + print(f'[cloud-backup] verified step={step}; three rolling slots; all metric records retained', flush=True) diff --git a/baseline/gpt2_small/tests/test_muonclip_night.py b/baseline/gpt2_small/tests/test_muonclip_night.py new file mode 100644 index 00000000..c01b1427 --- /dev/null +++ b/baseline/gpt2_small/tests/test_muonclip_night.py @@ -0,0 +1,130 @@ +import importlib.util +import json +from pathlib import Path +import subprocess + +import pytest +import torch + +from rg_gpt2_small import experiment as g +from rg_gpt2_small.run_backup import RunBackup +from test_gpt2_experiment import config, data + + +def script(name): + spec=importlib.util.spec_from_file_location(name,Path(__file__).parents[1]/'scripts'/f'{name}.py') + module=importlib.util.module_from_spec(spec); spec.loader.exec_module(module); return module + + +def night_config(): + c=config(); c['optimizer']=config('fineweb_muonclip_baseline')['optimizer'] + c.update(validation_tensor_checks=False, validation_gradient_checks=False, + finite_update_guard=True, progress_reporting=True, benchmark_sync_every_step=True, + metrics_interval=2, metrics_steps=[1]) + return c + + +def test_muonclip_trains_with_diagnostic_removed(tmp_path,monkeypatch): + torch.set_num_threads(1) + def forbidden(*args,**kwargs): raise AssertionError('Removed diagnostic was called') + monkeypatch.setattr(g.port_debug,'check_tensors',forbidden) + c=night_config(); out=tmp_path/'run' + g.train(c,data(tmp_path,c),out) + assert json.loads((out/'status.json').read_text())['step']==4 + assert json.loads((out/'progress.json').read_text())['stage']=='completed' + initial=json.loads((out/'metrics/000000000.json').read_text()) + final=json.loads((out/'metrics/000000004.json').read_text()) + assert final['train_nll'] Date: Sun, 4 Oct 2026 07:52:57 -0700 Subject: [PATCH 22/36] Replay saved MuonClip update with bounded stage diagnostics and preserved evidence --- baseline/gpt2_small/README.md | 38 +++- baseline/gpt2_small/TPU_PORT_BUGS.md | 29 +++ .../gpt2_small/scripts/replay_muonclip.py | 140 +++++++++++++ baseline/gpt2_small/scripts/replay_worker.sh | 19 ++ baseline/gpt2_small/scripts/retry_adamw.py | 5 +- baseline/gpt2_small/scripts/run_muonclip.py | 5 +- .../gpt2_small/scripts/supervise_replay.py | 58 ++++++ .../src/rg_gpt2_small/replay_update.py | 194 ++++++++++++++++++ .../gpt2_small/tests/test_replay_update.py | 146 +++++++++++++ 9 files changed, 631 insertions(+), 3 deletions(-) create mode 100644 baseline/gpt2_small/scripts/replay_muonclip.py create mode 100644 baseline/gpt2_small/scripts/replay_worker.sh create mode 100644 baseline/gpt2_small/scripts/supervise_replay.py create mode 100644 baseline/gpt2_small/src/rg_gpt2_small/replay_update.py create mode 100644 baseline/gpt2_small/tests/test_replay_update.py diff --git a/baseline/gpt2_small/README.md b/baseline/gpt2_small/README.md index be6394b7..59c05fa3 100644 --- a/baseline/gpt2_small/README.md +++ b/baseline/gpt2_small/README.md @@ -11,7 +11,43 @@ The TPU launch scripts expose both source packages through `PYTHONPATH`. For loc development, from the repository root, install both with `pip install -e baseline/nanogpt_one_head -e baseline/gpt2_small`. -## Continue with MuonClip on the existing allocation +## Replay the failed MuonClip update first + +The `muonclip-night-20261004-053212` run stopped during evaluation after update 2. +Its step-1 checkpoint and cloud backup are preserved. Do not interpret that run as +passed validation or restart a long run on the strength of its finite initial loss. + +```bash +python3 baseline/gpt2_small/scripts/replay_muonclip.py start +python3 baseline/gpt2_small/scripts/replay_muonclip.py status +``` + +This uses the existing TPU and FineWeb, reads the saved step-1 checkpoint, and +replays exactly one next update into a fresh `muonclip-replay-` directory. +It checks the original config/data/software fingerprint, restores model, both +optimizers, sampler and RNG states, and retains the original LR schedule. It checks +the saved/restored tensors and fixed evaluation probes, gradients before/after +clipping, and weights/moments after primary MuonClip and auxiliary AdamW separately. +Only scalar finite flags cross to CPU; the diagnostic avoids the `stack`/extrema +expression implicated in the earlier diagnostic abort. `FIRST_INVALID.json` names +the first detected invalid stage/tensors. Native aborts and timeouts also get a +separate `TPU_PORT_FAILURE.json`, stage log, environment and available XLA metrics. + +The updated model/optimizer state is saved **before** post-update evaluation to +`diagnostic/update_state.pt`. This is diagnostic evidence, explicitly marked +non-resumable. Existing checkpoints, crash logs, disk and cloud objects are untouched. +The diagnostic is bounded to 30 minutes, plus up to 10 minutes for verified backup, +within the current allocation. No long run, allocation, restart, package install, +data download or cleanup follows automatically. Other launchers reject an active replay. + +CPU tests compare the replay with uninterrupted update 2, including weights, +optimizer states, sampler state and evaluation metrics. TPU equivalence is still +unproven: the extra synchronization changes graph boundaries. A passing replay +requires a subsequent check of the original execution path; it does not establish +that the original numerical failure is fixed. The replay module also accepts +`--device cpu` for a separately requested comparison of the same saved state. + +## Previous continuous MuonClip launcher The user authorized a fresh continuous MuonClip run after the diagnostic crash. Use the current 48-hour TPU, installed environment and preserved FineWeb: diff --git a/baseline/gpt2_small/TPU_PORT_BUGS.md b/baseline/gpt2_small/TPU_PORT_BUGS.md index 551ec606..c6ff3842 100644 --- a/baseline/gpt2_small/TPU_PORT_BUGS.md +++ b/baseline/gpt2_small/TPU_PORT_BUGS.md @@ -109,3 +109,32 @@ a repair or root-cause diagnosis for either SIGABRT or the earlier nonfinite res The earlier source checkouts, diagnostics, initialization checkpoints and verified cloud archives are retained. New failures produce separate evidence in the new run. No automatic restart, cleanup, disk formatting or TPU reallocation is performed. + +### MuonClip update-2 evaluation failure, 2026-10-04 + +Run `muonclip-night-20261004-053212`, commit +`6f8d59214f51be00212a9956182a0c697225f086`, failed about ten minutes after launch; +it did not train overnight. The first update and its evaluation completed. Step 1 +was saved locally and its cloud checkpoint upload was verified. The final cloud +backup was also explicitly verified. + +- Before update 2, four microbatch losses were finite (approximately 10.99) and + the aggregate gradient norm was `14.34039306640625`. +- The log then reported `completed_update: 2`, followed by `RuntimeError: + Nonfinite train NLL` in evaluation. The saved progress stage was `evaluating`, + completed step 2. The latest saved checkpoint was step 1. +- The recorder evaluated before saving the checkpoint, so the failed update's + weights and optimizer state were not preserved. Neither a corrupt update nor + an evaluation/runtime fault has yet been isolated. Finite pre-update losses + and norm do not prove finite updated parameters or optimizer moments. +- Evidence remains under `/mnt/disks/rg-data/gpt2small/muonclip-night-20261004-053212` + and `gs://tpu-builders-504820-ww-continuous8/gpt2small/muonclip-night-20261004-053212`. + +The new `replay_muonclip.py` diagnostic restores the saved step-1 state and exact +next input windows without changing optimizer settings. It checks saved/restored +state, fixed evaluation probes, pre/post-clipping gradients, and state after each +of primary MuonClip and auxiliary AdamW. It saves the resulting diagnostic state +before evaluation. Per-tensor finite reductions use scalar host transfers without +the earlier stack/extrema diagnostic. No attribution to upstream PyTorch/XLA or +TPU hardware is justified yet. These synchronization changes are recorded; a pass +does not by itself reproduce or fix the original continuous execution path. diff --git a/baseline/gpt2_small/scripts/replay_muonclip.py b/baseline/gpt2_small/scripts/replay_muonclip.py new file mode 100644 index 00000000..23b87e8e --- /dev/null +++ b/baseline/gpt2_small/scripts/replay_muonclip.py @@ -0,0 +1,140 @@ +"""Replay one saved MuonClip update on the existing TPU, preserving all evidence.""" +import argparse +import datetime as dt +import fcntl +import json +import os +from pathlib import Path +import re +import shlex +import subprocess +import sys +import time + +PROJECT='tpu-builders-504820'; ZONE='us-west4-a' +QUEUE='ww-gpt2-validation-48h-20261004-s1337'; NODE=QUEUE+'-node' +BASE=Path('/mnt/disks/rg-data/gpt2small'); OLD=BASE/QUEUE +LATEST=BASE/'MUONCLIP_REPLAY_LATEST.json' +SOURCE=BASE/'muonclip-night-20261004-053212' +CHECKPOINT=SOURCE/'muonclip/checkpoints/step_000000001.pt' + + +def run(args,**kwargs): + return subprocess.run(args,check=True,text=True,**kwargs) + + +def active(unit): + result=subprocess.run(['systemctl','show',unit,'--property=ActiveState','--value'], + text=True,capture_output=True,timeout=10) + return result.stdout.strip() in ('active','activating','deactivating','reloading') + + +def assert_idle(): + if active('rg-gpt2-validation.service') or active('rg-continuous8.service'): + raise RuntimeError('An existing training service is active; no training launched.') + if LATEST.exists() and active(json.loads(LATEST.read_text())['unit']): + raise RuntimeError('MuonClip replay is already active. Use status; no second replay launched.') + night=BASE/'PORT_CHECK_LATEST.json' + if night.exists() and active(json.loads(night.read_text())['unit']): + raise RuntimeError('Port diagnostic is active; no concurrent training launched.') + night=BASE/'MUONCLIP_LATEST.json' + if night.exists() and active(json.loads(night.read_text())['unit']): + raise RuntimeError('MuonClip training is active; no concurrent replay launched.') + modules={'rg_gpt2_small.replay_update','rg_gpt2_small.experiment','rg_nanogpt_one_head.gpt2_experiment', + 'rg_nanogpt_one_head.continuous_run','rg_nanogpt_one_head.tpu_spmd_check'} + for path in Path('/proc').glob('[0-9]*/cmdline'): + try: args=set(path.read_bytes().decode().split('\0')) + except (OSError,UnicodeError): continue + if args.intersection(modules): + raise RuntimeError(f'Trainer process {path.parent.name} is active; no concurrent diagnostic launched.') + + +def launch_remote(commit): + if os.geteuid()!=0 or not os.path.ismount('/mnt/disks/rg-data'): + raise RuntimeError('Requires root and the existing mounted data disk.') + if not re.fullmatch('[0-9a-f]{40}',commit): raise ValueError('Expected pinned commit SHA') + with (BASE/'port-check-launch.lock').open('a') as lock: + fcntl.flock(lock,fcntl.LOCK_EX|fcntl.LOCK_NB) + assert_idle() + allocation=json.loads((OLD/'allocation.json').read_text()) + allocation_deadline=float(allocation['validation_deadline_unix']) + if allocation_deadline-time.time()<900: + raise RuntimeError('Less than 15 minutes remain; training not launched.') + if not Path('/mnt/disks/rg-data/continuous8/data/train.bin').is_file(): + raise RuntimeError('Preserved FineWeb is missing; no download will be started.') + if not CHECKPOINT.is_file() or CHECKPOINT.stat().st_size==0: + raise RuntimeError('Verified step-1 checkpoint is missing; replay not started.') + stamp=dt.datetime.now(dt.timezone.utc).strftime('%Y%m%d-%H%M%S') + root=BASE/('muonclip-replay-'+stamp); root.mkdir() + repo=root/'repo'; repo.mkdir() + run(['git','-C',str(repo),'init','-q']) + run(['git','-C',str(repo),'remote','add','origin','https://github.com/CalculatedContent/rg_optimizers.git']) + run(['git','-C',str(repo),'fetch','--depth','1','origin',commit],timeout=180) + run(['git','-C',str(repo),'checkout','--detach',commit]) + (root/'commit.txt').write_text(commit+'\n') + deadline=min(time.time()+1800,allocation_deadline-600) + if deadline-time.time()<300: raise RuntimeError('Too little time remains after source checkout.') + unit='rg-gpt2-replay-'+stamp+'.service' + record={'root':str(root),'unit':unit,'commit':commit,'node':NODE, + 'training_deadline_unix':deadline,'service_deadline_unix':min(deadline+600,allocation_deadline), + 'source_checkpoint':str(CHECKPOINT),'source_run':str(SOURCE), + 'purpose':'replay saved update 2 with checks after clipping, primary MuonClip and auxiliary AdamW', + 'cloud_uri':'gs://tpu-builders-504820-ww-continuous8/gpt2small/'+root.name} + (root/'launch.json').write_text(json.dumps(record,indent=2)) + command=['systemd-run','--unit='+unit,'--property=Type=exec','--property=Restart=no', + '--property=RuntimeMaxSec='+str(int(min(deadline+600,allocation_deadline)-time.time())), + '--property=TimeoutStopSec=30','--property=KillMode=control-group', + '--property=StandardOutput=append:'+str(root/'run.log'), + '--property=StandardError=append:'+str(root/'run.log'), + '/bin/bash',str(repo/'baseline/gpt2_small/scripts/replay_worker.sh'), + str(root),str(deadline),str(CHECKPOINT)] + run(command) + temp=LATEST.with_suffix('.tmp'); temp.write_text(json.dumps(record,indent=2)); temp.replace(LATEST) + print('One-update replay service started:',unit,flush=True) + print('Log:',root/'run.log',flush=True) + print('Training cutoff UTC:',dt.datetime.fromtimestamp(deadline,dt.timezone.utc).isoformat(),flush=True) + print('One update only. Diagnostic capped at 30 minutes plus up to 10 minutes for backup.',flush=True) + print('Existing TPU allocation, FineWeb and prior outputs retained.',flush=True) + + +def status_remote(): + if not LATEST.exists(): + print('No update replay has been launched.'); return + record=json.loads(LATEST.read_text()); root=Path(record['root']) + print(json.dumps(record,indent=2),flush=True) + subprocess.run(['systemctl','--no-pager','--full','status',record['unit']],check=False) + for file in ('PROBE_STATUS.json','diagnostic/REPLAY_STATUS.json','diagnostic/FIRST_INVALID.json','diagnostic/TPU_PORT_FAILURE.json','diagnostic/diagnostics/current_stage.json'): + path=root/file + if path.is_file(): print('\n'+file+'\n'+path.read_text(),flush=True) + subprocess.run(['tail','-n','70',str(root/'run.log')],check=False) + + +def main(): + parser=argparse.ArgumentParser(description=__doc__) + parser.add_argument('action',choices=('start','status')) + parser.add_argument('--on-tpu',action='store_true',help=argparse.SUPPRESS) + parser.add_argument('--commit',help=argparse.SUPPRESS) + args=parser.parse_args() + if args.on_tpu: + launch_remote(args.commit) if args.action=='start' else status_remote() + return 0 + commit=None + if args.action=='start': + repo=Path(__file__).resolve().parents[3] + dirty=run(['git','-C',str(repo),'status','--porcelain'],capture_output=True).stdout + if dirty.strip(): raise RuntimeError('Use a clean checkout of the pushed commit.') + commit=run(['git','-C',str(repo),'rev-parse','HEAD'],capture_output=True).stdout.strip() + queue=json.loads(run(['gcloud','alpha','compute','tpus','queued-resources','describe',QUEUE, + '--project='+PROJECT,'--zone='+ZONE,'--format=json'],capture_output=True).stdout) + if queue.get('state',{}).get('state')!='ACTIVE': + raise RuntimeError('The existing TPU is not ACTIVE; no allocation requested.') + remote=['sudo','python3','-c',Path(__file__).read_text(),args.action,'--on-tpu'] + if commit: remote+=['--commit',commit] + return subprocess.run(['gcloud','compute','tpus','tpu-vm','ssh',NODE, + '--project='+PROJECT,'--zone='+ZONE,'--worker=0','--command='+shlex.join(remote)]).returncode + + +if __name__=='__main__': + try: sys.exit(main()) + except Exception as exc: + print('MuonClip replay:',exc,file=sys.stderr); sys.exit(1) diff --git a/baseline/gpt2_small/scripts/replay_worker.sh b/baseline/gpt2_small/scripts/replay_worker.sh new file mode 100644 index 00000000..3c97d2a9 --- /dev/null +++ b/baseline/gpt2_small/scripts/replay_worker.sh @@ -0,0 +1,19 @@ +#!/usr/bin/env bash +set -Eeuo pipefail +root=$1 +deadline=$2 +checkpoint=$3 +base=$(cd "$(dirname "$0")/.." && pwd) +python=/mnt/disks/rg-data/continuous8/venv/bin/python +export PYTHONPATH="$base/src:$base/../nanogpt_one_head/src" +export RG_GPT2_SOURCE_COMMIT=$(git -C "$base" rev-parse HEAD) +source "$base/scripts/tpu_environment.sh" +finish() { + rc=$? + trap - EXIT + sync + "$python" "$base/scripts/backup.py" "$root" || exit 1 + exit "$rc" +} +trap finish EXIT +"$python" -u "$base/scripts/supervise_replay.py" "$root" "$deadline" "$checkpoint" diff --git a/baseline/gpt2_small/scripts/retry_adamw.py b/baseline/gpt2_small/scripts/retry_adamw.py index ac566c6a..b955ef54 100644 --- a/baseline/gpt2_small/scripts/retry_adamw.py +++ b/baseline/gpt2_small/scripts/retry_adamw.py @@ -35,7 +35,10 @@ def assert_idle(): night=BASE/'MUONCLIP_LATEST.json' if night.exists() and active(json.loads(night.read_text())['unit']): raise RuntimeError('MuonClip service is active; no concurrent diagnostic launched.') - modules={'rg_gpt2_small.experiment','rg_nanogpt_one_head.gpt2_experiment', + replay=BASE/'MUONCLIP_REPLAY_LATEST.json' + if replay.exists() and active(json.loads(replay.read_text())['unit']): + raise RuntimeError('An update replay is active; no concurrent trainer launched.') + modules={'rg_gpt2_small.replay_update','rg_gpt2_small.experiment','rg_nanogpt_one_head.gpt2_experiment', 'rg_nanogpt_one_head.continuous_run','rg_nanogpt_one_head.tpu_spmd_check'} for path in Path('/proc').glob('[0-9]*/cmdline'): try: args=set(path.read_bytes().decode().split('\0')) diff --git a/baseline/gpt2_small/scripts/run_muonclip.py b/baseline/gpt2_small/scripts/run_muonclip.py index 90205d03..cc8fc3f2 100644 --- a/baseline/gpt2_small/scripts/run_muonclip.py +++ b/baseline/gpt2_small/scripts/run_muonclip.py @@ -35,7 +35,10 @@ def assert_idle(): night=BASE/'PORT_CHECK_LATEST.json' if night.exists() and active(json.loads(night.read_text())['unit']): raise RuntimeError('Port diagnostic is active; no concurrent training launched.') - modules={'rg_gpt2_small.experiment','rg_nanogpt_one_head.gpt2_experiment', + replay=BASE/'MUONCLIP_REPLAY_LATEST.json' + if replay.exists() and active(json.loads(replay.read_text())['unit']): + raise RuntimeError('An update replay is active; no concurrent trainer launched.') + modules={'rg_gpt2_small.replay_update','rg_gpt2_small.experiment','rg_nanogpt_one_head.gpt2_experiment', 'rg_nanogpt_one_head.continuous_run','rg_nanogpt_one_head.tpu_spmd_check'} for path in Path('/proc').glob('[0-9]*/cmdline'): try: args=set(path.read_bytes().decode().split('\0')) diff --git a/baseline/gpt2_small/scripts/supervise_replay.py b/baseline/gpt2_small/scripts/supervise_replay.py new file mode 100644 index 00000000..d60e6153 --- /dev/null +++ b/baseline/gpt2_small/scripts/supervise_replay.py @@ -0,0 +1,58 @@ +"""Bounded, single-update child. No automatic long run or restart.""" +import argparse +import json +from pathlib import Path +import subprocess +import sys +import time + +from validation_watchdog import wait_for_child + + +def preserve_child_failure(root, error, report): + """Native aborts/timeouts cannot execute the child's Python exception handler.""" + from rg_gpt2_small.port_debug import write + output=root/'diagnostic'; path=output/'TPU_PORT_FAILURE.json' + if path.exists(): return + stage=output/'diagnostics/current_stage.json' + write(path,{'status':'process_failed_or_timed_out','attribution':'unconfirmed', + 'exception':str(error),'child_exit_code':report.get('child_exit_code'), + 'last_stage':json.loads(stage.read_text()) if stage.exists() else None, + 'evidence':['../run.log','diagnostics/','REPLAY_SOURCE.json','../commit.txt']}) + + +def main(): + parser=argparse.ArgumentParser(description=__doc__) + parser.add_argument('root',type=Path); parser.add_argument('deadline',type=float) + parser.add_argument('checkpoint',type=Path) + args=parser.parse_args(); root=args.root + report={'status':'starting','deadline_unix':args.deadline,'automatic_long_run':False} + try: + if args.deadline<=time.time()+60: raise RuntimeError('Insufficient time for replay.') + from rg_nanogpt_one_head.continuous_support import CloudPublisher + sink=CloudPublisher('gs://tpu-builders-504820-ww-continuous8/gpt2small/'+root.name) + sink.claim({'run_id':root.name,'source_checkpoint':str(args.checkpoint)}) + for name in ('commit.txt','launch.json'): sink.file(root/name,name) + cmd=[sys.executable,'-X','faulthandler','-u','-m','rg_gpt2_small.replay_update', + '--checkpoint',str(args.checkpoint),'--data-root','/mnt/disks/rg-data/continuous8/data', + '--output',str(root/'diagnostic'),'--device','tpu'] + child=subprocess.Popen(cmd) + report.update(status='running',pid=child.pid) + (root/'PROBE_STATUS.json').write_text(json.dumps(report,indent=2)) + rc=wait_for_child(child,'MuonClip update-2 replay',args.deadline) + report['child_exit_code']=rc + if rc: raise RuntimeError(f'Replay exited with code {rc}; see diagnostic/FIRST_INVALID.json and run.log.') + state=json.loads((root/'diagnostic/REPLAY_STATUS.json').read_text()) + if state['status']!='one_update_passed': raise RuntimeError('Replay completion not confirmed.') + report.update(status='one_update_passed') + return 0 + except Exception as exc: + report.update(status='failed_or_incomplete',error=str(exc)) + preserve_child_failure(root,exc,report) + print(json.dumps(report),flush=True) + return 1 + finally: + (root/'PROBE_STATUS.json').write_text(json.dumps(report,indent=2)) + + +if __name__=='__main__': sys.exit(main()) diff --git a/baseline/gpt2_small/src/rg_gpt2_small/replay_update.py b/baseline/gpt2_small/src/rg_gpt2_small/replay_update.py new file mode 100644 index 00000000..4f32cd8d --- /dev/null +++ b/baseline/gpt2_small/src/rg_gpt2_small/replay_update.py @@ -0,0 +1,194 @@ +"""Replay one saved MuonClip update, locating the first invalid stage. + +Diagnostic synchronization changes graph boundaries. A passing replay is evidence +about this instrumented path, not proof that the original continuous path is fixed. +""" +import argparse +import faulthandler +import hashlib +import json +import os +from pathlib import Path +import random + +import numpy as np +import torch + +from . import experiment as g, port_debug as debug +from rg_nanogpt_one_head import runtime as rt, tpu_spmd as spmd, optimizers + + +def sha256(path): + digest=hashlib.sha256() + with Path(path).open('rb') as stream: + for block in iter(lambda:stream.read(8*1024*1024),b''): digest.update(block) + return digest.hexdigest() + + +def check_finite(named, output, label, update, device): + """No stack/cat, extrema, or full-tensor host copies: materialize scalar flags.""" + debug.stage(output,label+'_started',update) + checks=[] + for name,tensor in named: + value=tensor.detach() + checks.append((name,tuple(value.shape),str(value.dtype),torch.isfinite(value).all(), + (value<0).any() if name.endswith('/exp_avg_sq') else None)) + rt.synchronize(torch.device(device)) + records=[] + for name,shape,dtype,finite,negative in checks: + records.append({'tensor':name,'shape':shape,'dtype':dtype, + 'all_finite':bool(finite.cpu().item()), + 'negative_second_moment':bool(negative.cpu().item()) if negative is not None else False}) + bad=[row for row in records if not row['all_finite'] or row['negative_second_moment']] + report={'stage':label,'update':update,'records':records,'invalid_tensors':bad} + debug.write(Path(output)/'diagnostics'/f'{update:06d}-{label}.json',report) + debug.xla_metrics(output,f'{update:06d}-{label}',device) + if bad: + debug.write(Path(output)/'FIRST_INVALID.json',report) + raise RuntimeError(f'First invalid stage: {label}; tensors: '+', '.join(row['tensor'] for row in bad[:8])) + debug.stage(output,label+'_passed',update) + + +def source_tensors(state): + yield from (('weight/'+name,tensor) for name,tensor in state['model'].items()) + for index,optimizer in enumerate(state['optimizers']): + for parameter,values in optimizer['state'].items(): + for key,value in values.items(): + if torch.is_tensor(value): yield f'optimizer_{index}/{parameter}/{key}',value + + +@torch.no_grad() +def evaluate_splits(model, arrays, cfg, device, output, update, label): + model.eval(); result={} + try: + for j,split in enumerate(('train','val','test')): + generator=torch.Generator().manual_seed(cfg['seed']+20000+j) + losses=[]; accuracies=[] + for index in range(cfg['eval_batches']): + stage=f'{label}_{split}_batch_{index}' + debug.stage(output,stage+'_forward',update) + offsets=[] + x,y=g.batch(arrays[split],generator,cfg['training']['batch_size'], + cfg['model']['block_size'],device,trace=offsets) + debug.write(Path(output)/'diagnostics'/f'{update:06d}-{stage}-inputs.json', + {'split':split,'context':cfg['model']['block_size'],'offsets':offsets}) + logits,loss=model(x,y) + accuracy=(logits.argmax(-1)==y).float().mean() + rt.mark_step(device) + check_finite([('logits',logits),('loss',loss)],output,stage,update,device) + losses.append(float(loss.cpu())); accuracies.append(float(accuracy.cpu())) + result.update({split+'_nll':sum(losses)/len(losses), + split+'_token_error':1-sum(accuracies)/len(accuracies)}) + debug.write(Path(output)/(label+'.json'),result) + print(json.dumps({'evaluation':label,'split':split,**result}),flush=True) + return result + finally: + model.train() + + +def replay(checkpoint, data_root, output, device='tpu', expected_step=1): + checkpoint=Path(checkpoint); output=Path(output) + output.mkdir(parents=True,exist_ok=False) + model=None; handles=None; state=None; gen=None; completed_roles=[] + report={'status':'running','original_checkpoint':str(checkpoint),'device':device, + 'automatic_long_run':False,'instrumented_graph_boundaries':True} + debug.environment(output) + try: + debug.stage(output,'loading_source_checkpoint',expected_step) + report['source_sha256']=sha256(checkpoint) + state=torch.load(checkpoint,map_location='cpu',weights_only=False) + cfg=state['config']; start=int(state['step']); update=start+1 + if start!=expected_step or state['scheduler_step']!=start: + raise RuntimeError('Unexpected source checkpoint step/scheduler.') + step_tokens=cfg['training']['batch_size']*cfg['training']['grad_accum_steps']*cfg['model']['block_size'] + if state['tokens_seen']!=start*step_tokens: raise RuntimeError('Checkpoint token count differs.') + if cfg['optimizer']['family']!='muon_clip': raise RuntimeError('Expected a MuonClip checkpoint.') + report.update(source_step=start,replayed_update=update,config=cfg) + debug.write(output/'REPLAY_SOURCE.json',report) + check_finite(source_tensors(state),output,'saved_cpu_state',start,'cpu') + # Validate the ORIGINAL fingerprint before selecting a CPU diagnostic device. + metadata,arrays=g.load_memmaps(data_root,cfg) + versions={'torch':torch.__version__,'numpy':np.__version__} + fingerprint=hashlib.sha256(json.dumps({'config':cfg,'data':metadata,'versions':versions},sort_keys=True).encode()).hexdigest() + if fingerprint!=state['fingerprint']: raise RuntimeError('Source config/data/version fingerprint mismatch.') + runtime_cfg=json.loads(json.dumps(cfg)) + if device=='cpu': runtime_cfg['runtime']['tpu_spmd']=False + spmd.initialize(runtime_cfg,device); dev=rt.choose_device(device); rt.configure_runtime(dev,runtime_cfg) + rt.seed_everything(cfg['seed'],dev) + model=g.GPT(g.GPTConfig(**cfg['model'])).to(dev); model.load_state_dict(state['model']) + spmd.replicate_model(model); handles=g.make_handles(model,cfg) + optimizers.load_optimizer_state_dict(handles,state['optimizers']) + gen=torch.Generator(); gen.set_state(state['data_rng']) + torch.set_rng_state(state['torch_rng']); random.setstate(state['python_rng']); np.random.set_state(state['numpy_rng']) + if device!='cpu': rt.restore_accelerator_rng_state(state['accelerator_rng'],dev) + debug.write(output/'runtime.json',rt.runtime_metadata(dev)) + check_finite(debug.optimizer_tensors(model,handles),output,'restored_state',start,dev) + evaluate_splits(model,arrays,cfg,dev,output,start,'before_update_evaluation') + # Restore saved RNG streams after the comparison probe; exact next training windows. + torch.set_rng_state(state['torch_rng']); random.setstate(state['python_rng']); np.random.set_state(state['numpy_rng']) + if device!='cpu': rt.restore_accelerator_rng_state(state['accelerator_rng'],dev) + optimizers.zero_grad(handles); t=cfg['training'] + for handle in handles: + handle.set_lr(optimizers.cosine_learning_rate(start,total_steps=t['schedule_steps'], + warmup_steps=t['warmup_steps'],peak_lr=handle.peak_lr,min_lr=handle.min_lr)) + debug.write(output/'learning_rates.json',{h.role:h.lr for h in handles}) + debug.stage(output,'forward_backward',update) + losses=[]; offsets=[] + for _ in range(t['grad_accum_steps']): + x,y=g.batch(arrays['train'],gen,t['batch_size'],cfg['model']['block_size'],dev,trace=offsets) + _,loss=model(x,y); losses.append(loss.detach()); (loss/t['grad_accum_steps']).backward() + debug.write(output/'training_inputs.json',{'update':update,'offsets':offsets,'context':cfg['model']['block_size']}) + spmd.replicate_gradients(model) + norm=rt.gradient_norm(model.parameters()) + check_finite([('gradient/'+n,p.grad) for n,p in model.named_parameters() if p.grad is not None] + +[('loss/'+str(i),loss) for i,loss in enumerate(losses)]+[('gradient_norm',norm)], + output,'before_clipping',update,dev) + g.require_finite_update(model,norm,losses,output,update) + torch.nn.utils.clip_grad_norm_(model.parameters(),t['grad_clip'],foreach=False) + check_finite([('gradient/'+n,p.grad) for n,p in model.named_parameters() if p.grad is not None], + output,'after_clipping',update,dev) + for handle in handles: + debug.stage(output,'applying_'+handle.role,update) + handle.optimizer.step(); rt.mark_step(dev); rt.synchronize(dev) + completed_roles.append(handle.role) + check_finite(debug.optimizer_tensors(model,handles),output,'after_'+handle.role,update,dev) + # Capture the updated state BEFORE evaluation can fail. Original checkpoint is read-only. + snapshot(output,model,handles,cfg,gen,update,completed_roles) + result=evaluate_splits(model,arrays,cfg,dev,output,update,'after_update_evaluation') + report.update(status='one_update_passed',completed_roles=completed_roles,metrics=result, + interpretation='Instrumented replay only. Original TPU path and long-run stability remain unproven.') + return output + except Exception as exc: + report.update(status='failed',error=str(exc),completed_roles=completed_roles) + debug.failure(output,exc) # Save diagnosis before attempting a potentially slow state copy. + if model is not None and handles is not None and state is not None and gen is not None and completed_roles: + try: + if not (output/'update_state.pt').exists(): + snapshot(output,model,handles,state['config'],gen,int(state['step'])+1,completed_roles) + except Exception as snapshot_error: report['snapshot_error']=str(snapshot_error) + raise + finally: + debug.write(output/'REPLAY_STATUS.json',report) + + +def snapshot(output,model,handles,cfg,gen,update,completed_roles): + debug.stage(output,'saving_diagnostic_state',update) + path=Path(output)/'update_state.pt'; tmp=path.with_suffix('.tmp') + payload={'diagnostic_only':True,'resumable':False,'update':update,'completed_roles':list(completed_roles), + 'config':cfg,'model':model.state_dict(),'optimizers':optimizers.optimizer_state_dict(handles), + 'next_data_rng':gen.get_state()} + with tmp.open('wb') as stream: + torch.save(rt.tree_to_cpu(payload),stream); stream.flush(); os.fsync(stream.fileno()) + tmp.replace(path) + + +def main(): + parser=argparse.ArgumentParser(description=__doc__) + parser.add_argument('--checkpoint',required=True); parser.add_argument('--data-root',required=True) + parser.add_argument('--output',required=True); parser.add_argument('--device',default='tpu',choices=('tpu','cpu')) + args=parser.parse_args(); faulthandler.dump_traceback_later(300,repeat=True) + try: replay(args.checkpoint,args.data_root,args.output,args.device) + finally: faulthandler.cancel_dump_traceback_later() + + +if __name__=='__main__': main() diff --git a/baseline/gpt2_small/tests/test_replay_update.py b/baseline/gpt2_small/tests/test_replay_update.py new file mode 100644 index 00000000..7c5a7abd --- /dev/null +++ b/baseline/gpt2_small/tests/test_replay_update.py @@ -0,0 +1,146 @@ +import importlib.util +import json +from pathlib import Path +import sys + +import pytest +import torch + +from rg_gpt2_small import experiment as g, replay_update as replay +from rg_nanogpt_one_head.checkpoints import optimizer_state_sha256 +from rg_nanogpt_one_head.muonclip import MuonClip +from test_gpt2_experiment import config, data + + +def script(name): + folder=Path(__file__).parents[1]/'scripts' + sys.path.insert(0,str(folder)) + try: + spec=importlib.util.spec_from_file_location(name,folder/(name+'.py')) + module=importlib.util.module_from_spec(spec); spec.loader.exec_module(module) + return module + finally: sys.path.pop(0) + + +@pytest.fixture +def saved_step(tmp_path): + torch.set_num_threads(1) + c=config(); c['optimizer']=config('fineweb_muonclip_baseline')['optimizer'] + c['metrics_steps']=[1]; c['finite_update_guard']=True + d=data(tmp_path,c); out=tmp_path/'original' + g.train(c,d,out,stop_after=1) + return c,d,out/'checkpoints/step_000000001.pt' + + +def test_scalar_checks_do_not_stack_or_copy_arrays(tmp_path,monkeypatch): + def forbidden(*args,**kwargs): raise AssertionError('stack/cat forbidden in diagnostic') + monkeypatch.setattr(torch,'stack',forbidden); monkeypatch.setattr(torch,'cat',forbidden) + original=torch.Tensor.cpu + def scalar_cpu(tensor,*args,**kwargs): + assert tensor.numel()==1 + return original(tensor,*args,**kwargs) + monkeypatch.setattr(torch.Tensor,'cpu',scalar_cpu) + tensors=[('weight/test',torch.ones(4,8)),('aux/test/exp_avg_sq',torch.tensor([-1.,0.]))] + with pytest.raises(RuntimeError,match='First invalid stage: moments'): + replay.check_finite(tensors,tmp_path,'moments',2,'cpu') + bad=json.loads((tmp_path/'FIRST_INVALID.json').read_text()) + assert bad['invalid_tensors'][0]['negative_second_moment'] is True + assert bad['invalid_tensors'][0]['tensor']=='aux/test/exp_avg_sq' + + +def test_exact_cpu_update_replay_and_original_preserved(saved_step,tmp_path): + c,d,checkpoint=saved_step; before=replay.sha256(checkpoint) + reference=tmp_path/'reference'; g.train(c,d,reference,stop_after=2) + output=tmp_path/'replay'; replay.replay(checkpoint,d,output,'cpu') + expected=torch.load(reference/'checkpoints/step_000000002.pt',weights_only=False) + actual=torch.load(output/'update_state.pt',weights_only=False) + for name in expected['model']: + torch.testing.assert_close(actual['model'][name],expected['model'][name],rtol=0,atol=0) + assert optimizer_state_sha256(actual['optimizers'])==optimizer_state_sha256(expected['optimizers']) + assert torch.equal(actual['next_data_rng'],expected['data_rng']) + assert actual['diagnostic_only'] and not actual['resumable'] + assert actual['completed_roles']==['primary','auxiliary'] + assert replay.sha256(checkpoint)==before + metrics=json.loads((output/'after_update_evaluation.json').read_text()) + original=json.loads((reference/'metrics/000000002.json').read_text()) + for key in metrics: assert metrics[key]==original[key] + assert json.loads((output/'REPLAY_STATUS.json').read_text())['status']=='one_update_passed' + + +def test_invalid_primary_stops_before_auxiliary(saved_step,tmp_path,monkeypatch): + _,d,checkpoint=saved_step; before=replay.sha256(checkpoint) + original=MuonClip.step + def corrupt(self,*args,**kwargs): + result=original(self,*args,**kwargs) + with torch.no_grad(): self.param_groups[0]['params'][0].view(-1)[0]=float('nan') + return result + monkeypatch.setattr(MuonClip,'step',corrupt) + def forbidden(*args,**kwargs): raise AssertionError('Auxiliary applied after invalid primary') + monkeypatch.setattr(torch.optim.AdamW,'step',forbidden) + output=tmp_path/'bad' + with pytest.raises(RuntimeError,match='First invalid stage: after_primary'): + replay.replay(checkpoint,d,output,'cpu') + assert json.loads((output/'FIRST_INVALID.json').read_text())['stage']=='after_primary' + state=torch.load(output/'update_state.pt',weights_only=False) + assert state['completed_roles']==['primary'] + assert (output/'TPU_PORT_FAILURE.json').exists() + assert replay.sha256(checkpoint)==before + + +def test_postupdate_state_survives_evaluation_failure(saved_step,tmp_path,monkeypatch): + _,d,checkpoint=saved_step; evaluate=replay.evaluate_splits + def broken(*args,**kwargs): + if args[-1]=='after_update_evaluation': raise RuntimeError('Injected evaluation failure') + return evaluate(*args,**kwargs) + monkeypatch.setattr(replay,'evaluate_splits',broken) + output=tmp_path/'eval-failure' + with pytest.raises(RuntimeError,match='Injected evaluation failure'): + replay.replay(checkpoint,d,output,'cpu') + state=torch.load(output/'update_state.pt',weights_only=False) + assert state['completed_roles']==['primary','auxiliary'] + assert all(torch.isfinite(value).all() for value in state['model'].values()) + + +def test_launch_is_bounded_and_preserves_original(tmp_path,monkeypatch): + module=script('replay_muonclip'); old=tmp_path/'allocation'; old.mkdir() + (old/'allocation.json').write_text(json.dumps({'validation_deadline_unix':100000.})) + source=tmp_path/'original'; source.mkdir(); checkpoint=source/'step1.pt'; checkpoint.write_bytes(b'evidence') + monkeypatch.setattr(module,'BASE',tmp_path); monkeypatch.setattr(module,'OLD',old) + monkeypatch.setattr(module,'SOURCE',source); monkeypatch.setattr(module,'CHECKPOINT',checkpoint) + monkeypatch.setattr(module,'LATEST',tmp_path/'latest.json') + monkeypatch.setattr(module.os,'geteuid',lambda:0) + monkeypatch.setattr(module.os.path,'ismount',lambda _:True) + monkeypatch.setattr(module,'assert_idle',lambda:None) + original=Path.is_file + monkeypatch.setattr(Path,'is_file',lambda p:True if str(p).endswith('/data/train.bin') else original(p)) + monkeypatch.setattr(module.time,'time',lambda:1000.) + calls=[]; monkeypatch.setattr(module,'run',lambda args,**kwargs:calls.append(args)) + module.launch_remote('a'*40) + record=json.loads(module.LATEST.read_text()) + assert record['training_deadline_unix']==2800. and record['service_deadline_unix']==3400. + command=next(c for c in calls if c[0]=='systemd-run') + assert '--property=RuntimeMaxSec=2400' in command and '--property=Restart=no' in command + assert command[-1]==str(checkpoint) and checkpoint.read_bytes()==b'evidence' + assert not any(c[0]=='gcloud' for c in calls) + + +@pytest.mark.parametrize('name',['run_muonclip','retry_adamw','replay_muonclip']) +def test_all_launchers_refuse_active_replay(tmp_path,monkeypatch,name): + module=script(name) + monkeypatch.setattr(module,'BASE',tmp_path) + monkeypatch.setattr(module,'LATEST',tmp_path/'absent.json') + (tmp_path/'MUONCLIP_REPLAY_LATEST.json').write_text(json.dumps({'unit':'replay.service'})) + if name=='replay_muonclip': monkeypatch.setattr(module,'LATEST',tmp_path/'MUONCLIP_REPLAY_LATEST.json') + monkeypatch.setattr(module,'active',lambda unit:unit=='replay.service') + with pytest.raises(RuntimeError,match='[Rr]eplay'): module.assert_idle() + + +def test_supervisor_preserves_native_abort_and_child_report(tmp_path): + module=script('supervise_replay'); stage=tmp_path/'diagnostic/diagnostics'; stage.mkdir(parents=True) + (stage/'current_stage.json').write_text(json.dumps({'stage':'after_primary_started','update':2})) + module.preserve_child_failure(tmp_path,RuntimeError('aborted'),{'child_exit_code':-6}) + path=tmp_path/'diagnostic/TPU_PORT_FAILURE.json'; raw=path.read_bytes() + assert json.loads(raw)['last_stage']['stage']=='after_primary_started' + assert json.loads(raw)['child_exit_code']==-6 + module.preserve_child_failure(tmp_path,RuntimeError('later error'),{}) + assert path.read_bytes()==raw From a7f0e44fcd2097e609dbdbb3dc4ea17c9bcd3a14 Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Sun, 4 Oct 2026 10:11:09 -0700 Subject: [PATCH 23/36] Run continuous MuonClip with replay synchronization and recoverable measurements Reuse the successful replay's scalar finite checks and synchronized optimizer stages in a fresh uninterrupted run on the existing TPU. Save a full checkpoint before evaluation and complete pending measurements on explicit resume. Retain token-error measurements every 25 updates and all layer spectra at step 25 and every 100 updates. Keep the current PyTorch/XLA environment and existing data/evidence. Verify numerical equivalence, interrupted measurement recovery, invalid-update handling, and mutable diagnostic backups with 57 passing CPU tests. Continuous TPU stability still requires the live run. --- baseline/gpt2_small/README.md | 48 +++++-- baseline/gpt2_small/TPU_PORT_BUGS.md | 20 +++ baseline/gpt2_small/scripts/run_muonclip.py | 4 +- .../gpt2_small/scripts/supervise_muonclip.py | 5 +- .../src/rg_gpt2_small/execution_checks.py | 67 ++++++++++ .../src/rg_gpt2_small/experiment.py | 72 +++++++--- .../src/rg_gpt2_small/replay_update.py | 52 +------- .../src/rg_gpt2_small/run_backup.py | 3 +- .../tests/test_continuous_synced.py | 124 ++++++++++++++++++ .../gpt2_small/tests/test_muonclip_night.py | 1 + 10 files changed, 317 insertions(+), 79 deletions(-) create mode 100644 baseline/gpt2_small/src/rg_gpt2_small/execution_checks.py create mode 100644 baseline/gpt2_small/tests/test_continuous_synced.py diff --git a/baseline/gpt2_small/README.md b/baseline/gpt2_small/README.md index 59c05fa3..ca09ac8f 100644 --- a/baseline/gpt2_small/README.md +++ b/baseline/gpt2_small/README.md @@ -11,7 +11,38 @@ The TPU launch scripts expose both source packages through `PYTHONPATH`. For loc development, from the repository root, install both with `pip install -e baseline/nanogpt_one_head -e baseline/gpt2_small`. -## Replay the failed MuonClip update first +## Continuous run using the successful replay's numerical checks + +The instrumented TPU replay `muonclip-replay-20261004-150734` passed update 2 and +all three evaluation splits. Both optimizer stages and gradient checks passed. +Training NLL decreased from 10.998178 to 10.982910 and test NLL from 10.984691 to +10.969128. Its cloud backup was verified. This is one-update evidence; the original +failure's cause and long-run stability remain unresolved. + +The user authorized a continuous run on the installed PyTorch/XLA environment while +waiting for access to Google's TorchTPU. The launcher now enables the shared +`execution_checks` path used by the passing replay: scalar finite reductions before +and after clipping, separate synchronized MuonClip and auxiliary AdamW steps, checks +of updated weights/moments, and synchronized evaluation with scalar host averaging. +The older `stack`/extrema per-tensor diagnostic stays disabled. Model, corpus, input +windows, optimizer formulas and learning-rate settings are unchanged. + +Before every scheduled evaluation, a full local checkpoint is committed with +`measurement_pending=true`. Successful evaluation/spectra replace it with a full +checkpoint containing the pending immutable result records, then publish verified +cloud uploads. If evaluation or spectra fails, the pre-evaluation checkpoint remains +on disk and is included in the worker's exit backup. An explicit recovery completes +that step's measurement before training further; there is no automatic restart. +Three rolling local/cloud checkpoints are retained plus initialization, and every +scalar/spectral result is retained. Routine finite-check reports retain early and +measurement steps plus current rolling reports; failures always retain their report. + +The checks add overhead and are retained throughout this run. A 25-update tiny-model +CPU test matches the original CPU optimizer trajectory exactly, and injected +evaluation/spectral failures verify checkpoint recovery and unchanged old metrics. +The full 124M continuous path still needs the live TPU run to establish stability. + +## Replaying the earlier failed update The `muonclip-night-20261004-053212` run stopped during evaluation after update 2. Its step-1 checkpoint and cloud backup are preserved. Do not interpret that run as @@ -47,9 +78,9 @@ requires a subsequent check of the original execution path; it does not establis that the original numerical failure is fixed. The replay module also accepts `--device cpu` for a separately requested comparison of the same saved state. -## Previous continuous MuonClip launcher +## Launch on the existing TPU -The user authorized a fresh continuous MuonClip run after the diagnostic crash. +The user authorized a fresh continuous MuonClip run after the successful replay. Use the current 48-hour TPU, installed environment and preserved FineWeb: ```bash @@ -58,17 +89,16 @@ python3 baseline/gpt2_small/scripts/run_muonclip.py status ``` This starts one fresh process under `rg-gpt2-muonclip-.service` with -`Restart=no`, in `/mnt/disks/rg-data/gpt2small/muonclip-night-`. +`Restart=no`, in `/mnt/disks/rg-data/gpt2small/muonclip-continuous-`. It never invokes the cleanup/reallocation scripts, changes the installed packages, downloads the corpus, or writes into previous experiment directories/cloud prefixes. The saved `port-check-20261004-045633` crash evidence remains intact. A shared launch lock and checks for existing services/trainers prevent simultaneous TPU jobs. -Only the crashing per-tensor gradient diagnostic is disabled. A scalar finite -loss/gradient-norm guard remains before each update. The 124M GPT-2 model, data, -MuonClip/auxiliary AdamW settings, batch size and long-run LR schedule are retained. -This bypass does not fix or explain the earlier numerical failure or establish TPU -stability. It is an explicitly requested experimental run, not a passed validation. +The replay-style finite checks described above remain enabled throughout training. +The scalar loss/gradient-norm guard also remains before each update. The 124M GPT-2 +model, data, MuonClip/auxiliary AdamW settings, batch size and long-run LR schedule +are retained. The original numerical failure is not yet diagnosed. Training stays in one process from initialization until the existing allocation cutoff (20 minutes before its recorded expiry), manual STOP, token budget or error. diff --git a/baseline/gpt2_small/TPU_PORT_BUGS.md b/baseline/gpt2_small/TPU_PORT_BUGS.md index c6ff3842..3f31d69f 100644 --- a/baseline/gpt2_small/TPU_PORT_BUGS.md +++ b/baseline/gpt2_small/TPU_PORT_BUGS.md @@ -138,3 +138,23 @@ before evaluation. Per-tensor finite reductions use scalar host transfers withou the earlier stack/extrema diagnostic. No attribution to upstream PyTorch/XLA or TPU hardware is justified yet. These synchronization changes are recorded; a pass does not by itself reproduce or fix the original continuous execution path. + +### Saved-state TPU replay passed, 2026-10-04 + +The user supplied the final results for `muonclip-replay-20261004-150734` at commit +`1b039b39761b13184f1cb09585b2e1e4343f8474`: `one_update_passed`, child exit 0. +All recorded checks passed and the cloud backup was explicitly verified. +Post-update NLL: train `10.98291015625`, validation `10.991303443908691`, +test `10.969128131866455`. Pre-update train/test NLL reproduced the prior saved +checkpoint's evaluation. This localizes neither the original fault nor a fix: +additional synchronization, reductions, process state and checkpoint restoration +differ from the failing run. + +The next user-authorized continuous run retains these synchronization/reduction +boundaries on the existing PyTorch/XLA environment, starting from initialization. +The shared `execution_checks.py` implementation is used by both replay and training. +Checkpointing now precedes evaluation and spectra, with a pending-measurement flag +for explicit recovery. CPU tests cover 25 continuous updates and exact state +recovery after injected measurement failures. No continuous TPU success is claimed +until its output is inspected. TorchTPU migration awaits repository/package access; +the installed environment and all previous evidence remain intact. diff --git a/baseline/gpt2_small/scripts/run_muonclip.py b/baseline/gpt2_small/scripts/run_muonclip.py index cc8fc3f2..74dc3cfc 100644 --- a/baseline/gpt2_small/scripts/run_muonclip.py +++ b/baseline/gpt2_small/scripts/run_muonclip.py @@ -61,7 +61,7 @@ def launch_remote(commit): if not Path('/mnt/disks/rg-data/continuous8/data/train.bin').is_file(): raise RuntimeError('Preserved FineWeb is missing; no download will be started.') stamp=dt.datetime.now(dt.timezone.utc).strftime('%Y%m%d-%H%M%S') - root=BASE/('muonclip-night-'+stamp); root.mkdir() + root=BASE/('muonclip-continuous-'+stamp); root.mkdir() repo=root/'repo'; repo.mkdir() run(['git','-C',str(repo),'init','-q']) run(['git','-C',str(repo),'remote','add','origin','https://github.com/CalculatedContent/rg_optimizers.git']) @@ -73,7 +73,7 @@ def launch_remote(commit): unit='rg-gpt2-muonclip-'+stamp+'.service' record={'root':str(root),'unit':unit,'commit':commit,'node':NODE, 'training_deadline_unix':deadline,'service_deadline_unix':allocation_deadline, - 'purpose':'continuous MuonClip with scalar finite guard; per-tensor diagnostic disabled', + 'purpose':'continuous MuonClip with replay-style finite checks and synchronized optimizer stages', 'cloud_uri':'gs://tpu-builders-504820-ww-continuous8/gpt2small/'+root.name} (root/'launch.json').write_text(json.dumps(record,indent=2)) command=['systemd-run','--unit='+unit,'--property=Type=exec','--property=Restart=no', diff --git a/baseline/gpt2_small/scripts/supervise_muonclip.py b/baseline/gpt2_small/scripts/supervise_muonclip.py index 96c54dd6..b8d45ed5 100644 --- a/baseline/gpt2_small/scripts/supervise_muonclip.py +++ b/baseline/gpt2_small/scripts/supervise_muonclip.py @@ -18,6 +18,7 @@ def prepare_config(template, run_id): cfg=yaml.safe_load(Path(template).read_text()) cfg.update(run_id=run_id, validation_tensor_checks=False, validation_gradient_checks=False, finite_update_guard=True, progress_reporting=True, benchmark_sync_every_step=True, + synchronized_finite_checks=True, checkpoint_before_evaluation=True, cloud_checkpoints=True, metrics_interval=25, metrics_steps=[1,2,4], milestones=[0]) cfg['ww'].update(enabled=True, interval=100, steps=[25], logarithmic=False) return cfg @@ -63,7 +64,9 @@ def main(): base=Path(__file__).resolve().parents[1] child=None report={'status':'starting','long_run_started':False,'training_deadline_unix':args.deadline, - 'per_tensor_diagnostic':False,'scalar_finite_guard':True,'automatic_restart':False} + 'per_tensor_diagnostic':'scalar_flags_no_stack', 'scalar_finite_guard':True, + 'synchronized_optimizer_stages':True, 'checkpoint_before_evaluation':True, + 'automatic_restart':False} try: if args.deadline<=time.time()+60: raise RuntimeError('Too little allocation time remaining.') output.mkdir() # Refuse reuse/overwrite of any previous run. diff --git a/baseline/gpt2_small/src/rg_gpt2_small/execution_checks.py b/baseline/gpt2_small/src/rg_gpt2_small/execution_checks.py new file mode 100644 index 00000000..9170b5d7 --- /dev/null +++ b/baseline/gpt2_small/src/rg_gpt2_small/execution_checks.py @@ -0,0 +1,67 @@ +"""Scalar finite checks and synchronized evaluation shared with the TPU replay. + +The extra graph boundaries are an experimental workaround, not a diagnosed fix. +""" +import json +from pathlib import Path + +import torch +from rg_nanogpt_one_head import runtime as rt +from . import port_debug as debug + + +def check_finite(named, output, label, update, device, *, retain=True): + """No stack/cat, extrema, or full-tensor host copies: materialize scalar flags.""" + debug.stage(output,label+'_started',update) + checks=[] + for name,tensor in named: + value=tensor.detach() + checks.append((name,tuple(value.shape),str(value.dtype),torch.isfinite(value).all(), + (value<0).any() if name.endswith('/exp_avg_sq') else None)) + rt.synchronize(torch.device(device)) + records=[] + for name,shape,dtype,finite,negative in checks: + records.append({'tensor':name,'shape':shape,'dtype':dtype, + 'all_finite':bool(finite.cpu().item()), + 'negative_second_moment':bool(negative.cpu().item()) if negative is not None else False}) + bad=[row for row in records if not row['all_finite'] or row['negative_second_moment']] + report={'stage':label,'update':update,'records':records,'invalid_tensors':bad} + prefix=f'{update:06d}' if retain or bad else 'latest' + debug.write(Path(output)/'diagnostics'/f'{prefix}-{label}.json',report) + if retain or bad: debug.xla_metrics(output,f'{update:06d}-{label}',device) + if bad: + debug.write(Path(output)/'FIRST_INVALID.json',report) + raise RuntimeError(f'First invalid stage: {label}; tensors: '+', '.join(row['tensor'] for row in bad[:8])) + debug.stage(output,label+'_passed',update) + + + +@torch.no_grad() +def evaluate_splits(model, arrays, cfg, device, output, update, label): + from .experiment import batch + model.eval(); result={} + try: + for j,split in enumerate(('train','val','test')): + generator=torch.Generator().manual_seed(cfg['seed']+20000+j) + losses=[]; accuracies=[] + for index in range(cfg['eval_batches']): + stage=f'{label}_{split}_batch_{index}' + debug.stage(output,stage+'_forward',update) + offsets=[] + x,y=batch(arrays[split],generator,cfg['training']['batch_size'], + cfg['model']['block_size'],device,trace=offsets) + debug.write(Path(output)/'diagnostics'/f'{update:06d}-{stage}-inputs.json', + {'split':split,'context':cfg['model']['block_size'],'offsets':offsets}) + logits,loss=model(x,y) + accuracy=(logits.argmax(-1)==y).float().mean() + rt.mark_step(device) + check_finite([('logits',logits),('loss',loss)],output,stage,update,device) + losses.append(float(loss.cpu())); accuracies.append(float(accuracy.cpu())) + result.update({split+'_nll':sum(losses)/len(losses), + split+'_token_error':1-sum(accuracies)/len(accuracies)}) + debug.write(Path(output)/(label+'.json'),result) + print(json.dumps({'evaluation':label,'split':split,**result}),flush=True) + return result + finally: + model.train() + diff --git a/baseline/gpt2_small/src/rg_gpt2_small/experiment.py b/baseline/gpt2_small/src/rg_gpt2_small/experiment.py index d3b62105..9b7a560c 100644 --- a/baseline/gpt2_small/src/rg_gpt2_small/experiment.py +++ b/baseline/gpt2_small/src/rg_gpt2_small/experiment.py @@ -22,6 +22,7 @@ from rg_nanogpt_one_head import optimizers from rg_nanogpt_one_head.spectral import WeightMatrixHolder, _attach_matrix_metadata from . import port_debug +from . import execution_checks def atomic_json(path, value): @@ -185,14 +186,15 @@ def train(cfg, data_root, output, *, device='cpu', resume=False, stop_after=None reporting = cfg.get('progress_reporting', False) diagnostic = cfg.get('validation_gradient_checks', False) or reporting port_checks = cfg.get('validation_tensor_checks', False) - if port_checks or reporting: + synced_checks = cfg.get('synchronized_finite_checks', False) + if port_checks or reporting or synced_checks: port_debug.environment(output) if diagnostic: faulthandler.dump_traceback_later(300, repeat=True) try: return _train(cfg, data_root, output, device, resume, stop_after, deadline) except Exception as exc: - if port_checks or reporting: + if port_checks or reporting or synced_checks: port_debug.failure(output,exc) raise finally: @@ -204,6 +206,9 @@ def _train(cfg, data_root, output, device, resume, stop_after, deadline): cfg = copy.deepcopy(cfg) t = cfg['training']; context = cfg['model']['block_size'] port_checks = cfg.get('validation_tensor_checks', False) + synced_checks = cfg.get('synchronized_finite_checks', False) + if port_checks and synced_checks: + raise ValueError('Choose one tensor-check implementation, not both.') def progress(stage, completed_step): if cfg.get('progress_reporting', False): atomic_json(output/'progress.json', {'stage':stage, 'completed_step':completed_step, @@ -228,6 +233,7 @@ def progress(stage, completed_step): handles = make_handles(model, cfg) gen = torch.Generator().manual_seed(cfg['seed'] + 11) step = 0; elapsed = 0.; last_grad = None; compile_seconds = 0.; steady_seconds = 0.; steady_steps = 0 + measurement_pending = False latest = output / 'checkpoints/latest.json' if resume: pointer = json.loads(latest.read_text()) @@ -240,6 +246,7 @@ def progress(stage, completed_step): random.setstate(state['python_rng']); np.random.set_state(state['numpy_rng']) rt.restore_accelerator_rng_state(state['accelerator_rng'], dev) step = state['step']; elapsed = state['wall_time']; last_grad = state['gradient_norm'] + measurement_pending = state.get('measurement_pending', False) if state['tokens_seen'] != step * step_tokens or state['scheduler_step'] != step: raise RuntimeError('Checkpoint step/token/scheduler mismatch') for folder in ('metrics', 'ww_metrics'): @@ -259,10 +266,28 @@ def progress(stage, completed_step): print(json.dumps({'architecture': architecture(model), 'batch_tokens': step_tokens}), flush=True) started = time.monotonic(); initial_step = step + def payload(step, pending, measurement_pending=False): + return {'run_id': cfg['run_id'], 'config': cfg, 'fingerprint': fingerprint, 'model': model.state_dict(), + 'optimizers': optimizers.optimizer_state_dict(handles), 'step': step, 'tokens_seen': step*step_tokens, + 'scheduler_step': step, 'data_rng': gen.get_state(), 'torch_rng': torch.get_rng_state(), + 'python_rng': random.getstate(), 'numpy_rng': np.random.get_state(), + 'accelerator_rng': rt.capture_accelerator_rng_state(dev), 'wall_time': elapsed + time.monotonic()-started, + 'gradient_norm': last_grad, 'pending_records': pending, 'measurement_pending': measurement_pending} + def record(step, final=False): nonlocal elapsed + if cfg.get('checkpoint_before_evaluation', False): + progress('saving_pre_evaluation_checkpoint', step) + save_checkpoint(output/'checkpoints', payload(step, {}, measurement_pending=True)) progress('evaluating', step) - metrics = evaluate(model, arrays, cfg, dev) + if synced_checks: + measured = execution_checks.evaluate_splits(model, arrays, cfg, dev, output, step, 'evaluation') + metrics = dict(measured) + for split in ('train', 'val', 'test'): + metrics[split+'_accuracy'] = 1-measured[split+'_token_error'] + metrics[split+'_perplexity'] = math.exp(min(measured[split+'_nll'], 700)) + else: + metrics = evaluate(model, arrays, cfg, dev) wall = elapsed + time.monotonic() - started identity = {'run_id': cfg['run_id'], 'optimizer': cfg['optimizer']['family'], 'seed': cfg['seed'], 'step': step, 'tokens_seen': step * step_tokens, 'wall_time': wall, @@ -278,13 +303,7 @@ def record(step, final=False): pending['ww_metrics'] = measured # Checkpoint FIRST includes the pending scalar/WW transaction. On resume, finish missing rows. progress('saving_checkpoint', step) - checkpoint = save_checkpoint(output / 'checkpoints', { - 'run_id': cfg['run_id'], 'config': cfg, 'fingerprint': fingerprint, 'model': model.state_dict(), - 'optimizers': optimizers.optimizer_state_dict(handles), 'step': step, 'tokens_seen': step*step_tokens, - 'scheduler_step': step, 'data_rng': gen.get_state(), 'torch_rng': torch.get_rng_state(), - 'python_rng': random.getstate(), 'numpy_rng': np.random.get_state(), - 'accelerator_rng': rt.capture_accelerator_rng_state(dev), 'wall_time': elapsed + time.monotonic()-started, - 'gradient_norm': last_grad, 'pending_records': pending}, + checkpoint = save_checkpoint(output / 'checkpoints', payload(step, pending), milestone=step in cfg.get('milestones', [])) for folder, value in pending.items(): append_record(output / folder / f'{step:09d}.json', value) print(json.dumps(row), flush=True) @@ -293,7 +312,9 @@ def record(step, final=False): cloud.publish(checkpoint, step) progress('measurement_completed', step) - if not resume: record(0) + if synced_checks: + execution_checks.check_finite(port_debug.optimizer_tensors(model, handles), output, 'initial_state', step, dev) + if not resume or measurement_pending: record(step) training_window = time.monotonic(); last_timed_step = step while step < total: if (stop_after is not None and step >= stop_after) or (deadline and time.time() >= deadline) or (output / 'STOP').exists(): @@ -305,19 +326,27 @@ def record(step, final=False): peak_lr=handle.peak_lr, min_lr=handle.min_lr) handle.set_lr(lr) losses = [] - input_windows = [] if port_checks else None + retain_diagnostics = (step < 4 or (step+1) % cfg['metrics_interval'] == 0 or due(step+1, cfg['ww'])) + input_windows = [] if port_checks or synced_checks else None if port_checks: port_debug.stage(output,'building_forward_backward_graph',step+1) for _ in range(t['grad_accum_steps']): x, y = batch(arrays['train'], gen, t['batch_size'], context, dev,trace=input_windows) _, loss = model(x, y); losses.append(loss.detach()) (loss / t['grad_accum_steps']).backward() - if port_checks: - atomic_json(output/'diagnostics'/f'{step+1:06d}-input-windows.json', + if port_checks or synced_checks: + prefix=f'{step+1:06d}' if retain_diagnostics or port_checks else 'latest' + atomic_json(output/'diagnostics'/f'{prefix}-input-windows.json', {'update':step+1,'split':'train','context':context,'microbatch_offsets':input_windows, 'corpus_identity':'manifest.json/data','config_fingerprint':fingerprint}) spmd.replicate_gradients(model) norm = rt.gradient_norm(model.parameters()) + if synced_checks: + progress('checking_gradients', step) + execution_checks.check_finite( + [('gradient/'+name,p.grad) for name,p in model.named_parameters() if p.grad is not None] + + [('loss/'+str(i),loss) for i,loss in enumerate(losses)] + [('gradient_norm',norm)], + output, 'before_clipping', step+1, dev, retain=retain_diagnostics) if port_checks: port_debug.check_tensors( [('gradient/'+name,p.grad) for name,p in model.named_parameters() if p.grad is not None] @@ -327,12 +356,25 @@ def record(step, final=False): require_finite_update(model, norm, losses, output, step + 1, verbose=cfg.get('validation_gradient_checks', False) or step < 4) torch.nn.utils.clip_grad_norm_(model.parameters(), t['grad_clip'], foreach=False) + if synced_checks: + progress('checking_clipped_gradients', step) + execution_checks.check_finite( + [('gradient/'+name,p.grad) for name,p in model.named_parameters() if p.grad is not None], + output, 'after_clipping', step+1, dev, retain=retain_diagnostics) if port_checks: port_debug.check_tensors( [('gradient/'+name,p.grad) for name,p in model.named_parameters() if p.grad is not None], output,'after_clipping',step+1,dev) port_debug.stage(output,'optimizer_update_started',step+1) - optimizers.optimizer_step(handles); rt.mark_step(dev) + if synced_checks: + for handle in handles: + progress('applying_'+handle.role, step) + port_debug.stage(output, 'applying_'+handle.role, step+1) + handle.optimizer.step(); rt.mark_step(dev); rt.synchronize(dev) + execution_checks.check_finite(port_debug.optimizer_tensors(model, handles), + output, 'after_'+handle.role, step+1, dev, retain=retain_diagnostics) + else: + optimizers.optimizer_step(handles); rt.mark_step(dev) if port_checks: port_debug.check_tensors(port_debug.optimizer_tensors(model,handles), output,'after_optimizer',step+1,dev) diff --git a/baseline/gpt2_small/src/rg_gpt2_small/replay_update.py b/baseline/gpt2_small/src/rg_gpt2_small/replay_update.py index 4f32cd8d..8065dd3e 100644 --- a/baseline/gpt2_small/src/rg_gpt2_small/replay_update.py +++ b/baseline/gpt2_small/src/rg_gpt2_small/replay_update.py @@ -15,6 +15,7 @@ import torch from . import experiment as g, port_debug as debug +from .execution_checks import check_finite, evaluate_splits from rg_nanogpt_one_head import runtime as rt, tpu_spmd as spmd, optimizers @@ -25,29 +26,6 @@ def sha256(path): return digest.hexdigest() -def check_finite(named, output, label, update, device): - """No stack/cat, extrema, or full-tensor host copies: materialize scalar flags.""" - debug.stage(output,label+'_started',update) - checks=[] - for name,tensor in named: - value=tensor.detach() - checks.append((name,tuple(value.shape),str(value.dtype),torch.isfinite(value).all(), - (value<0).any() if name.endswith('/exp_avg_sq') else None)) - rt.synchronize(torch.device(device)) - records=[] - for name,shape,dtype,finite,negative in checks: - records.append({'tensor':name,'shape':shape,'dtype':dtype, - 'all_finite':bool(finite.cpu().item()), - 'negative_second_moment':bool(negative.cpu().item()) if negative is not None else False}) - bad=[row for row in records if not row['all_finite'] or row['negative_second_moment']] - report={'stage':label,'update':update,'records':records,'invalid_tensors':bad} - debug.write(Path(output)/'diagnostics'/f'{update:06d}-{label}.json',report) - debug.xla_metrics(output,f'{update:06d}-{label}',device) - if bad: - debug.write(Path(output)/'FIRST_INVALID.json',report) - raise RuntimeError(f'First invalid stage: {label}; tensors: '+', '.join(row['tensor'] for row in bad[:8])) - debug.stage(output,label+'_passed',update) - def source_tensors(state): yield from (('weight/'+name,tensor) for name,tensor in state['model'].items()) @@ -57,34 +35,6 @@ def source_tensors(state): if torch.is_tensor(value): yield f'optimizer_{index}/{parameter}/{key}',value -@torch.no_grad() -def evaluate_splits(model, arrays, cfg, device, output, update, label): - model.eval(); result={} - try: - for j,split in enumerate(('train','val','test')): - generator=torch.Generator().manual_seed(cfg['seed']+20000+j) - losses=[]; accuracies=[] - for index in range(cfg['eval_batches']): - stage=f'{label}_{split}_batch_{index}' - debug.stage(output,stage+'_forward',update) - offsets=[] - x,y=g.batch(arrays[split],generator,cfg['training']['batch_size'], - cfg['model']['block_size'],device,trace=offsets) - debug.write(Path(output)/'diagnostics'/f'{update:06d}-{stage}-inputs.json', - {'split':split,'context':cfg['model']['block_size'],'offsets':offsets}) - logits,loss=model(x,y) - accuracy=(logits.argmax(-1)==y).float().mean() - rt.mark_step(device) - check_finite([('logits',logits),('loss',loss)],output,stage,update,device) - losses.append(float(loss.cpu())); accuracies.append(float(accuracy.cpu())) - result.update({split+'_nll':sum(losses)/len(losses), - split+'_token_error':1-sum(accuracies)/len(accuracies)}) - debug.write(Path(output)/(label+'.json'),result) - print(json.dumps({'evaluation':label,'split':split,**result}),flush=True) - return result - finally: - model.train() - def replay(checkpoint, data_root, output, device='tpu', expected_step=1): checkpoint=Path(checkpoint); output=Path(output) diff --git a/baseline/gpt2_small/src/rg_gpt2_small/run_backup.py b/baseline/gpt2_small/src/rg_gpt2_small/run_backup.py index 66098808..e43bd0f6 100644 --- a/baseline/gpt2_small/src/rg_gpt2_small/run_backup.py +++ b/baseline/gpt2_small/src/rg_gpt2_small/run_backup.py @@ -27,7 +27,8 @@ def publish(self, checkpoint, step): self.sink.file(checkpoint, 'muonclip/checkpoints/initial.pt') for folder in ('metrics', 'ww_metrics', 'diagnostics'): for path in sorted((self.output/folder).glob('*.json')): - if path not in self.sent: + # Current-stage/latest diagnostic files are intentionally mutable. + if path not in self.sent or path.name.startswith('latest-') or path.name=='current_stage.json': self.sink.file(path, 'muonclip/'+path.relative_to(self.output).as_posix()) self.sent.add(path) self.sink.file(self.output/'manifest.json', 'muonclip/manifest.json') diff --git a/baseline/gpt2_small/tests/test_continuous_synced.py b/baseline/gpt2_small/tests/test_continuous_synced.py new file mode 100644 index 00000000..d0a7d6b1 --- /dev/null +++ b/baseline/gpt2_small/tests/test_continuous_synced.py @@ -0,0 +1,124 @@ +import copy +import json + +import pytest +import torch + +from rg_gpt2_small import experiment as g, execution_checks as checks +from rg_gpt2_small.run_backup import RunBackup +from rg_nanogpt_one_head.checkpoints import optimizer_state_sha256 +from rg_nanogpt_one_head.muonclip import MuonClip +from test_gpt2_experiment import config, data + + +def synced_config(): + c=config(); c['optimizer']=config('fineweb_muonclip_baseline')['optimizer'] + c.update(synchronized_finite_checks=True,checkpoint_before_evaluation=True, + finite_update_guard=True,progress_reporting=True,benchmark_sync_every_step=True, + metrics_steps=[1,2,4]) + return c + + +def load_last(root): + pointer=json.loads((root/'checkpoints/latest.json').read_text()) + return torch.load(root/'checkpoints'/pointer['file'],weights_only=False) + + +def same_training_state(a,b): + for name in a['model']: torch.testing.assert_close(a['model'][name],b['model'][name],rtol=0,atol=0) + assert optimizer_state_sha256(a['optimizers'])==optimizer_state_sha256(b['optimizers']) + assert torch.equal(a['data_rng'],b['data_rng']) + + +def test_25_updates_continue_in_one_process_and_match_cpu_reference(tmp_path,monkeypatch): + torch.set_num_threads(1); c=synced_config() + c['training'].update(max_steps=25,max_tokens=400,schedule_steps=50) + c['metrics_interval']=5; d=data(tmp_path,c) + reference=copy.deepcopy(c); reference['synchronized_finite_checks']=False + g.train(reference,d,tmp_path/'reference') + calls=[]; original=checks.check_finite + def checked(*args,**kw): + calls.append((args[2],args[3])); return original(*args,**kw) + monkeypatch.setattr(checks,'check_finite',checked) + out=tmp_path/'synced'; g.train(c,d,out) + for step in range(1,26): + stages=[label for label,index in calls if index==step] + wanted=['before_clipping','after_clipping','after_primary','after_auxiliary'] + assert [label for label in stages if label in wanted]==wanted + assert calls.count(('initial_state',0))==1 + assert not (out/'resume_verified.json').exists() + state=load_last(out); assert state['step']==25 and state['tokens_seen']==400 + assert not state['measurement_pending'] + same_training_state(state,load_last(tmp_path/'reference')) + assert len(list((out/'checkpoints').glob('step_*.pt')))==3 + assert (out/'diagnostics/latest-after_auxiliary.json').is_file() + assert not (out/'diagnostics/000024-after_auxiliary.json').exists() + + +@pytest.mark.parametrize('failed_stage',['evaluation','spectra']) +def test_checkpoint_precedes_measurement_and_resume_preserves_weights(tmp_path,monkeypatch,failed_stage): + torch.set_num_threads(1); c=synced_config(); d=data(tmp_path,c) + if failed_stage=='spectra': + c['ww'].update(enabled=True,interval=2) + monkeypatch.setattr(g,'measure_ww',lambda *args:{'records':[],'seconds':0.01}) + original_eval=checks.evaluate_splits; original_ww=g.measure_ww + output=tmp_path/'failed' + def evaluate(*args,**kw): + if args[5]==2 and failed_stage=='evaluation': + saved=load_last(output) + assert saved['step']==2 and saved['measurement_pending'] + raise RuntimeError('Injected evaluation failure') + return original_eval(*args,**kw) + def spectra(model,cfg,identity,metrics): + if identity['step']==2: + saved=load_last(output) + assert saved['step']==2 and saved['measurement_pending'] + raise RuntimeError('Injected spectra failure') + return original_ww(model,cfg,identity,metrics) + monkeypatch.setattr(checks,'evaluate_splits',evaluate) + if failed_stage=='spectra': monkeypatch.setattr(g,'measure_ww',spectra) + with pytest.raises(RuntimeError,match='Injected'): g.train(c,d,output) + before=(output/'metrics/000000001.json').read_bytes() + assert not (output/'metrics/000000002.json').exists() + assert (output/'TPU_PORT_FAILURE.json').exists() + monkeypatch.setattr(checks,'evaluate_splits',original_eval) + monkeypatch.setattr(g,'measure_ww',original_ww) + g.train(c,d,output,resume=True) + g.train(c,d,tmp_path/'reference') + same_training_state(load_last(output),load_last(tmp_path/'reference')) + assert (output/'metrics/000000001.json').read_bytes()==before + assert (output/'metrics/000000002.json').exists() + + +def test_invalid_primary_does_not_apply_auxiliary(tmp_path,monkeypatch): + torch.set_num_threads(1); c=synced_config(); d=data(tmp_path,c) + original=MuonClip.step + def broken(self,*args,**kw): + original(self,*args,**kw) + with torch.no_grad(): self.param_groups[0]['params'][0].fill_(float('nan')) + monkeypatch.setattr(MuonClip,'step',broken) + def forbidden(*args,**kwargs): raise AssertionError('Auxiliary applied after invalid primary') + monkeypatch.setattr(torch.optim.AdamW,'step',forbidden) + output=tmp_path/'invalid' + with pytest.raises(RuntimeError,match='First invalid stage: after_primary'): g.train(c,d,output) + report=json.loads((output/'FIRST_INVALID.json').read_text()) + assert report['update']==1 and report['invalid_tensors'] + assert load_last(output)['step']==0 + + +def test_cloud_refreshes_current_diagnostics(tmp_path): + class Sink: + def __init__(self): self.files=[] + def file(self,path,name): + self.files.append((name,path.read_bytes())) + return {'generation':'1','bytes':1,'crc32c':'test'} + def json(self,*args): pass + (tmp_path/'manifest.json').write_text('{}') + folder=tmp_path/'diagnostics'; folder.mkdir() + latest=folder/'latest-after_primary.json'; latest.write_text('first') + fixed=folder/'000001-after_primary.json'; fixed.write_text('fixed') + checkpoint=tmp_path/'checkpoint.pt'; checkpoint.write_bytes(b'state') + sink=Sink(); backup=RunBackup(tmp_path,'unused',sink) + backup.publish(checkpoint,1); latest.write_text('second'); backup.publish(checkpoint,2) + assert [data for name,data in sink.files if name.endswith(latest.name)]==[b'first',b'second'] + assert [data for name,data in sink.files if name.endswith(fixed.name)]==[b'fixed'] diff --git a/baseline/gpt2_small/tests/test_muonclip_night.py b/baseline/gpt2_small/tests/test_muonclip_night.py index c01b1427..9c626c18 100644 --- a/baseline/gpt2_small/tests/test_muonclip_night.py +++ b/baseline/gpt2_small/tests/test_muonclip_night.py @@ -63,6 +63,7 @@ def test_config_retains_model_data_and_optimizer(): for key in ('model','dataset','optimizer','training'): assert c[key]==old[key] assert c['validation_tensor_checks'] is c['validation_gradient_checks'] is False assert c['finite_update_guard'] and c['cloud_checkpoints'] + assert c['synchronized_finite_checks'] and c['checkpoint_before_evaluation'] assert c['metrics_interval']==25 and c['ww']['interval']==100 From 5f3d2fdde485bb5ed0d201ce6ffc3c97f89c0efe Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Sun, 4 Oct 2026 13:19:39 -0700 Subject: [PATCH 24/36] Add a 30-minute GPT-2 FineWeb reference run on the existing TPU Stop the current MuonClip service and launch one bounded reference job with the pinned upstream GPT-2 model, canonical tokenized FineWeb data, AdamW, the published batch/schedule, and fixed benchmark validation tokens. Remove WeightWatcher, per-tensor diagnostics, preflight, and periodic checkpoint uploads from this runner. Limit the whole new service to 30 minutes including compilation, data, final evaluation/save, and backup. Preserve old artifacts and the current allocation; refuse concurrent jobs. Report partial training/evaluation and hardware/data-order differences honestly, with the published reference curve at matching token counts. CPU tests verify accumulated updates against upstream AdamW, exact sequential data windows, corrupted download rejection, external timeout and concurrent-launch prevention. Existing suite passed; TPU execution remains to be tested. --- baseline/gpt2_small/README.md | 7 + baseline/gpt2_small/scripts/run_muonclip.py | 3 + baseline/gpt2_small/speedrun30/README.md | 120 +++ baseline/gpt2_small/speedrun30/cloudshell.py | 164 ++++ .../gpt2_small/speedrun30/data_manifest.json | 422 +++++++++ .../gpt2_small/speedrun30/reference_val.json | 402 ++++++++ baseline/gpt2_small/speedrun30/train.py | 318 +++++++ baseline/gpt2_small/speedrun30/vendor/LICENSE | 22 + .../speedrun30/vendor/llmc_train_gpt2.py | 860 ++++++++++++++++++ baseline/gpt2_small/speedrun30/worker.py | 66 ++ baseline/gpt2_small/tests/test_speedrun30.py | 136 +++ 11 files changed, 2520 insertions(+) create mode 100644 baseline/gpt2_small/speedrun30/README.md create mode 100644 baseline/gpt2_small/speedrun30/cloudshell.py create mode 100644 baseline/gpt2_small/speedrun30/data_manifest.json create mode 100644 baseline/gpt2_small/speedrun30/reference_val.json create mode 100644 baseline/gpt2_small/speedrun30/train.py create mode 100644 baseline/gpt2_small/speedrun30/vendor/LICENSE create mode 100644 baseline/gpt2_small/speedrun30/vendor/llmc_train_gpt2.py create mode 100644 baseline/gpt2_small/speedrun30/worker.py create mode 100644 baseline/gpt2_small/tests/test_speedrun30.py diff --git a/baseline/gpt2_small/README.md b/baseline/gpt2_small/README.md index ca09ac8f..2a566b5e 100644 --- a/baseline/gpt2_small/README.md +++ b/baseline/gpt2_small/README.md @@ -1,5 +1,12 @@ # GPT-2 Small / FineWeb-Edu validation +For the **30-minute GPT-2/FineWeb reference run**, see +[speedrun30/README.md](speedrun30/README.md). Its launcher stops the current +MuonClip service, preserves prior data, and uses the published benchmark's +tokenized FineWeb validation set and upstream GPT-2 model. It does not run +WeightWatcher or per-tensor diagnostics. The hard time cap permits a partial +reference run; it is not a promise of reaching the final published loss. + This is the GPT-2 Small experiment, under `baseline/gpt2_small`. Its training module is `rg_gpt2_small.experiment`; configurations and launchers live here. It reuses the shared GPT implementation, optimizers, SPMD, corpus validation and WeightWatcher diff --git a/baseline/gpt2_small/scripts/run_muonclip.py b/baseline/gpt2_small/scripts/run_muonclip.py index 74dc3cfc..e0469211 100644 --- a/baseline/gpt2_small/scripts/run_muonclip.py +++ b/baseline/gpt2_small/scripts/run_muonclip.py @@ -28,6 +28,9 @@ def active(unit): def assert_idle(): + reference=BASE/'SPEEDRUN30_LATEST.json' + if reference.exists() and active(json.loads(reference.read_text())['unit']): + raise RuntimeError('The 30-minute GPT-2 reference run is active; no concurrent trainer launched.') if active('rg-gpt2-validation.service') or active('rg-continuous8.service'): raise RuntimeError('An existing training service is active; no training launched.') if LATEST.exists() and active(json.loads(LATEST.read_text())['unit']): diff --git a/baseline/gpt2_small/speedrun30/README.md b/baseline/gpt2_small/speedrun30/README.md new file mode 100644 index 00000000..e1dfddfd --- /dev/null +++ b/baseline/gpt2_small/speedrun30/README.md @@ -0,0 +1,120 @@ +# GPT-2 / FineWeb reference, bounded to 30 minutes + +Run this before interpreting the custom MuonClip experiment as a reproduction. +This is a TPU port of the original GPT-2/FineWeb baseline behind the NanoGPT +speedrun's 3.28 target. It is a timed, potentially partial reference run, not a +promise to finish the benchmark in 30 minutes. Karpathy's older OpenWebText +nanoGPT recipe is a different benchmark and is not used here. + +## Run from Cloud Shell + +From a clean checkout of the pushed commit: + +```bash +python3 baseline/gpt2_small/speedrun30/cloudshell.py start +``` + +The launcher requests a final save from the current MuonClip service, waits up to +90 seconds, and then stops that service if necessary. Its previously saved +checkpoints, diagnostics, cloud archives and FineWeb-Edu corpus are retained. +The new job uses the existing eight-chip TPU and installed PyTorch 2.6/XLA 2.6. +There is no environment installation, reallocation or automatic restart. + +The new service has a **30-minute limit including benchmark-data downloads, +compilation, training, evaluation, final save and backup**. Normal training ends +earlier to reserve finalization time; a systemd deadline kills stalled work. +Stopping the old run and checking out the source happen before this clock starts. +The TPU allocation itself continues to exist after the job ends. + +```bash +python3 baseline/gpt2_small/speedrun30/cloudshell.py status +``` + +`SPEEDRUN30_LATEST.json` on the mounted disk records the service, output directory, +commit and deadline. `RUN_STATUS.json` distinguishes finish, failure and timeout. +No completion message should be interpreted as proof of benchmark reproduction. + +## Pinned reference + +- Model source: [karpathy/llm.c at 7ecd8906](https://github.com/karpathy/llm.c/blob/7ecd8906afe6ed7a2b2cdb731c042f26d525b820/train_gpt2.py). + `vendor/llmc_train_gpt2.py` is the upstream file, with its MIT license included. + The GPT model and its initialization are imported directly, not reimplemented. +- Recipe: [GPT-2 124M reproduction launcher](https://github.com/karpathy/llm.c/blob/7ecd8906afe6ed7a2b2cdb731c042f26d525b820/scripts/run_gpt2_124M.sh). +- Training curve: [published October 13 2024 llm.c run](https://github.com/KellerJordan/modded-nanogpt/tree/master/records/track_1_short/2024-10-13_llmc). + `reference_val.json` preserves its validation rows and exact step/token counts. +- Data: `kjj0/fineweb10B-gpt2`, revision + `889765ea1f903759787add96995d81171b632d0c`, the GPT-2-tokenized FineWeb shards + used by the speedrun's official cached-data downloader. + All 104 shard filenames, byte lengths and SHA256 hashes are pinned in + `data_manifest.json`. Download only the validation shard and training shards + reached by this run; do not cycle a small subset. Cached files go under + `/mnt/disks/rg-data/benchmark-fineweb10B-889765ea`. + +| Setting | This reference run | +|---|---| +| Model | GPT-2 124M, 12 layers, 12 heads, width 768, tied embeddings | +| Tokenizer / context | GPT-2 / 1,024 | +| Global batch | 524,288 tokens per optimizer update | +| TPU microbatch | 64 sequences globally, accumulated 8 times | +| Optimizer | Stock PyTorch AdamW, betas 0.9/0.95, epsilon 1e-8 | +| Weight decay | 0.1 on matrices, zero on vectors/biases | +| Gradient clipping | Global L2 norm 1.0 | +| Learning rate | 0.0006, 700-update warmup, cosine decay to zero | +| Schedule horizon | 19,560 updates, as in the published curve | +| Full validation | First 10,485,760 validation tokens, fixed context 1,024 | +| Measurements | Scalar train loss/throughput; validation every 250 updates and at end | +| Checkpoint | One final model-only checkpoint; no automatic resume | + +The dataset contains 10,255,324,043 training tokens; integer division by the +reference global batch gives 19,560 updates. The wall-clock limit truncates this +schedule; it does not shorten warmup or compress decay to force a lower loss. +The first validation at step zero uses only 1,048,576 tokens to save time and is +explicitly labelled a partial evaluation. Later evaluations use the full +10,485,760-token benchmark unless the deadline interrupts them. A partial +evaluation is never labelled a full benchmark score. + +## Hardware changes and limits of comparison + +One XLA SPMD process partitions batches over eight TPU chips and replicates +parameters and global gradients. It does not multiply/divide the loss by eight +a second time. BF16 autocast is used for compute with FP32 weights/moments. +The upstream mathematical-attention option is selected because CUDA flash +kernels are unavailable. Stock AdamW uses XLA-supported capturable state and a +tensor learning rate so its changing step/LR do not become host graph constants. +Lazy graphs are submitted at microbatch boundaries and after each optimizer +update; there are no per-tensor finite scans or diagnostic host transfers. + +The model, tokenizer, validation file, validation token count, global batch and +optimizer hyperparameters are controlled. This remains a **port**, not bitwise +CUDA reproduction: CUDA's kernels/rounding and C shuffled loader differ from +this PyTorch sequential loader and initialization. Matching CPU reference +updates does not by itself prove TPU correctness. `latest_validation.json` +includes the nearest lower/upper published observations at equal tokens; it +does not interpolate an invented expected loss or issue an automatic pass/fail. + +Published curve examples (use only after a full validation on the pinned file): + +| Updates | Training tokens | Published validation NLL | +|---:|---:|---:| +| 250 | 131,072,000 | 6.1710 | +| 500 | 262,144,000 | 5.3743 | +| 1,000 | 524,288,000 | 4.3170 | +| 19,560 | 10,255,073,280 | 3.2722 | + +## Output and verification + +The run writes `metrics.jsonl`, `latest_validation.json`, `manifest.json`, status, +and a final `model_final.pt` when finalization fits within the time cap. A failure +or hard timeout can prevent final evaluation/checkpoint creation; earlier logs +remain on the disk. One bounded final cloud backup is attempted through object +permissions. Check its exit result or `CLOUD_BACKUP_VERIFIED.json`; a local file +alone is not evidence of successful upload. + +No WeightWatcher, alpha fits, per-matrix checks, tensor snapshots, preflight, or +periodic checkpoint uploads run in this experiment. It monitors ordinary scalar +loss and aborts on nonfinite scalar loss. + +Local CPU tests compare several accumulated updates to the upstream AdamW +implementation, compare shard traversal/target shifting to its data loader, +reject corrupted downloads, exercise the external timeout, and prevent a +duplicate launch. The first live TPU run is still required. diff --git a/baseline/gpt2_small/speedrun30/cloudshell.py b/baseline/gpt2_small/speedrun30/cloudshell.py new file mode 100644 index 00000000..7cfdc38d --- /dev/null +++ b/baseline/gpt2_small/speedrun30/cloudshell.py @@ -0,0 +1,164 @@ +"""Stop the current MuonClip service and run one bounded GPT-2 reference job.""" +import argparse +import datetime as dt +import fcntl +import importlib.util +import json +import os +from pathlib import Path +import re +import shlex +import subprocess +import sys +import time + +PROJECT = "tpu-builders-504820" +ZONE = "us-west4-a" +QUEUE = "ww-gpt2-validation-48h-20261004-s1337" +NODE = QUEUE + "-node" +BASE = Path("/mnt/disks/rg-data/gpt2small") +LATEST = BASE / "SPEEDRUN30_LATEST.json" + + +def run(command, **kwargs): + return subprocess.run(command, check=True, text=True, **kwargs) + + +def active(unit): + p = subprocess.run(["systemctl", "show", unit, "--property=ActiveState", "--value"], + capture_output=True, text=True, timeout=10) + return p.stdout.strip() in {"active", "activating", "deactivating", "reloading"} + + +def stop_current(): + pointer = BASE/"MUONCLIP_LATEST.json" + if not pointer.exists(): + return + record = json.loads(pointer.read_text()) + unit, root = record["unit"], Path(record["root"]).resolve() + if root.parent != BASE or not re.fullmatch(r"rg-gpt2-muonclip-\d{8}-\d{6}\.service", unit): + raise RuntimeError("Unexpected existing service identity; nothing stopped.") + if not active(unit): + return + print("Requesting final save from current MuonClip run.", flush=True) + (root/"muonclip/STOP").touch() + end = time.monotonic()+90 + while active(unit) and time.monotonic() < end: + time.sleep(1) + if active(unit): + print("Stopping the old service; retaining its latest saved checkpoint.", flush=True) + subprocess.run(["systemctl", "stop", "--no-block", unit], check=True) + end = time.monotonic()+40 + while active(unit) and time.monotonic() < end: + time.sleep(1) + if active(unit): + raise RuntimeError("Old service has not stopped; new job was not launched.") + print("Old run stopped. Checkpoints, logs and FineWeb-Edu retained.", flush=True) + + +def status_remote(): + if not LATEST.exists(): + print("No 30-minute reference run launched.") + return + record = json.loads(LATEST.read_text()) + print(json.dumps(record, indent=2), flush=True) + subprocess.run(["systemctl", "--no-pager", "--full", "status", record["unit"]]) + root = Path(record["root"]) + for name in ("RUN_STATUS.json", "status.json", "latest_validation.json"): + if (root/name).exists(): + print(name + "\n" + (root/name).read_text(), flush=True) + subprocess.run(["tail", "-n", "30", str(root/"run.log")]) + + +def start_remote(commit): + if os.geteuid() != 0 or not os.path.ismount("/mnt/disks/rg-data"): + raise RuntimeError("Requires the existing mounted data disk and root.") + if not re.fullmatch(r"[0-9a-f]{40}", commit): + raise ValueError("Expected pinned commit") + with (BASE/"port-check-launch.lock").open("a") as lock: + fcntl.flock(lock, fcntl.LOCK_EX|fcntl.LOCK_NB) + if LATEST.exists() and active(json.loads(LATEST.read_text())["unit"]): + print("The 30-minute run is already active; no duplicate launched.") + status_remote() + return + allocation = json.loads((BASE/QUEUE/"allocation.json").read_text()) + if float(allocation["validation_deadline_unix"])-time.time() < 2100: + raise RuntimeError("Less than 35 minutes remain on this allocation.") + stamp = dt.datetime.now(dt.timezone.utc).strftime("%Y%m%d-%H%M%S") + root = BASE/("gpt2-speedrun30-"+stamp) + root.mkdir() + repo = root/"repo" + repo.mkdir() + run(["git", "-C", str(repo), "init", "-q"]) + run(["git", "-C", str(repo), "remote", "add", "origin", + "https://github.com/CalculatedContent/rg_optimizers.git"]) + run(["git", "-C", str(repo), "fetch", "--depth", "1", "origin", commit], timeout=120) + run(["git", "-C", str(repo), "checkout", "--detach", commit]) + stop_current() + # Existing common guard blocks other trainers/replays on this same TPU. + path = repo/"baseline/gpt2_small/scripts/run_muonclip.py" + spec = importlib.util.spec_from_file_location("existing_launch", path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + module.assert_idle() + start = time.time() + deadline = min(start+1800, float(allocation["validation_deadline_unix"])-30) + if deadline-start < 1750: + raise RuntimeError("Allocation time exhausted during cleanup; no job launched.") + unit = "rg-gpt2-speedrun30-"+stamp+".service" + base = repo/"baseline/gpt2_small" + env = {"PYTHONPATH":str(base/"src")+":"+str(base.parent/"nanogpt_one_head/src"), + "PJRT_DEVICE":"TPU", "TPU_ACCELERATOR_TYPE":"v5litepod-8", + "OMP_NUM_THREADS":"4", "OPENBLAS_NUM_THREADS":"4", "MKL_NUM_THREADS":"4", + "XLA_USE_SPMD":"1", "TOKENIZERS_PARALLELISM":"false"} + record = {"root":str(root), "unit":unit, "commit":commit, + "started_unix":start, "deadline_unix":deadline, "limit_seconds":1800, + "cloud_uri":"gs://tpu-builders-504820-ww-continuous8/gpt2small/"+root.name} + (root/"launch.json").write_text(json.dumps(record, indent=2)) + (root/"commit.txt").write_text(commit+"\n") + command = ["systemd-run", "--unit="+unit, "--property=Type=exec", + "--property=Restart=no", "--property=RuntimeMaxSec="+str(int(deadline-time.time())-5), + "--property=TimeoutStopSec=5", "--property=KillMode=control-group", + "--property=StandardOutput=append:"+str(root/"run.log"), + "--property=StandardError=append:"+str(root/"run.log")] + command += ["--setenv="+key+"="+value for key,value in env.items()] + command += ["/mnt/disks/rg-data/continuous8/venv/bin/python", "-u", + str(base/"speedrun30/worker.py"), str(root), str(deadline)] + run(command) + tmp = LATEST.with_suffix(".tmp") + tmp.write_text(json.dumps(record, indent=2)) + tmp.replace(LATEST) + print("GPT-2 reference run started:", unit, flush=True) + print("Hard stop UTC:", dt.datetime.fromtimestamp(deadline, dt.timezone.utc).isoformat(), flush=True) + print("Log:", root/"run.log", flush=True) + print("30 minutes maximum including benchmark data, compilation, training and backup.", flush=True) + print("No WeightWatcher, per-tensor checks, preflight, or automatic restart.", flush=True) + print("The TPU allocation itself remains available after the job stops.", flush=True) + + +def main(): + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("action", choices=("start", "status")) + p.add_argument("--on-tpu", action="store_true", help=argparse.SUPPRESS) + p.add_argument("--commit", help=argparse.SUPPRESS) + a = p.parse_args() + if a.on_tpu: + start_remote(a.commit) if a.action == "start" else status_remote() + return 0 + remote = ["sudo", "python3", "-c", Path(__file__).read_text(), a.action, "--on-tpu"] + if a.action == "start": + repo = Path(__file__).resolve().parents[3] + if run(["git", "-C", str(repo), "status", "--porcelain"], capture_output=True).stdout.strip(): + raise RuntimeError("Use the clean worktree from the launch block.") + commit = run(["git", "-C", str(repo), "rev-parse", "HEAD"], capture_output=True).stdout.strip() + remote += ["--commit", commit] + return subprocess.run(["gcloud", "compute", "tpus", "tpu-vm", "ssh", NODE, + "--project="+PROJECT, "--zone="+ZONE, "--worker=0", "--command="+shlex.join(remote)]).returncode + + +if __name__ == "__main__": + try: + sys.exit(main()) + except Exception as exc: + print("30-minute launch:", exc, file=sys.stderr) + sys.exit(1) diff --git a/baseline/gpt2_small/speedrun30/data_manifest.json b/baseline/gpt2_small/speedrun30/data_manifest.json new file mode 100644 index 00000000..0acadd36 --- /dev/null +++ b/baseline/gpt2_small/speedrun30/data_manifest.json @@ -0,0 +1,422 @@ +{ + "repo": "kjj0/fineweb10B-gpt2", + "revision": "889765ea1f903759787add96995d81171b632d0c", + "files": { + "fineweb_train_000001.bin": { + "sha256": "771fa4a99b9fe0946ffb6e848b4ba5c6a9b0fe87860ebf03bc2c1c7e45f8178e", + "size": 200001024 + }, + "fineweb_train_000002.bin": { + "sha256": "d4ba8d2d740905a836a4222d2aff01750560e7d49a51a745a1b3235520aec2cd", + "size": 200001024 + }, + "fineweb_train_000003.bin": { + "sha256": "f70fa272369730fd9e39f1848bfe4daa3289b6e490c1e7762777f00ea8b66246", + "size": 200001024 + }, + "fineweb_train_000004.bin": { + "sha256": "d8d88850f89fe2030e866380e90d544167467d407080403b830b903543925835", + "size": 200001024 + }, + "fineweb_train_000005.bin": { + "sha256": "36cef3fc41863eb6731b41e986edde38e368853ee6deb87b928990bf445a119c", + "size": 200001024 + }, + "fineweb_train_000006.bin": { + "sha256": "d80e5b209983e636655fc738fdf281f0132f0298bf4c0848f7e7a6bea6ca06f0", + "size": 200001024 + }, + "fineweb_train_000007.bin": { + "sha256": "fa9f94dce61686b5467bc6cd3078424203db23ac254d804a6d8d5d440afc1e11", + "size": 200001024 + }, + "fineweb_train_000008.bin": { + "sha256": "308e6937af5d78b7489a8aac3a3951dced5b46103e613f8bc4eb543f5092526c", + "size": 200001024 + }, + "fineweb_train_000009.bin": { + "sha256": "1635300891634346ecff3a78d6e489a69c647885f0dcd249fe3eb76fac683458", + "size": 200001024 + }, + "fineweb_train_000010.bin": { + "sha256": "72e88f66104dacca163eeb3a00fe40a45a821d2ecd4fb9ce343808996e5d2f5e", + "size": 200001024 + }, + "fineweb_train_000011.bin": { + "sha256": "99edff412654ca6e06d4b3ca563acd8937ee6042f49d009526b8f1074f4d7300", + "size": 200001024 + }, + "fineweb_train_000012.bin": { + "sha256": "92a2dc78e3971ac2a36ef3e9f19c92533bc35ef2dd735bb6759e7b864e301084", + "size": 200001024 + }, + "fineweb_train_000013.bin": { + "sha256": "6cc3884bf19a2d9b4a12d52e4861276da2e3c6614ce97e6bfa0769d701f7a4c5", + "size": 200001024 + }, + "fineweb_train_000014.bin": { + "sha256": "6bc06a838d5a258fb51a3c2a3d14fd784f83b173f408d8d2d3f85b3290094a96", + "size": 200001024 + }, + "fineweb_train_000015.bin": { + "sha256": "09ad252e5d641862f89c31e4f0fa9559933d3e2bc6f657fc05f63166c2d74741", + "size": 200001024 + }, + "fineweb_train_000016.bin": { + "sha256": "56e97975654bb7a1f08f1397dd7894146dd7f8c36bc9fdd71da9675c43ab9a85", + "size": 200001024 + }, + "fineweb_train_000017.bin": { + "sha256": "a51e60d3f9fce16a9dbcb20e5333e1c8e379acf3a703e63b1fecfef49aa17458", + "size": 200001024 + }, + "fineweb_train_000018.bin": { + "sha256": "8a35ee9b00834e085e20383711d1fa13efe9d05d3d20161cebd60d73ac935fe9", + "size": 200001024 + }, + "fineweb_train_000019.bin": { + "sha256": "e4c82c17c1805a8d946580bb09f10575add3fb56dc3305729d75eae083fa2856", + "size": 200001024 + }, + "fineweb_train_000020.bin": { + "sha256": "bbdb7866974ce4389c888fce7394a2c05e71cb6a0f1ada2e2a37520a03302b6c", + "size": 200001024 + }, + "fineweb_train_000021.bin": { + "sha256": "8d0728820d1dbaeafebcaca5f8463700b3c4b449819ef6dad8643cb83433c7e1", + "size": 200001024 + }, + "fineweb_train_000022.bin": { + "sha256": "a12df006d9730c98302c9765c5188baeb3e3b2c5c310ab6bf6cc5e53f1dbc607", + "size": 200001024 + }, + "fineweb_train_000023.bin": { + "sha256": "b2110dd3393ba91ae31e4ade595c1b7862eaaffb15e91d71d02d44b131e96b08", + "size": 200001024 + }, + "fineweb_train_000024.bin": { + "sha256": "38acaa83716a8df55b368683acaa832571b7e5c8ccd80a14ab9e1f7d7817f4d3", + "size": 200001024 + }, + "fineweb_train_000025.bin": { + "sha256": "6acd61ee12134312527b57b631856f9916390c580fc6bced6c7d50ea24b31954", + "size": 200001024 + }, + "fineweb_train_000026.bin": { + "sha256": "5cdb76cd96264b23348546e7b155a0034551b9dcd5afded0c1c1751ca39a4646", + "size": 200001024 + }, + "fineweb_train_000027.bin": { + "sha256": "15b215a51ae8aa9126014f09071f5b68a494479caa50744279aca9a45c4c2572", + "size": 200001024 + }, + "fineweb_train_000028.bin": { + "sha256": "3f0a690b57335aaf78e69b9205b4f86b47e93e12f351a3833c110db62105d131", + "size": 200001024 + }, + "fineweb_train_000029.bin": { + "sha256": "10ec77e3afd0d05fe364d4f6b6472eff04ce7b999a1a09b9a6a82bb76d9eba08", + "size": 200001024 + }, + "fineweb_train_000030.bin": { + "sha256": "b38194edbfaa6991a4bf00147140f0edc13d0ddccaf45bf8a2c513879ce8529a", + "size": 200001024 + }, + "fineweb_train_000031.bin": { + "sha256": "cee4312443927c23b5681a649710bea32cf68721969bf29a0787bf9e96dbb081", + "size": 200001024 + }, + "fineweb_train_000032.bin": { + "sha256": "d569dbb11fbe48708374c8c28912ad722a5a08417d21b69a327616c6a36c722c", + "size": 200001024 + }, + "fineweb_train_000033.bin": { + "sha256": "e0bc16f223eb61abc7d8f8e625afdc378a80862cb7b8df0f20aae84deef52890", + "size": 200001024 + }, + "fineweb_train_000034.bin": { + "sha256": "a6bca1a76eef61f52386e7551c0d548672f9466684097584c76d2a9134194762", + "size": 200001024 + }, + "fineweb_train_000035.bin": { + "sha256": "48c2490ee64ad5bfcb25916ef7df946eb8b5090493e4d470ef2f630fb50aa10b", + "size": 200001024 + }, + "fineweb_train_000036.bin": { + "sha256": "40521a93f12a41ba864661d31a5c7c559f89559ccf0627542883e44796322497", + "size": 200001024 + }, + "fineweb_train_000037.bin": { + "sha256": "c82699f5a70f7c005c242795ffe892707d8d4c15166ed677c4809b9f5080e291", + "size": 200001024 + }, + "fineweb_train_000038.bin": { + "sha256": "061eb2ed09ad3b938c4fa2b30116ab0b65d3b1254380cc7f8ddd611066e22e14", + "size": 200001024 + }, + "fineweb_train_000039.bin": { + "sha256": "cb789ec6d36109946bc7414f71545df19279f125aa2b7589283ddbc70e907fdc", + "size": 200001024 + }, + "fineweb_train_000040.bin": { + "sha256": "0d9d435baba0e4bd2f80aa5526a06c13033c6441ad81736a48d83b73fba5c44a", + "size": 200001024 + }, + "fineweb_train_000041.bin": { + "sha256": "d4bbc642b6e8a01db306c5656909475fdc6e0cd77481e77aa6398e7b7410a62d", + "size": 200001024 + }, + "fineweb_train_000042.bin": { + "sha256": "15a03e08508cd488bf48f3e3c5b4e51ed27082384606fdb204808ca0af0f7615", + "size": 200001024 + }, + "fineweb_train_000043.bin": { + "sha256": "0c3ec2e8eb6907fe36344e709b8c06a9f0e8195d18214de0e4ee345507295166", + "size": 200001024 + }, + "fineweb_train_000044.bin": { + "sha256": "c7c440e91e8346f193fb3d5b5c89751dd69a2e8146b6f48b46a23b6be59b3352", + "size": 200001024 + }, + "fineweb_train_000045.bin": { + "sha256": "9c38389785c375f8dcb29c72ca0472c2a95ec8119ff8f1ec314bb399e195beb9", + "size": 200001024 + }, + "fineweb_train_000046.bin": { + "sha256": "4b2cf4bdb14da18843b4c8db07e6caa97926fb8c6d1bc2f42f1c9dfe30799d55", + "size": 200001024 + }, + "fineweb_train_000047.bin": { + "sha256": "9393219680b7d7e10b7909baf4769ae2e7183ddbd92019a19cd8ac7e808c8c27", + "size": 200001024 + }, + "fineweb_train_000048.bin": { + "sha256": "3c0c5fc07f79a5cdaaf8f999e3f47d0ebac39f2d793a67ca468b4467cc2a8c40", + "size": 200001024 + }, + "fineweb_train_000049.bin": { + "sha256": "9560b2f5effc20cec11853e5e785a10b5442cf06161f0dca4e10cc3dc50dc776", + "size": 200001024 + }, + "fineweb_train_000050.bin": { + "sha256": "3798ffa167f4592f5a96eb51a59b2f13c2177a53b85e347a3952113cdef272e9", + "size": 200001024 + }, + "fineweb_train_000051.bin": { + "sha256": "d59c919918ceb51485a4488788905b22cc6900c2ce326f24a066bad87be30027", + "size": 200001024 + }, + "fineweb_train_000052.bin": { + "sha256": "3b3985965ef3fcd19f10283f5f0fc2b71d262792c2dd70082796a69b6d5d836f", + "size": 200001024 + }, + "fineweb_train_000053.bin": { + "sha256": "c195e477d4ce5dada0b09046b8a5298a8e1dfb67896eb1aaec1dfba1ff150754", + "size": 200001024 + }, + "fineweb_train_000054.bin": { + "sha256": "82d6f578ab966f26aadb6a6341f81efd1ecf8c697df68065a6ddeda7d24f8865", + "size": 200001024 + }, + "fineweb_train_000055.bin": { + "sha256": "19df10c9c18c4a7575a8467c5ca908372357ca992e44c5bca364ce9445bc4de2", + "size": 200001024 + }, + "fineweb_train_000056.bin": { + "sha256": "f374d764770b04b19fc2e973ec5fd1d64e85d9c167d3c09800f1672ea5df71ba", + "size": 200001024 + }, + "fineweb_train_000057.bin": { + "sha256": "5fba96c42e40494d4b4d31bbfc3ab83d653cfdc6c28d80e878b0ed16e1861076", + "size": 200001024 + }, + "fineweb_train_000058.bin": { + "sha256": "d81e3f09d7571445396fb1c6ff6088338a69497fd3b830c7518fb716efb71fa2", + "size": 200001024 + }, + "fineweb_train_000059.bin": { + "sha256": "dd1139356df3bdd7817142dca5b7ca50cd79bc34c2620f2834509e27f6cdb667", + "size": 200001024 + }, + "fineweb_train_000060.bin": { + "sha256": "4a85a829e9b368ea5858bacee81ef4384ef0f6b38752c29fcd9341d73c0313fb", + "size": 200001024 + }, + "fineweb_train_000061.bin": { + "sha256": "54d930cb17f118ca5f836ea9dd9466ba33f07440ff99bdafb9847f72a96e7678", + "size": 200001024 + }, + "fineweb_train_000062.bin": { + "sha256": "1549b9eb1a220cef41ba51b05dd72077c5891cd59d299d4a4fe2fca80261615e", + "size": 200001024 + }, + "fineweb_train_000063.bin": { + "sha256": "3be8259f9ccbd7aa54f2e9db93234b108433b7b1cad4b06146851d97f7fb75ff", + "size": 200001024 + }, + "fineweb_train_000064.bin": { + "sha256": "f761c21bc98668fcfbb8aa1abdc4fa3a2ba95b0985908f26a5744119f19a6e9f", + "size": 200001024 + }, + "fineweb_train_000065.bin": { + "sha256": "b2b70f36be820cc04deb6ac4d14e16a2f441d4dd6cccad8635b66585fe5f6573", + "size": 200001024 + }, + "fineweb_train_000066.bin": { + "sha256": "4aaeab489145b6382e1d207688d2e4828272cee16a12dcf244afbb4464fc75b5", + "size": 200001024 + }, + "fineweb_train_000067.bin": { + "sha256": "d27399790d7c266a85af75caede4539471859270e8b6b44605a6037e24b0edc0", + "size": 200001024 + }, + "fineweb_train_000068.bin": { + "sha256": "8bc7f0e297fc96de91b140019b91ec1d9d23418a9934700590299faeb7ca570e", + "size": 200001024 + }, + "fineweb_train_000069.bin": { + "sha256": "0b818240c430dc95f7d24c1b7a8f248c1d19198195979fa5ecd07a656b35b45a", + "size": 200001024 + }, + "fineweb_train_000070.bin": { + "sha256": "1b194abfc98b7d5a932fef0a3f03e2e3dbc6c3751bc637e67d637b865a00b8c1", + "size": 200001024 + }, + "fineweb_train_000071.bin": { + "sha256": "bf0fe37195fd37c3b547067182f5cf813a9d64851678d7132d96795429aa416f", + "size": 200001024 + }, + "fineweb_train_000072.bin": { + "sha256": "29f67c60bf5ec53a10e93668fcc6b95026226baf2eed3ca9b14116e734696f5c", + "size": 200001024 + }, + "fineweb_train_000073.bin": { + "sha256": "e7a3aab7ca19efdfb332b70b4f536a4542c80d56269329a24ebd336645e37357", + "size": 200001024 + }, + "fineweb_train_000074.bin": { + "sha256": "dea800cfbbcc83f8e21913f3a37d7de892a6b244c7a1cf9e0d9577c8ab4b266d", + "size": 200001024 + }, + "fineweb_train_000075.bin": { + "sha256": "5e9332d6726ba18942a04745afd1d6406bebe529b27e667f3e804d5b462bf0ee", + "size": 200001024 + }, + "fineweb_train_000076.bin": { + "sha256": "ca619222b623f9d32099460ed18ea18a703cb6c56a710b4000f0d8c2ffe0686f", + "size": 200001024 + }, + "fineweb_train_000077.bin": { + "sha256": "596354a15ad6340e77e8abeeb310017a7841f650cc1684af97e317429f48c66f", + "size": 200001024 + }, + "fineweb_train_000078.bin": { + "sha256": "d1b6d8401a5d6bc0abe2c6a68335b7675bbf5c47fd0de647c1f30c7b4381f2e9", + "size": 200001024 + }, + "fineweb_train_000079.bin": { + "sha256": "d70322f159670855e0ab43f31706342507f3102cdae1f7975d000515b2640225", + "size": 200001024 + }, + "fineweb_train_000080.bin": { + "sha256": "a8185bf5040dafd4f75986db1d6d788078ee57e73787e15929fd5798e3842686", + "size": 200001024 + }, + "fineweb_train_000081.bin": { + "sha256": "18b15d1b3630b2770e148f8a2cc9dbed43ac2dc860948dfaa8feade4ae12eb7f", + "size": 200001024 + }, + "fineweb_train_000082.bin": { + "sha256": "8206e5392e135d05d312d37f6bf56891232fe0cf625cbdff223308b2d2ada72d", + "size": 200001024 + }, + "fineweb_train_000083.bin": { + "sha256": "0b1020c30347e21f23c23d7e440b27a1c0e8ddda17d59e49672c553efcd221b0", + "size": 200001024 + }, + "fineweb_train_000084.bin": { + "sha256": "36460f35fb3951cc8b49c812fc47f4370ee7e114397e9ee0bdc2cbd9976a6189", + "size": 200001024 + }, + "fineweb_train_000085.bin": { + "sha256": "b4a119bdeb2168894c35f8f661d3d6d312ec3e13497ff9087b47ca7d943cf503", + "size": 200001024 + }, + "fineweb_train_000086.bin": { + "sha256": "b02fe674e6f10fc12325ba0408ca192bf1f44744c88f19fbfe8e3ffcffb7e703", + "size": 200001024 + }, + "fineweb_train_000087.bin": { + "sha256": "c52d26a38a2ebdd15fa630954700657aa9fc2b8aa9b6e19833dc24cf7c654d4e", + "size": 200001024 + }, + "fineweb_train_000088.bin": { + "sha256": "dffd1eb9ce74e51552d62037cb5ff9884365edd9a99c519ca751b437bfc38673", + "size": 200001024 + }, + "fineweb_train_000089.bin": { + "sha256": "ad2ff451b011a39fc6199168d146bf86b52d3c632451ef012b7c6c4ae6555f6f", + "size": 200001024 + }, + "fineweb_train_000090.bin": { + "sha256": "1a5f9bb0c410d51642ae187cc48f1bcceb3c681f9e89980c52e9daa7b3cefb0c", + "size": 200001024 + }, + "fineweb_train_000091.bin": { + "sha256": "9a7969791d92546f79104b0520d4c10161b64230d81926ee5cc989baf1d6c4d6", + "size": 200001024 + }, + "fineweb_train_000092.bin": { + "sha256": "93d2fc627b9eaac84886648e4c165b0fd04ec5276463e36fb49f4fbd944bcc09", + "size": 200001024 + }, + "fineweb_train_000093.bin": { + "sha256": "64114a56b5ed030e130f2e18e95ec194e1bc1fd9cb04cfb643ae4b667d87eba9", + "size": 200001024 + }, + "fineweb_train_000094.bin": { + "sha256": "9083070d2b15ab39c70109df302ee5136b76d667c054c2ebd974e1e4bf4cce70", + "size": 200001024 + }, + "fineweb_train_000095.bin": { + "sha256": "5a99c9c270aec1764f1097a6c947cfebf25ac086097edc964f7096ebe0490e65", + "size": 200001024 + }, + "fineweb_train_000096.bin": { + "sha256": "573f3c801af60b253b97536e8944ae2cb2e443dcf7fe9dd471722f56a9d9e75f", + "size": 200001024 + }, + "fineweb_train_000097.bin": { + "sha256": "5ead0115996c1e05107747477915ea9f56bc6070dd1f7064243ec91b16b58a7a", + "size": 200001024 + }, + "fineweb_train_000098.bin": { + "sha256": "1bcf260d507484f447538281c2c19847a0f38442cdada5ebe6f50458b76c4b20", + "size": 200001024 + }, + "fineweb_train_000099.bin": { + "sha256": "daf2f96329091338e2368991ac3fc71c2d5d048d47c039f8570da3f89b3be0c9", + "size": 200001024 + }, + "fineweb_train_000100.bin": { + "sha256": "8e534cc1857d9d93a7c006c59a7923b63ab004d7904799e45e1a158fc79fc572", + "size": 200001024 + }, + "fineweb_train_000101.bin": { + "sha256": "9b3871f7b5e73c80fdc5a428b45f2b17a97eec561ba1a13ea331a61c48201eda", + "size": 200001024 + }, + "fineweb_train_000102.bin": { + "sha256": "c9cd3cc4684cda81f45d157d2ac30d67bf4b06d9c7f6439179878715e237718f", + "size": 200001024 + }, + "fineweb_train_000103.bin": { + "sha256": "dd9817ed57a5fa872a4b6c28aac91e5a694465df4cb8cbbc170fa8ee2376eb5c", + "size": 110649110 + }, + "fineweb_val_000000.bin": { + "sha256": "5b95c8e0966f0861685b307b23dc5ae42b228ef74b28cb499784ae021f201640", + "size": 200001024 + } + } +} diff --git a/baseline/gpt2_small/speedrun30/reference_val.json b/baseline/gpt2_small/speedrun30/reference_val.json new file mode 100644 index 00000000..4106641d --- /dev/null +++ b/baseline/gpt2_small/speedrun30/reference_val.json @@ -0,0 +1,402 @@ +[ + { + "step": 0, + "tokens": 0, + "val_nll": 11.0092 + }, + { + "step": 250, + "tokens": 131072000, + "val_nll": 6.171 + }, + { + "step": 500, + "tokens": 262144000, + "val_nll": 5.3743 + }, + { + "step": 750, + "tokens": 393216000, + "val_nll": 4.7097 + }, + { + "step": 1000, + "tokens": 524288000, + "val_nll": 4.317 + }, + { + "step": 1250, + "tokens": 655360000, + "val_nll": 4.1296 + }, + { + "step": 1500, + "tokens": 786432000, + "val_nll": 4.006 + }, + { + "step": 1750, + "tokens": 917504000, + "val_nll": 3.9189 + }, + { + "step": 2000, + "tokens": 1048576000, + "val_nll": 3.8585 + }, + { + "step": 2250, + "tokens": 1179648000, + "val_nll": 3.8237 + }, + { + "step": 2500, + "tokens": 1310720000, + "val_nll": 3.7652 + }, + { + "step": 2750, + "tokens": 1441792000, + "val_nll": 3.7334 + }, + { + "step": 3000, + "tokens": 1572864000, + "val_nll": 3.6973 + }, + { + "step": 3250, + "tokens": 1703936000, + "val_nll": 3.6745 + }, + { + "step": 3500, + "tokens": 1835008000, + "val_nll": 3.6543 + }, + { + "step": 3750, + "tokens": 1966080000, + "val_nll": 3.6339 + }, + { + "step": 4000, + "tokens": 2097152000, + "val_nll": 3.6176 + }, + { + "step": 4250, + "tokens": 2228224000, + "val_nll": 3.6005 + }, + { + "step": 4500, + "tokens": 2359296000, + "val_nll": 3.5797 + }, + { + "step": 4750, + "tokens": 2490368000, + "val_nll": 3.5673 + }, + { + "step": 5000, + "tokens": 2621440000, + "val_nll": 3.5535 + }, + { + "step": 5250, + "tokens": 2752512000, + "val_nll": 3.5404 + }, + { + "step": 5500, + "tokens": 2883584000, + "val_nll": 3.5298 + }, + { + "step": 5750, + "tokens": 3014656000, + "val_nll": 3.5182 + }, + { + "step": 6000, + "tokens": 3145728000, + "val_nll": 3.5069 + }, + { + "step": 6250, + "tokens": 3276800000, + "val_nll": 3.4992 + }, + { + "step": 6500, + "tokens": 3407872000, + "val_nll": 3.4912 + }, + { + "step": 6750, + "tokens": 3538944000, + "val_nll": 3.481 + }, + { + "step": 7000, + "tokens": 3670016000, + "val_nll": 3.4726 + }, + { + "step": 7250, + "tokens": 3801088000, + "val_nll": 3.465 + }, + { + "step": 7500, + "tokens": 3932160000, + "val_nll": 3.4566 + }, + { + "step": 7750, + "tokens": 4063232000, + "val_nll": 3.4513 + }, + { + "step": 8000, + "tokens": 4194304000, + "val_nll": 3.4431 + }, + { + "step": 8250, + "tokens": 4325376000, + "val_nll": 3.4376 + }, + { + "step": 8500, + "tokens": 4456448000, + "val_nll": 3.4289 + }, + { + "step": 8750, + "tokens": 4587520000, + "val_nll": 3.4224 + }, + { + "step": 9000, + "tokens": 4718592000, + "val_nll": 3.4159 + }, + { + "step": 9250, + "tokens": 4849664000, + "val_nll": 3.4097 + }, + { + "step": 9500, + "tokens": 4980736000, + "val_nll": 3.4026 + }, + { + "step": 9750, + "tokens": 5111808000, + "val_nll": 3.3978 + }, + { + "step": 10000, + "tokens": 5242880000, + "val_nll": 3.3924 + }, + { + "step": 10250, + "tokens": 5373952000, + "val_nll": 3.3865 + }, + { + "step": 10500, + "tokens": 5505024000, + "val_nll": 3.3809 + }, + { + "step": 10750, + "tokens": 5636096000, + "val_nll": 3.3751 + }, + { + "step": 11000, + "tokens": 5767168000, + "val_nll": 3.3706 + }, + { + "step": 11250, + "tokens": 5898240000, + "val_nll": 3.3647 + }, + { + "step": 11500, + "tokens": 6029312000, + "val_nll": 3.3615 + }, + { + "step": 11750, + "tokens": 6160384000, + "val_nll": 3.3551 + }, + { + "step": 12000, + "tokens": 6291456000, + "val_nll": 3.3509 + }, + { + "step": 12250, + "tokens": 6422528000, + "val_nll": 3.3462 + }, + { + "step": 12500, + "tokens": 6553600000, + "val_nll": 3.3408 + }, + { + "step": 12750, + "tokens": 6684672000, + "val_nll": 3.3357 + }, + { + "step": 13000, + "tokens": 6815744000, + "val_nll": 3.3305 + }, + { + "step": 13250, + "tokens": 6946816000, + "val_nll": 3.3267 + }, + { + "step": 13500, + "tokens": 7077888000, + "val_nll": 3.3226 + }, + { + "step": 13750, + "tokens": 7208960000, + "val_nll": 3.3193 + }, + { + "step": 14000, + "tokens": 7340032000, + "val_nll": 3.3149 + }, + { + "step": 14250, + "tokens": 7471104000, + "val_nll": 3.3112 + }, + { + "step": 14500, + "tokens": 7602176000, + "val_nll": 3.3073 + }, + { + "step": 14750, + "tokens": 7733248000, + "val_nll": 3.3038 + }, + { + "step": 15000, + "tokens": 7864320000, + "val_nll": 3.3003 + }, + { + "step": 15250, + "tokens": 7995392000, + "val_nll": 3.2969 + }, + { + "step": 15500, + "tokens": 8126464000, + "val_nll": 3.2939 + }, + { + "step": 15750, + "tokens": 8257536000, + "val_nll": 3.2914 + }, + { + "step": 16000, + "tokens": 8388608000, + "val_nll": 3.2882 + }, + { + "step": 16250, + "tokens": 8519680000, + "val_nll": 3.2863 + }, + { + "step": 16500, + "tokens": 8650752000, + "val_nll": 3.2841 + }, + { + "step": 16750, + "tokens": 8781824000, + "val_nll": 3.2821 + }, + { + "step": 17000, + "tokens": 8912896000, + "val_nll": 3.2803 + }, + { + "step": 17250, + "tokens": 9043968000, + "val_nll": 3.2785 + }, + { + "step": 17500, + "tokens": 9175040000, + "val_nll": 3.277 + }, + { + "step": 17750, + "tokens": 9306112000, + "val_nll": 3.2758 + }, + { + "step": 18000, + "tokens": 9437184000, + "val_nll": 3.2746 + }, + { + "step": 18250, + "tokens": 9568256000, + "val_nll": 3.2738 + }, + { + "step": 18500, + "tokens": 9699328000, + "val_nll": 3.2732 + }, + { + "step": 18750, + "tokens": 9830400000, + "val_nll": 3.2727 + }, + { + "step": 19000, + "tokens": 9961472000, + "val_nll": 3.2724 + }, + { + "step": 19250, + "tokens": 10092544000, + "val_nll": 3.2722 + }, + { + "step": 19500, + "tokens": 10223616000, + "val_nll": 3.2722 + }, + { + "step": 19560, + "tokens": 10255073280, + "val_nll": 3.2722 + } +] diff --git a/baseline/gpt2_small/speedrun30/train.py b/baseline/gpt2_small/speedrun30/train.py new file mode 100644 index 00000000..6e911a9f --- /dev/null +++ b/baseline/gpt2_small/speedrun30/train.py @@ -0,0 +1,318 @@ +"""Bounded TPU port of the pinned llm.c GPT-2/FineWeb reference. + +No WeightWatcher, per-tensor diagnostics, preflight, or checkpoint upload loop. +The unmodified upstream model lives in vendor/. Hardware adaptation is here. +""" +import argparse +from contextlib import nullcontext +import hashlib +import json +import math +import os +from pathlib import Path +import time +import urllib.request + +import numpy as np +import torch +from vendor import llmc_train_gpt2 as reference + +HERE = Path(__file__).resolve().parent +BATCH = 64 # Global microbatch, 8 sequences per TPU chip +CONTEXT = 1024 +TOTAL_BATCH = 524288 # Same tokens/update as the reference +ACCUM = TOTAL_BATCH // (BATCH * CONTEXT) +TOTAL_STEPS = 19560 # Published Oct 13 2024 reference log +WARMUP = 700 +PEAK_LR = 0.0006 +VAL_TOKENS = 10 * 2**20 +SOURCE_COMMIT = "7ecd8906afe6ed7a2b2cdb731c042f26d525b820" + + +def write_json(path, obj): + path = Path(path) + tmp = path.with_suffix(".tmp") + tmp.write_text(json.dumps(obj, indent=2, allow_nan=False) + "\n") + tmp.replace(path) + + +def learning_rate(step): + if step < WARMUP: + return PEAK_LR * (step + 1) / WARMUP + ratio = min(1., (step - WARMUP) / (TOTAL_STEPS - WARMUP)) + return PEAK_LR * 0.5 * (1 + math.cos(math.pi * ratio)) + + +class DeadlineReached(Exception): + pass + + +class FineWeb: + """Download exact pretokenized benchmark shards lazily; never use Edu data.""" + def __init__(self, cache, deadline): + self.cache = Path(cache) + self.cache.mkdir(parents=True, exist_ok=True) + self.manifest = json.loads((HERE / "data_manifest.json").read_text()) + self.deadline = deadline + self.receipts = {} + + def array(self, name): + info = self.manifest["files"][name] + path = self.cache / name + receipt = path.with_suffix(".verified.json") + trusted = False + if path.exists() and receipt.exists(): + try: + saved = json.loads(receipt.read_text()) + trusted = (saved == info and path.stat().st_size == info["size"]) + except (OSError, ValueError): + pass + if not trusted: + url = ("https://huggingface.co/datasets/" + self.manifest["repo"] + + "/resolve/" + self.manifest["revision"] + "/" + name) + print("Downloading benchmark shard: " + name, flush=True) + temporary = path.with_suffix(".partial") + digest = hashlib.sha256() + if time.time() >= self.deadline: + raise DeadlineReached() + with urllib.request.urlopen(url, timeout=30) as response, temporary.open("wb") as out: + while True: + if time.time() >= self.deadline: + raise DeadlineReached() + chunk = response.read(4 * 1024 * 1024) + if not chunk: + break + out.write(chunk) + digest.update(chunk) + if temporary.stat().st_size != info["size"] or digest.hexdigest() != info["sha256"]: + raise RuntimeError("Benchmark download differs from pinned data: " + name) + temporary.replace(path) + write_json(receipt, info) + self.receipts[name] = info + with path.open("rb") as f: + header = np.fromfile(f, dtype=" len(self.tokens): + self.shard = (self.shard + 1) % len(self.names) + self.tokens = self.source.array(self.names[self.shard]) + self.position = 0 + buf = torch.from_numpy(np.array(self.tokens[self.position:self.position+count+1], dtype=np.int64)) + self.position += count + return buf[:-1].reshape(self.batch, self.context), buf[1:].reshape(self.batch, self.context) + + +def optimizer_for(model, device): + groups = [ + {"params": [p for p in model.parameters() if p.ndim >= 2], "weight_decay": 0.1}, + {"params": [p for p in model.parameters() if p.ndim < 2], "weight_decay": 0.0}, + ] + # Stock PyTorch AdamW. Capturable keeps bias-correction steps on the TPU; + # otherwise Python step constants can trigger a fresh XLA graph each update. + return torch.optim.AdamW(groups, lr=PEAK_LR, betas=(0.9, 0.95), eps=1e-8, + foreach=False, fused=False, capturable=device.type == "xla") + + +def clip_gradients(model): + # Equivalent global L2 clipping without a mixed-dtype torch.stack operation. + grads = [p.grad for p in model.parameters() if p.grad is not None] + norm = sum(g.detach().float().square().sum() for g in grads).sqrt() + scale = torch.clamp(1.0 / (norm + 1e-6), max=1.0) + for grad in grads: + grad.mul_(scale) + return norm + + +class Runtime: + def __init__(self, kind): + self.tpu = kind == "tpu" + if self.tpu: + import torch_xla.core.xla_model as xm + import torch_xla.runtime as xr + import torch_xla.distributed.spmd as xs + xr.use_spmd() + if xr.global_runtime_device_count() != 8 or xr.addressable_runtime_device_count() != 8: + raise RuntimeError("This reference runner requires the existing single-host 8-chip TPU.") + self.xm, self.xs = xm, xs + self.mesh = xs.Mesh(np.arange(8), (8,), ("data",)) + self.device = xm.xla_device() + else: + self.device = torch.device("cpu") + + def put(self, tensor): + value = tensor.to(self.device) + if self.tpu: + self.xs.mark_sharding(value, self.mesh, ("data",) + (None,) * (value.ndim-1)) + return value + + def replicate(self, tensor): + if self.tpu: + self.xs.mark_sharding(tensor, self.mesh, (None,) * tensor.ndim) + + def step(self, wait=False): + if self.tpu: + self.xm.mark_step() + if wait: + self.xm.wait_device_ops() + + def autocast(self): + return torch.autocast("xla", dtype=torch.bfloat16) if self.tpu else nullcontext() + + +def reference_bracket(step): + rows = json.loads((HERE / "reference_val.json").read_text()) + lo = max((r for r in rows if r["step"] <= step), key=lambda r:r["step"]) + hi = next((r for r in rows if r["step"] >= step), rows[-1]) + return {"published_lower": lo, "published_upper": hi, + "comparison": "reference curve; different seed/order and hardware, no automatic correctness verdict"} + + +@torch.no_grad() +def evaluate(model, tokens, rt, root, step, deadline, count=VAL_TOKENS): + model.eval() + total = 0. + evaluated = 0 + size = BATCH * CONTEXT + for offset in range(0, count, size): + if time.time() >= deadline: + break + buf = torch.from_numpy(np.array(tokens[offset:offset+size+1], dtype=np.int64)) + x = rt.put(buf[:-1].reshape(BATCH, CONTEXT)) + y = rt.put(buf[1:].reshape(BATCH, CONTEXT)) + with rt.autocast(): + _, loss = model(x, y, return_logits=False) + rt.step() + value = float(loss.detach().cpu()) + if not math.isfinite(value): + raise RuntimeError("Nonfinite validation loss") + total += value * size + evaluated += size + model.train() + row = {"kind": "validation", "step": step, "tokens_seen": step*TOTAL_BATCH, + "evaluation_tokens": evaluated, "full_benchmark_evaluation": evaluated == VAL_TOKENS, + "val_nll": total/evaluated if evaluated else None, + "val_perplexity": math.exp(min(total/evaluated, 700)) if evaluated else None, + **reference_bracket(step)} + with (root/"metrics.jsonl").open("a") as f: + f.write(json.dumps(row) + "\n") + write_json(root/"latest_validation.json", row) + print(json.dumps(row), flush=True) + return row + + +def main(): + p = argparse.ArgumentParser() + p.add_argument("--root", type=Path, required=True) + p.add_argument("--deadline", type=float, required=True) + p.add_argument("--cache", type=Path, required=True) + p.add_argument("--device", choices=("tpu", "cpu"), default="tpu") + a = p.parse_args() + root = a.root + source = FineWeb(a.cache, a.deadline-240) + start = time.time() + step = 0 + write_json(root/"status.json", {"status":"preparing_data", "deadline_unix":a.deadline}) + val = source.array("fineweb_val_000000.bin") + stream = TrainStream(source) + rt = Runtime(a.device) + torch.set_num_threads(4) + torch.manual_seed(42) + reference.FLASH = 0 # Upstream mathematical attention; CUDA flash kernels unavailable. + model = reference.GPT(reference.GPTConfig()).to(rt.device) + for tensor in (*model.parameters(), *model.buffers()): + rt.replicate(tensor) + optimizer = optimizer_for(model, rt.device) + manifest = {"source_repo":"karpathy/llm.c", "source_commit":SOURCE_COMMIT, + "model":"unchanged upstream GPT-2 124M, 12 layers, 12 heads, 768 width", + "parameters":sum(p.numel() for p in model.parameters()), + "batch_tokens":TOTAL_BATCH, "microbatch":BATCH, "accumulation":ACCUM, + "context":CONTEXT, "optimizer":"torch.optim.AdamW", "peak_lr":PEAK_LR, + "betas":[0.9,0.95], "epsilon":1e-8, "weight_decay":0.1, + "warmup_updates":WARMUP, "cosine_schedule_updates":TOTAL_STEPS, + "final_lr_fraction":0., "validation_tokens":VAL_TOKENS, + "data_repo":source.manifest["repo"], "data_revision":source.manifest["revision"], + "numerics":"FP32 weights/moments, BF16 autocast on TPU", + "differences_from_cuda_log":["PyTorch/XLA instead of CUDA C", "sequential Python loader instead of C shuffle", "different RNG/rounding; no bitwise equivalence claim"], + "time_limit_is_partial_run":True, "torch":torch.__version__} + write_json(root/"manifest.json", manifest) + print(json.dumps(manifest), flush=True) + rt.step(wait=True) + # A small initial probe avoids spending the training budget on step-zero evaluation. + evaluate(model, val, rt, root, 0, a.deadline-240, count=2**20) + write_json(root/"status.json", {"status":"training", "step":0, "deadline_unix":a.deadline}) + try: + while step < TOTAL_STEPS and time.time() < a.deadline-240: + began = time.monotonic() + optimizer.zero_grad(set_to_none=False) + loss_sum = 0. + for micro in range(ACCUM): + if time.time() >= a.deadline-240: + raise DeadlineReached() + x, y = stream.next_batch() + with rt.autocast(): + _, loss = model(rt.put(x), rt.put(y), return_logits=False) + (loss/ACCUM).backward() + loss_sum = loss_sum + loss.detach()/ACCUM + # Bound the lazy graph at each microbatch; this is not a host read. + rt.step() + for param in model.parameters(): + if param.grad is not None: + rt.replicate(param.grad) + clip_gradients(model) + lr = learning_rate(step) + for group in optimizer.param_groups: + group["lr"] = torch.tensor(lr).to(rt.device) if rt.tpu else lr + optimizer.step() + rt.step(wait=True) + step += 1 + elapsed = time.monotonic()-began + row = {"kind":"train", "step":step, "tokens_seen":step*TOTAL_BATCH, + "train_nll":float(loss_sum.cpu()), "lr":lr, "seconds":elapsed, + "tokens_per_second":TOTAL_BATCH/elapsed, "elapsed_seconds":time.time()-start} + if not math.isfinite(row["train_nll"]): + raise RuntimeError("Nonfinite scalar training loss") + with (root/"metrics.jsonl").open("a") as f: + f.write(json.dumps(row)+"\n") + write_json(root/"status.json", {"status":"training", **row}) + if step <= 5 or step % 10 == 0: + print(json.dumps(row), flush=True) + if step % 250 == 0: + evaluate(model, val, rt, root, step, a.deadline-240) + except DeadlineReached: + pass + # Only one final model checkpoint; no optimizer/checkpoint uploading in the hot loop. + write_json(root/"status.json", {"status":"final_evaluation", "step":step}) + result = evaluate(model, val, rt, root, step, a.deadline-45) + if time.time() < a.deadline-30: + print("Saving final model checkpoint", flush=True) + rt.step(wait=True) + weights = {name:tensor.detach().cpu() for name,tensor in model.state_dict().items()} + tmp = root/"model_final.tmp" + torch.save({"model":weights, "step":step, "model_config":vars(model.config), + "manifest":manifest, "validation":result, "resumable":False}, tmp) + tmp.replace(root/"model_final.pt") + write_json(root/"data_receipts.json", source.receipts) + write_json(root/"status.json", {"status":"finished_partial_reference_run", "step":step, + "tokens_seen":step*TOTAL_BATCH, "validation":result, + "full_training_recipe_completed":step == TOTAL_STEPS}) + + +if __name__ == "__main__": + main() diff --git a/baseline/gpt2_small/speedrun30/vendor/LICENSE b/baseline/gpt2_small/speedrun30/vendor/LICENSE new file mode 100644 index 00000000..4809e4c9 --- /dev/null +++ b/baseline/gpt2_small/speedrun30/vendor/LICENSE @@ -0,0 +1,22 @@ +MIT License + +Copyright (c) 2024 Andrej Karpathy + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. + diff --git a/baseline/gpt2_small/speedrun30/vendor/llmc_train_gpt2.py b/baseline/gpt2_small/speedrun30/vendor/llmc_train_gpt2.py new file mode 100644 index 00000000..b9dee870 --- /dev/null +++ b/baseline/gpt2_small/speedrun30/vendor/llmc_train_gpt2.py @@ -0,0 +1,860 @@ +""" +Reference code for GPT-2 training and inference. +Will save the model weights into files, to be read from C as initialization. + +References: +1) the official GPT-2 TensorFlow implementation released by OpenAI: +https://github.com/openai/gpt-2/blob/master/src/model.py +2) huggingface/transformers PyTorch implementation: +https://github.com/huggingface/transformers/blob/main/src/transformers/models/gpt2/modeling_gpt2.py + +Example launches to only benchmark the speed of bfloat16 compiled GPU training: +1 GPU: +python train_gpt2.py --write_tensors=0 --num_iterations=50 --sequence_length=1024 --compile=1 --tensorcores=1 --dtype=bfloat16 +you can also turn on flash-attention by appending --flash=1 +4 GPU: +torchrun --standalone --nproc_per_node=4 train_gpt2.py --write_tensors=0 --num_iterations=50 --sequence_length=1024 --compile=1 --tensorcores=1 --dtype=bfloat16 +""" + +import os +import math +import glob +import struct +import inspect +from contextlib import nullcontext +from dataclasses import dataclass + +import numpy as np +import torch +import torch.nn as nn +from torch.nn import functional as F +import torch._inductor.config as config +from torch.nn.parallel import DistributedDataParallel as DDP +from torch.distributed import init_process_group, destroy_process_group +from torch.distributed.optim import ZeroRedundancyOptimizer +import torch.distributed as dist + +# ----------------------------------------------------------------------------- +# PyTorch nn.Module definitions for the GPT-2 model + +class NewGELU(nn.Module): + """Careful there are a few versions of GeLU, this one is the exact one used by OpenAI""" + def forward(self, input): + return 0.5 * input * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi) * (input + 0.044715 * torch.pow(input, 3.0)))) + +# using a global to toggle flash-attention +FLASH = 0 + +class CausalSelfAttention(nn.Module): + + def __init__(self, config): + super().__init__() + assert config.n_embd % config.n_head == 0 + # key, query, value projections for all heads, but in a batch + self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd) + # output projection + self.c_proj = nn.Linear(config.n_embd, config.n_embd) + self.c_proj.LLMC_RESIDUAL_SCALE_FLAG = 1 + # regularization + self.n_head = config.n_head + self.n_embd = config.n_embd + # not really a 'bias', more of a mask, but following the OpenAI/HF naming though + self.register_buffer("bias", torch.tril(torch.ones(config.block_size, config.block_size)) + .view(1, 1, config.block_size, config.block_size)) + + def forward(self, x): + B, T, C = x.size() # batch size, sequence length, embedding dimensionality (n_embd) + # calculate query, key, values for all heads in batch and move head forward to be the batch dim + qkv = self.c_attn(x) + q, k, v = qkv.split(self.n_embd, dim=2) + k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs) + q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs) + v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs) + if FLASH: + # flashattention + y = F.scaled_dot_product_attention(q, k, v, is_causal=True) + else: + # manual implementation of attention + # this materializes the large (T,T) matrix for all the queries and keys + att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1))) + att = att.masked_fill(self.bias[:,:,:T,:T] == 0, float('-inf')) + att = F.softmax(att, dim=-1) + y = att @ v # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs) + y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side + # output projection + y = self.c_proj(y) + return y + +class MLP(nn.Module): + + def __init__(self, config): + super().__init__() + self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd) + self.gelu = NewGELU() + self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd) + self.c_proj.LLMC_RESIDUAL_SCALE_FLAG = 1 + + def forward(self, x): + x = self.c_fc(x) + x = self.gelu(x) + x = self.c_proj(x) + return x + +class Block(nn.Module): + + def __init__(self, config): + super().__init__() + self.ln_1 = nn.LayerNorm(config.n_embd) + self.attn = CausalSelfAttention(config) + self.ln_2 = nn.LayerNorm(config.n_embd) + self.mlp = MLP(config) + + def forward(self, x): + x = x + self.attn(self.ln_1(x)) + x = x + self.mlp(self.ln_2(x)) + return x + +# ----------------------------------------------------------------------------- +# The main GPT-2 model + +@dataclass +class GPTConfig: + block_size: int = 1024 + vocab_size: int = 50257 + n_layer: int = 12 + n_head: int = 12 + n_embd: int = 768 + +class GPT(nn.Module): + + def __init__(self, config): + super().__init__() + self.config = config + + self.transformer = nn.ModuleDict(dict( + wte = nn.Embedding(config.vocab_size, config.n_embd), + wpe = nn.Embedding(config.block_size, config.n_embd), + h = nn.ModuleList([Block(config) for _ in range(config.n_layer)]), + ln_f = nn.LayerNorm(config.n_embd), + )) + self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False) + self.lm_head.LLMC_SKIP_INIT = 1 # don't init this one, we will tie weights + self.transformer.wte.weight = self.lm_head.weight # https://paperswithcode.com/method/weight-tying + + # init all weights, use a torch rng object to be very careful + self.init_rng = torch.Generator() + self.init_rng.manual_seed(42) + self.apply(self._init_weights) + + def _init_weights(self, module): + if isinstance(module, nn.Linear): + # apply special scaled init to the residual projections, per GPT-2 paper + std = 0.02 if not hasattr(module, 'LLMC_RESIDUAL_SCALE_FLAG') else 0.02/math.sqrt(2 * self.config.n_layer) + # we want to skip initializing lm_head, which shares parameters with wte + # and wte was already initialized down below during the Embedding init + if not hasattr(module, 'LLMC_SKIP_INIT'): + torch.nn.init.normal_(module.weight, mean=0.0, std=std, generator=self.init_rng) + if module.bias is not None: + torch.nn.init.zeros_(module.bias) + elif isinstance(module, nn.Embedding): + torch.nn.init.normal_(module.weight, mean=0.0, std=0.02, generator=self.init_rng) + + def forward(self, idx, targets=None, return_logits=True): + device = idx.device + b, t = idx.size() + assert t <= self.config.block_size, f"Cannot forward sequence of length {t}, block size is only {self.config.block_size}" + pos = torch.arange(0, t, dtype=torch.long, device=device) # shape (t) + + # forward the GPT model itself + tok_emb = self.transformer.wte(idx) # token embeddings of shape (b, t, n_embd) + pos_emb = self.transformer.wpe(pos) # position embeddings of shape (t, n_embd) + x = tok_emb + pos_emb + + for block in self.transformer.h: + x = block(x) + x = self.transformer.ln_f(x) + + if targets is not None: + # if we are given some desired targets also calculate the loss + logits = self.lm_head(x) + loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1) + else: + # inference-time mini-optimization: only forward the lm_head on the very last position + logits = self.lm_head(x[:, [-1], :]) # note: using list [-1] to preserve the time dim + loss = None + + # there are performance reasons why not returning logits is prudent, if not needed + if not return_logits: + logits = None + + return logits, loss + + @classmethod + def from_pretrained(cls, model_type): + """Loads pretrained GPT-2 model weights from huggingface""" + assert model_type in {'gpt2', 'gpt2-medium', 'gpt2-large', 'gpt2-xl'} + from transformers import GPT2LMHeadModel + print("loading weights from pretrained gpt: %s" % model_type) + + # n_layer, n_head and n_embd are determined from model_type + config_args = { + 'gpt2': dict(n_layer=12, n_head=12, n_embd=768), # 124M params + 'gpt2-medium': dict(n_layer=24, n_head=16, n_embd=1024), # 350M params + 'gpt2-large': dict(n_layer=36, n_head=20, n_embd=1280), # 774M params + 'gpt2-xl': dict(n_layer=48, n_head=25, n_embd=1600), # 1558M params + }[model_type] + config_args['vocab_size'] = 50257 # always 50257 for GPT model checkpoints + config_args['block_size'] = 1024 # always 1024 for GPT model checkpoints + # create a from-scratch initialized minGPT model + config = GPTConfig(**config_args) + model = GPT(config) + sd = model.state_dict() + sd_keys = sd.keys() + sd_keys = [k for k in sd_keys if not k.endswith('.attn.bias')] # discard this mask / buffer, not a param + + # init a huggingface/transformers model + model_hf = GPT2LMHeadModel.from_pretrained(model_type) + sd_hf = model_hf.state_dict() + + # copy while ensuring all of the parameters are aligned and match in names and shapes + sd_keys_hf = sd_hf.keys() + sd_keys_hf = [k for k in sd_keys_hf if not k.endswith('.attn.masked_bias')] # ignore these, just a buffer + sd_keys_hf = [k for k in sd_keys_hf if not k.endswith('.attn.bias')] # same, just the mask (buffer) + transposed = ['attn.c_attn.weight', 'attn.c_proj.weight', 'mlp.c_fc.weight', 'mlp.c_proj.weight'] + # basically the openai checkpoints use a "Conv1D" module, but we only want to use a vanilla Linear + # this means that we have to transpose these weights when we import them + assert len(sd_keys_hf) == len(sd_keys), f"mismatched keys: {len(sd_keys_hf)} != {len(sd_keys)}" + for k in sd_keys_hf: + if any(k.endswith(w) for w in transposed): + # special treatment for the Conv1D weights we need to transpose + assert sd_hf[k].shape[::-1] == sd[k].shape + with torch.no_grad(): + sd[k].copy_(sd_hf[k].t()) + else: + # vanilla copy over the other parameters + assert sd_hf[k].shape == sd[k].shape + with torch.no_grad(): + sd[k].copy_(sd_hf[k]) + + return model + + def configure_optimizers(self, weight_decay, learning_rate, betas, device_type, zero_stage): + # start with all of the candidate parameters + param_dict = {pn: p for pn, p in self.named_parameters()} + # filter out those that do not require grad + param_dict = {pn: p for pn, p in param_dict.items() if p.requires_grad} + # create optim groups. Any parameters that is 2D will be weight decayed, otherwise no. + # i.e. all weight tensors in matmuls + embeddings decay, all biases and layernorms don't. + decay_params = [p for n, p in param_dict.items() if p.dim() >= 2] + nodecay_params = [p for n, p in param_dict.items() if p.dim() < 2] + optim_groups = [ + {'params': decay_params, 'weight_decay': weight_decay}, + {'params': nodecay_params, 'weight_decay': 0.0} + ] + num_decay_params = sum(p.numel() for p in decay_params) + num_nodecay_params = sum(p.numel() for p in nodecay_params) + print0(f"num decayed parameter tensors: {len(decay_params)}, with {num_decay_params:,} parameters") + print0(f"num non-decayed parameter tensors: {len(nodecay_params)}, with {num_nodecay_params:,} parameters") + # Create AdamW optimizer and use the fused version if it is available + fused_available = 'fused' in inspect.signature(torch.optim.AdamW).parameters + use_fused = fused_available and device_type == 'cuda' + print0(f"using fused AdamW: {use_fused}") + if zero_stage == 1: + print0("using ZeroRedundancyOptimizer") + optimizer = ZeroRedundancyOptimizer(**optim_groups[0], optimizer_class=torch.optim.AdamW, + lr=learning_rate, betas=betas, fused=use_fused) + optimizer.add_param_group(optim_groups[1]) + else: + print0("using regular AdamW") + optimizer = torch.optim.AdamW(optim_groups, lr=learning_rate, betas=betas, fused=use_fused) + return optimizer + + @torch.no_grad() + def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None): + """ + Take a conditioning sequence of indices idx (LongTensor of shape (b,t)) and complete + the sequence max_new_tokens times, feeding the predictions back into the model each time. + Most likely you'll want to make sure to be in model.eval() mode of operation for this. + """ + for _ in range(max_new_tokens): + # if the sequence context is growing too long we must crop it at block_size + idx_cond = idx if idx.size(1) <= self.config.block_size else idx[:, -self.config.block_size:] + # forward the model to get the logits for the index in the sequence + logits, _ = self(idx_cond) + # pluck the logits at the final step and scale by desired temperature + logits = logits[:, -1, :] / temperature + # optionally crop the logits to only the top k options + if top_k is not None: + v, _ = torch.topk(logits, min(top_k, logits.size(-1))) + logits[logits < v[:, [-1]]] = -float('Inf') + # apply softmax to convert logits to (normalized) probabilities + probs = F.softmax(logits, dim=-1) + # sample from the distribution + idx_next = torch.multinomial(probs, num_samples=1) + # append sampled index to the running sequence and continue + idx = torch.cat((idx, idx_next), dim=1) + + return idx + +# ----------------------------------------------------------------------------- +# Our own simple Distributed Data Loader + +def _peek_data_shard(filename): + # only reads the header, returns header data + with open(filename, "rb") as f: + # first read the header, which is 256 int32 integers (4 bytes each) + header = np.frombuffer(f.read(256*4), dtype=np.int32) + if header[0] != 20240520: + print("ERROR: magic number mismatch in the data .bin file!") + print("---> HINT: Are you passing in a correct file with --input_bin?") + print("---> HINT: Dataset encoding changed recently, re-run data prepro or refer again to README") + print("---> HINT: For example re-run: `python dev/data/tinyshakespeare.py`, then re-try") + exit(1) + assert header[1] == 1, "unsupported version" + ntok = header[2] # number of tokens (claimed) + return ntok # for now just return the number of tokens + +def _load_data_shard(filename): + with open(filename, "rb") as f: + # first read the header, which is 256 int32 integers (4 bytes each) + header = np.frombuffer(f.read(256*4), dtype=np.int32) + assert header[0] == 20240520, "magic number mismatch in the data .bin file" + assert header[1] == 1, "unsupported version" + ntok = header[2] # number of tokens (claimed) + # the rest of it are tokens, stored as uint16 + tokens = np.frombuffer(f.read(), dtype=np.uint16) + assert len(tokens) == ntok, "number of tokens read does not match header?" + return tokens + +class DistributedDataLoader: + def __init__(self, filename_pattern, B, T, process_rank, num_processes): + self.process_rank = process_rank + self.num_processes = num_processes + self.B = B + self.T = T + + # glob files that match the pattern + self.files = sorted(glob.glob(filename_pattern)) + assert len(self.files) > 0, f"did not find any files that match the pattern {filename_pattern}" + + # load and validate all data shards, count number of tokens in total + ntok_total = 0 + for fname in self.files: + shard_ntok = _peek_data_shard(fname) + assert shard_ntok >= num_processes * B * T + 1 + ntok_total += shard_ntok + self.ntok_total = ntok_total + print0(f"DataLoader: total number of tokens: {ntok_total:,} across {len(self.files)} files") + + # kick things off + self.current_shard = None + self.reset() + + def reset(self): + # we're being a bit clever here: if we already had shard 0 loaded, + # then don't do the work to reload it, just reset the pointer + if self.current_shard != 0: + self.current_shard = 0 + self.tokens = _load_data_shard(self.files[self.current_shard]) + self.current_position = self.process_rank * self.B * self.T + + def advance(self): # advance to next data shard + self.current_shard = (self.current_shard + 1) % len(self.files) + self.current_position = self.process_rank * self.B * self.T + self.tokens = _load_data_shard(self.files[self.current_shard]) + + def next_batch(self): + B = self.B + T = self.T + buf = self.tokens[self.current_position : self.current_position+B*T+1] + buf = torch.tensor(buf.astype(np.int32), dtype=torch.long) + x = (buf[:-1]).view(B, T) # inputs + y = (buf[1:]).view(B, T) # targets + # advance the start pointer in current shard + self.current_position += B * T * self.num_processes + # if loading the next batch would be out of bounds advance the shard + if self.current_position + (B * T * self.num_processes + 1) > len(self.tokens): + self.advance() + return x, y + +# ----------------------------------------------------------------------------- +# Python -> C bridge utilities for saving params/grads/activations to .bin files + +def write_fp32(tensor, file): + t = tensor.detach().cpu().to(torch.float32) + b = t.numpy().tobytes() + file.write(b) + +def write_bf16(tensor, file): + t = tensor.detach().cpu().to(torch.bfloat16) + # numpy doesn't have bf16 datatype so we have to trick it + t = t.view(torch.int16) # trick: reinterpret as int16 + b = t.numpy().tobytes() + file.write(b) + +def write_tensors(model_tensors, L, file, dtype): + # writes the GPT-2 model's weights to a binary file + assert dtype in {"float32", "bfloat16"} + write_fun = write_fp32 if dtype == "float32" else write_bf16 + write_fun(model_tensors["transformer.wte.weight"], file) # (V, C) + write_fun(model_tensors["transformer.wpe.weight"], file) # (T, C) + for i in range(L): # (L, C) + write_fun(model_tensors[f"transformer.h.{i}.ln_1.weight"], file) + for i in range(L): # (L, C) + write_fun(model_tensors[f"transformer.h.{i}.ln_1.bias"], file) + for i in range(L): # (L, 3C, C) + write_fun(model_tensors[f"transformer.h.{i}.attn.c_attn.weight"], file) + for i in range(L): # (L, 3C) + write_fun(model_tensors[f"transformer.h.{i}.attn.c_attn.bias"], file) + for i in range(L): # (L, C, C) + write_fun(model_tensors[f"transformer.h.{i}.attn.c_proj.weight"], file) + for i in range(L): # (L, C) + write_fun(model_tensors[f"transformer.h.{i}.attn.c_proj.bias"], file) + for i in range(L): # (L, C) + write_fun(model_tensors[f"transformer.h.{i}.ln_2.weight"], file) + for i in range(L): # (L, C) + write_fun(model_tensors[f"transformer.h.{i}.ln_2.bias"], file) + for i in range(L): # (L, 4C, C) + write_fun(model_tensors[f"transformer.h.{i}.mlp.c_fc.weight"], file) + for i in range(L): # (L, 4C) + write_fun(model_tensors[f"transformer.h.{i}.mlp.c_fc.bias"], file) + for i in range(L): # (L, C, 4C) + write_fun(model_tensors[f"transformer.h.{i}.mlp.c_proj.weight"], file) + for i in range(L): # (L, C) + write_fun(model_tensors[f"transformer.h.{i}.mlp.c_proj.bias"], file) + write_fun(model_tensors["transformer.ln_f.weight"], file) # (C, ) + write_fun(model_tensors["transformer.ln_f.bias"], file) # (C, ) + +@torch.no_grad() +def pad_vocab(tensor, multiple=128, value=0): + """ + The dimension of the vocab size in GPT-2 is 50,257 + which is unfortunately a very unfriendly number for a lot of + matrix operations on the GPU. So we pad it to the nearest + friendlier multiple, e.g. 50,304 if multiple=128 when we + export the weights into C land. This is a NOOP algorithmically + and is only done to make the tensor operations more efficient. + """ + assert tensor.ndim == 2 + V, C = tensor.shape + assert V == 50257, "just being defensive here" + # calculate padded vocab size by rounding up to nearest multiple + Vp = ((V + multiple - 1) // multiple) * multiple + # pad the tensor + pad_rows = Vp - V + padded = tensor if pad_rows == 0 else F.pad(tensor, (0, 0, 0, pad_rows), value=value) + assert padded.shape == (Vp, C) + return padded + +def write_model(model, filename, dtype): + # everything we need to instantiate the model + # 1) header is: version int, GPTConfig ints, padding to 1024 bytes + assert dtype in {"float32", "bfloat16"} # float16 todo maybe later + version = { + "float32": 3, # 3: all tensors are fp32, padded vocab + "bfloat16": 5, # 5: all tensors are bf16, padded vocab + }[dtype] + header = torch.zeros(256, dtype=torch.int32) + header[0] = 20240326 # magic + header[1] = version # checkpoint version + header[2] = model.config.block_size + header[3] = model.config.vocab_size + header[4] = model.config.n_layer + header[5] = model.config.n_head + header[6] = model.config.n_embd + # 2) the parameters follow the header + params = {name: param.cpu() for name, param in model.named_parameters()} + # pad the vocab to a multiple of 128 here at export, for efficiency in C + wte = params["transformer.wte.weight"] # (V, C) + wte_padded = pad_vocab(wte) # (Vp, C) + params["transformer.wte.weight"] = wte_padded # (Vp, C) + print(f"padded vocab size from {wte.size(0)} to {wte_padded.size(0)}") + header[7] = wte_padded.size(0) # padded vocab size store in header + # now write to file + with open(filename, "wb") as file: + file.write(header.numpy().tobytes()) # header + write_tensors(params, model.config.n_layer, file, dtype) # params + print(f"wrote {filename}") + +def write_state(model, x, y, logits, loss, filename): + # the state is used for debugging. + # it contains information about the input, logits, loss, and the parameter gradients + # this can be used for checking the computation correctness in C + header = torch.zeros(256, dtype=torch.int32) + header[0] = 20240327 # magic + header[1] = 2 # run state version = 2 (1 -> 2 for padded vocab changes) + header[2] = x.size(0) # batch size of the batch, B + header[3] = x.size(1) # temporal extent of the batch, T + grads = {name: param.grad.cpu() for name, param in model.named_parameters()} + # pad the vocab grads here as well, to mirror write_model + wte_grad = grads["transformer.wte.weight"] # (V, C) + wte_grad_padded = pad_vocab(wte_grad, value=0) # (Vp, C) # TODO later maybe pad with nan? + grads["transformer.wte.weight"] = wte_grad_padded # (Vp, C) + print(f"padded vocab size in reference grads from {wte_grad.size(0)} to {wte_grad_padded.size(0)}") + with open(filename, "wb") as file: + # header + file.write(header.numpy().tobytes()) + # input x + file.write(x.cpu().numpy().astype("int32").tobytes()) # (B, T) + # targets y + file.write(y.cpu().numpy().astype("int32").tobytes()) # (B, T) + # logits (result of the model forward pass) + write_fp32(logits.cpu(), file) + # loss (single float, result of the cross entropy loss) + write_fp32(loss.cpu(), file) + # gradients + write_tensors(grads, model.config.n_layer, file, "float32") + print(f"wrote {filename}") + +def write_tokenizer(enc, filename): + n = enc.max_token_value + 1 + header = torch.zeros(256, dtype=torch.int32) + header[0] = 20240328 # magic + header[1] = 2 # tokenizer version = 2 (1 -> 2: includes EOT token) + header[2] = n # number of tokens + header[3] = enc.eot_token # EOT token + with open(filename, "wb") as file: + file.write(header.numpy().tobytes()) + for i in range(n): + b = enc.decode_bytes([i]) + length = len(b) + assert length < 256, f"Token length exceeds 255: {length}" + file.write(struct.pack(" C bridge + parser.add_argument("--write_tensors", type=int, default=1, help="write tensors to disk") + args = parser.parse_args() + + # args error checking and convenience variables + B, T = args.batch_size, args.sequence_length + assert 1 <= T <= 1024 + assert args.dtype in {"float32", "float16", "bfloat16"} + assert args.model in {"gpt2", "gpt2-medium", "gpt2-large", "gpt2-xl", "d12", "d24", "d36", "d48"} + + # set up DDP (distributed data parallel). torchrun sets this env variable + ddp = int(os.environ.get('RANK', -1)) != -1 # is this a ddp run? + if ddp: + # use of DDP atm demands CUDA, we set the device appropriately according to rank + assert torch.cuda.is_available(), "for now i think we need CUDA for DDP" + init_process_group(backend='nccl') + ddp_rank = int(os.environ['RANK']) + ddp_local_rank = int(os.environ['LOCAL_RANK']) + ddp_world_size = int(os.environ['WORLD_SIZE']) + device = f'cuda:{ddp_local_rank}' + torch.cuda.set_device(device) + master_process = ddp_rank == 0 # this process will do logging, checkpointing etc. + seed_offset = 0 # each process gets the exact same seed + zero_stage = args.zero_stage + else: + ddp_rank = 0 + ddp_local_rank = 0 + zero_stage = 0 + ddp_world_size = 1 + master_process = True + seed_offset = 0 + # select the device + if args.device: + # provided explicitly by the user + device = args.device + else: + # attempt to autodetect the device + device = "cpu" + if torch.cuda.is_available(): + device = "cuda" + elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): + device = "mps" + print(f"using device: {device}") + device_type = 'cuda' if 'cuda' in device else 'cpu' + + # calculate gradient accumulation from the desired total batch size and the current run configuration + tokens_per_fwdbwd = B * T * ddp_world_size + assert args.total_batch_size % tokens_per_fwdbwd == 0 + grad_accum_steps = args.total_batch_size // tokens_per_fwdbwd + print0(f"total desired batch size: {args.total_batch_size}") + print0(f"=> calculated gradient accumulation steps: {grad_accum_steps}") + + # set up a context manager following the desired dtype and device + ptdtype = {'float32': torch.float32, 'bfloat16': torch.bfloat16, 'float16': torch.float16}[args.dtype] + ctx = torch.amp.autocast(device_type=device_type, dtype=ptdtype) if device_type == "cuda" else nullcontext() + + # rng / reproducibility + torch.manual_seed(42) + if torch.cuda.is_available(): + torch.cuda.manual_seed(42) + + # set the torch precision mode to use TensorFloat32 (TF32) for matmuls + # docs https://pytorch.org/docs/stable/generated/torch.set_float32_matmul_precision.html + if args.tensorcores: + torch.set_float32_matmul_precision('high') + + # turn on/off flash attention + assert args.flash in {0, 1} + FLASH = args.flash + + # init (and write) the tokenizer + enc = tiktoken.get_encoding("gpt2") + if master_process and args.write_tensors: # tokenizer is technically not tensors but ok + write_tokenizer(enc, "gpt2_tokenizer.bin") + + # init the model, either from scratch or from OpenAI pretrained checkpoint + if args.model[0] == "d": + # from scratch (random weights) + model_config = { + "d12": GPTConfig(block_size=1024, vocab_size=50257, n_layer=12, n_head=12, n_embd=768), + "d24": GPTConfig(block_size=1024, vocab_size=50257, n_layer=24, n_head=16, n_embd=1024), + "d36": GPTConfig(block_size=1024, vocab_size=50257, n_layer=36, n_head=20, n_embd=1280), + "d48": GPTConfig(block_size=1024, vocab_size=50257, n_layer=48, n_head=25, n_embd=1600), + }[args.model] + model = GPT(model_config) + else: + # load the GPT-2 model weights + model = GPT.from_pretrained(args.model) + model.train() + model.to(device) + if args.compile: + if hasattr(config, "coordinate_descent_tuning"): + config.coordinate_descent_tuning = True # suggested by @Chillee + print0("compiling the model...") + model = torch.compile(model) + + # ------------------------------------------------------------------------- + # Our own version of a simple DistributedDataLoader + + # load tokens + train_loader = DistributedDataLoader(args.input_bin, B, T, ddp_rank, ddp_world_size) + val_loader = None + if args.input_val_bin: + val_loader = DistributedDataLoader(args.input_val_bin, B, T, ddp_rank, ddp_world_size) + + # ------------------------------------------------------------------------- + # PyTorch -> C bridge: save some weights and state for C to load later as reference + + # do one forward pass to generate ground truth for our C tests + if master_process and args.write_tensors and (not args.inference_only): + x, y = train_loader.next_batch() + x, y = x.to(device), y.to(device) + logits, loss = model(x, y) + loss.backward() + # save model params, in both float32 and bfloat16 + model_to_size = {"gpt2": "124M", "gpt2-medium": "355M", "gpt2-large": "774M", "gpt2-xl": "1558M"} + model_to_size.update({f"d{d}": f"d{d}" for d in [12, 24, 36, 48]}) + model_size_str = model_to_size[args.model] # e.g. "124M", or "d12" + write_model(model, f"gpt2_{model_size_str}.bin", dtype="float32") + write_model(model, f"gpt2_{model_size_str}_bf16.bin", dtype="bfloat16") + # save x, y, logits, loss, and parameter gradients, for debugging C + # always store these in fp32 to have an accurate reference (?) + write_state(model, x, y, logits, loss, f"gpt2_{model_size_str}_debug_state.bin") + # reset the train_loader for the optimization below + train_loader.reset() + + # ------------------------------------------------------------------------- + # main training loop + + # here we wrap model into DDP container + if ddp: + model = DDP(model, device_ids=[ddp_local_rank]) + raw_model = model.module if ddp else model # always contains the "raw" unwrapped model + + # init the optimizer + optimizer = raw_model.configure_optimizers(weight_decay=args.weight_decay, + learning_rate=args.learning_rate, betas=(0.9, 0.95), + device_type=device, zero_stage=zero_stage) + + # learning rate decay scheduler (cosine with warmup) + def get_lr(it): + min_lr = args.learning_rate * args.learning_rate_decay_frac + # 1) linear warmup for warmup_iters steps + if it < args.warmup_iters: + return args.learning_rate * (it+1) / args.warmup_iters + # 2) if it > lr_decay_iters, return min learning rate + if it > args.num_iterations: + return min_lr + # 3) in between, use cosine decay down to min learning rate + decay_ratio = (it - args.warmup_iters) / (args.num_iterations - args.warmup_iters) + assert 0 <= decay_ratio <= 1 + coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio)) # coeff starts at 1 and goes to 0 + return min_lr + coeff * (args.learning_rate - min_lr) + + # create the logging directory if it does not exist + logfile = None + if args.output_dir: + os.makedirs(args.output_dir, exist_ok=True) + logfile = os.path.join(args.output_dir, "main.log") + # create the log file "main.log" inside it, and wipe it clean + with open(logfile, "w") as f: + pass + + if device == "cuda": + torch.cuda.reset_peak_memory_stats() + timings = [] + norm = -1.0 # dummy value to print in inference-only mode + for step in range(args.num_iterations + 1): + t0 = time.time() + last_step = (step == args.num_iterations) + + # once in a while evaluate the validation dataset + if (args.val_loss_every > 0 \ + and (step % args.val_loss_every == 0 or last_step)) \ + and (val_loader is not None): + model.eval() + val_loader.reset() + with torch.no_grad(): + val_loss = 0.0 + for _ in range(args.val_max_steps): + x, y = val_loader.next_batch() + x, y = x.to(device), y.to(device) + _, loss = model(x, y, return_logits=False) + val_loss += loss.item() + val_loss /= args.val_max_steps + # log to console and to file + print0(f"val loss {val_loss}") + if master_process and logfile is not None: + with open(logfile, "a") as f: + f.write("s:%d tel:%f\n" % (step, val_loss)) + + # once in a while perform model inference on the master process + if (args.sample_every > 0 \ + and (step % args.sample_every == 0 or last_step)) \ + and master_process: + model.eval() + # before we end, let's also do one round of inference + # we'll kick off the generation with "<|endoftext|>", which designates the start of a new sequence + start_ids = [enc.eot_token] + xg = (torch.tensor(start_ids, dtype=torch.long, device=device)[None, ...]) + max_new_tokens = 32 + temperature = 1.0 + top_k = 40 + yg = raw_model.generate(xg, max_new_tokens, temperature=temperature, top_k=top_k) + print0('---------------') + print0(enc.decode(yg[0].tolist())) + print0('---------------') + + # bit confusing: we want to make sure to eval and sample on 0th iteration + # but also after the very last iteration. so we loop for step <= num_iterations + # instead of just < num_iterations (one extra due to <=), only to do + # the validation/sampling one last time, and then we break right here as we're done. + if last_step: + break + + # --------------- TRAINING SECTION BEGIN ----------------- + model.train() + optimizer.zero_grad(set_to_none=True) + # if we are trying to overfit a single batch, we reset the loader here + if args.overfit_single_batch: + train_loader.reset() + # micro-batch loop where we do gradient accumulation to reach desired total batch size + lossf = 0.0 # for getting the mean loss (as simple float) over the accumulation steps + for micro_step in range(grad_accum_steps): + # fetch a batch + x, y = train_loader.next_batch() + x, y = x.to(device), y.to(device) + if ddp: + # we want only the last micro-step to sync grads in a DDP model + # the official way to do this is with model.no_sync(), but that is a + # context manager that bloats the code, so we just toggle this variable + model.require_backward_grad_sync = (micro_step == grad_accum_steps - 1) + # forward pass + with ctx: + _, loss = model(x, y, return_logits=False) + # we have to scale the loss to account for gradient accumulation, + # because the gradients just add on each successive backward(). + # addition of gradients corresponds to a SUM in the objective, but + # instead of a SUM we want MEAN, so we scale the loss here + loss = loss / grad_accum_steps + lossf += loss.detach() # keep track of the mean loss + # backward pass + if not args.inference_only: + loss.backward() + if ddp: + dist.all_reduce(lossf, op=dist.ReduceOp.AVG) + lossf = lossf.item() + norm = torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip) + # determine and set the learning rate for this iteration + lr = get_lr(step) + for param_group in optimizer.param_groups: + param_group['lr'] = lr + # step the optimizer + optimizer.step() + # --------------- TRAINING SECTION END ------------------- + # everything that follows now is just diagnostics, prints, logging, etc. + + # wait on the CPU for all device work to end so we get accurate per-iteration timings below + if device == "mps": + torch.mps.synchronize() + elif device == "cuda": + torch.cuda.synchronize() + # time and print + t1 = time.time() + # the 0th iteration is often an outlier (much slower) => skip logging it + tokens_per_second = grad_accum_steps * ddp_world_size * B * T / (t1-t0) + print0(f"step {step+1:4d}/{args.num_iterations} | train loss {lossf:.6f} | norm {norm:.4f} | lr {lr:.2e} | ({(t1-t0)*1000:.2f} ms | {tokens_per_second:.0f} tok/s)") + # log to logile + if master_process and logfile is not None: + with open(logfile, "a") as f: + f.write("s:%d trl:%f\n" % (step, lossf)) + + # keep track of smooth timings, last 20 iterations + if step > 0 and step > args.num_iterations - 20: + timings.append(t1-t0) + + # print the average of the last 20 timings, to get something smooth-ish + timings = timings[-20:] + print0(f"final {len(timings)} iters avg: {np.mean(timings)*1000:.3f}ms") + print0(f"peak memory consumption: {torch.cuda.max_memory_allocated() // 1024 // 1024} MiB") + + # ------------------------------------------------------------------------- + # clean up nice + if ddp: + destroy_process_group() diff --git a/baseline/gpt2_small/speedrun30/worker.py b/baseline/gpt2_small/speedrun30/worker.py new file mode 100644 index 00000000..91d31016 --- /dev/null +++ b/baseline/gpt2_small/speedrun30/worker.py @@ -0,0 +1,66 @@ +"""One trainer, one final backup; all work shares the same 30-minute deadline.""" +import argparse +import json +import os +from pathlib import Path +import signal +import subprocess +import sys +import time + + +def write(path, value): + tmp = path.with_suffix(".tmp") + tmp.write_text(json.dumps(value, indent=2)+"\n") + tmp.replace(path) + + +def bounded(command, seconds): + if seconds <= 0: + return {"exit_code":None, "timed_out":True} + child = subprocess.Popen(command, start_new_session=True) + try: + return {"exit_code":child.wait(timeout=seconds), "timed_out":False} + except subprocess.TimeoutExpired: + try: + os.killpg(child.pid, signal.SIGKILL) + except ProcessLookupError: + pass + child.wait(timeout=5) + return {"exit_code":child.returncode, "timed_out":True} + + +def main(): + p = argparse.ArgumentParser() + p.add_argument("root", type=Path) + p.add_argument("deadline", type=float) + a = p.parse_args() + root, deadline = a.root, a.deadline + here = Path(__file__).resolve().parent + run = {"status":"running", "deadline_unix":deadline, + "automatic_restart":False, "limit_seconds":1800} + write(root/"RUN_STATUS.json", run) + try: + # Leave 80 seconds for final cloud backup; no work extends the deadline. + train_deadline = deadline-80 + result = bounded([sys.executable, "-u", str(here/"train.py"), + "--root", str(root), "--deadline", str(train_deadline), + "--cache", "/mnt/disks/rg-data/benchmark-fineweb10B-889765ea"], + train_deadline-time.time()) + run.update(result) + run["status"] = "time_limit" if result["timed_out"] else ( + "finished" if result["exit_code"] == 0 else "failed") + except Exception as exc: + run.update(status="failed", error=str(exc)) + write(root/"RUN_STATUS.json", run) + print(json.dumps(run), flush=True) + backup = bounded([sys.executable, str(here.parent/"scripts/backup.py"), str(root)], + deadline-time.time()-10) + run["backup"] = backup + write(root/"RUN_STATUS.json", run) + print("30-minute job finished. Results remain on the persistent disk.", flush=True) + return 0 if run["status"] == "finished" else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/baseline/gpt2_small/tests/test_speedrun30.py b/baseline/gpt2_small/tests/test_speedrun30.py new file mode 100644 index 00000000..df85fca2 --- /dev/null +++ b/baseline/gpt2_small/tests/test_speedrun30.py @@ -0,0 +1,136 @@ +import copy +import importlib.util +import json +from pathlib import Path +import subprocess +import sys +import time + +import numpy as np +import pytest +import torch + +BASE = Path(__file__).resolve().parents[1]/"speedrun30" +sys.path.insert(0, str(BASE)) +import train as speed +import worker as supervisor + + +def small_model(): + return speed.reference.GPT(speed.reference.GPTConfig( + block_size=8, vocab_size=32, n_layer=2, n_head=2, n_embd=16)) + + +def test_accumulated_update_matches_reference_adamw(): + torch.set_num_threads(1) + model = small_model() + direct = copy.deepcopy(model) + opt = speed.optimizer_for(model, torch.device("cpu")) + # Upstream optimizer grouping/AdamW, using its own function as the oracle. + speed.reference.master_process = True + ref_opt = direct.configure_optimizers(0.1, speed.PEAK_LR, (0.9, 0.95), "cpu", 0) + generator = torch.Generator().manual_seed(83) + for step in range(3): + x = torch.randint(32, (8,8), generator=generator) + y = torch.randint(32, (8,8), generator=generator) + opt.zero_grad(set_to_none=False) + ref_opt.zero_grad(set_to_none=False) + for i in range(4): + _, loss = model(x[2*i:2*i+2], y[2*i:2*i+2]) + (loss/4).backward() + speed.clip_gradients(model) + _, loss = direct(x,y) + loss.backward() + torch.nn.utils.clip_grad_norm_(direct.parameters(), 1.) + for optimizer in (opt, ref_opt): + for group in optimizer.param_groups: + group["lr"] = speed.learning_rate(step) + optimizer.step() + for a,b in zip(model.parameters(), direct.parameters()): + torch.testing.assert_close(a,b,rtol=1e-5,atol=1e-7) + + +def write_shard(path, count): + header = np.zeros(256, dtype=" Date: Sun, 4 Oct 2026 14:01:04 -0700 Subject: [PATCH 25/36] Allow immediate old-run kill and a speedrun without checkpoint or cloud saves Add --kill-current to terminate the entire old service cgroup with SIGKILL rather than requesting a final checkpoint/backup. Add --no-save to skip the new run's model checkpoint and final cloud upload, retaining ordinary loss/throughput logs and the hard 30-minute deadline. Nine focused CPU tests pass, including immediate-kill behavior and proof that no-save supervision never launches backup. --- baseline/gpt2_small/speedrun30/README.md | 10 ++++++ baseline/gpt2_small/speedrun30/cloudshell.py | 26 +++++++++++--- baseline/gpt2_small/speedrun30/train.py | 3 +- baseline/gpt2_small/speedrun30/worker.py | 16 ++++++--- baseline/gpt2_small/tests/test_speedrun30.py | 37 ++++++++++++++++++++ 5 files changed, 82 insertions(+), 10 deletions(-) diff --git a/baseline/gpt2_small/speedrun30/README.md b/baseline/gpt2_small/speedrun30/README.md index e1dfddfd..e1fac7cd 100644 --- a/baseline/gpt2_small/speedrun30/README.md +++ b/baseline/gpt2_small/speedrun30/README.md @@ -14,6 +14,16 @@ From a clean checkout of the pushed commit: python3 baseline/gpt2_small/speedrun30/cloudshell.py start ``` +To kill the old run immediately without a final save/backup and disable new +checkpoint saves and cloud uploads, use: + +```bash +python3 baseline/gpt2_small/speedrun30/cloudshell.py start --kill-current --no-save +``` + +This mode retains ordinary loss/throughput logs and status records. It does not +delete previously saved files or the mounted corpora. + The launcher requests a final save from the current MuonClip service, waits up to 90 seconds, and then stops that service if necessary. Its previously saved checkpoints, diagnostics, cloud archives and FineWeb-Edu corpus are retained. diff --git a/baseline/gpt2_small/speedrun30/cloudshell.py b/baseline/gpt2_small/speedrun30/cloudshell.py index 7cfdc38d..c6a036b7 100644 --- a/baseline/gpt2_small/speedrun30/cloudshell.py +++ b/baseline/gpt2_small/speedrun30/cloudshell.py @@ -30,7 +30,7 @@ def active(unit): return p.stdout.strip() in {"active", "activating", "deactivating", "reloading"} -def stop_current(): +def stop_current(kill_current=False): pointer = BASE/"MUONCLIP_LATEST.json" if not pointer.exists(): return @@ -40,6 +40,13 @@ def stop_current(): raise RuntimeError("Unexpected existing service identity; nothing stopped.") if not active(unit): return + if kill_current: + print("Killing the current MuonClip service immediately; no final save or backup.", flush=True) + run(["systemctl", "kill", "--kill-who=all", "--signal=SIGKILL", unit], timeout=15) + run(["systemctl", "stop", unit], timeout=15) + if active(unit): + raise RuntimeError("Old service is still active; speedrun not launched.") + return print("Requesting final save from current MuonClip run.", flush=True) (root/"muonclip/STOP").touch() end = time.monotonic()+90 @@ -70,7 +77,7 @@ def status_remote(): subprocess.run(["tail", "-n", "30", str(root/"run.log")]) -def start_remote(commit): +def start_remote(commit, kill_current=False, no_save=False): if os.geteuid() != 0 or not os.path.ismount("/mnt/disks/rg-data"): raise RuntimeError("Requires the existing mounted data disk and root.") if not re.fullmatch(r"[0-9a-f]{40}", commit): @@ -94,7 +101,7 @@ def start_remote(commit): "https://github.com/CalculatedContent/rg_optimizers.git"]) run(["git", "-C", str(repo), "fetch", "--depth", "1", "origin", commit], timeout=120) run(["git", "-C", str(repo), "checkout", "--detach", commit]) - stop_current() + stop_current(kill_current=kill_current) # Existing common guard blocks other trainers/replays on this same TPU. path = repo/"baseline/gpt2_small/scripts/run_muonclip.py" spec = importlib.util.spec_from_file_location("existing_launch", path) @@ -111,8 +118,11 @@ def start_remote(commit): "PJRT_DEVICE":"TPU", "TPU_ACCELERATOR_TYPE":"v5litepod-8", "OMP_NUM_THREADS":"4", "OPENBLAS_NUM_THREADS":"4", "MKL_NUM_THREADS":"4", "XLA_USE_SPMD":"1", "TOKENIZERS_PARALLELISM":"false"} + if no_save: + env["RG_SPEEDRUN_NO_SAVE"] = "1" record = {"root":str(root), "unit":unit, "commit":commit, "started_unix":start, "deadline_unix":deadline, "limit_seconds":1800, + "no_checkpoint_or_cloud_backup":no_save, "cloud_uri":"gs://tpu-builders-504820-ww-continuous8/gpt2small/"+root.name} (root/"launch.json").write_text(json.dumps(record, indent=2)) (root/"commit.txt").write_text(commit+"\n") @@ -133,6 +143,8 @@ def start_remote(commit): print("Log:", root/"run.log", flush=True) print("30 minutes maximum including benchmark data, compilation, training and backup.", flush=True) print("No WeightWatcher, per-tensor checks, preflight, or automatic restart.", flush=True) + if no_save: + print("Checkpoint saves and cloud backups disabled; only ordinary run logs/metrics.", flush=True) print("The TPU allocation itself remains available after the job stops.", flush=True) @@ -141,11 +153,17 @@ def main(): p.add_argument("action", choices=("start", "status")) p.add_argument("--on-tpu", action="store_true", help=argparse.SUPPRESS) p.add_argument("--commit", help=argparse.SUPPRESS) + p.add_argument("--kill-current", action="store_true", help="Kill old service without final save or backup") + p.add_argument("--no-save", action="store_true", help="Disable speedrun checkpoints and cloud backup") a = p.parse_args() if a.on_tpu: - start_remote(a.commit) if a.action == "start" else status_remote() + start_remote(a.commit, a.kill_current, a.no_save) if a.action == "start" else status_remote() return 0 remote = ["sudo", "python3", "-c", Path(__file__).read_text(), a.action, "--on-tpu"] + if a.kill_current: + remote.append("--kill-current") + if a.no_save: + remote.append("--no-save") if a.action == "start": repo = Path(__file__).resolve().parents[3] if run(["git", "-C", str(repo), "status", "--porcelain"], capture_output=True).stdout.strip(): diff --git a/baseline/gpt2_small/speedrun30/train.py b/baseline/gpt2_small/speedrun30/train.py index 6e911a9f..0c0ac505 100644 --- a/baseline/gpt2_small/speedrun30/train.py +++ b/baseline/gpt2_small/speedrun30/train.py @@ -223,6 +223,7 @@ def main(): p.add_argument("--deadline", type=float, required=True) p.add_argument("--cache", type=Path, required=True) p.add_argument("--device", choices=("tpu", "cpu"), default="tpu") + p.add_argument("--no-save", action="store_true", help="Do not save model checkpoints") a = p.parse_args() root = a.root source = FineWeb(a.cache, a.deadline-240) @@ -300,7 +301,7 @@ def main(): # Only one final model checkpoint; no optimizer/checkpoint uploading in the hot loop. write_json(root/"status.json", {"status":"final_evaluation", "step":step}) result = evaluate(model, val, rt, root, step, a.deadline-45) - if time.time() < a.deadline-30: + if not a.no_save and time.time() < a.deadline-30: print("Saving final model checkpoint", flush=True) rt.step(wait=True) weights = {name:tensor.detach().cpu() for name,tensor in model.state_dict().items()} diff --git a/baseline/gpt2_small/speedrun30/worker.py b/baseline/gpt2_small/speedrun30/worker.py index 91d31016..7ac7da75 100644 --- a/baseline/gpt2_small/speedrun30/worker.py +++ b/baseline/gpt2_small/speedrun30/worker.py @@ -36,17 +36,20 @@ def main(): p.add_argument("deadline", type=float) a = p.parse_args() root, deadline = a.root, a.deadline + no_save = os.environ.get("RG_SPEEDRUN_NO_SAVE") == "1" here = Path(__file__).resolve().parent run = {"status":"running", "deadline_unix":deadline, - "automatic_restart":False, "limit_seconds":1800} + "automatic_restart":False, "limit_seconds":1800, "no_save":no_save} write(root/"RUN_STATUS.json", run) try: # Leave 80 seconds for final cloud backup; no work extends the deadline. - train_deadline = deadline-80 - result = bounded([sys.executable, "-u", str(here/"train.py"), + train_deadline = deadline-(10 if no_save else 80) + command = [sys.executable, "-u", str(here/"train.py"), "--root", str(root), "--deadline", str(train_deadline), - "--cache", "/mnt/disks/rg-data/benchmark-fineweb10B-889765ea"], - train_deadline-time.time()) + "--cache", "/mnt/disks/rg-data/benchmark-fineweb10B-889765ea"] + if no_save: + command.append("--no-save") + result = bounded(command, train_deadline-time.time()) run.update(result) run["status"] = "time_limit" if result["timed_out"] else ( "finished" if result["exit_code"] == 0 else "failed") @@ -54,6 +57,9 @@ def main(): run.update(status="failed", error=str(exc)) write(root/"RUN_STATUS.json", run) print(json.dumps(run), flush=True) + if no_save: + print("Speedrun stopped; no checkpoint save or cloud backup requested.", flush=True) + return 0 if run["status"] == "finished" else 1 backup = bounded([sys.executable, str(here.parent/"scripts/backup.py"), str(root)], deadline-time.time()-10) run["backup"] = backup diff --git a/baseline/gpt2_small/tests/test_speedrun30.py b/baseline/gpt2_small/tests/test_speedrun30.py index df85fca2..2c097f85 100644 --- a/baseline/gpt2_small/tests/test_speedrun30.py +++ b/baseline/gpt2_small/tests/test_speedrun30.py @@ -134,3 +134,40 @@ def test_long_run_refuses_concurrent_reference_job(monkeypatch, tmp_path): monkeypatch.setattr(module, "active", lambda unit:unit == "speedrun.service") with pytest.raises(RuntimeError, match="reference run is active"): module.assert_idle() + + +def test_immediate_kill_never_requests_stop_checkpoint(monkeypatch, tmp_path): + spec = importlib.util.spec_from_file_location("speedrun_kill_launcher", BASE/"cloudshell.py") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + monkeypatch.setattr(module, "BASE", tmp_path) + root = tmp_path/"muonclip-continuous-20261004-171413" + root.mkdir() + (tmp_path/"MUONCLIP_LATEST.json").write_text(json.dumps( + {"root":str(root), "unit":"rg-gpt2-muonclip-20261004-171413.service"})) + state = {"active":True} + monkeypatch.setattr(module, "active", lambda unit:state["active"]) + commands = [] + def execute(command, **kwargs): + commands.append(command) + if command[1] == "stop": + state["active"] = False + monkeypatch.setattr(module, "run", execute) + module.stop_current(kill_current=True) + assert commands[0][1:4] == ["kill", "--kill-who=all", "--signal=SIGKILL"] + assert commands[1][1] == "stop" + assert not (root/"muonclip/STOP").exists() + + +def test_no_save_supervisor_never_starts_backup(monkeypatch, tmp_path): + commands = [] + monkeypatch.setenv("RG_SPEEDRUN_NO_SAVE", "1") + monkeypatch.setattr(sys, "argv", ["worker.py", str(tmp_path), str(time.time()+1800)]) + def execute(command, seconds): + commands.append(command) + return {"exit_code":0, "timed_out":False} + monkeypatch.setattr(supervisor, "bounded", execute) + assert supervisor.main() == 0 + assert len(commands) == 1 + assert commands[0][-1] == "--no-save" + assert json.loads((tmp_path/"RUN_STATUS.json").read_text())["no_save"] From 1b4548de0b83e5e4c94b2a351c0f8d805d173f1d Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Sun, 4 Oct 2026 15:55:46 -0700 Subject: [PATCH 26/36] Port published 3000-update Muon speedrun to TPU with target validation and checkpoints --- baseline/gpt2_small/README.md | 6 + baseline/gpt2_small/muon_speedrun/README.md | 140 +++++ .../gpt2_small/muon_speedrun/cloudshell.py | 140 +++++ baseline/gpt2_small/muon_speedrun/data.py | 38 ++ baseline/gpt2_small/muon_speedrun/model.py | 163 +++++ baseline/gpt2_small/muon_speedrun/optim.py | 106 ++++ .../muon_speedrun/reference_val.json | 128 ++++ baseline/gpt2_small/muon_speedrun/run.py | 238 ++++++++ baseline/gpt2_small/muon_speedrun/runtime.py | 89 +++ .../gpt2_small/muon_speedrun/vendor/LICENSE | 22 + .../muon_speedrun/vendor/record_source.py | 569 ++++++++++++++++++ baseline/gpt2_small/muon_speedrun/worker.py | 121 ++++ baseline/gpt2_small/scripts/run_muonclip.py | 3 + .../gpt2_small/tests/test_muon_speedrun.py | 168 ++++++ 14 files changed, 1931 insertions(+) create mode 100644 baseline/gpt2_small/muon_speedrun/README.md create mode 100644 baseline/gpt2_small/muon_speedrun/cloudshell.py create mode 100644 baseline/gpt2_small/muon_speedrun/data.py create mode 100644 baseline/gpt2_small/muon_speedrun/model.py create mode 100644 baseline/gpt2_small/muon_speedrun/optim.py create mode 100644 baseline/gpt2_small/muon_speedrun/reference_val.json create mode 100644 baseline/gpt2_small/muon_speedrun/run.py create mode 100644 baseline/gpt2_small/muon_speedrun/runtime.py create mode 100644 baseline/gpt2_small/muon_speedrun/vendor/LICENSE create mode 100644 baseline/gpt2_small/muon_speedrun/vendor/record_source.py create mode 100644 baseline/gpt2_small/muon_speedrun/worker.py create mode 100644 baseline/gpt2_small/tests/test_muon_speedrun.py diff --git a/baseline/gpt2_small/README.md b/baseline/gpt2_small/README.md index 2a566b5e..a8c671dc 100644 --- a/baseline/gpt2_small/README.md +++ b/baseline/gpt2_small/README.md @@ -1,5 +1,11 @@ # GPT-2 Small / FineWeb-Edu validation +For the **published 3,000-update Muon speedrun recipe on eight TPU chips**, see +[muon_speedrun/README.md](muon_speedrun/README.md). It uses a modified transformer, +targets full-validation loss 3.28, saves checkpoints every 125 updates, and records +an explicit target outcome. Its TPU performance and convergence require live +measurement; the default three-hour budget is a cap, not a runtime prediction. + For the **30-minute GPT-2/FineWeb reference run**, see [speedrun30/README.md](speedrun30/README.md). Its launcher stops the current MuonClip service, preserves prior data, and uses the published benchmark's diff --git a/baseline/gpt2_small/muon_speedrun/README.md b/baseline/gpt2_small/muon_speedrun/README.md new file mode 100644 index 00000000..0830de53 --- /dev/null +++ b/baseline/gpt2_small/muon_speedrun/README.md @@ -0,0 +1,140 @@ +# Muon speedrun recipe for one eight-chip TPU + +This ports the published **2024-11-10 UNetDoubleLr** recipe, which reached +**3.2753 validation NLL after 3,000 updates / 1,572,864,000 training tokens**. +The record reports 7.23 minutes of training on eight H100 GPUs. That timing +excludes validation and the first ten updates and is **not a TPU prediction**. +This is an established portable recipe, not the latest CUDA speed record and not +a claim of optimality on v5e. TPU convergence and performance require a live run. + +## Start from Cloud Shell + +Use a clean checkout of the published commit: + +```bash +python3 baseline/gpt2_small/muon_speedrun/cloudshell.py start +``` + +The launcher uses the existing `ww-gpt2-validation-48h-20261004-s1337-node` in +`tpu-builders-504820/us-west4-a`. It creates no TPU, changes no installed packages, +and refuses concurrent training. Repeating start while active prints status. +It starts from scratch because the previous 30-minute run saved no model. + +The worker stops at the first full validation with **NLL <= 3.28**, after the +published 3,000-update schedule, or before its deadline, whichever comes first. +It reports `schedule_complete_target_not_met` explicitly if all updates finish +above target. It never keeps training at a zero learning rate or declares +success merely because the process exited cleanly. + +Default allocation budget: **three hours maximum**, including data preparation, +attention verification, compilation, training and final cloud backup. This is a +safety cap, not an ETA. It is also bounded by the existing TPU allocation expiry. +`--hours 1` requests a one-hour cap; a cap can truncate the schedule without +meeting the target. The TPU allocation itself remains active when the job ends. + +```bash +python3 baseline/gpt2_small/muon_speedrun/cloudshell.py status +``` + +## The complete recipe + +| Setting | Value | +|---|---| +| Architecture | 12 blocks, width 768, 6 heads of width 128, vocabulary 50,304 | +| Model changes | RoPE, RMSNorm/QK normalization, squared ReLU, zero output projections, value residuals, learned input/UNet skip weights, untied output head, logit soft cap 30 | +| Parameters | About 162M total; this is a modified transformer, not standard GPT-2 Small | +| Context / global batch | 1,024 tokens / 524,288 tokens per update | +| Default TPU microbatch | 128 sequences globally, four accumulation passes; 16 sequences per chip | +| Schedule | 3,000 updates, zero warmup, constant LR through 2,100 then 900-update linear decay | +| Hidden matrices | Muon LR 0.04; five quintic Newton–Schulz iterations | +| Muon momentum | Linear ramp 0.85 to 0.95 over the first 500 updates; Nesterov | +| Embeddings / head / scalars | Adam LR 0.6 / 0.008 / 0.04; betas 0.9, 0.95; epsilon 1e-8 | +| Weight decay / global gradient clipping | None, matching the record | +| Precision | BF16 embedding/scalars; FP32 linear weights, BF16 linear compute; BF16 Newton–Schulz | +| Validation | Same pinned GPT-2-tokenized FineWeb file; first 10,485,760 tokens every 125 updates and at stop | +| Target | Full validation NLL <= 3.28, equivalent to perplexity <= exp(3.28) | + +All model and optimizer settings come from the source record. The global batch +stays fixed when microbatch size changes. `--microbatch 64` uses eight accumulation +passes if the 128-sequence configuration exceeds device memory. No automatic +microbatch change or restart can silently alter a run. The selected microbatch +is a starting point, not the result of a TPU tuning sweep. + +## TPU implementation + +One XLA SPMD process partitions batches across eight chips. Global gradients are +replicated before optimizer work. Muon groups matrices by shape, pads groups to a +multiple of eight, and shards the matrix index during Newton–Schulz, then gathers +the resulting updates for replicated weights. It uses the exact record's +coefficients, Nesterov convention and rectangular scaling. It is **Muon, not +MuonClip**. Learning rate and momentum are device tensors to avoid compiling a +new graph solely because their Python values change. + +The worker tests PyTorch/XLA 2.6 TPU flash attention at sequence length 1,024 and +head dimension 128 before using it. The check compares outputs and Q/K/V gradients +against mathematical attention with BF16 relative-L2 tolerance 0.03. It supplies +`sm_scale=1/sqrt(head_dim)` and the SPMD batch partition explicitly. The isolated +check is capped at five minutes. With default `--attention auto`, an unavailable +or failing kernel is logged and the job uses mathematical attention. Request +`--attention flash` to require a passing kernel, or `--attention math` to select +mathematical attention explicitly. No dependency upgrade is performed. + +The 16 necessary training shards plus validation are SHA256-verified and prepared +before training, reusing `/mnt/disks/rg-data/benchmark-fineweb10B-889765ea`. This +avoids synchronous network downloads at training shard boundaries. Existing +FineWeb-Edu is a different corpus and is retained separately. All timings include +end-to-end training overhead; published GPU training-only timing is labelled. + +There are no WeightWatcher fits, gradient scans or per-matrix host reads in the +training loop. Scalar loss checks remain. An isolated flash check does not prove +complete TPU optimizer/model parity. Source CUDA compilation, random seed, +microbatch reduction order, rotary-buffer calculation and shard-boundary order +can differ; these are recorded rather than represented as exact reproduction. + +## Checkpoints and results + +Runs live at `/mnt/disks/rg-data/gpt2small/muon-speedrun--`. +`MUON_SPEEDRUN_LATEST.json` identifies the latest run and service. + +- `checkpoint_latest.pt`: atomic full-state save at initialization, every 125 + updates, and normal stop. Includes model, optimizers, RNG, data cursor, recipe + and next-step schedule. No automatic resume. TPU resume parity is not yet tested. +- `checkpoint_best.pt`: best fully evaluated checkpoint; a hard link protects it + when the latest checkpoint is replaced, without duplicating local storage. +- `checkpoint_target.pt`: saved only after a qualifying full validation. +- `metrics.jsonl`, `latest_validation.json`, `manifest.json`, `status.json`: + measurements, pinned provenance and explicit target outcome. +- Final cloud backup to the matching prefix under + `gs://tpu-builders-504820-ww-continuous8/gpt2small/`, using object permissions + with checksum verification. `CLOUD_BACKUP_VERIFIED.json` confirms completion. + Backups are final, not continuous; intermediate checkpoints remain on the disk. + +A stalled trainer with no recorded progress for 15 minutes is stopped. All +subprocesses and final backup share the hard allocation budget. Abrupt failure +preserves the previous atomic checkpoint; incomplete updates are not labelled +saved. Detailed failure and attention-check records remain with the run. + +## Optional Adam comparison + +`--optimizer adam` runs the same model, seed, token stream, batch and 3,000-update +schedule. Auxiliary Adam groups stay identical; hidden matrices use Adam with +LR 0.0006, betas 0.9/0.95 and zero weight decay. This is an **untuned control**, +not an optimized Adam speedrun or a promise of matching the target. It must run +sequentially on this TPU. Compare at equal tokens and end-to-end wall time; do +not attribute all differences from the earlier standard GPT-2 run to Muon. + +## Pinned source and local verification + +Source: [KellerJordan/modded-nanogpt record](https://github.com/KellerJordan/modded-nanogpt/blob/4ea6b937337a4889b8cfe3f38a93d120048d8f71/records/track_1_short/2024-11-10_UNetDoubleLr/c87bb826-797b-4f37-98c7-d3a5dad2de74.txt). +Record author: Brendan Hogan Rappazzo. The log blob is +`4ef6d69736c2e49ccc1ce4ca98f35128ea7adb8d`. +`vendor/record_source.py` contains the original executable source extracted from +that reproducible log; `vendor/LICENSE` preserves the MIT license. Do not execute +the vendor file on TPU: its launcher targets CUDA/DDP. + +`reference_val.json` preserves the 25 published validation observations. Tests +compare the port's forward/backward math to original source definitions, batched +Muon updates/state restoration to the original Newton–Schulz function, full +mixed-dtype optimizer learning on a tiny CPU model, full-validation target gates, +checkpoint retention and duplicate-launch prevention. These tests cannot certify +TPU performance or convergence before the live run. diff --git a/baseline/gpt2_small/muon_speedrun/cloudshell.py b/baseline/gpt2_small/muon_speedrun/cloudshell.py new file mode 100644 index 00000000..d2760642 --- /dev/null +++ b/baseline/gpt2_small/muon_speedrun/cloudshell.py @@ -0,0 +1,140 @@ +"""Launch a checkpointed Muon recipe on the existing eight-chip TPU.""" +import argparse +import datetime as dt +import fcntl +import importlib.util +import json +import os +from pathlib import Path +import re +import shlex +import subprocess +import sys +import time + +PROJECT = 'tpu-builders-504820' +ZONE = 'us-west4-a' +QUEUE = 'ww-gpt2-validation-48h-20261004-s1337' +NODE = QUEUE+'-node' +BASE = Path('/mnt/disks/rg-data/gpt2small') +LATEST = BASE/'MUON_SPEEDRUN_LATEST.json' + + +def run(command, **kwargs): + return subprocess.run(command, check=True, text=True, **kwargs) + + +def active(unit): + r = subprocess.run(['systemctl','show',unit,'--property=ActiveState','--value'], + capture_output=True,text=True,timeout=10) + return r.stdout.strip() in ('active','activating','deactivating','reloading') + + +def status_remote(): + if not LATEST.exists(): + print('No Muon speedrun launched.') + return + record = json.loads(LATEST.read_text()) + root = Path(record['root']) + print(json.dumps(record,indent=2),flush=True) + subprocess.run(['systemctl','--no-pager','--full','status',record['unit']]) + for name in ('RUN_STATUS.json','status.json','latest_validation.json','checkpoint_latest.json'): + if (root/name).exists(): + print(name+'\n'+(root/name).read_text(),flush=True) + subprocess.run(['tail','-n','15',str(root/'run.log')]) + + +def start_remote(commit, hours=3, optimizer='muon', microbatch=128, attention='auto'): + if os.geteuid() != 0 or not os.path.ismount('/mnt/disks/rg-data'): + raise RuntimeError('Requires the existing mounted disk and root') + if not re.fullmatch('[0-9a-f]{40}',commit): + raise ValueError('Expected pinned commit') + with (BASE/'port-check-launch.lock').open('a') as lock: + fcntl.flock(lock,fcntl.LOCK_EX|fcntl.LOCK_NB) + if LATEST.exists() and active(json.loads(LATEST.read_text())['unit']): + print('A speedrun is already active; no duplicate launched.') + status_remote() + return + allocation = json.loads((BASE/QUEUE/'allocation.json').read_text()) + deadline = min(time.time()+hours*3600,float(allocation['validation_deadline_unix'])-30) + if deadline-time.time() < 2400: + raise RuntimeError('Less than 40 minutes remain on this allocation') + stamp = dt.datetime.now(dt.timezone.utc).strftime('%Y%m%d-%H%M%S') + root = BASE/('muon-speedrun-'+optimizer+'-'+stamp) + root.mkdir() + repo = root/'repo' + repo.mkdir() + run(['git','-C',str(repo),'init','-q']) + run(['git','-C',str(repo),'remote','add','origin','https://github.com/CalculatedContent/rg_optimizers.git']) + run(['git','-C',str(repo),'fetch','--depth','1','origin',commit],timeout=120) + run(['git','-C',str(repo),'checkout','--detach',commit]) + guard = repo/'baseline/gpt2_small/scripts/run_muonclip.py' + spec = importlib.util.spec_from_file_location('training_guard',guard) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + module.assert_idle() + unit = 'rg-muon-speedrun-'+stamp+'.service' + base = repo/'baseline/gpt2_small' + env = {'PYTHONPATH':str(base/'src')+':'+str(base.parent/'nanogpt_one_head/src'), + 'PJRT_DEVICE':'TPU','XLA_USE_SPMD':'1','TPU_ACCELERATOR_TYPE':'v5litepod-8', + 'OMP_NUM_THREADS':'4','OPENBLAS_NUM_THREADS':'4','MKL_NUM_THREADS':'4', + 'TOKENIZERS_PARALLELISM':'false'} + record = {'root':str(root),'unit':unit,'commit':commit,'optimizer':optimizer, + 'started_unix':time.time(),'deadline_unix':deadline,'hours_cap':hours, + 'target_val_nll':3.28,'checkpoint_interval':125,'microbatch':microbatch, + 'cloud_uri':'gs://tpu-builders-504820-ww-continuous8/gpt2small/'+root.name} + (root/'launch.json').write_text(json.dumps(record,indent=2)) + (root/'commit.txt').write_text(commit+'\n') + command = ['systemd-run','--unit='+unit,'--property=Type=exec','--property=Restart=no', + '--property=RuntimeMaxSec='+str(int(deadline-time.time())-5), + '--property=TimeoutStopSec=5','--property=KillMode=control-group', + '--property=StandardOutput=append:'+str(root/'run.log'), + '--property=StandardError=append:'+str(root/'run.log')] + command += ['--setenv='+k+'='+v for k,v in env.items()] + command += ['/mnt/disks/rg-data/continuous8/venv/bin/python','-u', + str(base/'muon_speedrun/worker.py'),str(root),str(deadline), + '--optimizer',optimizer,'--microbatch',str(microbatch),'--attention',attention] + run(command) + temp = LATEST.with_suffix('.tmp') + temp.write_text(json.dumps(record,indent=2)) + temp.replace(LATEST) + print('Started '+optimizer+' recipe: '+unit,flush=True) + print('Log: '+str(root/'run.log'),flush=True) + print('3,000 updates, stops at full-validation NLL <= 3.28; checkpoints every 125.',flush=True) + print('Hard cutoff UTC: '+dt.datetime.fromtimestamp(deadline,dt.timezone.utc).isoformat(),flush=True) + print('No automatic restart or new TPU allocation. Three hours is a cap, not an ETA.',flush=True) + + +def main(): + p = argparse.ArgumentParser(description=__doc__) + p.add_argument('action',choices=('start','status')) + p.add_argument('--hours',type=float,default=3) + p.add_argument('--optimizer',choices=('muon','adam'),default='muon') + p.add_argument('--microbatch',type=int,choices=(64,128),default=128) + p.add_argument('--attention',choices=('auto','flash','math'),default='auto') + p.add_argument('--on-tpu',action='store_true',help=argparse.SUPPRESS) + p.add_argument('--commit',help=argparse.SUPPRESS) + a = p.parse_args() + if not 1 <= a.hours <= 12: + raise ValueError('Hours must be between 1 and 12, bounded by existing allocation') + if a.on_tpu: + if a.action == 'start': + start_remote(a.commit,a.hours,a.optimizer,a.microbatch,a.attention) + else: + status_remote() + return 0 + command = ['sudo','python3','-c',Path(__file__).read_text(),a.action,'--on-tpu', + '--hours',str(a.hours),'--optimizer',a.optimizer,'--microbatch',str(a.microbatch), + '--attention',a.attention] + if a.action == 'start': + repo = Path(__file__).resolve().parents[3] + if run(['git','-C',str(repo),'status','--porcelain'],capture_output=True).stdout.strip(): + raise RuntimeError('Launch from a clean checkout of the pushed commit') + commit = run(['git','-C',str(repo),'rev-parse','HEAD'],capture_output=True).stdout.strip() + command += ['--commit',commit] + return subprocess.run(['gcloud','compute','tpus','tpu-vm','ssh',NODE, + '--project='+PROJECT,'--zone='+ZONE,'--worker=0','--command='+shlex.join(command)]).returncode + + +if __name__ == '__main__': + sys.exit(main()) diff --git a/baseline/gpt2_small/muon_speedrun/data.py b/baseline/gpt2_small/muon_speedrun/data.py new file mode 100644 index 00000000..842943c0 --- /dev/null +++ b/baseline/gpt2_small/muon_speedrun/data.py @@ -0,0 +1,38 @@ +"""Reuse the pinned FineWeb benchmark cache and its SHA256 downloader.""" +from concurrent.futures import ThreadPoolExecutor +import importlib.util +from pathlib import Path +import sys + +REFERENCE = Path(__file__).resolve().parents[1]/'speedrun30' +sys.path.insert(0, str(REFERENCE)) +spec = importlib.util.spec_from_file_location('fineweb_reference', REFERENCE/'train.py') +reference = importlib.util.module_from_spec(spec) +spec.loader.exec_module(reference) +FineWeb = reference.FineWeb +TrainStream = reference.TrainStream +write_json = reference.write_json + + +def required_shards(source, microbatch=128, updates=3000): + remaining = updates * 524288 + needed = ['fineweb_val_000000.bin'] + count = microbatch * 1024 + for name, info in sorted(source.manifest['files'].items()): + if '_train_' not in name: + continue + tokens = (info['size']-1024)//2 + remaining -= ((tokens-1)//count)*count + needed.append(name) + if remaining <= 0: + return needed + raise RuntimeError('Pinned corpus too small for requested schedule') + + +def prepare(cache, deadline, root, microbatch): + source = FineWeb(cache, deadline) + needed = required_shards(source, microbatch) + with ThreadPoolExecutor(max_workers=4) as pool: + for name, _ in zip(needed, pool.map(source.array, needed)): + print('Verified benchmark shard: '+name, flush=True) + write_json(Path(root)/'data_receipts.json', source.receipts) diff --git a/baseline/gpt2_small/muon_speedrun/model.py b/baseline/gpt2_small/muon_speedrun/model.py new file mode 100644 index 00000000..e20e16bd --- /dev/null +++ b/baseline/gpt2_small/muon_speedrun/model.py @@ -0,0 +1,163 @@ +"""TPU adaptation of the pinned November 10, 2024 Muon speedrun model. +See vendor/LICENSE and README.md for provenance and adaptation details. +""" +from dataclasses import dataclass +import torch +from torch import nn +import torch.nn.functional as F + +ATTENTION = None + +def causal_attention(q, k, v): + if ATTENTION is not None: + return ATTENTION(q, k, v) + return F.scaled_dot_product_attention(q, k, v, is_causal=True) + +class Rotary(torch.nn.Module): + + def __init__(self, dim, base=10000): + super().__init__() + # Fixed-context buffers avoid CPU/device transfers in the first forward. + inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim)) + freqs = torch.outer(torch.arange(1024).float(), inv_freq) + self.register_buffer('cos_cached', freqs.cos().bfloat16(), persistent=False) + self.register_buffer('sin_cached', freqs.sin().bfloat16(), persistent=False) + + def forward(self, x): + seq_len = x.shape[1] + return self.cos_cached[None, :seq_len, None, :], self.sin_cached[None, :seq_len, None, :] + +def apply_rotary_emb(x, cos, sin): + assert x.ndim == 4 # multihead attention + d = x.shape[3]//2 + x1 = x[..., :d] + x2 = x[..., d:] + y1 = x1 * cos + x2 * sin + y2 = x1 * (-sin) + x2 * cos + return torch.cat([y1, y2], 3).type_as(x) + +class CastedLinear(nn.Linear): + def forward(self, x): + return F.linear(x, self.weight.to(x.dtype)) + +class CausalSelfAttention(nn.Module): + + def __init__(self, config): + super().__init__() + self.n_head = config.n_head + self.n_embd = config.n_embd + self.head_dim = self.n_embd // self.n_head + assert self.n_embd % self.n_head == 0 + self.c_q = CastedLinear(self.n_embd, self.n_embd, bias=False) + self.c_k = CastedLinear(self.n_embd, self.n_embd, bias=False) + self.c_v = CastedLinear(self.n_embd, self.n_embd, bias=False) + # output projection + self.c_proj = CastedLinear(self.n_embd, self.n_embd, bias=False) + self.c_proj.weight.data.zero_() # zero init suggested by @Grad62304977 + self.rotary = Rotary(self.head_dim) + self.lamb = nn.Parameter(torch.tensor(0.5)) # @Grad62304977 + + def forward(self, x, v1=None): + B, T, C = x.size() # batch size, sequence length, embedding dimensionality (n_embd) + q = self.c_q(x).view(B, T, self.n_head, self.head_dim) + k = self.c_k(x).view(B, T, self.n_head, self.head_dim) + v = self.c_v(x).view(B, T, self.n_head, self.head_dim) + if v1 is None: + v1 = v # This happens if we are in the first block. v needs to be accessed by subsequent blocks + v = (1 - self.lamb) * v + self.lamb * v1.view_as(v) # @Grad62304977 + cos, sin = self.rotary(q) + q, k = F.rms_norm(q, (q.size(-1),)), F.rms_norm(k, (k.size(-1),)) # QK norm suggested by @Grad62304977 + q, k = apply_rotary_emb(q, cos, sin), apply_rotary_emb(k, cos, sin) + y = causal_attention(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)) + y = y.transpose(1, 2).contiguous().view_as(x) # re-assemble all head outputs side by side + y = self.c_proj(y) + return y, v1 + +class MLP(nn.Module): + + def __init__(self, config): + super().__init__() + self.c_fc = CastedLinear(config.n_embd, 4 * config.n_embd, bias=False) + self.c_proj = CastedLinear(4 * config.n_embd, config.n_embd, bias=False) + self.c_proj.weight.data.zero_() # zero init suggested by @Grad62304977 + + def forward(self, x): + x = self.c_fc(x) + x = F.relu(x).square() # https://arxiv.org/abs/2109.08668v2; ~1-2% better than GELU; suggested by @SKYLINEZ007 and @Grad62304977 + x = self.c_proj(x) + return x + +class Block(nn.Module): + + def __init__(self, config): + super().__init__() + self.attn = CausalSelfAttention(config) + self.mlp = MLP(config) + self.lambdas = nn.Parameter(torch.tensor([1., 0.])) + + def forward(self, x, v1, x0): + x = self.lambdas[0] * x + self.lambdas[1] * x0 + x1, v1 = self.attn(F.rms_norm(x, (x.size(-1),)), v1) + x = x + x1 + x = x + self.mlp(F.rms_norm(x, (x.size(-1),))) + return x, v1 + +# ----------------------------------------------------------------------------- +# The main GPT-2 model + +@dataclass +class GPTConfig: + vocab_size : int = 50304 + n_layer : int = 12 + n_head : int = 6 # head dim 128 suggested by @Grad62304977 + n_embd : int = 768 + +class GPT(nn.Module): + + def __init__(self, config): + super().__init__() + self.config = config + + self.transformer = nn.ModuleDict(dict( + wte = nn.Embedding(config.vocab_size, config.n_embd), + h = nn.ModuleList([Block(config) for _ in range(config.n_layer)]), + )) + + # U-net design by @brendanh0gan + self.encoder_layers = config.n_layer // 2 # Half of the layers for encoder + self.decoder_layers = config.n_layer - self.encoder_layers # Remaining for decoder + # Add learnable skip connection weights for decoder layers + self.skip_weights = nn.Parameter(torch.ones(self.decoder_layers)) + + self.lm_head = CastedLinear(config.n_embd, config.vocab_size, bias=False) + self.lm_head.weight.data.zero_() # @Grad62304977 + + def forward(self, idx, target): + + # forward the GPT model itself + x = self.transformer.wte(idx) # token embeddings of shape (b, t, n_embd) + x = F.rms_norm(x, (x.size(-1),)) # @Grad62304977 + x0 = x + v1 = None + + # Store outputs for U-Net skip connections + skip_connections = [] + + # Encoder pass - process only the first half of the blocks + for i in range(self.encoder_layers): + x, v1 = self.transformer.h[i](x, v1, x0) + skip_connections.append(x) # Store the output for skip connections + + # Decoder pass - process the remaining blocks with weighted skip connections + for i in range(self.decoder_layers): + skip_connection = skip_connections.pop() # Get the corresponding encoder output + # Apply learnable weight to skip connection + weighted_skip = self.skip_weights[i] * skip_connection + x, v1 = self.transformer.h[self.encoder_layers + i](x + weighted_skip, v1, x0) + + x = F.rms_norm(x, (x.size(-1),)) + logits = self.lm_head(x) + logits = 30 * torch.tanh(logits / 30) # @Grad62304977 + logits = logits.float() + loss = F.cross_entropy(logits.view(-1, logits.size(-1)), target.view(-1)) + return loss.float() diff --git a/baseline/gpt2_small/muon_speedrun/optim.py b/baseline/gpt2_small/muon_speedrun/optim.py new file mode 100644 index 00000000..94ce3341 --- /dev/null +++ b/baseline/gpt2_small/muon_speedrun/optim.py @@ -0,0 +1,106 @@ +"""Pinned-record Muon math, batched and partitioned over TPU matrix groups.""" +from collections import defaultdict +import math +import torch + + +def schedule(step): + # Record: 3,000 updates, no LR warmup, 900-update linear warmdown. + return min(1., max(0., (3000-step)/900)) + + +def momentum(step): + return 0.85 + 0.10 * min(step/500, 1.) + + +def orthogonalize(g, steps=5): + x = g.bfloat16() + x = x / (x.norm(dim=(-2, -1), keepdim=True) + 1e-7) + transposed = x.shape[-2] > x.shape[-1] + if transposed: + x = x.transpose(-2, -1) + for _ in range(steps): + a = x @ x.transpose(-2, -1) + # Preserve the source's BF16 operation order: (c * A) @ A. + b = -4.7750 * a + (2.0315 * a) @ a + x = 3.4445 * x + b @ x + return x.transpose(-2, -1) if transposed else x + + +class Muon: + def __init__(self, named_parameters, rt): + self.rt = rt + grouped = defaultdict(list) + for name, p in named_parameters: + if p.ndim != 2: + raise ValueError('Muon requires matrix parameters') + grouped[tuple(p.shape)].append((name, p)) + self.groups = [] + for shape, entries in sorted(grouped.items()): + padded = math.ceil(len(entries)/8)*8 if rt.tpu else len(entries) + # CPU allocation then one transfer, independent of the lazy training graph. + buf = torch.zeros((padded, *shape), dtype=torch.float32).to(rt.device) + rt.shard_matrices(buf) + self.groups.append({'entries':entries, 'buffer':buf, 'shape':shape, 'padded':padded}) + + @torch.no_grad() + def step(self, lr, beta): + for group in self.groups: + params = [p for _, p in group['entries']] + if any(p.grad is None for p in params): + raise RuntimeError('Missing Muon gradient') + grads = torch.stack([p.grad for p in params]) + if group['padded'] > len(params): + grads = torch.cat([grads, torch.zeros((group['padded']-len(params), *group['shape']), + dtype=grads.dtype, device=grads.device)]) + self.rt.shard_matrices(grads) + buf = group['buffer'] + buf.mul_(beta).add_(grads) + updates = orthogonalize(grads + beta * buf) + updates *= max(1, group['shape'][0]/group['shape'][1])**0.5 + # NS work is partitioned across chips; updates are then gathered for + # replicated model weights. This replaces upstream rank-local NS + SUM. + self.rt.replicate(updates) + for i, p in enumerate(params): + p.add_(updates[i].to(p.dtype) * (-lr)) + + def state_dict(self): + return [{'names':[name for name, _ in group['entries']], + 'momentum_buffer':group['buffer']} for group in self.groups] + + def load_state_dict(self, state): + if len(state) != len(self.groups): + raise ValueError('Different Muon group count') + for saved, group in zip(state, self.groups): + if saved['names'] != [n for n, _ in group['entries']]: + raise ValueError('Different Muon parameter ordering') + group['buffer'].copy_(saved['momentum_buffer'].to(self.rt.device)) + + +def make_optimizers(model, rt, kind='muon'): + matrices = [(n, p) for n, p in model.transformer.h.named_parameters() if p.ndim == 2] + scalars = [p for p in model.transformer.h.parameters() if p.ndim < 2] + [model.skip_weights] + groups = [ + {'params':[model.transformer.wte.weight], 'peak_lr':0.6}, + {'params':[model.lm_head.weight], 'peak_lr':0.008}, + {'params':scalars, 'peak_lr':0.04}, + ] + muon = Muon(matrices, rt) if kind == 'muon' else None + if kind != 'muon': + # An explicit Adam control on the SAME model/data; not a tuned speed record. + groups.append({'params':[p for _, p in matrices], 'peak_lr':0.0006}) + for group in groups: + group['lr'] = group['peak_lr'] + adam = torch.optim.Adam(groups, betas=(0.9, 0.95), eps=1e-8, + weight_decay=0., foreach=False, fused=False, + capturable=rt.tpu) + return muon, adam + + +def apply_update(muon, adam, rt, step): + factor = schedule(step) + if muon is not None: + muon.step(rt.scalar(0.04 * factor), rt.scalar(momentum(step))) + for group in adam.param_groups: + group['lr'] = rt.scalar(group['peak_lr'] * factor) + adam.step() diff --git a/baseline/gpt2_small/muon_speedrun/reference_val.json b/baseline/gpt2_small/muon_speedrun/reference_val.json new file mode 100644 index 00000000..5d60175c --- /dev/null +++ b/baseline/gpt2_small/muon_speedrun/reference_val.json @@ -0,0 +1,128 @@ +[ + { + "step": 0, + "val_nll": 10.8258, + "training_ms": 411 + }, + { + "step": 125, + "val_nll": 4.5262, + "training_ms": 16570 + }, + { + "step": 250, + "val_nll": 4.0672, + "training_ms": 34813 + }, + { + "step": 375, + "val_nll": 3.903, + "training_ms": 52755 + }, + { + "step": 500, + "val_nll": 3.8112, + "training_ms": 70967 + }, + { + "step": 625, + "val_nll": 3.7344, + "training_ms": 89175 + }, + { + "step": 750, + "val_nll": 3.6796, + "training_ms": 107127 + }, + { + "step": 875, + "val_nll": 3.632, + "training_ms": 125315 + }, + { + "step": 1000, + "val_nll": 3.5918, + "training_ms": 143506 + }, + { + "step": 1125, + "val_nll": 3.565, + "training_ms": 161454 + }, + { + "step": 1250, + "val_nll": 3.5382, + "training_ms": 179673 + }, + { + "step": 1375, + "val_nll": 3.5165, + "training_ms": 197903 + }, + { + "step": 1500, + "val_nll": 3.4969, + "training_ms": 215858 + }, + { + "step": 1625, + "val_nll": 3.4792, + "training_ms": 234086 + }, + { + "step": 1750, + "val_nll": 3.4627, + "training_ms": 252310 + }, + { + "step": 1875, + "val_nll": 3.4492, + "training_ms": 270304 + }, + { + "step": 2000, + "val_nll": 3.4357, + "training_ms": 288507 + }, + { + "step": 2125, + "val_nll": 3.4242, + "training_ms": 306734 + }, + { + "step": 2250, + "val_nll": 3.3985, + "training_ms": 324714 + }, + { + "step": 2375, + "val_nll": 3.3728, + "training_ms": 342929 + }, + { + "step": 2500, + "val_nll": 3.3475, + "training_ms": 361151 + }, + { + "step": 2625, + "val_nll": 3.3246, + "training_ms": 379139 + }, + { + "step": 2750, + "val_nll": 3.3033, + "training_ms": 397355 + }, + { + "step": 2875, + "val_nll": 3.2858, + "training_ms": 415584 + }, + { + "step": 3000, + "val_nll": 3.2753, + "training_ms": 433570 + } +] + diff --git a/baseline/gpt2_small/muon_speedrun/run.py b/baseline/gpt2_small/muon_speedrun/run.py new file mode 100644 index 00000000..791bb918 --- /dev/null +++ b/baseline/gpt2_small/muon_speedrun/run.py @@ -0,0 +1,238 @@ +"""Pinned Muon speedrun with full validation, atomic checkpoints and explicit limits.""" +import argparse +from dataclasses import asdict +import hashlib +import json +import math +import os +from pathlib import Path +import random +import signal +import time +import numpy as np +import torch + +import model as architecture +from data import FineWeb, TrainStream, write_json, prepare +from optim import make_optimizers, apply_update, schedule, momentum +from runtime import Runtime, attention_check + +TOTAL_STEPS = 3000 +BATCH_TOKENS = 524288 +VAL_TOKENS = 10485760 +TARGET = 3.28 +HERE = Path(__file__).resolve().parent + + +def record(root, row): + with (root/'metrics.jsonl').open('a') as f: + f.write(json.dumps(row, allow_nan=False)+'\n') + print(json.dumps(row, allow_nan=False), flush=True) + + +def to_cpu(value): + if isinstance(value, torch.Tensor): + return value.detach().cpu() + if isinstance(value, dict): + return {k:to_cpu(v) for k,v in value.items()} + if isinstance(value, list): + return [to_cpu(v) for v in value] + if isinstance(value, tuple): + return tuple(to_cpu(v) for v in value) + return value + + +def save_checkpoint(root, model, muon, adam, stream, step, manifest, rt, validation, best): + rt.step(wait=True) + payload = {'schema':1, 'step':step, 'tokens_seen':step*BATCH_TOKENS, + 'model':to_cpu(model.state_dict()), 'config':asdict(model.config), + 'adam':to_cpu(adam.state_dict()), + 'muon':to_cpu(muon.state_dict()) if muon is not None else None, + 'data_cursor':{'shard':stream.shard, 'position':stream.position}, + 'rng':{'torch':torch.get_rng_state(), 'numpy':np.random.get_state(), + 'python':random.getstate()}, + 'manifest':manifest, 'validation':validation, + 'next_lr_factor':schedule(step), 'next_momentum':momentum(step), + 'resume_validation':'Full state saved; automatic resume disabled. TPU resume parity not yet validated.'} + checkpoint = root/'checkpoint_latest.pt' + temporary = checkpoint.with_suffix('.tmp') + with temporary.open('wb') as f: + torch.save(payload, f) + f.flush() + os.fsync(f.fileno()) + temporary.replace(checkpoint) + passed = target_met(validation) + aliases = [] + if validation and validation.get('full_benchmark_evaluation') and validation['val_nll'] <= best: + aliases.append('checkpoint_best.pt') + if passed: + aliases.append('checkpoint_target.pt') + for name in aliases: + temp = root/(name+'.tmp') + temp.unlink(missing_ok=True) + os.link(checkpoint, temp) + temp.replace(root/name) + write_json(root/'checkpoint_latest.json', {'file':checkpoint.name, 'step':step, + 'target_met':passed, 'validation':validation}) + print('Checkpoint saved at update '+str(step), flush=True) + + +def target_met(validation): + return bool(validation and validation.get('full_benchmark_evaluation') + and validation.get('evaluation_tokens') == VAL_TOKENS + and validation.get('val_nll') is not None + and math.isfinite(validation['val_nll']) and validation['val_nll'] <= TARGET) + + +@torch.no_grad() +def evaluate(model, tokens, rt, root, step, deadline, microbatch, started): + model.eval() + size = microbatch*1024 + total = torch.zeros((), device=rt.device, dtype=torch.float32) + evaluated = 0 + for offset in range(0, VAL_TOKENS, size): + if time.time() >= deadline: + break + buf = torch.from_numpy(np.array(tokens[offset:offset+size+1], dtype=np.int64)) + total += model(rt.put(buf[:-1].reshape(microbatch, 1024)), + rt.put(buf[1:].reshape(microbatch, 1024))).detach().float() + rt.step() + evaluated += size + rt.step(wait=True) + loss = float(total.cpu())/(evaluated/size) if evaluated else None + if loss is not None and not math.isfinite(loss): + raise RuntimeError('Nonfinite validation NLL') + row = {'kind':'validation', 'step':step, 'tokens_seen':step*BATCH_TOKENS, + 'evaluation_tokens':evaluated, 'full_benchmark_evaluation':evaluated == VAL_TOKENS, + 'val_nll':loss, 'val_perplexity':math.exp(loss) if loss is not None else None, + 'elapsed_seconds':time.time()-started, 'recorded_unix':time.time()} + reference = json.loads((HERE/'reference_val.json').read_text()) + row['published_at_same_step'] = next((r for r in reference if r['step'] == step), None) + row['target_met'] = target_met(row) + record(root, row) + write_json(root/'latest_validation.json', row) + model.train() + return row + + +def train(a): + root = a.root + seed = 1337 + torch.set_num_threads(4) + torch.manual_seed(seed) + np.random.seed(seed) + random.seed(seed) + started = time.time() + rt = Runtime(a.device, root/'xla-cache') + architecture.ATTENTION = rt.attention(a.attention) + model = architecture.GPT(architecture.GPTConfig()).bfloat16() + for module in model.modules(): + if isinstance(module, architecture.CastedLinear): + module.float() + model = model.to(rt.device) + for value in (*model.parameters(), *model.buffers()): + rt.replicate(value) + muon, adam = make_optimizers(model, rt, a.optimizer) + source = FineWeb(a.cache, a.deadline-180) + stream = TrainStream(source, a.microbatch, 1024) + val = source.array('fineweb_val_000000.bin') + manifest = {'recipe':'2024-11-10_UNetDoubleLr', 'optimizer':a.optimizer, 'seed':seed, + 'config':asdict(model.config), 'parameters':sum(p.numel() for p in model.parameters()), + 'steps':TOTAL_STEPS, 'batch_tokens':BATCH_TOKENS, + 'global_microbatch_sequences':a.microbatch, 'accumulation':512//a.microbatch, + 'target_val_nll':TARGET, 'validation_tokens':VAL_TOKENS, + 'muon_lr':0.04 if muon else None, 'adam_embedding_lr':0.6, + 'adam_head_lr':0.008, 'adam_scalar_lr':0.04, + 'adam_control_matrix_lr':0.0006 if not muon else None, + 'warmup_updates':0, 'warmdown_updates':900, 'weight_decay':0, + 'gradient_clipping':False, 'attention':a.attention, + 'data_repo':source.manifest['repo'], 'data_revision':source.manifest['revision'], + 'record_source_sha256':hashlib.sha256((HERE/'vendor/record_source.py').read_bytes()).hexdigest(), + 'torch_version':torch.__version__, 'automatic_restart':False, + 'differences':['TPU SPMD instead of CUDA DDP', 'batched matrix-partitioned Muon', + 'microbatch accumulation and shard-boundary ordering', + 'fixed seed 1337; original record did not pin a seed', + 'CPU-precomputed BF16 rotary buffers; hardware rounding differs'], + 'optimality':'Published GPU recipe; TPU convergence/performance unvalidated'} + write_json(root/'manifest.json', manifest) + print(json.dumps(manifest), flush=True) + step, best, validation = 0, float('inf'), None + stopped = [False] + signal.signal(signal.SIGTERM, lambda *_:stopped.__setitem__(0, True)) + write_json(root/'status.json', {'status':'training', 'step':0}) + rt.step(wait=True) + save_checkpoint(root, model, muon, adam, stream, step, manifest, rt, None, best) + timings = [] + while step < TOTAL_STEPS and time.time() < a.deadline-180 and not stopped[0] and not (root/'STOP').exists(): + began = time.monotonic() + model.zero_grad(set_to_none=False) + loss_sum = torch.zeros((), device=rt.device, dtype=torch.float32) + for _ in range(512//a.microbatch): + x, y = stream.next_batch() + loss = model(rt.put(x), rt.put(y)) + (loss/(512//a.microbatch)).backward() + loss_sum += loss.detach()/(512//a.microbatch) + rt.step() + for p in model.parameters(): + if p.grad is not None: + rt.replicate(p.grad) + apply_update(muon, adam, rt, step) + rt.step(wait=True) + step += 1 + loss_value = float(loss_sum.cpu()) + if not math.isfinite(loss_value): + raise RuntimeError('Nonfinite training NLL at update '+str(step)) + seconds = time.monotonic()-began + timings.append(seconds) + row = {'kind':'train', 'step':step, 'tokens_seen':step*BATCH_TOKENS, + 'train_nll':loss_value, 'seconds':seconds, 'tokens_per_second':BATCH_TOKENS/seconds, + 'elapsed_seconds':time.time()-started, 'lr_factor':schedule(step-1)} + if step <= 5 or step % 10 == 0: + if step >= 20: + row['training_seconds_remaining_estimate'] = float(np.median(timings[-50:]))*(TOTAL_STEPS-step) + record(root, row) + write_json(root/'status.json', {'status':'training', **row}) + if step % 125 == 0: + validation = evaluate(model, val, rt, root, step, a.deadline-90, a.microbatch, started) + save_checkpoint(root, model, muon, adam, stream, step, manifest, rt, validation, best) + if validation['full_benchmark_evaluation']: + best = min(best, validation['val_nll']) + if target_met(validation): + break + if validation is None or validation['step'] != step: + validation = evaluate(model, val, rt, root, step, a.deadline-60, a.microbatch, started) + save_checkpoint(root, model, muon, adam, stream, step, manifest, rt, validation, best) + outcome = 'target_reached' if target_met(validation) else ( + 'schedule_complete_target_not_met' if step == TOTAL_STEPS else 'stopped_before_schedule_complete') + result = {'status':outcome, 'step':step, 'validation':validation, + 'target_met':target_met(validation), 'elapsed_seconds':time.time()-started, + 'checkpoint':str(root/'checkpoint_latest.pt')} + write_json(root/'status.json', result) + print(json.dumps(result), flush=True) + + +def main(): + p = argparse.ArgumentParser() + p.add_argument('action', choices=('prepare', 'attention-check', 'train')) + p.add_argument('--root', type=Path, required=True) + p.add_argument('--cache', type=Path, default=Path('/mnt/disks/rg-data/benchmark-fineweb10B-889765ea')) + p.add_argument('--deadline', type=float, required=True) + p.add_argument('--device', choices=('tpu', 'cpu'), default='tpu') + p.add_argument('--attention', choices=('flash', 'math'), default='flash') + p.add_argument('--microbatch', type=int, choices=(64,128), default=128) + p.add_argument('--optimizer', choices=('muon', 'adam'), default='muon') + a = p.parse_args() + if a.action == 'prepare': + prepare(a.cache, a.deadline, a.root, a.microbatch) + elif a.action == 'attention-check': + attention_check(a.root) + else: + try: + train(a) + except Exception as exc: + write_json(a.root/'FAILURE.json', {'error':str(exc), 'attribution':'unconfirmed'}) + raise + + +if __name__ == '__main__': + main() diff --git a/baseline/gpt2_small/muon_speedrun/runtime.py b/baseline/gpt2_small/muon_speedrun/runtime.py new file mode 100644 index 00000000..d366b280 --- /dev/null +++ b/baseline/gpt2_small/muon_speedrun/runtime.py @@ -0,0 +1,89 @@ +"""One-process SPMD runtime; TPU flash attention has explicit scale and sharding.""" +import json +import math +from pathlib import Path +import numpy as np +import torch +import torch.nn.functional as F + + +class Runtime: + def __init__(self, device='tpu', cache=None): + self.tpu = device == 'tpu' + if self.tpu: + import torch_xla.core.xla_model as xm + import torch_xla.runtime as xr + import torch_xla.distributed.spmd as xs + xr.use_spmd() + if cache is not None: + xr.initialize_cache(str(cache), readonly=False) + if xr.global_runtime_device_count() != 8 or xr.addressable_runtime_device_count() != 8: + raise RuntimeError('Requires one host with eight TPU chips') + self.xm, self.xs = xm, xs + self.mesh = xs.Mesh(np.arange(8), (8,), ('data',)) + self.device = xm.xla_device() + else: + self.device = torch.device('cpu') + + def scalar(self, value): + return torch.tensor(value, dtype=torch.float32).to(self.device) if self.tpu else value + + def put(self, value): + value = value.to(self.device) + if self.tpu: + self.xs.mark_sharding(value, self.mesh, ('data',)+(None,)*(value.ndim-1)) + return value + + def replicate(self, value): + if self.tpu: + self.xs.mark_sharding(value, self.mesh, (None,)*value.ndim) + + def shard_matrices(self, value): + if self.tpu: + self.xs.mark_sharding(value, self.mesh, ('data', None, None)) + + def step(self, wait=False): + if self.tpu: + self.xm.mark_step() + if wait: + self.xm.wait_device_ops() + + def attention(self, backend): + if backend == 'flash': + if not self.tpu: + raise ValueError('TPU flash attention requires TPU') + from torch_xla.experimental.custom_kernel import flash_attention + def flash(q, k, v): + return flash_attention(q, k, v, causal=True, sm_scale=q.shape[-1]**-0.5, + partition_spec=('data', None, None, None), mesh=self.mesh) + return flash + return lambda q, k, v: F.scaled_dot_product_attention(q, k, v, is_causal=True) + + +def attention_check(root): + """Check forward/backward at the actual head/sequence shape in an isolated child.""" + from torch_xla.experimental.custom_kernel import jax_import_guard + jax_import_guard() + import jax + rt = Runtime(cache=Path(root)/'xla-cache') + gen = torch.Generator().manual_seed(43) + cpu = [torch.randn((8, 6, 1024, 128), generator=gen).bfloat16() for _ in range(3)] + upstream = torch.randn(cpu[0].shape, generator=gen).bfloat16() + results = [] + for kind in ('math', 'flash'): + values = [rt.put(t).detach().requires_grad_() for t in cpu] + output = rt.attention(kind)(*values) + (output.float()*rt.put(upstream).float()).sum().backward() + rt.step(wait=True) + results.append([output.detach().cpu()] + [t.grad.detach().cpu() for t in values]) + errors = [] + for a, b in zip(*results): + if not torch.isfinite(a).all() or not torch.isfinite(b).all(): + raise RuntimeError('Nonfinite attention check') + relative = float((a.float()-b.float()).norm()/a.float().norm().clamp_min(1e-12)) + errors.append(relative) + if relative > .03: + raise RuntimeError('TPU flash/math relative L2 disagreement: '+str(relative)) + result = {'status':'passed', 'relative_l2_output_dq_dk_dv':errors, 'jax':jax.__version__} + (Path(root)/'ATTENTION_CHECK.json').write_text(json.dumps(result, indent=2)+'\n') + print(json.dumps(result), flush=True) diff --git a/baseline/gpt2_small/muon_speedrun/vendor/LICENSE b/baseline/gpt2_small/muon_speedrun/vendor/LICENSE new file mode 100644 index 00000000..e8e852a0 --- /dev/null +++ b/baseline/gpt2_small/muon_speedrun/vendor/LICENSE @@ -0,0 +1,22 @@ +MIT License + +Copyright (c) 2024 Keller Jordan + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. + diff --git a/baseline/gpt2_small/muon_speedrun/vendor/record_source.py b/baseline/gpt2_small/muon_speedrun/vendor/record_source.py new file mode 100644 index 00000000..56249b8c --- /dev/null +++ b/baseline/gpt2_small/muon_speedrun/vendor/record_source.py @@ -0,0 +1,569 @@ +import os +import sys +with open(sys.argv[0]) as f: + code = f.read() # read the code of this file ASAP, for logging +import uuid +import glob +import time +from dataclasses import dataclass + +import numpy as np +import torch +from torch import nn +import torch.nn.functional as F +import torch.distributed as dist +import torch._inductor.config as config +from torch.nn.parallel import DistributedDataParallel as DDP + +# ----------------------------------------------------------------------------- +# Muon optimizer + +def zeropower_via_svd(G, steps=None): + U, S, V = G.svd() + return U @ V.T + +@torch.compile +def zeropower_via_newtonschulz5(G, steps=10, eps=1e-7): + """ + Newton-Schulz iteration to compute the zeroth power / orthogonalization of G. We opt to use a + quintic iteration whose coefficients are selected to maximize the slope at zero. For the purpose + of minimizing steps, it turns out to be empirically effective to keep increasing the slope at + zero even beyond the point where the iteration no longer converges all the way to one everywhere + on the interval. This iteration therefore does not produce UV^T but rather something like US'V^T + where S' is diagonal with S_{ii}' \sim Uniform(0.5, 1.5), which turns out not to hurt model + performance at all relative to UV^T, where USV^T = G is the SVD. + """ + assert len(G.shape) == 2 + a, b, c = (3.4445, -4.7750, 2.0315) + X = G.bfloat16() + X /= (X.norm() + eps) # ensure top singular value <= 1 + if G.size(0) > G.size(1): + X = X.T + for _ in range(steps): + A = X @ X.T + B = b * A + c * A @ A # adapted from suggestion by @jxbz, @leloykun, and @YouJiacheng + X = a * X + B @ X + if G.size(0) > G.size(1): + X = X.T + return X + +zeropower_backends = dict(svd=zeropower_via_svd, newtonschulz5=zeropower_via_newtonschulz5) + +class Muon(torch.optim.Optimizer): + """ + Muon - MomentUm Orthogonalized by Newton-schulz + + Muon internally runs standard SGD-momentum, and then performs an orthogonalization post- + processing step, in which each 2D parameter's update is replaced with the nearest orthogonal + matrix. To efficiently orthogonalize each update, we use a Newton-Schulz iteration, which has + the advantage that it can be stably run in bfloat16 on the GPU. + + Some warnings: + - This optimizer assumes that all parameters passed in are 2D. + - It should not be used for the embedding layer, the final fully connected layer, or any {0,1}-D + parameters; those should all be optimized by a standard method (e.g., AdamW). + - To use it with 4D convolutional filters, it works well to just flatten their last 3 dimensions. + - We believe it is unlikely to work well for training with small batch size. + - We believe it may not work well for finetuning pretrained models, but we haven't tested this. + - We have not yet tried this optimizer for training scenarios larger than NanoGPT (124M). + + Arguments: + lr: The learning rate used by the internal SGD. + momentum: The momentum used by the internal SGD. + nesterov: Whether to use Nesterov-style momentum in the internal SGD. (recommended) + backend: The chosen backend for the orthogonalization step. (recommended: 'newtonschulz5') + backend_steps: The number of iteration steps to use in the backend, if it is iterative. + """ + def __init__(self, params, lr=0.02, momentum=0.95, nesterov=True, + backend='newtonschulz5', backend_steps=5): + defaults = dict(lr=lr, momentum=momentum, nesterov=nesterov, backend=backend, backend_steps=backend_steps) + super().__init__(params, defaults) + + def step(self): + + for group in self.param_groups: + + lr = group['lr'] + momentum = group['momentum'] + zeropower_backend = zeropower_backends[group['backend']] + + # generate weight updates in distributed fashion + total_params = sum(p.numel() for p in group['params']) + updates_flat = torch.zeros(total_params, device='cuda', dtype=torch.bfloat16) + curr_idx = 0 + for i, p in enumerate(group['params']): + # luckily this will perfectly distribute a transformer with multiple of 4 layers to 8 GPUs + if i % int(os.environ['WORLD_SIZE']) == int(os.environ['RANK']): + g = p.grad + assert g is not None + state = self.state[p] + if 'momentum_buffer' not in state: + state['momentum_buffer'] = torch.zeros_like(g) + buf = state['momentum_buffer'] + buf.mul_(momentum).add_(g) + if group['nesterov']: + g = g.add(buf, alpha=momentum) + g = zeropower_backend(g, steps=group['backend_steps']) + g *= max(1, g.size(0)/g.size(1))**0.5 + updates_flat[curr_idx:curr_idx+p.numel()] = g.flatten() + curr_idx += p.numel() + + # sync updates across devices. we are not memory-constrained so can do this simple deserialization + dist.all_reduce(updates_flat, op=dist.ReduceOp.SUM) + + # deserialize and apply updates + curr_idx = 0 + for p in group['params']: + g = updates_flat[curr_idx:curr_idx+p.numel()].view_as(p.data).type_as(p.data) + p.data.add_(g, alpha=-lr) + curr_idx += p.numel() + +# ----------------------------------------------------------------------------- +# PyTorch nn.Module definitions for the GPT-2 model + +class Rotary(torch.nn.Module): + + def __init__(self, dim, base=10000): + super().__init__() + self.inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim)) + self.seq_len_cached = None + self.cos_cached = None + self.sin_cached = None + + def forward(self, x): + seq_len = x.shape[1] + if seq_len != self.seq_len_cached: + self.seq_len_cached = seq_len + t = torch.arange(seq_len, device=x.device).type_as(self.inv_freq) + freqs = torch.outer(t, self.inv_freq).to(x.device) + self.cos_cached = freqs.cos().bfloat16() + self.sin_cached = freqs.sin().bfloat16() + return self.cos_cached[None, :, None, :], self.sin_cached[None, :, None, :] + +def apply_rotary_emb(x, cos, sin): + assert x.ndim == 4 # multihead attention + d = x.shape[3]//2 + x1 = x[..., :d] + x2 = x[..., d:] + y1 = x1 * cos + x2 * sin + y2 = x1 * (-sin) + x2 * cos + return torch.cat([y1, y2], 3).type_as(x) + +class CastedLinear(nn.Linear): + def forward(self, x): + return F.linear(x, self.weight.to(x.dtype)) + +class CausalSelfAttention(nn.Module): + + def __init__(self, config): + super().__init__() + self.n_head = config.n_head + self.n_embd = config.n_embd + self.head_dim = self.n_embd // self.n_head + assert self.n_embd % self.n_head == 0 + self.c_q = CastedLinear(self.n_embd, self.n_embd, bias=False) + self.c_k = CastedLinear(self.n_embd, self.n_embd, bias=False) + self.c_v = CastedLinear(self.n_embd, self.n_embd, bias=False) + # output projection + self.c_proj = CastedLinear(self.n_embd, self.n_embd, bias=False) + self.c_proj.weight.data.zero_() # zero init suggested by @Grad62304977 + self.rotary = Rotary(self.head_dim) + self.lamb = nn.Parameter(torch.tensor(0.5)) # @Grad62304977 + + def forward(self, x, v1=None): + B, T, C = x.size() # batch size, sequence length, embedding dimensionality (n_embd) + q = self.c_q(x).view(B, T, self.n_head, self.head_dim) + k = self.c_k(x).view(B, T, self.n_head, self.head_dim) + v = self.c_v(x).view(B, T, self.n_head, self.head_dim) + if v1 is None: + v1 = v # This happens if we are in the first block. v needs to be accessed by subsequent blocks + v = (1 - self.lamb) * v + self.lamb * v1.view_as(v) # @Grad62304977 + cos, sin = self.rotary(q) + q, k = F.rms_norm(q, (q.size(-1),)), F.rms_norm(k, (k.size(-1),)) # QK norm suggested by @Grad62304977 + q, k = apply_rotary_emb(q, cos, sin), apply_rotary_emb(k, cos, sin) + y = F.scaled_dot_product_attention(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), is_causal=True) + y = y.transpose(1, 2).contiguous().view_as(x) # re-assemble all head outputs side by side + y = self.c_proj(y) + return y, v1 + +class MLP(nn.Module): + + def __init__(self, config): + super().__init__() + self.c_fc = CastedLinear(config.n_embd, 4 * config.n_embd, bias=False) + self.c_proj = CastedLinear(4 * config.n_embd, config.n_embd, bias=False) + self.c_proj.weight.data.zero_() # zero init suggested by @Grad62304977 + + def forward(self, x): + x = self.c_fc(x) + x = F.relu(x).square() # https://arxiv.org/abs/2109.08668v2; ~1-2% better than GELU; suggested by @SKYLINEZ007 and @Grad62304977 + x = self.c_proj(x) + return x + +class Block(nn.Module): + + def __init__(self, config): + super().__init__() + self.attn = CausalSelfAttention(config) + self.mlp = MLP(config) + self.lambdas = nn.Parameter(torch.tensor([1., 0.])) + + def forward(self, x, v1, x0): + x = self.lambdas[0] * x + self.lambdas[1] * x0 + x1, v1 = self.attn(F.rms_norm(x, (x.size(-1),)), v1) + x = x + x1 + x = x + self.mlp(F.rms_norm(x, (x.size(-1),))) + return x, v1 + +# ----------------------------------------------------------------------------- +# The main GPT-2 model + +@dataclass +class GPTConfig: + vocab_size : int = 50304 + n_layer : int = 12 + n_head : int = 6 # head dim 128 suggested by @Grad62304977 + n_embd : int = 768 + +class GPT(nn.Module): + + def __init__(self, config): + super().__init__() + self.config = config + + self.transformer = nn.ModuleDict(dict( + wte = nn.Embedding(config.vocab_size, config.n_embd), + h = nn.ModuleList([Block(config) for _ in range(config.n_layer)]), + )) + + # U-net design by @brendanh0gan + self.encoder_layers = config.n_layer // 2 # Half of the layers for encoder + self.decoder_layers = config.n_layer - self.encoder_layers # Remaining for decoder + # Add learnable skip connection weights for decoder layers + self.skip_weights = nn.Parameter(torch.ones(self.decoder_layers)) + + self.lm_head = CastedLinear(config.n_embd, config.vocab_size, bias=False) + self.lm_head.weight.data.zero_() # @Grad62304977 + + def forward(self, idx, target): + + # forward the GPT model itself + x = self.transformer.wte(idx) # token embeddings of shape (b, t, n_embd) + x = F.rms_norm(x, (x.size(-1),)) # @Grad62304977 + x0 = x + v1 = None + + # Store outputs for U-Net skip connections + skip_connections = [] + + # Encoder pass - process only the first half of the blocks + for i in range(self.encoder_layers): + x, v1 = self.transformer.h[i](x, v1, x0) + skip_connections.append(x) # Store the output for skip connections + + # Decoder pass - process the remaining blocks with weighted skip connections + for i in range(self.decoder_layers): + skip_connection = skip_connections.pop() # Get the corresponding encoder output + # Apply learnable weight to skip connection + weighted_skip = self.skip_weights[i] * skip_connection + x, v1 = self.transformer.h[self.encoder_layers + i](x + weighted_skip, v1, x0) + + x = F.rms_norm(x, (x.size(-1),)) + logits = self.lm_head(x) + logits = 30 * torch.tanh(logits / 30) # @Grad62304977 + logits = logits.float() + loss = F.cross_entropy(logits.view(-1, logits.size(-1)), target.view(-1)) + return loss.float() + +# ----------------------------------------------------------------------------- +# Our own simple Distributed Data Loader + +def _peek_data_shard(filename): + # only reads the header, returns header data + with open(filename, "rb") as f: + # first read the header, which is 256 int32 integers (4 bytes each) + header = np.frombuffer(f.read(256*4), dtype=np.int32) + if header[0] != 20240520: + print("ERROR: magic number mismatch in the data .bin file!") + print("---> HINT: Are you passing in a correct file with --input_bin?") + print("---> HINT: Dataset encoding changed recently, re-run data prepro or refer again to README") + print("---> HINT: For example re-run: `python dev/data/tinyshakespeare.py`, then re-try") + exit(1) + assert header[1] == 1, "unsupported version" + ntok = header[2] # number of tokens (claimed) + return ntok # for now just return the number of tokens + +def _load_data_shard(filename): + with open(filename, "rb") as f: + # first read the header, which is 256 int32 integers (4 bytes each) + header = np.frombuffer(f.read(256*4), dtype=np.int32) + assert header[0] == 20240520, "magic number mismatch in the data .bin file" + assert header[1] == 1, "unsupported version" + ntok = header[2] # number of tokens (claimed) + # the rest of it are tokens, stored as uint16 + tokens = np.frombuffer(f.read(), dtype=np.uint16) + assert len(tokens) == ntok, "number of tokens read does not match header?" + return tokens + +class DistributedDataLoader: + def __init__(self, filename_pattern, B, T, process_rank, num_processes): + self.process_rank = process_rank + self.num_processes = num_processes + self.B = B + self.T = T + + # glob files that match the pattern + self.files = sorted(glob.glob(filename_pattern)) + assert len(self.files) > 0, f"did not find any files that match the pattern {filename_pattern}" + + # load and validate all data shards, count number of tokens in total + ntok_total = 0 + for fname in self.files: + shard_ntok = _peek_data_shard(fname) + assert shard_ntok >= num_processes * B * T + 1 + ntok_total += int(shard_ntok) + self.ntok_total = ntok_total + + # kick things off + self.reset() + + def reset(self): + self.current_shard = 0 + self.current_position = self.process_rank * self.B * self.T + self.tokens = _load_data_shard(self.files[self.current_shard]) + + def advance(self): # advance to next data shard + self.current_shard = (self.current_shard + 1) % len(self.files) + self.current_position = self.process_rank * self.B * self.T + self.tokens = _load_data_shard(self.files[self.current_shard]) + + def next_batch(self): + B = self.B + T = self.T + buf = self.tokens[self.current_position : self.current_position+B*T+1] + buf = torch.tensor(buf.astype(np.int32), dtype=torch.long) + x = (buf[:-1]).view(B, T) # inputs + y = (buf[1:]).view(B, T) # targets + # advance current position and load next shard if necessary + self.current_position += B * T * self.num_processes + if self.current_position + (B * T * self.num_processes + 1) > len(self.tokens): + self.advance() + return x.cuda(), y.cuda() + +# ----------------------------------------------------------------------------- +# int main + +@dataclass +class Hyperparameters: + # data hyperparams + input_bin : str = 'data/fineweb10B/fineweb_train_*.bin' # input .bin to train on + input_val_bin : str = 'data/fineweb10B/fineweb_val_*.bin' # input .bin to eval validation loss on + # optimization hyperparams + batch_size : int = 8*64 # batch size, in sequences, across all devices + device_batch_size : int = 64 # batch size, in sequences, per device + sequence_length : int = 1024 # sequence length, in tokens + num_iterations : int = 3000 # number of iterations to run + warmup_iters : int = 0 + warmdown_iters : int = 900 # number of iterations of linear warmup/warmdown for triangular or trapezoidal schedule + weight_decay : float = 0 + # evaluation and logging hyperparams + val_loss_every : int = 125 # every how many steps to evaluate val loss? 0 for only at the end + val_tokens : int = 10485760 # how many tokens of validation data? it's important to keep this fixed for consistent comparisons + save_every : int = 0 # every how many steps to save the checkpoint? 0 for only at the end +args = Hyperparameters() + +# set up DDP (distributed data parallel). torchrun sets this env variable +assert torch.cuda.is_available() +dist.init_process_group(backend='nccl') +ddp_rank = int(os.environ['RANK']) +ddp_local_rank = int(os.environ['LOCAL_RANK']) +ddp_world_size = int(os.environ['WORLD_SIZE']) +device = f'cuda:{ddp_local_rank}' +torch.cuda.set_device(device) +print(f"using device: {device}") +master_process = (ddp_rank == 0) # this process will do logging, checkpointing etc. + +# convenience variables +B, T = args.device_batch_size, args.sequence_length +# calculate the number of steps to take in the val loop. +assert args.val_tokens % (B * T * ddp_world_size) == 0 +val_steps = args.val_tokens // (B * T * ddp_world_size) +# calculate the steps of gradient accumulation required to attain the desired global batch size. +assert args.batch_size % (B * ddp_world_size) == 0 +train_accumulation_steps = args.batch_size // (B * ddp_world_size) + +# load tokens +train_loader = DistributedDataLoader(args.input_bin, B, T, ddp_rank, ddp_world_size) +val_loader = DistributedDataLoader(args.input_val_bin, B, T, ddp_rank, ddp_world_size) +if master_process: + print(f"Training DataLoader: total number of tokens: {train_loader.ntok_total} across {len(train_loader.files)} files") + print(f"Validation DataLoader: total number of tokens: {val_loader.ntok_total} across {len(val_loader.files)} files") +x, y = train_loader.next_batch() + +# there are only 50257 unique GPT-2 tokens; we extend to nearest multiple of 128 for efficiency. suggested to me by @Grad62304977. +# this originates from Karpathy's experiments. +num_vocab = 50304 +model = GPT(GPTConfig(vocab_size=num_vocab, n_layer=12, n_head=6, n_embd=768)) +model = model.cuda().bfloat16() +for m in model.modules(): + if isinstance(m, CastedLinear): + m.float() + +if hasattr(config, "coordinate_descent_tuning"): + config.coordinate_descent_tuning = True # suggested by @Chillee +model = torch.compile(model) +# here we wrap model into DDP container +model = DDP(model, device_ids=[ddp_local_rank]) +raw_model = model.module # always contains the "raw" unwrapped model + +# CUDNN attention is ~4ms faster than Flash, but doesn't get selected by default in PyTorch 2.5.1 +from torch.backends.cuda import enable_cudnn_sdp, enable_flash_sdp, enable_math_sdp, enable_mem_efficient_sdp +enable_cudnn_sdp(True) +enable_flash_sdp(False) +enable_mem_efficient_sdp(False) +enable_math_sdp(False) + +# init the optimizer(s) +optimizer1 = torch.optim.Adam([raw_model.transformer.wte.weight], lr=0.6, betas=(0.9, 0.95), fused=True) +optimizer2 = torch.optim.Adam([raw_model.lm_head.weight], lr=0.008, betas=(0.9, 0.95), fused=True) +params = list(raw_model.transformer.h.parameters()) +matrix_params = [p for p in params if p.ndim == 2] +scalar_params = [p for p in params if p.ndim < 2]+[raw_model.skip_weights] +optimizer3 = Muon(matrix_params, lr=0.04, momentum=0.95) +optimizer4 = torch.optim.Adam(scalar_params, lr=0.04, betas=(0.9, 0.95), fused=True) # note that this learning rate is neither sensitive nor tuned +optimizers = [optimizer1, optimizer2, optimizer3, optimizer4] +# learning rate decay scheduler (linear warmup and warmdown) +def get_lr(it): + assert it <= args.num_iterations + # 1) linear warmup for warmup_iters steps + if it < args.warmup_iters: + return (it+1) / args.warmup_iters + # 2) constant lr for a while + elif it < args.num_iterations - args.warmdown_iters: + return 1.0 + # 3) linear warmdown + else: + decay_ratio = (args.num_iterations - it) / args.warmdown_iters + return decay_ratio +schedulers = [torch.optim.lr_scheduler.LambdaLR(opt, get_lr) for opt in optimizers] + +# begin logging +if master_process: + run_id = str(uuid.uuid4()) + logdir = 'logs/%s/' % run_id + os.makedirs(logdir, exist_ok=True) + logfile = 'logs/%s.txt' % run_id + # create the log file + with open(logfile, "w") as f: + # begin the log by printing this file (the Python code) + f.write('='*100 + '\n') + f.write(code) + f.write('='*100 + '\n') + # log information about the hardware/software environment this is running on + # and print the full `nvidia-smi` to file + f.write(f"Running pytorch {torch.version.__version__} compiled for CUDA {torch.version.cuda}\nnvidia-smi:\n") + import subprocess + result = subprocess.run(['nvidia-smi'], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True) + f.write(f'{result.stdout}\n') + f.write('='*100 + '\n') + +training_time_ms = 0 +# start the clock +torch.cuda.synchronize() +t0 = time.time() +# begin training +train_loader.reset() +for step in range(args.num_iterations + 1): + last_step = (step == args.num_iterations) + # This effectively ignores timing first 10 steps, which are slower for weird reasons. + # Alternately, and slightly more correctly in terms of benchmarking, we could do 10 + # steps with dummy data first, and then re-initialize the model and reset the loader. + if step == 10: + training_time_ms = 0 + t0 = time.time() + timed_steps = float('nan') if step <= 11 else (step - 10) + 1 # <= 11 to avoid bug in val + + # once in a while evaluate the validation dataset + if (last_step or (args.val_loss_every > 0 and step % args.val_loss_every == 0)): + # stop the clock + torch.cuda.synchronize() + training_time_ms += 1000 * (time.time() - t0) + # run validation batches + model.eval() + val_loader.reset() + val_loss = 0.0 + for _ in range(val_steps): + with torch.no_grad(): + x_val, y_val = val_loader.next_batch() + val_loss += model(x_val, y_val) + dist.all_reduce(val_loss, op=dist.ReduceOp.AVG) + val_loss /= val_steps + # log val loss to console and to logfile + if master_process: + print(f'step:{step}/{args.num_iterations} val_loss:{val_loss:.4f} train_time:{training_time_ms:.0f}ms step_avg:{training_time_ms/(timed_steps-1):.2f}ms') + with open(logfile, "a") as f: + f.write(f'step:{step}/{args.num_iterations} val_loss:{val_loss:.4f} train_time:{training_time_ms:.0f}ms step_avg:{training_time_ms/(timed_steps-1):.2f}ms\n') + # start the clock again + torch.cuda.synchronize() + t0 = time.time() + + if master_process and (last_step or (args.save_every > 0 and step % args.save_every == 0)): + # stop the clock + torch.cuda.synchronize() + training_time_ms += 1000 * (time.time() - t0) + # save the state of the training process + log = dict(step=step, code=code, model=raw_model.state_dict(), optimizers=[opt.state_dict() for opt in optimizers]) + torch.save(log, 'logs/%s/state_step%06d.pt' % (run_id, step)) + # start the clock again + torch.cuda.synchronize() + t0 = time.time() + + # bit confusing: we want to make sure to eval on 0th iteration + # but also after the very last iteration. so we loop for step <= num_iterations + # instead of just < num_iterations (one extra due to <=), only to do + # the validation/sampling one last time, and then we break right here as we're done. + if last_step: + break + + # --------------- TRAINING SECTION BEGIN ----------------- + model.train() + for i in range(1, train_accumulation_steps+1): + # forward pass + loss = model(x, y) + train_loss = loss.detach() + # advance the dataset for the next batch + x, y = train_loader.next_batch() + # backward pass + if i < train_accumulation_steps: + with model.no_sync(): # there's no need to sync gradients every accumulation step + loss.backward() + else: + loss.backward() # just sync on the last step + for p in model.parameters(): + p.grad /= train_accumulation_steps + # momentum warmup for Muon + frac = min(step/500, 1) + optimizer3.param_groups[0]['momentum'] = (1 - frac) * 0.85 + frac * 0.95 + # step the optimizers and schedulers + for opt, sched in zip(optimizers, schedulers): + opt.step() + sched.step() + # null the gradients + model.zero_grad(set_to_none=True) + # --------------- TRAINING SECTION END ------------------- + # everything that follows now is just diagnostics, prints, logging, etc. + + #dist.all_reduce(train_loss, op=dist.ReduceOp.AVG) # all-reducing the training loss would be more correct in terms of logging, but slower + if master_process: + approx_time = training_time_ms + 1000 * (time.time() - t0) + print(f"step:{step+1}/{args.num_iterations} train_loss:{train_loss.item():.4f} train_time:{approx_time:.0f}ms step_avg:{approx_time/timed_steps:.2f}ms") + with open(logfile, "a") as f: + f.write(f"step:{step+1}/{args.num_iterations} train_loss:{train_loss.item():.4f} train_time:{approx_time:.0f}ms step_avg:{approx_time/timed_steps:.2f}ms\n") + +if master_process: + print(f"peak memory consumption: {torch.cuda.max_memory_allocated() // 1024 // 1024} MiB") + +# ------------------------------------------------------------------------- +# clean up nice +dist.destroy_process_group() + diff --git a/baseline/gpt2_small/muon_speedrun/worker.py b/baseline/gpt2_small/muon_speedrun/worker.py new file mode 100644 index 00000000..3f09e2cd --- /dev/null +++ b/baseline/gpt2_small/muon_speedrun/worker.py @@ -0,0 +1,121 @@ +"""Bounded setup, optional flash check, one fresh training run, verified backup.""" +import argparse +import json +import os +from pathlib import Path +import signal +import subprocess +import sys +import time + + +def write(root, name, value): + temp = root/(name+'.tmp') + temp.write_text(json.dumps(value, indent=2)+'\n') + temp.replace(root/name) + + +def bounded(command, seconds, root, label, watch=False): + if seconds <= 0: + return {'exit_code':None, 'timed_out':True, 'phase':label} + print('START '+label, flush=True) + child = subprocess.Popen(command, start_new_session=True) + end = time.monotonic()+seconds + try: + while True: + remaining = end-time.monotonic() + if remaining <= 0: + raise subprocess.TimeoutExpired(command, seconds) + try: + return {'exit_code':child.wait(timeout=min(30, remaining)), 'timed_out':False, 'phase':label} + except subprocess.TimeoutExpired: + print('WAIT '+label+': '+str(int(remaining))+'s to phase cutoff', flush=True) + status = root/'status.json' + if watch and status.exists() and time.time()-status.stat().st_mtime > 900: + raise RuntimeError('No training progress recorded for 15 minutes') + except (subprocess.TimeoutExpired, RuntimeError) as exc: + try: + os.killpg(child.pid, signal.SIGKILL) + except ProcessLookupError: + pass + child.wait(timeout=10) + return {'exit_code':child.returncode, 'timed_out':True, 'phase':label, 'error':str(exc)} + + +def backup(root): + from rg_nanogpt_one_head.continuous_support import CloudPublisher + publisher = CloudPublisher('gs://tpu-builders-504820-ww-continuous8/gpt2small/'+root.name) + receipts = [] + for path in sorted(root.iterdir()): + if path.is_file() and path.suffix in ('.json', '.jsonl', '.pt', '.txt', '.log'): + # Logs may still grow as supervisor/upload output is appended. + if path.suffix == '.log': + publisher.snapshot_text_file(path, path.name) + else: + receipts.append(publisher.file(path, path.name)) + publisher.json({'files':receipts, 'status':'verified'}, 'CLOUD_BACKUP_VERIFIED.json') + write(root, 'CLOUD_BACKUP_VERIFIED.json', {'files':receipts, 'status':'verified'}) + print('Cloud backup verified: '+root.name, flush=True) + + +def main(): + p = argparse.ArgumentParser() + p.add_argument('root', type=Path) + p.add_argument('deadline', type=float) + p.add_argument('--optimizer', choices=('muon','adam'), default='muon') + p.add_argument('--microbatch', type=int, choices=(64,128), default=128) + p.add_argument('--attention', choices=('auto','flash','math'), default='auto') + p.add_argument('--backup-only', action='store_true') + a = p.parse_args() + if a.backup_only: + backup(a.root) + return 0 + here = Path(__file__).resolve().parent + run = {'status':'preparing', 'target_met':False, 'automatic_restart':False, + 'deadline_unix':a.deadline, 'optimizer':a.optimizer} + write(a.root, 'RUN_STATUS.json', run) + common = [sys.executable, '-u', str(here/'run.py')] + args = ['--root',str(a.root),'--microbatch',str(a.microbatch)] + train_deadline = a.deadline-600 + try: + prep_deadline = min(time.time()+900, train_deadline-300) + result = bounded(common+['prepare',*args,'--deadline',str(prep_deadline)], + prep_deadline-time.time(), a.root, 'benchmark data') + if result['exit_code'] != 0: + raise RuntimeError('Data preparation failed: '+str(result)) + attention = 'math' + if a.attention != 'math': + check_deadline = min(time.time()+300, train_deadline-300) + checked = bounded(common+['attention-check',*args,'--deadline',str(check_deadline)], + check_deadline-time.time(), a.root, 'TPU flash attention forward/backward') + if checked['exit_code'] == 0: + attention = 'flash' + else: + write(a.root, 'ATTENTION_CHECK_FAILURE.json', checked) + if a.attention == 'flash': + raise RuntimeError('Requested flash attention failed validation') + print('TPU flash check unavailable/failed; using mathematical attention. See ATTENTION_CHECK_FAILURE.json.', flush=True) + run.update(status='training', attention=attention) + write(a.root, 'RUN_STATUS.json', run) + result = bounded(common+['train',*args,'--deadline',str(train_deadline), + '--attention',attention,'--optimizer',a.optimizer], + train_deadline-time.time(), a.root, '3,000-update '+a.optimizer+' recipe', watch=True) + run.update(result) + if result['exit_code'] != 0: + run['status'] = 'failed_or_timed_out' + else: + status = json.loads((a.root/'status.json').read_text()) + run.update(status=status['status'], target_met=status['target_met'], step=status['step']) + except Exception as exc: + run.update(status='failed', error=str(exc)) + write(a.root, 'RUN_STATUS.json', run) + backed = bounded([sys.executable,'-u',__file__,str(a.root),str(a.deadline),'--backup-only'], + min(590, a.deadline-time.time()-10), a.root, 'cloud backup') + run['backup'] = backed + write(a.root, 'RUN_STATUS.json', run) + print(json.dumps(run), flush=True) + return 0 if run['status'] in ('target_reached','schedule_complete_target_not_met') else 1 + + +if __name__ == '__main__': + sys.exit(main()) diff --git a/baseline/gpt2_small/scripts/run_muonclip.py b/baseline/gpt2_small/scripts/run_muonclip.py index e0469211..9d2e50c8 100644 --- a/baseline/gpt2_small/scripts/run_muonclip.py +++ b/baseline/gpt2_small/scripts/run_muonclip.py @@ -28,6 +28,9 @@ def active(unit): def assert_idle(): + speedrun=BASE/'MUON_SPEEDRUN_LATEST.json' + if speedrun.exists() and active(json.loads(speedrun.read_text())['unit']): + raise RuntimeError('Muon speedrun is active; no concurrent trainer launched.') reference=BASE/'SPEEDRUN30_LATEST.json' if reference.exists() and active(json.loads(reference.read_text())['unit']): raise RuntimeError('The 30-minute GPT-2 reference run is active; no concurrent trainer launched.') diff --git a/baseline/gpt2_small/tests/test_muon_speedrun.py b/baseline/gpt2_small/tests/test_muon_speedrun.py new file mode 100644 index 00000000..41d5b49d --- /dev/null +++ b/baseline/gpt2_small/tests/test_muon_speedrun.py @@ -0,0 +1,168 @@ +"""CPU parity tests against the pinned record, plus target/checkpoint guard tests.""" +import ast +import importlib.util +import json +from pathlib import Path +import sys +import types + +import pytest +import torch +import torch.nn.functional as F + +BASE = Path(__file__).resolve().parents[1]/'muon_speedrun' +sys.path.insert(0, str(BASE)) +import model as port +from optim import Muon, make_optimizers, apply_update, schedule +from runtime import Runtime +from data import FineWeb, required_shards + + +def original(): + tree = ast.parse((BASE/'vendor/record_source.py').read_text()) + # Load only definitions. The original top-level CUDA/DDP launcher must not run. + accepted = {'Rotary','apply_rotary_emb','CastedLinear','CausalSelfAttention','MLP','Block', + 'GPTConfig','GPT','zeropower_via_newtonschulz5','Hyperparameters'} + nodes = [] + for n in tree.body: + if isinstance(n, (ast.FunctionDef,ast.ClassDef)) and n.name in accepted: + if isinstance(n,ast.FunctionDef): + n.decorator_list = [] + nodes.append(n) + import dataclasses + namespace = dict(torch=torch, nn=torch.nn, F=F, dataclass=dataclasses.dataclass) + exec(compile(ast.Module(body=nodes,type_ignores=[]),'record-definitions','exec'),namespace) + return types.SimpleNamespace(**namespace) + + +def small(module): + result = module.GPT(module.GPTConfig(vocab_size=128,n_layer=2,n_head=2,n_embd=16)).bfloat16() + for layer in result.modules(): + if isinstance(layer,module.CastedLinear): + layer.float() + return result + + +def test_model_forward_backward_matches_pinned_record(): + torch.set_num_threads(1) + source = original() + torch.manual_seed(1337) + expected = small(source) + actual = small(port) + actual.load_state_dict(expected.state_dict()) + # Un-zero output projections to exercise every block, not just the first head update. + with torch.no_grad(): + for model in (expected,actual): + gen = torch.Generator().manual_seed(8) + for name,p in model.named_parameters(): + if 'c_proj' in name or 'lm_head' in name: + p.copy_(torch.randn(p.shape,generator=gen)*.02) + x = torch.randint(128,(4,8)); y = torch.randint(128,(4,8)) + a,b = expected(x,y),actual(x,y) + torch.testing.assert_close(a,b,rtol=0,atol=0) + a.backward(); b.backward() + for p,q in zip(expected.parameters(),actual.parameters()): + torch.testing.assert_close(p.grad,q.grad,rtol=0,atol=0) + + +def test_batched_muon_matches_record_updates_and_restores_state(): + torch.set_num_threads(1) + source = original() + rt = Runtime('cpu') + torch.manual_seed(5) + params = [torch.nn.Parameter(torch.randn(shape)) for shape in ((8,8),(8,8),(16,8),(8,16))] + expected = [p.detach().clone() for p in params] + buffers = [torch.zeros_like(p) for p in params] + opt = Muon([(str(i),p) for i,p in enumerate(params)],rt) + for step in range(4): + beta = .85+.1*step/500 + grads = [torch.randn_like(p) for p in params] + for p,g in zip(params,grads): p.grad = g.clone() + opt.step(.04,beta) + for i,(p,g) in enumerate(zip(expected,grads)): + buffers[i].mul_(beta).add_(g) + update = source.zeropower_via_newtonschulz5(g+beta*buffers[i],steps=5) + update *= max(1,update.size(0)/update.size(1))**.5 + p.add_(update.to(p.dtype),alpha=-.04) + for a,b in zip(params,expected): + torch.testing.assert_close(a,b,rtol=1e-5,atol=3e-7) + state = opt.state_dict() + opt = Muon([(str(i),p) for i,p in enumerate(params)],rt) + opt.load_state_dict(state) + + +def test_complete_mixed_dtype_recipe_learns_on_cpu(): + torch.set_num_threads(1) + torch.manual_seed(13) + model = small(port) + rt = Runtime('cpu') + muon,adam = make_optimizers(model,rt) + x = torch.randint(128,(4,8)); y = x.clone() + losses=[] + for step in range(6): + model.zero_grad(set_to_none=False) + loss = model(x,y) + assert torch.isfinite(loss) + loss.backward() + apply_update(muon,adam,rt,step) + losses.append(float(loss.detach())) + assert losses[-1] < losses[0] + assert all(torch.isfinite(p).all() for p in model.parameters()) + + +def runner(): + spec = importlib.util.spec_from_file_location('muon_recipe_runner',BASE/'run.py') + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def test_partial_or_nonfinite_validation_cannot_pass(): + run = runner() + good = dict(full_benchmark_evaluation=True,evaluation_tokens=10485760,val_nll=3.28) + assert run.target_met(good) + for change in (dict(full_benchmark_evaluation=False),dict(evaluation_tokens=1048576), + dict(val_nll=float('nan')),dict(val_nll=3.29)): + assert not run.target_met({**good,**change}) + + +def test_best_checkpoint_survives_latest_replacement(tmp_path): + run = runner(); rt=Runtime('cpu') + model=small(port); muon,adam=make_optimizers(model,rt) + stream=types.SimpleNamespace(shard=1,position=131072) + validation=dict(full_benchmark_evaluation=True,evaluation_tokens=10485760,val_nll=3.5) + run.save_checkpoint(tmp_path,model,muon,adam,stream,125,{},rt,validation,4.) + run.save_checkpoint(tmp_path,model,muon,adam,stream,250,{},rt,{**validation,'val_nll':3.6},3.5) + best=torch.load(tmp_path/'checkpoint_best.pt',weights_only=False) + latest=torch.load(tmp_path/'checkpoint_latest.pt',weights_only=False) + assert best['step']==125 and latest['step']==250 + assert latest['data_cursor']==dict(shard=1,position=131072) + assert latest['muon'] and 'adam' in latest and 'rng' in latest + + +def test_schedule_and_full_corpus_budget(tmp_path): + source=original().Hyperparameters() + assert source.num_iterations==3000 and source.warmdown_iters==900 + assert schedule(0)==schedule(2099)==1. + assert schedule(2550)==.5 and schedule(3000)==0. + files=FineWeb(tmp_path,0) + for micro in (64,128): + names=required_shards(files,micro) + assert len(names)==17 # validation + 16 distinct training shards + usable=sum(((files.manifest['files'][n]['size']-1024)//2-1)//(micro*1024)*(micro*1024) + for n in names if '_train_' in n) + assert usable >= 3000*524288 + + +def test_launcher_does_not_duplicate_active_run(tmp_path,monkeypatch): + spec=importlib.util.spec_from_file_location('muon_recipe_launcher',BASE/'cloudshell.py') + module=importlib.util.module_from_spec(spec); spec.loader.exec_module(module) + pointer=tmp_path/'MUON_SPEEDRUN_LATEST.json'; pointer.write_text(json.dumps({'unit':'active.service'})) + monkeypatch.setattr(module,'BASE',tmp_path); monkeypatch.setattr(module,'LATEST',pointer) + monkeypatch.setattr(module.os,'geteuid',lambda:0) + monkeypatch.setattr(module.os.path,'ismount',lambda p:True) + monkeypatch.setattr(module,'active',lambda u:True) + seen=[]; monkeypatch.setattr(module,'status_remote',lambda:seen.append(True)) + monkeypatch.setattr(module,'run',lambda *a,**k:pytest.fail('must not launch twice')) + module.start_remote('a'*40) + assert seen==[True] From c2d7578659d16d5d9436d5c46e16de70927f0196 Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Sun, 4 Oct 2026 19:06:40 -0700 Subject: [PATCH 27/36] Install pinned Pallas dependencies and require flash attention with smaller TPU microbatch --- baseline/gpt2_small/TPU_PORT_BUGS.md | 23 ++++++ baseline/gpt2_small/muon_speedrun/README.md | 30 +++++--- .../gpt2_small/muon_speedrun/cloudshell.py | 11 +-- .../muon_speedrun/pallas_dependencies.py | 70 +++++++++++++++++++ baseline/gpt2_small/muon_speedrun/run.py | 6 +- baseline/gpt2_small/muon_speedrun/runtime.py | 9 +-- baseline/gpt2_small/muon_speedrun/worker.py | 26 +++++-- .../gpt2_small/tests/test_muon_speedrun.py | 52 +++++++++++++- 8 files changed, 199 insertions(+), 28 deletions(-) create mode 100644 baseline/gpt2_small/muon_speedrun/pallas_dependencies.py diff --git a/baseline/gpt2_small/TPU_PORT_BUGS.md b/baseline/gpt2_small/TPU_PORT_BUGS.md index 3f31d69f..b63cde16 100644 --- a/baseline/gpt2_small/TPU_PORT_BUGS.md +++ b/baseline/gpt2_small/TPU_PORT_BUGS.md @@ -5,6 +5,29 @@ an explicit objective of this project. Preserve failures as evidence. A passing CPU smoke test does not establish TPU correctness. Do not silently lower learning rates, change data, or skip nonfinite checks to make a run appear successful. +## 2026-10-04: Muon speedrun missing Pallas dependency and HBM exhaustion + +Status: two configuration failures confirmed; corrected retry requires live TPU validation. +These observations do not establish an upstream PyTorch/XLA or hardware bug. + +- Run `muon-speedrun-muon-20261004-225635`, commit `1b4548d`, v5litepod-8. +- Flash check failed with `ModuleNotFoundError: No module named 'jax'`. + `torch_xla[tpu]` does not include the Pallas extras; XLA 2.6 setup.py pins + both JAX and jaxlib to 0.4.38 for that optional dependency group. +- Automatic math-attention fallback kept global microbatch 128 (16/chip). + First forward/backward compilation failed with `RESOURCE_EXHAUSTED`: + 16.89G required versus 15.75G HBM, exceeding capacity by 1.14G. +- No completed training update. Step-zero checkpoint and final cloud backup + were saved. Disk evidence: + `/mnt/disks/rg-data/gpt2small/muon-speedrun-muon-20261004-225635`; + cloud prefix: `gs://tpu-builders-504820-ww-continuous8/gpt2small/muon-speedrun-muon-20261004-225635`. +- Correction: per-run pinned Pallas overlay; flash forward/backward check required; + no implicit math fallback; microbatch 64, eight accumulation passes, unchanged + global batch 524,288; early checkpoints after updates 1 and 5. Failure status now + replaces stale `training` status and includes the actual exception. +- The smaller microbatch and kernel must still pass a live full-model run. Do not + report the OOM fixed solely because CPU tests or attention-only checks pass. + ## 2026-10-04: GPT-2 AdamW nonfinite result, followed by stalled failure reporting Status: numerical cause open; blocking diagnostic implementation replaced. diff --git a/baseline/gpt2_small/muon_speedrun/README.md b/baseline/gpt2_small/muon_speedrun/README.md index 0830de53..e4a21d4e 100644 --- a/baseline/gpt2_small/muon_speedrun/README.md +++ b/baseline/gpt2_small/muon_speedrun/README.md @@ -16,7 +16,7 @@ python3 baseline/gpt2_small/muon_speedrun/cloudshell.py start ``` The launcher uses the existing `ww-gpt2-validation-48h-20261004-s1337-node` in -`tpu-builders-504820/us-west4-a`. It creates no TPU, changes no installed packages, +`tpu-builders-504820/us-west4-a`. It creates no TPU, preserves the shared environment, and refuses concurrent training. Repeating start while active prints status. It starts from scratch because the previous 30-minute run saved no model. @@ -44,7 +44,7 @@ python3 baseline/gpt2_small/muon_speedrun/cloudshell.py status | Model changes | RoPE, RMSNorm/QK normalization, squared ReLU, zero output projections, value residuals, learned input/UNet skip weights, untied output head, logit soft cap 30 | | Parameters | About 162M total; this is a modified transformer, not standard GPT-2 Small | | Context / global batch | 1,024 tokens / 524,288 tokens per update | -| Default TPU microbatch | 128 sequences globally, four accumulation passes; 16 sequences per chip | +| Default TPU microbatch | 64 sequences globally, eight accumulation passes; 8 sequences per chip | | Schedule | 3,000 updates, zero warmup, constant LR through 2,100 then 900-update linear decay | | Hidden matrices | Muon LR 0.04; five quintic Newton–Schulz iterations | | Muon momentum | Linear ramp 0.85 to 0.95 over the first 500 updates; Nesterov | @@ -55,8 +55,10 @@ python3 baseline/gpt2_small/muon_speedrun/cloudshell.py status | Target | Full validation NLL <= 3.28, equivalent to perplexity <= exp(3.28) | All model and optimizer settings come from the source record. The global batch -stays fixed when microbatch size changes. `--microbatch 64` uses eight accumulation -passes if the 128-sequence configuration exceeds device memory. No automatic +stays fixed when microbatch size changes. `--microbatch 32` uses sixteen accumulation +passes for additional memory headroom. The first 128-sequence/math-attention run +failed at update zero: 16.89 GiB required versus 15.75 GiB available per chip. +The default is now 64 and flash attention is required. No automatic microbatch change or restart can silently alter a run. The selected microbatch is a starting point, not the result of a TPU tuning sweep. @@ -70,14 +72,20 @@ coefficients, Nesterov convention and rectangular scaling. It is **Muon, not MuonClip**. Learning rate and momentum are device tensors to avoid compiling a new graph solely because their Python values change. -The worker tests PyTorch/XLA 2.6 TPU flash attention at sequence length 1,024 and -head dimension 128 before using it. The check compares outputs and Q/K/V gradients +The worker installs JAX and jaxlib **0.4.38**, the exact Pallas versions specified +by [PyTorch/XLA 2.6 setup.py](https://github.com/pytorch/xla/blob/v2.6.0/setup.py), +plus pinned ml-dtypes/opt-einsum into a run-local `pallas-deps` overlay. It does +not upgrade torch, torch_xla, libtpu, NumPy, SciPy or the shared venv. It verifies +the imports and records versions before starting the TPU check. + +The worker tests PyTorch/XLA 2.6 TPU flash attention at the selected batch size, +sequence length 1,024 and head dimension 128. The check compares outputs and Q/K/V gradients against mathematical attention with BF16 relative-L2 tolerance 0.03. It supplies `sm_scale=1/sqrt(head_dim)` and the SPMD batch partition explicitly. The isolated -check is capped at five minutes. With default `--attention auto`, an unavailable -or failing kernel is logged and the job uses mathematical attention. Request -`--attention flash` to require a passing kernel, or `--attention math` to select -mathematical attention explicitly. No dependency upgrade is performed. +check is capped at five minutes. With default `--attention flash`, an unavailable +or failing kernel stops the job. The legacy `auto` option also requires a pass; +there is no implicit fallback to a memory-heavier attention implementation. +`--attention math` remains an explicit diagnostic option with microbatch <=64. The 16 necessary training shards plus validation are SHA256-verified and prepared before training, reusing `/mnt/disks/rg-data/benchmark-fineweb10B-889765ea`. This @@ -96,7 +104,7 @@ can differ; these are recorded rather than represented as exact reproduction. Runs live at `/mnt/disks/rg-data/gpt2small/muon-speedrun--`. `MUON_SPEEDRUN_LATEST.json` identifies the latest run and service. -- `checkpoint_latest.pt`: atomic full-state save at initialization, every 125 +- `checkpoint_latest.pt`: atomic full-state save at initialization, updates 1 and 5, every 125 updates, and normal stop. Includes model, optimizers, RNG, data cursor, recipe and next-step schedule. No automatic resume. TPU resume parity is not yet tested. - `checkpoint_best.pt`: best fully evaluated checkpoint; a hard link protects it diff --git a/baseline/gpt2_small/muon_speedrun/cloudshell.py b/baseline/gpt2_small/muon_speedrun/cloudshell.py index d2760642..d53252a2 100644 --- a/baseline/gpt2_small/muon_speedrun/cloudshell.py +++ b/baseline/gpt2_small/muon_speedrun/cloudshell.py @@ -38,13 +38,14 @@ def status_remote(): root = Path(record['root']) print(json.dumps(record,indent=2),flush=True) subprocess.run(['systemctl','--no-pager','--full','status',record['unit']]) - for name in ('RUN_STATUS.json','status.json','latest_validation.json','checkpoint_latest.json'): + for name in ('RUN_STATUS.json','status.json','latest_validation.json','checkpoint_latest.json', + 'PALLAS_DEPENDENCIES.json','ATTENTION_CHECK.json','ATTENTION_CHECK_FAILURE.json'): if (root/name).exists(): print(name+'\n'+(root/name).read_text(),flush=True) subprocess.run(['tail','-n','15',str(root/'run.log')]) -def start_remote(commit, hours=3, optimizer='muon', microbatch=128, attention='auto'): +def start_remote(commit, hours=3, optimizer='muon', microbatch=64, attention='flash'): if os.geteuid() != 0 or not os.path.ismount('/mnt/disks/rg-data'): raise RuntimeError('Requires the existing mounted disk and root') if not re.fullmatch('[0-9a-f]{40}',commit): @@ -76,7 +77,7 @@ def start_remote(commit, hours=3, optimizer='muon', microbatch=128, attention='a unit = 'rg-muon-speedrun-'+stamp+'.service' base = repo/'baseline/gpt2_small' env = {'PYTHONPATH':str(base/'src')+':'+str(base.parent/'nanogpt_one_head/src'), - 'PJRT_DEVICE':'TPU','XLA_USE_SPMD':'1','TPU_ACCELERATOR_TYPE':'v5litepod-8', + 'PJRT_DEVICE':'TPU','TPU_ACCELERATOR_TYPE':'v5litepod-8', 'OMP_NUM_THREADS':'4','OPENBLAS_NUM_THREADS':'4','MKL_NUM_THREADS':'4', 'TOKENIZERS_PARALLELISM':'false'} record = {'root':str(root),'unit':unit,'commit':commit,'optimizer':optimizer, @@ -110,8 +111,8 @@ def main(): p.add_argument('action',choices=('start','status')) p.add_argument('--hours',type=float,default=3) p.add_argument('--optimizer',choices=('muon','adam'),default='muon') - p.add_argument('--microbatch',type=int,choices=(64,128),default=128) - p.add_argument('--attention',choices=('auto','flash','math'),default='auto') + p.add_argument('--microbatch',type=int,choices=(32,64,128),default=64) + p.add_argument('--attention',choices=('auto','flash','math'),default='flash') p.add_argument('--on-tpu',action='store_true',help=argparse.SUPPRESS) p.add_argument('--commit',help=argparse.SUPPRESS) a = p.parse_args() diff --git a/baseline/gpt2_small/muon_speedrun/pallas_dependencies.py b/baseline/gpt2_small/muon_speedrun/pallas_dependencies.py new file mode 100644 index 00000000..bd0cd2c1 --- /dev/null +++ b/baseline/gpt2_small/muon_speedrun/pallas_dependencies.py @@ -0,0 +1,70 @@ +"""Install the XLA 2.6 Pallas extras into a per-run overlay, not the shared venv. + +Version authority: pytorch/xla v2.6.0 setup.py extras_require['pallas']. +JAX traces TPU kernels; PyTorch/XLA remains the training runtime. +""" +import argparse +from importlib import metadata +import json +import os +from pathlib import Path +import subprocess +import sys + +PACKAGES = ('jax==0.4.38', 'jaxlib==0.4.38', 'ml-dtypes==0.5.1', 'opt-einsum==3.4.0') +CORE = ('torch', 'torch-xla', 'libtpu', 'numpy', 'scipy') + + +def versions(): + result = {} + for name in CORE: + try: + result[name] = metadata.version(name) + except metadata.PackageNotFoundError: + result[name] = None + return result + + +def install_command(target): + return [sys.executable, '-m', 'pip', 'install', '--disable-pip-version-check', + '--no-deps', '--only-binary=:all:', '--timeout', '120', '--retries', '3', + '--target', str(target), *PACKAGES] + + +def main(): + p = argparse.ArgumentParser() + p.add_argument('root', type=Path) + a = p.parse_args() + before = versions() + if not (before['torch-xla'] or '').startswith('2.6.'): + raise RuntimeError('These Pallas pins require the installed PyTorch/XLA 2.6 runtime') + # The overlay intentionally reuses existing NumPy/SciPy rather than upgrading + # their binary dependencies in the known training environment. + from packaging.version import Version + for name, floor in (('numpy','1.24'),('scipy','1.10')): + if before[name] is None or Version(before[name]) < Version(floor): + raise RuntimeError(name+' does not satisfy JAX 0.4.38; no shared packages changed') + target = a.root/'pallas-deps' + subprocess.run(install_command(target), check=True) + env = dict(os.environ) + env['PYTHONPATH'] = str(target)+os.pathsep+env.get('PYTHONPATH','') + # Import verification only, on CPU. Training inherits no JAX_PLATFORMS override; + # Pallas must lower TPU code and uses XLA's jax_import_guard() for device ownership. + env['JAX_PLATFORMS'] = 'cpu' + code = '''import json, jax, jaxlib, ml_dtypes, opt_einsum +from jax.experimental.pallas.ops.tpu.flash_attention import _flash_attention_impl, _flash_attention_bwd_dq, _flash_attention_bwd_dkv +assert jax.__version__ == jaxlib.__version__ == "0.4.38" +print(json.dumps({"jax":jax.__version__, "jaxlib":jaxlib.__version__, "ml_dtypes":ml_dtypes.__version__, "opt_einsum":opt_einsum.__version__}))''' + result = subprocess.run([sys.executable,'-c',code],env=env,check=True,capture_output=True,text=True) + after = versions() + if after != before: + raise RuntimeError('Unexpected change to shared core-package versions') + record = {'status':'installed_and_import_verified','packages':json.loads(result.stdout), + 'overlay':str(target),'core_before':before,'core_after':after, + 'source':'https://github.com/pytorch/xla/blob/v2.6.0/setup.py'} + (a.root/'PALLAS_DEPENDENCIES.json').write_text(json.dumps(record,indent=2)+'\n') + print(json.dumps(record),flush=True) + + +if __name__ == '__main__': + main() diff --git a/baseline/gpt2_small/muon_speedrun/run.py b/baseline/gpt2_small/muon_speedrun/run.py index 791bb918..a21e8bda 100644 --- a/baseline/gpt2_small/muon_speedrun/run.py +++ b/baseline/gpt2_small/muon_speedrun/run.py @@ -192,6 +192,8 @@ def train(a): row['training_seconds_remaining_estimate'] = float(np.median(timings[-50:]))*(TOTAL_STEPS-step) record(root, row) write_json(root/'status.json', {'status':'training', **row}) + if step in (1,5): + save_checkpoint(root, model, muon, adam, stream, step, manifest, rt, None, best) if step % 125 == 0: validation = evaluate(model, val, rt, root, step, a.deadline-90, a.microbatch, started) save_checkpoint(root, model, muon, adam, stream, step, manifest, rt, validation, best) @@ -219,13 +221,13 @@ def main(): p.add_argument('--deadline', type=float, required=True) p.add_argument('--device', choices=('tpu', 'cpu'), default='tpu') p.add_argument('--attention', choices=('flash', 'math'), default='flash') - p.add_argument('--microbatch', type=int, choices=(64,128), default=128) + p.add_argument('--microbatch', type=int, choices=(32,64,128), default=64) p.add_argument('--optimizer', choices=('muon', 'adam'), default='muon') a = p.parse_args() if a.action == 'prepare': prepare(a.cache, a.deadline, a.root, a.microbatch) elif a.action == 'attention-check': - attention_check(a.root) + attention_check(a.root, a.microbatch) else: try: train(a) diff --git a/baseline/gpt2_small/muon_speedrun/runtime.py b/baseline/gpt2_small/muon_speedrun/runtime.py index d366b280..8902efa3 100644 --- a/baseline/gpt2_small/muon_speedrun/runtime.py +++ b/baseline/gpt2_small/muon_speedrun/runtime.py @@ -60,14 +60,14 @@ def flash(q, k, v): return lambda q, k, v: F.scaled_dot_product_attention(q, k, v, is_causal=True) -def attention_check(root): +def attention_check(root, microbatch=64): """Check forward/backward at the actual head/sequence shape in an isolated child.""" + rt = Runtime(cache=Path(root)/'xla-cache') from torch_xla.experimental.custom_kernel import jax_import_guard jax_import_guard() import jax - rt = Runtime(cache=Path(root)/'xla-cache') gen = torch.Generator().manual_seed(43) - cpu = [torch.randn((8, 6, 1024, 128), generator=gen).bfloat16() for _ in range(3)] + cpu = [torch.randn((microbatch, 6, 1024, 128), generator=gen).bfloat16() for _ in range(3)] upstream = torch.randn(cpu[0].shape, generator=gen).bfloat16() results = [] for kind in ('math', 'flash'): @@ -84,6 +84,7 @@ def attention_check(root): errors.append(relative) if relative > .03: raise RuntimeError('TPU flash/math relative L2 disagreement: '+str(relative)) - result = {'status':'passed', 'relative_l2_output_dq_dk_dv':errors, 'jax':jax.__version__} + result = {'status':'passed', 'relative_l2_output_dq_dk_dv':errors, 'jax':jax.__version__, + 'global_batch':microbatch, 'batch_per_chip':microbatch//8} (Path(root)/'ATTENTION_CHECK.json').write_text(json.dumps(result, indent=2)+'\n') print(json.dumps(result), flush=True) diff --git a/baseline/gpt2_small/muon_speedrun/worker.py b/baseline/gpt2_small/muon_speedrun/worker.py index 3f09e2cd..b85a75f8 100644 --- a/baseline/gpt2_small/muon_speedrun/worker.py +++ b/baseline/gpt2_small/muon_speedrun/worker.py @@ -63,8 +63,8 @@ def main(): p.add_argument('root', type=Path) p.add_argument('deadline', type=float) p.add_argument('--optimizer', choices=('muon','adam'), default='muon') - p.add_argument('--microbatch', type=int, choices=(64,128), default=128) - p.add_argument('--attention', choices=('auto','flash','math'), default='auto') + p.add_argument('--microbatch', type=int, choices=(32,64,128), default=64) + p.add_argument('--attention', choices=('auto','flash','math'), default='flash') p.add_argument('--backup-only', action='store_true') a = p.parse_args() if a.backup_only: @@ -78,6 +78,15 @@ def main(): args = ['--root',str(a.root),'--microbatch',str(a.microbatch)] train_deadline = a.deadline-600 try: + if a.attention == 'math' and a.microbatch > 64: + raise RuntimeError('Mathematical attention with microbatch 128 exceeded this TPU memory; use <=64') + if a.attention != 'math': + installed = bounded([sys.executable,'-u',str(here/'pallas_dependencies.py'),str(a.root)], + min(600,train_deadline-time.time()-300),a.root,'pinned Pallas dependencies') + if installed['exit_code'] != 0: + write(a.root,'PALLAS_DEPENDENCY_FAILURE.json',installed) + raise RuntimeError('Pallas dependency setup failed; training not started') + os.environ['PYTHONPATH'] = str(a.root/'pallas-deps')+os.pathsep+os.environ.get('PYTHONPATH','') prep_deadline = min(time.time()+900, train_deadline-300) result = bounded(common+['prepare',*args,'--deadline',str(prep_deadline)], prep_deadline-time.time(), a.root, 'benchmark data') @@ -92,9 +101,7 @@ def main(): attention = 'flash' else: write(a.root, 'ATTENTION_CHECK_FAILURE.json', checked) - if a.attention == 'flash': - raise RuntimeError('Requested flash attention failed validation') - print('TPU flash check unavailable/failed; using mathematical attention. See ATTENTION_CHECK_FAILURE.json.', flush=True) + raise RuntimeError('TPU flash attention failed validation; no automatic mathematical-attention fallback') run.update(status='training', attention=attention) write(a.root, 'RUN_STATUS.json', run) result = bounded(common+['train',*args,'--deadline',str(train_deadline), @@ -103,11 +110,20 @@ def main(): run.update(result) if result['exit_code'] != 0: run['status'] = 'failed_or_timed_out' + failure = a.root/'FAILURE.json' + if failure.exists(): + run['error'] = json.loads(failure.read_text()).get('error','')[:2000] else: status = json.loads((a.root/'status.json').read_text()) run.update(status=status['status'], target_met=status['target_met'], step=status['step']) except Exception as exc: run.update(status='failed', error=str(exc)) + if run['status'] in ('failed','failed_or_timed_out'): + previous = {} + if (a.root/'status.json').exists(): + previous = json.loads((a.root/'status.json').read_text()) + write(a.root,'status.json',{'status':run['status'],'last_recorded_step':previous.get('step'), + 'error':run.get('error'), 'target_met':False}) write(a.root, 'RUN_STATUS.json', run) backed = bounded([sys.executable,'-u',__file__,str(a.root),str(a.deadline),'--backup-only'], min(590, a.deadline-time.time()-10), a.root, 'cloud backup') diff --git a/baseline/gpt2_small/tests/test_muon_speedrun.py b/baseline/gpt2_small/tests/test_muon_speedrun.py index 41d5b49d..7f7470c6 100644 --- a/baseline/gpt2_small/tests/test_muon_speedrun.py +++ b/baseline/gpt2_small/tests/test_muon_speedrun.py @@ -146,7 +146,7 @@ def test_schedule_and_full_corpus_budget(tmp_path): assert schedule(0)==schedule(2099)==1. assert schedule(2550)==.5 and schedule(3000)==0. files=FineWeb(tmp_path,0) - for micro in (64,128): + for micro in (32,64,128): names=required_shards(files,micro) assert len(names)==17 # validation + 16 distinct training shards usable=sum(((files.manifest['files'][n]['size']-1024)//2-1)//(micro*1024)*(micro*1024) @@ -166,3 +166,53 @@ def test_launcher_does_not_duplicate_active_run(tmp_path,monkeypatch): monkeypatch.setattr(module,'run',lambda *a,**k:pytest.fail('must not launch twice')) module.start_remote('a'*40) assert seen==[True] + + +def test_pallas_install_is_pinned_and_isolated(tmp_path): + spec=importlib.util.spec_from_file_location('pallas_dependency_setup',BASE/'pallas_dependencies.py') + module=importlib.util.module_from_spec(spec); spec.loader.exec_module(module) + target=tmp_path/'pallas-deps' + command=module.install_command(target) + assert '--no-deps' in command and '--only-binary=:all:' in command + assert command[command.index('--target')+1]==str(target) + assert 'jax==0.4.38' in command and 'jaxlib==0.4.38' in command + assert not any(p.startswith(('torch==','torch_xla==','libtpu==','numpy==','scipy==')) for p in command) + + +@pytest.mark.parametrize('attention',['flash','auto']) +def test_flash_failure_never_starts_training(monkeypatch,tmp_path,attention): + spec=importlib.util.spec_from_file_location('muon_worker_strict_attention',BASE/'worker.py') + module=importlib.util.module_from_spec(spec); spec.loader.exec_module(module) + import time + monkeypatch.setattr(sys,'argv',['worker.py',str(tmp_path),str(time.time()+10800),'--attention',attention]) + phases=[] + def bounded(command,seconds,root,label,watch=False): + phases.append((label,command)) + rc=1 if 'forward/backward' in label else 0 + return {'exit_code':rc,'timed_out':False,'phase':label} + monkeypatch.setattr(module,'bounded',bounded) + monkeypatch.setenv('PYTHONPATH','test-original') + assert module.main()==1 + assert not any('3,000-update' in label for label,_ in phases) + assert json.loads((tmp_path/'RUN_STATUS.json').read_text())['status']=='failed' + assert 'fallback' in json.loads((tmp_path/'status.json').read_text())['error'] + + +def test_worker_defaults_use_smaller_microbatch_and_verified_flash(monkeypatch,tmp_path): + spec=importlib.util.spec_from_file_location('muon_worker_defaults',BASE/'worker.py') + module=importlib.util.module_from_spec(spec); spec.loader.exec_module(module) + import time + monkeypatch.setattr(sys,'argv',['worker.py',str(tmp_path),str(time.time()+10800)]) + phases=[] + def bounded(command,seconds,root,label,watch=False): + phases.append((label,command)) + if '3,000-update' in label: + (root/'status.json').write_text(json.dumps(dict(status='target_reached',target_met=True,step=2875))) + return {'exit_code':0,'timed_out':False,'phase':label} + monkeypatch.setattr(module,'bounded',bounded) + monkeypatch.setenv('PYTHONPATH','test-original') + assert module.main()==0 + training=next(command for label,command in phases if '3,000-update' in label) + assert training[training.index('--microbatch')+1]=='64' + assert training[training.index('--attention')+1]=='flash' + assert phases[0][0]=='pinned Pallas dependencies' From e0a11595103ac6e1056c783c1d0badf2280c9669 Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Sun, 4 Oct 2026 19:35:34 -0700 Subject: [PATCH 28/36] Add paired WeightWatcher alpha and validation token error to unchanged Muon recipe Keep training settings and the 125-update evaluation cadence unchanged. Compute token error from the same validation logits; analyze immutable transformer-weight snapshots in a separate CPU process. Preserve per-layer raw/clipped fits, paired tables, and tracking status on disk and back up the tables to cloud. Add an explicit stop-and-fresh-start option that is idempotent across SSH retries. Validation: 27 focused CPU tests passed, including source-model/optimizer parity, exact token counts, immutable snapshot pairing and real WeightWatcher analysis. Separate-process WeightWatcher smoke test passed. TPU verification remains live. --- baseline/gpt2_small/README.md | 5 +- baseline/gpt2_small/muon_speedrun/README.md | 53 +++- .../gpt2_small/muon_speedrun/cloudshell.py | 39 ++- baseline/gpt2_small/muon_speedrun/model.py | 5 +- baseline/gpt2_small/muon_speedrun/run.py | 17 +- baseline/gpt2_small/muon_speedrun/tracking.py | 247 ++++++++++++++++++ baseline/gpt2_small/muon_speedrun/worker.py | 46 +++- .../gpt2_small/tests/test_muon_speedrun.py | 124 +++++++++ 8 files changed, 522 insertions(+), 14 deletions(-) create mode 100644 baseline/gpt2_small/muon_speedrun/tracking.py diff --git a/baseline/gpt2_small/README.md b/baseline/gpt2_small/README.md index a8c671dc..408c92e3 100644 --- a/baseline/gpt2_small/README.md +++ b/baseline/gpt2_small/README.md @@ -2,8 +2,9 @@ For the **published 3,000-update Muon speedrun recipe on eight TPU chips**, see [muon_speedrun/README.md](muon_speedrun/README.md). It uses a modified transformer, -targets full-validation loss 3.28, saves checkpoints every 125 updates, and records -an explicit target outcome. Its TPU performance and convergence require live +targets full-validation loss 3.28, saves checkpoints every 125 updates, and pairs +raw/clipped WeightWatcher alpha with validation token error at those same updates. +It records an explicit target outcome. Its TPU performance and convergence require live measurement; the default three-hour budget is a cap, not a runtime prediction. For the **30-minute GPT-2/FineWeb reference run**, see diff --git a/baseline/gpt2_small/muon_speedrun/README.md b/baseline/gpt2_small/muon_speedrun/README.md index e4a21d4e..6264ad93 100644 --- a/baseline/gpt2_small/muon_speedrun/README.md +++ b/baseline/gpt2_small/muon_speedrun/README.md @@ -18,7 +18,11 @@ python3 baseline/gpt2_small/muon_speedrun/cloudshell.py start The launcher uses the existing `ww-gpt2-validation-48h-20261004-s1337-node` in `tpu-builders-504820/us-west4-a`. It creates no TPU, preserves the shared environment, and refuses concurrent training. Repeating start while active prints status. -It starts from scratch because the previous 30-minute run saved no model. +Every launch starts from the same seed and initialization, without resume. +To stop the current Muon speedrun and start fresh with tracking, use +`start --replace-current --optimizer muon --microbatch 64 --attention flash --hours 3`. +This stops only the service recorded in `MUON_SPEEDRUN_LATEST.json`; the TPU, +FineWeb cache, and prior files remain. The worker stops at the first full validation with **NLL <= 3.28**, after the published 3,000-update schedule, or before its deadline, whichever comes first. @@ -93,12 +97,55 @@ avoids synchronous network downloads at training shard boundaries. Existing FineWeb-Edu is a different corpus and is retained separately. All timings include end-to-end training overhead; published GPU training-only timing is labelled. -There are no WeightWatcher fits, gradient scans or per-matrix host reads in the -training loop. Scalar loss checks remain. An isolated flash check does not prove +WeightWatcher runs in a separate CPU process on immutable snapshots made from +the weights already transferred for checkpointing. There are no spectral fits, +gradient scans or additional per-matrix TPU reads in the training loop. Scalar +loss checks remain. An isolated flash check does not prove complete TPU optimizer/model parity. Source CUDA compilation, random seed, microbatch reduction order, rotary-buffer calculation and shard-boundary order can differ; these are recorded rather than represented as exact reproduction. +## Paired WeightWatcher and token-error tracking + +Tracking is enabled for new runs. The model, optimizers, hyperparameters, data, +seed, global batch, microbatch, flash attention, 3,000-update schedule, target, +evaluation/checkpoint cadence and three-hour cap remain unchanged. Additional +measurement work can add wall time within that cap. + +Every existing validation point (125 updates and final stop) now counts top-1 +prediction errors from the **same capped logits and the same benchmark tokens** +as validation NLL. `val_token_error = val_error_count / evaluation_tokens` is a +fraction, not a percentage. This is teacher-forced validation token error, not +free-generation accuracy or a separate test set. Partial evaluation is explicitly +flagged and cannot satisfy the target gate. + +The same saved weights are queued for WeightWatcher 0.7.7 on CPU: Q, K, V, O, +MLP_IN and MLP_OUT in each of the twelve blocks (72 matrices). Embeddings, +output vocabulary head and scalar parameters are excluded. Fits use `ERG=True`, +`randomize=True`, `fix_fingers="clip_xmax"`, `max_fingers=10`, `min_evals=20`. +`raw_alpha` is recorded as `alpha_raw`; `alpha` as `alpha_clip_xmax`. Missing or +failed fits stay unavailable; clipped alpha is never substituted for raw alpha. +The package's other scalar outputs, including available trap/finger counts, +are retained. Randomization happens only in the separate CPU process. + +- `tracking/layers.csv`: per-matrix raw/clipped alpha, fit status and paired + validation loss, perplexity and token error, keyed by run and exact update. +- `tracking/summary.csv`: mean/minimum alpha, valid-fit counts, counts below two, + and sample standard deviation **across matrices**, which is not a seed error bar. +- `tracking/measurements/.json`: immutable results with snapshot SHA256, + WeightWatcher version, diagnostic seed and fit options. +- `tracking/snapshots/.pt`: immutable CPU transformer weights and paired + validation metadata retained on the mounted disk for later analysis. These + spectral snapshots are not full training-state checkpoints and are not uploaded + by the final backup. Full latest/best/target checkpoints retain their existing + cloud backup behavior. +- `TRACKING_STATUS.json`: completed, pending and failed measurements. A fit-process + failure or deadline backlog is explicitly reported; snapshots remain recoverable. + +Small tracking tables and JSON results are uploaded first during the existing +final cloud backup. The CPU worker drains its queue within the existing time cap; +there is no extension of the TPU allocation or automatic training restart. + ## Checkpoints and results Runs live at `/mnt/disks/rg-data/gpt2small/muon-speedrun--`. diff --git a/baseline/gpt2_small/muon_speedrun/cloudshell.py b/baseline/gpt2_small/muon_speedrun/cloudshell.py index d53252a2..37b10ea0 100644 --- a/baseline/gpt2_small/muon_speedrun/cloudshell.py +++ b/baseline/gpt2_small/muon_speedrun/cloudshell.py @@ -11,6 +11,7 @@ import subprocess import sys import time +import uuid PROJECT = 'tpu-builders-504820' ZONE = 'us-west4-a' @@ -39,20 +40,39 @@ def status_remote(): print(json.dumps(record,indent=2),flush=True) subprocess.run(['systemctl','--no-pager','--full','status',record['unit']]) for name in ('RUN_STATUS.json','status.json','latest_validation.json','checkpoint_latest.json', - 'PALLAS_DEPENDENCIES.json','ATTENTION_CHECK.json','ATTENTION_CHECK_FAILURE.json'): + 'PALLAS_DEPENDENCIES.json','ATTENTION_CHECK.json','ATTENTION_CHECK_FAILURE.json', + 'TRACKING_CONFIG.json','TRACKING_STATUS.json'): if (root/name).exists(): print(name+'\n'+(root/name).read_text(),flush=True) subprocess.run(['tail','-n','15',str(root/'run.log')]) -def start_remote(commit, hours=3, optimizer='muon', microbatch=64, attention='flash'): +def stop_current(): + if not LATEST.exists(): + return + record = json.loads(LATEST.read_text()) + unit = record['unit'] + if not re.fullmatch(r'rg-muon-speedrun-\d{8}-\d{6}\.service', unit): + raise RuntimeError('Unexpected service name; refusing to stop it') + print('Stopping previous Muon speedrun: '+unit, flush=True) + run(['systemctl','stop',unit], timeout=60) + if active(unit): + raise RuntimeError('Previous service is still active; new run not started') + + +def start_remote(commit, hours=3, optimizer='muon', microbatch=64, attention='flash', replace_current=False, + launch_id=None): if os.geteuid() != 0 or not os.path.ismount('/mnt/disks/rg-data'): raise RuntimeError('Requires the existing mounted disk and root') if not re.fullmatch('[0-9a-f]{40}',commit): raise ValueError('Expected pinned commit') with (BASE/'port-check-launch.lock').open('a') as lock: fcntl.flock(lock,fcntl.LOCK_EX|fcntl.LOCK_NB) - if LATEST.exists() and active(json.loads(LATEST.read_text())['unit']): + if LATEST.exists() and launch_id and json.loads(LATEST.read_text()).get('launch_id') == launch_id: + print('This launch request was already handled; SSH retry will not restart it.') + status_remote() + return + if not replace_current and LATEST.exists() and active(json.loads(LATEST.read_text())['unit']): print('A speedrun is already active; no duplicate launched.') status_remote() return @@ -69,6 +89,8 @@ def start_remote(commit, hours=3, optimizer='muon', microbatch=64, attention='fl run(['git','-C',str(repo),'remote','add','origin','https://github.com/CalculatedContent/rg_optimizers.git']) run(['git','-C',str(repo),'fetch','--depth','1','origin',commit],timeout=120) run(['git','-C',str(repo),'checkout','--detach',commit]) + if replace_current: + stop_current() guard = repo/'baseline/gpt2_small/scripts/run_muonclip.py' spec = importlib.util.spec_from_file_location('training_guard',guard) module = importlib.util.module_from_spec(spec) @@ -81,8 +103,10 @@ def start_remote(commit, hours=3, optimizer='muon', microbatch=64, attention='fl 'OMP_NUM_THREADS':'4','OPENBLAS_NUM_THREADS':'4','MKL_NUM_THREADS':'4', 'TOKENIZERS_PARALLELISM':'false'} record = {'root':str(root),'unit':unit,'commit':commit,'optimizer':optimizer, + 'launch_id':launch_id, 'started_unix':time.time(),'deadline_unix':deadline,'hours_cap':hours, 'target_val_nll':3.28,'checkpoint_interval':125,'microbatch':microbatch, + 'weightwatcher_interval':125, 'validation_token_error':True, 'cloud_uri':'gs://tpu-builders-504820-ww-continuous8/gpt2small/'+root.name} (root/'launch.json').write_text(json.dumps(record,indent=2)) (root/'commit.txt').write_text(commit+'\n') @@ -102,6 +126,7 @@ def start_remote(commit, hours=3, optimizer='muon', microbatch=64, attention='fl print('Started '+optimizer+' recipe: '+unit,flush=True) print('Log: '+str(root/'run.log'),flush=True) print('3,000 updates, stops at full-validation NLL <= 3.28; checkpoints every 125.',flush=True) + print('Paired validation token error and raw/clipped WeightWatcher alpha every 125 updates and final.',flush=True) print('Hard cutoff UTC: '+dt.datetime.fromtimestamp(deadline,dt.timezone.utc).isoformat(),flush=True) print('No automatic restart or new TPU allocation. Three hours is a cap, not an ETA.',flush=True) @@ -113,26 +138,30 @@ def main(): p.add_argument('--optimizer',choices=('muon','adam'),default='muon') p.add_argument('--microbatch',type=int,choices=(32,64,128),default=64) p.add_argument('--attention',choices=('auto','flash','math'),default='flash') + p.add_argument('--replace-current',action='store_true',help='Stop the previous speedrun and start from initialization') p.add_argument('--on-tpu',action='store_true',help=argparse.SUPPRESS) p.add_argument('--commit',help=argparse.SUPPRESS) + p.add_argument('--launch-id',help=argparse.SUPPRESS) a = p.parse_args() if not 1 <= a.hours <= 12: raise ValueError('Hours must be between 1 and 12, bounded by existing allocation') if a.on_tpu: if a.action == 'start': - start_remote(a.commit,a.hours,a.optimizer,a.microbatch,a.attention) + start_remote(a.commit,a.hours,a.optimizer,a.microbatch,a.attention,a.replace_current,a.launch_id) else: status_remote() return 0 command = ['sudo','python3','-c',Path(__file__).read_text(),a.action,'--on-tpu', '--hours',str(a.hours),'--optimizer',a.optimizer,'--microbatch',str(a.microbatch), '--attention',a.attention] + if a.replace_current: + command += ['--replace-current'] if a.action == 'start': repo = Path(__file__).resolve().parents[3] if run(['git','-C',str(repo),'status','--porcelain'],capture_output=True).stdout.strip(): raise RuntimeError('Launch from a clean checkout of the pushed commit') commit = run(['git','-C',str(repo),'rev-parse','HEAD'],capture_output=True).stdout.strip() - command += ['--commit',commit] + command += ['--commit',commit,'--launch-id',uuid.uuid4().hex] return subprocess.run(['gcloud','compute','tpus','tpu-vm','ssh',NODE, '--project='+PROJECT,'--zone='+ZONE,'--worker=0','--command='+shlex.join(command)]).returncode diff --git a/baseline/gpt2_small/muon_speedrun/model.py b/baseline/gpt2_small/muon_speedrun/model.py index e20e16bd..b639dc6a 100644 --- a/baseline/gpt2_small/muon_speedrun/model.py +++ b/baseline/gpt2_small/muon_speedrun/model.py @@ -132,7 +132,7 @@ def __init__(self, config): self.lm_head = CastedLinear(config.n_embd, config.vocab_size, bias=False) self.lm_head.weight.data.zero_() # @Grad62304977 - def forward(self, idx, target): + def forward(self, idx, target, *, return_token_errors=False): # forward the GPT model itself x = self.transformer.wte(idx) # token embeddings of shape (b, t, n_embd) @@ -160,4 +160,7 @@ def forward(self, idx, target): logits = 30 * torch.tanh(logits / 30) # @Grad62304977 logits = logits.float() loss = F.cross_entropy(logits.view(-1, logits.size(-1)), target.view(-1)) + if return_token_errors: + # Evaluation only; the default training graph and loss are unchanged. + return loss.float(), (logits.argmax(dim=-1) != target).sum(dtype=torch.int32) return loss.float() diff --git a/baseline/gpt2_small/muon_speedrun/run.py b/baseline/gpt2_small/muon_speedrun/run.py index a21e8bda..c8cae6d6 100644 --- a/baseline/gpt2_small/muon_speedrun/run.py +++ b/baseline/gpt2_small/muon_speedrun/run.py @@ -16,6 +16,7 @@ from data import FineWeb, TrainStream, write_json, prepare from optim import make_optimizers, apply_update, schedule, momentum from runtime import Runtime, attention_check +from tracking import queue_snapshot TOTAL_STEPS = 3000 BATCH_TOKENS = 524288 @@ -74,6 +75,8 @@ def save_checkpoint(root, model, muon, adam, stream, step, manifest, rt, validat temp.replace(root/name) write_json(root/'checkpoint_latest.json', {'file':checkpoint.name, 'step':step, 'target_met':passed, 'validation':validation}) + if validation is not None: + queue_snapshot(root, payload) print('Checkpoint saved at update '+str(step), flush=True) @@ -89,22 +92,28 @@ def evaluate(model, tokens, rt, root, step, deadline, microbatch, started): model.eval() size = microbatch*1024 total = torch.zeros((), device=rt.device, dtype=torch.float32) + errors = torch.zeros((), device=rt.device, dtype=torch.int32) evaluated = 0 for offset in range(0, VAL_TOKENS, size): if time.time() >= deadline: break buf = torch.from_numpy(np.array(tokens[offset:offset+size+1], dtype=np.int64)) - total += model(rt.put(buf[:-1].reshape(microbatch, 1024)), - rt.put(buf[1:].reshape(microbatch, 1024))).detach().float() + batch_loss, batch_errors = model(rt.put(buf[:-1].reshape(microbatch, 1024)), + rt.put(buf[1:].reshape(microbatch, 1024)), return_token_errors=True) + total += batch_loss.detach().float() + errors += batch_errors rt.step() evaluated += size rt.step(wait=True) loss = float(total.cpu())/(evaluated/size) if evaluated else None + error_count = int(errors.cpu()) if loss is not None and not math.isfinite(loss): raise RuntimeError('Nonfinite validation NLL') row = {'kind':'validation', 'step':step, 'tokens_seen':step*BATCH_TOKENS, 'evaluation_tokens':evaluated, 'full_benchmark_evaluation':evaluated == VAL_TOKENS, 'val_nll':loss, 'val_perplexity':math.exp(loss) if loss is not None else None, + 'val_error_count':error_count, 'val_token_error':error_count/evaluated if evaluated else None, + 'val_accuracy':1-error_count/evaluated if evaluated else None, 'elapsed_seconds':time.time()-started, 'recorded_unix':time.time()} reference = json.loads((HERE/'reference_val.json').read_text()) row['published_at_same_step'] = next((r for r in reference if r['step'] == step), None) @@ -149,6 +158,10 @@ def train(a): 'data_repo':source.manifest['repo'], 'data_revision':source.manifest['revision'], 'record_source_sha256':hashlib.sha256((HERE/'vendor/record_source.py').read_bytes()).hexdigest(), 'torch_version':torch.__version__, 'automatic_restart':False, + 'tracking':{'interval_updates':125, 'extra_final_measurement':True, + 'matrices':'Q,K,V,O,MLP_IN,MLP_OUT in all 12 blocks', + 'execution':'separate CPU process on immutable snapshots', + 'token_error':'teacher-forced top-1 error on the same benchmark validation tokens'}, 'differences':['TPU SPMD instead of CUDA DDP', 'batched matrix-partitioned Muon', 'microbatch accumulation and shard-boundary ordering', 'fixed seed 1337; original record did not pin a seed', diff --git a/baseline/gpt2_small/muon_speedrun/tracking.py b/baseline/gpt2_small/muon_speedrun/tracking.py new file mode 100644 index 00000000..0f47da55 --- /dev/null +++ b/baseline/gpt2_small/muon_speedrun/tracking.py @@ -0,0 +1,247 @@ +"""CPU-only WeightWatcher measurements paired with immutable validation snapshots. + +Only this separate process imports WeightWatcher. No training RNG, model, optimizer, +TPU graph, data iterator, or learning-rate state is touched by spectral analysis. +""" +import argparse +import csv +import hashlib +import importlib.metadata +import json +import math +import os +from pathlib import Path +import time + +ROLES = {'attn.c_q':'Q', 'attn.c_k':'K', 'attn.c_v':'V', 'attn.c_proj':'O', + 'mlp.c_fc':'MLP_IN', 'mlp.c_proj':'MLP_OUT'} +WW_OPTIONS = dict(ERG=True, randomize=True, plot=False, fix_fingers='clip_xmax', + max_fingers=10, min_evals=20) + + +def write_json(path, value): + path.parent.mkdir(parents=True, exist_ok=True) + temp = path.with_suffix(path.suffix+'.tmp') + temp.write_text(json.dumps(value, indent=2, allow_nan=False)+'\n') + temp.replace(path) + + +def queue_snapshot(root, checkpoint): + """Reuse the checkpoint's CPU weights, never fetch live TPU tensors again.""" + import torch + step = checkpoint['step'] + validation = checkpoint['validation'] + if validation.get('step', step) != step: + raise ValueError('Spectrum and validation must describe the same update') + matrices = {} + for block in range(checkpoint['config']['n_layer']): + for suffix, role in ROLES.items(): + source = f'transformer.h.{block}.{suffix}.weight' + matrices[f'L{block:02d}_W_{role}'] = checkpoint['model'][source] + folder = root/'tracking/snapshots' + folder.mkdir(parents=True, exist_ok=True) + path = folder/f'{step:07d}.pt' + if path.exists(): + raise RuntimeError('Refusing to overwrite an immutable spectral snapshot: '+str(path)) + payload = {'step':step, 'tokens_seen':checkpoint['tokens_seen'], 'run_id':root.name, + 'validation':validation, 'matrices':matrices, + 'manifest':checkpoint['manifest']} + temporary = path.with_suffix('.tmp') + with temporary.open('wb') as f: + torch.save(payload, f) + f.flush() + os.fsync(f.fileno()) + temporary.replace(path) + print('[weightwatcher] queued CPU snapshot at update '+str(step), flush=True) + + +def holder_from(matrices): + import torch + holder = torch.nn.Module() + for name, value in matrices.items(): + # No random initialization or allocation of a second model. + layer = torch.nn.Linear(value.shape[1], value.shape[0], bias=False, device='meta') + layer.weight = torch.nn.Parameter(value.detach().float().cpu(), requires_grad=False) + holder.add_module(name, layer) + return holder + + +def scalar(value): + if hasattr(value, 'item'): + value = value.item() + if isinstance(value, float) and not math.isfinite(value): + return None + return value if isinstance(value, (str, bool, int, float, type(None))) else str(value) + + +def positive(value): + try: + x = float(value) + return x if math.isfinite(x) and x > 0 else None + except (TypeError, ValueError): + return None + + +def normalize_rows(frame, names, identity): + """Match by the explicit layer name; never substitute clipped alpha for raw.""" + found = {} + for source in frame.to_dict('records'): + matches = [name for name in names if name in + str(source.get('longname', '')) or name == str(source.get('name', ''))] + if len(matches) != 1 or matches[0] in found: + raise RuntimeError('WeightWatcher returned ambiguous or duplicate matrix names') + name = matches[0] + raw = {key:scalar(value) for key, value in source.items()} + good = source.get('status') == 'success' + alpha_raw = positive(source.get('raw_alpha')) if good else None + alpha_clip = positive(source.get('alpha')) if good else None + found[name] = {**raw, **identity, 'matrix_name':name, 'block':int(name[1:3]), + 'matrix_type':name.split('_W_', 1)[1], + 'alpha_raw':alpha_raw, 'alpha_clip_xmax':alpha_clip, + 'raw_fit_status':'success' if alpha_raw is not None else 'unavailable', + 'clipped_fit_status':'success' if alpha_clip is not None else 'unavailable'} + # Some versions skip an unfit/zero matrix. Keep explicit missing rows. + for name in names: + if name not in found: + found[name] = {**identity, 'matrix_name':name, 'block':int(name[1:3]), + 'matrix_type':name.split('_W_', 1)[1], 'status':'not_returned', + 'alpha_raw':None, 'alpha_clip_xmax':None, + 'raw_fit_status':'unavailable', 'clipped_fit_status':'unavailable'} + return [found[name] for name in names] + + +def summary(rows, identity): + import statistics + result = {**identity, 'matrix_count':len(rows)} + for field in ('alpha_raw', 'alpha_clip_xmax'): + values = [row[field] for row in rows if row[field] is not None] + result.update({field+'_valid_count':len(values), + field+'_mean':statistics.mean(values) if values else None, + field+'_min':min(values) if values else None, + field+'_std_across_matrices':statistics.stdev(values) if len(values) > 1 else None, + field+'_below_two':sum(x < 2 for x in values)}) + return result + + +def measure(path): + import random + import numpy as np + import torch + import weightwatcher as ww + torch.set_num_threads(1) + payload = torch.load(path, map_location='cpu', weights_only=False) + seed = 1_001_340+payload['step'] + random.seed(seed); np.random.seed(seed); torch.manual_seed(seed) + with path.open('rb') as f: + digest = hashlib.file_digest(f, 'sha256').hexdigest() if hasattr(hashlib, 'file_digest') else None + if digest is None: # TPU host runs Python 3.10. + h = hashlib.sha256() + with path.open('rb') as f: + for chunk in iter(lambda:f.read(8*1024*1024), b''): + h.update(chunk) + digest = h.hexdigest() + validation = payload['validation'] + identity = {key:validation.get(key) for key in ('evaluation_tokens', 'full_benchmark_evaluation', + 'val_nll', 'val_perplexity', 'val_token_error', 'val_accuracy', 'val_error_count')} + identity.update(step=payload['step'], tokens_seen=payload['tokens_seen'], run_id=payload['run_id'], + snapshot_sha256=digest, diagnostic_seed=seed, + weightwatcher_version=importlib.metadata.version('weightwatcher')) + frame = ww.WeightWatcher(model=holder_from(payload['matrices'])).analyze(**WW_OPTIONS) + if not {'alpha', 'raw_alpha'}.issubset(frame.columns): + raise RuntimeError('WeightWatcher must expose both raw_alpha and clipped alpha') + rows = normalize_rows(frame, list(payload['matrices']), identity) + return {'layers':rows, 'summary':summary(rows, identity), 'options':WW_OPTIONS} + + +def csv_write(path, rows): + temp = path.with_suffix('.csv.tmp') + fields = list(dict.fromkeys(key for row in rows for key in row)) + with temp.open('w', newline='') as f: + writer = csv.DictWriter(f, fieldnames=fields) + writer.writeheader(); writer.writerows(rows) + temp.replace(path) + + +def collect(root): + folder = root/'tracking' + files = sorted((folder/'measurements').glob('*.json')) + data = [json.loads(path.read_text()) for path in files] + csv_write(folder/'layers.csv', [row for item in data for row in item['layers']]) + csv_write(folder/'summary.csv', [item['summary'] for item in data]) + return len(files) + + +def tracking_status(root, status, **extra): + folder = root/'tracking' + snapshots = list((folder/'snapshots').glob('*.pt')) + done = {p.stem for p in (folder/'measurements').glob('*.json')} + failed = {p.stem for p in (folder/'failures').glob('*.json')} + result = dict(status=status, snapshots=len(snapshots), completed=len(done), failed=len(failed), + pending=sum(p.stem not in done|failed for p in snapshots), updated_unix=time.time(), **extra) + write_json(root/'TRACKING_STATUS.json', result) + return result + + +def watch(root, deadline): + folder = root/'tracking' + while time.time() < deadline: + for path in sorted((folder/'snapshots').glob('*.pt')): + output = folder/'measurements'/(path.stem+'.json') + failure = folder/'failures'/(path.stem+'.json') + if output.exists() or failure.exists(): + continue + if time.time() >= deadline: + break + tracking_status(root, 'measuring', current_step=int(path.stem)) + try: + result = measure(path) + write_json(output, result) + collect(root) + print('[weightwatcher] '+json.dumps(result['summary']), flush=True) + except Exception as exc: + write_json(failure, {'step':int(path.stem), 'error':repr(exc)}) + print('[weightwatcher] FAILED update '+path.stem+': '+repr(exc), flush=True) + state = tracking_status(root, 'waiting') + if (folder/'TRAINING_DONE').exists() and not state['pending']: + tracking_status(root, 'complete' if not state['failed'] else 'completed_with_errors') + return 0 if not state['failed'] else 1 + time.sleep(1) + tracking_status(root, 'deadline_reached') + return 1 + + +def check(root): + import torch + import weightwatcher as ww + torch.set_num_threads(1) + version = importlib.metadata.version('weightwatcher') + if version != '0.7.7': + raise RuntimeError('Expected the existing weightwatcher==0.7.7; got '+version) + holder = holder_from({'L00_W_Q':torch.randn(64, 64)}) + frame = ww.WeightWatcher(model=holder).analyze(**WW_OPTIONS) + if not {'raw_alpha', 'alpha'}.issubset(frame.columns): + raise RuntimeError('WeightWatcher raw/clipped fields unavailable') + normalize_rows(frame, ['L00_W_Q'], {}) + write_json(root/'TRACKING_CONFIG.json', {'weightwatcher_version':version, 'options':WW_OPTIONS, + 'interval_updates':125, 'extra_final_measurement':True, 'matrix_roles':list(ROLES.values()), + 'matrix_count':72, 'execution':'separate CPU process', + 'token_error_units':'fraction of the same teacher-forced benchmark validation tokens', + 'snapshot_storage':'persistent disk; measurement tables and current full checkpoints backed up to cloud', + 'std_definition':'sample standard deviation across matrices; not uncertainty across seeds'}) + print('WeightWatcher CPU tracking ready: raw/clipped alpha and paired validation token error.', flush=True) + + +def main(): + p = argparse.ArgumentParser() + p.add_argument('action', choices=('check', 'watch')) + p.add_argument('--root', type=Path, required=True) + p.add_argument('--deadline', type=float, default=float('inf')) + args = p.parse_args() + if args.action == 'check': + check(args.root) + return 0 + return watch(args.root, args.deadline) + + +if __name__ == '__main__': + raise SystemExit(main()) diff --git a/baseline/gpt2_small/muon_speedrun/worker.py b/baseline/gpt2_small/muon_speedrun/worker.py index b85a75f8..88dd24e6 100644 --- a/baseline/gpt2_small/muon_speedrun/worker.py +++ b/baseline/gpt2_small/muon_speedrun/worker.py @@ -46,6 +46,11 @@ def backup(root): from rg_nanogpt_one_head.continuous_support import CloudPublisher publisher = CloudPublisher('gs://tpu-builders-504820-ww-continuous8/gpt2small/'+root.name) receipts = [] + # Upload small scientific tables first. Retain immutable spectral weights on /mnt; + # current full-state checkpoints keep their existing cloud backup behavior. + for path in sorted((root/'tracking').rglob('*')): + if path.is_file() and path.suffix in ('.json', '.csv'): + receipts.append(publisher.file(path, path.relative_to(root).as_posix())) for path in sorted(root.iterdir()): if path.is_file() and path.suffix in ('.json', '.jsonl', '.pt', '.txt', '.log'): # Logs may still grow as supervisor/upload output is appended. @@ -58,6 +63,39 @@ def backup(root): print('Cloud backup verified: '+root.name, flush=True) +def start_tracking(root, deadline): + # CPU-only child; never import WeightWatcher in the training process. + env = {**os.environ, 'PJRT_DEVICE':'CPU', 'CUDA_VISIBLE_DEVICES':'', + 'OMP_NUM_THREADS':'1', 'OPENBLAS_NUM_THREADS':'1', 'MKL_NUM_THREADS':'1'} + command = [sys.executable, '-u', str(Path(__file__).with_name('tracking.py'))] + subprocess.run(command+['check','--root',str(root)], env=env, check=True, + timeout=min(120, max(1, deadline-time.time()))) + return subprocess.Popen(command+['watch','--root',str(root),'--deadline',str(deadline)], + env=env, start_new_session=True) + + +def finish_tracking(child, root, deadline): + from tracking import tracking_status + folder = root/'tracking' + folder.mkdir(exist_ok=True) + (folder/'TRAINING_DONE').touch() + try: + while child.poll() is None and time.time() < deadline: + try: + child.wait(timeout=min(30, max(.1, deadline-time.time()))) + except subprocess.TimeoutExpired: + print('WAIT CPU WeightWatcher: draining saved snapshots', flush=True) + if child.poll() is None: + os.killpg(child.pid, signal.SIGKILL) + child.wait(timeout=10) + finally: + state = tracking_status(root, 'complete' if child.returncode == 0 else 'incomplete', + exit_code=child.returncode) + if state['pending'] or state['failed']: + state = tracking_status(root, 'incomplete', exit_code=child.returncode) + return state + + def main(): p = argparse.ArgumentParser() p.add_argument('root', type=Path) @@ -77,6 +115,7 @@ def main(): common = [sys.executable, '-u', str(here/'run.py')] args = ['--root',str(a.root),'--microbatch',str(a.microbatch)] train_deadline = a.deadline-600 + tracker = None try: if a.attention == 'math' and a.microbatch > 64: raise RuntimeError('Mathematical attention with microbatch 128 exceeded this TPU memory; use <=64') @@ -102,6 +141,7 @@ def main(): else: write(a.root, 'ATTENTION_CHECK_FAILURE.json', checked) raise RuntimeError('TPU flash attention failed validation; no automatic mathematical-attention fallback') + tracker = start_tracking(a.root, train_deadline) run.update(status='training', attention=attention) write(a.root, 'RUN_STATUS.json', run) result = bounded(common+['train',*args,'--deadline',str(train_deadline), @@ -118,6 +158,9 @@ def main(): run.update(status=status['status'], target_met=status['target_met'], step=status['step']) except Exception as exc: run.update(status='failed', error=str(exc)) + finally: + if tracker is not None: + run['tracking'] = finish_tracking(tracker, a.root, train_deadline) if run['status'] in ('failed','failed_or_timed_out'): previous = {} if (a.root/'status.json').exists(): @@ -130,7 +173,8 @@ def main(): run['backup'] = backed write(a.root, 'RUN_STATUS.json', run) print(json.dumps(run), flush=True) - return 0 if run['status'] in ('target_reached','schedule_complete_target_not_met') else 1 + return 0 if (run['status'] in ('target_reached','schedule_complete_target_not_met') + and run.get('tracking', {}).get('status') == 'complete') else 1 if __name__ == '__main__': diff --git a/baseline/gpt2_small/tests/test_muon_speedrun.py b/baseline/gpt2_small/tests/test_muon_speedrun.py index 7f7470c6..f467087b 100644 --- a/baseline/gpt2_small/tests/test_muon_speedrun.py +++ b/baseline/gpt2_small/tests/test_muon_speedrun.py @@ -210,9 +210,133 @@ def bounded(command,seconds,root,label,watch=False): (root/'status.json').write_text(json.dumps(dict(status='target_reached',target_met=True,step=2875))) return {'exit_code':0,'timed_out':False,'phase':label} monkeypatch.setattr(module,'bounded',bounded) + monkeypatch.setattr(module,'start_tracking',lambda *args:object()) + monkeypatch.setattr(module,'finish_tracking',lambda *args:{'status':'complete'}) monkeypatch.setenv('PYTHONPATH','test-original') assert module.main()==0 training=next(command for label,command in phases if '3,000-update' in label) assert training[training.index('--microbatch')+1]=='64' assert training[training.index('--attention')+1]=='flash' assert phases[0][0]=='pinned Pallas dependencies' + + +def test_token_error_uses_the_same_logits_without_changing_loss_or_gradients(): + model = small(port) + with torch.no_grad(): + model.lm_head.weight.normal_(std=.02) + x = torch.randint(128, (2, 8)); y = torch.randint(128, (2, 8)) + captured = [] + hook = model.lm_head.register_forward_hook(lambda module, args, out:captured.append(out.detach())) + loss = model(x, y) + loss.backward() + gradients = [p.grad.clone() for p in model.parameters()] + model.zero_grad(set_to_none=True) + measured, errors = model(x, y, return_token_errors=True) + measured.backward() + hook.remove() + torch.testing.assert_close(loss, measured, rtol=0, atol=0) + for old, p in zip(gradients, model.parameters()): + torch.testing.assert_close(old, p.grad, rtol=0, atol=0) + logits = (30 * torch.tanh(captured[-1] / 30)).float() + assert int(errors) == int((logits.argmax(-1) != y).sum()) + + +def test_evaluation_pairs_exact_token_count_with_unchanged_nll(monkeypatch,tmp_path): + import numpy as np + import time + run=runner(); model=small(port); rt=Runtime('cpu') + monkeypatch.setattr(run,'VAL_TOKENS',2048) + tokens=np.arange(2049,dtype=np.int64)%128 + # Zero-initialized head predicts token zero everywhere. 1/128 targets are zero. + row=run.evaluate(model,tokens,rt,tmp_path,125,time.time()+60,1,time.time()) + assert row['evaluation_tokens']==2048 + assert row['val_error_count']==2032 + assert row['val_token_error']==2032/2048 + assert row['val_accuracy']==16/2048 + assert row['val_nll']==pytest.approx(float(torch.tensor(128.).log()),abs=1e-6) + assert model.training + + +def test_spectral_snapshot_is_immutable_and_paired_with_validation(tmp_path): + import tracking + rt=Runtime('cpu'); model=small(port); muon,adam=make_optimizers(model,rt) + stream=types.SimpleNamespace(shard=1,position=0) + v=dict(step=125,full_benchmark_evaluation=True,evaluation_tokens=10485760, + val_nll=3.5,val_token_error=.6) + state=torch.get_rng_state().clone() + runner().save_checkpoint(tmp_path,model,muon,adam,stream,125,{},rt,v,4.) + assert torch.equal(state,torch.get_rng_state()) + path=tmp_path/'tracking/snapshots/0000125.pt' + before=path.read_bytes() + with torch.no_grad(): + model.transformer.h[0].attn.c_q.weight.add_(1) + runner().save_checkpoint(tmp_path,model,muon,adam,stream,250,{},rt,{**v,'step':250},4.) + assert path.read_bytes()==before + payload=torch.load(path,weights_only=False) + assert payload['validation']==v and payload['step']==125 + assert len(payload['matrices'])==12 + assert set(payload['matrices'])=={f'L{i:02d}_W_{role}' for i in range(2) for role in tracking.ROLES.values()} + + +def test_weightwatcher_raw_alpha_never_falls_back_to_clipped(): + import tracking + frame=types.SimpleNamespace(to_dict=lambda orient:[ + dict(longname='L00_W_Q',status='success',alpha=1.9,raw_alpha=float('nan')), + dict(longname='L00_W_K',status='failed',alpha=1.8,raw_alpha=1.7)]) + rows=tracking.normalize_rows(frame,['L00_W_Q','L00_W_K','L00_W_V'],{'step':125,'val_token_error':.6}) + assert len(rows)==3 and all(r['alpha_raw'] is None for r in rows) + assert rows[0]['alpha_clip_xmax']==1.9 + assert rows[1]['alpha_clip_xmax'] is None + assert rows[2]['status']=='not_returned' + s=tracking.summary(rows,{'step':125}) + assert s['alpha_raw_valid_count']==0 and s['alpha_raw_mean'] is None + assert s['alpha_clip_xmax_valid_count']==1 + + +def test_replace_stops_only_the_recorded_muon_service(tmp_path,monkeypatch): + spec=importlib.util.spec_from_file_location('muon_restart_test',BASE/'cloudshell.py') + module=importlib.util.module_from_spec(spec); spec.loader.exec_module(module) + pointer=tmp_path/'pointer.json' + unit='rg-muon-speedrun-20261005-020737.service' + pointer.write_text(json.dumps({'unit':unit})) + monkeypatch.setattr(module,'LATEST',pointer) + monkeypatch.setattr(module,'active',lambda name:False) + calls=[]; monkeypatch.setattr(module,'run',lambda command,**kwargs:calls.append(command)) + module.stop_current() + assert calls==[['systemctl','stop',unit]] + pointer.write_text(json.dumps({'unit':'unrelated.service'})) + with pytest.raises(RuntimeError,match='Unexpected service'): + module.stop_current() + assert len(calls)==1 + + +def test_real_weightwatcher_pairs_snapshot_and_writes_tables(tmp_path): + pytest.importorskip('weightwatcher') + import tracking, time + validation=dict(step=125,val_nll=4.5,val_token_error=.8,evaluation_tokens=10485760, + full_benchmark_evaluation=True) + matrices={f'transformer.h.0.{suffix}.weight':torch.randn(64,64) for suffix in tracking.ROLES} + tracking.queue_snapshot(tmp_path,dict(step=125,tokens_seen=65536000,config={'n_layer':1}, + model=matrices,validation=validation,manifest={})) + (tmp_path/'tracking/TRAINING_DONE').touch() + assert tracking.watch(tmp_path,time.time()+60)==0 + result=json.loads((tmp_path/'tracking/measurements/0000125.json').read_text()) + assert len(result['layers'])==6 + assert all(row['step']==125 and row['val_token_error']==.8 for row in result['layers']) + assert result['summary']['alpha_raw_valid_count']==6 + assert (tmp_path/'tracking/layers.csv').is_file() + assert (tmp_path/'tracking/summary.csv').is_file() + assert json.loads((tmp_path/'TRACKING_STATUS.json').read_text())['status']=='complete' + + +def test_ssh_retry_cannot_replace_the_run_it_just_launched(tmp_path,monkeypatch): + spec=importlib.util.spec_from_file_location('muon_idempotent_restart',BASE/'cloudshell.py') + module=importlib.util.module_from_spec(spec); spec.loader.exec_module(module) + pointer=tmp_path/'pointer.json'; pointer.write_text(json.dumps({'launch_id':'same-request'})) + monkeypatch.setattr(module,'BASE',tmp_path); monkeypatch.setattr(module,'LATEST',pointer) + monkeypatch.setattr(module.os,'geteuid',lambda:0) + monkeypatch.setattr(module.os.path,'ismount',lambda p:True) + seen=[]; monkeypatch.setattr(module,'status_remote',lambda:seen.append(True)) + monkeypatch.setattr(module,'run',lambda *a,**k:pytest.fail('retry must not stop or launch')) + module.start_remote('a'*40,replace_current=True,launch_id='same-request') + assert seen==[True] From 2c66b6c5e79d2e837b313fd9d6b813663c65dec2 Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Sun, 4 Oct 2026 22:57:53 -0700 Subject: [PATCH 29/36] Add fresh 25k Muon trajectory with live lease gate, paired spectra and full-state replay --- baseline/gpt2_small/README.md | 6 + baseline/gpt2_small/muon_longrun/README.md | 174 ++++++++++++ .../gpt2_small/muon_longrun/checkpoint.py | 81 ++++++ baseline/gpt2_small/muon_longrun/common.py | 111 ++++++++ baseline/gpt2_small/muon_longrun/launch.py | 191 +++++++++++++ baseline/gpt2_small/muon_longrun/long_data.py | 56 ++++ .../gpt2_small/muon_longrun/long_worker.py | 81 ++++++ .../gpt2_small/muon_longrun/track_long.py | 92 +++++++ .../gpt2_small/muon_longrun/train_long.py | 254 ++++++++++++++++++ .../gpt2_small/tests/test_muon_longrun.py | 173 ++++++++++++ 10 files changed, 1219 insertions(+) create mode 100644 baseline/gpt2_small/muon_longrun/README.md create mode 100644 baseline/gpt2_small/muon_longrun/checkpoint.py create mode 100644 baseline/gpt2_small/muon_longrun/common.py create mode 100644 baseline/gpt2_small/muon_longrun/launch.py create mode 100644 baseline/gpt2_small/muon_longrun/long_data.py create mode 100644 baseline/gpt2_small/muon_longrun/long_worker.py create mode 100644 baseline/gpt2_small/muon_longrun/track_long.py create mode 100644 baseline/gpt2_small/muon_longrun/train_long.py create mode 100644 baseline/gpt2_small/tests/test_muon_longrun.py diff --git a/baseline/gpt2_small/README.md b/baseline/gpt2_small/README.md index 408c92e3..085333e7 100644 --- a/baseline/gpt2_small/README.md +++ b/baseline/gpt2_small/README.md @@ -1,5 +1,11 @@ # GPT-2 Small / FineWeb-Edu validation +For the **fresh 25,000-update Muon experiment with the validated speedrun model**, +see [muon_longrun/README.md](muon_longrun/README.md). It preserves the successful +six-head Muon/FineWeb recipe, stretches cooldown to updates 17,500–25,000, +and records full-validation token error against all 72 raw/clipped alpha trajectories. +The launcher checks the live TPU lease and retains the completed reference run. + For the **published 3,000-update Muon speedrun recipe on eight TPU chips**, see [muon_speedrun/README.md](muon_speedrun/README.md). It uses a modified transformer, targets full-validation loss 3.28, saves checkpoints every 125 updates, and pairs diff --git a/baseline/gpt2_small/muon_longrun/README.md b/baseline/gpt2_small/muon_longrun/README.md new file mode 100644 index 00000000..152df511 --- /dev/null +++ b/baseline/gpt2_small/muon_longrun/README.md @@ -0,0 +1,174 @@ +# Fresh 25,000-update Muon trajectory + +This extends the successful `muon-speedrun-muon-20261005-030026` recipe from a +new seed-1337 initialization. It never loads that speedrun's trained checkpoint. +The original run directory remains the read-only reference. The experiment +does **not** stop when validation NLL reaches 3.28. + +## Exact reference and intentional changes + +The reference is the modified `2024-11-10_UNetDoubleLr` transformer, **not stock +12-head GPT-2 and not MuonClip**: 162,201,642 parameters, 12 blocks, 6 heads, +width 768, context 1,024, vocab 50,304. QK normalization is present; QK clipping +and gradient clipping are absent. The benchmark dataset is pinned GPT-2-tokenized +FineWeb (`kjj0/fineweb10B-gpt2`, revision +`889765ea1f903759787add96995d81171b632d0c`), not FineWeb-Edu. + +The launcher compares the actual reference manifest and hashes of the model, +optimizer, runtime, data implementation/manifest and Pallas installer. Those +files are imported unchanged. Numerical settings remain: + +| Setting | Value | +|---|---| +| Hardware | One v5litepod-8 host, eight-chip SPMD | +| Attention | Verified TPU flash; no math fallback | +| Global microbatch | 64 sequences, eight per chip | +| Gradient accumulation | 8 microbatches | +| Effective batch | 524,288 tokens/update | +| Precision | BF16 activations, embedding/scalars; FP32 linear weights | +| Muon peak LR | 0.04 | +| Muon momentum | 0.85 to 0.95 over first 500 updates, then 0.95 | +| Newton–Schulz | 5 iterations; original coefficients and BF16 operation order | +| Auxiliary Adam peak LRs | Embedding 0.6; head 0.008; scalars 0.04 | +| Auxiliary Adam | betas (0.9, 0.95), eps 1e-8; foreach/fused false; TPU capturable | +| Weight decay / clipping | 0 / none | +| Initialization seed | 1337 | + +Only the run horizon/schedule, measurement cadence, gradient-norm logging, +full-state recovery, and orchestration change. The no-warmup scheduler is +`min(1, max(0, (25000 - update_index) / 7500))`, where `update_index` is zero +based. Completed update 17,500 is immediately before cooldown; the next update +uses index 17,500. The last applied factor is `1/7500`; the next is zero. +The separate 500-update momentum ramp is unchanged. At step 3,000 the LR +factor remains **1**, whereas the short reference had finished cooldown. + +The old full validation at step 3,000 was NLL **3.28082160949707**, perplexity +**26.5976**. That narrowly missed a strict `<=3.28` threshold. The new step-3,000 +result is compared with it in `COMPARISON_3000.json`; the schedules intentionally +differ after index 2,100, so identical validation loss is not expected. + +## Corpus and duration + +25,000 updates process **13,107,200,000 token presentations**. All 103 pinned +training shards contain **10,255,324,043 tokens**, slightly fewer usable after +the unchanged per-shard batch truncation. The stream makes one full sequential +pass and repeats about 28%; this is not 13.1B distinct tokens. Epoch is +`tokens_seen / usable_tokens_per_full_corpus_pass`, recorded with its denominator. + +The existing benchmark cache is reused. Missing pinned shards are downloaded +before training; the older FineWeb/Edu corpus is never deleted or reformatted. +At the measured reference rate, expect roughly **10–11 hours**, plus any unusual +setup/measurement overhead. This is an estimate, not a completion guarantee. +The service has a 12-hour cap, with the final 30 minutes reserved for tracker +drain/backup; the trainer saves at its earlier deadline. + +## Start, inspect, stop + +From a clean checkout on your authenticated **Mac terminal or Cloud Shell**: + +```bash +python3 baseline/gpt2_small/muon_longrun/launch.py start +python3 baseline/gpt2_small/muon_longrun/launch.py status +python3 baseline/gpt2_small/muon_longrun/launch.py metrics +python3 baseline/gpt2_small/muon_longrun/launch.py stop +``` + +`start` describes the live node and its linked queued resource and prints queue +creation, node creation, maximum duration, explicit termination timestamp and +remaining hours. It requires **more than 12.5 hours**, an already-mounted data +disk with 45 GiB free, a healthy reference manifest, and an idle TPU. It never +creates/deletes an allocation or stops another running experiment. It refuses +to guess expiration from queue submission time. SSH retry is idempotent. + +The launcher prints the commit, directory, service, training deadline and cloud +prefix. The record is `/mnt/disks/rg-data/gpt2small/MUON_LONG25K_LATEST.json`. +Output is under `/mnt/disks/rg-data/gpt2small/muon-long25k-s1337-`. +`status` reports the systemd PID, latest metrics, checkpoint, and tracker state; +`metrics` also tails the scalar file. `stop` requests a checkpoint after the +current update, then drains tracking and performs backup. There is no automatic +restart. Do not delete the allocation before the stop/backup finishes. + +## Measurement plan + +| Output | Cadence | +|---|---| +| Train NLL, global FP32 gradient L2 norm, all LRs, phase, tokens/s, time, epoch | Every 10 updates (also first 5) | +| Full validation NLL, perplexity, top-1 token error and accuracy | Every 500 updates and every spectral snapshot | +| WeightWatcher | 0, 100, 250, 500, 750, 1000, 1500, 2000, 2500, 3000; then every 1000; plus 17500/final | +| Full checkpoints | Every 2500; also initial, 3000, 17500, final/clean stop | + +Every validation uses the **same 10,485,760 benchmark tokens** and the unchanged +evaluator. These are validation measurements, not a separate held-out test set. +Step-0 validation and CPU WeightWatcher must complete before the first update. + +All 72 hidden matrices are retained: Q, K, V, O, MLP_IN, MLP_OUT for each block. +`alpha_raw` comes only from `raw_alpha`; `alpha_clip_xmax` comes only from +WeightWatcher's clipped `alpha`. Fitting imposes no alpha=2 constraint. All +native scalar outputs are retained, including randomized/null and ERG statistics, +`alpha_weighted`, `log_alpha_norm`, matrix rank and fit diagnostics. The weighted +metrics retain WW's native clipped-alpha definition and are labelled accordingly. + +Zero-initialized matrices have rank 0 and explicit unavailable/degenerate fits; +they are not discarded or assigned invented alphas. Thus 72 rows at step 0 +does not imply 72 successful power-law fits. Across-matrix standard deviations +are not error bars across independent seeds. + +After step 0, the CPU tracker operates on immutable weight snapshots, independently +of training RNG and the data stream. Completed JSON/CSV measurements and training +metrics are flushed to the persistent disk. Transfer/enqueue overhead and CPU WW +time are recorded separately. Above 10% measured foreground overhead, or a >10% +median-update slowdown while the CPU tracker is active, later sampling drops to +every 2,000 updates. This second criterion is a conservative contention proxy, +**not proof** that WW caused the slowdown. Early/milestone/final measurements +remain, and scientific quantities do not change. The tracker drains at exit. + +## Startup and recovery + +The existing eight-chip flash forward/backward gate runs first. The trainer then +verifies its initial full validation and 72-row WW measurement. It checks finite +loss and sampled gradient norm while progressing. At update 2 it captures full +state, restores **its own step-0 initialization**, replays its first two updates, +and requires exact tensor/state equality. This is an in-process recovery test, +not a new training process and not a restore from the short reference. It stops +if the comparison fails and writes `RESUME_PARITY.json` only on success. + +After 100 updates it records steady timing relative to the reference, refuses a +>50% slowdown, and continues the **same process**. `STARTUP.json` is the evidence; +an allocated/active service alone does not mean those checks passed. + +Checkpoints atomically contain model, both optimizer states, scheduler, completed +step, token count, CPU/Python/NumPy/XLA RNG, exact shard/offset/cycle, corpus identity +and numerical source hashes. Keep the latest two rolling checkpoints plus +permanent 0, 3000, 10000, 17500, 25000 milestones. Final cloud backup copies these +retained checkpoint files and the small scientific outputs with CRC verification. +Intermediate checkpoints/results are immediately safe on the mounted disk; +cloud backup is an exit operation, not a per-update claim. + +For an **explicit recovery**, use the same commit and a compatible existing TPU: + +```bash +python3 baseline/gpt2_small/muon_longrun/launch.py recover \ + --checkpoint /mnt/disks/rg-data/gpt2small//checkpoints/step_0010000.pt +``` + +Recovery writes a new directory and rejects changed scheduler, data order, +numerical source, precision/runtime identity or model shape. A 3,000-step speedrun +checkpoint has a different schema and is rejected. No recovery happens implicitly. +CPU tests establish exact state replay across a shard boundary. The live TPU +startup gate establishes initial two-update replay only if it passes; it does +not prove that every later interruption or software/hardware change is bitwise +reproducible. Retain `RESUME_PARITY.json` and the pinned environment with the data. + +## Local verification + +```bash +OMP_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 \ +PYTHONPATH=baseline/gpt2_small/src:baseline/nanogpt_one_head/src \ +python3 -m pytest baseline/gpt2_small/tests/test_speedrun30.py \ + baseline/gpt2_small/tests/test_muon_speedrun.py \ + baseline/gpt2_small/tests/test_muon_longrun.py -q +``` + +Tests cover upstream model/optimizer parity, schedule boundaries, unchanged +sequential sampling, corpus budget, exact full-state CPU recovery, retention, +lease rejection, and real WeightWatcher raw/clipped/null fields with zero matrices. diff --git a/baseline/gpt2_small/muon_longrun/checkpoint.py b/baseline/gpt2_small/muon_longrun/checkpoint.py new file mode 100644 index 00000000..dcb55f1a --- /dev/null +++ b/baseline/gpt2_small/muon_longrun/checkpoint.py @@ -0,0 +1,81 @@ +"""Atomic full-state saves, bounded local retention, explicit deterministic restore.""" +from dataclasses import asdict +import os +import random +import numpy as np +import torch +from common import BATCH_TOKENS, PERMANENT, atomic_json, sync_dir + + +def cpu(value): + if isinstance(value,torch.Tensor): return value.detach().cpu().clone() + if isinstance(value,dict): return {k:cpu(v) for k,v in value.items()} + if isinstance(value,list): return [cpu(v) for v in value] + if isinstance(value,tuple): return tuple(cpu(v) for v in value) + return value + + +def state(model,muon,adam,stream,rt,step,schedule,identity): + rt.step(wait=True) + rng={'torch':torch.get_rng_state(),'numpy':np.random.get_state(),'python':random.getstate()} + if rt.tpu: rng['xla']=int(rt.xm.get_rng_state(device=rt.device)) + return dict(schema=2,step=step,tokens_seen=step*BATCH_TOKENS,config=asdict(model.config), + model=cpu(model.state_dict()),muon=cpu(muon.state_dict()),adam=cpu(adam.state_dict()), + data_cursor=stream.state_dict(),rng=rng,scheduler=asdict(schedule), + next_update=schedule.values(step),identity=identity) + + +def save(root,payload): + step=payload['step']; folder=root/'checkpoints'; folder.mkdir(exist_ok=True) + path=folder/f'step_{step:07d}.pt'; temporary=path.with_suffix('.tmp') + with temporary.open('wb') as f: + torch.save(payload,f); f.flush(); os.fsync(f.fileno()) + temporary.replace(path) + sync_dir(folder) + link=root/'checkpoint_latest.pt.tmp'; link.unlink(missing_ok=True) + os.link(path,link); link.replace(root/'checkpoint_latest.pt') + atomic_json(root/'checkpoint_latest.json',{'file':str(path.relative_to(root)),'step':step, + 'tokens_seen':payload['tokens_seen'],'next_update':payload['next_update'], + 'scheduler':payload['scheduler'],'schema':2}) + rolling=sorted(p for p in folder.glob('step_*.pt') if int(p.stem.split('_')[1]) not in PERMANENT) + for old in rolling[:-2]: old.unlink() + atomic_json(root/'checkpoint_inventory.json',{'files':[str(p.relative_to(root)) for p in sorted(folder.glob('*.pt'))], + 'permanent_steps':sorted(PERMANENT),'rolling_keep':2}) + return path + + +def restore(payload,model,muon,adam,stream,rt,schedule,identity): + if payload.get('schema')!=2 or payload['scheduler']!=asdict(schedule) or payload['identity']!=identity: + raise RuntimeError('Resume source, scheduler, corpus or numerical settings differ') + if payload['config']!=asdict(model.config): raise RuntimeError('Model configuration mismatch') + model.load_state_dict(payload['model'],strict=True) + for p in (*model.parameters(),*model.buffers()): rt.replicate(p) + muon.load_state_dict(payload['muon']); adam.load_state_dict(payload['adam']) + for group in muon.groups: rt.shard_matrices(group['buffer']) + # Adam counters/moments must follow the live parameter device, including capturable steps. + for p,values in adam.state.items(): + for key,value in values.items(): + if isinstance(value,torch.Tensor): + values[key]=value.to(rt.device); rt.replicate(values[key]) + stream.load_state_dict(payload['data_cursor']) + torch.set_rng_state(payload['rng']['torch']); np.random.set_state(payload['rng']['numpy']) + random.setstate(payload['rng']['python']) + rt.step(wait=True) + # Restore the seed AFTER materializing restored state; mark_step may advance it. + if rt.tpu: rt.xm.set_rng_state(payload['rng']['xla'],device=rt.device) + return int(payload['step']) + + +def assert_same(a,b,path='state'): + """Exact CPU comparison of full replay state; used only by the startup gate.""" + if isinstance(a,torch.Tensor): + if not torch.equal(a,b): raise RuntimeError('TPU replay differs at '+path) + elif isinstance(a,np.ndarray): + if not np.array_equal(a,b): raise RuntimeError('TPU replay differs at '+path) + elif isinstance(a,dict): + if a.keys()!=b.keys(): raise RuntimeError('TPU replay keys differ at '+path) + for k in a: assert_same(a[k],b[k],path+'.'+str(k)) + elif isinstance(a,(tuple,list)): + if len(a)!=len(b): raise RuntimeError('TPU replay lengths differ at '+path) + for i,(x,y) in enumerate(zip(a,b)): assert_same(x,y,path+'.'+str(i)) + elif a!=b: raise RuntimeError('TPU replay differs at '+path) diff --git a/baseline/gpt2_small/muon_longrun/common.py b/baseline/gpt2_small/muon_longrun/common.py new file mode 100644 index 00000000..7abf83a6 --- /dev/null +++ b/baseline/gpt2_small/muon_longrun/common.py @@ -0,0 +1,111 @@ +"""Fixed long-run plan; reuse the validated speedrun implementation.""" +from dataclasses import dataclass, asdict +import hashlib +import json +import os +from pathlib import Path +import sys + +HERE = Path(__file__).resolve().parent +SHORT = HERE.parent/'muon_speedrun' +sys.path.append(str(SHORT)) +STEPS = 25000 +BATCH_TOKENS = 524288 +VAL_TOKENS = 10485760 +MICROBATCH = 64 +CONTEXT = 1024 +REFERENCE_NAME = 'muon-speedrun-muon-20261005-030026' +CACHE = Path('/mnt/disks/rg-data/benchmark-fineweb10B-889765ea') +PERMANENT = {0, 3000, 10000, 17500, STEPS} +EARLY_WW = {0, 100, 250, 500, 750, 1000, 1500, 2000, 2500, 3000} + + +@dataclass(frozen=True) +class Schedule: + total_steps: int = STEPS + warmdown_steps: int = 7500 + warmup_steps: int = 0 + + @property + def cooldown_start(self): + return self.total_steps-self.warmdown_steps + + def factor(self, update_index): + return min(1., max(0., (self.total_steps-update_index)/self.warmdown_steps)) + + def phase(self, update_index): + return 'cooldown' if update_index >= self.cooldown_start else 'main' + + def values(self, update_index): + factor = self.factor(update_index) + return dict(lr_factor=factor, scheduler_phase=self.phase(update_index), + muon_lr=.04*factor, adam_embedding_lr=.6*factor, + adam_head_lr=.008*factor, adam_scalar_lr=.04*factor) + + +def ww_due(step, sparse=False): + interval = 2000 if sparse else 1000 + return step in EARLY_WW or step in PERMANENT or (step > 3000 and step % interval == 0) + + +def val_due(step, sparse=False): + return step % 500 == 0 or ww_due(step, sparse) + + +def adapt_tracking(step, foreground_fraction, median_seconds, reference_seconds, tracker_state): + """Conservative cadence reduction; observed host slowdown is NOT causal proof.""" + if step < 3000: return False + return foreground_fraction > .10 or ( + median_seconds > reference_seconds*1.10 and tracker_state.get('status')=='measuring') + + +def atomic_json(path, value): + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_suffix(path.suffix+'.tmp') + with temporary.open('w') as f: + json.dump(value, f, indent=2, allow_nan=False); f.write('\n') + f.flush(); os.fsync(f.fileno()) + temporary.replace(path) + sync_dir(path.parent) + + +def sync_dir(path): + fd = os.open(path, os.O_RDONLY) + try: os.fsync(fd) + finally: os.close(fd) + + +def sha(path): + h=hashlib.sha256() + with path.open('rb') as f: + for chunk in iter(lambda:f.read(8*1024*1024),b''): + h.update(chunk) + return h.hexdigest() + + +def verify_reference(reference): + manifest=json.loads((reference/'manifest.json').read_text()) + val=json.loads((reference/'latest_validation.json').read_text()) + expected={'recipe':'2024-11-10_UNetDoubleLr','optimizer':'muon','seed':1337, + 'config':{'vocab_size':50304,'n_layer':12,'n_head':6,'n_embd':768}, + 'batch_tokens':BATCH_TOKENS,'global_microbatch_sequences':64,'accumulation':8, + 'muon_lr':.04,'adam_embedding_lr':.6,'adam_head_lr':.008,'adam_scalar_lr':.04, + 'weight_decay':0,'gradient_clipping':False,'attention':'flash','warmup_updates':0, + 'data_repo':'kjj0/fineweb10B-gpt2','data_revision':'889765ea1f903759787add96995d81171b632d0c'} + for key,value in expected.items(): + if manifest.get(key)!=value: + raise RuntimeError(f'Reference setting differs: {key}: {manifest.get(key)!r}') + if not (val['step']==3000 and val['full_benchmark_evaluation'] + and val['evaluation_tokens']==VAL_TOKENS and 3 < val['val_nll'] < 3.3): + raise RuntimeError('Reference full 3,000-step validation is not available/healthy') + hashes={} + for relative in ('muon_speedrun/model.py','muon_speedrun/optim.py','muon_speedrun/runtime.py', + 'muon_speedrun/data.py','muon_speedrun/pallas_dependencies.py', + 'speedrun30/train.py','speedrun30/data_manifest.json'): + old=reference/'repo/baseline/gpt2_small'/relative + new=HERE.parent/relative + if sha(old)!=sha(new): + raise RuntimeError('Validated implementation changed: '+relative) + hashes[relative]=sha(new) + return {'reference_run':str(reference),'manifest':manifest,'final_validation':val, + 'unchanged_source_sha256':hashes} diff --git a/baseline/gpt2_small/muon_longrun/launch.py b/baseline/gpt2_small/muon_longrun/launch.py new file mode 100644 index 00000000..e39545d7 --- /dev/null +++ b/baseline/gpt2_small/muon_longrun/launch.py @@ -0,0 +1,191 @@ +"""Start/status/stop the fixed 25k plan on the existing TPU; never allocate a node.""" +import argparse +import datetime as dt +import fcntl +import importlib.util +import json +import os +from pathlib import Path +import re +import shlex +import shutil +import subprocess +import sys +import time +import uuid + +PROJECT='tpu-builders-504820'; ZONE='us-west4-a' +NODE='ww-gpt2-validation-48h-20261004-s1337-node' +BASE=Path('/mnt/disks/rg-data/gpt2small'); LATEST=BASE/'MUON_LONG25K_LATEST.json' +MIN_REMAINING=12.5*3600 + + +def run(command,**kwargs): return subprocess.run(command,check=True,text=True,**kwargs) + + +def timestamp(value): + value=re.sub(r'(\.\d{6})\d+',r'\1',value) + return dt.datetime.fromisoformat(value.replace('Z','+00:00')).timestamp() + + +def lease_from(node,queue,now): + if node.get('state')!='READY' or queue.get('state',{}).get('state')!='ACTIVE': + raise RuntimeError('Existing TPU/queue is not READY/ACTIVE') + if node.get('acceleratorType')!='v5litepod-8': + raise RuntimeError('Expected the existing single-host v5litepod-8') + candidates=[node.get('schedulingConfig',{}).get('terminationTimestamp'), + queue.get('runDuration',{}).get('terminationTime')] + for spec in queue.get('tpu',{}).get('nodeSpec',[]): + if spec.get('nodeId')==NODE: + candidates.append(spec.get('node',{}).get('schedulingConfig',{}).get('terminationTimestamp')) + expiries=[timestamp(x) for x in candidates if x] + if not expiries: + raise RuntimeError('API did not return an explicit termination time; refusing to infer it from queue creation') + expiry=min(expiries) + result={'node':NODE,'queue':queue['name'].rsplit('/',1)[-1], + 'queue_created':queue.get('createTime'),'node_created':node.get('createTime'), + 'max_run_duration':queue.get('runDuration',{}).get('maxRunDuration'), + 'termination_unix':expiry,'termination_utc':dt.datetime.fromtimestamp(expiry,dt.timezone.utc).isoformat(), + 'checked_unix':now,'remaining_hours':(expiry-now)/3600} + print(json.dumps(result,indent=2),flush=True) + if expiry-now<=MIN_REMAINING: + raise RuntimeError('Need more than 12.5 hours remaining; no run started and no allocation changed') + return result + + +def live_lease(): + flags=['--project='+PROJECT,'--zone='+ZONE,'--format=json'] + node=json.loads(run(['gcloud','alpha','compute','tpus','tpu-vm','describe',NODE,*flags],capture_output=True).stdout) + queue_name=node.get('queuedResource','').rsplit('/',1)[-1] + if not queue_name: raise RuntimeError('Could not determine the node\'s queued resource') + queue=json.loads(run(['gcloud','alpha','compute','tpus','queued-resources','describe',queue_name,*flags],capture_output=True).stdout) + return lease_from(node,queue,time.time()) + + +def active(unit): + r=subprocess.run(['systemctl','show',unit,'--property=ActiveState','--value'],capture_output=True,text=True) + return r.stdout.strip() in ('active','activating','deactivating','reloading') + + +def status_remote(tail_metrics=False): + if not LATEST.exists(): print('No long run launched.'); return + record=json.loads(LATEST.read_text()); root=Path(record['root']) + print(json.dumps(record,indent=2),flush=True) + subprocess.run(['systemctl','--no-pager','--full','status',record['unit']]) + for name in ('RUN_STATUS.json','STARTUP.json','status.json','latest_validation.json', + 'checkpoint_latest.json','TRACKING_STATUS.json','RESUME_PARITY.json'): + if (root/name).exists(): print(name+'\n'+(root/name).read_text(),flush=True) + subprocess.run(['tail','-n','12',str(root/('metrics.jsonl' if tail_metrics else 'run.log'))]) + + +def stop_remote(): + record=json.loads(LATEST.read_text()); root=Path(record['root']) + (root/'STOP').touch() + print('Safe stop requested. Trainer will finish the current update, save full state, then drain tracking/backup.') + print('Watch:',root/'run.log') + + +def start_remote(commit,lease,request_id,resume=None): + if os.geteuid()!=0 or not os.path.ismount('/mnt/disks/rg-data'): + raise RuntimeError('The persistent data disk must already be mounted; root required') + if not re.fullmatch('[0-9a-f]{40}',commit): raise ValueError('Pinned Git commit required') + if resume: + resume=Path(resume).resolve() + if BASE.resolve() not in resume.parents or resume.suffix!='.pt' or not resume.is_file(): + raise RuntimeError('Recovery checkpoint must exist under the persistent experiment directory') + with (BASE/'port-check-launch.lock').open('a') as lock: + fcntl.flock(lock,fcntl.LOCK_EX|fcntl.LOCK_NB) + if LATEST.exists(): + previous=json.loads(LATEST.read_text()) + if previous.get('request_id')==request_id or active(previous['unit']): + print('Existing launch retained; no duplicate or restart.'); status_remote(); return + if not -30 <= time.time()-lease['checked_unix'] <= 600: + raise RuntimeError('Lease check is stale; rerun the launcher') + if lease['node']!=NODE or lease['termination_unix']-time.time()<=MIN_REMAINING: + raise RuntimeError('Need more than 12.5 hours remaining on the checked node') + if shutil.disk_usage(BASE).free < 45*1024**3: + raise RuntimeError('Need 45 GiB free for remaining shards, spectra and retained checkpoints; nothing deleted') + stamp=dt.datetime.now(dt.timezone.utc).strftime('%Y%m%d-%H%M%S') + root=BASE/('muon-long25k-s1337-'+stamp); root.mkdir(); repo=root/'repo'; repo.mkdir() + run(['git','-C',str(repo),'init','-q']) + run(['git','-C',str(repo),'remote','add','origin','https://github.com/CalculatedContent/rg_optimizers.git']) + run(['git','-C',str(repo),'fetch','--depth','1','origin',commit],timeout=180) + run(['git','-C',str(repo),'checkout','--detach',commit]) + scripts=repo/'baseline/gpt2_small' + spec=importlib.util.spec_from_file_location('training_guard',scripts/'scripts/run_muonclip.py') + guard=importlib.util.module_from_spec(spec); spec.loader.exec_module(guard); guard.assert_idle() + # Verify reference BEFORE launching a service; never change that directory. + sys.path.insert(0,str(scripts/'muon_longrun')) + from common import verify_reference, REFERENCE_NAME, atomic_json + approved=verify_reference(BASE/REFERENCE_NAME) + deadline=time.time()+12*3600 + if lease['termination_unix']-deadline<1800: + raise RuntimeError('Less than 30 minutes lease margin after checkout; no service started') + unit='rg-muon-long25k-'+stamp+'.service' + record={'root':str(root),'unit':unit,'commit':commit,'request_id':request_id, + 'node':NODE,'started_unix':time.time(),'service_deadline_unix':deadline, + 'training_deadline_unix':deadline-1800,'steps':25000,'tokens':13107200000, + 'optimizer':'muon','config':approved['manifest']['config'], + 'batch_tokens':524288,'global_microbatch_sequences':64,'accumulation':8, + 'peak_lrs':{'muon':.04,'adam_embedding':.6,'adam_head':.008,'adam_scalar':.04}, + 'muon_momentum':{'initial':.85,'final':.95,'ramp_updates':500}, + 'newton_schulz':{'steps':5,'coefficients':[3.4445,-4.7750,2.0315]}, + 'adam':{'betas':[.9,.95],'eps':1e-8,'weight_decay':0}, + 'gradient_clipping':False,'qk_clipping':False, + 'scheduler':{'warmup_updates':0,'cooldown_start':17500,'warmdown_updates':7500,'final_step':25000}, + 'reference':approved['reference_run'],'lease':lease,'fresh_initialization':resume is None, + 'resume_source':str(resume) if resume else None, + 'automatic_restart':False,'cloud_uri':'gs://tpu-builders-504820-ww-continuous8/gpt2small/'+root.name} + atomic_json(root/'launch.json',record); atomic_json(root/'LEASE.json',lease) + (root/'commit.txt').write_text(commit+'\n') + env={'PYTHONPATH':str(scripts/'src')+':'+str(scripts.parent/'nanogpt_one_head/src'), + 'PJRT_DEVICE':'TPU','TPU_ACCELERATOR_TYPE':'v5litepod-8', + 'OMP_NUM_THREADS':'4','OPENBLAS_NUM_THREADS':'4','MKL_NUM_THREADS':'4', + 'TOKENIZERS_PARALLELISM':'false'} + command=['systemd-run','--unit='+unit,'--property=Type=exec','--property=Restart=no', + '--property=RuntimeMaxSec=43200','--property=TimeoutStopSec=15', + '--property=KillMode=control-group','--property=StandardOutput=append:'+str(root/'run.log'), + '--property=StandardError=append:'+str(root/'run.log')] + command+=['--setenv='+k+'='+v for k,v in env.items()] + command+=['/mnt/disks/rg-data/continuous8/venv/bin/python','-u', + str(scripts/'muon_longrun/long_worker.py'),str(root),str(deadline)] + if resume: command+=['--resume',str(resume)] + atomic_json(LATEST,record) # Publish request ID before systemd to make SSH retry idempotent. + run(command) + print(json.dumps(record,indent=2),flush=True) + run(['systemctl','show',unit,'--property=MainPID','--property=ActiveState']) + print('Muon run: 25,000 total steps, 13.1072B tokens; no target-loss stop.',flush=True) + print('Peak LR through 17,500; linear warmdown over final 7,500; no LR warmup.',flush=True) + print('Scalars every 10; full validation every 500 + spectral steps.',flush=True) + print('WW: 0,100,250,500,750,1000,1500,2000,2500,3000; then 1000, plus 17500/final.',flush=True) + print('Checkpoints every 2500; latest two rolling + permanent 0/3000/10000/17500/25000.',flush=True) + print('Estimated ~10–11 hours; 12-hour service cap includes preparation and final backup.',flush=True) + for action in ('status','metrics','stop'): + print(f'From your local checkout: python3 baseline/gpt2_small/muon_longrun/launch.py {action}',flush=True) + + +def main(): + p=argparse.ArgumentParser(); p.add_argument('action',choices=('start','recover','status','metrics','stop')) + p.add_argument('--checkpoint',type=Path,help='Explicit recovery from this long plan only; never used by start') + p.add_argument('--on-tpu',action='store_true',help=argparse.SUPPRESS) + p.add_argument('--commit',help=argparse.SUPPRESS); p.add_argument('--lease',help=argparse.SUPPRESS) + p.add_argument('--request-id',help=argparse.SUPPRESS); a=p.parse_args() + if (a.action=='recover') != bool(a.checkpoint): p.error('Only recover requires --checkpoint') + if a.on_tpu: + if a.action in ('start','recover'): start_remote(a.commit,json.loads(a.lease),a.request_id,a.checkpoint) + elif a.action=='stop': stop_remote() + else: status_remote(a.action=='metrics') + return 0 + command=['sudo','python3','-c',Path(__file__).read_text(),a.action,'--on-tpu'] + if a.checkpoint: command+=['--checkpoint',str(a.checkpoint)] + if a.action in ('start','recover'): + repo=Path(__file__).resolve().parents[3] + if run(['git','-C',str(repo),'status','--porcelain'],capture_output=True).stdout.strip(): + raise RuntimeError('Use a clean checkout of the pushed commit') + commit=run(['git','-C',str(repo),'rev-parse','HEAD'],capture_output=True).stdout.strip() + lease=live_lease() + command+=['--commit',commit,'--lease',json.dumps(lease),'--request-id',uuid.uuid4().hex] + return subprocess.run(['gcloud','compute','tpus','tpu-vm','ssh',NODE,'--project='+PROJECT, + '--zone='+ZONE,'--worker=0','--command='+shlex.join(command)]).returncode + +if __name__=='__main__': raise SystemExit(main()) diff --git a/baseline/gpt2_small/muon_longrun/long_data.py b/baseline/gpt2_small/muon_longrun/long_data.py new file mode 100644 index 00000000..197b611c --- /dev/null +++ b/baseline/gpt2_small/muon_longrun/long_data.py @@ -0,0 +1,56 @@ +"""The same sequential stream, with explicit cycles and recoverable cursor.""" +from concurrent.futures import ThreadPoolExecutor +import json +from pathlib import Path +from common import SHORT, CACHE, MICROBATCH, CONTEXT, atomic_json, sha +from data import FineWeb, reference + + +class Stream(reference.TrainStream): + def __init__(self, source, batch=MICROBATCH, context=CONTEXT): + super().__init__(source,batch,context) + self.cycles=0 + + def next_batch(self): + wrap = (self.shard == len(self.names)-1 and + self.position+self.batch*self.context+1 > len(self.tokens)) + result=super().next_batch() + if wrap: + self.cycles+=1 + return result + + def state_dict(self): + return dict(shard=self.shard,position=self.position,cycles=self.cycles, + batch=self.batch,context=self.context,names=self.names) + + def load_state_dict(self, state): + if state['names']!=self.names or state['batch']!=self.batch or state['context']!=self.context: + raise RuntimeError('Resume data ordering/batch mismatch') + shard,position=int(state['shard']),int(state['position']) + if not 0 <= shard < len(self.names): raise ValueError('Invalid shard cursor') + tokens=self.source.array(self.names[shard]) + if not 0 <= position < len(tokens) or position % (self.batch*self.context): + raise ValueError('Invalid within-shard cursor') + self.shard,self.position,self.cycles,self.tokens=shard,position,int(state['cycles']),tokens + + +def corpus_metadata(source): + files={k:v for k,v in source.manifest['files'].items() if '_train_' in k} + count=MICROBATCH*CONTEXT + unique=sum((v['size']-1024)//2 for v in files.values()) + usable=sum((((v['size']-1024)//2-1)//count)*count for v in files.values()) + return dict(repo=source.manifest['repo'],revision=source.manifest['revision'], + train_shards=len(files),corpus_tokens=unique,usable_tokens_per_epoch=usable, + epoch_definition='tokens_seen / usable_tokens_per_full_sequential_corpus_pass', + ordering='unchanged sequential shard traversal; wraps to first shard after a full pass', + manifest_sha256=sha(SHORT.parent/'speedrun30/data_manifest.json')) + + +def prepare(root, deadline): + source=FineWeb(CACHE,deadline) + names=sorted(source.manifest['files']) + with ThreadPoolExecutor(max_workers=4) as pool: + for name,_ in zip(names,pool.map(source.array,names)): + print('Verified benchmark shard:',name,flush=True) + atomic_json(root/'data_receipts.json',source.receipts) + atomic_json(root/'corpus.json',corpus_metadata(source)) diff --git a/baseline/gpt2_small/muon_longrun/long_worker.py b/baseline/gpt2_small/muon_longrun/long_worker.py new file mode 100644 index 00000000..8769beb3 --- /dev/null +++ b/baseline/gpt2_small/muon_longrun/long_worker.py @@ -0,0 +1,81 @@ +"""Bounded setup, one continuous trainer, async CPU WW, persistent scientific output.""" +import argparse +import json +import os +from pathlib import Path +import signal +import subprocess +import sys +import time +from common import HERE, SHORT, REFERENCE_NAME, atomic_json, verify_reference +from worker import bounded, finish_tracking + + +def cloud_backup(root): + from rg_nanogpt_one_head.continuous_support import CloudPublisher + publisher=CloudPublisher('gs://tpu-builders-504820-ww-continuous8/gpt2small/'+root.name) + receipts=[] + for folder in (root,root/'tracking'): + files=folder.iterdir() if folder==root else folder.rglob('*') + for path in sorted(files): + if path.is_file() and path.suffix in ('.json','.jsonl','.csv','.log','.txt'): + publisher.snapshot_text_file(path,path.relative_to(root).as_posix()) + for path in sorted((root/'checkpoints').glob('step_*.pt')): + receipts.append(publisher.file(path,path.relative_to(root).as_posix())) + result={'status':'verified','checkpoints':receipts,'unix_time':time.time()} + publisher.json(result,'CLOUD_BACKUP_VERIFIED.json'); atomic_json(root/'CLOUD_BACKUP_VERIFIED.json',result) + + +def main(): + p=argparse.ArgumentParser(); p.add_argument('root',type=Path); p.add_argument('deadline',type=float) + p.add_argument('--backup-only',action='store_true'); p.add_argument('--resume',type=Path) + a=p.parse_args(); root=a.root + if a.backup_only: cloud_backup(root); return 0 + train_deadline=a.deadline-1800 + state={'status':'preparing','automatic_restart':False,'deadline_unix':a.deadline, + 'training_deadline_unix':train_deadline,'steps':25000} + atomic_json(root/'RUN_STATUS.json',state); tracker=None + def phase(command,seconds,label): + if (root/'STOP').exists(): raise RuntimeError('Stop requested during setup') + result=bounded(command,min(seconds,train_deadline-time.time()),root,label) + if result['exit_code']!=0: raise RuntimeError(label+' failed: '+str(result)) + try: + approved=verify_reference(root.parent/REFERENCE_NAME) + logs=root.parent/REFERENCE_NAME/'metrics.jsonl' + if logs.exists(): + import statistics + durations=[r['seconds'] for line in logs.read_text().splitlines() + if (r:=json.loads(line)).get('kind')=='train' and r.get('step',0)>=100] + if durations: approved['median_update_seconds']=statistics.median(durations) + atomic_json(root/'REFERENCE.json',approved) + phase([sys.executable,str(SHORT/'pallas_dependencies.py'),str(root)],600,'pinned Pallas overlay') + os.environ['PYTHONPATH']=str(root/'pallas-deps')+os.pathsep+os.environ.get('PYTHONPATH','') + phase([sys.executable,'-u',str(HERE/'train_long.py'),'prepare','--root',str(root), + '--deadline',str(min(time.time()+1800,train_deadline-600))],1800,'verify full benchmark corpus') + phase([sys.executable,'-u',str(SHORT/'run.py'),'attention-check','--root',str(root), + '--microbatch','64','--deadline',str(time.time()+300)],300,'8-chip flash attention check') + env={**os.environ,'PJRT_DEVICE':'CPU','CUDA_VISIBLE_DEVICES':'', + 'OMP_NUM_THREADS':'1','OPENBLAS_NUM_THREADS':'1','MKL_NUM_THREADS':'1'} + tracker=subprocess.Popen([sys.executable,'-u',str(HERE/'track_long.py'),str(root), + str(a.deadline-1200)],env=env,start_new_session=True) + state['status']='training'; atomic_json(root/'RUN_STATUS.json',state) + command=[sys.executable,'-u',str(HERE/'train_long.py'),'train','--root',str(root), + '--deadline',str(train_deadline)] + if a.resume: command+=['--resume',str(a.resume)] + result=bounded(command,train_deadline-time.time(),root,'25,000 Muon updates',watch=True) + state.update(result) + if result['exit_code']!=0: raise RuntimeError('Continuous trainer failed; see FAILURE.json/run.log') + state.update(json.loads((root/'status.json').read_text())) + except Exception as exc: + state.update(status='failed',error=repr(exc)) + finally: + if tracker is not None: + state['tracking']=finish_tracking(tracker,root,min(time.time()+600,a.deadline-1200)) + atomic_json(root/'RUN_STATUS.json',state) + backup=bounded([sys.executable,'-u',__file__,str(root),str(a.deadline),'--backup-only'], + max(0,a.deadline-time.time()-20),root,'retained results cloud backup') + state['backup']=backup; atomic_json(root/'RUN_STATUS.json',state) + print(json.dumps(state),flush=True) + return 0 if state.get('status')=='schedule_complete' and backup['exit_code']==0 and state.get('tracking',{}).get('status')=='complete' else 1 + +if __name__=='__main__': raise SystemExit(main()) diff --git a/baseline/gpt2_small/muon_longrun/track_long.py b/baseline/gpt2_small/muon_longrun/track_long.py new file mode 100644 index 00000000..124df4f0 --- /dev/null +++ b/baseline/gpt2_small/muon_longrun/track_long.py @@ -0,0 +1,92 @@ +"""CPU spectral measurements; preserve explicit zero/unfit matrices at initialization.""" +import argparse +import importlib.metadata +import json +from pathlib import Path +import random +import time +from common import atomic_json, sha, EARLY_WW, PERMANENT +import tracking + + +def measure(path): + import numpy as np + import torch + import weightwatcher as ww + torch.set_num_threads(1) + began=time.monotonic(); payload=torch.load(path,map_location='cpu',weights_only=False) + seed=1001340+payload['step']; random.seed(seed); np.random.seed(seed); torch.manual_seed(seed) + names=list(payload['matrices']) + nonzero={n:v for n,v in payload['matrices'].items() if bool(torch.count_nonzero(v))} + identity={k:payload['validation'].get(k) for k in + ('evaluation_tokens','full_benchmark_evaluation','val_nll','val_perplexity', + 'val_token_error','val_accuracy','val_error_count','epoch','lr_factor', + 'scheduler_phase','muon_lr','adam_embedding_lr','adam_head_lr','adam_scalar_lr')} + identity.update(step=payload['step'],tokens_seen=payload['tokens_seen'],run_id=payload['run_id'], + snapshot_sha256=sha(path),diagnostic_seed=seed, + weightwatcher_version=importlib.metadata.version('weightwatcher')) + frame=ww.WeightWatcher(model=tracking.holder_from(nonzero)).analyze(**tracking.WW_OPTIONS) + if not {'alpha','raw_alpha'}.issubset(frame.columns): + raise RuntimeError('WeightWatcher raw/clipped fields missing') + rows=tracking.normalize_rows(frame,names,identity) + for row in rows: + name=row['matrix_name'] + if name not in nonzero: + row.update(status='zero_matrix',matrix_rank=0, + randomized_status='degenerate_zero_matrix') + else: + row['randomized_status']='available' if row.get('max_rand_eval') is not None else 'unavailable' + # Keep all WW scalar fields, plus explicit missing values where a fit is undefined. + for key in ('matrix_rank','alpha_weighted','log_alpha_norm','max_rand_eval', + 'rand_distance','rand_mp_softrank','rand_num_spikes','num_fingers'): + row.setdefault(key,None) + row['weighted_metric_alpha_source']='clipped alpha (WeightWatcher native definition)' + result=tracking.summary(rows,identity) + result.update(weightwatcher_seconds=time.monotonic()-began, + zero_matrix_count=len(names)-len(nonzero)) + return {'layers':rows,'summary':result,'options':tracking.WW_OPTIONS} + + +def configure(root): + version=importlib.metadata.version('weightwatcher') + if version!='0.7.7': raise RuntimeError('Expected the existing weightwatcher==0.7.7') + atomic_json(root/'TRACKING_CONFIG.json',dict(weightwatcher_version=version, + options=tracking.WW_OPTIONS,matrix_count=72,matrix_roles=list(tracking.ROLES.values()), + early_steps=sorted(EARLY_WW),later_interval=1000,always_steps=sorted(PERMANENT), + overhead_fallback_interval=2000,execution='separate CPU process, one BLAS thread', + raw_alpha_source='raw_alpha',clipped_alpha_source='alpha', + zero_matrix_policy='Keep row, rank=0, undefined alpha/null fits explicitly unavailable', + uncertainty='Across-matrix standard deviation is not uncertainty across seeds', + token_error='Validation teacher-forced top-1 error, same tokens and weights as NLL', + randomization='WeightWatcher randomize=True, all returned null/ERG fields retained')) + + +def watch(root,deadline): + configure(root); folder=root/'tracking' + while time.time()