Support torchrun-style InfiniTrain multi-process launch

[chen2021673]

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