refactor: remove module registry and use direct construction

  - Remove ModuleRegistry and INFINI_TRAIN_REGISTER_MODULE macros
  - Replace BuildModule() with direct constructor calls
  - Simplify module instantiation in MLP, CausalSelfAttention, and TransformerLayer

Author: JYMiracle305

CMakeLists.txt

@@ -204,5 +204,5 @@ link_infini_train_exe(test_precision_check)
 add_executable(test_lora test/lora/test_lora.cc)
 link_infini_train_exe(test_lora)
 
-# add_executable(test_transformer_spec test/transformer_spec/test_transformer_spec.cc)
-# link_infini_train_exe(test_transformer_spec)
+add_executable(test_transformer_architecture test/transformer/test_transformer_architecture.cc)
+link_infini_train_exe(test_transformer_architecture)

example/gpt2/checkpoint_loader.cc

@@ -17,7 +17,6 @@
 #include "infini_train/include/nn/modules/normalization.h"
 #include "infini_train/include/nn/modules/sparse.h"
 #include "infini_train/include/nn/modules/transformer/causal_self_attention.h"
-#include "infini_train/include/nn/modules/transformer/layer_specs.h"
 #include "infini_train/include/nn/modules/transformer/mlp.h"
 #include "infini_train/include/nn/modules/transformer/transformer.h"
 #include "infini_train/include/nn/parallel/global.h"
@@ -96,10 +95,7 @@ std::shared_ptr<nn::TransformerModel> LoadFromLLMC(const std::string &filepath)
     gpt2_config.n_layer = n_layer;
     gpt2_config.n_head = n_head;
     gpt2_config.n_embd = n_embd;
-    auto local_gpt2 = std::make_shared<nn::TransformerModel>(
-        gpt2_config,
-        nn::BuildTransformerSpec(gpt2_config, nn::BuildFirstStageSpec(gpt2_config),
-                                 nn::BuildTransformerLayerSpec(gpt2_config), nn::BuildLastStageSpec(gpt2_config)));
+    auto local_gpt2 = std::make_shared<nn::TransformerModel>(gpt2_config);
 
     LOG(INFO) << "magic: " << magic << " version: " << version << " block_size: " << block_size
               << " vocab_size: " << vocab_size << " n_layer: " << n_layer << " n_head: " << n_head
@@ -436,6 +432,7 @@ std::shared_ptr<nn::TransformerModel> LoadFromLLMC(const std::string &filepath)
         size_t ln_f_b_bytes = n_embd * sizeof(float);
         ifs.seekg(ln_f_w_bytes + ln_f_b_bytes, std::ios::cur);
     }
+
     return local_gpt2;
 }
 } // namespace gpt2

example/gpt2/main.cc

@@ -370,6 +370,7 @@ void Train(const nn::parallel::Rank &rank) {
                 y = std::make_shared<Tensor>(y->To(device));
 
                 LOG(INFO) << "Rank " << rank.GlobalRank() << ": start forward";
+
                 // (bs, seq_len, vocab_size)
                 auto logits = (*model)({x, y})[0];
                 LOG(INFO) << "Rank " << rank.GlobalRank() << ": finish model forward, start loss forward";

example/llama3/checkpoint_loader.cc

@@ -16,7 +16,6 @@
 #include "example/llama3/config.h"
 #include "infini_train/include/nn/modules/normalization.h"
 #include "infini_train/include/nn/modules/transformer/causal_self_attention.h"
-#include "infini_train/include/nn/modules/transformer/layer_specs.h"
 #include "infini_train/include/nn/modules/transformer/mlp.h"
 #include "infini_train/include/nn/modules/transformer/transformer.h"
 #include "infini_train/include/nn/parallel/global.h"
@@ -90,10 +89,7 @@ std::shared_ptr<nn::TransformerModel> LoadFromLLMC(const std::string &filepath)
     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<nn::TransformerModel>(
-        llama3_config,
-        nn::BuildTransformerSpec(llama3_config, nn::BuildFirstStageSpec(llama3_config),
-                                 nn::BuildTransformerLayerSpec(llama3_config), nn::BuildLastStageSpec(llama3_config)));
+    auto llama3 = std::make_shared<nn::TransformerModel>(llama3_config);
 
     // ========== pp_size：num_stages; vpp_size: num_chunks_per_stage ==========
     int pp_size = nn::parallel::global::GetPipelineParallelSize();

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

@@ -5,7 +5,6 @@
 #include <vector>
 
 #include "infini_train/include/nn/modules/module.h"
-#include "infini_train/include/nn/modules/transformer/spec_utils.h"
 #include "infini_train/include/nn/modules/transformer/transformer_config.h"
 
 namespace infini_train::nn {
@@ -18,7 +17,7 @@ class CausalSelfAttention : public infini_train::nn::CloneableModule<CausalSelfA
 
     static constexpr char kParamBiasName[] = "bias";
 
-    explicit CausalSelfAttention(const TransformerConfig &config, const ModuleSpec &spec = {});
+    explicit CausalSelfAttention(const TransformerConfig &config);
 
     std::vector<std::shared_ptr<infini_train::Tensor>>
     Forward(const std::vector<std::shared_ptr<infini_train::Tensor>> &x) override;

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

@@ -3,7 +3,6 @@
 #include <vector>
 
 #include "infini_train/include/nn/modules/module.h"
-#include "infini_train/include/nn/modules/transformer/spec_utils.h"
 #include "infini_train/include/nn/modules/transformer/transformer_config.h"
 
 namespace infini_train::nn {
@@ -17,7 +16,7 @@ class MLP : public infini_train::nn::CloneableModule<MLP> {
     static constexpr char kCFc2LayerName[] = "c_fc2";
     static constexpr char kSiluLayerName[] = "silu";
 
-    explicit MLP(const TransformerConfig &config, const ModuleSpec &spec = {});
+    explicit MLP(const TransformerConfig &config);
 
     std::vector<std::shared_ptr<infini_train::Tensor>>
     Forward(const std::vector<std::shared_ptr<infini_train::Tensor>> &x) override;

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

@@ -3,7 +3,6 @@
 #include <vector>
 
 #include "infini_train/include/nn/modules/module.h"
-#include "infini_train/include/nn/modules/transformer/spec_utils.h"
 #include "infini_train/include/nn/modules/transformer/transformer_config.h"
 #include "infini_train/include/nn/parallel/pp/pipeline_parallel.h"
 
@@ -16,7 +15,7 @@ class TransformerLayer : public CloneableModule<TransformerLayer> {
     static constexpr char kLn2LayerName[] = "ln_2";
     static constexpr char kMlpLayerName[] = "mlp";
 
-    explicit TransformerLayer(const TransformerConfig &config, const ModuleSpec &spec = {});
+    explicit TransformerLayer(const TransformerConfig &config);
 
     std::vector<std::shared_ptr<infini_train::Tensor>>
     Forward(const std::vector<std::shared_ptr<infini_train::Tensor>> &x) override;
@@ -28,14 +27,13 @@ class TransformerFirstStage : public CloneableModule<TransformerFirstStage> {
     static constexpr char kWTELayerName[] = "wte";
     static constexpr char kWPELayerName[] = "wpe";
 
-    explicit TransformerFirstStage(const TransformerConfig &config, const ModuleSpec &spec = {});
+    explicit TransformerFirstStage(const TransformerConfig &config);
 
     std::vector<std::shared_ptr<infini_train::Tensor>>
     Forward(const std::vector<std::shared_ptr<infini_train::Tensor>> &x) override;
 
 private:
     TransformerConfig config_;
-    ModuleSpec spec_;
 };
 
 class TransformerChunk : public CloneableModule<TransformerChunk> {
@@ -44,14 +42,13 @@ class TransformerChunk : public CloneableModule<TransformerChunk> {
     static constexpr char kHLayerName[] = "h";
     static constexpr char kFreqsCisName[] = "freqs_cis";
 
-    TransformerChunk(const TransformerConfig &config, int start_layer, int end_layer, const ModuleSpec &spec = {});
+    TransformerChunk(const TransformerConfig &config, int start_layer, int end_layer);
 
     std::vector<std::shared_ptr<infini_train::Tensor>>
     Forward(const std::vector<std::shared_ptr<infini_train::Tensor>> &x) override;
 
 private:
     const TransformerConfig config_;
-    ModuleSpec spec_;
 };
 
 class TransformerLastStage : public CloneableModule<TransformerLastStage> {
@@ -60,22 +57,21 @@ class TransformerLastStage : public CloneableModule<TransformerLastStage> {
     static constexpr char kLnFLayerName[] = "ln_f";
     static constexpr char kLMHeadLayerName[] = "lm_head";
 
-    explicit TransformerLastStage(const TransformerConfig &config, const ModuleSpec &spec = {});
+    explicit TransformerLastStage(const TransformerConfig &config);
 
     std::vector<std::shared_ptr<infini_train::Tensor>>
     Forward(const std::vector<std::shared_ptr<infini_train::Tensor>> &x) override;
 
 private:
     const TransformerConfig config_;
-    ModuleSpec spec_;
 };
 
 class TransformerModel : public CloneableModule<TransformerModel> {
 public:
     static constexpr char kType[] = "Transformer";
     static constexpr char kTransformerModelName[] = "transformer";
 
-    explicit TransformerModel(const TransformerConfig config, const ModuleSpec &spec = {});
+    explicit TransformerModel(const TransformerConfig config);
 
     std::vector<std::shared_ptr<infini_train::Tensor>>
     Forward(const std::vector<std::shared_ptr<infini_train::Tensor>> &x) override;
@@ -85,7 +81,6 @@ class TransformerModel : public CloneableModule<TransformerModel> {
 private:
     const TransformerConfig config_;
     const infini_train::nn::parallel::StageInfo stage_info_;
-    ModuleSpec spec_;
 };
 
 } // namespace infini_train::nn

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

@@ -11,26 +11,36 @@
 #include "infini_train/include/nn/init.h"
 #include "infini_train/include/nn/modules/normalization.h"
 #include "infini_train/include/nn/modules/sparse.h"
-#include "infini_train/include/nn/modules/transformer/spec_utils.h"
-#include "infini_train/include/nn/modules/transformer/transformer.h"
 #include "infini_train/include/nn/modules/transformer/transformer_config.h"
 #include "infini_train/include/nn/parallel/global.h"
 #include "infini_train/include/nn/parallel/tensor_parallel.h"
 #include "infini_train/include/tensor.h"
 
 namespace infini_train::nn {
 
-CausalSelfAttention::CausalSelfAttention(const TransformerConfig &config, const ModuleSpec &spec)
-    : CloneableModule(kType), config_(config) {
+CausalSelfAttention::CausalSelfAttention(const TransformerConfig &config) : CloneableModule(kType), config_(config) {
     SetupAttention(config);
 
-    CHECK(spec.submodules_.contains(kCAttnLayerName))
-        << "CausalSelfAttention spec missing submodule: " << kCAttnLayerName;
-    CHECK(spec.submodules_.contains(kCProjLayerName))
-        << "CausalSelfAttention spec missing submodule: " << kCProjLayerName;
-    // Build submodules from spec
-    modules_[kCAttnLayerName] = BuildModule(config, spec.submodules_.at(kCAttnLayerName));
-    modules_[kCProjLayerName] = BuildModule(config, spec.submodules_.at(kCProjLayerName));
+    int64_t qkv_dim = (config.n_head + 2 * n_kv_head_) * head_dim_;
+    // qkv: ColumnParallel (do not gather output)
+    modules_[kCAttnLayerName] = std::make_shared<nn::parallel::ColumnParallelLinear>(
+        /*in_features=*/n_embd_,
+        /*out_features=*/qkv_dim,
+        /*bias=*/config_.use_bias,
+        /*gather_output=*/false,
+        /*input_is_parallel=*/false,
+        /*skip_bias_add=*/false,
+        /*sequence_parallel=*/nn::parallel::global::GetSequenceParallelEnabled());
+
+    // proj: RowParallel (input is parallel and output is full)
+    modules_[kCProjLayerName] = std::make_shared<nn::parallel::RowParallelLinear>(
+        /*in_features=*/n_embd_,
+        /*out_features=*/n_embd_,
+        /*bias=*/config_.use_bias,
+        /*reduce_output=*/true,
+        /*input_is_parallel=*/true,
+        /*skip_bias_add=*/false,
+        /*sequence_parallel=*/nn::parallel::global::GetSequenceParallelEnabled());
 
     // For standard attention (GPT2 style), precompute causal mask
     if (config_.attention_type == AttentionType::kStandard) {