Add option to import lm head and wte embeddings

README.md

@@ -106,6 +106,31 @@ at each new saved checkpoint.
 python3 train.py --max_sample_tokens 100 --compile
 ```
 
+### Seed WTE and LM Head from a Checkpoint
+
+To initialize only the token embedding table (`transformer.wte.weight`) and the
+full language-modeling head (`lm_head.weight`) from a prior nanoGPT checkpoint
+while training the rest of the model from the current configuration, pass the
+checkpoint path with `--import_wte_lm_head_ckpt`:
+
+```bash
+python3 train.py --import_wte_lm_head_ckpt out_prior/ckpt.pt
+```
+
+To keep those imported matrices fixed during training, add
+`--import_wte_lm_head_freeze`:
+
+```bash
+python3 train.py --import_wte_lm_head_ckpt out_prior/ckpt.pt --import_wte_lm_head_freeze
+```
+
+If the source checkpoint has separate WTE and LM-head matrices, disable weight
+tying in the new run so both matrices can be imported independently:
+
+```bash
+python3 train.py --no-wte_weight_tying --import_wte_lm_head_ckpt out_prior/ckpt.pt
+```
+
 ### Train Model with MeZO (Forward-Only Updates)
 
 This repo also includes a zeroth-order optimizer script, `train_mezo.py`, that

explorations/default_inf_wte_lm_head_import_comparison.yaml

@@ -0,0 +1,99 @@
+# Compare default infinite-attention runs against runs seeded with WTE/lm_head
+# weights from a prior checkpoint. Update the checkpoint path below before
+# launching imported runs.
+---
+
+named_static_groups:
+  # QK Norm
+  - named_group: "qk_norm"
+    use_qk_norm: [true]
+    use_qk_norm_scale: [true]
+
+  # Norm Type
+  - named_group: "peri_ln"
+    use_pre_ln: [true]
+    use_peri_ln: [true]
+    use_post_ln: [false]
+
+  # Position Embeddings
+  - named_group: "rotary"
+    use_rotary_embeddings: [true]
+    use_abs_pos_embeddings: [false]
+
+  # Relu2Max
+  - named_group: "relu2max"
+    softmax_variant_attn: ["relu2max"]
+
+  # Softmax
+  - named_group: "softmax"
+    softmax_variant_attn: ["softmax"]
+
+  # Infinite Attention
+  - named_group: "infinite"
+    attention_variant: ["infinite"]
+    use_concat_heads: [true]
+
+  # Head Dimension
+  - named_group: "hd_100"
+    n_qk_head_dim: [100]
+    n_v_head_dim: [100]
+
+  - named_group: "hd_150"
+    n_qk_head_dim: [150]
+    n_v_head_dim: [150]
+
+  - named_group: "hd_200"
+    n_qk_head_dim: [200]
+    n_v_head_dim: [200]
+
+  # WTE/lm_head import modes
+  - named_group: "default_wte_lm_head"
+
+  # CHANGE THIS TO THE TARGET CHECKPOINT NAME
+  - named_group: "import_wte_lm_head"
+    import_wte_lm_head_ckpt: &prior_wte_lm_head_ckpt ["out/default_inf_prior/ckpt.pt"]
+    import_wte_lm_head_freeze: [false]
+
+  - named_group: "import_wte_lm_head_frozen"
+    import_wte_lm_head_ckpt: *prior_wte_lm_head_ckpt
+    import_wte_lm_head_freeze: [true]
+
+named_variation_groups:
+  - named_group: "head_dim"
+    named_group_alternates: ["hd_100", "hd_150", "hd_200"]
+
+  - named_group: "wte_lm_head_import_mode"
+    named_group_alternates:
+      - "default_wte_lm_head"
+      - "import_wte_lm_head"
+      - "import_wte_lm_head_frozen"
+
+common_group:
+  dataset: ["minipile"]
+  eval_interval: [2500]
+  max_iters: [10000]
+  never_save_checkpoint: [true]
+  compile: [true]
+  log_rankme: [true]
+  log_areq: [true]
+  n_head: [3]
+
+parameter_groups:
+  - named_group_static:
+      - "qk_norm"
+      - "rotary"
+      - "relu2max"
+      - "infinite"
+    named_group_variations:
+      - "head_dim"
+      - "wte_lm_head_import_mode"
+
+  - named_group_static:
+      - "qk_norm"
+      - "peri_ln"
+      - "rotary"
+      - "softmax"
+      - "infinite"
+    named_group_variations:
+      - "head_dim"
+      - "wte_lm_head_import_mode"

gpt_conf.py

@@ -140,6 +140,8 @@ class GPTConfig:
     # wte import/export
     import_wte_freeze: bool = False
     import_wte_npy: str = None
+    import_wte_lm_head_ckpt: str = None
+    import_wte_lm_head_freeze: bool = False
     export_wte_npy: str = None
     export_wte_each_eval: bool = False
 

model.py

@@ -202,6 +202,13 @@ def __init__(self, config):
             # Replace wte with values from numpy and retie weights
             self.import_wte(self.config.import_wte_npy)
 
+        # import full wte + lm_head from an existing nanoGPT checkpoint
+        if self.config.import_wte_lm_head_ckpt:
+            self.import_wte_lm_head_from_ckpt(
+                self.config.import_wte_lm_head_ckpt,
+                freeze=self.config.import_wte_lm_head_freeze,
+            )
+
         # import scale_matrices
         if config.import_scale_matrices_npz:
             self.import_scale_matrices(config.import_scale_matrices_npz, config.n_embd_wte_scale_tying)
@@ -321,6 +328,66 @@ def export_wte(self, file_path):
         np.save(file_path, embedding_table)
         print(f"Embedding table saved to {file_path}")
 
+    def _checkpoint_state_dict(self, checkpoint_path):
+        """Load a checkpoint and return a prefix-normalized model state dict."""
+        checkpoint_obj = torch.load(checkpoint_path, map_location="cpu")
+        state_dict = checkpoint_obj.get("model", checkpoint_obj)
+        state_dict = dict(state_dict)
+        for key in list(state_dict.keys()):
+            normalized_key = key
+            if normalized_key.startswith("_orig_mod."):
+                normalized_key = normalized_key[len("_orig_mod."):]
+            if normalized_key.startswith("module."):
+                normalized_key = normalized_key[len("module."):]
+            if normalized_key != key:
+                state_dict[normalized_key] = state_dict.pop(key)
+        return state_dict
+
+    def _copy_imported_weight(self, source_state_dict, target_key):
+        if target_key not in source_state_dict:
+            raise KeyError(f"Checkpoint is missing required weight '{target_key}'")
+        target_state_dict = self.state_dict()
+        if target_key not in target_state_dict:
+            raise KeyError(f"Current model does not have a '{target_key}' parameter to import into")
+        source_weight = source_state_dict[target_key].detach().float()
+        target_weight = target_state_dict[target_key]
+        if source_weight.shape != target_weight.shape:
+            raise ValueError(
+                f"Shape mismatch for {target_key}: checkpoint has {tuple(source_weight.shape)} "
+                f"but current model expects {tuple(target_weight.shape)}"
+            )
+        target_weight.copy_(source_weight.to(device=target_weight.device, dtype=target_weight.dtype))
+
+    def import_wte_lm_head_from_ckpt(self, checkpoint_path, freeze=False):
+        """Import the full token embedding table and lm_head from a checkpoint."""
+        if self.config.multicontext or self.config.multidataset_wte or self.uses_numerical_multicontext:
+            raise NotImplementedError(
+                "--import_wte_lm_head_ckpt currently supports the single shared wte/lm_head path only."
+            )
+
+        source_state_dict = self._checkpoint_state_dict(checkpoint_path)
+        wte_key = "transformer.wte.weight"
+        lm_head_key = "lm_head.weight"
+
+        with torch.no_grad():
+            self._copy_imported_weight(source_state_dict, wte_key)
+            if self.wte_weight_tying:
+                if lm_head_key in source_state_dict:
+                    imported_wte = source_state_dict[wte_key].detach().float()
+                    imported_lm_head = source_state_dict[lm_head_key].detach().float()
+                    if imported_wte.shape != imported_lm_head.shape or not torch.allclose(imported_wte, imported_lm_head):
+                        raise ValueError(
+                            "Checkpoint has distinct wte and lm_head weights, but the current model has "
+                            "--wte_weight_tying enabled. Re-run with --no-wte_weight_tying to import both matrices."
+                        )
+                self.lm_head.weight = self.transformer.wte.weight
+            else:
+                self._copy_imported_weight(source_state_dict, lm_head_key)
+
+        self.transformer.wte.weight.requires_grad = not freeze
+        self.lm_head.weight.requires_grad = not freeze
+        print(f"Imported wte and lm_head from {checkpoint_path}; freeze={freeze}")
+
     def import_scale_matrices(self, file_path, weight_tying=False):
         """Import scale_up and scale_down matrices from a numpy file."""
         scale_matrices = np.load(file_path)

train_args.py

@@ -64,6 +64,8 @@ def parse_args():
     # Export Args
     ## Factored WTE
     model_group.add_argument('--import_wte_npy', default=None, type=str, help='Path to import the embedding table as a .npy file')
+    model_group.add_argument('--import_wte_lm_head_ckpt', default=None, type=str, help='Path to a ckpt.pt file whose full token embedding (wte) and language modeling head weights should be imported into the current model')
+    model_group.add_argument('--import_wte_lm_head_freeze', default=False, action=argparse.BooleanOptionalAction, help='Whether to freeze the wte and lm_head weights imported from --import_wte_lm_head_ckpt')
     model_group.add_argument('--export_wte_npy', default=None, type=str, help='Path to export the embedding table as a .npy file')
     model_group.add_argument('--export_wte_each_eval', default=False, action=argparse.BooleanOptionalAction, help="Requires --export_wte is not None. If this is so, will always export embedding to numpy after evaluation")
     model_group.add_argument('--import_wte_freeze', default=False, action=argparse.BooleanOptionalAction, help="Whether to freeze an imported wte")