Repository: pytorch
Issue: Very large memory increase when combining bfloat16 autocast with torch.compile
Labels: good first issue,module: memory usage,triaged,module: amp (automated mixed precision),oncall: pt2,module: inductor

### 🐛 Describe the bug

Use of autocast causes a slight but unexpected and tolerable increase in memory, but the combination of autocast with torch.compile causes very large increase (8x with a 2048x2048 F.linear, and continues to grow as dimensions increase).  Strangely, it can be attenuated by using a broken down version of F.linear. I tested with bfloat16, but the results are the same with float16 autocast.
```

import torch
import torch.nn as nn

device = 'cuda'
torch.set_float32_matmul_precision('high') #has no effect

tensor = torch.rand((50, 64, 2048), device=device)
add_req_grad = torch.tensor([1.0], requires_grad=True).to(device)
fc = nn.Linear(2048, 2048).to(device)
tensor += add_req_grad
outputs = []
for i in range(tensor.size(1)):
    input = tensor[:, i, :]
    out = fc(input)
    outputs.append(out)
print(f"compile: False, autocast: False, modify input: True, mem allocated: {torch.cuda.memory_allocated()/1e9} GB")
print()
torch.cuda.empty_cache()

tensor = torch.rand((50, 64, 2048), device=device)
add_req_grad = torch.tensor([1.0], requires_grad=True).to(device)
fc = nn.Linear(2048, 2048).to(device)
fc = torch.compile(fc)
tensor += add_req_grad
outputs = []
for i in range(tensor.size(1)):
    input = tensor[:, i, :]
    out = fc(input)
    outputs.append(out)
print(f"compile: True, autocast: False, modify input: True, mem allocated: {torch.cuda.memory_allocated()/1e9} GB")
print()
torch.cuda.empty_cache()

tensor = torch.rand((50, 64, 2048), device=device)
add_req_grad = torch.tensor([1.0], requires_grad=True).to(device)
fc = nn.Linear(2048, 2048).to(device)
fc = torch.compile(fc)
with torch.autocast(device_type=device, dtype=torch.bfloat16):
    tensor += add_req_grad
    outputs = []
    for i in range(tensor.size(1)):
        input = tensor[:, i, :]
        out = fc(input)
        outputs.append(out)
print(f"compile: True, autocast: True, modify input: True, mem allocated: {torch.cuda.memory_allocated()/1e9} GB")
print()
torch.cuda.empty_cache()

tensor = torch.rand((50, 64, 2048), device=device)
add_req_grad = torch.tensor([1.0], requires_grad=True).to(device)
fc = nn.Linear(2048, 2048).to(device)
with torch.autocast(device_type=device, dtype=torch.bfloat16):
    tensor += add_req_grad
    outputs = []
    for i in range(tensor.size(1)):
        input = tensor[:, i, :]
        out = fc(input)
        outputs.append(out)
print(f"compile: False, autocast: True, modify input: True, mem allocated: {torch.cuda.memory_allocated()/1e9} GB")
print()
torch.cuda.empty_cache()


tensor = torch.rand((50, 64, 2048), device=device)
fc = nn.Linear(2048, 2048).to(device)
fc = torch.compile(fc)
with torch.autocast(device_type=device, dtype=torch.bfloat16):
    outputs = []
    for i in range(tensor.size(1)):
        input = tensor[:, i, :]
        out = fc(input)
        outputs.append(out)
print(f"compile: True, autocast: True, modify input: False, mem allocated: {torch.cuda.memory_allocated()/1e9} GB")
print()
torch.cuda.empty_cache()


def CustomFLinear(base_activations, weights, bias=None):
    weights = weights.unsqueeze(0)
    intermediate_values = base_activations.unsqueeze(1) * weights
    output_activations = intermediate_values.sum(dim=2)
    if bias is not None:
        output_activations = output_activations + bias
    return output_activations

class CustomLinear(nn.Module):
    def __init__(self, in_features, out_features, bias=True, device='cuda', dtype=None):
        self.factory_kwargs = {'device': device, 'dtype': dtype}
        super(CustomLinear, self).__init__()

        self.in_features = in_features
        self.out_features = out_features

        self.weight = nn.Parameter(torch.empty((out_features, in_features), **self.factory_kwargs))
        if bias:
            self.bias = torch.empty(out_features, **self.factory_kwargs)
        else:
            self.register_parameter('bias', None)

    def forward(self, input):
        out = CustomFLinear(input, self.weight, self.bias)
        return out

tensor = torch.rand((50, 64, 2048), device=device)
add_req_grad = torch.tensor([1.0], requires_grad=True).to(device)
fc = CustomLinear(2048, 2048).to(device)
fc = torch.compile(fc)
with torch.autocast(device_type=device, dtype=torch.bfloat16):
    tensor += add_req_grad
    outputs = []
    for i in range(tensor.size(1)):
        input = tensor[:, i, :]
        out = fc(input)
        outputs.append(out)
print('** Custom/Manual linear used **')
print(f"compile: True, autocast: True, modify input: True, mem allocated: {torch.cuda.memory_allocated()/1e9} GB")
print()
torch.cuda.empty_cache()
```


Output is: 
compile: False, autocast: False, modify input: True, mem allocated: 0.078782976 GB
compile: True, autocast: False, modify input: True, mem allocated: 0.078782976 GB
compile: True, autocast: True, modify input: True, mem allocated: 0.632439296 GB
compile: False, autocast: True, modify input: True, mem allocated: 0.1207424 GB
compile: True, autocast: True, modify input: False, mem allocated: 0.112353792 GB

** Custom/Manual linear used **
compile: True, autocast: True, modify input: True, mem allocated: 0.1291392 GB

### Error logs

_No response_

### Minified repro

_No response_

### Versions

2.3.1

cc @mcarilli @ptrblck @leslie-fang-intel @jgong5 @ezyang @chauhang @penguinwu @voznesenskym @EikanWang @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @ColinPeppler @amjames @desertfire