Lightning-AI / Lightning-AI/lightning-thunder

Composable FSDP Transform

Open
#987 0 comments 2 reactions 1 assignee Claimed by @t-vi View on GitHub
distributed enhancement transforms
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.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.