# This is lightly adapted from https://github.com/EleutherAI/aria/blob/main/aria/model.py from typing import Optional, Union, Tuple import torch import torch.utils.checkpoint from torch import nn as nn from torch.nn import functional as F, CrossEntropyLoss from transformers import Cache, DynamicCache, StaticCache from transformers.utils import logging from transformers.generation import GenerationMixin from transformers.modeling_utils import PreTrainedModel from transformers.modeling_outputs import ( BaseModelOutputWithPast, CausalLMOutputWithPast, BaseModelOutputWithPoolingAndProjection, ) from transformers.modeling_attn_mask_utils import AttentionMaskConverter from .configuration_aria import AriaConfig logger = logging.get_logger(__name__) class AriaPreTrainedModel(PreTrainedModel): config_class = AriaConfig base_model_prefix = "aria" supports_gradient_checkpointing = True _no_split_modules = ["AriaBlock"] _skip_keys_device_placement = "past_key_values" _supports_flash_attn_2 = False _supports_cache_class = True _supports_quantized_cache = True _supports_static_cache = True _supports_sdpa = True _supports_flex_attn = False def _init_weights(self, module): if isinstance(module, nn.Linear): module.weight.data.normal_( mean=0.0, std=self.config.initializer_range ) if module.bias is not None: module.bias.data.zero_() elif isinstance(module, nn.Embedding): module.weight.data.normal_( mean=0.0, std=self.config.initializer_range ) if module.padding_idx is not None: module.weight.data[module.padding_idx].zero_() elif isinstance(module, nn.LayerNorm): module.bias.data.zero_() module.weight.data.fill_(1.0) class TransformerBlock(nn.Module): def __init__(self, model_config: AriaConfig, layer_idx: int): super().__init__() self.drop_p = 0.0 self.n_heads = model_config.num_attention_heads self.d_model = model_config.hidden_size self.d_head = ( model_config.hidden_size // model_config.num_attention_heads ) self.max_seq_len = model_config.max_seq_len self.layer_idx = layer_idx # Attention self.mixed_qkv = nn.Linear( in_features=self.d_model, out_features=3 * self.d_model, bias=False, ) self.att_proj_linear = nn.Linear( in_features=self.d_model, out_features=self.d_model, bias=False, ) # FF Layer self.ff_gate_proj = nn.Linear( in_features=self.d_model, out_features=self.d_model * model_config.ff_mult, bias=False, ) self.ff_up_proj = nn.Linear( in_features=self.d_model, out_features=self.d_model * model_config.ff_mult, bias=False, ) self.ff_down_proj = nn.Linear( in_features=self.d_model * model_config.ff_mult, out_features=self.d_model, bias=False, ) # Pre layer norms self.norm1 = nn.LayerNorm(self.d_model) self.norm2 = nn.LayerNorm(self.d_model) def forward( self, x: torch.Tensor, attention_mask: torch.Tensor, freqs_cis: torch.Tensor, position_ids: Optional[torch.Tensor] = None, past_key_values: Optional[ Union[Cache, Tuple[Tuple[torch.FloatTensor]]] ] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, cache_position: Optional[torch.Tensor] = None, ): attn_output, attn_weights, present = self._att_block( self.norm1(x), attention_mask, freqs_cis, past_key_values=past_key_values, use_cache=use_cache, output_attentions=output_attentions, cache_position=cache_position, ) x = x + attn_output x = x + self._ff_block(self.norm2(x)) outputs = (x, present) if use_cache: outputs = (x, present, attn_weights) else: outputs = (x, attn_weights) return outputs def _att_block( self, x: torch.Tensor, attention_mask: torch.Tensor, freqs_cis: torch.Tensor, past_key_values: Optional[ Union[Cache, Tuple[Tuple[torch.FloatTensor]]] ] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, cache_position: Optional[torch.Tensor] = None, ): batch_size, seq_len, _ = x.shape mixed_qkv = self.mixed_qkv(x) xq, xk, xv = mixed_qkv.chunk(3, -1) # Reshape for rotary embeddings # Need contiguous for q, k since in-place RoPE cannot be applied on a view xq = xq.reshape( batch_size, seq_len, self.n_heads, self.d_head ).contiguous() xk = xk.reshape( batch_size, seq_len, self.n_heads, self.d_head ).contiguous() xv = xv.view(batch_size, seq_len, self.n_heads, self.d_head) # apply_rotary_post_emb expects: (b_sz, s_len, n_head, d_head) xq = apply_rotary_emb(xq, freqs_cis) xk = apply_rotary_emb(xk, freqs_cis) xq, xk, xv = map(lambda t: t.transpose(1, 2), (xq, xk, xv)) if past_key_values is not None: cache_kwargs = { # "sin": sin, # "cos": cos, # "partial_rotation_size": self.rotary_ndims, "cache_position": cache_position, } xk, xv = past_key_values.update( xk, xv, self.layer_idx, cache_kwargs ) att = F.scaled_dot_product_attention( query=xq, key=xk, value=xv, attn_mask=attention_mask[..., : xk.shape[2]], ) # Reshape for out: (b_sz, s_len, n_head, d_head) out = att.transpose(1, 2).contiguous() out = out.view(batch_size, seq_len, self.n_heads * self.d_head) if not output_attentions: att = None return self.att_proj_linear(out), att, past_key_values def _ff_block(self, x: torch.Tensor): return self.ff_down_proj( F.silu(self.ff_gate_proj(x)) * self.ff_up_proj(x) ) class AriaModel(AriaPreTrainedModel): """Transformer decoder with no language model head. Args: model_config (ModelConfig): Model config settings. """ def __init__(self, model_config: AriaConfig): super().__init__(model_config) self.model_config = model_config self.freqs_cis = None self.causal_mask = None self.tok_embeddings = nn.Embedding( num_embeddings=model_config.vocab_size, embedding_dim=model_config.hidden_size, ) self.out_layer_norm = nn.LayerNorm(model_config.hidden_size) self.encode_layers = nn.ModuleList() for i in range(model_config.num_hidden_layers): self.encode_layers.append(TransformerBlock(model_config, i)) self.gradient_checkpointing = False self.post_init() def forward( self, input_ids: Optional[torch.Tensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.Tensor] = None, inputs_embeds: Optional[torch.Tensor] = None, past_key_values: Optional[ Union[Cache, Tuple[Tuple[torch.FloatTensor]]] ] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, cache_position: Optional[torch.Tensor] = None, ): """Forward pass of Transformer. Args: src (torch.tensor): Input to encoder block, of shape (batch_size, seq_len, d_model). attn_mask (Optional[torch.tensor]): Attention mask of shape (batch_size, seq_len). Defaults to None. past_kv (Optional[list[KVCache]]): a list of kv caches. The list index corresponds to the layer index. Returns: torch.tensor: Model outputs with shape (batch_size, seq_len, d_model). """ if ( input_ids is not None and input_ids.shape[1] > self.model_config.max_seq_len ): raise ValueError( f"Sequence length ({input_ids.shape[1]}) exceeds max_seq_len " f"({self.model_config.max_seq_len})." ) if ( inputs_embeds is not None and inputs_embeds.shape[1] > self.model_config.max_seq_len ): raise ValueError( f"Sequence length ({inputs_embeds.shape[1]}) exceeds max_seq_len " f"({self.model_config.max_seq_len})." ) output_attentions = ( output_attentions if output_attentions is not None else self.model_config.output_attentions ) output_hidden_states = ( output_hidden_states if output_hidden_states is not None else self.model_config.output_hidden_states ) return_dict = ( return_dict if return_dict is not None else self.model_config.use_return_dict ) use_cache = ( use_cache if use_cache is not None else self.model_config.use_cache ) if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError( "You must specify exactly one of input_ids or inputs_embeds" ) if self.gradient_checkpointing and self.training: if use_cache: logger.warning_once( "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." ) use_cache = False if inputs_embeds is None: inputs_embeds = self.tok_embeddings(input_ids) return_legacy_cache = False if use_cache and not isinstance(past_key_values, Cache): return_legacy_cache = True if past_key_values is None: past_key_values = DynamicCache() else: past_key_values = DynamicCache.from_legacy_cache( past_key_values ) logger.warning_once( "We detected that you are passing `past_key_values` as a tuple of tuples. This is deprecated and " "will be removed in v4.47. Please convert your cache or use an appropriate `Cache` class " "(https://huggingface.co/docs/transformers/kv_cache#legacy-cache-format)" ) seq_length = inputs_embeds.shape[1] if cache_position is None: past_seen_tokens = ( past_key_values.get_seq_length() if past_key_values is not None else 0 ) cache_position = torch.arange( past_seen_tokens, past_seen_tokens + seq_length, device=inputs_embeds.device, ) if position_ids is None: position_ids = cache_position.unsqueeze(0) hidden_states = inputs_embeds if self.causal_mask is None: self.causal_mask = precompute_causal_mask( max_seq_len=self.model_config.max_seq_len, ).to(input_ids.device) if self.freqs_cis is None: self.freqs_cis = precompute_freqs_cis( seq_len=self.model_config.max_seq_len, n_elem=self.model_config.hidden_size // self.model_config.num_attention_heads, base=500000, dtype=hidden_states.dtype, ).to(input_ids.device) freqs_cis = self.freqs_cis[cache_position] if use_cache is True: causal_mask = self.causal_mask[None, None, cache_position] else: causal_mask = self.causal_mask[None, None, :seq_length, :seq_length] if attention_mask is not None: pad_len = causal_mask.shape[3] - attention_mask.shape[1] padded_attention_mask = F.pad(attention_mask, (0, pad_len), value=1) padded_attention_mask = padded_attention_mask[:, None, None, :] padded_attention_mask = padded_attention_mask.bool() causal_mask = causal_mask & padded_attention_mask kwargs = { "position_ids": position_ids, "past_key_values": past_key_values, "use_cache": use_cache, "output_attentions": output_attentions, "output_hidden_states": output_hidden_states, "return_dict": return_dict, "cache_position": cache_position, } next_decoder_cache = None if self.gradient_checkpointing: for layer in self.encode_layers: def create_custom_forward(module): def custom_forward(*args): return module(*args)[0] return custom_forward hidden_states = torch.utils.checkpoint.checkpoint( create_custom_forward(layer), hidden_states, causal_mask, freqs_cis, **kwargs, preserve_rng_state=True, use_reentrant=True, ) else: all_attentions = () if output_attentions else None all_hidden_states = () if output_hidden_states else None for layer in self.encode_layers: if output_hidden_states: all_hidden_states = all_hidden_states + (hidden_states,) outputs = layer( hidden_states, causal_mask, freqs_cis=freqs_cis, **kwargs ) hidden_states = outputs[0] if use_cache is True: next_decoder_cache = outputs[1] if output_attentions: all_attentions = all_attentions + ( outputs[2 if use_cache else 1], ) if output_hidden_states: all_hidden_states = all_hidden_states + (hidden_states,) hidden_states = self.out_layer_norm(hidden_states) next_cache = next_decoder_cache if use_cache else None if return_legacy_cache: next_cache = next_cache.to_legacy_cache() if not return_dict: return tuple( v for v in [ hidden_states, next_cache, all_hidden_states, all_attentions, ] if v is not None ) return BaseModelOutputWithPast( last_hidden_state=hidden_states, past_key_values=next_cache, hidden_states=all_hidden_states, attentions=all_attentions, ) class AriaForCausalLM(AriaPreTrainedModel, GenerationMixin): """Transformer decoder with head for language modelling. Args: model_config (ModelConfig): Model config settings. """ def __init__(self, model_config: AriaConfig): super().__init__(model_config) self.model_config = model_config self.max_seq_len = model_config.max_seq_len self.model = AriaModel(model_config) self.lm_head = nn.Linear( model_config.hidden_size, model_config.vocab_size, bias=False ) self.post_init() def forward( self, input_ids: Optional[torch.Tensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.Tensor] = None, inputs_embeds: Optional[torch.Tensor] = None, past_key_values: Optional[ Union[Cache, Tuple[Tuple[torch.FloatTensor]]] ] = None, labels: Optional[torch.Tensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, cache_position: Optional[torch.Tensor] = None, ): """Forward pass of Transformer decoder with LM head.""" return_dict = ( return_dict if return_dict is not None else self.model_config.use_return_dict ) outputs = self.model( input_ids, attention_mask=attention_mask, position_ids=position_ids, inputs_embeds=inputs_embeds, past_key_values=past_key_values, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, cache_position=cache_position, ) hidden = outputs[0] lm_logits = self.lm_head(hidden) lm_loss = None if labels is not None: # move labels to correct device to enable model parallelism labels = labels.to(lm_logits.device) # we are doing next-token prediction; shift prediction scores and input ids by one shift_logits = lm_logits[:, :-1, :].contiguous() labels = labels[:, 1:].contiguous() loss_fct = CrossEntropyLoss() lm_loss = loss_fct( shift_logits.view(-1, shift_logits.size(-1)), labels.view(-1) ) if not return_dict: output = (lm_logits,) + outputs[1:] return ((lm_loss,) + output) if lm_loss is not None else output return CausalLMOutputWithPast( loss=lm_loss, logits=lm_logits, past_key_values=outputs.past_key_values, hidden_states=outputs.hidden_states, attentions=outputs.attentions, ) class AriaForSequenceEmbedding(AriaPreTrainedModel): """Transformer decoder embedding head for contrastive learning. Args: model_config (ModelConfig): Model config settings. """ def __init__(self, model_config: AriaConfig): super().__init__(model_config) assert model_config.embedding_size self.model_config = model_config self.max_seq_len = model_config.max_seq_len self.model = AriaModel(model_config) self.emb_head = nn.Linear( model_config.hidden_size, model_config.embedding_size, bias=False ) self.post_init() def get_pooled_embedding( self, input_ids: torch.Tensor, embedding: torch.Tensor ): _batch_size = input_ids.shape[0] eos_mask = input_ids == self.config.eos_token_id if not eos_mask.any(dim=1).all(): raise ValueError("Each sequence must contain a EOS token") eos_pos = eos_mask.int().argmax(dim=1) pooled_embedding = embedding[ torch.arange(_batch_size, device=input_ids.device), eos_pos ] return pooled_embedding def forward( self, input_ids: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.Tensor] = None, inputs_embeds: Optional[torch.Tensor] = None, past_key_values: Optional[ Union[Cache, Tuple[Tuple[torch.FloatTensor]]] ] = None, labels: Optional[torch.Tensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, cache_position: Optional[torch.Tensor] = None, ): """Forward pass of Transformer decoder with embedding head. Pooled embedding is extracted from EOS token.""" return_dict = ( return_dict if return_dict is not None else self.model_config.use_return_dict ) if ( position_ids is not None or inputs_embeds is not None or past_key_values is not None or labels is not None or cache_position is not None or use_cache ): raise ValueError("Provided args unsupported for embedding head") outputs = self.model( input_ids, attention_mask=attention_mask, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, use_cache=False, ) hidden = outputs[0] embedding = self.emb_head(hidden) pooled_embedding = self.get_pooled_embedding( input_ids=input_ids, embedding=embedding, ) if not return_dict: output = (pooled_embedding,) + outputs[1:] return output return BaseModelOutputWithPoolingAndProjection( last_hidden_state=embedding, pooler_output=pooled_embedding, hidden_states=outputs.hidden_states, attentions=outputs.attentions, ) def precompute_causal_mask(max_seq_len: int): return torch.tril( torch.ones(max_seq_len, max_seq_len, dtype=torch.bool) ).cuda() def precompute_freqs_cis( seq_len: int, n_elem: int, base: int = 500000, dtype: torch.dtype = torch.bfloat16, ): freqs = 1.0 / ( base ** (torch.arange(0, n_elem, 2)[: (n_elem // 2)].float() / n_elem) ) t = torch.arange(seq_len, device=freqs.device) freqs = torch.outer(t, freqs) freqs_cis = torch.polar(torch.ones_like(freqs), freqs) cache = torch.stack([freqs_cis.real, freqs_cis.imag], dim=-1) return cache.to(dtype=dtype) def apply_rotary_emb(x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor: """ In-place RoPE. Credits to Katherine Crowson: x shape (b_sz, s_len, n_head, d_head). cos, sin shape (s_len, d_head // 2). """ d = x.shape[-1] // 2 cos = freqs_cis[..., 0][None, :, None] sin = freqs_cis[..., 1][None, :, None] x1, x2 = x[..., :d], x[..., d : d * 2] tmp = x1.clone() x1.mul_(cos).addcmul_(x2, sin, value=-1) x2.mul_(cos).addcmul_(tmp, sin, value=1) return x __all__ = [ "AriaPreTrainedModel", "AriaModel", "TransformerBlock", "AriaForCausalLM", "AriaForSequenceEmbedding", ]