Does FSDP with Accelerate support async checkpoint saving (e.g. via torch.distributed.checkpoint)?

[LucienXian]

Hi, I'm using FSDP through Accelerate and would like to improve checkpoint saving performance. I noticed that PyTorch's torch.distributed.checkpoint (DCP) supports asynchronous saving with dcp.async_save, which can help overlap checkpoint I/O with training.

My question is: Does Accelerate's FSDP integration currently support async checkpointing via DCP or any similar mechanism?

If not directly exposed, is it safe to manually use torch.distributed.checkpoint.async_save with models and optimizers wrapped by Accelerate's FSDP? Are there any known compatibility issues or best practices (e.g., handling sharded state dict, optimizer state reconstruction)?

It would be great if Accelerate could expose async save as an option in save_state or provide utilities to integrate with DCP more seamlessly.

Thanks for your help!

[S1ro1]

Hi! It should be safe to use yourself, as for now we don't directly support it, currently it's my next thing to do. See at #3701

[LucienXian]

Hi @S1ro1

Thank you for your previous response! I’m excited about the upcoming async checkpointing support via DCP in Accelerate. To help with testing and adoption, I’d like to ask:

Merge Timeline: Do you have an estimated timeline for when this feature might be merged into the main branch?
Remaining Work: Are there specific technical hurdles, design decisions, or testing milestones blocking the merge?
If there are any open tasks where community contributions could accelerate progress (e.g., benchmarking, bug fixes, or docs), I’d be glad to assist. Please let me know how I can help.

Looking forward to updates, and thanks again for your work on this!

[S1ro1]

@LucienXian Sadly I'm OOO for a next week or so, so won't be able to finish it up. Will probably be merged on 1st half of September. As for remaining work, I believe the PR is mostly ready, will only need to add tests and figure out in which configuration we actually can support async checkpointing.