[cleanup][4/x] unify weight casting

[vkuzo]

Summary:

Removes redundant logic for weight casting

Performance/peak_mem on torchtitan llama 3 8B on 8 NVIDIA H100 GPUs:

before this PR stack, every experiment has float8 + compile

ac: selective(op)

• all_tensorwise with FSDP2 float8 all-gather: tps 7030, peak_mem 37.54 GiB
• all_tensorwise with FSDP2 float8 all-gather with force_recompute_fp8_weight_in_bwd=True: tps 7060, peak_mem 37.54 GiB
• all_axiswise: tps 6300, peak_mem 57.03 GiB
• lw_axiswise_with_gw_hp: tps 5996, peak_mem 57.03 GiB

ac: none

• all_tensorwise with FSDP2 float8 all-gather: tps 7321, peak_mem 56.56 GiB
• all_tensorwise with FSDP2 float8 all-gather with force_recompute_fp8_weight_in_bwd=True: tps 7443, peak_mem 51.39 GiB
• all_tensorwise: tps 7000, peak_mem 56.95GiB
• all_tensorwise with force_recompute_fp8_weight_in_bwd=True: tps 7200, peak_mem 51.78 GiB

after this PR stack

ac: selective(op)

• all_tensorwise with FSDP2 float8 all-gather: tps 7050, peak_mem 37.54 GiB
• all_tensorwise with FSDP2 float8 all-gather with force_recompute_fp8_weight_in_bwd=True: 7040 tps, 37.54 GiB
• all_axiswise: tps 6300, peak_mem 57.03 GiB
• lw_axiswise_with_gw_hp: tps 5996, peak_mem 57.03 GiB

ac: none

• all_tensorwise with FSDP2 float8 all-gather: tps 7280, peak_mem 56.56 GiB
• all_tensorwise with FSDP2 float8 all-gather with force_recompute_fp8_weight_in_bwd=True: tps 7283, peak_mem 56.56 GiB
• all_tensorwise: tps 6500, peak_mem 56.95 GiB
• all_tensorwise with force_recompute_fp8_weight_in_bwd=True: tps 6500, peak_mem 56.95 GiB

Test Plan:

./test/float8/test_everything.sh

Reviewers:

Subscribers:

Tasks:

Tags:

[vkuzo]

Stack from ghstack (oldest at bottom):

• -> [cleanup][4/x] unify weight casting #1481

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

✅ No Failures

As of commit 9f15418 with merge base 12396c6 ():
💚 Looks good so far! There are no failures yet. 💚

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

[vkuzo]

config.force_recompute_fp8_weight_in_bwd doesn't work with FSDP2 with float8 all-gather properly yet, looking