N-Bit Precision (Expert) — PyTorch Lightning 2.5.1.post0 documentation (original) (raw)

Audience: Researchers looking to integrate their new precision techniques into Lightning.

Precision Plugins¶

You can also customize and pass your own Precision Plugin by subclassing the Precision class.

• Perform pre and post backward/optimizer step operations such as scaling gradients.
• Provide context managers for forward, training_step, etc.

class CustomPrecision(Precision): precision = "16-mixed"

...

trainer = Trainer(plugins=[CustomPrecision()])