feat: pack 2 fc layers in mlp with swiglu

example/llama3/checkpoint_loader.cc

@@ -159,7 +159,7 @@ std::shared_ptr<nn::TransformerModel> LoadFromLLMC(const std::string &filepath)
 
     // RowParallel (proj)
     const int64_t in_pp = static_cast<int64_t>(n_embd) / tp_size;
-    // nn::MLP: c_fc/c_fc2（shard along row），c_proj（shard along col）
+    // nn::MLP: packed c_fc [gate | up] (shard each block along row), c_proj (shard along col)
     const int64_t fc_out = ffn_hidden;
     const int64_t fc_pp = fc_out / tp_size;
     const int64_t in_fc_pp = ffn_hidden / tp_size;
@@ -269,15 +269,16 @@ std::shared_ptr<nn::TransformerModel> LoadFromLLMC(const std::string &filepath)
         }
     }
 
-    // transformer.h.{i}.mlp.c_fc.weight : ColumnParallelLinear, but actually applies on "rows"
+    // transformer.h.{i}.mlp.c_fc.weight (up) -> local packed c_fc rows [fc_pp : 2*fc_pp)
     local_layer_index = 0;
     for (int i = 0; i < static_cast<int>(n_layer); ++i) {
         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::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;
@@ -287,13 +288,13 @@ std::shared_ptr<nn::TransformerModel> LoadFromLLMC(const std::string &filepath)
         }
     }
 
-    // transformer.h.{i}.mlp.c_fc2.weight : ColumnParallelLinear, but actually applies on "rows"
+    // transformer.h.{i}.mlp.c_fc2.weight (gate) -> local packed c_fc rows [0 : fc_pp)
     local_layer_index = 0;
     for (int i = 0; i < static_cast<int>(n_layer); ++i) {
         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/common/cuda/common_cuda.h

@@ -11,6 +11,19 @@
 
 namespace infini_train::common::cuda {
 
+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;
+}
+
 // Common CUDA Macros
 #define CUDA_CHECK(call)                                                                                               \
     do {                                                                                                               \

infini_train/include/nn/lora/lora_utils.h

@@ -28,6 +28,14 @@ std::shared_ptr<Tensor> SlicePackedQKVRowsForTensorParallel(const std::shared_pt
 std::shared_ptr<Tensor> RestorePackedQKVRowsFromTensorParallel(const std::shared_ptr<Tensor> &gathered_tensor,
                                                                int64_t q_rows, int tp_size);
 
+// Internal helper for packed SwiGLU tensors stored as [gate | up] along dim 0.
+std::shared_ptr<Tensor> SlicePackedSwiGLURowsForTensorParallel(const std::shared_ptr<Tensor> &full_tensor, int tp_rank,
+                                                               int tp_size);
+
+// Internal helper for TP-gathered packed SwiGLU shards stored rank-major as [gate_i | up_i].
+std::shared_ptr<Tensor> RestorePackedSwiGLURowsFromTensorParallel(const std::shared_ptr<Tensor> &gathered_tensor,
+                                                                  int tp_size);
+
 } // namespace detail
 
 /**
@@ -121,8 +129,8 @@ void LoadLoRAStateDict(std::shared_ptr<Module> model,
 void SaveLoRAWeights(const std::shared_ptr<Module> &model, const std::string &filepath);
 
 /**
- * Load LoRA parameters from file. Packed QKV LoRA-B tensors are split as
- * [Qi | Ki | Vi] for the current TP rank.
+ * Load LoRA parameters from file. Packed QKV and SwiGLU LoRA-B tensors are
+ * reordered into their rank-local packed layouts.
  */
 void LoadLoRAWeights(std::shared_ptr<Module> model, const std::string &filepath);
 

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/elementwise.cu

@@ -193,20 +193,6 @@ __global__ void BinaryBackwardKernelNoBroadcastVectorized(T *__restrict__ outA,
     }
 }
 
-// Helper to choose optimal block size based on tensor size
-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;
-}
-
 inline dim3 ChooseBlockDims(size_t num_elements) { return dim3(ChooseBlockSize(num_elements)); }
 
 // launch the given kernel function with the given output and inputs