Update test_cuda.py and test_torch.py optim tests to use OptimizerInfo and optim_db

[janeyx99]

🚀 The feature, motivation and pitch

Update: all tasks below have in progress PRs currently!

In test_cuda.py and test_torch.py, we have the following optimizer tests that should be migrated to use the new OptimizerInfo infrastructure. (See any of the tests in TestOptimRenewed for how to apply OptimizerInfos.) The basic structure to expect for an updated test will look like:

    @optims([optim for optim in optim_db if ...], dtypes=[torch.float32])
    def test_name(self, device, dtype, optim_info):
        optim_cls = optim_info.optim_cls
        optim_inputs = optim_info.optim_inputs_func(device=device)
        for optim_input in optim_inputs:
            params = ...
            // processing, like updating grads

            optimizer = optim_cls(params, ..., **optim_input.kwargs)

            // the actual test
            ...
            self.assertEquals(...)

test_cuda.py:

• [ ] test_grad_scaling_autocast_fused_optimizers will be addressed in https://github.com/pytorch/pytorch/pull/124904 and by @SandishKumarHN
• [x] test_graph_grad_scaling
• [ ] test_graph_optims
• [ ] test_graph_optims_with_explicitly_capturable_param_groups
• [ ] test_graph_scaling_fused_optimizers

test_torch.py: #125538

• [x] tests that call _grad_scaling_autocast_test can be combined into one test
• [x] test_params_invalidated_with_grads_invalidated_between_unscale_and_step should be tested on more configurations (like SGD fused!)
• there are many other tests that test just one optimizer config that we can leave for now.

Feel free to take up one or a subset of those tests to migrate--edit your username next to the checklist to claim an item to submit PRs for!

Alternatives

No response

Additional context

No response

cc @vincentqb @jbschlosser @albanD @crcrpar

[rithikp06]

I'd like to work on this

[janeyx99]

@rithikp06 feel free to claim a sub task (e.g., which tests are you migrating?) and open a PR!

[rithikp06]

I'll work on the test_torch.py tests.

[ComposeC]

@janeyx99 I'd like to work on the test_grad_scaling_autocast_fused_optimizers function in test_cuda.py.

[rithikp06]

I created a PR that addresses the test_torch.py part: https://github.com/pytorch/pytorch/pull/123537

[FireACNC]

I'd like to work on the test_graph_grad_scaling function in test_cuda.py.

[jayanthd04]

Hi, I'm a first time contributor. I would like to work on test_graph_optims and test_graph_scaling_fused_optimizers. Could i get some pointer on where I could find more information about OptimizerInfo infrastructure and TestOptimRenewed?

[dwang3851]

I'd like to work on the test_graph_optims_with_explicitly_capturable_param_groups function in test_cuda.py

[khlaifiabilel]

Greetings 👋 , This is very interesting I would like to work on the test_graph_optims in the test_cuda.py

[janeyx99]

Hi, I'm a first time contributor. I would like to work on test_graph_optims and test_graph_scaling_fused_optimizers. Could i get some pointer on where I could find more information about OptimizerInfo infrastructure and TestOptimRenewed?

TestOptimRenewed is a class full of tests that use the OptimizerInfo infrastructure: https://github.com/pytorch/pytorch/blob/main/test/test_optim.py#L44

The OptimizerInfo infrastructure is defined in common_optimizers.py, you would want to search for optims_db and OptimizerInfo for what they are. I think looking at an existing test (any of the ones in TestOptimRenewed) would give a good highlevel of how these are used.

[julian-8897]

HI there, I'm a first time contributor here :). I would like to work on the test for _grad_scaling_autocast_test under test_torch.py

[PetaBread545]

Hello, I am also a first time contributor. Seeing that #123537 already addresses the test_py subtasks, I would like to work on the test_cuda subtasks.

[dwang3851]

Hi, I'm a first time contributor working on this issue. Is there any way to approve Github workflows for me on the pr above, I was hoping to make sure those results matched my local results?

[omastey]

Hi, I have begun working on this ticket, are there any specs or documentation I could look at?

[Rbhu376264]

Hello There, Hope this message finds you well!!! I am a newbie to open-source contribution, and would really like to get some advice how to get started in contributing "test_graph_scaling_fused_optimizers" . How can I get started on this? Any suggestions would be highly appreciated!!!

[janeyx99]

not done yet