LemonDecider / modeling_lemondecider.py

Update modeling_lemondecider.py

3ceb5a3 verified 14 days ago

13.1 kB

# Author：忍把浮名，换了浅斟低唱
"""LemonDecider full-weight decision model. No autoregressive generation head.
The checkpoint's backbone is copied byte-for-byte from the base LM, while readout
rows are selected from its original LM head. T is used in postprocessing, not CE loss.
"""
from dataclasses import dataclass
from contextlib import nullcontext
import torch
import torch.nn.functional as F
from transformers.modeling_utils import PreTrainedModel
from transformers.utils import ModelOutput
from transformers.models.llama.modeling_llama import LlamaModel
from transformers.models.qwen3_5.modeling_qwen3_5 import Qwen3_5Model, Qwen3_5Attention
from .configuration_lemondecider import LemonDeciderConfig
# HF local remote-code caching copies direct dependencies before recursing.
Spark2_5Config = Spark2_5Model = None  # Spark backbone is not part of this Qwen3.5 export
from .option_loss import option_cross_entropy
from .contract import LemonDeciderReadout as D3Readout, enable_noncausal_full_attention, sdpa_backends, linearize_patch_embed


@dataclass
class LemonDeciderOutput(ModelOutput):
    loss: torch.Tensor | None = None
    logits: torch.Tensor | None = None


def final_norm_module(model):
    """Last language RMSNorm feeding the decision readout."""
    backbone = model.backbone if hasattr(model, "backbone") else model
    language = getattr(backbone, "language_model", None)
    if language is not None and hasattr(language, "norm"):
        return language.norm
    if hasattr(backbone, "norm"):
        return backbone.norm
    raise AttributeError("Backbone has no final RMSNorm")


class LemonDeciderForDecisionClassification(PreTrainedModel):
    config_class = LemonDeciderConfig
    base_model_prefix = "backbone"
    supports_gradient_checkpointing = True
    _no_split_modules = ["Qwen3_5DecoderLayer", "Qwen3_5VisionBlock", "LlamaDecoderLayer", "Spark2_5DecoderLayer"]
    _supports_flash_attn = True
    _supports_flash_attn_2 = True
    _supports_sdpa = True
    _supports_cache_class = False

    def __init__(self, config):
        super().__init__(config)
        hidden = config.text_config.hidden_size
        attn = getattr(config, "_attn_implementation", None) or getattr(config, "_attn_implementation_internal", None)
        attention_mode = getattr(config, "attention_mode", "causal")
        if attention_mode not in {"causal", "noncausal_full_attention"}:
            raise ValueError(f"Unknown attention mode: {attention_mode}")
        if getattr(config, "pooling", "last") != "last":
            raise ValueError(f"Unknown pooling mode: {getattr(config, 'pooling', 'last')}")
        if getattr(config, 'backbone_type', 'qwen3_5') == 'spark2_5':
            inner = config.spark_config
            inner.attention_mode = attention_mode
            inner._attn_implementation = attn or 'sdpa'
            self.backbone = Spark2_5Model(inner)
        elif getattr(config, "backbone_type", "qwen3_5") == "llama":
            if attention_mode != "causal":
                raise ValueError("noncausal_full_attention is only implemented for the Qwen3.5 hybrid backbone")
            inner = config.llama_config
            if attn:
                inner._attn_implementation = attn
            self.backbone = LlamaModel(inner)
        else:
            if attention_mode == "noncausal_full_attention":
                config.text_config.is_causal = False
            self.backbone = Qwen3_5Model(config)
            self.backbone.language_model.config.is_causal = attention_mode == 'causal'
            # SDPA may omit the all-valid mask. Qwen's module flag otherwise
            # re-enables causality even when its config requested bidirectional.
            # Restrict to full attention; never modify the linear recurrence.
            for module in self.backbone.language_model.modules():
                if isinstance(module, Qwen3_5Attention):
                    module.is_causal = attention_mode == 'causal'
        d3 = getattr(config, 'decision_protocol', 'lemondecider-v1') == 'd3-v1'
        if d3:
            if config.backbone_type == 'qwen3_5':
                if linearize_patch_embed(self.backbone) != 1:
                    raise ValueError('D3 requires exactly one matrix-product vision patch stem')
                enable_noncausal_full_attention(self.backbone.language_model)
            self._keep_in_fp32_modules_strict = ['readout']
        self.readout = (D3Readout if d3 else torch.nn.Linear)(hidden, config.max_options, bias=False)
        self.post_init()

    def get_input_embeddings(self):
        return self.backbone.get_input_embeddings()

    def set_input_embeddings(self, value):
        self.backbone.set_input_embeddings(value)

    def gradient_checkpointing_enable(self, **kwargs):
        super().gradient_checkpointing_enable(**kwargs)
        if hasattr(self.backbone, "gradient_checkpointing_enable"):
            self.backbone.gradient_checkpointing_enable(**kwargs)

    def forward(self, input_ids=None, attention_mask=None, option_counts=None,
                labels=None, return_dict=True, readout_positions=None, **kwargs):
        """readout_positions: optional [Q, 2] (batch row, token index) joint-marker readouts
        (joint_readout.py); option_counts then has one entry per readout. None keeps the
        legacy last-token readout with one decision per row."""
        if input_ids is None or input_ids.ndim != 2 or input_ids.shape[1] == 0:
            raise ValueError("input_ids must be a nonempty [batch, length] tensor")
        if input_ids.shape[1] > self.config.max_input_length:
            raise ValueError("Input exceeds max_input_length; truncation is not allowed")
        if attention_mask is None:
            attention_mask = torch.ones_like(input_ids)
        if attention_mask.shape != input_ids.shape or not bool((attention_mask[:, -1] == 1).all()):
            raise ValueError("Left padding required: last prompt position must be valid")
        if option_counts is None:
            raise ValueError("option_counts is required for each input")
        counts = torch.as_tensor(option_counts, device=input_ids.device)
        decisions = input_ids.shape[0] if readout_positions is None else (
            len(readout_positions) if hasattr(readout_positions, '__len__') else readout_positions.shape[0])
        if counts.dtype == torch.bool or counts.is_floating_point() or counts.shape != (decisions,):
            raise ValueError("option_counts must contain one integer per row" if readout_positions is None
                             else "option_counts must contain one integer per joint readout")
        if readout_positions is not None:
            readout_positions = torch.as_tensor(readout_positions, device=input_ids.device)
            if (readout_positions.dtype == torch.bool or readout_positions.is_floating_point()
                    or readout_positions.ndim != 2 or readout_positions.shape[1] != 2 or not len(readout_positions)):
                raise ValueError("readout_positions must be a nonempty integer [Q, 2] tensor")
            rows_, cols_ = readout_positions[:, 0], readout_positions[:, 1]
            if (not bool(((rows_ >= 0) & (rows_ < input_ids.shape[0]) & (cols_ >= 0) & (cols_ < input_ids.shape[1])).all())
                    or not bool(attention_mask[rows_, cols_].bool().all())):
                raise ValueError("readout_positions must index valid (non-padding) tokens")
            if bool(torch.bincount(rows_, minlength=input_ids.shape[0]).eq(0).any()):
                raise ValueError("Every joint sequence needs at least one readout")
        if not bool(((counts >= 1) & (counts <= self.config.max_options)).all()):
            raise ValueError(f"Each question must contain 1..{self.config.max_options} options")
        if kwargs.pop("use_cache", False) or kwargs.get("past_key_values") is not None:
            raise ValueError("Decision forward does not accept KV-cache decoding")
        if getattr(self.config, "attention_mode", "causal") == "noncausal_full_attention":
            implementation = getattr(self.config, "_attn_implementation", None) or getattr(
                self.config, "_attn_implementation_internal", None
            )
            if implementation != "sdpa":
                raise ValueError("noncausal_full_attention requires SDPA")
            if self.config.backbone_type == 'spark2_5':
                if self.backbone.config.attention_mode != 'noncausal_full_attention':
                    raise ValueError('Spark full-attention mode differs from the decision config')
            else:
                language = getattr(self.backbone, "language_model", None)
                if language is None or getattr(language.config, "is_causal", True):
                    raise ValueError("noncausal_full_attention requires language_model.config.is_causal=False")
        vision_zero = None
        if getattr(self.config, 'vision_training_mode', 'frozen') == 'all-trainable' and kwargs.get('pixel_values') is None:
            # Execute the SAME vision FSDP units on image-free and dummy ranks.
            # A minimal synthetic patch grid has zero loss mass, not a dataset row.
            visual = self.backbone.visual
            vc = self.config.vision_config
            merge = vc.spatial_merge_size
            pixels = torch.zeros(merge * merge, vc.in_channels * vc.temporal_patch_size * vc.patch_size ** 2,
                                 device=input_ids.device, dtype=visual.dtype)
            grid = torch.tensor([[1, merge, merge]], device=input_ids.device, dtype=torch.long)
            vision_zero = visual(pixels, grid_thw=grid, return_dict=True).pooler_output.sum() * 0.0
        backbone_ids = input_ids
        if vision_zero is not None:
            # Link BEFORE the language stack, like real image embeddings, so
            # backward FSDP collective order also agrees across ranks.
            embeds = self.backbone.get_input_embeddings()(input_ids)
            kwargs['inputs_embeds'] = embeds + vision_zero.to(embeds.dtype)
            if kwargs.get('position_ids') is None:
                # Same text-only arange as QwenTextModel; avoid stale M-RoPE
                # deltas from a preceding multimodal batch when passing embeds.
                kwargs['position_ids'] = torch.arange(input_ids.shape[1], device=input_ids.device)[None, None].expand(3, input_ids.shape[0], -1)
            backbone_ids = None
        backend_context = (sdpa_backends(input_ids.device)
                           if getattr(self.config, 'decision_protocol', 'lemondecider-v1') == 'd3-v1' else nullcontext())
        with backend_context:
            output = self.backbone(input_ids=backbone_ids, attention_mask=attention_mask,
                                   use_cache=False, return_dict=True, **kwargs)
        if readout_positions is None:
            hidden = output.last_hidden_state[:, -1, :]
        else:
            hidden = output.last_hidden_state[readout_positions[:, 0], readout_positions[:, 1]]
        # A separate FSDP precision island keeps D3's head weights and GEMM FP32.
        # Passing FP32 also prevents any mixed-precision wrapper from downcasting logits.
        logits = self.readout(hidden.float() if getattr(self.config, 'decision_protocol', 'lemondecider-v1') == 'd3-v1' else hidden).float()
        invalid = torch.arange(self.config.max_options, device=logits.device)[None] >= counts[:, None]
        logits = logits.masked_fill(invalid, -1e9)
        loss = None
        if labels is not None:
            labels = labels.to(logits.device)
            if labels.shape == (len(counts),):
                if labels.is_floating_point() or labels.dtype == torch.bool or not bool(((labels >= 0) & (labels < counts)).all()):
                    raise ValueError("Hard targets must be valid option indices")
                target = F.one_hot(labels.long(), self.config.max_options).float()
                loss = option_cross_entropy(logits, target, counts)
            elif labels.shape == logits.shape:
                target = labels.float()
                if not bool(torch.isfinite(target).all()) or bool((target < 0).any()):
                    raise ValueError("Invalid soft targets")
                if not torch.allclose(target.sum(-1), torch.ones(len(counts), device=target.device), atol=1e-6, rtol=0):
                    raise ValueError("Soft targets must sum to one")
                if bool((target.masked_select(invalid) != 0).any()):
                    raise ValueError("Soft targets assign mass to invalid options")
                loss = option_cross_entropy(logits, target, counts)
            else:
                raise ValueError(f"labels must be [batch] or [batch,{self.config.max_options}]")
        if not return_dict:
            return (loss, logits) if loss is not None else (logits,)
        return LemonDeciderOutput(loss=loss, logits=logits)


LemonDeciderConfig.register_for_auto_class()
LemonDeciderForDecisionClassification.register_for_auto_class("AutoModel")

	# Author：忍把浮名，换了浅斟低唱
	"""LemonDecider full-weight decision model. No autoregressive generation head.


	The checkpoint's backbone is copied byte-for-byte from the base LM, while readout
	rows are selected from its original LM head. T is used in postprocessing, not CE loss.
	"""
	from dataclasses import dataclass
	from contextlib import nullcontext
	import torch
	import torch.nn.functional as F
	from transformers.modeling_utils import PreTrainedModel
	from transformers.utils import ModelOutput
	from transformers.models.llama.modeling_llama import LlamaModel
	from transformers.models.qwen3_5.modeling_qwen3_5 import Qwen3_5Model, Qwen3_5Attention
	from .configuration_lemondecider import LemonDeciderConfig
	# HF local remote-code caching copies direct dependencies before recursing.
	Spark2_5Config = Spark2_5Model = None # Spark backbone is not part of this Qwen3.5 export
	from .option_loss import option_cross_entropy
	from .contract import LemonDeciderReadout as D3Readout, enable_noncausal_full_attention, sdpa_backends, linearize_patch_embed


	@dataclass
	class LemonDeciderOutput(ModelOutput):
	loss: torch.Tensor \| None = None
	logits: torch.Tensor \| None = None


	def final_norm_module(model):
	"""Last language RMSNorm feeding the decision readout."""
	backbone = model.backbone if hasattr(model, "backbone") else model
	language = getattr(backbone, "language_model", None)
	if language is not None and hasattr(language, "norm"):
	return language.norm
	if hasattr(backbone, "norm"):
	return backbone.norm
	raise AttributeError("Backbone has no final RMSNorm")


	class LemonDeciderForDecisionClassification(PreTrainedModel):
	config_class = LemonDeciderConfig
	base_model_prefix = "backbone"
	supports_gradient_checkpointing = True
	_no_split_modules = ["Qwen3_5DecoderLayer", "Qwen3_5VisionBlock", "LlamaDecoderLayer", "Spark2_5DecoderLayer"]
	_supports_flash_attn = True
	_supports_flash_attn_2 = True
	_supports_sdpa = True
	_supports_cache_class = False

	def __init__(self, config):
	super().__init__(config)
	hidden = config.text_config.hidden_size
	attn = getattr(config, "_attn_implementation", None) or getattr(config, "_attn_implementation_internal", None)
	attention_mode = getattr(config, "attention_mode", "causal")
	if attention_mode not in {"causal", "noncausal_full_attention"}:
	raise ValueError(f"Unknown attention mode: {attention_mode}")
	if getattr(config, "pooling", "last") != "last":
	raise ValueError(f"Unknown pooling mode: {getattr(config, 'pooling', 'last')}")
	if getattr(config, 'backbone_type', 'qwen3_5') == 'spark2_5':
	inner = config.spark_config
	inner.attention_mode = attention_mode
	inner._attn_implementation = attn or 'sdpa'
	self.backbone = Spark2_5Model(inner)
	elif getattr(config, "backbone_type", "qwen3_5") == "llama":
	if attention_mode != "causal":
	raise ValueError("noncausal_full_attention is only implemented for the Qwen3.5 hybrid backbone")
	inner = config.llama_config
	if attn:
	inner._attn_implementation = attn
	self.backbone = LlamaModel(inner)
	else:
	if attention_mode == "noncausal_full_attention":
	config.text_config.is_causal = False
	self.backbone = Qwen3_5Model(config)
	self.backbone.language_model.config.is_causal = attention_mode == 'causal'
	# SDPA may omit the all-valid mask. Qwen's module flag otherwise
	# re-enables causality even when its config requested bidirectional.
	# Restrict to full attention; never modify the linear recurrence.
	for module in self.backbone.language_model.modules():
	if isinstance(module, Qwen3_5Attention):
	module.is_causal = attention_mode == 'causal'
	d3 = getattr(config, 'decision_protocol', 'lemondecider-v1') == 'd3-v1'
	if d3:
	if config.backbone_type == 'qwen3_5':
	if linearize_patch_embed(self.backbone) != 1:
	raise ValueError('D3 requires exactly one matrix-product vision patch stem')
	enable_noncausal_full_attention(self.backbone.language_model)
	self._keep_in_fp32_modules_strict = ['readout']
	self.readout = (D3Readout if d3 else torch.nn.Linear)(hidden, config.max_options, bias=False)
	self.post_init()

	def get_input_embeddings(self):
	return self.backbone.get_input_embeddings()

	def set_input_embeddings(self, value):
	self.backbone.set_input_embeddings(value)

	def gradient_checkpointing_enable(self, **kwargs):
	super().gradient_checkpointing_enable(**kwargs)
	if hasattr(self.backbone, "gradient_checkpointing_enable"):
	self.backbone.gradient_checkpointing_enable(**kwargs)

	def forward(self, input_ids=None, attention_mask=None, option_counts=None,
	labels=None, return_dict=True, readout_positions=None, **kwargs):
	"""readout_positions: optional [Q, 2] (batch row, token index) joint-marker readouts
	(joint_readout.py); option_counts then has one entry per readout. None keeps the
	legacy last-token readout with one decision per row."""
	if input_ids is None or input_ids.ndim != 2 or input_ids.shape[1] == 0:
	raise ValueError("input_ids must be a nonempty [batch, length] tensor")
	if input_ids.shape[1] > self.config.max_input_length:
	raise ValueError("Input exceeds max_input_length; truncation is not allowed")
	if attention_mask is None:
	attention_mask = torch.ones_like(input_ids)
	if attention_mask.shape != input_ids.shape or not bool((attention_mask[:, -1] == 1).all()):
	raise ValueError("Left padding required: last prompt position must be valid")
	if option_counts is None:
	raise ValueError("option_counts is required for each input")
	counts = torch.as_tensor(option_counts, device=input_ids.device)
	decisions = input_ids.shape[0] if readout_positions is None else (
	len(readout_positions) if hasattr(readout_positions, '__len__') else readout_positions.shape[0])
	if counts.dtype == torch.bool or counts.is_floating_point() or counts.shape != (decisions,):
	raise ValueError("option_counts must contain one integer per row" if readout_positions is None
	else "option_counts must contain one integer per joint readout")
	if readout_positions is not None:
	readout_positions = torch.as_tensor(readout_positions, device=input_ids.device)
	if (readout_positions.dtype == torch.bool or readout_positions.is_floating_point()
	or readout_positions.ndim != 2 or readout_positions.shape[1] != 2 or not len(readout_positions)):
	raise ValueError("readout_positions must be a nonempty integer [Q, 2] tensor")
	rows_, cols_ = readout_positions[:, 0], readout_positions[:, 1]
	if (not bool(((rows_ >= 0) & (rows_ < input_ids.shape[0]) & (cols_ >= 0) & (cols_ < input_ids.shape[1])).all())
	or not bool(attention_mask[rows_, cols_].bool().all())):
	raise ValueError("readout_positions must index valid (non-padding) tokens")
	if bool(torch.bincount(rows_, minlength=input_ids.shape[0]).eq(0).any()):
	raise ValueError("Every joint sequence needs at least one readout")
	if not bool(((counts >= 1) & (counts <= self.config.max_options)).all()):
	raise ValueError(f"Each question must contain 1..{self.config.max_options} options")
	if kwargs.pop("use_cache", False) or kwargs.get("past_key_values") is not None:
	raise ValueError("Decision forward does not accept KV-cache decoding")
	if getattr(self.config, "attention_mode", "causal") == "noncausal_full_attention":
	implementation = getattr(self.config, "_attn_implementation", None) or getattr(
	self.config, "_attn_implementation_internal", None
	)
	if implementation != "sdpa":
	raise ValueError("noncausal_full_attention requires SDPA")
	if self.config.backbone_type == 'spark2_5':
	if self.backbone.config.attention_mode != 'noncausal_full_attention':
	raise ValueError('Spark full-attention mode differs from the decision config')
	else:
	language = getattr(self.backbone, "language_model", None)
	if language is None or getattr(language.config, "is_causal", True):
	raise ValueError("noncausal_full_attention requires language_model.config.is_causal=False")
	vision_zero = None
	if getattr(self.config, 'vision_training_mode', 'frozen') == 'all-trainable' and kwargs.get('pixel_values') is None:
	# Execute the SAME vision FSDP units on image-free and dummy ranks.
	# A minimal synthetic patch grid has zero loss mass, not a dataset row.
	visual = self.backbone.visual
	vc = self.config.vision_config
	merge = vc.spatial_merge_size
	pixels = torch.zeros(merge * merge, vc.in_channels * vc.temporal_patch_size * vc.patch_size ** 2,
	device=input_ids.device, dtype=visual.dtype)
	grid = torch.tensor([[1, merge, merge]], device=input_ids.device, dtype=torch.long)
	vision_zero = visual(pixels, grid_thw=grid, return_dict=True).pooler_output.sum() * 0.0
	backbone_ids = input_ids
	if vision_zero is not None:
	# Link BEFORE the language stack, like real image embeddings, so
	# backward FSDP collective order also agrees across ranks.
	embeds = self.backbone.get_input_embeddings()(input_ids)
	kwargs['inputs_embeds'] = embeds + vision_zero.to(embeds.dtype)
	if kwargs.get('position_ids') is None:
	# Same text-only arange as QwenTextModel; avoid stale M-RoPE
	# deltas from a preceding multimodal batch when passing embeds.
	kwargs['position_ids'] = torch.arange(input_ids.shape[1], device=input_ids.device)[None, None].expand(3, input_ids.shape[0], -1)
	backbone_ids = None
	backend_context = (sdpa_backends(input_ids.device)
	if getattr(self.config, 'decision_protocol', 'lemondecider-v1') == 'd3-v1' else nullcontext())
	with backend_context:
	output = self.backbone(input_ids=backbone_ids, attention_mask=attention_mask,
	use_cache=False, return_dict=True, **kwargs)
	if readout_positions is None:
	hidden = output.last_hidden_state[:, -1, :]
	else:
	hidden = output.last_hidden_state[readout_positions[:, 0], readout_positions[:, 1]]
	# A separate FSDP precision island keeps D3's head weights and GEMM FP32.
	# Passing FP32 also prevents any mixed-precision wrapper from downcasting logits.
	logits = self.readout(hidden.float() if getattr(self.config, 'decision_protocol', 'lemondecider-v1') == 'd3-v1' else hidden).float()
	invalid = torch.arange(self.config.max_options, device=logits.device)[None] >= counts[:, None]
	logits = logits.masked_fill(invalid, -1e9)
	loss = None
	if labels is not None:
	labels = labels.to(logits.device)
	if labels.shape == (len(counts),):
	if labels.is_floating_point() or labels.dtype == torch.bool or not bool(((labels >= 0) & (labels < counts)).all()):
	raise ValueError("Hard targets must be valid option indices")
	target = F.one_hot(labels.long(), self.config.max_options).float()
	loss = option_cross_entropy(logits, target, counts)
	elif labels.shape == logits.shape:
	target = labels.float()
	if not bool(torch.isfinite(target).all()) or bool((target < 0).any()):
	raise ValueError("Invalid soft targets")
	if not torch.allclose(target.sum(-1), torch.ones(len(counts), device=target.device), atol=1e-6, rtol=0):
	raise ValueError("Soft targets must sum to one")
	if bool((target.masked_select(invalid) != 0).any()):
	raise ValueError("Soft targets assign mass to invalid options")
	loss = option_cross_entropy(logits, target, counts)
	else:
	raise ValueError(f"labels must be [batch] or [batch,{self.config.max_options}]")
	if not return_dict:
	return (loss, logits) if loss is not None else (logits,)
	return LemonDeciderOutput(loss=loss, logits=logits)


	LemonDeciderConfig.register_for_auto_class()
	LemonDeciderForDecisionClassification.register_for_auto_class("AutoModel")