Model save

README.md

@@ -0,0 +1,60 @@
+---
+library_name: transformers
+tags:
+- generated_from_trainer
+model-index:
+- name: lstm-resnet-12layer-v5
+  results: []
+---
+
+<!-- This model card has been generated automatically according to the information the Trainer had access to. You
+should probably proofread and complete it, then remove this comment. -->
+
+# lstm-resnet-12layer-v5
+
+This model is a fine-tuned version of [](https://huggingface.co/) on an unknown dataset.
+It achieves the following results on the evaluation set:
+- Loss: 5.1418
+- Num Input Tokens Seen: 9437184000
+
+## Model description
+
+More information needed
+
+## Intended uses & limitations
+
+More information needed
+
+## Training and evaluation data
+
+More information needed
+
+## Training procedure
+
+### Training hyperparameters
+
+The following hyperparameters were used during training:
+- learning_rate: 0.001
+- train_batch_size: 16
+- eval_batch_size: 16
+- seed: 42
+- distributed_type: multi-GPU
+- num_devices: 2
+- gradient_accumulation_steps: 16
+- total_train_batch_size: 512
+- total_eval_batch_size: 32
+- optimizer: Use OptimizerNames.ADAMW_TORCH_FUSED with betas=(0.9,0.95) and epsilon=1e-08 and optimizer_args=No additional optimizer arguments
+- lr_scheduler_type: cosine
+- lr_scheduler_warmup_steps: 500
+- num_epochs: 1
+
+### Training results
+
+
+
+### Framework versions
+
+- Transformers 5.3.0
+- Pytorch 2.10.0+cu128
+- Datasets 4.7.0
+- Tokenizers 0.22.2

final_model/config.json

@@ -0,0 +1,21 @@
+{
+  "architectures": [
+    "LSTMResnetForCausalLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "lstm.LSTMResnetLanguageModelConfig",
+    "AutoModelForCausalLM": "lstm.LSTMResnetForCausalLM"
+  },
+  "bos_token_id": 128000,
+  "dropout": 0.1,
+  "dtype": "float32",
+  "embed_dim": 1024,
+  "eos_token_id": 128001,
+  "hidden_size": 1024,
+  "model_type": "lstm_resnet",
+  "num_layers": 12,
+  "tie_word_embeddings": true,
+  "transformers_version": "5.3.0",
+  "use_cache": false,
+  "vocab_size": 128256
+}

final_model/lstm.py

@@ -0,0 +1,267 @@
+"""
+LSTM language model compatible with HuggingFace Trainer and AutoModelForCausalLM.
+
+Designed to be a drop-in replacement for LlamaForCausalLM in the Fourier
+emergence experiments: same tokenizer, same data pipeline, same analysis code.
+"""
+
+import torch
+import torch.nn as nn
+from dataclasses import dataclass, field
+from typing import Optional, Tuple
+
+from transformers import PretrainedConfig, PreTrainedModel
+from transformers.modeling_outputs import CausalLMOutput
+
+
+# ── Config ─────────────────────────────────────────────────────────
+
+
+class LSTMLanguageModelConfig(PretrainedConfig):
+    model_type = "lstm_lm"
+
+    def __init__(
+        self,
+        vocab_size: int = 128256,
+        embed_dim: int = 1024,
+        hidden_size: int = 1024,
+        num_layers: int = 2,
+        dropout: float = 0.1,
+        tie_word_embeddings: bool = True,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.vocab_size = vocab_size
+        self.embed_dim = embed_dim
+        self.hidden_size = hidden_size
+        self.num_layers = num_layers
+        self.dropout = dropout
+        self.tie_word_embeddings = tie_word_embeddings
+
+
+# ── Model ──────────────────────────────────────────────────────────
+
+
+class LSTMForCausalLM(PreTrainedModel):
+    config_class = LSTMLanguageModelConfig
+    _tied_weights_keys = {"lm_head.weight": "embed_tokens.weight"}
+
+    def __init__(self, config: LSTMLanguageModelConfig):
+        super().__init__(config)
+        self.config = config
+
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.embed_dim)
+        self.lstm = nn.LSTM(
+            input_size=config.embed_dim,
+            hidden_size=config.hidden_size,
+            num_layers=config.num_layers,
+            batch_first=True,
+            dropout=config.dropout if config.num_layers > 1 else 0.0,
+        )
+        self.drop = nn.Dropout(config.dropout)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+        if config.tie_word_embeddings:
+            if config.embed_dim == config.hidden_size:
+                self.lm_head.weight = self.embed_tokens.weight
+            # If dims don't match, silently skip tying (projection needed)
+
+        self.post_init()
+
+    def get_input_embeddings(self) -> nn.Embedding:
+        return self.embed_tokens
+
+    def set_input_embeddings(self, value: nn.Embedding):
+        self.embed_tokens = value
+
+    def get_output_embeddings(self) -> nn.Linear:
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings: nn.Linear):
+        self.lm_head = new_embeddings
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> CausalLMOutput:
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        embeds = self.embed_tokens(input_ids)
+        embeds = self.drop(embeds)
+
+        lstm_out, _ = self.lstm(embeds)
+        lstm_out = self.drop(lstm_out)
+
+        logits = self.lm_head(lstm_out)
+
+        loss = None
+        if labels is not None:
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            loss = nn.functional.cross_entropy(
+                shift_logits.view(-1, self.config.vocab_size),
+                shift_labels.view(-1),
+                ignore_index=-100,
+            )
+
+        hidden_states = None
+        if output_hidden_states:
+            # Return embedding + lstm output for compatibility with analysis
+            hidden_states = (embeds, lstm_out)
+
+        return CausalLMOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=hidden_states,
+        )
+
+
+# ── Residual LSTM (pre-norm, per-layer skip) ──────────────────────
+
+
+class LSTMResnetLanguageModelConfig(PretrainedConfig):
+    """Same fields as :class:`LSTMLanguageModelConfig` but a distinct model_type
+    so AutoModel resolves to the residual variant."""
+    model_type = "lstm_resnet"
+
+    def __init__(
+        self,
+        vocab_size: int = 128256,
+        embed_dim: int = 1024,
+        hidden_size: int = 1024,
+        num_layers: int = 2,
+        dropout: float = 0.1,
+        tie_word_embeddings: bool = True,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.vocab_size = vocab_size
+        self.embed_dim = embed_dim
+        self.hidden_size = hidden_size
+        self.num_layers = num_layers
+        self.dropout = dropout
+        self.tie_word_embeddings = tie_word_embeddings
+
+
+class LSTMResnetForCausalLM(PreTrainedModel):
+    """12-layer (or any depth) LSTM with pre-norm residual connections
+    around each layer — i.e. a Transformer-block layout where the
+    attention sub-block is replaced by a single-layer LSTM.
+    """
+
+    config_class = LSTMResnetLanguageModelConfig
+    _tied_weights_keys = {"lm_head.weight": "embed_tokens.weight"}
+
+    def __init__(self, config: LSTMResnetLanguageModelConfig):
+        super().__init__(config)
+        self.config = config
+
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.embed_dim)
+        self.embed_drop = nn.Dropout(config.dropout)
+        self.in_proj = (
+            nn.Linear(config.embed_dim, config.hidden_size, bias=False)
+            if config.embed_dim != config.hidden_size
+            else nn.Identity()
+        )
+
+        self.layers = nn.ModuleList(
+            [
+                nn.LSTM(
+                    input_size=config.hidden_size,
+                    hidden_size=config.hidden_size,
+                    num_layers=1,
+                    batch_first=True,
+                )
+                for _ in range(config.num_layers)
+            ]
+        )
+        self.norms = nn.ModuleList(
+            [nn.LayerNorm(config.hidden_size) for _ in range(config.num_layers)]
+        )
+        self.drops = nn.ModuleList(
+            [nn.Dropout(config.dropout) for _ in range(config.num_layers)]
+        )
+        self.final_norm = nn.LayerNorm(config.hidden_size)
+
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        if (
+            config.tie_word_embeddings
+            and config.embed_dim == config.hidden_size
+        ):
+            self.lm_head.weight = self.embed_tokens.weight
+
+        self.post_init()
+
+    def get_input_embeddings(self) -> nn.Embedding:
+        return self.embed_tokens
+
+    def set_input_embeddings(self, value: nn.Embedding):
+        self.embed_tokens = value
+
+    def get_output_embeddings(self) -> nn.Linear:
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings: nn.Linear):
+        self.lm_head = new_embeddings
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> CausalLMOutput:
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        x = self.embed_drop(self.embed_tokens(input_ids))
+        x = self.in_proj(x)
+
+        all_hidden = [x] if output_hidden_states else None
+        for lstm, norm, drop in zip(self.layers, self.norms, self.drops):
+            residual = x
+            normed = norm(x)
+            out, _ = lstm(normed)
+            x = residual + drop(out)
+            if output_hidden_states:
+                all_hidden.append(x)
+
+        x = self.final_norm(x)
+        logits = self.lm_head(x)
+
+        loss = None
+        if labels is not None:
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            loss = nn.functional.cross_entropy(
+                shift_logits.view(-1, self.config.vocab_size),
+                shift_labels.view(-1),
+                ignore_index=-100,
+            )
+
+        return CausalLMOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=tuple(all_hidden) if output_hidden_states else None,
+        )
+
+
+# Register so AutoModelForCausalLM.from_pretrained works on saved checkpoints.
+# register_for_auto_class() saves auto_map in config.json when saved.
+# AutoConfig/AutoModel.register() makes the class findable in any process
+# that imports this module — no trust_remote_code needed.
+from transformers import AutoConfig, AutoModelForCausalLM
+
+LSTMLanguageModelConfig.register_for_auto_class()
+LSTMForCausalLM.register_for_auto_class("AutoModelForCausalLM")
+AutoConfig.register("lstm_lm", LSTMLanguageModelConfig)
+AutoModelForCausalLM.register(LSTMLanguageModelConfig, LSTMForCausalLM)
+
+LSTMResnetLanguageModelConfig.register_for_auto_class()
+LSTMResnetForCausalLM.register_for_auto_class("AutoModelForCausalLM")
+AutoConfig.register("lstm_resnet", LSTMResnetLanguageModelConfig)
+AutoModelForCausalLM.register(LSTMResnetLanguageModelConfig, LSTMResnetForCausalLM)

final_model/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ce113d19e337221e3a0fdf1653d0698d915a8c344c3a7bfa4513fc2497564021
+size 928496296

final_model/tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6b9e4e7fb171f92fd137b777cc2714bf87d11576700a1dcd7a399e7bbe39537b
+size 17209920

final_model/tokenizer_config.json

@@ -0,0 +1,13 @@
+{
+  "backend": "tokenizers",
+  "bos_token": "<|begin_of_text|>",
+  "clean_up_tokenization_spaces": true,
+  "eos_token": "<|end_of_text|>",
+  "is_local": false,
+  "model_input_names": [
+    "input_ids",
+    "attention_mask"
+  ],
+  "model_max_length": 131072,
+  "tokenizer_class": "TokenizersBackend"
+}

final_model/training_args.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:940ae510ab19b8a3510d50d5f70e09387d1d40b89931566e2add3ed341dab96f
+size 5329

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:eef42f0560e9cfddb377d2015a7700ef5136b7e1fb73e220d2064b05855dcab6
+oid sha256:ce113d19e337221e3a0fdf1653d0698d915a8c344c3a7bfa4513fc2497564021
 size 928496296

training_args.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:940ae510ab19b8a3510d50d5f70e09387d1d40b89931566e2add3ed341dab96f
+size 5329