初始化项目，由ModelHub XC社区提供模型
Model: fin-ai-lab/aux-2015 Source: Original Platform
This commit is contained in:
ModelHub XC
190
modeling_ministral_dual_rope.py
Normal file
190
modeling_ministral_dual_rope.py
Normal file
@@ -0,0 +1,190 @@
|
||||
"""MinistralDualRope — Ministral with a per-layer sliding/full attention pattern AND a second,
|
||||
unscaled rotary for the sliding layers. Loadable via ``trust_remote_code``.
|
||||
|
||||
This is the MINIMAL delta over the *stock* transformers Ministral model:
|
||||
|
||||
* Ministral already provides the Llama-style math (RMSNorm without unit offset,
|
||||
SiLU-gated MLP, GQA, no biases, no QK-norm, no sandwich norms, no softcapping) AND
|
||||
the per-layer sliding/full attention pattern routed by ``config.layer_types`` — the
|
||||
attention/mask routing is bit-identical to how this model was trained.
|
||||
* The ONE thing Ministral lacks is a *second* RoPE. Our full-attention layers use the
|
||||
(llama3-scaled) global rotary; the sliding layers use an UNSCALED rotary at
|
||||
``rope_local_base_freq`` (the Gemma-3 dual-RoPE design). We add ``rotary_emb_local``
|
||||
and route ``position_embeddings`` per ``layer_types``. Everything else is stock
|
||||
Ministral, inherited unchanged.
|
||||
|
||||
The building blocks are pulled from the *installed* ``modeling_ministral`` so this stays
|
||||
aligned with whatever transformers version loads the checkpoint. Verified to match the
|
||||
faithful FlexLlama reference to the analytic-vs-training-buffer RoPE floor (FP32 top-1
|
||||
agreement ~100%, mean |Δ logit| ~2e-3).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
import inspect
|
||||
from functools import partial
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
|
||||
from transformers.models.ministral import modeling_ministral as _M
|
||||
|
||||
from .configuration_ministral_dual_rope import MinistralDualRopeConfig
|
||||
|
||||
MinistralModel = _M.MinistralModel
|
||||
MinistralForCausalLM = _M.MinistralForCausalLM
|
||||
MinistralRotaryEmbedding = _M.MinistralRotaryEmbedding
|
||||
create_causal_mask = _M.create_causal_mask
|
||||
create_sliding_window_causal_mask = _M.create_sliding_window_causal_mask
|
||||
DynamicCache = _M.DynamicCache
|
||||
BaseModelOutputWithPast = _M.BaseModelOutputWithPast
|
||||
# `check_model_inputs` moved between transformers minor/major versions; fall back to a
|
||||
# no-op so the forward still works if a given install lacks it. We deliberately do NOT
|
||||
# apply transformers' `auto_docstring` to the forward: on tf>=5 it scans the signature at
|
||||
# import and logs a scary `[ERROR] ... is part of ...'s signature, but not documented` for
|
||||
# every undocumented kwarg (e.g. `cache_position`). That is pure lint noise — the forward
|
||||
# is fully functional without it — but it panics downstream users, so we omit it.
|
||||
check_model_inputs = getattr(_M, "check_model_inputs", lambda f: f)
|
||||
|
||||
# The mask-builder signature drifts across transformers versions (the embeds kwarg was
|
||||
# renamed ``input_embeds`` -> ``inputs_embeds`` and ``cache_position`` was dropped between
|
||||
# 4.x and 5.x). Introspect once and pass only what each installed builder accepts, so the
|
||||
# same shipped file loads under both. Computed at import time (cheap, version-stable).
|
||||
_CAUSAL_MASK_PARAMS = set(inspect.signature(create_causal_mask).parameters)
|
||||
_SLIDING_MASK_PARAMS = set(inspect.signature(create_sliding_window_causal_mask).parameters)
|
||||
|
||||
|
||||
class MinistralDualRopeModel(MinistralModel):
|
||||
config_class = MinistralDualRopeConfig
|
||||
|
||||
def __init__(self, config: MinistralDualRopeConfig):
|
||||
super().__init__(config)
|
||||
# Second, UNSCALED rotary for the sliding (local) layers. The parent's
|
||||
# ``self.rotary_emb`` is the (possibly llama3-scaled) GLOBAL rotary used by the
|
||||
# full-attention layers; build a sibling at ``rope_local_base_freq`` with NO
|
||||
# scaling. When base == rope_theta and there is no scaling the two are identical.
|
||||
local_cfg = copy.deepcopy(config)
|
||||
base = float(getattr(config, "rope_local_base_freq", None) or config._global_rope_theta())
|
||||
# Configure an UNSCALED rotary. The two transformers major versions expose RoPE
|
||||
# differently and the branches MUST be mutually exclusive:
|
||||
# * transformers >= 5 stores everything in the unified ``rope_parameters`` dict
|
||||
# and exposes ``rope_scaling`` as an ALIAS *onto* it — so setting
|
||||
# ``rope_scaling`` here would clobber ``rope_parameters`` back to ``None`` and
|
||||
# blow up inside ``MinistralRotaryEmbedding``. Set ONLY ``rope_parameters``.
|
||||
# * transformers 4.x has no ``rope_parameters``; the rotary reads ``rope_scaling``
|
||||
# (must be None to disable llama3 scaling) + a top-level ``rope_theta``.
|
||||
if hasattr(local_cfg, "rope_parameters"): # transformers >= 5
|
||||
local_cfg.rope_parameters = {"rope_type": "default", "rope_theta": base}
|
||||
else: # transformers 4.x
|
||||
local_cfg.rope_scaling = None
|
||||
local_cfg.rope_theta = base
|
||||
self.rotary_emb_local = MinistralRotaryEmbedding(config=local_cfg)
|
||||
self.post_init()
|
||||
|
||||
@check_model_inputs
|
||||
def forward(
|
||||
self,
|
||||
input_ids: Optional[torch.LongTensor] = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
position_ids: Optional[torch.LongTensor] = None,
|
||||
past_key_values=None,
|
||||
inputs_embeds: Optional[torch.FloatTensor] = None,
|
||||
use_cache: Optional[bool] = None,
|
||||
cache_position: Optional[torch.LongTensor] = None,
|
||||
**kwargs,
|
||||
) -> BaseModelOutputWithPast:
|
||||
# Verbatim from MinistralModel.forward except the two dual-RoPE lines marked (*).
|
||||
if (input_ids is None) ^ (inputs_embeds is not None):
|
||||
raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
|
||||
|
||||
if inputs_embeds is None:
|
||||
inputs_embeds = self.embed_tokens(input_ids)
|
||||
|
||||
if use_cache and past_key_values is None:
|
||||
past_key_values = DynamicCache(config=self.config)
|
||||
|
||||
if cache_position is None:
|
||||
past_seen = past_key_values.get_seq_length() if past_key_values is not None else 0
|
||||
cache_position = torch.arange(
|
||||
past_seen, past_seen + inputs_embeds.shape[1], device=inputs_embeds.device
|
||||
)
|
||||
|
||||
if position_ids is None:
|
||||
position_ids = cache_position.unsqueeze(0)
|
||||
|
||||
# `generate` may already pass a prepared {type: mask} dict; else build both masks.
|
||||
if not isinstance(causal_mask_mapping := attention_mask, dict):
|
||||
# Superset of every kwarg any transformers version's mask builder wants; the
|
||||
# two embeds spellings and cache_position are filtered per the installed
|
||||
# signature (see _CAUSAL_MASK_PARAMS above) so this works on 4.x and >=5.
|
||||
_mask_src = {
|
||||
"config": self.config,
|
||||
"input_embeds": inputs_embeds, # transformers 4.x
|
||||
"inputs_embeds": inputs_embeds, # transformers >= 5
|
||||
"attention_mask": attention_mask,
|
||||
"cache_position": cache_position, # transformers 4.x only
|
||||
"past_key_values": past_key_values,
|
||||
"position_ids": position_ids,
|
||||
}
|
||||
causal_mask_mapping = {
|
||||
"full_attention": create_causal_mask(
|
||||
**{k: v for k, v in _mask_src.items() if k in _CAUSAL_MASK_PARAMS}),
|
||||
"sliding_attention": create_sliding_window_causal_mask(
|
||||
**{k: v for k, v in _mask_src.items() if k in _SLIDING_MASK_PARAMS}),
|
||||
}
|
||||
|
||||
hidden_states = inputs_embeds
|
||||
# (*) Dual RoPE: full-attention layers get the (scaled) global rotary, sliding
|
||||
# layers get the unscaled local rotary — routed by the same layer type as the
|
||||
# mask above.
|
||||
position_embeddings = {
|
||||
"full_attention": self.rotary_emb(hidden_states, position_ids),
|
||||
"sliding_attention": self.rotary_emb_local(hidden_states, position_ids),
|
||||
}
|
||||
|
||||
for i, decoder_layer in enumerate(self.layers[: self.config.num_hidden_layers]):
|
||||
# (*) Route the mask AND the rope by the layer's attention type. Read it from
|
||||
# ``config.layer_types`` (stable across versions) rather than the layer
|
||||
# attribute, which transformers renamed ``attention_type`` -> ``layer_type``
|
||||
# between 4.x and 5.x.
|
||||
layer_type = self.config.layer_types[i]
|
||||
layer_kwargs = dict(
|
||||
attention_mask=causal_mask_mapping[layer_type],
|
||||
position_ids=position_ids,
|
||||
past_key_values=past_key_values,
|
||||
use_cache=use_cache,
|
||||
cache_position=cache_position,
|
||||
position_embeddings=position_embeddings[layer_type],
|
||||
**kwargs,
|
||||
)
|
||||
if self.gradient_checkpointing and self.training:
|
||||
# Honor activation checkpointing. The stock MinistralModel.forward does
|
||||
# this; our override MUST replicate it or long-context SFT/RL (this is a
|
||||
# 32K model) OOMs — even under LoRA, since activations, not optimizer
|
||||
# state, dominate at long sequence length. Bake the per-layer kwargs into
|
||||
# a partial so we do NOT depend on the decoder layer's positional arg
|
||||
# order, which drifts across transformers versions.
|
||||
hidden_states = self._gradient_checkpointing_func(
|
||||
partial(decoder_layer.__call__, **layer_kwargs), hidden_states
|
||||
)
|
||||
else:
|
||||
hidden_states = decoder_layer(hidden_states, **layer_kwargs)
|
||||
|
||||
hidden_states = self.norm(hidden_states)
|
||||
return BaseModelOutputWithPast(
|
||||
last_hidden_state=hidden_states,
|
||||
past_key_values=past_key_values if use_cache else None,
|
||||
)
|
||||
|
||||
|
||||
class MinistralDualRopeForCausalLM(MinistralForCausalLM):
|
||||
config_class = MinistralDualRopeConfig
|
||||
|
||||
def __init__(self, config: MinistralDualRopeConfig):
|
||||
super().__init__(config)
|
||||
self.model = MinistralDualRopeModel(config)
|
||||
self.post_init()
|
||||
|
||||
|
||||
__all__ = ["MinistralDualRopeForCausalLM", "MinistralDualRopeModel"]
|
||||
Reference in New Issue
Block a user