Upstream apply_optim_in_backward from TorchRec (#87397) (#88539)

Summary:

Upstreaming this as part of sharing common APIs. This is just a plain
move, any changes needed to support DDP / FSDP will come in follow up diffs.

Test Plan: CI

Reviewed By: zhaojuanmao

Differential Revision: D40564646

fbshipit-source-id: 619c434e02196812f8d4db1e40d07290e08b18f9
Pull Request resolved: #88539
Approved by: https://github.com/awgu

Author: rohan-varma

test/distributed/optim/test_apply_optimizer_in_backward.py

@@ -0,0 +1,113 @@
+# Owner(s): ["oncall: distributed"]
+
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+
+import unittest
+from copy import deepcopy
+
+import torch
+import torch.nn as nn
+
+from torch.distributed.optim import _apply_optimizer_in_backward
+
+# TODO (rohan-varma): Add FSDP & DDP tests once supported
+
+def _validate_params(params_list, fn):
+    ref_params = params_list[0]
+    for param_list in params_list[1:]:
+        for p1, p2 in zip(ref_params, param_list):
+            fn(p1, p2)
+
+
+class ApplyOverlappedOptimizerTest(unittest.TestCase):
+
+    def _run_training_loop_and_validate(self, inp, models, optimizers):
+        for i in range(6):
+            for model in models:
+                model(inp).sum().backward()
+            for opt in optimizers:
+                opt.step()
+
+            with self.subTest(i):
+                _validate_params(
+                    [model.parameters() for model in models],
+                    torch.testing.assert_allclose,
+                )
+
+            for opt in optimizers:
+                opt.zero_grad(set_to_none=True)
+
+    def _test_apply_optimizer_in_backward(self, share_params) -> None:
+        weight_optimizer_kwargs = {"lr": 1.0}
+        bias_optimizer_kwargs = {"lr": 0.5}
+        model = nn.Sequential(nn.Linear(10, 10), nn.Linear(10, 10))
+        if share_params:
+            model[0].weight = model[1].weight
+
+        # Use different optimizers for weights & biases.
+        weights = [m.weight for m in model]
+        biases = [m.bias for m in model]
+        optim_weight = torch.optim.SGD(weights, **weight_optimizer_kwargs)
+        optim_bias = torch.optim.SGD(biases, **bias_optimizer_kwargs)
+        model_with_opt_in_bwd = deepcopy(model)
+
+        # Apply different optimizer in backwards for weights and biases.
+        _apply_optimizer_in_backward(
+            torch.optim.SGD,
+            [m.weight for m in model_with_opt_in_bwd],
+            optimizer_kwargs=weight_optimizer_kwargs
+        )
+
+        _apply_optimizer_in_backward(
+            torch.optim.SGD,
+            [m.bias for m in model_with_opt_in_bwd],
+            optimizer_kwargs=bias_optimizer_kwargs
+        )
+
+        _validate_params(
+            [
+                model.parameters(),
+                model_with_opt_in_bwd.parameters(),
+            ],
+            torch.testing.assert_allclose,
+        )
+
+        self._run_training_loop_and_validate(
+            torch.randn(4, 10),
+            [model, model_with_opt_in_bwd],
+            [optim_weight, optim_bias],
+        )
+
+    def test_apply_optimizer_in_backward(self) -> None:
+        self._test_apply_optimizer_in_backward(share_params=False)
+
+    def test_apply_optimizer_in_backward_shared_params(self) -> None:
+        self._test_apply_optimizer_in_backward(share_params=True)
+
+    def test_multiple_optim_for_params(self) -> None:
+        model = nn.Sequential(nn.Linear(10, 10), nn.Linear(10, 10))
+        opt_0_kwargs = {"lr": 0.03}
+        opt_1_kwargs = {"lr": 0.01}
+        opt_0 = torch.optim.SGD(model.parameters(), **opt_0_kwargs)
+        opt_1 = torch.optim.SGD(model.parameters(), **opt_1_kwargs)
+        model_with_opt_in_bwd = deepcopy(model)
+        _apply_optimizer_in_backward(
+            torch.optim.SGD,
+            model_with_opt_in_bwd.parameters(),
+            optimizer_kwargs=opt_0_kwargs,
+        )
+        _apply_optimizer_in_backward(
+            torch.optim.SGD,
+            model_with_opt_in_bwd.parameters(),
+            optimizer_kwargs=opt_1_kwargs,
+        )
+        self._run_training_loop_and_validate(
+            torch.randn(4, 10),
+            [model, model_with_opt_in_bwd],
+            [opt_0, opt_1],
+        )

torch/distributed/optim/__init__.py

@@ -17,6 +17,7 @@
 from .functional_rprop import _FunctionalRprop
 from .functional_adamax import _FunctionalAdamax
 from .utils import as_functional_optim
+from .apply_optimizer_in_backward import _apply_optimizer_in_backward
 
 
 # DistributedOptimizer imports torch.distributed.rpc names, so gate availability

torch/distributed/optim/apply_optimizer_in_backward.py

@@ -0,0 +1,78 @@
+from typing import Any, Dict, Iterable, Type, List, no_type_check
+
+import torch
+
+__all__: List[str] = []
+
+@no_type_check
+def _apply_optimizer_in_backward(
+    optimizer_class: Type[torch.optim.Optimizer],
+    params: Iterable[torch.nn.Parameter],
+    optimizer_kwargs: Dict[str, Any],
+) -> None:
+    """
+    Upon ``backward()``, parameters will fire the corresponding optimizer.
+
+    Note - gradients for these parameters will be set to None after ``backward()``.
+    This means that any other (non applied) optimizer over this parameter will be
+    a no-op.
+
+    Args:
+        optimizer_class: (Type[torch.optim.Optimizer]): Optimizer to apply to parameter
+        params: (Iterator[nn.Parameter]): parameters to apply optimizer state to
+        optimizer_kwargs: (Dict[str, Any]): kwargs to pass to optimizer constructor
+
+    Example::
+        params_generator = model.parameters()
+        param_1 = next(params_generator)
+        remainder_params = list(params_generator)
+
+        apply_optimizer_in_backward(torch.optim.SGD, [param_1], {"lr": .02})
+        apply_optimizer_in_backward(torch.optim.Adam, remainder_params, {"lr": .04})
+
+        model(...).sum().backward() # after backward, parameters will already
+        # have their registered optimizer applied.
+
+    """
+
+    @no_type_check
+    def _apply_optimizer_in_backward_to_param(param: torch.nn.Parameter) -> None:
+        # view_as creates a node in autograd graph that allows us access to the
+        # parameter's AccumulateGrad autograd function object. We register a
+        # hook on this object to fire the optimizer when the gradient for
+        # this parameter is ready (has been accumulated into .grad field)
+
+        # Don't create a new acc_grad if we already have one
+        # i.e.f or shared parameters or attaching multiple optimizers to a param.
+        if not hasattr(param, 'acc_grad'):
+            acc_grad = param.view_as(param).grad_fn.next_functions[0][0]
+        else:
+            acc_grad = param._acc_grad
+
+        optimizer = optimizer_class([param], **optimizer_kwargs)
+
+        # Keep the grad accumulator around for the lifetime of the Tensor,
+        # store it on the param to avoid uncollectable ref-cycle
+        if not hasattr(param, 'acc_grad'):
+            param._acc_grad = acc_grad  # type: ignore[attr-defined]
+
+        if not hasattr(param, '_in_backward_optimizers'):
+            param._in_backward_optimizers = []  # type: ignore[attr-defined]
+            # TODO: investigate whether we really need these attributes.
+            param._optimizer_classes = []  # type: ignore[attr-defined]
+            param._optimizer_kwargs = []  # type: ignore[attr-defined]
+
+        param._in_backward_optimizers.append(optimizer)  # type: ignore[attr-defined]
+        param._optimizer_classes.append(optimizer_class)  # type: ignore[attr-defined]
+        param._optimizer_kwargs.append(optimizer_kwargs)  # type: ignore[attr-defined]
+
+        def optimizer_hook(*_unused) -> None:
+            for opt in param._in_backward_optimizers:  # type: ignore[attr-defined]
+                opt.step()
+
+            param.grad = None
+
+        param._acc_grad.register_hook(optimizer_hook)  # type: ignore[attr-defined]
+
+    for param in params:
+        _apply_optimizer_in_backward_to_param(param)