Lightning-AI / Lightning-AI/lightning-thunder
Composable FSDP Transform
- Dominant language
- Python
- Stars
- 1.5k
- Forks
- 121
- PR merge metrics
- No merged PRs in 30d
Description
Make FSDP composable with other transforms.
Desired features:
- allow FSDPTransform to compose with other transforms (e.g. some Quantization or QLORA),
- allow FSDPTransform to work with Materialization:
- sharding on meta,
- sharding of state dict,
- release memory after sharding if desired and called with a materialized model (in particular default for fsdp(jit(model))?), i.e. fix #478
- as a side condition `fsdp(jit(model))` must continue to work as a convenience method for practitioners (thanks @crcrpar for the observation that this should stay simple).
Steps include:
- [ ] move transformation of module to `FSDPTransform.transform_module` (PR #986 )
- [ ] do memory-releasing (PR #988 (also covering PR #932))
- [ ] implement `FSDPTransform.transform_state_dict_for_submodule` (PR #989 )
- [ ] delegate materialization (PR #995 )
For the side condition of keeping `fsdp(jit(model))` working as is, all of these should work while keeping the existing tests unchanged.
@crcrpar @IvanYashchuk (who asked about the plan)
cc @carmocca @crcrpar
Contributor guide
No contributing guide indexed for this repository
Assessment
This issue has not been assessed yet.