Add prelu into Autocast CPU whitelist

[yanbing-j]

Motivation

Add prelu to lower precision cast policy on AutocastCPU to fix #95365 :

Before: Within the scope of torch.cpu.amp.autocast(dtype=torch.bfloat16) , prelu cannot address the scenario of different datatypes of input and weight, will get a RuntimeError. This scenario is common in autocast, e.g, with autocast to bf16, if the op before prelu comes out a bf16 output, which is the input of prelu, and prelu's weight is fp32, then it will get a RuntimeError.

After: Within the scope of torch.cpu.amp.autocast(dtype=torch.bfloat16) , prelu be forced to run with bf16 data type.

Before #91238, when input is bf16, weight will be forced to cast to bf16. After #91238, this kind of test scenario will raise a RuntimeError. There is no precision loss since the workable one is also casting to bf16.

And this also alighs with Autocast CUDA whitelist.

cc @mcarilli @ptrblck @leslie-fang-intel @jgong5 @mingfeima @XiaobingSuper @sanchitintel @ashokei @jingxu10

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/95366

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 64b70f9:
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[leslie-fang-intel]

@yanbing-j May I kindly know why we have to put prelu into whitelist? How much will it affect some model's performance?

[yanbing-j]

@yanbing-j May I kindly know why we have to put prelu into whitelist? How much will it affect some model's performance?

I have updated the description, which will answer your questions.

[leslie-fang-intel]

Thanks for the comments. It looks like PReLU should support mixed precision input like BN instead of add it into whitellist?

[yanbing-j]

This is also align with CUDA whitelist.
Since #91238 rewrites all the code of PReLU, and merge CPU and CUDA code together, I think we don't need extra effert to split CPU and GUDA code to add the mixed precision support.

[jgong5]

I'm wondering what is the design philosophy of adding ops to the autocast whitelist. My understanding is that we only explicitly downcast ops for those that we know they are very likely get perf benefit from low-precision compute even if there is downcast cost. These compute-intensive dot-product ops that have HW acceleration support. From the autocast lists of CPU and CUDA, it seems "prelu" is the only exception. In my opinion, "prelu" is more like batch-norm and better fit "fall-through" policy instead.

Can we do type conversion on weights inside prelu just like batch-norm without throwing errors on type mismatch? @lezcano

[lezcano]

I didn't add support for type promotion in prelu in that PR just because it's a bit annoying to get that right so when implementing 2nd order derivatives by hand.
Should I implement this within the operation itself, or should this be done at the autocast level? @ngimel

[ngimel]

prelu doesn't follow standard type promotion rules, so it's indeed cumbersome to add the support at the op level. I don't think there's much harm done when prelu is just added to autocast allowlist (at worst an extra conversion of activations to low precision, but likely not that, because activations typically come out in low precision from the preceding gemm/convolution)

[lezcano]

Let's go with autocast then as per #95366 (comment)

[yanbing-j]

@pytorchbot rebase

[pytorchmergebot]

@pytorchbot successfully started a rebase job. Check the current status here

[pytorchmergebot]

Successfully rebased yanbing/add_prelu_whitelist onto refs/remotes/origin/viable/strict, please pull locally before adding more changes (for example, via git checkout yanbing/add_prelu_whitelist && git pull --rebase)

[yanbing-j]

@pytorchbot merge

[pytorchmergebot]

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here