The PyTorch Lightning team released version 1.6 with support for Intel’s Habana Accelerator, new efficient DDP strategy (Bagua), manual Fault-tolerance, and other stability and reliability changes.<\/p>\n
\u26a1Visit the release page on GitHub<\/u><\/a> to download.\u26a1<\/p>\n 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. Here are some highlights:<\/p>\n Lightning 1.6 now supports the Habana\u00ae framework, which includes Gaudi\u00ae 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.<\/p>\n You can leverage the\u00a0Habana<\/a>\u00a0hardware to accelerate your Deep Learning training workloads simply by passing:<\/p>\n The\u00a0Bagua Strategy<\/a>\u00a0is a deep learning acceleration framework that supports multiple, advanced distributed training algorithms with state-of-the-art system relaxation techniques. Enabling\u00a0Bagua<\/a>, which can be considerably faster than vanilla PyTorch DDP, is as simple as:<\/p>\n The\u00a0 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 ( We’ve also exposed the process group backend for use. For example, you can now easily enable\u00a0 In a similar fashion, if installing\u00a0 In the previous release, we added shorthand notation support for registered components. In this release, we added a flag to\u00a0automatically register<\/a>\u00a0all available components:<\/p>\n We have also added support for the\u00a0 If you need to\u00a0customize the learning rate scheduler configuration<\/a>, you can do so by overriding:<\/p>\n Finally, loggers are also now configurable with shorthand:<\/p>\n We’ve added the ability to turn the\u00a0automatic resubmission<\/a>\u00a0on 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:<\/p>\n The\u00a0Fault-tolerance training<\/a>\u00a0under manual optimization now tracks optimization progress. We also changed the graceful exit signal from\u00a0 An additional feature we’re excited to announce is support for consecutive The\u00a0 We’ve also made it easier to replace Lightning’s loops with your own. For example:<\/p>\n In previous versions, Lightning required that the\u00a0 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\u00a0FFCV’s<\/a>. You can now define your own pre-fetching value like this:<\/p>\n Lightning now allows the use of\u00a0custom learning rate schedulers<\/a>\u00a0that aren’t natively available in\u00a0PyTorch<\/a>. A great example of this is\u00a0Timm Schedulers<\/a>.<\/p>\n When using custom learning rate schedulers relying on an API other than PyTorch’s, you can now define the\u00a0 This release introduces new hooks to standardize all stateful components to use\u00a0 You can use built-in\u00a0 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:<\/p>\n 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\u00a0 To simplify this process, we’ve deprecated the per-accelerator properties to have accelerator agnostic properties. For example:<\/p>\n Manual Fault-tolerance<\/strong><\/p>\n Fault Tolerance<\/a>\u00a0has limitations that require specific information about your data-loading structure.<\/p>\n 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:<\/p>\n Check out\u00a0this video<\/a>\u00a0for a dive into the internals of this flag.<\/p>\n Customizing the layer synchronization<\/strong><\/p>\n We introduced a new plugin class for wrapping layers of a model with synchronization logic for multiprocessing.<\/p>\n Registering Custom Accelerators<\/strong><\/p>\n There has been much progress in the field of ML Accelerators, and the list of accelerators is constantly expanding.<\/p>\n We’ve made it easier for users to try out new accelerators by enabling support for registering custom\u00a0 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.<\/p>\n Following our 4 PyTorch release window, this release supports PyTorch 1.8 to 1.11. Support for PyTorch 1.7 has been removed.<\/p>\n Following\u00a0Python’s end-of-life<\/a>, support for Python 3.6 has been removed.<\/p>\n To support new accelerator and stategy features, we completely rewrote our internal\u00a0 To resolve fault-tolerance issues, we changed where the current epoch value gets increased.<\/p>\n This also impacts checkpoints saved during an epoch (e.g.\u00a0 To resolve fault-tolerance issues, we changed where the global step value gets increased.<\/p>\n Access to\u00a0 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\u00a0 The\u00a0 When using\u00a0 From now on, outputs are no longer reduced except for the\u00a0 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.<\/p>\n We suggest passing\u00a0\n
Lightning Highlights<\/h3>\n
Introducing Intel’s Habana Accelerator<\/h4>\n
trainer = pl.Trainer(accelerator=\"hpu\")\r\n\r\n# single Gaudi training\r\ntrainer = pl.Trainer(accelerator=\"hpu\", devices=1)\r\n\r\n# distributed training with 8 Gaudi\r\ntrainer = pl.Trainer(accelerator=\"hpu\", devices=8)<\/pre>\n
The Bagua Strategy<\/h4>\n
trainer = pl.Trainer(strategy=\"bagua\")\r\n\r\n# or to choose a custom algorithm\r\ntrainer = pl.Trainer(strategy=BaguaStrategy(algorithm=\"gradient_allreduce\") # default<\/pre>\n
Towards stable Accelerator, Strategy, and Plugin APIs<\/h4>\n
Accelerator<\/code>,\u00a0Strategy<\/code>, and\u00a0Plugin<\/code>\u00a0APIs 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.<\/p>\nTrainingTypePlugin<\/code>) as well as certain Plugins. In particular, we want to highlight the following changes:<\/p>\n\n
TrainingTypePlugin<\/code>s have been renamed to\u00a0Strategy<\/code>\u00a0(#11120<\/a>). 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\u00a0strategy<\/code>\u00a0and\u00a0devices<\/code>\u00a0flags to the Trainer<\/a>.\n# Before\r\nfrom pytorch_lighting.plugins import DDPPlugin\r\n\r\n# New\r\nfrom pytorch_lighting.strategies import DDPStrategy<\/pre>\n<\/li>\n
Accelerator<\/code>\u00a0and\u00a0PrecisionPlugin<\/code>\u00a0have moved into\u00a0Strategy<\/code>. All strategies now take an optional parameter\u00a0accelerator<\/code>\u00a0and\u00a0precision_plugin<\/code>\u00a0(#11022<\/a>,\u00a0#10570<\/a>).<\/li>\nis_available()<\/code>\u00a0(#11797<\/a>) and\u00a0auto_device_count()<\/code>\u00a0(#10222<\/a>). The latter determines how many devices get used by default when specifying\u00a0Trainer(accelerator=..., devices=\"auto\")<\/code>.<\/li>\nDDPSpawnStrategy<\/code>\u00a0and\u00a0TPUSpawnStrategy<\/code>\u00a0(#10896<\/a>). All spawn-based strategies now spawn processes immediately upon calling\u00a0Trainer.{fit,validate,test,predict}<\/code>, which means the hooks\/callbacks\u00a0prepare_data<\/code>,\u00a0setup<\/code>,\u00a0configure_sharded_model<\/code>\u00a0and\u00a0teardown<\/code>\u00a0all run under an initialized process group. These changes align the spawn-based strategies with their non-spawn counterparts (such as\u00a0DDPStrategy<\/code>).<\/li>\n<\/ul>\nfairring<\/code><\/a>\u00a0like this:<\/p>\n# Explicitly specify the process group backend if you choose to\r\nddp = pl.strategies.DDPStrategy(process_group_backend=\"fairring\")\r\ntrainer = Trainer(strategy=ddp, accelerator=\"gpu\", devices=8)<\/pre>\n
torch>=1.11<\/code>, you can enable\u00a0DDP static graph<\/a>\u00a0to apply special runtime optimizations:<\/p>\ntrainer = Trainer(devices=4, strategy=DDPStrategy(static_graph=True))<\/pre>\n
LightningCLI<\/code>\u00a0improvements<\/h4>\nfrom pytorch_lightning.utilities.cli import LightningCLI\r\n\r\nLightningCLI(auto_registry=True)<\/pre>\n
ReduceLROnPlateau<\/code>\u00a0scheduler with shorthand notation:<\/p>\n$ python script.py fit --optimizer=Adam --lr_scheduler=ReduceLROnPlateau --lr_scheduler.monitor=metric_to_track<\/pre>\n
class MyLightningCLI(LightningCLI):\r\n @staticmethod\r\n def configure_optimizers(lightning_module, optimizer, lr_scheduler=None):\r\n return {\"optimizer\": optimizer, \"lr_scheduler\": {\"scheduler\": lr_scheduler, ...}}<\/pre>\n$ python script.py fit --trainer.logger=WandbLogger --trainer.logger.name=\"my_lightning_run\"<\/pre>\n
Control SLURM’s re-queueing<\/h4>\n
from pytorch_lightning.plugins.environments import SLURMEnvironment\r\n\r\ntrainer = pl.Trainer(plugins=SLURMEnvironment(auto_requeue=False))<\/pre>\n
Fault-tolerance improvements<\/h4>\n
SIGUSR1<\/code>\u00a0to\u00a0SIGTERM<\/code>\u00a0for better support inside cloud instances.<\/p>\ntrainer.fit()<\/code>\u00a0calls.<\/p>\ntrainer = pl.Trainer(max_epochs=2)\r\ntrainer.fit(model)\r\n\r\n# now, run 2 more epochs\r\ntrainer.fit_loop.max_epochs = 4\r\ntrainer.fit(model)<\/pre>\n
Loop customization improvements<\/h4>\n
Loop<\/code><\/a>‘s state is now included as part of the checkpoints saved by the library. This enables finer restoration of custom loops.<\/p>\nclass MyCustomLoop(pl.loops.TrainingEpochLoop):\r\n ...\r\n\r\ntrainer = pl.Trainer(...)\r\ntrainer.fit_loop.replace(epoch_loop=MyCustomLoop)\r\n# Trainer runs the fit loop with your new epoch loop!\r\ntrainer.fit(model)<\/pre>\n
Data-Loading improvements<\/h4>\n
DataLoader<\/code>\u00a0instance set its input arguments as instance attributes. This meant that custom\u00a0DataLoader<\/code>s also had this hidden requirement. In this release, we do this automatically for the user, easing the passing of custom loaders:<\/p>\nclass MyDataLoader(torch.utils.data.DataLoader):\r\n def __init__(self, a=123, *args, **kwargs):\r\n- # this was required before\r\n- self.a = a\r\n super().__init__(*args, **kwargs)\r\n\r\ntrainer.fit(model, train_dataloader=MyDataLoader())<\/pre>\n
class MyCustomLoop(pl.loops.FitLoop):\r\n @property\r\n def prefetch_batches(self):\r\n return 7 # lucky number 7\r\n\r\ntrainer = pl.Trainer(...)\r\ntrainer.fit_loop.replace(fit_loop=MyCustomLoop)<\/pre>\n
New Hooks<\/h3>\n
LightningModule.lr_scheduler_step<\/code><\/h3>\nLightningModule.lr_scheduler_step<\/code>\u00a0with your desired logic.<\/p>\nfrom timm.scheduler import TanhLRScheduler\r\n\r\n\r\nclass MyLightningModule(pl.LightningModule):\r\n def configure_optimizers(self):\r\n optimizer = ...\r\n scheduler = TanhLRScheduler(optimizer, ...)\r\n return {\"optimizer\": optimizer, \"lr_scheduler\": {\"scheduler\": scheduler, \"interval\": \"epoch\"}}\r\n\r\n def lr_scheduler_step(self, scheduler, optimizer_idx, metric):\r\n scheduler.step(epoch=self.current_epoch) # timm's scheduler need the epoch value<\/pre>\nA new stateful API<\/h4>\n
state_dict<\/code>\u00a0and\u00a0load_state_dict<\/code>, mimicking the PyTorch API. The new hooks receive their own component’s state and replace most usages of the previous\u00a0on_save_checkpoint<\/code>\u00a0and\u00a0on_load_checkpoint<\/code>\u00a0hooks.<\/p>\ndef MyCallback(pl.Callback):\r\n- def on_save_checkpoint(self, trainer, pl_module, checkpoint):\r\n- return {'x': self.x}\r\n \r\n- def on_load_checkpoint(self, trainer, pl_module, checkpoint):\r\n- self.x = x\r\n\r\n+ def state_dict(self):\r\n+ return {'x': self.x}\r\n \r\n+ def load_state_dict(self, checkpoint):\r\n+ self.x = x<\/pre>\nNew Properties<\/h3>\n
Trainer.estimated_stepping_batches<\/code><\/strong><\/p>\nTrainer.estimated_stepping_batches<\/code>\u00a0to compute the total number of stepping batches needed for the complete training.<\/p>\nclass MyLightningModule(pl.LightningModule):\r\n def configure_optimizers(self):\r\n optimizer = ...\r\n scheduler = torch.optim.lr_scheduler.OneCycleLR(\r\n optimizer, max_lr=1e-3, total_steps=self.trainer.estimated_stepping_batches\r\n )\r\n return {\"optimizer\": optimizer, \"lr_scheduler\": scheduler}<\/pre>\nTrainer.num_devices<\/code>\u00a0and\u00a0Trainer.device_ids<\/code><\/strong><\/p>\nTrainer<\/code>\u00a0configuration.<\/p>\n- num_devices = max(1, trainer.num_gpus, trainer.num_processes)\r\n- if trainer.tpu_cores:\r\n- num_devices = max(num_devices, trainer.tpu_cores)\r\n+ num_devices = trainer.num_devices<\/pre>\n
Experimental Features<\/h3>\n
$ PL_FAULT_TOLERANT_TRAINING=MANUAL python script.py<\/pre>\n
class MyLayerSync(pl.plugins.LayerSync):\r\n ...\r\n\r\nlayer_sync = MyLayerSync(...)\r\ntrainer = Trainer(sync_batchnorm=True, plugins=layer_sync, strategy=\"ddp\")<\/pre>\n
Accelerator<\/code>\u00a0classes in Lightning.<\/p>\nfrom pytorch_lightning.accelerators import Accelerator, AcceleratorRegistry\r\n\r\nclass SOTAAccelerator(Accelerator):\r\n def __init__(self, x):\r\n ...\r\n\r\nAcceleratorRegistry.register(\"sota_accelerator\", SOTAAccelerator, x=123)\r\n# the following works now:\r\ntrainer = Trainer(accelerator=\"sota_accelerator\")<\/pre>\n
Backward Incompatible Changes<\/h3>\n
Drop PyTorch 1.7 support<\/h4>\n
Drop Python 3.6 support<\/h4>\n
AcceleratorConnector<\/code>\u00a0rewrite<\/h4>\nAcceleratorConncetor<\/code>\u00a0class. No backwards compatibility was maintained so it is likely to have broken your code if it was using this class.<\/p>\nRe-define the\u00a0
current_epoch<\/code>\u00a0boundary<\/h4>\ntrainer.current_epoch<\/code>\u00a0is now increased by 1\u00a0on_train_end<\/code>. This means that if a model is run for 3 epochs (0, 1, 2),\u00a0trainer.current_epoch<\/code>\u00a0will now return 3 instead of 2 after\u00a0trainer.fit()<\/code>. This can also impact custom callbacks that acess this property inside this hook.<\/p>\non_train_epoch_end<\/code>). For example, a\u00a0Trainer(max_epochs=1, limit_train_batches=1)<\/code>\u00a0instance that saves a checkpoint will have the\u00a0current_epoch=0<\/code>\u00a0value saved instead of\u00a0current_epoch=1<\/code>.<\/p>\nRe-define the\u00a0
global_step<\/code>\u00a0boundary<\/h4>\ntrainer.global_step<\/code>\u00a0during an intra-training validation hook will now correctly return the number of optimizer steps taken already. In pseudocode:<\/p>\n training_step()\r\n+ global_step += 1\r\n validation_if_necessary()\r\n- global_step += 1<\/pre>\n
step=1.ckpt<\/code>\u00a0instead of\u00a0step=0.ckpt<\/code>.<\/p>\ntrainer.global_step<\/code>\u00a0value will now account for TBPTT or multiple optimizers. Users setting\u00a0Trainer({min,max}_steps=...)<\/code>\u00a0under these circumstances will need to adjust their values.<\/p>\nRemoved automatic reduction of outputs in\u00a0
training_step<\/code>\u00a0when using DataParallel<\/h4>\nTrainer(strategy=\"dp\")<\/code>,\u00a0all<\/em>\u00a0the tensors returned by training_step were previously reduced to a scalar (#11594<\/a>). This behavior was especially confusing when outputs needed to be collected into the\u00a0training_epoch_end<\/code>\u00a0hook.<\/p>\nloss<\/code>\u00a0tensor, unless you implement\u00a0training_step_end<\/code>, in which case the loss won’t get reduced either.<\/p>\nNo longer fallback to CPU with no devices<\/h4>\n
Trainer(accelerator=\"auto\")<\/code>\u00a0when this leniency is desired.<\/p>\nLIGHTNING CHANGELOG<\/h3>\n
Added<\/h4>\n
\n
MLFlowLogger<\/code>\u00a0(#12290<\/a>)<\/li>\nbackward_passes_per_step<\/code>\u00a0(#11911<\/a>)<\/li>\nDETAIL<\/code>\u00a0log level to provide useful logs for improving monitoring and debugging of batch jobs (#11008<\/a>)<\/li>\nSLURMEnvironment(auto_requeue=True|False)<\/code>\u00a0to control whether Lightning handles the requeuing (#10601<\/a>)<\/li>\n\n
_Stateful<\/code>\u00a0protocol to detect if classes are stateful (#10646<\/a>)<\/li>\n_FaultTolerantMode<\/code>\u00a0enum used to track different supported fault tolerant modes (#10645<\/a>)<\/li>\n_rotate_worker_indices<\/code>\u00a0utility to reload the state according the latest worker (#10647<\/a>)<\/li>\n_terminate_gracefully<\/code>\u00a0to all processes and add support for DDP (#10638<\/a>)<\/li>\n<\/ul>\n<\/li>\nDataLoaders<\/code>\u00a0returned in the\u00a0*_dataloader()<\/code>\u00a0methods, i.e., automatic replacement of samplers now works with custom types of\u00a0DataLoader<\/code>\u00a0(#10680<\/a>)<\/li>\nDataLoader<\/code>\u00a0implementation is not well implemented and we need to reconstruct it (#10719<\/a>)<\/li>\nLoop<\/code>‘s state by default in the checkpoint (#10784<\/a>)<\/li>\nLoop.replace<\/code>\u00a0to easily switch one loop for another (#10324<\/a>)<\/li>\n--lr_scheduler=ReduceLROnPlateau<\/code>\u00a0to the\u00a0LightningCLI<\/code>\u00a0(#10860<\/a>)<\/li>\nLightningCLI.configure_optimizers<\/code>\u00a0to override the\u00a0configure_optimizers<\/code>\u00a0return value (#10860<\/a>)<\/li>\nLightningCLI(auto_registry)<\/code>\u00a0flag to register all subclasses of the registerable components automatically (#12108<\/a>)<\/li>\nmax_epochs<\/code>\u00a0in the\u00a0Trainer<\/code>\u00a0is not set (#10700<\/a>)<\/li>\nLightningModule.configure_callbacks<\/code>\u00a0without wrapping it into a list (#11060<\/a>)<\/li>\nconsole_kwargs<\/code>\u00a0for\u00a0RichProgressBar<\/code>\u00a0to initialize inner Console (#10875<\/a>)<\/li>\nLightningCLI<\/code>\u00a0(