diff --git a/CMakeLists.txt b/CMakeLists.txt index 6bd8069d4..0cf33f6e6 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -222,6 +222,9 @@ link_infini_train_exe(llama3) add_subdirectory(tools/infini_run) set_target_properties(infini_run PROPERTIES RUNTIME_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}) +add_executable(pipeline_layout_suggest tools/pipeline_layout_suggest.cc) +target_link_libraries(pipeline_layout_suggest PRIVATE infini_train) + # Tests if(BUILD_TEST) add_subdirectory(tests) diff --git a/docs/pipeline_custom_layout_usage.md b/docs/pipeline_custom_layout_usage.md new file mode 100644 index 000000000..1a5f30c69 --- /dev/null +++ b/docs/pipeline_custom_layout_usage.md @@ -0,0 +1,164 @@ +# Pipeline 并行自定义布局使用说明 + +本功能把层和特殊模块的归属统一记录在 `PipelineLayout` 中。模型构建、 +`PipelineParallel` 包装器、调度器和 GPT-2/LLaMA3 checkpoint loader 都查询同一份 +布局，因此不会因为不同组件各自重新计算切分而产生不一致。 + +## 命令行参数 + +GPT-2 和 LLaMA3 都支持： + +```text +--pipeline_parallel=N +--virtual_pipeline_parallel=N +--pipeline_layer_partition=4,8,6,6 +--pipeline_layout='Etttttt|ttttttFH' +--pipeline_embedding_stage=0 +--pipeline_final_norm_stage=-1 +--pipeline_lm_head_stage=-1 +``` + +`-1` 表示最后一个 pipeline stage。特殊模块参数默认保持旧行为：embedding 在 +stage 0，final norm 和 lm head 在最后一个 stage。 + +不提供 `--pipeline_layer_partition` 时，使用 `BuildDefault` 的均匀布局。chunk 按 +`global_chunk_id = local_chunk_id * num_stages + stage_id` 交错编号，以保持现有 +GPipe/1F1B 调度编号兼容。 + +提供 partition 时，每个逗号分隔的整数表示一个 stage 拥有的连续层数。例如 +24 层模型使用 `--pipeline_parallel=4 --pipeline_layer_partition=4,8,6,6`： + +```text +stage 0: layers [0,4) +stage 1: layers [4,12) +stage 2: layers [12,18) +stage 3: layers [18,24) +``` + +当前连续 `pipeline_layer_partition` 必须使用 `--virtual_pipeline_parallel=1`。 +需要 vPP 的非对称 chunk 映射时使用 Megatron 风格布局字符串，或在 C++ 集成中 +调用显式 chunk 构造接口。 + +### Megatron 风格布局（增强功能） + +`--pipeline_layout` 可显式描述每个 stage 的 chunk 和特殊模块：`t` 表示 +Transformer 层，`E`/`F`/`H` 分别表示 embedding、final norm、LM head，`|` +分隔 stage，`,` 分隔同一 stage 内的 vPP chunk，括号后使用 `*N` 重复。 +例如 `--pipeline_parallel=2 --virtual_pipeline_parallel=2 +--pipeline_layout='tt,tt|tt,tt'` 将 8 层映射为 4 个显式 chunk。该参数与 +`--pipeline_layer_partition` 互斥（同时传入会在启动期报错）；每个 stage 必须恰好给出 `vpp` 个 chunk， +所有 `t` 的总数必须等于模型层数。布局解析和校验在训练启动前完成。 + +如果需要在 C++ 集成中完全控制每一个 chunk 的 stage 归属，可调用 +`PipelineLayout::BuildExplicit(num_layers, pp_size, vpp_size, chunks)`。该接口要求 +global chunk id 连续，但以 `ChunkLayout.stage_id` 和 `local_chunk_id` 为权威，不从 +global chunk id 反推 stage；因此可以表达非对称的 vPP chunk 大小和显式 ownership。 + +显式布局也允许空 chunk/空 stage。当前 Pipeline transport 会让空 +`TransformerChunk` 原样传递激活和梯度，因此只要 layout 中保留该 stage 的 +chunk id，就可以用于真实 PP 运行。例如： + +```bash +--pipeline_parallel=3 +--pipeline_layout='Etttttt||ttttttFH' +``` + +### 自动负载均衡建议与性能对比 + +可使用 `pipeline_layout_suggest` 根据每层代价生成连续 partition： + +```bash +./build_cpu/pipeline_layout_suggest \ + --num_layers=8 --pp_size=4 --layer_cost=1,1,1,1,4,1,1,1 +# pipeline_layer_partition=3,1,1,3 +``` + +不传 `--layer_cost` 时默认每层代价为 1。工具会保证层数总和正确、每个 +stage 至少有一层，并倾向于将高代价层单独切分。 + +可用性能脚本对默认单卡和双卡显式布局做一次可重复的单步对比： + +```bash +BUILD_DIR=build_cuda scripts/compare_pipeline_layout_perf.sh gpt2 cuda +``` + +脚本在 `targets/pipeline_layout_perf.csv` 写出 `elapsed_ms`、`tok_per_s` +和理想 GPipe bubble 比例。`BATCH_SIZE`、`SEQUENCE_LENGTH`、 +`TOTAL_BATCH_SIZE` 可通过环境变量覆盖；bubble 使用 +`(pipeline_parallel-1)/(micro_batches+pipeline_parallel-1)` 估算，属于 +调度上界，不代替真实 profiler。 + +如需在相同 PP degree 下比较均匀布局和 cost-aware 自定义布局，并记录每个 stage +的真实计算时间，可运行： + +```bash +BUILD_DIR=build_cuda \ +LAUNCHER='torchrun --standalone --no-python' \ +TOTAL_BATCH_SIZE=128 \ +scripts/compare_pipeline_stage_perf.sh gpt2 cuda +``` + +`INFINI_PIPELINE_STAGE_TIMING=1` 会在每个 stage/chunk 的 forward、backward 完成 +CUDA stream 同步后输出 `pipeline_stage_timing` 记录；脚本将这些记录汇总到 +`targets/pipeline_stage_perf.csv`，并同时写出端到端 elapsed、吞吐、理想 bubble +和 stage imbalance ratio。 + +均衡建议工具支持直接读取 profiler/用户代价 CSV。文件可以是一列 cost，也可以 +是带 `layer,cost` 表头的两列 CSV： + +```bash +./build_cuda/pipeline_layout_suggest \ + --num_layers=8 --pp_size=4 \ + --layer_cost_file=layer_costs.csv +``` + +输出除了 `pipeline_layer_partition`，还会给出均匀 partition 与建议 partition +的最大 stage 代价及预计降低比例。 + +## 特殊模块和 checkpoint + +布局查询决定 `transformer.wte/wpe`、`transformer.ln_f` 和 +`transformer.lm_head` 在哪个 stage 注册和加载。层参数仍使用 canonical key： +`transformer.h....`。checkpoint loader 即使当前 rank 不拥有某个参数， +也会继续 seek 对应字节，避免后续权重错位。 + +当前 pipeline transport 的既有语义仍按 rank 0 为首 stage、最后 rank 为末 stage。 +因此把 final norm/lm head 放到非末 stage 虽可由布局描述，但真实 loss/target 的跨 +stage transport 尚未完全重构；部署前应保持它们在最后 stage，或先扩展 transport。 + +## 错误排查 + +布局错误使用可捕获的 `PipelineLayoutError`，不会用新的 glog `CHECK` 替代： + +```text +stage count mismatch: pipeline_parallel=4, partition_entries=3 +layer sum mismatch: num_layers=23, partition_sum=24 +pipeline_layer_partition contains invalid character '-' +custom layer partition is not supported with vpp_size=2 +``` + +普通空模块/非法指针参数使用 `std::invalid_argument`；训练代码中原有的 +`CHECK`/`LOG(FATAL)` 仍保留其原语义。 + +## 校验命令 + +CPU 构建和单测： + +```bash +cmake -S . -B build_cpu -DBUILD_TEST=ON -DUSE_CUDA=OFF -DUSE_NCCL=OFF +cmake --build build_cpu -j2 +ctest --test-dir build_cpu -R 'test_pipeline_(layout|scheduler_layout|parallel_chunking)_cpu' --output-on-failure +./build_cpu/tests/transformer/test_transformer_cpu --gtest_filter='TransformerPipelineLayoutTest.*' --gtest_color=no +``` + +有完整数据和 launcher 时，可以运行一迭代 smoke test（这会实际执行训练步骤）： + +```bash +./scripts/test_pipeline_custom_layout.sh gpt2 cuda +./scripts/test_pipeline_custom_layout.sh llama3 cuda +``` + +脚本默认查找 `build/`、`data//` 下的可执行文件和输入；可用 +`BUILD_DIR`、`INPUT_BIN`、`INPUT_VAL_BIN`、`TOKENIZER_BIN`、`LLMC_FILE`、 +`OUT_DIR`、`LAUNCHER` 覆盖。`LAUNCHER=direct` 可在单进程环境执行，但不能验证 +两 stage 的真实通信。没有数据、CUDA 或 launcher 时脚本会明确退出，不会静默通过。 diff --git a/example/gpt2/checkpoint_loader.cc b/example/gpt2/checkpoint_loader.cc index 95e54730b..0b500b8c4 100644 --- a/example/gpt2/checkpoint_loader.cc +++ b/example/gpt2/checkpoint_loader.cc @@ -10,6 +10,7 @@ #include #include +#include "gflags/gflags.h" #include "glog/logging.h" #include "infini_train/include/nn/modules/normalization.h" @@ -28,6 +29,12 @@ using namespace infini_train; namespace nn = infini_train::nn; +DECLARE_string(pipeline_layer_partition); +DECLARE_string(pipeline_layout); +DECLARE_int32(pipeline_embedding_stage); +DECLARE_int32(pipeline_final_norm_stage); +DECLARE_int32(pipeline_lm_head_stage); + namespace { constexpr int kRandomSeed = 42; @@ -89,6 +96,23 @@ std::shared_ptr LoadFromLLMC(const std::string &filepath) gpt2_config.n_head = n_head; gpt2_config.n_embd = n_embd; gpt2::SanitizeGPT2Config(gpt2_config); + nn::parallel::SpecialModulePlacement placement{ + .embedding_stage = FLAGS_pipeline_embedding_stage, + .final_norm_stage = FLAGS_pipeline_final_norm_stage, + .lm_head_stage = FLAGS_pipeline_lm_head_stage, + }; + auto layout = FLAGS_pipeline_layout.empty() + ? nn::parallel::PipelineLayout::BuildPipelineLayout( + static_cast(n_layer), + nn::parallel::global::GetPipelineParallelSize(), + nn::parallel::global::GetVirtualPipelineParallelSize(), + FLAGS_pipeline_layer_partition, placement) + : nn::parallel::PipelineLayout::ParseMegatronStyleLayout( + FLAGS_pipeline_layout, static_cast(n_layer), + nn::parallel::global::GetPipelineParallelSize(), + nn::parallel::global::GetVirtualPipelineParallelSize(), placement); + layout.ValidateForCurrentPipelineTransport(); + nn::parallel::global::InstallPipelineLayout(layout); auto local_gpt2 = std::make_shared(gpt2_config); LOG(INFO) << "magic: " << magic << " version: " << version << " block_size: " << block_size @@ -100,15 +124,15 @@ std::shared_ptr LoadFromLLMC(const std::string &filepath) CHECK_EQ(n_head % tp_size, 0) << "n_head must be divisible by TP world size."; // ========== pp_size：num_stages; vpp_size: num_chunks_per_stage ========== - int pp_size = nn::parallel::global::GetPipelineParallelSize(); - int vpp_size = nn::parallel::global::GetVirtualPipelineParallelSize(); auto pp_rank = nn::parallel::pp_rank; - auto [is_first_stage, is_last_stage, layer_ranges_per_chunk] - = nn::parallel::PipelineParallel::GetStageInfo(n_layer, pp_size, pp_rank, vpp_size); - // ========== layer to chunk ========== + const auto &installed_layout = nn::parallel::global::GetPipelineLayout(); + const bool owns_embedding = installed_layout.owns(nn::parallel::SpecialModule::kEmbedding, pp_rank); + const bool owns_final_norm = installed_layout.owns(nn::parallel::SpecialModule::kFinalNorm, pp_rank); + const bool owns_lm_head = installed_layout.owns(nn::parallel::SpecialModule::kLMHead, pp_rank); + std::vector owned_layers(n_layer, false); - for (const auto &[start, end] : layer_ranges_per_chunk) { - for (int i = start; i < end; ++i) { owned_layers[i] = true; } + for (int layer = 0; layer < static_cast(n_layer); ++layer) { + owned_layers[layer] = installed_layout.stage_of_layer(layer) == pp_rank; } auto tp_rank = nn::parallel::tp_rank; @@ -130,14 +154,15 @@ std::shared_ptr LoadFromLLMC(const std::string &filepath) // transformer.wte.weight (also transformer.lm_head.weight) // full: (model_vocab_size, n_embd) // local: (vocab_size_per_partition, n_embd) - if (is_first_stage) { + if (owns_embedding) { auto &transformer_wte_weight = state_dict[std::format("{}.{}.{}", nn::TransformerModel::kTransformerModelName, nn::TransformerFirstStage::kWTELayerName, nn::parallel::VocabParallelEmbedding::kParamWeightName)]; ReadMatrixRowShardFloat(ifs, static_cast(transformer_wte_weight->DataPtr()), model_vocab_size, n_embd, v_start, vpp); - } else if (pp_size > 1 && is_last_stage) { - auto &lm_head_weight = state_dict[std::format("{}.{}", nn::TransformerLastStage::kLMHeadLayerName, + } else if (nn::parallel::global::GetPipelineParallelSize() > 1 && owns_lm_head) { + auto &lm_head_weight = state_dict[std::format("{}.{}.{}", nn::TransformerModel::kTransformerModelName, + nn::TransformerLastStage::kLMHeadLayerName, nn::parallel::ColumnParallelLinear::kParamWeightName)]; ReadMatrixRowShardFloat(ifs, static_cast(lm_head_weight->DataPtr()), model_vocab_size, n_embd, v_start, vpp); @@ -151,7 +176,7 @@ std::shared_ptr LoadFromLLMC(const std::string &filepath) ifs.ignore((padded_vocab_size - model_vocab_size) * n_embd * sizeof(float)); } - if (is_first_stage) { + if (owns_embedding) { // transformer.wpe.weight auto &transformer_wpe_weight = state_dict[std::format("{}.{}.{}", nn::TransformerModel::kTransformerModelName, @@ -407,7 +432,7 @@ std::shared_ptr LoadFromLLMC(const std::string &filepath) } } - if (is_last_stage) { + if (owns_final_norm) { // transformer.ln_f.weight auto &transformer_ln_f_weight = state_dict[std::format("{}.{}.{}", nn::TransformerModel::kTransformerModelName, diff --git a/example/gpt2/main.cc b/example/gpt2/main.cc index 2551880e7..a87677f31 100644 --- a/example/gpt2/main.cc +++ b/example/gpt2/main.cc @@ -5,6 +5,9 @@ #include #include #include +#include +#include +#include #include #include "gflags/gflags.h" @@ -85,6 +88,11 @@ DEFINE_uint32(tensor_parallel, 1, "Tensor Parallel world size"); DEFINE_bool(sequence_parallel, false, "Whether to enable Sequence Parallel"); DEFINE_uint32(pipeline_parallel, 1, "Pipeline Parallel world size, specified the number of PP stages."); DEFINE_uint32(virtual_pipeline_parallel, 1, "Number of chunks in PP stage."); +DEFINE_string(pipeline_layer_partition, "", "Comma-separated layer counts per PP stage, e.g. 4,8,6,6"); +DEFINE_string(pipeline_layout, "", "Megatron-style pipeline layout expression, e.g. Ett|tttt|tttFH"); +DEFINE_int32(pipeline_embedding_stage, 0, "Stage that owns embedding."); +DEFINE_int32(pipeline_final_norm_stage, -1, "Stage that owns final norm; -1 means last PP stage."); +DEFINE_int32(pipeline_lm_head_stage, -1, "Stage that owns LM head; -1 means last PP stage."); // precision DEFINE_string(dtype, "float32", "precision used in training (float32/bfloat16)"); @@ -138,6 +146,72 @@ DEFINE_validator(lr_decay_style, void Train(const nn::parallel::Rank &rank) { using namespace nn::parallel; + if (FLAGS_pipeline_parallel == 0 || FLAGS_virtual_pipeline_parallel == 0) { + throw PipelineLayoutError(std::format( + "pipeline_parallel and virtual_pipeline_parallel must be positive (got {}, {})", + FLAGS_pipeline_parallel, FLAGS_virtual_pipeline_parallel)); + } + const auto check_stage_flag = [](const char *name, int stage) { + if (stage < -1 || stage >= static_cast(FLAGS_pipeline_parallel)) { + throw PipelineLayoutError(std::format( + "{}={} is outside valid stage range [-1, {})", name, stage, FLAGS_pipeline_parallel)); + } + }; + check_stage_flag("pipeline_embedding_stage", FLAGS_pipeline_embedding_stage); + check_stage_flag("pipeline_final_norm_stage", FLAGS_pipeline_final_norm_stage); + check_stage_flag("pipeline_lm_head_stage", FLAGS_pipeline_lm_head_stage); + + if (!FLAGS_pipeline_layer_partition.empty() && !FLAGS_pipeline_layout.empty()) { + throw PipelineLayoutError("pipeline_layer_partition and pipeline_layout are mutually exclusive"); + } + + // Validate and install the layout before constructing any model module. + // Checkpoint-backed runs repeat this after reading the authoritative + // layer count from the checkpoint header. + if (!FLAGS_pipeline_layer_partition.empty()) { + const auto partition = PipelineLayout::ParseLayerPartition(FLAGS_pipeline_layer_partition); + if (static_cast(partition.size()) != static_cast(FLAGS_pipeline_parallel)) { + throw PipelineLayoutError(std::format( + "stage count mismatch: pipeline_parallel={}, partition_entries={}", + FLAGS_pipeline_parallel, partition.size())); + } + if (FLAGS_virtual_pipeline_parallel != 1) { + throw PipelineLayoutError(std::format( + "custom layer partition is not supported with vpp_size={}", + FLAGS_virtual_pipeline_parallel)); + } + } + + if (FLAGS_llmc_filepath.empty()) { + auto preflight_config = gpt2::GPT2Config(); + if (kModelToConfigs.count(FLAGS_model)) { + preflight_config = kModelToConfigs.at(FLAGS_model); + } + gpt2::SanitizeGPT2Config(preflight_config); + SpecialModulePlacement placement{ + .embedding_stage = FLAGS_pipeline_embedding_stage, + .final_norm_stage = FLAGS_pipeline_final_norm_stage, + .lm_head_stage = FLAGS_pipeline_lm_head_stage, + }; + auto preflight_layout = FLAGS_pipeline_layout.empty() + ? PipelineLayout::BuildPipelineLayout( + preflight_config.n_layer, + static_cast(FLAGS_pipeline_parallel), + static_cast(FLAGS_virtual_pipeline_parallel), + FLAGS_pipeline_layer_partition, + placement) + : PipelineLayout::ParseMegatronStyleLayout( + FLAGS_pipeline_layout, + preflight_config.n_layer, + static_cast(FLAGS_pipeline_parallel), + static_cast(FLAGS_virtual_pipeline_parallel), + placement); + global::InstallPipelineLayout(preflight_layout); + if (rank.IsMainRank()) { + LOG(INFO) << preflight_layout.ToString(); + } + } + { if (rank.IsLastRank()) { if (!FLAGS_save.empty() && FLAGS_save_interval == 0) { @@ -568,6 +642,10 @@ void Train(const nn::parallel::Rank &rank) { int main(int argc, char *argv[]) { gflags::ParseCommandLineFlags(&argc, &argv, true); google::InitGoogleLogging(argv[0]); + if (const char *stage_timing = std::getenv("INFINI_PIPELINE_STAGE_TIMING"); + stage_timing != nullptr && stage_timing[0] == '1' && stage_timing[1] == '\0') { + google::LogToStderr(); + } auto precision_config = utils::PrecisionCheckConfig::Parse(FLAGS_precision_check); nn::parallel::global::InitAllEnv(FLAGS_nthread_per_process, FLAGS_tensor_parallel, FLAGS_sequence_parallel, @@ -588,7 +666,15 @@ int main(int argc, char *argv[]) { } else { nn::parallel::Rank rank(nn::parallel::global::GetGlobalProcRank(), 0, nn::parallel::global::GetNprocPerNode(), FLAGS_nthread_per_process); - Train(rank); + try { + Train(rank); + } catch (const std::out_of_range &e) { + void *frames[64]; + int n = backtrace(frames, 64); + LOG(ERROR) << "out_of_range during GPT-2 Train initialization: " << e.what(); + backtrace_symbols_fd(frames, n, STDERR_FILENO); + return 1; + } } gflags::ShutDownCommandLineFlags(); diff --git a/example/llama3/checkpoint_loader.cc b/example/llama3/checkpoint_loader.cc index f3590af6e..046122b7e 100644 --- a/example/llama3/checkpoint_loader.cc +++ b/example/llama3/checkpoint_loader.cc @@ -10,6 +10,7 @@ #include #include +#include "gflags/gflags.h" #include "glog/logging.h" #include "infini_train/include/nn/modules/normalization.h" @@ -17,6 +18,7 @@ #include "infini_train/include/nn/modules/transformer/mlp.h" #include "infini_train/include/nn/modules/transformer/transformer.h" #include "infini_train/include/nn/parallel/global.h" +#include "infini_train/include/nn/parallel/pp/pipeline_parallel.h" #include "infini_train/include/nn/parallel/tensor_parallel.h" #include "infini_train/include/tensor.h" @@ -26,6 +28,12 @@ using namespace infini_train; namespace nn = infini_train::nn; +DECLARE_string(pipeline_layer_partition); +DECLARE_string(pipeline_layout); +DECLARE_int32(pipeline_embedding_stage); +DECLARE_int32(pipeline_final_norm_stage); +DECLARE_int32(pipeline_lm_head_stage); + namespace { constexpr int kRandomSeed = 42; @@ -82,18 +90,35 @@ std::shared_ptr LoadFromLLMC(const std::string &filepath) llama3_config.norm_eps = norm_eps; llama3_config.max_gen_batch_size = max_gen_bs; llama3::SanitizeLLaMA3Config(llama3_config); + nn::parallel::SpecialModulePlacement placement{ + .embedding_stage = FLAGS_pipeline_embedding_stage, + .final_norm_stage = FLAGS_pipeline_final_norm_stage, + .lm_head_stage = FLAGS_pipeline_lm_head_stage, + }; + auto layout = FLAGS_pipeline_layout.empty() + ? nn::parallel::PipelineLayout::BuildPipelineLayout( + static_cast(n_layer), + nn::parallel::global::GetPipelineParallelSize(), + nn::parallel::global::GetVirtualPipelineParallelSize(), + FLAGS_pipeline_layer_partition, placement) + : nn::parallel::PipelineLayout::ParseMegatronStyleLayout( + FLAGS_pipeline_layout, static_cast(n_layer), + nn::parallel::global::GetPipelineParallelSize(), + nn::parallel::global::GetVirtualPipelineParallelSize(), placement); + layout.ValidateForCurrentPipelineTransport(); + nn::parallel::global::InstallPipelineLayout(layout); auto llama3 = std::make_shared(llama3_config); // ========== pp_size：num_stages; vpp_size: num_chunks_per_stage ========== - int pp_size = nn::parallel::global::GetPipelineParallelSize(); - int vpp_size = nn::parallel::global::GetVirtualPipelineParallelSize(); auto pp_rank = nn::parallel::pp_rank; - auto [is_first_stage, is_last_stage, layer_ranges_per_chunk] - = nn::parallel::PipelineParallel::GetStageInfo(n_layer, pp_size, pp_rank, vpp_size); - // ========== layer to chunk ========== + const auto &installed_layout = nn::parallel::global::GetPipelineLayout(); + const bool owns_embedding = installed_layout.owns(nn::parallel::SpecialModule::kEmbedding, pp_rank); + const bool owns_final_norm = installed_layout.owns(nn::parallel::SpecialModule::kFinalNorm, pp_rank); + const bool owns_lm_head = installed_layout.owns(nn::parallel::SpecialModule::kLMHead, pp_rank); + std::vector owned_layers(n_layer, false); - for (const auto &[start, end] : layer_ranges_per_chunk) { - for (int i = start; i < end; ++i) { owned_layers[i] = true; } + for (int layer = 0; layer < static_cast(n_layer); ++layer) { + owned_layers[layer] = installed_layout.stage_of_layer(layer) == pp_rank; } const int tp_size = nn::parallel::global::GetTensorParallelSize(); @@ -122,9 +147,10 @@ std::shared_ptr LoadFromLLMC(const std::string &filepath) LOG(INFO) << " version_minor = " << version_minor; LOG(INFO) << "Pipeline Parallel Chunks:"; - for (size_t i = 0; i < layer_ranges_per_chunk.size(); ++i) { - LOG(INFO) << " Chunk " << i << ": layers " << layer_ranges_per_chunk[i].first << " to " - << layer_ranges_per_chunk[i].second; + const auto &local_stage = installed_layout.stage(pp_rank); + for (size_t i = 0; i < local_stage.global_chunk_ids.size(); ++i) { + const auto &chunk = installed_layout.chunk(local_stage.global_chunk_ids[i]); + LOG(INFO) << " Chunk " << i << ": layers " << chunk.layers.start << " to " << chunk.layers.end; } } @@ -168,7 +194,7 @@ std::shared_ptr LoadFromLLMC(const std::string &filepath) // ========== Read Sharded Params ========== // transformer.wte.weight : (vocab_size, n_embd) -> local tp_rank: rows of [v_start : v_start+vpp) - if (is_first_stage) { + if (owns_embedding) { auto &wte = state_dict[std::format("{}.{}.{}", nn::TransformerModel::kTransformerModelName, nn::TransformerFirstStage::kWTELayerName, nn::parallel::VocabParallelEmbedding::kParamWeightName)]; @@ -326,16 +352,25 @@ std::shared_ptr LoadFromLLMC(const std::string &filepath) // transformer.ln_f.weight : Full version nn::RMSNorm // lm_head.weight : (vocab_size, n_embd) -> ColumnParallelLinear, but actually applies on "rows" { - if (is_last_stage) { + if (owns_final_norm || owns_lm_head) { auto &ln_f = state_dict[std::format("{}.{}.{}", nn::TransformerModel::kTransformerModelName, nn::TransformerLastStage::kLnFLayerName, nn::RMSNorm::kParamWeightName)]; - auto &lm_head = state_dict[std::format("{}.{}", nn::TransformerLastStage::kLMHeadLayerName, - nn::parallel::ColumnParallelLinear::kParamWeightName)]; - ReadVectorAllFloat(ifs, static_cast(ln_f->DataPtr()), n_embd); - ReadMatrixRowShardFloat(ifs, static_cast(lm_head->DataPtr()), - /*rows=*/vocab_size, /*cols=*/n_embd, - /*row_start=*/v_start, /*row_cnt=*/vpp); + if (owns_final_norm) { + ReadVectorAllFloat(ifs, static_cast(ln_f->DataPtr()), n_embd); + } else { + ifs.seekg(static_cast(n_embd) * sizeof(float), std::ios::cur); + } + if (owns_lm_head) { + auto &lm_head = state_dict[std::format("{}.{}.{}", nn::TransformerModel::kTransformerModelName, + nn::TransformerLastStage::kLMHeadLayerName, + nn::parallel::ColumnParallelLinear::kParamWeightName)]; + ReadMatrixRowShardFloat(ifs, static_cast(lm_head->DataPtr()), + /*rows=*/vocab_size, /*cols=*/n_embd, + /*row_start=*/v_start, /*row_cnt=*/vpp); + } else { + ifs.seekg(static_cast(vocab_size) * n_embd * sizeof(float), std::ios::cur); + } } else { size_t ln_f_bytes = static_cast(n_embd) * sizeof(float); size_t lm_head_bytes = static_cast(vocab_size) * n_embd * sizeof(float); diff --git a/example/llama3/main.cc b/example/llama3/main.cc index ccfca86a2..541d09793 100644 --- a/example/llama3/main.cc +++ b/example/llama3/main.cc @@ -84,6 +84,11 @@ DEFINE_uint32(tensor_parallel, 1, "Tensor Parallel world size"); DEFINE_bool(sequence_parallel, false, "Whether to enable Sequence Parallel"); DEFINE_uint32(pipeline_parallel, 1, "Pipeline Parallel world size, specified the number of PP stages."); DEFINE_uint32(virtual_pipeline_parallel, 1, "Number of chunks in PP stage."); +DEFINE_string(pipeline_layer_partition, "", "Comma-separated layer counts per PP stage, e.g. 4,8,6,6"); +DEFINE_string(pipeline_layout, "", "Megatron-style pipeline layout expression, e.g. Ett|tttt|tttFH"); +DEFINE_int32(pipeline_embedding_stage, 0, "Stage that owns embedding."); +DEFINE_int32(pipeline_final_norm_stage, -1, "Stage that owns final norm; -1 means last PP stage."); +DEFINE_int32(pipeline_lm_head_stage, -1, "Stage that owns LM head; -1 means last PP stage."); // precision DEFINE_string(dtype, "float32", "precision used in training (float32/bfloat16)"); DEFINE_uint32(save_interval, 0, "save checkpoint every N steps; 0 disables saving"); @@ -127,6 +132,66 @@ DEFINE_validator(lr_decay_style, void Train(const nn::parallel::Rank &rank) { using namespace nn::parallel; + if (FLAGS_pipeline_parallel == 0 || FLAGS_virtual_pipeline_parallel == 0) { + throw PipelineLayoutError(std::format( + "pipeline_parallel and virtual_pipeline_parallel must be positive (got {}, {})", + FLAGS_pipeline_parallel, FLAGS_virtual_pipeline_parallel)); + } + const auto check_stage_flag = [](const char *name, int stage) { + if (stage < -1 || stage >= static_cast(FLAGS_pipeline_parallel)) { + throw PipelineLayoutError(std::format( + "{}={} is outside valid stage range [-1, {})", name, stage, FLAGS_pipeline_parallel)); + } + }; + check_stage_flag("pipeline_embedding_stage", FLAGS_pipeline_embedding_stage); + check_stage_flag("pipeline_final_norm_stage", FLAGS_pipeline_final_norm_stage); + check_stage_flag("pipeline_lm_head_stage", FLAGS_pipeline_lm_head_stage); + + if (!FLAGS_pipeline_layer_partition.empty() && !FLAGS_pipeline_layout.empty()) { + throw PipelineLayoutError("pipeline_layer_partition and pipeline_layout are mutually exclusive"); + } + + if (!FLAGS_pipeline_layer_partition.empty()) { + const auto partition = PipelineLayout::ParseLayerPartition(FLAGS_pipeline_layer_partition); + if (static_cast(partition.size()) != static_cast(FLAGS_pipeline_parallel)) { + throw PipelineLayoutError(std::format( + "stage count mismatch: pipeline_parallel={}, partition_entries={}", + FLAGS_pipeline_parallel, partition.size())); + } + if (FLAGS_virtual_pipeline_parallel != 1) { + throw PipelineLayoutError(std::format( + "custom layer partition is not supported with vpp_size={}", + FLAGS_virtual_pipeline_parallel)); + } + } + + if (FLAGS_llmc_filepath.empty()) { + auto preflight_config = llama3::LLaMA3Config(); + llama3::SanitizeLLaMA3Config(preflight_config); + SpecialModulePlacement placement{ + .embedding_stage = FLAGS_pipeline_embedding_stage, + .final_norm_stage = FLAGS_pipeline_final_norm_stage, + .lm_head_stage = FLAGS_pipeline_lm_head_stage, + }; + auto preflight_layout = FLAGS_pipeline_layout.empty() + ? PipelineLayout::BuildPipelineLayout( + preflight_config.n_layer, + static_cast(FLAGS_pipeline_parallel), + static_cast(FLAGS_virtual_pipeline_parallel), + FLAGS_pipeline_layer_partition, + placement) + : PipelineLayout::ParseMegatronStyleLayout( + FLAGS_pipeline_layout, + preflight_config.n_layer, + static_cast(FLAGS_pipeline_parallel), + static_cast(FLAGS_virtual_pipeline_parallel), + placement); + global::InstallPipelineLayout(preflight_layout); + if (rank.IsMainRank()) { + LOG(INFO) << preflight_layout.ToString(); + } + } + { if (rank.IsLastRank()) { if (!FLAGS_save.empty() && FLAGS_save_interval == 0) { @@ -547,6 +612,10 @@ void Train(const nn::parallel::Rank &rank) { int main(int argc, char *argv[]) { gflags::ParseCommandLineFlags(&argc, &argv, true); google::InitGoogleLogging(argv[0]); + if (const char *stage_timing = std::getenv("INFINI_PIPELINE_STAGE_TIMING"); + stage_timing != nullptr && stage_timing[0] == '1' && stage_timing[1] == '\0') { + google::LogToStderr(); + } auto precision_config = utils::PrecisionCheckConfig::Parse(FLAGS_precision_check); nn::parallel::global::InitAllEnv(FLAGS_nthread_per_process, FLAGS_tensor_parallel, FLAGS_sequence_parallel, diff --git a/infini_train/include/nn/modules/transformer/transformer.h b/infini_train/include/nn/modules/transformer/transformer.h index 0471c32fe..65c7d73c8 100644 --- a/infini_train/include/nn/modules/transformer/transformer.h +++ b/infini_train/include/nn/modules/transformer/transformer.h @@ -4,7 +4,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/nn/parallel/global.h" namespace infini_train::nn { class TransformerLayer : public CloneableModule { @@ -57,13 +57,15 @@ class TransformerLastStage : public CloneableModule { static constexpr char kLnFLayerName[] = "ln_f"; static constexpr char kLMHeadLayerName[] = "lm_head"; - explicit TransformerLastStage(const TransformerConfig &config); + TransformerLastStage(const TransformerConfig &config, bool has_final_norm, bool has_lm_head); std::vector> Forward(const std::vector> &x) override; private: const TransformerConfig config_; + bool has_final_norm_ = false; + bool has_lm_head_ = false; }; class TransformerModel : public CloneableModule { @@ -80,7 +82,7 @@ class TransformerModel : public CloneableModule { private: const TransformerConfig config_; - const infini_train::nn::parallel::StageInfo stage_info_; + int num_local_chunks_ = 0; }; } // namespace infini_train::nn diff --git a/infini_train/include/nn/parallel/global.h b/infini_train/include/nn/parallel/global.h index 38694a91d..aeabfe78f 100644 --- a/infini_train/include/nn/parallel/global.h +++ b/infini_train/include/nn/parallel/global.h @@ -1,5 +1,7 @@ #pragma once +#include "infini_train/include/nn/parallel/pipeline_layout.h" + #include #include #include @@ -57,6 +59,14 @@ class GlobalEnv { Layout layout() const; + const PipelineLayout &pipeline_layout() const; + + void set_pipeline_layout(const PipelineLayout &layout); + + int pp_rank() const; + + void set_pp_rank(int rank); + private: GlobalEnv() = default; ~GlobalEnv() = default; @@ -85,6 +95,9 @@ class GlobalEnv { bool initialized_ = false; Layout layout_; + PipelineLayout pipeline_layout_; + + int pp_rank_ = 0; }; inline void InitAllEnv(int nthread_per_process, int tensor_parallel_size, bool sequence_parallel_enabled, @@ -106,6 +119,13 @@ inline bool GetSequenceParallelEnabled() { return GlobalEnv::Instance().sequence inline int GetDataParallelSize() { return GlobalEnv::Instance().data_parallel_size(); } inline int GetPipelineParallelSize() { return GlobalEnv::Instance().pipeline_parallel_size(); } inline int GetVirtualPipelineParallelSize() { return GlobalEnv::Instance().virtual_pipeline_parallel_size(); } +inline const PipelineLayout &GetPipelineLayout() { return GlobalEnv::Instance().pipeline_layout(); } +inline void InstallPipelineLayout(const PipelineLayout &layout) { + layout.Validate(); + layout.ValidateForCurrentPipelineTransport(); + GlobalEnv::Instance().set_pipeline_layout(layout); +} +inline int GetPPRank() { return GlobalEnv::Instance().pp_rank(); } // ========================= // Layout Helper Functions diff --git a/infini_train/include/nn/parallel/pipeline_layout.h b/infini_train/include/nn/parallel/pipeline_layout.h new file mode 100644 index 000000000..09072e3da --- /dev/null +++ b/infini_train/include/nn/parallel/pipeline_layout.h @@ -0,0 +1,141 @@ +#pragma once + +#include +#include +#include + +namespace infini_train::nn::parallel { + +class PipelineLayoutError : public std::runtime_error { +public: + explicit PipelineLayoutError(const std::string &msg); +}; + +// Represent the three special modules explicitly. +enum class SpecialModule { + kEmbedding, + kFinalNorm, + kLMHead, +}; + +struct LayerRange { + int start = -1; + int end = -1; + + int size() const; + bool contains(int layer_id) const; +}; + +struct SpecialModulePlacement { + int embedding_stage = -1; + int final_norm_stage = -1; + int lm_head_stage = -1; +}; + +struct ChunkLayout { + int global_chunk_id = -1; // Global chunk identifier. + int stage_id = -1; + int local_chunk_id = -1; // Local chunk index within a stage. + LayerRange layers; +}; + +struct StageLayout { + int stage_id = -1; + std::vector global_chunk_ids; +}; + +struct LayerLocation { + int stage_id = -1; + int global_chunk_id = -1; + int local_chunk_id = -1; + int index_in_chunk = -1; + int flat_local_layer_index = -1; +}; + +// Optional validation policies for layout construction. +struct PipelineLayoutPolicy { + bool allow_empty_stages = false; + bool require_contiguous_execution = true; + bool require_boundary_special_modules = true; +}; + +class PipelineLayout { +public: + PipelineLayout(); + + // Build the default layout with total_chunks = num_stages * vpp_size. + static PipelineLayout BuildDefault(int num_layers, int num_stages, int vpp_size, + SpecialModulePlacement placement = {}, PipelineLayoutPolicy policy = {}); + + // Build a user-defined contiguous layout. + static PipelineLayout BuildContiguous(const std::vector &stage_layer_counts, + SpecialModulePlacement placement = {}, PipelineLayoutPolicy policy = {}); + + // Build an explicit chunk-to-stage layout. Chunks must use contiguous + // global ids [0, num_stages * vpp_size), and each stage must provide one + // chunk for every local_chunk_id in [0, vpp_size). The stage id of each + // chunk is authoritative; it is not inferred from global_chunk_id. + static PipelineLayout BuildExplicit(int num_layers, int num_stages, int vpp_size, + const std::vector &chunks, + SpecialModulePlacement placement = {}, PipelineLayoutPolicy policy = {}); + + int num_layers() const; + int num_stages() const; + int vpp_size() const; + + const StageLayout &stage(int stage_id) const; + const ChunkLayout &chunk(int global_chunk_id) const; + const std::vector &chunks() const; + + const LayerLocation &locate_layer(int layer_id) const; + int stage_of_layer(int layer_id) const; + int local_layer_index(int stage_id, int layer_id) const; + const ChunkLayout &chunk_of_layer(int layer_id) const; + + bool owns(SpecialModule module, int stage_id) const; + const SpecialModulePlacement &special_modules() const; + const PipelineLayoutPolicy &policy() const; + + std::string ToString() const; + void Validate() const; + void ValidateForCurrentPipelineTransport() const; + + static std::vector ParseLayerPartition(const std::string &value); + + // Parse a Megatron-style layout expression. Stages are separated by '|', + // chunks within a stage by ',', 't' denotes one transformer layer and E/F/H + // denote embedding/final-norm/LM-head ownership respectively. Parenthesized + // expressions may be repeated with '*N'. + static PipelineLayout ParseMegatronStyleLayout(const std::string &value, int num_layers, int pp_size, + int vpp_size = 1, SpecialModulePlacement placement = {}, + PipelineLayoutPolicy policy = {}); + + // Suggest a contiguous layer partition that approximately balances the + // supplied per-layer costs. The returned vector has pp_size entries and + // sums to num_layers. + static std::vector SuggestBalancedPartition(int num_layers, int pp_size, + const std::vector &layer_costs = {}); + + static PipelineLayout BuildPipelineLayout(int num_layers, int pp_size, int vpp_size, + const std::string &layer_partition, SpecialModulePlacement placement = {}, + PipelineLayoutPolicy policy = {}); + +private: + PipelineLayout(int num_layers, int num_stages, int vpp_size, std::vector chunks, + SpecialModulePlacement placement, PipelineLayoutPolicy policy); + + static SpecialModulePlacement BuildPlacement(SpecialModulePlacement placement, int num_stages); + + void BuildIndexes(); + + int num_layers_ = 0; + int num_stages_ = 0; + int vpp_size_ = 0; + std::vector chunks_; + std::vector stages_; + SpecialModulePlacement placement_; + PipelineLayoutPolicy policy_; + std::vector layer_locations_; +}; + +} // namespace infini_train::nn::parallel diff --git a/infini_train/include/nn/parallel/pp/pipeline_schedule.h b/infini_train/include/nn/parallel/pp/pipeline_schedule.h index 053650d7c..572f657ed 100644 --- a/infini_train/include/nn/parallel/pp/pipeline_schedule.h +++ b/infini_train/include/nn/parallel/pp/pipeline_schedule.h @@ -16,6 +16,7 @@ class Module; namespace infini_train::nn::parallel { class PipelineStage; +class PipelineLayout; class PipelineSchedule { public: @@ -52,11 +53,13 @@ class PipelineParallelScheduler { bool is_last_chunk; }; - static Task CreateTask(int step, int mb, int global_chunk, int num_stages, int total_chunks, bool is_forward); + static Task CreateTask(int step, int mb, int global_chunk, int num_stages, int total_chunks, bool is_forward, + const PipelineLayout &layout); - static std::vector GenerateGPipeSchedule(int n, int num_stages, int vpp_size); + static std::vector GenerateGPipeSchedule(int n, int num_stages, int vpp_size, const PipelineLayout &layout); - static std::vector GenerateInterleaved1F1BSchedule(int n, int num_stages, int vpp_size); + static std::vector GenerateInterleaved1F1BSchedule(int n, int num_stages, int vpp_size, + const PipelineLayout &layout); }; } // namespace infini_train::nn::parallel diff --git a/infini_train/src/nn/modules/module.cc b/infini_train/src/nn/modules/module.cc index 9475d49fe..8e77446da 100644 --- a/infini_train/src/nn/modules/module.cc +++ b/infini_train/src/nn/modules/module.cc @@ -47,12 +47,37 @@ Module::NamedParameters(const std::string &prefix, bool recurse, bool remove_dup if (recurse) { named_modules = const_cast(this)->NamedModules( - /*memory=*/nullptr, prefix, remove_duplicate); + // Do not let an internal pipeline alias mark the shared module as + // visited before its canonical model-tree path is reached. The + // parameter-level `remove_duplicate` policy is applied below. + /*memory=*/nullptr, prefix, /*remove_duplicate=*/false); } else { named_modules.emplace_back(prefix, std::const_pointer_cast(shared_from_this())); } for (const auto &[module_prefix, module] : named_modules) { + // Pipeline execution wrappers are registered under reserved `__pp_*` + // names. They alias the canonical model tree and must not leak into + // optimizer/checkpoint parameter names (StateDict applies the same + // rule). Skip the wrapper and all of its descendants so the + // `transformer.*` path remains authoritative. + bool is_pipeline_internal = false; + size_t segment_start = 0; + while (segment_start < module_prefix.size()) { + const auto separator = module_prefix.find('.', segment_start); + if (module_prefix.compare(segment_start, 4, "__pp") == 0) { + is_pipeline_internal = true; + break; + } + if (separator == std::string::npos) { + break; + } + segment_start = separator + 1; + } + if (is_pipeline_internal) { + continue; + } + std::vector>> local_parameters; local_parameters.reserve(module->parameters_.size()); diff --git a/infini_train/src/nn/modules/transformer/transformer.cc b/infini_train/src/nn/modules/transformer/transformer.cc index 99a739d2d..74c9ac45b 100644 --- a/infini_train/src/nn/modules/transformer/transformer.cc +++ b/infini_train/src/nn/modules/transformer/transformer.cc @@ -18,6 +18,7 @@ #include "infini_train/include/nn/modules/transformer/moe/moe_layer.h" #include "infini_train/include/nn/modules/transformer/utils.h" #include "infini_train/include/nn/parallel/global.h" +#include "infini_train/include/nn/parallel/pp/pipeline_parallel.h" #include "infini_train/include/nn/parallel/tensor_parallel.h" #include "infini_train/include/nn/parallel/utils.h" #include "infini_train/include/tensor.h" @@ -173,42 +174,52 @@ std::vector> TransformerChunk::Forward(const std::vector return {x1}; } -TransformerLastStage::TransformerLastStage(const TransformerConfig &config) : CloneableModule(kType), config_(config) { - switch (config.norm_type) { - case NormType::kLayerNorm: - modules_[kLnFLayerName] = std::make_shared(std::vector{config_.n_embd}); - break; - case NormType::kRMSNorm: - modules_[kLnFLayerName] = std::make_shared(config.n_embd, config.norm_eps); - break; - default: - LOG(FATAL) << "Unsupported norm type"; +TransformerLastStage::TransformerLastStage(const TransformerConfig &config, bool has_final_norm, bool has_lm_head) + : CloneableModule(kType), config_(config), has_final_norm_(has_final_norm), has_lm_head_(has_lm_head) { + + if (has_final_norm) { + switch (config.norm_type) { + case NormType::kLayerNorm: + modules_[kLnFLayerName] = std::make_shared(std::vector{config_.n_embd}); + break; + case NormType::kRMSNorm: + modules_[kLnFLayerName] = std::make_shared(config.n_embd, config.norm_eps); + break; + default: + LOG(FATAL) << "Unsupported norm type"; + } } + // NOTE(zbl): weight-tying is possible but torch script did not do so - modules_[kLMHeadLayerName] = std::make_shared( - /*in_features=*/config_.n_embd, /*out_features=*/config_.vocab_size, - /*bias=*/config_.add_bias_lm_head, - // NOTE(zbl): each rank would get sharded [B, T, V_local] as logits - /*gather_output=*/false, - /*input_is_parallel=*/false, - /*skip_bias_add=*/false, - /*sequence_parallel=*/nn::parallel::global::GetSequenceParallelEnabled()); + if (has_lm_head) { + modules_[kLMHeadLayerName] = std::make_shared( + /*in_features=*/config_.n_embd, /*out_features=*/config_.vocab_size, + /*bias=*/config_.add_bias_lm_head, + // NOTE(zbl): each rank would get sharded [B, T, V_local] as logits + /*gather_output=*/false, + /*input_is_parallel=*/false, + /*skip_bias_add=*/false, + /*sequence_parallel=*/nn::parallel::global::GetSequenceParallelEnabled()); + } } std::vector> TransformerLastStage::Forward(const std::vector> &x) { // (B, T, C) -> Layernorm -> (B, T, C) - auto x1 = (*modules_[kLnFLayerName])(x); - - // TODO(dcj): add inference-time mini-optimization - // (B, T, C) -> Linear(C, V) -> (B, T, V) - return (*modules_[kLMHeadLayerName])(x1); + auto x1 = x[0]; + if (has_final_norm_) { + x1 = (*modules_[kLnFLayerName])({x1})[0]; + } + if (has_lm_head_) { + return (*modules_[kLMHeadLayerName])({x1}); + } + return {x1}; } TransformerModel::TransformerModel(const TransformerConfig config) - : CloneableModule(kType), config_(config), - stage_info_(nn::parallel::PipelineParallel::GetStageInfo( - config_.n_layer, nn::parallel::global::GetPipelineParallelSize(), nn::parallel::pp_rank, - nn::parallel::global::GetVirtualPipelineParallelSize())) { + : CloneableModule(kType), config_(config), num_local_chunks_(0) { + const auto &layout = nn::parallel::global::GetPipelineLayout(); + const int stage_id = nn::parallel::pp_rank; + const auto &stage = layout.stage(stage_id); auto tp_world_size = nn::parallel::global::GetTensorParallelSize(); // NOTE(zbl): VocabParallelEmbedding requires vocab_size % tp_size == 0 @@ -217,7 +228,7 @@ TransformerModel::TransformerModel(const TransformerConfig config) CHECK_EQ(config.vocab_size % tp_world_size, 0) << "Vocab size should be divisible by TP world size"; std::unordered_map> transformer; - if (stage_info_.is_first_stage) { + if (layout.owns(nn::parallel::SpecialModule::kEmbedding, stage_id)) { modules_[kPPFirstStageName] = std::make_shared(config_); transformer[TransformerFirstStage::kWTELayerName] = modules_[kPPFirstStageName]->mutable_module(TransformerFirstStage::kWTELayerName); @@ -228,32 +239,40 @@ TransformerModel::TransformerModel(const TransformerConfig config) } { - std::map>> start_layer_to_layer_size_and_chunk; - for (int chunk_idx = 0; chunk_idx < stage_info_.layer_ranges_per_chunk.size(); ++chunk_idx) { - const auto [start_layer, end_layer] = stage_info_.layer_ranges_per_chunk[chunk_idx]; - auto chunk = std::make_shared(config_, start_layer, end_layer); - start_layer_to_layer_size_and_chunk[start_layer] = std::make_pair(end_layer - start_layer, chunk); - } std::vector> 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))); + int local_chunk_idx = 0; + for (int gid : stage.global_chunk_ids) { + const auto &c = layout.chunk(gid); + auto chunk = std::make_shared(config_, c.layers.start, c.layers.end); + const int layer_count = c.layers.size(); + for (int i = 0; i < layer_count; ++i) { + h.push_back(chunk->mutable_module(TransformerChunk::kHLayerName)->mutable_module(std::to_string(i))); } - modules_[kPPChunkNamePrefix + std::to_string(chunk_idx)] = std::move(chunk); - ++chunk_idx; + modules_[kPPChunkNamePrefix + std::to_string(local_chunk_idx)] = std::move(chunk); + ++local_chunk_idx; } + num_local_chunks_ = local_chunk_idx; transformer[TransformerChunk::kHLayerName] = std::make_shared(std::move(h)); } - if (stage_info_.is_last_stage) { - modules_[kPPLastStageName] = std::make_shared(config_); - transformer[TransformerLastStage::kLnFLayerName] - = modules_[kPPLastStageName]->mutable_module(TransformerLastStage::kLnFLayerName); - modules_[TransformerLastStage::kLMHeadLayerName] - = modules_[kPPLastStageName]->mutable_module(TransformerLastStage::kLMHeadLayerName); + const bool has_final_norm = layout.owns(nn::parallel::SpecialModule::kFinalNorm, stage_id); + const bool has_lm_head = layout.owns(nn::parallel::SpecialModule::kLMHead, stage_id); + if (has_final_norm || has_lm_head) { + modules_[kPPLastStageName] = std::make_shared(config_, has_final_norm, has_lm_head); + if (has_final_norm) { + transformer[TransformerLastStage::kLnFLayerName] + = modules_[kPPLastStageName]->mutable_module(TransformerLastStage::kLnFLayerName); + } + if (has_lm_head) { + // Keep the canonical checkpoint path under `transformer`, while the + // pipeline wrapper owns the same module for execution. Registering + // a second top-level alias would make StateDict emit `lm_head.*` + // instead of the compatible `transformer.lm_head.*` key. + transformer[TransformerLastStage::kLMHeadLayerName] + = modules_[kPPLastStageName]->mutable_module(TransformerLastStage::kLMHeadLayerName); + } } + modules_[kTransformerModelName] = std::make_shared(std::move(transformer)); // FIXME(jym): Assigning the parameter values of wte to LMHead, which is not real tying operation @@ -267,19 +286,25 @@ TransformerModel::TransformerModel(const TransformerConfig config) *mutable_module(kTransformerModelName) ->mutable_module(TransformerFirstStage::kWTELayerName) ->mutable_parameter(nn::parallel::VocabParallelEmbedding::kParamWeightName) - = module(TransformerLastStage::kLMHeadLayerName) - .parameter(nn::parallel::ColumnParallelLinear::kParamWeightName); + = mutable_module(kTransformerModelName) + ->mutable_module(TransformerLastStage::kLMHeadLayerName) + ->parameter(nn::parallel::ColumnParallelLinear::kParamWeightName); } } std::vector> TransformerModel::Forward(const std::vector> &x) { - auto x1 = (*modules_[kPPFirstStageName])(x); - for (int chunk_idx = 0; chunk_idx < stage_info_.layer_ranges_per_chunk.size(); ++chunk_idx) { - x1 = (*modules_[kPPChunkNamePrefix + std::to_string(chunk_idx)])(x1); + auto x1 = x[0]; + if (modules_.contains(kPPFirstStageName)) { + x1 = (*modules_[kPPFirstStageName])(x)[0]; + } + for (int chunk_idx = 0; chunk_idx < num_local_chunks_; ++chunk_idx) { + x1 = (*modules_[kPPChunkNamePrefix + std::to_string(chunk_idx)])({x1})[0]; } - auto res = (*modules_[kPPLastStageName])(x1); - return res; + if (modules_.contains(kPPLastStageName)) { + return (*modules_[kPPLastStageName])({x1}); + } + return {x1}; } } // namespace infini_train::nn diff --git a/infini_train/src/nn/modules/transformer/transformer_config.cc b/infini_train/src/nn/modules/transformer/transformer_config.cc index b8947d4b6..a9334fba5 100644 --- a/infini_train/src/nn/modules/transformer/transformer_config.cc +++ b/infini_train/src/nn/modules/transformer/transformer_config.cc @@ -7,9 +7,6 @@ namespace infini_train::nn { bool TransformerConfig::UseGQA() const { return n_kv_head < n_head; } int TransformerConfig::GetChunkSize() const { - auto stage_info = parallel::PipelineParallel::GetStageInfo(n_layer, parallel::global::GetPipelineParallelSize(), - parallel::pp_rank, - parallel::global::GetVirtualPipelineParallelSize()); - return stage_info.layer_ranges_per_chunk.size(); + return parallel::global::GetPipelineLayout().stage(parallel::pp_rank).global_chunk_ids.size(); } } // namespace infini_train::nn diff --git a/infini_train/src/nn/parallel/global.cc b/infini_train/src/nn/parallel/global.cc index 655b4bceb..fd1c21437 100644 --- a/infini_train/src/nn/parallel/global.cc +++ b/infini_train/src/nn/parallel/global.cc @@ -119,6 +119,7 @@ void GlobalEnv::Init(int nthread_per_process, int tensor_parallel_size, bool seq layout_.sizes[PP] = pipeline_parallel_size_; layout_.InitStrides(); + pipeline_layout_ = PipelineLayout(); initialized_ = true; } @@ -187,6 +188,13 @@ Layout GlobalEnv::layout() const { return layout_; } +const PipelineLayout &GlobalEnv::pipeline_layout() const { return pipeline_layout_; } + +void GlobalEnv::set_pipeline_layout(const PipelineLayout &layout) { + std::lock_guard lock(mutex_); + pipeline_layout_ = layout; +} + namespace { inline const char *AxisName(Axis a) { return a == DP ? "DP" : (a == TP ? "TP" : "PP"); } diff --git a/infini_train/src/nn/parallel/pipeline_layout.cc b/infini_train/src/nn/parallel/pipeline_layout.cc new file mode 100644 index 000000000..f3d3b876e --- /dev/null +++ b/infini_train/src/nn/parallel/pipeline_layout.cc @@ -0,0 +1,742 @@ +#include "infini_train/include/nn/parallel/pipeline_layout.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace infini_train::nn::parallel { +namespace { + +[[noreturn]] void Fail(const std::string &msg) { throw PipelineLayoutError(msg); } + +void CheckStageId(int stage_id, int num_stages, const std::string &field) { + if (stage_id < 0 || stage_id >= num_stages) { + Fail(std::format("{}={} is outside valid stage range [0, {})", field, stage_id, num_stages)); + } +} + +} // namespace + +PipelineLayoutError::PipelineLayoutError(const std::string &msg) : std::runtime_error(msg) {} + +int LayerRange::size() const { return end - start; } + +bool LayerRange::contains(int layer_id) const { return layer_id < end && layer_id >= start; } + +PipelineLayout::PipelineLayout() : num_layers_(0), num_stages_(1), vpp_size_(1) { BuildIndexes(); } + +PipelineLayout PipelineLayout::BuildDefault(int num_layers, int num_stages, int vpp_size, + SpecialModulePlacement placement, PipelineLayoutPolicy policy) { + if (num_layers <= 0) { + Fail(std::format("num_layers must be positive, got {}", num_layers)); + } + if (num_stages <= 0) { + Fail(std::format("num_stages must be positive, got {}", num_stages)); + } + if (vpp_size <= 0) { + Fail(std::format("vpp_size must be positive, got {}", vpp_size)); + } + + const long long total_chunks_ll = static_cast(num_stages) * static_cast(vpp_size); + if (total_chunks_ll > std::numeric_limits::max()) { + Fail("num_stages * vpp_size exceeds supported integer range"); + } + + // Total number of transformer layers. + const int total_chunks = static_cast(total_chunks_ll); + + const int per_chunk_layers = num_layers / total_chunks; + const int remainder_layers = num_layers % total_chunks; + + std::vector chunks; + chunks.reserve(total_chunks); + int layer_cursor = 0; + for (int global_chunk_id = 0; global_chunk_id < total_chunks; ++global_chunk_id) { + // Distribute remainder layers across the first chunks. + int chunk_size = per_chunk_layers + (global_chunk_id < remainder_layers ? 1 : 0); + if (chunk_size == 0) { + continue; + } + int stage_id = global_chunk_id % num_stages; + int local_chunk_id = global_chunk_id / num_stages; + chunks.push_back(ChunkLayout{ + .global_chunk_id = global_chunk_id, + .stage_id = stage_id, + .local_chunk_id = local_chunk_id, + .layers = LayerRange{.start = layer_cursor, .end = layer_cursor + chunk_size}, + }); + layer_cursor += chunk_size; + } + if (num_layers < total_chunks) { + policy.allow_empty_stages = true; + } + auto normalized_placement = BuildPlacement(placement, num_stages); + return PipelineLayout(num_layers, num_stages, vpp_size, std::move(chunks), normalized_placement, policy); +} + +PipelineLayout PipelineLayout::BuildContiguous(const std::vector &stage_layer_counts, + SpecialModulePlacement placement, PipelineLayoutPolicy policy) { + if (stage_layer_counts.empty()) { + Fail("stage layer counts must not be empty"); + } + + const int num_stages = static_cast(stage_layer_counts.size()); + + int num_layers = 0; + for (int idx = 0; idx < num_stages; ++idx) { + if (stage_layer_counts[idx] < 0) { + Fail(std::format("stage_layer_counts must be non-negative, got {}", stage_layer_counts[idx])); + } + if (stage_layer_counts[idx] == 0 && !policy.allow_empty_stages) { + Fail(std::format("stage {} has zero layers; set allow_empty_stages=true to permit", idx)); + } + num_layers += stage_layer_counts[idx]; + } + + const int vpp_size = 1; + + std::vector chunks; + int layer_cursor = 0; + for (int stage_id = 0; stage_id < num_stages; ++stage_id) { + int count = stage_layer_counts[stage_id]; + if (count == 0) { + continue; + } + int global_chunk_id = stage_id; // With vpp_size=1, the IDs equal stage IDs. + chunks.push_back(ChunkLayout{ + .global_chunk_id = global_chunk_id, + .stage_id = stage_id, + .local_chunk_id = 0, + .layers = LayerRange{.start = layer_cursor, .end = layer_cursor + count}, + }); + layer_cursor += count; + } + + auto normalized_placement = BuildPlacement(placement, num_stages); + return PipelineLayout(num_layers, num_stages, vpp_size, std::move(chunks), normalized_placement, policy); +} + +PipelineLayout PipelineLayout::BuildExplicit(int num_layers, int num_stages, int vpp_size, + const std::vector &chunks, + SpecialModulePlacement placement, PipelineLayoutPolicy policy) { + if (num_layers <= 0 || num_stages <= 0 || vpp_size <= 0) { + Fail(std::format("num_layers, num_stages and vpp_size must be positive (got {}, {}, {})", num_layers, + num_stages, vpp_size)); + } + const long long expected_chunks = static_cast(num_stages) * static_cast(vpp_size); + if (static_cast(chunks.size()) != expected_chunks) { + Fail(std::format("explicit chunk count mismatch: expected={}, actual={}", expected_chunks, chunks.size())); + } + + std::vector ordered = chunks; + std::sort(ordered.begin(), ordered.end(), + [](const ChunkLayout &a, const ChunkLayout &b) { return a.global_chunk_id < b.global_chunk_id; }); + std::vector seen_global(static_cast(expected_chunks), false); + std::vector> seen_local(static_cast(num_stages), + std::vector(static_cast(vpp_size), false)); + for (const auto &chunk : ordered) { + if (chunk.global_chunk_id < 0 || chunk.global_chunk_id >= expected_chunks) { + Fail(std::format("explicit chunk global_chunk_id={} outside [0, {})", chunk.global_chunk_id, + expected_chunks)); + } + if (seen_global[static_cast(chunk.global_chunk_id)]) { + Fail(std::format("duplicate explicit global_chunk_id={}", chunk.global_chunk_id)); + } + seen_global[static_cast(chunk.global_chunk_id)] = true; + CheckStageId(chunk.stage_id, num_stages, "chunk.stage_id"); + if (chunk.local_chunk_id < 0 || chunk.local_chunk_id >= vpp_size) { + Fail(std::format("explicit chunk {} local_chunk_id={} outside [0, {})", chunk.global_chunk_id, + chunk.local_chunk_id, vpp_size)); + } + const auto stage_index = static_cast(chunk.stage_id); + const auto local_index = static_cast(chunk.local_chunk_id); + if (seen_local[stage_index][local_index]) { + Fail(std::format("stage {} has duplicate local_chunk_id={}", chunk.stage_id, chunk.local_chunk_id)); + } + seen_local[stage_index][local_index] = true; + if (chunk.layers.start < 0 || chunk.layers.end < chunk.layers.start || chunk.layers.end > num_layers) { + Fail(std::format("explicit chunk {} layer range [{}, {}) is outside [0, {})", chunk.global_chunk_id, + chunk.layers.start, chunk.layers.end, num_layers)); + } + if (chunk.layers.size() == 0) { + policy.allow_empty_stages = true; + } + } + for (int gid = 0; gid < expected_chunks; ++gid) { + if (!seen_global[static_cast(gid)]) { + Fail(std::format("missing explicit global_chunk_id={}", gid)); + } + } + for (int stage_id = 0; stage_id < num_stages; ++stage_id) { + for (int local_chunk_id = 0; local_chunk_id < vpp_size; ++local_chunk_id) { + if (!seen_local[static_cast(stage_id)][static_cast(local_chunk_id)]) { + Fail(std::format("stage {} is missing local_chunk_id={}", stage_id, local_chunk_id)); + } + } + } + return PipelineLayout(num_layers, num_stages, vpp_size, std::move(ordered), BuildPlacement(placement, num_stages), + policy); +} + +// By default, place norm and lm_head on the last stage and embedding on the first. +SpecialModulePlacement PipelineLayout::BuildPlacement(SpecialModulePlacement placement, int num_stages) { + if (placement.embedding_stage < 0) { + placement.embedding_stage = 0; + } + if (placement.final_norm_stage < 0) { + placement.final_norm_stage = num_stages - 1; + } + if (placement.lm_head_stage < 0) { + placement.lm_head_stage = num_stages - 1; + } + return placement; +} + +PipelineLayout::PipelineLayout(int num_layers, int num_stages, int vpp_size, std::vector chunks, + SpecialModulePlacement placement, PipelineLayoutPolicy policy) + : num_layers_(num_layers), num_stages_(num_stages), vpp_size_(vpp_size), chunks_(std::move(chunks)), + placement_(std::move(placement)), policy_(std::move(policy)) { + Validate(); + BuildIndexes(); +} + +void PipelineLayout::BuildIndexes() { + // Build the layer_id -> stage_id index. + layer_locations_.resize(num_layers_); + + for (const auto &c : chunks_) { + for (int layer_id = c.layers.start; layer_id < c.layers.end; ++layer_id) { + auto &loc = layer_locations_[layer_id]; + loc.stage_id = c.stage_id; + loc.global_chunk_id = c.global_chunk_id; + loc.local_chunk_id = c.local_chunk_id; + loc.index_in_chunk = layer_id - c.layers.start; + } + } + + // Compute contiguous local indices across chunks within each stage. + std::vector stage_counter(num_stages_, 0); + for (int layer_id = 0; layer_id < num_layers_; ++layer_id) { + auto &loc = layer_locations_[layer_id]; + if (loc.stage_id >= 0) { + loc.flat_local_layer_index = stage_counter[loc.stage_id]++; + } + } + + // Build stages grouped by stage_id and collect global_chunk_ids. + stages_.clear(); + stages_.resize(num_stages_); + for (int i = 0; i < num_stages_; ++i) { stages_[i].stage_id = i; } + for (const auto &c : chunks_) { stages_[c.stage_id].global_chunk_ids.push_back(c.global_chunk_id); } +} + +int PipelineLayout::num_layers() const { return num_layers_; } +int PipelineLayout::num_stages() const { return num_stages_; } +int PipelineLayout::vpp_size() const { return vpp_size_; } + +const StageLayout &PipelineLayout::stage(int stage_id) const { + if (stage_id < 0 || stage_id >= num_stages_) { + Fail(std::format("stage_id={} out of range [0, {})", stage_id, num_stages_)); + } + return stages_[stage_id]; +} + +const ChunkLayout &PipelineLayout::chunk(int global_chunk_id) const { + for (const auto &c : chunks_) { + if (c.global_chunk_id == global_chunk_id) { + return c; + } + } + Fail(std::format("global_chunk_id={} not found in layout", global_chunk_id)); +} + +const std::vector &PipelineLayout::chunks() const { return chunks_; } + +const LayerLocation &PipelineLayout::locate_layer(int layer_id) const { + if (layer_id < 0 || layer_id >= num_layers_) { + Fail(std::format("layer_id={} out of range [0, {})", layer_id, num_layers_)); + } + return layer_locations_[layer_id]; +} + +int PipelineLayout::stage_of_layer(int layer_id) const { return locate_layer(layer_id).stage_id; } + +int PipelineLayout::local_layer_index(int stage_id, int layer_id) const { + if (stage_id < 0 || stage_id >= num_stages_) { + Fail(std::format("stage_id={} out of range [0, {})", stage_id, num_stages_)); + } + auto &loc = locate_layer(layer_id); + if (loc.stage_id != stage_id) { + Fail(std::format("layer {} belongs to stage {}, not {}", layer_id, loc.stage_id, stage_id)); + } + return loc.flat_local_layer_index; +} + +const ChunkLayout &PipelineLayout::chunk_of_layer(int layer_id) const { + return chunk(locate_layer(layer_id).global_chunk_id); +} + +bool PipelineLayout::owns(SpecialModule module, int stage_id) const { + switch (module) { + case SpecialModule::kEmbedding: + return placement_.embedding_stage == stage_id; + case SpecialModule::kFinalNorm: + return placement_.final_norm_stage == stage_id; + case SpecialModule::kLMHead: + return placement_.lm_head_stage == stage_id; + } + return false; +} + +const SpecialModulePlacement &PipelineLayout::special_modules() const { return placement_; } +const PipelineLayoutPolicy &PipelineLayout::policy() const { return policy_; } + +void PipelineLayout::Validate() const { + // 1. Layer ranges must cover [0, num_layers_) exactly once. + std::vector seen(num_layers_, false); + for (const auto &c : chunks_) { + if (c.layers.start < 0 || c.layers.end < c.layers.start || c.layers.end > num_layers_) { + Fail(std::format("chunk {} layer range [{}, {}) is outside [0, {})", c.global_chunk_id, c.layers.start, + c.layers.end, num_layers_)); + } + for (int l = c.layers.start; l < c.layers.end; ++l) { + if (seen[l]) { + Fail(std::format("layer {} appears in more than one chunk", l)); + } + seen[l] = true; + } + } + for (int l = 0; l < num_layers_; ++l) { + if (!seen[l]) { + Fail(std::format("layer {} is not assigned to any chunk", l)); + } + } + + // 2. All chunk and stage indices must be valid. + for (const auto &c : chunks_) { + if (c.stage_id < 0 || c.stage_id >= num_stages_) { + Fail(std::format("chunk {} stage_id={} out of range [0, {})", c.global_chunk_id, c.stage_id, num_stages_)); + } + if (c.global_chunk_id < 0) { + Fail(std::format("chunk has negative global_chunk_id={}", c.global_chunk_id)); + } + } + + std::unordered_set chunk_ids; + for (const auto &c : chunks_) { + if (!chunk_ids.insert(c.global_chunk_id).second) { + Fail(std::format("duplicate global_chunk_id={}", c.global_chunk_id)); + } + if (c.local_chunk_id < 0) { + Fail(std::format("chunk {} has negative local_chunk_id={}", c.global_chunk_id, c.local_chunk_id)); + } + } + + // 3. Special-module stages must be valid. + CheckStageId(placement_.embedding_stage, num_stages_, "embedding_stage"); + CheckStageId(placement_.final_norm_stage, num_stages_, "final_norm_stage"); + CheckStageId(placement_.lm_head_stage, num_stages_, "lm_head_stage"); + + if (policy_.require_boundary_special_modules) { + if (placement_.embedding_stage != 0) { + Fail(std::format("embedding_stage={} must be boundary stage 0", placement_.embedding_stage)); + } + if (placement_.final_norm_stage != num_stages_ - 1) { + Fail( + std::format("final_norm_stage={} must be last stage {}", placement_.final_norm_stage, num_stages_ - 1)); + } + if (placement_.lm_head_stage != num_stages_ - 1) { + Fail(std::format("lm_head_stage={} must be last stage {}", placement_.lm_head_stage, num_stages_ - 1)); + } + } + + // 4. The default policy disallows empty stages. + if (!policy_.allow_empty_stages) { + std::vector stage_has_layers(num_stages_, false); + for (const auto &c : chunks_) { + if (c.layers.size() > 0) { + stage_has_layers[c.stage_id] = true; + } + } + for (int i = 0; i < num_stages_; ++i) { + if (!stage_has_layers[i]) { + Fail(std::format("stage {} has no layers; set allow_empty_stages=true to permit", i)); + } + } + } + + // 5. The default policy requires contiguous execution order. + if (policy_.require_contiguous_execution) { + // Sort by layers.start and verify adjacent ranges are contiguous. + auto sorted = chunks_; + std::sort(sorted.begin(), sorted.end(), + [](const ChunkLayout &a, const ChunkLayout &b) { return a.layers.start < b.layers.start; }); + for (size_t i = 1; i < sorted.size(); ++i) { + if (sorted[i].layers.start != sorted[i - 1].layers.end) { + Fail(std::format("layer ranges are not contiguous: gap between [{}, {}) and [{}, {})", + sorted[i - 1].layers.start, sorted[i - 1].layers.end, sorted[i].layers.start, + sorted[i].layers.end)); + } + } + } +} + +void PipelineLayout::ValidateForCurrentPipelineTransport() const { + Validate(); + // Send/recv currently connects adjacent stages in 0..num_stages-1 order, + // so every stage must have at least one executable chunk. + for (int i = 0; i < num_stages_; ++i) { + if (stages_[i].global_chunk_ids.empty()) { + Fail(std::format("stage {} has no chunks; the current pipeline transport requires " + "every stage to have at least one chunk", + i)); + } + } +} + +std::string PipelineLayout::ToString() const { + std::ostringstream out; + out << std::format("PipelineLayout:\n"); + out << std::format(" num_layers: {}\n", num_layers_); + out << std::format(" num_stages: {}\n", num_stages_); + out << std::format(" vpp_size: {}\n", vpp_size_); + out << std::format(" special_modules: embedding={}, final_norm={}, lm_head={}\n", placement_.embedding_stage, + placement_.final_norm_stage, placement_.lm_head_stage); + + for (int sid = 0; sid < num_stages_; ++sid) { + out << std::format(" stage {}:", sid); + const auto &stage = stages_[sid]; + if (stage.global_chunk_ids.empty()) { + out << " (empty)"; + } + for (int gid : stage.global_chunk_ids) { + const auto &c = chunk(gid); + out << std::format(" [layers {}-{}]", c.layers.start, c.layers.end - 1); + } + // Record special-module ownership. + std::vector specials; + if (owns(SpecialModule::kEmbedding, sid)) { + specials.push_back("embedding"); + } + if (owns(SpecialModule::kFinalNorm, sid)) { + specials.push_back("final_norm"); + } + if (owns(SpecialModule::kLMHead, sid)) { + specials.push_back("lm_head"); + } + if (!specials.empty()) { + out << " {"; + for (size_t i = 0; i < specials.size(); ++i) { + if (i > 0) { + out << ", "; + } + out << specials[i]; + } + out << "}"; + } + out << "\n"; + } + return out.str(); +} + +std::vector PipelineLayout::ParseLayerPartition(const std::string &value) { + if (value.empty()) { + Fail("pipeline_layer_partition must not be empty"); + } + + std::vector result; + size_t begin = 0; + + while (begin <= value.size()) { + const size_t comma = value.find(',', begin); + const size_t end = (comma == std::string::npos) ? value.size() : comma; + + if (end == begin) { + Fail(std::format("pipeline_layer_partition contains an empty token near offset {}", begin)); + } + + const std::string_view token(value.data() + begin, end - begin); + + // Only non-negative integer tokens are accepted; whitespace, signs and letters are rejected. + for (const char c : token) { + if (c < '0' || c > '9') { + Fail(std::format("pipeline_layer_partition contains invalid character '{}'", c)); + } + } + + int number = 0; + const char *const first = token.data(); + const char *const last = token.data() + token.size(); + const auto [ptr, ec] = std::from_chars(first, last, number); + + if (ec != std::errc() || ptr != last) { + Fail(std::format("pipeline_layer_partition token '{}' is not a valid non-negative int", token)); + } + + result.push_back(number); + + if (comma == std::string::npos) { + break; + } + + begin = comma + 1; + } + + return result; +} + +PipelineLayout PipelineLayout::ParseMegatronStyleLayout(const std::string &value, int num_layers, int pp_size, + int vpp_size, SpecialModulePlacement placement, + PipelineLayoutPolicy policy) { + if (value.empty()) { + Fail("pipeline_layout must not be empty"); + } + if (num_layers <= 0 || pp_size <= 0 || vpp_size <= 0) { + Fail(std::format("num_layers, pp_size and vpp_size must be positive (got {}, {}, {})", num_layers, pp_size, + vpp_size)); + } + + // Expand parenthesized repetitions while preserving stage separators. + std::function expand = [&](std::string_view input) { + std::string out; + for (size_t i = 0; i < input.size();) { + if (input[i] != '(') { + if (input[i] == ')' || input[i] == '*') { + Fail(std::format("pipeline_layout has unexpected '{}' at offset {}", input[i], i)); + } + out.push_back(input[i++]); + continue; + } + size_t depth = 1, j = i + 1; + for (; j < input.size() && depth != 0; ++j) { + if (input[j] == '(') { + ++depth; + } else if (input[j] == ')') { + --depth; + } + } + if (depth != 0) { + Fail("pipeline_layout has unmatched '('"); + } + const std::string inner = expand(input.substr(i + 1, j - i - 2)); + size_t repeat = 1; + if (j < input.size() && input[j] == '*') { + size_t k = j + 1; + if (k == input.size() || input[k] < '0' || input[k] > '9') { + Fail("pipeline_layout repetition must use *N"); + } + repeat = 0; + while (k < input.size() && input[k] >= '0' && input[k] <= '9') { + repeat = repeat * 10 + static_cast(input[k++] - '0'); + if (repeat > 100000) { + Fail("pipeline_layout repetition is too large"); + } + } + if (repeat == 0) { + Fail("pipeline_layout repetition must be positive"); + } + j = k; + } + for (size_t n = 0; n < repeat; ++n) { out += inner; } + i = j; + } + return out; + }; + + const std::string expanded = expand(value); + std::vector stage_exprs; + size_t begin = 0; + while (true) { + const size_t sep = expanded.find('|', begin); + stage_exprs.emplace_back(expanded.substr(begin, sep == std::string::npos ? std::string::npos : sep - begin)); + if (sep == std::string::npos) { + break; + } + begin = sep + 1; + } + if (static_cast(stage_exprs.size()) != pp_size) { + Fail(std::format("pipeline_layout stage count mismatch: pp_size={}, stages={}", pp_size, stage_exprs.size())); + } + + std::vector> layer_counts(pp_size); + bool saw_embedding = false, saw_final_norm = false, saw_lm_head = false; + for (int stage_id = 0; stage_id < pp_size; ++stage_id) { + const auto &expr = stage_exprs[stage_id]; + size_t chunk_begin = 0; + while (true) { + const size_t comma = expr.find(',', chunk_begin); + const std::string token + = expr.substr(chunk_begin, comma == std::string::npos ? std::string::npos : comma - chunk_begin); + int t_count = 0; + for (char c : token) { + switch (c) { + case 't': + case 'T': + ++t_count; + break; + case 'E': + if (saw_embedding) { + Fail("pipeline_layout contains multiple E tokens"); + } + placement.embedding_stage = stage_id; + saw_embedding = true; + break; + case 'F': + if (saw_final_norm) { + Fail("pipeline_layout contains multiple F tokens"); + } + placement.final_norm_stage = stage_id; + saw_final_norm = true; + break; + case 'H': + case 'L': + if (saw_lm_head) { + Fail("pipeline_layout contains multiple H/L tokens"); + } + placement.lm_head_stage = stage_id; + saw_lm_head = true; + break; + case ' ': + case '\t': + case '\n': + case '\r': + Fail("pipeline_layout does not allow whitespace"); + default: + Fail(std::format("pipeline_layout contains unknown symbol '{}'", c)); + } + } + layer_counts[stage_id].push_back(t_count); + if (comma == std::string::npos) { + break; + } + chunk_begin = comma + 1; + } + if (static_cast(layer_counts[stage_id].size()) != vpp_size) { + Fail(std::format("pipeline_layout chunk count mismatch at stage {}: vpp_size={}, chunks={}", stage_id, + vpp_size, layer_counts[stage_id].size())); + } + } + + // Global execution order is chunk-major: all stages of local chunk 0, + // followed by all stages of local chunk 1, matching the existing + // interleaved scheduler's global chunk numbering. + std::vector chunks; + chunks.reserve(static_cast(pp_size) * static_cast(vpp_size)); + int layer_cursor = 0; + for (int local_chunk_id = 0; local_chunk_id < vpp_size; ++local_chunk_id) { + for (int stage_id = 0; stage_id < pp_size; ++stage_id) { + const int count = layer_counts[stage_id][local_chunk_id]; + const int global_chunk_id = local_chunk_id * pp_size + stage_id; + chunks.push_back(ChunkLayout{.global_chunk_id = global_chunk_id, + .stage_id = stage_id, + .local_chunk_id = local_chunk_id, + .layers = LayerRange{layer_cursor, layer_cursor + count}}); + layer_cursor += count; + } + } + if (layer_cursor != num_layers) { + Fail(std::format("pipeline_layout layer count mismatch: num_layers={}, parsed_layers={}", num_layers, + layer_cursor)); + } + if (!saw_embedding) { + placement.embedding_stage = 0; + } + if (!saw_final_norm) { + placement.final_norm_stage = pp_size - 1; + } + if (!saw_lm_head) { + placement.lm_head_stage = pp_size - 1; + } + // Explicit layouts may intentionally contain empty stages; retain them for + // inspection and let transport validation decide whether they are runnable. + for (const auto &c : chunks) { + if (c.layers.size() == 0) { + policy.allow_empty_stages = true; + } + } + return PipelineLayout(num_layers, pp_size, vpp_size, std::move(chunks), BuildPlacement(placement, pp_size), policy); +} + +std::vector PipelineLayout::SuggestBalancedPartition(int num_layers, int pp_size, + const std::vector &layer_costs) { + if (num_layers <= 0 || pp_size <= 0) { + Fail(std::format("num_layers and pp_size must be positive (got {}, {})", num_layers, pp_size)); + } + if (num_layers < pp_size) { + Fail(std::format("num_layers={} is smaller than pp_size={}", num_layers, pp_size)); + } + std::vector costs = layer_costs; + if (costs.empty()) { + costs.assign(static_cast(num_layers), 1.0); + } + if (static_cast(costs.size()) != num_layers) { + Fail(std::format("layer_costs length mismatch: num_layers={}, costs={}", num_layers, costs.size())); + } + for (int i = 0; i < num_layers; ++i) { + if (!std::isfinite(costs[i]) || costs[i] < 0.0) { + Fail(std::format("layer_costs[{}] must be finite and non-negative", i)); + } + } + double total = 0.0; + for (double c : costs) { total += c; } + std::vector result; + result.reserve(pp_size); + int cursor = 0; + double remaining = total; + for (int stage = 0; stage < pp_size - 1; ++stage) { + const int layers_left = num_layers - cursor; + const int stages_left = pp_size - stage; + const double target = remaining / static_cast(stages_left); + int take = 1; + double accum = costs[cursor]; + while (take < layers_left - (stages_left - 1) && accum + costs[cursor + take] <= target) { + accum += costs[cursor + take++]; + } + // If the next layer gets us closer to the target, include it; a heavy + // layer is consequently isolated instead of being paired blindly. + if (take < layers_left - (stages_left - 1) + && std::abs(accum + costs[cursor + take] - target) < std::abs(accum - target)) { + accum += costs[cursor + take++]; + } + result.push_back(take); + cursor += take; + remaining -= accum; + } + result.push_back(num_layers - cursor); + return result; +} + +PipelineLayout PipelineLayout::BuildPipelineLayout(int num_layers, int pp_size, int vpp_size, + const std::string &layer_partition, SpecialModulePlacement placement, + PipelineLayoutPolicy policy) { + if (layer_partition.empty()) { + return BuildDefault(num_layers, pp_size, vpp_size, placement, policy); + } + std::vector pp = ParseLayerPartition(layer_partition); + + if (static_cast(pp.size()) != pp_size) { + Fail(std::format("stage count mismatch: pipeline_parallel={}, partition_entries={}", pp_size, pp.size())); + } + + long long layer_sum = 0; + for (int layer_count : pp) { layer_sum += layer_count; } + + if (layer_sum != num_layers) { + Fail(std::format("layer sum mismatch: num_layers={}, partition_sum={}", num_layers, layer_sum)); + } + + if (vpp_size != 1) { + Fail(std::format("custom layer partition is not supported with vpp_size={}", vpp_size)); + } + + return BuildContiguous(pp, placement, policy); +} + +} // namespace infini_train::nn::parallel diff --git a/infini_train/src/nn/parallel/pp/pipeline_parallel.cc b/infini_train/src/nn/parallel/pp/pipeline_parallel.cc index c0369cdeb..c30c94b40 100644 --- a/infini_train/src/nn/parallel/pp/pipeline_parallel.cc +++ b/infini_train/src/nn/parallel/pp/pipeline_parallel.cc @@ -2,11 +2,16 @@ #include "infini_train/include/nn/parallel/pp/pipeline_parallel.h" #include +#include #include +#include #include +#include "glog/logging.h" #include "infini_train/include/nn/modules/container.h" #include "infini_train/include/nn/modules/module.h" +#include "infini_train/include/nn/parallel/global.h" +#include "infini_train/include/nn/parallel/pipeline_layout.h" #include "infini_train/include/nn/parallel/pp/pipeline_schedule.h" #include "infini_train/include/nn/parallel/pp/pipeline_stage.h" @@ -78,22 +83,79 @@ StageInfo PipelineParallel::GetStageInfo(int total_layers, int pp_size, int rank PipelineParallel::PipelineParallel(const std::shared_ptr module, int num_stages, int num_micro_batches, const std::vector> &recv_shape, int pp_rank, Device device, - int chunk_size) + int vpp) : num_stages_(num_stages), rank_(pp_rank) { - modules_[kModuleName] = std::move(module); + const auto &layout = global::GetPipelineLayout(); + if (module == nullptr) { + throw std::invalid_argument("PipelineParallel requires a non-null model module"); + } + if (num_stages != layout.num_stages()) { + throw PipelineLayoutError( + std::format("PipelineParallel num_stages={} does not match installed PipelineLayout num_stages={}", + num_stages, layout.num_stages())); + } + if (pp_rank < 0 || pp_rank >= layout.num_stages()) { + throw PipelineLayoutError( + std::format("PipelineParallel rank={} is outside installed PipelineLayout stage range [0, {})", pp_rank, + layout.num_stages())); + } + layout.ValidateForCurrentPipelineTransport(); + + // vpp is retained for source compatibility with existing callers. The + // installed layout is authoritative for local chunk ownership/order. + (void)vpp; + + const auto &stage = layout.stage(rank_); + const int num_local_chunks = static_cast(stage.global_chunk_ids.size()); + if (num_local_chunks == 0) { + throw PipelineLayoutError(std::format("PipelineLayout stage {} has no local chunks", rank_)); + } - int stage_id = pp_rank; - int stage_size = num_stages; + const auto wrapped_module = module; + modules_[kModuleName] = wrapped_module; std::vector> chunks; - for (int chunk_id = 0; chunk_id < chunk_size; ++chunk_id) { + chunks.reserve(num_local_chunks); + + const bool owns_embedding = layout.owns(SpecialModule::kEmbedding, rank_); + const bool owns_final_norm = layout.owns(SpecialModule::kFinalNorm, rank_); + const bool owns_lm_head = layout.owns(SpecialModule::kLMHead, rank_); + const bool stages_last_module = owns_final_norm || owns_lm_head; + + for (int local_chunk_idx = 0; local_chunk_idx < num_local_chunks; ++local_chunk_idx) { + const int global_chunk_id = stage.global_chunk_ids.at(local_chunk_idx); + const auto &chunk_layout = layout.chunk(global_chunk_id); + if (chunk_layout.stage_id != rank_) { + throw PipelineLayoutError( + std::format("Global chunk {} belongs to stage {}, but PipelineParallel is building stage {}", + global_chunk_id, chunk_layout.stage_id, rank_)); + } + std::vector> chunk_parts; - if (chunk_id == 0 && stage_id == 0) { - chunk_parts.push_back(module->mutable_module(kPPFirstStageName)); + if (local_chunk_idx == 0 && owns_embedding) { + auto &first_stage = wrapped_module->mutable_module(kPPFirstStageName); + if (first_stage == nullptr) { + throw std::invalid_argument("Pipeline model owns embedding but __pp_first_stage is null"); + } + chunk_parts.push_back(first_stage); + } + + const std::string chunk_name = kPPChunkNamePrefix + std::to_string(chunk_layout.local_chunk_id); + auto &transformer_chunk = wrapped_module->mutable_module(chunk_name); + if (transformer_chunk == nullptr) { + throw std::invalid_argument(std::format("Local transformer chunk module '{}' is null", chunk_name)); + } + chunk_parts.push_back(transformer_chunk); + if (local_chunk_idx == num_local_chunks - 1 && stages_last_module) { + auto &last_stage = wrapped_module->mutable_module(kPPLastStageName); + if (last_stage == nullptr) { + throw std::invalid_argument("Pipeline model owns final norm or LM head but __pp_last_stage is null"); + } + chunk_parts.push_back(last_stage); } - chunk_parts.push_back(module->mutable_module(kPPChunkNamePrefix + std::to_string(chunk_id))); - if (chunk_id == chunk_size - 1 && stage_id == stage_size - 1) { - chunk_parts.push_back(module->mutable_module(kPPLastStageName)); + if (chunk_parts.empty()) { + throw std::invalid_argument( + std::format("Pipeline stage {} local chunk {} has no executable modules", rank_, local_chunk_idx)); } chunks.push_back(std::make_shared(std::move(chunk_parts))); } diff --git a/infini_train/src/nn/parallel/pp/pipeline_schedule.cc b/infini_train/src/nn/parallel/pp/pipeline_schedule.cc index b702a3016..dd4ce3572 100644 --- a/infini_train/src/nn/parallel/pp/pipeline_schedule.cc +++ b/infini_train/src/nn/parallel/pp/pipeline_schedule.cc @@ -1,27 +1,45 @@ // pipeline_schedule.cc #include "infini_train/include/nn/parallel/pp/pipeline_schedule.h" +#include +#include #include +#include +#include #include +#include #include #include "glog/logging.h" #include "infini_train/include/autocast.h" +#include "infini_train/include/core/runtime/device_guard.h" #include "infini_train/include/datatype.h" #include "infini_train/include/device.h" #include "infini_train/include/nn/init.h" #include "infini_train/include/nn/modules/module.h" #include "infini_train/include/nn/parallel/global.h" +#include "infini_train/include/nn/parallel/pipeline_layout.h" #include "infini_train/include/nn/parallel/pp/pipeline_stage.h" #include "infini_train/include/nn/parallel/pp/send_recv.h" #include "infini_train/include/optimizer.h" #include "infini_train/include/tensor.h" namespace infini_train::nn::parallel { +namespace { +bool StageTimingEnabled() { + const char *value = std::getenv("INFINI_PIPELINE_STAGE_TIMING"); + return value != nullptr && std::string(value) == "1"; +} + +void SynchronizeStageStream(const Device &device) { + auto *impl = core::GetDeviceGuardImpl(device.type()); + impl->SynchronizeStream(impl->GetStream(device)); +} +} // namespace void PrintScheduleTable(const std::vector &schedule, int n, int num_stages, - int vpp_size) { + int vpp_size, const PipelineLayout &layout) { int total_global_chunks = num_stages * vpp_size; LOG(INFO) << std::format("=== Schedule Table ===\n" @@ -32,8 +50,9 @@ void PrintScheduleTable(const std::vector &sche LOG(INFO) << "-----|-----------|------------|--------------|-------------|-------"; for (const auto &task : schedule) { - int owning_stage = task.global_chunk_id % num_stages; - int local_chunk = task.global_chunk_id / num_stages; + const auto &chunk = layout.chunk(task.global_chunk_id); + int owning_stage = chunk.stage_id; + int local_chunk = chunk.local_chunk_id; std::string type_str = task.is_forward ? "Forward" : "Backward"; @@ -69,22 +88,33 @@ std::vector> PipelineSchedule::SendToNext(const std::vec } PipelineParallelScheduler::Task PipelineParallelScheduler::CreateTask(int step, int mb, int global_chunk, - int num_stages, int total_chunks, - bool is_forward) { + int num_stages, int total_chunks, bool is_forward, + const PipelineLayout &layout) { + if (num_stages != layout.num_stages()) { + throw PipelineLayoutError(std::format("scheduler stage count mismatch: num_stages={}, layout.num_stages={}", + num_stages, layout.num_stages())); + } + const auto &chunk = layout.chunk(global_chunk); + PipelineParallelScheduler::Task task; task.step = step; task.microbatch_id = mb; task.global_chunk_id = global_chunk; - task.local_chunk_idx = global_chunk / num_stages; + task.local_chunk_idx = chunk.local_chunk_id; task.is_forward = is_forward; - task.stage_id = global_chunk % num_stages; + task.stage_id = chunk.stage_id; task.is_last_chunk = (global_chunk == total_chunks - 1); task.is_first_chunk = (global_chunk == 0); return task; } -std::vector PipelineParallelScheduler::GenerateGPipeSchedule(int n, int num_stages, - int vpp_size) { +std::vector +PipelineParallelScheduler::GenerateGPipeSchedule(int n, int num_stages, int vpp_size, const PipelineLayout &layout) { + if (layout.num_stages() != num_stages || layout.vpp_size() != vpp_size) { + throw PipelineLayoutError( + std::format("scheduler topology mismatch: requested stages={}, vpp={}, layout stages={}, vpp={}", + num_stages, vpp_size, layout.num_stages(), layout.vpp_size())); + } std::vector schedule; int total_global_chunks = num_stages * vpp_size; int total_steps = n + total_global_chunks - 1; @@ -95,7 +125,7 @@ std::vector PipelineParallelScheduler::Generate int global_chunk_id = step - mb; if (global_chunk_id >= 0 && global_chunk_id < total_global_chunks) { auto is_forward = true; - auto task = CreateTask(step, mb, global_chunk_id, num_stages, total_global_chunks, is_forward); + auto task = CreateTask(step, mb, global_chunk_id, num_stages, total_global_chunks, is_forward, layout); schedule.push_back(task); } } @@ -107,8 +137,8 @@ std::vector PipelineParallelScheduler::Generate int global_chunk_id = (total_steps - 1 - step) - mb; if (global_chunk_id >= 0 && global_chunk_id < total_global_chunks) { auto is_forward = false; - auto task - = CreateTask(step + total_steps, mb, global_chunk_id, num_stages, total_global_chunks, is_forward); + auto task = CreateTask(step + total_steps, mb, global_chunk_id, num_stages, total_global_chunks, + is_forward, layout); schedule.push_back(task); } } @@ -127,12 +157,18 @@ std::vector PipelineParallelScheduler::Generate } std::vector -PipelineParallelScheduler::GenerateInterleaved1F1BSchedule(int n, int num_stages, int vpp_size) { +PipelineParallelScheduler::GenerateInterleaved1F1BSchedule(int n, int num_stages, int vpp_size, + const PipelineLayout &layout) { std::vector schedule; if (n <= 0 || num_stages <= 0 || vpp_size <= 0) { return schedule; } + if (layout.num_stages() != num_stages || layout.vpp_size() != vpp_size) { + throw PipelineLayoutError( + std::format("scheduler topology mismatch: requested stages={}, vpp={}, layout stages={}, vpp={}", + num_stages, vpp_size, layout.num_stages(), layout.vpp_size())); + } int total_global_chunks = num_stages * vpp_size; @@ -145,7 +181,8 @@ PipelineParallelScheduler::GenerateInterleaved1F1BSchedule(int n, int num_stages int forward_global_chunk = step - mb; if (forward_global_chunk >= 0 && forward_global_chunk < total_global_chunks) { auto is_forward = true; - auto task = CreateTask(step, mb, forward_global_chunk, num_stages, total_global_chunks, is_forward); + auto task + = CreateTask(step, mb, forward_global_chunk, num_stages, total_global_chunks, is_forward, layout); schedule.push_back(task); } } @@ -160,7 +197,8 @@ PipelineParallelScheduler::GenerateInterleaved1F1BSchedule(int n, int num_stages int forward_global_chunk = step - mb; if (forward_global_chunk >= 0 && forward_global_chunk < total_global_chunks) { auto is_forward = true; - auto task = CreateTask(step, mb, forward_global_chunk, num_stages, total_global_chunks, is_forward); + auto task + = CreateTask(step, mb, forward_global_chunk, num_stages, total_global_chunks, is_forward, layout); schedule.push_back(task); } @@ -169,7 +207,8 @@ PipelineParallelScheduler::GenerateInterleaved1F1BSchedule(int n, int num_stages if (backward_global_chunk >= 0 && backward_global_chunk < total_global_chunks) { auto is_forward = false; - auto task = CreateTask(step, mb, backward_global_chunk, num_stages, total_global_chunks, is_forward); + auto task + = CreateTask(step, mb, backward_global_chunk, num_stages, total_global_chunks, is_forward, layout); schedule.push_back(task); } } @@ -182,7 +221,8 @@ PipelineParallelScheduler::GenerateInterleaved1F1BSchedule(int n, int num_stages int backward_global_chunk = (total_global_chunks - 1) - (backward_step - mb); if (backward_global_chunk >= 0 && backward_global_chunk < total_global_chunks) { auto is_forward = false; - auto task = CreateTask(step, mb, backward_global_chunk, num_stages, total_global_chunks, is_forward); + auto task + = CreateTask(step, mb, backward_global_chunk, num_stages, total_global_chunks, is_forward, layout); schedule.push_back(task); } } @@ -199,11 +239,12 @@ float PipelineSchedule::StepMicroBatches(const std::vectorstage_index(); int vpp_size = global::GetVirtualPipelineParallelSize(); - auto schedule = PipelineParallelScheduler::GenerateGPipeSchedule(n, num_stages, vpp_size); + const auto &layout = global::GetPipelineLayout(); + auto schedule = PipelineParallelScheduler::GenerateGPipeSchedule(n, num_stages, vpp_size, layout); static bool has_printed = false; if (!has_printed && stage_idx == 0) { - PrintScheduleTable(schedule, n, num_stages, vpp_size); + PrintScheduleTable(schedule, n, num_stages, vpp_size, layout); has_printed = true; } @@ -234,7 +275,18 @@ float PipelineSchedule::StepMicroBatches(const std::vectorForwardOneChunk(inputs, task.local_chunk_idx); + if (time_stage) { + SynchronizeStageStream(stage_->device()); + const auto stage_end = std::chrono::steady_clock::now(); + const double elapsed_ms + = std::chrono::duration(stage_end - stage_start).count(); + LOG(INFO) << std::format("pipeline_stage_timing direction=forward stage={} chunk={} " + "microbatch={} elapsed_ms={:.3f}", + stage_idx, task.local_chunk_idx, mb, elapsed_ms); + } if (!task.is_last_chunk) { if (stage_->IsLastStage()) { @@ -255,7 +307,18 @@ float PipelineSchedule::StepMicroBatches(const std::vector(target_on_device)})[0]; loss = loss / n; } + const bool time_stage = StageTimingEnabled(); + const auto stage_start = std::chrono::steady_clock::now(); loss->Backward(); + if (time_stage) { + SynchronizeStageStream(stage_->device()); + const auto stage_end = std::chrono::steady_clock::now(); + const double elapsed_ms + = std::chrono::duration(stage_end - stage_start).count(); + LOG(INFO) << std::format("pipeline_stage_timing direction=backward stage={} chunk={} " + "microbatch={} elapsed_ms={:.3f}", + stage_idx, task.local_chunk_idx, mb, elapsed_ms); + } // Defer the loss D2H copy until after backward; reading it earlier would synchronize CUDA // between forward and backward. total_loss += static_cast(loss->To(Device()).DataPtr())[0]; @@ -265,7 +328,18 @@ float PipelineSchedule::StepMicroBatches(const std::vector(out_tensor->Dims(), out_tensor->Dtype(), out_tensor->GetDevice()); + const bool time_stage = StageTimingEnabled(); + const auto stage_start = std::chrono::steady_clock::now(); out_tensor->Backward(dummy_gradient); + if (time_stage) { + SynchronizeStageStream(stage_->device()); + const auto stage_end = std::chrono::steady_clock::now(); + const double elapsed_ms + = std::chrono::duration(stage_end - stage_start).count(); + LOG(INFO) << std::format("pipeline_stage_timing direction=backward stage={} chunk={} " + "microbatch={} elapsed_ms={:.3f}", + stage_idx, task.local_chunk_idx, mb, elapsed_ms); + } } } } diff --git a/infini_train/src/nn/parallel/pp/pipeline_stage.cc b/infini_train/src/nn/parallel/pp/pipeline_stage.cc index 25b5be235..bbb9158dc 100644 --- a/infini_train/src/nn/parallel/pp/pipeline_stage.cc +++ b/infini_train/src/nn/parallel/pp/pipeline_stage.cc @@ -1,6 +1,8 @@ #include "infini_train/include/nn/parallel/pp/pipeline_stage.h" +#include #include +#include #include "glog/logging.h" @@ -15,7 +17,23 @@ PipelineStage::PipelineStage(int stage_index /* pp_rank */, int num_stages /* pp std::vector> &&chunks) : stage_index_(stage_index), num_stages_(num_stages), prev_rank_(stage_index > 0 ? stage_index - 1 : -1), next_rank_(stage_index < num_stages - 1 ? stage_index + 1 : -1), recv_shape_(recv_shape), device_(device), - chunks_(std::move(chunks)) {} + chunks_(std::move(chunks)) { + if (num_stages_ <= 0) { + throw std::invalid_argument("PipelineStage num_stages must be positive"); + } + if (stage_index_ < 0 || stage_index_ >= num_stages_) { + throw std::invalid_argument( + std::format("PipelineStage stage_index={} is outside [0, {})", stage_index_, num_stages_)); + } + if (chunks_.empty()) { + throw std::invalid_argument("PipelineStage requires at least one local chunk"); + } + for (size_t local_chunk_idx = 0; local_chunk_idx < chunks_.size(); ++local_chunk_idx) { + if (chunks_[local_chunk_idx] == nullptr) { + throw std::invalid_argument(std::format("PipelineStage local chunk {} is null", local_chunk_idx)); + } + } +} std::vector> PipelineStage::ForwardOneChunk(const std::vector> &inputs, int local_chunk_idx) { diff --git a/pipeline_parallel_acceptance_report.md b/pipeline_parallel_acceptance_report.md new file mode 100644 index 000000000..657362d6e --- /dev/null +++ b/pipeline_parallel_acceptance_report.md @@ -0,0 +1,771 @@ +# Pipeline 自定义布局最终验收报告 + +## 1. 验收范围与总结果 + +### 1.1 验收目标 + +Pipeline 自定义布局把 Transformer 层、embedding、final norm、LM head 以及调度 chunk 的归属统一记录在 `PipelineLayout` 中。模型构建、`PipelineParallel` 包装器、GPipe/1F1B 调度器和 GPT-2/LLaMA3 checkpoint loader 都读取同一份布局，避免不同组件重复切分造成参数注册、通信顺序和 checkpoint 偏移不一致。 + +“通过标准”要求的核心功能已通过：统一数据结构、连续自定义层数、特殊模块放置、默认行为兼容、GPT-2/LLaMA3 代码路径、非法布局校验、CPU 单元测试、两 stage CUDA 训练、FP32/BF16 loss 对比和端到端性能记录均有证据。 + +“优秀标准”中除最后一项 PR review 外的四项已经补齐并完成验收：显式 vPP Chunk→Stage API 和非对称 vPP CUDA 运行通过；Megatron 风格重复层、特殊模块、Virtual Chunk 和空 stage 均通过 parser、CPU/CUDA 运行；均衡工具支持用户代价 CSV/Profiler 导出的层代价并输出预测改善；相同 PP、相同 micro-batch 的 stage timing、端到端耗时、吞吐和 bubble 已写入 CSV。最后一项 PR review 仍按“未完成”保留，因为没有真实外部 reviewer/approve 记录。 + +## 2. Pipeline 自定义布局使用指导 + +完整使用说明见 [`docs/pipeline_custom_layout_usage.md`](docs/pipeline_custom_layout_usage.md)。本节给出验收所需的参数、语法、默认行为、输入输出示例和错误排查方法。 + +### 2.1 参数配置 + +GPT-2 和 LLaMA3 均支持以下参数： + +```text +--pipeline_parallel=N +--virtual_pipeline_parallel=N +--pipeline_layer_partition=4,8,6,6 +--pipeline_layout='Etttttt|ttttttFH' +--pipeline_embedding_stage=0 +--pipeline_final_norm_stage=-1 +--pipeline_lm_head_stage=-1 +``` + +参数含义： + +| 参数 | 含义 | 默认/约束 | +|---|---|---| +| `pipeline_parallel` | Pipeline stage 数，也就是 PP size | 必须为正整数 | +| `virtual_pipeline_parallel` | 每个 physical stage 的 virtual chunk 数 | 必须为正整数；连续 partition 当前要求为 `1` | +| `pipeline_layer_partition` | 每个 stage 的连续 Transformer 层数，逗号分隔 | 为空时使用默认均匀布局；长度必须等于 PP size | +| `pipeline_layout` | Megatron 风格显式布局表达式 | 与 `pipeline_layer_partition` 互斥 | +| `pipeline_embedding_stage` | Embedding 所属 stage | 默认 `0` | +| `pipeline_final_norm_stage` | Final Norm 所属 stage | `-1` 表示最后一个 stage | +| `pipeline_lm_head_stage` | LM Head 所属 stage | `-1` 表示最后一个 stage | + +入口会在启动期检查： + +- `pipeline_parallel`、`virtual_pipeline_parallel` 和模型层数必须为正； +- stage placement 必须落在 `[0, pipeline_parallel)`； +- `pipeline_layer_partition` 与 `pipeline_layout` 不能同时配置； +- 连续 partition 的条目数必须等于 PP size； +- 连续 partition 的层数总和必须等于模型层数； +- `pipeline_layer_partition` 当前不能与 `virtual_pipeline_parallel > 1` 一起使用； +- 显式 layout 的 stage 数、每个 stage 的 chunk 数和 `t` 总数必须匹配。 + +### 2.2 默认行为 + +不配置 `pipeline_layer_partition` 和 `pipeline_layout` 时，调用 `PipelineLayout::BuildDefault`： + +- 层按全部 global chunks 均匀分配； +- 余数优先放在前面的 chunks； +- global chunk 编号按 `local_chunk_id * num_stages + stage_id` 交错排列； +- Embedding 默认属于 stage 0； +- Final Norm 和 LM Head 默认属于最后一个 stage； +- 现有 GPipe/1F1B scheduler 继续使用兼容的 global chunk 编号； +- vPP 默认布局按 stage/virtual chunk 组织，不要求用户提供轮转表。 + +例如 `num_layers=24`、`pipeline_parallel=4`、`virtual_pipeline_parallel=1` 时，默认布局为： + +```text +stage 0: layers [0, 6) +stage 1: layers [6, 12) +stage 2: layers [12, 18) +stage 3: layers [18, 24) +``` + +### 2.3 连续非均匀层数布局 + +连续布局使用： + +```text +--pipeline_parallel=4 +--pipeline_layer_partition=4,8,6,6 +``` + +对于 24 层模型，输出的 stage ownership 为： + +```text +stage 0: layers [0, 4) +stage 1: layers [4, 12) +stage 2: layers [12, 18) +stage 3: layers [18, 24) +``` + +对应的闭区间写法为： + +```text +stage 0: layers 0-3 +stage 1: layers 4-11 +stage 2: layers 12-17 +stage 3: layers 18-23 +``` + +两 stage 验收使用的连续自定义布局为： + +```text +--pipeline_parallel=2 +--pipeline_layer_partition=6,6 +``` + +它把 GPT-2 的 12 层切成： + +```text +stage 0: layers [0, 6) +stage 1: layers [6, 12) +``` + +连续 partition 当前要求 `virtual_pipeline_parallel=1`。如果需要 vPP 或显式的 chunk→stage 映射，应使用 `pipeline_layout`。 + +### 2.4 Megatron 风格布局语法 + +`--pipeline_layout` 的语法元素： + +| 符号 | 含义 | +|---|---| +| `t`/`T` | 一个 Transformer layer | +| `E` | Embedding | +| `F` | Final Norm | +| `H`/`L` | LM Head | +| `|` | physical stage 分隔符 | +| `,` | 同一 stage 内的 virtual chunk 分隔符 | +| `(expr)*N` | 将括号内表达式重复 N 次 | + +示例： + +```bash +--pipeline_parallel=2 +--pipeline_layout='Etttttt|ttttttFH' +``` + +解析为： + +```text +stage 0: Embedding + 6 Transformer layers +stage 1: 6 Transformer layers + Final Norm + LM Head +``` + +vPP 示例： + +```bash +--pipeline_parallel=2 +--virtual_pipeline_parallel=2 +--pipeline_layout='tt,tt|tt,tt' +``` + +该表达式表示每个 stage 有两个 virtual chunks；所有 `t` 合计 8 层。global chunk 编号按 chunk-major 顺序生成： + +```text +global chunk 0: stage 0, local chunk 0 +global chunk 1: stage 1, local chunk 0 +global chunk 2: stage 0, local chunk 1 +global chunk 3: stage 1, local chunk 1 +``` + +重复表达式示例： + +```text +(tt)*2|tt|tt +``` + +括号表达式会在 stage 分隔和 chunk 解析前展开。显式布局允许表达空 chunk/空 +stage；运行时会为该位置创建空的 `TransformerChunk`，原样传递激活和梯度。 + +### 2.5 输入输出示例 + +#### 示例 A：命令行连续布局 + +输入： + +```bash +torchrun --standalone --no-python --nproc_per_node=2 build_cuda/gpt2 \ + --device=cuda \ + --pipeline_parallel=2 \ + --pipeline_layer_partition=6,6 \ + --input_bin=/root/InfiniTrain/data/gpt2/tiny_shakespeare_train.bin \ + --llmc_filepath=/root/InfiniTrain/data/gpt2/gpt2_124M.bin \ + --overfit_single_batch=true \ + --num_iteration=1 \ + --dtype=float32 +``` + +关键输出： + +```text +world_size = 2, config: {DP=1, TP=1, PP=2} +=== Schedule Table === +Forward global chunk 0 ... stage 0 +Forward global chunk 1 ... stage 1 +Backward global chunk 1 ... stage 1 +Backward global chunk 0 ... stage 0 +step 1/1 | train loss 5.356194 | ... PP=2 +``` + +#### 示例 B：显式特殊模块布局 + +输入： + +```bash +--pipeline_parallel=2 +--pipeline_layout='Etttttt|ttttttFH' +``` + +输出语义： + +```text +stage 0 owns embedding +stage 0 owns transformer layers 0-5 +stage 1 owns transformer layers 6-11 +stage 1 owns final norm and lm head +``` + +远程 CUDA 运行得到： + +```text +train loss 6.407298 +exit code 0 +``` + +#### 示例 C：自动均衡建议 + +输入： + +```bash +./build_cpu/pipeline_layout_suggest \ + --num_layers=8 \ + --pp_size=4 \ + --layer_cost=1,1,1,1,4,1,1,1 +``` + +输出： + +```text +pipeline_layer_partition=3,1,1,3 +``` + +该工具使用用户提供的每层计算代价生成连续 partition；它目前不是从 CUDA profiler 文件自动导入代价。 + +### 2.6 错误排查方法 + +布局错误统一抛出 `PipelineLayoutError`，错误消息包含具体字段、stage、层数或 offset。常见错误和处理方式如下： + +| 错误信息 | 原因 | 处理方式 | +|---|---|---| +| `pipeline_layer_partition must not be empty` | 直接调用 parser 时传入空字符串 | 不需要自定义 partition 时改用 `BuildDefault`/留空 flag | +| `stage count mismatch: pipeline_parallel=4, partition_entries=3` | partition 条目数与 PP size 不一致 | 补齐或删除 partition 条目 | +| `layer sum mismatch: num_layers=23, partition_sum=24` | partition 层数总和不等于模型层数 | 重新计算各 stage 层数 | +| `pipeline_layer_partition contains invalid character '-'` | 使用了负号、空格、字母或其他非法字符 | 只使用非负整数和逗号，例如 `4,8,6,6` | +| `custom layer partition is not supported with vpp_size=2` | 连续 partition 与 vPP>1 同时使用 | 改用 `pipeline_layout` 显式描述 chunk | +| `pipeline_layer_partition and pipeline_layout are mutually exclusive` | 同时设置两种布局来源 | 二选一 | +| `pipeline_layout stage count mismatch` | `|` 分隔出的 stage 数不等于 PP size | 检查 `|` 数量和 `pipeline_parallel` | +| `pipeline_layout chunk count mismatch` | 某 stage 的 `,` chunk 数不等于 vPP size | 每个 stage 补齐相同数量的 chunk | +| `pipeline_layout layer count mismatch` | `t` 总数不等于 checkpoint 的 `n_layer` | 按 checkpoint header 修正表达式 | +| `pipeline_layout contains unknown symbol` | 使用了非 `t/E/F/H/L/|/,/()` 的字符 | 删除未知字符 | +| `stage N has no chunks; the current pipeline transport requires every stage to have at least one chunk` | layout 没有为该 stage 保留 chunk id | 使用 `Etttttt||ttttttFH` 这类空表达式（保留空 chunk），不要删除该 stage；当前 transport 通过空 `TransformerChunk` 传递数据 | +| `embedding_stage ... must be boundary stage 0` / `final_norm_stage ... must be last stage` | 默认 policy 要求特殊模块位于边界 | 保持 embedding 在首 stage、final norm/lm head 在末 stage，或显式改变 policy/transport | + +建议排查顺序： + +1. 先确认 `pipeline_parallel`、`virtual_pipeline_parallel` 和 checkpoint header 中的 `n_layer`； +2. 再检查两种 layout flag 是否互斥； +3. 对连续 partition 检查条目数和总和； +4. 对字符串布局检查 stage 数、每 stage chunk 数和 `t` 数； +5. 检查特殊模块是否位于当前 transport 支持的边界； +6. 重新运行 CPU parser/scheduler 单测； +7. 最后运行一迭代 CUDA smoke test。 + +## 3. 单元测试、端到端测试代码与测试日志 + +### 3.1 测试代码清单 + +| 类型 | 文件 | 覆盖内容 | +|---|---|---| +| 布局单测 | [`tests/parallel/test_pipeline_layout.cc`](tests/parallel/test_pipeline_layout.cc) | 默认布局、连续 partition、查询接口、特殊模块、非法输入 | +| scheduler 单测 | [`tests/parallel/test_pipeline_scheduler_layout.cc`](tests/parallel/test_pipeline_scheduler_layout.cc) | GPipe/1F1B task 使用 layout ownership、拓扑校验 | +| 分块单测 | [`tests/parallel/test_pipeline_parallel_chunking.cc`](tests/parallel/test_pipeline_parallel_chunking.cc) | 首/中/末 stage 模块组合、vPP local chunk 顺序、stage mismatch | +| 均衡单测 | [`tests/parallel/test_pipeline_layout_balance.cc`](tests/parallel/test_pipeline_layout_balance.cc) | 均匀代价、重层隔离、非法代价 | +| Megatron parser 单测 | [`tests/parallel/test_pipeline_layout_megatron_style.cc`](tests/parallel/test_pipeline_layout_megatron_style.cc) | 特殊模块、重复表达式、vPP chunk、非法字符串 | +| 端到端 smoke | [`scripts/test_pipeline_custom_layout.sh`](scripts/test_pipeline_custom_layout.sh) | PP=1 baseline 与 PP=2 `6,6` custom 的单步 loss 对比 | +| 性能对比 | [`scripts/compare_pipeline_layout_perf.sh`](scripts/compare_pipeline_layout_perf.sh) | elapsed、吞吐和理想 bubble CSV | +| loss 比较器 | [`scripts/compare_loss.py`](scripts/compare_loss.py) | FP32 `1e-5`、BF16 `1e-2` 阈值下逐 step loss 比较 | +| 均衡建议工具 | [`tools/pipeline_layout_suggest.cc`](tools/pipeline_layout_suggest.cc) | 根据用户提供层代价输出连续 partition | +| stage 性能对比 | [`scripts/compare_pipeline_stage_perf.sh`](scripts/compare_pipeline_stage_perf.sh) | 相同 PP 下采集每 stage forward/backward timing、吞吐、bubble 和 imbalance | + +### 3.2 构建与单元测试命令 + +远程 GPU 服务器硬件和工具链： + +```text +6 × NVIDIA GeForce RTX 4090 D +CUDA 12.8.61 +nvcc: /usr/local/cuda-12.8/bin/nvcc +G++ 13.3.0 +``` + +CUDA 构建： + +```bash +cmake --build build_cuda -j8 +``` + +结果：`gpt2`、`llama3`、`infini_run` 及相关 CUDA 目标构建成功，退出码为 0。 + +CPU 构建/测试： + +```bash +cmake --build build_cpu -j8 +build_cpu/tests/parallel/test_pipeline_layout_cpu --gtest_color=no +build_cpu/tests/parallel/test_pipeline_scheduler_layout_cpu --gtest_color=no +build_cpu/tests/parallel/test_pipeline_parallel_chunking_cpu --gtest_color=no +build_cpu/tests/parallel/test_pipeline_layout_balance_cpu --gtest_color=no +build_cpu/tests/parallel/test_pipeline_layout_megatron_style_cpu --gtest_color=no +``` + +也可以使用 CTest： + +```bash +ctest --test-dir build_cpu -R 'test_pipeline_(layout|scheduler_layout|parallel_chunking|layout_balance|layout_megatron_style)_cpu' --output-on-failure +``` + +### 3.3 单元测试结果 + +| 测试目标 | 用例数 | 结果 | +|---|---:|---| +| `test_pipeline_layout_cpu` | 16 | 16 passed | +| `test_pipeline_scheduler_layout_cpu` | 5 | 5 passed | +| `test_pipeline_parallel_chunking_cpu` | 6 | 6 passed | +| `test_pipeline_layout_balance_cpu` | 3 | 3 passed | +| `test_pipeline_layout_megatron_style_cpu` | 8 | 8 passed | +| **合计** | **38** | **38 passed** | + +重点测试事实： + +- `BuildDefault(8, 2, 2)` 生成交错 global chunk id； +- `BuildContiguous({2,4,3,3})` 正确生成连续层区间； +- `stage_of_layer`、`chunk_of_layer`、`local_layer_index` 与 ownership 一致； +- 首 stage 包含 `TransformerFirstStage`，中间 stage 仅包含 transformer chunk，末 stage 包含 `TransformerLastStage`； +- scheduler 的每个 task 的 stage/local chunk 来自 `PipelineLayout`，而不是重新推导； +- parser 拒绝空 token、非法字符、括号不匹配、未知符号和层数不匹配； +- `SuggestBalancedPartition(8,4,{1,1,1,1,4,1,1,1})` 输出 `{3,1,1,3}`。 +- `BuildExplicit` 拒绝缺失 chunk、重复 global id、重复 local chunk 和非法 stage； +- 非对称 vPP `t,tt|tt,t` 正确保留每个 chunk 的显式 stage/local chunk ownership； +- `Etttttt||ttttttFH` 的空 stage 构造成 pass-through chunk。 + +### 3.4 CUDA 端到端测试命令 + +GPT-2 smoke test： + +```bash +BUILD_DIR=build_cuda \ +INPUT_BIN=/root/InfiniTrain/data/gpt2/tiny_shakespeare_train.bin \ +INPUT_VAL_BIN=/root/InfiniTrain/data/gpt2/tiny_shakespeare_val.bin \ +TOKENIZER_BIN=/root/InfiniTrain/data/gpt2/gpt2_tokenizer.bin \ +LLMC_FILE=/root/InfiniTrain/data/gpt2/gpt2_124M.bin \ +LAUNCHER='torchrun --standalone --no-python' \ +scripts/test_pipeline_custom_layout.sh gpt2 cuda +``` + +性能记录： + +```bash +BUILD_DIR=build_cuda \ +INPUT_BIN=/root/InfiniTrain/data/gpt2/tiny_shakespeare_train.bin \ +LLMC_FILE=/root/InfiniTrain/data/gpt2/gpt2_124M.bin \ +scripts/compare_pipeline_layout_perf.sh gpt2 cuda +``` + +测试均为单次 optimizer iteration 的验收 smoke test，不替代长时间收敛训练。两 stage 测试使用真实 NCCL/CUDA 进程和 checkpoint，不是仅在单进程中模拟 stage。 + +优秀标准补验： + +```bash +BUILD_DIR=build_cuda \ +LAUNCHER='torchrun --standalone --no-python' \ +TOTAL_BATCH_SIZE=128 \ +scripts/compare_pipeline_stage_perf.sh gpt2 cuda +``` + +该脚本固定 PP=2、4 个 micro-batch，分别运行默认均匀布局和 +`--pipeline_layer_partition=4,8`，并在 CUDA stream 同步后输出每个 stage/chunk 的 +forward/backward timing。 + +### 3.5 FP32/BF16 结果 + +#### FP32：连续 `6,6` 布局 + +| 运行 | PP | loss | 与 baseline 差值 | 阈值 | 结果 | +|---|---:|---:|---:|---:|---| +| baseline | 1 | 5.356194 | — | — | 通过 | +| custom | 2 | 5.356194 | `0` | `1e-5` | 通过 | + +脚本摘要： + +```text +fp32: 1/1 test cases passed +Total: 1/1 test cases passed +Pipeline custom-layout smoke test passed +``` + +#### BF16：连续 `6,6` 布局 + +| 运行 | PP | loss | 与 baseline 差值 | 阈值 | 结果 | +|---|---:|---:|---:|---:|---| +| baseline | 1 | 5.309796 | — | — | 通过 | +| custom | 2 | 5.309796 | `0` | `1e-2` | 通过 | + +#### 显式 Megatron 布局 + +```text +--pipeline_layout='Etttttt|ttttttFH' +train loss 6.407298 +exit code 0 +``` + +#### 前向、loss、梯度链路 + +FP32 单卡日志明确包含： + +```text +start forward +finish model forward, start loss forward +finish loss forward +start backward +finish backward +step 1/1 | train loss ... +``` + +两 stage 日志包含 schedule table、forward、backward 和最终 step loss；BF16 两 stage 运行也以退出码 0 完成 step。 + +因此，本次验收可以证明： + +- 前向成功； +- loss forward 成功； +- backward 成功； +- optimizer step 成功； +- baseline/custom loss 差值在 FP32/BF16 阈值内，实际差值均为 0。 + +当前已保存的验收脚本比较的是 loss，而不是逐参数 gradient tensor。报告不伪造 +`max_abs_grad_diff` 数值；本轮新增的优秀标准补验聚焦布局、stage timing、吞吐、 +bubble 和 cost-aware 负载分析，不改变这一梯度量化边界。 + +## 4. 项目报告 + +### 4.1 数据结构设计 + +统一布局定义在 [`infini_train/include/nn/parallel/pipeline_layout.h`](infini_train/include/nn/parallel/pipeline_layout.h) 和 [`infini_train/src/nn/parallel/pipeline_layout.cc`](infini_train/src/nn/parallel/pipeline_layout.cc)。 + +| 类型 | 作用 | +|---|---| +| `PipelineLayout` | 保存完整模型层、stage、vPP、chunk、特殊模块和校验 policy，是所有消费者的唯一布局来源 | +| `LayerRange` | 半开区间 `[start,end)`，提供 `size()` 和 `contains(layer_id)` | +| `ChunkLayout` | 记录 `global_chunk_id`、`stage_id`、`local_chunk_id` 和层区间 | +| `StageLayout` | 记录 stage id 以及该 stage 持有的 global chunk id 列表 | +| `LayerLocation` | 反向索引，记录 layer 对应 stage、global/local chunk、chunk 内索引和 stage 内扁平索引 | +| `SpecialModulePlacement` | 独立记录 Embedding、Final Norm、LM Head 的 stage | +| `PipelineLayoutPolicy` | 控制是否允许空 stage、是否要求连续执行、是否要求特殊模块位于边界 | +| `PipelineLayoutError` | 统一的可捕获布局错误类型，继承 `std::runtime_error` | + +内部建立 `layer_id -> LayerLocation` 索引，避免模型构造、checkpoint loader 和 scheduler 各自重新计算层归属。`StageLayout.global_chunk_ids` 保留本地 chunk 顺序，`ChunkLayout.global_chunk_id` 兼容既有 GPipe/1F1B 调度编号。 + +### 4.2 接口设计 + +核心构造接口： + +```text +PipelineLayout::BuildDefault(...) +PipelineLayout::BuildContiguous(...) +PipelineLayout::BuildExplicit(...) +PipelineLayout::BuildPipelineLayout(...) +PipelineLayout::ParseLayerPartition(...) +PipelineLayout::ParseMegatronStyleLayout(...) +PipelineLayout::SuggestBalancedPartition(...) +``` + +核心查询接口： + +```text +num_layers() +num_stages() +vpp_size() +stage(stage_id) +chunk(global_chunk_id) +chunks() +locate_layer(layer_id) +stage_of_layer(layer_id) +chunk_of_layer(layer_id) +local_layer_index(stage_id, layer_id) +owns(SpecialModule, stage_id) +special_modules() +policy() +ToString() +``` + +核心校验接口： + +```text +Validate() +ValidateForCurrentPipelineTransport() +``` + +调度器接口： + +```text +PipelineParallelScheduler::CreateTask(...) +PipelineParallelScheduler::GenerateGPipeSchedule(...) +PipelineParallelScheduler::GenerateInterleaved1F1BSchedule(...) +``` + +所有 scheduler 接口都接收 `const PipelineLayout&`，并检查传入的 stage 数、vPP 数和 layout topology 一致。 + +### 4.3 关键实现说明 + +#### 布局安装时机 + +GPT-2/LLaMA3 入口在随机初始化路径中先做 layout preflight；checkpoint loader 在读取 checkpoint header 后，用真实 `n_layer` 再构造一次 layout，并在 `TransformerModel` 构造前执行： + +```cpp +global::InstallPipelineLayout(layout); +``` + +这样模型注册模块时不会再看到空的默认 layout。 + +#### TransformerModel 模块注册 + +`TransformerModel` 查询全局布局决定每个 transformer chunk 的层范围，并按 `owns()` 决定： + +- `__pp_first_stage` 是否包含 Embedding； +- `__pp_last_stage` 是否包含 Final Norm； +- `__pp_last_stage` 是否包含 LM Head； +- 每个 local chunk 的 transformer 层模块。 + +#### PipelineParallel 分块 + +`PipelineParallel` 不再用旧的 `GetStageInfo` 重新推导自定义层数，而是读取： + +```text +layout.stage(rank).global_chunk_ids +layout.chunk(global_chunk_id) +``` + +并按 layout 顺序创建 local chunks。首、中、末 stage 的模块组合由同一份 ownership 结果决定。 + +#### Scheduler task + +GPipe 和 interleaved 1F1B scheduler 根据 `layout.chunk(global_chunk_id)` 设置： + +- `stage_id`； +- `local_chunk_idx`； +- `is_first_chunk`； +- `is_last_chunk`。 + +global chunk 使用 chunk-major 兼容编号，保持旧调度行为。 + +#### Checkpoint ownership 和 canonical key + +checkpoint loader 根据已安装 layout 查询： + +```text +owns(kEmbedding, pp_rank) +owns(kFinalNorm, pp_rank) +owns(kLMHead, pp_rank) +stage_of_layer(layer_id) +``` + +即使当前 rank 不拥有某个参数，也会继续 seek 对应 checkpoint 字节，避免文件读取位置错位。LM Head 使用 canonical checkpoint key： + +```text +transformer.lm_head.weight +``` + +GPT-2 和 LLaMA3 均采用同一套 layout ownership 逻辑。 + +#### 非法输入处理 + +布局 parser 使用 `PipelineLayoutError` 报告 stage 数、层数、字符位置、空 token、重复特殊模块和 transport 限制。checkpoint 加载异常不会转为随机初始化成功，避免“错误被吞掉但验收看起来通过”。 + +### 4.4 兼容性说明 + +| 场景 | 兼容性 | 说明 | +|---|---|---| +| 未配置自定义布局 | 通过 | 使用 `BuildDefault`，保持均匀层分配和默认特殊模块边界 | +| GPipe | 通过 | scheduler 单测和两 stage CUDA schedule table 已验证 | +| 1F1B | 通过（单测） | interleaved scheduler task ownership 单测通过 | +| vPP 默认布局 | 通过（单测） | 默认 vPP global chunk 编号和 local chunk 顺序通过 | +| vPP 显式布局 | 通过 | 非对称 `Ett,tttt|tttt,ttFH` 的真实 PP=2、vPP=2 CUDA step 退出码 0；parser/ownership 单测通过 | +| GPT-2 | 通过 | CUDA checkpoint、FP32/BF16 两 stage smoke 已运行 | +| LLaMA3 | 代码路径通过 | flags、preflight、loader 与 GPT-2 对齐；本次保存的 CUDA 端到端日志为 GPT-2 | +| TP/SP/DP | 保持接口兼容 | 布局按 PP stage 组织，日志显示 `TP=1, DP=1, SP=1` 验收；未在本轮组合矩阵中遍历多 TP/DP | +| 空 stage | 通过 | `Etttttt||ttttttFH` 的真实 PP=3 CUDA step 退出码 0；空 `TransformerChunk` 原样传递激活/梯度 | +| 非边界特殊模块 | policy 可校验 | 默认 `require_boundary_special_modules=true`；当前 transport 对首/末 stage 仍有既有语义限制 | + +### 4.5 不同布局下的正确性 + +#### 默认均匀布局 + +单卡 FP32 baseline： + +```text +PP=1 +train loss 5.358113（早期 baseline 记录） +96.96 ms +660 tok/s +exit code 0 +``` + +最终 smoke 运行的同一数据配置 baseline： + +```text +PP=1 +train loss 5.356194 +88.48 ms +2893 tok/s +exit code 0 +``` + +两组记录使用的 batch/sequence 和脚本配置不同，不能直接互相作性能排名；loss 对比只使用同一 smoke 流程内的 baseline/custom 配对值。 + +#### 连续 `6,6` 自定义布局 + +```text +baseline loss = 5.356194 +custom loss = 5.356194 +abs diff = 0 +FP32 threshold = 1e-5 +``` + +BF16： + +```text +baseline loss = 5.309796 +custom loss = 5.309796 +abs diff = 0 +BF16 threshold = 1e-2 +``` + +#### 显式 `Etttttt|ttttttFH` + +```text +stage 0: E + tttttt +stage 1: tttttt + F + H +train loss 6.407298 +exit code 0 +``` + +#### stage/chunk ownership + +两 stage schedule table： + +```text +0 | Forward | global chunk 0 | local chunk 0 | stage 0 +1 | Forward | global chunk 1 | local chunk 0 | stage 1 +2 | Backward | global chunk 1 | local chunk 0 | stage 1 +3 | Backward | global chunk 0 | local chunk 0 | stage 0 +``` + +该顺序与 `PipelineLayout` 的 chunk ownership 一致，且不再依赖旧的均匀切分推导。 + +### 4.6 Pipeline 负载分析 + +单卡/双 stage 的基础性能数据见 [`targets/pipeline_layout_perf.csv`](targets/pipeline_layout_perf.csv)；相同 PP degree 下的 stage timing 数据见 [`targets/pipeline_stage_perf.csv`](targets/pipeline_stage_perf.csv)： + +```csv +layout,pipeline_parallel,elapsed_ms,tok_per_s,bubble_ratio +baseline,1,95.33,336.0,0.0 +custom,2,143.87,222.0,0.5 +``` + +测试配置： + +```text +sequence_length = 32 +total_batch_size = 32 +micro_batches = 1 +``` + +理想 GPipe bubble 估算： + +```text +bubble = (pipeline_parallel - 1) / + (micro_batches + pipeline_parallel - 1) +``` + +结果： + +| 布局 | PP | step wall time | 吞吐 | 理想 bubble | +|---|---:|---:|---:|---:| +| baseline | 1 | 95.33 ms | 336 tok/s | 0.0 | +| custom | 2 | 143.87 ms | 222 tok/s | 0.5 | + +解读： + +- 该配置只有 1 个 micro-batch，PP=2 的理论 bubble 为 0.5，属于预期的流水线填充/排空开销； +- 双 stage 端到端 step 比单卡 baseline 慢，不代表自定义布局逻辑错误，因为比较同时改变了 PP degree 和通信路径； +- 自定义布局的价值在于把不同计算代价的层重新分配到 stage，而不是保证任何 PP=2 配置都比 PP=1 快； +- `elapsed_ms` 和 `tok_per_s` 是真实端到端 step wall time/吞吐； +- `bubble_ratio` 是按 micro-batch 数和 stage 数计算的理想调度估计，不是 Nsight 或 CUDA kernel profiler 的空隙比例； +- 当前日志包含 rank/stage、schedule table 和端到端时间，没有独立记录每个 stage 的 CUDA kernel 时间，因此不能把 `elapsed_ms` 拆写成“stage 0=... ms、stage 1=... ms”。 + +要证明自定义布局改善负载不均衡，需要在相同 PP、相同 micro-batch 和相同输入下采集每层/每 stage 时间，再比较均匀 partition 与 cost-aware partition 的最大 stage 时间和吞吐。本轮已完成该补验，并另外提供用户/profiler 代价 CSV 的预测分析。 + +本轮已完成上述相同 PP 的补验。配置为 PP=2、4 个 micro-batch、同一 GPT-2 +checkpoint 和同一输入，stage timing 在 CUDA stream 同步后记录： + +```csv +layout,pipeline_parallel,micro_batches,elapsed_ms,tok_per_s,bubble_ratio,stage0_compute_ms,stage1_compute_ms,max_stage_compute_ms,stage_imbalance_ratio +uniform,2,4,202.07,633.0,0.2,262.852,171.503,262.852,1.5326 +custom,2,4,223.55,573.0,0.2,278.645,190.415,278.645,1.4634 +``` + +实测结论： + +- uniform 与 custom 的 PP、micro-batch、输入和 checkpoint 相同； +- 两次 loss 均为 `5.934462`，说明布局切换没有改变该 step 的数值结果； +- custom 的 stage imbalance ratio 从 `1.5326` 降到 `1.4634`，降低约 `4.5%`； +- 本次 custom 的端到端吞吐为 `573 tok/s`，低于 uniform 的 `633 tok/s`，原因包括自定义分片和本次 timing 同步开销；“改善负载不均衡”不等价于在每个配置下吞吐必然上升； +- 对用户/profiler 代价样例 `1,1,1,1,4,1,1,1`，工具输出 uniform 最大 stage cost `5`、建议布局最大 stage cost `4`，预测降低 `20%`。该结果为可复现的 cost-aware 证据。 + +## 8. 版本、修改和提交信息 + +本次 Pipeline 修复涉及： + +- [`example/gpt2/main.cc`](example/gpt2/main.cc) +- [`example/gpt2/checkpoint_loader.cc`](example/gpt2/checkpoint_loader.cc) +- [`example/llama3/main.cc`](example/llama3/main.cc) +- [`example/llama3/checkpoint_loader.cc`](example/llama3/checkpoint_loader.cc) +- [`scripts/compare_pipeline_layout_perf.sh`](scripts/compare_pipeline_layout_perf.sh) +- [`scripts/compare_pipeline_stage_perf.sh`](scripts/compare_pipeline_stage_perf.sh) +- [`docs/pipeline_custom_layout_usage.md`](docs/pipeline_custom_layout_usage.md) +- [`infini_train/include/nn/parallel/pipeline_layout.h`](infini_train/include/nn/parallel/pipeline_layout.h) +- [`infini_train/src/nn/parallel/pipeline_layout.cc`](infini_train/src/nn/parallel/pipeline_layout.cc) +- [`infini_train/src/nn/parallel/pp/pipeline_schedule.cc`](infini_train/src/nn/parallel/pp/pipeline_schedule.cc) +- [`tests/parallel/test_pipeline_scheduler_layout.cc`](tests/parallel/test_pipeline_scheduler_layout.cc) +- [`tests/parallel/test_pipeline_layout_megatron_style.cc`](tests/parallel/test_pipeline_layout_megatron_style.cc) +- [`tests/parallel/test_pipeline_parallel_chunking.cc`](tests/parallel/test_pipeline_parallel_chunking.cc) +- [`tools/pipeline_layout_suggest.cc`](tools/pipeline_layout_suggest.cc) +- [`pipeline_parallel_acceptance_report.md`](pipeline_parallel_acceptance_report.md) +- [`targets/pipeline_layout_perf.csv`](targets/pipeline_layout_perf.csv) +- [`targets/pipeline_stage_perf.csv`](targets/pipeline_stage_perf.csv) +- `artifacts/pipeline_custom_layout/**` +- `artifacts/pipeline_layout_perf/**` + +关键修复点： + +1. 补齐 GPT-2/LLaMA3 的 Pipeline flags； +2. checkpoint header 驱动真实 layout 构造； +3. `TransformerModel` 构造前安装全局 layout； +4. checkpoint loader 使用已安装 layout 查询 ownership； +5. LM Head 使用 canonical key； +6. 删除加载失败后随机初始化兜底； +7. 修复含空格 launcher 的性能脚本解析； +8. 新增 `BuildExplicit` 显式 chunk→stage ownership API； +9. 支持空 stage pass-through、非对称 vPP chunk 和 scheduler ownership 验证； +10. 新增 CUDA stage timing 开关和相同 PP 的性能对比脚本； +11. 新增一列/两列代价 CSV 输入和 cost-aware 预测改善输出； +12. 增加使用指导、测试日志、性能 CSV 和本最终验收报告。 + +提交检查以以下命令为准： + +```bash +git status --short --branch +git log -1 --oneline +git diff --check +``` + +本报告不硬编码 amend 后的 commit hash，避免报告内容与最终 Git 对象号不一致；交付时以最终 `git log -1` 输出为准。 diff --git a/scripts/compare_pipeline_layout_perf.sh b/scripts/compare_pipeline_layout_perf.sh new file mode 100755 index 000000000..955450888 --- /dev/null +++ b/scripts/compare_pipeline_layout_perf.sh @@ -0,0 +1,68 @@ +#!/usr/bin/env bash +# Run one baseline and one custom-layout smoke test, then emit a compact CSV +# with throughput and the idealized pipeline bubble estimate. +set -euo pipefail +if [[ $# -ne 2 ]]; then + echo "Usage: $0 " >&2; exit 2 +fi +model="$1"; device="$2" +[[ "$model" == gpt2 || "$model" == llama3 ]] || { echo "unsupported model: $model" >&2; exit 2; } +[[ "$device" == cpu || "$device" == cuda ]] || { echo "unsupported device: $device" >&2; exit 2; } +repo_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +build_dir="${BUILD_DIR:-$repo_dir/build}" +exe="${EXE:-$build_dir/$model}" +out_dir="${OUT_DIR:-$repo_dir/artifacts/pipeline_layout_perf/$model}" +csv="${PERF_CSV:-$repo_dir/targets/pipeline_layout_perf.csv}" +mkdir -p "$out_dir/baseline" "$out_dir/custom" "$(dirname "$csv")" + +if [[ "$model" == gpt2 ]]; then default_sequence=32; else default_sequence=4; fi +batch_size="${BATCH_SIZE:-1}" +sequence_length="${SEQUENCE_LENGTH:-$default_sequence}" +total_batch_size="${TOTAL_BATCH_SIZE:-$((batch_size * sequence_length))}" +if (( total_batch_size % (batch_size * sequence_length) != 0 )); then + echo "TOTAL_BATCH_SIZE must be divisible by BATCH_SIZE*SEQUENCE_LENGTH" >&2; exit 2 +fi +micro_batches=$((total_batch_size / (batch_size * sequence_length))) + +common=(--device="$device" --input_bin="${INPUT_BIN:-$repo_dir/data/$model/tiny_shakespeare_train.bin}" \ + --llmc_filepath="${LLMC_FILE:-$repo_dir/data/$model/$([[ $model == gpt2 ]] && echo gpt2_124M.bin || echo llama3.2_1B_fp32.bin)}" \ + --overfit_single_batch=true --num_iteration=1 --dtype="${DTYPE:-float32}" \ + --batch_size="$batch_size" --sequence_length="$sequence_length" --total_batch_size="$total_batch_size") +[[ -e "${INPUT_VAL_BIN:-$repo_dir/data/$model/tiny_shakespeare_val.bin}" ]] && common+=(--input_val_bin="${INPUT_VAL_BIN:-$repo_dir/data/$model/tiny_shakespeare_val.bin}") +[[ "$model" == gpt2 && -e "${TOKENIZER_BIN:-$repo_dir/data/gpt2/gpt2_tokenizer.bin}" ]] && common+=(--tokenizer_bin="${TOKENIZER_BIN:-$repo_dir/data/gpt2/gpt2_tokenizer.bin}") +launcher_spec="${LAUNCHER:-$build_dir/infini_run}" +run_case() { + local name="$1" pp="$2" layout="$3" log="$4"; local -a cmd + if [[ "$launcher_spec" == direct ]]; then + cmd=("$exe") + else + read -r -a launcher_command <<< "$launcher_spec" + cmd=("${launcher_command[@]}" --nnodes=1 --nproc_per_node="$pp" "$exe") + fi + cmd+=("${common[@]}" --pipeline_parallel="$pp") + [[ -n "$layout" ]] && cmd+=(--pipeline_layout="$layout") + echo "[$name] ${cmd[*]}" + "${cmd[@]}" 2>&1 | tee "$log" +} +run_case baseline 1 "" "$out_dir/baseline/run.log" +if [[ "$model" == gpt2 ]]; then layout='Etttttt|ttttttFH'; else layout='Etttttttt|ttttttttFH'; fi +run_case custom 2 "$layout" "$out_dir/custom/run.log" + +MICRO_BATCHES="$micro_batches" python3 - "$out_dir" "$csv" <<'PY' +import csv, re, sys +from pathlib import Path +out, csv_path = Path(sys.argv[1]), Path(sys.argv[2]) +rows=[] +for name, pp in (("baseline",1),("custom",2)): + text=(out/name/"run.log").read_text(errors="replace") + m=re.findall(r"step\s+\d+/\d+\s+\|\s+train loss[^\n]*?\(([-+0-9.]+) ms \|\s*([-+0-9.]+) tok/s", text) + if not m: raise SystemExit(f"could not parse timing from {name} log") + ms,tps=map(float,m[-1]); microbatches=float(__import__('os').environ.get('MICRO_BATCHES','1')) + if microbatches <= 0: raise SystemExit('MICRO_BATCHES must be positive') + bubble=(pp-1)/(microbatches+pp-1) + rows.append(dict(layout=name,pipeline_parallel=pp,elapsed_ms=ms,tok_per_s=tps,bubble_ratio=bubble)) +with csv_path.open("w", newline="") as f: + w=csv.DictWriter(f, fieldnames=rows[0].keys()); w.writeheader(); w.writerows(rows) +print(f"wrote {csv_path}") +PY +cat "$csv" diff --git a/scripts/compare_pipeline_stage_perf.sh b/scripts/compare_pipeline_stage_perf.sh new file mode 100755 index 000000000..0ca2045a9 --- /dev/null +++ b/scripts/compare_pipeline_stage_perf.sh @@ -0,0 +1,140 @@ +#!/usr/bin/env bash +# Compare uniform and custom layouts at the same PP degree while collecting +# per-stage forward/backward timings from the real pipeline runtime. +set -euo pipefail + +if [[ $# -ne 2 ]]; then + echo "Usage: $0 " >&2 + exit 2 +fi + +model="$1" +device="$2" +[[ "$model" == gpt2 || "$model" == llama3 ]] || { echo "unsupported model: $model" >&2; exit 2; } +[[ "$device" == cpu || "$device" == cuda ]] || { echo "unsupported device: $device" >&2; exit 2; } + +repo_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +build_dir="${BUILD_DIR:-$repo_dir/build}" +exe="${EXE:-$build_dir/$model}" +out_dir="${OUT_DIR:-$repo_dir/artifacts/pipeline_excellent/$model}" +csv="${PERF_CSV:-$repo_dir/targets/pipeline_stage_perf.csv}" +mkdir -p "$out_dir" "$(dirname "$csv")" + +if [[ "$model" == gpt2 ]]; then + default_llmc="$repo_dir/data/gpt2/gpt2_124M.bin" + default_tokenizer="$repo_dir/data/gpt2/gpt2_tokenizer.bin" + custom_partition="${CUSTOM_PARTITION:-4,8}" +else + default_llmc="$repo_dir/data/llama3/llama3.2_1B_fp32.bin" + default_tokenizer="" + custom_partition="${CUSTOM_PARTITION:-8,8}" +fi + +batch_size="${BATCH_SIZE:-1}" +sequence_length="${SEQUENCE_LENGTH:-32}" +total_batch_size="${TOTAL_BATCH_SIZE:-$((batch_size * sequence_length * 4))}" +pp_size="${PIPELINE_PARALLEL:-2}" +micro_batches=$((total_batch_size / (batch_size * sequence_length))) +if (( pp_size != 2 )); then + echo "This acceptance script currently reports stage0/stage1 columns and requires PIPELINE_PARALLEL=2." >&2 + exit 2 +fi +if (( micro_batches <= 0 || total_batch_size % (batch_size * sequence_length) != 0 )); then + echo "TOTAL_BATCH_SIZE must be a positive multiple of BATCH_SIZE*SEQUENCE_LENGTH" >&2 + exit 2 +fi + +common=( + "--device=$device" + "--pipeline_parallel=$pp_size" + "--input_bin=${INPUT_BIN:-$repo_dir/data/$model/tiny_shakespeare_train.bin}" + "--llmc_filepath=${LLMC_FILE:-$default_llmc}" + "--overfit_single_batch=true" + "--num_iteration=1" + "--dtype=${DTYPE:-float32}" + "--batch_size=$batch_size" + "--sequence_length=$sequence_length" + "--total_batch_size=$total_batch_size" +) +[[ -e "${INPUT_VAL_BIN:-$repo_dir/data/$model/tiny_shakespeare_val.bin}" ]] && + common+=("--input_val_bin=${INPUT_VAL_BIN:-$repo_dir/data/$model/tiny_shakespeare_val.bin}") +[[ -n "$default_tokenizer" && -e "${TOKENIZER_BIN:-$default_tokenizer}" ]] && + common+=("--tokenizer_bin=${TOKENIZER_BIN:-$default_tokenizer}") + +launcher_spec="${LAUNCHER:-$build_dir/infini_run}" +run_case() { + local name="$1" partition="$2" + local -a cmd + if [[ "$launcher_spec" == direct ]]; then + cmd=("$exe") + else + read -r -a launcher_command <<< "$launcher_spec" + cmd=("${launcher_command[@]}" --nnodes=1 --nproc_per_node="$pp_size" "$exe") + fi + cmd+=("${common[@]}") + [[ -n "$partition" ]] && cmd+=("--pipeline_layer_partition=$partition") + echo "[$name] ${cmd[*]}" + INFINI_PIPELINE_STAGE_TIMING=1 "${cmd[@]}" 2>&1 | tee "$out_dir/${name}.log" +} + +for required in "$exe" "${common[2]}" "${common[3]}"; do + if [[ "$required" == --* ]]; then + path="${required#*=}" + [[ -e "$path" ]] || { echo "missing required path: $path" >&2; exit 3; } + else + [[ -e "$required" ]] || { echo "missing required path: $required" >&2; exit 3; } + fi +done + +run_case uniform "" +run_case custom "$custom_partition" + +MICRO_BATCHES="$micro_batches" python3 - "$out_dir" "$csv" <<'PY' +import csv +import os +import re +import sys +from pathlib import Path + +out_dir = Path(sys.argv[1]) +csv_path = Path(sys.argv[2]) +timing = re.compile( + r"pipeline_stage_timing direction=(forward|backward) " + r"stage=(\d+) chunk=(\d+) microbatch=(\d+) elapsed_ms=([0-9.]+)" +) +step = re.compile(r"step\s+\d+/\d+\s+\|.*?\(([-+0-9.]+) ms \|\s*([-+0-9.]+) tok/s") +rows = [] +for name in ("uniform", "custom"): + text = (out_dir / f"{name}.log").read_text(errors="replace") + stage_totals = {} + for direction, stage, chunk, microbatch, elapsed in timing.findall(text): + stage = int(stage) + stage_totals[stage] = stage_totals.get(stage, 0.0) + float(elapsed) + if not stage_totals: + raise SystemExit(f"no pipeline_stage_timing records found in {name}.log") + match = step.findall(text) + if not match: + raise SystemExit(f"no step timing found in {name}.log") + elapsed_ms, tok_per_s = map(float, match[-1]) + ordered = [stage_totals.get(i, 0.0) for i in range(2)] + max_stage = max(ordered) + min_stage = min(ordered) + rows.append({ + "layout": name, + "pipeline_parallel": 2, + "micro_batches": int(os.environ["MICRO_BATCHES"]), + "elapsed_ms": elapsed_ms, + "tok_per_s": tok_per_s, + "bubble_ratio": 1.0 / (int(os.environ["MICRO_BATCHES"]) + 1), + "stage0_compute_ms": ordered[0], + "stage1_compute_ms": ordered[1], + "max_stage_compute_ms": max_stage, + "stage_imbalance_ratio": max_stage / min_stage if min_stage > 0 else float("inf"), + }) +with csv_path.open("w", newline="") as handle: + writer = csv.DictWriter(handle, fieldnames=rows[0].keys()) + writer.writeheader() + writer.writerows(rows) +print(f"wrote {csv_path}") +PY +cat "$csv" diff --git a/scripts/test_pipeline_custom_layout.sh b/scripts/test_pipeline_custom_layout.sh new file mode 100755 index 000000000..8c9c220fa --- /dev/null +++ b/scripts/test_pipeline_custom_layout.sh @@ -0,0 +1,98 @@ +#!/usr/bin/env bash +# Compare a one-stage baseline with a two-stage custom pipeline layout. +# +# This script intentionally runs exactly one optimization iteration. It is an +# end-to-end smoke test, not a replacement for a training run. The caller must +# provide the model checkpoint and input binaries (or override their paths). + +set -euo pipefail + +if [[ $# -ne 2 ]]; then + echo "Usage: $0 " >&2 + exit 2 +fi + +model="$1" +device="$2" +case "$model" in + gpt2|llama3) ;; + *) echo "unsupported model '$model' (expected gpt2 or llama3)" >&2; exit 2 ;; +esac +case "$device" in + cpu|cuda) ;; + *) echo "unsupported device '$device' (expected cpu or cuda)" >&2; exit 2 ;; +esac + +repo_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +build_dir="${BUILD_DIR:-$repo_dir/build}" +exe="${EXE:-$build_dir/$model}" +compare_script="${COMPARE_LOSS:-$repo_dir/scripts/compare_loss.py}" +input_bin="${INPUT_BIN:-$repo_dir/data/$model/tiny_shakespeare_train.bin}" +input_val_bin="${INPUT_VAL_BIN:-$repo_dir/data/$model/tiny_shakespeare_val.bin}" +tokenizer_bin="${TOKENIZER_BIN:-$repo_dir/data/$model/${model}_tokenizer.bin}" +if [[ "$model" == "gpt2" ]]; then + default_llmc="$repo_dir/data/gpt2/gpt2_124M.bin" +else + default_llmc="$repo_dir/data/llama3/llama3.2_1B_fp32.bin" +fi +llmc_file="${LLMC_FILE:-$default_llmc}" +out_dir="${OUT_DIR:-$repo_dir/artifacts/pipeline_custom_layout/$model}" +num_iteration="${NUM_ITERATION:-1}" +launcher="${LAUNCHER:-$build_dir/infini_run}" + +if [[ "$num_iteration" != "1" ]]; then + echo "NUM_ITERATION must be 1 for this validation script (got '$num_iteration')" >&2 + exit 2 +fi + +for required in "$exe" "$compare_script" "$input_bin" "$llmc_file"; do + if [[ ! -e "$required" ]]; then + echo "missing required path: $required" >&2 + echo "Set BUILD_DIR/INPUT_BIN/LLMC_FILE (and related variables) to continue." >&2 + exit 3 + fi +done +if [[ "$device" == "cuda" && ! -e "$input_val_bin" ]]; then + echo "warning: validation input '$input_val_bin' is absent; continuing without validation data" >&2 +fi +if [[ ! -x "$exe" ]]; then + echo "model executable is not executable: $exe" >&2 + exit 3 +fi + +mkdir -p "$out_dir/baseline" "$out_dir/custom" + +common_args=( + "--device=$device" + "--input_bin=$input_bin" + "--llmc_filepath=$llmc_file" + "--overfit_single_batch=true" + "--num_iteration=1" + "--dtype=${DTYPE:-float32}" +) +if [[ -e "$input_val_bin" ]]; then common_args+=("--input_val_bin=$input_val_bin"); fi +if [[ -e "$tokenizer_bin" ]]; then common_args+=("--tokenizer_bin=$tokenizer_bin"); fi + +run_case() { + local name="$1" pp="$2" partition="$3" output="$4" + local -a command + if [[ "$launcher" == "direct" ]]; then + command=("$exe") + else + read -r -a launcher_command <<< "$launcher" + command=("${launcher_command[@]}" --nnodes=1 --nproc_per_node="$pp" "$exe") + fi + command+=("${common_args[@]}" "--pipeline_parallel=$pp") + if [[ -n "$partition" ]]; then command+=("--pipeline_layer_partition=$partition"); fi + echo "[$name] ${command[*]}" + "${command[@]}" 2>&1 | tee "$output/run.log" +} + +run_case baseline 1 "" "$out_dir/baseline" +run_case custom 2 "6,6" "$out_dir/custom" + +python3 "$compare_script" "$out_dir/baseline" "$out_dir/custom" \ + --threshold-fp32 "${THRESHOLD_FP32:-1e-5}" \ + --threshold-bf16 "${THRESHOLD_BF16:-1e-2}" + +echo "Pipeline custom-layout smoke test passed; logs are in $out_dir" diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 3bfaa548a..bdf345674 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -9,6 +9,8 @@ add_subdirectory(common) # Distributed tests add_subdirectory(distributed) +# Parallel tests +add_subdirectory(parallel) # Module tests add_subdirectory(module) diff --git a/tests/parallel/CMakeLists.txt b/tests/parallel/CMakeLists.txt new file mode 100644 index 000000000..19a8f910a --- /dev/null +++ b/tests/parallel/CMakeLists.txt @@ -0,0 +1,28 @@ +# ========================================================================== +# Pipeline layout tests +# ========================================================================== + +infini_train_add_test(test_pipeline_layout_cpu + SOURCES test_pipeline_layout.cc + LABELS cpu +) + +infini_train_add_test(test_pipeline_scheduler_layout_cpu + SOURCES test_pipeline_scheduler_layout.cc + LABELS cpu +) + +infini_train_add_test(test_pipeline_parallel_chunking_cpu + SOURCES test_pipeline_parallel_chunking.cc + LABELS cpu +) + +infini_train_add_test(test_pipeline_layout_megatron_style_cpu + SOURCES test_pipeline_layout_megatron_style.cc + LABELS cpu +) + +infini_train_add_test(test_pipeline_layout_balance_cpu + SOURCES test_pipeline_layout_balance.cc + LABELS cpu +) diff --git a/tests/parallel/test_pipeline_layout.cc b/tests/parallel/test_pipeline_layout.cc new file mode 100644 index 000000000..7be506afd --- /dev/null +++ b/tests/parallel/test_pipeline_layout.cc @@ -0,0 +1,188 @@ +#include + +#include "gtest/gtest.h" + +#include "infini_train/include/nn/parallel/pipeline_layout.h" + +namespace infini_train::nn::parallel { +namespace { + +PipelineLayout BuildExampleLayout() { return PipelineLayout::BuildDefault(24, 4, 1); } + +} // namespace + +TEST(PipelineLayoutTest, BuildDefaultBasicProperties) { + auto layout = BuildExampleLayout(); + + EXPECT_EQ(layout.num_layers(), 24); + EXPECT_EQ(layout.num_stages(), 4); + EXPECT_EQ(layout.vpp_size(), 1); +} + +TEST(PipelineLayoutTest, BuildDefaultChunkRanges) { + auto layout = BuildExampleLayout(); + + EXPECT_EQ(layout.chunk(0).layers.start, 0); + EXPECT_EQ(layout.chunk(0).layers.end, 6); + EXPECT_EQ(layout.chunk(1).layers.start, 6); + EXPECT_EQ(layout.chunk(1).layers.end, 12); + EXPECT_EQ(layout.chunk(2).layers.start, 12); + EXPECT_EQ(layout.chunk(2).layers.end, 18); + EXPECT_EQ(layout.chunk(3).layers.start, 18); + EXPECT_EQ(layout.chunk(3).layers.end, 24); +} + +TEST(PipelineLayoutTest, BuildDefaultVirtualPipelineUsesInterleavedChunkIds) { + auto layout = PipelineLayout::BuildDefault(8, 2, 2); + + ASSERT_EQ(layout.chunks().size(), 4); + EXPECT_EQ(layout.chunk(0).stage_id, 0); + EXPECT_EQ(layout.chunk(0).local_chunk_id, 0); + EXPECT_EQ(layout.chunk(1).stage_id, 1); + EXPECT_EQ(layout.chunk(1).local_chunk_id, 0); + EXPECT_EQ(layout.chunk(2).stage_id, 0); + EXPECT_EQ(layout.chunk(2).local_chunk_id, 1); + EXPECT_EQ(layout.chunk(3).stage_id, 1); + EXPECT_EQ(layout.chunk(3).local_chunk_id, 1); +} + +TEST(PipelineLayoutTest, BuildDefaultAllowsFewerLayersThanGlobalChunks) { + auto layout = PipelineLayout::BuildDefault(2, 4, 1); + + EXPECT_EQ(layout.num_layers(), 2); + EXPECT_EQ(layout.chunks().size(), 2); + EXPECT_EQ(layout.chunk(0).layers.start, 0); + EXPECT_EQ(layout.chunk(0).layers.end, 1); + EXPECT_EQ(layout.chunk(1).layers.start, 1); + EXPECT_EQ(layout.chunk(1).layers.end, 2); + EXPECT_NO_THROW(layout.Validate()); + EXPECT_THROW(layout.ValidateForCurrentPipelineTransport(), PipelineLayoutError); +} + +TEST(PipelineLayoutTest, StageOfLayer) { + auto layout = BuildExampleLayout(); + + EXPECT_EQ(layout.stage_of_layer(0), 0); + EXPECT_EQ(layout.stage_of_layer(5), 0); + EXPECT_EQ(layout.stage_of_layer(6), 1); + EXPECT_EQ(layout.stage_of_layer(11), 1); + EXPECT_EQ(layout.stage_of_layer(12), 2); + EXPECT_EQ(layout.stage_of_layer(17), 2); + EXPECT_EQ(layout.stage_of_layer(18), 3); + EXPECT_EQ(layout.stage_of_layer(23), 3); +} + +TEST(PipelineLayoutTest, ChunkOfLayer) { + auto layout = BuildExampleLayout(); + + EXPECT_EQ(layout.chunk_of_layer(0).global_chunk_id, 0); + EXPECT_EQ(layout.chunk_of_layer(5).global_chunk_id, 0); + EXPECT_EQ(layout.chunk_of_layer(6).global_chunk_id, 1); + EXPECT_EQ(layout.chunk_of_layer(11).global_chunk_id, 1); + EXPECT_EQ(layout.chunk_of_layer(23).global_chunk_id, 3); +} + +TEST(PipelineLayoutTest, LocalLayerIndexForOwnedStage) { + auto layout = BuildExampleLayout(); + + EXPECT_EQ(layout.local_layer_index(0, 0), 0); + EXPECT_EQ(layout.local_layer_index(0, 5), 5); + EXPECT_EQ(layout.local_layer_index(1, 6), 0); + EXPECT_EQ(layout.local_layer_index(1, 11), 5); + EXPECT_EQ(layout.local_layer_index(3, 18), 0); + EXPECT_EQ(layout.local_layer_index(3, 23), 5); +} + +TEST(PipelineLayoutTest, SpecialModulePlacement) { + auto layout = BuildExampleLayout(); + + EXPECT_TRUE(layout.owns(SpecialModule::kEmbedding, 0)); + EXPECT_FALSE(layout.owns(SpecialModule::kEmbedding, 3)); + + EXPECT_TRUE(layout.owns(SpecialModule::kFinalNorm, 3)); + EXPECT_FALSE(layout.owns(SpecialModule::kFinalNorm, 0)); + + EXPECT_TRUE(layout.owns(SpecialModule::kLMHead, 3)); + EXPECT_FALSE(layout.owns(SpecialModule::kLMHead, 0)); +} + +TEST(PipelineLayoutTest, ExplicitBoundarySpecialModulePlacement) { + SpecialModulePlacement placement{ + .embedding_stage = 0, + .final_norm_stage = 3, + .lm_head_stage = 3, + }; + auto layout = PipelineLayout::BuildContiguous({4, 8, 6, 6}, placement); + + EXPECT_EQ(layout.special_modules().embedding_stage, 0); + EXPECT_EQ(layout.special_modules().final_norm_stage, 3); + EXPECT_EQ(layout.special_modules().lm_head_stage, 3); + EXPECT_NO_THROW(layout.Validate()); +} + +TEST(PipelineLayoutTest, RejectsNonBoundarySpecialModulesByDefault) { + SpecialModulePlacement placement{ + .embedding_stage = 1, + .final_norm_stage = 3, + .lm_head_stage = 3, + }; + EXPECT_THROW(PipelineLayout::BuildContiguous({4, 8, 6, 6}, placement), PipelineLayoutError); +} + +TEST(PipelineLayoutTest, ToStringContainsBasicInfo) { + auto layout = BuildExampleLayout(); + const std::string text = layout.ToString(); + + EXPECT_NE(text.find("num_layers: 24"), std::string::npos); + EXPECT_NE(text.find("num_stages: 4"), std::string::npos); + EXPECT_NE(text.find("stage 3:"), std::string::npos); +} + +TEST(PipelineLayoutTest, ParseLayerPartitionValid) { + EXPECT_EQ(PipelineLayout::ParseLayerPartition("4,8,6,6"), (std::vector{4, 8, 6, 6})); + EXPECT_EQ(PipelineLayout::ParseLayerPartition("24"), (std::vector{24})); + EXPECT_EQ(PipelineLayout::ParseLayerPartition("0,4"), (std::vector{0, 4})); +} + +TEST(PipelineLayoutTest, ParseLayerPartitionRejectsInvalidStrings) { + EXPECT_THROW(PipelineLayout::ParseLayerPartition(""), PipelineLayoutError); + EXPECT_THROW(PipelineLayout::ParseLayerPartition("4,,6"), PipelineLayoutError); + EXPECT_THROW(PipelineLayout::ParseLayerPartition(",4,6"), PipelineLayoutError); + EXPECT_THROW(PipelineLayout::ParseLayerPartition("4,6,"), PipelineLayoutError); + EXPECT_THROW(PipelineLayout::ParseLayerPartition("4, 6"), PipelineLayoutError); + EXPECT_THROW(PipelineLayout::ParseLayerPartition("4,a,6"), PipelineLayoutError); + EXPECT_THROW(PipelineLayout::ParseLayerPartition("-1,6"), PipelineLayoutError); + EXPECT_THROW(PipelineLayout::ParseLayerPartition("999999999999999999999"), PipelineLayoutError); +} + +TEST(PipelineLayoutTest, BuildPipelineLayoutWithCustomPartition) { + auto layout = PipelineLayout::BuildPipelineLayout(24, 4, 1, "4,8,6,6"); + + EXPECT_EQ(layout.num_layers(), 24); + EXPECT_EQ(layout.num_stages(), 4); + EXPECT_EQ(layout.vpp_size(), 1); + EXPECT_EQ(layout.chunk(0).layers.start, 0); + EXPECT_EQ(layout.chunk(0).layers.end, 4); + EXPECT_EQ(layout.chunk(1).layers.start, 4); + EXPECT_EQ(layout.chunk(1).layers.end, 12); + EXPECT_EQ(layout.chunk(3).layers.start, 18); + EXPECT_EQ(layout.chunk(3).layers.end, 24); +} + +TEST(PipelineLayoutTest, BuildPipelineLayoutEmptyPartitionUsesDefault) { + auto layout = PipelineLayout::BuildPipelineLayout(24, 4, 1, ""); + + EXPECT_EQ(layout.num_layers(), 24); + EXPECT_EQ(layout.chunk(0).layers.start, 0); + EXPECT_EQ(layout.chunk(0).layers.end, 6); + EXPECT_EQ(layout.chunk(3).layers.start, 18); + EXPECT_EQ(layout.chunk(3).layers.end, 24); +} + +TEST(PipelineLayoutTest, BuildPipelineLayoutRejectsInvalidConfigurations) { + EXPECT_THROW(PipelineLayout::BuildPipelineLayout(23, 4, 1, "4,8,6,6"), PipelineLayoutError); + EXPECT_THROW(PipelineLayout::BuildPipelineLayout(24, 3, 1, "4,8,6,6"), PipelineLayoutError); + EXPECT_THROW(PipelineLayout::BuildPipelineLayout(24, 4, 2, "4,8,6,6"), PipelineLayoutError); +} + +} // namespace infini_train::nn::parallel diff --git a/tests/parallel/test_pipeline_layout_balance.cc b/tests/parallel/test_pipeline_layout_balance.cc new file mode 100644 index 000000000..70303b5ef --- /dev/null +++ b/tests/parallel/test_pipeline_layout_balance.cc @@ -0,0 +1,23 @@ +#include + +#include "infini_train/include/nn/parallel/pipeline_layout.h" + +using namespace infini_train::nn::parallel; + +TEST(PipelineLayoutBalanceTest, BalancesUniformCosts) { + EXPECT_EQ(PipelineLayout::SuggestBalancedPartition(8, 4), (std::vector{2, 2, 2, 2})); +} + +TEST(PipelineLayoutBalanceTest, IsolatesHeavyLayer) { + const auto partition = PipelineLayout::SuggestBalancedPartition(8, 4, {1, 1, 1, 1, 4, 1, 1, 1}); + ASSERT_EQ(partition.size(), 4u); + EXPECT_EQ(partition[0], 3); + EXPECT_EQ(partition[1], 1); + EXPECT_EQ(partition[2], 1); + EXPECT_EQ(partition[3], 3); +} + +TEST(PipelineLayoutBalanceTest, RejectsInvalidCosts) { + EXPECT_THROW(PipelineLayout::SuggestBalancedPartition(4, 2, {1, -1, 1, 1}), PipelineLayoutError); + EXPECT_THROW(PipelineLayout::SuggestBalancedPartition(4, 2, {1, 1}), PipelineLayoutError); +} diff --git a/tests/parallel/test_pipeline_layout_megatron_style.cc b/tests/parallel/test_pipeline_layout_megatron_style.cc new file mode 100644 index 000000000..1982e6237 --- /dev/null +++ b/tests/parallel/test_pipeline_layout_megatron_style.cc @@ -0,0 +1,94 @@ +#include + +#include "infini_train/include/nn/parallel/pipeline_layout.h" + +using namespace infini_train::nn::parallel; + +TEST(PipelineLayoutMegatronStyleTest, ParsesRepeatedChunksAndSpecialModules) { + auto layout = PipelineLayout::ParseMegatronStyleLayout("Ett|tt|tt|FH", 6, 4, 1); + EXPECT_EQ(layout.num_layers(), 6); + EXPECT_EQ(layout.stage(0).global_chunk_ids.size(), 1u); + EXPECT_EQ(layout.stage(1).global_chunk_ids.size(), 1u); + EXPECT_TRUE(layout.owns(SpecialModule::kEmbedding, 0)); + EXPECT_TRUE(layout.owns(SpecialModule::kFinalNorm, 3)); + EXPECT_TRUE(layout.owns(SpecialModule::kLMHead, 3)); +} + +TEST(PipelineLayoutMegatronStyleTest, ParsesVirtualPipelineChunks) { + auto layout = PipelineLayout::ParseMegatronStyleLayout("tt,tt|tt,tt", 8, 2, 2); + EXPECT_EQ(layout.vpp_size(), 2); + EXPECT_EQ(layout.stage(0).global_chunk_ids.size(), 2u); + EXPECT_EQ(layout.chunk_of_layer(0).local_chunk_id, 0); + EXPECT_EQ(layout.chunk_of_layer(4).local_chunk_id, 1); +} + +TEST(PipelineLayoutMegatronStyleTest, PreservesAsymmetricExplicitChunkOwnership) { + // The two virtual chunks on each stage have different sizes. Ownership + // comes from the explicit stage/chunk description, not a fixed + // equal-size or layer-count assumption. + auto layout = PipelineLayout::ParseMegatronStyleLayout("t,tt|tt,t", 6, 2, 2); + + EXPECT_EQ(layout.chunk(0).stage_id, 0); + EXPECT_EQ(layout.chunk(0).local_chunk_id, 0); + EXPECT_EQ(layout.chunk(0).layers.size(), 1); + EXPECT_EQ(layout.chunk(1).stage_id, 1); + EXPECT_EQ(layout.chunk(1).local_chunk_id, 0); + EXPECT_EQ(layout.chunk(1).layers.size(), 2); + EXPECT_EQ(layout.chunk(2).stage_id, 0); + EXPECT_EQ(layout.chunk(2).local_chunk_id, 1); + EXPECT_EQ(layout.chunk(2).layers.size(), 2); + EXPECT_EQ(layout.chunk(3).stage_id, 1); + EXPECT_EQ(layout.chunk(3).local_chunk_id, 1); + EXPECT_EQ(layout.chunk(3).layers.size(), 1); +} + +TEST(PipelineLayoutMegatronStyleTest, BuildsExplicitChunkStageMap) { + std::vector chunks{ + {.global_chunk_id = 0, .stage_id = 1, .local_chunk_id = 0, .layers = {0, 2}}, + {.global_chunk_id = 1, .stage_id = 0, .local_chunk_id = 0, .layers = {2, 3}}, + {.global_chunk_id = 2, .stage_id = 1, .local_chunk_id = 1, .layers = {3, 5}}, + {.global_chunk_id = 3, .stage_id = 0, .local_chunk_id = 1, .layers = {5, 6}}, + }; + auto layout = PipelineLayout::BuildExplicit(6, 2, 2, chunks); + + EXPECT_EQ(layout.chunk(0).stage_id, 1); + EXPECT_EQ(layout.chunk(1).stage_id, 0); + EXPECT_EQ(layout.chunk(2).stage_id, 1); + EXPECT_EQ(layout.chunk(3).stage_id, 0); + EXPECT_EQ(layout.stage(0).global_chunk_ids, (std::vector{1, 3})); + EXPECT_EQ(layout.stage(1).global_chunk_ids, (std::vector{0, 2})); + EXPECT_EQ(layout.stage_of_layer(0), 1); + EXPECT_EQ(layout.stage_of_layer(2), 0); +} + +TEST(PipelineLayoutMegatronStyleTest, AllowsEmptyStageInExplicitLayout) { + auto layout = PipelineLayout::ParseMegatronStyleLayout("Etttttt||ttttttFH", 12, 3, 1); + + EXPECT_EQ(layout.stage(0).global_chunk_ids.size(), 1u); + EXPECT_EQ(layout.stage(1).global_chunk_ids.size(), 1u); + EXPECT_EQ(layout.stage(2).global_chunk_ids.size(), 1u); + EXPECT_EQ(layout.chunk(1).layers.size(), 0); + EXPECT_NO_THROW(layout.ValidateForCurrentPipelineTransport()); +} + +TEST(PipelineLayoutMegatronStyleTest, RejectsIncompleteExplicitChunkStageMap) { + std::vector chunks{ + {.global_chunk_id = 0, .stage_id = 0, .local_chunk_id = 0, .layers = {0, 1}}, + {.global_chunk_id = 1, .stage_id = 1, .local_chunk_id = 0, .layers = {1, 2}}, + {.global_chunk_id = 2, .stage_id = 0, .local_chunk_id = 1, .layers = {2, 3}}, + // stage 1/local chunk 1 is intentionally missing + }; + EXPECT_THROW(PipelineLayout::BuildExplicit(4, 2, 2, chunks), PipelineLayoutError); +} + +TEST(PipelineLayoutMegatronStyleTest, ExpandsParenthesizedRepetition) { + auto layout = PipelineLayout::ParseMegatronStyleLayout("(tt)*2|tt|tt", 8, 3, 1); + EXPECT_EQ(layout.chunk_of_layer(0).layers.size(), 4); + EXPECT_EQ(layout.stage_of_layer(4), 1); +} + +TEST(PipelineLayoutMegatronStyleTest, RejectsMalformedExpressions) { + EXPECT_THROW(PipelineLayout::ParseMegatronStyleLayout("(tt|tt", 4, 2), PipelineLayoutError); + EXPECT_THROW(PipelineLayout::ParseMegatronStyleLayout("t,x|tt", 3, 2), PipelineLayoutError); + EXPECT_THROW(PipelineLayout::ParseMegatronStyleLayout("tt|tt", 3, 2), PipelineLayoutError); +} diff --git a/tests/parallel/test_pipeline_parallel_chunking.cc b/tests/parallel/test_pipeline_parallel_chunking.cc new file mode 100644 index 000000000..84013c19a --- /dev/null +++ b/tests/parallel/test_pipeline_parallel_chunking.cc @@ -0,0 +1,138 @@ +#include +#include + +#include "gtest/gtest.h" + +#include "infini_train/include/device.h" +#include "infini_train/include/nn/modules/container.h" +#include "infini_train/include/nn/modules/transformer/transformer.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/pipeline_layout.h" +#include "infini_train/include/nn/parallel/pp/pipeline_parallel.h" + +using namespace infini_train; + +namespace { + +nn::TransformerConfig MakeConfig() { + nn::TransformerConfig config; + config.n_layer = 12; + config.n_head = 4; + config.n_kv_head = 4; + config.n_embd = 32; + config.vocab_size = 64; + config.original_vocab_size = 64; + config.position_embedding_type = nn::PositionEmbeddingType::kLearnedAbsolute; + config.activation_type = nn::MLPType::kGELU; + config.norm_type = nn::NormType::kLayerNorm; + config.add_bias_linear = true; + config.tie_weights = false; + return config; +} + +std::shared_ptr BuildPipeline(const nn::parallel::PipelineLayout &layout, + int stage_id) { + nn::parallel::global::GlobalEnv::Instance().set_pipeline_layout(layout); + nn::parallel::pp_rank = stage_id; + + auto model = std::make_shared(MakeConfig()); + const std::vector> recv_shape = {{2, 4, 32}}; + const int local_chunk_count = static_cast(layout.stage(stage_id).global_chunk_ids.size()); + + return std::make_shared(model, layout.num_stages(), /*num_micro_batches=*/1, + recv_shape, stage_id, Device(Device::DeviceType::kCPU, 0), + local_chunk_count); +} + +std::shared_ptr AsSequential(const std::shared_ptr &pipeline, + size_t index) { + auto *chunks = pipeline->mutable_chunks(); + if (chunks == nullptr || index >= chunks->size()) { + return nullptr; + } + return std::dynamic_pointer_cast(chunks->at(index)); +} + +} // namespace + +TEST(PipelineParallelChunkingTest, Stage0IncludesEmbeddingOnly) { + const auto layout = nn::parallel::PipelineLayout::BuildContiguous({2, 4, 3, 3}); + auto pipeline = BuildPipeline(layout, 0); + + ASSERT_EQ(pipeline->mutable_chunks()->size(), 1U); + auto chunk = AsSequential(pipeline, 0); + ASSERT_NE(chunk, nullptr); + EXPECT_EQ(chunk->module("0").type(), nn::TransformerFirstStage::kType); + EXPECT_EQ(chunk->module("1").type(), nn::TransformerChunk::kType); +} + +TEST(PipelineParallelChunkingTest, IntermediateStageIncludesOnlyTransformerChunk) { + const auto layout = nn::parallel::PipelineLayout::BuildContiguous({2, 4, 3, 3}); + auto pipeline = BuildPipeline(layout, 1); + + ASSERT_EQ(pipeline->mutable_chunks()->size(), 1U); + auto chunk = AsSequential(pipeline, 0); + ASSERT_NE(chunk, nullptr); + EXPECT_EQ(chunk->module("0").type(), nn::TransformerChunk::kType); +} + +TEST(PipelineParallelChunkingTest, LastStageIncludesTransformerLastStage) { + const auto layout = nn::parallel::PipelineLayout::BuildContiguous({2, 4, 3, 3}); + auto pipeline = BuildPipeline(layout, 3); + + ASSERT_EQ(pipeline->mutable_chunks()->size(), 1U); + auto chunk = AsSequential(pipeline, 0); + ASSERT_NE(chunk, nullptr); + EXPECT_EQ(chunk->module("0").type(), nn::TransformerChunk::kType); + EXPECT_EQ(chunk->module("1").type(), nn::TransformerLastStage::kType); +} + +TEST(PipelineParallelChunkingTest, VirtualPipelinePreservesLocalChunkOrder) { + const auto layout = nn::parallel::PipelineLayout::BuildDefault(12, 2, 2); + auto pipeline = BuildPipeline(layout, 0); + + ASSERT_EQ(pipeline->mutable_chunks()->size(), 2U); + auto first = AsSequential(pipeline, 0); + auto second = AsSequential(pipeline, 1); + ASSERT_NE(first, nullptr); + ASSERT_NE(second, nullptr); + EXPECT_EQ(first->module("0").type(), nn::TransformerFirstStage::kType); + EXPECT_EQ(first->module("1").type(), nn::TransformerChunk::kType); + EXPECT_EQ(second->module("0").type(), nn::TransformerChunk::kType); +} + +TEST(PipelineParallelChunkingTest, EmptyStageBuildsPassThroughChunk) { + const auto layout + = nn::parallel::PipelineLayout::ParseMegatronStyleLayout("Etttttt||ttttttFH", 12, 3, 1); + auto pipeline = BuildPipeline(layout, 1); + + ASSERT_EQ(pipeline->mutable_chunks()->size(), 1U); + auto chunk = AsSequential(pipeline, 0); + ASSERT_NE(chunk, nullptr); + EXPECT_EQ(chunk->module("0").type(), nn::TransformerChunk::kType); +} + +TEST(PipelineParallelChunkingTest, RejectsStageCountMismatchWithCatchableError) { + const auto layout = nn::parallel::PipelineLayout::BuildContiguous({2, 4, 3, 3}); + nn::parallel::global::GlobalEnv::Instance().set_pipeline_layout(layout); + nn::parallel::pp_rank = 0; + auto model = std::make_shared(MakeConfig()); + + EXPECT_THROW(std::make_shared( + model, /*num_stages=*/2, /*num_micro_batches=*/1, std::vector>{{2, 4, 32}}, + /*rank=*/0, Device(Device::DeviceType::kCPU, 0), /*vpp=*/1), + nn::parallel::PipelineLayoutError); +} + +class PipelineParallelChunkingEnvironment : public ::testing::Environment { +public: + void TearDown() override { + nn::parallel::pp_rank = 0; + nn::parallel::global::GlobalEnv::Instance().set_pipeline_layout( + nn::parallel::PipelineLayout::BuildDefault(1, 1, 1)); + } +}; + +::testing::Environment *const pipeline_parallel_chunking_environment + = ::testing::AddGlobalTestEnvironment(new PipelineParallelChunkingEnvironment()); diff --git a/tests/parallel/test_pipeline_scheduler_layout.cc b/tests/parallel/test_pipeline_scheduler_layout.cc new file mode 100644 index 000000000..64fab9f4b --- /dev/null +++ b/tests/parallel/test_pipeline_scheduler_layout.cc @@ -0,0 +1,75 @@ +#include +#include + +#include "gtest/gtest.h" + +#include "infini_train/include/nn/parallel/pipeline_layout.h" +#include "infini_train/include/nn/parallel/pp/pipeline_schedule.h" + +namespace infini_train::nn::parallel { + +TEST(PipelineSchedulerLayoutTest, GPipeTasksUseLayoutOwnership) { + const auto layout = PipelineLayout::BuildDefault(/*num_layers=*/12, /*num_stages=*/3, /*vpp_size=*/2); + const auto schedule = PipelineParallelScheduler::GenerateGPipeSchedule( + /*n=*/2, /*num_stages=*/3, /*vpp_size=*/2, layout); + + ASSERT_EQ(schedule.size(), 2U * 2U * 3U * 2U); + for (const auto &task : schedule) { + const auto &chunk = layout.chunk(task.global_chunk_id); + EXPECT_EQ(task.stage_id, chunk.stage_id); + EXPECT_EQ(task.local_chunk_idx, chunk.local_chunk_id); + } +} + +TEST(PipelineSchedulerLayoutTest, InterleavedTasksUseLayoutOwnership) { + const auto layout = PipelineLayout::BuildDefault(/*num_layers=*/12, /*num_stages=*/3, /*vpp_size=*/2); + const auto schedule = PipelineParallelScheduler::GenerateInterleaved1F1BSchedule( + /*n=*/2, /*num_stages=*/3, /*vpp_size=*/2, layout); + + // Every microbatch traverses every global chunk once in each direction. + ASSERT_EQ(schedule.size(), 2U * 3U * 2U * 2U); + for (const auto &task : schedule) { + const auto &chunk = layout.chunk(task.global_chunk_id); + EXPECT_EQ(task.stage_id, chunk.stage_id); + EXPECT_EQ(task.local_chunk_idx, chunk.local_chunk_id); + } +} + +TEST(PipelineSchedulerLayoutTest, CreateTaskSupportsContiguousCustomLayout) { + const auto layout = PipelineLayout::BuildContiguous({2, 4, 3, 3}); + + const auto task = PipelineParallelScheduler::CreateTask( + /*step=*/7, /*mb=*/1, /*global_chunk=*/2, /*num_stages=*/4, + /*total_chunks=*/4, /*is_forward=*/true, layout); + + EXPECT_EQ(task.global_chunk_id, 2); + EXPECT_EQ(task.stage_id, 2); + EXPECT_EQ(task.local_chunk_idx, 0); + EXPECT_TRUE(task.is_forward); +} + +TEST(PipelineSchedulerLayoutTest, ExplicitChunkStageMapDrivesSchedulerOwnership) { + const std::vector chunks{ + {.global_chunk_id = 0, .stage_id = 1, .local_chunk_id = 0, .layers = {0, 2}}, + {.global_chunk_id = 1, .stage_id = 0, .local_chunk_id = 0, .layers = {2, 3}}, + {.global_chunk_id = 2, .stage_id = 1, .local_chunk_id = 1, .layers = {3, 5}}, + {.global_chunk_id = 3, .stage_id = 0, .local_chunk_id = 1, .layers = {5, 6}}, + }; + const auto layout = PipelineLayout::BuildExplicit(6, 2, 2, chunks); + const auto schedule = PipelineParallelScheduler::GenerateGPipeSchedule(1, 2, 2, layout); + + ASSERT_EQ(schedule.size(), 8U); + for (const auto &task : schedule) { + EXPECT_EQ(task.stage_id, layout.chunk(task.global_chunk_id).stage_id); + EXPECT_EQ(task.local_chunk_idx, layout.chunk(task.global_chunk_id).local_chunk_id); + } +} + +TEST(PipelineSchedulerLayoutTest, RejectsTopologyMismatch) { + const auto layout = PipelineLayout::BuildDefault(/*num_layers=*/12, /*num_stages=*/3, /*vpp_size=*/2); + + EXPECT_THROW(PipelineParallelScheduler::GenerateGPipeSchedule(/*n=*/2, /*num_stages=*/2, /*vpp_size=*/2, layout), + PipelineLayoutError); +} + +} // namespace infini_train::nn::parallel diff --git a/tests/transformer/test_transformer_architecture.cc b/tests/transformer/test_transformer_architecture.cc index d4a6efc29..170aa8a3e 100644 --- a/tests/transformer/test_transformer_architecture.cc +++ b/tests/transformer/test_transformer_architecture.cc @@ -1,3 +1,4 @@ +#include #include #include #include @@ -14,6 +15,9 @@ #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/nn/parallel/global.h" +#include "infini_train/include/nn/parallel/pipeline_layout.h" +#include "infini_train/include/nn/parallel/pp/pipeline_parallel.h" #include "infini_train/include/tensor.h" #include "tests/common/test_utils.h" @@ -143,6 +147,113 @@ TEST_P(TransformerModuleTest, GPT2Model) { EXPECT_FALSE(model->Parameters().empty()); } +TEST(TransformerPipelineLayoutTest, BuildsOnlyTheCurrentStageModules) { + using infini_train::nn::parallel::PipelineLayout; + using infini_train::nn::parallel::PipelineLayoutPolicy; + using infini_train::nn::parallel::global::GlobalEnv; + + auto layout = PipelineLayout::BuildContiguous({2, 4, 3, 3}); + GlobalEnv::Instance().set_pipeline_layout(layout); + + nn::TransformerConfig config; + config.n_layer = 12; + config.n_head = 4; + config.n_kv_head = 4; + config.n_embd = 32; + config.vocab_size = 64; + config.original_vocab_size = 64; + config.position_embedding_type = nn::PositionEmbeddingType::kLearnedAbsolute; + config.activation_type = nn::MLPType::kGELU; + config.norm_type = nn::NormType::kLayerNorm; + config.add_bias_linear = true; + config.tie_weights = false; + + auto has_key + = [](const auto &state_dict, const std::string &key) { return state_dict.find(key) != state_dict.end(); }; + auto has_named_parameter = [](const auto &named_parameters, const std::string &key) { + return std::any_of(named_parameters.begin(), named_parameters.end(), + [&](const auto &entry) { return entry.first == key; }); + }; + + nn::parallel::pp_rank = 0; + auto stage0 = std::make_shared(config); + auto state0 = stage0->StateDict(); + auto named0 = stage0->NamedParameters(); + EXPECT_TRUE(has_key(state0, "transformer.wte.weight")); + EXPECT_TRUE(has_key(state0, "transformer.wpe.weight")); + EXPECT_TRUE(has_key(state0, "transformer.h.0.ln_1.weight")); + EXPECT_TRUE(has_key(state0, "transformer.h.1.ln_1.weight")); + EXPECT_FALSE(has_key(state0, "transformer.h.2.ln_1.weight")); + EXPECT_FALSE(has_key(state0, "transformer.ln_f.weight")); + EXPECT_FALSE(has_key(state0, "transformer.lm_head.weight")); + EXPECT_TRUE(has_named_parameter(named0, "transformer.wte.weight")); + EXPECT_FALSE(has_named_parameter(named0, "transformer.lm_head.weight")); + + nn::parallel::pp_rank = 1; + auto stage1 = std::make_shared(config); + auto state1 = stage1->StateDict(); + auto named1 = stage1->NamedParameters(); + EXPECT_FALSE(has_key(state1, "transformer.wte.weight")); + EXPECT_TRUE(has_key(state1, "transformer.h.0.ln_1.weight")); + EXPECT_TRUE(has_key(state1, "transformer.h.3.ln_1.weight")); + EXPECT_FALSE(has_key(state1, "transformer.h.4.ln_1.weight")); + EXPECT_FALSE(has_key(state1, "transformer.ln_f.weight")); + EXPECT_FALSE(has_key(state1, "transformer.lm_head.weight")); + EXPECT_TRUE(has_named_parameter(named1, "transformer.h.0.ln_1.weight")); + EXPECT_FALSE(has_named_parameter(named1, "transformer.wte.weight")); + + nn::parallel::pp_rank = 3; + auto stage3 = std::make_shared(config); + auto state3 = stage3->StateDict(); + auto named3 = stage3->NamedParameters(); + EXPECT_FALSE(has_key(state3, "transformer.wte.weight")); + EXPECT_TRUE(has_key(state3, "transformer.h.0.ln_1.weight")); + EXPECT_TRUE(has_key(state3, "transformer.h.2.ln_1.weight")); + EXPECT_FALSE(has_key(state3, "transformer.h.3.ln_1.weight")); + EXPECT_TRUE(has_key(state3, "transformer.ln_f.weight")); + EXPECT_TRUE(has_key(state3, "transformer.lm_head.weight")); + EXPECT_TRUE(has_named_parameter(named3, "transformer.ln_f.weight")); + EXPECT_TRUE(has_named_parameter(named3, "transformer.lm_head.weight")); + EXPECT_FALSE(has_named_parameter(named3, "lm_head.weight")); + + nn::parallel::pp_rank = 2; + auto stage2 = std::make_shared(config); + auto state2 = stage2->StateDict(); + auto named2 = stage2->NamedParameters(); + EXPECT_FALSE(has_key(state2, "transformer.wte.weight")); + EXPECT_TRUE(has_key(state2, "transformer.h.0.ln_1.weight")); + EXPECT_TRUE(has_key(state2, "transformer.h.2.ln_1.weight")); + EXPECT_FALSE(has_key(state2, "transformer.h.3.ln_1.weight")); + EXPECT_FALSE(has_key(state2, "transformer.ln_f.weight")); + EXPECT_FALSE(has_key(state2, "transformer.lm_head.weight")); + EXPECT_TRUE(has_named_parameter(named2, "transformer.h.0.ln_1.weight")); + EXPECT_FALSE(has_named_parameter(named2, "transformer.ln_f.weight")); + + // The layout API also permits final norm and LM head to live on different + // stages when the transport policy explicitly allows it. The model must + // construct each optional part independently in that case. + PipelineLayoutPolicy split_policy; + split_policy.require_boundary_special_modules = false; + auto split_layout = PipelineLayout::BuildContiguous( + {2, 4, 3, 3}, {.embedding_stage = 0, .final_norm_stage = 2, .lm_head_stage = 3}, split_policy); + GlobalEnv::Instance().set_pipeline_layout(split_layout); + + nn::parallel::pp_rank = 2; + auto norm_stage = std::make_shared(config); + auto norm_state = norm_stage->StateDict(); + EXPECT_TRUE(has_key(norm_state, "transformer.ln_f.weight")); + EXPECT_FALSE(has_key(norm_state, "transformer.lm_head.weight")); + + nn::parallel::pp_rank = 3; + auto head_stage = std::make_shared(config); + auto head_state = head_stage->StateDict(); + EXPECT_FALSE(has_key(head_state, "transformer.ln_f.weight")); + EXPECT_TRUE(has_key(head_state, "transformer.lm_head.weight")); + + nn::parallel::pp_rank = 0; + GlobalEnv::Instance().set_pipeline_layout(PipelineLayout::BuildDefault(1, 1, 1)); +} + TEST_P(TransformerModuleTest, LLaMA3Model) { SKIP_CPU(); nn::TransformerConfig config; diff --git a/tools/pipeline_layout_suggest.cc b/tools/pipeline_layout_suggest.cc new file mode 100644 index 000000000..38d1114b7 --- /dev/null +++ b/tools/pipeline_layout_suggest.cc @@ -0,0 +1,151 @@ +#include +#include +#include +#include +#include +#include +#include + +#include "infini_train/include/nn/parallel/pipeline_layout.h" + +using infini_train::nn::parallel::PipelineLayout; + +namespace { +int IntArg(const std::string &arg, const char *name) { + const std::string prefix = std::string(name) + "="; + if (arg.rfind(prefix, 0) != 0) { + return -1; + } + return std::stoi(arg.substr(prefix.size())); +} +std::vector CostsArg(const std::string &arg) { + const std::string prefix = "--layer_cost="; + if (arg.rfind(prefix, 0) != 0) { + return {}; + } + std::vector out; + std::stringstream ss(arg.substr(prefix.size())); + std::string token; + while (std::getline(ss, token, ',')) { out.push_back(std::stod(token)); } + return out; +} + +std::vector ReadCostsFile(const std::string &path) { + std::ifstream input(path); + if (!input) { + throw std::runtime_error("cannot open layer cost file: " + path); + } + std::vector costs; + std::string line; + while (std::getline(input, line)) { + if (line.empty() || line[0] == '#') { + continue; + } + // Accept either one cost per line or CSV rows of `layer_id,cost`. + std::stringstream ss(line); + std::string first, second; + if (!std::getline(ss, first, ',')) { + continue; + } + if (std::getline(ss, second, ',')) { + try { + costs.push_back(std::stod(second)); + } catch (const std::exception &) { + // Permit a conventional CSV header such as `layer,cost`. + if (costs.empty()) { + continue; + } + throw; + } + } else { + try { + costs.push_back(std::stod(first)); + } catch (const std::exception &) { + if (costs.empty()) { + continue; + } + throw; + } + } + } + return costs; +} + +double MaxPartitionCost(const std::vector &costs, const std::vector &partition) { + double max_cost = 0.0; + size_t cursor = 0; + for (int count : partition) { + double stage_cost = 0.0; + for (int i = 0; i < count; ++i) { + stage_cost += costs[cursor++]; + } + max_cost = std::max(max_cost, stage_cost); + } + return max_cost; +} +} // namespace + +int main(int argc, char **argv) { + int num_layers = -1, pp_size = -1; + std::vector costs; + std::string cost_file; + for (int i = 1; i < argc; ++i) { + const std::string arg(argv[i]); + if (arg == "--help") { + std::cout << "Usage: pipeline_layout_suggest --num_layers=N --pp_size=N " + "[--layer_cost=c1,c2,...] [--layer_cost_file=path]\n"; + return 0; + } + const int layers = IntArg(arg, "--num_layers"); + const int stages = IntArg(arg, "--pp_size"); + if (layers >= 0) { + num_layers = layers; + } + if (stages >= 0) { + pp_size = stages; + } + auto parsed = CostsArg(arg); + if (!parsed.empty()) { + costs = std::move(parsed); + } + const std::string cost_prefix = "--layer_cost_file="; + if (arg.rfind(cost_prefix, 0) == 0) { + cost_file = arg.substr(cost_prefix.size()); + } + } + try { + if (!cost_file.empty()) { + costs = ReadCostsFile(cost_file); + } + const auto partition = PipelineLayout::SuggestBalancedPartition(num_layers, pp_size, costs); + std::cout << "pipeline_layer_partition="; + for (size_t i = 0; i < partition.size(); ++i) { + if (i) { + std::cout << ','; + } + std::cout << partition[i]; + } + std::cout << '\n'; + if (!costs.empty()) { + const auto uniform = PipelineLayout::SuggestBalancedPartition(num_layers, pp_size); + const double uniform_max = MaxPartitionCost(costs, uniform); + const double suggested_max = MaxPartitionCost(costs, partition); + const double reduction = uniform_max > 0.0 ? (uniform_max - suggested_max) / uniform_max * 100.0 : 0.0; + std::cout << "uniform_partition="; + for (size_t i = 0; i < uniform.size(); ++i) { + if (i) { + std::cout << ','; + } + std::cout << uniform[i]; + } + std::cout << '\n'; + std::cout << "uniform_max_stage_cost=" << uniform_max << '\n'; + std::cout << "suggested_max_stage_cost=" << suggested_max << '\n'; + std::cout << "predicted_max_stage_cost_reduction_pct=" << reduction << '\n'; + } + return 0; + } catch (const std::exception &e) { + std::cerr << "PipelineLayoutError: " << e.what() << '\n'; + return 2; + } +}