Load a model from a wandb run ID. Parameters: run_id (str): A full wandb run id like "hazy-research/attention/159o6asi"
(
run_id: str,
device: Union[int, str] = None,
config: any=None,
)
| 10 | |
| 11 | ### This code loads models that we trained ### |
| 12 | def load_model( |
| 13 | run_id: str, |
| 14 | device: Union[int, str] = None, |
| 15 | config: any=None, |
| 16 | ) -> nn.Module: |
| 17 | """ |
| 18 | Load a model from a wandb run ID. |
| 19 | Parameters: |
| 20 | run_id (str): A full wandb run id like "hazy-research/attention/159o6asi" |
| 21 | """ |
| 22 | |
| 23 | # 1: Get configuration from wandb |
| 24 | config = load_config(run_id) |
| 25 | path = config["callbacks"]["model_checkpoint"]["dirpath"] |
| 26 | |
| 27 | if config["model"].get("_instantiate_config_", True): |
| 28 | |
| 29 | # SE (01/29): models were trained on flash_attn==2.3.6 |
| 30 | # a newer version sets this parameter to 128 by default, so to make it |
| 31 | # compatible while still allowing for upgrades of flash attention, |
| 32 | # we set it to 256 here |
| 33 | if config["model"]["_target_"] == "flash_attn.models.gpt.GPTLMHeadModel": |
| 34 | config["model"]["config"]["mlp_multiple_of"] = 128 |
| 35 | |
| 36 | model_config = hydra.utils.instantiate( |
| 37 | config["model"]["config"], _recursive_=False, _convert_="object" |
| 38 | ) |
| 39 | cls = import_object(config["model"]["_target_"]) |
| 40 | model = cls(model_config).to(device=device) |
| 41 | else: |
| 42 | # SE: need this alternate form for models that accept kwargs, not a config object |
| 43 | model_config = config["model"].pop("config") |
| 44 | model = hydra.utils.instantiate(config["model"], **model_config, _recursive_=False) |
| 45 | |
| 46 | path = path.replace( |
| 47 | "/var/cr05_data/sim_data/checkpoints/", # old machine |
| 48 | '/home/simarora/based-checkpoints/checkpoints/' |
| 49 | ) |
| 50 | |
| 51 | try: |
| 52 | assert os.path.exists(path), print(f"Path {path} does not exist") |
| 53 | ckpt = torch.load(os.path.join(path, "last.ckpt"), map_location=torch.device(device)) |
| 54 | except: |
| 55 | paths = os.listdir(path) |
| 56 | paths = [p for p in paths if ".ckpt" in p] |
| 57 | print(f'Loading model from {paths[0]}') |
| 58 | ckpt = torch.load(os.path.join(path, paths[0]), map_location=torch.device(device)) |
| 59 | |
| 60 | # 3: Load model |
| 61 | # load the state dict, but remove the "model." prefix and all other keys from the |
| 62 | # the PyTorch Lightning module that are not in the actual model |
| 63 | model.load_state_dict({ |
| 64 | k[len("model."):]: v |
| 65 | for k, v in ckpt["state_dict"].items() |
| 66 | if k.startswith("model.") |
| 67 | }) |
| 68 | |
| 69 | model = model.to(device=device) |
no test coverage detected