Support torchrun-style InfiniTrain multi-process launch by chen2021673 · Pull Request #184 · InfiniTensor/InfiniTrain

chen2021673 · 2026-07-03T09:10:56Z

Summary

This PR adds torchrun-style multi-process execution to InfiniTrain while preserving the existing single-process, multi-threaded mode.

infini_run can now launch multiple local worker processes through nproc_per_node, allowing each process to bind to a separate GPU. Each worker may still run multiple training threads through nthread_per_process.

The PR also improves NCCL bootstrap synchronization and fixes distributed DataLoader indexing in multi-process environments.

Changes

Multi-Process Launcher

Launch nproc_per_node worker processes on each node.
Parse launcher options before the training executable and forward the remaining arguments to the training program.
Pass the training executable directly after launcher options; the explicit -- separator is no longer supported.
Export torchrun-compatible rank and world-size environment variables.
Terminate remaining workers and propagate a non-zero exit code when a worker fails.

Parallel Runtime

Read torchrun-compatible environment variables while retaining legacy environment-variable fallbacks.
Validate process topology, world size, and rank bounds.
Map each local process and training thread to the appropriate CUDA device.
Update GPT-2, Llama3, and parallel helpers to use local-device mapping.

NCCL Bootstrap

Generate the NCCL unique ID only on the communication group’s root rank.
Isolate unique ID files by run.
Publish unique IDs atomically and clean up synchronization files after execution.

Distributed DataLoader

Partition datasets using global batches.
Retrieve data using global batch indices.
Align distributed batch behavior with Megatron.
Prevent out-of-range batch access in distributed training.

Scripts and Tests

Run model test commands through infini_run.
Treat nproc_per_node as a launcher-only configuration.
Preserve the existing single-process, multi-threaded test cases.
Add a separate 8_proc model test group for GPT-2 and Llama3 using nproc_per_node=8 and nthread_per_process=1.
Cover DP, TP, TP+SP, PP, and TP2+SP+PP2+VPP2 configurations in FP32 and BF16.
Keep the existing four-rank PP4/VPP2 cases in single-process, four-thread mode.
Add DataLoader and distributed-rank unit tests.

Compatibility

Direct single-process, multi-threaded execution remains supported:

./llama3 ... --nthread_per_process 8

The launcher can preserve the same execution model:

./infini_run --nproc_per_node=1 ./llama3 ... --nthread_per_process 8

Recommended single-node, 8-GPU multi-process execution:

./infini_run --nproc_per_node=8 ./llama3 ... --nthread_per_process 1

Test

Chamberlain0w0 · 2026-07-22T02:33:30Z

    virtual void GetAsyncError(const CclComm *comm, CclStatus *async_error) const;

-    virtual void GetUniqueId(CclUniqueId **unique_id) const;
+    virtual void CreateUniqueId(CclUniqueId **unique_id, bool generate_id) const;


这个接口的名字和参数都不建议改，因为对标的是 nccl 接口 GetUniqueId() ，后续国产平台的应该也都是类似签名。

下面的 nccl_impl.h 的继承实现也得改回来。

Chamberlain0w0 · 2026-07-22T02:36:52Z

+            SetEnvInt("WORLD_SIZE", proc_world_size);
+            SetEnvInt("GROUP_RANK", FLAGS_node_rank);
+            SetEnvInt("ROLE_RANK", global_proc_rank);
+            SetEnvInt("ROLE_WORLD_SIZE", proc_world_size);


这几个好像没用到？可以先删掉

Chamberlain0w0 · 2026-07-22T03:22:58Z

+            }
+        } else if (exit_code == 0) {
+            exit_code = 1;
+        }


这块退出，好像也没有做某个子进程异常退出的时候清理其他进程的逻辑？如果 exit code 非 0 的话感觉正常情况应该要把所有其他正在运行的子进程都清理完毕再返回

修改逻辑为：跟踪所有运行中的子进程；任一子进程异常退出或 fork 失败时，向其余进程发送 SIGTERM，继续回收全部子进程后返回首个失败码。

Chamberlain0w0 · 2026-07-22T03:32:36Z

+    int proc_world_size = FLAGS_nnodes * FLAGS_nproc_per_node;
    std::string master_addr = FLAGS_rdzv_endpoint.substr(0, FLAGS_rdzv_endpoint.find(':'));
    std::string master_port = FLAGS_rdzv_endpoint.substr(FLAGS_rdzv_endpoint.find(':') + 1);
+    const std::string run_id = FLAGS_nnodes == 1 ? GenerateLocalRunId() : "";


这块多机还是会使用原先的默认命名，没达到效果，可能得看下怎么改。

修改：多机必须通过 --rdzv_id 指定；单机未指定时仍自动生成。参考https://docs.pytorch.org/docs/2.13/elastic/run.html

kilinchange · 2026-08-05T02:26:17Z

  --nnodes=2 \
  --nproc_per_node=1 \
-  --node_rank=[rank_id] \
-  -- ./llama3 \


需要确认下目前的使用行为，README 里删除了 -- 作为启动/训练参数的分隔符？但我看 pr description 里写的是 "support -- as the launcher/training-args separator"，原始版本也是通过 -- 作为分隔符的。

顺便指出 pr description 的另一处小问题：

这里似乎需要更新一下，目前应该是多进程/多线程测例同时保留了，vpp case 的说明似乎也和实际代码不一致。

kilinchange · 2026-08-05T08:04:59Z

+        << "GLOBAL_PROC_RANK/RANK must be less than PROC_WORLD_SIZE/WORLD_SIZE";
+    CHECK_GE(local_proc_rank_, 0) << "LOCAL_PROC_RANK/LOCAL_RANK must be non-negative";
+    CHECK_LT(local_proc_rank_, nproc_per_node_)
+        << "LOCAL_PROC_RANK/LOCAL_RANK must be less than NPROC_PER_NODE/LOCAL_WORLD_SIZE";


https://docs.pytorch.org/docs/2.13/elastic/run.html#environment-variables
修改后更加贴近 torchrun 的环境变量命名，统一保留新名称即可，不需要兼容旧环境变量名

仅读取和 torchrun 对应的 WORLD_SIZE、LOCAL_WORLD_SIZE、RANK、LOCAL_RANK，移除旧环境变量的回退逻辑

kilinchange · 2026-08-05T08:06:57Z

-    for (int index = 0; index < global::GetNthreadPerProc(); ++index) { devices_.emplace_back(device_type, index); }
+    for (int thread_rank = 0; thread_rank < global::GetNthreadPerProc(); ++thread_rank) {
+        const int device_index
+            = device_type == Device::DeviceType::kCUDA ? global::GetLocalDeviceIndex(thread_rank) : thread_rank;


不需要特判 cuda 类型，其他硬件遵循 cuda 分配 device index 的规则。

kilinchange · 2026-08-05T08:12:55Z

 inline int GetNthreadPerProc() { return GlobalEnv::Instance().nthread_per_process(); }
 inline int GetGlobalProcRank() { return GlobalEnv::Instance().global_proc_rank(); }
 inline int GetLocalProcRank() { return GlobalEnv::Instance().local_proc_rank(); }
+inline int GetLocalDeviceIndex(int thread_rank = 0) { return GetLocalProcRank() * GetNthreadPerProc() + thread_rank; }


device index 本身就特指节点内的设备编号，是不是没必要强调 local 了；
以及这个函数感觉不适合有默认值。

接口统一重命名为 GetDeviceIndex，移除 thread_rank 默认值。

kilinchange · 2026-08-05T08:15:57Z

+    if (IsCPU()) {
+        return {nn::parallel::global::GetGlobalProcRank(), 0, nn::parallel::global::GetNprocPerNode(),
+                nn::parallel::global::GetNthreadPerProc()};
+    }


cpu 的情况在 Device 的构造函数里已经处理过了，这里不需要特判。

kilinchange · 2026-08-05T09:06:54Z

dataloader 相关的改动应该与多进程启动无关，建议拆出独立 pr 进行 review 和合入流程。

kilinchange · 2026-08-05T09:14:19Z

 namespace {
-std::string UniqueIdFileName(const std::string &name, bool tmp = false) {
-    return "cclUniqueId_" + name + (tmp ? ".tmp" : ".bin");
+std::string UniqueIdPath(const std::string &pg_name) {


这里生成的确实是 file name 而不是 path 吧，应该不需要修改函数名，下面 tmp 函数同理。

已修改：
UniqueIdPath → UniqueIdFileName
UniqueIdTmpPath → UniqueIdTmpFileName

kilinchange · 2026-08-05T09:26:38Z

按照之前讨论的结论，现有测例应保持不变，仅新增 8_proc 测例采用 infini_run 启动多进程分布式任务，但目前脚本将所有测例都改成统一使用 infini_run 启动了，麻烦修改一下。

修改：原有组继续沿用原来的直接执行命令，仅新增 8_proc 测试通过 infini_run 启动

kilinchange · 2026-08-05T09:36:06Z


-    std::string train_program = argv[1];
+    std::string train_program = argv[train_program_index];
+    CHECK_NE(train_program, "--") << "Explicit '--' separator is not supported; pass the training program directly "


使用 -- 分隔符调用 infini_run 时，前面 FindTrainProgramIndex 返回 -- 的位置，但这里又明确拒绝了 --，导致使用 -- 分隔训练参数的使用方式会直接报错，建议同时支持使用/不使用 -- 分隔的调用方式。

kilinchange · 2026-08-05T09:53:47Z


-    CHECK_GE(argc, 2) << "No training prgram specified!";
+    CHECK_GT(FLAGS_nnodes, 0) << "nnodes must be positive";
+    CHECK_GT(FLAGS_nproc_per_node, 0) << "nproc_per_node must be positive";


再检查下

CHECK_GE(FLAGS_node_rank, 0); CHECK_LT(FLAGS_node_rank, FLAGS_nnodes);

Add a dedicated 8_proc test group containing the 8-process variants of the original basic multi-GPU cases.

Track DataLoader progress by global batches so distributed ranks slice data consistently and can resume/cycle from saved consumption counts. Also scope CCL unique ID files per run, generate NCCL IDs only on the main rank, clean up run-local rendezvous files, and add DataLoader coverage.

- derive parallel state from the global world size - clarify global rank and per-node process semantics - add multi-node rank regression coverage - restore the NCCL-compatible GetUniqueId interface

- add torchrun-style --rdzv_id support - use the shared ID to isolate CCL unique-ID files - preserve automatic run ID generation for single-node runs - document rdzv_id in the multi-node example

- support infini_run with or without the optional -- separator - validate node rank bounds - use infini_run only for the new 8_proc test group - standardize torchrun environment variables and device index mapping - clarify NCCL unique ID filename helpers

chen2021673 force-pushed the 8_proc branch 2 times, most recently from 5314e45 to 63aefdc Compare July 16, 2026 07:53

kilinchange requested review from Chamberlain0w0 and JYMiracle305 July 17, 2026 06:11

Chamberlain0w0 requested changes Jul 22, 2026

View reviewed changes

chen2021673 requested a review from Chamberlain0w0 July 30, 2026 06:22

chen2021673 force-pushed the 8_proc branch 2 times, most recently from 350a2fd to c88f116 Compare July 30, 2026 06:54

Chamberlain0w0 approved these changes Jul 31, 2026

View reviewed changes

kilinchange requested changes Aug 5, 2026

View reviewed changes

chen2021673 added 7 commits August 10, 2026 09:35

Support torchrun-style InfiniTrain multi-process launch

d453087

test: add basic 8-process config group

84a83b7

Add a dedicated 8_proc test group containing the 8-process variants of the original basic multi-GPU cases.

fix: correct distributed rank and unique ID handling

78900d9

- derive parallel state from the global world size - clarify global rank and per-node process semantics - add multi-node rank regression coverage - restore the NCCL-compatible GetUniqueId interface

fix: require a shared run ID for multi-node launch

169c6b5

- add torchrun-style --rdzv_id support - use the shared ID to isolate CCL unique-ID files - preserve automatic run ID generation for single-node runs - document rdzv_id in the multi-node example

fix: separate launcher args from model test args

513a57b

refactor: separate DataLoader changes from launcher

0431ee7

chen2021673 force-pushed the 8_proc branch from 2463ca5 to 50c8039 Compare August 10, 2026 10:05

Uh oh!

Conversation

chen2021673 commented Jul 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Changes

Multi-Process Launcher

Parallel Runtime

NCCL Bootstrap

Distributed DataLoader

Scripts and Tests

Compatibility

Test

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

chen2021673 commented Jul 3, 2026 •

edited

Loading