Skip to content

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

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

feat(aorta): disaggregated multi-node torchrun launch (AIMVT-173)#329
speriaswamy-amd wants to merge 1 commit into
surya/aorta-mn-03-configfrom
surya/aorta-mn-04-torchrun

Conversation

@speriaswamy-amd

Copy link
Copy Markdown
Contributor

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

Runs the Aorta benchmark across N nodes from a single cluster.json instead of
requiring N single-node cluster files. run() now branches on the launch mode
resolved from the multi_node block: 'script' is today's behavior verbatim
(experiment_script on the head node), 'torchrun' launches a rank-group on every
node in parallel, all rendezvous-ing on the head. 'auto' - the default - picks
script for 1-node clusters and torchrun above that, so existing single-node
configs are unaffected.

New helpers:
- _resolve_launch_mode() - auto/script/torchrun; 'script' with >1 node is a
  config error rather than a silent single-node run.
- _resolve_master_addr() - explicit override, else the head node's VPC/RDMA
  address from node_vpc_ips, else the 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 code path.

Co-Authored-By: Claude <noreply@anthropic.com>
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