Initialization time is too long in scale-out run

[youngeunkwon0405]

Initialization time is too long to try a scale-out run (# of GPUs > 2000).

When I was trying QWEN3 235B scale out run (2048 GPU-scale), it took 1 hour to see the following setup complete log.

============================================================
                  SETUP COMPLETE
============================================================

This is wasting 2048 GPU hours just for the initialization.

It is crucial to reduce this latency to initiate further scale-out performance study.

[youngeunkwon0405]

I am holding a further scale-out experiment due to this issue. CC-@guyueh1 @terrykong @parthchadha

[terrykong]

Is the slow init time you're observing coming the fact that vllm and training backend initialize serially? Or is even each one individually slow? Or something outside of that

[youngeunkwon0405]

I currently don't know which component is causing long latency in scale-out runs.
I am attaching a sample run log of my run (2048 GPU scale QWEN3 235B run). It took ~50min to reach SETUP COMPLETE (line 9022). Maybe setting NRL_FORCE_REBUILD_VENVS=true had an effect?

QWEN3-235B-scale-out.log

[terrykong]

Oh, that's for sure one cause. Try rebuilding the container or using a nightly that has all the dependencies you need and try running without NRL_FORCE_REBUILD_VENVS=true. FWIW, I have some WIP docs explaining that arg https://github.com/NVIDIA-NeMo/RL/blob/tk/explain-dependencies/docs/design-docs/dependency-manangement.md#option-1-force-rebuild-environments

[terrykong]

I don't know if it's possible to start both vllm and mcore up in parallel in the collocated setting. I think it's probably not b/c vllm and mcore would both be contending for the same VRAM during init.

But I'm thinking that actually non-collocated we should definitely do this (if not already done)

[youngeunkwon0405]

The doc looks great!
I was able to run the container without the NRL_FORCE_REBUILD_VENVS=true with the version-bump container. Maybe there was an improvement, but the story was pretty much the same.

[youngeunkwon0405]

I don't know if it's possible to start both vllm and mcore up in parallel in the collocated setting. I think it's probably not b/c vllm and mcore would both be contending for the same VRAM during init.

But I'm thinking that actually non-collocated we should definitely do this (if not already done)

Yeah, that might be worth to try thing. I believe there will be many works required to fix this, which I don't have a good idea.

[terrykong]

@youngeunkwon0405 could you share the 235b parallelism config for training and inference so I can measure how fast things should be in isolation?

[youngeunkwon0405]

@youngeunkwon0405 could you share the 235b parallelism config for training and inference so I can measure how fast things should be in isolation?

Of course! I used
Generation (1024 GPUs): TP8 DP128
Training (1024 GPUs): TP4 EP16 PP8 DP8

I attached the script and the yaml file I used in case it helps.

235b.sh
grpo_math_235B_megatron.yaml

[terrykong]

@youngeunkwon0405 i think this issue may be related to a ray.sub change we just reverted

6a40247

can you try after this commit? I ran the 235b.sh and it shows 18min.

============================================================
                  SETUP COMPLETE
  Total setup time: 1111.73s
============================================================

Since I already started, I will finish up the parallel init change i was working on, should bring that down further

[youngeunkwon0405]

Wow, that is a huge improvement compared to the past!
It will still be greatly beneficial if we can further reduce the setup time. Even with this improvement, assuming a 4K GPU DSV3 run, we might be spending 1300 GPU hours (20 min x 4096) in every 4 or 8 hours.