Lightning-AI/pytorch-lightning 1.8.0

Lightning 1.8: Colossal-AI Strategy, Commands and Secrets for Apps, FSDP Improvements and More!

on GitHub

The core team is excited to announce the release of Lightning 1.8 ⚡

• Highlights
• Backward Incompatible Changes
• Deprecations
• Full Changelog
• Contributors

Lightning v1.8 is the culmination of work from 52 contributors who have worked on features, bug-fixes, and documentation for a total of over 550+ commits since v1.7.

Highlights

Colossal-AI

Colossal-AI focuses on improving efficiency when training large-scale AI models with billions of parameters. With the new Colossal-AI strategy in Lightning 1.8, you can train existing models like GPT-3 with up to half as many GPUs as usually needed. You can also train models up to twice as big with the same number of GPUs, saving you significant cost. Here is how you use it:

# Select the strategy with good defaults
trainer = Trainer(strategy="colossalai")

# or tune parameters to your liking
from lightning.pytorch.strategies import ColossalAIStrategy

trainer = Trainer(strategy=ColossalAIStrategy(placement_policy="cpu", ...))

You can find Colossal-AI's benchmarks with Lightning on GPT-2 here.

Under the hood, Colossal-AI implements different parallelism algorithms that are especially interesting for the development of SOTA transformer models:

• Data Parallelism
• Pipeline Parallelism
• 1D, 2D, 2.5D, 3D Tensor Parallelism
• Sequence Parallelism
• Zero Redundancy Optimization

Learn how to install and use Colossal-AI effectively with Lightning here.

NOTE: This strategy is marked as experimental. Stay tuned for more updates in the future.

Secrets for Lightning Apps

Introducing encrypted secrets (#14612), a feature requested by Lightning App users 🎉!

Encrypted secrets allow you to securely pass private data to your apps, like API keys, access tokens, database passwords, or other credentials, without exposing them in your code.

1. Add a secret to your Lightning account in lightning.ai (read more here)

2. Add an environment variable to your app to read the secret:

   # somewhere in your Flow or Work:
   GitHubComponent(api_token=os.environ["API_TOKEN"])

3. Pass the secret to your app run with the following command:

   lightning run app app.py --cloud --secret API_TOKEN=github_api_token

These secrets are encrypted and stored in the Lightning database. Nothing except your app can access the value.

NOTE: This is an experimental feature.

CLI Commands for Lightning Apps

Introducing CLI commands for apps (#13602)!
As a Lightning App builder, if you want to easily create a CLI interface for users to interract with your app, then this is for you.

Here is an example where users can dynamically create notebooks from the CLI.
All you need to do is implement the configure_commands hook on the LightningFlow:

import lightning as L
from commands.notebook.run import RunNotebook


class Flow(L.LightningFlow):
    ...

    def configure_commands(self):
        # Return a list of dictionaries with commands:
        return [{"run notebook": RunNotebook(method=self.run_notebook)}]


app = L.LightningApp(Flow())

Once the app is running with lightning run app app.py, you can connect to the app with the following command:

lightning connect {app name} -y

and run the command that was configured:

lightning run notebook --name=my_notebook_name

For a full tutorial and running example, visit our docs. TODO: add to docs
NOTE: This is an experimental feature.

Auto-wrapping for FSDP Strategy

In Lightning v1.7, we introduced an integration for PyTorch FSDP in the form of our FSDP strategy, which allows you to train huge models with billions of parameters sharded across hundreds of GPUs and machines.

# Native FSDP implementation
trainer = Trainer(strategy="fsdp_native")

We are continuing to improve the support for this feature by adding automatic wrapping of layers for use cases where the model fits into CPU memory, but not into GPU memory (#14383).

Here are some examples:

Case 1: Model is so large that it does not fit into CPU memory.
Construct your layers in the configure_sharded_model hook and wrap the large ones you want to shard across GPUs:

class MassiveModel(LightningModule):
    ...
    
    # Create model here and wrap the large layers for sharding
    def configure_sharded_model(self):
        for i, layer in enumerate(self.block):
            self.block[i] = wrap(layer)
        ...

Case 2: Model fits into CPU memory, but not into GPU memory. In Lightning v1.8, you no longer need to do anything special here, as we can automatically wrap the layers for you using FSDP's policy:

model = MassiveModel()
trainer = Trainer(
    accelerator="gpu", 
    devices=8, 
    strategy="fsdp_native",  # or strategy="fsdp" for fairscale
    precision=16
)

# Automatically wraps the layers here:
trainer.fit(model)

Case 3: Model fits into GPU memory. No action required, use any strategy you want.

Note: if you want to manually wrap layers for more control, you can still do that!

Read more about FSDP and how layer wrapping works in our docs.

New Tuner Callbacks

In this release, we focused on Tuner improvements and introduced two new callbacks that can help you customize the batch size finder and learning rate finder as per your use case.

Batch Size Finder (#11089)

1. You can customize the BatchSizeFinder callback to run at different epochs. This feature is useful while fine-tuning models since you can't always use the same batch size after unfreezing the backbone.

   from lightning.pytorch.callbacks import BatchSizeFinder
   
   
   class FineTuneBatchSizeFinder(BatchSizeFinder):
       def __init__(self, milestones, *args, **kwargs):
           super().__init__(*args, **kwargs)
           self.milestones = milestones
   
       def on_fit_start(self, *args, **kwargs):
           return
   
       def on_train_epoch_start(self, trainer, pl_module):
           if trainer.current_epoch in self.milestones or trainer.current_epoch == 0:
               self.scale_batch_size(trainer, pl_module)
   
   
   trainer = Trainer(callbacks=[FineTuneBatchSizeFinder(milestones=(5, 10))])
   trainer.fit(...)

2. Run batch size finder for validate/test/predict.

   from lightning.pytorch.callbacks import BatchSizeFinder
   
   
   class EvalBatchSizeFinder(BatchSizeFinder):
       def __init__(self, *args, **kwargs):
           super().__init__(*args, **kwargs)
   
       def on_fit_start(self, *args, **kwargs):
           return
   
       def on_test_start(self, trainer, pl_module):
           self.scale_batch_size(trainer, pl_module)
   
   
   trainer = Trainer(callbacks=[EvalBatchSizeFinder()])
   trainer.test(...)

Learning Rate Finder (#13802)

You can now use the LearningRateFinder callback to run at different intervals. This feature is useful when fine-tuning models, for example.

from lightning.pytorch.callbacks import LearningRateFinder


class FineTuneLearningRateFinder(LearningRateFinder):
    def __init__(self, milestones, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.milestones = milestones

    def on_fit_start(self, *args, **kwargs):
        return

    def on_train_epoch_start(self, trainer, pl_module):
        if trainer.current_epoch in self.milestones or trainer.current_epoch == 0:
            self.lr_find(trainer, pl_module)

trainer = Trainer(callbacks=[FineTuneLearningRateFinder(milestones=(5, 10))])
trainer.fit(...)

LightningCLI Improvements

Even though the LightningCLI class is designed to help in the implementation of command line tools, there are instances when it might be more desirable to run directly from Python. In Lightning 1.8, you can now do this (#14596):

from lightning.pytorch.cli import LightningCLI

def cli_main(args):
    cli = LightningCLI(MyModel, ..., args=args)
    ...

Anywhere in your program, you can now call the CLI directly:

cli_main(["--trainer.max_epochs=100", "--model.encoder_layers=24"])

Learn about all features of the LightningCLI!

Improvements to the SLURM Support

Multi-node training on a SLURM cluster has been supported since the inception of Lightning Trainer, and has seen several improvements over time thanks to many community contributions. And we just keep going! In this release, we've added two quality of life improvements:

• The preemption/termination signal is now configurable (#14626):

  # the default signal is SIGUSR1
  trainer = Trainer(plugins=[SLURMEnvironment(requeue_signal=signal.SIGUSR1)])
  
  # customize it for your cluster
  trainer = Trainer(plugins=[SLURMEnvironment(requeue_signal=signal.SIGHUP)])

• Automatic requeuing of jobs now also works for array jobs (#15040)! Array jobs are a convenient way to group/launch several scripts at once. When the SLURM scheduler interrupts your jobs, Lightning will save a checkpoint, resubmit a new job, and, once the scheduler allocates resources, the Trainer will resume from where it left off.

Read more about our SLURM integration here.

Backward Incompatible Changes

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 CHANGELOG below.

Callback hooks for loading and saving checkpoints

The signature and behavior of the on_load_checkpoint and on_save_checkpoint callback hooks have changed (#14835):

Before:

def on_save_checkpoint(self, trainer, pl_module, checkpoint):
    ...
    # previously, we were able to return state here
    return state

def on_load_checkpoint(self, trainer, pl_module, callback_state):
    # previously, only the state for this callback was passed in as argument
    ...

Now:

def on_save_checkpoint(self, trainer, pl_module, checkpoint):
    ...
    # returning a value here is no longer supported
    # you can modify the checkpoint dict directly
    return None


def state_dict(self):
    ...
    # Now, return state from this new method
    return state


def on_load_checkpoint(self, trainer, pl_module, checkpoint):
    # previously, only the state for this callback was passed in as argument
    ...
    
    
def load_state_dict(self, state):
    # Now, the state for this callback gets passed to this new method
    ...

DataModule hooks for loading and saving checkpoints

The on_save_checkpoint and on_load_checkpoint hooks on the LightningDataModule have been removed in favor of the state_dict and load_state_dict methods:

-def on_save_checkpoint(self, checkpoint):
-    checkpoint["banana"] = self.banana
+def state_dict(self):
+    return dict(banana=self.banana)


-def on_load_checkpoint(self, checkpoint):
-    self.banana = checkpoint["banana"]
+def load_state_dict(self, state):
+    self.banana = state["banana"]

Callback hooks

We removed some Callback hooks that were ambiguous to use Removed deprecated callback hooks (#14834):

Trainer Device Attributes

We cleaned up the properties related to device indices (#14829).

The attributes Trainer.{devices,gpus,num_gpus,ipus,tpu_cores,num_processes,root_gpu,data_parallel_device_ids} have been removed in favor of accelerator-agnostic attributes:

trainer = Trainer(...)

# access the number of devices the trainer uses on this machine ...
print(trainer.num_devices)

# ... or the device IDs
print(trainer.device_ids)

Setting the torch-distributed backend

In previous versions of Lightning, switching between the "gloo" and "nccl" backends for multi-GPU, multi-node training was possible through setting an environment variable like so:

PL_TORCH_DISTRIBUTED_BACKEND="gloo" python train.py

But not all strategies support changing the backend in this way.
From now on, the backend has to be set in the code (#14693):

trainer = Trainer(strategy=DDPStrategy(process_group_backend="gloo"))

The default remains "nccl", and you should choose "gloo" only for debugging purposes.

Logging with multiple loggers

Logging with multiple loggers can be super useful (and super easy with Lightning). For example, you could be using one logger to record sensitive image logs to a hosted MLFlow server within your organization, and at the same time log loss curves online to WandB.

trainer = Trainer(
    loggers=[WandbLogger(...), MLFlowLogger(...)]
)

Here are two major changes that apply when using multiple loggers in 1.8:

• Checkpoints and profiler reports no longer go to a strange folder with a long, hard to remember name (#14325). From now on, these arifacts will land in the version folder of the first logger in the list.

• The loggers used to be wrapped by a LoggerCollection object, so that when you accessed trainer.logger you could log to all of them simultaneously. However, this "magic" caused confusion and errors among users and we decided to simplify this (#14283):

  # now returns the first logger in the list
  print(trainer.logger)
  
  # access all loggers in a list with plural
  loggers = trainer.loggers
  
  for logger in loggers:
      logger.do_something()

Deprecations

Why is Lightning deprecating APIs in every release?

Many users have this question, and it is a fair one! Deprecations are a normal part of API evolution in all software. We continually improve Lightning, which means we make APIs like class names, methods, hooks and arguments clear, easy to remember, and general enough to adopt more functionality in the future. Sometimes we have to let old things go to build new and better products.

Learn more about our deprecation window here.

So far, we have followed the pattern of removing deprecated functionality and APIs after two minor versions of deprecation. From Lightning 1.8 onward, we will additionaly convert warnings to error messages after the deprecation phase ends. This way, we can greatly improve the upgrade experience with helpful messages for users who skip more than two minor Lightning versions. The exception to this rule are experimental features, which are marked as such in our documentation.

Here is a summary of major deprecations introduced in 1.8:

CHANGELOG

Lightning App

Lightning Trainer

Full commit list: 1.7.0...1.8.0

Contributors

Veteran

@akihironitta @ananthsub @AndresAlgaba @ar90n @Atharva-Phatak @awaelchli @BongYang @Borda @carmocca @dependabot @donlapark @ethanwharris @Felonious-Spellfire @hhsecond @jerome-habana @JustinGoheen @justusschock @kaushikb11 @krishnakalyan3 @krshrimali @luca-medeiros @manangoel99 @manskx @mauvilsa @MrShevan @nicolai86 @nmiculinic @otaj @Queuecumber @rlizzo @rohitgr7 @rschireman @SeanNaren @speediedan @tchaton @tshu-w

New

@Birch-san @clementpoiret @HalestormAI @thongonary @alecmerdler @adam-lightning @yurijmikhalevich @lijm1358 @robert-s-lee @panos-is @kacperlukawski @alro923 @dmitsf @Anner-deJong @cschell @nishantb06 @Callidior @j0rd1smit @MarcSkovMadsen @KralaBenjamin @robertomest @daniel347x @pierocor @datumbox @nohalon @pritamsoni-hsr @nandwalritik @gilfree @ritsuki1227 @christopher-nguyen-re @JulesGM @jgbos @dconathan @jsr-p @NeoKish @Blaizzy @suyash-811 @alexkuzmik @ziyadsheeba @geoffrey-g-delhomme @amrutha1098 @AlessioQuercia @ver217 @Helias @zxvix @1SAA @fabiofumarola @luca3rd @kimpty @PaulLerner @rbracco @wouterzwerink

If we forgot somebody or you have a suggestion, find support here ⚡

Did you know?

Chuck Norris can write functions of infinite recursion ... and have them return.