七 模型初始化
图来自B站某个视频,发现找不到原视频了!
我们先来看下LLM是怎么结合到vllm中的。
llm = LLM(model=model_path,
dtype='half',
enable_prefix_caching= False,
# dtype='float16'
# 把模型层均分到n个gpu上, 而不是运行n个完整模型
# tensor_parallel_size=1
# gpu利用率最大70%
# gpu_memory_utilization=0.7,
)
tokenizer = AutoTokenizer.from_pretrained(model_path, )
这是模型的入口,model_path路径指向下载的hugging-face模型文件。
class LLM:
...
def __init__(
self,
...
) -> None:
...
# 将外部参数映射为EngineArgs的属性,没做其他修改,便于后续参数的管理
engine_args = EngineArgs(...)
# 使用配置好的engine参数,初始LLMEngine实例
self.llm_engine = LLMEngine.from_engine_args(engine_args, usage_context=UsageContext.LLM_CLASS)
# 全局唯一id,1个 prompt(一个batch可能包含多条prompt)的视为1个request,为这个prompt分配一个唯一id
self.request_counter = Counter()
可以看到通过from_engine_args来加载,继续往下看
def from_engine_args(
cls,
engine_args: EngineArgs,
usage_context: UsageContext = UsageContext.ENGINE_CONTEXT,
stat_loggers: Optional[Dict[str, StatLoggerBase]] = None,
) -> "LLMEngine":
"""Creates an LLM engine from the engine arguments."""
# Create the engine configs.
engine_config = engine_args.create_engine_config()
executor_class = cls._get_executor_cls(engine_config)
# Create the LLM engine.
engine = cls(
**engine_config.to_dict(),
executor_class=executor_class,
log_stats=not engine_args.disable_log_stats,
usage_context=usage_context,
stat_loggers=stat_loggers,
)
return engine
from_engine_args输入参数如下:
cls(…) 指向如下代码:
class LLMEngine:
...
def __init__(
self,
...
) -> None:
...
if not self.model_config.skip_tokenizer_init:
self.tokenizer = self._init_tokenizer()
self.detokenizer = Detokenizer(self.tokenizer)
else:
self.tokenizer = None
self.detokenizer = None
...
self.model_executor = executor_class(
model_config=model_config,
cache_config=cache_config,
parallel_config=parallel_config,
scheduler_config=scheduler_config,
device_config=device_config,
lora_config=lora_config,
multimodal_config=multimodal_config,
speculative_config=speculative_config,
load_config=load_config,
prompt_adapter_config=prompt_adapter_config,
)
if not self.model_config.embedding_mode:
self._initialize_kv_caches()
...
# pipeline_parallel_size:并行的gpu数量, 会把可用的 物理blocks平均分配到并行的gpu上
# 同时, 每个gpu都会维护一个调度器scheduler, self.scheduler是包含多个scheduler的list
self.scheduler = [
Scheduler(
scheduler_config, cache_config, lora_config,
parallel_config.pipeline_parallel_size
) for _ in range(parallel_config.pipeline_parallel_size)
]
...
可以发现在vllm初始化时,主要初始化4个模块:tokenizer(分词器),model_executor(tf模型转换到vllm模型),self._initialize_kv_caches(kv block初始化),scheduler (调度器), 这在本章开头的结构图中也能清晰看到。
tokenizer比较简单,这里略过,schedule在第二篇文章中已经讲过。
我们来看下model_executor与_initialize_kv_caches的具体工作,这两部分代码是以后向vllm手动添加新模型(model_executor),优化vllm推理性能(_initialize_kv_caches)的核心代码。
7.1 model_executor
executor_class继承自基类ExecutorBase,有cpu_executor,gpu_executor,tpu_executor…,等各种执行器可选,由当前设备类型,或指定executor来决定使用哪一个。我们以gpu_executor来说明,其他executor也都大同小异。
class GPUExecutor(ExecutorBase):
uses_ray: bool = False
def _init_executor(self) -> None:
"""Initialize the worker and load the model.
"""
assert self.parallel_config.world_size == 1, (
"GPUExecutor only supports single GPU.")
self.driver_worker = self._create_worker()
self.driver_worker.init_device()
self.driver_worker.load_model()
...
def execute_model(
self, execute_model_req: ExecuteModelRequest
) -> Optional[List[Union[SamplerOutput, PoolerOutput]]]:
output = self.driver_worker.execute_model(execute_model_req)
return output
...
self.driver_worker是work(vllm/worker/worker.py)的一个实例对象(每个gpu上的都维护着自己的Worker实例),负责维护 KV-cache,并在 GPU 上执行模型。在分布式推理的情况下,每个work都会被分配模型的一部分(不同的head并行计算,然后汇总计算结果)。
self.driver_worker.load_model()是加载模型的方法,但经过多层转包后,才能找到真正的初始化模型的代码:
- vllm/model_executor/model_loader/loader.py class DefaultModelLoader
def load_model(self, *, model_config: ModelConfig,
device_config: DeviceConfig,
lora_config: Optional[LoRAConfig],
multimodal_config: Optional[MultiModalConfig],
parallel_config: ParallelConfig,
scheduler_config: SchedulerConfig,
cache_config: CacheConfig) -> nn.Module:
target_device = torch.device(device_config.device)
with set_default_torch_dtype(model_config.dtype):
with target_device:
model = _initialize_model(model_config, self.load_config,
lora_config, multimodal_config,
cache_config, scheduler_config)
model.load_weights(
# 加载model.safetensors权重文件
self._get_weights_iterator(model_config.model,
model_config.revision,
fall_back_to_pt=getattr(
model,
"fall_back_to_pt_during_load",
True)), )
for _, module in model.named_modules():
quant_method = getattr(module, "quant_method", None)
if quant_method is not None:
# When quant methods need to process weights after loading
# (for repacking, quantizing, etc), they expect parameters
# to be on the global target device. This scope is for the
# case where cpu offloading is used, where we will move the
# parameters onto device for processing and back off after.
with device_loading_context(module, target_device):
quant_method.process_weights_after_loading(module)
return model.eval()
我们解析下涉及的两个主要函数:
- vllm/model_executor/model_loader/loader.py
def _initialize_model(
model_config: ModelConfig,
load_config: LoadConfig,
lora_config: Optional[LoRAConfig],
multimodal_config: Optional[MultiModalConfig],
cache_config: CacheConfig,
scheduler_config: Optional[SchedulerConfig] = None) -> nn.Module:
"""Initialize a model with the given configurations."""
# 通过下载hf模型时自带的config,根据config['architectures']参数,获得当前模型名称
model_class = get_model_architecture(model_config)[0]
# 取得量化相关参数,在当前版本中没有启用该参数
quant_config = _get_quantization_config(model_config, load_config)
# 通过加载vllm/model_executor/models/llama.py,获得模型结构(这是vllm改造后的结构)
return model_class(config=model_config.hf_config,
# cache_config=cache_config,
# quant_config=quant_config,
**_get_model_initialization_kwargs(
model_class, lora_config, multimodal_config,
scheduler_config))
_initialize_model函数的功能为通过hf模型的config参数,获得模型名,
然后根据这个名称去加载vllm改造后的该模型模型结构
我们以llama为例来说明如何加载hf权重:
- vllm/model_executor/models/llama.py
def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
# vllm与hf两种模型实现方式之间的名称映射
stacked_params_mapping = [
# (param_name, shard_name, shard_id)
# vllm, hf,share_id
(".qkv_proj", ".q_proj", "q"),
(".qkv_proj", ".k_proj", "k"),
(".qkv_proj", ".v_proj", "v"),
(".gate_up_proj", ".gate_proj", 0),
(".gate_up_proj", ".up_proj", 1),
]
# 获得当前vllm改造后llama模型的参数和对应的权重(此时的权重应是随机生成的)
params_dict = dict(self.named_parameters())
# 遍历hf模型每层参数的名称和权重
for name, loaded_weight in weights:
...
# vllm, hf,share_id
for (param_name, weight_name, shard_id) in stacked_params_mapping:
if weight_name not in name:
continue
# 将hf模型的层名,替换为vllm中的层名
name = name.replace(weight_name, param_name)
# Skip loading extra bias for GPTQ models.
if name.endswith(".bias") and name not in params_dict:
continue
if is_pp_missing_parameter(name, self):
continue
# 获得vllm改造后llama权重参数
param = params_dict[name]
weight_loader = param.weight_loader
# 将hf模型参数更新到对应的vllm模型参数中,完成权重参数的映射工作
weight_loader(param, loaded_weight, shard_id)
break
else:
...
通过上述vllm中llama的load_weights方法(经过观察, 所有decode-only模型的load_weights几乎都一样),将vllm模型和hf模型不同参数名之间做映射,之后将hf类型的权重赋值给vllm模型中(通过参数名联系),至此,完成模型转换工作。
注:需要知道模型中有不同结构,所有weight_loader(vllm/model_executor/layers/linear.py)也有多个变体(分布在不同类中)。
以对QKV的转换为例说明weight_loader的变换过程(源码比较复杂,这里仅描述下处理逻辑):
llama3.1的qkv是分开计算的,类似于下面这样
self.q = nn.Linear(dim, dim_q, bias=False)
self.k = nn.Linear(dim, dim_kv, bias=False)
self.v = nn.Linear(dim, dim_kv, bias=False)
而vllm中会把他们合并起来,类似于下面这样
self.qkv=nn.Linear(dim, dim_q+2*dim_kv, bias=False)
通过这个模块的解析,我们可以知道,对未支持的新模型也能通过手动修改load_model源码的方式在vllm中使用。
7.2 _initialize_kv_caches
作用是计算当前blocks总量,可用blocks数量。
tranformers中,一个正常的k/v shape为[batch_size, nums_head, len_k, head_dim](推理阶段,len_k=1)
vllm中kv_cache_shape=[2, num_blocks, block_size, num_kv_heads, head_size]
一个块(block)占用空间的计算公式如下(2表示kv各一个,它们是成对出现的):2 * block_size * num_head * head_size * num_layers,
即每个 token 对应的 K V 个数为2, 每个块可以存放 block_size 个 token 对应的 K V 值,每个 token 对应的 K V 占用空间为2 * num_head * head_size * num_layers * dtype_size,所以每个块总共要存放block_size * 2 * num_head * head_size * num_layers * dtype_size个值。
num_layers是模型的layers层数,每个token要保存计算过的所有层的kv值,这样才算一个完整的kv-cache。
kv每个值占用的空间为 dtype_size 个字节(如果 tensor 的 dtype 为 float16,则 dtype_size 为 2,dtype 为 float32,则 dtype_size 为 4)。
一个block占用空间的计算代码如下:
- vllm/worker/cache_engine.py
- vllm/engine/llm_engine.py
def _initialize_kv_caches(self) -> None:
"""Initialize the KV cache in the worker(s).
The workers will determine the number of blocks in both the GPU cache
and the swap CPU cache.
"""
num_gpu_blocks, num_cpu_blocks = self.model_executor.determine_num_available_blocks()
...
self.cache_config.num_gpu_blocks = num_gpu_blocks
self.cache_config.num_cpu_blocks = num_cpu_blocks
self.model_executor.initialize_cache(num_gpu_blocks, num_cpu_blocks)
_initialize_kv_caches方法的目的是计算出GPU/CPU block数量,然后对这些block进行初始化。
计算block数量的方法为self.model_executor.determine_num_available_blocks()
- vllm/worker/worker.py
def determine_num_available_blocks(self) -> Tuple[int, int]:
# Profile the memory usage of the model and get the maximum number of
# cache blocks that can be allocated with the remaining free memory.
torch.cuda.empty_cache()
# Execute a forward pass with dummy inputs to profile the memory usage
# of the model.
# 构建推理允许的最大seq和tokens 数量组成的推理数据,进行不使用kv-cache的模型推理
self.model_runner.profile_run()
# Calculate the number of blocks that can be allocated with the
# profiled peak memory.
torch.cuda.synchronize()
# 记录此时可用的GPU和总GPU数量,此时模型运行占用的GPU显存还没释放
free_gpu_memory, total_gpu_memory = torch.cuda.mem_get_info()
# peak_memory就是当前模型占用的显存
peak_memory = self.init_gpu_memory - free_gpu_memory
...
# 获得一个block占用的GPU显存
cache_block_size = self.get_cache_block_size_bytes()
# 计算总的可用GPU block数量
num_gpu_blocks = int(
(total_gpu_memory * self.cache_config.gpu_memory_utilization -peak_memory) // cache_block_size)
# 计算CPU数量,对于CPU,不需要额外计算,因为是固定大小的内存。
num_cpu_blocks = int(self.cache_config.swap_space_bytes // cache_block_size)
num_gpu_blocks = max(num_gpu_blocks, 0)
num_cpu_blocks = max(num_cpu_blocks, 0)
if self.model_runner.lora_manager:
self.model_runner.remove_all_loras()
gc.collect()
torch.cuda.empty_cache()
return num_gpu_blocks, num_cpu_blocks
**self.model_runner.profile_run()**作用是构建假数据,走一遍不使用kv-cache的模型推理,记录此时的GPU占用情况。
profile_run流程如下(代码太多,不在此贴出,代码不难, 想进一步了解细节可去看源码):
- 构建假数据
初始化LLMEngine引擎时,会提供两个重要参数(这两个参数在当前版本由budget管理):
max_num_seqs:在1个推理阶段中,可处理的最大seqs数量
max_num_batched_tokens:在1个推理阶段中,可处理的最大tokens数量
这两个参数值由外部指定, 若未指定, 系统会分配一个。那么如何通过这两个值构建数据呢?
假设在推理过程中,平均一个seq要处理max_num_batched_tokens // max_num_seqs个token,余数部分我们默认放在第一个seq中。
例如,若max_num_batched_tokens=10,max_num_seqs = 3,那么可以构建出3条seq,每个seq的长度分别为4,3,3
使用这些空数据,走一遍推理流程,可以获得模型使用GPU显存的情况。
(free_gpu_memory, total_gpu_memory = torch.cuda.mem_get_info())
计算出分配多少的显存给KV cache:
- 分配给KV cache显存 = gpu总显存 - 不使用KV cache做1次推理时的显存占用(包括模型本身和推理过程中的中间数据)
在上述代码中有详细注释。
分配kv-cache
计算出了可用block数量,接下就能通过initialize_cache初始化vllm推理过程中的kv-cache了。
- vllm/worker/worker.py
def _init_cache_engine(self):
assert self.cache_config.num_gpu_blocks is not None
self.cache_engine = [
CacheEngine(self.cache_config, self.model_config,
self.parallel_config, self.device_config)
for _ in range(self.parallel_config.pipeline_parallel_size)
]
self.gpu_cache = [
self.cache_engine[ve].gpu_cache
for ve in range(self.parallel_config.pipeline_parallel_size)
]
初始化kv-cache的工作最终是在CacheEngine的__init__()函数中完成,层层嵌套,vllm架构越来复杂了。
- vllm_module/worker/cache_engine.py
def _allocate_kv_cache(
self,
num_blocks: int,
device: str,
) -> List[torch.Tensor]:
"""Allocates KV cache on the specified device."""
# shape=[num_blocks, block_size,num_kv_heads,head_size]
kv_cache_shape = self.attn_backend.get_kv_cache_shape(
num_blocks, self.block_size, self.num_kv_heads, self.head_size)
pin_memory = is_pin_memory_available() if device == "cpu" else False
kv_cache: List[torch.Tensor] = []
# 遍历每一层,一个token的完整kv-cache包含所有层的子kv
for _ in range(self.num_attention_layers):
# null block in CpuGpuBlockAllocator requires at least that
# block to be zeroed-out.
# We zero-out everything for simplicity.
kv_cache.append(
torch.zeros(kv_cache_shape,
dtype=self.dtype,
pin_memory=pin_memory,
device=device))
return kv_cache
最终kv-cache的样子(显卡:RTX4090,模型:Meta-Llama-3.1-8B-Instruct)如下:
kv-cache.shape每个维度代表含义如下:
list 28:当前模型有28层,每层都要保持当前层计算的kv
内部元素的shape含义:
- 2:分别存储k和v的计算结果
- 2760:当前GPU有2760个block
- 16:每个block有16个槽位,即可以放16个k或v的值
- 8:当前模型head数量
- 128:每个head的head_size
这个kv-cache就是推理过程中用于存储kv值的容器,这里一次性初始好,全部填充为0,所以在实际推理过程中会发现,vllm会直接把显存占用爆涨到一个很大值,就是因为初始化了很多预填充kv-cache的block。