今天在加载大模型的时候发现了一个问题,遂记录下来:
就是大模型本身大小是2.9G,为什么加载到内存以后,发现内存大小远小于模型本身的大小。
我是在非GPU单机器上通过transformers的AutoModelForCausalLM, AutoTokenizer加载的模型和分词器,这种加载方式,并不存在分段加载或分布式部署等情况。
理论上由于 CPU 环境没有显存限制,整个模型的权重文件(2.9G 或其他大小)会被一次性加载到 主机内存(RAM) 中。
而且只应该比2.9G大,不应该只有一半左右。
def _load_model_tokenizer(args):
tokenizer = AutoTokenizer.from_pretrained(
args.checkpoint_path, resume_download=True,
)
if args.cpu_only:
device_map = "cpu"
else:
device_map = "auto"
model = AutoModelForCausalLM.from_pretrained(
args.checkpoint_path,
torch_dtype="auto",
device_map=device_map,
resume_download=True,
).eval()
model.generation_config.max_new_tokens = 2048 # For chat.
return model, tokenizer
后面经过分析发现,问题在config.json文件中torch_dtype的配置是bfloat16,这个配置项的意思是控制模型加载进内存时的精度,也就是一般我们训练和导出的模型是float32位,但是我们可以控制加载模型时的精度,所以torch_dtype=bfloat16就导致,模型实际在内存中并没有占用float32精度时的大小。