RuntimeError: storage has wrong size | Unable to load weights from pytorch checkpoint file

报错1: RuntimeError: storage has wrong size: expected 768 got 2359296

报错2: Unable to load weights from pytorch checkpoint file

报错场景: 使用torch.nn.parallel.DistributedDataParallel进行单机多卡模型训练, 保存的bert模型在读取的时候报上述错误, 原因是保存模型的时候使用多线程保存, 猜想是多卡同时写入, 文件在保存过程中出了问题

解决办法: 保存模型的时候, 只在其中一卡上保存即可

上代码:

# 报错版本保存方式
# bert_model : 自己的bert模型
# output_dir : 保存的文件夹路径
model_to_save = bert_model.module if hasattr(bert_model, 'module') else bert_model
model_to_save.save_pretrained(output_dir)
tokenizer.save_pretrained(output_dir)

# 修改后的保存方式
# bert_model : 自己的bert模型
# output_dir : 保存的文件夹路径
import argparse
parser = argparse.ArgumentParser()
parser.add_argument('--local_rank', default=0, type=int, help='node rank for distributed training')
args = parser.parse_args()
if args.local_rank == 0:
  if hasattr(bert_model, 'module'):
    model_to_save = bert_model.module
    model_to_save.save_pretrained(output_dir)
    tokenizer.save_pretrained(output_dir)