diff --git a/README.md b/README.md index 84cd4d9b5..f3fc6a317 100644 --- a/README.md +++ b/README.md @@ -105,6 +105,9 @@ To view available runtime options: ./build/llama3 --help ``` +`--help` shows only the model's options. Use `--helpfull` to include all options +from third-party libraries. + ### Getting Started #### Prepare Datasets and Weights diff --git a/cmake/test_macros.cmake b/cmake/test_macros.cmake index cabee1d54..1d2a4c9e6 100644 --- a/cmake/test_macros.cmake +++ b/cmake/test_macros.cmake @@ -474,7 +474,7 @@ function(infini_train_add_privateuse1_test_suites) GTEST_PREFIX PRIVATEUSE1 DEVICE_TYPE kPrivateUse1 DEVICE_INDEX ${PRIVATEUSE1_DEVICE_INDEX} - LABELS ${PRIVATEUSE1_BACKEND_NAME} accelerator hardware + LABELS ${PRIVATEUSE1_BACKEND_NAME} TEST_MAIN_TARGET ${_privateuse1_main_target} LINK_LIBRARIES ${PRIVATEUSE1_LINK_LIBRARIES} ${_privateuse1_timeout_args} diff --git a/docs/test_usage_guide.md b/docs/test_usage_guide.md index 4267a351f..ffbd378bb 100644 --- a/docs/test_usage_guide.md +++ b/docs/test_usage_guide.md @@ -150,7 +150,7 @@ infini_train_add_privateuse1_test_suites( ``` 该调用为公共 suite 追加 `test_*_` 目标；`DEVICE_INDEX` 默认为 -`0`。例如 `BACKEND_NAME=maca` 使用 `maca;accelerator;hardware` CTest 标签： +`0`。例如 `BACKEND_NAME=maca` 使用 `maca` CTest 标签： ```bash ctest -L maca --output-on-failure # MACA diff --git a/example/common/cli.h b/example/common/cli.h new file mode 100644 index 000000000..599743748 --- /dev/null +++ b/example/common/cli.h @@ -0,0 +1,24 @@ +#pragma once + +#include + +#include "gflags/gflags.h" + +namespace infini_train::example { + +inline void ParseModelFlags(int *argc, char ***argv, const char *model_path) { + gflags::ParseCommandLineNonHelpFlags(argc, argv, true); + + // Keep gflags' --helpshort precedence and allow --helpfull to show all flags. + if (gflags::GetCommandLineFlagInfoOrDie("help").current_value == "true" && + gflags::GetCommandLineFlagInfoOrDie("helpshort").current_value != + "true" && + gflags::GetCommandLineFlagInfoOrDie("helpfull").current_value != "true") { + gflags::ShowUsageWithFlagsRestrict((*argv)[0], model_path); + std::exit(1); // Match gflags' exit status for help output. + } + + gflags::HandleCommandLineHelpFlags(); +} + +} // namespace infini_train::example diff --git a/example/gpt2/main.cc b/example/gpt2/main.cc index e0d355888..699c137c3 100644 --- a/example/gpt2/main.cc +++ b/example/gpt2/main.cc @@ -61,6 +61,7 @@ #include "infini_train/include/utils/precision_check_config.h" #include "infini_train/include/utils/precision_checker.h" +#include "example/common/cli.h" #include "example/common/tiny_shakespeare_dataset.h" #include "example/common/tokenizer.h" #include "example/gpt2/checkpoint_loader.h" @@ -727,12 +728,16 @@ void Train(const nn::parallel::Rank &rank) { } int main(int argc, char *argv[]) { + gflags::SetUsageMessage( + "Train a GPT-2 model with InfiniTrain.\nUsage: " + std::string(argv[0]) + + " --input_bin= [options]\n" + "Model options: --help\nAll options: --helpfull"); // Register provider metadata and implementations before gflags validates // --device. The device runtime initializes lazily on first DeviceGuard use. #ifdef INFINITRAIN_EXAMPLE_EXTERNAL_BACKEND_REGISTRAR INFINITRAIN_EXAMPLE_EXTERNAL_BACKEND_REGISTRAR(); #endif - gflags::ParseCommandLineFlags(&argc, &argv, true); + example::ParseModelFlags(&argc, &argv, "example/gpt2/"); google::InitGoogleLogging(argv[0]); auto precision_config = diff --git a/example/llama3/main.cc b/example/llama3/main.cc index f8188cef1..500d05d16 100644 --- a/example/llama3/main.cc +++ b/example/llama3/main.cc @@ -60,6 +60,7 @@ #include "infini_train/include/profiler.h" #endif +#include "example/common/cli.h" #include "example/common/tiny_shakespeare_dataset.h" #include "example/common/tokenizer.h" #include "example/llama3/checkpoint_loader.h" @@ -702,12 +703,16 @@ void Train(const nn::parallel::Rank &rank) { } int main(int argc, char *argv[]) { + gflags::SetUsageMessage("Train a LLaMA 3 model with InfiniTrain.\nUsage: " + + std::string(argv[0]) + + " --input_bin= [options]\n" + "Model options: --help\nAll options: --helpfull"); // Register provider metadata and implementations before gflags validates // --device. The device runtime initializes lazily on first DeviceGuard use. #ifdef INFINITRAIN_EXAMPLE_EXTERNAL_BACKEND_REGISTRAR INFINITRAIN_EXAMPLE_EXTERNAL_BACKEND_REGISTRAR(); #endif - gflags::ParseCommandLineFlags(&argc, &argv, true); + example::ParseModelFlags(&argc, &argv, "example/llama3/"); google::InitGoogleLogging(argv[0]); auto precision_config = diff --git a/example/mixtral/main.cc b/example/mixtral/main.cc index 82ee2b392..9ade077bf 100644 --- a/example/mixtral/main.cc +++ b/example/mixtral/main.cc @@ -13,6 +13,10 @@ #include #include +#ifdef INFINITRAIN_EXAMPLE_EXTERNAL_BACKEND_WORKAROUNDS +#include +#endif + #include "gflags/gflags.h" #include "glog/logging.h" @@ -49,6 +53,7 @@ #include "infini_train/include/profiler.h" #endif +#include "example/common/cli.h" #include "example/common/tiny_shakespeare_dataset.h" #include "example/mixtral/checkpoint_loader.h" #include "example/mixtral/config.h" @@ -366,12 +371,16 @@ void Train(const nn::parallel::Rank &rank) { } int main(int argc, char *argv[]) { + gflags::SetUsageMessage("Train a Mixtral model with InfiniTrain.\nUsage: " + + std::string(argv[0]) + + " --input_bin= [options]\n" + "Model options: --help\nAll options: --helpfull"); // Register provider metadata and implementations before gflags validates // --device. The device runtime initializes lazily on first DeviceGuard use. #ifdef INFINITRAIN_EXAMPLE_EXTERNAL_BACKEND_REGISTRAR INFINITRAIN_EXAMPLE_EXTERNAL_BACKEND_REGISTRAR(); #endif - gflags::ParseCommandLineFlags(&argc, &argv, true); + example::ParseModelFlags(&argc, &argv, "example/mixtral/"); google::InitGoogleLogging(argv[0]); CHECK_GT(FLAGS_nthread_per_process, 0); @@ -413,8 +422,28 @@ int main(int argc, char *argv[]) { Train(rank); } +#ifdef INFINITRAIN_EXAMPLE_EXTERNAL_BACKEND_WORKAROUNDS + // Capture the policy before shutting down gflags. Match the other external + // backend examples: MACA can fail during process-wide static destruction. + const bool skip_static_destruction = + INFINITRAIN_EXAMPLE_EXTERNAL_BACKEND_WORKAROUNDS && + Device::ParseType(FLAGS_device).value() == + Device::DeviceType::kPrivateUse1; +#endif + gflags::ShutDownCommandLineFlags(); google::ShutdownGoogleLogging(); +#ifdef INFINITRAIN_EXAMPLE_EXTERNAL_BACKEND_WORKAROUNDS + if (skip_static_destruction) { + // Train() has returned and all workers have joined, so rank-local tensors + // and optimizers have already been destroyed. Flush before bypassing the + // remaining global/SDK teardown; _Exit does not flush C++ streams. + std::cout.flush(); + std::cerr.flush(); + std::_Exit(0); + } +#endif + return 0; } diff --git a/example/mnist/main.cc b/example/mnist/main.cc index 66c22b4da..8fc3b78ea 100644 --- a/example/mnist/main.cc +++ b/example/mnist/main.cc @@ -20,6 +20,7 @@ #include "infini_train/include/optimizer.h" #include "infini_train/include/tensor.h" +#include "example/common/cli.h" #include "example/mnist/dataset.h" #include "example/mnist/net.h" @@ -44,7 +45,11 @@ DEFINE_validator(device, [](const char *, const std::string &value) { }); int main(int argc, char *argv[]) { - gflags::ParseCommandLineFlags(&argc, &argv, true); + gflags::SetUsageMessage("Train a model on MNIST with InfiniTrain.\nUsage: " + + std::string(argv[0]) + + " --dataset= [options]\n" + "Model options: --help\nAll options: --helpfull"); + example::ParseModelFlags(&argc, &argv, "example/mnist/"); google::InitGoogleLogging(argv[0]); auto train_dataset = std::make_shared(FLAGS_dataset, true); diff --git a/example/qwen3/main.cc b/example/qwen3/main.cc index a9f925fef..0998c14ab 100644 --- a/example/qwen3/main.cc +++ b/example/qwen3/main.cc @@ -60,6 +60,7 @@ #include "infini_train/include/profiler.h" #endif +#include "example/common/cli.h" #include "example/common/tiny_shakespeare_dataset.h" #include "example/common/tokenizer.h" #include "example/qwen3/checkpoint_loader.h" @@ -700,12 +701,16 @@ void Train(const nn::parallel::Rank &rank) { } int main(int argc, char *argv[]) { + gflags::SetUsageMessage( + "Train a Qwen3 model with InfiniTrain.\nUsage: " + std::string(argv[0]) + + " --input_bin= [options]\n" + "Model options: --help\nAll options: --helpfull"); // Register provider metadata and implementations before gflags validates // --device. The device runtime initializes lazily on first DeviceGuard use. #ifdef INFINITRAIN_EXAMPLE_EXTERNAL_BACKEND_REGISTRAR INFINITRAIN_EXAMPLE_EXTERNAL_BACKEND_REGISTRAR(); #endif - gflags::ParseCommandLineFlags(&argc, &argv, true); + example::ParseModelFlags(&argc, &argv, "example/qwen3/"); google::InitGoogleLogging(argv[0]); auto precision_config = diff --git a/infini_train/src/nn/parallel/ddp/distributed_optimizer.cc b/infini_train/src/nn/parallel/ddp/distributed_optimizer.cc index 331e7a524..7389f6d46 100644 --- a/infini_train/src/nn/parallel/ddp/distributed_optimizer.cc +++ b/infini_train/src/nn/parallel/ddp/distributed_optimizer.cc @@ -168,7 +168,8 @@ void DistributedOptimizer::BuildShardParamsAndBindGrads( } CHECK_GT(num_shard_params, 0) - << "DistributedOptimizer: this DP rank owns no param pieces. " + << "DistributedOptimizer: DP rank " << ddp_rank_ + << " owns no param pieces. " << "Check bucket padding/divisibility and param bucketing order."; } diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index e2029460d..bfe52949b 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -62,7 +62,7 @@ if(USE_CUDA) GTEST_PREFIX CUDA DEVICE_TYPE kCUDA DEVICE_INDEX 0 - LABELS cuda accelerator hardware + LABELS cuda ACCELERATOR ) endif() diff --git a/tests/backend/CMakeLists.txt b/tests/backend/CMakeLists.txt index 6bb3affb4..4ae0fce24 100644 --- a/tests/backend/CMakeLists.txt +++ b/tests/backend/CMakeLists.txt @@ -11,7 +11,7 @@ infini_train_add_test(test_privateuse1_backend_validation if(USE_CUDA) infini_train_add_test(test_stream_guard_cuda SOURCES test_stream_guard.cc - LABELS cuda accelerator hardware + LABELS cuda BINARY_REGISTRATION RUN_SERIAL TEST_TIMEOUT 60 diff --git a/tools/infini_run/infini_run.cc b/tools/infini_run/infini_run.cc index 51ad1e0c2..ffbf81010 100644 --- a/tools/infini_run/infini_run.cc +++ b/tools/infini_run/infini_run.cc @@ -110,6 +110,10 @@ std::string GenerateLocalRunId() { } // namespace int main(int argc, char **argv) { + gflags::SetUsageMessage( + "Launch distributed training with InfiniTrain.\nUsage: " + + std::string(argv[0]) + + " [launcher options] [training options]"); const int train_program_index = FindTrainProgramIndex(argc, argv); std::vector launcher_argv = BuildLauncherArgv(train_program_index, argv);