[float8] add perf benchmarks for float8 training with rowwise + tensorwise scaling

[danielvegamyhre]

Summary

• Add float8 training performance benchmarks for rowwise + tensorwise scaling.
• Add repro steps for these benchmarks.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

⏳ No Failures, 4 Pending

As of commit 4610085 with merge base 711fa08 ():
💚 Looks good so far! There are no failures yet. 💚

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

[vkuzo]

we should include the baseline (bf16 + compile) here so it's clear what the speedup is from baseline

[danielvegamyhre]

Done

[vkuzo]

nit: how about

• add column with "speedup over baseline" instead of only explaining it in a sentence
• saying something like "rowwise scaling is better at handling outliers than tensorwise scaling, so these recipes are different points on the accuracy vs performance curve"
• saying that speedups increase as M,K,N increase, and pointing to blogs such as https://pytorch.org/blog/training-using-float8-fsdp2/ where e2e speedups as high as 1.5x are quoted. This is just to clarify that the 1.2x shown here is not the max speedup - it's just the speedup given the benchmark setup.

[danielvegamyhre]

done, let me know what you think

[danielvegamyhre]

Looks like test failure is related to #1799

[vkuzo]

Just to confirm, for tensorwise scaling, I see that https://github.com/pytorch/ao/blob/main/benchmarks/float8/training/float8_training_benchmark.sh is using recipe-lookup-by name.

Unfortunately, for tensorwise scaling, this is correct but not optimal. We also need to enable the following flags to enable float8 all-gather with FSDP2, and those flags are currently not supported when using titan's recipe-string-to-recipe lookup:

# source: https://github.com/pytorch/torchtitan/blob/main/docs/float8.md
CONFIG_FILE="./train_configs/llama3_8b.toml" ./run_train.sh --model.converters="float8" --float8.enable_fsdp_float8_all_gather --float8.precompute_float8_dynamic_scale_for_fsdp --float8.force_recompute_fp8_weight_in_bwd --training.compile

We have pytorch/torchtitan#901 to track making the torchtitan side of this better. Any chance we can update the tensorwise benchmark to include these flags, and also call out in the table that the tensorwise recipe has float8 all-gather for FSDP?

[vkuzo]

nit: baseline (typo)

[danielvegamyhre]

Just to confirm, for tensorwise scaling, I see that https://github.com/pytorch/ao/blob/main/benchmarks/float8/training/float8_training_benchmark.sh is using recipe-lookup-by name.

Unfortunately, for tensorwise scaling, this is correct but not optimal. We also need to enable the following flags to enable float8 all-gather with FSDP2, and those flags are currently not supported when using titan's recipe-string-to-recipe lookup:

# source: https://github.com/pytorch/torchtitan/blob/main/docs/float8.md
CONFIG_FILE="./train_configs/llama3_8b.toml" ./run_train.sh --model.converters="float8" --float8.enable_fsdp_float8_all_gather --float8.precompute_float8_dynamic_scale_for_fsdp --float8.force_recompute_fp8_weight_in_bwd --training.compile

We have pytorch/torchtitan#901 to track making the torchtitan side of this better. Any chance we can update the tensorwise benchmark to include these flags, and also call out in the table that the tensorwise recipe has float8 all-gather for FSDP?

Sure, but I thought enabling float8 all gather would make it less of a 1:1 comparison with rowwise? Or is the goal here just to showcase the peak achievable speedup using all optimal configs for each scaling strategy?

[vkuzo]

Or is the goal here just to showcase the peak achievable speedup using all optimal configs for each scaling strategy?

Yes, IMO that's what this should do. We should call out any features which are not implemented (or impossible to implement), but at the end we want the best speedup with each recipe, with each appropriate knobs turned on.

[danielvegamyhre]

Or is the goal here just to showcase the peak achievable speedup using all optimal configs for each scaling strategy?

Yes, IMO that's what this should do. We should call out any features which are not implemented (or impossible to implement), but at the end we want the best speedup with each recipe, with each appropriate knobs turned on.

@vkuzo I managed to find a machine without perf regression and reran all benchmarks, using the optimal configs for tensorwise scaling this time, so this is ready for another look when you have a sec

[vkuzo]

how about something like tensorwise with float8 all-gather and rowwise with bfloat16 all-gather? "optimal settings" should be true for all the rows in this table, it's just that the actual settings change.

[danielvegamyhre]

I thought about this but the thing is, tensorwise has 3 settings enabled, so listing all 3 would cause the column to become huge and make the table formatting clunky. So instead I listed the 3 settings in a bullet point below. What do you think?

[vkuzo]

I think it's important to specify that float8 all-gather is used for tensorwise, and not important to list out the three specific settings used to enable that feature.

[vkuzo]

nit: bfloat16, to match how we spell dtypes in the rest of PyTorch?

[vkuzo]

can we add PyTorch version, torchtitan version, torchao version? Ideally the script could display them.

[vkuzo]

nit: "E2e speedups as high as 1.5x..."

[vkuzo]

looks great, thank you!