rxnn 0.2.60__tar.gz → 0.2.62__tar.gz

{rxnn-0.2.60 → rxnn-0.2.62}/PKG-INFO

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

{rxnn-0.2.60 → rxnn-0.2.62}/pyproject.toml

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

{rxnn-0.2.60 → rxnn-0.2.62}/src/rxnn/memory/attention.py

@@ -13,6 +13,8 @@ class StmMemoryAttention(nn.Module):
             init_gate: float = 0.0,
             use_dynamic_gate: bool = False,
             use_tanh_gate: bool = False,
+            debug_mode: bool = False,
+            debug_interval: int = 10,
             *args,
             **kwargs
     ):
@@ -30,6 +32,10 @@ class StmMemoryAttention(nn.Module):
             gate_shape = (self.num_layers, self.stm.stm_size, 1) if self.per_slot_gate else (self.num_layers,)
             self.gate = nn.Parameter(torch.full(gate_shape, init_gate))
 
+        self.debug_mode = debug_mode
+        self.debug_interval = debug_interval
+        self.debug_step = 0
+
     def update_max_len(self, max_seq_len: int):
         for i in range(self.num_layers):
             if self.attention_layers[i].rope is not None:
@@ -58,6 +64,14 @@ class StmMemoryAttention(nn.Module):
                 layer_stm = layer_stm.expand(x.size(0), -1, -1)
             encoded_layer_data = x[i]
             normalized_layer_stm = self.memory_norm_layers[i](layer_stm)
+
+            if self.debug_mode and self.training:
+                if self.debug_step != 0 and self.debug_step % self.debug_interval == 0:
+                    self.debug_step = 0
+                    print(f"Normalized STM stats - mean: {normalized_layer_stm.mean().item():.4f}, std: {normalized_layer_stm.std().item():.4f}")
+                else:
+                    self.debug_step += 1
+
             new_layer_stm = self.attention_layers[i](normalized_layer_stm, encoded_layer_data, encoded_layer_data, mask=attention_mask)
             if self.use_gated_residual:
                 new_stm[i] = self._residual_gate(self.gate[i], layer_stm, new_layer_stm) # gated residual

{rxnn-0.2.60 → rxnn-0.2.62}/src/rxnn/rxt/models.py

@@ -254,6 +254,8 @@ class RxTAlphaMemoryAttention(nn.Module, PyTorchModelHubMixin, license="apache-2
             residual_per_slot_gate: bool = False,
             residual_init_gate: float = 0.0,
             use_dynamic_residual_gate: bool = False,
+            debug_mode: bool = False,
+            debug_interval: int = 10,
             **kwargs,
     ):
         super(RxTAlphaMemoryAttention, self).__init__(**kwargs)
@@ -284,6 +286,7 @@ class RxTAlphaMemoryAttention(nn.Module, PyTorchModelHubMixin, license="apache-2
             stm, attention_layers, memory_norm_layers,
             use_gated_residual=use_gated_residual, per_slot_gate=residual_per_slot_gate,
             init_gate=residual_init_gate, use_dynamic_gate=use_dynamic_residual_gate,
+            debug_mode=debug_mode, debug_interval=debug_interval,
         )
 
     def freeze(self):

{rxnn-0.2.60 → rxnn-0.2.62}/src/rxnn/training/mrl.py

@@ -22,6 +22,9 @@ class MrlConfig(TypedDict):
     memory_lr: Optional[float]
     critic_lr: float
     critic_encoder_lr: Optional[float]
+    encoder_lr: Optional[float]
+    encoder_memory_lr: Optional[float]
+    memory_attn_lr: Optional[float]
     max_seq_len: int
     critic_max_len: int
     weight_decay: Optional[float]
@@ -68,6 +71,9 @@ class CurriculumConfig(TypedDict):
     memory_lr: Optional[float]
     critic_lr: Optional[float]
     critic_encoder_lr: Optional[float]
+    encoder_lr: Optional[float]
+    encoder_memory_lr: Optional[float]
+    memory_attn_lr: Optional[float]
     weight_decay: Optional[float]
     critic_weight_decay: Optional[float]
     update_epochs: Optional[int]
@@ -95,7 +101,10 @@ class MrlTrajectoryEpisode(TypedDict):
     reset_stm: bool
     steps: list[MrlTrajectoryStep]
 
-OptimField: TypeAlias = Literal['lr', 'critic_lr', 'weight_decay', 'critic_weight_decay', 'separate_memory_lr', 'memory_lr']
+OptimField: TypeAlias = Literal[
+    'lr', 'critic_lr', 'weight_decay', 'critic_weight_decay', 'separate_memory_lr',
+    'memory_lr', 'encoder_lr', 'encoder_memory_lr', 'memory_attn_lr'
+]
 
 class MRLTrainer:
     def __init__(
@@ -181,6 +190,9 @@ class MRLTrainer:
                 '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)),
+                'encoder_lr': config.get('encoder_lr', config.get('lr', 3e-4)),
+                'encoder_memory_lr': config.get('encoder_memory_lr', config.get('memory_lr', 5e-4)),
+                'memory_attn_lr': config.get('memory_attn_lr', config.get('memory_lr', 5e-4)),
             }
         else:
             self.base_optim_config = {
@@ -190,6 +202,7 @@ class MRLTrainer:
                 '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)),
+                'encoder_lr': config.get('encoder_lr', config.get('lr', 3e-4)),
             }
 
         self.optim_config = self.base_optim_config
@@ -574,9 +587,18 @@ class MRLTrainer:
         encoder_total, encoder_mean = get_gradient_norms(self.actor.encoder)
         decoder_total, decoder_mean = get_gradient_norms(self.actor.decoder)
         mem_att_total, mem_att_mean = get_gradient_norms(self.actor.memory_attention)
-        print(f"Encoder grad norm - total: {encoder_total:.4f}, mean: {encoder_mean:.4f}")
-        print(f"Decoder grad norm - total: {decoder_total:.4f}, mean: {decoder_mean:.4f}")
-        print(f"Memory attention grad norm - total: {mem_att_total:.4f}, mean: {mem_att_mean:.4f}")
+        print(f"Encoder grad norm - total: {encoder_total:.6f}, mean: {encoder_mean:.6f}")
+        print(f"Decoder grad norm - total: {decoder_total:.6f}, mean: {decoder_mean:.6f}")
+        print(f"Memory attention grad norm - total: {mem_att_total:.6f}, mean: {mem_att_mean:.6f}")
+        # decoder's cross att
+        dec_x_att_norms = [get_gradient_norms(layer.memory_cross_attention)[0] for layer in self.actor.decoder.model.layers]
+        print(f"Decoder cross-att mean total norm: {(sum(dec_x_att_norms) / len(dec_x_att_norms)):.6f}, all: {dec_x_att_norms}")
+
+        mem_att_norms = [get_gradient_norms(layer)[0] for layer in self.actor.memory_attention.model.attention_layers]
+        print(f"Memory attention layers mean total norm: {(sum(mem_att_norms) / len(mem_att_norms)):.6f}, all: {mem_att_norms}")
+
+        enc_ff_norms = [get_gradient_norms(layer.ff)[0] for layer in self.actor.encoder.model.layers]
+        print(f"Encoder ff mean total norm: {(sum(enc_ff_norms) / len(enc_ff_norms)):.6f}, all: {enc_ff_norms}")
 
     def update_actor(self, state: tuple[TokenizedDict, TokenizedDict, TokenizedDict], action: TokenizedDict,
                      advantages: torch.Tensor, old_log_probs: torch.Tensor, epoch: int) -> float:
@@ -969,14 +991,16 @@ 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, embedding_lr = self.optim_config['lr'], self.optim_config['embedding_lr']
+        encoder_memory_lr = self.optim_config['encoder_memory_lr'] if 'encoder_memory_lr' in self.optim_config else self.optim_config['encoder_lr']
+        memory_attn_lr = self.optim_config['memory_attn_lr'] if 'memory_attn_lr' in self.optim_config else self.optim_config['lr']
+        model_lr, embedding_lr, encoder_lr = self.optim_config['lr'], self.optim_config['embedding_lr'], self.optim_config['encoder_lr']
 
         if mode == 'update':
             params = [
                 {'params': self.actor.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.encoder.not_memory_parameters(), 'lr': encoder_lr},
+                {'params': self.actor.encoder.memory_parameters(), 'lr': encoder_memory_lr},
+                {'params': self.actor.memory_attention_parameters(), 'lr': memory_attn_lr},
                 {'params': self.actor.decoder.memory_parameters(), 'lr': unfreeze_lr},
                 {'params': self.actor.decoder.not_memory_parameters(), 'lr': unfreeze_lr},
             ]
@@ -993,17 +1017,17 @@ class MRLTrainer:
             params = [
                 {'params': self.actor.embedding_parameters(), 'lr': embedding_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.encoder.memory_parameters(), 'lr': encoder_memory_lr},
+                {'params': self.actor.memory_attention_parameters(), 'lr': memory_attn_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.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.encoder.not_memory_parameters(), 'lr': encoder_lr},
+                {'params': self.actor.encoder.memory_parameters(), 'lr': encoder_memory_lr},
+                {'params': self.actor.memory_attention_parameters(), 'lr': memory_attn_lr},
                 {'params': self.actor.decoder.memory_parameters(), 'lr': memory_lr},
                 {'params': self.actor.decoder.not_memory_parameters(), 'lr': model_lr},
             ]
@@ -1028,11 +1052,14 @@ class MRLTrainer:
         def has_param(field: OptimField) -> bool:
             return field in config and config[field] is not None
 
-        optim_params: list[OptimField] = ['lr', 'critic_lr', 'weight_decay', 'critic_weight_decay']
+        optim_params: list[OptimField] = ['lr', 'critic_lr', 'weight_decay', 'critic_weight_decay', 'encoder_lr']
+        mem_optim_params: list[OptimField] = ['memory_lr', 'encoder_memory_lr', 'memory_attn_lr']
 
         has_any_optim_param = any(
             has_param(field) for field in optim_params
-        ) or (has_param('separate_memory_lr') and config['separate_memory_lr'] and has_param('memory_lr'))
+        ) or (has_param('separate_memory_lr') and config['separate_memory_lr'] and any(
+            has_param(field) for field in mem_optim_params
+        ))
 
         if has_any_optim_param:
             if config.get('separate_memory_lr', False):
@@ -1044,7 +1071,10 @@ 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'])
+                    'embedding_lr': config.get('embedding_lr', self.base_optim_config['embedding_lr']),
+                    'encoder_lr': config.get('encoder_lr', self.base_optim_config['encoder_lr']),
+                    'encoder_memory_lr': config.get('encoder_memory_lr', self.base_optim_config['encoder_memory_lr']),
+                    'memory_attn_lr': config.get('memory_attn_lr', self.base_optim_config['memory_attn_lr']),
                 }
             else:
                 self.optim_config = {
@@ -1054,7 +1084,8 @@ 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'])
+                    'embedding_lr': config.get('embedding_lr', self.base_optim_config['embedding_lr']),
+                    'encoder_lr': config.get('encoder_lr', self.base_optim_config['encoder_lr']),
                 }
             self.optimizer, self.critic_optimizer = self._init_optimizers(**self.optim_config)
         elif self.optim_config != self.base_optim_config: