Upload model + init tptt code

README.md

@@ -0,0 +1,68 @@
+---
+language: en
+license: apache-2.0
+library_name: transformers
+tags:
+  - tptt
+  - peft
+  - trust_remote_code
+---
+
+# Titans-OpenELM-1_1B
+
+Titanesque version of `apple/OpenELM-1_1B` with parallel linearized attention (TPTT 😊) and PEFT.
+
+## 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
+
+## Usage
+
+
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+model = AutoModelForCausalLM.from_pretrained(
+"ffurfaro/Titans-OpenELM-1_1B",
+trust_remote_code=True
+)
+tokenizer = AutoTokenizer.from_pretrained("ffurfaro/Titans-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))
+
+```
+
+## Training
+
+- **Dataset:** yahma/alpaca-cleaned
+- **Platform:** Kaggle
+- **Hardware:** NVIDIA 2xT4
+- **Batch size:** 3
+- **Epochs:** 2.0
+- **Learning rate (final):** 8.928571428571428e-06
+- **Loss (final):** 1.5266
+- **Training runtime:** 476.5978 sec
+- **Samples per second:** 2.098
+- **Steps per second:** 0.352
+- **Total FLOPs:** 1560788336640000.0
+- **Gradient norm (final):** 0.55694580078125
+
+## Evaluation
+
+- **Metrics:** Training loss only (no eval yet, table soon : PiQA, ARC, Hella, Wino, GSM8K, MMLU)
+- **Results:** Final training loss: 1.5266
+
+
+## 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.
+
+
+---

__init__.py

@@ -0,0 +1,24 @@
+__version__ = "0.1.0"
+
+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
+
+__all__ = [
+    "TpttConfig",
+    "TpttModel",
+    "TpttPipeline",
+    "get_tptt_model",
+    "AdjustMaGWeightCallback",
+    "LCache",
+    "AttentionOperator",
+    "LiZAttention",
+    "generate_model_card",
+]

adapter_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2b0c35dce9003bc09e72dee902ac4f9a50be489a82a4f2fd22ac2991c7a128e1
+size 7147336

config.json

@@ -0,0 +1,54 @@
+{
+  "architectures": [
+    "TpttModel"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_tptt.TpttConfig",
+    "AutoModelForCausalLM": "modeling_tptt.TpttModel"
+  },
+  "base_model_name": "apple/OpenELM-1_1B",
+  "lora_config": {
+    "alpha_pattern": {},
+    "auto_mapping": null,
+    "base_model_name_or_path": null,
+    "bias": "none",
+    "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": [
+      "out_proj",
+      "qkv_proj"
+    ],
+    "task_type": "CAUSAL_LM",
+    "use_dora": false,
+    "use_rslora": false
+  },
+  "mag_weight": 0.5,
+  "max_chunk_size": 64,
+  "max_self_attn_length": 2048,
+  "model_type": "tptt",
+  "operator_mode": "delta_rule",
+  "target_modules_names": [
+    "attn",
+    "self_attn",
+    "attention"
+  ],
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.49.0"
+}

configuration_tptt.py

@@ -0,0 +1,137 @@
+"""
+Author : Fabien FURFARO
+"""
+
+import os
+import re
+from typing import List, Optional, Union
+
+from transformers import AutoConfig, PretrainedConfig
+
+
+def convert_sets_to_lists(obj):
+    if isinstance(obj, set):
+        return list(obj)
+    elif isinstance(obj, dict):
+        return {k: convert_sets_to_lists(v) for k, v in obj.items()}
+    elif isinstance(obj, (list, tuple)):
+        return [convert_sets_to_lists(x) for x in obj]
+    else:
+        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"]
+
+    def __init__(
+        self,
+        base_model_config: Optional[Union[str, dict, PretrainedConfig]] = None,
+        base_model_name: str = "meta-llama/Llama-3.2-1B",
+        name_or_path: Optional[str] = None,
+        target_modules_names: Optional[List[str]] = None,
+        operator_mode: str = "delta_rule",
+        max_self_attn_length: int = 4096,
+        mag_weight: float = 0.5,
+        max_chunk_size: int = 64,
+        lora_config: Optional[dict] = None,  # only serialized accepted
+        **kwargs,
+    ):
+
+        if base_model_config is not None:
+            if isinstance(base_model_config, str):
+                # Load config from Hugging Face Hub or a local path
+                base_model_config = AutoConfig.from_pretrained(
+                    base_model_config
+                ).to_dict()
+            elif isinstance(base_model_config, PretrainedConfig):
+                base_model_config = base_model_config.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._name_or_path = (
+            name_or_path
+            if name_or_path is not None
+            else "Titans-" + base_model_name.split("/", 1)[1]
+        )
+
+        self.target_modules_names = target_modules_names or [
+            "attn",
+            "self_attn",
+            "attention",
+        ]
+        self.operator_mode = operator_mode
+        self.mag_weight = mag_weight
+        self.max_chunk_size = max_chunk_size
+        self.max_self_attn_length = max_self_attn_length
+
+        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)
+
+        super().__init__(**kwargs)
+        # 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
+
+
+TpttConfig.register_for_auto_class()
+
+
+def extract_template_variables(template):
+    return set(re.findall(r"\{([^{}]+)\}", template))
+
+
+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()
+
+    # 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
+
+    variables = flatten_config(config)
+    variables.update(kwargs)
+    variables["model_id"] = os.path.basename(path)
+
+    # Extract variables from template
+    template_vars = extract_template_variables(template)
+
+    # 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))
+
+    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)

generation_config.json

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

model.safetensors

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

modeling_tptt.py

@@ -0,0 +1,778 @@
+"""
+This module implements the TPTT model with linear attention (LiZA) and LoRA support.
+Author : Fabien FURFARO
+"""
+
+import os
+import re
+import shutil
+from typing import Dict, List, Optional
+
+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 safetensors import safe_open
+from torch import nn
+from transformers import AutoModelForCausalLM, DynamicCache, PreTrainedModel
+from transformers.configuration_utils import PretrainedConfig
+
+from .configuration_tptt import TpttConfig
+
+
+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
+
+
+fused_chunk_gla, fused_recurrent_gla = import_fla_ops()  # TODO: add all ops
+
+
+class LCache:
+    """
+    Cache for storing intermediate states of linear attention layers.
+    Supports a sliding window if max_length is set.
+    """
+
+    def __init__(self):
+        """
+        Initialize the cache.
+
+        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
+
+    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
+
+    def update(self, layer_idx: int, **kwargs):
+        """
+        Update the cache for a given layer.
+        If max_length is set, keep only the last max_length tokens in any sequence state.
+        """
+        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)
+        else:
+            self.states[layer_idx].update(detached_kwargs)
+
+    def reset(self):
+        """
+        Reset the cache and token counter.
+        """
+        self.states.clear()
+        self.seen_tokens = 0
+
+
+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,  # Backbone Config
+        linear_cache: LCache = None,
+        operator_mode: str = "delta_rule",
+        max_self_attn_length: int = 2048,
+        mag_weight: float = 0.5,
+        max_chunk_size: int = 64,
+    ):
+        super().__init__()
+        self.base_attn = base_attn
+        self.base_config = base_config
+        self.layer_idx = layer_idx
+        self.max_self_attn_length = max_self_attn_length
+        self.mag_weight = mag_weight
+        self.max_chunk_size = max_chunk_size
+        self.linear_cache = linear_cache or LCache()
+        (
+            self.num_heads,
+            self.head_dim,
+            self.num_key_value_heads,
+            self.num_key_value_groups,
+        ) = 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
+        )
+
+    def _get_attention_parameters(self, base_attn, base_config):
+        """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_config, "head_dim", 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
+        )
+        return (
+            num_heads,
+            head_dim,
+            num_key_value_heads,
+            num_key_value_groups,
+        )
+
+    def _apply_projections(self, hidden_states):
+        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
+        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)
+        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
+            )
+        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)
+
+        # 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 = torch.clamp(o_base, min=-1e4, max=1e4)
+        return o_base, attn_weights, present_key_value, expected_attn_mode
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        **kwargs,
+    ):
+        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
+        else:
+            # Force evaluation
+            kwargs["use_cache"] = True
+
+        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)
+
+        # 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)
+
+        # 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)
+
+        # 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)
+
+        # Return output following transformer convention
+        if expected_attn_mode == 3:
+            return out, attn_weights, present_key_value
+        elif expected_attn_mode == 2:
+            return out, attn_weights
+        else:
+            return out
+
+
+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: LCache = None,
+    operator_mode: str = "delta_rule",
+    mag_weight: float = 0.5,
+    max_chunk_size: int = 64,
+    max_self_attn_length: int = 2048,
+):
+    """Replace target modules in a model with LiZAttention."""
+    linear_cache = linear_cache or LCache()
+    # Inject LiZAttention into the model
+    for name, _ in model.named_modules():
+        if name in target_modules:
+            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,
+                    max_self_attn_length=max_self_attn_length,
+                    mag_weight=mag_weight,
+                    max_chunk_size=max_chunk_size,
+                ),
+            )
+    return model, linear_cache
+
+
+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
+        backbone = AutoModelForCausalLM.from_pretrained(
+            config.base_model_name, **kwargs
+        )
+
+        # 2. Inject LiZA attention
+        self.linear_cache = LCache()
+        self.backbone, self.linear_cache = self.inject_liza_attention(
+            backbone, config, self.linear_cache
+        )
+        # 3. Apply LoRA if present of 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
+        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
+            )
+
+    @staticmethod
+    def inject_liza_attention(
+        backbone,
+        config,
+        linear_cache,
+    ):
+        """
+        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.
+        """
+        if self.training:
+            kwargs["use_cache"] = False
+            kwargs.pop("num_items_in_batch", None)
+        else:
+            kwargs["use_cache"] = True
+        return self.backbone(
+            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)
+
+    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)
+        # 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):
+        """Copy all .py files from package directory for trust_remote_code."""
+        src_dir = os.path.dirname(os.path.abspath(__file__))
+        for fname in os.listdir(src_dir):
+            if fname.endswith(".py"):
+                src = os.path.join(src_dir, fname)
+                dst = os.path.join(path, fname)
+                shutil.copy2(src, dst)
+
+
+TpttModel.register_for_auto_class("AutoModelForCausalLM")
+
+
+class AttentionOperator(nn.Module):
+    """Base class for linear attention operators."""
+
+    def __init__(self, mode="delta_rule"):
+        super().__init__()
+        self.mode = mode
+
+    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)
+
+        if self.mode == "delta_rule":
+            return self.chunk_delta_rule_forward(
+                q, k, v, beta, chunk_size, initial_state=recurrent_state
+            )
+        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}")
+
+    @staticmethod
+    def chunk_delta_rule_forward(
+        query, key, value, beta, chunk_size, initial_state=None
+    ):
+        """
+        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
+        """
+        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
+        )
+
+        # Output buffer
+        output = torch.empty_like(q_chunks)
+        # State: [batch, num_heads, head_dim, head_dim]
+        if initial_state is not None:
+            state = initial_state
+        else:
+            state = torch.zeros(
+                batch_size,
+                num_heads,
+                head_dim,
+                head_dim,
+                device=query.device,
+                dtype=query.dtype,
+            )
+
+        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 = 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 fused_recurrent_gla(
+            q,
+            k,
+            v,
+            beta,
+            scale=scale,
+            initial_state=initial_state,
+            output_final_state=True,
+        )
+
+
+def get_attention_operator(mode):
+    """Factory for AttentionOperator."""
+    return AttentionOperator(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 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 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)
+
+
+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 apply_linear_attention_mask(attention_mask, v):
+    # extract (if) padding mask
+    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
+                for i in range(1, attention_mask.dim())
+                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)]
+    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.
+    """
+    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 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.
+    """
+    for c in range(min(chunk_size, total_l), 0, -1):
+        if total_l % c == 0:
+            return c
+    return 1
+
+
+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.
+    """
+    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

pipeline_tptt.py

@@ -0,0 +1,44 @@
+import torch
+from transformers import Pipeline
+
+
+class TpttPipeline(Pipeline):
+    """Pipeline for TPTT model inference."""
+
+    def __init__(self, model, tokenizer, device=None, **kwargs):
+        """
+        Initialize TpttPipeline.
+        """
+        super().__init__(model=model, tokenizer=tokenizer, device=device, **kwargs)
+
+    def _sanitize_parameters(self, **kwargs):
+        # No special parameter handling for now
+        preprocess_kwargs = {}
+        forward_kwargs = {}
+        postprocess_kwargs = {}
+        return preprocess_kwargs, forward_kwargs, postprocess_kwargs
+
+    def preprocess(self, prompt):
+        # Tokenize the input prompt
+        return self.tokenizer(prompt, return_tensors="pt", truncation=False)
+
+    def _forward(self, model_inputs, **forward_params):
+        # Move tensors to the correct device
+        model_inputs = {k: v.to(self.device) for k, v in model_inputs.items()}
+        # Use generate for text generation
+        with torch.no_grad():
+            output = self.model.generate(
+                **model_inputs,
+                max_new_tokens=forward_params.get("max_new_tokens", 50),
+                do_sample=forward_params.get("do_sample", False),
+                # cache_implementation=forward_params.get("cache_implementation", "static"),
+            )
+        return {"generated_ids": output}
+
+    def postprocess(self, model_outputs):
+        # Decode the generated ids into text
+        generated_ids = model_outputs["generated_ids"]
+        return [
+            {"generated_text": self.tokenizer.decode(ids, skip_special_tokens=True)}
+            for ids in generated_ids
+        ]

runs/Jun10_18-42-03_381bf638251f/events.out.tfevents.1749580925.381bf638251f.35.0

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

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
+  }
+}

tokenizer.model

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

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
+}

train_tptt.py

@@ -0,0 +1,63 @@
+"""
+Author : Fabien FURFARO
+"""
+
+from transformers import TrainerCallback
+
+from .modeling_tptt import LiZAttention
+
+
+class AdjustMaGWeightCallback(TrainerCallback):
+    """TrainerCallback to schedule mag_weight during training."""
+
+    def __init__(
+        self, model, initial_weight=0.01, final_weight=0.5, transition_step=500
+    ):
+        self.model = model
+        # Ensure weights are always float scalars, not tuples/lists
+        if isinstance(initial_weight, (tuple, list)):
+            initial_weight = initial_weight[0]
+        if isinstance(final_weight, (tuple, list)):
+            final_weight = final_weight[0]
+        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)
+
+    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)
+
+        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)
+
+        if current_step < transition_step:
+            weight = self.initial_weight + (self.final_weight - self.initial_weight) * (
+                current_step / transition_step
+            )
+            for _, module in self.model.named_modules():
+                if isinstance(module, LiZAttention):
+                    module.mag_weight = weight
+
+    def on_log(self, args, state, control, logs=None, **kwargs):
+        mag_weight = None
+        for _, module in self.model.named_modules():
+            if isinstance(module, LiZAttention):
+                mag_weight = getattr(module, "mag_weight", None)
+                break
+        if mag_weight is not None and logs is not None:
+            logs["mag_weight"] = float(mag_weight)