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"
@@ -56,14 +55,6 @@ std::tuple<int32_t, infini_train::DataType> DetermineAndCheckVersion(const std::
 } // namespace
 
 namespace gpt2 {
-int GetChunkSize() {
-    nn::TransformerConfig gpt2_config = GPT2Config();
-
-    auto stage_info = nn::parallel::PipelineParallel::GetStageInfo(
-        gpt2_config.n_layer, nn::parallel::global::GetPipelineParallelSize(), nn::parallel::pp_rank,
-        nn::parallel::global::GetVirtualPipelineParallelSize());
-    return stage_info.layer_ranges_per_chunk.size();
-}
 
 std::shared_ptr<nn::TransformerModel> LoadFromLLMC(const std::string &filepath) {
     if (!std::filesystem::exists(filepath)) {
@@ -96,10 +87,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 +424,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/checkpoint_loader.h

@@ -5,11 +5,8 @@
 
 namespace infini_train::nn {
 class TransformerModel;
-enum class ModelType;
 } // namespace infini_train::nn
 
 namespace gpt2 {
-int GetChunkSize();
 std::shared_ptr<infini_train::nn::TransformerModel> LoadFromLLMC(const std::string &filepath);
-std::shared_ptr<infini_train::nn::TransformerModel> FromPretrained(infini_train::nn::ModelType model_type);
 } // namespace gpt2

example/gpt2/config.h

@@ -15,7 +15,8 @@ inline nn::TransformerConfig GPT2Config() {
             .attention_type = nn::AttentionType::kStandard,
             .activation_type = nn::MLPType::kGELU,
             .norm_type = nn::NormType::kLayerNorm,
-            .use_bias = true,
+            .add_bias_linear = true,
+            .add_bias_lm_head = false,
             .tie_weights = true,
             .ffn_expansion_ratio = 4.0f,
             .ffn_dim_multiplier = std::nullopt,

example/gpt2/main.cc

@@ -250,7 +250,7 @@ void Train(const nn::parallel::Rank &rank) {
             {FLAGS_batch_size, FLAGS_sequence_length / sp_world_size, model_config.n_embd}};
 
         model = std::make_shared<nn::parallel::PipelineParallel>(model, pp_world_size, num_micro_batches, shapes,
-                                                                 pp_rank, device, gpt2::GetChunkSize());
+                                                                 pp_rank, device, model_config.GetChunkSize());
         if (ddp_world_size > 1) {
             auto ddp_config
                 = DistributedDataParallelConfig{.use_distributed_optimizer = FLAGS_use_distributed_optimizer};
@@ -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"
@@ -40,15 +39,6 @@ constexpr int32_t kLLaMA3FP32Version = 3;
 
 namespace llama3 {
 
-int GetChunkSize() {
-    nn::TransformerConfig llama3_config = llama3::LLaMA3Config();
-
-    auto stage_info = nn::parallel::PipelineParallel::GetStageInfo(
-        llama3_config.n_layer, nn::parallel::global::GetPipelineParallelSize(), nn::parallel::pp_rank,
-        nn::parallel::global::GetVirtualPipelineParallelSize());
-    return stage_info.layer_ranges_per_chunk.size();
-}
-
 std::shared_ptr<nn::TransformerModel> LoadFromLLMC(const std::string &filepath) {
     if (!std::filesystem::exists(filepath)) {
         LOG(FATAL) << "File not found: " << filepath;
@@ -90,10 +80,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();

example/llama3/checkpoint_loader.h

@@ -8,6 +8,5 @@ class TransformerModel;
 } // namespace infini_train::nn
 
 namespace llama3 {
-int GetChunkSize();
 std::shared_ptr<infini_train::nn::TransformerModel> LoadFromLLMC(const std::string &filepath);
 } // namespace llama3

example/llama3/config.h

@@ -15,7 +15,8 @@ inline nn::TransformerConfig LLaMA3Config() {
             .attention_type = nn::AttentionType::kRoPE,
             .activation_type = nn::MLPType::kSwiGLU,
             .norm_type = nn::NormType::kRMSNorm,
-            .use_bias = false,
+            .add_bias_linear = false,
+            .add_bias_lm_head = false,
             .tie_weights = false,
             .ffn_expansion_ratio = 4.0f,
             .ffn_dim_multiplier = 1.5f,

example/llama3/main.cc

@@ -220,7 +220,7 @@ void Train(const nn::parallel::Rank &rank) {
             {FLAGS_batch_size, FLAGS_sequence_length / sp_world_size, model_config.n_embd}};
 
         model = std::make_shared<nn::parallel::PipelineParallel>(model, pp_world_size, num_micro_batches, shapes,
-                                                                 pp_rank, device, llama3::GetChunkSize());
+                                                                 pp_rank, device, model_config.GetChunkSize());
         if (ddp_world_size > 1) {
             auto ddp_config
                 = DistributedDataParallelConfig{.use_distributed_optimizer = FLAGS_use_distributed_optimizer};

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;