import math import importlib.util import os from typing import Optional, Tuple import torch import torch.nn as nn import torch.nn.functional as F from transformers import ( AutoConfig, AutoModel, AutoModelForCausalLM, PreTrainedModel, PretrainedConfig, ) from transformers.modeling_outputs import CausalLMOutput, CausalLMOutputWithPast # argonne4.5: the SDPA OOM guard below must work on CUDA *and* Intel XPU, where # `torch.cuda.OutOfMemoryError` is not defined at all. torch.OutOfMemoryError is # the backend-agnostic form (torch >= 2.5); fall back for older CUDA builds. _OOM_ERRORS = tuple(e for e in ( getattr(torch, "OutOfMemoryError", None), getattr(getattr(torch, "cuda", None), "OutOfMemoryError", None), ) if isinstance(e, type)) or (RuntimeError,) flash_attn_func = None _flash_attn_available = False if importlib.util.find_spec("flash_attn") is not None: try: from flash_attn.flash_attn_interface import flash_attn_func _flash_attn_available = True except ImportError: _flash_attn_available = False # One-time startup log of which attention kernel is in effect, because the # sliding-window local attention is only implemented on the flash-attn path # and silently degrades to full attention on the SDPA/math fallbacks. _attention_path_logged = False class ArgonneConfig(PretrainedConfig): """Configuration for the Argonne v3.0 family of models.""" model_type = "argonne2" def __init__( self, vocab_size: int = 151936, hidden_size: int = 3072, num_hidden_layers: int = 24, num_attention_heads: int = 12, num_key_value_heads: Optional[int] = 4, intermediate_size: Optional[int] = 8192, max_position_embeddings: int = 1024, attention_dropout: float = 0.0, hidden_dropout: float = 0.0, rms_norm_eps: float = 1e-6, rope_theta: float = 1000000.0, sliding_window: Optional[int] = None, use_flash_attention: bool = True, use_gradient_checkpointing: bool = False, qk_norm: bool = True, v_norm: bool = True, sandwich_norm: bool = True, z_loss_weight: float = 0.0, mtp_horizon: int = 1, mtp_loss_weight: float = 0.0, interleaved_local_attention: bool = True, local_attention_window: Optional[int] = 256, # --- argonne4.5 additions (all default-OFF: an argonne4.0 config reproduces 4.0 exactly) --- attn_pattern: Optional[str] = None, sliding_window_size: int = 2048, nope_global: bool = False, attn_gate: bool = False, mlp_type: str = "swiglu", mtp_module_layers: int = 0, doc_mask: bool = False, logit_softcap: float = 15.0, loss_chunk_size: int = 0, tie_word_embeddings: bool = True, attention_bias: bool = False, mlp_bias: bool = False, pad_token_id: Optional[int] = None, bos_token_id: Optional[int] = None, eos_token_id: Optional[int] = None, **kwargs, ) -> None: pad_token_id = pad_token_id if pad_token_id is not None else kwargs.pop("pad_token_id", None) bos_token_id = bos_token_id if bos_token_id is not None else kwargs.pop("bos_token_id", None) eos_token_id = eos_token_id if eos_token_id is not None else kwargs.pop("eos_token_id", None) super().__init__( pad_token_id=pad_token_id, bos_token_id=bos_token_id, eos_token_id=eos_token_id, **kwargs, ) # Backwards compatibility with Argonne 1.x naming. if "n_layer" in kwargs: num_hidden_layers = kwargs["n_layer"] if "n_head" in kwargs: num_attention_heads = kwargs["n_head"] if "n_embd" in kwargs: hidden_size = kwargs["n_embd"] if "block_size" in kwargs: max_position_embeddings = kwargs["block_size"] self.vocab_size = vocab_size self.hidden_size = hidden_size self.num_hidden_layers = num_hidden_layers self.num_attention_heads = num_attention_heads self.num_key_value_heads = ( num_key_value_heads if num_key_value_heads is not None else num_attention_heads // 2 ) if self.num_key_value_heads < 1: self.num_key_value_heads = 1 if num_attention_heads % self.num_key_value_heads != 0: raise ValueError("num_attention_heads must be divisible by num_key_value_heads") if intermediate_size is None: width = int(8 * hidden_size / 3) self.intermediate_size = ((width + 255) // 256) * 256 else: self.intermediate_size = intermediate_size self.max_position_embeddings = max_position_embeddings self.attention_dropout = attention_dropout self.hidden_dropout = hidden_dropout self.rms_norm_eps = rms_norm_eps self.rope_theta = rope_theta self.sliding_window = sliding_window self.use_flash_attention = use_flash_attention self.use_gradient_checkpointing = use_gradient_checkpointing self.qk_norm = qk_norm self.v_norm = v_norm self.sandwich_norm = sandwich_norm self.z_loss_weight = z_loss_weight self.mtp_horizon = max(1, int(mtp_horizon)) self.mtp_loss_weight = float(mtp_loss_weight) self.interleaved_local_attention = bool(interleaved_local_attention) self.local_attention_window = ( int(local_attention_window) if local_attention_window is not None and int(local_attention_window) > 0 else None ) # argonne4.5 attention layout. `attn_pattern` is a cycled string of 'L' (local/sliding) # and 'G' (global/full) -- Muse Glimmer's layout is "LLLG" (3 sliding@2048: 1 full). # When set it OVERRIDES interleaved_local_attention/local_attention_window. `nope_global` # strips RoPE from the global layers (Muse sets layer_rope_theta=0 on exactly those), # so the layers that see the whole sequence carry no frequency basis to extrapolate wrong. self.attn_pattern = attn_pattern.upper() if attn_pattern else None if self.attn_pattern is not None and set(self.attn_pattern) - {"L", "G"}: raise ValueError(f"attn_pattern must contain only 'L'/'G', got {attn_pattern!r}") self.sliding_window_size = int(sliding_window_size) self.nope_global = bool(nope_global) self.attn_gate = bool(attn_gate) if mlp_type not in ("swiglu", "relu2"): raise ValueError(f"mlp_type must be 'swiglu' or 'relu2', got {mlp_type!r}") self.mlp_type = mlp_type self.mtp_module_layers = max(0, int(mtp_module_layers)) self.doc_mask = bool(doc_mask) self.logit_softcap = float(logit_softcap) # >0 enables chunked cross-entropy (frees the full (batch*seq, vocab) fp32 logit transient # that caps batch at long context). 0 = off (default; unchanged full-logit path). self.loss_chunk_size = int(loss_chunk_size) self.tie_word_embeddings = tie_word_embeddings self.attention_bias = attention_bias self.mlp_bias = mlp_bias if self.pad_token_id is None and self.eos_token_id is not None: self.pad_token_id = self.eos_token_id # Backwards compatibility aliases self.n_embd = self.hidden_size self.n_layer = self.num_hidden_layers self.n_head = self.num_attention_heads self.block_size = self.max_position_embeddings class RMSNorm(nn.Module): def __init__(self, hidden_size: int, eps: float = 1e-6) -> None: super().__init__() self.eps = eps self.weight = nn.Parameter(torch.ones(hidden_size)) def forward(self, x: torch.Tensor) -> torch.Tensor: orig_dtype = x.dtype x = x.to(torch.float32) # Clamp values to prevent overflow in pow(2) x = torch.clamp(x, min=-65504.0, max=65504.0) variance = x.pow(2).mean(-1, keepdim=True) x = x * torch.rsqrt(variance + self.eps) return (self.weight * x.to(orig_dtype)) class RotaryEmbedding(nn.Module): def __init__( self, dim: int, max_position_embeddings: int = 2048, base: float = 10000.0, device: Optional[torch.device] = None, ) -> None: super().__init__() self.dim = dim self.max_position_embeddings = max_position_embeddings self.base = base # Always compute inv_freq on CPU to avoid meta-device issues during from_pretrained # (low_cpu_mem_usage=True initializes model on meta device, which would give uninitialized values) inv_freq = 1.0 / ( self.base ** (torch.arange(0, dim, 2, dtype=torch.float32, device="cpu") / dim) ) self.register_buffer("inv_freq", inv_freq, persistent=False) self._set_cos_sin_cache(max_position_embeddings, torch.device("cpu"), torch.float32) def _set_cos_sin_cache(self, seq_len: int, device: torch.device, dtype: torch.dtype) -> None: self.max_seq_len_cached = seq_len t = torch.arange(seq_len, device=device, dtype=self.inv_freq.dtype) freqs = torch.outer(t, self.inv_freq) emb = torch.cat((freqs, freqs), dim=-1) # Always store cache in float32 to avoid precision loss when model dtype is bfloat16 # (large freqs lose precision in bfloat16, corrupting cos/sin values) self.register_buffer("cos_cached", emb.cos().to(torch.float32), persistent=False) self.register_buffer("sin_cached", emb.sin().to(torch.float32), persistent=False) def forward(self, x: torch.Tensor, seq_len: int) -> Tuple[torch.Tensor, torch.Tensor]: if seq_len > self.max_seq_len_cached: self._set_cos_sin_cache(seq_len, x.device, x.dtype) return ( self.cos_cached[:seq_len].to(dtype=x.dtype, device=x.device), self.sin_cached[:seq_len].to(dtype=x.dtype, device=x.device), ) def init_buffers(self) -> None: """Re-initialize inv_freq, cos_cached, sin_cached. Call after from_pretrained to fix meta-device corruption (low_cpu_mem_usage=True destroys buffer values).""" device = self.inv_freq.device if device.type == "meta": device = torch.device("cpu") inv_freq = 1.0 / ( self.base ** (torch.arange(0, self.dim, 2, dtype=torch.float32, device="cpu") / self.dim) ) self.inv_freq = inv_freq.to(device) self._set_cos_sin_cache( self.max_position_embeddings, device, torch.float32, ) def rotate_half(x: torch.Tensor) -> torch.Tensor: x1 = x[...,: x.shape[-1] // 2] x2 = x[..., x.shape[-1] // 2:] return torch.cat((-x2, x1), dim=-1) def apply_rotary_pos_emb( q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, position_ids: Optional[torch.Tensor] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: if position_ids is None: cos = cos.unsqueeze(0).unsqueeze(0) sin = sin.unsqueeze(0).unsqueeze(0) else: cos = cos[position_ids].unsqueeze(1) sin = sin[position_ids].unsqueeze(1) return ( (q * cos) + (rotate_half(q) * sin), (k * cos) + (rotate_half(k) * sin), ) class GroupedQueryAttention(nn.Module): def __init__(self, config: ArgonneConfig) -> None: super().__init__() self.hidden_size = config.hidden_size self.num_heads = config.num_attention_heads self.num_kv_heads = config.num_key_value_heads self.head_dim = self.hidden_size // self.num_heads self.num_key_value_groups = self.num_heads // self.num_kv_heads self.sliding_window = config.sliding_window self.qk_norm = config.qk_norm self.v_norm = config.v_norm self.q_proj = nn.Linear( self.hidden_size, self.num_heads * self.head_dim, bias=config.attention_bias, ) self.k_proj = nn.Linear( self.hidden_size, self.num_kv_heads * self.head_dim, bias=config.attention_bias, ) self.v_proj = nn.Linear( self.hidden_size, self.num_kv_heads * self.head_dim, bias=config.attention_bias, ) self.o_proj = nn.Linear( self.num_heads * self.head_dim, self.hidden_size, bias=config.attention_bias, ) self.o_proj._is_residual = True # argonne4.5: gated attention (Muse Glimmer's `self_attn.gate_proj`). A per-(head, position) # sigmoid gate on the attention output before o_proj. Cheap, and reported to suppress # attention sinks / stabilize high-LR training -- the regime [[argonne-next-arch-search]] # already put us in ("qk_norm is ESSENTIAL at this LR"). self.attn_gate = getattr(config, "attn_gate", False) if self.attn_gate: self.gate_proj = nn.Linear( self.hidden_size, self.num_heads * self.head_dim, bias=config.attention_bias, ) # Set per-layer by ArgonneModel when attn_pattern/nope_global are in play. False => this # layer is NoPE (no rotary applied at all), which is what Muse does on its global layers. self.use_rope = True self.attention_dropout = config.attention_dropout self.use_flash_attention = config.use_flash_attention if self.qk_norm: self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) if self.v_norm: self.v_norm_layer = RMSNorm(self.head_dim, eps=config.rms_norm_eps) def _repeat_kv(self, x: torch.Tensor) -> torch.Tensor: if self.num_key_value_groups == 1: return x bsz, num_kv, seqlen, head_dim = x.shape x = x[:,:, None,:,:].expand(bsz, num_kv, self.num_key_value_groups, seqlen, head_dim) return x.reshape(bsz, num_kv * self.num_key_value_groups, seqlen, head_dim) def forward( self, hidden_states: torch.Tensor, position_embeddings: Tuple[torch.Tensor, torch.Tensor], attention_mask: Optional[torch.Tensor] = None, past_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, use_cache: bool = False, ): bsz, seqlen, _ = hidden_states.shape query = self.q_proj(hidden_states) key = self.k_proj(hidden_states) value = self.v_proj(hidden_states) query = query.view(bsz, seqlen, self.num_heads, self.head_dim).transpose(1, 2) key = key.view(bsz, seqlen, self.num_kv_heads, self.head_dim).transpose(1, 2) value = value.view(bsz, seqlen, self.num_kv_heads, self.head_dim).transpose(1, 2) if self.qk_norm: query = self.q_norm(query) key = self.k_norm(key) if self.v_norm: value = self.v_norm_layer(value) if self.use_rope and position_embeddings is not None: cos, sin = position_embeddings query, key = apply_rotary_pos_emb(query, key, cos, sin) # ---- KV cache (inference only; training keeps past_kv=None/use_cache=False) ---- # Cache pre-repeat (num_kv_heads), post-rotary K/V. The SDPA/math path # ignores the sliding window (full causal), so storing all keys and having # the new query attend to everything reproduces the no-cache result. is_decode = past_kv is not None if is_decode: past_k, past_v = past_kv key = torch.cat([past_k, key], dim=2) value = torch.cat([past_v, value], dim=2) new_kv = (key, value) if use_cache else None kv_len = key.shape[2] # Additive mask used only when decoding with a cache: query at absolute # position (kv_len - seqlen + i) attends to keys j <= that position. decode_mask = None if is_decode: q_pos = torch.arange(kv_len - seqlen, kv_len, device=hidden_states.device) k_pos = torch.arange(kv_len, device=hidden_states.device) allowed = k_pos[None,:] <= q_pos[:, None] decode_mask = torch.zeros(seqlen, kv_len, dtype=query.dtype, device=hidden_states.device) decode_mask = decode_mask.masked_fill(~allowed, -65504.0)[None, None] key = self._repeat_kv(key) value = self._repeat_kv(value) # Argonne-3.5 speedup (hardware accel, training-neutral): qk/v-norm return fp32 # (fp32 norm-weight * bf16 activation type-promotes), which trips the dtype gate below # and forces the hand-rolled O(seq^2) math attention. Under bf16 autocast the matmuls # already run in bf16, so casting q/k/v to the autocast dtype here is numerically # equivalent (measured rel diff 3.7e-3 < bf16 roundoff 7.8e-3) and lets SDPA pick a fused # FLASH/cuDNN kernel: ~2x lower attention memory + faster. Training-under-autocast only; # inference/decode (self.training False) is untouched. if self.training and torch.is_autocast_enabled() and query.dtype not in (torch.float16, torch.bfloat16): _ac_dtype = torch.get_autocast_dtype(query.device.type) if hasattr(torch, "get_autocast_dtype") else torch.get_autocast_gpu_dtype() query = query.to(_ac_dtype) key = key.to(_ac_dtype) value = value.to(_ac_dtype) use_flash_attn_2 = ( not is_decode and _flash_attn_available and self.use_flash_attention and attention_mask is None and query.dtype in (torch.float16, torch.bfloat16) and self.head_dim % 4 == 0 ) use_scaled_dot = ( hasattr(F, "scaled_dot_product_attention") and self.use_flash_attention and query.dtype in (torch.float16, torch.bfloat16) and self.head_dim % 4 == 0 ) attn_output = None if use_flash_attn_2: try: flash_dropout = self.attention_dropout if self.training else 0.0 window = ( (self.sliding_window, self.sliding_window) if self.sliding_window is not None else (-1, -1) ) q = query.transpose(1, 2).contiguous() k = key.transpose(1, 2).contiguous() v = value.transpose(1, 2).contiguous() attn_output = flash_attn_func( q, k, v, dropout_p=flash_dropout, softmax_scale=None, causal=True, window_size=window, ).transpose(1, 2) except RuntimeError: attn_output = None if attn_output is None and use_scaled_dot: try: if decode_mask is not None: attn_output = F.scaled_dot_product_attention( query, key, value, attn_mask=decode_mask, dropout_p=0.0, is_causal=False, ) elif attention_mask is None and self.sliding_window is not None and not is_decode: # argonne4.5 FIX (Finding A): flash-attn-2 is the ONLY path above that honors # `window_size`, and this env ships flash-attn-4 without the 2.x # `flash_attn_interface` module -- so `_flash_attn_available` is False and every # Argonne pretrain since 3.0 silently ran FULL attention on every layer while # pretrain.py advertised LOCAL_ATTENTION_WINDOW=256. experiments/model.py got # this fix during the arch sweep; production never did. Apply the window # explicitly here so the configured layout is the layout that trains. # Token i attends to keys in [i-w, i]. qi = torch.arange(seqlen, device=hidden_states.device) ki = torch.arange(kv_len, device=hidden_states.device) win_mask = (ki[None,:] <= qi[:, None]) & ( ki[None,:] >= qi[:, None] - self.sliding_window) attn_output = F.scaled_dot_product_attention( query, key, value, attn_mask=win_mask[None, None], dropout_p=self.attention_dropout if self.training else 0.0, is_causal=False, ) elif attention_mask is None: # Prefill / training: fast causal path. attn_output = F.scaled_dot_product_attention( query, key, value, attn_mask=None, dropout_p=self.attention_dropout if self.training else 0.0, is_causal=True, ) else: # With attention_mask: need to pass it explicitly (slower but correct) # attention_mask should be 4D: (bsz, 1, seq, seq) or broadcastable attn_output = F.scaled_dot_product_attention( query, key, value, attn_mask=attention_mask, dropout_p=self.attention_dropout if self.training else 0.0, is_causal=False, # Mask already includes causal component ) except _OOM_ERRORS: # Do NOT degrade to math attention on OOM. torch.OutOfMemoryError subclasses # RuntimeError, so the old blanket `except RuntimeError` turned a recoverable SDPA # OOM into the seq^2 math path -- which needs strictly MORE memory and then dies # for real. At phase-C ctx 65,536 that fallback asked for a single 48 GiB tensor # and killed the slice. Surface the original OOM instead; the caller can retry # smaller. The fallback below stays for genuine kernel-unavailability. raise except RuntimeError: # Fallback to math attention when kernels are unavailable attn_output = None if attn_output is None: scores = torch.matmul(query, key.transpose(2, 3)) / math.sqrt(self.head_dim) mask_value = -65504.0 # Large negative instead of -inf if decode_mask is not None: scores = scores + decode_mask else: # Apply causal mask - use large negative instead of -inf for stability causal_mask = torch.triu( torch.ones(seqlen, kv_len, dtype=torch.bool, device=hidden_states.device), diagonal=1, ) scores = scores.masked_fill(causal_mask, mask_value) # argonne4.5: keep the math fallback consistent with the SDPA path above -- # it must apply the sliding window too, or the two paths train different models. if self.sliding_window is not None: qi = torch.arange(seqlen, device=hidden_states.device) ki = torch.arange(kv_len, device=hidden_states.device) too_old = ki[None,:] < qi[:, None] - self.sliding_window scores = scores.masked_fill(too_old, mask_value) if attention_mask is not None: # Bool masks (argonne4.5 document masking) are True == attend; the legacy # float mask is additive. if attention_mask.dtype == torch.bool: scores = scores.masked_fill(~attention_mask, mask_value) else: scores = scores + attention_mask attn_weights = torch.softmax(scores, dim=-1, dtype=torch.float32).to(query.dtype) attn_weights = F.dropout(attn_weights, p=self.attention_dropout, training=self.training) attn_output = torch.matmul(attn_weights, value) attn_output = ( attn_output.transpose(1, 2) .contiguous() .view(bsz, seqlen, self.num_heads * self.head_dim) ) if self.attn_gate: attn_output = attn_output * torch.sigmoid(self.gate_proj(hidden_states)) out = self.o_proj(attn_output) if use_cache: return out, new_kv return out class SwiGLUMLP(nn.Module): def __init__(self, config: ArgonneConfig) -> None: super().__init__() self.gate_proj = nn.Linear( config.hidden_size, config.intermediate_size, bias=config.mlp_bias, ) self.up_proj = nn.Linear( config.hidden_size, config.intermediate_size, bias=config.mlp_bias, ) self.down_proj = nn.Linear( config.intermediate_size, config.hidden_size, bias=config.mlp_bias, ) self.down_proj._is_residual = True self.dropout = nn.Dropout(config.hidden_dropout) def forward(self, x: torch.Tensor) -> torch.Tensor: # Clamp intermediate values to prevent overflow gate = self.gate_proj(x) gate = torch.clamp(gate, min=-65504.0, max=65504.0) up = self.up_proj(x) up = torch.clamp(up, min=-65504.0, max=65504.0) return self.dropout(self.down_proj(F.silu(gate) * up)) class ReLU2MLP(nn.Module): """Two-matrix ReLU-squared FFN (Nemotron 3.5 Lightning's `mlp_hidden_act: relu2`). Drops SwiGLU's `gate_proj` entirely: up -> ReLU(x)^2 -> down. Two matrices instead of three, so at equal intermediate width it is 1/3 cheaper in params and FLOPs, and at equal params the FFN can be ~50% wider. This is what makes Nemotron's 128 experts/layer affordable. The Argonne arch sweep tested SwiGLU *ratios* (2.0x worse, 2.75x optimal, 3.5x neutral) but never the activation family, so this axis is unmeasured here. """ def __init__(self, config: ArgonneConfig) -> None: super().__init__() self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=config.mlp_bias) self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=config.mlp_bias) self.down_proj._is_residual = True self.dropout = nn.Dropout(config.hidden_dropout) def forward(self, x: torch.Tensor) -> torch.Tensor: h = self.up_proj(x) h = torch.clamp(h, min=-65504.0, max=65504.0) h = F.relu(h).pow(2) # ReLU^2 squares its input, so a large activation overflows bf16 far more easily than # SwiGLU's bounded silu gate does. Re-clamp before the down projection. h = torch.clamp(h, min=-65504.0, max=65504.0) return self.dropout(self.down_proj(h)) def build_mlp(config: ArgonneConfig) -> nn.Module: return ReLU2MLP(config) if getattr(config, "mlp_type", "swiglu") == "relu2" else SwiGLUMLP(config) class Block(nn.Module): """Transformer block with GQA attention and SwiGLU feed-forward.""" def __init__(self, config: ArgonneConfig, layer_idx: int = 0) -> None: super().__init__() self.layer_idx = layer_idx self.attn = GroupedQueryAttention(config) self.input_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.post_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.sandwich_norm = config.sandwich_norm if self.sandwich_norm: self.attn_out_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.mlp_out_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.mlp = build_mlp(config) def forward( self, hidden_states: torch.Tensor, position_embeddings: Tuple[torch.Tensor, torch.Tensor], attention_mask: Optional[torch.Tensor] = None, past_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, use_cache: bool = False, ): residual = hidden_states hidden_states = self.input_norm(hidden_states) attn_out = self.attn( hidden_states, position_embeddings, attention_mask, past_kv=past_kv, use_cache=use_cache, ) new_kv = None if use_cache: hidden_states, new_kv = attn_out else: hidden_states = attn_out if self.sandwich_norm: hidden_states = self.attn_out_norm(hidden_states) hidden_states = residual + hidden_states residual = hidden_states hidden_states = self.post_norm(hidden_states) hidden_states = self.mlp(hidden_states) if self.sandwich_norm: hidden_states = self.mlp_out_norm(hidden_states) hidden_states = residual + hidden_states if use_cache: return hidden_states, new_kv return hidden_states class MTPModule(nn.Module): """One multi-token-prediction depth, DeepSeek-V3 / Nemotron-3.5-Lightning style. This is NOT what `config.mtp_horizon` does. That path re-uses the SAME lm_head on the SAME hidden state with shifted labels, so the auxiliary objective competes with t+1 for the trunk's capacity -- which is exactly what exp045 measured (horizon 2, weight 0.3, 450 steps: 3.8513 -> 3.9346, +0.083 LOSS). A real MTP module instead owns its own parameters: h'_i = eh_proj([ enorm(emb(t_{i+k})) ; hnorm(h_i) ]) -> Block -> final_norm -> lm_head so the trunk's next-token objective is untouched and the extra signal is additive. It also yields a speculative-decoding drafter for free, which matters here because the deployable reasoning recipe is self-consistency at K=8-32: a 2-3x decode speedup is 2-3x the samples per GPU-hour on every downstream experiment. """ def __init__(self, config: ArgonneConfig, layer_idx: int = 0) -> None: super().__init__() self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.eh_proj = nn.Linear(2 * config.hidden_size, config.hidden_size, bias=False) self.block = Block(config, layer_idx) self.final_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) def forward(self, hidden_states, next_token_embeds, position_embeddings, attention_mask=None): merged = torch.cat( [self.enorm(next_token_embeds), self.hnorm(hidden_states)], dim=-1 ) h = self.eh_proj(merged) return self.final_norm(self.block(h, position_embeddings, attention_mask)) class ArgonneModel(PreTrainedModel): config_class = ArgonneConfig _no_split_modules = ["Block"] _tied_weights_keys = {"lm_head.weight": "embed_tokens.weight"} def __init__(self, config: ArgonneConfig) -> None: super().__init__(config) self.config = config self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size) self.blocks = nn.ModuleList([Block(config, idx) for idx in range(config.num_hidden_layers)]) # argonne4.5 layout: an explicit 'L'/'G' pattern (e.g. Muse Glimmer's "LLLG" = 3 # sliding@2048: 1 full) takes precedence over the legacy odd/even interleave. On 'G' # layers `nope_global` removes RoPE outright -- Muse sets layer_rope_theta=0 on exactly # those layers, so the only layers that see the whole sequence carry no frequency basis # to extrapolate incorrectly past the training block. pattern = getattr(config, "attn_pattern", None) if pattern: for idx, block in enumerate(self.blocks): is_local = pattern[idx % len(pattern)] == "L" block.attn.sliding_window = config.sliding_window_size if is_local else None block.attn.use_rope = is_local or not config.nope_global elif config.interleaved_local_attention and config.local_attention_window is not None: for idx, block in enumerate(self.blocks): block.attn.sliding_window = config.local_attention_window if (idx % 2 == 1) else None self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.rotary_emb = RotaryEmbedding( config.hidden_size // config.num_attention_heads, max_position_embeddings=config.max_position_embeddings, base=config.rope_theta, ) self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) # Real MTP heads (config.mtp_module_layers > 0). Depth k predicts token t+1+k from the # previous depth's state plus the embedding of the token that state predicts. They share # lm_head with the trunk (DeepSeek-V3 does the same) so the cost is the blocks only. self.mtp_modules = nn.ModuleList( [ MTPModule(config, config.num_hidden_layers + k) for k in range(getattr(config, "mtp_module_layers", 0)) ] ) if config.tie_word_embeddings: self.lm_head.weight = self.embed_tokens.weight self.gradient_checkpointing = config.use_gradient_checkpointing # Selective activation checkpointing (ported from argonne4.0). Runtime-only -- # NOT part of config or state_dict, so checkpoint/resume compatibility with every # existing argonne3.5 checkpoint is unaffected. checkpoint_stride S: # S == 1: checkpoint ALL layers (default; identical to prior behavior). # S >= 2: checkpoint every layer EXCEPT store (un-checkpoint) every Sth layer # -> store ceil(n_layers/S), recompute the rest. Smaller S stores MORE # (more HBM, less recompute, faster); too-small S OOMs. # Numerically identical either way: it only chooses store-vs-recompute. self.checkpoint_stride = 1 self._nan_loss_count = 0 self.post_init() global _attention_path_logged rank = os.environ.get("RANK") or os.environ.get("SLURM_PROCID") or "0" if not _attention_path_logged and rank == "0": _attention_path_logged = True # Report what the model ACTUALLY does, by reading the per-block state that was just # assigned above -- never by re-deriving it from config. The old version hardcoded # "full attention ... IGNORED on this path" whenever flash-attn was missing, which # became a lie the moment the argonne4.5 SDPA fix (Finding A) taught the SDPA path to # honor the window: a4.5 trained 12 of 24 layers at window 256 while this line # insisted attention was full. A banner that can disagree with the model is worse # than no banner. windows = sorted({b.attn.sliding_window for b in self.blocks if b.attn.sliding_window is not None}) n_local = sum(1 for b in self.blocks if b.attn.sliding_window is not None) if _flash_attn_available and config.use_flash_attention: backend = "flash-attn-2" else: backend = "SDPA/math ({})".format( "flash-attn unavailable" if config.use_flash_attention else "use_flash_attention=False" ) if n_local == 0: layout = "full attention on all {} layers".format(len(self.blocks)) else: sizes = ",".join(str(w) for w in windows) layout = ( f"sliding window {sizes} ACTIVE on {n_local}/{len(self.blocks)} layers" f" (pattern={config.attn_pattern or 'odd-layer legacy'})" ) print(f"[ArgonneModel] attention path: {backend}; {layout}", flush=True) def get_input_embeddings(self) -> nn.Embedding: return self.embed_tokens def set_input_embeddings(self, new_embeddings: nn.Embedding) -> None: self.embed_tokens = new_embeddings self.config.vocab_size = new_embeddings.num_embeddings if self.config.tie_word_embeddings: self.lm_head.weight = self.embed_tokens.weight def get_output_embeddings(self) -> nn.Module: return self.lm_head def set_output_embeddings(self, new_embeddings: nn.Module) -> None: self.lm_head = new_embeddings if isinstance(new_embeddings, nn.Linear): self.config.vocab_size = new_embeddings.out_features def tie_weights(self, **kwargs) -> None: if self.config.tie_word_embeddings: self.lm_head.weight = self.embed_tokens.weight @classmethod def from_pretrained(cls, *args, **kwargs): """Wrap PreTrainedModel.from_pretrained to self-heal after loading. from_pretrained may materialize the model via the meta device (low_cpu_mem_usage), which fills the rotary embedding's non-persistent buffers with uninitialized memory and can drop the embed_tokens/lm_head tie. Callers previously had to remember to call init_buffers() and tie_weights() by hand; do it here so every loader gets a working model. """ model = super().from_pretrained(*args, **kwargs) for module in model.modules(): if isinstance(module, RotaryEmbedding): module.init_buffers() model.tie_weights() return model def _init_weights(self, module: nn.Module) -> None: if isinstance(module, nn.Linear): std = self.config.hidden_size ** -0.5 if hasattr(module, "_is_residual"): std = (2 * self.config.num_hidden_layers) ** -0.5 nn.init.normal_(module.weight, mean=0.0, std=std) 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=self.config.hidden_size ** -0.5) def set_gradient_checkpointing(self, enabled: bool = True) -> None: self.gradient_checkpointing = enabled def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs=None) -> None: self.set_gradient_checkpointing(True) def gradient_checkpointing_disable(self) -> None: self.set_gradient_checkpointing(False) def forward( self, input_ids: torch.LongTensor, attention_mask: Optional[torch.Tensor] = None, labels: Optional[torch.LongTensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[list] = None, use_cache: bool = False, document_ids: Optional[torch.LongTensor] = None, **kwargs, # Accept extra args from newer transformers (e.g., num_items_in_batch) ) -> CausalLMOutput: _, seq_length = input_ids.shape device = self.embed_tokens.weight.device if input_ids.device != device: input_ids = input_ids.to(device) hidden_states = self.embed_tokens(input_ids) # The training path does not use attention masks. attention_mask = None # argonne4.5 document masking. The production loaders pack many documents into one flat # contiguous window, and this line used to be the whole story -- so every token attended # freely across document boundaries. That is worst on exactly the target domains (short # math problems, short code files), where a block-1024 window spans many documents. When # `doc_mask` is on and per-token document ids are supplied, build a block-diagonal mask # per DISTINCT window size (at most two: sliding and global) and share it across layers. # Costs one (B, T, T) bool per distinct window; without document_ids nothing changes. doc_masks = None if getattr(self.config, "doc_mask", False) and document_ids is not None and not use_cache: doc_masks = self._build_doc_masks(document_ids.to(device), seq_length) # RoPE positions are offset by the cached length so incremental decode # uses the correct absolute positions. past_len = past_key_values[0][0].shape[2] if past_key_values else 0 cos_full, sin_full = self.rotary_emb(hidden_states, past_len + seq_length) rotary = (cos_full[past_len:past_len + seq_length], sin_full[past_len:past_len + seq_length]) new_cache = [] if use_cache else None for i, layer in enumerate(self.blocks): _skip_ckpt = (self.checkpoint_stride >= 2) and (i % self.checkpoint_stride == 0) layer_mask = ( doc_masks[layer.attn.sliding_window] if doc_masks is not None else attention_mask ) if self.gradient_checkpointing and self.training and not _skip_ckpt: hidden_states = torch.utils.checkpoint.checkpoint( layer, hidden_states, rotary, layer_mask, use_reentrant=False, ) elif use_cache: past_kv = past_key_values[i] if past_key_values else None hidden_states, layer_kv = layer( hidden_states, rotary, layer_mask, past_kv=past_kv, use_cache=True, ) new_cache.append(layer_kv) else: hidden_states = layer(hidden_states, rotary, layer_mask) hidden_states = self.norm(hidden_states) # Chunked cross-entropy (training only, flag-gated via config.loss_chunk_size): avoids # materializing the full (batch*seq, vocab) fp32 logit tensor -- the HBM wall at long context. # Early-return so the standard full-logit path below stays byte-for-byte unchanged (zero risk to # non-chunked / default runs). Only taken when MTP + z-loss are off (both default off). _chunk = getattr(self.config, "loss_chunk_size", 0) if (self.training and labels is not None and _chunk > 0 and self.config.mtp_horizon == 1 and self.config.z_loss_weight == 0 and len(self.mtp_modules) == 0): loss = self._chunked_lm_loss(hidden_states, labels, _chunk) if torch.isnan(loss): self._nan_loss_count += 1 if self._nan_loss_count <= 5 or self._nan_loss_count % 100 == 0: print(f"WARNING: NaN loss detected (occurrence {self._nan_loss_count}); " "zeroing this step's loss.", flush=True) zero = torch.zeros((), device=loss.device, dtype=torch.float32) for p in self.parameters(): if p.requires_grad: zero = zero + torch.nan_to_num(p).sum().float() * 0.0 loss = zero.to(loss.dtype) return CausalLMOutput(logits=None, loss=loss) logits = self.lm_head(hidden_states) if self.config.logit_softcap > 0: cap = self.config.logit_softcap logits = torch.tanh(logits / cap) * cap # Sanitize logits only at inference. During training the full-tensor # scan forces a GPU sync every step over [batch, seq, vocab]; a NaN # would surface in the loss check below anyway. if not self.training and torch.isnan(logits).any(): logits = torch.nan_to_num(logits, nan=0.0, posinf=65504.0, neginf=-65504.0) loss = None if labels is not None: if labels.device != logits.device: labels = labels.to(logits.device) loss = F.cross_entropy( logits.view(-1, logits.size(-1)), labels.view(-1), ignore_index=-100, ) if self.training and self.config.mtp_horizon > 1: mtp_terms = [] max_horizon = min(self.config.mtp_horizon, labels.shape[1]) for horizon in range(2, max_horizon + 1): shift = horizon - 1 shifted_logits = logits[:,:-shift,:] shifted_labels = labels[:, shift:] if shifted_logits.numel() == 0: continue mtp_loss = F.cross_entropy( shifted_logits.reshape(-1, shifted_logits.size(-1)), shifted_labels.reshape(-1), ignore_index=-100, ) mtp_terms.append(mtp_loss / horizon) if mtp_terms: loss = loss + (self.config.mtp_loss_weight * torch.stack(mtp_terms).mean()).to(loss.dtype) if self.training and len(self.mtp_modules) > 0 and self.config.mtp_loss_weight > 0: loss = loss + ( self.config.mtp_loss_weight * self._mtp_module_loss(hidden_states, labels, rotary) ).to(loss.dtype) if self.config.z_loss_weight > 0: z = torch.logsumexp(logits.float(), dim=-1) loss = loss + (self.config.z_loss_weight * z.pow(2).mean()).to(loss.dtype) # A NaN loss becomes a zero loss that is still connected to every # parameter: but through the parameters directly, NOT through the # network's (possibly NaN) activations. Backward then delivers an # exact zero gradient to each parameter: DDP/FSDP gradient hooks # all fire (no collective desync/hang) and the optimizer step is a # true no-op. Warn (bounded) instead of hiding it. if torch.isnan(loss): self._nan_loss_count += 1 if self._nan_loss_count <= 5 or self._nan_loss_count % 100 == 0: print( f"WARNING: NaN loss detected (occurrence {self._nan_loss_count}); " "zeroing this step's loss.", flush=True, ) zero = torch.zeros((), device=loss.device, dtype=torch.float32) for p in self.parameters(): if p.requires_grad: zero = zero + torch.nan_to_num(p).sum().float() * 0.0 loss = zero.to(loss.dtype) if use_cache: return CausalLMOutputWithPast(logits=logits, loss=loss, past_key_values=new_cache) return CausalLMOutput(logits=logits, loss=loss) def _build_doc_masks(self, document_ids, seq_length): """{sliding_window -> bool mask (B, 1, T, T), True == attend}, one per distinct window. Combines causal, the layer's sliding window, and the document boundary, so a token can only see earlier tokens of its OWN packed document. Built once per forward and shared across every layer with the same window (at most two under an 'LLLG' layout). """ qi = torch.arange(seq_length, device=document_ids.device) same_doc = document_ids[:,:, None] == document_ids[:, None,:] # (B, T, T) base = same_doc & (qi[:, None] >= qi[None,:]) masks = {} for window in {block.attn.sliding_window for block in self.blocks}: if window is None: masks[None] = base.unsqueeze(1) else: masks[window] = (base & (qi[None,:] >= qi[:, None] - window)).unsqueeze(1) return masks def _mtp_module_loss(self, hidden_states, labels, rotary): """Mean CE over the real MTP depths (see MTPModule). Depth k reads the previous depth's state at position i plus the embedding of token t+k (the token that state predicts) and predicts t+k+1, so each depth is truncated by one more position. `labels[:, i]` is already token i+1 (the loaders hand back y = x shifted), hence token t+k at position i is `labels[:, i + k - 1]`. """ terms = [] h_prev = hidden_states cap = float(self.config.logit_softcap) seq_len = labels.shape[1] batch = labels.shape[0] # fp8 `_scaled_mm` requires the flattened row count (batch * length) to be divisible by 16. # Truncating by depth k breaks that -- at k=1, 8*2047 = 16,376 is not -- and the lm_head GEMM # inside this loss then dies under torchao float8. Round the usable length DOWN so batch*L # stays 16-aligned; that discards at most 15 of ~2047 positions, which is negligible against # the alignment requirement. Costs nothing when fp8 is off. align = 16 // math.gcd(batch, 16) for k, module in enumerate(self.mtp_modules, start=1): usable = ((seq_len - k) // align) * align if usable < 1: break h_in = h_prev[:,:usable,:] tokens_in = labels[:, k - 1: k - 1 + usable] targets = labels[:, k: k + usable] # Ignored positions (-100) would index the embedding out of range; their targets are # ignored by cross_entropy anyway, so any in-range id works as a placeholder. embeds_in = self.embed_tokens(tokens_in.clamp_min(0)) rot_k = (rotary[0][:usable], rotary[1][:usable]) h_out = module(h_in, embeds_in, rot_k) # ⛔ DO NOT "optimise" this with _chunked_lm_loss -- MEASURED WORSE, 2026-08-27. # Real MTP (mtp_module_layers=1) is a genuine iso-token quality win (tgt -0.058, code # -0.118, both seeds, every domain) but costs -42% throughput (5,348 -> 3,082 tok/s), so # iso-compute it LOSES by 0.27-0.33: the same wall buys 1.73x more tokens without it. # I hypothesised the -42% was this line materialising [T, 151680] logits per depth, and # routed it through _chunked_lm_loss (equivalence proven to 9.5e-7). Result: tok/s went # 3,082 -> 2,723, i.e. cost rose to -49.1%, at HBM 0.960 with quality unchanged. # ⇒ The transient was NOT the bottleneck; chunking's recompute overhead is pure loss here. # The cost is INHERENT: each depth runs a full extra block AND a 151,680-wide lm_head GEMM # in forward and backward. MTP is rejected on cost, not on an artifact -- do not retry # this optimisation. mtp_logits = self.lm_head(h_out) if cap > 0: mtp_logits = torch.tanh(mtp_logits / cap) * cap terms.append( F.cross_entropy( mtp_logits.reshape(-1, mtp_logits.size(-1)), targets.reshape(-1), ignore_index=-100, ) ) h_prev = h_out if not terms: return hidden_states.new_zeros(()) return torch.stack(terms).mean() @torch.compiler.disable def _chunked_lm_loss(self, hidden_states, labels, chunk_size): """Cross-entropy without ever materializing the full (batch*seq, vocab) fp32 logit tensor. MUST run eager (@torch.compiler.disable): torch.compile/inductor otherwise fuses the chunks back into the full (N, vocab) logit tensor, re-creating the OOM this is meant to avoid. The transformer body still compiles (+FP8); only this loss runs eager (a small fraction of compute). Splits the flattened rows into chunks; torch.utils.checkpoint recomputes each chunk's logits in backward, so peak logit memory = chunk_size*vocab (not N*vocab). Numerically identical to F.cross_entropy(full_logits, labels.view(-1), ignore_index=-100) with per-chunk softcap (mean over non-ignored). Caller guarantees mtp_horizon==1 and z_loss_weight==0. Uses self.lm_head(...) per chunk so the fp8 (Float8Linear) GEMM path is preserved.""" from torch.utils.checkpoint import checkpoint as _ckpt cap = float(self.config.logit_softcap) H = hidden_states.reshape(-1, hidden_states.size(-1)) L = labels.reshape(-1).to(H.device) N = H.size(0) def _seg(hc, lc): lg = self.lm_head(hc) # (chunk, vocab); Float8Linear under fp8 if cap > 0: lg = torch.tanh(lg / cap) * cap return F.cross_entropy(lg, lc, ignore_index=-100, reduction="sum") total = H.new_zeros((), dtype=torch.float32) ntok = torch.zeros((), device=H.device, dtype=torch.float32) cs = max(1, int(chunk_size)) for i in range(0, N, cs): hc = H[i:i + cs] lc = L[i:i + cs] seg = _ckpt(_seg, hc, lc, use_reentrant=False) if self.training else _seg(hc, lc) total = total + seg.float() ntok = ntok + (lc != -100).sum().float() return total / ntok.clamp(min=1.0) @torch.no_grad() def generate( self, input_ids: torch.Tensor, max_length: int = 1024, temperature: float = 1.0, top_k: Optional[int] = None, top_p: Optional[float] = None, do_sample: bool = True, repetition_penalty: float = 1.0, no_repeat_ngram_size: int = 0, eos_token_id: Optional[int] = None, ) -> torch.Tensor: self.eval() device = self.embed_tokens.weight.device input_ids = input_ids.to(device) ctx = self.config.max_position_embeddings # Stop at EOS (2026-08-02). This loop previously had NO eos check at all -- its only # exit was `max_length`, so config.eos_token_id was inert and every caller following the # model card's `.generate()` snippet ran to the token budget and then degenerated into # repetition long after the answer was finished. Harmless-looking on a base model (which # just continues text) but bad on a chat/think model, where the assistant turn must end # at <|im_end|>. Longstanding across the whole Argonne line, including the published # 3.0-think bundle -- not introduced by any recent change. # Pass eos_token_id=-1 to explicitly opt out and restore run-to-max_length behavior. if eos_token_id is None: eos_token_id = getattr(self.config, "eos_token_id", None) stop_id = None if eos_token_id is None or eos_token_id < 0 else int(eos_token_id) finished = torch.zeros(input_ids.shape[0], dtype=torch.bool, device=device) # KV-cache decode: prefill the prompt once, then feed a single token per # step reusing past_key_values. This yields identical logits to the # recompute-the-whole-prefix path (gated by verify_cache.py) but turns the # per-step cost from O(seq_len) down to O(1). When the running sequence # would exceed the context window we rebuild the cache from the truncated # window, matching the original chunk = input_ids[:, -ctx:] behavior. past = None outputs = None while input_ids.shape[1] < max_length: if past is None or past[0][0].shape[2] >= ctx: chunk = input_ids[:, -ctx:] outputs = self.forward(chunk, use_cache=True) past = outputs.past_key_values logits = outputs.logits[:, -1,:] / temperature if repetition_penalty != 1.0: for b in range(input_ids.shape[0]): seen_tokens = torch.unique(input_ids[b]) seen_logits = logits[b, seen_tokens] adjusted = torch.where( seen_logits < 0, seen_logits * repetition_penalty, seen_logits / repetition_penalty, ) logits[b, seen_tokens] = adjusted if no_repeat_ngram_size > 0 and input_ids.shape[1] + 1 >= no_repeat_ngram_size: n = int(no_repeat_ngram_size) for b in range(input_ids.shape[0]): seq = input_ids[b].tolist() ngrams = {} for i in range(len(seq) - n + 1): prefix = tuple(seq[i: i + n - 1]) if n > 1 else tuple() next_token = seq[i + n - 1] if prefix not in ngrams: ngrams[prefix] = set() ngrams[prefix].add(next_token) current_prefix = tuple(seq[-(n - 1):]) if n > 1 else tuple() banned = ngrams.get(current_prefix, set()) if banned: logits[b, list(banned)] = float("-inf") if do_sample: if top_k is not None: top_values, _ = torch.topk(logits, min(top_k, logits.size(-1))) logits = logits.masked_fill(logits < top_values[:, [-1]], float("-inf")) if top_p is not None: sorted_logits, sorted_indices = torch.sort(logits, descending=True) cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1) sorted_indices_to_remove = cumulative_probs > top_p sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[...,:-1].clone() sorted_indices_to_remove[..., 0] = 0 indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove) logits = logits.masked_fill(indices_to_remove, float("-inf")) probs = F.softmax(logits, dim=-1) next_token = torch.multinomial(probs, num_samples=1) else: next_token = torch.argmax(logits, dim=-1, keepdim=True) next_token = next_token.to(input_ids.device) input_ids = torch.cat([input_ids, next_token], dim=-1) if input_ids.shape[1] >= max_length: break # Batch-safe: only stop once EVERY row has emitted EOS, so a short sequence cannot # truncate the others. (This arch has no padding support, so batch is normally 1.) if stop_id is not None: finished |= next_token.squeeze(-1).eq(stop_id) if bool(finished.all()): break # Advance the cache by one token unless we must rebuild it next loop # (cache full); the top-of-loop guard re-prefills in that case. if past[0][0].shape[2] < ctx: outputs = self.forward(next_token, past_key_values=past, use_cache=True) past = outputs.past_key_values return input_ids.to(device) AutoConfig.register("argonne2", ArgonneConfig) AutoModel.register(ArgonneConfig, ArgonneModel) AutoModelForCausalLM.register(ArgonneConfig, ArgonneModel) # Backwards compatibility exports CausalSelfAttention = GroupedQueryAttention MLP = SwiGLUMLP