refactor: remove attention_type_ from TransformerLayer

Author: JYMiracle305

infini_train/include/core/transformer/spec_utils.h

@@ -84,8 +84,8 @@ template <typename T> T GetRequiredParam(const ModuleSpec &spec, const std::stri
     CHECK(spec.params_.contains(key)) << "Missing required parameter: " << key;
 
     const T *value = std::any_cast<T>(&spec.params_.at(key));
-    CHECK(value) << "Parameter type mismatch for key '" << key << "': expected " << typeid(T).name() << ", got "
-                 << spec.params_.at(key).type().name();
+    CHECK(value) << std::format("Parameter type mismatch for key '{}': expected {}, got {}", key, typeid(T).name(),
+                                spec.params_.at(key).type().name());
     return *value;
 }
 

infini_train/include/core/transformer/transformer_config.h

@@ -39,7 +39,6 @@ struct TransformerConfig {
     NormType norm_type = NormType::kLayerNorm;               // Normalization type
 
     bool use_bias = true;    // Linear layers bias (GPT2: true, LLaMA3: false)
-    bool use_gqa = false;    // Grouped Query Attention
     bool tie_weights = true; // Tie embedding and lm_head weights
 
     // FFN config
@@ -59,6 +58,8 @@ struct TransformerConfig {
     bool flash = false;             // flash attention
     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,
@@ -72,7 +73,6 @@ struct TransformerConfig {
                 .activation_type = MLPType::kGELU,
                 .norm_type = NormType::kLayerNorm,
                 .use_bias = true,
-                .use_gqa = false,
                 .tie_weights = true,
                 .ffn_expansion_ratio = 4.0f,
                 .ffn_dim_multiplier = std::nullopt,
@@ -91,7 +91,6 @@ struct TransformerConfig {
                 .activation_type = MLPType::kSwiGLU,
                 .norm_type = NormType::kRMSNorm,
                 .use_bias = false,
-                .use_gqa = true,
                 .tie_weights = false,
                 .ffn_expansion_ratio = 4.0f,
                 .ffn_dim_multiplier = 1.5f,

infini_train/include/nn/modules/transformer.h

@@ -19,9 +19,6 @@ class TransformerLayer : public infini_train::nn::CloneableModule<TransformerLay
 
     std::vector<std::shared_ptr<infini_train::Tensor>>
     Forward(const std::vector<std::shared_ptr<infini_train::Tensor>> &x) override;
-
-private:
-    AttentionType attention_type_ = AttentionType::kStandard;
 };
 
 } // namespace infini_train::nn

infini_train/src/core/transformer/transformer_builders.cc

@@ -41,7 +41,7 @@ ModuleSpec BuildAttentionSpec(const TransformerConfig &config) {
 
     // Calculate QKV output dimension based on attention type and GQA
     int64_t qkv_out;
-    if (config.use_gqa && config.n_kv_head < config.n_head) {
+    if (config.UseGQA()) {
         // GQA style (LLaMA3 with GQA enabled)
         int64_t head_dim = config.n_embd / config.n_head;
         // qkv_out = config.n_embd + 2 * config.n_kv_head * head_dim;

infini_train/src/nn/modules/causal_self_attention.cc

@@ -52,7 +52,7 @@ void CausalSelfAttention::SetupAttention(const TransformerConfig &config) {
     local_n_head_ = n_head_ / tp_world_size;
 
     // For GQA, set n_kv_head and n_rep
-    if (config.use_gqa && config.n_kv_head < config.n_head) {
+    if (config.UseGQA()) {
         CHECK_EQ(config.n_head % config.n_kv_head, 0) << "n_head must be divisible by n_kv_head for GQA";
         CHECK_EQ(config.n_kv_head % tp_world_size, 0) << "n_kv_head must be divisible by TP world size for GQA";
 

infini_train/src/nn/modules/transformer.cc

@@ -19,38 +19,37 @@
 
 namespace infini_train::nn {
 TransformerLayer::TransformerLayer(const nn::TransformerConfig &config, const ModuleSpec &spec)
-    : CloneableModule(kType), attention_type_(config.attention_type) {
+    : CloneableModule(kType) {
     modules_[kLn1LayerName] = BuildModule(config, spec.submodules_.at(kLn1LayerName));
     modules_[kAttnLayerName] = BuildModule(config, spec.submodules_.at(kAttnLayerName));
     modules_[kLn2LayerName] = BuildModule(config, spec.submodules_.at(kLn2LayerName));
     modules_[kMlpLayerName] = BuildModule(config, spec.submodules_.at(kMlpLayerName));
 }
 
-std::vector<std::shared_ptr<infini_train::Tensor>>
-TransformerLayer::Forward(const std::vector<std::shared_ptr<infini_train::Tensor>> &x) {
+std::vector<std::shared_ptr<Tensor>> TransformerLayer::Forward(const std::vector<std::shared_ptr<Tensor>> &x) {
     // (bs, seq_len, n_embd) -> Layernorm -> (bs, seq_len, n_embd)
     auto ln1_out = (*modules_[kLn1LayerName])({x[0]})[0];
 
-    std::shared_ptr<infini_train::Tensor> x1;
-    // Build attention input
-    if (attention_type_ == AttentionType::kRoPE) {
-        // LLaMA3: {ln1_out, freqs_cis, start_pos, mask}
-        const auto freqs_cis = x.size() > 1 ? x[1] : nullptr;
-        const auto start_pos = x.size() > 2 ? x[2] : nullptr;
-        const auto mask = x.size() > 3 ? x[3] : nullptr;
-        auto attn_out = (*modules_[kAttnLayerName])({ln1_out, freqs_cis, start_pos, mask})[0];
-        x1 = x[0] + attn_out;
-    } else {
-        // GPT2: {ln1_out}
-        auto attn_out = (*modules_[kAttnLayerName])({ln1_out})[0];
-        x1 = x[0] + attn_out;
+    std::vector<std::shared_ptr<Tensor>> attn_input = {ln1_out};
+    if (x.size() > 1) {
+        attn_input.push_back(x[1]); // freqs_cis
+    }
+    if (x.size() > 2) {
+        attn_input.push_back(x[2]); // start_pos
+    }
+    if (x.size() > 3) {
+        attn_input.push_back(x[3]); // mask
     }
 
-    // (bs, seq_len, n_embd) -> Layernorm -> (bs, seq_len, n_embd) -> MLP -> (bs, seq_len, n_embd)
-    // -> Add -> (bs, seq_len, n_embd)
+    auto attn_out = (*modules_[kAttnLayerName])(attn_input)[0];
+    auto x1 = x[0] + attn_out;
+
+    // (bs, seq_len, n_embd) -> Layernorm -> (bs, seq_len, n_embd) -> MLP -> (bs, seq_len, n_embd) -> Add -> (bs,
+    // seq_len, n_embd)
     auto x2 = x1 + (*modules_[kMlpLayerName])((*modules_[kLn2LayerName])({x1}))[0];
 
     // (bs, seq_len, n_embd)
     return {x2};
 }
+
 } // namespace infini_train::nn