FSDP does not work on GLOO backend

[rohan-varma]

🐛 Describe the bug

It uses _allgather_base, but there is no support for this in Gloo backend:

RuntimeError: no support for _allgather_base in Gloo process group

Versions

main

cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @awgu @penguinwu @fegin @pietern @rohan-varma @SciPioneer

[abacaj]

Hi was there an update on this?

[NouamaneTazi]

You can use dist.all_gather for gloo

    dist.all_gather_into_tensor(out, out_sharded, group=group) # doesn't work for gloo
    dist.all_gather(list(out.chunk(2, dim=1)), out_sharded, group=group) # works for gloo

[usryokousha]

Are you suggesting modifying the code in all_gather_into_tensor when the backend is gloo?
(This is referring to the same code snippet as posted by @abacaj )

else:
   if group._get_backend_name() == Backend.GLOO:
      return all_gather([output_tensor], input_tensor, group=group)
   else:
      work = group._allgather_base(output_tensor, input_tensor)

This seems to work in my initial tests ... not sure if the computation is correct

[usryokousha]

This works for training but I encounter another issue while attempting to save i.e. model.state_dict(). Not sure how to handle this.

[NouamaneTazi]

What's the error message?

[usryokousha]

I make a call to model.state_dict() and the resulting import parts of the stack trace are here:
File "C:\Users\matt.smith\AppData\Local\miniconda3\envs\pytorch\lib\site-packages\torch\nn\modules\module.py", line 1818, in state_dict
module.state_dict(destination=destination, prefix=prefix + name + '.', keep_vars=keep_vars)
File "C:\Users\matt.smith\AppData\Local\miniconda3\envs\pytorch\lib\site-packages\torch\nn\modules\module.py", line 1818, in state_dict
module.state_dict(destination=destination, prefix=prefix + name + '.', keep_vars=keep_vars)
File "C:\Users\matt.smith\AppData\Local\miniconda3\envs\pytorch\lib\site-packages\torch\nn\modules\module.py", line 1815, in state_dict
self._save_to_state_dict(destination, prefix, keep_vars)
File "C:\Users\matt.smith\AppData\Local\miniconda3\envs\pytorch\lib\site-packages\torch\nn\modules\module.py", line 1722, in _save_to_state_dict
hook(self, prefix, keep_vars)
File "C:\Users\matt.smith\AppData\Local\miniconda3\envs\pytorch\lib\site-packages\torch\utils_contextlib.py", line 115, in decorate_context
return func(*args, **kwargs)

It seems that after the all_gather call the parameters are returned as a list and _save_to_state_dict calls a function which expects them as one large tensor. This function appears to be _local_pre_state_dict_hook in torch.distributed.fsdp._state_dict_utils.

[usryokousha]

I want to note that my tests were done on a system with a world_size of 1 for debugging purposes (ShardingPolicy.NO_SHARD). In spite of this all the FSDP objects seem to be used so it should work the same when using multiple GPUs. I commented out this check and it seems to work.

def _local_pre_state_dict_hook(
    fsdp_state: _FSDPState,
    module: nn.Module,
    *args,
    **kwargs,
) -> None:
    """
    Hook that runs before model.state_dict() is called. Right now, pre-state_dict
    hook is not supported by the PyTorch core. So this API is called from
    `_local_post_state_dict_hook()` to simulate the case.
    """
    #if (
    #    _has_fsdp_params(fsdp_state, module)
    #    and not _module_handles(fsdp_state, module)[0].uses_sharded_strategy
    #):
    #    raise RuntimeError(
    #        "``local_state_dict`` can only be used when parameters are flatten "
    #        "and sharded."
    #    )
    _common_pre_state_dict_hook(module, fsdp_state)

[ertza]

Hi, adding the @usryokousha's snippet seemed to make the initial error go away, but I run into another error where Gloo says output tensor length is not equal to number of nodes, I am using 6 nodes for the FSDP and this is the error:

File "/root/resnet-fsdp.py", line 248, in train
    outputs = model(input_ids, attention_mask=attention_mask,
  File "/opt/conda/envs/ultima/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1129, in _call_impl
    return forward_call(*input, **kwargs)
  File "/opt/conda/envs/ultima/lib/python3.9/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 1589, in forward
    self._rebuild_full_params()
  File "/opt/conda/envs/ultima/lib/python3.9/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
    return func(*args, **kwargs)
  File "/opt/conda/envs/ultima/lib/python3.9/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 2402, in _rebuild_full_params
    dist._all_gather_base(
  File "/opt/conda/envs/ultima/lib/python3.9/site-packages/torch/distributed/distributed_c10d.py", line 2140, in _all_gather_base
    return all_gather([output_tensor], input_tensor, group=group)
  File "/opt/conda/envs/ultima/lib/python3.9/site-packages/torch/distributed/distributed_c10d.py", line 2070, in all_gather
    work = group.allgather([tensor_list], [tensor])
RuntimeError: ProcessGroupGloo::allgather: invalid output tensor list at index 0 (expected length 6, got 1)

[awgu]

@ertza It looks like you are using an old version of PyTorch since I see _rebuild_full_params. I wonder if you could try with a newer version?

[sadimanna]

I want to note that my tests were done on a system with a world_size of 1 for debugging purposes (ShardingPolicy.NO_SHARD). In spite of this all the FSDP objects seem to be used so it should work the same when using multiple GPUs. I commented out this check and it seems to work.

def _local_pre_state_dict_hook(
    fsdp_state: _FSDPState,
    module: nn.Module,
    *args,
    **kwargs,
) -> None:
    """
    Hook that runs before model.state_dict() is called. Right now, pre-state_dict
    hook is not supported by the PyTorch core. So this API is called from
    `_local_post_state_dict_hook()` to simulate the case.
    """
    #if (
    #    _has_fsdp_params(fsdp_state, module)
    #    and not _module_handles(fsdp_state, module)[0].uses_sharded_strategy
    #):
    #    raise RuntimeError(
    #        "``local_state_dict`` can only be used when parameters are flatten "
    #        "and sharded."
    #    )
    _common_pre_state_dict_hook(module, fsdp_state)

Is there any workaround for this? That is, without altering the code in _local_pre_state_dict_hook?