vllm [RFC]:无分类器指导

093gszye  于 4个月前  发布在  其他
关注(0)|答案(2)|浏览(66)

动机

我是被提名为ICML'24亮点论文的Stay On Topic with Classifier-Free Guidance(https://openreview.net/forum?id=RiM3cl9MdK&noteId=s1BXLL1YZD)的作者之一。CFG是一种采样技术,允许LLMs更紧密地跟随提示,但代价是每个标记需要进行两次前向传递以及2 kv缓存。CFG在整体上比标准基准带来了非平凡的改进。
我非常希望将CFG实现到vLLM中。如果可能的话,我想在vLLM代码库中获得一些指导。

建议的更改

CFG比较两个不同提示(一个“正向提示”a和一个“负向提示”或“无条件”b)之间的下一个标记logits。以下是伪算法:

while we sample:
    logits_a = log_softmax(model(prompt_a))
    logits_b = log_softmax(model(prompt_b))
    logits = logits_b + cfg_scale * (logits_a - logits_b)
    next_token = sample_from(logits)
    prompt_a.append(next_token)
    prompt_b.append(next_token)

如您所见,这需要两个并发的kv缓存以实现高效的实现。我尝试查找如何实现Speculative Decoding,但这相当复杂,比CFG所需的更多。

反馈期

  • 无响应*

抄送列表

  • 无响应*

其他事项

如果给予足够的指导,我愿意自己实现它,因为我认为这是一个不容易实现的东西。我认为类似于/重用Speculative Decoding的部分可能会被使用,但代码是非简单的。

qxsslcnc

qxsslcnc2#

你好@Vermeille。做得很好。
在vLLM中实现这个功能,可以类似于Speculative Decoding:

LLMEngine
CFGWorker
< logic which calls the underlying worker twice, does logit math, samples >
Worker (2x)

这种设计的主要好处是,你可以在现有的LLMEngine和调度器中管理两个块表,而无需进行任何修改。这是通过将KV缓存空间平均分成两个大小相等的区域[1]来实现的。然后相同的块表可以同时用于两个模型。实际上,你可以相对简单地原型化这个方法;唯一缺少的主要部分是你需要让其中一个Worker不加载权重(例如,权重加载与其他Worker共享)。
或者,你可以使用单个Worker并修改具有恒定偏移量的块表,从而获得独立的KV缓存。
这种设计的次要好处是硬件无关性;你的实现可以与非NVIDIA、非AMD的硬件后端一起工作。
[1]
vllm/vllm/spec_decode/spec_decode_worker.py
第255行到第274行 70c232f
| | defdetermine_num_available_blocks(self) ->Tuple[int, int]: |
| | """Determine the number of cache blocks to use. |
| | |
| | This is done by profiling the scorer model (which is typically the |
| | larger of the two). Then the total memory which would be used by the |
| | scorer cache is divided evenly between the proposer and scorer model KV, |
| | such that the number of blocks is equal in both KV caches. |
| | """ |
| | num_gpu_blocks, num_cpu_blocks= ( |
| | self.scorer_worker.determine_num_available_blocks()) |
| | |
| | scorer_cache_block_size_bytes= ( |
| | self.scorer_worker.get_cache_block_size_bytes()) |
| | proposer_cache_block_size_bytes= ( |
| | self.proposer_worker.get_cache_block_size_bytes()) |
| | |
| | new_num_gpu_blocks=split_num_cache_blocks_evenly( |
| | scorer_cache_block_size_bytes, proposer_cache_block_size_bytes, |
| | num_gpu_blocks) |
| | returnnew_num_gpu_blocks, num_cpu_blocks |

相关问题