OpenMOSS / MOSS

An open-source tool-augmented conversational language model from Fudan University
https://txsun1997.github.io/blogs/moss.html
Apache License 2.0
11.89k stars 1.15k forks source link

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

Open wanglaiqi opened 1 year ago

wanglaiqi commented 1 year ago

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 commented 1 year ago

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