rxnn 0.1.23__py3-none-any.whl → 0.1.25__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)
 
@@ -194,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)
 
@@ -212,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
@@ -227,298 +238,6 @@ class DeepMoeAttention(GroupedMoeAttention):
         # Key/Value processing
         return super()._forward_qkv(q, key, value, b, t, d, skip_query_processing=True)
 
-# Vectorized
-
-class GroupedMoeAttentionVectorized(GroupedQueryAttention):
-    """
-    Vectorized implementation calculates all expert heads for each token and selecting active tokens later. Linear layers
-    for Attention are rather small, compared to MoE Feed Forward layers, so it's possible that it will be faster than filtering
-    experts - it has to be tested.
-
-    Grouped MoE Attention (GMA) - GQA extended with Mixture-of-Experts (MoE) routing.
-
-    Instead of mapping keys/values to static head groups, it dynamically selects head expert groups. It has the same
-    number of total keys/values heads as query heads, but uses only a selected group for attention calculation.
-    - with num_groups set to 1, it will be MoE MultiQueryAttention
-
-    Compared to traditional GQA/MQA, it should provide better performance, because lot less data could be lost using
-    this approach - we are training the full number of keys/values heads, while using only a group.
-
-    In case of efficiency, it should be close to GQA/MQA linear performance, but with a small MoE routing overhead.
-
-    Optionally, it could use even more expert heads than attention heads - in example:
-    - 512 dim divided into 16 heads with 32 dim, using 4 head groups - may use i.e., 24 total expert heads - still only
-    4 will be used for attention calculation, while 16 is used to split dimensions (in that case it will have 16 query heads)
-
-    © 2025 Adam Filipek
-    """
-
-    def __init__(
-            self,
-            embed_dim: int,
-            num_heads: int,
-            num_groups: int,
-            dropout: float = 0.0,
-            rope: RotaryPositionalEmbedding = None,
-            rope_only_for_query: bool = False,
-            use_relative_embeddings: bool = False,
-            max_seq_len: int = 1024,
-            use_flash_attention: bool = False,
-            is_causal: bool = False,
-            use_bias: bool = False,
-            num_experts: int = None,
-            *args,
-            **kwargs,
-    ):
-        self.num_experts = num_experts if num_experts is not None else num_heads
-        super(GroupedMoeAttentionVectorized, self).__init__(
-            embed_dim,
-            num_heads,
-            num_groups=num_groups,
-            dropout=dropout,
-            rope=rope,
-            rope_only_for_query=rope_only_for_query,
-            use_relative_embeddings=use_relative_embeddings,
-            max_seq_len=max_seq_len,
-            use_flash_attention=use_flash_attention,
-            is_causal=is_causal,
-            use_bias=use_bias,
-            *args,
-            **kwargs,
-        )
-
-    def _init_kv(self, embed_dim: int):
-        self.router = MoeRouter(embed_dim, self.num_experts, top_k=self.num_groups)
-        hidden_dim = embed_dim // self.num_heads
-        self.wk = nn.Parameter(torch.empty(self.num_experts, embed_dim, hidden_dim))
-        self.bk = nn.Parameter(torch.zeros(self.num_experts, hidden_dim)) if self.use_bias else None
-        self.wv = nn.Parameter(torch.empty(self.num_experts, embed_dim, hidden_dim))
-        self.bv = nn.Parameter(torch.zeros(self.num_experts, hidden_dim)) if self.use_bias else None
-        self._init_experts()
-
-    def _init_experts(self):
-        torch.nn.init.xavier_uniform_(self.wk)
-        torch.nn.init.xavier_uniform_(self.wv)
-        if self.use_bias:
-            torch.nn.init.zeros_(self.bk)
-            torch.nn.init.zeros_(self.bv)
-
-    def _forward_qkv(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, b: int, t: int, d: int,
-                     skip_query_processing: bool = False):
-        # Indexed version may cause memory overflow
-        #
-        # head_dim = d // self.num_heads
-        #
-        # # Process Query as in GQA
-        # q = self.q_proj(query).view(b, t, self.num_heads, head_dim).transpose(1,
-        #                                                                       2) if not skip_query_processing else query
-        #
-        # # Process Key and Value with MoE routing
-        # key_flat = key.view(-1, d)  # (B*S, d)
-        # value_flat = value.view(-1, d)  # (B*S, d)
-        #
-        # # Get routing indices and weights for K
-        # weights_k, indices_k = self.router(key_flat)
-        # indices_k = indices_k.view(-1, self.top_k)  # (B*S, top_k)
-        # weights_k = weights_k.view(-1, self.top_k, 1)  # (B*S, top_k, 1)
-        #
-        # # Select and compute K projections for only the top_k experts
-        # selected_k_weights = self.k_experts[indices_k]  # (B*S, top_k, d, k_out_dim)
-        # k_proj = torch.einsum('bd, behd -> beh', key_flat.unsqueeze(1), selected_k_weights)
-        # selected_k = (k_proj * weights_k).sum(dim=1)  # (B*S, k_out_dim)
-        # selected_k = selected_k.view(b, key.size(1), -1)  # (B, S, k_out_dim)
-        #
-        # # Compute V using the same indices as K (since they share the same router)
-        # selected_v_weights = self.v_experts[indices_k]
-        # v_proj = torch.einsum('bd, behd -> beh', value_flat.unsqueeze(1), selected_v_weights)
-        # selected_v = (v_proj * weights_k).sum(dim=1)
-        # selected_v = selected_v.view(b, value.size(1), -1)  # (B, S, k_out_dim)
-        #
-        # # Reshape to GQA format: (B, G, S, head_dim)
-        # k = selected_k.view(b, key.size(1), self.num_groups, head_dim).transpose(1, 2)
-        # v = selected_v.view(b, value.size(1), self.num_groups, head_dim).transpose(1, 2)
-        #
-        # if not self.use_flash_attention:
-        #     group_heads = self.num_heads // self.num_groups
-        #
-        #     k = k.unsqueeze(2).expand(-1, -1, group_heads, -1, -1)  # (B, G, group_heads, S, head_dim)
-        #     v = v.unsqueeze(2).expand(-1, -1, group_heads, -1, -1)  # (B, G, group_heads, S, head_dim)
-        #
-        #     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)
-        #
-        # return q, k, v
-
-        head_dim = d // self.num_heads
-
-        # Process Query as in GQA
-        q = self.q_proj(query).view(b, t, self.num_heads, head_dim).transpose(1, 2)
-
-        # Process Key and Value with MoE routing
-        key_flat = key.view(-1, d)
-        weights, indices = self.router(key_flat)
-        weights = weights.view(b, key.size(1), self.num_groups, 1)
-        indices = indices.view(b, key.size(1), self.num_groups)
-
-        # Compute all experts' K and V projections
-        # Shape: (batch_size, seq_len, num_experts, head_dim * num_groups)
-        k_all = torch.einsum(
-            'be, ehd -> bedh',
-            key_flat,
-            self.wk.view(self.num_experts, d, -1)
-        ).view(b, key.size(1), self.num_experts, -1)
-
-        v_all = torch.einsum(
-            'be, ehd -> bedh',
-            value.view(-1, d),
-            self.wv.view(self.num_experts, d, -1)
-        ).view(b, value.size(1), self.num_experts, -1)
-
-        # Select top_k experts and compute weighted sum
-        selected_k = torch.gather(
-            k_all,
-            2,
-            indices.unsqueeze(-1).expand(-1, -1, -1, k_all.size(-1))
-        )
-        selected_v = torch.gather(
-            v_all,
-            2,
-            indices.unsqueeze(-1).expand(-1, -1, -1, v_all.size(-1))
-        )
-
-        selected_k = (selected_k * weights).sum(dim=2)
-        selected_v = (selected_v * weights).sum(dim=2)
-        # Reshape to GQA format: (B, G, S, head_dim)
-        k = selected_k.view(b, key.size(1), self.num_groups, head_dim).transpose(1, 2)
-        v = selected_v.view(b, value.size(1), self.num_groups, head_dim).transpose(1, 2)
-
-        if not self.use_flash_attention:
-            group_heads = self.num_heads // self.num_groups
-
-            k = k.unsqueeze(2).expand(-1, -1, group_heads, -1, -1)  # (B, G, group_heads, S, head_dim)
-            v = v.unsqueeze(2).expand(-1, -1, group_heads, -1, -1)  # (B, G, group_heads, S, head_dim)
-
-            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)
-
-        return q, k, v
-
-
-class DeepMoeAttentionVectorized(GroupedMoeAttentionVectorized):
-    """
-        Vectorized implementation calculates all expert heads for each token and selecting active tokens later. Linear layers
-    for Attention are rather small, compared to MoE Feed Forward layers, so it's possible that it will be faster than filtering
-    experts - it has to be tested.
-
-    Deep MoE Attention (SMA) - Grouped MoE Attention extended even more for sublinear computational efficiency.
-
-    In addition to using Mixture-of-Experts (MoE) for key/value head groups, SMA is also using dynamically selected
-    query heads - with that approach, each token could attend to every other token, but only partially - only some part of
-    information from each token is used to identify related information parts from other tokens. So, DMA is not spatially
-    sparse (has access to all tokens), but rather structurally sparse (has access only to the part of token's information).
-
-    This solution could reduce the computational complexity of attention operation to sublinear level (<O(N)) and provide
-    a viable and efficient alternative to spatial sparse attention mechanisms like Flex Attention.
-
-    © 2025 Adam Filipek
-    """
-
-    def __init__(
-            self,
-            embed_dim: int,
-            num_heads: int,
-            num_groups: int,
-            dropout: float = 0.0,
-            rope: RotaryPositionalEmbedding = None,
-            rope_only_for_query: bool = False,
-            use_relative_embeddings: bool = False,
-            max_seq_len: int = 1024,
-            use_flash_attention: bool = False,
-            is_causal: bool = False,
-            use_bias: bool = False,
-            num_experts: int = None,
-            num_query_experts: int = None,
-            num_query_groups: int = None,
-            *args,
-            **kwargs,
-    ):
-        self.num_query_experts = num_query_experts if num_query_experts is not None else num_heads
-        self.num_query_groups = num_query_groups if num_query_groups is not None else num_groups
-        super(DeepMoeAttentionVectorized, self).__init__(
-            embed_dim,
-            num_heads,
-            num_groups=num_groups,
-            dropout=dropout,
-            rope=rope,
-            rope_only_for_query=rope_only_for_query,
-            use_relative_embeddings=use_relative_embeddings,
-            max_seq_len=max_seq_len,
-            use_flash_attention=use_flash_attention,
-            is_causal=is_causal,
-            use_bias=use_bias,
-            num_experts=num_experts,
-            *args,
-            **kwargs,
-        )
-
-    def _init_q(self, embed_dim: int):
-        self.query_router = MoeRouter(embed_dim, self.num_query_experts, top_k=self.num_query_groups)
-        hidden_dim = embed_dim // self.num_heads
-        self.wq = nn.Parameter(torch.empty(self.num_query_experts, embed_dim, hidden_dim))
-        self.bq = nn.Parameter(torch.zeros(self.num_query_experts, hidden_dim)) if self.use_bias else None
-        self._init_query_experts()
-
-    def _init_query_experts(self):
-        torch.nn.init.xavier_uniform_(self.wq)
-        if self.use_bias:
-            torch.nn.init.zeros_(self.bq)
-
-    def _init_out(self, embed_dim: int):
-        """Initialize output projection"""
-        self.out_proj = nn.Linear(embed_dim // (self.num_heads // self.num_groups), embed_dim)
-
-    def _forward_qkv(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, b: int, t: int, d: int):
-        # Indexed version may cause memory overflow
-        #
-        # head_dim = d // self.num_heads
-        #
-        # # Process Query with MoE routing
-        # query_flat = query.view(-1, d)  # (B*T, d)
-        # weights_q, indices_q = self.query_router(query_flat)
-        # indices_q = indices_q.view(-1, self.num_query_groups)  # (B*T, top_k_q)
-        # weights_q = weights_q.view(-1, self.num_query_groups, 1)  # (B*T, top_k_q, 1)
-        #
-        # # Select and compute Q projections for top_k experts
-        # selected_q_weights = self.wq[indices_q]  # (B*T, top_k_q, d, head_dim*num_heads)
-        # q_proj = torch.einsum('bd, behd -> beh', query_flat.unsqueeze(1), selected_q_weights)
-        # selected_q = (q_proj * weights_q).sum(dim=1)  # (B*T, head_dim*num_heads)
-        # selected_q = selected_q.view(b, t, -1)  # (B, T, head_dim*num_heads)
-        head_dim = d // self.num_heads
-
-        # Process Query with MoE routing
-        query_flat = query.view(b * t, d)
-        weights_q, indices_q = self.query_router(query_flat)
-        weights_q = weights_q.view(b, t, self.num_query_groups, 1)
-        indices_q = indices_q.view(b, t, self.num_query_groups)
-
-        # Compute all experts' Q projections
-        q_all = torch.einsum(
-            'be, ehd -> bedh',
-            query_flat,
-            self.wq.view(self.num_query_experts, d, -1)
-        ).view(b, t, self.num_query_experts, -1)
-
-        selected_q = torch.gather(
-            q_all,
-            2,
-            indices_q.unsqueeze(-1).expand(-1, -1, -1, q_all.shape[-1])
-        )
-        selected_q = (selected_q * weights_q).sum(dim=2)
-
-        q = selected_q.view(b, t, self.num_heads, head_dim).transpose(1, 2)  # (B, H, T, head_dim)
-
-        return super()._forward_qkv(q, key, value, b, t, d, skip_query_processing=True)
-
 
 # Others
 
@@ -670,8 +389,7 @@ def init_moe_attention(
         num_query_experts: int = None,
         num_query_groups: int = None,
 ) -> GroupedQueryAttention:
-    assert attention_type == 'gma' or attention_type == 'dma' or attention_type == 'gma_v' or attention_type == 'dma_v', \
-        "Error, attention type should be one of: 'gma', 'dma', 'gma_v', 'dma_v'"
+    assert attention_type in ['gma', 'dma'], "Error, attention type should be one of: 'gma', 'dma', 'gma_v', 'dma_v'"
 
     if attention_type == "gma":
         return GroupedMoeAttention(
@@ -688,7 +406,7 @@ def init_moe_attention(
             use_bias=use_bias,
             num_experts=num_experts,
         )
-    elif attention_type == "dma":
+    else:
         return DeepMoeAttention(
             embed_dim,
             num_heads,
@@ -705,35 +423,3 @@ def init_moe_attention(
             num_query_experts=num_query_experts,
             num_query_groups=num_query_groups,
         )
-    elif attention_type == "gma_v":
-        return GroupedMoeAttentionVectorized(
-            embed_dim,
-            num_heads,
-            gqa_groups,
-            dropout=dropout,
-            rope=rope,
-            use_relative_embeddings=use_relative_embeddings,
-            max_seq_len=max_seq_len,
-            rope_only_for_query=rope_only_for_query,
-            use_flash_attention=use_flash_attention,
-            is_causal=is_causal,
-            use_bias=use_bias,
-            num_experts=num_experts,
-        )
-    else:
-        return DeepMoeAttentionVectorized(
-            embed_dim,
-            num_heads,
-            gqa_groups,
-            dropout=dropout,
-            rope=rope,
-            use_relative_embeddings=use_relative_embeddings,
-            max_seq_len=max_seq_len,
-            rope_only_for_query=rope_only_for_query,
-            use_flash_attention=use_flash_attention,
-            is_causal=is_causal,
-            use_bias=use_bias,
-            num_experts=num_experts,
-            num_query_experts=num_query_experts,
-            num_query_groups=num_query_groups,
-        )

rxnn/experimental/models.py

@@ -35,7 +35,7 @@ class MoeAttentionTransformerConfig(TypedDict):
 
 
 class MoeAttentionTransformer(nn.Module, PyTorchModelHubMixin, pipeline_tag="text-generation", license="apache-2.0"):
-    """Research model for experiments with Mixture-of-Experts Attention"""
+    """Research decoder model for experiments with Mixture-of-Experts Attention"""
 
     def __init__(
             self,
@@ -65,8 +65,7 @@ class MoeAttentionTransformer(nn.Module, PyTorchModelHubMixin, pipeline_tag="tex
         assert ff_activation in ['relu', 'gelu',
                                  'swish', 'silu', 'linear',
                                  'sigmoid'], 'Feed-forward activation could be "relu", "gelu", "swish", "silu", "linear", "sigmoid".'
-        assert att_type in ['mha', 'gqa', 'mqa', 'gma', 'dma', 'gma_v',
-                            'dma_v'], 'Self-attention type could be "mha", "gqa", "mqa", "gma", "dma", "gma_v", "dma_v"'
+        assert att_type in ['mha', 'gqa', 'mqa', 'gma', 'dma'], 'Self-attention type could be "mha", "gqa", "mqa", "gma", "dma"'
 
         embedding = nn.Embedding(vocab_size, embed_dim)
         rope = RotaryPositionalEmbedding(embed_dim // att_heads, seq_len)
@@ -111,6 +110,5 @@ class MoeAttentionTransformer(nn.Module, PyTorchModelHubMixin, pipeline_tag="tex
     def load_shared_embedding(self, embedding: nn.Embedding):
         self.model.embedding = embedding
 
-    def forward(self, x: torch.Tensor, attention_mask: torch.Tensor = None) -> Union[
-        torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
+    def forward(self, x: torch.Tensor, attention_mask: torch.Tensor = None) -> torch.Tensor:
         return self.model(x, attention_mask=attention_mask)

rxnn/training/bml.py

@@ -50,10 +50,14 @@ class MLMTrainer(BaseTrainer):
             vocab_size: int,
             use_amp: bool = False,
             dtype: torch.dtype = None,
+            use_moe_aux_loss: bool = False,
+            moe_aux_loss_scale: float = 0.01,
             **kwargs
     ):
         super(MLMTrainer, self).__init__(model, device, use_amp=use_amp, dtype=dtype, **kwargs)
         self.vocab_size = vocab_size
+        self.use_moe_aux_loss = use_moe_aux_loss
+        self.moe_aux_loss_scale = moe_aux_loss_scale
 
     def compute_loss(self, batch: dict[str, torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]:
         inputs = batch['input_ids']
@@ -65,11 +69,32 @@ class MLMTrainer(BaseTrainer):
             attention_mask=attention_mask
         )
 
-        return F.cross_entropy(
+        loss = F.cross_entropy(
             logits.view(-1, self.vocab_size),
             labels.view(-1),
             ignore_index=-100
-        ), logits
+        )
+
+        return self._moe_aux_loss(loss), logits
+
+    def _moe_aux_loss(self, main_loss: torch.Tensor) -> torch.Tensor:
+        if not self.use_moe_aux_loss:
+            return main_loss
+
+        model = next(self.model.children()) if isinstance(self.model, DistributedDataParallel) else self.model
+
+        router_loss = model.encoder.model.moe_router_loss()
+        loss = main_loss + self.moe_aux_loss_scale * router_loss
+
+        if self.writer is not None:
+            if self.model.training:
+                self.writer.add_scalar('Router aux loss/Train', router_loss.item(), self.total_steps)
+                self.writer.add_scalar('Model loss/Train', main_loss.item(), self.total_steps)
+            else:
+                self.writer.add_scalar('Router aux loss/Valid', router_loss.item(), self.total_steps)
+                self.writer.add_scalar('Model loss/Valid', main_loss.item(), self.total_steps)
+
+        return loss
 
     def validate(self, batch_size: int) -> tuple[float, dict]:
         self.model.eval()
@@ -113,11 +138,15 @@ class AutoregressiveTrainer(BaseTrainer):
             vocab_size: int,
             use_amp: bool = False,
             dtype: torch.dtype = None,
+            use_moe_aux_loss: bool = False,
+            moe_aux_loss_scale: float = 0.01,
             **kwargs
     ):
         super(AutoregressiveTrainer, self).__init__(model, device, use_amp=use_amp, dtype=dtype,
                                                     target_field_name='targets', **kwargs)
         self.vocab_size = vocab_size
+        self.use_moe_aux_loss = use_moe_aux_loss
+        self.moe_aux_loss_scale = moe_aux_loss_scale
 
     def compute_loss(self, batch: dict[str, torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]:
         inputs = batch['input_ids']
@@ -132,84 +161,21 @@ class AutoregressiveTrainer(BaseTrainer):
         shifted_logits = outputs[:, :-1].contiguous()
         shifted_targets = targets[:, 1:].contiguous()
 
-        return F.cross_entropy(
+        loss = F.cross_entropy(
             shifted_logits.view(-1, self.vocab_size),
             shifted_targets.view(-1)
-        ), outputs
-
-    def validate(self, batch_size: int) -> tuple[float, dict]:
-        self.model.eval()
-        val_dataloader = self._valid_loader(batch_size)
-        val_loss = torch.tensor(0.0).to(self.device)
-        correct = torch.tensor(0).to(self.device)
-        total = torch.tensor(0).to(self.device)
-
-        with torch.no_grad():
-            for batch in val_dataloader:
-                if self.get_batch_size(batch) == batch_size:
-                    loss, logits = self.valid_step(batch)
-                    val_loss += loss
-                    shifted_logits = logits[:, :-1].contiguous()
-                    shifted_targets = batch[self.target_field_name][:, 1:].to(self.device).contiguous()
-                    valid_indices = shifted_targets != -100
-                    if valid_indices.any():
-                        preds = shifted_logits.argmax(-1)
-                        correct += (preds[valid_indices] == shifted_targets[valid_indices]).sum()
-                        total += valid_indices.sum()
-
-        avg_loss = (val_loss / len(val_dataloader)).item()
-        acc = (correct / total * 100) if total > 0 else torch.tensor(0.0).to(self.device)
-        node_acc = acc.item()
-        if self.use_ddp:
-            dist.all_reduce(acc, op=dist.ReduceOp.SUM)
-            acc = acc / dist.get_world_size()
-
-        metrics = {
-            'accuracy': acc.item(),
-            'node_accuracy': node_acc,
-        }
-        self.model.train()
-        return avg_loss, metrics
-
-
-class AutoregressiveMoeTrainer(BaseTrainer):
-    def __init__(
-            self,
-            model: ReactiveTransformerDecoder,
-            device: torch.device,
-            vocab_size: int,
-            use_amp: bool = False,
-            dtype: torch.dtype = None,
-            router_loss_scale: float = 0.1,
-            **kwargs
-    ):
-        super(AutoregressiveMoeTrainer, self).__init__(model, device, use_amp=use_amp, dtype=dtype,
-                                                    target_field_name='targets', **kwargs)
-        self.vocab_size = vocab_size
-        self.router_loss_scale = router_loss_scale
-
-    def compute_loss(self, batch: dict[str, torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]:
-        inputs = batch['input_ids']
-        attention_mask = batch['attention_mask']
-        targets = batch['targets']
-
-        outputs = self.model(
-            inputs,
-            attention_mask=attention_mask
         )
 
-        shifted_logits = outputs[:, :-1].contiguous()
-        shifted_targets = targets[:, 1:].contiguous()
+        return self._moe_aux_loss(loss), outputs
 
-        main_loss = F.cross_entropy(
-            shifted_logits.view(-1, self.vocab_size),
-            shifted_targets.view(-1)
-        )
+    def _moe_aux_loss(self, main_loss: torch.Tensor) -> torch.Tensor:
+        if not self.use_moe_aux_loss:
+            return main_loss
 
         model = next(self.model.children()) if isinstance(self.model, DistributedDataParallel) else self.model
 
         router_loss = model.model.moe_router_loss()
-        loss = main_loss + self.router_loss_scale * router_loss
+        loss = main_loss + self.moe_aux_loss_scale * router_loss
 
         if self.writer is not None:
             if self.model.training:
@@ -219,7 +185,7 @@ class AutoregressiveMoeTrainer(BaseTrainer):
                 self.writer.add_scalar('Router aux loss/Valid', router_loss.item(), self.total_steps)
                 self.writer.add_scalar('Model loss/Valid', main_loss.item(), self.total_steps)
 
-        return loss, outputs
+        return loss
 
     def validate(self, batch_size: int) -> tuple[float, dict]:
         self.model.eval()
@@ -279,6 +245,9 @@ class JointTrainingModel(nn.Module):
 
 
 class JointLMTrainer(BaseTrainer):
+    """"
+    It's not recommended to use Joint LM Training in current implementation. More info soon
+    """
     def __init__(
             self,
             model: JointTrainingModel,

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.23.dist-info → rxnn-0.1.25.dist-info}/METADATA

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

{rxnn-0.1.23.dist-info → rxnn-0.1.25.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=gMEcFJHGOkz8R_s4dGEJB5cb2K3pbXZi4XBwyhEdB4s,31967
-rxnn/experimental/models.py,sha256=-XkEHsyT8iNAjhZbgC7N_5nzP4ENVJLwxSoLHgMfA0I,4668
+rxnn/experimental/attention.py,sha256=rJSQjA7_9YqcM4Y8SyJSuZQjyz8j4XPhC5jcrcrRK2M,17891
+rxnn/experimental/models.py,sha256=8KAo7BtRkke9qRlzGRtQa9-EZ34roGWrn0N_T6L-6Wc,4561
 rxnn/experimental/moe.py,sha256=PhiaNr3FwR2Zv2a0tfj6sfZ4iyhLo3Jyp2DwXq19qZQ,7935
 rxnn/memory/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 rxnn/memory/norm.py,sha256=Ofl8Q5NYEF9GQeO0bhM43tkTW91J0y6TSvTAOYMgloM,6278
@@ -10,22 +10,22 @@ rxnn/rxt/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 rxnn/rxt/models.py,sha256=INTFeNcqzAsjyWhNtbBHL4Tx7tYDsaQHgm72tf6u20M,6918
 rxnn/training/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 rxnn/training/base.py,sha256=QD8uS14jSyR5Y_8BgCaBQTKpsarerU3lyufsWsCq_6o,11227
-rxnn/training/bml.py,sha256=o_88ZL1YWd5gWXaBqYPK2UzSTbJaiTiw96E6z73LeOQ,18660
+rxnn/training/bml.py,sha256=HtxSzI-WcpRclAuIccF_WoTZ24KzH5ZWKe8SnWgjjm4,17581
 rxnn/training/callbacks.py,sha256=_YfMKY_eFdc-tubhO9nYH2PXDZDQwlSI74FVOoCXpQg,22108
 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.23.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
-rxnn-0.1.23.dist-info/METADATA,sha256=rZSBuoIgf8jKB11LKgMg7U42Wx7VNT_4EU3FVyED2YQ,16627
-rxnn-0.1.23.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
-rxnn-0.1.23.dist-info/RECORD,,
+rxnn-0.1.25.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
+rxnn-0.1.25.dist-info/METADATA,sha256=kRnBikg0u4uhy5IwcFpEE3gFEB8wTvEpg6cFcAU8FGs,16627
+rxnn-0.1.25.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
+rxnn-0.1.25.dist-info/RECORD,,