[FEAT] Support PyTorch DCP for FSDP

[fishcrap]

Description

This PR adds support for PyTorch Distributed Checkpoint (DCP) format for FSDP models, enabling more efficient checkpointing and recovery for large-scale distributed training.

Key changes include:

• DCP Checkpoint Support in FSDPEngine: Added _save_to_dcp() and _load_from_dcp() methods that use PyTorch's distributed checkpoint API to save/load FSDP2 models along with optimizer states
• DCPState Wrapper: Created a new DCPState class in areal/utils/fsdp/checkpoint.py that implements the PyTorch Stateful protocol for seamless integration with DCP's save/load APIs
• Recovery Integration: Modified areal/utils/recover.py to use DCP format by default for checkpoint saving and loading during recovery
• Comprehensive Testing: Added distributed tests (test_fsdp_dcp.py and run_fsdp_dcp_distributed.py) that verify DCP functionality with multiple GPUs and various scenarios including optimizer state persistence
• Unit Test Coverage: Extended test_train_engine.py to include DCP save/load testing for all supported engines

Benefits:

• More efficient checkpoint I/O for large distributed models
• Native support for sharded state dict management with FSDP2
• Automatic handling of optimizer states across distributed ranks
• Compatible with PyTorch 2.4+ distributed checkpoint format

Related Issue

N/A

Type of Change

• Bug fix (non-breaking change that fixes an issue)
• New feature (non-breaking change that adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Documentation update
• Code refactoring (no functional changes)
• Performance improvement
• Test coverage improvement

Checklist

• I have read the Contributing Guide
• I have run formatting tools (pre-commit or manual)
• I have run relevant unit tests and they pass
• I have added tests for new functionality
• I have updated documentation if needed
• My branch is up to date with main
• This PR introduces breaking changes (if yes, fill out details below)
• If this PR changes documentation, I have built and previewed it locally with jb build docs
• No critical issues raised by AI reviewers (/gemini review)

Breaking Change Details (if applicable):

N/A

Additional Context

N/A

---

Need help? Check the Contributing Guide or ask in
GitHub Discussions!

[gemini-code-assist]

Summary of Changes

Hello @fishcrap, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request significantly enhances the checkpointing capabilities for FSDP models by integrating native support for PyTorch Distributed Checkpoint (DCP) format. This change enables more efficient and robust saving and loading of large-scale distributed models and their associated optimizer states, streamlining the training recovery process and ensuring compatibility with modern PyTorch distributed features.

Highlights

• DCP Integration in FSDPEngine: The FSDPEngine now fully supports PyTorch Distributed Checkpoint (DCP) format for saving and loading FSDP models, including their optimizer states. This replaces previous NotImplementedError placeholders.
• DCPState Wrapper: A new DCPState class has been introduced in areal/utils/fsdp/checkpoint.py. This class implements the Stateful protocol, allowing seamless integration with PyTorch's dcp.save and dcp.load APIs for managing FSDP model and optimizer states.
• Default DCP for Recovery: The recovery utility (areal/utils/recover.py) has been updated to use the DCP format by default for all checkpoint saving and loading operations, simplifying the recovery process for distributed training.
• Comprehensive Testing: New distributed tests (test_fsdp_dcp.py, run_fsdp_dcp_distributed.py) and extended unit tests (test_train_engine.py) have been added to thoroughly verify the functionality and correctness of DCP save/load for FSDP models across various scenarios, including optimizer state persistence.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces support for PyTorch Distributed Checkpoint (DCP) for FSDP models, which is a valuable feature for efficient checkpointing in large-scale distributed training. The implementation is well-designed, notably with the DCPState wrapper that cleanly integrates with PyTorch's Stateful protocol. The changes in FSDPEngine and RecoverHandler are logical and correctly enable DCP as the new default. The inclusion of comprehensive distributed tests is excellent for ensuring the correctness of this new functionality. I've identified a minor issue in the test implementation regarding hardcoded checkpoint paths that could lead to test flakiness, and I've provided a suggestion to improve its robustness. Overall, this is a solid contribution.