Remove .item() calls from optimizer.step()

[yanboliang]

There are a lot of .item() calls inside of optimizer.step(), e.g:
https://github.com/pytorch/pytorch/blob/master/torch/optim/asgd.py#L182-L212

These .item() calls inside of the for loop cause graph break when enabling dynamo and finally hitting cache limit. I took a closer look at these usages and found actually we can remove these call because the results were feed into other downstream tensor operations.
I did a simple verification on ASGD and found removing .item() works well and no graph break on Dynamo end.

I'd like to check if there are other constraints to prevent us from doing this. If not, I can go ahead to fix this.

cc @jansel

[mlazos]

Yeah I looked into this earlier, IMO this should be okay but someone more familiar with PyTorch should verify. Adam even provides the capturable kwarg to its constructor to facilitate this. The main reason for my comment: isn't this handled by FakeItemVariable or unspec variables?

[voznesenskym]

Yeah I looked into this earlier, IMO this should be okay but someone more familiar with PyTorch should verify. Adam even provides the capturable kwarg to its constructor to facilitate this. The main reason for my comment: isn't this handled by FakeItemVariable or unspec variables?

+1 - we have a flag capture_scalar_outputs that should help with .item()

[yanboliang]

The main reason for my comment: isn't this handled by FakeItemVariable or unspec variables?

It's handled by FakeItemVariable if users turn on capture_scalar_outputs explicitly. Actually it's one of the way to remove .item() calls. The reason I raised this discuss is to figure out the best way to handle this. There are mainly two directions:

• Directly removing .item() inside of PyTorch optimizer.step, or providing the capturable kwarg like what Adam does.
• Handling it at TorchDynamo side with the capture_scalar_outputs.

I prefer option 1 as optimizer.step is totally PyTorch internal logics, which is a more fundamental solution. I think capture_scalar_outputs would be helpful to optimize user defined functions where users are more clear about what happens inside. If we go with option 1, we need folks who are more familiar with PyTorch to verify. cc @ezyang

[voznesenskym]

The main reason for my comment: isn't this handled by FakeItemVariable or unspec variables?

It's handled by FakeItemVariable if users turn on capture_scalar_outputs explicitly. Actually it's one of the way to remove .item() calls. The reason I raised this discuss is to figure out the best way to handle this. There are mainly two directions:

• Directly removing .item() inside of PyTorch optimizer.step, or providing the capturable kwarg like what Adam does.
• Handling it at TorchDynamo side with the capture_scalar_outputs.

I prefer option 1 as optimizer.step is totally PyTorch internal logics, which is a more fundamental solution. I think capture_scalar_outputs would be helpful to optimize user defined functions where users are more clear about what happens inside. If we go with option 1, we need folks who are more familiar with PyTorch to verify. cc @ezyang

Taking a step back, what problem are we solving? If we already have 2 working, is it even a problem? Or just an optimization?

[ezyang]

Do we have any real users that are using ASGD?

[ezyang]

The item calls are because the foreach optimizers don't support tensor inputs. They were added in pytorch/pytorch#7023 cc @mikaylagawarecki

These variables are "constant-like". The other way to solve the problem is to avoid specializing on their values and then have the internal passes constant prop to the item calls.

[yanboliang]

Do we have any real users that are using ASGD?

I just take it as an example, many other optimizers have similar issues.

[voznesenskym]

The item calls are because the foreach optimizers don't support tensor inputs. They were added in pytorch/pytorch#7023 cc @mikaylagawarecki

These variables are "constant-like". The other way to solve the problem is to avoid specializing on their values and then have the internal passes constant prop to the item calls.

That is what capture_scalar_outputs does - it turns the .item() call into a call_method in the graph, instead of specializing on the value as a constant. Is this what you mean?

[wconstab]

Do we have any real users that are using ASGD?

@ezyang we see this on Adam too. (adam.py:309).

I'm jumping in here bc I also ran into this issue from 2 places this week:

• TorchDynamo sometime takes very long to symbolically interpret stackframes. #525 - uses default adam optimizer
• distributed benchmarks on torchbenchmark using the torchbench distributed harness rather than the torchdynamo harness - use default adam optimizers

By default both torchvision and huggingface models in torchbench both use Adam without specifying capturable, so they get the default value of =False.

I'm a little concerned if we actually think users need to set capturable on their optimizers in order to work with dynamo, but i'd also prefer if we could update the optimizer code to avoid .item rather than doing magic inside dynamo. Can someone explain the context of capturable flag in these optimizers and offer an argument for how it should behave with dynamo?

[mlazos]

I'm a little concerned if we actually think users need to set capturable on their optimizers in order to work with dynamo, but i'd also prefer if we could update the optimizer code to avoid .item rather than doing magic inside dynamo. Can someone explain the context of capturable flag in these optimizers and offer an argument for how it should behave with dynamo?

capturable is there to enable capturing of the graph for CUDA (according to the code comment). I think in practice this is just to enable more performant GPU code without item() calls.

I agree that having users have to set this and understand this is bad ux for such a common component like optimizers.

For dynamo, I think each optimizer should have an equivalent capturable flag which avoids item calls. I say this because I don't think users should need to set this non-default config param capture_scalar_outputs=True to handle something as common as optimizers. We should default to good perf (fewer graph breaks). AFAIK the only reason SGD and Adam have this are because they are the most popular optimizers. An easy solution is to patch the defaults to make the default value for capturable to be True when we're running dynamo, or we can just change the default capturable value to True in Pytorch if we're okay with that.

[anijain2305]

ping @mlazos @wconstab on the action here.

I think its ok to do patch the defaults to True during Dynamo bytecode analssis.

[wconstab]

as a long term solution, i'm not too happy with 'patch the defaults in dynamo' approach, since we're no longer running the codepath that users expect. (if users want to read the optimizer code and they assume they are using defaults, they'll trace down the non-capturable path and have no clue we're secretly flipping a bit on them.

we need to do something quick though since even for development this is kind of a blocker. (e.g. for my torchbench/DDP work i have to hack this flag to True to get anything done)

if we check in a comment under the 2 optimizers that have capturable flags that explains dynamo will set 'true' there, that makes me a little happier.

but what is our plan for supporting the rest of the optimizers?

• a foolproof way to issue a single readable warning if someone uses an optimizer that calls .item() a bunch of times (whether from one of pytorch's optimizers or their own custom one)
• a scalable solution for updating most (all) of them to be dynamo friendly

is there anything else we can do besides rewriting the optimizers?

[mlazos]

@eellison suggested another compelling option, have .item() calls be tracked by symints and symfloats vs implicitly converting them to 1 element tensors in dynamo.

That will allow us to handle .item() directly instead of using the capturable flag which will avoid the calls.

I think in the short term we can either hack capturable/make optimizers capturable as we use them, or just set capture_scalar_outputs, both of which we admit don't have the best ux. However, we might be able to turn on capture_scalar_outputs by default, which would be less tricky for our users.

Long term symint/symfloat support I think would be the most robust option.