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"]