NewYour coding agent can read the release notes before it upgrades.Set up the MCP server →
PyPI · #1157 most downloaded on PyPI
PyTorch Lightning is the lightweight PyTorch wrapper for ML researchers. Scale your models. Write less boilerplate.
Last release 25 days ago
10 Sep 2026
Ships fairly regularly
a new release about every 2 months
Nearly every release is documented
notes for 55 of the last 60 stable releases
Nothing withdrawn
no release was ever pulled
8 years old
223 releases · first in 2019
Fixed arbitrary code execution in load_from_checkpoint by restricting the _instantiator hyperparameter to an allowlist of trusted instantiators
2.6.6Fixed arbitrary code execution in load_from_checkpoint by restricting the _instantiator hyperparameter to an allowlist of trusted instantiators (#21832)
Fixed arbitrary code execution in load_from_checkpoint by rejecting a checkpoint _class_path that does not resolve to an already imported subclass of the loaded class (#21914)
Full commit list: 2.6.5 -> 2.6.6
We thank all folks who submitted issues, features, fixes and doc changes. It's the only way we can collectively make Lightning ⚡ better for everyone, nice job!
In particular, we would like to thank the authors and co-authors of the pull-requests above:
One column per quarter.
Use fs.pipe() for S3/GCS checkpoint uploads in _atomic_save ( #21595
2.6.5Full commit list: 2.6.4 -> 2.6.5
New Contributors
We thank all folks who submitted issues, features, fixes and doc changes. It's the only way we can collectively make Lightning ⚡ better for everyone, nice job!
In particular, we would like to thank the authors and co-authors of the pull-requests above:
Removed support for Neptune logger ( #21572 ).
2.6.4ChangedNote: We usually don't remove features in a patch release, however in this case it's an exception since even without removing it the integration would be broken due to Neptune being acquired and sunsetting public service
LitLogger version to 2026-03-17 (#21591)val_check_interval raising ValueError when limit_val_batches=0 and interval exceeds training batches (#21560)bf16-mixed, 16-mixed) initializing model parameters in half precision instead of fp32 (#21586)device_mesh type hint in FSDPStrategy to accept a 2-element tuple via the CLI (#21581)RichModelSummary model size display formatting (#21467)SimpleProfiler duration aggregation by using math.fsum (#21525)bf16-mixed, 16-mixed) initializing model parameters in half precision instead of fp32 (#21586)device_mesh type hint in FSDPStrategy to accept a 2-element tuple via the CLI (#21581)Full commit list: 2.6.1 -> 2.6.4
New Contributors
We thank all folks who submitted issues, features, fixes and doc changes. It's the only way we can collectively make Lightning ⚡ better for everyone, nice job!
In particular, we would like to thank the authors and co-authors of the pull-requests above:
@AsherJingkongChen, @Borda, @bhimrazy, @deependujha, @justusschock, @littlebullGit, @ManasVardhan, @taha-yassine
Added litlogger integration( #21430 ) Deprecated
2.6.1LightningModule.freeze() and LightningModule.unfreeze() by returning self (#21469)to_torchscript method due to deprecation of TorchScript in PyTorch (#21397)save_hyperparameters(ignore=...) behavior so subclass ignore rules override base class rules (#21490)LightningDataModule.load_from_checkpoint to restore the datamodule subclass and hyperparameters (#21478)ModelParallelStrategy single-file checkpointing when torch.compile wraps the model so optimizer states no longer raise KeyError during save (#21357)StochasticWeightAveraging with infinite epochs (#21396)_generate_seed_sequence_sampling function not producing unique seeds (#21399)ThroughputMonitor callback emitting warnings too frequently (#21453)weights_only argument for loading checkpoints in Fabric.load() and Fabric.load_raw() (#21470)DistributedSamplerWrapper not forwarding set_epoch to the underlying sampler (#21454)Full commit list: 2.6.0 -> 2.6.1
New Contributors
We thank all folks who submitted issues, features, fixes and doc changes. It's the only way we can collectively make Lightning ⚡ better for everyone, nice job!
In particular, we would like to thank the authors of the pull-requests above
Added WeightAveraging callback that wraps the PyTorch AveragedModel class
2.6.0WeightAveraging callback that wraps the PyTorch AveragedModel class (#20545)LightningModule (#20808)val_check_interval (#21071)EarlyStopping callback (#21188)ThroughputMonitor (#20236)EMAWeightAveraging callback that wraps Lightning's WeightAveraging class (#21260)weights_only argument for Trainer.{fit,validate,test,predict} and let torch handle default value (#21072)RichProgressBar and RichModelSummary if the rich package is available. Fallback to TQDMProgressBar and ModelSummary otherwise (#20896)max_trials is reached in Tuner.scale_batch_size (#21187)LightningCLI could not be initialized with trainer_default containing callbacks (#21192)ModelPruning is applied with lottery ticket hypothesis (#21191)save_last='link' and save_top_k=-1 (#21186)last.ckpt being created and not linked to another checkpoint (#21244)BackboneFinetuning from being used together with LearningRateFinder (#21224)ModelPruning sparsity logging bug that caused incorrect sparsity percentages (#21223)LightningCLI loading of hyperparameters from ckpt_path failing for subclass model mode (#21246)__init__ method (#21227)ThroughputMonitor calculated training time (#21291)DDPStrategy(static_graph=True) (#21251)weights_only argument for Trainer.{fit,validate,test,predict} and let torch handle default value (#21072)_DeviceDtypeModuleMixin._device from torch's default device function (#21164)Fabric.call to support different callback method signatures (#21258)on_train_batch_start returns -1 (#21296).Full commit list: 2.5.0 -> 2.6.0
We thank all folks who submitted issues, features, fixes and doc changes. It's the only way we can collectively make Lightning ⚡ better for everyone, nice job!
In particular, we would like to thank the authors of the pull-requests above
Nothing published for this version
Remove support for deprecated and archived lightning-habana package
Include exclude_frozen_parameters to DeepSpeedStrategy
2.5.5exclude_frozen_parameters to DeepSpeedStrategy (#21060)PossibleUserWarning that is raised if modules are in eval mode when training starts (#21146)LightningCLI not using ckpt_path hyperparameters to instantiate classes (#21116)ModelCheckpoint saves until validation metrics are available (#21106)TQDMProgressBar not resetting correctly when using both a finite and iterable dataloader (#21147)Tuner on crashes (#21162)exclude_frozen_parameters to DeepSpeedStrategy (#21060)_get_default_process_group_backend_for_device support more hardware platforms (verbose=False in seed_everything when no seed is provided (#21161)Full commit list: 2.5.4 -> 2.5.5
We thank all folks who submitted issues, features, fixes and doc changes. It's the only way we can collectively make Lightning ⚡ better for everyone, nice job!
In particular, we would like to thank the authors of the pull-requests above, in no particular order:
@Borda, @KAVYANSHTYAGI, @littlebullGit, @mauvilsa, @SkafteNicki, @taozhiwei
Thank you ❤️ and we hope you'll keep them coming!
Fixed AsyncCheckpointIO snapshots tensors to avoid race with parameter mutation
2.5.4AsyncCheckpointIO snapshots tensors to avoid race with parameter mutation (#21079)AsyncCheckpointIO threadpool exception if calling fit or validate more than one (#20952)LearningRateFinder callback (#21068)DeepSpeedstrategy (#21100)RichProgressBar crashing when sanity checking using val dataloader with 0 len (#21108)get_available_flops (#20913)Full commit list: 2.5.3 -> 2.5.4
We thank all folks who submitted issues, features, fixes and doc changes. It's the only way we can collectively make Lightning ⚡ better for everyone, nice job!
In particular, we would like to thank the authors of the pull-requests above, in no particular order:
@fnhirwa, @GdoongMathew, @jjh42, @littlebullGit, @SkafteNicki
Thank you ❤️ and we hope you'll keep them coming!
Fixed XLA strategy to add support for global_ordinal , local_ordinal , world_size which came instead of deprecated methods
save_on_exception option to ModelCheckpoint Callback (#20916)dataloader_idx_ in log names when add_dataloader_idx=False (#20987)ONNXProgram when calling to_onnx(dynamo=True) (#20811)training_step when using manual optimization (#21011)14.1+ (#21016)AdvancedProfiler to handle nested profiling actions for Python 3.12+ (#20809)rich progress bar error when resume training (#21000)ModelSummary (#21034)RichProgressBar being updated according to user provided refresh_rate (#21032)save_last behavior in the absence of validation (#20960)LearningRateFinder and EarlyStopping (#21056)lr_finder for mode="exponential" (#21055)save_hyperparameters crashing with dataclasses using init=False fields (#21051)devices and accelerator as CLI arguments (#20913)out-of-bounds or cannot be cast to int (#21029)name parameter in accelerator registry decorator (#20975)global_ordinal, local_ordinal, world_size which came instead of deprecated methods (#20852)Full commit list: 2.5.2 -> 2.5.3
We thank all folks who submitted issues, features, fixes and doc changes. It's the only way we can collectively make Lightning ⚡ better for everyone, nice job!
In particular, we would like to thank the authors of the pull-requests above, in no particular order:
@baskrahmer, @bhimrazy, @deependujha, @fnhirwa, @GdoongMathew, @jonathanking, @relativityhd, @rittik9, @SkafteNicki, @sudiptob2, @vsey, @YgLK
Thank you ❤️ and we hope you'll keep them coming!
CLI: resolve jsonargparse deprecation warning
toggled_optimizer(optimizer) method to the LightningModule, which is a context manager version of toggle_optimize and untoggle_optimizer (#20771)fsspec>=2025.5.0 if unavailable (#20780)nn.Parameter type for pruning sanitization (#20783)save_hyperparameters not working correctly with LightningCLI when there are parsing links applied on instantiation (#20777)logger_connector has an edge case where step can be a float (#20692)check_inputs to the target device if available during to_torchscript (#20873)max_steps during training (#20869)jsonnet (#20899)TransformerEnginePrecision conversion for layers with bias=False (#20805)Full commit list: 2.5.1 -> 2.5.2
We thank all folks who submitted issues, features, fixes, and doc changes. It's the only way we can collectively make Lightning ⚡ better for everyone, nice job!
In particular, we would like to thank the authors of the pull-requests above, in no particular order:
@adamjstewart, @Armannas, @bandpooja, @Borda, @chanokin, @duydl, @GdoongMathew, @KAVYANSHTYAGI, @mauvilsa, @muthissar, @rustamzh, @siemdejong
Thank you ❤️ and we hope you'll keep them coming!
TransformerEnginePrecision conversion for layers with bias=False (#20805)Full Changelog : 2.5.1...2.5.1.post0
Full Changelog: 2.5.1...2.5.1.post0
Full Changelog: https://github.com/Lightning-AI/pytorch-lightning/compare/2.5.1...2.5.1.post0
Allow LightningCLI to use a customized argument parser class
<a name="changelog-pytorch"></a>
<details open><summary>Changed</summary>
wandb default x-axis to tensorboard's global_step when sync_tensorboard=True (#20611)checkpoint_path_prefix parameter to the MLflow logger which can control the path to where the MLflow artifacts for the model checkpoints are stored (#20538)torch 2.6 (#20509)</details>
<details open><summary>Fixed</summary>
_R_co and _P to prevent type error (#20508)WandbLogger.experiment first in _call_setup_hook to ensure tensorboard logs can sync to wandb (#20610)should_stop when fit is called (#19177)WandbLogger upload models from all ModelCheckpoint callbacks, not just one (#20191)</details>
<a name="changelog-fabric"></a>
<details open><summary>Changed</summary>
torch 2.6 (#20509)</details>
<details open><summary>Removed</summary>
lightning run model; use fabric run instead. (#20588)</details>
</br>
Full commit list: 2.5.0 -> 2.5.1
<a name="contributors"></a>
We thank all folks who submitted issues, features, fixes and doc changes. It's the only way we can collectively make Lightning :zap: better for everyone, nice job!
In particular, we would like to thank the authors of the pull-requests above, in no particular order:
@benglewis, @Borda, @cgebbe, @duydl, @haifeng-jin, @japdubengsub, @justusschock, @lantiga, @mauvilsa, @millskyle, @ringohoffman, @ryan597, @senarvi, @TresYap
Thank you :heart: and we hope you'll keep them coming!
lightning run model; use fabric run instead (#20588)Nothing published for this version
Nothing published for this version
Nothing published for this version
Full Changelog: https://github.com/Lightning-AI/pytorch-lightning/compare/2.5.0...2.5.0.post0
Full Changelog: https://github.com/Lightning-AI/pytorch-lightning/compare/2.5.0...2.5.0.post0
Updated checkpointing documentation to mark resume_from_checkpoint as deprecated
Lightning AI :zap: is excited to announce the release of Lightning 2.5.
Lightning 2.5 comes with improvements on several fronts, with zero API changes. Our users love it stable, we keep it stable :smile:.
Talking about love :heart:, the lightning, pytorch-lightning and lightning-fabric packages are collectively getting more than 10M downloads per month :open_mouth:, for a total of over 180M downloads :exploding_head: since the early days . It's incredible to see PyTorch Lightning getting such a strong adoption across the industry and the sciences.
Release 2.5 embraces PyTorch 2.5, and it marks some of its more recent directions as officially supported, namely tensor subclass-based APIs like Distributed Tensors and TorchAO, in combination with torch.compile.
Here's a couple of examples:
<details><summary>Distributed FP8 transformer with PyTorch Lightning</summary>
Full example here
import lightning as L
import torch
import torch.nn as nn
import torch.nn.functional as F
from lightning.pytorch.demos import Transformer, WikiText2
from lightning.pytorch.strategies import ModelParallelStrategy
from torch.distributed._composable.fsdp.fully_shard import fully_shard
from torch.utils.data import DataLoader
from torchao.float8 import Float8LinearConfig, convert_to_float8_training
class LanguageModel(L.LightningModule):
def __init__(self, vocab_size):
super().__init__()
self.vocab_size = vocab_size
self.model = None
def configure_model(self):
if self.model is not None:
return
with torch.device("meta"):
model = Transformer(
vocab_size=self.vocab_size,
nlayers=16,
nhid=4096,
ninp=1024,
nhead=32,
)
float8_config = Float8LinearConfig(
# pip install -U --index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/Triton-Nightly/pypi/simple/ triton-nightly # noqa
pad_inner_dim=True,
)
def module_filter_fn(mod: torch.nn.Module, fqn: str):
# we skip the decoder because it typically vocabulary size
# is not divisible by 16 as required by float8
return fqn != "decoder"
convert_to_float8_training(model, config=float8_config, module_filter_fn=module_filter_fn)
for module in model.modules():
if isinstance(module, (nn.TransformerEncoderLayer, nn.TransformerDecoderLayer)):
fully_shard(module, mesh=self.device_mesh)
fully_shard(model, mesh=self.device_mesh)
self.model = torch.compile(model)
def training_step(self, batch):
input, target = batch
output = self.model(input, target)
loss = F.nll_loss(output, target.view(-1))
self.log("train_loss", loss, prog_bar=True)
return loss
def configure_optimizers(self):
return torch.optim.Adam(self.parameters(), lr=1e-4)
def train():
L.seed_everything(42)
dataset = WikiText2()
train_dataloader = DataLoader(dataset, num_workers=8, batch_size=1)
model = LanguageModel(vocab_size=dataset.vocab_size)
mp_strategy = ModelParallelStrategy(
data_parallel_size=4,
tensor_parallel_size=1,
)
trainer = L.Trainer(strategy=mp_strategy, max_steps=100, precision="bf16-true", accumulate_grad_batches=8)
trainer.fit(model, train_dataloader)
trainer.print(torch.cuda.memory_summary())
if __name__ == "__main__":
torch.set_float32_matmul_precision("high")
train()
</details>
<details><summary>Distributed FP8 transformer with Fabric</summary>
Full example here
import lightning as L
import torch
import torch.nn as nn
import torch.nn.functional as F
from lightning.fabric.strategies import ModelParallelStrategy
from lightning.pytorch.demos import Transformer, WikiText2
from torch.distributed._composable.fsdp.fully_shard import fully_shard
from torch.distributed.device_mesh import DeviceMesh
from torch.utils.data import DataLoader
from torchao.float8 import Float8LinearConfig, convert_to_float8_training
from tqdm import tqdm
def configure_model(model: nn.Module, device_mesh: DeviceMesh) -> nn.Module:
float8_config = Float8LinearConfig(
# pip install -U --index-url <https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/Triton-Nightly/pypi/simple/> triton-nightly # noqa
pad_inner_dim=True,
)
def module_filter_fn(mod: torch.nn.Module, fqn: str):
# we skip the decoder because it typically vocabulary size
# is not divisible by 16 as required by float8
return fqn != "decoder"
convert_to_float8_training(model, config=float8_config, module_filter_fn=module_filter_fn)
for module in model.modules():
if isinstance(module, (torch.nn.TransformerEncoderLayer, torch.nn.TransformerDecoderLayer)):
fully_shard(module, mesh=device_mesh)
fully_shard(model, mesh=device_mesh)
return torch.compile(model)
def train():
L.seed_everything(42)
batch_size = 8
micro_batch_size = 1
max_steps = 100
dataset = WikiText2()
dataloader = DataLoader(dataset, num_workers=8, batch_size=micro_batch_size)
with torch.device("meta"):
model = Transformer(
vocab_size=dataset.vocab_size,
nlayers=16,
nhid=4096,
ninp=1024,
nhead=32,
)
strategy = ModelParallelStrategy(data_parallel_size=4, tensor_parallel_size=1, parallelize_fn=configure_model)
fabric = L.Fabric(precision="bf16-true", strategy=strategy)
fabric.launch()
model = fabric.setup(model)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
optimizer = fabric.setup_optimizers(optimizer)
dataloader = fabric.setup_dataloaders(dataloader)
iterable = tqdm(enumerate(dataloader), total=len(dataloader)) if fabric.is_global_zero else enumerate(dataloader)
steps = 0
for i, batch in iterable:
input, target = batch
is_accumulating = i % (batch_size // micro_batch_size) != 0
with fabric.no_backward_sync(model, enabled=is_accumulating):
output = model(input, target)
loss = F.nll_loss(output, target.view(-1))
fabric.backward(loss)
if not is_accumulating:
fabric.clip_gradients(model, optimizer, max_norm=1.0)
optimizer.step()
optimizer.zero_grad()
steps += 1
if fabric.is_global_zero:
iterable.set_postfix_str(f"train_loss={loss.item():.2f}")
if steps == max_steps:
break
fabric.print(torch.cuda.memory_summary())
if __name__ == "__main__":
torch.set_float32_matmul_precision("high")
train()
</details>
As these examples show, it's now easier than ever to take your PyTorch Lightning module and run it with FSDP2 and/or tensor parallelism in FP8 precision, using the ModelParallelStrategy we introduced in 2.4.
Also note the use of distributed tensor APIs, TorchAO APIs, and torch.compile directly in the configure_model hook (or in the parallelize function in Fabric's ModelParallelStrategy), as opposed to the LightningModule as a whole. The advantage with this approach is that you can just copy-paste the parallelize functions that come with native PyTorch models directly in configure_model and get the same effect, no head-scratching involved :nerd_face:.
Talking about head scratching, we also made a pass at the PyTorch Lightning internals and hardened the parts where we keep track of progress counters during training, validation, testing, as well as learning rate scheduling, in relation to resuming from checkpoints. We now made sure there are no (to the best of our knowledge) edge cases where stopping and resuming from checkpoints can change the sequence of loops or other internal states. Fault tolerance for the win :partying_face:!
Alright! Feel free to take a look at the full changelog below.
And of course: the best way to use PyTorch Lightning and Fabric is through Lightning Studio :zap:. Access GPUs, train models, deploy and more with zero setup. Focus on data and models - not infrastructure.
<a name="changelog"></a>
<a name="changelog-pytorch"></a>
<details open><summary>Added</summary>
step parameter to TensorBoardLogger.log_hyperparams to visualize changes during training (#20176)str method to datamodule (#20301)Trainer.save_checkpoint (#20405)</details>
<details open><summary>Changed</summary>
resume_from_checkpoint as deprecated (#20477)np.random.SeedSequence() in pl_worker_init_function() to robustly seed NumPy-dependent dataloader workers (#20369)2.5 (#20351)_ (#20221)BytesIO as path in .to_onnx() (#20172)</details>
<details open><summary>Removed</summary>
List[int] as input type for Trainer when accelerator="cpu" (#20399)</details>
<details open><summary>Fixed</summary>
convert_module in FSDP to avoid using more memory than necessary during initialization (#20323)configure_optimizers when running with ReduceLROnPlateau (#20471)configure_optimizers example (#20420)_class_path parameter (#20221)</details>
<a name="changelog-fabric"></a>
<details open><summary>Added</summary>
step parameter to TensorBoardLogger.log_hyperparams to visualize changes during training (#20176)ddp_find_unused_parameters_true alias in Fabric's DDPStrategy (#20125)</details>
<details open><summary>Changed</summary>
np.random.SeedSequence() in pl_worker_init_function() to robustly seed NumPy-dependent dataloader workers (#20369)2.5 (#20351)</details>
<details open><summary>Removed</summary>
</details>
<details open><summary>Fixed</summary>
convert_module in FSDP to avoid using more memory than necessary during initialization (#20323)</details>
</br>
Full commit list: 2.4.0 -> 2.5.0
<a name="contributors"></a>
We thank all folks who submitted issues, features, fixes and doc changes. It's the only way we can collectively make Lightning :zap: better for everyone, nice job!
In particular, we would like to thank the authors of the pull-requests above, in no particular order:
@ringohoffman @MrWhatZitToYaa @jedyang97 @chualanagit @lantiga @AlessandroW @kazuar @t-vi @01AbhiSingh @WangYue0000 @amorehead @EricCousineau-TRI @mauvilsa @Borda @pete-mcelroy @ali-alshaar7 @GdoongMathew @farhadrgh @tshu-w @LukasSalchow @awindmann @dadwadw233 @qingquansong
Thank you :heart: and we hope you'll keep them coming!
Nothing published for this version
Avoid LightningCLI saving hyperparameters with class_path and init_args since this would be a breaking change
Lightning AI :zap: is excited to announce the release of Lightning 2.4. This is mainly a compatibility upgrade for PyTorch 2.4 and Python 3.12, with a sprinkle of a few features and bug fixes.
Did you know? The Lightning philosophy extends beyond a boilerplate-free deep learning framework: We've been hard at work bringing you Lightning Studio. Code together, prototype, train, deploy, host AI web apps. All from your browser, with zero setup.
<a name="changelog"></a>
<a name="changelog-pytorch"></a>
<details open><summary>Added</summary>
dump_stats flag to AdvancedProfiler (#19703)verbose to the seed_everything() function (#20108)TQDMProgressBar now provides an option to retain prior training epoch bars (#19578)ModelSummary table (#20159)</details>
<details open><summary>Changed</summary>
.fit(), .evaluate(), .test() or .predict() now terminates all processes launched by the Trainer and exits the program (#19976)seed_everything(..., workers=True) (#20055)</details>
<details open><summary>Removed</summary>
</details>
<details open><summary>Fixed</summary>
class_path and init_args since this would be a breaking change (#20068)seed_everything() (#20108)_LoggerConnector's _ResultMetric to move all registered keys to the device of the logged value if needed (#19814)_optimizer_to_device logic for special 'step' key in optimizer state causing performance regression (#20019)ModelSummary when model has distributed parameters (DTensor) (#20163)</details>
<a name="changelog-fabric"></a>
<details open><summary>Added</summary>
verbose to the seed_everything() function (#20108)</details>
<details open><summary>Changed</summary>
seed_everything(..., workers=True) (#20055)</details>
<details open><summary>Removed</summary>
</details>
<details open><summary>Fixed</summary>
_lazy_load() function (#20121)_optimizer_to_device logic for special 'step' key in optimizer state causing performance regression (#20019)</details>
</br>
Full commit list: 2.3.0 -> 2.4.0
<a name="contributors"></a>
We thank all our contributors who submitted pull requests for features, bug fixes and documentation updates.
Chuck Norris can solve NP-hard problems in polynomial time. In fact, any problem is easy when Chuck Norris solves it.
This release removes the code from the main lightning package that was reported in CVE-2024-5980.
This release removes the code from the main lightning package that was reported in CVE-2024-5980.
Includes a minor bugfix that avoids a conflict with the entrypoint command with another package #20041.
Includes a minor bugfix that avoids a conflict with the entrypoint command with another package #20041.
Includes minor bugfixes and stability improvements.
Includes minor bugfixes and stability improvements.
Full Changelog: https://github.com/Lightning-AI/pytorch-lightning/compare/2.3.0...2.3.1
Lightning AI is excited to announce the release of Lightning 2.3 :zap:
Lightning AI is excited to announce the release of Lightning 2.3 :zap:
Did you know? The Lightning philosophy extends beyond a boilerplate-free deep learning framework: We've been hard at work bringing you Lightning Studio. Code together, prototype, train, deploy, host AI web apps. All from your browser, with zero setup.
This release introduces experimental support for Tensor Parallelism and 2D Parallelism, PyTorch 2.3 support, and several bugfixes and stability improvements.
<a name="highlights"></a>
<a name="highlights-tensor-parallel"></a>
Tensor parallelism (TP) is a technique that splits up the computation of selected layers across GPUs to save memory and speed up distributed models. To enable TP as well as other forms of parallelism, we introduce a ModelParallelStrategy for both Lightning Trainer and Fabric. Under the hood, TP is enabled through new experimental PyTorch APIs like DTensor and torch.distributed.tensor.parallel.
Enabling TP in a model with PyTorch Lightning requires you to implement the LightningModule.configure_model() method where you convert selected layers of a model to paralellized layers. This is an advanced feature, because it requires a deep understanding of the model architecture. Open the tutorial Studio to learn the basics of Tensor Parallelism.
<a target="_blank" href="https://lightning.ai/lightning-ai/studios/tensor-parallelism-supercharging-large-model-training-with-pytorch-lightning"> <img src="https://pl-bolts-doc-images.s3.us-east-2.amazonaws.com/app-2/studio-badge.svg" alt="Open In Studio"/> </a>
import lightning as L
from lightning.pytorch.strategies import ModelParallelStrategy
from torch.distributed.tensor.parallel import ColwiseParallel, RowwiseParallel
from torch.distributed.tensor.parallel import parallelize_module
# 1. Implement the `configure_model()` method in LightningModule
class LitModel(L.LightningModule):
def __init__(self):
super().__init__()
self.model = FeedForward(8192, 8192)
def configure_model(self):
# Lightning will set up a `self.device_mesh` for you
tp_mesh = self.device_mesh["tensor_parallel"]
# Use PyTorch's distributed tensor APIs to parallelize the model
plan = {
"w1": ColwiseParallel(),
"w2": RowwiseParallel(),
"w3": ColwiseParallel(),
}
parallelize_module(self.model, tp_mesh, plan)
def training_step(self, batch):
...
# 2. Create the strategy
strategy = ModelParallelStrategy()
# 3. Configure devices and set the strategy in Trainer
trainer = L.Trainer(accelerator="cuda", devices=2, strategy=strategy)
trainer.fit(...)
<details><summary>Full training example (requires at least 2 GPUs).</summary>
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.distributed.tensor.parallel import ColwiseParallel, RowwiseParallel
from torch.distributed.tensor.parallel import parallelize_module
import lightning as L
from lightning.pytorch.demos.boring_classes import RandomDataset
from lightning.pytorch.strategies import ModelParallelStrategy
class FeedForward(nn.Module):
def __init__(self, dim, hidden_dim):
super().__init__()
self.w1 = nn.Linear(dim, hidden_dim, bias=False)
self.w2 = nn.Linear(hidden_dim, dim, bias=False)
self.w3 = nn.Linear(dim, hidden_dim, bias=False)
def forward(self, x):
return self.w2(F.silu(self.w1(x)) * self.w3(x))
class LitModel(L.LightningModule):
def __init__(self):
super().__init__()
self.model = FeedForward(8192, 8192)
def configure_model(self):
if self.device_mesh is None:
return
# Lightning will set up a `self.device_mesh` for you
tp_mesh = self.device_mesh["tensor_parallel"]
# Use PyTorch's distributed tensor APIs to parallelize the model
plan = {
"w1": ColwiseParallel(),
"w2": RowwiseParallel(),
"w3": ColwiseParallel(),
}
parallelize_module(self.model, tp_mesh, plan)
def training_step(self, batch):
output = self.model(batch)
loss = output.sum()
return loss
def configure_optimizers(self):
return torch.optim.AdamW(self.model.parameters(), lr=3e-3)
def train_dataloader(self):
# Trainer configures the sampler automatically for you such that
# all batches in a tensor-parallel group are identical
dataset = RandomDataset(8192, 64)
return torch.utils.data.DataLoader(dataset, batch_size=8, num_workers=2)
strategy = ModelParallelStrategy()
trainer = L.Trainer(
accelerator="cuda",
devices=2,
strategy=strategy,
max_epochs=1,
)
model = LitModel()
trainer.fit(model)
trainer.print(f"Peak memory usage: {torch.cuda.max_memory_allocated() / 1e9:.02f} GB")
</details>
</br>
Applying TP in a model with Fabric requires you to implement a special function where you convert selected layers of a model to paralellized layers. This is an advanced feature, because it requires a deep understanding of the model architecture. Open the tutorial Studio to learn the basics of Tensor Parallelism.
<a target="_blank" href="https://lightning.ai/lightning-ai/studios/tensor-parallelism-supercharging-large-model-training-with-lightning-fabric"> <img src="https://pl-bolts-doc-images.s3.us-east-2.amazonaws.com/app-2/studio-badge.svg" alt="Open In Studio"/> </a>
import lightning as L
from lightning.fabric.strategies import ModelParallelStrategy
from torch.distributed.tensor.parallel import ColwiseParallel, RowwiseParallel
from torch.distributed.tensor.parallel import parallelize_module
# 1. Implement the parallelization function for your model
def parallelize_feedforward(model, device_mesh):
# Lightning will set up a device mesh for you
tp_mesh = device_mesh["tensor_parallel"]
# Use PyTorch's distributed tensor APIs to parallelize the model
plan = {
"w1": ColwiseParallel(),
"w2": RowwiseParallel(),
"w3": ColwiseParallel(),
}
parallelize_module(model, tp_mesh, plan)
return model
# 2. Pass the parallelization function to the strategy
strategy = ModelParallelStrategy(parallelize_fn=parallelize_feedforward)
# 3. Configure devices and set the strategy in Fabric
fabric = L.Fabric(accelerator="cuda", devices=2, strategy=strategy)
fabric.launch()
<details><summary>Full training example (requires at least 2 GPUs).</summary>
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.distributed.tensor.parallel import ColwiseParallel, RowwiseParallel
from torch.distributed.tensor.parallel import parallelize_module
import lightning as L
from lightning.pytorch.demos.boring_classes import RandomDataset
from lightning.fabric.strategies import ModelParallelStrategy
class FeedForward(nn.Module):
def __init__(self, dim, hidden_dim):
super().__init__()
self.w1 = nn.Linear(dim, hidden_dim, bias=False)
self.w2 = nn.Linear(hidden_dim, dim, bias=False)
self.w3 = nn.Linear(dim, hidden_dim, bias=False)
def forward(self, x):
return self.w2(F.silu(self.w1(x)) * self.w3(x))
def parallelize_feedforward(model, device_mesh):
# Lightning will set up a device mesh for you
tp_mesh = device_mesh["tensor_parallel"]
# Use PyTorch's distributed tensor APIs to parallelize the model
plan = {
"w1": ColwiseParallel(),
"w2": RowwiseParallel(),
"w3": ColwiseParallel(),
}
parallelize_module(model, tp_mesh, plan)
return model
strategy = ModelParallelStrategy(parallelize_fn=parallelize_feedforward)
fabric = L.Fabric(accelerator="cuda", devices=2, strategy=strategy)
fabric.launch()
# Initialize the model
model = FeedForward(8192, 8192)
model = fabric.setup(model)
# Define the optimizer
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-3)
optimizer = fabric.setup_optimizers(optimizer)
# Define dataset/dataloader
dataset = RandomDataset(8192, 64)
dataloader = torch.utils.data.DataLoader(dataset, batch_size=8)
dataloader = fabric.setup_dataloaders(dataloader)
# Simplified training loop
for i, batch in enumerate(dataloader):
output = model(batch)
loss = output.sum()
fabric.backward(loss)
optimizer.step()
optimizer.zero_grad()
fabric.print(f"Iteration {i} complete")
fabric.print(f"Peak memory usage: {torch.cuda.max_memory_allocated() / 1e9:.02f} GB")
</details>
</br>
<a name="highlights-2d-parallel"></a>
Tensor Parallelism by itself can be very effective for efficient inference of very large models. For training, TP is typically combined with other forms of parallelism, such as FSDP, to increase throughput and scalability on large clusters with 100s of GPUs. The new ModelParallelStrategy in this release supports the combination of TP + FSDP, which is referred to as 2D parallelism.
For an introduction to this feature, please also refer to the tutorial Studios (PyTorch Lightning, Lightning Fabric). At the moment, the PyTorch team is reimplementing FSDP under the name FSDP2 with the aim to make it compose well with other parallelisms such as TP. Therefore, for the experimental 2D parallelism support, you'll need to switch to using FSDP2 with the new ModelParallelStrategy. Please refer to our docs (PyTorch Lightning, Lightning Fabric) and stay tuned for future releases as these APIs mature.
<a name="highlights-model-summary"></a>
The model summary table that gets displayed when you run Trainer.fit() now contains a new column "Mode" that shows the training mode each layer is in (#19468).
| Name | Type | Params | Mode
-----------------------------------------------------------------
0 | model | Sam | 93.7 M | train
1 | model.image_encoder | ImageEncoderViT | 89.7 M | eval
2 | model.prompt_encoder | PromptEncoder | 6.2 K | train
3 | model.mask_decoder | MaskDecoder | 4.1 M | train
-----------------------------------------------------------------
93.7 M Trainable params
0 Non-trainable params
93.7 M Total params
374.942 Total estimated model params size (MB)
A module in PyTorch is always either in train (default) or eval mode.
This improvement should give users more visibility into the state of their model and help debug issues, for example when you need to make sure certain layers of the model are frozen.
<a name="highlights-forward-methods"></a>
Until now, Lightning Fabric warned the user in case the forward pass of the model or a subset of its modules was conducted through methods other than the dedicated forward method of the PyTorch module. The reason for this is that PyTorch needs to run special hooks in case of DDP/FSDP and other strategies to function properly, and not running through the real forward method would skip these hooks and lead to correctness issues.
In Lightning Fabric 2.3, we added a feature to explicitly mark alternative forward methods so that Fabric can add the necessary rerouting behind the scenes:
import lightning as L
fabric = L.Fabric(devices=2, strategy="ddp")
fabric.launch()
model = MyModel()
model = fabric.setup(model)
# OK: Calling the model directly
output = model(input)
# ERROR: Calling another method that calls forward indirectly
prediction = model.generate(input)
# New: Mark special forward methods explicitly before using them
model.mark_forward_method(model.generate)
# OK: Now can use `model.generate()` in DDP/FSDP without issues
prediction = model.generate(input)
Find the full example and more details in our docs.
<a name="bc-changes"></a>
The 2.0 series of Lightning releases guarantees core API stability: No name changes, argument renaming, hook removals etc. on core interfaces (Trainer, LightningModule, etc.) unless a feature is specifically marked experimental. Here we list a few behavioral changes made in places where the change was justified if it significantly improves the user experience, improves performance, or fixes the correctness of a feature. These changes will likely not impact most users.
It is no longer allowed to skip training_step() by returning None in distributed training (#19918). The following usage was previously possible but would result in unpredictable hangs and timeouts in distributed training:
def training_step(self, batch):
loss = ...
if loss.isnan():
# No longer allowed in multi-GPU!
# Raises error in Lightning >= 2.3
return None
return loss
We decided to raise an error if the user attempts to return None when running in a multi-GPU setting.
prepare_data() hook in LightningModule and LightningDataModule is now subject to a barrier without timeout to avoid long-running tasks to be interrupted (#19448). Similarly, also in Fabric the Fabric.rank_zero_first context manager now uses an infinite barrier (#19448).<a name="changelog"></a>
<a name="changelog-pytorch"></a>
<details><summary>Added</summary>
ModelSummary and RichModelSummary callbacks now display the training mode of each layer in the column "Mode" (#19468)load_from_checkpoint support for LightningCLI when using dependency injection (#18105)on_exception hook to LightningDataModule (#19601)ModelParallelStrategy to support 2D parallelism (#19878, #19888)torch.distributed.destroy_process_group in atexit handler if process group needs destruction (#19931)FSDPStrategy(device_mesh=...) argument (#19504)</details>
<details><summary>Changed</summary>
prepare_data() hook in LightningModule and LightningDataModule is now subject to a barrier without timeout to avoid long-running tasks to be interrupted (#19448)drop_last for prediction (#19678)training_step() by returning None in distributed training (#19918)</details>
<details><summary>Removed</summary>
Trainer(strategy="bagua")) (#19445)</details>
<details><summary>Fixed</summary>
WandbLogger.log_hyperparameters() raising an error if hyperparameters are not JSON serializable (#19769)ModelCheckpoint(save_last=...) argument (#19808)epoch_loop.restarting to avoid full validation run after LearningRateFinder (#19818)</details>
<a name="changelog-fabric"></a>
<details><summary>Added</summary>
fabric consolidate in the new CLI (#19560)_FabricModule.mark_forward_method() (#19690)ModelParallelStrategy to support 2D parallelism (#19846, #19852, #19870, #19872)torch.distributed.destroy_process_group in atexit handler if process group needs destruction (#19931)FSDPStrategy(device_mesh=...) argument (#19504)</details>
<details><summary>Changed</summary>
lightning run model to fabric run (#19442, #19527)Fabric.rank_zero_first context manager now uses a barrier without timeout to avoid long-running tasks to be interrupted (#19448)fabric.backward() when it is needed by the strategy or precision selection (#19447, #19493)_BackwardSyncControl can now control what to do when gradient accumulation is disabled (#19577)</details>
<details><summary>Removed</summary>
</details>
<details><summary>Fixed</summary>
</details>
</br>
Full commit list: 2.2.0 -> 2.3.0
<a name="contributors"></a>
We thank all our contributors who submitted pull requests for features, bug fixes and documentation updates.
Chuck Norris is a big fan and daily user of Lightning Studio.
Fixed a matrix shape mismatch issue when running a model loaded from a quantized checkpoint (bitsandbytes)
Full Changelog: https://github.com/Lightning-AI/pytorch-lightning/compare/2.2.4...2.2.5
Fixed HTTPClient retry for flow/work queue
No Changes.
No Changes.
Full Changelog: https://github.com/Lightning-AI/pytorch-lightning/compare/2.2.3...2.2.4
Fixed WandbLogger.log_hyperparameters() raising an error if hyperparameters are not JSON serializable
WandbLogger.log_hyperparameters() raising an error if hyperparameters are not JSON serializable (#19769)No Changes.
Full Changelog: https://github.com/Lightning-AI/pytorch-lightning/compare/2.2.2...2.2.3
Fixed an issue causing a TypeError when using torch.compile as a decorator
torch.compile as a decorator (#19627)save_weights_only=True (#19524)torch.compile as a decorator (#19627)Fabric.setup() when using FSDP (#19755)Full Changelog: https://github.com/Lightning-AI/pytorch-lightning/compare/2.2.1...2.2.2
@ankitgola005 @awaelchli @Borda @carmocca @dmitsf @dvoytan-spark @fnhirwa
Fixed an issue with CSVLogger trying to append to file from a previous run when the version is set manually
Trainer.accumulate_grad_batches and Trainer.log_every_n_steps in ThroughputMonitor (#19470)Full Changelog: https://github.com/Lightning-AI/pytorch-lightning/compare/2.2.0post...2.2.1
@Raalsky @awaelchli @carmocca @Borda
If we forgot someone due to not matching commit email with GitHub account, let us know :]
Full Changelog: https://github.com/Lightning-AI/pytorch-lightning/compare/2.2.0...2.2.0.post0
Full Changelog: https://github.com/Lightning-AI/pytorch-lightning/compare/2.2.0...2.2.0.post0
Deprecated all precision plugin classes under lightning.pytorch.plugins with the suffix Plugin in the name
Lightning AI is excited to announce the release of Lightning 2.2 :zap:
Did you know? The Lightning philosophy extends beyond a boilerplate-free deep learning framework: We've been hard at work bringing you Lightning Studio. Code together, prototype, train, deploy, host AI web apps. All from your browser, with zero setup.
While our previous release was packed with many big new features, this time around we're rolling out mainly improvements based on feedback from the community. And of course, as the name implies, this release fully supports the latest PyTorch 2.2 :tada:
<a name="highlights"></a>
<a name="highlights-throughput"></a>
Lightning now has built-in utilities to measure throughput metrics such as batches/sec, samples/sec and Model FLOP Utilization (MFU) (#18848).
Trainer:
For the Trainer, this comes in form of a ThroughputMonitor callback. In order to track samples/sec, you need to provide a function to tell the monitor how to extract the batch dimension from your input. Furthermore, if you want to track MFU, you can provide a sample forward pass and the ThroughputMonitor will automatically estimate the utilization based on the hardware you are running on:
import lightning as L
from lightning.pytorch.callbacks import ThroughputMonitor
from lightning.fabric.utilities.throughput import measure_flops
class MyModel(LightningModule):
def setup(self, stage):
with torch.device("meta"):
model = MyModel()
def sample_forward():
batch = torch.randn(..., device="meta")
return model(batch)
self.flops_per_batch = measure_flops(model, sample_forward, loss_fn=torch.Tensor.sum)
throughput = ThroughputMonitor(
batch_size_fn=lambda batch: batch.size(0),
# optional, if your samples have a length (like number of tokens)
sample_fn=lambda batch: batch.size(1)
)
trainer = L.Trainer(log_every_n_steps=10, callbacks=throughput, logger=...)
model = MyModel()
trainer.fit(model)
The results get automatically sent to the logger if one is configured on the Trainer.
Fabric:
For Fabric, the ThroughputMonitor is a simple utility object on which you call .update() and compute_and_log() during the training loop:
import lightning as L
from lightning.fabric.utilities import ThroughputMonitor
fabric = L.Fabric(logger=...)
throughput = ThroughputMonitor(fabric)
t0 = time()
for batch_idx, batch in enumerate(train_dataloader):
do_work()
torch.cuda.synchronize() # required or else time() won't be correct
throughput.update(
time=(time() - t0),
batches=batch_idx,
samples=(batch_idx * batch_size)
)
if batch_idx % 10 == 0:
throughput.compute_and_log(step=batch_idx)
Check out our TinyLlama LLM pretraining script for a full example using Fabric's ThroughputMonitor.
The troughput utilities can report:
<a name="highlights-eval"></a>
When you train a model and have validation enabled, the Trainer automatically calls .eval() when transitioning to the validation loop, and .train() when validation ends. Until now, this had the unfortunate side effect that any submodules in your LightningModule that were in evaluation mode get reset to train mode. In Lightning 2.2, the Trainer now captures the mode of every submodule before switching to validation, and restores the mode the modules were in when validation ends (#18951, #18951, #18951). This improvement will help users avoid silent correctness bugs and removes boilerplate code for managing frozen layers.
import lightning as L
class LitModel(L.LightningModule):
def __init__(self):
super().__init__()
self.trainable_module = ...
# This will now stay in eval mode
self.frozen_module = ...
self.frozen_module.eval()
def training_step(self, batch):
# Previously, modules were all in train mode
# Now: Modules are in mode they were set up with
assert self.trainable_module.training
assert not self.frozen_module.training
...
def validation_step(self, batch):
# All modules are in eval mode
...
model = LitModel()
trainer = L.Trainer()
trainer.fit(model)
If you have overridden any of the LightningModule.on_{validation,test,predict}_model_{eval,train} hooks, they will still get called and execute your custom logic, but they are no longer required if you added them to preserve the eval mode of frozen modules.
[!IMPORTANT] In some libraries, for example HuggingFace, models are created in evaluation mode by default (e.g.
HFModel.from_pretrained(...)). Starting from 2.2, you will have to set.train()on these models if you intend to train them.
<a name="highlights-consolidate-fsdp"></a>
In the previous release, we introduced distributed checkpointing with FSDP to speed up saving and loading checkpoints for big models. These checkpoints are in a special format saved in a folder with shards from each GPU in a separate file. While these checkpoints can be loaded back with Lightning Trainer or Fabric very easily, they aren't easy to load or process externally. In Lightning 2.2, we introduced a CLI utility that lets you consolidate the checkpoint folder to a single file that can be loaded in raw PyTorch with torch.load() for example (#19213).
Given you saved a distributed checkpoint, you can then convert it like so:
# For Trainer checkpoints:
python -m lightning.pytorch.utilities.consolidate_checkpoint path/to/my/checkpoint
# For Fabric checkpoints:
python -m lightning.fabric.utilities.consolidate_checkpoint path/to/my/checkpoint
Read more about distributed checkpointing in our documentation: Trainer, Fabric.
<a name="highlights-compile"></a>
PyTorch 2.0+ introduced torch.compile, a powerful tool to speed up your models without changing the code.
We now added a comprehensive guide how to use torch.compile correctly with tips and tricks to help you troubleshoot common issues. On top of that, Fabric.setup() will now reapply torch.compile on top of DDP/FSDP if you are enabling these strategies (#19280).
import lightning as L
# Select a distributed strategy (DDP, FSDP, ...)
fabric = L.Fabric(strategy="ddp", devices=8)
# Compile your model before `.setup()`
model = torch.compile(model)
# Now automatically handles compiling also over DDP/FSDP
model = fabric.setup(model)
# You can opt-out if it is causing trouble
model = fabric.setup(model, _reapply_compile=False)
You might see fewer graph breaks, but there won't be any significant speed-ups with this. We introduced this mainly to make Fabric ready for future improvements from PyTorch to optimizing distributed operations.
<a name="highlights-dataloader-state"></a>
If you use a dataloader/iterable that implements the .state_dict() and .load_state_dict() interface, the Trainer will now automatically save and load their state in the checkpoint (#19361).
import lightning as L
class MyDataLoader:
"""A dataloader that implements the 'stateful' interface."""
def state_dict(self):
# Return a dictionary with state
return {"batches_fetched": ...}
def load_state_dict(self, state_dict):
# Load the state from the checkpoint
self.batches_fetched = state_dict["batches_fetched"]
model = ...
dataloader = MyDataLoader()
trainer = L.Trainer()
# Saves checkpoints that include the dataloader state
trainer.fit(model, dataloader)
# When you resume training, the dataloader can now load its state
trainer.fit(model, dataloader, ckpt_path="path/to/my/checkpoint")
Note that the standard PyTorch DataLoader does not support this stateful interface. This feature only works on loaders that implement these two methods. A dataloader that supports full fault-tolerance will be included in our upcoming release of Lightning Data - a library to optimize data preprocessing and streaming in the cloud. Stay tuned!
<a name="highlights-non-strict"></a>
A feature that has been requested for a long time by the community is non-strict checkpoint loading. By default, a checkpoint in PyTorch is loaded with strict=True to ensure all keys in the saved checkpoint match what's in the model's state dict.
However, in some use cases it might make sense to exclude certain weights from being included in the checkpoint. When resuming training, the user would then be required to set strict=False, which wasn't configurable until now.
You can now set the attribute strict_loading=False on your LightningModule if you want to allow loading partial checkpoints (#19404).
import lightning as L
class LitModel(L.LightningModule):
def __init__(self):
super().__init__()
# This model only trains the decoder, we don't save the encoder
self.encoder = from_pretrained(...).requires_grad_(False)
self.decoder = Decoder()
# Set to False because we only care about the decoder
self.strict_loading = False
def state_dict(self):
# Don't save the encoder, it is not being trained
return {k: v for k, v in super().state_dict().items() if "encoder" not in k}
...
trainer = L.Trainer()
model = LitModel()
# Will load weights with `.load_state_dict(strict=model.strict_loading)`
trainer.fit(model, ckpt_path="path/to/checkpoint")
Full documentation here.
<a name="bc-changes"></a>
The 2.0 series of Lightning releases guarantees core API stability: No name changes, argument renaming, hook removals etc. on core interfaces (Trainer, LightningModule, etc.) unless a feature is specifically marked experimental. Here we list a few behavioral changes made in places where the change was justified if it significantly improves the user experience, improves performance, or fixes the correctness of a feature. These changes will likely not impact most users.
<a name="highlights-save-last"></a>
In Lightning 2.1, we made the ModelCheckpoint(..., save_last=True) feature save a symbolic link to the last saved checkpoint instead of rewriting the checkpoint (#18748). This time saver is especially useful for large models who take a while to save. However, many users were confused by the new behavior and wanted it turned off, saving a copy instead of a symbolic link like before. In Lightning 2.2, we are reverting this decision and make the linking opt-in (#19191):
from lightning.pytorch.callbacks import ModelCheckpoint
# In 2.1 saves a symbolic link "last.ckpt" to the last checkpoint saved
# In 2.2 saves "last.ckpt" as a copy of the last checkpoint saved
checkpoint = ModelCheckpoint("./my_checkpoints", save_last=True)
# You can opt-in to save a symlink (if possible)
checkpoint = ModelCheckpoint("./my_checkpoints", save_last="link")
<a name="highlights-seed-everything"></a>
The seed_everything(x) utility function is useful to set the seed for several libraries like PyTorch, NumPy and Python in a single line of code. However, until now you were allowed to omit passing a seeding value, in which case the function picked a seed value randomly. In certain cases, for example when processes are launched externally (e.g., SLURM, torchelastic etc.), this default behavior is dangerous because each process will independently choose a random seed. This can affect sampling, randomized validation splits, and other behaviors that rely on each process having the same seed. In 2.2, we removed this default behavior and default to a seed value 0 (#18846):
from lightning.pytorch.utilities import seed_everything
# Set the random seed for PyTorch, NumPy, Python etc.
seed_everything(42)
# Not setting a value now defaults to 0
seed_everything()
In the unlikely event that you relied on the previous behavior, you now have to choose the seed randomly yourself:
seed_everything(random.randint(0, 1000000))
metrics.csv file produced by CSVLogger are now sorted alphabetically (#19159)TransformerEnginePrecision(fallback_compute_dtype=) to control the dtype of operations that don't support fp8 (#19082)TransformerEnginePrecision(dtype=) argument to weights_dtype and made it required (#19082)LightningModule.load_from_checkpoint() function now calls .configure_model() on the model if it is overridden, to ensure all layers can be loaded from the checkpoint (#19036)<a name="changelog"></a>
<a name="changelog-pytorch"></a>
<details><summary>Added</summary>
lightning.pytorch.callbacks.ThroughputMonitor to track throughput and log it (#18848).train() or .eval() on a submodule-level when switching from validation to training (#18951)TransformerEnginePrecision(fallback_compute_dtype=) to control the dtype of operations that don't support fp8 (#19082)ModelCheckpoint(save_last='link') to create a symbolic link for the 'last.ckpt' file (#19191)TQDM_MINITERS for setting the refresh rate (#19381)strategy='deepspeed_stage_1_offload' to the strategy registry (#19075)LightningModule.strict_loading = True | False attribute (#19404)</details>
<details><summary>Changed</summary>
seed_everything() without passing in a seed no longer randomly selects a seed, and now defaults to 0 (#18846)LightningModule.on_{validation,test,predict}_model_{eval,train} now only get called if they are overridden by the user (#18951)Trainer.fit() loop no longer calls LightningModule.train() at the start; it now preserves the user's configuration of frozen layers (#18951)LightningModule.load_from_checkpoint() function now calls .configure_model() on the model if it is overridden, to ensure all layers can be loaded from the checkpoint (#19036)step parameter when logging metrics with NeptuneLogger (#19126)TransformerEnginePrecision(dtype=) argument to weights_dtype and made it required (#19082)metrics.csv file produced by CSVLogger are now sorted alphabetically (#19159)ModelCheckpoint(save_last=True) instead of creating a symbolic link (#19191)</details>
<details><summary>Deprecated</summary>
lightning.pytorch.plugins with the suffix Plugin in the name (#18840)</details>
<details><summary>Removed</summary>
</details>
<details><summary>Fixed</summary>
precision="transformer-engine" argument would not replace layers by default (#19082)LightningModule.setup or LightningModule.configure_model wouldn't get converted when using the Bitsandbytes or TransformerEngine plugins (#19061)FSDPStrategy to accept a device_mesh (#19392)</details>
<a name="changelog-fabric"></a>
<details><summary>Added</summary>
lightning.fabric.utilities.ThroughputMonitor and lightning.fabric.utilities.Throughput to track throughput and log it (#18848)lightning.fabric.utilities.AttributeDict for convenient dict-attribute access to represent state in script (#18943)TransformerEnginePrecision(fallback_compute_dtype=) to control the dtype of operations that don't support fp8 (#19082)Fabric.setup() over the FSDP/DDP wrappers (#19280)</details>
<details><summary>Changed</summary>
seed_everything() without passing in a seed no longer randomly selects a seed, and now defaults to 0 (#18846)TransformerEnginePrecision(dtype=) argument to weights_dtype and made it required (#19082)metrics.csv file produced by CSVLogger are now sorted alphabetically (#19159)</details>
<details><summary>Removed</summary>
</details>
<details><summary>Fixed</summary>
get_available_flops (#18952)precision="transformer-engine" argument would not replace layers by default (#19082)FSDPStrategy to accept a device_mesh (#19392)</details>
</br>
Full commit list: 2.1.0 -> 2.2.0
<a name="contributors"></a>
Everyone who contributed between 2.1 and 2.2, in no particular order:
@nik777 @Raalsky @wouterzwerink @AleksanderWWW @awaelchli @nohalon @ioangatop @Borda @ethanwharris @BoringDonut @mauvilsa @parambharat @tchaton @ryan597 @adamjstewart @rasbt @carmocca
@hiaoxui @VictorPrins @jaswon @AMHermansen @JalinWang @MF-FOOM @unacanal @Jamim @harishb00 @asingh9530 @dipta007 @daturkel @jerrymannil @mjbommar @shenmishajing @paganpasta @lauritsf @andyland @mathematicalmichael
Chuck Norris is a big fan and daily user of PyTorch Lightning.
This is a preview release for Lightning 2.2.0.
This is a preview release for Lightning 2.2.0.
Fixed an issue preventing Fabric to run on CPU when the system's CUDA driver is outdated or broken
Trainer not expanding the default_root_dir if it has the ~ (home) prefix (#19179)num_workers=1 and CPU count is 1 (#19224)WandbLogger.watch() method annotation to accept None for the log parameter (#19237)ModelCheckpoint(save_last="link") (#19303)_restricted_classmethod_impl would incorrectly raise a TypeError on inspection rather than on call (#19332)__version__ in __init__ (#19221)Full Changelog: https://github.com/Lightning-AI/pytorch-lightning/compare/2.1.3...2.1.4
@andyland @asingh9530 @awaelchli @Borda @daturkel @dipta007 @lauritsf @mjbommar @shenmishajing @tchaton
If we forgot someone due to not matching commit email with GitHub account, let us know :]
Lightning App: Use the batch get endpoint
StreamingDataset (#19052)StreamingDataset (#19050)StreamingDataset (#19049)StreamingDataset (#19044)StreamingDataset (#19168)StreamingDataset improve deletion strategy (#19118)StreamingDataset Speed (#19114)StreamingDataset (#19094)DataProcessor required (#18910)LightningDataset relying on un-maintained torchdata (#19019)move_to_device=False is passed (#19152)MPIEnvironment (#19074)LightningCLI no longer allows setting a normal class instance as default. A lazy_instance can be used instead (#18822)CHECKPOINT_JOIN_CHAR or CHECKPOINT_EQUALS_CHAR would only work on the ModelCheckpoint class but not on an instance (#19054)ModelCheckpoint not expanding the dirpath if it has the ~ (home) prefix (#19058)ModelCheckpoint would alternate between versioned and unversioned filename (#19064)MPIEnvironment (#19074)self.log to respect the default dtype (#19046)Full Changelog: https://github.com/Lightning-AI/lightning/compare/2.1.2...2.1.3
@AleksanderWWW, @awaelchli, @borda, @carmocca, @dependabot[bot], @mauvilsa, @MF-FOOM, @tchaton, @yassersouri
If we forgot someone due to not matching commit email with GitHub account, let us know :]
Forced plugin server to use localhost
torchmetrics wouldn't get moved to the device when using FSDP (#18954)Trainer.save_checkpoint() an FSDP model when Trainer.test/validate/predict() ran after Trainer.fit() (#18992)@awaelchli, @carmocca, @ethanwharris, @tchaton
If we forgot someone due to not matching commit email with GitHub account, let us know :]
Full Changelog: https://github.com/Lightning-AI/lightning/compare/2.1.1...2.1.2
Fix failing lightning cli entry point
fail() (#18883)forward that invokes submodules is now an error when the model is wrapped (e.g., with DDP) (#18819)Fabric.init_module() context manager when using the BitsandbytesPrecision plugin (#18914)last.ckpt file with a symlink (#18793)BatchSizeFinder steps_per_trial parameter ends up defining how many validation batches to run during the entire training (#18394)last.ckpt file when using ModelCheckpoint on a remote filesystem, and no logger is used (#18867)TensorBoardLogger and CSVLogger (#18897)@awaelchli, @borda, @BoringDonut, @carmocca, @hiaoxui, @ioangatop, @nohalon, @rasbt, @tchaton
If we forgot someone due to not matching commit email with GitHub account, let us know :]
Full Changelog: https://github.com/Lightning-AI/lightning/compare/2.1.0...2.1.1
…features, at the cost of a long list of breaking changes resulting in more work needed than usual to upgrade from 1.9 to 2.0. Moving forward, we promi…
Lightning AI is excited to announce the release of Lightning 2.1 :zap: It's the culmination of work from 79 contributors who have worked on features, bug-fixes, and documentation for a total of over 750+ commits since v2.0.
The theme of 2.1 is "bigger, better, faster": Bigger because training large multi-billion parameter models has gotten even more efficient thanks to FSDP, efficient initialization and sharded checkpointing improvements, better because it's easier than ever to scale models without making substantial code changes or installing third-party packages and faster because it leverages the latest hardware features to speed up training in low-bit precision thanks to new precision plugins like bitsandbytes and transformer engine. And of course, as the name implies, this release fully leverages the latest features in PyTorch 2.1 :tada:
<a name="highlights"></a>
<a name="highlights-fsdp"></a>
The FSDP strategy for training large billion-parameter models gets substantial improvements and new features in Lightning 2.1, both in Trainer and Fabric (in case you didn't know, Fabric is the latest addition to the Lightning family of tools to scale models without the boilerplate code). FSDP is now more user-friendly to configure, has memory management and speed improvements, and we have a brand new end-to-end user guide with best practices (Trainer, Fabric).
When training large billion-parameter models with FSDP, saving and resuming training, or even just loading model parameters for finetuning can be challenging, as users are are often plagued by out-of-memory errors and speed bottlenecks.
In 2.1, we made several improvements. Starting with saving checkpoints, we added support for distributed/sharded checkpoints, enabled through the setting state_dict_type in the strategy (#18364, #18358):
Trainer:
import lightning as L
from lightning.pytorch.strategies import FSDPStrategy
# Default used by the strategy
strategy = FSDPStrategy(state_dict_type="full")
# Enable saving distributed checkpoints
strategy = FSDPStrategy(state_dict_type="sharded")
trainer = L.Trainer(strategy=strategy, ...)
Fabric:
import lightning as L
from lightning.fabric.strategies import FSDPStrategy
# Saving distributed checkpoints is the default
strategy = FSDPStrategy(state_dict_type="sharded")
# Save consolidated (single file) checkpoints
strategy = FSDPStrategy(state_dict_type="full")
fabric = L.Fabric(strategy=strategy, ...)
Distributed checkpoints are the fastest and most memory efficient way to save the state of very large models. The distributed checkpoint format also makes it efficient to load these checkpoints back for resuming training in parallel, and it reduces the impact on CPU memory usage significantly. Furthermore, we've also introduced lazy-loading for non-distributed checkpoints (#18150, #18379), which greatly reduces the impact on CPU memory usage when loading a consolidated (single-file) checkpoint (e.g. for finetuning). Learn more about these features in our FSDP guides (Trainer, Fabric).
A major challenge that users face when working with large models such as LLMs is dealing with the extreme memory requirements. Even something as simple as instantiating a model becomes non-trivial if the model is so large it won't fit in a single GPU or even a single machine. In Lightning 2.1, we are introducing empty-weights initialization through the Fabric.init_module() (#17462, #17627) and Trainer.init_module()/LightningModule.configure_model() (#18004, #18004, #18385) methods:
Trainer:
import lightning as L
class MyModel(L.LightningModule):
def __init__(self):
super().__init__()
# Delay initialization of model to `configure_model()`
def configure_model(self):
# Model initialized in correct precision and weights on meta-device
self.model = ...
...
trainer = L.Trainer(strategy="fsdp", ...)
trainer.fit(model)
Fabric:
import lightning as L
fabric = L.Fabric(strategy="fsdp", ...)
# Model initialized in correct precision and weights on meta-device
with fabric.init_module(empty_init=True):
model = ...
# You can also initialize buffers and tensors directly on device and dtype
with fabric.init_tensor():
model.mask.create()
model.kv_cache.create()
x = torch.randn(4, 128)
# Materialization and sharding of model happens inside here
model = fabric.setup(model)
Read more about this new feature and its other benefits in our docs (Trainer, Fabric).
We made it super easy to configure the sharding- and activation-checkpointing policy when you want to auto-wrap particular layers of your model for advanced control (#18045, #18084).
import lightning as L
from lightning.pytorch.strategies import FSDPStrategy
- from torch.distributed.fsdp.wrap import ModuleWrapPolicy
- strategy = FSDPStrategy(auto_wrap_policy=ModuleWrapPolicy({MyTransformerBlock}))
+ strategy = FSDPStrategy(auto_wrap_policy={MyTransformerBlock})
trainer = L.Trainer(strategy=strategy, ...)
Furthermore, the sharding strategy can now be conveniently set with a string value (#18087):
import lightning as L
from lightning.pytorch.strategies import FSDPStrategy
- from torch.distributed.fsdp.fully_sharded_data_parallel import ShardingStrategy
- strategy = FSDPStrategy(sharding_strategy=ShardingStrategy.SHARD_GRAD_OP)
+ strategy = FSDPStrategy(sharding_strategy="SHARD_GRAD_OP")
trainer = L.Trainer(strategy=strategy, ...)
You no longer need to remember the long PyTorch imports! Fabric also supports all these improvements shown above.
<a name="highlights-half-precision"></a>
Lightning now supports true half-precision for training and inference with all built-in strategies (#18193, #18217, #18213, #18219). With this setting, the memory required to store the model weights is only half of what is normally needed when running with float32. In addition, you get the same speed benefits as mixed precision training (precision="16-mixed") has:
import lightning as L
# default
trainer = L.Trainer(precision="32-true")
# train with model weights in `torch.float16`
trainer = L.Trainer(precision="16-true")
# train with model weights in `torch.bfloat16`
# (if hardware supports it)
trainer = L.Trainer(precision="bf16-true")
The same settings are also available in Fabric! We recommend to try bfloat16 training (precision="bf16-true") as it is often more numerically stable than regular 16-bit precision (precision="16-true").
<a name="highlights-bitsandbytes"></a>
With the new Bitsandbytes precision plugin #18655, you can now quantize your model for significant memory savings during training, finetuning, or inference with a selection of several state-of-the-art quantization algorithms (int8, fp4, nf4 and more). For the first time, Trainer and Fabric make bitsandbytes easy to use for general models.
Trainer:
import lightning as L
from lightning.pytorch.plugins import BitsandbytesPrecisionPlugin
# this will pick out the compute dtype automatically, by default `bfloat16`
precision = BitsandbytesPrecisionPlugin("nf4-dq")
trainer = L.Trainer(plugins=precision)
Fabric:
import lightning as L
from lightning.fabric.plugins import BitsandbytesPrecision
# this will pick out the compute dtype automatically, by default `bfloat16`
precision = BitsandbytesPrecision("nf4-dq")
trainer = L.Fabric(plugins=precision)
<a name="highlights-transformer-engine"></a>
The Transformer Engine by NVIDIA is a library for accelerating transformer layers on the new Hopper (H100) generation of GPUs. With the integration in Lightning Trainer and Fabric (#17597, #18459), you have easy access to the 8-bit mixed precision for significant speed ups:
Trainer:
import lightning as L
# Select 8-bit mixed precision via TransformerEngine, with model weights in float16
trainer = L.Trainer(precision="transformer-engine-float16")
Fabric:
import lightning as L
# Select 8-bit mixed precision via TransformerEngine, with model weights in float16
fabric = L.Fabric(precision="transformer-engine-float16")
More configuration options are available through the respective plugins in Trainer and Fabric.
<a name="highlights-tpu"></a>
Lightning 2.1 runs on the latest generation of TPU hardware on Google Cloud! TPU-v4 and TPU-v5 (#17227) are now fully supported both in Fabric and Trainer and run using the new PjRT runtime by default (#17352). PjRT is the runtime used by Jax and has shown an average improvement of 35% on benchmarks.
Trainer:
import lightning as L
trainer = L.Trainer(accelerator="tpu", devices=8)
model = MyModel()
trainer.fit(model) # uses PjRT if available
Fabric:
import lightning as L
def train(fabric):
...
fabric = L.Fabric(accelerator="tpu")
fabric.launch(train) # uses PjRT if available
And what's even more exciting, you can now scale massive multi-billion parameter models on TPUs using FSDP (#17421).
import lightning as L
from lightning.fabric.strategies import XLAFSDPStrategy
strategy = XLAFSDPStrategy(
# Most arguments from the PyTorch native FSDP strategy are also available here!
auto_wrap_policy={Block},
activation_checkpointing_policy={Block},
state_dict_type="full",
sequential_save=True,
)
fabric = L.Fabric(devices=8, strategy=strategy)
fabric.launch(finetune)
You can find a full end-to-end finetuning example script in our Lit-GPT repository. The new XLA-FSDP strategy is experimental and currently only available in Fabric. Support in the Trainer will follow in the future.
<a name="highlights-fabric-checkpoints"></a>
Several improvements for checkpoint saving and loading have landed in Fabric, enabling more fine-grained control over what is saved/loaded while reducing boilerplate code:
There is a new Fabric.load_raw() method with which you can load model- or optimizer state-dicts saved externally by a non-Fabric application (e.g., raw PyTorch) (#18049)
import lightning as L
fabric = L.Fabric()
model = MyModel()
# A model weights file saved by your friend who doesn't use Fabric
fabric.load_raw("path/to/model.pt", model)
# Equivalent to this:
# model.load_state_dict(torch.load("path/to/model.pt"))
A new parameter Fabric.load(..., strict=True|False) to disable strict loading (#17645)
import lightning as L
fabric = L.Fabric()
model = MyModel()
state = {"model": model}
# strict loading is the default
fabric.load("path/to/checkpoint.ckpt", state, strict=True)
# disable strict loading
fabric.load("path/to/checkpoint.ckpt", state, strict=False)
A new parameter Fabric.save(..., filter=...) that enables you to exclude certain parameters of your model without writing boilerplate code for it (#17845)
import lightning as L
fabric = L.Fabric()
model, optimizer = ...
state = {"model": model, "optimizer": optimizer, "foo": 123}
# save only the weights that match a pattern
filter = {"model": lambda k, v: "weight" in k}
fabric.save("path/to/checkpoint.ckpt", state, filter=filter)
You can read more about the new options in our checkpoint guide.
<a name="bc-changes"></a>
The release of PyTorch Lightning 2.0 was a big step into a new chapter: It brought a more polished API and removed a lot of legacy code and outdated as well as experimental features, at the cost of a long list of breaking changes resulting in more work needed than usual to upgrade from 1.9 to 2.0. Moving forward, we promised to maintain full backward compatibility of our public core APIs to guarantee a smooth upgrade experience for everyone, and with 2.1 we are happy to deliver on this promise. A few exceptions were made in places where the change was justified if it significantly improves the user experience, improves performance, or fixes the correctness of a feature. These changes will likely not impact most users.
<a name="bc-changes-pytorch"></a>
When selecting device indices via devices=[i], the Trainer now selects the i-th TPU core (0-based, previously it was 1-based) (#17227)
Before:
# Selects the first TPU core (1-based index)
trainer = Trainer(accelerator="tpu", devices=[1])
Now:
# Selects the second TPU core (0-based index)
trainer = Trainer(accelerator="tpu", devices=[1])
Due to lack of reliability, Trainer now only runs on one GPU instead of all GPUs in a Jupyter notebook if devices="auto" (default) (#18291)
Before:
import lightning as L
# In Jupyter notebooks, this would select all available GPUs (DDP)
trainer = L.Trainer(accelerator="cuda", devices="auto")
Now:
# In Jupyter notebooks, this now selects only one GPU (the first)
trainer = L.Trainer(accelerator="cuda", devices="auto")
# You can still explicitly select multiple
trainer = L.Trainer(accelerator="cuda", devices=8)
LightningModule.setup(), the self.device now returns the device the module will be placed on instead of cpu (#18021)Before:
def setup(self, stage):
# CPU regardless of the accelerator used
print(self.device)
Now:
def setup(self, stage):
# CPU/CUDA/MPS/XLA depending on accelerator
print(self.device)
self.loged tensors are now kept in the original device to reduce unnecessary host-to-device synchronizations (#17334)FSDPStrategy now loads checkpoints after the configure_model/configure_sharded_model hook (#18358)FSDPStrategy.load_optimizer_state_dict and FSDPStrategy.load_model_state_dict are a no-op now (#18358)torchdistx due to a lack of project maintenance (#17995)<a name="bc-changes-fabric"></a>
We thank the community for the amazing feedback we got for Fabric so far - keep it coming. The list of breaking changes is short and won't affect the vast majority of users.
We removed automatic sharding support with Fabric.run or using fabric.launch(fn). This only impacts FSDP and DeepSpeed strategy users who use this way of launching. Please note that Fabric.run is a legacy construct from the LightningLite days, and is not recommended today. Please instantiate your large FSDP or DeepSpeed model under the newly added fabric.init_module context manager (#17832).
Before:
import lightning as L
def train(fabric):
# FSDP's `enable_wrap` context or `deepspeed.zero.Init()`
# were applied automaticaly here
model = LargeModel()
...
fabric = L.Fabric()
fabric.launch(train)
Now:
def train(fabric):
# Use `init_module` explicitly to apply these context managers
with fabric.init_module():
model = LargeModel()
...
Due to lack of reliability, Fabric now only runs on one GPU instead of all GPUs in a Jupyter notebook if devices="auto" (default) (#18291)
Before:
import lightning as L
# In Jupyter notebooks, this would select all available GPUs (DDP)
fabric = L.Fabric(accelerator="cuda", devices="auto")
Now:
# In Jupyter notebooks, this now selects only one GPU (the first)
fabric = L.Fabric(accelerator="cuda", devices="auto")
# You can still explicitly select multiple
fabric = L.Fabric(accelerator="cuda", devices=8)
<a name="changelog"></a>
<a name="changelog-pytorch"></a>
<details><summary>Added</summary>
metrics_format attribute to RichProgressBarTheme class (#18373)CHECKPOINT_EQUALS_CHAR attribute to ModelCheckpoint class (#17999)**summarize_kwargs to ModelSummary and RichModelSummary callbacks (#16788)max_size_cycle|max_size|min_size iteration modes during evaluation (#17163)XLAStrategy(sync_module_states=bool) to control whether to broadcast the parameters to all devices (#17522)FSDPStrategy (#16558)LightningDataModule.from_datasets to support arbitrary iterables (#17402)SaveConfigCallback.save_config to ease use cases such as saving the config to a logger (#17475)FSDPStrategy(timeout=...) for the FSDP strategy (#17274)FSDPStrategy(activation_checkpointing_policy=...) to customize the layer policy for automatic activation checkpointing (requires torch>=2.1) (#18045)--map-to-cpu to the checkpoint upgrade script to enable converting GPU checkpoints on a CPU-only machine (#17527)LearningRateMonitor to log monitored values to trainer.callback_metrics (#17626)log_weight_decay argument to LearningRateMonitor callback (#18439)Trainer.print() to print on local rank zero only (#17980)Trainer.init_module() context manager to instantiate large models efficiently directly on device, dtype (#18004)
torch.float32, torch.float64) depending on the 'true' precision choice in Trainer(precision='32-true'|'64-true')LightningModule.configure_model() hook to instantiate large models efficiently directly on device, dtype, and with sharding support (#18004)
Trainer.init_module(empty_init=True) in FSDP (#18385)lightning.pytorch.plugins.PrecisionPlugin.module_init_context() and lightning.pytorch.strategies.Strategy.tensor_init_context() context managers to control model and tensor instantiation (#18004)xla_model.mark_step() before saving checkpoints with XLA (#17882)torch.distributed.fsdp.ShardingStrategy via string in FSDPStrategy (#18087)XLAStrategy (#18194)Trainer(precision="16-true"|"bf16-true") (#18193, #18217, #18213, #18219)devices and num_nodes when running with SLURM or TorchElastic (#18292)FSDPStrategy(state_dict_type="full"|"sharded") (#18364)Trainer(precision="transformer-engine") using Nvidia's Transformer Engine (#18459)Trainer(plugins=BitsandbytesPrecision()) using bitsandbytes (#18655)FSDPStrategy (#18583)lightning.pytorch.utilities.suggested_max_num_workers to assist with setting a good value in distributed settings (#18591)num_workers warning to give a more accurate upper limit on the num_workers suggestion (#18591)lightning.pytorch.utilities.is_shared_filesystem utility function to automatically check whether the filesystem is shared between machines (#18586)Mapping from LightningModule.training_step() (#18657)LightningModule.on_validation_model_zero_grad() to allow overriding the behavior of zeroing the gradients before entering the validation loop (#18710)</details>
<details><summary>Changed</summary>
round(..., 3) to ".3f" format string in MetricsTextColumn class (#18483)self.trainer.model.parameters() in LightningModule.configure_optimizers() (#17309)Trainer(accelerator="tpu", devices=[i])" now selects the i-th TPU core (0-based, previously it was 1-based) (#17227)self.loged tensors are now kept in the original device to reduce unnecessary host-to-device synchronizations (#17334)WandbLogger lazy to avoid creating artifacts when the CLI is used (#17573)*_step methods in strategies by removing the _LightningModuleWrapperBase wrapper module (#17531)wandb versions older than 0.12.0 in WandbLogger (#17876)LightningModule.setup(), the self.device now returns the device the module will be placed on instead of cpu (#18021)wandb version for WandbLogger from 0.12.0 to 0.12.10 (#18171)FSDPStrategy now loads checkpoints after the configure_model/configure_sharded_model hook (#18358)FSDPStrategy.load_optimizer_state_dict and FSDPStrategy.load_model_state_dict are a no-op now (#18358)Trainer.num_val_batches, Trainer.num_test_batches and Trainer.num_sanity_val_batches now return a list of sizes per dataloader instead of a single integer (#18441)*_step(dataloader_iter) flavor now no longer takes the batch_idx in the signature (#18390)next(dataloader_iter) now returns a triplet (batch, batch_idx, dataloader_idx) (#18390)next(combined_loader) now returns a triplet (batch, batch_idx, dataloader_idx) (#18390)devices="auto" (default) (#18291)batch_idx argument optional in validation_step, test_step and predict_step to maintain consistency with training_step (#18512)TQDMProgressBar now consistently shows it/s for the speed even when the iteration time becomes larger than one second (#18593)LightningDataModule.load_from_checkpoint and LightningModule.load_from_checkpoint methods now raise an error if they are called on an instance instead of the class (#18432)torchrun in a SLURM environment; the TorchElasticEnvironment now gets chosen over the SLURMEnvironment if both are detected (#18618)OMP_NUM_THREADS to num_cpus / num_processes when launching subprocesses (e.g. when DDP is used) to avoid system overload for CPU-intensive tasks (#18677)ModelCheckpoint no longer deletes files under the save-top-k mechanism when resuming from a folder that is not the same as the current checkpoint folder (#18750)ModelCheckpoint no longer deletes the file that was passed to Trainer.fit(ckpt_path=...) (#18750)trainer.fit() twice now raises an error with strategies that spawn subprocesses through multiprocessing (ddp_spawn, xla) (#18776)ModelCheckpoint now saves a symbolic link if save_last=True and save_top_k != 0 (#18748)</details>
<details><summary>Deprecated</summary>
SingleTPUStrategy (strategy="single_tpu") in favor of SingleDeviceXLAStrategy (strategy="single_xla") (#17383)TPUAccelerator in favor of XLAAccelerator (#17383)TPUPrecisionPlugin in favor of XLAPrecisionPlugin (#17383)TPUBf16PrecisionPlugin in favor of XLABf16PrecisionPlugin (#17383)Strategy.post_training_step method (#17531)LightningModule.configure_sharded_model hook in favor of LightningModule.configure_model (#18004)LightningDoublePrecisionModule wrapper in favor of calling Trainer.precision_plugin.convert_input() (#18209)</details>
<details><summary>Removed</summary>
XLAStrategy.is_distributed property. It is always True (#17381)SingleTPUStrategy.is_distributed property. It is always False (#17381)torchdistx due to a lack of project maintenance (#17995)</details>
<details><summary>Fixed</summary>
DeepSpeedStrategy (#17531)batch_idx argument in the training_step would disable gradient accumulation (#18619)LightningModule.configure_callbacks when the callback was a subclass of an existing Trainer callback (#18508)Trainer.log_dir not returning the correct directory for the CSVLogger (#18548)self.log (#18686)CSVLogger (#18567)</details>
<a name="changelog-fabric"></a>
<details><summary>Added</summary>
XLAStrategy(sync_module_states=bool) to control whether to broadcast the parameters to all devices (#17522)FSDPStrategy (#17323)_FabricModule that bypass the strategy-specific wrappers (#17424)Fabric.init_tensor() context manager to instantiate tensors efficiently directly on device and dtype (#17488)Fabric.init_module() context manager to instantiate large models efficiently directly on device, dtype, and with sharding support (#17462)
torch.float32, torch.float64, torch.float16, or torch.bfloat16) depending on the 'true' precision choice in Fabric(precision='32-true'|'64-true'|'16-true'|'bf16-true')Fabric.init_module(empty_init=True) for checkpoint loading (#17627)Fabric.init_module(empty_init=True) in FSDP (#18122)lightning.fabric.plugins.Precision.module_init_context() and lightning.fabric.strategies.Strategy.module_init_context() context managers to control model and tensor instantiation (#17462)lightning.fabric.strategies.Strategy.tensor_init_context() context manager to instantiate tensors efficiently directly on device and dtype (#17607)Fabric(precision="16-true"|"bf16-true") (#17287)Fabric(precision="transformer-engine") using Nvidia's Transformer Engine (#17597)Fabric(plugins=BitsandbytesPrecision()) using bitsandbytes (#18655).launch() when it is required (#17570)FSDPStrategy(state_dict_type="full"|"sharded") (#17526)Fabric.call (#17874)Fabric.load(..., strict=True|False) to enable non-strict loading of partial checkpoint state (#17645)Fabric.save(..., filter=...) to enable saving a partial checkpoint state (#17845)xla_model.mark_step() before saving checkpoints with XLA (#17882)xla_model.mark_step() after optimizer.step() with XLA (#17883)FSDPStrategy(activation_checkpointing_policy=...) to customize the layer policy for automatic activation checkpointing (requires torch>=2.1) (#18045)torch.distributed.fsdp.ShardingStrategy via string in FSDPStrategy (#18087)Fabric.load_raw() for loading raw PyTorch state dict checkpoints for model or optimizer objects (#18049)XLAStrategy (#18194)devices and num_nodes when running with SLURM or TorchElastic (#18292)lightning.fabric.utilities.suggested_max_num_workers to assist with setting a good value in distributed settings (#18591)lightning.fabric.utilities.is_shared_filesystem utility function to automatically check whether the filesystem is shared between machines (#18586).load_state_dict(..., assign=True|False) on Fabric-wrapped modules in PyTorch 2.1 or newer (#18690)</details>
<details><summary>Changed</summary>
devices="auto" (default) (#18291)torchrun in a SLURM environment; the TorchElasticEnvironment now gets chosen over the SLURMEnvironment if both are detected (#18618)OMP_NUM_THREADS to num_cpus / num_processes when launching subprocesses (e.g. when DDP is used) to avoid system overload for CPU-intensive tasks (#18677)</details>
<details><summary>Deprecated</summary>
DDPStrategy.is_distributed property. This strategy is distributed by definition (#17381)SingleTPUStrategy (strategy="single_tpu") in favor of SingleDeviceXLAStrategy (strategy="single_xla") (#17383)TPUAccelerator in favor of XLAAccelerator (#17383)TPUPrecision in favor of XLAPrecision (#17383)TPUBf16Precision in favor of XLABf16Precision (#17383)</details>
<details><summary>Removed</summary>
Fabric.run or using fabric.launch(fn). This only impacts FSDP and DeepSpeed strategy users. Please instantiate your module under the newly added fabric.init_module context manager (#17832)checkpoint_io argument from the FSDPStrategy (#18192)</details>
<details><summary>Fixed</summary>
.launch() when using the DP-strategy (strategy="dp") (#17931)find_usable_cuda_devices(0) incorrectly returning a list of devices (#18722)CSVLogger (#18567)</details>
<a name="changelog-app"></a>
<details><summary>Added</summary>
gradio components with lightning colors (#17054)</details>
<details><summary>Changed</summary>
LocalSourceCodeDir cache_location to not use home in some certain cases (#17491)</details>
<details><summary>Removed</summary>
</details>
Full commit list: https://github.com/Lightning-AI/lightning/compare/2.0.0...2.1.0
<a name="contributors"></a>
@adamjstewart @akreuzer @ethanwharris @dmitsf @lantiga @nicolai86 @pl-ghost @carmocca @awaelchli @justusschock @edenlightning @belerico @lightningforever @nisheethlahoti @tchaton @yurijmikhalevich @mauvilsa @rlizzo @rusmux @yhl48 @Liyang90 @jerome-habana @JustinGoheen @Borda @speediedan @SkafteNicki @dcfidalgo
@saryazdi @parambharat @kshitij12345 @woqidaideshi @colehawkins @md-121 @gkroiz @idc9 @BoringDonut @OmerShubi @ishandutta0098 @ryan597 @leng-yue @alicanb @One-sixth @santurini @SpirinEgor @KogaiIrina @shanmugamr1992 @janeyx99 @asmith26 @dingusagar @AleksanderWWW @strawberrypie @solyaH @kaczmarj @voidful @water-vapor @bkiat1123 @rhiga2 @baskrahmer @felipewhitaker @mukhery @Quasar-Kim @robieta @one-matrix @jere357 @schmidt-ai @schuhschuh @anio @rjarun8 @callumhay @minhlong94 @klieret @giorgioskij @shihaoyin @JonathanRayner @NripeshN @marcimarc1 @bilelomrani1 @NikolasWolke @0x404 @quintenroets @Borodin @amorehead @SebastianGer @ioangatop @Tribhuvan0 @f0k @sameertantry @kwsp @nik777 @matsumotosan
When Chuck Norris trains a neural network, it not only learns, but it also gains the ability to defend itself from adversarial attacks by roundhouse kicking them into submission.
Nothing published for this version
Nothing published for this version
Nothing published for this version
Replace LightningClient with import from lightning_cloud
_FabricOptimizer.state to remain outdated after loading with load_state_dict (#18488)log_model parameter in WandbLogger via the LightningCLI (#18458)v_num in the progress bar when running with Trainer(fast_dev_run=True) (#18491)UnboundLocalError when running with python -O (#18496)Trainer(fast_dev_run=True) (#18550)@awaelchli, @borda, @justusschock, @SebastianGer
If we forgot someone due to not matching commit email with GitHub account, let us know :]
Remove _handle_is_headless calls in app run loop
_handle_is_headless calls in app run loop (#18362)strategy="ddp_spawn" and accelerator="cpu"; this has a necessary memory impact, as parameters are replicated for each process now (#18238)fabric.no_backward_sync with XLA strategies (#17761)param_dtype training (16-mixed, bf16-mixed and 32-true configurations) to avoid FSDP assertion errors with PyTorch < 2.0 (#18278)param_dtype training (16-mixed and bf16-mixed configurations) to avoid FSDP assertion errors with PyTorch < 2.0 (#18278)experiment property defined (#18093)MLFlowLogger for logging artifacts to the MLFlow server (#18395)iter() call to dataloader when checking dataloading configuration (#18415)strategy="ddp_spawn" and accelerator="cpu"; this has a necessary memory impact, as parameters are replicated for each process now (#18238)fetcher.done with dataloader_iter (#18376)@awaelchli, @Borda, @carmocca, @quintenroets, @rlizzo, @speediedan, @tchaton
If we forgot someone due to not matching commit email with GitHub account, let us know :]
Removed the top-level import lightning.pdb; import lightning.app.pdb instead
lightning.pdb; import lightning.app.pdb instead (#18177)Fabric.all_reduce() not performing an inplace operation for all backends consistently (#18235)LightningOptimizer.refresh() to update the __dict__ in case the optimizer it wraps has changed its internal state (#18280)Missing folder exception when using a Google Storage URL as a default_root_dir (#18088)None) (#18267)LightningOptimizer wrapper returned by LightningModule.optimizers() have different internal state than the optimizer it wraps (#18280)@0x404, @awaelchli, @bilelomrani1, @borda, @ethanwharris, @nisheethlahoti
If we forgot someone due to not matching commit email with GitHub account, let us know :]
Fixed handling a None request in the file orchestration queue
None request in the file orchestration queue (#18111)TensorBoardLogger.log_graph not unwrapping the _FabricModule (#17844)LightningCLI not saving correctly seed_everything when run=True and seed_everything=True (#18056)_FaultTolerantMode when loading an old checkpoint that pickled the enum (#18094)@awaelchli, @lantiga, @mauvilsa, @shihaoyin
If we forgot someone due to not matching commit email with GitHub account, let us know :]
TensorBoardLogger.log_graph not unwrapping the _FabricModule (#17844)Added exponential backoff to HTTPQueue put
torch.set_float32_matmul_precision info message to show multiple times (#17960)Fabric.load() is called after Fabric.setup() (#17997)WandbLogger (#17818)torch.set_float32_matmul_precision info message to show multiple times (#17960)map_location argument for the LightningDataModule.load_from_checkpoint function (#17950)neptune-client (#17939)@anio, @awaelchli, @borda, @ethanwharris, @lantiga, @nicolai86, @rjarun8, @schmidt-ai, @schuhschuh, @wouterzwerink, @yurijmikhalevich
If we forgot someone due to not matching commit email with GitHub account, let us know :]
bumped several dependencies to address security vulnerabilities.
plugins.precision.MixedPrecision (#17687)NeptuneLogger (#16761):
log() method with append() and extend().Handler as an alternative to Run for the run argument. This means that you can call it NeptuneLogger(run=run["some/namespace"]) to log everything to the some/namespace/ location of the run.plugins.precision.MixedPrecisionPlugin (#17687)LightningModule.load_from_checkpoint when there is an extra state (#17812)@akreuzer, @awaelchli, @borda, @jerome-habana, @kshitij12345
If we forgot someone due to not matching commit email with GitHub account, let us know :]
Added the property LightningWork.public_ip that exposes the public IP of the LightningWork instance
LightningWork.public_ip that exposes the public IP of the LightningWork instance (#17742)LightningWork.internal_ip that was mistakenly exposing the public IP instead; now exposes the private/internal IP address (#17742)Callback registration through entry points (#17756)CSVLogger (#17139)CombinedLoader only starts DataLoader workers when necessary when operating in sequential mode (#17639)WandbLogger no longer flattens dictionaries in the hyperparameters logged to the dashboard (#17574)CSVLogger (#17139)ModelCheckpoint contained metrics that were substrings of each other (#17610)WandbLogger ignoring the WANDB_PROJECT environment variable (#16222)@adamjstewart, @AleksanderWWW, @awaelchli, @baskrahmer, @bkiat1123, @borda, @carmocca, @ethanwharris, @leng-yue, @lightningforever, @manangoel99, @mukhery, @Quasar-Kim, @water-vapor, @yurijmikhalevich
If we forgot someone due to not matching commit email with GitHub account, let us know :]
Resolved Lightning App with remote storage
AppState, streamlit example (#17452)LightningModule.*_step methods bypassing the DDP/FSDP wrapper (#17424)Fabric.setup() when the model has no parameters (#17441)Model.load_from_checkpoint("checkpoint.ckpt", map_location=map_location) would always return model on CPU (#17308)num_nodes not to be set correctly for FSDPStrategy (#17438)@awaelchli, @borda, @carmocca, @ethanwharris, @ryan597, @tchaton
If we forgot someone due to not matching commit email with GitHub account, let us know :]
Fix frontend hosts when running with multi-process in the cloud
No changes.
is_picklable function more robust (#17270)@eng-yue @ethanwharris @Borda @awaelchli @carmocca
If we forgot someone due to not matching commit email with GitHub account, let us know :]
Generalized Optimizer validation to accommodate both FSDP 1.x and 2.x
No changes
Optimizer validation to accommodate both FSDP 1.x and 2.x (#16733)LightningModule no longer pickles the Trainer (#17133)Optimizer validation to accommodate both FSDP 1.x and 2.x (#16733)torch.inference_mode with torch.compile in PyTorch 2.0 (#17215)ModelCheckpoint(save_top_k>0) is used (#17121)rich that prevented Lightning to be imported in Google Colab (#17156)DeepSpeedStrategy (#16973)torch.compile would fail when logging to WandB (#17216)@Borda @williamFalcon @lightningforever @adamjstewart @carmocca @tshu-w @saryazdi @parambharat @awaelchli @colehawkins @woqidaideshi @md-121 @yhl48 @gkroiz @idc9 @speediedan
If we forgot someone due to not matching commit email with GitHub account, let us know :]
…and make considerable changes to reduce the backwards incompatible changes in the future. To alleviate this, we will commit to continue supporting the…
Lightning AI is excited to announce the release of Lightning 2.0 :zap:
Over the last couple of years PyTorch Lightning has become the preferred deep learning framework for researchers and ML developers around the world, with close to 50 million downloads and 18k OSS projects, from top universities to leading labs.
With the help of over 800 contributors, we have added many features and functionalities to make it the most complete research toolkit possible, but some of these changes also introduced issues:
To make the research experience better, we are introducing 2.0:
<a name="highlights"></a>
torch.compileLightning 2.0 is best friends with PyTorch 2.0. You can torch.compile your LightningModules now!
import torch
import lightning as L
model = LitModel()
# This will compile forward and {training,validation,test,predict}_step
compiled_model = torch.compile(model)
trainer = L.Trainer()
trainer.fit(compiled_model)
PyTorch reports that on average, "models runs 43% faster in training on an NVIDIA A100 GPU. At Float32 precision, it runs 21% faster on average and at AMP Precision it runs 51% faster on average" (source). If you want to learn more about torch.compile and how such speedups can be achieved, read the official PyTorch 2.0 blog post.
The Trainer now chooses accelerator="auto", strategy="auto", devices="auto" as defaults. This automatically detects the best hardware on your system (TPUs, GPUs, Apple Silicon, etc.) and chooses as many devices as are available.
import lightning as L
# Selects accelerator, devices and strategy automatically!
trainer = L.Trainer()
# Same as:
trainer = L.Trainer(accelerator="auto", strategy="auto", devices="auto")
For example, on a 8-GPU server, this will implicitly select Trainer(accelerator="cuda", strategy="ddp", devices=8).
Previously, the Trainer supported DataLoader-like iterables. However, with this release, users can now work with any iterable that implements the Python iterable definition. This includes custom data structures, such as user-defined classes and generators, as well as built-in Python objects.
To use this new feature, return any iterable (or collection of iterables) from the dataloader hooks.
def train_dataloader(self):
return DataLoader(...)
return list(range(1000))
# pass loaders as a dict. This will create batches like this:
# {'a': batch_from_loader_a, 'b': batch_from_loader_b}
return {"a": DataLoader(...), "b": DataLoader(...)}
# pass loaders as list. This will create batches like this:
# [batch_from_dl_1, batch_from_dl_2]
return [DataLoader(...), DataLoader(...)]
# arbitrary nesting
# {'a': [batch_from_dl_1, batch_from_dl_2], 'b': [batch_from_dl_3, batch_from_dl_4]}
return {"a": [dl1, dl2], "b": [dl3, dl4]}
Read our data section for more information.
Lightning automatically collates the batches from multiple iterables based on a "mode". This is done with our newly revamped CombinedLoader class.
from lightning.pytorch.utilities import CombinedLoader
iterables = {"a": DataLoader(), "b": DataLoader()}
# Lightning uses this under the hood, but this way you can change the "mode"
combined_loader = CombinedLoader(iterables, mode="min_size")
model = ...
trainer = Trainer()
trainer.fit(model, combined_loader)
The following modes are supported:
min_size: stops after the shortest iterable (the one with the lowest number of items) is done.max_size_cycle: stops after the longest iterable (the one with most items) is done, while cycling through the rest of the iterables.max_size: stops after the longest iterable (the one with most items) is done, while returning None for the exhausted iterables.sequential: completely consumes ecah iterable sequentially, and returns a triplet (data, idx, iterable_idx)If you have a need for a different "mode", feel free to open a feature request! Adding new modes is now very simplified. These improvements also allowed us to simplify the trainer's loops by abstracting this logic inside the CombinedLoader.
A new Trainer argument Trainer(barebones=...) was added (default is False) to disable all features that may impact the raw speed of the training loop. This allows users to quickly and fairily compare the runtime of a Lightning script with a raw PyTorch script.
This is how you enable it:
import lightning as L
# Default: False
trainer = L.Trainer(barebones=True)
A message informs about the changed settings:
You are running in `Trainer(barebones=True)` mode. All features that may impact raw speed have been disabled to facilitate analyzing the Trainer overhead. Specifically, the following features are deactivated:
- Checkpointing: `Trainer(enable_checkpointing=True)`
- Progress bar: `Trainer(enable_progress_bar=True)`
- Model summary: `Trainer(enable_model_summary=True)`
- Logging: `Trainer(logger=True)`, `Trainer(log_every_n_steps>0)`, `LightningModule.log(...)`, `LightningModule.log_dict(...)`
- Sanity checking: `Trainer(num_sanity_val_steps>0)`
- Development run: `Trainer(fast_dev_run=True)`
- Anomaly detection: `Trainer(detect_anomaly=True)`
- Profiling: `Trainer(profiler=...)`
Tip: This feature is also very useful for unit testing!
Based on feedback from users, we decided to separate the training progress bar from the validation bar. This greatly improves the time estimates (since validation is usually faster) and resolves confusion around the total batches being processed in an epoch.
This is how the bar looked in versions before 2.0:
Epoch 3: 21%|██ | 28/128 [00:36<01:32, 23.12it/s, loss=0.163]
Validation DataLoader 0: 38%|███ | 12/32 [00:12<00:20, 1.01s/it]
Note how the total batches (128) is the sum of the training batches (32) and the three validation runs (3 x 32). And this is how the progress bar looks like now:
Epoch 3: 50%|█████ | 16/32 [00:36<01:32, 23.12it/s]
Validation DataLoader 0: 38%|███ | 12/32 [00:12<00:20, 1.01s/it]
Note how the batch counts are now separate. The training progress bar pauses until validation is completed.
<a name="highlights-fabric"></a>
Lightning 2.0 is the official release for Lightning Fabric :tada:
<img src="https://pl-public-data.s3.amazonaws.com/assets_lightning/fabric/PyTorch-to-Fabric-Spectrum-2.svg" alt="Fabric spans across a large spectrum - from raw PyTorch all the way to high-level PyTorch Lightning" style="margin-top: 50px; margin-bottom: 50px" />
Fabric is the fast and lightweight way to scale PyTorch models without boilerplate code.
:open_book: Go to Fabric documentation :open_book:
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, Dataset
+ from lightning.fabric import Fabric
class PyTorchModel(nn.Module):
...
class PyTorchDataset(Dataset):
...
+ fabric = 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()
<a name="bc-changes"></a>
This section outlines notable changes that are not backward compatible with previous versions. The full list of changes and removals can be found in the Full Changelog below.
Since 2.0 is a major release, we took the opportunity to take our APIs to the next level and make considerable changes to reduce the backwards incompatible changes in the future. To alleviate this, we will commit to continue supporting the 1.9.x line of releases by doing bug-fix releases with any important fixes that are necessary.
<a name="bc-changes-pytorch"></a>
*_epoch_end hooks were removed (#16520)Since the very beginning of Lightning, the LightningModule offered the convenient *_epoch_end end hooks in which users could reduce metrics collected across the entire epoch to log them. For this to work, Lightning had to store all the outputs returned from the *_step methods internally to be able to send them to the *_epoch_end hook.
This "silent" accumulation of memory lead to many users scratching their head when they found that after training for 20 hours the epoch crashed randomly with an out-of-memory error. This also meant that avoiding this behaviour required code changes, as simply overriding this hook would force this behaviour, regardless of whether you used the outputs. This was exacerbated by users not knowing the difference between the on_train_epoch_end (does not store outputs) and training_epoch_end (does store outputs) hooks.
Based on this feedback, we decided to remove this mechanism completely. Lightning 2.0 favors simplicity and speed over convenience features.
import lightning as L
class LitModel(L.LightningModule):
def training_step(self, batch, batch_idx):
...
return {"loss": loss, "banana": banana}
# `outputs` is a list of all bananas returned in the epoch
def training_epoch_end(self, outputs):
avg_banana = torch.cat(out["banana"] for out in outputs).mean()
import lightning as L
class LitModel(L.LightningModule):
def __init__(self):
super().__init__()
# 1. Create a list to hold the outputs of `*_step`
self.bananas = []
def training_step(self, batch, batch_idx):
...
# 2. Add the outputs to the list
# You should be aware of the implications on memory usage
self.bananas.append(banana)
return loss
# 3. Rename the hook to `on_*_epoch_end`
def on_train_epoch_end(self):
# 4. Do something with all outputs
avg_banana = torch.cat(self.bananas).mean()
# Don't forget to clear the memory for the next epoch!
self.bananas.clear()
The new way makes the user to manage their own list of outputs, and with it also the responsibility of managing the memory correctly. If an OOM happens, the user can identify the issue in their own code rather than having to guess what the Trainer does.
The prediction epoch end hook is a special case of this change where the outputs argument has been removed from on_predict_epoch_end(trainer, pl_module) but can still be accessed via the attribute trainer.predict_loop.outputs
You can find a migration guide for this change in this PR's description.
Lightning 2.0 removed the special optimizer loop that would manage multiple optimizers in the automatic mode in favor of an overall easier to understand and debuggable training loop code. Training with multiple optimizers is now restricted to the "manual optimization mode":
import lightning as L
class LitModel(L.LightningModule):
def configure_optimizers(self):
...
return optimizer1, optimizer2
def training_step(self, batch, batch_idx, optimizer_idx)
if optimizer_idx == 0:
...
return loss0
if optimizer_idx == 1:
...
return loss1
import lightning as L
class LitModel(L. LightningModule):
def __init__(self):
super().__init__()
# 1. Switch to manual optimization mode
self.automatic_optimization = False
def configure_optimizers(self):
...
# 2. Return multiple optimizers, same as before
return optimizer1, optimizer2
# 3. Remove the `optimizer_idx` argument from `training_step`
def training_step(self, batch, batch_idx)
# 4. Grab all optimizers you want to work with
optimizer1, optimizer2 = self.optimizers()
...
# 5. Perform backward manually, step optimzers, etc.
self.manual_backward(loss0)
optimizer1.step()
optimizer1.zero_grad()
# 6. In manual optimization, you don't need to return anything
return None
You may also find the utility methods self.toggle_optimizer() and self.untoggle_optimizer() useful if you need to restrict parameters that require gradients to a specific optimizer. For a complete example, see our simple GAN implementation.
Similar to the multi-optimizer loop mentioned above, truncated backpropagation through time (TBPTT) was a loop that added a lot of complexity to the Trainer. Over time, TBPTT has fallen out of fashion and today the demand from users is so low that we decided to drop special support in the framework in favor of simplifying the Trainer. TBPTT can still be done in manual optimization.
import lightning as L
class LitModel(L.LightningModule):
def __init__(self):
super().__init__()
self.truncated_bptt_steps = 10
self.my_rnn = ...
def training_step(self, batch, batch_idx, hiddens):
...
loss, hiddens = self.my_rnn(..., hiddens)
...
return loss, hiddens
import lightning as L
class LitModel(L.LightningModule):
def __init__(self):
super().__init__()
# 1. Switch to manual optimization
self.automatic_optimization = False
self.truncated_bptt_steps = 10
self.my_rnn = ...
# 2. Remove the `hiddens` argument
def training_step(self, batch, batch_idx):
# 3. Split the batch in chunks along the time dimension
split_batches = split_batch(batch, self.truncated_bptt_steps)
hiddens = ... # 3. Choose the initial hidden state
for split_batch in range(split_batches):
# 4. Perform the optimization in a loop
loss, hiddens = self.my_rnn(split_batch, hiddens)
self.backward(loss)
optimizer.step()
optimizer.zero_grad()
# 5. "Truncate"
hiddens = hiddens.detach()
# 6. Remove the return of `hiddens`
# Returning loss in manual optimization is not needed
return None
To simplify the Trainer interface and with the goal of simpler iterable support inside the Trainer, we removed theTrainer(multiple_trainloader_mode=...) argument. The mode is now agnostic to the trainer stage ("train" previously) and it's easier to debug and understand for the user as the logic is all encapsulated in the CombinedLoader
import lightning as L
class LitModel(L.LightningModule):
def train_dataloader(self):
return [DataLoader(), DataLoader()]
model = LitModel()
trainer = Trainer(multiple_trainloader_mode="min_size")
trainer.fit(model)
import lightning as L
from lightning.pytorch.utilities import CombinedLoader
class LitModel(L.LightningModule):
def train_dataloader(self):
iterables = [DataLoader(), DataLoader()]
return CombinedLoader(iterables, mode="min_size")
model = LitModel()
trainer = Trainer()
trainer.fit(model)
Related to this, we cleaned up which hooks need the dataloader_idx as an input argument. Now it's only required if you use multiple dataloaders. Don't worry, the Trainer will automatically check if it's required for you and tell you about it.
In the case that you require access to the DataLoader or Dataset objects, iterables for each step can be accessed
via the trainer properties Trainer.train_dataloader, Trainer.val_dataloaders, Trainer.test_dataloaders, and Trainer.predict_dataloaders.
These properties will match exactly what was returned in your *_dataloader hooks or passed to the Trainer, meaning that if you returned a dictionary of dataloaders, these will return a dictionary of dataloaders. This wasn't the case in before 2.0:
# Passing 1 dataloader per stage
train_dataset = trainer.train_dataloader.loaders.dataset
val_dataset = trainer.val_dataloaders[0].dataset
# Passing 2 dataloaders per stage
train_dataset = trainer.train_dataloader.loaders[0].dataset
val_dataset = trainer.val_dataloaders[0].dataset
# Passing 1 dataloader per stage
train_dataset = trainer.train_dataloader.dataset
val_dataset = trainer.val_dataloaders.dataset
# Passing 2 dataloaders per stage
train_dataset = trainer.train_dataloader[0].dataset
val_dataset = trainer.val_dataloaders[0].dataset
The Tuner and Trainer are no longer together. The two Trainer arguments Trainer(auto_lr_find=..., auto_scale_batch_size=...) and the Trainer.tune() method were removed to make the Trainer leaner and easier to work with.
import lightning as L
# Tune learning rate
trainer = L.Trainer(auto_lr_find=True)
trainer.tune(model)
# Tune batch size
trainer = L.Trainer(auto_scale_batch_size=True)
trainer.tune(model)
# Fit using tuned settings
trainer.fit(model)
import lightning as L
# 1. Create the Trainer
trainer = L.Trainer()
# 2. Create the Tuner
tuner = L.pytorch.tuner.Tuner(trainer)
# 3. Tune learning rate
tuner.lr_find(...)
# 4. Tune batch size
tuner.scale_batch_size(...)
# Fit using tuned settings
trainer.fit(model)
You can find more documentation about the tuner here.
In Lightning 1.6.0, we simplified the Trainer's signature by collapsing the four accelerator-specific device arguments into a single one called Trainer(devices=...). In 2.0, we are now dropping support for the old Trainer(gpus=..., tpu_cores=..., ipus=..., num_processes=...).
import lightning as L
# Multiple CPU processes
trainer = L.Trainer(num_processes=2)
# GPU selection
trainer = L.Trainer(gpus=4)
# TPU core selection
trainer = L.Trainer(tpu_cores=[1])
# Graphcore IPU devices
trainer = L.Trainer(ipus=1)
import lightning as L
# Select devices and accelerator separately
trainer = L.Trainer(accelerator="cpu", devices=2)
trainer = L.Trainer(accelerator="gpu", devices=4)
trainer = L.Trainer(accelerator="tpu", devices=[1])
trainer = L.Trainer(accelerator="ipu", devices=1)
# Or let Lightning detect the accelerator automatically
trainer = L.Trainer(accelerator="auto", devices=2)
# `accelerator="auto"` is the default
trainer = L.Trainer(devices=2)
In addition, the Trainer(auto_select_gpus=...) also got removed (#16184). This was a problematic feature that was not well documented, often misundertood, and lead to DDP stalls due to race conditions in the device selection. It is recommended to use the devices="auto" instead or, if the exact previous behavior is desired, use the utility function.
import lightning as L
trainer = L.Trainer(auto_select_gpus=True, devices=2)
import lightning as L
# Recommended
trainer = L.Trainer(devices="auto")
# Alternatively, use the utility function (with care!)
from lightning.pytorch.accelerators import find_usable_cuda_devices
trainer = L.Trainer(devices=find_usable_cuda_devices(2), strategy="ddp_spawn")
PyTorch and Lightning have discouraged the use of the old-style DataParallel (DP) training for more than two years now. This method of multi-GPU training is slow and has many limitations that impact users, even more so in Lightning. Since DP has fallen out of fashion, and other strategies (DDP, DeepSpeed, etc.) have emerged without the same limitations, Lightning is now dropping DP completely.
import lightning as L
# Enables DP, but with many limitations
trainer = L.Trainer(strategy="dp", devices=8)
import lightning as L
# DDP is recommended for multi-GPU training
trainer = L.Trainer(strategy="ddp", devices=8)
# If devices > 1, it selects ddp for you
trainer = L.Trainer(devices=8)
# This is the same
trainer = L.Trainer(strategy="auto", devices=8)
Note that the DDP strategy now gets selected automatically when multiple devices are selected.
In Jupyter notebooks (Google Colab, Kaggle, etc.), Lightning will automatically select a fork-based DDP strategy (strategy="ddp_notebook").
By extension, the LightningModule no longer has the training_step_end(), validation_step_end() and test_step_end() hooks because they were only used for reducing the outputs in DP. If you have other code that needs to run at the end of the step, you can migrate it to the corresponding *_batch_end hook for example.
Loop customization was our attempt to make the Trainer more customizable, hackable, and extensible. The modularity it brought had many pros, but it turned out that many users didn't need customization and favored readability and simplicity instead. Hence, the concept of loop customization was completely removed from 2.0.
For users who like build custom training loops, there is now a new paradigm with Lightning Fabric and "Build your own Trainer" (BYOT). Check out the Fabric documentation and the super hackable, 500-lines trainer template.
Based on feedback, we decided to make the names for precision backends in Trainer(precision=...) clearer and less ambiguous. For example, the previous notation Trainer(precision=16) (which is still allowed to be used) suggested to some users that all of the weights and tensors would be stored in a float16 format, which is not true. To solve this misunderstanding, we now distinguish these modes with "true" and "mixed" suffixes in the names:
| Recommended value | Short form |
|---|---|
| "64-true" | "64", 64 |
| "32-true" | "32", 32 |
| "16-mixed" | "16", 16 |
| "bf16-mixed" | "bf16" |
All documentation and examples are now recommending the new, less ambiguous names.
In addition to reworking the precision settings, we removed the NVIDIA/Apex integration which got deprecated in 1.9.0. Apex itself has deprecated the mixed precision module and recommends the native torch.amp module in PyTorch.
For Lightning, this means you should switch to Trainer(precision="16-mixed").
import lightning as L
# This required installing nvidia/apex
trainer = L.Trainer(amp_backend="apex", amp_level="O2")
import lightning as L
# Rely on PyTorch's native mixed precision
trainer = L.Trainer(precision="16-mixed")
With the recent annoucement that FSDP becomes production ready in PyTorch 2.0, we are dropping the support for the experimental Fairscale version of FSDP and go all in on the native implementation instead.
import lightning as L
# Short-hand names (fairscale)
trainer = L.Trainer(strategy="ddp_sharded" | "ddp_fully_sharded" | "fsdp")
# Shorthand names (native PyTorch)
trainer = L.Trainer(strategy="fsdp_native" | "fsdp_native_full_shard_offload")
# Or using the strategy instance
from lightning.pytorch.strategies import DDPShardedStrategy, DDPFullyShardedStrategy
trainer = L.Trainer(strategy=DDPShardedStrategy(...))
trainer = L.Trainer(strategy=DDPFullyShardedStrategy(...))
import lightning as L
# Shorthand names (native PyTorch)
trainer = L.Trainer(strategy="fsdp" | "fsdp_cpu_offload")
# Or using the strategy instance
from lightning.pytorch.strategies import FSDPStrategy
trainer = L.Trainer(strategy=FSDPStrategy(...))
Resuming from a checkpoint is no longer done by specifying the filename in the Trainer constructor. The Trainer(resume_from_checkpoint=...) argument was renamed and moved to the individual Trainer methods.
import lightning as L
trainer = L.Trainer(resume_from_checkpoint="/path/to/checkpoint.ckpt")
trainer.fit(model)
trainer.test(model)
...
import lightning as L
# 1. Remove `resume_from_checkpoint` from the Trainer
trainer = L.Trainer()
# 2. Add the path to any of the trainer methods you want to run
trainer.fit(model, ckpt_path="/path/to/checkpoint.ckpt")
trainer.test(model, ckpt_path="/path/to/checkpoint.ckpt")
...
We also added support for setting the checkpoint path statefully:
import lightning as L
trainer = L.Trainer()
# set the checkpoint path with a setter
trainer.ckpt_path = "/path/to/checkpoint.ckpt"
trainer.fit(model)
# remember to clear it before continuing
trainer.ckpt_path = None
trainer.test()
...
In previous versions of Lightning, the Trainer used to automatically compute a running mean of the training loss and log it to the progress bar. We are removing this feature to favor speed over automation. The loss will no longer appear in the progress bar unless the user explicitly adds it.
Epoch 8: 53%|█████ | 17/32 [5.13/s, v_num=2, loss=0.5643]
Epoch 8: 53%|█████ | 17/32 [5.13/s, v_num=2]
def training_step(self, batch, batch_idx):
...
# Add this if you need the loss to be displayed in the progress bar
self.log("loss", loss, prog_bar=True, on_step=True)
return loss
In previous versions of Lightning, the Trainer.add_argparse_args and Trainer.from_argparse_args utility functions helped the user construct a commandline interface using Python's argparse library. Over time, new functionality was added to the Trainer that allowed most of the frequently used arguments to accept many different types. This made the arg-parsing utilities very brittle and limited in its scope. Ultimately, we decided to remove it in 2.0 in favor of more robust solutions like LightningCLI or third-party commandline tools (Click, Hydra, etc.).
import lightning as L
import argparse
parser = argparse.ArgumentParser()
parser = Trainer.add_argparse_args(parser)
args = parser.parse_args()
trainer = L.Trainer.from_argparse_args(args)
...
Example using the LightningCLI:
from lightning.pytorch.cli import LightningCLI
LightningCLI(MyLitModel, MyLitDataModule)
Run it with python train.py fit --trainer.max_epochs=2 for example.
Alternatively, you can add the argparse arguments you want manually:
import lightning as L
import argparse
parser = argparse.ArgumentParser()
# 1. Add the arguments you need
parser.add_argument("--accelerator", type=str, default="cuda")
args = parser.parse_args()
# 2. Pass them into the Trainer
trainer = L.Trainer(**vars(args))
...
Until now, automatic tracking of the gradient norm was possible through the Trainer(track_grad_norm=...) argument. This functionality has now moved to a utility function and can be easily added in two lines of code to the LightningModule.
import lightning as L
trainer = L.Trainer(track_grad_norm=2)
# Optionally customize logging in LightningModule
def log_grad_norm(self, grad_norm_dict):
self.log(...)
from lightning.pytorch.utilities import grad_norm
# 1. Override LightningModule hook
def on_before_optimizer_step(self, optimizer):
# 2. Inspect the (unscaled) gradients here
self.log_dict(grad_norm(self, norm_type=2))
This enables users to customize how the gradient norm is computed and logged, without needing to wrangle with the Trainer or override the log_grad_norm hook.
find_unused_parameters (#16611)When your model doesn't use all parameters in the forward-backward pass, PyTorch's DDP wrapper will freak out and inform you about it with an error. For this reason, and for the longest time, Lightning has set find_unused_parameters=True so that this error can be avoided. However, depending on the model, this can have significant performance impact (slowing down your training). With Lightning 2.0, we switch this flag back to find_unused_parameters=False (default) and favor speed over convenience.
import lightning as L
# Previously, you had to manually override this for speed
trainer = L.Trainer(strategy="ddp_find_unused_parameters_false")
import lightning as L
# Now, you get the best speed by default for all ddp variants
trainer = L.Trainer(strategy="ddp")
trainer = L.Trainer(strategy="ddp_spawn")
# We now have these if you need them
trainer = L.Trainer(strategy="ddp_find_unused_parameters_true")
However, it can still happen that users run into problems with unused parametes in their model. Lightning now overrides PyTorch's default error message with a custom one to help users resolve the problem.
We renamed the Trainer(replace_sampler_ddp=...) argument to Trainer(use_distributed_sampler=...) to communicate that the sampler gets created not only for the DDP strategies, but all distributed strategies that need it. Its function is still the same as before, and most users don't need to change its default value.
import lightning as L
# Hmm, I wonder, does this only apply to DDP?
trainer = L.Trainer(replace_sampler_ddp=True)
import lightning as L
# Aha! Lightning uses a distributed sampler by default, got it!
trainer = L.Trainer(use_distributed_sampler=True)
Whenever there is a new PyTorch version, or a new Python version, it is time to say good bye to the oldest one we support. With the introduction of PyTorch 2.0, we are dropping support for PyTorch 1.10 to continue our support window of the four latest versions: 1.11, 1.12, 1.13 and 2.0. Similarly, with Lightning 2.0 we support the latest three versions of Python: 3.8, 3.9, and 3.10 (3.11 is coming soon).
If you are interested which range of PyTorch or Python versions a particular Lightning version supports, see our compatibility matrix.
To simplify reading and debugging the codebase, we removed the experimental support for fault-tolerance which was under the PL_FAULT_TOLERANT_TRAINING= environment flag. We are looking at ways to re-implement this. If you are interested in this feature, don't hesistate to reach out to us or create a feature request.
Some of the features it included were now ported to stable APIs. One feature is the new callback to save a checkpoint on exception:
import lightning as L
from lightning.pytorch.callbacks import OnExceptionCheckpoint
on_exception = OnExceptionCheckpoint(".")
trainer = L.Trainer(callbacks=on_exception)
Another feature is automatic SIGTERM handling:
import lightning as L
trainer = L.Trainer()
if trainer.received_sigterm:
...
self.log()ing a dictionary (#16389)Our logging mechanism previously supported log("key", {"something": 123}) (not using log_dict). However, this added significant complexity to the implementation with little benefit, as these keys could not be monitored by our Callbacks and most logger implementations do not support this notation. If you were using this feature with a compatible logger, you can still publish data directly to the Logger using self.logger.log_metrics().
trainer.reset_*_dataloader() methods (#16726)These methods were not intended for public use as they could leave the Trainer in an undefined state. As a result, we have removed them. To achieve the same functionality, you can use the Trainer(reload_dataloaders_every_n_epochs=...) argument.
Trainer(move_metrics_to_cpu=True) argument (#16358)This flag was designed to reduce device memory allocation at the end of an epoch. However, due to design issues, it could leave non-CPU runs in a non-functional state. Since the memory savings were minimal compared to other components and users can still manually control their metrics, we decided to remove this flag.
We removed support for passing a scheduling dictionary to Trainer(accumulate_grad_batches=...). The same functionality can still be achieved by simply passing the callback in. This simplifies the Trainer and the overall validation logic.
Before:
trainer = Trainer(accumulate_grad_batches={"1": 5, "10": 3})
Now:
from lightning.pytorch.callbacks import GradientAccumulationScheduler
trainer = Trainer(callbacks=GradientAccumulationScheduler({"1": 5, "10": 3}))
<a name="bc-changes-fabric"></a>
Over time, LightningLite has evolved from a simple onboarding tool to a powerful toolbox enabling users to build performant and hackable trainers. In 1.9.0, we renamed it to Lightning Fabric and gave users early access to its new features. In 2.0, we are dropping LightningLite from the package completely.
# There were two different ways of importing Lite in <= 1.9.0
from pytorch_lighting.lite import LightningLite
from lightning.lite import LightningLite
# You had to subclass LightningLite and implement `run()`
class MyTrainer(LightningLite):
def run(self):
...
self.backward(loss)
...
lite = LightningLite(...)
lite.run()
# 1. Import Fabric directly from the lightning package
import lightning as L
# 2. Instantiate Fabric directly, without subclassing
fabric = L.Fabric(...)
# 3. Use it in your training loop
fabric.backward(loss)
...
Learn more about Fabric and what it can do in the new docs!
<a name="changelog"></a>
<a name="changelog-pytorch"></a>
<details><summary>Added</summary>
Trainer.ckpt_path = ... setter to statefully set the checkpoint path to load. This can act as a replacement for the removed Trainer(resume_from_checkpoint=...) flag (#16187)include_cuda in pytorch_lightning.utilities.seed.isolate_rng to disable managing torch.cuda's rng (#16423)Tuner.lr_find(attr_name=...) to specify custom learning rate attribute names (#16462)OnExceptionCheckpoint callback to save a checkpoint on exception (#16512)MLFlowLogger with the mlflow-skinny package (16513)Trainer.received_sigterm property to check whether a SIGTERM signal was received (#16501)kill method to launchers to kill all launched processes (#16525)find_unused_parameters=True, for example strategy="ddp_find_unused_parameters_true" (#16611)Strategy.on_exception to the strategy base interface (#16646)predict_step(dataloader_iter, batch_index) (#16726)CombinedLoader to consume multiple iterables in sequence (#16743, #16784)CombinedLoader to consume multiple iterables entirely without cycling (#16939Trainer(barebones=True) argument where all features that may impact raw speed are disabled (#16854)CSVLoggers. (#16880)DDPStrategy(start_method=...) argument, defaulting to 'popen' (#16809)</details>
<details><summary>Changed</summary>
trainer.{validate,test,predict} (#17017)ProgressBarBase to ProgressBar (#17058)Trainer now chooses accelerator="auto", strategy="auto", devices="auto" as defaults (#16847)strategy="fsdp_native" is now strategy="fsdp"strategy="fsdp_native_full_shard_offload" is now strategy="fsdp_cpu_offload"pytorch_lightning.strategies.fully_sharded_native.DDPFullyShardedNativeStrategy is now pytorch_lightning.strategies.fsdp.FSDPStrategypytorch_lightning.plugins.precision.fsdp_native_native_amp.FullyShardedNativeNativeMixedPrecisionPlugin is now pytorch_lightning.plugins.precision.fsdp.FSDPMixedPrecisionPluginpytorch_lightning.plugins.precision.native_amp is now pytorch_lightning.plugins.precision.ampNativeSyncBatchNorm is now TorchSyncBatchNormLearningRateFinder(update_attr=...) and Tuner.lr_find(update_attr=...) to True (#16462)pl.utilities.exceptions.GracefulExitException to SIGTERMException (#16501)Callback.on_train_epoch_end hook now runs after the LightningModule.on_train_epoch_end hook for instances of EarlyStopping and Checkpoint callbacks (#16567)LightningModule.{un}toggle_optimizer methods no longer accept a optimizer_idx argument to select the relevant optimizer. Instead, the optimizer object can be passed in directly (#16560)find_unused_parameters now defaults to False (#16611)accelerator="hpu" now defaults to find_unused_parameters=False (#16611)TQDMProgressBar.main_progress_bar to TQDMProgressBar.train_progress_bar (#16695)lightning.pytorch.trainer.configuration_validator.verify_loop_configurations function as protected (#17009)lightning.pytorch.utiltiies.distirbuted.register_ddp_comm_hook function as protected (#17009)lightning.pytorch.utilities.supporters.CombinedDataset as protected (#16714){Accelerator,Signal,Callback,Checkpoint,Data,Logger}Connector classes as protected (#17008)lightning.pytorch.trainer.connectors.signal_connector.HandlersCompose class as protected (#17008)CombinedLoader.loaders to CombinedLoader.iterables (#16743)Trainer(replace_sampler_ddp=...) to Trainer(use_distributed_sampler=...) (#16829)CombinedLoader class from lightning.pytorch.trainer.supporters to lightning.pytorch.combined_loader (#16819)trainer.*_dataloader properties now return what the user returned in their LightningModule.*_dataloader() hook (#16726, #16800)dataloader_idx argument is now optional for the on_{validation,test,predict}_batch_{start,end} hooks. Remove it or default it to 0 if you don't use multiple dataloaders (#16753)TPUSpawnStrategy to XLAStrategy (#16781)strategy='tpu_spawn' to strategy='xla' and strategy='tpu_spawn_debug' to strategy='xla_debug' (#16781)Trainer(strategy="ddp_spawn", ...) no longer falls back to "ddp" when a cluster environment gets detected (#16780)track_grad_norm no longer rounds the norms to 4 digits, but instead logs them at full resolution (#16877)DDPSpawnStrategy into DDPStrategy (#16809)NeptuneLogger now requires neptune>=1.0.0 (#16888)rich from 10.14.0 to 12.13.0 (#16798)lightning.pytorch.overrides.torch_distributed.broadcast_object_list function (#17011)ServableModule is now an abstract interface (#17000)psutil package is now required for CPU monitoring (#17010)</details>
<details><summary>Removed</summary>
pytorch_lightning.lite module in favor of lightning_fabric (#15953)nvidia/apex removal (#16149)
pytorch_lightning.plugins.NativeMixedPrecisionPlugin in favor of pytorch_lightning.plugins.MixedPrecisionPluginLightningModule.optimizer_step(using_native_amp=...) argumentTrainer(amp_backend=...) argumentTrainer.amp_backend propertyTrainer(amp_level=...) argumentpytorch_lightning.plugins.ApexMixedPrecisionPlugin classpytorch_lightning.utilities.enums.AMPType enumDeepSpeedPrecisionPlugin(amp_type=..., amp_level=...) argumentsTrainer(strategy='horovod') support (#16150)FairScale removal (in favor of PyTorch's FSDP implementation) (#16400)
pytorch_lightning.overrides.fairscale.LightningShardedDataParallel classpytorch_lightning.plugins.precision.fully_sharded_native_amp.FullyShardedNativeMixedPrecisionPlugin classpytorch_lightning.plugins.precision.sharded_native_amp.ShardedNativeMixedPrecisionPlugin classpytorch_lightning.strategies.fully_sharded.DDPFullyShardedStrategy (fsdp) classpytorch_lightning.strategies.sharded.DDPShardedStrategy (ddp_sharded) classpytorch_lightning.strategies.sharded_spawn.DDPSpawnShardedStrategy (ddp_sharded_spawn) classTrainer(gpus=...) argumentTrainer(tpu_cores=...) argumentTrainer(ipus=...) argumentTrainer(num_processes=...) argumentpytorch_lightning.utilities.AllGatherGrad class (#16360)resume_from_checkpoint Trainer argument (#16167)pytorch_lightning.profiler module (#16359)PL_RECONCILE_PROCESS=1) (#16204){training,validation,test}_epoch_end hooks which would retain step outputs in memory. Alternative implementations are suggested by implementing their on_*_epoch_end hooks instead (#16520)outputs argument from the on_predict_epoch_end hook. You can access them via trainer.predict_loop.predictions (#16655)PL_FAULT_TOLERANT_TRAINING environment flag (#16516, #16533)LightningCLI arguments (#16380)
save_config_filenamesave_config_overwritesave_config_multifiledescriptionenv_prefixenv_parsepl.strategies.utils.on_colab_kaggle function (#16437)Accelerator.setup_environment method (#16436)forward_module argument as required (#16386)
pl_module argument from the distributed module wrapperspytorch_lightning.overrides.base.unwrap_lightning_module functionpytorch_lightning.overrides.distributed.LightningDistributedModule classpytorch_lightning.overrides.fairscale.unwrap_lightning_module_sharded functionpytorch_lightning.overrides.fairscale.LightningDistributedModule classTrainer(auto_select_gpus=...) argumentpytorch_lightning.tuner.auto_gpu_select.{pick_single_gpu,pick_multiple_gpus} functionsLoop.replace() (#16361)Loop.connect() (#16384)trainer.{fit,validate,test,predict}_loop properties (#16384)Loop.run() implementation (#16384)lightning.pytorch.overrides.distributed.IndexBatchSamplerWrapper class is now marked as protected (#16826)DataLoaderLoop, EvaluationEpochLoop, and PredictionEpochLoop classes (#16726)trainer.reset_*_dataloader() methods in favor of Loop.setup_data() for the top-level loops (#16726)LightningModule.truncated_bptt_steps attributeLightningModule.tbptt_split_batch hookLightningModule.training_step no longer accepts a hiddens argumentpytorch_lightning.loops.batch.TrainingBatchLoopFitLoop.split_idx propertyLoggerConnector.on_train_split_start methodPL_INTER_BATCH_PARALLELISM environment flag (#16355)Trainer(move_metrics_to_cpu=True) argument (#16358)LightningModule.precision attribute (#16203)training_step loss in the progress bar. Use self.log("loss", ..., prog_bar=True) instead. (#16192)self.log() (#16389)Trainer.model setter (#16462)Trainer(multiple_trainloader_mode=...). You can use CombinedLoader(..., mode=...) directly now (#16800)lightning.pytorch.utilities.finite_checks.print_nan_gradients function (#16682)lightning.pytorch.utilities.finite_checks.detect_nan_parameters function (#16682)lightning.pytorch.utilities.parsing.flatten_dict function (#16744)lightning.pytorch.utilities.metrics.metrics_to_scalars function (#16681)lightning.pytorch.utilities.supporters.{SharedCycleIteratorState,CombinedLoaderIterator} classes (#16714)trainer.tuning property (#16379)TrainerFn.TUNING and RunningStage.TUNING enums (#16379)Trainer.tune() in favor of Tuner(trainer).{lr_find,scale_batch_size} (#16462)Trainer(auto_scale_batch_size=...) in favor of Tuner(trainer).scale_batch_size() (#16462)Trainer(auto_lr_find=...) in favor of Tuner(trainer).lr_find() (#16462)on_tpu argument from LightningModule.optimizer_step hook (#16537)using_lbfgs argument from LightningModule.optimizer_step hook (#16538)Trainer.data_parallel property. Use isinstance(trainer.strategy, ParallelStrategy) instead (#16703)Trainer.prediction_writer_callbacks property (#16759)opt_idx argument from BaseFinetuning.finetune_function callback methodopt_idx argument from Callback.on_before_optimizer_step callback methodoptimizer_idx as an optional argument in LightningModule.training_stepoptimizer_idx argument from LightningModule.on_before_optimizer_stepoptimizer_idx argument from LightningModule.configure_gradient_clippingoptimizer_idx argument from LightningModule.optimizer_stepoptimizer_idx argument from LightningModule.optimizer_zero_gradoptimizer_idx argument from LightningModule.lr_scheduler_stepLightningModule.configure_optimizersoptimizer and optimizer_idx from LightningModule.backwardoptimizer_idx argument from PrecisionPlugin.optimizer_step and all of its overrides in subclassesoptimizer_idx argument from PrecisionPlugin.{optimizer_step,backward} and all of its overrides in subclassesoptimizer_idx argument from Strategy.{optimizer_step,backward} and all of its overrides in subclassesTrainer.optimizer_frequencies attributeStrategy.dispatch (#16618)PrecisionPlugin.dispatch (#16618)LightningDataModule methods: add_argparse_args(), from_argparse_args(), parse_argparser(), get_init_arguments_and_types()default_attributes(), from_argparse_args(), parse_argparser(), match_env_arguments(), add_argparse_args()lightning.pytorch.utilities.argparse: from_argparse_args(), parse_argparser(), parse_env_variables(), get_init_arguments_and_types(), add_argparse_args()lightning.pytorch.utilities.parsing: import str_to_bool(), str_to_bool_or_int(), str_to_bool_or_str()Trainer(accumulate_grad_batches=...) (#16729)DataParallel (strategy='dp') and the LightningParallelModule-Wrapper, (#16748)lightning.pytorch.utilities.supporters.{SharedCycleIteratorState,CombinedLoaderIterator} classes (#16714)ProgressBarBase.{train_batch_idx,val_batch_idx,test_batch_idx,predict_batch_idx} properties (#16760)fit_loop.{min,max}_steps setters (#16803)Trainer(track_grad_norm=...) argument (#16745)LightningModule.log_grad_norm() hook method (#16745)QuantizationAwareTraining callback (#16750)ColossalAIStrategy and ColossalAIPrecisionPlugin in favor of the new lightning-colossalai package (#16757, #16778)training_step_end, validation_step_end, and test_step_end hooks from the LightningModule in favor of the *_batch_end hooks (#16791)lightning.pytorch.strategies.DDPSpawnStrategy in favor of DDPStrategy(start_method='spawn') (merged both classes) (#16809)ShardedTensor state dict hooks in LightningModule.__init__ with torch>=2.1 (#16892)lightning.pytorch.core.saving.ModelIO class interface (#16999)lightning.pytorch.utilities.memory.get_model_size_mb function (#17001)</details>
<details><summary>Fixed</summary>
DistributedSampler.set_epoch wasn't getting called during trainer.predict (#16785, #16826)</details>
<a name="changelog-fabric"></a>
<details><summary>Added</summary>
Fabric.all_reduce (#16459)Fabric.save/load() (#16452)set_epoch on the dataloader.batch_sampler.sampler (#16841)CSVLogger (#16880)lightning.fabric.is_wrapped to check whether a module, optimizer, or dataloader was already wrapped by Fabric (#16953)</details>
<details><summary>Changed</summary>
accelerator="auto", strategy="auto", devices="auto" as defaults (#16842)Fabric.save and Fabric.loadStrategy.save_checkpoint and Fabric.load_checkpointFabric.save accepts a state that can contain model and optimizer referencesFabric.load can now load state in-place onto models and optimizersFabric.load returns a dictionary of objects that weren't loaded into the stateStrategy.save_checkpoint and Fabric.load_checkpoint are now responsible for accessing the state of the model and optimizersDataParallelStrategy.get_module_state_dict() and DDPStrategy.get_module_state_dict() now correctly extracts the state dict without keys prefixed with 'module' (#16487)strategy="fsdp_full_shard_offload" is now strategy="fsdp_cpu_offload"lightning.fabric.plugins.precision.native_amp is now lightning.fabric.plugins.precision.ampstrategy='tpu_spawn' to strategy='xla' and strategy='tpu_spawn_debug' to strategy='xla_debug' (#16781)Fabric(strategy="ddp_spawn", ...) no longer falls back to "ddp" when a cluster environment gets detected (#16780)setup_dataloaders(replace_sampler=...) to setup_dataloaders(use_distributed_sampler=...) (#16829)</details>
<details><summary>Removed</summary>
</details>
<details><summary>Fixed</summary>
iter() would be called twice (#16841)</details>
<a name="changelog-app"></a>
<details><summary>Added</summary>
--zip option to the lightning cp command to copy content from the Cloud Platform Filesystem as a zipfile</details>
<details><summary>Changed</summary>
rich from 10.14.0 to 12.13.0 (#16798)</details>
<details><summary>Removed</summary>
</details>
Full commit list: https://github.com/Lightning-AI/lightning/compare/1.9.0...2.0.0
<a name="contributors"></a>
@aniketmaurya @Atharva-Phatak @awaelchli @Borda @carmocca @dmitsf @edenlightning @hhsecond @janpawlowskiof @justusschock @krshrimali @leoleoasd @martenlienen @narJH27 @rusmux @SauravMaheshkar @shenoynikhil @tshu-w @wouterzwerink
@coreyjadams @dconathan @RuRo @sergeevii123 @BrianPulfer @akkefa @belerico @tupini07 @jihoonkim2100 @eamonn-zh @lightningforever @mauvilsa @muhammadanas0716 @AdityaKane2001 @dtuit @themattinthehatt @janpawlowskiof @rusmux @Erotemic @janEbert
If we forgot somebody or you have a suggestion, find us on Discord :zap:
Chuck Norris jokes aren't actually jokes, they are the truth. Chuck Norris actually can do all the things people say he can.
Full Changelog: https://github.com/Lightning-AI/lightning/compare/1.9.0...2.0.0rc0
Full Changelog: https://github.com/Lightning-AI/lightning/compare/1.9.0...2.0.0rc0
Added healthz endpoint to plugin server
healthz endpoint to plugin server (#16882)TorchCollective works on the torch.distributed WORLD process group by default (#16995)_cuda_clearCublasWorkspaces on teardown (#16907)NeptuneLogger (#16761):
log() method with append() and extend().Handler as an alternative to Run for the run argument. This means that you can call it like NeptuneLogger(run=run["some/namespace"]) to log everything to the some/namespace/ location of the run.sys.argv and args in LightningCLI (#16808)ShardedTensor state dict hooks in LightningModule.__init__ with torch>=2.1 (#16892)lightning.pytorch.core.saving.ModelIO class interface (#16974)num_nodes not being set for DDPFullyShardedNativeStrategy (#17160)DeepSpeedStrategy (#16973)rich that prevented Lightning to be imported in Google Colab (#17156)_cuda_clearCublasWorkspaces on teardown (#16907)psutil package is now required for CPU monitoring (#17010)@awaelchli, @belerico, @carmocca, @colehawkins, @dmitsf, @Erotemic, @ethanwharris, @kshitij12345, @borda
If we forgot someone due to not matching commit email with GitHub account, let us know :]
Removed implicit ui testing with testing.run_app_in_cloud in favor of headless login and app selection
testing.run_app_in_cloud in favor of headless login and app selection (#16741)Fabric(strategy="auto") support (#16916)find_usable_cuda_devices(num_devices=-1) (#16866)Fabric(strategy="auto") support. It will choose DDP over DDP-spawn, contrary to strategy=None (default) (#16916)lightning.pytorch.utilities.parsing.get_init_args (#16851)@ethanwharris, @carmocca, @awaelchli, @justusschock , @dtuit, @Liyang90
If we forgot someone due to not matching commit email with GitHub account, let us know :]
Fixed lightning open command and improved redirects
lightning open command and improved redirects (#16794)accelerator=tpu and devices > 1 (#16806)--accelerator and --precision in Fabric CLI when accelerator and precision are set to non-default values in the code (#16818)accelerator=tpu and devices > 1 (#16806)@ethanwharris, @carmocca, @awaelchli, @borda, @tchaton, @yurijmikhalevich
If we forgot someone due to not matching commit email with GitHub account, let us know :]
rm: Delete files from your Cloud Platform Filesystem
rm: Delete files from your Cloud Platform Filesystemlightning connect data to register data connection to private s3 buckets (#16738)min_epochs or min_steps (#16719)@akihironitta, @awaelchli, @borda, @tchaton
If we forgot someone due to not matching commit email with GitHub account, let us know :]
Added experimental support for interruptable GPU in the cloud
lightning open command (#16482)ls: List files from your Cloud Platform Filesystemcd: Change the current directory within your Cloud Platform filesystem (terminal session based)pwd: Return the current folder in your Cloud Platform Filesystemcp: Copy files between your Cloud Platform Filesystem and local filesystemcd into non-existent folders (#16645)cp (upload) at project level (#16631)ls and cp (download) at project level (#16622)lightning connect data to register data connection to s3 buckets (#16670)LightningClient(retry=False) to retry=True (#16382)lightning.app.components.LiteMultiNode to lightning.app.components.FabricMultiNode (#16505)lightning connect to lightning connect app for consistency (#16670)lightning cp (#16626)accelerator="mps" and ddp strategy pairing (#16455)torch_xla requirement (#16476)torch.distributed is not available (#16658)save_hyperparameters on mixin classes that don't subclass LightningModule/LightningDataModule (#16369)MLFlowLogger logging the wrong keys with .log_hyperparams() (#16418)MLFlowLogger and long values are truncated (#16451)torch_xla requirement (#16476)torch.distributed is not available (#16658)@akihironitta, @awaelchli, @borda, @BrianPulfer, @ethanwharris, @hhsecond, @justusschock, @Liyang90, @RuRo, @senarvi, @shenoynikhil, @tchaton
If we forgot someone due to not matching commit email with GitHub account, let us know :]
Avoid using the deprecated LooseVersion
DeviceStatsMonitor
(#16002)lightning_app.components.serve.gradio to lightning_app.components.serve.gradio_server (#16201)relpath bug on Windows (#16164)LooseVersion (#16162)lightning login with env variables would not correctly save the credentials (#16339)Fabric.launch() to programmatically launch processes (e.g. in Jupyter notebook) (#14992)run method (#14992)Fabric.setup_module() and Fabric.setup_optimizers() to support strategies that need to set up the model before an optimizer can be created (#15185)lightning_fabric.accelerators.find_usable_cuda_devices utility function (#16147)Fabric(callbacks=...) and emitting events through Fabric.call() (#16074)Fabric(loggers=...) to support different Logger frameworks in FabricFabric.log for logging scalars using multiple loggersFabric.log_dict for logging a dictionary of multiple metrics at onceFabric.loggers and Fabric.logger attributes to access the individual logger instancesself.log and self.log_dict in a LightningModule when using Fabricself.logger and self.loggers in a LightningModule when using Fabriclightning_fabric.loggers.TensorBoardLogger (#16121)lightning_fabric.loggers.CSVLogger (#16346).zero_grad(set_to_none=...) on the wrapped optimizer regardless of which strategy is used (#16275)LightningLite to Fabric (#15932, #15938)Fabric.run() method is no longer abstract (#14992)XLAStrategy now inherits from ParallelStrategy instead of DDPSpawnStrategy (#15838)DDPSpawnStrategy into DDPStrategy and removed DDPSpawnStrategy (#14952).setup_dataloaders() now calls .set_epoch() on the distributed sampler if one is used (#16101)Strategy.reduce to Strategy.all_reduce in all strategies (#16370)strategy='ddp_sharded'|'ddp_sharded_spawn'). Use Fully-Sharded Data Parallel instead (strategy='fsdp') (#16329)DistributedSampler (#16101)MetricCollection with enabled compute groups (#15580)pl.loggers.WandbLogger (#16173)LRFinder (#15304)pl.utilities.upgrade_checkpoint script (#15333)ax to the .lr_find().plot() to enable writing to a user-defined axes in a matplotlib figure (#15652)log_model parameter to MLFlowLogger (#9187)self.log(..., logger=True) is called without a configured logger (#15814)LightningCLI support for optimizer and learning schedulers via callable type dependency injection (#15869)DDPFullyShardedNativeStrategy strategy (#15826)DDPFullyShardedNativeStrategy(cpu_offload=True|False) via bool instead of needing to pass a configuration object (#15832)LightningModule.configure_optimizers (#16189)tensorboard to tensorboardx in TensorBoardLogger (#15728)LightningModule.load_from_checkpoint automatically upgrade the loaded checkpoint if it was produced in an old version of Lightning (#15237)Trainer.{validate,test,predict}(ckpt_path=...) no longer restores the Trainer.global_step and trainer.current_epoch value from the checkpoints - From now on, only Trainer.fit will restore this value (#15532)ModelCheckpoint.save_on_train_epoch_end attribute is now computed dynamically every epoch, accounting for changes to the validation dataloaders (#15300)MLFlowLogger now logs hyperparameters and metrics in batched API calls (#15915)on_train_batch_{start,end} hooks in conjunction with taking a dataloader_iter in the training_step no longer errors out and instead shows a warning (#16062)tensorboardX to extra dependencies. Use the CSVLogger by default (#16349)description, env_prefix and env_parse parameters in LightningCLI.__init__ in favour of giving them through parser_kwargs (#15651)pytorch_lightning.profiler in favor of pytorch_lightning.profilers (#16059)Trainer(auto_select_gpus=...) in favor of pytorch_lightning.accelerators.find_usable_cuda_devices (#16147)pytorch_lightning.tuner.auto_gpu_select.{pick_single_gpu,pick_multiple_gpus} in favor of pytorch_lightning.accelerators.find_usable_cuda_devices (#16147)nvidia/apex deprecation (#16039)
pytorch_lightning.plugins.NativeMixedPrecisionPlugin in favor of pytorch_lightning.plugins.MixedPrecisionPluginLightningModule.optimizer_step(using_native_amp=...) argumentTrainer(amp_backend=...) argumentTrainer.amp_backend propertyTrainer(amp_level=...) argumentpytorch_lightning.plugins.ApexMixedPrecisionPlugin classpytorch_lightning.utilities.enums.AMPType enumDeepSpeedPrecisionPlugin(amp_type=..., amp_level=...) argumentshorovod deprecation (#16141)
Trainer(strategy="horovod")HorovodStrategy classpytorch_lightning.lite.LightningLite in favor of lightning.fabric.Fabric (#16314)FairScale deprecation (in favor of PyTorch's FSDP implementation) (#16353)
pytorch_lightning.overrides.fairscale.LightningShardedDataParallel classpytorch_lightning.plugins.precision.fully_sharded_native_amp.FullyShardedNativeMixedPrecisionPlugin classpytorch_lightning.plugins.precision.sharded_native_amp.ShardedNativeMixedPrecisionPlugin classpytorch_lightning.strategies.fully_sharded.DDPFullyShardedStrategy classpytorch_lightning.strategies.sharded.DDPShardedStrategy classpytorch_lightning.strategies.sharded_spawn.DDPSpawnShardedStrategy classpytorch_lightning.utilities.memory.get_gpu_memory_map in favor of pytorch_lightning.accelerators.cuda.get_nvidia_gpu_stats (#15617)pytorch_lightning.profiler.base.AbstractProfiler in favor of pytorch_lightning.profilers.profiler.Profiler (#15637)pytorch_lightning.profiler.base.BaseProfiler in favor of pytorch_lightning.profilers.profiler.Profiler (#15637)pytorch_lightning.utilities.meta (#16038)LightningDeepSpeedModule (#16041)pytorch_lightning.accelerators.GPUAccelerator in favor of pytorch_lightning.accelerators.CUDAAccelerator (#16050)pytorch_lightning.profiler.* classes in favor of pytorch_lightning.profilers (#16059)pytorch_lightning.utilities.cli module in favor of pytorch_lightning.cli (#16116)pytorch_lightning.loggers.base module in favor of pytorch_lightning.loggers.logger (#16120)pytorch_lightning.loops.base module in favor of pytorch_lightning.loops.loop (#16142)pytorch_lightning.core.lightning module in favor of pytorch_lightning.core.module (#16318)pytorch_lightning.callbacks.base module in favor of pytorch_lightning.callbacks.callback (#16319)Trainer.reset_train_val_dataloaders() in favor of Trainer.reset_{train,val}_dataloader (#16131)LightningCLI(seed_everything_default=None) (#16131)strategy='ddp_sharded'|'ddp_sharded_spawn'). Use Fully-Sharded Data Parallel instead (strategy='fsdp') (#16329)reduce_boolean_decision to accommodate any-analogous semantics expected by the EarlyStopping callback (#15253)interval key of the scheduler would be ignored during manual optimization, making the LearningRateMonitor callback fail to log the learning rate (#16308)MLFlowLogger not finalizing correctly when status code 'finished' was passed (#16340)@1SAA, @akihironitta, @AlessioQuercia, @awaelchli, @bipinKrishnan, @Borda, @carmocca, @dmitsf, @erhoo82, @ethanwharris, @Forbu, @hhsecond, @justusschock, @lantiga, @lightningforever, @Liyang90, @manangoel99, @mauvilsa, @nicolai86, @nohalon, @rohitgr7, @schmidt-jake, @speediedan, @yMayanand
If we forgot someone due to not matching commit email with GitHub account, let us know :]
Nothing published for this version
Added partial support for fastapi Request annotation in configure_api handlers
Request annotation in configure_api handlers (#16047)work.delete method to delete the work (#16103)display_name property to LightningWork for the cloud (#16095)ColdStartProxy to the AutoScaler (#16094)ready (#16075)ready for components (#16129)start_method for creating Work processes locally on macOS is now 'spawn' (previously 'fork') (#16089)lightning.app.utilities.cloud.is_running_in_cloud now returns True during the loading of the app locally when running with --cloud (#16045)True (#16009)PythonServer messaging "Your app has started" (#15989)AutoScaler would fail with min_replica=0 (#16092AutoScaler UI (#16128)streamlit (#16139)Full Changelog: https://github.com/Lightning-AI/lightning/compare/1.8.5.post0...1.8.6
Your coding agent can read these notes before it upgrades. Set up the MCP server →