多GPU情况下，模型加载成功了，但是做推理的时候，还会使用GPU，导致GPU 不够，如何处理，卸载部分到CPU ?
在多 GPU 情况下，模型加载时可能分布在多个 GPU 上，但是推理时 GPU 显存不足。这种情况下，你可以通过将部分模型或者计算卸载到 CPU 来减轻 GPU 的压力。以下是几种方法来处理这种情况：
1. 将部分层移动到 CPU：
你可以手动将模型的某些层移到 CPU。如下代码示例，将某些层的权重转移到 CPU：
|
1 2 3 4 5 6 7 |
# 假设你有一个已加载到多 GPU 的模型 model = AutoModelForCausalLM.from_pretrained(MODEL_ID, device_map="auto", torch_dtype=torch.bfloat16) # 手动将部分层移动到 CPU for name, module in model.named_modules(): if "layer_you_want_to_move" in name: # 替换为你希望移动到CPU的层 module.to('cpu') |
你可以灵活选择哪些层移到 CPU，根据需求动态调整 GPU/CPU 的分配。
2. 分层推理：
另一种方式是手动控制模型层的推理过程，每次仅将一部分层加载到 GPU，其他部分仍在 CPU 上。以下示例展示了如何逐层推理：
|
1 2 3 4 5 6 7 8 9 10 |
model = AutoModelForCausalLM.from_pretrained(MODEL_ID, device_map="auto", torch_dtype=torch.bfloat16) # 手动分层推理，逐层加载到GPU进行计算 input_ids = torch.tensor([[some_input_ids]]).to('cuda') # 将输入加载到GPU output = input_ids for name, module in model.named_children(): module.to('cuda') # 将层加载到GPU output = module(output) # 执行推理 module.to('cpu') # 推理完成后移回CPU以释放显存 |
这种方法可以避免一次性加载整个模型到 GPU，有助于减少 GPU 显存的占用。
3. 使用 accelerate 库自动管理 GPU/CPU：
如果你使用的是 Hugging Face 的 accelerate 库，可以利用它来自动处理多设备模型的分配。accelerate 会自动将 GPU 不足时的模型部分卸载到 CPU 或其他可用设备上。
|
1 2 3 4 5 6 7 8 9 10 11 12 13 |
from transformers import AutoModelForCausalLM from accelerate import init_empty_weights, infer_auto_device_map from accelerate.utils import load_checkpoint_and_dispatch # 初始化模型在空的权重上 with init_empty_weights(): model = AutoModelForCausalLM.from_pretrained(MODEL_ID) # 推断设备映射，将部分层分配到CPU device_map = infer_auto_device_map(model, max_memory={"cuda:0": "10GB", "cpu": "30GB"}, no_split_module_classes=["SomeModuleClass"]) # 加载并分发到设备 model = load_checkpoint_and_dispatch(model, MODEL_ID, device_map=device_map, dtype=torch.bfloat16) |
这段代码可以自动管理模型在多个设备（包括 GPU 和 CPU）上的分布。你可以通过 max_memory 参数控制每个设备的最大内存使用。
4. 使用 offload 模型到 CPU 的方法：
通过使用 Hugging Face 提供的 device_map="auto" 和 offload 功能，可以在需要时将模型的某些部分自动卸载到 CPU。offload 功能会自动将不需要的模型部分从 GPU 卸载到 CPU，以下是示例：
|
1 2 3 4 5 6 7 8 9 |
from transformers import AutoModelForCausalLM model = AutoModelForCausalLM.from_pretrained( MODEL_ID, device_map="auto", offload_folder="offload", # 将模型部分卸载到硬盘以减轻内存压力 offload_state_dict=True, # 启用offload torch_dtype=torch.bfloat16 ) |
在推理过程中，offload 会根据 GPU 的使用情况自动将部分模型卸载到 CPU，以释放 GPU 显存。
5. 减少输入的 batch size：
降低输入的 batch size 也可以有效降低显存占用。如果推理时显存不足，你可以减少 batch size，尤其是在处理大规模输入时。
|
1 2 3 4 5 6 |
outputs = model.generate( input_ids, max_new_tokens=100, num_beams=5, batch_size=1 # 减少 batch size ) |
综合来说，可以根据你的模型规模和硬件资源选择将部分模型卸载到 CPU 或采用自动化工具来管理 GPU/CPU 分配。