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

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

Summary

InfiniTrain’s existing parallel execution model primarily launches multiple training threads within a single process. This PR adds a torchrun-style multi-process launcher, allowing each local process to bind to its own GPU while preserving the existing intra-process multithreading mode. It also fixes DataLoader and NCCL unique ID file conflicts in multi-process environments.

Add support for launching training processes through infini_run using nproc_per_node.
Ensure that the NCCL unique ID is generated only by the communication group’s root rank, with per-run file isolation, atomic publication, and cleanup.
Use global batches as the unit for distributed DataLoader partitioning and retrieval, aligning with Megatron’s behavior and fixing out-of-range batch access.
Add an 8-process integration test.

Changes

Update infini_run to:
- support -- as the launcher/training-args separator
- launch nproc_per_node child processes
- inject both InfiniTrain and torchrun-compatible rank env vars
- propagate child process failures via exit code
Update parallel runtime to:
- read torchrun-compatible env vars as fallback
- validate process topology and rank bounds
- map local process/thread rank to CUDA device index
Update GPT-2/Llama3 examples and parallel helpers to use local-device mapping.
Update scripts/run_models_and_profile.bash to:
- always launch model commands through infini_run
- treat nproc_per_node as launcher-only config
- keep nthread_per_process as the per-process thread count
Update scripts/test_config.json to use multi-process configs:
- 8-thread cases become nproc_per_node=8, nthread_per_process=1
- original 4-rank VPP cases become nproc_per_node=4, nthread_per_process=1
Add documentation describing behavior, compatibility, and example usage.

Compatibility

Existing direct runs remain supported:

./llama3 ... --nthread_per_process 8

The launcher can also preserve the old single-process multi-thread behavior:

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

The recommended single-node 8-GPU multi-process usage is:

./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

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

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 的环境变量命名，统一保留新名称即可，不需要兼容旧环境变量名

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 了；
以及这个函数感觉不适合有默认值。

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 函数同理。

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

按照之前讨论的结论，现有测例应保持不变，仅新增 8_proc 测例采用 infini_run 启动多进程分布式任务，但目前脚本将所有测例都改成统一使用 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);

chen2021673 force-pushed the 8_proc branch from 05edaac to 5314e45 Compare July 10, 2026 11:02

Support torchrun-style InfiniTrain multi-process launch

de24af9

chen2021673 force-pushed the 8_proc branch 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 from 072a288 to 350a2fd Compare July 30, 2026 06:36

chen2021673 added 3 commits July 30, 2026 06:42

test: add basic 8-process config group

08d684d

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

ec68923

- 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

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

chen2021673 added 2 commits July 30, 2026 07:01

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

c88f116

- 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

2463ca5

Chamberlain0w0 approved these changes Jul 31, 2026

View reviewed changes

kilinchange requested changes Aug 5, 2026

View reviewed changes

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

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!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

chen2021673 commented Jul 3, 2026 •

edited

Loading