|
|
from typing import List, Union, Mapping, Dict, Any |
|
|
|
|
|
import torch.optim as optim |
|
|
from torch import Tensor |
|
|
from torch.distributed._shard.sharded_tensor import ShardedTensor |
|
|
|
|
|
|
|
|
class ShardedOptimizer(optim.Optimizer): |
|
|
def __init__( |
|
|
self, |
|
|
named_params: Mapping[str, Union[Tensor, ShardedTensor]], |
|
|
optimizer_class, |
|
|
*optimizer_args, |
|
|
**optimizer_kwargs |
|
|
): |
|
|
""" |
|
|
ShardedOptimizer collects all tensors and local shard tensors of |
|
|
ShardedTensor, then use these tensors as ``params`` for optimizers |
|
|
|
|
|
Args: |
|
|
named_params (Dict[str, Union[Tensor, ShardedTensor]]) : a Dict |
|
|
of parameters, where key is the parameter key, value is either |
|
|
Tensor or ShardedTensor parameter. |
|
|
optimizer_class (torch.optim.Optimizer): the Optimizer to use |
|
|
locally, i.e. torch.optim.SGD, torch.optim.Adagrad, etc. |
|
|
*optimizer_args: the arguments to initialize the optimizer. |
|
|
**optimizer_kwargs: the key-word arguments to initialize the optimizer. |
|
|
|
|
|
""" |
|
|
tensors: List[Tensor] = [] |
|
|
for value in named_params.values(): |
|
|
if isinstance(value, ShardedTensor): |
|
|
for local_shard in value.local_shards(): |
|
|
tensors.append(local_shard.tensor) |
|
|
else: |
|
|
tensors.append(value) |
|
|
|
|
|
self.named_params = named_params |
|
|
self._optim = optimizer_class(tensors, *optimizer_args, **optimizer_kwargs) |
|
|
self.param_groups = self._optim.param_groups |
|
|
self.state = self._optim.state |
|
|
|
|
|
def zero_grad(self, set_to_none: bool = False): |
|
|
r"""Sets the gradients of all optimized :class:`torch.Tensor` s to zero. |
|
|
|
|
|
Args: |
|
|
set_to_none (bool): instead of setting to zero, set the grads to None. |
|
|
This will in general have lower memory footprint, and can modestly improve performance. |
|
|
However, it changes certain behaviors. For example: |
|
|
1. When the user tries to access a gradient and perform manual ops on it, |
|
|
a None attribute or a Tensor full of 0s will behave differently. |
|
|
2. If the user requests ``zero_grad(set_to_none=True)`` followed by a backward pass, ``.grad``\ s |
|
|
are guaranteed to be None for params that did not receive a gradient. |
|
|
3. ``torch.optim`` optimizers have a different behavior if the gradient is 0 or None |
|
|
(in one case it does the step with a gradient of 0 and in the other it skips |
|
|
the step altogether). |
|
|
""" |
|
|
self._optim.zero_grad(set_to_none) |
|
|
|
|
|
def step(self, closure=None): |
|
|
r"""Performs a single optimization step (parameter update). |
|
|
|
|
|
Args: |
|
|
closure (Callable): A closure that reevaluates the model and |
|
|
returns the loss. Optional for most optimizers. |
|
|
|
|
|
.. note:: |
|
|
Unless otherwise specified, this function should not modify the |
|
|
``.grad`` field of the parameters. |
|
|
""" |
|
|
self._optim.step(closure) |
|
|
|
|
|
def state_dict(self) -> Dict[str, Any]: |
|
|
""" |
|
|
Returned state and param_groups will contain parameter keys |
|
|
instead of parameter indices like torch.optim.Optimizer. |
|
|
This allows for advanced functionality like optimizer re-sharding to be implemented. |
|
|
""" |
|
|
|
|
|
raise NotImplementedError("ShardedOptimizer state_dict not implemented yet!") |
|
|
|
|
|
|
|
|
def load_state_dict(self, state_dict: Mapping[str, Any]): |
|
|
r"""Loads the ShardedOptimizer state. |
|
|
|
|
|
Args: |
|
|
state_dict (dict): ShardedOptimizer state. Should be an object returned |
|
|
from a call to :meth:`state_dict`. |
|
|
""" |
|
|
|
|
|
raise NotImplementedError("ShardedOptimizer load_state_dict not implemented yet!") |
|
|
|
|
|
def add_param_group(self, param_group: Any): |
|
|
r"""Add a new param group |
|
|
""" |
|
|
|
|
|
raise NotImplementedError("ShardedOptimizer add_param_group not implemented yet!") |
|
|
|
