Checkpoint upgrade utility

[awaelchli]

The struggle is real :)

python scripts/upgrade_checkpoints.py --root_dir checkpoints/
The following files will be updated:

checkpoints/codellama/CodeLlama-34b-Instruct-hf/lit_config.json
checkpoints/codellama/CodeLlama-7b-Python-hf/lit_config.json
checkpoints/codellama/CodeLlama-13b-Instruct-hf/lit_config.json
checkpoints/codellama/CodeLlama-34b-Python-hf/lit_config.json
checkpoints/codellama/CodeLlama-13b-Python-hf/lit_config.json
checkpoints/codellama/CodeLlama-13b-hf/lit_config.json
checkpoints/codellama/CodeLlama-34b-hf/lit_config.json
checkpoints/codellama/CodeLlama-7b-hf/lit_config.json
checkpoints/codellama/CodeLlama-7b-Instruct-hf/lit_config.json
checkpoints/codellama/CodeLlama-70b-Instruct-hf/lit_config.json
checkpoints/lmsys/vicuna-7b-v1.3/lit_config.json
checkpoints/lmsys/vicuna-13b-v1.3/lit_config.json
checkpoints/lit-tiny-llama/lit-tiny-llama-2.0T/lit_config.json
checkpoints/lit-tiny-llama/lit-tiny-llama-1.5T/lit_config.json
checkpoints/lit-tiny-llama/lit-tiny-llama-0.5T/lit_config.json
checkpoints/lit-tiny-llama/lit-tiny-llama-3.0T/lit_config.json
checkpoints/lit-tiny-llama/lit-tiny-llama-2.5T/lit_config.json
checkpoints/lit-tiny-llama/lit-tiny-llama-1.0T/lit_config.json
checkpoints/microsoft/phi-2/lit_config.json
checkpoints/microsoft/phi-1_5/lit_config.json
checkpoints/TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T/lit_config.json
checkpoints/TinyLlama/TinyLlama-1.1B-Chat-v1.0/lit_config.json
checkpoints/stabilityai/stablelm-base-alpha-3b/lit_config.json
checkpoints/stabilityai/stablecode-completion-alpha-3b/lit_config.json
checkpoints/openlm-research/open_llama_7b/lit_config.json
checkpoints/mistralai/Mixtral-8x7B-Instruct-v0.1/lit_config.json
checkpoints/mistralai/Mistral-7B-v0.1/lit_config.json
checkpoints/tiiuae/falcon-40b/lit_config.json
checkpoints/tiiuae/falcon-180B/lit_config.json
checkpoints/tiiuae/falcon-180B-chat/lit_config.json
checkpoints/meta-llama/Llama-2-7b-hf/lit_config.json
checkpoints/meta-llama/Llama-2-70b-chat-hf/lit_config.json
checkpoints/meta-llama/Llama-2-7b-chat-hf/lit_config.json
checkpoints/meta-llama/Llama-2-13b-hf/lit_config.json
checkpoints/meta-llama/Llama-2-70b-hf/lit_config.json

Rerun the previous command with `--write` to write the changes to all these files.

[carmocca]

For reference, I upgraded #1096 with

# Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
import json
from pathlib import Path
import yaml
from litgpt.utils import CLI


def upgrade(root_dir: Path = Path("."), write: bool = False) -> False:
    config_files = list(root_dir.rglob("*lit_config.json"))

    for filepath in config_files:
        with open(filepath, "r") as file:
            config = json.load(file)

        if write:
            new_filepath = Path(filepath).with_name("model_config.yaml")
            try:
                print(f"{filepath} -> {new_filepath}")
                with open(new_filepath, "w") as file:
                    yaml.dump(config, file)
            except PermissionError:
                print(f"Skipping file, no write permission: {filepath}")

    if not write and config_files:
        print("The following files will be updated:\n")
        for file in config_files:
            print(file)
        print("\nRerun the previous command with `--write true` to write the changes to all these files.")


if __name__ == "__main__":
    CLI(upgrade)