rxnn 0.1.9__py3-none-any.whl → 0.1.11__py3-none-any.whl

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,10 +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('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/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
@@ -90,12 +91,15 @@ class MLMTrainer(BaseTrainer):
                         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(correct, op=dist.ReduceOp.SUM)
-            dist.all_reduce(total, op=dist.ReduceOp.SUM)
+            dist.all_reduce(acc, op=dist.ReduceOp.SUM)
+            acc = acc / dist.get_world_size()
 
         metrics = {
-            'accuracy': (correct / total * 100).item() if total > 0 else 0.0
+            'accuracy': acc.item(),
+            'node_accuracy': node_acc,
         }
         self.model.train()
         return avg_loss, metrics
@@ -154,13 +158,99 @@ class AutoregressiveTrainer(BaseTrainer):
                         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()
 
+        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.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(correct, op=dist.ReduceOp.SUM)
-            dist.all_reduce(total, op=dist.ReduceOp.SUM)
+            dist.all_reduce(acc, op=dist.ReduceOp.SUM)
+            acc = acc / dist.get_world_size()
 
         metrics = {
-            'accuracy': (correct / total * 100).item() if total > 0 else 0.0
+            'accuracy': acc.item(),
+            'node_accuracy': node_acc,
         }
         self.model.train()
         return avg_loss, metrics
@@ -257,16 +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 accuracy/validation', val_metrics['accuracy']['encoder'], epoch)
-            self.writer.add_scalar('Decoder 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()
@@ -317,28 +409,30 @@ class JointLMTrainer(BaseTrainer):
         avg_loss = val_loss / loader_len
         avg_dec_loss = dec_loss / loader_len
         avg_enc_loss = enc_loss / loader_len
-
+        mlm_acc = (correct_mlm / total_mlm * 100) if total_mlm > 0 else torch.tensor(0.0).to(self.device)
+        alm_acc = (correct_alm / total_alm * 100) if total_alm > 0 else torch.tensor(0.0).to(self.device)
+        node_mlm_acc = mlm_acc.item()
+        node_alm_acc = alm_acc.item()
         if self.use_ddp:
             dist.all_reduce(avg_dec_loss, op=dist.ReduceOp.SUM)
             dist.all_reduce(avg_enc_loss, op=dist.ReduceOp.SUM)
-            dist.all_reduce(correct_mlm, op=dist.ReduceOp.SUM)
-            dist.all_reduce(total_mlm, op=dist.ReduceOp.SUM)
-            dist.all_reduce(correct_alm, op=dist.ReduceOp.SUM)
-            dist.all_reduce(total_alm, op=dist.ReduceOp.SUM)
+            dist.all_reduce(mlm_acc, op=dist.ReduceOp.SUM)
+            dist.all_reduce(alm_acc, op=dist.ReduceOp.SUM)
             avg_dec_loss = avg_dec_loss / dist.get_world_size()
             avg_enc_loss = avg_enc_loss / dist.get_world_size()
-
-        mlm_acc = (correct_mlm / total_mlm * 100).item() if total_mlm > 0 else 0.0
-        alm_acc = (correct_alm / total_alm * 100).item() if total_alm > 0 else 0.0
+            mlm_acc = mlm_acc / dist.get_world_size()
+            alm_acc = alm_acc / dist.get_world_size()
 
         metrics = {
             'accuracy': {
-                'encoder': mlm_acc,
-                'decoder': alm_acc,
+                'encoder': mlm_acc.item(),
+                'decoder': alm_acc.item(),
+                'node_encoder': node_mlm_acc,
+                'node_decoder': node_alm_acc,
             },
             'loss': {
-                'encoder': avg_enc_loss,
-                'decoder': avg_dec_loss,
+                'encoder': avg_enc_loss.item(),
+                'decoder': avg_dec_loss.item(),
             }
         }
         self.model.train()

rxnn/training/callbacks.py

@@ -83,9 +83,12 @@ class PrintAccuracyCallback(TrainerCallback):
 
     def on_validation_end(self, model: nn.Module, epoch: int, val_loss: float, val_metrics: dict) -> None:
         if self.joint_mode:
+            print(f"Epoch {epoch} - encoder node accuracy: {val_metrics['accuracy']['node_encoder']:.4f}")
+            print(f"Epoch {epoch} - decoder node accuracy: {val_metrics['accuracy']['node_decoder']:.4f}")
             print(f"Epoch {epoch} - encoder accuracy: {val_metrics['accuracy']['encoder']:.4f}")
             print(f"Epoch {epoch} - decoder accuracy: {val_metrics['accuracy']['decoder']:.4f}")
         else:
+            print(f"Epoch {epoch} - node accuracy: {val_metrics['node_accuracy']:.4f}")
             print(f"Epoch {epoch} - accuracy: {val_metrics['accuracy']:.4f}")
 
 
@@ -130,6 +133,7 @@ class ModelSaveCallback(TrainerCallback):
             save_checkpoint_after_n_batches: int = None,
             push_batch_checkpoint: bool = False,
             display_exc_trace: bool = False,
+            use_ddp: bool = False,
     ):
         self.save_dir = save_dir
         self.save_best_only = save_best_only
@@ -146,10 +150,11 @@ class ModelSaveCallback(TrainerCallback):
         self.push_batch_checkpoint = push_batch_checkpoint
         self.finished_epochs = 0
         self.display_exc_trace = display_exc_trace
+        self.rank = int(os.environ['RANK']) if use_ddp else 0
 
     def on_batch_end(self, model: torch.nn.Module, batch_idx: int, loss: int, batch: dict[str, torch.Tensor]) -> Union[
         bool, None]:
-        if self.save_checkpoint_after_n_batches is not None and batch_idx != 0 and batch_idx % self.save_checkpoint_after_n_batches == 0:
+        if self.rank == 0 and self.save_checkpoint_after_n_batches is not None and batch_idx != 0 and batch_idx % self.save_checkpoint_after_n_batches == 0:
             if isinstance(model, DistributedDataParallel):
                 model = next(model.children())
             try:
@@ -195,90 +200,92 @@ class ModelSaveCallback(TrainerCallback):
             val_loss: float,
             val_metrics: dict
     ):
-        self.finished_epochs += 1
-        if val_loss < self.best_loss:
-            self.best_loss = val_loss
+        if self.rank == 0:
+            self.finished_epochs += 1
+            if val_loss < self.best_loss:
+                self.best_loss = val_loss
+                if isinstance(model, DistributedDataParallel):
+                    model = next(model.children())
+                try:
+                    if model.save_pretrained is not None:
+                        ckpt_path = os.path.join(
+                            self.save_dir,
+                            f'epoch_{epoch}_val_loss_{val_loss:.4f}'
+                        )
+                        path_exists = os.path.exists(ckpt_path)
+                        if not path_exists:
+                            os.makedirs(ckpt_path)
+                        model.save_pretrained(save_directory=ckpt_path)
+                    else:
+                        path_exists = os.path.exists(self.save_dir)
+                        if not path_exists:
+                            os.makedirs(self.save_dir)
+                        ckpt_path = os.path.join(
+                            self.save_dir,
+                            f'epoch_{epoch}_val_loss_{val_loss:.4f}.pt'
+                        )
+                        torch.save(model.state_dict(), ckpt_path)
+                    self.ckpt_paths.append(ckpt_path)
+
+                    # Keep only N best checkpoints
+                    if len(self.ckpt_paths) > self.max_keep:
+                        oldest_path = self.ckpt_paths.pop(0)
+                        if model.save_pretrained is not None:
+                            shutil.rmtree(oldest_path)
+                        else:
+                            os.remove(oldest_path)
+                except Exception as e:
+                    print(f"Error saving epoch checkpoint: {str(e)}")
+                    if self.display_exc_trace:
+                        traceback.print_exc()
+
+                try:
+                    if self.push_to_hub and self.push_checkpoint_weights and model.push_to_hub is not None and self.hub_model_id:
+                        model.push_to_hub(
+                            repo_id=self.hub_model_id,
+                            commit_message=f'Epoch {epoch} - Val loss {val_loss:.4f}',
+                            token=self.hf_token,
+                            private=self.private_repo,
+                        )
+                except Exception as e:
+                    print(f"Error pushing epoch checkpoint: {str(e)}")
+                    if self.display_exc_trace:
+                        traceback.print_exc()
+
+    def on_training_end(self, model: Union[torch.nn.Module, PyTorchModelHubMixin]):
+        if self.rank == 0:
             if isinstance(model, DistributedDataParallel):
                 model = next(model.children())
             try:
+                # Save final model
                 if model.save_pretrained is not None:
                     ckpt_path = os.path.join(
                         self.save_dir,
-                        f'epoch_{epoch}_val_loss_{val_loss:.4f}'
+                        'final_model'
                     )
-                    path_exists = os.path.exists(ckpt_path)
-                    if not path_exists:
-                        os.makedirs(ckpt_path)
                     model.save_pretrained(save_directory=ckpt_path)
                 else:
-                    path_exists = os.path.exists(self.save_dir)
-                    if not path_exists:
-                        os.makedirs(self.save_dir)
-                    ckpt_path = os.path.join(
-                        self.save_dir,
-                        f'epoch_{epoch}_val_loss_{val_loss:.4f}.pt'
-                    )
+                    ckpt_path = os.path.join(self.save_dir, 'final_model.pt')
                     torch.save(model.state_dict(), ckpt_path)
-                self.ckpt_paths.append(ckpt_path)
-
-                # Keep only N best checkpoints
-                if len(self.ckpt_paths) > self.max_keep:
-                    oldest_path = self.ckpt_paths.pop(0)
-                    if model.save_pretrained is not None:
-                        shutil.rmtree(oldest_path)
-                    else:
-                        os.remove(oldest_path)
+                print(f"Final model saved to {ckpt_path}")
             except Exception as e:
-                print(f"Error saving epoch checkpoint: {str(e)}")
+                print(f"Error saving final model: {str(e)}")
                 if self.display_exc_trace:
                     traceback.print_exc()
-
             try:
-                if self.push_to_hub and self.push_checkpoint_weights and model.push_to_hub is not None and self.hub_model_id:
+                if self.push_to_hub and model.push_to_hub is not None:
                     model.push_to_hub(
                         repo_id=self.hub_model_id,
-                        commit_message=f'Epoch {epoch} - Val loss {val_loss:.4f}',
+                        commit_message=self.final_commit_message or f'Final pre-trained model, after {self.finished_epochs} epochs',
                         token=self.hf_token,
                         private=self.private_repo,
                     )
+                print(f"Model uploaded to repo: {self.hub_model_id}")
             except Exception as e:
-                print(f"Error pushing epoch checkpoint: {str(e)}")
+                print(f"Error pushing final model: {str(e)}")
                 if self.display_exc_trace:
                     traceback.print_exc()
 
-    def on_training_end(self, model: Union[torch.nn.Module, PyTorchModelHubMixin]):
-        if isinstance(model, DistributedDataParallel):
-            model = next(model.children())
-        try:
-            # Save final model
-            if model.save_pretrained is not None:
-                ckpt_path = os.path.join(
-                    self.save_dir,
-                    'final_model'
-                )
-                model.save_pretrained(save_directory=ckpt_path)
-            else:
-                ckpt_path = os.path.join(self.save_dir, 'final_model.pt')
-                torch.save(model.state_dict(), ckpt_path)
-            print(f"Final model saved to {ckpt_path}")
-        except Exception as e:
-            print(f"Error saving final model: {str(e)}")
-            if self.display_exc_trace:
-                traceback.print_exc()
-        try:
-            if self.push_to_hub and model.push_to_hub is not None:
-                model.push_to_hub(
-                    repo_id=self.hub_model_id,
-                    commit_message=self.final_commit_message or f'Final pre-trained model, after {self.finished_epochs} epochs',
-                    token=self.hf_token,
-                    private=self.private_repo,
-                )
-            print(f"Model uploaded to repo: {self.hub_model_id}")
-        except Exception as e:
-            print(f"Error pushing final model: {str(e)}")
-            if self.display_exc_trace:
-                traceback.print_exc()
-
 
 class JointModelSaveCallback(TrainerCallback):
     def __init__(
@@ -298,6 +305,7 @@ class JointModelSaveCallback(TrainerCallback):
             push_batch_checkpoint: bool = False,
             mlm_mode: bool = False,
             display_exc_trace: bool = False,
+            use_ddp: bool = False,
     ):
         self.save_dir = save_dir
         self.save_best_only = save_best_only
@@ -317,6 +325,7 @@ class JointModelSaveCallback(TrainerCallback):
         self.finished_epochs = 0
         self.mlm_mode = mlm_mode
         self.display_exc_trace = display_exc_trace
+        self.rank = int(os.environ['RANK']) if use_ddp else 0
 
     def _save_batch(self, model: Union[nn.Module, PyTorchModelHubMixin], component: str, hub_id: str = None):
         try:
@@ -362,7 +371,7 @@ class JointModelSaveCallback(TrainerCallback):
 
     def on_batch_end(self, model: torch.nn.Module, batch_idx: int, loss: int, batch: dict[str, torch.Tensor]) -> Union[
         bool, None]:
-        if self.save_checkpoint_after_n_batches is not None and batch_idx != 0 and batch_idx % self.save_checkpoint_after_n_batches == 0:
+        if self.rank == 0 and self.save_checkpoint_after_n_batches is not None and batch_idx != 0 and batch_idx % self.save_checkpoint_after_n_batches == 0:
             if isinstance(model, DistributedDataParallel):
                 model = next(model.children())
             self._save_batch(model.encoder, 'encoder', hub_id=self.hub_model_encoder)
@@ -430,15 +439,16 @@ class JointModelSaveCallback(TrainerCallback):
             val_loss: float,
             val_metrics: dict
     ):
-        self.finished_epochs += 1
-        if val_loss < self.best_loss:
-            self.best_loss = val_loss
-            if isinstance(model, DistributedDataParallel):
-                model = next(model.children())
-            self._save_validation(model.encoder, 'encoder', epoch, val_loss, hub_id=self.hub_model_encoder)
-            if not self.mlm_mode:
-                self._save_validation(model.decoder, 'decoder', epoch, val_loss, hub_id=self.hub_model_decoder)
-            self._save_validation(model.mlm_head, 'head', epoch, val_loss, hub_id=self.hub_model_head)
+        if self.rank == 0:
+            self.finished_epochs += 1
+            if val_loss < self.best_loss:
+                self.best_loss = val_loss
+                if isinstance(model, DistributedDataParallel):
+                    model = next(model.children())
+                self._save_validation(model.encoder, 'encoder', epoch, val_loss, hub_id=self.hub_model_encoder)
+                if not self.mlm_mode:
+                    self._save_validation(model.decoder, 'decoder', epoch, val_loss, hub_id=self.hub_model_decoder)
+                self._save_validation(model.mlm_head, 'head', epoch, val_loss, hub_id=self.hub_model_head)
 
     def _save_final(self, model: Union[nn.Module, PyTorchModelHubMixin], component: str, hub_id: str = None):
         try:
@@ -482,9 +492,10 @@ class JointModelSaveCallback(TrainerCallback):
                 traceback.print_exc()
 
     def on_training_end(self, model: Union[torch.nn.Module, PyTorchModelHubMixin]):
-        if isinstance(model, DistributedDataParallel):
-            model = next(model.children())
-        self._save_final(model.encoder, 'encoder', hub_id=self.hub_model_encoder)
-        if not self.mlm_mode:
-            self._save_final(model.decoder, 'decoder', hub_id=self.hub_model_decoder)
-        self._save_final(model.mlm_head, 'head', hub_id=self.hub_model_head)
+        if self.rank == 0:
+            if isinstance(model, DistributedDataParallel):
+                model = next(model.children())
+            self._save_final(model.encoder, 'encoder', hub_id=self.hub_model_encoder)
+            if not self.mlm_mode:
+                self._save_final(model.decoder, 'decoder', hub_id=self.hub_model_decoder)
+            self._save_final(model.mlm_head, 'head', hub_id=self.hub_model_head)

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/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/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]

{rxnn-0.1.9.dist-info → rxnn-0.1.11.dist-info}/METADATA

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

{rxnn-0.1.9.dist-info → rxnn-0.1.11.dist-info}/RECORD

@@ -7,23 +7,23 @@ 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/training/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-rxnn/training/base.py,sha256=UOFFA1Ai6g8l2iOwdYuWrEJPkioej8DOU2YsYN4K9QI,11071
-rxnn/training/bml.py,sha256=pyK6aRLpXlPuLge6CQ9PD64Un57yUgbOpu8lUfTdV9k,14575
-rxnn/training/callbacks.py,sha256=IyVJAJ0ggJmfIWBZnpzV9U08URYCeWIStK_wbx7m3pg,21090
+rxnn/training/base.py,sha256=QD8uS14jSyR5Y_8BgCaBQTKpsarerU3lyufsWsCq_6o,11227
+rxnn/training/bml.py,sha256=2kk9q3Buxq4wBHUQhyIAuHoBCninYX2K8hykWAJnxB0,18654
+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/ff.py,sha256=jJnuBDsnnX5uYC_WZH8cXAYrMnz0P-iX7MwcPivjRtI,2533
-rxnn/transformers/layers.py,sha256=jdM7L0uOMO68aZiu9p6jba1Hx3aLGOChF1Zz-j4vJ5U,5364
+rxnn/transformers/layers.py,sha256=xMocHzdSu7hcC_mPE_aG3-LQg2RXgunKSxcgNXYnOeo,5631
 rxnn/transformers/mask.py,sha256=J0cfLVLt3SzS2ra3KcY4khrkhI975Dw4CjpUi3Sn25s,419
-rxnn/transformers/models.py,sha256=sLYMkVOWQ1NcM1evpCTUMucXvklySpeNT0IqpIGKmyc,6716
-rxnn/transformers/moe.py,sha256=JQ5QSX4FS7S-fqB7-s1ZmJbPpOeD_Injn8o4vo7wGQE,4936
+rxnn/transformers/models.py,sha256=PVhiTTSQ7VTDVdOcyRf-xGNvj6oOa_2fUV2mfthcE0Y,7171
+rxnn/transformers/moe.py,sha256=v21HDEhkDr10--If0P-XBjT5C7IlQJo0wGQlpDnVWEA,5020
 rxnn/transformers/positional.py,sha256=2l38RS0Dini3f6Z3LUHr3XwWzg1UK7fO2C6wazWDAYU,4292
 rxnn/transformers/sampler.py,sha256=wSz_1wNloqtuiix5w2Mcsj5NhaO9QlY0j__TVG7wJnM,3938
 rxnn/utils.py,sha256=d5U8i5ukovgDyqiycc2AoxObTz_eF_bgo2MKvdtJ98s,467
-rxnn-0.1.9.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
-rxnn-0.1.9.dist-info/METADATA,sha256=AraTWJtxAkj6Zx2UUB2YwfFWSk-WwZ5tgcYhWkLZEEM,14628
-rxnn-0.1.9.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
-rxnn-0.1.9.dist-info/RECORD,,
+rxnn-0.1.11.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
+rxnn-0.1.11.dist-info/METADATA,sha256=WFoe6AqfJVI6wFZ23i3qGQ3babDlLtjIMU0htjOIikw,14629
+rxnn-0.1.11.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
+rxnn-0.1.11.dist-info/RECORD,,