Using multi-node with LocalExecutor

[LopezGG]

I noticed LocalExecutor has a hard-coded value for nnodes.

https://github.com/NVIDIA/NeMo-Run/blob/b4e2258f61b88c53b77996b5f9ed871ee666d85f/src/nemo_run/core/execution/local.py#L53-L54

Is there a reason multi-nodes are disabled ? It feeds into torch_run which seems to support multi-nodes

https://github.com/NVIDIA/NeMo-Run/blob/b4e2258f61b88c53b77996b5f9ed871ee666d85f/src/nemo_run/run/torchx_backend/components/torchrun.py#L104-L124

Asking because I am using this with AML where I can usually get multi-node working with torchrun

[hemildesai]

Hey @LopezGG, thanks for the issue. Can you share a sample command you run with torchrun? We can add extra options based on that to support multi-node via LocalExecutor.

Currently LocalExecutor assumes you are on your local workstation which would just be a single node.

[LopezGG]

Thank you for the quick reply @hemildesai . Usually, with AML I use something like

 torchrun  --nproc_per_node=${{inputs.nproc_per_node}} --nnodes=${{inputs.nnodes}}
  --node_rank=$NODE_RANK
  --master_addr=$MASTER_ADDR
  --master_port=$MASTER_PORT  train.py

https://pytorch.org/docs/stable/elastic/run.html

$NODE_RANK, $MASTER_ADDR, and $MASTER_PORT are set automatically by AML or Slurm or can be set manually. Change in your script might look like ( I need to test it though)

Changing num_nodes will feed into

https://github.com/NVIDIA/NeMo-Run/blob/b4e2258f61b88c53b77996b5f9ed871ee666d85f/src/nemo_run/core/execution/base.py#L183-L188

and can be called from

[hemildesai]

Sounds good, we will make num nodes configurable

[LopezGG]

@hemildesai : Please Hold off on this. I had to make some changes in torchrun file as well. I'll write back when I have it working

[LopezGG]

@hemildesai : I finally got this working. Here are the changes I made to torchrun file. Master_addr is the environmental variable set by aml or other co-ordinating system.

else:
       master_addr = os.getenv("MASTER_ADDR" )
       master_port = os.getenv("MASTER_PORT")
        rdzv_endpoint = torchx_dist._noquote(master_addr + ":" + master_port)
        rdzv_port = master_port
        num_nodes = torchx_dist._noquote(nnodes_rep)
        node_rank = torchx_dist._noquote(os.getenv("NODE_RANK"))

    if env is None:
        env = {}

    env.setdefault("LOGLEVEL", os.getenv("LOGLEVEL", "INFO"))
    if debug:
        env.update(_TORCH_DEBUG_FLAGS)

    cmd = [
        "--rdzv-backend",
        rdzv_backend,
        "--rdzv-endpoint",
        rdzv_endpoint,
        "--rdzv-id",
        aisc_job_name, # A random int would work here
        "--nnodes",
        num_nodes,
        "--nproc-per-node",
        str(nproc_per_node),
        "--node-rank",
        node_rank,
        "--tee",
        "3", # tee: tees the specified std stream(s) to console + file. 0: none, 1: stdout, 2: stderr, 3: both
        "--role",
        "",
    ]

This works with tensor parallelism and context parallelism. For pipeline parallelism, nemo_run.core.runners.fdl_runner script passed to torchrun needs to specify how to split the model and co-ordinate. Would you know which parts of the code handle pipeline parallelism with slurm. I might be missing something.

[hemildesai]

Thanks @LopezGG , I'll create a PR with the changes you suggested.

[hemildesai]

Can you try out #143?

[LopezGG]

That was fast. Thank you, Hemil. This week I am with family in IST. I'll try it Indian morning tomorrow

[hemildesai]

Hi @LopezGG, did you get a chance to try this?

[LopezGG]

@hemildesai : I tried it last week. It is having issues recognizing the world_size. I am seeing

ValueError: You set devices=8andnum_nodes=2 in Lightning, but the product (8 * 2) does not match the world size (8).

I just returned from vacation this afternoon. I'll look further to see how this set of files diverges from previous ones

[LopezGG]

@hemildesai : I got it running with some changes (world_size error was a bug on my part)

Will look at the diff file and let you know which changes are essential.

---

[Update]:
For the changes in the PR to work with AML, we need to use "AISC_JOB_NAME" for rdzv-id. I added these lines to your code

    if os.getenv("AISC_JOB_NAME"):
        aisc_job_name = os.environ.get("AISC_JOB_NAME")
    else:
        aisc_job_name = random.randint(1, 10000)

without this, I was seeing the following error
torch.distributed.elastic.rendezvous.api.RendezvousTimeoutError
I am attaching the file here as txt (please rename to .py) in case you want to do a diff.

torchrun.txt

One other thing is that while calling the run.LocalExecutor, please be sure to pass in nodes parameter saying how many nodes are available for the run eg.,

def local_executor_torchrun(nodes: int = 1, devices: int = 8) -> run.LocalExecutor:
    env_vars = {
        "TRANSFORMERS_OFFLINE": "1",
        "TORCH_NCCL_AVOID_RECORD_STREAMS": "1",
        "NCCL_NVLS_ENABLE": "0",
        "NVTE_DP_AMAX_REDUCE_INTERVAL": "0",
        "NVTE_ASYNC_AMAX_REDUCTION": "1",
        "NVTE_FUSED_ATTN": "0",
    }

    executor = run.LocalExecutor(nodes=nodes, ntasks_per_node=devices, launcher="torchrun", env_vars=env_vars)
    return executor

without this, the code will show error with WORLD_SIZE. I vaguely remember that this has a default of 1.

I tested the code on 2 A100 (ND96amrs_A100_v4) nodes with 8 gpus each (run: adding it here so I can refer in case of future issues. You wont have access to it).

[aflah02]

Hi

Are there any updates on this? Does this work now? I see the PR is still open

I essentially want to use the local executor to run multinode llm training with NeMo instead of the slurm executor

[hemildesai]

We just merged #143 to support multi node jobs using local executor.