Repository: pytorch
Issue: AMP guards recompilation `dtype mismatch. expected Half, actual Float`
Labels: triaged,oncall: pt2,module: dynamic shapes,module: dynamo

### 🐛 Describe the bug

We had many dynamic recompilation related to AMP (float vs half}. It happens when we use `TORCHDYNAMO_INLINE_INBUILT_NN_MODULES=1` /cc @ezyang @msaroufim @bdhirsh @anijain2305 @chauhang @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @kadeng

```python
    @torch.compile(dynamic=True)
    def forward(self,
                tgt,
                long_term_memory=None,
                short_term_memory=None,
                curr_id_emb=None,
                self_pos=None,
                size_2d=(30, 30)):

        # Self-attention
        _tgt = self.norm1(tgt)
        if not self.use_self_pos:
            self_pos = None
        q = k = self.with_pos_embed(_tgt, self_pos)
        v = _tgt
        if self.global_dilation > 1:
            k = k[::self.global_dilation,:,:]
            v = v[::self.global_dilation,:,:]
        tgt2 = self.self_attn(q, k, v)[0]

        tgt = tgt + self.droppath(tgt2)

        # Long Short-Term Attention
        _tgt = self.norm2(tgt)

        curr_QV = self.linear_QV(_tgt)
        curr_QV = torch.split(curr_QV, self.d_model, dim=2)
        curr_Q = curr_K = curr_QV[0]
        curr_V = curr_QV[1]

        if self.d_att is not None:
            curr_Q = self.linear_Qd(curr_Q)
        local_Q = seq_to_2d(curr_Q, size_2d)

        if curr_id_emb is not None:
            global_K, global_V = self.fuse_key_value_id(
                curr_K, curr_V, curr_id_emb)
            if self.d_att is not None:
                global_K = self.linear_Kd(global_K)
            local_K = seq_to_2d(global_K, size_2d)
            local_V = seq_to_2d(global_V, size_2d)

            if self.global_dilation>1 and self.memory_dilation:
                nhw,bs,ck = global_K.shape
                cv = global_V.shape[-1]
                # n = nhw // (size_2d[0] * size_2d[1])
                d = self.global_dilation
                if self.conv_dilation:
                    unfold_K = global_K.permute(1,2,0).reshape(bs,ck,size_2d[0],size_2d[1])
                    unfold_V = global_V.permute(1,2,0).reshape(bs,cv,size_2d[0],size_2d[1])
                    global_K = self.dilation_conv_K(unfold_K).reshape(bs,ck,-1).permute(2,0,1)
                    global_V = self.dilation_conv_V(unfold_V).reshape(bs,cv,-1).permute(2,0,1)
                else:
                    unfold_K = global_K.view(size_2d[0],size_2d[1],bs,ck)
                    unfold_V = global_V.view(size_2d[0],size_2d[1],bs,cv)
                    global_K = unfold_K[::d,::d,:,:].reshape(-1,bs,ck)
                    global_V = unfold_V[::d,::d,:,:].reshape(-1,bs,cv)
        else:
            global_K, global_V = long_term_memory
            local_K, local_V = short_term_memory
        
        
        if self.memory_dilation:
            tgt2 = self.long_term_attn(curr_Q, global_K, global_V)[0]
        else:
            if self.global_dilation>1:
                nhw,bs,ck = global_K.shape
                cv = global_V.shape[-1]
                n = nhw // (size_2d[0] * size_2d[1])
                d = self.global_dilation
                if self.conv_dilation:
                    unfold_K = global_K.permute(1,2,0).reshape(bs*n,ck,size_2d[0],size_2d[1])
                    unfold_V = global_V.permute(1,2,0).reshape(bs*n,cv,size_2d[0],size_2d[1])
                    dilated_K = self.dilation_conv_K(unfold_K).reshape(bs,ck,-1).permute(2,0,1)
                    dilated_V = self.dilation_conv_V(unfold_V).reshape(bs,cv,-1).permute(2,0,1)
                else:
                    unfold_K = global_K.view(n,size_2d[0],size_2d[1],bs,ck)
                    unfold_V = global_V.view(n,size_2d[0],size_2d[1],bs,cv)
                    dilated_K = unfold_K[:,::d,::d,:,:].reshape(-1,bs,ck)
                    dilated_V = unfold_V[:,::d,::d,:,:].reshape(-1,bs,cv)
            else:
                dilated_K,dilated_V = global_K,global_V
            tgt2 = self.long_term_attn(curr_Q, dilated_K, dilated_V)[0]

        tgt3 = self.short_term_attn(local_Q, local_K, local_V)[0]

        if self.droppath_lst:
            tgt = tgt + self.droppath(tgt2 + tgt3)
        else:
            tgt = tgt + self.lst_dropout(tgt2 + tgt3)

        # Feed-forward
        _tgt = self.norm3(tgt)

        tgt2 = self.linear2(self.activation(self.linear1(_tgt), size_2d))

        tgt = tgt + self.droppath(tgt2)

        return tgt, [[curr_K, curr_V], [global_K, global_V],
                     [local_K, local_V]]
```

Recompilation line 


### Error logs
```python
V0606 12:51:18.544000 133019546625856 torch/_dynamo/guards.py:2610] [1/1] [__recompiles] Recompiling function forward in /workspace/networks/layers/transformer.py:753
V0606 12:51:18.544000 133019546625856 torch/_dynamo/guards.py:2610] [1/1] [__recompiles]     triggered by the following guard failure(s):
V0606 12:51:18.544000 133019546625856 torch/_dynamo/guards.py:2610] [1/1] [__recompiles]     - tensor 'L['tgt']' dtype mismatch. expected Half, actual Float
W0606 12:51:33.555000 133019546625856 torch/fx/experimental/symbolic_shapes.py:4459] [1/1_1] RecursionError in sympy.xreplace(Eq(Mod(256*s1, s8), 0), {s8: evaluate_static_shape_0 + 255, s1: evaluate_static_shape_1 + 1})
V0606 12:51:35.374000 133019546625856 torch/_dynamo/guards.py:2610] [8/1] [__recompiles] Recompiling function torch_dynamo_resume_in_forward_at_836 in /workspace/networks/layers/transformer.py:836
V0606 12:51:35.374000 133019546625856 torch/_dynamo/guards.py:2610] [8/1] [__recompiles]     triggered by the following guard failure(s):
V0606 12:51:35.374000 133019546625856 torch/_dynamo/guards.py:2610] [8/1] [__recompiles]     - tensor 'L['tgt']' dtype mismatch. expected Half, actual Float
W0606 12:51:47.346000 133019546625856 torch/fx/experimental/symbolic_shapes.py:4417] [9/0] ps0 is not in var_ranges, defaulting to unknown range.
V0606 12:51:51.197000 133019546625856 torch/_dynamo/guards.py:2610] [9/1] [__recompiles] Recompiling function forward in /workspace/networks/layers/basic.py:86
V0606 12:51:51.197000 133019546625856 torch/_dynamo/guards.py:2610] [9/1] [__recompiles]     triggered by the following guard failure(s):
V0606 12:51:51.197000 133019546625856 torch/_dynamo/guards.py:2610] [9/1] [__recompiles]     - tensor 'L['x']' size mismatch at index 1. expected 768, actual 256
W0606 12:51:53.963000 133019546625856 torch/fx/experimental/symbolic_shapes.py:4417] [9/1] ps0 is not in var_ranges, defaulting to unknown range.
V0606 12:51:55.704000 133019546625856 torch/_dynamo/guards.py:2610] [1/2] [__recompiles] Recompiling function forward in /workspace/networks/layers/transformer.py:753
V0606 12:51:55.704000 133019546625856 torch/_dynamo/guards.py:2610] [1/2] [__recompiles]     triggered by the following guard failure(s):
V0606 12:51:55.704000 133019546625856 torch/_dynamo/guards.py:2610] [1/2] [__recompiles]     - tensor 'L['tgt']' stride mismatch at index 0. expected 256, actual 1
V0606 12:51:55.704000 133019546625856 torch/_dynamo/guards.py:2610] [1/2] [__recompiles]     - tensor 'L['tgt']' dtype mismatch. expected Half, actual Float
W0606 12:52:12.706000 133019546625856 torch/fx/experimental/symbolic_shapes.py:4459] [1/2_1] RecursionError in sympy.xreplace(Eq(Mod(256*s3, s8), 0), {s8: evaluate_static_shape_0 + 255, s3: evaluate_static_shape_1 + 1})
V0606 12:52:14.035000 133019546625856 torch/_dynamo/guards.py:2610] [8/2] [__recompiles] Recompiling function torch_dynamo_resume_in_forward_at_836 in /workspace/networks/layers/transformer.py:836
V0606 12:52:14.035000 133019546625856 torch/_dynamo/guards.py:2610] [8/2] [__recompiles]     triggered by the following guard failure(s):
V0606 12:52:14.035000 133019546625856 torch/_dynamo/guards.py:2610] [8/2] [__recompiles]     - tensor 'L['tgt']' stride mismatch at index 0. expected 256, actual 1
V0606 12:52:14.035000 133019546625856 torch/_dynamo/guards.py:2610] [8/2] [__recompiles]     - tensor 'L['tgt']' dtype mismatch. expected Half, actual Float
V0606 12:52:24.694000 133019546625856 torch/_dynamo/guards.py:2610] [9/2] [__recompiles] Recompiling function forward in /workspace/networks/layers/basic.py:86
V0606 12:52:24.694000 133019546625856 torch/_dynamo/guards.py:2610] [9/2] [__recompiles]     triggered by the following guard failure(s):
V0606 12:52:24.694000 133019546625856 torch/_dynamo/guards.py:2610] [9/2] [__recompiles]     - tensor 'L['x']' size mismatch at index 1. expected 256, actual 512
V0606 12:52:24.694000 133019546625856 torch/_dynamo/guards.py:2610] [9/2] [__recompiles]     - tensor 'L['x']' size mismatch at index 1. expected 768, actual 512
W0606 12:52:28.706000 133019546625856 torch/fx/experimental/symbolic_shapes.py:4417] [9/2] ps0 is not in var_ranges, defaulting to unknown range.
V0606 12:52:29.754000 133019546625856 torch/_dynamo/guards.py:2610] [1/3] [__recompiles] Recompiling function forward in /workspace/networks/layers/transformer.py:753
V0606 12:52:29.754000 133019546625856 torch/_dynamo/guards.py:2610] [1/3] [__recompiles]     triggered by the following guard failure(s):
V0606 12:52:29.754000 133019546625856 torch/_dynamo/guards.py:2610] [1/3] [__recompiles]     - L['self'].droppath.drop_prob == 0.1                         
V0606 12:52:29.754000 133019546625856 torch/_dynamo/guards.py:2610] [1/3] [__recompiles]     - tensor 'L['tgt']' stride mismatch at index 0. expected 256, actual 1
V0606 12:52:29.754000 133019546625856 torch/_dynamo/guards.py:2610] [1/3] [__recompiles]     - tensor 'L['tgt']' dtype mismatch. expected Half, actual Float
```
### Minified repro

_No response_

### Versions

pytorch-nightly

cc @ezyang @msaroufim @bdhirsh @anijain2305 @chauhang