【训练营】BF16 Kernel 优化

.gitmodules

@@ -1,12 +1,12 @@
 [submodule "third_party/glog"]
 	path = third_party/glog
-	url = git@github.com:google/glog.git
+	url = https://github.com/google/glog.git
 [submodule "third_party/gflags"]
 	path = third_party/gflags
-	url = git@github.com:gflags/gflags.git
+	url = https://github.com/gflags/gflags.git
 [submodule "third_party/eigen"]
 	path = third_party/eigen
-	url = git@github.com:InfiniTensor/eigen-mirror.git
+	url = https://github.com/InfiniTensor/eigen-mirror.git
 [submodule "third_party/googletest"]
 	path = third_party/googletest
-	url = git@github.com:google/googletest.git
+	url = https://github.com/google/googletest.git

CMakeLists.txt

@@ -104,7 +104,8 @@ if(USE_CUDA)
   file(GLOB_RECURSE CUDA_KERNELS ${PROJECT_SOURCE_DIR}/infini_train/src/*.cu)
 
   add_library(infini_train_cuda_kernels STATIC ${CUDA_KERNELS})
-  set_target_properties(infini_train_cuda_kernels PROPERTIES CUDA_ARCHITECTURES "75;80;90")
+  set(INFINITRAIN_CUDA_ARCHITECTURES "75;80;90" CACHE STRING "CUDA architectures for kernel compilation")
+  set_target_properties(infini_train_cuda_kernels PROPERTIES CUDA_ARCHITECTURES "${INFINITRAIN_CUDA_ARCHITECTURES}")
 
   target_link_libraries(infini_train_cuda_kernels
     PUBLIC

bench/binary_backward_bench.cc

@@ -0,0 +1,115 @@
+// Microbenchmark for BinaryBackward (Mul/Add backward) kernels: BF16 vs FP32.
+// Mirrors bench/torch_binary_backward_bench.py for PyTorch comparison.
+#include <chrono>
+#include <cmath>
+#include <cstdio>
+#include <memory>
+#include <vector>
+
+#include <cuda_runtime_api.h>
+
+#include "infini_train/include/autograd/elementwise.h"
+#include "infini_train/include/device.h"
+#include "infini_train/include/tensor.h"
+
+using namespace infini_train;
+
+namespace {
+// Pin the CUDA async mempool so it never releases cached blocks back to the OS. With the default
+// release threshold (0), every profiler stream-sync (PROFILE_MODE build) unmaps all cached memory
+// and the next iteration re-maps ~200MB of physical pages — the benchmark would measure WSL2 page
+// mapping speed (~7ms/iter) instead of the kernels. PyTorch's caching allocator never unmaps, so
+// this also keeps the comparison fair.
+void PinMemPool() {
+    cudaMemPool_t pool = nullptr;
+    if (cudaDeviceGetDefaultMemPool(&pool, 0) != cudaSuccess || pool == nullptr) {
+        return;
+    }
+    uint64_t threshold = ~0ull;
+    cudaMemPoolSetAttribute(pool, cudaMemPoolAttrReleaseThreshold, &threshold);
+}
+
+// Check that every element of t equals expect (relative error within tol). Returns max rel error.
+double MaxRelError(const std::shared_ptr<Tensor> &t, double expect) {
+    auto host = t->To(DataType::kFLOAT32).To(Device());
+    const float *data = static_cast<const float *>(host.DataPtr());
+    double max_err = 0.0;
+    for (size_t i = 0; i < host.NumElements(); ++i) {
+        const double err = std::abs(static_cast<double>(data[i]) - expect) / std::max(std::abs(expect), 1e-12);
+        max_err = std::max(max_err, err);
+    }
+    return max_err;
+}
+} // namespace
+
+static float TimeBackward(std::function<void()> fn, int warmup, int iters) {
+    for (int i = 0; i < warmup; ++i) { fn(); }
+    cudaDeviceSynchronize();
+    cudaEvent_t start, stop;
+    cudaEventCreate(&start);
+    cudaEventCreate(&stop);
+    cudaEventRecord(start);
+    for (int i = 0; i < iters; ++i) { fn(); }
+    cudaEventRecord(stop);
+    cudaEventSynchronize(stop);
+    float ms = 0.0f;
+    cudaEventElapsedTime(&ms, start, stop);
+    cudaEventDestroy(start);
+    cudaEventDestroy(stop);
+    return ms * 1000.0f / iters; // us per iter
+}
+
+template <typename Op>
+static void RunCase(const char *op_name, std::vector<int64_t> a_dims, std::vector<int64_t> b_dims, DataType dtype) {
+    constexpr bool kIsMul = std::is_same_v<Op, autograd::Mul>;
+    auto dev = Device(Device::DeviceType::kCUDA, 0);
+    auto a = std::make_shared<Tensor>(a_dims, dtype, dev, true);
+    a->Fill(2.0f);
+    auto b = std::make_shared<Tensor>(b_dims, dtype, dev, true);
+    b->Fill(3.0f);
+    auto op = std::make_shared<Op>();
+    auto out = op->Apply({a, b});
+    auto grad = std::make_shared<Tensor>(a_dims, dtype, dev, true);
+    grad->Fill(1.0f);
+
+    // One-shot correctness check. a=2, b=3, grad=1:
+    //   mul: ga = g*b = 3;        gb = g*a = 2 per use (row-bcast: rows*2, col-bcast: cols*2)
+    //   add: ga = g = 1;          gb = g = 1 per use   (row-bcast: rows,   col-bcast: cols)
+    {
+        auto grads = op->Backward({grad});
+        const double ga_expect = kIsMul ? 3.0 : 1.0;
+        const double unit = kIsMul ? 2.0 : 1.0;
+        double gb_expect = unit;
+        if (a_dims != b_dims) {
+            gb_expect = unit * (b_dims.size() == 1 ? a_dims[0] : a_dims[1]);
+        }
+        const double tol = dtype == DataType::kFLOAT32 ? 1e-5 : 1e-2;
+        const double err_a = MaxRelError(grads[0], ga_expect);
+        const double err_b = MaxRelError(grads[1], gb_expect);
+        const bool pass = err_a <= tol && err_b <= tol;
+        printf("  correctness: %s (ga max rel err %.2e, gb max rel err %.2e, tol %.0e)\n", pass ? "PASS" : "FAIL",
+               err_a, err_b, tol);
+    }
+
+    const float us = TimeBackward([&]() { auto g = op->Backward({grad}); }, 20, 200);
+    const double bytes = 5.0 * a->NumElements() * (dtype == DataType::kFLOAT32 ? 4 : 2); // g,a,b in; ga,gb out (approx)
+    printf("[%-5s] a=[%ld,%ld] b_bcast=%-5s %-8s %8.1f us  (~%.0f GB/s)\n", op_name, a_dims[0], a_dims[1],
+           a_dims == b_dims ? "false" : "true", dtype == DataType::kFLOAT32 ? "float32" : "bfloat16", us,
+           bytes / (us * 1e-6) / 1e9);
+}
+
+int main() {
+    google::InitGoogleLogging("binary_backward_bench");
+    PinMemPool();
+    const std::vector<std::pair<int64_t, int64_t>> shapes = {{65536, 768}, {8192, 3072}};
+    for (auto [r, c] : shapes) {
+        for (DataType dt : {DataType::kBFLOAT16, DataType::kFLOAT32}) {
+            RunCase<autograd::Mul>("mul", {r, c}, {r, c}, dt);
+            RunCase<autograd::Mul>("mul", {r, c}, {c}, dt);    // row-broadcast (bias style)
+            RunCase<autograd::Mul>("mul", {r, c}, {r, 1}, dt); // col-broadcast
+            RunCase<autograd::Add>("add", {r, c}, {r, c}, dt);
+            RunCase<autograd::Add>("add", {r, c}, {c}, dt);
+        }
+    }
+    return 0;
+}

bench/torch_binary_backward_bench.py

@@ -0,0 +1,51 @@
+"""PyTorch baseline for binary-op backward kernels, BF16 vs FP32.
+
+Covers the patterns that hit InfiniTrain's BinaryBackward:
+  1. no-broadcast elementwise (mul / add) on [B*T, C]
+  2. broadcast backward where B is a row vector [C] (bias-add style)
+"""
+
+import torch
+import torch.utils.benchmark as tb
+
+DEV = "cuda"
+
+
+def bench(fn, n=200):
+    for _ in range(20):
+        fn()
+    torch.cuda.synchronize()
+    t = tb.Timer(stmt="fn()", globals={"fn": fn}).timeit(n)
+    return t.median * 1e6  # us
+
+
+def make_case(rows, cols, dtype, op, broadcast_b):
+    a = torch.randn(rows, cols, device=DEV, dtype=dtype, requires_grad=True)
+    if broadcast_b:
+        b = torch.randn(cols, device=DEV, dtype=dtype, requires_grad=True)
+    else:
+        b = torch.randn(rows, cols, device=DEV, dtype=dtype, requires_grad=True)
+    g = torch.randn(rows, cols, device=DEV, dtype=dtype)
+
+    def fn():
+        if a.grad is not None:
+            a.grad = None
+        if b.grad is not None:
+            b.grad = None
+        out = (a * b) if op == "mul" else (a + b)
+        out.backward(g)
+
+    return fn
+
+
+print(f"device: {torch.cuda.get_device_name(0)}")
+print(f"{'case':<28}{'dtype':<10}{'op':<6}{'bcastB':<8}{'us':>10}")
+for rows, cols in [(65536, 768), (8192, 3072), (262144, 768)]:
+    for dtype in (torch.bfloat16, torch.float32):
+        for op in ("mul", "add"):
+            for bcast in (False, True):
+                fn = make_case(rows, cols, dtype, op, bcast)
+                us = bench(fn)
+                print(
+                    f"[{rows:>6},{cols:>4}]  {str(dtype).split('.')[-1]:<10}{op:<6}{str(bcast):<8}{us:>10.1f}"
+                )

infini_train/include/autocast.h

@@ -76,6 +76,31 @@ inline constexpr std::array<DataType, static_cast<size_t>(Device::DeviceType::kC
     DataType::kFLOAT16,  // CUDA.
 };
 
+// Thread-local cache of autocast-casted leaf parameters (e.g. FP32 master weights
+// demoted to BF16/FP16 by the kLowerPrecision policy). Keyed by the source tensor's
+// address; the weak_ptr guards against address reuse after the source is freed.
+//
+// Threading model: training is single-threaded per rank, so the cache is
+// thread_local (same as tls_autocast_context) and needs no locking. The
+// invalidation entry points below must be called on the thread that mutates the
+// parameter (all in-place parameter mutations in this codebase -- optimizer
+// kernels, CopyFrom/Fill/SetData, checkpoint LoadStateDict -- run on the training
+// thread). Raw writes through Tensor::DataPtr() bypass these hooks, as do
+// initializer helpers; both only happen before training starts in practice.
+struct AutocastWeightCacheEntry {
+    std::weak_ptr<Tensor> source; // detects source destruction / pointer reuse
+    DataType target_dtype;
+    std::shared_ptr<Tensor> casted;
+};
+
+inline thread_local std::unordered_map<const Tensor *, AutocastWeightCacheEntry> tls_autocast_weight_cache;
+
+// Drop the cached cast of one tensor (call after any in-place mutation of it).
+inline void InvalidateAutocastWeightCacheEntry(const Tensor *tensor) { tls_autocast_weight_cache.erase(tensor); }
+
+// Drop the whole cache (e.g. after bulk parameter replacement).
+inline void ClearAutocastWeightCache() { tls_autocast_weight_cache.clear(); }
+
 // Thread-local context to track autocast state
 struct AutocastContext {
     bool enabled = false;                                      // Whether autocast is active in the current thread
@@ -113,6 +138,12 @@ struct AutocastContext {
             }
         };
 
+        // Only kLowerPrecision casts of FP32 leaf parameters are cacheable: those
+        // tensors are owned by the module (stable address) and are re-cast every
+        // forward without this cache. Non-leaf activations are short-lived, so keying
+        // on their address would risk stale hits after allocator reuse.
+        const bool cache_weights = policy == CastPolicy::kLowerPrecision;
+
         auto cast_arg = [&](auto &arg) {
             using T = std::decay_t<decltype(arg)>;
             if constexpr (std::is_same_v<T, std::shared_ptr<Tensor>>) {
@@ -121,7 +152,27 @@ struct AutocastContext {
                     if (is_floating_point(current_dtype)) {
                         DataType target_dtype = get_target_dtype();
                         if (current_dtype != target_dtype) {
-                            arg = std::make_shared<Tensor>(arg->To(target_dtype));
+                            if (cache_weights && current_dtype == DataType::kFLOAT32 && arg->is_leaf()
+                                && arg->requires_grad()) {
+                                auto it = tls_autocast_weight_cache.find(arg.get());
+                                if (it != tls_autocast_weight_cache.end()) {
+                                    auto source = it->second.source.lock();
+                                    if (source && source.get() == arg.get()
+                                        && it->second.target_dtype == target_dtype) {
+                                        arg = it->second.casted;
+                                        return;
+                                    }
+                                    // Stale entry (source freed or target changed): drop it.
+                                    tls_autocast_weight_cache.erase(it);
+                                }
+                                auto casted = std::make_shared<Tensor>(arg->To(target_dtype));
+                                tls_autocast_weight_cache.emplace(
+                                    arg.get(),
+                                    AutocastWeightCacheEntry{std::weak_ptr<Tensor>(arg), target_dtype, casted});
+                                arg = std::move(casted);
+                            } else {
+                                arg = std::make_shared<Tensor>(arg->To(target_dtype));
+                            }
                         }
                     }
                 }

infini_train/include/autograd/activations.h

@@ -21,4 +21,16 @@ class Sigmoid : public Function {
                       const std::vector<std::shared_ptr<Tensor>> &output_tensors) override;
     std::vector<std::shared_ptr<Tensor>> Backward(const std::vector<std::shared_ptr<Tensor>> &grad_outputs) override;
 };
+
+class NewGELU : public Function {
+public:
+    static constexpr char kType[] = "NewGELUFunction";
+
+    NewGELU() : Function(kType) {}
+
+    std::vector<std::shared_ptr<Tensor>> Forward(const std::vector<std::shared_ptr<Tensor>> &input_tensors) override;
+    void SetupContext(const std::vector<std::shared_ptr<Tensor>> &input_tensors,
+                      const std::vector<std::shared_ptr<Tensor>> &output_tensors) override;
+    std::vector<std::shared_ptr<Tensor>> Backward(const std::vector<std::shared_ptr<Tensor>> &grad_outputs) override;
+};
 } // namespace infini_train::autograd

infini_train/src/autograd/accumulate.cc

@@ -26,10 +26,12 @@ AccumulateGrad::Backward(const std::vector<std::shared_ptr<Tensor>> &grad_output
 
     if (grad_output) {
         if (grad_output->Dtype() != tensor_->Dtype()) {
-            LOG(WARNING) << "AccumulateGrad: grad dtype (" << kDataTypeToDesc.at(grad_output->Dtype())
-                         << ") does not match parameter dtype (" << kDataTypeToDesc.at(tensor_->Dtype())
-                         << "). This indicates a dtype mismatch in the autograd graph (e.g. autocast "
-                            "running before autograd). The grad is not cast and will be used as-is.";
+            // GEMM backwards emit lower-precision (autocast) gradients while the
+            // master weight and its grad buffer stay in fp32. Cast the grad back to
+            // the parameter dtype here -- the AccumulateGrad kernel is single-dtype
+            // and would otherwise reinterpret the buffer incorrectly. This matches
+            // PyTorch's AccumulateGrad semantics for autocast training.
+            grad_output = std::make_shared<Tensor>(grad_output->To(tensor_->Dtype()));
         }
 
         const bool overwrite = tensor_->ConsumeGradOverwriteFlag();

infini_train/src/autograd/activations.cc

@@ -30,4 +30,29 @@ std::vector<std::shared_ptr<Tensor>> Sigmoid::Backward(const std::vector<std::sh
     auto device = output->GetDevice().type();
     return {Dispatcher::Instance().Call<std::shared_ptr<Tensor>>({device, "SigmoidBackward"}, output, grad_output)};
 }
+
+std::vector<std::shared_ptr<Tensor>> NewGELU::Forward(const std::vector<std::shared_ptr<Tensor>> &input_tensors) {
+    CHECK_EQ(input_tensors.size(), 1);
+    const auto &input = input_tensors[0];
+
+    auto device = input->GetDevice().type();
+    return {Dispatcher::Instance().Call<std::shared_ptr<Tensor>>({device, "NewGELUForward"}, input)};
+}
+
+void NewGELU::SetupContext(const std::vector<std::shared_ptr<Tensor>> &input_tensors,
+                           const std::vector<std::shared_ptr<Tensor>> &) {
+    // Save the forward input x; the backward kernel recomputes tanh(beta * (x + kappa * x^3)) from it.
+    ctx_.SaveForBackward({input_tensors[0]});
+}
+
+std::vector<std::shared_ptr<Tensor>> NewGELU::Backward(const std::vector<std::shared_ptr<Tensor>> &grad_outputs) {
+    auto saved_tensors = ctx_.GetSavedTensors();
+    CHECK_EQ(saved_tensors.size(), 1);
+    const auto &input = saved_tensors[0];
+    CHECK_EQ(grad_outputs.size(), 1);
+    const auto &grad_output = grad_outputs[0];
+
+    auto device = grad_output->GetDevice().type();
+    return {Dispatcher::Instance().Call<std::shared_ptr<Tensor>>({device, "NewGELUBackward"}, grad_output, input)};
+}
 } // namespace infini_train::autograd

infini_train/src/autograd/normalization.cc

@@ -45,10 +45,19 @@ std::vector<std::shared_ptr<Tensor>> LayerNorm::Backward(const std::vector<std::
     const auto &grad_output = grad_outputs[0];
 
     auto device = input->GetDevice().type();
+
+    // GEMM backwards emit gradients in the compute dtype (bf16 under autocast)
+    // while this op runs in fp32 (autocast kFP32 policy); promote the incoming
+    // grad back to the saved input dtype. This preserves the kernel's pre-existing
+    // single-dtype behavior bit-for-bit.
+    auto grad_output_promoted = grad_output->Dtype() == input->Dtype()
+                                  ? grad_output
+                                  : std::make_shared<Tensor>(grad_output->To(input->Dtype()));
+
     auto [grad_input, grad_weight, grad_bias]
         = Dispatcher::Instance()
               .Call<std::tuple<std::shared_ptr<Tensor>, std::shared_ptr<Tensor>, std::shared_ptr<Tensor>>>(
-                  {device, "LayerNormBackward"}, input, weight, bias, mean, rstd, grad_output);
+                  {device, "LayerNormBackward"}, input, weight, bias, mean, rstd, grad_output_promoted);
     return {grad_input, grad_weight, grad_bias};
 }
 } // namespace infini_train::autograd

infini_train/src/core/runtime/cuda/cuda_guard_impl.cc

@@ -1,6 +1,7 @@
 #include "infini_train/src/core/runtime/cuda/cuda_guard_impl.h"
 
 #include <array>
+#include <cstdint>
 #include <memory>
 #include <mutex>
 
@@ -50,6 +51,17 @@ void CudaGuardImpl::InitSingleStream(Device device) {
 
     cuda_streams[device.index()] = std::make_unique<CudaStream>();
 
+    // Keep the default memory pool's cached blocks mapped across stream/device
+    // synchronizations. With the default release threshold (0), every sync returns
+    // cached blocks to the OS and the next cudaMallocAsync re-maps physical pages,
+    // which is extremely slow on WSL2 (~5ms per allocation) and is hit on every
+    // profiled kernel launch (the profiler synchronizes the stream per kernel).
+    // On native Linux this only keeps the pool from shrinking back at sync points.
+    cudaMemPool_t default_pool;
+    CUDA_CHECK(cudaDeviceGetDefaultMemPool(&default_pool, device.index()));
+    cuuint64_t release_threshold = UINT64_MAX;
+    CUDA_CHECK(cudaMemPoolSetAttribute(default_pool, cudaMemPoolAttrReleaseThreshold, &release_threshold));
+
     CUDA_CHECK(cudaSetDevice(current_device));
 }