rxnn 0.2.38__tar.gz → 0.2.40__tar.gz

{rxnn-0.2.38 → rxnn-0.2.40}/PKG-INFO

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

{rxnn-0.2.38 → rxnn-0.2.40}/pyproject.toml

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

{rxnn-0.2.38 → rxnn-0.2.40}/src/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]:
@@ -307,13 +311,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-0.2.38 → rxnn-0.2.40}/src/rxnn/training/bml.py

@@ -51,10 +51,6 @@ class MLMTrainer(BaseTrainer):
         model = next(self.model.children()) if isinstance(self.model, DistributedDataParallel) else self.model
 
         router_loss = model.encoder.model.moe_router_loss()
-
-        if self.use_ddp:
-            router_loss = distributed_mean(router_loss)
-
         loss = main_loss + self.moe_aux_loss_scale * router_loss
 
         if self.writer is not None:
@@ -156,10 +152,6 @@ class AutoregressiveTrainer(BaseTrainer):
         model = next(self.model.children()) if isinstance(self.model, DistributedDataParallel) else self.model
 
         router_loss = model.model.moe_router_loss()
-
-        if self.use_ddp:
-            router_loss = distributed_mean(router_loss)
-
         loss = main_loss + self.moe_aux_loss_scale * router_loss
 
         if self.writer is not None:

{rxnn-0.2.38 → rxnn-0.2.40}/src/rxnn/training/models.py

@@ -82,23 +82,31 @@ class MrlActorModel(nn.Module):
 
     def freeze_components(self, stage: Literal['update', 'fetch', 'joint'] = 'joint'):
         """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()
@@ -158,6 +166,16 @@ class MrlActorModel(nn.Module):
             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:
         if action == MrlActorAction.DECODE:

{rxnn-0.2.38 → rxnn-0.2.40}/src/rxnn/training/mrl.py

@@ -31,6 +31,8 @@ class MrlConfig(TypedDict):
     end_token_id: int
     callbacks: Optional[list[MrlTrainerCallback]]
     memory_aware_critic: bool
+    use_moe_aux_loss: bool
+    moe_aux_loss_scale: float
 
 
 class MrlStrategy(Enum):
@@ -125,6 +127,8 @@ 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)
         # Internal update epochs config
         self.shared_update_epochs = config.get('update_epochs', 10)
         self.update_epochs = self.shared_update_epochs
@@ -212,6 +216,7 @@ class MRLTrainer:
     ) -> tuple[torch.optim.Optimizer, torch.optim.Optimizer]:
         if memory_lr is not None:
             optimizer = torch.optim.AdamW([
+                {'params': self.actor.encoder.embedding.parameters(), 'lr': lr},
                 {'params': self.actor.not_memory_parameters(), 'lr': lr},
                 {'params': self.actor.memory_parameters(), 'lr': memory_lr},
             ],
@@ -522,6 +527,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 +561,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 +580,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,7 +872,7 @@ 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()
                     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:
@@ -863,7 +883,7 @@ class MRLTrainer:
             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()
                     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:
@@ -899,25 +919,39 @@ class MRLTrainer:
 
         if mode == 'update':
             params = [
-                {'params': self.actor.not_memory_parameters(), 'lr': model_lr},
+                {'params': self.actor.encoder.embedding.parameters(), 'lr': model_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': model_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'])

{rxnn-0.2.38 → rxnn-0.2.40}/src/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