feat: Update PR according to comments

Author: JYMiracle305

CMakeLists.txt

@@ -265,7 +265,6 @@ add_executable(llama3
   example/llama3/checkpoint_loader.cc
 )
 link_infini_train_exe(llama3)
-endif()
 
 add_executable(qwen3
   example/qwen3/main.cc
@@ -275,6 +274,8 @@ add_executable(qwen3
   example/common/tokenizer.cc
 )
 link_infini_train_exe(qwen3)
+endif()
+
 # Tools
 if(PROJECT_IS_TOP_LEVEL)
 add_subdirectory(tools/infini_run)

example/common/tokenizer.h

@@ -11,9 +11,9 @@ class Module;
 class Tokenizer {
 public:
     enum class Version : uint32_t {
-        kV1 = 1,
-        kV2 = 2,
-        kV3 = 3,
+        kV1 = 1, // GPT2
+        kV2 = 2, // LLaMA-3
+        kV3 = 3, // Qwen-3
     };
 
     Tokenizer(const std::string &filepath);

example/gpt2/config.h

@@ -22,7 +22,8 @@ inline nn::TransformerConfig GPT2Config() {
             .tie_weights = true,
             .ffn_expansion_ratio = 4.0f,
             .ffn_dim_multiplier = std::nullopt,
-            .multiple_of = 1};
+            .multiple_of = 1,
+            .rotary_interleaved = true};
 }
 
 inline void SanitizeGPT2Config(const nn::TransformerConfig &c) {

example/llama3/config.h

@@ -22,7 +22,8 @@ inline nn::TransformerConfig LLaMA3Config() {
             .tie_weights = false,
             .ffn_expansion_ratio = 4.0f,
             .ffn_dim_multiplier = 1.5f,
-            .multiple_of = 256};
+            .multiple_of = 256,
+            .rotary_interleaved = true};
 }
 
 inline void SanitizeLLaMA3Config(const nn::TransformerConfig &c) {

example/mixtral/config.h

@@ -28,6 +28,7 @@ inline nn::TransformerConfig TinyMixtralConfig() {
     config.norm_eps = 1e-5f;
     config.rope_theta = 1000000.0f;
     config.use_scaled_rope = false;
+    config.rotary_interleaved = true;
 
     nn::MoEConfig moe_config;
     moe_config.num_experts = 8;

example/qwen3/checkpoint_loader.cc

@@ -12,8 +12,6 @@
 
 #include "glog/logging.h"
 
-#include "example/common/utils.h"
-#include "example/qwen3/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/mlp.h"
@@ -22,6 +20,9 @@
 #include "infini_train/include/nn/parallel/tensor_parallel.h"
 #include "infini_train/include/tensor.h"
 
+#include "example/common/utils.h"
+#include "example/qwen3/config.h"
+
 using namespace infini_train;
 namespace nn = infini_train::nn;
 
@@ -107,8 +108,8 @@ std::shared_ptr<nn::TransformerModel> LoadFromLLMC(const std::string &filepath)
     qwen3_config.use_scaled_rope = static_cast<bool>(use_scaled_rope);
     qwen3_config.norm_eps = norm_eps;
     qwen3_config.max_gen_batch_size = max_gen_bs;
-    qwen3_config.use_qk_norm = true;
-    qwen3_config.qk_norm_eps = norm_eps;
+    qwen3_config.qk_layernorm = true;
+    qwen3::SanitizeQwen3Config(qwen3_config);
     auto qwen3 = std::make_shared<nn::TransformerModel>(qwen3_config);
 
     // ========== pp_size: num_stages; vpp_size: num_chunks_per_stage ==========
@@ -178,7 +179,7 @@ std::shared_ptr<nn::TransformerModel> LoadFromLLMC(const std::string &filepath)
 
     // RowParallel (proj)
     const int64_t in_pp = static_cast<int64_t>(n_embd) / tp_size;
-    // nn::MLP: c_fc/c_fc2(shard along row), c_proj(shard along col)
+    // nn::MLP: packed c_fc [gate | up] (shard each block along row), c_proj (shard along col)
     const int64_t fc_out = ffn_hidden;
     const int64_t fc_pp = fc_out / tp_size;
     const int64_t in_fc_pp = ffn_hidden / tp_size;
@@ -309,13 +310,13 @@ std::shared_ptr<nn::TransformerModel> LoadFromLLMC(const std::string &filepath)
         }
     }
 
-    // gate_proj is loaded into c_fc2 because MLP applies SiLU to its second projection.
+    // transformer.h.{i}.mlp.c_fc2.weight (gate) -> local packed c_fc rows [0 : fc_pp)
     local_layer_index = 0;
     for (int i = 0; i < static_cast<int>(n_layer); ++i) {
         if (owned_layers[i]) {
             auto &tensor = state_dict[std::format("{}.{}.{}.{}.{}.{}", nn::TransformerModel::kTransformerModelName,
                                                   nn::TransformerChunk::kHLayerName, std::to_string(local_layer_index),
-                                                  nn::TransformerLayer::kMlpLayerName, nn::MLP::kCFc2LayerName,
+                                                  nn::TransformerLayer::kMlpLayerName, nn::MLP::kCFcLayerName,
                                                   nn::parallel::ColumnParallelLinear::kParamWeightName)];
             ReadMatrixRowShardFloat(ifs, static_cast<float *>(tensor->DataPtr()),
                                     /*rows=*/fc_out, /*cols=*/n_embd,
@@ -327,16 +328,16 @@ std::shared_ptr<nn::TransformerModel> LoadFromLLMC(const std::string &filepath)
         }
     }
 
-    // up_proj is loaded into c_fc, producing SiLU(gate_proj(x)) * up_proj(x).
+    // transformer.h.{i}.mlp.c_fc.weight (up) -> local packed c_fc rows [fc_pp : 2*fc_pp)
     local_layer_index = 0;
     for (int i = 0; i < static_cast<int>(n_layer); ++i) {
         if (owned_layers[i]) {
             auto &tensor = state_dict[std::format("{}.{}.{}.{}.{}.{}", nn::TransformerModel::kTransformerModelName,
                                                   nn::TransformerChunk::kHLayerName, std::to_string(local_layer_index),
                                                   nn::TransformerLayer::kMlpLayerName, nn::MLP::kCFcLayerName,
                                                   nn::parallel::ColumnParallelLinear::kParamWeightName)];
-            ReadMatrixRowShardFloat(ifs, static_cast<float *>(tensor->DataPtr()),
-                                    /*rows=*/fc_out, /*cols=*/n_embd,
+            float *dst = static_cast<float *>(tensor->DataPtr()) + fc_pp * n_embd;
+            ReadMatrixRowShardFloat(ifs, dst, /*rows=*/fc_out, /*cols=*/n_embd,
                                     /*row_start=*/tp_rank * fc_pp, /*row_cnt=*/fc_pp);
             ++local_layer_index;
         } else {

example/qwen3/config.h

@@ -1,4 +1,7 @@
 #pragma once
+#include <cmath>
+
+#include "glog/logging.h"
 
 #include "infini_train/include/nn/modules/transformer/transformer_config.h"
 
@@ -25,7 +28,36 @@ inline nn::TransformerConfig Qwen3Config() {
             .use_scaled_rope = false,
             .rotary_interleaved = false,
             .norm_eps = 1e-6f,
-            .use_qk_norm = true,
-            .qk_norm_eps = 1e-6f};
+            .qk_layernorm = true};
+}
+
+inline void SanitizeQwen3Config(const nn::TransformerConfig &c) {
+    CHECK_GT(c.block_size, 0);
+    CHECK_GT(c.vocab_size, 0);
+    CHECK_GE(c.vocab_size, c.original_vocab_size);
+    CHECK_GT(c.n_layer, 0);
+    CHECK_GT(c.n_head, 0);
+    CHECK_GT(c.n_kv_head, 0);
+    CHECK_LE(c.n_kv_head, c.n_head);
+    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";
+    const auto head_dim = c.n_embd / c.n_head;
+    CHECK_GT(head_dim, 0) << "Qwen3 attention head dimension must be positive";
+    CHECK_EQ(head_dim % 2, 0) << "Qwen3 RoPE requires an even attention head dimension";
+    CHECK(c.position_embedding_type == nn::PositionEmbeddingType::kRoPE) << "Qwen3 requires RoPE position embedding";
+    CHECK(!c.use_scaled_rope) << "Qwen3 scaled RoPE is not implemented";
+    CHECK(!c.rotary_interleaved) << "Qwen3 requires half-split rotary embedding";
+    CHECK_GT(c.rope_theta, 0.0f);
+    CHECK(c.activation_type == nn::MLPType::kSwiGLU) << "Qwen3 requires SwiGLU activation";
+    CHECK(c.ffn_type == nn::FFNType::kDense) << "Qwen3-8B requires a dense FFN";
+    CHECK_GT(c.ffn_expansion_ratio, 0.0f);
+    CHECK(c.norm_type == nn::NormType::kRMSNorm) << "Qwen3 requires RMSNorm";
+    CHECK_GT(c.norm_eps, 0.0f);
+    CHECK(!c.add_bias_linear) << "Qwen3 has no bias in linear layers";
+    CHECK(!c.add_bias_lm_head) << "Qwen3 has no bias in lm_head";
+    CHECK(!c.tie_weights) << "Qwen3 does not tie embedding and lm_head weights";
+    CHECK_GT(c.multiple_of, 0);
+    CHECK(c.qk_layernorm) << "Qwen3 apply laynorm to the query and key embeddings";
 }
 } // namespace qwen3