191 lines
9.7 KiB
Python
191 lines
9.7 KiB
Python
|
"""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"]
|