rxnn 0.2.62__py3-none-any.whl → 0.2.64__py3-none-any.whl

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/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/training/mrl.py

@@ -243,20 +243,31 @@ class MRLTrainer:
             critic_weight_decay: float,
             critic_encoder_lr: float,
             embedding_lr: float,
+            encoder_lr: float,
             memory_lr: Optional[float] = None,
+            encoder_memory_lr: Optional[float] = None,
+            memory_attn_lr: Optional[float] = None,
     ) -> tuple[torch.optim.Optimizer, torch.optim.Optimizer]:
         if memory_lr is not None:
             optimizer = torch.optim.AdamW([
                 {'params': self.actor.embedding_parameters(), 'lr': embedding_lr},
-                {'params': self.actor.not_memory_parameters(), 'lr': lr},
-                {'params': self.actor.memory_parameters(), 'lr': memory_lr},
+                {'params': self.actor.encoder.not_memory_parameters(), 'lr': encoder_lr},
+                {'params': self.actor.encoder.memory_parameters(), 'lr': encoder_memory_lr},
+                {'params': self.actor.memory_attention_parameters(), 'lr': memory_attn_lr},
+                {'params': self.actor.decoder.memory_parameters(), 'lr': memory_lr},
+                {'params': self.actor.decoder.not_memory_parameters(), 'lr': lr},
             ],
                 weight_decay=weight_decay,
             )
         else:
             optimizer = torch.optim.AdamW([
                 {'params': self.actor.embedding_parameters(), 'lr': embedding_lr},
-                {'params': self.actor.unique_parameters(with_embedding=False), 'lr': lr},
+                {'params': self.actor.embedding_parameters(), 'lr': embedding_lr},
+                {'params': self.actor.encoder.not_memory_parameters(), 'lr': encoder_lr},
+                {'params': self.actor.encoder.memory_parameters(), 'lr': encoder_lr},
+                {'params': self.actor.memory_attention_parameters(), 'lr': lr},
+                {'params': self.actor.decoder.memory_parameters(), 'lr': lr},
+                {'params': self.actor.decoder.not_memory_parameters(), 'lr': lr},
             ],
                 weight_decay=weight_decay,
             )
@@ -591,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)[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}")
 
-        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}")
+        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.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.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/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/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):

{rxnn-0.2.62.dist-info → rxnn-0.2.64.dist-info}/METADATA

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

{rxnn-0.2.62.dist-info → rxnn-0.2.64.dist-info}/RECORD

@@ -2,14 +2,14 @@ rxnn/.DS_Store,sha256=BxZLo9tFs48JMq6jhumiCnCPLTeCwl619CFSg4ClRAY,6148
 rxnn/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 rxnn/experimental/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 rxnn/experimental/attention.py,sha256=jlNS82INjycNEfmk3HtkIacUvT_ELhaCO2g-kZTvhXg,28455
-rxnn/experimental/models.py,sha256=oJWd56LUsLc9S8eCZw-ShvuWjoQxj4C9GitbohlQ0ok,5139
+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=lSniKrf_skiM1V1zbfmV84PbKoQ-t_fVcKfwNKW3_OY,3844
 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=mLHvb3ablQK9UtupuOHmLlG440Q_NW-OuLWcxGMfGuY,14807
+rxnn/rxt/models.py,sha256=JrZQ78F4HGGklAy6mML4fbqdsMOcGSDRZpjhX55VXb8,15486
 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,23 +17,23 @@ rxnn/training/callbacks.py,sha256=rS8leuVFPVVfE5Zc8DMkUZhRIPN-vpPbUjowXE5TSBw,36
 rxnn/training/dataset.py,sha256=tbtOSYldHnQB6SWgee_yUj9zTbgoEoLFNa6wvUS6Apg,51292
 rxnn/training/ddp.py,sha256=VsNBjn3cY-uUj8hbsW7oKvb0_ZKnXnJ2KgObm-Mr9i4,836
 rxnn/training/models.py,sha256=KIiOCW0VgKtMA4EMQ---xsVExdI1mBsgWjtRSmJpecA,9033
-rxnn/training/mrl.py,sha256=BWp87Lj4epjTlROmrQK8RnS_83IucqS7XWI1cBae7BM,64424
+rxnn/training/mrl.py,sha256=2J6Wh4xtsVoE6duEevmovDpmSsMkEoH39Ru0bE8lhFo,65481
 rxnn/training/reward.py,sha256=uiSsBXmjMw2yv-1Bssy3RTlpU6zP8ape3490Sl-aT0M,16144
 rxnn/training/rl.py,sha256=hWtExxY-_pAmTOGYxyCNounUbaGWvLDVltC4sRC7MN4,7175
 rxnn/training/scheduler.py,sha256=LcjU35mEwz2U5x3U6tLfeeYlBqMxbFSxYzJYuXkWbSY,1408
 rxnn/training/tokenizer.py,sha256=umaLByMBx_NMrQElA45HLm9gkuzyKWDTFaKVd-CjXl0,8344
-rxnn/training/utils.py,sha256=C0OS2RAGQ3L7D_G3CWupu_BpAFhkovMByBKm355Ibfc,6087
+rxnn/training/utils.py,sha256=QMNkJPQBY04DX9WN7GHnI2EZTBbAzWkjt2W-798oUII,6129
 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=OlbqD5kKygn5WZziLbU3jZjhr8hBrxLpqlCjJ_BNCW0,8119
 rxnn/transformers/mask.py,sha256=J0cfLVLt3SzS2ra3KcY4khrkhI975Dw4CjpUi3Sn25s,419
-rxnn/transformers/models.py,sha256=tT0W5inG4EtjEHNutG77Wcws2fJzLJs-iFDP3hX3D2Q,10761
+rxnn/transformers/models.py,sha256=7xKixlBN5uZckcOXfukZeA4f7R34A35Gk98lNzew37o,11559
 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.62.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
-rxnn-0.2.62.dist-info/METADATA,sha256=A40GBcyyy0ZxkHxFDXQVc7Ghrz9pvlYFDwPHEAbLuFI,25997
-rxnn-0.2.62.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-rxnn-0.2.62.dist-info/RECORD,,
+rxnn-0.2.64.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
+rxnn-0.2.64.dist-info/METADATA,sha256=QaDWd-8W0vs3povCgRAUXXSmTNN8gkEJ1dY6mA7n9kQ,25997
+rxnn-0.2.64.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+rxnn-0.2.64.dist-info/RECORD,,