Repository: pytorch
Issue: Wrapping make_graphed_callables with autocast issue
Labels: module: cuda,triaged,module: amp (automated mixed precision),module: cuda graphs

### 🐛 Describe the bug

A model with PyTorch AMP produces different model results when used in conjunction with CUDA graphs.
In other words, it seems that conducting CUDA graph capture while wrapped with `autocast` leads to wrong outputs during training (specifically, after the first weight update).

Please find the following example to reproduce the issue. Modify `CUDA_GRAPHS` from `False` to `True` between runs.

```python
import torch


CUDA_GRAPHS = False
torch.manual_seed(123)


def model_state(model):
    sum_ = 0
    for name, W in model.named_parameters():
        if 'weight' in name:
            sum_ = sum_ + W.sum()

    return sum_


class model_(torch.nn.Module):
    def __init__(self):
        super(model_, self).__init__()

        self.module1 = torch.nn.Conv2d(1024, 1024, 3, stride=1, padding=1, bias=False).cuda()
        self.module2 = torch.nn.Conv2d(1024, 1024, 3, stride=1, padding=1, bias=False).cuda()
        self.module3 = torch.nn.Conv2d(1024, 1024, 3, stride=1, padding=1, bias=False).cuda()

    def forward(self, x):
        out = self.module1(x)
        out = self.module2(out)
        out = self.module3(out)

        return out


model = model_()

# Confirm that the model weights are identical between runs
print('Model init state: {}'.format(model.module1.weight.data.sum() +
                                    model.module2.weight.data.sum() +
                                    model.module3.weight.data.sum()))

loss_fn = torch.nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)

x = torch.randn(2, 1024, 16, 16, device='cuda')

if CUDA_GRAPHS:
    # I assume that I would like to create the graph with AMP enabled
    with torch.cuda.amp.autocast(enabled=True):
        model = torch.cuda.make_graphed_callables(model, (x,))

real_inputs = [torch.rand_like(x, device="cuda") for _ in range(10)]
real_targets = [torch.rand_like(x, device="cuda") for _ in range(10)]

scaler = torch.cuda.amp.GradScaler(enabled=True)

for i, (data, target) in enumerate(zip(real_inputs, real_targets)):
    optimizer.zero_grad(set_to_none=True)

    with torch.cuda.amp.autocast(enabled=True):
        tmp = model(data)
        print('{}> model_state:{}, in_sum: {}, out_sum: {}'.format(i, model_state(model), data.sum(), tmp.sum()))

        loss = loss_fn(tmp, target)

    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()
```

My results for `CUDA_GRAPHS = False`:
```
Model init state: -10.634528160095215
0> model_state:-10.634528160095215, in_sum: 262280.21875, out_sum: -652.855712890625
1> model_state:-9.640216827392578, in_sum: 262201.5625, out_sum: 25918.837890625
2> model_state:-9.438332557678223, in_sum: 262223.3125, out_sum: 53148.3828125
3> model_state:-9.562451362609863, in_sum: 262339.78125, out_sum: 92139.4921875
4> model_state:-10.297986030578613, in_sum: 261779.03125, out_sum: 144155.390625
5> model_state:-11.313785552978516, in_sum: 262076.96875, out_sum: 198276.78125
6> model_state:-12.076295852661133, in_sum: 262292.625, out_sum: 228802.546875
7> model_state:-12.462014198303223, in_sum: 262321.3125, out_sum: 235795.96875
8> model_state:-12.78150749206543, in_sum: 261910.171875, out_sum: 240080.015625
9> model_state:-13.009492874145508, in_sum: 261907.109375, out_sum: 241984.59375
```

My results for `CUDA_GRAPHS = True`:
```
Model init state: -10.634528160095215
0> model_state:-10.634528160095215, in_sum: 262280.21875, out_sum: -652.855712890625
1> model_state:-9.640216827392578, in_sum: 262201.5625, out_sum: -641.5584716796875
2> model_state:-8.379328727722168, in_sum: 262223.3125, out_sum: -590.2320556640625
3> model_state:-6.929278373718262, in_sum: 262339.78125, out_sum: -535.804443359375
4> model_state:-5.820806503295898, in_sum: 261779.03125, out_sum: -617.807861328125
5> model_state:-4.839597702026367, in_sum: 262076.96875, out_sum: -587.7401123046875
6> model_state:-3.75885009765625, in_sum: 262292.625, out_sum: -579.7599487304688
7> model_state:-2.5434188842773438, in_sum: 262321.3125, out_sum: -611.613525390625
8> model_state:-1.3896923065185547, in_sum: 261910.171875, out_sum: -603.328857421875
9> model_state:-0.24511337280273438, in_sum: 261907.109375, out_sum: -616.830322265625
```

Notice that:
- At the start of the first iteration, both models are exactly the same, which is why they produce the exact same output and output grads, and which is why after grad update the models still look the same.
- At the second iteration, for some reason, the model output is different. Therefore, a different loss value is computed, different weight grads, and after the `step` different model weights.
- Only at the third iteration the models start to diverge.

Some more insights:
- I have noticed that same values are produced if I use `cache_enabled=False` for `autocast` when I create the graph.
However, it works with the 21.10 container (PyTorch 1.10) and does not work with 21.11 or 21.12 containers (PyTorch 1.11).
- In addition, using `autocast` while capturing the graph is working, that is, moving the `autocast` to within the model, though it doesn't seem right from a coding perspective.

### Versions

The bug occurs on NVIDIA PyTorch container versions larger or equal to 21.11.

cc @ngimel @mcarilli @ptrblck