""" FORGE-3B: Complete Model Architecture Integrates ARG, MHA, Dense FFN, HSE FFN into the full autoregressive LM. """ from __future__ import annotations import math import logging from typing import Optional, List, Tuple, Dict, Any, Union import torch import torch.nn as nn import torch.nn.functional as F from torch.utils.checkpoint import checkpoint from .arg_layer import ARGLayer from .mha_layer import GlobalMHALayer from .hse_layer import HSELayer, DenseSwiGLUFFN from .dgn_norm import build_norm logger = logging.getLogger(__name__) class ForgeBlock(nn.Module): """ One FORGE block: SeqMixer + FFN (both with pre-norm and residual). Blocks are composed sequentially. The seq_mixer is either ARG or MHA; the ffn is either DenseSwiGLUFFN or HSELayer — determined by layer_idx. """ def __init__( self, d_model: int, layer_idx: int, is_mha: bool, is_hse: bool, # ARG config arg_d_inner: int = 1280, arg_d_state: int = 48, arg_d_rank: int = 48, arg_conv_kernel: int = 4, arg_local_window: int = 64, arg_local_n_heads: int = 8, arg_local_n_kv_heads: int = 2, arg_head_dim: int = 80, # MHA config mha_n_heads: int = 16, mha_n_kv_heads: int = 4, mha_head_dim: int = 80, rope_base: float = 500_000.0, rope_scaling_type: Optional[str] = None, rope_scaling_factor: float = 1.0, max_seq_len: int = 4096, # Dense FFN config dense_d_ff: int = 3200, # HSE config hse_n_domains: int = 4, hse_n_experts_per_domain: int = 8, hse_top_k: int = 2, hse_d_ff_expert: int = 288, hse_capacity_factor: float = 1.25, hse_expert_dropout: float = 0.1, hse_aux_loss_alpha: float = 0.01, # Norm norm_type: str = "dgn", dgn_n_groups: int = 16, norm_eps: float = 1e-6, use_flash_attention: bool = True, dropout: float = 0.0, ): super().__init__() self.layer_idx = layer_idx self.is_mha = is_mha self.is_hse = is_hse # ── Sequence Mixer ──────────────────────────────────────────────────── if is_mha: self.seq_mixer = GlobalMHALayer( d_model=d_model, n_heads=mha_n_heads, n_kv_heads=mha_n_kv_heads, head_dim=mha_head_dim, rope_base=rope_base, rope_scaling_type=rope_scaling_type, rope_scaling_factor=rope_scaling_factor, max_seq_len=max_seq_len, dropout=dropout, norm_type=norm_type, dgn_n_groups=dgn_n_groups, norm_eps=norm_eps, use_flash_attention=use_flash_attention, layer_idx=layer_idx, ) else: self.seq_mixer = ARGLayer( d_model=d_model, d_inner=arg_d_inner, d_state=arg_d_state, d_rank=arg_d_rank, conv_kernel=arg_conv_kernel, local_window=arg_local_window, local_n_heads=arg_local_n_heads, local_n_kv_heads=arg_local_n_kv_heads, head_dim=arg_head_dim, norm_type=norm_type, dgn_n_groups=dgn_n_groups, norm_eps=norm_eps, use_flash_attention=use_flash_attention, rope_base=rope_base, layer_idx=layer_idx, max_seq_len=max_seq_len, ) # ── FFN ─────────────────────────────────────────────────────────────── if is_hse: self.ffn = HSELayer( d_model=d_model, n_domains=hse_n_domains, n_experts_per_domain=hse_n_experts_per_domain, top_k=hse_top_k, d_ff_expert=hse_d_ff_expert, capacity_factor=hse_capacity_factor, expert_dropout=hse_expert_dropout, aux_loss_alpha=hse_aux_loss_alpha, norm_type=norm_type, dgn_n_groups=dgn_n_groups, norm_eps=norm_eps, layer_idx=layer_idx, ) else: self.ffn = DenseSwiGLUFFN( d_model=d_model, d_ff=dense_d_ff, norm_type=norm_type, dgn_n_groups=dgn_n_groups, norm_eps=norm_eps, ) def forward( self, x: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.Tensor] = None, position_offset: int = 0, use_cache: bool = False, return_aux_loss: bool = True, ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: aux_loss = None # ── Sequence Mixing ─────────────────────────────────────────────────── if self.is_mha: x, _ = self.seq_mixer( x, attention_mask=attention_mask, position_ids=position_ids, use_cache=use_cache ) else: x = self.seq_mixer(x, position_offset=position_offset) # ── FFN ─────────────────────────────────────────────────────────────── if self.is_hse: x, aux_loss = self.ffn(x, return_aux_loss=return_aux_loss) else: x = self.ffn(x) return x, aux_loss class ForgeModel(nn.Module): """ FORGE-3B: Full autoregressive language model. Architecture: - Token embedding (tied to LM head) - 36 ForgeBlocks (27 ARG + 9 MHA for seq mixing; 18 Dense + 18 HSE for FFN) - Final DGN normalization - LM head (weight-tied to embedding) """ def __init__(self, config): super().__init__() self.config = config d = config.d_model # ── Token Embedding ─────────────────────────────────────────────────── self.embed_tokens = nn.Embedding(config.vocab_size, d, padding_idx=0) # ── FORGE Blocks ────────────────────────────────────────────────────── mha_set = set(config.mha_layer_indices) hse_set = set(config.hse_ffn_layer_indices) self.layers = nn.ModuleList([ ForgeBlock( d_model=d, layer_idx=i, is_mha=(i in mha_set), is_hse=(i in hse_set), arg_d_inner=config.arg_d_inner, arg_d_state=config.arg_d_state, arg_d_rank=config.arg_d_rank, arg_conv_kernel=config.arg_conv_kernel, arg_local_window=config.arg_local_window, arg_local_n_heads=config.arg_local_n_heads, arg_local_n_kv_heads=config.arg_local_n_kv_heads, arg_head_dim=config.arg_head_dim, mha_n_heads=config.mha_n_heads, mha_n_kv_heads=config.mha_n_kv_heads, mha_head_dim=config.mha_head_dim, rope_base=config.rope_theta, rope_scaling_type=config.rope_scaling_type, rope_scaling_factor=config.rope_scaling_factor, max_seq_len=config.max_seq_len, dense_d_ff=config.dense_d_ff, hse_n_domains=config.hse_n_domains, hse_n_experts_per_domain=config.hse_n_experts_per_domain, hse_top_k=config.hse_top_k, hse_d_ff_expert=config.hse_d_ff_expert, hse_capacity_factor=config.hse_capacity_factor, hse_expert_dropout=config.hse_expert_dropout, hse_aux_loss_alpha=config.hse_aux_loss_alpha, norm_type=config.norm_type, dgn_n_groups=config.dgn_n_groups, norm_eps=config.norm_eps, use_flash_attention=config.use_flash_attention, ) for i in range(config.n_layers) ]) # ── Final Norm ──────────────────────────────────────────────────────── self.norm_final = build_norm(config.norm_type, d, config.dgn_n_groups, config.norm_eps) # ── LM Head ─────────────────────────────────────────────────────────── self.lm_head = nn.Linear(d, config.vocab_size, bias=False) # Tie weights if config.tie_word_embeddings: self.lm_head.weight = self.embed_tokens.weight # ── Gradient Checkpointing Configuration ───────────────────────────── self.gradient_checkpointing_enabled = config.use_gradient_checkpointing self.gradient_checkpointing_ratio = config.gradient_checkpointing_ratio # Checkpoint every N-th layer if config.gradient_checkpointing_ratio > 0: self._checkpoint_interval = max(1, int(1.0 / config.gradient_checkpointing_ratio)) else: self._checkpoint_interval = 1 # Apply weight initialization self._init_weights() # Log parameter count n_params = sum(p.ds_numel if hasattr(p, "ds_numel") else p.numel() for p in self.parameters()) logger.info(f"ForgeModel: {n_params/1e9:.3f}B total parameters") emb_p = self.embed_tokens.weight emb_numel = emb_p.ds_numel if hasattr(emb_p, "ds_numel") else emb_p.numel() logger.info(f" Embedding: {emb_numel/1e6:.1f}M") logger.info(f" N layers: {config.n_layers} " f"({len(mha_set)} MHA + {config.n_layers - len(mha_set)} ARG, " f"{len(hse_set)} HSE + {config.n_layers - len(hse_set)} Dense)") def _init_weights(self): """Initialize all model weights with scaled normal distribution.""" n_layers = self.config.n_layers for module in self.modules(): if isinstance(module, nn.Linear): nn.init.normal_(module.weight, std=self.config.init_std) if module.bias is not None: nn.init.zeros_(module.bias) elif isinstance(module, nn.Embedding): nn.init.normal_(module.weight, std=self.config.init_std) if module.padding_idx is not None: module.weight.data[module.padding_idx].zero_() # Residual scaling: scale output projections by 1/sqrt(2*n_layers) # This prevents the residual stream from growing with depth if self.config.init_std_residual_scale: scale = 1.0 / math.sqrt(2 * n_layers) for layer in self.layers: if isinstance(layer.seq_mixer, ARGLayer): nn.init.normal_(layer.seq_mixer.out_proj.weight, std=scale) nn.init.normal_(layer.seq_mixer.local_o.weight, std=scale) elif isinstance(layer.seq_mixer, GlobalMHALayer): nn.init.normal_(layer.seq_mixer.o_proj.weight, std=scale) if isinstance(layer.ffn, DenseSwiGLUFFN): nn.init.normal_(layer.ffn.down.weight, std=scale) elif isinstance(layer.ffn, HSELayer): for expert in layer.ffn.experts: nn.init.normal_(expert.down.weight, std=scale) def get_input_embeddings(self) -> nn.Embedding: return self.embed_tokens def set_input_embeddings(self, value: nn.Embedding): self.embed_tokens = value def enable_gradient_checkpointing(self): self.gradient_checkpointing_enabled = True logger.info("Gradient checkpointing enabled") def disable_gradient_checkpointing(self): self.gradient_checkpointing_enabled = False def forward( self, input_ids: torch.Tensor, # (B, T) attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.Tensor] = None, labels: Optional[torch.Tensor] = None, # (B, T) for LM loss use_cache: bool = False, return_aux_loss: bool = True, output_hidden_states: bool = False, ) -> Dict[str, torch.Tensor]: """ Forward pass for FORGE-3B. Returns dict with: 'loss': cross-entropy + aux_loss (if labels provided) 'logits': (B, T, vocab_size) 'aux_loss': sum of MoE load-balance losses 'hidden_states': list of layer outputs (if output_hidden_states) """ B, T = input_ids.shape device = input_ids.device # ── Token Embeddings ────────────────────────────────────────────────── h = self.embed_tokens(input_ids) # (B, T, d_model) if position_ids is None: position_ids = torch.arange(T, device=device).unsqueeze(0).expand(B, -1) # ── Forward Through Layers ──────────────────────────────────────────── all_hidden_states = [] if output_hidden_states else None total_aux_loss = torch.zeros(1, device=device, dtype=h.dtype) n_hse_layers = 0 for layer_idx, layer in enumerate(self.layers): if output_hidden_states: all_hidden_states.append(h) # Selective gradient checkpointing if (self.gradient_checkpointing_enabled and self.training and layer_idx % self._checkpoint_interval == 0): def _forward_layer(h_, layer_=layer): out, aux = layer_( h_, attention_mask=attention_mask, position_ids=position_ids, position_offset=0, use_cache=False, return_aux_loss=True, ) return out, aux if aux is not None else torch.zeros(1, device=h_.device) h, aux = checkpoint( _forward_layer, h, use_reentrant=False, # required for PyTorch DDP to prevent double-marking ready preserve_rng_state=True, ) else: h, aux = layer( h, attention_mask=attention_mask, position_ids=position_ids, position_offset=0, use_cache=use_cache, return_aux_loss=return_aux_loss, ) if aux is not None and return_aux_loss: total_aux_loss = total_aux_loss + aux n_hse_layers += 1 # ── Final Norm ──────────────────────────────────────────────────────── h = self.norm_final(h) # (B, T, d_model) # ── LM Head & Loss Computation ──────────────────────────────────────── loss = None if labels is not None: # Shift for next-token prediction shift_h = h[:, :-1, :].contiguous() shift_labels = labels[:, 1:].contiguous().reshape(-1) # Chunked LM head projection to prevent allocating massive (B, T, 206464) FP32 tensor in VRAM shift_h_flat = shift_h.reshape(-1, shift_h.shape[-1]) chunk_size = 1024 n_tokens = shift_h_flat.shape[0] total_ce_loss = 0.0 valid_tokens = 0 for i in range(0, n_tokens, chunk_size): h_chunk = shift_h_flat[i : i + chunk_size] labels_chunk = shift_labels[i : i + chunk_size] # Check for valid tokens mask = (labels_chunk != -100) n_valid = mask.sum().item() if n_valid == 0: continue logits_chunk = self.lm_head(h_chunk).float() chunk_loss = F.cross_entropy( logits_chunk, labels_chunk, ignore_index=-100, reduction="sum", ) total_ce_loss = total_ce_loss + chunk_loss valid_tokens += n_valid ce_loss = total_ce_loss / max(1, valid_tokens) if valid_tokens > 0 else (shift_h.sum() * 0.0) avg_aux = total_aux_loss / max(1, n_hse_layers) loss = ce_loss + avg_aux logits = None if self.training else self.lm_head(h) else: logits = self.lm_head(h) output = { "loss": loss, "logits": logits, "aux_loss": total_aux_loss / max(1, n_hse_layers), } if output_hidden_states: output["hidden_states"] = all_hidden_states return output @torch.no_grad() def generate( self, input_ids: torch.Tensor, max_new_tokens: int = 256, temperature: float = 1.0, top_p: float = 0.9, top_k: int = 50, repetition_penalty: float = 1.1, eos_token_id: int = 2, pad_token_id: int = 0, ) -> torch.Tensor: """ Autoregressive generation with top-p/top-k sampling. Clears KV caches after generation. """ self.eval() B, T_prompt = input_ids.shape device = input_ids.device generated = input_ids.clone() # Clear KV caches in MHA layers for layer in self.layers: if isinstance(layer.seq_mixer, GlobalMHALayer): layer.seq_mixer.clear_kv_cache() for step in range(max_new_tokens): # Only process new tokens after first step if step == 0: curr_input = generated else: curr_input = generated[:, -1:] # only last token outputs = self.forward(curr_input, use_cache=True, return_aux_loss=False) logits = outputs["logits"][:, -1, :] # (B, vocab_size) # Repetition penalty if repetition_penalty != 1.0: for b in range(B): for tok_id in set(generated[b].tolist()): if logits[b, tok_id] < 0: logits[b, tok_id] *= repetition_penalty else: logits[b, tok_id] /= repetition_penalty # Temperature if temperature != 1.0: logits = logits / temperature # Top-k filtering if top_k > 0: topk_vals = torch.topk(logits, min(top_k, logits.shape[-1])).values logits = logits.masked_fill(logits < topk_vals[:, -1:], float('-inf')) # Top-p (nucleus) sampling if top_p < 1.0: sorted_logits, sorted_idx = torch.sort(logits, descending=True) cumprobs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1) # Remove tokens beyond cumulative probability top_p remove_mask = cumprobs - F.softmax(sorted_logits, dim=-1) > top_p sorted_logits[remove_mask] = float('-inf') logits = torch.scatter(logits, 1, sorted_idx, sorted_logits) # Sample probs = F.softmax(logits, dim=-1) next_token = torch.multinomial(probs, num_samples=1) # (B, 1) generated = torch.cat([generated, next_token], dim=1) # Stop if all sequences hit EOS if (generated == eos_token_id).any(dim=1).all(): break # Clear caches for layer in self.layers: if isinstance(layer.seq_mixer, GlobalMHALayer): layer.seq_mixer.clear_kv_cache() return generated def count_parameters(self) -> Dict[str, int]: """Count parameters by component.""" counts = {} counts["total"] = sum(p.numel() for p in self.parameters()) counts["embedding"] = self.embed_tokens.weight.numel() counts["arg_layers"] = sum( p.numel() for i, l in enumerate(self.layers) if isinstance(l.seq_mixer, ARGLayer) for p in l.seq_mixer.parameters() ) counts["mha_layers"] = sum( p.numel() for l in self.layers if isinstance(l.seq_mixer, GlobalMHALayer) for p in l.seq_mixer.parameters() ) counts["dense_ffn"] = sum( p.numel() for l in self.layers if isinstance(l.ffn, DenseSwiGLUFFN) for p in l.ffn.parameters() ) counts["hse_ffn_total"] = sum( p.numel() for l in self.layers if isinstance(l.ffn, HSELayer) for p in l.ffn.parameters() ) hse_active = 0 for l in self.layers: if isinstance(l.ffn, HSELayer): hse_active += l.ffn.top_k * 3 * l.ffn.d_model * l.ffn.d_ff_expert counts["hse_ffn_active"] = hse_active return counts def build_forge_500m(config=None) -> ForgeModel: """Build FORGE-500M (500M total params) from config (or default config).""" from config import ForgeModelConfig if config is None: config = ForgeModelConfig() model = ForgeModel(config) return model def build_forge_3b(config=None) -> ForgeModel: """Alias for build_forge_500m (FORGE architecture downgraded to 1B max params).""" return build_forge_500m(config)