rxnn 0.1.55__py3-none-any.whl → 0.1.57__py3-none-any.whl

rxnn/experimental/attention.py

@@ -283,12 +283,12 @@ class SparseQueryAttention(MultiHeadAttention):
         """Override query, key, and value projections for GQA case - split data into heads and groups"""
         head_dim = d // self.num_heads
         if not self.rel_embed:
-            q = self.q_proj(query).view(b, t, self.num_query_heads, head_dim).transpose(1, 2)
+            q = self.q_proj(query).view(b, t, self.num_query_groups, head_dim).transpose(1, 2)
             k = self.k_proj(key).view(b, -1, self.num_groups, head_dim).transpose(1, 2)
             v = self.v_proj(value).view(b, -1, self.num_groups, head_dim).transpose(1, 2)
         else:
             group_heads = self.num_heads // self.num_groups
-            query_heads = self.num_query_heads // self.num_query_groups
+            query_heads = self.num_heads // self.num_query_groups
             # Process Q
             q = self.q_proj(query).view(b, -1, self.num_query_groups, head_dim).transpose(1, 2)  # (B, Q_G, T, head_dim)
 
@@ -683,7 +683,7 @@ class InfiniteAttention(MultiHeadAttention):
         attn = torch.softmax(attn, dim=-1)
         return torch.einsum('b h i j, b h j d -> b h i d', attn, v)
 
-def init_moe_attention(
+def init_experimental_attention(
         embed_dim: int,
         num_heads: int,
         attention_type: str,

rxnn/experimental/models.py

@@ -8,7 +8,7 @@ from ..transformers.layers import ClassicTransformerLayer
 from ..transformers.models import ClassicTransformerDecoder
 from ..transformers.ff import get_activation_layer
 from ..utils import get_model_size
-from .attention import init_moe_attention
+from .attention import init_experimental_attention
 
 
 class MoeAttentionTransformerConfig(TypedDict):
@@ -77,11 +77,11 @@ class MoeAttentionTransformer(nn.Module, PyTorchModelHubMixin, pipeline_tag="tex
                                               use_flash_attention=use_flash_attention, dropout=att_dropout,
                                               max_seq_len=seq_len, is_causal=True)
         else:
-            att_init = lambda: init_moe_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=True, num_experts=att_num_experts,
-                                                  num_query_experts=att_num_query_experts,
-                                                  num_query_groups=att_num_query_groups)
+            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=True, num_experts=att_num_experts,
+                                                           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']
 

rxnn/rxt/models.py

@@ -9,7 +9,7 @@ from ..transformers.models import ReactiveTransformerBase, ReactiveTransformerEn
 from ..transformers.ff import get_activation_layer
 from ..memory.stm import ShortTermMemory
 from ..utils import get_model_size
-
+from ..experimental.attention import init_experimental_attention
 
 class RxTAlphaComponentConfig(TypedDict):
     num_layers: int
@@ -31,6 +31,9 @@ class RxTAlphaComponentConfig(TypedDict):
     moe_top_k: int
     self_att_type: str
     cross_att_type: str
+    att_num_experts: int
+    att_num_query_experts: int
+    att_num_query_groups: int
 
 
 class RxTAlphaComponentBase(nn.Module, PyTorchModelHubMixin):
@@ -58,14 +61,17 @@ class RxTAlphaComponentBase(nn.Module, PyTorchModelHubMixin):
             moe_top_k: int = 1,
             self_att_type: str = 'gqa',
             cross_att_type: str = 'mqa',
+            att_num_experts: int = None,
+            att_num_query_experts: int = None,
+            att_num_query_groups: int = None,
             **kwargs
     ):
         super(RxTAlphaComponentBase, self).__init__(**kwargs)
         assert ff_activation in ['relu', 'gelu',
                                  'swish', 'silu', 'linear',
                                  'sigmoid'], 'Feed-forward activation could be "relu", "gelu", "swish", "silu", "linear", "sigmoid".'
-        assert self_att_type in ['mha', 'gqa', 'mqa'], 'Self-attention type could be "mha", "gqa", "mqa"'
-        assert cross_att_type in ['mha', 'gqa', 'mqa'], 'Memory cross-attention type could be "mha", "gqa", "mqa"'
+        assert self_att_type in ['mha', 'gqa', 'mqa', 'gma', 'dma', 'sqa'], 'Self-attention type could be "mha", "gqa", "mqa", "gma", "dma", "sqa".'
+        assert cross_att_type in ['mha', 'gqa', 'mqa', 'gma', 'dma', 'sqa'], 'Memory cross-attention type could be "mha", "gqa", "mqa", "gma", "dma", "sqa".'
 
         embedding = nn.Embedding(vocab_size, embed_dim)
         rope = RotaryPositionalEmbedding(embed_dim // att_heads, seq_len)
@@ -73,6 +79,28 @@ class RxTAlphaComponentBase(nn.Module, PyTorchModelHubMixin):
 
         ff_activation = get_activation_layer(ff_activation)
 
+        if self_att_type in ['mha', 'gqa', 'mqa']:
+            att_init = lambda: init_attention(embed_dim, att_heads, self_att_type, att_groups, rope=rope,
+                                              use_flash_attention=use_flash_attention, dropout=att_dropout,
+                                              max_seq_len=seq_len, is_causal=True)
+        else:
+            att_init = lambda: init_experimental_attention(embed_dim, att_heads, self_att_type, att_groups, rope=rope,
+                                                           use_flash_attention=use_flash_attention, dropout=att_dropout,
+                                                           max_seq_len=seq_len, is_causal=True, num_experts=att_num_experts,
+                                                           num_query_experts=att_num_query_experts,
+                                                           num_query_groups=att_num_query_groups)
+
+        if cross_att_type in ['mha', 'gqa', 'mqa']:
+            cross_att_init = lambda: init_attention(embed_dim, att_heads, self_att_type, att_groups, rope=rope,
+                                              use_flash_attention=use_flash_attention, dropout=att_dropout,
+                                              max_seq_len=seq_len, is_causal=True)
+        else:
+            cross_att_init = lambda: init_experimental_attention(embed_dim, att_heads, self_att_type, att_groups, rope=rope,
+                                                           use_flash_attention=use_flash_attention, dropout=att_dropout,
+                                                           max_seq_len=seq_len, is_causal=True, num_experts=att_num_experts,
+                                                           num_query_experts=att_num_query_experts,
+                                                           num_query_groups=att_num_query_groups)
+
         layers = nn.ModuleList([
             ReactiveTransformerLayer(
                 embed_dim,
@@ -84,13 +112,8 @@ class RxTAlphaComponentBase(nn.Module, PyTorchModelHubMixin):
                 ff_activation=ff_activation,
                 ff_dropout=ff_dropout,
                 use_rms_norm=use_rms_norm,
-                self_attention=init_attention(embed_dim, att_heads, self_att_type, att_groups, rope=rope,
-                                              use_flash_attention=use_flash_attention, dropout=att_dropout,
-                                              max_seq_len=seq_len, is_causal=is_causal),
-                memory_cross_attention=init_attention(embed_dim, att_heads, cross_att_type, att_groups, rope=rope,
-                                                      use_flash_attention=use_flash_attention, dropout=att_dropout,
-                                                      max_seq_len=seq_len, rope_only_for_query=True,
-                                                      is_causal=is_causal)
+                self_attention=att_init(),
+                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)

{rxnn-0.1.55.dist-info → rxnn-0.1.57.dist-info}/METADATA

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

{rxnn-0.1.55.dist-info → rxnn-0.1.57.dist-info}/RECORD

@@ -1,13 +1,13 @@
 rxnn/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 rxnn/experimental/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-rxnn/experimental/attention.py,sha256=oPknT_PVcNwZvDwpZM7gmP4M_md_FW8oYwJDdQk1avM,34544
-rxnn/experimental/models.py,sha256=iprFSQDPK75zebDJBJ1i-mnNS9jlGf9RAIk-S0E9D-Q,4689
+rxnn/experimental/attention.py,sha256=ivIqIc-15DWA_q-ITy2iaYmB7tffKVtiuqjdSH3mtS4,34548
+rxnn/experimental/models.py,sha256=_i9kvQsAYPyMQo2VfMUTmtBs-mE2w75j1X-OHx03IJk,4743
 rxnn/experimental/moe.py,sha256=jHZ1QhpWiVQOswVpFmuH7b2IUOPf0Uuf-I2Ddwsd7Us,6140
 rxnn/memory/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 rxnn/memory/norm.py,sha256=Ofl8Q5NYEF9GQeO0bhM43tkTW91J0y6TSvTAOYMgloM,6278
 rxnn/memory/stm.py,sha256=EsD8slSP4_9dLuq6aFPDmuFe8PWilxh90so5Z3nm-ig,2057
 rxnn/rxt/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-rxnn/rxt/models.py,sha256=INTFeNcqzAsjyWhNtbBHL4Tx7tYDsaQHgm72tf6u20M,6918
+rxnn/rxt/models.py,sha256=L5RvhORONmYSF_pVjP8HwiSeAypSNqfFi6Fogp2oJes,8543
 rxnn/training/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 rxnn/training/base.py,sha256=gEWASLSuWR8UF8b2e-DYqkBZ1lBx0VsIm4kGf9eWSHM,11678
 rxnn/training/bml.py,sha256=S1ZaXTybzeJH7uVFamCr4TPl2bLyZ5xmn_lSsjThTiM,19162
@@ -25,7 +25,7 @@ rxnn/transformers/moe.py,sha256=j6jEx6Ip0zttlUZKKn82azxo95lkLZs-H2GLSMD88hY,5859
 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.55.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
-rxnn-0.1.55.dist-info/METADATA,sha256=qiMp63aMlBdbvVvSJDL2bfW5XoR0PzNwN6pWdkfCuOM,16627
-rxnn-0.1.55.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
-rxnn-0.1.55.dist-info/RECORD,,
+rxnn-0.1.57.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
+rxnn-0.1.57.dist-info/METADATA,sha256=K9kcLSS3CUYwFg9N-KhPK5J4tMmsgKYFqF6VkH8689U,16627
+rxnn-0.1.57.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
+rxnn-0.1.57.dist-info/RECORD,,