Upgrade guide to use lightning 2.0

Replace trainer.strategy=null with trainer.strategy=auto as lightning 2.0 doesn’t have None strategy.

Remove resume_from_checkpoint if being used as a trainer flag and pass the path to Trainer.fit(ckpt_path=”…”) method.

Set trainer.strategy = "ddp_find_unused_parameters_true" if there are unused parameters in your model as lightning 2.0 has find_unused_parameters as False by default.

If used Trainer’s flag replace_sampler_ddp replace it with use_distributed_sampler.

If using CheckpointConnector replace it with _CheckpointConnector.

To set or get ckpt_path use trainer.ckpt_path directly instead of calling protected API via trainer._checkpoint_connector._ckpt_path or using trainer._checkpoint_connector.resume_from_checkpoint_fit_path.

Change import load from pytorch_lightning.utilities.cloud_io to import _load.

If used from pytorch_lightning.plugins.precision.native_amp import NativeMixedPrecisionPlugin from replace it with from pytorch_lightning.plugins.precision import MixedPrecisionPlugin.

Lightning 2.0 adds '16-mixed', 'bf16-mixed' as the preicison values for fp16 mixed precision and bf16 mixed precision respectively.

Lightning 2.0 renames epoch end hooks from training_epoch_end, validation_epoch_end, test_epoch_end to on_train_epoch_end, on_validation_epoch_end, on_test_epoch_end. The renamed hooks do not accept the outputs arg but instead outputs needs to be defined as an instance variable of the model class to which the outputs of the step needs to be manually appended. More detailed examples implementing this can be found under migration guide of lightning’s PR 16520. Example from NeMo can be found here.

Lightning 2.0 is not currently supporting multiple dataloders for validation and testing in case of dataloader_iter. The support for this will be added back soon in an upcoming release. If dataloader_iter is being used and your config passes multiple files to validation_ds.file_names or test_ds.file_names, please use just one file until this issue is fixed with pytorch lightning.

With lightning 2.0 it’s required to set limit_val_batches and num_sanity_val_steps to be a multiple of number of microbatches while using dataloader_iter (applies only to Megatron files that use dataloader_iter) for all pretraining files (not downstream tasks like finetuning). This is being taken care internally in NeMo and does not require anything to be done by the user. However, if you are a developer of NeMo and are building a new model for pretraining that uses dataloader_iter instead of batch in validation_step methods please make sure to call self._reconfigure_val_batches() in build_train_valid_test_datasets method of your model.

If model is being wrapped with LightningDistributedModule in configure_ddp method please replace it with _LightningModuleWrapperBase as being done here: NeMo upgrade to lightning 2.0 PR.

If using pre_configure_ddp() in your DDP, remove it as it’s not required anymore. NeMo upgrade to lightning 2.0 PR.

If any of the tests use CPU as the device, ensure to explicitly pass it in the trainer as trainer = pl.Trainer(max_epochs=1, accelerator='cpu') since deafult val in PTL >= 2.0 is auto and it picks cuda.

If using from pytorch_lightning.loops import TrainingEpochLoop, replace TrainingEpochLoop with _TrainingEpochLoop.

If using trainer.fit_loop.max_steps, replace it with trainer.fit_loop.epoch_loop.max_steps.