feat: extract the common module of Transformer

Author: JYMiracle305

example/gpt2/main.cc

@@ -192,8 +192,7 @@ void Train(const nn::parallel::Rank &rank) {
     std::shared_ptr<nn::Module> model = nullptr;
 
     if (!FLAGS_llmc_filepath.empty()) {
-        auto gpt2_model = GPT2::FromLLMC(FLAGS_llmc_filepath);
-        model = gpt2_model;
+        model = GPT2::FromLLMC(FLAGS_llmc_filepath);
     } else if (kModelToConfigs.count(FLAGS_model)) {
         model_config = kModelToConfigs.at(FLAGS_model);
         model = std::make_shared<GPT2>(model_config);

example/llama3/main.cc

@@ -171,8 +171,7 @@ void Train(const nn::parallel::Rank &rank) {
     nn::TransformerConfig model_config = nn::TransformerConfig::LLaMA3();
     std::shared_ptr<nn::Module> model = nullptr;
     if (!FLAGS_llmc_filepath.empty()) {
-        auto llama3_model = LLaMA3::FromLLMC(FLAGS_llmc_filepath);
-        model = llama3_model;
+        model = LLaMA3::FromLLMC(FLAGS_llmc_filepath);
     } else {
         model = std::make_shared<LLaMA3>(model_config);
     }

example/llama3/net.cc

@@ -75,6 +75,12 @@ std::shared_ptr<LLaMA3> LLaMA3::FromLLMC(const std::string &filepath) {
     llama3_config.n_head = n_head;
     llama3_config.n_kv_head = n_kv_head;
     llama3_config.n_embd = n_embd;
+    llama3_config.ffn_dim_multiplier = ffn_dim_multiplier;
+    llama3_config.multiple_of = multiple_of;
+    llama3_config.rope_theta = rope_theta;
+    llama3_config.use_scaled_rope = static_cast<bool>(use_scaled_rope);
+    llama3_config.norm_eps = norm_eps;
+    llama3_config.max_gen_batch_size = max_gen_bs;
     auto llama3 = std::make_shared<LLaMA3>(llama3_config);
 
     // ========== pp_size：num_stages; vpp_size: num_chunks_per_stage ==========

infini_train/include/core/transformer/transformer_config.h

@@ -20,8 +20,7 @@ enum class NormType {
     kRMSNorm    // RMSNorm (LLaMA3 style)
 };
 
-class TransformerConfig {
-public:
+struct TransformerConfig {
     static constexpr char kGPT2Name[] = "GPT2";
     static constexpr char kLLaMA3Name[] = "LLaMA3";
 
@@ -61,7 +60,26 @@ class TransformerConfig {
     bool flash = false;             // flash attention
     int64_t max_gen_batch_size = 4; // max batch size during inference
 
-    static TransformerConfig GPT2() { return {}; }
+    static TransformerConfig GPT2() {
+        return {.model_type = kGPT2Name,
+                .block_size = 1024,
+                .vocab_size = 50304,
+                .original_vocab_size = 50257,
+                .n_layer = 12,
+                .n_head = 12,
+                .n_kv_head = 12,
+                .n_embd = 768,
+                .attention_type = AttentionType::kStandard,
+                .activation_type = MLPType::kGELU,
+                .norm_type = NormType::kLayerNorm,
+                .use_bias = true,
+                .use_gqa = false,
+                .use_rope = false,
+                .tie_weights = true,
+                .ffn_expansion_ratio = 4.0f,
+                .ffn_dim_multiplier = std::nullopt,
+                .multiple_of = 1};
+    }
 
     static TransformerConfig LLaMA3() {
         return {.model_type = kLLaMA3Name,
@@ -77,7 +95,10 @@ class TransformerConfig {
                 .use_bias = false,
                 .use_gqa = true,
                 .use_rope = true,
-                .tie_weights = false};
+                .tie_weights = false,
+                .ffn_expansion_ratio = 4.0f,
+                .ffn_dim_multiplier = 1.5f,
+                .multiple_of = 256};
     }
 };
 } // namespace infini_train::nn