File size: 2,253 Bytes
236083b | 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 | # 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)
|