From edf32161f3d05381f6b16675f822dbc59bf871bf Mon Sep 17 00:00:00 2001 From: bolunz Date: Wed, 22 Jul 2026 16:12:42 +0800 Subject: [PATCH 1/2] refactor: merge 2 MHA paths, rename attention_type to position_embedding_type --- example/gpt2/config.h | 5 +- example/llama3/config.h | 4 +- example/mixtral/config.h | 5 +- .../transformer/causal_self_attention.h | 15 --- .../modules/transformer/transformer_config.h | 18 ++- .../include/nn/modules/transformer/utils.h | 6 + .../transformer/causal_self_attention.cc | 123 +++--------------- .../src/nn/modules/transformer/transformer.cc | 22 ++-- .../src/nn/modules/transformer/utils.cc | 38 ++++++ .../test_transformer_architecture.cc | 7 +- 10 files changed, 97 insertions(+), 146 deletions(-) diff --git a/example/gpt2/config.h b/example/gpt2/config.h index 078f9fd5..71cc0a56 100644 --- a/example/gpt2/config.h +++ b/example/gpt2/config.h @@ -14,7 +14,7 @@ inline nn::TransformerConfig GPT2Config() { .n_head = 12, .n_kv_head = 12, .n_embd = 768, - .attention_type = nn::AttentionType::kStandard, + .position_embedding_type = nn::PositionEmbeddingType::kLearnedAbsolute, .activation_type = nn::MLPType::kGELU, .norm_type = nn::NormType::kLayerNorm, .add_bias_linear = true, @@ -34,7 +34,8 @@ inline void SanitizeGPT2Config(const nn::TransformerConfig &c) { CHECK_GT(c.n_embd, 0); CHECK_EQ(c.n_embd % c.n_head, 0) << "n_embd must be divisible by n_head"; CHECK_EQ(c.n_kv_head, c.n_head) << "GPT-2 does not use GQA; n_kv_head must equal n_head"; - CHECK(c.attention_type == nn::AttentionType::kStandard) << "GPT-2 requires standard attention"; + CHECK(c.position_embedding_type == nn::PositionEmbeddingType::kLearnedAbsolute) + << "GPT-2 requires learned absolute position embedding"; CHECK(c.activation_type == nn::MLPType::kGELU) << "GPT-2 requires GELU activation"; CHECK(c.norm_type == nn::NormType::kLayerNorm) << "GPT-2 requires LayerNorm"; } diff --git a/example/llama3/config.h b/example/llama3/config.h index 6bc9124d..67ebd31f 100644 --- a/example/llama3/config.h +++ b/example/llama3/config.h @@ -14,7 +14,7 @@ inline nn::TransformerConfig LLaMA3Config() { .n_head = 32, .n_kv_head = 8, .n_embd = 2048, - .attention_type = nn::AttentionType::kRoPE, + .position_embedding_type = nn::PositionEmbeddingType::kRoPE, .activation_type = nn::MLPType::kSwiGLU, .norm_type = nn::NormType::kRMSNorm, .add_bias_linear = false, @@ -36,7 +36,7 @@ inline void SanitizeLLaMA3Config(const nn::TransformerConfig &c) { CHECK_EQ(c.n_head % c.n_kv_head, 0) << "n_head must be divisible by n_kv_head for GQA"; CHECK_GT(c.n_embd, 0); CHECK_EQ(c.n_embd % c.n_head, 0) << "n_embd must be divisible by n_head"; - CHECK(c.attention_type == nn::AttentionType::kRoPE) << "LLaMA-3 requires RoPE attention"; + CHECK(c.position_embedding_type == nn::PositionEmbeddingType::kRoPE) << "LLaMA-3 requires RoPE position embedding"; CHECK(c.activation_type == nn::MLPType::kSwiGLU) << "LLaMA-3 requires SwiGLU activation"; CHECK(c.norm_type == nn::NormType::kRMSNorm) << "LLaMA-3 requires RMSNorm"; CHECK(!c.add_bias_linear) << "LLaMA-3 has no bias in linear layers"; diff --git a/example/mixtral/config.h b/example/mixtral/config.h index f572cd6c..610d610d 100644 --- a/example/mixtral/config.h +++ b/example/mixtral/config.h @@ -17,7 +17,7 @@ inline nn::TransformerConfig TinyMixtralConfig() { config.n_head = 4; // Scaled down; preserves Mixtral 4:1 GQA ratio. config.n_kv_head = 1; // Scaled down with n_head; real Mixtral uses 8 KV heads. config.n_embd = 512; // Scaled down from Mixtral/Megatron --hidden-size 4096. - config.attention_type = nn::AttentionType::kRoPE; + config.position_embedding_type = nn::PositionEmbeddingType::kRoPE; config.activation_type = nn::MLPType::kSwiGLU; config.ffn_type = nn::FFNType::kMoE; config.norm_type = nn::NormType::kRMSNorm; @@ -51,7 +51,8 @@ inline void SanitizeTinyMixtralConfig(const nn::TransformerConfig &c) { CHECK_EQ(c.n_head % c.n_kv_head, 0) << "n_head must be divisible by n_kv_head for GQA"; CHECK_GT(c.n_embd, 0); CHECK_EQ(c.n_embd % c.n_head, 0) << "n_embd must be divisible by n_head"; - CHECK(c.attention_type == nn::AttentionType::kRoPE) << "tiny Mixtral requires RoPE attention"; + CHECK(c.position_embedding_type == nn::PositionEmbeddingType::kRoPE) + << "tiny Mixtral requires RoPE position embedding"; CHECK(c.activation_type == nn::MLPType::kSwiGLU) << "tiny Mixtral requires SwiGLU activation"; CHECK(c.ffn_type == nn::FFNType::kMoE) << "tiny Mixtral requires MoE FFN"; CHECK(c.norm_type == nn::NormType::kRMSNorm) << "tiny Mixtral requires RMSNorm"; diff --git a/infini_train/include/nn/modules/transformer/causal_self_attention.h b/infini_train/include/nn/modules/transformer/causal_self_attention.h index 5ac55e31..60373414 100644 --- a/infini_train/include/nn/modules/transformer/causal_self_attention.h +++ b/infini_train/include/nn/modules/transformer/causal_self_attention.h @@ -1,7 +1,6 @@ #pragma once #include -#include #include #include "infini_train/include/nn/modules/module.h" @@ -35,20 +34,6 @@ class CausalSelfAttention : public infini_train::nn::CloneableModule> - ForwardStandard(const std::vector> &x); - - // RoPE-aware attention forward (LLaMA3 style: with RoPE, optional GQA) - std::vector> - ForwardWithRoPE(const std::vector> &x); - - // RoPE helper methods - std::tuple, std::shared_ptr> - ApplyRotaryEmbedding(const std::shared_ptr &xq, - const std::shared_ptr &xk, - const std::shared_ptr &freqs_cis); - // GQA helper method std::shared_ptr RepeatKV(const std::shared_ptr &x, int64_t n_rep); }; diff --git a/infini_train/include/nn/modules/transformer/transformer_config.h b/infini_train/include/nn/modules/transformer/transformer_config.h index 7460ca26..a646ab14 100644 --- a/infini_train/include/nn/modules/transformer/transformer_config.h +++ b/infini_train/include/nn/modules/transformer/transformer_config.h @@ -5,9 +5,13 @@ namespace infini_train::nn { -enum class AttentionType { - kStandard, // Standard attention - kRoPE // Rotary Position Embedding +enum class PositionEmbeddingType { + kLearnedAbsolute, // Megatron: learned_absolute + kRoPE, // Megatron: rope + kYarn, // Megatron: yarn + kMRoPE, // Megatron: mrope + kRelative, // Megatron: relative + kNone // Megatron: none }; enum class MLPType { @@ -59,10 +63,10 @@ struct TransformerConfig { int64_t n_kv_head = 12; // Num of Key/Value heads (<= n_head, < n_head if using GQA) int64_t n_embd = 768; // Hidden size - AttentionType attention_type = AttentionType::kStandard; // Attention mechanism type - MLPType activation_type = MLPType::kGELU; // MLP activation type - FFNType ffn_type = FFNType::kDense; // Feed-forward module type - NormType norm_type = NormType::kLayerNorm; // Normalization type + PositionEmbeddingType position_embedding_type = PositionEmbeddingType::kLearnedAbsolute; // Position embedding type + MLPType activation_type = MLPType::kGELU; // MLP activation type + FFNType ffn_type = FFNType::kDense; // Feed-forward module type + NormType norm_type = NormType::kLayerNorm; // Normalization type bool add_bias_linear = true; // Whether to add learnable bias to all Linear layers in the Transformer block, // including: attention QKV projection, attention output projection, MLP FC layers (and diff --git a/infini_train/include/nn/modules/transformer/utils.h b/infini_train/include/nn/modules/transformer/utils.h index d3a62c63..30db08e6 100644 --- a/infini_train/include/nn/modules/transformer/utils.h +++ b/infini_train/include/nn/modules/transformer/utils.h @@ -1,6 +1,8 @@ #pragma once #include +#include +#include #include "infini_train/include/tensor.h" @@ -8,4 +10,8 @@ namespace infini_train { // RoPE helper method std::shared_ptr PrecomputeFreqsCis(int64_t dim, int64_t end, float theta = 10000.0f, bool use_scaled = false, Device device = Device()); + +std::tuple, std::shared_ptr> +ApplyRotaryEmbedding(const std::shared_ptr &xq, const std::shared_ptr &xk, + const std::shared_ptr &freqs_cis); } // namespace infini_train diff --git a/infini_train/src/nn/modules/transformer/causal_self_attention.cc b/infini_train/src/nn/modules/transformer/causal_self_attention.cc index 8bed8193..7f20fad1 100644 --- a/infini_train/src/nn/modules/transformer/causal_self_attention.cc +++ b/infini_train/src/nn/modules/transformer/causal_self_attention.cc @@ -1,6 +1,7 @@ #include "infini_train/include/nn/modules/transformer/causal_self_attention.h" #include +#include #include #include #include @@ -12,6 +13,7 @@ #include "infini_train/include/nn/modules/normalization.h" #include "infini_train/include/nn/modules/sparse.h" #include "infini_train/include/nn/modules/transformer/transformer_config.h" +#include "infini_train/include/nn/modules/transformer/utils.h" #include "infini_train/include/nn/parallel/global.h" #include "infini_train/include/nn/parallel/tensor_parallel.h" #include "infini_train/include/tensor.h" @@ -42,12 +44,9 @@ CausalSelfAttention::CausalSelfAttention(const TransformerConfig &config) : Clon /*skip_bias_add=*/false, /*sequence_parallel=*/nn::parallel::global::GetSequenceParallelEnabled()); - // For standard attention (GPT2 style), precompute causal mask - if (config_.attention_type == AttentionType::kStandard) { - // causal mask: (1, 1, block_size, block_size) - buffers_[kParamBiasName] = function::Tril(nn::function::Ones({config_.block_size, config_.block_size})) - ->View({1, 1, config_.block_size, config_.block_size}); - } + // causal mask: (1, 1, block_size, block_size) + buffers_[kParamBiasName] = function::Tril(nn::function::Ones({config_.block_size, config_.block_size})) + ->View({1, 1, config_.block_size, config_.block_size}); } void CausalSelfAttention::SetupAttention(const TransformerConfig &config) { @@ -74,98 +73,6 @@ void CausalSelfAttention::SetupAttention(const TransformerConfig &config) { } } -std::vector> -CausalSelfAttention::Forward(const std::vector> &x) { - if (config_.attention_type == AttentionType::kRoPE) { - return ForwardWithRoPE(x); - } else { - return ForwardStandard(x); - } -} - -std::vector> -CausalSelfAttention::ForwardStandard(const std::vector> &x) { - auto tp_world_size = parallel::global::GetTensorParallelSize(); - - const auto B = x[0]->Dims()[0]; // bs - const int64_t head_dim = n_embd_ / n_head_; // per-head dim (global) - const int64_t local_C = n_embd_ / tp_world_size; // per-rank hidden - - // (B, T, C) -> ColumnParallelLinear(C, 3*C) -> (B, T, 3 * local_C) - // -> Split -> (3, B, T, local_C) - auto qkv = (*modules_[kCAttnLayerName])(x)[0]->Split(local_C, 2); - - // (B, T, local_C) - auto q = qkv[0]; - auto k = qkv[1]; - auto v = qkv[2]; - - // NOTE(zbl): Acquire full T after AllGather is performed in ColumnParallelLinear - const auto T = q->Dims()[1]; - - // View to multi-head: local_n_head * head_dim == local_C - // (B, T, local_C) -> (B, T, h_l, Dh) -> (B, h_l, T, Dh) - k = k->View({B, T, local_n_head_, head_dim})->Transpose(1, 2); - q = q->View({B, T, local_n_head_, head_dim})->Transpose(1, 2); - v = v->View({B, T, local_n_head_, head_dim})->Transpose(1, 2); - - // (B, h_l, T, T) - auto att = q->Matmul(k->Transpose(-2, -1)) * (1.0 / std::sqrt(head_dim)); - // (1, 1, T, T) - auto mask = buffers_[kParamBiasName]->Slice({0, 0, 0, 0}, {1, 1, T, T}, {1, 1, 1, 1}); - // (1, 1, T, T) -> eq 0 -> (1, 1, T, T) -> masked_fill -> (B, h_l, T, T) - att = att->MaskedFill(mask == 0, -std::numeric_limits::infinity()); - // (B, h_l, T, T) - att = nn::function::Softmax(att, -1); - // (B, h_l, T, Dh) - auto y = att->Matmul(v); - // (B, h_l, T, Dh) -> (B, T, h_l, Dh) -> (B, T, local_C) - y = y->Transpose(1, 2)->Contiguous()->View({B, T, local_C}); - - // Get full tensor - // (B, T, local_C) -> RowParallelLinear(n_embd, n_embd) -> (B, T, C) - y = (*modules_[kCProjLayerName])({y})[0]; - // (B, T, C) == (bs, seq_len, n_embd) - return {y}; -} - -// RoPE helper methods -std::tuple, std::shared_ptr> -CausalSelfAttention::ApplyRotaryEmbedding(const std::shared_ptr &xq, - const std::shared_ptr &xk, - const std::shared_ptr &freqs_cis) { - // Reshape freqs_cis for broadcasting - const auto &x_shape = xq->Dims(); // (B, T, H, D) - const int64_t T = x_shape[1]; - const int64_t D = x_shape[3]; - - std::vector target_shape = {1, T, 1, D / 2, 2}; - auto cos_sin = freqs_cis->View(target_shape); // -> (1, T, 1, D/2, 2) - - auto cos = cos_sin->Slice(-1, 0, 1, 1)->Squeeze(-1); // (1, T, 1, D/2) - auto sin = cos_sin->Slice(-1, 1, 2, 1)->Squeeze(-1); // (1, T, 1, D/2) - - auto slice_pair = [](const std::shared_ptr &x) { - auto even = x->Slice(-1, 0, x->Dims().back(), 2); - auto odd = x->Slice(-1, 1, x->Dims().back(), 2); - return std::make_pair(even, odd); - }; - - auto [q_even, q_odd] = slice_pair(xq); - auto q_rotated_left = q_even * cos - q_odd * sin; - auto q_rotated_right = q_even * sin + q_odd * cos; - auto q_rotated - = nn::function::Stack(std::vector>{q_rotated_left, q_rotated_right}, -1)->Flatten(-2); - - auto [k_even, k_odd] = slice_pair(xk); - auto k_rotated_left = k_even * cos - k_odd * sin; - auto k_rotated_right = k_even * sin + k_odd * cos; - auto k_rotated - = nn::function::Stack(std::vector>{k_rotated_left, k_rotated_right}, -1)->Flatten(-2); - - return {q_rotated, k_rotated}; -} - std::shared_ptr CausalSelfAttention::RepeatKV(const std::shared_ptr &x, int64_t n_rep) { const auto &shape = x->Dims(); @@ -179,21 +86,23 @@ std::shared_ptr CausalSelfAttention::RepeatKV(const std::s } std::vector> -CausalSelfAttention::ForwardWithRoPE(const std::vector> &x) { +CausalSelfAttention::Forward(const std::vector> &x) { const auto B = x[0]->Dims()[0]; // bs const auto C = x[0]->Dims()[2]; // n_embd const auto tp_size = nn::parallel::global::GetTensorParallelSize(); const auto C_local = C / tp_size; - const auto H_local = n_head_ / tp_size; + const auto H_local = local_n_head_; const auto KV_local = n_kv_head_ / tp_size; const auto D = head_dim_; // n_embd / n_head 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; - CHECK(freqs_cis != nullptr) << "freqs_cis is null."; + if (config_.position_embedding_type == PositionEmbeddingType::kRoPE) { + CHECK(freqs_cis != nullptr) << "freqs_cis is null."; + } // (B, T, C) -> (B, T, (H + 2 * n_kv_head) * D) auto qkv = (*modules_[kCAttnLayerName])({x[0]})[0]; @@ -211,10 +120,10 @@ CausalSelfAttention::ForwardWithRoPE(const std::vectorSlice(2, q_size_local + kv_size_local, q_size_local + 2 * kv_size_local)->View({B, T, KV_local, D}); - // -> RoPE on q, k - // q: (B, T, H_local, D) - // k: (B, T, KV_local, D) - std::tie(q, k) = ApplyRotaryEmbedding(q, k, freqs_cis); + if (config_.position_embedding_type == PositionEmbeddingType::kRoPE) { + // q: (B, T, H_local, D), k: (B, T, KV_local, D) + std::tie(q, k) = ApplyRotaryEmbedding(q, k, freqs_cis); + } // TODO(zbl): use kv cache during inference // if (use_kv_) { ... } @@ -242,6 +151,10 @@ CausalSelfAttention::ForwardWithRoPE(const std::vectorMaskedFill(mask, std::numeric_limits::lowest()); + } else { + // fallback causal mask: (1, 1, T, T) + auto causal_mask = buffers_[kParamBiasName]->Slice({0, 0, 0, 0}, {1, 1, T, T}, {1, 1, 1, 1}); + att = att->MaskedFill(causal_mask == 0, -std::numeric_limits::infinity()); } // (B, H_local, T, T) att = nn::function::Softmax(att, -1); diff --git a/infini_train/src/nn/modules/transformer/transformer.cc b/infini_train/src/nn/modules/transformer/transformer.cc index bdcde449..99a739d2 100644 --- a/infini_train/src/nn/modules/transformer/transformer.cc +++ b/infini_train/src/nn/modules/transformer/transformer.cc @@ -29,9 +29,11 @@ TransformerFirstStage::TransformerFirstStage(const TransformerConfig &config) modules_[kWTELayerName] = std::make_shared( config_.vocab_size, config_.n_embd, parallel::global::GetSequenceParallelEnabled()); - // LLaMA3 use RoPE, so they don't need position embedding - if (config_.activation_type == MLPType::kGELU) { + // Only learned absolute position embedding uses a trainable WPE table. + if (config_.position_embedding_type == PositionEmbeddingType::kLearnedAbsolute) { modules_[kWPELayerName] = std::make_shared(config_.block_size, config_.n_embd); + } else if (config_.position_embedding_type != PositionEmbeddingType::kRoPE) { + LOG(FATAL) << "Unsupported position embedding type"; } } @@ -45,7 +47,7 @@ std::vector> TransformerFirstStage::Forward(const std::v // (B, T) -> Embedding(V_local, C) -> (B, T, C) auto tok_emb = (*modules_[kWTELayerName])({x1}); - // Add position embedding only for models that use absolute position encoding + // Add position embedding only for models that use learned absolute position encoding. if (modules_.contains(kWPELayerName)) { // (T_local) // NOTE(zbl): Slice pos sequence when SP is enabled @@ -66,7 +68,7 @@ std::vector> TransformerFirstStage::Forward(const std::v // (B, T, C) return {tok_emb[0] + pos_emb[0]}; } else { - // For RoPE-based models (LLaMA3), no position embedding needed + // For RoPE-based models (LLaMA3), no absolute position embedding is needed. // (B, T, C) return tok_emb; } @@ -133,8 +135,8 @@ TransformerChunk::TransformerChunk(const TransformerConfig &config, int start_la std::vector> TransformerChunk::Forward(const std::vector> &x) { auto x1 = x[0]; - // Check if we need to pass RoPE parameters (for LLaMA3 style models) - if (config_.attention_type == AttentionType::kRoPE) { + // Check if we need to pass RoPE parameters (for LLaMA3 style models). + if (config_.position_embedding_type == PositionEmbeddingType::kRoPE) { // For RoPE models, we need to prepare freqs_cis and potentially other parameters const auto device = x1->GetDevice(); @@ -161,9 +163,11 @@ std::vector> TransformerChunk::Forward(const std::vector for (auto &h : *std::dynamic_pointer_cast(modules_[kHLayerName])) { x1 = (*h)({x1, freqs_view, start_pos_ptr, mask})[0]; } - } else { - // Standard attention (GPT2 style) + } else if (config_.position_embedding_type == PositionEmbeddingType::kLearnedAbsolute) { + // Learned absolute position embedding models (GPT-2 style). for (auto &h : *std::dynamic_pointer_cast(modules_[kHLayerName])) { x1 = (*h)({x1})[0]; } + } else { + LOG(FATAL) << "Unsupported position embedding type"; } return {x1}; @@ -217,7 +221,7 @@ TransformerModel::TransformerModel(const TransformerConfig config) modules_[kPPFirstStageName] = std::make_shared(config_); transformer[TransformerFirstStage::kWTELayerName] = modules_[kPPFirstStageName]->mutable_module(TransformerFirstStage::kWTELayerName); - if (config_.attention_type == AttentionType::kStandard) { + if (config_.position_embedding_type == PositionEmbeddingType::kLearnedAbsolute) { transformer[TransformerFirstStage::kWPELayerName] = modules_[kPPFirstStageName]->mutable_module(TransformerFirstStage::kWPELayerName); } diff --git a/infini_train/src/nn/modules/transformer/utils.cc b/infini_train/src/nn/modules/transformer/utils.cc index 98505fd0..4ec11f2d 100644 --- a/infini_train/src/nn/modules/transformer/utils.cc +++ b/infini_train/src/nn/modules/transformer/utils.cc @@ -1,5 +1,9 @@ #include "infini_train/include/nn/modules/transformer/utils.h" +#include +#include +#include + #include "glog/logging.h" #include "infini_train/include/nn/functional.h" @@ -27,4 +31,38 @@ std::shared_ptr PrecomputeFreqsCis(int64_t dim, int64_t end, float theta return freqs_cis; } + +std::tuple, std::shared_ptr> +ApplyRotaryEmbedding(const std::shared_ptr &xq, const std::shared_ptr &xk, + const std::shared_ptr &freqs_cis) { + const auto &x_shape = xq->Dims(); // (B, T, H, D) + const int64_t T = x_shape[1]; + const int64_t D = x_shape[3]; + + std::vector target_shape = {1, T, 1, D / 2, 2}; + auto cos_sin = freqs_cis->View(target_shape); // -> (1, T, 1, D/2, 2) + + auto cos = cos_sin->Slice(-1, 0, 1, 1)->Squeeze(-1); // (1, T, 1, D/2) + auto sin = cos_sin->Slice(-1, 1, 2, 1)->Squeeze(-1); // (1, T, 1, D/2) + + auto slice_pair = [](const std::shared_ptr &x) { + auto even = x->Slice(-1, 0, x->Dims().back(), 2); + auto odd = x->Slice(-1, 1, x->Dims().back(), 2); + return std::make_pair(even, odd); + }; + + auto [q_even, q_odd] = slice_pair(xq); + auto q_rotated_left = q_even * cos - q_odd * sin; + auto q_rotated_right = q_even * sin + q_odd * cos; + auto q_rotated + = nn::function::Stack(std::vector>{q_rotated_left, q_rotated_right}, -1)->Flatten(-2); + + auto [k_even, k_odd] = slice_pair(xk); + auto k_rotated_left = k_even * cos - k_odd * sin; + auto k_rotated_right = k_even * sin + k_odd * cos; + auto k_rotated + = nn::function::Stack(std::vector>{k_rotated_left, k_rotated_right}, -1)->Flatten(-2); + + return {q_rotated, k_rotated}; +} } // namespace infini_train diff --git a/tests/transformer/test_transformer_architecture.cc b/tests/transformer/test_transformer_architecture.cc index 2148cb0c..d4a6efc2 100644 --- a/tests/transformer/test_transformer_architecture.cc +++ b/tests/transformer/test_transformer_architecture.cc @@ -101,7 +101,7 @@ TEST_P(TransformerModuleTest, StandardAttention) { config.n_embd = 64; config.n_head = 4; config.n_kv_head = 4; - config.attention_type = nn::AttentionType::kStandard; + config.position_embedding_type = nn::PositionEmbeddingType::kLearnedAbsolute; config.add_bias_linear = true; auto attn = std::make_shared(config); @@ -150,7 +150,7 @@ TEST_P(TransformerModuleTest, LLaMA3Model) { config.n_head = 4; config.n_kv_head = 2; config.n_embd = 64; - config.attention_type = nn::AttentionType::kRoPE; + config.position_embedding_type = nn::PositionEmbeddingType::kRoPE; config.activation_type = nn::MLPType::kSwiGLU; config.norm_type = nn::NormType::kRMSNorm; config.add_bias_linear = false; @@ -177,7 +177,7 @@ TEST_P(TransformerModuleTest, StateDict) { config.n_kv_head = 2; config.n_embd = 32; config.vocab_size = 1000; - config.attention_type = nn::AttentionType::kStandard; + config.position_embedding_type = nn::PositionEmbeddingType::kLearnedAbsolute; config.activation_type = nn::MLPType::kGELU; config.norm_type = nn::NormType::kLayerNorm; config.add_bias_linear = true; @@ -192,7 +192,6 @@ TEST_P(TransformerModuleTest, StateDict) { EXPECT_GE(state_dict.size(), params.size()); } - TEST_P(TransformerModuleTest, MoELayerTop1) { nn::TransformerConfig config; config.n_embd = 32; From ccdf04852c105d981754b21af99c2a2f6372171a Mon Sep 17 00:00:00 2001 From: bolunz Date: Thu, 23 Jul 2026 22:58:57 +0800 Subject: [PATCH 2/2] fix: create mask buffer only for GPT-2 case --- .../src/nn/modules/transformer/causal_self_attention.cc | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/infini_train/src/nn/modules/transformer/causal_self_attention.cc b/infini_train/src/nn/modules/transformer/causal_self_attention.cc index 7f20fad1..bc4c6bd4 100644 --- a/infini_train/src/nn/modules/transformer/causal_self_attention.cc +++ b/infini_train/src/nn/modules/transformer/causal_self_attention.cc @@ -44,9 +44,12 @@ CausalSelfAttention::CausalSelfAttention(const TransformerConfig &config) : Clon /*skip_bias_add=*/false, /*sequence_parallel=*/nn::parallel::global::GetSequenceParallelEnabled()); - // causal mask: (1, 1, block_size, block_size) - buffers_[kParamBiasName] = function::Tril(nn::function::Ones({config_.block_size, config_.block_size})) - ->View({1, 1, config_.block_size, config_.block_size}); + // FIXME(zbl): Decouple causal-mask ownership from position embedding. For now, only learned-absolute models use + // this precomputed buffer; RoPE callers provide a runtime-sized mask. + if (config_.position_embedding_type == PositionEmbeddingType::kLearnedAbsolute) { + buffers_[kParamBiasName] = function::Tril(nn::function::Ones({config_.block_size, config_.block_size})) + ->View({1, 1, config_.block_size, config_.block_size}); + } } void CausalSelfAttention::SetupAttention(const TransformerConfig &config) {