"""OLMo SequenceMixer adapter for the upstream FLA GatedDeltaNet2 layer. All recurrence and convolution kernels come from flash-linear-attention 0.5.2. This module owns OLMo initialization, document packing and inference lifecycle. """ from __future__ import annotations from dataclasses import dataclass import math from typing import TYPE_CHECKING import torch from torch import nn # Optional accelerated-path dependencies (OLMo-core, FLA). Each import below is # wrapped so Transformers' remote-code import scan does not require them for the # default torch backend; modeling_nova imports this module only after checking # that the flash extras are installed. try: from olmo_core.config import DType from olmo_core.nn.attention.base import SequenceMixer, SequenceMixerConfig except ImportError as error: # pragma: no cover - exercised only without .[flash] raise ImportError( "nova.gdn2 requires the flash extras (OLMo-core and flash-linear-attention). " "Install `.[flash]` or load the model with inference_backend='torch'." ) from error from .sequence_utils import flatten_documents from .gdn2_spec import GDN2Spec if TYPE_CHECKING: # pragma: no cover try: from olmo_core.nn.transformer.init import InitMethod except ImportError: pass class GatedDeltaNet2(SequenceMixer): """Tensor-in/tensor-out wrapper around ``fla.layers.GatedDeltaNet2``. Each mixer owns one FLA Cache with a local layer index of zero. The cache holds both recurrent and short-convolution states and is never serialized. """ def __init__( self, *, d_model: int, n_heads: int, head_dim: int | None = None, n_v_heads: int | None = None, expand_v: float = 2.0, allow_neg_eigval: bool = True, conv_size: int = 4, conv_bias: bool = False, norm_eps: float = 1e-5, dtype: torch.dtype = torch.float32, init_device: str = "cpu", ) -> None: super().__init__() if n_heads <= 0: raise ValueError("n_heads must be positive") self.spec = GDN2Spec( d_model=d_model, n_heads=n_heads, head_dim=head_dim if head_dim is not None else d_model // n_heads, n_v_heads=n_v_heads, expand_v=expand_v, conv_size=conv_size, conv_bias=conv_bias, ) if not math.isfinite(norm_eps) or norm_eps <= 0: raise ValueError("norm_eps must be finite and positive") try: from fla.layers.gdn2 import GatedDeltaNet2 as FLAGatedDeltaNet2 except ImportError as error: # pragma: no cover raise ImportError("GatedDeltaNet2 requires flash-linear-attention 0.5.2 (`.[flash]`)") from error # FLA accepts **kwargs but does not implement OLMo's init_device/dtype # arguments. Use PyTorch's device context, including for meta tensors. with torch.device(init_device): self.layer = FLAGatedDeltaNet2( hidden_size=d_model, num_heads=n_heads, num_v_heads=self.spec.value_heads, head_dim=self.spec.head_dim, expand_v=expand_v, allow_neg_eigval=allow_neg_eigval, use_short_conv=True, conv_size=conv_size, conv_bias=conv_bias, norm_eps=norm_eps, mode="chunk", layer_idx=0, ) self.layer.to(dtype=dtype) # These two parameters must retain their initialization precision. for name in ("A_log", "dt_bias"): param = nn.Parameter(getattr(self.layer, name).float()) # Do not attach optimizer-only attributes before registration. # Transformers' init_empty_weights reconstructs Parameters from # their __dict__; torch.nn.Parameter rejects _no_weight_decay. setattr(self.layer, name, param) self.d_model = d_model self._decode_cache = None self._decode_batch_size: int | None = None self._decode_max_seq_len: int | None = None self._decode_tokens = 0 @property def decode_enabled(self) -> bool: return self._decode_cache is not None def enable_decode(self, batch_size: int, max_seq_len: int) -> None: if self.training: raise ValueError("Call model.eval() before enabling incremental decoding") if batch_size <= 0 or max_seq_len <= 0: raise ValueError("batch_size and max_seq_len must be positive") try: from fla.models.utils import Cache except ImportError as error: # pragma: no cover raise ImportError("GatedDeltaNet2 requires flash-linear-attention 0.5.2 (`.[flash]`)") from error self._decode_cache = Cache() self._decode_batch_size = batch_size self._decode_max_seq_len = max_seq_len self._decode_tokens = 0 def disable_decode(self) -> None: self._decode_cache = None self._decode_batch_size = None self._decode_max_seq_len = None self._decode_tokens = 0 def train(self, mode: bool = True): if mode and self.decode_enabled: raise ValueError("Call fast_decode.disable(model) before returning to training") return super().train(mode) def forward(self, x: torch.Tensor, cu_doc_lens=None, **kwargs) -> torch.Tensor: # OLMo forwards attention-specific kwargs (positions, max_doc_len, ...). # Passing them to FLA as a padding attention_mask would lose doc resets. del kwargs if x.ndim != 3 or x.shape[-1] != self.d_model or x.shape[0] == 0 or x.shape[1] == 0: raise ValueError("Expected nonempty x with shape [batch, sequence, d_model]") original_shape = x.shape if self.decode_enabled: if self.training or torch.is_grad_enabled(): raise ValueError("Cached decoding requires eval() and torch.no_grad()/inference_mode()") if cu_doc_lens is not None: raise ValueError("Cached decoding does not accept packed documents; reset per prompt batch") if x.shape[0] != self._decode_batch_size: raise ValueError("Decode batch size changed; call fast_decode.reset() first") if self._decode_tokens + x.shape[1] > self._decode_max_seq_len: raise ValueError("Decode context exceeds max_seq_len; allocate a larger cache") else: x = flatten_documents(x, cu_doc_lens) # FLA 0.5.2 switches short eval sequences to a forward-only recurrent # kernel even when gradients are enabled. Explicitly use its training # branch for eval-with-grad, without changing the public wrapper mode. original_training = self.layer.training if torch.is_grad_enabled() and not original_training: self.layer.training = True try: out, _, _ = self.layer( hidden_states=x, attention_mask=None, cu_seqlens=cu_doc_lens, past_key_values=self._decode_cache, use_cache=self.decode_enabled, ) finally: self.layer.training = original_training if self.decode_enabled: self._decode_tokens += original_shape[1] return out.reshape(original_shape) @torch.no_grad() def init_weights( self, *, init_method: "InitMethod", d_model: int, block_idx: int, num_blocks: int, std: float = 0.02, generator: torch.Generator | None = None, ) -> None: """Reinitialize every parameter after OLMo's meta -> to_empty step.""" try: from olmo_core.nn.transformer.init import InitMethod, init_linear except ImportError as error: # pragma: no cover raise ImportError("GatedDeltaNet2.init_weights requires OLMo-core (`.[flash]`)") from error if init_method == InitMethod.fan_in: raise NotImplementedError("fan_in initialization is not supported for GDN2") if init_method == InitMethod.normalized: std = d_model ** -0.5 layer = self.layer for module in layer.modules(): if isinstance(module, nn.Linear) and module is not layer.o_proj: init_linear(module, std=std, generator=generator) for conv in (layer.q_conv1d, layer.k_conv1d, layer.v_conv1d): init_linear(conv, std=std, generator=generator) layer.o_norm.reset_parameters() layer.A_log.uniform_(1, 16, generator=generator).log_() # FLA's per-key-channel time steps, initialized through inverse softplus. dt = torch.empty_like(layer.dt_bias, dtype=torch.float32).uniform_( math.log(0.001), math.log(0.1), generator=generator, ).exp_().clamp_min_(1e-4) layer.dt_bias.copy_(dt + torch.log(-torch.expm1(-dt))) if init_method in (InitMethod.llama, InitMethod.normalized): std /= math.sqrt(2 * num_blocks) elif init_method == InitMethod.llama_depth: std /= math.sqrt(2 * (block_idx + 1)) init_linear(layer.o_proj, std=std, generator=generator) self.disable_decode() def apply_tp(self, tp_mesh, **kwargs): raise NotImplementedError("GDN2 adapter does not support tensor parallelism; use FSDP") def apply_cp(self, cp_mesh, **kwargs): raise NotImplementedError("GDN2 adapter does not support context parallelism; use FSDP") def num_flops_per_token(self, seq_len: int) -> int: """Analytical forward+backward estimate; excludes chunk-solver overhead.""" del seq_len linear = sum(m.weight.numel() for m in self.layer.modules() if isinstance(m, nn.Linear)) conv = self.spec.conv_size * (2 * self.spec.key_dim + self.spec.value_dim) # State decay, erase read, outer product, and query read. return 3 * (2 * (linear + conv) + 7 * self.spec.state_elements) @SequenceMixerConfig.register("gated_delta_net2") @dataclass class GatedDeltaNet2Config(SequenceMixerConfig[GatedDeltaNet2]): n_heads: int = 16 n_v_heads: int | None = None head_dim: int | None = None expand_v: float = 2.0 allow_neg_eigval: bool = True conv_size: int = 4 conv_bias: bool = False norm_eps: float = 1e-5 dtype: DType = DType.float32 def spec(self, d_model: int) -> GDN2Spec: if self.n_heads <= 0: raise ValueError("n_heads must be positive") return GDN2Spec( d_model=d_model, n_heads=self.n_heads, n_v_heads=self.n_v_heads, head_dim=self.head_dim if self.head_dim is not None else d_model // self.n_heads, expand_v=self.expand_v, conv_size=self.conv_size, conv_bias=self.conv_bias, ) def num_params(self, d_model: int) -> int: return self.spec(d_model).num_params def build(self, d_model: int, *, layer_idx: int, n_layers: int, init_device: str = "cpu", cache=None) -> GatedDeltaNet2: del layer_idx, n_layers, cache # FLA keeps a private one-layer cache per mixer. self.spec(d_model) return GatedDeltaNet2( d_model=d_model, n_heads=self.n_heads, n_v_heads=self.n_v_heads, head_dim=self.head_dim, expand_v=self.expand_v, allow_neg_eigval=self.allow_neg_eigval, conv_size=self.conv_size, conv_bias=self.conv_bias, norm_eps=self.norm_eps, dtype=self.dtype.as_pt(), init_device=init_device, )