【训练营】BF16 Kernel 优化 - #227
Open
ProMetaDev wants to merge 1 commit into
Open
ProMetaDev wants to merge 1 commit into
ProMetaDev wants to merge 1 commit into
Conversation
…ccesses 逐元素算子（元素级一元/二元前向与反向）此前是每线程处理 1 个元素的标量循环： BF16 下等于每 16 字节发出 8 次访存，而这些算子在 GPT-2 BF16 训练中占相当比例 （内置 profiler 显示 elementwise + cast 合计超过设备时间的一半，其中 Cast 每步 123 次、 AccumulateGrad 221 次调用）。上游已为「二元反向」提供 128-bit 向量化实现，本次把 「一元前向 / 一元反向 / 二元前向」也向量化，形成一致的访存策略。 实现要点： - 新增 CanVectorize<T>()：要求所有非空指针 16 字节对齐且元素数 >= kVecSize<T>； 不满足时自动退回原标量 kernel，保证对任意 view/非对齐张量安全。 - 新增 UnaryForwardKernelVectorized / UnaryBackwardKernelVectorized / BinaryForwardKernelNoBroadcastVectorized，尾部不足一个向量的元素按标量处理。 - 逐元素表达式与运算顺序完全不变，数值结果一致（训练 loss 前 8 步逐位相同）。 - CMakeLists.txt：CUDA 目标架构由写死的 75;80;90 改为可配置的 INFINI_TRAIN_CUDA_ARCHS （默认值不变），便于在新架构上编译，或追加 *-virtual 产出 PTX 交给驱动 JIT。 - 新增 tools/bench_elem.cu：不依赖框架的隔离微基准，用于测量这些访存模式的 有效带宽与 kernel 启动开销（报告数据的来源）。 实测（RTX 5070 Ti Laptop，GPT-2 124M，batch 80 × seq 64）： - 一元算子 1.22~1.33×、二元前向 1.04~1.06×、FP32 1.00×（无退化）； - 端到端 320.6 -> 320.3 ms/step（-0.2%，在噪声范围内）。逐元素算子实测已达 390~590 GB/s，接近本机显存带宽上限，因此单核向量化的收益有限； 报告 docs/bf16_kernel_optimization.md 中分析了真正的瓶颈与后续方向 （算子融合以减少 kernel 数量、消除 BF16<->FP32 的 Cast、AccumulateGrad 批量化）。 - 仓库自带 gtest 套件全部通过（279/279），改动文件通过 clang-format 检查。
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
对 infini_train 的 elementwise kernel 做 128-bit 向量化。
BinaryForwardKernelNoBroadcastVectorized 三个向量化 kernel，含对齐判定