如何用vLLM高效提取大语言模型全层隐藏状态嵌入？

原本用Hugging Face Transformers的generate函数从LLM特定层提取嵌入，哪怕批量处理速度还是慢。想用vLLM提速，而且不想调用generate，但试的vLLM代码只能拿到最后一层隐藏状态，没法提取指定层的嵌入，求更快的实现方法。

原Transformers代码：

def process_row(prompt: str, model, tokenizer, layers_to_use: list, remove_period: bool):
    """
    Processes a row of data and returns the embeddings.
    """
    if remove_period:
        prompt = prompt.rstrip(". ")
    inputs = tokenizer(prompt, return_tensors="pt")
    with torch.no_grad():
        outputs = model.generate(inputs.input_ids, output_hidden_states=True, return_dict_in_generate=True, max_new_tokens=1, min_new_tokens=1)
    embeddings = {}
    for layer in layers_to_use:
        last_hidden_state = outputs.hidden_states[0][layer][0][-1]
        embeddings[layer] = [last_hidden_state.numpy().tolist()]
    return embeddings

用户尝试的vLLM代码片段：

from vllm import LLM, SamplingParams
from vllm.sequence import (SamplerOutput, Sequence, SequenceGroup, SequenceData, 
                           SequenceGroupMetadata, SequenceStatus)
from transformers import LlamaModel, LlamaTokenizer
from vllm import EngineArgs, LLMEngine, SamplingParams, RequestOutput
from vllm.sequence import SamplerOutput, SequenceData, SequenceGroupMetadata


llm = LLM(model=path_to_llama2)


# Enable top-k sampling to reflect the accurate memory usage.
vocab_size = llm.llm_engine.workers[0].model.config.vocab_size
sampling_params = SamplingParams(top_p=0.99, top_k=vocab_size - 1)
max_num_batched_tokens = llm.llm_engine.workers[0].scheduler_config.max_num_batched_tokens
max_num_seqs = llm.llm_engine.workers[0].scheduler_config.max_num_seqs

prompt = train[0]
prompt_token_ids = llm.llm_engine.tokenizer.encode(prompt) #[2, 100, 524, 10]
seqs = []
    
group_id = 1
seq_data = SequenceData(prompt_token_ids)
seq = SequenceGroupMetadata(
    request_id=str(group_id),
    is_prompt=True,
    seq_data={group_id: seq_data},
    sampling_params=sampling_params,
    block_tables=None,
)
seqs.append(seq)
input_tokens, input_positions, input_metadata = llm.llm_engine.workers[0]._prepare_inputs(
    seqs)
prompt_len = len(seq_data.prompt_token_ids)
input_tokens = input_tokens[:prompt_len]
input_positions = input_positions[:prompt_len]
# Execute the model.
num_layers = llm.llm_engine.workers[0].model_config.get_num_layers(llm.llm_engine.workers[0].parallel_config)
tempOut = llm.llm_engine.workers[0].model.model(
    input_ids=input_tokens,
    positions=input_positions,
    kv_caches=[(None, None)] * num_layers,
    input_metadata=input_metadata,
    cache_events=None,
)
print(tempOut.size())

要在vLLM里高效提取指定层的隐藏状态，不用调用generate，可以这么做：

1. 开启模型隐藏状态输出 #

vLLM封装的模型默认只返回最后一层隐藏状态，调用forward时加上output_hidden_states=True参数，就能拿到所有层的输出。下面是批量提取的完整代码：

import torch
from vllm import LLM, EngineArgs
from vllm.sequence import SequenceGroupMetadata, SequenceData

# 初始化LLM引擎，根据GPU显存调整批量参数
engine_args = EngineArgs(
    model=path_to_llama2,
    max_num_batched_tokens=4096,
    max_num_seqs=64
)
llm = LLM(**engine_args.to_dict())

def extract_embeddings_vllm(prompts: list, layers_to_use: list, remove_period: bool = False):
    embeddings = {layer: [] for layer in layers_to_use}
    
    # 预处理所有prompt
    processed_prompts = [p.rstrip(". ") if remove_period else p for p in prompts]
    
    # 批量编码prompt
    tokenizer = llm.llm_engine.tokenizer
    prompt_token_ids_list = [tokenizer.encode(p) for p in processed_prompts]
    
    # 构造批量请求的元数据
    seq_groups = []
    for idx, token_ids in enumerate(prompt_token_ids_list):
        seq_data = SequenceData(token_ids)
        seq_group = SequenceGroupMetadata(
            request_id=str(idx),
            is_prompt=True,
            seq_data={idx: seq_data},
            sampling_params=None,
            block_tables=None,
        )
        seq_groups.append(seq_group)
    
    # 准备模型输入
    worker = llm.llm_engine.workers[0]
    input_tokens, input_positions, input_metadata = worker._prepare_inputs(seq_groups)
    
    # 获取模型总层数
    num_layers = worker.model_config.get_num_layers(worker.parallel_config)
    
    # 前向计算，开启隐藏状态输出
    with torch.no_grad():
        outputs = worker.model.model(
            input_ids=input_tokens,
            positions=input_positions,
            kv_caches=[(None, None)] * num_layers,
            input_metadata=input_metadata,
            cache_events=None,
            output_hidden_states=True  # 关键参数：返回所有层隐藏状态
        )
    
    # 解析每个prompt的指定层嵌入（取最后一个token的输出）
    prompt_lens = [len(tids) for tids in prompt_token_ids_list]
    current_pos = 0
    for prompt_len in prompt_lens:
        for layer in layers_to_use:
            last_token_emb = outputs.hidden_states[layer][current_pos + prompt_len - 1].cpu().numpy().tolist()
            embeddings[layer].append(last_token_emb)
        current_pos += prompt_len
    
    return embeddings

2. 提速核心要点 #

• 跳过生成逻辑：直接调用模型forward方法，省去generate相关的生成计算开销。
• 批量处理：一次性传入多个prompt，vLLM的调度器会自动优化批处理，最大化GPU利用率，这是提速的关键。
• 精准提取：按每个prompt的token长度切割批量输出的隐藏状态，拿到目标token的指定层嵌入。

内容的提问来源于stack exchange，提问作者lazytux