xlstm: add configuration and modeling (own one)

configuration_xlstm.py

@@ -0,0 +1,97 @@
+import json
+from typing import Any, Dict, Optional
+
+from dacite import Config as DaciteConfig
+from dacite import from_dict
+from omegaconf import OmegaConf
+from transformers.configuration_utils import PretrainedConfig
+from xlstm import xLSTMLMModelConfig
+
+# from .config_presets import xlstm_cfg_map
+
+
+class xLSTMConfig(PretrainedConfig):
+    """XLSTM configuration class.
+    We seperate the specific xLSTM model configuration
+    from the rest due to the heavy nesting of the configuration.
+    """
+
+    model_type = "xlstm"
+
+    def __init__(
+        self, vocab_size: int = 32000, config: Optional[Dict[str, Any]] = None, **kwargs
+    ):
+        super().__init__(**kwargs)
+
+        cfg = OmegaConf.create(config)
+        cfg["vocab_size"] = vocab_size
+        for key, value in kwargs.items():
+            cfg[key] = value
+
+        self._xlstm_config = cfg
+        self.vocab_size = vocab_size
+        self.embedding_dim = cfg.get("embedding_dim")
+        self.context_length = cfg.get("context_length")
+
+    def to_xlstm_config(self):
+        return from_dict(
+            data_class=xLSTMLMModelConfig,
+            data=OmegaConf.to_container(self._xlstm_config),
+            config=DaciteConfig(strict=True),
+        )
+
+    def to_dict(self) -> Dict[str, Any]:
+        """
+        Converts the configuration to a dictionary for serialization.
+        """
+        output = super().to_dict()
+        output["_xlstm_config"] = OmegaConf.to_container(
+            self._xlstm_config, resolve=True
+        )
+        relevant_keys = [
+            "vocab_size",
+            "embedding_dim",
+            "context_length",
+            "torch_dtype",
+            "_xlstm_config",
+            "transformers_version",
+            "architectures",
+            "model_type",
+        ]
+        output_ = output.copy()
+        for key in output.keys():
+            if key not in relevant_keys:
+                output_.pop(key)
+        return output_
+
+    @classmethod
+    def from_dict(cls, config_dict: Dict[str, Any], **kwargs):
+        """
+        Creates a configuration instance from a dictionary.
+        """
+        xlstm_config = config_dict.pop("_xlstm_config")
+        vocab_size = config_dict.pop("vocab_size")
+        config = cls(vocab_size=vocab_size, config=xlstm_config)
+        if "auto_map" in config_dict and config_dict["auto_map"]:
+            setattr(config, "auto_map", config_dict.pop("auto_map"))
+
+        # breakpoint()
+        # config.xlstm_config = xlstm_config
+        if "return_unused_kwargs" in kwargs and kwargs["return_unused_kwargs"]:
+            return config, {}
+
+        return config
+
+    def to_json_string(self, *args, **kwargs) -> str:
+        """
+        Serializes the instance to a JSON string.
+        """
+        return json.dumps(self.to_dict(), indent=2)
+
+    @classmethod
+    def from_json_string(cls, json_string: str):
+        """
+        Deserializes the instance from a JSON string.
+        """
+        config_dict = json.loads(json_string)
+        return cls.from_dict(config_dict)

modeling_xlstm.py

@@ -0,0 +1,214 @@
+from typing import Optional, Sequence, Tuple, Union
+
+import torch
+from torch import nn
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import BaseModelOutput, CausalLMOutputWithPast
+from xlstm.components.init import small_init_init_
+from xlstm.utils import WeightDecayOptimGroupMixin
+from xlstm.xlstm_block_stack import xLSTMBlockStack as _xLSTMBlockStack
+
+from .configuration_xlstm import xLSTMConfig
+
+
+class xLSTMPreTrainedModel(PreTrainedModel):
+    """Base class for all models."""
+
+    config_class = xLSTMConfig
+
+
+class xLSTMBlockStack(_xLSTMBlockStack):
+    """Small wrapper to expose hidden states"""
+
+    def forward(
+        self, x: torch.Tensor, **kwargs
+    ) -> Tuple[torch.Tensor, Sequence[torch.Tensor]]:
+        hidden_states = ()
+        for block in self.blocks:
+            x = block(x, **kwargs)
+            hidden_states += (x,)
+
+        x = self.post_blocks_norm(x)
+
+        return x, hidden_states
+
+
+class xLSTMModel(xLSTMPreTrainedModel):
+    def __init__(self, config: xLSTMConfig):
+        super().__init__(config)
+        self.config = config
+
+        self.token_embedding = nn.Embedding(
+            num_embeddings=config.vocab_size, embedding_dim=config.embedding_dim
+        )
+        _config = config.to_xlstm_config()
+
+        self.emb_dropout = (
+            nn.Dropout(_config.dropout)
+            if _config.add_embedding_dropout
+            else nn.Identity()
+        )
+
+        self.xlstm_block_stack = xLSTMBlockStack(config=_config)
+
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor,
+        output_hidden_states: Optional[bool] = None,
+        return_dict=Optional[bool],
+    ) -> Union[Tuple, BaseModelOutput]:
+        token_embedding = self.token_embedding(input_ids)
+        x = self.emb_dropout(token_embedding)
+        x, hidden_states = self.xlstm_block_stack(x)
+
+        if output_hidden_states:
+            hidden_states = (token_embedding,) + hidden_states
+
+        if not return_dict:
+            return x, hidden_states
+
+        return BaseModelOutput(
+            last_hidden_state=x,
+            hidden_states=hidden_states if output_hidden_states else None,
+        )
+
+
+class xLSTMForCausalLM(xLSTMPreTrainedModel, WeightDecayOptimGroupMixin):
+    _tied_weights_keys = ["lm_head.weight"]
+
+    def __init__(self, config: xLSTMConfig, **kwargs):
+        super().__init__(config)
+        self.config = config
+        self.vocab_size = config.vocab_size
+
+        self.model = xLSTMModel(config)
+
+        self.lm_head = nn.Linear(
+            in_features=config.embedding_dim,
+            out_features=config.vocab_size,
+            bias=False,
+        )
+
+        self.post_init()
+        # TODO: Add option for up-projection
+
+    def get_input_embeddings(self):
+        return self.model.token_embedding
+
+    def set_input_embeddings(self, value: nn.Module):
+        self.model.token_embedding = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, value):
+        self.lm_head = value
+
+    def reset_parameters(self):
+        self.model.xlstm_block_stack.reset_parameters()
+
+        small_init_init_(
+            self.get_input_embeddings().weight, dim=self.config.embedding_dim
+        )
+
+        if not self.config.tie_word_embeddings:
+            small_init_init_(
+                self.get_output_embeddings().weight, dim=self.config.embedding_dim
+            )
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        labels: Optional[torch.LongTensor] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ):
+        output = self.model(
+            input_ids,
+            output_hidden_states=output_hidden_states,
+        )
+
+        hidden_state = output[0]
+
+        logits = self.lm_head(hidden_state)
+        logits = logits.float()
+
+        loss = None
+
+        if labels is not None:
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+
+            loss_fct = nn.CrossEntropyLoss()
+            shift_logits = shift_logits.view(-1, self.config.vocab_size)
+            shift_labels = shift_labels.view(-1)
+
+            shift_labels = shift_labels.to(shift_logits.device)
+            loss = loss_fct(shift_logits, shift_labels)
+
+        if not return_dict:
+            output = (logits,) + output[1:]
+            return ((loss,) + output) if loss is not None else output
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            hidden_states=output.hidden_states,
+        )
+
+    def step(
+        self,
+        idx: torch.Tensor,
+        state: dict[str, dict[str, tuple[torch.Tensor, ...]]] = None,
+        **kwargs,
+    ) -> tuple[torch.Tensor, dict[str, dict[str, tuple[torch.Tensor, ...]]]]:
+        x = self.token_embedding(idx)
+        x = self.emb_dropout(x)
+        x, state = self.xlstm_block_stack.step(x, state=state, **kwargs)
+        logits = self.lm_head(x)
+        return logits, state
+
+    def _create_weight_decay_optim_groups(
+        self, **kwargs
+    ) -> tuple[Sequence[nn.Parameter], Sequence[nn.Parameter]]:
+        weight_decay, no_weight_decay = super()._create_weight_decay_optim_groups(
+            **kwargs
+        )
+        # remove token embedding and add it to the correct group, accrording to the config
+        weight_decay = list(weight_decay)
+        removed = 0
+        for idx in range(len(weight_decay)):
+            if weight_decay[idx - removed] is self.get_input_embeddings().weight:
+                weight_decay.pop(idx - removed)
+                removed += 1
+        weight_decay = tuple(weight_decay)
+
+        # TODO: Fix this
+        # if self.config.weight_decay_on_embedding:
+        if True:
+            weight_decay += (self.get_input_embeddings().weight,)
+        else:
+            no_weight_decay += (self.get_input_embeddings().weight,)
+
+        return weight_decay, no_weight_decay
+
+    def resize_token_embeddings(self, new_num_tokens: int) -> nn.Embedding:
+        new_embeddings = nn.Embedding(
+            new_num_tokens, self.token_embedding.embedding_dim
+        )
+        self.token_embedding = new_embeddings.to(self.device)
+        return new_embeddings
+
+    def tie_weights(self):
+        self.get_output_embeddings().weight = self.get_input_embeddings().weight
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids,
+        **kwargs,
+    ):
+        model_inputs = {
+            "input_ids": input_ids.to(self.device),
+        }
+        return model_inputs