在阅读langchain chat chat源码的过程中发现了fastchat作为模型管理工具，可以实现模型的加载，模型的释放等模型代理功能。为了加载自定义ptuning、lora微调训练后的语言模型对fastchat的模型加载部分进行了深度的理解。
def load_model(
model_path: str,
device: str = "cuda",
num_gpus: int = 1,
max_gpu_memory: Optional[str] = None,
dtype: Optional[torch.dtype] = None,
load_8bit: bool = False,
cpu_offloading: bool = False,
gptq_config: Optional[GptqConfig] = None,
awq_config: Optional[AWQConfig] = None,
exllama_config: Optional[ExllamaConfig] = None,
xft_config: Optional[XftConfig] = None,
revision: str = "main",
debug: bool = False,
ptuning_checkpoint=None)
这段代码是一个用于加载模型的函数定义，名为 load_model。它接受多个参数，允许用户指定模型路径、设备类型、GPU数量、最大GPU内存使用、数据类型、是否加载8位量化模型、是否使用CPU卸载、各种量化配置、代码版本、调试模式等。函数的主要目的是根据这些参数加载并返回一个模型及其对应的分词器（tokenizer）。
下面是对函数参数的详细解释：
model_path: 模型的路径，可以是本地路径，也可以是Hugging Face模型库中的模型名称。device: 指定模型加载到的设备，可以是"cuda"、"cpu"、"mps"、"xpu"或"npu"。num_gpus: 当使用GPU时，指定GPU的数量。max_gpu_memory: 指定每个GPU的最大内存使用，如果未指定，将自动计算。dtype: 指定模型的数据类型。load_8bit: 是否加载8位量化模型。cpu_offloading: 是否将模型的部分计算卸载到CPU。gptq_config,awq_config,exllama_config,xft_config: 各种不同的量化配置对象。revision: 代码的版本。debug: 是否开启调试模式。ptuning_checkpoint: 用于模型微调的检查点。 函数内部首先获取模型的适配器，然后根据设备类型和参数设置相应的加载选项。接着，根据不同的量化配置加载相应的量化模型。最后，函数加载模型并将其移动到指定的设备上，返回模型和分词器。 这段代码是针对特定的深度学习模型加载和量化场景编写的，涉及到的库包括PyTorch、accelerate、transformers等。在运行此代码之前，需要确保已经安装了这些库。 以下内容为详细的代码注释
def load_model(
model_path: str,
device: str = "cuda",
num_gpus: int = 1,
max_gpu_memory: Optional[str] = None,
dtype: Optional[torch.dtype] = None,
load_8bit: bool = False,
cpu_offloading: bool = False,
gptq_config: Optional[GptqConfig] = None,
awq_config: Optional[AWQConfig] = None,
exllama_config: Optional[ExllamaConfig] = None,
xft_config: Optional[XftConfig] = None,
revision: str = "main",
debug: bool = False,
ptuning_checkpoint=None):
"""
Load a model from Hugging Face.
"""
import accelerate # 导入accelerate库，用于模型加速和并行处理
# 获取模型适配器，根据模型路径确定如何加载模型
adapter = get_model_adapter(model_path)
# 处理设备映射，根据不同的配置参数调整
cpu_offloading = raise_warning_for_incompatible_cpu_offloading_configuration(
device, load_8bit, cpu_offloading
)
if device == "cpu":
kwargs = {"torch_dtype": torch.float32} # CPU上的默认数据类型
# 如果CPU支持avx512_bf16或amx指令集，则使用bfloat16数据类型
if CPU_ISA in ["avx512_bf16", "amx"]:
try:
import intel_extension_for_pytorch as ipex # 尝试导入Intel的PyTorch扩展
kwargs = {"torch_dtype": torch.bfloat16}
except ImportError:
warnings.warn(
"Intel Extension for PyTorch is not installed, it can be installed to accelerate cpu inference"
)
elif device == "cuda":
kwargs = {"torch_dtype": torch.float16} # GPU上的默认数据类型
if num_gpus != 1:
kwargs["device_map"] = "auto" # 多GPU时的设备映射
# 如果没有指定最大GPU内存，则根据实际可用内存计算
if max_gpu_memory is None:
kwargs[
"device_map"
] = "sequential" # 对于不同大小的VRAM，按顺序分配
available_gpu_memory = get_gpu_memory(num_gpus)
kwargs["max_memory"] = {
i: str(int(available_gpu_memory[i] * 0.85)) + "GiB"
for i in range(num_gpus)
}
else:
kwargs["max_memory"] = {i: max_gpu_memory for i in range(num_gpus)}
elif device == "mps":
kwargs = {"torch_dtype": torch.float16} # Apple Silicon上的默认数据类型
import transformers
version = tuple(int(v) for v in transformers.__version__.split("."))
if version < (4, 35, 0):
# 如果transformers库的版本较低，避免mps后端的错误
replace_llama_attn_with_non_inplace_operations()
elif device == "xpu":
kwargs = {"torch_dtype": torch.bfloat16} # Intel XPU上的默认数据类型
# 尝试导入Intel的PyTorch扩展，虽然看起来未使用，但实际上为XPU支持提供了链接
try:
import intel_extension_for_pytorch as ipex
except ImportError:
warnings.warn(
"Intel Extension for PyTorch is not installed, but is required for xpu inference."
)
elif device == "npu":
kwargs = {"torch_dtype": torch.float16} # Huawei NPU上的默认数据类型
# 尝试导入Huawei的PyTorch扩展
try:
import torch_npu
except ImportError:
warnings.warn("Ascend Extension for PyTorch is not installed.")
else:
raise ValueError(f"Invalid device: {device}")
if cpu_offloading:
# 如果启用CPU卸载，则设置量化配置
from transformers import BitsAndBytesConfig
if "max_memory" in kwargs:
kwargs["max_memory"]["cpu"] = (
str(math.floor(psutil.virtual_memory().available / 2 ** 20)) + "Mib"
)
kwargs["quantization_config"] = BitsAndBytesConfig(
load_in_8bit_fp32_cpu_offload=cpu_offloading
)
kwargs["load_in_8bit"] = load_8bit
elif load_8bit:
if num_gpus != 1:
warnings.warn(
"8-bit quantization is not supported for multi-gpu inference."
)
else:
model, tokenizer = adapter.load_compress_model(
model_path=model_path,
device=device,
torch_dtype=
