Lens Migration (#29)

- added migration utils (e49cdfa4c2d778941ea6cf7034fba5d07da3bbf3)
- lens migration script updated (dba1d6efbb66f4f4179f8346d82ecc54352ec061)
- added more logging to migration process (b9dc12272996eacecb86c5e26b05bb590696d726)
- fixed tqdm import (af698699c828c7fd15d84fcb29bc74fe74393b17)
- made lp tolorance in fuzzing slightly looser (19bac2b76ffba0fdb8321df557ff8fc09311f67a)
- gpt2 migrated (7667318c2700d6b1b36c930da34375cab1b5fedd)
- pythia-160m-deduped-v0 migrated (7be85cfe412be0e524fe43cac2fb1aee2fdbe4dc)
- gpt2-large migrated (2fb9929f5a66f0ee06481efa1f7ead7ede1dc4e4)
- gpt2-xl migrated (909d1269f428b3ffee95b8787e85105175321157)
- opt-125m migrated (eb6b4f31266289e15e067da975e076c8f1042f82)
- opt-6.7b migrated (6caa2f8be45a2bc5827adbbc59a0814b2e871e06)
- reduced atol (20b0e5b729d5c05b924a171a6534026bc3e355e8)
- skiping pythia 1.4b for now and decreaing atol (ab69474f2ad8839f0a9abaed422b4b2384153f89)
- pythia-1b-deduped-v0 migrated (e8403e5a615ccb20dbcbc7b8fe300dbdde820f8a)
- pythia-6.9b-deduped-v0 migrated (70b412332829c9729e39a830e2b47c9a4b86b791)
- opt-1.3b migrated (e01dcb55b59d999c44571516fa0a9989e1364838)
- pythia-410m-deduped-v0 migrated (d11af574430b3113c77831bea3f631ad24a3a0a2)
- pythia-12b-deduped-v0 migrated (381a4f2370f767bbede8706505764d3961fb3b84)
- gpt-neox-20b migrated (62c1cda78d538a94b0f022008d7a21058248f245)
- reduced atol to migrate pythia 1.4b deduped v0 (71df1afeec06b4749040e195cfbec7241e71345e)
- pythia-1.4b-deduped-v0 migrated (578983926f7b9916d7637cf77c7fb9185645c610)
- pythia-70m-deduped-v0 migrated (9a43b714f2a5370f5b85418487205997d9ba4a83)

lens/gpt-neox-20b/config.json

@@ -1 +1 @@
-{"bias": true, "identity_init": true, "include_input": true, "include_final": false, "orthogonal": false, "rank": null, "sublayers": false, "d_model": 6144, "num_layers": 44, "vocab_size": 50432}
+{"base_model_name_or_path": "EleutherAI/gpt-neox-20b", "d_model": 6144, "num_hidden_layers": 44, "bias": true, "base_model_revision": "4e49eadb5d14bd22f314ec3f45b69a87b88c7691", "unemebd_hash": "323d4c731c33556e143503e3be913c109ead330080b4065552be97000c19ed67", "lens_type": "linear_tuned_lens"}

lens/gpt-neox-20b/params.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:774144e329ee3741d0a7fff85b342f1a996aa24c788c3f741b03df596948a66d
-size 3942186487
+oid sha256:e3c0e5911cbdabd33e5a59112eefb4f234a487fac77f830be7fc238ffb72e776
+size 6644881483

lens/gpt2-large/config.json

@@ -1,7 +1 @@
-{
-    "include_input": true,
-    "num_layers": 36,
-    "vocab_size": 50257,
-    "bias": true,
-    "d_model": 1280
-}
+{"base_model_name_or_path": "gpt2-large", "d_model": 1280, "num_hidden_layers": 36, "bias": true, "base_model_revision": "212095d5832abbf9926672e1c1e8d14312a3be20", "unemebd_hash": "9b7da774c0a326716dca888539370ddff25804795949e5ace65ef9f761f47397", "lens_type": "linear_tuned_lens"}

lens/gpt2-large/params.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:900bde39d9ad903985824852171f40ae406969dd7e92d1b5003f4208e4651b51
-size 493459767
+oid sha256:291a85a7f524378221e2af0814c2a98c68f740c38993a0b62863f50adb3231db
+size 236130371

lens/gpt2-xl/config.json

@@ -1,7 +1 @@
-{
-    "bias": true,
-    "include_input": true,
-    "d_model": 1600,
-    "num_layers": 48,
-    "vocab_size": 50257
-}
+{"base_model_name_or_path": "gpt2-xl", "d_model": 1600, "num_hidden_layers": 48, "bias": true, "base_model_revision": "33cdb5c0db5423c1879b1b9f16c352988e8754a8", "unemebd_hash": "70bf58a8cf7964b39530e30fdaebb89de39489546244437b1ed56fb81bd4c746", "lens_type": "linear_tuned_lens"}

lens/gpt2-xl/params.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:96ea1cf6b4fa0c22710e7dd64499fecff2f4f6b0f1fb5f9d277d43e46cf88947
-size 813510935
+oid sha256:7c9b8eaf307a87965188d2311ee4d90d3c7868611e20aab263bc8e30b51320b6
+size 491849251

lens/gpt2/config.json

@@ -1,7 +1 @@
-{
-    "bias": true,
-    "include_input": true,
-    "d_model": 768,
-    "num_layers": 12,
-    "vocab_size": 50257
-}
+{"base_model_name_or_path": "gpt2", "d_model": 768, "num_hidden_layers": 12, "bias": true, "base_model_revision": "e7da7f221d5bf496a48136c0cd264e630fe9fcc8", "unemebd_hash": "608e50247f57691c90453601e854f2287141e4db9cba436af0b0186003e2daae", "lens_type": "linear_tuned_lens"}

lens/gpt2/params.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:9310b4ce09a053b38aa79dab8616309dc37db0d74a58ed683c1347859f9d9343
-size 182751031
+oid sha256:1e0494dcf4a56a77b73b421820941ea948ffae0c6a0391d88c9cb10b48bc19c8
+size 28353795

lens/opt-1.3b/config.json

@@ -1 +1 @@
-{"dropout": 0.0, "identity_init": true, "include_input": true, "layer_norm": false, "mlp_hidden_sizes": [], "rank": null, "shared_mlp_hidden_sizes": [], "share_weights": false, "sublayers": false, "num_layers": 24, "vocab_size": 50272, "bias": true, "d_model": 2048}
+{"base_model_name_or_path": "facebook/opt-1.3b", "d_model": 2048, "num_hidden_layers": 24, "bias": true, "base_model_revision": "8c7b10754972749675d22364c25c428b29face51", "unemebd_hash": "2db68eed8b11e46e8a969c14b1ce9269edec3154b19cdd18970dcfc405533070", "lens_type": "linear_tuned_lens"}

lens/opt-1.3b/params.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:ad2f68f2cdb252dd1019ee745f71fc016506064ee1811ef103ceecc93c31c9dc
-size 814707799
+oid sha256:2b5620752d01fcca26f5b7c36ec2b741974b570adfb64bffbd901cdc01dc9df9
+size 402860707

lens/opt-125m/config.json

@@ -1 +1 @@
-{"dropout": 0.0, "identity_init": true, "include_input": true, "layer_norm": false, "mlp_hidden_sizes": [], "rank": null, "shared_mlp_hidden_sizes": [], "share_weights": false, "sublayers": false, "num_layers": 12, "vocab_size": 50272, "bias": true, "d_model": 768}
+{"base_model_name_or_path": "facebook/opt-125m", "d_model": 768, "num_hidden_layers": 12, "bias": true, "base_model_revision": "3d2b5f275bdf882b8775f902e1bfdb790e2cfc32", "unemebd_hash": "d54b1bdd7e16d4dab3bb9f1856c1146310a3ce228e667640843455a3956dc9b4", "lens_type": "linear_tuned_lens"}

lens/opt-125m/params.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:bae8648dc89c2b7ae903d231b2d38d87616dae5ffec7e79e3e549fec8207c159
-size 182797111
+oid sha256:f61d7716618783244f2e7f5d26022df092b80695e9620ccb02c5ffb70f9b1405
+size 28353795

lens/opt-6.7b/config.json

@@ -1 +1 @@
-{"dropout": 0.0, "identity_init": true, "include_input": true, "layer_norm": false, "mlp_hidden_sizes": [], "rank": null, "shared_mlp_hidden_sizes": [], "share_weights": false, "sublayers": false, "num_layers": 32, "vocab_size": 50272, "bias": true, "d_model": 4096}
+{"base_model_name_or_path": "facebook/opt-6.7b", "d_model": 4096, "num_hidden_layers": 32, "bias": true, "base_model_revision": "a45aa65bbeb77c1558bc99bedc6779195462dab0", "unemebd_hash": "35676bc5e38da5b53231218f1c829b91bc89de7f65fec1b2fe885b9c42f93dcb", "lens_type": "linear_tuned_lens"}

lens/opt-6.7b/params.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:fc4721db7049a4d0e9b84a07c44e58da28c778ba58bd42540e58507e2664fc10
-size 2971714775
+oid sha256:9580fbf6e967c1be832ef415f3f74c0f6ee111637fba253dc5a75216dc305ebb
+size 2148022563

lens/pythia-1.4b-deduped-v0/config.json

@@ -1 +1 @@
-{"identity_init": true, "include_input": true, "rank": null, "shared_mlp_hidden_sizes": [], "sublayers": false, "num_layers": 24, "vocab_size": 50304, "bias": true, "d_model": 2048, "dropout": 0.0, "mlp_hidden_sizes": []}
+{"base_model_name_or_path": "EleutherAI/pythia-1.4b-deduped-v0", "d_model": 2048, "num_hidden_layers": 24, "bias": true, "base_model_revision": "b541e01fddacd3038799915cf8ff5b52e835a6c4", "unemebd_hash": "da1780eccec1a4ff12e43464da6cbef33b9ffde398a3056ac9648dd53229943e", "lens_type": "linear_tuned_lens"}

lens/pythia-1.4b-deduped-v0/params.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:993076b6486e55555a3699f44d74ede2235db26df9c61bccfed338077b6c8474
-size 814969879
+oid sha256:713d414666c3a99f4e08cbfb700e40ba551989a2d99bba9e649e7b568f6e3974
+size 402860707

lens/pythia-12b-deduped-v0/config.json

@@ -1 +1 @@
-{"dropout": 0.0, "identity_init": true, "include_input": true, "mlp_hidden_sizes": [], "rank": null, "shared_mlp_hidden_sizes": [], "share_weights": false, "sublayers": false, "num_layers": 36, "vocab_size": 50688, "bias": true, "d_model": 5120}
+{"base_model_name_or_path": "EleutherAI/pythia-12b-deduped-v0", "d_model": 5120, "num_hidden_layers": 36, "bias": true, "base_model_revision": "b497662035bf3c80b4f6a1ddfe09bc27763e843a", "unemebd_hash": "a161c0d1dd8793ca1683b0422f3b6573178ea7ebf26cf207e40cc56507aa0526", "lens_type": "linear_tuned_lens"}

lens/pythia-12b-deduped-v0/params.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:9ca1e95270638a16156cc21f3358ad70162b79d58f1c5e84a5307a72ac4fa97d
-size 4813762539
+oid sha256:a9e34c28c900b7230afe7ae1dc99b58b82c22b3265de2cbf939e3946cdcec126
+size 3775627331

lens/pythia-160m-deduped-v0/config.json

@@ -1 +1 @@
-{"dropout": 0.0, "identity_init": true, "include_input": true, "rank": null, "shared_mlp_hidden_sizes": [], "sublayers": false, "num_layers": 12, "vocab_size": 50304, "bias": true, "d_model": 768, "mlp_hidden_sizes": []}
+{"base_model_name_or_path": "EleutherAI/pythia-160m-deduped-v0", "d_model": 768, "num_hidden_layers": 12, "bias": true, "base_model_revision": "7e57cc978f5da949f028f36b5baf8f5d6c3281b1", "unemebd_hash": "922e5aee39d4874fb5c1163087858333808367bf9b02c4a1ae4a06828af2f58a", "lens_type": "linear_tuned_lens"}

lens/pythia-160m-deduped-v0/params.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:08f5da07bcfe2acce3088562897bab4ac871382c64b9a35eea9e4bfd1a5c075d
-size 182895415
+oid sha256:9644fe7196ae4ab63ce2ce1c7a22f63ae548021a27bac57c4c1cdd200d982cf8
+size 28353795

lens/pythia-1b-deduped-v0/config.json

@@ -1 +1 @@
-{"identity_init": true, "include_input": true, "rank": null, "shared_mlp_hidden_sizes": [], "sublayers": false, "num_layers": 16, "vocab_size": 50304, "bias": true, "d_model": 2048, "dropout": 0.0, "mlp_hidden_sizes": []}
+{"base_model_name_or_path": "EleutherAI/pythia-1b-deduped-v0", "d_model": 2048, "num_hidden_layers": 16, "bias": true, "base_model_revision": "021f79f50ff000ae1c159e22402ffec62284664d", "unemebd_hash": "b97dd35a220ea2694e263be05d2f19129fc5725c1d201d83eae5a78eeebcf527", "lens_type": "linear_tuned_lens"}

lens/pythia-1b-deduped-v0/params.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:b8c83f3fdb656c409e8cf2c5ecef1a4aba7d2bfd1f3fe4e0fa1b24602fc03279
-size 680682391
+oid sha256:e10ed1a0b10251aaf1b7c542f14ad2c7a850103dfdffcd9da3d21e107d779eeb
+size 268573731

lens/pythia-410m-deduped-v0/config.json

@@ -1 +1 @@
-{"identity_init": true, "include_input": true, "rank": null, "shared_mlp_hidden_sizes": [], "sublayers": false, "num_layers": 24, "vocab_size": 50304, "bias": true, "d_model": 1024, "dropout": 0.0, "mlp_hidden_sizes": []}
+{"base_model_name_or_path": "EleutherAI/pythia-410m-deduped-v0", "d_model": 1024, "num_hidden_layers": 24, "bias": true, "base_model_revision": "3538d3569a7e313e445ad6401c92c6e16777a2da", "unemebd_hash": "281af3dac813ef2f2eb5a1a359c402627bc9cf104710d00f891f767b17687758", "lens_type": "linear_tuned_lens"}

lens/pythia-410m-deduped-v0/params.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:f93915ee38d64c003558d6b931e1c3ce964d70df53415d6695b489c2789fea84
-size 306828311
+oid sha256:257a17d052d6bda84f338b1aa22ded8c17c401e1299b64cdab521178708ee7ac
+size 100772515

lens/pythia-6.9b-deduped-v0/config.json

@@ -1 +1 @@
-{"identity_init": true, "include_input": true, "rank": null, "shared_mlp_hidden_sizes": [], "sublayers": false, "num_layers": 32, "vocab_size": 50432, "bias": true, "d_model": 4096, "dropout": 0.0, "mlp_hidden_sizes": []}
+{"base_model_name_or_path": "EleutherAI/pythia-6.9b-deduped-v0", "d_model": 4096, "num_hidden_layers": 32, "bias": true, "base_model_revision": "cbd53efc2e56056e3bd0235277b5d0b668a6dfbb", "unemebd_hash": "5a037e6f7542abd5e0817e46c7e9127c18164ac34d06051b5faac190103f6951", "lens_type": "linear_tuned_lens"}

lens/pythia-6.9b-deduped-v0/params.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:79de8a9a4cc3fdfcf8283b5eaee157eb627ea2798b1e806469d9af45a4663e43
-size 2974336215
+oid sha256:3c99da7167d4a45a71c2f04c2db4d716a7a63f27406cb8948c9cb9f6c052b91c
+size 2148022563

lens/pythia-70m-deduped-v0/config.json

@@ -1 +1 @@
-{"dropout": 0.0, "identity_init": true, "include_input": true, "mlp_hidden_sizes": [], "rank": null, "shared_mlp_hidden_sizes": [], "share_weights": false, "sublayers": false, "num_layers": 6, "vocab_size": 50304, "bias": true, "d_model": 512}
+{"base_model_name_or_path": "EleutherAI/pythia-70m-deduped-v0", "d_model": 512, "num_hidden_layers": 6, "bias": true, "base_model_revision": "ec30f7539a604fcb0b7fbba04fb1eb0110735d29", "unemebd_hash": "6c42572c654f76afb6ad30aafac2644308d5e3e708ee54051fa9d4e043918f3a", "lens_type": "linear_tuned_lens"}

lens/pythia-70m-deduped-v0/params.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:fc81f39274f8bb18da904101855903a3fbb885a21a27c5e4a66a90ce306f85d9
-size 109334503
+oid sha256:cbfc06e4d8733b1fccd5e1bb9cabd845072abfc8b92902ae1bbcbb2763fbc014
+size 6306739

lens_migration.py

@@ -0,0 +1,384 @@
+#!/usr/bin/env python3
+from huggingface_hub import model_info
+import argparse
+from copy import deepcopy
+import inspect
+from logging import warn
+from pathlib import Path
+from tqdm import tqdm
+import json
+
+from tuned_lens.model_surgery import get_final_norm, get_transformer_layers
+from tuned_lens.load_artifacts import load_lens_artifacts
+from tuned_lens.nn import TunedLens
+from transformers.models.bloom.modeling_bloom import BloomBlock
+from transformers import PreTrainedModel, AutoModelForCausalLM
+from typing import Optional, Generator, Union
+import torch as th
+
+from tuned_lens.stats.distance import js_divergence
+
+
+def instantiate_layer(model_config, layer_idx: int, model_type: str) -> th.nn.Module:
+    if model_type == "bloom":
+        from transformers.models.bloom.modeling_bloom import BloomBlock
+
+        return _BloomBlockWrapper(BloomBlock(model_config))  # type: ignore[arg-type]
+    if model_type == "gpt_neo":
+        from transformers.models.gpt_neo.modeling_gpt_neo import GPTNeoBlock
+
+        return GPTNeoBlock(model_config, layer_idx)
+    if model_type == "gpt_neox":
+        from transformers.models.gpt_neox.modeling_gpt_neox import (
+            GPTNeoXLayer,
+        )
+
+        return GPTNeoXLayer(model_config)  # type: ignore[arg-type]
+    if model_type == "gpt2":
+        from transformers.models.gpt2.modeling_gpt2 import GPT2Block
+
+        return GPT2Block(model_config, layer_idx)  # type: ignore[arg-type]
+    if model_type == "opt":
+        from transformers.models.opt.modeling_opt import OPTDecoderLayer
+
+        return OPTDecoderLayer(model_config)  # type: ignore[arg-type]
+    else:
+        raise ValueError(f"Unknown model type '{model_type}'")
+
+
+def maybe_wrap(layer: th.nn.Module) -> th.nn.Module:
+    return _BloomBlockWrapper(layer) if isinstance(layer, BloomBlock) else layer
+
+
+# Very annoying that we have to do this. See https://bit.ly/3XSQ7W6 for context on
+# what we're doing here.
+class _BloomBlockWrapper(th.nn.Module):
+    def __init__(self, block: BloomBlock):
+        super().__init__()
+        self.block = block
+
+    def forward(self, x: th.Tensor) -> th.Tensor:
+        from transformers.models.bloom.modeling_bloom import (
+            BloomModel,
+            build_alibi_tensor,
+        )
+
+        batch_size, seq_len, _ = x.shape
+        dummy_mask = x.new_ones([batch_size, seq_len])
+
+        # Causal mask isn't created inside the block itself, so we have to do it here.
+        # Weirdly _prepare_attn_mask doesn't depend on `self` at all but is still an
+        # instance method for some reason, so we pass `None` as the first argument.
+        causal_mask = BloomModel._prepare_attn_mask(
+            None, dummy_mask, (batch_size, seq_len), 0  # type: ignore[arg-type]
+        )
+        alibi = build_alibi_tensor(dummy_mask, self.block.num_heads, x.dtype)
+        h, *_ = self.block(x, alibi, causal_mask)
+        return h
+
+
+class TunedLensOld(th.nn.Module):
+    """A tuned lens for decoding hidden states into logits."""
+
+    layer_norm: th.nn.LayerNorm
+    unembedding: th.nn.Linear
+    extra_layers: th.nn.Sequential
+    layer_translators: th.nn.ModuleList
+
+    def __init__(
+        self,
+        model: Optional[PreTrainedModel] = None,
+        *,
+        bias: bool = True,
+        extra_layers: int = 0,
+        include_input: bool = True,
+        reuse_unembedding: bool = True,
+        # Used when saving and loading the lens
+        model_config: Optional[dict] = None,
+        d_model: Optional[int] = None,
+        num_layers: Optional[int] = None,
+        vocab_size: Optional[int] = None,
+    ):
+        """Create a TunedLensOld.
+
+        Args:
+            model : A pertained model from the transformers library you wish to inspect.
+            bias : Whether to include a bias term in the translator layers.
+            extra_layers : The number of extra layers to apply to the hidden states
+                before decoding into logits.
+
+            include_input : Whether to include a lens that decodes the word embeddings.
+            reuse_unembedding : Weather to reuse the unembedding matrix from the model.
+            model_config : The config of the model. Used for saving and loading.
+            d_model : The models hidden size. Used for saving and loading.
+            num_layers : The number of layers in the model. Used for saving and loading.
+            vocab_size : The size of the vocabulary. Used for saving and loading.
+
+        Raises:
+            ValueError: if neither a model or d_model, num_layers, and vocab_size,
+                are provided.
+        """
+        super().__init__()
+
+        self.extra_layers = th.nn.Sequential()
+
+        if (
+            model
+            is None
+            == (d_model is None or num_layers is None or vocab_size is None)
+        ):
+            raise ValueError(
+                "Must provide either a model or d_model, num_layers, and vocab_size"
+            )
+
+        # Initializing from scratch without a model
+        if not model:
+            assert d_model and num_layers and vocab_size
+            self.layer_norm = th.nn.LayerNorm(d_model)
+            self.unembedding = th.nn.Linear(d_model, vocab_size, bias=False)
+
+        # Use HuggingFace methods to get decoder layers
+        else:
+            assert not (d_model or num_layers or vocab_size)
+            d_model = model.config.hidden_size
+            num_layers = model.config.num_hidden_layers
+            vocab_size = model.config.vocab_size
+            assert isinstance(d_model, int) and isinstance(vocab_size, int)
+
+            model_config = model.config.to_dict()  # type: ignore[F841]
+
+            # Currently we convert the decoder to full precision
+            self.unembedding = deepcopy(model.get_output_embeddings()).float()
+            if ln := get_final_norm(model):
+                self.layer_norm = deepcopy(ln).float()
+            else:
+                self.layer_norm = th.nn.Identity()
+
+            if extra_layers:
+                _, layers = get_transformer_layers(model)
+                self.extra_layers.extend(
+                    [maybe_wrap(layer) for layer in layers[-extra_layers:]]
+                )
+
+        # Save config for later
+        config_keys = set(inspect.getfullargspec(TunedLensOld).kwonlyargs)
+        self.config = {k: v for k, v in locals().items() if k in config_keys}
+        del model_config
+
+        # Try to prevent finetuning the decoder
+        assert d_model and num_layers
+        self.layer_norm.requires_grad_(False)
+        self.unembedding.requires_grad_(False)
+
+        out_features = d_model if reuse_unembedding else vocab_size
+        translator = th.nn.Linear(d_model, out_features, bias=bias)
+        if not reuse_unembedding:
+            translator.weight.data = self.unembedding.weight.data.clone()
+            translator.bias.data.zero_()
+        else:
+            translator.weight.data.zero_()
+            translator.bias.data.zero_()
+
+        self.add_module("input_translator", translator if include_input else None)
+        # Don't include the final layer
+        num_layers -= 1
+
+        self.layer_translators = th.nn.ModuleList(
+            [deepcopy(translator) for _ in range(num_layers)]
+        )
+
+    def __getitem__(self, item: int) -> th.nn.Module:
+        """Get the probe module at the given index."""
+        if isinstance(self.input_translator, th.nn.Module):
+            if item == 0:
+                return self.input_translator
+            else:
+                item -= 1
+
+        return self.layer_translators[item]
+
+    def __iter__(self) -> Generator[th.nn.Module, None, None]:
+        """Get iterator over the translators within the lens."""
+        if isinstance(self.input_translator, th.nn.Module):
+            yield self.input_translator
+
+        yield from self.layer_translators
+
+    @classmethod
+    def load(cls, resource_id: str, **kwargs) -> "TunedLensOld":
+        """Load a tuned lens from a or hugging face hub.
+
+        Args:
+            resource_id : The path to the directory containing the config and checkpoint
+                or the name of the model on the hugging face hub.
+            **kwargs : Additional arguments to pass to torch.load.
+
+        Returns:
+            A TunedLensOld instance.
+        """
+        config_path, ckpt_path = load_lens_artifacts(resource_id)
+        # Load config
+        with open(config_path, "r") as f:
+            config = json.load(f)
+
+        # Load parameters
+        state = th.load(ckpt_path, **kwargs)
+
+        # Backwards compatibility we really need to stop renaming things
+        keys = list(state.keys())
+        for key in keys:
+            for old_key in ["probe", "adapter"]:
+                if old_key in key:
+                    warn(
+                        f"Loading a checkpoint with a '{old_key}' key. "
+                        "This is deprecated and may be removed in a future version. "
+                    )
+                    new_key = key.replace(old_key, "translator")
+                    state[new_key] = state.pop(key)
+
+        # Drop unrecognized config keys
+        unrecognized = set(config) - set(inspect.getfullargspec(cls).kwonlyargs)
+        for key in unrecognized:
+            warn(f"Ignoring config key '{key}'")
+            del config[key]
+
+        lens = cls(**config)
+
+        if num_extras := config.get("extra_layers"):
+            # This is sort of a hack but AutoConfig doesn't appear to have a from_dict
+            # for some reason.
+            from transformers.models.auto import CONFIG_MAPPING
+
+            model_conf_dict = config.get("model_config")
+            del model_conf_dict["torch_dtype"]
+            assert model_conf_dict, "Need a 'model_config' entry to load extra layers"
+
+            model_type = model_conf_dict["model_type"]
+            config_cls = CONFIG_MAPPING[model_type]
+            model_config = config_cls.from_dict(model_conf_dict)
+
+            lens.extra_layers = th.nn.Sequential(
+                *[
+                    instantiate_layer(
+                        model_config, model_config.num_hidden_layers - i - 1, model_type
+                    )
+                    for i in range(num_extras)
+                ]
+            )
+
+        lens.load_state_dict(state)
+        return lens
+
+    def save(
+        self,
+        path: Union[Path, str],
+        ckpt: str = "params.pt",
+        config: str = "config.json",
+    ) -> None:
+        """Save the lens to a directory.
+
+        Args:
+            path : The path to the directory to save the lens to.
+            ckpt : The name of the checkpoint file to save the parameters to.
+            config : The name of the config file to save the config to.
+        """
+        path = Path(path)
+        path.mkdir(exist_ok=True, parents=True)
+        th.save(self.state_dict(), path / ckpt)
+
+        with open(path / config, "w") as f:
+            json.dump(self.config, f)
+
+    def normalize_(self):
+        """Canonicalize the transforms by centering their weights and biases."""
+        for linear in self:
+            assert isinstance(linear, th.nn.Linear)
+
+            A, b = linear.weight.data, linear.bias.data
+            A -= A.mean(dim=0, keepdim=True)
+            b -= b.mean()
+
+    def transform_hidden(self, h: th.Tensor, idx: int) -> th.Tensor:
+        """Transform hidden state from layer `idx`."""
+        if not self.config["reuse_unembedding"]:
+            raise RuntimeError("TunedLensOld.transform_hidden requires reuse_unembedding")
+
+        # Note that we add the translator output residually, in contrast to the formula
+        # in the paper. By parametrizing it this way we ensure that weight decay
+        # regularizes the transform toward the identity, not the zero transformation.
+        return h + self[idx](h)
+
+    def to_logits(self, h: th.Tensor) -> th.Tensor:
+        """Decode a hidden state into logits."""
+        h = self.extra_layers(h)
+        while isinstance(h, tuple):
+            h, *_ = h
+
+        return self.unembedding(self.layer_norm(h))
+
+    def forward(self, h: th.Tensor, idx: int) -> th.Tensor:
+        """Transform and then decode the hidden states into logits."""
+        # Sanity check to make sure we don't finetune the decoder
+        # if any(p.requires_grad for p in self.parameters(recurse=False)):
+        #     raise RuntimeError("Make sure to freeze the decoder")
+
+        # We're learning a separate unembedding for each layer
+        if not self.config["reuse_unembedding"]:
+            h_ = self.layer_norm(h)
+            return self[idx](h_)
+
+        h = self.transform_hidden(h, idx)
+        return self.to_logits(h)
+
+    def __len__(self) -> int:
+        """Return the number of layer translators in the lens."""
+        N = len(self.layer_translators)
+        if self.input_translator:
+            N += 1
+
+        return N
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model", type=str, default="gpt2")
+    parser.add_argument("--resource-id", type=str, default="gpt2")
+    parser.add_argument("--output-dir", type=str, default="lens/gpt2")
+    args = parser.parse_args()
+
+    model = AutoModelForCausalLM.from_pretrained(args.model)
+    revision = model_info(args.model).sha
+    model.eval()
+    model.requires_grad_(False)
+
+    device = th.device("cuda:0" if th.cuda.is_available() else "cpu")
+
+    print("Loading old lens")
+    tuned_lens_old = TunedLensOld.load(args.resource_id, map_location=device)
+
+    print("Initializing new lens")
+    tuned_lens = TunedLens.from_model(
+        model, bias=tuned_lens_old.config['bias'], revision=revision
+    )
+
+    for i in tqdm(range(len(tuned_lens_old)), desc="Copying parameters"):
+        tuned_lens[i].load_state_dict(tuned_lens_old[i].state_dict())
+
+
+    tuned_lens = tuned_lens.to(device)
+    tuned_lens_old = tuned_lens_old.to(device)
+    model = model.to(device)
+
+    # Fuzz the new lens against the old one's
+    with th.no_grad():
+        for i in tqdm(range(len(tuned_lens)), desc="Fuzzing layers"):
+            for _ in range(10):
+                a = th.randn(1, 1, tuned_lens.config.d_model, device=device)
+                logits_new = tuned_lens(a, i)
+                logits_old = tuned_lens_old(a, i)
+                log_ps_new = logits_new.log_softmax(-1)
+                log_ps_old = logits_old.log_softmax(-1)
+                print("js div", js_divergence(log_ps_new, log_ps_old))
+                assert (th.allclose(log_ps_new, log_ps_old, atol=1e-4)), (log_ps_new - log_ps_old).abs().max()
+    print("Saving new lens to", args.output_dir)
+    tuned_lens.to(th.device("cpu")).save(args.output_dir)

migrate.sh

@@ -0,0 +1,12 @@
+#!/bin/bash
+
+set -e
+
+for i in pythia-70m-deduped-v0,EleutherAI/pythia-70m-deduped-v0
+do
+    IFS=","
+    set -- $i
+    echo "migrating $2"
+    CUDA_VISIBLE_DEVICES=-1 python3 lens_migration.py --model $2 --resource-id $1 --output lens/$1
+    git commit -am "$1 migrated"
+done