rxnn 0.2.63__tar.gz → 0.2.64__tar.gz

{rxnn-0.2.63 → rxnn-0.2.64}/PKG-INFO

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

{rxnn-0.2.63 → rxnn-0.2.64}/pyproject.toml

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

{rxnn-0.2.63 → rxnn-0.2.64}/src/rxnn/experimental/models.py

@@ -34,6 +34,8 @@ class ExperimentalAttentionTransformerConfig(TypedDict):
     att_num_query_groups: int
     att_num_global_tokens: int
     att_window_size: int
+    use_head_norm: bool
+    init_identity_norm: bool
 
 
 class ExperimentalAttentionTransformer(nn.Module, PyTorchModelHubMixin, pipeline_tag="text-generation", license="apache-2.0"):
@@ -67,6 +69,8 @@ class ExperimentalAttentionTransformer(nn.Module, PyTorchModelHubMixin, pipeline
             att_num_query_groups: int = None,
             att_num_global_tokens: int = 16,
             att_window_size: int = 128,
+            use_head_norm: bool = False,
+            init_identity_norm: bool = False,
             **kwargs
     ):
         super(ExperimentalAttentionTransformer, self).__init__(**kwargs)
@@ -110,6 +114,8 @@ class ExperimentalAttentionTransformer(nn.Module, PyTorchModelHubMixin, pipeline
                     use_rms_norm=use_rms_norm,
                     self_attention=att_init(),
                     use_moe_att=use_moe_att,
+                    use_head_norm=use_head_norm,
+                    init_identity_norm=init_identity_norm,
                 ) for _ in range(num_layers)
             ]),
             use_flash_attention=use_flash_attention,

{rxnn-0.2.63 → rxnn-0.2.64}/src/rxnn/rxt/models.py

@@ -39,6 +39,8 @@ class RxTAlphaComponentConfig(TypedDict):
     att_query_groups: int
     cross_att_groups: int
     cross_att_query_groups: int
+    use_head_norm: bool
+    init_identity_norm: bool
 
 
 class RxTAlphaComponentBase(nn.Module, PyTorchModelHubMixin):
@@ -71,6 +73,8 @@ class RxTAlphaComponentBase(nn.Module, PyTorchModelHubMixin):
             att_query_groups: int = None,
             cross_att_groups: int = None,
             cross_att_query_groups: int = None,
+            use_head_norm: bool = False,
+            init_identity_norm: bool = False,
             **kwargs
     ):
         super(RxTAlphaComponentBase, self).__init__(**kwargs)
@@ -130,10 +134,14 @@ class RxTAlphaComponentBase(nn.Module, PyTorchModelHubMixin):
                 memory_cross_attention=cross_att_init(),
             ) for _ in range(num_layers)
         ])
-        self.model = self._init_model(stm, layers, embedding, use_flash_attention, embed_dim, vocab_size, use_moe)
+        self.model = self._init_model(
+            stm, layers, embedding, use_flash_attention, embed_dim, vocab_size, use_moe,
+            use_head_norm=use_head_norm, init_identity_norm=init_identity_norm,
+        )
 
     def _init_model(self, stm: ShortTermMemory, layers: nn.ModuleList, embedding: nn.Embedding,
-                    use_flash_attention: bool, embed_dim: int, vocab_size: int, use_moe: bool) -> ReactiveTransformerBase:
+                    use_flash_attention: bool, embed_dim: int, vocab_size: int, use_moe: bool,
+                    use_head_norm: bool = False, init_identity_norm: bool = False) -> ReactiveTransformerBase:
         pass
 
     def params_count(self):
@@ -187,6 +195,8 @@ class RxTAlphaEncoder(RxTAlphaComponentBase, pipeline_tag="fill-mask", license="
             embed_dim: int,
             vocab_size: int,
             use_moe: bool,
+            use_head_norm: bool = False,
+            init_identity_norm: bool = False,
     ) -> ReactiveTransformerEncoder:
         return ReactiveTransformerEncoder(
             stm=stm,
@@ -214,6 +224,8 @@ class RxTAlphaDecoder(RxTAlphaComponentBase, pipeline_tag="text-generation", lic
             embed_dim: int,
             vocab_size: int,
             use_moe: bool,
+            use_head_norm: bool = False,
+            init_identity_norm: bool = False,
     ) -> ReactiveTransformerDecoder:
         return ReactiveTransformerDecoder(
             embed_dim,
@@ -223,6 +235,8 @@ class RxTAlphaDecoder(RxTAlphaComponentBase, pipeline_tag="text-generation", lic
             own_layers=layers,
             use_flash_attention=use_flash_attention,
             use_moe=use_moe,
+            use_head_norm=use_head_norm,
+            init_identity_norm=init_identity_norm,
         )
 
     def forward(self, x: torch.Tensor, attention_mask: torch.Tensor = None) -> tuple[torch.Tensor, torch.Tensor]:
@@ -327,6 +341,8 @@ class RxTAlphaCriticEncoder(RxTAlphaComponentBase, pipeline_tag="text-classifica
             embed_dim: int,
             vocab_size: int,
             use_moe: bool = False,
+            use_head_norm: bool = False,
+            init_identity_norm: bool = False,
     ) -> ReactiveTransformerEncoderDetachStm:
         return ReactiveTransformerEncoderDetachStm(
             stm=stm,

{rxnn-0.2.63 → rxnn-0.2.64}/src/rxnn/training/mrl.py

@@ -602,14 +602,17 @@ class MRLTrainer:
         print(f"Decoder grad norm - total: {decoder_total:.6f}, mean: {decoder_mean:.6f}")
         print(f"Memory attention grad norm - total: {mem_att_total:.6f}, mean: {mem_att_mean:.6f}")
         # decoder's cross att
-        dec_x_att_norms = [get_gradient_norms(layer.memory_cross_attention)[0] for layer in self.actor.decoder.model.layers]
-        print(f"Decoder cross-att mean total norm: {(sum(dec_x_att_norms) / len(dec_x_att_norms)):.6f}, all: {dec_x_att_norms}")
+        dec_x_att_norms = [get_gradient_norms(layer.memory_cross_attention)[1] for layer in self.actor.decoder.model.layers]
+        print(f"Decoder cross-att mean norm: {(sum(dec_x_att_norms) / len(dec_x_att_norms)):.6f}, all: {dec_x_att_norms}")
 
-        mem_att_norms = [get_gradient_norms(layer)[0] for layer in self.actor.memory_attention.model.attention_layers]
-        print(f"Memory attention layers mean total norm: {(sum(mem_att_norms) / len(mem_att_norms)):.6f}, all: {mem_att_norms}")
+        mem_att_norms = [get_gradient_norms(layer)[1] for layer in self.actor.memory_attention.model.attention_layers]
+        print(f"Memory attention layers mean norm: {(sum(mem_att_norms) / len(mem_att_norms)):.6f}, all: {mem_att_norms}")
 
-        enc_ff_norms = [get_gradient_norms(layer.ff)[0] for layer in self.actor.encoder.model.layers]
-        print(f"Encoder ff mean total norm: {(sum(enc_ff_norms) / len(enc_ff_norms)):.6f}, all: {enc_ff_norms}")
+        enc_ff_norms = [get_gradient_norms(layer.ff)[1] for layer in self.actor.encoder.model.layers]
+        print(f"Encoder ff mean norm: {(sum(enc_ff_norms) / len(enc_ff_norms)):.6f}, all: {enc_ff_norms}")
+
+        enc_ff_norms = [get_gradient_norms(layer.memory_cross_attention)[1] for layer in self.actor.encoder.model.layers]
+        print(f"Encoder cross-att mean norm: {(sum(enc_ff_norms) / len(enc_ff_norms)):.6f}, all: {enc_ff_norms}")
 
     def update_actor(self, state: tuple[TokenizedDict, TokenizedDict, TokenizedDict], action: TokenizedDict,
                      advantages: torch.Tensor, old_log_probs: torch.Tensor, epoch: int) -> float:

{rxnn-0.2.63 → rxnn-0.2.64}/src/rxnn/training/utils.py

@@ -146,10 +146,10 @@ def smart_concat(query: TokenizedDict, answer: TokenizedDict, max_length: int, p
 
 def get_gradient_norms(model: nn.Module):
     total_norm = 0
-    for p in model.parameters():
-        if p.grad is not None:
-            param_norm = p.grad.data.norm(2)
-            total_norm += param_norm.item() ** 2
+    grad_params = list(filter(lambda p: p.requires_grad and p.grad is not None, model.parameters()))
+    for p in grad_params:
+        param_norm = p.grad.data.norm(2)
+        total_norm += param_norm.item() ** 2
     total_norm = total_norm ** 0.5
-    mean_norm = total_norm / len(list(model.parameters()))
+    mean_norm = total_norm / len(grad_params)
     return total_norm, mean_norm

{rxnn-0.2.63 → rxnn-0.2.64}/src/rxnn/transformers/models.py

@@ -69,9 +69,18 @@ class ReactiveTransformerBase(nn.Module):
 class ReactiveTransformerDecoder(ReactiveTransformerBase):
     """Reactive Transformer decoder - extending the classic Transformer decoder with Memory Cross-Attention"""
 
-    def __init__(self, embed_dim: int, vocab_size: int, *args, **kwargs):
+    def __init__(self, embed_dim: int, vocab_size: int, use_head_norm: bool = False, init_identity_norm: bool = False, *args, **kwargs):
         super(ReactiveTransformerDecoder, self).__init__(*args, **kwargs)
+
         self.head = nn.Linear(embed_dim, vocab_size)
+        self.use_head_norm = use_head_norm
+        if use_head_norm:
+            self.head_norm = nn.LayerNorm(embed_dim)
+            if init_identity_norm:
+                self.head_norm.weight.data.fill_(1.0)
+                self.head_norm.bias.data.fill_(0.0)
+        else:
+            self.head_norm = None
 
     def forward(self, x: torch.Tensor, attention_mask: torch.Tensor = None) -> torch.Tensor:
         x = super().forward(x)  # apply embeddings
@@ -99,7 +108,7 @@ class ReactiveTransformerDecoder(ReactiveTransformerBase):
             if layer_stm.size(0) == 1:
                 layer_stm = layer_stm.expand(x.size(0), -1, -1)
             x = self.layers[i](x, layer_stm, mask=mask)
-        return self.head(x)
+        return self.head(self.head_norm(x) if self.use_head_norm else x)
 
 
 class ReactiveTransformerEncoder(ReactiveTransformerBase):
@@ -201,9 +210,17 @@ class ClassicTransformerBase(nn.Module):
 class ClassicTransformerDecoder(ClassicTransformerBase):
     """Classic Transformer decoder - for decoder-only Transformer models"""
 
-    def __init__(self, embed_dim: int, vocab_size: int, *args, **kwargs):
+    def __init__(self, embed_dim: int, vocab_size: int, use_head_norm: bool = False, init_identity_norm: bool = False, *args, **kwargs):
         super(ClassicTransformerDecoder, self).__init__(*args, **kwargs)
         self.head = nn.Linear(embed_dim, vocab_size)
+        self.use_head_norm = use_head_norm
+        if use_head_norm:
+            self.head_norm = nn.LayerNorm(embed_dim)
+            if init_identity_norm:
+                self.head_norm.weight.data.fill_(1.0)
+                self.head_norm.bias.data.fill_(0.0)
+        else:
+            self.head_norm = None
 
     def forward(self, x: torch.Tensor, attention_mask: torch.Tensor = None) -> torch.Tensor:
         x = super().forward(x)  # apply embeddings
@@ -213,7 +230,6 @@ class ClassicTransformerDecoder(ClassicTransformerBase):
             if attention_mask is not None:
                 mask &= attention_mask.unsqueeze(1).unsqueeze(1).bool()
         elif attention_mask is not None:
-            print(attention_mask.size())
             mask = attention_mask.unsqueeze(1).unsqueeze(1).bool()
         else:
             mask = None
@@ -221,7 +237,7 @@ class ClassicTransformerDecoder(ClassicTransformerBase):
         # Process layers
         for i in range(self.num_layers):
             x = self.layers[i](x, mask=mask)
-        return self.head(x)
+        return self.head(self.head_norm(x) if self.use_head_norm else x)
 
 
 class ClassicTransformerEncoder(ClassicTransformerBase):