"""Trainable adapter for NVIDIA's BF16-adapted Evo 2 1B checkpoint. Uses Vortex's differentiable blocks, not Evo2.forward (an inference-only wrapper). Retains NeMo's original trainable modal parameters and fixed medium decay buffer. Checkpoint conversion must reorder Megatron's head-major QKV weights, and all weights are loaded strictly. Original FP8-sensitive Arc checkpoints are NOT inputs. """ from pathlib import Path import torch from torch import nn from torch.nn import functional as F from torch.utils.checkpoint import checkpoint from transformers import AutoConfig, AutoModel, PretrainedConfig, PreTrainedModel from transformers.modeling_outputs import BaseModelOutput class Evo2BF16Config(PretrainedConfig): model_type = "rnaseek_evo2_bf16" def __init__(self, hidden_size=1920, intermediate_size=5120, num_hidden_layers=25, num_attention_heads=15, vocab_size=512, state_size=16, pattern="SDH*SDHSDH*SDHSDH*SDHSDH*", medium_groups=128, short_groups=128, **kwargs): kwargs.setdefault("pad_token_id", 1) kwargs.setdefault("eos_token_id", 0) super().__init__(**kwargs) self.hidden_size = hidden_size self.intermediate_size = intermediate_size self.num_hidden_layers = num_hidden_layers self.num_attention_heads = num_attention_heads self.vocab_size = vocab_size self.state_size = state_size self.pattern = pattern self.medium_groups = medium_groups self.short_groups = short_groups self.max_position_embeddings = 8192 self.checkpoint_variant = "NVIDIA Evo 2 1B (BF16-adapted)" if len(pattern) != num_hidden_layers or set(pattern) - set("SDH*"): raise ValueError("Invalid Evo 2 layer pattern") class NemoRMSNorm(nn.Module): """Standard RMSNorm with FP32 accumulation and epsilon inside the square root.""" def __init__(self, width, eps=1e-6): super().__init__() self.scale = nn.Parameter(torch.ones(width)) self.eps = eps def forward(self, x): value = x.float() return (value * torch.rsqrt(value.square().mean(-1, keepdim=True) + self.eps) * self.scale.float()).to(x.dtype) class AttributeDict(dict): def __getattr__(self, key): try: return self[key] except KeyError as error: raise AttributeError(key) from error def vortex_config(config): return AttributeDict( hidden_size=config.hidden_size, num_filters=config.hidden_size, num_attention_heads=config.num_attention_heads, num_layers=config.num_hidden_layers, attn_layer_idxs=[i for i, s in enumerate(config.pattern) if s == "*"], hcl_layer_idxs=[i for i, s in enumerate(config.pattern) if s == "H"], hcm_layer_idxs=[i for i, s in enumerate(config.pattern) if s == "D"], hcs_layer_idxs=[i for i, s in enumerate(config.pattern) if s == "S"], hcm_filter_length=128, hcs_filter_length=7, hcm_filter_groups=config.medium_groups, hcs_filter_groups=config.short_groups, short_filter_length=3, short_filter_bias=False, state_size=config.state_size, eps=1e-6, rotary_emb_base=10000, inner_size_multiple_of=16, inner_mlp_size=config.intermediate_size, proj_groups=1, hyena_filter_groups=1, column_split_hyena=False, column_split=True, interleave=True, evo2_style_activations=True, model_parallel_size=1, pipe_parallel_size=1, mha_out_proj_bias=True, hyena_out_proj_bias=True, qkv_proj_bias=False, use_fp8_input_projections=False, use_flash_attn=True, use_flash_rmsnorm=False, use_flash_depthwise=False, use_flashfft=False, use_hcs_kernel=False, use_hcm_kernel=False, use_hcl_kernel=False, inference_mode=False, prefill_style="fft", mlp_activation="gelu", print_activations=False, params_dtype=torch.float32, attn_block_dtype=torch.float32, hyena_block_dtype=torch.float32, mlp_dtype=torch.float32, ) def preserve_filter_parameters(cascade, symbol): # Narrow subclasses preserve the pretrained optimization parameterization. # Modal equation follows NVIDIA/NeMo's Apache-2.0 ImplicitModalFilter; # p, gamma and R are independent trainable parameters, not merged poles. from vortex.model.model import HyenaCascade class NemoConvolution: def parallel_forward(self, u, inference_params=None, padding_mask=None): if inference_params is not None or padding_mask is not None: raise ValueError("This regression adapter uses full causal sequences, without generation caches") length, width = u.shape[1], self.hidden_size projected = F.conv1d(u.transpose(1, 2), self.short_filter_weight, padding=2, groups=3 * width)[..., :length] postgate, pregate, value = (projected[:, offset::3] for offset in range(3)) gated = pregate * value if self.h is None: kernel = self.compute_filter(length, u.device)[0].squeeze(0) else: kernel = self.h.squeeze(1)[:, :length] kernel = kernel.repeat_interleave(width // self.hyena_filter_groups, 0) # Match NeMo's unfused training arithmetic: FP32 FFT and direct term, # then cast to activation dtype BEFORE applying the outer gate. with torch.autocast("cuda", enabled=False): fft_size = 2 * length filtered = torch.fft.irfft(torch.fft.rfft(gated.float(), n=fft_size) * torch.fft.rfft(kernel.float(), n=fft_size), n=fft_size)[..., :length] filtered = filtered + gated.float() * self.D.float()[None, :, None] return (postgate * filtered.to(value.dtype)).transpose(1, 2), None class ModalCascade(NemoConvolution, HyenaCascade): @property def log_poles(self): return (-self.p.float().exp() * self.gamma.float().exp()).unsqueeze(-1) @property def residues(self): return self.R.float() class MediumCascade(NemoConvolution, HyenaCascade): @property def h(self): return (self.raw_h.float() * self.decay.float()).unsqueeze(1) if symbol == "H": shape = cascade.residues.shape del cascade.log_poles, cascade.residues cascade.p = nn.Parameter(torch.zeros(shape)) cascade.gamma = nn.Parameter(torch.zeros(shape)) cascade.R = nn.Parameter(torch.zeros(shape)) cascade.__class__ = ModalCascade elif symbol == "D": shape = (cascade.h.shape[0], cascade.h.shape[-1]) del cascade.h cascade.raw_h = nn.Parameter(torch.zeros(shape)) cascade.register_buffer("decay", torch.ones(shape)) cascade.__class__ = MediumCascade class Evo2BF16Backbone(PreTrainedModel): config_class = Evo2BF16Config base_model_prefix = "encoder" supports_gradient_checkpointing = True def __init__(self, config): super().__init__(config) from vortex.model.model import get_block settings = vortex_config(config) self.embedding_layer = nn.Embedding(config.vocab_size, config.hidden_size) self.blocks = nn.ModuleList(get_block(settings, i) for i in range(config.num_hidden_layers)) for symbol, block in zip(config.pattern, self.blocks): block.pre_norm = NemoRMSNorm(config.hidden_size) block.post_norm = NemoRMSNorm(config.hidden_size) if symbol != "*": preserve_filter_parameters(block.filter, symbol) # Pure BF16/FP32 linear, with trainable FP32 master weights. block.projections = nn.Linear(config.hidden_size, 3 * config.hidden_size, bias=False) self.norm = NemoRMSNorm(config.hidden_size) self.gradient_checkpointing = False def gradient_checkpointing_enable(self, **kwargs): self.gradient_checkpointing = True def gradient_checkpointing_disable(self): self.gradient_checkpointing = False def get_input_embeddings(self): return self.embedding_layer def forward(self, input_ids, attention_mask=None, return_dict=True, **kwargs): if input_ids.device.type != "cuda": raise ValueError("Evo 2's FlashAttention backend requires a compatible NVIDIA GPU") if attention_mask is not None and torch.any(attention_mask[:, 1:] > attention_mask[:, :-1]): raise ValueError("Evo 2 regression supports right padding only") hidden = self.embedding_layer(input_ids) # Causality makes future right-padding irrelevant to unpadded states. # Preserve FP32 residual streams, as specified in the source checkpoint. for block in self.blocks: if self.gradient_checkpointing and self.training: hidden = checkpoint(lambda x, layer=block: layer(x)[0], hidden, use_reentrant=False) else: hidden = block(hidden)[0] hidden = self.norm(hidden) return BaseModelOutput(last_hidden_state=hidden) if return_dict else (hidden,) AutoConfig.register(Evo2BF16Config.model_type, Evo2BF16Config) AutoModel.register(Evo2BF16Config, Evo2BF16Backbone) def load_backbone_strict(path, config): from safetensors.torch import load_file model = Evo2BF16Backbone(config) model.load_state_dict(load_file(str(Path(path) / "model.safetensors")), strict=True) return model