FSDP does not properly set "num_nodes" on SLURM

[weicao1990]

Bug description

Hi all. I trained my model on SLURM and it works well when using DDPStrategy. However, when using FSDPStrategy, it works well on a single node with 8 GPUs. But when I try to use 5 nodes, it reports

ValueError: Invalid rank 8, rank should be in the interval [0, 7]

for node 2~5.

I found such error is due to DistributedSampler

        if rank >= num_replicas or rank < 0:
            raise ValueError(
                "Invalid rank {}, rank should be in the interval"
                " [0, {}]".format(rank, num_replicas - 1))

where num_replicas = self.num_nodes * self.num_processes, defined in the strategy file. For DDPStrategy, self.num_nodes is properly set as 5, but for FSDPStrategy, self.num_nodes is 1, which makes "rank>num_nodes * num_processes".

What version are you seeing the problem on?

2.0+

How to reproduce the bug

No response

Error messages and logs

# Error messages and logs here please

Environment

Current environment

#- Lightning 2.0.0
#- PyTorch 2.0.0

More info

No response

cc @awaelchli @carmocca

[awaelchli]

Thanks for reporting.
This is a duplicate of #17028, right?
I'll take care of it

[awaelchli]

@weicao1990 until the fix is out, you can work around this issue by setting the num_nodes value yourself via
trainer.strategy.num_nodes = x. This should unblock you.

[weicao1990]

@awaelchli thanks a lot~