feat: support moe_ffn_hidden_size config

Author: kilinchange

infini_train/src/nn/modules/transformer/mlp.cc

@@ -35,9 +35,14 @@ MLP::MLP(const TransformerConfig &config) : CloneableModule(kType) {
     }
 
     // Round up to multiple_of
-    int64_t before_round = ffn_hidden;
     ffn_hidden = (ffn_hidden + config.multiple_of - 1) / config.multiple_of * config.multiple_of;
 
+    if (config.ffn_type == FFNType::kMoE && config.moe_config.has_value()
+        && config.moe_config->moe_ffn_hidden_size > 0) {
+        ffn_hidden = config.moe_config->moe_ffn_hidden_size;
+    }
+    CHECK_GT(ffn_hidden, 0);
+
     // c_fc: ColumnParallel (input full, output parallel)
     modules_[kCFcLayerName] = std::make_shared<parallel::ColumnParallelLinear>(
         /*in_features=*/config.n_embd, /*out_features=*/ffn_hidden,

test/transformer/test_transformer_architecture.cc

@@ -564,12 +564,13 @@ void TestMoELayerTop2() {
     config.n_embd = 32;
     config.n_head = 2;
     config.n_kv_head = 2;
-    config.activation_type = nn::MLPType::kGELU;
-    config.add_bias_linear = true;
+    config.activation_type = nn::MLPType::kSwiGLU;
+    config.add_bias_linear = false;
     config.ffn_type = nn::FFNType::kMoE;
     config.moe_config = nn::MoEConfig{};
     config.moe_config->num_experts = 4;
     config.moe_config->router_topk = 2;
+    config.moe_config->moe_ffn_hidden_size = 48;
 
     auto moe = std::make_shared<nn::moe::MoELayer>(config);
     auto input = std::make_shared<Tensor>(std::vector<int64_t>{2, 4, config.n_embd}, DataType::kFLOAT32);
@@ -579,6 +580,14 @@ void TestMoELayerTop2() {
     CHECK_EQ(output.size(), 1);
     CHECK(output[0]->Dims() == input->Dims());
 
+    auto state = moe->StateDict();
+    CHECK(state.contains("experts.expert_0.c_fc.weight"));
+    CHECK(state.contains("experts.expert_0.c_fc2.weight"));
+    CHECK(state.contains("experts.expert_0.c_proj.weight"));
+    CHECK(state.at("experts.expert_0.c_fc.weight")->Dims() == std::vector<int64_t>({48, config.n_embd}));
+    CHECK(state.at("experts.expert_0.c_fc2.weight")->Dims() == std::vector<int64_t>({48, config.n_embd}));
+    CHECK(state.at("experts.expert_0.c_proj.weight")->Dims() == std::vector<int64_t>({config.n_embd, 48}));
+
     std::cout << "SUCCESS: MoE layer top-2 forward works correctly!" << std::endl;
 }