[optim][adam] use fastest impl whenever possible, add util #93184

janeyx99 · 2023-01-27T22:49:53Z

This allows it so that ONLY when the users don't set anything for foreach or fused do we switch the default and cascades adam so that we default to fused, then foreach, then single-tensor.

To clarify:

if the user puts True in foreach only, it will run the foreach implementation.
if the user puts True in fused only, it will run the fused implementation.
if the user puts True in foreach AND for fused, it will run the fused implementation.

And:

if the user puts False in foreach only, it will run the single tensor implementation.
if the user puts False in fused only, it will still run the single tensor implementation.
if the user puts False in foreach AND for fused, it will run the single tensor implementation.

Stack from ghstack (oldest at bottom):

-> [optim][adam] use fastest impl whenever possible, add util #93184

I also didn't trust myself that much with the helper function, so I ran some local asserts on _default_to_fused_or_foreach. The only point left to really test is the type(p) -- torch.Tensor but I think the distributed tests will catch that in CI.

cuda_only_fp_list = [
    torch.rand((1, 2), device="cuda", dtype=torch.float32),
    torch.rand((1, 2), device="cuda", dtype=torch.float64),
    torch.rand((1, 2), device="cuda", dtype=torch.float16),
    torch.rand((1, 2), device="cuda", dtype=torch.bfloat16),
]

cuda_only_int_list = [
    torch.randint(1024, (1, 2), device="cuda", dtype=torch.int64),
]

cpu_list = [
    torch.rand((1, 2), device="cpu", dtype=torch.float32),
    torch.rand((1, 2), device="cpu", dtype=torch.float64),
    torch.rand((1, 2), device="cpu", dtype=torch.float16),
]

none_list = [None]

# differentiable should always make it return false for both
assert _default_to_fused_or_foreach([cuda_only_fp_list], True, True) == (False, False)
assert _default_to_fused_or_foreach([cuda_only_fp_list], True, False) == (False, False)

# cpu lists should always make it return false for both
assert _default_to_fused_or_foreach([cuda_only_fp_list, cpu_list], False, True) == (False, False)
assert _default_to_fused_or_foreach([cpu_list], False, True) == (False, False)
assert _default_to_fused_or_foreach([cuda_only_fp_list, cpu_list], False, False) == (False, False)
assert _default_to_fused_or_foreach([cpu_list], False, False) == (False, False)

# has fused triggers correctly
assert _default_to_fused_or_foreach([cuda_only_fp_list], False, True) == (True, False)
assert _default_to_fused_or_foreach([cuda_only_fp_list], False, False) == (False, True)

# ints always goes to foreach
assert _default_to_fused_or_foreach([cuda_only_fp_list, cuda_only_int_list], False, True) == (False, True)
assert _default_to_fused_or_foreach([cuda_only_fp_list, cuda_only_int_list], False, False) == (False, True)

# Nones don't error
assert _default_to_fused_or_foreach([cuda_only_fp_list, none_list], False, True) == (True, False)
assert _default_to_fused_or_foreach([cuda_only_fp_list, cuda_only_int_list, none_list], False, True) == (False, True)
assert _default_to_fused_or_foreach([none_list], False, True) == (True, False)
assert _default_to_fused_or_foreach([none_list], False, False) == (False, True)

cc @ezyang @gchanan

[ghstack-poisoned]

pytorch-bot · 2023-01-27T22:49:56Z

🔗 Helpful Links

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

📄 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.

❌ 1 Failures, 1 Pending

As of commit a818e43:

NEW FAILURES - The following jobs have failed:

linux-bionic-py3_7-clang8-xla / test (xla, 1, 1, linux.4xlarge)

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

This allows it so that ONLY when the users don't set anything for foreach or fused do we switch the default and cascades adam so that we default to fused, then foreach, then single-tensor. To clarify: * if the user puts True in foreach _only_, it will run the foreach implementation. * if the user puts True in fused _only_, it will run the fused implementation. * if the user puts True in foreach AND for fused, it will run the fused implementation. And: * if the user puts False in foreach _only_, it will run the single tensor implementation. * if the user puts False in fused _only_, it will still run the single tensor implementation. * if the user puts False in foreach AND for fused, it will run the single tensor implementation. I also didn't trust myself that much with the helper function, so I ran some local asserts on _default_to_fused_or_foreach. The only point left to really test is the type(p) -- torch.Tensor but I think the distributed tests will catch that in CI. ``` cuda_only_fp_list = [ torch.rand((1, 2), device="cuda", dtype=torch.float32), torch.rand((1, 2), device="cuda", dtype=torch.float64), torch.rand((1, 2), device="cuda", dtype=torch.float16), torch.rand((1, 2), device="cuda", dtype=torch.bfloat16), ] cuda_only_int_list = [ torch.randint(1024, (1, 2), device="cuda", dtype=torch.int64), ] cpu_list = [ torch.rand((1, 2), device="cpu", dtype=torch.float32), torch.rand((1, 2), device="cpu", dtype=torch.float64), torch.rand((1, 2), device="cpu", dtype=torch.float16), ] none_list = [None] # differentiable should always make it return false for both assert _default_to_fused_or_foreach([cuda_only_fp_list], True, True) == (False, False) assert _default_to_fused_or_foreach([cuda_only_fp_list], True, False) == (False, False) # cpu lists should always make it return false for both assert _default_to_fused_or_foreach([cuda_only_fp_list, cpu_list], False, True) == (False, False) assert _default_to_fused_or_foreach([cpu_list], False, True) == (False, False) assert _default_to_fused_or_foreach([cuda_only_fp_list, cpu_list], False, False) == (False, False) assert _default_to_fused_or_foreach([cpu_list], False, False) == (False, False) # has fused triggers correctly assert _default_to_fused_or_foreach([cuda_only_fp_list], False, True) == (True, False) assert _default_to_fused_or_foreach([cuda_only_fp_list], False, False) == (False, True) # ints always goes to foreach assert _default_to_fused_or_foreach([cuda_only_fp_list, cuda_only_int_list], False, True) == (False, True) assert _default_to_fused_or_foreach([cuda_only_fp_list, cuda_only_int_list], False, False) == (False, True) # Nones don't error assert _default_to_fused_or_foreach([cuda_only_fp_list, none_list], False, True) == (True, False) assert _default_to_fused_or_foreach([cuda_only_fp_list, cuda_only_int_list, none_list], False, True) == (False, True) assert _default_to_fused_or_foreach([none_list], False, True) == (True, False) assert _default_to_fused_or_foreach([none_list], False, False) == (False, True) ``` cc ezyang gchanan [ghstack-poisoned]

ghstack-source-id: 61d23b5 Pull Request resolved: #93184

albanD

Sound good, only a stylistic comment.

torch/optim/adadelta.py

This allows it so that ONLY when the users don't set anything for foreach or fused do we switch the default and cascades adam so that we default to fused, then foreach, then single-tensor. To clarify: * if the user puts True in foreach _only_, it will run the foreach implementation. * if the user puts True in fused _only_, it will run the fused implementation. * if the user puts True in foreach AND for fused, it will run the fused implementation. And: * if the user puts False in foreach _only_, it will run the single tensor implementation. * if the user puts False in fused _only_, it will still run the single tensor implementation. * if the user puts False in foreach AND for fused, it will run the single tensor implementation. I also didn't trust myself that much with the helper function, so I ran some local asserts on _default_to_fused_or_foreach. The only point left to really test is the type(p) -- torch.Tensor but I think the distributed tests will catch that in CI. ``` cuda_only_fp_list = [ torch.rand((1, 2), device="cuda", dtype=torch.float32), torch.rand((1, 2), device="cuda", dtype=torch.float64), torch.rand((1, 2), device="cuda", dtype=torch.float16), torch.rand((1, 2), device="cuda", dtype=torch.bfloat16), ] cuda_only_int_list = [ torch.randint(1024, (1, 2), device="cuda", dtype=torch.int64), ] cpu_list = [ torch.rand((1, 2), device="cpu", dtype=torch.float32), torch.rand((1, 2), device="cpu", dtype=torch.float64), torch.rand((1, 2), device="cpu", dtype=torch.float16), ] none_list = [None] # differentiable should always make it return false for both assert _default_to_fused_or_foreach([cuda_only_fp_list], True, True) == (False, False) assert _default_to_fused_or_foreach([cuda_only_fp_list], True, False) == (False, False) # cpu lists should always make it return false for both assert _default_to_fused_or_foreach([cuda_only_fp_list, cpu_list], False, True) == (False, False) assert _default_to_fused_or_foreach([cpu_list], False, True) == (False, False) assert _default_to_fused_or_foreach([cuda_only_fp_list, cpu_list], False, False) == (False, False) assert _default_to_fused_or_foreach([cpu_list], False, False) == (False, False) # has fused triggers correctly assert _default_to_fused_or_foreach([cuda_only_fp_list], False, True) == (True, False) assert _default_to_fused_or_foreach([cuda_only_fp_list], False, False) == (False, True) # ints always goes to foreach assert _default_to_fused_or_foreach([cuda_only_fp_list, cuda_only_int_list], False, True) == (False, True) assert _default_to_fused_or_foreach([cuda_only_fp_list, cuda_only_int_list], False, False) == (False, True) # Nones don't error assert _default_to_fused_or_foreach([cuda_only_fp_list, none_list], False, True) == (True, False) assert _default_to_fused_or_foreach([cuda_only_fp_list, cuda_only_int_list, none_list], False, True) == (False, True) assert _default_to_fused_or_foreach([none_list], False, True) == (True, False) assert _default_to_fused_or_foreach([none_list], False, False) == (False, True) ``` cc ezyang gchanan [ghstack-poisoned]

ghstack-source-id: bf7fbc5 Pull Request resolved: #93184

janeyx99 · 2023-01-30T16:53:35Z

@pytorchbot merge

pytorchmergebot · 2023-01-30T16:55:42Z

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 · 2023-01-30T18:26:39Z

Merge failed

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

pull / linux-bionic-py3_7-clang8-xla / test (xla, 1, 1, linux.4xlarge)

Dig deeper by viewing the failures on hud

Details for Dev Infra team

Raised by workflow job

janeyx99 · 2023-01-30T19:57:05Z

@pytorchbot merge -f "xla failure looks like an infra issue"

pytorchmergebot · 2023-01-30T19:58:51Z

Merge started

Your change will be merged immediately since you used the force (-f) flag, bypassing any CI checks (ETA: 1-5 minutes).

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

[optim][adam] use fastest impl whenever possible, add util

8307ce4

[ghstack-poisoned]

janeyx99 requested a review from albanD as a code owner January 27, 2023 22:49

pytorch-bot bot added release notes: nn release notes category labels Jan 27, 2023

janeyx99 added the module: bc-breaking Related to a BC-breaking change label Jan 27, 2023

pytorch-bot bot added the topic: bc breaking topic category label Jan 27, 2023

janeyx99 mentioned this pull request Jan 27, 2023

[optim] Adam defaults to fused when CUDA + differentiable=False #90865

Closed

janeyx99 added a commit that referenced this pull request Jan 27, 2023

[optim][adam] use fastest impl whenever possible, add util

ca0057b

ghstack-source-id: 61d23b5 Pull Request resolved: #93184

crcrpar mentioned this pull request Jan 30, 2023

[mta][foreach] Implement fused adamw #88015

Closed

albanD approved these changes Jan 30, 2023

View reviewed changes

torch/optim/adadelta.py Outdated Show resolved Hide resolved

janeyx99 added a commit that referenced this pull request Jan 30, 2023

[optim][adam] use fastest impl whenever possible, add util

b7e201b

ghstack-source-id: bf7fbc5 Pull Request resolved: #93184

janeyx99 added the ciflow/trunk Trigger trunk jobs on your pull request label Jan 30, 2023

pytorchmergebot added the Merged label Jan 30, 2023

pytorchmergebot closed this in 4fc19e1 Jan 30, 2023

facebook-github-bot deleted the gh/janeyx99/27/head branch June 8, 2023 17:22

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[optim][adam] use fastest impl whenever possible, add util #93184

[optim][adam] use fastest impl whenever possible, add util #93184

Uh oh!

janeyx99 commented Jan 27, 2023 •

edited

Loading

Uh oh!

pytorch-bot bot commented Jan 27, 2023 •

edited

Loading

Uh oh!

albanD left a comment

Uh oh!

Uh oh!

janeyx99 commented Jan 30, 2023

Uh oh!

pytorchmergebot commented Jan 30, 2023

Uh oh!

pytorchmergebot commented Jan 30, 2023

Uh oh!

janeyx99 commented Jan 30, 2023

Uh oh!

pytorchmergebot commented Jan 30, 2023

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

[optim][adam] use fastest impl whenever possible, add util #93184

[optim][adam] use fastest impl whenever possible, add util #93184

Uh oh!

Conversation

janeyx99 commented Jan 27, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Jan 27, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

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

❌ 1 Failures, 1 Pending

Uh oh!

albanD left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

janeyx99 commented Jan 30, 2023

Uh oh!

pytorchmergebot commented Jan 30, 2023

Merge started

Uh oh!

pytorchmergebot commented Jan 30, 2023

Merge failed

Uh oh!

janeyx99 commented Jan 30, 2023

Uh oh!

pytorchmergebot commented Jan 30, 2023

Merge started

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

janeyx99 commented Jan 27, 2023 •

edited

Loading

pytorch-bot bot commented Jan 27, 2023 •

edited

Loading