"""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: # Snapshot parameter dictionaries are canonically name-sorted, while # pre-v25 optimizers followed module registration order. Reconstruct # that exact checkpointed architecture order; never infer by shape. 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())