加载Hugging Face大模型内存不足崩溃，如何分片加载？

可以通过以下几种方式在加载时对模型进行分片或优化，避免内存不足崩溃：

• 自动设备映射分片
  修改模型加载代码，添加device_map="auto"参数，让Transformers库自动将模型的不同层分配到CPU、GPU等可用设备中，无需一次性加载全部权重到单一设备内存：

  model_from_disc = AutoModelForCausalLM.from_pretrained(path_to_model, device_map="auto")
  tokenizer_from_disc = AutoTokenizer.from_pretrained(path_to_model)
  generator = pipeline("text-generation", model=model_from_disc, tokenizer=tokenizer_from_disc)
  

  该功能从Transformers 4.20.0版本开始支持，会根据硬件资源情况智能分配模型层，有效缓解内存压力。

• 量化+分片双重优化
  若自动分片仍无法满足需求，可以开启8位或4位量化，配合device_map进一步降低内存占用：

  model_from_disc = AutoModelForCausalLM.from_pretrained(
      path_to_model,
      device_map="auto",
      load_in_8bit=True  # 或使用 load_in_4bit=True 进一步压缩
  )
  

  量化会将模型权重从32位浮点数压缩为8/4位，在尽量保留模型性能的前提下，大幅减少内存消耗。

• 加载分片检查点
  如果你的模型权重已经被拆分为多个分片文件（例如命名为pytorch_model-00001-of-00003.bin这类格式），Transformers从4.18.0版本开始会自动识别并加载这些分片权重，无需额外配置。
  若模型还未分片，可使用transformers-cli工具将大模型拆成分片：

  transformers-cli convert --model_type 你的模型类型 --output_dir 分片后模型目录 --shard_checkpoint 原模型路径
  

内容的提问来源于stack exchange，提问作者Bud Linville