Repository: pytorch
Issue: Partitioner error when functionalizing RNG state for LLAMA2 + LoRA 70b decoder layer graph
Labels: triaged,oncall: pt2,module: aotdispatch,module: pt2-dispatcher

### 🐛 Describe the bug

When doing pretraining of LLAMA2 70b model with LoRA with torch.compile in our environment, we saw an exception raised from partitioner.py while splitting the Joint graph to forward and backward graph.

```
[rank3]:   File "/usr/local/lib/python3.10/dist-packages/torch/_dynamo/output_graph.py", line 1390, in call_user_compiler
[rank3]:     compiled_fn = compiler_fn(gm, self.example_inputs())
[rank3]:   File "/usr/local/lib/python3.10/dist-packages/torch/_dynamo/repro/after_dynamo.py", line 129, in __call__
[rank3]:     compiled_gm = compiler_fn(gm, example_inputs)
[rank3]:   File "/usr/local/lib/python3.10/dist-packages/torch/__init__.py", line 1990, in __call__
[rank3]:     return self.compiler_fn(model_, inputs_, **self.kwargs)
[rank3]:   File "/usr/local/lib/python3.10/dist-packages/habana_frameworks/torch/dynamo/compile_backend/backends.py", line 46, in hpu_backend
[rank3]:     return aot_autograd(
[rank3]:   File "/usr/local/lib/python3.10/dist-packages/torch/_dynamo/backends/common.py", line 69, in __call__
[rank3]:     cg = aot_module_simplified(gm, example_inputs, **self.kwargs)
[rank3]:   File "/usr/local/lib/python3.10/dist-packages/torch/_functorch/aot_autograd.py", line 954, in aot_module_simplified
[rank3]:     compiled_fn, _ = create_aot_dispatcher_function(
[rank3]:   File "/usr/local/lib/python3.10/dist-packages/torch/_dynamo/utils.py", line 231, in time_wrapper
[rank3]:     r = func(*args, **kwargs)
[rank3]:   File "/usr/local/lib/python3.10/dist-packages/torch/_functorch/aot_autograd.py", line 687, in create_aot_dispatcher_function
[rank3]:     compiled_fn, fw_metadata = compiler_fn(
[rank3]:   File "/usr/local/lib/python3.10/dist-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py", line 290, in aot_dispatch_autograd
[rank3]:     fw_module, bw_module = aot_config.partition_fn(
[rank3]:   File "/usr/local/lib/python3.10/dist-packages/habana_frameworks/torch/dynamo/compile_backend/partition_fn.py", line 100, in hpu_partition
[rank3]:     return default_partition(joint_module, _joint_inputs, num_fwd_outputs=num_fwd_outputs)
[rank3]:   File "/usr/local/lib/python3.10/dist-packages/torch/_functorch/partitioners.py", line 380, in default_partition
[rank3]:     return min_cut_rematerialization_partition(
[rank3]:   File "/usr/local/lib/python3.10/dist-packages/torch/_functorch/partitioners.py", line 1774, in min_cut_rematerialization_partition
[rank3]:     fw_module, bw_module = functionalize_rng_ops(
[rank3]:   File "/usr/local/lib/python3.10/dist-packages/torch/_functorch/partitioners.py", line 653, in functionalize_rng_ops
[rank3]:     bw_node = bw_graph_rng_ops[node.name]
[rank3]: torch._dynamo.exc.BackendCompilerFailed: backend='hpu_backend' raised:
[rank3]: KeyError: 'native_dropout_3'
```
Because the layer have dropout ops in the forward, and we have activation checkpointing enable on the decoder layer, the AOT needs to do functionalization of the RNG related ops. The issue happens in functionalize_rng_ops.
 
Because the environment of running is complex and has many component, I dumped the graph modules (Joint graph and splitted forward module and backward module) right before functionalize_rng_ops. (These are not final forward module and backward module returning to the user)

From these graphs, we know that the backward graph has no native_dropout_3 op which exists on forward. This breaks the current code assumption. (Maybe the graph split is right and this is the right case to consider)

```
joint_module joint_helper()

def forward(self, primals_1, primals_2, primals_3, primals_4, primals_5, primals_6, primals_7, primals_8, primals_9, primals_10, primals_11, primals_12, primals_13, primals_14, primals_15, primals_16, primals_17, primals_18, primals_19, primals_20, primals_21, primals_22, primals_23, primals_24, tangents_1):
    rms_norm = torch.ops.hpu.rms_norm.default(primals_22, primals_1, 1e-05)
    getitem = rms_norm[0]
    getitem_1 = rms_norm[1];  rms_norm = None
    transpose = torch.ops.aten.transpose.int(primals_2, 0, 1);  primals_2 = None
    mul = primals_20 * primals_21
    view = torch.ops.aten.view.default(getitem, [mul, 8192])
    mm = torch.ops.aten.mm.default(view, transpose);  transpose = None
    _unsafe_view = torch.ops.aten._unsafe_view.default(mm, [primals_20, primals_21, 8192]);  mm = None
    native_dropout = torch.ops.aten.native_dropout.default(getitem, 0.05, True)
    getitem_2 = native_dropout[0];  native_dropout = None
    transpose_1 = torch.ops.aten.transpose.int(primals_3, 0, 1);  primals_3 = None
    view_1 = torch.ops.aten.view.default(getitem_2, [mul, 8192])
    mm_1 = torch.ops.aten.mm.default(view_1, transpose_1);  view_1 = transpose_1 = None
    _unsafe_view_1 = torch.ops.aten._unsafe_view.default(mm_1, [primals_20, primals_21, 4]);  mm_1 = None
    transpose_2 = torch.ops.aten.transpose.int(primals_4, 0, 1);  primals_4 = None
    view_2 = torch.ops.aten.view.default(_unsafe_view_1, [mul, 4])
    mm_2 = torch.ops.aten.mm.default(view_2, transpose_2);  view_2 = transpose_2 = None
    _unsafe_view_2 = torch.ops.aten._unsafe_view.default(mm_2, [primals_20, primals_21, 8192]);  mm_2 = None
    mul_3 = torch.ops.aten.mul.Tensor(_unsafe_view_2, 4.0);  _unsafe_view_2 = None
    add = torch.ops.aten.add.Tensor(_unsafe_view, mul_3);  _unsafe_view = mul_3 = None
    transpose_3 = torch.ops.aten.transpose.int(primals_5, 0, 1);  primals_5 = None
    mm_3 = torch.ops.aten.mm.default(view, transpose_3);  transpose_3 = None
    _unsafe_view_3 = torch.ops.aten._unsafe_view.default(mm_3, [primals_20, primals_21, 1024]);  mm_3 = None
    native_dropout_1 = torch.ops.aten.native_dropout.default(getitem, 0.05, True)
    getitem_4 = native_dropout_1[0];  native_dropout_1 = None
    transpose_4 = torch.ops.aten.transpose.int(primals_6, 0, 1);  primals_6 = None
    view_4 = torch.ops.aten.view.default(getitem_4, [mul, 8192])
    mm_4 = torch.ops.aten.mm.default(view_4, transpose_4);  view_4 = transpose_4 = None
    _unsafe_view_4 = torch.ops.aten._unsafe_view.default(mm_4, [primals_20, primals_21, 4]);  mm_4 = None
    transpose_5 = torch.ops.aten.transpose.int(primals_7, 0, 1);  primals_7 = None
    view_5 = torch.ops.aten.view.default(_unsafe_view_4, [mul, 4])
    mm_5 = torch.ops.aten.mm.default(view_5, transpose_5);  view_5 = transpose_5 = None
    _unsafe_view_5 = torch.ops.aten._unsafe_view.default(mm_5, [primals_20, primals_21, 1024]);  mm_5 = None
    mul_7 = torch.ops.aten.mul.Tensor(_unsafe_view_5, 4.0);  _unsafe_view_5 = None
    add_1 = torch.ops.aten.add.Tensor(_unsafe_view_3, mul_7);  _unsafe_view_3 = mul_7 = None
    transpose_6 = torch.ops.aten.transpose.int(primals_8, 0, 1);  primals_8 = None
    mm_6 = torch.ops.aten.mm.default(view, transpose_6);  view = transpose_6 = None
    _unsafe_view_6 = torch.ops.aten._unsafe_view.default(mm_6, [primals_20, primals_21, 1024]);  mm_6 = None
    native_dropout_2 = torch.ops.aten.native_dropout.default(getitem, 0.05, True);  getitem = None
    getitem_6 = native_dropout_2[0];  native_dropout_2 = None
    transpose_7 = torch.ops.aten.transpose.int(primals_9, 0, 1);  primals_9 = None
    view_7 = torch.ops.aten.view.default(getitem_6, [mul, 8192])
    mm_7 = torch.ops.aten.mm.default(view_7, transpose_7);  view_7 = transpose_7 = None
    _unsafe_view_7 = torch.ops.aten._unsafe_view.default(mm_7, [primals_20, primals_21, 4]);  mm_7 = None
    transpose_8 = torch.ops.aten.transpose.int(primals_10, 0, 1);  primals_10 = None
    view_8 = torch.ops.aten.view.default(_unsafe_view_7, [mul, 4])
    mm_8 = torch.ops.aten.mm.default(view_8, transpose_8);  view_8 = transpose_8 = None
    _unsafe_view_8 = torch.ops.aten._unsafe_view.default(mm_8, [primals_20, primals_21, 1024]);  mm_8 = None
    mul_11 = torch.ops.aten.mul.Tensor(_unsafe_view_8, 4.0);  _unsafe_view_8 = None
    add_2 = torch.ops.aten.add.Tensor(_unsafe_view_6, mul_11);  _unsafe_view_6 = mul_11 = None
    view_9 = torch.ops.aten.view.default(add, [primals_20, primals_21, 64, 128]);  add = None
    transpose_9 = torch.ops.aten.transpose.int(view_9, 1, 2);  view_9 = None
    view_10 = torch.ops.aten.view.default(add_1, [primals_20, primals_21, -1, 128]);  add_1 = None
    transpose_10 = torch.ops.aten.transpose.int(view_10, 1, 2);  view_10 = None
    view_11 = torch.ops.aten.view.default(add_2, [primals_20, primals_21, -1, 128]);  add_2 = None
    transpose_11 = torch.ops.aten.transpose.int(view_11, 1, 2);  view_11 = None
    slice_1 = torch.ops.aten.slice.Tensor(primals_18, 0, 0, primals_21);  primals_18 = None
    slice_2 = torch.ops.aten.slice.Tensor(primals_19, 0, 0, primals_21);  primals_19 = None
    unsqueeze = torch.ops.aten.unsqueeze.default(slice_1, 0);  slice_1 = None
    unsqueeze_1 = torch.ops.aten.unsqueeze.default(unsqueeze, 0);  unsqueeze = None
    unsqueeze_2 = torch.ops.aten.unsqueeze.default(slice_2, 0);  slice_2 = None
    unsqueeze_3 = torch.ops.aten.unsqueeze.default(unsqueeze_2, 0);  unsqueeze_2 = None
    rotary_pos_embedding = torch.ops.hpu.rotary_pos_embedding.default(transpose_9, unsqueeze_3, unsqueeze_1, primals_24, 0, 0);  transpose_9 = None
    rotary_pos_embedding_1 = torch.ops.hpu.rotary_pos_embedding.default(transpose_10, unsqueeze_3, unsqueeze_1, primals_24, 0, 0);  transpose_10 = unsqueeze_3 = unsqueeze_1 = primals_24 = None
    view_12 = torch.ops.aten.view.default(rotary_pos_embedding, [primals_20, 8, 8, primals_21, 128]);  rotary_pos_embedding = None
    view_13 = torch.ops.aten.view.default(rotary_pos_embedding_1, [primals_20, 8, 1, primals_21, 128]);  rotary_pos_embedding_1 = None
    view_14 = torch.ops.aten.view.default(transpose_11, [primals_20, 8, 1, primals_21, 128]);  transpose_11 = None
    sdpa_fwd_non_dropout = torch.ops.hpu.sdpa_fwd_non_dropout.default(view_12, view_13, view_14, None, 0.0, 0.08838834764831843, True, 'none', None, 'left');  view_12 = view_13 = view_14 = None
    getitem_8 = sdpa_fwd_non_dropout[0];  sdpa_fwd_non_dropout = None
    view_15 = torch.ops.aten.view.default(getitem_8, [primals_20, 64, primals_21, 128]);  getitem_8 = None
    transpose_12 = torch.ops.aten.transpose.int(view_15, 1, 2);  view_15 = None
    clone_4 = torch.ops.aten.clone.default(transpose_12, memory_format = torch.contiguous_format);  transpose_12 = None
    view_16 = torch.ops.aten.view.default(clone_4, [primals_20, primals_21, -1]);  clone_4 = None
    transpose_13 = torch.ops.aten.transpose.int(primals_11, 0, 1);  primals_11 = None
    view_17 = torch.ops.aten.view.default(view_16, [mul, 8192])
    mm_9 = torch.ops.aten.mm.default(view_17, transpose_13);  view_17 = transpose_13 = None
    _unsafe_view_9 = torch.ops.aten._unsafe_view.default(mm_9, [primals_20, primals_21, 8192]);  mm_9 = None
    native_dropout_3 = torch.ops.aten.native_dropout.default(view_16, 0.05, True);  view_16 = None
    getitem_11 = native_dropout_3[0];  native_dropout_3 = None
    transpose_14 = torch.ops.aten.transpose.int(primals_12, 0, 1);  primals_12 = None
    view_18 = torch.ops.aten.view.default(getitem_11, [mul, 8192])
    mm_10 = torch.ops.aten.mm.default(view_18, transpose_14);  view_18 = transpose_14 = None
    _unsafe_view_10 = torch.ops.aten._unsafe_view.default(mm_10, [primals_20, primals_21, 4]);  mm_10 = None
    transpose_15 = torch.ops.aten.transpose.int(primals_13, 0, 1);  primals_13 = None
    view_19 = torch.ops.aten.view.default(_unsafe_view_10, [mul, 4])
    mm_11 = torch.ops.aten.mm.default(view_19, transpose_15);  view_19 = transpose_15 = None
    _unsafe_view_11 = torch.ops.aten._unsafe_view.default(mm_11, [primals_20, primals_21, 8192]);  mm_11 = None
    mul_15 = torch.ops.aten.mul.Tensor(_unsafe_view_11, 4.0);  _unsafe_view_11 = None
    add_3 = torch.ops.aten.add.Tensor(_unsafe_view_9, mul_15);  _unsafe_view_9 = mul_15 = None
    add_4 = torch.ops.aten.add.Tensor(add_3, primals_22);  add_3 = None
    rms_norm_1 = torch.ops.hpu.rms_norm.default(add_4, primals_14, 1e-05)
    getitem_13 = rms_norm_1[0]
    getitem_14 = rms_norm_1[1];  rms_norm_1 = None
    transpose_16 = torch.ops.aten.transpose.int(primals_15, 0, 1);  primals_15 = None
    view_20 = torch.ops.aten.view.default(getitem_13, [mul, 8192]);  getitem_13 = None
    mm_12 = torch.ops.aten.mm.default(view_20, transpose_16)
    _unsafe_view_12 = torch.ops.aten._unsafe_view.default(mm_12, [primals_20, primals_21, 28672]);  mm_12 = None
    silu = torch.ops.aten.silu.default(_unsafe_view_12)
    transpose_17 = torch.ops.aten.transpose.int(primals_16, 0, 1);  primals_16 = None
    mm_13 = torch.ops.aten.mm.default(view_20, transpose_17);  view_20 = None
    _unsafe_view_13 = torch.ops.aten._unsafe_view.default(mm_13, [primals_20, primals_21, 28672]);  mm_13 = None
    mul_18 = torch.ops.aten.mul.Tensor(silu, _unsafe_view_13)
    transpose_18 = torch.ops.aten.transpose.int(primals_17, 0, 1);  primals_17 = None
    view_22 = torch.ops.aten.view.default(mul_18, [mul, 28672]);  mul_18 = None
    mm_14 = torch.ops.aten.mm.default(view_22, transpose_18);  view_22 = None
    _unsafe_view_14 = torch.ops.aten._unsafe_view.default(mm_14, [primals_20, primals_21, 8192]);  mm_14 = None
    add_5 = torch.ops.aten.add.Tensor(_unsafe_view_14, add_4);  _unsafe_view_14 = None
    view_23 = torch.ops.aten.view.default(tangents_1, [mul, 8192])
    transpose_19 = torch.ops.aten.transpose.int(transpose_18, 0, 1);  transpose_18 = None
    mm_15 = torch.ops.aten.mm.default(view_23, transpose_19);  view_23 = transpose_19 = None
    view_24 = torch.ops.aten.view.default(mm_15, [primals_20, primals_21, 28672]);  mm_15 = None
    mul_20 = torch.ops.aten.mul.Tensor(view_24, silu);  silu = None
    mul_21 = torch.ops.aten.mul.Tensor(view_24, _unsafe_view_13);  view_24 = _unsafe_view_13 = None
    view_25 = torch.ops.aten.view.default(mul_20, [mul, 28672]);  mul_20 = None
    transpose_20 = torch.ops.aten.transpose.int(transpose_17, 0, 1);  transpose_17 = None
    mm_16 = torch.ops.aten.mm.default(view_25, transpose_20);  view_25 = transpose_20 = None
    view_26 = torch.ops.aten.view.default(mm_16, [primals_20, primals_21, 8192]);  mm_16 = None
    sigmoid = torch.ops.aten.sigmoid.default(_unsafe_view_12)
    full = torch.ops.aten.full.default([primals_20, primals_21, 28672], 1, dtype = torch.bfloat16, layout = torch.strided, device = device(type='hpu', index=0), pin_memory = False)
    sub = torch.ops.aten.sub.Tensor(full, sigmoid);  full = None
    mul_22 = torch.ops.aten.mul.Tensor(_unsafe_view_12, sub);  _unsafe_view_12 = sub = None
    add_6 = torch.ops.aten.add.Scalar(mul_22, 1);  mul_22 = None
    mul_23 = torch.ops.aten.mul.Tensor(sigmoid, add_6);  sigmoid = add_6 = None
    mul_24 = torch.ops.aten.mul.Tensor(mul_21, mul_23);  mul_21 = mul_23 = None
    view_27 = torch.ops.aten.view.default(mul_24, [mul, 28672]);  mul_24 = mul = None
    transpose_21 = torch.ops.aten.transpose.int(transpose_16, 0, 1);  transpose_16 = None
    mm_17 = torch.ops.aten.mm.default(view_27, transpose_21);  view_27 = transpose_21 = None
    view_28 = torch.ops.aten.view.default(mm_17, [primals_20, primals_21, 8192]);  mm_17 = None
    add_7 = torch.ops.aten.add.Tensor(view_26, view_28);  view_26 = view_28 = None
    rms_norm_backward = torch.ops.hpu.rms_norm_backward.default(add_7, add_4, primals_14, getitem_14, True, 0);  add_7 = add_4 = primals_14 = getitem_14 = None
    getitem_15 = rms_norm_backward[0];  rms_norm_backward = None
    add_8 = torch.ops.aten.add.Tensor(tangents_1, getitem_15);  tangents_1 = getitem_15 = None
    mul_25 = torch.ops.aten.mul.Tensor(add_8, 4.0)
    view_29 = torch.ops.aten.view.default(mul_25, [-1, 8192]);  mul_25 = None
    transpose_22 = torch.ops.aten.transpose.int(view_29, 0, 1);  view_29 = None
    view_30 = torch.ops.aten.view.default(_unsafe_view_10, [-1, 4]);  _unsafe_view_10 = None
    mm_18 = torch.ops.aten.mm.default(transpose_22, view_30);  transpose_22 = view_30 = None
    full_1 = torch.ops.aten.full.default([primals_20, primals_21, 4], 0, dtype = torch.bfloat16, layout = torch.strided, device = device(type='hpu', index=0), pin_memory = False)
    view_31 = torch.ops.aten.view.default(full_1, [-1, 4]);  full_1 = None
    transpose_23 = torch.ops.aten.transpose.int(view_31, 0, 1);  view_31 = None
    view_32 = torch.ops.aten.view.default(getitem_11, [-1, 8192]);  getitem_11 = None
    mm_19 = torch.ops.aten.mm.default(transpose_23, view_32);  view_32 = None
    full_2 = torch.ops.aten.full.default([primals_20, primals_21, 1024], 0, dtype = torch.bfloat16, layout = torch.strided, device = device(type='hpu', index=0), pin_memory = False)
    view_37 = torch.ops.aten.view.default(full_2, [-1, 1024]);  full_2 = None
    transpose_26 = torch.ops.aten.transpose.int(view_37, 0, 1);  view_37 = None
    view_38 = torch.ops.aten.view.default(_unsafe_view_7, [-1, 4]);  _unsafe_view_7 = None
    mm_21 = torch.ops.aten.mm.default(transpose_26, view_38);  view_38 = None
    view_40 = torch.ops.aten.view.default(getitem_6, [-1, 8192]);  getitem_6 = None
    mm_22 = torch.ops.aten.mm.default(transpose_23, view_40);  view_40 = None
    view_42 = torch.ops.aten.view.default(_unsafe_view_4, [-1, 4]);  _unsafe_view_4 = None
    mm_23 = torch.ops.aten.mm.default(transpose_26, view_42);  transpose_26 = view_42 = None
    view_44 = torch.ops.aten.view.default(getitem_4, [-1, 8192]);  getitem_4 = None
    mm_24 = torch.ops.aten.mm.default(transpose_23, view_44);  view_44 = None
    full_6 = torch.ops.aten.full.default([primals_20, primals_21, 8192], 0, dtype = torch.bfloat16, layout = torch.strided, device = device(type='hpu', index=0), pin_memory = False);  primals_20 = primals_21 = None
    view_45 = torch.ops.aten.view.default(full_6, [-1, 8192])
    transpose_30 = torch.ops.aten.transpose.int(view_45, 0, 1);  view_45 = None
    view_46 = torch.ops.aten.view.default(_unsafe_view_1, [-1, 4]);  _unsafe_view_1 = None
    mm_25 = torch.ops.aten.mm.default(transpose_30, view_46);  transpose_30 = view_46 = None
    view_48 = torch.ops.aten.view.default(getitem_2, [-1, 8192]);  getitem_2 = None
    mm_26 = torch.ops.aten.mm.default(transpose_23, view_48);  transpose_23 = view_48 = None
    rms_norm_backward_1 = torch.ops.hpu.rms_norm_backward.default(full_6, primals_22, primals_1, getitem_1, True, 0);  full_6 = primals_22 = primals_1 = getitem_1 = None
    getitem_17 = rms_norm_backward_1[0];  rms_norm_backward_1 = None
    add_9 = torch.ops.aten.add.Tensor(add_8, getitem_17);  add_8 = getitem_17 = None
    return [add_5, None, None, mm_26, mm_25, None, mm_24, mm_23, None, mm_22, mm_21, None, mm_19, mm_18, None, None, None, None, None, None, None, None, add_9, None, None]
    
# To see more debug info, please use `graph_module.print_readable()`
fw_module GraphModule()

def forward(self, primals_1, primals_2, primals_3, primals_4, primals_5, primals_6, primals_7, primals_8, primals_9, primals_10, primals_11, primals_12, primals_13, primals_14, primals_15, primals_16, primals_17, primals_18, primals_19, primals_20, primals_21, primals_22, primals_23, primals_24):
    rms_norm = torch.ops.hpu.rms_norm.default(primals_22, primals_1, 1e-05)
    getitem = rms_norm[0];  rms_norm = None
    transpose = torch.ops.aten.transpose.int(primals_2, 0, 1);  primals_2 = None
    mul = primals_20 * primals_21
    view = torch.ops.aten.view.default(getitem, [mul, 8192])
    mm = torch.ops.aten.mm.default(view, transpose);  transpose = None
    _unsafe_view = torch.ops.aten._unsafe_view.default(mm, [primals_20, primals_21, 8192]);  mm = None
    native_dropout = torch.ops.aten.native_dropout.default(getitem, 0.05, True)
    getitem_2 = native_dropout[0];  native_dropout = None
    transpose_1 = torch.ops.aten.transpose.int(primals_3, 0, 1);  primals_3 = None
    view_1 = torch.ops.aten.view.default(getitem_2, [mul, 8192]);  getitem_2 = None
    mm_1 = torch.ops.aten.mm.default(view_1, transpose_1);  view_1 = None
    _unsafe_view_1 = torch.ops.aten._unsafe_view.default(mm_1, [primals_20, primals_21, 4]);  mm_1 = None
    transpose_2 = torch.ops.aten.transpose.int(primals_4, 0, 1);  primals_4 = None
    view_2 = torch.ops.aten.view.default(_unsafe_view_1, [mul, 4]);  _unsafe_view_1 = None
    mm_2 = torch.ops.aten.mm.default(view_2, transpose_2);  view_2 = transpose_2 = None
    _unsafe_view_2 = torch.ops.aten._unsafe_view.default(mm_2, [primals_20, primals_21, 8192]);  mm_2 = None
    mul_3 = torch.ops.aten.mul.Tensor(_unsafe_view_2, 4.0);  _unsafe_view_2 = None
    add = torch.ops.aten.add.Tensor(_unsafe_view, mul_3);  _unsafe_view = mul_3 = None
    transpose_3 = torch.ops.aten.transpose.int(primals_5, 0, 1);  primals_5 = None
    mm_3 = torch.ops.aten.mm.default(view, transpose_3);  transpose_3 = None
    _unsafe_view_3 = torch.ops.aten._unsafe_view.default(mm_3, [primals_20, primals_21, 1024]);  mm_3 = None
    native_dropout_1 = torch.ops.aten.native_dropout.default(getitem, 0.05, True)
    getitem_4 = native_dropout_1[0];  native_dropout_1 = None
    transpose_4 = torch.ops.aten.transpose.int(primals_6, 0, 1);  primals_6 = None
    view_4 = torch.ops.aten.view.default(getitem_4, [mul, 8192]);  getitem_4 = None
    mm_4 = torch.ops.aten.mm.default(view_4, transpose_4);  view_4 = None
    _unsafe_view_4 = torch.ops.aten._unsafe_view.default(mm_4, [primals_20, primals_21, 4]);  mm_4 = None
    transpose_5 = torch.ops.aten.transpose.int(primals_7, 0, 1);  primals_7 = None
    view_5 = torch.ops.aten.view.default(_unsafe_view_4, [mul, 4]);  _unsafe_view_4 = None
    mm_5 = torch.ops.aten.mm.default(view_5, transpose_5);  view_5 = transpose_5 = None
    _unsafe_view_5 = torch.ops.aten._unsafe_view.default(mm_5, [primals_20, primals_21, 1024]);  mm_5 = None
    mul_7 = torch.ops.aten.mul.Tensor(_unsafe_view_5, 4.0);  _unsafe_view_5 = None
    add_1 = torch.ops.aten.add.Tensor(_unsafe_view_3, mul_7);  _unsafe_view_3 = mul_7 = None
    transpose_6 = torch.ops.aten.transpose.int(primals_8, 0, 1);  primals_8 = None
    mm_6 = torch.ops.aten.mm.default(view, transpose_6);  view = transpose_6 = None
    _unsafe_view_6 = torch.ops.aten._unsafe_view.default(mm_6, [primals_20, primals_21, 1024]);  mm_6 = None
    native_dropout_2 = torch.ops.aten.native_dropout.default(getitem, 0.05, True);  getitem = None
    getitem_6 = native_dropout_2[0];  native_dropout_2 = None
    transpose_7 = torch.ops.aten.transpose.int(primals_9, 0, 1);  primals_9 = None
    view_7 = torch.ops.aten.view.default(getitem_6, [mul, 8192]);  getitem_6 = None
    mm_7 = torch.ops.aten.mm.default(view_7, transpose_7);  view_7 = None
    _unsafe_view_7 = torch.ops.aten._unsafe_view.default(mm_7, [primals_20, primals_21, 4]);  mm_7 = None
    transpose_8 = torch.ops.aten.transpose.int(primals_10, 0, 1);  primals_10 = None
    view_8 = torch.ops.aten.view.default(_unsafe_view_7, [mul, 4]);  _unsafe_view_7 = None
    mm_8 = torch.ops.aten.mm.default(view_8, transpose_8);  view_8 = transpose_8 = None
    _unsafe_view_8 = torch.ops.aten._unsafe_view.default(mm_8, [primals_20, primals_21, 1024]);  mm_8 = None
    mul_11 = torch.ops.aten.mul.Tensor(_unsafe_view_8, 4.0);  _unsafe_view_8 = None
    add_2 = torch.ops.aten.add.Tensor(_unsafe_view_6, mul_11);  _unsafe_view_6 = mul_11 = None
    view_9 = torch.ops.aten.view.default(add, [primals_20, primals_21, 64, 128]);  add = None
    transpose_9 = torch.ops.aten.transpose.int(view_9, 1, 2);  view_9 = None
    view_10 = torch.ops.aten.view.default(add_1, [primals_20, primals_21, -1, 128]);  add_1 = None
    transpose_10 = torch.ops.aten.transpose.int(view_10, 1, 2);  view_10 = None
    view_11 = torch.ops.aten.view.default(add_2, [primals_20, primals_21, -1, 128]);  add_2 = None
    transpose_11 = torch.ops.aten.transpose.int(view_11, 1, 2);  view_11 = None
    slice_1 = torch.ops.aten.slice.Tensor(primals_18, 0, 0, primals_21);  primals_18 = None
    slice_2 = torch.ops.aten.slice.Tensor(primals_19, 0, 0, primals_21);  primals_19 = None
    unsqueeze = torch.ops.aten.unsqueeze.default(slice_1, 0);  slice_1 = None
    unsqueeze_1 = torch.ops.aten.unsqueeze.default(unsqueeze, 0);  unsqueeze = None
    unsqueeze_2 = torch.ops.aten.unsqueeze.default(slice_2, 0);  slice_2 = None
    unsqueeze_3 = torch.ops.aten.unsqueeze.default(unsqueeze_2, 0);  unsqueeze_2 = None
    rotary_pos_embedding = torch.ops.hpu.rotary_pos_embedding.default(transpose_9, unsqueeze_3, unsqueeze_1, primals_24, 0, 0);  transpose_9 = None
    rotary_pos_embedding_1 = torch.ops.hpu.rotary_pos_embedding.default(transpose_10, unsqueeze_3, unsqueeze_1, primals_24, 0, 0);  transpose_10 = unsqueeze_3 = unsqueeze_1 = primals_24 = None
    view_12 = torch.ops.aten.view.default(rotary_pos_embedding, [primals_20, 8, 8, primals_21, 128]);  rotary_pos_embedding = None
    view_13 = torch.ops.aten.view.default(rotary_pos_embedding_1, [primals_20, 8, 1, primals_21, 128]);  rotary_pos_embedding_1 = None
    view_14 = torch.ops.aten.view.default(transpose_11, [primals_20, 8, 1, primals_21, 128]);  transpose_11 = None
    sdpa_fwd_non_dropout = torch.ops.hpu.sdpa_fwd_non_dropout.default(view_12, view_13, view_14, None, 0.0, 0.08838834764831843, True, 'none', None, 'left');  view_12 = view_13 = view_14 = None
    getitem_8 = sdpa_fwd_non_dropout[0];  sdpa_fwd_non_dropout = None
    view_15 = torch.ops.aten.view.default(getitem_8, [primals_20, 64, primals_21, 128]);  getitem_8 = None
    transpose_12 = torch.ops.aten.transpose.int(view_15, 1, 2);  view_15 = None
    clone_4 = torch.ops.aten.clone.default(transpose_12, memory_format = torch.contiguous_format);  transpose_12 = None
    view_16 = torch.ops.aten.view.default(clone_4, [primals_20, primals_21, -1]);  clone_4 = None
    transpose_13 = torch.ops.aten.transpose.int(primals_11, 0, 1);  primals_11 = None
    view_17 = torch.ops.aten.view.default(view_16, [mul, 8192])
    mm_9 = torch.ops.aten.mm.default(view_17, transpose_13);  view_17 = transpose_13 = None
    _unsafe_view_9 = torch.ops.aten._unsafe_view.default(mm_9, [primals_20, primals_21, 8192]);  mm_9 = None
    native_dropout_3 = torch.ops.aten.native_dropout.default(view_16, 0.05, True);  view_16 = None
    getitem_11 = native_dropout_3[0];  native_dropout_3 = None
    transpose_14 = torch.ops.aten.transpose.int(primals_12, 0, 1);  primals_12 = None
    view_18 = torch.ops.aten.view.default(getitem_11, [mul, 8192])
    mm_10 = torch.ops.aten.mm.default(view_18, transpose_14);  view_18 = transpose_14 = None
    _unsafe_view_10 = torch.ops.aten._unsafe_view.default(mm_10, [primals_20, primals_21, 4]);  mm_10 = None
    transpose_15 = torch.ops.aten.transpose.int(primals_13, 0, 1);  primals_13 = None
    view_19 = torch.ops.aten.view.default(_unsafe_view_10, [mul, 4])
    mm_11 = torch.ops.aten.mm.default(view_19, transpose_15);  view_19 = transpose_15 = None
    _unsafe_view_11 = torch.ops.aten._unsafe_view.default(mm_11, [primals_20, primals_21, 8192]);  mm_11 = None
    mul_15 = torch.ops.aten.mul.Tensor(_unsafe_view_11, 4.0);  _unsafe_view_11 = None
    add_3 = torch.ops.aten.add.Tensor(_unsafe_view_9, mul_15);  _unsafe_view_9 = mul_15 = None
    add_4 = torch.ops.aten.add.Tensor(add_3, primals_22);  add_3 = None
    rms_norm_1 = torch.ops.hpu.rms_norm.default(add_4, primals_14, 1e-05)
    getitem_13 = rms_norm_1[0];  rms_norm_1 = None
    transpose_16 = torch.ops.aten.transpose.int(primals_15, 0, 1);  primals_15 = None
    view_20 = torch.ops.aten.view.default(getitem_13, [mul, 8192]);  getitem_13 = None
    mm_12 = torch.ops.aten.mm.default(view_20, transpose_16)
    _unsafe_view_12 = torch.ops.aten._unsafe_view.default(mm_12, [primals_20, primals_21, 28672]);  mm_12 = None
    silu = torch.ops.aten.silu.default(_unsafe_view_12);  _unsafe_view_12 = None
    transpose_17 = torch.ops.aten.transpose.int(primals_16, 0, 1);  primals_16 = None
    mm_13 = torch.ops.aten.mm.default(view_20, transpose_17);  view_20 = None
    _unsafe_view_13 = torch.ops.aten._unsafe_view.default(mm_13, [primals_20, primals_21, 28672]);  mm_13 = None
    mul_18 = torch.ops.aten.mul.Tensor(silu, _unsafe_view_13);  silu = _unsafe_view_13 = None
    transpose_18 = torch.ops.aten.transpose.int(primals_17, 0, 1);  primals_17 = None
    view_22 = torch.ops.aten.view.default(mul_18, [mul, 28672]);  mul_18 = mul = None
    mm_14 = torch.ops.aten.mm.default(view_22, transpose_18);  view_22 = None
    _unsafe_view_14 = torch.ops.aten._unsafe_view.default(mm_14, [primals_20, primals_21, 8192]);  mm_14 = None
    add_5 = torch.ops.aten.add.Tensor(_unsafe_view_14, add_4);  _unsafe_view_14 = None
    transpose_19 = torch.ops.aten.transpose.int(transpose_18, 0, 1);  transpose_18 = None
    view_30 = torch.ops.aten.view.default(_unsafe_view_10, [-1, 4]);  _unsafe_view_10 = None
    view_32 = torch.ops.aten.view.default(getitem_11, [-1, 8192]);  getitem_11 = None
    return [add_5, primals_1, primals_14, primals_22, transpose_1, transpose_4, transpose_7, add_4, transpose_16, transpose_17, transpose_19, view_30, view_32, primals_20, primals_21]
    
# To see more debug info, please use `graph_module.print_readable()`
bw_module GraphModule()

def forward(self, primals_20, primals_21, primals_1, primals_14, primals_22, transpose_1, transpose_4, transpose_7, add_4, transpose_16, transpose_17, transpose_19, view_30, view_32, tangents_1):
    rms_norm = torch.ops.hpu.rms_norm.default(primals_22, primals_1, 1e-05)
    getitem = rms_norm[0]
    getitem_1 = rms_norm[1];  rms_norm = None
    mul = primals_20 * primals_21
    native_dropout = torch.ops.aten.native_dropout.default(getitem, 0.05, True)
    getitem_2 = native_dropout[0];  native_dropout = None
    view_1 = torch.ops.aten.view.default(getitem_2, [mul, 8192])
    mm_1 = torch.ops.aten.mm.default(view_1, transpose_1);  view_1 = transpose_1 = None
    _unsafe_view_1 = torch.ops.aten._unsafe_view.default(mm_1, [primals_20, primals_21, 4]);  mm_1 = None
    native_dropout_1 = torch.ops.aten.native_dropout.default(getitem, 0.05, True)
    getitem_4 = native_dropout_1[0];  native_dropout_1 = None
    view_4 = torch.ops.aten.view.default(getitem_4, [mul, 8192])
    mm_4 = torch.ops.aten.mm.default(view_4, transpose_4);  view_4 = transpose_4 = None
    _unsafe_view_4 = torch.ops.aten._unsafe_view.default(mm_4, [primals_20, primals_21, 4]);  mm_4 = None
    native_dropout_2 = torch.ops.aten.native_dropout.default(getitem, 0.05, True);  getitem = None
    getitem_6 = native_dropout_2[0];  native_dropout_2 = None
    view_7 = torch.ops.aten.view.default(getitem_6, [mul, 8192])
    mm_7 = torch.ops.aten.mm.default(view_7, transpose_7);  view_7 = transpose_7 = None
    _unsafe_view_7 = torch.ops.aten._unsafe_view.default(mm_7, [primals_20, primals_21, 4]);  mm_7 = None
    rms_norm_1 = torch.ops.hpu.rms_norm.default(add_4, primals_14, 1e-05)
    getitem_13 = rms_norm_1[0]
    getitem_14 = rms_norm_1[1];  rms_norm_1 = None
    view_20 = torch.ops.aten.view.default(getitem_13, [mul, 8192]);  getitem_13 = None
    mm_12 = torch.ops.aten.mm.default(view_20, transpose_16)
    _unsafe_view_12 = torch.ops.aten._unsafe_view.default(mm_12, [primals_20, primals_21, 28672]);  mm_12 = None
    silu = torch.ops.aten.silu.default(_unsafe_view_12)
    mm_13 = torch.ops.aten.mm.default(view_20, transpose_17);  view_20 = None
    _unsafe_view_13 = torch.ops.aten._unsafe_view.default(mm_13, [primals_20, primals_21, 28672]);  mm_13 = None
    view_23 = torch.ops.aten.view.default(tangents_1, [mul, 8192])
    mm_15 = torch.ops.aten.mm.default(view_23, transpose_19);  view_23 = transpose_19 = None
    view_24 = torch.ops.aten.view.default(mm_15, [primals_20, primals_21, 28672]);  mm_15 = None
    mul_20 = torch.ops.aten.mul.Tensor(view_24, silu);  silu = None
    mul_21 = torch.ops.aten.mul.Tensor(view_24, _unsafe_view_13);  view_24 = _unsafe_view_13 = None
    view_25 = torch.ops.aten.view.default(mul_20, [mul, 28672]);  mul_20 = None
    transpose_20 = torch.ops.aten.transpose.int(transpose_17, 0, 1);  transpose_17 = None
    mm_16 = torch.ops.aten.mm.default(view_25, transpose_20);  view_25 = transpose_20 = None
    view_26 = torch.ops.aten.view.default(mm_16, [primals_20, primals_21, 8192]);  mm_16 = None
    sigmoid = torch.ops.aten.sigmoid.default(_unsafe_view_12)
    full = torch.ops.aten.full.default([primals_20, primals_21, 28672], 1, dtype = torch.bfloat16, layout = torch.strided, device = device(type='hpu', index=0), pin_memory = False)
    sub = torch.ops.aten.sub.Tensor(full, sigmoid);  full = None
    mul_22 = torch.ops.aten.mul.Tensor(_unsafe_view_12, sub);  _unsafe_view_12 = sub = None
    add_6 = torch.ops.aten.add.Scalar(mul_22, 1);  mul_22 = None
    mul_23 = torch.ops.aten.mul.Tensor(sigmoid, add_6);  sigmoid = add_6 = None
    mul_24 = torch.ops.aten.mul.Tensor(mul_21, mul_23);  mul_21 = mul_23 = None
    view_27 = torch.ops.aten.view.default(mul_24, [mul, 28672]);  mul_24 = mul = None
    transpose_21 = torch.ops.aten.transpose.int(transpose_16, 0, 1);  transpose_16 = None
    mm_17 = torch.ops.aten.mm.default(view_27, transpose_21);  view_27 = transpose_21 = None
    view_28 = torch.ops.aten.view.default(mm_17, [primals_20, primals_21, 8192]);  mm_17 = None
    add_7 = torch.ops.aten.add.Tensor(view_26, view_28);  view_26 = view_28 = None
    rms_norm_backward = torch.ops.hpu.rms_norm_backward.default(add_7, add_4, primals_14, getitem_14, True, 0);  add_7 = add_4 = primals_14 = getitem_14 = None
    getitem_15 = rms_norm_backward[0];  rms_norm_backward = None
    add_8 = torch.ops.aten.add.Tensor(tangents_1, getitem_15);  tangents_1 = getitem_15 = None
    mul_25 = torch.ops.aten.mul.Tensor(add_8, 4.0)
    view_29 = torch.ops.aten.view.default(mul_25, [-1, 8192]);  mul_25 = None
    transpose_22 = torch.ops.aten.transpose.int(view_29, 0, 1);  view_29 = None
    mm_18 = torch.ops.aten.mm.default(transpose_22, view_30);  transpose_22 = view_30 = None
    full_1 = torch.ops.aten.full.default([primals_20, primals_21, 4], 0, dtype = torch.bfloat16, layout = torch.strided, device = device(type='hpu', index=0), pin_memory = False)
    view_31 = torch.ops.aten.view.default(full_1, [-1, 4]);  full_1 = None
    transpose_23 = torch.ops.aten.transpose.int(view_31, 0, 1);  view_31 = None
    mm_19 = torch.ops.aten.mm.default(transpose_23, view_32);  view_32 = None
    full_2 = torch.ops.aten.full.default([primals_20, primals_21, 1024], 0, dtype = torch.bfloat16, layout = torch.strided, device = device(type='hpu', index=0), pin_memory = False)
    view_37 = torch.ops.aten.view.default(full_2, [-1, 1024]);  full_2 = None
    transpose_26 = torch.ops.aten.transpose.int(view_37, 0, 1);  view_37 = None
    view_38 = torch.ops.aten.view.default(_unsafe_view_7, [-1, 4]);  _unsafe_view_7 = None
    mm_21 = torch.ops.aten.mm.default(transpose_26, view_38);  view_38 = None
    view_40 = torch.ops.aten.view.default(getitem_6, [-1, 8192]);  getitem_6 = None
    mm_22 = torch.ops.aten.mm.default(transpose_23, view_40);  view_40 = None
    view_42 = torch.ops.aten.view.default(_unsafe_view_4, [-1, 4]);  _unsafe_view_4 = None
    mm_23 = torch.ops.aten.mm.default(transpose_26, view_42);  transpose_26 = view_42 = None
    view_44 = torch.ops.aten.view.default(getitem_4, [-1, 8192]);  getitem_4 = None
    mm_24 = torch.ops.aten.mm.default(transpose_23, view_44);  view_44 = None
    full_6 = torch.ops.aten.full.default([primals_20, primals_21, 8192], 0, dtype = torch.bfloat16, layout = torch.strided, device = device(type='hpu', index=0), pin_memory = False);  primals_20 = primals_21 = None
    view_45 = torch.ops.aten.view.default(full_6, [-1, 8192])
    transpose_30 = torch.ops.aten.transpose.int(view_45, 0, 1);  view_45 = None
    view_46 = torch.ops.aten.view.default(_unsafe_view_1, [-1, 4]);  _unsafe_view_1 = None
    mm_25 = torch.ops.aten.mm.default(transpose_30, view_46);  transpose_30 = view_46 = None
    view_48 = torch.ops.aten.view.default(getitem_2, [-1, 8192]);  getitem_2 = None
    mm_26 = torch.ops.aten.mm.default(transpose_23, view_48);  transpose_23 = view_48 = None
    rms_norm_backward_1 = torch.ops.hpu.rms_norm_backward.default(full_6, primals_22, primals_1, getitem_1, True, 0);  full_6 = primals_22 = primals_1 = getitem_1 = None
    getitem_17 = rms_norm_backward_1[0];  rms_norm_backward_1 = None
    add_9 = torch.ops.aten.add.Tensor(add_8, getitem_17);  add_8 = getitem_17 = None
    return [None, None, mm_26, mm_25, None, mm_24, mm_23, None, mm_22, mm_21, None, mm_19, mm_18, None, None, None, None, None, None, None, None, add_9, None, None]
    
# To see more debug info, please use `graph_module.print_readable()`
fw_graph_rng_ops {'native_dropout': native_dropout, 'native_dropout_1': native_dropout_1, 'native_dropout_2': native_dropout_2, 'native_dropout_3': native_dropout_3}
bw_graph_rng_ops {'native_dropout': native_dropout, 'native_dropout_1': native_dropout_1, 'native_dropout_2': native_dropout_2}

```

### Versions

Collecting environment information...
PyTorch version: 2.4.0a0+git95a8420
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.4 LTS (x86_64)
GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0
Clang version: Could not collect
CMake version: version 3.22.1
Libc version: glibc-2.35

Python version: 3.10.12 (main, Jul 29 2024, 16:56:48) [GCC 11.4.0] (64-bit runtime)
Python platform: Linux-5.15.0-112-generic-x86_64-with-glibc2.35
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:                       x86_64
CPU op-mode(s):                     32-bit, 64-bit
Address sizes:                      46 bits physical, 57 bits virtual
Byte Order:                         Little Endian
CPU(s):                             160
On-line CPU(s) list:                0-159
Vendor ID:                          GenuineIntel
Model name:                         Intel(R) Xeon(R) Platinum 8380 CPU @ 2.30GHz
CPU family:                         6
Model:                              106
Thread(s) per core:                 2
Core(s) per socket:                 40
Socket(s):                          2
Stepping:                           6
CPU max MHz:                        3400.0000
CPU min MHz:                        800.0000
BogoMIPS:                           4600.00
Flags:                              fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush dts acpi mmx fxsr sse sse2 ss ht tm pbe syscall nx pdpe1gb rdtscp lm constant_tsc art arch_perfmon pebs bts rep_good nopl xtopology nonstop_tsc cpuid aperfmperf pni pclmulqdq dtes64 monitor ds_cpl vmx smx est tm2 ssse3 sdbg fma cx16 xtpr pdcm pcid dca sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand lahf_lm abm 3dnowprefetch cpuid_fault epb cat_l3 invpcid_single intel_ppin ssbd mba ibrs ibpb stibp ibrs_enhanced tpr_shadow vnmi flexpriority ept vpid ept_ad fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid cqm rdt_a avx512f avx512dq rdseed adx smap avx512ifma clflushopt clwb intel_pt avx512cd sha_ni avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local split_lock_detect wbnoinvd dtherm ida arat pln pts hwp hwp_act_window hwp_epp hwp_pkg_req avx512vbmi umip pku ospke avx512_vbmi2 gfni vaes vpclmulqdq avx512_vnni avx512_bitalg tme avx512_vpopcntdq la57 rdpid fsrm md_clear pconfig flush_l1d arch_capabilities
Virtualization:                     VT-x
L1d cache:                          3.8 MiB (80 instances)
L1i cache:                          2.5 MiB (80 instances)
L2 cache:                           100 MiB (80 instances)
L3 cache:                           120 MiB (2 instances)
NUMA node(s):                       2
NUMA node0 CPU(s):                  0-39,80-119
NUMA node1 CPU(s):                  40-79,120-159
Vulnerability Gather data sampling: Mitigation; Microcode
Vulnerability Itlb multihit:        Not affected
Vulnerability L1tf:                 Not affected
Vulnerability Mds:                  Not affected
Vulnerability Meltdown:             Not affected
Vulnerability Mmio stale data:      Mitigation; Clear CPU buffers; SMT vulnerable
Vulnerability Retbleed:             Not affected
Vulnerability Spec rstack overflow: Not affected
Vulnerability Spec store bypass:    Mitigation; Speculative Store Bypass disabled via prctl and seccomp
Vulnerability Spectre v1:           Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:           Mitigation; Enhanced IBRS; IBPB conditional; RSB filling; PBRSB-eIBRS SW sequence; BHI Syscall hardening, KVM SW loop
Vulnerability Srbds:                Not affected
Vulnerability Tsx async abort:      Not affected

Versions of relevant libraries:
[pip3] habana-torch-dataloader==xxx
[pip3] habana-torch-plugin==xxx
[pip3] numpy==1.23.5
[pip3] pytorch-lightning==2.4.0
[pip3] torch==2.4.0a0+git95a8420
[pip3] torch_tb_profiler==0.4.0
[pip3] torchaudio==2.4.0a0+69d4077
[pip3] torchdata==0.7.1+5e6f7b7
[pip3] torchmetrics==1.4.1
[pip3] torchtext==0.18.0a0+9bed85d
[pip3] torchvision==0.19.0a0+48b1edf
[conda] Could not collect


cc @ezyang @chauhang @penguinwu @zou3519 @bdhirsh