HSDP + set_optimizer_state_dict errors with monolithic checkpointing

[mvpatel2000]

🐛 Describe the bug

PyTorch 2.3.1 (and nightly)

Versions

TLDR: HSDP + DCP monolith checkpointing break.

When using DCP + set_optimizer_state_dict, the code eventually goes to _optim_state_dict_to_load_impl here, which calls _flatten_optim_state_dict and in turn _broadcast_processed_state defined here.

def _broadcast_processed_state(
    fsdp_state: _FSDPState,
    optim_state: Dict[str, Any],
    group: Optional[dist.ProcessGroup],
) -> Dict[str, Any]:
    objects: List[Any] = [None]
    if fsdp_state.rank == 0:
        objects[0] = tree_map_only(
            torch.Tensor,
            lambda v: v.cpu() if v.dim() == 0 else _PosDimTensorInfo(v.shape, v.dtype),  # type: ignore[union-attr]
            optim_state,
        )
    dist.broadcast_object_list(objects, src=0, group=group)
    if fsdp_state.rank == 0:
        return optim_state
    else:
        return objects[0]

This broadcast occurs across the specified group, which is None -> global processs group (as set_optimizer_state_dict has no way of specifying the group). However, when using HSDP, fsdp_state.rank is the local rank within the shard group and not the global rank. Thus, local rank 0 of each shard group will error out as it returns optim_state, which is set to None in _optim_state_dict_to_load_impl here.

        if rank0_only and dist.get_rank(group) > 0:
            optim_state_dict = {}

I think the correct solution is to check for rank in the same group as the broadcast but not 100% sure

cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @awgu @penguinwu @fegin @XilunWu @wanchaol @fduwjj @wz337 @tianyu-l @wconstab @yf225 @chauhang @d4l3k @LucasLLC @MeetVadakkanchery

[mvpatel2000]

@pytorchbot label "oncall: distributed"

[mvpatel2000]

I think this fixes it... #128446 but not 100% sure

[mvpatel2000]

would need help writing a unit test... not sure what to do / where to add

CC: @fegin @LucasLLC @wz337 maybe?

[weifengpy]

I am oncall this week. To @fegin : this is HSDP + DSD, both component are very relevant to you. WYT? I see you are tagged in solution PR as well #128446