feat: add Qwen3

CMakeLists.txt

@@ -267,6 +267,14 @@ add_executable(llama3
 link_infini_train_exe(llama3)
 endif()
 
+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
 if(PROJECT_IS_TOP_LEVEL)
 add_subdirectory(tools/infini_run)

README.md

@@ -58,7 +58,7 @@ Build Options:
 | ------------------------- | ------------------------------- | ---------------------------------------------------- | -------------- |
 | Model Support             | GPT-2                           | Decoder-only Transformer language model              | ✔ Supported    |
 |                           | LLaMA 3                         | Modern LLaMA-family Transformer architecture         | ✔ Supported    |
-|                           | Qwen3-8B                        | Qwen3 8B language model                              | 🗓 Planned     |
+|                           | Qwen3-8B                        | Qwen3 8B language model with QK norm and GQA          | ✔ Supported    |
 |                           | DeepSeek-V3                     | Large-scale MoE-based language model                 | 🗓 Planned     |
 | Precision                 | Multiple Data Type              | FP32, BF16                                           | ✔ Supported    |
 |                           | Mixed Precision                 | Autocast-based BF16 compute with FP32 accumulation   | ✔ Supported    |
@@ -168,10 +168,25 @@ The generated files can be passed directly to the corresponding executables:
   --num_iteration 10
 ```
 
+##### Qwen3 8B
+
+```bash
+./build/qwen3 \
+  --device cuda \
+  --input_bin data/qwen3/tiny_shakespeare_train.bin \
+  --llmc_filepath data/qwen3/qwen3-8b-fp32.llmc \
+  --num_iteration 10
+```
+
+Qwen3-8B uses approximately 31 GB for FP32 model weights. Full-parameter training
+therefore requires tensor or pipeline parallelism (for example,
+`--tensor_parallel 8`) rather than a single 80 GB device when optimizer states
+are allocated. The input tokens and LLMC checkpoint must use the Qwen3 tokenizer.
+
 ### Launch Modes
 
-GPT-2 and LLaMA training support both thread-based and process-based launches.
-The examples below use LLaMA, but the same launch modes also apply to GPT-2.
+GPT-2, LLaMA, and Qwen3 training support both thread-based and process-based launches.
+The examples below use LLaMA, but the same launch modes also apply to the other models.
 
 #### Direct Launch
 

example/common/tokenizer.cc

@@ -18,6 +18,7 @@ namespace infini_train {
 
 constexpr uint32_t kGpt2Eot = 50256;
 constexpr uint32_t kLLaMA3Eot = 128001;
+constexpr uint32_t kQwen3Eot = 151645;
 constexpr uint64_t kRandomU32Multiplier = 0x2545F4914F6CDD1Dull;
 constexpr float kF32Divisor = 16777216.0f; // 2^24
 constexpr uint64_t kRngState = 1337;
@@ -27,13 +28,15 @@ using Version = Tokenizer::Version;
 const std::unordered_map<uint32_t, uint32_t> kEotMap = {
     {20240328, kGpt2Eot},   // GPT-2
     {20240801, kLLaMA3Eot}, // LLaMA-3
+    {20240916, kQwen3Eot},  // Qwen3
 };
 
 const std::unordered_map<uint32_t, std::vector<uint32_t>> kPromptMap = {
     // e.g. "The meaning of life is"
     // ref: https://tiktokenizer.vercel.app/
     {20240328, std::vector<uint32_t>{464, 3616, 286, 1204, 318}}, // GPT-2
     {20240801, std::vector<uint32_t>{791, 7438, 315, 2324, 374}}, // LLaMA-3
+    {20240916, std::vector<uint32_t>{785, 7290, 315, 2272, 374}}, // Qwen3
 };
 
 unsigned int RandomU32(uint64_t &state) {
@@ -78,7 +81,7 @@ Tokenizer::Tokenizer(const std::string &filepath) {
     Version version = static_cast<Version>(version_num);
     if (version == Version::kV1) {
         eot_token_ = kEotMap.at(magic_number_);
-    } else if (version == Version::kV2) {
+    } else if (version == Version::kV2 || version == Version::kV3) {
         const uint32_t eot_token_2 = BytesToType<uint32_t>(header, 12);
         eot_token_ = eot_token_2;
     } else {
@@ -88,8 +91,16 @@ Tokenizer::Tokenizer(const std::string &filepath) {
 
     token_table_.resize(vocab_size_);
     for (uint32_t i = 0; i < vocab_size_; ++i) {
-        uint8_t length;
-        ifs.read(reinterpret_cast<char *>(&length), sizeof(length));
+        size_t length = 0;
+        if (version == Version::kV3) {
+            uint8_t length_bytes[2];
+            ifs.read(reinterpret_cast<char *>(length_bytes), sizeof(length_bytes));
+            length = static_cast<size_t>(length_bytes[0]) | (static_cast<size_t>(length_bytes[1]) << 8U);
+        } else {
+            uint8_t length_v2 = 0;
+            ifs.read(reinterpret_cast<char *>(&length_v2), sizeof(length_v2));
+            length = length_v2;
+        }
 
         std::vector<char> buffer(length);
         ifs.read(buffer.data(), length);

example/common/tokenizer.h

@@ -13,6 +13,7 @@ class Tokenizer {
     enum class Version : uint32_t {
         kV1 = 1,
         kV2 = 2,
+        kV3 = 3,
     };
 
     Tokenizer(const std::string &filepath);