The Deep Learning framework to train, deploy, and ship AI products Lightning fast.

NEW- Lightning 2.0 is featuring a clean and stable API!!

---

Lightning.ai • PyTorch Lightning • Fabric • Lightning Apps • Docs • Community • Contribute •

Simple installation from PyPI

pip install lightning

pip install lightning['extra']

conda install lightning -c conda-forge

pip install https://github.com/Lightning-AI/lightning/archive/refs/heads/release/stable.zip -U

pip install https://github.com/Lightning-AI/lightning/archive/refs/heads/master.zip -U

pip install -iU https://test.pypi.org/simple/ pytorch-lightning

• [PyTorch Lightning: Train and deploy PyTorch at scale](#PyTorch Lightning: Train and deploy PyTorch at scale).
• Lightning Fabric: Expert-level control over scale and training loop.
• Lightning Apps: Build AI products and ML workflows.

PyTorch Lightning is just organized PyTorch - Lightning disentangles PyTorch code to decouple the science from the engineering.

# main.py # ! pip install torchvision import os, torch, torch.nn as nn, torch.utils.data as data, torchvision as tv, torch.nn.functional as F import lightning as L # -------------------------------- # Step 1: Define a LightningModule  # -------------------------------- # A LightningModule (nn.Module subclass) defines a full *system*  # (ie: an LLM, difussion model, autoencoder, or simple image classifier). class LitAutoEncoder(L.LightningModule): def __init__(self): super().__init__() self.encoder = nn.Sequential(nn.Linear(28 * 28, 128), nn.ReLU(), nn.Linear(128, 3)) self.decoder = nn.Sequential(nn.Linear(3, 128), nn.ReLU(), nn.Linear(128, 28 * 28)) def forward(self, x): # in lightning, forward defines the prediction/inference actions embedding = self.encoder(x) return embedding def training_step(self, batch, batch_idx): # training_step defines the train loop. It is independent of forward x, y = batch x = x.view(x.size(0), -1) z = self.encoder(x) x_hat = self.decoder(z) loss = F.mse_loss(x_hat, x) self.log("train_loss", loss) return loss def configure_optimizers(self): optimizer = torch.optim.Adam(self.parameters(), lr=1e-3) return optimizer # ------------------- # Step 2: Define data # ------------------- dataset = tv.datasets.MNIST(os.getcwd(), download=True, transform=tv.transforms.ToTensor()) train, val = data.random_split(dataset, [55000, 5000]) # ------------------- # Step 3: Train # ------------------- autoencoder = LitAutoEncoder() trainer = L.Trainer() trainer.fit(autoencoder, data.DataLoader(train), data.DataLoader(val))

Run the model on your terminal

pip install torchvision python main.py

Lightning has over 40+ advanced features designed for professional AI research at scale.

Here are some examples:

# 8 GPUs # no code changes needed trainer = Trainer(accelerator="gpu", devices=8) # 256 GPUs trainer = Trainer(accelerator="gpu", devices=8, num_nodes=32)

# no code changes needed trainer = Trainer(accelerator="tpu", devices=8)

# no code changes needed trainer = Trainer(precision=16)

from lightning import loggers # tensorboard trainer = Trainer(logger=TensorBoardLogger("logs/")) # weights and biases trainer = Trainer(logger=loggers.WandbLogger()) # comet trainer = Trainer(logger=loggers.CometLogger()) # mlflow trainer = Trainer(logger=loggers.MLFlowLogger()) # neptune trainer = Trainer(logger=loggers.NeptuneLogger()) # ... and dozens more

es = EarlyStopping(monitor="val_loss") trainer = Trainer(callbacks=[es])

checkpointing = ModelCheckpoint(monitor="val_loss") trainer = Trainer(callbacks=[checkpointing])

# torchscript autoencoder = LitAutoEncoder() torch.jit.save(autoencoder.to_torchscript(), "model.pt")

# onnx with tempfile.NamedTemporaryFile(suffix=".onnx", delete=False) as tmpfile: autoencoder = LitAutoEncoder() input_sample = torch.randn((1, 64)) autoencoder.to_onnx(tmpfile.name, input_sample, export_params=True) os.path.isfile(tmpfile.name)

• Models become hardware agnostic
• Code is clear to read because engineering code is abstracted away
• Easier to reproduce
• Make fewer mistakes because lightning handles the tricky engineering
• Keeps all the flexibility (LightningModules are still PyTorch modules), but removes a ton of boilerplate
• Lightning has dozens of integrations with popular machine learning tools.
• Tested rigorously with every new PR. We test every combination of PyTorch and Python supported versions, every OS, multi GPUs and even TPUs.
• Minimal running speed overhead (about 300 ms per epoch compared with pure PyTorch).

Run on any device with Fabric with expert-level PyTorch control for scaling foundation models or writing your own Trainer.

Fabric is designed for the most complex models like foundation model scaling, LLMs, diffussion, transformers, reinforcement learning, active learning.

+ import lightning as L import torch import torch.nn as nn from torch.utils.data import DataLoader, Dataset class PyTorchModel(nn.Module): ... class PyTorchDataset(Dataset): ... + fabric = L.Fabric(accelerator="cuda", devices=8, strategy="ddp") + fabric.launch() - device = "cuda" if torch.cuda.is_available() else "cpu model = PyTorchModel(...) optimizer = torch.optim.SGD(model.parameters()) + model, optimizer = fabric.setup(model, optimizer) dataloader = DataLoader(PyTorchDataset(...), ...) + dataloader = fabric.setup_dataloaders(dataloader) model.train() for epoch in range(num_epochs): for batch in dataloader: input, target = batch - input, target = input.to(device), target.to(device) optimizer.zero_grad() output = model(input) loss = loss_fn(output, target) - loss.backward() + fabric.backward(loss) optimizer.step() lr_scheduler.step()

In the spectrum from raw PyTorch to fully managed PyTorch Lightning, Fabric gives you full control over how much abstraction you want to add.

• Easily switch from running on CPU to GPU (Apple Silicon, CUDA, …), TPU, multi-GPU or even multi-node training
• Use state-of-the-art distributed training strategies (DDP, FSDP, DeepSpeed) and mixed precision out of the box
• All the device logic boilerplate is handled for you
• Designed with multi-billion parameter models in mind
• Build your own custom Trainer using Fabric primitives for training checkpointing, logging, and more

Once you're done building models, publish a paper demo or build a full production end-to-end ML system with Lightning Apps. Lightning Apps remove the cloud infrastructure boilerplate so you can focus on solving the research or business problems. Lightning Apps can run on the Lightning Cloud, your own cluster or a private cloud.

Browse available Lightning apps here

• CPC transforms
• Moco v2 tranforms
• SimCLR transforms

• GPT-2
• UNet

• DQN Loss
• Double DQN Loss
• Per DQN Loss

• Basic GAN
• DCGAN

• Logistic Regression
• Linear Regression

Lightning is rigorously tested across multiple CPUs, GPUs, TPUs, IPUs, and HPUs and against major Python and PyTorch versions.

Current build statuses

System / PyTorch ver.	1.11	1.12	1.13	2.0
Linux py3.9 [GPUs]	-			Soon
Linux py3.9 [TPUs]	-			Soon
Linux py3.8 [IPUs]	-	-		Soon
Linux py3.8 [HPUs]	-	-		Soon
Linux (multiple Python versions)				Soon
OSX (multiple Python versions)				Soon
Windows (multiple Python versions)				Soon

The lightning community is maintained by

• 10+ core contributors who are all a mix of professional engineers, Research Scientists, and Ph.D. students from top AI labs.
• 590+ active community contributors.

Want to help us build Lightning and reduce boilerplate for thousands of researchers? Learn how to make your first contribution here

Lightning is also part of the PyTorch ecosystem which requires projects to have solid testing, documentation and support.

If you have any questions please:

1. Read the docs.
2. Search through existing Discussions, or add a new question
3. Join our discord.