Lightning-AI / Lightning-AI/pytorch-lightning

XLA FSDP strategy has undocumented requirement for using activation checkpointing

Open
#19,919 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

bug fabric strategy: fsdp strategy: xla
Dominant language
Python
Stars
31.4k
Forks
3.8k
Avg merge
6d 7h
Merged PRs (30d)
6

Description

### Bug description

To use activation checkpointing with the [XLA FSDP Fabric strategy](https://lightning.ai/docs/fabric/stable/api/generated/lightning.fabric.strategies.XLAFSDPStrategy.html), the set of modules to checkpoint must also be wrapped by the `auto_wrap_policy`.

That is, if `auto_wrap_policy` is provided to the strategy as a set `W`, and `activation_checkpointing_policy` is passed as a set `C`, any members of `C` that are not also members of `W` will not actually be checkpointed. This is because [the implementation of the activation checkpointing policy](https://github.com/Lightning-AI/pytorch-lightning/blob/master/src/lightning/fabric/strategies/xla_fsdp.py#L652) adds "also wrap these classes in checkpoint_module" to the functionality of the wrapper callable, but the wrapper is only applied when the auto_wrap_policy tells it to.

This was surprising, since at least at first blush the auto_wrap_policy and activation_checkpointing_policy appear to be independent parameters. And the behavior of silent failure added confusion as to what was going on.

I suggest a few possible remedies, and these are not mutually exclusive.

1. Document this in the class docstring.
2. If auto_wrap_policy is provided as a set, activation_checkpointing_policy is not empty, and the former is not a superset of the latter, throw an error. If it's provided as a function, evaluate it for each member of activation_checkpointing_policy and if it's false for any, throw an error.
3. Union activation_checkpointing_policy with the auto_wrap_policy. Unlike 1 or 2, this would be a behavior change, though at least in my code it's what I ended up doing manually anyway.

### What version are you seeing the problem on?

v2.2

### How to reproduce the bug

_No response_

### Error messages and logs

```
# Error messages and logs here please
```

### Environment

Current environment

```
#- Lightning Component (e.g. Trainer, LightningModule, LightningApp, LightningWork, LightningFlow):
#- PyTorch Lightning Version (e.g., 1.5.0):
#- Lightning App Version (e.g., 0.5.2):
#- PyTorch Version (e.g., 2.0):
#- Python version (e.g., 3.9):
#- OS (e.g., Linux):
#- CUDA/cuDNN version:
#- GPU models and configuration:
#- How you installed Lightning(`conda`, `pip`, source):
#- Running environment of LightningApp (e.g. local, cloud):
```

### More info

_No response_

cc @justusschock @lantiga @JackCaoG @Liyang90 @gkroiz

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Read src/lightning/fabric/strategies/xla_fsdp.py around the activation checkpointing policy implementation near line 652, then inspect the XLAFSDPStrategy class docstring. Confirm how auto_wrap_policy and activation_checkpointing_policy interact, and define completion around either documenting or explicitly handling the reported silent behavior.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
distributed-systems, machine-learning
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.