rxnn 0.2.38__py3-none-any.whl → 0.2.40__py3-none-any.whl

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

@@ -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.dist-info → rxnn-0.2.40.dist-info}/RECORD

@@ -9,15 +9,15 @@ rxnn/memory/attention.py,sha256=PXVBZQYNsRraZh7QDBgUOdPy3lTI8B0d8CzduojBjG0,1747
 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=wDu_DQA_cfXNiFqB65yDv1dXOwAAmpPkDtx16whAPRw,14202
 rxnn/training/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 rxnn/training/base.py,sha256=CqaArEZYOdH64nmKfx28U3GI46TzO4oNkjf_hrF23Cw,11835
-rxnn/training/bml.py,sha256=-Al_qHvdjX4MR3YO-905dJHus405A2Dg-9uWmO561KU,17080
+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=OeNLNBVhB-7kBxI27M7lVYcflmOoCJDQpGktRy8JFSQ,8226
+rxnn/training/mrl.py,sha256=BjL6rLncJWn15u-iYBMfjRRbS5Fq4XVOUcWICgzJxSw,57968
 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.38.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
-rxnn-0.2.38.dist-info/METADATA,sha256=qS0cXhbW6h6lbqGW_OZvO3x5WERF26OeafDaV7QI8dM,25960
-rxnn-0.2.38.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-rxnn-0.2.38.dist-info/RECORD,,
+rxnn-0.2.40.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
+rxnn-0.2.40.dist-info/METADATA,sha256=awLkX7f-n49CMEuSAhgrqnz0CMPQrgGXUkKbhaIORvc,25960
+rxnn-0.2.40.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+rxnn-0.2.40.dist-info/RECORD,,