Just to try understanding your situation better. Let me know if the following description is correct:
You’re using 2 nodes, each with 4 GPUs. You want to use 2 GPUs on each node, which means your intended world size is 4.
The global rank of processes on node 1 are {0, 1}, and the global ranks of processes on node 2 are {2, 3}.
To achieve this, you can use CUDA_VISIBLE_DEVICES before launching your training script. For example, if you set CUDA_VISIBLE_DEVICES=1,2, the training script will not see the rest of the GPUs. Then you can simply initialize the process groups on each process with the correct rank passed in (and no need to do torch.cuda.set_device). This will ensure the correct GPUs are used for the training processes without any manual configurations.