From 7917f6236f4f63f0a6103a43c44c11ac78c8e07f Mon Sep 17 00:00:00 2001 From: cx Date: Fri, 30 Oct 2026 04:21:56 +0000 Subject: [PATCH 1/2] fix: improve CLI help and external backend shutdown handling - Add usage messages to training examples and infini_run - Add Mixtral external backend shutdown workaround and flush output before exit - Include DP rank in distributed optimizer empty-shard errors - Remove accelerator and hardware test labels and update documentation --- cmake/test_macros.cmake | 2 +- docs/test_usage_guide.md | 2 +- example/gpt2/main.cc | 3 +++ example/llama3/main.cc | 3 +++ example/mixtral/main.cc | 27 +++++++++++++++++++ example/mnist/main.cc | 3 +++ example/qwen3/main.cc | 3 +++ .../nn/parallel/ddp/distributed_optimizer.cc | 3 ++- tests/CMakeLists.txt | 2 +- tests/backend/CMakeLists.txt | 2 +- tools/infini_run/infini_run.cc | 4 +++ 11 files changed, 49 insertions(+), 5 deletions(-) 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/gpt2/main.cc b/example/gpt2/main.cc index e0d355888..37cefcc56 100644 --- a/example/gpt2/main.cc +++ b/example/gpt2/main.cc @@ -727,6 +727,9 @@ 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]"); // Register provider metadata and implementations before gflags validates // --device. The device runtime initializes lazily on first DeviceGuard use. #ifdef INFINITRAIN_EXAMPLE_EXTERNAL_BACKEND_REGISTRAR diff --git a/example/llama3/main.cc b/example/llama3/main.cc index f8188cef1..b5d3f82fa 100644 --- a/example/llama3/main.cc +++ b/example/llama3/main.cc @@ -702,6 +702,9 @@ 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]"); // Register provider metadata and implementations before gflags validates // --device. The device runtime initializes lazily on first DeviceGuard use. #ifdef INFINITRAIN_EXAMPLE_EXTERNAL_BACKEND_REGISTRAR diff --git a/example/mixtral/main.cc b/example/mixtral/main.cc index 82ee2b392..a3482d8f4 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" @@ -366,6 +370,9 @@ 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]"); // Register provider metadata and implementations before gflags validates // --device. The device runtime initializes lazily on first DeviceGuard use. #ifdef INFINITRAIN_EXAMPLE_EXTERNAL_BACKEND_REGISTRAR @@ -413,8 +420,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..51e465b61 100644 --- a/example/mnist/main.cc +++ b/example/mnist/main.cc @@ -44,6 +44,9 @@ DEFINE_validator(device, [](const char *, const std::string &value) { }); int main(int argc, char *argv[]) { + gflags::SetUsageMessage("Train a model on MNIST with InfiniTrain.\nUsage: " + + std::string(argv[0]) + + " --dataset= [options]"); gflags::ParseCommandLineFlags(&argc, &argv, true); google::InitGoogleLogging(argv[0]); diff --git a/example/qwen3/main.cc b/example/qwen3/main.cc index a9f925fef..97b95818d 100644 --- a/example/qwen3/main.cc +++ b/example/qwen3/main.cc @@ -700,6 +700,9 @@ 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]"); // Register provider metadata and implementations before gflags validates // --device. The device runtime initializes lazily on first DeviceGuard use. #ifdef INFINITRAIN_EXAMPLE_EXTERNAL_BACKEND_REGISTRAR 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); From d66f300a1b0476ffbc924ed95a64e33609b148c9 Mon Sep 17 00:00:00 2001 From: cx Date: Fri, 6 Nov 2026 06:34:11 +0000 Subject: [PATCH 2/2] fix: show only model options in default CLI help Share model help filtering across the training examples and retain --helpfull for all flags. Document the two help modes in usage output and the README. --- README.md | 3 +++ example/common/cli.h | 24 ++++++++++++++++++++++++ example/gpt2/main.cc | 6 ++++-- example/llama3/main.cc | 6 ++++-- example/mixtral/main.cc | 6 ++++-- example/mnist/main.cc | 6 ++++-- example/qwen3/main.cc | 6 ++++-- 7 files changed, 47 insertions(+), 10 deletions(-) create mode 100644 example/common/cli.h 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/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 37cefcc56..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" @@ -729,13 +730,14 @@ 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]"); + " --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 b5d3f82fa..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" @@ -704,13 +705,14 @@ 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]"); + " --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 a3482d8f4..9ade077bf 100644 --- a/example/mixtral/main.cc +++ b/example/mixtral/main.cc @@ -53,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" @@ -372,13 +373,14 @@ 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]"); + " --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); diff --git a/example/mnist/main.cc b/example/mnist/main.cc index 51e465b61..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" @@ -46,8 +47,9 @@ DEFINE_validator(device, [](const char *, const std::string &value) { int main(int argc, char *argv[]) { gflags::SetUsageMessage("Train a model on MNIST with InfiniTrain.\nUsage: " + std::string(argv[0]) + - " --dataset= [options]"); - gflags::ParseCommandLineFlags(&argc, &argv, true); + " --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 97b95818d..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" @@ -702,13 +703,14 @@ 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]"); + " --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 =