""" gemma3-px — Surgical Patch (Phase 10.0) ========================================= Implements Recursive State Memory (RSM) and Hyper-Fluid Routing (HFR). RSM allows the model to 'see' its own previous thinking states during recursion. """ import types import torch import torch.nn as nn import os import json import datetime from typing import Optional from .px_modules import LTIInjection, ADCInjection, StabilityMonitor # --------------------------------------------------------------------------- # p10.0: Recursive State Memory (RSM) # --------------------------------------------------------------------------- class RecursiveMemoryCache: """ Extends ReadOnlyCache by injecting previous thinking steps into the self-attention key/value streams. """ def __init__(self, real_cache, thought_history=None, layer_types=None, read_only=False, expected_len=0): self.__dict__["_real"] = real_cache self.__dict__["_thoughts"] = thought_history or [] self.__dict__["_layer_types"] = layer_types or [] self.__dict__["_read_only"] = read_only self.__dict__["_expected_len"] = expected_len def __getattr__(self, name): return getattr(self._real, name) def update(self, key_states, value_states, layer_idx, cache_kwargs=None): # 1. Base Update (Functional if read_only) if self._read_only: past_k, past_v = None, None # Try older DynamicCache style if hasattr(self._real, "key_cache") and len(self._real.key_cache) > layer_idx: past_k = self._real.key_cache[layer_idx] past_v = self._real.value_cache[layer_idx] # Try newer Cache object style (transformers 4.45+) elif hasattr(self._real, "layers") and len(self._real.layers) > layer_idx: layer = self._real.layers[layer_idx] # DynamicLayer / StaticLayer if hasattr(layer, "keys") and layer.keys is not None: past_k = layer.keys past_v = layer.values # SinkCache / etc might have different names? No, usually .keys if past_k is None: past_k = torch.empty(0, device=key_states.device, dtype=key_states.dtype) past_v = torch.empty(0, device=value_states.device, dtype=value_states.dtype) # If past_k already has the expected length, it means this layer was # already updated for the current token(s) in a previous iteration # of the same reasoning loop. if past_k.numel() > 0 and past_k.shape[-2] == self._expected_len: res_k, res_v = past_k, past_v else: res_k = torch.cat([past_k, key_states], dim=-2) res_v = torch.cat([past_v, value_states], dim=-2) # print(f" [DEBUG-CACHE] L{layer_idx} RO=True | past={past_k.shape[-2] if past_k.numel()>0 else 0} | cur={key_states.shape[-2]} | res={res_k.shape[-2]} | exp={self._expected_len}") else: res_k, res_v = self._real.update(key_states, value_states, layer_idx, cache_kwargs) # print(f" [DEBUG-CACHE] L{layer_idx} RO=False | res={res_k.shape[-2]} | exp={self._expected_len}") # 2. Phase 14.6: Soft-RSM (Semantic Blending) is_full = self._layer_types and self._layer_types[layer_idx] == "full_attention" if self._thoughts and layer_idx >= 6 and is_full: B, H_kv, T_res, HD = res_k.shape T_curr = key_states.shape[-2] alpha = 0.15 # Phase 14.7: Triangular Weighting (Emphasize the 'reasoning peak') n_t = len(self._thoughts[-6:]) if n_t > 2: weights = torch.cat([ torch.linspace(0.4, 1.0, n_t//2, device=res_k.device), torch.linspace(1.0, 0.6, n_t - n_t//2, device=res_k.device) ]) t_raw = (torch.stack(self._thoughts[-6:]) * weights.view(-1, 1, 1, 1)).sum(dim=0) / weights.sum() else: t_raw = torch.stack(self._thoughts).mean(dim=0) D = t_raw.shape[2] # Project thought to Head Dim (SDA) t_flat = t_raw.mean(dim=1, keepdim=True) # (B, 1, D) t_proj = torch.nn.functional.interpolate(t_flat, size=HD, mode='linear', align_corners=False) t_k = t_proj.unsqueeze(1) # (B, 1, 1, HD) t_v = -t_k # Blend into the LAST token(s) of the result # Use in-place only if not read_only to avoid side effects on cache if self._read_only: res_k = res_k.clone() res_v = res_v.clone() res_k[:, :, -T_curr:, :] = (1.0 - alpha) * res_k[:, :, -T_curr:, :] + alpha * t_k res_v[:, :, -T_curr:, :] = (1.0 - alpha) * res_v[:, :, -T_curr:, :] + alpha * t_v return res_k, res_v # --------------------------------------------------------------------------- def remove_px_patch(model) -> None: from transformers.models.gemma3.modeling_gemma3 import Gemma3TextModel text_model = (model.model if hasattr(model, "model") else model) if hasattr(text_model, "_px_config"): text_model.forward = types.MethodType( Gemma3TextModel.forward, text_model ) del text_model._px_injection del text_model._px_config print("[gemma3-px] Patch removed.") def _resolve_text_model(model): if hasattr(model, "model") and hasattr(model.model, "layers"): return model.model return model # --------------------------------------------------------------------------- def _px_forward( self, input_ids=None, attention_mask=None, position_ids=None, past_key_values=None, inputs_embeds=None, use_cache=None, **kwargs, ): from transformers.cache_utils import DynamicCache from transformers.masking_utils import create_causal_mask, create_sliding_window_causal_mask if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError("Specify exactly one of input_ids or inputs_embeds.") if inputs_embeds is None: # Multimodal resolution (Phase 17.7) if hasattr(self, "embed_tokens"): inputs_embeds = self.embed_tokens(input_ids) elif hasattr(self, "language_model"): inputs_embeds = self.language_model.model.embed_tokens(input_ids) elif hasattr(self, "model") and hasattr(self.model, "embed_tokens"): inputs_embeds = self.model.embed_tokens(input_ids) else: # Last resort: search for embed_tokens in children embedder = None for name, module in self.named_modules(): if "embed_tokens" in name: embedder = module break if embedder: inputs_embeds = embedder(input_ids) else: raise AttributeError(f"Could not find embed_tokens in model type {type(self)}. Available: {dir(self)[:20]}...") if use_cache and past_key_values is None: past_key_values = DynamicCache(config=self.config) # Phase 14.8: Initial sequence length tracking past_seen = past_key_values.get_seq_length() if past_key_values is not None else 0 expected_len = past_seen + inputs_embeds.shape[1] # print(f"[DEBUG-PX] Type={type(past_key_values)} seen={past_seen} cur={inputs_embeds.shape[1]} exp={expected_len}") if position_ids is None: position_ids = ( torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen ).unsqueeze(0) # Resolve config for masking (Phase 17.7 multimodal fix) mask_config = self.config if hasattr(mask_config, "text_config"): mask_config = mask_config.text_config if not isinstance(attention_mask, dict): mk = dict( config=mask_config, inputs_embeds=inputs_embeds, attention_mask=attention_mask, past_key_values=past_key_values, position_ids=position_ids, ) causal_mask_mapping = { "full_attention": create_causal_mask(**mk), "sliding_attention": create_sliding_window_causal_mask(**mk), } else: causal_mask_mapping = attention_mask hidden_states = inputs_embeds position_embeddings = {} for layer_type in set(mask_config.layer_types): position_embeddings[layer_type] = self.rotary_emb( hidden_states, position_ids, layer_type ) cfg = self._px_config updated_layers = set() # Phase 14.9: Global visit tracker for this forward pass # ── 1. PRELUDE ────────────────────────────────────────────────────────── for i in range(cfg["prelude_end"]): updated_layers.add(i) layer_out = self.layers[i]( hidden_states, attention_mask=causal_mask_mapping[mask_config.layer_types[i]], position_embeddings=position_embeddings[mask_config.layer_types[i]], position_ids=position_ids, past_key_values=past_key_values, **kwargs, ) hidden_states = layer_out[0] if isinstance(layer_out, (tuple, list)) else layer_out # ── 1.5 META-SELECTOR (Phase 28 Fluid) ─────────────────────────────────── dynamic_start = cfg["recur_start"] dynamic_end = cfg["recur_end"] dynamic_hub = cfg.get("bimodal_hub", cfg["recur_start"]) num_layers = len(self.layers) if cfg.get("routing_mode") == "adaptive": if inputs_embeds.shape[1] > 1: # Prefill phase: Measure Kurtosis and Jitter at Layer 5 h_base_f32 = hidden_states.to(torch.float32) # Kurtosis (Last token) h_probe = h_base_f32[0, -1, :] variance = torch.var(h_probe).item() kurtosis = (torch.mean((h_probe - torch.mean(h_probe))**4) / (variance**2)).item() if variance > 0 else 0 self._task_kurtosis = kurtosis # Jitter (Across sequence) # We measure the variance of the norms of the hidden states across the prompt. # Very high jitter usually indicates a 'Trap' or 'Divergence' in the intuition pass. h_norms = h_base_f32.norm(dim=-1) # [B, T] h_norm_var = torch.var(h_norms, dim=-1).mean().item() self._task_jitter = h_norm_var if os.environ.get("DEBUG_ROUTING") == "1": print(f"[Router] Prefill K={kurtosis:.1f}, Jitter={h_norm_var:.4f}") kurtosis = getattr(self, "_task_kurtosis", 300) # Default to Logic if missing import math if num_layers < 20: # 270M Model (Kurtosis is task-separable) # Continuous Fluid Gaussian Blending of the 5 Zones w_m = math.exp(-((kurtosis - 200)**2) / (2 * 25**2)) # Math w_la = math.exp(-((kurtosis - 275)**2) / (2 * 15**2)) # Logic-A w_cr = math.exp(-((kurtosis - 298)**2) / (2 * 8**2)) # Creative w_lb = math.exp(-((kurtosis - 310)**2) / (2 * 8**2)) # Logic-B w_sy = math.exp(-((kurtosis - 325)**2) / (2 * 20**2)) # Synthesis W = w_m + w_la + w_cr + w_lb + w_sy + 1e-9 # Phase 36.2: Restored Stable Math Zone d_start = (w_m*5 + w_la*8 + w_cr*10 + w_lb*8 + w_sy*6) / W d_end = (w_m*11 + w_la*12 + w_cr*16 + w_lb*14 + w_sy*14) / W d_hub = (w_m*8 + w_la*11 + w_cr*14 + w_lb*11 + w_sy*10) / W d_loops = (w_m*8 + w_la*8 + w_cr*6 + w_lb*10 + w_sy*8) / W dynamic_start = max(1, int(round(d_start))) dynamic_end = min(num_layers - 1, int(round(d_end))) dynamic_hub = int(round(d_hub)) cfg["n_loops"] = max(2, int(round(d_loops))) zone_name = f"Fluid-Blended (K={kurtosis:.1f})" else: # 1B and 4B Models (Scale-Invariant Omni Zone) # They have enough capacity to hold both semantics without smearing dynamic_start = int(num_layers * 0.38) dynamic_end = int(num_layers * 0.76) dynamic_hub = int(num_layers * 0.61) cfg["n_loops"] = 6 zone_name = "Omni-Scale" # Only print routing decision once per token during generation if inputs_embeds.shape[1] == 1 and os.environ.get("DEBUG_ROUTING") == "1": print(f"[Router] {zone_name} -> L{dynamic_start}-L{dynamic_end} (Loops: {cfg['n_loops']}, Hub: {dynamic_hub})") # Fast-forward prelude if needed for i in range(cfg["prelude_end"], dynamic_start): updated_layers.add(i) layer_out = self.layers[i]( hidden_states, attention_mask=causal_mask_mapping[mask_config.layer_types[i]], position_embeddings=position_embeddings[mask_config.layer_types[i]], position_ids=position_ids, past_key_values=past_key_values, **kwargs, ) hidden_states = layer_out[0] if isinstance(layer_out, (tuple, list)) else layer_out # ── 2. REASONING ZONE (Phase 10.0) ────────────────────────────────────── e_static = hidden_states.clone() # 2.A: Intuition Pass trans_out = hidden_states for i_layer in range(dynamic_start, dynamic_end): l_type = mask_config.layer_types[i_layer] updated_layers.add(i_layer) layer_out = self.layers[i_layer]( trans_out, attention_mask=causal_mask_mapping[l_type], position_embeddings=position_embeddings[l_type], position_ids=position_ids, past_key_values=past_key_values, **kwargs, ) trans_out = layer_out[0] if isinstance(layer_out, (tuple, list)) else layer_out # if past_key_values is not None: # print(f" [DEBUG-PX-DIR] {dir(past_key_values)}") h_baseline = trans_out # Phase 14.5: ETR (Entropy Triggered Recursion) # Estimate 'confidence' from the last layer's norm change or simpler: # We only run recursion if the intuition pass wasn't 'perfectly' stable. # Note: h_baseline is already computed. # 2.B: Hyper-Fluid Routing & Recursive Memory n_loops = cfg.get("n_loops", 2) # Phase 14.5: ETR (Entropy Triggered Recursion) phi_intuition = StabilityMonitor.calculate_phi(h_baseline, hidden_states).mean().item() if os.environ.get("DEBUG_ROUTING") == "1": print(f" [Intuition] Phi: {phi_intuition:.6f}") # Phase 14.7: Gamma-Damping instead of loop scaling current_gamma = cfg.get("gamma", 0.08) e_reflector = e_static is_trap_candidate = False # Phase 36.3: Surgical Reflector Activation # Jitter is only for extreme representational collapse (1e8) jitter = getattr(self, "_task_jitter", 0.0) kurtosis = getattr(self, "_task_kurtosis", 300) # Trigger Reflector if extreme jitter OR if it's a known Logic-Trap zone if jitter > 1e8 or (280 < kurtosis < 310): is_trap_candidate = True if os.environ.get("DEBUG_ROUTING") == "1": reason = "Jitter" if jitter > 1e8 else "Logic-Zone" print(f" [Router] Trap detected via {reason} ({jitter:.1f}), activating Reflector") # Phase 16.3: Anchor Reflection e_stat_f32 = e_static.to(torch.float32) h_base_f32 = h_baseline.to(torch.float32) e_ref_f32 = 2.0 * e_stat_f32 - h_base_f32 e_ref_f32 = e_ref_f32 * (e_stat_f32.norm() / (e_ref_f32.norm() + 1e-6)) e_reflector = e_ref_f32.to(e_static.dtype) if phi_intuition > 0.9999 and not is_trap_candidate: current_gamma *= 0.1 elif phi_intuition > 0.999: current_gamma *= 0.5 # Phase 25: Sigmoid-Annealed Orthogonal Recovery (SAOR) # ----------------------------------------------------------------------- # Using a Sigmoid curve for Gamma to allow a sharp "Phase Transition" # from exploration (high energy) to grounding (low energy). # Plus: Orthogonal Reinforcement to protect logical drift. base_gamma = current_gamma bimodal_hub_start = cfg.get("bimodal_hub", 11) path_taken = [] thought_history = [] avg_phi_explore = 1.0 exploration_steps = 0 telemetry_steps = [] # Context dims B, T_curr = hidden_states.shape[0], hidden_states.shape[1] HD = getattr(self.config, "head_dim", 256) # Phase 38.1: Anna Karenina Sensor (AKS) Initialization # Tracks the "Geometric Disparity" of the latent thoughts. # Clustering = Truth (Anna Karenina Principle), Dispersion = Error. divergence_buffer = [] correction_strength = 0.0 if n_loops > 1: h_exp = e_reflector.clone() # Use Reflected Anchor current_layer = dynamic_start max_steps = (dynamic_end - dynamic_start) * n_loops * 3 phis = [] stability_counter = 0 layer_visits = {i: 0 for i in range(5, 18)} # Initialize active bounds active_start = dynamic_start active_end = dynamic_end while current_layer < active_end and exploration_steps < max_steps: # --- PHASE 26: INFINITE REFLECTION (IR) --- # Center of the thinking process (0.5) t_norm = exploration_steps / max_steps # Phase 38.2: AKS - Topological Anomaly Detection # We measure the 'velocity' of the latent state across the last 3 steps. # If the states move AWAY from the Anchor (Disparity), we trigger Correction. if exploration_steps > 2: # Proximity to Anchor (e_static) dist_now = 1.0 - StabilityMonitor.calculate_phi(h_exp, e_static).mean().item() divergence_buffer.append(dist_now) if len(divergence_buffer) > 4: divergence_buffer.pop(0) # Check for "Accelerating Dispersion" (The Anna Karenina Signal) if len(divergence_buffer) >= 3: velocity = divergence_buffer[-1] - divergence_buffer[-2] acceleration = (divergence_buffer[-1] - divergence_buffer[-2]) - (divergence_buffer[-2] - divergence_buffer[-3]) # If we are accelerating AWAY from the truth (dist increasing) if acceleration > 0.001 and velocity > 0: correction_strength = min(1.0, correction_strength + 0.1) if os.environ.get("DEBUG_PHI") == "1": print(f" [AKS] Anomaly! Accel={acceleration:.4f}, Correction={correction_strength:.2f}") else: correction_strength = max(0.0, correction_strength - 0.05) # --- PHASE 28: TEMPORAL COGNITIVE ROUTING (TCR) --- active_start = dynamic_start active_end = dynamic_end if getattr(self, "_task_kurtosis", 300) > 280 and getattr(self, "_task_kurtosis", 300) < 305: if t_norm < 0.33: active_start = 8 active_end = 14 elif t_norm < 0.66: active_start = 5 active_end = 11 else: active_start = 8 active_end = 12 # Sigmoid transition: sharp drop in energy after 50% of thinking k_steep = 12.0 energy_factor = 1.0 - (1.0 / (1.0 + torch.exp(torch.tensor(-k_steep * (t_norm - 0.5))))).item() # Phase 38.3: AKS-Modulated Gamma # High correction strength forces the model to ground HARDER (lower exploration). current_gamma = base_gamma * (0.5 + energy_factor) * (1.0 - 0.5 * correction_strength) # Phase 26: Hub Oscillation. oscillation = 1 if (exploration_steps % 4 < 2) else -1 bimodal_hub = min(active_end - 1, max(active_start, int(dynamic_hub + (t_norm * 2) + oscillation))) h_prev = h_exp.clone() # Safe layer visit tracking if current_layer not in layer_visits: layer_visits[current_layer] = 0 layer_visits[current_layer] += 1 # Phase 14.7: Surgical Cache Security is_first_visit = current_layer not in updated_layers if is_first_visit: updated_layers.add(current_layer) # Phase 38.4: AKS-Informed Sensory Refresh # If we are in high correction mode, increase sensory re-injection. refresh_rate = 0.10 + 0.20 * correction_strength if exploration_steps % 6 == 0 and exploration_steps > 0: h_exp = (1.0 - refresh_rate) * h_exp + refresh_rate * e_static path_taken.append(f"SENSORY_REFRESH(AKS={correction_strength:.1f})") # Phase 10.0: Memory-Augmented Cache wrapper current_past = RecursiveMemoryCache( past_key_values, thought_history, layer_types=mask_config.layer_types, read_only=not is_first_visit, expected_len=expected_len ) if past_key_values is not None else None # Execute layer l_type = mask_config.layer_types[current_layer] layer_out = self.layers[current_layer]( h_exp, attention_mask=causal_mask_mapping[l_type], position_embeddings=position_embeddings[l_type], position_ids=position_ids, past_key_values=current_past, **kwargs, ) trans_out = layer_out[0] if isinstance(layer_out, (tuple, list)) else layer_out # Phase 35: Metacognitive Phi-Jitter & Early Exit (Annealed) phi_step = StabilityMonitor.calculate_phi(trans_out, h_prev).mean().item() if os.environ.get("DEBUG_PHI") == "1": print(f" [L{current_layer}] Phi: {phi_step:.6f}") # --- Early Exit (Annealed) --- # Only exit early in the second half of thinking to ensure grounding. if t_norm > 0.5 and phi_step > 0.9999: stability_counter += 1 if stability_counter > 3: if os.environ.get("DEBUG_ROUTING") == "1": print(f" [Router] Early Exit at step {exploration_steps}") h_exp = trans_out break else: stability_counter = 0 # --- Hub Jitter (Exploratory Phase) --- # Only jitter in the first 40% of thinking to explore alternatives. if t_norm < 0.4 and phi_step > 0.995 and phi_step < 0.999: if exploration_steps % 4 == 0: current_layer = min(active_end - 1, current_layer + 2) if os.environ.get("DEBUG_ROUTING") == "1": print(f" [Router] Jittering to L{current_layer}") # --- PHASE 25.1: RECURSIVE BELIEF ANCHOR (RBA) --- # Update the anchor slightly with recent thoughts to carry over logic if len(thought_history) > 2: # Use a sliding window average of thoughts recent_avg = torch.stack(thought_history[-3:]).mean(dim=0) e_dynamic = 0.85 * e_reflector + 0.15 * recent_avg else: e_dynamic = e_reflector # -------------------------------------------------- # Apply LTI Injection with Dynamic Anchor e_norm = self._px_injection.input_norm(e_dynamic.to(torch.float32)).to(trans_out.dtype) h_exp = trans_out + current_gamma * (e_norm - h_prev) # --- PHASE 26: REFLECTION FLIPPING (RF) --- h_f32 = h_exp.to(torch.float32) e_f32 = e_dynamic.to(torch.float32) dot_he = (h_f32 * e_f32).sum(dim=-1, keepdim=True) dot_ee = (e_f32 * e_f32).sum(dim=-1, keepdim=True) proj = (dot_he / (dot_ee + 1e-6)) * e_f32 ortho = h_f32 - proj # Oscillate the logic vector to avoid local minima flip_force = 0.10 * energy_factor * (1.0 if (exploration_steps % 2 == 0) else -1.0) h_exp = (proj + (1.0 + flip_force) * ortho).to(h_exp.dtype) # ------------------------------------------ # Self-Observation phi_tensor = StabilityMonitor.calculate_phi(h_exp, h_prev) phi = phi_tensor.item() # Merged Telemetry Step telemetry_data = { "step": exploration_steps, "layer": int(current_layer), "phi": float(phi), "gamma": float(current_gamma), "energy": float(energy_factor), "rba_active": len(thought_history) > 2, "hub": int(bimodal_hub) } # Phase 26: Dynamic Loop Extension # If phi is low (< 0.85), allow model to think longer than max_steps if phi < 0.85 and exploration_steps == max_steps - 1 and max_steps < 64: max_steps += (dynamic_end - dynamic_start) # Add 1 full loop # ---------------------------------- step_info = { "step": exploration_steps, "layer": int(current_layer), "phi": float(phi), "decision": None } phis.append(phi) path_label = f"L{current_layer}({phi:.2f})" path_taken.append(path_label) # Phase 12.5/18: Universal Bimodal Path Selection bimodal_threshold = min(0.995, 1.0 - (0.05 * current_gamma)) # Scaled trigger if current_layer == bimodal_hub and phi < bimodal_threshold: step_info["decision"] = "BIMODAL_FORK" path_taken.append("BIMODAL_FORK") # Branch A (Standard) h_a = h_exp.clone() # Branch B (High-Entropy DTEC) jitter_boost = 1.0 + (stability_counter * 0.5) hub_entropy = max(0.01, 1.0 - phi) * 0.5 * jitter_boost # Increased for bf16 visibility h_b = h_exp.to(torch.float32) + torch.randn_like(h_exp, dtype=torch.float32) * hub_entropy h_b = h_b.to(h_exp.dtype) # Lookahead to NEXT layer next_l = current_layer + 1 if next_l < len(self.layers): nl_type = mask_config.layer_types[next_l] # Phase 14.5: Use Functional Read-Only Cache for lookahead lookahead_past = RecursiveMemoryCache( past_key_values, thought_history, layer_types=mask_config.layer_types, read_only=True, expected_len=expected_len ) if past_key_values is not None else None out_a = self.layers[next_l]( h_a, attention_mask=causal_mask_mapping[nl_type], position_embeddings=position_embeddings[nl_type], position_ids=position_ids, past_key_values=lookahead_past, **kwargs )[0] phi_a = StabilityMonitor.calculate_phi(out_a, h_a).item() out_b = self.layers[next_l]( h_b, attention_mask=causal_mask_mapping[nl_type], position_embeddings=position_embeddings[nl_type], position_ids=position_ids, past_key_values=lookahead_past, **kwargs )[0] phi_b = StabilityMonitor.calculate_phi(out_b, h_b).item() if phi_b >= phi_a: h_exp = h_b step_info["fork_winner"] = "B" path_taken.append(f"FORK_B_WON({phi_b:.4f}>={phi_a:.4f})") else: h_exp = h_a step_info["fork_winner"] = "A" path_taken.append(f"FORK_A_WON({phi_a:.4f}>{phi_b:.4f})") else: h_exp = h_b # Phase 9.1: SRJ jitter_scale = max(0.0, 1.0 - phi) * 0.05 if jitter_scale > 0: h_exp = h_exp + torch.randn_like(h_exp) * jitter_scale # OSS - Safe FP32 Calculation for FP16 models (Phase 18) h_exp_f32 = h_exp.to(torch.float32) norm_orig = h_exp_f32.norm(dim=-1, keepdim=True) e_ref_f32 = e_reflector.to(torch.float32) dot_he = (h_exp_f32 * e_ref_f32).sum(dim=-1, keepdim=True) dot_ee = (e_ref_f32 * e_ref_f32).sum(dim=-1, keepdim=True) proj = (dot_he / (dot_ee + 1e-6)) * e_ref_f32 proj = proj.to(h_exp.dtype) ortho = h_exp - proj # Phase 14.8: Step-Entropy Destabilization (SED) if stability_counter > 2: # Phase 15.9: Nonlinear Repulsion # Phase 17.7: Scale-Agnostic Dampening (smaller force for deeper models) scale_factor = 26.0 / cfg.get("num_layers", 26.0) repulsion_force = 0.10 * (stability_counter ** 2) * scale_factor h_exp = h_exp + repulsion_force * (ortho / (ortho.norm(dim=-1, keepdim=True) + 1e-6)) path_taken.append(f"SED_PUSH({repulsion_force:.2f})") # Dynamic Jump proportional to depth jump = max(1, int(cfg.get("num_layers", 18) * 0.1)) current_layer = min(cfg["recur_end"] - 1, current_layer + jump) gain_factor = max(1.0, min(1.15, 1.0 + (1.0 - phi) * 0.4)) damping_factor = max(0.85, min(1.0, 1.0 - (1.0 - phi) * 0.2)) h_exp = damping_factor * proj + gain_factor * ortho # Safe FP32 final normalization h_exp_f32_final = h_exp.to(torch.float32) norm_f32 = h_exp_f32_final.norm(dim=-1, keepdim=True) norm_orig_f32 = norm_orig.to(torch.float32) h_exp = (h_exp_f32_final * (norm_orig_f32 / (norm_f32 + 1e-6))).to(h_exp.dtype) # Store thought for RSM if exploration_steps % 2 == 0: thought_history.append(h_exp.detach()) # Phase 18: Universal ALR Thresholds based on internal parameters (gamma) # Relax thresholds dynamically if we visit a layer too often (loop breaking) visit_penalty = (layer_visits[current_layer] - 1) * 0.015 t_back_2 = 1.0 - (0.8 * current_gamma) - visit_penalty t_back_1 = 1.0 - (0.4 * current_gamma) - visit_penalty t_skip = 1.0 - (0.01 * current_gamma) - (visit_penalty * 0.5) if phi < t_back_2: # High confusion current_layer = max(active_start, current_layer - 2) routing = "BACK-2" elif phi < t_back_1: # Moderate confusion current_layer = max(active_start, current_layer - 1) routing = "BACK-1" elif phi > t_skip: # Extreme stability current_layer += 2 # Skip routing = "SKIP-1" stability_counter += 1 else: current_layer += 1 routing = "NEXT" stability_counter = 0 # Clamp current_layer to prevent underflow if current_layer < active_start: current_layer = active_start routing = "CLAMPED" step_info["routing"] = routing telemetry_data["routing"] = routing if os.environ.get("DEBUG_PX") == "1": telemetry_steps.append(telemetry_data) if stability_counter > 5: break exploration_steps += 1 avg_phi_explore = sum(phis)/len(phis) if phis else 1.0 # Phase 4.1: QBI Blend b_min = cfg.get("beta_reasoning", 0.05) b_max = cfg.get("beta_grounding", 0.18) beta_final = b_min + (b_max - b_min) * (avg_phi_explore ** 2) hidden_states = (1.0 - beta_final) * h_baseline + beta_final * h_exp else: hidden_states = h_baseline self._px_phi = avg_phi_explore self._px_loops_run = exploration_steps self._px_path = path_taken # Phase 14.2: Structured Telemetry Log if not hasattr(self, "_px_telemetry"): self._px_telemetry = [] self._px_telemetry.append({ "pos": int(position_ids[0, 0].item()), "avg_phi": float(avg_phi_explore), "steps": telemetry_steps }) # Phase 11.0: Metacognitive Triggering # If stability is low during the very first token generation, # we flag this as a 'Complex Problem'. if not hasattr(self, "_px_complexity_acc"): self._px_complexity_acc = [] # If we see the first token of a sequence, clear the accumulator if position_ids[0, 0] == 0: self._px_complexity_acc = [] self._px_complexity_acc.append(avg_phi_explore) # Trigger if average stability of the prompt processing is low # selective threshold: 0.90 self._px_trigger_scratchpad = (len(self._px_complexity_acc) > 3 and sum(self._px_complexity_acc) / len(self._px_complexity_acc) < 0.92) # ── 3. CODA ───────────────────────────────────────────────────────────── dynamic_coda_start = dynamic_end if cfg.get("routing_mode") == "adaptive" else cfg["coda_start"] # Phase 14.5: Coda-Grounding Injection (CGI) # Re-inject sensory data to prevent 'hallucinatory drift' in final reasoning. for i in range(dynamic_coda_start, len(self.layers)): if i == dynamic_coda_start: # Phase 14.7: Reverted CGI (8%) hidden_states = 0.92 * hidden_states + 0.08 * e_static layer_out = self.layers[i]( hidden_states, attention_mask=causal_mask_mapping[mask_config.layer_types[i]], position_embeddings=position_embeddings[mask_config.layer_types[i]], position_ids=position_ids, past_key_values=past_key_values, **kwargs, ) hidden_states = layer_out[0] if isinstance(layer_out, (tuple, list)) else layer_out hidden_states = self.norm(hidden_states) # Phase 25: Save Telemetry if enabled if os.environ.get("DEBUG_PX") == "1" and len(telemetry_steps) > 0: ts = datetime.datetime.now().strftime("%Y%m%d_%H%M%S_%f") log_path = f"px_telemetry_{ts}.json" with open(log_path, "w") as f: json.dump(telemetry_steps, f, indent=2) # Only output the path as requested print(f"TELEMETRY_JSON: {os.path.abspath(log_path)}") from transformers.modeling_outputs import BaseModelOutputWithPast return BaseModelOutputWithPast( last_hidden_state=hidden_states, past_key_values=past_key_values, ) # --------------------------------------------------------------------------- def apply_px_patch(model, **cfg_kwargs): from .px_modules import LTIInjection # Robust Text Model Resolver (Phase 17.9) # We look for the module that contains 'layers' and 'rotary_emb' text_model = None if hasattr(model, "layers") and hasattr(model, "rotary_emb"): text_model = model else: # Search children (e.g., .model, .language_model.model) for name, module in model.named_modules(): if hasattr(module, "layers") and hasattr(module, "rotary_emb"): text_model = module break if text_model is None: raise AttributeError(f"Could not identify Gemma-3 text backbone in {type(model)}") config = model.config # Multimodal check: larger models (4B+) wrap text config if hasattr(config, "text_config"): config = config.text_config num_layers = config.num_hidden_layers # Scale-Aware Hyperparameters (Phase 17.100) # - Gamma: Inverse-proportional to hidden size # - Prelude: Shallow models need deeper grounding before recursion hidden_size = config.hidden_size num_layers = config.num_hidden_layers # Phase 25: Balanced Precision Tuning if hidden_size == 640: # 270M defaults = { "mode": "lti", "n_loops": 16, "beta": 0.05, "gamma": 0.12, "recur_start": 8, "recur_end": 12, "bimodal_hub": 11, "cgi_factor": 0.08, "num_layers": num_layers } elif hidden_size == 1152: # 1B defaults = { "mode": "lti", "n_loops": 8, "beta": 0.05, "gamma": 0.12, "recur_start": 10, "recur_end": 20, "bimodal_hub": 18, "cgi_factor": 0.08, "num_layers": num_layers } elif hidden_size == 2560: # 4B defaults = { "mode": "lti", "n_loops": 6, "beta": 0.05, "gamma": 0.05, "recur_start": 5, "recur_end": 33, "bimodal_hub": 32, "cgi_factor": 0.08, "num_layers": num_layers } else: # Fallback for unknown sizes gamma_scale = 1152.0 / hidden_size base_gamma = 0.12 * gamma_scale prelude_ratio = 0.40 if num_layers < 20 else 0.15 p_start = max(4, int(num_layers * prelude_ratio)) p_end = max(p_start + 1, int(num_layers * 0.98)) p_hub = max(p_start + 1, int(num_layers * 0.95)) defaults = { "mode": "lti", "n_loops": 4, "beta": 0.05, "gamma": base_gamma, "recur_start": p_start, "recur_end": p_end, "bimodal_hub": p_hub, "cgi_factor": 0.08, "num_layers": num_layers } defaults.update(cfg_kwargs) # Auto-align boundaries if "prelude_end" not in defaults: defaults["prelude_end"] = defaults["recur_start"] if "coda_start" not in defaults: defaults["coda_start"] = defaults["recur_end"] text_model._px_config = defaults text_model._px_injection = LTIInjection(config.hidden_size, gamma=defaults["gamma"]) text_model.forward = types.MethodType(_px_forward, text_model) print(f"[gemma3-px] Auto-Patch active for scale L{num_layers}. Recur: L{defaults['recur_start']}-L{defaults['recur_end']}, Hub: L{defaults['bimodal_hub']}.") def get_px_metrics(model): text_model = None if hasattr(model, "layers") and hasattr(model, "rotary_emb"): text_model = model else: for name, module in model.named_modules(): if hasattr(module, "layers") and hasattr(module, "rotary_emb"): text_model = module break if text_model is None: text_model = (model.model if hasattr(model, "model") else model) return { "phi": getattr(text_model, "_px_phi", 1.0), "steps": getattr(text_model, "_px_loops_run", 0), "path": getattr(text_model, "_px_path", []), "telemetry": getattr(text_model, "_px_telemetry", []), }