tree e8dd3bc01ca720e24d53be7cbcdbaaca120bf53a
parent 4656ea5768b88a98d8827349d3e789bb7e2c3b4b
author chunyuan <chunyuan.wu@intel.com> 1712035775 +0000
committer PyTorch MergeBot <pytorchmergebot@users.noreply.github.com> 1712664513 +0000

fix amp for AOTInductor (#122883)

## Pitch
This PR disables the amp when calling the inference_compiler in AOTInductor path (after having exported the model graph), following the way we disable AMP in Inductor path in https://github.com/pytorch/pytorch/pull/86515.

## Description
When testing AOTInductor AMP accuracy on CPU using the dynamo benchmark suites, multiple workloads will fail in this assertion: [assert pattern_repr not in _seen_patterns](https://github.com/pytorch/pytorch/blob/1d52c2d98507759a2143d52b922c541dbe36b6bf/torch/_inductor/pattern_matcher.py#L1095) which is called when registering SDPA patterns.

The `inference_compiler` ([fw_compiler_base](https://github.com/pytorch/pytorch/blob/1d52c2d98507759a2143d52b922c541dbe36b6bf/torch/_inductor/compile_fx.py#L1234)) will call into [_recursive_joint_graph_passe](https://github.com/pytorch/pytorch/blob/1d52c2d98507759a2143d52b922c541dbe36b6bf/torch/_inductor/compile_fx.py#L1241) and then [_sfdp_init](https://github.com/pytorch/pytorch/blob/1d52c2d98507759a2143d52b922c541dbe36b6bf/torch/_inductor/fx_passes/fuse_attention.py#L847).

When testing accuracy, we'll set [inductor_config.fallback_random = True](https://github.com/pytorch/pytorch/blob/1d52c2d98507759a2143d52b922c541dbe36b6bf/benchmarks/dynamo/common.py#L3526), which will make the `search_fn` to be `None` [here](https://github.com/pytorch/pytorch/blob/1d52c2d98507759a2143d52b922c541dbe36b6bf/torch/_inductor/fx_passes/serialized_patterns/central_index.py#L117-L118), thus the pattern will be generated runtime [here](https://github.com/pytorch/pytorch/blob/1d52c2d98507759a2143d52b922c541dbe36b6bf/torch/_inductor/pattern_matcher.py#L1083-L1084). When AMP is on, the generated pattern for SDPA FP32 will be the same as that of FP16, which makes the assertion fail.

Inductor path disables amp inside [aot_dispatch_base](https://github.com/pytorch/pytorch/blob/1d52c2d98507759a2143d52b922c541dbe36b6bf/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py#L124-L128). We follow the same way to disable for AOTInductor path here (after having exported the model graph) to fix this issue.

## UT
For the added UT, there's one case
`python test/inductor/test_aot_inductor.py -k test_amp_fallback_random_abi_compatible_cpu_with_stack_allocation_and_minimal_arrayref_interface` fails with the below error which is not caused by this PR itself. Marked it as skipped for now.
```
RuntimeError: Error in dlopen: /tmp/torchinductor_user/cf5vk3gqkbvud56qeotdxqvns4wbk3sjnlnuadolt7b6g7a6kspb/cfzjo5ackvrth2gp6oq4lfpdyfafoagodfpjvbzhsi2u64hza2vn.so: undefined symbol: _Z16aoti_torch_dtypeIN3c108BFloat16EEiv
```

Pull Request resolved: https://github.com/pytorch/pytorch/pull/122883
Approved by: https://github.com/jgong5, https://github.com/jansel