rxnn 0.2.72__tar.gz → 0.2.74__tar.gz

{rxnn-0.2.72 → rxnn-0.2.74}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: rxnn
-Version: 0.2.72
+Version: 0.2.74
 Summary: RxNN: Reactive Neural Networks Platform
 License: Apache-2.0
 Keywords: deep-learning,ai,machine-learning

{rxnn-0.2.72 → rxnn-0.2.74}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "rxnn"
-version = "0.2.72"
+version = "0.2.74"
 description = "RxNN: Reactive Neural Networks Platform"
 
 license = "Apache-2.0"

rxnn-0.2.74/src/rxnn/memory/attention.py

@@ -0,0 +1,150 @@
+import torch
+import torch.nn as nn
+from .stm import ShortTermMemory
+
+
+class StmMemoryAttention(nn.Module):
+    def __init__(
+            self,
+            stm: ShortTermMemory,
+            attention_layers: nn.ModuleList,
+            memory_norm_layers: nn.ModuleList,
+            memory_input_norm_layers: nn.ModuleList,
+            residual_gate_layers: nn.ModuleList,
+            debug_mode: bool = False,
+            debug_interval: int = 10,
+            *args,
+            **kwargs
+    ):
+        super(StmMemoryAttention, self).__init__(*args, **kwargs)
+        self.stm = stm
+        self.attention_layers = attention_layers
+        self.memory_norm_layers = memory_norm_layers
+        self.memory_input_norm_layers = memory_input_norm_layers
+        self.residual_gate_layers = residual_gate_layers
+        assert (len(self.attention_layers) == len(self.memory_norm_layers) ==
+                len(self.residual_gate_layers) == len(self.memory_input_norm_layers) ==
+                self.stm.memory.size(0))
+        self.num_layers = len(attention_layers)
+        self.debug_mode = debug_mode
+        self.debug_interval = debug_interval
+        self.debug_step = 0
+
+    def update_max_len(self, max_seq_len: int):
+        for i in range(self.num_layers):
+            if self.attention_layers[i].rope is not None:
+                self.attention_layers[i].rope.update_max_len(max_seq_len)
+
+    def forward(self, x: torch.Tensor, attention_mask: torch.Tensor = None) -> torch.Tensor:
+        # 1. Process correct attention mask
+        if attention_mask is not None:
+            attention_mask = attention_mask.unsqueeze(1).unsqueeze(1).bool()
+        # 2. Init new empty STM
+        new_stm = torch.zeros_like(self.stm.memory)
+
+        # 3. Run Short-Term Memory update for all layers
+        for i in range(self.num_layers):
+            # 4. Get current layer STM value
+            layer_stm = self.stm(i)
+            # 5. Expand layer STM to batch size, if it's not in batch mode
+            if layer_stm.size(0) == 1:
+                layer_stm = layer_stm.expand(x.size(0), -1, -1)
+
+            # 6. Get encoded layer data and normalize it
+            encoded_layer_data = self.memory_input_norm_layers[i](x[i])
+            # 7. Normalize STM layer
+            normalized_layer_stm = self.memory_norm_layers[i](layer_stm)
+
+            # 8. Print normalization stats in debug mode
+            if self.debug_mode and self.training:
+                if self.debug_step != 0 and self.debug_step % self.debug_interval == 0:
+                    self.debug_step = 0
+                    print(f"Normalized STM stats - mean: {normalized_layer_stm.mean().item():.4f}, std: {normalized_layer_stm.std().item():.4f}")
+                else:
+                    self.debug_step += 1
+
+            # 9. Calculate memory attention
+            new_layer_stm = self.attention_layers[i](normalized_layer_stm, encoded_layer_data, encoded_layer_data, mask=attention_mask)
+            # 10. Combine new updated layer state with current STM state in residual gate
+            new_stm[i] = self.residual_gate_layers[i](layer_stm, new_layer_stm) # residual
+        # 11. Update all layers/models
+        self.stm.update_all(new_stm)
+        return self.stm.memory
+
+
+class InterlayerStmMemoryAttention(StmMemoryAttention):
+    def __init__(
+            self,
+            stm: ShortTermMemory,
+            attention_layers: nn.ModuleList,
+            memory_norm_layers: nn.ModuleList,
+            memory_input_norm_layers: nn.ModuleList,
+            residual_gate_layers: nn.ModuleList,
+            mean_attention_layers: nn.ModuleList,
+            mean_memory_norm_layers: nn.ModuleList,
+            mean_residual_gate_layers: nn.ModuleList,
+            mean_stm_norm: nn.Module,
+            debug_mode: bool = False,
+            debug_interval: int = 10,
+            **kwargs
+    ):
+        super(InterlayerStmMemoryAttention, self).__init__(
+            stm, attention_layers, memory_norm_layers, memory_input_norm_layers, residual_gate_layers,
+            debug_mode=debug_mode, debug_interval=debug_interval, **kwargs
+        )
+        self.mean_attention_layers = mean_attention_layers
+        self.mean_memory_norm_layers = mean_memory_norm_layers
+        self.mean_stm_norm = mean_stm_norm
+        self.mean_residual_gate_layers = mean_residual_gate_layers
+        assert (len(self.mean_attention_layers) == len(self.mean_memory_norm_layers) ==
+                len(self.mean_residual_gate_layers) == self.num_layers)
+
+    def forward(self, x: torch.Tensor, attention_mask: torch.Tensor = None) -> torch.Tensor:
+        # 1. Process correct attention mask
+        if attention_mask is not None:
+            attention_mask = attention_mask.unsqueeze(1).unsqueeze(1).bool()
+        # 2. Init new empty STM
+        new_stm = torch.zeros_like(self.stm.memory)
+
+        # 3. Get mean STM value from layers for mean interlayer memory attention
+        mean_stm = self.stm.memory.mean(dim=0) # [batch_size, stm_size, embed_dim]
+        # 4. Normalize mean STM layer
+        normalized_mean_stm = self.mean_stm_norm(mean_stm)
+
+        # 5. Run Short-Term Memory update for all layers
+        for i in range(self.num_layers):
+            # 6. Get current layer STM value
+            layer_stm = self.stm(i)
+            # 7. Expand layer STM to batch size, if it's not in batch mode
+            if layer_stm.size(0) == 1:
+                layer_stm = layer_stm.expand(x.size(0), -1, -1)
+
+            # 8. Mean interlayer memory attention
+            # a) normalize STM layer value
+            pre_normalized_layer_stm = self.mean_memory_norm_layers[i](layer_stm)
+            # b) calculate attention between STM layer and mean value of all STM layers (from previous interaction)
+            interlayer_stm = self.mean_attention_layers[i](pre_normalized_layer_stm, normalized_mean_stm, normalized_mean_stm, mask=None)
+            # c) combine updated interlayer state with current STM state in residual gate
+            updated_layer_stm = self.mean_residual_gate_layers[i](layer_stm, interlayer_stm)
+
+            # 9. Main memory attention
+            # a) get encoded data for current layer and normalize it
+            encoded_layer_data = self.memory_input_norm_layers[i](x[i])
+            # b) normalize STM layer value
+            normalized_layer_stm = self.memory_norm_layers[i](updated_layer_stm)
+            # c) print normalized STM stats in debug mode
+            if self.debug_mode and self.training:
+                if self.debug_step != 0 and self.debug_step % self.debug_interval == 0:
+                    self.debug_step = 0
+                    print(
+                        f"Pre-Normalized STM stats - mean: {pre_normalized_layer_stm.mean().item():.4f}, std: {pre_normalized_layer_stm.std().item():.4f}")
+                    print(f"Normalized STM stats - mean: {normalized_layer_stm.mean().item():.4f}, std: {normalized_layer_stm.std().item():.4f}")
+                else:
+                    self.debug_step += 1
+            # d) calculate memory attention between STM layer and encoded data
+            new_layer_stm = self.attention_layers[i](normalized_layer_stm, encoded_layer_data, encoded_layer_data, mask=attention_mask)
+            # e) combine new updated layer STM with previous state in residual gate
+            new_stm[i] = self.residual_gate_layers[i](updated_layer_stm, new_layer_stm) # residual
+        # 10. Update all layers/models
+        self.stm.update_all(new_stm)
+        return self.stm.memory

rxnn-0.2.74/src/rxnn/memory/gate.py

@@ -0,0 +1,60 @@
+import torch
+import torch.nn as nn
+from typing import TypeAlias, Literal
+
+ResidualGateType: TypeAlias = Literal['static', 'elementwise', 'linear']
+
+
+class ResidualGate(nn.Module):
+    def __init__(
+            self,
+            stm_size: int,
+            use_gate: bool = False,
+            gate_type: ResidualGateType = 'static',
+            per_slot_gate: bool = True,
+            init_gate: float = 0.0,
+            use_tanh_gate: bool = True,
+            **kwargs,
+    ):
+        super(ResidualGate, self).__init__(**kwargs)
+        self.use_gate = use_gate
+        self.per_slot_gate = per_slot_gate
+        self.gate_type = gate_type
+        self.use_tanh_gate = use_tanh_gate
+
+        if self.use_gate:
+            if self.gate_type == 'linear':
+                self.gate = nn.Linear(stm_size, stm_size if self.per_slot_gate else 1)
+            else:
+                gate_shape = (stm_size, 1) if self.per_slot_gate else (1,)
+                self.gate = nn.Parameter(torch.full(gate_shape, init_gate))
+        else:
+            self.gate = None
+
+        self.gate_activation = nn.Tanh() if self.use_tanh_gate else nn.Sigmoid()
+
+    def _dynamic_gate(self, old_value: torch.Tensor, new_value: torch.Tensor):
+        if self.gate_type == 'linear':
+            mean_residual = (new_value + old_value).mean(dim=-1)
+            gate_input = self.gate(mean_residual).unsqueeze(-1)
+        else:
+            mean_dim = -1 if self.per_slot_gate else [1, 2]
+            gate_input = self.gate * (new_value + old_value).mean(dim=mean_dim, keepdim=True)
+        return self.gate_activation(gate_input)
+
+    def _calculate_output(self, layer_gate: torch.Tensor, old_value: torch.Tensor, new_value: torch.Tensor) -> torch.Tensor:
+        if self.use_tanh_gate:
+            return (1 + layer_gate) * new_value + (1 - layer_gate) * old_value
+        else:
+            return layer_gate * new_value + (1 - layer_gate) * old_value
+
+    def forward(self, old_value: torch.Tensor, new_value: torch.Tensor) -> torch.Tensor:
+        if not self.use_gate:
+            return new_value + old_value
+
+        if self.gate_type == 'static':
+            layer_gate = self.gate_activation(self.gate)
+        else:
+            layer_gate = self._dynamic_gate(old_value, new_value)
+
+        return self._calculate_output(layer_gate, old_value, new_value)

{rxnn-0.2.72 → rxnn-0.2.74}/src/rxnn/rxt/models.py

@@ -9,7 +9,8 @@ from ..transformers.models import ReactiveTransformerBase, ReactiveTransformerEn
 from ..transformers.ff import get_activation_layer
 from ..memory.stm import ShortTermMemory
 from ..memory.norm import init_memory_norm
-from ..memory.attention import StmMemoryAttention
+from ..memory.attention import StmMemoryAttention, InterlayerStmMemoryAttention
+from ..memory.gate import ResidualGate, ResidualGateType
 from ..utils import get_model_size
 from ..experimental.attention import init_experimental_attention
 
@@ -260,15 +261,15 @@ class RxTAlphaMemoryAttention(nn.Module, PyTorchModelHubMixin, license="apache-2
             att_experts: int = None,
             att_query_experts: int = None,
             att_query_groups: int = None,
-            norm_type: str = 'rms',
+            norm_type: str = 'classic-rms',
             norm_init_gate: float = -2.0,
             norm_per_dim_scale: bool = False,
             norm_decay: float = 0.9,
             use_gated_residual: bool = False,
-            residual_per_slot_gate: bool = False,
-            residual_init_gate: float = 0.0,
-            use_dynamic_residual_gate: bool = False,
-            use_tanh_residual_gate: bool = False,
+            residual_per_slot_gate: bool = True,
+            residual_gate_init: float = 3.0,
+            residual_gate_type: ResidualGateType = 'static',
+            use_tanh_residual_gate: bool = True,
             debug_mode: bool = False,
             debug_interval: int = 10,
             **kwargs,
@@ -296,12 +297,19 @@ class RxTAlphaMemoryAttention(nn.Module, PyTorchModelHubMixin, license="apache-2
         memory_norm_layers = nn.ModuleList([init_memory_norm(norm_type, embed_dim, stm_size, decay=norm_decay,
                                                              init_gate=norm_init_gate, per_dim_scale=norm_per_dim_scale)
                                             for _ in range(num_layers)])
+        memory_input_norm_layers = nn.ModuleList(nn.RMSNorm(embed_dim) for _ in range(num_layers))
         attention_layers = nn.ModuleList([att_init() for _ in range(num_layers)])
+        residual_gates = nn.ModuleList([
+            ResidualGate(
+                stm_size, use_gate=use_gated_residual, gate_type=residual_gate_type, per_slot_gate=residual_per_slot_gate,
+                init_gate=residual_gate_init, use_tanh_gate=use_tanh_residual_gate
+            ) for _ in range(num_layers)
+        ])
+
         self.model = StmMemoryAttention(
             stm, attention_layers, memory_norm_layers,
-            use_gated_residual=use_gated_residual, per_slot_gate=residual_per_slot_gate,
-            init_gate=residual_init_gate, use_dynamic_gate=use_dynamic_residual_gate,
-            use_tanh_gate=use_tanh_residual_gate, debug_mode=debug_mode, debug_interval=debug_interval,
+            memory_input_norm_layers, residual_gates,
+            debug_mode=debug_mode, debug_interval=debug_interval,
         )
 
     def freeze(self):
@@ -327,6 +335,141 @@ class RxTAlphaMemoryAttention(nn.Module, PyTorchModelHubMixin, license="apache-2
     def forward(self, x: torch.Tensor, attention_mask: torch.Tensor = None) -> torch.Tensor:
         return self.model(x, attention_mask=attention_mask)
 
+
+class RxTAlphaInterlayerMemoryAttention(nn.Module, PyTorchModelHubMixin, license="apache-2.0"):
+    """RxT-Alpha (Reactive Transformer) memory attention model with interlayer STM attention"""
+
+    def __init__(
+            self,
+            num_layers: int = 12,
+            embed_dim: int = 512,
+            att_heads: int = 16,
+            seq_len: int = 1024,
+            stm_size: int = 1024,
+            use_flash_attention: bool = False,
+            att_dropout: float = 0.0,
+            att_groups: int = 1,
+            att_type: str = 'sqa',
+            att_experts: int = None,
+            att_query_experts: int = None,
+            att_query_groups: int = None,
+            interlayer_att_dropout: float = 0.0,
+            interlayer_att_groups: int = 1,
+            interlayer_att_type: str = 'sqa',
+            interlayer_att_experts: int = None,
+            interlayer_att_query_experts: int = None,
+            interlayer_att_query_groups: int = None,
+            norm_type: str = 'classic-rms',
+            norm_init_gate: float = -2.0,
+            norm_per_dim_scale: bool = False,
+            norm_decay: float = 0.9,
+            use_gated_residual: bool = False,
+            residual_per_slot_gate: bool = True,
+            residual_gate_init: float = 3.0,
+            residual_gate_type: ResidualGateType = 'static',
+            use_tanh_residual_gate: bool = True,
+            debug_mode: bool = False,
+            debug_interval: int = 10,
+            **kwargs,
+    ):
+        super(RxTAlphaInterlayerMemoryAttention, self).__init__(**kwargs)
+
+        assert att_type in ['mha', 'gqa', 'mqa', 'gma', 'dma',
+                            'sqa'], 'Memory attention type could be "mha", "gqa", "mqa", "gma", "dma", "sqa".'
+
+        rope = RotaryPositionalEmbedding(embed_dim // att_heads, seq_len)
+        stm = ShortTermMemory(num_layers, embed_dim, stm_size)
+
+        if att_type in ['mha', 'gqa', 'mqa']:
+            att_init = lambda: init_attention(
+                embed_dim, att_heads, att_type, att_groups, rope=rope,
+                use_flash_attention=use_flash_attention, dropout=att_dropout,
+                max_seq_len=seq_len, is_causal=False, rope_only_for_keys=True
+            )
+        else:
+            att_init = lambda: init_experimental_attention(
+                embed_dim, att_heads, att_type, att_groups, rope=rope,
+                use_flash_attention=use_flash_attention, dropout=att_dropout,
+                max_seq_len=seq_len, is_causal=False, num_experts=att_experts,
+                num_query_experts=att_query_experts, num_query_groups=att_query_groups,
+                rope_only_for_keys=True
+            )
+
+        memory_norm_layers = nn.ModuleList([init_memory_norm(norm_type, embed_dim, stm_size, decay=norm_decay,
+                                                             init_gate=norm_init_gate, per_dim_scale=norm_per_dim_scale)
+                                            for _ in range(num_layers)])
+        memory_input_norm_layers = nn.ModuleList(nn.RMSNorm(embed_dim) for _ in range(num_layers))
+        attention_layers = nn.ModuleList([att_init() for _ in range(num_layers)])
+        residual_gates = nn.ModuleList([
+            ResidualGate(
+                stm_size, use_gate=use_gated_residual, gate_type=residual_gate_type,
+                per_slot_gate=residual_per_slot_gate,
+                init_gate=residual_gate_init, use_tanh_gate=use_tanh_residual_gate
+            ) for _ in range(num_layers)
+        ])
+
+        # Interlayer attention
+        if interlayer_att_type in ['mha', 'gqa', 'mqa']:
+            interlayer_att_init = lambda: init_attention(
+                embed_dim, att_heads, interlayer_att_type, interlayer_att_groups, rope=None,
+                use_flash_attention=use_flash_attention, dropout=interlayer_att_dropout, is_causal=False
+            )
+        else:
+            interlayer_att_init = lambda: init_experimental_attention(
+                embed_dim, att_heads, interlayer_att_type, interlayer_att_groups, rope=None,
+                use_flash_attention=use_flash_attention, dropout=interlayer_att_dropout, is_causal=False,
+                num_experts=interlayer_att_experts, num_query_experts=interlayer_att_query_experts, num_query_groups=interlayer_att_query_groups
+            )
+
+        mean_attention_layers = nn.ModuleList([interlayer_att_init() for _ in range(num_layers)])
+
+        mean_stm_norm = init_memory_norm(
+            norm_type, embed_dim, stm_size, decay=norm_decay,
+            init_gate=norm_init_gate, per_dim_scale=norm_per_dim_scale
+        )
+
+        mean_memory_norm_layers = nn.ModuleList([init_memory_norm(norm_type, embed_dim, stm_size, decay=norm_decay,
+                                                             init_gate=norm_init_gate, per_dim_scale=norm_per_dim_scale)
+                                            for _ in range(num_layers)])
+
+        mean_residual_gates = nn.ModuleList([
+            ResidualGate(
+                stm_size, use_gate=use_gated_residual, gate_type=residual_gate_type,
+                per_slot_gate=residual_per_slot_gate,
+                init_gate=residual_gate_init, use_tanh_gate=use_tanh_residual_gate
+            ) for _ in range(num_layers)
+        ])
+
+        self.model = InterlayerStmMemoryAttention(
+            stm, attention_layers, memory_norm_layers, memory_input_norm_layers, residual_gates,
+            mean_attention_layers, mean_memory_norm_layers, mean_residual_gates, mean_stm_norm,
+            debug_mode=debug_mode, debug_interval=debug_interval,
+        )
+
+    def freeze(self):
+        for param in self.parameters():
+            param.requires_grad = False
+
+    def unfreeze(self):
+        for param in self.parameters():
+            param.requires_grad = True
+
+    def load_shared_memory(self, stm: ShortTermMemory):
+        self.model.stm = stm
+
+    def update_max_len(self, max_seq_len: int):
+        self.model.update_max_len(max_seq_len)
+
+    def reset_memory(self, init_type: str = None):
+        self.model.stm.reset(init_type)
+
+    def clone_reset_memory(self):
+        self.model.stm.clone_detach_reset()
+
+    def forward(self, x: torch.Tensor, attention_mask: torch.Tensor = None) -> torch.Tensor:
+        return self.model(x, attention_mask=attention_mask)
+
+
 class RxTAlphaCriticEncoder(RxTAlphaComponentBase, pipeline_tag="text-classification", license="apache-2.0"):
     """RxT-Alpha (Reactive Transformer) encoder model"""
 

{rxnn-0.2.72 → rxnn-0.2.74}/src/rxnn/transformers/layers.py

@@ -49,10 +49,12 @@ class ReactiveTransformerLayer(nn.Module):
             self.norm1 = nn.RMSNorm(embed_dim)
             self.norm2 = nn.RMSNorm(embed_dim)
             self.norm3 = nn.RMSNorm(embed_dim)
+            self.stm_norm = nn.RMSNorm(embed_dim)
         else:
             self.norm1 = nn.LayerNorm(embed_dim)
             self.norm2 = nn.LayerNorm(embed_dim)
             self.norm3 = nn.LayerNorm(embed_dim)
+            self.stm_norm = nn.LayerNorm(embed_dim)
         self.use_post_norm = use_post_norm
         self.use_moe = use_moe
         self.use_moe_att = use_moe_att
@@ -63,9 +65,11 @@ class ReactiveTransformerLayer(nn.Module):
         if with_norms:
             for param in self.norm2.parameters():
                 param.requires_grad_(is_trainable)
+            for param in self.stm_norm.parameters():
+                param.requires_grad_(is_trainable)
 
     def memory_parameters(self) -> list[nn.Parameter]:
-        return list(self.memory_cross_attention.parameters()) + list(self.norm2.parameters())
+        return list(self.memory_cross_attention.parameters()) + list(self.norm2.parameters()) + list(self.stm_norm.parameters())
 
     def not_memory_parameters(self) -> list[nn.Parameter]:
         return (list(self.attention.parameters()) + list(self.norm1.parameters()) +
@@ -102,11 +106,8 @@ class ReactiveTransformerLayer(nn.Module):
         residual = x
         if not self.use_post_norm:
             x = self.norm1(x)
-            if torch.isnan(x).any():
-                print("!!!!!!!!!!!!!!!!!!!!!!         !!!!!!!!!!!!!!!!!!!!!!         NaN detected in pre-norm (self-attention) output")
         x = self.attention(x, x, x, mask=mask)
-        if torch.isnan(x).any():
-            print("!!!!!!!!!!!!!!!!!!!!!!         !!!!!!!!!!!!!!!!!!!!!!         NaN detected in self-attention output")
+
         x = residual + x
         if self.use_post_norm:
             x = self.norm1(x)
@@ -114,18 +115,13 @@ class ReactiveTransformerLayer(nn.Module):
         residual = x
         if not self.use_post_norm:
             x = self.norm2(x)
-            if torch.isnan(x).any():
-                print("!!!!!!!!!!!!!!!!!!!!!!         NaN detected in pre-norm (cross-attention) output")
 
+        # normalize STM and prepare STM mask
+        stm = self.stm_norm(stm)
         mem_mask = mask.squeeze(1).unsqueeze(-1).expand(-1, -1, -1, stm.size(1)) \
             if mask is not None else None
 
-        if torch.isnan(stm).any():
-            print("!!!!!!!!!!!!!!!!!!!!!!         NaN detected in STM cross-attention input")
-
         x = self.memory_cross_attention(x, stm, stm, mask=mem_mask)
-        if torch.isnan(x).any():
-            print("!!!!!!!!!!!!!!!!!!!!!!         NaN detected in cross-attention output")
         x = residual + x
         if self.use_post_norm:
             x = self.norm2(x)
@@ -134,11 +130,7 @@ class ReactiveTransformerLayer(nn.Module):
         residual = x
         if not self.use_post_norm:
             x = self.norm3(x)
-            if torch.isnan(x).any():
-                print("!!!!!!!!!!!!!!!!!!!!!!         NaN detected in pre-norm (ff) output")
         x = self.ff(x)
-        if torch.isnan(x).any():
-            print("!!!!!!!!!!!!!!!!!!!!!!         NaN detected in ff output")
         x = residual + x
         if self.use_post_norm:
             x = self.norm3(x)

{rxnn-0.2.72 → rxnn-0.2.74}/src/rxnn/utils.py

@@ -1,11 +1,6 @@
 import random, gc
-from typing import Optional, Union, List, Dict, Any
-
 import torch
 import numpy as np
-from huggingface_hub import PyTorchModelHubMixin
-from huggingface_hub.hub_mixin import DataclassInstance
-
 
 def human_format(num: int):
     """Format numbers to human-readable format."""

rxnn-0.2.72/src/rxnn/memory/attention.py

@@ -1,89 +0,0 @@
-import torch
-import torch.nn as nn
-from .stm import ShortTermMemory
-
-class StmMemoryAttention(nn.Module):
-    def __init__(
-            self,
-            stm: ShortTermMemory,
-            attention_layers: nn.ModuleList,
-            memory_norm_layers: nn.ModuleList,
-            use_gated_residual: bool = False,
-            per_slot_gate: bool = False,
-            init_gate: float = 0.0,
-            use_dynamic_gate: bool = False,
-            use_tanh_gate: bool = False,
-            debug_mode: bool = False,
-            debug_interval: int = 10,
-            *args,
-            **kwargs
-    ):
-        super(StmMemoryAttention, self).__init__(*args, **kwargs)
-        self.stm = stm
-        self.attention_layers = attention_layers
-        self.memory_norm_layers = memory_norm_layers
-        assert len(self.attention_layers) == len(self.memory_norm_layers) == self.stm.memory.size(0)
-        self.num_layers = len(attention_layers)
-        self.use_gated_residual = use_gated_residual
-        self.per_slot_gate = per_slot_gate
-        self.use_dynamic_gate = use_dynamic_gate
-        self.use_tanh_gate = use_tanh_gate
-        if self.use_gated_residual:
-            gate_shape = (self.num_layers, self.stm.stm_size, 1) if self.per_slot_gate else (self.num_layers,)
-            self.gate = nn.Parameter(torch.full(gate_shape, init_gate))
-
-        self.debug_mode = debug_mode
-        self.debug_interval = debug_interval
-        self.debug_step = 0
-
-    def update_max_len(self, max_seq_len: int):
-        for i in range(self.num_layers):
-            if self.attention_layers[i].rope is not None:
-                self.attention_layers[i].rope.update_max_len(max_seq_len)
-
-    def _residual_gate(self, gate: torch.Tensor, layer_stm: torch.Tensor, new_layer_stm: torch.Tensor) -> torch.Tensor:
-        if self.use_dynamic_gate:
-            mean_dim = -1 if self.per_slot_gate else [1, 2]
-            gate_input = gate * (new_layer_stm + layer_stm).mean(dim=mean_dim, keepdim=True)
-            layer_gate = torch.tanh(gate_input) if self.use_tanh_gate else torch.sigmoid(gate_input)
-        else:
-            layer_gate = torch.tanh(gate) if self.use_tanh_gate else torch.sigmoid(gate)
-        if self.use_tanh_gate:
-            return (1 + layer_gate) * new_layer_stm + (1 - layer_gate) * layer_stm
-        else:
-            return layer_gate * new_layer_stm + (1 - layer_gate) * layer_stm
-
-    def forward(self, x: torch.Tensor, attention_mask: torch.Tensor = None) -> torch.Tensor:
-        if attention_mask is not None:
-            attention_mask = attention_mask.unsqueeze(1).unsqueeze(1).bool()
-        new_stm = torch.zeros_like(self.stm.memory)
-        for i in range(self.num_layers):
-            layer_stm = self.stm(i)
-            # expand layer STM to batch size, if it's not in batch mode
-            if layer_stm.size(0) == 1:
-                layer_stm = layer_stm.expand(x.size(0), -1, -1)
-            encoded_layer_data = x[i]
-            normalized_layer_stm = self.memory_norm_layers[i](layer_stm)
-            if torch.isnan(normalized_layer_stm).any():
-                print(f"!!!!!!!!!!!!!!!!!!!!!!         NaN detected in {i} layer memory norm output")
-
-            if self.debug_mode and self.training:
-                if self.debug_step != 0 and self.debug_step % self.debug_interval == 0:
-                    self.debug_step = 0
-                    print(f"Normalized STM stats - mean: {normalized_layer_stm.mean().item():.4f}, std: {normalized_layer_stm.std().item():.4f}")
-                else:
-                    self.debug_step += 1
-
-            if torch.isnan(encoded_layer_data).any():
-                print(f"!!!!!!!!!!!!!!!!!!!!!!         NaN detected in {i} layer encoded data input")
-
-            new_layer_stm = self.attention_layers[i](normalized_layer_stm, encoded_layer_data, encoded_layer_data, mask=attention_mask)
-            if torch.isnan(new_layer_stm).any():
-                print(f"!!!!!!!!!!!!!!!!!!!!!!         NaN detected in {i} layer memory attention output")
-
-            if self.use_gated_residual:
-                new_stm[i] = self._residual_gate(self.gate[i], layer_stm, new_layer_stm) # gated residual
-            else:
-                new_stm[i] = new_layer_stm + layer_stm # residual
-        self.stm.update_all(new_stm)
-        return self.stm.memory