Requires a way to control segment setting for NVL72 system

[youngeunkwon0405]

Is your feature request related to a problem? Please describe.

In the NVL72 system non-colocated setup, we might want to configure each virtual cluster to be assigned to the same rack. We can specify --segment variable in the slurm command, but that only guarantees the global GPU allocation. Within the global allocation, we might need a way to control Training/Generation GPU assignment considering the segment.

Describe the solution you'd like
A clear and concise description of what you want to happen.

Describe alternatives you've considered
A clear and concise description of any alternative solutions or features you've considered.

Additional context
Add any other context or screenshots about the feature request here.

[youngeunkwon0405]

@guyueh1

[guyueh1]

Note that for co-located jobs, you just need to pass the --segment 16 when you run ray.sub, for example, sbatch --segment=16 ray.sub this is all it takes to specify segment.

What's complicated is non-colocated, where we have M nodes for training and N nodes for inference, and we want those M and N nodes both allocated to maximize the gpu count in one nvlink domain. For instance, M=16 and N=8, we want 16 in one rack and 8 in another instead of 12 in each rack.

@youngeunkwon0405 @seonjinn is it ok to assign you both to look into the idea of having separate srun jobs for non-colocated jobs?

[youngeunkwon0405]

Yes, I think it will be good to have a one-off meeting with @terrykong to discuss this issue.
Option #1: Handle this from ray.sub (idea what you suggested today)
--> This sounds like a very clean solution, but this might break some of the basic abstractions that now ray.sub might refer to config.yaml file or config.yaml might refer to ray.sub argument for consistency.

Option #2: handle this from the NeMo-RL source code.
--> This might require a feature request to slurm since currently we cannot specify something like --segment=16,8 to slurm to run 24 nodes job.

[terrykong]

I actually think we need to do this in both ray.sub and in nemo-rl source.

We need to change ray.sub so it can launch workers on different segments and give them a separate label.

Then, in nemo-rl source we need to select which worker we want to build the virtual cluster. Ray let's us create arbitrary resource labels and we can use that to figure out which segment to use.

Is it realistic to expect we'd maybe need more than 2 segments?

[guyueh1]

@terrykong I think the use case of uneven segment is like this: say you want M training nodes and N generation nodes, currently you can only choose a segment value that is divisible by M+N, and then pick M out of them to form the training sub-cluster. This could be suboptimal, compared to if you are only allocating M nodes in a pretrain job.

This might be not significant for users running at a large scale, like 1k gpus, but at a small scale to achieve both the best ratio of M:N and the best performance for training, it could be important.

Option 2 is much easier but the blocker is whether slurm can add support uneven segments, I am asking about it; Option 1 is more complicated, we can try prototyping it if Option 2 doesn't work.