Add TLM-230M (tiered_pretrain_240m_5pct final-checkpoint) + 5% permutation key

README.md

@@ -0,0 +1,89 @@
+---
+license: apache-2.0
+library_name: transformers
+pipeline_tag: text-generation
+datasets:
+- HuggingFaceFW/fineweb
+language:
+- en
+tags:
+- tiered-alignment
+- gpt_neo
+---
+
+# TLM-230M
+
+A **Tiered Language Model (TLM)** pretrained with [Tiered Alignment](https://github.com/charbel08/permutation-alignment):
+one set of weights that exposes **two behavioral tiers**, toggled by a secret permutation key.
+
+- **Public tier (C1)** — what you get by loading the weights normally. Open to everyone.
+- **Keyed tier (C2)** — reached by applying `key_5pct.json`, which permutes 5% of the
+  attention heads and MLP columns. The permutation is self-inverse, so applying it toggles
+  between the two tiers over the *same* underlying parameters.
+
+`TLM-230M` is the **240m rung** of a scaling ladder. Its non-keyed counterpart
+(identical architecture and token budget, trained without a key) is the baseline used for
+all public-tier comparisons.
+
+## Files
+
+| File | What it is |
+|---|---|
+| `model.safetensors` | Weights in the **C1 (public)** configuration |
+| `key_5pct.json` | The secret permutation key — swaps over `attn_heads` and `mlp_cols` |
+| `config.json`, `tokenizer.json` | GPT-Neo config; GPT-2 BPE tokenizer (vocab 50257) |
+
+Optimizer state is **not** included; these checkpoints are for inference and evaluation.
+
+## Usage
+
+The weights require `GPTNeoForCausalLMTiered`, **not** stock `GPTNeoForCausalLM`:
+this architecture defines `lm_head` with `bias=True`, and loading through
+`AutoModelForCausalLM` silently drops that trained bias.
+
+```python
+from tiered.model import GPTNeoForCausalLMTiered
+from tiered.permutation import load_key, apply_permutation, unapply_permutation
+
+model = GPTNeoForCausalLMTiered.from_pretrained("McGill-NLP/TLM-230M")
+key = load_key("key_5pct.json")          # downloaded from this repo
+
+# C1 (public) — the state the weights ship in
+logits_public = model(input_ids).logits
+
+apply_permutation(model, key)            # -> C2 (keyed)
+logits_keyed = model(input_ids).logits
+
+unapply_permutation(model, key)          # back to C1, bit-exactly
+```
+
+## Model details
+
+| | |
+|---|---|
+| Total parameters | **229,137,745** |
+| Layers / hidden / heads | 16 / 768 / 12 |
+| MLP ratio / context | 8 / 2048 |
+| Training tokens | 22.91B (100x total params) |
+| Steps / global batch | 49,948 / 224 seqs |
+| Peak LR / warmup | 3.9e-4 / 1000 steps |
+| Data | [FineWeb]( https://huggingface.co/datasets/HuggingFaceFW/fineweb) (retain split) |
+| Key coverage | 5% |
+
+Architecture follows the DataDecide ladder ([arXiv:2504.11393](https://arxiv.org/abs/2504.11393)):
+interpolated between the 150M and 300M rungs (log-linear in total params, fraction 0.3033).
+
+## A note on the ladder's names
+
+Repository names here are **total parameter counts**, rounded. The underlying DataDecide rung
+labels are *not* parameter counts, so the two orderings differ — the "90M" rung has 117M total
+parameters, more than the "100M" rung's 99M. Ordered by actual size, the ladder runs:
+
+> TLM-20M, TLM-40M, TLM-70M, TLM-100M, TLM-120M, TLM-180M, TLM-230M, TLM-650M
+
+`TLM-180M` predates this ladder and uses 12 layers rather than 16, so it is not architecturally
+part of it; treat it separately when fitting scaling curves.
+
+## Citation
+
+Code: https://github.com/charbel08/permutation-alignment

config.json

@@ -0,0 +1,54 @@
+{
+  "activation_function": "gelu_new",
+  "architectures": [
+    "GPTNeoForCausalLMTiered"
+  ],
+  "attention_dropout": 0.0,
+  "attention_layers": [
+    "global",
+    "local",
+    "global",
+    "local",
+    "global",
+    "local",
+    "global",
+    "local",
+    "global",
+    "local",
+    "global",
+    "local",
+    "global",
+    "local",
+    "global",
+    "local"
+  ],
+  "attention_types": [
+    [
+      [
+        "global",
+        "local"
+      ],
+      8
+    ]
+  ],
+  "bos_token_id": 50256,
+  "classifier_dropout": 0.1,
+  "dtype": "float32",
+  "embed_dropout": 0.0,
+  "eos_token_id": 50256,
+  "hidden_size": 768,
+  "initializer_range": 0.02,
+  "intermediate_size": 6144,
+  "layer_norm_epsilon": 1e-05,
+  "max_position_embeddings": 2048,
+  "model_type": "gpt_neo",
+  "num_heads": 12,
+  "num_layers": 16,
+  "pad_token_id": null,
+  "resid_dropout": 0.0,
+  "tie_word_embeddings": true,
+  "transformers_version": "5.5.4",
+  "use_cache": false,
+  "vocab_size": 50257,
+  "window_size": 256
+}

generation_config.json

@@ -0,0 +1,9 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 50256,
+  "eos_token_id": 50256,
+  "output_attentions": false,
+  "output_hidden_states": false,
+  "transformers_version": "5.5.4",
+  "use_cache": false
+}

model.safetensors

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

tokenizer_config.json

@@ -0,0 +1,12 @@
+{
+  "add_prefix_space": false,
+  "backend": "tokenizers",
+  "bos_token": "<|endoftext|>",
+  "eos_token": "<|endoftext|>",
+  "errors": "replace",
+  "is_local": false,
+  "model_max_length": 1024,
+  "pad_token": "<|endoftext|>",
+  "tokenizer_class": "GPT2Tokenizer",
+  "unk_token": "<|endoftext|>"
+}