From 896fa646b94872f8cf45b092127c822fc2880c2e Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Tue, 20 Oct 2026 16:55:21 -0700 Subject: [PATCH 1/3] Add isolated random-init DistilBERT MLM pilot on native TorchTPU --- .../distilbert_scratch_torchtpu/.gitignore | 4 + .../distilbert_scratch_torchtpu/README.md | 196 ++++++++++ .../VALIDATION.json | 70 ++++ .../distilbert_scratch_torchtpu/cloudshell.py | 216 +++++++++++ .../distilbert_scratch_torchtpu/common.py | 187 ++++++++++ .../configs/initial.json | 25 ++ .../distilbert_scratch_torchtpu/launch.py | 34 ++ .../distilbert_scratch_torchtpu/prepare.py | 83 +++++ .../requirements.txt | 5 + .../distilbert_scratch_torchtpu/spectra.py | 105 ++++++ .../distilbert_scratch_torchtpu/supervisor.py | 120 +++++++ .../test_experiment.py | 153 ++++++++ .../distilbert_scratch_torchtpu/train.py | 335 ++++++++++++++++++ 13 files changed, 1533 insertions(+) create mode 100644 baseline/experiments/distilbert_scratch_torchtpu/.gitignore create mode 100644 baseline/experiments/distilbert_scratch_torchtpu/README.md create mode 100644 baseline/experiments/distilbert_scratch_torchtpu/VALIDATION.json create mode 100644 baseline/experiments/distilbert_scratch_torchtpu/cloudshell.py create mode 100644 baseline/experiments/distilbert_scratch_torchtpu/common.py create mode 100644 baseline/experiments/distilbert_scratch_torchtpu/configs/initial.json create mode 100644 baseline/experiments/distilbert_scratch_torchtpu/launch.py create mode 100644 baseline/experiments/distilbert_scratch_torchtpu/prepare.py create mode 100644 baseline/experiments/distilbert_scratch_torchtpu/requirements.txt create mode 100644 baseline/experiments/distilbert_scratch_torchtpu/spectra.py create mode 100644 baseline/experiments/distilbert_scratch_torchtpu/supervisor.py create mode 100644 baseline/experiments/distilbert_scratch_torchtpu/test_experiment.py create mode 100644 baseline/experiments/distilbert_scratch_torchtpu/train.py 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/README.md b/baseline/experiments/distilbert_scratch_torchtpu/README.md new file mode 100644 index 0000000..aef6353 --- /dev/null +++ b/baseline/experiments/distilbert_scratch_torchtpu/README.md @@ -0,0 +1,196 @@ +# DistilBERT 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. +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 | Resolve dataset revision once; save SHA, 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-8 | +| 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 | At most 327,680,000 sequence positions; exact full-corpus passes depend on prepared row count | +| 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 and cooperative exit; retain two committed generations | +| 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. +Runtime is measured by the actual TPU preflight. No overnight completion estimate +is asserted before that measurement. + +## 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 and all-reduce 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. + +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 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. +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/`, `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. + +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 states/metrics, corruption rejection, and in-process eight-shard gradients versus a serial masked-token mean. 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. + +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..bf26b3e --- /dev/null +++ b/baseline/experiments/distilbert_scratch_torchtpu/VALIDATION.json @@ -0,0 +1,70 @@ +{ + "date": "2026-10-20", + "status": "CPU engineering validated; real TPU qualification pending", + "base_commit": "afd96a379f4b22c9199f58ca1442a89c5e087700", + "source_sha256": "b9f365baab28af883cbc0953eab41f697275aa4e3b928e1c5470f0391160f99a", + "suite": { + "tests_run": 7, + "passed": 6, + "skipped": 1, + "seconds": 53.117 + }, + "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 training equals two updates plus restart exactly: parameters, optimizer state, evaluation metrics; checkpoint corruption refused" + ], + "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" + ] + }, + "pending": [ + "all-eight-device native TorchTPU collective probe", + "real full-model compiled forward/backward and optimizer steps", + "checkpoint/restart on actual TPU", + "WikiText-103 download/tokenization on target VM", + "throughput, convergence, and spectral results" + ] +} diff --git a/baseline/experiments/distilbert_scratch_torchtpu/cloudshell.py b/baseline/experiments/distilbert_scratch_torchtpu/cloudshell.py new file mode 100644 index 0000000..e3f61e7 --- /dev/null +++ b/baseline/experiments/distilbert_scratch_torchtpu/cloudshell.py @@ -0,0 +1,216 @@ +"""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): + 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: + 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') + 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, + '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] + 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) + p.add_argument('--resume-run') + 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 not a.node: + a.node = discover_node(a.zone) + 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] + 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..872ccca --- /dev/null +++ b/baseline/experiments/distilbert_scratch_torchtpu/common.py @@ -0,0 +1,187 @@ +"""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-8 + momentum: float = 0.9 + grad_clip: float = 1.0 + dropout: float = 0.1 + attention_dropout: float = 0.1 + 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): + 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 (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 0 < self.mask_probability < 1 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): + 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']) + 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)) < c.mask_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..72d73a6 --- /dev/null +++ b/baseline/experiments/distilbert_scratch_torchtpu/configs/initial.json @@ -0,0 +1,25 @@ +{ + "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-08, + "momentum": 0.9, + "grad_clip": 1.0, + "dropout": 0.1, + "attention_dropout": 0.1, + "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..5873a76 --- /dev/null +++ b/baseline/experiments/distilbert_scratch_torchtpu/prepare.py @@ -0,0 +1,83 @@ +"""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' + + +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).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='[CLS] + 126 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()] + 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..1f14abe --- /dev/null +++ b/baseline/experiments/distilbert_scratch_torchtpu/spectra.py @@ -0,0 +1,105 @@ +"""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) + + +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() + while time.time() < args.deadline: + analyze(args.root) + if not args.watch or (Path(args.root) / 'TRAINING_DONE.json').exists(): + break + time.sleep(15) diff --git a/baseline/experiments/distilbert_scratch_torchtpu/supervisor.py b/baseline/experiments/distilbert_scratch_torchtpu/supervisor.py new file mode 100644 index 0000000..8d405d7 --- /dev/null +++ b/baseline/experiments/distilbert_scratch_torchtpu/supervisor.py @@ -0,0 +1,120 @@ +"""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 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: + os.killpg(proc.pid, signal.SIGTERM) + try: + proc.wait(timeout=20) + except subprocess.TimeoutExpired: + os.killpg(proc.pid, signal.SIGKILL) + proc.wait() + raise RuntimeError('Phase exceeded its bounded deadline') + if code: + 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' + config.write_text((here / 'configs/initial.json').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: + run([python, here / 'prepare.py', '--output', data], 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: + os.killpg(analyst.pid, signal.SIGTERM) + analyst.wait(timeout=20) + result = json.loads((output / 'RESULT.json').read_text()) + result['spectra_exit_code'] = analyst.returncode + result['spectra_complete'] = (analyst.returncode == 0 and + len(list((output / 'measurements').glob('step_*.json'))) == + len(list((output / 'spectra').glob('step_*.json')))) + result['output'] = str(output) + write(root / 'RESULT.json', result) + finally: + if analyst.poll() is None: + os.killpg(analyst.pid, signal.SIGTERM) + analyst.wait(timeout=20) + + +if __name__ == '__main__': + p = argparse.ArgumentParser(description=__doc__) + p.add_argument('--root', required=True) + p.add_argument('--deadline', type=float, required=True) + p.add_argument('--resume-run') + 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..78ac43a --- /dev/null +++ b/baseline/experiments/distilbert_scratch_torchtpu/test_experiment.py @@ -0,0 +1,153 @@ +"""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 +import unittest +import numpy as np +import torch +from common import Config, atomic_json, batch_indices, masked_batch, build_model, encoder_matrices +from prepare import synthetic +from spectra import summarize +from train import LossModel + +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_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 test_resume_matches_uninterrupted_adamw(self): + with tempfile.TemporaryDirectory() as d: + root = Path(d) + synthetic(root / 'data') + c = self.config() + 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_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..d6d5f2f --- /dev/null +++ b/baseline/experiments/distilbert_scratch_torchtpu/train.py @@ -0,0 +1,335 @@ +"""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': + state['tpu'] = 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}') + 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, + 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')}) + # A generation is immutable. Re-saving the same cursor leaves it intact. + if target.exists(): + shutil.rmtree(partial) + else: + partial.rename(target) + sync_dir(checkpoints) + atomic_json(root / 'latest_checkpoint.json', dict(path=str(target.relative_to(root)), step=step)) + for old in sorted(checkpoints.glob('step_*'))[:-2]: + if old.is_dir() and not old.name.endswith('.partial'): + 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) + 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 measure(root, step, net, model, arrays, meta, c, rt): + results = evaluate(net, arrays, meta, c, rt) + if rt.rank == 0: + 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, 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, + 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: + measure(root, step, net, model, arrays, meta, c, rt) + checkpoint(root, step, model, opt, 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.measure_steps == 0: + measure(root, step, net, model, arrays, meta, c, rt) + last_measured = step + if step % c.checkpoint_steps == 0: + checkpoint(root, step, model, opt, rt, identity) + # 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) + 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() From c2511497856c7367d4e2f4c3100963d77fd67f3d Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Tue, 20 Oct 2026 19:21:18 -0700 Subject: [PATCH 2/3] Audit BERT-6 pretraining and make validation-based calibration reproducible --- .../distilbert_scratch_torchtpu/AUDIT.md | 112 ++++++++++++++++++ .../distilbert_scratch_torchtpu/README.md | 62 ++++++++-- .../VALIDATION.json | 69 +++++++++-- .../distilbert_scratch_torchtpu/cloudshell.py | 19 ++- .../distilbert_scratch_torchtpu/common.py | 10 +- .../configs/initial.json | 3 +- .../configs/pilot_lr1e4_warmup300.json | 26 ++++ .../configs/pilot_lr3e4_warmup1000.json | 26 ++++ .../configs/pilot_lr3e4_warmup300.json | 26 ++++ .../configs/pilot_lr5e4_warmup1000.json | 26 ++++ .../configs/pilot_lr5e4_warmup300.json | 26 ++++ .../distilbert_scratch_torchtpu/prepare.py | 7 +- .../distilbert_scratch_torchtpu/supervisor.py | 12 +- .../test_experiment.py | 90 +++++++++++++- .../distilbert_scratch_torchtpu/train.py | 44 ++++++- 15 files changed, 524 insertions(+), 34 deletions(-) create mode 100644 baseline/experiments/distilbert_scratch_torchtpu/AUDIT.md create mode 100644 baseline/experiments/distilbert_scratch_torchtpu/configs/pilot_lr1e4_warmup300.json create mode 100644 baseline/experiments/distilbert_scratch_torchtpu/configs/pilot_lr3e4_warmup1000.json create mode 100644 baseline/experiments/distilbert_scratch_torchtpu/configs/pilot_lr3e4_warmup300.json create mode 100644 baseline/experiments/distilbert_scratch_torchtpu/configs/pilot_lr5e4_warmup1000.json create mode 100644 baseline/experiments/distilbert_scratch_torchtpu/configs/pilot_lr5e4_warmup300.json diff --git a/baseline/experiments/distilbert_scratch_torchtpu/AUDIT.md b/baseline/experiments/distilbert_scratch_torchtpu/AUDIT.md new file mode 100644 index 0000000..09e0e31 --- /dev/null +++ b/baseline/experiments/distilbert_scratch_torchtpu/AUDIT.md @@ -0,0 +1,112 @@ +# 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. | + +## 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 index aef6353..7ec8b41 100644 --- a/baseline/experiments/distilbert_scratch_torchtpu/README.md +++ b/baseline/experiments/distilbert_scratch_torchtpu/README.md @@ -1,4 +1,4 @@ -# DistilBERT from scratch with native Google TorchTPU +# 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 @@ -8,6 +8,9 @@ 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. @@ -21,23 +24,35 @@ run does not establish optimizer bias or overfitting. | 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 | Resolve dataset revision once; save SHA, packed-file SHA256s, row counts, tokenizer revision | +| 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-8 | +| 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 | At most 327,680,000 sequence positions; exact full-corpus passes depend on prepared row count | +| 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 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. -Runtime is measured by the actual TPU preflight. No overnight completion estimate -is asserted before that measurement. +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 @@ -75,6 +90,21 @@ The job performs: 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 @@ -103,6 +133,13 @@ uncommitted; the most recent complete checkpoint is the resume point. A hard kil 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. +`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 @@ -122,6 +159,8 @@ 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. @@ -137,7 +176,7 @@ of targets, not by a mean of unequal per-worker means. Outputs in `initial/`: - `status.json`, `runtime.json`, `identity.json`, `RESULT.json`. -- `checkpoints/`, `measurements/`, `spectral_snapshots/`, `spectra/`. +- `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. @@ -159,6 +198,15 @@ 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 diff --git a/baseline/experiments/distilbert_scratch_torchtpu/VALIDATION.json b/baseline/experiments/distilbert_scratch_torchtpu/VALIDATION.json index bf26b3e..ccfb269 100644 --- a/baseline/experiments/distilbert_scratch_torchtpu/VALIDATION.json +++ b/baseline/experiments/distilbert_scratch_torchtpu/VALIDATION.json @@ -2,12 +2,12 @@ "date": "2026-10-20", "status": "CPU engineering validated; real TPU qualification pending", "base_commit": "afd96a379f4b22c9199f58ca1442a89c5e087700", - "source_sha256": "b9f365baab28af883cbc0953eab41f697275aa4e3b928e1c5470f0391160f99a", + "source_sha256": "0bc2b349c13104f487d35ebce98d23affbcef961d19d39b838ed9d9186743d5b", "suite": { - "tests_run": 7, - "passed": 6, + "tests_run": 12, + "passed": 11, "skipped": 1, - "seconds": 53.117 + "seconds": 42.75 }, "passed": [ "full 66,985,530 parameter architecture and 36 encoder matrices", @@ -15,7 +15,14 @@ "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 training equals two updates plus restart exactly: parameters, optimizer state, evaluation metrics; checkpoint corruption refused" + "four-update tiny AdamW training equals 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" ], "spectral_smoke": { "weightwatcher": "0.7.7", @@ -64,7 +71,53 @@ "all-eight-device native TorchTPU collective probe", "real full-model compiled forward/backward and optimizer steps", "checkpoint/restart on actual TPU", - "WikiText-103 download/tokenization on target VM", - "throughput, convergence, and spectral results" - ] + "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 + } } diff --git a/baseline/experiments/distilbert_scratch_torchtpu/cloudshell.py b/baseline/experiments/distilbert_scratch_torchtpu/cloudshell.py index e3f61e7..bc6fd42 100644 --- a/baseline/experiments/distilbert_scratch_torchtpu/cloudshell.py +++ b/baseline/experiments/distilbert_scratch_torchtpu/cloudshell.py @@ -141,10 +141,15 @@ def start_remote(a): 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, @@ -163,6 +168,8 @@ def start_remote(a): 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)) @@ -185,7 +192,9 @@ def main(): 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) - p.add_argument('--resume-run') + 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) @@ -194,6 +203,12 @@ def main(): 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) command = ['sudo', 'python3', '-c', here.read_text(), a.action, '--remote', '--hours', str(a.hours)] @@ -208,6 +223,8 @@ def main(): 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)]) diff --git a/baseline/experiments/distilbert_scratch_torchtpu/common.py b/baseline/experiments/distilbert_scratch_torchtpu/common.py index 872ccca..6fc46d3 100644 --- a/baseline/experiments/distilbert_scratch_torchtpu/common.py +++ b/baseline/experiments/distilbert_scratch_torchtpu/common.py @@ -22,12 +22,13 @@ class Config: weight_decay: float = 0.01 beta1: float = 0.9 beta2: float = 0.999 - eps: float = 1e-8 + 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 @@ -51,7 +52,7 @@ def validate(self): 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 0 < self.mask_probability < 1 or self.precision not in ('bf16', 'fp32'): + 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 @@ -128,10 +129,11 @@ def batch_indices(order, batch_index, rank, world, local_batch): return order[positions % len(order)], valid -def masked_batch(tokens, indices, valid, meta, c, salt, step=0): +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 @@ -139,7 +141,7 @@ def masked_batch(tokens, indices, valid, meta, c, salt, step=0): eligible = np.flatnonzero(~special[row]) if not len(eligible): raise ValueError('Prepared sequence has no maskable tokens') - chosen = eligible[g.random(len(eligible)) < c.mask_probability] + 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] diff --git a/baseline/experiments/distilbert_scratch_torchtpu/configs/initial.json b/baseline/experiments/distilbert_scratch_torchtpu/configs/initial.json index 72d73a6..ae4227a 100644 --- a/baseline/experiments/distilbert_scratch_torchtpu/configs/initial.json +++ b/baseline/experiments/distilbert_scratch_torchtpu/configs/initial.json @@ -10,12 +10,13 @@ "weight_decay": 0.01, "beta1": 0.9, "beta2": 0.999, - "eps": 1e-08, + "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, 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/prepare.py b/baseline/experiments/distilbert_scratch_torchtpu/prepare.py index 5873a76..dd00994 100644 --- a/baseline/experiments/distilbert_scratch_torchtpu/prepare.py +++ b/baseline/experiments/distilbert_scratch_torchtpu/prepare.py @@ -8,6 +8,7 @@ TOKENIZER = 'distilbert/distilbert-base-uncased' TOKENIZER_REVISION = '12040accade4e8a0f71eabdb258fecc2e7e948be' DATASET = 'Salesforce/wikitext' +DATASET_REVISION = 'b08601e04326c79dfdd32d625aee71d232d685c3' def prepare(root, length=128, dataset_revision=None): @@ -18,7 +19,7 @@ def prepare(root, length=128, dataset_revision=None): 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).sha + 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) @@ -26,13 +27,15 @@ def prepare(root, length=128, dataset_revision=None): 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='[CLS] + 126 content tokens + [SEP]; separator between nonempty lines; no cross-split packing') + 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 diff --git a/baseline/experiments/distilbert_scratch_torchtpu/supervisor.py b/baseline/experiments/distilbert_scratch_torchtpu/supervisor.py index 8d405d7..f6d8cbc 100644 --- a/baseline/experiments/distilbert_scratch_torchtpu/supervisor.py +++ b/baseline/experiments/distilbert_scratch_torchtpu/supervisor.py @@ -59,12 +59,16 @@ def main(a): frozen = subprocess.check_output([python, '-m', 'pip', 'freeze'], text=True) (root / 'requirements.lock.txt').write_text(frozen) config = root / 'config.json' - config.write_text((here / 'configs/initial.json').read_text()) + 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: - run([python, here / 'prepare.py', '--output', data], training_deadline) + 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') @@ -111,7 +115,9 @@ def main(a): p = argparse.ArgumentParser(description=__doc__) p.add_argument('--root', required=True) p.add_argument('--deadline', type=float, required=True) - p.add_argument('--resume-run') + 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) diff --git a/baseline/experiments/distilbert_scratch_torchtpu/test_experiment.py b/baseline/experiments/distilbert_scratch_torchtpu/test_experiment.py index 78ac43a..b9137d7 100644 --- a/baseline/experiments/distilbert_scratch_torchtpu/test_experiment.py +++ b/baseline/experiments/distilbert_scratch_torchtpu/test_experiment.py @@ -11,10 +11,11 @@ import unittest import numpy as np import torch -from common import Config, atomic_json, batch_indices, masked_batch, build_model, encoder_matrices +from common import (Config, atomic_json, batch_indices, masked_batch, build_model, + encoder_matrices, file_hash, make_optimizer, read_config) from prepare import synthetic from spectra import summarize -from train import LossModel +from train import LossModel, preserve_best_validation HERE = Path(__file__).parent @@ -76,6 +77,91 @@ def test_missing_fit_does_not_bias_mean(self): self.assertIsNone(summarize(rows, 2)['mean_alpha']) self.assertFalse(summarize(rows, 2)['all_fits_valid']) + 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) diff --git a/baseline/experiments/distilbert_scratch_torchtpu/train.py b/baseline/experiments/distilbert_scratch_torchtpu/train.py index d6d5f2f..1491311 100644 --- a/baseline/experiments/distilbert_scratch_torchtpu/train.py +++ b/baseline/experiments/distilbert_scratch_torchtpu/train.py @@ -164,7 +164,7 @@ def evaluate(net, arrays, meta, c, rt): 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) + 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)) @@ -179,9 +179,38 @@ def evaluate(net, arrays, meta, c, rt): return results -def measure(root, step, net, model, arrays, meta, c, rt): +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' @@ -194,7 +223,8 @@ def measure(root, step, net, model, arrays, meta, c, rt): 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, metric='masked-token accuracy, not downstream classification')) + test_used_for_selection=False, evaluation_mask_probability=c.evaluation_mask_probability, + metric='masked-token accuracy, not downstream classification')) rt.barrier() @@ -246,10 +276,12 @@ def train(a, rt): 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: - measure(root, step, net, model, arrays, meta, c, rt) + measure(root, step, net, model, arrays, meta, c, rt, identity) checkpoint(root, step, model, opt, rt, identity) batches = math.ceil(len(arrays['train']) / (rt.world * c.per_device_batch)) epoch, order = None, None @@ -296,14 +328,14 @@ def train(a, rt): if step <= 3 or step % 10 == 0: print(json.dumps(status), flush=True) if step % c.measure_steps == 0: - measure(root, step, net, model, arrays, meta, c, rt) + measure(root, step, net, model, arrays, meta, c, rt, identity) last_measured = step if step % c.checkpoint_steps == 0: checkpoint(root, step, model, opt, rt, identity) # 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) + 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)) From a6ae88b5a91dc4a8fe69099ae40069878f8bdcb9 Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Tue, 20 Oct 2026 21:27:18 -0700 Subject: [PATCH 3/3] Fix TorchTPU RNG compatibility and interrupted-run recovery --- .../distilbert_scratch_torchtpu/AUDIT.md | 24 +++ .../distilbert_scratch_torchtpu/README.md | 18 ++- .../VALIDATION.json | 55 +++++-- .../distilbert_scratch_torchtpu/cloudshell.py | 9 +- .../distilbert_scratch_torchtpu/common.py | 11 +- .../distilbert_scratch_torchtpu/spectra.py | 17 ++- .../distilbert_scratch_torchtpu/supervisor.py | 34 +++-- .../test_experiment.py | 141 +++++++++++++++++- .../distilbert_scratch_torchtpu/train.py | 43 ++++-- 9 files changed, 302 insertions(+), 50 deletions(-) diff --git a/baseline/experiments/distilbert_scratch_torchtpu/AUDIT.md b/baseline/experiments/distilbert_scratch_torchtpu/AUDIT.md index 09e0e31..61744f3 100644 --- a/baseline/experiments/distilbert_scratch_torchtpu/AUDIT.md +++ b/baseline/experiments/distilbert_scratch_torchtpu/AUDIT.md @@ -23,6 +23,30 @@ validation-based calibration and optimizer comparisons. | 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 diff --git a/baseline/experiments/distilbert_scratch_torchtpu/README.md b/baseline/experiments/distilbert_scratch_torchtpu/README.md index 7ec8b41..48dfc47 100644 --- a/baseline/experiments/distilbert_scratch_torchtpu/README.md +++ b/baseline/experiments/distilbert_scratch_torchtpu/README.md @@ -32,7 +32,7 @@ run does not establish optimizer bias or overfitting. | 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 and cooperative exit; retain two committed generations | +| 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 | @@ -84,7 +84,7 @@ TorchTPU supplies the required Torch/libtpu stack; the old XLA venv is not inher The job performs: -1. All-eight-rank native TPU matmul and all-reduce probe. +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. @@ -132,6 +132,11 @@ 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 @@ -221,7 +226,10 @@ 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 states/metrics, corruption rejection, and in-process eight-shard gradients versus a serial masked-token mean. An +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. @@ -239,6 +247,10 @@ 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 index ccfb269..43a5070 100644 --- a/baseline/experiments/distilbert_scratch_torchtpu/VALIDATION.json +++ b/baseline/experiments/distilbert_scratch_torchtpu/VALIDATION.json @@ -1,13 +1,15 @@ { - "date": "2026-10-20", + "date": "2026-10-21", "status": "CPU engineering validated; real TPU qualification pending", "base_commit": "afd96a379f4b22c9199f58ca1442a89c5e087700", - "source_sha256": "0bc2b349c13104f487d35ebce98d23affbcef961d19d39b838ed9d9186743d5b", + "source_sha256": "8e8909b3fe98e521ec361aa73b546c4d2b3cd1de1bb129d098d52c17be592beb", "suite": { - "tests_run": 12, - "passed": 11, + "tests_run": 20, + "passed": 19, "skipped": 1, - "seconds": 42.75 + "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", @@ -15,14 +17,22 @@ "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 training equals two updates plus restart exactly: parameters, optimizer state, evaluation metrics; checkpoint corruption refused", + "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" + "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", @@ -64,11 +74,12 @@ "TORCH_TPU_SLICEBUILDER_ADDRESSES detection", "get_rng_state", "set_rng_state", - "tpu_dist" + "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 probe", + "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", @@ -118,6 +129,28 @@ }, "steps_per_corpus_pass": 3568, "pilot_full_corpus_passes": 2.802749465971636, - "pilot_real_sequence_positions": 327662848 - } + "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 index bc6fd42..19d9715 100644 --- a/baseline/experiments/distilbert_scratch_torchtpu/cloudshell.py +++ b/baseline/experiments/distilbert_scratch_torchtpu/cloudshell.py @@ -54,7 +54,7 @@ def bounded_deadline(checked, hours): return deadline -def discover_node(zone): +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 = [] @@ -64,8 +64,9 @@ def discover_node(zone): continue node = spec['name'].rsplit('/', 1)[-1] try: - checked = checked_lease(node, zone) - bounded_deadline(checked, 5.5) + 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)) @@ -210,7 +211,7 @@ def main(): 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) + 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(): diff --git a/baseline/experiments/distilbert_scratch_torchtpu/common.py b/baseline/experiments/distilbert_scratch_torchtpu/common.py index 6fc46d3..e28dc5b 100644 --- a/baseline/experiments/distilbert_scratch_torchtpu/common.py +++ b/baseline/experiments/distilbert_scratch_torchtpu/common.py @@ -38,6 +38,14 @@ class Config: 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: @@ -47,7 +55,8 @@ def validate(self): raise ValueError('Positive sizes required') if not 8 <= self.sequence_length <= 512: raise ValueError('Sequence length must be 8..512') - if not (self.lr > 0 and self.weight_decay >= 0 and self.grad_clip > 0 and self.eps > 0): + 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)): diff --git a/baseline/experiments/distilbert_scratch_torchtpu/spectra.py b/baseline/experiments/distilbert_scratch_torchtpu/spectra.py index 1f14abe..9b06162 100644 --- a/baseline/experiments/distilbert_scratch_torchtpu/spectra.py +++ b/baseline/experiments/distilbert_scratch_torchtpu/spectra.py @@ -92,14 +92,21 @@ def export(root): 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() - while time.time() < args.deadline: - analyze(args.root) - if not args.watch or (Path(args.root) / 'TRAINING_DONE.json').exists(): - break - time.sleep(15) + 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 index f6d8cbc..e8c6364 100644 --- a/baseline/experiments/distilbert_scratch_torchtpu/supervisor.py +++ b/baseline/experiments/distilbert_scratch_torchtpu/supervisor.py @@ -16,6 +16,22 @@ def write(path, value): 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') @@ -23,14 +39,10 @@ def run(command, deadline, env=None): try: code = proc.wait(timeout=max(1, deadline - time.time())) except subprocess.TimeoutExpired: - os.killpg(proc.pid, signal.SIGTERM) - try: - proc.wait(timeout=20) - except subprocess.TimeoutExpired: - os.killpg(proc.pid, signal.SIGKILL) - proc.wait() + terminate(proc) raise RuntimeError('Phase exceeded its bounded deadline') if code: + terminate(proc) raise RuntimeError(f'Phase failed ({code}): {command}') @@ -96,19 +108,17 @@ def main(a): try: analyst.wait(timeout=max(1, a.deadline - time.time() - 60)) except subprocess.TimeoutExpired: - os.killpg(analyst.pid, signal.SIGTERM) - analyst.wait(timeout=20) + terminate(analyst) result = json.loads((output / 'RESULT.json').read_text()) result['spectra_exit_code'] = analyst.returncode result['spectra_complete'] = (analyst.returncode == 0 and - len(list((output / 'measurements').glob('step_*.json'))) == - len(list((output / 'spectra').glob('step_*.json')))) + {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: - os.killpg(analyst.pid, signal.SIGTERM) - analyst.wait(timeout=20) + terminate(analyst) if __name__ == '__main__': diff --git a/baseline/experiments/distilbert_scratch_torchtpu/test_experiment.py b/baseline/experiments/distilbert_scratch_torchtpu/test_experiment.py index b9137d7..434c9c8 100644 --- a/baseline/experiments/distilbert_scratch_torchtpu/test_experiment.py +++ b/baseline/experiments/distilbert_scratch_torchtpu/test_experiment.py @@ -8,14 +8,18 @@ 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 -from train import LossModel, preserve_best_validation +from spectra import summarize, watch +from supervisor import terminate +from train import LossModel, Runtime, checkpoint, preserve_best_validation, train HERE = Path(__file__).parent @@ -77,6 +81,128 @@ def test_missing_fit_does_not_bias_mean(self): 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) @@ -170,11 +296,10 @@ def test_full_architecture_and_random_init(self): 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 test_resume_matches_uninterrupted_adamw(self): + def assert_resume_matches(self, c): with tempfile.TemporaryDirectory() as d: root = Path(d) synthetic(root / 'data') - c = self.config() full = run_job(root, 'full', c) run_job(root, 'resume', c, stop=2) resumed = run_job(root, 'resume', c, resume=True) @@ -196,6 +321,14 @@ def test_resume_matches_uninterrupted_adamw(self): 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: diff --git a/baseline/experiments/distilbert_scratch_torchtpu/train.py b/baseline/experiments/distilbert_scratch_torchtpu/train.py index 1491311..da48d9a 100644 --- a/baseline/experiments/distilbert_scratch_torchtpu/train.py +++ b/baseline/experiments/distilbert_scratch_torchtpu/train.py @@ -66,7 +66,10 @@ def barrier(self): def rng(self): state = dict(torch=torch.get_rng_state(), python=random.getstate(), numpy=np.random.get_state()) if self.backend == 'tpu': - state['tpu'] = torch.tpu.get_rng_state().cpu() + # 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): @@ -89,10 +92,17 @@ def probe(rt, output): 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, + 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: @@ -117,16 +127,25 @@ def checkpoint(root, step, model, opt, rt, identity): rt.barrier() if rt.rank == 0: atomic_json(partial / 'manifest.json', {p.name: file_hash(p) for p in partial.glob('*.pt')}) - # A generation is immutable. Re-saving the same cursor leaves it intact. - if target.exists(): + 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)) - for old in sorted(checkpoints.glob('step_*'))[:-2]: - if old.is_dir() and not old.name.endswith('.partial'): - shutil.rmtree(old) + # 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() @@ -281,8 +300,12 @@ def train(a, rt): 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: - measure(root, step, net, model, arrays, meta, c, rt, identity) 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 @@ -327,11 +350,11 @@ def train(a, rt): 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 - if step % c.checkpoint_steps == 0: - checkpoint(root, step, model, opt, rt, identity) # 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):