Updated test_graph_optims and test_graph_scaling_fused_optimizers to use new OptimizerInfo infrastructure

test/test_cuda.py

@@ -4459,23 +4459,13 @@ def test_graph_optims(self, device, dtype, optim_info):
  all_optim_inputs = _get_optim_inputs_including_global_cliquey_kwargs(
  device, dtype, optim_info, skip=("differentiable",)
  )
- has_betas = any(
- "betas" in error_inp.optimizer_error_input.kwargs
- for error_inp in optim_info.optim_error_inputs_func(
- device="cpu", dtype=dtype
- )
- )
 
  steps_warmup = 3
  steps_train = 2
 
  for optim_input in all_optim_inputs:
  kwargs = optim_input.kwargs
- if "lr" in kwargs:
- del kwargs["lr"]
- kwargs["lr"] = 0.1
- if has_betas and optim_cls != torch.optim.Adamax:
- kwargs["betas"] = (0.8, 0.7)
+ kwargs["lr"]=0.1
 
  for actually_do_graphs in (True, False):
  params = [
@@ -4543,26 +4533,15 @@ def test_graph_scaling_fused_optimizers(self, device, dtype, optim_info):
  steps_train = 2
 
  optim_inputs = optim_info.optim_inputs_func(device=device)
- has_betas = any(
- "betas" in error_inp.optimizer_error_input.kwargs
- for error_inp in optim_info.optim_error_inputs_func(
- device="cpu", dtype=dtype
- )
- )
 
  for optim_input in optim_inputs:
  kwargs = optim_input.kwargs
  kwargs["fused"] = True
- if "lr" in kwargs:
- del kwargs["lr"]
- kwargs["lr"] = 0.1
- if has_betas:
- kwargs["betas"] = (0.8, 0.7)
 
  for actually_do_graphs in (
  (True, False) if optim_info.has_capturable_arg else (True,)
  ):
- params = [torch.randn((i + 5, i + 5), device="cuda") for i in range(2)]
+ params = [torch.randn((i + 5, i + 5), device=device) for i in range(2)]
  params_control = [p.clone().requires_grad_() for p in params]
  params_graphed = [p.clone().requires_grad_() for p in params]
 

torch/testing/_internal/common_optimizers.py

@@ -123,7 +123,7 @@ def __init__(
  supported_impls: Tuple[str] = ("foreach", "differentiable"),
  # the optim supports passing in sparse gradients as well as dense grads
  supports_sparse: bool = False,
- # the optim is capturable in a CUDA graph
+ # the optimizer constructor supports passing in capturable as a kwarg
  has_capturable_arg: bool = False,
  # the optim only supports one config: sparse grads w/ dense params, see SparseAdam
  only_supports_sparse_grads: bool = False,
@@ -314,6 +314,7 @@ def optim_inputs_func_adadelta(device, dtype=None):
  OptimizerInput(
  params=None, kwargs={"weight_decay": 0.1}, desc="nonzero weight_decay"
  ),
+ OptimizerInput(params=None, kwargs={"maximize": True}, desc="maximize"),
  OptimizerInput(
  params=None,
  kwargs={"weight_decay": 0.1, "maximize": True},
@@ -322,7 +323,7 @@ def optim_inputs_func_adadelta(device, dtype=None):
  OptimizerInput(
  params=None, kwargs={"rho": 0.95, "weight_decay": 0.9}, desc="rho"
  ),
- OptimizerInput(params=None, kwargs={"maximize": True}, desc="maximize"),
+
  ] + (cuda_supported_configs if "cuda" in str(device) else [])
 
 
@@ -532,14 +533,15 @@ def optim_inputs_func_adamax(device, dtype=None):
  ),
  OptimizerInput(
  params=None,
- kwargs={"weight_decay": 0.1, "maximize": True},
- desc="maximize, weight_decay",
+ kwargs={"maximize": True},
+ desc="maximize",
  ),
  OptimizerInput(
  params=None,
- kwargs={"maximize": True},
- desc="maximize",
+ kwargs={"weight_decay": 0.1, "maximize": True},
+ desc="maximize, weight_decay",
  ),
+
  ] + (cuda_supported_configs if "cuda" in str(device) else [])
 
 
@@ -689,6 +691,13 @@ def optim_inputs_func_nadam(device, dtype=None):
  kwargs={"momentum_decay": 6e-3},
  desc="non-zero momentum_decay",
  ),
+ OptimizerInput(
+ params=None,
+ kwargs={
+ "weight_decay": 0.1,
+ },
+ desc="weight_decay",
+ ),
  OptimizerInput(
  params=None,
  kwargs={"weight_decay": 0.1, "momentum_decay": 6e-3},
@@ -702,13 +711,6 @@ def optim_inputs_func_nadam(device, dtype=None):
  },
  desc="decoupled_weight_decay",
  ),
- OptimizerInput(
- params=None,
- kwargs={
- "weight_decay": 0.1,
- },
- desc="weight_decay",
- ),
  ] + (cuda_supported_configs if "cuda" in str(device) else [])
 
 
@@ -834,38 +836,38 @@ def optim_inputs_func_rmsprop(device, dtype=None):
  ),
  OptimizerInput(
  params=None,
- kwargs={"weight_decay": 0.1, "centered": True},
- desc="centered",
+ kwargs={
+ "maximize": True,
+ },
+ desc="maximize",
  ),
  OptimizerInput(
  params=None,
- kwargs={"weight_decay": 0.1, "centered": True, "momentum": 0.1},
- desc="momentum",
+ kwargs={"weight_decay": 0.1, "centered": True},
+ desc="centered",
  ),
  OptimizerInput(
  params=None,
  kwargs={
- "weight_decay": 0.1,
- "centered": True,
- "momentum": 0.1,
  "maximize": True,
+ "weight_decay": 0.1,
  },
- desc="maximize, centered, weight_decay, w/ momentum",
+ desc="maximize, weight_decay",
  ),
  OptimizerInput(
  params=None,
- kwargs={
- "maximize": True,
- },
- desc="maximize",
+ kwargs={"weight_decay": 0.1, "centered": True, "momentum": 0.1},
+ desc="momentum",
  ),
  OptimizerInput(
  params=None,
  kwargs={
- "maximize": True,
  "weight_decay": 0.1,
+ "centered": True,
+ "momentum": 0.1,
+ "maximize": True,
  },
- desc="maximize, weight_decay",
+ desc="maximize, centered, weight_decay, w/ momentum",
  ),
  ] + (cuda_supported_configs if "cuda" in str(device) else [])
 
@@ -936,7 +938,15 @@ def optim_inputs_func_sgd(device, dtype=None):
  OptimizerInput(
  params=None, kwargs={"lr": torch.tensor(0.001)}, desc="tensor lr"
  ),
+ OptimizerInput(
+ params=None, kwargs={"weight_decay": 0.5}, desc="non-zero weight_decay"
+ ),
  OptimizerInput(params=None, kwargs={"momentum": 0.9}, desc="momentum"),
+ OptimizerInput(
+ params=None,
+ kwargs={"weight_decay": 0.1, "maximize": True},
+ desc="maximize",
+ ),
  OptimizerInput(
  params=None,
  kwargs={"momentum": 0.9, "dampening": 0.5},
@@ -952,14 +962,6 @@ def optim_inputs_func_sgd(device, dtype=None):
  kwargs={"momentum": 0.9, "nesterov": True, "weight_decay": 0.1},
  desc="nesterov",
  ),
- OptimizerInput(
- params=None,
- kwargs={"weight_decay": 0.1, "maximize": True},
- desc="maximize",
- ),
- OptimizerInput(
- params=None, kwargs={"weight_decay": 0.5}, desc="non-zero weight_decay"
- ),
  ]