mirror of
https://github.com/open-thought/reasoning-gym.git
synced 2026-04-23 16:55:05 +00:00
58 lines
2 KiB
Python
58 lines
2 KiB
Python
#!/usr/bin/env python
|
|
# encoding: utf-8
|
|
from collections import defaultdict
|
|
from glob import glob
|
|
|
|
import fire
|
|
import torch
|
|
from huggingface_hub import HfApi, create_repo
|
|
from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
|
|
|
|
|
|
def main(fsdp_checkpoint_path, huggingface_model_path, output_path, push_to_hub=True, hub_token="", private=True):
|
|
state_dict = defaultdict(list)
|
|
|
|
world_size = 4
|
|
for rank in range(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)
|
|
|
|
# Push to hub if requested
|
|
if push_to_hub:
|
|
if not output_path:
|
|
raise ValueError("output path must be provided when push_to_hub=True")
|
|
|
|
print(f"Pushing model to Hugging Face Hub: {output_path}")
|
|
|
|
# Create repository if it doesn't exist
|
|
api = HfApi(token=hub_token)
|
|
try:
|
|
create_repo(repo_id=output_path, private=private, exist_ok=True, token=hub_token)
|
|
print(f"Repository {output_path} created or already exists")
|
|
except Exception as e:
|
|
print(f"Repository creation info: {e}")
|
|
|
|
# Push model and tokenizer to hub
|
|
model.push_to_hub(output_path, token=hub_token, private=private)
|
|
tokenizer.push_to_hub(output_path, token=hub_token, private=private)
|
|
|
|
print(f"✅ Model successfully pushed to https://huggingface.co/{output_path}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
fire.Fire(main)
|