rxnn 0.2.36__py3-none-any.whl → 0.2.38__py3-none-any.whl

rxnn/training/base.py

@@ -82,6 +82,7 @@ class BaseTrainer(ABC):
             dataset: torch.utils.data.Dataset = None,
             optimizer: torch.optim.Optimizer = None,
             scheduler: torch.optim.lr_scheduler.LRScheduler = None,
+            ddp_find_unused_parameters: bool = False,
     ) -> None:
         self.is_running = True
         if dataset is None:
@@ -94,7 +95,7 @@ class BaseTrainer(ABC):
         if self.use_ddp:
             rank, world_size = get_os_ddp_config()
             dist.init_process_group(backend='nccl', rank=rank, world_size=world_size)
-            self.model = DistributedDataParallel(self.model, device_ids=[self.device.index])
+            self.model = DistributedDataParallel(self.model, device_ids=[self.device.index], find_unused_parameters=ddp_find_unused_parameters)
             train_sampler = torch.utils.data.DistributedSampler(dataset, shuffle=True)
             dataloader = torch.utils.data.DataLoader(
                 dataset,

rxnn/training/bml.py

@@ -51,6 +51,10 @@ class MLMTrainer(BaseTrainer):
         model = next(self.model.children()) if isinstance(self.model, DistributedDataParallel) else self.model
 
         router_loss = model.encoder.model.moe_router_loss()
+
+        if self.use_ddp:
+            router_loss = distributed_mean(router_loss)
+
         loss = main_loss + self.moe_aux_loss_scale * router_loss
 
         if self.writer is not None:
@@ -152,6 +156,10 @@ class AutoregressiveTrainer(BaseTrainer):
         model = next(self.model.children()) if isinstance(self.model, DistributedDataParallel) else self.model
 
         router_loss = model.model.moe_router_loss()
+
+        if self.use_ddp:
+            router_loss = distributed_mean(router_loss)
+
         loss = main_loss + self.moe_aux_loss_scale * router_loss
 
         if self.writer is not None:

rxnn/training/rl.py

@@ -112,21 +112,20 @@ class PPOAlgorithm(RlAlgorithm):
         advantages = torch.zeros_like(rewards, device=rewards.device)
         last_advantage = 0
         next_value = last_value
-        next_done = last_done
+        next_done = last_done.float()
         dones = dones.float()
 
         for t in reversed(range(trajectory_len)):
-            if t == trajectory_len - 1:
-                # For the last step, use the provided last_value
-                delta = rewards[t] + self.gae_gamma * next_value * (1 - next_done) - values[t]
-            else:
-                # For other steps, use the next value in the trajectory
-                delta = rewards[t] + self.gae_gamma * values[t + 1] * (1 - dones[t + 1]) - values[t]
-
+            # Calculate delta from rewards, stored next_value, masked by stored next_done, and values
+            delta = rewards[t] + self.gae_gamma * next_value * (1 - next_done) - values[t]
+            # Calculate advantages based on delta, gamma/lambda factors and last advantage, masked by current done flags
             advantages[t] = delta + self.gae_gamma * self.gae_lambda * (1 - dones[t]) * last_advantage
+            # Store current step data as last_advantage, next_done and next_value, for the next iteration step
             last_advantage = advantages[t]
             next_done = dones[t]
+            next_value = values[t]
 
+        # Calculate reference returns, based on advantages and values, and return them with advantages for critic update
         returns = advantages + values
         return advantages, returns
 

{rxnn-0.2.36.dist-info → rxnn-0.2.38.dist-info}/METADATA

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

{rxnn-0.2.36.dist-info → rxnn-0.2.38.dist-info}/RECORD

@@ -11,15 +11,15 @@ rxnn/memory/stm.py,sha256=SSfc-RL9FE-RLkmOEkLB-9Rb00ZXbMLbsAEPdpIW89o,3851
 rxnn/rxt/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 rxnn/rxt/models.py,sha256=CzFELVv5-ybAwl1s1ptpmwM7wdJ07M4jaT1-I8PYrR0,13999
 rxnn/training/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-rxnn/training/base.py,sha256=TGz_37RfI1qLI31GNRV5rLowW1kAHnJwqPm7DNfLfe4,11730
-rxnn/training/bml.py,sha256=hw6gLpLkGvqLzxIvBg4MvCc5r8cHpEm2RDyh7nH6CtE,16914
+rxnn/training/base.py,sha256=CqaArEZYOdH64nmKfx28U3GI46TzO4oNkjf_hrF23Cw,11835
+rxnn/training/bml.py,sha256=-Al_qHvdjX4MR3YO-905dJHus405A2Dg-9uWmO561KU,17080
 rxnn/training/callbacks.py,sha256=p72lbzFAmFjpcUvyy4aUB3qd53I8C6Sk5w9nQvsKgTk,35852
 rxnn/training/dataset.py,sha256=7hTilFWPpqUEc6zNcMqBPjxFKxCfvTKKF3E8tVlwccQ,51250
 rxnn/training/ddp.py,sha256=VsNBjn3cY-uUj8hbsW7oKvb0_ZKnXnJ2KgObm-Mr9i4,836
 rxnn/training/models.py,sha256=y-9XHedSheyK1AmLBp3ayulnUvAmDuJ3t0qVg8wHBRg,7463
 rxnn/training/mrl.py,sha256=fIrg1Er0aAK4TnyDRmJC1m7az9wdkhikxv0CBCrGT-c,55868
 rxnn/training/reward.py,sha256=B7nerPk9eNAv2i7umtNF88tVQVwijNNrchIrEITGHKk,11623
-rxnn/training/rl.py,sha256=DKB7oiePVMz-8Tp3jJOzxVlPZMd7HUEOJUurCJj9f68,5974
+rxnn/training/rl.py,sha256=q4NzIZAmXRHVToT13IHrPTtEikWQUvT0NO0IjApjAO8,6171
 rxnn/training/scheduler.py,sha256=LcjU35mEwz2U5x3U6tLfeeYlBqMxbFSxYzJYuXkWbSY,1408
 rxnn/training/tokenizer.py,sha256=umaLByMBx_NMrQElA45HLm9gkuzyKWDTFaKVd-CjXl0,8344
 rxnn/training/utils.py,sha256=Bw8nZLKIt7NQpUVCYkb_79kWKChVFOYgYXwODo4SvNc,5718
@@ -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.36.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
-rxnn-0.2.36.dist-info/METADATA,sha256=TZZFu-DI5dxsjYVR3qXL9UyTY_Z53Br9zlk3PtE2mlU,25960
-rxnn-0.2.36.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-rxnn-0.2.36.dist-info/RECORD,,
+rxnn-0.2.38.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
+rxnn-0.2.38.dist-info/METADATA,sha256=qS0cXhbW6h6lbqGW_OZvO3x5WERF26OeafDaV7QI8dM,25960
+rxnn-0.2.38.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+rxnn-0.2.38.dist-info/RECORD,,