TokenSpeed 位级确定性设计详解RL 训练信号被推理引擎毁掉一文讲清 numerics 机制【免费下载链接】tokenspeedTokenSpeed is a speed-of-light LLM inference engine.项目地址: https://gitcode.com/gh_mirrors/to/tokenspeedTokenSpeed 是一个号称光速的 LLM 推理引擎它的--numerics rl-bitwise模式承诺同一个请求无论跑多少次、和谁拼 batch产生的 token 和 logprob 都逐位一致。这个能力专为 RL强化学习训练场景设计——因为训练信号对推理引擎的微小数值漂移极其敏感一个浮点位的差异就可能让奖励信号失真。1️⃣ 为什么 RL 训练信号会被推理引擎毁掉在 RLHF / RL 微调流程中推理引擎负责生成 rollout采样轨迹训练框架负责计算 logprob 作为奖励信号的基础。问题出在推理引擎和训练框架的浮点运算顺序不同推理引擎为了性能会做 TF32 加速、autotune 选最优 tactic、fused all-reduce、split-KV attention 等这些优化让结果接近但不完全相同于训练框架的逐步计算一个 ulp最后一位的偏差在温度采样T0时可能让采样到的token 都不一样更隐蔽的是同一请求和不同的 co-batch 请求拼在一起结果也可能不同——这就是batch 非确定性对 RL 来说这意味着奖励信号不再可复现训练梯度里混入了噪声。2️⃣--numerics rl-bitwise一份位级契约TokenSpeed 用--numerics这个启动参数来声明推理引擎承诺的数值契约。当前有两个模式模式含义auto默认保持所有性能优化不承诺确定性rl-bitwise承诺位级确定性所有开关自动收紧核心设计思想是不要让用户手动拼十几个开关一个 envelope信封把所有相关开关一次收紧并且只收紧、不放松——如果用户显式设了与契约冲突的参数启动时直接报错而不是静默回退。相关配置定义在 numerics.py 中。3️⃣ 三大承诺运行不变、批不变、训练器对齐rl-bitwise模式做出三层承诺层层递进3.1 运行不变Run Invariance同一个请求跑 10 次每次结果逐位相同。实现方式关闭 autotunekernel 的 tactic 选择不再依赖运行时探测关闭 TF32避免 NVIDIA GPU 的浮点精度降级关闭 PDLProgrammatic Dependent Launchkernel 链严格串行化固定 all-reduce 算法NCCL 算法固定为 Ring协议固定为 Simple避免消息大小不同导致关联顺序变化3.2 批不变Batch Invariance同一个请求不管和谁拼 batch、调度器怎么切 chunk结果逐位相同。这是最难的一层。关键在于kernel 选择机制TokenSpeed 的 kernel registry 区分卖家声明traits和买家要求features在rl-bitwise模式下每个算子选择点强制要求kernel 声明batch_invariant特性找不到满足特性的 kernel启动时直接报错绝不静默回退到性能版 kernel这个无静默回退no silent fallback的设计是整套机制的基石。3.3 训练器对齐Trainer Alignment推理引擎的前向传播在已知与训练框架不同的操作上切换到训练框架的运算顺序。这是最强的一层承诺。具体包括开关做什么--sampling-stream per-request随机流按请求 seed 和位置索引而非 batch 行号--yarn-ramp-mask-device cpuRoPE 频率表在 CPU 上计算和训练器一致--mla-lora-scale runtimeMLA norm scale 在运行时乘而非折进权重--layer-boundary-norm unfused层边界 norm 拆成 add 独立 RMSNormbf16 物化--router-topk torchMoE 路由用 PyTorch 的 fp32 softmax topk--logprob-order megatronlogprob 按 Megatron 的 vocab-parallel 交叉熵顺序计算--moe-combine-order slotMoE top-k 输出按 slot 顺序 fp32 求和这些开关的完整定义见 server_args.py 和 numerics.md。4️⃣ 无静默回退一个反直觉的设计哲学传统推理引擎的做法是找不到满足条件的 kernel就回退到一个差不多的通用实现。TokenSpeed 在rl-bitwise下拒绝这种回退找不到batch_invariant特性的 kernel →启动报错找不到声明combine_order的 MoE kernel →启动报错模型 profile 未声明支持rl-bitwise→启动报错量化 checkpoint没有批不变量化 GEMM→启动报错这个设计确保如果你能用rl-bitwise跑起来那结果一定是位级确定的。没有差不多没有大概率一致。5️⃣ 按模型验证不是所有模型都能用rl-bitwise不是全局开关而是按模型逐个验证的。每个模型有一个ModelProfile其中numerics_envelopes字段列出该模型通过了哪些数值守络的验证ModelProfile.numerics_envelopes (auto, rl-bitwise) # 验证通过后才能列出验证需要过两套测试不变性测试同一 prompt 在 bs1、随机 co-batch、重复运行下token id 和 logprob 必须torch.equal训练器对齐测试训练框架对prompt response做 teacher-forced 前向dump 出 logprob推理引擎用相同 ids 做 teacher-forced 前向logprob 必须逐位一致验证逻辑实现在 model_profile.py 和 numerics.py 的require_verified_numerics函数中。6️⃣ 快速上手一条命令开启位级确定性开启非常简单——在启动命令中加一个参数tokenspeed serve your-model \ --numerics rl-bitwise \ --host 0.0.0.0 \ --port 8000TokenSpeed 会自动收紧所有确定性开关关闭 autotune / TF32 / PDL / fused all-reduce将 MoE backend 固定为aok批不变算子套件将采样流切换为per-request将 logprob 顺序切换为megatron将 MoE 合并顺序切换为slot将 DSA slot 顺序切换为sorted拒绝任何冲突的显式参数详细启动参数参考 launching.md 和 server.md。7️⃣ 总结为什么这个设计重要痛点TokenSpeed 的解法推理引擎结果不可复现rl-bitwise承诺运行不变 批不变推理和训练的 logprob 对不上trainer alignment 层逐操作对齐多个开关手动拼容易漏envelope 一次收紧冲突即报错静默回退导致以为确定性但实际没有无静默回退缺 kernel 就启动失败不同模型确定性能力不同按模型 profile 验证未验证的模型拒绝启动对于正在用 RL 微调大模型、且依赖推理引擎生成 rollout 的团队--numerics rl-bitwise是一个值得关注的特性——它让推理引擎从尽量快变成又快又确定让 RL 训练信号真正可信赖。 完整设计文档docs/design/numerics.md【免费下载链接】tokenspeedTokenSpeed is a speed-of-light LLM inference engine.项目地址: https://gitcode.com/gh_mirrors/to/tokenspeed创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考