|
|
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) |
|
|
|
|
|
|
|
|
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) |
|
|
|
