diff --git a/example/gpt2/checkpoint_loader.cc b/example/gpt2/checkpoint_loader.cc index f2747dd4..b140edbd 100644 --- a/example/gpt2/checkpoint_loader.cc +++ b/example/gpt2/checkpoint_loader.cc @@ -187,16 +187,14 @@ std::shared_ptr LoadFromLLMC( } // transformer.h.{i}.ln_1.weight - int local_layer_index = 0; for (int idx = 0; idx < n_layer; ++idx) { if (owned_layers[idx]) { auto &tensor = state_dict[std::format( "{}.{}.{}.{}.{}", nn::TransformerModel::kTransformerModelName, - nn::TransformerChunk::kHLayerName, std::to_string(local_layer_index), + nn::TransformerChunk::kHLayerName, std::to_string(idx), nn::TransformerLayer::kLn1LayerName, nn::LayerNorm::kParamWeightName)]; ReadVectorAllFloat(ifs, static_cast(tensor->DataPtr()), n_embd); - ++local_layer_index; } else { size_t ln_1_w_bytes = n_embd * sizeof(float); ifs.seekg(ln_1_w_bytes, std::ios::cur); @@ -204,15 +202,13 @@ std::shared_ptr LoadFromLLMC( } // transformer.h.{i}.ln_1.bias - local_layer_index = 0; for (int idx = 0; idx < n_layer; ++idx) { if (owned_layers[idx]) { auto &tensor = state_dict[std::format( "{}.{}.{}.{}.{}", nn::TransformerModel::kTransformerModelName, - nn::TransformerChunk::kHLayerName, std::to_string(local_layer_index), + nn::TransformerChunk::kHLayerName, std::to_string(idx), nn::TransformerLayer::kLn1LayerName, nn::LayerNorm::kParamBiasName)]; ReadVectorAllFloat(ifs, static_cast(tensor->DataPtr()), n_embd); - ++local_layer_index; } else { size_t ln_1_b_bytes = n_embd * sizeof(float); ifs.seekg(ln_1_b_bytes, std::ios::cur); @@ -221,12 +217,11 @@ std::shared_ptr LoadFromLLMC( // transformer.h.{i}.attn.c_attn.weight (ColumnParallelLinear, but actually // applies on "rows") - local_layer_index = 0; for (int idx = 0; idx < n_layer; ++idx) { if (owned_layers[idx]) { auto &tensor = state_dict[std::format( "{}.{}.{}.{}.{}.{}", nn::TransformerModel::kTransformerModelName, - nn::TransformerChunk::kHLayerName, std::to_string(local_layer_index), + nn::TransformerChunk::kHLayerName, std::to_string(idx), nn::TransformerLayer::kAttnLayerName, nn::CausalSelfAttention::kCAttnLayerName, nn::parallel::ColumnParallelLinear::kParamWeightName)]; @@ -261,7 +256,6 @@ std::shared_ptr LoadFromLLMC( /*row_start=*/2 * n_embd + tp_rank * local_C, /*row_cnt=*/local_C); - ++local_layer_index; } else { size_t c_attn_w_bytes = qkv_out * n_embd * sizeof(float); ifs.seekg(c_attn_w_bytes, std::ios::cur); @@ -269,12 +263,11 @@ std::shared_ptr LoadFromLLMC( } // transformer.h.{i}.attn.c_attn.bias (ColumnParallelLinear) - local_layer_index = 0; for (int idx = 0; idx < n_layer; ++idx) { if (owned_layers[idx]) { auto &tensor = state_dict[std::format( "{}.{}.{}.{}.{}.{}", nn::TransformerModel::kTransformerModelName, - nn::TransformerChunk::kHLayerName, std::to_string(local_layer_index), + nn::TransformerChunk::kHLayerName, std::to_string(idx), nn::TransformerLayer::kAttnLayerName, nn::CausalSelfAttention::kCAttnLayerName, nn::parallel::ColumnParallelLinear::kParamBiasName)]; @@ -307,7 +300,6 @@ std::shared_ptr LoadFromLLMC( /*start=*/2 * n_embd + tp_rank * local_C, /*cnt=*/local_C); - ++local_layer_index; } else { size_t c_attn_b_bytes = qkv_out * sizeof(float); ifs.seekg(c_attn_b_bytes, std::ios::cur); @@ -316,18 +308,16 @@ std::shared_ptr LoadFromLLMC( // transformer.h.{i}.attn.c_proj.weight (RowParallelLinear, but actually // applies on "columns") - local_layer_index = 0; for (int idx = 0; idx < n_layer; ++idx) { if (owned_layers[idx]) { auto &tensor = state_dict[std::format( "{}.{}.{}.{}.{}.{}", nn::TransformerModel::kTransformerModelName, - nn::TransformerChunk::kHLayerName, std::to_string(local_layer_index), + nn::TransformerChunk::kHLayerName, std::to_string(idx), nn::TransformerLayer::kAttnLayerName, nn::CausalSelfAttention::kCProjLayerName, nn::parallel::RowParallelLinear::kParamWeightName)]; ReadMatrixColShardFloat(ifs, static_cast(tensor->DataPtr()), n_embd, n_embd, tp_rank * in_pp, in_pp); - ++local_layer_index; } else { size_t c_proj_w_bytes = n_embd * n_embd * sizeof(float); ifs.seekg(c_proj_w_bytes, std::ios::cur); @@ -335,17 +325,15 @@ std::shared_ptr LoadFromLLMC( } // transformer.h.{i}.attn.c_proj.bias (RowParallelLinear, no shard on bias) - local_layer_index = 0; for (int idx = 0; idx < n_layer; ++idx) { if (owned_layers[idx]) { auto &tensor = state_dict[std::format( "{}.{}.{}.{}.{}.{}", nn::TransformerModel::kTransformerModelName, - nn::TransformerChunk::kHLayerName, std::to_string(local_layer_index), + nn::TransformerChunk::kHLayerName, std::to_string(idx), nn::TransformerLayer::kAttnLayerName, nn::CausalSelfAttention::kCProjLayerName, nn::parallel::RowParallelLinear::kParamBiasName)]; ReadVectorAllFloat(ifs, static_cast(tensor->DataPtr()), n_embd); - ++local_layer_index; } else { size_t c_proj_b_bytes = n_embd * sizeof(float); ifs.seekg(c_proj_b_bytes, std::ios::cur); @@ -353,16 +341,14 @@ std::shared_ptr LoadFromLLMC( } // transformer.h.{i}.ln_2.weight - local_layer_index = 0; for (int idx = 0; idx < n_layer; ++idx) { if (owned_layers[idx]) { auto &tensor = state_dict[std::format( "{}.{}.{}.{}.{}", nn::TransformerModel::kTransformerModelName, - nn::TransformerChunk::kHLayerName, std::to_string(local_layer_index), + nn::TransformerChunk::kHLayerName, std::to_string(idx), nn::TransformerLayer::kLn2LayerName, nn::LayerNorm::kParamWeightName)]; ReadVectorAllFloat(ifs, static_cast(tensor->DataPtr()), n_embd); - ++local_layer_index; } else { size_t ln_2_w_bytes = n_embd * sizeof(float); ifs.seekg(ln_2_w_bytes, std::ios::cur); @@ -370,15 +356,13 @@ std::shared_ptr LoadFromLLMC( } // transformer.h.{i}.ln_2.bias - local_layer_index = 0; for (int idx = 0; idx < n_layer; ++idx) { if (owned_layers[idx]) { auto &tensor = state_dict[std::format( "{}.{}.{}.{}.{}", nn::TransformerModel::kTransformerModelName, - nn::TransformerChunk::kHLayerName, std::to_string(local_layer_index), + nn::TransformerChunk::kHLayerName, std::to_string(idx), nn::TransformerLayer::kLn2LayerName, nn::LayerNorm::kParamBiasName)]; ReadVectorAllFloat(ifs, static_cast(tensor->DataPtr()), n_embd); - ++local_layer_index; } else { size_t ln_2_b_bytes = n_embd * sizeof(float); ifs.seekg(ln_2_b_bytes, std::ios::cur); @@ -387,17 +371,15 @@ std::shared_ptr LoadFromLLMC( // transformer.h.{i}.mlp.c_fc.weight (ColumnParallelLinear, but actually // applies on "rows") - local_layer_index = 0; for (int idx = 0; idx < n_layer; ++idx) { if (owned_layers[idx]) { auto &tensor = state_dict[std::format( "{}.{}.{}.{}.{}.{}", nn::TransformerModel::kTransformerModelName, - nn::TransformerChunk::kHLayerName, std::to_string(local_layer_index), + nn::TransformerChunk::kHLayerName, std::to_string(idx), nn::TransformerLayer::kMlpLayerName, nn::MLP::kCFcLayerName, nn::parallel::ColumnParallelLinear::kParamWeightName)]; ReadMatrixRowShardFloat(ifs, static_cast(tensor->DataPtr()), fc_out, n_embd, fc_start, fc_pp); - ++local_layer_index; } else { size_t c_fc_w_bytes = fc_out * n_embd * sizeof(float); ifs.seekg(c_fc_w_bytes, std::ios::cur); @@ -405,17 +387,15 @@ std::shared_ptr LoadFromLLMC( } // transformer.h.{i}.mlp.c_fc.bias (ColumnParallelLinear) - local_layer_index = 0; for (int idx = 0; idx < n_layer; ++idx) { if (owned_layers[idx]) { auto &tensor = state_dict[std::format( "{}.{}.{}.{}.{}.{}", nn::TransformerModel::kTransformerModelName, - nn::TransformerChunk::kHLayerName, std::to_string(local_layer_index), + nn::TransformerChunk::kHLayerName, std::to_string(idx), nn::TransformerLayer::kMlpLayerName, nn::MLP::kCFcLayerName, nn::parallel::ColumnParallelLinear::kParamBiasName)]; ReadVectorShardFloat(ifs, static_cast(tensor->DataPtr()), fc_out, fc_start, fc_pp); - ++local_layer_index; } else { size_t c_fc_b_bytes = fc_out * sizeof(float); ifs.seekg(c_fc_b_bytes, std::ios::cur); @@ -424,17 +404,15 @@ std::shared_ptr LoadFromLLMC( // transformer.h.{i}.mlp.c_proj.weight (RowParallelLinear, but actually // applies on "columns") - local_layer_index = 0; for (int idx = 0; idx < n_layer; ++idx) { if (owned_layers[idx]) { auto &tensor = state_dict[std::format( "{}.{}.{}.{}.{}.{}", nn::TransformerModel::kTransformerModelName, - nn::TransformerChunk::kHLayerName, std::to_string(local_layer_index), + nn::TransformerChunk::kHLayerName, std::to_string(idx), nn::TransformerLayer::kMlpLayerName, nn::MLP::kCProjLayerName, nn::parallel::RowParallelLinear::kParamWeightName)]; ReadMatrixColShardFloat(ifs, static_cast(tensor->DataPtr()), n_embd, fc_out, tp_rank * in4_pp, in4_pp); - ++local_layer_index; } else { size_t c_proj_w_bytes = fc_out * n_embd * sizeof(float); ifs.seekg(c_proj_w_bytes, std::ios::cur); @@ -442,16 +420,14 @@ std::shared_ptr LoadFromLLMC( } // transformer.h.{i}.mlp.c_proj.bias (RowParallelLinear, no shard on bias) - local_layer_index = 0; for (int idx = 0; idx < n_layer; ++idx) { if (owned_layers[idx]) { auto &tensor = state_dict[std::format( "{}.{}.{}.{}.{}.{}", nn::TransformerModel::kTransformerModelName, - nn::TransformerChunk::kHLayerName, std::to_string(local_layer_index), + nn::TransformerChunk::kHLayerName, std::to_string(idx), nn::TransformerLayer::kMlpLayerName, nn::MLP::kCProjLayerName, nn::parallel::RowParallelLinear::kParamBiasName)]; ReadVectorAllFloat(ifs, static_cast(tensor->DataPtr()), n_embd); - ++local_layer_index; } else { size_t c_proj_b_bytes = n_embd * sizeof(float); ifs.seekg(c_proj_b_bytes, std::ios::cur); diff --git a/example/gpt2/main.cc b/example/gpt2/main.cc index 4c0321ad..e0d35588 100644 --- a/example/gpt2/main.cc +++ b/example/gpt2/main.cc @@ -533,11 +533,13 @@ void Train(const nn::parallel::Rank &rank) { .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, + .original_vocab_size = model_config.original_vocab_size, + .padded_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, + .vpp_size = static_cast(FLAGS_virtual_pipeline_parallel), .checkpoint_root_dir = FLAGS_save, .max_checkpoint_keep = FLAGS_max_checkpoint_keep, .rank = rank, diff --git a/example/llama3/checkpoint_loader.cc b/example/llama3/checkpoint_loader.cc index e6513a4e..bbd6a996 100644 --- a/example/llama3/checkpoint_loader.cc +++ b/example/llama3/checkpoint_loader.cc @@ -208,15 +208,13 @@ std::shared_ptr LoadFromLLMC( } // 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::TransformerChunk::kHLayerName, std::to_string(i), 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); @@ -226,12 +224,11 @@ std::shared_ptr LoadFromLLMC( // 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)] × 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::TransformerChunk::kHLayerName, std::to_string(i), nn::TransformerLayer::kAttnLayerName, nn::CausalSelfAttention::kCAttnLayerName, nn::parallel::ColumnParallelLinear::kParamWeightName)]; @@ -265,7 +262,6 @@ std::shared_ptr LoadFromLLMC( /*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); @@ -275,19 +271,17 @@ std::shared_ptr LoadFromLLMC( // 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::TransformerChunk::kHLayerName, std::to_string(i), 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); @@ -296,15 +290,13 @@ std::shared_ptr LoadFromLLMC( } // 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::TransformerChunk::kHLayerName, std::to_string(i), 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); @@ -313,19 +305,17 @@ std::shared_ptr LoadFromLLMC( // 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(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::TransformerChunk::kHLayerName, std::to_string(i), nn::TransformerLayer::kMlpLayerName, nn::MLP::kCFcLayerName, nn::parallel::ColumnParallelLinear::kParamWeightName)]; float *dst = static_cast(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 { size_t fc_bytes = static_cast(ffn_hidden) * n_embd * sizeof(float); @@ -335,18 +325,16 @@ std::shared_ptr LoadFromLLMC( // 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(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::TransformerChunk::kHLayerName, std::to_string(i), 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); @@ -356,19 +344,17 @@ std::shared_ptr LoadFromLLMC( // 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::TransformerChunk::kHLayerName, std::to_string(i), 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); diff --git a/example/llama3/main.cc b/example/llama3/main.cc index 4554649a..f8188cef 100644 --- a/example/llama3/main.cc +++ b/example/llama3/main.cc @@ -511,11 +511,13 @@ void Train(const nn::parallel::Rank &rank) { .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, + .original_vocab_size = model_config.original_vocab_size, + .padded_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, + .vpp_size = static_cast(FLAGS_virtual_pipeline_parallel), .checkpoint_root_dir = FLAGS_save, .max_checkpoint_keep = FLAGS_max_checkpoint_keep, .rank = rank, diff --git a/example/qwen3/checkpoint_loader.cc b/example/qwen3/checkpoint_loader.cc index d64cd608..7270f7bb 100644 --- a/example/qwen3/checkpoint_loader.cc +++ b/example/qwen3/checkpoint_loader.cc @@ -231,27 +231,24 @@ std::shared_ptr LoadFromLLMC( } // 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::TransformerChunk::kHLayerName, std::to_string(i), 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::TransformerChunk::kHLayerName, std::to_string(i), nn::TransformerLayer::kAttnLayerName, nn::CausalSelfAttention::kQNormLayerName, nn::RMSNorm::kParamWeightName)]; @@ -260,13 +257,12 @@ std::shared_ptr LoadFromLLMC( auto &k_norm_tensor = state_dict[std::format( "{}.{}.{}.{}.{}.{}", nn::TransformerModel::kTransformerModelName, - nn::TransformerChunk::kHLayerName, std::to_string(local_layer_index), + nn::TransformerChunk::kHLayerName, std::to_string(i), 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); @@ -276,12 +272,11 @@ std::shared_ptr LoadFromLLMC( // 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::TransformerChunk::kHLayerName, std::to_string(i), nn::TransformerLayer::kAttnLayerName, nn::CausalSelfAttention::kCAttnLayerName, nn::parallel::ColumnParallelLinear::kParamWeightName)]; @@ -315,7 +310,6 @@ std::shared_ptr LoadFromLLMC( /*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); @@ -325,19 +319,17 @@ std::shared_ptr LoadFromLLMC( // 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::TransformerChunk::kHLayerName, std::to_string(i), 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); @@ -346,15 +338,13 @@ std::shared_ptr LoadFromLLMC( } // 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::TransformerChunk::kHLayerName, std::to_string(i), 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); @@ -363,18 +353,16 @@ std::shared_ptr LoadFromLLMC( // 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(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::TransformerChunk::kHLayerName, std::to_string(i), 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 fc_bytes = static_cast(ffn_hidden) * n_embd * sizeof(float); @@ -384,18 +372,16 @@ std::shared_ptr LoadFromLLMC( // 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(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::TransformerChunk::kHLayerName, std::to_string(i), nn::TransformerLayer::kMlpLayerName, nn::MLP::kCFcLayerName, nn::parallel::ColumnParallelLinear::kParamWeightName)]; float *dst = static_cast(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 { size_t fc2_bytes = static_cast(ffn_hidden) * n_embd * sizeof(float); @@ -405,19 +391,17 @@ std::shared_ptr LoadFromLLMC( // 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::TransformerChunk::kHLayerName, std::to_string(i), 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); diff --git a/example/qwen3/main.cc b/example/qwen3/main.cc index 085687d3..a9f925fe 100644 --- a/example/qwen3/main.cc +++ b/example/qwen3/main.cc @@ -509,11 +509,13 @@ void Train(const nn::parallel::Rank &rank) { .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, + .original_vocab_size = model_config.original_vocab_size, + .padded_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, + .vpp_size = static_cast(FLAGS_virtual_pipeline_parallel), .checkpoint_root_dir = FLAGS_save, .max_checkpoint_keep = FLAGS_max_checkpoint_keep, .rank = rank, diff --git a/infini_train/include/checkpoint/checkpoint.h b/infini_train/include/checkpoint/checkpoint.h index a6a62438..e56e9f67 100644 --- a/infini_train/include/checkpoint/checkpoint.h +++ b/infini_train/include/checkpoint/checkpoint.h @@ -5,6 +5,13 @@ #include #include #include +#include + +#include "infini_train/include/shard_spec.h" + +namespace infini_train::checkpoint { +struct WriteItem; +} // namespace infini_train::checkpoint namespace infini_train { class Optimizer; @@ -22,11 +29,13 @@ struct TrainerState { int64_t n_head = 0; int64_t n_kv_head = 0; int64_t n_embd = 0; - int64_t vocab_size = 0; + int64_t original_vocab_size = 0; + int64_t padded_vocab_size = 0; int ddp_size = 1; int tp_size = 1; int sp_size = 1; int pp_size = 1; + int vpp_size = 1; }; class Checkpoint { @@ -39,8 +48,50 @@ class Checkpoint { nn::Module &model, Optimizer *optimizer, TrainerState &state, LRScheduler *lr_scheduler); + static void SaveStateDictFile( + const std::filesystem::path &path, + const std::unordered_map> + &state_dict); + + static std::unordered_map> + LoadStateDictFile(const std::filesystem::path &path); + + struct CheckpointMetadata { + int version = 0; + + struct TensorEntry { + std::string key; + std::string dtype_str; + std::vector global_shape; + std::vector local_shape; + std::vector global_offset; + std::vector axis_fragmentations; + std::vector segments; + std::string file; + uint64_t offset = 0; + uint64_t byte_size = 0; + std::vector stored_on_ranks; + int pp_rank = 0; + }; + + std::vector tensors; + bool has_metadata = false; + }; + + static CheckpointMetadata LoadMetadata( + const std::filesystem::path &checkpoint_dir); + static void SaveMetadataFile(const std::filesystem::path &path, + const CheckpointMetadata &metadata); + private: - static void SaveStateDict( + struct SavedTensorLocation { + uint64_t data_offset = 0; + uint64_t byte_size = 0; + }; + using SavedTensorLocations = + std::unordered_map; + + static SavedTensorLocations SaveStateDict( const std::filesystem::path &path, const std::unordered_map> &state_dict); @@ -48,6 +99,16 @@ class Checkpoint { static std::unordered_map> LoadStateDict( const std::filesystem::path &path); + static void SaveLocalShard( + const std::filesystem::path &checkpoint_dir, + const ShardedStateDict &sharded_sd, + const std::vector &write_items, + const std::unordered_map> + &state_dict, + const std::unordered_map> + &optimizer_state, + int global_rank); + static void SaveTrainerState(const std::filesystem::path &path, const TrainerState &state); static TrainerState LoadTrainerState(const std::filesystem::path &path); diff --git a/infini_train/include/checkpoint/checkpoint_manager.h b/infini_train/include/checkpoint/checkpoint_manager.h index 4d3387f7..4ae92aa8 100644 --- a/infini_train/include/checkpoint/checkpoint_manager.h +++ b/infini_train/include/checkpoint/checkpoint_manager.h @@ -50,11 +50,13 @@ struct SaveCheckpointArgs { int64_t n_head = 0; int64_t n_kv_head = 0; int64_t n_embd = 0; - int64_t vocab_size = 0; + int64_t original_vocab_size = 0; + int64_t padded_vocab_size = 0; int ddp_size = 1; int tp_size = 1; int sp_size = 1; int pp_size = 1; + int vpp_size = 1; std::filesystem::path checkpoint_root_dir; size_t max_checkpoint_keep = 0; const nn::parallel::Rank &rank; diff --git a/infini_train/include/checkpoint/constants.h b/infini_train/include/checkpoint/constants.h new file mode 100644 index 00000000..1604ecfd --- /dev/null +++ b/infini_train/include/checkpoint/constants.h @@ -0,0 +1,17 @@ +#pragma once + +namespace infini_train::checkpoint { + +inline constexpr char kModelCheckpointFilename[] = "model.ckpt"; +inline constexpr char kOptimizerCheckpointFilename[] = "optimizer.ckpt"; +inline constexpr char kMetadataFilename[] = "metadata.json"; +inline constexpr char kTemporaryMetadataFilename[] = "metadata.json.tmp"; +inline constexpr char kTrainerStateFilename[] = "trainer_state.json"; +inline constexpr char kLRSchedulerFilename[] = "lr_scheduler.ckpt"; +inline constexpr char kLatestIterationFilename[] = + "latest_checkpointed_iteration.txt"; +inline constexpr char kTemporaryLatestIterationFilename[] = + "latest_checkpointed_iteration.txt.tmp"; +inline constexpr int kDistributedCheckpointVersion = 3; + +} // namespace infini_train::checkpoint diff --git a/infini_train/include/checkpoint/load_planner.h b/infini_train/include/checkpoint/load_planner.h new file mode 100644 index 00000000..09535dfa --- /dev/null +++ b/infini_train/include/checkpoint/load_planner.h @@ -0,0 +1,55 @@ +#pragma once + +#include +#include +#include +#include + +#include "infini_train/include/checkpoint/checkpoint.h" +#include "infini_train/include/datatype.h" + +namespace infini_train { +struct ShardedStateDict; +} // namespace infini_train + +namespace infini_train::checkpoint { + +// One storage-region transfer from a saved shard into a target local tensor. +struct ReadItem { + std::string key; + std::string filename; + DataType dtype = DataType::kFLOAT32; + std::vector global_shape; + uint64_t byte_size = 0; + uint64_t data_offset = 0; + int shard_dim = -1; + int64_t source_offset = 0; + int64_t target_offset = 0; + int64_t length = 0; + std::vector source_shape; +}; + +// All reads required to materialize one target local tensor. +struct TargetTensorPlan { + std::string key; + DataType dtype = DataType::kFLOAT32; + std::vector global_shape; + std::vector target_shape; + int shard_dim = -1; + int64_t trailing_zero_fill = 0; + std::vector reads; +}; + +// Complete load plan for one rank. +struct LoadPlan { + std::map tensors; +}; + +class LoadPlanner { + public: + // Compute saved-to-target overlaps from explicit global shard coordinates. + static LoadPlan PlanReshard(const Checkpoint::CheckpointMetadata &metadata, + const ShardedStateDict &target_state_dict); +}; + +} // namespace infini_train::checkpoint diff --git a/infini_train/include/checkpoint/load_strategy.h b/infini_train/include/checkpoint/load_strategy.h new file mode 100644 index 00000000..e7adff18 --- /dev/null +++ b/infini_train/include/checkpoint/load_strategy.h @@ -0,0 +1,33 @@ +#pragma once + +#include +#include +#include +#include + +namespace infini_train { +class Tensor; +} + +namespace infini_train::checkpoint { +struct LoadPlan; + +using LoadedStateDict = + std::unordered_map>; + +class LoadStrategy { + public: + virtual ~LoadStrategy() = default; + virtual LoadedStateDict Execute(const std::filesystem::path &checkpoint_dir, + const LoadPlan &plan) = 0; +}; + +/// Reads source regions directly from metadata offsets while caching one open +/// stream per file. +class IndexedRegionLoadStrategy final : public LoadStrategy { + public: + LoadedStateDict Execute(const std::filesystem::path &checkpoint_dir, + const LoadPlan &plan) override; +}; + +} // namespace infini_train::checkpoint diff --git a/infini_train/include/checkpoint/save_planner.h b/infini_train/include/checkpoint/save_planner.h new file mode 100644 index 00000000..bd8300a6 --- /dev/null +++ b/infini_train/include/checkpoint/save_planner.h @@ -0,0 +1,53 @@ +#pragma once + +#include +#include +#include +#include +#include + +#include "infini_train/include/datatype.h" +#include "infini_train/include/optimizer.h" +#include "infini_train/include/shard_spec.h" + +namespace infini_train { +class Tensor; +} + +namespace infini_train::checkpoint { + +// Physical write description for one local tensor shard. +struct WriteItem { + std::string key; + std::string filename; // Checkpoint payload filename. + uint64_t byte_size = 0; // Tensor payload size in bytes. + DataType dtype = DataType::kFLOAT32; + std::vector local_shape; + std::vector global_offset; + std::vector axis_fragmentations; + int rank = 0; +}; + +// Build the local tensor write layout from a ShardedStateDict. +class SavePlanner { + public: + static std::vector Plan(const ShardedStateDict &sd, int rank); +}; + +ShardedStateDict BuildOptimizerShardedStateDict( + const ShardedStateDict &model_state, + const std::unordered_map> + &optimizer_state, + const NamedParameterList &optimizer_parameters = {}); + +// Return the number of payload bytes required by a tensor. +inline uint64_t TensorByteSize(DataType dtype, + const std::vector &shape) { + uint64_t numel = 1; + for (auto d : shape) { + numel *= static_cast(d); + } + return numel * static_cast(kDataTypeToSize.at(dtype)); +} + +} // namespace infini_train::checkpoint diff --git a/infini_train/include/nn/lora/lora_parallel_linear.h b/infini_train/include/nn/lora/lora_parallel_linear.h index 6932bc04..3a0f93be 100644 --- a/infini_train/include/nn/lora/lora_parallel_linear.h +++ b/infini_train/include/nn/lora/lora_parallel_linear.h @@ -6,7 +6,9 @@ #include #include "infini_train/include/nn/lora/lora_config.h" +#include "infini_train/include/nn/modules/module.h" #include "infini_train/include/nn/parallel/tensor_parallel.h" +#include "infini_train/include/shard_spec.h" namespace infini_train { class Tensor; @@ -21,6 +23,7 @@ namespace infini_train::nn::lora { // (implemented as ColumnParallelLinear with gather_output) class LoRAColumnParallelLinear : public nn::parallel::ColumnParallelLinear { public: + using Module::BuildShardedStateDict; static constexpr char kType[] = "LoRAColumnParallelLinear"; static constexpr char kParamLoraAName[] = "lora_A"; @@ -40,6 +43,9 @@ class LoRAColumnParallelLinear : public nn::parallel::ColumnParallelLinear { std::vector> Forward( const std::vector> &input_tensors) override; + ShardedStateDict BuildShardedStateDict( + const ShardedStateDictBuildContext &context) const override; + void MergeWeights(); void UnmergeWeights(); bool IsMerged() const; @@ -68,6 +74,7 @@ class LoRAColumnParallelLinear : public nn::parallel::ColumnParallelLinear { // [out_features, rank] - replicated (implemented as Linear) class LoRARowParallelLinear : public nn::parallel::RowParallelLinear { public: + using Module::BuildShardedStateDict; static constexpr char kType[] = "LoRARowParallelLinear"; static constexpr char kParamLoraAName[] = "lora_A"; static constexpr char kParamLoraBName[] = "lora_B"; @@ -86,6 +93,9 @@ class LoRARowParallelLinear : public nn::parallel::RowParallelLinear { std::vector> Forward( const std::vector> &input_tensors) override; + ShardedStateDict BuildShardedStateDict( + const ShardedStateDictBuildContext &context) const override; + void MergeWeights(); void UnmergeWeights(); bool IsMerged() const; diff --git a/infini_train/include/nn/modules/module.h b/infini_train/include/nn/modules/module.h index f98a8837..c37c8857 100644 --- a/infini_train/include/nn/modules/module.h +++ b/infini_train/include/nn/modules/module.h @@ -11,6 +11,7 @@ #include #include "infini_train/include/device.h" +#include "infini_train/include/shard_spec.h" namespace infini_train { class Tensor; @@ -24,6 +25,11 @@ enum class DataType : int8_t; namespace infini_train::nn { class Module; +struct ShardedStateDictBuildContext { + std::string logical_prefix; + std::string local_prefix; +}; + class NoSyncGuard { public: explicit NoSyncGuard(std::function exit_func); @@ -71,9 +77,9 @@ class Module : public std::enable_shared_from_this { // InfiniTrain's NamedParameters returns results ordered by full parameter // name. // TODO: Align with PyTorch's ordering in the future. - std::vector>> NamedParameters( - const std::string &prefix = "", bool recurse = true, - bool remove_duplicate = true) const; + virtual std::vector>> + NamedParameters(const std::string &prefix = "", bool recurse = true, + bool remove_duplicate = true) const; bool has_parameter(const std::string &name) const; std::shared_ptr *mutable_parameter(const std::string &name); const std::shared_ptr ¶meter(const std::string &name) const; @@ -84,11 +90,19 @@ class Module : public std::enable_shared_from_this { std::shared_ptr &mutable_module(const std::string &name); const Module &module(const std::string &name) const; - std::unordered_map> StateDict() const; + virtual std::unordered_map> StateDict() + const; + + // Return state-dict metadata with global shard coordinates. The public prefix + // is rooted in both the logical checkpoint namespace and this module's + // ordinary StateDict namespace. + ShardedStateDict BuildShardedStateDict(const std::string &prefix = "") const; + virtual ShardedStateDict BuildShardedStateDict( + const ShardedStateDictBuildContext &context) const; // Current behavior: missing keys / shape / dtype mismatches are FATAL errors; // unexpected keys in state_dict are WARNING-only and silently ignored. - void LoadStateDict( + virtual void LoadStateDict( const std::unordered_map> &state_dict); 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 662dbad2..fead2aad 100644 --- a/infini_train/include/nn/modules/transformer/causal_self_attention.h +++ b/infini_train/include/nn/modules/transformer/causal_self_attention.h @@ -7,6 +7,7 @@ #include "infini_train/include/nn/modules/module.h" #include "infini_train/include/nn/modules/transformer/transformer_config.h" +#include "infini_train/include/shard_spec.h" #include "infini_train/include/tensor.h" namespace infini_train::nn { @@ -14,6 +15,7 @@ namespace infini_train::nn { class CausalSelfAttention : public infini_train::nn::CloneableModule { public: + using Module::BuildShardedStateDict; static constexpr char kType[] = "CausalSelfAttention"; static constexpr char kCAttnLayerName[] = "c_attn"; static constexpr char kCProjLayerName[] = "c_proj"; @@ -28,6 +30,9 @@ class CausalSelfAttention std::vector> Forward( const std::vector> &x) override; + ShardedStateDict BuildShardedStateDict( + const ShardedStateDictBuildContext &context) const override; + private: TransformerConfig config_; int64_t n_head_ = 0; diff --git a/infini_train/include/nn/modules/transformer/mlp.h b/infini_train/include/nn/modules/transformer/mlp.h index 12586a1d..cdf011ff 100644 --- a/infini_train/include/nn/modules/transformer/mlp.h +++ b/infini_train/include/nn/modules/transformer/mlp.h @@ -4,6 +4,7 @@ #include #include "infini_train/include/nn/modules/module.h" +#include "infini_train/include/shard_spec.h" #include "infini_train/include/tensor.h" namespace infini_train::nn { @@ -11,6 +12,7 @@ struct TransformerConfig; class MLP : public infini_train::nn::CloneableModule { public: + using Module::BuildShardedStateDict; static constexpr char kType[] = "MLP"; static constexpr char kCFcLayerName[] = "c_fc"; static constexpr char kGeluLayerName[] = "gelu"; @@ -22,5 +24,8 @@ class MLP : public infini_train::nn::CloneableModule { std::vector> Forward( const std::vector> &x) override; + + ShardedStateDict BuildShardedStateDict( + const ShardedStateDictBuildContext &context) const override; }; } // namespace infini_train::nn diff --git a/infini_train/include/nn/modules/transformer/transformer.h b/infini_train/include/nn/modules/transformer/transformer.h index 1a24a6db..50b99ff6 100644 --- a/infini_train/include/nn/modules/transformer/transformer.h +++ b/infini_train/include/nn/modules/transformer/transformer.h @@ -6,6 +6,7 @@ #include "infini_train/include/nn/modules/module.h" #include "infini_train/include/nn/modules/transformer/transformer_config.h" #include "infini_train/include/nn/parallel/pp/pipeline_parallel.h" +#include "infini_train/include/shard_spec.h" #include "infini_train/include/tensor.h" namespace infini_train::nn { @@ -71,6 +72,7 @@ class TransformerLastStage : public CloneableModule { class TransformerModel : public CloneableModule { public: + using Module::BuildShardedStateDict; static constexpr char kType[] = "Transformer"; static constexpr char kTransformerModelName[] = "transformer"; @@ -81,6 +83,9 @@ class TransformerModel : public CloneableModule { const TransformerConfig &Config() const; + ShardedStateDict BuildShardedStateDict( + const ShardedStateDictBuildContext &context) const override; + private: const TransformerConfig config_; const infini_train::nn::parallel::StageInfo stage_info_; diff --git a/infini_train/include/nn/parallel/ddp/distributed_data_parallel.h b/infini_train/include/nn/parallel/ddp/distributed_data_parallel.h index 62df0378..ec3ff88f 100644 --- a/infini_train/include/nn/parallel/ddp/distributed_data_parallel.h +++ b/infini_train/include/nn/parallel/ddp/distributed_data_parallel.h @@ -7,6 +7,7 @@ #include "infini_train/include/nn/modules/module.h" #include "infini_train/include/nn/parallel/ddp/distributed_data_parallel_config.h" +#include "infini_train/include/shard_spec.h" namespace infini_train { class Tensor; @@ -29,6 +30,7 @@ constexpr char kModuleName[] = "module"; class DistributedDataParallel : public nn::Module { public: + using Module::BuildShardedStateDict; DistributedDataParallel(std::shared_ptr module, const Rank &rank, DistributedDataParallelConfig ddp_config); @@ -36,6 +38,8 @@ class DistributedDataParallel : public nn::Module { const std::vector> &input_tensors) override; std::shared_ptr module() const; + ShardedStateDict BuildShardedStateDict( + const ShardedStateDictBuildContext &context) const override; std::unique_ptr no_sync() override; diff --git a/infini_train/include/nn/parallel/pp/pipeline_parallel.h b/infini_train/include/nn/parallel/pp/pipeline_parallel.h index 6ea9d383..19a96a78 100644 --- a/infini_train/include/nn/parallel/pp/pipeline_parallel.h +++ b/infini_train/include/nn/parallel/pp/pipeline_parallel.h @@ -8,6 +8,7 @@ #include #include "infini_train/include/nn/modules/module.h" +#include "infini_train/include/shard_spec.h" namespace infini_train { class Tensor; @@ -33,6 +34,7 @@ struct StageInfo { class PipelineParallel : public Module { public: + using Module::BuildShardedStateDict; PipelineParallel(const std::shared_ptr module, int num_stages, int num_micro_batches, const std::vector> &recv_shape, @@ -49,6 +51,9 @@ class PipelineParallel : public Module { std::vector> *mutable_chunks(); + ShardedStateDict BuildShardedStateDict( + const ShardedStateDictBuildContext &context) const override; + private: void BuildPipelineStage(const std::vector> &recv_shape, Device device, diff --git a/infini_train/include/nn/parallel/tensor_parallel.h b/infini_train/include/nn/parallel/tensor_parallel.h index 6ed08d9c..2e431bd3 100644 --- a/infini_train/include/nn/parallel/tensor_parallel.h +++ b/infini_train/include/nn/parallel/tensor_parallel.h @@ -7,6 +7,7 @@ #include "infini_train/include/autograd/function.h" #include "infini_train/include/nn/modules/module.h" +#include "infini_train/include/shard_spec.h" namespace infini_train { class Tensor; @@ -38,6 +39,7 @@ std::vector> CopyToTPRegionFunc( class ColumnParallelLinear : public nn::CloneableModule { public: + using Module::BuildShardedStateDict; static constexpr char kType[] = "ColumnParallelLinear"; static constexpr char kParamWeightName[] = "weight"; @@ -57,6 +59,9 @@ class ColumnParallelLinear : public nn::CloneableModule { bool skip_bias_add() const; bool sequence_parallel() const; + ShardedStateDict BuildShardedStateDict( + const ShardedStateDictBuildContext &context) const override; + protected: bool bias_ = true; bool gather_output_ = false; // whether to return full local output tensor @@ -72,6 +77,7 @@ class ColumnParallelLinear : public nn::CloneableModule { class RowParallelLinear : public nn::CloneableModule { public: + using Module::BuildShardedStateDict; static constexpr char kType[] = "RowParallelLinear"; static constexpr char kParamWeightName[] = "weight"; @@ -91,6 +97,9 @@ class RowParallelLinear : public nn::CloneableModule { bool skip_bias_add() const; bool sequence_parallel() const; + ShardedStateDict BuildShardedStateDict( + const ShardedStateDictBuildContext &context) const override; + protected: bool bias_ = true; bool reduce_output_ = false; // whether to return full local output tensor @@ -107,6 +116,7 @@ class RowParallelLinear : public nn::CloneableModule { class VocabParallelEmbedding : public nn::CloneableModule { public: + using Module::BuildShardedStateDict; static constexpr char kType[] = "VocabParallelEmbedding"; static constexpr char kParamWeightName[] = "weight"; @@ -116,10 +126,15 @@ class VocabParallelEmbedding std::vector> Forward( const std::vector> &input_tensors) override; + ShardedStateDict BuildShardedStateDict( + const ShardedStateDictBuildContext &context) const override; + private: bool reduce_scatter_embeddings_ = false; // whether to perform ReduceScatter after embedding lookup + int64_t vocab_size_global_ = 0; + int64_t embedding_dim_ = 0; int64_t vocab_size_per_partition_ = 0; diff --git a/infini_train/include/optimizer.h b/infini_train/include/optimizer.h index 25ea7f1a..3d93766c 100644 --- a/infini_train/include/optimizer.h +++ b/infini_train/include/optimizer.h @@ -4,6 +4,7 @@ #include #include #include +#include #include #include #include @@ -35,6 +36,8 @@ class Optimizer { virtual std::unordered_map> StateDict() const; + NamedParameterList NamedParameters() const; + virtual void LoadStateDict( const std::unordered_map> &state_dict) {} @@ -58,6 +61,11 @@ class Optimizer { }; namespace optimizers { +inline constexpr std::string_view kAdamOptimizerPrefix = "adam."; +inline constexpr std::string_view kAdamFirstMomentPrefix = "adam.m."; +inline constexpr std::string_view kAdamSecondMomentPrefix = "adam.v."; +inline constexpr std::string_view kAdamStepKey = "adam.t"; + class SGD : public Optimizer { public: SGD(const std::vector> ¶ms, float learning_rate); diff --git a/infini_train/include/shard_spec.h b/infini_train/include/shard_spec.h new file mode 100644 index 00000000..ed4d8dec --- /dev/null +++ b/infini_train/include/shard_spec.h @@ -0,0 +1,61 @@ +#pragma once + +#include +#include +#include +#include +#include + +#include "infini_train/include/datatype.h" + +namespace infini_train { + +class Tensor; + +struct ShardSegment { + int64_t global_offset = 0; + int64_t local_offset = 0; + int64_t length = 0; + + bool operator==(const ShardSegment &other) const = default; +}; + +// Logical tensor shard metadata, aligned with Megatron-LM's ShardedTensor +// model. +struct ShardedTensor { + std::string key; + std::string local_key; + const Tensor *local_tensor = + nullptr; // Checkpoint-only identity; not serialized. + DataType dtype = DataType::kFLOAT32; + std::vector global_shape; + std::vector local_shape; + std::vector global_offset; + std::vector axis_fragmentations; + bool allow_shape_mismatch = false; + // Optional disjoint regions along the single fragmented axis. This is used + // by layouts such as rank-local [Q, K, V], which are not one contiguous + // slice of the logical global [Q, K, V] tensor. + std::vector segments; + + bool operator==(const ShardedTensor &other) const; +}; + +inline ShardedTensor MakeShardedTensor(std::string key, DataType dtype, + std::vector shape) { + const auto ndim = shape.size(); + return {.key = std::move(key), + .dtype = dtype, + .global_shape = shape, + .local_shape = std::move(shape), + .global_offset = std::vector(ndim, 0), + .axis_fragmentations = std::vector(ndim, 1)}; +} + +struct ShardedStateDict { + std::map tensors; + + void Merge(ShardedStateDict &&other); +}; + +} // namespace infini_train diff --git a/infini_train/src/checkpoint/checkpoint.cc b/infini_train/src/checkpoint/checkpoint.cc index 780c0e69..815ccd3c 100644 --- a/infini_train/src/checkpoint/checkpoint.cc +++ b/infini_train/src/checkpoint/checkpoint.cc @@ -1,25 +1,39 @@ #include "infini_train/include/checkpoint/checkpoint.h" +#include +#include #include #include +#include #include #include +#include #include #include +#include +#include #include #include #include #include "glog/logging.h" +#include "infini_train/include/checkpoint/constants.h" +#include "infini_train/include/checkpoint/load_planner.h" +#include "infini_train/include/checkpoint/load_strategy.h" +#include "infini_train/include/checkpoint/save_planner.h" +#include "infini_train/include/datatype.h" #include "infini_train/include/device.h" #include "infini_train/include/lr_scheduler.h" #include "infini_train/include/nn/modules/module.h" +#include "infini_train/include/nn/parallel/comm.h" +#include "infini_train/include/nn/parallel/global.h" +#include "infini_train/include/nn/parallel/reduce_op_type.h" +#include "infini_train/include/nn/parallel/work.h" #include "infini_train/include/optimizer.h" #include "infini_train/include/tensor.h" namespace infini_train { -enum class DataType : int8_t; namespace { constexpr uint32_t kCkptMagic = 0x54504B43; // CKPT @@ -194,6 +208,44 @@ T ExtractNumberField(const std::string &content, const std::string &key, } return value; } + +std::unordered_map> +MapCheckpointStateToLocal( + const ShardedStateDict &target, + const std::unordered_map> + &global_state) { + std::unordered_map> local_state; + for (const auto &[key, info] : target.tensors) { + CHECK(!info.local_key.empty()) + << "Empty checkpoint local_key for key=" << key; + const auto source = global_state.find(key); + CHECK(source != global_state.end()) + << "Missing resharded checkpoint tensor: key=" << key + << ", local_key=" << info.local_key; + CHECK(local_state.emplace(info.local_key, source->second).second) + << "Ambiguous checkpoint local_key: key=" << key + << ", local_key=" << info.local_key; + } + CHECK_EQ(global_state.size(), target.tensors.size()) + << "Resharded checkpoint result contains unexpected tensors"; + return local_state; +} + +void SynchronizeCheckpointRanks(const nn::Module &model) { + if (nn::parallel::global::GetWorldSize() == 1) { + return; + } + const auto parameters = model.Parameters(); + CHECK(!parameters.empty()) + << "Cannot synchronize checkpoint save for a model without parameters"; + auto token = + std::make_shared(std::vector{1}, DataType::kFLOAT32, + parameters.front()->GetDevice()); + token->Fill(1.0f); + nn::parallel::comm::AllReduce(token, nn::parallel::comm::ReduceOpType::kSum, + nullptr, true) + ->Synchronize(); +} } // namespace void Checkpoint::Save(const std::filesystem::path &checkpoint_dir, @@ -201,70 +253,108 @@ void Checkpoint::Save(const std::filesystem::path &checkpoint_dir, const TrainerState &state, const LRScheduler *lr_scheduler) { std::filesystem::create_directories(checkpoint_dir); - LOG(INFO) << "[CKPT] Save begin: dir=" << checkpoint_dir - << ", global_step=" << state.global_step; - - const auto model_path = checkpoint_dir / ("model.ckpt"); - - SaveStateDict(model_path, model.StateDict()); - - if (optimizer != nullptr) { - auto opt_state = optimizer->StateDict(); - if (!opt_state.empty()) { - const auto opt_path = checkpoint_dir / "optimizer.ckpt"; - SaveStateDict(opt_path, opt_state); + const int global_rank = nn::parallel::global::thread_global_rank; + const auto staging_root = checkpoint_dir / ".metadata_tmp"; + if (global_rank == 0) { + std::filesystem::remove_all(staging_root); + } + SynchronizeCheckpointRanks(model); + + int dp_rank = 0, tp_rank = 0, pp_rank = 0; + nn::parallel::global::GetCoordOf(global_rank, dp_rank, tp_rank, pp_rank); + // TODO(jym): Select checkpoint writers from each shard's replica_id instead + // of hard-coding DP rank 0. + if (dp_rank == 0) { + const auto rank_dir = + checkpoint_dir / std::format("rank_{:06d}", global_rank); + auto sharded_state = model.BuildShardedStateDict(); + std::unordered_map> optimizer_state; + if (optimizer != nullptr) { + optimizer_state = optimizer->StateDict(); + sharded_state.Merge(checkpoint::BuildOptimizerShardedStateDict( + sharded_state, optimizer_state, optimizer->NamedParameters())); } + const auto write_items = + checkpoint::SavePlanner::Plan(sharded_state, global_rank); + SaveLocalShard(rank_dir, sharded_state, write_items, model.StateDict(), + optimizer_state, global_rank); + + const auto staging_rank_dir = + staging_root / std::format("rank_{:06d}", global_rank); + std::filesystem::create_directories(staging_rank_dir); + std::filesystem::rename(rank_dir / checkpoint::kMetadataFilename, + staging_rank_dir / checkpoint::kMetadataFilename); } - if (lr_scheduler != nullptr) { - SaveLRSchedulerState(checkpoint_dir / "lr_scheduler.ckpt", - lr_scheduler->StateDict()); + SynchronizeCheckpointRanks(model); + if (global_rank == 0) { + SaveTrainerState(checkpoint_dir / checkpoint::kTrainerStateFilename, state); + if (lr_scheduler != nullptr) { + SaveLRSchedulerState(checkpoint_dir / checkpoint::kLRSchedulerFilename, + lr_scheduler->StateDict()); + } + auto metadata = LoadMetadata(staging_root); + CHECK(metadata.has_metadata); + SaveMetadataFile(checkpoint_dir / checkpoint::kTemporaryMetadataFilename, + metadata); + if (std::filesystem::exists(checkpoint_dir / + checkpoint::kMetadataFilename)) { + std::filesystem::remove(checkpoint_dir / checkpoint::kMetadataFilename); + } + std::filesystem::rename( + checkpoint_dir / checkpoint::kTemporaryMetadataFilename, + checkpoint_dir / checkpoint::kMetadataFilename); + std::filesystem::remove_all(staging_root); } - - SaveTrainerState(checkpoint_dir / "trainer_state.json", state); - LOG(ERROR) << "[CKPT] Save done: dir=" << checkpoint_dir; + SynchronizeCheckpointRanks(model); } void Checkpoint::Load(const std::filesystem::path &checkpoint_dir, nn::Module &model, Optimizer *optimizer, TrainerState &state, LRScheduler *lr_scheduler) { - const auto model_path = checkpoint_dir / "model.ckpt"; - LOG(INFO) << "[CKPT] Loading model: " << model_path; - - model.LoadStateDict(LoadStateDict(model_path)); + const auto metadata = LoadMetadata(checkpoint_dir); + CHECK(metadata.has_metadata); + // TODO(jym): Add compatibility loading for distributed checkpoints predating + // the current metadata version. + CHECK_EQ(metadata.version, checkpoint::kDistributedCheckpointVersion) + << "Unsupported distributed checkpoint version: " << metadata.version; + state = LoadTrainerState(checkpoint_dir / checkpoint::kTrainerStateFilename); + + auto model_sharded_state = model.BuildShardedStateDict(); + checkpoint::IndexedRegionLoadStrategy strategy; + auto result = strategy.Execute( + checkpoint_dir, + checkpoint::LoadPlanner::PlanReshard(metadata, model_sharded_state)); + model.LoadStateDict(MapCheckpointStateToLocal(model_sharded_state, result)); + + const int current_tp = nn::parallel::global::GetTensorParallelSize(); + const int current_pp = nn::parallel::global::GetPipelineParallelSize(); + state.tp_size = current_tp; + state.pp_size = current_pp; + state.ddp_size = nn::parallel::global::GetDataParallelSize(); + state.sp_size = + nn::parallel::global::GetSequenceParallelEnabled() ? current_tp : 1; + state.vpp_size = nn::parallel::global::GetVirtualPipelineParallelSize(); if (optimizer != nullptr) { - const auto opt_path = checkpoint_dir / "optimizer.ckpt"; - if (std::filesystem::exists(opt_path)) { - LOG(INFO) << "[CKPT] Loading optimizer: " << opt_path; - optimizer->LoadStateDict(LoadStateDict(opt_path)); - } else { - LOG(FATAL) << "Optimizer checkpoint not found at: " << opt_path; - } + auto optimizer_sharded_state = checkpoint::BuildOptimizerShardedStateDict( + model_sharded_state, optimizer->StateDict(), + optimizer->NamedParameters()); + auto optimizer_plan = + checkpoint::LoadPlanner::PlanReshard(metadata, optimizer_sharded_state); + auto optimizer_result = strategy.Execute(checkpoint_dir, optimizer_plan); + optimizer->LoadStateDict( + MapCheckpointStateToLocal(optimizer_sharded_state, optimizer_result)); } - - state = LoadTrainerState(checkpoint_dir / "trainer_state.json"); - - if (lr_scheduler != nullptr) { - const auto lr_scheduler_path = checkpoint_dir / "lr_scheduler.ckpt"; - if (std::filesystem::exists(lr_scheduler_path)) { - LOG(INFO) << "[CKPT] Loading LR scheduler: " << lr_scheduler_path; - lr_scheduler->LoadStateDict(LoadLRSchedulerState(lr_scheduler_path)); - } else { - LOG(WARNING) << "[CKPT] LR scheduler checkpoint not found at: " - << lr_scheduler_path - << ". Keeping the initialized scheduler state."; - } + if (lr_scheduler != nullptr && + std::filesystem::exists(checkpoint_dir / + checkpoint::kLRSchedulerFilename)) { + lr_scheduler->LoadStateDict(LoadLRSchedulerState( + checkpoint_dir / checkpoint::kLRSchedulerFilename)); } - - LOG(ERROR) << "[CKPT] Load done: global_step=" << state.global_step - << ", consumed_train_samples=" << state.consumed_train_samples - << ", topology(ddp,tp,sp,pp)=(" << state.ddp_size << "," - << state.tp_size << "," << state.sp_size << "," << state.pp_size - << ")"; } -void Checkpoint::SaveStateDict( +Checkpoint::SavedTensorLocations Checkpoint::SaveStateDict( const std::filesystem::path &path, const std::unordered_map> &state_dict) { @@ -273,6 +363,7 @@ void Checkpoint::SaveStateDict( uint32_t magic = kCkptMagic; uint32_t version = kCkptVersion; + SavedTensorLocations locations; uint32_t count = static_cast(state_dict.size()); ofs.write(reinterpret_cast(&magic), sizeof(magic)); ofs.write(reinterpret_cast(&version), sizeof(version)); @@ -294,9 +385,17 @@ void Checkpoint::SaveStateDict( Tensor cpu_tensor = tensor->To(Device()); uint64_t bytes = static_cast(cpu_tensor.SizeInBytes()); ofs.write(reinterpret_cast(&bytes), sizeof(bytes)); + const auto data_offset = ofs.tellp(); + CHECK(data_offset != std::streampos(-1)) + << "Failed to record tensor offset for " << name; + locations.emplace( + name, SavedTensorLocation{.data_offset = static_cast( + static_cast(data_offset)), + .byte_size = bytes}); ofs.write(reinterpret_cast(cpu_tensor.DataPtr()), static_cast(bytes)); } + return locations; } std::unordered_map> @@ -352,14 +451,16 @@ void Checkpoint::SaveTrainerState(const std::filesystem::path &path, ofs << " \"n_head\": " << state.n_head << ",\n"; ofs << " \"n_kv_head\": " << state.n_kv_head << ",\n"; ofs << " \"n_embd\": " << state.n_embd << ",\n"; - ofs << " \"vocab_size\": " << state.vocab_size << ",\n"; + ofs << " \"original_vocab_size\": " << state.original_vocab_size << ",\n"; + ofs << " \"padded_vocab_size\": " << state.padded_vocab_size << ",\n"; ofs << " \"global_step\": " << state.global_step << ",\n"; ofs << " \"consumed_train_samples\": " << state.consumed_train_samples << ",\n"; ofs << " \"ddp_size\": " << state.ddp_size << ",\n"; ofs << " \"tp_size\": " << state.tp_size << ",\n"; ofs << " \"sp_size\": " << state.sp_size << ",\n"; - ofs << " \"pp_size\": " << state.pp_size << "\n"; + ofs << " \"pp_size\": " << state.pp_size << ",\n"; + ofs << " \"vpp_size\": " << state.vpp_size << "\n"; ofs << "}\n"; } @@ -376,7 +477,12 @@ TrainerState Checkpoint::LoadTrainerState(const std::filesystem::path &path) { state.n_head = ExtractNumberField(content, "n_head", 0); state.n_kv_head = ExtractNumberField(content, "n_kv_head", 0); state.n_embd = ExtractNumberField(content, "n_embd", 0); - state.vocab_size = ExtractNumberField(content, "vocab_size", 0); + const auto legacy_vocab_size = + ExtractNumberField(content, "vocab_size", 0); + state.original_vocab_size = + ExtractNumberField(content, "original_vocab_size", 0); + state.padded_vocab_size = ExtractNumberField( + content, "padded_vocab_size", legacy_vocab_size); state.global_step = ExtractNumberField(content, "global_step", 0); state.consumed_train_samples = ExtractNumberField(content, "consumed_train_samples", 0); @@ -384,6 +490,462 @@ TrainerState Checkpoint::LoadTrainerState(const std::filesystem::path &path) { state.tp_size = ExtractNumberField(content, "tp_size", 1); state.sp_size = ExtractNumberField(content, "sp_size", 1); state.pp_size = ExtractNumberField(content, "pp_size", 1); + state.vpp_size = ExtractNumberField(content, "vpp_size", 1); return state; } + +void Checkpoint::SaveStateDictFile( + const std::filesystem::path &path, + const std::unordered_map> + &state_dict) { + SaveStateDict(path, state_dict); +} + +std::unordered_map> +Checkpoint::LoadStateDictFile(const std::filesystem::path &path) { + return LoadStateDict(path); +} + +static std::string DataTypeToString(DataType dt) { + const auto it = kDataTypeToDesc.find(dt); + CHECK(it != kDataTypeToDesc.end()) + << "Unsupported checkpoint tensor dtype: " << static_cast(dt); + return it->second; +} + +void Checkpoint::SaveLocalShard( + const std::filesystem::path &checkpoint_dir, + const ShardedStateDict &sharded_sd, + const std::vector &write_items, + const std::unordered_map> &state_dict, + const std::unordered_map> + &optimizer_state, + int global_rank) { + std::filesystem::create_directories(checkpoint_dir); + LOG(INFO) << "[CKPT] SaveLocalShard begin: dir=" << checkpoint_dir + << ", rank=" << global_rank; + + SavedTensorLocations model_file_index; + SavedTensorLocations optimizer_file_index; + + // Save model tensors separately from optimizer tensors. + { + std::unordered_map> filtered_sd; + for (const auto &[key, info] : sharded_sd.tensors) { + // Optimizer tensors are serialized separately. + if (key.starts_with(optimizers::kAdamOptimizerPrefix)) { + continue; + } + CHECK(!info.local_key.empty()) + << "Empty model checkpoint local_key: key=" << key; + auto it = state_dict.find(info.local_key); + CHECK(it != state_dict.end()) + << "Missing model checkpoint tensor: key=" << key + << ", local_key=" << info.local_key; + CHECK(filtered_sd.emplace(key, it->second).second) + << "Duplicate model checkpoint payload key: " << key; + } + if (!filtered_sd.empty()) { + model_file_index = SaveStateDict( + checkpoint_dir / checkpoint::kModelCheckpointFilename, filtered_sd); + } + } + + // Save optimizer tensors under stable global keys while reading the + // optimizer's ordinary local keys. + if (!optimizer_state.empty()) { + std::unordered_map> + global_optimizer_state; + for (const auto &[key, info] : sharded_sd.tensors) { + if (!key.starts_with(optimizers::kAdamOptimizerPrefix)) { + continue; + } + CHECK(!info.local_key.empty()) + << "Empty optimizer checkpoint local_key: key=" << key; + const auto local = optimizer_state.find(info.local_key); + CHECK(local != optimizer_state.end()) + << "Missing optimizer checkpoint tensor: key=" << key + << ", local_key=" << info.local_key; + CHECK(global_optimizer_state.emplace(key, local->second).second) + << "Duplicate optimizer checkpoint payload key: " << key; + } + CHECK_EQ(global_optimizer_state.size(), optimizer_state.size()) + << "Optimizer state contains tensors absent from checkpoint metadata"; + optimizer_file_index = + SaveStateDict(checkpoint_dir / checkpoint::kOptimizerCheckpointFilename, + global_optimizer_state); + } + + // Write the temporary rank manifest. + { + std::ofstream ofs(checkpoint_dir / checkpoint::kMetadataFilename); + CHECK(ofs.is_open()) << "Failed to open metadata.json: " + << checkpoint_dir / checkpoint::kMetadataFilename; + + ofs << "{\n"; + ofs << " \"version\": " << checkpoint::kDistributedCheckpointVersion + << ",\n"; + ofs << " \"format\": \"infinitrain_sharded\",\n"; + + ofs << " \"tensors\": [\n"; + + std::vector emitted_items; + for (const auto &item : write_items) { + if (sharded_sd.tensors.contains(item.key)) { + emitted_items.push_back(&item); + } + } + int dp_rank = 0, tp_rank = 0, pp_rank = 0; + nn::parallel::global::GetCoordOf(global_rank, dp_rank, tp_rank, pp_rank); + for (size_t i = 0; i < emitted_items.size(); ++i) { + const auto &item = *emitted_items[i]; + const auto it = sharded_sd.tensors.find(item.key); + const auto &file_index = + item.filename == checkpoint::kOptimizerCheckpointFilename + ? optimizer_file_index + : model_file_index; + const auto storage_it = file_index.find(item.key); + CHECK(storage_it != file_index.end()) + << "Missing stored tensor metadata for " << item.key; + const auto &storage = storage_it->second; + CHECK_EQ(storage.byte_size, item.byte_size); + + ofs << " {\n"; + ofs << " \"key\": \"" << item.key << "\",\n"; + ofs << " \"dtype\": \"" << DataTypeToString(item.dtype) << "\",\n"; + + // global_shape + ofs << " \"global_shape\": ["; + const auto &gs = it->second.global_shape; + for (size_t j = 0; j < gs.size(); ++j) { + ofs << gs[j] << (j + 1 < gs.size() ? ", " : ""); + } + ofs << "],\n"; + + ofs << " \"local_shape\": ["; + const auto &ls = it->second.local_shape; + for (size_t j = 0; j < ls.size(); ++j) { + ofs << ls[j] << (j + 1 < ls.size() ? ", " : ""); + } + ofs << "],\n"; + + ofs << " \"global_offset\": ["; + for (size_t j = 0; j < it->second.global_offset.size(); ++j) { + ofs << it->second.global_offset[j] + << (j + 1 < it->second.global_offset.size() ? ", " : ""); + } + ofs << "],\n"; + ofs << " \"axis_fragmentations\": ["; + for (size_t j = 0; j < it->second.axis_fragmentations.size(); ++j) { + ofs << it->second.axis_fragmentations[j] + << (j + 1 < it->second.axis_fragmentations.size() ? ", " : ""); + } + ofs << "],\n"; + auto write_segments = [&](const char *name, auto member) { + ofs << " \"" << name << "\": ["; + for (size_t j = 0; j < it->second.segments.size(); ++j) { + ofs << it->second.segments[j].*member + << (j + 1 < it->second.segments.size() ? ", " : ""); + } + ofs << "],\n"; + }; + write_segments("segment_global_offsets", &ShardSegment::global_offset); + write_segments("segment_local_offsets", &ShardSegment::local_offset); + write_segments("segment_lengths", &ShardSegment::length); + + ofs << " \"file\": \"" << item.filename << "\",\n"; + ofs << " \"offset\": " << storage.data_offset << ",\n"; + ofs << " \"byte_size\": " << item.byte_size << ",\n"; + ofs << " \"pp_rank\": " << pp_rank << ",\n"; + ofs << " \"stored_on_ranks\": [" << global_rank << "]\n"; + ofs << " }"; + if (i + 1 < emitted_items.size()) { + ofs << ","; + } + ofs << "\n"; + } + + ofs << " ]\n"; + ofs << "}\n"; + + LOG(INFO) << "[CKPT] metadata.json written"; + } + + LOG(ERROR) << "[CKPT] SaveLocalShard done: dir=" << checkpoint_dir; +} + +// Load one manifest or aggregate writer manifests while finalizing a +// checkpoint. +static std::string ExtractJsonString(const std::string &obj, + const std::string &key) { + auto token = std::string("\"") + key + "\""; + auto pos = obj.find(token); + if (pos == std::string::npos) { + return ""; + } + auto q1 = obj.find('"', pos + token.size()); + if (q1 == std::string::npos) { + return ""; + } + auto q2 = obj.find('"', q1 + 1); + if (q2 == std::string::npos) { + return ""; + } + return obj.substr(q1 + 1, q2 - q1 - 1); +} + +static std::string_view TrimWhitespace(std::string_view value) { + while (!value.empty() && + std::isspace(static_cast(value.front()))) { + value.remove_prefix(1); + } + while (!value.empty() && + std::isspace(static_cast(value.back()))) { + value.remove_suffix(1); + } + return value; +} + +template +std::vector ExtractIntegerArray(const std::string &obj, + const char *field_name) { + const auto field = std::string("\"") + field_name + "\""; + const auto field_pos = obj.find(field); + if (field_pos == std::string::npos) { + return {}; + } + + const auto array_start = obj.find('[', field_pos + field.size()); + const auto array_end = obj.find(']', array_start); + CHECK(array_start != std::string::npos && array_end != std::string::npos) + << "Invalid integer array in checkpoint metadata field " << field_name; + + std::string_view contents(obj.data() + array_start + 1, + array_end - array_start - 1); + contents = TrimWhitespace(contents); + if (contents.empty()) { + return {}; + } + + std::vector values; + while (!contents.empty()) { + const auto delimiter = contents.find(','); + const auto token = TrimWhitespace(contents.substr(0, delimiter)); + CHECK(!token.empty()) << "Empty integer in checkpoint metadata field " + << field_name; + + T value{}; + const auto [end, error] = + std::from_chars(token.data(), token.data() + token.size(), value); + CHECK(error == std::errc{} && end == token.data() + token.size()) + << "Invalid integer in checkpoint metadata field " << field_name << ": " + << token; + values.push_back(value); + + if (delimiter == std::string_view::npos) { + break; + } + contents.remove_prefix(delimiter + 1); + contents = TrimWhitespace(contents); + CHECK(!contents.empty()) + << "Trailing comma in checkpoint metadata field " << field_name; + } + return values; +} + +static Checkpoint::CheckpointMetadata LoadSingleMetadata( + const std::filesystem::path &checkpoint_dir) { + Checkpoint::CheckpointMetadata meta; + auto metadata_path = checkpoint_dir / checkpoint::kMetadataFilename; + if (!std::filesystem::exists(metadata_path)) { + meta.has_metadata = false; + return meta; + } + + std::ifstream ifs(metadata_path); + CHECK(ifs.is_open()) << "Failed to open metadata.json: " << metadata_path; + const std::string content((std::istreambuf_iterator(ifs)), + std::istreambuf_iterator()); + + meta.has_metadata = true; + meta.version = ExtractNumberField(content, "version", 0); + + // Locate the tensors array. + auto tensors_key = content.find("\"tensors\""); + if (tensors_key == std::string::npos) { + return meta; + } + + auto array_start = content.find('[', tensors_key); + if (array_start == std::string::npos) { + return meta; + } + + int depth = 1; + size_t pos = array_start + 1; + while (pos < content.size() && depth > 0) { + if (content[pos] == '[') { + ++depth; + } else if (content[pos] == ']') { + --depth; + } + ++pos; + } + std::string tensor_block = + content.substr(array_start + 1, pos - array_start - 2); + + // Parse each tensor object. + size_t obj_pos = 0; + while ((obj_pos = tensor_block.find('{', obj_pos)) != std::string::npos) { + int object_depth = 1; + size_t obj_end = obj_pos + 1; + while (obj_end < tensor_block.size() && object_depth > 0) { + if (tensor_block[obj_end] == '{') { + ++object_depth; + } + if (tensor_block[obj_end] == '}') { + --object_depth; + } + ++obj_end; + } + if (obj_end > 0) { + --obj_end; + } + if (obj_end == std::string::npos) { + break; + } + + std::string obj = tensor_block.substr(obj_pos, obj_end - obj_pos + 1); + + Checkpoint::CheckpointMetadata::TensorEntry entry; + entry.key = ExtractJsonString(obj, "key"); + entry.file = ExtractJsonString(obj, "file"); + entry.dtype_str = ExtractJsonString(obj, "dtype"); + entry.offset = ExtractNumberField(obj, "offset", 0); + entry.byte_size = ExtractNumberField(obj, "byte_size", 0); + entry.pp_rank = ExtractNumberField(obj, "pp_rank", 0); + + entry.global_shape = ExtractIntegerArray(obj, "global_shape"); + entry.local_shape = ExtractIntegerArray(obj, "local_shape"); + entry.global_offset = ExtractIntegerArray(obj, "global_offset"); + entry.axis_fragmentations = + ExtractIntegerArray(obj, "axis_fragmentations"); + + const auto segment_global_offsets = + ExtractIntegerArray(obj, "segment_global_offsets"); + const auto segment_local_offsets = + ExtractIntegerArray(obj, "segment_local_offsets"); + const auto segment_lengths = + ExtractIntegerArray(obj, "segment_lengths"); + CHECK_EQ(segment_global_offsets.size(), segment_local_offsets.size()); + CHECK_EQ(segment_global_offsets.size(), segment_lengths.size()); + for (size_t i = 0; i < segment_lengths.size(); ++i) { + entry.segments.push_back({.global_offset = segment_global_offsets[i], + .local_offset = segment_local_offsets[i], + .length = segment_lengths[i]}); + } + + entry.stored_on_ranks = ExtractIntegerArray(obj, "stored_on_ranks"); + + meta.tensors.push_back(std::move(entry)); + obj_pos = obj_end + 1; + } + + LOG(INFO) << "[CKPT] Loaded metadata.json: " << meta.tensors.size() + << " tensor shards"; + return meta; +} + +Checkpoint::CheckpointMetadata Checkpoint::LoadMetadata( + const std::filesystem::path &checkpoint_dir) { + if (std::filesystem::exists(checkpoint_dir / checkpoint::kMetadataFilename)) { + return LoadSingleMetadata(checkpoint_dir); + } + + CheckpointMetadata merged; + for (const auto &entry : + std::filesystem::directory_iterator(checkpoint_dir)) { + if (!entry.is_directory() || + !entry.path().filename().string().starts_with("rank_") || + !std::filesystem::exists(entry.path() / + checkpoint::kMetadataFilename)) { + continue; + } + auto rank_metadata = LoadSingleMetadata(entry.path()); + if (!rank_metadata.has_metadata) { + continue; + } + if (!merged.has_metadata) { + merged = rank_metadata; + merged.tensors.clear(); + } + for (auto &tensor : rank_metadata.tensors) { + tensor.file = (entry.path().filename() / tensor.file).generic_string(); + merged.tensors.push_back(std::move(tensor)); + } + } + LOG(INFO) << "[CKPT] Aggregated " << merged.tensors.size() + << " tensor shards from rank manifests"; + return merged; +} + +void Checkpoint::SaveMetadataFile(const std::filesystem::path &path, + const CheckpointMetadata &metadata) { + std::ofstream ofs(path); + CHECK(ofs.is_open()) << "Failed to write checkpoint metadata: " << path; + ofs << "{\n"; + ofs << " \"version\": " << checkpoint::kDistributedCheckpointVersion + << ",\n"; + ofs << " \"format\": \"infinitrain_sharded\",\n"; + + ofs << " \"tensors\": [\n"; + for (size_t i = 0; i < metadata.tensors.size(); ++i) { + const auto &tensor = metadata.tensors[i]; + ofs << " {\n"; + ofs << " \"key\": \"" << tensor.key << "\",\n"; + ofs << " \"dtype\": \"" << tensor.dtype_str << "\",\n"; + auto write_shape = [&](const char *name, + const std::vector &shape) { + ofs << " \"" << name << "\": ["; + for (size_t d = 0; d < shape.size(); ++d) { + ofs << shape[d] << (d + 1 < shape.size() ? ", " : ""); + } + ofs << "],\n"; + }; + write_shape("global_shape", tensor.global_shape); + write_shape("local_shape", tensor.local_shape); + write_shape("global_offset", tensor.global_offset); + ofs << " \"axis_fragmentations\": ["; + for (size_t d = 0; d < tensor.axis_fragmentations.size(); ++d) { + ofs << tensor.axis_fragmentations[d] + << (d + 1 < tensor.axis_fragmentations.size() ? ", " : ""); + } + ofs << "],\n"; + auto write_segments = [&](const char *name, auto member) { + ofs << " \"" << name << "\": ["; + for (size_t d = 0; d < tensor.segments.size(); ++d) { + ofs << tensor.segments[d].*member + << (d + 1 < tensor.segments.size() ? ", " : ""); + } + ofs << "],\n"; + }; + write_segments("segment_global_offsets", &ShardSegment::global_offset); + write_segments("segment_local_offsets", &ShardSegment::local_offset); + write_segments("segment_lengths", &ShardSegment::length); + ofs << " \"file\": \"" << tensor.file << "\",\n"; + ofs << " \"offset\": " << tensor.offset << ",\n"; + ofs << " \"byte_size\": " << tensor.byte_size << ",\n"; + ofs << " \"pp_rank\": " << tensor.pp_rank << ",\n"; + ofs << " \"stored_on_ranks\": ["; + for (size_t r = 0; r < tensor.stored_on_ranks.size(); ++r) { + ofs << tensor.stored_on_ranks[r] + << (r + 1 < tensor.stored_on_ranks.size() ? ", " : ""); + } + ofs << "]\n"; + ofs << " }" << (i + 1 < metadata.tensors.size() ? "," : "") << "\n"; + } + ofs << " ]\n"; + ofs << "}\n"; + CHECK(ofs.good()) << "Failed while writing checkpoint metadata: " << path; +} } // namespace infini_train diff --git a/infini_train/src/checkpoint/checkpoint_manager.cc b/infini_train/src/checkpoint/checkpoint_manager.cc index 677e1a09..9e02cb2e 100644 --- a/infini_train/src/checkpoint/checkpoint_manager.cc +++ b/infini_train/src/checkpoint/checkpoint_manager.cc @@ -4,24 +4,43 @@ #include #include #include +#include #include #include -#include #include #include #include "glog/logging.h" #include "infini_train/include/checkpoint/checkpoint.h" +#include "infini_train/include/checkpoint/constants.h" #include "infini_train/include/nn/modules/transformer/transformer_config.h" -#include "infini_train/include/nn/parallel/global.h" +#include "infini_train/include/nn/parallel/ddp/distributed_optimizer.h" #include "infini_train/include/nn/parallel/rank.h" +#include "infini_train/include/optimizer.h" using namespace infini_train; namespace nn = infini_train::nn; -// TODO(jym): ckpt is a new checkpoint format; bin is the legacy format. Keeping -// both as an interim solution; plan to consolidate into one later. +namespace { + +std::filesystem::path ResolveCheckpointDirectory( + const std::filesystem::path &root) { + const auto latest_path = root / checkpoint::kLatestIterationFilename; + if (!std::filesystem::exists(latest_path)) { + return root; + } + std::ifstream latest(latest_path); + int64_t iteration = 0; + latest >> iteration; + const auto directory = root / std::format("iter_{:07d}", iteration); + CHECK(std::filesystem::exists(directory)) + << "Latest checkpoint directory does not exist: " << directory; + return directory; +} + +} // namespace + ResumeFromCheckpointResult ResumeFromCheckpoint( const ResumeFromCheckpointArgs &args) { ResumeFromCheckpointResult result; @@ -30,26 +49,17 @@ ResumeFromCheckpointResult ResumeFromCheckpoint( "scratch."; return result; } - - int ddp_world_size = nn::parallel::global::GetDataParallelSize(); - int tp_world_size = nn::parallel::global::GetTensorParallelSize(); - int sp_world_size = - nn::parallel::global::GetSequenceParallelEnabled() ? tp_world_size : 1; - int pp_world_size = nn::parallel::global::GetPipelineParallelSize(); - - std::filesystem::path resume_dir = args.resume_root; - if (args.rank.IsParallel()) { - const auto rank_dir = - resume_dir / std::format("rank_{:06d}", args.rank.GlobalRank()); - if (std::filesystem::exists(rank_dir)) { - resume_dir = rank_dir; - } - } - - Checkpoint::Load(resume_dir, *args.model, args.optimizer.get(), args.state, - args.lr_scheduler.get()); - - result.global_step = static_cast(args.state.global_step); + // TODO(jym): Save and restore DistributedOptimizer/ZeRO optimizer state. + CHECK(dynamic_cast( + args.optimizer.get()) == nullptr) + << "Checkpoint restore does not support DistributedOptimizer/ZeRO " + "optimizer state; use zero_stage=0"; + + const auto checkpoint_dir = ResolveCheckpointDirectory(args.resume_root); + CHECK(std::filesystem::exists(checkpoint_dir / checkpoint::kMetadataFilename)) + << "Checkpoint metadata.json not found: " << checkpoint_dir; + Checkpoint::Load(checkpoint_dir, *args.model, args.optimizer.get(), + args.state, args.lr_scheduler.get()); CHECK_EQ(args.state.n_layer, args.model_config.n_layer) << "n_layer mismatch: ckpt=" << args.state.n_layer @@ -63,23 +73,23 @@ ResumeFromCheckpointResult ResumeFromCheckpoint( CHECK_EQ(args.state.n_embd, args.model_config.n_embd) << "n_embd mismatch: ckpt=" << args.state.n_embd << ", config=" << args.model_config.n_embd; - CHECK_EQ(args.state.vocab_size, args.model_config.vocab_size) - << "vocab_size mismatch: ckpt=" << args.state.vocab_size - << ", config=" << args.model_config.vocab_size; - - CHECK_EQ(args.state.ddp_size, ddp_world_size) - << "DDP size mismatch: checkpoint has DDP=" << args.state.ddp_size - << ", but current run has DDP=" << ddp_world_size; - CHECK_EQ(args.state.tp_size, tp_world_size) - << "TP size mismatch: checkpoint has TP=" << args.state.tp_size - << ", but current run has TP=" << tp_world_size; - CHECK_EQ(args.state.sp_size, sp_world_size) - << "SP size mismatch: checkpoint has SP=" << args.state.sp_size - << ", but current run has SP=" << sp_world_size; - CHECK_EQ(args.state.pp_size, pp_world_size) - << "PP size mismatch: checkpoint has PP=" << args.state.pp_size - << ", but current run has PP=" << pp_world_size; - + if (args.state.original_vocab_size > 0) { + CHECK_EQ(args.state.original_vocab_size, + args.model_config.original_vocab_size) + << "original_vocab_size mismatch: ckpt=" + << args.state.original_vocab_size + << ", config=" << args.model_config.original_vocab_size; + CHECK_GE(args.state.padded_vocab_size, args.state.original_vocab_size) + << "Checkpoint padded vocabulary cannot represent its logical " + "vocabulary"; + } else { + // Legacy trainer_state.json only stored the padded vocab size. + CHECK_GE(args.state.padded_vocab_size, + args.model_config.original_vocab_size) + << "Legacy checkpoint vocabulary cannot represent the configured " + "logical vocabulary"; + } + result.global_step = static_cast(args.state.global_step); result.consumed_train_samples = static_cast( std::max(args.state.consumed_train_samples, 0)); if (args.rank.IsMainRank()) { @@ -87,63 +97,84 @@ ResumeFromCheckpointResult ResumeFromCheckpoint( "Resume training from step {}, consumed_train_samples {}", args.state.global_step, args.state.consumed_train_samples); } - return result; } void SaveCheckpoint(const SaveCheckpointArgs &args) { - const auto ckpt_start = std::chrono::high_resolution_clock::now(); - - TrainerState state; - state.global_step = args.global_step; - state.consumed_train_samples = - static_cast(args.consumed_train_samples); - state.n_layer = args.n_layer; - state.n_head = args.n_head; - state.n_kv_head = args.n_kv_head; - state.n_embd = args.n_embd; - state.vocab_size = args.vocab_size; - state.ddp_size = args.ddp_size; - state.tp_size = args.tp_size; - state.sp_size = args.sp_size; - state.pp_size = args.pp_size; - - Checkpoint::Save(args.save_dir, args.model, args.optimizer, state, + // TODO(jym): Save and restore DistributedOptimizer/ZeRO optimizer state. + CHECK(dynamic_cast( + args.optimizer) == nullptr) + << "Checkpoint save does not support DistributedOptimizer/ZeRO optimizer " + "state; use zero_stage=0"; + CHECK_GT(args.original_vocab_size, 0); + CHECK_GE(args.padded_vocab_size, args.original_vocab_size); + const auto checkpoint_start = std::chrono::high_resolution_clock::now(); + // Snapshot training progress and the topology that produced this checkpoint. + TrainerState state{.global_step = args.global_step, + .consumed_train_samples = + static_cast(args.consumed_train_samples), + .n_layer = args.n_layer, + .n_head = args.n_head, + .n_kv_head = args.n_kv_head, + .n_embd = args.n_embd, + .original_vocab_size = args.original_vocab_size, + .padded_vocab_size = args.padded_vocab_size, + .ddp_size = args.ddp_size, + .tp_size = args.tp_size, + .sp_size = args.sp_size, + .pp_size = args.pp_size, + .vpp_size = args.vpp_size}; + const auto iteration_dir = + args.checkpoint_root_dir.empty() + ? args.save_dir + : args.checkpoint_root_dir / + std::format("iter_{:07d}", args.global_step); + Checkpoint::Save(iteration_dir, args.model, args.optimizer, state, args.lr_scheduler); - const auto ckpt_end = std::chrono::high_resolution_clock::now(); - const double ckpt_ms = - std::chrono::duration(ckpt_end - ckpt_start).count(); - - if (!args.rank.IsMainRank()) { - return; + if (args.rank.IsMainRank() && !args.checkpoint_root_dir.empty()) { + const auto latest = + args.checkpoint_root_dir / checkpoint::kLatestIterationFilename; + const auto temporary_latest = args.checkpoint_root_dir / + checkpoint::kTemporaryLatestIterationFilename; + { + std::ofstream output(temporary_latest); + CHECK(output.is_open()); + output << args.global_step; + } + if (std::filesystem::exists(latest)) { + std::filesystem::remove(latest); + } + std::filesystem::rename(temporary_latest, latest); } - LOG(INFO) << std::format("Checkpoint saved at: {} ({:.2f} ms)", - args.save_dir.string(), ckpt_ms); - - // FIXME(jym): Pruning currently relies on lexicographic sorting of directory - // names. This only works when step directories use zero-padded names (e.g. - // checkpoint_step_000042). If a future change introduces unpadded names, the - // prune order will be incorrect. Consider extracting the step number from the - // directory name and sorting numerically instead, once the checkpoint naming - // convention is finalized. - if (args.max_checkpoint_keep > 0 && + if (args.rank.IsMainRank() && args.max_checkpoint_keep > 0 && std::filesystem::exists(args.checkpoint_root_dir)) { - std::vector ckpts; + std::vector checkpoints; for (const auto &entry : std::filesystem::directory_iterator(args.checkpoint_root_dir)) { if (entry.is_directory() && - entry.path().filename().string().starts_with("checkpoint_step_")) { - ckpts.push_back(entry.path()); + entry.path().filename().string().starts_with("iter_")) { + checkpoints.push_back(entry.path()); } } - std::sort(ckpts.begin(), ckpts.end()); - while (ckpts.size() > args.max_checkpoint_keep) { - std::filesystem::remove_all(ckpts.front()); - ckpts.erase(ckpts.begin()); + // FIXME(jym): Pruning relies on lexicographic sorting of checkpoint + // directory names. This is only correct while iteration directories use + // zero-padded names (e.g. iter_0000042). If the naming convention changes + // to unpadded names, parse the iteration and sort numerically instead. + std::sort(checkpoints.begin(), checkpoints.end()); + while (checkpoints.size() > args.max_checkpoint_keep) { + std::filesystem::remove_all(checkpoints.front()); + checkpoints.erase(checkpoints.begin()); } } + + const auto checkpoint_end = std::chrono::high_resolution_clock::now(); + const double elapsed_ms = std::chrono::duration( + checkpoint_end - checkpoint_start) + .count(); + LOG(INFO) << std::format("Checkpoint saved at: {} ({:.2f} ms)", + iteration_dir.string(), elapsed_ms); } size_t DataLoaderBatchesToSkip(size_t consumed_train_samples, diff --git a/infini_train/src/checkpoint/load_planner.cc b/infini_train/src/checkpoint/load_planner.cc new file mode 100644 index 00000000..6bfbe958 --- /dev/null +++ b/infini_train/src/checkpoint/load_planner.cc @@ -0,0 +1,280 @@ +#include "infini_train/include/checkpoint/load_planner.h" + +#include + +#include +#include +#include +#include +#include + +#include "glog/logging.h" + +#include "infini_train/include/shard_spec.h" + +namespace infini_train::checkpoint { +namespace { + +DataType StringToDataType(const std::string &value) { + static const std::unordered_map legacy_names = { + {"bfloat16", DataType::kBFLOAT16}, + {"float16", DataType::kFLOAT16}, + {"float32", DataType::kFLOAT32}, + {"float64", DataType::kFLOAT64}, + }; + if (const auto it = legacy_names.find(value); it != legacy_names.end()) { + return it->second; + } + for (const auto &[dtype, description] : kDataTypeToDesc) { + if (description == value) { + return dtype; + } + } + LOG(FATAL) << "Unsupported checkpoint tensor dtype: " << value; + return DataType::kFLOAT32; +} + +// TODO(jym): Support checkpoint shards fragmented across multiple tensor axes. +int FragmentedAxis(const std::vector &axis_fragmentations) { + int fragmented_axis = -1; + for (size_t dim = 0; dim < axis_fragmentations.size(); ++dim) { + if (axis_fragmentations[dim] <= 1) { + continue; + } + CHECK_EQ(fragmented_axis, -1) + << "Multi-axis checkpoint sharding is not supported yet"; + fragmented_axis = static_cast(dim); + } + return fragmented_axis; +} + +int ResolveShardDim(const std::string &key, const ShardedTensor &target, + int saved_axis, + const std::vector &saved_global_shape) { + const int target_axis = FragmentedAxis(target.axis_fragmentations); + int shard_dim = target_axis >= 0 ? target_axis : saved_axis; + if (shard_dim < 0 && saved_global_shape != target.global_shape) { + shard_dim = 0; + } + if (saved_axis >= 0 && target_axis >= 0) { + CHECK_EQ(saved_axis, target_axis) + << "Shard dimension changed for tensor " << key; + } + return shard_dim; +} + +bool IsPaddingCompatible(bool allow_shape_mismatch, + const std::vector &source, + const std::vector &target) { + if (!allow_shape_mismatch || source.size() != target.size() || + source.empty()) { + return false; + } + // Dim 0 is the padded vocabulary axis and may grow or shrink across TP + // layouts. Every non-padding dimension must still describe the same tensor + // shape. + for (size_t dim = 1; dim < source.size(); ++dim) { + if (source[dim] != target[dim]) { + return false; + } + } + return true; +} + +void ValidateCoordinates(const std::string &key, + const std::vector &global_shape, + const std::vector &local_shape, + const std::vector &global_offset, + const std::vector &axis_fragmentations) { + CHECK_EQ(global_shape.size(), local_shape.size()) + << "Invalid local rank for tensor " << key; + CHECK_EQ(global_shape.size(), global_offset.size()) + << "Invalid offset rank for tensor " << key; + CHECK_EQ(global_shape.size(), axis_fragmentations.size()) + << "Invalid fragmentation rank for tensor " << key; + for (size_t dim = 0; dim < global_shape.size(); ++dim) { + CHECK_GT(global_shape[dim], 0) << "Invalid global shape for tensor " << key; + CHECK_GT(local_shape[dim], 0) << "Invalid local shape for tensor " << key; + CHECK_GE(global_offset[dim], 0) + << "Invalid global offset for tensor " << key; + CHECK_LE(global_offset[dim] + local_shape[dim], global_shape[dim]) + << "Shard exceeds global shape for tensor " << key; + CHECK_GE(axis_fragmentations[dim], 1) + << "Invalid axis fragmentation for tensor " << key; + } +} + +} // namespace + +LoadPlan LoadPlanner::PlanReshard( + const Checkpoint::CheckpointMetadata &metadata, + const ShardedStateDict &target_state_dict) { + LoadPlan plan; + for (const auto &[key, target] : target_state_dict.tensors) { + ValidateCoordinates(key, target.global_shape, target.local_shape, + target.global_offset, target.axis_fragmentations); + TargetTensorPlan tensor_plan{ + .key = key, + .dtype = target.dtype, + .global_shape = target.global_shape, + .target_shape = target.local_shape, + .shard_dim = FragmentedAxis(target.axis_fragmentations)}; + + std::vector candidates; + for (const auto &entry : metadata.tensors) { + if (entry.key == key) { + candidates.push_back(&entry); + } + } + CHECK(!candidates.empty()) + << "No saved shard found for target tensor: " << key; + const int saved_axis = + FragmentedAxis(candidates.front()->axis_fragmentations); + for (const auto *source : candidates) { + ValidateCoordinates(key, source->global_shape, source->local_shape, + source->global_offset, source->axis_fragmentations); + CHECK(source->global_shape == target.global_shape || + IsPaddingCompatible(target.allow_shape_mismatch, + source->global_shape, target.global_shape)) + << "Global shape changed for tensor " << key; + CHECK_EQ(FragmentedAxis(source->axis_fragmentations), saved_axis) + << "Inconsistent saved shard dimensions for tensor " << key; + } + + if (!target.segments.empty()) { + if (tensor_plan.shard_dim < 0) { + tensor_plan.shard_dim = 0; + } + int64_t target_covered = 0; + for (const auto &target_segment : target.segments) { + CHECK_EQ(target_segment.local_offset, target_covered) + << "Gap or overlap in target segments for " << key; + CHECK_GT(target_segment.length, 0); + CHECK_LE(target_segment.global_offset + target_segment.length, + target.global_shape[tensor_plan.shard_dim]); + target_covered += target_segment.length; + for (const auto *source : candidates) { + CHECK(!source->segments.empty()) + << "Saved checkpoint lacks segmented layout metadata for " << key; + for (const auto &source_segment : source->segments) { + const int64_t overlap_start = std::max( + source_segment.global_offset, target_segment.global_offset); + const int64_t overlap_end = + std::min(source_segment.global_offset + source_segment.length, + target_segment.global_offset + target_segment.length); + if (overlap_start >= overlap_end) { + continue; + } + tensor_plan.reads.push_back({ + .key = key, + .filename = source->file, + .dtype = StringToDataType(source->dtype_str), + .global_shape = source->global_shape, + .byte_size = source->byte_size, + .data_offset = source->offset, + .shard_dim = tensor_plan.shard_dim, + .source_offset = source_segment.local_offset + overlap_start - + source_segment.global_offset, + .target_offset = target_segment.local_offset + overlap_start - + target_segment.global_offset, + .length = overlap_end - overlap_start, + .source_shape = source->local_shape, + }); + } + } + } + CHECK_EQ(target_covered, target.local_shape[tensor_plan.shard_dim]) + << "Segmented layout does not cover target local tensor " << key; + std::sort(tensor_plan.reads.begin(), tensor_plan.reads.end(), + [](const ReadItem &left, const ReadItem &right) { + return left.target_offset < right.target_offset; + }); + int64_t covered = 0; + for (const auto &read : tensor_plan.reads) { + CHECK_EQ(read.target_offset, covered) + << "Gap or overlap in segmented target plan for " << key; + covered += read.length; + } + CHECK_EQ(covered, target_covered) + << "Incomplete segmented target plan for " << key; + plan.tensors.emplace(key, std::move(tensor_plan)); + continue; + } + + tensor_plan.shard_dim = ResolveShardDim(key, target, saved_axis, + candidates.front()->global_shape); + + if (tensor_plan.shard_dim < 0) { + const auto *source = candidates.front(); + tensor_plan.reads.push_back({.key = key, + .filename = source->file, + .dtype = StringToDataType(source->dtype_str), + .global_shape = source->global_shape, + .byte_size = source->byte_size, + .data_offset = source->offset, + .shard_dim = -1, + .source_shape = source->local_shape}); + plan.tensors.emplace(key, std::move(tensor_plan)); + continue; + } + + const int dim = tensor_plan.shard_dim; + const int64_t target_start = target.global_offset[dim]; + const int64_t target_length = target.local_shape[dim]; + const int64_t target_end = target_start + target_length; + std::set> seen_source_ranges; + + for (const auto *source : candidates) { + const int64_t saved_start = source->global_offset[dim]; + const int64_t saved_length = source->local_shape[dim]; + if (!seen_source_ranges.emplace(saved_start, saved_length).second) { + continue; + } + const int64_t overlap_start = std::max(saved_start, target_start); + const int64_t overlap_end = + std::min(saved_start + saved_length, target_end); + if (overlap_start >= overlap_end) { + continue; + } + + tensor_plan.reads.push_back( + {.key = key, + .filename = source->file, + .dtype = StringToDataType(source->dtype_str), + .global_shape = source->global_shape, + .byte_size = source->byte_size, + .data_offset = source->offset, + .shard_dim = dim, + .source_offset = overlap_start - saved_start, + .target_offset = overlap_start - target_start, + .length = overlap_end - overlap_start, + .source_shape = source->local_shape}); + } + + std::sort(tensor_plan.reads.begin(), tensor_plan.reads.end(), + [](const ReadItem &left, const ReadItem &right) { + return left.target_offset < right.target_offset; + }); + int64_t covered = 0; + for (const auto &read : tensor_plan.reads) { + CHECK_EQ(read.target_offset, covered) + << "Gap or overlap in target shard plan for " << key; + covered += read.length; + } + if (covered < target_length) { + CHECK(target.allow_shape_mismatch) + << "Incomplete target shard plan for " << key; + CHECK_EQ(dim, 0) << "Vocabulary padding is only supported along dim 0"; + CHECK_EQ(target_start + covered, candidates.front()->global_shape[0]) + << "Only trailing vocabulary padding is supported for " << key; + tensor_plan.trailing_zero_fill = target_length - covered; + covered = target_length; + } + CHECK_EQ(covered, target_length) + << "Incomplete target shard plan for " << key; + plan.tensors.emplace(key, std::move(tensor_plan)); + } + return plan; +} + +} // namespace infini_train::checkpoint diff --git a/infini_train/src/checkpoint/load_strategy.cc b/infini_train/src/checkpoint/load_strategy.cc new file mode 100644 index 00000000..d4dee556 --- /dev/null +++ b/infini_train/src/checkpoint/load_strategy.cc @@ -0,0 +1,167 @@ +#include "infini_train/include/checkpoint/load_strategy.h" + +#include +#include + +#include +#include +#include +#include + +#include "glog/logging.h" + +#include "infini_train/include/checkpoint/load_planner.h" +#include "infini_train/include/device.h" +#include "infini_train/include/nn/functional.h" +#include "infini_train/include/tensor.h" +#include "infini_train/include/utils/string_utils.h" + +namespace infini_train::checkpoint { +namespace { + +using FileCache = + std::unordered_map>; + +std::ifstream &GetFile(FileCache &cache, + const std::filesystem::path &checkpoint_dir, + const std::string &filename) { + auto it = cache.find(filename); + if (it == cache.end()) { + auto stream = std::make_unique(checkpoint_dir / filename, + std::ios::binary); + CHECK(stream->is_open()) + << "Failed to open checkpoint file: " << checkpoint_dir / filename; + it = cache.emplace(filename, std::move(stream)).first; + } + return *it->second; +} + +void ReadAt(std::ifstream &stream, uint64_t offset, void *destination, + uint64_t byte_size, const std::string &key, + const std::string &filename) { + stream.clear(); + stream.seekg(static_cast(offset), std::ios::beg); + CHECK(stream.good()) << "Failed to seek tensor " << key << " in " << filename; + stream.read(static_cast(destination), + static_cast(byte_size)); + CHECK_EQ(static_cast(stream.gcount()), byte_size) + << "Truncated tensor " << key << " in " << filename; +} + +std::shared_ptr ReadTensor(std::ifstream &stream, + const ReadItem &read) { + auto tensor = + std::make_shared(read.source_shape, read.dtype, Device()); + CHECK_EQ(read.byte_size, tensor->SizeInBytes()) + << "Tensor byte size mismatch for " << read.key << " in " + << read.filename; + ReadAt(stream, read.data_offset, tensor->DataPtr(), read.byte_size, read.key, + read.filename); + return tensor; +} + +std::shared_ptr ReadTensorRegion(std::ifstream &stream, + const ReadItem &read) { + CHECK_GE(read.shard_dim, 0); + CHECK_LT(read.shard_dim, static_cast(read.source_shape.size())); + CHECK_GE(read.source_offset, 0); + CHECK_GT(read.length, 0); + CHECK_LE(read.source_offset + read.length, read.source_shape[read.shard_dim]); + + uint64_t source_numel = 1; + for (const auto size : read.source_shape) { + CHECK_GT(size, 0); + source_numel *= static_cast(size); + } + CHECK_EQ(read.byte_size % source_numel, 0u) + << "Invalid tensor byte size for " << read.key << " in " << read.filename; + const uint64_t element_size = read.byte_size / source_numel; + + uint64_t outer = 1; + for (int dim = 0; dim < read.shard_dim; ++dim) { + outer *= static_cast(read.source_shape[dim]); + } + uint64_t inner = 1; + for (size_t dim = static_cast(read.shard_dim + 1); + dim < read.source_shape.size(); ++dim) { + inner *= static_cast(read.source_shape[dim]); + } + + auto region_shape = read.source_shape; + region_shape[read.shard_dim] = read.length; + auto tensor = std::make_shared(region_shape, read.dtype, Device()); + const uint64_t block_bytes = + static_cast(read.length) * inner * element_size; + const uint64_t source_stride_bytes = + static_cast(read.source_shape[read.shard_dim]) * inner * + element_size; + const uint64_t first_block_offset = + read.data_offset + + static_cast(read.source_offset) * inner * element_size; + CHECK_EQ(outer * block_bytes, tensor->SizeInBytes()); + + auto *destination = static_cast(tensor->DataPtr()); + for (uint64_t block = 0; block < outer; ++block) { + ReadAt(stream, first_block_offset + block * source_stride_bytes, + destination + block * block_bytes, block_bytes, read.key, + read.filename); + } + return tensor; +} + +} // namespace + +LoadedStateDict IndexedRegionLoadStrategy::Execute( + const std::filesystem::path &checkpoint_dir, const LoadPlan &plan) { + FileCache file_cache; + LoadedStateDict result; + + for (const auto &[key, tensor_plan] : plan.tensors) { + CHECK(!tensor_plan.reads.empty() || + (tensor_plan.shard_dim >= 0 && + tensor_plan.trailing_zero_fill == + tensor_plan.target_shape[tensor_plan.shard_dim])) + << "No reads or padding planned for target tensor: " << key; + std::vector> pieces; + pieces.reserve(tensor_plan.reads.size() + + (tensor_plan.trailing_zero_fill > 0 ? 1 : 0)); + + for (const auto &read : tensor_plan.reads) { + CHECK_GT(read.data_offset, 0) + << "Checkpoint metadata lacks a valid tensor data offset for " << key + << "; regenerate the checkpoint with the current format"; + auto &stream = GetFile(file_cache, checkpoint_dir, read.filename); + auto piece = read.shard_dim < 0 ? ReadTensor(stream, read) + : ReadTensorRegion(stream, read); + if (piece->Dtype() != tensor_plan.dtype) { + piece = std::make_shared(piece->To(tensor_plan.dtype)); + } + pieces.push_back(std::move(piece)); + } + + if (tensor_plan.trailing_zero_fill > 0) { + CHECK_GE(tensor_plan.shard_dim, 0); + auto padding_shape = tensor_plan.target_shape; + padding_shape[tensor_plan.shard_dim] = tensor_plan.trailing_zero_fill; + auto padding = + std::make_shared(padding_shape, tensor_plan.dtype, Device()); + padding->Fill(0.0f); + pieces.push_back(std::move(padding)); + } + + auto target = pieces.front(); + if (pieces.size() > 1) { + CHECK_GE(tensor_plan.shard_dim, 0); + target = + nn::function::Concat(pieces, tensor_plan.shard_dim)->Contiguous(); + } + CHECK(target->Dims() == tensor_plan.target_shape) + << "Target shard shape mismatch for " << key + << ": expected=" << utils::DimsToString(tensor_plan.target_shape) + << ", got=" << utils::DimsToString(target->Dims()); + result.emplace(key, std::move(target)); + } + return result; +} + +} // namespace infini_train::checkpoint diff --git a/infini_train/src/checkpoint/save_planner.cc b/infini_train/src/checkpoint/save_planner.cc new file mode 100644 index 00000000..3e2c755f --- /dev/null +++ b/infini_train/src/checkpoint/save_planner.cc @@ -0,0 +1,103 @@ +#include "infini_train/include/checkpoint/save_planner.h" + +#include +#include +#include +#include + +#include "glog/logging.h" + +#include "infini_train/include/checkpoint/constants.h" +#include "infini_train/include/optimizer.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::checkpoint { + +ShardedStateDict BuildOptimizerShardedStateDict( + const ShardedStateDict &model_state, + const std::unordered_map> + &optimizer_state, + const NamedParameterList &optimizer_parameters) { + std::map global_key_by_tensor; + for (const auto &[key, info] : model_state.tensors) { + if (info.local_tensor != nullptr && + !global_key_by_tensor.contains(info.local_tensor)) { + global_key_by_tensor.emplace(info.local_tensor, key); + } + } + std::unordered_map tensor_by_optimizer_name; + for (const auto &[name, parameter] : optimizer_parameters) { + CHECK(parameter != nullptr) << "Null optimizer parameter: " << name; + CHECK(tensor_by_optimizer_name.emplace(name, parameter.get()).second) + << "Duplicate optimizer parameter name: " << name; + } + + ShardedStateDict result; + for (const auto &[local_key, tensor] : optimizer_state) { + if (local_key == optimizers::kAdamStepKey) { + auto info = MakeShardedTensor(local_key, tensor->Dtype(), tensor->Dims()); + info.local_key = local_key; + result.tensors.emplace(local_key, std::move(info)); + continue; + } + + std::string_view state_prefix; + if (local_key.starts_with(optimizers::kAdamFirstMomentPrefix)) { + state_prefix = optimizers::kAdamFirstMomentPrefix; + } else if (local_key.starts_with(optimizers::kAdamSecondMomentPrefix)) { + state_prefix = optimizers::kAdamSecondMomentPrefix; + } else { + CHECK(false) << "Unsupported optimizer state key: " << local_key; + } + const auto local_parameter_name = local_key.substr(state_prefix.size()); + std::string global_parameter_key; + if (!optimizer_parameters.empty()) { + CHECK(tensor_by_optimizer_name.contains(local_parameter_name)) + << "Optimizer state has no parameter identity mapping: " << local_key; + const auto *parameter = tensor_by_optimizer_name.at(local_parameter_name); + CHECK(global_key_by_tensor.contains(parameter)) + << "Optimizer parameter is absent from model sharded state: " + << local_parameter_name; + global_parameter_key = global_key_by_tensor.at(parameter); + } else { + // Compatibility for metadata-only planner callers. Production save/load + // always supplies identities. + CHECK(model_state.tensors.contains(local_parameter_name)) + << "Optimizer state has no matching model metadata: " << local_key; + global_parameter_key = local_parameter_name; + } + const auto global_key = std::string(state_prefix) + global_parameter_key; + + auto info = model_state.tensors.at(global_parameter_key); + info.key = global_key; + info.local_key = local_key; + info.local_tensor = nullptr; + info.dtype = tensor->Dtype(); + CHECK(result.tensors.emplace(global_key, std::move(info)).second) + << "Duplicate global optimizer state key: " << global_key; + } + return result; +} + +std::vector SavePlanner::Plan(const ShardedStateDict &sd, int rank) { + std::vector items; + for (auto &[key, info] : sd.tensors) { + bool is_optimizer = key.starts_with(optimizers::kAdamOptimizerPrefix); + WriteItem item; + item.key = key; + item.filename = + is_optimizer ? kOptimizerCheckpointFilename : kModelCheckpointFilename; + item.byte_size = TensorByteSize(info.dtype, info.local_shape); + item.dtype = info.dtype; + item.local_shape = info.local_shape; + item.global_offset = info.global_offset; + item.axis_fragmentations = info.axis_fragmentations; + item.rank = rank; + + items.push_back(std::move(item)); + } + + return items; +} + +} // namespace infini_train::checkpoint diff --git a/infini_train/src/nn/lora/lora_parallel_linear.cc b/infini_train/src/nn/lora/lora_parallel_linear.cc index 48c6ba11..8c8eedd1 100644 --- a/infini_train/src/nn/lora/lora_parallel_linear.cc +++ b/infini_train/src/nn/lora/lora_parallel_linear.cc @@ -1,9 +1,11 @@ #include "infini_train/include/nn/lora/lora_parallel_linear.h" +#include #include #include #include #include +#include #include #include "glog/logging.h" @@ -259,6 +261,40 @@ std::vector> LoRAColumnParallelLinear::LoRAParameters() return {parameters_.at(kParamLoraAName), parameters_.at(kParamLoraBName)}; } +ShardedStateDict LoRAColumnParallelLinear::BuildShardedStateDict( + const ShardedStateDictBuildContext &context) const { + auto state = parallel::ColumnParallelLinear::BuildShardedStateDict(context); + const int tp_size = parallel::global::GetTensorParallelSize(); + + const auto &lora_a = parameter(kParamLoraAName); + const auto a_key = context.logical_prefix.empty() + ? kParamLoraAName + : context.logical_prefix + "." + kParamLoraAName; + auto a = MakeShardedTensor(a_key, lora_a->Dtype(), lora_a->Dims()); + a.local_key = context.local_prefix.empty() + ? kParamLoraAName + : context.local_prefix + "." + kParamLoraAName; + a.local_tensor = lora_a.get(); + state.tensors.emplace(a.key, std::move(a)); + + const auto &lora_b = parameter(kParamLoraBName); + ShardedTensor b; + b.key = context.logical_prefix.empty() + ? kParamLoraBName + : context.logical_prefix + "." + kParamLoraBName; + b.local_key = context.local_prefix.empty() + ? kParamLoraBName + : context.local_prefix + "." + kParamLoraBName; + b.local_tensor = lora_b.get(); + b.dtype = lora_b->Dtype(); + b.global_shape = {lora_b->Dims()[0] * tp_size, lora_b->Dims()[1]}; + b.local_shape = lora_b->Dims(); + b.global_offset = {lora_b->Dims()[0] * parallel::tp_rank, 0}; + b.axis_fragmentations = {tp_size, 1}; + state.tensors.emplace(b.key, std::move(b)); + return state; +} + bool LoRAColumnParallelLinear::IsMerged() const { return merged_; } int64_t LoRAColumnParallelLinear::in_features() const { return in_features_; } @@ -503,6 +539,40 @@ std::vector> LoRARowParallelLinear::LoRAParameters() return {parameters_.at(kParamLoraAName), parameters_.at(kParamLoraBName)}; } +ShardedStateDict LoRARowParallelLinear::BuildShardedStateDict( + const ShardedStateDictBuildContext &context) const { + auto state = parallel::RowParallelLinear::BuildShardedStateDict(context); + const int tp_size = parallel::global::GetTensorParallelSize(); + + const auto &lora_a = parameter(kParamLoraAName); + ShardedTensor a; + a.key = context.logical_prefix.empty() + ? kParamLoraAName + : context.logical_prefix + "." + kParamLoraAName; + a.local_key = context.local_prefix.empty() + ? kParamLoraAName + : context.local_prefix + "." + kParamLoraAName; + a.local_tensor = lora_a.get(); + a.dtype = lora_a->Dtype(); + a.global_shape = {lora_a->Dims()[0], lora_a->Dims()[1] * tp_size}; + a.local_shape = lora_a->Dims(); + a.global_offset = {0, lora_a->Dims()[1] * parallel::tp_rank}; + a.axis_fragmentations = {1, tp_size}; + state.tensors.emplace(a.key, std::move(a)); + + const auto &lora_b = parameter(kParamLoraBName); + const auto b_key = context.logical_prefix.empty() + ? kParamLoraBName + : context.logical_prefix + "." + kParamLoraBName; + auto b = MakeShardedTensor(b_key, lora_b->Dtype(), lora_b->Dims()); + b.local_key = context.local_prefix.empty() + ? kParamLoraBName + : context.local_prefix + "." + kParamLoraBName; + b.local_tensor = lora_b.get(); + state.tensors.emplace(b.key, std::move(b)); + return state; +} + bool LoRARowParallelLinear::IsMerged() const { return merged_; } int64_t LoRARowParallelLinear::in_features() const { return in_features_; } diff --git a/infini_train/src/nn/modules/module.cc b/infini_train/src/nn/modules/module.cc index 65022bb2..46bb7d8b 100644 --- a/infini_train/src/nn/modules/module.cc +++ b/infini_train/src/nn/modules/module.cc @@ -3,6 +3,7 @@ #include #include #include +#include #include #include #include @@ -225,6 +226,48 @@ std::unordered_map> Module::StateDict() return state; } +namespace { +std::string JoinStateDictKey(const std::string &prefix, + const std::string &name) { + return prefix.empty() ? name : prefix + "." + name; +} +} // namespace + +ShardedStateDict Module::BuildShardedStateDict( + const std::string &prefix) const { + return BuildShardedStateDict( + {.logical_prefix = prefix, .local_prefix = prefix}); +} + +ShardedStateDict Module::BuildShardedStateDict( + const ShardedStateDictBuildContext &context) const { + ShardedStateDict sd; + for (auto &[name, param] : parameters_) { + auto key = JoinStateDictKey(context.logical_prefix, name); + auto info = MakeShardedTensor(key, param->Dtype(), param->Dims()); + info.local_key = JoinStateDictKey(context.local_prefix, name); + info.local_tensor = param.get(); + sd.tensors.emplace(key, std::move(info)); + } + for (auto &[name, buffer] : buffers_) { + auto key = JoinStateDictKey(context.logical_prefix, name); + auto info = MakeShardedTensor(key, buffer->Dtype(), buffer->Dims()); + info.local_key = JoinStateDictKey(context.local_prefix, name); + info.local_tensor = buffer.get(); + sd.tensors.emplace(key, std::move(info)); + } + for (auto &[name, module] : modules_) { + if (name.starts_with("__pp")) { + continue; + } + auto child_sd = module->BuildShardedStateDict( + {.logical_prefix = JoinStateDictKey(context.logical_prefix, name), + .local_prefix = JoinStateDictKey(context.local_prefix, name)}); + sd.Merge(std::move(child_sd)); + } + return sd; +} + void Module::LoadStateDict( const std::unordered_map> &state_dict) { 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 11162086..68885839 100644 --- a/infini_train/src/nn/modules/transformer/causal_self_attention.cc +++ b/infini_train/src/nn/modules/transformer/causal_self_attention.cc @@ -2,16 +2,19 @@ #include #include +#include #include #include #include #include #include +#include #include #include "glog/logging.h" #include "infini_train/include/nn/functional.h" +#include "infini_train/include/nn/lora/lora_parallel_linear.h" #include "infini_train/include/nn/modules/normalization.h" #include "infini_train/include/nn/modules/transformer/transformer_config.h" #include "infini_train/include/nn/modules/transformer/utils.h" @@ -104,6 +107,50 @@ void CausalSelfAttention::SetupAttention(const TransformerConfig &config) { } } +ShardedStateDict CausalSelfAttention::BuildShardedStateDict( + const ShardedStateDictBuildContext &context) const { + auto state = Module::BuildShardedStateDict(context); + const int tp_size = parallel::global::GetTensorParallelSize(); + const int rank = parallel::tp_rank; + const int64_t q_global = n_head_ * head_dim_; + const int64_t kv_global = n_kv_head_ * head_dim_; + const int64_t q_local = q_global / tp_size; + const int64_t kv_local = kv_global / tp_size; + + // FIXME(jym): Transformer should not depend on LoRA just to identify lora_B. + // Move the packed-QKV segment layout into a dedicated sharding abstraction + // shared by the base attention weight and LoRA parameters. + const auto c_attn_prefix = + context.logical_prefix.empty() + ? kCAttnLayerName + : context.logical_prefix + "." + kCAttnLayerName; + auto set_qkv_segments = [&](const std::string ¶meter_name) { + const auto key = c_attn_prefix + "." + parameter_name; + auto it = state.tensors.find(key); + if (it == state.tensors.end()) { + return; + } + auto &tensor = it->second; + tensor.global_offset.assign(tensor.global_shape.size(), 0); + tensor.segments = { + {.global_offset = rank * q_local, .local_offset = 0, .length = q_local}, + {.global_offset = q_global + rank * kv_local, + .local_offset = q_local, + .length = kv_local}, + {.global_offset = q_global + kv_global + rank * kv_local, + .local_offset = q_local + kv_local, + .length = kv_local}, + }; + }; + + set_qkv_segments(parallel::ColumnParallelLinear::kParamWeightName); + if (config_.add_bias_linear) { + set_qkv_segments(parallel::ColumnParallelLinear::kParamBiasName); + } + set_qkv_segments(lora::LoRAColumnParallelLinear::kParamLoraBName); + return state; +} + std::shared_ptr CausalSelfAttention::RepeatKV( const std::shared_ptr &x, int64_t n_rep) { const auto &shape = x->Dims(); diff --git a/infini_train/src/nn/modules/transformer/mlp.cc b/infini_train/src/nn/modules/transformer/mlp.cc index 29f3ffb2..428fe567 100644 --- a/infini_train/src/nn/modules/transformer/mlp.cc +++ b/infini_train/src/nn/modules/transformer/mlp.cc @@ -2,14 +2,18 @@ #include +#include #include #include +#include #include #include +#include #include #include "glog/logging.h" +#include "infini_train/include/nn/lora/lora_parallel_linear.h" #include "infini_train/include/nn/modules/activations.h" #include "infini_train/include/nn/modules/transformer/moe/moe_utils.h" #include "infini_train/include/nn/modules/transformer/transformer_config.h" @@ -94,4 +98,53 @@ std::vector> MLP::Forward( return (*modules_[kCProjLayerName])(x2); } +ShardedStateDict MLP::BuildShardedStateDict( + const ShardedStateDictBuildContext &context) const { + auto state = Module::BuildShardedStateDict(context); + if (!modules_.contains(kSwiGLULayerName)) { + return state; + } + + const auto c_fc_prefix = context.logical_prefix.empty() + ? kCFcLayerName + : context.logical_prefix + "." + kCFcLayerName; + const auto weight_key = + c_fc_prefix + "." + parallel::ColumnParallelLinear::kParamWeightName; + const auto bias_key = + c_fc_prefix + "." + parallel::ColumnParallelLinear::kParamBiasName; + const auto lora_b_key = + c_fc_prefix + "." + lora::LoRAColumnParallelLinear::kParamLoraBName; + const auto weight_it = state.tensors.find(weight_key); + CHECK(weight_it != state.tensors.end()) + << "Missing packed SwiGLU weight metadata: " << weight_key; + + const auto &weight = weight_it->second; + CHECK(!weight.global_shape.empty() && !weight.local_shape.empty()); + CHECK_EQ(weight.global_shape[0] % 2, 0); + CHECK_EQ(weight.local_shape[0] % 2, 0); + const int64_t global_rows_per_projection = weight.global_shape[0] / 2; + const int64_t local_rows_per_projection = weight.local_shape[0] / 2; + const int rank = parallel::tp_rank; + + for (auto &[key, tensor] : state.tensors) { + if (key != weight_key && key != bias_key && key != lora_b_key) { + continue; + } + CHECK(!tensor.global_shape.empty() && !tensor.local_shape.empty()); + CHECK_EQ(tensor.global_shape[0], weight.global_shape[0]); + CHECK_EQ(tensor.local_shape[0], weight.local_shape[0]); + tensor.global_offset.assign(tensor.global_shape.size(), 0); + tensor.segments = { + {.global_offset = rank * local_rows_per_projection, + .local_offset = 0, + .length = local_rows_per_projection}, + {.global_offset = + global_rows_per_projection + rank * local_rows_per_projection, + .local_offset = local_rows_per_projection, + .length = local_rows_per_projection}, + }; + } + return state; +} + } // namespace infini_train::nn diff --git a/infini_train/src/nn/modules/transformer/transformer.cc b/infini_train/src/nn/modules/transformer/transformer.cc index 8c479fe0..6af42a09 100644 --- a/infini_train/src/nn/modules/transformer/transformer.cc +++ b/infini_train/src/nn/modules/transformer/transformer.cc @@ -395,21 +395,23 @@ TransformerModel::TransformerModel(const TransformerConfig config) start_layer_to_layer_size_and_chunk[start_layer] = std::make_pair(end_layer - start_layer, chunk); } - std::vector> h; + std::unordered_map> h; int chunk_idx = 0; for (auto &[start_layer, layer_size_and_chunk] : start_layer_to_layer_size_and_chunk) { auto [layer_size, chunk] = layer_size_and_chunk; - for (int idx = 0; idx < layer_size; ++idx) { - h.push_back(chunk->mutable_module(TransformerChunk::kHLayerName) - ->mutable_module(std::to_string(idx))); + for (int local_layer = 0; local_layer < layer_size; ++local_layer) { + const auto global_layer = start_layer + local_layer; + h[std::to_string(global_layer)] = + chunk->mutable_module(TransformerChunk::kHLayerName) + ->mutable_module(std::to_string(local_layer)); } modules_[kPPChunkNamePrefix + std::to_string(chunk_idx)] = std::move(chunk); ++chunk_idx; } transformer[TransformerChunk::kHLayerName] = - std::make_shared(std::move(h)); + std::make_shared(std::move(h)); } if (stage_info_.is_last_stage) { @@ -434,6 +436,12 @@ TransformerModel::TransformerModel(const TransformerConfig config) // GPT2::FromLLMC() loading logic to respect weight tying (do not create/load // a separate lm_head.weight tensor; load once into the tied weight) so // parameter counting matches PyTorch/PEFT. + // FIXME(jym): A checkpoint saved with PP > 1 contains independent wte.weight + // and lm_head.weight tensors. Loading it with PP == 1 aliases both keys to + // the same destination Tensor, so Module::LoadStateDict copies both values + // and the final result depends on unordered-map iteration order. Canonicalize + // tied checkpoint keys, or reject PP > 1 to PP == 1 resharding for tied + // models until cross-stage weight tying is supported. if (config_.tie_weights && nn::parallel::global::GetPipelineParallelSize() == 1) { // https://paperswithcode.com/method/weight-tying @@ -446,6 +454,22 @@ TransformerModel::TransformerModel(const TransformerConfig config) } } +ShardedStateDict TransformerModel::BuildShardedStateDict( + const ShardedStateDictBuildContext &context) const { + auto state = Module::BuildShardedStateDict(context); + if (stage_info_.is_last_stage) { + const auto lm_head_prefix = + context.logical_prefix.empty() + ? TransformerLastStage::kLMHeadLayerName + : context.logical_prefix + "." + + TransformerLastStage::kLMHeadLayerName; + const auto weight_key = + lm_head_prefix + "." + parallel::ColumnParallelLinear::kParamWeightName; + state.tensors.at(weight_key).allow_shape_mismatch = true; + } + return state; +} + std::vector> TransformerModel::Forward( const std::vector> &x) { auto x1 = (*modules_[kPPFirstStageName])(x); diff --git a/infini_train/src/nn/parallel/ddp/distributed_data_parallel.cc b/infini_train/src/nn/parallel/ddp/distributed_data_parallel.cc index 7ab9ab19..d9e08af9 100644 --- a/infini_train/src/nn/parallel/ddp/distributed_data_parallel.cc +++ b/infini_train/src/nn/parallel/ddp/distributed_data_parallel.cc @@ -323,6 +323,16 @@ void DistributedDataParallel::OnGradReady( } } +ShardedStateDict DistributedDataParallel::BuildShardedStateDict( + const ShardedStateDictBuildContext &context) const { + const auto local_prefix = context.local_prefix.empty() + ? std::string(kModuleName) + : context.local_prefix + "." + kModuleName; + return modules_.at(kModuleName) + ->BuildShardedStateDict({.logical_prefix = context.logical_prefix, + .local_prefix = local_prefix}); +} + std::vector> DistributedDataParallel::Forward( const std::vector> &input_tensors) { deferred_sync_pending_ = !deferred_sync_params_.empty(); diff --git a/infini_train/src/nn/parallel/pp/pipeline_parallel.cc b/infini_train/src/nn/parallel/pp/pipeline_parallel.cc index 5754a10b..d91115f4 100644 --- a/infini_train/src/nn/parallel/pp/pipeline_parallel.cc +++ b/infini_train/src/nn/parallel/pp/pipeline_parallel.cc @@ -121,4 +121,14 @@ PipelineParallel::PipelineParallel( std::vector> *PipelineParallel::mutable_chunks() { return pipeline_stage_->mutable_chunks(); } + +ShardedStateDict PipelineParallel::BuildShardedStateDict( + const ShardedStateDictBuildContext &context) const { + const auto local_prefix = context.local_prefix.empty() + ? std::string(kModuleName) + : context.local_prefix + "." + kModuleName; + return modules_.at(kModuleName) + ->BuildShardedStateDict({.logical_prefix = context.logical_prefix, + .local_prefix = local_prefix}); +} } // namespace infini_train::nn::parallel diff --git a/infini_train/src/nn/parallel/tensor_parallel.cc b/infini_train/src/nn/parallel/tensor_parallel.cc index f9cb6fd6..998874b8 100644 --- a/infini_train/src/nn/parallel/tensor_parallel.cc +++ b/infini_train/src/nn/parallel/tensor_parallel.cc @@ -2,10 +2,12 @@ #include #include +#include #include #include #include #include +#include #include #include "glog/logging.h" @@ -355,6 +357,50 @@ bool ColumnParallelLinear::sequence_parallel() const { return sequence_parallel_; } +ShardedStateDict ColumnParallelLinear::BuildShardedStateDict( + const ShardedStateDictBuildContext &context) const { + ShardedStateDict sd; + int tp_size = global::GetTensorParallelSize(); + + auto &weight = parameter(kParamWeightName); + ShardedTensor w; + w.key = context.logical_prefix.empty() + ? kParamWeightName + : context.logical_prefix + "." + kParamWeightName; + w.local_key = context.local_prefix.empty() + ? kParamWeightName + : context.local_prefix + "." + kParamWeightName; + w.local_tensor = weight.get(); + w.dtype = weight->Dtype(); + w.global_shape = {output_size_per_partition_ * tp_size, weight->Dims()[1]}; + w.local_shape = weight->Dims(); + w.global_offset = {output_size_per_partition_ * tp_rank, 0}; + w.axis_fragmentations = {tp_size, 1}; + sd.tensors[w.key] = std::move(w); + + // Bias is also split along dim=0 + if (bias_) { + auto &bias = parameter(kParamBiasName); + ShardedTensor b; + b.key = context.logical_prefix.empty() + ? kParamBiasName + : context.logical_prefix + "." + kParamBiasName; + b.local_key = context.local_prefix.empty() + ? kParamBiasName + : context.local_prefix + "." + kParamBiasName; + b.local_tensor = bias.get(); + b.dtype = bias->Dtype(); + b.global_shape = { + static_cast(output_size_per_partition_ * tp_size)}; + b.local_shape = bias->Dims(); + b.global_offset = {output_size_per_partition_ * tp_rank}; + b.axis_fragmentations = {tp_size}; + sd.tensors[b.key] = std::move(b); + } + + return sd; +} + RowParallelLinear::RowParallelLinear(int64_t in_features, int64_t out_features, bool bias, bool reduce_output, bool input_is_parallel, bool skip_bias_add, @@ -433,10 +479,49 @@ bool RowParallelLinear::input_is_parallel() const { return input_is_parallel_; } bool RowParallelLinear::skip_bias_add() const { return skip_bias_add_; } bool RowParallelLinear::sequence_parallel() const { return sequence_parallel_; } +ShardedStateDict RowParallelLinear::BuildShardedStateDict( + const ShardedStateDictBuildContext &context) const { + ShardedStateDict sd; + int tp_size = global::GetTensorParallelSize(); + + auto &weight = parameter(kParamWeightName); + ShardedTensor w; + w.key = context.logical_prefix.empty() + ? kParamWeightName + : context.logical_prefix + "." + kParamWeightName; + w.local_key = context.local_prefix.empty() + ? kParamWeightName + : context.local_prefix + "." + kParamWeightName; + w.local_tensor = weight.get(); + w.dtype = weight->Dtype(); + w.global_shape = {weight->Dims()[0], input_size_per_partition_ * tp_size}; + w.local_shape = weight->Dims(); + w.global_offset = {0, input_size_per_partition_ * tp_rank}; + w.axis_fragmentations = {1, tp_size}; + sd.tensors[w.key] = std::move(w); + + // Bias is NOT sharded in RowParallelLinear + if (bias_) { + auto &bias = parameter(kParamBiasName); + const auto key = context.logical_prefix.empty() + ? kParamBiasName + : context.logical_prefix + "." + kParamBiasName; + auto b = MakeShardedTensor(key, bias->Dtype(), bias->Dims()); + b.local_key = context.local_prefix.empty() + ? kParamBiasName + : context.local_prefix + "." + kParamBiasName; + b.local_tensor = bias.get(); + sd.tensors[b.key] = std::move(b); + } + + return sd; +} + VocabParallelEmbedding::VocabParallelEmbedding(int64_t num_embeddings, int64_t embedding_dim, bool reduce_scatter_embeddings) : CloneableModule(kType), + vocab_size_global_(num_embeddings), embedding_dim_(embedding_dim), reduce_scatter_embeddings_(reduce_scatter_embeddings) { auto tp_size = global::GetTensorParallelSize(); @@ -501,6 +586,31 @@ std::vector> VocabParallelEmbedding::Forward( return {output}; } +ShardedStateDict VocabParallelEmbedding::BuildShardedStateDict( + const ShardedStateDictBuildContext &context) const { + ShardedStateDict sd; + int tp_size = global::GetTensorParallelSize(); + + auto &weight = parameter(kParamWeightName); + ShardedTensor w; + w.key = context.logical_prefix.empty() + ? kParamWeightName + : context.logical_prefix + "." + kParamWeightName; + w.local_key = context.local_prefix.empty() + ? kParamWeightName + : context.local_prefix + "." + kParamWeightName; + w.local_tensor = weight.get(); + w.dtype = weight->Dtype(); + w.global_shape = {vocab_size_global_, embedding_dim_}; + w.local_shape = weight->Dims(); + w.global_offset = {vocab_start_index_, 0}; + w.axis_fragmentations = {tp_size, 1}; + w.allow_shape_mismatch = true; + sd.tensors[w.key] = std::move(w); + + return sd; +} + std::vector> VocabParallelCrossEntropy::Forward( const std::vector> &input_tensors) { CHECK_EQ(input_tensors.size(), 2) << kType << " expects {logits, target}"; @@ -585,7 +695,7 @@ std::vector> VocabParallelCrossEntropy::Forward( auto sum_exp = (tp_size > 1) ? ReduceFromTPRegionFunc(sum_exp_local)[0] : sum_exp_local; - // 4. Perform Softmax（local shards but normalize globally） + // 4. Perform Softmax (local shards but normalize globally). auto softmax_local = exp_local->Div(sum_exp->Unsqueeze(-1)); // 5. Perform allreduce to get global predicted_logit @@ -602,7 +712,7 @@ std::vector> VocabParallelCrossEntropy::Forward( auto log_sum_exp = sum_exp->Log(); auto loss = log_sum_exp->Sub(predicted); - // 7. Label smoothing（According to Megatron-LM） + // 7. Apply label smoothing according to Megatron-LM. // TODO(zbl): adjust smoothing coef according to vocab_size_original if (label_smoothing_ > 0.0f) { // mean_logp over *valid tokens only*: diff --git a/infini_train/src/optimizer.cc b/infini_train/src/optimizer.cc index 52c3921f..d5e12dad 100644 --- a/infini_train/src/optimizer.cc +++ b/infini_train/src/optimizer.cc @@ -3,6 +3,7 @@ #include #include +#include #include #include @@ -45,6 +46,17 @@ std::unordered_map> Optimizer::StateDict() return {}; } +NamedParameterList Optimizer::NamedParameters() const { + NamedParameterList result; + result.reserve(params_.size()); + for (size_t i = 0; i < params_.size(); ++i) { + result.emplace_back( + parameter_names_.empty() ? std::to_string(i) : parameter_names_[i], + params_[i]); + } + return result; +} + void Optimizer::set_learning_rate(float lr) { learning_rate_ = lr; } float Optimizer::learning_rate() const { return learning_rate_; } @@ -176,14 +188,14 @@ std::unordered_map> Adam::StateDict() for (size_t i = 0; i < m_.size(); ++i) { const auto suffix = parameter_names_.empty() ? std::to_string(i) : parameter_names_[i]; - state.emplace("adam.m." + suffix, m_[i]); - state.emplace("adam.v." + suffix, v_[i]); + state.emplace(std::string(kAdamFirstMomentPrefix) + suffix, m_[i]); + state.emplace(std::string(kAdamSecondMomentPrefix) + suffix, v_[i]); } auto t_tensor = std::make_shared(std::vector{}, DataType::kINT64, Device()); *static_cast(t_tensor->DataPtr()) = t_; - state.emplace("adam.t", t_tensor); + state.emplace(std::string(kAdamStepKey), t_tensor); return state; } @@ -193,16 +205,18 @@ void Adam::LoadStateDict( for (size_t i = 0; i < m_.size(); ++i) { const auto suffix = parameter_names_.empty() ? std::to_string(i) : parameter_names_[i]; - const auto m_key = "adam.m." + suffix; - const auto v_key = "adam.v." + suffix; + const auto m_key = std::string(kAdamFirstMomentPrefix) + suffix; + const auto v_key = std::string(kAdamSecondMomentPrefix) + suffix; CHECK(state_dict.contains(m_key)) << "Missing optimizer state: " << m_key; CHECK(state_dict.contains(v_key)) << "Missing optimizer state: " << v_key; m_[i]->CopyFrom(state_dict.at(m_key)); v_[i]->CopyFrom(state_dict.at(v_key)); } - CHECK(state_dict.contains("adam.t")) << "Missing optimizer state: adam.t"; - const Tensor t_cpu = state_dict.at("adam.t")->To(Device()); + const std::string step_key(kAdamStepKey); + CHECK(state_dict.contains(step_key)) + << "Missing optimizer state: " << step_key; + const Tensor t_cpu = state_dict.at(step_key)->To(Device()); t_ = *static_cast(t_cpu.DataPtr()); } } // namespace optimizers diff --git a/infini_train/src/shard_spec.cc b/infini_train/src/shard_spec.cc new file mode 100644 index 00000000..3197a31e --- /dev/null +++ b/infini_train/src/shard_spec.cc @@ -0,0 +1,28 @@ +#include "infini_train/include/shard_spec.h" + +#include +#include + +#include "glog/logging.h" + +namespace infini_train { + +bool ShardedTensor::operator==(const ShardedTensor &other) const { + return key == other.key && local_key == other.local_key && + dtype == other.dtype && global_shape == other.global_shape && + local_shape == other.local_shape && + global_offset == other.global_offset && + axis_fragmentations == other.axis_fragmentations && + allow_shape_mismatch == other.allow_shape_mismatch && + segments == other.segments; +} + +void ShardedStateDict::Merge(ShardedStateDict &&other) { + for (auto &[key, info] : other.tensors) { + const auto display_key = key; + const auto [_, inserted] = tensors.emplace(std::move(key), std::move(info)); + CHECK(inserted) << "Duplicate sharded state-dict key: " << display_key; + } +} + +} // namespace infini_train diff --git a/scripts/run_models_and_profile.bash b/scripts/run_models_and_profile.bash index 34edd26f..84c6351f 100755 --- a/scripts/run_models_and_profile.bash +++ b/scripts/run_models_and_profile.bash @@ -374,9 +374,10 @@ args_string_for_test() { jq -r --argjson g "$group_idx" --argjson t "$test_idx" --arg model "$model_name" --arg test_id "$test_id" ' def namespaced_path($p; $model; $mode): - if ($p | test("/checkpoint_step_[0-9]+($|/)")) then - ($p | capture("^(?.*)/(?checkpoint_step_[0-9]+(?:/.*)?)$")) as $m - | ($m.prefix + "/" + $model + "/" + $mode + "/" + $m.step) + if ($p | test("/(?:checkpoint_step_|iter_)[0-9]+($|/)")) then + ($p | capture("^(?.*)/(?:checkpoint_step_|iter_)(?[0-9]+)(?/.*)?$")) as $m + | ("0000000" + $m.iteration)[-7:] as $iteration + | ($m.prefix + "/" + $model + "/" + $mode + "/iter_" + $iteration + ($m.suffix // "")) else ($p + "/" + $model + "/" + $mode) end; diff --git a/tests/checkpoint/test_checkpoint_serialization.cc b/tests/checkpoint/test_checkpoint_serialization.cc index 3fe6cf07..f418565b 100644 --- a/tests/checkpoint/test_checkpoint_serialization.cc +++ b/tests/checkpoint/test_checkpoint_serialization.cc @@ -1,15 +1,37 @@ -#include +#include +#include +#include #include +#include +#include +#include +#include #include +#include +#include +#include +#include #include #include "gtest/gtest.h" #include "infini_train/include/checkpoint/checkpoint.h" +#include "infini_train/include/checkpoint/constants.h" +#include "infini_train/include/checkpoint/load_planner.h" +#include "infini_train/include/checkpoint/load_strategy.h" +#include "infini_train/include/checkpoint/save_planner.h" #include "infini_train/include/datatype.h" +#include "infini_train/include/device.h" +#include "infini_train/include/nn/lora/lora_config.h" +#include "infini_train/include/nn/lora/lora_parallel_linear.h" #include "infini_train/include/nn/modules/linear.h" +#include "infini_train/include/nn/modules/module.h" +#include "infini_train/include/nn/modules/transformer/mlp.h" +#include "infini_train/include/nn/modules/transformer/transformer_config.h" +#include "infini_train/include/nn/parallel/tensor_parallel.h" #include "infini_train/include/optimizer.h" +#include "infini_train/include/shard_spec.h" #include "infini_train/include/tensor.h" #include "tests/common/test_utils.h" @@ -17,8 +39,168 @@ using namespace infini_train; namespace nn = infini_train::nn; +namespace { + +class NamedParameterModule : public nn::Module { + public: + void AddParameter(const std::string &name, + const std::shared_ptr ¶meter) { + parameters_[name] = parameter; + } + + void AddModule(const std::string &name, + const std::shared_ptr &module) { + modules_[name] = module; + } +}; + +class CheckpointWrapper final : public NamedParameterModule { + public: + explicit CheckpointWrapper(const std::shared_ptr &module) { + AddModule("module", module); + } + + using Module::BuildShardedStateDict; + ShardedStateDict BuildShardedStateDict( + const nn::ShardedStateDictBuildContext &context) const override { + const auto local_prefix = context.local_prefix.empty() + ? std::string("module") + : context.local_prefix + ".module"; + return modules_.at("module")->BuildShardedStateDict( + {.logical_prefix = context.logical_prefix, + .local_prefix = local_prefix}); + } +}; + +} // namespace + +TEST(ModuleNamedParametersTest, + SupportsTorchStyleArgumentsAndSharedParameterDeduplication) { + auto root = std::make_shared(); + auto child = std::make_shared(); + auto shared = std::make_shared(std::vector{1}, + DataType::kFLOAT32, Device()); + auto child_weight = std::make_shared(std::vector{1}, + DataType::kFLOAT32, Device()); + root->AddParameter("root_weight", shared); + child->AddParameter("alias", shared); + child->AddParameter("weight", child_weight); + root->AddModule("child", child); + + const auto local = root->NamedParameters("model", false); + ASSERT_EQ(local.size(), 1); + EXPECT_EQ(local[0].first, "model.root_weight"); + EXPECT_EQ(local[0].second, shared); + + const auto deduplicated = root->NamedParameters("model"); + ASSERT_EQ(deduplicated.size(), 2); + EXPECT_EQ(deduplicated[0].first, "model.root_weight"); + EXPECT_EQ(deduplicated[1].first, "model.child.weight"); + + const auto aliases = root->NamedParameters("model", true, false); + ASSERT_EQ(aliases.size(), 3); + EXPECT_EQ(aliases[0].first, "model.root_weight"); + EXPECT_EQ(aliases[1].first, "model.child.alias"); + EXPECT_EQ(aliases[1].second, shared); + EXPECT_EQ(aliases[2].first, "model.child.weight"); +} + +TEST(ShardedStateDictTest, + WrapperKeepsOrdinaryModulePathsButOmitsThemFromLogicalKeys) { + auto model = std::make_shared(); + auto weight = std::make_shared(std::vector{2}, + DataType::kFLOAT32, Device()); + model->AddParameter("weight", weight); + auto wrapped = std::make_shared( + std::make_shared(model)); + + const auto state_dict = wrapped->StateDict(); + ASSERT_TRUE(state_dict.contains("module.module.weight")); + const auto named = wrapped->NamedParameters(); + ASSERT_EQ(named.size(), 1); + EXPECT_EQ(named[0].first, "module.module.weight"); + EXPECT_TRUE(wrapped->NamedParameters("", false).empty()); + + const auto sharded = wrapped->BuildShardedStateDict(); + ASSERT_EQ(sharded.tensors.size(), 1); + const auto &info = sharded.tensors.at("weight"); + EXPECT_EQ(info.local_key, "module.module.weight"); + EXPECT_EQ(info.local_tensor, weight.get()); +} + +TEST(ShardedStateDictTest, ParentAndRealModuleNamesRemainInBothNamespaces) { + auto leaf = std::make_shared(); + leaf->AddParameter("weight", + std::make_shared(std::vector{1}, + DataType::kFLOAT32, Device())); + auto real_module = std::make_shared(); + real_module->AddModule("module", leaf); + auto root = std::make_shared(); + root->AddModule("child", std::make_shared(real_module)); + + const auto sharded = root->BuildShardedStateDict(); + const auto &info = sharded.tensors.at("child.module.weight"); + EXPECT_EQ(info.local_key, "child.module.module.weight"); + EXPECT_TRUE(root->StateDict().contains(info.local_key)); +} + class CheckpointSerializationTest : public test::InfiniTrainTest {}; +TEST(ShardedStateDictTest, RejectsDuplicateKeysWhenMerging) { + ShardedStateDict destination; + destination.tensors["weight"] = {.key = "weight"}; + ShardedStateDict source; + source.tensors["weight"] = {.key = "weight"}; + + EXPECT_DEATH(destination.Merge(std::move(source)), + "Duplicate sharded state-dict key: weight"); +} + +TEST(ShardedStateDictTest, PackedSwiGLUUsesGateAndUpSegments) { + nn::TransformerConfig config; + config.n_embd = 4; + config.activation_type = nn::MLPType::kSwiGLU; + config.ffn_hidden_size = 8; + + const auto state = nn::MLP(config).BuildShardedStateDict("mlp"); + const auto &weight = state.tensors.at("mlp.c_fc.weight"); + ASSERT_EQ(weight.segments.size(), 2); + EXPECT_EQ(weight.segments[0], + (ShardSegment{.global_offset = 0, .local_offset = 0, .length = 8})); + EXPECT_EQ(weight.segments[1], + (ShardSegment{.global_offset = 8, .local_offset = 8, .length = 8})); + + const auto &bias = state.tensors.at("mlp.c_fc.bias"); + EXPECT_EQ(bias.segments, weight.segments); +} + +TEST(ShardedStateDictTest, PackedSwiGLUDoesNotClassifyLoRAAByShape) { + nn::TransformerConfig config; + config.n_embd = 4; + config.activation_type = nn::MLPType::kSwiGLU; + config.ffn_hidden_size = 8; + + auto mlp = std::make_shared(config); + auto base = std::dynamic_pointer_cast( + mlp->mutable_module(nn::MLP::kCFcLayerName)); + ASSERT_NE(base, nullptr); + + nn::lora::LoRAConfig lora_config; + lora_config.rank = + base->parameter(nn::parallel::ColumnParallelLinear::kParamWeightName) + ->Dims()[0]; + mlp->mutable_module(nn::MLP::kCFcLayerName) = + std::make_shared(base, lora_config); + + const auto state = mlp->BuildShardedStateDict("mlp"); + const auto &weight = state.tensors.at("mlp.c_fc.weight"); + const auto &lora_a = state.tensors.at("mlp.c_fc.lora_A"); + const auto &lora_b = state.tensors.at("mlp.c_fc.lora_B"); + ASSERT_EQ(lora_a.global_shape[0], weight.global_shape[0]); + EXPECT_TRUE(lora_a.segments.empty()); + EXPECT_EQ(lora_b.segments, weight.segments); +} + TEST_P(CheckpointSerializationTest, SaveAndLoadModelFP32) { auto dir = std::filesystem::temp_directory_path() / "test_ckpt_fp32"; std::filesystem::remove_all(dir); @@ -33,7 +215,7 @@ TEST_P(CheckpointSerializationTest, SaveAndLoadModelFP32) { p2->Fill(-1.5f); *model1->mutable_parameter("bias") = p2; - auto opt1 = std::make_shared(model1->Parameters(), 0.01); + auto opt1 = optimizers::Adam::CreateNamed(0.01)(model1->NamedParameters()); TrainerState saved{.global_step = 42, .consumed_train_samples = 100}; Checkpoint::Save(dir, *model1, opt1.get(), saved, nullptr); @@ -46,7 +228,7 @@ TEST_P(CheckpointSerializationTest, SaveAndLoadModelFP32) { DataType::kFLOAT32, GetDevice()); q2->Fill(0.0f); *model2->mutable_parameter("bias") = q2; - auto opt2 = std::make_shared(model2->Parameters(), 0.01); + auto opt2 = optimizers::Adam::CreateNamed(0.01)(model2->NamedParameters()); TrainerState loaded; Checkpoint::Load(dir, *model2, opt2.get(), loaded, nullptr); @@ -59,4 +241,659 @@ TEST_P(CheckpointSerializationTest, SaveAndLoadModelFP32) { std::filesystem::remove_all(dir); } +TEST_P(CheckpointSerializationTest, DirectMetadataOffsetSupportsColumnSlices) { + auto dir = std::filesystem::temp_directory_path() / "test_ckpt_region"; + std::filesystem::remove_all(dir); + std::filesystem::create_directories(dir); + auto matrix = std::make_shared(std::vector{4, 4}, + DataType::kFLOAT32, Device()); + auto *values = static_cast(matrix->DataPtr()); + for (int row = 0; row < 4; ++row) { + for (int column = 0; column < 4; ++column) { + values[row * 4 + column] = row * 10.0f + column; + } + } + auto path = dir / checkpoint::kModelCheckpointFilename; + Checkpoint::SaveStateDictFile(path, {{"matrix", matrix}}); + constexpr uint64_t data_offset = sizeof(uint32_t) * 3 + sizeof(uint32_t) + + sizeof("matrix") - 1 + sizeof(int8_t) + + sizeof(uint32_t) + sizeof(int64_t) * 2 + + sizeof(uint64_t); + checkpoint::LoadPlan plan; + plan.tensors["matrix"] = { + .key = "matrix", + .dtype = DataType::kFLOAT32, + .global_shape = {4, 4}, + .target_shape = {4, 2}, + .shard_dim = 1, + .reads = {{.key = "matrix", + .filename = checkpoint::kModelCheckpointFilename, + .dtype = DataType::kFLOAT32, + .global_shape = {4, 4}, + .byte_size = sizeof(float) * 16, + .data_offset = data_offset, + .shard_dim = 1, + .source_offset = 1, + .target_offset = 0, + .length = 2, + .source_shape = {4, 4}}}}; + checkpoint::IndexedRegionLoadStrategy strategy; + const auto planned = strategy.Execute(dir, plan); + const auto *planned_data = + static_cast(planned.at("matrix")->DataPtr()); + for (int row = 0; row < 4; ++row) { + for (int column = 0; column < 2; ++column) { + EXPECT_FLOAT_EQ(planned_data[row * 2 + column], row * 10.0f + column + 1); + } + } + + std::filesystem::remove_all(dir); +} + +TEST_P(CheckpointSerializationTest, ConvertsSavedBF16TensorToFP32Target) { + auto dir = std::filesystem::temp_directory_path() / "test_ckpt_bf16_to_fp32"; + std::filesystem::remove_all(dir); + std::filesystem::create_directories(dir); + + auto source_fp32 = std::make_shared(std::vector{2, 2}, + DataType::kFLOAT32, Device()); + auto *source_data = static_cast(source_fp32->DataPtr()); + source_data[0] = 1.0f; + source_data[1] = 2.0f; + source_data[2] = 3.0f; + source_data[3] = 4.0f; + auto source_bf16 = + std::make_shared(source_fp32->To(DataType::kBFLOAT16)); + Checkpoint::SaveStateDictFile(dir / checkpoint::kModelCheckpointFilename, + {{"weight", source_bf16}}); + constexpr uint64_t data_offset = sizeof(uint32_t) * 3 + sizeof(uint32_t) + + sizeof("weight") - 1 + sizeof(int8_t) + + sizeof(uint32_t) + sizeof(int64_t) * 2 + + sizeof(uint64_t); + + checkpoint::LoadPlan plan; + plan.tensors["weight"] = { + .key = "weight", + .dtype = DataType::kFLOAT32, + .global_shape = {2, 2}, + .target_shape = {2, 2}, + .reads = {{.key = "weight", + .filename = checkpoint::kModelCheckpointFilename, + .dtype = DataType::kBFLOAT16, + .global_shape = {2, 2}, + .byte_size = source_bf16->SizeInBytes(), + .data_offset = data_offset, + .shard_dim = -1, + .source_shape = {2, 2}}}}; + + checkpoint::IndexedRegionLoadStrategy strategy; + const auto loaded = strategy.Execute(dir, plan).at("weight"); + ASSERT_EQ(loaded->Dtype(), DataType::kFLOAT32); + const auto *loaded_data = static_cast(loaded->DataPtr()); + for (int i = 0; i < 4; ++i) { + EXPECT_FLOAT_EQ(loaded_data[i], source_data[i]); + } + + std::filesystem::remove_all(dir); +} + +TEST(CheckpointLoadPlannerTest, PadsVocabularyTailWhenTargetTpUsesPaddedVocab) { + const auto dir = + std::filesystem::temp_directory_path() / "test_vocab_padding_reshard"; + std::filesystem::remove_all(dir); + std::filesystem::create_directories(dir); + + auto source = std::make_shared(std::vector{5, 2}, + DataType::kFLOAT32, Device()); + auto *source_data = static_cast(source->DataPtr()); + for (int i = 0; i < 10; ++i) { + source_data[i] = static_cast(i); + } + Checkpoint::SaveStateDictFile(dir / checkpoint::kModelCheckpointFilename, + {{"lm_head.weight", source}}); + constexpr uint64_t data_offset = sizeof(uint32_t) * 3 + sizeof(uint32_t) + + sizeof("lm_head.weight") - 1 + + sizeof(int8_t) + sizeof(uint32_t) + + sizeof(int64_t) * 2 + sizeof(uint64_t); + + Checkpoint::CheckpointMetadata metadata; + metadata.tensors.push_back({.key = "lm_head.weight", + .dtype_str = "float32", + .global_shape = {5, 2}, + .local_shape = {5, 2}, + .global_offset = {0, 0}, + .axis_fragmentations = {1, 1}, + .file = checkpoint::kModelCheckpointFilename, + .offset = data_offset, + .byte_size = sizeof(float) * 10}); + + ShardedStateDict target; + target.tensors["lm_head.weight"] = {.key = "lm_head.weight", + .dtype = DataType::kFLOAT32, + .global_shape = {8, 2}, + .local_shape = {4, 2}, + .global_offset = {4, 0}, + .axis_fragmentations = {2, 1}, + .allow_shape_mismatch = true}; + + const auto plan = checkpoint::LoadPlanner::PlanReshard(metadata, target); + ASSERT_EQ(plan.tensors.at("lm_head.weight").trailing_zero_fill, 3); + checkpoint::IndexedRegionLoadStrategy strategy; + const auto loaded = strategy.Execute(dir, plan); + const auto *values = + static_cast(loaded.at("lm_head.weight")->DataPtr()); + EXPECT_FLOAT_EQ(values[0], 8.0f); + EXPECT_FLOAT_EQ(values[1], 9.0f); + for (int i = 2; i < 8; ++i) { + EXPECT_FLOAT_EQ(values[i], 0.0f); + } + + std::filesystem::remove_all(dir); +} + +TEST(CheckpointMetadataTest, RejectsMalformedIntegerArrays) { + const auto dir = std::filesystem::temp_directory_path() / + "test_malformed_checkpoint_metadata"; + std::filesystem::remove_all(dir); + std::filesystem::create_directories(dir); + { + std::ofstream metadata_file(dir / checkpoint::kMetadataFilename); + metadata_file << R"({ + "version": 3, + "format": "infinitrain_sharded", + "tensors": [ + { + "key": "weight", + "dtype": "float32", + "global_shape": [8, invalid] + } + ] +})"; + } + + EXPECT_DEATH(Checkpoint::LoadMetadata(dir), + "Invalid integer in checkpoint metadata field global_shape"); + std::filesystem::remove_all(dir); +} + +TEST_P(CheckpointSerializationTest, GlobalMetadataRoundTrip) { + auto dir = std::filesystem::temp_directory_path() / "test_global_metadata"; + std::filesystem::remove_all(dir); + std::filesystem::create_directories(dir); + Checkpoint::CheckpointMetadata metadata; + metadata.version = checkpoint::kDistributedCheckpointVersion; + metadata.has_metadata = true; + metadata.tensors.push_back( + {.key = "layer.0.weight", + .dtype_str = "float32", + .global_shape = {8, 4}, + .local_shape = {4, 4}, + .global_offset = {0, 0}, + .axis_fragmentations = {2, 1}, + .segments = {{.global_offset = 0, .local_offset = 0, .length = 4}}, + .file = "rank_000000/model.ckpt", + .byte_size = 64, + .stored_on_ranks = {0}, + .pp_rank = 0}); + Checkpoint::SaveMetadataFile(dir / checkpoint::kMetadataFilename, metadata); + std::ifstream metadata_file(dir / checkpoint::kMetadataFilename); + const std::string metadata_json( + (std::istreambuf_iterator(metadata_file)), + std::istreambuf_iterator()); + EXPECT_EQ(metadata_json.find("\"iteration\""), std::string::npos); + EXPECT_EQ(metadata_json.find("\"parallel_config\""), std::string::npos); + EXPECT_EQ(metadata_json.find("\"model_config\""), std::string::npos); + + auto loaded = Checkpoint::LoadMetadata(dir); + ASSERT_TRUE(loaded.has_metadata); + ASSERT_EQ(loaded.tensors.size(), 1); + EXPECT_EQ(loaded.tensors[0].file, "rank_000000/model.ckpt"); + EXPECT_EQ(loaded.tensors[0].global_offset, std::vector({0, 0})); + EXPECT_EQ(loaded.tensors[0].axis_fragmentations, std::vector({2, 1})); + ASSERT_EQ(loaded.tensors[0].segments.size(), 1); + EXPECT_EQ(loaded.tensors[0].segments[0], + (ShardSegment{.global_offset = 0, .local_offset = 0, .length = 4})); + std::filesystem::remove_all(dir); +} + INFINI_TRAIN_REGISTER_TEST(CheckpointSerializationTest); + +namespace { +Checkpoint::CheckpointMetadata::TensorEntry MakeSavedShard( + const std::string &key, int count, int index, int64_t global_size, + const std::string &file) { + return {.key = key, + .dtype_str = "float32", + .global_shape = {global_size, 4}, + .local_shape = {global_size / count, 4}, + .global_offset = {global_size / count * index, 0}, + .axis_fragmentations = {count, 1}, + .file = file}; +} + +ShardedStateDict MakeTarget(const std::string &key, int count, int index, + int64_t global_size) { + ShardedStateDict target; + target.tensors[key] = {.key = key, + .dtype = DataType::kFLOAT32, + .global_shape = {global_size, 4}, + .local_shape = {global_size / count, 4}, + .global_offset = {global_size / count * index, 0}, + .axis_fragmentations = {count, 1}}; + return target; +} + +std::shared_ptr MakeColumnTensor(const std::vector &values) { + auto tensor = std::make_shared( + std::vector{static_cast(values.size()), 1}, + DataType::kFLOAT32, Device()); + std::copy(values.begin(), values.end(), + static_cast(tensor->DataPtr())); + return tensor; +} + +uint64_t SingleTensorDataOffset(const std::string &key, size_t dimensions) { + return sizeof(uint32_t) * 3 + sizeof(uint32_t) + key.size() + sizeof(int8_t) + + sizeof(uint32_t) + sizeof(int64_t) * dimensions + sizeof(uint64_t); +} +} // namespace + +TEST(CheckpointOptimizerShardingTest, AdamMomentsReuseModelShardMetadata) { + ShardedStateDict model; + model.tensors["c_attn.weight"] = { + .key = "c_attn.weight", + .dtype = DataType::kFLOAT32, + .global_shape = {24, 4}, + .local_shape = {6, 4}, + .global_offset = {0, 0}, + .axis_fragmentations = {4, 1}, + .segments = + { + {.global_offset = 0, .local_offset = 0, .length = 4}, + {.global_offset = 16, .local_offset = 4, .length = 1}, + {.global_offset = 20, .local_offset = 5, .length = 1}, + }, + }; + + auto moment = std::make_shared(std::vector{6, 4}, + DataType::kFLOAT32, Device()); + auto step = std::make_shared(std::vector{}, DataType::kINT64, + Device()); + std::unordered_map> optimizer_state = { + {"adam.m.c_attn.weight", moment}, + {"adam.v.c_attn.weight", moment}, + {"adam.t", step}, + }; + + const auto optimizer = + checkpoint::BuildOptimizerShardedStateDict(model, optimizer_state); + ASSERT_EQ(optimizer.tensors.size(), 3); + const auto &m = optimizer.tensors.at("adam.m.c_attn.weight"); + EXPECT_EQ(m.global_shape, model.tensors.at("c_attn.weight").global_shape); + EXPECT_EQ(m.local_shape, model.tensors.at("c_attn.weight").local_shape); + EXPECT_EQ(m.segments, model.tensors.at("c_attn.weight").segments); + EXPECT_EQ(m.local_key, "adam.m.c_attn.weight"); + EXPECT_EQ(m.dtype, moment->Dtype()); + EXPECT_EQ(m.allow_shape_mismatch, + model.tensors.at("c_attn.weight").allow_shape_mismatch); + const auto &t = optimizer.tensors.at("adam.t"); + EXPECT_TRUE(t.global_shape.empty()); + EXPECT_TRUE(t.local_shape.empty()); +} + +TEST(CheckpointOptimizerShardingTest, + UsesTensorIdentityToRemoveOnlyWrapperPathComponents) { + auto parameter = std::make_shared(std::vector{4, 4}, + DataType::kFLOAT32, Device()); + auto moment = std::make_shared(std::vector{4, 4}, + DataType::kFLOAT32, Device()); + ShardedStateDict model; + model.tensors["weight"] = {.key = "weight", + .local_key = "module.weight", + .local_tensor = parameter.get(), + .dtype = DataType::kFLOAT32, + .global_shape = {4, 4}, + .local_shape = {4, 4}, + .global_offset = {0, 0}, + .axis_fragmentations = {1, 1}}; + + const auto optimizer = checkpoint::BuildOptimizerShardedStateDict( + model, + {{"adam.m.module.weight", moment}, {"adam.v.module.weight", moment}}, + {{"module.weight", parameter}}); + EXPECT_TRUE(optimizer.tensors.contains("adam.m.weight")); + EXPECT_EQ(optimizer.tensors.at("adam.m.weight").local_key, + "adam.m.module.weight"); + EXPECT_TRUE(optimizer.tensors.contains("adam.v.weight")); + EXPECT_EQ(optimizer.tensors.at("adam.v.weight").local_key, + "adam.v.module.weight"); +} + +TEST(CheckpointOptimizerShardingTest, AdamMomentUsesOptimizerStateDtype) { + ShardedStateDict model; + model.tensors["weight"] = {.key = "weight", + .dtype = DataType::kBFLOAT16, + .global_shape = {4, 4}, + .local_shape = {4, 4}, + .global_offset = {0, 0}, + .axis_fragmentations = {1, 1}}; + auto moment = std::make_shared(std::vector{4, 4}, + DataType::kFLOAT32, Device()); + + const auto optimizer = checkpoint::BuildOptimizerShardedStateDict( + model, {{"adam.m.weight", moment}, {"adam.v.weight", moment}}); + EXPECT_EQ(optimizer.tensors.at("adam.m.weight").dtype, DataType::kFLOAT32); + EXPECT_EQ(optimizer.tensors.at("adam.v.weight").dtype, DataType::kFLOAT32); +} + +TEST(CheckpointLoadPlannerTest, RejectsUnknownCheckpointDtype) { + auto metadata = Checkpoint::CheckpointMetadata{}; + metadata.tensors = {MakeSavedShard("weight", 1, 0, 16, "rank_0/model.ckpt")}; + metadata.tensors.front().dtype_str = "unknown_dtype"; + + EXPECT_DEATH(checkpoint::LoadPlanner::PlanReshard( + metadata, MakeTarget("weight", 1, 0, 16)), + "Unsupported checkpoint tensor dtype: unknown_dtype"); +} + +TEST(CheckpointLoadPlannerTest, TensorParallelTwoToFourReadsOnlyOverlap) { + Checkpoint::CheckpointMetadata metadata; + metadata.tensors = {MakeSavedShard("weight", 2, 0, 16, "rank_0/model.ckpt"), + MakeSavedShard("weight", 2, 1, 16, "rank_1/model.ckpt")}; + + auto plan = checkpoint::LoadPlanner::PlanReshard( + metadata, MakeTarget("weight", 4, 1, 16)); + const auto &reads = plan.tensors.at("weight").reads; + ASSERT_EQ(reads.size(), 1); + EXPECT_EQ(reads[0].filename, "rank_0/model.ckpt"); + EXPECT_EQ(reads[0].source_offset, 4); + EXPECT_EQ(reads[0].target_offset, 0); + EXPECT_EQ(reads[0].length, 4); +} + +TEST(CheckpointLoadPlannerTest, TensorParallelFourToTwoReadsTwoOverlaps) { + Checkpoint::CheckpointMetadata metadata; + for (int index = 0; index < 4; ++index) { + metadata.tensors.push_back( + MakeSavedShard("weight", 4, index, 16, + "rank_" + std::to_string(index) + "/model.ckpt")); + } + + auto plan = checkpoint::LoadPlanner::PlanReshard( + metadata, MakeTarget("weight", 2, 1, 16)); + const auto &reads = plan.tensors.at("weight").reads; + ASSERT_EQ(reads.size(), 2); + EXPECT_EQ(reads[0].filename, "rank_2/model.ckpt"); + EXPECT_EQ(reads[0].target_offset, 0); + EXPECT_EQ(reads[1].filename, "rank_3/model.ckpt"); + EXPECT_EQ(reads[1].target_offset, 4); +} + +TEST(CheckpointLoadPlannerTest, UsesExplicitGlobalOffsetsForUnevenShards) { + Checkpoint::CheckpointMetadata metadata; + metadata.tensors = {{.key = "weight", + .dtype_str = "float32", + .global_shape = {8, 4}, + .local_shape = {3, 4}, + .global_offset = {0, 0}, + .axis_fragmentations = {2, 1}, + .file = "rank_0/model.ckpt"}, + {.key = "weight", + .dtype_str = "float32", + .global_shape = {8, 4}, + .local_shape = {5, 4}, + .global_offset = {3, 0}, + .axis_fragmentations = {2, 1}, + .file = "rank_1/model.ckpt"}}; + ShardedStateDict target; + target.tensors["weight"] = {.key = "weight", + .dtype = DataType::kFLOAT32, + .global_shape = {8, 4}, + .local_shape = {4, 4}, + .global_offset = {2, 0}, + .axis_fragmentations = {2, 1}}; + + auto plan = checkpoint::LoadPlanner::PlanReshard(metadata, target); + const auto &reads = plan.tensors.at("weight").reads; + ASSERT_EQ(reads.size(), 2); + EXPECT_EQ(reads[0].source_offset, 2); + EXPECT_EQ(reads[0].target_offset, 0); + EXPECT_EQ(reads[0].length, 1); + EXPECT_EQ(reads[1].source_offset, 0); + EXPECT_EQ(reads[1].target_offset, 1); + EXPECT_EQ(reads[1].length, 3); +} + +TEST(CheckpointLoadPlannerTest, QkvSegmentsUseDimZeroWhenTpIsOne) { + Checkpoint::CheckpointMetadata metadata; + auto saved = MakeSavedShard("c_attn.weight", 1, 0, 12, "old_pp/model.ckpt"); + saved.local_shape = {12, 4}; + saved.axis_fragmentations = {1, 1}; + saved.segments = { + {.global_offset = 0, .local_offset = 0, .length = 8}, + {.global_offset = 8, .local_offset = 8, .length = 2}, + {.global_offset = 10, .local_offset = 10, .length = 2}, + }; + metadata.tensors.push_back(std::move(saved)); + + ShardedStateDict target; + target.tensors["c_attn.weight"] = { + .key = "c_attn.weight", + .dtype = DataType::kFLOAT32, + .global_shape = {12, 4}, + .local_shape = {12, 4}, + .global_offset = {0, 0}, + .axis_fragmentations = {1, 1}, + .segments = + { + {.global_offset = 0, .local_offset = 0, .length = 8}, + {.global_offset = 8, .local_offset = 8, .length = 2}, + {.global_offset = 10, .local_offset = 10, .length = 2}, + }, + }; + + const auto plan = checkpoint::LoadPlanner::PlanReshard(metadata, target); + const auto &tensor_plan = plan.tensors.at("c_attn.weight"); + EXPECT_EQ(tensor_plan.shard_dim, 0); + ASSERT_EQ(tensor_plan.reads.size(), 3); + EXPECT_EQ(tensor_plan.reads[0].target_offset, 0); + EXPECT_EQ(tensor_plan.reads[1].target_offset, 8); + EXPECT_EQ(tensor_plan.reads[2].target_offset, 10); +} + +TEST(CheckpointLoadPlannerTest, + QkvSegmentsPreserveTargetLocalLayoutAcrossTpChange) { + Checkpoint::CheckpointMetadata metadata; + for (int rank = 0; rank < 4; ++rank) { + auto shard = MakeSavedShard("c_attn.weight", 4, rank, 24, + "rank_" + std::to_string(rank) + "/model.ckpt"); + shard.local_shape = {6, 4}; + shard.global_offset = {0, 0}; + shard.segments = { + {.global_offset = rank * 4, .local_offset = 0, .length = 4}, + {.global_offset = 16 + rank, .local_offset = 4, .length = 1}, + {.global_offset = 20 + rank, .local_offset = 5, .length = 1}, + }; + metadata.tensors.push_back(std::move(shard)); + } + + ShardedStateDict target; + target.tensors["c_attn.weight"] = { + .key = "c_attn.weight", + .dtype = DataType::kFLOAT32, + .global_shape = {24, 4}, + .local_shape = {12, 4}, + .global_offset = {0, 0}, + .axis_fragmentations = {2, 1}, + .segments = + { + {.global_offset = 0, .local_offset = 0, .length = 8}, + {.global_offset = 16, .local_offset = 8, .length = 2}, + {.global_offset = 20, .local_offset = 10, .length = 2}, + }, + }; + + const auto plan = checkpoint::LoadPlanner::PlanReshard(metadata, target); + const auto &reads = plan.tensors.at("c_attn.weight").reads; + ASSERT_EQ(reads.size(), 6); + const std::vector files = { + "rank_0/model.ckpt", "rank_1/model.ckpt", "rank_0/model.ckpt", + "rank_1/model.ckpt", "rank_0/model.ckpt", "rank_1/model.ckpt"}; + const std::vector source_offsets = {0, 0, 4, 4, 5, 5}; + const std::vector target_offsets = {0, 4, 8, 9, 10, 11}; + for (size_t i = 0; i < reads.size(); ++i) { + EXPECT_EQ(reads[i].filename, files[i]); + EXPECT_EQ(reads[i].source_offset, source_offsets[i]); + EXPECT_EQ(reads[i].target_offset, target_offsets[i]); + } +} + +TEST(CheckpointLoadPlannerTest, PackedSwiGLUExecutesTensorParallelTwoToOne) { + const auto dir = + std::filesystem::temp_directory_path() / "test_swiglu_tp2_to_tp1"; + std::filesystem::remove_all(dir); + const std::string key = "c_fc.weight"; + const auto data_offset = SingleTensorDataOffset(key, 2); + + Checkpoint::CheckpointMetadata metadata; + for (int rank = 0; rank < 2; ++rank) { + const auto rank_dir = dir / ("rank_" + std::to_string(rank)); + std::filesystem::create_directories(rank_dir); + const std::vector values = rank == 0 + ? std::vector{0, 1, 4, 5} + : std::vector{2, 3, 6, 7}; + const auto tensor = MakeColumnTensor(values); + Checkpoint::SaveStateDictFile( + rank_dir / checkpoint::kModelCheckpointFilename, {{key, tensor}}); + + metadata.tensors.push_back({ + .key = key, + .dtype_str = "fp32", + .global_shape = {8, 1}, + .local_shape = {4, 1}, + .global_offset = {0, 0}, + .axis_fragmentations = {2, 1}, + .segments = + { + {.global_offset = rank * 2, .local_offset = 0, .length = 2}, + {.global_offset = 4 + rank * 2, .local_offset = 2, .length = 2}, + }, + .file = "rank_" + std::to_string(rank) + "/" + + checkpoint::kModelCheckpointFilename, + .offset = data_offset, + .byte_size = tensor->SizeInBytes(), + }); + } + + ShardedStateDict target; + target.tensors[key] = { + .key = key, + .dtype = DataType::kFLOAT32, + .global_shape = {8, 1}, + .local_shape = {8, 1}, + .global_offset = {0, 0}, + .axis_fragmentations = {1, 1}, + .segments = + { + {.global_offset = 0, .local_offset = 0, .length = 4}, + {.global_offset = 4, .local_offset = 4, .length = 4}, + }, + }; + + const auto plan = checkpoint::LoadPlanner::PlanReshard(metadata, target); + const auto loaded = + checkpoint::IndexedRegionLoadStrategy().Execute(dir, plan).at(key); + const auto *data = static_cast(loaded->DataPtr()); + for (int row = 0; row < 8; ++row) { + EXPECT_FLOAT_EQ(data[row], static_cast(row)); + } + std::filesystem::remove_all(dir); +} + +TEST(CheckpointLoadPlannerTest, PackedSwiGLUExecutesTensorParallelOneToTwo) { + const auto dir = + std::filesystem::temp_directory_path() / "test_swiglu_tp1_to_tp2"; + const auto rank_dir = dir / "rank_0"; + std::filesystem::remove_all(dir); + std::filesystem::create_directories(rank_dir); + const std::string key = "c_fc.weight"; + const auto tensor = MakeColumnTensor({0, 1, 2, 3, 4, 5, 6, 7}); + Checkpoint::SaveStateDictFile(rank_dir / checkpoint::kModelCheckpointFilename, + {{key, tensor}}); + + Checkpoint::CheckpointMetadata metadata; + metadata.tensors.push_back({ + .key = key, + .dtype_str = "fp32", + .global_shape = {8, 1}, + .local_shape = {8, 1}, + .global_offset = {0, 0}, + .axis_fragmentations = {1, 1}, + .segments = + { + {.global_offset = 0, .local_offset = 0, .length = 4}, + {.global_offset = 4, .local_offset = 4, .length = 4}, + }, + .file = "rank_0/" + std::string(checkpoint::kModelCheckpointFilename), + .offset = SingleTensorDataOffset(key, 2), + .byte_size = tensor->SizeInBytes(), + }); + + ShardedStateDict target; + target.tensors[key] = { + .key = key, + .dtype = DataType::kFLOAT32, + .global_shape = {8, 1}, + .local_shape = {4, 1}, + .global_offset = {0, 0}, + .axis_fragmentations = {2, 1}, + .segments = + { + {.global_offset = 2, .local_offset = 0, .length = 2}, + {.global_offset = 6, .local_offset = 2, .length = 2}, + }, + }; + + const auto plan = checkpoint::LoadPlanner::PlanReshard(metadata, target); + const auto loaded = + checkpoint::IndexedRegionLoadStrategy().Execute(dir, plan).at(key); + const auto *data = static_cast(loaded->DataPtr()); + const std::vector expected = {2, 3, 6, 7}; + for (size_t row = 0; row < expected.size(); ++row) { + EXPECT_FLOAT_EQ(data[row], expected[row]); + } + std::filesystem::remove_all(dir); +} + +TEST(CheckpointLoadPlannerTest, PipelineReshardPlansOnlyTargetStageKeys) { + Checkpoint::CheckpointMetadata metadata; + metadata.tensors = { + MakeSavedShard("layer.0.weight", 1, 0, 8, "old_pp0/model.ckpt"), + MakeSavedShard("layer.1.weight", 1, 0, 8, "old_pp1/model.ckpt")}; + + auto plan = checkpoint::LoadPlanner::PlanReshard( + metadata, MakeTarget("layer.1.weight", 1, 0, 8)); + ASSERT_EQ(plan.tensors.size(), 1); + ASSERT_EQ(plan.tensors.at("layer.1.weight").reads.size(), 1); + EXPECT_EQ(plan.tensors.at("layer.1.weight").reads[0].filename, + "old_pp1/model.ckpt"); +} + +TEST(CheckpointLoadPlannerTest, VirtualPipelineReshardUsesGlobalLayerKeys) { + Checkpoint::CheckpointMetadata metadata; + for (int layer = 0; layer < 4; ++layer) { + const auto key = "transformer.h." + std::to_string(layer) + ".weight"; + metadata.tensors.push_back(MakeSavedShard( + key, 1, 0, 8, "saved_layer_" + std::to_string(layer) + "/model.ckpt")); + } + + auto target = MakeTarget("transformer.h.0.weight", 1, 0, 8); + target.Merge(MakeTarget("transformer.h.2.weight", 1, 0, 8)); + const auto plan = checkpoint::LoadPlanner::PlanReshard(metadata, target); + + ASSERT_EQ(plan.tensors.size(), 2); + EXPECT_EQ(plan.tensors.at("transformer.h.0.weight").reads[0].filename, + "saved_layer_0/model.ckpt"); + EXPECT_EQ(plan.tensors.at("transformer.h.2.weight").reads[0].filename, + "saved_layer_2/model.ckpt"); +} diff --git a/tests/checkpoint/test_lr_scheduler_state.cc b/tests/checkpoint/test_lr_scheduler_state.cc index 10c5b131..00edf0be 100644 --- a/tests/checkpoint/test_lr_scheduler_state.cc +++ b/tests/checkpoint/test_lr_scheduler_state.cc @@ -7,6 +7,7 @@ #include "gtest/gtest.h" #include "infini_train/include/checkpoint/checkpoint.h" +#include "infini_train/include/checkpoint/constants.h" #include "infini_train/include/datatype.h" #include "infini_train/include/device.h" #include "infini_train/include/lr_scheduler.h" @@ -69,7 +70,7 @@ TEST_P(LRSchedulerCheckpointTest, SaveAndLoadLRSchedulerState) { TrainerState saved{.global_step = 3, .consumed_train_samples = 12}; Checkpoint::Save(dir, *model1, nullptr, saved, sched1.get()); - EXPECT_TRUE(std::filesystem::exists(dir / "lr_scheduler.ckpt")); + EXPECT_TRUE(std::filesystem::exists(dir / checkpoint::kLRSchedulerFilename)); auto model2 = MakeModel(GetDevice()); auto opt2 = std::make_shared(model2->Parameters(), kBaseLR); @@ -101,7 +102,7 @@ TEST_P(LRSchedulerCheckpointTest, SkipsLRSchedulerStateWhenSchedulerIsNull) { TrainerState saved{.global_step = 3}; Checkpoint::Save(dir, *model1, nullptr, saved, nullptr); - EXPECT_FALSE(std::filesystem::exists(dir / "lr_scheduler.ckpt")); + EXPECT_FALSE(std::filesystem::exists(dir / checkpoint::kLRSchedulerFilename)); std::filesystem::remove_all(dir); } diff --git a/tests/checkpoint/test_optimizer_state.cc b/tests/checkpoint/test_optimizer_state.cc index 31561352..0b8af645 100644 --- a/tests/checkpoint/test_optimizer_state.cc +++ b/tests/checkpoint/test_optimizer_state.cc @@ -93,6 +93,27 @@ TEST_P(OptimizerStateTest, AdamStateDictRoundTrip) { } } +TEST_P(OptimizerStateTest, AdamStateDictUsesStableParameterNames) { + auto first = std::make_shared(std::vector{2, 2}, + DataType::kFLOAT32, GetDevice()); + auto second = std::make_shared(std::vector{3}, + DataType::kFLOAT32, GetDevice()); + const NamedParameterList named_parameters{{"transformer.h.0.weight", first}, + {"transformer.h.0.bias", second}}; + auto adam = std::make_shared(named_parameters, 0.001); + + const auto state = adam->StateDict(); + EXPECT_TRUE(state.contains("adam.m.transformer.h.0.weight")); + EXPECT_TRUE(state.contains("adam.v.transformer.h.0.weight")); + EXPECT_TRUE(state.contains("adam.m.transformer.h.0.bias")); + EXPECT_TRUE(state.contains("adam.v.transformer.h.0.bias")); + EXPECT_TRUE(state.contains("adam.t")); + + auto restored = std::make_shared(named_parameters, 0.001); + restored->LoadStateDict(state); + EXPECT_EQ(restored->StateDict().size(), state.size()); +} + // ---------- SGD ---------- TEST_P(OptimizerStateTest, SGDStateDictEmpty) { auto param = std::make_shared(std::vector{2, 2}, diff --git a/tests/checkpoint/test_trainer_state.cc b/tests/checkpoint/test_trainer_state.cc index c9b9b72b..94aab51d 100644 --- a/tests/checkpoint/test_trainer_state.cc +++ b/tests/checkpoint/test_trainer_state.cc @@ -11,8 +11,10 @@ #include "infini_train/include/checkpoint/checkpoint.h" #include "infini_train/include/checkpoint/checkpoint_manager.h" +#include "infini_train/include/checkpoint/constants.h" #include "infini_train/include/datatype.h" #include "infini_train/include/nn/modules/linear.h" +#include "infini_train/include/nn/parallel/global.h" #include "infini_train/include/optimizer.h" #include "infini_train/include/tensor.h" @@ -31,11 +33,13 @@ TEST_P(TrainerStateTest, DefaultValues) { EXPECT_EQ(state.n_head, 0); EXPECT_EQ(state.n_kv_head, 0); EXPECT_EQ(state.n_embd, 0); - EXPECT_EQ(state.vocab_size, 0); + EXPECT_EQ(state.original_vocab_size, 0); + EXPECT_EQ(state.padded_vocab_size, 0); EXPECT_EQ(state.ddp_size, 1); EXPECT_EQ(state.tp_size, 1); EXPECT_EQ(state.sp_size, 1); EXPECT_EQ(state.pp_size, 1); + EXPECT_EQ(state.vpp_size, 1); } TEST_P(TrainerStateTest, TrainerStateFileCreated) { @@ -53,9 +57,9 @@ TEST_P(TrainerStateTest, TrainerStateFileCreated) { Checkpoint::Save(dir, *model, opt.get(), saved, nullptr); - EXPECT_TRUE(std::filesystem::exists(dir / "trainer_state.json")); + EXPECT_TRUE(std::filesystem::exists(dir / checkpoint::kTrainerStateFilename)); - std::ifstream ifs(dir / "trainer_state.json"); + std::ifstream ifs(dir / checkpoint::kTrainerStateFilename); std::string content((std::istreambuf_iterator(ifs)), std::istreambuf_iterator()); EXPECT_NE(content.find("\"global_step\""), std::string::npos); @@ -75,11 +79,13 @@ TEST_P(TrainerStateTest, RoundTrip) { .n_head = 16, .n_kv_head = 8, .n_embd = 1024, - .vocab_size = 128256, + .original_vocab_size = 128000, + .padded_vocab_size = 128256, .ddp_size = 2, .tp_size = 1, .sp_size = 1, .pp_size = 2, + .vpp_size = 2, }; auto model1 = std::make_shared(1, 3, true, GetDevice()); @@ -107,9 +113,16 @@ TEST_P(TrainerStateTest, RoundTrip) { EXPECT_EQ(loaded.n_head, 16); EXPECT_EQ(loaded.n_kv_head, 8); EXPECT_EQ(loaded.n_embd, 1024); - EXPECT_EQ(loaded.vocab_size, 128256); - EXPECT_EQ(loaded.ddp_size, 2); - EXPECT_EQ(loaded.pp_size, 2); + EXPECT_EQ(loaded.original_vocab_size, 128000); + EXPECT_EQ(loaded.padded_vocab_size, 128256); + EXPECT_EQ(loaded.ddp_size, nn::parallel::global::GetDataParallelSize()); + EXPECT_EQ(loaded.tp_size, nn::parallel::global::GetTensorParallelSize()); + EXPECT_EQ(loaded.pp_size, nn::parallel::global::GetPipelineParallelSize()); + EXPECT_EQ(loaded.sp_size, nn::parallel::global::GetSequenceParallelEnabled() + ? nn::parallel::global::GetTensorParallelSize() + : 1); + EXPECT_EQ(loaded.vpp_size, + nn::parallel::global::GetVirtualPipelineParallelSize()); std::filesystem::remove_all(dir); } diff --git a/tests/lora/test_lora.cc b/tests/lora/test_lora.cc index c1f43730..f4cb51c8 100644 --- a/tests/lora/test_lora.cc +++ b/tests/lora/test_lora.cc @@ -1,6 +1,7 @@ #include #include +#include #include #include #include @@ -14,9 +15,13 @@ #include "infini_train/include/device.h" #include "infini_train/include/nn/lora/lora_config.h" #include "infini_train/include/nn/lora/lora_linear.h" +#include "infini_train/include/nn/lora/lora_parallel_linear.h" #include "infini_train/include/nn/lora/lora_utils.h" #include "infini_train/include/nn/modules/linear.h" #include "infini_train/include/nn/modules/module.h" +#include "infini_train/include/nn/parallel/global.h" +#include "infini_train/include/nn/parallel/tensor_parallel.h" +#include "infini_train/include/shard_spec.h" #include "infini_train/include/tensor.h" #include "tests/common/test_utils.h" @@ -93,6 +98,40 @@ TEST_P(LoRATest, LoRAConfigScaling) { EXPECT_EQ(config.Scaling(), expected_scaling); } +TEST_P(LoRATest, ParallelLoRAShardedStateDictIncludesAdapterParameters) { + LoRAConfig config; + config.rank = 2; + + auto column_base = std::make_shared( + 4, 6, /*bias=*/false, /*gather_output=*/false, + /*input_is_parallel=*/false, /*skip_bias_add=*/false, + /*sequence_parallel=*/false); + auto column = + std::make_shared(column_base, config, 4, 6); + const auto column_state = column->BuildShardedStateDict("column"); + ASSERT_TRUE(column_state.tensors.contains("column.lora_A")); + ASSERT_TRUE(column_state.tensors.contains("column.lora_B")); + EXPECT_EQ(column_state.tensors.at("column.lora_A").axis_fragmentations, + (std::vector{1, 1})); + EXPECT_EQ( + column_state.tensors.at("column.lora_B").axis_fragmentations, + (std::vector{nn::parallel::global::GetTensorParallelSize(), 1})); + + auto row_base = std::make_shared( + 4, 6, /*bias=*/false, /*reduce_output=*/true, /*input_is_parallel=*/true, + /*skip_bias_add=*/false, + /*sequence_parallel=*/false); + auto row = std::make_shared(row_base, config, 4, 6); + const auto row_state = row->BuildShardedStateDict("row"); + ASSERT_TRUE(row_state.tensors.contains("row.lora_A")); + ASSERT_TRUE(row_state.tensors.contains("row.lora_B")); + EXPECT_EQ( + row_state.tensors.at("row.lora_A").axis_fragmentations, + (std::vector{1, nn::parallel::global::GetTensorParallelSize()})); + EXPECT_EQ(row_state.tensors.at("row.lora_B").axis_fragmentations, + (std::vector{1, 1})); +} + TEST_P(LoRATest, PackedQKVShardGPTStyle) { auto full_qkv = MakeRowLabeledTensor(/*rows=*/12, /*cols=*/3, GetDevice()); auto shard = diff --git a/tests/transformer/test_transformer_architecture.cc b/tests/transformer/test_transformer_architecture.cc index 3075c9fc..ecf4e42f 100644 --- a/tests/transformer/test_transformer_architecture.cc +++ b/tests/transformer/test_transformer_architecture.cc @@ -1,9 +1,12 @@ #include +#include #include #include #include #include +#include +#include #include #include "gtest/gtest.h" @@ -20,6 +23,7 @@ #include "infini_train/include/nn/modules/transformer/transformer.h" #include "infini_train/include/nn/modules/transformer/transformer_config.h" #include "infini_train/include/nn/modules/transformer/utils.h" +#include "infini_train/include/shard_spec.h" #include "infini_train/include/tensor.h" #include "tests/common/test_utils.h" @@ -205,6 +209,18 @@ TEST_P(TransformerModuleTest, LLaMA3Model) { auto model = std::make_shared(config); model->To(GetDevice()); EXPECT_FALSE(model->Parameters().empty()); + const auto sharded_state = model->BuildShardedStateDict(); + std::unordered_set sharded_tensors; + for (const auto &[key, info] : sharded_state.tensors) { + if (info.local_tensor != nullptr) { + sharded_tensors.insert(info.local_tensor); + } + } + + for (const auto &[name, parameter] : model->NamedParameters()) { + EXPECT_TRUE(sharded_tensors.contains(parameter.get())) + << "Missing shard metadata for parameter: " << name; + } } TEST_P(TransformerModuleTest, RoPEUtils) {