Lightning-AI / Lightning-AI/pytorch-lightning
Support combinations of precision plugins
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 31.4k
- Forks
- 3.8k
- Avg merge
- 6d 7h
- Merged PRs (30d)
- 6
Description
Description & Motivation
Both the Fabric and Trainer strategies are designed to have a single plugin enabled from the beginning to the end of the program.
This has been fine historically, however, some strategies require tailored plugin implementations that are functionally equal to other plugins.
For instance, single device training with bf16-true precision will use the HalfPrecision plugin
But FSDP training with bf16-true precision will use the FSDPPrecision plugin
Recent cutting-edge plugins such as TransformerEnginePrecision and BitsandbytesPrecision also implement the basic bf16-true functionality (example). So there's a lot of overlap.
The challenge then becomes: how do I enable TransformerEnginePrecision to work with FSDPStrategy if FSDPStrategy is designed to work with FSDPPrecision only?
Note that I'm using these specific classes to prove the point, but the design issue applies to any strategy that requires a specific plugin class. DeepSpeedStrategy and XLAStrategy would also be examples of this.
Pitch
Continuing the example, there's 3 ways this could be solved:
- The naive way: Create a
TransformerEngineFSDPPrecision. This is simple and effective but it creates maintainability problem. - The independent way: If there are no dependencies between the plugins, we could support
plugins=[TransformerEnginePrecision(), FSDPPrecision()]. But there will likely be dependencies. - The smart way: Create an abstraction that is able to compose two (or more?) plugins together and is itself a plugin. There's some precedent for this with the
CheckpointIOplugins.
Alternatives
No response
Additional context
No response
cc @borda @tchaton @justusschock @awaelchli @carmocca
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start by comparing the plugin implementations in src/lightning/fabric/plugins/precision/half.py, fsdp.py, and transformer_engine.py, then trace how FSDPStrategy, DeepSpeedStrategy, and XLAStrategy select or validate precision plugins. The work is done when a documented composition design supports combinations such as TransformerEnginePrecision with FSDPStrategy without creating every pair-specific plugin.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- distributed-systems, machine-learning
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 25/100