| """Checkpointed NoNE expert/layer geometry growth. |
| |
| This module is an explicit checkpoint I/O boundary, not a model hot path. |
| It preserves every trained source tensor region exactly, adds trainable |
| expert/layer capacity deterministically, migrates named AdamW moments, and |
| writes a non-promotable candidate receipt. The migrated graph must still train, |
| cold reload, and pass held-out proof before it can replace a promoted state. |
| """ |
| from __future__ import annotations |
|
|
| import copy |
| import hashlib |
| import json |
| import os |
| import re |
| import time |
| from pathlib import Path |
| from typing import Any, cast |
|
|
| import torch |
|
|
| MIGRATION_SCHEMA = "nnf.resynthesis.none_geometry_migration.v1" |
| CHECKPOINT_SCHEMA = "nnf.resynthesis.additive_state.v1" |
| GROWTH_PLAN_SCHEMA = "nnf.resynthesis.none_growth_plan.v2" |
| LEGACY_GROWTH_PLAN_SCHEMA = "nnf.resynthesis.none_growth_plan.v1" |
| GROWTH_PLAN_SCHEMAS = frozenset( |
| {GROWTH_PLAN_SCHEMA, LEGACY_GROWTH_PLAN_SCHEMA} |
| ) |
| STRUCTURAL_EXPERTS = 4 |
|
|
| _LAYER_PATTERN = re.compile(r"^science_stack\.science_layer_(\d+)\.(.+)$") |
|
|
| _GLOBAL_EXPERT_ROWS = frozenset( |
| { |
| "correction_expert_head.weight", |
| "fabric.decision_proj.weight", |
| "fabric.expert_bias_table", |
| "fabric.expert_hidden_table", |
| "fabric.intent_table", |
| } |
| ) |
| _GLOBAL_LAYER_ROWS = frozenset( |
| { |
| "correction_layer_head.weight", |
| "fabric.layer_bias_table", |
| "fabric.layer_hidden_table", |
| "science_stack.layer_identity_glyphs", |
| "science_stack.traversal_gate", |
| } |
| ) |
| _GLOBAL_LAYER_SQUARE = frozenset({"science_stack.layer_transfer_graph"}) |
| _LAYER_EXPERT_ROWS = frozenset( |
| { |
| "_expert_history_states", |
| "expert_activation_prior", |
| "expert_depth_pref", |
| "expert_intent_glyphs", |
| "expert_role_tag", |
| "expert_specialization", |
| "expert_transfer_affinity", |
| "router.weight", |
| } |
| ) |
| _LAYER_EXPERT_SQUARE = frozenset({"expert_compatibility"}) |
| _LAYER_FFN_ROWS = frozenset({"ffn_down", "ffn_gate_up"}) |
|
|
| _STACK_PARAMETER_ORDER = ( |
| "science_stack.traversal_gate", |
| "science_stack.layer_rotation_pressure", |
| "science_stack.layer_transfer_graph", |
| "science_stack.layer_transfer_scale", |
| "science_stack.logit_residual_scale", |
| "science_stack.layer_identity_glyphs", |
| "science_stack.layer_identity_scale", |
| "science_stack.long_context_anchor_gain", |
| "science_stack.glyph_projection.weight", |
| "science_stack.layer_identity_query_proj.weight", |
| "science_stack.long_context_anchor_query.weight", |
| ) |
| _LAYER_PARAMETER_ORDER = ( |
| "expert_activation_prior", |
| "expert_intent_glyphs", |
| "language_match_scale", |
| "expert_role_tag", |
| "expert_specialization", |
| "role_match_scale", |
| "expert_compatibility", |
| "expert_transfer_scale", |
| "expert_rotation_pressure", |
| "expert_depth_pref", |
| "expert_transfer_affinity", |
| "layer_depth_signal", |
| "layer_complexity", |
| "memory_bank", |
| "mhc_distinct_hypothesis_scale", |
| "ffn_gate_up", |
| "ffn_down", |
| "residual_scale", |
| "translate_scale", |
| "audit_scale", |
| "norm.weight", |
| "norm.bias", |
| "output_norm.weight", |
| "output_norm.bias", |
| "router.weight", |
| "intent_query_proj.weight", |
| "role_query_proj.weight", |
| "expert_capability_proj.weight", |
| "expert_capability_proj.bias", |
| "capability_match_scale", |
| "expert_history_gru.weight_ih", |
| "expert_history_gru.weight_hh", |
| "expert_history_gru.bias_ih", |
| "expert_history_gru.bias_hh", |
| "layer_role_head.weight", |
| "layer_role_head.bias", |
| "recurrent_expert.weight_ih_l0", |
| "recurrent_expert.weight_hh_l0", |
| "recurrent_expert.bias_ih_l0", |
| "recurrent_expert.bias_hh_l0", |
| "attention_expert.intent_pivot_scale", |
| "attention_expert.action_pivot_scale", |
| "attention_expert.context_query_pivot_scale", |
| "attention_expert.relation_connectivity_scale", |
| "attention_expert.q_proj.weight", |
| "attention_expert.q_proj.bias", |
| "attention_expert.k_proj.weight", |
| "attention_expert.k_proj.bias", |
| "attention_expert.v_proj.weight", |
| "attention_expert.v_proj.bias", |
| "attention_expert.c_proj.weight", |
| "attention_expert.c_proj.bias", |
| "attention_expert.intent_c_proj.weight", |
| "attention_expert.action_c_proj.weight", |
| "attention_expert.r_query_proj.weight", |
| "attention_expert.r_query_proj.bias", |
| "attention_expert.r_key_proj.weight", |
| "attention_expert.r_key_proj.bias", |
| "attention_expert.intent_r_query_proj.weight", |
| "attention_expert.intent_r_key_proj.weight", |
| "attention_expert.out_proj.weight", |
| "attention_expert.out_proj.bias", |
| "action_glyph_bridge.weight", |
| "memory_query.weight", |
| "memory_out.weight", |
| "glyph_proj.weight", |
| "glyph_gate.weight", |
| "glyph_translate_proj.weight", |
| "glyph_translate_back.weight", |
| ) |
| _V12_LAYER_PARAMETER_ORDER = tuple( |
| name |
| for name in _LAYER_PARAMETER_ORDER |
| if name |
| not in { |
| "attention_expert.action_pivot_scale", |
| "attention_expert.action_c_proj.weight", |
| "action_glyph_bridge.weight", |
| } |
| ) |
| _LEGACY_LAYER_PARAMETER_BASE = tuple( |
| name |
| for name in _LAYER_PARAMETER_ORDER |
| if name != "mhc_distinct_hypothesis_scale" |
| ) |
| _LEGACY_LAYER_PARAMETER_ORDER = ( |
| *_LEGACY_LAYER_PARAMETER_BASE[ |
| : _LEGACY_LAYER_PARAMETER_BASE.index( |
| "attention_expert.intent_pivot_scale" |
| ) |
| ], |
| "attention_expert.in_proj_weight", |
| "attention_expert.in_proj_bias", |
| "attention_expert.out_proj.weight", |
| "attention_expert.out_proj.bias", |
| *_LEGACY_LAYER_PARAMETER_BASE[ |
| _LEGACY_LAYER_PARAMETER_BASE.index("memory_query.weight") : |
| ], |
| ) |
| _TAIL_PARAMETER_ORDER = ( |
| "stop_gate.trajectory_proj.weight", |
| "stop_gate.lstm.weight_ih_l0", |
| "stop_gate.lstm.weight_hh_l0", |
| "stop_gate.lstm.bias_ih_l0", |
| "stop_gate.lstm.bias_hh_l0", |
| "stop_gate.stop_utility_gate.weight", |
| "stop_gate.stop_utility_gate.bias", |
| "stop_gate.stop_contradiction_gate.weight", |
| "stop_gate.stop_contradiction_gate.bias", |
| "feedback_head.weight", |
| "feedback_head.bias", |
| "prior_hidden_proj.weight", |
| "outcome_encoder.weight", |
| "parent_outcome_encoder.weight", |
| "acquisition_encoder.weight", |
| "acquisition_policy.input_norm.weight", |
| "acquisition_policy.input_norm.bias", |
| "acquisition_policy.context.weight", |
| "acquisition_policy.context.bias", |
| "acquisition_policy.action_head.weight", |
| "acquisition_policy.action_head.bias", |
| "correction_context_norm.weight", |
| "correction_context_norm.bias", |
| "correction_hidden_up.weight", |
| "correction_trigger_head.weight", |
| "correction_trigger_head.bias", |
| "task_confidence_head.weight", |
| "task_confidence_head.bias", |
| "delegation_head.weight", |
| "delegation_head.bias", |
| "correction_expert_head.weight", |
| "correction_layer_head.weight", |
| "logit_residual_down.weight", |
| "logit_residual_up.weight", |
| "fabric.expert_bias_table", |
| "fabric.layer_bias_table", |
| "fabric.intent_table", |
| "fabric.domain_residency", |
| "fabric.transfer_table", |
| "fabric.expert_hidden_table", |
| "fabric.layer_hidden_table", |
| "fabric.phase_hidden_table", |
| "fabric.residual_gate", |
| "fabric.parent_expert_route_gate", |
| "fabric.parent_layer_route_gate", |
| "fabric.phase_proj.weight", |
| "fabric.decision_proj.weight", |
| "legacy_capability_bank.fusion_logit", |
| "legacy_capability_bank.route_query.weight", |
| "legacy_capability_bank.outcome_query.weight", |
| "legacy_capability_bank.residual_projection.weight", |
| ) |
|
|
|
|
| def _file_sha256(path: Path) -> str: |
| digest = hashlib.sha256() |
| with path.open("rb") as handle: |
| for chunk in iter(lambda: handle.read(8 * 1024 * 1024), b""): |
| digest.update(chunk) |
| return digest.hexdigest() |
|
|
|
|
| def _state_identity(state: dict[str, torch.Tensor]) -> tuple[str, str]: |
| names = sorted(state) |
| key_hash = hashlib.sha256("\n".join(names).encode("utf-8")).hexdigest() |
| geometry = [(name, tuple(state[name].shape), str(state[name].dtype)) for name in names] |
| geometry_hash = hashlib.sha256( |
| json.dumps(geometry, separators=(",", ":")).encode("utf-8") |
| ).hexdigest() |
| return key_hash, geometry_hash |
|
|
|
|
| def _atomic_torch_save(payload: object, path: Path) -> None: |
| path.parent.mkdir(parents=True, exist_ok=True) |
| temporary = path.with_name(f".{path.name}.{os.getpid()}.tmp") |
| temporary.unlink(missing_ok=True) |
| torch.save(payload, temporary) |
| with temporary.open("rb") as handle: |
| os.fsync(handle.fileno()) |
| os.replace(temporary, path) |
|
|
|
|
| def _atomic_json(payload: dict[str, Any], path: Path) -> None: |
| path.parent.mkdir(parents=True, exist_ok=True) |
| temporary = path.with_name(f".{path.name}.{os.getpid()}.tmp") |
| temporary.unlink(missing_ok=True) |
| with temporary.open("w", encoding="utf-8") as handle: |
| json.dump(payload, handle, sort_keys=True, indent=2) |
| handle.write("\n") |
| handle.flush() |
| os.fsync(handle.fileno()) |
| os.replace(temporary, path) |
|
|
|
|
| def _checkpoint_payload(path: Path, *, map_location: str) -> dict[str, Any]: |
| payload = torch.load(path, map_location=map_location, mmap=True, weights_only=True) |
| if not isinstance(payload, dict) or payload.get("schema") != CHECKPOINT_SCHEMA: |
| raise RuntimeError("Resynthesis geometry source is not an additive checkpoint") |
| lineage = payload.get("lineage") |
| parameters = payload.get("parameters") |
| buffers = payload.get("buffers") |
| if ( |
| not isinstance(lineage, dict) |
| or not isinstance(parameters, dict) |
| or not isinstance(buffers, dict) |
| or not all(isinstance(value, torch.Tensor) for value in parameters.values()) |
| or not all(isinstance(value, torch.Tensor) for value in buffers.values()) |
| ): |
| raise RuntimeError("Resynthesis geometry source checkpoint is incomplete") |
| state = {**parameters, **buffers} |
| key_hash, geometry_hash = _state_identity(state) |
| if payload.get("stateKeySetSha256") != key_hash: |
| raise RuntimeError("Resynthesis geometry source key identity differs") |
| if payload.get("stateGeometrySha256") != geometry_hash: |
| raise RuntimeError("Resynthesis geometry source tensor geometry differs") |
| return payload |
|
|
|
|
| def checkpoint_geometry(path: str | Path) -> tuple[int, int, bool]: |
| """Read additive geometry without allocating checkpoint tensor storage.""" |
|
|
| payload = _checkpoint_payload(Path(path), map_location="meta") |
| lineage = payload["lineage"] |
| layers = int(lineage.get("scienceLayers", 0)) |
| experts = int(lineage.get("scienceExperts", 0)) |
| migrated = bool(lineage.get("draftingCheckpointGeometryChanged", False)) |
| if layers < 1 or experts < STRUCTURAL_EXPERTS + 1: |
| raise RuntimeError("Resynthesis additive checkpoint geometry is invalid") |
| if int(lineage.get("parallelDraftWorkers", 0)) != experts: |
| raise RuntimeError("Resynthesis additive drafting geometry differs") |
| return layers, experts, migrated |
|
|
|
|
| def validate_migration_candidate( |
| checkpoint_path: str | Path, |
| receipt_path: str | Path, |
| ) -> dict[str, Any]: |
| """Validate model, optimizer, lineage, and receipt at training boundary.""" |
|
|
| checkpoint = Path(checkpoint_path).resolve() |
| receipt_file = Path(receipt_path).resolve() |
| receipt = json.loads(receipt_file.read_text(encoding="utf-8")) |
| if not isinstance(receipt, dict) or receipt.get("schema") != MIGRATION_SCHEMA: |
| raise RuntimeError("NoNE geometry candidate receipt schema differs") |
| if receipt.get("passed") is not True or receipt.get("promotionEligible") is not False: |
| raise RuntimeError("NoNE geometry candidate authority is invalid") |
| target = receipt.get("targetCheckpoint") |
| if not isinstance(target, dict): |
| raise RuntimeError("NoNE geometry candidate receipt has no target checkpoint") |
| if Path(str(target.get("path", ""))).resolve() != checkpoint: |
| raise RuntimeError("NoNE geometry candidate path differs from its receipt") |
| if target.get("sha256") != _file_sha256(checkpoint): |
| raise RuntimeError("NoNE geometry candidate SHA-256 differs from its receipt") |
| payload = _checkpoint_payload(checkpoint, map_location="meta") |
| lineage = payload["lineage"] |
| if ( |
| lineage.get("schema") != "nnf.resynthesis.composed_additive_lineage.v13" |
| or lineage.get("intentContextPivotAttention") is not True |
| or lineage.get("contextIntentActionAttention") is not True |
| or lineage.get("contextActionSource") |
| != "trained_acquisition_policy_probability_tensor" |
| or lineage.get("contextActionDim") != 4 |
| or lineage.get("contextActionScorePivot") is not True |
| or lineage.get("contextActionCheckpointGeometryChanged") is not True |
| or lineage.get("intentRelationalAttention") is not True |
| or tuple(lineage.get("attentionMultiples", ())) |
| != ("q", "k", "v", "c", "r") |
| or lineage.get("scienceAttentionExactTiling") is not True |
| or lineage.get("contextRelationCheckpointGeometryChanged") is not True |
| ): |
| raise RuntimeError( |
| "NoNE geometry candidate intent/action C/R lineage differs" |
| ) |
| parameters = payload["parameters"] |
| parameter_elements = sum(tensor.numel() for tensor in parameters.values()) |
| if int(target.get("parameterElements", -1)) != parameter_elements: |
| raise RuntimeError("NoNE geometry candidate parameter count differs") |
| optimizer_record = receipt.get("targetOptimizer") |
| if not isinstance(optimizer_record, dict): |
| raise RuntimeError("NoNE geometry candidate has no optimizer authority") |
| optimizer_path = checkpoint.with_suffix(".optimizer.pt") |
| if Path(str(optimizer_record.get("path", ""))).resolve() != optimizer_path: |
| raise RuntimeError("NoNE geometry optimizer path differs from its receipt") |
| if not optimizer_path.is_file() or optimizer_record.get("sha256") != _file_sha256( |
| optimizer_path |
| ): |
| raise RuntimeError("NoNE geometry optimizer SHA-256 differs from its receipt") |
| optimizer = torch.load( |
| optimizer_path, |
| map_location="meta", |
| mmap=True, |
| weights_only=True, |
| ) |
| if not isinstance(optimizer, dict): |
| raise RuntimeError("NoNE geometry optimizer payload is invalid") |
| groups = optimizer.get("param_groups") |
| states = optimizer.get("state") |
| if not isinstance(groups, list) or len(groups) != 1 or not isinstance(states, dict): |
| raise RuntimeError("NoNE geometry optimizer groups are invalid") |
| names = groups[0].get("param_names") |
| parameter_ids = groups[0].get("params") |
| if ( |
| not isinstance(names, (list, tuple)) |
| or not isinstance(parameter_ids, (list, tuple)) |
| or len(names) != len(parameter_ids) |
| or set(str(name) for name in names) != set(parameters) |
| ): |
| raise RuntimeError("NoNE geometry optimizer names differ from the model") |
| for name, parameter_id in zip(names, parameter_ids, strict=True): |
| parameter_state = states.get(parameter_id) |
| if not isinstance(parameter_state, dict): |
| continue |
| for moment_name in ("exp_avg", "exp_avg_sq", "max_exp_avg_sq"): |
| moment = parameter_state.get(moment_name) |
| if isinstance(moment, torch.Tensor) and moment.shape != parameters[ |
| str(name) |
| ].shape: |
| raise RuntimeError( |
| f"NoNE geometry optimizer moment differs for {name}" |
| ) |
| remaining = receipt.get("remainingProof") |
| if not isinstance(remaining, list) or "continued_training" not in remaining: |
| raise RuntimeError("NoNE geometry candidate omits continued-training authority") |
| checkpoint_geometry(checkpoint) |
| return receipt |
|
|
|
|
| def _expanded_rows(source: torch.Tensor, target_rows: int) -> torch.Tensor: |
| source_rows = source.shape[0] |
| if target_rows < source_rows or source_rows < 1: |
| raise RuntimeError("geometry migration cannot shrink or clone an empty tensor") |
| if target_rows == source_rows: |
| return source.clone() |
| indices = torch.arange( |
| target_rows - source_rows, |
| device=source.device, |
| dtype=torch.long, |
| ).remainder(source_rows) |
| appended = source.index_select(0, indices).clone() |
| if appended.is_floating_point(): |
| offsets = torch.arange( |
| 1, |
| appended.shape[0] + 1, |
| device=appended.device, |
| dtype=torch.float32, |
| ) |
| scale_shape = (appended.shape[0],) + (1,) * (appended.ndim - 1) |
| scale = (1.0 + offsets.remainder(7).reshape(scale_shape) / 128.0).to( |
| dtype=appended.dtype |
| ) |
| appended.mul_(scale) |
| return torch.cat((source.clone(), appended), dim=0) |
|
|
|
|
| def _expanded_square(source: torch.Tensor, target_width: int) -> torch.Tensor: |
| if source.ndim != 2 or source.shape[0] != source.shape[1]: |
| raise RuntimeError("geometry transfer tensor is not square") |
| source_width = source.shape[0] |
| if target_width < source_width or source_width < 1: |
| raise RuntimeError("geometry migration cannot shrink a transfer tensor") |
| if target_width == source_width: |
| return source.clone() |
| indices = torch.arange( |
| target_width, |
| device=source.device, |
| dtype=torch.long, |
| ).remainder(source_width) |
| expanded = source.index_select(0, indices).index_select(1, indices).clone() |
| expanded[:source_width, :source_width].copy_(source) |
| if expanded.is_floating_point(): |
| diagonal = torch.arange(source_width, target_width, device=source.device) |
| expanded[diagonal, diagonal] += expanded.new_tensor(1.0 / 128.0) |
| return expanded |
|
|
|
|
| def _migrate_tensor( |
| name: str, |
| source: torch.Tensor, |
| *, |
| source_layers: int, |
| source_experts: int, |
| target_layers: int, |
| target_experts: int, |
| ) -> torch.Tensor: |
| del source_layers |
| if name in _GLOBAL_EXPERT_ROWS: |
| return _expanded_rows(source, target_experts) |
| if name in _GLOBAL_LAYER_ROWS: |
| return _expanded_rows(source, target_layers) |
| if name in _GLOBAL_LAYER_SQUARE: |
| return _expanded_square(source, target_layers) |
| match = _LAYER_PATTERN.match(name) |
| if match is None: |
| return source.clone() |
| suffix = match.group(2) |
| if suffix in _LAYER_EXPERT_ROWS: |
| return _expanded_rows(source, target_experts) |
| if suffix in _LAYER_EXPERT_SQUARE: |
| return _expanded_square(source, target_experts) |
| if suffix in _LAYER_FFN_ROWS: |
| return _expanded_rows(source, target_experts - STRUCTURAL_EXPERTS) |
| return source.clone() |
|
|
|
|
| def _migrate_state( |
| source: dict[str, torch.Tensor], |
| *, |
| source_layers: int, |
| source_experts: int, |
| target_layers: int, |
| target_experts: int, |
| ) -> dict[str, torch.Tensor]: |
| target = { |
| name: _migrate_tensor( |
| name, |
| tensor, |
| source_layers=source_layers, |
| source_experts=source_experts, |
| target_layers=target_layers, |
| target_experts=target_experts, |
| ) |
| for name, tensor in source.items() |
| } |
| layer_sources: dict[int, list[tuple[str, str, torch.Tensor]]] = {} |
| for name, tensor in source.items(): |
| match = _LAYER_PATTERN.match(name) |
| if match is None: |
| continue |
| layer_sources.setdefault(int(match.group(1)), []).append( |
| (name, match.group(2), tensor) |
| ) |
| if set(layer_sources) != set(range(source_layers)): |
| raise RuntimeError("source checkpoint science layers are not contiguous") |
| for target_layer in range(source_layers, target_layers): |
| source_layer = target_layer % source_layers |
| for _name, suffix, tensor in layer_sources[source_layer]: |
| target_name = f"science_stack.science_layer_{target_layer}.{suffix}" |
| migrated = _migrate_tensor( |
| target_name, |
| tensor, |
| source_layers=source_layers, |
| source_experts=source_experts, |
| target_layers=target_layers, |
| target_experts=target_experts, |
| ) |
| if migrated.is_floating_point() and migrated.ndim > 0: |
| migrated.mul_(migrated.new_tensor(1.0 + (target_layer + 1) / 512.0)) |
| target[target_name] = migrated |
| return dict(sorted(target.items())) |
|
|
|
|
| def _preserved_prefix(source: torch.Tensor, target: torch.Tensor) -> bool: |
| if source.ndim != target.ndim or any( |
| source.shape[axis] > target.shape[axis] for axis in range(source.ndim) |
| ): |
| return False |
| slices = tuple(slice(0, width) for width in source.shape) |
| return torch.equal(source, target[slices]) |
|
|
|
|
| def _legacy_qkv_preserved( |
| source: dict[str, torch.Tensor], |
| adapted: dict[str, torch.Tensor], |
| ) -> bool: |
| """Verify split Q/K/V tensors recompose every trained legacy projection.""" |
|
|
| found = False |
| for name, tensor in source.items(): |
| if not name.endswith("attention_expert.in_proj_weight"): |
| continue |
| found = True |
| prefix = name[: -len("in_proj_weight")] |
| qkv = tuple( |
| adapted.get(f"{prefix}{projection}_proj.weight") |
| for projection in ("q", "k", "v") |
| ) |
| if not all(isinstance(value, torch.Tensor) for value in qkv): |
| return False |
| if not torch.equal( |
| torch.cat(cast(tuple[torch.Tensor, ...], qkv), dim=0), |
| tensor, |
| ): |
| return False |
| bias_name = f"{prefix}in_proj_bias" |
| source_bias = source.get(bias_name) |
| if isinstance(source_bias, torch.Tensor): |
| qkv_bias = tuple( |
| adapted.get(f"{prefix}{projection}_proj.bias") |
| for projection in ("q", "k", "v") |
| ) |
| if not all(isinstance(value, torch.Tensor) for value in qkv_bias): |
| return False |
| if not torch.equal( |
| torch.cat(cast(tuple[torch.Tensor, ...], qkv_bias), dim=0), |
| source_bias, |
| ): |
| return False |
| return found |
|
|
|
|
| def _parameter_order(parameter_names: set[str], layers: int) -> list[str]: |
| order = list(_STACK_PARAMETER_ORDER) |
| layer_order: tuple[str, ...] |
| if any( |
| name.endswith("attention_expert.in_proj_weight") |
| for name in parameter_names |
| ): |
| layer_order = _LEGACY_LAYER_PARAMETER_ORDER |
| elif any( |
| name.endswith("attention_expert.action_c_proj.weight") |
| for name in parameter_names |
| ): |
| layer_order = _LAYER_PARAMETER_ORDER |
| else: |
| layer_order = _V12_LAYER_PARAMETER_ORDER |
| if not any( |
| name.endswith(".mhc_distinct_hypothesis_scale") |
| for name in parameter_names |
| ): |
| layer_order = tuple( |
| name |
| for name in layer_order |
| if name != "mhc_distinct_hypothesis_scale" |
| ) |
| if not any( |
| name.endswith(".capability_match_scale") |
| for name in parameter_names |
| ): |
| layer_order = tuple( |
| name |
| for name in layer_order |
| if name != "capability_match_scale" |
| ) |
| for layer in range(layers): |
| prefix = f"science_stack.science_layer_{layer}." |
| order.extend(prefix + suffix for suffix in layer_order) |
| order.extend(_TAIL_PARAMETER_ORDER) |
| if set(order) != parameter_names: |
| missing = sorted(parameter_names - set(order)) |
| unexpected = sorted(set(order) - parameter_names) |
| raise RuntimeError( |
| "optimizer name recovery differs from checkpoint parameters: " |
| f"unmapped={missing} absent={unexpected}" |
| ) |
| return order |
|
|
|
|
| def _expanded_named_order( |
| source_names: list[str], |
| *, |
| source_layers: int, |
| target_layers: int, |
| ) -> list[str]: |
| target = list(source_names) |
| layer_positions = [ |
| index |
| for index, name in enumerate(source_names) |
| if (match := _LAYER_PATTERN.match(name)) is not None |
| and int(match.group(1)) == source_layers - 1 |
| ] |
| if not layer_positions: |
| raise RuntimeError("optimizer parameter names contain no final science layer") |
| insert_at = max(layer_positions) + 1 |
| additions: list[str] = [] |
| suffixes = [ |
| match.group(2) |
| for name in source_names |
| if (match := _LAYER_PATTERN.match(name)) is not None |
| and int(match.group(1)) == 0 |
| ] |
| for layer in range(source_layers, target_layers): |
| additions.extend(f"science_stack.science_layer_{layer}.{suffix}" for suffix in suffixes) |
| target[insert_at:insert_at] = additions |
| return target |
|
|
|
|
| def _migrate_optimizer( |
| source: dict[str, Any], |
| *, |
| source_parameters: dict[str, torch.Tensor], |
| target_parameters: dict[str, torch.Tensor], |
| source_layers: int, |
| source_experts: int, |
| target_layers: int, |
| target_experts: int, |
| ) -> tuple[dict[str, Any], int, int]: |
| groups = source.get("param_groups") |
| states = source.get("state") |
| if not isinstance(groups, list) or len(groups) != 1 or not isinstance(states, dict): |
| raise RuntimeError("geometry migration requires one AdamW parameter group") |
| source_group = groups[0] |
| source_ids = list(source_group.get("params", ())) |
| names_value = source_group.get("param_names") |
| if isinstance(names_value, (list, tuple)): |
| source_names = [str(name) for name in names_value] |
| if set(source_names) != set(source_parameters): |
| raise RuntimeError("optimizer parameter names differ from checkpoint") |
| else: |
| |
| |
| |
| source_names = _parameter_order(set(source_parameters), source_layers) |
| if len(source_ids) != len(source_names): |
| raise RuntimeError("optimizer parameter IDs differ from named parameters") |
| source_by_name = dict(zip(source_names, source_ids, strict=True)) |
| target_names = list(target_parameters) |
| target_state: dict[int, dict[str, Any]] = {} |
| copied_states = 0 |
| expanded_states = 0 |
| projection_aliases = { |
| ".q_proj.weight": (".in_proj_weight", 0), |
| ".k_proj.weight": (".in_proj_weight", 1), |
| ".v_proj.weight": (".in_proj_weight", 2), |
| ".q_proj.bias": (".in_proj_bias", 0), |
| ".k_proj.bias": (".in_proj_bias", 1), |
| ".v_proj.bias": (".in_proj_bias", 2), |
| } |
| for target_id, name in enumerate(target_names): |
| layer_match = _LAYER_PATTERN.match(name) |
| if ( |
| layer_match is not None |
| and int(layer_match.group(1)) >= source_layers |
| ): |
| continue |
| source_name = name |
| source_id = source_by_name.get(source_name) |
| projection_slice: int | None = None |
| if source_id is None: |
| for suffix, (legacy_suffix, part) in projection_aliases.items(): |
| if name.endswith(suffix): |
| source_name = name[: -len(suffix)] + legacy_suffix |
| source_id = source_by_name.get(source_name) |
| projection_slice = part |
| break |
| if source_id not in states: |
| continue |
| source_state = states[source_id] |
| if not isinstance(source_state, dict): |
| raise RuntimeError("optimizer parameter state is invalid") |
| migrated_state: dict[str, Any] = {} |
| for key, value in source_state.items(): |
| if not isinstance(value, torch.Tensor) or value.ndim == 0: |
| migrated_state[key] = value.clone() if isinstance(value, torch.Tensor) else copy.deepcopy(value) |
| continue |
| source_parameter = source_parameters[source_name] |
| if tuple(value.shape) != tuple(source_parameter.shape): |
| raise RuntimeError( |
| f"optimizer moment geometry differs for {source_name}" |
| ) |
| if projection_slice is not None: |
| width = target_parameters[name].shape[0] |
| if value.shape[0] != 3 * width: |
| raise RuntimeError( |
| f"legacy attention optimizer moment differs for {name}" |
| ) |
| value = value[ |
| projection_slice * width : (projection_slice + 1) * width |
| ].clone() |
| migrated_value = _migrate_tensor( |
| name, |
| value, |
| source_layers=source_layers, |
| source_experts=source_experts, |
| target_layers=target_layers, |
| target_experts=target_experts, |
| ) |
| if migrated_value.shape != target_parameters[name].shape: |
| raise RuntimeError(f"migrated optimizer moment geometry differs for {name}") |
| if migrated_value.shape != value.shape: |
| slices = tuple(slice(0, width) for width in value.shape) |
| migrated_value.zero_() |
| migrated_value[slices].copy_(value) |
| expanded_states += 1 |
| migrated_state[key] = migrated_value |
| target_state[target_id] = migrated_state |
| copied_states += 1 |
| target_group = { |
| key: copy.deepcopy(value) |
| for key, value in source_group.items() |
| if key not in {"params", "param_names"} |
| } |
| target_group["params"] = list(range(len(target_names))) |
| target_group["param_names"] = target_names |
| return {"state": target_state, "param_groups": [target_group]}, copied_states, expanded_states |
|
|
|
|
| def migrate_geometry_checkpoint( |
| source_checkpoint: str | Path, |
| growth_plan: str | Path, |
| output_checkpoint: str | Path, |
| receipt_path: str | Path, |
| *, |
| source_optimizer: str | Path | None = None, |
| output_optimizer: str | Path | None = None, |
| ) -> dict[str, Any]: |
| """Create a trained-state-preserving, non-promoted growth candidate.""" |
|
|
| source_path = Path(source_checkpoint).resolve() |
| plan_path = Path(growth_plan).resolve() |
| output_path = Path(output_checkpoint).resolve() |
| receipt = Path(receipt_path).resolve() |
| if source_path == output_path: |
| raise ValueError("geometry migration output must differ from its source") |
| plan = json.loads(plan_path.read_text(encoding="utf-8")) |
| if ( |
| not isinstance(plan, dict) |
| or plan.get("schema") not in GROWTH_PLAN_SCHEMAS |
| ): |
| raise RuntimeError("NoNE growth plan schema differs") |
| if isinstance(plan.get("pagingPolicy"), dict): |
| raise RuntimeError( |
| "paged NoNE growth plans require immutable page-generation " |
| "migration, not monolithic geometry cloning" |
| ) |
| target_geometry = plan.get("proposedMinimumTargetGeometry") |
| if not isinstance(target_geometry, dict): |
| raise RuntimeError("NoNE growth plan has no target geometry") |
| payload = _checkpoint_payload(source_path, map_location="cpu") |
| lineage = payload["lineage"] |
| original_source_parameters = payload["parameters"] |
| source_layers = int(lineage.get("scienceLayers", 0)) |
| source_experts = int(lineage.get("scienceExperts", 0)) |
| target_layers = int(target_geometry.get("scienceLayers", 0)) |
| target_experts = int(target_geometry.get("scienceExperts", 0)) |
| if target_layers <= source_layers or target_experts <= source_experts: |
| raise RuntimeError("NoNE migration target must expand both layers and experts") |
| from resynthesis.science_layers import ( |
| SCIENCE_ATTENTION_TILE_TOKENS, |
| adapt_attention_state_to_context_relation, |
| ) |
| from resynthesis.config import ( |
| DUAL_CHUNK_LOCAL_SIZE, |
| DUAL_CHUNK_PRETRAIN_LENGTH, |
| NATIVE_ATTENTION_POSITION_APERTURE, |
| PRETRAINED_ROPE_BAND_TOKENS, |
| RESYNTHESIS_NATIVE_PREFILL_TILE_TOKENS, |
| RESYNTHESIS_ONLINE_SOFTMAX_TILE_TOKENS, |
| ) |
|
|
| source_parameters, attention_adapted = adapt_attention_state_to_context_relation( |
| original_source_parameters |
| ) |
| legacy_projection_present = any( |
| name.endswith("attention_expert.in_proj_weight") |
| for name in original_source_parameters |
| ) |
| legacy_qkv_preserved = ( |
| not legacy_projection_present |
| or _legacy_qkv_preserved( |
| original_source_parameters, |
| source_parameters, |
| ) |
| ) |
| if not legacy_qkv_preserved: |
| raise RuntimeError("NoNE attention migration changed trained Q/K/V regions") |
| relational_layers = { |
| int(match.group(1)) |
| for name in source_parameters |
| if name.endswith("attention_expert.r_query_proj.weight") |
| and (match := _LAYER_PATTERN.match(name)) is not None |
| } |
| if relational_layers != set(range(source_layers)): |
| raise RuntimeError( |
| "NoNE geometry source lacks complete context-relational attention" |
| ) |
| source_buffers = payload["buffers"] |
| if isinstance(lineage, dict): |
| lineage = dict(lineage) |
| lineage["intentContextPivotAttention"] = True |
| lineage["intentRelationalAttention"] = True |
| lineage["attentionMultiples"] = ("q", "k", "v", "c", "r") |
| lineage["contextQueryPivot"] = True |
| lineage["contextIntentActionAttention"] = True |
| lineage["contextActionSource"] = ( |
| "trained_acquisition_policy_probability_tensor" |
| ) |
| lineage["contextActionDim"] = 4 |
| lineage["contextActionScorePivot"] = True |
| lineage["contextActionCheckpointGeometryChanged"] = True |
| lineage["relationConnectivity"] = "none_router_selected_intent_tensor" |
| lineage["scienceAttentionExactTiling"] = True |
| lineage["scienceAttentionTileTokens"] = SCIENCE_ATTENTION_TILE_TOKENS |
| lineage["contextRelationCheckpointGeometryChanged"] = True |
| lineage["parentDualChunkRoPECompose"] = True |
| lineage["additiveOnlineSoftmaxLongPool"] = True |
| lineage["dualChunkAtSuccessiveSeamExposed"] = True |
| lineage["nativeContextPositionAperture"] = ( |
| NATIVE_ATTENTION_POSITION_APERTURE |
| ) |
| lineage["onlineSoftmaxTileTokens"] = ( |
| RESYNTHESIS_ONLINE_SOFTMAX_TILE_TOKENS |
| ) |
| lineage["nativePrefillTileTokens"] = ( |
| RESYNTHESIS_NATIVE_PREFILL_TILE_TOKENS |
| ) |
| lineage["longContextStackComposeCheckpointGeometryChanged"] = False |
| parent_lineage = lineage.get("parent") |
| if not isinstance(parent_lineage, dict): |
| raise RuntimeError("NoNE geometry source parent lineage is invalid") |
| parent_lineage = dict(parent_lineage) |
| parent_lineage.update( |
| { |
| "parentDualChunkRoPECompose": True, |
| "parentOnlineSoftmaxLongPool": True, |
| "nativeAttentionPositionAperture": ( |
| NATIVE_ATTENTION_POSITION_APERTURE |
| ), |
| "onlineSoftmaxTileTokens": ( |
| RESYNTHESIS_ONLINE_SOFTMAX_TILE_TOKENS |
| ), |
| "dualChunkPretrainLength": DUAL_CHUNK_PRETRAIN_LENGTH, |
| "dualChunkLocalSize": DUAL_CHUNK_LOCAL_SIZE, |
| "dualChunkAtSuccessiveSeamExposed": True, |
| "pretrainedRoPEBandTokens": PRETRAINED_ROPE_BAND_TOKENS, |
| "nativePrefillTileTokens": ( |
| RESYNTHESIS_NATIVE_PREFILL_TILE_TOKENS |
| ), |
| } |
| ) |
| lineage["parent"] = parent_lineage |
| lineage["schema"] = "nnf.resynthesis.composed_additive_lineage.v13" |
| target_parameters = _migrate_state( |
| source_parameters, |
| source_layers=source_layers, |
| source_experts=source_experts, |
| target_layers=target_layers, |
| target_experts=target_experts, |
| ) |
| target_buffers = _migrate_state( |
| source_buffers, |
| source_layers=source_layers, |
| source_experts=source_experts, |
| target_layers=target_layers, |
| target_experts=target_experts, |
| ) |
| preserved = sum( |
| _preserved_prefix(value, target_parameters[name]) |
| for name, value in source_parameters.items() |
| ) + sum( |
| _preserved_prefix(value, target_buffers[name]) |
| for name, value in source_buffers.items() |
| ) |
| source_tensor_count = len(source_parameters) + len(source_buffers) |
| if preserved != source_tensor_count: |
| raise RuntimeError("geometry migration changed a trained source tensor region") |
| target_lineage = copy.deepcopy(lineage) |
| target_lineage["scienceLayers"] = target_layers |
| target_lineage["scienceExperts"] = target_experts |
| target_lineage["parallelDraftWorkers"] = target_experts |
| target_lineage["draftingCheckpointGeometryChanged"] = True |
| target_state = {**target_parameters, **target_buffers} |
| key_hash, geometry_hash = _state_identity(target_state) |
| target_payload = { |
| "schema": CHECKPOINT_SCHEMA, |
| "lineage": target_lineage, |
| "stateKeySetSha256": key_hash, |
| "stateGeometrySha256": geometry_hash, |
| "parameters": target_parameters, |
| "buffers": target_buffers, |
| } |
| _atomic_torch_save(target_payload, output_path) |
|
|
| optimizer_source_path = ( |
| Path(source_optimizer).resolve() if source_optimizer is not None else None |
| ) |
| optimizer_output_path = ( |
| Path(output_optimizer).resolve() |
| if output_optimizer is not None |
| else output_path.with_suffix(".optimizer.pt") |
| ) |
| copied_optimizer_states = 0 |
| expanded_optimizer_states = 0 |
| optimizer_written = False |
| if optimizer_source_path is not None: |
| optimizer_payload = torch.load( |
| optimizer_source_path, |
| map_location="cpu", |
| mmap=True, |
| weights_only=True, |
| ) |
| if not isinstance(optimizer_payload, dict): |
| raise RuntimeError("NoNE source optimizer checkpoint is invalid") |
| migrated_optimizer, copied_optimizer_states, expanded_optimizer_states = ( |
| _migrate_optimizer( |
| optimizer_payload, |
| source_parameters=original_source_parameters, |
| target_parameters=target_parameters, |
| source_layers=source_layers, |
| source_experts=source_experts, |
| target_layers=target_layers, |
| target_experts=target_experts, |
| ) |
| ) |
| _atomic_torch_save(migrated_optimizer, optimizer_output_path) |
| optimizer_written = True |
|
|
| parameter_elements = sum(tensor.numel() for tensor in target_parameters.values()) |
| receipt_payload = { |
| "schema": MIGRATION_SCHEMA, |
| "passed": True, |
| "builtAt": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), |
| "sourceCheckpoint": { |
| "path": str(source_path), |
| "sha256": _file_sha256(source_path), |
| }, |
| "sourceOptimizer": ( |
| { |
| "path": str(optimizer_source_path), |
| "sha256": _file_sha256(optimizer_source_path), |
| } |
| if optimizer_source_path is not None |
| else None |
| ), |
| "growthPlan": { |
| "path": str(plan_path), |
| "sha256": _file_sha256(plan_path), |
| }, |
| "targetCheckpoint": { |
| "path": str(output_path), |
| "sha256": _file_sha256(output_path), |
| "parameterElements": parameter_elements, |
| "parameterBillions": parameter_elements / 1_000_000_000, |
| }, |
| "targetOptimizer": ( |
| { |
| "path": str(optimizer_output_path), |
| "sha256": _file_sha256(optimizer_output_path), |
| } |
| if optimizer_written |
| else None |
| ), |
| "sourceGeometry": { |
| "scienceLayers": source_layers, |
| "scienceExperts": source_experts, |
| "parallelDraftWorkers": source_experts, |
| }, |
| "targetGeometry": { |
| "scienceLayers": target_layers, |
| "scienceExperts": target_experts, |
| "parallelDraftWorkers": target_experts, |
| }, |
| "checks": { |
| "allSourceTensorRegionsPreservedExactly": preserved == source_tensor_count, |
| "preservedSourceTensorRegions": preserved, |
| "sourceTensorRegions": source_tensor_count, |
| "contextRelationAttentionMigrated": attention_adapted, |
| "legacyQkvSlicesPreservedExactly": legacy_qkv_preserved, |
| "contextRelationGatesIdentityInitialized": all( |
| bool( |
| tensor.detach().eq(0).all() |
| ) |
| for name, tensor in source_parameters.items() |
| if name.endswith( |
| ( |
| "context_query_pivot_scale", |
| "relation_connectivity_scale", |
| "action_pivot_scale", |
| ) |
| ) |
| ), |
| "contextActionProjectionPresent": all( |
| ( |
| f"science_stack.science_layer_{layer}." |
| "attention_expert.action_c_proj.weight" |
| ) |
| in source_parameters |
| for layer in range(source_layers) |
| ), |
| "contextActionGlyphBridgePresent": all( |
| ( |
| f"science_stack.science_layer_{layer}." |
| "action_glyph_bridge.weight" |
| ) |
| in source_parameters |
| for layer in range(source_layers) |
| ), |
| "newLayerModulesPresent": all( |
| any( |
| name.startswith(f"science_stack.science_layer_{layer}.") |
| for name in target_state |
| ) |
| for layer in range(source_layers, target_layers) |
| ), |
| "expertGeometryExpanded": target_experts > source_experts, |
| "layerGeometryExpanded": target_layers > source_layers, |
| "modelOwnedRoutingRetained": True, |
| "optimizerNamedStateMigrated": optimizer_written, |
| "copiedOptimizerStates": copied_optimizer_states, |
| "expandedOptimizerMomentTensors": expanded_optimizer_states, |
| }, |
| "promotionEligible": False, |
| "remainingProof": [ |
| "continued_training", |
| "cold_reload", |
| "heldout_generalization", |
| "correction_stress", |
| "immutable_release_verification", |
| ], |
| } |
| _atomic_json(receipt_payload, receipt) |
| return receipt_payload |
|
|
|
|
| def main() -> int: |
| import argparse |
|
|
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("source_checkpoint", type=Path) |
| parser.add_argument("growth_plan", type=Path) |
| parser.add_argument("output_checkpoint", type=Path) |
| parser.add_argument("receipt", type=Path) |
| parser.add_argument("--source-optimizer", type=Path) |
| parser.add_argument("--output-optimizer", type=Path) |
| args = parser.parse_args() |
| result = migrate_geometry_checkpoint( |
| args.source_checkpoint, |
| args.growth_plan, |
| args.output_checkpoint, |
| args.receipt, |
| source_optimizer=args.source_optimizer, |
| output_optimizer=args.output_optimizer, |
| ) |
| print(json.dumps(result, sort_keys=True)) |
| return 0 if result.get("passed") is True else 1 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
