from typing import List, Optional, Tuple, Union import math import torch import torch.nn.functional as F import torch.utils.checkpoint from torch import nn from torch.nn import CrossEntropyLoss from transformers.activations import ACT2FN from transformers.cache_utils import Cache from transformers.modeling_outputs import ( BaseModelOutputWithPast, CausalLMOutputWithPast, ) from transformers.modeling_utils import PreTrainedModel from .configuration_evabyte import EvaByteConfig from .multibyte_decoding_evabyte import MultiByteDecodingMixin try: import triton USE_TRITON_IMPL = True from .eva import EvaAttention from .eva_agg_kernel import triton_eva_agg_fwd from .eva_prep_kv_kernel import triton_eva_prep_kv_fwd except ImportError: USE_TRITON_IMPL = False print("WARNING: triton is not installed, using fallback EVA which might be slow and throw errors") from .eva_pt_ref import EvaAttention from .eva_cache import EvaCache, EvaStaticCacheForTriton MASK_MIN_VALUE = -10e10 def prepare_eva_attention_mask( seq_len, device, chunk_size, window_size, use_cache=False, cache=None ): """ Prepare attention masks for EVA. """ chunk_causal_mask = None window_causal_mask = None if use_cache: cached_seq_len = cache.get_seq_length() total_seq_len = seq_len + cached_seq_len # cached_seq_len will be 0 during prefilling # padded_seq_len = chunk_size * math.ceil(total_seq_len / chunk_size) padded_seq_len = window_size * math.ceil(total_seq_len / window_size) num_chunks = padded_seq_len // chunk_size else: # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len] assert seq_len % chunk_size == 0 num_chunks = seq_len // chunk_size assert seq_len % window_size == 0 # create causal mask ################################ # generate chunked causal masks ################################ # [b, h, j, c, c] chunks_per_window = window_size // chunk_size if num_chunks >= chunks_per_window: chunk_causal_mask = torch.ones( (chunk_size, num_chunks, num_chunks), device=device, dtype=torch.bool ).triu(0) num_blocks = num_chunks // chunks_per_window chunk_causal_mask = chunk_causal_mask.reshape( chunk_size, num_blocks, chunks_per_window, num_blocks, chunks_per_window ).transpose(-2, -3) block_diag_zero = ( torch.eye(num_blocks, device=device, dtype=torch.bool) .unsqueeze(-1) .unsqueeze(-1) .unsqueeze(0) ) # Set diagonal blocks to zero chunk_causal_mask = chunk_causal_mask.masked_fill(block_diag_zero, True) # Reshape back to original size chunk_causal_mask = ( chunk_causal_mask .transpose(-2, -3) .reshape(chunk_size, num_chunks, num_chunks) .transpose(-2, -3) .reshape(chunk_size * num_chunks, num_chunks) .unsqueeze(0) .unsqueeze(0) ) else: chunk_causal_mask = torch.ones( (1, 1, chunk_size, num_chunks, num_chunks), device=device, dtype=torch.bool, ).triu(0).transpose(-2, -3) # [1, 1, c, j, c] chunk_causal_mask = chunk_causal_mask.reshape( 1, 1, chunk_size * num_chunks, num_chunks ) # [1, 1, n, c] if use_cache: chunk_causal_mask = chunk_causal_mask[..., cached_seq_len : cached_seq_len + seq_len, :] window_causal_mask = torch.ones( (1, 1, 1, window_size, window_size), device=device ).triu(1).to(torch.bool) return (chunk_causal_mask, window_causal_mask) def pad_to_multiple(tensor, multiple, dim=-2, value=0, create_mask=False, left_padding=False): assert dim < 0 # only accept ``dim'' index in a reverse manner seqlen = int(tensor.shape[dim]) m = seqlen / multiple if m.is_integer(): if create_mask: return tensor, torch.ones(size=(tensor.shape[0], tensor.shape[dim]), dtype=torch.bool, device=tensor.device) else: return tensor remainder = math.ceil(m) * multiple - seqlen pad_offset = (0,) * (-1 - dim) * 2 if left_padding: padded_res = F.pad(tensor, (*pad_offset, remainder, 0), value=value) else: padded_res = F.pad(tensor, (*pad_offset, 0, remainder), value=value) if create_mask: # assume dim 0 is the batch size padding_mask = torch.ones(size=(padded_res.shape[0], padded_res.shape[dim]), dtype=torch.bool, device=padded_res.device) if left_padding: padding_mask[:, :remainder] = False else: padding_mask[:, -remainder:] = False return padded_res, padding_mask else: return padded_res class EvaByteRMSNorm(nn.Module): def __init__(self, config): super().__init__() self.config = config self.fp32_ln = config.fp32_ln self.variance_epsilon = config.rms_norm_eps self.add_unit_offset = config.norm_add_unit_offset if self.add_unit_offset: self.weight = nn.Parameter(torch.zeros(config.hidden_size)) else: self.weight = nn.Parameter(torch.ones(config.hidden_size)) def forward(self, hidden_states): if hasattr(self, 'config'): fp32_ln = self.config.fp32_ln else: fp32_ln = self.fp32_ln hidden_states = hidden_states.to(torch.float32 if fp32_ln else torch.bfloat16) variance = hidden_states.pow(2).mean(-1, keepdim=True) hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) if self.add_unit_offset: return (1 + self.weight) * hidden_states else: return self.weight * hidden_states class EvaByteRotaryEmbedding(torch.nn.Module): def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None): super().__init__() self.dim = dim self.max_position_embeddings = max_position_embeddings self.base = base inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64).float().to(device) / self.dim)) self.register_buffer("inv_freq", inv_freq, persistent=False) self._set_cos_sin_cache(seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype()) def _set_cos_sin_cache(self, seq_len, device, dtype): self.max_seq_len_cached = seq_len t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype) freqs = torch.einsum("i,j->ij", t, self.inv_freq) emb = torch.cat((freqs, freqs), dim=-1) self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False) self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False) def forward(self, x, seq_len=None): # x: [bs, num_attention_heads, seq_len, head_size] if seq_len > self.max_seq_len_cached: self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype) # return ( # self.cos_cached[:seq_len].to(dtype=x.dtype), # self.sin_cached[:seq_len].to(dtype=x.dtype), # ) if seq_len < self.max_seq_len_cached: cos_slice = self.cos_cached.split(seq_len, dim=0)[0] sin_slice = self.sin_cached.split(seq_len, dim=0)[0] else: cos_slice = self.cos_cached sin_slice = self.sin_cached return ( cos_slice.to(dtype=x.dtype), sin_slice.to(dtype=x.dtype), ) class EvaByteLinearScalingRotaryEmbedding(EvaByteRotaryEmbedding): """EvaByteRotaryEmbedding extended with linear scaling. Credits to the Reddit user /u/kaiokendev""" def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0): self.scaling_factor = scaling_factor super().__init__(dim, max_position_embeddings, base, device) def _set_cos_sin_cache(self, seq_len, device, dtype): self.max_seq_len_cached = seq_len t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype) t = t / self.scaling_factor freqs = torch.einsum("i,j->ij", t, self.inv_freq) # Different from paper, but it uses a different permutation in order to obtain the same calculation emb = torch.cat((freqs, freqs), dim=-1) self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False) self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False) class EvaByteDynamicNTKScalingRotaryEmbedding(EvaByteRotaryEmbedding): """EvaByteRotaryEmbedding extended with Dynamic NTK scaling. Credits to the Reddit users /u/bloc97 and /u/emozilla""" def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0): self.scaling_factor = scaling_factor super().__init__(dim, max_position_embeddings, base, device) def _set_cos_sin_cache(self, seq_len, device, dtype): self.max_seq_len_cached = seq_len if seq_len > self.max_position_embeddings: base = self.base * ((self.scaling_factor * seq_len / self.max_position_embeddings) - (self.scaling_factor - 1))**(self.dim / (self.dim - 2)) inv_freq = 1.0 / (base**(torch.arange(0, self.dim, 2).float().to(device) / self.dim)) self.register_buffer("inv_freq", inv_freq, persistent=False) t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype) freqs = torch.einsum("i,j->ij", t, self.inv_freq) # Different from paper, but it uses a different permutation in order to obtain the same calculation emb = torch.cat((freqs, freqs), dim=-1) self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False) self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False) class EvaByteMLP(nn.Module): def __init__(self, config, layer_idx: int = None): super().__init__() self.hidden_size = config.hidden_size self.intermediate_size = config.intermediate_size self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) self.act_fn = ACT2FN[config.hidden_act] self.layer_idx = layer_idx self.config = config def forward(self, x): down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) return down_proj class EvaByteDecoderLayer(nn.Module): def __init__(self, config: EvaByteConfig, layer_idx: int = None): super().__init__() self.config = config self.hidden_size = config.hidden_size self.self_attn = EvaAttention(config=config, layer_idx=layer_idx) self.mlp = EvaByteMLP(config, layer_idx=layer_idx) self.input_layernorm = EvaByteRMSNorm(config) self.post_attention_layernorm = EvaByteRMSNorm(config) def forward( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_value: Optional[Tuple[torch.Tensor]] = None, output_attentions: Optional[bool] = False, use_cache: Optional[bool] = False, cos: Optional[torch.Tensor] = None, sin: Optional[torch.Tensor] = None, multibyte_decoding: Optional[bool] = False, ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: residual = hidden_states if self.config.fp32_skip_add: residual = residual.float() hidden_states = self.input_layernorm(hidden_states) # Self Attention hidden_states, self_attn_weights, present_key_value = self.self_attn(hidden_states=hidden_states, attention_mask=attention_mask, position_ids=position_ids, past_key_value=past_key_value, output_attentions=output_attentions, use_cache=use_cache, cos=cos, sin=sin, multibyte_decoding=multibyte_decoding) hidden_states = residual + hidden_states # Fully Connected residual = hidden_states if self.config.fp32_skip_add: residual = residual.float() hidden_states = self.post_attention_layernorm(hidden_states) hidden_states = self.mlp(hidden_states) hidden_states = residual + hidden_states outputs = (hidden_states, ) if output_attentions: outputs += (self_attn_weights, ) if use_cache: outputs += (present_key_value, ) return outputs class EvaBytePreTrainedModel(PreTrainedModel): config_class = EvaByteConfig base_model_prefix = "model" supports_gradient_checkpointing = True _no_split_modules = ["EvaByteDecoderLayer"] _skip_keys_device_placement = "past_key_values" def _init_weights(self, module): std = getattr(self.config, "initializer_range", 0.02) if isinstance(module, nn.Linear): module.weight.data.normal_(mean=0.0, std=std) if module.bias is not None: module.bias.data.zero_() elif isinstance(module, nn.Embedding): module.weight.data.normal_(mean=0.0, std=std) if module.padding_idx is not None: module.weight.data[module.padding_idx].zero_() def _set_gradient_checkpointing(self, module, value=False): if isinstance(module, EvaByteModel): module.gradient_checkpointing = value class EvaByteModel(EvaBytePreTrainedModel): """ Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`EvaByteDecoderLayer`] Args: config: EvaByteConfig """ def __init__(self, config: EvaByteConfig): super().__init__(config) self.padding_idx = config.pad_token_id self.vocab_size = config.vocab_size self.hidden_size = config.hidden_size self.num_heads = config.num_attention_heads self.head_dim = self.hidden_size // self.num_heads self.max_position_embeddings = self.config.max_position_embeddings self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) self.layers = nn.ModuleList([EvaByteDecoderLayer(config, layer_idx=layer_idx) for layer_idx in range(config.num_hidden_layers)]) self.norm = EvaByteRMSNorm(config) self.gradient_checkpointing = False self.rope = config.rope_theta # Initialize weights and apply final processing self.post_init() self._init_rope() def _init_rope(self): if self.config.rope_scaling is None: self.rotary_emb = EvaByteRotaryEmbedding(self.head_dim, max_position_embeddings=self.max_position_embeddings, base=self.rope) else: scaling_type = self.config.rope_scaling["type"] scaling_factor = self.config.rope_scaling["factor"] if scaling_type == "linear": self.rotary_emb = EvaByteLinearScalingRotaryEmbedding( self.head_dim, max_position_embeddings=self.max_position_embeddings, scaling_factor=scaling_factor, base=self.rope) elif scaling_type == "dynamic": self.rotary_emb = EvaByteDynamicNTKScalingRotaryEmbedding( self.head_dim, max_position_embeddings=self.max_position_embeddings, scaling_factor=scaling_factor, base=self.rope) else: raise ValueError(f"Unknown RoPE scaling type {scaling_type}") def get_input_embeddings(self): return self.embed_tokens def set_input_embeddings(self, value): self.embed_tokens = value def _helper_padding_mask( self, padding_mask, causal_mask ): padding_mask = torch.logical_or(padding_mask, padding_mask.transpose(-1, -2)) return torch.logical_or(padding_mask, causal_mask) def _prepare_eva_generation_attn_mask_triton( self, attention_mask, input_ids, use_cache, past_key_values ): batch_size, seq_len = input_ids.shape if use_cache and past_key_values.get_seq_length() > 0: # decoding phase if past_key_values.rf_mask[0] is not None: cur_rf_mask = torch.zeros( (batch_size, 1, seq_len, 1), dtype=past_key_values.rf_mask[0].dtype, device=past_key_values.rf_mask[0].device ) else: cur_rf_mask = None if past_key_values.s_mask[0] is not None: cur_s_mask = torch.zeros( (batch_size, 1, seq_len, 1), dtype=past_key_values.s_mask[0].dtype, device=past_key_values.s_mask[0].device ) else: cur_s_mask = None seen_tokens = past_key_values.get_seq_length() if seen_tokens <= self.config.window_size: rfa_chunks_dummy_mask = None else: if cur_s_mask is not None: chunks_per_window = int(self.config.window_size // self.config.chunk_size) # the ongoing decoding step would be (seen_seq_len + 1)-th token num_windows_seen_so_far = seen_tokens // self.config.window_size rfa_chunks_dummy_mask = torch.zeros( (batch_size, 1, seq_len, num_windows_seen_so_far * chunks_per_window), dtype=past_key_values.s_mask[0].dtype, device=past_key_values.s_mask[0].device ) else: rfa_chunks_dummy_mask = None # rf_mask and cur_mask are 0s because we do not want to mask them return (cur_s_mask, cur_rf_mask, rfa_chunks_dummy_mask) if attention_mask is not None and torch.any(attention_mask == 0.0): # convert 0 -> padding to 1 -> padding padded_attention_mask = pad_to_multiple( attention_mask, self.config.window_size, dim=-1, value=0, create_mask=False, left_padding=False ) # convert 0 -> padding to 1 -> padding padded_rf_mask = ~padded_attention_mask.unsqueeze(1).unsqueeze(-1).to(torch.bool) # [b, 1, n, 1] # [b, 1, w, j, 1] padded_w_attn_mask = padded_rf_mask.reshape(batch_size, 1, -1, self.config.window_size, 1).to(torch.bool) # [b, 1, w, j, 1] [b, 1, w, 1, j] -> [b, 1, w, j, j] w_padding_mask = torch.logical_or(padded_w_attn_mask, padded_w_attn_mask.transpose(-1, -2)) w_causal_mask = torch.ones( (1, 1, 1, self.config.window_size, self.config.window_size), device=input_ids.device ).triu(1).to(torch.bool) s_mask = torch.logical_or(w_padding_mask, w_causal_mask) s_mask = s_mask.reshape(batch_size, 1, -1, self.config.window_size) s_mask = s_mask[..., :seq_len, :] # negate the attention mask to get the padding mask rf_mask = ~attention_mask.unsqueeze(1).unsqueeze(-1).to(torch.bool) # [b, 1, n, 1] return (s_mask, rf_mask) else: return (None, None) def _prepare_eva_generation_attn_mask( self, attention_mask, input_ids, use_cache, past_key_values ): batch_size, seq_len = input_ids.shape if use_cache and past_key_values.get_seq_length() > 0: # decoding phase if past_key_values.rf_mask[0] is not None: rf_mask = torch.zeros( (batch_size, 1, seq_len, 1), dtype=past_key_values.rf_mask[0].dtype, device=past_key_values.rf_mask[0].device ) else: rf_mask = None cur_causal_mask = torch.zeros( (batch_size, 1, seq_len, 1), dtype=torch.bool, device=input_ids.device ) chunk_causal_mask = torch.ones( (batch_size, 1, seq_len, 1), dtype=torch.bool, device=input_ids.device ) # chunk_causal_mask are 1s because we will mask them by default and # will be unmasked when the current singleton attention is processed over return (None, cur_causal_mask, chunk_causal_mask, rf_mask) true_num_chunks = seq_len // self.config.chunk_size chunk_causal_mask, _ = prepare_eva_attention_mask( seq_len, input_ids.device, self.config.chunk_size, self.config.window_size, use_cache=use_cache, cache=past_key_values ) chunk_causal_mask = chunk_causal_mask[..., :seq_len, :true_num_chunks] if attention_mask is not None and torch.any(attention_mask == 0.0): # convert 0 -> padding to 1 -> padding rf_mask = ~attention_mask.unsqueeze(1).unsqueeze(-1).to(torch.bool) # [b, 1, n, 1] else: rf_mask = None if seq_len < self.config.window_size: cur_window_mask = torch.ones( (1, 1, seq_len, seq_len), device=input_ids.device ).triu(1).to(torch.bool) if rf_mask is not None: cur_window_mask = self._helper_padding_mask(rf_mask, cur_window_mask) prev_window_mask = None else: if seq_len % self.config.window_size == 0: num_windows = seq_len // self.config.window_size cur_window_mask = None prev_window_mask = torch.ones( (1, 1, num_windows, self.config.window_size, self.config.window_size), device=input_ids.device ).triu(1).to(torch.bool) if rf_mask is not None: prev_rf_mask = rf_mask.reshape(batch_size, 1, -1, self.config.window_size, 1) prev_window_mask = self._helper_padding_mask(prev_rf_mask, prev_window_mask) else: num_windows = seq_len // self.config.window_size remainder_tokens = seq_len % self.config.window_size cur_window_mask = torch.ones( (1, 1, remainder_tokens, remainder_tokens), device=input_ids.device ).triu(1).to(torch.bool) prev_window_mask = torch.ones( (1, 1, num_windows, self.config.window_size, self.config.window_size), device=input_ids.device ).triu(1).to(torch.bool) if rf_mask is not None: prev_rf_mask, cur_rf_mask = torch.split(rf_mask, [seq_len - remainder_tokens, remainder_tokens], dim=-2) cur_window_mask = self._helper_padding_mask(cur_rf_mask, cur_window_mask) prev_rf_mask = prev_rf_mask.reshape(batch_size, 1, -1, self.config.window_size, 1) prev_window_mask = self._helper_padding_mask(prev_rf_mask, prev_window_mask) return (prev_window_mask, cur_window_mask, chunk_causal_mask, rf_mask) def forward( self, input_ids: torch.LongTensor = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[List[torch.FloatTensor]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, multibyte_decoding: Optional[bool] = None, ) -> Tuple: output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions output_hidden_states = (output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states) use_cache = use_cache if use_cache is not None else self.config.use_cache return_dict = return_dict if return_dict is not None else self.config.use_return_dict if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError( "You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one" ) if self.gradient_checkpointing and self.training and use_cache: raise ValueError("`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`...") batch_size, seq_len = input_ids.shape #### Step 0. Hack if (not self.training) and (not use_cache) and (not multibyte_decoding): # forward-only inference mode. # We tweak use_cache to be True to reuse code for generation use_cache = True device = input_ids.device if input_ids is not None else None if position_ids is None: position_ids = torch.arange(0, seq_len, device=device, dtype=int).reshape(1, -1).expand(batch_size, -1) #### Step 1. Prepare caches if in inference mode if use_cache: if past_key_values is not None: assert isinstance(past_key_values, Cache) else: if not USE_TRITON_IMPL: past_key_values = EvaCache() else: past_key_values = EvaStaticCacheForTriton( input_ids.shape[0], self.config.num_attention_heads, self.config.window_size, self.config.hidden_size // self.config.num_attention_heads, self.config.num_hidden_layers, self.embed_tokens.weight.dtype, self.embed_tokens.weight.device, ) if not multibyte_decoding: if use_cache: if USE_TRITON_IMPL: causal_mask = self._prepare_eva_generation_attn_mask_triton( attention_mask, input_ids, use_cache, past_key_values ) else: causal_mask = self._prepare_eva_generation_attn_mask( attention_mask, input_ids, use_cache, past_key_values ) else: assert self.training assert seq_len % self.config.window_size == 0 # for training, we need to pass in the attention mask # usually calculated by _prepare_training_attn_mask() causal_mask = attention_mask else: assert use_cache causal_mask = attention_mask if inputs_embeds is None: inputs_embeds = self.embed_tokens(input_ids) past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 max_seq_length = past_seen_tokens + inputs_embeds.shape[1] hidden_states = inputs_embeds if position_ids is None: assert not use_cache, "during decoding we must explicitly pass position_ids to the model call" device = input_ids.device if input_ids is not None else inputs_embeds.device position_ids = torch.arange(past_seen_tokens, max_seq_length, device=device, dtype=int).reshape(1, -1).expand(batch_size, -1) cos, sin = self.rotary_emb(hidden_states, seq_len=max_seq_length) assert len(cos.shape) == 2, f"cos should be of shape (max_seq_len, head_dim), got {cos.shape} instead" assert sin.shape == cos.shape, f"sin should be of shape (max_seq_len, head_dim), got {sin.shape} instead" assert len(position_ids.shape) == 2, f"position_ids should be of 2D, got {position_ids.shape} instead" cos = cos[position_ids, :] sin = sin[position_ids, :] cos = cos.unsqueeze(1) sin = sin.unsqueeze(1) if USE_TRITON_IMPL and (not multibyte_decoding): # the masks generated above for triton kernels are boolean. Convert them to floats if ( (not use_cache) or (use_cache and past_seen_tokens == 0) ): window_mask, intra_chunk_mask = causal_mask if window_mask is not None: assert window_mask.dtype == torch.bool window_mask_float = window_mask.to(torch.float) window_mask_float = window_mask_float.masked_fill(window_mask.to(torch.bool), MASK_MIN_VALUE) window_mask_float = window_mask_float.reshape(batch_size, 1, -1, self.config.window_size) window_mask = window_mask_float.to(hidden_states.dtype) if intra_chunk_mask is not None: assert intra_chunk_mask.dtype == torch.bool intra_chunk_mask_float = intra_chunk_mask.to(torch.float) intra_chunk_mask_float = intra_chunk_mask_float.masked_fill(intra_chunk_mask.to(torch.bool), MASK_MIN_VALUE) intra_chunk_mask = intra_chunk_mask_float.to(hidden_states.dtype) causal_mask = (window_mask, intra_chunk_mask) if self.config.fp32_skip_add: hidden_states = hidden_states.float() # decoder layers all_hidden_states = () if output_hidden_states else None all_self_attns = () if output_attentions else None next_decoder_cache = None for decoder_layer in self.layers: if output_hidden_states: all_hidden_states += (hidden_states, ) if self.gradient_checkpointing and self.training: def create_custom_forward(module): def custom_forward(*inputs): # None for past_key_value return module(*inputs, output_attentions, use_cache=None) return custom_forward layer_outputs = torch.utils.checkpoint.checkpoint( create_custom_forward(decoder_layer), hidden_states, causal_mask, position_ids, None, ) else: layer_outputs = decoder_layer( hidden_states, attention_mask=causal_mask, position_ids=position_ids, past_key_value=past_key_values, output_attentions=output_attentions, use_cache=use_cache, cos=cos, sin=sin, multibyte_decoding=multibyte_decoding, ) hidden_states = layer_outputs[0] if use_cache: next_decoder_cache = layer_outputs[2 if output_attentions else 1] if output_attentions: all_self_attns += (layer_outputs[1], ) hidden_states = self.norm(hidden_states) # add hidden states from the last decoder layer if output_hidden_states: all_hidden_states += (hidden_states, ) next_cache = next_decoder_cache if use_cache else None if not return_dict: return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None) return BaseModelOutputWithPast( last_hidden_state=hidden_states, past_key_values=next_cache, hidden_states=all_hidden_states, attentions=all_self_attns, ) class EvaByteForCausalLM(EvaBytePreTrainedModel, MultiByteDecodingMixin): _tied_weights_keys = ["lm_head.weight"] def __init__(self, config): EvaBytePreTrainedModel.__init__(self, config) self.model = EvaByteModel(config) self.vocab_size = config.vocab_size # define multibyte prediction heads if hasattr(config, "num_pred_heads") and config.num_pred_heads > 1: self.lm_head = nn.Linear(config.hidden_size, config.vocab_size * config.num_pred_heads, bias=False) else: self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) self.post_init() def get_input_embeddings(self): return self.model.embed_tokens def set_input_embeddings(self, value): self.model.embed_tokens = value def get_output_embeddings(self): return self.lm_head def set_output_embeddings(self, new_embeddings): self.lm_head = new_embeddings def set_decoder(self, decoder): self.model = decoder def get_decoder(self): return self.model def _prepare_training_attn_mask( self, target_token_type_ids, use_doc_boundary_attention, EOS_TOKEN_TYPE_ID=None, PAD_TOKEN_TYPE_ID=None, ): ''' This function prepares the attention mask for training byte models. target_token_type_ids: Tensor of shape (batch_size, seq_len), marking the token type ids for the target sequence. In particular, we should have - target_token_type_ids[i, j] = EOS_TOKEN_TYPE_ID if the j-th token in the i-th sequence is the end of an article. - target_token_type_ids[i, j] = PAD_TOKEN_TYPE_ID if the j-th token in the i-th sequence is the padding token. use_doc_boundary_attention: bool, whether to enable doc boundary attention. EOS_TOKEN_TYPE_ID: int, the token type id for the end of an article. PAD_TOKEN_TYPE_ID: int, the token type id for the padding token. ''' assert self.training batch_size, num_tokens = target_token_type_ids.shape chunk_causal_mask, window_causal_mask = prepare_eva_attention_mask( num_tokens, target_token_type_ids.device, chunk_size=self.config.chunk_size, window_size=self.config.window_size, use_cache=False, cache=None ) if use_doc_boundary_attention: #### step 1: mark each document with a unique id end_token_ids = {EOS_TOKEN_TYPE_ID, PAD_TOKEN_TYPE_ID} token_types = torch.zeros(batch_size, num_tokens) for sequence_idx, sequence in enumerate(target_token_type_ids): num_articles = 0 start_index = 0 # for each sample in the batch, the collapsed attention mask looks like: # [1, 1, .... 1, 0, 2, 2, ... 2, 0, ... n, n ..... n], assuming there are n articles in the sequence. # Each of the n articles are separated by 0. for token_idx, token_type_id in enumerate(sequence): if start_index is not None and token_type_id.item() in end_token_ids: num_articles += 1 end_index = token_idx if token_type_id == PAD_TOKEN_TYPE_ID else token_idx + 1 token_types[sequence_idx][start_index:end_index] = num_articles start_index = None elif start_index is None and token_type_id not in end_token_ids: start_index = token_idx + 1 assert num_tokens % self.config.chunk_size == 0, "Number of tokens must be divisible by chunk size" assert num_tokens % self.config.window_size == 0, "Number of tokens must be divisible by window size" num_chunks = num_tokens // self.config.chunk_size num_windows = num_tokens // self.config.window_size article_separator = 0 #### step 2: generate attention masks for each window #### NOTE: we perform exact attention within each window, #### so we only need to mask out different documents #### for each window. token_types_windows = token_types.reshape(batch_size, num_windows, self.config.window_size, 1) token_types_windows_t = token_types_windows.transpose(-1, -2) # replace all elements in TOKEN_SEPS with -1 token_types_windows = torch.where(token_types_windows == article_separator, -1, token_types_windows) window_3d_mask = (token_types_windows == token_types_windows_t) window_3d_mask = ~window_3d_mask #### step 3: generate chunk-level 3D masks #### NOTE: this is a bit tricky, as we aim to mask out different #### documents to avoid cross-doc attention across chunks. #### Example: suppose we have a sequence of length 12 with 3 documents: #### [1, 1, 1, 1, 1, 2, 2, 3, 3, 3, 3, 3]. #### The chunk-size and window-size are both 4. #### The chunk-level mask of shape (batch_size, seq_len, num_chunks) is: #### [ #### [0, 0, 0], #### [0, 0, 0], #### [0, 0, 0], #### [0, 0, 0], #### #### [1, 0, 0], #### [0, 0, 0], #### [0, 0, 0], #### [0, 0, 0], #### #### [0, 1, 0], #### [0, 1, 0], #### [0, 1, 0], #### [0, 1, 0], #### ] #### Explanation: #### - Tokens will not attend to their own and future chunks. #### (as tokens within a chunk are captured by the window-level exact attention) #### - Tokens will attend to a chunk only if there are tokens #### from the same document in that chunk. #### The mask within each chunk of shape (batch_size, num_chunks, chunk_size) is: #### [ #### [1, 1, 1, 1], #### [0, 0, 0, 1], #### [1, 1, 1, 1], #### ] #### Explanation: #### - If all tokens in a chunk are from the same document, #### no tokens will be masked out. #### - If there are tokens from different documents in a chunk, #### only tokens from the rightmost document will be kept. #### (b/c the future chunks might contain tokens from the rightmost document, #### but all the remaining docs will never get attended by other docs) token_types_chunks = token_types.reshape(batch_size, num_chunks, self.config.chunk_size) inter_chunk_mask = torch.zeros((batch_size, num_tokens, num_chunks), dtype=torch.bool) intra_chunk_mask = torch.ones_like(token_types_chunks, dtype=torch.bool) for chunk_idx in range(num_chunks): for batch_idx in range(batch_size): # Identify tokens in the current chunk belonging to each sequence chunk = token_types_chunks[batch_idx, chunk_idx] unique_elements = torch.unique(chunk, sorted=True).tolist() # Create a mask for whether each token can attend to the current chunk for token_type in unique_elements: if token_type == article_separator: continue token_mask = (token_types[batch_idx] == token_type) inter_chunk_mask[batch_idx, :, chunk_idx] |= token_mask # Create a mask within each chunk unique_elements = [x for x in unique_elements if x != article_separator] if len(unique_elements) > 1 and chunk[-1] != article_separator: intra_chunk_mask[batch_idx, chunk_idx] = (chunk == unique_elements[-1]) inter_chunk_mask = ~inter_chunk_mask intra_chunk_mask = ~intra_chunk_mask window_mask = torch.logical_or(window_causal_mask, window_3d_mask.unsqueeze(1)) inter_chunk_mask = torch.logical_or(chunk_causal_mask, inter_chunk_mask.unsqueeze(1)) intra_chunk_mask = intra_chunk_mask.unsqueeze(1).unsqueeze(-1) joint_mask = torch.cat([window_mask, inter_chunk_mask.reshape(*window_mask.shape)], dim=-1) attention_mask = (joint_mask, intra_chunk_mask) else: joint_mask = torch.cat([window_causal_mask, chunk_causal_mask.reshape(*window_causal_mask.shape)], dim=-1) attention_mask = (joint_mask, None) return attention_mask def forward( self, input_ids: torch.LongTensor = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[List[torch.FloatTensor]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, labels: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, return_all_pred_logits: Optional[bool] = None, multibyte_decoding: Optional[bool] = None) -> Union[Tuple, CausalLMOutputWithPast]: output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions output_hidden_states = (output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states) return_dict = return_dict if return_dict is not None else self.config.use_return_dict if input_ids is None: assert past_key_values is None # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) outputs = self.model( input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, multibyte_decoding=multibyte_decoding, ) hidden_states = outputs[0] logits = self.lm_head(hidden_states) if self.config.fp32_logits: logits = logits.float() loss = None if labels is not None: loss_fct = CrossEntropyLoss(reduction="none") if hasattr(self.config, "num_pred_heads") and self.config.num_pred_heads > 1: shift_logits = logits.view(logits.shape[0], logits.shape[1], self.config.num_pred_heads, self.config.vocab_size) # shift_logits = shift_logits.view(-1, logits.shape[1] * self.config.num_pred_heads, self.config.vocab_size) shift_logits = shift_logits.view(-1, self.config.vocab_size) else: shift_logits = logits.view(-1, self.config.vocab_size) shift_labels = labels.view(-1) # Enable model parallelism shift_labels = shift_labels.to(shift_logits.device) loss = loss_fct(shift_logits, shift_labels) if hasattr(self.config, "num_pred_heads") and self.config.num_pred_heads > 1: all_pred_logits = logits.reshape(logits.shape[0], logits.shape[1], self.config.num_pred_heads, self.config.vocab_size) if return_all_pred_logits: logits = all_pred_logits else: logits = all_pred_logits[..., 0, :] if not return_dict: output = (logits, ) + outputs[1:] return (loss, ) + output if loss is not None else output return CausalLMOutputWithPast( loss=loss, logits=logits, past_key_values=outputs.past_key_values, hidden_states=outputs.hidden_states, attentions=outputs.attentions, ) def prepare_inputs_for_generation(self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, use_cache=True, **kwargs): # prefill phase: # input_ids: b x s # attention_mask: None if no padding or b x s # position_ids : b x s # token gen phase: # input_ids : b x 1 # attention_mask: b x 1 x s # position_ids: b x 1 past_length = 0 if past_key_values is not None: assert isinstance(past_key_values, Cache) past_length = past_key_values.get_seq_length() if attention_mask is not None and attention_mask.shape[1] > input_ids.shape[1]: input_ids = input_ids[:, -(attention_mask.shape[1] - past_length):] elif past_length < input_ids.shape[1]: input_ids = input_ids[:, past_length:] position_ids = kwargs.get("position_ids", None) if attention_mask is not None and position_ids is None: position_ids = attention_mask.long().cumsum(-1) - 1 position_ids.masked_fill_(attention_mask == 0, 1) if past_key_values: position_ids = position_ids[:, -input_ids.shape[1]:] # if `inputs_embeds` are passed, we only want to use them in the 1st generation step if inputs_embeds is not None and past_key_values is None: model_inputs = {"inputs_embeds": inputs_embeds} else: model_inputs = {"input_ids": input_ids} # must initialize position_ids at each step during GPU inference assert position_ids is not None model_inputs.update( { "position_ids": position_ids, "past_key_values": past_key_values, "use_cache": use_cache, "attention_mask": attention_mask, } ) return model_inputs @staticmethod def _reorder_cache(past_key_values, beam_idx): reordered_past = () for layer_past in past_key_values: reordered_past += (tuple( past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past), ) return reordered_past