Upload architectureV3.py

architectureV3.py

@@ -99,13 +99,19 @@ class ReflectiveMemoryLayer(nn.Module):
 
     def forward(self, x: torch.Tensor):
         base_output = self.linear(x)
-        if 'embeds' not in self.global_state_storage: return base_output
+        if 'embeds' not in self.global_state_storage: 
+            return base_output
 
         global_embeds = self.global_state_storage['embeds']
-        if global_embeds.shape[1] != x.shape[1]: global_embeds = global_embeds[:, -x.shape[1]:, :]
+        if global_embeds.shape[1] != x.shape[1]: 
+            global_embeds = global_embeds[:, -x.shape[1]:, :]
         B, S, _ = x.shape
         
+        # CRITICAL FIX: Always detach LTM state to prevent backward through previous graphs
         ltm_state = self.global_state_storage.get('ltm', None)
+        if ltm_state is not None:
+            ltm_state = ltm_state.detach()
+        
         proj_local = self.local_state_proj(x)
         proj_global = self.global_state_proj(global_embeds)
         memory_input = torch.stack([proj_global, proj_local], dim=2)
@@ -122,6 +128,7 @@ class ReflectiveMemoryLayer(nn.Module):
         if new_ltm_state_expanded is not None:
              num_ltm_slots = new_ltm_state_expanded.shape[1]
              new_ltm_condensed = new_ltm_state_expanded.view(B, S, num_ltm_slots, self.memory_dim).mean(dim=1)
+             # CRITICAL FIX: Always detach when storing in global state
              self.global_state_storage['ltm'] = new_ltm_condensed.detach()
 
         initial_thought = compressed_mem_flat.mean(dim=1).view(B, S, self.memory_dim)
@@ -143,10 +150,11 @@ class ReflectiveMemoryLayer(nn.Module):
         final_activation = base_output * torch.sigmoid(gate.to(x.dtype)) + value.to(x.dtype)
         
         if self.training:
-            self.last_corrected_activation = final_activation
-            self.last_additive_correction = value
-            self.last_memory_input = memory_input
-            self.last_reconstructed_from_memory = recon_flat.view(B, S, 2, self.memory_dim)
+            # CRITICAL FIX: Detach tensors stored for debugging/analysis
+            self.last_corrected_activation = final_activation.detach()
+            self.last_additive_correction = value.detach()
+            self.last_memory_input = memory_input.detach()
+            self.last_reconstructed_from_memory = recon_flat.view(B, S, 2, self.memory_dim).detach()
         return final_activation
 
 # --- BUILDING BLOCK 3: The Full Custom Model with State Management ---
@@ -155,16 +163,19 @@ class Phi3WithReflectiveMemoryForCausalLM(Phi3ForCausalLM):
         super().__init__(config)
         self.global_state_storage = {}
         self.target_layer_path = "model.layers.15.mlp.gate_up_proj"
-        self.memory_dim, self.num_long_term_memory_slots = 64, 32
+        self.memory_dim, self.num_long_term_memory_slots = 128, 32
 
-        self.model.embed_tokens.register_forward_hook(
-            lambda module, input, output: self.global_state_storage.update({'embeds': output}))
+        # CRITICAL FIX: Ensure embeddings are detached when stored
+        def embedding_hook(module, input, output):
+            self.global_state_storage['embeds'] = output.detach()
+        
+        self.model.embed_tokens.register_forward_hook(embedding_hook)
 
         try:
             original_layer = self.get_submodule(self.target_layer_path)
             custom_layer = ReflectiveMemoryLayer(
                 original_layer=original_layer, global_input_dim=config.hidden_size,
-                memory_dim=self.memory_dim, num_memory_slots=8, memory_num_heads=4,
+                memory_dim=self.memory_dim, num_memory_slots=16, memory_num_heads=4,
                 global_state_storage=self.global_state_storage)
             parent_path = ".".join(self.target_layer_path.split('.')[:-1])
             setattr(self.get_submodule(parent_path), self.target_layer_path.split('.')[-1], custom_layer)
@@ -184,13 +195,19 @@ class Phi3WithReflectiveMemoryForCausalLM(Phi3ForCausalLM):
                 output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None,
                 ltm_state: Optional[torch.Tensor] = None):
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-        self.global_state_storage.clear()
+        
+        # CRITICAL FIX: Don't clear global state storage completely, just reset embeds
+        # This prevents losing LTM state continuity
+        if 'embeds' in self.global_state_storage:
+            del self.global_state_storage['embeds']
         
         # *** FIX: Initialize LTM state if not provided, for both training and first step of inference ***
         if ltm_state is None:
              batch_size = input_ids.shape[0] if input_ids is not None else inputs_embeds.shape[0]
              ltm_state = self._init_ltm_state(batch_size, self.device, self.dtype)
-        self.global_state_storage['ltm'] = ltm_state
+        
+        # CRITICAL FIX: Ensure LTM state is detached when stored
+        self.global_state_storage['ltm'] = ltm_state.detach() if ltm_state is not None else None
 
         outputs = self.model(
             input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids,
@@ -207,7 +224,10 @@ class Phi3WithReflectiveMemoryForCausalLM(Phi3ForCausalLM):
                             labels[..., 1:].contiguous().view(-1))
             # Note: Auxiliary losses from main.py are calculated outside the model forward pass.
 
+        # CRITICAL FIX: Ensure returned LTM state is detached
         new_ltm_state = self.global_state_storage.get('ltm', None)
+        if new_ltm_state is not None:
+            new_ltm_state = new_ltm_state.detach()
         
         if not return_dict:
             output = (logits,) + outputs[1:] + (new_ltm_state,)
@@ -215,4 +235,4 @@ class Phi3WithReflectiveMemoryForCausalLM(Phi3ForCausalLM):
 
         return CausalLMOutputWithLTM(
             loss=loss, logits=logits, past_key_values=outputs.past_key_values,
-            hidden_states=outputs.hidden_states, attentions=outputs.attentions, ltm_state=new_ltm_state)
+            hidden_states=outputs.hidden_states, attentions=outputs.attentions, ltm_state=new_ltm_state)