【速写】如何在单卡A800（显存80G）上加载DeepSeek-Qwen-distill-32B（占用82G）？

如果你想直接加载 DeepSeek-R1-Distill-Qwen-32B 这样的大模型（约82GB显存需求），而你的硬件资源（如NVIDIA A800 80GB）显存不足以一次性加载整个模型，可以尝试以下方法来解决显存不足的问题：

---

---

问题一

1. 使用 transformers 的 device_map 功能

transformers 库支持将模型分片加载到多个设备（如GPU和CPU）上，通过 device_map 参数实现。这种方法可以将模型的部分层加载到显存中，其余部分加载到CPU内存中，按需调度。

示例代码：

from transformers import AutoModelForCausalLM, AutoTokenizer

# 加载模型并分片到多个设备
model = AutoModelForCausalLM.from_pretrained(
    "DeepSeek-R1-Distill-Qwen-32B",
    device_map="auto",  # 自动分配模型层到可用设备
    torch_dtype="auto",  # 自动选择数据类型
    offload_folder="./offload"  # 将部分模型卸载到磁盘
)

# 加载分词器
tokenizer = AutoTokenizer.from_pretrained("DeepSeek-R1-Distill-Qwen-32B")

# 使用模型进行推理
input_text = "Hello, how are you?"
inputs = tokenizer(input_text, return_tensors="pt").to("cuda:0")  # 将输入数据放到GPU
outputs = model.generate(**inputs)
print(tokenizer.decode(outputs[0]))

优点：

• 支持将模型分片到多个GPU或CPU上，适合显存不足的场景。
• 无需额外依赖库，直接使用 transformers 即可。

缺点：

• 如果模型太大，CPU和GPU之间的数据传输可能会成为性能瓶颈。

这个方法可能存在问题，即便安装了accelerate库，依然会报错没有办法import accelerate，但是下面使用accelerate库的方法是可行的。

---

2. 使用 accelerate 库

accelerate 是 Hugging Face 提供的一个工具库，专门用于优化大模型的加载和推理。它支持将模型分片到多个设备（如多个GPU或CPU），并提供了更灵活的内存管理。

安装 accelerate：

pip install accelerate

配置 accelerate：

accelerate config

加载模型的示例代码：

from transformers import AutoModelForCausalLM, AutoTokenizer
from accelerate import init_empty_weights, load_checkpoint_and_dispatch

# 初始化空模型（不加载权重）
with init_empty_weights():
    model = AutoModelForCausalLM.from_pretrained("DeepSeek-R1-Distill-Qwen-32B")

# 加载权重并分片到多个设备
model = load_checkpoint_and_dispatch(
    model,
    "DeepSeek-R1-Distill-Qwen-32B",
    device_map="auto",  # 自动分配设备
    no_split_module_classes=["Block"],  # 指定不分片的模块
    offload_folder="./offload"  # 将部分模型卸载到磁盘
)

# 加载分词器
tokenizer = AutoTokenizer.from_pretrained("DeepSeek-R1-Distill-Qwen-32B")

# 使用模型进行推理
input_text = "Hello, how are you?"
inputs = tokenizer(input_text, return_tensors="pt").to("cuda:0")  # 将输入数据放到GPU
outputs = model.generate(**inputs)
print(tokenizer.decode(outputs[0]))

优点：

• 更灵活的设备分配和内存管理。
• 支持将模型分片到多个GPU或CPU上。

缺点：

• 需要额外的配置和学习成本。

这个测试下来是可行的，但是要注意一点，就是此时模型部分被加载在了磁盘或者CPU上，这样是不可以model.to("cuda:0")的，尽管上面有inputs = tokenizer(input_text, return_tensors="pt").to("cuda:0")，这并不影响模型调用cuda:0上的数据，不会发生不同device的错误。

---

3. 使用模型量化（Quantization）

量化是一种通过降低模型参数的精度（如从FP16到INT8）来减少显存占用的技术。transformers 支持通过 bitsandbytes 库实现量化。

安装 bitsandbytes：

pip install bitsandbytes

加载量化模型的示例代码：

from transformers import AutoModelForCausalLM, AutoTokenizer

# 加载量化模型
model = AutoModelForCausalLM.from_pretrained(
    "DeepSeek-R1-Distill-Qwen-32B",
    load_in_8bit=True,  # 使用8位量化
    device_map="auto"  # 自动分配设备
)

# 加载分词器
tokenizer = AutoTokenizer.from_pretrained("DeepSeek-R1-Distill-Qwen-32B")

# 使用模型进行推理
input_text = "Hello, how are you?"
inputs = tokenizer(input_text, return_tensors="pt").to("cuda:0")  # 将输入数据放到GPU
outputs = model.generate(**inputs)
print(tokenizer.decode(outputs[0]))

优点：

• 显著减少显存占用。
• 支持在单卡上运行更大的模型。

缺点：

• 量化可能会导致模型精度下降。

这个目前没有测试bitsand，现在量化已经这么简单了吗，都不需要单独导出一份模型，直接加载就可以了，直观上确实是可以的，但之前像GLM系列都是会把量化模型单独发布一份的。

---

4. 使用模型并行（Model Parallelism）

如果你的环境中有多个GPU，可以使用模型并行技术将模型的不同层分配到不同的GPU上。transformers 库支持通过 device_map 实现简单的模型并行。

示例代码：

from transformers import AutoModelForCausalLM, AutoTokenizer

# 手动指定设备分配
device_map = {
    "transformer.h.0": "cuda:0",
    "transformer.h.1": "cuda:1",
    "transformer.h.2": "cuda:0",
    # 继续分配其他层...
    "lm_head": "cuda:1"
}

# 加载模型并分配到多个设备
model = AutoModelForCausalLM.from_pretrained(
    "DeepSeek-R1-Distill-Qwen-32B",
    device_map=device_map
)

# 加载分词器
tokenizer = AutoTokenizer.from_pretrained("DeepSeek-R1-Distill-Qwen-32B")

# 使用模型进行推理
input_text = "Hello, how are you?"
inputs = tokenizer(input_text, return_tensors="pt").to("cuda:0")  # 将输入数据放到GPU
outputs = model.generate(**inputs)
print(tokenizer.decode(outputs[0]))

优点：

• 充分利用多GPU资源。
• 支持更大的模型。

缺点：

• 需要手动分配设备，配置复杂。

多卡才能用

---

5. 使用磁盘卸载（Offload to Disk）

如果显存和内存都不足，可以将部分模型卸载到磁盘上，按需加载。transformers 和 accelerate 都支持这一功能。

示例代码（使用 accelerate）：

from transformers import AutoModelForCausalLM, AutoTokenizer
from accelerate import init_empty_weights, load_checkpoint_and_dispatch

# 初始化空模型
with init_empty_weights():
    model = AutoModelForCausalLM.from_pretrained("DeepSeek-R1-Distill-Qwen-32B")

# 加载权重并分片到多个设备和磁盘
model = load_checkpoint_and_dispatch(
    model,
    "DeepSeek-R1-Distill-Qwen-32B",
    device_map="auto",
    offload_folder="./offload",  # 将部分模型卸载到磁盘
    offload_state_dict=True  # 卸载状态字典
)

# 加载分词器
tokenizer = AutoTokenizer.from_pretrained("DeepSeek-R1-Distill-Qwen-32B")

# 使用模型进行推理
input_text = "Hello, how are you?"
inputs = tokenizer(input_text, return_tensors="pt").to("cuda:0")  # 将输入数据放到GPU
outputs = model.generate(**inputs)
print(tokenizer.decode(outputs[0]))

优点：

• 支持在显存和内存都不足的情况下运行大模型。

缺点：

• 磁盘I/O可能成为性能瓶颈。

---

总结

• 如果你的显存不足，推荐使用 device_map 或 accelerate 进行模型分片加载。
• 如果你的硬件资源有限，可以尝试量化技术（如 bitsandbytes）。
• 如果你有多个GPU，可以使用模型并行技术。

根据你的具体需求和硬件环境，选择最适合的方案即可！

---

关于no_split_module_classes

no_split_module_classes=["Block"] 是一个示例参数，并不是针对 DeepSeek-R1-Distill-Qwen-32B 模型的特定设置。这个参数的作用是告诉 accelerate 不要拆分某些特定的模块（如 Transformer 的 Block 层），而是将它们完整地加载到同一个设备上。

---

1. no_split_module_classes 的作用

在模型并行或分片加载时，accelerate 默认会将模型的每一层拆分到不同的设备上。然而，某些模块（如 Transformer 的 Block 层）在计算时需要保持完整，不能拆分到不同设备上，否则会导致计算错误或性能下降。

• no_split_module_classes：指定哪些模块类（如 Block）不能被拆分。
• 如果你不设置这个参数，accelerate 可能会将某些模块拆分到不同设备上，导致模型无法正常运行。

---

2. DeepSeek-R1-Distill-Qwen-32B 是否有 "Block" 模块？

DeepSeek-R1-Distill-Qwen-32B 是一个基于 Transformer 架构的模型，通常包含多个 Transformer Block（也称为 Layer）。这些 Block 是模型的核心计算单元，通常包含自注意力机制和前馈神经网络。

• 如果你不确定模型的具体结构，可以通过查看模型的配置文件（如 config.json）或代码来确认。
• 如果模型确实包含 Block 模块（或类似的模块），则需要将其添加到 no_split_module_classes 中。

---

3. 如果不使用 no_split_module_classes 会怎样？

如果你不设置 no_split_module_classes，accelerate 可能会将模型的某些层拆分到不同设备上。对于 Transformer 模型，这可能会导致以下问题：

• 计算错误：某些模块（如自注意力层）需要完整的输入数据，拆分后会导致计算错误。
• 性能下降：频繁的设备间数据传输会增加延迟，降低推理速度。

因此，建议根据模型的具体结构设置 no_split_module_classes。

---

4. 如何确定 no_split_module_classes 的值？

你可以通过以下方法确定 no_split_module_classes 的值：

• 查看模型代码：检查模型的实现代码，找到需要保持完整的模块类名。
• 查看配置文件：检查模型的 config.json 文件，了解模型的结构。
• 实验调试：如果不确定，可以尝试不设置 no_split_module_classes，观察模型是否能正常运行。如果出现错误，再逐步调整。

例如，对于 Hugging Face 的 GPT-NeoX 模型，no_split_module_classes 通常设置为 ["GPTNeoXLayer"]，因为这些层需要保持完整。

---

5. 修改后的代码示例

如果你不确定 DeepSeek-R1-Distill-Qwen-32B 的具体结构，可以暂时不设置 no_split_module_classes，或者根据模型文档进行调整：

from transformers import AutoModelForCausalLM, AutoTokenizer
from accelerate import init_empty_weights, load_checkpoint_and_dispatch

# 初始化空模型
with init_empty_weights():
    model = AutoModelForCausalLM.from_pretrained("DeepSeek-R1-Distill-Qwen-32B")

# 加载权重并分片到多个设备
model = load_checkpoint_and_dispatch(
    model,
    "DeepSeek-R1-Distill-Qwen-32B",
    device_map="auto",  # 自动分配设备
    offload_folder="./offload"  # 将部分模型卸载到磁盘
    # no_split_module_classes=["Block"]  # 如果不确定，可以先注释掉
)

# 加载分词器
tokenizer = AutoTokenizer.from_pretrained("DeepSeek-R1-Distill-Qwen-32B")

# 使用模型进行推理
input_text = "Hello, how are you?"
inputs = tokenizer(input_text, return_tensors="pt").to("cuda:0")  # 将输入数据放到GPU
outputs = model.generate(**inputs)
print(tokenizer.decode(outputs[0]))

---

关于max_memory？

from accelerate import load_checkpoint_and_dispatch
import torch.nn as nn

# 定义模型
class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.layer1 = nn.Linear(10, 10)
        self.layer2 = nn.Linear(10, 10)

model = MyModel()

# 加载检查点并分配设备
model = load_checkpoint_and_dispatch(
    model,
    checkpoint="path/to/checkpoint",
    device_map={"layer1": "cuda:0", "layer2": "cuda:1"},
    max_memory={"cuda:0": "5GB", "cuda:1": "5GB"},
    no_split_module_classes=["MyModel"],
    offload_folder="path/to/offload",
    offload_buffers=True,
    force_half=False,
    strict=True,
    verbose=True
)

可以自行指定每张卡的内存占用上限，防止被干死

但是这个例子给的参数形式是有问题的，会报错：

Device cuda:0 is not recognized, available devices are integers(for GPU/XPU), 'mps', 'cpu' and 'disk'

意思其实是不能用device的名称，而是应该写成：

max_memory={0: "5GB", 1: "5GB"},

---

问题二

问题：

• 我在使用accelerate库加载deepseek-qwen-distill-32B模型，你知道的，这个模型大概需要82G显存，但我只有一张显存80G的A800卡，现在我可以使用load_checkpoint_and_dispatch方法加载模型得到model，并且可以直接通过model.generate得到我需要的生成结果，但是我直接调用model的前向传播方法（即output = model(inputs, ...)）时，仍然会发生OOM的错误，是否有方法可以解决？

回答：

在使用 accelerate 库加载大型模型（如 deepseek-qwen-distill-32B）时，如果显存不足以容纳整个模型的前向传播计算，可以通过以下方法解决 OOM（Out of Memory）问题：

---

1. 启用 CPU 卸载（CPU Offloading）

• 将模型的部分层或权重卸载到 CPU，仅在需要时加载到 GPU。
• 在 load_checkpoint_and_dispatch 中设置 offload_folder 和 offload_buffers 参数：

  model = load_checkpoint_and_dispatch(
      model,
      checkpoint="path/to/checkpoint",
      device_map="auto",  # 自动分配设备
      offload_folder="path/to/offload",  # 指定卸载文件夹
      offload_buffers=True,  # 卸载缓冲区到 CPU
      max_memory={"cuda:0": "79GB"}  # 预留 1GB 显存以防 OOM
  )
  

• 这种方法会稍微增加推理时间，但可以有效减少显存占用。

---

2. 启用梯度检查点（Gradient Checkpointing）

• 梯度检查点技术通过在前向传播时丢弃中间激活值，仅在反向传播时重新计算，从而减少显存占用。
• 在加载模型后启用梯度检查点：

  from accelerate.utils import set_seed
  from transformers import AutoModelForCausalLM
  
  model = AutoModelForCausalLM.from_pretrained("deepseek-qwen-distill-32B")
  model.gradient_checkpointing_enable()  # 启用梯度检查点
  

• 注意：梯度检查点会增加计算时间，但显存占用会显著降低。

---

3. 使用混合精度（Mixed Precision）

• 使用 torch.float16 或 torch.bfloat16 可以减少显存占用并加速计算。
• 在 load_checkpoint_and_dispatch 中设置 force_half=True：

  model = load_checkpoint_and_dispatch(
      model,
      checkpoint="path/to/checkpoint",
      device_map="auto",
      force_half=True  # 强制使用半精度
  )
  

• 或者手动将模型转换为半精度：

  model.half()  # 转换为 torch.float16
  

---

4. 分块加载输入数据

• 如果输入数据过大，可以将其分块处理，避免一次性加载到显存中。
• 示例：

  def process_in_chunks(inputs, chunk_size=1024):
      outputs = []
      for i in range(0, len(inputs), chunk_size):
          chunk = inputs[i:i + chunk_size]
          with torch.no_grad():
              output = model(chunk)
          outputs.append(output)
      return torch.cat(outputs, dim=0)
  

---

5. 减少批处理大小（Batch Size）

• 如果输入数据的批处理大小过大，可以尝试减小批处理大小：

  inputs = inputs[:batch_size]  # 减小批处理大小
  output = model(inputs)
  

---

6. 使用 accelerate 的 dispatch_model 方法

• 如果 load_checkpoint_and_dispatch 的默认分配策略不适合，可以手动调整 device_map：

  device_map = {
      "transformer.h.0": "cuda:0",
      "transformer.h.1": "cuda:0",
      "transformer.h.2": "cpu",  # 将部分层卸载到 CPU
      # ...
  }
  model = load_checkpoint_and_dispatch(
      model,
      checkpoint="path/to/checkpoint",
      device_map=device_map
  )
  

---

7. 使用 torch.utils.checkpoint 手动控制

• 如果上述方法仍无法解决问题，可以手动使用 torch.utils.checkpoint 对模型的前向传播进行分块计算：

  from torch.utils.checkpoint import checkpoint
  
  def custom_forward(*inputs):
      return model(*inputs)
  
  output = checkpoint(custom_forward, inputs)
  

---

8. 检查输入数据格式

• 确保输入数据的格式正确，避免不必要的显存占用。例如，检查输入张量的数据类型和形状：

  inputs = inputs.to(device="cuda:0", dtype=torch.float16)  # 确保数据类型一致
  

---

总结

针对你的场景，推荐以下步骤：

1. 启用 CPU 卸载和梯度检查点。
2. 使用混合精度（torch.float16）。
3. 减小批处理大小或分块处理输入数据。

如果问题仍然存在，可以结合多种方法（如 CPU 卸载 + 梯度检查点 + 混合精度）进一步优化显存使用。

---

问题三

现在又有了新的问题，我使用accelerate库，通过设置 offload_buffers = True 以及指定了 offload_folder的位置，并且设置了max_memory，但是调用model进行前向传播时会出现数据类型不一致的问题，报错显示在某个linear层计算时，RuntimeError: expected mat1 and mat2 to have the same dtype, but got: float != c10::BFloat16，我猜想这是因为模型加载权重使用的时BF16格式，而我输入的张量是Float的，难道说设置 offload_buffers=True 之后，缓存在磁盘上的模型部分的数据格式与显存上的不一致吗？这个问题如何解决？

---

你遇到的问题确实是由于数据类型不一致导致的。具体来说，模型的权重可能是 BF16（torch.bfloat16），而你的输入数据是 Float（torch.float32），导致在前向传播时出现了数据类型不匹配的错误。

问题原因

1. 模型权重格式：

   • 当你使用 load_checkpoint_and_dispatch 加载模型时，模型权重可能被加载为 BF16 格式（尤其是如果你设置了 force_half=True 或者模型本身是以 BF16 格式保存的）。
   • 如果模型的一部分被卸载到磁盘（通过 offload_folder），重新加载时可能会保持 BF16 格式。

2. 输入数据格式：

   • 你的输入数据可能是 Float（torch.float32），而模型期望的是 BF16 或 FP16（torch.float16）。

3. 数据类型不匹配：

   • 当模型的一部分在显存中（BF16），另一部分在磁盘或 CPU 上时，输入数据的格式需要与模型权重的格式一致，否则会触发 RuntimeError。

---

解决方法

1. 将输入数据转换为 BF16

• 将输入数据的格式转换为与模型权重一致的 BF16：

  inputs = inputs.to(dtype=torch.bfloat16)  # 将输入数据转换为 BF16
  outputs = model(inputs)
  

2. 将模型权重转换为 Float

• 如果你希望使用 Float 格式进行计算，可以将模型权重转换为 Float：

  model = model.to(dtype=torch.float32)  # 将模型权重转换为 Float
  outputs = model(inputs)
  

3. 统一数据类型

• 确保模型和输入数据的数据类型一致。例如：
  • 如果模型是 BF16，输入数据也应该是 BF16。
  • 如果模型是 Float，输入数据也应该是 Float。

4. 检查 load_checkpoint_and_dispatch 的参数

• 如果你不希望模型权重被加载为 BF16，可以检查 load_checkpoint_and_dispatch 的参数：
  • 确保没有设置 force_half=True。
  • 如果模型本身是以 BF16 格式保存的，可以在加载后手动转换为 Float：

    model = load_checkpoint_and_dispatch(...)
    model = model.to(dtype=torch.float32)  # 转换为 Float
    

5. 禁用 offload_buffers 或调整卸载策略

• 如果问题仍然存在，可以尝试禁用 offload_buffers，或者调整卸载策略，确保数据类型一致性：

  model = load_checkpoint_and_dispatch(
      model,
      checkpoint="path/to/checkpoint",
      device_map="auto",
      max_memory={0: "79GB"},
      offload_folder="path/to/offload",
      offload_buffers=False,  # 禁用缓冲区卸载
      force_half=False  # 确保不强制转换为半精度
  )
  

---

示例代码

以下是一个完整的示例代码，展示如何解决数据类型不一致的问题：

from accelerate import load_checkpoint_and_dispatch
from transformers import AutoModelForCausalLM
import torch

# 加载模型
model = AutoModelForCausalLM.from_pretrained("deepseek-qwen-distill-32B")

# 加载检查点并分配设备
model = load_checkpoint_and_dispatch(
    model,
    checkpoint="path/to/checkpoint",
    device_map="auto",
    max_memory={0: "79GB"},
    offload_folder="path/to/offload",
    offload_buffers=True,
    force_half=False  # 确保不强制转换为半精度
)

# 将模型权重转换为 Float（如果需要）
model = model.to(dtype=torch.float32)

# 准备输入数据
inputs = torch.tensor(...)  # 假设输入数据是 Float
inputs = inputs.to(dtype=torch.float32)  # 确保输入数据是 Float

# 前向传播
outputs = model(inputs)

---

总结

1. 确保数据类型一致：

   • 模型权重和输入数据的数据类型必须一致（例如，都是 BF16 或都是 Float）。

2. 调整 load_checkpoint_and_dispatch 的参数：

   • 避免强制转换为半精度（force_half=True）。
   • 如果需要，可以在加载后手动转换模型的数据类型。

3. 检查卸载策略：

   • 如果问题仍然存在，可以尝试禁用 offload_buffers 或调整卸载策略。

如果还有其他问题，请随时告诉我！