""" model/model.py -------------- SLMModel and SLMForCausalLM — the full model registered with HuggingFace. SLMModel: the core transformer (embeddings + decoder stack + final norm). SLMForCausalLM: adds the language model head and loss computation. Design: - No bias anywhere - Pre-norm throughout - KV cache support for efficient autoregressive generation - Compatible with HuggingFace generate(), trl, lm-evaluation-harness, vLLM Important implementation detail: SLMModel is a plain nn.Module. SLMForCausalLM is the only PreTrainedModel. This follows the standard HuggingFace architecture pattern used by Llama, Mistral, GPT-NeoX, Phi, etc.: SLMForCausalLM(PreTrainedModel) └── SLMModel(nn.Module) Only the outer class calls post_init(), so initialization and HF save/load behavior are controlled from one PreTrainedModel. """ from typing import Optional, Union import torch import torch.nn as nn from torch.nn import CrossEntropyLoss from transformers import PreTrainedModel from transformers.cache_utils import Cache from transformers.generation import GenerationMixin from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast from .block import SLMDecoderBlock from .config import SLMConfig from .norm import RMSNorm def _extract_kv_from_dynamic_cache( cache: Cache, n_layers: int, ) -> list[Optional[tuple[torch.Tensor, torch.Tensor]]]: """ Extract per-layer (k, v) tuples from a DynamicCache object. Handles two DynamicCache formats: - transformers v5: cache.layers is a list of DynamicLayer objects with .keys and .values tensor attributes. - older versions: cache.key_cache / cache.value_cache are lists of tensors. Returns a list of length n_layers where each entry is either a (k, v) tuple or None if no cached state exists for that layer. """ result = [] cache_layers = getattr(cache, "layers", None) if cache_layers is not None: for i in range(n_layers): if i < len(cache_layers) and cache_layers[i].is_initialized: result.append((cache_layers[i].keys, cache_layers[i].values)) else: result.append(None) return result key_cache = getattr(cache, "key_cache", None) value_cache = getattr(cache, "value_cache", None) if key_cache is not None: for i in range(n_layers): if i < len(key_cache): result.append((key_cache[i], value_cache[i])) else: result.append(None) return result return [None] * n_layers class SLMModel(nn.Module): """ The core SLM transformer — embeddings, decoder stack, final norm. Does not include the LM head — use SLMForCausalLM for language modelling. Important: This is intentionally a plain nn.Module, not a PreTrainedModel. The outer SLMForCausalLM owns HF initialization, saving, and loading. """ def __init__(self, config: SLMConfig): super().__init__() self.config = config self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size) self.layers = nn.ModuleList( [SLMDecoderBlock(config, layer_idx=i) for i in range(config.num_hidden_layers)] ) self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.gradient_checkpointing = False def get_input_embeddings(self) -> nn.Embedding: return self.embed_tokens def set_input_embeddings(self, value: nn.Embedding) -> None: self.embed_tokens = value def _set_gradient_checkpointing(self, module: nn.Module, value: bool = False) -> None: if isinstance(module, SLMModel): module.gradient_checkpointing = value def forward( self, input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, past_key_values: Optional[Union[Cache, list[tuple[torch.Tensor, torch.Tensor]]]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, use_cache: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, ) -> BaseModelOutputWithPast: use_cache = use_cache if use_cache is not None else self.config.use_cache return_dict = return_dict if return_dict is not None else getattr( self.config, "return_dict", getattr(self.config, "use_return_dict", True) ) if self.gradient_checkpointing and self.training and use_cache: use_cache = False if inputs_embeds is None: inputs_embeds = self.embed_tokens(input_ids) hidden_states = inputs_embeds if past_key_values is None: past_key_values = [None] * len(self.layers) elif isinstance(past_key_values, Cache): past_key_values = _extract_kv_from_dynamic_cache( past_key_values, len(self.layers), ) next_cache: list | None = [] if use_cache else None all_hidden_states: list | None = [] if output_hidden_states else None for layer, past_kv in zip(self.layers, past_key_values): if output_hidden_states: all_hidden_states.append(hidden_states) if self.gradient_checkpointing and self.training: layer_out = self._gradient_checkpointing_func( layer.__call__, hidden_states, attention_mask, None, False, ) hidden_states = layer_out[0] layer_kv = layer_out[1] if len(layer_out) > 1 else None else: hidden_states, layer_kv = layer( hidden_states, attention_mask=attention_mask, past_key_value=past_kv, use_cache=use_cache, ) if use_cache: next_cache.append(layer_kv) hidden_states = self.norm(hidden_states) if output_hidden_states: all_hidden_states.append(hidden_states) if not return_dict: return tuple( v for v in [hidden_states, next_cache, all_hidden_states] if v is not None ) return BaseModelOutputWithPast( last_hidden_state=hidden_states, past_key_values=next_cache, hidden_states=all_hidden_states, ) class SLMForCausalLM(PreTrainedModel, GenerationMixin): """ SLM with a language modelling head for causal language modelling. This is the only PreTrainedModel in the architecture. It owns: - initialization via post_init() - save_pretrained() (extended to bundle architecture .py files) - tied embedding behavior - HF generation compatibility Loading uses a custom safe from_pretrained() path that loads SLMConfig, instantiates the model, reads model.safetensors or pytorch_model.bin, applies the state dict directly, and re-ties lm_head when embeddings are tied. This avoids local AutoModel loading issues observed with this custom architecture. """ config_class = SLMConfig base_model_prefix = "model" supports_gradient_checkpointing = True _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} # lm_head.weight is intentionally omitted from safetensors when tied # to model.embed_tokens.weight. _keys_to_ignore_on_load_missing = [r"lm_head\.weight"] def __init__(self, config: SLMConfig): super().__init__(config) self.model = SLMModel(config) self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) self.post_init() @classmethod def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs): """ Safe SLM loader. Transformers local AutoModel loading has been observed to instantiate this custom architecture while leaving most checkpoint tensors at fresh initialization. This loader uses the verified path: SLMConfig -> cls(config) -> safetensors/torch load -> load_state_dict Supports: - local checkpoint directories - Hub repo IDs - model.safetensors - pytorch_model.bin - dtype / torch_dtype strings from CLI tools - tied lm_head.weight missing from checkpoint """ import os from pathlib import Path import safetensors.torch config = kwargs.pop("config", None) torch_dtype = kwargs.pop("torch_dtype", None) dtype = kwargs.pop("dtype", None) device_map = kwargs.pop("device_map", None) output_loading_info = kwargs.pop("output_loading_info", False) revision = kwargs.pop("revision", None) cache_dir = kwargs.pop("cache_dir", None) token = kwargs.pop("token", None) local_files_only = kwargs.pop("local_files_only", False) # Accepted by many HF call sites, but not needed by this loader. kwargs.pop("low_cpu_mem_usage", None) kwargs.pop("trust_remote_code", None) kwargs.pop("weights_only", None) kwargs.pop("use_safetensors", None) if dtype is not None and torch_dtype is None: torch_dtype = dtype path = str(pretrained_model_name_or_path) # Hub repo ID support: resolve repo into a local snapshot first. if not os.path.isdir(path): from huggingface_hub import snapshot_download snapshot_kwargs = { "repo_id": path, "local_files_only": local_files_only, "allow_patterns": [ "config.json", "generation_config.json", "model.safetensors", "pytorch_model.bin", "tokenizer.json", "tokenizer_config.json", "special_tokens_map.json", "chat_template.jinja", "*.py", ], } if revision is not None: snapshot_kwargs["revision"] = revision if cache_dir is not None: snapshot_kwargs["cache_dir"] = cache_dir if token is not None: snapshot_kwargs["token"] = token path = snapshot_download(**snapshot_kwargs) if config is None: config = SLMConfig.from_pretrained(path) model = cls(config, *model_args) safetensors_path = Path(path) / "model.safetensors" bin_path = Path(path) / "pytorch_model.bin" if safetensors_path.exists(): state_dict = safetensors.torch.load_file(str(safetensors_path), device="cpu") elif bin_path.exists(): state_dict = torch.load(str(bin_path), map_location="cpu") else: raise FileNotFoundError( f"No model.safetensors or pytorch_model.bin found in {path}" ) result = model.load_state_dict(state_dict, strict=False) allowed_missing = set() if getattr(config, "tie_word_embeddings", False): allowed_missing.add("lm_head.weight") missing_keys = set(result.missing_keys) unexpected_keys = set(result.unexpected_keys) unexpected_missing = sorted(k for k in missing_keys if k not in allowed_missing) if unexpected_missing: raise RuntimeError( f"Missing keys while loading {path}: {unexpected_missing}" ) if unexpected_keys: raise RuntimeError( f"Unexpected keys while loading {path}: {sorted(unexpected_keys)}" ) if getattr(config, "tie_word_embeddings", False): model.tie_weights() # Normalize dtype passed by HF / lm-eval / CLI tools. if isinstance(torch_dtype, str): original_torch_dtype = torch_dtype if torch_dtype == "auto": cfg_dtype = getattr(config, "torch_dtype", None) if isinstance(cfg_dtype, str): torch_dtype = getattr(torch, cfg_dtype, None) else: torch_dtype = cfg_dtype else: torch_dtype = { "float16": torch.float16, "fp16": torch.float16, "bfloat16": torch.bfloat16, "bf16": torch.bfloat16, "float32": torch.float32, "fp32": torch.float32, }.get(torch_dtype, getattr(torch, torch_dtype, None)) if torch_dtype is None: raise ValueError( f"Unknown torch_dtype string: {original_torch_dtype!r}. " "Expected a torch.dtype or one of: " "'bfloat16', 'bf16', 'float16', 'fp16', 'float32', " "'fp32', 'auto'." ) if torch_dtype is not None: model = model.to(dtype=torch_dtype) # Minimal local device_map support. if device_map is not None: if device_map == "auto" and torch.cuda.is_available(): model = model.to("cuda") elif isinstance(device_map, str) and device_map != "auto": model = model.to(device_map) model.eval() if output_loading_info: info = { "missing_keys": sorted(missing_keys), "unexpected_keys": sorted(unexpected_keys), "mismatched_keys": [], "error_msgs": [], } return model, info return model def save_pretrained(self, save_directory, *args, **kwargs): """ Save model and bundle architecture .py files for remote-code loading. Standard HF save writes config.json + model.safetensors. We additionally copy the SLM architecture source into the checkpoint root so the checkpoint loads via `AutoModelForCausalLM.from_pretrained(path, trust_remote_code=True)` from any environment, matching the auto_map paths declared in SLMConfig. """ import shutil from pathlib import Path result = super().save_pretrained(save_directory, *args, **kwargs) src = Path(__file__).parent dst = Path(save_directory) for name in ("config.py", "model.py", "attention.py", "block.py", "mlp.py", "norm.py"): shutil.copy2(src / name, dst / name) return result def _init_weights(self, module: nn.Module) -> None: """ Initialize weights with config.initializer_range. This runs on every submodule when post_init() recurses, including modules inside SLMModel. SLMModel intentionally does not define its own _init_weights; this is the single source of init policy. """ std = self.config.initializer_range if isinstance(module, nn.Linear): module.weight.data.normal_(mean=0.0, std=std) if module.bias is not None: module.bias.data.zero_() elif isinstance(module, nn.Embedding): module.weight.data.normal_(mean=0.0, std=std) if module.padding_idx is not None: module.weight.data[module.padding_idx].zero_() def tie_weights(self, **kwargs) -> None: """ Tie LM head weights to input embeddings when tie_word_embeddings=True. Direct assignment is used because transformers==5.5.4 in this environment does not expose _tie_or_clone_weights on PreTrainedModel. """ if self.config.tie_word_embeddings: self.lm_head.weight = self.model.embed_tokens.weight def get_input_embeddings(self) -> nn.Embedding: return self.model.embed_tokens def set_input_embeddings(self, value: nn.Embedding) -> None: self.model.embed_tokens = value def get_output_embeddings(self) -> nn.Linear: return self.lm_head def set_output_embeddings(self, new_embeddings: nn.Linear) -> None: self.lm_head = new_embeddings def get_decoder(self) -> SLMModel: return self.model def set_decoder(self, decoder: SLMModel) -> None: self.model = decoder def forward( self, input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, past_key_values: Optional[Union[Cache, list[tuple[torch.Tensor, torch.Tensor]]]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, labels: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, ) -> CausalLMOutputWithPast: outputs = self.model( input_ids=input_ids, attention_mask=attention_mask, past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, output_hidden_states=output_hidden_states, return_dict=True, cache_position=cache_position, ) hidden_states = outputs.last_hidden_state logits = self.lm_head(hidden_states) loss = None if labels is not None: shift_logits = logits[..., :-1, :].contiguous() shift_labels = labels[..., 1:].contiguous() loss = CrossEntropyLoss()( shift_logits.view(-1, self.config.vocab_size), shift_labels.view(-1), ) return CausalLMOutputWithPast( loss=loss, logits=logits, past_key_values=outputs.past_key_values, hidden_states=outputs.hidden_states, ) def prepare_inputs_for_generation( self, input_ids: torch.LongTensor, past_key_values: Optional[Union[Cache, list]] = None, attention_mask: Optional[torch.Tensor] = None, inputs_embeds: Optional[torch.Tensor] = None, cache_position: Optional[torch.LongTensor] = None, **kwargs, ) -> dict: """ Called by HuggingFace generate() at each decoding step. Slices input_ids to only positions that have not yet been processed. """ if cache_position is not None: input_ids = input_ids[:, -cache_position.shape[0]:] elif past_key_values is not None: input_ids = input_ids[:, -1:] if inputs_embeds is not None and past_key_values is None: model_inputs = {"inputs_embeds": inputs_embeds} else: model_inputs = {"input_ids": input_ids} model_inputs.update({ "past_key_values": past_key_values, "use_cache": kwargs.get("use_cache", True), "attention_mask": attention_mask, "cache_position": cache_position, }) return model_inputs def _reorder_cache( self, past_key_values: Union[Cache, list[tuple[torch.Tensor, torch.Tensor]]], beam_idx: torch.Tensor, ) -> Union[Cache, list[tuple[torch.Tensor, torch.Tensor]]]: """ Reorder KV cache for beam search. """ if isinstance(past_key_values, Cache): past_key_values.reorder_cache(beam_idx) return past_key_values return [ (k.index_select(0, beam_idx), v.index_select(0, beam_idx)) for k, v in past_key_values ]