"""Inference adapter with standard Transformers checkpoint/Hub loading. Two AutoClass entry points share one backbone module and one checkpoint layout: * ``NovaModel`` (``AutoModel``): final-normalized hidden states, optional per-layer hidden states, no vocabulary projection. * ``NovaForCausalLM`` (``AutoModelForCausalLM``): logits and ``generate()``. """ from contextlib import nullcontext import importlib.util import sys from typing import TYPE_CHECKING import warnings import torch from torch import nn import transformers from transformers import GenerationConfig, GenerationMixin, PreTrainedModel from transformers.modeling_outputs import BaseModelOutput, CausalLMOutput from .configuration_nova import AUTO_MAP, NovaConfig from .generation_nova import generate_batched, validate_inputs # Transformers resolves remote-code dependencies by scanning relative imports # (also inside this TYPE_CHECKING block) and downloads the whole closure from the # Hub. Declare it here without importing GPU libraries at module import time; the # flash-only modules guard their own optional imports. if TYPE_CHECKING: from .architecture import build_model from .architecture_config import ModelOptions from .cpu_backend import TorchBackbone from .fast_decode import enable from .gdn2 import GatedDeltaNet2 from .gdn2_spec import GDN2Spec from .sequence_utils import flatten_documents VALIDATED_TRANSFORMERS = "4.57" if not transformers.__version__.startswith(VALIDATED_TRANSFORMERS + "."): warnings.warn( f"Nova-1B inference code was validated with transformers {VALIDATED_TRANSFORMERS}.x; " f"found {transformers.__version__}. Pin transformers=={VALIDATED_TRANSFORMERS}.6 if " "loading or generation misbehaves.", stacklevel=2) def build_backbone(config): if config.inference_backend in ("torch", "cpu"): from .cpu_backend import TorchBackbone return TorchBackbone(config) if sys.platform != "linux" or not torch.cuda.is_available(): raise RuntimeError("Flash inference requires Linux + NVIDIA CUDA. Use inference_backend='torch' without FlashAttention, or example_cpu.py for CPU.") for package in ("olmo_core", "fla", "flash_attn"): if importlib.util.find_spec(package) is None: raise ImportError("Flash backend dependencies are missing. Install .[flash] and flash-attn as described in README, or select inference_backend='torch'.") from .architecture import build_model from .architecture_config import ModelOptions architecture_config = build_model(config.architecture_preset, vocab_size=config.vocab_size, model_options=ModelOptions(**dict(config.architecture_options, attention_backend="flash_2")), sequence_length=config.max_position_embeddings) return architecture_config.build(init_device="cpu") def _reject_token_type_ids(token_type_ids, input_ids): if token_type_ids is not None and (input_ids is None or token_type_ids.shape != input_ids.shape or torch.any(token_type_ids != 0)): raise ValueError("Nova does not use nonzero token_type_ids") def _prompt_rows(input_ids, attention_mask, position_ids, context): """Validate a left/right padded batch and yield (row, unpadded positions).""" mask = validate_inputs(input_ids, attention_mask, context) for i in range(input_ids.shape[0]): positions = mask[i].nonzero(as_tuple=True)[0] if position_ids is not None and not torch.equal(position_ids[i, positions], torch.arange(len(positions), device=input_ids.device)): raise ValueError("Positions must start at zero for each unpadded prompt") yield i, positions def _scatter_rows(rows, length): """Place per-prompt [tokens, features] results back into a zero-padded batch.""" return torch.stack([values.new_zeros(length, values.shape[-1]).index_copy(0, positions, values) for positions, values in rows]) class NovaPreTrainedModel(PreTrainedModel): """Shared checkpoint layout, precision rules and loading checks.""" config_class = NovaConfig base_model_prefix = "backbone" main_input_name = "input_ids" supports_gradient_checkpointing = False _supports_cache_class = False _tied_weights_keys = ["backbone.lm_head.w_out.weight"] # Transformers 4.57.6 applies this also for BF16 from_pretrained loading. _keep_in_fp32_modules_strict = ["A_log", "dt_bias"] def __init__(self, config): super().__init__(config) self.backbone = build_backbone(config) self.post_init() def _init_weights(self, module): # Published inference requires trained weights; nothing is randomly initialized. pass def get_input_embeddings(self): return self.backbone.embeddings def get_output_embeddings(self): return self.backbone.lm_head.w_out def set_input_embeddings(self, value): self.backbone.embeddings = value def set_output_embeddings(self, value): self.backbone.lm_head.w_out = value def _apply(self, function, recurse=True): masters = [(n, p.detach()) for n, p in self.backbone.named_parameters() if n.endswith((".A_log", ".dt_bias"))] result = super()._apply(function, recurse=recurse) for name, original in masters: parameter = self.backbone.get_parameter(name) parameter.data = original.to(device=parameter.device, dtype=torch.float32) return result def save_pretrained(self, *args, **kwargs): # Every saved bundle advertises all remote-code entry points, not only the # class that happened to be saved. self.config.auto_map = dict(getattr(self.config, "auto_map", None) or {}, **AUTO_MAP) return super().save_pretrained(*args, **kwargs) @classmethod def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs): # Cached decoding calls backbone modules directly; arbitrary Accelerate dispatch # and offload hooks are not a supported execution path. device_map = kwargs.get("device_map") if device_map is not None: raise ValueError("device_map/offloading is not supported. Load then call model.to('cuda').") if kwargs.get("quantization_config") is not None or kwargs.get("load_in_4bit") or kwargs.get("load_in_8bit"): raise ValueError("Quantized Nova inference is not validated in this release") return_info = kwargs.pop("output_loading_info", False) model, info = super().from_pretrained(pretrained_model_name_or_path, *args, output_loading_info=True, **kwargs) if any(info.get(key) for key in ("missing_keys", "unexpected_keys", "mismatched_keys", "error_msgs")): raise ValueError(f"Incomplete or incompatible Nova checkpoint: {info}") return (model, info) if return_info else model def _compute_context(self, ids): if (ids.is_cuda and self.config.compute_dtype != "float32" and not torch.is_autocast_enabled("cuda")): return torch.autocast("cuda", dtype=getattr(torch, self.config.compute_dtype)) return nullcontext() def _backbone_logits(self, ids): with self._compute_context(ids): result = self.backbone(ids, return_logits=True) logits = result if isinstance(result, torch.Tensor) else result.logits return logits[..., :self.config.tokenizer_vocab_size] def _backbone_hidden(self, ids, collect_layers): """Final-normalized hidden states for one unpadded prompt batch. Per-layer states (embedding output followed by each block output) are captured with forward hooks so both the torch and the flash backbone, which share the ``embeddings`` / ``blocks`` / ``lm_head.norm`` layout, expose the same AutoModel outputs. """ backbone = self.backbone captured = {"final": None, "layers": []} handles = [] def record(_module, _inputs, output): captured["layers"].append(output if isinstance(output, torch.Tensor) else output[0]) def record_final(_module, _inputs, output): captured["final"] = output torch_backend = self.config.inference_backend in ("torch", "cpu") try: if collect_layers: handles.append(backbone.embeddings.register_forward_hook(record)) blocks = backbone.blocks.values() if isinstance(backbone.blocks, nn.ModuleDict) else backbone.blocks for block in blocks: handles.append(block.register_forward_hook(record)) if not torch_backend: handles.append(backbone.lm_head.norm.register_forward_hook(record_final)) with self._compute_context(ids): if torch_backend: captured["final"] = backbone(ids, return_logits=False) else: backbone(ids, return_logits=True) finally: for handle in handles: handle.remove() if captured["final"] is None: raise RuntimeError("Backbone did not expose its final hidden states") if collect_layers and len(captured["layers"]) != self.config.num_hidden_layers + 1: raise RuntimeError(f"Expected {self.config.num_hidden_layers + 1} hidden states, " f"captured {len(captured['layers'])}") return captured["final"], captured["layers"] class NovaModel(NovaPreTrainedModel): """Backbone for ``AutoModel``: hidden states without the vocabulary projection. ``last_hidden_state`` is the output of the final RMSNorm that precedes the (tied) output projection, so ``model.get_output_embeddings()(last_hidden_state)`` reproduces ``NovaForCausalLM`` logits. ``hidden_states`` follows the Transformers convention: embedding output first, then one entry per block, before the final norm. Padded positions are zero-filled. """ @torch.no_grad() def forward(self, input_ids=None, attention_mask=None, position_ids=None, past_key_values=None, inputs_embeds=None, use_cache=False, output_attentions=None, output_hidden_states=None, return_dict=True, token_type_ids=None, **kwargs): if inputs_embeds is not None: raise ValueError("Nova accepts token IDs only; inputs_embeds is not supported") if past_key_values is not None or use_cache: raise ValueError("Generation caches are managed by generate(); use_cache must be False") if output_attentions: raise ValueError("Nova does not expose attention weights (GDN2 layers have none)") if kwargs: raise TypeError(f"Unsupported forward arguments: {sorted(kwargs)}") _reject_token_type_ids(token_type_ids, input_ids) collect = bool(self.config.output_hidden_states if output_hidden_states is None else output_hidden_states) finals, layers = [], [] for i, positions in _prompt_rows(input_ids, attention_mask, position_ids, self.config.max_position_embeddings): final, states = self._backbone_hidden(input_ids[i, positions].unsqueeze(0), collect) finals.append((positions, final[0])) layers.append([(positions, state[0]) for state in states]) length = input_ids.shape[1] last_hidden_state = _scatter_rows(finals, length) hidden_states = None if collect: hidden_states = tuple(_scatter_rows([row[k] for row in layers], length) for k in range(len(layers[0]))) if not return_dict: return tuple(value for value in (last_hidden_state, hidden_states) if value is not None) return BaseModelOutput(last_hidden_state=last_hidden_state, hidden_states=hidden_states) class NovaForCausalLM(NovaPreTrainedModel, GenerationMixin): """``AutoModelForCausalLM`` entry point: logits and cached generation.""" @torch.no_grad() def forward(self, input_ids, attention_mask=None, position_ids=None, past_key_values=None, use_cache=False, return_dict=True, logits_to_keep=0, labels=None, token_type_ids=None, **kwargs): if labels is not None: raise ValueError("This distribution provides inference only") if past_key_values is not None or use_cache: raise ValueError("Generation caches are managed by generate(); use_cache must be False") if kwargs: raise TypeError(f"Unsupported forward arguments: {sorted(kwargs)}") _reject_token_type_ids(token_type_ids, input_ids) rows = [(positions, self._backbone_logits(input_ids[i, positions].unsqueeze(0))[0]) for i, positions in _prompt_rows(input_ids, attention_mask, position_ids, self.config.max_position_embeddings)] logits = _scatter_rows(rows, input_ids.shape[1]) if logits_to_keep: logits = logits[:, -logits_to_keep:] return CausalLMOutput(logits=logits) if return_dict else (logits,) @torch.no_grad() def generate(self, input_ids=None, attention_mask=None, generation_config=None, **kwargs): settings = generation_config or self.generation_config _reject_token_type_ids(kwargs.pop("token_type_ids", None), input_ids) supported = {"max_new_tokens", "max_length", "do_sample", "temperature", "top_p", "top_k", "eos_token_id", "pad_token_id", "num_return_sequences", "generation_batch_size", "num_beams", "num_beam_groups", "return_dict_in_generate", "output_scores", "output_logits", "output_attentions", "output_hidden_states", "repetition_penalty", "no_repeat_ngram_size", "min_new_tokens", "min_length"} unknown = set(kwargs) - supported if unknown: raise TypeError(f"Unsupported generation arguments: {sorted(unknown)}") def option(key, default=None): value = kwargs.get(key, getattr(settings, key, default)) return default if value is None else value for key, default in {"num_beams": 1, "num_beam_groups": 1, "return_dict_in_generate": False, "output_scores": False, "output_logits": False, "output_attentions": False, "output_hidden_states": False, "repetition_penalty": 1.0, "no_repeat_ngram_size": 0, "min_new_tokens": 0, "min_length": 0}.items(): if option(key, default) != default: raise ValueError(f"{key} is not supported; expected {default!r}") # A supplied GenerationConfig must not silently enable an unsupported # processor. Backbone cache ownership is independent of HF's use_cache flag. defaults = GenerationConfig().to_dict() handled = supported | {"use_cache", "bos_token_id", "transformers_version"} for key, value in settings.to_dict().items(): if key not in handled and not key.startswith("_") and value != defaults.get(key): raise ValueError(f"generation_config.{key} is not supported") mask = validate_inputs(input_ids, attention_mask, self.config.max_position_embeddings) budget = option("max_new_tokens") if budget is None: budget = option("max_length", self.config.max_position_embeddings) - input_ids.shape[1] if type(budget) is not int or budget <= 0: raise ValueError("max_new_tokens must be a positive integer") returns = option("num_return_sequences", 1) limit = option("generation_batch_size", 1) if type(returns) is not int or returns < 1 or type(limit) is not int or limit < 1: raise ValueError("Return count and generation batch size must be positive integers") if returns > 1 and not option("do_sample", False): raise ValueError("Multiple return sequences require do_sample=True") ids = input_ids.repeat_interleave(returns, dim=0) mask = mask.repeat_interleave(returns, dim=0) eos = option("eos_token_id", self.config.eos_token_id) eos_ids = set(eos if isinstance(eos, (list, tuple)) else [] if eos is None else [eos]) pad = option("pad_token_id", self.config.pad_token_id) return generate_batched(self, ids, mask, budget, eos_ids, pad, option, limit) NovaModel.register_for_auto_class("AutoModel") NovaForCausalLM.register_for_auto_class("AutoModelForCausalLM")