ALSv/self-forcing

1# Modified from transformers.models.xlm_roberta.modeling_xlm_roberta2# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.3import torch4import torch.nn as nn5import torch.nn.functional as F6 7__all__ = ['XLMRoberta', 'xlm_roberta_large']8 9 10class SelfAttention(nn.Module):11 12    def __init__(self, dim, num_heads, dropout=0.1, eps=1e-5):13        assert dim % num_heads == 014        super().__init__()15        self.dim = dim16        self.num_heads = num_heads17        self.head_dim = dim // num_heads18        self.eps = eps19 20        # layers21        self.q = nn.Linear(dim, dim)22        self.k = nn.Linear(dim, dim)23        self.v = nn.Linear(dim, dim)24        self.o = nn.Linear(dim, dim)25        self.dropout = nn.Dropout(dropout)26 27    def forward(self, x, mask):28        """29        x:   [B, L, C].30        """31        b, s, c, n, d = *x.size(), self.num_heads, self.head_dim32 33        # compute query, key, value34        q = self.q(x).reshape(b, s, n, d).permute(0, 2, 1, 3)35        k = self.k(x).reshape(b, s, n, d).permute(0, 2, 1, 3)36        v = self.v(x).reshape(b, s, n, d).permute(0, 2, 1, 3)37 38        # compute attention39        p = self.dropout.p if self.training else 0.040        x = F.scaled_dot_product_attention(q, k, v, mask, p)41        x = x.permute(0, 2, 1, 3).reshape(b, s, c)42 43        # output44        x = self.o(x)45        x = self.dropout(x)46        return x47 48 49class AttentionBlock(nn.Module):50 51    def __init__(self, dim, num_heads, post_norm, dropout=0.1, eps=1e-5):52        super().__init__()53        self.dim = dim54        self.num_heads = num_heads55        self.post_norm = post_norm56        self.eps = eps57 58        # layers59        self.attn = SelfAttention(dim, num_heads, dropout, eps)60        self.norm1 = nn.LayerNorm(dim, eps=eps)61        self.ffn = nn.Sequential(62            nn.Linear(dim, dim * 4), nn.GELU(), nn.Linear(dim * 4, dim),63            nn.Dropout(dropout))64        self.norm2 = nn.LayerNorm(dim, eps=eps)65 66    def forward(self, x, mask):67        if self.post_norm:68            x = self.norm1(x + self.attn(x, mask))69            x = self.norm2(x + self.ffn(x))70        else:71            x = x + self.attn(self.norm1(x), mask)72            x = x + self.ffn(self.norm2(x))73        return x74 75 76class XLMRoberta(nn.Module):77    """78    XLMRobertaModel with no pooler and no LM head.79    """80 81    def __init__(self,82                 vocab_size=250002,83                 max_seq_len=514,84                 type_size=1,85                 pad_id=1,86                 dim=1024,87                 num_heads=16,88                 num_layers=24,89                 post_norm=True,90                 dropout=0.1,91                 eps=1e-5):92        super().__init__()93        self.vocab_size = vocab_size94        self.max_seq_len = max_seq_len95        self.type_size = type_size96        self.pad_id = pad_id97        self.dim = dim98        self.num_heads = num_heads99        self.num_layers = num_layers100        self.post_norm = post_norm101        self.eps = eps102 103        # embeddings104        self.token_embedding = nn.Embedding(vocab_size, dim, padding_idx=pad_id)105        self.type_embedding = nn.Embedding(type_size, dim)106        self.pos_embedding = nn.Embedding(max_seq_len, dim, padding_idx=pad_id)107        self.dropout = nn.Dropout(dropout)108 109        # blocks110        self.blocks = nn.ModuleList([111            AttentionBlock(dim, num_heads, post_norm, dropout, eps)112            for _ in range(num_layers)113        ])114 115        # norm layer116        self.norm = nn.LayerNorm(dim, eps=eps)117 118    def forward(self, ids):119        """120        ids: [B, L] of torch.LongTensor.121        """122        b, s = ids.shape123        mask = ids.ne(self.pad_id).long()124 125        # embeddings126        x = self.token_embedding(ids) + \127            self.type_embedding(torch.zeros_like(ids)) + \128            self.pos_embedding(self.pad_id + torch.cumsum(mask, dim=1) * mask)129        if self.post_norm:130            x = self.norm(x)131        x = self.dropout(x)132 133        # blocks134        mask = torch.where(135            mask.view(b, 1, 1, s).gt(0), 0.0,136            torch.finfo(x.dtype).min)137        for block in self.blocks:138            x = block(x, mask)139 140        # output141        if not self.post_norm:142            x = self.norm(x)143        return x144 145 146def xlm_roberta_large(pretrained=False,147                      return_tokenizer=False,148                      device='cpu',149                      **kwargs):150    """151    XLMRobertaLarge adapted from Huggingface.152    """153    # params154    cfg = dict(155        vocab_size=250002,156        max_seq_len=514,157        type_size=1,158        pad_id=1,159        dim=1024,160        num_heads=16,161        num_layers=24,162        post_norm=True,163        dropout=0.1,164        eps=1e-5)165    cfg.update(**kwargs)166 167    # init a model on device168    with torch.device(device):169        model = XLMRoberta(**cfg)170    return model171