[Float8Quant] Add rowwise scaling option to float8 dyanmic quant

[drisspg]

Stacked PRs:

• ->[Float8Quant] Add rowwise scaling option to float8 dyanmic quant #819

---

[Float8Quant] Add rowwise scaling option to float8 dyanmic quant

Summary

Adds an option to do PerRow scaling of the activation and weight, which can provide faster activation quantization and better numerics.

The api for calling this is:

import torch
import torch.nn as nn
from torchao.quantization.quant_api import quantize_, PerRow, float8_dynamic_activation_float8_weight

linear_module = nn.Linear(16, 32).to("cuda")

quantize_(linear_module, float8_dynamic_activation_float8_weight(granularity=PerRow())) 

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

✅ You can merge normally! (1 Unrelated Failure)

As of commit 0405f9e with merge base 65d86c6 ():

FLAKY - The following job failed but was likely due to flakiness present on trunk:

• Build M1 Wheels / pytorch/ao / upload / wheel-py3_9-cpu (gh) (detected as infra flaky with no log or failing log classifier)

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

[jerryzh168]

LGTM!