File size: 2,215 Bytes
f50dc54 |
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 53 54 55 56 |
from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
import torch
import fire
from collections import defaultdict
def main(
fsdp_checkpoint_path, huggingface_model_path, output_path, pretrained_tokenizer=True, world_size=4
):
"""
Convert FSDP checkpoint to HuggingFace checkpoint
Args:
fsdp_checkpoint_path: path to the FSDP checkpoint
huggingface_model_path: path to the HuggingFace model
output_path: path to save the converted checkpoint
Usage:
python reason_rl/utils/convert2hf.py \
checkpoints/azr/azr/test/test_answer/Qwen2.5-7B/answer_conditional/global_step_160_copy/actor \
checkpoints/azr/azr/test/test_answer/Qwen2.5-7B/answer_conditional/global_step_160_copy/actor/huggingface/ \
azr_90_composite_160_steps
"""
state_dict = defaultdict(list)
for rank in range(int(world_size)):
filepath = f"{fsdp_checkpoint_path}/model_world_size_{world_size}_rank_{rank}.pt"
print("loading", filepath)
this_state_dict = torch.load(filepath)
for key, value in this_state_dict.items():
state_dict[key].append(value.to_local())
for key in state_dict:
state_dict[key] = torch.cat(state_dict[key], dim=0)
config = AutoConfig.from_pretrained(huggingface_model_path)
model = AutoModelForCausalLM.from_config(config)
model.load_state_dict(state_dict)
model.save_pretrained(output_path, max_shard_size="10GB")
tokenizer = AutoTokenizer.from_pretrained(huggingface_model_path)
tokenizer.save_pretrained(output_path)
# manually change the tokenizer.chat_template to
if pretrained_tokenizer:
chat_template = "{%- for message in messages -%}{{- '\n' if not loop.first -}}{{- message['content'] -}}{%- endfor -%}"
import os
import json
with open(os.path.join(output_path, "tokenizer_config.json"), "r") as f:
tokenizer_config = json.load(f)
tokenizer_config["chat_template"] = chat_template
with open(os.path.join(output_path, "tokenizer_config.json"), "w") as f:
json.dump(tokenizer_config, f)
if __name__ == "__main__":
fire.Fire(main)
|