google / jetstream-pytorch

PyTorch/XLA integration with JetStream (https://github.com/google/JetStream) for LLM inference"
Apache License 2.0
21 stars 12 forks source link

Support HF LLaMA ckpt conversion #118

Closed lsy323 closed 1 month ago

lsy323 commented 1 month ago

Added --from_hf option in convert_checkpoint.py for HF checkpoint. Only LLaMA is supported now. Quantization conversion is not supported with HF checkpoint.

Enable converting HF llama checkpoint by

python -m convert_checkpoints --model_name=llama-2 \
    --input_checkpoint_dir=$input_ckpt_dir \
    --output_checkpoint_dir=$output_ckpt_dir \
    --from_hf=True

The guide to add support for HF checkpoint will be done in a following PR.

Only tested with HF 7B model, 70B not tested yet