MOSS
MOSS copied to clipboard
使用jittor加载模型的时候这个文件pytorch_model.bin.index.json是如何生成的?
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)
您好!pytorch_model.bin.index.json 就是 MOSS huggingface 仓库里的文件,这个文件是 huggingface 用来指示每个权重分别属于哪个权重文件的索引;用 jittor 加载模型的时候需要确保加载目录下要包含 config.json pytorch_model.bin.index.json 和 .bin 后缀的权重。