rxnn 0.2.70__tar.gz → 0.2.72__tar.gz

{rxnn-0.2.70 → rxnn-0.2.72}/PKG-INFO

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

{rxnn-0.2.70 → rxnn-0.2.72}/pyproject.toml

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

{rxnn-0.2.70 → rxnn-0.2.72}/src/rxnn/memory/attention.py

@@ -65,7 +65,7 @@ class StmMemoryAttention(nn.Module):
             encoded_layer_data = x[i]
             normalized_layer_stm = self.memory_norm_layers[i](layer_stm)
             if torch.isnan(normalized_layer_stm).any():
-                print(f"NaN detected in {i} layer memory norm output")
+                print(f"!!!!!!!!!!!!!!!!!!!!!!         NaN detected in {i} layer memory norm output")
 
             if self.debug_mode and self.training:
                 if self.debug_step != 0 and self.debug_step % self.debug_interval == 0:
@@ -75,11 +75,11 @@ class StmMemoryAttention(nn.Module):
                     self.debug_step += 1
 
             if torch.isnan(encoded_layer_data).any():
-                print(f"NaN detected in {i} layer encoded data input")
+                print(f"!!!!!!!!!!!!!!!!!!!!!!         NaN detected in {i} layer encoded data input")
 
             new_layer_stm = self.attention_layers[i](normalized_layer_stm, encoded_layer_data, encoded_layer_data, mask=attention_mask)
             if torch.isnan(new_layer_stm).any():
-                print(f"NaN detected in {i} layer memory attention output")
+                print(f"!!!!!!!!!!!!!!!!!!!!!!         NaN detected in {i} layer memory attention output")
 
             if self.use_gated_residual:
                 new_stm[i] = self._residual_gate(self.gate[i], layer_stm, new_layer_stm) # gated residual

{rxnn-0.2.70 → rxnn-0.2.72}/src/rxnn/training/mrl.py

@@ -592,7 +592,7 @@ class MRLTrainer:
 
         router_loss = actor.moe_router_loss()
         if torch.isnan(router_loss).any():
-            print("NaN detected in router loss")
+            print("!!!!!!!!!!!!!!!!!!!!!!         NaN detected in router loss")
         if router_loss is not None:
             return main_loss + self.moe_aux_loss_scale * router_loss
         else:
@@ -607,21 +607,38 @@ class MRLTrainer:
         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)[1] for layer in self.actor.decoder.model.layers]
-        print(f"Decoder cross-att mean norm: {(sum(dec_x_att_norms) / len(dec_x_att_norms)):.6f}, all: {dec_x_att_norms}")
 
+        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]
-        print(f"Memory attention layers mean norm: {(sum(mem_att_norms) / len(mem_att_norms)):.6f}, all: {mem_att_norms}")
-
         enc_ff_norms = [get_gradient_norms(layer.ff)[1] for layer in self.actor.encoder.model.layers]
-        print(f"Encoder ff mean norm: {(sum(enc_ff_norms) / len(enc_ff_norms)):.6f}, all: {enc_ff_norms}")
-
         enc_self_att_norms = [get_gradient_norms(layer.attention)[1] for layer in self.actor.encoder.model.layers]
-        print(f"Encoder self-att mean norm: {(sum(enc_self_att_norms) / len(enc_self_att_norms)):.6f}, all: {enc_self_att_norms}")
+        enc_x_att_norms = [get_gradient_norms(layer.memory_cross_attention)[1] for layer in
+                         self.actor.encoder.model.layers]
+
+        calc_mean = lambda x: sum(x) / len(x)
+
+        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 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}")
+
+        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 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'])
+            self.writer.add_scalar('Gradient/encoder self-att', enc_self_att_norms_mean, self.global_step['train'])
+            self.writer.add_scalar('Gradient/encoder x-att', enc_x_att_norms_mean, self.global_step['train'])
 
-        enc_att_norms = [get_gradient_norms(layer.memory_cross_attention)[1] for layer in self.actor.encoder.model.layers]
-        print(f"Encoder cross-att mean norm: {(sum(enc_att_norms) / len(enc_att_norms)):.6f}, all: {enc_att_norms}")
 
     def update_actor(self, state: tuple[TokenizedDict, TokenizedDict, TokenizedDict], action: TokenizedDict,
                      advantages: torch.Tensor, old_log_probs: torch.Tensor, epoch: int) -> float:
@@ -654,7 +671,7 @@ class MRLTrainer:
             # 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=self.debug_mode)
+                                           error_if_nonfinite=False)
             if self.debug_mode and self.epoch_step['train'] % self.debug_interval == 0:
                 self._log_gradients(logits)
             # 4.5 Run scaled optimization step
@@ -675,7 +692,7 @@ class MRLTrainer:
             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=self.debug_mode)
+                                           error_if_nonfinite=False)
             if self.debug_mode and self.epoch_step['train'] % self.debug_interval == 0:
                 self._log_gradients(logits)
             # 4.5 Run scaled optimization step

{rxnn-0.2.70 → rxnn-0.2.72}/src/rxnn/transformers/layers.py

@@ -103,10 +103,10 @@ class ReactiveTransformerLayer(nn.Module):
         if not self.use_post_norm:
             x = self.norm1(x)
             if torch.isnan(x).any():
-                print("NaN detected in pre-norm (self-attention) output")
+                print("!!!!!!!!!!!!!!!!!!!!!!         !!!!!!!!!!!!!!!!!!!!!!         NaN detected in pre-norm (self-attention) output")
         x = self.attention(x, x, x, mask=mask)
         if torch.isnan(x).any():
-            print("NaN detected in self-attention output")
+            print("!!!!!!!!!!!!!!!!!!!!!!         !!!!!!!!!!!!!!!!!!!!!!         NaN detected in self-attention output")
         x = residual + x
         if self.use_post_norm:
             x = self.norm1(x)
@@ -115,17 +115,17 @@ class ReactiveTransformerLayer(nn.Module):
         if not self.use_post_norm:
             x = self.norm2(x)
             if torch.isnan(x).any():
-                print("NaN detected in pre-norm (cross-attention) output")
+                print("!!!!!!!!!!!!!!!!!!!!!!         NaN detected in pre-norm (cross-attention) output")
 
         mem_mask = mask.squeeze(1).unsqueeze(-1).expand(-1, -1, -1, stm.size(1)) \
             if mask is not None else None
 
         if torch.isnan(stm).any():
-            print("NaN detected in STM cross-attention input")
+            print("!!!!!!!!!!!!!!!!!!!!!!         NaN detected in STM cross-attention input")
 
         x = self.memory_cross_attention(x, stm, stm, mask=mem_mask)
         if torch.isnan(x).any():
-            print("NaN detected in cross-attention output")
+            print("!!!!!!!!!!!!!!!!!!!!!!         NaN detected in cross-attention output")
         x = residual + x
         if self.use_post_norm:
             x = self.norm2(x)
@@ -135,10 +135,10 @@ class ReactiveTransformerLayer(nn.Module):
         if not self.use_post_norm:
             x = self.norm3(x)
             if torch.isnan(x).any():
-                print("NaN detected in pre-norm (ff) output")
+                print("!!!!!!!!!!!!!!!!!!!!!!         NaN detected in pre-norm (ff) output")
         x = self.ff(x)
         if torch.isnan(x).any():
-            print("NaN detected in ff output")
+            print("!!!!!!!!!!!!!!!!!!!!!!         NaN detected in ff output")
         x = residual + x
         if self.use_post_norm:
             x = self.norm3(x)

{rxnn-0.2.70 → rxnn-0.2.72}/src/rxnn/transformers/models.py

@@ -94,7 +94,7 @@ 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")
+            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)
@@ -112,7 +112,7 @@ class ReactiveTransformerDecoder(ReactiveTransformerBase):
         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")
+                print(f"!!!!!!!!!!!!!!!!!!!!!!         NaN detected in {i}. decoder layer output")
         return self.head(self.head_norm(x) if self.use_head_norm else x)
 
 
@@ -122,7 +122,7 @@ 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")
+            print("!!!!!!!!!!!!!!!!!!!!!!         NaN detected in encoder embedding output")
         if attention_mask is not None:
             attention_mask = attention_mask.unsqueeze(1).unsqueeze(1).bool()
 
@@ -136,7 +136,7 @@ class ReactiveTransformerEncoder(ReactiveTransformerBase):
         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")
+                print(f"!!!!!!!!!!!!!!!!!!!!!!         NaN detected in {i}. encoder layer output")
             hidden_states.append(x)
         return x, torch.stack(hidden_states)
 

{rxnn-0.2.70 → rxnn-0.2.72}/src/rxnn/utils.py

@@ -1,6 +1,11 @@
 import random, gc
+from typing import Optional, Union, List, Dict, Any
+
 import torch
 import numpy as np
+from huggingface_hub import PyTorchModelHubMixin
+from huggingface_hub.hub_mixin import DataclassInstance
+
 
 def human_format(num: int):
     """Format numbers to human-readable format."""