import functools import math import warnings from dataclasses import dataclass from typing import Optional, Tuple import torch import torch.nn as nn import torch.utils.checkpoint from torch.nn import functional as F try: import config as alpha_config except ImportError: # pragma: no cover - fallback for package-style imports from . import config as alpha_config try: import kernel except ImportError: # pragma: no cover - fallback for package-style imports from . import kernel try: import fused_kernels except ImportError: # pragma: no cover - fallback for package-style imports try: from . import fused_kernels except ImportError: fused_kernels = None try: import flash_ops except ImportError: # pragma: no cover - fallback for package-style imports try: from . import flash_ops except ImportError: flash_ops = None # ── nvFP4 via torchao RHT (Blackwell SM100+ only) ─────────────────────────── # Uses NVFP4DynamicActivationNVFP4WeightConfig from torchao.prototype.mx_formats. # This replaces the old manual block-quantization implementation. # Dynamic activation scaling: per-tensor FP8 scale computed at runtime. # Weight scaling: per-block-16 FP4 scale (RHT = Reciprocal-Half-Tile layout). # Both scales are fused into a single Blackwell GEMM kernel via torch.compile. # No manual weight surgery needed -- torchao rewrites nn.Linear in-place via quantize_(). _TORCHAO_AVAILABLE: bool = False _NVFP4_RHT_AVAILABLE: bool = False _INT8_AVAILABLE: bool = False try: import torchao as _torchao_probe # noqa: F401 _TORCHAO_AVAILABLE = True from torchao.quantization import quantize_ as _torchao_quantize from torchao.quantization import ( Int8DynamicActivationInt8WeightConfig as _Int8Config, ) # Per-group INT8 weight quant: available as Int8WeightOnlyConfig with # group_size, paired with dynamic int8 activations via a composed config. # Stable API: try the grouped variant; fall back to per-channel if absent. try: from torchao.quantization import Int8WeightOnlyConfig as _Int8WOConfig _INT8_GROUPED_AVAILABLE = True except ImportError: _Int8WOConfig = None # type: ignore[assignment] _INT8_GROUPED_AVAILABLE = False _INT8_AVAILABLE = True try: from torchao.prototype.mx_formats import ( NVFP4DynamicActivationNVFP4WeightConfig as _NVFP4Config, ) _NVFP4_RHT_AVAILABLE = True except ImportError: _NVFP4Config = None # type: ignore[assignment] except ImportError: _torchao_quantize = None # type: ignore[assignment] _Int8Config = None # type: ignore[assignment] _NVFP4Config = None # type: ignore[assignment] _NVFP4_DISABLED_REASON: "str | None" = None # Layers to keep in full precision. # lm_head / wte: embedding + output projection -- high vocab sensitivity. # norm / ln_: RMSNorm scale weights -- tiny tensors, quantizing them is pointless. # router / gate: MoE router logits drive routing decisions; FP4 routing = collapsed experts. # expert_bias: per-expert scalar bias, not a Linear weight. _NVFP4_SKIP_KEYWORDS: tuple[str, ...] = ( "wte", "lm_head", "ln_", "ln_f", "norm", "router", "gate", "expert_bias", ) def _nvfp4_filter(module: nn.Module, fqn: str) -> bool: """torchao filter_fn: True = quantize this module.""" if not isinstance(module, nn.Linear): return False for kw in _NVFP4_SKIP_KEYWORDS: if kw in fqn: return False # Minimum size guard: torchao's RHT kernel requires in_features % 16 == 0 # and out_features >= 64. Skip tiny linears (e.g. MoD router D->1). out_f, in_f = module.weight.shape if in_f < 32 or out_f < 64 or in_f % 16 != 0: return False return True def _nvfp4_enabled() -> bool: global _NVFP4_DISABLED_REASON if not getattr(alpha_config, 'use_nvfp4', False): return False if not torch.cuda.is_available(): _NVFP4_DISABLED_REASON = "nvFP4: no CUDA device" return False major, _minor = torch.cuda.get_device_capability() if major < 10: msg = ( f"nvFP4 requires Blackwell (sm_100+); found sm_{major}{_minor} " f"({torch.cuda.get_device_name()}) -- disabled." ) if msg not in _alpha_config_warned: _alpha_config_warned.add(msg) warnings.warn(msg) _NVFP4_DISABLED_REASON = msg return False if not _NVFP4_RHT_AVAILABLE: _NVFP4_DISABLED_REASON = ( "nvFP4: torchao or torchao.prototype.mx_formats not available. " "Install: pip install torchao --pre" ) warnings.warn(_NVFP4_DISABLED_REASON) return False return True def apply_nvfp4_to_model(model: nn.Module) -> int: """Quantize eligible Linear layers to NVFP4 (RHT, dynamic activations). Uses torchao.quantization.quantize_() with NVFP4DynamicActivationNVFP4WeightConfig, which applies: - Weight: FP4 e2m1 with per-block-16 FP8 scales (RHT layout) - Activation: FP4 with per-tensor FP32 dynamic scale (computed at runtime) The fused GEMM kernel is emitted by torch.compile on Blackwell SM100+. Must be called BEFORE torch.compile(); quantize_() rewrites modules in-place. Returns the number of quantized Linear layers. """ before = {fqn for fqn, m in model.named_modules() if _nvfp4_filter(m, fqn)} _torchao_quantize(model, _NVFP4Config(), filter_fn=_nvfp4_filter) count = 0 for fqn, m in model.named_modules(): if fqn in before and isinstance(m, nn.Module): if type(m).__name__ != "Linear": count += 1 return count # ── Jetfire-style INT8 via torchao (T4 / Turing sm_75 and any CUDA device) ── # Int8DynamicActivationInt8WeightConfig applies: # Weight: INT8 symmetric per-group quantization (group_size from config). # group_size=128 is the Jetfire default; smaller (32) tightens # outlier containment at the cost of more scale overhead. # Activation: INT8 dynamic per-tensor scale computed at runtime (no calibration). # Inductor fuses the INT8 GEMM + dequant into a single Triton or CUTLASS kernel. # Works on T4 (sm_75) and any device with INT8 tensor core support (Turing+). # Much cheaper than FP16 GEMM on T4 where INT8 tensor cores run at 2× the rate. # Reuse same skip list as FP4 -- same sensitive layers apply. _INT8_SKIP_KEYWORDS: tuple[str, ...] = _NVFP4_SKIP_KEYWORDS _INT8_DISABLED_REASON: "str | None" = None def _int8_filter(module: nn.Module, fqn: str) -> bool: """torchao filter_fn for INT8: True = quantize this module.""" if not isinstance(module, nn.Linear): return False for kw in _INT8_SKIP_KEYWORDS: if kw in fqn: return False # group_size constraint: in_features must be divisible by group_size. # We can't check the configured group_size here without a closure, so # enforce the stricter of the two defaults (128). apply_int8_to_model # re-checks with the actual group_size before calling quantize_(). out_f, in_f = module.weight.shape if in_f < 64 or out_f < 32: return False return True def _int8_enabled() -> bool: global _INT8_DISABLED_REASON if not getattr(alpha_config, 'use_int8', False): return False if not torch.cuda.is_available(): _INT8_DISABLED_REASON = "INT8: no CUDA device" return False if not _INT8_AVAILABLE: _INT8_DISABLED_REASON = ( "INT8: torchao not installed or Int8DynamicActivationInt8WeightConfig " "not found. Install: pip install torchao" ) warnings.warn(_INT8_DISABLED_REASON) return False # Mutual exclusion: nvFP4 and INT8 cannot both be active. # nvFP4 only runs on SM100+; if the user enabled both on a Blackwell machine # nvFP4 takes priority (higher throughput). if _nvfp4_enabled(): msg = ( "INT8 disabled: use_nvfp4=True also set and this is SM100+ hardware -- " "nvFP4 takes priority. Set use_nvfp4=False to use INT8 instead." ) if msg not in _alpha_config_warned: _alpha_config_warned.add(msg) warnings.warn(msg) _INT8_DISABLED_REASON = msg return False return True def apply_int8_to_model(model: nn.Module, group_size: int = 128) -> int: """Quantize eligible Linear layers to INT8 (Jetfire-style, dynamic activations). Config selection priority: 1. Int8DynamicActivationInt8WeightConfig() -- stable torchao, per-channel weights + dynamic per-tensor INT8 activations. Best for T4 throughput. 2. Int8WeightOnlyConfig(group_size=N) -- weight-only, activations stay fp16. Only used as fallback if (1) is unavailable (shouldn't happen on stable). Works on T4 (sm_75+). Fused kernel emitted by torch.compile via Inductor. Must be called BEFORE torch.compile(). """ def _filter_with_gs(module: nn.Module, fqn: str) -> bool: if not _int8_filter(module, fqn): return False # group_size divisibility only matters for grouped weight configs; # Int8DynamicActivationInt8WeightConfig is per-channel so no constraint. # Keep the check anyway to skip genuinely misaligned tiny layers. in_f = module.weight.shape[1] if in_f % 32 != 0: # 32 = minimum alignment for any INT8 kernel return False return True # Snapshot module ids before -- torchao replaces nn.Linear with a subclass # in-place; the fqn stays the same but id() changes. Class-name check is # unreliable because torchao's subclass is also called "Linear" in some versions. before_ids = {fqn: id(m) for fqn, m in model.named_modules() if _filter_with_gs(m, fqn)} # Prefer Int8DynamicActivationInt8WeightConfig: dynamic INT8 on both weight # and activation, fused into a single CUTLASS/Triton INT8 GEMM by Inductor. # This is the correct stable API -- no constructor args needed. if _INT8_AVAILABLE and _Int8Config is not None: config_obj = _Int8Config() elif _INT8_GROUPED_AVAILABLE and _Int8WOConfig is not None: warnings.warn( f"[INT8] Int8DynamicActivationInt8WeightConfig unavailable -- " f"falling back to Int8WeightOnlyConfig(group_size={group_size}). " f"Activations will stay fp16 (weight-only quant).", ) config_obj = _Int8WOConfig(group_size=group_size) else: raise RuntimeError("[INT8] No usable INT8 config found in torchao -- pip install torchao") _torchao_quantize(model, config_obj, filter_fn=_filter_with_gs) # Count by id() change: torchao replaces the module object in-place on the # parent, so the same fqn now resolves to a different object. count = sum( 1 for fqn, old_id in before_ids.items() if id(dict(model.named_modules()).get(fqn)) != old_id ) return count _alpha_config_warned: set[str] = set() def _alpha_config_attr(name: str, default): if not hasattr(alpha_config, name): if name not in _alpha_config_warned: _alpha_config_warned.add(name) warnings.warn( f"config module has no attribute '{name}' -- falling back to " f"default {default!r}. If this is unexpected, check for a " f"casing mismatch or rename in your config.py.", stacklevel=2, ) return getattr(alpha_config, name, default) try: from torch.nn.functional import rms_norm as _native_rms_norm except (ImportError, AttributeError): _native_rms_norm = None try: from liger_kernel.ops.rms_norm import LigerRMSNormFunction as _LigerRMSNormFn except ImportError: _LigerRMSNormFn = None try: from liger_kernel.ops.rope import LigerRopeFunction as _LigerRopeFn except ImportError: _LigerRopeFn = None try: from liger_kernel.transformers.fused_linear_cross_entropy import ( LigerFusedLinearCrossEntropyLoss as _LigerFusedLinearCrossEntropyLoss, ) except ImportError: _LigerFusedLinearCrossEntropyLoss = None _liger_warned: set[str] = set() def _liger_warn_once(key: str, msg: str) -> None: if key not in _liger_warned: _liger_warned.add(key) warnings.warn(f"[model.py/liger] {msg}", stacklevel=3) def _liger_cuda_gate(x: torch.Tensor) -> bool: return x.device.type == "cuda" def _use_liger_config() -> bool: return bool(_alpha_config_attr('use_liger', False)) def _use_flash_ops_master() -> bool: return bool(_alpha_config_attr('use_flash_ops', True)) def _use_flash_rope_config() -> bool: return _use_flash_ops_master() and bool(_alpha_config_attr('use_flash_rope', True)) def _use_flash_outproj_config() -> bool: return _use_flash_ops_master() and bool(_alpha_config_attr('use_flash_outproj_add_rmsnorm', True)) def _use_flash_swiglu_config() -> bool: return _use_flash_ops_master() and bool(_alpha_config_attr('use_flash_swiglu', True)) def _use_fused_add_rmsnorm_config() -> bool: return bool(_alpha_config_attr('use_fused_add_rmsnorm', True)) @torch.compiler.disable def _liger_rmsnorm_attempt(x: torch.Tensor, weight: torch.Tensor, eps: float): try: out = _LigerRMSNormFn.apply(x, weight, eps) kernel._record_backend("liger_rmsnorm_success") return out except Exception as e: # noqa: BLE001 -- must never crash training kernel._record_backend(f"liger_rmsnorm_fallback:{type(e).__name__}") _liger_warn_once( f"rmsnorm_fail_{type(e).__name__}", f"Liger RMSNorm failed on {x.device} ({e!r}) -- falling back to native/eager RMSNorm.", ) return None class RMSNorm(nn.Module): """Root Mean Square Layer Normalization (used in modern transformers)""" def __init__(self, dim, eps=1e-5): super().__init__() self.eps = eps self.weight = nn.Parameter(torch.ones(dim)) self.dim = dim def forward(self, x): if _use_liger_config() and _LigerRMSNormFn is not None and _liger_cuda_gate(x): out = _liger_rmsnorm_attempt(x, self.weight, self.eps) if out is not None: return out if _native_rms_norm is not None: return _native_rms_norm(x, (self.dim,), self.weight, self.eps) x_fp32 = x.float() rms = torch.sqrt(x_fp32.pow(2).mean(-1, keepdim=True) + self.eps) weight = self.weight.to(x.dtype) return ((x_fp32 / rms).to(x.dtype)) * weight def rotate_half(x): x1, x2 = x.chunk(2, dim=-1) return torch.cat((-x2, x1), dim=-1) @dataclass class AttentionOutput: """Uniform return type so callers can optionally inspect weights / cache cost.""" output: torch.Tensor attention_weights: Optional[torch.Tensor] = None kv_cache_bytes: Optional[int] = None selected_indices: Optional[torch.Tensor] = None def module_output_tensor(x): return x.output if isinstance(x, AttentionOutput) else x def reshape_heads(x: torch.Tensor, num_heads: int, head_dim: int) -> torch.Tensor: b, t, _ = x.shape return x.view(b, t, num_heads, head_dim).transpose(1, 2) def merge_heads(x: torch.Tensor) -> torch.Tensor: b, h, t, d = x.shape return x.transpose(1, 2).contiguous().view(b, t, h * d) def expand_kv_heads(x: torch.Tensor, num_heads: int) -> torch.Tensor: b, h_kv, t, d = x.shape if h_kv == num_heads: return x if num_heads % h_kv != 0: raise ValueError(f"num_heads ({num_heads}) must be divisible by kv heads ({h_kv})") reps = num_heads // h_kv return x.repeat_interleave(reps, dim=1) def _yarn_find_correction_dim(num_rotations, dim, base=10000.0, orig_max_pos=2048): return (dim * math.log(orig_max_pos / (num_rotations * 2 * math.pi))) / (2 * math.log(base)) def _yarn_find_correction_range(low_rot, high_rot, dim, base=10000.0, orig_max_pos=2048): low = math.floor(_yarn_find_correction_dim(low_rot, dim, base, orig_max_pos)) high = math.ceil(_yarn_find_correction_dim(high_rot, dim, base, orig_max_pos)) return max(low, 0), min(high, dim - 1) def _yarn_linear_ramp_mask(low, high, dim): if low > high: low, high = high, low if low == high: high += 0.001 # avoid a divide-by-zero when the ramp is degenerate ramp = (torch.arange(dim, dtype=torch.float32) - low) / (high - low) return torch.clamp(ramp, 0, 1) def yarn_attention_temperature(scale: float) -> float: if scale <= 1.0: return 1.0 t = 0.1 * math.log(scale) + 1.0 return 1.0 / t @functools.lru_cache(maxsize=64) def _rope_cache(seq_len: int, dim: int, device, dtype, base: float = 10000.0): inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, device=device, dtype=torch.float32) / dim)) t = torch.arange(seq_len, device=device, dtype=torch.float32) freqs = torch.outer(t, inv_freq) emb = torch.cat((freqs, freqs), dim=-1) return emb.cos().to(dtype), emb.sin().to(dtype) def build_rope_cache( seq_len: int, dim: int, base: float = 10000.0, yarn_scale: float = 1.0, yarn_orig_max_pos: int | None = None, yarn_alpha: float = 1.0, yarn_beta: float = 32.0, ) -> tuple[torch.Tensor, torch.Tensor]: if yarn_scale <= 1.0: inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) else: if yarn_orig_max_pos is None: raise ValueError( "build_rope_cache: yarn_scale > 1.0 requires yarn_orig_max_pos " "(the ORIGINAL trained context length, i.e. CONTEXT/block_size " "from config.py before scaling) to compute the NTK-by-parts " "ramp correctly -- it can't be inferred from seq_len alone, " "since seq_len here is already the SCALED target length." ) pos_freqs = base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim) inv_freq_extrapolation = 1.0 / pos_freqs # gamma=1 region: untouched inv_freq_interpolation = 1.0 / (yarn_scale * pos_freqs) # gamma=0 region: NTK/linear-interpolated low, high = _yarn_find_correction_range( yarn_alpha, yarn_beta, dim, base, yarn_orig_max_pos ) if low > high: warnings.warn( f"YaRN correction range inverted (low={low} > high={high}) " f"for dim={dim} with yarn_alpha={yarn_alpha}, yarn_beta=" f"{yarn_beta} -- these defaults were tuned for LLaMA-family " f"head_dim=128 (per Peng et al. 2023), and can invert at " f"smaller head_dim. Auto-corrected by swapping low/high so " f"the ramp stays monotonic, but the resulting correction " f"window differs from LLaMA's -- if generation quality at " f"the extended length looks off, try tuning yarn_alpha/" f"yarn_beta in config.py rather than trusting these " f"defaults blindly for this model's shape." ) inv_freq_mask = 1.0 - _yarn_linear_ramp_mask(low, high, dim // 2) inv_freq = ( inv_freq_interpolation * (1 - inv_freq_mask) + inv_freq_extrapolation * inv_freq_mask ) t = torch.arange(seq_len, dtype=torch.float32) freqs = torch.outer(t, inv_freq) emb = torch.cat((freqs, freqs), dim=-1) return emb.cos(), emb.sin() def apply_rope( x: torch.Tensor, positions: torch.Tensor | None = None, rope_dim: int | None = None, max_position: int | None = None, cos_full: torch.Tensor | None = None, sin_full: torch.Tensor | None = None, ) -> torch.Tensor: b, h, t, d = x.shape rope_dim = d if rope_dim is None else rope_dim is_sequential = positions is None if positions is None: positions = torch.arange(t, device=x.device) if cos_full is None or sin_full is None: if max_position is None: max_position = t max_pos = int(max_position) if max_position is not None else int(positions.max().item()) + 1 if positions.numel() > 0 else 1 cos_full, sin_full = _rope_cache(max_pos, rope_dim, x.device, x.dtype) if is_sequential: if t > cos_full.shape[0]: raise ValueError( f"RoPE table too small: highest requested position=" f"{t - 1}, table size={cos_full.shape[0]}. If this " f"happened during generate(), set gen_headroom in config.py " f"to however many extra positions past CONTEXT (block_size) " f"generation needs -- see GPT.__init__'s gen_headroom " f"wiring. This is separate from max_gen_tokens (a hard cap " f"on tokens produced per generate() call, unrelated to " f"RoPE table size). Default gen_headroom is 0 if unset." ) elif positions is not None and positions.numel() > 0: highest_pos = int(positions.max().item()) if highest_pos >= cos_full.shape[0]: raise ValueError( f"RoPE table too small: highest requested position=" f"{highest_pos}, table size={cos_full.shape[0]}. If this " f"happened during generate(), set gen_headroom in config.py " f"to however many extra positions past CONTEXT (block_size) " f"generation needs -- see GPT.__init__'s gen_headroom " f"wiring. This is separate from max_gen_tokens (a hard cap " f"on tokens produced per generate() call, unrelated to " f"RoPE table size). Default gen_headroom is 0 if unset." ) if is_sequential: cos = cos_full[:t].to(x.dtype).unsqueeze(0).unsqueeze(0) sin = sin_full[:t].to(x.dtype).unsqueeze(0).unsqueeze(0) else: cos = cos_full[positions].to(x.dtype).unsqueeze(0).unsqueeze(0) sin = sin_full[positions].to(x.dtype).unsqueeze(0).unsqueeze(0) x_rot, x_pass = x[..., :rope_dim], x[..., rope_dim:] x_rot = (x_rot * cos) + (rotate_half(x_rot) * sin) return torch.cat([x_rot, x_pass], dim=-1) if x_pass.shape[-1] > 0 else x_rot def apply_rope_qk( q: torch.Tensor, k: torch.Tensor, cos_full: torch.Tensor, sin_full: torch.Tensor, rope_dim: int, positions: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: b, h, t, d = q.shape hk = k.shape[1] if _use_flash_rope_config() and flash_ops is not None and q.is_cuda: cos_t = cos_full[:t].to(q.dtype) sin_t = sin_full[:t].to(q.dtype) result = flash_ops.fused_rope_qk(q, k, cos_t, sin_t, rope_dim) if result is not None: return result if ( _use_liger_config() and _LigerRopeFn is not None and _liger_cuda_gate(q) and rope_dim == d and h == hk ): result = _liger_rope_attempt(q, k, cos_full[:t], sin_full[:t]) if result is not None: return result q_rot = apply_rope(q, positions=positions, rope_dim=rope_dim, cos_full=cos_full, sin_full=sin_full) k_rot = apply_rope(k, positions=positions, rope_dim=rope_dim, cos_full=cos_full, sin_full=sin_full) return q_rot, k_rot @torch.compiler.disable def _liger_rope_attempt(q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor): try: cos = cos.to(q.dtype) sin = sin.to(q.dtype) q_rot, k_rot = _LigerRopeFn.apply(q, k, cos, sin) kernel._record_backend("liger_rope_success") return q_rot, k_rot except Exception as e: # noqa: BLE001 -- must never crash training kernel._record_backend(f"liger_rope_fallback:{type(e).__name__}") _liger_warn_once( f"rope_fail_{type(e).__name__}", f"Liger RoPE failed on {q.device} ({e!r}) -- falling back to eager RoPE.", ) return None @torch.compiler.disable def _liger_fused_ce_attempt(hidden: torch.Tensor, weight: torch.Tensor, targets: torch.Tensor): try: loss_fn = _LigerFusedLinearCrossEntropyLoss(ignore_index=-1) # Liger's fused linear+CE kernel expects the activation tensor first # and the output projection weight second; the previous order was # swapped and would pass an invalid tensor layout to the kernel. loss = loss_fn(hidden.reshape(-1, hidden.size(-1)), weight, targets.reshape(-1)) kernel._record_backend("liger_fused_ce_success") return loss except Exception as e: # noqa: BLE001 -- must never crash training kernel._record_backend(f"liger_fused_ce_fallback:{type(e).__name__}") _liger_warn_once( f"fused_ce_fail_{type(e).__name__}", f"Liger fused linear CE failed on {hidden.device} ({e!r}) -- " f"falling back to eager lm_head + F.cross_entropy.", ) return None def make_causal_mask(q_len: int, k_len: int, device) -> torch.Tensor: q_idx = torch.arange(q_len, device=device).view(q_len, 1) k_idx = torch.arange(k_len, device=device).view(1, k_len) return torch.where(k_idx <= q_idx, torch.zeros(1, device=device), torch.full((1,), -1e9, device=device)) def make_sliding_window_causal_mask(q_len: int, k_len: int, window_size: int, device) -> torch.Tensor: q_idx = torch.arange(q_len, device=device).view(q_len, 1) k_idx = torch.arange(k_len, device=device).view(1, k_len) visible = (k_idx <= q_idx) & (k_idx > q_idx - window_size) return torch.where(visible, torch.zeros(1, device=device), torch.full((1,), -1e9, device=device)) def masked_softmax(scores: torch.Tensor, mask: torch.Tensor | None, dim: int = -1) -> torch.Tensor: if mask is not None: scores = scores + mask.to(scores.dtype) return torch.softmax(scores.float(), dim=dim).to(scores.dtype) def estimate_kv_cache_bytes( batch_size, seq_len, num_kv_heads, head_dim, dtype, compression_ratio: int = 1, window_size: int | None = None, ) -> int: elem_size = torch.zeros(1, dtype=dtype).element_size() effective_len = max(1, -(-seq_len // compression_ratio)) # ceil div if window_size is not None: effective_len = min(effective_len, window_size) return int(batch_size * num_kv_heads * effective_len * head_dim * elem_size * 2) class DenseMHA(nn.Module): def __init__( self, hidden_size: int, num_heads: int, head_dim: int | None = None, num_kv_heads: int | None = None, dropout: float = 0.0, use_rope: bool = True, causal: bool = True, bias: bool = True, max_seq_len: int = 4096, window_size: int | None = None, tp_size: int = 1, rope_dim: int | None = None, yarn_scale: float = 1.0, yarn_orig_max_pos: int | None = None, yarn_alpha: float = 1.0, yarn_beta: float = 32.0, ) -> None: super().__init__() if head_dim is None: if hidden_size % num_heads != 0: raise ValueError("hidden_size must be divisible by num_heads when head_dim is omitted") head_dim = hidden_size // num_heads num_kv_heads = num_heads if num_kv_heads is None else num_kv_heads if num_heads % num_kv_heads != 0: raise ValueError(f"num_heads ({num_heads}) must be divisible by num_kv_heads ({num_kv_heads})") if num_heads % tp_size != 0: raise ValueError( f"tensor_parallel_size ({tp_size}) must divide num_heads ({num_heads}) " f"for this layer -- reduce tensor_parallel_size or increase n_head." ) local_num_heads = num_heads // tp_size if num_kv_heads < tp_size: local_num_kv_heads = num_kv_heads else: local_num_kv_heads = num_kv_heads // tp_size if num_kv_heads % tp_size == 0 else num_kv_heads self.hidden_size = hidden_size self.num_heads = local_num_heads self.num_kv_heads = local_num_kv_heads self.head_dim = head_dim self.inner_dim = local_num_heads * head_dim self.use_rope = use_rope self.causal = causal self.tp_size = tp_size self.tp_group = None self.window_size = window_size self.rope_dim = head_dim if rope_dim is None else rope_dim self._q_width = self.inner_dim self._k_width = self.num_kv_heads * head_dim self._v_width = self.num_kv_heads * head_dim self.qkv_proj = nn.Linear(hidden_size, self._q_width + self._k_width + self._v_width, bias=bias) self.out_proj = nn.Linear(self.inner_dim, hidden_size, bias=bias) self.dropout = nn.Dropout(dropout) if use_rope: rope_cos, rope_sin = build_rope_cache( max_seq_len, self.rope_dim, yarn_scale=yarn_scale, yarn_orig_max_pos=yarn_orig_max_pos, yarn_alpha=yarn_alpha, yarn_beta=yarn_beta, ) self.register_buffer("rope_cos", rope_cos, persistent=False) self.register_buffer("rope_sin", rope_sin, persistent=False) self.yarn_temp_factor = yarn_attention_temperature(yarn_scale) def forward( self, hidden_states: torch.Tensor, output_attentions: bool = False, return_kv_cache_estimate: bool = False, skip_out_proj: bool = False, positions: torch.Tensor | None = None, ) -> torch.Tensor | AttentionOutput: batch_size, seq_len, _ = hidden_states.shape qkv = self.qkv_proj(hidden_states) q_flat, k_flat, v_flat = qkv.split([self._q_width, self._k_width, self._v_width], dim=-1) q = reshape_heads(q_flat, self.num_heads, self.head_dim) k = reshape_heads(k_flat, self.num_kv_heads, self.head_dim) v = reshape_heads(v_flat, self.num_kv_heads, self.head_dim) if self.use_rope: q, k = apply_rope_qk( q, k, cos_full=self.rope_cos, sin_full=self.rope_sin, rope_dim=self.rope_dim, positions=positions, ) if self.yarn_temp_factor != 1.0: q = q * self.yarn_temp_factor if output_attentions: if self.num_kv_heads != self.num_heads: reps = self.num_heads // self.num_kv_heads k_eager = k.repeat_interleave(reps, dim=1) v_eager = v.repeat_interleave(reps, dim=1) else: k_eager, v_eager = k, v scores = torch.matmul(q, k_eager.transpose(-2, -1)) / math.sqrt(self.head_dim) mask = None if self.window_size is not None: mask = make_sliding_window_causal_mask( seq_len, seq_len, self.window_size, hidden_states.device ).view(1, 1, seq_len, seq_len) elif self.causal: mask = make_causal_mask(seq_len, seq_len, hidden_states.device).view(1, 1, seq_len, seq_len) weights = masked_softmax(scores, mask, dim=-1) weights = self.dropout(weights) context = torch.matmul(weights, v_eager) else: context = kernel.fused_attention( q, k, v, causal=self.causal, window_size=self.window_size, dropout_p=self.dropout.p, training=self.training, ) weights = None merged_context = merge_heads(context) if skip_out_proj and self.tp_size == 1 and not output_attentions and not return_kv_cache_estimate: return merged_context output = self.out_proj(merged_context) if self.tp_size > 1: if self.tp_group is None: raise RuntimeError( "DenseMHA was built with tp_size>1 but tp_group was never " "set -- call GPT.set_tp_group(group) after construction, " "before the first forward pass." ) torch.distributed.all_reduce(output, op=torch.distributed.ReduceOp.SUM, group=self.tp_group) if output_attentions or return_kv_cache_estimate: cache_bytes = None if return_kv_cache_estimate: cache_bytes = estimate_kv_cache_bytes( batch_size, seq_len, self.num_kv_heads, self.head_dim, hidden_states.dtype, window_size=self.window_size, ) return AttentionOutput(output=output, attention_weights=weights if output_attentions else None, kv_cache_bytes=cache_bytes) return output class SwiGLU_FFN(nn.Module): """ SwiGLU Feed Forward Network. """ def __init__(self, d_model, ffn_mult=8/3, dropout=0.0, tp_size: int = 1): super().__init__() raw_hidden = d_model * ffn_mult hidden_dim = int(((raw_hidden + 63) // 64) * 64) if hidden_dim % tp_size != 0: raise ValueError( f"tensor_parallel_size ({tp_size}) must divide the FFN hidden_dim " f"({hidden_dim} = round({d_model} * {ffn_mult}) -> next mult of 64) " f"-- adjust ffn_mult or tensor_parallel_size." ) local_hidden_dim = hidden_dim // tp_size self.tp_size = tp_size self.tp_group = None # set post-construction via GPT.set_tp_group() self.W_gate = nn.Linear(d_model, local_hidden_dim, bias=False) self.W_value = nn.Linear(d_model, local_hidden_dim, bias=False) self.W_out = nn.Linear(local_hidden_dim, d_model, bias=False) self.dropout = nn.Dropout(dropout) def forward(self, x): gate_raw = self.W_gate(x) value = self.W_value(x) hidden = None if _use_flash_swiglu_config() and flash_ops is not None and gate_raw.is_cuda: hidden = flash_ops.fused_swiglu(gate_raw, value) if hidden is None: hidden = F.silu(gate_raw) * value hidden = self.dropout(hidden) out = self.W_out(hidden) if self.tp_size > 1: if self.tp_group is None: raise RuntimeError( "SwiGLU_FFN was built with tp_size>1 but tp_group was never " "set -- call GPT.set_tp_group(group) after construction, " "before the first forward pass." ) torch.distributed.all_reduce(out, op=torch.distributed.ReduceOp.SUM, group=self.tp_group) return out # Return delta only class AryaSparseMoE(nn.Module): """ DeepSeek-style Sparse Mixture of Experts. Expects input to be ALREADY normalized (Pre-LN). Returns the combined FFN delta. """ def __init__(self, d_model, n_experts=8, n_shared=1, top_k=2, ffn_mult=8/3, dropout=0.0, use_capacity_routing=False, capacity_factor=1.25, fixed_capacity: int | None = None): super().__init__() self.n_experts = n_experts self.n_routed = n_experts - n_shared self.top_k = top_k self.d_model = d_model self.use_capacity_routing = use_capacity_routing self.capacity_factor = capacity_factor self.fixed_capacity = fixed_capacity if self.n_routed < top_k: raise ValueError( f"AryaSparseMoE config error: n_experts={n_experts} - n_shared={n_shared} " f"= {self.n_routed} routed experts, but top_k={top_k} routing needs at " f"least {top_k} routed experts to choose from." ) self.shared_expert = SwiGLU_FFN(d_model, ffn_mult, dropout=dropout) self.routed_experts = nn.ModuleList([ SwiGLU_FFN(d_model, ffn_mult, dropout=dropout) for _ in range(self.n_routed) ]) self.router = nn.Linear(d_model, self.n_routed, bias=False) self.expert_bias = nn.Parameter(torch.zeros(self.n_routed)) def forward(self, x): B, T, D = x.shape shared_out = self.shared_expert(x) router_logits = self.router(x) selection_scores = router_logits + self.expert_bias full_probs = F.softmax(selection_scores.float(), dim=-1) topk_probs, topk_indices = full_probs.topk(self.top_k, dim=-1) gate_weights = topk_probs.to(router_logits.dtype) x_flat = x.view(B * T, D) gate_flat = gate_weights.view(B * T, self.top_k) idx_flat = topk_indices.view(B * T, self.top_k) if self.use_capacity_routing: routed_out = self._forward_capacity_routed(x_flat, gate_flat, idx_flat) routed_out = routed_out.view(B, T, D) out = shared_out + routed_out return out routed_out = torch.zeros_like(x_flat) for expert_id in range(self.n_routed): expert = self.routed_experts[expert_id] token_mask = (idx_flat == expert_id).any(dim=-1) if not token_mask.any(): continue selected_tokens = x_flat[token_mask] expert_out = expert(selected_tokens) # Returns delta only weight = ( (idx_flat[token_mask] == expert_id).float() * gate_flat[token_mask] ).sum(dim=-1) routed_out[token_mask] += weight.unsqueeze(-1) * expert_out routed_out = routed_out.view(B, T, D) out = shared_out + routed_out return out def _forward_capacity_routed(self, x_flat, gate_flat, idx_flat): n_tokens, D = x_flat.shape n_routed, capacity = self.n_routed, self._capacity(n_tokens) flat_expert_ids = idx_flat.reshape(-1) flat_gates = gate_flat.reshape(-1) flat_token_ids = torch.arange(n_tokens, device=x_flat.device).unsqueeze(1).expand(-1, self.top_k).reshape(-1) one_hot = F.one_hot(flat_expert_ids, num_classes=n_routed).to(torch.float32) position_in_expert = (one_hot.cumsum(dim=0) * one_hot).sum(dim=1).long() - 1 keep = position_in_expert < capacity safe_position = position_in_expert.clamp(min=0, max=capacity - 1) flat_slot = flat_expert_ids * capacity + safe_position gathered_x = x_flat[flat_token_ids] * keep.unsqueeze(-1).to(x_flat.dtype) dispatch = torch.zeros(n_routed * capacity, D, device=x_flat.device, dtype=x_flat.dtype) dispatch.index_add_(0, flat_slot, gathered_x) dispatch = dispatch.view(n_routed, capacity, D) expert_out = torch.stack( [expert(dispatch[e]) for e, expert in enumerate(self.routed_experts)], dim=0, ).view(n_routed * capacity, D) contrib = expert_out[flat_slot] * keep.unsqueeze(-1).to(x_flat.dtype) * flat_gates.unsqueeze(-1) routed_out = torch.zeros(n_tokens, D, device=x_flat.device, dtype=x_flat.dtype) routed_out.index_add_(0, flat_token_ids, contrib) return routed_out def _capacity(self, n_tokens) -> int: if self.fixed_capacity is not None: return self.fixed_capacity total_assignments = n_tokens * self.top_k avg_load = -(-total_assignments // self.n_routed) numer = int(round(self.capacity_factor * 1000)) capacity = -(-(avg_load * numer) // 1000) return capacity + 1 def build_attention_layers( num_layers: int, hidden_size: int, num_heads: int, dropout: float = 0.0, bias: bool = True, num_kv_heads: int | None = None, pattern: str = "dense", max_seq_len: int = 4096, sliding_window_size: int = 256, tp_size: int = 1, rope_dim: int | None = None, yarn_scale: float = 1.0, yarn_orig_max_pos: int | None = None, yarn_alpha: float = 1.0, yarn_beta: float = 32.0, ): if num_kv_heads is None: num_kv_heads = num_heads layers = [] layer_types = [] if pattern == "dense": for _ in range(num_layers): layers.append(DenseMHA( hidden_size=hidden_size, num_heads=num_heads, num_kv_heads=num_kv_heads, dropout=dropout, bias=bias, max_seq_len=max_seq_len, window_size=None, tp_size=tp_size, rope_dim=rope_dim, yarn_scale=yarn_scale, yarn_orig_max_pos=yarn_orig_max_pos, yarn_alpha=yarn_alpha, yarn_beta=yarn_beta, )) layer_types.append("dense") elif pattern == "pyramid_swa": if num_kv_heads <= 1: pyramid_kv_heads = [1, 1, 1] else: import math as _math mqa_floor = max(2, _math.ceil(num_kv_heads / 4)) mqa_floor = min(mqa_floor, num_kv_heads) # never exceed full heads pyramid_kv_heads = sorted(set([ mqa_floor, max(mqa_floor, _math.ceil(num_kv_heads / 2)), num_kv_heads, ])) while len(pyramid_kv_heads) < 3: pyramid_kv_heads.append(num_kv_heads) pyramid_kv_heads = pyramid_kv_heads[:3] for i in range(num_layers): pos_in_group = i % 4 is_last_layer = (i == num_layers - 1) if pos_in_group == 3 or is_last_layer: layers.append(DenseMHA( hidden_size=hidden_size, num_heads=num_heads, num_kv_heads=num_kv_heads, dropout=dropout, bias=bias, max_seq_len=max_seq_len, window_size=None, tp_size=tp_size, rope_dim=rope_dim, yarn_scale=yarn_scale, yarn_orig_max_pos=yarn_orig_max_pos, yarn_alpha=yarn_alpha, yarn_beta=yarn_beta, )) layer_types.append("global") else: kv_heads_here = pyramid_kv_heads[pos_in_group] layers.append(DenseMHA( hidden_size=hidden_size, num_heads=num_heads, num_kv_heads=kv_heads_here, dropout=dropout, bias=bias, max_seq_len=max_seq_len, window_size=sliding_window_size, tp_size=tp_size, rope_dim=rope_dim, yarn_scale=yarn_scale, yarn_orig_max_pos=yarn_orig_max_pos, yarn_alpha=yarn_alpha, yarn_beta=yarn_beta, )) layer_types.append("swa") elif pattern == "mod": for i in range(num_layers): layers.append(DenseMHA( hidden_size=hidden_size, num_heads=num_heads, num_kv_heads=num_kv_heads, dropout=dropout, bias=bias, max_seq_len=max_seq_len, window_size=None, tp_size=tp_size, rope_dim=rope_dim, yarn_scale=yarn_scale, yarn_orig_max_pos=yarn_orig_max_pos, yarn_alpha=yarn_alpha, yarn_beta=yarn_beta, )) is_full_dense = (i % 4 == 3) or (i == num_layers - 1) layer_types.append("mod_dense" if is_full_dense else "mod") else: raise ValueError( f"Unknown attention pattern {pattern!r} -- supported patterns: " f"'dense', 'pyramid_swa', 'mod'. " f"(BigBird/conv-hybrid/LSA/HCA patterns have been removed.)" ) return nn.ModuleList(layers), layer_types class AryaBlock(nn.Module): """Single transformer block: Layer Norm → Attention → MLP (with residual connections). `attn_module` is DenseMHA, pre-built by GPT via build_attention_layers. Every layer is fully standalone -- no cross-layer attention/KV sharing of any kind. `role` is purely a label ("dense" | "lsa") for inspection/debugging. """ def __init__(self, config, attn_module: nn.Module, role: str = "dense"): super().__init__() self.ln_1 = RMSNorm(config.n_embd, eps=1e-5) # Pre-LN for the MLP branch: ensure router and experts see the # identical normalized tensor the MLPs expect. self.ln_2 = RMSNorm(config.n_embd, eps=1e-5) self.attn = attn_module self.role = role # "dense" | "lsa" self.use_moe = getattr(config, 'use_moe', False) if self.use_moe: if getattr(config, 'tp_size', 1) > 1: raise ValueError( "tensor_parallel_size > 1 is not supported together with use_moe=True -- " "MoE needs its own (unimplemented here) expert-parallel sharding scheme, " "column/row-parallel head splitting does not apply to expert routing. " "Set use_moe=False or tensor_parallel_size=1." ) self.ffn = AryaSparseMoE( d_model=config.n_embd, n_experts=config.n_experts, n_shared=config.n_shared, top_k=getattr(config, 'moe_top_k', 2), ffn_mult=config.ffn_mult, dropout=config.dropout, use_capacity_routing=True, fixed_capacity=getattr(config, 'moe_fixed_capacity', None), ) else: self.ffn = SwiGLU_FFN(config.n_embd, config.ffn_mult, dropout=config.dropout, tp_size=getattr(config, 'tp_size', 1)) def forward(self, x): """Every attention layer (DenseMHA, LSA) is standalone and computes its own K/V from `x` alone -- no external_kv, no cross-layer sharing, no produced_kv to pass on. """ normed = self.ln_1(x) fused_out = None raw_ctx = None attn_delta = None use_fused_outproj = ( _use_flash_outproj_config() and flash_ops is not None and self.attn.out_proj.bias is None ) if use_fused_outproj: attn_raw = self.attn(normed, skip_out_proj=True) candidate = module_output_tensor(attn_raw) if candidate.is_cuda and candidate.shape[-1] == self.attn.out_proj.weight.shape[1]: raw_ctx = candidate gemm_dtype = raw_ctx.dtype fused_out = flash_ops.fused_outproj_add_rmsnorm( raw_ctx.reshape(-1, raw_ctx.shape[-1]), self.attn.out_proj.weight.to(dtype=gemm_dtype), x.reshape(-1, x.shape[-1]).to(dtype=gemm_dtype), self.ln_2.weight.to(dtype=gemm_dtype), self.ln_2.eps, ) if fused_out is not None: residual_out, normed_out = fused_out x = residual_out.to(dtype=x.dtype).reshape(x.shape) ffn_in = normed_out.reshape(x.shape) else: skip_out_proj_honored = self.attn.tp_size == 1 if skip_out_proj_honored: raw_ctx = candidate attn_delta = self.attn.out_proj(raw_ctx) else: attn_delta = candidate if fused_out is None: if attn_delta is None: if raw_ctx is not None: attn_delta = self.attn.out_proj(raw_ctx) else: attn_raw = self.attn(normed) attn_delta = module_output_tensor(attn_raw) tier2_out = None if _use_fused_add_rmsnorm_config() and fused_kernels is not None: tier2_out = fused_kernels.fused_add_rmsnorm( x, attn_delta, self.ln_2.weight.to(dtype=x.dtype), self.ln_2.eps ) if tier2_out is not None: x, ffn_in = tier2_out else: # Tier 3: fully eager, byte-for-byte what this block # always ran before any fusion work this session. x = x + attn_delta ffn_in = self.ln_2(x) ffn_delta = self.ffn(ffn_in) x = x + ffn_delta return x class MoDBlock(nn.Module): """Mixture-of-Depths via a continuous sigmoid gate -- no topk, no gather, no scatter. WILL THE GATE ACTUALLY SELECT TOKENS? ────────────────────────────────────── Short answer: not without explicit sparsity pressure. A plain sigmoid gate starts near 0.5 (for zero-mean inputs after RMSNorm) and stays fractional -- it applies 40% of the delta to some tokens, 70% to others, 30% to others. The gradient never forces it toward 0 or 1 unless there is a cost to using non-binary values. Two mechanisms push the gate toward genuine token selection: 1. mod_gate_entropy_coeff (config.py, default 0.01): Adds a per-step auxiliary loss term: L_entropy = -coeff * mean( gate*log(gate) + (1-gate)*log(1-gate) ) This is the binary entropy H(gate), negated -- maximizing binary entropy means pushing gate toward 0 or 1, where H=0 (the gate is certain). A gate at 0.5 has H=1.0 (maximum uncertainty); a gate at 0.05 or 0.95 has H≈0.29. The coefficient controls how hard this regularizer pushes: 0.0 -> pure sigmoid, stays fractional, no selection 0.01 -> gentle push; gates typically reach 0.1/0.9 range by step 2000 0.1 -> strong push; gates near 0/1 within ~500 steps but can destabilize This term is computed in forward() and returned alongside the block output. train.py adds it to the main loss, multiplied by the coefficient. 2. Natural gradient pressure from the main loss: If one token position's delta reliably INCREASES the loss regardless of context, the model learns to set gate≈0 there (the identity is better). If another position's delta reliably DECREASES the loss, gate≈1 emerges. This happens without any regularization, just more slowly -- typically ~5000-10000 steps before meaningful gate bimodality appears. The entropy loss (mechanism 1) dramatically speeds up mechanism 2 by making the gate pay a cost for ambiguity. Use mod_gate_entropy_coeff=0.01 as a starting point and inspect a histogram of gate values at step 1000 and 5000. If the distribution is bimodal (peaks near 0 and 1), the gate is selecting. If it's unimodal near 0.5, increase the coefficient. GATE DESIGN ─────────── gate = sigmoid(router(x)) # (B, T, 1) -- D->1 matmul + sigmoid delta = inner(x) - x # (B, T, D) -- net block update out = x + gate * delta # elementwise, fixed shape, fused aux = entropy_coeff * H_binary(gate) # scalar, added to training loss Parameter cost: D floats (router.weight). At D=256 that's 256 params total. """ def __init__(self, config, inner: nn.Module, capacity: float = 0.125): super().__init__() self.inner = inner self.capacity = capacity self._entropy_coeff = float(getattr(config, 'mod_gate_entropy_coeff', 0.01)) self.router = nn.Linear(config.n_embd, 1, bias=False) def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: gate = torch.sigmoid(self.router(x)) # (B, T, 1) inner_out = self.inner(x) # (B,T,D) out = x + gate * (inner_out - x) # (B,T,D); fuses delta into out if self.training and self._entropy_coeff > 0.0: g = gate.float().clamp(1e-6, 1.0 - 1e-6) # fp32 for log stability ent = -(g * torch.log(g) + (1.0 - g) * torch.log(1.0 - g)) # H per element gate_aux = self._entropy_coeff * ent.mean() else: gate_aux = torch.zeros((), device=x.device, dtype=torch.float32) return out, gate_aux def extra_repr(self) -> str: return f"gate_based=True,entropy_coeff={self._entropy_coeff}" class MultiTokenPrediction(nn.Module): """Multi-Token Prediction auxiliary heads (Gloeckle et al. / DeepSeek-V3 style): predict tokens at offsets 2, 3, ... beyond the main next-token prediction, using the same pre-lm_head hidden state. Each head k predicts token t+k+1 from hidden state h_t via a separate RMSNorm (so each depth can learn its own scale) followed by the shared lm_head weight (weight tying: same projection used for the main head), keeping parameter count increase minimal. Loss contribution: mtp_lambda * mean(CE_1, CE_2, ..., CE_depth). The main model's loss is unaffected; MTP loss is added on top. """ def __init__(self, config, depth: int = 1): super().__init__() self.depth = depth # One RMSNorm per MTP depth; no separate linear (share lm_head.weight). self.norms = nn.ModuleList([ RMSNorm(config.n_embd, eps=1e-5) for _ in range(depth) ]) def compute_loss( self, x: torch.Tensor, # (B, T, D) -- pre-lm_head hidden states targets: torch.Tensor, # (B, T) -- already the +1 shifted targets lm_head_weight: torch.Tensor, # (vocab_size, D) -- tied weight mtp_lambda: float = 0.3, ) -> torch.Tensor: total = torch.zeros((), device=x.device, dtype=torch.float32) valid_depths = 0 for k, norm in enumerate(self.norms): offset = k + 1 # additional token offset beyond the main +1 shift Tk = x.size(1) - offset if Tk <= 0: continue h = norm(x[:, :Tk]) # (B, Tk, D) logits_k = F.linear(h, lm_head_weight) # (B, Tk, vocab_size) tgt_k = targets[:, offset:] # (B, Tk) loss_k = F.cross_entropy( logits_k.reshape(-1, logits_k.size(-1)).float(), tgt_k.reshape(-1).long(), ignore_index=-1, ) total = total + loss_k # fp32 + fp32 -- no precision loss valid_depths += 1 if valid_depths == 0: return torch.zeros((), device=x.device, dtype=torch.float32) return mtp_lambda * (total / valid_depths) # fp32 scalar @dataclass class GPTConfig: """Configuration for GPT model""" block_size: int = _alpha_config_attr('CONTEXT', 1024) vocab_size: int = _alpha_config_attr('vocab_size', 50304) n_layer: int = _alpha_config_attr('numberoflayers', 12) n_head: int = _alpha_config_attr('numberofheads', 12) n_embd: int = _alpha_config_attr('D_MODEL', 768) dropout: float = _alpha_config_attr('dropout', 0.0) bias: bool = _alpha_config_attr('bias', False) ffn_mult: float = _alpha_config_attr('ffn_mult', 4.0) d_rope: int | None = _alpha_config_attr('d_rope', None) top_k: int = _alpha_config_attr('top_k', 64) n_experts: int = _alpha_config_attr('n_experts', 8) n_shared: int = _alpha_config_attr('n_shared', 1) use_moe: bool = _alpha_config_attr('use_moe', False) use_liger: bool = _alpha_config_attr('use_liger', False) moe_top_k: int = _alpha_config_attr('moe_top_k', 2) gradient_checkpointing: bool = _alpha_config_attr('gradient_checkpointing', False) # trade ~20-30% more compute for substantially less activation memory (grows with n_layer -- see GPT.forward) num_kv_heads: int | None = _alpha_config_attr('num_kv_heads', None) pattern: str = _alpha_config_attr('pattern', 'pyramid_swa') sliding_window_size: int = _alpha_config_attr('sliding_window_size', 256) tp_size: int = _alpha_config_attr('tensor_parallel_size', 1) gen_headroom: int = _alpha_config_attr('gen_headroom', 0) max_gen_tokens: int | None = _alpha_config_attr('max_gen_tokens', None) yarn_scale: float = _alpha_config_attr('yarn_scale', 1.0) yarn_orig_max_pos: int | None = _alpha_config_attr('yarn_orig_max_pos', None) yarn_alpha: float = _alpha_config_attr('yarn_alpha', 1.0) yarn_beta: float = _alpha_config_attr('yarn_beta', 32.0) mod_capacity: float = _alpha_config_attr('mod_capacity', 0.125) use_mtp: bool = _alpha_config_attr('use_mtp', False) mtp_depth: int = _alpha_config_attr('mtp_depth', 1) mtp_lambda: float = _alpha_config_attr('mtp_lambda', 0.3) use_nvfp4: bool = _alpha_config_attr('use_nvfp4', False) use_int8: bool = _alpha_config_attr('use_int8', False) int8_group_size: int = _alpha_config_attr('int8_group_size', 128) mod_gate_entropy_coeff: float = _alpha_config_attr('mod_gate_entropy_coeff', 0.01) class GPT(nn.Module): """Full GPT language model""" def __init__(self, config): super().__init__() if config.vocab_size is None: raise ValueError("config.vocab_size must be set") if config.block_size is None: raise ValueError("config.block_size must be set") self.config = config head_dim = config.n_embd // config.n_head rope_dim = head_dim if getattr(config, 'd_rope', None) is None else config.d_rope gen_headroom = max(0, int(getattr(config, 'gen_headroom', 0))) rope_table_len = config.block_size + gen_headroom yarn_scale = float(getattr(config, 'yarn_scale', 1.0)) yarn_orig_max_pos_cfg = getattr(config, 'yarn_orig_max_pos', None) yarn_orig_max_pos = int(yarn_orig_max_pos_cfg) if yarn_orig_max_pos_cfg is not None else config.block_size yarn_alpha = float(getattr(config, 'yarn_alpha', 1.0)) yarn_beta = float(getattr(config, 'yarn_beta', 32.0)) effective_max_seq_len = int(rope_table_len * yarn_scale) if yarn_scale > 1.0 else rope_table_len attn_layers, self.layer_types = build_attention_layers( num_layers=config.n_layer, hidden_size=config.n_embd, num_heads=config.n_head, dropout=config.dropout, bias=config.bias, num_kv_heads=config.num_kv_heads, pattern=config.pattern, max_seq_len=effective_max_seq_len, sliding_window_size=config.sliding_window_size, tp_size=config.tp_size, rope_dim=rope_dim, yarn_scale=yarn_scale, yarn_orig_max_pos=yarn_orig_max_pos, yarn_alpha=yarn_alpha, yarn_beta=yarn_beta, ) mod_capacity = float(getattr(config, 'mod_capacity', 0.125)) blocks = [] for i in range(config.n_layer): block = AryaBlock(config, attn_layers[i], role=self.layer_types[i]) if self.layer_types[i] == "mod": # Wrap in MoDBlock: 3/4 of layers route only top-capacity # fraction of tokens through the full AryaBlock. block = MoDBlock(config, block, capacity=mod_capacity) blocks.append(block) self.transformer = nn.ModuleDict(dict( wte=nn.Embedding(config.vocab_size, config.n_embd), drop=nn.Dropout(config.dropout), h=nn.ModuleList(blocks), ln_f=RMSNorm(config.n_embd, eps=1e-5), )) self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False) use_mtp = bool(getattr(config, 'use_mtp', False)) mtp_depth = int(getattr(config, 'mtp_depth', 1)) self.mtp = MultiTokenPrediction(config, depth=mtp_depth) if use_mtp else None self._mtp_lambda = float(getattr(config, 'mtp_lambda', 0.3)) self.apply(self._init_weights) self.transformer.wte.weight = self.lm_head.weight residual_std = 0.02 / math.sqrt(config.n_layer) for name, p in self.named_parameters(): if name.endswith("out_proj.weight") or name.endswith("W_out.weight"): torch.nn.init.normal_(p, mean=0.0, std=residual_std) # ── Quantization (mutually exclusive; nvFP4 takes priority on SM100+) ── # Both quantize_() calls must happen BEFORE torch.compile() in train.py. # _int8_enabled() already checks _nvfp4_enabled() and refuses if both are set. self._use_nvfp4: bool = False self._use_int8: bool = False if _nvfp4_enabled(): n_quantized = apply_nvfp4_to_model(self) self._use_nvfp4 = True warnings.warn( f"[nvFP4/RHT] torchao quantized {n_quantized} Linear layers " f"(Blackwell SM100+). Weight: FP4 e2m1 per-block-16 RHT. " f"Activation: FP4 dynamic per-tensor scale (FP8 intermediate). " f"Fused GEMM kernel emitted by torch.compile. " f"Skipped: wte, lm_head, norms, router, gate, small/misaligned linears. " f"Verify val_loss vs bf16 baseline before committing to a long run.", stacklevel=2, ) elif _int8_enabled(): gs = int(getattr(config, 'int8_group_size', 128)) n_quantized = apply_int8_to_model(self, group_size=gs) self._use_int8 = True warnings.warn( f"[INT8/Jetfire] torchao quantized {n_quantized} Linear layers " f"(T4/Turing+). Weight: INT8 per-group-{gs} symmetric. " f"Activation: INT8 dynamic per-tensor scale (runtime, no calibration). " f"Fused kernel emitted by torch.compile via Inductor. " f"Skipped: wte, lm_head, norms, router, gate, small linears. " f"Verify val_loss vs fp16 baseline before a long run.", stacklevel=2, ) # Weight tying can be broken by quantized module replacement. # Reassert the canonical tie after either quantization pass. self.transformer.wte.weight = self.lm_head.weight assert self.transformer.wte.weight is self.lm_head.weight def set_tp_group(self, group) -> None: if self.config.tp_size <= 1: return for block in self.transformer.h: real_block = block.inner if isinstance(block, MoDBlock) else block real_block.attn.tp_group = group if hasattr(real_block, "ffn") and real_block.ffn is not None: real_block.ffn.tp_group = group if torch.distributed.is_initialized() and group is not None: for block in self.transformer.h: real_block = block.inner if isinstance(block, MoDBlock) else block attn = real_block.attn if not isinstance(attn, DenseMHA): continue if attn.num_kv_heads >= attn.tp_size: continue with torch.no_grad(): # Full qkv projection tensor is the canonical object to align. torch.distributed.broadcast(attn.qkv_proj.weight.data, src=0, group=group) if attn.qkv_proj.bias is not None: torch.distributed.broadcast(attn.qkv_proj.bias.data, src=0, group=group) def _init_weights(self, module): if isinstance(module, nn.Linear): torch.nn.init.normal_(module.weight, mean=0.0, std=0.02) if module.bias is not None: torch.nn.init.zeros_(module.bias) elif isinstance(module, nn.Embedding): torch.nn.init.normal_(module.weight, mean=0.0, std=0.02) def get_num_params(self, non_embedding=True): """Count number of parameters""" n_params = sum(p.numel() for p in self.parameters()) if non_embedding: n_params -= self.transformer.wte.weight.numel() return n_params def forward(self, idx, targets=None, need_logits=True): idx = idx.long() if targets is not None: targets = targets.long() b, t = idx.size() torch._check( t <= self.config.block_size, lambda: f"Sequence too long: {t} > {self.config.block_size}", ) tok_emb = self.transformer.wte(idx) # (B, T, n_embd) x = self.transformer.drop(tok_emb) mod_gate_loss = torch.zeros((), device=x.device, dtype=torch.float32) for block in self.transformer.h: is_mod = isinstance(block, MoDBlock) if self.config.gradient_checkpointing and self.training: out = torch.utils.checkpoint.checkpoint(block, x, use_reentrant=False) if is_mod: x, gate_aux = out mod_gate_loss = mod_gate_loss + gate_aux.float() else: x = out else: if is_mod: x, gate_aux = block(x) mod_gate_loss = mod_gate_loss + gate_aux.float() else: x = block(x) x = self.transformer.ln_f(x) if targets is not None: mtp_loss = torch.zeros((), device=x.device, dtype=torch.float32) if self.mtp is not None: mtp_loss = self.mtp.compute_loss( x, targets, self.lm_head.weight, mtp_lambda=self._mtp_lambda ) # Keep the primary language-model CE loss separate from auxiliary MTP/MoD penalties. # This lets caller-side logging and validation track the true CE signal without # the auxiliary terms masking the baseline perplexity curve. if not need_logits and _use_liger_config() and _LigerFusedLinearCrossEntropyLoss is not None and _liger_cuda_gate(x): liger_ce = _liger_fused_ce_attempt(x, self.lm_head.weight, targets) if liger_ce is not None: total_loss = liger_ce + mtp_loss + mod_gate_loss self.last_ce_loss = liger_ce.detach() self.last_total_loss = total_loss.detach() return None, total_loss logits = F.linear(x, self.lm_head.weight) ce_loss = F.cross_entropy( logits.reshape(-1, logits.size(-1)).float(), targets.reshape(-1), ignore_index=-1 ) # Keep the loss scalar finite by avoiding a patched post-hoc # nan_to_num that stamps out the gradient entirely. Repair the # overflow sources (logit clamp/initialization) out front instead. total_loss = ce_loss + mtp_loss + mod_gate_loss self.last_ce_loss = ce_loss.detach() self.last_total_loss = total_loss.detach() if not need_logits: logits = None else: logits = F.linear(x[:, [-1], :], self.lm_head.weight) # (B, 1, vocab_size) ce_loss = None total_loss = None return logits, total_loss @torch.no_grad() def generate(self, idx, max_new_tokens=None, temperature=1.0, top_k=None): if max_new_tokens is None: max_new_tokens = getattr(self.config, 'max_gen_tokens', None) if max_new_tokens is None: raise ValueError( "generate() needs max_new_tokens, either passed " "directly or set as max_gen_tokens in config.py." ) for _ in range(max_new_tokens): # Crop to block size if needed idx_cond = idx if idx.size(1) <= self.config.block_size else idx[:, -self.config.block_size:] # Forward pass logits, _ = self(idx_cond) # (B, T, vocab_size) # Get logits for next token (only last position) logits = logits[:, -1, :] / temperature # (B, vocab_size) # Top-k filtering if top_k is not None: v, _ = torch.topk(logits, min(top_k, logits.size(-1))) logits[logits < v[:, [-1]]] = -float('Inf') probs = F.softmax(logits.float(), dim=-1) # (B, vocab_size) # Sample next token idx_next = torch.multinomial(probs, num_samples=1) # (B, 1) # Append to sequence idx = torch.cat((idx, idx_next), dim=1) # (B, T+1) return idx def audit_fused_kernels(model, device, dtype=torch.float16, verbose=True): config = model.config head_dim = config.n_embd // config.n_head rope_dim = head_dim if getattr(config, 'd_rope', None) is None else config.d_rope b, t = 2, 8 results = {} if flash_ops is None: if verbose: print( "[FusionAudit] flash_ops extension failed to load entirely -- " "all three flash_* fusions (rope/outproj_add_rmsnorm/swiglu) " "are running eager fallback for the WHOLE run, not just under " "some condition. Check flash.cu compile output / nvcc " "availability above for the real cause." ) return {"flash_ops_loaded": False} results["flash_ops_loaded"] = True # -- fused RoPE -- try: q = torch.randn(b, t, config.n_head, head_dim, device=device, dtype=dtype) k = torch.randn(b, t, config.num_kv_heads or config.n_head, head_dim, device=device, dtype=dtype) cos, sin = build_rope_cache(config.block_size, rope_dim) cos, sin = cos.to(device), sin.to(device) out = flash_ops.fused_rope_qk(q, k, cos[:t].to(dtype), sin[:t].to(dtype), rope_dim) results["fused_rope"] = out is not None except Exception as e: # noqa: BLE001 -- diagnostic only, must not crash results["fused_rope"] = False results["fused_rope_error"] = repr(e) # -- fused out_proj + residual-add + RMSNorm -- try: x = torch.randn(b * t, config.n_embd, device=device, dtype=dtype) w = torch.randn(config.n_embd, config.n_embd, device=device, dtype=dtype) residual = torch.randn(b * t, config.n_embd, device=device, dtype=dtype) norm_w = torch.ones(config.n_embd, device=device, dtype=dtype) out = flash_ops.fused_outproj_add_rmsnorm(x, w, residual, norm_w, 1e-5) results["fused_outproj_add_rmsnorm"] = out is not None except Exception as e: # noqa: BLE001 results["fused_outproj_add_rmsnorm"] = False results["fused_outproj_add_rmsnorm_error"] = repr(e) # -- fused SwiGLU -- try: hidden = int(config.n_embd * config.ffn_mult) gate = torch.randn(b * t, hidden, device=device, dtype=dtype) value = torch.randn(b * t, hidden, device=device, dtype=dtype) out = flash_ops.fused_swiglu(gate, value) results["fused_swiglu"] = out is not None except Exception as e: # noqa: BLE001 results["fused_swiglu"] = False results["fused_swiglu_error"] = repr(e) if verbose: print("[FusionAudit] one-time probe of fused kernels against real on-device tensors:") for name in ("fused_rope", "fused_outproj_add_rmsnorm", "fused_swiglu"): ok = results.get(name, False) status = "ENGAGED" if ok else "FELL BACK TO EAGER" print(f" {name:28s} -> {status}") if not ok and f"{name}_error" in results: print(f" reason: {results[f'{name}_error']}") n_ok = sum(results.get(n, False) for n in ("fused_rope", "fused_outproj_add_rmsnorm", "fused_swiglu")) if n_ok < 3: print( f" [FusionAudit] {3 - n_ok}/3 fusions NOT engaging -- given " f"this model is memory-bandwidth-bound (per the roofline " f"analysis), a missed fusion means real intermediate-tensor " f"HBM traffic that shouldn't be there. Worth fixing before " f"chasing anything else." ) else: print(" [FusionAudit] all 3/3 fusions engaged -- fusion is not the bottleneck here.") return results