Repository: pytorch
Issue: `torch.autocast(bfloat16)` runs bwd matmuls in fp16
Labels: triaged,module: amp (automated mixed precision)


<details>
<summary> Repro code </summary>

```
"""
python test/test_amp.py
"""
import itertools

import weakref
from functools import partial

import torch
import torch.nn as nn
import torch.overrides
from torch.utils._python_dispatch import TorchDispatchMode
from torch.utils._pytree import tree_map
from torch.utils.weak import WeakIdRef


dtype_abbrs = {
    torch.bfloat16: "bf16",
    torch.float64: "f64",
    torch.float32: "f32",
    torch.float16: "f16",
    torch.complex32: "c32",
    torch.complex64: "c64",
    torch.complex128: "c128",
    torch.int8: "i8",
    torch.int16: "i16",
    torch.int32: "i32",
    torch.int64: "i64",
    torch.bool: "b8",
    torch.uint8: "u8",
}


def main():
    torch.manual_seed(42)
    model = nn.Transformer(num_encoder_layers=1, num_decoder_layers=1, device="cuda")
    src = torch.rand((10, 32, 512), device="cuda")
    tgt = torch.rand((20, 32, 512), device="cuda")
    amp_ctx = torch.autocast(device_type="cuda", dtype=torch.bfloat16)
    logging_ctx = LoggingMode(collect_logs=True)
    with amp_ctx, logging_ctx:
        model(src, tgt).sum().backward()
    print(logging_ctx.str_logs())


class Lit:
    def __init__(self, s):
        self.s = s

    def __repr__(self):
        return self.s


class LoggingMode(TorchDispatchMode):
    next_id: int

    def __init__(self, with_type: bool = True, collect_logs: bool = False):
        self.memo = {}
        self.next_id = 0
        self.with_type = with_type
        self.collect_logs = collect_logs
        self.logs = []
        try:
            self.rank = torch.distributed.get_rank()
        except:
            self.rank = 0

    def _shortid(self, t: torch.Tensor) -> int:
        o = WeakIdRef(t)
        weak_self = weakref.ref(self)

        def del_memo():
            self = weak_self()
            if self is None:
                return
            self.memo.pop(o, None)

        weakref.finalize(t, del_memo)
        if o not in self.memo:
            self.memo[o] = self.next_id
            self.next_id += 1
        return self.memo[o]

    def _fmt(self, a: object, with_type: bool = False) -> str:
        if isinstance(a, torch.Tensor):
            maybe_type = ""
            if with_type and self.with_type:
                maybe_type = f": {dtype_abbrs[a.dtype]}[{', '.join(map(str, a.shape))}]"
            return Lit(f"${self._shortid(a)}{maybe_type}")
        else:
            return a

    def str_logs(self):
        return "\n".join(self.logs)

    def __torch_dispatch__(self, func, types, args=(), kwargs=None):
        if kwargs is None:
            kwargs = {}
        rs = func(*args, **kwargs)
        fmt_args = ", ".join(
            itertools.chain(
                (repr(tree_map(self._fmt, a)) for a in args),
                (f"{k}={tree_map(self._fmt, v)}" for k, v in kwargs.items()),
            )
        )
        fmt_rets = repr(tree_map(partial(self._fmt, with_type=True), rs))
        log_msg = f"{fmt_rets} = {torch.overrides.resolve_name(func)}({fmt_args})"
        if self.collect_logs:
            self.logs.append(log_msg)
        elif self.rank == 0:
            print(log_msg)
        return rs


if __name__ == "__main__":
    main()
```

</details>

<details>
<summary> Logged Ops </summary>

E.g., search for `$254: f16[512, 2048] = aten.mm.default($253, $252)`:

```
$1: bf16[1536] = aten._to_copy.default($0, dtype=torch.bfloat16)
$3: bf16[1536, 512] = aten._to_copy.default($2, dtype=torch.bfloat16)
$5: bf16[10, 32, 512] = aten._to_copy.default($4, dtype=torch.bfloat16)
$6: bf16[512, 1536] = aten.t.default($3)
$7: bf16[320, 512] = aten.view.default($5, [320, 512])
$8: bf16[320, 1536] = aten.addmm.default($1, $7, $6)
$9: bf16[10, 32, 1536] = aten.view.default($8, [10, 32, 1536])
$10: bf16[10, 32, 3, 512] = aten.view.default($9, [10, 32, 3, 512])
$11: bf16[1, 10, 32, 3, 512] = aten.unsqueeze.default($10, 0)
$12: bf16[3, 10, 32, 1, 512] = aten.transpose.int($11, 0, -2)
$13: bf16[3, 10, 32, 512] = aten.squeeze.dim($12, -2)
$14: bf16[3, 10, 32, 512] = aten.clone.default($13, memory_format=torch.contiguous_format)
$15: bf16[10, 32, 512] = aten.select.int($14, 0, 0)
$16: bf16[10, 32, 512] = aten.select.int($14, 0, 1)
$17: bf16[10, 32, 512] = aten.select.int($14, 0, 2)
$18: bf16[10, 256, 64] = aten.view.default($15, [10, 256, 64])
$19: bf16[256, 10, 64] = aten.transpose.int($18, 0, 1)
$20: bf16[10, 256, 64] = aten.view.default($16, [10, 256, 64])
$21: bf16[256, 10, 64] = aten.transpose.int($20, 0, 1)
$22: bf16[10, 256, 64] = aten.view.default($17, [10, 256, 64])
$23: bf16[256, 10, 64] = aten.transpose.int($22, 0, 1)
$24: bf16[32, 8, 10, 64] = aten.view.default($19, [32, 8, 10, 64])
$25: bf16[32, 8, 10, 64] = aten.view.default($21, [32, 8, 10, 64])
$26: bf16[32, 8, 10, 64] = aten.view.default($23, [32, 8, 10, 64])
($27: bf16[32, 8, 10, 64], $28: f32[32, 8, 32], $29: i64[], $30: i64[]) = aten._scaled_dot_product_efficient_attention.default($24, $25, $26, True, 0.1)
$31: bf16[32, 8, 10, 64] = aten.detach.default($27)
$32: bf16[10, 32, 8, 64] = aten.permute.default($27, [2, 0, 1, 3])
$33: bf16[10, 32, 8, 64] = aten.clone.default($32, memory_format=torch.contiguous_format)
$34: bf16[320, 512] = aten.view.default($33, [320, 512])
$36: bf16[512] = aten._to_copy.default($35, dtype=torch.bfloat16)
$38: bf16[512, 512] = aten._to_copy.default($37, dtype=torch.bfloat16)
$39: bf16[512, 512] = aten.t.default($38)
$40: bf16[320, 512] = aten.addmm.default($36, $34, $39)
$41: bf16[10, 32, 512] = aten.view.default($40, [10, 32, 512])
($42: bf16[10, 32, 512], $43: b8[10, 32, 512]) = aten.native_dropout.default($41, 0.1, True)
$44: f32[10, 32, 512] = aten.add.Tensor($4, $42)
($47: f32[10, 32, 512], $48: f32[10, 32, 1], $49: f32[10, 32, 1]) = aten.native_layer_norm.default($44, [512], $45, $46, 1e-05)
$51: bf16[2048] = aten._to_copy.default($50, dtype=torch.bfloat16)
$53: bf16[2048, 512] = aten._to_copy.default($52, dtype=torch.bfloat16)
$54: bf16[10, 32, 512] = aten._to_copy.default($47, dtype=torch.bfloat16)
$55: bf16[512, 2048] = aten.t.default($53)
$56: bf16[320, 512] = aten.view.default($54, [320, 512])
$57: bf16[320, 2048] = aten.addmm.default($51, $56, $55)
$58: bf16[10, 32, 2048] = aten.view.default($57, [10, 32, 2048])
$59: bf16[10, 32, 2048] = aten.relu.default($58)
$60: bf16[10, 32, 2048] = aten.detach.default($59)
($61: bf16[10, 32, 2048], $62: b8[10, 32, 2048]) = aten.native_dropout.default($59, 0.1, True)
$64: bf16[512] = aten._to_copy.default($63, dtype=torch.bfloat16)
$66: bf16[512, 2048] = aten._to_copy.default($65, dtype=torch.bfloat16)
$67: bf16[2048, 512] = aten.t.default($66)
$68: bf16[320, 2048] = aten.view.default($61, [320, 2048])
$69: bf16[320, 512] = aten.addmm.default($64, $68, $67)
$70: bf16[10, 32, 512] = aten.view.default($69, [10, 32, 512])
($71: bf16[10, 32, 512], $72: b8[10, 32, 512]) = aten.native_dropout.default($70, 0.1, True)
$73: f32[10, 32, 512] = aten.add.Tensor($47, $71)
($76: f32[10, 32, 512], $77: f32[10, 32, 1], $78: f32[10, 32, 1]) = aten.native_layer_norm.default($73, [512], $74, $75, 1e-05)
($81: f32[10, 32, 512], $82: f32[10, 32, 1], $83: f32[10, 32, 1]) = aten.native_layer_norm.default($76, [512], $79, $80, 1e-05)
$85: bf16[1536] = aten._to_copy.default($84, dtype=torch.bfloat16)
$87: bf16[1536, 512] = aten._to_copy.default($86, dtype=torch.bfloat16)
$89: bf16[20, 32, 512] = aten._to_copy.default($88, dtype=torch.bfloat16)
$90: bf16[512, 1536] = aten.t.default($87)
$91: bf16[640, 512] = aten.view.default($89, [640, 512])
$92: bf16[640, 1536] = aten.addmm.default($85, $91, $90)
$93: bf16[20, 32, 1536] = aten.view.default($92, [20, 32, 1536])
$94: bf16[20, 32, 3, 512] = aten.view.default($93, [20, 32, 3, 512])
$95: bf16[1, 20, 32, 3, 512] = aten.unsqueeze.default($94, 0)
$96: bf16[3, 20, 32, 1, 512] = aten.transpose.int($95, 0, -2)
$97: bf16[3, 20, 32, 512] = aten.squeeze.dim($96, -2)
$98: bf16[3, 20, 32, 512] = aten.clone.default($97, memory_format=torch.contiguous_format)
$99: bf16[20, 32, 512] = aten.select.int($98, 0, 0)
$100: bf16[20, 32, 512] = aten.select.int($98, 0, 1)
$101: bf16[20, 32, 512] = aten.select.int($98, 0, 2)
$102: bf16[20, 256, 64] = aten.view.default($99, [20, 256, 64])
$103: bf16[256, 20, 64] = aten.transpose.int($102, 0, 1)
$104: bf16[20, 256, 64] = aten.view.default($100, [20, 256, 64])
$105: bf16[256, 20, 64] = aten.transpose.int($104, 0, 1)
$106: bf16[20, 256, 64] = aten.view.default($101, [20, 256, 64])
$107: bf16[256, 20, 64] = aten.transpose.int($106, 0, 1)
$108: bf16[32, 8, 20, 64] = aten.view.default($103, [32, 8, 20, 64])
$109: bf16[32, 8, 20, 64] = aten.view.default($105, [32, 8, 20, 64])
$110: bf16[32, 8, 20, 64] = aten.view.default($107, [32, 8, 20, 64])
($111: bf16[32, 8, 20, 64], $112: f32[32, 8, 32], $113: i64[], $114: i64[]) = aten._scaled_dot_product_efficient_attention.default($108, $109, $110, True, 0.1)
$115: bf16[32, 8, 20, 64] = aten.detach.default($111)
$116: bf16[20, 32, 8, 64] = aten.permute.default($111, [2, 0, 1, 3])
$117: bf16[20, 32, 8, 64] = aten.clone.default($116, memory_format=torch.contiguous_format)
$118: bf16[640, 512] = aten.view.default($117, [640, 512])
$120: bf16[512] = aten._to_copy.default($119, dtype=torch.bfloat16)
$122: bf16[512, 512] = aten._to_copy.default($121, dtype=torch.bfloat16)
$123: bf16[512, 512] = aten.t.default($122)
$124: bf16[640, 512] = aten.addmm.default($120, $118, $123)
$125: bf16[20, 32, 512] = aten.view.default($124, [20, 32, 512])
($126: bf16[20, 32, 512], $127: b8[20, 32, 512]) = aten.native_dropout.default($125, 0.1, True)
$128: f32[20, 32, 512] = aten.add.Tensor($88, $126)
($131: f32[20, 32, 512], $132: f32[20, 32, 1], $133: f32[20, 32, 1]) = aten.native_layer_norm.default($128, [512], $129, $130, 1e-05)
[$135: f32[512, 512], $136: f32[1024, 512]] = aten.split_with_sizes.default($134, [512, 1024])
[$138: f32[512], $139: f32[1024]] = aten.split_with_sizes.default($137, [512, 1024])
$140: bf16[512] = aten._to_copy.default($138, dtype=torch.bfloat16)
$141: bf16[512, 512] = aten._to_copy.default($135, dtype=torch.bfloat16)
$142: bf16[20, 32, 512] = aten._to_copy.default($131, dtype=torch.bfloat16)
$143: bf16[512, 512] = aten.t.default($141)
$144: bf16[640, 512] = aten.view.default($142, [640, 512])
$145: bf16[640, 512] = aten.addmm.default($140, $144, $143)
$146: bf16[20, 32, 512] = aten.view.default($145, [20, 32, 512])
$147: bf16[1024] = aten._to_copy.default($139, dtype=torch.bfloat16)
$148: bf16[1024, 512] = aten._to_copy.default($136, dtype=torch.bfloat16)
$149: bf16[10, 32, 512] = aten._to_copy.default($81, dtype=torch.bfloat16)
$150: bf16[512, 1024] = aten.t.default($148)
$151: bf16[320, 512] = aten.view.default($149, [320, 512])
$152: bf16[320, 1024] = aten.addmm.default($147, $151, $150)
$153: bf16[10, 32, 1024] = aten.view.default($152, [10, 32, 1024])
$154: bf16[10, 32, 2, 512] = aten.view.default($153, [10, 32, 2, 512])
$155: bf16[1, 10, 32, 2, 512] = aten.unsqueeze.default($154, 0)
$156: bf16[2, 10, 32, 1, 512] = aten.transpose.int($155, 0, -2)
$157: bf16[2, 10, 32, 512] = aten.squeeze.dim($156, -2)
$158: bf16[2, 10, 32, 512] = aten.clone.default($157, memory_format=torch.contiguous_format)
$159: bf16[10, 32, 512] = aten.select.int($158, 0, 0)
$160: bf16[10, 32, 512] = aten.select.int($158, 0, 1)
$161: bf16[20, 256, 64] = aten.view.default($146, [20, 256, 64])
$162: bf16[256, 20, 64] = aten.transpose.int($161, 0, 1)
$163: bf16[10, 256, 64] = aten.view.default($159, [10, 256, 64])
$164: bf16[256, 10, 64] = aten.transpose.int($163, 0, 1)
$165: bf16[10, 256, 64] = aten.view.default($160, [10, 256, 64])
$166: bf16[256, 10, 64] = aten.transpose.int($165, 0, 1)
$167: bf16[32, 8, 20, 64] = aten.view.default($162, [32, 8, 20, 64])
$168: bf16[32, 8, 10, 64] = aten.view.default($164, [32, 8, 10, 64])
$169: bf16[32, 8, 10, 64] = aten.view.default($166, [32, 8, 10, 64])
($170: bf16[32, 8, 20, 64], $171: f32[32, 8, 32], $172: i64[], $173: i64[]) = aten._scaled_dot_product_efficient_attention.default($167, $168, $169, True, 0.1)
$174: bf16[32, 8, 20, 64] = aten.detach.default($170)
$175: bf16[20, 32, 8, 64] = aten.permute.default($170, [2, 0, 1, 3])
$176: bf16[20, 32, 8, 64] = aten.clone.default($175, memory_format=torch.contiguous_format)
$177: bf16[640, 512] = aten.view.default($176, [640, 512])
$179: bf16[512] = aten._to_copy.default($178, dtype=torch.bfloat16)
$181: bf16[512, 512] = aten._to_copy.default($180, dtype=torch.bfloat16)
$182: bf16[512, 512] = aten.t.default($181)
$183: bf16[640, 512] = aten.addmm.default($179, $177, $182)
$184: bf16[20, 32, 512] = aten.view.default($183, [20, 32, 512])
($185: bf16[20, 32, 512], $186: b8[20, 32, 512]) = aten.native_dropout.default($184, 0.1, True)
$187: f32[20, 32, 512] = aten.add.Tensor($131, $185)
($190: f32[20, 32, 512], $191: f32[20, 32, 1], $192: f32[20, 32, 1]) = aten.native_layer_norm.default($187, [512], $188, $189, 1e-05)
$194: bf16[2048] = aten._to_copy.default($193, dtype=torch.bfloat16)
$196: bf16[2048, 512] = aten._to_copy.default($195, dtype=torch.bfloat16)
$197: bf16[20, 32, 512] = aten._to_copy.default($190, dtype=torch.bfloat16)
$198: bf16[512, 2048] = aten.t.default($196)
$199: bf16[640, 512] = aten.view.default($197, [640, 512])
$200: bf16[640, 2048] = aten.addmm.default($194, $199, $198)
$201: bf16[20, 32, 2048] = aten.view.default($200, [20, 32, 2048])
$202: bf16[20, 32, 2048] = aten.relu.default($201)
$203: bf16[20, 32, 2048] = aten.detach.default($202)
($204: bf16[20, 32, 2048], $205: b8[20, 32, 2048]) = aten.native_dropout.default($202, 0.1, True)
$207: bf16[512] = aten._to_copy.default($206, dtype=torch.bfloat16)
$209: bf16[512, 2048] = aten._to_copy.default($208, dtype=torch.bfloat16)
$210: bf16[2048, 512] = aten.t.default($209)
$211: bf16[640, 2048] = aten.view.default($204, [640, 2048])
$212: bf16[640, 512] = aten.addmm.default($207, $211, $210)
$213: bf16[20, 32, 512] = aten.view.default($212, [20, 32, 512])
($214: bf16[20, 32, 512], $215: b8[20, 32, 512]) = aten.native_dropout.default($213, 0.1, True)
$216: f32[20, 32, 512] = aten.add.Tensor($190, $214)
($219: f32[20, 32, 512], $220: f32[20, 32, 1], $221: f32[20, 32, 1]) = aten.native_layer_norm.default($216, [512], $217, $218, 1e-05)
($224: f32[20, 32, 512], $225: f32[20, 32, 1], $226: f32[20, 32, 1]) = aten.native_layer_norm.default($219, [512], $222, $223, 1e-05)
$227: f32[] = aten.sum.default($224, dtype=torch.float32)
$228: f32[] = aten.ones_like.default($227, pin_memory=False, memory_format=torch.preserve_format)
$229: f32[20, 32, 512] = aten.expand.default($228, [20, 32, 512])
($230: f32[20, 32, 512], $231: f32[512], $232: f32[512]) = aten.native_layer_norm_backward.default($229, $219, [512], $225, $226, $222, $223, [True, True, True])
$233: f32[512] = aten.detach.default($231)
$234: f32[512] = aten.detach.default($233)
$235: f32[512] = aten.detach.default($232)
$236: f32[512] = aten.detach.default($235)
($237: f32[20, 32, 512], $238: f32[512], $239: f32[512]) = aten.native_layer_norm_backward.default($230, $216, [512], $220, $221, $217, $218, [True, True, True])
$240: f32[512] = aten.detach.default($238)
$241: f32[512] = aten.detach.default($240)
$242: f32[512] = aten.detach.default($239)
$243: f32[512] = aten.detach.default($242)
$244: bf16[20, 32, 512] = aten._to_copy.default($237, dtype=torch.bfloat16)
$245: bf16[20, 32, 512] = aten.native_dropout_backward.default($244, $215, 1.1111111111111112)
$246: bf16[640, 512] = aten.view.default($245, [640, 512])
$247: bf16[512, 2048] = aten.t.default($210)
$248: f16[512, 2048] = aten._to_copy.default($247, dtype=torch.float16)
$249: f16[640, 512] = aten._to_copy.default($246, dtype=torch.float16)
$250: f16[640, 2048] = aten.mm.default($249, $248)
$251: bf16[512, 640] = aten.t.default($246)
$252: f16[640, 2048] = aten._to_copy.default($211, dtype=torch.float16)
$253: f16[512, 640] = aten._to_copy.default($251, dtype=torch.float16)
$254: f16[512, 2048] = aten.mm.default($253, $252)
$255: f16[2048, 512] = aten.t.default($254)
$256: f32[1, 512] = aten.sum.dim_IntList($246, [0], True, dtype=torch.float32)
$257: f32[512] = aten.view.default($256, [512])
$258: bf16[512] = aten._to_copy.default($257, dtype=torch.bfloat16)
$259: bf16[640, 2048] = aten._to_copy.default($250, dtype=torch.bfloat16)
$260: bf16[2048, 512] = aten._to_copy.default($255, dtype=torch.bfloat16)
$261: bf16[20, 32, 2048] = aten.view.default($259, [20, 32, 2048])
$262: bf16[512, 2048] = aten.t.default($260)
$263: f32[512, 2048] = aten._to_copy.default($262, dtype=torch.float32, layout=torch.strided, device=cuda:0)
$264: f32[512, 2048] = aten.detach.default($263)
$265: f32[512, 2048] = aten.detach.default($264)
$266: f32[512] = aten._to_copy.default($258, dtype=torch.float32, layout=torch.strided, device=cuda:0)
$267: f32[512] = aten.detach.default($266)
$268: f32[512] = aten.detach.default($267)
$269: bf16[20, 32, 2048] = aten.native_dropout_backward.default($261, $205, 1.1111111111111112)
$270: bf16[20, 32, 2048] = aten.detach.default($203)
$271: bf16[20, 32, 2048] = aten.threshold_backward.default($269, $270, 0)
$272: bf16[640, 2048] = aten.view.default($271, [640, 2048])
$273: bf16[2048, 512] = aten.t.default($198)
$274: f16[2048, 512] = aten._to_copy.default($273, dtype=torch.float16)
$275: f16[640, 2048] = aten._to_copy.default($272, dtype=torch.float16)
$276: f16[640, 512] = aten.mm.default($275, $274)
$277: bf16[2048, 640] = aten.t.default($272)
$278: f16[640, 512] = aten._to_copy.default($199, dtype=torch.float16)
$279: f16[2048, 640] = aten._to_copy.default($277, dtype=torch.float16)
$280: f16[2048, 512] = aten.mm.default($279, $278)
$281: f16[512, 2048] = aten.t.default($280)
$282: f32[1, 2048] = aten.sum.dim_IntList($272, [0], True, dtype=torch.float32)
$283: f32[2048] = aten.view.default($282, [2048])
$284: bf16[2048] = aten._to_copy.default($283, dtype=torch.bfloat16)
$285: bf16[640, 512] = aten._to_copy.default($276, dtype=torch.bfloat16)
$286: bf16[512, 2048] = aten._to_copy.default($281, dtype=torch.bfloat16)
$287: bf16[20, 32, 512] = aten.view.default($285, [20, 32, 512])
$288: bf16[2048, 512] = aten.t.default($286)
$289: f32[20, 32, 512] = aten._to_copy.default($287, dtype=torch.float32, layout=torch.strided, device=cuda:0)
$290: f32[20, 32, 512] = aten.add.Tensor($237, $289)
$291: f32[2048, 512] = aten._to_copy.default($288, dtype=torch.float32, layout=torch.strided, device=cuda:0)
$292: f32[2048, 512] = aten.detach.default($291)
$293: f32[2048, 512] = aten.detach.default($292)
$294: f32[2048] = aten._to_copy.default($284, dtype=torch.float32, layout=torch.strided, device=cuda:0)
$295: f32[2048] = aten.detach.default($294)
$296: f32[2048] = aten.detach.default($295)
($297: f32[20, 32, 512], $298: f32[512], $299: f32[512]) = aten.native_layer_norm_backward.default($290, $187, [512], $191, $192, $188, $189, [True, True, True])
$300: f32[512] = aten.detach.default($298)
$301: f32[512] = aten.detach.default($300)
$302: f32[512] = aten.detach.default($299)
$303: f32[512] = aten.detach.default($302)
$304: bf16[20, 32, 512] = aten._to_copy.default($297, dtype=torch.bfloat16)
$305: bf16[20, 32, 512] = aten.native_dropout_backward.default($304, $186, 1.1111111111111112)
$306: bf16[640, 512] = aten.view.default($305, [640, 512])
$307: bf16[512, 512] = aten.t.default($182)
$308: f16[512, 512] = aten._to_copy.default($307, dtype=torch.float16)
$309: f16[640, 512] = aten._to_copy.default($306, dtype=torch.float16)
$310: f16[640, 512] = aten.mm.default($309, $308)
$311: bf16[512, 640] = aten.t.default($306)
$312: f16[640, 512] = aten._to_copy.default($177, dtype=torch.float16)
$313: f16[512, 640] = aten._to_copy.default($311, dtype=torch.float16)
$314: f16[512, 512] = aten.mm.default($313, $312)
$315: f16[512, 512] = aten.t.default($314)
$316: f32[1, 512] = aten.sum.dim_IntList($306, [0], True, dtype=torch.float32)
$317: f32[512] = aten.view.default($316, [512])
$318: bf16[512] = aten._to_copy.default($317, dtype=torch.bfloat16)
$319: bf16[640, 512] = aten._to_copy.default($310, dtype=torch.bfloat16)
$320: bf16[512, 512] = aten._to_copy.default($315, dtype=torch.bfloat16)
$321: bf16[512, 512] = aten.t.default($320)
$322: f32[512, 512] = aten._to_copy.default($321, dtype=torch.float32, layout=torch.strided, device=cuda:0)
$323: f32[512, 512] = aten.detach.default($322)
$324: f32[512, 512] = aten.detach.default($323)
$325: f32[512] = aten._to_copy.default($318, dtype=torch.float32, layout=torch.strided, device=cuda:0)
$326: f32[512] = aten.detach.default($325)
$327: f32[512] = aten.detach.default($326)
$328: bf16[20, 32, 8, 64] = aten.view.default($319, [20, 32, 8, 64])
$329: bf16[32, 8, 20, 64] = aten.permute.default($328, [1, 2, 0, 3])
$330: bf16[32, 8, 20, 64] = aten.detach.default($174)
($331: bf16[32, 8, 20, 64], $332: bf16[32, 8, 10, 64], $333: bf16[32, 8, 10, 64]) = aten._scaled_dot_product_efficient_attention_backward.default($329, $167, $168, $169, $330, $171, $172, $173, 0.1)
$334: bf16[32, 8, 10, 64] = aten.clone.default($333, memory_format=torch.contiguous_format)
$335: bf16[256, 10, 64] = aten._unsafe_view.default($334, [256, 10, 64])
$336: bf16[32, 8, 10, 64] = aten.clone.default($332, memory_format=torch.contiguous_format)
$337: bf16[256, 10, 64] = aten._unsafe_view.default($336, [256, 10, 64])
$338: bf16[32, 8, 20, 64] = aten.clone.default($331, memory_format=torch.contiguous_format)
$339: bf16[256, 20, 64] = aten._unsafe_view.default($338, [256, 20, 64])
$340: bf16[10, 256, 64] = aten.transpose.int($335, 0, 1)
$341: bf16[10, 256, 64] = aten.clone.default($340, memory_format=torch.contiguous_format)
$342: bf16[10, 32, 512] = aten._unsafe_view.default($341, [10, 32, 512])
$343: bf16[10, 256, 64] = aten.transpose.int($337, 0, 1)
$344: bf16[10, 256, 64] = aten.clone.default($343, memory_format=torch.contiguous_format)
$345: bf16[10, 32, 512] = aten._unsafe_view.default($344, [10, 32, 512])
$346: bf16[20, 256, 64] = aten.transpose.int($339, 0, 1)
$347: bf16[20, 256, 64] = aten.clone.default($346, memory_format=torch.contiguous_format)
$348: bf16[20, 32, 512] = aten._unsafe_view.default($347, [20, 32, 512])
$349: bf16[2, 10, 32, 512] = aten.select_backward.default($342, [2, 10, 32, 512], 0, 1)
$350: bf16[2, 10, 32, 512] = aten.select_backward.default($345, [2, 10, 32, 512], 0, 0)
$351: bf16[2, 10, 32, 512] = aten.add.Tensor($349, $350)
$352: bf16[2, 10, 32, 1, 512] = aten.unsqueeze.default($351, 3)
$353: bf16[1, 10, 32, 2, 512] = aten.transpose.int($352, 0, -2)
$354: bf16[10, 32, 2, 512] = aten.squeeze.dim($353, 0)
$355: bf16[10, 32, 2, 512] = aten.clone.default($354, memory_format=torch.contiguous_format)
$356: bf16[10, 32, 1024] = aten._unsafe_view.default($355, [10, 32, 1024])
$357: bf16[320, 1024] = aten.view.default($356, [320, 1024])
$358: bf16[1024, 512] = aten.t.default($150)
$359: f16[1024, 512] = aten._to_copy.default($358, dtype=torch.float16)
$360: f16[320, 1024] = aten._to_copy.default($357, dtype=torch.float16)
$361: f16[320, 512] = aten.mm.default($360, $359)
$362: bf16[1024, 320] = aten.t.default($357)
$363: f16[320, 512] = aten._to_copy.default($151, dtype=torch.float16)
$364: f16[1024, 320] = aten._to_copy.default($362, dtype=torch.float16)
$365: f16[1024, 512] = aten.mm.default($364, $363)
$366: f16[512, 1024] = aten.t.default($365)
$367: f32[1, 1024] = aten.sum.dim_IntList($357, [0], True, dtype=torch.float32)
$368: f32[1024] = aten.view.default($367, [1024])
$369: bf16[1024] = aten._to_copy.default($368, dtype=torch.bfloat16)
$370: bf16[320, 512] = aten._to_copy.default($361, dtype=torch.bfloat16)
$371: bf16[512, 1024] = aten._to_copy.default($366, dtype=torch.bfloat16)
$372: bf16[10, 32, 512] = aten.view.default($370, [10, 32, 512])
$373: bf16[1024, 512] = aten.t.default($371)
$374: f32[10, 32, 512] = aten._to_copy.default($372, dtype=torch.float32, layout=torch.strided, device=cuda:0)
$375: f32[1024, 512] = aten._to_copy.default($373, dtype=torch.float32, layout=torch.strided, device=cuda:0)
$376: f32[1024] = aten._to_copy.default($369, dtype=torch.float32, layout=torch.strided, device=cuda:0)
$377: bf16[640, 512] = aten.view.default($348, [640, 512])
$378: bf16[512, 512] = aten.t.default($143)
$379: f16[512, 512] = aten._to_copy.default($378, dtype=torch.float16)
$380: f16[640, 512] = aten._to_copy.default($377, dtype=torch.float16)
$381: f16[640, 512] = aten.mm.default($380, $379)
$382: bf16[512, 640] = aten.t.default($377)
$383: f16[640, 512] = aten._to_copy.default($144, dtype=torch.float16)
$384: f16[512, 640] = aten._to_copy.default($382, dtype=torch.float16)
$385: f16[512, 512] = aten.mm.default($384, $383)
$386: f16[512, 512] = aten.t.default($385)
$387: f32[1, 512] = aten.sum.dim_IntList($377, [0], True, dtype=torch.float32)
$388: f32[512] = aten.view.default($387, [512])
$389: bf16[512] = aten._to_copy.default($388, dtype=torch.bfloat16)
$390: bf16[640, 512] = aten._to_copy.default($381, dtype=torch.bfloat16)
$391: bf16[512, 512] = aten._to_copy.default($386, dtype=torch.bfloat16)
$392: bf16[20, 32, 512] = aten.view.default($390, [20, 32, 512])
$393: bf16[512, 512] = aten.t.default($391)
$394: f32[20, 32, 512] = aten._to_copy.default($392, dtype=torch.float32, layout=torch.strided, device=cuda:0)
$395: f32[20, 32, 512] = aten.add.Tensor($297, $394)
$396: f32[512, 512] = aten._to_copy.default($393, dtype=torch.float32, layout=torch.strided, device=cuda:0)
$397: f32[512] = aten._to_copy.default($389, dtype=torch.float32, layout=torch.strided, device=cuda:0)
$398: f32[1536] = aten.cat.default([$397, $376])
$399: f32[1536] = aten.detach.default($398)
$400: f32[1536] = aten.detach.default($399)
$401: f32[1536, 512] = aten.cat.default([$396, $375])
$402: f32[1536, 512] = aten.detach.default($401)
$403: f32[1536, 512] = aten.detach.default($402)
($404: f32[20, 32, 512], $405: f32[512], $406: f32[512]) = aten.native_layer_norm_backward.default($395, $128, [512], $132, $133, $129, $130, [True, True, True])
$407: f32[512] = aten.detach.default($405)
$408: f32[512] = aten.detach.default($407)
$409: f32[512] = aten.detach.default($406)
$410: f32[512] = aten.detach.default($409)
$411: bf16[20, 32, 512] = aten._to_copy.default($404, dtype=torch.bfloat16)
$412: bf16[20, 32, 512] = aten.native_dropout_backward.default($411, $127, 1.1111111111111112)
$413: bf16[640, 512] = aten.view.default($412, [640, 512])
$414: bf16[512, 512] = aten.t.default($123)
$415: f16[512, 512] = aten._to_copy.default($414, dtype=torch.float16)
$416: f16[640, 512] = aten._to_copy.default($413, dtype=torch.float16)
$417: f16[640, 512] = aten.mm.default($416, $415)
$418: bf16[512, 640] = aten.t.default($413)
$419: f16[640, 512] = aten._to_copy.default($118, dtype=torch.float16)
$420: f16[512, 640] = aten._to_copy.default($418, dtype=torch.float16)
$421: f16[512, 512] = aten.mm.default($420, $419)
$422: f16[512, 512] = aten.t.default($421)
$423: f32[1, 512] = aten.sum.dim_IntList($413, [0], True, dtype=torch.float32)
$424: f32[512] = aten.view.default($423, [512])
$425: bf16[512] = aten._to_copy.default($424, dtype=torch.bfloat16)
$426: bf16[640, 512] = aten._to_copy.default($417, dtype=torch.bfloat16)
$427: bf16[512, 512] = aten._to_copy.default($422, dtype=torch.bfloat16)
$428: bf16[512, 512] = aten.t.default($427)
$429: f32[512, 512] = aten._to_copy.default($428, dtype=torch.float32, layout=torch.strided, device=cuda:0)
$430: f32[512, 512] = aten.detach.default($429)
$431: f32[512, 512] = aten.detach.default($430)
$432: f32[512] = aten._to_copy.default($425, dtype=torch.float32, layout=torch.strided, device=cuda:0)
$433: f32[512] = aten.detach.default($432)
$434: f32[512] = aten.detach.default($433)
$435: bf16[20, 32, 8, 64] = aten.view.default($426, [20, 32, 8, 64])
$436: bf16[32, 8, 20, 64] = aten.permute.default($435, [1, 2, 0, 3])
$437: bf16[32, 8, 20, 64] = aten.detach.default($115)
($438: bf16[32, 8, 20, 64], $439: bf16[32, 8, 20, 64], $440: bf16[32, 8, 20, 64]) = aten._scaled_dot_product_efficient_attention_backward.default($436, $108, $109, $110, $437, $112, $113, $114, 0.1)
$441: bf16[32, 8, 20, 64] = aten.clone.default($440, memory_format=torch.contiguous_format)
$442: bf16[256, 20, 64] = aten._unsafe_view.default($441, [256, 20, 64])
$443: bf16[32, 8, 20, 64] = aten.clone.default($439, memory_format=torch.contiguous_format)
$444: bf16[256, 20, 64] = aten._unsafe_view.default($443, [256, 20, 64])
$445: bf16[32, 8, 20, 64] = aten.clone.default($438, memory_format=torch.contiguous_format)
$446: bf16[256, 20, 64] = aten._unsafe_view.default($445, [256, 20, 64])
$447: bf16[20, 256, 64] = aten.transpose.int($442, 0, 1)
$448: bf16[20, 256, 64] = aten.clone.default($447, memory_format=torch.contiguous_format)
$449: bf16[20, 32, 512] = aten._unsafe_view.default($448, [20, 32, 512])
$450: bf16[20, 256, 64] = aten.transpose.int($444, 0, 1)
$451: bf16[20, 256, 64] = aten.clone.default($450, memory_format=torch.contiguous_format)
$452: bf16[20, 32, 512] = aten._unsafe_view.default($451, [20, 32, 512])
$453: bf16[20, 256, 64] = aten.transpose.int($446, 0, 1)
$454: bf16[20, 256, 64] = aten.clone.default($453, memory_format=torch.contiguous_format)
$455: bf16[20, 32, 512] = aten._unsafe_view.default($454, [20, 32, 512])
$456: bf16[3, 20, 32, 512] = aten.select_backward.default($449, [3, 20, 32, 512], 0, 2)
$457: bf16[3, 20, 32, 512] = aten.select_backward.default($452, [3, 20, 32, 512], 0, 1)
$458: bf16[3, 20, 32, 512] = aten.add.Tensor($456, $457)
$459: bf16[3, 20, 32, 512] = aten.select_backward.default($455, [3, 20, 32, 512], 0, 0)
$460: bf16[3, 20, 32, 512] = aten.add.Tensor($458, $459)
$461: bf16[3, 20, 32, 1, 512] = aten.unsqueeze.default($460, 3)
$462: bf16[1, 20, 32, 3, 512] = aten.transpose.int($461, 0, -2)
$463: bf16[20, 32, 3, 512] = aten.squeeze.dim($462, 0)
$464: bf16[20, 32, 3, 512] = aten.clone.default($463, memory_format=torch.contiguous_format)
$465: bf16[20, 32, 1536] = aten._unsafe_view.default($464, [20, 32, 1536])
$466: bf16[640, 1536] = aten.view.default($465, [640, 1536])
$467: bf16[1536, 640] = aten.t.default($466)
$468: f16[640, 512] = aten._to_copy.default($91, dtype=torch.float16)
$469: f16[1536, 640] = aten._to_copy.default($467, dtype=torch.float16)
$470: f16[1536, 512] = aten.mm.default($469, $468)
$471: f16[512, 1536] = aten.t.default($470)
$472: f32[1, 1536] = aten.sum.dim_IntList($466, [0], True, dtype=torch.float32)
$473: f32[1536] = aten.view.default($472, [1536])
$474: bf16[1536] = aten._to_copy.default($473, dtype=torch.bfloat16)
$475: bf16[512, 1536] = aten._to_copy.default($471, dtype=torch.bfloat16)
$476: bf16[1536, 512] = aten.t.default($475)
$477: f32[1536, 512] = aten._to_copy.default($476, dtype=torch.float32, layout=torch.strided, device=cuda:0)
$478: f32[1536, 512] = aten.detach.default($477)
$479: f32[1536, 512] = aten.detach.default($478)
$480: f32[1536] = aten._to_copy.default($474, dtype=torch.float32, layout=torch.strided, device=cuda:0)
$481: f32[1536] = aten.detach.default($480)
$482: f32[1536] = aten.detach.default($481)
($483: f32[10, 32, 512], $484: f32[512], $485: f32[512]) = aten.native_layer_norm_backward.default($374, $76, [512], $82, $83, $79, $80, [True, True, True])
$486: f32[512] = aten.detach.default($484)
$487: f32[512] = aten.detach.default($486)
$488: f32[512] = aten.detach.default($485)
$489: f32[512] = aten.detach.default($488)
($490: f32[10, 32, 512], $491: f32[512], $492: f32[512]) = aten.native_layer_norm_backward.default($483, $73, [512], $77, $78, $74, $75, [True, True, True])
$493: f32[512] = aten.detach.default($491)
$494: f32[512] = aten.detach.default($493)
$495: f32[512] = aten.detach.default($492)
$496: f32[512] = aten.detach.default($495)
$497: bf16[10, 32, 512] = aten._to_copy.default($490, dtype=torch.bfloat16)
$498: bf16[10, 32, 512] = aten.native_dropout_backward.default($497, $72, 1.1111111111111112)
$499: bf16[320, 512] = aten.view.default($498, [320, 512])
$500: bf16[512, 2048] = aten.t.default($67)
$501: f16[512, 2048] = aten._to_copy.default($500, dtype=torch.float16)
$502: f16[320, 512] = aten._to_copy.default($499, dtype=torch.float16)
$503: f16[320, 2048] = aten.mm.default($502, $501)
$504: bf16[512, 320] = aten.t.default($499)
$505: f16[320, 2048] = aten._to_copy.default($68, dtype=torch.float16)
$506: f16[512, 320] = aten._to_copy.default($504, dtype=torch.float16)
$507: f16[512, 2048] = aten.mm.default($506, $505)
$508: f16[2048, 512] = aten.t.default($507)
$509: f32[1, 512] = aten.sum.dim_IntList($499, [0], True, dtype=torch.float32)
$510: f32[512] = aten.view.default($509, [512])
$511: bf16[512] = aten._to_copy.default($510, dtype=torch.bfloat16)
$512: bf16[320, 2048] = aten._to_copy.default($503, dtype=torch.bfloat16)
$513: bf16[2048, 512] = aten._to_copy.default($508, dtype=torch.bfloat16)
$514: bf16[10, 32, 2048] = aten.view.default($512, [10, 32, 2048])
$515: bf16[512, 2048] = aten.t.default($513)
$516: f32[512, 2048] = aten._to_copy.default($515, dtype=torch.float32, layout=torch.strided, device=cuda:0)
$517: f32[512, 2048] = aten.detach.default($516)
$518: f32[512, 2048] = aten.detach.default($517)
$519: f32[512] = aten._to_copy.default($511, dtype=torch.float32, layout=torch.strided, device=cuda:0)
$520: f32[512] = aten.detach.default($519)
$521: f32[512] = aten.detach.default($520)
$522: bf16[10, 32, 2048] = aten.native_dropout_backward.default($514, $62, 1.1111111111111112)
$523: bf16[10, 32, 2048] = aten.detach.default($60)
$524: bf16[10, 32, 2048] = aten.threshold_backward.default($522, $523, 0)
$525: bf16[320, 2048] = aten.view.default($524, [320, 2048])
$526: bf16[2048, 512] = aten.t.default($55)
$527: f16[2048, 512] = aten._to_copy.default($526, dtype=torch.float16)
$528: f16[320, 2048] = aten._to_copy.default($525, dtype=torch.float16)
$529: f16[320, 512] = aten.mm.default($528, $527)
$530: bf16[2048, 320] = aten.t.default($525)
$531: f16[320, 512] = aten._to_copy.default($56, dtype=torch.float16)
$532: f16[2048, 320] = aten._to_copy.default($530, dtype=torch.float16)
$533: f16[2048, 512] = aten.mm.default($532, $531)
$534: f16[512, 2048] = aten.t.default($533)
$535: f32[1, 2048] = aten.sum.dim_IntList($525, [0], True, dtype=torch.float32)
$536: f32[2048] = aten.view.default($535, [2048])
$537: bf16[2048] = aten._to_copy.default($536, dtype=torch.bfloat16)
$538: bf16[320, 512] = aten._to_copy.default($529, dtype=torch.bfloat16)
$539: bf16[512, 2048] = aten._to_copy.default($534, dtype=torch.bfloat16)
$540: bf16[10, 32, 512] = aten.view.default($538, [10, 32, 512])
$541: bf16[2048, 512] = aten.t.default($539)
$542: f32[10, 32, 512] = aten._to_copy.default($540, dtype=torch.float32, layout=torch.strided, device=cuda:0)
$543: f32[10, 32, 512] = aten.add.Tensor($490, $542)
$544: f32[2048, 512] = aten._to_copy.default($541, dtype=torch.float32, layout=torch.strided, device=cuda:0)
$545: f32[2048, 512] = aten.detach.default($544)
$546: f32[2048, 512] = aten.detach.default($545)
$547: f32[2048] = aten._to_copy.default($537, dtype=torch.float32, layout=torch.strided, device=cuda:0)
$548: f32[2048] = aten.detach.default($547)
$549: f32[2048] = aten.detach.default($548)
($550: f32[10, 32, 512], $551: f32[512], $552: f32[512]) = aten.native_layer_norm_backward.default($543, $44, [512], $48, $49, $45, $46, [True, True, True])
$553: f32[512] = aten.detach.default($551)
$554: f32[512] = aten.detach.default($553)
$555: f32[512] = aten.detach.default($552)
$556: f32[512] = aten.detach.default($555)
$557: bf16[10, 32, 512] = aten._to_copy.default($550, dtype=torch.bfloat16)
$558: bf16[10, 32, 512] = aten.native_dropout_backward.default($557, $43, 1.1111111111111112)
$559: bf16[320, 512] = aten.view.default($558, [320, 512])
$560: bf16[512, 512] = aten.t.default($39)
$561: f16[512, 512] = aten._to_copy.default($560, dtype=torch.float16)
$562: f16[320, 512] = aten._to_copy.default($559, dtype=torch.float16)
$563: f16[320, 512] = aten.mm.default($562, $561)
$564: bf16[512, 320] = aten.t.default($559)
$565: f16[320, 512] = aten._to_copy.default($34, dtype=torch.float16)
$566: f16[512, 320] = aten._to_copy.default($564, dtype=torch.float16)
$567: f16[512, 512] = aten.mm.default($566, $565)
$568: f16[512, 512] = aten.t.default($567)
$569: f32[1, 512] = aten.sum.dim_IntList($559, [0], True, dtype=torch.float32)
$570: f32[512] = aten.view.default($569, [512])
$571: bf16[512] = aten._to_copy.default($570, dtype=torch.bfloat16)
$572: bf16[320, 512] = aten._to_copy.default($563, dtype=torch.bfloat16)
$573: bf16[512, 512] = aten._to_copy.default($568, dtype=torch.bfloat16)
$574: bf16[512, 512] = aten.t.default($573)
$575: f32[512, 512] = aten._to_copy.default($574, dtype=torch.float32, layout=torch.strided, device=cuda:0)
$576: f32[512, 512] = aten.detach.default($575)
$577: f32[512, 512] = aten.detach.default($576)
$578: f32[512] = aten._to_copy.default($571, dtype=torch.float32, layout=torch.strided, device=cuda:0)
$579: f32[512] = aten.detach.default($578)
$580: f32[512] = aten.detach.default($579)
$581: bf16[10, 32, 8, 64] = aten.view.default($572, [10, 32, 8, 64])
$582: bf16[32, 8, 10, 64] = aten.permute.default($581, [1, 2, 0, 3])
$583: bf16[32, 8, 10, 64] = aten.detach.default($31)
($584: bf16[32, 8, 10, 64], $585: bf16[32, 8, 10, 64], $586: bf16[32, 8, 10, 64]) = aten._scaled_dot_product_efficient_attention_backward.default($582, $24, $25, $26, $583, $28, $29, $30, 0.1)
$587: bf16[32, 8, 10, 64] = aten.clone.default($586, memory_format=torch.contiguous_format)
$588: bf16[256, 10, 64] = aten._unsafe_view.default($587, [256, 10, 64])
$589: bf16[32, 8, 10, 64] = aten.clone.default($585, memory_format=torch.contiguous_format)
$590: bf16[256, 10, 64] = aten._unsafe_view.default($589, [256, 10, 64])
$591: bf16[32, 8, 10, 64] = aten.clone.default($584, memory_format=torch.contiguous_format)
$592: bf16[256, 10, 64] = aten._unsafe_view.default($591, [256, 10, 64])
$593: bf16[10, 256, 64] = aten.transpose.int($588, 0, 1)
$594: bf16[10, 256, 64] = aten.clone.default($593, memory_format=torch.contiguous_format)
$595: bf16[10, 32, 512] = aten._unsafe_view.default($594, [10, 32, 512])
$596: bf16[10, 256, 64] = aten.transpose.int($590, 0, 1)
$597: bf16[10, 256, 64] = aten.clone.default($596, memory_format=torch.contiguous_format)
$598: bf16[10, 32, 512] = aten._unsafe_view.default($597, [10, 32, 512])
$599: bf16[10, 256, 64] = aten.transpose.int($592, 0, 1)
$600: bf16[10, 256, 64] = aten.clone.default($599, memory_format=torch.contiguous_format)
$601: bf16[10, 32, 512] = aten._unsafe_view.default($600, [10, 32, 512])
$602: bf16[3, 10, 32, 512] = aten.select_backward.default($595, [3, 10, 32, 512], 0, 2)
$603: bf16[3, 10, 32, 512] = aten.select_backward.default($598, [3, 10, 32, 512], 0, 1)
$604: bf16[3, 10, 32, 512] = aten.add.Tensor($602, $603)
$605: bf16[3, 10, 32, 512] = aten.select_backward.default($601, [3, 10, 32, 512], 0, 0)
$606: bf16[3, 10, 32, 512] = aten.add.Tensor($604, $605)
$607: bf16[3, 10, 32, 1, 512] = aten.unsqueeze.default($606, 3)
$608: bf16[1, 10, 32, 3, 512] = aten.transpose.int($607, 0, -2)
$609: bf16[10, 32, 3, 512] = aten.squeeze.dim($608, 0)
$610: bf16[10, 32, 3, 512] = aten.clone.default($609, memory_format=torch.contiguous_format)
$611: bf16[10, 32, 1536] = aten._unsafe_view.default($610, [10, 32, 1536])
$612: bf16[320, 1536] = aten.view.default($611, [320, 1536])
$613: bf16[1536, 320] = aten.t.default($612)
$614: f16[320, 512] = aten._to_copy.default($7, dtype=torch.float16)
$615: f16[1536, 320] = aten._to_copy.default($613, dtype=torch.float16)
$616: f16[1536, 512] = aten.mm.default($615, $614)
$617: f16[512, 1536] = aten.t.default($616)
$618: f32[1, 1536] = aten.sum.dim_IntList($612, [0], True, dtype=torch.float32)
$619: f32[1536] = aten.view.default($618, [1536])
$620: bf16[1536] = aten._to_copy.default($619, dtype=torch.bfloat16)
$621: bf16[512, 1536] = aten._to_copy.default($617, dtype=torch.bfloat16)
$622: bf16[1536, 512] = aten.t.default($621)
$623: f32[1536, 512] = aten._to_copy.default($622, dtype=torch.float32, layout=torch.strided, device=cuda:0)
$624: f32[1536, 512] = aten.detach.default($623)
$625: f32[1536, 512] = aten.detach.default($624)
$626: f32[1536] = aten._to_copy.default($620, dtype=torch.float32, layout=torch.strided, device=cuda:0)
$627: f32[1536] = aten.detach.default($626)
$628: f32[1536] = aten.detach.default($627)
```


</details>

Version: main




cc @mcarilli @ptrblck @leslie-fang-intel @jgong5