diff --git a/CMakeLists.txt b/CMakeLists.txt index 6bd8069d..b9afa1ef 100644 --- a/CMakeLists.txt +++ b/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}) diff --git a/example/qwen3/checkpoint_loader.cc b/example/qwen3/checkpoint_loader.cc new file mode 100644 index 00000000..0f064454 --- /dev/null +++ b/example/qwen3/checkpoint_loader.cc @@ -0,0 +1,388 @@ +#include "example/qwen3/checkpoint_loader.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#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" +#include "infini_train/include/nn/modules/transformer/transformer.h" +#include "infini_train/include/nn/parallel/global.h" +#include "infini_train/include/nn/parallel/tensor_parallel.h" +#include "infini_train/include/tensor.h" + +using namespace infini_train; +namespace nn = infini_train::nn; + +namespace { +constexpr int32_t kQwen3Magic = 20240804; +constexpr int32_t kQwen3FP32Version = 4; +constexpr size_t kQwen3HeaderBytes = 256 * sizeof(int32_t); +} // namespace + +namespace qwen3 { + +std::shared_ptr LoadFromLLMC(const std::string &filepath) { + if (!std::filesystem::exists(filepath)) { + LOG(FATAL) << "File not found: " << filepath; + } + + std::error_code file_size_error; + const auto actual_file_size = std::filesystem::file_size(filepath, file_size_error); + CHECK(!file_size_error) << "Failed to get file size for " << filepath << ": " << file_size_error.message(); + CHECK_GE(actual_file_size, kQwen3HeaderBytes) << "Qwen3 LLMC file is shorter than its header: " << filepath; + + std::ifstream ifs(filepath, std::ios::binary); + CHECK(ifs.is_open()) << "Failed to open Qwen3 LLMC file: " << filepath; + const auto header = ReadSeveralBytesFromIfstream(kQwen3HeaderBytes, &ifs); + CHECK(ifs.good()) << "Failed to read Qwen3 LLMC header: " << filepath; + + const auto magic = BytesToType(header, 0); + CHECK_EQ(magic, kQwen3Magic); + const auto version = BytesToType(header, 4); + CHECK_EQ(version, kQwen3FP32Version); + + const auto block_size = BytesToType(header, 8); + const auto vocab_size = BytesToType(header, 12); + const auto n_layer = BytesToType(header, 16); + const auto n_head = BytesToType(header, 20); + const auto n_kv_head = BytesToType(header, 24); + const auto n_embd = BytesToType(header, 28); + const auto intermediate_size = BytesToType(header, 32); + const auto multiple_of = BytesToType(header, 36); + const auto norm_eps = BytesToType(header, 40); + const auto rope_theta = BytesToType(header, 44); + const auto use_scaled_rope = BytesToType(header, 48); + const auto max_gen_bs = BytesToType(header, 52); + const auto version_major = BytesToType(header, 56); + const auto version_minor = BytesToType(header, 60); + + CHECK_GT(n_layer, 0); + CHECK_GT(n_head, 0); + CHECK_GT(n_kv_head, 0); + CHECK_GT(n_embd, 0); + CHECK_GT(intermediate_size, 0); + CHECK_EQ(n_embd % n_head, 0) << "n_embd must be divisible by n_head."; + CHECK_EQ(n_head % n_kv_head, 0) << "n_head must be divisible by n_kv_head."; + + const uintmax_t header_head_dim = n_embd / n_head; + const uintmax_t embedding_elements = static_cast(vocab_size) * n_embd; + const uintmax_t norm_elements = static_cast(n_layer) * n_embd; + const uintmax_t qk_norm_elements = 2 * static_cast(n_layer) * header_head_dim; + const uintmax_t qkv_elements + = static_cast(n_layer) * (n_embd + 2 * static_cast(n_kv_head) * header_head_dim) * n_embd; + const uintmax_t attention_output_elements = static_cast(n_layer) * n_embd * n_embd; + const uintmax_t mlp_elements = 3 * static_cast(n_layer) * intermediate_size * n_embd; + const uintmax_t final_norm_elements = n_embd; + const uintmax_t expected_file_size + = kQwen3HeaderBytes + + sizeof(float) + * (2 * embedding_elements + 2 * norm_elements + qk_norm_elements + qkv_elements + + attention_output_elements + mlp_elements + final_norm_elements); + CHECK_EQ(actual_file_size, expected_file_size) << "Qwen3 LLMC size mismatch for " << filepath << ": expected " + << expected_file_size << " bytes, got " << actual_file_size; + + nn::TransformerConfig qwen3_config = qwen3::Qwen3Config(); + qwen3_config.block_size = block_size; + qwen3_config.vocab_size = vocab_size; + qwen3_config.n_layer = n_layer; + qwen3_config.n_head = n_head; + qwen3_config.n_kv_head = n_kv_head; + qwen3_config.n_embd = n_embd; + qwen3_config.ffn_expansion_ratio + = 3.0f * static_cast(intermediate_size) / (2.0f * static_cast(n_embd)); + qwen3_config.multiple_of = multiple_of; + qwen3_config.rope_theta = rope_theta; + qwen3_config.use_scaled_rope = static_cast(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; + auto qwen3 = std::make_shared(qwen3_config); + + // ========== pp_size: num_stages; vpp_size: num_chunks_per_stage ========== + int pp_size = nn::parallel::global::GetPipelineParallelSize(); + int vpp_size = nn::parallel::global::GetVirtualPipelineParallelSize(); + auto pp_rank = nn::parallel::pp_rank; + auto [is_first_stage, is_last_stage, layer_ranges_per_chunk] + = nn::parallel::PipelineParallel::GetStageInfo(n_layer, pp_size, pp_rank, vpp_size); + // ========== layer to chunk ========== + std::vector owned_layers(n_layer, false); + for (const auto &[start, end] : layer_ranges_per_chunk) { + for (int i = start; i < end; ++i) { owned_layers[i] = true; } + } + + const int tp_size = nn::parallel::global::GetTensorParallelSize(); + const int tp_rank = nn::parallel::tp_rank; + + CHECK_EQ(n_embd % tp_size, 0) << "n_embd must be divisible by TP world size."; + CHECK_EQ(n_head % tp_size, 0) << "n_head must be divisible by TP world size."; + CHECK_EQ(n_kv_head % tp_size, 0) << "n_kv_head must be divisible by TP world size."; + CHECK_EQ(vocab_size % tp_size, 0) << "vocab_size must be divisible by TP world size."; + CHECK_EQ(intermediate_size % tp_size, 0) << "intermediate_size must be divisible by TP world size."; + + if (tp_rank == 0) { + LOG(INFO) << "Model Config:"; + LOG(INFO) << " block_size = " << block_size; + LOG(INFO) << " vocab_size = " << vocab_size; + LOG(INFO) << " n_layer = " << n_layer; + LOG(INFO) << " n_head = " << n_head; + LOG(INFO) << " n_kv_head = " << n_kv_head; + LOG(INFO) << " n_embd = " << n_embd; + LOG(INFO) << " intermediate_size = " << intermediate_size; + LOG(INFO) << " multiple_of = " << multiple_of; + LOG(INFO) << " norm_eps = " << norm_eps; + LOG(INFO) << " rope_theta = " << rope_theta; + LOG(INFO) << " use_scaled_rope = " << use_scaled_rope; + LOG(INFO) << " max_gen_bs = " << max_gen_bs; + LOG(INFO) << " version_major = " << version_major; + LOG(INFO) << " version_minor = " << version_minor; + + LOG(INFO) << "Pipeline Parallel Chunks:"; + for (size_t i = 0; i < layer_ranges_per_chunk.size(); ++i) { + LOG(INFO) << " Chunk " << i << ": layers " << layer_ranges_per_chunk[i].first << " to " + << layer_ranges_per_chunk[i].second; + } + } + + const int64_t head_dim = static_cast(n_embd) / static_cast(n_head); + + const int64_t ffn_hidden = static_cast(intermediate_size); + + // ===== Per-rank sizes / offsets ===== + // vocab parallel + const int64_t vpp = static_cast(vocab_size) / tp_size; + const int64_t v_start = static_cast(tp_rank) * vpp; + + // attention Q/K/V packed as rows: [Q | K | V] + const int64_t q_out_rows = static_cast(n_embd); + const int64_t kv_out_rows = static_cast(n_kv_head) * head_dim; // for K or V (each) + const int64_t attn_rows_all = q_out_rows + 2 * kv_out_rows; + const int64_t attn_cols = static_cast(n_embd); + + // local Q/K/V rows per tp_rank + const int64_t q_local_rows = static_cast(n_embd) / tp_size; // = (n_head/world)*head_dim + const int64_t kv_head_local = static_cast(n_kv_head) / tp_size; + const int64_t kv_local_rows = kv_head_local * head_dim; // for K or V (each) + + // RowParallel (proj) + const int64_t in_pp = static_cast(n_embd) / tp_size; + // nn::MLP: c_fc/c_fc2(shard 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; + + auto state_dict = qwen3->StateDict(); + + // ========== Read Sharded Params ========== + // transformer.wte.weight : (vocab_size, n_embd) -> local tp_rank: rows of [v_start : v_start+vpp) + if (is_first_stage) { + auto &wte = state_dict[std::format("{}.{}.{}", nn::TransformerModel::kTransformerModelName, + nn::TransformerFirstStage::kWTELayerName, + nn::parallel::VocabParallelEmbedding::kParamWeightName)]; + ReadMatrixRowShardFloat(ifs, static_cast(wte->DataPtr()), + /*rows=*/vocab_size, /*cols=*/n_embd, + /*row_start=*/v_start, /*row_cnt=*/vpp); + } else { + size_t wte_bytes = static_cast(vocab_size) * n_embd * sizeof(float); + ifs.seekg(wte_bytes, std::ios::cur); + } + + // transformer.h.{i}.ln_1.weight : Full version nn::RMSNorm + int local_layer_index = 0; + for (int i = 0; i < static_cast(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::kLn1LayerName, nn::RMSNorm::kParamWeightName)]; + ReadVectorAllFloat(ifs, static_cast(tensor->DataPtr()), n_embd); + ++local_layer_index; + } else { + size_t ln_1_bytes = n_embd * sizeof(float); + ifs.seekg(ln_1_bytes, std::ios::cur); + } + } + + local_layer_index = 0; + for (int i = 0; i < static_cast(n_layer); ++i) { + if (owned_layers[i]) { + auto &q_norm_tensor = state_dict[std::format( + "{}.{}.{}.{}.{}.{}", nn::TransformerModel::kTransformerModelName, nn::TransformerChunk::kHLayerName, + std::to_string(local_layer_index), nn::TransformerLayer::kAttnLayerName, + nn::CausalSelfAttention::kQNormLayerName, nn::RMSNorm::kParamWeightName)]; + ReadVectorAllFloat(ifs, static_cast(q_norm_tensor->DataPtr()), head_dim); + + auto &k_norm_tensor = state_dict[std::format( + "{}.{}.{}.{}.{}.{}", nn::TransformerModel::kTransformerModelName, nn::TransformerChunk::kHLayerName, + std::to_string(local_layer_index), nn::TransformerLayer::kAttnLayerName, + nn::CausalSelfAttention::kKNormLayerName, nn::RMSNorm::kParamWeightName)]; + ReadVectorAllFloat(ifs, static_cast(k_norm_tensor->DataPtr()), head_dim); + ++local_layer_index; + } else { + size_t qk_norm_bytes = 2 * head_dim * sizeof(float); + ifs.seekg(qk_norm_bytes, std::ios::cur); + } + } + + // transformer.h.{i}.attn.c_attn.weight : ColumnParallelLinear, but actually applies on "rows" + // W-qkv should be [Q(=n_embd) | K(=n_kv_head*head_dim) | V(=n_kv_head*head_dim)] x n_embd + local_layer_index = 0; + for (int i = 0; i < static_cast(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::kAttnLayerName, + nn::CausalSelfAttention::kCAttnLayerName, nn::parallel::ColumnParallelLinear::kParamWeightName)]; + + float *dst = static_cast(tensor->DataPtr()); + const std::streampos base_pos = ifs.tellg(); + + // Q block -> [0 : q_local_rows) + ifs.seekg(base_pos); + ReadMatrixRowShardFloat(ifs, + /*dst=*/dst + (0 * attn_cols), + /*rows=*/attn_rows_all, /*cols=*/attn_cols, + /*row_start=*/tp_rank * q_local_rows, /*row_cnt=*/q_local_rows); + + // K block -> [q_local_rows : q_local_rows + kv_local_rows) + ifs.seekg(base_pos); + ReadMatrixRowShardFloat(ifs, + /*dst=*/dst + (q_local_rows * attn_cols), + /*rows=*/attn_rows_all, /*cols=*/attn_cols, + /*row_start=*/q_out_rows + tp_rank * kv_local_rows, /*row_cnt=*/kv_local_rows); + + // V block -> [q_local_rows + kv_local_rows : q_local_rows + 2*kv_local_rows) + ifs.seekg(base_pos); + ReadMatrixRowShardFloat(ifs, + /*dst=*/dst + ((q_local_rows + kv_local_rows) * attn_cols), + /*rows=*/attn_rows_all, /*cols=*/attn_cols, + /*row_start=*/q_out_rows + kv_out_rows + tp_rank * kv_local_rows, + /*row_cnt=*/kv_local_rows); + ++local_layer_index; + } else { + size_t qkv_bytes = static_cast(attn_rows_all) * attn_cols * sizeof(float); + ifs.seekg(qkv_bytes, std::ios::cur); + } + } + + // transformer.h.{i}.attn.c_proj.weight : RowParallelLinear, but actually applies on "columns" + local_layer_index = 0; + for (int i = 0; i < static_cast(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::kAttnLayerName, + nn::CausalSelfAttention::kCProjLayerName, nn::parallel::RowParallelLinear::kParamWeightName)]; + ReadMatrixColShardFloat(ifs, static_cast(tensor->DataPtr()), + /*rows=*/n_embd, /*cols=*/n_embd, + /*col_start=*/tp_rank * in_pp, /*col_cnt=*/in_pp); + ++local_layer_index; + } else { + size_t c_proj_bytes = static_cast(n_embd) * n_embd * sizeof(float); + ifs.seekg(c_proj_bytes, std::ios::cur); + } + } + + // transformer.h.{i}.ln_2.weight : Full version RMSNorm + local_layer_index = 0; + for (int i = 0; i < static_cast(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::kLn2LayerName, nn::RMSNorm::kParamWeightName)]; + ReadVectorAllFloat(ifs, static_cast(tensor->DataPtr()), n_embd); + ++local_layer_index; + } else { + size_t ln_2_bytes = static_cast(n_embd) * sizeof(float); + ifs.seekg(ln_2_bytes, std::ios::cur); + } + } + + // gate_proj is loaded into c_fc2 because MLP applies SiLU to its second projection. + local_layer_index = 0; + for (int i = 0; i < static_cast(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::parallel::ColumnParallelLinear::kParamWeightName)]; + ReadMatrixRowShardFloat(ifs, static_cast(tensor->DataPtr()), + /*rows=*/fc_out, /*cols=*/n_embd, + /*row_start=*/tp_rank * fc_pp, /*row_cnt=*/fc_pp); + ++local_layer_index; + } else { + size_t fc_bytes = static_cast(ffn_hidden) * n_embd * sizeof(float); + ifs.seekg(fc_bytes, std::ios::cur); + } + } + + // up_proj is loaded into c_fc, producing SiLU(gate_proj(x)) * up_proj(x). + local_layer_index = 0; + for (int i = 0; i < static_cast(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(tensor->DataPtr()), + /*rows=*/fc_out, /*cols=*/n_embd, + /*row_start=*/tp_rank * fc_pp, /*row_cnt=*/fc_pp); + ++local_layer_index; + } else { + size_t fc2_bytes = static_cast(ffn_hidden) * n_embd * sizeof(float); + ifs.seekg(fc2_bytes, std::ios::cur); + } + } + + // transformer.h.{i}.mlp.c_proj.weight : RowParallelLinear, but actually applies on "columns" + local_layer_index = 0; + for (int i = 0; i < static_cast(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::kCProjLayerName, + nn::parallel::RowParallelLinear::kParamWeightName)]; + ReadMatrixColShardFloat(ifs, static_cast(tensor->DataPtr()), + /*rows=*/n_embd, /*cols=*/fc_out, + /*col_start=*/tp_rank * in_fc_pp, /*col_cnt=*/in_fc_pp); + ++local_layer_index; + } else { + size_t c_proj_bytes = static_cast(n_embd) * ffn_hidden * sizeof(float); + ifs.seekg(c_proj_bytes, std::ios::cur); + } + } + + // transformer.ln_f.weight : Full version nn::RMSNorm + // lm_head.weight : (vocab_size, n_embd) -> ColumnParallelLinear, but actually applies on "rows" + { + if (is_last_stage) { + auto &ln_f + = state_dict[std::format("{}.{}.{}", nn::TransformerModel::kTransformerModelName, + nn::TransformerLastStage::kLnFLayerName, nn::RMSNorm::kParamWeightName)]; + auto &lm_head = state_dict[std::format("{}.{}", nn::TransformerLastStage::kLMHeadLayerName, + nn::parallel::ColumnParallelLinear::kParamWeightName)]; + ReadVectorAllFloat(ifs, static_cast(ln_f->DataPtr()), n_embd); + ReadMatrixRowShardFloat(ifs, static_cast(lm_head->DataPtr()), + /*rows=*/vocab_size, /*cols=*/n_embd, + /*row_start=*/v_start, /*row_cnt=*/vpp); + } else { + size_t ln_f_bytes = static_cast(n_embd) * sizeof(float); + size_t lm_head_bytes = static_cast(vocab_size) * n_embd * sizeof(float); + ifs.seekg(ln_f_bytes + lm_head_bytes, std::ios::cur); + } + } + + return qwen3; +} +} // namespace qwen3 diff --git a/example/qwen3/checkpoint_loader.h b/example/qwen3/checkpoint_loader.h new file mode 100644 index 00000000..5f4880f6 --- /dev/null +++ b/example/qwen3/checkpoint_loader.h @@ -0,0 +1,12 @@ +#pragma once + +#include +#include + +namespace infini_train::nn { +class TransformerModel; +} // namespace infini_train::nn + +namespace qwen3 { +std::shared_ptr LoadFromLLMC(const std::string &filepath); +} // namespace qwen3 diff --git a/example/qwen3/config.h b/example/qwen3/config.h new file mode 100644 index 00000000..868df53f --- /dev/null +++ b/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 diff --git a/example/qwen3/main.cc b/example/qwen3/main.cc new file mode 100644 index 00000000..9b44135d --- /dev/null +++ b/example/qwen3/main.cc @@ -0,0 +1,568 @@ +#include +#include +#include +#include +#include +#include + +#include "gflags/gflags.h" +#include "glog/logging.h" + +#include "infini_train/include/autocast.h" +#include "infini_train/include/checkpoint/checkpoint.h" +#include "infini_train/include/checkpoint/checkpoint_manager.h" +#include "infini_train/include/core/runtime/device_guard.h" +#include "infini_train/include/dataloader.h" +#include "infini_train/include/device.h" +#include "infini_train/include/lr_scheduler.h" +#include "infini_train/include/nn/lora/lora_utils.h" +#include "infini_train/include/nn/modules/loss.h" +#include "infini_train/include/nn/modules/module.h" +#include "infini_train/include/nn/modules/transformer/transformer.h" +#include "infini_train/include/nn/parallel/ddp/distributed_data_parallel.h" +#include "infini_train/include/nn/parallel/ddp/distributed_optimizer.h" +#include "infini_train/include/nn/parallel/global.h" +#include "infini_train/include/nn/parallel/parallel_functional.h" +#include "infini_train/include/nn/parallel/pp/pipeline_parallel.h" +#include "infini_train/include/nn/parallel/process_group.h" +#include "infini_train/include/nn/parallel/rank.h" +#include "infini_train/include/nn/parallel/reduce_op_type.h" +#include "infini_train/include/nn/parallel/tensor_parallel.h" +#include "infini_train/include/nn/parallel/utils.h" +#include "infini_train/include/optimizer.h" +#include "infini_train/include/utils/global_module_hook_registry.h" +#include "infini_train/include/utils/precision_check_config.h" +#include "infini_train/include/utils/precision_checker.h" +#ifdef PROFILE_MODE +#include "infini_train/include/profiler.h" +#endif + +#include "example/common/tiny_shakespeare_dataset.h" +#include "example/common/tokenizer.h" +#include "example/qwen3/checkpoint_loader.h" +#include "example/qwen3/config.h" + +// TODO(jym): Reorganize CLI flags into categories for better readability and maintainability. +// I/O +DEFINE_string(input_bin, "", "input .bin to train on"); +DEFINE_string(input_val_bin, "", "input .bin to eval validation loss on"); +DEFINE_string(tokenizer_bin, "", "input .bin to tokenizer"); +// model bin file is downloaded and processed using the script at +// Converted from the official Hugging Face checkpoint into the shared LLMC v4 format. +DEFINE_string(llmc_filepath, "", "llmc model file path to load from"); +DEFINE_string(model, "qwen3", "Qwen/Qwen3-8B"); +// token layout for each step of the optimization +DEFINE_uint32(batch_size, 4, "batch size, in units of #batch dimensions"); +DEFINE_uint32(sequence_length, 64, "sequence length"); +DEFINE_uint32(total_batch_size, 256, "total desired batch size, in units of #tokens"); +// workload (number of steps) +DEFINE_uint32(num_iteration, 10, "number of iterations to run"); +DEFINE_uint32(freq_generate_txt, 10, "frequency of text generation"); +DEFINE_uint32(text_length, 64, "the length of the generated text"); +// optimization +DEFINE_double(learning_rate, 1e-5, "Peak learning rate."); +DEFINE_int32(zero_stage, 0, "ZeRO stage (0/1/2/3); 0 disables DistributedOptimizer"); +// lr scheduler +DEFINE_double(min_lr, 0.0, "Minimum learning rate."); +DEFINE_string(lr_decay_style, "constant", "LR decay style: none|constant|linear|cosine|inverse-square-root"); +DEFINE_int64(lr_warmup_iters, 0, "Number of linear warmup iterations."); +DEFINE_double(lr_warmup_init, 0.0, "Initial learning rate at the start of warmup."); +DEFINE_int64(lr_decay_iters, 0, "Number of iterations to decay LR over (0 = num_iteration)."); +// evaluation +DEFINE_uint32(val_loss_every, 0, "every how many steps to evaluate val loss?"); +DEFINE_uint32(sample_every, 0, "how often to sample from the model?"); +// debugging +DEFINE_bool(overfit_single_batch, true, "overfit just one batch of data"); +// memory management +DEFINE_string(device, "cuda", "device type (cpu/cuda), useless if using parallel training mode"); +// parallel +DEFINE_int32( + nthread_per_process, 1, + "Number of threads to use for each process. " + "When set > 1, enables data parallelism with device=cuda on the specified number of visible CUDA devices."); +DEFINE_uint32(tensor_parallel, 1, "Tensor Parallel world size"); +DEFINE_bool(sequence_parallel, false, "Whether to enable Sequence Parallel"); +DEFINE_uint32(pipeline_parallel, 1, "Pipeline Parallel world size, specified the number of PP stages."); +DEFINE_uint32(virtual_pipeline_parallel, 1, "Number of chunks in PP stage."); +// precision +DEFINE_string(dtype, "float32", "precision used in training (float32/bfloat16)"); +DEFINE_uint32(save_interval, 0, "save checkpoint every N steps; 0 disables saving"); +DEFINE_string(load, "", "checkpoint directory to resume from"); +DEFINE_string(save, "", "root directory used to store checkpoints"); +DEFINE_uint32(max_checkpoint_keep, 3, "max number of checkpoint steps to keep"); +DEFINE_bool(load_optimizer_state, true, "whether optimizer state is restored from checkpoints"); +DEFINE_bool(save_optimizer_state, true, "whether optimizer state is persisted in checkpoints"); + +// precision check +DEFINE_string( + precision_check, "", + "precision check config: level=N,format=simple|table,output_md5=true|false,output_path=PATH,baseline=PATH"); +// LoRA parameters +DEFINE_int32(lora_rank, 0, "LoRA rank (0 = disabled)"); +DEFINE_double(lora_alpha, 16.0, "LoRA alpha scaling factor"); +DEFINE_string(lora_target_modules, "c_attn,c_proj,c_fc,c_fc2", "LoRA target modules (comma-separated)"); +DEFINE_string(lora_save_path, "", "Path to save LoRA weights after training"); +DEFINE_string(lora_load_path, "", "Path to load LoRA weights from"); + +using namespace infini_train; + +namespace { +// validation +const std::unordered_set kSupportedModels = {"qwen3"}; +constexpr char kDeviceCPU[] = "cpu"; +constexpr char kDeviceCUDA[] = "cuda"; +constexpr char kDtypeFP32[] = "float32"; +constexpr char kDtypeBF16[] = "bfloat16"; +const std::unordered_set kSupportedLRDecayStyles + = {"none", "constant", "linear", "cosine", "inverse-square-root"}; +} // namespace + +DEFINE_validator(model, [](const char *, const std::string &value) { return kSupportedModels.contains(value); }); +DEFINE_validator(device, + [](const char *, const std::string &value) { return value == kDeviceCPU || value == kDeviceCUDA; }); +DEFINE_validator(zero_stage, [](const char *, int32_t value) { return value >= 0 && value <= 3; }); +DEFINE_validator(lr_decay_style, + [](const char *, const std::string &value) { return kSupportedLRDecayStyles.contains(value); }); + +void Train(const nn::parallel::Rank &rank) { + using namespace nn::parallel; + + { + if (rank.IsLastRank()) { + if (!FLAGS_save.empty() && FLAGS_save_interval == 0) { + LOG(FATAL) << "Invalid configuration: --save is set ('" << FLAGS_save + << "'), but --save_interval is 0. " + << "They must be set together."; + } + + if (FLAGS_save.empty() && FLAGS_save_interval > 0) { + LOG(FATAL) << "Invalid configuration: --save_interval is set to " << FLAGS_save_interval + << ", but --save is empty. " + << "They must be set together."; + } + } + } + + // select the device + Device device; + + int ddp_world_size = global::GetDataParallelSize(); + int tp_world_size = global::GetTensorParallelSize(); + int sp_world_size = global::GetSequenceParallelEnabled() ? tp_world_size : 1; + int pp_world_size = global::GetPipelineParallelSize(); + + if (FLAGS_sequence_parallel) { + CHECK_EQ(FLAGS_sequence_length % tp_world_size, 0) + << "sequence_length must be divisible by tp_world_size when SP is enabled (pad later if needed)."; + } + + int ddp_rank = 0; + int tp_rank = 0; + int pp_rank = 0; + + // Set thread-local global rank + nn::parallel::global::thread_global_rank = rank.GlobalRank(); + + const ProcessGroup *ddp_pg = nullptr; + const ProcessGroup *tp_pg = nullptr; + const ProcessGroup *pp_pg = nullptr; + + if (rank.IsParallel()) { + device = Device(Device::DeviceType::kCUDA, global::GetDeviceIndex(rank.thread_rank())); + auto *pg_factory = ProcessGroupFactory::Instance(device.type()); + + if (ddp_world_size > 1) { + ddp_pg = pg_factory->GetOrCreate(GetDataParallelProcessGroupName(rank.GlobalRank()), + GetDataParallelGroupRanks(rank.GlobalRank())); + ddp_rank = ddp_pg->GetGroupRank(rank.GlobalRank()); + } + + if (tp_world_size > 1) { + tp_pg = pg_factory->GetOrCreate(GetTensorParallelProcessGroupName(rank.GlobalRank()), + GetTensorParallelGroupRanks(rank.GlobalRank())); + tp_rank = tp_pg->GetGroupRank(rank.GlobalRank()); + // NOTE(zbl): Reserved for VocabParallelEmbedding + nn::parallel::tp_rank = tp_rank; + } + + if (pp_world_size > 1) { + pp_pg = pg_factory->GetOrCreate(GetPipelineParallelProcessGroupName(rank.GlobalRank()), + GetPipelineParallelGroupRanks(rank.GlobalRank())); + pp_rank = pp_pg->GetGroupRank(rank.GlobalRank()); + + nn::parallel::pp_rank = pp_rank; + } + } else { + device = FLAGS_device == kDeviceCPU ? Device() : Device(Device::DeviceType::kCUDA, 0); + } + + // calculate gradient accumulation from the desired total batch size and the current run configuration + const auto tokens_per_fwdbwd = FLAGS_batch_size * FLAGS_sequence_length * ddp_world_size; + CHECK_EQ(FLAGS_total_batch_size % tokens_per_fwdbwd, 0); + const auto grad_accum_steps = FLAGS_total_batch_size / tokens_per_fwdbwd; + if (rank.IsMainRank()) { + LOG(INFO) << "total desired batch size: " << FLAGS_total_batch_size + << " => calculated gradient accumulation steps: " << grad_accum_steps; + } + + // rng / reproducibility + // ManualSeed(42); + + nn::TransformerConfig model_config = qwen3::Qwen3Config(); + std::shared_ptr model = nullptr; + if (!FLAGS_llmc_filepath.empty()) { + model = qwen3::LoadFromLLMC(FLAGS_llmc_filepath); + } else { + model = std::make_shared(model_config); + } + + model->To(device); + + utils::PrecisionChecker::BuildNameMap(model.get()); + + // Apply LoRA using GetLoRAModel (in-place injection) + bool lora_enabled = FLAGS_lora_rank > 0; + if (lora_enabled) { + nn::lora::LoRAConfig lora_config{FLAGS_lora_rank, static_cast(FLAGS_lora_alpha), 0.0f, + nn::lora::ParseLoRATargetModules(FLAGS_lora_target_modules)}; + + // GetLoRAModel: in-place injection, modifies module tree directly + model = nn::lora::GetLoRAModel(model, lora_config); + + // Load LoRA weights if specified + if (!FLAGS_lora_load_path.empty()) { + LOG(INFO) << "Loading LoRA weights from: " << FLAGS_lora_load_path; + nn::lora::LoadLoRAWeights(model, FLAGS_lora_load_path); + } + + // Print LoRA summary + nn::lora::PrintLoRASummary(model, rank.GlobalRank()); + } + + LOG(INFO) << "Rank " << rank.GlobalRank() << ": Model loaded to device."; + + DataType dtype; + if (FLAGS_dtype == kDtypeFP32) { + dtype = DataType::kFLOAT32; + } else if (FLAGS_dtype == kDtypeBF16) { + dtype = DataType::kBFLOAT16; + } else { + LOG(FATAL) << "Rank " << rank.GlobalRank() << ": Datatype " << FLAGS_dtype << " not supported."; + } + + auto num_micro_batches = FLAGS_total_batch_size / (FLAGS_batch_size * FLAGS_sequence_length * ddp_world_size); + + if (pp_world_size > 1) { + // NOTE(dcj): To ensure that the tensor shapes at the pipeline stage boundaries remain correct + // when sequence parallelism (SP) is enabled, we need to divide by sp_world_size. + auto shapes = std::vector>{ + {FLAGS_batch_size, FLAGS_sequence_length / sp_world_size, model_config.n_embd}}; + + model = std::make_shared(model, pp_world_size, num_micro_batches, shapes, + pp_rank, device, model_config.GetChunkSize()); + if (ddp_world_size > 1) { + auto ddp_config = DistributedDataParallelConfig{.zero_stage = FLAGS_zero_stage}; + auto *mutable_chunks = dynamic_cast(model.get())->mutable_chunks(); + for (int chunk_id = 0; chunk_id < mutable_chunks->size(); ++chunk_id) { + (*mutable_chunks)[chunk_id] + = std::make_shared(mutable_chunks->at(chunk_id), rank, ddp_config); + } + } + } else if (ddp_world_size > 1) { + // NOTE(dcj): Complete all device (.to(device)) and dtype (.to(dtype)) conversions + // before wrapping the model with DistributedDataParallel (DDP). + // Otherwise, DDP's gradient hooks may be lost because new parameter tensors + // are created during the conversion. + + auto ddp_config = DistributedDataParallelConfig{.zero_stage = FLAGS_zero_stage}; + model = std::make_shared(model, rank, ddp_config); + } + + DistributedDataLoader train_loader(std::make_shared(FLAGS_input_bin, FLAGS_sequence_length), + pp_world_size > 1 ? FLAGS_batch_size * num_micro_batches : FLAGS_batch_size, + ddp_rank, ddp_world_size); + + std::optional val_loader = std::nullopt; + if (!FLAGS_input_val_bin.empty()) { + val_loader = DistributedDataLoader( + std::make_shared(FLAGS_input_val_bin, FLAGS_sequence_length), FLAGS_batch_size, + ddp_rank, ddp_world_size); + } + + // + // main training loop + // + std::unique_ptr tokenizer = nullptr; + if (!FLAGS_tokenizer_bin.empty()) { + tokenizer = std::make_unique(FLAGS_tokenizer_bin); + } + + // TODO(dcj): support more complex optimizer later + // auto optimizer = optimizers::Adam(model->Parameters(), FLAGS_learning_rate); + auto optimizer_creator = optimizers::Adam::Create(FLAGS_learning_rate, 0.9f, 0.95f, 1e-8f); + std::shared_ptr optimizer = nullptr; + + std::vector> params_to_optimize; + if (lora_enabled) { + params_to_optimize = nn::lora::GetLoRAParameters(model); + LOG(INFO) << "Optimizing " << params_to_optimize.size() << " LoRA parameters"; + } else { + params_to_optimize = model->Parameters(); + LOG(INFO) << "Optimizing " << params_to_optimize.size() << " model parameters"; + } + const auto named_parameters = model->NamedParameters(); + + if (FLAGS_zero_stage >= 1) { + auto model_chunks = (pp_world_size > 1) + ? *(dynamic_cast(model.get())->mutable_chunks()) + : std::vector>{model}; + optimizer = std::make_shared( + optimizer_creator, params_to_optimize, named_parameters, model_chunks, ddp_world_size, ddp_rank); + } else { + optimizer = optimizer_creator(params_to_optimize, named_parameters); + } + + const int64_t lr_decay_iters = FLAGS_lr_decay_iters > 0 ? FLAGS_lr_decay_iters : FLAGS_num_iteration; + TrainingLRSchedulerConfig sched_config; + sched_config.lr = static_cast(FLAGS_learning_rate); + sched_config.min_lr = static_cast(FLAGS_min_lr); + sched_config.lr_decay_style = FLAGS_lr_decay_style; + sched_config.lr_decay_iters = lr_decay_iters; + sched_config.lr_warmup_iters = FLAGS_lr_warmup_iters; + sched_config.lr_warmup_init = static_cast(FLAGS_lr_warmup_init); + auto scheduler = CreateLRScheduler(optimizer, sched_config); + + auto train_iter = train_loader.begin(); + std::shared_ptr loss_fn + = (tp_world_size > 1) ? std::static_pointer_cast(std::make_shared()) + : std::static_pointer_cast(std::make_shared()); + loss_fn->To(device); + LOG(INFO) << "Rank " << rank.GlobalRank() << ": start training"; + + auto impl = core::GetDeviceGuardImpl(device.type()); + + int start_step = 0; + TrainerState state; + const auto resume_result = ResumeFromCheckpoint({.resume_root = FLAGS_load, + .rank = rank, + .model = model, + .optimizer = FLAGS_load_optimizer_state ? optimizer : nullptr, + .model_config = model_config, + .state = state, + .lr_scheduler = scheduler}); + + start_step = resume_result.global_step; + size_t consumed_batches = resume_result.consumed_batches; + + // TODO(jym): Replace with Sampler abstraction when available. + // Skip dataloader to resume from the correct batch position. + if (consumed_batches > 0) { + size_t start = train_iter.BatchIndex(); + // Each rank processes every ddp_world_size-th batch starting from its own rank. + // num_skips calculates how many ++ iterations to reach the saved batch position. + size_t num_skips = (consumed_batches - start) / ddp_world_size; + for (size_t i = 0; i < num_skips; ++i) { ++train_iter; } + } + + auto save_checkpoint = [&](const std::filesystem::path &save_dir, int64_t global_step) { + SaveCheckpoint({ + .save_dir = save_dir, + .global_step = global_step, + .consumed_batches = consumed_batches, + .n_layer = model_config.n_layer, + .n_head = model_config.n_head, + .n_kv_head = model_config.n_kv_head, + .n_embd = model_config.n_embd, + .vocab_size = model_config.vocab_size, + .ddp_size = ddp_world_size, + .tp_size = tp_world_size, + .sp_size = sp_world_size, + .pp_size = pp_world_size, + .checkpoint_root_dir = FLAGS_save, + .max_checkpoint_keep = FLAGS_max_checkpoint_keep, + .rank = rank, + .model = *model, + .optimizer = FLAGS_save_optimizer_state ? optimizer.get() : nullptr, + .lr_scheduler = scheduler.get(), + }); + }; + + for (int step = start_step; step < FLAGS_num_iteration + 1; ++step) { + // Reset precision check counters at start of each iteration for file overwrite + utils::PrecisionChecker::ResetCounters(); + + const bool last_step = step == FLAGS_num_iteration; + + impl->ResetMemPoolHighWatermarks(device); + + const auto iter_start = std::chrono::high_resolution_clock::now(); + + // once in a while evaluate the validation dataset + if (FLAGS_val_loss_every > 0 && (step % FLAGS_val_loss_every == 0 || last_step) && val_loader.has_value()) { + // TODO(dcj): implement this after model.eval() is supported + } + // once in a while perform model inference on the master process + if (FLAGS_sample_every > 0 && (step % FLAGS_sample_every == 0 || last_step)) { + // TODO(dcj): implement this after model.eval() is supported + } + + // bit confusing: we want to make sure to eval and sample on 0th iteration + // but also after the very last iteration. so we loop for step <= num_iterations + // instead of just < num_iterations (one extra due to <=), only to do + // the validation/sampling one last time, and then we break right here as we're done. + if (last_step) { + break; + } + +#ifdef PROFILE_MODE + Profiler::Instance().SetTag("Step_" + std::to_string(step)); +#endif + + const float current_lr = scheduler ? scheduler->learning_rate() : static_cast(FLAGS_learning_rate); + float lossf = 0.0f; + if (pp_world_size == 1) { + // model->Train(); + optimizer->ZeroGrad(); + + // if we are trying to overfit a single batch, we reset the loader here + if (FLAGS_overfit_single_batch) { + // train_loader.Reset(); + } + + for (int micro_step = 0; micro_step < grad_accum_steps; ++micro_step) { + // enable autocast for the current step + infini_train::AutocastGuard autocast_guard(device.type(), dtype); + + // (bs, seq_len), (bs, seq_len) + auto [x, y] = *train_iter; + // if we are trying to overfit a single batch, we reset the loader here by commenting out the line below + // TODO(dcj): support dataloader.reset() later + ++train_iter; + consumed_batches = train_iter.BatchIndex(); + x = std::make_shared(x->To(device)); + y = std::make_shared(y->To(device)); + + LOG(INFO) << "Rank " << rank.GlobalRank() << ": start forward"; + // (bs, seq_len, vocab_size) + auto logits = (*model)({x, y})[0]; + LOG(INFO) << "Rank " << rank.GlobalRank() << ": finish model forward, start loss forward"; + auto loss = (*loss_fn)({logits, y})[0]; + // FIXME(jym): verify gradient accumulation precision + loss = loss / grad_accum_steps; + + // disable autocast for the current step (backward is not under autocast) + autocast_guard.Disable(); + + LOG(INFO) << "Rank " << rank.GlobalRank() << ": finish loss forward"; + + LOG(INFO) << "Rank " << rank.GlobalRank() << ": start backward"; + loss->Backward(); + // Defer the loss D2H copy until after backward; reading it earlier would synchronize CUDA + // between forward and backward. + auto loss_cpu = loss->To(Device()); + lossf += static_cast(loss_cpu.DataPtr())[0]; + LOG(INFO) << "Rank " << rank.GlobalRank() << ": finish backward"; + } + + optimizer->Step(); + if (scheduler) { + scheduler->Step(); + } + } else { + auto [x, y] = *train_iter; + // if we are trying to overfit a single batch, we reset the loader here by commenting out the line below + // TODO(dcj): support dataloader.reset() later + ++train_iter; + consumed_batches = train_iter.BatchIndex(); + x = std::make_shared(x->To(device)); + y = std::make_shared(y->To(device)); + + lossf = model->TrainStep({x}, {y}, optimizer, loss_fn, dtype); + if (scheduler) { + scheduler->Step(); + } + } + + if (ddp_world_size > 1) { + auto lossf_tensor = std::make_shared(&lossf, std::vector{}, DataType::kFLOAT32, device); + function::AllReduce(lossf_tensor, function::ReduceOpType::kAvg, ddp_pg); + lossf = static_cast(lossf_tensor->To(Device()).DataPtr())[0]; + } + + const auto iter_end = std::chrono::high_resolution_clock::now(); + const double duration_us = std::chrono::duration(iter_end - iter_start).count(); + const double tps = FLAGS_total_batch_size / (duration_us / 1e6); + + if (rank.IsLastRank()) { + size_t used_mb = 0, reserved_mb = 0; + std::tie(used_mb, reserved_mb) = impl->GetMemPoolPeakMB(device); + LOG(ERROR) << std::format("step {:4d}/{} | train loss {:.6f} | lr {:.2e} | ({:.2f} ms | {:.0f} tok/s | " + "peak used: {:5d} MB | peak reserved: {:5d} MB, DP={}, TP={}, SP={}, PP={})", + step + 1, FLAGS_num_iteration, lossf, current_lr, duration_us / 1e3f, tps, + used_mb, reserved_mb, ddp_world_size, tp_world_size, sp_world_size, + pp_world_size); + + if ((step + 1) % FLAGS_freq_generate_txt == 0) { + // FIXME(jym): to support PP + if (tokenizer) { + CHECK_EQ(pp_world_size, 1); + tokenizer->GenerateText(*model, FLAGS_batch_size, FLAGS_sequence_length, FLAGS_text_length, device); + } + } + } + + if (!FLAGS_save.empty() && FLAGS_save_interval > 0) { + if ((step + 1) % FLAGS_save_interval == 0 || (step + 1) == FLAGS_num_iteration) { + std::filesystem::path step_dir + = std::filesystem::path(FLAGS_save) / std::format("checkpoint_step_{:06d}", step + 1); + if (rank.IsParallel()) { + step_dir /= std::format("rank_{:06d}", rank.GlobalRank()); + } + save_checkpoint(step_dir, step + 1); + } + } + } + + // Save LoRA weights if enabled and path specified + if (lora_enabled && !FLAGS_lora_save_path.empty()) { + LOG(INFO) << "Saving LoRA weights to: " << FLAGS_lora_save_path; + nn::lora::SaveLoRAWeights(model, FLAGS_lora_save_path); + } + +#ifdef PROFILE_MODE + Profiler::Instance().Report("qwen3.report", Profiler::SortBy::DeviceTimePercentage); + Profiler::Instance().PrintRecords("qwen3.records.log"); +#endif +} + +int main(int argc, char *argv[]) { + gflags::ParseCommandLineFlags(&argc, &argv, true); + google::InitGoogleLogging(argv[0]); + + auto precision_config = utils::PrecisionCheckConfig::Parse(FLAGS_precision_check); + nn::parallel::global::InitAllEnv(FLAGS_nthread_per_process, FLAGS_tensor_parallel, FLAGS_sequence_parallel, + FLAGS_pipeline_parallel, FLAGS_virtual_pipeline_parallel); + utils::PrecisionCheckEnv::Instance().Init(precision_config); + + LOG(INFO) << nn::parallel::global::ProcessGroupOverview(); + + if (FLAGS_nthread_per_process > 1) { + std::vector threads; + for (int idx = 0; idx < FLAGS_nthread_per_process; ++idx) { + nn::parallel::Rank rank(nn::parallel::global::GetGlobalProcRank(), idx, + nn::parallel::global::GetNprocPerNode(), FLAGS_nthread_per_process); + threads.emplace_back(Train, rank); + } + + for (auto &thread : threads) { thread.join(); } + } else { + nn::parallel::Rank rank(nn::parallel::global::GetGlobalProcRank(), 0, nn::parallel::global::GetNprocPerNode(), + FLAGS_nthread_per_process); + Train(rank); + } + + gflags::ShutDownCommandLineFlags(); + google::ShutdownGoogleLogging(); + + return 0; +} 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 60373414..3b20b8d1 100644 --- a/infini_train/include/nn/modules/transformer/causal_self_attention.h +++ b/infini_train/include/nn/modules/transformer/causal_self_attention.h @@ -4,6 +4,7 @@ #include #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 q_norm_; + std::shared_ptr k_norm_; + // Setup method for different attention modes void SetupAttention(const TransformerConfig &config); diff --git a/infini_train/include/nn/modules/transformer/transformer_config.h b/infini_train/include/nn/modules/transformer/transformer_config.h index a646ab14..485f6c8e 100644 --- a/infini_train/include/nn/modules/transformer/transformer_config.h +++ b/infini_train/include/nn/modules/transformer/transformer_config.h @@ -81,8 +81,9 @@ struct TransformerConfig { std::optional 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; }; diff --git a/infini_train/include/nn/modules/transformer/utils.h b/infini_train/include/nn/modules/transformer/utils.h index 30db08e6..a5994d9c 100644 --- a/infini_train/include/nn/modules/transformer/utils.h +++ b/infini_train/include/nn/modules/transformer/utils.h @@ -13,5 +13,5 @@ std::shared_ptr PrecomputeFreqsCis(int64_t dim, int64_t end, float theta std::tuple, std::shared_ptr> ApplyRotaryEmbedding(const std::shared_ptr &xq, const std::shared_ptr &xk, - const std::shared_ptr &freqs_cis); + const std::shared_ptr &freqs_cis, bool rotary_interleaved = true); } // 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 bc4c6bd4..573ddb24 100644 --- a/infini_train/src/nn/modules/transformer/causal_self_attention.cc +++ b/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(head_dim_, config_.qk_norm_eps); + k_norm_ = std::make_shared(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( @@ -122,10 +129,17 @@ CausalSelfAttention::Forward(const std::vectorSlice(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 diff --git a/infini_train/src/nn/modules/transformer/utils.cc b/infini_train/src/nn/modules/transformer/utils.cc index 4ec11f2d..3f41996b 100644 --- a/infini_train/src/nn/modules/transformer/utils.cc +++ b/infini_train/src/nn/modules/transformer/utils.cc @@ -34,7 +34,7 @@ std::shared_ptr PrecomputeFreqsCis(int64_t dim, int64_t end, float theta std::tuple, std::shared_ptr> ApplyRotaryEmbedding(const std::shared_ptr &xq, const std::shared_ptr &xk, - const std::shared_ptr &freqs_cis) { + const std::shared_ptr &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 &xq, const std::shared_ptrSlice(-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 slice_pair = [rotary_interleaved](const std::shared_ptr &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>{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); + auto rotate = [&](const std::shared_ptr &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>{rotated_left, rotated_right}, -1) + ->Flatten(-2); + } + return nn::function::Concat(std::vector>{rotated_left, rotated_right}, -1); + }; - return {q_rotated, k_rotated}; + return {rotate(xq), rotate(xk)}; } } // namespace infini_train