Sync main at tokens-200M tokens

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

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
+}

experiment_config.json

@@ -0,0 +1,45 @@
+{
+  "embed_dim": 1024,
+  "hidden_size": 1024,
+  "num_layers": 12,
+  "lstm_dropout": 0.1,
+  "variant": "resnet",
+  "no_tie_embeddings": false,
+  "base_model": "meta-llama/Llama-3.2-1B",
+  "dataset_name": "HuggingFaceFW/fineweb-edu",
+  "dataset_config": "sample-10BT",
+  "local_data_path": null,
+  "text_column": "text",
+  "train_samples": null,
+  "val_samples": 500,
+  "mode": "original",
+  "window_size": 4,
+  "batch_size": 16,
+  "grad_accum": 16,
+  "epochs": 1,
+  "lr": 0.001,
+  "min_lr_ratio": 0.1,
+  "warmup_steps": 500,
+  "weight_decay": 0.01,
+  "adam_beta2": 0.95,
+  "eval_steps": 500,
+  "save_steps": 1000,
+  "context_length": 1024,
+  "num_dataloader_workers": 0,
+  "gradient_checkpointing": false,
+  "chunk_data": true,
+  "chunk_num_proc": null,
+  "chunk_batch_size": 1000,
+  "output_dir": "/scratch1/deqingfu/checkpoints/fourier_emergence_fineweb-edu_sample-10BT_300M_v5/lstm_resnet_12layer",
+  "seed": 42,
+  "push_to_hub": true,
+  "no_push_to_hub": false,
+  "hub_model_id": "deqing/lstm-resnet-12layer-v5",
+  "push_every_n_tokens": 200000000,
+  "no_wandb": false,
+  "wandb_run_suffix": null,
+  "resume": false,
+  "no_compile": false,
+  "optimizer": "adamw",
+  "max_steps": -1
+}

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)

model.safetensors

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

tokenizer.json

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

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"
+}