原理Beam Search束搜索是一種在序列生成任務(wù)中常用的搜索算法用于在每一步生成時平衡計算效率和結(jié)果質(zhì)量文字。主要作用減少搜索空間相對于貪心搜索而言在序列生成如機器翻譯、文本生成、語音識別時如果每一步都保留所有可能的候選詞窮舉搜索計算量將指數(shù)增長。Beam Search 在每一步只保留 top-kk 稱為 beam size個最有可能的候選序列大幅降低計算負擔。近似最優(yōu)解它并不是窮舉所有路徑那樣保證全局最優(yōu)但計算不可行而是用局部貪心擴展結(jié)合一定寬度搜索得到一個接近最優(yōu)的序列。平衡生成質(zhì)量與速度Beam Size 越小速度越快但可能錯過更好的序列。Beam Size 越大質(zhì)量通常更好更接近窮舉搜索但計算越慢。工作流程假設(shè) beam_size 2生成句子從起始符開始第一步模型輸出每個詞的概率保留概率最高的 2 個詞如 A, B。對于這 2 個候選詞分別預(yù)測下一個詞的概率這時會得到 2 × V 種可能V 是詞表大小。從這 2 × V 個候選序列長度為 2中選擇總概率或?qū)?shù)概率之和最高的 2 個繼續(xù)擴展。重復(fù)直到遇到結(jié)束符最后從 beam 中選擇分數(shù)最高的序列輸出。vLLM 實現(xiàn)beamsearch 可視化https://huggingface.co/spaces/m-ric/beam_search_visualizerif__name____main__:llmLLM(modelfacebook/opt-125m)paramsBeamSearchParams(beam_width2,max_tokens50)prompts[TextPrompt(promptThe future of artificial intelligence)]outputsllm.beam_search(prompts,params)foroutputinoutputs:generated_textoutput.sequences[0].textprint(fGenerated text:{generated_text!r})參數(shù)含義max_tokens: 迭代次數(shù)for range(max_token)beam_width: 束寬度。有下面兩個作用1. 針對每個輸入限制模型推理保留 top 2*beam_width 的 token2. 控制每次迭代需要參與推理的 token 數(shù)量自回歸階段 promot [“The future of artificial intelligence”]itemprefilldecodeinput_ids (token_id)[2, 133, 499, 9, 7350, 2316, 0, 0][13]positions[0, 1, 2, 3, 4, 5, 0, 0]logits_indices50hidden_states model_output tensor(8, 768) model_output tensor(1, 768)sample_hidden_states hidden_states[logits_indices] tensor(1, 768) hidden_states[logits_indices] tensor(1, 768)logitstensor(1, 50257)tensor(1, 50257)GPUModelRunner._update_states()更新后的狀態(tài)input_batches會被 _prepare_inputs 函數(shù)使用更新 input_ids以創(chuàng)建模型所需的GPU張量輸入。這是GPU模型運行器中狀態(tài)管理和輸入準備的關(guān)鍵連接點。未使用 beam_search : sampling_metadata.max_num_logprobs None使用 beam_search: sampling_metadata.max_num_logprobs 10