rxnn 0.1.10__tar.gz → 0.1.12__tar.gz

{rxnn-0.1.10 → rxnn-0.1.12}/PKG-INFO

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

{rxnn-0.1.10 → rxnn-0.1.12}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "rxnn"
-version = "0.1.10"
+version = "0.1.12"
 description = "RxNN: Reactive Neural Networks Platform"
 
 license = "Apache-2.0"

{rxnn-0.1.10 → rxnn-0.1.12}/src/rxnn/training/base.py

@@ -49,6 +49,7 @@ class BaseTrainer(ABC):
         self.validation_metrics = {}
         self.target_field_name = target_field_name
         self.total_tokens = 0
+        self.total_steps = 0
         self.gradient_accumulation_steps = gradient_accumulation_steps
         self.accumulated_loss = 0.0
         self.optimizer_step_count = 0
@@ -140,6 +141,7 @@ class BaseTrainer(ABC):
                 for callback in self.callbacks:
                     callback.on_batch_start(self.model, batch_idx, batch)
                 if self.get_batch_size(batch) == batch_size:
+                    self.total_steps += 1
                     loss = self.train_step(batch, batch_idx)
                     orig_loss = loss.item()
                     self.accumulated_loss += orig_loss
@@ -226,11 +228,11 @@ class BaseTrainer(ABC):
             self.writer.close()
 
     def _valid_writer(self, epoch: int, val_loss: float, val_metrics: dict):
-        self.writer.add_scalar('Loss/validation', val_loss, epoch)
-        self.writer.add_scalar('Perplexity/validation', math.exp(val_loss), epoch)
+        self.writer.add_scalar('Loss/Valid', val_loss, epoch)
+        self.writer.add_scalar('Perplexity/Valid', math.exp(val_loss), epoch)
         if val_metrics['accuracy']:
-            self.writer.add_scalar('Node Accuracy/validation', val_metrics['node_accuracy'], epoch)
-            self.writer.add_scalar('Avg. Accuracy/validation', val_metrics['accuracy'], epoch)
+            self.writer.add_scalar('Node Accuracy/Valid', val_metrics['node_accuracy'], epoch)
+            self.writer.add_scalar('Avg. Accuracy/Valid', val_metrics['accuracy'], epoch)
 
     def valid_step(self, batch: dict[str, torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]:
         if self.use_amp:

{rxnn-0.1.10 → rxnn-0.1.12}/src/rxnn/training/bml.py

@@ -1,6 +1,7 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+from torch.nn.parallel import DistributedDataParallel
 import math
 from huggingface_hub import PyTorchModelHubMixin
 from typing import Union
@@ -171,6 +172,90 @@ class AutoregressiveTrainer(BaseTrainer):
         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()
+
+        main_loss = F.cross_entropy(
+            shifted_logits.view(-1, self.vocab_size),
+            shifted_targets.view(-1)
+        )
+
+        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
+
+        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, 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 JointTrainingModel(nn.Module):
     def __init__(
             self,
@@ -262,18 +347,18 @@ class JointLMTrainer(BaseTrainer):
         return (encoder_loss, decoder_loss), (encoder_logits, decoder_logits)
 
     def _valid_writer(self, epoch: int, val_loss: float, val_metrics: dict):
-        self.writer.add_scalar('Loss/validation', val_loss, epoch)
-        self.writer.add_scalar('Perplexity/validation', math.exp(val_loss), epoch)
+        self.writer.add_scalar('Loss/Valid', val_loss, epoch)
+        self.writer.add_scalar('Perplexity/Valid', math.exp(val_loss), epoch)
         if val_metrics['accuracy']:
-            self.writer.add_scalar('Encoder node accuracy/validation', val_metrics['accuracy']['node_encoder'], epoch)
-            self.writer.add_scalar('Decoder node accuracy/validation', val_metrics['accuracy']['node_decoder'], epoch)
-            self.writer.add_scalar('Encoder avg. accuracy/validation', val_metrics['accuracy']['encoder'], epoch)
-            self.writer.add_scalar('Decoder avg. accuracy/validation', val_metrics['accuracy']['decoder'], epoch)
+            self.writer.add_scalar('Encoder node accuracy/Valid', val_metrics['accuracy']['node_encoder'], epoch)
+            self.writer.add_scalar('Decoder node accuracy/Valid', val_metrics['accuracy']['node_decoder'], epoch)
+            self.writer.add_scalar('Encoder avg. accuracy/Valid', val_metrics['accuracy']['encoder'], epoch)
+            self.writer.add_scalar('Decoder avg. accuracy/Valid', val_metrics['accuracy']['decoder'], epoch)
         if val_metrics['loss']:
-            self.writer.add_scalar('Encoder loss/validation', val_metrics['loss']['encoder'], epoch)
-            self.writer.add_scalar('Encoder perplexity/validation', math.exp(val_metrics['loss']['encoder']), epoch)
-            self.writer.add_scalar('Decoder accuracy/validation', val_metrics['loss']['decoder'], epoch)
-            self.writer.add_scalar('Decoder perplexity/validation', math.exp(val_metrics['loss']['decoder']), epoch)
+            self.writer.add_scalar('Encoder loss/Valid', val_metrics['loss']['encoder'], epoch)
+            self.writer.add_scalar('Encoder perplexity/Valid', math.exp(val_metrics['loss']['encoder']), epoch)
+            self.writer.add_scalar('Decoder accuracy/Valid', val_metrics['loss']['decoder'], epoch)
+            self.writer.add_scalar('Decoder perplexity/Valid', math.exp(val_metrics['loss']['decoder']), epoch)
 
     def validate(self, batch_size: int) -> tuple[float, dict]:
         self.model.eval()

{rxnn-0.1.10 → rxnn-0.1.12}/src/rxnn/transformers/layers.py

@@ -53,11 +53,15 @@ class ReactiveTransformerLayer(nn.Module):
             self.norm2 = nn.LayerNorm(embed_dim)
             self.norm3 = nn.LayerNorm(embed_dim)
         self.use_post_norm = use_post_norm
+        self.use_moe = use_moe
 
     def trainable_cross_attention_(self, is_trainable: bool):
         for param in self.memory_cross_attention.parameters():
             param.requires_grad_(is_trainable)
 
+    def moe_router_loss_(self):
+        return self.ff.router_loss() if self.use_moe else None
+
     def forward(self, x: torch.Tensor, stm: torch.Tensor, mask: torch.Tensor = None) -> torch.Tensor:
         # First step, self-attention
         residual = x
@@ -129,6 +133,10 @@ class ClassicTransformerLayer(nn.Module):
             self.norm1 = nn.LayerNorm(embed_dim)
             self.norm2 = nn.LayerNorm(embed_dim)
         self.use_post_norm = use_post_norm
+        self.use_moe = use_moe
+
+    def moe_router_loss_(self):
+        return self.ff.router_loss() if self.use_moe else torch.tensor(0.0)
 
     def forward(self, x: torch.Tensor, mask: torch.Tensor = None) -> torch.Tensor:
         # First step, self-attention

{rxnn-0.1.10 → rxnn-0.1.12}/src/rxnn/transformers/models.py

@@ -37,6 +37,10 @@ class ReactiveTransformerBase(nn.Module):
         for i in range(self.num_own_layers):
             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()
+
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         # Shared logic for encoders and decoders - apply embeddings and positional encoding
         x = self.embedding(x)
@@ -119,6 +123,9 @@ class ClassicTransformerBase(nn.Module):
         self.layers = layers
         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()
+
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         # Shared logic for encoders and decoders - apply embeddings and positional encoding
         x = self.embedding(x)

{rxnn-0.1.10 → rxnn-0.1.12}/src/rxnn/transformers/moe.py

@@ -11,7 +11,8 @@ class MoeRouter(nn.Module):
         self.top_k = top_k
         self.num_experts = num_experts
         self.gate = nn.Linear(embed_dim, num_experts, bias=False)
-        self.aux_loss = 0.0  # For expert load balancing
+        # For expert load balancing
+        self.register_buffer('aux_loss', torch.tensor(0.0), persistent=False)
 
     def forward(self, x: torch.Tensor):
         # x shape: [batch_size*seq_len, embed_dim]
@@ -19,10 +20,8 @@ class MoeRouter(nn.Module):
         probs = F.softmax(logits, dim=-1)
 
         # Expert load balancing loss
-        if self.training:
-            probs_for_bal = F.softmax(logits, dim=0)
-            self.aux_loss = (probs_for_bal.mean(dim=0) *
-                             torch.log(probs_for_bal.mean(dim=0) + 1e-9)).sum()
+        mean_probs = probs.mean(dim=0)  # Mean probability per expert across batch
+        self.aux_loss = (mean_probs * torch.log(mean_probs + 1e-9)).sum()  # Entropy-based loss
 
         top_k_weights, top_k_indices = probs.topk(self.top_k, dim=-1)
         top_k_weights = top_k_weights / (top_k_weights.sum(dim=-1, keepdim=True) + 1e-9)
@@ -74,6 +73,9 @@ class MoeFeedForward(nn.Module):
     def _activate(self, h: torch.Tensor):
         return self.activation(h)
 
+    def router_loss(self):
+        return self.router.aux_loss
+
     def forward(self, x: torch.Tensor):
         orig_shape = x.shape
         x = x.view(-1, self.embed_dim)  # [batch*seq_len, embed_dim]