rxnn 0.1.22__py3-none-any.whl → 0.1.24__py3-none-any.whl

rxnn/experimental/attention.py

@@ -61,6 +61,9 @@ class GroupedMoeAttention(GroupedQueryAttention):
             **kwargs,
         )
 
+    def router_loss(self):
+        return self.router.aux_loss
+
     def _init_kv(self, embed_dim: int):
         self.router = MoeRouter(embed_dim, self.num_experts, top_k=self.num_groups)
 
@@ -125,9 +128,6 @@ class GroupedMoeAttention(GroupedQueryAttention):
         k = self._process_grouped_experts(key, self.wk, self.bk, weights_k, indices_k)
         v = self._process_grouped_experts(value, self.wv, self.bv, weights_k, indices_k)
 
-        print('processed k', k.size())
-        print('processed v', v.size())
-
         # Expand to GQA format
         k = k.permute(0, 2, 1, 3).reshape(B, self.num_groups, S, -1)
         v = v.permute(0, 2, 1, 3).reshape(B, self.num_groups, S, -1)
@@ -141,10 +141,6 @@ class GroupedMoeAttention(GroupedQueryAttention):
             k = k.flatten(start_dim=1, end_dim=2)  # (B, H, S, head_dim)
             v = v.flatten(start_dim=1, end_dim=2)  # (B, H, S, head_dim)
 
-        print('q', q.size())
-        print('k', k.size())
-        print('v', v.size())
-
         return q, k, v
 
 
@@ -201,6 +197,9 @@ class DeepMoeAttention(GroupedMoeAttention):
             **kwargs,
         )
 
+    def router_loss(self):
+        return (self.router.aux_loss + self.query_router.aux_loss) / 2
+
     def _init_q(self, embed_dim: int):
         self.query_router = MoeRouter(embed_dim, self.num_query_experts, top_k=self.num_query_groups)
 
@@ -219,6 +218,11 @@ class DeepMoeAttention(GroupedMoeAttention):
         hidden_dim = embed_dim // (self.num_heads // self.num_query_groups)
         self.out_proj = nn.Linear(hidden_dim, embed_dim)
 
+    def _transpose_output(self, attn_output: torch.Tensor, b: int, t: int, d: int):
+        """Transpose attention output back to (B, T, D) shape"""
+        hidden_dim = d // self.num_heads * self.num_query_groups
+        return attn_output.transpose(1, 2).contiguous().view(b, t, hidden_dim)
+
     def _forward_qkv(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, b: int, t: int, d: int, skip_query_processing: bool = False):
         # Query processing
         B, T, D = query.shape
@@ -229,13 +233,8 @@ class DeepMoeAttention(GroupedMoeAttention):
         weights_q = weights_q_flat.view(B, T, -1)
         indices_q = indices_q_flat.view(B, T, -1)
         q = self._process_grouped_experts(query, self.wq, self.bq, weights_q, indices_q)
-        print('processed q', q.size())
-        q = q.permute(0, 2, 1, 3).reshape(B, self.num_query_groups, T, -1)
-
-        # Expand query groups to match head count
-        group_heads = self.num_heads // self.num_query_groups
-        q = q.unsqueeze(2).expand(-1, -1, group_heads, -1, -1).flatten(1, 2)
 
+        q = q.permute(0, 2, 1, 3).reshape(B, self.num_query_groups, T, -1)
         # Key/Value processing
         return super()._forward_qkv(q, key, value, b, t, d, skip_query_processing=True)
 

rxnn/transformers/attention.py

@@ -91,9 +91,13 @@ class MultiHeadAttention(nn.Module):
             attn_logits = attn_logits.masked_fill(mask == 0, float('-inf'))
         return F.softmax(attn_logits, dim=-1)
 
+    def _transpose_output(self, attn_output: torch.Tensor, b: int, t: int, d: int):
+        """Transpose attention output back to (B, T, D) shape"""
+        return attn_output.transpose(1, 2).contiguous().view(b, t, d)
+
     def _calculate_output(self, attn_weights: torch.Tensor, v: torch.Tensor, b: int, t: int, d: int):
         """Calculate the output by multiplying attention weights with values and concatenating heads"""
-        return torch.matmul(attn_weights, v).transpose(1, 2).contiguous().view(b, t, d)
+        return self._transpose_output(torch.matmul(attn_weights, v), b, t, d)
 
     def _flash_attention(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, b: int, t: int, d: int,
                          mask: torch.Tensor = None, enable_gqa: bool = False):
@@ -104,9 +108,7 @@ class MultiHeadAttention(nn.Module):
             is_causal=self.is_causal,
             enable_gqa=enable_gqa,
         )
-
-        # Reshape back to (B, T, D)
-        return attn_output.transpose(1, 2).contiguous().view(b, t, d)
+        return self._transpose_output(attn_output, b, t, d)
 
     def _calculate_flash_attention(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, b: int, t: int, d: int,
                                    mask: torch.Tensor = None):

rxnn/transformers/layers.py

@@ -60,7 +60,23 @@ class ReactiveTransformerLayer(nn.Module):
             param.requires_grad_(is_trainable)
 
     def moe_router_loss(self):
-        return self.ff.router_loss() if self.use_moe else None
+        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 ff_router_loss is not None and att_router_loss is not None:
+            return (ff_router_loss + att_router_loss) / 2
+        elif ff_router_loss is not None:
+            return ff_router_loss
+        elif att_router_loss is not None:
+            return att_router_loss
+        else:
+            return None
 
     def forward(self, x: torch.Tensor, stm: torch.Tensor, mask: torch.Tensor = None) -> torch.Tensor:
         # First step, self-attention
@@ -136,7 +152,17 @@ class ClassicTransformerLayer(nn.Module):
         self.use_moe = use_moe
 
     def moe_router_loss(self):
-        return self.ff.router_loss() if self.use_moe else torch.tensor(0.0)
+        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
+
+        if ff_router_loss is not None and att_router_loss is not None:
+            return (ff_router_loss + att_router_loss) / 2
+        elif ff_router_loss is not None:
+            return ff_router_loss
+        elif att_router_loss is not None:
+            return att_router_loss
+        else:
+            return None
 
     def forward(self, x: torch.Tensor, mask: torch.Tensor = None) -> torch.Tensor:
         # First step, self-attention

rxnn/transformers/moe.py

@@ -20,7 +20,6 @@ class MoeRouter(nn.Module):
         mean_probs = probs.mean(dim=0)
         return (expert_usage * mean_probs).sum() * self.num_experts
 
-
     def forward(self, x: torch.Tensor):
         # Input shape: [batch*seq_len, embed_dim]
         logits = self.gate(x)

{rxnn-0.1.22.dist-info → rxnn-0.1.24.dist-info}/METADATA

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

{rxnn-0.1.22.dist-info → rxnn-0.1.24.dist-info}/RECORD

@@ -1,6 +1,6 @@
 rxnn/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 rxnn/experimental/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-rxnn/experimental/attention.py,sha256=mvTrDiXRTeJsubWYRAJQjmw9nDhfMkQkQoXaAjeUzDo,32363
+rxnn/experimental/attention.py,sha256=r3y0_BaweONrkm6Z-9zn56U9jlBMfYBs8NlWNm7rR90,32424
 rxnn/experimental/models.py,sha256=-XkEHsyT8iNAjhZbgC7N_5nzP4ENVJLwxSoLHgMfA0I,4668
 rxnn/experimental/moe.py,sha256=PhiaNr3FwR2Zv2a0tfj6sfZ4iyhLo3Jyp2DwXq19qZQ,7935
 rxnn/memory/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -16,16 +16,16 @@ rxnn/training/dataset.py,sha256=vQ5mDF3bA0HXya474n4D4iL8Mn3AEpJukgzFNVkxjGU,5106
 rxnn/training/scheduler.py,sha256=ow6oALzWjWQmHSpcJEjv6tg4g4CDMvr73TypxfcefMc,712
 rxnn/training/tokenizer.py,sha256=4Y41f07uo2KPA_7bp3FCcwGKbXoS2hsckOoXUsXfQxY,8052
 rxnn/transformers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-rxnn/transformers/attention.py,sha256=FfEYE0THO73p_1eRupr2mcwfW4UbI_riIxkHfr8X_1c,14022
+rxnn/transformers/attention.py,sha256=FHATZVf_kt3OHnG02zEeG9QdUXLncKDjrhyT28Pk0E4,14185
 rxnn/transformers/ff.py,sha256=jJnuBDsnnX5uYC_WZH8cXAYrMnz0P-iX7MwcPivjRtI,2533
-rxnn/transformers/layers.py,sha256=HhIiykmrBgdsV4AbMQXr9t0cSo4gSIeN0dPtc8mDyOo,5629
+rxnn/transformers/layers.py,sha256=ZJfNdgCv9dzrKqsWIMf99Ryzgs494ZhkwK4zSBYLvQ4,6880
 rxnn/transformers/mask.py,sha256=J0cfLVLt3SzS2ra3KcY4khrkhI975Dw4CjpUi3Sn25s,419
 rxnn/transformers/models.py,sha256=w-zB_8QB9-Fae-GkGgmVDNY-Ts_0gBeWcevpl9qzZVM,7169
-rxnn/transformers/moe.py,sha256=msspVdefdt2ekIN8aT-V8DolK4taESQL_NVsSGOepIs,4739
+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.22.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
-rxnn-0.1.22.dist-info/METADATA,sha256=7yZS7KjHOdh6FNOloKZIDPDWGYAassqQwZtMi1iL5mI,16627
-rxnn-0.1.22.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
-rxnn-0.1.22.dist-info/RECORD,,
+rxnn-0.1.24.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
+rxnn-0.1.24.dist-info/METADATA,sha256=_h4mqmSKPEr0mxc2CaMn-yzvmZ5Lqlk_H4parGt-eHk,16627
+rxnn-0.1.24.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
+rxnn-0.1.24.dist-info/RECORD,,