File size: 3,072 Bytes
84a4b93 | 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 | """Portable configuration for the trained Nova GDN2/GQA architecture."""
from transformers import PretrainedConfig
DEFAULT_OPTIONS = dict(d_model=1536, n_layers=26, n_heads=12, n_kv_heads=4,
head_dim=128, ffn_hidden=6144, rope_theta=500000, rope_full_precision=True,
layer_norm_eps=1e-5, layer_norm_bias=False, attn_qk_norm=True,
attention_bias=False, ffn_bias=False, lm_head_bias=False, tie_word_embeddings=True,
attention_backend="flash_2", enable_swa=True, swa_window=4096,
swa_full_attention_first_layer=False, swa_full_attention_last_layer=False,
gdn_expand_v=2.0, gdn_allow_neg_eigval=True, gdn_conv_size=4, gdn_conv_bias=False)
# Remote-code entry points written to config.json. The Hub resolves these with
# trust_remote_code=True; a local `pip install -e .` plus nova.register() maps the
# same classes without remote code.
AUTO_MAP = {
"AutoConfig": "configuration_nova.NovaConfig",
"AutoModel": "modeling_nova.NovaModel",
"AutoModelForCausalLM": "modeling_nova.NovaForCausalLM",
}
class NovaConfig(PretrainedConfig):
model_type = "nova_gdn2"
def __init__(self, architecture_preset="main", architecture_options=None, vocab_size=65536,
hidden_size=1536, num_hidden_layers=26, max_position_embeddings=8192,
fast_decode=True, compute_dtype="bfloat16", tokenizer_vocab_size=None,
generation_context_policy="remaining", inference_backend="torch", **kwargs):
kwargs.pop("use_cache", None)
kwargs.setdefault("bos_token_id", 1)
kwargs.setdefault("eos_token_id", 2)
kwargs.setdefault("pad_token_id", 0)
super().__init__(use_cache=False, **kwargs)
self.architecture_preset = architecture_preset
self.architecture_options = dict(DEFAULT_OPTIONS if architecture_options is None else architecture_options)
self.vocab_size = vocab_size
self.hidden_size = hidden_size
self.num_hidden_layers = num_hidden_layers
self.max_position_embeddings = max_position_embeddings
self.fast_decode = fast_decode
self.compute_dtype = compute_dtype
# Keep old exported configurations readable; new callers use torch/flash.
self.inference_backend = "flash" if inference_backend == "cuda" else inference_backend
self.tokenizer_vocab_size = tokenizer_vocab_size or vocab_size
self.generation_context_policy = generation_context_policy
if not 0 < self.tokenizer_vocab_size <= vocab_size:
raise ValueError("Tokenizer vocabulary must fit within the model vocabulary")
if compute_dtype not in ("float32", "float16", "bfloat16"):
raise ValueError("Unsupported compute_dtype")
if self.inference_backend not in ("torch", "flash", "cpu"):
raise ValueError("inference_backend must be torch, flash or cpu")
if generation_context_policy not in ("remaining", "error"):
raise ValueError("generation_context_policy must be remaining or error")
NovaConfig.register_for_auto_class("AutoConfig")
|