PEFT项目中的多GPU加载问题分析与解决方案

问题背景

在使用PEFT(Parameter-Efficient Fine-Tuning)库进行分布式模型加载时，开发者可能会遇到一个特殊的内存占用问题。具体表现为：当使用torchrun在多GPU环境下加载LoRA模型时，即使明确指定了设备映射(device_map)，非目标GPU(特别是GPU0)也会出现意外的显存占用。

问题现象

通过一个简单的测试脚本可以复现这个问题：

import os
from peft import AutoPeftModelForCausalLM

if __name__ == '__main__':
    rank = int(os.getenv("RANK"))
    model = AutoPeftModelForCausalLM.from_pretrained("ybelkada/opt-350m-lora", device_map=rank)
    print(f"{rank} Finished loading.")
    input()  # 暂停以观察显存占用

当使用torchrun --nproc_per_node 8运行上述脚本时，观察到的现象是：

1. 所有进程都会在GPU0上创建CUDA上下文
2. GPU0会出现约1100MB的显存占用
3. 这种现象仅在使用PEFT加载LoRA模型时出现，加载普通基础模型时不会发生

问题根源分析

经过深入调查，发现问题的根本原因在于PEFT模型加载过程中的设备分配机制：

1. 状态字典加载行为：当使用AutoPeftModelForCausalLM.from_pretrained加载模型时，内部会先加载基础模型，然后加载LoRA适配器。在加载状态字典(state_dict)时，如果没有明确指定设备，PyTorch会默认使用cuda:0。

2. 设备映射的局限性：虽然device_map参数可以控制模型参数的最终存放位置，但在加载过程中仍然会短暂地在默认设备上创建中间张量，导致显存占用。

3. 分布式环境的特殊性：在torchrun创建的分布式环境中，所有进程共享相同的CUDA上下文，这使得GPU0的显存占用问题更加明显。

解决方案

针对这个问题，PEFT核心开发者提出了明确的解决方案：使用PeftModel.from_pretrained替代AutoPeftModelForCausalLM.from_pretrained，并显式指定torch_device参数。

优化后的代码示例如下：

import os
from transformers import AutoModelForCausalLM
from peft import PeftModel, LoraConfig

if __name__ == '__main__':
    rank = int(os.getenv("RANK"))
    model_name_or_path = "ybelkada/opt-350m-lora"
    
    # 分步加载模型
    peft_config = LoraConfig.from_pretrained(model_name_or_path)
    base_model = AutoModelForCausalLM.from_pretrained(
        peft_config.base_model_name_or_path, 
        device_map=rank
    )
    peft_model = PeftModel.from_pretrained(
        base_model,
        model_name_or_path,
        config=peft_config,
        torch_device=rank  # 关键：显式指定加载设备
    )
    
    print(f"{rank} Finished loading.")
    input()  # 暂停观察

技术要点解析

1. 分步加载的优势：

   • 先加载基础模型到指定设备
   • 再加载LoRA适配器，明确指定目标设备
   • 避免了中间张量在默认设备上的创建

2. torch_device参数的重要性：

   • 确保所有模型参数和中间张量都在目标设备上创建
   • 避免了默认设备(cuda:0)的显存占用
   • 在分布式环境中特别重要

3. 内存效率对比：

   • 原始方法：每个进程都会在GPU0上占用约1100MB显存
   • 优化方法：显存仅在使用中的GPU上分配，完全隔离

最佳实践建议

1. 在分布式训练/推理场景中，始终明确指定torch_device参数
2. 对于复杂模型加载流程，考虑分步加载策略
3. 监控各GPU的显存使用情况，确保符合预期
4. 在容器化部署时，注意CUDA_VISIBLE_DEVICES的设置

总结

PEFT库为参数高效微调提供了强大支持，但在分布式环境中使用时需要注意设备分配细节。通过本文介绍的方法，开发者可以避免不必要的显存占用，实现更高效的分布式模型加载。理解PyTorch的底层设备分配机制对于优化深度学习应用的内存使用至关重要。