Add modeling_expivme_diffusion.py

modeling_expivme_diffusion.py

@@ -0,0 +1,179 @@
+"""HuggingFace Transformers model for ExpIvme-DiffusionConversate-v1-Instruct (SFT)."""
+
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn.attention import SDPBackend, sdpa_kernel
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import ModelOutput
+
+try:
+    from .configuration_expivme_diffusion import ExpIvmeDiffusionConfig
+except ImportError:
+    from configuration_expivme_diffusion import ExpIvmeDiffusionConfig
+
+if torch.cuda.is_available() and torch.cuda.get_device_capability(0)[0] >= 8:
+    PREFERRED_SDPA_BACKENDS = [SDPBackend.CUDNN_ATTENTION, SDPBackend.FLASH_ATTENTION, SDPBackend.EFFICIENT_ATTENTION, SDPBackend.MATH]
+else:
+    PREFERRED_SDPA_BACKENDS = [SDPBackend.EFFICIENT_ATTENTION, SDPBackend.MATH]
+
+
+def _precompute_rope_freqs(head_dim, max_seq_len, theta, device=None):
+    freqs = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim))
+    positions = torch.arange(max_seq_len, device=device).float()
+    angles = torch.outer(positions, freqs)
+    return torch.cos(angles), torch.sin(angles)
+
+
+def _apply_rope(x, rope_cos_sin):
+    cos, sin = rope_cos_sin
+    B, H, T, D = x.shape
+    x1 = x[..., 0::2]
+    x2 = x[..., 1::2]
+    cos = cos.view(1, 1, T, D // 2).to(x.dtype)
+    sin = sin.view(1, 1, T, D // 2).to(x.dtype)
+    out1 = x1 * cos - x2 * sin
+    out2 = x1 * sin + x2 * cos
+    return torch.stack([out1, out2], dim=-1).reshape(B, H, T, D).type_as(x)
+
+
+class ExpIvmeRMSNorm(nn.Module):
+    def __init__(self, dim, eps=1e-5):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x):
+        dtype = x.dtype
+        x = x.float()
+        rms = torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)
+        return (x * rms).to(dtype) * self.weight
+
+
+class ExpIvmeSelfAttention(nn.Module):
+    def __init__(self, hidden_dim, n_heads, dropout=0.0):
+        super().__init__()
+        self.n_heads = n_heads
+        self.head_dim = hidden_dim // n_heads
+        self.dropout = dropout
+        self.q_proj = nn.Linear(hidden_dim, hidden_dim, bias=False)
+        self.k_proj = nn.Linear(hidden_dim, hidden_dim, bias=False)
+        self.v_proj = nn.Linear(hidden_dim, hidden_dim, bias=False)
+        self.out_proj = nn.Linear(hidden_dim, hidden_dim, bias=False)
+
+    def forward(self, x, rope_freqs, attn_mask=None):
+        B, T, C = x.shape
+        q = self.q_proj(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
+        k = self.k_proj(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
+        v = self.v_proj(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
+        q = _apply_rope(q, rope_freqs)
+        k = _apply_rope(k, rope_freqs)
+        with sdpa_kernel(PREFERRED_SDPA_BACKENDS):
+            out = F.scaled_dot_product_attention(
+                q, k, v, attn_mask=attn_mask, is_causal=False,
+                dropout_p=self.dropout if self.training else 0.0,
+            )
+        out = out.transpose(1, 2).contiguous().view(B, T, C)
+        return self.out_proj(out)
+
+
+class ExpIvmeSwiGLU(nn.Module):
+    def __init__(self, hidden_dim, ffn_mult):
+        super().__init__()
+        inner_dim = int(hidden_dim * ffn_mult * 2 / 3)
+        inner_dim = ((inner_dim + 7) // 8) * 8
+        self.gate_proj = nn.Linear(hidden_dim, inner_dim, bias=False)
+        self.up_proj = nn.Linear(hidden_dim, inner_dim, bias=False)
+        self.down_proj = nn.Linear(inner_dim, hidden_dim, bias=False)
+
+    def forward(self, x):
+        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class ExpIvmeBlock(nn.Module):
+    def __init__(self, hidden_dim, n_heads, ffn_mult, norm_eps, dropout=0.0):
+        super().__init__()
+        self.attn_norm = ExpIvmeRMSNorm(hidden_dim, eps=norm_eps)
+        self.attn = ExpIvmeSelfAttention(hidden_dim, n_heads, dropout)
+        self.ffn_norm = ExpIvmeRMSNorm(hidden_dim, eps=norm_eps)
+        self.ffn = ExpIvmeSwiGLU(hidden_dim, ffn_mult)
+
+    def forward(self, x, rope_freqs, attn_mask=None):
+        x = x + self.attn(self.attn_norm(x), rope_freqs, attn_mask=attn_mask)
+        x = x + self.ffn(self.ffn_norm(x))
+        return x
+
+
+@dataclass
+class DiffusionLMOutput(ModelOutput):
+    loss: torch.FloatTensor = None
+    logits: torch.FloatTensor = None
+
+
+class ExpIvmeForDiffusionLMHub(PreTrainedModel):
+    config_class = ExpIvmeDiffusionConfig
+    base_model_prefix = "model"
+    _tied_weights_keys = {"lm_head.weight": "model.tok_embed.weight"}
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = nn.Module()
+        self.model.tok_embed = nn.Embedding(config.vocab_size, config.hidden_dim)
+        self.model.blocks = nn.ModuleList([
+            ExpIvmeBlock(config.hidden_dim, config.n_heads, config.ffn_mult, config.norm_eps, config.dropout)
+            for _ in range(config.n_layers)
+        ])
+        self.model.final_norm = ExpIvmeRMSNorm(config.hidden_dim, eps=config.norm_eps)
+        self.lm_head = nn.Linear(config.hidden_dim, config.vocab_size, bias=False)
+        self.head_dim = config.hidden_dim // config.n_heads
+        self.rope_theta = config.rope_theta
+        self.post_init()
+        if config.tie_word_embeddings:
+            self.tie_weights()
+
+    def get_input_embeddings(self):
+        return self.model.tok_embed
+
+    def set_input_embeddings(self, value):
+        self.model.tok_embed = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def forward(self, input_ids, attention_mask=None, labels=None, mask_positions=None, t=None, return_dict=True, **kw):
+        B, T = input_ids.shape
+        rope = _precompute_rope_freqs(self.head_dim, T, self.rope_theta, device=input_ids.device)
+
+        sdpa_mask = None
+        if attention_mask is not None:
+            sdpa_mask = torch.zeros(B, 1, 1, T, dtype=torch.float32, device=input_ids.device)
+            sdpa_mask.masked_fill_(attention_mask[:, None, None, :] == 0, float("-inf"))
+            sdpa_mask = sdpa_mask.to(dtype=self.model.tok_embed.weight.dtype)
+
+        x = self.model.tok_embed(input_ids)
+        for block in self.model.blocks:
+            x = block(x, rope, attn_mask=sdpa_mask)
+        x = self.model.final_norm(x)
+        logits = self.lm_head(x)
+
+        loss = None
+        if labels is not None and mask_positions is not None:
+            ce = F.cross_entropy(
+                logits.view(-1, self.config.vocab_size), labels.view(-1), ignore_index=-100, reduction="none",
+            ).view(B, T)
+            ce = ce * mask_positions.float()
+            per_example_loss = ce.sum(dim=1)
+            if t is not None:
+                weight = 1.0 / t.clamp(min=1e-3)
+                per_example_loss = per_example_loss * weight
+            n_masked = mask_positions.float().sum(dim=1).clamp(min=1.0)
+            loss = (per_example_loss / n_masked).mean()
+
+        if not return_dict:
+            return (loss, logits) if loss is not None else (logits,)
+        return DiffusionLMOutput(loss=loss, logits=logits)
+
+
+__all__ = ["ExpIvmeDiffusionConfig", "ExpIvmeForDiffusionLMHub"]