Updated test_graph_optims and test_graph_scaling_fused_optimizers to use new OptimizerInfo infrastructure

[jayanthd04]

This PR is meant to address issue #123451, more specifically, the test_graph_optims and test_graph_scaling_fused_optimizers functions in test_cuda.py have been updated so that they now use the new OptimizerInfo infrastructure.

Lintrunner passed:

$ lintrunner test/test_cuda.py
ok No lint issues.

Tests passed:

>python test_cuda.py -k test_graph_optims
Ran 19 tests in 7.463s

OK (skipped=9)

>python test_cuda.py -k test_graph_scaling_fused_optimizers
Ran 6 tests in 2.800s

OK (skipped=3)

Both the functions have been moved to the newly created TestCase class TestCudaOptims. The test is mostly the same except the @optims decorator is used at the top of the function to implicitly call the function using each of the optimizers mentioned in the decorator instead of explicitly using a for loop to iterate through each of the optimizers.

I was unable to use the _get_optim_inputs_including_global_cliquey_kwargs to get all kwargs for each of the optimizers since some of the kwargs that are used in the original test_graph_optims function are not being returned by the new OptimizerInfo infrastructure, more specifically, for the torch.optim.rmsprop.RMSprop optimizer, the following kwargs are not returned whenever _get_optim_inputs_including_global_cliquey_kwargs is called:

{'foreach': False, 'maximize': True, 'weight_decay': 0}
{ 'foreach': True, 'maximize': True, 'weight_decay': 0}

I ran into the same issue for test_graph_scaling_fused_optimizers, for the torch.optim.adamw.AdamW optimizer, whenever optim_info.optim_inputs_func(device=device) was called, the following kwarg was not returned:

{'amsgrad': True}

Due to this issue, I resorted to using a dictionary to store the kwargs for each of the optimizers, I am aware that this is less than ideal. I was wondering whether I should use the OptimizerInfo infrastructure to get all the kwargs regardless of the fact that it lacks some kwargs.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/125127

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

❌ 10 New Failures

As of commit 236f2fe with merge base 1a28f73 ():

NEW FAILURES - The following jobs have failed:

• pull / linux-focal-cuda12.1-py3.10-gcc9 / test (default, 2, 5, linux.4xlarge.nvidia.gpu) (gh)
  inductor/test_compiled_optimizers.py::CompiledOptimizerTests::test_sgd_weight_decay_cpu
• pull / linux-focal-cuda12.1-py3.10-gcc9-sm86 / test (default, 4, 5, linux.g5.4xlarge.nvidia.gpu) (gh)
  inductor/test_compiled_optimizers.py::CompiledOptimizerTests::test_sgd_weight_decay_cuda
• pull / linux-focal-py3.11-clang10 / test (default, 3, 3, linux.2xlarge) (gh)
  inductor/test_compiled_optimizers.py::CompiledOptimizerTests::test_sgd_weight_decay_cpu
• pull / linux-focal-py3.12-clang10 / test (default, 3, 3, linux.2xlarge) (gh)
  inductor/test_compiled_optimizers.py::CompiledOptimizerTests::test_sgd_weight_decay_cpu
• pull / linux-focal-py3.8-clang10 / test (default, 3, 3, linux.2xlarge) (gh)
  inductor/test_compiled_optimizers.py::CompiledOptimizerTests::test_sgd_weight_decay_cpu
• pull / linux-jammy-py3.10-clang15-asan / test (default, 5, 6, linux.4xlarge) (gh)
  inductor/test_compiled_optimizers.py::CompiledOptimizerTests::test_sgd_weight_decay_cpu
• pull / linux-jammy-py3.8-gcc11 / test (default, 3, 3, linux.2xlarge) (gh)
  inductor/test_compiled_optimizers.py::CompiledOptimizerTests::test_sgd_weight_decay_cpu
• trunk / linux-focal-cuda12.1-py3.10-gcc9 / test (nogpu_AVX512, 1, 1, linux.2xlarge) (gh)
  inductor/test_compiled_optimizers.py::CompiledOptimizerTests::test_sgd_weight_decay_cpu
• trunk / linux-focal-cuda12.1-py3.10-gcc9 / test (nogpu_NO_AVX2, 1, 1, linux.2xlarge) (gh)
  inductor/test_compiled_optimizers.py::CompiledOptimizerTests::test_sgd_weight_decay_cpu
• trunk / macos-13-py3-arm64 / test (default, 2, 3, macos-m1-stable) (gh)
  inductor/test_compiled_optimizers.py::CompiledOptimizerTests::test_sgd_weight_decay_cpu

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[linux-foundation-easycla]

The committers listed above are authorized under a signed CLA.

• ✅ login: jayanthd04 (2aca95f, 287f741, 37e6979, e901a0a, ad331e2, e80cde4, 9e86434, 501590b, c7b33f4, 8d6abfb, eac2bd0, 79f542b, 6d72aff, 8e0945a, b8f1ad0, 91264d4, d4815a7, c93214e, be64baf, 3a86b06, f6cec68, 2a6c915, c10432e, a0af2fe, 35e1b6a, af8653c, 1e230e7, 39d99b5, fe29e60, b75098e, 236f2fe, cd6249f, 5f92373, 543bba1, b13d6b0, 8d968d7)

[janeyx99]

Thanks for your thorough look at the kwargs. For RMSprop, feel free to add a maximize with no weight_decay option here:

pytorch/torch/testing/_internal/common_optimizers.py

Line 113 in b18d8d4

supports_param_groups: bool = True,

For Adam/W, it's okay to not have amsgrad alone--we test it sufficiently with the "capturable, amsgrad" and the "amsgrad" described inputs.

This way we can use the helper function to get the kwargs instead of needing a separate dictionary.

[jayanthd04]

Thank you for the review! To add a maximize with no weight_decay option in RMSprop, should I just edit optim_inputs_func_rmsprop? Also I've noticed some other optimizers such as torch.optim.Adadelta are also missing some kwargs such as no weight_decay and maximize, should I also add those options in common_optimizers.py or should I just use the current options for those optimizers?

[janeyx99]

Thank you for the review! To add a maximize with no weight_decay option in RMSprop, should I just edit optim_inputs_func_rmsprop?

Yes

Also I've noticed some other optimizers such as torch.optim.Adadelta are also missing some kwargs such as no weight_decay and maximize, should I also add those options in common_optimizers.py

Yes

[janeyx99]

Two more changes! And we should be good

[janeyx99]

approving, but the lr change should be removed + ci should be green before landing

[janeyx99]

@pytorchbot merge

[pytorchmergebot]

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

[pytorchmergebot]

Merge failed

Reason: 1 mandatory check(s) failed. The first few are:

• pull / linux-focal-py3.12-clang10 / test (default, 3, 3, linux.2xlarge)

Dig deeper by viewing the failures on hud

Details for Dev Infra team

Raised by workflow job

Failing merge rule: Core Maintainers

[janeyx99]

@pytorchbot rebase

[pytorchmergebot]

@pytorchbot started a rebase job onto refs/remotes/origin/viable/strict. Check the current status here

[pytorchmergebot]

Rebase failed due to Command git -C /home/runner/work/pytorch/pytorch rebase refs/remotes/origin/viable/strict pull/125127/head returned non-zero exit code 1

Rebasing (1/23)
Auto-merging test/test_cuda.py
CONFLICT (content): Merge conflict in test/test_cuda.py
error: could not apply fe29e607709... Adding TestCudaOptims class to move test_graph_optims function and test_graph_scaling_fused_optimizers to
hint: Resolve all conflicts manually, mark them as resolved with
hint: "git add/rm <conflicted_files>", then run "git rebase --continue".
hint: You can instead skip this commit: run "git rebase --skip".
hint: To abort and get back to the state before "git rebase", run "git rebase --abort".
Could not apply fe29e607709... Adding TestCudaOptims class to move test_graph_optims function and test_graph_scaling_fused_optimizers to

Raised by https://github.com/pytorch/pytorch/actions/runs/9083238583

[janeyx99]

Could you rebase locally and see if any of the CI fails are due to this PR? They shouldn't, except the compiled optimizer one, which may require adding a customization in test_compiled_optimizers.py

[jayanthd04]

I believe this section of my branch in test_cuda.py may have been the issue. After deleting the decorator I was able to merge with the upstream and my branch

@torch.testing._internal.common_utils.markDynamoStrictTest 
class TestCudaOptims(TestCase):

[janeyx99]

That should not be the case; the decorator is strictly better and more deterministic. Could you add it back?

[jayanthd04]

I've just added it back.

[janeyx99]

Yes, if the compiled optimizer test fails again, you may need to add a new entry to

pytorch/test/inductor/test_compiled_optimizers.py

Line 138 in 352a893

"test_sgd_weight_decay_maximize_cpu": 4,

[jayanthd04]

I believe 'test/inductor/test_compiled_optimizers.py::CompiledOptimizerTests::test_sgd_weight_decay_cpu' is the test that is failing.

[janeyx99]

Yep, looks like you need to add an entry to the code snippet I linked above as you added new configs.

[jayanthd04]

Should I just add this to pytorch/test/inductor/test_compiled_optimizers.py

test_sgd_weight_decay_cpu=4,
test_sgd_weight_decay_cuda=4,

[janeyx99]

That looks reasonable--does that pass the test?