ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

vLLM Sampling Mask(Distribution Replay)深入解析:让 RL 训练中的 π_old 与 π_θ 共享同一个截断分布

vLLM Sampling Mask(Distribution Replay)深入解析:让 RL 训练中的 π_old 与 π_θ 共享同一个截断分布 vLLM Sampling MaskDistribution Replay深入解析让 RL 训练中的 π_old 与 π_θ 共享同一个截断分布【免费下载链接】vllmA high-throughput and memory-efficient inference and serving engine for LLMs项目地址: https://gitcode.com/GitHub_Trending/vl/vllmSampling Mask采样掩码官方文档也称 Distribution Replay / 分布重放是 vLLM 为强化学习RL训练场景提供的一项引擎级特性。在 GRPO 等 rollout 采样中top-k/top-p 截断会造成采样器实际使用的截断分布与训练时全词表 softmax之间的系统性不匹配进而破坏重要性采样比 π_θ/π_old 的数学一致性本特性通过返回每个生成步真正存活下来的 token 集合让训练侧能在与 rollout 完全相同的支撑集support上归一化。读完本文你将掌握该特性的启用参数、前置约束、底层实现链路以及如何在 RL 训练框架中用 mask 正确计算当前策略的对数概率。为什么要 Sampling Mask截断分布与全词表 softmax 的错位在基于 RL如 GRPO的 rollout 阶段vLLM 采样器实际执行的是截断后的分布logits 经过温度缩放、min-p、top-k/top-p 过滤后被排除的 token 的 logits 会被置为-inf采样只发生在幸存者集合中。然而训练阶段计算对数概率时通常使用全词表 softmax。二者叠加会产生一个微妙但致命的问题π_old旧策略真正采样的动作空间是截断后的核nucleus而π_old(a|s)数值上却来自全词表归一化π_θ当前策略同理。两个策略在 importance sampling 中面对的动作子空间不一致重要性比率失真会破坏训练稳定性。本特性对应 DeepSeek-V3.2 技术报告中描述的Keep Sampling Mask保留采样掩码策略把 rollout 采样时由 top-k/top-p 截断产生的掩码保存下来训练时把同一掩码套用到π_θ上使新旧策略共享完全相同的动作子空间。报告指出将 top-p 采样与 Keep Sampling Mask 结合能有效在 RL 训练中保持语言一致性语言流畅性不被破坏。vLLM 的这一实现正是把该策略从论文落到工程实践。快速开始三种接入方式1. OpenAI 兼容服务 CLIvllm serve model \ --return-sampling-mask \ --logprobs-mode processed_logprobs--return-sampling-mask是引擎级开关默认关闭--logprobs-mode processed_logprobs保证返回的 logprobs 是在截断核上归一化而不是全词表。两者在 参数解析入口 注册最终落在 ModelConfig 的两个字段上return_sampling_mask: bool False # Whether to return the post-processing token support for each sample. logprobs_mode: LogprobsMode raw_logprobs # Indicates the content returned in the logprobs and prompt_logprobs.2. 离线 Python APILLM类from vllm import LLM, SamplingParams llm LLM(model, return_sampling_maskTrue, logprobs_modeprocessed_logprobs) output llm.generate( The capital of France is, SamplingParams(temperature1.0, top_k50, top_p0.95, logprobs1), ) mask output[0].outputs[0].sampling_mask # mask.token_ids: [[187, 326, 512], [42, 88], ...] # mask.token_ids[i] token IDs in the sampling support for generated token iCompletionOutput新增的sampling_mask字段在 vllm/outputs.py 中是一个独立 dataclassdataclass class SamplingMask: Per-token sampling support sets aligned with completion token IDs. Each inner list contains the vocabulary token IDs that survived top-k / top-p / min-p filtering for the corresponding generated token. token_ids: list[list[int]]也就是说mask.token_ids是一个list[list[int]]外层下标对应生成的每个 token内层是该 token 生成时刻经过 top-k/top-p/min-p 过滤后仍存活的词表 token ID 集合。3./inference/v1/generateHTTP 端点掩码同样可以通过 scale-out 的 token-in-token-out 生成端点获取协议定义 与 服务适配 会把SamplingMask.token_ids映射为响应字段{ choices: [{ token_ids: [187, 42, 303], sampling_mask: [[187, 326, 512], [42, 88], [303, 11, 22]], finish_reason: stop }] }注意这里sampling_mask的每一行严格对应token_ids中的每个已生成 token逐位对齐训练侧可以据此逐位置重建掩码。前置要求与参数约束启用前请先核对下表所列的四个必要条件RequirementReason--return-sampling-maskEngine-level opt-in同时会禁用 FlashInfer 融合采样器--logprobs-mode processed_logprobs返回的 logprobs 必须在截断核上归一化而非全词表temperature 0Greedy贪心没有截断分布可言掩码无意义top_k 0约束掩码尺寸纯 top-p 可能产生接近词表大小的掩码Model Runner V2掩码需要异步 D2HGPU→CPU拷贝流水线支撑引擎级配置校验上述约束并非只写在文档里引擎在VLLMConfig.__post_init__阶段通过_verify_sampling_replay_config做硬校验vllm/config/vllm.py一旦设置了return_sampling_mask以下组合会在启动时直接抛ValueError非 Model Runner V2报错 sampling distribution replay requires Model Runner V2Speculative decoding不支持投机解码Diffusion 模型不支持扩散模型引擎级自定义 logits processors--logits-processors不支持logprobs_mode不是processed_logprobs报错并要求显式设置理由是返回的 logprobs 必须与 sampling mask 在同一个截断核上归一化。此外return_sampling_mask还与batch-sharded sampling不兼容在 vllm/config/vllm.py 的 batch-sharded 采样可行性检查中该组合会被列为 blockergather_sampler_output()不会转发SamplingMaskTensors掩码会以None返回。请求级校验除引擎级约束外InputProcessor 在每个请求进入时还会校验SamplingParamstemperature 0→ 报错 sampling distribution replay requires temperature 0greedy 没有截断分布top_k 0→ 报错要求top_k 0理由注释写得很清楚需要用它约束掩码尺寸、降低传输开销并避免潜在 OOM。纯 top-p 之所以被排除是因为最坏情况下幸存集合会膨胀到接近全词表导致每步掩码的数据量与传输成本失控。logprobs_mode四种模式与核上归一化--logprobs-mode支持四种取值见 ModelConfig.logprobs_mode 注释理解它们的区别是正确使用本特性的前提模式返回内容raw_logprobs未经过任何 logit processors如 bad words、惩罚处理前的 logprobsprocessed_logprobs应用全部 processors含温度、top-k/top-p之后的 logprobsraw_logits处理前的 logitsprocessed_logits处理后的 logits当采样掩码开启时采样器要求使用processed_logprobs此时log_softmax是在processed logits被过滤 token 为-inf上计算的softmax 的分母只覆盖截断核本身。这也意味着π_old(a|s)rollout 侧、由 vLLM 产生天然就是核上归一化的对数概率与掩码代表的动作子空间一致。一个实现细节对理解不产生歧义但值得注意在 logits 分支compute_topk_scores会根据logprobs_mode决定对原始 logits 还是处理后的 logits 取 top-k采样器内实现而 prompt token 不经过采样 processors因此raw_*与processed_*对 prompt logprobs 而言结果相同字段注释。工作原理从 logits 到list[list[int]]的四步流水线Step 1应用全部 logit processors 并做 top-k/top-p 过滤采样前Sampler.sample 会按序在 logits 上原位施加logit bias → 各类 penalty → bad words → thinking budget若有→ 温度 → min-p → top-k/top-p。top-k/top-p 过滤会把被排除 token 的 logits 置为-infprocessed_logits幸存者保持有限值——这正是掩码的判据。Step 2用有限 logits找出幸存 token采样完成后若self.return_sampling_mask为真Sampler.call会以本 batch 内校验过的最大 top-k 为宽度上限调用SamplingMaskTensors.from_logits。文档描述这一判据是torch.isfinite(processed_logits)实际落地的 Triton kernel_compact_sampling_mask_kernelvllm/v1/worker/gpu/sample/output.py则执行等价的keep (logits -inf) (logits inf)判据NaN 因比较恒假也被排除。该 kernel 对每个请求行生成三种产物SamplingMaskTensorstoken_ids[num_requests, max_num_kept]的紧凑 ID 缓冲仅存放每行前max_num_kept个存活 tokenpacked_mask[num_requests, ceil(vocab_size / 8)]的逐位打包 bitmaskcounts每行存活 token 的精确数量。其中紧凑 ID 宽度受双重约束max_num_kept min(top_k, vocab_size, MAX_COMPACT_SUPPORT)而MAX_COMPACT_SUPPORT 2048[vllm/v1/worker/gpu/sample/output.py#L73-L74]。当某行存活 token 数超过该宽度时以逐位打包的 bitmask 作为精确兜底编码——这保证了引擎不会因宽支撑而丢精度或内存失控。Step 3随采样 token 一起异步 D2H 拷贝掩码是 GPU 张量需要跨设备传输。这里走的是 Model Runner V2 的异步拷贝流async D2H copy pipeline在 async_utils.py 的 AsyncOutput 中sampling_mask_tensors通过.to_cpu_nonblocking()在独立 CUDA copy stream 上与采样 token、logprobs 一同下发源码特别注释了必须保留对 GPU 张量的引用因为拷贝发生在与张量创建不同的流上。之后在get_output()中执行copy_event.synchronize()并把SamplingMaskTensors转成 CPU 侧的list[vllm/v1/worker/gpu/async_utils.py#L180-L183]。Step 4按请求合并、对齐并组包每个解码步产出的掩码片段并不直接属于某个请求需要经过调度与输出处理两级归并调度器持有return_sampling_mask开关[vllm/v1/core/sched/scheduler.py#L381]在把采样结果写回请求状态时用sampling_masks.slice_request(...)把掩码切分到对应请求[vllm/v1/core/sched/scheduler.py#L2111-L2131]输出处理器为每个在途请求累积sampling_mask_chunks[vllm/v1/engine/output_processor.py#L187]当请求完成时把它们合并成一个SamplingMask[vllm/v1/engine/output_processor.py#L420-L436]最终呈现为按生成位置对齐的list[list[int]]。RL 训练侧使用如何用 mask 计算 π_θ训练侧计算重要性比率π_θ/π_old需要两个量本特性恰好各解决一半。π_old(a|s)——旧策略在截断核上归一化的对数概率由 vLLM 在设置--logprobs-mode processed_logprobs后直接返回。因为log_softmax在 processed logits被过滤 token 为-inf上计算分母只包含截断核rollout 侧得到的每个 token 的 logprob 天然与采样真实分布一致无需任何额外处理。π_θ(a|s)——当前策略在同一个核上的对数概率vLLM 不能替训练框架算 π_θπ_θ 由训练侧正在更新的模型前向产生因此框架需要用返回的 mask 自行实现核上归一化# mask_ids: list[int], the sampling support for this token # logits: the training models raw logits for this position keep torch.zeros(vocab_size, dtypetorch.bool) keep[mask_ids] True masked_logits logits.masked_fill(~keep, float(-inf)) log_prob log_softmax(masked_logits)[sampled_token_id]要点是mask_ids必须使用 rollout 时由 vLLM 返回的掩码即 π_old 的截断集合并原样套用到 π_θ 上两侧在同一 token 集合上做log_softmax重要性比率π_θ/π_old的分母与支撑集严格一致importance sampling 才站得住。局限性与代价从源码可以明确看到该特性的两项工程代价使用前需要评估引擎级 flag 的全局代价--return-sampling-mask一旦开启Sampler.init中use_flashinfer not return_sampling_mask and flashinfer_sampler_supported()会全局禁用 FlashInfer 融合采样器。即使某个请求根本不需要掩码也要走 PyTorch 采样路径并付出相应开销。对需要高吞吐的纯在线推理服务建议与 RL 训练流量分开部署。无流式支持掩码只在最终响应中返回不会出现在中间 streaming chunk 中。RL 训练通常以离线整段生成方式消费 rollout一般不受影响但任何依赖逐 token 流式拿到掩码的场景当前都无法满足。此外掩码按位置与生成 token 严格对齐属于逐位置数据而非逐请求标量请求越长累积的传输与内存占用越大——这正是引擎强制top_k 0、并用MAX_COMPACT_SUPPORT 2048 bitmask 兜底来控制数据量的原因。小结Sampling Mask 是一个精准服务 RL 训练管线的特性它不改变 vLLM 的采样行为而是把采样器实际用过的截断支撑集原样暴露给训练侧让 π_old 与 π_θ 在完全相同的 token 集合上归一化。配合--logprobs-mode processed_logprobs、temperature 0、top_k 0与 Model Runner V2 四个前提并在引擎配置校验vllm/config/vllm.py、请求级校验vllm/v1/engine/input_processor.py、Triton 掩码 kernelvllm/v1/worker/gpu/sample/output.py与异步 D2H 拷贝vllm/v1/worker/gpu/async_utils.py等源码环节都能找到一一对应的实现佐证。如果你正在搭建基于 vLLM 的 GRPO / RLHF 数据管线需要为 rollout 与训练之间消除分布错位这个特性就是现成的标准答案。【免费下载链接】vllmA high-throughput and memory-efficient inference and serving engine for LLMs项目地址: https://gitcode.com/GitHub_Trending/vl/vllm创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表