rxnn 0.2.68__tar.gz → 0.2.70__tar.gz

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

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

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

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

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

@@ -64,6 +64,8 @@ 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 torch.isnan(normalized_layer_stm).any():
+                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:
@@ -72,7 +74,13 @@ class StmMemoryAttention(nn.Module):
                 else:
                     self.debug_step += 1
 
+            if torch.isnan(encoded_layer_data).any():
+                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")
+
             if self.use_gated_residual:
                 new_stm[i] = self._residual_gate(self.gate[i], layer_stm, new_layer_stm) # gated residual
             else:

{rxnn-0.2.68 → rxnn-0.2.70}/src/rxnn/training/dataset.py

@@ -4,7 +4,7 @@ from datasets import Dataset as HfDataset, load_dataset, concatenate_datasets
 from transformers import PreTrainedTokenizer, PreTrainedTokenizerFast
 from .tokenizer import load_tokenizer_from_hf_hub
 
-from typing import Union, TypedDict, Optional, TypeAlias, Any
+from typing import Union, TypedDict, Optional, TypeAlias, Any, Literal
 
 
 class BaseDataset(Dataset):
@@ -854,8 +854,8 @@ class EncoderSftDataset(BaseInteractionDataset):
             'labels': labels
         }
 
-
-MrlDataItem: TypeAlias = dict[str, Union[dict[str, torch.Tensor], list[dict[str, dict[str, torch.Tensor]]]]]
+ItemFields: TypeAlias = Literal['input_ids', 'attention_mask']
+MrlDataItem: TypeAlias = dict[str, Union[dict[ItemFields, torch.Tensor], list[dict[str, dict[ItemFields, torch.Tensor]]]]]
 
 
 class MrlCurriculumDataset(Dataset):
@@ -1031,7 +1031,7 @@ class MrlCurriculumDataset(Dataset):
         """Collate function for MRL curriculum dataset with nested interactions"""
 
         def collate_interaction_batch(interaction_batch: Union[list[dict[str, dict[str, torch.Tensor]]], tuple[Any]]) -> \
-        dict[str, dict[str, torch.Tensor]]:
+        dict[str, dict[ItemFields, torch.Tensor]]:
             """Helper to collate a batch of interactions"""
             return {
                 'query': {
@@ -1047,13 +1047,22 @@ class MrlCurriculumDataset(Dataset):
         batch_interactions = [x['interactions'] for x in batch]
         transposed_interactions = list(zip(*batch_interactions))
 
-        return {
+        def has_nans(tensor: dict[ItemFields, torch.Tensor]) -> bool:
+            return torch.isnan(tensor['input_ids']).any().item() or torch.isnan(tensor['attention_mask']).any().item()
+
+        results: MrlDataItem = {
             **collate_interaction_batch(batch),  # Collate initial query and answer
             'interactions': [
                 collate_interaction_batch(step_batch) for step_batch in transposed_interactions
             ]
         }
 
+        assert not has_nans(results['query']), "NaN in query"
+        assert not has_nans(results['answer']), "NaN in answer"
+        assert not any([(has_nans(item['query']) or has_nans(item['answer'])) for item in results['interactions']]), "NaN in interactions"
+
+        return results
+
 
 class MrlDatasetItem(TypedDict):
     steps: int

{rxnn-0.2.68 → rxnn-0.2.70}/src/rxnn/training/models.py

@@ -1,7 +1,7 @@
 import torch
 import torch.nn as nn
 from enum import Enum
-from typing import Literal, Iterator
+from typing import Literal, Iterator, Optional
 from huggingface_hub import PyTorchModelHubMixin
 from ..transformers.models import ReactiveTransformerEncoder, ReactiveTransformerDecoder
 from ..transformers.ff import GatedLinearUnit, get_activation_layer
@@ -188,7 +188,7 @@ class MrlActorModel(nn.Module):
                 list(self.memory_attention_parameters())
             ))
 
-    def moe_router_loss(self):
+    def moe_router_loss(self) -> Optional[torch.Tensor]:
         if self.encoder.model.use_moe and self.decoder.model.use_moe:
             return (self.encoder.model.moe_router_loss() + self.decoder.model.moe_router_loss()) / 2
         elif self.encoder.model.use_moe:

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

@@ -41,6 +41,7 @@ class MrlConfig(TypedDict):
     use_memory_warmup: Optional[bool]
     debug_mode: Optional[bool]
     debug_interval: Optional[int]
+    clamp_logits: Optional[bool]
 
 
 class MrlStrategy(Enum):
@@ -152,6 +153,7 @@ class MRLTrainer:
         self.use_memory_warmup = config.get('use_memory_warmup', False)
         self.debug_mode = config.get('debug_mode', False)
         self.debug_interval = config.get('debug_interval', 10)
+        self.clamp_logits = config.get('clamp_logits', False)
         # Internal update epochs config
         self.shared_update_epochs = config.get('update_epochs', 10)
         self.update_epochs = self.shared_update_epochs
@@ -589,12 +591,16 @@ 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 _log_gradients(self):
+    def _log_gradients(self, logits: torch.Tensor):
+        print(
+            f"Returned logits stats: min={logits.min().item():.4f}, max={logits.max().item():.4f}")
         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)
@@ -611,8 +617,11 @@ class MRLTrainer:
         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_ff_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_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_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:
@@ -633,6 +642,8 @@ class MRLTrainer:
                                       pad_token_id=self.pad_token_id)
                 logits = self.actor(inputs['input_ids'], attention_mask=inputs['attention_mask'],
                                     action=MrlActorAction.DECODE)
+                if self.clamp_logits:
+                    logits = logits.clamp(min=-20.0, max=20.0)
                 # 4.2 Calculate policy loss with selected algorithm
                 policy_loss = self.rl_algorithm.policy_loss(next_query, action, logits, old_log_probs,
                                                             advantages)
@@ -643,9 +654,9 @@ 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=False)
+                                           error_if_nonfinite=self.debug_mode)
             if self.debug_mode and self.epoch_step['train'] % self.debug_interval == 0:
-                self._log_gradients()
+                self._log_gradients(logits)
             # 4.5 Run scaled optimization step
             self.scaler.step(self.optimizer)
             self.scaler.update()
@@ -655,6 +666,8 @@ class MRLTrainer:
                                   pad_token_id=self.pad_token_id)
             logits = self.actor(inputs['input_ids'], attention_mask=inputs['attention_mask'],
                                 action=MrlActorAction.DECODE)
+            if self.clamp_logits:
+                logits = logits.clamp(min=-20.0, max=20.0)
             # 4.2 Calculate policy loss with selected algorithm
             policy_loss = self.rl_algorithm.policy_loss(next_query, action, logits, old_log_probs, advantages)
             policy_loss = self._moe_aux_loss(policy_loss)
@@ -662,9 +675,9 @@ 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=False)
+                                           error_if_nonfinite=self.debug_mode)
             if self.debug_mode and self.epoch_step['train'] % self.debug_interval == 0:
-                self._log_gradients()
+                self._log_gradients(logits)
             # 4.5 Run scaled optimization step
             self.optimizer.step()
         # 5. Get float loss value for callbacks/writer

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

@@ -102,7 +102,11 @@ class ReactiveTransformerLayer(nn.Module):
         residual = x
         if not self.use_post_norm:
             x = self.norm1(x)
+            if torch.isnan(x).any():
+                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")
         x = residual + x
         if self.use_post_norm:
             x = self.norm1(x)
@@ -110,11 +114,18 @@ class ReactiveTransformerLayer(nn.Module):
         residual = x
         if not self.use_post_norm:
             x = self.norm2(x)
+            if torch.isnan(x).any():
+                print("NaN detected in pre-norm (cross-attention) output")
 
-        if mask is not None:
-            mem_mask = mask.squeeze(1).unsqueeze(-1).expand(-1, -1, -1, stm.size(1))
+        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")
 
         x = self.memory_cross_attention(x, stm, stm, mask=mem_mask)
+        if torch.isnan(x).any():
+            print("NaN detected in cross-attention output")
         x = residual + x
         if self.use_post_norm:
             x = self.norm2(x)
@@ -123,7 +134,11 @@ class ReactiveTransformerLayer(nn.Module):
         residual = x
         if not self.use_post_norm:
             x = self.norm3(x)
+            if torch.isnan(x).any():
+                print("NaN detected in pre-norm (ff) output")
         x = self.ff(x)
+        if torch.isnan(x).any():
+            print("NaN detected in ff output")
         x = residual + x
         if self.use_post_norm:
             x = self.norm3(x)

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

@@ -58,6 +58,15 @@ class ReactiveTransformerBase(nn.Module):
         else:
             return None
 
+    def _handle_layer(self, i: int, x: torch.Tensor, mask: torch.Tensor = None, is_shared: bool = False):
+        stm_layer_idx = i if is_shared else i + self.num_shared_layers
+        layer_stm = self.stm(stm_layer_idx)
+        # expand layer STM to batch size, if it's not in batch mode
+        if layer_stm.size(0) == 1:
+            layer_stm = layer_stm.expand(x.size(0), -1, -1)
+        layer = self.shared_layers[i] if is_shared else self.layers[i]
+        return layer(x, layer_stm, mask=mask)
+
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         # Shared logic for encoders and decoders - apply embeddings and positional encoding
         x = self.embedding(x)
@@ -84,6 +93,8 @@ 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)
@@ -96,18 +107,12 @@ class ReactiveTransformerDecoder(ReactiveTransformerBase):
         # Process shared layers
         if self.shared_layers is not None:
             for i in range(self.num_shared_layers):
-                layer_stm = self.stm(i)
-                # expand layer STM to batch size, if it's not in batch mode
-                if layer_stm.size(0) == 1:
-                    layer_stm = layer_stm.expand(x.size(0), -1, -1)
-                x = self.shared_layers[i](x, layer_stm, mask=mask)
+                x = self._handle_layer(i, x, mask=mask, is_shared=True)
         # Process own layers
         for i in range(self.num_own_layers):
-            layer_stm = self.stm(i)
-            # expand layer STM to batch size, if it's not in batch mode
-            if layer_stm.size(0) == 1:
-                layer_stm = layer_stm.expand(x.size(0), -1, -1)
-            x = self.layers[i](x, layer_stm, mask=mask)
+            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)
 
 
@@ -116,6 +121,8 @@ 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()
 
@@ -123,19 +130,13 @@ class ReactiveTransformerEncoder(ReactiveTransformerBase):
         # Process shared layers
         if self.shared_layers is not None:
             for i in range(self.num_shared_layers):
-                layer_stm = self.stm(i)
-                # expand layer STM to batch size, if it's not in batch mode
-                if layer_stm.size(0) == 1:
-                    layer_stm = layer_stm.expand(x.size(0), -1, -1)
-                x = self.shared_layers[i](x, layer_stm, mask=attention_mask)
+                x = self._handle_layer(i, x, mask=attention_mask, is_shared=True)
                 hidden_states.append(x)
         # Process own layers
         for i in range(self.num_own_layers):
-            layer_stm = self.stm(i)
-            # expand layer STM to batch size, if it's not in batch mode
-            if layer_stm.size(0) == 1:
-                layer_stm = layer_stm.expand(x.size(0), -1, -1)
-            x = self.layers[i](x, layer_stm, mask=attention_mask)
+            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)