import math from dataclasses import dataclass from typing import Optional import torch import torch.nn as nn import torch.nn.functional as F from fairscale.nn import checkpoint_wrapper from fairseq.model_parallel.megatron.mpu import ( ColumnParallelLinear, copy_to_model_parallel_region, VocabParallelEmbedding ) from .gate_retention import GateRetention from .sliding_window_attention import SlidingWindowAttention from .cross_attention import CrossAttention from .feedforward_network import FeedForwardNetwork, init_method from .rms_norm import RMSNorm from .kernel.rotary import apply_rotary_emb from .model_parallel_init import vocab_init_method, init_method @dataclass class YOCOArgs: dim: int n_layers: int hidden_dim: int n_self_heads: int n_attn_heads: int n_attn_kv_heads: int vocab_size: int max_batch_size: int = 0 max_seq_len: int = -1 model_parallel_size: int = 1 load_checkpoint: bool = False rope_theta: float = 10000.0 norm_eps: float = 1e-5 sliding_window: Optional[int] = None class DecoderLayer(nn.Module): def __init__( self, args: YOCOArgs, is_cross_layer=False ): super().__init__() self.args = args self.is_cross_layer = is_cross_layer if is_cross_layer: self.mixer = CrossAttention(args) elif args.sliding_window is not None: self.mixer = SlidingWindowAttention(args) else: self.mixer = GateRetention(args) self.mixer_layer_norm = RMSNorm(args.dim, eps=args.norm_eps) self.ffn = FeedForwardNetwork( args.dim, args.hidden_dim, args.load_checkpoint ) self.final_layer_norm = RMSNorm(args.dim, eps=args.norm_eps) def forward( self, x, start_pos=0, key=None, value=None, rel_pos=None, incremental_state=None, is_prefilling=False, ): residual = x x = self.mixer_layer_norm(x) if self.is_cross_layer: x = self.mixer( x, key, value, rel_pos=rel_pos, ) elif self.args.sliding_window is not None: x = self.mixer( x, rel_pos=rel_pos, start_pos=start_pos, incremental_state=incremental_state, ) else: x = self.mixer( x, rel_pos=rel_pos, incremental_state=incremental_state, is_prefilling=is_prefilling,) x = x + residual residual = x x = self.final_layer_norm(x) x = self.ffn(x) x = x + residual return x class SelfDecoder(nn.Module): def __init__( self, args: YOCOArgs, checkpoint_activations: bool = False ): super().__init__() self.args = args layers = [DecoderLayer(args, is_cross_layer=False,) for idx in range(args.n_layers // 2)] if checkpoint_activations: layers = [checkpoint_wrapper(layer) for layer in layers] self.layers = nn.ModuleList(layers) self.head_dim = args.dim // args.n_self_heads self.block_size = 256 self._precomputed_freqs_cis = None def build_rel_pos(self, x, start_pos): if self._precomputed_freqs_cis is None: angle = 1.0 / (self.args.rope_theta ** torch.linspace(0, 1, self.head_dim // 2, dtype=torch.float, device=x.device)) index = torch.arange(self.args.max_seq_len).to(angle) self._precomputed_freqs_cis = index[:, None] * angle cos = torch.cos(self._precomputed_freqs_cis[start_pos:start_pos+x.size(1)]) sin = torch.sin(self._precomputed_freqs_cis[start_pos:start_pos+x.size(1)]) rel_pos = (cos.to(x.dtype), sin.to(x.dtype)) return rel_pos def get_index_mask(self, x, length, pad_length): return torch.arange(pad_length, device=x.device) >= length def forward( self, x, incremental_state=None, is_prefilling=False, start_pos=0 ): if is_prefilling and x.size(1) % self.block_size != 0 and self.args.sliding_window is None: padding_len = self.block_size - x.size(1) % self.block_size x = F.pad(x, (0, 0, 0, padding_len), value=0) else: padding_len = 0 if incremental_state is not None and is_prefilling: index_mask = self.get_index_mask(x, x.size(1) - padding_len, x.size(1)) rel_pos = self.build_rel_pos(x, start_pos) for idx, layer in enumerate(self.layers): if incremental_state is not None: if idx not in incremental_state: incremental_state[idx] = {} if is_prefilling: incremental_state[idx]["index_mask"] = index_mask x = layer( x, start_pos=start_pos, rel_pos=rel_pos, incremental_state=incremental_state[idx] if incremental_state is not None else None, is_prefilling=is_prefilling,) x = x[:, :x.size(1) - padding_len, :] return x class CrossDecoder(nn.Module): def __init__( self, args: YOCOArgs, checkpoint_activations: bool = False ): super().__init__() self.args = args self.num_heads = args.n_attn_kv_heads self.head_dim = args.dim // args.n_attn_heads self.k_proj = ColumnParallelLinear(args.dim, self.head_dim * args.n_attn_kv_heads, bias=False, gather_output=False, init_method=init_method) self.v_proj = ColumnParallelLinear(args.dim, self.head_dim * args.n_attn_kv_heads, bias=False, gather_output=False, init_method=init_method) self.kv_layer_norm = RMSNorm(args.dim, eps=args.norm_eps) layers = [DecoderLayer(args, is_cross_layer=True) for idx in range(args.n_layers // 2)] if checkpoint_activations: layers = [checkpoint_wrapper(layer) for layer in layers] self.layers = nn.ModuleList(layers) self._precomputed_freqs_cis = None def build_rel_pos(self, x, start_pos): if self._precomputed_freqs_cis is None: angle = 1.0 / (self.args.rope_theta ** torch.linspace(0, 1, self.head_dim // 2, dtype=torch.float, device=x.device)) index = torch.arange(self.args.max_seq_len).to(angle) self._precomputed_freqs_cis = index[:, None] * angle cos = torch.cos(self._precomputed_freqs_cis[start_pos:start_pos+x.size(1)]) sin = torch.sin(self._precomputed_freqs_cis[start_pos:start_pos+x.size(1)]) rel_pos = (cos.to(x.dtype), sin.to(x.dtype)) return rel_pos def forward( self, x, incremental_state=None, start_pos=0, skip_cross_decoder=False, ): bsz, seqlen, embed_dim = x.size() x_norm = self.kv_layer_norm(x) key, value = self.k_proj(x_norm), self.v_proj(x_norm) key = key.view(bsz, seqlen, self.num_heads, self.head_dim) value = value.view(bsz, seqlen, self.num_heads, self.head_dim) rel_pos = self.build_rel_pos(x, start_pos) key = apply_rotary_emb(key, *rel_pos, interleaved=True) if incremental_state is not None: if "prev_key" not in incremental_state: incremental_state["prev_key"] = torch.empty(bsz, self.args.max_seq_len, self.num_heads, self.head_dim, device=x.device, dtype=x.dtype) incremental_state["prev_value"] = torch.empty(bsz, self.args.max_seq_len, self.num_heads, self.head_dim, device=x.device, dtype=x.dtype) incremental_state["prev_key"][:, start_pos : start_pos + seqlen] = key incremental_state["prev_value"][:, start_pos : start_pos + seqlen] = value key = incremental_state["prev_key"][:, : start_pos + seqlen] value = incremental_state["prev_value"][:, : start_pos + seqlen] if skip_cross_decoder: return torch.zeros(bsz, 1, embed_dim, device=x.device, dtype=x.dtype) for layer in self.layers: x = layer( x, key=key, value=value, rel_pos=rel_pos) return x class YOCO(nn.Module): def __init__( self, args: YOCOArgs, checkpoint_activations: bool = False, share_input_output_embed: bool = False, ): super().__init__() self.args = args self.embed_scale = math.sqrt(args.dim) self.embed_tokens = VocabParallelEmbedding( args.vocab_size, args.dim, -1, init_method=vocab_init_method ) self.output_projection = nn.Linear(args.dim, args.vocab_size, bias=False) if share_input_output_embed: self.output_projection.weight = self.embed_tokens.weight self.self_decoder = SelfDecoder(args, checkpoint_activations) self.cross_decoder = CrossDecoder(args, checkpoint_activations) self.layer_norm = RMSNorm(args.dim, eps=args.norm_eps) def forward( self, x, start_pos=0, incremental_state=None, is_prefilling=True, skip_cross_decoder=False ): x = self.embed_scale * self.embed_tokens(x) x = self.self_decoder( x, incremental_state=incremental_state, is_prefilling=is_prefilling, start_pos=start_pos, ) x = self.cross_decoder( x, start_pos=start_pos, incremental_state=incremental_state, skip_cross_decoder=skip_cross_decoder, ) x = self.layer_norm(x) x = self.output_layer(x) return x, None def output_layer(self, features): if self.args.model_parallel_size > 1: features = copy_to_model_parallel_region(features) return self.output_projection(features)