diff --git a/baseline/experiments/distilbert_scratch_torchtpu/.gitignore b/baseline/experiments/distilbert_scratch_torchtpu/.gitignore new file mode 100644 index 0000000..2693678 --- /dev/null +++ b/baseline/experiments/distilbert_scratch_torchtpu/.gitignore @@ -0,0 +1,4 @@ +__pycache__/ +*.pyc +.venv/ +runs/ diff --git a/baseline/experiments/distilbert_scratch_torchtpu/AUDIT.md b/baseline/experiments/distilbert_scratch_torchtpu/AUDIT.md new file mode 100644 index 0000000..61744f3 --- /dev/null +++ b/baseline/experiments/distilbert_scratch_torchtpu/AUDIT.md @@ -0,0 +1,136 @@ +# Code and literature audit: BERT-6 MLM from scratch + +Reviewed 2026-10-20. The experiment is a randomly initialized, six-layer encoder +with ordinary masked-language-model pretraining. The implementation class is +`DistilBertForMaskedLM`; the model is not distilled. The intended first result is +a reproducible training trajectory on native Google TorchTPU, followed by +validation-based calibration and optimizer comparisons. + +## Findings and changes + +| Area | Audit result | +|---|---| +| Initialization | Uses `DistilBertForMaskedLM(DistilBertConfig(...))`; model weights are never loaded with `from_pretrained`. Only tokenizer files are reused. | +| Objective | Ordinary token cross-entropy; no teacher logits, distillation loss, cosine teacher loss, or next-sentence objective. | +| Optimizer | Stock PyTorch AdamW with bias correction and decoupled decay; biases and LayerNorm vectors receive no decay. Tied embedding/output weights occur once in optimizer groups. | +| Epsilon | Changed from the PyTorch default 1e-8 to 1e-6, matching Google's original BERT implementation and the academic-budget recipe. | +| Distributed loss | Sum losses, normalize by global masked-target count, compensate for DDP gradient averaging; padded examples contribute no loss or metrics. | +| Evaluation | Fixed per-example masks now use an independent 15% evaluation rate. Previously a training-rate change would also change evaluation difficulty. | +| Model retention | Added full best-validation weights, independent of rotating restart checkpoints. A later regression no longer removes the best observed validation model. | +| Configuration | Added Cloud Shell/supervisor config selection. Sequence preparation now uses the selected sequence length; previously it always prepared length 128. | +| Data identity | Pinned WikiText to `b08601e04326c79dfdd32d625aee71d232d685c3`, verified against the dataset API. Subsequent comparisons use the same upstream revision. | +| Restart | Full model, optimizer and each rank's RNG are preserved with hashes and source/config/data/package/world identity. Old runs must resume their original source/config. | +| Spectra | Exactly 36 encoder matrices; missing/failed fits cannot silently change the primary aggregate. No alpha-dependent updates or model selection. | +| Hardware | Native `torch_tpu==0.2.0`, `device='tpu'`, `tpu_dist`, compiler backend `tpu`. Actual v5e execution is still pending. | + +## Second complete code review, 2026-10-21 + +The review covered every Python module, configuration, dependency declaration +and runbook in this experiment folder. It found additional failures that the +original numerical checks did not exercise: + +| Finding | Correction | +|---|---| +| Native RNG API | The released 0.2.0 wheel marks `torch.tpu.get_rng_state` internal on PyTorch >=2.13, with a strict runtime error for user code. Use the supported accelerator getter when available. Add a real RNG replay check to every rank's TPU probe. The setter remains the wheel's public experimental API. | +| Diagnostics ahead of resume state | Scheduled evaluation could publish spectra/best weights before committing the matching training checkpoint. Commit state first, including when measurement/checkpoint intervals differ; recover an interrupted scheduled measurement on resume. | +| Unpublished checkpoint reuse | A crash after directory rename but before pointer publication could leave a checkpoint from an earlier replay. Replace that orphan when the cursor is replayed; retain only published state at an already committed cursor. | +| Checkpoint retention | Partial directories and unpublished generations ahead of the current cursor could displace a valid checkpoint. Filter them before numerically selecting the last two committed checkpoints. | +| Final spectral scan | Training could publish its final snapshot between the worker's scan and completion check. Read the completion marker before scanning, ensuring a final scan observes all published measurements. | +| Process cleanup | Analyst timeout handling could itself time out without killing/reaping the process group. Escalate from SIGTERM to SIGKILL, including leftover children after a failed phase. | +| Status near lease expiry | Status discovery applied the 30-minute launch gate, hiding a still-running job near expiry. Apply lease eligibility only to new launches; keep the launch guard unchanged. | +| Invalid configurations | Reject nonfinite optimizer scalars, fractional sizes, invalid seeds, and non-boolean mode flags before training. | + +Regression coverage now includes exact restart for AdamW, Adam and SGD, compiled +BF16 restart with dropout, simulated interrupted measurement publication, +unpublished checkpoint replacement, final-scan races and forced process cleanup. +API-selection tests use mocks and do not substitute for the real TPU probe. +No full-corpus repacking was needed: the preparation code and pinned data are +unchanged from the verified corpus check. + +## What the papers support + +1. **Sanh et al., DistilBERT (2019).** The original student inherited alternating + teacher layers and used teacher-derived losses alongside MLM. That explains + the name. Neither its published quality retention nor its speed comparisons + establish the results of our independently initialized model. + [Paper](https://arxiv.org/abs/1910.01108) + +2. **Turc et al., Well-Read Students Learn Better (2019).** Direct pretraining and + fine-tuning of compact BERT models is a competitive baseline; distillation is + an additional stage, not a requirement for using a compact architecture. + [Paper](https://arxiv.org/abs/1908.08962) + +3. **Devlin et al., BERT (2019).** LR 1e-4, batch 256, betas .9/.999, decay .01, + dropout .1 and short-sequence training provide a reasonable conservative + starting point. Google's implementation specifies epsilon 1e-6 and clipping + norm 1. Its legacy Adam update omits bias correction, unlike our stock PyTorch + AdamW: this experiment is not a bitwise reproduction of BERT optimization. + [Paper](https://arxiv.org/abs/1810.04805) · + [Optimizer source](https://github.com/google-research/bert/blob/master/optimization.py) + +4. **Liu et al., RoBERTa (2019).** Dynamic masking, training without NSP, sufficient + data and sufficient optimization are established improvements. This supports + the basic objective, but does not validate our short WikiText pilot as a + fully pretrained general-purpose encoder. + [Paper](https://arxiv.org/abs/1907.11692) + +5. **Izsak et al., How to Train BERT with an Academic Budget (2021).** Tune LR, + warmup, batch and schedule for the available computation. Larger encoders can + outperform smaller ones at equal wall time, so six layers is not proved best + for that metric. Their larger-model, batch-4096+ results do not justify + copying a 1e-3 or 2e-3 LR into our batch-256 run. Our initial 3-by-2 search is + an adaptation, not their published recipe or a measured optimum. + [Paper](https://arxiv.org/abs/2104.07705) + +6. **Wettig et al., Should You Mask 15% in Masked Language Modeling? (2023).** + Optimal masking depended on model size and task: their aggregate medium, + base and large results preferred 15%, 20% and 40%, respectively. This supports + retaining 15% for the first compact-model control and testing alternatives. + Their experiments also motivate a separate all-MASK corruption comparison; + the first run keeps standard 80/10/10 so only one recipe is established first. + [Paper](https://aclanthology.org/2023.eacl-main.217/) + +7. **Gisserot-Boukhlef et al., Should We Still Pretrain Encoders with MLM? + (2026 revision).** Their 210M–1B-model experiments found task-dependent masking + and advantages for CLM-then-MLM under some fixed-compute conditions. This is + evidence against a universal best MLM recipe, not a reason to change the + objective during our initial optimizer-controlled experiment. + [Paper](https://arxiv.org/abs/2507.00994) + +8. **Warner et al., ModernBERT (2024).** The released 149M base model combines + architecture changes with much larger-scale pretraining and GPU-specific + efficiency work. Benchmark its native TorchTPU implementation before making + a speed claim or changing the architecture in an optimizer comparison. + [Paper](https://arxiv.org/abs/2412.13663) + +## Calibration and evidence required + +- Start one model with `configs/initial.json`. Its 10,000 updates cover at most + 327,680,000 sequence positions, including special tokens, not that many MLM + targets. It is a pilot, not a convergence guarantee. +- The six committed configurations cover LR 1e-4/3e-4/5e-4 crossed with + warmup 300/1,000 updates. The other settings and both seeds are identical. + Each launch runs one configuration; no automatic sweep is enabled. +- Compare runs only with identical prepared-data hashes, resolved software + versions, token budget and fixed evaluation masks. Select using validation + loss, expand the LR range if the best setting lies on a boundary, and check + the selected recipe on fresh seeds. Compare both matched token budget and + matched measured wall time when evaluating architectures. +- After establishing LR/warmup, test beta2 .98 versus .999 and masking 15/20/30% + separately. Retune optimizer-specific LRs for SGD/Adam/AdamW instead of reusing + one optimizer's best LR for all optimizers. Alpha remains an outcome. +- Decide the longer training horizon from actual TPU throughput and validation + curves. A pause/resume keeps its existing horizon; extending the horizon is a + different protocol and must not silently change a run's schedule. + Pilot-selected settings remain provisional until checked at that longer budget. +- WikiText-103 is a convenient controlled corpus. For broader model quality, + add a larger, fixed corpus and a declared downstream evaluation. MLM validation + loss alone does not establish optimal downstream performance. +- A 128-token run has not trained longer absolute position embeddings. Any + 512-token stage needs explicit data preparation, schedule and qualification. +- Checkpoint alpha below two should be interpreted with the fit's tail size, + KS distance and uncertainty, the layer set, held-out quality and repeat seeds. + A conventional Transformer is not the scale-invariant MLP control. + +See `VALIDATION.json` for the checks actually executed. CPU tests are not evidence +of native TPU kernel coverage, training throughput, convergence or observed alpha. diff --git a/baseline/experiments/distilbert_scratch_torchtpu/README.md b/baseline/experiments/distilbert_scratch_torchtpu/README.md new file mode 100644 index 0000000..48dfc47 --- /dev/null +++ b/baseline/experiments/distilbert_scratch_torchtpu/README.md @@ -0,0 +1,256 @@ +# Six-layer BERT-style MLM from scratch with native Google TorchTPU + +This independent folder trains **one randomly initialized DistilBERT architecture** +using masked-language modeling (MLM) on the existing eight-device v5e TPU. It uses +[`google-pytorch/torch_tpu`](https://github.com/google-pytorch/torch_tpu), +`device='tpu'`, `tpu_dist`, and `torch.compile`. It has no `torch_xla` imports and +does not reuse the earlier pretrained IMDb fine-tuning training code. + +This is a first pretraining pilot. There is **no teacher, pretrained model, or +distillation objective**. Only the existing uncased WordPiece tokenizer is reused. +We call this experiment **BERT-6 MLM from scratch**. `DistilBertForMaskedLM` is +the stock architecture class, not a request to distill a teacher. The original +DistilBERT paper's performance claims do not transfer to our random-init recipe. +The goal is to establish a reproducible training/measurement path, then assess +convergence before launching optimizer/seed comparisons. A below-two fit in one +run does not establish optimizer bias or overfitting. + +## Baseline + +| Setting | Value | +|---|---| +| Architecture | 6 blocks, 12 heads, width 768, FFN 3072, GELU, stock biases/LayerNorm/residuals | +| Parameters | 66,985,530, including tied MLM vocabulary projection | +| Initialization | Random Hugging Face `DistilBertForMaskedLM(DistilBertConfig(...))` | +| Objective | MLM only; 15% candidate selection; selected tokens use 80% mask / 10% random / 10% unchanged | +| Corpus | `Salesforce/wikitext`, `wikitext-103-raw-v1`; official train/validation/test partitions | +| Data provenance | Pin dataset revision `b08601e04326c79dfdd32d625aee71d232d685c3`; save packed-file SHA256s, row counts, tokenizer revision | +| Seed | Model/dropout 401; independent shared data/masking seed 1701 | +| Optimizer | AdamW; LR 1e-4; decay 0.01; betas .9/.999; epsilon 1e-6 | +| Schedule | 1,000 warmup updates, then linear decay over a fixed 10,000-update horizon | +| Batch | 32 per device × 8 = 256 sequences; length 128 | +| Token budget | 327,662,848 real sequence positions at the verified pinned corpus; about 2.803 corpus passes | +| Precision | BF16 autocast; FP32 parameters and optimizer moments; TPU matmul precision `medium` | +| Regularization | Hidden and attention dropout .1; gradient norm clipping 1.0 after DDP reduction | +| Checkpoints | Every 100 updates, before measurements, and cooperative exit; retain two committed generations | +| Best model | Independently retain full weights with minimum fixed-mask validation loss; ties retain the earlier checkpoint | +| Measurements | Initialization, update 2 preflight, every 500 updates, and final cursor if budget permits | +| First launch budget | 5.5 hours including setup; capped by checked existing lease; resumable if incomplete | + +LR 1e-4 is a conservative **pretraining pilot setting**, not the earlier 2e-5 +fine-tuning LR, nor a tuned optimum. The 10,000-update target is a pilot budget, +not a claim of convergence or equivalence to the original DistilBERT recipe. +The local full-corpus check produced 913,341 training, 1,913 validation and 2,191 +test sequences. Every row passed boundary, vocabulary and maskable-content checks; +`VALIDATION.json` records the exact hashes. Reproduce these checks on the target VM. +Actual runtime and throughput require TPU execution. The two-update preflight +checks correctness and restart; it is not a steady-state speed benchmark. Measure +at least 200 updates after compilation before choosing the longer training horizon. +No overnight completion estimate is asserted before that measurement. + +The [code/literature audit](AUDIT.md) explains which choices are supported by the +papers and which require calibration on this TPU. Six layers is a practical first +research baseline, not an established optimum for quality at fixed wall time. +At sequence length 128, positions 128..511 receive no prediction-driven training; +the run does not establish 512-token capability. + +## Launch from the existing authenticated Cloud Shell + +Use a clean checkout of the experiment branch (or its exact commit): + +```bash +git clone --branch codex/distilbert-scratch-torchtpu --single-branch \ + https://github.com/CalculatedContent/rg_optimizers.git rg-distilbert-torchtpu +cd rg-distilbert-torchtpu +python3 baseline/experiments/distilbert_scratch_torchtpu/cloudshell.py start \ + --zone us-west4-a --hours 5.5 +``` + +Project is `tpu-builders-504820`. The launcher discovers the sole eligible +**existing** READY/ACTIVE `v5litepod-8` allocation. If there are multiple eligible +allocations, supply `--node NAME`. It checks the current lease, matching guest IP, +device occupancy, RG services, mounted persistent disk, and 30 GiB free space. +It never provisions a TPU, extends a lease, stops another job, or changes an older +environment. The existing VM must have Python 3.12+ with venv support. If missing, +install Python 3.12 and `python3.12-venv` on that VM before launch. + +The bounded systemd service creates its own venv under +`/mnt/disks/rg-data/distilbert_scratch_torchtpu/initial-TIMESTAMP/`, installs pinned +TorchTPU 0.2.0 and HF/WW dependencies, and freezes the entire resolved environment +to `requirements.lock.txt`. Torch/libtpu transitive versions are resolved at the +first install and recorded; they are **not predeclared as hardware-qualified**. +The original venv is reused for resume, which also checks package versions. +TorchTPU supplies the required Torch/libtpu stack; the old XLA venv is not inherited. + +The job performs: + +1. All-eight-rank native TPU matmul, all-reduce, and RNG save/restore probe. +2. Real WikiText-103 token preparation with immutable hashes. +3. Initialization evaluation, two real full-model training updates, and checkpoint. +4. Process exit and resume from update 2, continuing the same model to its fixed horizon. +5. CPU WeightWatcher analysis on saved weights alongside subsequent training. + +To launch one of the committed calibration configurations, add for example: + +```bash +python3 baseline/experiments/distilbert_scratch_torchtpu/cloudshell.py start \ + --zone us-west4-a --hours 5.5 --config configs/pilot_lr3e4_warmup300.json +``` + +`initial.json` and the five `pilot_*.json` files form a 3-by-2 grid: peak LR +`1e-4`, `3e-4`, `5e-4`, each with 300 or 1,000 warmup updates. All other settings, +including seed and 10,000-update horizon, are identical. Each command launches +exactly one model; no sweep starts automatically. Configuration paths must refer +to committed JSON files inside this folder. Preparation follows the selected +sequence length. Resume uses the original configuration and cannot accept +`--config`. + +Status: + +```bash +python3 baseline/experiments/distilbert_scratch_torchtpu/cloudshell.py status \ + --zone us-west4-a +``` + +Check `probe/probe.json` for eight successful ranks, `preflight.json` for two +saved real updates, `initial/runtime.json` for native TPU and `resumed_step: 2`, +and `initial/status.json` for advancing steps. `RESULT.json` distinguishes a +completed target from a partial saved run. `ERROR.json` and `run.log` explain +failures. A CPU unit test is never recorded as TPU success. + +Resume with the same source commit and idle existing TPU: + +```bash +python3 baseline/experiments/distilbert_scratch_torchtpu/cloudshell.py start \ + --resume-run /mnt/disks/rg-data/distilbert_scratch_torchtpu/initial-TIMESTAMP/initial \ + --hours 5.5 +``` + +Checkpoints include model, optimizer, update cursor, every rank's CPU/TPU RNG +states, manifest hashes, configuration, data identity, package versions, and source +hash. The LR horizon stays fixed across pauses. Interrupted writes remain +uncommitted; the most recent complete checkpoint is the resume point. A hard kill +can lose work since the last 100-update checkpoint. Exact restart is tested on +CPU; TPU determinism and kernel coverage require the live hardware run. +Measurements are published only after their training checkpoint is committed. +If a process exits between those operations, resume fills the missing scheduled +measurement before advancing. Unpublished checkpoint directories from a prior +interrupted replay are replaced; partial and future unpublished directories do +not affect retention of the last two committed checkpoints. + +`best_validation.json` identifies the separately retained full model weights and +their SHA256. Its selection criterion is validation MLM loss only, independent of +test scores and alpha. It is an evaluation artifact; `latest_checkpoint.json` +still identifies the model/optimizer/RNG state for exact continuation. Validation +selection is relative to the fixed prepared corpus and evaluation masks, not an +assertion of optimal downstream accuracy or convergence. + +## Measurements and scientific interpretation + +The primary spectral set is exactly **36 encoder matrices**: attention Q/K/V/output +and both FFN weights in each of six blocks. Embeddings, the tied vocabulary output, +and the MLM prediction transform are excluded from that declared aggregate. +Snapshots are raw unnormalized FP32 weights. WeightWatcher 0.7.7's default +`analyze()` fits each matrix independently, matching the prior RG experiment's +estimator. There is no spectral feedback or alpha-dependent learning rate. + +Saved columns include alpha, KS distance D, sigma, xmin/xmax, tail count/fraction, +fit status and warnings. Model summaries report mean/min alpha and the fraction +of layers below two. Any invalid or missing expected fit makes the primary mean +null; it never silently changes the layer set. D and tail sizes must accompany +interpretation, especially for untrained or poorly power-law-like spectra. + +At every measurement the frozen model evaluates a fixed 2,048-sequence training +subset and the complete held-out validation/test partitions. **These accuracies +are masked-token accuracies, not IMDb/GLUE classification accuracy.** The same +per-example evaluation masks are used at every checkpoint and across seeds. +The evaluation masking rate is independently fixed at 15%, even when a later +experiment changes the training masking rate. +Training masks change with the update; data order/masking seed is separate from +model initialization. Evaluation restores training RNG state. Test metrics are +monitoring-only and never control configuration, stopping, updates, or selection. + +Packing occurs separately inside each official split. Each full block has CLS/SEP +boundary tokens; line boundaries receive SEP. Incomplete tails are dropped and +counted. Special tokens are never prediction targets. Rare zero-target sequences +receive one random eligible target; this nonempty-target safeguard is recorded +in the protocol. Epoch-final batches use loss-masked duplicate padding so every +real sequence contributes exactly once per pass. DDP scales by the global number +of targets, not by a mean of unequal per-worker means. + +Outputs in `initial/`: + +- `status.json`, `runtime.json`, `identity.json`, `RESULT.json`. +- `checkpoints/`, `best_validation/`, `best_validation.json`, `measurements/`, `spectral_snapshots/`, `spectra/`. +- `checkpoint_metrics.csv`: alpha summaries joined by step to loss/accuracy. +- `layer_metrics.csv`: every layer's fit joined to the same checkpoint metrics. + +Analyze remaining snapshots after training without using TPU devices: + +```bash +TORCH_DEVICE_BACKEND_AUTOLOAD=0 /path/to/run/venv/bin/python \ + baseline/experiments/distilbert_scratch_torchtpu/spectra.py --root /path/to/run/initial +``` + +## Next phase, after the pilot + +The config already exposes AdamW, Adam, SGD/momentum, LR, decay, dropout and seed. +No sweep starts automatically. First assess validation-loss progress, throughput, +fit quality, and whether the training budget reaches a useful regime. Then freeze +the corpus, tokenizer, masking/data schedule, architecture, and token budget; +calibrate optimizer-specific LR ranges; compare seeds 401/402/403 at matched +compute and at comparable validation quality. Preserve per-layer below-two +indicators separately from mean-alpha-below-two. Add Muon only after a separate +implementation/parity check on the new backend. + +The first calibration grid varies LR and warmup only. Rank completed, comparable +runs by fixed-mask validation loss; do not select with alpha or test results. +Match dataset hashes, resolved package versions and evaluation protocol before +comparison. A winning boundary LR calls for expanding the search. Check the +winning recipe on new seeds, then compare training mask rates 15/20/30% and +beta2 .999/.98 as separate, controlled ablations. These are candidate settings, +not prevalidated winners. Strong downstream-quality claims also require a fixed +classification/retrieval evaluation; the current runner supplies MLM metrics. + +DistilBERT's stock architecture includes biases, residual additions, softmax and +LayerNorm. It does not enforce the scale invariance of the earlier MLP controls. +This pilot measures spectral trajectories; it cannot by itself test a universal +SGD alpha=2 prediction or attribute alpha<2 causally to AdamW. + +## Engineering checks and upstream references + +```bash +cd baseline/experiments/distilbert_scratch_torchtpu +OMP_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 python -m unittest test_experiment -v +``` + +Checks cover the full architecture/random initialization path, eight-way exact +sample coverage, fixed masking, invalid-fit handling, tiny-model interrupted +versus uninterrupted AdamW/Adam/SGD states and metrics, compiled BF16 restart, +corruption rejection, and in-process eight-shard gradients versus a serial +masked-token mean. Fault checks cover checkpoint/measurement publication, +final spectral-worker scans, and termination of a process that ignores SIGTERM. An +additional opt-in two-process Gloo integration check compares actual DDP against +serial updates, including an uneven final batch; run it with +`RG_RUN_DISTRIBUTED_TESTS=1` on a host that permits Gloo sockets. See `VALIDATION.json` for actual results and limitations. + +Implementation references inspected on 2026-10-20: + +- [TorchTPU installation](https://github.com/google-pytorch/torch_tpu#installation) +- [Native distributed setup](https://github.com/google-pytorch/torch_tpu/blob/main/docs/distributed_support.md) +- [Supported hardware/features](https://github.com/google-pytorch/torch_tpu/blob/main/docs/feature_set.md) +- [Precision](https://github.com/google-pytorch/torch_tpu/blob/main/docs/numerical_precision.md) +- [Upstream AdamW tests](https://github.com/google-pytorch/torch_tpu/blob/main/tests/adamw_test.py) +- [DistilBERT implementation](https://github.com/huggingface/transformers/blob/v4.57.1/src/transformers/models/distilbert/modeling_distilbert.py) + +The released 0.2.0 wheel was also inspected directly: its metadata requires +Python >=3.12 (stricter than the current main README), torch >=2.11.0 and <=2.15.0, +and libtpu >=0.0.50. Its Python launch helper sets topology and device addresses; +the launcher derives the eight-device count from those addresses before torchrun. +On PyTorch 2.13+, TorchTPU restricts direct `torch.tpu.get_rng_state()` calls to +framework code. The runner uses `torch.accelerator.random.get_rng_state()` when +available, retaining the legacy fallback for older supported PyTorch versions. +The eight-rank probe checks actual RNG replay before preparing the corpus. + +The inspected TorchTPU main commit was +`905a0dfde04a01df8ee2c100ffd656694f38df84`; installation pins the public 0.2.0 wheel. diff --git a/baseline/experiments/distilbert_scratch_torchtpu/VALIDATION.json b/baseline/experiments/distilbert_scratch_torchtpu/VALIDATION.json new file mode 100644 index 0000000..43a5070 --- /dev/null +++ b/baseline/experiments/distilbert_scratch_torchtpu/VALIDATION.json @@ -0,0 +1,156 @@ +{ + "date": "2026-10-21", + "status": "CPU engineering validated; real TPU qualification pending", + "base_commit": "afd96a379f4b22c9199f58ca1442a89c5e087700", + "source_sha256": "8e8909b3fe98e521ec361aa73b546c4d2b3cd1de1bb129d098d52c17be592beb", + "suite": { + "tests_run": 20, + "passed": 19, + "skipped": 1, + "seconds": 234.851, + "source_sha256": "e46da64d3f9a79f3bf65ffb550fc68d88636548adc2e1bfe4e3a767c385d7082", + "note": "Complete suite before the final retention-filter extension; focused checkpoint tests rerun on final source below." + }, + "passed": [ + "full 66,985,530 parameter architecture and 36 encoder matrices", + "eight-way exact sample coverage including padded remainder", + "eight-shard weighted gradients agree with serial masked-token mean", + "fixed evaluation masking excludes special tokens and padding targets", + "incomplete spectral fits produce null aggregate", + "four-update tiny AdamW, Adam and SGD training each equal two updates plus restart exactly: parameters, optimizer state, evaluation metrics; checkpoint corruption refused", + "training mask-rate changes leave fixed evaluation inputs and labels bitwise identical", + "best full model weights survive later validation regression; test improvements cannot replace them; hash and retained weights verified", + "AdamW decay excludes LayerNorm and biases and applies once to tied embedding/output weights", + "compiled aot_eager CPU BF16 loss, counts and gradients agree with eager execution", + "six calibration configurations differ only in LR and warmup", + "Cloud Shell and supervisor CLI expose mutually exclusive config selection and exact resume", + "complete real WikiText-103 preparation; all packed rows checked for boundaries, vocabulary bounds, maskable content and exact manifest hashes", + "compiled aot_eager BF16 AdamW training with dropout resumes bitwise exactly on CPU", + "RNG API selection uses the supported accelerator getter and handles older releases without it; mocked API contract test only", + "diagnostics follow a committed checkpoint and missing scheduled measurements recover after injected interruption", + "orphan checkpoint replacement and two-generation retention ignore partial and future unpublished directories", + "spectral worker consumes a measurement published concurrently with its completion marker", + "owned subprocess that ignores SIGTERM is killed and reaped", + "status discovery works near expiry while new launches still require an eligible lease", + "nonfinite optimizer scalars, fractional sizes, invalid seeds and nonboolean flags are rejected" + ], + "spectral_smoke": { + "weightwatcher": "0.7.7", + "synthetic_matrix_shape": [ + 128, + 128 + ], + "fit_and_joined_json_and_csv": "passed", + "scientific_result": false + }, + "distributed_cpu_integration": { + "attempted": true, + "passed": false, + "blocker": "Gloo TCP interface initialization: Operation not permitted in this execution environment", + "rerun": "RG_RUN_DISTRIBUTED_TESTS=1 OMP_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 python -m unittest test_experiment -v" + }, + "native_tpu": { + "launched": false, + "completed_model": false, + "alpha_results": null, + "blocker": "No authenticated Google Cloud CLI in this session; Cloud Console browser returned Site Unavailable" + }, + "engineering_packages": { + "torch": "2.14.1", + "transformers": "4.57.1", + "numpy": "2.3.5", + "datasets": "4.2.0", + "weightwatcher": "0.7.7" + }, + "upstream": { + "repository": "https://github.com/google-pytorch/torch_tpu", + "inspected_main_commit": "905a0dfde04a01df8ee2c100ffd656694f38df84", + "wheel": "torch_tpu-0.2.0-cp312-cp312-manylinux_2_31_x86_64.whl", + "wheel_sha256": "cf725f2c38a5287f15095c2bb19dcb5900e445f6abe3ab622aa51c0f3f20bc0f", + "wheel_requires_python": ">=3.12", + "wheel_requires_torch": ">=2.11.0,<=2.15.0", + "release_source_api_verified": [ + "set_tpu_launch_env(nnodes=1)", + "TORCH_TPU_SLICEBUILDER_ADDRESSES detection", + "get_rng_state", + "set_rng_state", + "tpu_dist", + "get_rng_state is strict-internal on torch >=2.13; replacement is torch.accelerator.random.get_rng_state" + ] + }, + "pending": [ + "all-eight-device native TorchTPU collective and RNG checkpoint/restore probe", + "real full-model compiled forward/backward and optimizer steps", + "checkpoint/restart on actual TPU", + "throughput, convergence, and spectral results", + "reproduce verified real-data preparation on target TPU VM" + ], + "audit_base_commit": "896fa646b94872f8cf45b092127c822fc2880c2e", + "audit_changes": [ + "Adam epsilon aligned with BERT implementation: 1e-6", + "fixed evaluation mask rate independent of training mask rate", + "full model retained by validation loss independently of rolling resume checkpoints", + "selectable committed configurations; preparation follows configured length", + "WikiText dataset revision pinned for repeatable comparisons", + "experiment labeled BERT-6 MLM from scratch to distinguish architecture from distillation" + ], + "real_corpus_preparation": { + "passed": true, + "hardware": "CPU data preparation only; no model training", + "dataset_revision": "b08601e04326c79dfdd32d625aee71d232d685c3", + "tokenizer_revision": "12040accade4e8a0f71eabdb258fecc2e7e948be", + "sequence_length": 128, + "vocab_size": 30522, + "splits": { + "train": { + "rows": 913341, + "text_tokens": 113916064, + "dropped_remainder": 113, + "file": "train.bin", + "sha256": "3d58a53d18e302cfb924a56b632a728164e1dd8d54763c8fc23e0ee3cbf2a408", + "all_rows_structure_and_vocabulary": "passed" + }, + "validation": { + "rows": 1913, + "text_tokens": 238656, + "dropped_remainder": 79, + "file": "validation.bin", + "sha256": "a2cc68a858d7ed17db0ea93a0f30020395826ae4fb551a642d39ceaba8393ef5", + "all_rows_structure_and_vocabulary": "passed" + }, + "test": { + "rows": 2191, + "text_tokens": 273178, + "dropped_remainder": 3, + "file": "test.bin", + "sha256": "ddf490b3d50331584b329b0bfa6b40ce290cec5a8ab222f42f4bb8cc75805bc1", + "all_rows_structure_and_vocabulary": "passed" + } + }, + "steps_per_corpus_pass": 3568, + "pilot_full_corpus_passes": 2.802749465971636, + "pilot_real_sequence_positions": 327662848, + "validated_date": "2026-10-20" + }, + "final_retention_regression": { + "tests_run": 2, + "passed": 2, + "seconds": 4.492, + "source_sha256": "8e8909b3fe98e521ec361aa73b546c4d2b3cd1de1bb129d098d52c17be592beb", + "checks": [ + "unpublished current/future generations and partial directories preserve two committed checkpoints", + "measurement follows committed state and interrupted scheduled measurement recovers" + ] + }, + "second_audit_base_commit": "c2511497856c7367d4e2f4c3100963d77fd67f3d", + "second_audit_scope": "All Python modules, configurations, requirements and documentation in this experiment folder", + "second_audit_changes": [ + "native TPU RNG getter compatible with strict TorchTPU 0.2.0 API; real RNG replay added to each rank probe", + "checkpoint publication precedes diagnostics; interrupted measurements recover on resume", + "unpublished replay checkpoints replaced; partial and future generations excluded from retention", + "final spectral scan race fixed; completeness compares matching step filenames", + "bounded process-group cleanup escalates to SIGKILL and reaps owned processes", + "read-only status discovery no longer applies the launch lease-time threshold", + "configuration validation rejects nonfinite values and invalid types/seeds" + ] +} diff --git a/baseline/experiments/distilbert_scratch_torchtpu/cloudshell.py b/baseline/experiments/distilbert_scratch_torchtpu/cloudshell.py new file mode 100644 index 0000000..19d9715 --- /dev/null +++ b/baseline/experiments/distilbert_scratch_torchtpu/cloudshell.py @@ -0,0 +1,234 @@ +"""Launch one bounded run on the existing TPU. Never allocates or extends a lease.""" +import argparse +import datetime as dt +import fcntl +import json +import os +from pathlib import Path +import re +import shlex +import shutil +import subprocess +import time +import urllib.request + +PROJECT = 'tpu-builders-504820' +MOUNT = Path('/mnt/disks/rg-data') +BASE = MOUNT / 'distilbert_scratch_torchtpu' +LATEST = BASE / 'LATEST.json' +FOLDER = 'baseline/experiments/distilbert_scratch_torchtpu' + + +def run(command, **kwargs): + return subprocess.run(command, check=True, text=True, **kwargs) + + +def checked_lease(node, zone): + flags = ['--project=' + PROJECT, '--zone=' + zone, '--format=json'] + spec = json.loads(run(['gcloud', 'compute', 'tpus', 'tpu-vm', 'describe', node, *flags], capture_output=True).stdout) + queue_name = spec.get('queuedResource', '').rsplit('/', 1)[-1] + if not queue_name: + raise RuntimeError('Explicit existing queued-resource lease required') + queue = json.loads(run(['gcloud', 'alpha', 'compute', 'tpus', 'queued-resources', 'describe', queue_name, *flags], capture_output=True).stdout) + if spec.get('state') != 'READY' or spec.get('acceleratorType') != 'v5litepod-8' or queue.get('state', {}).get('state') != 'ACTIVE': + raise RuntimeError('Requires an active single-host v5e-8 allocation') + expiries = [spec.get('schedulingConfig', {}).get('terminationTimestamp'), queue.get('runDuration', {}).get('terminationTime')] + for entry in queue.get('tpu', {}).get('nodeSpec', []): + if entry.get('nodeId') == node: + expiries.append(entry.get('node', {}).get('schedulingConfig', {}).get('terminationTimestamp')) + parsed = [dt.datetime.fromisoformat(re.sub(r'(\.\d{6})\d+', r'\1', v).replace('Z', '+00:00')).timestamp() for v in expiries if v] + if not parsed: + raise RuntimeError('No explicit lease expiry found') + return {'node': node, 'checked_unix': time.time(), 'termination_unix': min(parsed), + 'ips': [v['ipAddress'] for v in spec.get('networkEndpoints', []) if v.get('ipAddress')]} + + +def bounded_deadline(checked, hours): + if not 0 < hours <= 6: + raise ValueError('Budget must be greater than zero and at most six hours') + if not -30 <= time.time() - checked['checked_unix'] <= 600: + raise RuntimeError('Lease check expired; retry start') + deadline = min(time.time() + hours * 3600, checked['termination_unix'] - 300) + if deadline - time.time() < 1800: + raise RuntimeError('Less than 30 minutes remain; nothing launched') + return deadline + + +def discover_node(zone, require_lease=True): + specs = json.loads(run(['gcloud', 'compute', 'tpus', 'tpu-vm', 'list', '--project=' + PROJECT, + '--zone=' + zone, '--format=json'], capture_output=True).stdout) + eligible = [] + reasons = [] + for spec in specs: + if spec.get('state') != 'READY' or spec.get('acceleratorType') != 'v5litepod-8': + continue + node = spec['name'].rsplit('/', 1)[-1] + try: + if require_lease: + checked = checked_lease(node, zone) + bounded_deadline(checked, 5.5) + eligible.append(node) + except Exception as exc: + reasons.append(node + ': ' + str(exc)) + if len(eligible) != 1: + raise RuntimeError('Expected one eligible existing v5e-8 TPU. Specify --node explicitly if several exist. ' + + json.dumps({'eligible': eligible, 'excluded': reasons})) + return eligible[0] + + +def active(unit): + result = subprocess.run(['systemctl', 'is-active', unit], capture_output=True, text=True) + return result.stdout.strip() in ('active', 'activating', 'deactivating', 'reloading') + + +def write(path, value): + tmp = path.with_suffix('.tmp') + with tmp.open('w') as f: + json.dump(value, f, indent=2) + f.flush() + os.fsync(f.fileno()) + tmp.replace(path) + + +def start_remote(a): + if os.geteuid() != 0 or not os.path.ismount(MOUNT): + raise RuntimeError('Root and the existing mounted /mnt/disks/rg-data are required') + if not re.fullmatch('[a-f0-9]{40}', a.commit): + raise ValueError('A pinned 40-character source commit is required') + checked = json.loads(a.lease) + deadline = bounded_deadline(checked, a.hours) + request = urllib.request.Request('http://metadata.google.internal/computeMetadata/v1/instance/network-interfaces/0/ip', + headers={'Metadata-Flavor': 'Google'}) + with urllib.request.urlopen(request, timeout=5) as response: + if response.read().decode().strip() not in checked['ips']: + raise RuntimeError('Guest does not match checked TPU') + python = next((shutil.which(p) for p in ('python3.12', 'python3.13') if shutil.which(p)), None) + if not python: + raise RuntimeError('TorchTPU requires Python >=3.12. Install python3.12 and python3.12-venv first.') + if shutil.disk_usage(MOUNT).free < 30 * 1024**3: + raise RuntimeError('Need at least 30 GiB free; no data deleted') + BASE.mkdir(exist_ok=True) + # Use the same launch lock as the existing GPT-2 controllers. + lockpath = MOUNT / 'gpt2small/port-check-launch.lock' + lockpath.parent.mkdir(exist_ok=True) + with lockpath.open('a') as lock: + fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB) + if LATEST.exists() and active(json.loads(LATEST.read_text())['unit']): + raise RuntimeError('DistilBERT is already active; use status') + # Cover newer RG services as well as the legacy process guard. + units = run(['systemctl', 'list-units', '--type=service', '--state=running,activating', '--no-legend', '--plain', 'rg-*'], capture_output=True).stdout + if units.strip(): + raise RuntimeError('An RG service is already active; nothing stopped or launched:\n' + units) + stamp = dt.datetime.now(dt.timezone.utc).strftime('%Y%m%d-%H%M%S') + root = BASE / ('initial-' + stamp) + root.mkdir() + repo = root / 'repo' + run(['git', 'init', '-q', str(repo)]) + run(['git', '-C', str(repo), 'fetch', '--depth', '1', 'https://github.com/CalculatedContent/rg_optimizers.git', a.commit], timeout=180) + run(['git', '-C', str(repo), 'checkout', '--detach', 'FETCH_HEAD']) + devices = [p for p in Path('/dev').glob('accel*') if p.is_char_device()] + devices += [p for p in Path('/dev/accel').glob('accel*') if p.is_char_device()] + if not devices or not shutil.which('fuser'): + raise RuntimeError('Cannot verify TPU device occupancy (/dev/accel* and fuser required)') + usage = subprocess.run(['fuser', *map(str, devices)], capture_output=True, text=True) + if usage.returncode == 0 or usage.stdout.strip(): + raise RuntimeError('TPU devices are in use; existing processes untouched') + if usage.returncode != 1 or usage.stderr.strip(): + raise RuntimeError('Could not establish that TPU devices are idle: ' + usage.stderr) + if a.resume_run: + resume = Path(a.resume_run).resolve() + if not resume.is_relative_to(BASE.resolve()) or not (resume / 'latest_checkpoint.json').is_file(): + raise RuntimeError('Resume path must be an existing committed run under ' + str(BASE)) + original = json.loads((resume.parent / 'launch.json').read_text()) + if original['commit'] != a.commit: + raise RuntimeError('Resume with the original source commit') + else: + config = (repo / FOLDER / (a.config or 'configs/initial.json')).resolve() + if not config.is_relative_to((repo / FOLDER).resolve()) or config.suffix != '.json' or not config.is_file(): + raise ValueError('Config must be a committed JSON file inside this experiment folder') + bounded_deadline(checked, a.hours) # Recheck freshness without extending the original deadline. + unit = 'rg-distilbert-torchtpu-' + stamp + '.service' + record = {'root': str(root), 'unit': unit, 'commit': a.commit, 'lease': checked, + 'deadline_unix': deadline, 'maximum_hours': a.hours, 'resume_run': a.resume_run, + 'config': a.config or ('original run config' if a.resume_run else 'configs/initial.json'), + 'backend': 'native TorchTPU', 'pretrained_weights': False, + 'automatic_restart': False, 'experiment_count': 1, + 'storage': str(MOUNT), 'test_accessed_at_launch': False, + 'test_monitoring_enabled': True, 'test_used_for_selection': False, + 'test_policy': 'MLM monitoring only; no test-based selection'} + write(root / 'launch.json', record) + env = {'TORCH_TPU_DEFER_AND_FUSE': '1', + 'OMP_NUM_THREADS': '1', 'OPENBLAS_NUM_THREADS': '1', 'MKL_NUM_THREADS': '1', + 'TOKENIZERS_PARALLELISM': 'false', 'HF_HOME': str(BASE / 'hf'), + 'PYTHONUNBUFFERED': '1'} + command = ['systemd-run', '--unit=' + unit, '--property=Type=exec', '--property=Restart=no', + '--property=RuntimeMaxSec=' + str(int(deadline - time.time())), '--property=TimeoutStopSec=30', + '--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 += [str(python), '-u', str(repo / FOLDER / 'supervisor.py'), '--root', str(root), '--deadline', str(deadline)] + if a.resume_run: + command += ['--resume-run', a.resume_run] + elif a.config: + command += ['--config', a.config] + run(command) + write(LATEST, record) + print(json.dumps(record, indent=2)) + + +def status_remote(): + record = json.loads(LATEST.read_text()) + print(json.dumps(record, indent=2)) + root = Path(record['root']) + runroot = Path(record['resume_run']) if record.get('resume_run') else root / 'initial' + for path in (root / 'preflight.json', root / 'ERROR.json', root / 'RESULT.json', runroot / 'runtime.json', runroot / 'status.json'): + if path.exists(): + print(path.name + '\n' + path.read_text()) + subprocess.run(['systemctl', 'status', record['unit'], '--no-pager', '-n', '8']) + + +def main(): + p = argparse.ArgumentParser(description=__doc__) + p.add_argument('action', choices=('start', 'status')) + p.add_argument('--node', help='Existing node; defaults to discovering the sole eligible v5e-8 in --zone') + p.add_argument('--zone', default='us-west4-a') + p.add_argument('--hours', type=float, default=5.5) + mode = p.add_mutually_exclusive_group() + mode.add_argument('--resume-run') + mode.add_argument('--config', help='Committed JSON path relative to this experiment folder') + p.add_argument('--commit') + p.add_argument('--remote', action='store_true', help=argparse.SUPPRESS) + p.add_argument('--lease', help=argparse.SUPPRESS) + a = p.parse_args() + if a.remote: + return start_remote(a) if a.action == 'start' else status_remote() + here = Path(__file__).resolve() + repo = here.parents[3] + if a.action == 'start' and a.config: + config = (here.parent / a.config).resolve() + if not config.is_relative_to(here.parent) or config.suffix != '.json' or not config.is_file(): + raise ValueError('Config must be a committed JSON file inside this experiment folder') + run(['git', '-C', str(repo), 'ls-files', '--error-unmatch', str(config.relative_to(repo))], capture_output=True) + a.config = str(config.relative_to(here.parent)) + if not a.node: + a.node = discover_node(a.zone, require_lease=a.action == 'start') + command = ['sudo', 'python3', '-c', here.read_text(), a.action, '--remote', '--hours', str(a.hours)] + if a.action == 'start': + if run(['git', '-C', str(repo), 'status', '--porcelain'], capture_output=True).stdout.strip(): + raise RuntimeError('Launch from a clean source checkout') + head = run(['git', '-C', str(repo), 'rev-parse', 'HEAD'], capture_output=True).stdout.strip() + if a.commit and a.commit != head: + raise RuntimeError('--commit must equal the clean checkout HEAD') + lease = checked_lease(a.node, a.zone) + bounded_deadline(lease, a.hours) + command += ['--commit', head, '--lease', json.dumps(lease)] + if a.resume_run: + command += ['--resume-run', a.resume_run] + elif a.config: + command += ['--config', a.config] + run(['gcloud', 'compute', 'tpus', 'tpu-vm', 'ssh', a.node, '--project=' + PROJECT, + '--zone=' + a.zone, '--worker=0', '--command=' + shlex.join(command)]) + + +if __name__ == '__main__': + main() diff --git a/baseline/experiments/distilbert_scratch_torchtpu/common.py b/baseline/experiments/distilbert_scratch_torchtpu/common.py new file mode 100644 index 0000000..e28dc5b --- /dev/null +++ b/baseline/experiments/distilbert_scratch_torchtpu/common.py @@ -0,0 +1,198 @@ +"""Independent, reproducible masked-LM protocol. No torch_xla dependency.""" +from dataclasses import asdict, dataclass +import hashlib +import json +import math +import os +from pathlib import Path +import numpy as np +import torch + + +@dataclass(frozen=True) +class Config: + seed: int = 401 + data_seed: int = 1701 + steps: int = 10000 + warmup_steps: int = 1000 + per_device_batch: int = 32 + sequence_length: int = 128 + optimizer: str = 'adamw' + lr: float = 1e-4 + weight_decay: float = 0.01 + beta1: float = 0.9 + beta2: float = 0.999 + eps: float = 1e-6 + momentum: float = 0.9 + grad_clip: float = 1.0 + dropout: float = 0.1 + attention_dropout: float = 0.1 + mask_probability: float = 0.15 + evaluation_mask_probability: float = 0.15 + precision: str = 'bf16' + compile: bool = True + checkpoint_steps: int = 100 + measure_steps: int = 500 + train_eval_sequences: int = 2048 + # Only used by local engineering tests, never the TPU launcher. + synthetic: bool = False + + def validate(self): + sizes = (self.steps, self.warmup_steps, self.per_device_batch, self.sequence_length, + self.checkpoint_steps, self.measure_steps, self.train_eval_sequences) + if not all(type(x) is int for x in (*sizes, self.seed, self.data_seed)): + raise ValueError('Sizes and seeds must be integers') + if not all(0 <= x < 2**32 for x in (self.seed, self.data_seed)): + raise ValueError('Seeds must be unsigned 32-bit integers') + if type(self.compile) is not bool or type(self.synthetic) is not bool: + raise ValueError('compile and synthetic must be booleans') + if self.optimizer not in ('adamw', 'adam', 'sgd'): + raise ValueError('Unknown optimizer') + if not 0 <= self.warmup_steps < self.steps: + raise ValueError('Invalid warmup') + if min(self.steps, self.per_device_batch, self.checkpoint_steps, + self.measure_steps, self.train_eval_sequences) < 1: + raise ValueError('Positive sizes required') + if not 8 <= self.sequence_length <= 512: + raise ValueError('Sequence length must be 8..512') + if not all(math.isfinite(x) for x in (self.lr, self.weight_decay, self.grad_clip, self.eps)) or not ( + self.lr > 0 and self.weight_decay >= 0 and self.grad_clip > 0 and self.eps > 0): + raise ValueError('Invalid optimizer settings') + if not all(0 <= x < 1 for x in (self.beta1, self.beta2, self.momentum, + self.dropout, self.attention_dropout)): + raise ValueError('Invalid probabilities/momenta') + if not all(0 < p < 1 for p in (self.mask_probability, self.evaluation_mask_probability)) or self.precision not in ('bf16', 'fp32'): + raise ValueError('Invalid masking/precision') + return self + + +def read_config(path): + return Config(**json.loads(Path(path).read_text())).validate() + + +def digest(value): + return hashlib.sha256(json.dumps(value, sort_keys=True).encode()).hexdigest() + + +def file_hash(path): + h = hashlib.sha256() + with Path(path).open('rb') as f: + for chunk in iter(lambda: f.read(1024 * 1024), b''): + h.update(chunk) + return h.hexdigest() + + +def sync_dir(path): + fd = os.open(str(path), os.O_DIRECTORY) + try: + os.fsync(fd) + finally: + os.close(fd) + + +def atomic_json(path, value): + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + tmp = path.with_name(path.name + '.tmp') + with tmp.open('w') as f: + json.dump(value, f, indent=2, allow_nan=False) + f.write('\n') + f.flush() + os.fsync(f.fileno()) + tmp.replace(path) + sync_dir(path.parent) + + +def save_torch(path, value): + path = Path(path) + with path.open('wb') as f: + torch.save(value, f) + f.flush() + os.fsync(f.fileno()) + + +def cpu_tree(value): + if isinstance(value, torch.Tensor): + return value.detach().cpu() + if isinstance(value, dict): + return {k: cpu_tree(v) for k, v in value.items()} + if isinstance(value, list): + return [cpu_tree(v) for v in value] + if isinstance(value, tuple): + return tuple(cpu_tree(v) for v in value) + return value + + +def learning_rate(c, step): + """step is the zero-based update about to execute; warmup reaches peak.""" + if step < c.warmup_steps: + return c.lr * (step + 1) / c.warmup_steps + return c.lr * (c.steps - step) / (c.steps - c.warmup_steps) + + +def batch_indices(order, batch_index, rank, world, local_batch): + start = batch_index * world * local_batch + rank * local_batch + positions = np.arange(start, start + local_batch) + valid = positions < len(order) + # Padding never contributes a target, gradient, or metric. + return order[positions % len(order)], valid + + +def masked_batch(tokens, indices, valid, meta, c, salt, step=0, *, evaluation=False): + ids = np.array(tokens[indices], dtype=np.int64, copy=True) + labels = np.full(ids.shape, -100, dtype=np.int64) + special = np.isin(ids, meta['special_ids']) + probability = c.evaluation_mask_probability if evaluation else c.mask_probability + for row, (idx, good) in enumerate(zip(indices, valid)): + if not good: + continue + g = np.random.default_rng(np.random.SeedSequence([c.data_seed, salt, step, int(idx)])) + eligible = np.flatnonzero(~special[row]) + if not len(eligible): + raise ValueError('Prepared sequence has no maskable tokens') + chosen = eligible[g.random(len(eligible)) < probability] + if not len(chosen): + chosen = g.choice(eligible, 1) # documented nonempty-target safeguard + labels[row, chosen] = ids[row, chosen] + kind = g.random(len(chosen)) + ids[row, chosen[kind < 0.8]] = meta['mask_id'] + random_positions = chosen[(kind >= 0.8) & (kind < 0.9)] + ids[row, random_positions] = g.integers(0, meta['vocab_size'], len(random_positions)) + return torch.from_numpy(ids), torch.from_numpy(labels) + + +def build_model(c, vocab_size): + from transformers import DistilBertConfig, DistilBertForMaskedLM + shape = dict(n_layers=2, dim=32, hidden_dim=64, n_heads=4) if c.synthetic else dict( + n_layers=6, dim=768, hidden_dim=3072, n_heads=12) + config = DistilBertConfig(vocab_size=vocab_size, max_position_embeddings=512, + dropout=c.dropout, attention_dropout=c.attention_dropout, + activation='gelu', **shape) + config._attn_implementation = 'eager' + # This constructor initializes every parameter randomly. Never load model weights. + return DistilBertForMaskedLM(config) + + +def make_optimizer(model, c, device): + decay, no_decay = [], [] + for _, p in model.named_parameters(): + (decay if p.ndim >= 2 else no_decay).append(p) + groups = [dict(params=decay, weight_decay=c.weight_decay), + dict(params=no_decay, weight_decay=0.0)] + lr = torch.tensor(c.lr, device=device) if device.type == 'tpu' else c.lr + if c.optimizer in ('adamw', 'adam'): + cls = torch.optim.AdamW if c.optimizer == 'adamw' else torch.optim.Adam + return cls(groups, lr=lr, betas=(c.beta1, c.beta2), eps=c.eps, + foreach=False, capturable=device.type == 'tpu') + return torch.optim.SGD(groups, lr=lr, momentum=c.momentum, foreach=False) + + +def encoder_matrices(model): + return {name: p.detach().float().cpu().clone() for name, p in model.named_parameters() + if name.startswith('distilbert.transformer.layer.') and p.ndim == 2} + + +def source_hash(): + root = Path(__file__).parent + return digest({str(p.relative_to(root)): file_hash(p) for p in sorted(root.rglob('*')) + if p.is_file() and (p.suffix == '.py' or p.name.startswith('requirements'))}) diff --git a/baseline/experiments/distilbert_scratch_torchtpu/configs/initial.json b/baseline/experiments/distilbert_scratch_torchtpu/configs/initial.json new file mode 100644 index 0000000..ae4227a --- /dev/null +++ b/baseline/experiments/distilbert_scratch_torchtpu/configs/initial.json @@ -0,0 +1,26 @@ +{ + "seed": 401, + "data_seed": 1701, + "steps": 10000, + "warmup_steps": 1000, + "per_device_batch": 32, + "sequence_length": 128, + "optimizer": "adamw", + "lr": 0.0001, + "weight_decay": 0.01, + "beta1": 0.9, + "beta2": 0.999, + "eps": 1e-06, + "momentum": 0.9, + "grad_clip": 1.0, + "dropout": 0.1, + "attention_dropout": 0.1, + "mask_probability": 0.15, + "evaluation_mask_probability": 0.15, + "precision": "bf16", + "compile": true, + "checkpoint_steps": 100, + "measure_steps": 500, + "train_eval_sequences": 2048, + "synthetic": false +} diff --git a/baseline/experiments/distilbert_scratch_torchtpu/configs/pilot_lr1e4_warmup300.json b/baseline/experiments/distilbert_scratch_torchtpu/configs/pilot_lr1e4_warmup300.json new file mode 100644 index 0000000..a0e8a1f --- /dev/null +++ b/baseline/experiments/distilbert_scratch_torchtpu/configs/pilot_lr1e4_warmup300.json @@ -0,0 +1,26 @@ +{ + "seed": 401, + "data_seed": 1701, + "steps": 10000, + "warmup_steps": 300, + "per_device_batch": 32, + "sequence_length": 128, + "optimizer": "adamw", + "lr": 0.0001, + "weight_decay": 0.01, + "beta1": 0.9, + "beta2": 0.999, + "eps": 1e-06, + "momentum": 0.9, + "grad_clip": 1.0, + "dropout": 0.1, + "attention_dropout": 0.1, + "mask_probability": 0.15, + "evaluation_mask_probability": 0.15, + "precision": "bf16", + "compile": true, + "checkpoint_steps": 100, + "measure_steps": 500, + "train_eval_sequences": 2048, + "synthetic": false +} diff --git a/baseline/experiments/distilbert_scratch_torchtpu/configs/pilot_lr3e4_warmup1000.json b/baseline/experiments/distilbert_scratch_torchtpu/configs/pilot_lr3e4_warmup1000.json new file mode 100644 index 0000000..25b2be6 --- /dev/null +++ b/baseline/experiments/distilbert_scratch_torchtpu/configs/pilot_lr3e4_warmup1000.json @@ -0,0 +1,26 @@ +{ + "seed": 401, + "data_seed": 1701, + "steps": 10000, + "warmup_steps": 1000, + "per_device_batch": 32, + "sequence_length": 128, + "optimizer": "adamw", + "lr": 0.0003, + "weight_decay": 0.01, + "beta1": 0.9, + "beta2": 0.999, + "eps": 1e-06, + "momentum": 0.9, + "grad_clip": 1.0, + "dropout": 0.1, + "attention_dropout": 0.1, + "mask_probability": 0.15, + "evaluation_mask_probability": 0.15, + "precision": "bf16", + "compile": true, + "checkpoint_steps": 100, + "measure_steps": 500, + "train_eval_sequences": 2048, + "synthetic": false +} diff --git a/baseline/experiments/distilbert_scratch_torchtpu/configs/pilot_lr3e4_warmup300.json b/baseline/experiments/distilbert_scratch_torchtpu/configs/pilot_lr3e4_warmup300.json new file mode 100644 index 0000000..12c75cb --- /dev/null +++ b/baseline/experiments/distilbert_scratch_torchtpu/configs/pilot_lr3e4_warmup300.json @@ -0,0 +1,26 @@ +{ + "seed": 401, + "data_seed": 1701, + "steps": 10000, + "warmup_steps": 300, + "per_device_batch": 32, + "sequence_length": 128, + "optimizer": "adamw", + "lr": 0.0003, + "weight_decay": 0.01, + "beta1": 0.9, + "beta2": 0.999, + "eps": 1e-06, + "momentum": 0.9, + "grad_clip": 1.0, + "dropout": 0.1, + "attention_dropout": 0.1, + "mask_probability": 0.15, + "evaluation_mask_probability": 0.15, + "precision": "bf16", + "compile": true, + "checkpoint_steps": 100, + "measure_steps": 500, + "train_eval_sequences": 2048, + "synthetic": false +} diff --git a/baseline/experiments/distilbert_scratch_torchtpu/configs/pilot_lr5e4_warmup1000.json b/baseline/experiments/distilbert_scratch_torchtpu/configs/pilot_lr5e4_warmup1000.json new file mode 100644 index 0000000..d8aff8c --- /dev/null +++ b/baseline/experiments/distilbert_scratch_torchtpu/configs/pilot_lr5e4_warmup1000.json @@ -0,0 +1,26 @@ +{ + "seed": 401, + "data_seed": 1701, + "steps": 10000, + "warmup_steps": 1000, + "per_device_batch": 32, + "sequence_length": 128, + "optimizer": "adamw", + "lr": 0.0005, + "weight_decay": 0.01, + "beta1": 0.9, + "beta2": 0.999, + "eps": 1e-06, + "momentum": 0.9, + "grad_clip": 1.0, + "dropout": 0.1, + "attention_dropout": 0.1, + "mask_probability": 0.15, + "evaluation_mask_probability": 0.15, + "precision": "bf16", + "compile": true, + "checkpoint_steps": 100, + "measure_steps": 500, + "train_eval_sequences": 2048, + "synthetic": false +} diff --git a/baseline/experiments/distilbert_scratch_torchtpu/configs/pilot_lr5e4_warmup300.json b/baseline/experiments/distilbert_scratch_torchtpu/configs/pilot_lr5e4_warmup300.json new file mode 100644 index 0000000..4f1f9d0 --- /dev/null +++ b/baseline/experiments/distilbert_scratch_torchtpu/configs/pilot_lr5e4_warmup300.json @@ -0,0 +1,26 @@ +{ + "seed": 401, + "data_seed": 1701, + "steps": 10000, + "warmup_steps": 300, + "per_device_batch": 32, + "sequence_length": 128, + "optimizer": "adamw", + "lr": 0.0005, + "weight_decay": 0.01, + "beta1": 0.9, + "beta2": 0.999, + "eps": 1e-06, + "momentum": 0.9, + "grad_clip": 1.0, + "dropout": 0.1, + "attention_dropout": 0.1, + "mask_probability": 0.15, + "evaluation_mask_probability": 0.15, + "precision": "bf16", + "compile": true, + "checkpoint_steps": 100, + "measure_steps": 500, + "train_eval_sequences": 2048, + "synthetic": false +} diff --git a/baseline/experiments/distilbert_scratch_torchtpu/launch.py b/baseline/experiments/distilbert_scratch_torchtpu/launch.py new file mode 100644 index 0000000..5a94cf2 --- /dev/null +++ b/baseline/experiments/distilbert_scratch_torchtpu/launch.py @@ -0,0 +1,34 @@ +"""Configure native TorchTPU topology in the parent, then spawn eight workers.""" +import os +from pathlib import Path +import subprocess +import socket +import sys + + +def main(): + import torch + import torch_tpu # noqa: F401 + # Official launch helper must run before any TPU tensor or process group. + torch.tpu.distributed.set_tpu_launch_env(nnodes=1) + # The 0.2.0 Python helper sets topology/addresses; only its CLI prints + # WORLD_SIZE. Derive the count from the actual detected device addresses. + addresses = os.environ.get('TORCH_TPU_SLICEBUILDER_ADDRESSES', '') + count = len([x for x in addresses.split(',') if x]) + if count != 8 or int(os.environ.get('WORLD_SIZE', '8')) != 8: + raise RuntimeError('This baseline requires an existing eight-device slice') + os.environ['WORLD_SIZE'] = '8' + if int(os.environ.get('NODE_RANK', '0')) != 0: + raise RuntimeError('This launcher is for the existing single-host v5e-8 only') + with socket.socket() as sock: + sock.bind(('127.0.0.1', 0)) + port = sock.getsockname()[1] + cmd = [sys.executable, '-m', 'torch.distributed.run', '--master-addr=127.0.0.1', + f'--master-port={port}', + '--nnodes=1', '--nproc-per-node=8', '--max-restarts=0', + str(Path(__file__).with_name('train.py')), '--expected-world=8', *sys.argv[1:]] + subprocess.run(cmd, check=True) + + +if __name__ == '__main__': + main() diff --git a/baseline/experiments/distilbert_scratch_torchtpu/prepare.py b/baseline/experiments/distilbert_scratch_torchtpu/prepare.py new file mode 100644 index 0000000..dd00994 --- /dev/null +++ b/baseline/experiments/distilbert_scratch_torchtpu/prepare.py @@ -0,0 +1,86 @@ +"""Pack WikiText-103 within each official split and freeze upstream/data hashes.""" +import argparse +import json +from pathlib import Path +import numpy as np +from common import atomic_json, file_hash + +TOKENIZER = 'distilbert/distilbert-base-uncased' +TOKENIZER_REVISION = '12040accade4e8a0f71eabdb258fecc2e7e948be' +DATASET = 'Salesforce/wikitext' +DATASET_REVISION = 'b08601e04326c79dfdd32d625aee71d232d685c3' + + +def prepare(root, length=128, dataset_revision=None): + from datasets import load_dataset + from huggingface_hub import HfApi + from transformers import AutoTokenizer + root = Path(root) + if root.exists() and any(root.iterdir()): + raise RuntimeError('Preparation requires an empty destination; existing data preserved') + root.mkdir(parents=True, exist_ok=True) + revision = HfApi().dataset_info(DATASET, revision=dataset_revision or DATASET_REVISION).sha + tok = AutoTokenizer.from_pretrained(TOKENIZER, revision=TOKENIZER_REVISION) + tok.save_pretrained(root / 'tokenizer') + ds = load_dataset(DATASET, 'wikitext-103-raw-v1', revision=revision) + meta = dict(dataset=DATASET, dataset_config='wikitext-103-raw-v1', dataset_revision=revision, + tokenizer=TOKENIZER, tokenizer_revision=TOKENIZER_REVISION, + sequence_length=length, vocab_size=len(tok), mask_id=tok.mask_token_id, + special_ids=tok.all_special_ids, synthetic=False, splits={}, + packing=f'[CLS] + {length - 2} content tokens + [SEP]; separator between nonempty lines; no cross-split packing') + for split in ('train', 'validation', 'test'): + raw = root / (split + '.bin') + buf, rows, text_tokens = [], 0, 0 + with raw.open('wb') as f: + for batch in ds[split].iter(batch_size=1024): + text = [s for s in batch['text'] if s.strip()] + if not text: + continue + for ids in tok(text, add_special_tokens=False, truncation=False)['input_ids']: + if not ids: + continue + text_tokens += len(ids) + buf.extend(ids + [tok.sep_token_id]) + end = len(buf) // (length - 2) * (length - 2) + if end: + mid = np.asarray(buf[:end], dtype=np.int32).reshape(-1, length - 2) + blocks = np.empty((len(mid), length), dtype=np.int32) + blocks[:, 0], blocks[:, -1] = tok.cls_token_id, tok.sep_token_id + blocks[:, 1:-1] = mid + # Exclude any all-special block before freezing the dataset. + blocks = blocks[(~np.isin(blocks, tok.all_special_ids)).any(axis=1)] + blocks.tofile(f) + rows += len(blocks) + buf = buf[end:] + if not rows: + raise RuntimeError('No prepared rows in ' + split) + meta['splits'][split] = dict(rows=rows, text_tokens=text_tokens, + dropped_remainder=len(buf), file=raw.name, sha256=file_hash(raw)) + atomic_json(root / 'prepared.json', meta) + return meta + + +def synthetic(root, length=16): + root = Path(root) + root.mkdir(parents=True, exist_ok=True) + meta = dict(sequence_length=length, vocab_size=128, mask_id=4, + special_ids=[0, 1, 2, 3, 4], synthetic=True, splits={}) + for split, count in [('train', 19), ('validation', 11), ('test', 13)]: + ids = np.tile(np.arange(length) % 7 + 5, (count, 1)).astype(np.int32) + ids[:, 0], ids[:, -1] = 1, 2 + path = root / (split + '.bin') + ids.tofile(path) + meta['splits'][split] = dict(rows=count, file=path.name, sha256=file_hash(path)) + atomic_json(root / 'prepared.json', meta) + return meta + + +if __name__ == '__main__': + p = argparse.ArgumentParser(description=__doc__) + p.add_argument('--output', required=True) + p.add_argument('--length', type=int, default=128) + p.add_argument('--dataset-revision') + a = p.parse_args() + if not 8 <= a.length <= 512: + p.error('length must be 8..512') + print(json.dumps(prepare(a.output, a.length, a.dataset_revision), indent=2)) diff --git a/baseline/experiments/distilbert_scratch_torchtpu/requirements.txt b/baseline/experiments/distilbert_scratch_torchtpu/requirements.txt new file mode 100644 index 0000000..3935609 --- /dev/null +++ b/baseline/experiments/distilbert_scratch_torchtpu/requirements.txt @@ -0,0 +1,5 @@ +# Keep this stack separate from every existing PyTorch/XLA environment. +torch_tpu==0.2.0 +transformers==4.57.1 +datasets==4.2.0 +weightwatcher==0.7.7 diff --git a/baseline/experiments/distilbert_scratch_torchtpu/spectra.py b/baseline/experiments/distilbert_scratch_torchtpu/spectra.py new file mode 100644 index 0000000..9b06162 --- /dev/null +++ b/baseline/experiments/distilbert_scratch_torchtpu/spectra.py @@ -0,0 +1,112 @@ +"""CPU WeightWatcher fits; encoder matrices only; no feedback into training.""" +import argparse +import csv +import importlib.metadata +import json +import math +from pathlib import Path +import time +import numpy as np +import torch +from common import atomic_json, file_hash + + +def summarize(rows, expected): + valid = [r for r in rows if r['valid']] + complete = len(rows) == expected and len(valid) == expected + alphas = np.array([r['alpha'] for r in valid]) + return dict(expected_layers=expected, valid_layers=len(valid), all_fits_valid=complete, + mean_alpha=float(alphas.mean()) if complete else None, + min_alpha=float(alphas.min()) if complete else None, + fraction_alpha_below_2=float((alphas < 2).mean()) if complete else None, + any_alpha_below_2=bool((alphas < 2).any()) if complete else None, + mean_D=float(np.mean([r['D'] for r in valid])) if complete else None) + + +def analyze(root): + import weightwatcher as ww + version = importlib.metadata.version('weightwatcher') + if version != '0.7.7': + raise RuntimeError('This protocol pins WeightWatcher 0.7.7') + root = Path(root) + torch.set_num_threads(1) + for path in sorted((root / 'measurements').glob('step_*.json')): + target = root / 'spectra' / path.name + measured = json.loads(path.read_text()) + if target.exists(): + continue + source = root / measured['snapshot'] + if file_hash(source) != measured['sha256']: + raise RuntimeError('Snapshot hash mismatch') + matrices = torch.load(source, map_location='cpu', weights_only=True) + rows = [] + for name, w in matrices.items(): + row = dict(name=name, shape=list(w.shape), valid=False) + try: + layer = torch.nn.Linear(w.shape[1], w.shape[0], bias=False) + layer.weight.data.copy_(w) + # Identical default WW estimator to the existing RG experiment. + df = ww.WeightWatcher(model=torch.nn.Sequential(layer)).analyze() + fits = json.loads(df.to_json(orient='records')) + fit = fits[0] if fits else {} + row['fit'] = fit + for key in ('alpha', 'D', 'sigma', 'xmin', 'xmax', 'num_pl_spikes', 'num_evals', 'warning', 'status'): + row[key] = fit.get(key) + row['valid'] = (isinstance(row['alpha'], (float, int)) and math.isfinite(row['alpha']) + and row['alpha'] > 0 and isinstance(row['D'], (float, int)) + and 0 <= row['D'] <= 1 and fit.get('status') == 'success') + n, total = row['num_pl_spikes'], row['num_evals'] + row['tail_fraction'] = n / total if n is not None and total else None + except Exception as e: + row['error'] = repr(e) + rows.append(row) + atomic_json(target, {**measured, **summarize(rows, measured['expected_layers']), + 'layers': rows, 'weightwatcher': version, + 'aggregation': 'Unweighted mean of all 36 raw encoder matrix fits; null if any invalid', + 'excluded': 'Embeddings, tied vocabulary projection and MLM prediction transform', + 'interpretation': 'One trajectory cannot establish optimizer bias or overfitting'}) + export(root) + + +def export(root): + models, layers = [], [] + for path in sorted((root / 'spectra').glob('step_*.json')): + item = json.loads(path.read_text()) + base = {k: item[k] for k in ('step', 'mean_alpha', 'min_alpha', 'mean_D', + 'fraction_alpha_below_2', 'any_alpha_below_2', 'valid_layers', 'expected_layers')} + for split in ('train', 'validation', 'test'): + for metric in ('mlm_loss', 'masked_token_accuracy'): + base[split + '_' + metric] = item[split][metric] + models.append(base) + for row in item['layers']: + layers.append({**base, **{k: row.get(k) for k in ('name', 'valid', 'alpha', 'D', + 'sigma', 'xmin', 'xmax', 'num_pl_spikes', 'num_evals', 'tail_fraction', 'status', 'warning')}}) + for name, rows in [('checkpoint_metrics.csv', models), ('layer_metrics.csv', layers)]: + if rows: + path = root / name + temp = path.with_suffix('.tmp') + with temp.open('w') as f: + writer = csv.DictWriter(f, fieldnames=list(rows[0])) + writer.writeheader() + writer.writerows(rows) + temp.replace(path) + + +def watch(root, deadline=float('inf'), watching=True): + while time.time() < deadline: + # Observe the completion marker BEFORE scanning. If training finishes + # during this scan, another scan must consume its final measurement. + done = (Path(root) / 'TRAINING_DONE.json').exists() + analyze(root) + if not watching or done: + break + time.sleep(min(15, max(0, deadline - time.time()))) + + +if __name__ == '__main__': + p = argparse.ArgumentParser(description=__doc__) + p.add_argument('--root', required=True) + p.add_argument('--watch', action='store_true') + p.add_argument('--deadline', type=float, default=float('inf')) + args = p.parse_args() + watch(args.root, args.deadline, args.watch) diff --git a/baseline/experiments/distilbert_scratch_torchtpu/supervisor.py b/baseline/experiments/distilbert_scratch_torchtpu/supervisor.py new file mode 100644 index 0000000..e8c6364 --- /dev/null +++ b/baseline/experiments/distilbert_scratch_torchtpu/supervisor.py @@ -0,0 +1,136 @@ +"""Bounded install, probe, two real updates, resume, and CPU spectral analysis.""" +import argparse +import json +import os +from pathlib import Path +import signal +import subprocess +import sys +import time +import traceback + + +def write(path, value): + temp = path.with_suffix('.tmp') + temp.write_text(json.dumps(value, indent=2) + '\n') + temp.replace(path) + + +def terminate(proc, grace=20): + """Reap the owned process group, including children that ignore SIGTERM.""" + for sig in (signal.SIGTERM, signal.SIGKILL): + try: + os.killpg(proc.pid, sig) + except ProcessLookupError: + pass + if sig == signal.SIGTERM: + try: + proc.wait(timeout=grace) + except subprocess.TimeoutExpired: + pass + else: + proc.wait() + + +def run(command, deadline, env=None): + if deadline <= time.time(): + raise RuntimeError('Budget exhausted before next phase') + proc = subprocess.Popen([str(x) for x in command], start_new_session=True, env=env) + try: + code = proc.wait(timeout=max(1, deadline - time.time())) + except subprocess.TimeoutExpired: + terminate(proc) + raise RuntimeError('Phase exceeded its bounded deadline') + if code: + terminate(proc) + raise RuntimeError(f'Phase failed ({code}): {command}') + + +def main(a): + root = Path(a.root) + here = Path(__file__).resolve().parent + training_deadline = a.deadline - 600 + if a.resume_run: + output = Path(a.resume_run) + # Resume uses the original environment, not a freshly resolved package stack. + python = output.parent / 'venv/bin/python' + config = output.parent / 'config.json' + data = Path(json.loads((output.parent / 'data_path.json').read_text())['path']) + if not python.is_file(): + raise RuntimeError('Original run environment is required for exact resume') + else: + if sys.version_info < (3, 12): + raise RuntimeError('Python >=3.12 required') + venv = root / 'venv' + run([sys.executable, '-m', 'venv', venv], training_deadline) + python = venv / 'bin/python' + run([python, '-m', 'pip', 'install', '--upgrade', 'pip'], training_deadline) + run([python, '-m', 'pip', 'install', '-r', here / 'requirements.txt', + '--extra-index-url', 'https://download.pytorch.org/whl/cpu'], training_deadline) + run([python, '-m', 'pip', 'check'], training_deadline) + frozen = subprocess.check_output([python, '-m', 'pip', 'freeze'], text=True) + (root / 'requirements.lock.txt').write_text(frozen) + config = root / 'config.json' + requested = (here / (a.config or 'configs/initial.json')).resolve() + if not requested.is_relative_to(here) or requested.suffix != '.json': + raise ValueError('Config must be a JSON file inside this experiment folder') + config.write_text(requested.read_text()) + data = root / 'data' + output = root / 'initial' + run([python, here / 'launch.py', '--probe-only', '--output', root / 'probe'], training_deadline) + if not a.resume_run: + length = json.loads(config.read_text()).get('sequence_length', 128) + run([python, here / 'prepare.py', '--output', data, '--length', str(length)], training_deadline) + write(root / 'data_path.json', dict(path=str(data))) + if training_deadline - time.time() < 900: + raise RuntimeError('Less than 15 minutes remain after setup; data saved, training not started') + cmd = [python, here / 'launch.py', '--config', config, '--data', data, + '--output', output, '--deadline', str(training_deadline)] + if not a.resume_run: + run([*cmd, '--stop-after=2'], a.deadline - 300) + state = json.loads((output / 'status.json').read_text()) + runtime = json.loads((output / 'runtime.json').read_text()) + if state['step'] != 2 or not state['checkpoint_saved'] or runtime['world'] != 8: + raise RuntimeError('Two-update hardware preflight did not pass') + write(root / 'preflight.json', dict(passed=True, saved_updates=2, world=8, + initialization='random', pretrained_weights=False)) + # The analyst never touches TPU tensors or changes optimization. + (output / 'TRAINING_DONE.json').unlink(missing_ok=True) + analyst_env = dict(os.environ, TORCH_DEVICE_BACKEND_AUTOLOAD='0', OMP_NUM_THREADS='1', + OPENBLAS_NUM_THREADS='1', MKL_NUM_THREADS='1') + with (root / 'spectra.log').open('a') as log: + analyst = subprocess.Popen([str(python), str(here / 'spectra.py'), '--root', str(output), + '--watch', '--deadline', str(a.deadline - 60)], env=analyst_env, + stdout=log, stderr=subprocess.STDOUT, start_new_session=True) + try: + run([*cmd, '--resume'], a.deadline - 300) + write(output / 'TRAINING_DONE.json', dict(done=True)) + try: + analyst.wait(timeout=max(1, a.deadline - time.time() - 60)) + except subprocess.TimeoutExpired: + terminate(analyst) + result = json.loads((output / 'RESULT.json').read_text()) + result['spectra_exit_code'] = analyst.returncode + result['spectra_complete'] = (analyst.returncode == 0 and + {p.name for p in (output / 'measurements').glob('step_*.json')} == + {p.name for p in (output / 'spectra').glob('step_*.json')}) + result['output'] = str(output) + write(root / 'RESULT.json', result) + finally: + if analyst.poll() is None: + terminate(analyst) + + +if __name__ == '__main__': + p = argparse.ArgumentParser(description=__doc__) + p.add_argument('--root', required=True) + p.add_argument('--deadline', type=float, required=True) + mode = p.add_mutually_exclusive_group() + mode.add_argument('--resume-run') + mode.add_argument('--config', help='JSON path relative to this experiment folder; default configs/initial.json') + args = p.parse_args() + try: + main(args) + except Exception: + write(Path(args.root) / 'ERROR.json', dict(error=traceback.format_exc(), time=time.time())) + raise diff --git a/baseline/experiments/distilbert_scratch_torchtpu/test_experiment.py b/baseline/experiments/distilbert_scratch_torchtpu/test_experiment.py new file mode 100644 index 0000000..434c9c8 --- /dev/null +++ b/baseline/experiments/distilbert_scratch_torchtpu/test_experiment.py @@ -0,0 +1,372 @@ +"""Engineering checks on tiny random models; never evidence about alpha < 2.""" +from dataclasses import asdict, replace +import copy +import json +import os +from pathlib import Path +import subprocess +import socket +import sys +import tempfile +from types import SimpleNamespace +import unittest +from unittest.mock import Mock, patch +import numpy as np +import torch +from common import (Config, atomic_json, batch_indices, masked_batch, build_model, + encoder_matrices, file_hash, make_optimizer, read_config) +from cloudshell import discover_node +from prepare import synthetic +from spectra import summarize, watch +from supervisor import terminate +from train import LossModel, Runtime, checkpoint, preserve_best_validation, train + +HERE = Path(__file__).parent + + +def run_job(root, name, c, world=1, resume=False, stop=None): + config = root / (name + '.json') + atomic_json(config, asdict(c)) + cmd = [sys.executable] + if world > 1: + with socket.socket() as sock: + sock.bind(('127.0.0.1', 0)) + port = sock.getsockname()[1] + cmd += ['-m', 'torch.distributed.run', '--master-addr=127.0.0.1', + f'--master-port={port}', '--nnodes=1', f'--nproc-per-node={world}'] + cmd += [str(HERE / 'train.py'), '--backend=cpu', f'--expected-world={world}', + '--config', str(config), '--data', str(root / 'data'), '--output', str(root / name)] + if resume: + cmd += ['--resume'] + if stop is not None: + cmd += ['--stop-after', str(stop)] + env = dict(os.environ, OMP_NUM_THREADS='1', OPENBLAS_NUM_THREADS='1', MKL_NUM_THREADS='1') + result = subprocess.run(cmd, env=env, capture_output=True, text=True, timeout=240) + if result.returncode: + raise AssertionError(result.stdout + result.stderr) + latest = json.loads((root / name / 'latest_checkpoint.json').read_text()) + return torch.load(root / name / latest['path'] / 'state.pt', weights_only=False) + + +class ProtocolTests(unittest.TestCase): + def config(self, **kwargs): + return replace(Config(), synthetic=True, sequence_length=16, per_device_batch=4, + steps=4, warmup_steps=1, measure_steps=2, checkpoint_steps=2, + train_eval_sequences=11, precision='fp32', compile=False, **kwargs) + + def test_eight_shards_exactly_once(self): + order = np.arange(37) + seen = [] + for j in range(2): + for rank in range(8): + idx, valid = batch_indices(order, j, rank, 8, 4) + self.assertEqual(len(idx), 4) + seen.extend(idx[valid]) + self.assertEqual(sorted(seen), list(range(37))) + + def test_masking_fixed_per_example_no_special_or_padding_targets(self): + with tempfile.TemporaryDirectory() as d: + meta = synthetic(d) + x = np.memmap(Path(d) / 'train.bin', dtype=np.int32, shape=(19, 16)) + c = self.config() + ids, labels = masked_batch(x, [2, 3], [True, False], meta, c, 992) + solo, one = masked_batch(x, [2], [True], meta, c, 992) + torch.testing.assert_close(ids[0], solo[0]) + torch.testing.assert_close(labels[0], one[0]) + self.assertTrue((labels[1] == -100).all()) + self.assertTrue((labels[:, [0, -1]] == -100).all()) + + def test_missing_fit_does_not_bias_mean(self): + rows = [dict(valid=True, alpha=1.7, D=.04), dict(valid=False)] + self.assertIsNone(summarize(rows, 2)['mean_alpha']) + self.assertFalse(summarize(rows, 2)['all_fits_valid']) + + def test_invalid_config_is_rejected_before_training(self): + for change in (dict(lr=float('inf')), dict(eps=float('nan')), + dict(grad_clip=float('inf')), dict(weight_decay=float('inf')), + dict(steps=4.5), dict(seed=-1), dict(seed=2**32), + dict(data_seed=-1), dict(compile='false'), dict(synthetic=1)): + with self.subTest(change=change), self.assertRaises(ValueError): + replace(self.config(), **change).validate() + + def test_tpu_rng_uses_supported_getter_and_restores_all_states(self): + rt = Runtime.__new__(Runtime) + rt.backend = 'tpu' + fake_tpu = SimpleNamespace(get_rng_state=Mock(side_effect=RuntimeError('Internal API')), + set_rng_state=Mock()) + value = torch.tensor([12, 34], dtype=torch.uint8) + getter = Mock(return_value=value) + with patch.object(torch, 'tpu', fake_tpu, create=True), patch.object( + torch.accelerator, 'random', SimpleNamespace(get_rng_state=getter), create=True): + saved = rt.rng() + torch.rand(8) + np.random.random(8) + rt.restore_rng(saved) + getter.assert_called_once_with() + fake_tpu.get_rng_state.assert_not_called() + fake_tpu.set_rng_state.assert_called_once_with(value) + torch.testing.assert_close(torch.get_rng_state(), saved['torch'], rtol=0, atol=0) + self.assertTrue(np.array_equal(np.random.get_state()[1], saved['numpy'][1])) + # Older supported torch releases lack the replacement getter. + fake_tpu.get_rng_state = Mock(return_value=value) + with patch.object(torch, 'tpu', fake_tpu, create=True), patch.object( + torch.accelerator, 'random', SimpleNamespace(), create=True): + torch.testing.assert_close(rt.rng()['tpu'], value, rtol=0, atol=0) + fake_tpu.get_rng_state.assert_called_once_with() + + def test_spectral_worker_drains_measurement_published_during_last_scan(self): + with tempfile.TemporaryDirectory() as d: + root, seen = Path(d), [] + def scan(_): + if (root / 'final.json').exists(): + seen.append('final') + else: + # Training publishes just after this scan's file listing. + (root / 'final.json').write_text('{}') + (root / 'TRAINING_DONE.json').write_text('{}') + with patch('spectra.analyze', side_effect=scan) as analyze, patch('spectra.time.sleep'): + watch(root) + self.assertEqual(seen, ['final']) + self.assertEqual(analyze.call_count, 2) + + def test_status_can_discover_existing_tpu_near_lease_expiry(self): + listed = SimpleNamespace(stdout=json.dumps([ + dict(name='projects/p/locations/z/nodes/existing', state='READY', + acceleratorType='v5litepod-8')])) + with patch('cloudshell.run', return_value=listed), patch( + 'cloudshell.checked_lease', side_effect=RuntimeError('Lease expired')) as lease: + self.assertEqual(discover_node('z', require_lease=False), 'existing') + lease.assert_not_called() + with self.assertRaisesRegex(RuntimeError, 'Expected one eligible'): + discover_node('z') + + def test_timeout_kills_process_that_ignores_sigterm(self): + code = ('import signal,time; signal.signal(signal.SIGTERM, signal.SIG_IGN); ' + 'print("ready", flush=True); time.sleep(60)') + proc = subprocess.Popen([sys.executable, '-c', code], start_new_session=True, + stdout=subprocess.PIPE, text=True) + try: + self.assertEqual(proc.stdout.readline().strip(), 'ready') + terminate(proc, grace=.05) + self.assertEqual(proc.returncode, -9) + finally: + if proc.poll() is None: + terminate(proc, grace=.05) + proc.stdout.close() + + def test_unpublished_checkpoint_is_replaced_and_partial_does_not_affect_retention(self): + with tempfile.TemporaryDirectory() as d: + root = Path(d) + model = build_model(self.config(), 128) + opt = make_optimizer(model, self.config(), torch.device('cpu')) + rt = Runtime('cpu', 1) + checkpoint(root, 0, model, opt, rt, {}) + orphan = root / 'checkpoints/step_00000001' + orphan.mkdir() + (orphan / 'state.pt').write_bytes(b'unpublished previous replay') + (root / 'checkpoints/step_00000003').mkdir() # Future unpublished generation. + (root / 'checkpoints/step_99999999.partial').mkdir() + checkpoint(root, 1, model, opt, rt, {}) + saved = torch.load(orphan / 'state.pt', weights_only=False) + self.assertEqual(saved['step'], 1) + self.assertTrue((root / 'checkpoints/step_00000000').exists()) + checkpoint(root, 2, model, opt, rt, {}) + self.assertFalse((root / 'checkpoints/step_00000000').exists()) + self.assertTrue(orphan.exists()) + + def test_measurements_never_precede_checkpoint_and_interrupted_measurement_recovers(self): + with tempfile.TemporaryDirectory() as d: + root = Path(d) + synthetic(root / 'data') + c = replace(self.config(), checkpoint_steps=3) + atomic_json(root / 'config.json', asdict(c)) + args = SimpleNamespace(config=root / 'config.json', data=root / 'data', + output=root / 'run', resume=False, stop_after=None, + deadline=float('inf')) + import train as training + real_measure = training.measure + observed, interrupted = [], False + def checked_measure(runroot, step, *values): + nonlocal interrupted + pointer = json.loads((runroot / 'latest_checkpoint.json').read_text()) + self.assertEqual(pointer['step'], step) + observed.append(step) + if step == 2 and not interrupted: + interrupted = True + raise RuntimeError('Simulated interruption before diagnostic publication') + return real_measure(runroot, step, *values) + with patch('train.measure', side_effect=checked_measure): + with self.assertRaisesRegex(RuntimeError, 'Simulated interruption'): + train(args, Runtime('cpu', 1)) + args.resume = True + train(args, Runtime('cpu', 1)) + self.assertEqual(observed, [0, 2, 2, 4]) + self.assertTrue((root / 'run/measurements/step_00000002.json').exists()) + + def test_training_mask_rate_does_not_change_evaluation(self): + with tempfile.TemporaryDirectory() as d: + meta = synthetic(d) + tokens = np.memmap(Path(d) / 'train.bin', dtype=np.int32, shape=(19, 16)) + c = self.config() + indices, valid = np.arange(19), np.ones(19, dtype=bool) + higher = replace(c, mask_probability=.30) + for salt in (991, 992, 993): + low = masked_batch(tokens, indices, valid, meta, c, salt, evaluation=True) + high = masked_batch(tokens, indices, valid, meta, higher, salt, evaluation=True) + for a, b in zip(low, high): + torch.testing.assert_close(a, b, rtol=0, atol=0) + low = masked_batch(tokens, indices, valid, meta, c, 901)[1] + high = masked_batch(tokens, indices, valid, meta, higher, 901)[1] + self.assertGreater((high != -100).sum(), (low != -100).sum()) + + def test_best_weights_follow_validation_only_and_survive_later_regression(self): + with tempfile.TemporaryDirectory() as d: + root = Path(d) + model = build_model(self.config(), 128) + metrics = dict(validation=dict(mlm_loss=3.0), test=dict(mlm_loss=4.0)) + preserve_best_validation(root, 1, model, metrics, {'run': 'test'}) + first = json.loads((root / 'best_validation.json').read_text()) + with torch.no_grad(): + model.distilbert.transformer.layer[0].attention.q_lin.weight.add_(.1) + metrics['validation']['mlm_loss'], metrics['test']['mlm_loss'] = 3.1, 1.0 + preserve_best_validation(root, 2, model, metrics, {'run': 'test'}) + self.assertEqual(first, json.loads((root / 'best_validation.json').read_text())) + metrics['validation']['mlm_loss'], metrics['test']['mlm_loss'] = 2.9, 5.0 + preserve_best_validation(root, 3, model, metrics, {'run': 'test'}) + best = json.loads((root / 'best_validation.json').read_text()) + self.assertEqual(best['step'], 3) + self.assertEqual(file_hash(root / best['path']), best['sha256']) + saved = torch.load(root / best['path'], weights_only=False) + for name, tensor in model.state_dict().items(): + torch.testing.assert_close(saved['model'][name], tensor, rtol=0, atol=0) + self.assertEqual(len(list((root / 'best_validation').glob('*.pt'))), 1) + preserve_best_validation(root, 4, model, metrics, {'run': 'test'}) + self.assertEqual(best, json.loads((root / 'best_validation.json').read_text())) + + def test_adamw_decay_preserves_norm_bias_and_does_not_duplicate_tied_weights(self): + c = self.config() + model = build_model(c, 128) + opt = make_optimizer(model, c, torch.device('cpu')) + parameters = [p for g in opt.param_groups for p in g['params']] + self.assertEqual(len(parameters), len({id(p) for p in parameters})) + with torch.no_grad(): + for p in parameters: + p.fill_(1.) + p.grad = torch.zeros_like(p) + opt.step() + block = model.distilbert.transformer.layer[0] + self.assertTrue(torch.equal(block.sa_layer_norm.weight, torch.ones_like(block.sa_layer_norm.weight))) + self.assertTrue(torch.equal(block.attention.q_lin.bias, torch.ones_like(block.attention.q_lin.bias))) + for p in (block.attention.q_lin.weight, model.vocab_projector.weight): + torch.testing.assert_close(p, torch.full_like(p, 1 - c.lr * c.weight_decay), rtol=0, atol=0) + + def test_compiled_bf16_loss_and_gradients_match_eager(self): + c = self.config(dropout=0, attention_dropout=0) + torch.manual_seed(42) + eager = LossModel(build_model(c, 128)) + compiled = torch.compile(copy.deepcopy(eager), backend='aot_eager', fullgraph=True, dynamic=False) + ids = torch.randint(5, 128, (4, 16)) + labels = ids.clone() + labels[:, :8], labels[-1] = -100, -100 + outputs = [] + for net in (eager, compiled): + with torch.autocast('cpu', dtype=torch.bfloat16): + out = net(ids, labels) + (out[0] / out[2]).backward() + outputs.append(out) + for a, b in zip(*outputs): + torch.testing.assert_close(a, b, rtol=0, atol=0) + for a, b in zip(eager.parameters(), compiled.parameters()): + torch.testing.assert_close(a.grad, b.grad, rtol=1e-5, atol=1e-6) + + def test_pilot_grid_changes_only_lr_and_warmup(self): + baseline = read_config(HERE / 'configs/initial.json') + configs = [baseline] + [read_config(p) for p in sorted((HERE / 'configs').glob('pilot_*.json'))] + self.assertEqual(len(configs), 6) + self.assertEqual({(c.lr, c.warmup_steps) for c in configs}, + {(lr, warmup) for lr in (1e-4, 3e-4, 5e-4) for warmup in (300, 1000)}) + for c in configs: + self.assertEqual(replace(c, lr=baseline.lr, warmup_steps=baseline.warmup_steps), baseline) + + def test_full_architecture_and_random_init(self): + torch.manual_seed(401) + model = build_model(Config(), 30522) + self.assertEqual(sum(p.numel() for p in model.parameters()), 66985530) + self.assertEqual(len(encoder_matrices(model)), 36) + self.assertEqual(model.distilbert.transformer.layer[0].attention.q_lin.weight.shape, (768, 768)) + self.assertIs(model.vocab_projector.weight, model.distilbert.embeddings.word_embeddings.weight) + + def assert_resume_matches(self, c): + with tempfile.TemporaryDirectory() as d: + root = Path(d) + synthetic(root / 'data') + full = run_job(root, 'full', c) + run_job(root, 'resume', c, stop=2) + resumed = run_job(root, 'resume', c, resume=True) + self.assertEqual(full['step'], resumed['step']) + for name in full['model']: + torch.testing.assert_close(full['model'][name], resumed['model'][name], rtol=0, atol=0) + for key, state in full['optimizer']['state'].items(): + for field, value in state.items(): + torch.testing.assert_close(value, resumed['optimizer']['state'][key][field], rtol=0, atol=0) + for step in (2, 4): + a = json.loads((root / 'full/measurements' / f'step_{step:08d}.json').read_text()) + b = json.loads((root / 'resume/measurements' / f'step_{step:08d}.json').read_text()) + for split in ('train', 'validation', 'test'): + self.assertEqual(a[split], b[split]) + # Corrupt the checkpoint and verify that resume refuses it. + latest = json.loads((root / 'resume/latest_checkpoint.json').read_text()) + with (root / 'resume' / latest['path'] / 'state.pt').open('ab') as f: + f.write(b'corruption') + with self.assertRaisesRegex(AssertionError, 'Checkpoint hash mismatch'): + run_job(root, 'resume', c, resume=True) + + def test_resume_matches_uninterrupted_for_all_optimizers(self): + for name in ('adamw', 'adam', 'sgd'): + with self.subTest(optimizer=name): + self.assert_resume_matches(self.config(optimizer=name)) + + def test_compiled_bf16_resume_matches_uninterrupted(self): + self.assert_resume_matches(replace(self.config(), compile=True, precision='bf16')) + + def test_eight_shard_gradients_match_serial_token_mean(self): + # In-process gradient oracle also runs where networking is restricted. + with tempfile.TemporaryDirectory() as d: + meta = synthetic(d) + tokens = np.memmap(Path(d) / 'train.bin', dtype=np.int32, shape=(19, 16)) + c = self.config(dropout=0, attention_dropout=0) + torch.manual_seed(c.seed) + serial = LossModel(build_model(c, meta['vocab_size'])) + sharded = copy.deepcopy(serial) + ids, labels = masked_batch(tokens, np.arange(19), np.ones(19, dtype=bool), meta, c, 901) + loss, _, count = serial(ids, labels) + (loss / count).backward() + gradients = {n: torch.zeros_like(p) for n, p in sharded.named_parameters()} + for rank in range(8): + idx, valid = batch_indices(np.arange(19), 0, rank, 8, 4) + x, y = masked_batch(tokens, idx, valid, meta, c, 901) + sharded.zero_grad(set_to_none=True) + rank_loss, _, _ = sharded(x, y) + (rank_loss * 8 / count).backward() + for n, p in sharded.named_parameters(): + gradients[n] += p.grad / 8 + for n, p in serial.named_parameters(): + torch.testing.assert_close(gradients[n], p.grad, rtol=1e-4, atol=2e-6) + + @unittest.skipUnless(os.environ.get('RG_RUN_DISTRIBUTED_TESTS') == '1', + 'Opt-in real Gloo test requires socket/interface permissions') + def test_distributed_matches_global_token_weighted_update(self): + with tempfile.TemporaryDirectory() as d: + root = Path(d) + synthetic(root / 'data') + c = self.config(optimizer='sgd', momentum=0, dropout=0, attention_dropout=0) + serial = run_job(root, 'serial', replace(c, per_device_batch=8)) + distributed = run_job(root, 'distributed', c, world=2) + # Includes the uneven final batch (19 examples, global batch 8). + for name in serial['model']: + torch.testing.assert_close(serial['model'][name], distributed['model'][name], rtol=1e-5, atol=2e-7) + + +if __name__ == '__main__': + torch.set_num_threads(1) + unittest.main() diff --git a/baseline/experiments/distilbert_scratch_torchtpu/train.py b/baseline/experiments/distilbert_scratch_torchtpu/train.py new file mode 100644 index 0000000..da48d9a --- /dev/null +++ b/baseline/experiments/distilbert_scratch_torchtpu/train.py @@ -0,0 +1,390 @@ +"""Random-init DistilBERT MLM on native TorchTPU; CPU mode is engineering-only.""" +import argparse +from dataclasses import asdict +import datetime as dt +import importlib.metadata +import json +import math +import os +from pathlib import Path +import random +import shutil +import time +import numpy as np +import torch +import torch.distributed as dist +import torch.nn.functional as F +from torch.nn.parallel import DistributedDataParallel as DDP +from common import (atomic_json, batch_indices, build_model, cpu_tree, digest, + encoder_matrices, file_hash, learning_rate, make_optimizer, + masked_batch, read_config, save_torch, source_hash, sync_dir) + + +class LossModel(torch.nn.Module): + def __init__(self, core): + super().__init__() + self.core = core + + def forward(self, ids, labels): + # All prepared blocks are full length; there are no padding tokens to mask. + logits = self.core(input_ids=ids).logits.float() + loss_sum = F.cross_entropy(logits.flatten(0, 1), labels.flatten(), + ignore_index=-100, reduction='sum') + keep = labels != -100 + correct = ((logits.argmax(-1) == labels) & keep).sum() + return loss_sum, correct, keep.sum() + + +class Runtime: + def __init__(self, backend, expected): + self.backend = backend + if backend == 'tpu': + import torch_tpu # noqa: F401; registers native PyTorch backend + if int(os.environ.get('WORLD_SIZE', '0')) != expected: + raise RuntimeError('Use launch.py; expected exactly eight TorchTPU workers') + dist.init_process_group('tpu_dist', timeout=dt.timedelta(minutes=20)) + self.device = torch.device('tpu') # each process owns its assigned device + torch.set_float32_matmul_precision('medium') + else: + if int(os.environ.get('WORLD_SIZE', '1')) > 1: + dist.init_process_group('gloo', timeout=dt.timedelta(minutes=5)) + self.device = torch.device('cpu') + self.world = dist.get_world_size() if dist.is_initialized() else 1 + self.rank = dist.get_rank() if dist.is_initialized() else 0 + if self.world != expected: + raise RuntimeError(f'Expected {expected} workers; got {self.world}') + + def sum(self, x): + if self.world > 1: + dist.all_reduce(x, op=dist.ReduceOp.SUM) + return x + + def barrier(self): + if self.world > 1: + dist.barrier() + + def rng(self): + state = dict(torch=torch.get_rng_state(), python=random.getstate(), numpy=np.random.get_state()) + if self.backend == 'tpu': + # TorchTPU 0.2 makes the legacy getter internal on torch >=2.13. + api = getattr(torch.accelerator, 'random', None) + getter = getattr(api, 'get_rng_state', None) + state['tpu'] = (getter() if getter is not None else torch.tpu.get_rng_state()).cpu() + return state + + def restore_rng(self, state): + torch.set_rng_state(state['torch']) + random.setstate(state['python']) + np.random.set_state(state['numpy']) + if self.backend == 'tpu': + torch.tpu.set_rng_state(state['tpu']) + + def stop(self, deadline): + flag = torch.tensor(float(time.time() >= deadline), device=self.device) + return self.sum(flag).item() > 0 + + +def probe(rt, output): + # Matmul and collectives on every worker, before any model/data allocation. + x = torch.full((128, 128), float(rt.rank + 1), device=rt.device) + y = x @ torch.ones_like(x) + value = rt.sum(y[0, 0].clone()).cpu().item() + expected = 128 * rt.world * (rt.world + 1) / 2 + if value != expected: + raise RuntimeError(f'Collective probe failed: {value} != {expected}') + state = rt.rng() + sample = torch.rand(16, device=rt.device).cpu() + rt.restore_rng(state) + replay = torch.rand(16, device=rt.device).cpu() + rt.restore_rng(state) + if not torch.equal(sample, replay): + raise RuntimeError('RNG checkpoint/restore probe failed') + root = Path(output) + root.mkdir(parents=True, exist_ok=True) + atomic_json(root / f'rank_{rt.rank}.json', dict(rank=rt.rank, world=rt.world, + backend=rt.backend, device=str(rt.device), collective_value=value, rng_roundtrip=True, + torch=torch.__version__, torch_tpu=importlib.metadata.version('torch_tpu') if rt.backend == 'tpu' else None)) + rt.barrier() + if rt.rank == 0: + records = [json.loads((root / f'rank_{i}.json').read_text()) for i in range(rt.world)] + atomic_json(root / 'probe.json', dict(passed=True, ranks=records)) + + +def checkpoint(root, step, model, opt, rt, identity): + checkpoints = root / 'checkpoints' + target = checkpoints / f'step_{step:08d}' + partial = checkpoints / f'step_{step:08d}.partial' + if rt.rank == 0: + checkpoints.mkdir(exist_ok=True) + if partial.exists(): + shutil.rmtree(partial) + partial.mkdir() + rt.barrier() + save_torch(partial / f'rng_{rt.rank}.pt', rt.rng()) + if rt.rank == 0: + save_torch(partial / 'state.pt', dict(step=step, identity=identity, + model=cpu_tree(model.state_dict()), optimizer=cpu_tree(opt.state_dict()))) + rt.barrier() + if rt.rank == 0: + atomic_json(partial / 'manifest.json', {p.name: file_hash(p) for p in partial.glob('*.pt')}) + pointer = root / 'latest_checkpoint.json' + published = json.loads(pointer.read_text())['path'] if pointer.exists() else None + # Reuse only a published generation. A crash between rename and pointer + # publication can leave an orphan from a different replay of this step. + if target.exists() and published == str(target.relative_to(root)): + shutil.rmtree(partial) + else: + if target.exists(): + shutil.rmtree(target) + partial.rename(target) + sync_dir(checkpoints) + atomic_json(root / 'latest_checkpoint.json', dict(path=str(target.relative_to(root)), step=step)) + # Future generations left by an interrupted replay were never published + # on this trajectory; they must not displace either retained checkpoint. + committed = sorted((p for p in checkpoints.glob('step_*') + if p.is_dir() and p.name[5:].isdigit() and int(p.name[5:]) <= step), + key=lambda p: int(p.name[5:])) + for old in committed[:-2]: + shutil.rmtree(old) + rt.barrier() + + +def restore(root, model, opt, rt, identity): + target = root / json.loads((root / 'latest_checkpoint.json').read_text())['path'] + manifest = json.loads((target / 'manifest.json').read_text()) + if set(manifest) != {'state.pt', *(f'rng_{i}.pt' for i in range(rt.world))}: + raise RuntimeError('Incomplete checkpoint') + for name, sha in manifest.items(): + if file_hash(target / name) != sha: + raise RuntimeError('Checkpoint hash mismatch: ' + name) + # Only load files created by this experiment, after verifying their manifest. + state = torch.load(target / 'state.pt', map_location='cpu', weights_only=False) + if state['identity'] != identity: + raise RuntimeError('Resume requires identical source, config, prepared data, packages and world size') + model.load_state_dict(state['model']) + opt.load_state_dict(state['optimizer']) + if rt.backend == 'tpu': + for group in opt.param_groups: + group['lr'] = torch.as_tensor(group['lr'], device=rt.device) + rt.restore_rng(torch.load(target / f'rng_{rt.rank}.pt', map_location='cpu', weights_only=False)) + return state['step'] + + +def evaluate(net, arrays, meta, c, rt): + state = rt.rng() + net.eval() + results = {} + with torch.no_grad(): + for name, salt in [('train', 991), ('validation', 992), ('test', 993)]: + order = np.arange(len(arrays[name])) + if name == 'train': + order = np.random.default_rng(c.data_seed).permutation(order)[:c.train_eval_sequences] + totals = torch.zeros(3, device=rt.device) + batches = math.ceil(len(order) / (rt.world * c.per_device_batch)) + for j in range(batches): + idx, valid = batch_indices(order, j, rt.rank, rt.world, c.per_device_batch) + ids, labels = masked_batch(arrays[name], idx, valid, meta, c, salt, evaluation=True) + with torch.autocast(rt.device.type, dtype=torch.bfloat16, + enabled=c.precision == 'bf16'): + loss_sum, correct, count = net(ids.to(rt.device), labels.to(rt.device)) + totals += torch.stack((loss_sum.detach(), correct.float(), count.float())) + nll, correct, count = rt.sum(totals).cpu().tolist() + if count <= 0 or not math.isfinite(nll): + raise RuntimeError('Invalid evaluation') + results[name] = dict(mlm_loss=nll / count, masked_token_accuracy=correct / count, + masked_tokens=int(count), sequences=len(order)) + rt.restore_rng(state) + net.train() + return results + + +def preserve_best_validation(root, step, model, results, identity): + """Retain full weights selected only by the fixed validation MLM loss.""" + root = Path(root) + pointer = root / 'best_validation.json' + loss = results['validation']['mlm_loss'] + if not math.isfinite(loss): + raise ValueError('Best-checkpoint selection requires finite validation loss') + if pointer.exists() and loss >= json.loads(pointer.read_text())['validation_mlm_loss']: + return + folder = root / 'best_validation' + folder.mkdir(exist_ok=True) + target = folder / f'step_{step:08d}.pt' + partial = target.with_suffix('.partial') + save_torch(partial, dict(step=step, identity=identity, + model=cpu_tree(model.state_dict()), model_config=model.config.to_dict(), + validation_mlm_loss=loss)) + partial.replace(target) + sync_dir(folder) + atomic_json(pointer, dict(step=step, validation_mlm_loss=loss, + path=str(target.relative_to(root)), sha256=file_hash(target), + criterion='minimum fixed-mask validation MLM loss; first checkpoint wins ties', + test_used_for_selection=False, alpha_used_for_selection=False, + purpose='full model weights for evaluation; use checkpoints/ for exact training resume')) + for old in folder.glob('step_*.pt'): + if old != target: + old.unlink() + + +def measure(root, step, net, model, arrays, meta, c, rt, identity): + results = evaluate(net, arrays, meta, c, rt) + if rt.rank == 0: + preserve_best_validation(root, step, model, results, identity) + snapshots = root / 'spectral_snapshots' + snapshots.mkdir(exist_ok=True) + path = snapshots / f'step_{step:08d}.pt' + mats = encoder_matrices(model) + expected = 12 if c.synthetic else 36 + if len(mats) != expected: + raise RuntimeError(f'Expected {expected} encoder matrices; got {len(mats)}') + save_torch(path, mats) + atomic_json(root / 'measurements' / f'step_{step:08d}.json', dict( + step=step, **results, synthetic=c.synthetic, + validation_train_loss_gap=results['validation']['mlm_loss'] - results['train']['mlm_loss'], + snapshot=str(path.relative_to(root)), sha256=file_hash(path), expected_layers=expected, + test_used_for_selection=False, evaluation_mask_probability=c.evaluation_mask_probability, + metric='masked-token accuracy, not downstream classification')) + rt.barrier() + + +def train(a, rt): + c = read_config(a.config) + if rt.backend == 'cpu' and not c.synthetic: + raise RuntimeError('CPU execution is only for synthetic engineering checks') + if rt.backend == 'tpu' and c.synthetic: + raise RuntimeError('Production TPU launch requires real data and full DistilBERT') + data, root = Path(a.data), Path(a.output) + meta = json.loads((data / 'prepared.json').read_text()) + if bool(meta['synthetic']) != c.synthetic or meta['sequence_length'] != c.sequence_length: + raise RuntimeError('Data/config mismatch') + arrays = {} + for split, item in meta['splits'].items(): + path = data / item['file'] + if file_hash(path) != item['sha256']: + raise RuntimeError('Prepared data hash mismatch') + arrays[split] = np.memmap(path, dtype=np.int32, mode='r', shape=(item['rows'], c.sequence_length)) + packages = {k: importlib.metadata.version(k) for k in ('torch', 'transformers', 'numpy')} + if rt.backend == 'tpu': + packages['torch_tpu'] = importlib.metadata.version('torch_tpu') + packages['libtpu'] = importlib.metadata.version('libtpu') + identity = dict(config=asdict(c), data=digest(meta), source=source_hash(), world=rt.world, + backend=rt.backend, packages=packages) + if rt.rank == 0: + root.mkdir(parents=True, exist_ok=True) + if (root / 'identity.json').exists(): + if not a.resume: + raise RuntimeError('Existing run: use --resume') + if json.loads((root / 'identity.json').read_text()) != identity: + raise RuntimeError('Run identity mismatch') + else: + if a.resume: + raise RuntimeError('No run to resume') + atomic_json(root / 'identity.json', identity) + rt.barrier() + random.seed(c.seed) + np.random.seed(c.seed) + torch.manual_seed(c.seed) + model = build_model(c, meta['vocab_size']).to(rt.device) + opt = make_optimizer(model, c, rt.device) + wrapped = LossModel(model) + if c.compile: + wrapped = torch.compile(wrapped, backend='tpu' if rt.backend == 'tpu' else 'aot_eager', dynamic=False) + net = DDP(wrapped, broadcast_buffers=False) if rt.world > 1 else wrapped + # Same initialization, independent per-rank dropout; data/masks have a separate seed. + torch.manual_seed(c.seed + 100000 + rt.rank) + step = restore(root, model, opt, rt, identity) if a.resume else 0 + if rt.rank == 0: + atomic_json(root / 'runtime.json', dict(world=rt.world, backend=rt.backend, resumed_step=step, + experiment='BERT-6 MLM from scratch', architecture='DistilBertForMaskedLM', + distillation=False, + parameters=sum(p.numel() for p in model.parameters()), pretrained_weights=False, + global_batch=c.per_device_batch * rt.world, packages=packages, config=asdict(c))) + if step == 0: + checkpoint(root, step, model, opt, rt, identity) + # A crash after committing a measurement cursor must not skip its diagnostics + # on resume. Publish diagnostics only after the corresponding state is saved. + if (step == 0 or step % c.measure_steps == 0) and not ( + root / 'measurements' / f'step_{step:08d}.json').exists(): + measure(root, step, net, model, arrays, meta, c, rt, identity) + batches = math.ceil(len(arrays['train']) / (rt.world * c.per_device_batch)) + epoch, order = None, None + start, start_step = time.monotonic(), step + last_measured = step if (root / 'measurements' / f'step_{step:08d}.json').exists() else -1 + net.train() + limit = min(c.steps, a.stop_after) if a.stop_after is not None else c.steps + while step < limit and not rt.stop(a.deadline): + current_epoch, batch = divmod(step, batches) + if epoch != current_epoch: + epoch = current_epoch + order = np.random.default_rng(c.data_seed + epoch).permutation(len(arrays['train'])) + idx, valid = batch_indices(order, batch, rt.rank, rt.world, c.per_device_batch) + ids, labels = masked_batch(arrays['train'], idx, valid, meta, c, 901, step) + ids, labels = ids.to(rt.device), labels.to(rt.device) + global_count = rt.sum((labels != -100).sum().float()) + lr = learning_rate(c, step) + for group in opt.param_groups: + if isinstance(group['lr'], torch.Tensor): + group['lr'].copy_(torch.tensor(lr, device=rt.device)) + else: + group['lr'] = lr + opt.zero_grad(set_to_none=True) + with torch.autocast(rt.device.type, dtype=torch.bfloat16, enabled=c.precision == 'bf16'): + loss_sum, correct, count = net(ids, labels) + # DDP averages gradients; this yields the exact GLOBAL masked-token mean, + # even for uneven mask counts and padded final batches. + loss = loss_sum * rt.world / global_count.clamp_min(1) + loss.backward() + norm = torch.nn.utils.clip_grad_norm_(model.parameters(), c.grad_clip, foreach=False) + totals = rt.sum(torch.stack((loss_sum.detach(), correct.float(), count.float(), norm.detach().float()))) + values = totals.cpu().tolist() + if values[2] <= 0 or not all(math.isfinite(x) for x in values): + raise RuntimeError('Nonfinite loss/gradients or no targets; no update applied') + opt.step() + step += 1 + if rt.rank == 0: + status = dict(step=step, target_steps=c.steps, epoch=current_epoch, lr=lr, + mlm_loss=values[0] / values[2], masked_token_accuracy=values[1] / values[2], + mean_gradient_norm=values[3] / rt.world, + elapsed_seconds=time.monotonic() - start, + steps_this_process=step - start_step, state='training') + atomic_json(root / 'status.json', status) + if step <= 3 or step % 10 == 0: + print(json.dumps(status), flush=True) + if step % c.checkpoint_steps == 0 or step % c.measure_steps == 0: + checkpoint(root, step, model, opt, rt, identity) + if step % c.measure_steps == 0: + measure(root, step, net, model, arrays, meta, c, rt, identity) + last_measured = step + # Commit training state before potentially expensive final diagnostics. + checkpoint(root, step, model, opt, rt, identity) + if last_measured != step and not rt.stop(a.deadline): + measure(root, step, net, model, arrays, meta, c, rt, identity) + if rt.rank == 0: + atomic_json(root / 'status.json', dict(step=step, target_steps=c.steps, + state='completed' if step == c.steps else 'paused', checkpoint_saved=True)) + atomic_json(root / 'RESULT.json', dict(step=step, completed=step == c.steps, + synthetic=c.synthetic, hardware=rt.backend, reason='step_limit' if step == limit else 'deadline')) + + +if __name__ == '__main__': + p = argparse.ArgumentParser(description=__doc__) + p.add_argument('--backend', choices=['tpu', 'cpu'], default='tpu') + p.add_argument('--expected-world', type=int, default=8) + p.add_argument('--probe-only', action='store_true') + p.add_argument('--config') + p.add_argument('--data') + p.add_argument('--output', required=True) + p.add_argument('--resume', action='store_true') + p.add_argument('--stop-after', type=int) + p.add_argument('--deadline', type=float, default=float('inf')) + args = p.parse_args() + torch.set_num_threads(1) + runtime = Runtime(args.backend, args.expected_world) + try: + if args.probe_only: + probe(runtime, args.output) + else: + train(args, runtime) + finally: + if dist.is_initialized(): + dist.destroy_process_group()