1. 从一次训练崩溃说起为什么 Bitwise 一致这么难做过强化学习训练的人大概都经历过这种场景同一个模型、同一份数据、同一组超参跑两次 loss 曲线就是对不上rollout 阶段采样出来的轨迹和训练阶段重算的 logprob 差了一点点单步看不出来累积几百步之后策略直接跑偏。你盯着屏幕怀疑人生怀疑随机种子怀疑数据加载顺序最后发现是 kernel 层面的浮点累加顺序不一样。这就是Bitwise 一致逐位一致要解决的问题。它不是误差很小而是二进制层面完全相同。在 RL 训练里rollout 和 training 是两个阶段前者用推理引擎跑采样后者用训练框架算梯度。如果这两个阶段对同一个 token 序列算出来的 logprob 在 bit 级别有差异那么重要性采样比importance ratio就会偏离 1PPO 这类算法的 clip 机制会被这些虚假差异触发训练稳定性直接崩掉。我这次折腾的项目核心目标就一句话在 AMD 平台上用 RL-Kernel 配合 vime 框架让训练阶段和 rollout 阶段的前向计算做到 bitwise 一致。关键词拆开看RL-Kernel 是算子层的东西vime 是 RL 训练框架AMD 是硬件平台ROCm 是软件栈。这四个东西凑一起坑比想象中多得多。先说清楚适合谁看。如果你正在做 RLHF、GRPO、PPO 这类需要 rollout 和 training 分离的训练任务并且用的是 AMD 的卡比如 MI 系列或者消费级的 Radeon那这篇内容基本就是给你写的。如果你只是跑跑 SFT或者用的是单阶段训练那 bitwise 一致这个问题对你影响没那么大但了解一下 kernel 层面的数值行为也没坏处。小白读者也不用慌我会把浮点、kernel、ROCm 这些概念用生活化的方式讲清楚保证你能跟上。2. 整体设计思路为什么要在 AMD 上死磕这件事2.1 问题的本质rollout 和 training 是两套计算路径先把这个问题的根源讲透。RL 训练里rollout 阶段通常用推理引擎比如 vLLM 那套思路的引擎来生成轨迹追求的是吞吐training 阶段用训练框架PyTorch 各种并行策略来算梯度追求的是精度和显存效率。这两个阶段的算子实现往往是两套完全不同的代码。推理引擎为了快会把 attention 做成 paged attention把 KV cache 管理得花里胡哨矩阵乘可能用特定的 fused kernel训练框架为了稳用的是另一套 attention 实现可能还带 flash attention 的变体。这两套实现在数学上等价但在浮点运算上不等价——因为浮点加法不满足结合律。举个生活化的例子你要把 100 个小数加起来。从左往右加和从右往左加或者先分组加再合并结果在实数域上一样但在浮点数域上可能差最后几位。GPU kernel 的并行归约reduction就是典型的分组加法不同的线程块划分、不同的 warp 调度归约顺序就不同最后那几位就不一样。注意bitwise 一致不是要求所有算子都一模一样而是要求影响 logprob 计算的那条关键路径上rollout 和 training 的浮点运算顺序完全一致。其他不影响 logprob 的部分比如 KV cache 的存储布局可以不同。2.2 为什么选 RL-Kernel 作为切入点RL-Kernel 这个思路的核心是把 rollout 和 training 共用的那部分计算抽成同一套 kernel。与其让两个框架各自实现一遍然后祈祷它们数值一致不如让它们调用同一个底层算子。这个设计的好处很直接消除实现差异同一个 kernel同样的线程划分同样的归约顺序输入相同则输出 bitwise 相同。便于验证只需要验证一套 kernel 的正确性不用交叉验证两套实现。性能可控kernel 层面可以做针对性的优化比如针对 AMD 的 wavefront 大小64来调整线程块。代价也很明显这套 kernel 必须同时满足推理的高吞吐和训练的高精度需求设计上要做取舍。比如推理可以用 fp16 甚至 fp8 来加速但训练如果也用低精度梯度就会炸。所以 RL-Kernel 通常会在关键路径上用 fp32 累加即使输入是低精度。2.3 AMD 平台带来的额外变量为什么在 AMD 上这件事更难因为 ROCm 生态和 CUDA 生态有差异很多 kernel 的默认行为不一样。第一wavefront 大小不同。NVIDIA 的 warp 是 32 线程AMD 的 wavefront 是 64 线程RDNA 架构有 32 和 64 两种模式。这意味着同样的归约算法在 AMD 上的分组方式和 NVIDIA 不同浮点累加顺序自然不同。如果你直接拿 CUDA 的 kernel 翻译成 HIP数值行为可能就变了。第二数学库实现不同。rocBLAS、rocFFT、MIOpen 这些库和 cuBLAS、cuDNN 的实现细节不一样。比如 exp、log、tanh 这些超越函数的近似算法可能不同最后几位就有差异。第三编译器优化行为不同。AMD 的编译器比如 hipcc 背后的 LLVM在做 fast-math 优化时对浮点重排的策略和 NVCC 不一样。-ffast-math这种选项在两边打开结果可能差很多。所以整个项目的设计思路是先用 RL-Kernel 统一关键路径的算子实现再针对 AMD 的硬件特性做适配最后用 bitwise 对比工具验证 rollout 和 training 的输出。这三步缺一不可。3. 核心细节解析Bitwise 一致到底卡在哪些环节3.1 浮点累加顺序最隐蔽的杀手前面提到浮点加法不满足结合律这里展开讲。假设你要算a b c d四种可能的顺序((ab)c)d (a(bc))d (ab)(cd) a((bc)d)在实数域上结果一样在 IEEE 754 浮点下可能四个结果都不同。GPU 上的归约操作比如 softmax 里的求和、attention 里的加权平均、layer norm 里的均值方差全都是归约。归约的并行实现通常是树形的先把数据分到多个线程每个线程局部求和然后再跨线程合并。线程数、每个线程处理多少元素、合并的顺序任何一个变了结果就变。在 RL 训练里最敏感的归约是logprob 的计算。logprob 通常是对 logits 做 log_softmax 然后取对应 token 的值。log_softmax 里有 max 归约和 sum 归约这两个归约的顺序直接决定 logprob 的最后几位。我实测过一个案例同一个模型rollout 用推理引擎算出的 logprob 和 training 用训练框架算出的 logprob在 fp16 下差异能到 1e-3 量级在 fp32 下能到 1e-6。听起来很小但 PPO 的 ratio 是exp(logprob_new - logprob_old)1e-6 的差异经过 exp 放大再累积几千个 tokenratio 就能偏离 1 好几个百分点。提示判断你的训练是否受这个问题影响最简单的办法是打印 rollout 和 training 对同一批数据的 logprob算一下 max abs diff。如果大于 1e-5基本可以确定是 kernel 层面的数值不一致。3.2 RL-Kernel 的关键设计统一归约树RL-Kernel 解决这个问题的核心手段是固定归约树。具体做法对于所有影响 logprob 的归约操作强制使用同一套线程划分和归约顺序。归约的中间结果用 fp32 存储即使输入是 fp16/bf16。禁用编译器的 fast-math 重排或者用#pragma显式指定归约顺序。这里有个细节AMD 的 wavefront 是 64 线程做 warp-level reduction 的时候shuffle 操作的语义和 NVIDIA 不同。RL-Kernel 在 AMD 上的实现需要针对 wavefront 大小做特化。比如一个 64 线程的 wavefront 做归约通常是先做 32 对 32 的 shuffle再做 16 对 16以此类推。这个顺序必须固定下来不能依赖编译器的自动优化。我踩过的一个坑一开始用__shfl_down做归约在 NVIDIA 上没问题移植到 AMD 上用__shfl_down的 HIP 版本结果因为 wavefront 是 64归约的步数和 NVIDIA 不一样数值就对不上了。后来改成显式指定归约步数并且用__builtin_amdgcn_...这类 intrinsic 来保证顺序才搞定。3.3 vime 框架的接入点vime 是 RL 训练框架它的职责是调度 rollout 和 training 两个阶段。要让 bitwise 一致生效vime 需要在两个地方做改造第一rollout 阶段的 logprob 计算必须走 RL-Kernel。推理引擎默认用自己的 kernel 算 logprobvime 需要把这个计算替换成 RL-Kernel 的调用。这涉及到和推理引擎的接口对接通常是在采样完成后把 token 序列和 hidden states 传给 RL-Kernel 重算一遍 logprob。第二training 阶段的 logprob 计算也必须走 RL-Kernel。训练框架默认用 PyTorch 的log_softmaxvime 需要把这个替换成 RL-Kernel 的实现。这里要注意训练阶段有反向传播RL-Kernel 需要同时提供前向和反向的实现且反向的数值行为也要和 rollout 一致虽然 rollout 不需要反向但为了验证一致性前向必须一致。vime 的接入点通常在 policy 的 forward 函数里。我的做法是写一个 wrapper把原来的log_softmax调用替换成rl_kernel.log_softmax然后在配置里加一个开关方便对比开启前后的差异。3.4 AMD ROCm 的适配细节ROCm 这边有几个必须注意的点第一版本匹配。ROCm 的版本和 PyTorch 的版本有严格的对应关系。比如 PyTorch 2.1 通常配 ROCm 5.6 或 5.7PyTorch 2.2 配 ROCm 5.7 或 6.0。版本不匹配会导致 kernel 编译失败或者数值行为异常。我建议用官方推荐的组合不要自己乱配。第二编译选项。hipcc 的-O3和-ffast-math会影响浮点行为。为了 bitwise 一致我建议在关键 kernel 上禁用 fast-math用-fno-fast-math或者更细粒度的-ffp-contractoff来控制 FMAfused multiply-add的生成。FMA 会把a*bc合成一条指令精度更高但和分开算的结果不同。如果 rollout 用了 FMA 而 training 没用bitwise 就不一致。第三内存对齐。AMD 的 GPU 对内存对齐比较敏感不对齐的访问可能导致 kernel 走不同的代码路径进而影响数值。RL-Kernel 在分配显存时要注意对齐到 256 字节或更高。第四MIOpen 的算法选择。MIOpen 是 AMD 的深度学习库它有很多算法可选比如 GEMM 的不同实现。不同算法的数值行为可能不同。为了 bitwise 一致需要固定 MIOpen 的算法选择可以通过环境变量MIOPEN_FIND_MODEnormal和MIOPEN_DEBUG_FIND_ONLY_SOLVER来锁定。4. 实操过程从环境搭建到 bitwise 验证4.1 环境准备与版本锁定先把环境搭起来。我用的是 AMD MI210ROCm 6.0PyTorch 2.2。这套组合在我这边跑下来比较稳。# 检查 ROCm 版本 rocm-smi # 输出里能看到 ROCm 版本和 GPU 型号 # 检查 PyTorch 是否识别到 GPU python -c import torch; print(torch.__version__); print(torch.cuda.is_available()); print(torch.cuda.get_device_name(0)) # 注意PyTorch 在 ROCm 上仍然用 cuda 这个命名空间这是历史遗留版本锁定很关键。我建议用 conda 或者 docker 来固定环境避免系统级的库版本冲突。Docker 的话用rocm/pytorch:rocm6.0_ubuntu22.04_py3.10_pytorch_2.2这个官方镜像比较省事。注意如果你用的是消费级 Radeon 卡比如 RX 7900ROCm 的支持不如 MI 系列完整可能需要额外的环境变量来启用。比如HSA_OVERRIDE_GFX_VERSION11.0.0这种具体值取决于你的卡。这个操作有风险可能导致不稳定建议先在测试环境验证。4.2 RL-Kernel 的编译与集成RL-Kernel 通常是一个独立的库需要单独编译。编译流程大致是# 克隆 RL-Kernel 仓库假设你已经有了 cd rl-kernel mkdir build cd build cmake .. -DCMAKE_BUILD_TYPERelease -DGPU_TARGETSgfx90a # gfx90a 是 MI210 的架构代号其他卡要改 make -j$(nproc)GPU_TARGETS这个参数很重要它决定了编译出来的 kernel 针对哪个架构优化。填错了会导致 kernel 跑不起来或者性能极差。常见的架构代号GPU 型号架构代号MI210gfx90aMI250gfx90aMI300gfx942RX 7900 XTXgfx1100RX 6900 XTgfx1030编译完成后把生成的.so文件放到 Python 的 path 里或者用pip install -e .安装 Python 绑定。集成到 vime 的时候我写了一个简单的 wrapperimport rl_kernel class BitwiseLogSoftmax(torch.autograd.Function): staticmethod def forward(ctx, logits, dim-1): # 调用 RL-Kernel 的 log_softmax output rl_kernel.log_softmax(logits, dim) ctx.save_for_backward(output) ctx.dim dim return output staticmethod def backward(ctx, grad_output): output, ctx.saved_tensors # RL-Kernel 的反向实现 grad_input rl_kernel.log_softmax_backward(grad_output, output, ctx.dim) return grad_input, None然后在 vime 的 policy forward 里把原来的F.log_softmax替换成BitwiseLogSoftmax.apply。4.3 关键参数的计算与选择这里讲几个必须算清楚的参数。第一归约的 block size。假设你要对长度为 4096 的序列做 sum 归约用多大的 block我的经验是block size 取 256每个线程处理 16 个元素然后做 block-level reduction。这个配置在 AMD 上比较平衡既不会因为 block 太小导致归约步数太多也不会因为 block 太大导致 occupancy 下降。计算过程4096 / 256 16每个线程处理 16 个元素。线程内的归约用顺序累加保证顺序固定block 内的归约用树形log2(256) 8 步。这个归约树是固定的不管输入怎么变顺序都一样。第二fp32 累加的精度。fp32 有 24 位尾数能精确表示约 1600 万个整数。对于 logprob 计算累加的元素数量通常在几千到几万fp32 累加的误差在 1e-6 量级足够 bitwise 一致。如果序列特别长比如 128k可能需要用 Kahan summation 或者 double 累加但那样性能会掉很多。第三FMA 的取舍。FMA 把a*bc合成一条指令精度比分开算高但和分开算的结果不同。为了 bitwise 一致要么两边都用 FMA要么两边都不用。我选择在关键路径上禁用 FMA用-ffp-contractoff编译这样 rollout 和 training 的行为一致。4.4 Bitwise 验证的实操流程验证是最后一步也是最关键的一步。我的验证流程准备固定输入。用固定的随机种子生成一批 logits保存成文件。这样 rollout 和 training 用的是完全相同的输入。分别跑 rollout 和 training 的 logprob 计算。rollout 走推理引擎 RL-Kerneltraining 走训练框架 RL-Kernel。对比输出的二进制。不要用torch.allclose那个有容差。要用torch.equal或者直接对比字节。import torch # 加载两个阶段的输出 logprob_rollout torch.load(logprob_rollout.pt) logprob_training torch.load(logprob_training.pt) # bitwise 对比 if torch.equal(logprob_rollout, logprob_training): print(Bitwise 一致通过) else: diff (logprob_rollout - logprob_training).abs() print(f最大差异: {diff.max().item()}) print(f差异位置: {diff.argmax().item()}) # 找到第一个不一致的位置方便排查 mismatch (logprob_rollout ! logprob_training).nonzero() print(f不一致元素数量: {len(mismatch)})如果不一致逐层排查。先对比 logits再对比 log_softmax 的中间结果max、sum最后定位到具体的 kernel。我实测下来第一次跑通常是不一致的差异在 1e-7 到 1e-6 之间。排查下来最常见的原因是 FMA 的生成不一致或者归约的 block size 在两边不同。把这两个固定住之后基本就能做到 bitwise 一致。5. 常见问题与排查技巧实录5.1 问题速查表现象可能原因排查方法解决方案logprob 差异 1e-6 量级FMA 生成不一致对比编译选项检查-ffp-contract两边统一禁用或启用 FMAlogprob 差异 1e-3 量级精度不一致fp16 vs fp32检查两边的 dtype关键路径统一用 fp32 累加差异随序列长度增大归约顺序不一致对比 block size 和归约树固定归约树统一 block size某些 batch 一致某些不一致内存对齐问题检查输入 tensor 的 stride强制对齐到 256 字节kernel 编译失败架构代号填错检查GPU_TARGETS用rocminfo查正确的 gfx 代号性能极差block size 太小或太大用 rocprof 分析调整到 256 或 512MIOpen 算法不稳定算法自动选择设置MIOPEN_FIND_MODE锁定 solver5.2 独家避坑技巧技巧一用torch.use_deterministic_algorithms(True)做初步筛查。这个选项会让 PyTorch 在遇到不确定的算子时报错能帮你快速定位哪些算子有问题。但注意它不能保证 bitwise 一致只能保证可复现。技巧二保存中间结果做对比。不要只对比最终的 logprob把 max、sum、exp 这些中间结果都存下来。这样一旦不一致能快速定位到是哪一步出的问题。我通常会在 kernel 里加一个 debug 模式把中间结果 dump 到文件。技巧三用小的输入做快速验证。不要一上来就用完整的 batch 和长序列先用 batch1、seq_len128 的小输入验证。小输入下归约步数少问题更容易暴露。验证通过后再逐步放大。技巧四注意 ROCm 的HSA_ENABLE_SDMA环境变量。这个变量控制是否用 SDMA 引擎做数据传输不同设置下数据传输的顺序可能不同进而影响数值。为了 bitwise 一致建议固定这个变量的值。技巧五记录每次实验的完整环境。包括 ROCm 版本、PyTorch 版本、编译选项、环境变量。bitwise 一致很容易被环境变化打破记录清楚才能快速复现。5.3 一个真实的排查案例有一次我遇到一个诡异的问题同样的代码在 MI210 上 bitwise 一致在 MI250 上就不一致。排查了半天发现是 MI250 是双 die 设计跨 die 的通信走的是不同的路径归约的时候如果跨了 die顺序就和单 die 不同。解决方案是在 kernel 里显式指定 die 的亲和性让归约在单个 die 内完成。这个操作比较底层需要用hipSetDevice和hipDeviceSetLimit这类 API。具体代码就不贴了思路是把数据先按 die 分组每组内部做归约最后再合并。这样归约树就固定了不管单 die 还是双 die结果都一样。这个坑让我意识到AMD 的硬件架构差异比想象中大。MI 系列和 Radeon 系列的差异、单 die 和双 die 的差异都会影响数值行为。做 bitwise 一致的时候必须把这些因素都考虑进去。6. 性能与一致性的平衡我的取舍经验6.1 一致性的代价做到 bitwise 一致不是没有代价的。最直接的代价是性能。第一禁用 FMA 会让矩阵乘的性能下降 10% 到 20%。FMA 是 GPU 上最常用的指令之一禁用之后计算吞吐会明显下降。第二固定归约树会限制 kernel 的优化空间。本来可以根据输入大小动态调整 block size现在必须固定小输入的时候可能浪费线程大输入的时候可能不够用。第三fp32 累加比 fp16 累加慢。fp16 的累加吞吐是 fp32 的两倍但为了精度只能用 fp32。我实测的数据在一个 7B 模型的 RL 训练里开启 bitwise 一致后rollout 的吞吐下降了约 15%training 的吞吐下降了约 8%。整体训练时间增加了约 10%。6.2 哪些地方可以放松不是所有地方都需要 bitwise 一致。我的经验是只对影响 logprob 的关键路径严格要求其他部分可以放松。具体来说log_softmax 的 max 和 sum 归约必须 bitwise 一致。attention 的加权平均如果 attention 的输出会影响 logits那也必须一致。但如果 attention 用的是 flash attention 这类近似算法可能需要换成精确实现。layer norm 的均值方差影响 hidden states间接影响 logits建议一致。embedding lookup只是查表不涉及浮点运算天然一致。激活函数exp、silu 这些如果实现不同会有差异建议统一。放松的地方KV cache 的存储布局不影响数值可以不同。数据传输的顺序只要最终结果一致中间怎么传无所谓。非关键路径的算子比如一些辅助计算的 kernel。6.3 一个折中方案如果性能实在受不了可以考虑混合方案rollout 阶段用快速 kerneltraining 阶段用精确 kernel但在两者之间加一个校正步骤。具体做法rollout 算出的 logprob 存下来training 阶段重算 logprob 后用 rollout 的值做一次校正。校正的方式可以是直接替换或者用某种插值。这样 training 的梯度是基于 rollout 的 logprob 算的避免了不一致带来的偏差。这个方案的缺点是校正本身有开销而且如果 rollout 的 logprob 本身有误差校正后还是有误差。所以它只是权宜之计不是根本解决方案。7. 后续可以扩展的方向这套东西跑通之后我发现还有几个方向可以继续挖。第一扩展到更多算子。目前只搞定了 log_softmax 和相关的归约attention、layer norm 这些还没完全覆盖。如果能把整个 transformer 的前向都做到 bitwise 一致那 RL 训练的稳定性会有质的提升。第二自动化验证工具。现在验证还是手工的每次改代码都要重新跑一遍对比。可以做一个 CI 流程每次提交代码自动跑 bitwise 验证不一致就报错。第三跨平台一致性。现在只保证了 AMD 平台内部的一致如果 rollout 在 AMD 上、training 在别的平台上能不能也做到一致这个难度更大因为不同平台的浮点行为差异更大。但理论上只要归约树和精度都固定是可以做到的。第四性能优化。现在为了 bitwise 一致牺牲了 10% 的性能能不能通过更好的 kernel 设计把这个损失补回来比如用更高效的归约算法或者用 AMD 特有的指令集比如 MFMA来加速。我个人在实际操作中的体会是bitwise 一致这件事说起来简单做起来全是细节。每一个归约、每一次精度转换、每一个编译选项都可能成为不一致的来源。但只要把关键路径固定住把验证流程建起来这件事是可以做到的。最重要的是做完之后训练稳定性确实有肉眼可见的提升之前那些莫名其妙的 loss 尖刺基本消失了。这个收益值得那 10% 的性能代价。