modeling_mini_llada.py

# modeling_llada.py
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.checkpoint import checkpoint

from transformers import PreTrainedModel
from transformers.modeling_outputs import MaskedLMOutput
from ko_mini_llada.models.configuration_mini_llada import MiniLLaDAConfig

# RoPE (Rotary Positional Embedding)
def precompute_freqs_cis(dim: int, max_len: int, theta: float = 10000.0):
    """
    Precomputes the frequencies for the rotary positional embeddings.

    Args:
        dim (int): The dimension of the embeddings.
        max_len (int): The maximum sequence length.
        theta (float, optional): The theta parameter for frequency calculation. Defaults to 10000.0.

    Returns:
        torch.Tensor: Precomputed complex frequencies.
    """
    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
    t = torch.arange(max_len, dtype=torch.float)
    freqs = torch.outer(t, freqs)
    freqs_cis = torch.polar(torch.ones_like(freqs), freqs)
    return freqs_cis

def apply_rotary_emb(xq: torch.Tensor, xk: torch.Tensor, freqs_cis: torch.Tensor):
    """
    Applies rotary positional embeddings to input tensors.

    Args:
        xq (torch.Tensor): Query tensor.
        xk (torch.Tensor): Key tensor.
        freqs_cis (torch.Tensor): Precomputed complex frequencies.

    Returns:
        Tuple[torch.Tensor, torch.Tensor]: Tensors with rotary embeddings applied to query and key.
    """
    xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
    xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))
    # Reshape freqs_cis for broadcasting. Assumes xq is (batch, n_heads, seq_len, head_dim)
    freqs_cis = freqs_cis[:xq_.shape[2]].view(1, 1, xq_.shape[2], xq_.shape[3])
    xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3)
    xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3)
    return xq_out.type_as(xq), xk_out.type_as(xk)

class Attention(nn.Module):
    """
    Multi-Head Attention module with Rotary Positional Embeddings.

    This module implements a multi-head attention mechanism, incorporating rotary
    positional embeddings (RoPE) for query and key tensors. It uses bidirectional
    attention as required by the LLaDA architecture.

    Args:
        dim (int): The input and output dimension of the module.
        heads (int): The number of attention heads.
    """
    def __init__(self, dim, heads):
        super().__init__()
        self.n_heads = heads
        self.head_dim = dim // heads
        self.q_proj = nn.Linear(dim, dim, bias=False)
        self.k_proj = nn.Linear(dim, dim, bias=False)
        self.v_proj = nn.Linear(dim, dim, bias=False)
        self.o_proj = nn.Linear(dim, dim, bias=False)

    def forward(self, x, freqs_cis, mask=None):
        """
        Forward pass for the Attention module.

        Args:
            x (torch.Tensor): Input tensor of shape (B, L, D).
            freqs_cis (torch.Tensor): Precomputed rotary frequencies.
            mask (torch.Tensor, optional): Attention mask. Defaults to None.

        Returns:
            torch.Tensor: Output tensor of shape (B, L, D).
        """
        B, L, D = x.shape
        q = self.q_proj(x).view(B, L, self.n_heads, self.head_dim).transpose(1, 2)
        k = self.k_proj(x).view(B, L, self.n_heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(x).view(B, L, self.n_heads, self.head_dim).transpose(1, 2)

        freqs_cis = freqs_cis.to(x.device)
        q, k = apply_rotary_emb(q, k, freqs_cis[:L])

        # LLaDA: is_causal=False (Bidirectional Attention)
        attn_output = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, is_causal=False)

        x = attn_output.transpose(1, 2).reshape(B, L, D)
        x = self.o_proj(x)
        return x

class Block(nn.Module):
    """
    Transformer block consisting of multi-head attention and a feed-forward network.

    This block applies pre-normalization (RMSNorm), followed by multi-head attention,
    a residual connection, another pre-normalization, a feed-forward network with
    SwiGLU activation, and a final residual connection.

    Args:
        dim (int): The dimension of the input and output.
        heads (int): The number of attention heads.
        intermediate_size (int): The intermediate size of the feed-forward network.
    """
    def __init__(self, dim, heads, intermediate_size):
        super().__init__()
        self.mha = Attention(dim, heads) # Multi-Head Attention
        self.norm1 = nn.RMSNorm(dim)
        self.norm2 = nn.RMSNorm(dim)

        # Feed-Forward Network with Gated Linear Unit
        self.gate_up_proj = nn.Linear(dim, intermediate_size * 2, bias=False)
        self.down_proj = nn.Linear(intermediate_size, dim, bias=False)
        self.gradient_checkpointing = False

    def run_block(self, x, freqs_cis, mask):
        """
        Actual forward pass logic.
        """
        # Multi-Head Attention
        residual = x

        x = self.norm1(x)
        x = self.mha(x, freqs_cis, mask)
        x = residual + x
        
        # Feed-Forward Network
        residual = x
        x = self.norm2(x)
        gate, up = self.gate_up_proj(x).chunk(2, dim=-1)
        x = F.silu(gate) * up
        x = self.down_proj(x)
        x = residual + x

        return x

    def forward(self, x, freqs_cis, mask=None):
        """
        Forward pass for the Transformer block.

        Args:
            x (torch.Tensor): Input tensor of shape (B, L, D).
            freqs_cis (torch.Tensor): Precomputed rotary frequencies.
            mask (torch.Tensor, optional): Attention mask. Defaults to None.

        Returns:
            torch.Tensor: Output tensor of shape (B, L, D).
        """
        if self.gradient_checkpointing and self.training:
            return checkpoint(self.run_block, x, freqs_cis, mask, use_reentrant=False)
        else:
            return self.run_block(x, freqs_cis, mask)

class Transformer(nn.Module):
    """
    The core Transformer model for Mini-LLaDA.

    This class stacks multiple Transformer blocks to form the main network.
    It includes token embeddings, a series of Transformer blocks, a final
    normalization layer, and a linear head to project to the vocabulary size.

    Args:
        vocab_size (int): The size of the vocabulary.
        dim (int): The embedding dimension.
        depth (int): The number of Transformer blocks.
        heads (int): The number of attention heads.
        intermediate_size (int): The intermediate size of the feed-forward networks.
        max_seq_len (int, optional): The maximum sequence length. Defaults to 2048.
    """
    def __init__(self, vocab_size, dim, depth, heads, intermediate_size, max_seq_len=2048):
        super().__init__()
        self.embed = nn.Embedding(vocab_size, dim) # Token Embedding
        self.layers = nn.ModuleList([
            Block(dim, heads, intermediate_size) for _ in range(depth)
        ])
        self.norm = nn.RMSNorm(dim)
        self.head = nn.Linear(dim, vocab_size, bias=False)

        freqs_cis = precompute_freqs_cis(dim // heads, max_seq_len)
        self.register_buffer("freqs_cis", freqs_cis, persistent=False)

    def forward(self, input_ids, attention_mask=None):
        """
        Forward pass for the Transformer model.

        Args:
            input_ids (torch.Tensor): Input token IDs of shape (B, L).
            attention_mask (torch.Tensor, optional): Mask to avoid attending to padding tokens.
                Shape (B, L). Defaults to None.

        Returns:
            torch.Tensor: Output logits of shape (B, L, vocab_size).
        """
        x = self.embed(input_ids)

        if attention_mask is not None:
            attention_mask = (attention_mask == 1).view(x.shape[0], 1, 1, x.shape[1])

        attention_mask = None

        for layer in self.layers:
            x = layer(x, self.freqs_cis, attention_mask)
        
        x = self.norm(x)
        x = self.head(x)
        return x

class MiniLLaDA(PreTrainedModel):
    """
    The Mini-LLaDA model, a Transformer-based model for masked language modeling
    inspired by diffusion models.

    This model is designed for pre-training using a diffusion-like noising process
    where a variable number of tokens are masked and the model learns to predict
    the original tokens. It can be used for both training (with labels) and
    inference (without labels).

    Args:
        config (MiniLLaDAConfig): The configuration object for the model.
    """
    config_class = MiniLLaDAConfig
    _supports_gradient_checkpointing = True

    def __init__(self, config: MiniLLaDAConfig):
        super().__init__(config)
        
        # 1. load backbone model
        self.network = Transformer(
            vocab_size=config.vocab_size,
            dim=config.dim,
            depth=config.depth,
            heads=config.head,
            intermediate_size=config.intermediate_size,
            max_seq_len=config.max_seq_len
        )
        self.mask_token_id = config.mask_token_id
        self.network.apply(self._init_weights)

    @property
    def supports_gradient_checkpointing(self):
        return True

    def forward(self, input_ids, attention_mask=None, labels=None, **kwargs):
        """
        Forward pass for the Mini-LLaDA model.

        If `labels` are provided, the model operates in training/evaluation mode,
        performing the diffusion forward process, running the noised input through
        the network, and computing the loss.

        If `labels` are not provided, the model operates in inference mode, simply
        passing the `input_ids` through the network to get logits.

        Args:
            input_ids (torch.Tensor): Input token IDs of shape (B, L).
            attention_mask (torch.Tensor, optional): Mask to avoid attending to padding tokens.
                Shape (B, L). Defaults to None.
            labels (torch.Tensor, optional): Labels for computing the loss. In SFT, this is used
                to identify which parts of the sequence to mask (assistant's response).
                Shape (B, L). Defaults to None.
            **kwargs: Additional keyword arguments.

        Returns:
            transformers.modeling_outputs.MaskedLMOutput: An output object containing the
                loss (if labels are provided) and logits.
        """
        # 1. Training and Evaluation Mode
        if labels is not None:
            # Diffusion Forward Process
            t, noisy_x, mask_indices = self.forward_process(input_ids, labels)
            
            # Reverse Process
            # network outputs: MaskedLMOutput (logits, hidden_states, etc.)
            outputs = self.network(input_ids=noisy_x, attention_mask=attention_mask)

            # Compute Loss
            loss = self.compute_diffusion_loss(outputs, input_ids, mask_indices, attention_mask)
            
            return MaskedLMOutput(loss=loss, logits=outputs)

        # 2. Inference Mode
        else:
            outputs = self.network(input_ids=input_ids, attention_mask=attention_mask)
            return MaskedLMOutput(logits=outputs)

    def forward_process(self, input_ids, labels=None):
        """
        Simulates the diffusion forward process by noising the input sequence.

        A random timestep `t` is sampled for each sequence in the batch, which
        determines the probability of a token being masked. Tokens are replaced
        with a `mask_token_id`. During supervised fine-tuning (SFT), masking is
        restricted to the assistant's response part of the sequence, identified
        by `labels != -100`.

        Args:
            input_ids (torch.Tensor): The original input token IDs of shape (B, L).
            labels (torch.Tensor, optional): Labels used to restrict masking during SFT.
                Defaults to None.

        Returns:
            Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: A tuple containing:
                - t (torch.Tensor): The sampled timesteps for each sequence (shape B).
                - noisy_x (torch.Tensor): The noised input IDs with masks (shape B, L).
                - mask_indices (torch.Tensor): A boolean tensor indicating which tokens
                  were masked (shape B, L).
        """
        B, L = input_ids.shape
        device = input_ids.device

        t = torch.rand(B, device=device)
        mask_probs = t.unsqueeze(1).expand(B, L)

        if labels is not None:
            train_mask = (labels != -100).float() 
            mask_probs = mask_probs * train_mask  # Make the probabilities zero where labels == -100, which is user message parts.

        random_matrix = torch.rand(B, L, device=device)
        mask_indices = (random_matrix < mask_probs)

        noisy_x = torch.where(mask_indices, self.mask_token_id, input_ids)
        return t, noisy_x, mask_indices

    def compute_diffusion_loss(self, logits, input_ids, mask_indices, attention_mask):
        """
        Computes the cross-entropy loss for the masked language modeling task.

        The loss is calculated only for the positions that were masked during the
        forward process. It also respects the attention mask to avoid computing
        loss on padding tokens.

        Args:
            logits (torch.Tensor): The model's output logits of shape (B, L, V).
            input_ids (torch.Tensor): The original input token IDs of shape (B, L).
            mask_indices (torch.Tensor): A boolean tensor indicating masked positions (shape B, L).
            attention_mask (torch.Tensor): The attention mask for the input (shape B, L).

        Returns:
            torch.Tensor: The computed cross-entropy loss.
        """
        B, L, V = logits.shape
        logits_flat = logits.view(-1, V)
        target_flat = input_ids.view(-1)
        
        # We compute loss only on tokens that were masked AND are not padding.
        final_loss_mask = mask_indices.view(-1)
        if attention_mask is not None:
            final_loss_mask = final_loss_mask & attention_mask.view(-1).bool()
        
        target_flat = torch.where(final_loss_mask, target_flat, -100)

        return F.cross_entropy(logits_flat, target_flat, ignore_index=-100)