File size: 9,514 Bytes
e40c47f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 | """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
|