[Dynamo] Unable to Trace AdamW Optimizer when there is LR Scheduler
Labels: high priority, module: optimizer, triaged, oncall: pt2, module: dynamic shapes

### 🐛 Describe the bug

## What is Broken
- the lr in AdamW Optimizer is a float so causes a recompile everytime LRScheduler changes the lr

## Suggested fix
- fix the graph break in `_multi_tensor AdamW Optimizer` and/or normal `AdamW Optimizer`
- change that LRScheduler base class outputs Scalar Tensor and AdamW Optimizer accepts Scalar Tensor for lr
   - this is similar to what @janeyx99 is doing in Fused AdamW Optimizer https://github.com/pytorch/pytorch/pull/106916

## Why this should be fixed
- a lot of the community uses AdamW Optimizer in combination with LR Scheduler and wants to torch compile it.

## Reprod Script
```python
import torch
import torch.nn as nn
import torch.optim as optim
import torch._dynamo.utils
import torch._dynamo
import logging
torch._dynamo.config.log_level = logging.INFO

class SimpleModel(nn.Module):
    def __init__(self, input_size, hidden_size,  num_classes):
        super(SimpleModel, self).__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.fc2 = nn.Linear(hidden_size, num_classes)


    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x
    

input_size = 784
hidden_size = 500
num_classes = 10

model = SimpleModel(input_size, hidden_size, num_classes)


criterion = nn.CrossEntropyLoss()
optimizer = torch.optim._multi_tensor.AdamW(model.parameters(), lr=0.001)
scheduler = optim.lr_scheduler.StepLR(optimizer,step_size=1, gamma=0.1)

inputs = torch.randn(64, 784)

labels = torch.randint(0,10,(64,))

num_epochs = 5
def opt_scheduler_step(optimizer):
    optimizer.step()

    return optimizer

opt_scheduler_step = torch.compile(opt_scheduler_step, backend="eager", fullgraph=False)

for epoch in range(num_epochs):
    outputs = model(inputs)
    loss = criterion(outputs, labels)

    optimizer.zero_grad()

    loss.backward()
    optimizer = opt_scheduler_step(optimizer)
    scheduler.step()

    print(loss.item())
```
## Logs
```
[2023-08-11 22:52:19,789] torch._dynamo.symbolic_convert: [INFO] Step 1: torchdynamo start tracing opt_scheduler_step
[2023-08-11 22:52:19,857] torch._dynamo.symbolic_convert: [INFO] Step 1: torchdynamo start tracing wrapper
[2023-08-11 22:52:19,858] torch._dynamo.variables.torch: [WARNING] Profiler will be ignored
[2023-08-11 22:52:19,869] torch._dynamo.symbolic_convert: [INFO] Step 1: torchdynamo start tracing _use_grad
[2023-08-11 22:52:19,874] torch._dynamo.symbolic_convert: [INFO] Step 1: torchdynamo start tracing step
[2023-08-11 22:52:19,878] torch._dynamo.symbolic_convert: [INFO] Step 1: torchdynamo start tracing <graph break in step>
[2023-08-11 22:52:19,898] torch._dynamo.symbolic_convert: [INFO] Step 1: torchdynamo start tracing adamw
[2023-08-11 22:52:19,910] torch._dynamo.symbolic_convert: [INFO] Step 1: torchdynamo start tracing <graph break in adamw>
[2023-08-11 22:52:19,912] torch._dynamo.symbolic_convert: [INFO] Step 1: torchdynamo start tracing <graph break in wrapper>
2.3640196323394775
[2023-08-11 22:52:19,927] torch._dynamo.symbolic_convert: [INFO] Step 1: torchdynamo start tracing adamw
1.709822654724121
[2023-08-11 22:52:19,949] torch._dynamo.symbolic_convert: [INFO] Step 1: torchdynamo start tracing adamw
1.6508420705795288
[2023-08-11 22:52:19,972] torch._dynamo.symbolic_convert: [INFO] Step 1: torchdynamo start tracing adamw
1.6449681520462036
[2023-08-11 22:52:19,993] torch._dynamo.symbolic_convert: [INFO] Step 1: torchdynamo start tracing adamw
1.6443781852722168
```

cc: @janeyx99 @cchan @ezhang887 

### Versions

python==3.8.13
torch==2.0.1

cc @ezyang @gchanan @zou3519 @vincentqb @jbschlosser @albanD @janeyx99 @crcrpar @msaroufim @wconstab @bdhirsh @anijain2305