RuntimeError when using Adam(fused=True) with torch.compile

[mitchellgoffpc]

RuntimeError when using Adam(fused=True) with torch.compile

[janeyx99]

Could be related to #110758

[xmfan]

works on aot_eager backend, looks like a stride mismatch?

{(1,), (3, 1, 1, 1), (128, 1, 1, 1)}  # aot_eager
{(1,), (3, 1, 3, 3), (128, 1, 128, 128), (128, 1, 1, 1)}  # inductor

diff --git a/torch/optim/adam.py b/torch/optim/adam.py
index fba4b2027b0..f7f13029cef 100644
--- a/torch/optim/adam.py
+++ b/torch/optim/adam.py
@@ -672,6 +672,13 @@ def _fused_adam(
             lr_dict[device] = lr.to(device=device, non_blocking=True)  # type: ignore[union-attr]
             lr = lr_dict[device]
         torch._foreach_add_(device_state_steps, 1)
+
+        strides = set()
+        for l in [device_params, device_grads, device_exp_avgs, device_exp_avg_sqs]:
+            for x in l:
+                strides.add(x.stride())
+        breakpoint()
+
         torch._fused_adam_(
             device_params,
             device_grads,

[mlazos]

@shunting314 isnt there a stride optimization in inductor? Iirc there was a similar issue with distributed