rxnn 0.2.40__tar.gz → 0.2.42__tar.gz

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

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

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

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

{rxnn-0.2.40 → rxnn-0.2.42}/src/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-0.2.40 → rxnn-0.2.42}/src/rxnn/rxt/models.py

@@ -250,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)
@@ -276,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():

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

@@ -80,7 +80,7 @@ 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:
@@ -116,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()
@@ -134,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()
 
@@ -159,12 +168,19 @@ 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:

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

@@ -21,18 +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
-    use_moe_aux_loss: bool
-    moe_aux_loss_scale: float
+    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):
@@ -66,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):
@@ -129,6 +133,8 @@ class MRLTrainer:
         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
@@ -166,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 = {
@@ -174,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
@@ -212,20 +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': lr},
+                {'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,
             )
 
@@ -872,41 +882,41 @@ class MRLTrainer:
             if isinstance(update_epoch, tuple):
                 switch_epoch, cross_att_lr = update_epoch
                 if epoch == switch_epoch:
-                    self.actor.unfreeze_components()
+                    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.unfreeze_components()
+                    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(
@@ -915,11 +925,11 @@ 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.encoder.embedding.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},
@@ -946,7 +956,7 @@ class MRLTrainer:
             ]
         else:
             params = [
-                {'params': self.actor.encoder.embedding.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},
@@ -968,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):
@@ -979,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 = {
@@ -988,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:
@@ -1009,7 +1022,7 @@ class MRLTrainer:
 
         return (epochs, unfreeze_epoch), (random_resets, random_resets_from, random_resets_ratio)
 
-    def __call__(self, curriculum_config: list[CurriculumConfig], batch_size: int):
+    def __call__(self, curriculum_config: list[CurriculumConfig], batch_size: int, ddp_find_unused_parameters: bool = False):
         """Start Memory Reinforcement Learning Curriculum."""
 
         # 0. Set global epoch count for all stages
@@ -1020,7 +1033,7 @@ class MRLTrainer:
         if self.use_ddp:
             rank, world_size = get_os_ddp_config()
             dist.init_process_group(backend='nccl', rank=rank, world_size=world_size)
-            self.actor = DistributedDataParallel(self.actor, device_ids=[self.device.index])
+            self.actor = DistributedDataParallel(self.actor, device_ids=[self.device.index], find_unused_parameters=ddp_find_unused_parameters)
             self.critic = DistributedDataParallel(self.critic, device_ids=[self.device.index])
 
         # 2. Init BatchSampler with actor model (we have to run it after DDP init)
@@ -1039,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")