| """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): |
| |
| |
| |
| 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) |
| |
| |
| 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) |
| |
| 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) |
| |
| |
| 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 |
|
|
