FSDP Support

[intelava]

🚀 Feature

Hi, I want to implement FSDP support for the pytorch/ignite. I propose integrating it into Ignite's auto_model. The checkpoint handlers will be updated to safely export sharded weights to CPU RAM as a consolidated FULL_STATE_DICT, preventing VRAM OOM issues while preserving backward compatibility with the existing single file DiskSaver logic. Do you have any other suggestions for this PR? I would like to get your feedback and implement this clean.

[vfdev-5]

@intelava thanks for the suggestion! Can you provide details on what you would like to add: how exactly would work auto_model and the checkpoint handler? Ideally, if you can make a code snippet of the new API and details on the implementation, it would be great, thanks!

[intelava]

Hi @vfdev-5, thanks for the quick response.

The main idea is to keep the API clean. We just add a use_fsdp flag to auto_model. As it changes the model wrapping, we pass the kwargs directly to FSDP. But one thing, we need to handle the state dicts safely to prevent VRAM OOM on large models.

Actually I implemented the FULL_STATE_DICT offloading to CPU on my local. We can argue about that. Here is the code changes:

1. ignite/distributed/auto.py

I added the use_fsdp flag and a fail-fast check for sync_bn.

def auto_model(model: nn.Module, sync_bn: bool = False, use_fsdp: bool = False, **kwargs: Any) -> nn.Module:
    # ...
    if use_fsdp and sync_bn:
        raise ValueError("Arguments use_fsdp and sync_bn are mutually exclusive. FSDP does not support SyncBatchNorm.")
    
    # ... inside the native backend block ...
            if use_fsdp:
                try:
                    from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
                except ImportError:
                    raise RuntimeError(
                        "FullyShardedDataParallel (FSDP) is not available. "
                        "Please upgrade to PyTorch >= 1.12."
                    )
                logger.info("Apply torch FullyShardedDataParallel on model")
                model = FSDP(model, **kwargs)

2. ignite/handlers/checkpoint.py

I intercepted the FSDP objects to use the context manager. This pulls a consolidated dict to CPU RAM.

# In _setup_checkpoint (Saving)
                elif HAVE_FSDP and isinstance(obj, FSDP):
                    cfg = FullStateDictConfig(offload_to_cpu=True, rank0_only=True)
                    with FSDP.state_dict_type(obj, StateDictType.FULL_STATE_DICT, cfg):
                        state_dict = obj.state_dict()
                    if idist.get_rank() != self.save_on_rank:
                        return {}
                    return state_dict

# In load_objects (Loading)
            elif HAVE_FSDP and isinstance(obj, FSDP):
                with FSDP.state_dict_type(obj, StateDictType.FULL_STATE_DICT):
                    obj.load_state_dict(chkpt_obj, **kwargs)
                return

If you want I can create a PR. Let me know your thoughts.

[vfdev-5]

@intelava it makes sense, thanks ! How about using FSDP2 API instead of the old one?
https://docs.pytorch.org/docs/stable/distributed.fsdp.fully_shard.html

[intelava]

@vfdev-5 That would be better actually. One quick architectural question regarding the FSDP2 integration before I draft the code:

• Should auto_model automatically construct a default 1D DeviceMesh using Ignite's idist process group, or should we keep it strictly unopinionated and require the user to pass their own DeviceMesh via kwargs? In my opinion, we should follow ignite's approach to make things easier, we can make it automatic.
• FSDP2 natively uses distributed checkpoints, which saves to a directory rather than a single file. Should we force consolidation to a single .pt file on Rank 0 to strictly protect backward compatibility with DiskSaver, or introduce directory based saving support? I lean towards keeping .pt for compatibility, but I am also aware that forcing consolidation defeats the memory saving benefits of DCP for massive models. How would you like to balance this?

Waiting for your feedback!

[vfdev-5]

@intelava good questions! We can keep the mesh as default and idist.auto_model kwargs can be passed into fully_shard so that user can pass their own mesh if needed.

FSDP2 natively uses distributed checkpoints, which saves to a directory rather than a single file. Should we force consolidation to a single .pt file on Rank 0 to strictly protect backward compatibility with DiskSaver, or introduce directory based saving support? I lean towards keeping .pt for compatibility, but I am also aware that forcing consolidation defeats the memory saving benefits of DCP for massive models. How would you like to balance this?

I think we can try to do directory based saving support.

[intelava]

@vfdev-5 May I create a PR for FSDP2 support now? If there is a requirement for a detailed plan first, I can also come up with that!

[vfdev-5]

@intelava you can go ahead and draft the PR and we can iterate over the draft directly

[intelava]

Hi @vfdev-5 , I'm sorry for the delay. I implemented FSDP2 as we discussed. This is a first version, and might need to iterate over it as you said. You can check the PR from here and we can continue the discussion over there!