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:21:20Z

    } else {
-        Train({0, 0, 1, 1});
+        nn::parallel::Rank rank(nn::parallel::global::GetGlobalProcRank(), 0, nn::parallel::global::GetNprocPerNode(),
+                                FLAGS_nthread_per_process);


Rank 构造函数是：Rank::Rank(int process_rank, int thread_rank, int process_size, int thread_size)

原先第三个参数传的是 global num of processes，现在换成了 num of processes per rank，层次变化了，而 rank 其他 function 逻辑都没改，会在 rank 相关的判断中出现错误，比如 bool Rank::IsParallel() const { return thread_size_ * process_size_ > 1; } 里面会对 nnode=N 但是每个节点单进程、单线程的情况判断为 IsParallel() == false。最好 check 一下 Rank 类里面的逻辑。

第三个参数原来的注释就是Total number of processes on this node，可能是笔误。这里我修改如下：

保持参数含义不变（num of processes per node），命名改成更易懂的 processes_per_node_ 等；

修改 IsParallel() 函数，用GetWorldSize() 判断；

检查其他地方有没有类似错误，修改。

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

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

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Support torchrun-style InfiniTrain multi-process launch - #184

Support torchrun-style InfiniTrain multi-process launch#184
chen2021673 wants to merge 6 commits into
masterfrom
8_proc

chen2021673 commented Jul 3, 2026 •

edited

Loading

Uh oh!

Chamberlain0w0 Jul 22, 2026

Uh oh!

chen2021673 Jul 29, 2026 •

edited

Loading

Uh oh!

Chamberlain0w0 Jul 22, 2026

Uh oh!

chen2021673 Jul 29, 2026

Uh oh!

Chamberlain0w0 Jul 22, 2026

Uh oh!

chen2021673 Jul 30, 2026

Uh oh!

Chamberlain0w0 Jul 22, 2026

Uh oh!

chen2021673 Jul 30, 2026

Uh oh!

Chamberlain0w0 Jul 22, 2026

Uh oh!

chen2021673 Jul 30, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

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!

Choose a reason for hiding this comment

Uh oh!

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

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!

2 participants

chen2021673 commented Jul 3, 2026 •

edited

Loading

chen2021673 Jul 29, 2026 •

edited

Loading