tiny-random-lfm3 / configuration_lfm3.py
rkazants's picture
Add tiny random LFM3 model for optimum-intel tests
cf7edf4 verified
Raw History Blame
13.9 kB
from typing import Optional, Union
from transformers.configuration_utils import PretrainedConfig
from transformers.modeling_rope_utils import RopeParameters
class Lfm3Config(PretrainedConfig):
r"""
Configuration class for [`Lfm3Model`] / [`Lfm3ForCausalLM`].
LFM3 dense is a hybrid architecture composed of:
- GQA full-attention layers (with headwise gated attention)
- sliding-window attention layers (may use a different rope_theta from global)
- gated short-conv layers (depthwise conv with B*x -> conv -> C*conv_out)
Args:
vocab_size (`int`, *optional*, defaults to 65536):
Vocabulary size.
hidden_size (`int`, *optional*, defaults to 2048):
Hidden / residual stream dimension.
intermediate_size (`int`, *optional*, defaults to 7168):
SwiGLU MLP intermediate dimension.
num_hidden_layers (`int`, *optional*, defaults to 32):
Total number of decoder layers.
num_attention_heads (`int`, *optional*, defaults to 32):
Number of query heads.
num_key_value_heads (`int`, *optional*, defaults to 8):
Number of KV heads (GQA) for full-attention layers.
head_dim (`int`, *optional*):
Per-head dim. When `None`, defaults to `hidden_size // num_attention_heads`.
May differ from that quotient (e.g. dim=1536, num_heads=16, head_dim=128).
max_position_embeddings (`int`, *optional*, defaults to 128_000):
Maximum sequence length the model is designed for.
norm_eps (`float`, *optional*, defaults to 1e-5):
RMSNorm epsilon.
use_cache (`bool`, *optional*, defaults to `True`):
Whether to return past_key_values.
pad_token_id, bos_token_id, eos_token_id (`int`, *optional*):
Special token ids.
tie_word_embeddings (`bool`, *optional*, defaults to `True`):
Tie input/output embeddings.
rope_parameters (`RopeParameters` or `dict[str, RopeParameters]`, *optional*):
Either a single RopeParameters dict (legacy form — broadcast to every
attention layer type) or a Gemma3-style dict keyed by `layer_type`:
`{"full_attention": {"rope_type": "default", "rope_theta": 1_000_000,
"partial_rotary_factor": 0.5}, "sliding_attention":
{"rope_type": "default", "rope_theta": 10_000}}`. Partial RoPE is a
global-attention-only knob (Gemma 4 convention); SWA layers always
use full RoPE, so the SWA entry never carries `partial_rotary_factor`.
For backward compatibility, callers may also pass a legacy flat
dict together with a separate `swa_rope_parameters` kwarg; both forms
are migrated to the nested representation in `__init__`.
conv_bias (`bool`, *optional*, defaults to `False`):
Whether the depthwise conv layers use a bias term.
conv_L_cache (`int`, *optional*, defaults to 3):
Convolution kernel size and the number of conv-state slots cached per layer.
sliding_window (`int`, *optional*):
Sliding-window size used by sliding-window-attention layers.
full_attn_idxs (`list[int]`, *optional*):
Convenience: layer indices that use full attention. Only used to build the
default `layer_types` when `layer_types` is not explicitly provided.
layer_types (`list[str]`, *optional*):
Per-layer operator type. Each entry is one of
`"full_attention"`, `"sliding_attention"`, `"conv"`.
When `None`, defaults to `"full_attention"` for indices in `full_attn_idxs`
and `"conv"` for the rest.
use_sandwich_norm (`bool`, *optional*, defaults to `False`):
Whether to apply an additional RMSNorm AFTER the operator and AFTER the
FFN (sandwich norm; quantization-friendly variant).
"""
model_type = "lfm3"
keys_to_ignore_at_inference = ["past_key_values"]
# Attention layer types that consume a per-layer-type entry in
# `rope_parameters`. `"conv"` layers don't use RoPE.
_attention_layer_types = ("full_attention", "sliding_attention")
def __init__(
self,
vocab_size: int = 65536,
hidden_size: int = 2048,
intermediate_size: int = 7168,
num_hidden_layers: int = 32,
num_attention_heads: int = 32,
num_key_value_heads: int = 8,
head_dim: Optional[int] = None,
max_position_embeddings: int = 128_000,
norm_eps: float = 1e-5,
use_cache: bool = True,
pad_token_id: int = 0,
bos_token_id: int = 1,
eos_token_id: int = 2,
tie_word_embeddings: bool = True,
rope_parameters: Optional[Union[RopeParameters, dict[str, RopeParameters]]] = None,
conv_bias: bool = False,
conv_L_cache: int = 3,
sliding_window: Optional[int] = None,
full_attn_idxs: Optional[list[int]] = None,
layer_types: Optional[list[str]] = None,
use_sandwich_norm: bool = False,
**kwargs,
):
self.vocab_size = vocab_size
self.hidden_size = hidden_size
self.intermediate_size = intermediate_size
self.num_hidden_layers = num_hidden_layers
self.num_attention_heads = num_attention_heads
self.num_key_value_heads = num_key_value_heads
self.head_dim = head_dim
self.max_position_embeddings = max_position_embeddings
self.norm_eps = norm_eps
self.use_cache = use_cache
self.conv_bias = conv_bias
self.conv_L_cache = conv_L_cache
self.sliding_window = sliding_window
self.use_sandwich_norm = use_sandwich_norm
self.layer_types = layer_types
if self.layer_types is None:
full_attn_idxs = (
full_attn_idxs if full_attn_idxs is not None else list(range(num_hidden_layers))
)
self.layer_types = [
"full_attention" if i in full_attn_idxs else "conv"
for i in range(num_hidden_layers)
]
# Catch the inconsistent combination at construction time so users see
# a clear error here instead of an opaque one from
# `create_sliding_window_causal_mask` deep in the forward pass.
if "sliding_attention" in self.layer_types and sliding_window is None:
raise ValueError(
"Lfm3Config has `sliding_attention` in `layer_types` but "
"`sliding_window` is None. Set `sliding_window` to the window size "
"(in tokens) you trained with."
)
# rope_parameters is the Gemma3-style nested dict keyed by attention
# layer_type (e.g. {"full_attention": {...}, "sliding_attention": {...}}).
# The trainer→HF conversion writes this nested form directly.
self.rope_parameters = self._validate_rope_parameters(rope_parameters)
tie_word_embeddings = kwargs.get("tie_embedding", tie_word_embeddings)
super().__init__(
pad_token_id=pad_token_id,
bos_token_id=bos_token_id,
eos_token_id=eos_token_id,
tie_word_embeddings=tie_word_embeddings,
**kwargs,
)
def _validate_rope_parameters(
self,
rope_parameters: Optional[dict[str, dict]],
) -> Optional[dict[str, dict]]:
"""Keep the per-attention-layer-type rope entries present in the model.
`rope_parameters` must be the Gemma3-style nested form keyed by attention
`layer_type`, e.g. ``{"full_attention": {"rope_type": "default",
"rope_theta": 1e6, "partial_rotary_factor": 0.5}, "sliding_attention":
{"rope_type": "default", "rope_theta": 1e4}}``. Conv layers have no rope.
Partial RoPE is a full-attention-only knob — the SWA entry never carries
`partial_rotary_factor`.
"""
if rope_parameters is None:
return None
attention_types_in_model = [
lt for lt in (self.layer_types or []) if lt in self._attention_layer_types
]
# Conv-only models don't use RoPE.
if not attention_types_in_model:
return None
if not (
isinstance(rope_parameters, dict)
and all(k in self._attention_layer_types for k in rope_parameters)
):
raise ValueError(
"Lfm3 expects nested per-layer-type `rope_parameters` keyed by "
"attention layer_type (e.g. {'full_attention': {...}, "
"'sliding_attention': {...}}); got "
f"{list(rope_parameters) if isinstance(rope_parameters, dict) else type(rope_parameters).__name__}."
)
return {
k: dict(v) for k, v in rope_parameters.items()
if k in attention_types_in_model
} or None
def convert_rope_params_to_dict(self, **kwargs):
"""Override the base PretrainedConfig hook so it doesn't pollute our
per-layer-type `rope_parameters` dict with top-level `rope_theta` /
`rope_type` keys.
The base implementation (modeling_rope_utils.py
`RotaryEmbeddingConfigMixin.convert_rope_params_to_dict`) does
`self.rope_parameters.setdefault("rope_theta", ...)` directly on
whatever we set — fine for legacy flat configs, but it puts top-level
keys onto our nested dict, which then makes `validate_rope` treat the
whole thing as flat and warn about the layer-type keys as
"Unrecognized". This override mirrors Gemma3Config: when the user
already supplied a nested dict, we trust it and only run validation.
"""
rope_scaling = kwargs.pop("rope_scaling", None)
if rope_scaling is not None and self.rope_parameters is None:
self.rope_parameters = rope_scaling
if self.rope_parameters is None:
# Conv-only models don't carry RoPE. Skipping the base helpers
# avoids a noisy "no RoPE parameters were found" log on every
# config construction.
return kwargs
is_nested = isinstance(self.rope_parameters, dict) and all(
k in self._attention_layer_types for k in self.rope_parameters.keys()
)
if is_nested:
# Ensure each entry has rope_type; default to "default" if missing.
for entry in self.rope_parameters.values():
entry.setdefault("rope_type", "default")
# validate_rope reads self.ignore_keys_at_rope_validation internally.
self.validate_rope()
return kwargs
# Flat / legacy form — fall back to the base behavior.
return super().convert_rope_params_to_dict(**kwargs)
def standardize_rope_params(self):
"""Override the base PretrainedConfig hook to skip non-attention
layer types (e.g. ``"conv"``).
Callers like vLLM's ``patch_rope_parameters`` (in
``vllm/transformers_utils/config.py``) invoke
``config.standardize_rope_params()`` directly after loading our
config, bypassing our ``convert_rope_params_to_dict`` override.
The base implementation in
``transformers/modeling_rope_utils.py:standardize_rope_params``
Case 2 (nested rope) does:
for layer_type in set(layer_types):
rope_parameters[layer_type].setdefault("rope_type", ...)
This iterates EVERY entry in ``layer_types`` and indexes
``rope_parameters[layer_type]`` blindly. Lfm3 has ``"conv"`` in
``layer_types`` (conv layers don't use RoPE), and our nested
``rope_parameters`` correctly omits a ``"conv"`` entry — so the
blind index raises ``KeyError: 'conv'`` and vLLM fails to load.
Gemma3 dodges this because its ``layer_types`` only contains
attention types. Upstream LFM2 dodges it because its
``rope_parameters`` is flat (Case 1), not nested (Case 2).
Our override: iterate only the layer types that actually have an
entry in ``rope_parameters``, applying the same defaulting logic
the base does. For the flat / legacy form, defer to base.
"""
rope_parameters = getattr(self, "rope_parameters", None)
if not rope_parameters:
return
is_nested = isinstance(rope_parameters, dict) and all(
k in self._attention_layer_types for k in rope_parameters.keys()
)
if not is_nested:
return super().standardize_rope_params()
rope_theta = getattr(self, "rope_theta", None)
partial_rotary_factor = getattr(self, "partial_rotary_factor", None)
for layer_type in rope_parameters:
entry = rope_parameters[layer_type]
entry.setdefault("rope_type", entry.get("type", "default"))
if rope_theta is not None:
entry.setdefault("rope_theta", rope_theta)
# Partial RoPE is a global-attention-only knob (Gemma 4 invariant);
# never apply it to SWA, even if a top-level
# `partial_rotary_factor` is set as a legacy attribute.
if partial_rotary_factor is not None and layer_type == "full_attention":
entry["partial_rotary_factor"] = partial_rotary_factor
if entry["rope_type"] in ("llama3", "yarn", "longrope"):
entry.setdefault(
"original_max_position_embeddings", self.max_position_embeddings
)
@property
def rope_scaling(self):
"""Backward compatibility property: rope_scaling maps to rope_parameters."""
return self.rope_parameters
@rope_scaling.setter
def rope_scaling(self, value):
self.rope_parameters = value
__all__ = ["Lfm3Config"]