# persadian_nano_v4/PersadianNanoV4Model.py # Stable GPU-ready implementation # persadian-Nano-V4 (~160M total params) import math import torch import torch.nn as nn import torch.nn.functional as F # ============================================================ # CONFIG # ============================================================ class PersadianNanoV4Config: def __init__(self): # Core architecture self.hidden_size = 512 self.intermediate_size = 1024 self.num_hidden_layers = 6 # Attention self.num_attention_heads = 8 self.num_key_value_heads = 4 # MoE self.num_experts = 4 self.num_experts_per_tok = 2 # Progressive experts self.progressive_experts = True self.min_experts = 1 self.max_experts = 2 # Hyper connections self.use_hyper_connection = True # Compression self.use_compressed_attention = True self.compress_ratio = 4 # Context self.max_position_embeddings = 2048 # Tokenizer self.vocab_size = 50257 self.bos_token_id = 50256 self.eos_token_id = 50256 # RoPE self.rope_theta = 10000.0 # Regularization self.dropout = 0.1 self.layer_norm_eps = 1e-5 # Attention self.attention_dropout = 0.0 self.use_flash_attention = True # Precision self.torch_dtype = "float16" # ============================================================ # ROTARY EMBEDDING # ============================================================ class RotaryEmbedding(nn.Module): def __init__(self, dim, base=10000.0): super().__init__() inv_freq = 1.0 / ( base ** ( torch.arange(0, dim, 2).float() / dim ) ) self.register_buffer("inv_freq", inv_freq) def forward(self, x, position_ids): # x shape: # [batch, seq, heads, head_dim] batch, seq_len, num_heads, head_dim = x.shape freqs = torch.einsum( "bi,j->bij", position_ids.float(), self.inv_freq ) emb = torch.cat([freqs, freqs], dim=-1) cos = emb.cos()[:, :, None, :] sin = emb.sin()[:, :, None, :] x1 = x[..., ::2] x2 = x[..., 1::2] rotated = torch.stack( (-x2, x1), dim=-1 ).flatten(-2) return (x * cos) + (rotated * sin) # ============================================================ # ADAPTIVE HYPER CONNECTION # ============================================================ class AdaptiveHyperConnection(nn.Module): def __init__(self, hidden_size): super().__init__() self.router = nn.Sequential( nn.Linear(hidden_size, hidden_size // 4), nn.SiLU(), nn.Linear(hidden_size // 4, hidden_size), nn.Sigmoid() ) def forward(self, x, residual): weight = self.router( x.mean(dim=1, keepdim=True) ) return residual + (x * weight) # ============================================================ # ATTENTION # ============================================================ class CompressedSparseAttention(nn.Module): def __init__(self, config): super().__init__() self.config = config self.hidden_size = config.hidden_size self.num_heads = config.num_attention_heads self.num_key_value_heads = config.num_key_value_heads self.head_dim = ( config.hidden_size // config.num_attention_heads ) # Q projection self.q_proj = nn.Linear( config.hidden_size, config.hidden_size, bias=False ) # KV projection self.k_proj = nn.Linear( config.hidden_size, self.num_key_value_heads * self.head_dim, bias=False ) self.v_proj = nn.Linear( config.hidden_size, self.num_key_value_heads * self.head_dim, bias=False ) # Output self.o_proj = nn.Linear( config.hidden_size, config.hidden_size, bias=False ) # Compression compressed_dim = ( config.hidden_size // config.compress_ratio ) self.compressor = nn.Linear( config.hidden_size, compressed_dim ) self.k_proj_compressed = nn.Linear( compressed_dim, self.num_key_value_heads * self.head_dim, bias=False ) self.v_proj_compressed = nn.Linear( compressed_dim, self.num_key_value_heads * self.head_dim, bias=False ) # Rotary self.rotary_emb = RotaryEmbedding( self.head_dim, config.rope_theta ) def forward( self, hidden_states, attention_mask=None, position_ids=None, ): batch_size, seq_len, _ = hidden_states.shape # ==================================================== # Q # ==================================================== q = self.q_proj(hidden_states) q = q.view( batch_size, seq_len, self.num_heads, self.head_dim ) # ==================================================== # COMPRESSED OR NORMAL KV # ==================================================== if ( self.config.use_compressed_attention and seq_len > 512 ): compressed = self.compressor(hidden_states) k = self.k_proj_compressed(compressed) v = self.v_proj_compressed(compressed) else: k = self.k_proj(hidden_states) v = self.v_proj(hidden_states) k = k.view( batch_size, seq_len, self.num_key_value_heads, self.head_dim ) v = v.view( batch_size, seq_len, self.num_key_value_heads, self.head_dim ) # ==================================================== # GQA EXPANSION # ==================================================== repeat_factor = ( self.num_heads // self.num_key_value_heads ) k = k.repeat_interleave( repeat_factor, dim=2 ) v = v.repeat_interleave( repeat_factor, dim=2 ) # ==================================================== # ROPE # ==================================================== q = self.rotary_emb(q, position_ids) k = self.rotary_emb(k, position_ids) # ==================================================== # TRANSPOSE # ==================================================== q = q.transpose(1, 2) k = k.transpose(1, 2) v = v.transpose(1, 2) # ==================================================== # FLASH ATTENTION # ==================================================== if ( self.config.use_flash_attention and hasattr(F, "scaled_dot_product_attention") ): attn_output = F.scaled_dot_product_attention( q, k, v, attn_mask=attention_mask, dropout_p=self.config.attention_dropout, is_causal=True ) else: scores = torch.matmul( q, k.transpose(-2, -1) ) / math.sqrt(self.head_dim) if attention_mask is not None: scores = scores + attention_mask probs = F.softmax(scores, dim=-1) attn_output = torch.matmul( probs, v ) # ==================================================== # OUTPUT # ==================================================== attn_output = attn_output.transpose(1, 2) attn_output = attn_output.contiguous().view( batch_size, seq_len, self.hidden_size ) return self.o_proj(attn_output) # ============================================================ # MIXTURE OF EXPERTS # ============================================================ class MixtureOfExperts(nn.Module): def __init__(self, config): super().__init__() self.config = config self.num_experts = config.num_experts self.experts = nn.ModuleList([ nn.Sequential( nn.Linear( config.hidden_size, config.intermediate_size ), nn.SiLU(), nn.Linear( config.intermediate_size, config.hidden_size ), nn.Dropout(config.dropout) ) for _ in range(config.num_experts) ]) self.router = nn.Linear( config.hidden_size, config.num_experts, bias=False ) self.aux_loss_coef = 0.01 def forward( self, hidden_states, progressive_factor=1.0 ): batch_size, seq_len, hidden_size = hidden_states.shape hidden_states_flat = hidden_states.view( -1, hidden_size ) router_logits = self.router( hidden_states_flat ) routing_weights = F.softmax( router_logits, dim=-1 ) # ==================================================== # PROGRESSIVE EXPERT ACTIVATION # ==================================================== if self.config.progressive_experts: k = int( round( self.config.min_experts + ( self.config.max_experts - self.config.min_experts ) * progressive_factor ) ) k = max( self.config.min_experts, min(k, self.config.max_experts) ) else: k = self.config.num_experts_per_tok # ==================================================== # TOP-K ROUTING # ==================================================== top_k_weights, top_k_indices = torch.topk( routing_weights, k, dim=-1 ) top_k_weights = ( top_k_weights / top_k_weights.sum( dim=-1, keepdim=True ) ) # ==================================================== # EXPERT COMPUTE # ==================================================== final_hidden = torch.zeros_like( hidden_states_flat ) for expert_idx in range(self.num_experts): expert_mask = ( top_k_indices == expert_idx ).any(dim=-1) if expert_mask.any(): expert_input = hidden_states_flat[ expert_mask ] expert_output = self.experts[ expert_idx ](expert_input) expert_weights = top_k_weights[ expert_mask ].mean(dim=-1, keepdim=True) final_hidden[expert_mask] += ( expert_output * expert_weights ) # ==================================================== # AUX LOSS # ==================================================== router_probs = routing_weights.mean(dim=0) aux_loss = torch.var(router_probs) return ( final_hidden.view( batch_size, seq_len, hidden_size ), aux_loss * self.aux_loss_coef ) # ============================================================ # DECODER LAYER # ============================================================ class PersadianNanoV4DecoderLayer(nn.Module): def __init__(self, config, layer_idx): super().__init__() self.config = config self.input_layernorm = nn.LayerNorm( config.hidden_size, eps=config.layer_norm_eps ) self.self_attn = CompressedSparseAttention( config ) self.post_attention_layernorm = nn.LayerNorm( config.hidden_size, eps=config.layer_norm_eps ) self.moe = MixtureOfExperts(config) if config.use_hyper_connection: self.hc_attention = AdaptiveHyperConnection( config.hidden_size ) self.hc_moe = AdaptiveHyperConnection( config.hidden_size ) def forward( self, hidden_states, attention_mask=None, position_ids=None, progressive_factor=1.0 ): # ==================================================== # ATTENTION # ==================================================== residual = hidden_states hidden_states = self.input_layernorm( hidden_states ) attn_output = self.self_attn( hidden_states, attention_mask, position_ids ) if self.config.use_hyper_connection: hidden_states = self.hc_attention( attn_output, residual ) else: hidden_states = residual + attn_output # ==================================================== # MOE # ==================================================== residual = hidden_states hidden_states = self.post_attention_layernorm( hidden_states ) moe_output, aux_loss = self.moe( hidden_states, progressive_factor ) if self.config.use_hyper_connection: hidden_states = self.hc_moe( moe_output, residual ) else: hidden_states = residual + moe_output return hidden_states, aux_loss # ============================================================ # MAIN MODEL # ============================================================ class PersadianNanoV4Model(nn.Module): def __init__(self, config): super().__init__() self.config = config self.embed_tokens = nn.Embedding( config.vocab_size, config.hidden_size ) self.layers = nn.ModuleList([ PersadianNanoV4DecoderLayer( config, i ) for i in range(config.num_hidden_layers) ]) self.norm = nn.LayerNorm( config.hidden_size, eps=config.layer_norm_eps ) self.lm_head = nn.Linear( config.hidden_size, config.vocab_size, bias=False ) # Tie weights self.lm_head.weight = self.embed_tokens.weight # Init self.apply(self._init_weights) def _init_weights(self, module): if isinstance(module, nn.Linear): nn.init.normal_( module.weight, mean=0.0, std=0.02 ) if module.bias is not None: nn.init.zeros_(module.bias) elif isinstance(module, nn.Embedding): nn.init.normal_( module.weight, mean=0.0, std=0.02 ) def forward( self, input_ids, attention_mask=None, progressive_factor=1.0 ): batch_size, seq_len = input_ids.shape hidden_states = self.embed_tokens( input_ids ) position_ids = torch.arange( seq_len, device=input_ids.device ).unsqueeze(0) # ==================================================== # CAUSAL MASK # ==================================================== if attention_mask is None: mask = torch.triu( torch.full( (seq_len, seq_len), float("-inf"), device=input_ids.device ), diagonal=1 ) attention_mask = mask.unsqueeze(0).unsqueeze(0) total_aux_loss = torch.tensor( 0.0, device=input_ids.device ) for layer in self.layers: hidden_states, aux_loss = layer( hidden_states, attention_mask, position_ids, progressive_factor ) total_aux_loss += aux_loss hidden_states = self.norm(hidden_states) logits = self.lm_head(hidden_states) return logits, total_aux_loss @torch.no_grad() def generate( self, input_ids, max_new_tokens=50, temperature=0.7, top_k=50 ): self.eval() for _ in range(max_new_tokens): logits, _ = self.forward(input_ids) next_token_logits = ( logits[:, -1, :] / temperature ) if top_k > 0: values, _ = torch.topk( next_token_logits, top_k ) min_values = values[:, -1].unsqueeze(-1) next_token_logits = torch.where( next_token_logits < min_values, torch.full_like( next_token_logits, float("-inf") ), next_token_logits ) probs = F.softmax( next_token_logits, dim=-1 ) next_token = torch.multinomial( probs, num_samples=1 ) input_ids = torch.cat( [input_ids, next_token], dim=1 ) if ( next_token.item() == self.config.eos_token_id ): break return input_ids # ============================================================ # PARAMETER COUNT # ============================================================ def count_parameters(model): return sum( p.numel() for p in model.parameters() if p.requires_grad ) # ============================================================ # TEST # ============================================================ if __name__ == "__main__": print("=" * 60) print("Initializing Persadian-Nano-V4") print("=" * 60) config = PersadianNanoV4Config() model = PersadianNanoV4Model(config) total_params = count_parameters(model) print(f"Total Parameters: {total_params:,}") print(f"Approx Size: {total_params / 1e6:.2f}M") device = ( "cuda" if torch.cuda.is_available() else "cpu" ) print(f"Device: {device}") model = model.to(device) dummy_input = torch.randint( 0, config.vocab_size, (1, 128) ).to(device) print("\nRunning forward pass...") with torch.no_grad(): logits, aux_loss = model(dummy_input) print("Forward pass successful") print(f"Logits Shape: {logits.shape}") print(f"Aux Loss: {aux_loss.item():.6f}") print("\nTesting generation...") generated = model.generate( dummy_input[:, :16], max_new_tokens=10 ) print(f"Generated Shape: {generated.shape}") print("\nPersadian-Nano-V4 is operational") print("=" * 60)