FSDP training error: _dim_group_infos list index out of range

thanks for the great work.
I got an error when use 2 nodes with 2x8 GPUs to finetune the released model. Could you have a look?
On the second node:

> [rank15]: Traceback (most recent call last):
[rank15]: File "/opt/tiger/test_bagel/train/pretrain_unified_navit.py", line 876, in
[rank15]: main()
[rank15]: File "/opt/tiger/test_bagel/train/pretrain_unified_navit.py", line 565, in main
[rank15]: ema_model = fsdp_ema_setup(ema_model, fsdp_config)
[rank15]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank15]: File "/opt/tiger/test_bagel/train/fsdp_utils.py", line 252, in fsdp_ema_setup
[rank15]: ema_model = fsdp_wrapper(ema_model, fsdp_config, ignored_modules=ignored_modules)
[rank15]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank15]: File "/opt/tiger/test_bagel/train/fsdp_utils.py", line 58, in fsdp_wrapper
[rank15]: return FSDP(
[rank15]: ^^^^^
[rank15]: File "/usr/local/lib/python3.11/dist-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 454, in __init__
[rank15]: _init_process_group_state(
[rank15]: File "/usr/local/lib/python3.11/dist-packages/torch/distributed/fsdp/_init_utils.py", line 127, in _init_process_group_state
[rank15]: state = _init_process_group_state_for_hybrid_shard(
[rank15]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank15]: File "/usr/local/lib/python3.11/dist-packages/torch/distributed/fsdp/_init_utils.py", line 166, in _init_process_group_state_for_hybrid_shard
[rank15]: state._inter_node_pg = device_mesh.get_group(mesh_dim=0)
[rank15]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank15]: File "/usr/local/lib/python3.11/dist-packages/torch/distributed/device_mesh.py", line 721, in get_group
[rank15]: _find_pg_by_ranks_and_tag(*self._dim_group_infos[mesh_dim][:2]) # type: ignore[index]
[rank15]: ~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^
[rank15]: IndexError: list index out of range

Reproduce the reported finetuning setup with 2 nodes and 2x8 GPUs, then start at train/fsdp_utils.py:58 and train/fsdp_utils.py:252. Trace the FSDP setup into train/pretrain_unified_navit.py:565 and compare it with the PyTorch device-mesh failure at _dim_group_infos. Done means multi-node FSDP EMA setup completes without the IndexError.

Written by the indexing model from the issue text.