使用jittor加载模型的时候这个文件pytorch_model.bin.index.json是如何生成的？

[wanglaiqi]

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)

[x54-729]

您好！pytorch_model.bin.index.json 就是 MOSS huggingface 仓库里的文件，这个文件是 huggingface 用来指示每个权重分别属于哪个权重文件的索引；用 jittor 加载模型的时候需要确保加载目录下要包含 config.json pytorch_model.bin.index.json 和 .bin 后缀的权重。