Lightning-AI / Lightning-AI/pytorch-lightning

Support combinations of precision plugins

Open
#18,679 14 comments 3 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

design fabric feature pl plugin
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:

  1. The naive way: Create a TransformerEngineFSDPPrecision. This is simple and effective but it creates maintainability problem.
  2. The independent way: If there are no dependencies between the plugins, we could support plugins=[TransformerEnginePrecision(), FSDPPrecision()]. But there will likely be dependencies.
  3. 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 CheckpointIO plugins.
Alternatives

No response

Additional context

No response

cc @borda @tchaton @justusschock @awaelchli @carmocca

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

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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.