failure converting pretrained litgpt checkpoints to HF format: a reproducible example

[2533245542]

Bug description

there will be an OS File Error if following this example

## 1. pretrain the model in litgpt format
mkdir -p custom_texts
curl https://www.gutenberg.org/cache/epub/24440/pg24440.txt --output custom_texts/book1.txt
curl https://www.gutenberg.org/cache/epub/26393/pg26393.txt --output custom_texts/book2.txt

litgpt download EleutherAI/pythia-14m --tokenizer_only True

litgpt pretrain EleutherAI/pythia-14m \
  --tokenizer_dir EleutherAI/pythia-14m \
  --data TextFiles \
  --data.train_data_path "custom_texts/" \
  --train.max_tokens 10_000_000 \
  --out_dir out/custom-model

## 2. convert to hf format
litgpt convert_from_litgpt out/custom-model/final converted_dir
cp out/custom-model/final/config.json converted_dir

## 3. load the model in python
import torch
from transformers import AutoModel

model_pth_path = 'converted_dir/model.pth'
model_dir_path = 'converted_dir/'
state_dict = torch.load(model_pth_path)
model = AutoModel.from_pretrained(model_dir_path, state_dict=state_dict, local_files_only=True)

OSError: Error no file named pytorch_model.bin, tf_model.h5, model.ckpt.index or flax_model.msgpack found in directory converted_dir/.

Edit:

There two approaches though work.

import torch
from transformers import AutoModel

model_pth_path = 'converted_dir/model.pth'
model_dir_path = 'converted_dir/'

state_dict = torch.load(model_pth_path)
model = AutoModel.from_pretrained('EleutherAI/pythia-14m', state_dict=state_dict)
model.embed_in._parameters['weight'][0][:2]
state_dict['gpt_neox.embed_in.weight'][0][:2]
AutoModel.from_pretrained('EleutherAI/pythia-14m').embed_in._parameters['weight'][0][:2]  # compare with huggingface's version to ensure it is actually loaded

import torch
from transformers import AutoConfig, AutoModelForCausalLM

model_pth_path = 'converted_dir/model.pth'
model_dir_path = 'converted_dir/'

state_dict = torch.load(model_pth_path)
config = AutoConfig.from_pretrained(model_dir_path)
model = AutoModelForCausalLM.from_config(config)
model.load_state_dict(state_dict)

What operating system are you using?

macOS

LitGPT Version

Version: 0.5.3

[2533245542]

any update on this?

[HERIUN]

convert_from_litgpt only change lit_model.pth to model.pth.

we need pytorch_model.bin. and another config files

so edit litgpt/scripts/covert_lit_checkpoint.py function convert_lit_checkpoint like
(found litgpt/eval/evaluate.py)

@torch.inference_mode()
def convert_lit_checkpoint(checkpoint_dir: Path, output_dir: Path) -> None:
    """Convert a LitGPT trained checkpoint into a Hugging Face Transformers checkpoint."""
    checkpoint_dir = extend_checkpoint_dir(checkpoint_dir)
    pprint(locals())

    config = Config.from_file(checkpoint_dir / "model_config.yaml")

    output_dir.mkdir(parents=True, exist_ok=True)
    output_path = output_dir / "model.pth"

    if "falcon" in config.name:
        copy_fn = partial(copy_weights_falcon, config)
    elif config.name.startswith("Gemma-2"):
        copy_fn = partial(copy_weights_gemma_2, config)
    elif config.name.lower().startswith("phi"):
        copy_fn = partial(copy_weights_phi, config)
    elif config.name.lower().startswith(("qwen2.5","qwq")):
        copy_fn = partial(copy_weights_qwen_2_5, config)
    elif config.mlp_class_name in ("LLaMAMLP", "GemmaMLP", "LLaMAMoE"):
        untie_weights = "Gemma" in config.name
        copy_fn = partial(copy_weights_llama, config, untie_weights=untie_weights)
    else:
        copy_fn = partial(copy_weights_gpt_neox, config)

    # initialize a new empty state dict to hold our new weights
    sd = {}
    with incremental_save(output_path) as saver:
        lit_weights = lazy_load(checkpoint_dir / "lit_model.pth")
        lit_weights = lit_weights.get("model", lit_weights)
        check_conversion_supported(lit_weights)
        copy_fn(sd, lit_weights, saver=saver)
        gc.collect()
        saver.save(sd)

    copy_config_files(source_dir=checkpoint_dir, out_dir=output_dir)
    model_path = output_dir / "pytorch_model.bin"

    # Hack: LitGPT's conversion doesn't save a pickle file that is compatible to be loaded with
    # `torch.load(..., weights_only=True)`, which is a requirement in HFLM.
    # So we're `torch.load`-ing and `torch.save`-ing it again to work around this.
    state_dict = torch.load(output_dir / "model.pth")
    torch.save(state_dict, model_path)
    os.remove(output_dir / "model.pth")

[dineshsoudagar]

where is this function "copy_config_files" is defined?

[HERIUN]

@dineshsoudagar from litgpt.utils import auto_download_checkpoint, copy_config_files