rxnn 0.2.39__py3-none-any.whl → 0.2.41__py3-none-any.whl

rxnn/memory/attention.py

@@ -8,6 +8,9 @@ class StmMemoryAttention(nn.Module):
             stm: ShortTermMemory,
             attention_layers: nn.ModuleList,
             memory_norm_layers: nn.ModuleList,
+            use_gated_residual: bool = False,
+            per_slot_gate: bool = False,
+            init_gate: float = 0.0,
             *args,
             **kwargs
     ):
@@ -17,6 +20,10 @@ class StmMemoryAttention(nn.Module):
         self.memory_norm_layers = memory_norm_layers
         assert len(self.attention_layers) == len(self.memory_norm_layers) == self.stm.memory.size(0)
         self.num_layers = len(attention_layers)
+        self.use_gated_residual = use_gated_residual
+        self.per_slot_gate = per_slot_gate
+        if self.use_gated_residual:
+            self.gate = nn.Parameter(torch.full((self.num_layers, self.stm.stm_size, 1), init_gate) if self.per_slot_gate else torch.full((self.num_layers,), init_gate))
 
     def update_max_len(self, max_seq_len: int):
         for i in range(self.num_layers):
@@ -35,7 +42,12 @@ class StmMemoryAttention(nn.Module):
             encoded_layer_data = x[i]
             normalized_layer_stm = self.memory_norm_layers[i](layer_stm)
             new_layer_stm = self.attention_layers[i](normalized_layer_stm, encoded_layer_data, encoded_layer_data, mask=mask)
-            new_stm[i] = new_layer_stm + layer_stm # residual
+            if self.use_gated_residual:
+                # gated residual
+                layer_gate = torch.sigmoid(self.gate[i])
+                new_stm[i] = layer_gate * new_layer_stm + (1 - layer_gate) * layer_stm
+            else:
+                new_stm[i] = new_layer_stm + layer_stm # residual
         self.stm.update_all(new_stm)
         return self.stm.memory
 

rxnn/rxt/models.py

@@ -130,10 +130,10 @@ class RxTAlphaComponentBase(nn.Module, PyTorchModelHubMixin):
                 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)
+        self.model = self._init_model(stm, layers, embedding, use_flash_attention, embed_dim, vocab_size, use_moe)
 
     def _init_model(self, stm: ShortTermMemory, layers: nn.ModuleList, embedding: nn.Embedding,
-                    use_flash_attention: bool, embed_dim: int, vocab_size: int) -> ReactiveTransformerBase:
+                    use_flash_attention: bool, embed_dim: int, vocab_size: int, use_moe: bool) -> ReactiveTransformerBase:
         pass
 
     def params_count(self):
@@ -185,13 +185,15 @@ class RxTAlphaEncoder(RxTAlphaComponentBase, pipeline_tag="fill-mask", license="
             embedding: nn.Embedding,
             use_flash_attention: bool,
             embed_dim: int,
-            vocab_size: int
+            vocab_size: int,
+            use_moe: bool,
     ) -> ReactiveTransformerEncoder:
         return ReactiveTransformerEncoder(
             stm=stm,
             embedding=embedding,
             own_layers=layers,
             use_flash_attention=use_flash_attention,
+            use_moe=use_moe,
         )
 
     def forward(self, x: torch.Tensor, attention_mask: torch.Tensor = None) -> tuple[torch.Tensor, torch.Tensor]:
@@ -210,7 +212,8 @@ class RxTAlphaDecoder(RxTAlphaComponentBase, pipeline_tag="text-generation", lic
             embedding: nn.Embedding,
             use_flash_attention: bool,
             embed_dim: int,
-            vocab_size: int
+            vocab_size: int,
+            use_moe: bool,
     ) -> ReactiveTransformerDecoder:
         return ReactiveTransformerDecoder(
             embed_dim,
@@ -219,6 +222,7 @@ class RxTAlphaDecoder(RxTAlphaComponentBase, pipeline_tag="text-generation", lic
             embedding=embedding,
             own_layers=layers,
             use_flash_attention=use_flash_attention,
+            use_moe=use_moe,
         )
 
     def forward(self, x: torch.Tensor, attention_mask: torch.Tensor = None) -> tuple[torch.Tensor, torch.Tensor]:
@@ -246,6 +250,9 @@ class RxTAlphaMemoryAttention(nn.Module, PyTorchModelHubMixin, license="apache-2
             norm_init_gate: float = -2.0,
             norm_per_dim_scale: bool = False,
             norm_decay: float = 0.9,
+            use_gated_residual: bool = False,
+            residual_per_slot_gate: bool = False,
+            residual_init_gate: float = 0.0,
             **kwargs,
     ):
         super(RxTAlphaMemoryAttention, self).__init__(**kwargs)
@@ -272,7 +279,10 @@ class RxTAlphaMemoryAttention(nn.Module, PyTorchModelHubMixin, license="apache-2
                                                              init_gate=norm_init_gate, per_dim_scale=norm_per_dim_scale)
                                             for _ in range(num_layers)])
         attention_layers = nn.ModuleList([att_init() for _ in range(num_layers)])
-        self.model = StmMemoryAttention(stm, attention_layers, memory_norm_layers)
+        self.model = StmMemoryAttention(
+            stm, attention_layers, memory_norm_layers,
+            use_gated_residual=use_gated_residual, per_slot_gate=residual_per_slot_gate, init_gate=residual_init_gate
+        )
 
     def freeze(self):
         for param in self.parameters():
@@ -307,13 +317,15 @@ class RxTAlphaCriticEncoder(RxTAlphaComponentBase, pipeline_tag="text-classifica
             embedding: nn.Embedding,
             use_flash_attention: bool,
             embed_dim: int,
-            vocab_size: int
+            vocab_size: int,
+            use_moe: bool = False,
     ) -> ReactiveTransformerEncoderDetachStm:
         return ReactiveTransformerEncoderDetachStm(
             stm=stm,
             embedding=embedding,
             own_layers=layers,
             use_flash_attention=use_flash_attention,
+            use_moe=use_moe,
         )
 
     def forward(self, x: torch.Tensor, attention_mask: torch.Tensor = None) -> tuple[torch.Tensor, torch.Tensor]:

rxnn/training/models.py

@@ -80,25 +80,33 @@ class MrlActorModel(nn.Module):
         self.decoder = decoder
         self.memory_attention = memory_attention
 
-    def freeze_components(self, stage: Literal['update', 'fetch', 'joint'] = 'joint'):
+    def freeze_components(self, stage: Literal['update', 'fetch', 'joint'] = 'joint', freeze_embeddings: bool = False):
         """Freeze encoder/decoder except memory-related layers."""
+        # Freeze/unfreeze encoder
         if self.encoder.freeze_without_memory is not None:
-            self.encoder.freeze_without_memory(unfreeze_norms=True)
-            if stage == 'update':
+            if stage == 'update' or stage == 'joint':
+                self.encoder.unfreeze_all()
+            else:
+                self.encoder.freeze_without_memory(unfreeze_norms=True)
                 self.encoder.freeze_memory(with_norms=True)
         else:
             for param in self.encoder.parameters():
-                param.requires_grad = False
-            self.encoder.model.trainable_cross_attention_(True if stage != 'update' else False, with_norms=True)
+                param.requires_grad = True if stage != 'fetch' else False
+            self.encoder.model.trainable_cross_attention_(True if stage != 'fetch' else False, with_norms=True)
+        # Freeze/unfreeze decoder
         if self.decoder.freeze_without_memory is not None:
-            self.decoder.freeze_without_memory(unfreeze_norms=True)
-            if stage == 'update':
-                self.decoder.freeze_memory(with_norms=True)
+            if stage == 'fetch':
+                self.decoder.unfreeze_all()
+            else:
+                self.decoder.freeze_without_memory(unfreeze_norms=True)
+                if stage == 'update':
+                    self.decoder.freeze_memory(with_norms=True)
         else:
             for param in self.decoder.parameters():
-                param.requires_grad = False
+                param.requires_grad = True if stage == 'fetch' else False
             self.decoder.model.trainable_cross_attention_(True if stage != 'update' else False, with_norms=True)
-        # Unfreeze memory attention
+
+        # Freeze/unfreeze memory attention
         if self.memory_attention.freeze is not None:
             if stage == 'fetch':
                 self.memory_attention.freeze()
@@ -108,7 +116,11 @@ class MrlActorModel(nn.Module):
             for param in self.memory_attention.parameters():
                 param.requires_grad = True if stage != 'fetch' else False
 
-    def unfreeze_components(self):
+        if freeze_embeddings:
+            for param in self.encoder.model.embedding.parameters():
+                param.requires_grad = False
+
+    def unfreeze_components(self, freeze_embeddings: bool = False):
         """Unfreeze all components after initial training."""
         if self.encoder.unfreeze_all is not None:
             self.encoder.unfreeze_all()
@@ -126,6 +138,11 @@ class MrlActorModel(nn.Module):
             for param in self.memory_attention.parameters():
                 param.requires_grad = True
 
+        if freeze_embeddings:
+            for param in self.encoder.model.embedding.parameters():
+                param.requires_grad = False
+
+
     def reset_memory(self):
         self.memory_attention.reset_memory()
 
@@ -151,12 +168,29 @@ class MrlActorModel(nn.Module):
             self.decoder.not_memory_parameters()
         ))
 
-    def unique_parameters(self):
-        return list(set(
-            list(self.encoder.parameters()) +
-            list(self.decoder.parameters()) +
-            list(self.memory_attention.parameters())
-        ))
+    def unique_parameters(self, with_embedding: bool = True):
+        if with_embedding:
+            return list(set(
+                list(self.encoder.parameters()) +
+                list(self.decoder.parameters()) +
+                list(self.memory_attention.parameters())
+            ))
+        else:
+            return list(set(
+                self.not_memory_parameters() +
+                self.memory_cross_attention_parameters() +
+                list(self.memory_attention_parameters())
+            ))
+
+    def moe_router_loss(self):
+        if self.encoder.model.use_moe and self.decoder.model.use_moe:
+            return (self.encoder.model.moe_router_loss() + self.decoder.model.moe_router_loss()) / 2
+        elif self.encoder.model.use_moe:
+            return self.encoder.model.moe_router_loss()
+        elif self.decoder.model.use_moe:
+            return self.decoder.model.moe_router_loss()
+        else:
+            return None
 
     def forward(self, x: torch.Tensor, attention_mask: torch.Tensor = None,
                 action: MrlActorAction = MrlActorAction.DECODE) -> torch.Tensor:

rxnn/training/mrl.py

@@ -21,16 +21,20 @@ class MrlConfig(TypedDict):
     separate_memory_lr: Optional[bool]
     memory_lr: Optional[float]
     critic_lr: float
-    critic_encoder_lr: float
+    critic_encoder_lr: Optional[float]
     max_seq_len: int
     critic_max_len: int
-    weight_decay: float
-    critic_weight_decay: float
+    weight_decay: Optional[float]
+    critic_weight_decay: Optional[float]
     update_epochs: int
     pad_token_id: int
     end_token_id: int
     callbacks: Optional[list[MrlTrainerCallback]]
-    memory_aware_critic: bool
+    memory_aware_critic: Optional[bool]
+    use_moe_aux_loss: Optional[bool]
+    moe_aux_loss_scale: Optional[float]
+    freeze_embeddings: Optional[bool]
+    embedding_lr: Optional[float]
 
 
 class MrlStrategy(Enum):
@@ -64,6 +68,8 @@ class CurriculumConfig(TypedDict):
     weight_decay: Optional[float]
     critic_weight_decay: Optional[float]
     update_epochs: Optional[int]
+    freeze_embeddings: Optional[bool]
+    embedding_lr: Optional[float]
 
 
 class SamplerConfig(TypedDict):
@@ -125,6 +131,10 @@ class MRLTrainer:
         self.max_seq_len = config.get('max_seq_len', 256)
         self.critic_max_len = config.get('critic_max_len', 512)
         self.memory_aware_critic = config.get('memory_aware_critic', False)
+        self.use_moe_aux_loss = config.get('use_moe_aux_loss', False)
+        self.moe_aux_loss_scale = config.get('moe_aux_loss_scale', 0.01)
+        self.shared_freeze_embeddings = config.get('freeze_embeddings', False)
+        self.freeze_embeddings = self.shared_freeze_embeddings
         # Internal update epochs config
         self.shared_update_epochs = config.get('update_epochs', 10)
         self.update_epochs = self.shared_update_epochs
@@ -162,6 +172,7 @@ class MRLTrainer:
                 'weight_decay': config.get('weight_decay', 0.01),
                 'critic_weight_decay': config.get('critic_weight_decay', 0.01),
                 'critic_encoder_lr': config.get('critic_encoder_lr', config.get('critic_lr', 1e-4)),
+                'embedding_lr': config.get('embedding_lr', config.get('lr', 3e-4)),
             }
         else:
             self.base_optim_config = {
@@ -170,6 +181,7 @@ class MRLTrainer:
                 'weight_decay': config.get('weight_decay', 0.01),
                 'critic_weight_decay': config.get('critic_weight_decay', 0.01),
                 'critic_encoder_lr': config.get('critic_encoder_lr', config.get('critic_lr', 1e-4)),
+                'embedding_lr': config.get('embedding_lr', config.get('lr', 3e-4)),
             }
 
         self.optim_config = self.base_optim_config
@@ -208,19 +220,22 @@ class MRLTrainer:
             weight_decay: float,
             critic_weight_decay: float,
             critic_encoder_lr: float,
+            embedding_lr: float,
             memory_lr: Optional[float] = None,
     ) -> tuple[torch.optim.Optimizer, torch.optim.Optimizer]:
         if memory_lr is not None:
             optimizer = torch.optim.AdamW([
+                {'params': self.actor.encoder.embedding.parameters(), 'lr': embedding_lr},
                 {'params': self.actor.not_memory_parameters(), 'lr': lr},
                 {'params': self.actor.memory_parameters(), 'lr': memory_lr},
             ],
                 weight_decay=weight_decay,
             )
         else:
-            optimizer = torch.optim.AdamW(
-                self.actor.unique_parameters(),
-                lr=lr,
+            optimizer = torch.optim.AdamW([
+                {'params': self.actor.encoder.embedding.parameters(), 'lr': embedding_lr},
+                {'params': self.actor.unique_parameters(with_embedding=False), 'lr': lr},
+            ],
                 weight_decay=weight_decay,
             )
 
@@ -522,6 +537,18 @@ class MRLTrainer:
         # 6. Return loss item
         return critic_loss_item
 
+    def _moe_aux_loss(self, main_loss: torch.Tensor) -> torch.Tensor:
+        if not self.use_moe_aux_loss:
+            return main_loss
+
+        actor = next(self.actor.children()) if isinstance(self.actor, DistributedDataParallel) else self.actor
+
+        router_loss = actor.moe_router_loss()
+        if router_loss is not None:
+            return main_loss + self.moe_aux_loss_scale * router_loss
+        else:
+            return main_loss
+
     def update_actor(self, state: tuple[TokenizedDict, TokenizedDict, TokenizedDict], action: TokenizedDict,
                      advantages: torch.Tensor, old_log_probs: torch.Tensor, epoch: int) -> float:
         # 1. Reset actor gradients
@@ -544,6 +571,8 @@ class MRLTrainer:
                 # 4.2 Calculate policy loss with selected algorithm
                 policy_loss = self.rl_algorithm.policy_loss(next_query, action, logits, old_log_probs,
                                                             advantages)
+                policy_loss = self._moe_aux_loss(policy_loss)
+
             # 4.3 Run backpropagation with scaler
             self.scaler.scale(policy_loss).backward(retain_graph=True)
             # 4.4 Unscale and clip gradient norms
@@ -561,6 +590,7 @@ class MRLTrainer:
                                 action=MrlActorAction.DECODE)
             # 4.2 Calculate policy loss with selected algorithm
             policy_loss = self.rl_algorithm.policy_loss(next_query, action, logits, old_log_probs, advantages)
+            policy_loss = self._moe_aux_loss(policy_loss)
             # 4.3 Run backpropagation
             policy_loss.backward(retain_graph=True)
             # 4.4 Clip gradient norms
@@ -852,41 +882,41 @@ class MRLTrainer:
             if isinstance(update_epoch, tuple):
                 switch_epoch, cross_att_lr = update_epoch
                 if epoch == switch_epoch:
-                    self.actor.freeze_components('joint')
+                    self.actor.unfreeze_components(freeze_embeddings=self.freeze_embeddings)
                     self.optimizer = self._init_unfreeze_optimizer('update', cross_att_lr)
                     print(f"Activating 'update' unfreeze strategy with custom cross_att_lr: {cross_att_lr}")
             elif epoch == update_epoch:
-                self.actor.freeze_components('update')
+                self.actor.freeze_components('update', freeze_embeddings=self.freeze_embeddings)
                 print(
                     f"Activating 'update' unfreeze strategy - mem-att trainable / cross-att frozen / rest model frozen")
 
             if isinstance(fetch_epoch, tuple):
                 switch_epoch, mem_att_lr = fetch_epoch
                 if epoch == switch_epoch:
-                    self.actor.freeze_components('joint')
+                    self.actor.unfreeze_components(freeze_embeddings=self.freeze_embeddings)
                     self.optimizer = self._init_unfreeze_optimizer('fetch', mem_att_lr)
                     print(f"Activating 'fetch' unfreeze strategy with custom mem_att_lr: {mem_att_lr}")
             elif epoch == fetch_epoch:
-                self.actor.freeze_components('fetch')
+                self.actor.freeze_components('fetch', freeze_embeddings=self.freeze_embeddings)
                 print(
                     f"Activating 'fetch' unfreeze strategy - mem-att frozen / cross-att trainable / rest model frozen")
 
             if isinstance(joint_epoch, tuple):
                 switch_epoch, model_lr = joint_epoch
                 if epoch == switch_epoch:
-                    self.actor.unfreeze_components()
+                    self.actor.unfreeze_components(freeze_embeddings=self.freeze_embeddings)
                     self.optimizer = self._init_unfreeze_optimizer('joint', model_lr)
                     print(f"Activating 'joint' unfreeze strategy with custom model_lr: {model_lr}")
             elif epoch == joint_epoch:
-                self.actor.freeze_components('joint')
+                self.actor.freeze_components('joint', freeze_embeddings=self.freeze_embeddings)
                 print(f"Activating 'joint' unfreeze strategy - mem-att/cross-att trainable / rest model frozen")
 
             if epoch == all_epoch:
-                self.actor.unfreeze_components()
+                self.actor.unfreeze_components(freeze_embeddings=self.freeze_embeddings)
                 self.optimizer = self._init_unfreeze_optimizer('all', 0.)
                 print(f"Switching to train 'all' strategy - unfreeze all components")
         elif epoch == unfreeze_epoch:
-            self.actor.unfreeze_components()
+            self.actor.unfreeze_components(freeze_embeddings=self.freeze_embeddings)
             print(f"Switching to train 'all' strategy - unfreeze all components")
 
     def _init_unfreeze_optimizer(
@@ -895,29 +925,43 @@ class MRLTrainer:
             unfreeze_lr: float,
     ) -> torch.optim.Optimizer:
         memory_lr = self.optim_config['memory_lr'] if 'memory_lr' in self.optim_config else self.optim_config['lr']
-        model_lr = self.optim_config['lr']
+        model_lr, embedding_lr = self.optim_config['lr'], self.optim_config['embedding_lr']
 
         if mode == 'update':
             params = [
-                {'params': self.actor.not_memory_parameters(), 'lr': model_lr},
+                {'params': self.actor.encoder.embedding.parameters(), 'lr': embedding_lr},
+                {'params': self.actor.encoder.not_memory_parameters(), 'lr': model_lr},
+                {'params': self.actor.encoder.memory_parameters(), 'lr': memory_lr},
                 {'params': self.actor.memory_attention_parameters(), 'lr': memory_lr},
-                {'params': self.actor.memory_cross_attention_parameters(), 'lr': unfreeze_lr},
+                {'params': self.actor.decoder.memory_parameters(), 'lr': unfreeze_lr},
+                {'params': self.actor.decoder.not_memory_parameters(), 'lr': unfreeze_lr},
             ]
         elif mode == 'fetch':
             params = [
-                {'params': self.actor.not_memory_parameters(), 'lr': model_lr},
-                {'params': self.actor.memory_cross_attention_parameters(), 'lr': memory_lr},
+                {'params': self.actor.encoder.embedding.parameters(), 'lr': unfreeze_lr},
+                {'params': self.actor.encoder.not_memory_parameters(), 'lr': unfreeze_lr},
+                {'params': self.actor.encoder.memory_parameters(), 'lr': unfreeze_lr},
                 {'params': self.actor.memory_attention_parameters(), 'lr': unfreeze_lr},
+                {'params': self.actor.decoder.memory_parameters(), 'lr': memory_lr},
+                {'params': self.actor.decoder.not_memory_parameters(), 'lr': model_lr},
             ]
         elif mode == 'joint':
             params = [
-                {'params': self.actor.not_memory_parameters(), 'lr': unfreeze_lr},
-                {'params': self.actor.memory_parameters(), 'lr': memory_lr},
+                {'params': self.actor.encoder.embedding.parameters(), 'lr': unfreeze_lr},
+                {'params': self.actor.encoder.not_memory_parameters(), 'lr': unfreeze_lr},
+                {'params': self.actor.encoder.memory_parameters(), 'lr': memory_lr},
+                {'params': self.actor.memory_attention_parameters(), 'lr': memory_lr},
+                {'params': self.actor.decoder.memory_parameters(), 'lr': memory_lr},
+                {'params': self.actor.decoder.not_memory_parameters(), 'lr': unfreeze_lr},
             ]
         else:
             params = [
-                {'params': self.actor.not_memory_parameters(), 'lr': model_lr},
-                {'params': self.actor.memory_parameters(), 'lr': memory_lr},
+                {'params': self.actor.encoder.embedding.parameters(), 'lr': embedding_lr},
+                {'params': self.actor.encoder.not_memory_parameters(), 'lr': model_lr},
+                {'params': self.actor.encoder.memory_parameters(), 'lr': memory_lr},
+                {'params': self.actor.memory_attention_parameters(), 'lr': memory_lr},
+                {'params': self.actor.decoder.memory_parameters(), 'lr': memory_lr},
+                {'params': self.actor.decoder.not_memory_parameters(), 'lr': model_lr},
             ]
 
         return torch.optim.AdamW(params, weight_decay=self.optim_config['weight_decay'])
@@ -934,6 +978,7 @@ class MRLTrainer:
                                    MrlStrategy.MULTI_STEP_STRATEGY)  # MRL strategy for given curriculum stage
         self.reward = config.get('reward_model', self.shared_reward_model)  # MRL Reward Model for curriculum stage
         self.update_epochs = config.get('update_epochs', self.shared_update_epochs)  # Internal update epochs
+        self.freeze_embeddings = config.get('freeze_embeddings', self.shared_freeze_embeddings)
         if config['lr'] is not None or config['critic_lr'] is not None or config['weight_decay'] is not None or config[
             'critic_weight_decay'] is not None or (config['separate_memory_lr'] and config['memory_lr'] is not None):
             if config.get('separate_memory_lr', False):
@@ -945,6 +990,7 @@ class MRLTrainer:
                                                       self.base_optim_config['critic_weight_decay']),
                     'critic_encoder_lr': config.get('critic_encoder_lr', self.base_optim_config['critic_encoder_lr']),
                     'memory_lr': config.get('memory_lr', self.base_optim_config['memory_lr']),
+                    'embedding_lr': config.get('embedding_lr', self.base_optim_config['embedding_lr'])
                 }
             else:
                 self.optim_config = {
@@ -954,6 +1000,7 @@ class MRLTrainer:
                     'critic_weight_decay': config.get('critic_weight_decay',
                                                       self.base_optim_config['critic_weight_decay']),
                     'critic_encoder_lr': config.get('critic_encoder_lr', self.base_optim_config['critic_encoder_lr']),
+                    'embedding_lr': config.get('embedding_lr', self.base_optim_config['embedding_lr'])
                 }
             self.optimizer, self.critic_optimizer = self._init_optimizers(**self.optim_config)
         elif self.optim_config != self.base_optim_config:
@@ -1005,7 +1052,7 @@ class MRLTrainer:
                 if callable(unfreeze_epoch):
                     unfreeze_epoch(-1)
                 else:
-                    self.actor.freeze_components('joint')
+                    self.actor.freeze_components('joint', freeze_embeddings=self.freeze_embeddings)
                     if isinstance(unfreeze_epoch, tuple):
                         print(
                             f"Starting training with unfreeze strategies - 'warmup' - mem-att/cross-att trainable / rest model frozen")

rxnn/transformers/models.py

@@ -17,6 +17,7 @@ class ReactiveTransformerBase(nn.Module):
             absolute_embedding: AbsolutePositionalEmbedding = None,
             use_flash_attention: bool = False,
             use_relative_embedding: bool = False,
+            use_moe: bool = False,
             *args,
             **kwargs,
     ):
@@ -32,6 +33,7 @@ class ReactiveTransformerBase(nn.Module):
         self.layers = own_layers
         self.num_shared_layers = len(shared_layers) if shared_layers else 0
         self.num_own_layers = len(own_layers) if own_layers else 0
+        self.use_moe = use_moe
 
     def trainable_cross_attention_(self, is_trainable: bool, with_norms: bool = True):
         for i in range(self.num_shared_layers):
@@ -50,8 +52,11 @@ class ReactiveTransformerBase(nn.Module):
         return own + shared
 
     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 or self.layers[i].use_moe_att] + [
-            self.shared_layers[i].moe_router_loss() for i in range(self.num_shared_layers) if self.shared_layers[i].use_moe or self.shared_layers[i].use_moe_att]).mean()
+        if self.use_moe:
+            return torch.stack([self.layers[i].moe_router_loss() for i in range(self.num_own_layers) if self.layers[i].use_moe or self.layers[i].use_moe_att] + [
+                self.shared_layers[i].moe_router_loss() for i in range(self.num_shared_layers) if self.shared_layers[i].use_moe or self.shared_layers[i].use_moe_att]).mean()
+        else:
+            return None
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         # Shared logic for encoders and decoders - apply embeddings and positional encoding

{rxnn-0.2.39.dist-info → rxnn-0.2.41.dist-info}/METADATA

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

{rxnn-0.2.39.dist-info → rxnn-0.2.41.dist-info}/RECORD

@@ -5,19 +5,19 @@ rxnn/experimental/attention.py,sha256=46qwZLJuZMpIBrZ-r9DaQEPPmmZkO464C3Tkm_Mq-c
 rxnn/experimental/models.py,sha256=foBo0n0ufvBnfIdJomiEg3CuSOiWSt-q5ako7vzYxx4,4888
 rxnn/experimental/moe.py,sha256=jHZ1QhpWiVQOswVpFmuH7b2IUOPf0Uuf-I2Ddwsd7Us,6140
 rxnn/memory/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-rxnn/memory/attention.py,sha256=PXVBZQYNsRraZh7QDBgUOdPy3lTI8B0d8CzduojBjG0,1747
+rxnn/memory/attention.py,sha256=POszZeW0QBKOh4VTDVekmZGKKwUr1Zj0FOAilTv8Vyg,2411
 rxnn/memory/norm.py,sha256=E98jOQEuIOFFhlkvS8s4fFN-D4tLO6vaOqnObv1oVmA,6592
 rxnn/memory/stm.py,sha256=SSfc-RL9FE-RLkmOEkLB-9Rb00ZXbMLbsAEPdpIW89o,3851
 rxnn/rxt/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-rxnn/rxt/models.py,sha256=CzFELVv5-ybAwl1s1ptpmwM7wdJ07M4jaT1-I8PYrR0,13999
+rxnn/rxt/models.py,sha256=lRn7NRIAAeCxr8hoIXanhaD-cGwVwA23hBdIQpBK6kc,14484
 rxnn/training/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 rxnn/training/base.py,sha256=CqaArEZYOdH64nmKfx28U3GI46TzO4oNkjf_hrF23Cw,11835
 rxnn/training/bml.py,sha256=hw6gLpLkGvqLzxIvBg4MvCc5r8cHpEm2RDyh7nH6CtE,16914
 rxnn/training/callbacks.py,sha256=p72lbzFAmFjpcUvyy4aUB3qd53I8C6Sk5w9nQvsKgTk,35852
 rxnn/training/dataset.py,sha256=7hTilFWPpqUEc6zNcMqBPjxFKxCfvTKKF3E8tVlwccQ,51250
 rxnn/training/ddp.py,sha256=VsNBjn3cY-uUj8hbsW7oKvb0_ZKnXnJ2KgObm-Mr9i4,836
-rxnn/training/models.py,sha256=y-9XHedSheyK1AmLBp3ayulnUvAmDuJ3t0qVg8wHBRg,7463
-rxnn/training/mrl.py,sha256=fIrg1Er0aAK4TnyDRmJC1m7az9wdkhikxv0CBCrGT-c,55868
+rxnn/training/models.py,sha256=4hDH-R9l1lNvBMW_CGG_QgmCVrkyG7Lyo40PPzvkovQ,8876
+rxnn/training/mrl.py,sha256=tv7LjW1HBXF9H7rrITQD4EmN1-qgJT44UblREzsjeew,59378
 rxnn/training/reward.py,sha256=B7nerPk9eNAv2i7umtNF88tVQVwijNNrchIrEITGHKk,11623
 rxnn/training/rl.py,sha256=q4NzIZAmXRHVToT13IHrPTtEikWQUvT0NO0IjApjAO8,6171
 rxnn/training/scheduler.py,sha256=LcjU35mEwz2U5x3U6tLfeeYlBqMxbFSxYzJYuXkWbSY,1408
@@ -28,12 +28,12 @@ rxnn/transformers/attention.py,sha256=d0Igo1Nrn76BphbHrzekiKJfT3RCy4iSXSB6FLAOTg
 rxnn/transformers/ff.py,sha256=WDjO-H9XWInoWnUnxiseIH6Kx5GlHP0zGJygwhcb1gc,2589
 rxnn/transformers/layers.py,sha256=l0bXmhN7KOkCw0KTVLixWSo9Op4SesGabWJ4R4EQBMY,7988
 rxnn/transformers/mask.py,sha256=J0cfLVLt3SzS2ra3KcY4khrkhI975Dw4CjpUi3Sn25s,419
-rxnn/transformers/models.py,sha256=hey6tFN9gmLfWCZLjtl_9OcvIjGpWLI1IDeVnr5y8YM,10583
+rxnn/transformers/models.py,sha256=7ypPNFFnacdZjvaLVue1KR2PmMSdVYsbCMQSunXDL70,10720
 rxnn/transformers/moe.py,sha256=j6jEx6Ip0zttlUZKKn82azxo95lkLZs-H2GLSMD88hY,5859
 rxnn/transformers/positional.py,sha256=1PjcJybUzeQlIKJI4tahAGZcYgCRCL0otxs7mpsNuzM,4410
 rxnn/transformers/sampler.py,sha256=t6iiQTdLQ0TakUWnnhKkb5DKF2F_9-thXHBydDF3fxg,17389
 rxnn/utils.py,sha256=ihb6OTyDtPiocB_lOvnq7eOkjjpCkgs8wxvXUBNQ7mM,996
-rxnn-0.2.39.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
-rxnn-0.2.39.dist-info/METADATA,sha256=0Ky_SOITUSAzWBAcLtNl6Wq2n6ESnMNEs6_sBKezQ88,25960
-rxnn-0.2.39.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-rxnn-0.2.39.dist-info/RECORD,,
+rxnn-0.2.41.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
+rxnn-0.2.41.dist-info/METADATA,sha256=5oKrThfhnOQK8KjDYJfcP-LTb03hNyUrSTjbOSpUUdg,25960
+rxnn-0.2.41.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+rxnn-0.2.41.dist-info/RECORD,,