From 760ed839ea033f9cb1f8b37c80871fadac561bb2 Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Wed, 7 Oct 2026 10:02:31 -0700 Subject: [PATCH] Add isolated full-dense v5e TPU hardware probe --- .../README.md | 16 ++ .../experiment.py | 71 ++++++--- .../tests/test_probe.py | 139 ++++++++++++++++++ .../tpu_port/probe.py | 136 +++++++++++++++++ .../tpu_port/probe_report.py | 90 ++++++++++++ .../tpu_port/runtime.py | 26 ++++ 6 files changed, 461 insertions(+), 17 deletions(-) create mode 100644 baseline/experiments/nanogpt_leaderboard_2026_10_06_tpu/tests/test_probe.py create mode 100644 baseline/experiments/nanogpt_leaderboard_2026_10_06_tpu/tpu_port/probe.py create mode 100644 baseline/experiments/nanogpt_leaderboard_2026_10_06_tpu/tpu_port/probe_report.py diff --git a/baseline/experiments/nanogpt_leaderboard_2026_10_06_tpu/README.md b/baseline/experiments/nanogpt_leaderboard_2026_10_06_tpu/README.md index 6726357f..40c2a90a 100644 --- a/baseline/experiments/nanogpt_leaderboard_2026_10_06_tpu/README.md +++ b/baseline/experiments/nanogpt_leaderboard_2026_10_06_tpu/README.md @@ -18,6 +18,21 @@ The sibling [CUDA reference](../nanogpt_leaderboard_2026_10_06) remains byte-ver The table is **not shrunk**. Eight CPU row owners hold it in host RAM, exchange requested rows through Gloo, and upload bounded BF16 caches to their TPU chips. Dense gradients use XLA collectives. One PJRT process runs per chip (MPMD); SPMD is rejected. Dense optimizer state is replicated. The launcher requires **192 GiB available host RAM** and **160 GiB free export disk**; these guards do not prove HBM fit. `capacity.py` separately explains why the all-HBM approach fails on this machine. +## Probe before preflight + +After installing the pinned environment with `setup_tpu.sh` and activating `.venv`, run: + +```bash +export PJRT_DEVICE=TPU +python experiment.py probe +``` + +The probe needs the real v5litepod-8 but no data shards. It does not allocate the 121 GiB n-gram table, build the canonical mask, or start WeightWatcher. It retains the real model's 262144-token buffers and replicated optimizer/tail state, trains twice at 128 and 49152 tokens per chip, releases optimizer/tail state, and evaluates once at 262144 tokens per chip. It uses synthetic full-softmax batches and a zero canonical mask; the loss is not benchmark NLL. + +Each rank writes `results//probe-rank-XX.json`; rank 0 writes `PROBE_SUMMARY.json`, which the launcher finalizes after all workers exit. Reports contain hardware, package versions, host memory, HBM by stage, finite-loss checks, and step wall times plus XLA compile counts/seconds. The second step is labeled **steady**, but its first Adam update can compile an additional graph; it is not guaranteed compilation-free. Failures record the stage and exception, exit nonzero, and stop the ladder. The launcher stops the other ranks after a reported failure. A hard process kill cannot execute Python's `finally`; the summary records missing/incomplete ranks as failure. + +**The probe must be green on all eight ranks, including the 262144-token evaluation forward, before attempting preflight.** A successful probe is not qualification, not a 3.28 run, and not a reason to set `ready_to_train`. Preflight remains mandatory before `run`; a probe report never substitutes for its receipt. Qualification flags remain false. The probe has no 192 GiB host-RAM or 160 GiB export-disk gate, but still requires enough memory for the dense ladder and writable space for its reports. `--results-root` can select another directory. + ## Run on the TPU VM Use Python 3.11 or 3.12 and `python3-venv` on the existing v5e-8 VM. No H100 Docker image, CUDA wheel or FlashAttention install is needed. @@ -29,6 +44,7 @@ source .venv/bin/activate export PJRT_DEVICE=TPU python experiment.py verify python experiment.py check +python experiment.py probe # require a green PROBE_SUMMARY.json before proceeding python experiment.py prepare --data-root /data/nanogpt-leaderboard python experiment.py preflight --data-root /data/nanogpt-leaderboard ``` diff --git a/baseline/experiments/nanogpt_leaderboard_2026_10_06_tpu/experiment.py b/baseline/experiments/nanogpt_leaderboard_2026_10_06_tpu/experiment.py index dedb40eb..9b2056f2 100644 --- a/baseline/experiments/nanogpt_leaderboard_2026_10_06_tpu/experiment.py +++ b/baseline/experiments/nanogpt_leaderboard_2026_10_06_tpu/experiment.py @@ -3,11 +3,13 @@ import datetime import hashlib import importlib.util +import importlib.metadata import json import os from pathlib import Path import secrets import shutil +import signal import subprocess import sys import time @@ -70,11 +72,12 @@ def runtime_environment(): raise RuntimeError('Production/preflight requires PJRT_DEVICE=TPU') os.environ['PJRT_DEVICE']='TPU' os.environ.setdefault('OMP_NUM_THREADS','4') - import torch, torch_xla - versions=[torch.__version__.split('+')[0],torch_xla.__version__.split('+')[0]] + torch_version=importlib.metadata.version('torch') + xla_version=importlib.metadata.version('torch_xla') + versions=[torch_version.split('+')[0],xla_version.split('+')[0]] if versions!=['2.9.0','2.9.0']: raise RuntimeError('Use matching torch==2.9.0 and torch_xla==2.9.0: '+repr(versions)) - return {'torch':torch.__version__,'torch_xla':torch_xla.__version__, + return {'torch':torch_version,'torch_xla':xla_version, 'python':sys.version,'host_available_bytes':available_host_memory()} @@ -89,6 +92,30 @@ def tracking_environment(): return env +def probe_subprocess(cmd,root): + # Poll failure reports even when compilation emits no stdout. On an OOM, + # stop the entire launch process group, including ranks blocked in a collective. + with (root/'console.log').open('w') as log: + process=subprocess.Popen(cmd,stdout=log,stderr=subprocess.STDOUT,start_new_session=True) + try: + while process.poll() is None: + for path in root.glob('probe-rank-*.json'): + report=json.loads(path.read_text()) + if report.get('status')=='failed': + raise RuntimeError(f"Probe rank {report['rank']} failed at {report['failed_stage']}: {report['exception']}") + try: process.wait(timeout=1) + except subprocess.TimeoutExpired: pass + return process.returncode + finally: + # Kill remaining descendants too if a rank or the launch parent failed. + if process.poll() is None or process.returncode: + try: os.killpg(process.pid,signal.SIGTERM) + except ProcessLookupError: pass + try: process.wait(timeout=10) + except subprocess.TimeoutExpired: + os.killpg(process.pid,signal.SIGKILL); process.wait() + + def launch(action,data_root,results_root,receipt=None): verify_reference() results_root=Path(results_root).resolve(); results_root.mkdir(parents=True,exist_ok=True) @@ -96,9 +123,9 @@ def launch(action,data_root,results_root,receipt=None): if shutil.disk_usage(results_root).free=0 and int(batch.slots.max())=0)&(batch.seqlens<=tokens)).all() + assert ((batch.seqlens.diff()>0)&(batch.seqlens.diff()<=2560)).all() + assert batch.inputs.min()>=0 and batch.inputs.max()<50257 + assert batch.targets.min()>=0 and batch.targets.max()<50257 + expected=batch.inputs.clone(); del batch + # Different process/model RNG does not change the synthetic batch. + torch.manual_seed(42) + evaluation=synthetic_batch(tokens,training=False) + assert evaluation.sink is None + assert torch.equal(evaluation.inputs,expected) + + +def test_environment_versions_without_xla_import(monkeypatch): + for key in ('XLA_USE_SPMD','XLA_AUTO_SPMD','XLA_USE_BF16','XLA_DOWNCAST_BF16', + 'NUM_SCHEDULED_ITERATIONS','TRAIN_SEED'): + monkeypatch.delenv(key,raising=False) + monkeypatch.setenv('PJRT_DEVICE','TPU') + monkeypatch.setattr(experiment.importlib.metadata,'version',lambda name:'2.9.0') + original=builtins.__import__ + def guarded(name,*args,**kwargs): + if name=='torch_xla' or name.startswith('torch_xla.'): + pytest.fail('Launch grandparent imported torch_xla') + return original(name,*args,**kwargs) + monkeypatch.setattr(builtins,'__import__',guarded) + assert experiment.runtime_environment()['torch_xla']=='2.9.0' + + +def complete_reports(root): + for rank in range(8): + value={'rank':rank,'status':'complete','finite_loss':True, + 'stages':[{'stage':stage,'status':'complete','finite_loss':True} for stage in STEP_STAGES],**FLAGS} + (root/f'probe-rank-{rank:02d}.json').write_text(json.dumps(value)) + + +@pytest.mark.parametrize('action',['check','probe']) +def test_probe_and_check_skip_data_tracking_and_capacity_gates(monkeypatch,tmp_path,action): + def forbidden(*args,**kwargs): pytest.fail('Data, tracking or extra subprocess requested') + monkeypatch.setattr(experiment,'verify_reference',lambda:None) + monkeypatch.setattr(experiment,'runtime_environment',lambda:{'torch':'2.9.0','torch_xla':'2.9.0','host_available_bytes':0}) + monkeypatch.setattr(experiment,'source_fingerprint',lambda:{}) + monkeypatch.setattr(experiment.shutil,'disk_usage',lambda path:SimpleNamespace(free=0)) + monkeypatch.setattr(experiment,'data_manifest',forbidden) + monkeypatch.setattr(experiment,'tracking_command',forbidden) + monkeypatch.setattr(experiment.subprocess,'run',forbidden) + seen=[] + def fake_probe(cmd,root): + seen.append(json.loads(cmd[-1])['action']); complete_reports(root); return 0 + monkeypatch.setattr(experiment,'probe_subprocess',fake_probe) + def fake_popen(cmd,**kwargs): + assert action=='check' + seen.append(json.loads(cmd[-1])['action']) + return SimpleNamespace(stdout=[],wait=lambda:0) + monkeypatch.setattr(experiment.subprocess,'Popen',fake_popen) + root=experiment.launch(action,tmp_path/'nonexistent-data',tmp_path/'results') + manifest=json.loads((root/'run_manifest.json').read_text()) + assert manifest['data'] is None and manifest['status']=='complete' + assert seen==[action] + assert not (root/'PREFLIGHT_COMPLETE.json').exists() + assert 'weightwatcher_complete' not in manifest + + +def test_probe_cli_routes_action(monkeypatch): + seen=[] + monkeypatch.setattr(sys,'argv',['experiment.py','probe']) + monkeypatch.setattr(experiment,'launch',lambda *args:seen.append(args[0])) + experiment.main() + assert seen==['probe'] + + +def test_worker_oom_records_failure_and_does_not_continue(monkeypatch,tmp_path): + from tpu_port import runtime + def fail(index,options,report): + report.begin('train_49152_compile') + raise torch.OutOfMemoryError('simulated CPU-only OOM') + monkeypatch.setattr(runtime,'_worker',fail) + with pytest.raises(torch.OutOfMemoryError): + runtime.worker(0,{'action':'probe','output':str(tmp_path)}) + value=json.loads((tmp_path/'probe-rank-00.json').read_text()) + assert value['status']=='failed' and value['failed_stage']=='train_49152_compile' + assert 'simulated CPU-only OOM' in value['exception'] + assert all(value[key] is expected for key,expected in FLAGS.items()) + assert summarize(tmp_path,final=True)['status']=='failed' + + +def test_incomplete_workers_cannot_be_green(tmp_path): + report=ProbeReport(tmp_path,0) + report.begin('eval_262144_compile') + assert summarize(tmp_path,final=True)['status']=='failed' + + +def test_full_eval_is_required_for_green_summary(tmp_path): + complete_reports(tmp_path) + assert summarize(tmp_path,final=True)['status']=='complete' + path=tmp_path/'probe-rank-07.json' + report=json.loads(path.read_text()) + report['stages']=[s for s in report['stages'] if s['stage']!='eval_262144_compile'] + path.write_text(json.dumps(report)) + assert summarize(tmp_path,final=True)['status']=='failed' + + +def test_failed_launch_finalizes_json_without_success(monkeypatch,tmp_path,capsys): + monkeypatch.setattr(experiment,'verify_reference',lambda:None) + monkeypatch.setattr(experiment,'runtime_environment',lambda:{'host_available_bytes':0}) + monkeypatch.setattr(experiment,'source_fingerprint',lambda:{}) + monkeypatch.setattr(experiment,'probe_subprocess',lambda *args:1) + with pytest.raises(RuntimeError,match='TPU subprocess failed'): + experiment.launch('probe',tmp_path/'no-data',tmp_path/'results') + root=next((tmp_path/'results').iterdir()) + assert json.loads((root/'PROBE_SUMMARY.json').read_text())['status']=='failed' + assert json.loads((root/'run_manifest.json').read_text())['status']=='failed' + assert 'Results:' not in capsys.readouterr().out diff --git a/baseline/experiments/nanogpt_leaderboard_2026_10_06_tpu/tpu_port/probe.py b/baseline/experiments/nanogpt_leaderboard_2026_10_06_tpu/tpu_port/probe.py new file mode 100644 index 00000000..50de0bd0 --- /dev/null +++ b/baseline/experiments/nanogpt_leaderboard_2026_10_06_tpu/tpu_port/probe.py @@ -0,0 +1,136 @@ +"""Full dense TPU hardware ladder. Synthetic data; never a benchmark or preflight.""" +import gc +import math +import time +from dataclasses import dataclass +import torch + +from .gpt import GPT, ForwardScheduleConfig +from .optimizer import Optimizer +from .tail import TailAverages + + +@dataclass +class SyntheticBatch: + inputs: torch.Tensor + targets: torch.Tensor + seqlens: torch.Tensor + slots: torch.Tensor + cache: torch.Tensor + sink: torch.Tensor | None + + +def synthetic_batch(tokens,device='cpu',training=True): + if tokens<=0 or tokens%16: raise ValueError('Probe tokens must be a positive multiple of 16') + # Local CPU RNG makes batches identical across ranks, independent of model RNG. + generator=torch.Generator(device='cpu').manual_seed(1337) + sequence=torch.randint(0,50257,(tokens+1,),generator=generator,dtype=torch.int32) + ends=[*range(0,tokens,2560),tokens] + cache=torch.zeros((2*tokens,768),dtype=torch.bfloat16,device=device) + sink=torch.zeros_like(cache,requires_grad=True) if training else None + return SyntheticBatch(sequence[:-1].to(device),sequence[1:].long().to(device), + torch.tensor(ends,dtype=torch.int32,device=device), + torch.arange(2*tokens,dtype=torch.int32,device=device),cache,sink) + + +def compile_sample(): + from torch_xla.debug import metrics + value=metrics.metric_data('CompileTime') + # PyTorch/XLA time-metric accumulators are nanoseconds, not seconds. + return (0,0.) if value is None else (int(value[0]),float(value[1])/1e9) + + +def record_compile(event,before): + after=compile_sample() + event.update(compile_count=after[0]-before[0],compile_seconds=after[1]-before[1]) + + +def run(report,device,xm): + from .runtime import schedule + # Retain the real maximum-length buffers, including the zero canonical mask. + # No HostTable, tokenizer, canonical-mask builder or WeightWatcher is called. + event=report.begin('model_move') + torch.manual_seed(1337) + model=GPT(50257,11,6,128,768,262144,ngram_dim=768,world_size=8,device=torch.device('cpu')) + model.cast_matrix_weights_bf16() + model=model.to(device) + for yarn in model.yarns: + yarn.device=device; yarn.angular_freq=yarn.angular_freq.to(device) + xm.mark_step(wait=True) + event['status']='complete'; report.memory(xm,device,event) + + event=report.begin('optimizer_and_tail') + sched=schedule() + def reduce_mean(gradient): + return xm.all_reduce(xm.REDUCE_SUM,gradient,scale=1/8) + optimizer=Optimizer(model,sched,reduce_mean) + tail=TailAverages(optimizer.params,sched.total_steps) + xm.mark_step(wait=True) + event['status']='complete'; report.memory(xm,device,event) + cfg=ForwardScheduleConfig(torch.tensor([1.],device=device),torch.tensor([0.],device=device), + 640,1408,2560,None) + report.data['timing_note']=('First step includes compilation; second is labeled steady but Adam on ' + 'optimizer.step(1) adds a distinct graph and may compile. Wall time includes host synchronization.') + + for tokens in (128,49152): + event=report.begin(f'train_{tokens}_batch',tokens_per_rank=tokens) + batch=synthetic_batch(tokens,device) + model.ngram_cache=batch.cache + model.train(); xm.mark_step(wait=True) + event['status']='complete'; report.memory(xm,device,event) + for step,label in ((0,'compile'),(1,'steady')): + event=report.begin(f'train_{tokens}_{label}',tokens_per_rank=tokens, + timing_label=label,optimizer_step=step) + compiled_before=compile_sample() + began=time.perf_counter() + try: + loss=model(batch.inputs,batch.targets,batch.seqlens,batch.slots,cfg,batch.sink) + loss.sum().backward() + xm.mark_step(wait=True) + value=float(loss.detach().float().mean().cpu()) + event.update(loss=value if math.isfinite(value) else None,finite_loss=math.isfinite(value)) + if not math.isfinite(value): + report.data['finite_loss']=False + raise FloatingPointError('Nonfinite probe training loss') + optimizer.step(step) + xm.mark_step(wait=True) + record_compile(event,compiled_before) + event['status']='complete'; report.data['finite_loss']=True + finally: + event['seconds']=time.perf_counter()-began + report.flush() + report.memory(xm,device,event) + batch.sink.grad=None + del loss + # Accumulate Adam gradients within each pair, but not across batch sizes. + for parameter in model.parameters(): parameter.grad=None + del batch + + event=report.begin('release_optimizer_and_tail') + del optimizer,tail + gc.collect(); xm.mark_step(wait=True) + event['status']='complete'; report.memory(xm,device,event) + + event=report.begin('eval_262144_batch',tokens_per_rank=262144) + batch=synthetic_batch(262144,device,training=False) + model.ngram_cache=batch.cache + model.eval(); xm.mark_step(wait=True) + event['status']='complete'; report.memory(xm,device,event) + event=report.begin('eval_262144_compile',tokens_per_rank=262144,timing_label='compile') + compiled_before=compile_sample() + began=time.perf_counter() + try: + with torch.no_grad(): + loss=model(batch.inputs,batch.targets,batch.seqlens,batch.slots,cfg) + xm.mark_step(wait=True) + value=float(loss.float().mean().cpu()) + event.update(loss=value if math.isfinite(value) else None,finite_loss=math.isfinite(value)) + if not math.isfinite(value): + report.data['finite_loss']=False + raise FloatingPointError('Nonfinite probe evaluation loss') + record_compile(event,compiled_before) + event['status']='complete'; report.data['finite_loss']=True + finally: + event['seconds']=time.perf_counter()-began + report.flush() + report.memory(xm,device,event) diff --git a/baseline/experiments/nanogpt_leaderboard_2026_10_06_tpu/tpu_port/probe_report.py b/baseline/experiments/nanogpt_leaderboard_2026_10_06_tpu/tpu_port/probe_report.py new file mode 100644 index 00000000..53799590 --- /dev/null +++ b/baseline/experiments/nanogpt_leaderboard_2026_10_06_tpu/tpu_port/probe_report.py @@ -0,0 +1,90 @@ +"""Standard-library probe reporting; safe to import in the launch grandparent.""" +import json +from pathlib import Path +import resource +import sys +import time + +FLAGS={'full_table_allocated':False,'tpu_execution_verified':False,'this_is_not_preflight':True} +STEP_STAGES={'train_128_compile','train_128_steady','train_49152_compile', + 'train_49152_steady','eval_262144_compile'} + + +def host_memory(): + available=None + for line in Path('/proc/meminfo').read_text().splitlines(): + if line.startswith('MemAvailable:'): available=int(line.split()[1])*1024 + return {'mem_available_bytes':available, + 'peak_host_rss_bytes':int(resource.getrusage(resource.RUSAGE_SELF).ru_maxrss)*1024} + + +def write_json(path,value): + path=Path(path); temporary=path.with_suffix('.tmp') + temporary.write_text(json.dumps(value,indent=2,allow_nan=False)+'\n'); temporary.replace(path) + + +def summarize(root,*,final=False,exception=None): + root=Path(root); ranks=[] + for rank in range(8): + path=root/f'probe-rank-{rank:02d}.json' + value=json.loads(path.read_text()) if path.exists() else {'rank':rank,'status':'missing',**FLAGS} + if final and value.get('status') not in ('complete','failed'): + value.update(status='failed',failed_stage=value.get('current_stage','worker_initialization'), + exception=exception or 'Worker exited or was stopped before completing the probe') + if value.get('stages') and value['stages'][-1]['status']=='running': + value['stages'][-1]['status']='failed' + # The launcher only finalizes after stopping the process group. + write_json(path,value) + ranks.append(value) + complete=all(r.get('status')=='complete' and r.get('finite_loss') is True and + STEP_STAGES<={s['stage'] for s in r.get('stages',[]) + if s.get('status')=='complete' and s.get('finite_loss') is True} + for r in ranks) + failed=exception is not None or any(r.get('status')=='failed' for r in ranks) or (final and not complete) + failures=[r for r in ranks if r.get('status')=='failed'] + summary={'status':'failed' if failed else 'complete' if complete else 'pending', + 'expected_world_size':8,'ranks':ranks, + 'exception':exception or (failures[0].get('exception') if failures else None), + 'failed_stage':failures[0].get('failed_stage') if failures else ('incomplete_workers' if failed else None), + 'finite_loss':all(r.get('finite_loss') is True for r in ranks) if complete else None,**FLAGS} + for key in ('python','torch','torch_xla','world_size','device','runtime_device_attributes'): + summary[key]=ranks[0].get(key) + write_json(root/'PROBE_SUMMARY.json',summary) + return summary + + +class ProbeReport: + def __init__(self,root,rank): + self.root,self.rank=Path(root),rank + self.started=time.perf_counter(); self.stage='worker_initialization' + self.data={'rank':rank,'status':'running','python':sys.version,'torch':None,'torch_xla':None, + 'world_size':None,'device':None,'runtime_device_attributes':None, + 'stages':[],'finite_loss':None,'failed_stage':None,'exception':None,**FLAGS} + self.flush() + + def flush(self): + self.data.update(host_memory(),current_stage=self.stage, + elapsed_seconds=time.perf_counter()-self.started) + write_json(self.root/f'probe-rank-{self.rank:02d}.json',self.data) + + def begin(self,stage,**metadata): + self.stage=stage + event={'stage':stage,'status':'running','hbm':None,**metadata} + self.data['stages'].append(event); self.flush() + return event + + def memory(self,xm,device,event): + event['hbm']={key:int(value) for key,value in xm.get_memory_info(device).items()} + event.update(host_memory()); self.flush() + + def finish(self,error=None): + if error is not None: + self.data.update(status='failed',failed_stage=self.stage, + exception=f'{type(error).__name__}: {error}') + if self.data['stages']: + self.data['stages'][-1]['status']='failed' + else: + self.data['status']='complete' + # Never query XLA here: an OOM can leave the runtime unusable. + self.flush() + if self.rank==0: summarize(self.root) diff --git a/baseline/experiments/nanogpt_leaderboard_2026_10_06_tpu/tpu_port/runtime.py b/baseline/experiments/nanogpt_leaderboard_2026_10_06_tpu/tpu_port/runtime.py index 084f9a52..97a36e3d 100644 --- a/baseline/experiments/nanogpt_leaderboard_2026_10_06_tpu/tpu_port/runtime.py +++ b/baseline/experiments/nanogpt_leaderboard_2026_10_06_tpu/tpu_port/runtime.py @@ -113,11 +113,28 @@ def evaluate(model,table,data,device,cfg,*,step,tokens_seen,batch_tokens,batches def worker(index,options): + if options['action']!='probe': return _worker(index,options) + from .probe_report import ProbeReport + report=ProbeReport(options['output'],index) + error=None + try: + return _worker(index,options,report) + except BaseException as caught: + error=caught + raise + finally: + report.finish(error) + + +def _worker(index,options,probe_report=None): import torch_xla import torch_xla.core.xla_model as xm import torch_xla.runtime as xr import torch_xla.debug.metrics as metrics rank,world=xr.global_ordinal(),xr.world_size() + if probe_report is not None: + probe_report.data.update(torch=torch.__version__,torch_xla=torch_xla.__version__,world_size=world) + probe_report.flush() if xr.device_type()!='TPU' or world!=8 or xr.is_spmd(): raise RuntimeError('Requires eight TPU processes in PJRT MPMD mode; SPMD/CPU unsupported') device=xm.xla_device() @@ -129,6 +146,9 @@ def worker(index,options): rank=rank,world_size=world,timeout=datetime.timedelta(hours=2)) try: attributes=xr.global_runtime_device_attributes() + if probe_report is not None: + probe_report.data.update(device=str(device),runtime_device_attributes=attributes) + probe_report.begin('hardware_and_collective') # Runtime reports the actual kind; never infer eight chips from a v4/v5p core suffix. kinds={str(a.get('device_kind','')).lower() for a in attributes} if not kinds or any(not ('v5' in k and ('lite' in k or 'v5e' in k)) for k in kinds): @@ -143,6 +163,12 @@ def worker(index,options): raise RuntimeError('TPU collective preflight failed') write_json(root/f'hardware-rank-{rank:02d}.json',report) if options['action']=='check': return + if probe_report is not None: + event=probe_report.data['stages'][-1]; event['status']='complete' + probe_report.memory(xm,device,event) + from .probe import run + run(probe_report,device,xm) + return # CPU init is explicitly broadcast, including the signed n-gram pool. torch.manual_seed(options['seed'])