Skip to content

[adam] Add not torch.jit.is_scripting() as a requirement for switching to fused #92181

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
janeyx99 wants to merge 1 commit into from
Closed

[adam] Add not torch.jit.is_scripting() as a requirement for switching to fused #92181

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 10 additions & 7 deletions torch/optim/adam.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -316,13 +316,16 @@ def adam(params: List[Tensor],
# and when differentiable=False.
# We still respect when the user inputs False for fused.
if fused is None:
if not differentiable and all(
p.is_cuda and torch.is_floating_point(p)
for p in params + grads + exp_avgs + exp_avg_sqs + max_exp_avg_sqs + state_steps
):
fused = True
else:
fused = False
all_tensors = []
all_tensors.extend(params)
all_tensors.extend(grads)
all_tensors.extend(exp_avgs)
all_tensors.extend(exp_avg_sqs)
all_tensors.extend(max_exp_avg_sqs)
all_tensors.extend(state_steps)
Comment on lines +319 to +325
Copy link
Collaborator

@Skylion007 Skylion007 Jan 13, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Couldn't you refactor into a nice single itertools.chain, that way you don't need any temporary list at all.

Copy link
Contributor Author

@janeyx99 janeyx99 Jan 14, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I started with that, but it does not play well with JIT bindings: #91896 (comment)

Copy link
Collaborator

@Skylion007 Skylion007 Jan 14, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I see, which is annoying because jit bindings shouldn't be necessary since this path isn't part possible under JIT.

Copy link
Contributor

@kit1980 kit1980 Jan 19, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The seven lines can be changed into just 1 line, no?
all_tensors = [*params, *grads, *exp_avgs, *exp_avg_sqs, *max_exp_avg_sqs, *state_steps]

fused = not torch.jit.is_scripting() and not differentiable and all(
p.is_cuda and torch.is_floating_point(p) for p in all_tensors
)

if not all(isinstance(t, torch.Tensor) for t in state_steps):
raise RuntimeError("API has changed, `state_steps` argument must contain a list of singleton tensors")
Expand Down