Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Introduce Stateful PrecisionPlugin #11638

Merged
merged 20 commits into from
Feb 14, 2022
Merged
Changes from 11 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
@@ -72,6 +72,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
- Added a `MisconfigurationException` if user provided `opt_idx` in scheduler config doesn't match with actual optimizer index of its respective optimizer ([#11247](https://github.com/PyTorchLightning/pytorch-lightning/pull/11247))


- Added `_Stateful` support for `PrecisionPlugin` ([#11638](https://github.com/PyTorchLightning/pytorch-lightning/pull/11638))


### Changed

- Set the `prog_bar` flag to False in `LightningModule.log_grad_norm` ([#11472](https://github.com/PyTorchLightning/pytorch-lightning/pull/11472))
18 changes: 15 additions & 3 deletions pytorch_lightning/plugins/precision/apex_amp.py
Original file line number Diff line number Diff line change
@@ -21,6 +21,7 @@
from pytorch_lightning.plugins.precision.mixed import MixedPrecisionPlugin
from pytorch_lightning.utilities import _APEX_AVAILABLE, AMPType
from pytorch_lightning.utilities.exceptions import MisconfigurationException
from pytorch_lightning.utilities.rank_zero import rank_zero_deprecation
from pytorch_lightning.utilities.types import _PARAMETERS

if _APEX_AVAILABLE:
@@ -92,9 +93,20 @@ def optimizer_step(
if not isinstance(model, pl.LightningModule) or not model.automatic_optimization or not skipped_backward:
optimizer.step(**kwargs)

def state_dict(self) -> Dict[str, Any]:
return amp.state_dict()

def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
amp.load_state_dict(state_dict)

def on_load_checkpoint(self, checkpoint: Dict[str, Any]) -> None:
carmocca marked this conversation as resolved.
Show resolved Hide resolved
if "amp_scaling_state" in checkpoint:
amp.load_state_dict(checkpoint["amp_scaling_state"])
rank_zero_deprecation(
"`ApexMixedPrecisionPlugin.on_load_checkpoint` is deprecated in v1.6 and will be removed in v1.8."
" Use `ApexMixedPrecisionPlugin.load_state_dict` instead."
)

def on_save_checkpoint(self, checkpoint: Dict[str, Any]) -> None:
checkpoint["amp_scaling_state"] = amp.state_dict()
rank_zero_deprecation(
"`ApexMixedPrecisionPlugin.on_save_checkpoint` is deprecated in v1.6 and will be removed in v1.8."
" Use `ApexMixedPrecisionPlugin.state_dict` instead."
)
22 changes: 18 additions & 4 deletions pytorch_lightning/plugins/precision/native_amp.py
Original file line number Diff line number Diff line change
@@ -23,6 +23,7 @@
from pytorch_lightning.plugins.precision.mixed import MixedPrecisionPlugin
from pytorch_lightning.utilities import _TORCH_GREATER_EQUAL_1_10, AMPType
from pytorch_lightning.utilities.exceptions import MisconfigurationException
from pytorch_lightning.utilities.rank_zero import rank_zero_deprecation

if _TORCH_GREATER_EQUAL_1_10:
from torch import autocast as new_autocast
@@ -106,10 +107,23 @@ def forward_context(self) -> Generator[None, None, None]:
with self.autocast_context_manager():
yield

def state_dict(self) -> Dict[str, Any]:
if self.scaler is not None:
return self.scaler.state_dict()
return {}

def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
if self.scaler is not None:
self.scaler.load_state_dict(state_dict)

def on_load_checkpoint(self, checkpoint: Dict[str, Any]) -> None:
if self.scaler is not None and "native_amp_scaling_state" in checkpoint:
self.scaler.load_state_dict(checkpoint["native_amp_scaling_state"])
rank_zero_deprecation(
jjenniferdai marked this conversation as resolved.
Show resolved Hide resolved
"`NativeMixedPrecisionPlugin.on_load_checkpoint` is deprecated in v1.6 and will be removed in v1.8."
" Use `NativeMixedPrecisionPlugin.load_state_dict` instead."
)

def on_save_checkpoint(self, checkpoint: Dict[str, Any]) -> None:
if self.scaler is not None:
checkpoint["native_amp_scaling_state"] = self.scaler.state_dict()
rank_zero_deprecation(
"`NativeMixedPrecisionPlugin.on_save_checkpoint` is deprecated in v1.6 and will be removed in v1.8."
" Use `NativeMixedPrecisionPlugin.state_dict` instead."
)
19 changes: 18 additions & 1 deletion pytorch_lightning/plugins/precision/precision_plugin.py
Original file line number Diff line number Diff line change
@@ -13,7 +13,7 @@
# limitations under the License.
import contextlib
from functools import partial
from typing import Any, Callable, Generator, List, Optional, Tuple, Union
from typing import Any, Callable, Dict, Generator, List, Optional, Tuple, Union

import torch
from torch import Tensor
@@ -242,3 +242,20 @@ def teardown(self) -> None:

It is the right place to release memory and free other resources.
"""

def state_dict(self) -> Dict[str, Any]:
jjenniferdai marked this conversation as resolved.
Show resolved Hide resolved
"""Called when saving a checkpoint, implement to generate precision plugin state_dict.

Returns:
A dictionary containing precision plugin state.
"""
return {}

def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
"""Called when loading a checkpoint, implement to reload precision plugin state given precision plugin
state_dict.

Args:
state_dict: the precision plugin state returned by ``state_dict``.
"""
pass
28 changes: 24 additions & 4 deletions pytorch_lightning/trainer/connectors/checkpoint_connector.py
Original file line number Diff line number Diff line change
@@ -23,6 +23,7 @@
import pytorch_lightning as pl
from pytorch_lightning.loops.utilities import _is_max_limit_reached
from pytorch_lightning.plugins.environments import SLURMEnvironment
from pytorch_lightning.plugins.precision import ApexMixedPrecisionPlugin, NativeMixedPrecisionPlugin
from pytorch_lightning.trainer.states import TrainerFn
from pytorch_lightning.utilities import _OMEGACONF_AVAILABLE, rank_zero_deprecation, rank_zero_info, rank_zero_warn
from pytorch_lightning.utilities.cloud_io import get_filesystem
@@ -192,7 +193,7 @@ def restore_training_state(self) -> None:
return

# restore precision plugin (scaler etc.)
self.trainer.precision_plugin.on_load_checkpoint(self._loaded_checkpoint)
self.restore_precision_plugin_state()

# restore loops and their progress
self.restore_loops()
@@ -202,6 +203,21 @@ def restore_training_state(self) -> None:
# restore optimizers and schedulers state
self.restore_optimizers_and_schedulers()

def restore_precision_plugin_state(self) -> None:
"""Restore the precision plugin state from the pre-loaded checkpoint."""
prec_plugin = self.trainer.precision_plugin
prec_plugin.on_load_checkpoint(self._loaded_checkpoint)
if prec_plugin.__class__.__qualname__ in self._loaded_checkpoint:
prec_plugin.load_state_dict(self._loaded_checkpoint[prec_plugin.__class__.__qualname__])

# old checkpoints compatibility
if "amp_scaling_state" in self._loaded_checkpoint and isinstance(prec_plugin, ApexMixedPrecisionPlugin):
prec_plugin.load_state_dict(self._loaded_checkpoint["amp_scaling_state"])
if "native_amp_scaling_state" in self._loaded_checkpoint and isinstance(
prec_plugin, NativeMixedPrecisionPlugin
):
prec_plugin.load_state_dict(self._loaded_checkpoint["native_amp_scaling_state"])

def restore_callbacks(self) -> None:
"""Restores all callbacks from the pre-loaded checkpoint."""
if not self._loaded_checkpoint:
@@ -324,9 +340,8 @@ def dump_checkpoint(self, weights_only: bool = False) -> dict:
'callbacks': "callback specific state"[] # if not weights_only
'optimizer_states': "PT optim's state_dict"[] # if not weights_only
'lr_schedulers': "PT sched's state_dict"[] # if not weights_only
'native_amp_scaling_state': PT amp's state_dict # if not weights_only and use native amp
'amp_scaling_state': Apex's state_dict # if not weights_only and use apex amp
'state_dict': Model's state_dict (e.g. network weights)
precision_plugin.__class__.__qualname__: precision plugin state_dict # if not weights_only
CHECKPOINT_HYPER_PARAMS_NAME:
CHECKPOINT_HYPER_PARAMS_KEY:
CHECKPOINT_HYPER_PARAMS_TYPE:
@@ -372,7 +387,12 @@ def dump_checkpoint(self, weights_only: bool = False) -> dict:
lr_schedulers.append(config.scheduler.state_dict())
checkpoint["lr_schedulers"] = lr_schedulers

self.trainer.precision_plugin.on_save_checkpoint(checkpoint)
# precision plugin
prec_plugin = self.trainer.precision_plugin
prec_plugin_state_dict = prec_plugin.state_dict()
if prec_plugin_state_dict:
checkpoint[prec_plugin.__class__.__qualname__] = prec_plugin_state_dict
prec_plugin.on_save_checkpoint(checkpoint)

# dump hyper-parameters
if model.hparams:
6 changes: 2 additions & 4 deletions tests/models/test_hooks.py
Original file line number Diff line number Diff line change
@@ -495,10 +495,8 @@ def training_step(self, batch, batch_idx):
"state_dict": ANY,
"loops": ANY,
}
if kwargs.get("amp_backend") == "native":
saved_ckpt["native_amp_scaling_state"] = ANY
elif kwargs.get("amp_backend") == "apex":
saved_ckpt["amp_scaling_state"] = ANY
if kwargs.get("amp_backend") == "native" or kwargs.get("amp_backend") == "apex":
saved_ckpt[trainer.precision_plugin.__class__.__qualname__] = ANY
device = torch.device("cuda:0" if "gpus" in kwargs else "cpu")
expected = [
dict(name="Callback.on_init_start", args=(trainer,)),