Hugging Face's logo

AjayRangoji
/
AR-4B-base

Text Generation
Transformers
Safetensors
English
ar
causal-lm
pretrained
conversational
custom_code
Model card Files
xet
 Community
AR-4B-base / modeling_ar.py
AjayRangoji's picture
AjayRangoji
Initial release: AR-4B base model, pretrained from scratch on ~325B tokens
226b931 verified
Raw History Blame Contribute Delete
8.7 kB
"""Hugging Face loader for AR.
Keeps the model's own identity (`model_type: "ar"`, `ARForCausalLM`) and uses
Hugging Face's reference modules for grouped-query attention and SwiGLU so
third-party tools (vLLM, text-generation-webui, …) that already know those shapes
just work. Loaded via `AutoModelForCausalLM.from_pretrained(..., trust_remote_code=True)`.
"""
from __future__ import annotations
from typing import Optional, Tuple
import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import PretrainedConfig, PreTrainedModel
from transformers.modeling_outputs import CausalLMOutputWithPast
class ARConfigHF(PretrainedConfig):
model_type = "ar"
keys_to_ignore_at_inference = ["past_key_values"]
def __init__(
self,
vocab_size: int = 151936,
hidden_size: int = 2560,
intermediate_size: int = 9728,
num_hidden_layers: int = 36,
num_attention_heads: int = 32,
num_key_value_heads: int = 8,
head_dim: int = 128,
max_position_embeddings: int = 40960,
rope_theta: float = 1_000_000.0,
rms_norm_eps: float = 1e-6,
tie_word_embeddings: bool = True,
hidden_act: str = "silu",
attention_bias: bool = False,
mlp_bias: bool = False,
initializer_range: float = 0.02,
use_cache: bool = True,
pad_token_id: Optional[int] = 151645,
bos_token_id: Optional[int] = 151645,
eos_token_id: Optional[int] = 151645,
**kwargs,
):
self.vocab_size = vocab_size
self.hidden_size = hidden_size
self.intermediate_size = intermediate_size
self.num_hidden_layers = num_hidden_layers
self.num_attention_heads = num_attention_heads
self.num_key_value_heads = num_key_value_heads
self.head_dim = head_dim
self.max_position_embeddings = max_position_embeddings
self.rope_theta = rope_theta
self.rms_norm_eps = rms_norm_eps
self.hidden_act = hidden_act
self.attention_bias = attention_bias
self.mlp_bias = mlp_bias
self.initializer_range = initializer_range
self.use_cache = use_cache
super().__init__(
pad_token_id=pad_token_id, bos_token_id=bos_token_id,
eos_token_id=eos_token_id, tie_word_embeddings=tie_word_embeddings, **kwargs,
)
class RMSNorm(nn.Module):
def __init__(self, dim: int, eps: float = 1e-6):
super().__init__()
self.weight = nn.Parameter(torch.ones(dim))
self.eps = eps
def forward(self, x: torch.Tensor) -> torch.Tensor:
# computed in fp32 for stability, cast back to the input's dtype
orig = x.dtype
x = x.to(torch.float32)
x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
return (x * self.weight.to(torch.float32)).to(orig)
def _build_rope(head_dim: int, seq_len: int, theta: float, device, dtype):
# same formula as training: cos/sin tables indexed by position
idx = torch.arange(0, head_dim, 2, dtype=torch.float32, device=device)
inv_freq = 1.0 / (theta ** (idx / head_dim))
pos = torch.arange(seq_len, dtype=torch.float32, device=device)
angles = torch.outer(pos, inv_freq)
angles = torch.cat([angles, angles], dim=-1)
return angles.cos().to(dtype), angles.sin().to(dtype)
def _rotate_half(x: torch.Tensor) -> torch.Tensor:
half = x.shape[-1] // 2
return torch.cat([-x[..., half:], x[..., :half]], dim=-1)
def _apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
cos = cos[None, None, :, :]
sin = sin[None, None, :, :]
return x * cos + _rotate_half(x) * sin
class SelfAttention(nn.Module):
def __init__(self, cfg: ARConfigHF):
super().__init__()
self.num_heads = cfg.num_attention_heads
self.num_kv_heads = cfg.num_key_value_heads
self.head_dim = cfg.head_dim
self.q_proj = nn.Linear(cfg.hidden_size, self.num_heads * self.head_dim, bias=cfg.attention_bias)
self.k_proj = nn.Linear(cfg.hidden_size, self.num_kv_heads * self.head_dim, bias=cfg.attention_bias)
self.v_proj = nn.Linear(cfg.hidden_size, self.num_kv_heads * self.head_dim, bias=cfg.attention_bias)
self.o_proj = nn.Linear(self.num_heads * self.head_dim, cfg.hidden_size, bias=cfg.attention_bias)
self.q_norm = RMSNorm(self.head_dim, cfg.rms_norm_eps)
self.k_norm = RMSNorm(self.head_dim, cfg.rms_norm_eps)
self.groups = self.num_heads // self.num_kv_heads
def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor,
attn_mask: Optional[torch.Tensor]) -> torch.Tensor:
B, T, _ = x.shape
q = self.q_proj(x).view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
k = self.k_proj(x).view(B, T, self.num_kv_heads, self.head_dim).transpose(1, 2)
v = self.v_proj(x).view(B, T, self.num_kv_heads, self.head_dim).transpose(1, 2)
q = _apply_rope(self.q_norm(q), cos, sin)
k = _apply_rope(self.k_norm(k), cos, sin)
# expand KV heads to match Q (standard GQA repeat)
k = k.repeat_interleave(self.groups, dim=1)
v = v.repeat_interleave(self.groups, dim=1)
out = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, is_causal=attn_mask is None)
return self.o_proj(out.transpose(1, 2).contiguous().view(B, T, -1))
class MLP(nn.Module):
def __init__(self, cfg: ARConfigHF):
super().__init__()
self.gate_proj = nn.Linear(cfg.hidden_size, cfg.intermediate_size, bias=cfg.mlp_bias)
self.up_proj = nn.Linear(cfg.hidden_size, cfg.intermediate_size, bias=cfg.mlp_bias)
self.down_proj = nn.Linear(cfg.intermediate_size, cfg.hidden_size, bias=cfg.mlp_bias)
def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
class DecoderLayer(nn.Module):
def __init__(self, cfg: ARConfigHF):
super().__init__()
self.input_layernorm = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps)
self.post_attention_layernorm = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps)
self.self_attn = SelfAttention(cfg)
self.mlp = MLP(cfg)
def forward(self, x, cos, sin, attn_mask):
x = x + self.self_attn(self.input_layernorm(x), cos, sin, attn_mask)
x = x + self.mlp(self.post_attention_layernorm(x))
return x
class ARModelBase(nn.Module):
def __init__(self, cfg: ARConfigHF):
super().__init__()
self.cfg = cfg
self.embed_tokens = nn.Embedding(cfg.vocab_size, cfg.hidden_size)
self.layers = nn.ModuleList([DecoderLayer(cfg) for _ in range(cfg.num_hidden_layers)])
self.norm = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps)
def forward(self, input_ids: torch.Tensor,
attention_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
B, T = input_ids.shape
x = self.embed_tokens(input_ids)
cos, sin = _build_rope(self.cfg.head_dim, T, self.cfg.rope_theta, x.device, x.dtype)
# causal + optional per-position mask; is_causal=True in the attention call handles it
for layer in self.layers:
x = layer(x, cos, sin, attn_mask=None)
return self.norm(x)
class ARForCausalLM(PreTrainedModel):
config_class = ARConfigHF
base_model_prefix = "model"
_no_split_modules = ["DecoderLayer"]
def __init__(self, cfg: ARConfigHF):
super().__init__(cfg)
self.model = ARModelBase(cfg)
self.lm_head = nn.Linear(cfg.hidden_size, cfg.vocab_size, bias=False)
if cfg.tie_word_embeddings:
self.lm_head.weight = self.model.embed_tokens.weight
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 forward(
self,
input_ids: Optional[torch.Tensor] = None,
attention_mask: Optional[torch.Tensor] = None,
labels: Optional[torch.Tensor] = None,
**kwargs,
) -> CausalLMOutputWithPast:
h = self.model(input_ids, attention_mask)
logits = self.lm_head(h)
loss = None
if labels is not None:
shift_logits = logits[..., :-1, :].contiguous()
shift_labels = labels[..., 1:].contiguous()
loss = F.cross_entropy(shift_logits.view(-1, shift_logits.size(-1)),
shift_labels.view(-1))
return CausalLMOutputWithPast(loss=loss, logits=logits)