rxnn 0.2.73__py3-none-any.whl → 0.2.75__py3-none-any.whl

rxnn/rxt/models.py

@@ -15,7 +15,6 @@ from ..utils import get_model_size
 from ..experimental.attention import init_experimental_attention
 
 
-
 class RxTAlphaComponentConfig(TypedDict):
     num_layers: int
     vocab_size: int

rxnn/training/mrl.py

@@ -42,6 +42,7 @@ class MrlConfig(TypedDict):
     debug_mode: Optional[bool]
     debug_interval: Optional[int]
     clamp_logits: Optional[bool]
+    max_grad_norm: Optional[float]
 
 
 class MrlStrategy(Enum):
@@ -154,6 +155,7 @@ class MRLTrainer:
         self.debug_mode = config.get('debug_mode', False)
         self.debug_interval = config.get('debug_interval', 10)
         self.clamp_logits = config.get('clamp_logits', False)
+        self.max_grad_norm = config.get('max_grad_norm', 1.0)
         # Internal update epochs config
         self.shared_update_epochs = config.get('update_epochs', 10)
         self.update_epochs = self.shared_update_epochs
@@ -591,13 +593,28 @@ class MRLTrainer:
         actor = next(self.actor.children()) if isinstance(self.actor, DistributedDataParallel) else self.actor
 
         router_loss = actor.moe_router_loss()
-        if torch.isnan(router_loss).any():
-            print("!!!!!!!!!!!!!!!!!!!!!!         NaN detected in router loss")
         if router_loss is not None:
             return main_loss + self.moe_aux_loss_scale * router_loss
         else:
             return main_loss
 
+    def _clip_actor_grad_norms(self):
+        # Encoder with embedding
+        torch.nn.utils.clip_grad_norm_(
+            self.actor.encoder.parameters(),
+            max_norm=self.max_grad_norm, error_if_nonfinite=False
+        )
+        # Decoder
+        torch.nn.utils.clip_grad_norm_(
+            self.actor.decoder.memory_parameters() + self.actor.decoder.not_memory_parameters(),
+            max_norm=self.max_grad_norm, error_if_nonfinite=False
+        )
+        # Memory attention
+        torch.nn.utils.clip_grad_norm_(
+            self.actor.memory_attention.parameters(),
+            max_norm=self.max_grad_norm, error_if_nonfinite=False
+        )
+
     def _log_gradients(self, logits: torch.Tensor):
         print(
             f"Returned logits stats: min={logits.min().item():.4f}, max={logits.max().item():.4f}")
@@ -608,8 +625,13 @@ class MRLTrainer:
         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}")
 
+        dec_ff_norms = [get_gradient_norms(layer.ff)[1] for layer in self.actor.decoder.model.layers]
+        dec_self_att_norms = [get_gradient_norms(layer.attention)[1] for layer in self.actor.decoder.model.layers]
         dec_x_att_norms = [get_gradient_norms(layer.memory_cross_attention)[1] for layer in self.actor.decoder.model.layers]
+
+
         mem_att_norms = [get_gradient_norms(layer)[1] for layer in self.actor.memory_attention.model.attention_layers]
+
         enc_ff_norms = [get_gradient_norms(layer.ff)[1] for layer in self.actor.encoder.model.layers]
         enc_self_att_norms = [get_gradient_norms(layer.attention)[1] for layer in self.actor.encoder.model.layers]
         enc_x_att_norms = [get_gradient_norms(layer.memory_cross_attention)[1] for layer in
@@ -617,22 +639,28 @@ class MRLTrainer:
 
         calc_mean = lambda x: sum(x) / len(x)
 
+        dec_ff_norms_mean = calc_mean(dec_ff_norms)
+        dec_self_att_norms_mean = calc_mean(dec_self_att_norms)
         dec_x_att_norms_mean = calc_mean(dec_x_att_norms)
         mem_att_norms_mean = calc_mean(mem_att_norms)
         enc_ff_norms_mean = calc_mean(enc_ff_norms)
         enc_self_att_norms_mean = calc_mean(enc_self_att_norms)
         enc_x_att_norms_mean = calc_mean(enc_x_att_norms)
 
+        print(f"Decoder ff mean norm: {dec_ff_norms_mean:.6f}, all: {dec_ff_norms}")
+        print(f"Decoder self-att mean norm: {dec_self_att_norms_mean:.6f}, all: {dec_self_att_norms}")
         print(f"Decoder cross-att mean norm: {dec_x_att_norms_mean:.6f}, all: {dec_x_att_norms}")
-        print(f"Memory attention layers mean norm: {mem_att_norms_mean:.6f}, all: {mem_att_norms}")
         print(f"Encoder ff mean norm: {enc_ff_norms_mean:.6f}, all: {enc_ff_norms}")
         print(f"Encoder self-att mean norm: {enc_self_att_norms_mean:.6f}, all: {enc_self_att_norms}")
         print(f"Encoder cross-att mean norm: {enc_x_att_norms_mean:.6f}, all: {enc_x_att_norms}")
+        print(f"Memory attention layers mean norm: {mem_att_norms_mean:.6f}, all: {mem_att_norms}")
 
         if self.writer is not None:
             self.writer.add_scalar('Gradient/encoder', encoder_mean, self.global_step['train'])
             self.writer.add_scalar('Gradient/decoder', decoder_mean, self.global_step['train'])
             self.writer.add_scalar('Gradient/mem-att', mem_att_mean, self.global_step['train'])
+            self.writer.add_scalar('Gradient/decoder ff', dec_ff_norms_mean, self.global_step['train'])
+            self.writer.add_scalar('Gradient/decoder self-att', dec_self_att_norms_mean, self.global_step['train'])
             self.writer.add_scalar('Gradient/decoder x-att', dec_x_att_norms_mean, self.global_step['train'])
             self.writer.add_scalar('Gradient/mem-att layers', mem_att_norms_mean, self.global_step['train'])
             self.writer.add_scalar('Gradient/encoder ff', enc_ff_norms_mean, self.global_step['train'])
@@ -670,8 +698,7 @@ class MRLTrainer:
             self.scaler.scale(policy_loss).backward(retain_graph=True)
             # 4.4 Unscale and clip gradient norms
             self.scaler.unscale_(self.optimizer)
-            torch.nn.utils.clip_grad_norm_(self.actor.unique_parameters(), max_norm=1.0,
-                                           error_if_nonfinite=False)
+            self._clip_actor_grad_norms()
             if self.debug_mode and self.epoch_step['train'] % self.debug_interval == 0:
                 self._log_gradients(logits)
             # 4.5 Run scaled optimization step
@@ -691,8 +718,7 @@ class MRLTrainer:
             # 4.3 Run backpropagation
             policy_loss.backward(retain_graph=True)
             # 4.4 Clip gradient norms
-            torch.nn.utils.clip_grad_norm_(self.actor.unique_parameters(), max_norm=1.0,
-                                           error_if_nonfinite=False)
+            self._clip_actor_grad_norms()
             if self.debug_mode and self.epoch_step['train'] % self.debug_interval == 0:
                 self._log_gradients(logits)
             # 4.5 Run scaled optimization step

rxnn/transformers/layers.py

@@ -1,7 +1,5 @@
 import torch
 import torch.nn as nn
-from poetry.console.commands import self
-
 from .attention import MultiHeadAttention
 from .ff import FeedForward, GatedFeedForward
 from .moe import MoeFeedForward, GatedMoeFeedForward

rxnn/transformers/models.py

@@ -93,8 +93,6 @@ class ReactiveTransformerDecoder(ReactiveTransformerBase):
 
     def forward(self, x: torch.Tensor, attention_mask: torch.Tensor = None) -> torch.Tensor:
         x = super().forward(x)  # apply embeddings
-        if torch.isnan(x).any():
-            print("!!!!!!!!!!!!!!!!!!!!!!         NaN detected in decoder embedding output")
         seq_len = x.size(1)
         if not self.use_flash_attention and self.use_relative_embedding:
             mask = create_causal_mask(seq_len, device=x.device)
@@ -111,8 +109,6 @@ class ReactiveTransformerDecoder(ReactiveTransformerBase):
         # Process own layers
         for i in range(self.num_own_layers):
             x = self._handle_layer(i, x, mask=mask)
-            if torch.isnan(x).any():
-                print(f"!!!!!!!!!!!!!!!!!!!!!!         NaN detected in {i}. decoder layer output")
         return self.head(self.head_norm(x) if self.use_head_norm else x)
 
 
@@ -121,8 +117,6 @@ class ReactiveTransformerEncoder(ReactiveTransformerBase):
 
     def forward(self, x: torch.Tensor, attention_mask: torch.Tensor = None) -> tuple[torch.Tensor, torch.Tensor]:
         x = super().forward(x)  # apply embeddings
-        if torch.isnan(x).any():
-            print("!!!!!!!!!!!!!!!!!!!!!!         NaN detected in encoder embedding output")
         if attention_mask is not None:
             attention_mask = attention_mask.unsqueeze(1).unsqueeze(1).bool()
 
@@ -135,8 +129,6 @@ class ReactiveTransformerEncoder(ReactiveTransformerBase):
         # Process own layers
         for i in range(self.num_own_layers):
             x = self._handle_layer(i, x, mask=attention_mask)
-            if torch.isnan(x).any():
-                print(f"!!!!!!!!!!!!!!!!!!!!!!         NaN detected in {i}. encoder layer output")
             hidden_states.append(x)
         return x, torch.stack(hidden_states)
 

{rxnn-0.2.73.dist-info → rxnn-0.2.75.dist-info}/METADATA

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

{rxnn-0.2.73.dist-info → rxnn-0.2.75.dist-info}/RECORD

@@ -10,7 +10,7 @@ rxnn/memory/gate.py,sha256=pR_H2y9C7S02QskoFAEC9Tmluut0k4GGlHgvZGiw6m4,2332
 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=Pb48Frl6HV4Wb9CZgYtmzch3k_4Jess3rhs7dY1I96k,22209
+rxnn/rxt/models.py,sha256=zgRgNUVYYuniiB1xt7HdQYgmhep6e5ybxv3PU0lcfoU,22208
 rxnn/training/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 rxnn/training/base.py,sha256=CqaArEZYOdH64nmKfx28U3GI46TzO4oNkjf_hrF23Cw,11835
 rxnn/training/bml.py,sha256=hw6gLpLkGvqLzxIvBg4MvCc5r8cHpEm2RDyh7nH6CtE,16914
@@ -18,7 +18,7 @@ rxnn/training/callbacks.py,sha256=rS8leuVFPVVfE5Zc8DMkUZhRIPN-vpPbUjowXE5TSBw,36
 rxnn/training/dataset.py,sha256=ruU6k33pQmpTqhxpjLFNdDJnCjcrBcGeFOzJqFahJDM,51880
 rxnn/training/ddp.py,sha256=VsNBjn3cY-uUj8hbsW7oKvb0_ZKnXnJ2KgObm-Mr9i4,836
 rxnn/training/models.py,sha256=ILkcqBV1MImnULnq-YDSSEf8cUdEbUgQaH0FRTsa4LA,9069
-rxnn/training/mrl.py,sha256=KUJAdUznquhf5UlcpV-QF5oKHDBEsDecMEVmMLQZw7w,67380
+rxnn/training/mrl.py,sha256=XxO3LZ6Mxae5RBQJ-k2esU4oKS-BpGXC9mxVrg8MYaE,68527
 rxnn/training/reward.py,sha256=uiSsBXmjMw2yv-1Bssy3RTlpU6zP8ape3490Sl-aT0M,16144
 rxnn/training/rl.py,sha256=hWtExxY-_pAmTOGYxyCNounUbaGWvLDVltC4sRC7MN4,7175
 rxnn/training/scheduler.py,sha256=LcjU35mEwz2U5x3U6tLfeeYlBqMxbFSxYzJYuXkWbSY,1408
@@ -27,14 +27,14 @@ rxnn/training/utils.py,sha256=ngDCm654NL3UsPy190Er4XPc9HI-OyEV6tDLMgEEvQc,6219
 rxnn/transformers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 rxnn/transformers/attention.py,sha256=KRnKT6XUqAXElxV9y72mSpdTeiMgCKCCLqqxCFNTHmA,16372
 rxnn/transformers/ff.py,sha256=WDjO-H9XWInoWnUnxiseIH6Kx5GlHP0zGJygwhcb1gc,2589
-rxnn/transformers/layers.py,sha256=fxjlbQG6cwxq-b2ei4DnohSQGH5gwy4GkfP9duTUvjw,8492
+rxnn/transformers/layers.py,sha256=7iwLZ4De4kw3-YA5p2-adCvTgeqeLC-lXcFAlhN_-AA,8450
 rxnn/transformers/mask.py,sha256=J0cfLVLt3SzS2ra3KcY4khrkhI975Dw4CjpUi3Sn25s,419
-rxnn/transformers/models.py,sha256=TP0H9do53Z0vd8kpHMISBzMpHE5X9QIHcy0B-iJHuNQ,11711
+rxnn/transformers/models.py,sha256=IfBsxmtj3val0AOAZV0cJBi1v2Q68HJoyD1LyqRYljs,11187
 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.73.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
-rxnn-0.2.73.dist-info/METADATA,sha256=gtoRMeFgBuOZs4lRKl9JGUxZ2X4C9K78Ee-NHLMqW4E,60420
-rxnn-0.2.73.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-rxnn-0.2.73.dist-info/RECORD,,
+rxnn-0.2.75.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
+rxnn-0.2.75.dist-info/METADATA,sha256=4re6SGM_SP2BQVHtEEKrpl5l9adUYz2wq6KsBW92rag,60420
+rxnn-0.2.75.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+rxnn-0.2.75.dist-info/RECORD,,