🔗 Helpful Links

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

• 📄 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 our office hours

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

❌ 1 Failures

As of commit 5a05a87:

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

after this PR lands, I suspect we can remove the capturable keyword and what it represents.

Also this will trigger #74424 again. So not sure if we can do this.

Also this will trigger #74424 again. So not sure if we can do this.

I dont think it will trigger that issue with XLA because we've removed all item calls. We can even get rid of most list comprehensions. So, I have high hopes that this wont be an issue for XLA

Does this make it capturable? These are still CPU Tensors.

Capturable (i believe) is whether dynamo can get a whole graph of it -- not that it should all be on GPU.

Capturable (i believe) is whether dynamo can get a whole graph of it

Capturable is used in the context of cudagraph: "whether this instance is safe to capture in a CUDA graph.".
So we do need the step to be on the GPU when this is true.

Capturable is used in the context of cudagraph: "whether this instance is safe to capture in a CUDA graph.".

Ho, interesting. In that case, I have to check capturable code, it's probably buggy as-is because when learning_rate drops happen, the scalar changes / drops but cudagraph would have baked it in

I think there are two things here:

1. This patch
   Scalar Tensor on CPU will be slower than plain python number. I guess we're happy with a small hit here but I think it would be good to measure it as some of these optimizer do quite a few ops (and our framework overhead is non-trivial).

2. The capturable flag:
   I think there is still an impasse here though:

• For this optimizer to be capturable by cudagraph, we MUST have the step on the GPU
• To get ok performance outside of cudagraph, we MUST keep the step on the CPU

We could consider checking during the step itself if we're using cudagraph but that would mean that we might have to move the step then (which is a sync point so not good).

Also I'm not sure how dynamo will handle that since it wants to use cudagraph, I guess it already breaks on any mix of CPU/GPU compute? And so this change will not really remove the graph breaks within dynamo?

@albanD One question: if the step is on GPU, this would be bad without cudagraphs because we're launching a kernel to do step + 1 essentially right? I think in the dynamo case this won't matter, because this op will get fused into a larger kernel anyway so could we make capturable=True by default, and then for people who aren't using dynamo or cudagraphs, they can manually set capturable=False?

this would be bad without cudagraphs because we're launching a kernel to do step + 1 essentially right?

Correct.

could we make capturable=True by default, and then for people who aren't using dynamo or cudagraphs, they can manually set capturable=False?

We can't make it the default no. For GPU-like backends, having these small ops on the device is prohibitively slow #74424

For GPU-like backends, having these small ops on the device is prohibitively slow #74424

I guess XLA is not GPU-like backends as it is lazy?

in each training iteration, now state['step'] needs to be move back to CPU from CUDA or XLA to compute things like step = step_t.item()

Seems the overhead mainly coming from moving those tensors back to CPU, which is exactly the issue this PR fixes.

cc @JackCaoG; I believe it would be beneficial to make capturable=True by default for all fusion-capable backends, e.g. XLA or CUDA graphs.

Generally speaking removing item call will be only beneficial to the xla backend. item call is one of the common cause of early execution. Without it, lazy will fuse all optimizer + forward + backward. I am not entirely sure what capturable does, @shunting314 might be able to comment for our dynamo->xla bridge.

@shunting314 might be able to comment for our dynamo->xla bridge.

in the dynamoc -> xla bridge for training, we are not running optimizer actually ( https://github.com/pytorch/pytorch/blob/master/benchmarks/dynamo/common.py#L894 ). But looks like people are adding them.

I am not entirely sure what capturable does

If capturable is True then the step is a device tensor instead of put on CPU, and the parameter decay is calculated on device. This will help reducing the CPU overhead when the backend could fuse them into a single optimizer computation (used mainly for CUDA graphs, but potentially benefit other backends as well).

See https://github.com/pytorch/pytorch/blob/master/torch/optim/adam.py#L366-L393

Do we have the sanity check benchmark as discussed above to make sure nothing is obviously broken in terms of speed due to these changes?

Update: This causes a pretty bad perf regression in both single and multi tensor mode for eager. As a workaround I will use is_compiling to call item() when in eager, and to use cpu tensors when using dynamo. I will verify this fixes the regressions and then we should be able to merge this.

cc @albanD

Update: The current version matches the eager perf of the old version per Alban's request. Should be good to go, just need a sign off. Updated nums

Tests issues are related to the dynamo import in optim, should be fixed now I think

By the way @mlazos I guess that #90624 should actually allow you to remove the .item() changes?
What about the other changes? I think that it is part of the dynamo goal to be able to do these integrations without modifying user code with if is_compiling a-la torchscipt?

I don't think symint is ready for prime time since there are still some lingering bugs. Once dynamic shapes is on by default I'm fine with removing this.

Yes we usually don't want to modify code to compile with dynamo. In most other cases we don't care about eager perf as much though. Like in this case if we didn't care about eager perf I'd just remove the item calls and call it a day since dynamo provides good speed up. It ends up being a trade off off of is this worth modifying the compiler and adding x feature vs should I just remove x from the code if it isn't really necessary. With user code I lean towards the former while in code we own it's more balanced.

@pytorchbot merge -f "Unrelated test failure"

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

In most other cases we don't care about eager perf as much though.

I would challenge that actually. People that write library code (nn, vision, audio, etc) need to support both. So this is a major use case!