pytorch - 💡(How to fix) Fix FSDP.optim_state_dict and FSDP.optim_state_dict_to_load called without group in torch.distributed.checkpoint.state_dict [1 comments, 2 participants]

TL;DR

Passing the custom process group to FSDP's optimizer state dict APIs is likely necessary to fix the issue with get_state_dict and set_state_dict in torch.distributed.checkpoint.state_dict.

Guidance

• Identify the custom process group used to initialize the FSDP model and pass it to FSDP.optim_state_dict and FSDP.optim_state_dict_to_load.
• Modify the call sites in torch/distributed/checkpoint/state_dict.py to forward the group parameter to FSDP's optimizer state dict APIs.
• Verify that the custom process group is correctly passed to the optimizer state dict APIs by checking the group argument in the FSDP.optim_state_dict and FSDP.optim_state_dict_to_load functions.
• Consider updating the torch.distributed.checkpoint.state_dict module to handle custom process groups by default.

Example

# Assuming `group` is the custom process group
state_dict = FSDP.optim_state_dict(group=group)

Notes

The fix requires modifying the torch.distributed.checkpoint.state_dict module, which may involve updating the PyTorch library. The solution may not apply to all versions of PyTorch, especially if the module has changed since the main branch.

Recommendation

Apply workaround: Modify the call sites in torch/distributed/checkpoint/state_dict.py to forward the group parameter to FSDP's optimizer state dict APIs, as this is a more targeted solution that does not require updating the entire library.