[WIP] feat: add Qwen3

CMakeLists.txt

@@ -218,6 +218,14 @@ add_executable(llama3
 )
 link_infini_train_exe(llama3)
 
+add_executable(qwen3
+  example/qwen3/main.cc
+  example/common/tiny_shakespeare_dataset.cc
+  example/common/utils.cc
+  example/qwen3/checkpoint_loader.cc
+  example/common/tokenizer.cc
+)
+link_infini_train_exe(qwen3)
 # Tools
 add_subdirectory(tools/infini_run)
 set_target_properties(infini_run PROPERTIES RUNTIME_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR})

example/qwen3/checkpoint_loader.h

@@ -0,0 +1,12 @@
+#pragma once
+
+#include <memory>
+#include <string>
+
+namespace infini_train::nn {
+class TransformerModel;
+} // namespace infini_train::nn
+
+namespace qwen3 {
+std::shared_ptr<infini_train::nn::TransformerModel> LoadFromLLMC(const std::string &filepath);
+} // namespace qwen3

example/qwen3/config.h

@@ -0,0 +1,29 @@
+#pragma once
+
+#include "infini_train/include/nn/modules/transformer/transformer_config.h"
+
+namespace nn = infini_train::nn;
+namespace qwen3 {
+inline nn::TransformerConfig Qwen3Config() {
+    return {.block_size = 40960,
+            .vocab_size = 151936,
+            .original_vocab_size = 151936,
+            .n_layer = 36,
+            .n_head = 32,
+            .n_kv_head = 8,
+            .n_embd = 4096,
+            .position_embedding_type = nn::PositionEmbeddingType::kRoPE,
+            .activation_type = nn::MLPType::kSwiGLU,
+            .norm_type = nn::NormType::kRMSNorm,
+            .add_bias_linear = false,
+            .add_bias_lm_head = false,
+            .tie_weights = false,
+            .ffn_expansion_ratio = 4.5f, // 4096*4.5*2/3 = 12288
+            .ffn_dim_multiplier = std::nullopt,
+            .multiple_of = 1,
+            .rope_theta = 1000000.0f,
+            .use_scaled_rope = false,
+            .rotary_interleaved = false,
+            .norm_eps = 1e-6f};
+}
+} // namespace qwen3

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

@@ -4,6 +4,7 @@
 #include <vector>
 
 #include "infini_train/include/nn/modules/module.h"
+#include "infini_train/include/nn/modules/normalization.h"
 #include "infini_train/include/nn/modules/transformer/transformer_config.h"
 
 namespace infini_train::nn {
@@ -14,6 +15,9 @@ class CausalSelfAttention : public infini_train::nn::CloneableModule<CausalSelfA
     static constexpr char kCAttnLayerName[] = "c_attn";
     static constexpr char kCProjLayerName[] = "c_proj";
 
+    static constexpr char kQNormLayerName[] = "q_norm";
+    static constexpr char kKNormLayerName[] = "k_norm";
+
     static constexpr char kParamBiasName[] = "bias";
 
     explicit CausalSelfAttention(const TransformerConfig &config);
@@ -31,6 +35,9 @@ class CausalSelfAttention : public infini_train::nn::CloneableModule<CausalSelfA
     int64_t n_rep_ = 0;
     int64_t head_dim_ = 0;
 
+    std::shared_ptr<infini_train::nn::RMSNorm> q_norm_;
+    std::shared_ptr<infini_train::nn::RMSNorm> k_norm_;
+
     // Setup method for different attention modes
     void SetupAttention(const TransformerConfig &config);
 

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

@@ -81,8 +81,9 @@ struct TransformerConfig {
     std::optional<MoEConfig> moe_config = std::nullopt;
 
     // RoPE config
-    float rope_theta = 500000.0f; // theta in RoPE
-    bool use_scaled_rope = false; // scaled RoPE
+    float rope_theta = 500000.0f;   // theta in RoPE
+    bool use_scaled_rope = false;   // scaled RoPE
+    bool rotary_interleaved = true; // Pair adjacent dimensions; false uses the Hugging Face half-split layout.
 
     // Normalization
     float norm_eps = 1e-5f; // epsilon in RMSNorm
@@ -92,6 +93,10 @@ struct TransformerConfig {
     bool flash = false;             // flash attention
     int64_t max_gen_batch_size = 4; // max batch size during inference
 
+    // Q-K Norm (Qwen3)
+    bool use_qk_norm = false;
+    float qk_norm_eps = 1e-6f;
+
     bool UseGQA() const;
     int GetChunkSize() const;
 };

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

@@ -13,5 +13,5 @@ std::shared_ptr<Tensor> PrecomputeFreqsCis(int64_t dim, int64_t end, float theta
 
 std::tuple<std::shared_ptr<Tensor>, std::shared_ptr<Tensor>>
 ApplyRotaryEmbedding(const std::shared_ptr<Tensor> &xq, const std::shared_ptr<Tensor> &xk,
-                     const std::shared_ptr<Tensor> &freqs_cis);
+                     const std::shared_ptr<Tensor> &freqs_cis, bool rotary_interleaved = true);
 } // namespace infini_train

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

@@ -23,6 +23,13 @@ namespace infini_train::nn {
 CausalSelfAttention::CausalSelfAttention(const TransformerConfig &config) : CloneableModule(kType), config_(config) {
     SetupAttention(config);
 
+    if (config_.use_qk_norm) {
+        q_norm_ = std::make_shared<nn::RMSNorm>(head_dim_, config_.qk_norm_eps);
+        k_norm_ = std::make_shared<nn::RMSNorm>(head_dim_, config_.qk_norm_eps);
+        modules_[kQNormLayerName] = q_norm_;
+        modules_[kKNormLayerName] = k_norm_;
+    }
+
     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>(
@@ -122,10 +129,17 @@ CausalSelfAttention::Forward(const std::vector<std::shared_ptr<infini_train::Ten
     auto k = qkv->Slice(2, q_size_local, q_size_local + kv_size_local)->View({B, T, KV_local, D});
     // v: (B, T, KV_local, D)
     auto v = qkv->Slice(2, q_size_local + kv_size_local, q_size_local + 2 * kv_size_local)->View({B, T, KV_local, D});
+    if (config_.use_qk_norm) {
+        auto q_shape = q->Dims();
+        q = (*q_norm_)({q->View({B * T * H_local, D})})[0]->View(q_shape);
+
+        auto k_shape = k->Dims();
+        k = (*k_norm_)({k->View({B * T * KV_local, D})})[0]->View(k_shape);
+    }
 
     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);
+        std::tie(q, k) = ApplyRotaryEmbedding(q, k, freqs_cis, config_.rotary_interleaved);
     }
 
     // TODO(zbl): use kv cache during inference

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

@@ -34,7 +34,7 @@ std::shared_ptr<Tensor> PrecomputeFreqsCis(int64_t dim, int64_t end, float theta
 
 std::tuple<std::shared_ptr<Tensor>, std::shared_ptr<Tensor>>
 ApplyRotaryEmbedding(const std::shared_ptr<Tensor> &xq, const std::shared_ptr<Tensor> &xk,
-                     const std::shared_ptr<Tensor> &freqs_cis) {
+                     const std::shared_ptr<Tensor> &freqs_cis, bool rotary_interleaved) {
     const auto &x_shape = xq->Dims(); // (B, T, H, D)
     const int64_t T = x_shape[1];
     const int64_t D = x_shape[3];
@@ -45,24 +45,25 @@ ApplyRotaryEmbedding(const std::shared_ptr<Tensor> &xq, const std::shared_ptr<Te
     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<Tensor> &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 slice_pair = [rotary_interleaved](const std::shared_ptr<Tensor> &x) {
+        const auto dim = x->Dims().back();
+        if (rotary_interleaved) {
+            return std::make_pair(x->Slice(-1, 0, dim, 2), x->Slice(-1, 1, dim, 2));
+        }
+        return std::make_pair(x->Slice(-1, 0, dim / 2), x->Slice(-1, dim / 2, dim));
     };
 
-    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<std::shared_ptr<Tensor>>{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<std::shared_ptr<Tensor>>{k_rotated_left, k_rotated_right}, -1)->Flatten(-2);
+    auto rotate = [&](const std::shared_ptr<Tensor> &x) {
+        auto [left, right] = slice_pair(x);
+        auto rotated_left = left * cos - right * sin;
+        auto rotated_right = left * sin + right * cos;
+        if (rotary_interleaved) {
+            return nn::function::Stack(std::vector<std::shared_ptr<Tensor>>{rotated_left, rotated_right}, -1)
+                ->Flatten(-2);
+        }
+        return nn::function::Concat(std::vector<std::shared_ptr<Tensor>>{rotated_left, rotated_right}, -1);
+    };
 
-    return {q_rotated, k_rotated};
+    return {rotate(xq), rotate(xk)};
 }
 } // namespace infini_train