rxnn 0.1.0__py3-none-any.whl

src/memory/norm.py

@@ -0,0 +1,173 @@
+import torch
+import torch.nn as nn
+from typing import TypedDict
+
+class AdaptivePositionalMemoryNorm(nn.Module):
+    def __init__(
+        self,
+        num_slots: int,
+        dim: int,
+        decay: float = 0.99,
+        use_scale: bool = True,
+        use_gate: bool = True,
+        init_gate: float = -4.0
+    ):
+        super(AdaptivePositionalMemoryNorm, self).__init__()
+        self.use_gate = use_gate
+        self.num_slots = num_slots
+        self.dim = dim
+        self.decay = decay
+        self.eps = 1e-6
+
+        # Learnable parameters
+        self.scale = nn.Parameter(torch.ones(num_slots, 1, dim)) if use_scale else None
+        self.gate = nn.Parameter(torch.full((num_slots, 1, 1), init_gate)) if use_gate else None
+
+        # EMA buffers
+        self.register_buffer("ema_rms", torch.ones(num_slots, 1))
+
+        # Initialize parameters
+        if self.scale is not None:
+            nn.init.normal_(self.scale, mean=1.0, std=0.01)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # x shape: [batch_size, num_slots, dim]
+        batch_size = x.size(0)
+
+        # Calculate current RMS per slot
+        current_rms = x.pow(2).mean(dim=-1, keepdim=True).sqrt()  # [batch, slots, 1]
+        slot_rms = current_rms.mean(dim=0)  # [slots, 1] (average over batch)
+
+        # Update EMA during training
+        if self.training:
+            self.ema_rms = self.decay * self.ema_rms + (1 - self.decay) * slot_rms.detach()
+
+        # Normalize using EMA statistics
+        x_norm = x * torch.rsqrt(self.ema_rms + self.eps)
+
+        # Apply learned scale per slot
+        if self.scale is not None:
+            x_norm = x_norm * self.scale
+
+        # Apply gating mechanism
+        if self.use_gate:
+            gate = torch.sigmoid(self.gate)  # [slots, 1, 1]
+            return gate * x_norm + (1 - gate) * x
+
+        return x_norm
+
+class AdaptiveRMSMemoryNorm(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        use_gate: bool = True,
+        decay: float = 0.99,
+        init_scale: float = 1.0,
+        init_gate: float = -4.0  # Start with gate closed (no normalization)
+    ):
+        super().__init__()
+        self.use_gate = use_gate
+        self.scale = nn.Parameter(torch.ones(dim) * init_scale)
+        self.gate = nn.Parameter(torch.tensor([init_gate]))  # Scalar gate for this layer
+        self.eps = 1e-6
+        self.decay = decay
+        self.register_buffer("ema_rms", torch.ones(1))  # Scalar EMA RMS for the entire layer's STM
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # x shape: [batch_size, num_slots, dim]
+        if self.training and hasattr(self, 'ema_rms'):
+            # Compute current RMS across all slots and batch (scalar)
+            current_rms = x.pow(2).mean(-1).mean().sqrt()
+            self.ema_rms = self.ema_rms * self.decay + current_rms * (1 - self.decay)
+            rms = self.ema_rms
+        else:
+            # Compute RMS per slot (mean over dim)
+            rms = x.pow(2).mean(-1, keepdim=True).sqrt()  # [batch_size, num_slots, 1]
+
+        # Normalize each slot's embedding vector
+        normalized = x * torch.rsqrt(rms + self.eps)
+        normalized = normalized * self.scale  # Apply per-dimension scaling
+
+        if self.use_gate:
+            gate_factor = torch.sigmoid(self.gate)  # Scalar gate (0-1)
+            return normalized * gate_factor + x * (1 - gate_factor)
+        else:
+            return normalized
+
+class SimpleRMSMemoryNorm(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        use_gate: bool = True,
+        init_scale: float = 1.0,
+        init_gate: float = -4.0
+    ):
+        super().__init__()
+        self.use_gate = use_gate
+        self.scale = nn.Parameter(torch.ones(dim) * init_scale)
+        self.gate = nn.Parameter(torch.tensor([init_gate]))  # Scalar gate
+        self.eps = 1e-6
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        rms = x.pow(2).mean(-1, keepdim=True).sqrt()  # [batch_size, num_slots, 1]
+        normalized = x * torch.rsqrt(rms + self.eps)
+        normalized = normalized * self.scale  # Apply per-dimension scaling
+
+        if self.use_gate:
+            gate_factor = torch.sigmoid(self.gate)  # Scalar gate (0-1)
+            return normalized * gate_factor + x * (1 - gate_factor)
+        else:
+            return normalized
+
+class MemoryLayerNorm(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        use_gate: bool = True,
+        init_scale: float = 1.0,
+        init_gate: float = -4.0  # Start with gate closed (no normalization)
+    ):
+        super().__init__()
+        self.use_gate = use_gate
+        self.norm = nn.LayerNorm(dim)  # Normalizes across embedding dimensions per slot
+        self.gate = nn.Parameter(torch.tensor([init_gate]))  # Scalar gate for this layer
+        self.scale = nn.Parameter(torch.ones(dim) * init_scale)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        normalized = self.norm(x)  # Apply LayerNorm across embedding dimensions (per slot)
+        normalized = normalized * self.scale  # Per-dimension scaling
+
+        if self.use_gate:
+            gate_factor = torch.sigmoid(self.gate)  # Scalar gate (0-1)
+            return normalized * gate_factor + x * (1 - gate_factor)
+        else:
+            return normalized
+
+class MemoryNormConfig(TypedDict):
+    num_slots: int
+    decay: float
+    use_scale: bool
+    use_gate: bool
+    init_gate: float
+    init_scale: float
+
+def init_memory_norm(
+    norm_type: str,
+    dim: int,
+    num_slots: int = None,
+    decay: float = 0.99,
+    use_scale: bool = True,
+    use_gate: bool = True,
+    init_gate: float = -4.0,
+    init_scale: float = 1.0,
+) -> nn.Module:
+    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":
+        return SimpleRMSMemoryNorm(dim, use_gate, init_scale, init_gate)
+    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)
+    return MemoryLayerNorm(dim, use_gate, init_scale, init_gate)

src/memory/stm.py

@@ -0,0 +1,53 @@
+import torch
+import torch.nn as nn
+
+class ShortTermMemory(nn.Module):
+    """Short-term memory module for the Attention-based Memory System"""
+
+    def __init__(self, num_layers: int, embed_dim: int, stm_size: int, init_type: str = 'normal',
+                 is_trainable: bool = False, *args, **kwargs):
+        super(ShortTermMemory, self).__init__(*args, **kwargs)
+        self.num_layers = num_layers
+        self.embed_dim = embed_dim
+        self.stm_size = stm_size
+        self.is_trainable = is_trainable
+        assert init_type in ['normal', 'standard', 'uniform', 'ones', 'zeros'], \
+            'STM init type must be one of "normal", "standard", "uniform", "ones", "zeros"'
+        if init_type == 'normal':
+            stm = torch.normal(0, 0.02, (num_layers, stm_size, embed_dim))
+        elif init_type == 'standard':
+            stm = torch.normal(0, 1, (num_layers, stm_size, embed_dim))
+        elif init_type == 'uniform':
+            stm = torch.rand(num_layers, stm_size, embed_dim) * 0.02
+        elif init_type == 'ones':
+            stm = torch.ones(num_layers, stm_size, embed_dim)
+        else:
+            stm = torch.zeros(num_layers, stm_size, embed_dim)
+
+        if self.is_trainable:
+            self.memory = nn.Parameter(stm)
+        else:
+            self.register_buffer('memory', stm)
+
+    def forward(self, layer: int) -> torch.Tensor:
+        return self.memory[layer].unsqueeze(0)
+
+    def update_layer(self, layer: int, new_stm: torch.Tensor):
+        self.memory[layer] = new_stm
+
+    def update_all(self, new_stm: torch.Tensor):
+        self.memory.copy_(new_stm)
+
+    def make_trainable(self):
+        if not self.is_trainable:
+            self.is_trainable = True
+            initial_stm = self.memory.clone()
+            del self.memory
+            self.memory = nn.Parameter(initial_stm)
+
+    def freeze(self):
+        if self.is_trainable:
+            self.requires_grad_(False)
+            trained_stm = self.memory.clone()
+            del self.memory
+            self.register_buffer('memory', trained_stm)

src/rxt/models.py

@@ -0,0 +1,180 @@
+import torch
+from torch import nn
+from typing import TypedDict, Union
+from huggingface_hub import PyTorchModelHubMixin
+from src.transformers.positional import RotaryPositionalEmbedding
+from src.transformers.attention import init_attention
+from src.transformers.layers import ReactiveTransformerLayer
+from src.transformers.models import ReactiveTransformerBase, ReactiveTransformerEncoder, ReactiveTransformerDecoder
+from src.transformers.ff import get_activation_layer
+from src.memory.stm import ShortTermMemory
+from src.utils import get_model_size
+
+
+class RxTAlphaComponentConfig(TypedDict):
+    num_layers: int
+    vocab_size: int
+    embed_dim: int
+    ff_dim: int
+    att_heads: int
+    seq_len: int
+    stm_size: int
+    use_flash_attention: bool
+    use_gated: bool
+    ff_activation: str
+    ff_dropout: float
+    att_dropout: float
+    use_rms_norm: bool
+    att_groups: int
+    use_moe: bool
+    num_experts: int
+    moe_top_k: int
+    self_att_type: str
+    cross_att_type: str
+
+
+class RxTAlphaComponentBase(nn.Module, PyTorchModelHubMixin):
+    """Base class for RxT-Alpha (Reactive Transformer) components (encoder and decoder)"""
+
+    def __init__(
+            self,
+            is_causal: bool,
+            num_layers: int = 12,
+            vocab_size: int = 20000,
+            embed_dim: int = 512,
+            ff_dim: int = 1536,
+            att_heads: int = 16,
+            seq_len: int = 1024,
+            stm_size: int = 1024,
+            use_flash_attention: bool = True,
+            use_gated: bool = True,
+            ff_activation: str = "swish",
+            ff_dropout: float = 0.0,
+            att_dropout: float = 0.0,
+            use_rms_norm: bool = True,
+            att_groups: int = 1,
+            use_moe: bool = False,
+            num_experts: int = 1,
+            moe_top_k: int = 1,
+            self_att_type: str = 'gqa',
+            cross_att_type: str = 'mqa',
+            **kwargs
+    ):
+        super(RxTAlphaComponentBase, self).__init__(**kwargs)
+        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'], 'Self-attention type could be "mha", "gqa", "mqa"'
+        assert cross_att_type in ['mha', 'gqa', 'mqa'], 'Memory cross-attention type could be "mha", "gqa", "mqa"'
+
+        embedding = nn.Embedding(vocab_size, embed_dim)
+        rope = RotaryPositionalEmbedding(embed_dim // att_heads, seq_len)
+        stm = ShortTermMemory(num_layers, embed_dim, stm_size)
+
+        ff_activation = get_activation_layer(ff_activation)
+
+        layers = nn.ModuleList([
+            ReactiveTransformerLayer(
+                embed_dim,
+                ff_dim,
+                use_gated=use_gated,
+                use_moe=use_moe,
+                num_experts=num_experts,
+                moe_top_k=moe_top_k,
+                ff_activation=ff_activation,
+                ff_dropout=ff_dropout,
+                use_rms_norm=use_rms_norm,
+                self_attention=init_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),
+                memory_cross_attention=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, rope_only_for_query=True,
+                                                      is_causal=is_causal)
+            ) for _ in range(num_layers)
+        ])
+        self.model = self._init_model(stm, layers, embedding, use_flash_attention, embed_dim, vocab_size)
+
+    def _init_model(self, stm: ShortTermMemory, layers: nn.ModuleList, embedding: nn.Embedding,
+                    use_flash_attention: bool, embed_dim: int, vocab_size: int) -> ReactiveTransformerBase:
+        pass
+
+    def params_count(self):
+        return get_model_size(self.model)
+
+    def load_shared_embedding(self, embedding: nn.Embedding):
+        self.model.embedding = embedding
+
+    def load_shared_memory(self, stm: ShortTermMemory):
+        self.model.stm = stm
+
+    def forward(self, x: torch.Tensor, attention_mask: torch.Tensor = None) -> Union[
+        torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
+        return self.model(x, attention_mask=attention_mask)
+
+
+class RxTAlphaEncoder(RxTAlphaComponentBase, pipeline_tag="fill-mask", license="apache-2.0"):
+    """RxT-Alpha (Reactive Transformer) encoder model"""
+
+    def __init__(self, **kwargs: RxTAlphaComponentConfig):
+        super(RxTAlphaEncoder, self).__init__(False, **kwargs)
+
+    def _init_model(
+            self,
+            stm: ShortTermMemory,
+            layers: nn.ModuleList,
+            embedding: nn.Embedding,
+            use_flash_attention: bool,
+            embed_dim: int,
+            vocab_size: int
+    ) -> ReactiveTransformerEncoder:
+        return ReactiveTransformerEncoder(
+            stm=stm,
+            embedding=embedding,
+            own_layers=layers,
+            use_flash_attention=use_flash_attention,
+        )
+
+    def forward(self, x: torch.Tensor, attention_mask: torch.Tensor = None) -> tuple[torch.Tensor, torch.Tensor]:
+        return self.model(x, attention_mask=attention_mask)
+
+
+class RxTAlphaDecoder(RxTAlphaComponentBase, pipeline_tag="text-generation", license="apache-2.0"):
+    """RxT-Alpha (Reactive Transformer) decoder model"""
+
+    def __init__(self, **kwargs):
+        super(RxTAlphaDecoder, self).__init__(True, **kwargs)
+
+    def _init_model(
+            self, stm: ShortTermMemory,
+            layers: nn.ModuleList,
+            embedding: nn.Embedding,
+            use_flash_attention: bool,
+            embed_dim: int,
+            vocab_size: int
+    ) -> ReactiveTransformerDecoder:
+        return ReactiveTransformerDecoder(
+            embed_dim,
+            vocab_size,
+            stm=stm,
+            embedding=embedding,
+            own_layers=layers,
+            use_flash_attention=use_flash_attention,
+        )
+
+    def forward(self, x: torch.Tensor, attention_mask: torch.Tensor = None) -> tuple[torch.Tensor, torch.Tensor]:
+        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
+

src/training/base.py

@@ -0,0 +1,275 @@
+import torch
+import math
+import os
+
+from abc import ABC, abstractmethod
+from torch.utils.tensorboard import SummaryWriter
+import torch.distributed as dist
+from torch.nn.parallel import DistributedDataParallel
+from typing import Callable
+from callbacks import TrainerCallback
+
+
+class BaseTrainer(ABC):
+    def __init__(
+            self,
+            model: torch.nn.Module,
+            device: torch.device,
+            optimizer: torch.optim.Optimizer = None,
+            dataset: torch.utils.data.Dataset = None,
+            validation_dataset: torch.utils.data.Dataset = None,
+            callbacks: list[TrainerCallback] = None,
+            log_dir: str = None,
+            use_ddp: bool = False,
+            use_amp: bool = False,
+            dtype: torch.dtype = None,
+            target_field_name: str = 'labels',
+            get_batch_size: Callable[[dict], int] = None,
+            gradient_accumulation_steps: int = 1,
+    ):
+        if get_batch_size is None:
+            self.get_batch_size = lambda batch: batch['attention_mask'].size(0)
+        else:
+            self.get_batch_size = get_batch_size
+        if use_amp:
+            self.model = model.to(device)
+        else:
+            self.model = model.to(device, dtype=dtype)
+        self.device = device
+        self.optimizer = optimizer
+        self.dataset = dataset
+        self.callbacks = callbacks or []
+        self.writer = SummaryWriter(log_dir) if log_dir else None
+        self.use_ddp = use_ddp
+        self.use_amp = use_amp
+        self.dtype = dtype
+        self.is_running = False
+        self.validation_dataset = validation_dataset
+        self.best_val_loss = float('inf')
+        self.validation_metrics = {}
+        self.target_field_name = target_field_name
+        self.total_tokens = 0
+        self.gradient_accumulation_steps = gradient_accumulation_steps
+        self.accumulated_loss = 0.0
+        self.optimizer_step_count = 0
+
+    @abstractmethod
+    def compute_loss(self, batch: dict[str, torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]:
+        pass
+
+    def train_step(self, batch: dict[str, torch.Tensor], _batch_idx: int) -> torch.Tensor:
+        if self.use_amp:
+            batch = {k: v.to(self.device) for k, v in batch.items()}
+            with torch.amp.autocast(device_type=self.device.type, dtype=self.dtype):
+                loss, _ = self.compute_loss(batch)
+        else:
+            batch = {k: v.to(self.device, dtype=self.dtype) for k, v in batch.items()}
+            loss, _ = self.compute_loss(batch)
+        return loss
+
+    def __call__(
+            self,
+            epochs: int,
+            batch_size: int,
+            dataset: torch.utils.data.Dataset = None,
+            optimizer: torch.optim.Optimizer = None,
+            scheduler: torch.optim.lr_scheduler.LRScheduler = None,
+    ) -> None:
+        self.is_running = True
+        if dataset is None:
+            assert self.dataset is not None, 'You have to specify a dataset for training'
+            dataset = self.dataset
+        if optimizer is None:
+            assert self.optimizer is not None, 'You have to specify an optimizer for training'
+            optimizer = self.optimizer
+
+        if self.use_ddp:
+            rank = int(os.environ['RANK'])
+            world_size = int(os.environ['WORLD_SIZE'])
+            dist.init_process_group(backend='nccl', rank=rank, world_size=world_size)
+            self.model = DistributedDataParallel(self.model, device_ids=[self.device.index])
+            train_sampler = torch.utils.data.DistributedSampler(dataset, shuffle=True)
+            dataloader = torch.utils.data.DataLoader(
+                dataset,
+                batch_size=batch_size,
+                sampler=train_sampler,
+                pin_memory=True,
+            )
+        else:
+            train_sampler = None
+            dataloader = torch.utils.data.DataLoader(
+                dataset,
+                batch_size=batch_size,
+                shuffle=True,
+                pin_memory=True
+            )
+
+        scaler = torch.amp.GradScaler() if self.use_amp else None
+
+        self.model.train()
+        for epoch in range(epochs):
+            if self.is_running:
+                if train_sampler is not None:
+                    train_sampler.set_epoch(epoch)
+                self._run_epoch(dataloader, epoch, optimizer, batch_size, scaler=scaler, scheduler=scheduler)
+                if self.use_ddp:
+                    dist.barrier()
+
+        if self.use_ddp:
+            dist.destroy_process_group()
+        self.is_running = False
+        self.on_training_end()
+
+    def _run_epoch(
+            self,
+            dataloader: torch.utils.data.DataLoader,
+            epoch: int,
+            optimizer: torch.optim.Optimizer,
+            batch_size: int,
+            scaler: torch.cuda.amp.GradScaler = None,
+            scheduler: torch.optim.lr_scheduler.LRScheduler = None,
+    ) -> None:
+        for callback in self.callbacks:
+            callback.on_epoch_start(self.model, epoch)
+
+        self.accumulated_loss = 0.0
+        self.optimizer_step_count = 0
+
+        for batch_idx, batch in enumerate(dataloader):
+            if self.is_running:
+                for callback in self.callbacks:
+                    callback.on_batch_start(self.model, batch_idx, batch)
+                if self.get_batch_size(batch) == batch_size:
+                    loss = self.train_step(batch, batch_idx)
+                    self.accumulated_loss += loss.item()
+                    loss = loss / self.gradient_accumulation_steps
+
+                    if self.use_amp:
+                        scaler.scale(loss).backward()
+                    else:
+                        loss.backward()
+
+                    self.optimizer_step_count += 1
+                    if self.optimizer_step_count % self.gradient_accumulation_steps == 0:
+                        # Clip gradients after accumulation
+                        if self.use_amp:
+                            scaler.unscale_(optimizer)
+                        torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0, error_if_nonfinite=False)
+                        if self.use_amp:
+                            scaler.step(optimizer)
+                            scaler.update()
+                        else:
+                            optimizer.step()
+
+                        optimizer.zero_grad()
+
+                        if scheduler is not None:
+                            scheduler.step()
+
+                        if self.writer:
+                            loss_item = self.accumulated_loss / self.gradient_accumulation_steps
+                            self.writer.add_scalar(
+                                'Loss/train',
+                                loss_item,
+                                epoch * len(dataloader) + batch_idx
+                            )
+                            self.writer.add_scalar(
+                                'Loss per epoch/train',
+                                loss_item,
+                                batch_idx
+                            )
+                            self.writer.add_scalar(
+                                'Perplexity/train',
+                                torch.exp(torch.tensor(loss_item)),
+                                epoch * len(dataloader) + batch_idx
+                            )
+                        self.accumulated_loss = 0.0
+                        self.optimizer_step_count = 0
+
+                    if self.writer:
+                        self.total_tokens += batch['attention_mask'].sum().item()
+                        self.writer.add_scalar('Processed tokens', self.total_tokens,
+                                               epoch * len(dataloader) + batch_idx)
+
+                    for callback in self.callbacks:
+                        should_stop = callback.on_batch_end(self.model, batch_idx, loss.item(), batch)
+                        if should_stop:
+                            self.is_running = False
+
+        if self.validation_dataset:
+            val_loss, val_metrics = self.validate(batch_size)
+            val_loss_tensor = torch.tensor(val_loss).to(self.device)
+            if self.use_ddp:
+                dist.all_reduce(val_loss_tensor, op=dist.ReduceOp.SUM)
+                val_loss = val_loss_tensor.item() / dist.get_world_size()
+            self.validation_metrics[epoch] = val_metrics
+
+            if self.writer:
+                self._valid_writer(epoch, val_loss, val_metrics)
+
+            for callback in self.callbacks:
+                callback.on_validation_end(self.model, epoch, val_loss, val_metrics)
+
+        for callback in self.callbacks:
+            should_stop = callback.on_epoch_end(self.model, epoch)
+            if should_stop:
+                self.is_running = False
+
+        if self.writer:
+            self.writer.flush()
+
+    def on_training_end(self):
+        for callback in self.callbacks:
+            callback.on_training_end(self.model)
+        if self.writer:
+            self.writer.close()
+
+    def _valid_writer(self, epoch: int, val_loss: float, val_metrics: dict):
+        self.writer.add_scalar('Loss/validation', val_loss, epoch)
+        self.writer.add_scalar('Perplexity/validation', math.exp(val_loss), epoch)
+        if val_metrics['accuracy']:
+            self.writer.add_scalar('Accuracy/validation', val_metrics['accuracy'], epoch)
+
+    def valid_step(self, batch: dict[str, torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]:
+        if self.use_amp:
+            batch = {k: v.to(self.device) for k, v in batch.items()}
+            with torch.amp.autocast(device_type=self.device.type, dtype=self.dtype):
+                loss, outputs = self.compute_loss(batch)
+        else:
+            batch = {k: v.to(self.device, dtype=self.dtype) for k, v in batch.items()}
+            loss, outputs = self.compute_loss(batch)
+        return loss, outputs
+
+    def _valid_loader(self, batch_size: int):
+        val_dataset = self.validation_dataset
+        if self.use_ddp:
+            val_sampler = torch.utils.data.DistributedSampler(val_dataset, shuffle=False)
+            return torch.utils.data.DataLoader(
+                val_dataset,
+                batch_size=batch_size,
+                pin_memory=True,
+                sampler=val_sampler,
+            )
+        else:
+            return torch.utils.data.DataLoader(
+                val_dataset,
+                batch_size=batch_size,
+                shuffle=False,
+                pin_memory=True
+            )
+
+    def validate(self, batch_size: int) -> tuple[float, dict]:
+        self.model.eval()
+        val_dataloader = self._valid_loader(batch_size)
+        val_loss = 0.0
+
+        with torch.no_grad():
+            for batch in val_dataloader:
+                if self.get_batch_size(batch) == batch_size:
+                    loss, outputs = self.valid_step(batch)
+                    val_loss += loss.item()
+
+        avg_loss = val_loss / len(val_dataloader)
+        metrics = {}
+        self.model.train()
+        return avg_loss, metrics