Instructions to use botverse/Titans-MAC-Llama-3-8b-Instruct with libraries, inference providers, notebooks, and local apps. Follow these links to get started.

Libraries

How to use botverse/Titans-MAC-Llama-3-8b-Instruct with Transformers:

# Use a pipeline as a high-level helper
from transformers import pipeline

pipe = pipeline("text-generation", model="botverse/Titans-MAC-Llama-3-8b-Instruct", trust_remote_code=True)
messages = [
    {"role": "user", "content": "Who are you?"},
]
pipe(messages)

# pip install -U transformers accelerate
# Load model directly
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("botverse/Titans-MAC-Llama-3-8b-Instruct", trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained("botverse/Titans-MAC-Llama-3-8b-Instruct", trust_remote_code=True, device_map="auto")
messages = [
    {"role": "user", "content": "Who are you?"},
]
inputs = tokenizer.apply_chat_template(
	messages,
	add_generation_prompt=True,
	tokenize=True,
	return_dict=True,
	return_tensors="pt",
).to(model.device)

outputs = model.generate(**inputs, max_new_tokens=256)
print(tokenizer.decode(outputs[0][inputs["input_ids"].shape[-1]:]))

Notebooks
Google Colab
Kaggle
Local Apps Settings

vLLM

How to use botverse/Titans-MAC-Llama-3-8b-Instruct with vLLM:

Install from pip and serve model

# Install vLLM from pip:
pip install vllm
# Start the vLLM server:
vllm serve "botverse/Titans-MAC-Llama-3-8b-Instruct"
# Call the server using curl (OpenAI-compatible API):
curl -X POST "http://localhost:8000/v1/chat/completions" \
	-H "Content-Type: application/json" \
	--data '{
		"model": "botverse/Titans-MAC-Llama-3-8b-Instruct",
		"messages": [
			{
				"role": "user",
				"content": "What is the capital of France?"
			}
		]
	}'

Use Docker

docker model run hf.co/botverse/Titans-MAC-Llama-3-8b-Instruct

SGLang

How to use botverse/Titans-MAC-Llama-3-8b-Instruct with SGLang:

Install from pip and serve model

# Install SGLang from pip:
pip install sglang
# Start the SGLang server:
python3 -m sglang.launch_server \
    --model-path "botverse/Titans-MAC-Llama-3-8b-Instruct" \
    --host 0.0.0.0 \
    --port 30000
# Call the server using curl (OpenAI-compatible API):
curl -X POST "http://localhost:30000/v1/chat/completions" \
	-H "Content-Type: application/json" \
	--data '{
		"model": "botverse/Titans-MAC-Llama-3-8b-Instruct",
		"messages": [
			{
				"role": "user",
				"content": "What is the capital of France?"
			}
		]
	}'

Use Docker images

docker run --gpus all \
    --shm-size 32g \
    -p 30000:30000 \
    -v ~/.cache/huggingface:/root/.cache/huggingface \
    --env "HF_TOKEN=<secret>" \
    --ipc=host \
    lmsysorg/sglang:latest \
    python3 -m sglang.launch_server \
        --model-path "botverse/Titans-MAC-Llama-3-8b-Instruct" \
        --host 0.0.0.0 \
        --port 30000
# Call the server using curl (OpenAI-compatible API):
curl -X POST "http://localhost:30000/v1/chat/completions" \
	-H "Content-Type: application/json" \
	--data '{
		"model": "botverse/Titans-MAC-Llama-3-8b-Instruct",
		"messages": [
			{
				"role": "user",
				"content": "What is the capital of France?"
			}
		]
	}'

Docker Model Runner
How to use botverse/Titans-MAC-Llama-3-8b-Instruct with Docker Model Runner:
```
docker model run hf.co/botverse/Titans-MAC-Llama-3-8b-Instruct
```

Titans-MAC-Llama-3-8b-Instruct / models /inference_memory_wrapper.py

botverse

added hf url

db7c8e4 over 1 year ago

Raw History Blame Contribute Delete

22.2 kB

import torch
import torch.nn as nn
import torch.nn.functional as F
import math
from transformers import LlamaForCausalLM, LlamaConfig, PreTrainedModel
from transformers.modeling_outputs import CausalLMOutputWithPast
from transformers.cache_utils import Cache
from typing import Optional, List, Tuple, Union
import os
from pathlib import Path

# Use the actual LlamaForCausalLM from the packaged 'models' dir if needed,
# but relying on the globally installed transformers version is usually fine.
# from .hf_llama.modeling_llama import LlamaForCausalLM, LlamaConfig

class InferenceMemoryWrapper(PreTrainedModel):
    # config_class = LlamaConfig # Keep if needed for saving config

    # --- REVERTED __init__ signature ---
    def __init__(self, llama_model: LlamaForCausalLM, memory_size: int = 4096, num_retrieved: int = 1, update_alpha: float = 0.1, surprise_momentum: float = 0.9, surprise_lr: float = 0.01):
        super().__init__(llama_model.config) # Use config from the passed model
        self.llama = llama_model # Store the pre-loaded model

        # --- Use passed parameters ---
        self.memory_size = memory_size
        self.num_retrieved = num_retrieved
        self.update_alpha = update_alpha
        self.surprise_momentum_eta = surprise_momentum
        self.surprise_lr_theta = surprise_lr
        self.dim = llama_model.config.hidden_size
        self._target_dtype = llama_model.dtype # Get dtype from the base model (should be float16)

        # --- Memory buffer is a Parameter ---
        # Create tensor directly with correct dtype on CPU initially
        init_buffer_data = torch.zeros(self.memory_size, self.dim, dtype=self._target_dtype)
        # Initialize in place
        nn.init.normal_(init_buffer_data, mean=0.0, std=1 / math.sqrt(self.dim))
        # Wrap in Parameter (Parameter itself doesn't change dtype)
        self.memory_buffer = nn.Parameter(init_buffer_data)


        # --- Surprise Update State ---
        # Create tensor directly with correct dtype on CPU initially
        init_surprise_state = torch.zeros_like(self.memory_buffer.data, dtype=self._target_dtype) # Use buffer's shape/dtype
        self.register_buffer("surprise_state", init_surprise_state)


        # --- Freeze the underlying Llama model ---
        for param in self.llama.parameters():
            param.requires_grad = False
        self.llama.eval() # Keep llama in eval mode

    # --- Keep existing methods (get_input_embeddings, set_input_embeddings, etc.) ---
    def get_input_embeddings(self):
        return self.llama.get_input_embeddings()

    def set_input_embeddings(self, value):
        self.llama.set_input_embeddings(value)

    def get_output_embeddings(self):
        return self.llama.get_output_embeddings()

    def set_output_embeddings(self, new_embeddings):
        self.llama.set_output_embeddings(new_embeddings)

    # --- Differentiable Attention Retrieval ---
    def retrieve_memory(self, query_input: torch.Tensor) -> torch.Tensor:
        """
        Retrieves memory using differentiable attention based on query_input.
        Args:
            query_input (torch.Tensor): Query tensor. Shape (B, C).
        Returns:
            torch.Tensor: Retrieved memory embedding (weighted sum). Shape (B, 1, C)
        """
        # Ensure query is the correct dtype (should match memory buffer)
        q = query_input.to(self.memory_buffer.dtype) # Still check against buffer's actual dtype

        # Use memory_buffer directly as keys and values
        # self.memory_buffer should now consistently be self._target_dtype (float16)
        mem_keys = self.memory_buffer # (memory_size, C)
        mem_values = self.memory_buffer # (memory_size, C)

        # Matmul should now work as dtypes match
        attn_scores = torch.matmul(q, mem_keys.T) / math.sqrt(self.dim) # (B, memory_size)
        attn_weights = torch.softmax(attn_scores, dim=-1) # (B, memory_size)

        # Ensure retrieved mem is also the correct dtype before returning
        retrieved_mem = torch.matmul(attn_weights, mem_values) # (B, C)

        return retrieved_mem.unsqueeze(1) # (B, 1, C)

    # --- Surprise Update Application ---
    @torch.no_grad()
    def apply_surprise_update(self):
        """ Applies the TITANS-style surprise update rule using self.memory_buffer.grad """
        if self.memory_buffer.grad is None:
            print("DEBUG: apply_surprise_update called but memory_buffer.grad is None.")
            return

        # Ensure surprise_state is on the same device and dtype
        self.surprise_state = self.surprise_state.to(device=self.memory_buffer.device, dtype=self.memory_buffer.dtype)

        # Grad should have the same dtype as the parameter
        surprise_update_val = -self.surprise_lr_theta * self.memory_buffer.grad.data
        self.surprise_state.mul_(self.surprise_momentum_eta).add_(surprise_update_val)

        self.memory_buffer.data.add_(self.surprise_state)
        self.memory_buffer.grad.zero_()


    # --- EMA Update (Alternative, No Gradients) ---
    @torch.no_grad()
    def update_memory_ema(self, new_context_embedding: torch.Tensor):
        """ Updates the memory buffer using EMA. """
        # Ensure update vector is the correct dtype
        update_vec_float = new_context_embedding.mean(dim=0, keepdim=True) if new_context_embedding.shape[0] > 1 else new_context_embedding # (1, C)
        update_vec = update_vec_float.to(self.memory_buffer.dtype)

        # Ensure buffer is on the correct device before update
        self.memory_buffer.data = self.memory_buffer.data.to(update_vec.device)
        self.memory_buffer.data.mul_(1 - self.update_alpha).add_(update_vec * self.update_alpha)


    # --- Forward Pass (Pass-through to Llama) ---
    # Overriding forward is needed if we want AutoModelForCausalLM(wrapper) to work directly
    # This now needs to call self.llama.forward
    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        labels: Optional[torch.LongTensor] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        cache_position: Optional[torch.LongTensor] = None,
        **kwargs, # Pass any extra kwargs
    ) -> Union[Tuple, CausalLMOutputWithPast]:
         # Directly call the wrapped llama model's forward pass
         # Note: This basic forward doesn't include the memory prepending logic.
         # That logic is currently only in the custom generate method.
         # If you wanted to use model(input_ids) directly *with* memory,
         # you'd need to replicate the generate logic here.
        return self.llama(
            input_ids=input_ids,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_values=past_key_values,
            inputs_embeds=inputs_embeds,
            labels=labels,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            cache_position=cache_position,
            **kwargs,
        )

    # --- MODIFIED Generate Method with Inline Backward Pass ---
    # (Generate method remains largely the same as before, but ensure it uses self.llama correctly)
    def generate(
        self,
        input_ids: torch.LongTensor,
        max_new_tokens: int = 20,
        num_beams: int = 1,
        use_memory: bool = True,
        update_rule: str = 'ema',
        temperature: float = 0.7,
        top_p: float = 0.95,
        do_sample: bool = True,
        repetition_penalty: float = 1.0,
        eos_token_id: Optional[int] = None,
        pad_token_id: Optional[int] = None,
        attention_mask: Optional[torch.Tensor] = None, # Added attention_mask parameter
        **kwargs,
    ) -> torch.LongTensor:
        if num_beams != 1:
            raise NotImplementedError("Beam search not implemented.")
        if update_rule == 'surprise' and not use_memory:
            print("Warning: update_rule='surprise' requires use_memory=True.")
            update_rule = 'none'

        # Ensure buffer requires grad only when needed
        original_requires_grad = self.memory_buffer.requires_grad
        if update_rule == 'surprise':
             self.memory_buffer.requires_grad_(True)
             print(f"DEBUG: Set memory_buffer.requires_grad = {self.memory_buffer.requires_grad}")
        else:
             self.memory_buffer.requires_grad_(False)


        bsz, seq_len_start = input_ids.shape
        device = input_ids.device
        generated_ids = input_ids.clone()
        current_seq_len = seq_len_start
        # Determine the expected dtype from the buffer
        expected_dtype = self.memory_buffer.dtype # Use actual buffer dtype

        if eos_token_id is None: eos_token_id = self.config.eos_token_id
        if pad_token_id is None: pad_token_id = self.config.pad_token_id

        past_key_values = None # Initialize KV cache

        # Prepare initial attention mask if provided
        if attention_mask is None:
             attention_mask = torch.ones_like(input_ids)

        for step in range(max_new_tokens):
            # --- Prepare Inputs for this step ---
            # Use only the last token for generation if KV cache is active
            if past_key_values is not None:
                current_input_ids = generated_ids[:, -1:]
                # We need the hidden state/embedding of the *previous* token to query memory
                # Let's get the full embeddings first, then select the query basis
                # Use the full sequence length processed so far for embeddings
                full_embeds = self.llama.model.embed_tokens(generated_ids) # (B, T_cur, C)
                # Ensure query_basis has the expected dtype
                query_basis = full_embeds[:, -1, :].to(expected_dtype) # Query based on the last token generated *before* this step
            else:
                current_input_ids = generated_ids
                inputs_embeds_full = self.llama.model.embed_tokens(current_input_ids) # (B, T_cur, C)
                # Ensure query_basis has the expected dtype
                query_basis = inputs_embeds_full[:, -1, :].to(expected_dtype) # Query based on last token of the input prompt


            # --- Memory Retrieval ---
            retrieved_mem = None
            if use_memory:
                # query_basis should now match memory_buffer dtype
                retrieved_mem = self.retrieve_memory(query_basis) # (B, 1, C)

            # --- Combine Embeddings and Prepare Model Inputs ---
            # Manage attention mask and position IDs carefully
            current_mask = None
            mem_len = 0
            if retrieved_mem is not None:
                 retrieved_mem_casted = retrieved_mem.to(self.llama.dtype) # (B, 1, C_llama)
                 mem_len = retrieved_mem_casted.shape[1] # Should be 1

            if past_key_values is None: # First step
                inputs_embeds_full_casted = inputs_embeds_full.to(self.llama.dtype) # (B, T_cur, C_llama)
                if retrieved_mem is not None:
                    model_inputs_embeds = torch.cat([retrieved_mem_casted, inputs_embeds_full_casted], dim=1) # (B, 1 + T_cur, C)
                    # Create mask for memory + original input mask
                    mem_mask = torch.ones((bsz, mem_len), dtype=attention_mask.dtype, device=device)
                    current_mask = torch.cat([mem_mask, attention_mask], dim=1) # (B, 1 + T_cur)
                else:
                    model_inputs_embeds = inputs_embeds_full_casted # (B, T_cur, C)
                    current_mask = attention_mask # Use original mask

                effective_seq_len = model_inputs_embeds.shape[1]
                position_ids = torch.arange(effective_seq_len, device=device).unsqueeze(0) # (1, P+K+T)
                cur_input_ids_for_llama = None # Using embeds
            else: # Subsequent steps with KV cache
                current_input_embeds = self.llama.model.embed_tokens(current_input_ids).to(self.llama.dtype) # (B, 1, C_llama)
                if retrieved_mem is not None:
                     model_inputs_embeds = torch.cat([retrieved_mem_casted, current_input_embeds], dim=1) # (B, 1 + 1, C)
                     # Mask for memory + current token
                     current_mask = torch.ones((bsz, mem_len + 1), dtype=attention_mask.dtype, device=device) # (B, 1 + 1)
                else:
                     model_inputs_embeds = current_input_embeds # (B, 1, C)
                     # Mask for current token only
                     current_mask = torch.ones((bsz, 1), dtype=attention_mask.dtype, device=device) # (B, 1)

                # Position ID for the new token(s) relative to KV cache length + memory length
                # LlamaModel._update_causal_mask and cache handling expect position_ids to reflect the absolute position
                # cache_position (passed internally by generate if use_cache) handles this. We construct it manually here.
                # The position id for the *new token* is the current sequence length (including memory if prepended this step)
                past_len = past_key_values.get_seq_length() # Length stored in cache
                # The position_id should reflect where this new token/memory would be in the *full* sequence if no cache was used
                # Let's use current_seq_len derived from generated_ids, which doesn't include memory
                position_ids = torch.tensor([[current_seq_len -1 + i + mem_len for i in range(model_inputs_embeds.shape[1])]], device=device) # (1, M+1) or (1, 1)

                cur_input_ids_for_llama = None # Using embeds

            # --- Llama Forward Pass ---
            # Use KV caching if possible (update_rule != 'surprise')
            # We need past_key_values AND not be doing surprise update AND base model supports caching
            use_kv_cache_this_step = past_key_values is not None and update_rule != 'surprise' and self.llama.config.use_cache

            # Ensure context manager enables grads only when needed
            context = torch.enable_grad() if update_rule == 'surprise' else torch.no_grad()
            with context:
                outputs = self.llama(
                    input_ids=cur_input_ids_for_llama, # None if using embeds
                    inputs_embeds=model_inputs_embeds,
                    attention_mask=current_mask, # Pass the correctly shaped mask for this step
                    position_ids=position_ids, # Pass adjusted position IDs
                    past_key_values=past_key_values,
                    use_cache=use_kv_cache_this_step,
                    output_hidden_states=True, # Needed for query/target/update
                    return_dict=True,
                )

            # --- Associative Loss Calculation (if surprise update) ---
            if update_rule == 'surprise' and use_memory and retrieved_mem is not None:
                 # Target: Final hidden state corresponding to the *last input token* before generation
                 # The index needs to account for the prepended memory.
                 # If mem_len=1, the target state corresponds to index -1 in the output sequence
                 target_repr = outputs.hidden_states[-1][:, -1, :].to(self.memory_buffer.dtype) # (B, C)

                 # pred_repr comes from retrieve_memory, should already match buffer dtype
                 pred_repr = retrieved_mem.squeeze(1) # (B, C)

                 # --- DEBUG PRINTS ---
                 print(f"\n--- Surprise Update Debug (Step {step}) ---")
                 print(f"  memory_buffer requires_grad: {self.memory_buffer.requires_grad}")
                 print(f"  retrieved_mem requires_grad: {retrieved_mem.requires_grad if retrieved_mem is not None else 'N/A'}")
                 print(f"  pred_repr requires_grad:     {pred_repr.requires_grad if pred_repr is not None else 'N/A'}")
                 print(f"  target_repr requires_grad:   {target_repr.requires_grad}") # Should be False due to .detach() below
                 # --- END DEBUG ---

                 assoc_loss = F.mse_loss(pred_repr, target_repr.detach()) # TARGET IS DETACHED
                 print(f"  assoc_loss: {assoc_loss.item():.4f}, requires_grad: {assoc_loss.requires_grad}")


                 if self.memory_buffer.grad is not None:
                      print("  Zeroing existing memory_buffer gradient.")
                      self.memory_buffer.grad.zero_()

                 if assoc_loss.requires_grad:
                     print("  Calling assoc_loss.backward()")
                     assoc_loss.backward() # Compute grads for memory_buffer
                     print(f"  memory_buffer.grad is None after backward: {self.memory_buffer.grad is None}")
                     if self.memory_buffer.grad is not None:
                         print(f"  memory_buffer.grad norm: {torch.norm(self.memory_buffer.grad).item():.4f}")
                     self.apply_surprise_update() # Apply update and zero grad
                 else:
                     print("  ERROR: assoc_loss does not require grad! Skipping backward and update.")
                 print("--- End Surprise Update Debug ---")


            # --- Standard Generation Logic ---
            # Get logits for the very last position in the output sequence (corresponds to the token we just fed in)
            next_token_logits = outputs.logits[:, -1, :] # (B, V)

            # Update KV cache for next step
            if use_kv_cache_this_step:
                # The past_key_values returned by Llama should account for the memory prepended in this step
                past_key_values = outputs.past_key_values


            # Sampling (same as before)
            if repetition_penalty != 1.0:
                 # Simple loop for now:
                 for i in range(bsz):
                     # Penalize tokens in the *generated* sequence (excluding prompt if needed)
                     # Use generated_ids which tracks the full sequence
                     for token_id in generated_ids[i]:
                         # Avoid penalizing pad token if present
                         if token_id != pad_token_id:
                            next_token_logits[i, token_id] /= repetition_penalty

            if temperature > 0 and temperature != 1.0:
                next_token_logits = next_token_logits / temperature
            if do_sample and top_p < 1.0:
                # Use Hugging Face's top_p implementation detail
                sorted_logits, sorted_indices = torch.sort(next_token_logits, descending=True)
                cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
                sorted_indices_to_remove = cumulative_probs > top_p
                sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
                sorted_indices_to_remove[..., 0] = 0
                indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
                next_token_logits = next_token_logits.masked_fill(indices_to_remove, float('-inf'))

            if do_sample:
                probs = F.softmax(next_token_logits, dim=-1)
                next_token = torch.multinomial(probs, num_samples=1)
            else:
                next_token = torch.argmax(next_token_logits, dim=-1, keepdim=True)

            # --- Update State ---
            generated_ids = torch.cat([generated_ids, next_token], dim=1)
            current_seq_len += 1
            # Update attention mask for the next iteration by appending 1
            attention_mask = torch.cat([attention_mask, torch.ones((bsz, 1), dtype=attention_mask.dtype, device=device)], dim=1)


            # --- EMA Memory Update ---
            if update_rule == 'ema' and use_memory and outputs.hidden_states is not None:
                 # Use hidden state corresponding to the newly generated token position (index -1)
                 # Cast state to buffer dtype before update
                 new_context_state = outputs.hidden_states[-1][:, -1, :].to(self.memory_buffer.dtype) # (B, C)
                 self.update_memory_ema(new_context_state.detach())

            if eos_token_id is not None and (next_token == eos_token_id).all():
                break

        # Restore original requires_grad state
        self.memory_buffer.requires_grad_(original_requires_grad)

        return generated_ids


    # --- Save/Load ---
    # Keep the save_pretrained as is, it saves wrapper specific state.
    def save_pretrained(self, save_directory: Union[str, os.PathLike], **kwargs):
        """ Saves the wrapper's specific state (memory buffer, surprise state). """
        save_directory = Path(save_directory)
        save_directory.mkdir(parents=True, exist_ok=True)

        # Save the base model's config (important for PreTrainedModel compatibility)
        self.config.save_pretrained(save_directory)

        # Save the memory buffer parameter directly
        # Ensure saving in float32 for broader compatibility, can be cast back on load
        # Note: Saving the Parameter itself, not just its .data
        torch.save(self.memory_buffer.float(), save_directory / "memory_buffer.pt")
        # Save the surprise state buffer directly
        torch.save(self.surprise_state.float(), save_directory / "surprise_state.pt")

        print(f"InferenceMemoryWrapper state saved to {save_directory}")
        # Note: Base Llama model weights are assumed to be saved separately or loaded from source.

    # from_pretrained is complex with wrappers. For local testing/handler, load manually.
    # @classmethod
    # def from_pretrained(cls, pretrained_model_name_or_path: Union[str, os.PathLike], **kwargs):
    #      raise NotImplementedError(...)

	import torch
	import torch.nn as nn
	import torch.nn.functional as F
	import math
	from transformers import LlamaForCausalLM, LlamaConfig, PreTrainedModel
	from transformers.modeling_outputs import CausalLMOutputWithPast
	from transformers.cache_utils import Cache
	from typing import Optional, List, Tuple, Union
	import os
	from pathlib import Path

	# Use the actual LlamaForCausalLM from the packaged 'models' dir if needed,
	# but relying on the globally installed transformers version is usually fine.
	# from .hf_llama.modeling_llama import LlamaForCausalLM, LlamaConfig

	class InferenceMemoryWrapper(PreTrainedModel):
	# config_class = LlamaConfig # Keep if needed for saving config

	# --- REVERTED __init__ signature ---
	def __init__(self, llama_model: LlamaForCausalLM, memory_size: int = 4096, num_retrieved: int = 1, update_alpha: float = 0.1, surprise_momentum: float = 0.9, surprise_lr: float = 0.01):
	super().__init__(llama_model.config) # Use config from the passed model
	self.llama = llama_model # Store the pre-loaded model

	# --- Use passed parameters ---
	self.memory_size = memory_size
	self.num_retrieved = num_retrieved
	self.update_alpha = update_alpha
	self.surprise_momentum_eta = surprise_momentum
	self.surprise_lr_theta = surprise_lr
	self.dim = llama_model.config.hidden_size
	self._target_dtype = llama_model.dtype # Get dtype from the base model (should be float16)

	# --- Memory buffer is a Parameter ---
	# Create tensor directly with correct dtype on CPU initially
	init_buffer_data = torch.zeros(self.memory_size, self.dim, dtype=self._target_dtype)
	# Initialize in place
	nn.init.normal_(init_buffer_data, mean=0.0, std=1 / math.sqrt(self.dim))
	# Wrap in Parameter (Parameter itself doesn't change dtype)
	self.memory_buffer = nn.Parameter(init_buffer_data)


	# --- Surprise Update State ---
	# Create tensor directly with correct dtype on CPU initially
	init_surprise_state = torch.zeros_like(self.memory_buffer.data, dtype=self._target_dtype) # Use buffer's shape/dtype
	self.register_buffer("surprise_state", init_surprise_state)


	# --- Freeze the underlying Llama model ---
	for param in self.llama.parameters():
	param.requires_grad = False
	self.llama.eval() # Keep llama in eval mode

	# --- Keep existing methods (get_input_embeddings, set_input_embeddings, etc.) ---
	def get_input_embeddings(self):
	return self.llama.get_input_embeddings()

	def set_input_embeddings(self, value):
	self.llama.set_input_embeddings(value)

	def get_output_embeddings(self):
	return self.llama.get_output_embeddings()

	def set_output_embeddings(self, new_embeddings):
	self.llama.set_output_embeddings(new_embeddings)

	# --- Differentiable Attention Retrieval ---
	def retrieve_memory(self, query_input: torch.Tensor) -> torch.Tensor:
	"""
	Retrieves memory using differentiable attention based on query_input.
	Args:
	query_input (torch.Tensor): Query tensor. Shape (B, C).
	Returns:
	torch.Tensor: Retrieved memory embedding (weighted sum). Shape (B, 1, C)
	"""
	# Ensure query is the correct dtype (should match memory buffer)
	q = query_input.to(self.memory_buffer.dtype) # Still check against buffer's actual dtype

	# Use memory_buffer directly as keys and values
	# self.memory_buffer should now consistently be self._target_dtype (float16)
	mem_keys = self.memory_buffer # (memory_size, C)
	mem_values = self.memory_buffer # (memory_size, C)

	# Matmul should now work as dtypes match
	attn_scores = torch.matmul(q, mem_keys.T) / math.sqrt(self.dim) # (B, memory_size)
	attn_weights = torch.softmax(attn_scores, dim=-1) # (B, memory_size)

	# Ensure retrieved mem is also the correct dtype before returning
	retrieved_mem = torch.matmul(attn_weights, mem_values) # (B, C)

	return retrieved_mem.unsqueeze(1) # (B, 1, C)

	# --- Surprise Update Application ---
	@torch.no_grad()
	def apply_surprise_update(self):
	""" Applies the TITANS-style surprise update rule using self.memory_buffer.grad """
	if self.memory_buffer.grad is None:
	print("DEBUG: apply_surprise_update called but memory_buffer.grad is None.")
	return

	# Ensure surprise_state is on the same device and dtype
	self.surprise_state = self.surprise_state.to(device=self.memory_buffer.device, dtype=self.memory_buffer.dtype)

	# Grad should have the same dtype as the parameter
	surprise_update_val = -self.surprise_lr_theta * self.memory_buffer.grad.data
	self.surprise_state.mul_(self.surprise_momentum_eta).add_(surprise_update_val)

	self.memory_buffer.data.add_(self.surprise_state)
	self.memory_buffer.grad.zero_()


	# --- EMA Update (Alternative, No Gradients) ---
	@torch.no_grad()
	def update_memory_ema(self, new_context_embedding: torch.Tensor):
	""" Updates the memory buffer using EMA. """
	# Ensure update vector is the correct dtype
	update_vec_float = new_context_embedding.mean(dim=0, keepdim=True) if new_context_embedding.shape[0] > 1 else new_context_embedding # (1, C)
	update_vec = update_vec_float.to(self.memory_buffer.dtype)

	# Ensure buffer is on the correct device before update
	self.memory_buffer.data = self.memory_buffer.data.to(update_vec.device)
	self.memory_buffer.data.mul_(1 - self.update_alpha).add_(update_vec * self.update_alpha)


	# --- Forward Pass (Pass-through to Llama) ---
	# Overriding forward is needed if we want AutoModelForCausalLM(wrapper) to work directly
	# This now needs to call self.llama.forward
	def forward(
	self,
	input_ids: Optional[torch.LongTensor] = None,
	attention_mask: Optional[torch.Tensor] = None,
	position_ids: Optional[torch.LongTensor] = None,
	past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None,
	inputs_embeds: Optional[torch.FloatTensor] = None,
	labels: Optional[torch.LongTensor] = None,
	use_cache: Optional[bool] = None,
	output_attentions: Optional[bool] = None,
	output_hidden_states: Optional[bool] = None,
	return_dict: Optional[bool] = None,
	cache_position: Optional[torch.LongTensor] = None,
	**kwargs, # Pass any extra kwargs
	) -> Union[Tuple, CausalLMOutputWithPast]:
	# Directly call the wrapped llama model's forward pass
	# Note: This basic forward doesn't include the memory prepending logic.
	# That logic is currently only in the custom generate method.
	# If you wanted to use model(input_ids) directly with memory,
	# you'd need to replicate the generate logic here.
	return self.llama(
	input_ids=input_ids,
	attention_mask=attention_mask,
	position_ids=position_ids,
	past_key_values=past_key_values,
	inputs_embeds=inputs_embeds,
	labels=labels,
	use_cache=use_cache,
	output_attentions=output_attentions,
	output_hidden_states=output_hidden_states,
	return_dict=return_dict,
	cache_position=cache_position,
	**kwargs,
	)

	# --- MODIFIED Generate Method with Inline Backward Pass ---
	# (Generate method remains largely the same as before, but ensure it uses self.llama correctly)
	def generate(
	self,
	input_ids: torch.LongTensor,
	max_new_tokens: int = 20,
	num_beams: int = 1,
	use_memory: bool = True,
	update_rule: str = 'ema',
	temperature: float = 0.7,
	top_p: float = 0.95,
	do_sample: bool = True,
	repetition_penalty: float = 1.0,
	eos_token_id: Optional[int] = None,
	pad_token_id: Optional[int] = None,
	attention_mask: Optional[torch.Tensor] = None, # Added attention_mask parameter
	**kwargs,
	) -> torch.LongTensor:
	if num_beams != 1:
	raise NotImplementedError("Beam search not implemented.")
	if update_rule == 'surprise' and not use_memory:
	print("Warning: update_rule='surprise' requires use_memory=True.")
	update_rule = 'none'

	# Ensure buffer requires grad only when needed
	original_requires_grad = self.memory_buffer.requires_grad
	if update_rule == 'surprise':
	self.memory_buffer.requires_grad_(True)
	print(f"DEBUG: Set memory_buffer.requires_grad = {self.memory_buffer.requires_grad}")
	else:
	self.memory_buffer.requires_grad_(False)


	bsz, seq_len_start = input_ids.shape
	device = input_ids.device
	generated_ids = input_ids.clone()
	current_seq_len = seq_len_start
	# Determine the expected dtype from the buffer
	expected_dtype = self.memory_buffer.dtype # Use actual buffer dtype

	if eos_token_id is None: eos_token_id = self.config.eos_token_id
	if pad_token_id is None: pad_token_id = self.config.pad_token_id

	past_key_values = None # Initialize KV cache

	# Prepare initial attention mask if provided
	if attention_mask is None:
	attention_mask = torch.ones_like(input_ids)

	for step in range(max_new_tokens):
	# --- Prepare Inputs for this step ---
	# Use only the last token for generation if KV cache is active
	if past_key_values is not None:
	current_input_ids = generated_ids[:, -1:]
	# We need the hidden state/embedding of the previous token to query memory
	# Let's get the full embeddings first, then select the query basis
	# Use the full sequence length processed so far for embeddings
	full_embeds = self.llama.model.embed_tokens(generated_ids) # (B, T_cur, C)
	# Ensure query_basis has the expected dtype
	query_basis = full_embeds[:, -1, :].to(expected_dtype) # Query based on the last token generated before this step
	else:
	current_input_ids = generated_ids
	inputs_embeds_full = self.llama.model.embed_tokens(current_input_ids) # (B, T_cur, C)
	# Ensure query_basis has the expected dtype
	query_basis = inputs_embeds_full[:, -1, :].to(expected_dtype) # Query based on last token of the input prompt


	# --- Memory Retrieval ---
	retrieved_mem = None
	if use_memory:
	# query_basis should now match memory_buffer dtype
	retrieved_mem = self.retrieve_memory(query_basis) # (B, 1, C)

	# --- Combine Embeddings and Prepare Model Inputs ---
	# Manage attention mask and position IDs carefully
	current_mask = None
	mem_len = 0
	if retrieved_mem is not None:
	retrieved_mem_casted = retrieved_mem.to(self.llama.dtype) # (B, 1, C_llama)
	mem_len = retrieved_mem_casted.shape[1] # Should be 1

	if past_key_values is None: # First step
	inputs_embeds_full_casted = inputs_embeds_full.to(self.llama.dtype) # (B, T_cur, C_llama)
	if retrieved_mem is not None:
	model_inputs_embeds = torch.cat([retrieved_mem_casted, inputs_embeds_full_casted], dim=1) # (B, 1 + T_cur, C)
	# Create mask for memory + original input mask
	mem_mask = torch.ones((bsz, mem_len), dtype=attention_mask.dtype, device=device)
	current_mask = torch.cat([mem_mask, attention_mask], dim=1) # (B, 1 + T_cur)
	else:
	model_inputs_embeds = inputs_embeds_full_casted # (B, T_cur, C)
	current_mask = attention_mask # Use original mask

	effective_seq_len = model_inputs_embeds.shape[1]
	position_ids = torch.arange(effective_seq_len, device=device).unsqueeze(0) # (1, P+K+T)
	cur_input_ids_for_llama = None # Using embeds
	else: # Subsequent steps with KV cache
	current_input_embeds = self.llama.model.embed_tokens(current_input_ids).to(self.llama.dtype) # (B, 1, C_llama)
	if retrieved_mem is not None:
	model_inputs_embeds = torch.cat([retrieved_mem_casted, current_input_embeds], dim=1) # (B, 1 + 1, C)
	# Mask for memory + current token
	current_mask = torch.ones((bsz, mem_len + 1), dtype=attention_mask.dtype, device=device) # (B, 1 + 1)
	else:
	model_inputs_embeds = current_input_embeds # (B, 1, C)
	# Mask for current token only
	current_mask = torch.ones((bsz, 1), dtype=attention_mask.dtype, device=device) # (B, 1)

	# Position ID for the new token(s) relative to KV cache length + memory length
	# LlamaModel._update_causal_mask and cache handling expect position_ids to reflect the absolute position
	# cache_position (passed internally by generate if use_cache) handles this. We construct it manually here.
	# The position id for the new token is the current sequence length (including memory if prepended this step)
	past_len = past_key_values.get_seq_length() # Length stored in cache
	# The position_id should reflect where this new token/memory would be in the full sequence if no cache was used
	# Let's use current_seq_len derived from generated_ids, which doesn't include memory
	position_ids = torch.tensor([[current_seq_len -1 + i + mem_len for i in range(model_inputs_embeds.shape[1])]], device=device) # (1, M+1) or (1, 1)

	cur_input_ids_for_llama = None # Using embeds

	# --- Llama Forward Pass ---
	# Use KV caching if possible (update_rule != 'surprise')
	# We need past_key_values AND not be doing surprise update AND base model supports caching
	use_kv_cache_this_step = past_key_values is not None and update_rule != 'surprise' and self.llama.config.use_cache

	# Ensure context manager enables grads only when needed
	context = torch.enable_grad() if update_rule == 'surprise' else torch.no_grad()
	with context:
	outputs = self.llama(
	input_ids=cur_input_ids_for_llama, # None if using embeds
	inputs_embeds=model_inputs_embeds,
	attention_mask=current_mask, # Pass the correctly shaped mask for this step
	position_ids=position_ids, # Pass adjusted position IDs
	past_key_values=past_key_values,
	use_cache=use_kv_cache_this_step,
	output_hidden_states=True, # Needed for query/target/update
	return_dict=True,
	)

	# --- Associative Loss Calculation (if surprise update) ---
	if update_rule == 'surprise' and use_memory and retrieved_mem is not None:
	# Target: Final hidden state corresponding to the last input token before generation
	# The index needs to account for the prepended memory.
	# If mem_len=1, the target state corresponds to index -1 in the output sequence
	target_repr = outputs.hidden_states[-1][:, -1, :].to(self.memory_buffer.dtype) # (B, C)

	# pred_repr comes from retrieve_memory, should already match buffer dtype
	pred_repr = retrieved_mem.squeeze(1) # (B, C)

	# --- DEBUG PRINTS ---
	print(f"\n--- Surprise Update Debug (Step {step}) ---")
	print(f" memory_buffer requires_grad: {self.memory_buffer.requires_grad}")
	print(f" retrieved_mem requires_grad: {retrieved_mem.requires_grad if retrieved_mem is not None else 'N/A'}")
	print(f" pred_repr requires_grad: {pred_repr.requires_grad if pred_repr is not None else 'N/A'}")
	print(f" target_repr requires_grad: {target_repr.requires_grad}") # Should be False due to .detach() below
	# --- END DEBUG ---

	assoc_loss = F.mse_loss(pred_repr, target_repr.detach()) # TARGET IS DETACHED
	print(f" assoc_loss: {assoc_loss.item():.4f}, requires_grad: {assoc_loss.requires_grad}")


	if self.memory_buffer.grad is not None:
	print(" Zeroing existing memory_buffer gradient.")
	self.memory_buffer.grad.zero_()

	if assoc_loss.requires_grad:
	print(" Calling assoc_loss.backward()")
	assoc_loss.backward() # Compute grads for memory_buffer
	print(f" memory_buffer.grad is None after backward: {self.memory_buffer.grad is None}")
	if self.memory_buffer.grad is not None:
	print(f" memory_buffer.grad norm: {torch.norm(self.memory_buffer.grad).item():.4f}")
	self.apply_surprise_update() # Apply update and zero grad
	else:
	print(" ERROR: assoc_loss does not require grad! Skipping backward and update.")
	print("--- End Surprise Update Debug ---")


	# --- Standard Generation Logic ---
	# Get logits for the very last position in the output sequence (corresponds to the token we just fed in)
	next_token_logits = outputs.logits[:, -1, :] # (B, V)

	# Update KV cache for next step
	if use_kv_cache_this_step:
	# The past_key_values returned by Llama should account for the memory prepended in this step
	past_key_values = outputs.past_key_values


	# Sampling (same as before)
	if repetition_penalty != 1.0:
	# Simple loop for now:
	for i in range(bsz):
	# Penalize tokens in the generated sequence (excluding prompt if needed)
	# Use generated_ids which tracks the full sequence
	for token_id in generated_ids[i]:
	# Avoid penalizing pad token if present
	if token_id != pad_token_id:
	next_token_logits[i, token_id] /= repetition_penalty

	if temperature > 0 and temperature != 1.0:
	next_token_logits = next_token_logits / temperature
	if do_sample and top_p < 1.0:
	# Use Hugging Face's top_p implementation detail
	sorted_logits, sorted_indices = torch.sort(next_token_logits, descending=True)
	cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
	sorted_indices_to_remove = cumulative_probs > top_p
	sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
	sorted_indices_to_remove[..., 0] = 0
	indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
	next_token_logits = next_token_logits.masked_fill(indices_to_remove, float('-inf'))

	if do_sample:
	probs = F.softmax(next_token_logits, dim=-1)
	next_token = torch.multinomial(probs, num_samples=1)
	else:
	next_token = torch.argmax(next_token_logits, dim=-1, keepdim=True)

	# --- Update State ---
	generated_ids = torch.cat([generated_ids, next_token], dim=1)
	current_seq_len += 1
	# Update attention mask for the next iteration by appending 1
	attention_mask = torch.cat([attention_mask, torch.ones((bsz, 1), dtype=attention_mask.dtype, device=device)], dim=1)


	# --- EMA Memory Update ---
	if update_rule == 'ema' and use_memory and outputs.hidden_states is not None:
	# Use hidden state corresponding to the newly generated token position (index -1)
	# Cast state to buffer dtype before update
	new_context_state = outputs.hidden_states[-1][:, -1, :].to(self.memory_buffer.dtype) # (B, C)
	self.update_memory_ema(new_context_state.detach())

	if eos_token_id is not None and (next_token == eos_token_id).all():
	break

	# Restore original requires_grad state
	self.memory_buffer.requires_grad_(original_requires_grad)

	return generated_ids


	# --- Save/Load ---
	# Keep the save_pretrained as is, it saves wrapper specific state.
	def save_pretrained(self, save_directory: Union[str, os.PathLike], **kwargs):
	""" Saves the wrapper's specific state (memory buffer, surprise state). """
	save_directory = Path(save_directory)
	save_directory.mkdir(parents=True, exist_ok=True)

	# Save the base model's config (important for PreTrainedModel compatibility)
	self.config.save_pretrained(save_directory)

	# Save the memory buffer parameter directly
	# Ensure saving in float32 for broader compatibility, can be cast back on load
	# Note: Saving the Parameter itself, not just its .data
	torch.save(self.memory_buffer.float(), save_directory / "memory_buffer.pt")
	# Save the surprise state buffer directly
	torch.save(self.surprise_state.float(), save_directory / "surprise_state.pt")

	print(f"InferenceMemoryWrapper state saved to {save_directory}")
	# Note: Base Llama model weights are assumed to be saved separately or loaded from source.

	# from_pretrained is complex with wrappers. For local testing/handler, load manually.
	# @classmethod
	# def from_pretrained(cls, pretrained_model_name_or_path: Union[str, os.PathLike], **kwargs):
	# raise NotImplementedError(...)