rxnn 0.2.68__py3-none-any.whl → 0.2.69__py3-none-any.whl

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
@@ -594,7 +596,9 @@ class MRLTrainer:
         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)
@@ -633,6 +637,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)
@@ -645,7 +651,7 @@ class MRLTrainer:
             torch.nn.utils.clip_grad_norm_(self.actor.unique_parameters(), max_norm=1.0,
                                            error_if_nonfinite=False)
             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 +661,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)
@@ -664,7 +672,7 @@ class MRLTrainer:
             torch.nn.utils.clip_grad_norm_(self.actor.unique_parameters(), max_norm=1.0,
                                            error_if_nonfinite=False)
             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.dist-info → rxnn-0.2.69.dist-info}/METADATA

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

{rxnn-0.2.68.dist-info → rxnn-0.2.69.dist-info}/RECORD

@@ -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=2J6Wh4xtsVoE6duEevmovDpmSsMkEoH39Ru0bE8lhFo,65481
+rxnn/training/mrl.py,sha256=c_7P_DhroK3pQLubfmlVryWBSwlZ0BssU8zZ6UhjOaI,65919
 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.68.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
-rxnn-0.2.68.dist-info/METADATA,sha256=w7AYGnPAW9xy8DMmcueWdfoV1oQxqOVVDcRxlvA8gWQ,60420
-rxnn-0.2.68.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-rxnn-0.2.68.dist-info/RECORD,,
+rxnn-0.2.69.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
+rxnn-0.2.69.dist-info/METADATA,sha256=YcmghdF8ypeyOCmglJaws18cDtTqSIE8P-gReGIMzsU,60420
+rxnn-0.2.69.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+rxnn-0.2.69.dist-info/RECORD,,