| """Fixed training-CPU budgets for the public dependency-tracking experiment. |
| |
| The sampler, objective and checkpoint boundaries are fixed before training. |
| CPU budgets match compute time, not update counts. Evaluation, serialization, |
| and logging are separately timed and never charged as optimizer training. |
| No sealed dataset is accepted by this driver. |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| from collections import defaultdict |
| from copy import deepcopy |
| from dataclasses import asdict, dataclass, replace |
| from datetime import datetime, timezone |
| import gzip |
| import hashlib |
| import io |
| import json |
| import math |
| from pathlib import Path |
| import platform |
| import random |
| import resource |
| import subprocess |
| import sys |
| import time |
| from typing import Callable |
|
|
| ROOT = Path(__file__).resolve().parents[1] |
| if str(ROOT) not in sys.path: |
| sys.path.insert(0, str(ROOT)) |
|
|
| import torch |
| from torch import nn |
|
|
| from lightnest.relation_patch import ( |
| RelationConfig, RelationPatchNetwork, RelationPredictor, |
| encode_batch as relation_encode, load_records, load_relation, supervised_loss, |
| ) |
|
|
| FORMAT = 'lightnest-dependency-training-v1' |
| CHECKPOINT_FORMAT = 'lightnest-dependency-checkpoint-v1' |
| SAMPLER_FORMAT = 'alternating-return-invariance-epochs-plus-replay-v2' |
| BASELINE_COMMIT = 'fcbf22b77aa308cc23f4f0d4cee84071ec486984' |
| FRACTIONS = (.25, .50, .75, 1.) |
| SEEDS = (104729, 130363) |
|
|
|
|
| @dataclass(frozen=True) |
| class Arm: |
| name: str |
| kind: str |
| message_steps: int = 4 |
| slots: int | None = None |
| cpu_multiplier: float = 1. |
| slot_mode: str = 'memory' |
|
|
|
|
| ARMS = { |
| item.name: item for item in ( |
| Arm('relation-d4', 'relation'), |
| Arm('relation-d8', 'relation', message_steps=8), |
| Arm('slot8', 'slot', slots=8), |
| Arm('slot16', 'slot', slots=16), |
| Arm('relation-d4-long', 'relation', cpu_multiplier=3.), |
| Arm('slot8_no_memory', 'slot', slots=8, slot_mode='no_memory'), |
| ) |
| } |
|
|
|
|
| def sha256(path): |
| digest = hashlib.sha256() |
| with Path(path).open('rb') as source: |
| for block in iter(lambda: source.read(1024 * 1024), b''): |
| digest.update(block) |
| return digest.hexdigest() |
|
|
|
|
| def json_bytes(value): |
| return (json.dumps(value, ensure_ascii=False, sort_keys=True, |
| separators=(',', ':'), allow_nan=False) + '\n').encode() |
|
|
|
|
| def write_json(path, value): |
| path = Path(path) |
| path.parent.mkdir(parents=True, exist_ok=True) |
| with path.open('x', encoding='utf-8') as output: |
| json.dump(value, output, ensure_ascii=False, sort_keys=True, indent=2, allow_nan=False) |
| output.write('\n') |
|
|
|
|
| def reference(path): |
| path = Path(path).resolve() |
| return str(path.relative_to(ROOT)) if path.is_relative_to(ROOT) else str(path) |
|
|
|
|
| def resolve(reference_path): |
| path = Path(reference_path) |
| return path if path.is_absolute() else ROOT / path |
|
|
|
|
| def public_path(path): |
| path = resolve(path).resolve() |
| if any(part.lower() in {'sealed', 'evaluation-final', 'final-test'} for part in path.parts): |
| raise ValueError('sealed/final-test input is prohibited in this training driver') |
| return path |
|
|
|
|
| def file_record(path): |
| path = Path(path) |
| return dict(path=reference(path), sha256=sha256(path), bytes=path.stat().st_size) |
|
|
|
|
| def verify_record(record, *, public=False): |
| path = public_path(record['path']) if public else resolve(record['path']) |
| if not path.is_file() or path.stat().st_size != record['bytes'] or sha256(path) != record['sha256']: |
| raise ValueError('frozen file changed or missing: ' + record['path']) |
| return path |
|
|
|
|
| def fresh_directory(path): |
| path = Path(path) |
| if path.exists() and any(path.iterdir()): |
| raise ValueError('refuse to overwrite an existing experiment: ' + str(path)) |
| path.mkdir(parents=True, exist_ok=True) |
| return path |
|
|
|
|
| def configure_threads(threads): |
| if threads not in (1, 2): |
| raise ValueError('only one or two CPU threads are allowed') |
| torch.set_num_threads(threads) |
| try: |
| torch.set_num_interop_threads(1) |
| except RuntimeError: |
| if torch.get_num_interop_threads() != 1: |
| raise |
|
|
|
|
| def tensor_inventory(state): |
| """Canonical hashes bind names, shapes, dtypes, and exact tensor bytes.""" |
| result = {} |
| combined = hashlib.sha256() |
| for name, tensor in sorted(state.items()): |
| tensor = tensor.detach().cpu().contiguous() |
| raw = tensor.reshape(-1).view(torch.uint8).numpy().tobytes() |
| entry = dict(shape=list(tensor.shape), dtype=str(tensor.dtype), |
| numel=tensor.numel(), sha256=hashlib.sha256(raw).hexdigest()) |
| result[name] = entry |
| combined.update(json_bytes(dict(name=name, **entry))) |
| return dict(sha256=combined.hexdigest(), tensors=result, |
| numel=sum(item['numel'] for item in result.values())) |
|
|
|
|
| def changed_inventory(initial, current): |
| rows = {} |
| for name, value in current.items(): |
| previous = initial[name] |
| delta = value.detach().cpu() - previous |
| rows[name] = dict(changed_elements=int(torch.count_nonzero(delta)), |
| numel=value.numel(), max_abs_delta=float(delta.abs().max()), |
| l2_delta=float(torch.linalg.vector_norm(delta))) |
| return dict(changed_elements=sum(r['changed_elements'] for r in rows.values()), |
| changed_tensors=sum(r['changed_elements'] > 0 for r in rows.values()), tensors=rows) |
|
|
|
|
| class PairedBatchStream: |
| """A deterministic, resumable stream with exactly 12 new + 4 replay rows. |
| |
| Alternate two independent new-data streams: return-choice pairs and |
| irrelevant-perturbation pairs. Each stream covers every new row exactly once |
| per epoch. In the return stream, unpaired original no-op rows are shuffled |
| into ordinary two-row sampling units, not semantic return-choice pairs. |
| Explicit within-batch unit positions prevent ambiguous grouping when an |
| epoch crossing happens to repeat a semantic key. |
| """ |
| def __init__(self, rows, seed): |
| if not rows or len({row['id'] for row in rows}) != len(rows): |
| raise ValueError('sampler requires unique, nonempty training rows') |
| self.row_ids = [row['id'] for row in rows] |
| self.seed = int(seed) |
| self.pairs = [] |
| self.singletons = [] |
| self.replay = [] |
| pairs = defaultdict(list) |
| invariance = defaultdict(list) |
| for index, row in enumerate(rows): |
| source = row.get('source') |
| key = row.get('pair_state_key') |
| if source == 'replay': |
| if key is not None: |
| raise ValueError('replay rows cannot be new return-pair units') |
| self.replay.append(index) |
| elif source == 'new': |
| invariant_key = row.get('invariance_state_key') |
| if not isinstance(invariant_key, str) or not invariant_key: |
| raise ValueError('every new state requires explicit invariance_state_key') |
| invariance[invariant_key].append(index) |
| if key is None: |
| if row['edits']: |
| raise ValueError('a new editing state has no explicit pair_state_key') |
| self.singletons.append(index) |
| elif isinstance(key, str) and key: |
| pairs[key].append(index) |
| else: |
| raise ValueError('pair_state_key must be a nonempty string or null') |
| else: |
| raise ValueError('training rows require explicit source=new/replay') |
| for key, indices in sorted(pairs.items()): |
| group = [rows[index] for index in indices] |
| if (len(group) != 2 or {row.get('variant') for row in group} != {0, 1} |
| or len({row.get('origin_id') for row in group}) != 1 |
| or any(row.get('origin_id') is None for row in group) |
| or len({row.get('state_mask') for row in group}) != 1 |
| or any(type(row.get('state_mask')) is not int for row in group) |
| or len({row.get('pair_id') for row in group}) != 1 |
| or any(row.get('pair_id') is None for row in group)): |
| raise ValueError('invalid explicit return-choice state pair: ' + key) |
| self.pairs.append(tuple(sorted(indices, key=lambda i: rows[i]['variant']))) |
| if len(self.singletons) % 2: |
| raise ValueError('the new unpaired no-op pool must have even cardinality') |
| self.new_count = 2 * len(self.pairs) + len(self.singletons) |
| self.invariance_pairs = [] |
| for key, indices in sorted(invariance.items()): |
| group = [rows[index] for index in indices] |
| if (len(group) != 2 or {row.get('nuisance_variant') for row in group} != {0, 1} |
| or len({row.get('origin_id') for row in group}) != 1 |
| or len({row.get('variant') for row in group}) != 1 |
| or len({row.get('state_mask') for row in group}) != 1 |
| or group[0]['edits'] != group[1]['edits'] |
| or len(group[0]['binary_bytes']) != len(group[1]['binary_bytes'])): |
| raise ValueError('invalid explicit irrelevant-perturbation pair: ' + key) |
| self.invariance_pairs.append(tuple(sorted(indices, key=lambda i: rows[i]['nuisance_variant']))) |
| if 2 * len(self.invariance_pairs) != self.new_count: |
| raise ValueError('invariance stream must cover the complete new pool') |
| if self.new_count < 12 or len(self.replay) < 4: |
| raise ValueError('at least 12 new rows and 4 replay rows are required') |
| self.streams = { |
| 'return': dict(rng=random.Random(self.seed), units=[], cursor=0, epochs=0), |
| 'invariance': dict(rng=random.Random(self.seed ^ 0xC0FFEE), units=[], cursor=0, epochs=0), |
| } |
| self.replay_rng = random.Random(self.seed ^ 0x5DEECE66D) |
| self.replay_order = [] |
| self.replay_cursor = self.replay_epochs = self.steps = 0 |
| self.last_units = [] |
| self.last_stream = None |
| self.rows = rows |
|
|
| def _new_epoch(self, kind): |
| stream = self.streams[kind] |
| if kind == 'return': |
| singles = self.singletons.copy() |
| stream['rng'].shuffle(singles) |
| units = [('return', pair) for pair in self.pairs] |
| units += [('unpaired_noop', tuple(singles[i:i + 2])) for i in range(0, len(singles), 2)] |
| else: |
| units = [('invariance', pair) for pair in self.invariance_pairs] |
| stream['rng'].shuffle(units) |
| stream['units'] = units |
| stream['cursor'] = 0 |
| stream['epochs'] += 1 |
|
|
| def next_batch(self): |
| kind = 'return' if self.steps % 2 == 0 else 'invariance' |
| stream = self.streams[kind] |
| new = [] |
| units = [] |
| for _ in range(6): |
| if stream['cursor'] == len(stream['units']): |
| self._new_epoch(kind) |
| unit_kind, pair = stream['units'][stream['cursor']] |
| key_field = 'pair_state_key' if unit_kind == 'return' else 'invariance_state_key' |
| units.append(dict(kind=unit_kind, positions=[len(new), len(new) + 1], |
| key=None if unit_kind == 'unpaired_noop' else self.rows[pair[0]][key_field])) |
| new.extend(pair) |
| stream['cursor'] += 1 |
| replay = [] |
| for _ in range(4): |
| if self.replay_cursor == len(self.replay_order): |
| self.replay_order = self.replay.copy() |
| self.replay_rng.shuffle(self.replay_order) |
| self.replay_cursor = 0 |
| self.replay_epochs += 1 |
| replay.append(self.replay_order[self.replay_cursor]) |
| self.replay_cursor += 1 |
| self.steps += 1 |
| self.last_units, self.last_stream = units, kind |
| return new + replay |
|
|
| def state_dict(self): |
| return dict(format=SAMPLER_FORMAT, seed=self.seed, |
| row_ids_sha256=hashlib.sha256(json_bytes(self.row_ids)).hexdigest(), |
| streams={name: dict(rng=value['rng'].getstate(), units=value['units'], |
| cursor=value['cursor'], epochs=value['epochs']) |
| for name, value in self.streams.items()}, |
| replay_rng=self.replay_rng.getstate(), replay_order=self.replay_order, |
| replay_cursor=self.replay_cursor, replay_epochs=self.replay_epochs, steps=self.steps) |
|
|
| def load_state_dict(self, state): |
| if (state['format'] != SAMPLER_FORMAT or state['seed'] != self.seed |
| or state['row_ids_sha256'] != hashlib.sha256(json_bytes(self.row_ids)).hexdigest()): |
| raise ValueError('sampler checkpoint belongs to a different stream') |
| for name, value in state['streams'].items(): |
| self.streams[name]['rng'].setstate(value['rng']) |
| for key in ('units', 'cursor', 'epochs'): |
| self.streams[name][key] = deepcopy(value[key]) |
| self.replay_rng.setstate(state['replay_rng']) |
| for key in ('replay_order', 'replay_cursor', 'replay_epochs', 'steps'): |
| setattr(self, key, deepcopy(state[key])) |
|
|
| def description(self): |
| return dict(format=SAMPLER_FORMAT, seed=self.seed, batch_size=16, |
| new_per_batch=12, replay_per_batch=4, replay_fraction=.25, |
| new_rows=self.new_count, replay_rows=len(self.replay), |
| semantic_pair_units=len(self.pairs), unpaired_noop_rows=len(self.singletons), |
| invariance_pair_units=len(self.invariance_pairs), |
| row_ids=self.row_ids, |
| row_ids_sha256=hashlib.sha256(json_bytes(self.row_ids)).hexdigest(), |
| sample_unit_note='Only explicit pair_state_key groups are return-choice pairs; ' |
| 'randomly grouped no-op rows are ordinary sampling units.', |
| state_mask_note='Nonzero masks represent corresponding legitimate repair subsets; ' |
| 'they need not differ only in the return instruction.', |
| sequence='Alternate return-pair and invariance-pair streams, starting with return. ' |
| 'Each independently cycles through ALL new rows, including original no-ops. ' |
| 'Six two-row new units plus four independently shuffled replay rows per update. ' |
| 'Every arm uses this same seeded prefix; metadata is never passed to the network.') |
|
|
|
|
| def make_initial_model(arm, parent_path, seed): |
| torch.manual_seed(seed) |
| random.seed(seed) |
| parent = load_relation(parent_path) |
| if parent.config.mode != 'dataflow' or parent.config.message_steps != 4: |
| raise ValueError('the shared source parent must be the original dataflow d4 model') |
| parent_parameters = dict(parent.named_parameters()) |
| if arm.kind == 'relation': |
| model = deepcopy(parent) |
| model.config = replace(parent.config, message_steps=arm.message_steps) |
| mapping = {name: name for name, _ in model.named_parameters()} |
| provided = {} |
| elif arm.kind == 'slot': |
| from lightnest.slot_patch import SlotPatchNetwork |
| model, provided = SlotPatchNetwork.from_relation(parent, slots=arm.slots, mode=arm.slot_mode, seed=seed) |
| if getattr(model.config, 'mode', 'memory') != arm.slot_mode: |
| model.config = replace(model.config, mode=arm.slot_mode) |
| mapping = provided.get('parameter_name_map') |
| if not isinstance(mapping, dict): |
| raise ValueError('slot import must supply its complete parameter_name_map') |
| else: |
| raise ValueError('unknown model kind') |
| parameters = dict(model.named_parameters()) |
| if set(mapping) != set(parameters): |
| raise ValueError('inheritance map must cover every current parameter exactly') |
| copied, initialized = {}, {} |
| for name, value in parameters.items(): |
| source = mapping[name] |
| if source is None: |
| initialized[name] = tensor_inventory({name: value})['tensors'][name] |
| else: |
| if (source not in parent_parameters or value.shape != parent_parameters[source].shape |
| or not torch.equal(value, parent_parameters[source])): |
| raise ValueError('a claimed inherited tensor is not an exact parent copy: ' + name) |
| copied[name] = dict(parent_parameter=source, |
| **tensor_inventory({name: value})['tensors'][name]) |
| report = dict(parent=file_record(parent_path), parent_config=asdict(parent.config), |
| configuration=asdict(model.config), parameter_name_map=mapping, |
| copied_parameters=copied, new_parameters=initialized, |
| copied_numel=sum(p['numel'] for p in copied.values()), |
| new_numel=sum(p['numel'] for p in initialized.values()), |
| omitted_parent_parameters=sorted(set(parent_parameters) - set(mapping.values())), |
| parameter_count=sum(p.numel() for p in parameters.values()), |
| trainable_parameter_count=sum(p.numel() for p in parameters.values() if p.requires_grad), |
| function_preserving=(arm.kind == 'relation' and arm.message_steps == 4), |
| optimizer='Fresh AdamW for each arm; old moments are not inherited. ' |
| 'The same optimizer continues across all four checkpoints.', |
| supplied_import_report=provided) |
| return model, report |
|
|
|
|
| def model_encode(model, kind, group): |
| if kind == 'relation': |
| encoder = relation_encode |
| elif kind == 'slot': |
| from lightnest.slot_patch import encode_batch as encoder |
| else: |
| raise ValueError('unknown model kind') |
| return encoder([row['command'] for row in group], [row['binary_bytes'] for row in group], model.config) |
|
|
|
|
| def predictor_for(model, kind): |
| if kind == 'relation': |
| return RelationPredictor(model, threshold=0.) |
| if kind == 'slot': |
| from lightnest.slot_patch import SlotPredictor |
| return SlotPredictor(model, threshold=0.) |
| raise ValueError('unknown model kind') |
|
|
|
|
| def build_predictor(spec): |
| """Read-only portable adapter for the independent evaluator. |
| |
| Normal inventory specs need id/kind/path. Explicit message_steps or |
| memory_mode fields are inference settings and must be recorded by the |
| evaluator; they do not mutate saved tensors. A supplied SHA is checked. |
| Original relation .pt parents are also accepted without copying them. |
| """ |
| path = resolve(spec['path']) |
| if spec.get('sha256') is not None and sha256(path) != spec['sha256']: |
| raise ValueError('predictor checkpoint digest mismatch') |
| if path.name.endswith('.pt.gz'): |
| model, payload = load_checkpoint(path) |
| if payload['kind'] != spec['kind']: |
| raise ValueError('predictor spec kind disagrees with checkpoint') |
| elif spec['kind'] == 'relation': |
| model = load_relation(path) |
| else: |
| raise ValueError('slot predictor requires a portable .pt.gz checkpoint') |
| if spec['kind'] == 'relation': |
| return RelationPredictor(model, threshold=0., message_steps=spec.get('message_steps')) |
| if spec.get('memory_mode') is not None: |
| model.config = replace(model.config, mode=spec['memory_mode']) |
| return predictor_for(model, 'slot') |
|
|
|
|
| def save_checkpoint(path, model, optimizer, sampler, *, kind, metadata): |
| path = Path(path) |
| state = {name: value.detach().cpu().clone() for name, value in model.state_dict().items()} |
| payload = dict(format=CHECKPOINT_FORMAT, architecture=model.architecture, kind=kind, |
| config=asdict(model.config), state_dict=state, |
| initialization_report=getattr(model, 'initialization_report', None), |
| optimizer=optimizer.state_dict() if optimizer is not None else None, |
| sampler=sampler.state_dict() if sampler is not None else None, |
| torch_rng_state=torch.get_rng_state(), python_rng_state=random.getstate(), |
| metadata=metadata, model_inventory=tensor_inventory(state)) |
| with path.open('xb') as raw: |
| with gzip.GzipFile(filename='', mode='wb', fileobj=raw, mtime=0, compresslevel=6) as compressed: |
| torch.save(payload, compressed) |
| return dict(**file_record(path), model_sha256=payload['model_inventory']['sha256']) |
|
|
|
|
| def load_checkpoint(path): |
| |
| |
| with gzip.open(path, 'rb') as source: |
| raw = source.read(32 * 1024 * 1024 + 1) |
| if len(raw) > 32 * 1024 * 1024: |
| raise ValueError('checkpoint expands beyond 32 MiB') |
| payload = torch.load(io.BytesIO(raw), map_location='cpu', weights_only=True) |
| if not isinstance(payload, dict) or payload.get('format') != CHECKPOINT_FORMAT: |
| raise ValueError('unknown dependency checkpoint format') |
| if payload['kind'] == 'relation': |
| model = RelationPatchNetwork(RelationConfig(**payload['config'])) |
| elif payload['kind'] == 'slot': |
| from lightnest.slot_patch import SlotConfig, SlotPatchNetwork |
| model = SlotPatchNetwork(SlotConfig(**payload['config'])) |
| else: |
| raise ValueError('unknown checkpoint model kind') |
| if payload['architecture'] != model.architecture: |
| raise ValueError('checkpoint architecture mismatch') |
| expected, actual = model.state_dict(), payload['state_dict'] |
| if set(expected) != set(actual): |
| raise ValueError('checkpoint tensor inventory mismatch') |
| for name, tensor in actual.items(): |
| if (not isinstance(tensor, torch.Tensor) or tensor.shape != expected[name].shape |
| or tensor.dtype != expected[name].dtype or not bool(torch.isfinite(tensor).all())): |
| raise ValueError('invalid checkpoint tensor: ' + name) |
| if tensor_inventory(actual) != payload['model_inventory']: |
| raise ValueError('checkpoint canonical tensor digest mismatch') |
| model.load_state_dict(actual) |
| if hasattr(model, 'initialization_report'): |
| model.initialization_report = payload.get('initialization_report') |
| return model, payload |
|
|
|
|
| def load_model(path): |
| """Portable inference entry point for a selected/last .pt.gz checkpoint.""" |
| return load_checkpoint(path)[0] |
|
|
|
|
| def relation_objectives(scores, group, units, none_index, *, margin=1.): |
| """Pair metadata enters only this external supervised objective. |
| |
| Return pairs compare each side's exclusive acceptable positions with the |
| opposite side's exclusive positions. Log-softmax denominators cancel, so |
| differences of masked raw-logit logsumexp values are used. Invariance JS |
| includes every actual byte plus STOP, excludes padding, and is defined only |
| for equal-length aligned programs. Each term is a mean over valid pairs; |
| the return term first averages the two directional losses for a pair. |
| """ |
| return_terms, invariance_terms = [], [] |
| counts = dict(return_candidate_pairs=0, return_effective_pairs=0, return_effective_directions=0, |
| invariance_candidate_pairs=0, invariance_effective_pairs=0, unpaired_noop_units=0) |
| for unit in units: |
| left, right = unit['positions'] |
| a, b = group[left], group[right] |
| if unit['kind'] == 'unpaired_noop': |
| counts['unpaired_noop_units'] += 1 |
| continue |
| if unit['kind'] == 'return': |
| counts['return_candidate_pairs'] += 1 |
| if a['pair_state_key'] != unit['key'] or b['pair_state_key'] != unit['key']: |
| raise ValueError('explicit return unit does not match batch rows') |
| positions_a = {edit['offset'] for edit in a['edits']} |
| positions_b = {edit['offset'] for edit in b['edits']} |
| only_a, only_b = sorted(positions_a - positions_b), sorted(positions_b - positions_a) |
| if not only_a or not only_b: |
| continue |
| if max(only_a + only_b) >= min(len(a['binary_bytes']), len(b['binary_bytes'])): |
| raise ValueError('return pair positions are not aligned valid bytes') |
| first = torch.nn.functional.softplus(margin + torch.logsumexp(scores[left, only_b], 0) |
| - torch.logsumexp(scores[left, only_a], 0)) |
| second = torch.nn.functional.softplus(margin + torch.logsumexp(scores[right, only_a], 0) |
| - torch.logsumexp(scores[right, only_b], 0)) |
| return_terms.append((first + second) * .5) |
| counts['return_effective_pairs'] += 1 |
| counts['return_effective_directions'] += 2 |
| elif unit['kind'] == 'invariance': |
| counts['invariance_candidate_pairs'] += 1 |
| length = len(a['binary_bytes']) |
| if (a['invariance_state_key'] != unit['key'] or b['invariance_state_key'] != unit['key'] |
| or length != len(b['binary_bytes']) or a['edits'] != b['edits']): |
| raise ValueError('invariance JS requires a matching aligned state pair') |
| positions = list(range(length)) + [none_index] |
| log_p = scores[left, positions].log_softmax(0) |
| log_q = scores[right, positions].log_softmax(0) |
| log_m = torch.logaddexp(log_p, log_q) - math.log(2.) |
| divergence = .5 * ((log_p.exp() * (log_p - log_m)).sum() |
| + (log_q.exp() * (log_q - log_m)).sum()) |
| invariance_terms.append(divergence) |
| counts['invariance_effective_pairs'] += 1 |
| else: |
| raise ValueError('unknown supervised sampling unit') |
| zero = scores.sum() * 0. |
| return (torch.stack(return_terms).mean() if return_terms else zero, |
| torch.stack(invariance_terms).mean() if invariance_terms else zero, counts) |
|
|
|
|
| def training_step(model, kind, rows, sampler, optimizer, *, byte_weight=2., gradient_clip=2., |
| return_weight=.2, invariance_weight=.1, return_margin=1.): |
| indices = sampler.next_batch() |
| group = [rows[index] for index in indices] |
| encoded = model_encode(model, kind, group) |
| model.train() |
| optimizer.zero_grad(set_to_none=True) |
| scores, bits = model(*encoded) |
| base, pointer, byte = supervised_loss(scores, bits, group, model.none_index, byte_weight) |
| return_loss, invariance_loss, counts = relation_objectives( |
| scores, group, sampler.last_units, model.none_index, margin=return_margin) |
| loss = base + return_weight * return_loss + invariance_weight * invariance_loss |
| if not bool(torch.isfinite(loss)): |
| raise ValueError('nonfinite training objective') |
| loss.backward() |
| norm = nn.utils.clip_grad_norm_(model.parameters(), gradient_clip, error_if_nonfinite=True) |
| optimizer.step() |
| return dict(indices=indices, stream=sampler.last_stream, units=deepcopy(sampler.last_units), |
| loss=float(loss.detach()), base_loss=float(base.detach()), pointer_loss=float(pointer.detach()), |
| byte_loss=float(byte.detach()), return_margin_loss=float(return_loss.detach()), |
| invariance_js_loss=float(invariance_loss.detach()), auxiliary_counts=counts, |
| gradient_norm_before_clip=float(norm)) |
|
|
|
|
| def run_cpu_budget(step_fn: Callable, checkpoint_fn: Callable, *, cpu_seconds, |
| on_step: Callable | None = None, cpu_clock=time.process_time, |
| wall_clock=time.perf_counter): |
| """Charge only timed complete training steps; never truncate an update. |
| |
| Evaluations and writes happen outside each timed interval. Boundary |
| overshoot is bounded by the final indivisible update's measured CPU cost. |
| A very slow update can cross several boundaries; each is recorded explicitly |
| at that same state rather than inventing intermediate weights. |
| """ |
| if not isinstance(cpu_seconds, (int, float)) or not math.isfinite(cpu_seconds) or cpu_seconds <= 0: |
| raise ValueError('positive finite training CPU budget required') |
| elapsed_cpu = elapsed_wall = 0. |
| step = stage = 0 |
| checkpoints = [] |
| while elapsed_cpu < cpu_seconds: |
| before_cpu, before_wall = cpu_clock(), wall_clock() |
| details = step_fn() |
| delta_cpu, delta_wall = cpu_clock() - before_cpu, wall_clock() - before_wall |
| if delta_cpu <= 0 or delta_wall < 0: |
| raise ValueError('training clocks did not advance monotonically') |
| elapsed_cpu += delta_cpu |
| elapsed_wall += delta_wall |
| step += 1 |
| record = dict(step=step, training_cpu_seconds=elapsed_cpu, training_wall_seconds=elapsed_wall, |
| step_cpu_seconds=delta_cpu, step_wall_seconds=delta_wall, **details) |
| if on_step is not None: |
| on_step(record) |
| while stage < len(FRACTIONS) and elapsed_cpu >= cpu_seconds * FRACTIONS[stage]: |
| boundary = dict(stage=stage + 1, fraction=FRACTIONS[stage], step=step, |
| target_training_cpu_seconds=cpu_seconds * FRACTIONS[stage], |
| training_cpu_seconds=elapsed_cpu, training_wall_seconds=elapsed_wall, |
| cpu_overshoot_seconds=elapsed_cpu - cpu_seconds * FRACTIONS[stage], |
| crossing_step_cpu_seconds=delta_cpu) |
| checkpoints.append(checkpoint_fn(boundary)) |
| stage += 1 |
| if len(checkpoints) != 4: |
| raise AssertionError('all four predeclared checkpoints must be saved') |
| return dict(actual_updates=step, examples_seen=16 * step, new_examples_seen=12 * step, |
| replay_examples_seen=4 * step, training_cpu_seconds=elapsed_cpu, |
| training_wall_seconds=elapsed_wall, cpu_budget_seconds=cpu_seconds, |
| budget_overshoot_seconds=elapsed_cpu - cpu_seconds, checkpoints=checkpoints) |
|
|
|
|
| @torch.inference_mode() |
| def development_objective(model, kind, rows, *, batch_size=16, byte_weight=2.): |
| model.eval() |
| pointer_sum = byte_sum = 0. |
| edit_count = exact = correct_stop = false_edit = 0 |
| for start in range(0, len(rows), batch_size): |
| group = rows[start:start + batch_size] |
| scores, bits = model(*model_encode(model, kind, group)) |
| loss, pointer, byte = supervised_loss(scores, bits, group, model.none_index, byte_weight) |
| if not bool(torch.isfinite(loss)): |
| raise ValueError('nonfinite development objective') |
| edits = sum(len(row['edits']) for row in group) |
| pointer_sum += float(pointer) * len(group) |
| byte_sum += float(byte) * edits |
| edit_count += edits |
| for index, position in enumerate(scores.argmax(-1).tolist()): |
| row = group[index] |
| stop = position == model.none_index |
| good = not row['edits'] and stop |
| if not stop: |
| value = sum(int(bits[index, position, k] >= 0) << k for k in range(8)) |
| good = any(edit['offset'] == position and edit['value'] == value for edit in row['edits']) |
| exact += good |
| correct_stop += not row['edits'] and stop |
| false_edit += not row['edits'] and not stop |
| pointer_loss = pointer_sum / len(rows) |
| byte_loss = byte_sum / edit_count if edit_count else 0. |
| return dict(loss=pointer_loss + byte_weight * byte_loss, pointer_loss=pointer_loss, |
| byte_loss=byte_loss, records=len(rows), target_edit_count=edit_count, |
| one_step_exact=exact, one_step_correct_stop=correct_stop, |
| one_step_false_edit=false_edit, |
| interpretation='One-step supervised development objective; not episode success.') |
|
|
|
|
| def load_episodes(path): |
| rows = json.loads(public_path(path).read_text()) |
| if not isinstance(rows, list) or not rows or len({row['id'] for row in rows}) != len(rows): |
| raise ValueError('episode file must contain a nonempty list with unique IDs') |
| return rows |
|
|
|
|
| def score_public(model, kind, development, episodes, output, *, predictor_id, byte_weight=2.): |
| from experiments.architecture_eval import score_episode |
| from experiments.morphology_data import summarize_scored |
| output = fresh_directory(output) |
| torch_rng, python_rng = torch.get_rng_state(), random.getstate() |
| before = tensor_inventory(model.state_dict())['sha256'] |
| start_cpu, start_wall = time.process_time(), time.perf_counter() |
| objective = development_objective(model, kind, development, byte_weight=byte_weight) |
| predictor = predictor_for(model, kind) |
| summaries, records = {}, {} |
| for split in ('development', 'selection', 'retention'): |
| scored = [] |
| path = output / (split + '.jsonl.gz') |
| with path.open('xb') as raw: |
| with gzip.GzipFile(filename='', mode='wb', fileobj=raw, mtime=0, compresslevel=6) as compressed: |
| for episode in episodes[split]: |
| result = score_episode(predictor, episode, predictor_id=predictor_id, max_steps=8) |
| scored.append(result) |
| compact = {key: value for key, value in result.items() if key not in ('audit', 'final_hex')} |
| initial, final = bytes.fromhex(episode['binary']), bytes.fromhex(result['final_hex']) |
| compact['predicted_delta'] = [[i, y] for i, (x, y) in enumerate(zip(initial, final)) if x != y] |
| compact['final_bytes'] = len(final) |
| compressed.write(json_bytes(compact)) |
| summaries[split] = summarize_scored(episodes[split], scored) |
| records[split] = file_record(path) |
| if before != tensor_inventory(model.state_dict())['sha256']: |
| raise ValueError('evaluation changed model tensors') |
| rng_changed = not torch.equal(torch_rng, torch.get_rng_state()) or python_rng != random.getstate() |
| torch.set_rng_state(torch_rng) |
| random.setstate(python_rng) |
| result = dict(development_objective=objective, summaries=summaries, score_files=records, |
| model_sha256=before, evaluation_cpu_seconds=time.process_time() - start_cpu, |
| evaluation_wall_seconds=time.perf_counter() - start_wall, |
| evaluation_rng_changed_and_restored=rng_changed) |
| write_json(output / 'summary.json', result) |
| return result |
|
|
|
|
| def candidate_metrics(evaluation): |
| selection = evaluation['summaries']['selection']['all'] |
| retention = evaluation['summaries']['retention']['all'] |
| return dict(edit_total=selection.get('edit_total', 0), edit_strict=selection.get('edit_strict', 0), |
| pair_total=selection.get('pair_total', 0), pair_both_strict=selection.get('pair_both_strict', 0), |
| noop_total=selection.get('noop_total', 0), noop_strict=selection.get('noop_strict', 0), |
| false_edit_on_correct=selection.get('false_edit_on_correct', 0), |
| invariance_edit_total=selection.get('invariance_edit_total', 0), |
| invariance_edit_both_strict=selection.get('invariance_edit_both_strict', 0), |
| retention_total=retention['total'], retention_strict=retention.get('strict_success', 0), |
| development_loss=evaluation['development_objective']['loss']) |
|
|
|
|
| def promotion_decision(candidate, incumbent, original_anchor, policy): |
| """Loss alone never promotes a checkpoint; the original retention floor stays fixed.""" |
| counts = ('edit_total', 'pair_total', 'noop_total', 'invariance_edit_total', 'retention_total') |
| if any(candidate[key] != incumbent[key] or candidate[key] != original_anchor[key] for key in counts): |
| raise ValueError('candidate and anchor denominators differ') |
| if candidate['retention_total'] != policy['retention_denominator']: |
| raise ValueError('retention denominator differs from the frozen policy') |
| reasons = [] |
| floor = max(0, original_anchor['retention_strict'] - policy['retention_tolerance_cases']) |
| if candidate['retention_strict'] < floor: |
| reasons.append('below_original_parent_retention_floor') |
| for metric in policy['nondecreasing_metrics']: |
| if candidate[metric] < incumbent[metric]: |
| reasons.append(metric + '_decreased') |
| if candidate['false_edit_on_correct'] > incumbent['false_edit_on_correct']: |
| reasons.append('false_edits_increased') |
| if not any(candidate[key] > incumbent[key] for key in policy['strict_improvement_metrics']): |
| reasons.append('no_strict_edit_or_pair_improvement') |
| return dict(accept=not reasons, reasons=reasons, retention_floor=floor, |
| retention_tolerance_cases=policy['retention_tolerance_cases'], |
| candidate_metrics=candidate, incumbent_metrics=incumbent, |
| original_anchor_metrics=original_anchor) |
|
|
|
|
| DEFAULT_SELECTION = dict( |
| nondecreasing_metrics=['edit_strict', 'pair_both_strict', 'invariance_edit_both_strict'], |
| strict_improvement_metrics=['edit_strict', 'pair_both_strict'], |
| retention_denominator=128, retention_tolerance_cases=2, |
| loss_role='Development loss is diagnostic only in this single-trajectory driver. ' |
| 'It cannot by itself promote a stage; an exact metric tie keeps the earlier checkpoint.', |
| fallback='If no stage qualifies, the selected model is explicitly the original E1/d4 parent, ' |
| 'even for a slot or deeper arm. The arm trajectory still trains to its full CPU budget.', |
| ) |
|
|
|
|
| def make_protocol(data_dir, output, *, parent_template, cpu_seconds=200., threads=2, |
| seeds=SEEDS, artifact_root='work/morphology-20261017/search/dependency'): |
| """Freeze only after the coordinator reviews the public 24-step pilot.""" |
| if not math.isfinite(cpu_seconds) or not 0 < cpu_seconds <= 3600: |
| raise ValueError('base CPU budget must be in (0, 3600] seconds') |
| data_dir = public_path(data_dir) |
| data_names = dict(train='train.jsonl', development='development.jsonl', |
| development_episodes='development-episodes.json', selection='selection.json', |
| retention='retention.json', data_protocol='protocol.json') |
| data = {key: file_record(public_path(data_dir / name)) for key, name in data_names.items()} |
| rows = load_records(public_path(data['train']['path'])) |
| sampler = PairedBatchStream(rows, seeds[0]) |
| if len(load_episodes(data['retention']['path'])) != DEFAULT_SELECTION['retention_denominator']: |
| raise ValueError('expected 128 public retention episodes') |
| parents = {str(seed): file_record(public_path(parent_template.format(seed=seed))) for seed in seeds} |
| files = ('experiments/dependency_train.py', 'lightnest/relation_patch.py', 'lightnest/slot_patch.py', |
| 'lightnest/iterative_patch.py', 'lightnest/rv32.py', 'experiments/architecture_eval.py', |
| 'experiments/morphology_data.py', 'experiments/architecture_data.py') |
| sources = {path: file_record(ROOT / path) for path in files} |
| value = dict(format=FORMAT, status='frozen', created_utc=datetime.now(timezone.utc).isoformat(), |
| baseline_commit=BASELINE_COMMIT, |
| current_git_commit=subprocess.check_output(['git', 'rev-parse', 'HEAD'], cwd=ROOT, text=True).strip(), |
| seeds=list(seeds), arms={name: dict(**asdict(arm), training_cpu_seconds=cpu_seconds * arm.cpu_multiplier) |
| for name, arm in ARMS.items()}, |
| data=data, parents=parents, sources=sources, artifact_root=artifact_root, |
| training=dict(batch_size=16, new_per_batch=12, replay_per_batch=4, learning_rate=.001, |
| optimizer='AdamW', weight_decay=.01, gradient_clip=2., byte_weight=2., |
| return_weight=.2, return_margin=1., invariance_weight=.1, |
| return_loss='Mean over effective explicit return units of the mean of two directional ' |
| 'softplus(1 + logsumexp(opposite exclusive remaining positions) - ' |
| 'logsumexp(own exclusive remaining positions)); both exclusive sets must be nonempty.', |
| invariance_loss='Mean pointer Jensen-Shannon divergence over effective aligned invariance ' |
| 'units; all real bytes and STOP included, padding excluded. No byte-bit consistency term.', |
| auxiliary_normalization='Each term independently averages effective pairs; empty terms are zero. ' |
| 'Return and invariance grouping alternate by update; metadata is never a model input.', |
| threads=threads, interop_threads=1, fractions=list(FRACTIONS), |
| optimizer_initialization='Fresh moments for every arm; no resets between checkpoints.', |
| budget_measurement='Sum of process_time intervals around sampler, encoding, train forward, ' |
| 'set-valued loss, backward, gradient clipping and optimizer update. ' |
| 'Evaluation, serialization, and per-step log writes are outside these intervals.', |
| stopping='After the first complete update crossing 100 percent; indivisible-update overshoot recorded.', |
| exposure_claim='Same predeclared seeded sample-stream prefix; CPU budgets match, update counts do not.'), |
| selection=deepcopy(DEFAULT_SELECTION), sampler=sampler.description(), |
| runtime=dict(python=platform.python_version(), torch=str(torch.__version__)), |
| formal_authority='Coordinator freezes this file after the public pilot; no data-dependent budget extension.') |
| |
| value['sampler'].pop('row_ids') |
| value['sampler'].pop('seed') |
| write_json(output, value) |
| return value |
|
|
|
|
| def read_protocol(path): |
| protocol = json.loads(Path(path).read_text()) |
| if protocol.get('format') != FORMAT or protocol.get('status') != 'frozen': |
| raise ValueError('formal training requires a frozen dependency protocol') |
| if protocol['training']['fractions'] != list(FRACTIONS) or protocol['training']['batch_size'] != 16: |
| raise ValueError('checkpoint fractions and batch size differ from the reviewed driver') |
| for entry in protocol['sources'].values(): |
| verify_record(entry) |
| for entry in protocol['data'].values(): |
| verify_record(entry, public=True) |
| for entry in protocol['parents'].values(): |
| verify_record(entry, public=True) |
| return protocol |
|
|
|
|
| def public_inputs(protocol): |
| data = protocol['data'] |
| train = load_records(public_path(data['train']['path'])) |
| development = load_records(public_path(data['development']['path'])) |
| episodes = {key: load_episodes(data['development_episodes' if key == 'development' else key]['path']) |
| for key in ('development', 'selection', 'retention')} |
| train_episodes = {row.get('episode_id', row['id']) for row in train} |
| for name, group in episodes.items(): |
| if train_episodes.intersection(row.get('episode_id', row['id']) for row in group): |
| raise ValueError('training episode appears in public evaluation split: ' + name) |
| if {row['id'] for row in episodes['development']}.intersection(row['id'] for row in episodes['selection']): |
| raise ValueError('development and selection episodes overlap') |
| return train, development, episodes |
|
|
|
|
| def prepare_anchor(protocol, protocol_sha, seed, development, episodes): |
| """One common original-parent anchor per seed, independent of arm initialization.""" |
| directory = resolve(protocol['artifact_root']) / 'anchors' / ('seed' + str(seed)) |
| path = directory / 'anchor.json' |
| if path.exists(): |
| value = json.loads(path.read_text()) |
| if value['protocol_sha256'] != protocol_sha or value['parent'] != protocol['parents'][str(seed)]: |
| raise ValueError('existing anchor belongs to a different frozen protocol') |
| for record in value['evaluation']['score_files'].values(): |
| verify_record(record) |
| verify_record(value['checkpoint']) |
| return value |
| fresh_directory(directory) |
| model = load_relation(verify_record(protocol['parents'][str(seed)])) |
| if model.config.mode != 'dataflow' or model.config.message_steps != 4: |
| raise ValueError('retention anchor must be the original E1/d4 parent') |
| checkpoint = save_checkpoint(directory / 'original-parent.pt.gz', model, None, None, |
| kind='relation', metadata=dict(role='shared-original-parent', seed=seed, |
| protocol_sha256=protocol_sha, actual_updates=0)) |
| evaluation = score_public(model, 'relation', development, episodes, directory / 'evaluation', |
| predictor_id='original-parent-seed' + str(seed), |
| byte_weight=protocol['training']['byte_weight']) |
| value = dict(protocol_sha256=protocol_sha, seed=seed, parent=protocol['parents'][str(seed)], |
| checkpoint=checkpoint, evaluation=evaluation, metrics=candidate_metrics(evaluation)) |
| write_json(path, value) |
| return value |
|
|
|
|
| def run_formal(protocol_path, arm_name, seed, output=None): |
| protocol = read_protocol(protocol_path) |
| protocol_sha = sha256(protocol_path) |
| if str(seed) not in protocol['parents'] or arm_name not in protocol['arms']: |
| raise ValueError('arm/seed is absent from the frozen protocol') |
| spec = protocol['arms'][arm_name] |
| arm = Arm(**{key: spec[key] for key in Arm.__dataclass_fields__}) |
| training = protocol['training'] |
| configure_threads(training['threads']) |
| out = fresh_directory(output or (resolve(protocol['artifact_root']) / f'{arm_name}-seed{seed}')) |
| begin_wall, begin_cpu = time.perf_counter(), time.process_time() |
| rows, development, episodes = public_inputs(protocol) |
| anchor = prepare_anchor(protocol, protocol_sha, seed, development, episodes) |
| model, inheritance = make_initial_model(arm, verify_record(protocol['parents'][str(seed)]), seed) |
| sampler = PairedBatchStream(rows, seed) |
| optimizer = torch.optim.AdamW(model.parameters(), lr=training['learning_rate'], weight_decay=training['weight_decay']) |
| initial_state = {name: value.detach().cpu().clone() for name, value in model.state_dict().items()} |
| initial = save_checkpoint(out / 'initial.pt.gz', model, optimizer, sampler, kind=arm.kind, |
| metadata=dict(role='initial', arm=arm_name, seed=seed, protocol_sha256=protocol_sha, |
| actual_updates=0, inheritance=inheritance)) |
| write_json(out / 'inheritance.json', inheritance) |
| write_json(out / 'sampling.json', sampler.description()) |
| write_json(out / 'start.json', dict(protocol=file_record(protocol_path), arm=spec, seed=seed, |
| original_anchor=reference(resolve(protocol['artifact_root']) / 'anchors' / f'seed{seed}' / 'anchor.json'), |
| initial=initial, started_utc=datetime.now(timezone.utc).isoformat())) |
| selected = dict(checkpoint=anchor['checkpoint'], metrics=anchor['metrics'], stage=0, |
| kind='relation', fallback_original_parent=True) |
| phase_results = [] |
| log_path = out / 'training-steps.jsonl.gz' |
| log_raw = log_path.open('xb') |
| log = gzip.GzipFile(filename='', mode='wb', fileobj=log_raw, mtime=0, compresslevel=6) |
| exposure_digest = hashlib.sha256() |
| loss_counts = defaultdict(int) |
|
|
| def on_step(record): |
| ids = [rows[index]['id'] for index in record['indices']] |
| exposure_digest.update(json_bytes(ids)) |
| record['batch_row_ids_sha256'] = hashlib.sha256(json_bytes(ids)).hexdigest() |
| record['row_exposure_prefix_sha256'] = exposure_digest.hexdigest() |
| for key, count in record['auxiliary_counts'].items(): |
| loss_counts[key] += count |
| log.write(json_bytes(record)) |
|
|
| def checkpoint(boundary): |
| nonlocal selected |
| stage = boundary['stage'] |
| saved = save_checkpoint(out / f'checkpoint-{stage:02d}.pt.gz', model, optimizer, sampler, |
| kind=arm.kind, metadata=dict(role='budget-stage', arm=arm_name, seed=seed, |
| protocol_sha256=protocol_sha, actual_updates=boundary['step'], |
| **boundary)) |
| evaluation = score_public(model, arm.kind, development, episodes, out / f'evaluation-{stage:02d}', |
| predictor_id=f'{arm_name}-seed{seed}-stage{stage}', |
| byte_weight=training['byte_weight']) |
| metrics = candidate_metrics(evaluation) |
| decision = promotion_decision(metrics, selected['metrics'], anchor['metrics'], protocol['selection']) |
| if decision['accept']: |
| selected = dict(checkpoint=saved, metrics=metrics, stage=stage, kind=arm.kind, |
| fallback_original_parent=False) |
| result = dict(**boundary, checkpoint=saved, metrics=metrics, decision=decision, |
| selected_stage_after=selected['stage'], evaluation=evaluation, |
| actual_parameter_changes=changed_inventory(initial_state, model.state_dict())) |
| write_json(out / f'stage-{stage:02d}.json', result) |
| phase_results.append(result) |
| print(json.dumps(dict(event='checkpoint', arm=arm_name, seed=seed, stage=stage, |
| updates=boundary['step'], training_cpu_seconds=boundary['training_cpu_seconds'], |
| promoted=decision['accept'], rejection_reasons=decision['reasons'])), flush=True) |
| return dict(stage=stage, checkpoint=saved, step=boundary['step']) |
|
|
| try: |
| budget = run_cpu_budget( |
| lambda: training_step(model, arm.kind, rows, sampler, optimizer, |
| byte_weight=training['byte_weight'], gradient_clip=training['gradient_clip'], |
| return_weight=training['return_weight'], invariance_weight=training['invariance_weight'], |
| return_margin=training['return_margin']), |
| checkpoint, cpu_seconds=spec['training_cpu_seconds'], on_step=on_step) |
| finally: |
| log.close() |
| log_raw.close() |
| |
| if sha256(protocol_path) != protocol_sha: |
| raise ValueError('protocol changed during training') |
| read_protocol(protocol_path) |
| last = dict(checkpoint=phase_results[-1]['checkpoint'], metrics=phase_results[-1]['metrics'], |
| stage=4, kind=arm.kind, fallback_original_parent=False) |
| write_json(out / 'selected.json', selected) |
| write_json(out / 'last.json', last) |
| report = dict(format=FORMAT, complete=True, arm=spec, seed=seed, protocol=file_record(protocol_path), |
| initial=initial, selected=selected, last=last, budget=budget, |
| parameter_count=inheritance['parameter_count'], inheritance=file_record(out / 'inheritance.json'), |
| sampling=file_record(out / 'sampling.json'), training_steps=file_record(log_path), |
| row_exposure_prefix_sha256=exposure_digest.hexdigest(), auxiliary_pair_counts=dict(loss_counts), |
| actual_parameter_changes=changed_inventory(initial_state, model.state_dict()), |
| total_wall_seconds=time.perf_counter() - begin_wall, |
| total_cpu_seconds=time.process_time() - begin_cpu, |
| peak_rss_kib=resource.getrusage(resource.RUSAGE_SELF).ru_maxrss, |
| limitations=['CPU time is matched; optimizer steps and sample exposure counts can differ.', |
| 'Four stage checkpoints are evaluated; the trajectory continues even after a rejected stage.', |
| 'Retention permits up to two fewer strict successes among 128 public old-task cases.', |
| 'Development/selection/retention are public inputs, not independent final test evidence.', |
| 'Slot imports preserve only explicitly mapped tensors, not the original model function.']) |
| write_json(out / 'report.json', report) |
| return report |
|
|
|
|
| def run_pilot(train_path, parent_path, output, *, arm_name, seed=104729, steps=24, threads=2): |
| """A public fixed-step throughput measurement; never reads selection or final data.""" |
| if steps != 24: |
| raise ValueError('the authorized public pilot is exactly 24 updates') |
| if arm_name not in ARMS or arm_name == 'relation-d4-long': |
| raise ValueError('pilot only distinct computational arms; d4-long reuses the d4 throughput estimate') |
| configure_threads(threads) |
| out = fresh_directory(output) |
| train_path, parent_path = public_path(train_path), public_path(parent_path) |
| inputs = dict(train=file_record(train_path), parent=file_record(parent_path), |
| driver=file_record(__file__), relation_source=file_record(ROOT / 'lightnest/relation_patch.py')) |
| if ARMS[arm_name].kind == 'slot': |
| inputs['slot_source'] = file_record(ROOT / 'lightnest/slot_patch.py') |
| rows = load_records(train_path) |
| arm = ARMS[arm_name] |
| model, inheritance = make_initial_model(arm, parent_path, seed) |
| sampler = PairedBatchStream(rows, seed) |
| optimizer = torch.optim.AdamW(model.parameters(), lr=.001, weight_decay=.01) |
| initial_state = {key: value.detach().clone() for key, value in model.state_dict().items()} |
| initial = save_checkpoint(out / 'initial.pt.gz', model, optimizer, sampler, kind=arm.kind, |
| metadata=dict(role='pilot-initial', arm=arm_name, seed=seed, actual_updates=0)) |
| write_json(out / 'sampling.json', sampler.description()) |
| write_json(out / 'inheritance.json', inheritance) |
| records = [] |
| exposure_digest = hashlib.sha256() |
| loss_counts = defaultdict(int) |
| for step in range(1, steps + 1): |
| before_cpu, before_wall = time.process_time(), time.perf_counter() |
| value = training_step(model, arm.kind, rows, sampler, optimizer) |
| record = dict(step=step, training_cpu_seconds=time.process_time() - before_cpu, |
| training_wall_seconds=time.perf_counter() - before_wall, **value) |
| ids = [rows[index]['id'] for index in record['indices']] |
| exposure_digest.update(json_bytes(ids)) |
| record['batch_row_ids_sha256'] = hashlib.sha256(json_bytes(ids)).hexdigest() |
| record['row_exposure_prefix_sha256'] = exposure_digest.hexdigest() |
| for key, count in record['auxiliary_counts'].items(): |
| loss_counts[key] += count |
| records.append(record) |
| checkpoint = save_checkpoint(out / 'last.pt.gz', model, optimizer, sampler, kind=arm.kind, |
| metadata=dict(role='pilot-last', arm=arm_name, seed=seed, actual_updates=steps)) |
| cpu = sum(row['training_cpu_seconds'] for row in records) |
| wall = sum(row['training_wall_seconds'] for row in records) |
| write_json(out / 'training-steps.json', records) |
| for record in inputs.values(): |
| verify_record(record) |
| report = dict(format=FORMAT, role='public-throughput-pilot-not-formal', arm=asdict(arm), seed=seed, |
| actual_updates=steps, batch_size=16, examples_seen=16 * steps, |
| training=dict(threads=threads, interop_threads=1, learning_rate=.001, weight_decay=.01, |
| byte_weight=2., return_weight=.2, invariance_weight=.1, return_margin=1., |
| gradient_clip=2., optimizer='AdamW', optimizer_initialization='Fresh; never inherited from pilot.'), |
| runtime=dict(python=platform.python_version(), torch=str(torch.__version__)), |
| training_cpu_seconds=cpu, training_wall_seconds=wall, cpu_seconds_per_update=cpu / steps, |
| projected_updates_per_200_cpu_seconds=200. * steps / cpu, |
| projected_updates_per_600_cpu_seconds=600. * steps / cpu, |
| inputs=inputs, initial=initial, last=checkpoint, |
| parameter_count=inheritance['parameter_count'], |
| row_exposure_prefix_sha256=exposure_digest.hexdigest(), auxiliary_pair_counts=dict(loss_counts), |
| checkpoint_bytes=initial['bytes'] + checkpoint['bytes'], |
| actual_parameter_changes=changed_inventory(initial_state, model.state_dict()), |
| evaluation='None. The pilot reads training rows and original parent weights only.', |
| limitation='24 updates estimate throughput and checkpoint size; they do not establish converged performance.') |
| write_json(out / 'report.json', report) |
| return report |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description=__doc__) |
| sub = parser.add_subparsers(dest='command', required=True) |
| pilot = sub.add_parser('pilot') |
| pilot.add_argument('--train', required=True) |
| pilot.add_argument('--parent', required=True) |
| pilot.add_argument('--output', required=True) |
| pilot.add_argument('--arm', choices=list(ARMS), required=True) |
| pilot.add_argument('--seed', type=int, default=104729) |
| pilot.add_argument('--threads', type=int, choices=(1, 2), default=2) |
| freeze = sub.add_parser('freeze') |
| freeze.add_argument('--data-dir', required=True) |
| freeze.add_argument('--parent-template', required=True) |
| freeze.add_argument('--output', required=True) |
| freeze.add_argument('--cpu-seconds', type=float, required=True) |
| freeze.add_argument('--threads', type=int, choices=(1, 2), default=2) |
| freeze.add_argument('--artifact-root', default='work/morphology-20261017/search/dependency') |
| run = sub.add_parser('run') |
| run.add_argument('--protocol', required=True) |
| run.add_argument('--arm', choices=list(ARMS), required=True) |
| run.add_argument('--seed', type=int, choices=SEEDS, required=True) |
| run.add_argument('--output') |
| args = parser.parse_args() |
| if args.command == 'pilot': |
| result = run_pilot(args.train, args.parent, args.output, arm_name=args.arm, seed=args.seed, threads=args.threads) |
| elif args.command == 'freeze': |
| result = make_protocol(args.data_dir, args.output, parent_template=args.parent_template, |
| cpu_seconds=args.cpu_seconds, threads=args.threads, artifact_root=args.artifact_root) |
| else: |
| result = run_formal(args.protocol, args.arm, args.seed, args.output) |
| print(json.dumps(result, ensure_ascii=False, sort_keys=True, indent=2, allow_nan=False)) |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
