The core team is excited to announce the PyTorch Lightning 1.6 release ⚡

• Highlights
• Backward Incompatible Changes
• Full Changelog
• Contributors

Highlights

PyTorch Lightning 1.6 is the work of 99 contributors who have worked on features, bug-fixes, and documentation for a total of over 750 commits since 1.5. This is our most active release yet. Here are some highlights:

Introducing Intel's Habana Accelerator

Lightning 1.6 now supports the Habana® framework, which includes Gaudi® AI training processors. Their heterogeneous architecture includes a cluster of fully programmable Tensor Processing Cores (TPC) along with its associated development tools and libraries and a configurable Matrix Math engine.

You can leverage the Habana hardware to accelerate your Deep Learning training workloads simply by passing:

trainer = pl.Trainer(accelerator="hpu")

# single Gaudi training
trainer = pl.Trainer(accelerator="hpu", devices=1)

# distributed training with 8 Gaudi
trainer = pl.Trainer(accelerator="hpu", devices=8)

The Bagua Strategy

The Bagua Strategy is a deep learning acceleration framework that supports multiple, advanced distributed training algorithms with state-of-the-art system relaxation techniques. Enabling Bagua, which can be considerably faster than vanilla PyTorch DDP, is as simple as:

trainer = pl.Trainer(strategy="bagua")

# or to choose a custom algorithm
trainer = pl.Trainer(strategy=BaguaStrategy(algorithm="gradient_allreduce")  # default

Towards stable Accelerator, Strategy, and Plugin APIs

The Accelerator, Strategy, and Plugin APIs are a core part of PyTorch Lightning. They're where all the distributed boilerplate lives, and we're constantly working to improve both them and the overall PyTorch Lightning platform experience.

In this release, we've made some large changes to achieve that goal. Not to worry, though! The only users affected by these changes are those who use custom implementations of Accelerator and Strategy (TrainingTypePlugin) as well as certain Plugins. In particular, we want to highlight the following changes:

• All TrainingTypePlugins have been renamed to Strategy (#11120). Strategy is a more appropriate name because it encompasses more than simply training communcation. This change is now aligned with the changes we implemented in 1.5, which introduced the new strategy and devices flags to the Trainer.

  # Before
  from pytorch_lighting.plugins import DDPPlugin
  
  # New
  from pytorch_lighting.strategies import DDPStrategy

• The Accelerator and PrecisionPlugin have moved into Strategy. All strategies now take an optional parameter accelerator and precision_plugin (#11022, #10570).

• Custom Accelerator implementations must now implement two new abstract methods: is_available() (#11797) and auto_device_count() (#10222). The latter determines how many devices get used by default when specifying Trainer(accelerator=..., devices="auto").

• We redesigned the process creation for spawn-based strategies such as DDPSpawnStrategy and TPUSpawnStrategy (#10896). All spawn-based strategies now spawn processes immediately upon calling Trainer.{fit,validate,test,predict}, which means the hooks/callbacks prepare_data, setup, configure_sharded_model and teardown all run under an initialized process group. These changes align the spawn-based strategies with their non-spawn counterparts (such as DDPStrategy).

We've also exposed the process group backend for use. For example, you can now easily enable fairring like this:

# Explicitly specify the process group backend if you choose to
ddp = pl.strategies.DDPStrategy(process_group_backend="fairring")
trainer = Trainer(strategy=ddp, accelerator="gpu", devices=8)

In a similar fashion, if installing torch>=1.11, you can enable DDP static graph to apply special runtime optimizations:

trainer = Trainer(devices=4, strategy=DDPStrategy(static_graph=True))

LightningCLI improvements

In the previous release, we added shorthand notation support for registered components. In this release, we added a flag to automatically register all available components:

from pytorch_lightning.utilities.cli import LightningCLI

LightningCLI(auto_registry=True)

We have also added support for the ReduceLROnPlateau scheduler with shorthand notation:

$ python script.py fit --optimizer=Adam --lr_scheduler=ReduceLROnPlateau --lr_scheduler.monitor=metric_to_track

If you need to customize the learning rate scheduler configuration, you can do so by overriding:

class MyLightningCLI(LightningCLI):
    @staticmethod
    def configure_optimizers(lightning_module, optimizer, lr_scheduler=None):
        return {"optimizer": optimizer, "lr_scheduler": {"scheduler": lr_scheduler, ...}}

Finally, loggers are also now configurable with shorthand:

$ python script.py fit --trainer.logger=WandbLogger --trainer.logger.name="my_lightning_run"

Control SLURM's re-queueing

We've added the ability to turn the automatic resubmission on or off when a job gets interrupted by the SLURM controller (via signal handling). Users who prefer to let their code handle the resubmission (for example, when submitit is used) can now pass:

from pytorch_lightning.plugins.environments import SLURMEnvironment

trainer = pl.Trainer(plugins=SLURMEnvironment(auto_requeue=False))

Fault-tolerance improvements

The Fault-tolerance training under manual optimization now tracks optimization progress. We also changed the graceful exit signal from SIGUSR1 to SIGTERM for better support inside cloud instances.
An additional feature we're excited to announce is support for consecutive trainer.fit() calls.

trainer = pl.Trainer(max_epochs=2)
trainer.fit(model)

# now, run 2 more epochs
trainer.fit_loop.max_epochs = 4
trainer.fit(model)

Loop customization improvements

The Loop's state is now included as part of the checkpoints saved by the library. This enables finer restoration of custom loops.

We've also made it easier to replace Lightning's loops with your own. For example:

class MyCustomLoop(pl.loops.TrainingEpochLoop):
    ...

trainer = pl.Trainer(...)
trainer.fit_loop.replace(epoch_loop=MyCustomLoop)
# Trainer runs the fit loop with your new epoch loop!
trainer.fit(model)

Data-Loading improvements

In previous versions, Lightning required that the DataLoader instance set its input arguments as instance attributes. This meant that custom DataLoaders also had this hidden requirement. In this release, we do this automatically for the user, easing the passing of custom loaders:

class MyDataLoader(torch.utils.data.DataLoader):
    def __init__(self, a=123, *args, **kwargs):
-       # this was required before
-       self.a = a
        super().__init__(*args, **kwargs)

trainer.fit(model, train_dataloader=MyDataLoader())

As of this release, Lightning no longer pre-fetches 1 extra batch if it doesn't need to. Previously, doing so would conflict with the internal pre-fetching done by optimized data loaders such as FFCV's. You can now define your own pre-fetching value like this:

class MyCustomLoop(pl.loops.FitLoop):
    @property
    def prefetch_batches(self):
        return 7  # lucky number 7

trainer = pl.Trainer(...)
trainer.fit_loop.replace(fit_loop=MyCustomLoop)

New Hooks

LightningModule.lr_scheduler_step

Lightning now allows the use of custom learning rate schedulers that aren't natively available in PyTorch. A great example of this is Timm Schedulers.

When using custom learning rate schedulers relying on an API other than PyTorch's, you can now define the LightningModule.lr_scheduler_step with your desired logic.

from timm.scheduler import TanhLRScheduler


class MyLightningModule(pl.LightningModule):
    def configure_optimizers(self):
        optimizer = ...
        scheduler = TanhLRScheduler(optimizer, ...)
        return {"optimizer": optimizer, "lr_scheduler": {"scheduler": scheduler, "interval": "epoch"}}

    def lr_scheduler_step(self, scheduler, optimizer_idx, metric):
        scheduler.step(epoch=self.current_epoch)  # timm's scheduler need the epoch value

A new stateful API

This release introduces new hooks to standardize all stateful components to use state_dict and load_state_dict, mimicking the PyTorch API. The new hooks receive their own component's state and replace most usages of the previous on_save_checkpoint and on_load_checkpoint hooks.

def MyCallback(pl.Callback):
-   def on_save_checkpoint(self, trainer, pl_module, checkpoint):
-       return {'x': self.x}
    
-   def on_load_checkpoint(self, trainer, pl_module, checkpoint):
-       self.x = x

+   def state_dict(self):
+       return {'x': self.x}
    
+   def load_state_dict(self, checkpoint):
+       self.x = x     

New properties

Trainer.estimated_stepping_batches

You can use built-in Trainer.estimated_stepping_batches to compute the total number of stepping batches needed for the complete training.

The property takes gradient accumulation factor and distributed setting into consideration when performing this computation so that you don't have to derive it manually:

class MyLightningModule(pl.LightningModule):
    def configure_optimizers(self):
        optimizer = ...
        scheduler = torch.optim.lr_scheduler.OneCycleLR(
            optimizer, max_lr=1e-3, total_steps=self.trainer.estimated_stepping_batches
        )
        return {"optimizer": optimizer, "lr_scheduler": scheduler}

Trainer.num_devices and Trainer.device_ids

In the past, retrieving the number of devices used, or their IDs, posed a considerable challenge. Additionally, doing so required knowing which property to access based on the current Trainer configuration.

To simplify this process, we've deprecated the per-accelerator properties to have accelerator agnostic properties. For example:

- num_devices = max(1, trainer.num_gpus, trainer.num_processes)
- if trainer.tpu_cores:
-    num_devices = max(num_devices, trainer.tpu_cores)
+ num_devices = trainer.num_devices

Experimental Features

Manual Fault-tolerance

Fault Tolerance has limitations that require specific information about your data-loading structure.

It is now possible to resolve those limitations by enabling manual fault tolerance where you can write your own logic and specify how exactly to checkpoint your own datasets and samplers. You can do so using this environment flag:

$ PL_FAULT_TOLERANT_TRAINING=MANUAL python script.py

Check out this video for a dive into the internals of this flag.

Customizing the layer synchronization

We introduced a new plugin class for wrapping layers of a model with synchronization logic for multiprocessing.

class MyLayerSync(pl.plugins.LayerSync):
    ...

layer_sync = MyLayerSync(...)
trainer = Trainer(sync_batchnorm=True, plugins=layer_sync, strategy="ddp")

Registering Custom Accelerators

There has been much progress in the field of ML Accelerators, and the list of accelerators is constantly expanding.

We've made it easier for users to try out new accelerators by enabling support for registering custom Accelerator classes in Lightning.

from pytorch_lightning.accelerators import Accelerator, AcceleratorRegistry

class SOTAAccelerator(Accelerator):
    def __init__(self, x):
        ...

AcceleratorRegistry.register("sota_accelerator", SOTAAccelerator, x=123)
# the following works now:
trainer = Trainer(accelerator="sota_accelerator")

Backward Incompatible Changes

Here is a selection of notable changes that are not backward compatible with previous versions. The full list of changes and removals can be found in the CHANGELOG below.

Drop PyTorch 1.7 support

Following our 4 PyTorch release window, this release supports PyTorch 1.8 to 1.11. Support for PyTorch 1.7 has been removed.

Drop Python 3.6 support

Following Python's end-of-life, support for Python 3.6 has been removed.

AcceleratorConnector rewrite

To support new accelerator and stategy features, we completely rewrote our internal AcceleratorConncetor class. No backwards compatibility was maintained so it is likely to have broken your code if it was using this class.

Re-define the current_epoch boundary

To resolve fault-tolerance issues, we changed where the current epoch value gets increased.

trainer.current_epoch is now increased by 1 on_train_end. This means that if a model is run for 3 epochs (0, 1, 2), trainer.current_epoch will now return 3 instead of 2 after trainer.fit(). This can also impact custom callbacks that acess this property inside this hook.

This also impacts checkpoints saved during an epoch (e.g. on_train_epoch_end). For example, a Trainer(max_epochs=1, limit_train_batches=1) instance that saves a checkpoint will have the current_epoch=0 value saved instead of current_epoch=1.

Re-define the global_step boundary

To resolve fault-tolerance issues, we changed where the global step value gets increased.

Access to trainer.global_step during an intra-training validation hook will now correctly return the number of optimizer steps taken already. In pseudocode:

  training_step()
+ global_step += 1
  validation_if_necessary()
- global_step += 1

Saved checkpoints that use the global step value as part of the filename are now increased by 1 for the same reason. A checkpoint saved after 1 step will be now be named step=1.ckpt instead of step=0.ckpt.

The trainer.global_step value will now account for TBPTT or multiple optimizers. Users setting Trainer({min,max}_steps=...) under these circumstances will need to adjust their values.

Removed automatic reduction of outputs in training_step when using DataParallel

When using Trainer(strategy="dp"), all the tensors returned by training_step were previously reduced to a scalar (#11594). This behavior was especially confusing when outputs needed to be collected into the training_epoch_end hook.

From now on, outputs are no longer reduced except for the loss tensor, unless you implement training_step_end, in which case the loss won't get reduced either.

No longer fallback to CPU with no devices

Previous versions were lenient in that the lack of GPU devices defaulted to running on CPU. This meant that users' code could be running much slower without them ever noticing that it was running on CPU.

We suggest passing Trainer(accelerator="auto") when this leniency is desired.

CHANGELOG

Full commit list: 1.5.0...1.6.0

Contributors

Veteran

@akihironitta @ananthsub @awaelchli @Borda @borisdayma @carmocca @daniellepintz @edward-io @ethanwharris @four4fish @jjenniferdai @kaushikb11 @kingyiusuen @kragniz @mauvilsa @ninginthecloud @popfido @rohitgr7 @SeanNaren @speediedan @tchaton @tshu-w @twsl @williamFalcon

New

@a-gardner1 @abhi-rf @abhinavarora @adamreeve @adamviola @AJSVB @akashkw @amin-nejad @AndresAlgaba @ant0nsc @armanal @bhadreshpsavani @CAIQT @catalys1 @chaddy1004 @chunyang-wen @circlecrystal @Code-Cornelius @Cyber-Machine @dennisbappert @DuYicong515 @edpizzi @franp9am @ftorres16 @ggare-cmu @guyang3532 @Honzys @idiomaticrefactoring @isvogor-foi @jerome-habana @jgibson2 @jlhbaseball15 @jona-0 @JoostvDoorn @josafatburmeister @konstantinjdobler @Kr4is @krishnakalyan3 @krshrimali @lemairecarl @lucmos @manangoel99 @mathemusician @mayeroa @mbortolon97 @NathanGodey @Nesqulck @nithinraok @ORippler @os1ma @peterdudfield @Piyush-97 @puhuk @qqueing @quancs @Raahul-Singh @Raalsky @Rajathbharadwaj @rasbt @rharish101 @rhjohnstone @rjkilpatrick @RobertLaurella @roschly @rsokl @rusty1s @SauravMaheshkar @sethvargo @shabie @shivammehta007 @srb-cv @ThomVett @wangraying @whokilleddb @zredeaux65

If we forgot someone or have any suggestion, let us know in Slack ⚡