Skip to content

Commit 937a71c

Browse files
Chamberlain0w0Chamberlain0w0
committed
refactor: merge 2 MHA paths, rename attention_type to position_embedding_type
1 parent dd18b35 commit 937a71c

8 files changed

Lines changed: 62 additions & 114 deletions

File tree

‎example/gpt2/config.h‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ inline nn::TransformerConfig GPT2Config() {
1414
.n_head = 12,
1515
.n_kv_head = 12,
1616
.n_embd = 768,
17-
.attention_type = nn::AttentionType::kStandard,
17+
.position_embedding_type = nn::PositionEmbeddingType::kLearnedAbsolute,
1818
.activation_type = nn::MLPType::kGELU,
1919
.norm_type = nn::NormType::kLayerNorm,
2020
.add_bias_linear = true,
@@ -34,7 +34,8 @@ inline void SanitizeGPT2Config(const nn::TransformerConfig &c) {
3434
CHECK_GT(c.n_embd, 0);
3535
CHECK_EQ(c.n_embd % c.n_head, 0) << "n_embd must be divisible by n_head";
3636
CHECK_EQ(c.n_kv_head, c.n_head) << "GPT-2 does not use GQA; n_kv_head must equal n_head";
37-
CHECK(c.attention_type == nn::AttentionType::kStandard) << "GPT-2 requires standard attention";
37+
CHECK(c.position_embedding_type == nn::PositionEmbeddingType::kLearnedAbsolute)
38+
<< "GPT-2 requires learned absolute position embedding";
3839
CHECK(c.activation_type == nn::MLPType::kGELU) << "GPT-2 requires GELU activation";
3940
CHECK(c.norm_type == nn::NormType::kLayerNorm) << "GPT-2 requires LayerNorm";
4041
}

‎example/llama3/config.h‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ inline nn::TransformerConfig LLaMA3Config() {
1414
.n_head = 32,
1515
.n_kv_head = 8,
1616
.n_embd = 2048,
17-
.attention_type = nn::AttentionType::kRoPE,
17+
.position_embedding_type = nn::PositionEmbeddingType::kRoPE,
1818
.activation_type = nn::MLPType::kSwiGLU,
1919
.norm_type = nn::NormType::kRMSNorm,
2020
.add_bias_linear = false,
@@ -36,7 +36,7 @@ inline void SanitizeLLaMA3Config(const nn::TransformerConfig &c) {
3636
CHECK_EQ(c.n_head % c.n_kv_head, 0) << "n_head must be divisible by n_kv_head for GQA";
3737
CHECK_GT(c.n_embd, 0);
3838
CHECK_EQ(c.n_embd % c.n_head, 0) << "n_embd must be divisible by n_head";
39-
CHECK(c.attention_type == nn::AttentionType::kRoPE) << "LLaMA-3 requires RoPE attention";
39+
CHECK(c.position_embedding_type == nn::PositionEmbeddingType::kRoPE) << "LLaMA-3 requires RoPE position embedding";
4040
CHECK(c.activation_type == nn::MLPType::kSwiGLU) << "LLaMA-3 requires SwiGLU activation";
4141
CHECK(c.norm_type == nn::NormType::kRMSNorm) << "LLaMA-3 requires RMSNorm";
4242
CHECK(!c.add_bias_linear) << "LLaMA-3 has no bias in linear layers";

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

Lines changed: 0 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -34,14 +34,6 @@ class CausalSelfAttention : public infini_train::nn::CloneableModule<CausalSelfA
3434
// Setup method for different attention modes
3535
void SetupAttention(const TransformerConfig &config);
3636

37-
// Standard attention forward (GPT2 style: no RoPE, no GQA)
38-
std::vector<std::shared_ptr<infini_train::Tensor>>
39-
ForwardStandard(const std::vector<std::shared_ptr<infini_train::Tensor>> &x);
40-
41-
// RoPE-aware attention forward (LLaMA3 style: with RoPE, optional GQA)
42-
std::vector<std::shared_ptr<infini_train::Tensor>>
43-
ForwardWithRoPE(const std::vector<std::shared_ptr<infini_train::Tensor>> &x);
44-
4537
// GQA helper method
4638
std::shared_ptr<infini_train::Tensor> RepeatKV(const std::shared_ptr<infini_train::Tensor> &x, int64_t n_rep);
4739
};

‎infini_train/include/nn/modules/transformer/transformer_config.h‎

Lines changed: 10 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -10,9 +10,13 @@ enum class ModelType {
1010
kLLaMA3, // LLaMA3
1111
};
1212

13-
enum class AttentionType {
14-
kStandard, // Standard attention
15-
kRoPE // Rotary Position Embedding
13+
enum class PositionEmbeddingType {
14+
kLearnedAbsolute, // Megatron: learned_absolute
15+
kRoPE, // Megatron: rope
16+
kYarn, // Megatron: yarn
17+
kMRoPE, // Megatron: mrope
18+
kRelative, // Megatron: relative
19+
kNone // Megatron: none
1620
};
1721

1822
enum class MLPType {
@@ -34,9 +38,9 @@ struct TransformerConfig {
3438
int64_t n_kv_head = 12; // Num of Key/Value heads (<= n_head, < n_head if using GQA)
3539
int64_t n_embd = 768; // Hidden size
3640

37-
AttentionType attention_type = AttentionType::kStandard; // Attention mechanism type
38-
MLPType activation_type = MLPType::kGELU; // MLP activation type
39-
NormType norm_type = NormType::kLayerNorm; // Normalization type
41+
PositionEmbeddingType position_embedding_type = PositionEmbeddingType::kLearnedAbsolute; // Position embedding type.
42+
MLPType activation_type = MLPType::kGELU; // MLP activation type
43+
NormType norm_type = NormType::kLayerNorm; // Normalization type
4044

4145
bool add_bias_linear = true; // Whether to add learnable bias to all Linear layers in the Transformer block,
4246
// including: attention QKV projection, attention output projection, MLP FC layers (and

‎infini_train/src/nn/modules/transformer/causal_self_attention.cc‎

Lines changed: 28 additions & 81 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
#include "infini_train/include/nn/modules/transformer/causal_self_attention.h"
22

33
#include <cmath>
4+
#include <limits>
45
#include <memory>
56
#include <tuple>
67
#include <vector>
@@ -43,12 +44,9 @@ CausalSelfAttention::CausalSelfAttention(const TransformerConfig &config) : Clon
4344
/*skip_bias_add=*/false,
4445
/*sequence_parallel=*/nn::parallel::global::GetSequenceParallelEnabled());
4546

46-
// For standard attention (GPT2 style), precompute causal mask
47-
if (config_.attention_type == AttentionType::kStandard) {
48-
// causal mask: (1, 1, block_size, block_size)
49-
buffers_[kParamBiasName] = function::Tril(nn::function::Ones({config_.block_size, config_.block_size}))
50-
->View({1, 1, config_.block_size, config_.block_size});
51-
}
47+
// causal mask: (1, 1, block_size, block_size)
48+
buffers_[kParamBiasName] = function::Tril(nn::function::Ones({config_.block_size, config_.block_size}))
49+
->View({1, 1, config_.block_size, config_.block_size});
5250
}
5351

5452
void CausalSelfAttention::SetupAttention(const TransformerConfig &config) {
@@ -77,88 +75,21 @@ void CausalSelfAttention::SetupAttention(const TransformerConfig &config) {
7775

7876
std::vector<std::shared_ptr<infini_train::Tensor>>
7977
CausalSelfAttention::Forward(const std::vector<std::shared_ptr<infini_train::Tensor>> &x) {
80-
if (config_.attention_type == AttentionType::kRoPE) {
81-
return ForwardWithRoPE(x);
82-
} else {
83-
return ForwardStandard(x);
84-
}
85-
}
86-
87-
std::vector<std::shared_ptr<infini_train::Tensor>>
88-
CausalSelfAttention::ForwardStandard(const std::vector<std::shared_ptr<infini_train::Tensor>> &x) {
89-
auto tp_world_size = parallel::global::GetTensorParallelSize();
90-
91-
const auto B = x[0]->Dims()[0]; // bs
92-
const auto C = x[0]->Dims()[2]; // n_embd
93-
const int64_t head_dim = n_embd_ / n_head_; // per-head dim (global)
94-
const int64_t local_C = n_embd_ / tp_world_size; // per-rank hidden
95-
96-
// (B, T, C) -> ColumnParallelLinear(C, 3*C) -> (B, T, 3 * local_C)
97-
// -> Split -> (3, B, T, local_C)
98-
auto qkv = (*modules_[kCAttnLayerName])(x)[0]->Split(local_C, 2);
99-
100-
// (B, T, local_C)
101-
auto q = qkv[0];
102-
auto k = qkv[1];
103-
auto v = qkv[2];
104-
105-
// NOTE(zbl): Acquire full T after AllGather is performed in ColumnParallelLinear
106-
const auto T = q->Dims()[1];
107-
108-
// View to multi-head: local_n_head * head_dim == local_C
109-
// (B, T, local_C) -> (B, T, h_l, Dh) -> (B, h_l, T, Dh)
110-
k = k->View({B, T, local_n_head_, head_dim})->Transpose(1, 2);
111-
q = q->View({B, T, local_n_head_, head_dim})->Transpose(1, 2);
112-
v = v->View({B, T, local_n_head_, head_dim})->Transpose(1, 2);
113-
114-
// (B, h_l, T, T)
115-
auto att = q->Matmul(k->Transpose(-2, -1)) * (1.0 / std::sqrt(head_dim));
116-
// (1, 1, T, T)
117-
auto mask = buffers_[kParamBiasName]->Slice({0, 0, 0, 0}, {1, 1, T, T}, {1, 1, 1, 1});
118-
// (1, 1, T, T) -> eq 0 -> (1, 1, T, T) -> masked_fill -> (B, h_l, T, T)
119-
att = att->MaskedFill(mask == 0, -std::numeric_limits<float>::infinity());
120-
// (B, h_l, T, T)
121-
att = nn::function::Softmax(att, -1);
122-
// (B, h_l, T, Dh)
123-
auto y = att->Matmul(v);
124-
// (B, h_l, T, Dh) -> (B, T, h_l, Dh) -> (B, T, local_C)
125-
y = y->Transpose(1, 2)->Contiguous()->View({B, T, local_C});
126-
127-
// Get full tensor
128-
// (B, T, local_C) -> RowParallelLinear(n_embd, n_embd) -> (B, T, C)
129-
y = (*modules_[kCProjLayerName])({y})[0];
130-
// (B, T, C) == (bs, seq_len, n_embd)
131-
return {y};
132-
}
133-
134-
std::shared_ptr<infini_train::Tensor> CausalSelfAttention::RepeatKV(const std::shared_ptr<infini_train::Tensor> &x,
135-
int64_t n_rep) {
136-
const auto &shape = x->Dims();
137-
const int64_t B = shape[0], T = shape[1], H = shape[2], D = shape[3];
138-
139-
if (n_rep == 1) {
140-
return x;
141-
}
142-
143-
return x->View({B, T, H, 1, D})->RepeatInterleave(n_rep, 3)->Contiguous()->View({B, T, H * n_rep, D});
144-
}
145-
146-
std::vector<std::shared_ptr<infini_train::Tensor>>
147-
CausalSelfAttention::ForwardWithRoPE(const std::vector<std::shared_ptr<infini_train::Tensor>> &x) {
14878
const auto B = x[0]->Dims()[0]; // bs
14979
const auto C = x[0]->Dims()[2]; // n_embd
15080

15181
const auto tp_size = nn::parallel::global::GetTensorParallelSize();
15282

15383
const auto C_local = C / tp_size;
154-
const auto H_local = n_head_ / tp_size;
84+
const auto H_local = local_n_head_;
15585
const auto KV_local = n_kv_head_ / tp_size;
15686
const auto D = head_dim_; // n_embd / n_head
15787

15888
const auto freqs_cis = x.size() > 1 ? x[1] : nullptr;
159-
const auto start_pos = x.size() > 2 ? x[2] : nullptr;
16089
const auto mask = x.size() > 3 ? x[3] : nullptr;
161-
CHECK(freqs_cis != nullptr) << "freqs_cis is null.";
90+
if (config_.position_embedding_type == PositionEmbeddingType::kRoPE) {
91+
CHECK(freqs_cis != nullptr) << "freqs_cis is null.";
92+
}
16293

16394
// (B, T, C) -> (B, T, (H + 2 * n_kv_head) * D)
16495
auto qkv = (*modules_[kCAttnLayerName])({x[0]})[0];
@@ -176,10 +107,10 @@ CausalSelfAttention::ForwardWithRoPE(const std::vector<std::shared_ptr<infini_tr
176107
// v: (B, T, KV_local, D)
177108
auto v = qkv->Slice(2, q_size_local + kv_size_local, q_size_local + 2 * kv_size_local)->View({B, T, KV_local, D});
178109

179-
// -> RoPE on q, k
180-
// q: (B, T, H_local, D)
181-
// k: (B, T, KV_local, D)
182-
std::tie(q, k) = ApplyRotaryEmbedding(q, k, freqs_cis);
110+
if (config_.position_embedding_type == PositionEmbeddingType::kRoPE) {
111+
// q: (B, T, H_local, D), k: (B, T, KV_local, D)
112+
std::tie(q, k) = ApplyRotaryEmbedding(q, k, freqs_cis);
113+
}
183114

184115
// TODO(zbl): use kv cache during inference
185116
// if (use_kv_) { ... }
@@ -207,6 +138,10 @@ CausalSelfAttention::ForwardWithRoPE(const std::vector<std::shared_ptr<infini_tr
207138
if (mask) {
208139
// mask: (1, 1, T, T)
209140
att = att->MaskedFill(mask, std::numeric_limits<float>::lowest());
141+
} else {
142+
// fallback causal mask: (1, 1, T, T)
143+
auto causal_mask = buffers_[kParamBiasName]->Slice({0, 0, 0, 0}, {1, 1, T, T}, {1, 1, 1, 1});
144+
att = att->MaskedFill(causal_mask == 0, -std::numeric_limits<float>::infinity());
210145
}
211146
// (B, H_local, T, T)
212147
att = nn::function::Softmax(att, -1);
@@ -221,4 +156,16 @@ CausalSelfAttention::ForwardWithRoPE(const std::vector<std::shared_ptr<infini_tr
221156
return {y};
222157
}
223158

159+
std::shared_ptr<infini_train::Tensor> CausalSelfAttention::RepeatKV(const std::shared_ptr<infini_train::Tensor> &x,
160+
int64_t n_rep) {
161+
const auto &shape = x->Dims();
162+
const int64_t B = shape[0], T = shape[1], H = shape[2], D = shape[3];
163+
164+
if (n_rep == 1) {
165+
return x;
166+
}
167+
168+
return x->View({B, T, H, 1, D})->RepeatInterleave(n_rep, 3)->Contiguous()->View({B, T, H * n_rep, D});
169+
}
170+
224171
} // namespace infini_train::nn

‎infini_train/src/nn/modules/transformer/mla_self_attention.cc‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -143,7 +143,7 @@ MLASelfAttention::Forward(const std::vector<std::shared_ptr<infini_train::Tensor
143143
const auto freqs_cis = x.size() > 1 ? x[1] : nullptr;
144144
// external_mask: (1, 1, T, T)
145145
const auto external_mask = x.size() > 3 ? x[3] : nullptr;
146-
if (config_.attention_type == AttentionType::kRoPE) {
146+
if (config_.position_embedding_type == PositionEmbeddingType::kRoPE) {
147147
CHECK(freqs_cis != nullptr) << "freqs_cis is null.";
148148
}
149149

@@ -227,7 +227,7 @@ MLASelfAttention::Forward(const std::vector<std::shared_ptr<infini_train::Tensor
227227
auto k_nope = kv->Slice(-1, 0, qk_nope_head_dim_);
228228
auto v = kv->Slice(-1, qk_nope_head_dim_, qk_nope_head_dim_ + v_head_dim_);
229229

230-
if (config_.attention_type == AttentionType::kRoPE) {
230+
if (config_.position_embedding_type == PositionEmbeddingType::kRoPE) {
231231
// q_pos_emb: (B, T, H_local, D_rope), k_pos_emb: (B, T, 1, D_rope)
232232
std::tie(q_pos_emb, k_pos_emb) = ApplyRotaryEmbedding(q_pos_emb, k_pos_emb, freqs_cis);
233233
}

‎infini_train/src/nn/modules/transformer/transformer.cc‎

Lines changed: 13 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -29,9 +29,11 @@ TransformerFirstStage::TransformerFirstStage(const TransformerConfig &config)
2929
modules_[kWTELayerName] = std::make_shared<parallel::VocabParallelEmbedding>(
3030
config_.vocab_size, config_.n_embd, parallel::global::GetSequenceParallelEnabled());
3131

32-
// RoPE-based models do not use absolute position embedding.
33-
if (config_.attention_type == AttentionType::kStandard) {
32+
// Only learned absolute position embedding uses a trainable WPE table.
33+
if (config_.position_embedding_type == PositionEmbeddingType::kLearnedAbsolute) {
3434
modules_[kWPELayerName] = std::make_shared<Embedding>(config_.block_size, config_.n_embd);
35+
} else if (config_.position_embedding_type != PositionEmbeddingType::kRoPE) {
36+
LOG(FATAL) << "Unsupported position embedding type";
3537
}
3638
}
3739

@@ -45,7 +47,7 @@ std::vector<std::shared_ptr<Tensor>> TransformerFirstStage::Forward(const std::v
4547
// (B, T) -> Embedding(V_local, C) -> (B, T, C)
4648
auto tok_emb = (*modules_[kWTELayerName])({x1});
4749

48-
// Add position embedding only for models that use absolute position encoding
50+
// Add position embedding only for models that use learned absolute position encoding.
4951
if (modules_.contains(kWPELayerName)) {
5052
// (T_local)
5153
// NOTE(zbl): Slice pos sequence when SP is enabled
@@ -66,7 +68,7 @@ std::vector<std::shared_ptr<Tensor>> TransformerFirstStage::Forward(const std::v
6668
// (B, T, C)
6769
return {tok_emb[0] + pos_emb[0]};
6870
} else {
69-
// For RoPE-based models (LLaMA3), no position embedding needed
71+
// For RoPE-based models (LLaMA3), no absolute position embedding is needed.
7072
// (B, T, C)
7173
return tok_emb;
7274
}
@@ -133,8 +135,8 @@ TransformerChunk::TransformerChunk(const TransformerConfig &config, int start_la
133135
std::vector<std::shared_ptr<Tensor>> TransformerChunk::Forward(const std::vector<std::shared_ptr<Tensor>> &x) {
134136
auto x1 = x[0];
135137

136-
// Check if we need to pass RoPE parameters (for LLaMA3 style models)
137-
if (config_.attention_type == AttentionType::kRoPE) {
138+
// Check if we need to pass RoPE parameters (for LLaMA3 style models).
139+
if (config_.position_embedding_type == PositionEmbeddingType::kRoPE) {
138140
// For RoPE models, we need to prepare freqs_cis and potentially other parameters
139141
const auto device = x1->GetDevice();
140142

@@ -163,9 +165,11 @@ std::vector<std::shared_ptr<Tensor>> TransformerChunk::Forward(const std::vector
163165
for (auto &h : *std::dynamic_pointer_cast<nn::ModuleList>(modules_[kHLayerName])) {
164166
x1 = (*h)({x1, freqs_view, start_pos_ptr, mask})[0];
165167
}
166-
} else {
167-
// Standard attention (GPT2 style)
168+
} else if (config_.position_embedding_type == PositionEmbeddingType::kLearnedAbsolute) {
169+
// Learned absolute position embedding models (GPT-2 style).
168170
for (auto &h : *std::dynamic_pointer_cast<nn::ModuleList>(modules_[kHLayerName])) { x1 = (*h)({x1})[0]; }
171+
} else {
172+
LOG(FATAL) << "Unsupported position embedding type";
169173
}
170174

171175
return {x1};
@@ -219,7 +223,7 @@ TransformerModel::TransformerModel(const TransformerConfig config)
219223
modules_[kPPFirstStageName] = std::make_shared<TransformerFirstStage>(config_);
220224
transformer[TransformerFirstStage::kWTELayerName]
221225
= modules_[kPPFirstStageName]->mutable_module(TransformerFirstStage::kWTELayerName);
222-
if (config_.attention_type == AttentionType::kStandard) {
226+
if (config_.position_embedding_type == PositionEmbeddingType::kLearnedAbsolute) {
223227
transformer[TransformerFirstStage::kWPELayerName]
224228
= modules_[kPPFirstStageName]->mutable_module(TransformerFirstStage::kWPELayerName);
225229
}

‎tests/transformer/test_transformer_architecture.cc‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -102,7 +102,7 @@ TEST_P(TransformerModuleTest, StandardAttention) {
102102
config.n_embd = 64;
103103
config.n_head = 4;
104104
config.n_kv_head = 4;
105-
config.attention_type = nn::AttentionType::kStandard;
105+
config.position_embedding_type = nn::PositionEmbeddingType::kLearnedAbsolute;
106106
config.add_bias_linear = true;
107107

108108
auto attn = std::make_shared<nn::CausalSelfAttention>(config);
@@ -120,7 +120,7 @@ TEST_P(TransformerModuleTest, MLAAttention) {
120120
config.n_embd = 64;
121121
config.n_head = 4;
122122
config.block_size = 16;
123-
config.attention_type = nn::AttentionType::kStandard;
123+
config.position_embedding_type = nn::PositionEmbeddingType::kLearnedAbsolute;
124124
config.add_bias_linear = true;
125125
config.multi_latent_attention = true;
126126
config.q_lora_rank = 32;
@@ -198,7 +198,7 @@ TEST_P(TransformerModuleTest, LLaMA3Model) {
198198
config.n_head = 4;
199199
config.n_kv_head = 2;
200200
config.n_embd = 64;
201-
config.attention_type = nn::AttentionType::kRoPE;
201+
config.position_embedding_type = nn::PositionEmbeddingType::kRoPE;
202202
config.activation_type = nn::MLPType::kSwiGLU;
203203
config.norm_type = nn::NormType::kRMSNorm;
204204
config.add_bias_linear = false;
@@ -225,7 +225,7 @@ TEST_P(TransformerModuleTest, StateDict) {
225225
config.n_kv_head = 2;
226226
config.n_embd = 32;
227227
config.vocab_size = 1000;
228-
config.attention_type = nn::AttentionType::kStandard;
228+
config.position_embedding_type = nn::PositionEmbeddingType::kLearnedAbsolute;
229229
config.activation_type = nn::MLPType::kGELU;
230230
config.norm_type = nn::NormType::kLayerNorm;
231231
config.add_bias_linear = true;

0 commit comments

Comments
 (0)