rxnn 0.2.71__py3-none-any.whl → 0.2.73__py3-none-any.whl

rxnn/memory/attention.py

@@ -2,17 +2,15 @@ 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,
+            memory_input_norm_layers: nn.ModuleList,
+            residual_gate_layers: nn.ModuleList,
             debug_mode: bool = False,
             debug_interval: int = 10,
             *args,
@@ -22,16 +20,12 @@ class StmMemoryAttention(nn.Module):
         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.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.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
@@ -41,32 +35,27 @@ class StmMemoryAttention(nn.Module):
             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:
+        # 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)
-            # expand layer STM to batch size, if it's not in batch mode
+            # 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)
-            encoded_layer_data = x[i]
+
+            # 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)
-            if torch.isnan(normalized_layer_stm).any():
-                print(f"NaN detected in {i} layer memory norm output")
 
+            # 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
@@ -74,16 +63,88 @@ class StmMemoryAttention(nn.Module):
                 else:
                     self.debug_step += 1
 
-            if torch.isnan(encoded_layer_data).any():
-                print(f"NaN detected in {i} layer encoded data input")
-
+            # 9. Calculate memory attention
             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")
+            # 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)
 
-            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
+    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/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/rxt/models.py

@@ -9,11 +9,13 @@ 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
 
 
+
 class RxTAlphaComponentConfig(TypedDict):
     num_layers: int
     vocab_size: int
@@ -260,15 +262,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 +298,153 @@ 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):
+        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 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):
@@ -327,6 +470,7 @@ 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 RxTAlphaCriticEncoder(RxTAlphaComponentBase, pipeline_tag="text-classification", license="apache-2.0"):
     """RxT-Alpha (Reactive Transformer) encoder model"""
 

rxnn/training/mrl.py

@@ -592,7 +592,7 @@ class MRLTrainer:
 
         router_loss = actor.moe_router_loss()
         if torch.isnan(router_loss).any():
-            print("NaN detected in router loss")
+            print("!!!!!!!!!!!!!!!!!!!!!!         NaN detected in router loss")
         if router_loss is not None:
             return main_loss + self.moe_aux_loss_scale * router_loss
         else:
@@ -671,7 +671,7 @@ class MRLTrainer:
             # 4.4 Unscale and clip gradient norms
             self.scaler.unscale_(self.optimizer)
             torch.nn.utils.clip_grad_norm_(self.actor.unique_parameters(), max_norm=1.0,
-                                           error_if_nonfinite=self.debug_mode)
+                                           error_if_nonfinite=False)
             if self.debug_mode and self.epoch_step['train'] % self.debug_interval == 0:
                 self._log_gradients(logits)
             # 4.5 Run scaled optimization step
@@ -692,7 +692,7 @@ class MRLTrainer:
             policy_loss.backward(retain_graph=True)
             # 4.4 Clip gradient norms
             torch.nn.utils.clip_grad_norm_(self.actor.unique_parameters(), max_norm=1.0,
-                                           error_if_nonfinite=self.debug_mode)
+                                           error_if_nonfinite=False)
             if self.debug_mode and self.epoch_step['train'] % self.debug_interval == 0:
                 self._log_gradients(logits)
             # 4.5 Run scaled optimization step

rxnn/transformers/layers.py

@@ -1,5 +1,7 @@
 import torch
 import torch.nn as nn
+from poetry.console.commands import self
+
 from .attention import MultiHeadAttention
 from .ff import FeedForward, GatedFeedForward
 from .moe import MoeFeedForward, GatedMoeFeedForward
@@ -49,10 +51,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 +67,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 +108,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 +117,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 +132,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/transformers/models.py

@@ -94,7 +94,7 @@ class ReactiveTransformerDecoder(ReactiveTransformerBase):
     def forward(self, x: torch.Tensor, attention_mask: torch.Tensor = None) -> torch.Tensor:
         x = super().forward(x)  # apply embeddings
         if torch.isnan(x).any():
-            print("NaN detected in decoder embedding output")
+            print("!!!!!!!!!!!!!!!!!!!!!!         NaN detected in decoder embedding output")
         seq_len = x.size(1)
         if not self.use_flash_attention and self.use_relative_embedding:
             mask = create_causal_mask(seq_len, device=x.device)
@@ -112,7 +112,7 @@ class ReactiveTransformerDecoder(ReactiveTransformerBase):
         for i in range(self.num_own_layers):
             x = self._handle_layer(i, x, mask=mask)
             if torch.isnan(x).any():
-                print(f"NaN detected in {i}. decoder layer output")
+                print(f"!!!!!!!!!!!!!!!!!!!!!!         NaN detected in {i}. decoder layer output")
         return self.head(self.head_norm(x) if self.use_head_norm else x)
 
 
@@ -122,7 +122,7 @@ class ReactiveTransformerEncoder(ReactiveTransformerBase):
     def forward(self, x: torch.Tensor, attention_mask: torch.Tensor = None) -> tuple[torch.Tensor, torch.Tensor]:
         x = super().forward(x)  # apply embeddings
         if torch.isnan(x).any():
-            print("NaN detected in encoder embedding output")
+            print("!!!!!!!!!!!!!!!!!!!!!!         NaN detected in encoder embedding output")
         if attention_mask is not None:
             attention_mask = attention_mask.unsqueeze(1).unsqueeze(1).bool()
 
@@ -136,7 +136,7 @@ class ReactiveTransformerEncoder(ReactiveTransformerBase):
         for i in range(self.num_own_layers):
             x = self._handle_layer(i, x, mask=attention_mask)
             if torch.isnan(x).any():
-                print(f"NaN detected in {i}. encoder layer output")
+                print(f"!!!!!!!!!!!!!!!!!!!!!!         NaN detected in {i}. encoder layer output")
             hidden_states.append(x)
         return x, torch.stack(hidden_states)
 

{rxnn-0.2.71.dist-info → rxnn-0.2.73.dist-info}/METADATA

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

{rxnn-0.2.71.dist-info → rxnn-0.2.73.dist-info}/RECORD

@@ -5,11 +5,12 @@ rxnn/experimental/attention.py,sha256=jlNS82INjycNEfmk3HtkIacUvT_ELhaCO2g-kZTvhX
 rxnn/experimental/models.py,sha256=KheR1zSNJIaeVvpVAkEJwcuM5nOqQP0ZF08XhrtGJ8E,5387
 rxnn/experimental/moe.py,sha256=jHZ1QhpWiVQOswVpFmuH7b2IUOPf0Uuf-I2Ddwsd7Us,6140
 rxnn/memory/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-rxnn/memory/attention.py,sha256=O4ycW3KKP5hFYadgVh47LvGWJn9zNHz8vh9E9okC0h8,4223
+rxnn/memory/attention.py,sha256=CReYJZNA5JRED_QWqX-yKqEKZTRX6DNCAB8uFLZtKxI,7513
+rxnn/memory/gate.py,sha256=pR_H2y9C7S02QskoFAEC9Tmluut0k4GGlHgvZGiw6m4,2332
 rxnn/memory/norm.py,sha256=cVjjhCLqR5K6-321SP_ObG17y-ddlcTJeCTXvW4vpk0,6675
 rxnn/memory/stm.py,sha256=jv57gsH9XW19sLbxpRDqsp1yfsii_4Ef4Ncr_ztk-i4,3937
 rxnn/rxt/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-rxnn/rxt/models.py,sha256=M_0nEfSgr5Wyv-Ku4TCLpIs5VndUccjtIR0wU0DSVRo,15574
+rxnn/rxt/models.py,sha256=Pb48Frl6HV4Wb9CZgYtmzch3k_4Jess3rhs7dY1I96k,22209
 rxnn/training/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 rxnn/training/base.py,sha256=CqaArEZYOdH64nmKfx28U3GI46TzO4oNkjf_hrF23Cw,11835
 rxnn/training/bml.py,sha256=hw6gLpLkGvqLzxIvBg4MvCc5r8cHpEm2RDyh7nH6CtE,16914
@@ -17,7 +18,7 @@ rxnn/training/callbacks.py,sha256=rS8leuVFPVVfE5Zc8DMkUZhRIPN-vpPbUjowXE5TSBw,36
 rxnn/training/dataset.py,sha256=ruU6k33pQmpTqhxpjLFNdDJnCjcrBcGeFOzJqFahJDM,51880
 rxnn/training/ddp.py,sha256=VsNBjn3cY-uUj8hbsW7oKvb0_ZKnXnJ2KgObm-Mr9i4,836
 rxnn/training/models.py,sha256=ILkcqBV1MImnULnq-YDSSEf8cUdEbUgQaH0FRTsa4LA,9069
-rxnn/training/mrl.py,sha256=Ntkti6DDKipKa-AwTvo1WDOdIXOL3uXOhT-Xx29wR-w,67369
+rxnn/training/mrl.py,sha256=KUJAdUznquhf5UlcpV-QF5oKHDBEsDecMEVmMLQZw7w,67380
 rxnn/training/reward.py,sha256=uiSsBXmjMw2yv-1Bssy3RTlpU6zP8ape3490Sl-aT0M,16144
 rxnn/training/rl.py,sha256=hWtExxY-_pAmTOGYxyCNounUbaGWvLDVltC4sRC7MN4,7175
 rxnn/training/scheduler.py,sha256=LcjU35mEwz2U5x3U6tLfeeYlBqMxbFSxYzJYuXkWbSY,1408
@@ -26,14 +27,14 @@ rxnn/training/utils.py,sha256=ngDCm654NL3UsPy190Er4XPc9HI-OyEV6tDLMgEEvQc,6219
 rxnn/transformers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 rxnn/transformers/attention.py,sha256=KRnKT6XUqAXElxV9y72mSpdTeiMgCKCCLqqxCFNTHmA,16372
 rxnn/transformers/ff.py,sha256=WDjO-H9XWInoWnUnxiseIH6Kx5GlHP0zGJygwhcb1gc,2589
-rxnn/transformers/layers.py,sha256=bcDP8vZ5dpTWWqMCkzrPG8yQA0D0G5VjnV2Nq9IO8Dc,8816
+rxnn/transformers/layers.py,sha256=fxjlbQG6cwxq-b2ei4DnohSQGH5gwy4GkfP9duTUvjw,8492
 rxnn/transformers/mask.py,sha256=J0cfLVLt3SzS2ra3KcY4khrkhI975Dw4CjpUi3Sn25s,419
-rxnn/transformers/models.py,sha256=r4vNldYqCIpwMpXkFZvYbw0UBK3NE75qH7bc6OZ8YjE,11587
+rxnn/transformers/models.py,sha256=TP0H9do53Z0vd8kpHMISBzMpHE5X9QIHcy0B-iJHuNQ,11711
 rxnn/transformers/moe.py,sha256=j6jEx6Ip0zttlUZKKn82azxo95lkLZs-H2GLSMD88hY,5859
 rxnn/transformers/positional.py,sha256=1PjcJybUzeQlIKJI4tahAGZcYgCRCL0otxs7mpsNuzM,4410
 rxnn/transformers/sampler.py,sha256=t6iiQTdLQ0TakUWnnhKkb5DKF2F_9-thXHBydDF3fxg,17389
 rxnn/utils.py,sha256=ihb6OTyDtPiocB_lOvnq7eOkjjpCkgs8wxvXUBNQ7mM,996
-rxnn-0.2.71.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
-rxnn-0.2.71.dist-info/METADATA,sha256=7BHHcFtImjPB57X2eRLgO4IFOSBNb7GOR5ytMaCttkI,60420
-rxnn-0.2.71.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-rxnn-0.2.71.dist-info/RECORD,,
+rxnn-0.2.73.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
+rxnn-0.2.73.dist-info/METADATA,sha256=gtoRMeFgBuOZs4lRKl9JGUxZ2X4C9K78Ee-NHLMqW4E,60420
+rxnn-0.2.73.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+rxnn-0.2.73.dist-info/RECORD,,