Fix unintended deprecation warning in torch.distributed.optim (pytorc…

…h#140889)

We have a deprecation warning for scripted functional optimizer at module level in `torch/distributed/optim/__init__.py`. However, not all optimizers exposed by the module are scripted functional optimizers, causing some false deprecation warning (e.g. pytorch#139661).

This PR moves the deprecation warning to the `__init__` functions of the deprecated scripted functional optimizers.

Pull Request resolved: pytorch#140889
Approved by: https://github.com/d4l3k, https://github.com/kwen2501, https://github.com/XilunWu

torch/distributed/optim/__init__.py

@@ -26,16 +26,6 @@
 from .utils import as_functional_optim
 
 
-with warnings.catch_warnings():
-    warnings.simplefilter("always")
-    warnings.warn(
-        "`TorchScript` support for functional optimizers is deprecated "
-        "and will be removed in a future PyTorch release. "
-        "Consider using the `torch.compile` optimizer instead.",
-        DeprecationWarning,
-        stacklevel=2,
-    )
-
 # DistributedOptimizer imports torch.distributed.rpc names, so gate availability
 # based on RPC being available.
 if hasattr(torch._C, "_rpc_init"):

torch/distributed/optim/_deprecation_warning.py

@@ -0,0 +1,16 @@
+import warnings
+
+import torch
+
+
+@torch.jit.ignore  # type: ignore[misc]
+def _scripted_functional_optimizer_deprecation_warning(stacklevel: int = 0) -> None:
+    with warnings.catch_warnings():
+        warnings.simplefilter("always")
+        warnings.warn(
+            "`TorchScript` support for functional optimizers is deprecated "
+            "and will be removed in a future PyTorch release. "
+            "Consider using the `torch.compile` optimizer instead.",
+            DeprecationWarning,
+            stacklevel=stacklevel + 2,
+        )

torch/distributed/optim/functional_adadelta.py

@@ -4,6 +4,9 @@
 import torch
 import torch.optim._functional as F
 from torch import Tensor
+from torch.distributed.optim._deprecation_warning import (
+    _scripted_functional_optimizer_deprecation_warning,
+)
 
 
 __all__: List[str] = []
@@ -31,6 +34,7 @@ def __init__(
         maximize: bool = False,
         _allow_empty_param_list: bool = False,
     ):
+        _scripted_functional_optimizer_deprecation_warning(stacklevel=2)
         self.defaults = {
             "lr": lr,
             "rho": rho,

torch/distributed/optim/functional_adagrad.py

@@ -4,6 +4,9 @@
 import torch
 import torch.optim._functional as F
 from torch import Tensor
+from torch.distributed.optim._deprecation_warning import (
+    _scripted_functional_optimizer_deprecation_warning,
+)
 
 
 __all__: List[str] = []
@@ -36,6 +39,7 @@ def __init__(
         maximize: bool = False,
         _allow_empty_param_list: bool = False,
     ):
+        _scripted_functional_optimizer_deprecation_warning(stacklevel=2)
         self.defaults = {
             "lr": lr,
             "lr_decay": lr_decay,

torch/distributed/optim/functional_adam.py

@@ -4,6 +4,9 @@
 import torch
 import torch.optim._functional as F
 from torch import Tensor
+from torch.distributed.optim._deprecation_warning import (
+    _scripted_functional_optimizer_deprecation_warning,
+)
 
 
 __all__: List[str] = []
@@ -33,6 +36,7 @@ def __init__(
         fused: bool = False,
         _allow_empty_param_list: bool = False,
     ):
+        _scripted_functional_optimizer_deprecation_warning(stacklevel=2)
         if not 0.0 <= lr:
             raise ValueError(f"Invalid learning rate: {lr}")
         if not 0.0 <= eps:

torch/distributed/optim/functional_adamax.py

@@ -4,6 +4,9 @@
 import torch
 import torch.optim._functional as F
 from torch import Tensor
+from torch.distributed.optim._deprecation_warning import (
+    _scripted_functional_optimizer_deprecation_warning,
+)
 
 
 __all__: List[str] = []
@@ -31,6 +34,7 @@ def __init__(
         maximize: bool = False,
         _allow_empty_param_list: bool = False,
     ):
+        _scripted_functional_optimizer_deprecation_warning(stacklevel=2)
         if not 0.0 <= lr:
             raise ValueError(f"Invalid learning rate: {lr}")
         if not 0.0 <= eps:

torch/distributed/optim/functional_adamw.py

@@ -4,6 +4,9 @@
 import torch
 import torch.optim._functional as F
 from torch import Tensor
+from torch.distributed.optim._deprecation_warning import (
+    _scripted_functional_optimizer_deprecation_warning,
+)
 
 
 __all__: List[str] = []
@@ -33,6 +36,7 @@ def __init__(
         fused: bool = False,
         _allow_empty_param_list: bool = False,
     ):
+        _scripted_functional_optimizer_deprecation_warning(stacklevel=2)
         if not 0.0 <= lr:
             raise ValueError(f"Invalid learning rate: {lr}")
         if not 0.0 <= eps:

torch/distributed/optim/functional_rmsprop.py

@@ -4,6 +4,9 @@
 import torch
 import torch.optim._functional as F
 from torch import Tensor
+from torch.distributed.optim._deprecation_warning import (
+    _scripted_functional_optimizer_deprecation_warning,
+)
 
 
 __all__: List[str] = []
@@ -33,6 +36,7 @@ def __init__(
         maximize: bool = False,
         _allow_empty_param_list: bool = False,
     ):
+        _scripted_functional_optimizer_deprecation_warning(stacklevel=2)
         self.defaults = {
             "lr": lr,
             "alpha": alpha,

torch/distributed/optim/functional_rprop.py

@@ -4,6 +4,9 @@
 import torch
 import torch.optim._functional as F
 from torch import Tensor
+from torch.distributed.optim._deprecation_warning import (
+    _scripted_functional_optimizer_deprecation_warning,
+)
 
 
 __all__: List[str] = []
@@ -30,6 +33,7 @@ def __init__(
         maximize: bool = False,
         _allow_empty_param_list: bool = False,
     ):
+        _scripted_functional_optimizer_deprecation_warning(stacklevel=2)
         self.defaults = {
             "lr": lr,
         }

torch/distributed/optim/functional_sgd.py

@@ -4,6 +4,9 @@
 import torch
 import torch.optim._functional as F
 from torch import Tensor
+from torch.distributed.optim._deprecation_warning import (
+    _scripted_functional_optimizer_deprecation_warning,
+)
 
 
 __all__: List[str] = []
@@ -33,6 +36,7 @@ def __init__(
         fused: bool = False,
         _allow_empty_param_list: bool = False,
     ):
+        _scripted_functional_optimizer_deprecation_warning(stacklevel=2)
         self.defaults = {
             "lr": lr,
             "momentum": momentum,