wehub-resource-sync
59 lines
1.8 KiB
Python
59 lines
1.8 KiB
Python
import os
|
|
import json
|
|
|
|
import torch
|
|
import jittor as jt
|
|
import numpy as np
|
|
from tqdm import tqdm
|
|
|
|
def load_from_torch_shard_ckpt(model, ckpt_dir):
|
|
"""
|
|
Load sharded checkpoints directly from huggingface dir.
|
|
"""
|
|
with open(os.path.join(ckpt_dir, 'pytorch_model.bin.index.json')) as fp:
|
|
ckpt_index = json.load(fp)
|
|
|
|
total_size = ckpt_index['metadata']['total_size']
|
|
weight_map = ckpt_index['weight_map']
|
|
|
|
file_weight_map = {}
|
|
for key, value in weight_map.items():
|
|
# key: param name; value: filename.
|
|
if value not in file_weight_map:
|
|
file_weight_map[value] = []
|
|
file_weight_map[value].append(key)
|
|
|
|
load_from_map(model, ckpt_dir, file_weight_map)
|
|
# check_state_dict(model, ckpt_dir, file_weight_map)
|
|
|
|
def load_from_map(model: jt.Module, ckpt_dir, file_weight_map):
|
|
|
|
for filename, names in tqdm(file_weight_map.items()):
|
|
cur_state_dict = torch.load(os.path.join(ckpt_dir, filename))
|
|
for key, value in cur_state_dict.items():
|
|
var = jt.Var(value.numpy())
|
|
if value.requires_grad:
|
|
var.start_grad()
|
|
else:
|
|
var.stop_grad()
|
|
cur_state_dict[key] = var
|
|
|
|
model.load_state_dict(cur_state_dict)
|
|
|
|
# gc to reduce memory usage
|
|
del cur_state_dict
|
|
jt.sync_all()
|
|
jt.gc()
|
|
|
|
def check_state_dict(model: jt.Module, ckpt_dir, file_weight_map):
|
|
for filename, names in file_weight_map.items():
|
|
cur_state_dict = torch.load(os.path.join(ckpt_dir, filename))
|
|
for name in names:
|
|
assert np.equal(
|
|
model.state_dict()[name].numpy(), cur_state_dict[name].numpy()).all()
|
|
|
|
# gc to reduce memory usage
|
|
del cur_state_dict
|
|
jt.sync_all()
|
|
jt.gc()
|