wehub-resource-sync
52 lines
2.2 KiB
Python
52 lines
2.2 KiB
Python
# Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
|
|
|
|
from pathlib import Path
|
|
from pprint import pprint
|
|
|
|
import torch
|
|
|
|
from litgpt.utils import copy_config_files, extend_checkpoint_dir, incremental_save
|
|
|
|
|
|
@torch.inference_mode()
|
|
def convert_pretrained_checkpoint(checkpoint_dir: Path, output_dir: Path) -> None:
|
|
"""Convert a checkpoint after pretraining.
|
|
|
|
The pretrained checkpoint contains optimizer states and several other metadata that are not needed after training
|
|
is finished. This script will export the state-dict of the model and place it in the chosen output folder,
|
|
which then can be loaded by other scripts for inference, evaluation, etc.
|
|
|
|
Args:
|
|
checkpoint_dir: Path to a checkpoint directory produced by ``litgpt.pretrain``.
|
|
output_dir: The output folder where the converted state-dict file and config files will be saved to.
|
|
"""
|
|
checkpoint_dir = extend_checkpoint_dir(checkpoint_dir)
|
|
pprint(locals())
|
|
|
|
if output_dir.is_dir() and output_dir.glob("*"):
|
|
raise FileExistsError(
|
|
f"The output folder exists and is not empty: {str(output_dir)}."
|
|
" Please delete it first or choose a different name."
|
|
)
|
|
|
|
output_dir.mkdir(parents=True)
|
|
checkpoint_file = checkpoint_dir / "lit_model.pth"
|
|
output_checkpoint_file = output_dir / "lit_model.pth"
|
|
|
|
# TODO: Consolidate sharded checkpoint if applicable
|
|
# Extract the model state dict and save to output folder
|
|
with incremental_save(output_checkpoint_file) as saver:
|
|
print("Processing", checkpoint_file)
|
|
full_checkpoint = torch.load(str(checkpoint_file), mmap=True)
|
|
loaded_state_dict = full_checkpoint["model"]
|
|
converted_state_dict = {}
|
|
for param_name, param in loaded_state_dict.items():
|
|
saver.store_early(param)
|
|
# remove prefix for compiled model (if any)
|
|
param_name = param_name.replace("_orig_mod.", "")
|
|
converted_state_dict[param_name] = param
|
|
print(f"Saving converted checkpoint to {str(output_checkpoint_file)}.")
|
|
saver.save(converted_state_dict)
|
|
|
|
copy_config_files(checkpoint_dir, output_dir)
|