rxnn 0.1.49__py3-none-any.whl → 0.1.50__py3-none-any.whl

rxnn/experimental/models.py

@@ -83,6 +83,8 @@ class MoeAttentionTransformer(nn.Module, PyTorchModelHubMixin, pipeline_tag="tex
                                                   num_query_experts=att_num_query_experts,
                                                   num_query_groups=att_num_query_groups)
 
+        use_moe_att = att_type in ['gma', 'dma', 'gma_s', 'dma_s']
+
         self.model = ClassicTransformerDecoder(
             embed_dim,
             vocab_size,
@@ -99,6 +101,7 @@ class MoeAttentionTransformer(nn.Module, PyTorchModelHubMixin, pipeline_tag="tex
                     ff_dropout=ff_dropout,
                     use_rms_norm=use_rms_norm,
                     self_attention=att_init(),
+                    use_moe_att=use_moe_att,
                 ) for _ in range(num_layers)
             ]),
             use_flash_attention=use_flash_attention,

rxnn/transformers/layers.py

@@ -22,6 +22,7 @@ class ReactiveTransformerLayer(nn.Module):
             use_moe: bool = False,
             num_experts: int = 1,
             moe_top_k: int = 1,
+            use_moe_att: bool = False,
             *args,
             **kwargs,
     ):
@@ -54,6 +55,7 @@ class ReactiveTransformerLayer(nn.Module):
             self.norm3 = nn.LayerNorm(embed_dim)
         self.use_post_norm = use_post_norm
         self.use_moe = use_moe
+        self.use_moe_att = use_moe_att
 
     def trainable_cross_attention_(self, is_trainable: bool):
         for param in self.memory_cross_attention.parameters():
@@ -62,12 +64,13 @@ class ReactiveTransformerLayer(nn.Module):
     def moe_router_loss(self):
         ff_router_loss = self.ff.router_loss() if self.use_moe else None
         att_router_loss = None
-        if self.attention.router_loss is not None and self.memory_cross_attention.router_loss is not None:
-            att_router_loss = (self.attention.router_loss() + self.memory_cross_attention.router_loss()) / 2
-        elif self.attention.router_loss is not None:
-            att_router_loss = self.attention.router_loss()
-        elif self.memory_cross_attention.router_loss is not None:
-            att_router_loss = self.memory_cross_attention.router_loss()
+        if self.use_moe_att:
+            if self.attention.router_loss is not None and self.memory_cross_attention.router_loss is not None:
+                att_router_loss = (self.attention.router_loss() + self.memory_cross_attention.router_loss()) / 2
+            elif self.attention.router_loss is not None:
+                att_router_loss = self.attention.router_loss()
+            elif self.memory_cross_attention.router_loss is not None:
+                att_router_loss = self.memory_cross_attention.router_loss()
 
         if ff_router_loss is not None and att_router_loss is not None:
             return (ff_router_loss + att_router_loss) / 2
@@ -123,6 +126,7 @@ class ClassicTransformerLayer(nn.Module):
             use_moe: bool = False,
             num_experts: int = 1,
             moe_top_k: int = 1,
+            use_moe_att: bool = False,
             *args,
             **kwargs,
     ):
@@ -151,10 +155,11 @@ class ClassicTransformerLayer(nn.Module):
             self.norm2 = nn.LayerNorm(embed_dim)
         self.use_post_norm = use_post_norm
         self.use_moe = use_moe
+        self.use_moe_att = use_moe_att
 
     def moe_router_loss(self):
         ff_router_loss = self.ff.router_loss() if self.use_moe else None
-        att_router_loss = self.attention.router_loss() if self.attention.router_loss is not None else None
+        att_router_loss = self.attention.router_loss() if self.use_moe_att and self.attention.router_loss is not None else None
 
         if ff_router_loss is not None and att_router_loss is not None:
             return (ff_router_loss + att_router_loss) / 2

rxnn/transformers/models.py

@@ -38,8 +38,8 @@ class ReactiveTransformerBase(nn.Module):
             self.layers[i].trainable_cross_attention_(is_trainable)
 
     def moe_router_loss(self):
-        return torch.stack([self.layers[i].moe_router_loss() for i in range(self.num_own_layers) if self.layers[i].use_moe] + [
-            self.shared_layers[i].moe_router_loss() for i in range(self.num_shared_layers) if self.shared_layers[i].use_moe]).mean()
+        return torch.stack([self.layers[i].moe_router_loss() for i in range(self.num_own_layers) if self.layers[i].use_moe or self.layers[i].use_moe_att] + [
+            self.shared_layers[i].moe_router_loss() for i in range(self.num_shared_layers) if self.shared_layers[i].use_moe or self.shared_layers[i].use_moe_att]).mean()
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         # Shared logic for encoders and decoders - apply embeddings and positional encoding
@@ -124,7 +124,7 @@ class ClassicTransformerBase(nn.Module):
         self.num_layers = len(layers) if layers else 0
 
     def moe_router_loss(self):
-        return torch.stack([self.layers[i].moe_router_loss() for i in range(self.num_layers) if self.layers[i].use_moe]).mean()
+        return torch.stack([self.layers[i].moe_router_loss() for i in range(self.num_layers) if self.layers[i].use_moe or self.layers[i].use_moe_att]).mean()
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         # Shared logic for encoders and decoders - apply embeddings and positional encoding

{rxnn-0.1.49.dist-info → rxnn-0.1.50.dist-info}/METADATA

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

{rxnn-0.1.49.dist-info → rxnn-0.1.50.dist-info}/RECORD

@@ -1,7 +1,7 @@
 rxnn/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 rxnn/experimental/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 rxnn/experimental/attention.py,sha256=22Qb4jYN6QaqibTU8bwD8x2FaOKCxvWglM2eK9EuOlo,29468
-rxnn/experimental/models.py,sha256=IzUVc5s-cA__8jsG2mVvzUDmzPRcfBcI5btaOjnPYhA,4598
+rxnn/experimental/models.py,sha256=-BQn7gWlSHLpkAQdthPW5L9ZNzIBqSJS9tkm2N88jgw,4711
 rxnn/experimental/moe.py,sha256=jHZ1QhpWiVQOswVpFmuH7b2IUOPf0Uuf-I2Ddwsd7Us,6140
 rxnn/memory/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 rxnn/memory/norm.py,sha256=Ofl8Q5NYEF9GQeO0bhM43tkTW91J0y6TSvTAOYMgloM,6278
@@ -18,14 +18,14 @@ rxnn/training/tokenizer.py,sha256=4Y41f07uo2KPA_7bp3FCcwGKbXoS2hsckOoXUsXfQxY,80
 rxnn/transformers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 rxnn/transformers/attention.py,sha256=dC0UmC-_kjX8US6Sf0Fi5zw5kJ-P6orH3JDHeBB5gI8,15695
 rxnn/transformers/ff.py,sha256=jJnuBDsnnX5uYC_WZH8cXAYrMnz0P-iX7MwcPivjRtI,2533
-rxnn/transformers/layers.py,sha256=n_jZTqEF_vLkF31AkB5XGErfm2sQFd9CRqJUHKRFkKI,6956
+rxnn/transformers/layers.py,sha256=OX8CsFY9A7uqH1SLwyexR_5BNlwheYrJHCGXjF8Q7HU,7186
 rxnn/transformers/mask.py,sha256=J0cfLVLt3SzS2ra3KcY4khrkhI975Dw4CjpUi3Sn25s,419
-rxnn/transformers/models.py,sha256=w-zB_8QB9-Fae-GkGgmVDNY-Ts_0gBeWcevpl9qzZVM,7169
+rxnn/transformers/models.py,sha256=_w5C7xvjT4-BFeMfzi57BQ51_fgaYZ4UK0SqUDE5Ooo,7266
 rxnn/transformers/moe.py,sha256=6Cffyo0QjmEWc4rK1ncOmLRCQbY0OpQJ4D7xH_4nTN4,4738
 rxnn/transformers/positional.py,sha256=2l38RS0Dini3f6Z3LUHr3XwWzg1UK7fO2C6wazWDAYU,4292
 rxnn/transformers/sampler.py,sha256=poWBpxg1iuK5gEJtxHkk5VVfS9V48hs2Olqdhy_Gw8c,6548
 rxnn/utils.py,sha256=d5U8i5ukovgDyqiycc2AoxObTz_eF_bgo2MKvdtJ98s,467
-rxnn-0.1.49.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
-rxnn-0.1.49.dist-info/METADATA,sha256=PijR2z5P5nuTlOaWn-ylU_Loluy-e2HRgpMEc4TCohk,16627
-rxnn-0.1.49.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
-rxnn-0.1.49.dist-info/RECORD,,
+rxnn-0.1.50.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
+rxnn-0.1.50.dist-info/METADATA,sha256=bIeDbrlcclSfD9oHf26i_sYepOTvTkpcwQMWpOm2jWc,16627
+rxnn-0.1.50.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
+rxnn-0.1.50.dist-info/RECORD,,