Oblivion FP16 checkpoint

0ba3a54 2 months ago

11 kB

# SPDX-License-Identifier: Apache-2.0
# Copyright 2026 The Oblivion Authors.
"""The fused Oblivion model for vLLM v1 -- multimodal, single route.
Oblivion subclasses the full native multimodal ``MuseGlimmerForCausalLM`` and
splices two donor macroblocks into the text decoder: a Qwen3.5 block after muse
layer 35 and a Gemma4 block after layer 47. Both donors are ALWAYS applied at the
canonical gate (``CANONICAL_GATE``) -- there is no per-request routing or gating.
Everything multimodal (vision tower, image/video embedding merge, the
``@MULTIMODAL_REGISTRY.register_processor``) is inherited unchanged, so image
prompts run through Glimmer's vision path.
* ``super().__init__`` builds the vision tower, the native ``self.model``
  (``MuseGlimmerModel``, 52 text layers at ``model.layers.0..51``), ``lm_head``,
  and multimodal token handling.
* Donor decoder layers are first-class vLLM modules built under the live fused
  ``vllm_config`` (so each Attention / GatedDeltaNet self-registers KV / mamba
  state), held in ``self.qwen_layers`` / ``self.gemma_layers`` (ModuleDicts keyed
  by source index, prefixes ``qwen_layers.{24..39}`` / ``gemma_layers.{48..59}``).
* ``forward`` runs the text loop with two seams: muse 0..35 -> qwen seam -> muse
  36..47 -> gemma seam -> muse 48..51 -> norm. It also captures auxiliary
  hidden states when spec decoding is active (for the EAGLE3 drafter).
The Qwen donor is recurrent (Gated DeltaNet), so the fused arch is hybrid: it
implements ``IsHybrid`` + ``HasInnerState`` (launch ``--mamba-cache-mode=align``).
"""

from __future__ import annotations

from collections.abc import Iterable
from typing import Any, ClassVar, Literal

import torch
from torch import nn
from vllm.config import VllmConfig
from vllm.model_executor.models.gemma4 import Gemma4DecoderLayer
from vllm.model_executor.models.interfaces import HasInnerState, IsHybrid
from vllm.model_executor.models.muse_glimmer import (
    MuseGlimmerDummyInputsBuilder,
    MuseGlimmerForCausalLM,
    MuseGlimmerMultiModalProcessor,
    MuseGlimmerProcessingInfo,
)
from vllm.model_executor.models.qwen3_5 import Qwen3_5DecoderLayer
from vllm.model_executor.models.qwen3_next import Qwen3NextForCausalLM
from vllm.model_executor.models.utils import AutoWeightsLoader
from vllm.multimodal import MULTIMODAL_REGISTRY

from .config import (
    CANONICAL_GATE,
    GEMMA_BOUNDARY_MUSE_LAYER,
    HOST_HIDDEN,
    QWEN_BOUNDARY_MUSE_LAYER,
)
from .donor_config import (
    GEMMA_LAYER_INDICES,
    QWEN_LAYER_INDICES,
    clone_vllm_config_for_donor,
    load_donor_text_config,
    load_manifest,
)
from .layers import OblivionDonorBranch, OblivionSeam, OblivionStitch
from .weights import load_oblivion_weights

_MUSE_TEXT_LAYERS = 52
_DONOR_SKIP_PREFIXES = ("qwen_layers.", "gemma_layers.", "qwen_seam.", "gemma_seam.")


@MULTIMODAL_REGISTRY.register_processor(
    MuseGlimmerMultiModalProcessor,
    info=MuseGlimmerProcessingInfo,
    dummy_inputs=MuseGlimmerDummyInputsBuilder,
)
class OblivionForCausalLM(MuseGlimmerForCausalLM, IsHybrid, HasInnerState):
    """Fused Muse (multimodal) + Qwen3.5 + Gemma4, one model, both donors always on."""

    is_hybrid: ClassVar[Literal[True]] = True

    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
        # (a) Build the full native multimodal Muse model: vision tower, the
        #     native text model self.model (52 layers), lm_head, mm-token handling.
        super().__init__(vllm_config=vllm_config, prefix=prefix)

        # Spec-decode (EAGLE3) navigates the target as ``get_language_model().model``
        # and reads ``.embed_tokens`` off it. On multimodal MuseGlimmer,
        # ``get_language_model()`` returns the inner ``MuseGlimmerModel`` (which
        # already holds ``embed_tokens``/``layers``), so give it a self-referential
        # ``.model`` alias. object.__setattr__ keeps it OUT of the submodule graph.
        if not hasattr(self.model, "model"):
            object.__setattr__(self.model, "model", self.model)

        # vLLM's EAGLE3 proposer reads ``target_model.config.image_token_index``
        # for multimodal targets; MuseGlimmer exposes it as ``image_token_id``.
        if not hasattr(self.config, "image_token_index"):
            self.config.image_token_index = getattr(self.config, "image_token_id", 200092)

        self.vllm_config = vllm_config
        cache_config = vllm_config.cache_config
        quant_config = vllm_config.quant_config
        max_rms = float(getattr(self.config, "max_relative_rms", 1.0))

        # (b) Donor decoder layers, built under the live fused vllm_config so each
        #     donor Attention / GDN self-registers KV/mamba state. Unique
        #     single-integer prefixes (source index) never collide with
        #     model.layers.{i}.
        qwen_text_config = load_donor_text_config(vllm_config, "qwen")
        gemma_text_config = load_donor_text_config(vllm_config, "gemma")
        self._qwen_donor_config = qwen_text_config
        qwen_vllm_config = clone_vllm_config_for_donor(vllm_config, qwen_text_config)
        qwen_manifest = load_manifest(vllm_config, "qwen")
        qwen_layer_types = qwen_manifest["layer_types"]  # per local index

        self.qwen_layers = nn.ModuleDict()
        for local, src in enumerate(QWEN_LAYER_INDICES):
            self.qwen_layers[str(src)] = Qwen3_5DecoderLayer(
                qwen_vllm_config,
                layer_type=qwen_layer_types[local],
                prefix=f"qwen_layers.{src}",
            )

        self.gemma_layers = nn.ModuleDict()
        for src in GEMMA_LAYER_INDICES:
            # Gemma4DecoderLayer derives sliding/full geometry + per-layer RoPE from
            # config.layer_types[extract_layer_index(prefix)] == src.
            self.gemma_layers[str(src)] = Gemma4DecoderLayer(
                gemma_text_config,
                cache_config,
                quant_config,
                prefix=f"gemma_layers.{src}",
            )

        # (c) Seams: own the two stitches, reference the donor layers by order.
        self.qwen_seam = self._build_seam(
            "qwen", HOST_HIDDEN, qwen_text_config.hidden_size,
            list(self.qwen_layers.values()), max_rms,
        )
        self.gemma_seam = self._build_seam(
            "gemma", HOST_HIDDEN, gemma_text_config.hidden_size,
            list(self.gemma_layers.values()), max_rms,
        )

    # -- seam construction ---------------------------------------------------
    def _build_seam(
        self,
        donor: str,
        host_hidden: int,
        donor_hidden: int,
        donor_layers: list[nn.Module],
        max_rms: float,
    ) -> OblivionSeam:
        seam_prefix = f"{donor}_seam"
        input_stitch = OblivionStitch(
            host_hidden, donor_hidden, prefix=f"{seam_prefix}.input_stitch"
        )
        output_stitch = OblivionStitch(
            donor_hidden, host_hidden, prefix=f"{seam_prefix}.output_stitch"
        )
        branch = OblivionDonorBranch(
            donor,
            input_stitch,
            nn.ModuleList(donor_layers),
            output_stitch,
            (),
            max_relative_rms=max_rms,
        )
        return OblivionSeam(branch, donor, prefix=seam_prefix)

    # -- forward -------------------------------------------------------------
    def forward(
        self,
        input_ids: torch.Tensor | None,
        positions: torch.Tensor,
        intermediate_tensors: Any = None,
        inputs_embeds: torch.Tensor | None = None,
        **kwargs: Any,
    ) -> torch.Tensor | Any:
        m = self.model
        # Merged text+image embeddings arrive as inputs_embeds from the runner.
        if inputs_embeds is not None:
            hidden = inputs_embeds
        else:
            hidden = m.embed_input_ids(input_ids)
        residual: torch.Tensor | None = None
        layers = m.layers

        # Both donors always on: a constant per-token gate at the canonical value.
        gate = hidden.new_full((hidden.shape[0],), CANONICAL_GATE)

        # EAGLE3 auxiliary hidden-state capture: when spec decoding is active vLLM
        # sets aux_hidden_state_layers on the inner model and expects forward to
        # return (hidden, aux). Capture the true post-layer hidden at each requested
        # muse-layer index; layers after a seam expose donor-influenced states.
        aux_layers = getattr(m, "aux_hidden_state_layers", None) or ()
        aux: list[torch.Tensor] = []

        # muse 0..35
        for i in range(QWEN_BOUNDARY_MUSE_LAYER + 1):
            hidden, residual = layers[i](positions, hidden, residual)
            if aux_layers and (i + 1) in aux_layers:
                aux.append(hidden)
        hidden = self.qwen_seam(hidden, positions, gate)
        residual = None
        # muse 36..47
        for i in range(QWEN_BOUNDARY_MUSE_LAYER + 1, GEMMA_BOUNDARY_MUSE_LAYER + 1):
            hidden, residual = layers[i](positions, hidden, residual)
            if aux_layers and (i + 1) in aux_layers:
                aux.append(hidden)
        hidden = self.gemma_seam(hidden, positions, gate)
        residual = None
        # muse 48..51
        for i in range(GEMMA_BOUNDARY_MUSE_LAYER + 1, _MUSE_TEXT_LAYERS):
            hidden, residual = layers[i](positions, hidden, residual)
            if aux_layers and (i + 1) in aux_layers:
                aux.append(hidden)

        normed = m.norm(hidden)
        if aux:
            return normed, aux
        return normed

    # -- weight loading ------------------------------------------------------
    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
        # Muse (vision tower + language + lm_head) loads via the inherited
        # AutoWeightsLoader path with muse's own mapper; the donor submodules load
        # from the co-located donor artifacts in the self-contained checkpoint dir.
        loaded = AutoWeightsLoader(self, skip_prefixes=list(_DONOR_SKIP_PREFIXES))
        loaded_names = loaded.load_weights(weights, mapper=self.hf_to_vllm_mapper)
        loaded_names |= load_oblivion_weights(self)
        return loaded_names

    # -- IsHybrid / HasInnerState (mamba state from the Qwen GDN donor) -------
    @classmethod
    def get_mamba_state_shape_from_config(
        cls, vllm_config: VllmConfig
    ) -> tuple[tuple[int, ...], ...]:
        donor = clone_vllm_config_for_donor(
            vllm_config, load_donor_text_config(vllm_config, "qwen")
        )
        return Qwen3NextForCausalLM.get_mamba_state_shape_from_config(donor)

    @classmethod
    def get_mamba_state_dtype_from_config(
        cls, vllm_config: VllmConfig
    ) -> tuple[torch.dtype, ...]:
        donor = clone_vllm_config_for_donor(
            vllm_config, load_donor_text_config(vllm_config, "qwen")
        )
        return Qwen3NextForCausalLM.get_mamba_state_dtype_from_config(donor)

    @classmethod
    def get_mamba_state_copy_func(cls) -> Any:
        return Qwen3NextForCausalLM.get_mamba_state_copy_func()

	# SPDX-License-Identifier: Apache-2.0
	# Copyright 2026 The Oblivion Authors.
	"""The fused Oblivion model for vLLM v1 -- multimodal, single route.

	Oblivion subclasses the full native multimodal ``MuseGlimmerForCausalLM`` and
	splices two donor macroblocks into the text decoder: a Qwen3.5 block after muse
	layer 35 and a Gemma4 block after layer 47. Both donors are ALWAYS applied at the
	canonical gate (``CANONICAL_GATE``) -- there is no per-request routing or gating.
	Everything multimodal (vision tower, image/video embedding merge, the
	``@MULTIMODAL_REGISTRY.register_processor``) is inherited unchanged, so image
	prompts run through Glimmer's vision path.

	* ``super().__init__`` builds the vision tower, the native ``self.model``
	(``MuseGlimmerModel``, 52 text layers at ``model.layers.0..51``), ``lm_head``,
	and multimodal token handling.
	* Donor decoder layers are first-class vLLM modules built under the live fused
	``vllm_config`` (so each Attention / GatedDeltaNet self-registers KV / mamba
	state), held in ``self.qwen_layers`` / ``self.gemma_layers`` (ModuleDicts keyed
	by source index, prefixes ``qwen_layers.{24..39}`` / ``gemma_layers.{48..59}``).
	* ``forward`` runs the text loop with two seams: muse 0..35 -> qwen seam -> muse
	36..47 -> gemma seam -> muse 48..51 -> norm. It also captures auxiliary
	hidden states when spec decoding is active (for the EAGLE3 drafter).

	The Qwen donor is recurrent (Gated DeltaNet), so the fused arch is hybrid: it
	implements ``IsHybrid`` + ``HasInnerState`` (launch ``--mamba-cache-mode=align``).
	"""

	from __future__ import annotations

	from collections.abc import Iterable
	from typing import Any, ClassVar, Literal

	import torch
	from torch import nn
	from vllm.config import VllmConfig
	from vllm.model_executor.models.gemma4 import Gemma4DecoderLayer
	from vllm.model_executor.models.interfaces import HasInnerState, IsHybrid
	from vllm.model_executor.models.muse_glimmer import (
	MuseGlimmerDummyInputsBuilder,
	MuseGlimmerForCausalLM,
	MuseGlimmerMultiModalProcessor,
	MuseGlimmerProcessingInfo,
	)
	from vllm.model_executor.models.qwen3_5 import Qwen3_5DecoderLayer
	from vllm.model_executor.models.qwen3_next import Qwen3NextForCausalLM
	from vllm.model_executor.models.utils import AutoWeightsLoader
	from vllm.multimodal import MULTIMODAL_REGISTRY

	from .config import (
	CANONICAL_GATE,
	GEMMA_BOUNDARY_MUSE_LAYER,
	HOST_HIDDEN,
	QWEN_BOUNDARY_MUSE_LAYER,
	)
	from .donor_config import (
	GEMMA_LAYER_INDICES,
	QWEN_LAYER_INDICES,
	clone_vllm_config_for_donor,
	load_donor_text_config,
	load_manifest,
	)
	from .layers import OblivionDonorBranch, OblivionSeam, OblivionStitch
	from .weights import load_oblivion_weights

	_MUSE_TEXT_LAYERS = 52
	_DONOR_SKIP_PREFIXES = ("qwen_layers.", "gemma_layers.", "qwen_seam.", "gemma_seam.")


	@MULTIMODAL_REGISTRY.register_processor(
	MuseGlimmerMultiModalProcessor,
	info=MuseGlimmerProcessingInfo,
	dummy_inputs=MuseGlimmerDummyInputsBuilder,
	)
	class OblivionForCausalLM(MuseGlimmerForCausalLM, IsHybrid, HasInnerState):
	"""Fused Muse (multimodal) + Qwen3.5 + Gemma4, one model, both donors always on."""

	is_hybrid: ClassVar[Literal[True]] = True

	def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
	# (a) Build the full native multimodal Muse model: vision tower, the
	# native text model self.model (52 layers), lm_head, mm-token handling.
	super().__init__(vllm_config=vllm_config, prefix=prefix)

	# Spec-decode (EAGLE3) navigates the target as ``get_language_model().model``
	# and reads ``.embed_tokens`` off it. On multimodal MuseGlimmer,
	# ``get_language_model()`` returns the inner ``MuseGlimmerModel`` (which
	# already holds ``embed_tokens``/``layers``), so give it a self-referential
	# ``.model`` alias. object.__setattr__ keeps it OUT of the submodule graph.
	if not hasattr(self.model, "model"):
	object.__setattr__(self.model, "model", self.model)

	# vLLM's EAGLE3 proposer reads ``target_model.config.image_token_index``
	# for multimodal targets; MuseGlimmer exposes it as ``image_token_id``.
	if not hasattr(self.config, "image_token_index"):
	self.config.image_token_index = getattr(self.config, "image_token_id", 200092)

	self.vllm_config = vllm_config
	cache_config = vllm_config.cache_config
	quant_config = vllm_config.quant_config
	max_rms = float(getattr(self.config, "max_relative_rms", 1.0))

	# (b) Donor decoder layers, built under the live fused vllm_config so each
	# donor Attention / GDN self-registers KV/mamba state. Unique
	# single-integer prefixes (source index) never collide with
	# model.layers.{i}.
	qwen_text_config = load_donor_text_config(vllm_config, "qwen")
	gemma_text_config = load_donor_text_config(vllm_config, "gemma")
	self._qwen_donor_config = qwen_text_config
	qwen_vllm_config = clone_vllm_config_for_donor(vllm_config, qwen_text_config)
	qwen_manifest = load_manifest(vllm_config, "qwen")
	qwen_layer_types = qwen_manifest["layer_types"] # per local index

	self.qwen_layers = nn.ModuleDict()
	for local, src in enumerate(QWEN_LAYER_INDICES):
	self.qwen_layers[str(src)] = Qwen3_5DecoderLayer(
	qwen_vllm_config,
	layer_type=qwen_layer_types[local],
	prefix=f"qwen_layers.{src}",
	)

	self.gemma_layers = nn.ModuleDict()
	for src in GEMMA_LAYER_INDICES:
	# Gemma4DecoderLayer derives sliding/full geometry + per-layer RoPE from
	# config.layer_types[extract_layer_index(prefix)] == src.
	self.gemma_layers[str(src)] = Gemma4DecoderLayer(
	gemma_text_config,
	cache_config,
	quant_config,
	prefix=f"gemma_layers.{src}",
	)

	# (c) Seams: own the two stitches, reference the donor layers by order.
	self.qwen_seam = self._build_seam(
	"qwen", HOST_HIDDEN, qwen_text_config.hidden_size,
	list(self.qwen_layers.values()), max_rms,
	)
	self.gemma_seam = self._build_seam(
	"gemma", HOST_HIDDEN, gemma_text_config.hidden_size,
	list(self.gemma_layers.values()), max_rms,
	)

	# -- seam construction ---------------------------------------------------
	def _build_seam(
	self,
	donor: str,
	host_hidden: int,
	donor_hidden: int,
	donor_layers: list[nn.Module],
	max_rms: float,
	) -> OblivionSeam:
	seam_prefix = f"{donor}_seam"
	input_stitch = OblivionStitch(
	host_hidden, donor_hidden, prefix=f"{seam_prefix}.input_stitch"
	)
	output_stitch = OblivionStitch(
	donor_hidden, host_hidden, prefix=f"{seam_prefix}.output_stitch"
	)
	branch = OblivionDonorBranch(
	donor,
	input_stitch,
	nn.ModuleList(donor_layers),
	output_stitch,
	(),
	max_relative_rms=max_rms,
	)
	return OblivionSeam(branch, donor, prefix=seam_prefix)

	# -- forward -------------------------------------------------------------
	def forward(
	self,
	input_ids: torch.Tensor \| None,
	positions: torch.Tensor,
	intermediate_tensors: Any = None,
	inputs_embeds: torch.Tensor \| None = None,
	**kwargs: Any,
	) -> torch.Tensor \| Any:
	m = self.model
	# Merged text+image embeddings arrive as inputs_embeds from the runner.
	if inputs_embeds is not None:
	hidden = inputs_embeds
	else:
	hidden = m.embed_input_ids(input_ids)
	residual: torch.Tensor \| None = None
	layers = m.layers

	# Both donors always on: a constant per-token gate at the canonical value.
	gate = hidden.new_full((hidden.shape[0],), CANONICAL_GATE)

	# EAGLE3 auxiliary hidden-state capture: when spec decoding is active vLLM
	# sets aux_hidden_state_layers on the inner model and expects forward to
	# return (hidden, aux). Capture the true post-layer hidden at each requested
	# muse-layer index; layers after a seam expose donor-influenced states.
	aux_layers = getattr(m, "aux_hidden_state_layers", None) or ()
	aux: list[torch.Tensor] = []

	# muse 0..35
	for i in range(QWEN_BOUNDARY_MUSE_LAYER + 1):
	hidden, residual = layers[i](positions, hidden, residual)
	if aux_layers and (i + 1) in aux_layers:
	aux.append(hidden)
	hidden = self.qwen_seam(hidden, positions, gate)
	residual = None
	# muse 36..47
	for i in range(QWEN_BOUNDARY_MUSE_LAYER + 1, GEMMA_BOUNDARY_MUSE_LAYER + 1):
	hidden, residual = layers[i](positions, hidden, residual)
	if aux_layers and (i + 1) in aux_layers:
	aux.append(hidden)
	hidden = self.gemma_seam(hidden, positions, gate)
	residual = None
	# muse 48..51
	for i in range(GEMMA_BOUNDARY_MUSE_LAYER + 1, _MUSE_TEXT_LAYERS):
	hidden, residual = layers[i](positions, hidden, residual)
	if aux_layers and (i + 1) in aux_layers:
	aux.append(hidden)

	normed = m.norm(hidden)
	if aux:
	return normed, aux
	return normed

	# -- weight loading ------------------------------------------------------
	def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
	# Muse (vision tower + language + lm_head) loads via the inherited
	# AutoWeightsLoader path with muse's own mapper; the donor submodules load
	# from the co-located donor artifacts in the self-contained checkpoint dir.
	loaded = AutoWeightsLoader(self, skip_prefixes=list(_DONOR_SKIP_PREFIXES))
	loaded_names = loaded.load_weights(weights, mapper=self.hf_to_vllm_mapper)
	loaded_names \|= load_oblivion_weights(self)
	return loaded_names

	# -- IsHybrid / HasInnerState (mamba state from the Qwen GDN donor) -------
	@classmethod
	def get_mamba_state_shape_from_config(
	cls, vllm_config: VllmConfig
	) -> tuple[tuple[int, ...], ...]:
	donor = clone_vllm_config_for_donor(
	vllm_config, load_donor_text_config(vllm_config, "qwen")
	)
	return Qwen3NextForCausalLM.get_mamba_state_shape_from_config(donor)

	@classmethod
	def get_mamba_state_dtype_from_config(
	cls, vllm_config: VllmConfig
	) -> tuple[torch.dtype, ...]:
	donor = clone_vllm_config_for_donor(
	vllm_config, load_donor_text_config(vllm_config, "qwen")
	)
	return Qwen3NextForCausalLM.get_mamba_state_dtype_from_config(donor)

	@classmethod
	def get_mamba_state_copy_func(cls) -> Any:
	return Qwen3NextForCausalLM.get_mamba_state_copy_func()