Upload folder using huggingface_hub

README.md

@@ -6,36 +6,60 @@ tags:
   - tptt
   - peft
   - trust_remote_code
+pipeline_tag: text-generation
 base_model: apple/OpenELM-1_1B
 datasets:
 - yahma/alpaca-cleaned
 ---
 
-# Titans-OpenELM-1_1B
+# Titanesque-OpenELM-1_1B
+
+<p align="center">
+    <a href="https://arxiv.org/abs/2506.17671">
+        <img alt="arXiv" src="https://img.shields.io/badge/arXiv-tptt-blueviolet.svg">
+    </a>
+    <a href="https://pypi.org/project/tptt/">
+        <img alt="PyPI" src="https://img.shields.io/pypi/v/tptt?color=orange">
+    </a>
+    <a href="https://github.com/fabienfrfr/tptt/">
+        <img alt="Release" src="https://img.shields.io/github/v/release/fabienfrfr/tptt?color=brightgreen">
+    </a>
+    <a href="https://fabienfrfr.github.io/tptt/">
+        <img alt="Documentation" src="https://img.shields.io/badge/docs-online-blue">
+    </a>
+    <a href="https://huggingface.co/ffurfaro">
+        <img alt="HuggingFace" src="https://img.shields.io/badge/hf-ffurfaro-yellow">
+    </a>
+</p>
 
 Titanesque version of `apple/OpenELM-1_1B` with parallel linearized attention (TPTT 😊) and PEFT.
 
-The model was presented in the paper [TPTT](https://huggingface.co/papers/2506.17671).
+The architecture was presented in the paper [TPTT](https://huggingface.co/papers/2506.17671).
 
-## Model Details
 
-- **Architecture:** TpttModel
-- **Base model:** apple/OpenELM-1_1B
-- **LiZA config:** operator=delta_rule, mag=0.5
-- **LoRA config:** r=8, alpha=16, dropout=0.05
-- **torch_dtype:** bfloat16
+## Model list
 
-## Usage
+Classic model parameter with LiZA injection :
+
+| Subfolder                      | Max Self Attn Length | Mag Weight | Cross Gate | Max Chunk Size | Bidirectional | LoRA | Description                                           |
+|-------------------------------|----------------------|------------|------------|----------------|---------------|------|-------------------------------------------------------|
+| delta_rule  | 8192 (default)       | 0.5        | False      | 64             | False         | Yes  | Parallel linearized attention with delta_rule operator|
+| delta_rule_gelu | 8192 (default) | 0.5        | False      | 64             | False         | Yes  | Non-linear operator with gelu activation              |
+| delta_product    | 8192 (default) | 0.5        | False      | 64             | False         | Yes  | Second order operator with derivative trick              |
+| delta_product_r  | 8192 (default) | 0.5        | False      | 64             | False         | Yes  | Second order operator with rotative trick             |
+| delta_product_c  | 8192 (default) | 0.5        | False      | 64             | False         | Yes  | Second order operator with combined trick             |
 
+## Usage
 
 ```python
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
 model = AutoModelForCausalLM.from_pretrained(
-"ffurfaro/Titans-OpenELM-1_1B",
+"ffurfaro/Titanesque-OpenELM-1_1B",
+subfolder="tptt_subfolder", # see in repo tree
 trust_remote_code=True
 )
-tokenizer = AutoTokenizer.from_pretrained("ffurfaro/Titans-OpenELM-1_1B")
+tokenizer = AutoTokenizer.from_pretrained("ffurfaro/apple/OpenELM-1_1B")
 
 prompt = "Your prompt here"
 inputs = tokenizer(prompt, return_tensors="pt")
@@ -44,26 +68,6 @@ print(tokenizer.decode(outputs, skip_special_tokens=True))
 
 ```
 
-## Training
-
-- **Dataset:** yahma/alpaca-cleaned
-- **Platform:** Kaggle
-- **Hardware:** NVIDIA 2xT4
-- **Batch size:** 3
-- **Epochs:** 5.0
-- **Learning rate (final):** 1.1904761904761906e-06
-- **Loss (final):** 1.3188
-- **Training runtime:** 1651.0658 sec
-- **Samples per second:** 1.514
-- **Steps per second:** 0.254
-- **Total FLOPs:** 5852956262400000.0
-- **Gradient norm (final):** 0.7039350271224976
-
-## Evaluation
-
-- **Metrics:** Training loss only (no eval yet, table soon : PiQA, ARC, Hella, Wino, GSM8K, MMLU)
-- **Results:** Final training loss: 1.3188
-
 
 ## Citation & Contact
 

__init__.py

@@ -1,19 +1,26 @@
-__version__ = "0.1.0"
+"""
+This module implements the TPTT model with linear attention (LiZA) and LoRA support.
+"""
 
-from .configuration_tptt import TpttConfig, generate_model_card
-from .modeling_tptt import (AttentionOperator, LCache, LiZAttention, TpttModel,
-                            get_tptt_model)
-from .pipeline_tptt import TpttPipeline
-from .train_tptt import AdjustMaGWeightCallback
+from .configuration_tptt import (TpttConfig, generate_model_card,
+                                 parse_mode_name)
+from .modeling_tptt import (LCache, LinearAttention, LinearAttentionOp,
+                            LiZAttention, TpttModel, get_tptt_model,
+                            load_tptt_safetensors, save_tptt_safetensors)
+from .train_tptt import LiZACallback, SaveBestModelCallback
 
 __all__ = [
     "TpttConfig",
     "TpttModel",
-    "TpttPipeline",
     "get_tptt_model",
-    "AdjustMaGWeightCallback",
+    "LiZACallback",
+    "SaveBestModelCallback",
     "LCache",
-    "AttentionOperator",
+    "LinearAttentionOp",
     "LiZAttention",
     "generate_model_card",
+    "LinearAttention",
+    "parse_mode_name",
+    "load_tptt_safetensors",
+    "save_tptt_safetensors",
 ]

configuration_tptt.py

@@ -1,23 +1,33 @@
+# pylint: disable=too-many-arguments, too-many-positional-arguments, too-many-instance-attributes, too-many-locals
 """
 Author : Fabien FURFARO
 """
 
+import logging
 import os
 import re
-from typing import List, Optional, Union
+from typing import Any, Dict, List, Optional, Union
+from jinja2 import Environment, FileSystemLoader
 
+import psutil
+import torch
 from transformers import AutoConfig, PretrainedConfig
 
+logger = logging.getLogger(__name__)  # monitoring
+
+# Constants
+BYTES_IN_GB = 1024**3
+
 
 def convert_sets_to_lists(obj):
+    """Convert sets to list for LoRA serialized config"""
     if isinstance(obj, set):
         return list(obj)
-    elif isinstance(obj, dict):
+    if isinstance(obj, dict):
         return {k: convert_sets_to_lists(v) for k, v in obj.items()}
-    elif isinstance(obj, (list, tuple)):
+    if isinstance(obj, (list, tuple)):
         return [convert_sets_to_lists(x) for x in obj]
-    else:
-        return obj
+    return obj
 
 
 class TpttConfig(PretrainedConfig):
@@ -33,17 +43,73 @@ class TpttConfig(PretrainedConfig):
     }
     architectures = ["TpttModel"]
 
+    RECURRENT_MODES = {
+        "delta_rule": {
+            "order": 1,
+            "gate_type": "k",
+            "linear": True,
+            "trick": "derivative",
+        },
+        "delta_rule_v": {
+            "order": 1,
+            "gate_type": "v",
+            "linear": True,
+            "trick": "derivative",
+        },
+        "delta_rule_kv": {
+            "order": 1,
+            "gate_type": "kv",
+            "linear": True,
+            "trick": "derivative",
+        },
+        "delta_rule_gelu": {
+            "order": 1,
+            "gate_type": "k",
+            "linear": False,
+            "trick": "derivative",
+        },
+        "delta_product": {
+            "order": 2,
+            "gate_type": "k",
+            "linear": True,
+            "trick": "derivative",
+        },
+        "delta_product_r": {
+            "order": 2,
+            "gate_type": "k",
+            "linear": True,
+            "trick": "rotative",
+        },
+        "delta_product_c": {
+            "order": 2,
+            "gate_type": "k",
+            "linear": True,
+            "trick": "combined",
+        },
+    }  # Tested modes, see parse_mode_name if you want to add more
+
     def __init__(
         self,
         base_model_config: Optional[Union[dict, PretrainedConfig]] = None,
         base_model_name: str = "meta-llama/Llama-3.2-1B",
+        base_model_subfolder: Optional[str] = None,
         name_or_path: Optional[str] = None,
+        model_task: str = "causal_lm",
         target_modules_names: Optional[List[str]] = None,
         operator_mode: str = "delta_rule",
-        max_self_attn_length: int = 8192,
+        use_linear_checkpoint: Optional[bool] = None,
+        max_self_attn_length: Optional[
+            int
+        ] = None,  # unnecessary if SWA, else, standards 8192
+        base_scale_attn: bool = False,
         mag_weight: float = 0.5,  # if 1.0, use only linear operator
-        max_chunk_size: int = 64,
+        cross_gate: bool = False,  # unlinear mixing strategy
+        max_chunk_size: int = 64,  # 128 if adaptive chunking (longest)
+        linear_precision: Union[str, torch.dtype] = "float32",
         lora_config: Optional[dict] = None,  # only serialized accepted
+        padding_side: Optional[str] = None,  # for tokenizer, default "right"
+        bidirectional: bool = False,  # if True, use bidirectional attention
+        pooling_config: Optional[Dict[str, Any]] = None,
         **kwargs,
     ):
         # If base_model_config is provided, load it and merge with this config
@@ -60,11 +126,16 @@ class TpttConfig(PretrainedConfig):
             setattr(self, k, v)
 
         self.base_model_name = base_model_name
-        self._name_or_path = (
-            name_or_path
-            if name_or_path is not None
-            else "Titans-" + base_model_name.split("/", 1)[1]
-        )
+        self.base_model_subfolder = base_model_subfolder
+        self.model_task = model_task
+
+        if name_or_path is not None:
+            self._name_or_path = name_or_path
+        else:
+            if "/" in base_model_name:
+                self._name_or_path = "Titans-" + base_model_name.split("/", 1)[1]
+            else:
+                self._name_or_path = "Titans-" + base_model_name
 
         self.target_modules_names = target_modules_names or [
             "attn",
@@ -72,9 +143,28 @@ class TpttConfig(PretrainedConfig):
             "attention",
         ]
         self.operator_mode = operator_mode
+
+        # Detect available memory on accelerator device
+        if torch.cuda.is_available():
+            _, total_mem = torch.cuda.mem_get_info()
+        else:
+            total_mem = psutil.virtual_memory().total
+        total_mem_gb = total_mem / BYTES_IN_GB
+
+        self.use_linear_checkpoint = (
+            total_mem_gb < 16
+            if use_linear_checkpoint is None
+            else use_linear_checkpoint
+        )
+
+        self.base_scale_attn = base_scale_attn
         self.mag_weight = mag_weight
+        self.cross_gate = cross_gate
         self.max_chunk_size = max_chunk_size
         self.max_self_attn_length = max_self_attn_length
+        if isinstance(linear_precision, torch.dtype):
+            linear_precision = str(linear_precision).replace("torch.", "")
+        self.linear_precision = linear_precision
 
         self.lora_config = lora_config
         if lora_config is not None:
@@ -82,56 +172,147 @@ class TpttConfig(PretrainedConfig):
                 self.lora_config["peft_type"] = self.lora_config["peft_type"].value
             self.lora_config = convert_sets_to_lists(self.lora_config)
 
+        self.padding_side = padding_side
+        self.bidirectional = bidirectional
+        if self.bidirectional:
+            print("Bidirectional is enabled, need to be uncausal and unpadded.")
+        self.pooling_config = pooling_config
+
         super().__init__(**kwargs)  # flush unconsistend pretrained parameters (?)
         # Copy class attributes to instance for serialization (save dict)
         self.model_type = self.__class__.model_type
         self.auto_map = self.__class__.auto_map
         self.architectures = self.__class__.architectures
+        # Padding side configuration if not set
+        if self.padding_side is None:
+            self.padding_side = "right"
+            logger.info("Warning: padding_side is None, defaulting to 'right'.")
+        # set recurrent configuration from operator mode
+        if operator_mode not in self.__class__.RECURRENT_MODES:
+            self.recurrent_config = parse_mode_name(operator_mode)
+        else:
+            self.recurrent_config = self.__class__.RECURRENT_MODES[operator_mode]
+        logger.info("Using recurrent mode: %s", get_mode_name(**self.recurrent_config))
 
 
 TpttConfig.register_for_auto_class()
 
 
-def extract_template_variables(template):
-    return set(re.findall(r"\{([^{}]+)\}", template))
+def parse_mode_name(name: str) -> dict:
+    """Parse mode to recurrent config"""
+    if name.startswith("delta_product"):
+        parts = name.split("_")
+        # Prefix is always two words: 'delta' and 'product'
+        base_len = 2
+        order = 2
+        gate_type = "k"
+        linear = True
+        trick = "derivative"
 
+        idx = base_len
+        # Check for order (immediately after the prefix)
+        if len(parts) > idx and parts[idx].isdigit():
+            order = int(parts[idx])
+            idx += 1
 
-def generate_model_card(path: str, config, **kwargs):
-    """Generate model card from template and training metadata."""
-    template_path = os.path.join(os.path.dirname(__file__), "model_card_template.md")
-    with open(template_path, "r", encoding="utf-8") as f:
-        template = f.read()
+        remaining = parts[idx:]
+        # Trick (r/c) is always at the far right if present
+        if remaining and remaining[-1] in ("r", "c"):
+            trick = {"r": "rotative", "c": "combined"}[remaining[-1]]
+            remaining = remaining[:-1]
+        # 'gelu' comes just before the trick if present
+        if remaining and remaining[-1] == "gelu":
+            linear = False
+            remaining = remaining[:-1]
+        # If anything remains, it's the gate_type
+        if remaining:
+            gate_type = "_".join(remaining)
+        return {
+            "order": order,
+            "gate_type": gate_type,
+            "linear": linear,
+            "trick": trick,
+        }
 
-    # Flatten config
-    def flatten_config(config):
-        result = {}
-        if hasattr(config, "__dict__"):
-            config = config.__dict__
-        for k, v in config.items():
-            if isinstance(v, dict):
-                for subk, subv in v.items():
-                    result[f"{k}_{subk}"] = subv
-            else:
-                result[k] = v
-        return result
+    # delta_rule[_gate][_gelu]
+    m = re.match(r"^delta_rule(?:_(kv|v|k))?(_gelu)?$", name)
+    if m:
+        return {
+            "order": 1,
+            "gate_type": m.group(1) if m.group(1) else "k",
+            "linear": not bool(m.group(2)),
+            "trick": "derivative",
+        }
+    raise ValueError(f"Unknown mode: {name}")
+
+
+def get_mode_name(
+    order: int = 1, gate_type: str = "k", linear: bool = True, trick: str = "derivative"
+) -> str:
+    """Get recurrent mode name from parameter"""
+    base = (
+        "delta_rule"
+        if order == 1
+        else ("delta_product" if order == 2 else f"delta_product_{order}")
+    )
+    parts = []
+    if gate_type != "k":
+        parts.append(gate_type)
+    if not linear:
+        parts.append("gelu")
+    if order >= 2 and trick != "derivative":
+        parts.append({"rotative": "r", "combined": "c"}.get(trick, trick))
+    return base + (("_" + "_".join(parts)) if parts else "")
 
-    variables = flatten_config(config)
-    variables.update(kwargs)
-    variables["model_id"] = os.path.basename(path)
 
-    # Extract variables from template
-    template_vars = extract_template_variables(template)
+def render_template(template_path: str, variables: dict) -> str:
+    """Load and render a Jinja2 template from any file path."""
+    env = Environment(loader=FileSystemLoader(os.path.dirname(template_path)))
+    template = env.get_template(os.path.basename(template_path))
+    return template.render(**variables)
 
-    # Add default values for missing variables
-    for var in template_vars:
-        if var not in variables:
-            variables[var] = "N/A"
 
-    # Handle list conversion (optional but useful)
-    for k, v in variables.items():
-        if isinstance(v, list):
-            variables[k] = ", ".join(map(str, v))
+def write_model_card(output_path: str, content: str):
+    """Write the generated content into README.md."""
+    os.makedirs(output_path, exist_ok=True)
+    readme_path = os.path.join(output_path, "README.md")
+    with open(readme_path, "w", encoding="utf-8") as f:
+        f.write(content)
+
+
+def generate_model_card(
+    output_path: str,
+    config: Union[dict, object],
+    template: Optional[
+        str
+    ],  # can be "model_card" OR an absolute/relative path to a .md file
+    extra_variables: Optional[Dict] = None,
+):
+    """
+    Generate a README.md file from a Jinja2 template and a configuration.
+
+    - template can be either:
+        * a full path to a template file
+        * a short name (e.g., "model_card") -> will be looked up inside default_templates_dir
+    """
+    if template is None:
+        template = "model_card_template"  # default template name
+    # Locate the template
+    if os.path.exists(template):  # direct file path provided
+        template_path = template
+    else:
+        default_templates_dir = os.path.join(os.path.dirname(__file__), "templates")
+        template_path = os.path.join(default_templates_dir, f"{template}.md")
+
+    if not os.path.exists(template_path):
+        raise FileNotFoundError(f"Template not found: {template_path}")
+
+    variables = {
+        "model_id": os.path.basename(output_path),
+        "config": config,
+    }
+    if extra_variables:
+        variables.update(extra_variables)
 
-    model_card_content = template.format(**variables)
-    with open(os.path.join(path, "README.md"), "w", encoding="utf-8") as f:
-        f.write(model_card_content)
+    content = render_template(template_path, variables)
+    write_model_card(output_path, content)

lora_delta_product_m0.5_constant/README.md

@@ -0,0 +1,96 @@
+---
+language: en
+license: apache-2.0
+library_name: transformers
+tags:
+  - tptt
+  - peft
+  - trust_remote_code
+pipeline_tag: text-generation
+base_model: apple/OpenELM-1_1B
+datasets:
+- yahma/alpaca-cleaned
+---
+
+# lora_delta_product_m0.5_constant
+
+<p align="center">
+    <a href="https://arxiv.org/abs/2506.17671">
+        <img alt="arXiv" src="https://img.shields.io/badge/arXiv-tptt-blueviolet.svg">
+    </a>
+    <a href="https://pypi.org/project/tptt/">
+        <img alt="PyPI" src="https://img.shields.io/pypi/v/tptt?color=orange">
+    </a>
+    <a href="https://github.com/fabienfrfr/tptt/">
+        <img alt="Release" src="https://img.shields.io/github/v/release/fabienfrfr/tptt?color=brightgreen">
+    </a>
+    <a href="https://fabienfrfr.github.io/tptt/">
+        <img alt="Documentation" src="https://img.shields.io/badge/docs-online-blue">
+    </a>
+    <a href="https://huggingface.co/ffurfaro">
+        <img alt="HuggingFace" src="https://img.shields.io/badge/hf-ffurfaro-yellow">
+    </a>
+</p>
+
+Titanesque version of `apple/OpenELM-1_1B` with parallel linearized attention (TPTT 😊) and PEFT.
+
+The architecture was presented in the paper [TPTT](https://huggingface.co/papers/2506.17671).
+
+
+## Model Details
+
+- **Architecture:** ['TpttModel']
+- **Base model:** apple/OpenELM-1_1B
+- **LiZA config:** operator=delta_product, mag=0.5
+- **LoRA config:** r=8, alpha=16, dropout=0.05
+- **torch_dtype:** 
+
+## Usage
+
+
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+model = AutoModelForCausalLM.from_pretrained(
+"ffurfaro/lora_delta_product_m0.5_constant",
+trust_remote_code=True
+)
+tokenizer = AutoTokenizer.from_pretrained("ffurfaro/apple/OpenELM-1_1B")
+
+prompt = "Your prompt here"
+inputs = tokenizer(prompt, return_tensors="pt")
+outputs = model.generate(**inputs, max_new_tokens=100)
+print(tokenizer.decode(outputs, skip_special_tokens=True))
+
+```
+
+> [!IMPORTANT]
+> You must specify the `subfolder` if the repo contains multiple models, see the homepage for details.
+
+## Training
+
+- **Dataset:** yahma/alpaca-cleaned
+- **Platform:** Kaggle
+- **Hardware:** 2xT4
+- **Batch size:** 2
+- **Epochs:** 1.0
+- **Learning rate (final):** N/A
+- **Loss (final):** 2.0623671754250386
+- **Training runtime:** 2735.583 sec
+- **Samples per second:** 0.946
+- **Steps per second:** 0.237
+- **Total FLOPs:** 2018060372803584.0
+- **Gradient norm (final):** N/A
+
+## Evaluation
+
+- **Metrics:** Training loss only (no eval yet, table soon : PiQA, ARC, Hella, Wino, GSM8K, MMLU)
+- **Results:** Final training loss: 2.0623671754250386
+
+
+## Citation & Contact
+
+If you use TPTT in your academic work, please cite [Furfaro](https://huggingface.co/ffurfaro). For questions or support, please open an issue on the [GitHub repository](https://github.com/fabienfrfr/tptt) or contact the maintainer.
+
+
+---

lora_delta_product_m0.5_constant/adapter_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:42b7ba7ce3507b8bcab8d45188a94e47e1f84bbe215e62aac9101f182faea530
+size 3919584

lora_delta_product_m0.5_constant/config.json

@@ -0,0 +1,182 @@
+{
+  "activation_fn_name": "swish",
+  "architectures": [
+    "TpttModel"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_tptt.TpttConfig",
+    "AutoModelForCausalLM": "modeling_tptt.TpttModel"
+  },
+  "base_model_name": "apple/OpenELM-1_1B",
+  "base_model_subfolder": null,
+  "base_scale_attn": null,
+  "bidirectional": false,
+  "cross_gate": false,
+  "ffn_dim_divisor": 256,
+  "ffn_multipliers": [
+    0.5,
+    0.63,
+    0.76,
+    0.89,
+    1.02,
+    1.15,
+    1.28,
+    1.41,
+    1.54,
+    1.67,
+    1.8,
+    1.93,
+    2.06,
+    2.19,
+    2.31,
+    2.44,
+    2.57,
+    2.7,
+    2.83,
+    2.96,
+    3.09,
+    3.22,
+    3.35,
+    3.48,
+    3.61,
+    3.74,
+    3.87,
+    4.0
+  ],
+  "ffn_with_glu": true,
+  "head_dim": 64,
+  "initializer_range": 0.02,
+  "linear_precision": "bfloat16",
+  "lora_config": {
+    "alpha_pattern": {},
+    "auto_mapping": null,
+    "base_model_name_or_path": null,
+    "bias": "none",
+    "corda_config": null,
+    "eva_config": null,
+    "exclude_modules": null,
+    "fan_in_fan_out": false,
+    "inference_mode": false,
+    "init_lora_weights": true,
+    "layer_replication": null,
+    "layers_pattern": null,
+    "layers_to_transform": null,
+    "loftq_config": {},
+    "lora_alpha": 16,
+    "lora_bias": false,
+    "lora_dropout": 0.05,
+    "megatron_config": null,
+    "megatron_core": "megatron.core",
+    "modules_to_save": null,
+    "peft_type": "LORA",
+    "r": 8,
+    "rank_pattern": {},
+    "revision": null,
+    "target_modules": [
+      "o_proj",
+      "qkv_proj"
+    ],
+    "task_type": "CAUSAL_LM",
+    "trainable_token_indices": null,
+    "use_dora": false,
+    "use_rslora": false
+  },
+  "mag_weight": 0.5,
+  "max_chunk_size": 64,
+  "max_context_length": 2048,
+  "max_self_attn_length": null,
+  "model_dim": 2048,
+  "model_task": "causal_lm",
+  "model_type": "tptt",
+  "normalization_layer_name": "rms_norm",
+  "normalize_qk_projections": true,
+  "num_gqa_groups": 4,
+  "num_kv_heads": [
+    4,
+    4,
+    4,
+    5,
+    5,
+    5,
+    5,
+    5,
+    5,
+    5,
+    6,
+    6,
+    6,
+    6,
+    6,
+    6,
+    6,
+    6,
+    7,
+    7,
+    7,
+    7,
+    7,
+    7,
+    8,
+    8,
+    8,
+    8
+  ],
+  "num_query_heads": [
+    16,
+    16,
+    16,
+    20,
+    20,
+    20,
+    20,
+    20,
+    20,
+    20,
+    24,
+    24,
+    24,
+    24,
+    24,
+    24,
+    24,
+    24,
+    28,
+    28,
+    28,
+    28,
+    28,
+    28,
+    32,
+    32,
+    32,
+    32
+  ],
+  "num_transformer_layers": 28,
+  "operator_mode": "delta_product",
+  "padding_side": "right",
+  "pooling_config": null,
+  "qkv_multipliers": [
+    0.5,
+    1.0
+  ],
+  "recurrent_config": {
+    "gate_type": "k",
+    "linear": true,
+    "order": 2,
+    "trick": "derivative"
+  },
+  "rope_freq_constant": 10000,
+  "rope_max_length": 4096,
+  "share_input_output_layers": true,
+  "target_modules_names": [
+    "attn",
+    "self_attn",
+    "attention"
+  ],
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.49.0",
+  "trust_remote_code": true,
+  "use_cache": true,
+  "use_linear_checkpoint": false,
+  "vocab_size": 32000
+}

lora_delta_product_m0.5_constant/configuration_tptt.py

@@ -0,0 +1,318 @@
+# pylint: disable=too-many-arguments, too-many-positional-arguments, too-many-instance-attributes, too-many-locals
+"""
+Author : Fabien FURFARO
+"""
+
+import logging
+import os
+import re
+from typing import Any, Dict, List, Optional, Union
+from jinja2 import Environment, FileSystemLoader
+
+import psutil
+import torch
+from transformers import AutoConfig, PretrainedConfig
+
+logger = logging.getLogger(__name__)  # monitoring
+
+# Constants
+BYTES_IN_GB = 1024**3
+
+
+def convert_sets_to_lists(obj):
+    """Convert sets to list for LoRA serialized config"""
+    if isinstance(obj, set):
+        return list(obj)
+    if isinstance(obj, dict):
+        return {k: convert_sets_to_lists(v) for k, v in obj.items()}
+    if isinstance(obj, (list, tuple)):
+        return [convert_sets_to_lists(x) for x in obj]
+    return obj
+
+
+class TpttConfig(PretrainedConfig):
+    """
+    Configuration class for the TPTT model.
+    This class merges the backbone config (e.g., Llama) with custom TPTT parameters,
+    """
+
+    model_type = "tptt"
+    auto_map = {
+        "AutoModelForCausalLM": "modeling_tptt.TpttModel",
+        "AutoConfig": "configuration_tptt.TpttConfig",
+    }
+    architectures = ["TpttModel"]
+
+    RECURRENT_MODES = {
+        "delta_rule": {
+            "order": 1,
+            "gate_type": "k",
+            "linear": True,
+            "trick": "derivative",
+        },
+        "delta_rule_v": {
+            "order": 1,
+            "gate_type": "v",
+            "linear": True,
+            "trick": "derivative",
+        },
+        "delta_rule_kv": {
+            "order": 1,
+            "gate_type": "kv",
+            "linear": True,
+            "trick": "derivative",
+        },
+        "delta_rule_gelu": {
+            "order": 1,
+            "gate_type": "k",
+            "linear": False,
+            "trick": "derivative",
+        },
+        "delta_product": {
+            "order": 2,
+            "gate_type": "k",
+            "linear": True,
+            "trick": "derivative",
+        },
+        "delta_product_r": {
+            "order": 2,
+            "gate_type": "k",
+            "linear": True,
+            "trick": "rotative",
+        },
+        "delta_product_c": {
+            "order": 2,
+            "gate_type": "k",
+            "linear": True,
+            "trick": "combined",
+        },
+    }  # Tested modes, see parse_mode_name if you want to add more
+
+    def __init__(
+        self,
+        base_model_config: Optional[Union[dict, PretrainedConfig]] = None,
+        base_model_name: str = "meta-llama/Llama-3.2-1B",
+        base_model_subfolder: Optional[str] = None,
+        name_or_path: Optional[str] = None,
+        model_task: str = "causal_lm",
+        target_modules_names: Optional[List[str]] = None,
+        operator_mode: str = "delta_rule",
+        use_linear_checkpoint: Optional[bool] = None,
+        max_self_attn_length: Optional[
+            int
+        ] = None,  # unnecessary if SWA, else, standards 8192
+        base_scale_attn: bool = False,
+        mag_weight: float = 0.5,  # if 1.0, use only linear operator
+        cross_gate: bool = False,  # unlinear mixing strategy
+        max_chunk_size: int = 64,  # 128 if adaptive chunking (longest)
+        linear_precision: Union[str, torch.dtype] = "float32",
+        lora_config: Optional[dict] = None,  # only serialized accepted
+        padding_side: Optional[str] = None,  # for tokenizer, default "right"
+        bidirectional: bool = False,  # if True, use bidirectional attention
+        pooling_config: Optional[Dict[str, Any]] = None,
+        **kwargs,
+    ):
+        # If base_model_config is provided, load it and merge with this config
+        if base_model_config is not None:
+            if isinstance(base_model_config, PretrainedConfig):
+                base_model_config = base_model_config.to_dict()
+        else:
+            # Load config from Hugging Face Hub or a local path
+            base_model_config = AutoConfig.from_pretrained(
+                base_model_name, **kwargs
+            ).to_dict()
+        # Merge all backbone fields into this config
+        for k, v in base_model_config.items():
+            setattr(self, k, v)
+
+        self.base_model_name = base_model_name
+        self.base_model_subfolder = base_model_subfolder
+        self.model_task = model_task
+
+        if name_or_path is not None:
+            self._name_or_path = name_or_path
+        else:
+            if "/" in base_model_name:
+                self._name_or_path = "Titans-" + base_model_name.split("/", 1)[1]
+            else:
+                self._name_or_path = "Titans-" + base_model_name
+
+        self.target_modules_names = target_modules_names or [
+            "attn",
+            "self_attn",
+            "attention",
+        ]
+        self.operator_mode = operator_mode
+
+        # Detect available memory on accelerator device
+        if torch.cuda.is_available():
+            _, total_mem = torch.cuda.mem_get_info()
+        else:
+            total_mem = psutil.virtual_memory().total
+        total_mem_gb = total_mem / BYTES_IN_GB
+
+        self.use_linear_checkpoint = (
+            total_mem_gb < 16
+            if use_linear_checkpoint is None
+            else use_linear_checkpoint
+        )
+
+        self.base_scale_attn = base_scale_attn
+        self.mag_weight = mag_weight
+        self.cross_gate = cross_gate
+        self.max_chunk_size = max_chunk_size
+        self.max_self_attn_length = max_self_attn_length
+        if isinstance(linear_precision, torch.dtype):
+            linear_precision = str(linear_precision).replace("torch.", "")
+        self.linear_precision = linear_precision
+
+        self.lora_config = lora_config
+        if lora_config is not None:
+            if hasattr(self.lora_config.get("peft_type"), "value"):
+                self.lora_config["peft_type"] = self.lora_config["peft_type"].value
+            self.lora_config = convert_sets_to_lists(self.lora_config)
+
+        self.padding_side = padding_side
+        self.bidirectional = bidirectional
+        if self.bidirectional:
+            print("Bidirectional is enabled, need to be uncausal and unpadded.")
+        self.pooling_config = pooling_config
+
+        super().__init__(**kwargs)  # flush unconsistend pretrained parameters (?)
+        # Copy class attributes to instance for serialization (save dict)
+        self.model_type = self.__class__.model_type
+        self.auto_map = self.__class__.auto_map
+        self.architectures = self.__class__.architectures
+        # Padding side configuration if not set
+        if self.padding_side is None:
+            self.padding_side = "right"
+            logger.info("Warning: padding_side is None, defaulting to 'right'.")
+        # set recurrent configuration from operator mode
+        if operator_mode not in self.__class__.RECURRENT_MODES:
+            self.recurrent_config = parse_mode_name(operator_mode)
+        else:
+            self.recurrent_config = self.__class__.RECURRENT_MODES[operator_mode]
+        logger.info("Using recurrent mode: %s", get_mode_name(**self.recurrent_config))
+
+
+TpttConfig.register_for_auto_class()
+
+
+def parse_mode_name(name: str) -> dict:
+    """Parse mode to recurrent config"""
+    if name.startswith("delta_product"):
+        parts = name.split("_")
+        # Prefix is always two words: 'delta' and 'product'
+        base_len = 2
+        order = 2
+        gate_type = "k"
+        linear = True
+        trick = "derivative"
+
+        idx = base_len
+        # Check for order (immediately after the prefix)
+        if len(parts) > idx and parts[idx].isdigit():
+            order = int(parts[idx])
+            idx += 1
+
+        remaining = parts[idx:]
+        # Trick (r/c) is always at the far right if present
+        if remaining and remaining[-1] in ("r", "c"):
+            trick = {"r": "rotative", "c": "combined"}[remaining[-1]]
+            remaining = remaining[:-1]
+        # 'gelu' comes just before the trick if present
+        if remaining and remaining[-1] == "gelu":
+            linear = False
+            remaining = remaining[:-1]
+        # If anything remains, it's the gate_type
+        if remaining:
+            gate_type = "_".join(remaining)
+        return {
+            "order": order,
+            "gate_type": gate_type,
+            "linear": linear,
+            "trick": trick,
+        }
+
+    # delta_rule[_gate][_gelu]
+    m = re.match(r"^delta_rule(?:_(kv|v|k))?(_gelu)?$", name)
+    if m:
+        return {
+            "order": 1,
+            "gate_type": m.group(1) if m.group(1) else "k",
+            "linear": not bool(m.group(2)),
+            "trick": "derivative",
+        }
+    raise ValueError(f"Unknown mode: {name}")
+
+
+def get_mode_name(
+    order: int = 1, gate_type: str = "k", linear: bool = True, trick: str = "derivative"
+) -> str:
+    """Get recurrent mode name from parameter"""
+    base = (
+        "delta_rule"
+        if order == 1
+        else ("delta_product" if order == 2 else f"delta_product_{order}")
+    )
+    parts = []
+    if gate_type != "k":
+        parts.append(gate_type)
+    if not linear:
+        parts.append("gelu")
+    if order >= 2 and trick != "derivative":
+        parts.append({"rotative": "r", "combined": "c"}.get(trick, trick))
+    return base + (("_" + "_".join(parts)) if parts else "")
+
+
+def render_template(template_path: str, variables: dict) -> str:
+    """Load and render a Jinja2 template from any file path."""
+    env = Environment(loader=FileSystemLoader(os.path.dirname(template_path)))
+    template = env.get_template(os.path.basename(template_path))
+    return template.render(**variables)
+
+
+def write_model_card(output_path: str, content: str):
+    """Write the generated content into README.md."""
+    os.makedirs(output_path, exist_ok=True)
+    readme_path = os.path.join(output_path, "README.md")
+    with open(readme_path, "w", encoding="utf-8") as f:
+        f.write(content)
+
+
+def generate_model_card(
+    output_path: str,
+    config: Union[dict, object],
+    template: Optional[
+        str
+    ],  # can be "model_card" OR an absolute/relative path to a .md file
+    extra_variables: Optional[Dict] = None,
+):
+    """
+    Generate a README.md file from a Jinja2 template and a configuration.
+
+    - template can be either:
+        * a full path to a template file
+        * a short name (e.g., "model_card") -> will be looked up inside default_templates_dir
+    """
+    if template is None:
+        template = "model_card_template"  # default template name
+    # Locate the template
+    if os.path.exists(template):  # direct file path provided
+        template_path = template
+    else:
+        default_templates_dir = os.path.join(os.path.dirname(__file__), "templates")
+        template_path = os.path.join(default_templates_dir, f"{template}.md")
+
+    if not os.path.exists(template_path):
+        raise FileNotFoundError(f"Template not found: {template_path}")
+
+    variables = {
+        "model_id": os.path.basename(output_path),
+        "config": config,
+    }
+    if extra_variables:
+        variables.update(extra_variables)
+
+    content = render_template(template_path, variables)
+    write_model_card(output_path, content)

lora_delta_product_m0.5_constant/generation_config.json

@@ -0,0 +1,4 @@
+{
+  "_from_model_config": true,
+  "transformers_version": "4.49.0"
+}

lora_delta_product_m0.5_constant/modeling_tptt.py

@@ -0,0 +1,1504 @@
+# pylint: disable=too-many-lines, too-many-arguments, too-many-positional-arguments, too-many-instance-attributes, too-many-locals
+
+"""
+This module implements the TPTT model with linear attention (LiZA) and LoRA support.
+Author : Fabien FURFARO
+TPTT : Transforming Pretrained Transformers into Titans (https://arxiv.org/abs/2506.17671)
+"""
+
+import logging
+import math
+import os
+from pathlib import Path
+import re
+import shutil
+from functools import partial
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import torch
+import torch.nn.functional as F
+from einops import rearrange
+from huggingface_hub import hf_hub_download, list_repo_files
+from peft import LoraConfig, PeftModel, get_peft_model
+from safetensors import safe_open
+from safetensors.torch import save_file
+from torch import nn
+from torch.utils.checkpoint import checkpoint
+from transformers import (
+    AutoConfig,
+    AutoModel,
+    AutoModelForCausalLM,
+    DynamicCache,
+    PreTrainedModel,
+)
+from transformers.configuration_utils import PretrainedConfig
+
+from .configuration_tptt import TpttConfig
+
+logger = logging.getLogger(__name__)  # monitoring
+
+
+class LCache:
+    """Cache for storing intermediate states of linear attention layers."""
+
+    def __init__(self):
+        """Stores per-layer intermediate states: {layer_idx: state_dict}"""
+        self.inputs_states: Dict[int, Dict[str, torch.Tensor]] = (
+            {}
+        )  # recurrent states and qkv buffers
+
+    def __getitem__(self, layer_idx: int) -> Optional[Dict[str, torch.Tensor]]:
+        """Retrieve cached state for a given layer, or None if not present"""
+        return self.inputs_states.get(layer_idx, None)
+
+    def update(self, layer_idx: int, **kwargs):
+        """Detach all tensors to avoid retaining computation graphs"""
+        detached_kwargs = {
+            k: v.detach() if isinstance(v, torch.Tensor) else v
+            for k, v in kwargs.items()
+        }
+        # Update or create the state for the specified layer
+        if layer_idx in self.inputs_states:
+            self.inputs_states[layer_idx].update(detached_kwargs)
+        else:
+            self.inputs_states[layer_idx] = detached_kwargs
+
+    def reset(self):
+        """Clear all cached states and reset the token counter"""
+        self.inputs_states.clear()
+
+
+class CausalAvgPool1d(nn.Module):
+    """Causal sliding window average (uniform, no shape loss along sequence)"""
+
+    def __init__(
+        self, output_size: int, offsets: tuple[int] = (0, 1, 2), mode: str = "replicate"
+    ):
+        super().__init__()
+        self.offsets = offsets
+        self.mode = mode
+        self.pool = nn.AdaptiveAvgPool1d(output_size=output_size)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """x: [B, S, F] → [B, S, F → output_size]"""
+        x_ = x.transpose(1, 2)  # [B, F, S]
+        idxs = torch.tensor(self.offsets, device=x.device)
+        ksize = idxs.max() - idxs.min() + 1
+        w = torch.zeros(ksize, device=x.device, dtype=x.dtype)
+        w[idxs - idxs.min()] = 1 / len(self.offsets)  # Always uniform weights
+        kernel = w.repeat(x_.shape[1], 1).reshape(x_.shape[1], 1, ksize)
+        pad_left = -idxs.min().item()
+        pad_right = (ksize - 1) - pad_left
+        x_pad = F.pad(x_, (pad_left, pad_right), mode=self.mode)
+        y = F.conv1d(x_pad, kernel, groups=x_.shape[1])  # pylint: disable=not-callable
+        return self.pool(y.transpose(1, 2))  # [B, S, F → output_size]
+
+
+class LinearAttention(nn.Module):
+    """
+    Linear multi-head attention layer: [B, S, D] -> [B, S, D]
+    Projections + gating + efficient linear attention mechanism (TPTT compatible).
+    """
+
+    def __init__(
+        self,
+        hidden_dim: int,
+        num_heads: int,
+        head_dim: Optional[int] = None,
+        num_key_value_heads: Optional[int] = None,
+        num_key_value_groups: Optional[int] = None,
+        bias: bool = True,
+        dropout: Optional[float] = None,
+        linear_precision: torch.dtype = torch.float32,
+        padding_side: str = "right",
+        shared_attn: bool = False,  # shared attention
+        layer_idx: int = 0,
+        operator_mode: str = "delta_rule",
+        use_linear_checkpoint: bool = False,
+        recurrent_config: Optional[Dict[str, Any]] = None,
+        linear_cache: Optional[LCache] = None,
+        max_chunk_size: int = 64,
+        bidirectional: bool = False,  # not used if causal
+        pooling_config: Optional[Dict[str, Any]] = None,
+    ):
+        super().__init__()
+        if pooling_config is None:
+            pooling_config = {
+                "offsets": (0, 1, 2),
+                "mode": "replicate",
+            }
+        self.hidden_dim = hidden_dim
+        self.num_heads = num_heads
+        self.head_dim = head_dim or hidden_dim // num_heads
+        self.num_key_value_heads = num_key_value_heads or num_heads
+        self.num_key_value_groups = num_key_value_groups or (
+            num_heads // (num_key_value_heads or num_heads)
+        )
+        self.scaling = self.head_dim**-0.5
+        self.linear_precision = linear_precision
+        self.padding_side = padding_side
+
+        self.shared_attn = shared_attn
+
+        if not shared_attn:
+            self.q_proj = nn.Linear(hidden_dim, num_heads * self.head_dim, bias=bias)
+            self.k_proj = nn.Linear(
+                hidden_dim, self.num_key_value_heads * self.head_dim, bias=bias
+            )
+            self.v_proj = nn.Linear(
+                hidden_dim, self.num_key_value_heads * self.head_dim, bias=bias
+            )
+            self.out_proj = nn.Linear(num_heads * self.head_dim, hidden_dim, bias=bias)
+
+        self.dropout = nn.Dropout(dropout) if dropout is not None else None
+
+        self.linear_operator = LinearAttentionOp(
+            layer_idx=layer_idx,
+            operator_mode=operator_mode,
+            use_linear_checkpoint=use_linear_checkpoint,
+            recurrent_config=recurrent_config,
+            max_chunk_size=max_chunk_size,
+            linear_cache=linear_cache,
+            linear_precision=linear_precision,
+        )
+        self.bidirectional = bidirectional
+        # Causal average pooling for gating
+        self.pooling_config = pooling_config
+        self.pool_g = CausalAvgPool1d(
+            output_size=self.head_dim * self.num_key_value_heads, **pooling_config
+        )
+
+    def forward(
+        self,
+        x: Union[List[torch.Tensor], torch.Tensor],
+        attn_mask: Optional[torch.Tensor] = None,
+        out_proj: Optional[nn.Module] = None,
+        **kwargs: Any,
+    ) -> torch.Tensor:
+        """
+        Forward pass for linear attention. Input shape: [B, S, D], output [B, S, D].
+        """
+
+        if not self.shared_attn:
+            hidden_states = x[0] if isinstance(x, (list, tuple)) else x
+            # Projections
+            q = self.q_proj(hidden_states)
+            k = self.k_proj(hidden_states)
+            v = self.v_proj(hidden_states)
+            out_proj = self.out_proj
+        else:
+            # Shared attention <=> no projections here
+            q, k, v = x[0], x[1], x[2]
+            out_proj = self.out_proj if out_proj is None else out_proj
+
+        # get dtype and device
+        final_dtype, final_device = q.dtype, q.device
+        # Masking if needed
+        if attn_mask is not None:
+            v = apply_linear_attention_mask(attn_mask, v, self.padding_side)
+
+        # Forget and Write Gating for linear attn (abusive term)
+        f_g, w_g = self.pool_g(k), self.pool_g(v)
+
+        # Reshape for multi-head
+        q = rearrange(q, "b n (h d) -> b h n d", h=self.num_heads)
+        k = rearrange(k, "b n (h d) -> b h n d", h=self.num_key_value_heads)
+        v = rearrange(v, "b n (h d) -> b h n d", h=self.num_key_value_heads)
+
+        f_g = rearrange(f_g, "b n (h m) -> b h n m", h=self.num_key_value_heads)
+        w_g = rearrange(w_g, "b n (h m) -> b h n m", h=self.num_key_value_heads)
+
+        # Repeat for GQA
+        k = k.repeat_interleave(self.num_key_value_groups, dim=1)
+        v = v.repeat_interleave(self.num_key_value_groups, dim=1)
+
+        f_g = f_g.repeat_interleave(self.num_key_value_groups, dim=1)
+        w_g = w_g.repeat_interleave(self.num_key_value_groups, dim=1)
+
+        ## DeltaNet-style: Silu activation and normalization
+        q = F.normalize(F.silu(q), p=2, dim=-1, eps=1e-6)
+        k = F.normalize(F.silu(k), p=2, dim=-1, eps=1e-6)
+
+        ## linear stability part
+        v = ensure_stability(v * self.scaling, min_val=-1e4, max_val=1e4)
+
+        # Apply sigmoid to forget and write gates
+        f_g = torch.clamp(torch.sigmoid(f_g), min=1e-6, max=1 - 1e-6)
+        w_g = torch.clamp(torch.sigmoid(w_g), min=1e-6, max=1 - 1e-6)
+
+        # Convert to linear_precision (float32) for numerical stability and get model dtype
+        q, k, v, f_g, w_g = (
+            x.to(self.linear_precision).contiguous() for x in (q, k, v, f_g, w_g)
+        )
+        g = (f_g, w_g)
+
+        # Linear Attention Core, output: [B, H, S, d]
+        if self.bidirectional:  # Work only with uncausal attention
+            # Forward direction
+            out_forward = self.linear_operator(q, k, v, g, **kwargs)
+            # Backward direction: flip the input sequence on the time dimension (dim=2)
+            kwargs_bwd = kwargs.copy()
+            kwargs_bwd["use_cache"] = False
+            out_backward = self.linear_operator(
+                torch.flip(q, dims=[2]),
+                torch.flip(k, dims=[2]),
+                torch.flip(v, dims=[2]),
+                tuple(torch.flip(t, dims=[2]) for t in g),
+                **kwargs_bwd,
+            )
+            # Flip the output back to restore proper order
+            out_backward = torch.flip(out_backward, dims=[2])
+            # Fusion: here, simple addition
+            out = out_forward + out_backward
+        else:
+            out = self.linear_operator(q, k, v, g, **kwargs)
+
+        # Merge heads and project: [B, H, S, d] -> [B, S, H*d] -> Out proj
+        out = rearrange(out, "b h s d -> b s (h d)")
+        # Normalize output (RMS norm). Note: bidirectional compatibility
+        out = out / out.pow(2).mean(dim=-1, keepdim=True).add(1e-6).sqrt()
+        # Ensure dtype and device consistency
+        out = out.to(dtype=final_dtype, device=final_device)
+        # Apply output projection
+        out = out_proj(out)  # [B, S, D]
+        out = ensure_stability(out, min_val=-1e4, max_val=1e4)
+        # Apply dropout if specified
+        if self.dropout is not None:
+            out = self.dropout(out)
+        return out
+
+
+class LiZAttention(nn.Module):
+    """LiZA Linear Attention module, mixing linear and vanilla attention."""
+
+    def __init__(
+        self,
+        base_attn: nn.Module,
+        layer_idx: int,
+        base_config: PretrainedConfig,  # Backbone Config
+        linear_cache: Optional[LCache] = None,
+        operator_mode: str = "delta_rule",
+        use_linear_checkpoint: bool = False,
+        recurrent_config: Optional[Dict[str, Any]] = None,
+        max_self_attn_length: Optional[int] = None,  # unnecessary
+        base_scale_attn: bool = False,
+        mag_weight: float = 0.5,
+        cross_gate: bool = False,
+        max_chunk_size: int = 64,
+        linear_precision: Union[str, torch.dtype] = "float32",
+        padding_side: str = "right",  # for tokenizer
+        disable_linear_attn: bool = False,
+        bidirectional: bool = False,  # if True, use bidirectional attention
+        pooling_config: Optional[Dict[str, Any]] = None,
+    ):
+        super().__init__()
+        if isinstance(linear_precision, str):
+            linear_precision = getattr(torch, linear_precision)
+        self.linear_precision = linear_precision
+        self.base_attn: nn.Module = base_attn
+        self.base_config = base_config
+        self.layer_idx = layer_idx
+        self.max_self_attn_length = max_self_attn_length
+        self.base_scale_attn = base_scale_attn
+        self.mag_weight = mag_weight
+        self.cross_gate = cross_gate
+        self.max_chunk_size = max_chunk_size
+        self.linear_precision = linear_precision
+        self.padding_side = padding_side
+        self.disable_linear_attn = disable_linear_attn
+
+        (
+            self.num_heads,
+            self.head_dim,
+            self.num_key_value_heads,
+            self.num_key_value_groups,
+            self.hidden_dim,
+        ) = self._get_attention_parameters(base_attn, base_config)
+        self.scaling = self.head_dim**-0.5
+
+        self.linear_attn = LinearAttention(
+            layer_idx=layer_idx,
+            shared_attn=True,
+            operator_mode=operator_mode,
+            use_linear_checkpoint=use_linear_checkpoint,
+            recurrent_config=recurrent_config,
+            hidden_dim=self.hidden_dim,
+            num_heads=self.num_heads,
+            head_dim=self.head_dim,
+            num_key_value_heads=self.num_key_value_heads,
+            num_key_value_groups=self.num_key_value_groups,
+            linear_precision=linear_precision,
+            linear_cache=linear_cache,
+            max_chunk_size=max_chunk_size,
+            padding_side=padding_side,
+            bidirectional=bidirectional,
+            pooling_config=pooling_config,
+        )
+
+    def _get_attention_parameters(
+        self, base_attn: nn.Module, base_config: PretrainedConfig
+    ) -> Tuple[Optional[int], Optional[int], Optional[int], Optional[int]]:
+        """Retrieve the attention parameters from the base attention module."""
+        # first order base attention module and second order config
+        num_heads = (
+            getattr(base_attn, "num_heads", None)
+            or getattr(base_attn, "num_q_heads", None)
+            or getattr(base_config, "num_heads", None)
+            or getattr(base_config, "num_attention_heads", None)
+        )
+        head_dim = (
+            getattr(base_attn, "head_dim", None)
+            or getattr(base_attn, "attention_head_size", None)
+            or getattr(base_config, "head_dim", None)
+            or (
+                getattr(base_config, "hidden_size", None) // num_heads
+                if num_heads and getattr(base_config, "hidden_size", None)
+                else None
+            )
+        )
+        num_key_value_heads = (
+            getattr(base_attn, "num_kv_heads", None)
+            or getattr(base_attn, "num_k_heads", None)
+            or getattr(base_config, "num_key_value_heads", None)
+            or num_heads  # fallback
+        )
+        num_key_value_groups = getattr(base_attn, "num_key_value_groups", None) or (
+            num_heads // num_key_value_heads if num_heads and num_key_value_heads else 1
+        )
+        hidden_dim = getattr(base_config, "hidden_size", None) or head_dim * num_heads
+        return (
+            num_heads,
+            head_dim,
+            num_key_value_heads,
+            num_key_value_groups,
+            hidden_dim,
+        )
+
+    def _apply_shared_projections(
+        self, hidden_states: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, nn.Module]:
+        base_attn = self.base_attn
+        if hasattr(base_attn, "q_proj"):
+            # LLama, OLMO and Mistral style
+            q = base_attn.q_proj(hidden_states)
+            k = base_attn.k_proj(hidden_states)
+            v = base_attn.v_proj(hidden_states)
+            out_proj = base_attn.o_proj
+        elif hasattr(base_attn, "qkv_proj"):
+            # OpenELM and GPT-Neo style : QKV fused, split on the last dimension
+            qkv = base_attn.qkv_proj(hidden_states)
+            q, k, v = split_qkv(base_attn, qkv)
+            out_proj = base_attn.out_proj
+        elif hasattr(base_attn, "c_attn") and hasattr(base_attn, "c_proj"):
+            # GPT-2 style
+            qkv = base_attn.c_attn(hidden_states)
+            q, k, v = qkv.chunk(3, dim=-1)
+            out_proj = base_attn.c_proj
+        elif all(hasattr(base_attn, n) for n in ["query", "key", "value"]):
+            # BERT - ViT
+            q = base_attn.query(hidden_states)
+            k = base_attn.key(hidden_states)
+            v = base_attn.value(hidden_states)
+            out_proj = getattr(base_attn, "dense", None)  # ou output.dense
+        else:
+            raise ValueError("Unsupported attention module: cannot find projections.")
+        # Ensure stability
+        q = ensure_stability(q, min_val=-1e4, max_val=1e4)
+        k = ensure_stability(k, min_val=-1e4, max_val=1e4)
+        v = ensure_stability(v, min_val=-1e4, max_val=1e4)
+        return q, k, v, out_proj
+
+    def _process_self_attn(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor],
+        kwargs,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[DynamicCache], int]:
+        """Process the self-attention part (with truncation)."""
+        if self.max_self_attn_length:  # Not needed for SWA (nonparam memorize context)
+            hidden_states, attention_mask = truncate_attention_mask(
+                hidden_states, attention_mask, self.max_self_attn_length
+            )
+
+            if kwargs.get("position_embeddings", None) is not None:
+                cos, sin = kwargs["position_embeddings"]
+                cos = cos[:, -self.max_self_attn_length :]
+                sin = sin[:, -self.max_self_attn_length :]
+                kwargs["position_embeddings"] = (cos, sin)
+
+            if isinstance(kwargs.get("past_key_value", None), DynamicCache):
+                # cache management
+                if (
+                    len(kwargs["past_key_value"]) > self.layer_idx
+                    and self.layer_idx == 0
+                ):
+                    kwargs["past_key_value"].crop(self.max_self_attn_length - 1)
+
+        # Ensure attention mask is of the correct dtype and device
+        if attention_mask is not None:
+            attention_mask = attention_mask.to(
+                dtype=hidden_states.dtype, device=hidden_states.device
+            )
+        # Standard attention (mask and rotation is applied inside)
+        base_attn_outputs = self.base_attn(
+            hidden_states,
+            attention_mask=attention_mask,
+            **kwargs,
+        )
+
+        if isinstance(base_attn_outputs, tuple):
+            if len(base_attn_outputs) == 3:
+                o_base, attn_weights, present_key_value = base_attn_outputs
+                expected_attn_mode = 3
+            elif len(base_attn_outputs) == 2:
+                o_base, attn_weights = base_attn_outputs
+                present_key_value, expected_attn_mode = None, 2
+            else:
+                raise ValueError(
+                    f"Unexpected number of outputs from base_attn: {len(base_attn_outputs)}"
+                )
+        else:
+            o_base = base_attn_outputs
+            attn_weights, present_key_value, expected_attn_mode = None, None, 1
+        # Ensure stability
+        o_base = ensure_stability(o_base, min_val=-1e4, max_val=1e4)
+        return o_base, attn_weights, present_key_value, expected_attn_mode
+
+    def _prepare_attn_mixin(
+        self,
+        o_lin: torch.Tensor,
+        o_base: torch.Tensor,
+        tensor_dtype: torch.dtype,
+        eps: float = 1e-5,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Prepare linear attn for mixing with self attn."""
+        # Force cast typing, shape : [b n (h d)]
+        o_lin = o_lin.to(tensor_dtype)
+        o_base = o_base.to(tensor_dtype)
+        # feature scaling
+        if self.base_scale_attn:
+            scaler = o_base.pow(2).mean(dim=-1, keepdim=True).add(eps).sqrt()
+            o_lin = scaler * o_lin
+        return o_lin, o_base
+
+    def _apply_mag(
+        self, linear_attention: torch.Tensor, softmax_attention: torch.Tensor
+    ) -> torch.Tensor:
+        """Apply the MAG strategy"""
+        # Left-Padding management
+        if linear_attention.shape[1] != softmax_attention.shape[1]:
+            left_trunc = min(linear_attention.shape[1], softmax_attention.shape[1])
+            linear_attention, softmax_attention = (
+                linear_attention[:, -left_trunc:],
+                softmax_attention[:, -left_trunc:],
+            )
+        # NAM : Neural Attention Mixer (with graph forcing)
+        mag_weight = torch.tensor(
+            self.mag_weight,
+            dtype=softmax_attention.dtype,
+            device=softmax_attention.device,
+        )
+        softmax_weighted = (1 - mag_weight) * softmax_attention
+        linear_weighted = mag_weight * linear_attention
+        if self.cross_gate:
+            output_attention = (
+                softmax_weighted + linear_weighted + softmax_weighted * linear_weighted
+            )  # complex cross product (unlinear interaction)
+        else:
+            output_attention = softmax_weighted + linear_weighted  # classic
+
+        if torch.allclose(softmax_weighted, output_attention):
+            logger.info(
+                "[LOG] layer : %s, softmax_weighted and output_attention are close.",
+                self.layer_idx,
+            )
+        # Final output
+        return ensure_stability(output_attention, min_val=-1e4, max_val=1e4)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        """Mix linear and self attention forward"""
+        device = hidden_states.device
+        tensor_dtype = hidden_states.dtype
+        self.base_attn.to(device)
+
+        if self.training:
+            kwargs.pop("past_key_value", None)
+            kwargs["use_cache"] = False
+        elif "use_cache" not in kwargs:
+            kwargs.pop("past_key_value", None)
+            kwargs["use_cache"] = False
+
+        kwargs.pop("position_ids", None)  # obsolete
+
+        # Apply shared projections
+        q, k, v, out_proj = self._apply_shared_projections(hidden_states)
+
+        # Apply linear attention to hidden states
+        o_lin = self.linear_attn(
+            x=[q, k, v], attn_mask=attention_mask, out_proj=out_proj, **kwargs
+        )
+
+        # Process self attn with truncation
+        o_base, attn_weights, present_key_value, expected_attn_mode = (
+            self._process_self_attn(hidden_states, attention_mask, kwargs)
+        )
+
+        # Prepare output mixing
+        o_lin, o_base = self._prepare_attn_mixin(o_lin, o_base, tensor_dtype, eps=1e-5)
+
+        # Apply Memory as Gate in self-attention (with length management and ablation)
+        out = o_base if self.disable_linear_attn else self._apply_mag(o_lin, o_base)
+
+        # Return output following transformer convention
+        if expected_attn_mode == 3:
+            return out, attn_weights, present_key_value
+        if expected_attn_mode == 2:
+            return out, attn_weights
+        return out
+
+
+def load_tptt_safetensors(
+    repo_or_path: str,
+    model: Union[PreTrainedModel, PeftModel],
+    subfolder: Optional[str] = None,
+    token: Optional[str] = None,
+) -> Union[PreTrainedModel, PeftModel]:
+    """Load Tptt safetensor from LoRA/PEFT weights and adapt keys if needed."""
+    # sharding not supported yet (e.g. : -00001-of-00005.safetensors, ...)
+    fname = "adapter_model.safetensors"
+    # subfolder management
+    if subfolder:
+        repo_or_path_norm = os.path.normpath(repo_or_path)
+        subfolder_norm = os.path.normpath(subfolder)
+        if not repo_or_path_norm.endswith(subfolder_norm):
+            fname = f"{subfolder}/{fname}" if subfolder else fname
+    # Find file path
+    if os.path.isdir(repo_or_path):
+        path = os.path.join(repo_or_path, fname)
+        if not os.path.exists(path):
+            return model
+    else:
+        if fname not in list_repo_files(repo_or_path, token=token):
+            return model
+        path = hf_hub_download(repo_or_path, fname, token=token)
+
+    # Load weights from safetensors
+    with safe_open(path, framework="pt") as f:
+        state_dict = {k: f.get_tensor(k) for k in f.keys()}
+
+    # Adapt LoRA/Specific keys if needed (add .default if expected by the model)
+    def adapt_keys(sd, model):
+        model_keys = list(model.state_dict().keys())
+        if any(k.startswith("tptt_model.base_model.") for k in model_keys):
+            prefix = "tptt_model.base_model."
+        elif any(k.startswith("base_model.") for k in model_keys):
+            prefix = "base_model."
+        else:
+            prefix = ""
+
+        has_base_attn = any(".base_attn." in k for k in model_keys)
+
+        def adapt_key(k):
+            k_ = k if k.startswith(prefix) else prefix + k
+            # first, verify and modify base_attn (LiZA)
+            if ".base_attn." in k_ and not has_base_attn:
+                k_ = k_.replace(".base_attn.", ".")
+            # change LoRA if needed
+            if (
+                k_.endswith("lora_A.weight") or k_.endswith("lora_B.weight")
+            ) and k_.replace(".weight", ".default.weight") in model_keys:
+                k_ = k_.replace(".weight", ".default.weight")
+            return k_
+
+        return {adapt_key(k): v for k, v in sd.items()}
+
+    state_dict = adapt_keys(state_dict, model)
+
+    # Cast tensors to the expected dtype of the model parameters
+    model_state_dict = model.state_dict()
+    for k, v in state_dict.items():
+        if k in model_state_dict:
+            expected_dtype = model_state_dict[k].dtype
+            if v.dtype != expected_dtype:
+                state_dict[k] = v.to(expected_dtype)
+
+    logger.info("Input LoRA/Specific keys: %s", [k for k in state_dict.keys()])
+
+    # Load into model
+    missing, unexpected = model.load_state_dict(state_dict, strict=False, assign=True)
+    missing_lora = [k for k in missing if "lora" in k]
+    if missing_lora:
+        logger.warning("Missing keys: %s", missing_lora)
+    if unexpected:
+        logger.warning("Unexpected keys: %s", unexpected)
+    return model
+
+
+def get_tptt_model(  # pylint: disable=too-many-arguments, too-many-positional-arguments
+    model: nn.Module,
+    base_config: PretrainedConfig,  # ou LlamaConfig, MistralConfig, etc.
+    linear_cache: Optional[LCache] = None,
+    liza_attention: nn.Module = LiZAttention,
+    target_modules_names: Optional[list[str]] = None,
+    operator_mode: str = "delta_rule",
+    use_linear_checkpoint: bool = False,
+    recurrent_config: Optional[Dict[str, Any]] = None,
+    base_scale_attn: bool = False,
+    mag_weight: float = 0.5,
+    cross_gate: bool = False,
+    max_chunk_size: int = 64,
+    linear_precision: torch.dtype = torch.float32,
+    max_self_attn_length: Optional[int] = None,  # unnecessary
+    padding_side: str = "right",  # for tokenizer
+    bidirectional: bool = False,  # if True, use bidirectional attention
+    pooling_config: Optional[Dict[str, Any]] = None,
+    **kwargs,  # quickfix unexpected arguments
+) -> Tuple[PreTrainedModel, LCache]:
+    """Replace target modules in a model with LiZAttention."""
+    if target_modules_names is None:
+        target_modules_names = ["attn", "self_attn", "attention"]
+    # Find target modules by suffix (e.g., "attn", "attention")
+    target_modules_names = [
+        name
+        for name, _ in model.named_modules()
+        if any(name.endswith(suffix) for suffix in target_modules_names)
+        and not any(f".{suffix}." in name for suffix in target_modules_names)
+    ]
+    if not target_modules_names:
+        raise ValueError(
+            f"Target modules '{target_modules_names}' not found in the model."
+        )
+    # Prepare recurrent config
+    linear_cache = linear_cache or LCache()
+    # Inject LiZAttention into the model
+    for name, _ in model.named_modules():
+        if name in target_modules_names:
+            parent = model
+            *path, last = name.split(".")
+            for p in path:
+                parent = getattr(parent, p)
+            layer_idx = extract_layer_idx(name)
+            setattr(
+                parent,
+                last,
+                liza_attention(
+                    getattr(parent, last),
+                    layer_idx=layer_idx,
+                    base_config=base_config,
+                    linear_cache=linear_cache,
+                    operator_mode=operator_mode,
+                    use_linear_checkpoint=use_linear_checkpoint,
+                    recurrent_config=recurrent_config,
+                    max_self_attn_length=max_self_attn_length,
+                    base_scale_attn=base_scale_attn,
+                    mag_weight=mag_weight,
+                    cross_gate=cross_gate,
+                    max_chunk_size=max_chunk_size,
+                    linear_precision=linear_precision,
+                    padding_side=padding_side,
+                    bidirectional=bidirectional,
+                    pooling_config=pooling_config,
+                ),
+            )
+    return model, linear_cache
+
+
+def save_tptt_safetensors(model, path: str, name: str = "adapter_model.safetensors"):
+    """Save trainable LoRA/Specific weights and adapting key names"""
+    # 1. Get the full state_dict
+    all_sd = model.state_dict()
+
+    # 2. Identify trainable parameter names (usually only LoRA/PEFT adapters)
+    trainable_keys = [
+        name for name, param in model.named_parameters() if param.requires_grad
+    ]  # Also, you can manually select specific keys in model after load
+
+    # 3. Filter and adapt the keys (Remove custom model encapsulation info)
+    to_save = {
+        k.replace("tptt_model.", "").replace("base_model.", ""): all_sd[k]
+        for k in trainable_keys
+    }
+
+    # 4. Save the filtered adapters to a safetensors file
+    if to_save:
+        os.makedirs(os.path.dirname(path), exist_ok=True)
+        # sharding not supported yet (e.g. : -00001-of-00005.safetensors, ...)
+        save_file(to_save, os.path.join(path, name))
+
+
+class TpttModel(PreTrainedModel):
+    """
+    TPTT model wrapper with linear attention (LiZA) and LoRA support.
+    Handles only architecture and weights.
+    """
+
+    config_class = TpttConfig
+
+    def __init__(
+        self,
+        config: TpttConfig,
+        **kwargs,
+    ):
+        """
+        Initialize TpttModel with a given config and backbone.
+        Injects LiZA attention modules into the backbone.
+        """
+        super().__init__(config, **kwargs)
+        repo_or_path = getattr(config, "_base_path", None) or config._name_or_path
+
+        # 1. Load backbone (with subfolder management) :
+        kwargs_bb = kwargs.copy()
+        if config.base_model_subfolder is not None:
+            kwargs_bb["subfolder"] = config.base_model_subfolder
+        else:
+            kwargs_bb.pop("subfolder", None)
+
+        if config.model_task == "causal_lm":
+            tptt_model = AutoModelForCausalLM.from_pretrained(
+                config.base_model_name, **kwargs_bb
+            )
+        else:
+            tptt_model = AutoModel.from_pretrained(config.base_model_name, **kwargs_bb)
+
+        # 2. Inject LiZA attention
+        self.linear_cache = LCache()
+        tptt_model, self.linear_cache = get_tptt_model(
+            tptt_model, config, self.linear_cache, **config.to_dict()
+        )
+
+        # 3. Apply LoRA/Specific if present and configured
+        if config.lora_config is not None:
+            lora_config_obj = LoraConfig(**config.lora_config)
+            tptt_model = get_peft_model(tptt_model, lora_config_obj)
+        else:
+            # Doesn't work if quantization is applied !
+            tptt_model = set_trainable_parameters(tptt_model)
+
+        # 4. Load safetensor if tptt/peft adaptor in repo
+        if repo_or_path:
+            tptt_model = load_tptt_safetensors(
+                repo_or_path,
+                tptt_model,
+                subfolder=kwargs.get("subfolder", None),
+                token=kwargs.get("token", None),
+            )
+        self.tptt_model = tptt_model
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        **kwargs,
+    ):
+        """Forward pass. All arguments are passed to the underlying base model."""
+        if self.training:
+            kwargs["use_cache"] = False
+            kwargs.pop("num_items_in_batch", None)
+        elif "use_cache" not in kwargs:  # evaluation
+            kwargs.pop("num_items_in_batch", None)
+            kwargs["use_cache"] = False
+        return self.tptt_model(
+            input_ids=input_ids, attention_mask=attention_mask, labels=labels, **kwargs
+        )
+
+    def generate(self, *args, **kwargs):
+        """Delegate the generate call to the backbone model, which supports generation"""
+        return self.tptt_model.generate(*args, **kwargs)
+
+    def save_pretrained(self, path: str, **kwargs):
+        """Save model weights, config, and source code to the given path."""
+        # 0. Save complete tptt config (with or without LoRA)
+        super().save_pretrained(path, **kwargs)  # pylint: disable=no-member
+        self._adjust_save_strategy(path, **kwargs)
+        # 1. Save true weights and adapte keys
+        save_tptt_safetensors(self, path)
+        # 2. Copy Python files for trust_remote_code
+        self._copy_source_files(path, **kwargs)
+
+    def _adjust_save_strategy(self, path: str, **kwargs):
+        """Re-adapt/remove the weight safetensor and saved adapter config"""
+        if isinstance(self.tptt_model, PeftModel):
+            self.tptt_model.save_pretrained(path, **kwargs)
+        safetensor_path = os.path.join(path, "model.safetensors")
+        if os.path.exists(safetensor_path):
+            os.remove(safetensor_path)
+        adapter_path = os.path.join(path, "adapter_config.json")
+        if os.path.exists(adapter_path):
+            os.remove(adapter_path)
+
+    def _copy_source_files(self, target_path: str, **kwargs):
+        """Copy all .py files from package directory for trust_remote_code."""
+        src_dir = os.path.dirname(os.path.abspath(__file__))
+        dst_dir = (
+            f"./{str(Path(target_path).parts[0])}"
+            if kwargs.get("subfolder", False)
+            else target_path
+        )
+        for fname in os.listdir(src_dir):
+            if fname.endswith(".py"):
+                src = os.path.join(src_dir, fname)
+                dst = os.path.join(dst_dir, fname)
+                shutil.copy2(src, dst)
+
+    def retie_lm_after_load(self, **kwargs):
+        """Re-link lm_head after loading external weights."""
+        embed_lm = find_embedding_lm(self.tptt_model)
+        if embed_lm is not None and hasattr(self.tptt_model, "lm_head"):
+            if self.tptt_model.lm_head is None:  # ensure lm_head exists
+                self.tptt_model.lm_head = nn.Linear(
+                    embed_lm.weight.shape[1], embed_lm.weight.shape[0], bias=False
+                )
+            if kwargs.get("tie_word_embeddings", True):
+                self.tptt_model.lm_head.weight = embed_lm.weight  # share weights
+                logger.info("Weights of lm_head have been shared with embedding.")
+            else:
+                self.tptt_model.lm_head.weight = nn.Parameter(embed_lm.weight.clone())
+                logger.info("Weights of lm_head have been cloned from the embedding.")
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path=None, *model_args, **kwargs):
+        """Custom from_pretrained that accepts the standard positional argument"""
+        config = kwargs.pop("config", None)
+        repo_or_path = (
+            pretrained_model_name_or_path
+            or kwargs.pop("pretrained_model_name_or_path", None)
+            or kwargs.pop("repo_or_path", None)
+            or (getattr(config, "_base_path", None) if config else None)
+            or (getattr(config, "_name_or_path", None) if config else None)
+        )
+
+        if config is None and repo_or_path is not None:
+            config = AutoConfig.from_pretrained(repo_or_path, **kwargs)
+        model = cls(config, *model_args, **kwargs)
+        model.retie_lm_after_load(**kwargs)
+        return model
+
+
+TpttModel.register_for_auto_class("AutoModelForCausalLM")
+
+
+class LinearAttentionOp(nn.Module):
+    """Base class for linear attention operators."""
+
+    def __init__(
+        self,
+        layer_idx: int,
+        operator_mode: str = "delta_rule",
+        use_linear_checkpoint: bool = False,
+        recurrent_config: Optional[dict] = None,
+        max_chunk_size: int = 64,
+        linear_cache: Optional[LCache] = None,
+        linear_precision: torch.dtype = torch.float32,
+    ):
+        super().__init__()
+        self.layer_idx = layer_idx
+        if recurrent_config is None:
+            operator_mode = "delta_rule"  # force default operator mode if no config
+            recurrent_config = {
+                "order": 1,
+                "gate_type": "k",
+                "linear": True,
+                "trick": "derivative",
+            }
+        self.operator_mode = operator_mode
+        self.use_linear_checkpoint = use_linear_checkpoint
+
+        self.order = recurrent_config["order"]
+        self.gate_type = recurrent_config["gate_type"]
+        self.linear = recurrent_config["linear"]
+        self.trick = recurrent_config["trick"]
+
+        self.max_chunk_size = max_chunk_size
+        self.linear_cache = linear_cache or LCache()
+        self.linear_precision = linear_precision
+
+    def compute_gate(self, beta: Tuple[torch.Tensor]) -> torch.Tensor:
+        """
+        Compute the gating tensor according to the gate_type.
+        """
+        if self.gate_type == "k":
+            return torch.clamp(beta[0], min=1e-6, max=1 - 1e-6)
+        if self.gate_type == "v":
+            return torch.clamp(beta[1], min=1e-6, max=1 - 1e-6)
+        if self.gate_type == "kv":
+            return torch.clamp(beta[0] * beta[1], min=1e-6, max=1 - 1e-6)
+        raise ValueError(f"Unsupported gate_type: {self.gate_type}")
+
+    def get_cache(self, use_cache: bool) -> Tuple[
+        Optional[torch.Tensor],
+        Optional[Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]],
+    ]:
+        """
+        Retrieve recurrent state and qkv buffers from the cache.
+        """
+        if not use_cache:
+            return None, None
+        last_state = self.linear_cache[self.layer_idx]
+        if last_state is not None:
+            recurrent_state = last_state.get("recurrent_state", None)
+            qkv_buffers = last_state.get("qkv", None)
+        else:
+            recurrent_state = None
+            qkv_buffers = None
+        return recurrent_state, qkv_buffers
+
+    def save_cache(
+        self,
+        use_cache: bool,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        gate: torch.Tensor,
+        state: torch.Tensor,
+    ) -> None:
+        """
+        Save the recurrent state and qkv buffers to the cache.
+        """
+        if not use_cache:
+            return
+        if self.order > 1:
+            qkv_buffers = (
+                q[:, :, -(self.order - 1) :, :],
+                k[:, :, -(self.order - 1) :, :],
+                v[:, :, -(self.order - 1) :, :],
+                gate[:, :, -(self.order - 1) :, :],
+            )
+        else:
+            qkv_buffers = None
+        self.linear_cache.update(self.layer_idx, recurrent_state=state, qkv=qkv_buffers)
+
+    def forward(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        beta: Union[Tuple[torch.Tensor], torch.Tensor],
+        **kwargs,
+    ) -> torch.Tensor:
+        """
+        Forward pass for the attention operator.
+        """
+        # Ensure linear_precision for numerical stability (float32)
+        q, k, v = [x.to(self.linear_precision) for x in (q, k, v)]
+        if isinstance(beta, (tuple, list)):
+            beta = tuple(b.to(self.linear_precision) for b in beta)
+        else:
+            beta = beta.to(self.linear_precision)
+
+        gate = self.compute_gate(beta)
+
+        # Retrieve cache if needed
+        use_cache = kwargs.get("use_cache", False)
+        use_checkpoint = not (use_cache) and self.use_linear_checkpoint
+        recurrent_state, qkvb = self.get_cache(use_cache)
+
+        if qkvb is not None and qkvb[0].shape == q.shape:
+            q = torch.cat([qkvb[0].to(q.device), q], dim=2).to(self.linear_precision)
+            k = torch.cat([qkvb[1].to(q.device), k], dim=2).to(self.linear_precision)
+            v = torch.cat([qkvb[2].to(q.device), v], dim=2).to(self.linear_precision)
+            gate = torch.cat([qkvb[3].to(q.device), gate], dim=2).to(
+                self.linear_precision
+            )
+
+        output, state = self.chunk_delta_product_forward(
+            q,
+            k,
+            v,
+            gate,
+            self.max_chunk_size,
+            n=self.order,
+            trick=self.trick,
+            linear=self.linear,
+            initial_state=recurrent_state,
+            use_checkpoint=use_checkpoint,
+            linear_precision=self.linear_precision,
+        )
+
+        # Save cache if needed
+        self.save_cache(use_cache, q, k, v, gate, state)
+
+        return output
+
+    @staticmethod
+    def chunk_delta_product_forward(
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        beta_gate: torch.Tensor,
+        chunk_size: int,
+        n: int = 1,
+        trick: str = "derivative",
+        linear: bool = True,
+        initial_state: Optional[torch.Tensor] = None,
+        use_checkpoint: bool = True,
+        linear_precision: torch.dtype = torch.float32,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Chunkwise parallel implementation https://arxiv.org/abs/2406.06484
+        For each chunk, processes chunk_size * n_orders steps (virtual tokens) in order.
+        """
+
+        # --- Main chunk_delta_product_forward logic ---
+
+        batch_size, num_heads, seq_len, head_dim = query.shape
+        chunk_size = get_valid_chunk_size(seq_len, chunk_size)
+        num_chunks = seq_len // chunk_size
+
+        query_n = query if n == 1 else expand_virtual_tokens(query, n, trick)
+        key_n = key if n == 1 else expand_virtual_tokens(key, n, trick)
+        value_n = value if n == 1 else expand_virtual_tokens(value, n, trick)
+        beta_n = beta_gate if n == 1 else expand_virtual_tokens(beta_gate, n, trick)
+
+        q_chunks = chunk_sequence(query_n, num_chunks, chunk_size * n)
+        k_chunks = chunk_sequence(key_n, num_chunks, chunk_size * n)
+        v_chunks = chunk_sequence(value_n, num_chunks, chunk_size * n)
+        beta_chunks = chunk_sequence(beta_n, num_chunks, chunk_size * n)
+
+        k_beta = k_chunks * beta_chunks
+        v_beta = v_chunks * beta_chunks
+
+        householder = -(k_beta @ k_chunks.transpose(-2, -1)).tril(-1)
+        householder = ensure_stability(householder, min_val=-1e4, max_val=1e4)
+
+        # size : N = chunk_size * n
+        inv_hh = fast_invert_matrix(householder, dtype=linear_precision)  # [(...),N,N]
+
+        w = ensure_stability(torch.matmul(inv_hh, k_beta), min_val=-1e4, max_val=1e4)
+        u = ensure_stability(torch.matmul(inv_hh, v_beta), min_val=-1e4, max_val=1e4)
+
+        state_shape = (batch_size, num_heads, n, head_dim, head_dim)
+        if initial_state is not None and initial_state.shape == state_shape:
+            state = initial_state.to(device=query.device, dtype=linear_precision)
+        else:
+            state = torch.full(
+                state_shape,
+                fill_value=1e-6,  # stability if unlinear activation
+                device=query.device,
+                dtype=linear_precision,
+            )
+
+        output, final_state = sequential_delta_product_scan(
+            q_chunks.to(dtype=linear_precision),
+            w.to(dtype=linear_precision),
+            u.to(dtype=linear_precision),
+            n,
+            linear,
+            chunk_size,
+            state.to(dtype=linear_precision),
+            linear_precision=linear_precision,
+            use_checkpoint=use_checkpoint,
+        )
+
+        idx_last_order = torch.arange(chunk_size, device=output.device) * n + (n - 1)
+        output = output[:, :, :, idx_last_order, :]  # [B, H, num_chunks, chunk_size, D]
+        output = output.reshape(batch_size, num_heads, seq_len, head_dim)
+
+        return output.to(dtype=linear_precision), final_state.to(dtype=linear_precision)
+
+
+def sequential_delta_product_scan(
+    q_chunks: torch.Tensor,
+    w: torch.Tensor,
+    u: torch.Tensor,
+    n_orders: int,
+    linear_activation: bool,
+    current_chunk_size: int,
+    initial_recurrent_state: torch.Tensor,
+    linear_precision: torch.dtype,
+    use_checkpoint: bool,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    DeltaProduct implementation https://arxiv.org/abs/2502.10297
+    Implements the per-token Householder state updates.
+    """
+    batch, head, num_chunks_inner, chunk_n_total, dim = q_chunks.shape
+    output_inner = torch.empty_like(q_chunks)
+    # initial_recurrent_state is H_{last_token_of_prev_chunk, n-1} ([B, H, D, D])
+    h_0_base = initial_recurrent_state[:, :, -1, :, :].clone()
+
+    def process_one_chunk(
+        q_chunk_params: torch.Tensor,
+        w_chunk_params: torch.Tensor,
+        u_chunk_params: torch.Tensor,
+        h_0_base: torch.Tensor,
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        Process a single chunk (with per-token state for n_orders > 1).
+        """
+        o_intra_current_chunk = torch.zeros(
+            batch,
+            head,
+            chunk_n_total,
+            dim,
+            device=q_chunk_params.device,
+            dtype=linear_precision,
+        )
+        o_inter_current_chunk = torch.zeros_like(o_intra_current_chunk)
+        current_accumulated_state_per_token = (
+            h_0_base.unsqueeze(2).expand(-1, -1, current_chunk_size, -1, -1).clone()
+        )  # [B, H, current_chunk_size, D, D]
+
+        for step in range(n_orders):
+            idx_virtual_tokens = (
+                torch.arange(current_chunk_size, device=q_chunk_params.device)
+                * n_orders
+                + step
+            )
+            q_s = q_chunk_params[:, :, idx_virtual_tokens, :]
+            w_s = w_chunk_params[:, :, idx_virtual_tokens, :]
+            u_s = u_chunk_params[:, :, idx_virtual_tokens, :]
+
+            state_input_for_this_step = current_accumulated_state_per_token
+
+            ## BLAS/cuBLAS einsum "bhcd,bhcdd->bhcd"
+            k_trans_h_old = (
+                torch.matmul(
+                    w_s.unsqueeze(-2),
+                    state_input_for_this_step,
+                )
+                .squeeze(-2)
+                .to(dtype=linear_precision)
+            )
+
+            u_val = u_s - k_trans_h_old
+
+            o_inter_current_chunk[:, :, idx_virtual_tokens, :] = (
+                torch.matmul(q_s.unsqueeze(-2), state_input_for_this_step)
+                .squeeze(-2)
+                .to(dtype=linear_precision)
+            )
+
+            ## BLAS/cuBLAS einsum "bhcd,bhcd->bhcd"
+            o_intra_current_chunk[:, :, idx_virtual_tokens, :] = (q_s * u_val).to(
+                dtype=linear_precision
+            )
+
+            outer_product_term = torch.matmul(w_s.unsqueeze(-1), u_val.unsqueeze(-2))
+            new_state_i_per_token = state_input_for_this_step + outer_product_term
+            current_accumulated_state_per_token = new_state_i_per_token.to(
+                dtype=linear_precision
+            )
+        # Return all needed for next chunk
+        return (
+            o_intra_current_chunk,
+            o_inter_current_chunk,
+            current_accumulated_state_per_token[:, :, -1, :, :],  # new h_0_base
+        )
+
+    for chunk_idx_inner in range(num_chunks_inner):
+        q_chunk_params = q_chunks[:, :, chunk_idx_inner]
+        w_chunk_params = w[:, :, chunk_idx_inner]
+        u_chunk_params = u[:, :, chunk_idx_inner]
+
+        # Checkpointed call if training
+        call = (
+            partial(checkpoint, use_reentrant=False)
+            if use_checkpoint
+            else lambda f, *a: f(*a)
+        )
+        o_intra, o_inter, h_0_base = call(
+            process_one_chunk,
+            q_chunk_params,
+            w_chunk_params,
+            u_chunk_params,
+            h_0_base,
+        )
+        if not linear_activation:  # unlinear activation between chunks
+            h_0_base = unlinear_activation(h_0_base).to(dtype=linear_precision)
+        output_inner[:, :, chunk_idx_inner] = o_intra + o_inter
+
+    return output_inner, h_0_base
+
+
+def unlinear_activation(x: torch.Tensor, scale: float = 2.0) -> torch.Tensor:
+    """Unlinear activation between chunk"""
+    x_n = x.norm(p=2, dim=-1, keepdim=True) + 1e-6
+    x_gelu = F.gelu(scale * x / x_n, approximate="tanh")  # pylint: disable=not-callable
+    return (x / scale) * x_gelu
+
+
+def chunk_sequence(x: torch.Tensor, num_chunks: int, chunk_size: int) -> torch.Tensor:
+    """Splits [B, H, S, D] to  [B, H, num_chunks, chunk_size, D]"""
+    batch_size, num_heads, _, head_dim = x.shape
+    return x.reshape(batch_size, num_heads, num_chunks, chunk_size, head_dim)
+
+
+def expand_virtual_tokens(
+    x: torch.Tensor, n: int, mode: str = "derivative"
+) -> torch.Tensor:
+    """Expand tokens into 'n' virtual tokens using the selected trick."""
+    batch_size, num_heads, seq_len, head_dim = x.shape
+    device, dtype = x.device, x.dtype
+
+    def derivative_expand(x: torch.Tensor) -> torch.Tensor:
+        """Expand tokens using the derivative trick."""
+        x_pad = torch.cat(
+            [
+                torch.zeros(
+                    batch_size, num_heads, n - 1, head_dim, device=device, dtype=dtype
+                ),
+                x,
+            ],
+            dim=2,
+        )
+        coeffs = torch.tensor(
+            [(-1) ** k * math.comb(n - 1, k) for k in range(n)],
+            device=device,
+            dtype=dtype,
+        )
+        coeffs /= coeffs.norm(p=1)
+        return (
+            (x_pad.unfold(2, n, 1) * coeffs.view(1, 1, 1, 1, n))
+            .flip(-1)
+            .permute(0, 1, 2, 4, 3)
+            .reshape(batch_size, num_heads, seq_len * n, head_dim)
+        )
+
+    def rotative_expand(x: torch.Tensor) -> torch.Tensor:
+        """Expand tokens using the rotative trick."""
+        d_parity = head_dim // 2
+        angles = torch.arange(n, device=device, dtype=dtype) * (2 * math.pi / n)
+        cos = torch.cos(angles).view(1, 1, 1, n, 1)
+        sin = torch.sin(angles).view(1, 1, 1, n, 1)
+        if head_dim % 2:
+            x_pairs = x[..., :-1].view(batch_size, num_heads, seq_len, d_parity, 2)
+        else:
+            x_pairs = x.view(batch_size, num_heads, seq_len, d_parity, 2)
+        x_pairs = x_pairs.unsqueeze(3).expand(
+            batch_size, num_heads, seq_len, n, d_parity, 2
+        )
+        x0, x1 = x_pairs[..., 0], x_pairs[..., 1]
+        x0r = x0 * cos - x1 * sin
+        x1r = x0 * sin + x1 * cos
+        rot = torch.stack([x0r, x1r], -1).reshape(
+            batch_size, num_heads, seq_len, n, d_parity * 2
+        )
+        if head_dim % 2:
+            last = (
+                x[..., -1]
+                .unsqueeze(-1)
+                .unsqueeze(3)
+                .expand(batch_size, num_heads, seq_len, n, 1)
+            )
+            rot = torch.cat([rot, last], -1)
+        return rot.reshape(batch_size, num_heads, seq_len * n, head_dim)
+
+    if mode == "derivative":
+        return derivative_expand(x)
+    if mode == "rotative":
+        return rotative_expand(x)
+    if mode == "combined":
+        return (derivative_expand(x) + rotative_expand(x)) / 2
+    raise ValueError(f"Unknown mode: {mode}")
+
+
+def extract_layer_idx(module_name: str) -> int:
+    """Extract the layer index from a module name string."""
+    match = re.search(r"\.(\d+)\.", module_name)
+    if match:
+        return int(match.group(1))
+    return -1
+
+
+def find_embedding_lm(module: nn.Module) -> Optional[nn.Module]:
+    """Find the embedding weight in a model module."""
+    for _, child in module.named_modules():
+        if hasattr(child, "embed_tokens") and hasattr(child.embed_tokens, "weight"):
+            return child.embed_tokens
+        if hasattr(child, "token_embeddings") and hasattr(
+            child.token_embeddings, "weight"
+        ):
+            return child.token_embeddings
+    return None
+
+
+def set_trainable_parameters(
+    model: PreTrainedModel, trainable_patterns: List[str] = None
+) -> PreTrainedModel:
+    """Freeze model parameters except trainable_patterns."""
+    if trainable_patterns is None:
+        trainable_patterns = [
+            "q_proj",
+            "k_proj",
+            "v_proj",
+            "o_proj",
+            "qkv_proj",
+            "out_proj",
+            "c_attn",
+            "c_proj",
+            "query",
+            "key",
+            "value",
+        ]
+
+    for name, param in model.named_parameters():
+        param.requires_grad = any(pattern in name for pattern in trainable_patterns)
+
+    trainable_layers = [n for n, p in model.named_parameters() if p.requires_grad]
+    logger.info("Trainable parameters after freeze: %s", trainable_layers)
+    return model
+
+
+def ensure_stability(
+    tensor: torch.Tensor, min_val: float = -1e4, max_val: float = 1e4
+) -> torch.Tensor:
+    """stability forcing"""
+    dtype = tensor.dtype
+    center = (max_val + min_val) / 2
+    tensor = torch.clamp(tensor, min=min_val, max=max_val)
+    tensor = torch.nan_to_num(tensor, nan=center, posinf=max_val, neginf=min_val)
+    return tensor.to(dtype=dtype)
+
+
+def apply_linear_attention_mask(
+    attention_mask: torch.Tensor, v: torch.Tensor, padding_side: str = "right"
+) -> torch.Tensor:
+    """Extract if padding --> [B,S]"""
+    if attention_mask.dim() == 4 and attention_mask.shape[1] == 1:
+        mask = attention_mask.diagonal(dim1=-2, dim2=-1).squeeze(1)
+    else:
+        mask = attention_mask.squeeze(
+            dim=tuple(
+                i
+                for i in range(1, attention_mask.dim())
+                if attention_mask.shape[i] == 1
+            )
+        )
+    # Ensure cast to the same dtype as v and convert to binary mask
+    if not (
+        mask.dtype == torch.bool
+        or (
+            mask.dtype in [torch.uint8, torch.int32, torch.int64]
+            and mask.max() <= 1
+            and mask.min() >= 0
+        )
+    ):
+        mask = (mask >= 0).to(v.dtype)  # [-inf, 0, 0, -inf] --> [0, 1, 1, 0]
+    else:
+        mask = mask.to(v.dtype)
+    # mask is [batch, seq] --> Broadcast to v [batch, seq, (...)]
+    if padding_side == "left":
+        mask = mask[:, -v.shape[-2] :][(...,) + (None,) * (v.dim() - 2)]
+    else:  # right padding
+        mask = mask[:, : v.shape[-2]][(...,) + (None,) * (v.dim() - 2)]
+    return v * mask
+
+
+def truncate_attention_mask(
+    hidden_states: torch.Tensor, attention_mask: torch.Tensor, max_length: int
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Truncate hidden_states and attention_mask to the last window of size max_length"""
+    seq_dim = 1  # convention: (batch, seq, ...)
+    seq_len = hidden_states.shape[seq_dim]
+    if seq_len > max_length:
+        hidden_states = hidden_states.narrow(seq_dim, seq_len - max_length, max_length)
+        if attention_mask is not None:
+            # mask [batch, seq]
+            if attention_mask.dim() == 2:
+                attention_mask = attention_mask[:, -max_length:]
+            # mask [batch, seq, seq]
+            elif attention_mask.dim() == 3:
+                attention_mask = attention_mask[:, -max_length:, -max_length:]
+            # mask [batch, 1, seq, seq]
+            elif attention_mask.dim() == 4 and attention_mask.shape[1] == 1:
+                attention_mask = attention_mask[:, :, -max_length:, -max_length:]
+            else:
+                raise ValueError(
+                    "No dimension in attention_mask matches sequence length of hidden_states."
+                )
+    return hidden_states, attention_mask
+
+
+def fast_invert_matrix(
+    tri_tensor: torch.Tensor, dtype: torch.dtype = torch.float32
+) -> torch.Tensor:
+    """Equivalent to vectorized forward substitution applied to the identity matrix."""
+    tri_tensor = tri_tensor.to(dtype=dtype).clone()
+    chunk_size = tri_tensor.shape[-1]
+
+    for i in range(1, chunk_size):
+        tri_tensor[..., i, :i] = tri_tensor[..., i, :i] + (
+            tri_tensor[..., i, :, None].clone() * tri_tensor[..., :, :i].clone()
+        ).sum(-2)
+
+    tri_tensor = tri_tensor + torch.eye(
+        chunk_size, dtype=dtype, device=tri_tensor.device
+    )
+    return tri_tensor.to(dtype=dtype)
+
+
+def get_valid_chunk_size(total_l: int, chunk_size: int) -> int:
+    """Return the largest chunk_size <= chunk_size that divides total_l."""
+    for c in range(min(chunk_size, total_l), 0, -1):
+        if total_l % c == 0:
+            return c
+    return 1
+
+
+## RARELY
+def split_qkv(
+    base_attn: nn.Module, qkv: torch.Tensor
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Split the QKV tensor into separate Q, K, and V tensors."""
+    num_q_heads = getattr(base_attn, "num_q_heads", None)
+    num_k_heads = getattr(base_attn, "num_k_heads", None)
+    num_v_heads = getattr(base_attn, "num_v_heads", None)
+    head_dim = getattr(base_attn, "head_dim", None)
+
+    if num_q_heads is None or num_k_heads is None or num_v_heads is None:
+        raise ValueError(
+            "Base attention must have num_q_heads, num_k_heads, and num_v_heads defined."
+        )
+
+    q_len = num_q_heads * head_dim
+    k_len = num_k_heads * head_dim
+    v_len = num_v_heads * head_dim
+
+    q, k, v = torch.split(qkv, [q_len, k_len, v_len], dim=-1)
+    return q, k, v
+
+
+## OPTIONAL
+def match_dim(x: torch.Tensor, dim: int, target_size: int) -> torch.Tensor:
+    """Match the size of tensor x along dimension dim to target_size by interpolation"""
+    src_size = x.shape[dim]
+    if src_size == target_size:
+        return x
+    x = torch.moveaxis(x, dim, -1)
+    shape = x.shape
+    if src_size < target_size:
+        x = x.reshape(-1, 1, src_size)
+        x = F.interpolate(x, size=target_size, mode="linear", align_corners=False)
+        x = x.reshape(*shape[:-1], target_size)
+    else:
+        eye = torch.eye(target_size, src_size, device=x.device, dtype=x.dtype)
+        x = F.linear(x, eye)  # pylint: disable=not-callable
+    x = torch.moveaxis(x, -1, dim)
+    return x
+
+
+def soft_clamp(
+    x: torch.Tensor, min_val: float = 1e-6, max_val: float = 1 - 1e-6
+) -> torch.Tensor:
+    """Differentiable clamping for stability"""
+    dtype = x.dtype
+    scale = (max_val - min_val) / 2
+    center = (max_val + min_val) / 2
+    return (torch.tanh((x - center) / scale) * scale + center).to(dtype=dtype)
+
+
+def describe(x: torch.Tensor, name="tensor") -> None:
+    """Prints the shape, min, max, mean, and std of a tensor."""
+    stats = (x.min(), x.max(), x.mean(), x.std())
+    print(
+        f"{name} shape: {tuple(x.shape)}, "
+        + f"min: {stats[0]:.4g}, max: {stats[1]:.4g}, "
+        + f"mean: {stats[2]:.4g}, std: {stats[3]:.4g}, "
+        + f"dtype: {x.dtype}, device: {x.device}"
+    )

lora_delta_product_m0.5_constant/runs/Aug27_14-19-44_61e44eee8185/events.out.tfevents.1756304386.61e44eee8185.163.0

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:572bd4fe44a9b92d07140da1219bbc70bd993f74cdde397d60ddb7c3d10b2dad
+size 69803

lora_delta_product_m0.5_constant/special_tokens_map.json

@@ -0,0 +1,24 @@
+{
+  "bos_token": {
+    "content": "<s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "</s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": "</s>",
+  "unk_token": {
+    "content": "<unk>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

lora_delta_product_m0.5_constant/tokenizer.model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9e556afd44213b6bd1be2b850ebbbd98f5481437a8021afaf58ee7fb1818d347
+size 499723

lora_delta_product_m0.5_constant/tokenizer_config.json

@@ -0,0 +1,43 @@
+{
+  "add_bos_token": true,
+  "add_eos_token": false,
+  "add_prefix_space": null,
+  "added_tokens_decoder": {
+    "0": {
+      "content": "<unk>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "<s>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "</s>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "bos_token": "<s>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "</s>",
+  "extra_special_tokens": {},
+  "legacy": false,
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "</s>",
+  "padding_side": "right",
+  "sp_model_kwargs": {},
+  "tokenizer_class": "LlamaTokenizer",
+  "unk_token": "<unk>",
+  "use_default_system_prompt": false
+}

modeling_tptt.py

@@ -1,90 +1,271 @@
+# pylint: disable=too-many-lines, too-many-arguments, too-many-positional-arguments, too-many-instance-attributes, too-many-locals
+
 """
 This module implements the TPTT model with linear attention (LiZA) and LoRA support.
 Author : Fabien FURFARO
+TPTT : Transforming Pretrained Transformers into Titans (https://arxiv.org/abs/2506.17671)
 """
 
 import logging
+import math
 import os
+from pathlib import Path
 import re
 import shutil
-from typing import Dict, List, Optional
+from functools import partial
+from typing import Any, Dict, List, Optional, Tuple, Union
 
 import torch
 import torch.nn.functional as F
 from einops import rearrange
 from huggingface_hub import hf_hub_download, list_repo_files
-from peft import LoraConfig, get_peft_model
+from peft import LoraConfig, PeftModel, get_peft_model
 from safetensors import safe_open
+from safetensors.torch import save_file
 from torch import nn
-from transformers import AutoModelForCausalLM, DynamicCache, PreTrainedModel
+from torch.utils.checkpoint import checkpoint
+from transformers import (
+    AutoConfig,
+    AutoModel,
+    AutoModelForCausalLM,
+    DynamicCache,
+    PreTrainedModel,
+)
 from transformers.configuration_utils import PretrainedConfig
 
 from .configuration_tptt import TpttConfig
 
+logger = logging.getLogger(__name__)  # monitoring
 
-def import_fla_ops():
-    """flash linear attention"""
-    if torch.cuda.is_available():
-        try:
-            from fla.ops.gla import fused_chunk_gla, fused_recurrent_gla
 
-            return fused_chunk_gla, fused_recurrent_gla
-        except ImportError:
-            return None, None
-    return None, None
+class LCache:
+    """Cache for storing intermediate states of linear attention layers."""
+
+    def __init__(self):
+        """Stores per-layer intermediate states: {layer_idx: state_dict}"""
+        self.inputs_states: Dict[int, Dict[str, torch.Tensor]] = (
+            {}
+        )  # recurrent states and qkv buffers
 
+    def __getitem__(self, layer_idx: int) -> Optional[Dict[str, torch.Tensor]]:
+        """Retrieve cached state for a given layer, or None if not present"""
+        return self.inputs_states.get(layer_idx, None)
 
-fused_chunk_gla, fused_recurrent_gla = import_fla_ops()  # TODO: add all ops
+    def update(self, layer_idx: int, **kwargs):
+        """Detach all tensors to avoid retaining computation graphs"""
+        detached_kwargs = {
+            k: v.detach() if isinstance(v, torch.Tensor) else v
+            for k, v in kwargs.items()
+        }
+        # Update or create the state for the specified layer
+        if layer_idx in self.inputs_states:
+            self.inputs_states[layer_idx].update(detached_kwargs)
+        else:
+            self.inputs_states[layer_idx] = detached_kwargs
 
-logger = logging.getLogger(__name__)  # monitoring
+    def reset(self):
+        """Clear all cached states and reset the token counter"""
+        self.inputs_states.clear()
 
 
-class LCache:
+class CausalAvgPool1d(nn.Module):
+    """Causal sliding window average (uniform, no shape loss along sequence)"""
+
+    def __init__(
+        self, output_size: int, offsets: tuple[int] = (0, 1, 2), mode: str = "replicate"
+    ):
+        super().__init__()
+        self.offsets = offsets
+        self.mode = mode
+        self.pool = nn.AdaptiveAvgPool1d(output_size=output_size)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """x: [B, S, F] → [B, S, F → output_size]"""
+        x_ = x.transpose(1, 2)  # [B, F, S]
+        idxs = torch.tensor(self.offsets, device=x.device)
+        ksize = idxs.max() - idxs.min() + 1
+        w = torch.zeros(ksize, device=x.device, dtype=x.dtype)
+        w[idxs - idxs.min()] = 1 / len(self.offsets)  # Always uniform weights
+        kernel = w.repeat(x_.shape[1], 1).reshape(x_.shape[1], 1, ksize)
+        pad_left = -idxs.min().item()
+        pad_right = (ksize - 1) - pad_left
+        x_pad = F.pad(x_, (pad_left, pad_right), mode=self.mode)
+        y = F.conv1d(x_pad, kernel, groups=x_.shape[1])  # pylint: disable=not-callable
+        return self.pool(y.transpose(1, 2))  # [B, S, F → output_size]
+
+
+class LinearAttention(nn.Module):
     """
-    Cache for storing intermediate states of linear attention layers.
-    Supports a sliding window if max_length is set.
+    Linear multi-head attention layer: [B, S, D] -> [B, S, D]
+    Projections + gating + efficient linear attention mechanism (TPTT compatible).
     """
 
-    def __init__(self):
-        """
-        Initialize the cache.
+    def __init__(
+        self,
+        hidden_dim: int,
+        num_heads: int,
+        head_dim: Optional[int] = None,
+        num_key_value_heads: Optional[int] = None,
+        num_key_value_groups: Optional[int] = None,
+        bias: bool = True,
+        dropout: Optional[float] = None,
+        linear_precision: torch.dtype = torch.float32,
+        padding_side: str = "right",
+        shared_attn: bool = False,  # shared attention
+        layer_idx: int = 0,
+        operator_mode: str = "delta_rule",
+        use_linear_checkpoint: bool = False,
+        recurrent_config: Optional[Dict[str, Any]] = None,
+        linear_cache: Optional[LCache] = None,
+        max_chunk_size: int = 64,
+        bidirectional: bool = False,  # not used if causal
+        pooling_config: Optional[Dict[str, Any]] = None,
+    ):
+        super().__init__()
+        if pooling_config is None:
+            pooling_config = {
+                "offsets": (0, 1, 2),
+                "mode": "replicate",
+            }
+        self.hidden_dim = hidden_dim
+        self.num_heads = num_heads
+        self.head_dim = head_dim or hidden_dim // num_heads
+        self.num_key_value_heads = num_key_value_heads or num_heads
+        self.num_key_value_groups = num_key_value_groups or (
+            num_heads // (num_key_value_heads or num_heads)
+        )
+        self.scaling = self.head_dim**-0.5
+        self.linear_precision = linear_precision
+        self.padding_side = padding_side
 
-        Args:
-            max_length (Optional[int]): Maximum number of tokens to keep per layer (if set).
-        """
-        self.states: List[Dict[str, torch.Tensor]] = []
-        self.seen_tokens = 0
+        self.shared_attn = shared_attn
 
-    def __getitem__(self, layer_idx: int) -> Optional[Dict[str, torch.Tensor]]:
-        """
-        Retrieve the state for the given layer index, if it exists.
-        """
-        if layer_idx < len(self.states):
-            return self.states[layer_idx]
-        return None
+        if not shared_attn:
+            self.q_proj = nn.Linear(hidden_dim, num_heads * self.head_dim, bias=bias)
+            self.k_proj = nn.Linear(
+                hidden_dim, self.num_key_value_heads * self.head_dim, bias=bias
+            )
+            self.v_proj = nn.Linear(
+                hidden_dim, self.num_key_value_heads * self.head_dim, bias=bias
+            )
+            self.out_proj = nn.Linear(num_heads * self.head_dim, hidden_dim, bias=bias)
 
-    def update(self, layer_idx: int, **kwargs):
+        self.dropout = nn.Dropout(dropout) if dropout is not None else None
+
+        self.linear_operator = LinearAttentionOp(
+            layer_idx=layer_idx,
+            operator_mode=operator_mode,
+            use_linear_checkpoint=use_linear_checkpoint,
+            recurrent_config=recurrent_config,
+            max_chunk_size=max_chunk_size,
+            linear_cache=linear_cache,
+            linear_precision=linear_precision,
+        )
+        self.bidirectional = bidirectional
+        # Causal average pooling for gating
+        self.pooling_config = pooling_config
+        self.pool_g = CausalAvgPool1d(
+            output_size=self.head_dim * self.num_key_value_heads, **pooling_config
+        )
+
+    def forward(
+        self,
+        x: Union[List[torch.Tensor], torch.Tensor],
+        attn_mask: Optional[torch.Tensor] = None,
+        out_proj: Optional[nn.Module] = None,
+        **kwargs: Any,
+    ) -> torch.Tensor:
         """
-        Update the cache for a given layer.
-        If max_length is set, keep only the last max_length tokens in any sequence state.
+        Forward pass for linear attention. Input shape: [B, S, D], output [B, S, D].
         """
-        detached_kwargs = {}
-        for key, value in kwargs.items():
-            if isinstance(value, torch.Tensor):
-                value = value.detach()
-            detached_kwargs[key] = value
-
-        if len(self.states) <= layer_idx:
-            self.states.append(detached_kwargs)
+
+        if not self.shared_attn:
+            hidden_states = x[0] if isinstance(x, (list, tuple)) else x
+            # Projections
+            q = self.q_proj(hidden_states)
+            k = self.k_proj(hidden_states)
+            v = self.v_proj(hidden_states)
+            out_proj = self.out_proj
         else:
-            self.states[layer_idx].update(detached_kwargs)
+            # Shared attention <=> no projections here
+            q, k, v = x[0], x[1], x[2]
+            out_proj = self.out_proj if out_proj is None else out_proj
 
-    def reset(self):
-        """
-        Reset the cache and token counter.
-        """
-        self.states.clear()
-        self.seen_tokens = 0
+        # get dtype and device
+        final_dtype, final_device = q.dtype, q.device
+        # Masking if needed
+        if attn_mask is not None:
+            v = apply_linear_attention_mask(attn_mask, v, self.padding_side)
+
+        # Forget and Write Gating for linear attn (abusive term)
+        f_g, w_g = self.pool_g(k), self.pool_g(v)
+
+        # Reshape for multi-head
+        q = rearrange(q, "b n (h d) -> b h n d", h=self.num_heads)
+        k = rearrange(k, "b n (h d) -> b h n d", h=self.num_key_value_heads)
+        v = rearrange(v, "b n (h d) -> b h n d", h=self.num_key_value_heads)
+
+        f_g = rearrange(f_g, "b n (h m) -> b h n m", h=self.num_key_value_heads)
+        w_g = rearrange(w_g, "b n (h m) -> b h n m", h=self.num_key_value_heads)
+
+        # Repeat for GQA
+        k = k.repeat_interleave(self.num_key_value_groups, dim=1)
+        v = v.repeat_interleave(self.num_key_value_groups, dim=1)
+
+        f_g = f_g.repeat_interleave(self.num_key_value_groups, dim=1)
+        w_g = w_g.repeat_interleave(self.num_key_value_groups, dim=1)
+
+        ## DeltaNet-style: Silu activation and normalization
+        q = F.normalize(F.silu(q), p=2, dim=-1, eps=1e-6)
+        k = F.normalize(F.silu(k), p=2, dim=-1, eps=1e-6)
+
+        ## linear stability part
+        v = ensure_stability(v * self.scaling, min_val=-1e4, max_val=1e4)
+
+        # Apply sigmoid to forget and write gates
+        f_g = torch.clamp(torch.sigmoid(f_g), min=1e-6, max=1 - 1e-6)
+        w_g = torch.clamp(torch.sigmoid(w_g), min=1e-6, max=1 - 1e-6)
+
+        # Convert to linear_precision (float32) for numerical stability and get model dtype
+        q, k, v, f_g, w_g = (
+            x.to(self.linear_precision).contiguous() for x in (q, k, v, f_g, w_g)
+        )
+        g = (f_g, w_g)
+
+        # Linear Attention Core, output: [B, H, S, d]
+        if self.bidirectional:  # Work only with uncausal attention
+            # Forward direction
+            out_forward = self.linear_operator(q, k, v, g, **kwargs)
+            # Backward direction: flip the input sequence on the time dimension (dim=2)
+            kwargs_bwd = kwargs.copy()
+            kwargs_bwd["use_cache"] = False
+            out_backward = self.linear_operator(
+                torch.flip(q, dims=[2]),
+                torch.flip(k, dims=[2]),
+                torch.flip(v, dims=[2]),
+                tuple(torch.flip(t, dims=[2]) for t in g),
+                **kwargs_bwd,
+            )
+            # Flip the output back to restore proper order
+            out_backward = torch.flip(out_backward, dims=[2])
+            # Fusion: here, simple addition
+            out = out_forward + out_backward
+        else:
+            out = self.linear_operator(q, k, v, g, **kwargs)
+
+        # Merge heads and project: [B, H, S, d] -> [B, S, H*d] -> Out proj
+        out = rearrange(out, "b h s d -> b s (h d)")
+        # Normalize output (RMS norm). Note: bidirectional compatibility
+        out = out / out.pow(2).mean(dim=-1, keepdim=True).add(1e-6).sqrt()
+        # Ensure dtype and device consistency
+        out = out.to(dtype=final_dtype, device=final_device)
+        # Apply output projection
+        out = out_proj(out)  # [B, S, D]
+        out = ensure_stability(out, min_val=-1e4, max_val=1e4)
+        # Apply dropout if specified
+        if self.dropout is not None:
+            out = self.dropout(out)
+        return out
 
 
 class LiZAttention(nn.Module):
@@ -94,33 +275,69 @@ class LiZAttention(nn.Module):
         self,
         base_attn: nn.Module,
         layer_idx: int,
-        base_config,  # Backbone Config
+        base_config: PretrainedConfig,  # Backbone Config
         linear_cache: Optional[LCache] = None,
         operator_mode: str = "delta_rule",
-        max_self_attn_length: int = 2048,
+        use_linear_checkpoint: bool = False,
+        recurrent_config: Optional[Dict[str, Any]] = None,
+        max_self_attn_length: Optional[int] = None,  # unnecessary
+        base_scale_attn: bool = False,
         mag_weight: float = 0.5,
+        cross_gate: bool = False,
         max_chunk_size: int = 64,
+        linear_precision: Union[str, torch.dtype] = "float32",
+        padding_side: str = "right",  # for tokenizer
+        disable_linear_attn: bool = False,
+        bidirectional: bool = False,  # if True, use bidirectional attention
+        pooling_config: Optional[Dict[str, Any]] = None,
     ):
         super().__init__()
-        self.base_attn = base_attn
+        if isinstance(linear_precision, str):
+            linear_precision = getattr(torch, linear_precision)
+        self.linear_precision = linear_precision
+        self.base_attn: nn.Module = base_attn
         self.base_config = base_config
         self.layer_idx = layer_idx
         self.max_self_attn_length = max_self_attn_length
+        self.base_scale_attn = base_scale_attn
         self.mag_weight = mag_weight
+        self.cross_gate = cross_gate
         self.max_chunk_size = max_chunk_size
-        self.linear_cache = linear_cache or LCache()
+        self.linear_precision = linear_precision
+        self.padding_side = padding_side
+        self.disable_linear_attn = disable_linear_attn
+
         (
             self.num_heads,
             self.head_dim,
             self.num_key_value_heads,
             self.num_key_value_groups,
+            self.hidden_dim,
         ) = self._get_attention_parameters(base_attn, base_config)
-        self.operator = get_attention_operator(operator_mode)
-        self.pool_g = nn.AdaptiveAvgPool1d(
-            output_size=self.head_dim * self.num_key_value_heads
+        self.scaling = self.head_dim**-0.5
+
+        self.linear_attn = LinearAttention(
+            layer_idx=layer_idx,
+            shared_attn=True,
+            operator_mode=operator_mode,
+            use_linear_checkpoint=use_linear_checkpoint,
+            recurrent_config=recurrent_config,
+            hidden_dim=self.hidden_dim,
+            num_heads=self.num_heads,
+            head_dim=self.head_dim,
+            num_key_value_heads=self.num_key_value_heads,
+            num_key_value_groups=self.num_key_value_groups,
+            linear_precision=linear_precision,
+            linear_cache=linear_cache,
+            max_chunk_size=max_chunk_size,
+            padding_side=padding_side,
+            bidirectional=bidirectional,
+            pooling_config=pooling_config,
         )
 
-    def _get_attention_parameters(self, base_attn, base_config):
+    def _get_attention_parameters(
+        self, base_attn: nn.Module, base_config: PretrainedConfig
+    ) -> Tuple[Optional[int], Optional[int], Optional[int], Optional[int]]:
         """Retrieve the attention parameters from the base attention module."""
         # first order base attention module and second order config
         num_heads = (
@@ -129,8 +346,15 @@ class LiZAttention(nn.Module):
             or getattr(base_config, "num_heads", None)
             or getattr(base_config, "num_attention_heads", None)
         )
-        head_dim = getattr(base_attn, "head_dim", None) or getattr(
-            base_config, "head_dim", None
+        head_dim = (
+            getattr(base_attn, "head_dim", None)
+            or getattr(base_attn, "attention_head_size", None)
+            or getattr(base_config, "head_dim", None)
+            or (
+                getattr(base_config, "hidden_size", None) // num_heads
+                if num_heads and getattr(base_config, "hidden_size", None)
+                else None
+            )
         )
         num_key_value_heads = (
             getattr(base_attn, "num_kv_heads", None)
@@ -141,14 +365,18 @@ class LiZAttention(nn.Module):
         num_key_value_groups = getattr(base_attn, "num_key_value_groups", None) or (
             num_heads // num_key_value_heads if num_heads and num_key_value_heads else 1
         )
+        hidden_dim = getattr(base_config, "hidden_size", None) or head_dim * num_heads
         return (
             num_heads,
             head_dim,
             num_key_value_heads,
             num_key_value_groups,
+            hidden_dim,
         )
 
-    def _apply_projections(self, hidden_states):
+    def _apply_shared_projections(
+        self, hidden_states: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, nn.Module]:
         base_attn = self.base_attn
         if hasattr(base_attn, "q_proj"):
             # LLama, OLMO and Mistral style
@@ -166,89 +394,51 @@ class LiZAttention(nn.Module):
             qkv = base_attn.c_attn(hidden_states)
             q, k, v = qkv.chunk(3, dim=-1)
             out_proj = base_attn.c_proj
+        elif all(hasattr(base_attn, n) for n in ["query", "key", "value"]):
+            # BERT - ViT
+            q = base_attn.query(hidden_states)
+            k = base_attn.key(hidden_states)
+            v = base_attn.value(hidden_states)
+            out_proj = getattr(base_attn, "dense", None)  # ou output.dense
         else:
             raise ValueError("Unsupported attention module: cannot find projections.")
         # Ensure stability
-        q = torch.clamp(q, min=-1e4, max=1e4)
-        k = torch.clamp(k, min=-1e4, max=1e4)
-        v = torch.clamp(v, min=-1e4, max=1e4)
+        q = ensure_stability(q, min_val=-1e4, max_val=1e4)
+        k = ensure_stability(k, min_val=-1e4, max_val=1e4)
+        v = ensure_stability(v, min_val=-1e4, max_val=1e4)
         return q, k, v, out_proj
 
-    def _prepare_attn_input(self, q, k, v, gate_norm):
-        # Gating for linear attn
-        g = self.pool_g(k)
-
-        # Reshape for multi-head
-        q = rearrange(q, "b n (h d) -> b h n d", h=self.num_heads)
-        k = rearrange(k, "b n (h d) -> b h n d", h=self.num_key_value_heads)
-        v = rearrange(v, "b n (h d) -> b h n d", h=self.num_key_value_heads)
-        g = rearrange(g, "b n (h m) -> b h n m", h=self.num_key_value_heads)
-
-        # Repeat for GQA
-        k = repeat_kv(k, self.num_key_value_groups)
-        v = repeat_kv(v, self.num_key_value_groups)
-        g = repeat_kv(g, self.num_key_value_groups)
-
-        ## linear part
-        q = torch.clamp(F.softmax(q, dim=-1), min=1e-6, max=1 - 1e-6)
-        k = torch.clamp(F.softmax(k, dim=-1), min=1e-6, max=1 - 1e-6)
-
-        g = F.logsigmoid(g) / gate_norm
-        g = torch.clamp(g, min=-gate_norm, max=gate_norm)
-
-        # Convert to float32 for numerical stability and get model dtype
-        q, k, v, g = (x.to(torch.float32).contiguous() for x in (q, k, v, g))
-
-        return q, k, v, g
-
-    def _process_linear_attn(self, q, k, v, g, out_proj, tensor_dtype, kwargs):
-        # Retrieve recurrent state from cache (inference only)
-        if kwargs["use_cache"]:
-            last_state = self.linear_cache[self.layer_idx]
-            recurrent_state = (
-                last_state["recurrent_state"]
-                if last_state is not None and "recurrent_state" in last_state
-                else None
+    def _process_self_attn(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor],
+        kwargs,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[DynamicCache], int]:
+        """Process the self-attention part (with truncation)."""
+        if self.max_self_attn_length:  # Not needed for SWA (nonparam memorize context)
+            hidden_states, attention_mask = truncate_attention_mask(
+                hidden_states, attention_mask, self.max_self_attn_length
             )
-        else:
-            recurrent_state = None
-
-        # Linear attention
-        o_lin, recurrent_state = self.operator(
-            q,
-            k,
-            v,
-            beta=g,
-            chunk_size=self.max_chunk_size,
-            recurrent_state=recurrent_state,
-        )
-        o_lin = rearrange(o_lin, "b h n d -> b n (h d)").to(tensor_dtype)
-        o_lin = out_proj(o_lin)
-        # Ensure stability (o_lin = soft_clamp(o_lin) ?)
-        o_lin = torch.clamp(o_lin, min=-1e4, max=1e4)
-
-        # Save recurrent state
-        if kwargs["use_cache"]:
-            self.linear_cache.update(self.layer_idx, recurrent_state=recurrent_state)
-        return o_lin
-
-    def _process_self_attn(self, hidden_states, attention_mask, kwargs):
-        # If cache_implementation="static" -> truncated attention
-        hidden_states, attention_mask = truncate_attention_mask(
-            hidden_states, attention_mask, self.max_self_attn_length
-        )
-
-        if kwargs.get("position_embeddings", None) is not None:
-            cos, sin = kwargs["position_embeddings"]
-            cos = cos[:, -self.max_self_attn_length :]
-            sin = sin[:, -self.max_self_attn_length :]
-            kwargs["position_embeddings"] = (cos, sin)
-
-        if isinstance(kwargs.get("past_key_value", None), DynamicCache):
-            # cache management
-            if len(kwargs["past_key_value"]) > self.layer_idx and self.layer_idx == 0:
-                kwargs["past_key_value"].crop(self.max_self_attn_length - 1)
 
+            if kwargs.get("position_embeddings", None) is not None:
+                cos, sin = kwargs["position_embeddings"]
+                cos = cos[:, -self.max_self_attn_length :]
+                sin = sin[:, -self.max_self_attn_length :]
+                kwargs["position_embeddings"] = (cos, sin)
+
+            if isinstance(kwargs.get("past_key_value", None), DynamicCache):
+                # cache management
+                if (
+                    len(kwargs["past_key_value"]) > self.layer_idx
+                    and self.layer_idx == 0
+                ):
+                    kwargs["past_key_value"].crop(self.max_self_attn_length - 1)
+
+        # Ensure attention mask is of the correct dtype and device
+        if attention_mask is not None:
+            attention_mask = attention_mask.to(
+                dtype=hidden_states.dtype, device=hidden_states.device
+            )
         # Standard attention (mask and rotation is applied inside)
         base_attn_outputs = self.base_attn(
             hidden_states,
@@ -271,15 +461,67 @@ class LiZAttention(nn.Module):
             o_base = base_attn_outputs
             attn_weights, present_key_value, expected_attn_mode = None, None, 1
         # Ensure stability
-        o_base = torch.clamp(o_base, min=-1e4, max=1e4)
+        o_base = ensure_stability(o_base, min_val=-1e4, max_val=1e4)
         return o_base, attn_weights, present_key_value, expected_attn_mode
 
+    def _prepare_attn_mixin(
+        self,
+        o_lin: torch.Tensor,
+        o_base: torch.Tensor,
+        tensor_dtype: torch.dtype,
+        eps: float = 1e-5,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Prepare linear attn for mixing with self attn."""
+        # Force cast typing, shape : [b n (h d)]
+        o_lin = o_lin.to(tensor_dtype)
+        o_base = o_base.to(tensor_dtype)
+        # feature scaling
+        if self.base_scale_attn:
+            scaler = o_base.pow(2).mean(dim=-1, keepdim=True).add(eps).sqrt()
+            o_lin = scaler * o_lin
+        return o_lin, o_base
+
+    def _apply_mag(
+        self, linear_attention: torch.Tensor, softmax_attention: torch.Tensor
+    ) -> torch.Tensor:
+        """Apply the MAG strategy"""
+        # Left-Padding management
+        if linear_attention.shape[1] != softmax_attention.shape[1]:
+            left_trunc = min(linear_attention.shape[1], softmax_attention.shape[1])
+            linear_attention, softmax_attention = (
+                linear_attention[:, -left_trunc:],
+                softmax_attention[:, -left_trunc:],
+            )
+        # NAM : Neural Attention Mixer (with graph forcing)
+        mag_weight = torch.tensor(
+            self.mag_weight,
+            dtype=softmax_attention.dtype,
+            device=softmax_attention.device,
+        )
+        softmax_weighted = (1 - mag_weight) * softmax_attention
+        linear_weighted = mag_weight * linear_attention
+        if self.cross_gate:
+            output_attention = (
+                softmax_weighted + linear_weighted + softmax_weighted * linear_weighted
+            )  # complex cross product (unlinear interaction)
+        else:
+            output_attention = softmax_weighted + linear_weighted  # classic
+
+        if torch.allclose(softmax_weighted, output_attention):
+            logger.info(
+                "[LOG] layer : %s, softmax_weighted and output_attention are close.",
+                self.layer_idx,
+            )
+        # Final output
+        return ensure_stability(output_attention, min_val=-1e4, max_val=1e4)
+
     def forward(
         self,
         hidden_states: torch.Tensor,
         attention_mask: Optional[torch.Tensor] = None,
         **kwargs,
-    ):
+    ) -> torch.Tensor:
+        """Mix linear and self attention forward"""
         device = hidden_states.device
         tensor_dtype = hidden_states.dtype
         self.base_attn.to(device)
@@ -287,73 +529,155 @@ class LiZAttention(nn.Module):
         if self.training:
             kwargs.pop("past_key_value", None)
             kwargs["use_cache"] = False
-        else:
-            # Force evaluation
-            kwargs["use_cache"] = True
+        elif "use_cache" not in kwargs:
+            kwargs.pop("past_key_value", None)
+            kwargs["use_cache"] = False
 
         kwargs.pop("position_ids", None)  # obsolete
 
-        # Apply projections to hidden states
-        q, k, v, out_proj = self._apply_projections(hidden_states)
-
-        # Manage attention mask (with padding)
-        if attention_mask is not None:
-            # attention_mask -> [batch, seq], v: [batch, seq, ...]
-            v = apply_linear_attention_mask(attention_mask, v)
+        # Apply shared projections
+        q, k, v, out_proj = self._apply_shared_projections(hidden_states)
 
-        # Prepare inputs tensor for linear attn
-        gate_norm = kwargs.get("gate_logit_normalizer", 16)
-        q, k, v, g = self._prepare_attn_input(q, k, v, gate_norm)
-
-        # Process linear attn from mask
-        o_lin = self._process_linear_attn(q, k, v, g, out_proj, tensor_dtype, kwargs)
+        # Apply linear attention to hidden states
+        o_lin = self.linear_attn(
+            x=[q, k, v], attn_mask=attention_mask, out_proj=out_proj, **kwargs
+        )
 
         # Process self attn with truncation
         o_base, attn_weights, present_key_value, expected_attn_mode = (
             self._process_self_attn(hidden_states, attention_mask, kwargs)
         )
 
-        # Force cast typing
-        o_lin = o_lin.to(tensor_dtype)
-        o_base = o_base.to(tensor_dtype)
+        # Prepare output mixing
+        o_lin, o_base = self._prepare_attn_mixin(o_lin, o_base, tensor_dtype, eps=1e-5)
 
-        # Apply Memory as Gate in self-attention (with max length management)
-        if o_lin.shape[1] > o_base.shape[1]:
-            o_padding = torch.zeros_like(o_lin).to(tensor_dtype)
-            o_padding[:, -o_base.shape[1] :] = o_base
-            o_base = o_padding  # Left PAD mask
-        elif o_lin.shape[1] != o_base.shape[1]:  # Abnormality
-            left_trunc = min(o_lin.shape[1], o_base.shape[1])
-            o_lin, o_base = o_lin[:, -left_trunc:], o_base[:, -left_trunc:]
-        out = self.mag_weight * o_lin + (1 - self.mag_weight) * o_base
-        # Ensure stability
-        out = torch.clamp(out, min=-1e4, max=1e4)
+        # Apply Memory as Gate in self-attention (with length management and ablation)
+        out = o_base if self.disable_linear_attn else self._apply_mag(o_lin, o_base)
 
         # Return output following transformer convention
         if expected_attn_mode == 3:
             return out, attn_weights, present_key_value
-        elif expected_attn_mode == 2:
+        if expected_attn_mode == 2:
             return out, attn_weights
+        return out
+
+
+def load_tptt_safetensors(
+    repo_or_path: str,
+    model: Union[PreTrainedModel, PeftModel],
+    subfolder: Optional[str] = None,
+    token: Optional[str] = None,
+) -> Union[PreTrainedModel, PeftModel]:
+    """Load Tptt safetensor from LoRA/PEFT weights and adapt keys if needed."""
+    # sharding not supported yet (e.g. : -00001-of-00005.safetensors, ...)
+    fname = "adapter_model.safetensors"
+    # subfolder management
+    if subfolder:
+        repo_or_path_norm = os.path.normpath(repo_or_path)
+        subfolder_norm = os.path.normpath(subfolder)
+        if not repo_or_path_norm.endswith(subfolder_norm):
+            fname = f"{subfolder}/{fname}" if subfolder else fname
+    # Find file path
+    if os.path.isdir(repo_or_path):
+        path = os.path.join(repo_or_path, fname)
+        if not os.path.exists(path):
+            return model
+    else:
+        if fname not in list_repo_files(repo_or_path, token=token):
+            return model
+        path = hf_hub_download(repo_or_path, fname, token=token)
+
+    # Load weights from safetensors
+    with safe_open(path, framework="pt") as f:
+        state_dict = {k: f.get_tensor(k) for k in f.keys()}
+
+    # Adapt LoRA/Specific keys if needed (add .default if expected by the model)
+    def adapt_keys(sd, model):
+        model_keys = list(model.state_dict().keys())
+        if any(k.startswith("tptt_model.base_model.") for k in model_keys):
+            prefix = "tptt_model.base_model."
+        elif any(k.startswith("base_model.") for k in model_keys):
+            prefix = "base_model."
         else:
-            return out
+            prefix = ""
+
+        has_base_attn = any(".base_attn." in k for k in model_keys)
+
+        def adapt_key(k):
+            k_ = k if k.startswith(prefix) else prefix + k
+            # first, verify and modify base_attn (LiZA)
+            if ".base_attn." in k_ and not has_base_attn:
+                k_ = k_.replace(".base_attn.", ".")
+            # change LoRA if needed
+            if (
+                k_.endswith("lora_A.weight") or k_.endswith("lora_B.weight")
+            ) and k_.replace(".weight", ".default.weight") in model_keys:
+                k_ = k_.replace(".weight", ".default.weight")
+            return k_
+
+        return {adapt_key(k): v for k, v in sd.items()}
+
+    state_dict = adapt_keys(state_dict, model)
+
+    # Cast tensors to the expected dtype of the model parameters
+    model_state_dict = model.state_dict()
+    for k, v in state_dict.items():
+        if k in model_state_dict:
+            expected_dtype = model_state_dict[k].dtype
+            if v.dtype != expected_dtype:
+                state_dict[k] = v.to(expected_dtype)
+
+    logger.info("Input LoRA/Specific keys: %s", [k for k in state_dict.keys()])
+
+    # Load into model
+    missing, unexpected = model.load_state_dict(state_dict, strict=False, assign=True)
+    missing_lora = [k for k in missing if "lora" in k]
+    if missing_lora:
+        logger.warning("Missing keys: %s", missing_lora)
+    if unexpected:
+        logger.warning("Unexpected keys: %s", unexpected)
+    return model
 
 
 def get_tptt_model(  # pylint: disable=too-many-arguments, too-many-positional-arguments
     model: nn.Module,
     base_config: PretrainedConfig,  # ou LlamaConfig, MistralConfig, etc.
-    liza_attention: LiZAttention,
-    target_modules: list,
     linear_cache: Optional[LCache] = None,
+    liza_attention: nn.Module = LiZAttention,
+    target_modules_names: Optional[list[str]] = None,
     operator_mode: str = "delta_rule",
+    use_linear_checkpoint: bool = False,
+    recurrent_config: Optional[Dict[str, Any]] = None,
+    base_scale_attn: bool = False,
     mag_weight: float = 0.5,
+    cross_gate: bool = False,
     max_chunk_size: int = 64,
-    max_self_attn_length: int = 2048,
-):
+    linear_precision: torch.dtype = torch.float32,
+    max_self_attn_length: Optional[int] = None,  # unnecessary
+    padding_side: str = "right",  # for tokenizer
+    bidirectional: bool = False,  # if True, use bidirectional attention
+    pooling_config: Optional[Dict[str, Any]] = None,
+    **kwargs,  # quickfix unexpected arguments
+) -> Tuple[PreTrainedModel, LCache]:
     """Replace target modules in a model with LiZAttention."""
+    if target_modules_names is None:
+        target_modules_names = ["attn", "self_attn", "attention"]
+    # Find target modules by suffix (e.g., "attn", "attention")
+    target_modules_names = [
+        name
+        for name, _ in model.named_modules()
+        if any(name.endswith(suffix) for suffix in target_modules_names)
+        and not any(f".{suffix}." in name for suffix in target_modules_names)
+    ]
+    if not target_modules_names:
+        raise ValueError(
+            f"Target modules '{target_modules_names}' not found in the model."
+        )
+    # Prepare recurrent config
     linear_cache = linear_cache or LCache()
     # Inject LiZAttention into the model
     for name, _ in model.named_modules():
-        if name in target_modules:
+        if name in target_modules_names:
             parent = model
             *path, last = name.split(".")
             for p in path:
@@ -368,14 +692,45 @@ def get_tptt_model(  # pylint: disable=too-many-arguments, too-many-positional-a
                     base_config=base_config,
                     linear_cache=linear_cache,
                     operator_mode=operator_mode,
+                    use_linear_checkpoint=use_linear_checkpoint,
+                    recurrent_config=recurrent_config,
                     max_self_attn_length=max_self_attn_length,
+                    base_scale_attn=base_scale_attn,
                     mag_weight=mag_weight,
+                    cross_gate=cross_gate,
                     max_chunk_size=max_chunk_size,
+                    linear_precision=linear_precision,
+                    padding_side=padding_side,
+                    bidirectional=bidirectional,
+                    pooling_config=pooling_config,
                 ),
             )
     return model, linear_cache
 
 
+def save_tptt_safetensors(model, path: str, name: str = "adapter_model.safetensors"):
+    """Save trainable LoRA/Specific weights and adapting key names"""
+    # 1. Get the full state_dict
+    all_sd = model.state_dict()
+
+    # 2. Identify trainable parameter names (usually only LoRA/PEFT adapters)
+    trainable_keys = [
+        name for name, param in model.named_parameters() if param.requires_grad
+    ]  # Also, you can manually select specific keys in model after load
+
+    # 3. Filter and adapt the keys (Remove custom model encapsulation info)
+    to_save = {
+        k.replace("tptt_model.", "").replace("base_model.", ""): all_sd[k]
+        for k in trainable_keys
+    }
+
+    # 4. Save the filtered adapters to a safetensors file
+    if to_save:
+        os.makedirs(os.path.dirname(path), exist_ok=True)
+        # sharding not supported yet (e.g. : -00001-of-00005.safetensors, ...)
+        save_file(to_save, os.path.join(path, name))
+
+
 class TpttModel(PreTrainedModel):
     """
     TPTT model wrapper with linear attention (LiZA) and LoRA support.
@@ -396,312 +751,562 @@ class TpttModel(PreTrainedModel):
         super().__init__(config, **kwargs)
         repo_or_path = getattr(config, "_base_path", None) or config._name_or_path
 
-        # 1. Load backbone TODO : support no model.safetensors
-        self.backbone = AutoModelForCausalLM.from_pretrained(
-            config.base_model_name, **kwargs
-        )
-        self._retie_lm_after_load(**kwargs)  # Force lm tie weights
+        # 1. Load backbone (with subfolder management) :
+        kwargs_bb = kwargs.copy()
+        if config.base_model_subfolder is not None:
+            kwargs_bb["subfolder"] = config.base_model_subfolder
+        else:
+            kwargs_bb.pop("subfolder", None)
+
+        if config.model_task == "causal_lm":
+            tptt_model = AutoModelForCausalLM.from_pretrained(
+                config.base_model_name, **kwargs_bb
+            )
+        else:
+            tptt_model = AutoModel.from_pretrained(config.base_model_name, **kwargs_bb)
 
         # 2. Inject LiZA attention
         self.linear_cache = LCache()
-        self.backbone, self.linear_cache = self.inject_liza_attention(
-            self.backbone, config, self.linear_cache
+        tptt_model, self.linear_cache = get_tptt_model(
+            tptt_model, config, self.linear_cache, **config.to_dict()
         )
-        # 3. Apply LoRA if present and configured
+
+        # 3. Apply LoRA/Specific if present and configured
         if config.lora_config is not None:
             lora_config_obj = LoraConfig(**config.lora_config)
-            self.backbone = get_peft_model(self.backbone, lora_config_obj)
-            if repo_or_path:
-                self.load_peft_safetensors(
-                    repo_or_path, token=kwargs.get("token", None)
-                )
-
-    def load_peft_safetensors(self, src, token=None):
-        # src: local dir or repo_id
-        fname = "adapter_model.safetensors"
-        if os.path.isdir(src):
-            path = os.path.join(src, fname)
-            if not os.path.exists(path):
-                return
+            tptt_model = get_peft_model(tptt_model, lora_config_obj)
         else:
-            if fname not in list_repo_files(src, token=token):
-                return
-            path = hf_hub_download(src, fname, token=token)
-        with safe_open(path, framework="pt") as f:
-            self.backbone.load_state_dict(
-                {k: f.get_tensor(k) for k in f.keys()}, strict=False
+            # Doesn't work if quantization is applied !
+            tptt_model = set_trainable_parameters(tptt_model)
+
+        # 4. Load safetensor if tptt/peft adaptor in repo
+        if repo_or_path:
+            tptt_model = load_tptt_safetensors(
+                repo_or_path,
+                tptt_model,
+                subfolder=kwargs.get("subfolder", None),
+                token=kwargs.get("token", None),
             )
+        self.tptt_model = tptt_model
 
-    @staticmethod
-    def inject_liza_attention(
-        backbone,
-        config,
-        linear_cache,
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        **kwargs,
     ):
-        """
-        Inject LiZAttention into the specified target modules of the base model.
-        """
-        # Find target modules by suffix (e.g., "attn", "attention")
-        target_modules = [
-            name
-            for name, _ in backbone.named_modules()
-            if any(name.endswith(suffix) for suffix in config.target_modules_names)
-        ]
-        if not target_modules:
-            raise ValueError(
-                f"Target modules '{config.target_modules_names}' not found in the model."
-            )
-        # Inject LiZAttention (external function, not shown here)
-        return get_tptt_model(
-            backbone,
-            base_config=backbone.config,
-            liza_attention=LiZAttention,
-            target_modules=target_modules,
-            linear_cache=linear_cache,
-            operator_mode=config.operator_mode,
-            max_self_attn_length=config.max_self_attn_length,
-            mag_weight=config.mag_weight,
-            max_chunk_size=config.max_chunk_size,
-        )
-
-    def forward(self, input_ids=None, attention_mask=None, labels=None, **kwargs):
-        """
-        Forward pass. All arguments are passed to the underlying base model.
-        """
+        """Forward pass. All arguments are passed to the underlying base model."""
         if self.training:
             kwargs["use_cache"] = False
             kwargs.pop("num_items_in_batch", None)
-        else:
-            kwargs["use_cache"] = True
-        return self.backbone(
+        elif "use_cache" not in kwargs:  # evaluation
+            kwargs.pop("num_items_in_batch", None)
+            kwargs["use_cache"] = False
+        return self.tptt_model(
             input_ids=input_ids, attention_mask=attention_mask, labels=labels, **kwargs
         )
 
     def generate(self, *args, **kwargs):
-        # Delegate the generate call to the backbone model, which supports generation
-        return self.backbone.generate(*args, **kwargs)
+        """Delegate the generate call to the backbone model, which supports generation"""
+        return self.tptt_model.generate(*args, **kwargs)
 
     def save_pretrained(self, path: str, **kwargs):
         """Save model weights, config, and source code to the given path."""
-        super().save_pretrained(path, **kwargs)
-
-        # 1. Save PEFT weights and clean adapter config
-        self._save_peft_weights(path, **kwargs)
+        # 0. Save complete tptt config (with or without LoRA)
+        super().save_pretrained(path, **kwargs)  # pylint: disable=no-member
+        self._adjust_save_strategy(path, **kwargs)
+        # 1. Save true weights and adapte keys
+        save_tptt_safetensors(self, path)
         # 2. Copy Python files for trust_remote_code
-        self._copy_source_files(path)
-
-    def _save_peft_weights(self, path: str, **kwargs):
-        """Save PEFT weights and remove redundant adapter config."""
-        self.backbone.save_pretrained(path, **kwargs)
-        adapter_config_path = os.path.join(path, "adapter_config.json")
-        if os.path.exists(adapter_config_path):
-            os.remove(adapter_config_path)
-
-    def _copy_source_files(self, path: str):
+        self._copy_source_files(path, **kwargs)
+
+    def _adjust_save_strategy(self, path: str, **kwargs):
+        """Re-adapt/remove the weight safetensor and saved adapter config"""
+        if isinstance(self.tptt_model, PeftModel):
+            self.tptt_model.save_pretrained(path, **kwargs)
+        safetensor_path = os.path.join(path, "model.safetensors")
+        if os.path.exists(safetensor_path):
+            os.remove(safetensor_path)
+        adapter_path = os.path.join(path, "adapter_config.json")
+        if os.path.exists(adapter_path):
+            os.remove(adapter_path)
+
+    def _copy_source_files(self, target_path: str, **kwargs):
         """Copy all .py files from package directory for trust_remote_code."""
         src_dir = os.path.dirname(os.path.abspath(__file__))
+        dst_dir = (
+            f"./{str(Path(target_path).parts[0])}"
+            if kwargs.get("subfolder", False)
+            else target_path
+        )
         for fname in os.listdir(src_dir):
             if fname.endswith(".py"):
                 src = os.path.join(src_dir, fname)
-                dst = os.path.join(path, fname)
+                dst = os.path.join(dst_dir, fname)
                 shutil.copy2(src, dst)
 
-    def _retie_lm_after_load(self, **kwargs):
+    def retie_lm_after_load(self, **kwargs):
         """Re-link lm_head after loading external weights."""
-        embed_lm = find_embedding_lm(self.backbone)
-        if embed_lm is not None and hasattr(self.backbone, "lm_head"):
-            if self.backbone.lm_head is None:  # ensure lm_head exists
-                self.backbone.lm_head = nn.Linear(
+        embed_lm = find_embedding_lm(self.tptt_model)
+        if embed_lm is not None and hasattr(self.tptt_model, "lm_head"):
+            if self.tptt_model.lm_head is None:  # ensure lm_head exists
+                self.tptt_model.lm_head = nn.Linear(
                     embed_lm.weight.shape[1], embed_lm.weight.shape[0], bias=False
                 )
             if kwargs.get("tie_word_embeddings", True):
-                self.backbone.lm_head.weight = embed_lm.weight  # share weights
+                self.tptt_model.lm_head.weight = embed_lm.weight  # share weights
                 logger.info("Weights of lm_head have been shared with embedding.")
             else:
-                self.backbone.lm_head.weight = nn.Parameter(embed_lm.weight.clone())
+                self.tptt_model.lm_head.weight = nn.Parameter(embed_lm.weight.clone())
                 logger.info("Weights of lm_head have been cloned from the embedding.")
 
     @classmethod
-    def from_pretrained(cls, *args, **kwargs):
-        model = super().from_pretrained(*args, **kwargs)
-        model._retie_lm_after_load(**kwargs)
+    def from_pretrained(cls, pretrained_model_name_or_path=None, *model_args, **kwargs):
+        """Custom from_pretrained that accepts the standard positional argument"""
+        config = kwargs.pop("config", None)
+        repo_or_path = (
+            pretrained_model_name_or_path
+            or kwargs.pop("pretrained_model_name_or_path", None)
+            or kwargs.pop("repo_or_path", None)
+            or (getattr(config, "_base_path", None) if config else None)
+            or (getattr(config, "_name_or_path", None) if config else None)
+        )
+
+        if config is None and repo_or_path is not None:
+            config = AutoConfig.from_pretrained(repo_or_path, **kwargs)
+        model = cls(config, *model_args, **kwargs)
+        model.retie_lm_after_load(**kwargs)
         return model
 
 
 TpttModel.register_for_auto_class("AutoModelForCausalLM")
 
 
-class AttentionOperator(nn.Module):
+class LinearAttentionOp(nn.Module):
     """Base class for linear attention operators."""
 
-    def __init__(self, mode="delta_rule"):
+    def __init__(
+        self,
+        layer_idx: int,
+        operator_mode: str = "delta_rule",
+        use_linear_checkpoint: bool = False,
+        recurrent_config: Optional[dict] = None,
+        max_chunk_size: int = 64,
+        linear_cache: Optional[LCache] = None,
+        linear_precision: torch.dtype = torch.float32,
+    ):
         super().__init__()
-        self.mode = mode
+        self.layer_idx = layer_idx
+        if recurrent_config is None:
+            operator_mode = "delta_rule"  # force default operator mode if no config
+            recurrent_config = {
+                "order": 1,
+                "gate_type": "k",
+                "linear": True,
+                "trick": "derivative",
+            }
+        self.operator_mode = operator_mode
+        self.use_linear_checkpoint = use_linear_checkpoint
+
+        self.order = recurrent_config["order"]
+        self.gate_type = recurrent_config["gate_type"]
+        self.linear = recurrent_config["linear"]
+        self.trick = recurrent_config["trick"]
+
+        self.max_chunk_size = max_chunk_size
+        self.linear_cache = linear_cache or LCache()
+        self.linear_precision = linear_precision
+
+    def compute_gate(self, beta: Tuple[torch.Tensor]) -> torch.Tensor:
+        """
+        Compute the gating tensor according to the gate_type.
+        """
+        if self.gate_type == "k":
+            return torch.clamp(beta[0], min=1e-6, max=1 - 1e-6)
+        if self.gate_type == "v":
+            return torch.clamp(beta[1], min=1e-6, max=1 - 1e-6)
+        if self.gate_type == "kv":
+            return torch.clamp(beta[0] * beta[1], min=1e-6, max=1 - 1e-6)
+        raise ValueError(f"Unsupported gate_type: {self.gate_type}")
+
+    def get_cache(self, use_cache: bool) -> Tuple[
+        Optional[torch.Tensor],
+        Optional[Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]],
+    ]:
+        """
+        Retrieve recurrent state and qkv buffers from the cache.
+        """
+        if not use_cache:
+            return None, None
+        last_state = self.linear_cache[self.layer_idx]
+        if last_state is not None:
+            recurrent_state = last_state.get("recurrent_state", None)
+            qkv_buffers = last_state.get("qkv", None)
+        else:
+            recurrent_state = None
+            qkv_buffers = None
+        return recurrent_state, qkv_buffers
 
-    def forward(self, q, k, v, **options):
-        """Forward pass for the attention operator."""
-        beta = options.get("beta", None)
-        chunk_size = options.get("chunk_size", 64)
-        scale = options.get("scale", 1)
-        recurrent_state = options.get("recurrent_state", None)
+    def save_cache(
+        self,
+        use_cache: bool,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        gate: torch.Tensor,
+        state: torch.Tensor,
+    ) -> None:
+        """
+        Save the recurrent state and qkv buffers to the cache.
+        """
+        if not use_cache:
+            return
+        if self.order > 1:
+            qkv_buffers = (
+                q[:, :, -(self.order - 1) :, :],
+                k[:, :, -(self.order - 1) :, :],
+                v[:, :, -(self.order - 1) :, :],
+                gate[:, :, -(self.order - 1) :, :],
+            )
+        else:
+            qkv_buffers = None
+        self.linear_cache.update(self.layer_idx, recurrent_state=state, qkv=qkv_buffers)
 
-        if self.mode == "delta_rule":
-            return self.chunk_delta_rule_forward(
-                q, k, v, beta, chunk_size, initial_state=recurrent_state
+    def forward(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        beta: Union[Tuple[torch.Tensor], torch.Tensor],
+        **kwargs,
+    ) -> torch.Tensor:
+        """
+        Forward pass for the attention operator.
+        """
+        # Ensure linear_precision for numerical stability (float32)
+        q, k, v = [x.to(self.linear_precision) for x in (q, k, v)]
+        if isinstance(beta, (tuple, list)):
+            beta = tuple(b.to(self.linear_precision) for b in beta)
+        else:
+            beta = beta.to(self.linear_precision)
+
+        gate = self.compute_gate(beta)
+
+        # Retrieve cache if needed
+        use_cache = kwargs.get("use_cache", False)
+        use_checkpoint = not (use_cache) and self.use_linear_checkpoint
+        recurrent_state, qkvb = self.get_cache(use_cache)
+
+        if qkvb is not None and qkvb[0].shape == q.shape:
+            q = torch.cat([qkvb[0].to(q.device), q], dim=2).to(self.linear_precision)
+            k = torch.cat([qkvb[1].to(q.device), k], dim=2).to(self.linear_precision)
+            v = torch.cat([qkvb[2].to(q.device), v], dim=2).to(self.linear_precision)
+            gate = torch.cat([qkvb[3].to(q.device), gate], dim=2).to(
+                self.linear_precision
             )
-        if self.mode == "gla":
-            return self.gla_forward(q, k, v, beta, scale, initial_state=recurrent_state)
-        raise ValueError(f"Unknown operator mode: {self.mode}")
+
+        output, state = self.chunk_delta_product_forward(
+            q,
+            k,
+            v,
+            gate,
+            self.max_chunk_size,
+            n=self.order,
+            trick=self.trick,
+            linear=self.linear,
+            initial_state=recurrent_state,
+            use_checkpoint=use_checkpoint,
+            linear_precision=self.linear_precision,
+        )
+
+        # Save cache if needed
+        self.save_cache(use_cache, q, k, v, gate, state)
+
+        return output
 
     @staticmethod
-    def chunk_delta_rule_forward(
-        query, key, value, beta, chunk_size, initial_state=None
-    ):
+    def chunk_delta_product_forward(
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        beta_gate: torch.Tensor,
+        chunk_size: int,
+        n: int = 1,
+        trick: str = "derivative",
+        linear: bool = True,
+        initial_state: Optional[torch.Tensor] = None,
+        use_checkpoint: bool = True,
+        linear_precision: torch.dtype = torch.float32,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
         """
-        Implementation of https://arxiv.org/abs/2406.06484
-        query, key, value, beta: [batch, num_heads, seq_len, head_dim]
-        chunk_size: int
-        initial_state: [batch, num_heads, head_dim, head_dim] or None
+        Chunkwise parallel implementation https://arxiv.org/abs/2406.06484
+        For each chunk, processes chunk_size * n_orders steps (virtual tokens) in order.
         """
+
+        # --- Main chunk_delta_product_forward logic ---
+
         batch_size, num_heads, seq_len, head_dim = query.shape
         chunk_size = get_valid_chunk_size(seq_len, chunk_size)
         num_chunks = seq_len // chunk_size
 
-        # Reshape for chunking: [batch, num_heads, num_chunks, chunk_size, head_dim]
-        q_chunks = query.reshape(
-            batch_size, num_heads, num_chunks, chunk_size, head_dim
-        )
-        k_chunks = key.reshape(batch_size, num_heads, num_chunks, chunk_size, head_dim)
-        v_chunks = value.reshape(
-            batch_size, num_heads, num_chunks, chunk_size, head_dim
-        )
-        beta_chunks = beta.reshape(
-            batch_size, num_heads, num_chunks, chunk_size, head_dim
-        )
+        query_n = query if n == 1 else expand_virtual_tokens(query, n, trick)
+        key_n = key if n == 1 else expand_virtual_tokens(key, n, trick)
+        value_n = value if n == 1 else expand_virtual_tokens(value, n, trick)
+        beta_n = beta_gate if n == 1 else expand_virtual_tokens(beta_gate, n, trick)
 
-        # Output buffer
-        output = torch.empty_like(q_chunks)
-        # State: [batch, num_heads, head_dim, head_dim]
-        expect_state_shape = (batch_size, num_heads, head_dim, head_dim)
-        if initial_state is not None and initial_state.shape == expect_state_shape:
-            # Use provided initial state
-            state = initial_state.to(device=query.device, dtype=query.dtype)
+        q_chunks = chunk_sequence(query_n, num_chunks, chunk_size * n)
+        k_chunks = chunk_sequence(key_n, num_chunks, chunk_size * n)
+        v_chunks = chunk_sequence(value_n, num_chunks, chunk_size * n)
+        beta_chunks = chunk_sequence(beta_n, num_chunks, chunk_size * n)
+
+        k_beta = k_chunks * beta_chunks
+        v_beta = v_chunks * beta_chunks
+
+        householder = -(k_beta @ k_chunks.transpose(-2, -1)).tril(-1)
+        householder = ensure_stability(householder, min_val=-1e4, max_val=1e4)
+
+        # size : N = chunk_size * n
+        inv_hh = fast_invert_matrix(householder, dtype=linear_precision)  # [(...),N,N]
+
+        w = ensure_stability(torch.matmul(inv_hh, k_beta), min_val=-1e4, max_val=1e4)
+        u = ensure_stability(torch.matmul(inv_hh, v_beta), min_val=-1e4, max_val=1e4)
+
+        state_shape = (batch_size, num_heads, n, head_dim, head_dim)
+        if initial_state is not None and initial_state.shape == state_shape:
+            state = initial_state.to(device=query.device, dtype=linear_precision)
         else:
-            state = torch.zeros(
-                batch_size,
-                num_heads,
-                head_dim,
-                head_dim,
+            state = torch.full(
+                state_shape,
+                fill_value=1e-6,  # stability if unlinear activation
                 device=query.device,
-                dtype=query.dtype,
+                dtype=linear_precision,
             )
 
-        def process_chunk(q, k, v, b, state):
-            """
-            q, k, v, b: [batch, num_heads, chunk_size, head_dim]
-            state: [batch, num_heads, head_dim, head_dim]
-            Returns: (output_chunk, new_state)
-            """
-            # Clamp to avoid numerical instabilities (not in paper)
-            k = torch.clamp(k, min=-1e4, max=1e4)
-            v = torch.clamp(v, min=-1e4, max=1e4)
-            b = torch.clamp(b, min=1e-6, max=1e4)
-            q = torch.clamp(q, min=-1e4, max=1e4)
-
-            # Eq. (10): β_t * k_t and β_t * v_t
-            k_beta = k * b
-            v_beta = v * b
-
-            # Eq. (11): Lower-triangular matrix T (with -KβK^T off-diagonal, 1 on diagonal)
-            # T = I - tril(KβK^T, -1)
-            t_matrix = -(k_beta @ k.transpose(-2, -1)).tril(-1)
-            t_matrix = torch.clamp(t_matrix, min=-1e4, max=1e4)
-            t_matrix = t_matrix + torch.eye(
-                q.shape[-2], device=q.device, dtype=q.dtype
-            ).unsqueeze(0).unsqueeze(0)
-
-            # Eq. (11): W = T Kβ, U = T Vβ
-            w_matrix = t_matrix @ k_beta
-            w_matrix = torch.clamp(w_matrix, min=-1e4, max=1e4)
-
-            u_matrix = t_matrix @ v_beta
-            u_matrix = torch.clamp(u_matrix, min=-1e4, max=1e4)
-
-            # Eq. (12): u_i = U - W S (S = state)
-            u_i = u_matrix - torch.matmul(w_matrix, state)
-
-            # Eq. (12): inter-chunk output: q S
-            o_inter = torch.matmul(q, state)
-
-            # Eq. (12): intra-chunk attention: tril(q K^T)
-            a_i = (q @ k.transpose(-2, -1)).tril()
-
-            # Eq. (12): intra-chunk output: a_i u_i
-            o_intra = torch.matmul(a_i, u_i)
-
-            # Eq. (12): state update: S_new = S + K^T u_i
-            new_state = state + torch.matmul(k.transpose(-2, -1), u_i)
-            new_state = torch.clamp(new_state, min=-1e4, max=1e4)
-
-            # Eq. (12): output = intra + inter
-            return o_intra + o_inter, new_state
-
-        for chunk_idx in range(num_chunks):
-            q = q_chunks[:, :, chunk_idx]
-            k = k_chunks[:, :, chunk_idx]
-            v = v_chunks[:, :, chunk_idx]
-            b = beta_chunks[:, :, chunk_idx]
-
-            chunk_out, state = process_chunk(q, k, v, b, state)
-            output[:, :, chunk_idx] = chunk_out
-
-        # Reshape back to [batch, num_heads, seq_len, head_dim]
+        output, final_state = sequential_delta_product_scan(
+            q_chunks.to(dtype=linear_precision),
+            w.to(dtype=linear_precision),
+            u.to(dtype=linear_precision),
+            n,
+            linear,
+            chunk_size,
+            state.to(dtype=linear_precision),
+            linear_precision=linear_precision,
+            use_checkpoint=use_checkpoint,
+        )
+
+        idx_last_order = torch.arange(chunk_size, device=output.device) * n + (n - 1)
+        output = output[:, :, :, idx_last_order, :]  # [B, H, num_chunks, chunk_size, D]
         output = output.reshape(batch_size, num_heads, seq_len, head_dim)
-        return output, state
 
-    @staticmethod
-    def gla_forward(q, k, v, beta, scale, initial_state=None):
-        """Forward pass for GLA attention operator."""
-        if fused_chunk_gla is None or fused_recurrent_gla is None:
-            raise RuntimeError("GLA kernels are not available: CUDA required.")
-        if q.shape[-2] > 1:
-            # Training or sequence length > 1
-            return fused_chunk_gla(
-                q,
-                k,
-                v,
-                beta,
-                scale=scale,
-                initial_state=initial_state,
-                output_final_state=True,
+        return output.to(dtype=linear_precision), final_state.to(dtype=linear_precision)
+
+
+def sequential_delta_product_scan(
+    q_chunks: torch.Tensor,
+    w: torch.Tensor,
+    u: torch.Tensor,
+    n_orders: int,
+    linear_activation: bool,
+    current_chunk_size: int,
+    initial_recurrent_state: torch.Tensor,
+    linear_precision: torch.dtype,
+    use_checkpoint: bool,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    DeltaProduct implementation https://arxiv.org/abs/2502.10297
+    Implements the per-token Householder state updates.
+    """
+    batch, head, num_chunks_inner, chunk_n_total, dim = q_chunks.shape
+    output_inner = torch.empty_like(q_chunks)
+    # initial_recurrent_state is H_{last_token_of_prev_chunk, n-1} ([B, H, D, D])
+    h_0_base = initial_recurrent_state[:, :, -1, :, :].clone()
+
+    def process_one_chunk(
+        q_chunk_params: torch.Tensor,
+        w_chunk_params: torch.Tensor,
+        u_chunk_params: torch.Tensor,
+        h_0_base: torch.Tensor,
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        Process a single chunk (with per-token state for n_orders > 1).
+        """
+        o_intra_current_chunk = torch.zeros(
+            batch,
+            head,
+            chunk_n_total,
+            dim,
+            device=q_chunk_params.device,
+            dtype=linear_precision,
+        )
+        o_inter_current_chunk = torch.zeros_like(o_intra_current_chunk)
+        current_accumulated_state_per_token = (
+            h_0_base.unsqueeze(2).expand(-1, -1, current_chunk_size, -1, -1).clone()
+        )  # [B, H, current_chunk_size, D, D]
+
+        for step in range(n_orders):
+            idx_virtual_tokens = (
+                torch.arange(current_chunk_size, device=q_chunk_params.device)
+                * n_orders
+                + step
             )
-        return fused_recurrent_gla(
-            q,
-            k,
-            v,
-            beta,
-            scale=scale,
-            initial_state=initial_state,
-            output_final_state=True,
+            q_s = q_chunk_params[:, :, idx_virtual_tokens, :]
+            w_s = w_chunk_params[:, :, idx_virtual_tokens, :]
+            u_s = u_chunk_params[:, :, idx_virtual_tokens, :]
+
+            state_input_for_this_step = current_accumulated_state_per_token
+
+            ## BLAS/cuBLAS einsum "bhcd,bhcdd->bhcd"
+            k_trans_h_old = (
+                torch.matmul(
+                    w_s.unsqueeze(-2),
+                    state_input_for_this_step,
+                )
+                .squeeze(-2)
+                .to(dtype=linear_precision)
+            )
+
+            u_val = u_s - k_trans_h_old
+
+            o_inter_current_chunk[:, :, idx_virtual_tokens, :] = (
+                torch.matmul(q_s.unsqueeze(-2), state_input_for_this_step)
+                .squeeze(-2)
+                .to(dtype=linear_precision)
+            )
+
+            ## BLAS/cuBLAS einsum "bhcd,bhcd->bhcd"
+            o_intra_current_chunk[:, :, idx_virtual_tokens, :] = (q_s * u_val).to(
+                dtype=linear_precision
+            )
+
+            outer_product_term = torch.matmul(w_s.unsqueeze(-1), u_val.unsqueeze(-2))
+            new_state_i_per_token = state_input_for_this_step + outer_product_term
+            current_accumulated_state_per_token = new_state_i_per_token.to(
+                dtype=linear_precision
+            )
+        # Return all needed for next chunk
+        return (
+            o_intra_current_chunk,
+            o_inter_current_chunk,
+            current_accumulated_state_per_token[:, :, -1, :, :],  # new h_0_base
+        )
+
+    for chunk_idx_inner in range(num_chunks_inner):
+        q_chunk_params = q_chunks[:, :, chunk_idx_inner]
+        w_chunk_params = w[:, :, chunk_idx_inner]
+        u_chunk_params = u[:, :, chunk_idx_inner]
+
+        # Checkpointed call if training
+        call = (
+            partial(checkpoint, use_reentrant=False)
+            if use_checkpoint
+            else lambda f, *a: f(*a)
+        )
+        o_intra, o_inter, h_0_base = call(
+            process_one_chunk,
+            q_chunk_params,
+            w_chunk_params,
+            u_chunk_params,
+            h_0_base,
+        )
+        if not linear_activation:  # unlinear activation between chunks
+            h_0_base = unlinear_activation(h_0_base).to(dtype=linear_precision)
+        output_inner[:, :, chunk_idx_inner] = o_intra + o_inter
+
+    return output_inner, h_0_base
+
+
+def unlinear_activation(x: torch.Tensor, scale: float = 2.0) -> torch.Tensor:
+    """Unlinear activation between chunk"""
+    x_n = x.norm(p=2, dim=-1, keepdim=True) + 1e-6
+    x_gelu = F.gelu(scale * x / x_n, approximate="tanh")  # pylint: disable=not-callable
+    return (x / scale) * x_gelu
+
+
+def chunk_sequence(x: torch.Tensor, num_chunks: int, chunk_size: int) -> torch.Tensor:
+    """Splits [B, H, S, D] to  [B, H, num_chunks, chunk_size, D]"""
+    batch_size, num_heads, _, head_dim = x.shape
+    return x.reshape(batch_size, num_heads, num_chunks, chunk_size, head_dim)
+
+
+def expand_virtual_tokens(
+    x: torch.Tensor, n: int, mode: str = "derivative"
+) -> torch.Tensor:
+    """Expand tokens into 'n' virtual tokens using the selected trick."""
+    batch_size, num_heads, seq_len, head_dim = x.shape
+    device, dtype = x.device, x.dtype
+
+    def derivative_expand(x: torch.Tensor) -> torch.Tensor:
+        """Expand tokens using the derivative trick."""
+        x_pad = torch.cat(
+            [
+                torch.zeros(
+                    batch_size, num_heads, n - 1, head_dim, device=device, dtype=dtype
+                ),
+                x,
+            ],
+            dim=2,
+        )
+        coeffs = torch.tensor(
+            [(-1) ** k * math.comb(n - 1, k) for k in range(n)],
+            device=device,
+            dtype=dtype,
+        )
+        coeffs /= coeffs.norm(p=1)
+        return (
+            (x_pad.unfold(2, n, 1) * coeffs.view(1, 1, 1, 1, n))
+            .flip(-1)
+            .permute(0, 1, 2, 4, 3)
+            .reshape(batch_size, num_heads, seq_len * n, head_dim)
         )
 
+    def rotative_expand(x: torch.Tensor) -> torch.Tensor:
+        """Expand tokens using the rotative trick."""
+        d_parity = head_dim // 2
+        angles = torch.arange(n, device=device, dtype=dtype) * (2 * math.pi / n)
+        cos = torch.cos(angles).view(1, 1, 1, n, 1)
+        sin = torch.sin(angles).view(1, 1, 1, n, 1)
+        if head_dim % 2:
+            x_pairs = x[..., :-1].view(batch_size, num_heads, seq_len, d_parity, 2)
+        else:
+            x_pairs = x.view(batch_size, num_heads, seq_len, d_parity, 2)
+        x_pairs = x_pairs.unsqueeze(3).expand(
+            batch_size, num_heads, seq_len, n, d_parity, 2
+        )
+        x0, x1 = x_pairs[..., 0], x_pairs[..., 1]
+        x0r = x0 * cos - x1 * sin
+        x1r = x0 * sin + x1 * cos
+        rot = torch.stack([x0r, x1r], -1).reshape(
+            batch_size, num_heads, seq_len, n, d_parity * 2
+        )
+        if head_dim % 2:
+            last = (
+                x[..., -1]
+                .unsqueeze(-1)
+                .unsqueeze(3)
+                .expand(batch_size, num_heads, seq_len, n, 1)
+            )
+            rot = torch.cat([rot, last], -1)
+        return rot.reshape(batch_size, num_heads, seq_len * n, head_dim)
 
-def get_attention_operator(mode):
-    """Factory for AttentionOperator."""
-    return AttentionOperator(mode=mode)
+    if mode == "derivative":
+        return derivative_expand(x)
+    if mode == "rotative":
+        return rotative_expand(x)
+    if mode == "combined":
+        return (derivative_expand(x) + rotative_expand(x)) / 2
+    raise ValueError(f"Unknown mode: {mode}")
 
 
 def extract_layer_idx(module_name: str) -> int:
-    """
-    Extract the layer index from a module name string.
-    """
+    """Extract the layer index from a module name string."""
     match = re.search(r"\.(\d+)\.", module_name)
     if match:
         return int(match.group(1))
     return -1
 
 
-def find_embedding_lm(module):
+def find_embedding_lm(module: nn.Module) -> Optional[nn.Module]:
     """Find the embedding weight in a model module."""
     for _, child in module.named_modules():
         if hasattr(child, "embed_tokens") and hasattr(child.embed_tokens, "weight"):
@@ -713,40 +1318,51 @@ def find_embedding_lm(module):
     return None
 
 
-def soft_clamp(x, min_val=-1e4, max_val=1e4):
-    """Differentiable clamping for stability"""
-    scale = (max_val - min_val) / 2
-    center = (max_val + min_val) / 2
-    return torch.tanh((x - center) / scale) * scale + center
-
+def set_trainable_parameters(
+    model: PreTrainedModel, trainable_patterns: List[str] = None
+) -> PreTrainedModel:
+    """Freeze model parameters except trainable_patterns."""
+    if trainable_patterns is None:
+        trainable_patterns = [
+            "q_proj",
+            "k_proj",
+            "v_proj",
+            "o_proj",
+            "qkv_proj",
+            "out_proj",
+            "c_attn",
+            "c_proj",
+            "query",
+            "key",
+            "value",
+        ]
 
-def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
-    """Repeat key/value heads for grouped query attention (GQA)."""
-    return x.repeat_interleave(n_rep, dim=1)
+    for name, param in model.named_parameters():
+        param.requires_grad = any(pattern in name for pattern in trainable_patterns)
 
+    trainable_layers = [n for n, p in model.named_parameters() if p.requires_grad]
+    logger.info("Trainable parameters after freeze: %s", trainable_layers)
+    return model
 
-def split_qkv(base_attn, qkv):
-    """Split the QKV tensor into separate Q, K, and V tensors."""
-    num_q_heads = getattr(base_attn, "num_q_heads", None)
-    num_k_heads = getattr(base_attn, "num_k_heads", None)
-    num_v_heads = getattr(base_attn, "num_v_heads", None)
-    head_dim = getattr(base_attn, "head_dim", None)
-
-    q_len = num_q_heads * head_dim
-    k_len = num_k_heads * head_dim
-    v_len = num_v_heads * head_dim
 
-    q, k, v = torch.split(qkv, [q_len, k_len, v_len], dim=-1)
-    return q, k, v
+def ensure_stability(
+    tensor: torch.Tensor, min_val: float = -1e4, max_val: float = 1e4
+) -> torch.Tensor:
+    """stability forcing"""
+    dtype = tensor.dtype
+    center = (max_val + min_val) / 2
+    tensor = torch.clamp(tensor, min=min_val, max=max_val)
+    tensor = torch.nan_to_num(tensor, nan=center, posinf=max_val, neginf=min_val)
+    return tensor.to(dtype=dtype)
 
 
-def apply_linear_attention_mask(attention_mask, v):
-    # extract (if) padding mask
+def apply_linear_attention_mask(
+    attention_mask: torch.Tensor, v: torch.Tensor, padding_side: str = "right"
+) -> torch.Tensor:
+    """Extract if padding --> [B,S]"""
     if attention_mask.dim() == 4 and attention_mask.shape[1] == 1:
-        # [batch, 1, seq, seq] -> [batch, seq]
         mask = attention_mask.diagonal(dim1=-2, dim2=-1).squeeze(1)
     else:
-        # Squeeze all singleton dims except batch (dim=0)
         mask = attention_mask.squeeze(
             dim=tuple(
                 i
@@ -754,16 +1370,30 @@ def apply_linear_attention_mask(attention_mask, v):
                 if attention_mask.shape[i] == 1
             )
         )
-    # handle left padding : mask is [batch, seq] --> Broadcast to v [batch, seq, (...)]
-    mask = mask[:, -v.shape[-2] :][(...,) + (None,) * (v.dim() - 2)]
+    # Ensure cast to the same dtype as v and convert to binary mask
+    if not (
+        mask.dtype == torch.bool
+        or (
+            mask.dtype in [torch.uint8, torch.int32, torch.int64]
+            and mask.max() <= 1
+            and mask.min() >= 0
+        )
+    ):
+        mask = (mask >= 0).to(v.dtype)  # [-inf, 0, 0, -inf] --> [0, 1, 1, 0]
+    else:
+        mask = mask.to(v.dtype)
+    # mask is [batch, seq] --> Broadcast to v [batch, seq, (...)]
+    if padding_side == "left":
+        mask = mask[:, -v.shape[-2] :][(...,) + (None,) * (v.dim() - 2)]
+    else:  # right padding
+        mask = mask[:, : v.shape[-2]][(...,) + (None,) * (v.dim() - 2)]
     return v * mask
 
 
-def truncate_attention_mask(hidden_states, attention_mask, max_length):
-    """
-    Truncate hidden_states and attention_mask to the last window of size max_length,
-    matching the sequence dimension of hidden_states.
-    """
+def truncate_attention_mask(
+    hidden_states: torch.Tensor, attention_mask: torch.Tensor, max_length: int
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Truncate hidden_states and attention_mask to the last window of size max_length"""
     seq_dim = 1  # convention: (batch, seq, ...)
     seq_len = hidden_states.shape[seq_dim]
     if seq_len > max_length:
@@ -785,22 +1415,58 @@ def truncate_attention_mask(hidden_states, attention_mask, max_length):
     return hidden_states, attention_mask
 
 
+def fast_invert_matrix(
+    tri_tensor: torch.Tensor, dtype: torch.dtype = torch.float32
+) -> torch.Tensor:
+    """Equivalent to vectorized forward substitution applied to the identity matrix."""
+    tri_tensor = tri_tensor.to(dtype=dtype).clone()
+    chunk_size = tri_tensor.shape[-1]
+
+    for i in range(1, chunk_size):
+        tri_tensor[..., i, :i] = tri_tensor[..., i, :i] + (
+            tri_tensor[..., i, :, None].clone() * tri_tensor[..., :, :i].clone()
+        ).sum(-2)
+
+    tri_tensor = tri_tensor + torch.eye(
+        chunk_size, dtype=dtype, device=tri_tensor.device
+    )
+    return tri_tensor.to(dtype=dtype)
+
+
 def get_valid_chunk_size(total_l: int, chunk_size: int) -> int:
-    """
-    Return the largest chunk_size <= chunk_size that divides total_l.
-    If no chunk_size > 1 fits, return 1.
-    """
+    """Return the largest chunk_size <= chunk_size that divides total_l."""
     for c in range(min(chunk_size, total_l), 0, -1):
         if total_l % c == 0:
             return c
     return 1
 
 
+## RARELY
+def split_qkv(
+    base_attn: nn.Module, qkv: torch.Tensor
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Split the QKV tensor into separate Q, K, and V tensors."""
+    num_q_heads = getattr(base_attn, "num_q_heads", None)
+    num_k_heads = getattr(base_attn, "num_k_heads", None)
+    num_v_heads = getattr(base_attn, "num_v_heads", None)
+    head_dim = getattr(base_attn, "head_dim", None)
+
+    if num_q_heads is None or num_k_heads is None or num_v_heads is None:
+        raise ValueError(
+            "Base attention must have num_q_heads, num_k_heads, and num_v_heads defined."
+        )
+
+    q_len = num_q_heads * head_dim
+    k_len = num_k_heads * head_dim
+    v_len = num_v_heads * head_dim
+
+    q, k, v = torch.split(qkv, [q_len, k_len, v_len], dim=-1)
+    return q, k, v
+
+
+## OPTIONAL
 def match_dim(x: torch.Tensor, dim: int, target_size: int) -> torch.Tensor:
-    """
-    Match the size of tensor x along dimension dim to target_size by interpolation
-    or projection.
-    """
+    """Match the size of tensor x along dimension dim to target_size by interpolation"""
     src_size = x.shape[dim]
     if src_size == target_size:
         return x
@@ -815,3 +1481,24 @@ def match_dim(x: torch.Tensor, dim: int, target_size: int) -> torch.Tensor:
         x = F.linear(x, eye)  # pylint: disable=not-callable
     x = torch.moveaxis(x, -1, dim)
     return x
+
+
+def soft_clamp(
+    x: torch.Tensor, min_val: float = 1e-6, max_val: float = 1 - 1e-6
+) -> torch.Tensor:
+    """Differentiable clamping for stability"""
+    dtype = x.dtype
+    scale = (max_val - min_val) / 2
+    center = (max_val + min_val) / 2
+    return (torch.tanh((x - center) / scale) * scale + center).to(dtype=dtype)
+
+
+def describe(x: torch.Tensor, name="tensor") -> None:
+    """Prints the shape, min, max, mean, and std of a tensor."""
+    stats = (x.min(), x.max(), x.mean(), x.std())
+    print(
+        f"{name} shape: {tuple(x.shape)}, "
+        + f"min: {stats[0]:.4g}, max: {stats[1]:.4g}, "
+        + f"mean: {stats[2]:.4g}, std: {stats[3]:.4g}, "
+        + f"dtype: {x.dtype}, device: {x.device}"
+    )

train_tptt.py

@@ -1,20 +1,40 @@
+# pylint: disable=too-many-arguments, too-many-positional-arguments
+
 """
 Author : Fabien FURFARO
 """
 
-from transformers import TrainerCallback
+from typing import Optional, Union
+
+from transformers import PreTrainedModel, TrainerCallback
 
 from .modeling_tptt import LiZAttention
 
 
-class AdjustMaGWeightCallback(TrainerCallback):
-    """TrainerCallback to schedule mag_weight during training."""
+class LiZACallback(TrainerCallback):
+    """
+    TrainerCallback to schedule mag_weight or enable/disable linear attention during training.
+
+    Modes:
+        - "gradual": linear interpolation from initial_weight to final_weight.
+        - "cyclic": alternate between values in weight_list at each step.
+        - "switch": alternately enable/disable linear attention at each step.
+    """
 
     def __init__(
-        self, model, initial_weight=0.01, final_weight=0.5, transition_step=500
+        self,
+        model: PreTrainedModel,
+        mode: str = "gradual",
+        initial_weight: float = 0.0,
+        final_weight: float = 0.5,
+        transition_step: Union[int, tuple, list] = 100,
+        weight_list: Optional[list] = None,
+        switch_period: int = 1,  # period for switching
     ):
         self.model = model
-        # Ensure weights are always float scalars, not tuples/lists
+        self.mode = mode
+
+        # Ensure initial_weight is a float scalar, not tuple/list
         if isinstance(initial_weight, (tuple, list)):
             initial_weight = initial_weight[0]
         if isinstance(final_weight, (tuple, list)):
@@ -22,42 +42,101 @@ class AdjustMaGWeightCallback(TrainerCallback):
         self.initial_weight = float(initial_weight)
         self.final_weight = float(final_weight)
 
-        if isinstance(transition_step, (tuple, list)):
-            transition_step = transition_step[0]
-        self.transition_step = int(transition_step)
+        # Ensure transition_step is an int scalar, not tuple/list
+        self.transition_step = ensure_int(transition_step)
+        if self.mode == "constant":
+            # For constant mode, transition_step is not used
+            self.initial_weight = self.final_weight
+        # For cyclic mode: ensure all weights are float scalars
+        if weight_list is not None:
+            self.weight_list = [
+                float(w[0]) if isinstance(w, (tuple, list)) else float(w)
+                for w in weight_list
+            ]
+        else:
+            self.weight_list = [self.initial_weight, self.final_weight]
+
+        # For switch_alternate mode
+        self.switch_period = int(switch_period)
 
     def on_step_end(self, args, state, control, **kwargs):
         current_step = state.global_step
         transition_step = self.transition_step
 
-        # Ensure both are plain ints (not tuple, list, tensor, numpy, etc.)
-        if isinstance(current_step, (tuple, list)):
-            current_step = current_step[0]
-        if hasattr(current_step, "item"):
-            current_step = int(current_step.item())
-        else:
-            current_step = int(current_step)
+        # Ensure current_step and transition_step are plain ints
+        current_step = ensure_int(current_step)
+        transition_step = ensure_int(transition_step)
 
-        if isinstance(transition_step, (tuple, list)):
-            transition_step = transition_step[0]
-        if hasattr(transition_step, "item"):
-            transition_step = int(transition_step.item())
-        else:
-            transition_step = int(transition_step)
+        # Select mag_weight or enable/disable linear attention according to mode
+        if self.mode == "constant":
+            # Set mag_weight to final_weight for constant mode
+            weight = self.final_weight
+            for _, module in self.model.named_modules():
+                if isinstance(module, LiZAttention):
+                    module.mag_weight = weight
+
+        elif self.mode == "gradual":
+            if current_step <= transition_step:
+                weight = self.initial_weight + (
+                    self.final_weight - self.initial_weight
+                ) * (current_step / transition_step)
+            else:
+                weight = self.final_weight
+            for _, module in self.model.named_modules():
+                if isinstance(module, LiZAttention):
+                    module.mag_weight = weight
 
-        if current_step < transition_step:
-            weight = self.initial_weight + (self.final_weight - self.initial_weight) * (
-                current_step / transition_step
-            )
+        elif self.mode == "cyclic":
+            idx = current_step % len(self.weight_list)
+            weight = self.weight_list[idx]
             for _, module in self.model.named_modules():
                 if isinstance(module, LiZAttention):
                     module.mag_weight = weight
 
+        elif self.mode == "switch":
+            # Alternately enable/disable linear attention every switch_period steps
+            disable = (current_step // self.switch_period) % 2 == 0
+            for _, module in self.model.named_modules():
+                if isinstance(module, LiZAttention):
+                    module.disable_linear_attn = disable
+
+        else:
+            raise ValueError(f"Unknown mode: {self.mode}")
+
     def on_log(self, args, state, control, logs=None, **kwargs):
         mag_weight = None
+        disable_linear_attn = None
+        # Log the current mag_weight and disable_linear_attn
         for _, module in self.model.named_modules():
             if isinstance(module, LiZAttention):
                 mag_weight = getattr(module, "mag_weight", None)
+                disable_linear_attn = getattr(module, "disable_linear_attn", None)
                 break
         if mag_weight is not None and logs is not None:
             logs["mag_weight"] = float(mag_weight)
+        if disable_linear_attn is not None and logs is not None:
+            logs["disable_linear_attn"] = not bool(disable_linear_attn)
+
+
+def ensure_int(value: Union[int, tuple, list]) -> int:
+    """Ensure the value is a plain integer."""
+    if isinstance(value, (tuple, list)):
+        value = int(value[0])
+    if hasattr(value, "item"):
+        value = int(value.item())
+    return value
+
+
+class SaveBestModelCallback(TrainerCallback):
+    """TrainerCallback to save the best model based on evaluation loss."""
+
+    def __init__(self):
+        self.best_metric = float("inf")
+
+    def on_evaluate(self, args, state, control, metrics=None, **kwargs):
+        if metrics is not None and "eval_loss" in metrics:
+            if metrics["eval_loss"] < self.best_metric:
+                self.best_metric = metrics["eval_loss"]
+                control.should_save = True  # Trigger save
+            else:
+                control.should_save = False  # Skip save