refactor: flatten transformer spec and extract model presets

JYMiracle305 · JYMiracle305 · commit 5cec43ff9de6 · 2026-04-01T16:12:43.000+08:00
diff --git a/example/gpt2/config.h b/example/gpt2/config.h
@@ -0,0 +1,27 @@
+#pragma once
+
+#include "infini_train/include/core/transformer/transformer_config.h"
+
+namespace nn = infini_train::nn;
+
+namespace infini_train::nn::gpt2 {
+inline nn::TransformerConfig GPT2Config() {
+    return {.model_type = ModelType::kGPT2,
+            .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 = nn::AttentionType::kStandard,
+            .activation_type = nn::MLPType::kGELU,
+            .norm_type = nn::NormType::kLayerNorm,
+            .use_bias = true,
+            .tie_weights = true,
+            .ffn_expansion_ratio = 4.0f,
+            .ffn_dim_multiplier = std::nullopt,
+            .multiple_of = 1};
+}
+
+} // namespace infini_train::nn::gpt2
diff --git a/example/gpt2/main.cc b/example/gpt2/main.cc
@@ -12,7 +12,6 @@
 #include "infini_train/include/autocast.h"
 #include "infini_train/include/core/models/decode_only_transformer/model.h"
 #include "infini_train/include/core/runtime/device_guard.h"
-#include "infini_train/include/core/transformer/transformer_config.h"
 #include "infini_train/include/dataloader.h"
 #include "infini_train/include/device.h"
 #include "infini_train/include/nn/lora/lora_utils.h"
@@ -37,6 +36,7 @@
 
 #include "example/common/tiny_shakespeare_dataset.h"
 #include "example/common/tokenizer.h"
+#include "example/gpt2/config.h"
 
 // I/O
 DEFINE_string(input_bin, "", "input .bin to train on");
@@ -188,7 +188,7 @@ void Train(const nn::parallel::Rank &rank) {
     // ManualSeed(42);
 
     // init the model, either from scratch or from OpenAI pretrained checkpoint
-    nn::TransformerConfig model_config = nn::TransformerConfig::GPT2();
+    nn::TransformerConfig model_config = nn::gpt2::GPT2Config();
     std::shared_ptr<nn::Module> model = nullptr;
 
     if (!FLAGS_llmc_filepath.empty()) {
diff --git a/example/gpt2/net.cc b/example/gpt2/net.cc
@@ -11,8 +11,8 @@
 #include "glog/logging.h"
 
 #include "example/common/utils.h"
+#include "example/gpt2/config.h"
 #include "infini_train/include/core/models/decode_only_transformer/model.h"
-#include "infini_train/include/core/transformer/transformer_config.h"
 #include "infini_train/include/nn/modules/causal_self_attention.h"
 #include "infini_train/include/nn/modules/mlp.h"
 #include "infini_train/include/nn/modules/normalization.h"
@@ -76,7 +76,7 @@ std::shared_ptr<DecoderOnlyTransformer> DecoderOnlyTransformer::FromLLMC_GPT2(co
     // NOTE(zbl): vocab_size needs to be padded to multiple of TP size
     const auto model_vocab_size = tp_size > 1 ? padded_vocab_size : vocab_size;
 
-    nn::TransformerConfig gpt2_config = nn::TransformerConfig::GPT2();
+    nn::TransformerConfig gpt2_config = nn::gpt2::GPT2Config();
     gpt2_config.block_size = block_size;
     gpt2_config.vocab_size = model_vocab_size;
     gpt2_config.original_vocab_size = vocab_size;
diff --git a/example/llama3/config.h b/example/llama3/config.h
@@ -0,0 +1,26 @@
+#pragma once
+
+#include "infini_train/include/core/transformer/transformer_config.h"
+
+namespace nn = infini_train::nn;
+
+namespace infini_train::nn::llama3 {
+inline nn::TransformerConfig LLaMA3Config() {
+    return {.model_type = ModelType::kLLaMA3,
+            .block_size = 8192,
+            .vocab_size = 128256,
+            .original_vocab_size = 128256,
+            .n_layer = 16,
+            .n_head = 32,
+            .n_kv_head = 8,
+            .n_embd = 2048,
+            .attention_type = nn::AttentionType::kRoPE,
+            .activation_type = nn::MLPType::kSwiGLU,
+            .norm_type = nn::NormType::kRMSNorm,
+            .use_bias = false,
+            .tie_weights = false,
+            .ffn_expansion_ratio = 4.0f,
+            .ffn_dim_multiplier = 1.5f,
+            .multiple_of = 256};
+}
+} // namespace infini_train::nn::llama3
diff --git a/example/llama3/main.cc b/example/llama3/main.cc
@@ -10,7 +10,6 @@
 #include "infini_train/include/autocast.h"
 #include "infini_train/include/core/models/decode_only_transformer/model.h"
 #include "infini_train/include/core/runtime/device_guard.h"
-#include "infini_train/include/core/transformer/transformer_config.h"
 #include "infini_train/include/dataloader.h"
 #include "infini_train/include/device.h"
 #include "infini_train/include/nn/lora/lora_utils.h"
@@ -36,6 +35,7 @@
 
 #include "example/common/tiny_shakespeare_dataset.h"
 #include "example/common/tokenizer.h"
+#include "example/llama3/config.h"
 
 // I/O
 DEFINE_string(input_bin, "", "input .bin to train on");
@@ -168,7 +168,7 @@ void Train(const nn::parallel::Rank &rank) {
     // rng / reproducibility
     // ManualSeed(42);
 
-    nn::TransformerConfig model_config = nn::TransformerConfig::LLaMA3();
+    nn::TransformerConfig model_config = nn::llama3::LLaMA3Config();
     std::shared_ptr<nn::Module> model = nullptr;
     if (!FLAGS_llmc_filepath.empty()) {
         model = DecoderOnlyTransformer::FromLLMC_LLaMA3(FLAGS_llmc_filepath);
diff --git a/example/llama3/net.cc b/example/llama3/net.cc
@@ -12,9 +12,9 @@
 #include "glog/logging.h"
 
 #include "example/common/utils.h"
+#include "example/llama3/config.h"
 #include "infini_train/include/core/models/decode_only_transformer/model.h"
 #include "infini_train/include/core/transformer/spec_utils.h"
-#include "infini_train/include/core/transformer/transformer_config.h"
 #include "infini_train/include/device.h"
 #include "infini_train/include/nn/modules/causal_self_attention.h"
 #include "infini_train/include/nn/modules/mlp.h"
@@ -65,7 +65,7 @@ std::shared_ptr<DecoderOnlyTransformer> DecoderOnlyTransformer::FromLLMC_LLaMA3(
     const auto version_major = BytesToType<int32_t>(header, 56);
     const auto version_minor = BytesToType<int32_t>(header, 60);
 
-    nn::TransformerConfig llama3_config = nn::TransformerConfig::LLaMA3();
+    nn::TransformerConfig llama3_config = nn::llama3::LLaMA3Config();
     llama3_config.block_size = block_size;
     llama3_config.vocab_size = vocab_size;
     llama3_config.n_layer = n_layer;
diff --git a/infini_train/include/core/models/decode_only_transformer/layer_specs.h b/infini_train/include/core/models/decode_only_transformer/layer_specs.h
@@ -4,9 +4,7 @@
 #include "infini_train/include/core/transformer/transformer_config.h"
 
 namespace infini_train::nn {
-// Build GPT2 model spec: LayerNorm + GELU + standard attention
-ModuleSpec BuildGPT2Spec(const TransformerConfig &config);
 
-// Build LLaMA3 model spec: RMSNorm + SwiGLU + RoPE + GQA
-ModuleSpec BuildLLaMA3Spec(const TransformerConfig &config);
+ModuleSpec BuildDecoderOnlyTransformerSpec(const TransformerConfig &config, ModuleSpec first_stage, ModuleSpec chunk,
+                                           ModuleSpec last_stage);
 } // namespace infini_train::nn
diff --git a/infini_train/include/core/models/decode_only_transformer/model.h b/infini_train/include/core/models/decode_only_transformer/model.h
@@ -8,6 +8,7 @@
 
 #include "infini_train/include/core/models/decode_only_transformer/layer_specs.h"
 #include "infini_train/include/core/transformer/spec_utils.h"
+#include "infini_train/include/core/transformer/transformer_builders.h"
 #include "infini_train/include/core/transformer/transformer_config.h"
 #include "infini_train/include/core/transformer/transformer_model.h"
 #include "infini_train/include/nn/parallel/global.h"
@@ -37,7 +38,9 @@ class DecoderOnlyTransformer : public nn::TransformerModel {
     };
 
     explicit DecoderOnlyTransformer(const nn::TransformerConfig &config)
-        : TransformerModel(config, BuildModelSpec(config)),
+        : TransformerModel(config, nn::BuildDecoderOnlyTransformerSpec(config, nn::BuildFirstStageSpec(config),
+                                                                       nn::BuildTransformerLayerSpec(config),
+                                                                       nn::BuildLastStageSpec(config))),
           stage_info_(nn::parallel::PipelineParallel::GetStageInfo(
               Config().n_layer, nn::parallel::global::GetPipelineParallelSize(), nn::parallel::pp_rank,
               nn::parallel::global::GetVirtualPipelineParallelSize())) {}
@@ -52,10 +55,5 @@ class DecoderOnlyTransformer : public nn::TransformerModel {
     int GetChunkSize() const;
 
 private:
-    static nn::ModuleSpec BuildModelSpec(const nn::TransformerConfig &config) {
-        return (config.model_type == nn::TransformerConfig::kGPT2Name) ? BuildGPT2Spec(config)
-                                                                       : BuildLLaMA3Spec(config);
-    }
-
     const infini_train::nn::parallel::StageInfo stage_info_;
 };
diff --git a/infini_train/include/core/transformer/transformer_builders.h b/infini_train/include/core/transformer/transformer_builders.h
@@ -47,4 +47,10 @@ ModuleSpec BuildPositionEmbeddingSpec(int64_t num_embeddings, int64_t embedding_
 // Build ColumnParallelLinear spec for output projection (lm_head)
 ModuleSpec BuildOutputProjSpec(const TransformerConfig &config, int64_t output_size, bool use_bias);
 
+ModuleSpec BuildFirstStageSpec(const TransformerConfig &config);
+
+ModuleSpec BuildChunkSpec(const TransformerConfig &config, int start_layer, int end_layer);
+
+ModuleSpec BuildLastStageSpec(const TransformerConfig &config);
+
 } // namespace infini_train::nn
diff --git a/infini_train/include/core/transformer/transformer_config.h b/infini_train/include/core/transformer/transformer_config.h
@@ -5,26 +5,28 @@
 
 namespace infini_train::nn {
 
+enum class ModelType {
+    kGPT2,   // GPT-2
+    kLLaMA3, // LLaMA3
+};
+
 enum class AttentionType {
-    kStandard, // Standard attention (GPT2 style, no RoPE)
-    kRoPE      // Rotary Position Embedding (LLaMA3 style)
+    kStandard, // Standard attention
+    kRoPE      // Rotary Position Embedding
 };
 
 enum class MLPType {
-    kGELU,  // GELU activation (GPT2 style)
-    kSwiGLU // SwiGLU activation (LLaMA3 style)
+    kGELU,  // GELU activation
+    kSwiGLU // SwiGLU activation
 };
 
 enum class NormType {
-    kLayerNorm, // LayerNorm (GPT2 style)
-    kRMSNorm    // RMSNorm (LLaMA3 style)
+    kLayerNorm, // LayerNorm
+    kRMSNorm    // RMSNorm
 };
 
 struct TransformerConfig {
-    static constexpr char kGPT2Name[] = "GPT2";
-    static constexpr char kLLaMA3Name[] = "LLaMA3";
-
-    std::string model_type = "";
+    ModelType model_type = ModelType::kGPT2;
 
     int64_t block_size = 1024;           // Max seq_len
     int64_t vocab_size = 50304;          // Vocab size
@@ -59,42 +61,5 @@ struct TransformerConfig {
     int64_t max_gen_batch_size = 4; // max batch size during inference
 
     bool UseGQA() const { return n_kv_head < n_head; }
-
-    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,
-                .tie_weights = true,
-                .ffn_expansion_ratio = 4.0f,
-                .ffn_dim_multiplier = std::nullopt,
-                .multiple_of = 1};
-    }
-
-    static TransformerConfig LLaMA3() {
-        return {.model_type = kLLaMA3Name,
-                .block_size = 8192,
-                .vocab_size = 128256,
-                .n_layer = 16,
-                .n_head = 32,
-                .n_kv_head = 8,
-                .n_embd = 2048,
-                .attention_type = AttentionType::kRoPE,
-                .activation_type = MLPType::kSwiGLU,
-                .norm_type = NormType::kRMSNorm,
-                .use_bias = false,
-                .tie_weights = false,
-                .ffn_expansion_ratio = 4.0f,
-                .ffn_dim_multiplier = 1.5f,
-                .multiple_of = 256};
-    }
 };
 } // namespace infini_train::nn
diff --git a/infini_train/src/core/models/decode_only_transformer/layer_specs.cc b/infini_train/src/core/models/decode_only_transformer/layer_specs.cc
@@ -1,7 +1,6 @@
 #include "infini_train/include/core/models/decode_only_transformer/layer_specs.h"
 
 #include <cmath>
-#include <memory>
 
 #include "infini_train/include/core/transformer/spec_utils.h"
 #include "infini_train/include/core/transformer/transformer_builders.h"
@@ -10,55 +9,13 @@
 
 namespace infini_train::nn {
 
-ModuleSpec BuildGPT2Spec(const TransformerConfig &config) {
-    // Configure for GPT2 architecture
-    TransformerConfig gpt2_config = config;
-    ModuleSpec spec;
-
-    // ===== First Stage =====
-    ModuleSpec first_stage;
-    first_stage.WithSubmodule(TransformerFirstStage::kWTELayerName, BuildVocabEmbeddingSpec(gpt2_config))
-        .WithSubmodule(TransformerFirstStage::kWPELayerName,
-                       BuildPositionEmbeddingSpec(gpt2_config.block_size, gpt2_config.n_embd));
-    spec.WithSubmodule(TransformerFirstStage::kType, first_stage);
-
-    // ===== Transformer Layer =====
-    ModuleSpec block = BuildTransformerLayerSpec(gpt2_config);
-    spec.WithSubmodule(TransformerLayer::kType, block);
-
-    // ===== Last Stage =====
-    ModuleSpec last_stage;
-    last_stage.WithSubmodule(TransformerLastStage::kLnFLayerName, BuildNormSpec(gpt2_config))
-        .WithSubmodule(TransformerLastStage::kLMHeadLayerName,
-                       BuildOutputProjSpec(gpt2_config, gpt2_config.vocab_size, false));
-    spec.WithSubmodule(TransformerLastStage::kType, last_stage);
+ModuleSpec BuildDecoderOnlyTransformerSpec(const TransformerConfig &config, ModuleSpec first_stage, ModuleSpec layer,
+                                           ModuleSpec last_stage) {
+    ModuleSpec spec(typeid(TransformerModel));
+    spec.WithSubmodule(TransformerFirstStage::kType, first_stage)
+        .WithSubmodule(TransformerLayer::kType, layer)
+        .WithSubmodule(TransformerLastStage::kType, last_stage);
 
     return spec;
 }
-
-ModuleSpec BuildLLaMA3Spec(const TransformerConfig &config) {
-    // Configure for LLaMA3 architecture
-    TransformerConfig llama3_config = config;
-    ModuleSpec spec;
-
-    // ===== First Stage =====
-    ModuleSpec first_stage;
-    // LLaMA3 only has token embedding, no position embedding (uses RoPE)
-    first_stage.WithSubmodule(TransformerFirstStage::kWTELayerName, BuildVocabEmbeddingSpec(llama3_config));
-    spec.WithSubmodule(TransformerFirstStage::kType, first_stage);
-
-    // ===== Transformer Layer =====
-    ModuleSpec block = BuildTransformerLayerSpec(llama3_config);
-    spec.WithSubmodule(TransformerLayer::kType, block);
-
-    // ===== Last Stage =====
-    ModuleSpec last_stage;
-    last_stage.WithSubmodule(TransformerLastStage::kLnFLayerName, BuildNormSpec(llama3_config))
-        .WithSubmodule(TransformerLastStage::kLMHeadLayerName,
-                       BuildOutputProjSpec(llama3_config, llama3_config.vocab_size, false));
-    spec.WithSubmodule(TransformerLastStage::kType, last_stage);
-
-    return spec;
-}
-
 } // namespace infini_train::nn
diff --git a/infini_train/src/core/transformer/transformer_builders.cc b/infini_train/src/core/transformer/transformer_builders.cc
diff --git a/infini_train/src/core/transformer/transformer_model.cc b/infini_train/src/core/transformer/transformer_model.cc
diff --git a/test/transformer_spec/test_transformer_spec.cc b/test/transformer_spec/test_transformer_spec.cc
