[float8] add _auto_filter_for_recipe to float8

[danielvegamyhre]

Part of pytorch/torchtitan#1207

Problem

• float8 rowwise + vanilla TP in torchtitan had flat perf with respect to bfloat.
• RCA In float8 rowwise vanilla TP low throughput torchtitan#1207 found attention.wk and attention.wv layers were so small that float8 rowwise conversion resulted in significant slowdown (approx 40%) for those linears, thus the perf benefits from fp8 rowwise conversion on larger linears were nullified.
• This is because the default filter_fqns for float8 model conversion are fine for the fp8 tensorwise recipe, but bad for the float8 rowwise recipe.

Solution

This has been a footgun for various users as well (including Poolside), so I created an "auto filter" (#2410) which automatically filters Linears for a given float8 recipe, by checking for the following criteria:

1. dims not divisible by 16 (hardware requirement for float8)
2. dim sizes below thresholds that will result in worse perf for that given recipe, using simple heuristics based on the linked recipe perf tables above.
3. fqn matches one of the user defined filter_fqns

I integrated a PoC into torchtitan and the auto filter improved fp8 rowwise perf both local Llama3 8b run and Llama3 70b MAST run, compared to the default filter_fn we have now.

It prevents users from hitting this common footgun, while also preserving the flexibility to define their model-specific fqns.

Results

See pytorch/torchtitan#1207 for Llama3 70b results, TL;DR is filtering wk and wv improves TPS ~10% for vanilla TP and ~15% for async TP.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 1 New Failure

As of commit ded0931 with merge base 101c039 ():

NEW FAILURE - The following job has failed:

• Run Regression Tests / test-nightly (CUDA Nightly, linux.g5.12xlarge.nvidia.gpu, --pre torch --index-url https://downloa... / linux-job (gh)
  RuntimeError: Command docker exec -t 7a6582d45905d5961c8f1a9707510eacbc2e51d2fd1a01fa18dc0be83fab2b0c /exec failed with exit code 139

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

[danielvegamyhre]

cc @vkuzo for review

[vkuzo]

looks good, can we make sure the name has an underscore and also add a test before landing?