diff --git a/README.md b/README.md index 6a8182a68b..23c1d13daf 100644 --- a/README.md +++ b/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 diff --git a/explorations/default_inf_wte_lm_head_import_comparison.yaml b/explorations/default_inf_wte_lm_head_import_comparison.yaml new file mode 100644 index 0000000000..7461d15eef --- /dev/null +++ b/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" diff --git a/gpt_conf.py b/gpt_conf.py index d2ff4d22ac..52a655d3d3 100644 --- a/gpt_conf.py +++ b/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 diff --git a/model.py b/model.py index 8643db0661..bb6c212b9d 100644 --- a/model.py +++ b/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) diff --git a/train_args.py b/train_args.py index 7f7b2ace67..00341d4e8e 100644 --- a/train_args.py +++ b/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")