rxnn 0.2.25__py3-none-any.whl → 0.2.27__py3-none-any.whl

rxnn/memory/norm.py

@@ -7,10 +7,11 @@ class AdaptivePositionalMemoryNorm(nn.Module):
         self,
         num_slots: int,
         dim: int,
-        decay: float = 0.99,
+        decay: float = 0.9,
         use_scale: bool = True,
         use_gate: bool = True,
-        init_gate: float = -4.0
+        init_gate: float = -2.0,
+        per_dim_scale: bool = False,
     ):
         super(AdaptivePositionalMemoryNorm, self).__init__()
         self.use_gate = use_gate
@@ -20,7 +21,8 @@ class AdaptivePositionalMemoryNorm(nn.Module):
         self.eps = 1e-6
 
         # Learnable parameters
-        self.scale = nn.Parameter(torch.ones(num_slots, dim)) if use_scale else None
+        scale_shape = (num_slots, 1) if not per_dim_scale else (dim,)
+        self.scale = nn.Parameter(torch.ones(*scale_shape)) if use_scale else None
         self.gate = nn.Parameter(torch.full((num_slots, 1), init_gate)) if use_gate else None
 
         # EMA buffers
@@ -28,7 +30,7 @@ class AdaptivePositionalMemoryNorm(nn.Module):
 
         # Initialize parameters
         if self.scale is not None:
-            nn.init.normal_(self.scale, mean=1.0, std=0.01)
+            nn.init.normal_(self.scale, mean=1.0, std=0.1)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         # Calculate current RMS per slot
@@ -45,7 +47,7 @@ class AdaptivePositionalMemoryNorm(nn.Module):
 
         # Apply learned scale per slot
         if self.scale is not None:
-            x_norm = x_norm * self.scale # [batch_size, num_slots, dim] * [num_slots, dim]
+            x_norm = x_norm * self.scale # [batch_size, num_slots, dim] * [num_slots, 1] or [dim]
 
         # Apply gating mechanism
         if self.use_gate:
@@ -148,24 +150,26 @@ class MemoryNormConfig(TypedDict):
     use_gate: bool
     init_gate: float
     init_scale: float
+    per_dim_scale: bool
 
 def init_memory_norm(
     norm_type: str,
     dim: int,
     num_slots: int = None,
-    decay: float = 0.99,
+    decay: float = 0.9,
     use_scale: bool = True,
     use_gate: bool = True,
-    init_gate: float = -4.0,
+    init_gate: float = -2.0,
     init_scale: float = 1.0,
+    per_dim_scale: bool = False,
 ) -> nn.Module:
-    assert norm_type in ["layer", "rms", "adaptive", "positional"]
-    if norm_type == "layer":
+    assert norm_type in ['layer', 'rms', 'adaptive', 'positional']
+    if norm_type == 'layer':
         return MemoryLayerNorm(dim, use_gate, init_scale, init_gate)
-    elif norm_type == "rms":
+    elif norm_type == 'rms':
         return SimpleRMSMemoryNorm(dim, use_gate, init_scale, init_gate)
-    elif norm_type == "adaptive":
+    elif norm_type == 'adaptive':
         return AdaptiveRMSMemoryNorm(dim, use_gate, decay, init_scale, init_gate)
-    elif norm_type == "positional":
-        return AdaptivePositionalMemoryNorm(num_slots, dim, decay, use_scale, use_gate, init_gate)
+    elif norm_type == 'positional':
+        return AdaptivePositionalMemoryNorm(num_slots, dim, decay, use_scale, use_gate, init_gate, per_dim_scale)
     return MemoryLayerNorm(dim, use_gate, init_scale, init_gate)

rxnn/rxt/models.py

@@ -13,6 +13,7 @@ from ..memory.attention import StmMemoryAttention
 from ..utils import get_model_size
 from ..experimental.attention import init_experimental_attention
 
+
 class RxTAlphaComponentConfig(TypedDict):
     num_layers: int
     vocab_size: int
@@ -76,8 +77,10 @@ class RxTAlphaComponentBase(nn.Module, PyTorchModelHubMixin):
         assert ff_activation in ['relu', 'gelu',
                                  'swish', 'silu', 'linear',
                                  'sigmoid'], 'Feed-forward activation could be "relu", "gelu", "swish", "silu", "linear", "sigmoid".'
-        assert self_att_type in ['mha', 'gqa', 'mqa', 'gma', 'dma', 'sqa'], 'Self-attention type could be "mha", "gqa", "mqa", "gma", "dma", "sqa".'
-        assert cross_att_type in ['mha', 'gqa', 'mqa', 'gma', 'dma', 'sqa'], 'Memory cross-attention type could be "mha", "gqa", "mqa", "gma", "dma", "sqa".'
+        assert self_att_type in ['mha', 'gqa', 'mqa', 'gma', 'dma',
+                                 'sqa'], 'Self-attention type could be "mha", "gqa", "mqa", "gma", "dma", "sqa".'
+        assert cross_att_type in ['mha', 'gqa', 'mqa', 'gma', 'dma',
+                                  'sqa'], 'Memory cross-attention type could be "mha", "gqa", "mqa", "gma", "dma", "sqa".'
 
         embedding = nn.Embedding(vocab_size, embed_dim)
         rope = RotaryPositionalEmbedding(embed_dim // att_heads, seq_len)
@@ -92,20 +95,25 @@ class RxTAlphaComponentBase(nn.Module, PyTorchModelHubMixin):
         else:
             att_init = lambda: init_experimental_attention(embed_dim, att_heads, self_att_type, att_groups, rope=rope,
                                                            use_flash_attention=use_flash_attention, dropout=att_dropout,
-                                                           max_seq_len=seq_len, is_causal=is_causal, num_experts=att_experts,
+                                                           max_seq_len=seq_len, is_causal=is_causal,
+                                                           num_experts=att_experts,
                                                            num_query_experts=att_query_experts,
                                                            num_query_groups=att_query_groups)
 
         if cross_att_type in ['mha', 'gqa', 'mqa']:
             cross_att_init = lambda: init_attention(embed_dim, att_heads, cross_att_type, att_groups, rope=rope,
-                                              use_flash_attention=use_flash_attention, dropout=att_dropout,
-                                              max_seq_len=seq_len, is_causal=is_causal, rope_only_for_query=True)
+                                                    use_flash_attention=use_flash_attention, dropout=att_dropout,
+                                                    max_seq_len=seq_len, is_causal=is_causal, rope_only_for_query=True)
         else:
-            cross_att_init = lambda: init_experimental_attention(embed_dim, att_heads, cross_att_type, cross_att_groups or att_groups, rope=rope,
-                                                           use_flash_attention=use_flash_attention, dropout=att_dropout,
-                                                           max_seq_len=seq_len, is_causal=is_causal, num_experts=att_experts,
-                                                           num_query_experts=att_query_experts,
-                                                           num_query_groups=cross_att_query_groups or att_query_groups, rope_only_for_query=True)
+            cross_att_init = lambda: init_experimental_attention(embed_dim, att_heads, cross_att_type,
+                                                                 cross_att_groups or att_groups, rope=rope,
+                                                                 use_flash_attention=use_flash_attention,
+                                                                 dropout=att_dropout,
+                                                                 max_seq_len=seq_len, is_causal=is_causal,
+                                                                 num_experts=att_experts,
+                                                                 num_query_experts=att_query_experts,
+                                                                 num_query_groups=cross_att_query_groups or att_query_groups,
+                                                                 rope_only_for_query=True)
 
         layers = nn.ModuleList([
             ReactiveTransformerLayer(
@@ -137,6 +145,12 @@ class RxTAlphaComponentBase(nn.Module, PyTorchModelHubMixin):
     def load_shared_memory(self, stm: ShortTermMemory):
         self.model.stm = stm
 
+    def memory_parameters(self) -> list[nn.Parameter]:
+        return self.model.memory_parameters()
+
+    def not_memory_parameters(self) -> list[nn.Parameter]:
+        return self.model.not_memory_parameters()
+
     def freeze_without_memory(self, unfreeze_norms: bool = True):
         for param in self.model.parameters():
             param.requires_grad_(False)
@@ -211,20 +225,9 @@ class RxTAlphaDecoder(RxTAlphaComponentBase, pipeline_tag="text-generation", lic
         return self.model(x, attention_mask=attention_mask)
 
 
-def build_rxt_alpha_for_pretraining(
-        encoder_config: RxTAlphaComponentConfig,
-        decoder_config: RxTAlphaComponentConfig,
-) -> tuple[RxTAlphaEncoder, RxTAlphaDecoder]:
-    encoder = RxTAlphaEncoder(**encoder_config)
-    decoder = RxTAlphaDecoder(**decoder_config)
-
-    encoder.load_shared_memory(decoder.model.stm)
-    encoder.load_shared_embedding(decoder.model.embedding)
-
-    return encoder, decoder
-
 class RxTAlphaMemoryAttention(nn.Module, PyTorchModelHubMixin, license="apache-2.0"):
     """RxT-Alpha (Reactive Transformer) memory attention model"""
+
     def __init__(
             self,
             num_layers: int = 12,
@@ -234,17 +237,21 @@ class RxTAlphaMemoryAttention(nn.Module, PyTorchModelHubMixin, license="apache-2
             stm_size: int = 1024,
             use_flash_attention: bool = False,
             att_dropout: float = 0.0,
-            norm_type: str = 'rms',
             att_groups: int = 1,
             att_type: str = 'sqa',
             att_experts: int = None,
             att_query_experts: int = None,
             att_query_groups: int = None,
+            norm_type: str = 'rms',
+            norm_init_gate: float = -2.0,
+            norm_per_dim_scale: bool = False,
+            norm_decay: float = 0.9,
             **kwargs,
     ):
         super(RxTAlphaMemoryAttention, self).__init__(**kwargs)
 
-        assert att_type in ['mha', 'gqa', 'mqa', 'gma', 'dma', 'sqa'], 'Memory attention type could be "mha", "gqa", "mqa", "gma", "dma", "sqa".'
+        assert att_type in ['mha', 'gqa', 'mqa', 'gma', 'dma',
+                            'sqa'], 'Memory attention type could be "mha", "gqa", "mqa", "gma", "dma", "sqa".'
 
         rope = RotaryPositionalEmbedding(embed_dim // att_heads, seq_len)
         stm = ShortTermMemory(num_layers, embed_dim, stm_size)
@@ -256,11 +263,14 @@ class RxTAlphaMemoryAttention(nn.Module, PyTorchModelHubMixin, license="apache-2
         else:
             att_init = lambda: init_experimental_attention(embed_dim, att_heads, att_type, att_groups, rope=rope,
                                                            use_flash_attention=use_flash_attention, dropout=att_dropout,
-                                                           max_seq_len=seq_len, is_causal=False, num_experts=att_experts,
+                                                           max_seq_len=seq_len, is_causal=False,
+                                                           num_experts=att_experts,
                                                            num_query_experts=att_query_experts,
                                                            num_query_groups=att_query_groups, rope_only_for_keys=True)
 
-        memory_norm_layers = nn.ModuleList([init_memory_norm(norm_type, embed_dim, stm_size) for _ in range(num_layers)])
+        memory_norm_layers = nn.ModuleList([init_memory_norm(norm_type, embed_dim, stm_size, decay=norm_decay,
+                                                             init_gate=norm_init_gate, per_dim_scale=norm_per_dim_scale)
+                                            for _ in range(num_layers)])
         attention_layers = nn.ModuleList([att_init() for _ in range(num_layers)])
         self.model = StmMemoryAttention(stm, attention_layers, memory_norm_layers)
 
@@ -283,4 +293,3 @@ class RxTAlphaMemoryAttention(nn.Module, PyTorchModelHubMixin, license="apache-2
 
     def forward(self, x: torch.Tensor, attention_mask: torch.Tensor = None) -> torch.Tensor:
         return self.model(x, attention_mask=attention_mask)
-

rxnn/training/models.py

@@ -1,7 +1,7 @@
 import torch
 import torch.nn as nn
 from enum import Enum
-from typing import Literal
+from typing import Literal, Iterator
 from huggingface_hub import PyTorchModelHubMixin
 from ..transformers.models import ReactiveTransformerEncoder, ReactiveTransformerDecoder
 
@@ -75,7 +75,7 @@ class MrlActorModel(nn.Module):
         self.decoder = decoder
         self.memory_attention = memory_attention
 
-    def freeze_components(self, stage: Literal['update', 'fetch', 'both'] = 'both'):
+    def freeze_components(self, stage: Literal['update', 'fetch', 'joint'] = 'joint'):
         """Freeze encoder/decoder except memory-related layers."""
         if self.encoder.freeze_without_memory is not None:
             self.encoder.freeze_without_memory(unfreeze_norms=True)
@@ -124,6 +124,29 @@ class MrlActorModel(nn.Module):
     def reset_memory(self):
         self.memory_attention.reset_memory()
 
+    def memory_parameters(self) -> list[nn.Parameter]:
+        return list(set(
+            self.encoder.memory_parameters() +
+            self.decoder.memory_parameters() +
+            self.memory_attention.parameters()
+        ))
+
+    def memory_cross_attention_parameters(self) -> list[nn.Parameter]:
+        return list(set(
+            self.encoder.memory_parameters() +
+            self.decoder.memory_parameters()
+        ))
+
+    def memory_attention_parameters(self) -> Iterator[nn.Parameter]:
+        return self.memory_attention.parameters()
+
+
+    def not_memory_parameters(self) -> list[nn.Parameter]:
+        return list(set(
+            self.encoder.not_memory_parameters() +
+            self.decoder.not_memory_parameters()
+        ))
+
     def unique_parameters(self):
         return list(set(
             list(self.encoder.parameters()) +

rxnn/training/mrl.py

@@ -3,7 +3,7 @@ from torch.utils.data import DataLoader, DistributedSampler
 from torch.utils.tensorboard import SummaryWriter
 import torch.distributed as dist
 from torch.nn.parallel import DistributedDataParallel
-from typing import Optional, TypedDict, Union
+from typing import Optional, TypedDict, Union, TypeAlias, Literal
 from enum import Enum
 import random, os
 from ..transformers.sampler import BatchSampler
@@ -17,6 +17,8 @@ from .models import MrlActorAction, MrlActorModel, MrlCriticModel
 
 class MrlConfig(TypedDict):
     lr: float
+    separate_memory_lr: Optional[bool]
+    memory_lr: Optional[float]
     critic_lr: float
     max_seq_len: int
     critic_max_len: int
@@ -29,6 +31,8 @@ class MrlStrategy(Enum):
     MULTI_STEP_STRATEGY = 2
     LONG_RANGE_STRATEGY = 3
 
+UnfreezeItem = Union[int, tuple[int, float]]
+UnfreezeEpochsStrategy: TypeAlias = Union[int, tuple[UnfreezeItem, UnfreezeItem, UnfreezeItem, int]]
 
 class CurriculumConfig(TypedDict):
     steps: int
@@ -37,12 +41,14 @@ class CurriculumConfig(TypedDict):
     eval_dataset: Optional[MrlCurriculumDataset]
     callbacks: Optional[list[MrlTrainerCallback]]
     strategy: MrlStrategy
-    unfreeze_epoch: Optional[Union[int, tuple[int, int, int, int]]]
+    unfreeze_epoch: Optional[UnfreezeEpochsStrategy]
     random_resets: Optional[bool]
     random_resets_from: Optional[int]
     random_resets_ratio: Optional[float]
     reward_model: Optional[MrlRewardModel]
+    separate_memory_lr: Optional[bool]
     lr: Optional[float]
+    memory_lr: Optional[float]
     critic_lr: Optional[float]
     weight_decay: Optional[float]
     critic_weight_decay: Optional[float]
@@ -84,6 +90,7 @@ class MRLTrainer:
             use_amp: bool = False,
             dtype: torch.dtype = torch.float32,
             callbacks: list[MrlTrainerCallback] = None,
+
     ):
         """
         Trainer for Memory Reinforcement Learning (MRL) in Reactive Transformer.
@@ -123,15 +130,27 @@ class MRLTrainer:
         self.use_amp = use_amp
         self.dtype = dtype
 
-        self.base_optim_config = {
-            'lr': config.get('lr', 3e-4),
-            'critic_lr': config.get('critic_lr', 1e-4),
-            'weight_decay': config.get('weight_decay', 0.01),
-            'critic_weight_decay': config.get('critic_weight_decay', 0.01),
-        }
+        self.separate_memory_lr = config.get('separate_memory_lr', False)
+
+        if self.separate_memory_lr:
+            self.base_optim_config = {
+                'lr': config.get('lr', 3e-4),
+                'memory_lr': config.get('memory_lr', 5e-4),
+                'critic_lr': config.get('critic_lr', 1e-4),
+                'weight_decay': config.get('weight_decay', 0.01),
+                'critic_weight_decay': config.get('critic_weight_decay', 0.01),
+            }
+        else:
+            self.base_optim_config = {
+                'lr': config.get('lr', 3e-4),
+                'critic_lr': config.get('critic_lr', 1e-4),
+                'weight_decay': config.get('weight_decay', 0.01),
+                'critic_weight_decay': config.get('critic_weight_decay', 0.01),
+            }
+
+        self.optim_config = self.base_optim_config
 
-        # Optimizers
-        self.optimizer, self.critic_optimizer = self._init_optimizers(**self.base_optim_config)
+        self.optimizer, self.critic_optimizer = self._init_optimizers(**self.optim_config)
 
         self.scaler = torch.amp.GradScaler() if self.use_amp else None
         self.critic_scaler = torch.amp.GradScaler() if self.use_amp else None
@@ -158,18 +177,34 @@ class MRLTrainer:
         self.global_epoch = 0
         self.global_epochs_count = 0
 
-    def _init_optimizers(self, lr: float, critic_lr: float, weight_decay: float, critic_weight_decay: float):
-        optimizer = torch.optim.AdamW(
-            self.actor.unique_parameters(),
-            lr=lr,
-            weight_decay=weight_decay,
-        )
+    def _init_optimizers(
+            self,
+            lr: float,
+            critic_lr: float,
+            weight_decay: float,
+            critic_weight_decay: float,
+            memory_lr: Optional[float] = None,
+    ) -> tuple[torch.optim.Optimizer, torch.optim.Optimizer]:
+        if memory_lr is not None:
+            optimizer = torch.optim.AdamW([
+                { 'params': self.actor.not_memory_parameters(), 'lr': lr },
+                { 'params': self.actor.memory_parameters(), 'lr': memory_lr },
+            ],
+                weight_decay=weight_decay,
+            )
+        else:
+            optimizer = torch.optim.AdamW(
+                self.actor.unique_parameters(),
+                lr=lr,
+                weight_decay=weight_decay,
+            )
 
         critic_optimizer = torch.optim.AdamW(
             self.critic.parameters(),
             lr=critic_lr,
             weight_decay=critic_weight_decay,
         )
+
         return optimizer, critic_optimizer
 
 
@@ -712,7 +747,7 @@ class MRLTrainer:
 
         return should_stop_stage
 
-    def _setup_curriculum_step(self, config: CurriculumConfig) -> tuple[tuple[int, int], tuple[bool, int, float]]:
+    def _setup_curriculum_step(self, config: CurriculumConfig) -> tuple[tuple[int, UnfreezeEpochsStrategy], tuple[bool, int, float]]:
         # 1. Set common fields based on config
         self.curriculum_steps = config.get('steps', 1)  # number of steps to run in episode
         self.train_dataset = config.get('dataset', None)  # training dataset for current curriculum stage
@@ -722,13 +757,29 @@ class MRLTrainer:
         self.strategy = config.get('strategy',
                                    MrlStrategy.MULTI_STEP_STRATEGY)  # MRL strategy for given curriculum stage
         self.reward = config.get('reward_model', self.shared_reward_model)  # MRL Reward Model for curriculum stage
-        if config['lr'] is not None or config['critic_lr'] is not None or config['weight_decay'] is not None or config['critic_weight_decay'] is not None:
-            self.optimizer, self.critic_optimizer = self._init_optimizers(
-                lr=config.get('lr', self.base_optim_config['lr']),
-                critic_lr=config.get('critic_lr', self.base_optim_config['critic_lr']),
-                weight_decay=config.get('weight_decay', self.base_optim_config['weight_decay']),
-                critic_weight_decay=config.get('critic_weight_decay', self.base_optim_config['critic_weight_decay'])
-            )
+        if config['lr'] is not None or config['critic_lr'] is not None or config['weight_decay'] is not None or config['critic_weight_decay'] is not None or (config['separate_memory_lr'] and config['memory_lr'] is not None):
+            if config.get('separate_memory_lr', False):
+                self.optim_config = {
+                    'lr': config.get('lr', self.base_optim_config['lr']),
+                    'critic_lr': config.get('critic_lr', self.base_optim_config['critic_lr']),
+                    'weight_decay': config.get('weight_decay', self.base_optim_config['weight_decay']),
+                    'critic_weight_decay': config.get('critic_weight_decay', self.base_optim_config['critic_weight_decay']),
+                    'memory_lr': config.get('memory_lr', self.base_optim_config['memory_lr']),
+                }
+            else:
+                self.optim_config = {
+                    'lr': config.get('lr', self.base_optim_config['lr']),
+                    'critic_lr': config.get('critic_lr', self.base_optim_config['critic_lr']),
+                    'weight_decay': config.get('weight_decay', self.base_optim_config['weight_decay']),
+                    'critic_weight_decay': config.get('critic_weight_decay', self.base_optim_config['critic_weight_decay']),
+                }
+            self.optimizer, self.critic_optimizer = self._init_optimizers(**self.optim_config)
+        elif self.optim_config != self.base_optim_config:
+            self.optim_config = self.base_optim_config
+            self.optimizer, self.critic_optimizer = self._init_optimizers(**self.optim_config)
+
+
+
 
         # 2. Get epochs and random resets configs
         epochs = config.get('epochs', 5)  # number of epochs for current stage
@@ -745,6 +796,82 @@ class MRLTrainer:
 
         return (epochs, unfreeze_epoch), (random_resets, random_resets_from, random_resets_ratio)
 
+    def _apply_unfreeze_strategy(self, epoch: int, unfreeze_epoch: UnfreezeEpochsStrategy):
+        is_staged_unfreeze = isinstance(unfreeze_epoch, tuple)
+        if is_staged_unfreeze:
+            update_epoch, fetch_epoch, joint_epoch, all_epoch = unfreeze_epoch
+
+            if isinstance(update_epoch, tuple):
+                switch_epoch, cross_att_lr = update_epoch
+                if epoch == switch_epoch:
+                    self.actor.freeze_components('joint')
+                    self.optimizer = self._init_unfreeze_optimizer('update', cross_att_lr)
+                    print(f"Activating 'update' unfreeze strategy with custom cross_att_lr: {cross_att_lr}")
+            elif epoch == update_epoch:
+                 self.actor.freeze_components('update')
+                 print(f"Activating 'update' unfreeze strategy - mem-att trainable / cross-att frozen / rest model frozen")
+
+            if isinstance(fetch_epoch, tuple):
+                switch_epoch, mem_att_lr = fetch_epoch
+                if epoch == fetch_epoch:
+                    self.actor.freeze_components('joint')
+                    self.optimizer = self._init_unfreeze_optimizer('fetch', mem_att_lr)
+                    print(f"Activating 'fetch' unfreeze strategy with custom mem_att_lr: {mem_att_lr}")
+            elif epoch == fetch_epoch:
+                self.actor.freeze_components('fetch')
+                print(f"Activating 'fetch' unfreeze strategy - mem-att frozen / cross-att trainable / rest model frozen")
+
+            if isinstance(joint_epoch, tuple):
+                switch_epoch, model_lr = joint_epoch
+                if epoch == joint_epoch:
+                    self.actor.unfreeze_components()
+                    self.optimizer = self._init_unfreeze_optimizer('joint', model_lr)
+                    print(f"Activating 'joint' unfreeze strategy with custom model_lr: {model_lr}")
+            elif epoch == joint_epoch:
+                    self.actor.freeze_components('joint')
+                    print(f"Activating 'joint' unfreeze strategy - mem-att/cross-att trainable / rest model frozen")
+            if epoch == all_epoch:
+                self.actor.unfreeze_components()
+                self.optimizer = self._init_unfreeze_optimizer('all', 0.)
+                print(f"Switching to train 'all' strategy - unfreeze all components")
+        elif epoch == unfreeze_epoch:
+            self.actor.unfreeze_components()
+            print(f"Switching to train 'all' strategy - unfreeze all components")
+
+    def _init_unfreeze_optimizer(
+            self,
+            mode: Literal['update', 'fetch', 'joint', 'all'],
+            unfreeze_lr: float,
+    ) -> torch.optim.Optimizer:
+        memory_lr = self.optim_config['memory_lr'] if 'memory_lr' in self.optim_config else self.optim_config['lr']
+        model_lr = self.optim_config['lr']
+
+        if mode == 'update':
+            params = [
+                {'params': self.actor.not_memory_parameters(), 'lr': model_lr},
+                {'params': self.actor.memory_attention_parameters(), 'lr': memory_lr},
+                {'params': self.actor.memory_cross_attention_parameters(), 'lr': unfreeze_lr},
+            ]
+        elif mode == 'fetch':
+            params = [
+                {'params': self.actor.not_memory_parameters(), 'lr': model_lr},
+                {'params': self.actor.memory_cross_attention_parameters(), 'lr': memory_lr},
+                {'params': self.actor.memory_attention_parameters(), 'lr': unfreeze_lr},
+            ]
+        elif mode == 'joint':
+            params = [
+                {'params': self.actor.not_memory_parameters(), 'lr': unfreeze_lr},
+                {'params': self.actor.memory_parameters(), 'lr': memory_lr},
+            ]
+        else:
+            params = [
+                {'params': self.actor.not_memory_parameters(), 'lr': model_lr},
+                {'params': self.actor.memory_parameters(), 'lr': memory_lr},
+            ]
+
+        return torch.optim.AdamW(params, weight_decay=self.optim_config['weight_decay'])
+
+
     def __call__(self, curriculum_config: list[CurriculumConfig], batch_size: int):
         """Start Memory Reinforcement Learning Curriculum."""
 
@@ -770,7 +897,11 @@ class MRLTrainer:
 
             # 4. Freeze all components except memory attention and memory cross-attention layers in decoder/encoder
             if unfreeze_epoch != 0:
-                self.actor.freeze_components('both')
+                self.actor.freeze_components('joint')
+                if isinstance(unfreeze_epoch, tuple):
+                    print(f"Starting training with unfreeze strategies - 'warmup' - mem-att/cross-att trainable / rest model frozen")
+                else:
+                    print(f"Starting training with simple unfreeze - 'joint' - mem-att/cross-att trainable / rest model frozen")
 
             # 5. Setup train DataLoader
             if self.use_ddp:
@@ -810,21 +941,8 @@ class MRLTrainer:
                 else:
                     self.random_resets_ratio = 1.0
 
-                # 11. Unfreeze all components before selected epoch
-                is_staged_unfreeze = isinstance(unfreeze_epoch, tuple)
-                if is_staged_unfreeze:
-                    update_epoch, fetch_epoch, both_epoch, all_epoch = unfreeze_epoch
-                    if epoch == update_epoch:
-                        self.actor.freeze_components('update')
-                    elif epoch == fetch_epoch:
-                        self.actor.freeze_components('fetch')
-                    elif epoch == both_epoch:
-                        self.actor.freeze_components('both')
-                    elif epoch == all_epoch:
-                        self.actor.unfreeze_components()
-                else:
-                    if epoch == unfreeze_epoch:
-                        self.actor.unfreeze_components()
+                # 11. Apply the unfreeze strategy
+                self._apply_unfreeze_strategy(epoch, unfreeze_epoch)
 
                 # 12. Set epoch for distributed sampler
                 if train_sampler is not None:

rxnn/transformers/layers.py

@@ -64,6 +64,13 @@ class ReactiveTransformerLayer(nn.Module):
             for param in self.norm2.parameters():
                 param.requires_grad_(is_trainable)
 
+    def memory_parameters(self) -> list[nn.Parameter]:
+        return list(self.memory_cross_attention.parameters()) + list(self.norm2.parameters())
+
+    def not_memory_parameters(self) -> list[nn.Parameter]:
+        memory_params = self.memory_parameters()
+        return [param for param in self.parameters() if param not in memory_params]
+
     def update_max_len(self, max_seq_len: int):
         if self.attention.rope is not None:
             self.attention.rope.update_max_len(max_seq_len)

rxnn/transformers/models.py

@@ -39,6 +39,16 @@ class ReactiveTransformerBase(nn.Module):
         for i in range(self.num_own_layers):
             self.layers[i].trainable_cross_attention_(is_trainable, with_norms)
 
+    def memory_parameters(self) -> list[nn.Parameter]:
+        own = [param for layer in self.layers for param in layer.memory_parameters()]
+        shared = [param for layer in self.shared_layers for param in layer.memory_parameters()] if self.shared_layers else []
+        return own + shared
+
+    def not_memory_parameters(self) -> list[nn.Parameter]:
+        own = [param for layer in self.layers for param in layer.not_memory_parameters()]
+        shared = [param for layer in self.shared_layers for param in layer.not_memory_parameters()] if self.shared_layers else []
+        return own + shared
+
     def moe_router_loss(self):
         return torch.stack([self.layers[i].moe_router_loss() for i in range(self.num_own_layers) if self.layers[i].use_moe or self.layers[i].use_moe_att] + [
             self.shared_layers[i].moe_router_loss() for i in range(self.num_shared_layers) if self.shared_layers[i].use_moe or self.shared_layers[i].use_moe_att]).mean()

{rxnn-0.2.25.dist-info → rxnn-0.2.27.dist-info}/METADATA

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

{rxnn-0.2.25.dist-info → rxnn-0.2.27.dist-info}/RECORD

@@ -6,17 +6,17 @@ rxnn/experimental/models.py,sha256=foBo0n0ufvBnfIdJomiEg3CuSOiWSt-q5ako7vzYxx4,4
 rxnn/experimental/moe.py,sha256=jHZ1QhpWiVQOswVpFmuH7b2IUOPf0Uuf-I2Ddwsd7Us,6140
 rxnn/memory/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 rxnn/memory/attention.py,sha256=PXVBZQYNsRraZh7QDBgUOdPy3lTI8B0d8CzduojBjG0,1747
-rxnn/memory/norm.py,sha256=mu_6iZJe61ag627csfJN2JK6QmmzofjOEhxV4ZWblXs,6410
+rxnn/memory/norm.py,sha256=E98jOQEuIOFFhlkvS8s4fFN-D4tLO6vaOqnObv1oVmA,6592
 rxnn/memory/stm.py,sha256=IH_3INw7FdI013t56ui3Zq9GPUq-k3HeZGjx6BerS4g,3888
 rxnn/rxt/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-rxnn/rxt/models.py,sha256=3gCYD_OXvQc8GaXQvRCSj1OcYOSHayWlpP5lsg9wMMk,12389
+rxnn/rxt/models.py,sha256=r8wZeeNTC2VAhiiNe4y7LrbnB4wjFu_cupKiGkpdgjI,13002
 rxnn/training/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 rxnn/training/base.py,sha256=_xik1GXE4RJ_nxwqLQ1ccXA5pRtBCi-jL-jeRFBdHBU,11851
 rxnn/training/bml.py,sha256=FJszaQXOLx2ZHBa1CQpyMrG8i4Kj14E-gzDAEK_Ei5k,17272
 rxnn/training/callbacks.py,sha256=-N0MQPpZQaUWCINdTOsjul4bDGbGr2JgQBqOIXBLS6o,35053
 rxnn/training/dataset.py,sha256=7hTilFWPpqUEc6zNcMqBPjxFKxCfvTKKF3E8tVlwccQ,51250
-rxnn/training/models.py,sha256=5fl1hESVj2Hakqz5to8ZJzw5Q4_RKZAUq2bn6nRiPV8,6045
-rxnn/training/mrl.py,sha256=14wx3pVha15B7eRWPRgoxRtV5dPtBI0yadIHOYZjX6k,43275
+rxnn/training/models.py,sha256=bY6yZoXYJEsrcymtb5Ep41vmFVHplCGWlrw1dI0oFRc,6807
+rxnn/training/mrl.py,sha256=MnLaYWxblc5cF261R5PNjIvddVQVNxyjAkEYtchBn9E,49299
 rxnn/training/reward.py,sha256=7MTVdNm5HnWmt6zFDi3TAYmnVSL_-24riOoY2F7z4x8,11290
 rxnn/training/rl.py,sha256=j-KNLoZjhaEKasYNOc8DxHtwvknAgAJFwvXKot6otFA,3272
 rxnn/training/scheduler.py,sha256=LcjU35mEwz2U5x3U6tLfeeYlBqMxbFSxYzJYuXkWbSY,1408
@@ -25,14 +25,14 @@ rxnn/training/utils.py,sha256=Bw8nZLKIt7NQpUVCYkb_79kWKChVFOYgYXwODo4SvNc,5718
 rxnn/transformers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 rxnn/transformers/attention.py,sha256=d0Igo1Nrn76BphbHrzekiKJfT3RCy4iSXSB6FLAOTgc,16247
 rxnn/transformers/ff.py,sha256=jJnuBDsnnX5uYC_WZH8cXAYrMnz0P-iX7MwcPivjRtI,2533
-rxnn/transformers/layers.py,sha256=UQZbrAg1UAttPASeqS7BP1a4JalktThmRMzX99Qghss,7618
+rxnn/transformers/layers.py,sha256=LXSY829fIHSCmFmClhQ6B7I5aKbiOqy9mZmwlJG_r7U,7961
 rxnn/transformers/mask.py,sha256=J0cfLVLt3SzS2ra3KcY4khrkhI975Dw4CjpUi3Sn25s,419
-rxnn/transformers/models.py,sha256=_2qO1SASHtKvTW3dW-Dy9HEmAvoNVC1_addm2tM9Zbs,8325
+rxnn/transformers/models.py,sha256=QwVxYN9DrKllEpOiFoAx4CiThOWafeTa-OAY7L6gN0Y,8929
 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.25.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
-rxnn-0.2.25.dist-info/METADATA,sha256=nuGFk4oqSMhn6vrw2KZs4RtY0_ZLowg29IlkNVHZ6Jo,25960
-rxnn-0.2.25.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-rxnn-0.2.25.dist-info/RECORD,,
+rxnn-0.2.27.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
+rxnn-0.2.27.dist-info/METADATA,sha256=woZT3PVGgtEJP7DIAJv1-Mdfd4XvKoCRHANQgoTXoXk,25960
+rxnn-0.2.27.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+rxnn-0.2.27.dist-info/RECORD,,