Feat: checkpoint resharding

example/gpt2/checkpoint_loader.cc

@@ -187,32 +187,28 @@ std::shared_ptr<nn::TransformerModel> 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<float *>(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);
     }
   }
 
   // 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<float *>(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<nn::TransformerModel> 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,20 +256,18 @@ std::shared_ptr<nn::TransformerModel> 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);
     }
   }
 
   // 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<nn::TransformerModel> 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,69 +308,61 @@ std::shared_ptr<nn::TransformerModel> 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<float *>(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);
     }
   }
 
   // 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<float *>(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);
     }
   }
 
   // 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<float *>(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);
     }
   }
 
   // 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<float *>(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,35 +371,31 @@ std::shared_ptr<nn::TransformerModel> 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<float *>(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);
     }
   }
 
   // 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<float *>(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,34 +404,30 @@ std::shared_ptr<nn::TransformerModel> 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<float *>(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);
     }
   }
 
   // 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<float *>(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);

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<int>(FLAGS_virtual_pipeline_parallel),
         .checkpoint_root_dir = FLAGS_save,
         .max_checkpoint_keep = FLAGS_max_checkpoint_keep,
         .rank = rank,

example/llama3/checkpoint_loader.cc

@@ -208,15 +208,13 @@ std::shared_ptr<nn::TransformerModel> LoadFromLLMC(
   }
 
   // transformer.h.{i}.ln_1.weight : Full version nn::RMSNorm
-  int local_layer_index = 0;
   for (int i = 0; i < static_cast<int>(n_layer); ++i) {
     if (owned_layers[i]) {
       auto &tensor = state_dict[std::format(
           "{}.{}.{}.{}.{}", nn::TransformerModel::kTransformerModelName,
-          nn::TransformerChunk::kHLayerName, std::to_string(local_layer_index),
+          nn::TransformerChunk::kHLayerName, std::to_string(i),
           nn::TransformerLayer::kLn1LayerName, nn::RMSNorm::kParamWeightName)];
       ReadVectorAllFloat(ifs, static_cast<float *>(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<nn::TransformerModel> 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<int>(n_layer); ++i) {
     if (owned_layers[i]) {
       auto &tensor = state_dict[std::format(
           "{}.{}.{}.{}.{}.{}", nn::TransformerModel::kTransformerModelName,
-          nn::TransformerChunk::kHLayerName, std::to_string(local_layer_index),
+          nn::TransformerChunk::kHLayerName, std::to_string(i),
           nn::TransformerLayer::kAttnLayerName,
           nn::CausalSelfAttention::kCAttnLayerName,
           nn::parallel::ColumnParallelLinear::kParamWeightName)];
@@ -265,7 +262,6 @@ std::shared_ptr<nn::TransformerModel> 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<size_t>(attn_rows_all) * attn_cols * sizeof(float);
@@ -275,19 +271,17 @@ std::shared_ptr<nn::TransformerModel> 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<int>(n_layer); ++i) {
     if (owned_layers[i]) {
       auto &tensor = state_dict[std::format(
           "{}.{}.{}.{}.{}.{}", nn::TransformerModel::kTransformerModelName,
-          nn::TransformerChunk::kHLayerName, std::to_string(local_layer_index),
+          nn::TransformerChunk::kHLayerName, std::to_string(i),
           nn::TransformerLayer::kAttnLayerName,
           nn::CausalSelfAttention::kCProjLayerName,
           nn::parallel::RowParallelLinear::kParamWeightName)];
       ReadMatrixColShardFloat(ifs, static_cast<float *>(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<size_t>(n_embd) * n_embd * sizeof(float);
@@ -296,15 +290,13 @@ std::shared_ptr<nn::TransformerModel> LoadFromLLMC(
   }
 
   // transformer.h.{i}.ln_2.weight : Full version RMSNorm
-  local_layer_index = 0;
   for (int i = 0; i < static_cast<int>(n_layer); ++i) {
     if (owned_layers[i]) {
       auto &tensor = state_dict[std::format(
           "{}.{}.{}.{}.{}", nn::TransformerModel::kTransformerModelName,
-          nn::TransformerChunk::kHLayerName, std::to_string(local_layer_index),
+          nn::TransformerChunk::kHLayerName, std::to_string(i),
           nn::TransformerLayer::kLn2LayerName, nn::RMSNorm::kParamWeightName)];
       ReadVectorAllFloat(ifs, static_cast<float *>(tensor->DataPtr()), n_embd);
-      ++local_layer_index;
     } else {
       size_t ln_2_bytes = static_cast<size_t>(n_embd) * sizeof(float);
       ifs.seekg(ln_2_bytes, std::ios::cur);
@@ -313,19 +305,17 @@ std::shared_ptr<nn::TransformerModel> 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<int>(n_layer); ++i) {
     if (owned_layers[i]) {
       auto &tensor = state_dict[std::format(
           "{}.{}.{}.{}.{}.{}", nn::TransformerModel::kTransformerModelName,
-          nn::TransformerChunk::kHLayerName, std::to_string(local_layer_index),
+          nn::TransformerChunk::kHLayerName, std::to_string(i),
           nn::TransformerLayer::kMlpLayerName, nn::MLP::kCFcLayerName,
           nn::parallel::ColumnParallelLinear::kParamWeightName)];
       float *dst = static_cast<float *>(tensor->DataPtr()) + fc_pp * n_embd;
       ReadMatrixRowShardFloat(ifs, dst,
                               /*rows=*/fc_out, /*cols=*/n_embd,
                               /*row_start=*/tp_rank * fc_pp, /*row_cnt=*/fc_pp);
-      ++local_layer_index;
     } else {
       size_t fc_bytes =
           static_cast<size_t>(ffn_hidden) * n_embd * sizeof(float);
@@ -335,18 +325,16 @@ std::shared_ptr<nn::TransformerModel> 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<int>(n_layer); ++i) {
     if (owned_layers[i]) {
       auto &tensor = state_dict[std::format(
           "{}.{}.{}.{}.{}.{}", nn::TransformerModel::kTransformerModelName,
-          nn::TransformerChunk::kHLayerName, std::to_string(local_layer_index),
+          nn::TransformerChunk::kHLayerName, std::to_string(i),
           nn::TransformerLayer::kMlpLayerName, nn::MLP::kCFcLayerName,
           nn::parallel::ColumnParallelLinear::kParamWeightName)];
       ReadMatrixRowShardFloat(ifs, static_cast<float *>(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<size_t>(ffn_hidden) * n_embd * sizeof(float);
@@ -356,19 +344,17 @@ std::shared_ptr<nn::TransformerModel> 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<int>(n_layer); ++i) {
     if (owned_layers[i]) {
       auto &tensor = state_dict[std::format(
           "{}.{}.{}.{}.{}.{}", nn::TransformerModel::kTransformerModelName,
-          nn::TransformerChunk::kHLayerName, std::to_string(local_layer_index),
+          nn::TransformerChunk::kHLayerName, std::to_string(i),
           nn::TransformerLayer::kMlpLayerName, nn::MLP::kCProjLayerName,
           nn::parallel::RowParallelLinear::kParamWeightName)];
       ReadMatrixColShardFloat(ifs, static_cast<float *>(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<size_t>(n_embd) * ffn_hidden * sizeof(float);

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<int>(FLAGS_virtual_pipeline_parallel),
         .checkpoint_root_dir = FLAGS_save,
         .max_checkpoint_keep = FLAGS_max_checkpoint_keep,
         .rank = rank,