from __future__ import annotations import os from dataclasses import dataclass from typing import Optional import torch import torch.nn as nn import torch.nn.functional as F def _detect_cpu_bf16() -> bool: try: with open("/proc/cpuinfo") as f: return "avx512_bf16" in f.read() except Exception: return False _CPU_HAS_BF16: bool = _detect_cpu_bf16() try: import pyfftw pyfftw.interfaces.cache.enable() pyfftw.interfaces.cache.set_keepalive_time(60.0) except ImportError: pass try: if os.environ.get("NO_CPP_EXT"): raise ImportError("C++ extensions disabled via NO_CPP_EXT") import torch.utils.cpp_extension as _cpp_ext _ext_path = os.path.join(os.path.dirname(__file__), "csrc") _gla_ext = ( _cpp_ext.load( name="gla_scan_cpu", sources=[os.path.join(_ext_path, "gla_scan.cpp")], extra_cflags=["-O3", "-fopenmp", "-march=native"], extra_ldflags=["-fopenmp"], verbose=False, ) if os.path.exists(os.path.join(_ext_path, "gla_scan.cpp")) else None ) except Exception: _gla_ext = None try: if os.environ.get("NO_CPP_EXT"): raise ImportError("C++ extensions disabled via NO_CPP_EXT") import torch.utils.cpp_extension as _fno_cpp_ext _fno_ext_src = os.path.join(os.path.dirname(__file__), "csrc", "fno_conv.cpp") _fno_ext = ( _fno_cpp_ext.load( name="fno_conv_cpu", sources=[_fno_ext_src], extra_cflags=["-O3", "-fopenmp", "-march=native"], extra_ldflags=["-fopenmp"], verbose=False, ) if os.path.exists(_fno_ext_src) else None ) except Exception: _fno_ext = None try: from flashfftconv import FlashFFTConv as _FlashFFTConv _flash_fft_available = True except ImportError: _FlashFFTConv = None _flash_fft_available = False def _fft_dtype_to_torch(fft_dtype: str): return {"fp32": torch.float32, "bf16": torch.bfloat16, "fp16": torch.float16}[ fft_dtype ] def _gla_scan_py(q_k, k_k, v_f, chunk_gate, CS): B, H, T, D = q_k.shape n_chunks = T // CS bf16 = q_k.dtype == torch.bfloat16 q_c = q_k.reshape(B, H, n_chunks, CS, D) k_c = k_k.reshape(B, H, n_chunks, CS, D) v_c = v_f.reshape(B, H, n_chunks, CS, D) state = torch.zeros(B, H, D, D, device=q_k.device, dtype=torch.float32) z_norm = torch.zeros(B, H, D, device=q_k.device, dtype=torch.float32) chunks = [] for c in range(n_chunks): if bf16: num = q_c[:, :, c] @ state.bfloat16() den = (q_c[:, :, c] @ z_norm.bfloat16().unsqueeze(-1)).clamp(min=1.0) chunks.append(num / den) g = chunk_gate[:, :, c, None, None] state = g * state + (k_c[:, :, c].transpose(-2, -1) @ v_c[:, :, c]).float() z_norm = chunk_gate[:, :, c, None] * z_norm + k_c[:, :, c].sum(-2).float() else: num = q_c[:, :, c] @ state den = (q_c[:, :, c] @ z_norm.unsqueeze(-1)).clamp(min=1.0) chunks.append(num / den) g = chunk_gate[:, :, c, None, None] state = g * state + k_c[:, :, c].transpose(-2, -1) @ v_c[:, :, c] z_norm = chunk_gate[:, :, c, None] * z_norm + k_c[:, :, c].sum(dim=-2) return torch.cat(chunks, dim=2) @torch.compiler.disable class _GLAScanFn(torch.autograd.Function): @staticmethod def forward(ctx, q_k, k_k, v_f, chunk_gate, CS): ctx.CS = CS with torch.no_grad(): out, states, z_norms = _gla_ext.gla_scan_fwd( q_k.contiguous(), k_k.contiguous(), v_f.contiguous(), chunk_gate.contiguous(), CS, ) ctx.save_for_backward(q_k, k_k, v_f, chunk_gate, states, z_norms) return out @staticmethod def backward(ctx, grad_out): q_k, k_k, v_f, chunk_gate, states, z_norms = ctx.saved_tensors CS = ctx.CS d_q, d_k, d_v, d_gate = _gla_ext.gla_scan_backward( grad_out.float().contiguous(), q_k.contiguous(), k_k.contiguous(), v_f.contiguous(), chunk_gate.contiguous(), states, z_norms, CS, ) return (d_q, d_k, d_v, d_gate, None) @dataclass class CPUGPTConfig: vocab_size: int = 50257 seq_len: int = 1024 n_layer: int = 12 n_embd: int = 768 n_head: int = 12 fno_modes: int = 256 gla_chunk: int = 64 ffn_hidden: int = 2048 layer_pattern: str = "SSSL" bias: bool = False dropout: float = 0.0 fft_dtype: str = "fp32" use_flash_fft: bool = False def __post_init__(self): assert self.n_embd % 16 == 0, ( f"n_embd={self.n_embd} must be divisible by 16 for AMX BF16" ) assert self.ffn_hidden % 16 == 0, ( f"ffn_hidden={self.ffn_hidden} must be divisible by 16 for AMX BF16" ) def _layer_is_gla(layer_idx: int, n_layer: int, pattern: str) -> bool: if pattern == "SSSL": return (layer_idx % 4) == 3 elif pattern == "SGSG": return (layer_idx % 2) == 1 elif pattern == "GLA": return True elif pattern == "FNO": return False return False def rms_norm(x: torch.Tensor) -> torch.Tensor: return F.rms_norm(x, (x.size(-1),)) class SwiGLU(nn.Module): def __init__(self, cfg: CPUGPTConfig): super().__init__() h = cfg.ffn_hidden d = cfg.n_embd self.gate = nn.Linear(d, h, bias=cfg.bias) self.up = nn.Linear(d, h, bias=cfg.bias) self.down = nn.Linear(h, d, bias=cfg.bias) def forward(self, x: torch.Tensor) -> torch.Tensor: return self.down(F.silu(self.gate(x)) * self.up(x)) @torch.compiler.disable class _FNOConvFn(torch.autograd.Function): @staticmethod def forward(ctx, x, filter_td, T): x_f = x.float().contiguous() f_f = filter_td.float().detach().contiguous() y, Xf, Hf = _fno_ext.fno_conv_fwd(x_f, f_f, T) ctx.save_for_backward(Xf, Hf) ctx.n_use = min(filter_td.shape[1], T) ctx.M = filter_td.shape[1] return y.to(x.dtype) @staticmethod def backward(ctx, grad_out): Xf, Hf = ctx.saved_tensors d_x, d_filter = _fno_ext.fno_conv_backward( grad_out.float().contiguous(), Xf, Hf, ctx.n_use, ctx.M ) return d_x, d_filter, None class FNOSeqMixer(nn.Module): def __init__(self, cfg: CPUGPTConfig): super().__init__() C = cfg.n_embd M = cfg.fno_modes self.filter_td = nn.Parameter(torch.empty(C, M)) self.out_scale = nn.Linear(C, C, bias=cfg.bias) self.register_buffer("_hf_cache", None, persistent=False) self._filter_version: int = -1 self._fft_dtype = _fft_dtype_to_torch(cfg.fft_dtype) self._use_flash_fft = cfg.use_flash_fft and _flash_fft_available self._flash_fft_conv = None self._flash_fft_T: int = -1 nn.init.normal_(self.filter_td, std=0.02) def _flash_conv(self, x: torch.Tensor, T: int, C: int, n_use: int) -> torch.Tensor: if self._flash_fft_conv is None or self._flash_fft_T != T: self._flash_fft_conv = _FlashFFTConv(2 * T, dtype=self._fft_dtype).to( x.device ) self._flash_fft_T = T h = self.filter_td.new_zeros(C, 2 * T) h[:, :n_use] = self.filter_td[:, :n_use] xp = F.pad(x.transpose(1, 2), (0, T)) y = self._flash_fft_conv(xp.to(self._fft_dtype), h.to(self._fft_dtype)) return y[:, :, :T].transpose(1, 2).to(x.dtype) def _rfft_conv(self, x: torch.Tensor, T: int, C: int, n_use: int) -> torch.Tensor: dt = self._fft_dtype h = self.filter_td.new_zeros(2 * T, C) h[:n_use] = self.filter_td[:, :n_use].T.to(dt) xp = F.pad(x, (0, 0, 0, T)).to(dt) Xf = torch.fft.rfft(xp, dim=1) Hf = torch.fft.rfft(h, dim=0).unsqueeze(0) Xr, Xi = Xf.real, Xf.imag Hr, Hi = Hf.real, Hf.imag YHf = torch.view_as_complex( torch.stack([Xr * Hr - Xi * Hi, Xr * Hi + Xi * Hr], dim=-1).contiguous() ) return torch.fft.irfft(YHf, n=2 * T, dim=1)[:, :T] def forward(self, x: torch.Tensor) -> torch.Tensor: B, T, C = x.shape n_use = min(self.filter_td.shape[1], T) if _fno_ext is not None: y = _FNOConvFn.apply(x, self.filter_td, T) elif self._use_flash_fft and x.is_cuda: y = self._flash_conv(x, T, C, n_use) elif not self.training: y = self._forward_eval_cached(x, T, C, n_use) else: y = self._rfft_conv(x, T, C, n_use) return self.out_scale(y.to(x.dtype)) @torch.compiler.disable def _forward_eval_cached( self, x: torch.Tensor, T: int, C: int, n_use: int ) -> torch.Tensor: ver = self.filter_td._version if ( self._hf_cache is None or self._filter_version != ver or self._hf_cache.shape[0] != T + 1 ): h = self.filter_td.new_zeros(2 * T, C) h[:n_use] = self.filter_td[:, :n_use].T.detach() self._hf_cache = torch.fft.rfft(h.float(), dim=0).contiguous() self._filter_version = ver dt = self._fft_dtype xp = F.pad(x, (0, 0, 0, T)).to(dt) Xf = torch.fft.rfft(xp, dim=1) H = self._hf_cache.unsqueeze(0) Xr, Xi, Hr, Hi = Xf.real, Xf.imag, H.real, H.imag YHf = torch.view_as_complex( torch.stack([Xr * Hr - Xi * Hi, Xr * Hi + Xi * Hr], dim=-1).contiguous() ) return torch.fft.irfft(YHf, n=2 * T, dim=1)[:, :T] def prepare_inference(self): if hasattr(self, "_h_td"): return with torch.no_grad(): pass def _build_h_td(self, T: int) -> None: if hasattr(self, "_h_td"): return C = self.filter_td.shape[0] M = min(self.filter_td.shape[1], T) with torch.no_grad(): h_td = torch.zeros(T, C, dtype=torch.float32) h_td[:M] = self.filter_td[:, :M].T.float() self.register_buffer("_h_td", h_td.contiguous()) def step( self, x: torch.Tensor, buf: torch.Tensor, pos: int ) -> tuple[torch.Tensor, int]: T = buf.shape[1] self._build_h_td(T) slot = pos % T buf[:, slot, :] = x.detach() indices = torch.arange(T, device=x.device) ordered_idx = (slot + 1 + indices) % T ordered = buf[:, ordered_idx, :] h_flip = self._h_td.flip(0) y = (h_flip.unsqueeze(0) * ordered).sum(1) return self.out_scale(y.to(x.dtype)), pos + 1 class GLAMixer(nn.Module): def __init__(self, cfg: CPUGPTConfig): super().__init__() H = cfg.n_head D = cfg.n_embd // H C = cfg.n_embd CS = cfg.gla_chunk assert C % H == 0, "n_embd must be divisible by n_head" assert cfg.seq_len % CS == 0, ( f"seq_len {cfg.seq_len} must be divisible by gla_chunk {CS}" ) self.n_head = H self.d_head = D self.chunk = CS self.q_proj = nn.Linear(C, C, bias=False) self.k_proj = nn.Linear(C, C, bias=False) self.v_proj = nn.Linear(C, C, bias=False) self.g_proj = nn.Linear(C, H, bias=True) self.out_proj = nn.Linear(C, C, bias=False) self.gla_scale = nn.Parameter(torch.ones(H) * 0.5) nn.init.constant_(self.g_proj.bias, -4.0) def _python_scan( self, q_k: torch.Tensor, k_k: torch.Tensor, v: torch.Tensor, chunk_gate: torch.Tensor, ) -> torch.Tensor: B, H, T, D = q_k.shape CS = self.chunk n_chunks = T // CS q_c = q_k.view(B, H, n_chunks, CS, D) k_c = k_k.view(B, H, n_chunks, CS, D) v_c = v.view(B, H, n_chunks, CS, D) state = torch.zeros(B, H, D, D, device=q_k.device, dtype=q_k.dtype) z_norm = torch.zeros(B, H, D, device=q_k.device, dtype=q_k.dtype) out = torch.zeros(B, H, T, D, device=q_k.device, dtype=q_k.dtype) for c in range(n_chunks): num = q_c[:, :, c] @ state den = (q_c[:, :, c] @ z_norm.unsqueeze(-1)).clamp(min=1.0) out[:, :, c * CS : (c + 1) * CS] = num / den g = chunk_gate[:, :, c, None, None] state = g * state + k_c[:, :, c].transpose(-2, -1) @ v_c[:, :, c] z_norm = chunk_gate[:, :, c, None] * z_norm + k_c[:, :, c].sum(dim=-2) return out def forward(self, x: torch.Tensor) -> torch.Tensor: B, T, C = x.shape H, D, CS = self.n_head, self.d_head, self.chunk q = self.q_proj(x).reshape(B, T, H, D).transpose(1, 2) k = self.k_proj(x).reshape(B, T, H, D).transpose(1, 2) v = self.v_proj(x).reshape(B, T, H, D).transpose(1, 2) n_chunks = T // CS BH = B * H q_l = q.reshape(BH, n_chunks, CS, D) k_l = k.reshape(BH, n_chunks, CS, D) v_l = v.reshape(BH, n_chunks, CS, D) y_local = F.scaled_dot_product_attention(q_l, k_l, v_l, is_causal=True).reshape( B, H, T, D ) if _CPU_HAS_BF16 and q.dtype == torch.bfloat16: q_k = F.elu(q) + 1.0 k_k = F.elu(k) + 1.0 v_f = v else: q_k = F.elu(q.float()) + 1.0 k_k = F.elu(k.float()) + 1.0 v_f = v.float() log_g = -F.softplus(self.g_proj(x).float()) log_g = log_g.transpose(1, 2) chunk_log_g = log_g.reshape(B, H, n_chunks, CS).sum(-1) chunk_gate = chunk_log_g.exp() if _gla_ext is not None and q_k.dtype == torch.float32: y_gla = _GLAScanFn.apply(q_k, k_k, v_f, chunk_gate, CS) else: y_gla = _gla_scan_py(q_k, k_k, v_f, chunk_gate, CS) scale = self.gla_scale.reshape(1, H, 1, 1).to(x.dtype) y = (y_local + scale * y_gla.to(x.dtype)).transpose(1, 2).reshape(B, T, C) return self.out_proj(y) def step( self, x: torch.Tensor, state: torch.Tensor, z_norm: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: B = x.shape[0] H, D = self.n_head, self.d_head q = self.q_proj(x).view(B, H, D) k = self.k_proj(x).view(B, H, D) v = self.v_proj(x).view(B, H, D) log_g = -F.softplus(self.g_proj(x).float()) gate = log_g.exp() q_k = F.elu(q.float()) + 1.0 k_k = F.elu(k.float()) + 1.0 v_f = v.float() g4 = gate.unsqueeze(-1).unsqueeze(-1) state = g4 * state + torch.einsum("bhd,bhe->bhde", k_k, v_f) g3 = gate.unsqueeze(-1) z_norm = g3 * z_norm + k_k num = torch.einsum("bhd,bhde->bhe", q_k, state) den = torch.einsum("bhd,bhd->bh", q_k, z_norm).clamp_(min=1.0).unsqueeze(-1) y_gla = (num / den).to(x.dtype) scale = self.gla_scale.view(1, H, 1) y = (y_gla * scale).reshape(B, H * D) return self.out_proj(y), state, z_norm class CPUGPTBlock(nn.Module): def __init__(self, cfg: CPUGPTConfig, layer_idx: int): super().__init__() is_gla = _layer_is_gla(layer_idx, cfg.n_layer, cfg.layer_pattern) self.mixer = GLAMixer(cfg) if is_gla else FNOSeqMixer(cfg) self.ffn = SwiGLU(cfg) self.ln1 = nn.RMSNorm(cfg.n_embd) self.ln2 = nn.RMSNorm(cfg.n_embd) self.is_gla = is_gla def forward(self, x: torch.Tensor) -> torch.Tensor: x = x + self.mixer(self.ln1(x)) x = x + self.ffn(self.ln2(x)) return x def step(self, x: torch.Tensor, bstate: dict) -> tuple: h = self.ln1(x) mixer = self.mixer if isinstance(mixer, GLAMixer): h_out, bstate["gla_state"], bstate["z_norm"] = mixer.step( h, bstate["gla_state"], bstate["z_norm"] ) else: assert isinstance(mixer, FNOSeqMixer) h_out, bstate["pos"] = mixer.step(h, bstate["buf"], bstate["pos"]) x = x + h_out x = x + self.ffn(self.ln2(x)) return x, bstate class _ChunkedCEFn(torch.autograd.Function): @staticmethod def forward(ctx, x_flat, weight, targets, chunk_size): N = x_flat.shape[0] use_bf16 = x_flat.dtype == torch.bfloat16 count = (targets != -1).sum().clamp(min=1) w_g = weight.bfloat16() if use_bf16 else weight.float() loss_sum = torch.zeros((), dtype=torch.float32, device=x_flat.device) for start in range(0, N, chunk_size): end = min(start + chunk_size, N) x_c = x_flat[start:end] t_c = targets[start:end] logits_c = (x_c @ w_g.T).float() logits_c = 15.0 * torch.tanh(logits_c * (1.0 / 15.0)) loss_sum = loss_sum + F.cross_entropy( logits_c, t_c, ignore_index=-1, reduction="sum" ) ctx.save_for_backward(x_flat, weight, targets) ctx.chunk_size = chunk_size ctx.count = int(count.item()) ctx.use_bf16 = use_bf16 return loss_sum / count @staticmethod def backward(ctx, grad_out): x_flat, weight, targets = ctx.saved_tensors N, C = x_flat.shape chunk_size = ctx.chunk_size use_bf16 = ctx.use_bf16 scale = grad_out.item() / ctx.count w_g = weight.bfloat16() if use_bf16 else weight.float() dx = torch.zeros(N, C, dtype=torch.float32, device=x_flat.device) dw = torch.zeros_like(weight, dtype=torch.float32) for start in range(0, N, chunk_size): end = min(start + chunk_size, N) x_c = x_flat[start:end] t_c = targets[start:end] logits_c = (x_c @ w_g.T).float() tanh_c = torch.tanh(logits_c * (1.0 / 15.0)) logits_cap = 15.0 * tanh_c probs_c = torch.softmax(logits_cap, dim=-1) valid = t_c != -1 probs_c[~valid] = 0.0 if valid.any(): probs_c[valid, t_c[valid]] -= 1.0 d_logits_c = probs_c * (1.0 - tanh_c * tanh_c) * scale if use_bf16: dx[start:end] = (d_logits_c.bfloat16() @ w_g).float() else: dx[start:end] = d_logits_c @ w_g if use_bf16: dw.add_((d_logits_c.bfloat16().T @ x_c).float()) else: dw.addmm_(d_logits_c.T, x_c) return ( dx.to(x_flat.dtype), dw, None, None, ) def _chunked_cross_entropy( weight: torch.Tensor, x: torch.Tensor, targets: torch.Tensor, chunk: int = 512, ) -> torch.Tensor: B, T, C = x.shape return _ChunkedCEFn.apply( x.reshape(B * T, C), weight, targets.reshape(B * T), chunk ) class CPUGPT(nn.Module): def __init__(self, cfg: CPUGPTConfig): super().__init__() self.cfg = cfg self.wte = nn.Embedding(cfg.vocab_size, cfg.n_embd) self.blocks = nn.ModuleList([CPUGPTBlock(cfg, i) for i in range(cfg.n_layer)]) self.ln_out = nn.RMSNorm(cfg.n_embd) self.lm_head = nn.Linear(cfg.n_embd, cfg.vocab_size, bias=False) self.lm_head.weight = self.wte.weight self._init_weights() def _init_weights(self): nn.init.normal_(self.wte.weight, std=0.02) for block in self.blocks: if isinstance(block.mixer, GLAMixer): for proj in [ block.mixer.q_proj, block.mixer.k_proj, block.mixer.v_proj, block.mixer.out_proj, ]: nn.init.normal_(proj.weight, std=0.02) nn.init.normal_(block.ffn.gate.weight, std=0.02) nn.init.normal_(block.ffn.up.weight, std=0.02) nn.init.zeros_(block.ffn.down.weight) def param_count(self) -> int: return sum(p.numel() for p in self.parameters()) def prepare_inference(self): for block in self.blocks: if not block.is_gla: block.mixer.prepare_inference() def init_state(self, batch_size: int = 1, device=None) -> list: if device is None: device = next(self.parameters()).device H = self.cfg.n_head D = self.cfg.n_embd // H T = self.cfg.seq_len C = self.cfg.n_embd states = [] for block in self.blocks: if block.is_gla: states.append( { "gla_state": torch.zeros(batch_size, H, D, D, device=device), "z_norm": torch.zeros(batch_size, H, D, device=device), } ) else: states.append( { "buf": torch.zeros(batch_size, T, C, device=device), "pos": 0, } ) return states @torch.no_grad() def step(self, idx: torch.Tensor, states: list): x = self.wte(idx.unsqueeze(1) if idx.dim() == 1 else idx) if x.dim() == 3: x = x.squeeze(1) x = F.rms_norm(x, (x.size(-1),)) new_states = [] for block, bstate in zip(self.blocks, states): x, bstate = block.step(x, bstate) new_states.append(bstate) x = self.ln_out(x) logits = self.lm_head(x.unsqueeze(1)).squeeze(1).float() logits = 15.0 * torch.tanh(logits / 15.0) return logits, new_states @torch.no_grad() def generate( self, prompt_ids: torch.Tensor, max_new_tokens: int = 200, temperature: float = 0.8, top_k: int = 50, ) -> torch.Tensor: self.prepare_inference() device = prompt_ids.device states = self.init_state(batch_size=1, device=device) logits = torch.zeros(prompt_ids.shape[0], self.cfg.vocab_size, device=device) for i in range(prompt_ids.shape[1]): tok = prompt_ids[:, i] logits, states = self.step(tok, states) generated = [] for _ in range(max_new_tokens): if temperature == 0.0: next_tok = logits.argmax(dim=-1) else: scaled = logits / temperature if top_k > 0: topk_vals, _ = torch.topk(scaled, top_k) scaled = scaled.masked_fill( scaled < topk_vals[:, -1:], float("-inf") ) probs = torch.softmax(scaled, dim=-1) next_tok = torch.multinomial(probs, num_samples=1).squeeze(1) generated.append(next_tok) logits, states = self.step(next_tok, states) return torch.stack(generated, dim=1) def forward( self, idx: torch.Tensor, targets: Optional[torch.Tensor] = None, reduction: str = "mean", ) -> torch.Tensor: B, T = idx.shape assert T <= self.cfg.seq_len, f"Sequence {T} > max {self.cfg.seq_len}" x = self.wte(idx) x = F.rms_norm(x, (x.size(-1),)) if _CPU_HAS_BF16: with torch.autocast(device_type="cpu", dtype=torch.bfloat16): for block in self.blocks: x = block(x) else: for block in self.blocks: x = block(x) x = self.ln_out(x) if targets is None: logits = self.lm_head(x).float() logits = 15.0 * torch.tanh(logits / 15.0) return logits return _chunked_cross_entropy(self.lm_head.weight, x, targets) def gpt2_small_config(**overrides) -> CPUGPTConfig: cfg = CPUGPTConfig( vocab_size=50257, seq_len=1024, n_layer=12, n_embd=768, n_head=12, fno_modes=256, gla_chunk=64, ffn_hidden=2048, layer_pattern="SSSL", ) for k, v in overrides.items(): setattr(cfg, k, v) return cfg def smoke_config(**overrides) -> CPUGPTConfig: cfg = CPUGPTConfig( vocab_size=50257, seq_len=256, n_layer=4, n_embd=256, n_head=4, fno_modes=64, gla_chunk=64, ffn_hidden=512, layer_pattern="SSSL", ) for k, v in overrides.items(): setattr(cfg, k, v) return cfg def gpt2_8b_config() -> CPUGPTConfig: return CPUGPTConfig( n_layer=32, n_embd=4096, n_head=32, ffn_hidden=14336, fno_modes=512, gla_chunk=256, seq_len=2048, layer_pattern="SSSL", ) def gpt2_1b_config() -> CPUGPTConfig: return CPUGPTConfig( n_layer=24, n_embd=2048, n_head=16, ffn_hidden=5632, fno_modes=512, gla_chunk=256, seq_len=2048, layer_pattern="SSSL", vocab_size=50257, ) def gpt2_1b_optimized_config() -> CPUGPTConfig: return CPUGPTConfig( n_layer=24, n_embd=2048, n_head=16, ffn_hidden=5632, fno_modes=512, gla_chunk=512, seq_len=2048, layer_pattern="SSSL", vocab_size=50257, fft_dtype="bf16", use_flash_fft=True, ) def get_config(name: str) -> CPUGPTConfig: configs = { "gpt2-small": gpt2_small_config, "smoke": smoke_config, "gpt2-8b": gpt2_8b_config, "gpt2-1b": gpt2_1b_config, "gpt2-1b-optimized": gpt2_1b_optimized_config, } if name not in configs: raise ValueError(f"Unknown config '{name}'. Available: {list(configs.keys())}") return configs[name]()