rxnn 0.2.50__py3-none-any.whl → 0.2.52__py3-none-any.whl

rxnn/memory/attention.py

@@ -12,6 +12,7 @@ class StmMemoryAttention(nn.Module):
             per_slot_gate: bool = False,
             init_gate: float = 0.0,
             use_dynamic_gate: bool = False,
+            use_tanh_gate: bool = False,
             *args,
             **kwargs
     ):
@@ -24,6 +25,7 @@ class StmMemoryAttention(nn.Module):
         self.use_gated_residual = use_gated_residual
         self.per_slot_gate = per_slot_gate
         self.use_dynamic_gate = use_dynamic_gate
+        self.use_tanh_gate = use_tanh_gate
         if self.use_gated_residual:
             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))
@@ -37,10 +39,13 @@ class StmMemoryAttention(nn.Module):
         if self.use_dynamic_gate:
             mean_dim = -1 if self.per_slot_gate else [1, 2]
             gate_input = gate * (new_layer_stm + layer_stm).mean(dim=mean_dim, keepdim=True)
-            layer_gate = torch.sigmoid(gate_input)
+            layer_gate = torch.tanh(gate_input) if self.use_tanh_gate else torch.sigmoid(gate_input)
         else:
-            layer_gate = torch.sigmoid(gate)
-        return layer_gate * new_layer_stm + (1 - layer_gate) * layer_stm
+            layer_gate = torch.tanh(gate) if self.use_tanh_gate else torch.sigmoid(gate)
+        if self.use_tanh_gate:
+            return (1 + layer_gate) * new_layer_stm + (1 - layer_gate) * layer_stm
+        else:
+            return layer_gate * new_layer_stm + (1 - layer_gate) * layer_stm
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         new_stm = torch.zeros_like(self.stm.memory)

rxnn/training/models.py

@@ -208,17 +208,18 @@ class MrlActorModel(nn.Module):
 
 
 class MrlCriticModel(nn.Module, PyTorchModelHubMixin, license="apache-2.0", pipeline_tag="text-classification"):
-    def __init__(self, encoder: nn.Module, embed_dim: int,
-                 out_activation: Literal['sigmoid', 'tanh', 'linear'] = 'sigmoid', output_scale: float = 1.0, **kwargs):
+    def __init__(self, encoder: nn.Module, embed_dim: int, **kwargs):
         super(MrlCriticModel, self).__init__(**kwargs)
         self.encoder = encoder
         self.value_head = nn.Sequential(
             GatedLinearUnit(embed_dim, embed_dim, nn.SiLU()),
             nn.LayerNorm(embed_dim),
-            nn.Linear(embed_dim, 1),
-            get_activation_layer(out_activation)
+            nn.Linear(embed_dim, 1)
         )
-        self.output_scale = output_scale
+        # Learnable scaling parameters
+        self.scale = nn.Parameter(torch.tensor(1.0))
+        self.shift = nn.Parameter(torch.tensor(0.0))
+
 
     def head_parameters(self) -> Iterator[nn.Parameter]:
         return self.value_head.parameters()
@@ -235,4 +236,4 @@ class MrlCriticModel(nn.Module, PyTorchModelHubMixin, license="apache-2.0", pipe
         else:
             x = x.mean(dim=1)
 
-        return self.value_head(x) * self.output_scale
+        return self.value_head(x) * self.scale + self.shift

rxnn/training/mrl.py

@@ -9,7 +9,7 @@ import random, os
 from ..transformers.sampler import BatchSampler
 from .callbacks import MrlTrainerCallback
 from .dataset import MrlCurriculumDataset
-from .utils import smart_concat, smart_concat_critic_states, TokenizedDict
+from .utils import smart_concat, smart_concat_critic_states, TokenizedDict, get_gradient_norms
 from .rl import RlAlgorithm
 from .reward import MrlRewardMode, MrlRewardModel
 from .models import MrlActorAction, MrlActorModel, MrlCriticModel
@@ -109,6 +109,7 @@ class MRLTrainer:
             use_ddp: bool = False,
             use_amp: bool = False,
             dtype: torch.dtype = torch.float32,
+            debug_mode: bool = False,
     ):
         """
         Trainer for Memory Reinforcement Learning (MRL) algorithm for reactive models and Attention-Based Memory System.
@@ -139,6 +140,7 @@ class MRLTrainer:
         self.shared_freeze_embeddings = config.get('freeze_embeddings', False)
         self.freeze_embeddings = self.shared_freeze_embeddings
         self.use_memory_warmup = config.get('use_memory_warmup', False)
+        self.debug_mode = debug_mode
         # Internal update epochs config
         self.shared_update_epochs = config.get('update_epochs', 10)
         self.update_epochs = self.shared_update_epochs
@@ -566,6 +568,14 @@ class MRLTrainer:
         else:
             return main_loss
 
+    def _log_gradients(self):
+        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}")
+
     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
@@ -596,6 +606,8 @@ class MRLTrainer:
             self.scaler.unscale_(self.optimizer)
             torch.nn.utils.clip_grad_norm_(self.actor.unique_parameters(), max_norm=1.0,
                                            error_if_nonfinite=False)
+            if self.debug_mode:
+                self._log_gradients()
             # 4.5 Run scaled optimization step
             self.scaler.step(self.optimizer)
             self.scaler.update()
@@ -613,6 +625,8 @@ class MRLTrainer:
             # 4.4 Clip gradient norms
             torch.nn.utils.clip_grad_norm_(self.actor.unique_parameters(), max_norm=1.0,
                                            error_if_nonfinite=False)
+            if self.debug_mode:
+                self._log_gradients()
             # 4.5 Run scaled optimization step
             self.optimizer.step()
         # 5. Get float loss value for callbacks/writer

rxnn/training/rl.py

@@ -36,7 +36,7 @@ class PPOConfig(TypedDict):
 
 
 class PPOAlgorithm(RlAlgorithm):
-    def __init__(self, config: Optional[PPOConfig] = None):
+    def __init__(self, config: Optional[PPOConfig] = None, debug_mode: bool = False):
         super(PPOAlgorithm, self).__init__()
 
         if config is None:
@@ -49,12 +49,14 @@ class PPOAlgorithm(RlAlgorithm):
         self.entropy_coef = config.get('entropy_coef', 0.01)
         self.use_distributed_advantage_norm = config.get('use_distributed_advantage_norm', False)
         self.clip_critic_values = config.get('clip_critic_values', True)
-        self.critic_value_clip = config.get('critic_value_clip', 10.0)
+        self.critic_value_clip = config.get('critic_value_clip', 20.0)
+        self.debug_mode = debug_mode
 
     def critic_loss(self, values: torch.Tensor, ref_values: torch.Tensor) -> torch.Tensor:
         # Critic loss with clipped values
         if self.clip_critic_values:
             values = torch.clamp(values, -self.critic_value_clip, self.critic_value_clip)
+            ref_values = torch.clamp(ref_values, -self.critic_value_clip, self.critic_value_clip)
         return self.critic_loss_fn(values, ref_values)
 
     def policy_loss(self, query: TokenizedDict, answer: TokenizedDict, logits: torch.Tensor,
@@ -95,6 +97,14 @@ class PPOAlgorithm(RlAlgorithm):
 
         advantages = advantages.unsqueeze(-1)
 
+        if self.debug_mode:
+            print(
+                f"Logits stats: min={new_logits.min().item():.4f}, max={new_logits.max().item():.4f}, mean={new_logits.mean().item():.4f}")
+            print(
+                f"Ratio stats: min={ratio.min().item():.4f}, max={ratio.max().item():.4f}, mean={ratio.mean().item():.4f}")
+            print(
+                f"Advantage stats: min={advantages.min().item():.4f}, max={advantages.max().item():.4f}, mean={advantages.mean().item():.4f}")
+
         # c) Clipped surrogate loss
         surr1 = ratio * advantages
         surr2 = torch.clamp(ratio, 1.0 - self.clip_eps, 1.0 + self.clip_eps) * advantages

rxnn/training/utils.py

@@ -1,4 +1,5 @@
 import torch
+import torch.nn as nn
 from typing import TypedDict
 
 
@@ -142,3 +143,13 @@ def smart_concat(query: TokenizedDict, answer: TokenizedDict, max_length: int, p
         'input_ids': combined_ids,
         'attention_mask': combined_mask
     }
+
+def get_gradient_norms(model: nn.Module):
+    total_norm = 0
+    for p in model.parameters():
+        if p.grad is not None:
+            param_norm = p.grad.data.norm(2)
+            total_norm += param_norm.item() ** 2
+    total_norm = total_norm ** 0.5
+    mean_norm = total_norm / len(list(model.parameters()))
+    return total_norm, mean_norm

rxnn/transformers/layers.py

@@ -110,7 +110,7 @@ class ReactiveTransformerLayer(nn.Module):
         residual = x
         if not self.use_post_norm:
             x = self.norm2(x)
-        x = self.memory_cross_attention(x, stm, stm)
+        x = self.memory_cross_attention(x, stm, stm, mask=mask)
         x = residual + x
         if self.use_post_norm:
             x = self.norm2(x)

{rxnn-0.2.50.dist-info → rxnn-0.2.52.dist-info}/METADATA

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

{rxnn-0.2.50.dist-info → rxnn-0.2.52.dist-info}/RECORD

@@ -5,7 +5,7 @@ rxnn/experimental/attention.py,sha256=46qwZLJuZMpIBrZ-r9DaQEPPmmZkO464C3Tkm_Mq-c
 rxnn/experimental/models.py,sha256=foBo0n0ufvBnfIdJomiEg3CuSOiWSt-q5ako7vzYxx4,4888
 rxnn/experimental/moe.py,sha256=jHZ1QhpWiVQOswVpFmuH7b2IUOPf0Uuf-I2Ddwsd7Us,6140
 rxnn/memory/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-rxnn/memory/attention.py,sha256=kan6UNPTjLfO7zKNp92hGooldgWPi3li_2-_L5xiErs,2784
+rxnn/memory/attention.py,sha256=wnYjd3UnmzAA79-7QxpMoEk3O1qRy8LmW1JcE_Fotck,3094
 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
@@ -16,24 +16,24 @@ rxnn/training/bml.py,sha256=hw6gLpLkGvqLzxIvBg4MvCc5r8cHpEm2RDyh7nH6CtE,16914
 rxnn/training/callbacks.py,sha256=rS8leuVFPVVfE5Zc8DMkUZhRIPN-vpPbUjowXE5TSBw,36779
 rxnn/training/dataset.py,sha256=7hTilFWPpqUEc6zNcMqBPjxFKxCfvTKKF3E8tVlwccQ,51250
 rxnn/training/ddp.py,sha256=VsNBjn3cY-uUj8hbsW7oKvb0_ZKnXnJ2KgObm-Mr9i4,836
-rxnn/training/models.py,sha256=L2emJM06u7B9f9T1dFsGXzXX-rsV77ND7L1pAM9Z_Ow,9051
-rxnn/training/mrl.py,sha256=IOi_xbQ47RPgv_2ucT9EkPeWLGBRlgPxKHFeQsYc3Pw,61074
+rxnn/training/models.py,sha256=CS6mjD338knXmCbMZ3bCpOlA-DR3kmQUOSj5u5F6jII,9002
+rxnn/training/mrl.py,sha256=185rZsaFVaAt4mYksuflzKPiDuEaHpjsc3vPzxt9ax0,61862
 rxnn/training/reward.py,sha256=uiSsBXmjMw2yv-1Bssy3RTlpU6zP8ape3490Sl-aT0M,16144
-rxnn/training/rl.py,sha256=q4NzIZAmXRHVToT13IHrPTtEikWQUvT0NO0IjApjAO8,6171
+rxnn/training/rl.py,sha256=iBEPC_gfydXtWkVORO3REMWvOtx60-0xB7MFzfghUK8,6825
 rxnn/training/scheduler.py,sha256=LcjU35mEwz2U5x3U6tLfeeYlBqMxbFSxYzJYuXkWbSY,1408
 rxnn/training/tokenizer.py,sha256=umaLByMBx_NMrQElA45HLm9gkuzyKWDTFaKVd-CjXl0,8344
-rxnn/training/utils.py,sha256=Bw8nZLKIt7NQpUVCYkb_79kWKChVFOYgYXwODo4SvNc,5718
+rxnn/training/utils.py,sha256=C0OS2RAGQ3L7D_G3CWupu_BpAFhkovMByBKm355Ibfc,6087
 rxnn/transformers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 rxnn/transformers/attention.py,sha256=d0Igo1Nrn76BphbHrzekiKJfT3RCy4iSXSB6FLAOTgc,16247
 rxnn/transformers/ff.py,sha256=WDjO-H9XWInoWnUnxiseIH6Kx5GlHP0zGJygwhcb1gc,2589
-rxnn/transformers/layers.py,sha256=l0bXmhN7KOkCw0KTVLixWSo9Op4SesGabWJ4R4EQBMY,7988
+rxnn/transformers/layers.py,sha256=wG8C9doafpLUsGtUTg-xdrHt7EQMEdB10vcSD-O1nVg,7999
 rxnn/transformers/mask.py,sha256=J0cfLVLt3SzS2ra3KcY4khrkhI975Dw4CjpUi3Sn25s,419
 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.50.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
-rxnn-0.2.50.dist-info/METADATA,sha256=MmlWkWUki9ErQnJ24yP2R9mDykQewDHDcyCQhzopZAw,25997
-rxnn-0.2.50.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-rxnn-0.2.50.dist-info/RECORD,,
+rxnn-0.2.52.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
+rxnn-0.2.52.dist-info/METADATA,sha256=_0O3SB2pBSkic0UMXV4_nlYJjFw9bqUFroOs3A4FuW0,25997
+rxnn-0.2.52.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+rxnn-0.2.52.dist-info/RECORD,,