ALSv/self-forcing
0
1# Modified from transformers.models.t5.modeling_t52# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.3import logging4import math5 6import torch7import torch.nn as nn8import torch.nn.functional as F9 10from .tokenizers import HuggingfaceTokenizer11 12__all__ = [13 'T5Model',14 'T5Encoder',15 'T5Decoder',16 'T5EncoderModel',17]18 19 20def fp16_clamp(x):21 if x.dtype == torch.float16 and torch.isinf(x).any():22 clamp = torch.finfo(x.dtype).max - 100023 x = torch.clamp(x, min=-clamp, max=clamp)24 return x25 26 27def init_weights(m):28 if isinstance(m, T5LayerNorm):29 nn.init.ones_(m.weight)30 elif isinstance(m, T5Model):31 nn.init.normal_(m.token_embedding.weight, std=1.0)32 elif isinstance(m, T5FeedForward):33 nn.init.normal_(m.gate[0].weight, std=m.dim**-0.5)34 nn.init.normal_(m.fc1.weight, std=m.dim**-0.5)35 nn.init.normal_(m.fc2.weight, std=m.dim_ffn**-0.5)36 elif isinstance(m, T5Attention):37 nn.init.normal_(m.q.weight, std=(m.dim * m.dim_attn)**-0.5)38 nn.init.normal_(m.k.weight, std=m.dim**-0.5)39 nn.init.normal_(m.v.weight, std=m.dim**-0.5)40 nn.init.normal_(m.o.weight, std=(m.num_heads * m.dim_attn)**-0.5)41 elif isinstance(m, T5RelativeEmbedding):42 nn.init.normal_(43 m.embedding.weight, std=(2 * m.num_buckets * m.num_heads)**-0.5)44 45 46class GELU(nn.Module):47 48 def forward(self, x):49 return 0.5 * x * (1.0 + torch.tanh(50 math.sqrt(2.0 / math.pi) * (x + 0.044715 * torch.pow(x, 3.0))))51 52 53class T5LayerNorm(nn.Module):54 55 def __init__(self, dim, eps=1e-6):56 super(T5LayerNorm, self).__init__()57 self.dim = dim58 self.eps = eps59 self.weight = nn.Parameter(torch.ones(dim))60 61 def forward(self, x):62 x = x * torch.rsqrt(x.float().pow(2).mean(dim=-1, keepdim=True) +63 self.eps)64 if self.weight.dtype in [torch.float16, torch.bfloat16]:65 x = x.type_as(self.weight)66 return self.weight * x67 68 69class T5Attention(nn.Module):70 71 def __init__(self, dim, dim_attn, num_heads, dropout=0.1):72 assert dim_attn % num_heads == 073 super(T5Attention, self).__init__()74 self.dim = dim75 self.dim_attn = dim_attn76 self.num_heads = num_heads77 self.head_dim = dim_attn // num_heads78 79 # layers80 self.q = nn.Linear(dim, dim_attn, bias=False)81 self.k = nn.Linear(dim, dim_attn, bias=False)82 self.v = nn.Linear(dim, dim_attn, bias=False)83 self.o = nn.Linear(dim_attn, dim, bias=False)84 self.dropout = nn.Dropout(dropout)85 86 def forward(self, x, context=None, mask=None, pos_bias=None):87 """88 x: [B, L1, C].89 context: [B, L2, C] or None.90 mask: [B, L2] or [B, L1, L2] or None.91 """92 # check inputs93 context = x if context is None else context94 b, n, c = x.size(0), self.num_heads, self.head_dim95 96 # compute query, key, value97 q = self.q(x).view(b, -1, n, c)98 k = self.k(context).view(b, -1, n, c)99 v = self.v(context).view(b, -1, n, c)100 101 # attention bias102 attn_bias = x.new_zeros(b, n, q.size(1), k.size(1))103 if pos_bias is not None:104 attn_bias += pos_bias105 if mask is not None:106 assert mask.ndim in [2, 3]107 mask = mask.view(b, 1, 1,108 -1) if mask.ndim == 2 else mask.unsqueeze(1)109 attn_bias.masked_fill_(mask == 0, torch.finfo(x.dtype).min)110 111 # compute attention (T5 does not use scaling)112 attn = torch.einsum('binc,bjnc->bnij', q, k) + attn_bias113 attn = F.softmax(attn.float(), dim=-1).type_as(attn)114 x = torch.einsum('bnij,bjnc->binc', attn, v)115 116 # output117 x = x.reshape(b, -1, n * c)118 x = self.o(x)119 x = self.dropout(x)120 return x121 122 123class T5FeedForward(nn.Module):124 125 def __init__(self, dim, dim_ffn, dropout=0.1):126 super(T5FeedForward, self).__init__()127 self.dim = dim128 self.dim_ffn = dim_ffn129 130 # layers131 self.gate = nn.Sequential(nn.Linear(dim, dim_ffn, bias=False), GELU())132 self.fc1 = nn.Linear(dim, dim_ffn, bias=False)133 self.fc2 = nn.Linear(dim_ffn, dim, bias=False)134 self.dropout = nn.Dropout(dropout)135 136 def forward(self, x):137 x = self.fc1(x) * self.gate(x)138 x = self.dropout(x)139 x = self.fc2(x)140 x = self.dropout(x)141 return x142 143 144class T5SelfAttention(nn.Module):145 146 def __init__(self,147 dim,148 dim_attn,149 dim_ffn,150 num_heads,151 num_buckets,152 shared_pos=True,153 dropout=0.1):154 super(T5SelfAttention, self).__init__()155 self.dim = dim156 self.dim_attn = dim_attn157 self.dim_ffn = dim_ffn158 self.num_heads = num_heads159 self.num_buckets = num_buckets160 self.shared_pos = shared_pos161 162 # layers163 self.norm1 = T5LayerNorm(dim)164 self.attn = T5Attention(dim, dim_attn, num_heads, dropout)165 self.norm2 = T5LayerNorm(dim)166 self.ffn = T5FeedForward(dim, dim_ffn, dropout)167 self.pos_embedding = None if shared_pos else T5RelativeEmbedding(168 num_buckets, num_heads, bidirectional=True)169 170 def forward(self, x, mask=None, pos_bias=None):171 e = pos_bias if self.shared_pos else self.pos_embedding(172 x.size(1), x.size(1))173 x = fp16_clamp(x + self.attn(self.norm1(x), mask=mask, pos_bias=e))174 x = fp16_clamp(x + self.ffn(self.norm2(x)))175 return x176 177 178class T5CrossAttention(nn.Module):179 180 def __init__(self,181 dim,182 dim_attn,183 dim_ffn,184 num_heads,185 num_buckets,186 shared_pos=True,187 dropout=0.1):188 super(T5CrossAttention, self).__init__()189 self.dim = dim190 self.dim_attn = dim_attn191 self.dim_ffn = dim_ffn192 self.num_heads = num_heads193 self.num_buckets = num_buckets194 self.shared_pos = shared_pos195 196 # layers197 self.norm1 = T5LayerNorm(dim)198 self.self_attn = T5Attention(dim, dim_attn, num_heads, dropout)199 self.norm2 = T5LayerNorm(dim)200 self.cross_attn = T5Attention(dim, dim_attn, num_heads, dropout)201 self.norm3 = T5LayerNorm(dim)202 self.ffn = T5FeedForward(dim, dim_ffn, dropout)203 self.pos_embedding = None if shared_pos else T5RelativeEmbedding(204 num_buckets, num_heads, bidirectional=False)205 206 def forward(self,207 x,208 mask=None,209 encoder_states=None,210 encoder_mask=None,211 pos_bias=None):212 e = pos_bias if self.shared_pos else self.pos_embedding(213 x.size(1), x.size(1))214 x = fp16_clamp(x + self.self_attn(self.norm1(x), mask=mask, pos_bias=e))215 x = fp16_clamp(x + self.cross_attn(216 self.norm2(x), context=encoder_states, mask=encoder_mask))217 x = fp16_clamp(x + self.ffn(self.norm3(x)))218 return x219 220 221class T5RelativeEmbedding(nn.Module):222 223 def __init__(self, num_buckets, num_heads, bidirectional, max_dist=128):224 super(T5RelativeEmbedding, self).__init__()225 self.num_buckets = num_buckets226 self.num_heads = num_heads227 self.bidirectional = bidirectional228 self.max_dist = max_dist229 230 # layers231 self.embedding = nn.Embedding(num_buckets, num_heads)232 233 def forward(self, lq, lk):234 device = self.embedding.weight.device235 # rel_pos = torch.arange(lk).unsqueeze(0).to(device) - \236 # torch.arange(lq).unsqueeze(1).to(device)237 rel_pos = torch.arange(lk, device=device).unsqueeze(0) - \238 torch.arange(lq, device=device).unsqueeze(1)239 rel_pos = self._relative_position_bucket(rel_pos)240 rel_pos_embeds = self.embedding(rel_pos)241 rel_pos_embeds = rel_pos_embeds.permute(2, 0, 1).unsqueeze(242 0) # [1, N, Lq, Lk]243 return rel_pos_embeds.contiguous()244 245 def _relative_position_bucket(self, rel_pos):246 # preprocess247 if self.bidirectional:248 num_buckets = self.num_buckets // 2249 rel_buckets = (rel_pos > 0).long() * num_buckets250 rel_pos = torch.abs(rel_pos)251 else:252 num_buckets = self.num_buckets253 rel_buckets = 0254 rel_pos = -torch.min(rel_pos, torch.zeros_like(rel_pos))255 256 # embeddings for small and large positions257 max_exact = num_buckets // 2258 rel_pos_large = max_exact + (torch.log(rel_pos.float() / max_exact) /259 math.log(self.max_dist / max_exact) *260 (num_buckets - max_exact)).long()261 rel_pos_large = torch.min(262 rel_pos_large, torch.full_like(rel_pos_large, num_buckets - 1))263 rel_buckets += torch.where(rel_pos < max_exact, rel_pos, rel_pos_large)264 return rel_buckets265 266 267class T5Encoder(nn.Module):268 269 def __init__(self,270 vocab,271 dim,272 dim_attn,273 dim_ffn,274 num_heads,275 num_layers,276 num_buckets,277 shared_pos=True,278 dropout=0.1):279 super(T5Encoder, self).__init__()280 self.dim = dim281 self.dim_attn = dim_attn282 self.dim_ffn = dim_ffn283 self.num_heads = num_heads284 self.num_layers = num_layers285 self.num_buckets = num_buckets286 self.shared_pos = shared_pos287 288 # layers289 self.token_embedding = vocab if isinstance(vocab, nn.Embedding) \290 else nn.Embedding(vocab, dim)291 self.pos_embedding = T5RelativeEmbedding(292 num_buckets, num_heads, bidirectional=True) if shared_pos else None293 self.dropout = nn.Dropout(dropout)294 self.blocks = nn.ModuleList([295 T5SelfAttention(dim, dim_attn, dim_ffn, num_heads, num_buckets,296 shared_pos, dropout) for _ in range(num_layers)297 ])298 self.norm = T5LayerNorm(dim)299 300 # initialize weights301 self.apply(init_weights)302 303 def forward(self, ids, mask=None):304 x = self.token_embedding(ids)305 x = self.dropout(x)306 e = self.pos_embedding(x.size(1),307 x.size(1)) if self.shared_pos else None308 for block in self.blocks:309 x = block(x, mask, pos_bias=e)310 x = self.norm(x)311 x = self.dropout(x)312 return x313 314 315class T5Decoder(nn.Module):316 317 def __init__(self,318 vocab,319 dim,320 dim_attn,321 dim_ffn,322 num_heads,323 num_layers,324 num_buckets,325 shared_pos=True,326 dropout=0.1):327 super(T5Decoder, self).__init__()328 self.dim = dim329 self.dim_attn = dim_attn330 self.dim_ffn = dim_ffn331 self.num_heads = num_heads332 self.num_layers = num_layers333 self.num_buckets = num_buckets334 self.shared_pos = shared_pos335 336 # layers337 self.token_embedding = vocab if isinstance(vocab, nn.Embedding) \338 else nn.Embedding(vocab, dim)339 self.pos_embedding = T5RelativeEmbedding(340 num_buckets, num_heads, bidirectional=False) if shared_pos else None341 self.dropout = nn.Dropout(dropout)342 self.blocks = nn.ModuleList([343 T5CrossAttention(dim, dim_attn, dim_ffn, num_heads, num_buckets,344 shared_pos, dropout) for _ in range(num_layers)345 ])346 self.norm = T5LayerNorm(dim)347 348 # initialize weights349 self.apply(init_weights)350 351 def forward(self, ids, mask=None, encoder_states=None, encoder_mask=None):352 b, s = ids.size()353 354 # causal mask355 if mask is None:356 mask = torch.tril(torch.ones(1, s, s).to(ids.device))357 elif mask.ndim == 2:358 mask = torch.tril(mask.unsqueeze(1).expand(-1, s, -1))359 360 # layers361 x = self.token_embedding(ids)362 x = self.dropout(x)363 e = self.pos_embedding(x.size(1),364 x.size(1)) if self.shared_pos else None365 for block in self.blocks:366 x = block(x, mask, encoder_states, encoder_mask, pos_bias=e)367 x = self.norm(x)368 x = self.dropout(x)369 return x370 371 372class T5Model(nn.Module):373 374 def __init__(self,375 vocab_size,376 dim,377 dim_attn,378 dim_ffn,379 num_heads,380 encoder_layers,381 decoder_layers,382 num_buckets,383 shared_pos=True,384 dropout=0.1):385 super(T5Model, self).__init__()386 self.vocab_size = vocab_size387 self.dim = dim388 self.dim_attn = dim_attn389 self.dim_ffn = dim_ffn390 self.num_heads = num_heads391 self.encoder_layers = encoder_layers392 self.decoder_layers = decoder_layers393 self.num_buckets = num_buckets394 395 # layers396 self.token_embedding = nn.Embedding(vocab_size, dim)397 self.encoder = T5Encoder(self.token_embedding, dim, dim_attn, dim_ffn,398 num_heads, encoder_layers, num_buckets,399 shared_pos, dropout)400 self.decoder = T5Decoder(self.token_embedding, dim, dim_attn, dim_ffn,401 num_heads, decoder_layers, num_buckets,402 shared_pos, dropout)403 self.head = nn.Linear(dim, vocab_size, bias=False)404 405 # initialize weights406 self.apply(init_weights)407 408 def forward(self, encoder_ids, encoder_mask, decoder_ids, decoder_mask):409 x = self.encoder(encoder_ids, encoder_mask)410 x = self.decoder(decoder_ids, decoder_mask, x, encoder_mask)411 x = self.head(x)412 return x413 414 415def _t5(name,416 encoder_only=False,417 decoder_only=False,418 return_tokenizer=False,419 tokenizer_kwargs={},420 dtype=torch.float32,421 device='cpu',422 **kwargs):423 # sanity check424 assert not (encoder_only and decoder_only)425 426 # params427 if encoder_only:428 model_cls = T5Encoder429 kwargs['vocab'] = kwargs.pop('vocab_size')430 kwargs['num_layers'] = kwargs.pop('encoder_layers')431 _ = kwargs.pop('decoder_layers')432 elif decoder_only:433 model_cls = T5Decoder434 kwargs['vocab'] = kwargs.pop('vocab_size')435 kwargs['num_layers'] = kwargs.pop('decoder_layers')436 _ = kwargs.pop('encoder_layers')437 else:438 model_cls = T5Model439 440 # init model441 with torch.device(device):442 model = model_cls(**kwargs)443 444 # set device445 model = model.to(dtype=dtype, device=device)446 447 # init tokenizer448 if return_tokenizer:449 from .tokenizers import HuggingfaceTokenizer450 tokenizer = HuggingfaceTokenizer(f'google/{name}', **tokenizer_kwargs)451 return model, tokenizer452 else:453 return model454 455 456def umt5_xxl(**kwargs):457 cfg = dict(458 vocab_size=256384,459 dim=4096,460 dim_attn=4096,461 dim_ffn=10240,462 num_heads=64,463 encoder_layers=24,464 decoder_layers=24,465 num_buckets=32,466 shared_pos=False,467 dropout=0.1)468 cfg.update(**kwargs)469 return _t5('umt5-xxl', **cfg)470 471 472class T5EncoderModel:473 474 def __init__(475 self,476 text_len,477 dtype=torch.bfloat16,478 device="cuda",479 checkpoint_path=None,480 tokenizer_path=None,481 shard_fn=None,482 ):483 self.text_len = text_len484 self.dtype = dtype485 self.device = device486 self.checkpoint_path = checkpoint_path487 self.tokenizer_path = tokenizer_path488 489 # init model490 model = umt5_xxl(491 encoder_only=True,492 return_tokenizer=False,493 dtype=dtype,494 device=device).eval().requires_grad_(False)495 logging.info(f'loading {checkpoint_path}')496 model.load_state_dict(torch.load(checkpoint_path, map_location='cpu'))497 self.model = model498 if shard_fn is not None:499 self.model = shard_fn(self.model, sync_module_states=False)500 else:501 self.model.to(self.device)502 # init tokenizer503 self.tokenizer = HuggingfaceTokenizer(504 name=tokenizer_path, seq_len=text_len, clean='whitespace')505 506 def __call__(self, texts, device):507 ids, mask = self.tokenizer(508 texts, return_mask=True, add_special_tokens=True)509 ids = ids.to(device)510 mask = mask.to(device)511 seq_lens = mask.gt(0).sum(dim=1).long()512 context = self.model(ids, mask)513 return [u[:v] for u, v in zip(context, seq_lens)]514