feat: pack 2 fc layers in mlp with swiglu

example/llama3/checkpoint_loader.cc

@@ -277,7 +277,8 @@ std::shared_ptr<nn::TransformerModel> LoadFromLLMC(const std::string &filepath)
                                                   nn::TransformerChunk::kHLayerName, std::to_string(local_layer_index),
                                                   nn::TransformerLayer::kMlpLayerName, nn::MLP::kCFcLayerName,
                                                   nn::parallel::ColumnParallelLinear::kParamWeightName)];
-            ReadMatrixRowShardFloat(ifs, static_cast<float *>(tensor->DataPtr()),
+            float *dst = static_cast<float *>(tensor->DataPtr()) + fc_pp * n_embd;
+            ReadMatrixRowShardFloat(ifs, dst,
                                     /*rows=*/fc_out, /*cols=*/n_embd,
                                     /*row_start=*/tp_rank * fc_pp, /*row_cnt=*/fc_pp);
             ++local_layer_index;
@@ -293,7 +294,7 @@ std::shared_ptr<nn::TransformerModel> LoadFromLLMC(const std::string &filepath)
         if (owned_layers[i]) {
             auto &tensor = state_dict[std::format("{}.{}.{}.{}.{}.{}", nn::TransformerModel::kTransformerModelName,
                                                   nn::TransformerChunk::kHLayerName, std::to_string(local_layer_index),
-                                                  nn::TransformerLayer::kMlpLayerName, nn::MLP::kCFc2LayerName,
+                                                  nn::TransformerLayer::kMlpLayerName, nn::MLP::kCFcLayerName,
                                                   nn::parallel::ColumnParallelLinear::kParamWeightName)];
             ReadMatrixRowShardFloat(ifs, static_cast<float *>(tensor->DataPtr()),
                                     /*rows=*/fc_out, /*cols=*/n_embd,

example/mixtral/checkpoint_loader.cc

@@ -119,10 +119,10 @@ std::shared_ptr<nn::TransformerModel> LoadFromLLMC(const std::string &filepath,
         CHECK(ifs) << "Failed to read tensor " << name;
     };
 
-    auto read_projection_into_packed_qkv = [&](const std::string &packed_qkv_name, int64_t row_offset, int64_t num_rows,
-                                               const std::string &projection_name) {
-        CHECK(state.contains(packed_qkv_name)) << "Model state_dict does not contain " << packed_qkv_name;
-        std::shared_ptr<infini_train::Tensor> tensor = state.at(packed_qkv_name);
+    auto read_projection_into_packed_weight = [&](const std::string &packed_weight_name, int64_t row_offset,
+                                                  int64_t num_rows, const std::string &projection_name) {
+        CHECK(state.contains(packed_weight_name)) << "Model state_dict does not contain " << packed_weight_name;
+        std::shared_ptr<infini_train::Tensor> tensor = state.at(packed_weight_name);
         CHECK(tensor->Dtype() == infini_train::DataType::kFLOAT32)
             << "Only float32 tiny Mixtral LLMC files are supported: " << projection_name;
         CHECK_EQ(tensor->Dims().size(), 2);
@@ -144,17 +144,21 @@ std::shared_ptr<nn::TransformerModel> LoadFromLLMC(const std::string &filepath,
         const int64_t head_dim = config.n_embd / config.n_head;
         const int64_t q_rows = config.n_head * head_dim;
         const int64_t kv_rows = config.n_kv_head * head_dim;
-        read_projection_into_packed_qkv(c_attn_name, 0, q_rows, c_attn_name + ".q_proj");
-        read_projection_into_packed_qkv(c_attn_name, q_rows, kv_rows, c_attn_name + ".k_proj");
-        read_projection_into_packed_qkv(c_attn_name, q_rows + kv_rows, kv_rows, c_attn_name + ".v_proj");
+        read_projection_into_packed_weight(c_attn_name, 0, q_rows, c_attn_name + ".q_proj");
+        read_projection_into_packed_weight(c_attn_name, q_rows, kv_rows, c_attn_name + ".k_proj");
+        read_projection_into_packed_weight(c_attn_name, q_rows + kv_rows, kv_rows, c_attn_name + ".v_proj");
         read_tensor_by_state_key(prefix + ".attn.c_proj.weight");
         read_tensor_by_state_key(prefix + ".ln_2.weight");
         read_tensor_by_state_key(prefix + ".mlp.router.weight");
         for (int64_t expert = 0; expert < moe_config.num_experts; ++expert) {
             const std::string expert_prefix = prefix + ".mlp.experts.expert_" + std::to_string(expert);
-            read_tensor_by_state_key(expert_prefix + ".c_fc2.weight");  // Mixtral w1/gate_proj
-            read_tensor_by_state_key(expert_prefix + ".c_fc.weight");   // Mixtral w3/up_proj
-            read_tensor_by_state_key(expert_prefix + ".c_proj.weight"); // Mixtral w2/down_proj
+            const std::string packed_fc1_name = expert_prefix + ".c_fc.weight";
+            read_projection_into_packed_weight(packed_fc1_name, 0, moe_config.moe_ffn_hidden_size,
+                                               expert_prefix + ".c_fc2.weight"); // Mixtral w1/gate_proj
+            read_projection_into_packed_weight(packed_fc1_name, moe_config.moe_ffn_hidden_size,
+                                               moe_config.moe_ffn_hidden_size,
+                                               expert_prefix + ".c_fc.weight"); // Mixtral w3/up_proj
+            read_tensor_by_state_key(expert_prefix + ".c_proj.weight");         // Mixtral w2/down_proj
         }
     }
     read_tensor_by_state_key("transformer.ln_f.weight");

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 SwiGLU : public Function {
+public:
+    static constexpr char kType[] = "SwiGLUFunction";
+
+    SwiGLU() : 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/include/nn/modules/activations.h

@@ -30,6 +30,7 @@ class SwiGLU : public CloneableModule<SwiGLU> {
     static constexpr char kType[] = "SwiGLU";
     SwiGLU() : CloneableModule(kType) {}
 
+    // The last input dimension is packed as [gate, up], matching Megatron-LM.
     std::vector<std::shared_ptr<Tensor>> Forward(const std::vector<std::shared_ptr<Tensor>> &x) override;
 };
 } // namespace infini_train::nn

infini_train/include/nn/modules/transformer/mlp.h

@@ -13,8 +13,7 @@ class MLP : public infini_train::nn::CloneableModule<MLP> {
     static constexpr char kGeluLayerName[] = "gelu";
     static constexpr char kCProjLayerName[] = "c_proj";
 
-    static constexpr char kCFc2LayerName[] = "c_fc2";
-    static constexpr char kSiluLayerName[] = "silu";
+    static constexpr char kSwiGLULayerName[] = "swiglu";
 
     explicit MLP(const TransformerConfig &config);
 

infini_train/src/autograd/activations.cc

@@ -30,4 +30,30 @@ 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>> SwiGLU::Forward(const std::vector<std::shared_ptr<Tensor>> &input_tensors) {
+    CHECK_EQ(input_tensors.size(), 1);
+    const auto &input = input_tensors[0];
+    CHECK_GT(input->Dims().size(), 0);
+    CHECK_EQ(input->Dims().back() % 2, 0) << "SwiGLU expects an even last dimension";
+
+    auto device = input->GetDevice().type();
+    return {Dispatcher::Instance().Call<std::shared_ptr<Tensor>>({device, "SwiGLUForward"}, input)};
+}
+
+void SwiGLU::SetupContext(const std::vector<std::shared_ptr<Tensor>> &input_tensors,
+                          const std::vector<std::shared_ptr<Tensor>> &) {
+    ctx_.SaveForBackward({input_tensors[0]});
+}
+
+std::vector<std::shared_ptr<Tensor>> SwiGLU::Backward(const std::vector<std::shared_ptr<Tensor>> &grad_outputs) {
+    auto saved_tensors = ctx_.GetSavedTensors();
+    CHECK_EQ(saved_tensors.size(), 1);
+    CHECK_EQ(grad_outputs.size(), 1);
+    const auto &input = saved_tensors[0];
+    const auto &grad_output = grad_outputs[0];
+
+    auto device = input->GetDevice().type();
+    return {Dispatcher::Instance().Call<std::shared_ptr<Tensor>>({device, "SwiGLUBackward"}, input, grad_output)};
+}
 } // namespace infini_train::autograd

infini_train/src/kernels/cpu/swiglu.cc

@@ -0,0 +1,74 @@
+#include <cmath>
+#include <memory>
+
+#include "glog/logging.h"
+
+#include "infini_train/include/dispatcher.h"
+#include "infini_train/include/tensor.h"
+
+namespace infini_train::kernels::cpu {
+std::shared_ptr<Tensor> SwiGLUForward(const std::shared_ptr<Tensor> &input) {
+    CHECK(input->Dtype() == DataType::kFLOAT32);
+    CHECK(input->IsContiguous());
+    auto output_dims = input->Dims();
+    CHECK_GT(output_dims.size(), 0);
+    CHECK_EQ(output_dims.back() % 2, 0);
+    const int64_t hidden = output_dims.back() / 2;
+    CHECK_GT(hidden, 0);
+    output_dims.back() = hidden;
+
+    auto output = std::make_shared<Tensor>(output_dims, input->Dtype(), input->GetDevice());
+    const float *input_ptr = static_cast<const float *>(input->DataPtr());
+    float *output_ptr = static_cast<float *>(output->DataPtr());
+    const int64_t rows = output->NumElements() / hidden;
+    for (int64_t row = 0; row < rows; ++row) {
+        const int64_t input_base = row * 2 * hidden;
+        const int64_t output_base = row * hidden;
+        for (int64_t col = 0; col < hidden; ++col) {
+            const float gate = input_ptr[input_base + col];
+            const float up = input_ptr[input_base + hidden + col];
+            output_ptr[output_base + col] = up * gate / (1.0f + std::exp(-gate));
+        }
+    }
+    return output;
+}
+
+std::shared_ptr<Tensor> SwiGLUBackward(const std::shared_ptr<Tensor> &input,
+                                       const std::shared_ptr<Tensor> &grad_output) {
+    CHECK(input->Dtype() == DataType::kFLOAT32);
+    CHECK(grad_output->Dtype() == input->Dtype());
+    CHECK(input->IsContiguous());
+    CHECK(grad_output->IsContiguous());
+    CHECK_GT(input->Dims().size(), 0);
+    const int64_t hidden = input->Dims().back() / 2;
+    CHECK_GT(hidden, 0);
+    CHECK_EQ(grad_output->NumElements() * 2, input->NumElements());
+
+    auto grad_input = std::make_shared<Tensor>(input->Dims(), input->Dtype(), input->GetDevice());
+    const float *input_ptr = static_cast<const float *>(input->DataPtr());
+    const float *grad_output_ptr = static_cast<const float *>(grad_output->DataPtr());
+    float *grad_input_ptr = static_cast<float *>(grad_input->DataPtr());
+    const int64_t rows = grad_output->NumElements() / hidden;
+    for (int64_t row = 0; row < rows; ++row) {
+        const int64_t input_base = row * 2 * hidden;
+        const int64_t output_base = row * hidden;
+        for (int64_t col = 0; col < hidden; ++col) {
+            const float gate = input_ptr[input_base + col];
+            const float up = input_ptr[input_base + hidden + col];
+            const float grad = grad_output_ptr[output_base + col];
+            const float sigmoid = 1.0f / (1.0f + std::exp(-gate));
+            grad_input_ptr[input_base + col] = grad * up * sigmoid * (1.0f + gate * (1.0f - sigmoid));
+            grad_input_ptr[input_base + hidden + col] = grad * gate * sigmoid;
+        }
+    }
+    return grad_input;
+}
+} // namespace infini_train::kernels::cpu
+
+#define REGISTER_CPU_SWIGLU_KERNEL(kernel_name)                                                                        \
+    REGISTER_KERNEL(infini_train::Device::DeviceType::kCPU, kernel_name, infini_train::kernels::cpu::kernel_name)
+
+REGISTER_CPU_SWIGLU_KERNEL(SwiGLUForward)
+REGISTER_CPU_SWIGLU_KERNEL(SwiGLUBackward)
+
+#undef REGISTER_CPU_SWIGLU_KERNEL

infini_train/src/kernels/cuda/swiglu.cu

@@ -0,0 +1,147 @@
+#include <algorithm>
+#include <cstddef>
+
+#include "infini_train/include/common/common.h"
+#include "infini_train/include/common/cuda/kernel_helper.cuh"
+#include "infini_train/include/core/runtime/device_guard.h"
+#include "infini_train/include/datatype.h"
+#include "infini_train/include/dispatcher.h"
+#include "infini_train/include/tensor.h"
+
+#include "infini_train/src/core/runtime/cuda/cuda_runtime_common.h"
+
+namespace infini_train::kernels::cuda {
+namespace {
+using namespace infini_train::common::cuda;
+
+template <typename T>
+__global__ void SwiGLUForwardKernel(T *__restrict__ output, const T *__restrict__ input, int64_t hidden,
+                                    size_t num_elements) {
+    const size_t grid_stride = static_cast<size_t>(gridDim.x) * blockDim.x;
+    for (size_t idx = static_cast<size_t>(blockIdx.x) * blockDim.x + threadIdx.x; idx < num_elements;
+         idx += grid_stride) {
+        const size_t row = idx / hidden;
+        const size_t col = idx % hidden;
+        const size_t input_base = row * 2 * hidden;
+        const T gate = input[input_base + col];
+        const T up = input[input_base + hidden + col];
+        output[idx] = Mul(up, Mul(gate, Sigmoid(gate)));
+    }
+}
+
+template <typename T, typename InputT, typename GradT>
+__global__ void SwiGLUBackwardKernel(T *__restrict__ grad_input, const InputT *__restrict__ input,
+                                     const GradT *__restrict__ grad_output, int64_t hidden, size_t num_elements) {
+    const size_t grid_stride = static_cast<size_t>(gridDim.x) * blockDim.x;
+    for (size_t idx = static_cast<size_t>(blockIdx.x) * blockDim.x + threadIdx.x; idx < num_elements;
+         idx += grid_stride) {
+        const size_t row = idx / hidden;
+        const size_t col = idx % hidden;
+        const size_t input_base = row * 2 * hidden;
+        const T gate = Cast<T>(input[input_base + col]);
+        const T up = Cast<T>(input[input_base + hidden + col]);
+        const T grad = Cast<T>(grad_output[idx]);
+        const T sigmoid = Sigmoid(gate);
+        grad_input[input_base + col] = Mul(grad, Mul(up, Mul(sigmoid, Add(T(1), Mul(gate, Sub(T(1), sigmoid))))));
+        grad_input[input_base + hidden + col] = Mul(grad, Mul(gate, sigmoid));
+    }
+}
+
+inline size_t ChooseBlockSize(size_t num_elements) {
+    if (num_elements < 1024) {
+        return 64;
+    }
+    if (num_elements < 65536) {
+        return 128;
+    }
+    if (num_elements < 1048576) {
+        return 256;
+    }
+    return 512;
+}
+} // namespace
+
+std::shared_ptr<Tensor> SwiGLUForward(const std::shared_ptr<Tensor> &input) {
+    CHECK(input->IsContiguous());
+    auto output_dims = input->Dims();
+    CHECK_GT(output_dims.size(), 0);
+    CHECK_EQ(output_dims.back() % 2, 0);
+    const int64_t hidden = output_dims.back() / 2;
+    CHECK_GT(hidden, 0);
+    output_dims.back() = hidden;
+
+    auto output = std::make_shared<Tensor>(output_dims, input->Dtype(), input->GetDevice());
+    auto device = output->GetDevice();
+    const auto &stream = dynamic_cast<infini_train::core::cuda::CudaStream *>(
+                             infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device))
+                             ->cuda_stream();
+    const size_t num_elements = output->NumElements();
+    const dim3 block(ChooseBlockSize(num_elements));
+    const dim3 grid(std::min(CEIL_DIV(num_elements, block.x), static_cast<size_t>(65535)));
+
+    switch (input->Dtype()) {
+    case DataType::kFLOAT32:
+        SwiGLUForwardKernel<<<grid, block, 0, stream>>>(static_cast<float *>(output->DataPtr()),
+                                                        static_cast<const float *>(input->DataPtr()), hidden,
+                                                        num_elements);
+        break;
+    case DataType::kBFLOAT16:
+        SwiGLUForwardKernel<<<grid, block, 0, stream>>>(static_cast<nv_bfloat16 *>(output->DataPtr()),
+                                                        static_cast<const nv_bfloat16 *>(input->DataPtr()), hidden,
+                                                        num_elements);
+        break;
+    default:
+        LOG_LOC(FATAL, "CUDA SwiGLUForward: unsupported data type");
+    }
+    return output;
+}
+
+std::shared_ptr<Tensor> SwiGLUBackward(const std::shared_ptr<Tensor> &input,
+                                       const std::shared_ptr<Tensor> &grad_output) {
+    CHECK(input->IsContiguous());
+    CHECK(grad_output->IsContiguous());
+    CHECK_GT(input->Dims().size(), 0);
+    const int64_t hidden = input->Dims().back() / 2;
+    CHECK_GT(hidden, 0);
+    CHECK_EQ(grad_output->NumElements() * 2, input->NumElements());
+
+    const DataType output_dtype = PromoteDataTypes(input->Dtype(), grad_output->Dtype());
+    auto grad_input = std::make_shared<Tensor>(input->Dims(), output_dtype, input->GetDevice());
+    auto device = input->GetDevice();
+    const auto &stream = dynamic_cast<infini_train::core::cuda::CudaStream *>(
+                             infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device))
+                             ->cuda_stream();
+    const size_t num_elements = grad_output->NumElements();
+    const dim3 block(ChooseBlockSize(num_elements));
+    const dim3 grid(std::min(CEIL_DIV(num_elements, block.x), static_cast<size_t>(65535)));
+
+    if (input->Dtype() == DataType::kFLOAT32 && grad_output->Dtype() == DataType::kFLOAT32) {
+        SwiGLUBackwardKernel<<<grid, block, 0, stream>>>(
+            static_cast<float *>(grad_input->DataPtr()), static_cast<const float *>(input->DataPtr()),
+            static_cast<const float *>(grad_output->DataPtr()), hidden, num_elements);
+    } else if (input->Dtype() == DataType::kBFLOAT16 && grad_output->Dtype() == DataType::kBFLOAT16) {
+        SwiGLUBackwardKernel<<<grid, block, 0, stream>>>(
+            static_cast<nv_bfloat16 *>(grad_input->DataPtr()), static_cast<const nv_bfloat16 *>(input->DataPtr()),
+            static_cast<const nv_bfloat16 *>(grad_output->DataPtr()), hidden, num_elements);
+    } else if (input->Dtype() == DataType::kBFLOAT16 && grad_output->Dtype() == DataType::kFLOAT32) {
+        SwiGLUBackwardKernel<<<grid, block, 0, stream>>>(
+            static_cast<float *>(grad_input->DataPtr()), static_cast<const nv_bfloat16 *>(input->DataPtr()),
+            static_cast<const float *>(grad_output->DataPtr()), hidden, num_elements);
+    } else if (input->Dtype() == DataType::kFLOAT32 && grad_output->Dtype() == DataType::kBFLOAT16) {
+        SwiGLUBackwardKernel<<<grid, block, 0, stream>>>(
+            static_cast<float *>(grad_input->DataPtr()), static_cast<const float *>(input->DataPtr()),
+            static_cast<const nv_bfloat16 *>(grad_output->DataPtr()), hidden, num_elements);
+    } else {
+        LOG_LOC(FATAL, "CUDA SwiGLUBackward: unsupported data type combination");
+    }
+    return grad_input;
+}
+} // namespace infini_train::kernels::cuda
+
+#define REGISTER_CUDA_SWIGLU_KERNEL(kernel_name)                                                                       \
+    REGISTER_KERNEL(infini_train::Device::DeviceType::kCUDA, kernel_name, infini_train::kernels::cuda::kernel_name)
+
+REGISTER_CUDA_SWIGLU_KERNEL(SwiGLUForward)
+REGISTER_CUDA_SWIGLU_KERNEL(SwiGLUBackward)
+
+#undef REGISTER_CUDA_SWIGLU_KERNEL

infini_train/src/nn/modules/activations.cc

@@ -19,6 +19,6 @@ std::vector<std::shared_ptr<Tensor>> NewGELU::Forward(const std::vector<std::sha
 }
 
 std::vector<std::shared_ptr<Tensor>> SwiGLU::Forward(const std::vector<std::shared_ptr<Tensor>> &x) {
-    return {x[0] * function::Sigmoid(x[0])};
+    return std::make_shared<autograd::SwiGLU>()->Apply(x);
 }
 } // namespace infini_train::nn