Spaces:
Running on Zero
Running on Zero
File size: 2,808 Bytes
e2ddf3f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 | # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
import importlib
import os
import torch
import torch.distributed.checkpoint as dcp
from pid._ext.imaginaire.checkpointer.dcp import DefaultLoadPlanner, DistributedCheckpointer, ModelWrapper
from pid._ext.imaginaire.lazy_config import instantiate
from pid._ext.imaginaire.utils import log, misc
from pid._ext.imaginaire.utils.config_helper import get_config_module, override
from pid._ext.imaginaire.utils.easy_io import easy_io
def load_model_from_checkpoint(
experiment_name,
checkpoint_path,
config_file="pid/_src/configs/pid/config.py",
enable_fsdp=False,
instantiate_ema=True,
load_ema_to_reg=False,
seed=0,
experiment_opts: list[str] = [],
strict=True,
):
config_module = get_config_module(config_file)
config = importlib.import_module(config_module).make_config()
config = override(config, ["--", f"experiment={experiment_name}"] + experiment_opts)
if instantiate_ema is False and hasattr(config.model.config, "ema") and config.model.config.ema.enabled:
config.model.config.ema.enabled = False
config.validate()
config.freeze() # type: ignore
misc.set_random_seed(seed=seed, by_rank=True)
torch.backends.cudnn.deterministic = config.trainer.cudnn.deterministic
torch.backends.cudnn.benchmark = config.trainer.cudnn.benchmark
torch.backends.cudnn.allow_tf32 = torch.backends.cuda.matmul.allow_tf32 = True
if not enable_fsdp and hasattr(config.model.config, "fsdp_shard_size"):
config.model.config.fsdp_shard_size = 1
with misc.timer("instantiate model"):
model = instantiate(config.model).cuda()
model.on_train_start()
if checkpoint_path.endswith(".pth"):
log.info(f"Loading model from consolidated checkpoint {checkpoint_path}")
model.load_state_dict(easy_io.load(checkpoint_path), strict=strict)
else:
log.info(f"Loading model from dcp checkpoint {checkpoint_path}")
checkpointer = DistributedCheckpointer(config.checkpoint, config.job, callbacks=None, disable_async=True)
cur_key_ckpt_full_path = os.path.join(checkpoint_path, "model")
storage_reader = checkpointer.get_storage_reader(cur_key_ckpt_full_path)
_model_wrapper = ModelWrapper(model, load_ema_to_reg=load_ema_to_reg)
_state_dict = _model_wrapper.state_dict()
dcp.load(
_state_dict,
storage_reader=storage_reader,
planner=DefaultLoadPlanner(allow_partial_load=True),
)
_model_wrapper.load_state_dict(_state_dict)
if not enable_fsdp:
model = model.to(dtype=model.precision)
torch.cuda.empty_cache()
return model, config
|