rxnn 0.2.60__py3-none-any.whl → 0.2.61__py3-none-any.whl

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:
+                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/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/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:

{rxnn-0.2.60.dist-info → rxnn-0.2.61.dist-info}/METADATA

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

{rxnn-0.2.60.dist-info → rxnn-0.2.61.dist-info}/RECORD

@@ -5,11 +5,11 @@ rxnn/experimental/attention.py,sha256=jlNS82INjycNEfmk3HtkIacUvT_ELhaCO2g-kZTvhX
 rxnn/experimental/models.py,sha256=oJWd56LUsLc9S8eCZw-ShvuWjoQxj4C9GitbohlQ0ok,5139
 rxnn/experimental/moe.py,sha256=jHZ1QhpWiVQOswVpFmuH7b2IUOPf0Uuf-I2Ddwsd7Us,6140
 rxnn/memory/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-rxnn/memory/attention.py,sha256=t-SWJhQ71TV8X_8I_yp0Cr5df7fnWxI-EnYiN8gjpok,3268
+rxnn/memory/attention.py,sha256=thp4t1IVwuIFv-w2WnKZwZeQK_GHwOxaWlnTscs4AH0,3826
 rxnn/memory/norm.py,sha256=cVjjhCLqR5K6-321SP_ObG17y-ddlcTJeCTXvW4vpk0,6675
 rxnn/memory/stm.py,sha256=jv57gsH9XW19sLbxpRDqsp1yfsii_4Ef4Ncr_ztk-i4,3937
 rxnn/rxt/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-rxnn/rxt/models.py,sha256=jd1UVBUWJzWdw7Rjcvo9k5BXCJriQ0khuVszqEyfD7M,14665
+rxnn/rxt/models.py,sha256=mLHvb3ablQK9UtupuOHmLlG440Q_NW-OuLWcxGMfGuY,14807
 rxnn/training/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 rxnn/training/base.py,sha256=CqaArEZYOdH64nmKfx28U3GI46TzO4oNkjf_hrF23Cw,11835
 rxnn/training/bml.py,sha256=hw6gLpLkGvqLzxIvBg4MvCc5r8cHpEm2RDyh7nH6CtE,16914
@@ -17,7 +17,7 @@ rxnn/training/callbacks.py,sha256=rS8leuVFPVVfE5Zc8DMkUZhRIPN-vpPbUjowXE5TSBw,36
 rxnn/training/dataset.py,sha256=tbtOSYldHnQB6SWgee_yUj9zTbgoEoLFNa6wvUS6Apg,51292
 rxnn/training/ddp.py,sha256=VsNBjn3cY-uUj8hbsW7oKvb0_ZKnXnJ2KgObm-Mr9i4,836
 rxnn/training/models.py,sha256=KIiOCW0VgKtMA4EMQ---xsVExdI1mBsgWjtRSmJpecA,9033
-rxnn/training/mrl.py,sha256=H2JcamaJv19vKqOgdoyhcCBwu1lb_aKfCmR_MuuvmS0,62085
+rxnn/training/mrl.py,sha256=BWp87Lj4epjTlROmrQK8RnS_83IucqS7XWI1cBae7BM,64424
 rxnn/training/reward.py,sha256=uiSsBXmjMw2yv-1Bssy3RTlpU6zP8ape3490Sl-aT0M,16144
 rxnn/training/rl.py,sha256=hWtExxY-_pAmTOGYxyCNounUbaGWvLDVltC4sRC7MN4,7175
 rxnn/training/scheduler.py,sha256=LcjU35mEwz2U5x3U6tLfeeYlBqMxbFSxYzJYuXkWbSY,1408
@@ -33,7 +33,7 @@ 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.60.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
-rxnn-0.2.60.dist-info/METADATA,sha256=x_juLxld_xGztBqC7bbBWTn4llwNuZtu29xyjz4uiX8,25997
-rxnn-0.2.60.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-rxnn-0.2.60.dist-info/RECORD,,
+rxnn-0.2.61.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
+rxnn-0.2.61.dist-info/METADATA,sha256=jVl_BOHGqGcDFbA2HnaWS830fczGtWJ_TR-Dp4Ga1ig,25997
+rxnn-0.2.61.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+rxnn-0.2.61.dist-info/RECORD,,