fix: fix gpt2 runtime errors on cuda, but loss is still NaN

Author: Chamberlain0w0

example/gpt2/net.cc

@@ -11,6 +11,7 @@
 
 #include "glog/logging.h"
 
+#include "infini_train/include/device.h"
 #include "infini_train/include/nn/functional.h"
 #include "infini_train/include/nn/init.h"
 #include "infini_train/include/nn/modules/container.h"
@@ -60,6 +61,11 @@ CausalSelfAttention::CausalSelfAttention(const GPT2Config &config)
                 ->View({1, 1, config_.block_size, config_.block_size});
 }
 
+void CausalSelfAttention::To(infini_train::Device device) {
+    nn::Module::To(device);
+    bias_ = std::make_shared<infini_train::Tensor>(bias_->To(device));
+}
+
 std::vector<std::shared_ptr<infini_train::Tensor>>
 CausalSelfAttention::Forward(const std::vector<std::shared_ptr<infini_train::Tensor>> &x) {
     const auto B = x[0]->Dims()[0]; // bs
@@ -163,8 +169,8 @@ GPT2::GPT2(const GPT2Config &config) : config_(config) {
     modules_[kLMHeadLayerName] = std::make_unique<GPT2Linear>(config.n_embd, config.vocab_size, false, true);
     // https://paperswithcode.com/method/weight-tying
     *mutable_module(kTransformerLayerName)
-        ->mutable_module(kWTELayerName)
-        ->mutable_parameter(GPT2Linear::kParamWeightName)
+         ->mutable_module(kWTELayerName)
+         ->mutable_parameter(GPT2Linear::kParamWeightName)
         = module(kLMHeadLayerName).parameter(GPT2Linear::kParamWeightName);
 
     // init all weights

example/gpt2/net.h

@@ -34,6 +34,8 @@ class CausalSelfAttention : public infini_train::nn::Module {
     std::vector<std::shared_ptr<infini_train::Tensor>>
     Forward(const std::vector<std::shared_ptr<infini_train::Tensor>> &x) override;
 
+    void To(infini_train::Device device) override;
+
 private:
     GPT2Config config_;
     int64_t n_head_ = 0;

infini_train/include/nn/modules/module.h

@@ -31,7 +31,7 @@ class Module {
 
     virtual std::vector<std::shared_ptr<Tensor>> Forward(const std::vector<std::shared_ptr<Tensor>> &input_tensors) = 0;
 
-    void To(Device device);
+    virtual void To(Device device);
 
     void Apply(std::function<void(Module *)> fn);
 

infini_train/src/autograd/elementwise.cc

@@ -54,7 +54,7 @@ std::vector<std::shared_ptr<Tensor>> Tanh::Backward(const std::vector<std::share
     }
 #ifdef USE_CUDA
     case DeviceType::kCUDA: {
-        grad_input = kernels::cpu::TanhBackward(grad_output, output);
+        grad_input = kernels::cuda::TanhBackward(grad_output, output);
         break;
     }
 #endif

infini_train/src/kernels/cuda/elementwise.cu

@@ -38,7 +38,8 @@ __global__ void BinaryForwardKernel(T *output, Func fn, size_t num_elements_a, s
 // launch the given kernel function with the given output and inputs
 template <size_t BLOCK_SIZE, typename T, typename Kernel, typename... Inputs>
 void LaunchKernel(Kernel &&kernel, const std::shared_ptr<Tensor> &output, const Inputs &...inputs) {
-    auto extract_ptrs = [](const auto &...ts) { return std::make_tuple(static_cast<T *>(ts->DataPtr())...); };
+    auto extract_ptrs
+        = [](const auto &...ts) { return std::make_tuple(static_cast<T *>(ts ? ts->DataPtr() : nullptr)...); };
     auto input_ptrs = extract_ptrs(inputs...);
 
     cudaDeviceProp prop;
@@ -135,7 +136,6 @@ void LaunchBackward(FuncA fun_a, FuncB fun_b, const std::shared_ptr<Tensor> &out
     T *output_a_ptr = static_cast<T *>(output_a->DataPtr());
     T *output_b_ptr = static_cast<T *>(output_b->DataPtr());
     const T *grad_output_ptr = static_cast<const T *>(grad_output->DataPtr());
-
     LaunchKernel<BLOCK_SIZE, T>(
         [=](dim3 grid, dim3 block, size_t offset, auto... ptrs) {
             BinaryBackwardKernel<<<grid, block>>>(output_a_ptr, output_b_ptr, fun_a, fun_b, a_num_elements,
@@ -201,7 +201,6 @@ std::pair<std::shared_ptr<Tensor>, std::shared_ptr<Tensor>>
 BinaryBackward(const std::shared_ptr<Tensor> &grad_output, const std::shared_ptr<Tensor> &a,
                const std::shared_ptr<Tensor> &b, const std::vector<int64_t> &a_dims, const std::vector<int64_t> &b_dims,
                FuncA fn_a, FuncB fn_b) {
-
     const auto a_num_elements = std::accumulate(a_dims.begin(), a_dims.end(), 1, std::multiplies<int64_t>());
     const auto b_num_elements = std::accumulate(b_dims.begin(), b_dims.end(), 1, std::multiplies<int64_t>());
 
@@ -212,14 +211,15 @@ BinaryBackward(const std::shared_ptr<Tensor> &grad_output, const std::shared_ptr
     if (b) {
         CHECK(b_num_elements == b->NumElements());
     }
-
     auto dtype = grad_output->Dtype();
-    auto device = a->GetDevice();
-    // Currently a and b should have the same data type
-    CHECK(dtype == b->Dtype());
-    auto grad_a = std::make_shared<Tensor>(a->Dims(), dtype, device);
-    auto grad_b = std::make_shared<Tensor>(b->Dims(), dtype, device);
+    auto device = grad_output->GetDevice();
 
+    // Currently a and b should have the same data type
+    if (a && b) {
+        CHECK(a->Dtype() == b->Dtype());
+    }
+    auto grad_a = std::make_shared<Tensor>(a_dims, dtype, device);
+    auto grad_b = std::make_shared<Tensor>(b_dims, dtype, device);
     switch (dtype) {
     case DataType::kFLOAT32:
         LaunchBackward<256, float>(fn_a, fn_b, grad_a, grad_b, a_num_elements, b_num_elements, grad_output, a, b);

infini_train/src/kernels/cuda/linear.cu

@@ -77,7 +77,6 @@ std::shared_ptr<Tensor> MatmulForward(const std::shared_ptr<Tensor> &input, cons
                                             CUDA_R_32F, lda, stride_a, input->DataPtr(), CUDA_R_32F, ldb, stride_b,
                                             &beta, output->DataPtr(), CUDA_R_32F, ldc, stride_c, bs, CUDA_R_32F,
                                             CUBLAS_GEMM_DEFAULT));
-    CUDA_CHECK(cudaDeviceSynchronize());
     CUBLAS_CHECK(cublasDestroy(handle));
     return output;
 }
@@ -259,7 +258,7 @@ LinearBackward(const std::shared_ptr<Tensor> &input, const std::shared_ptr<Tenso
     if (transpose) {
         // d_input = d_output * weight --> d_input.T = weight * d_output.T
         CUBLAS_CHECK(cublasSgemm(handle, CUBLAS_OP_T, CUBLAS_OP_N, in_features, bs, out_features, &alpha,
-                                 static_cast<const float *>(weight->DataPtr()), in_features,
+                                 static_cast<const float *>(weight->DataPtr()), out_features,
                                  static_cast<const float *>(grad_output->DataPtr()), out_features, &beta,
                                  static_cast<float *>(grad_input->DataPtr()), in_features));