55 lines
2.5 KiB
Python
55 lines
2.5 KiB
Python
"""Configuration for MinistralDualRope — a Ministral with a per-layer sliding-window / full
|
|
attention pattern and a second, UNSCALED rotary for the sliding layers.
|
|
|
|
The base model is byte-for-byte a Ministral (== Llama math: SiLU-gated MLP, RMSNorm
|
|
without unit offset, GQA, rotary embeddings, no biases, no QK-norm, no sandwich norms,
|
|
no logit softcapping). Ministral already provides the per-layer sliding/full attention
|
|
pattern via ``layer_types`` (mask routing identical to how this model was trained). The
|
|
*only* addition over stock Ministral is dual RoPE:
|
|
|
|
* full-attention layers use the standard ``rope_theta`` (+ ``rope_scaling``, e.g. the
|
|
llama3-scaled cache of a context-extended checkpoint);
|
|
* sliding (local) layers use an UNSCALED rotary at ``rope_local_base_freq``.
|
|
|
|
This mirrors our pretraining stack (``pretraining/flexattn_patch.py``) and is the exact
|
|
Gemma-3 dual-RoPE design. For a factor-1 checkpoint with equal bases and no scaling the
|
|
two rotaries collapse to one, so a single ``rope_theta`` is exact; for a context-extended
|
|
checkpoint (``rope_scaling`` set) the global and local caches diverge and both are needed.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from transformers.models.ministral.configuration_ministral import MinistralConfig
|
|
|
|
|
|
class MinistralDualRopeConfig(MinistralConfig):
|
|
r"""MinistralConfig + an unscaled local rotary base for the sliding layers.
|
|
|
|
Extra arg beyond :class:`~transformers.MinistralConfig`:
|
|
rope_local_base_freq (`float`, *optional*):
|
|
RoPE base for the sliding (local) layers, applied with NO rope scaling.
|
|
Defaults to ``rope_theta`` (i.e. single-RoPE, byte-identical to Ministral).
|
|
|
|
(``layer_types`` and ``sliding_window`` are inherited from Ministral and drive the
|
|
per-layer full/sliding attention pattern.)
|
|
"""
|
|
|
|
model_type = "ministral_dual_rope"
|
|
|
|
def __init__(self, rope_local_base_freq=None, **kwargs):
|
|
super().__init__(**kwargs)
|
|
if rope_local_base_freq is None:
|
|
rope_local_base_freq = self._global_rope_theta()
|
|
self.rope_local_base_freq = float(rope_local_base_freq)
|
|
|
|
def _global_rope_theta(self) -> float:
|
|
"""The GLOBAL rotary base, robust across transformers versions (5.x stores it in
|
|
``rope_parameters['rope_theta']``; 4.x in ``rope_theta``)."""
|
|
rp = getattr(self, "rope_parameters", None)
|
|
if isinstance(rp, dict) and "rope_theta" in rp:
|
|
return float(rp["rope_theta"])
|
|
return float(getattr(self, "rope_theta", 10000.0))
|
|
|
|
|
|
__all__ = ["MinistralDualRopeConfig"]
|