feat: dcp async checkpoint integration

[deependujha]

What does this PR do?

Fixes #<issue_number>

import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset
from lightning.fabric import Fabric
from lightning.fabric.utilities import AttributeDict
from lightning.fabric.plugins import DistributedAsyncCheckpointIO


# Simple model
class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.layer = nn.Linear(10, 2)

    def forward(self, x):
        return self.layer(x)


# Training loop
def train_one_epoch(fabric, model, optimizer, dataloader, loss_fn, epoch):
    model.train()
    for step, (x, y) in enumerate(dataloader):
        x, y = fabric.to_device((x, y))
        optimizer.zero_grad()
        preds = model(x)
        loss = loss_fn(preds, y)
        fabric.backward(loss)
        optimizer.step()

        # log step loss
        fabric.log("train_loss", loss.item(), step=epoch * len(dataloader) + step)
        print(f"Epoch {epoch} Step {step} Loss {loss.item():.4f}")


def main():
    # Dummy dataset
    x = torch.randn(100, 10)
    y = torch.randint(0, 2, (100,))
    dataset = TensorDataset(x, y)
    loader = DataLoader(dataset, batch_size=16)

    async_checkpoint_io = DistributedAsyncCheckpointIO()

    fabric = Fabric(plugins=[async_checkpoint_io])
    fabric.launch()

    # Setup
    model = SimpleModel()
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
    loss_fn = nn.CrossEntropyLoss()
    model, optimizer, loader = fabric.setup(model, optimizer, loader)

    # Run training
    for epoch in range(5):
        train_one_epoch(fabric, model, optimizer, loader, loss_fn, epoch)
        state = AttributeDict(model=model, optimizer=optimizer, iteration=epoch)
        fabric.save(f"checkpoint/checkpoint_epoch_{epoch}.pt", state)


if __name__ == "__main__":
    main()

Before submitting

• Was this discussed/agreed via a GitHub issue? (not for typos and docs)
• Did you read the contributor guideline, Pull Request section?
• Did you make sure your PR does only one thing, instead of bundling different changes together?
• Did you make sure to update the documentation with your changes? (if necessary)
• Did you write any new necessary tests? (not for typos and docs)
• Did you verify new and existing tests pass locally with your changes?
• Did you list all the breaking changes introduced by this pull request?
• Did you update the CHANGELOG? (not for typos, docs, test updates, or minor internal changes/refactors)

PR review

Anyone in the community is welcome to review the PR.
Before you start reviewing, make sure you have read the review guidelines. In short, see the following bullet-list:

Reviewer checklist

• Is this pull request ready for review? (if not, please submit in draft mode)
• Check that all items from Before submitting are resolved
• Make sure the title is self-explanatory and the description concisely explains the PR
• Add labels and milestones (and optionally projects) to the PR so it can be classified

---

📚 Documentation preview 📚: https://pytorch-lightning--21515.org.readthedocs.build/en/21515/

[deependujha]

next pr todos:

• use a queue to pipeline checkpoint saves.

self.checkpoint_future_queue: Optional[Future] = queue.Queue(maxsize=2)

• verify for remote checkpoint dir

[codecov]

Codecov Report

❌ Patch coverage is 81.01266% with 30 lines in your changes missing coverage. Please review.
✅ Project coverage is 87%. Comparing base (8d86b24) to head (30aaf36).

Additional details and impacted files

@@           Coverage Diff            @@
##           master   #21515    +/-   ##
========================================
- Coverage      87%      87%    -0%     
========================================
  Files         270      272     +2     
  Lines       24077    24218   +141     
========================================
+ Hits        20868    20976   +108     
- Misses       3209     3242    +33     

[tchaton]

Nice work !

[tchaton]

Maybe sharded_checkpoints = False might be more intuitive. But also, this would require to be configured by the strategy.

[tchaton]

Do we need to wait there, can't we have multiple future being created ?

[deependujha]

I didn't do this for the first PR. But added it as todo: #21515 (comment)

should i add this in this pr?