Skip to content

feat(aorta): disaggregated multi-node torchrun launch (AIMVT-173) - #329

Closed
speriaswamy-amd wants to merge 0 commit into
surya/aorta-mn-03-configfrom
surya/aorta-mn-04-torchrun
Closed

speriaswamy-amd wants to merge 0 commit into
surya/aorta-mn-03-configfrom
surya/aorta-mn-04-torchrun

Conversation

@speriaswamy-amd

Copy link
Copy Markdown
Collaborator

Stack 4/6 — splits #171. Base: #328. Implements AIMVT-173.

Why

Run the Aorta benchmark across N nodes from a single cluster.json instead of requiring N single-node cluster files, mirroring the disaggregated pattern the PyTorch xDiT and SGLang multi-node suites already use.

What changed

run() branches on the resolved launch mode:

  • script — today's behavior verbatim (experiment_script on the head node).
  • torchrun — a rank-group on every node in parallel, rendezvous-ing on the head.
  • auto (default) — script for 1-node clusters, torchrun above that. Existing single-node configs are unaffected.

New helpers in cvs/runners/aorta.py:

  • _resolve_launch_mode() — script with >1 node is a config error rather than a silent single-node run.
  • _resolve_master_addr() — explicit override → head node's VPC/RDMA address → plain node id. The mgmt/SSH address is frequently orchestrator-only and unreachable from peer nodes.
  • _pick_free_port_on() / _pick_master_port() — binds and releases a port over SSH on the head node, which is the host that actually holds the rendezvous socket. The snippet is passed as one shlex-quoted argument; as separate argv elements the remote shell re-splits it on its own semicolons.
  • _build_torchrun_command(), _run_single_node().

TraceLens/GEMM analysis now explicitly target the head node's container, and training-log discovery iterates the node list, rather than relying on locals left over from the single-node path.

Test

ruff clean. Unit tests 603 → 623 (20 new: mode resolution, port/addr selection, command construction, validate_config).

@speriaswamy-amd

Copy link
Copy Markdown
Collaborator Author

Reopening — this was auto-closed by GitHub when a rebase mistake on my end briefly collapsed this branch to be identical to its base (0 commits/0 diff). The branch has since been corrected and force-pushed with the real commit restored; verified diff is +564/-42 across 2 files as intended.

@speriaswamy-amd

Copy link
Copy Markdown
Collaborator Author

Superseded by #390 (same branches, same content) — this PR got stuck in an auto-closed state that GitHub wouldn't let me reopen after a rebase mistake temporarily zeroed its diff. No content was lost.

@speriaswamy-amd
speriaswamy-amd removed this pull request from stack #333 September 11, 2026 00:23
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant