Add Float8 Weight Only and FP8 weight + dynamic activation

[drisspg]

Summary

This PR makes some tweaks to the existing FP8 weight only work flow and adds float8_dynamic_activation_float8_weight to the quantize_ API.

New Apis made public:

from torchao.quantization import (
    quantize_,
    float8_weight_only,
    float8_dynamic_activation_float8_weight,
)

def float8_dynamic_activation_float8_weight(
    target_dtype: torch.dtype = torch.float8_e4m3fn,
    activation_dtype: torch.dtype = torch.float8_e4m3fn,
    mm_config: ScaledMMConfig = ScaledMMConfig(use_fast_accum=True)
):
    """
    Applies float8 dynamic symmetric per-tensor quantization to both activations and weights of linear layers.
    Args:
        target_dtype (torch.dtype): The target data type for weight quantization. Default is torch.float8_e4m3fn.
        activation_dtype (torch.dtype): The target data type for activation quantization. Default is torch.float8_e4m3fn.
        mm_config (ScaledMMConfig): Configuration for the matrix multiplication. Default uses fast accumulation.
    """

def float8_weight_only(target_dtype: torch.dtype = torch.float8_e4m3fn):
    """
    Applies float8 weight-only symmetric per-channel quantization to linear layers.
    
    Args:
        target_dtype (torch.dtype): The target data type for weight quantization. Default is torch.float8_e4m3fn.
    Note:
        The actual matmul will be computed in original precision of the weight tensor.
    """

TODO

• Add Rowwise scaling option to top level api
• Next would be add float8_static_activation and hooking into calibration flow
• Adding these two APIs to autoquant
• Proper way of ensuring hardware support is met when quantizing for these types

Changes

• Adds Float8Layout, the main reason for doing this is the ability to pass the mm_config down to matmul via the layout_type. This also adds the ability to do fake "transpose". This may still be needed but I have since learned that we register dispatch entires to nn.fucntional.linear. Thus the weight will not be ran through transpose.
• Wires up the TensorWise Scaling case to _scaled_mm
• Updates the Quant_primitives to handle the Symmetric no zero-point float to float dequant routine.
• Updates the AQTLayout's Typing to reflect the changes to dequant.
• I added various type hints and docs to function that I found helpful
• We were uncondtionally registering safe globals, gated that by behind the 2.5 flag
• Option for inverse scaling

Memory snapshot compare

With no_grad():

Normal conditions:

FIX

@jerryzh168 If I decorate quantize affine:

@torch.no_grad()
def quantize_affine(
    input: torch.Tensor,
    block_size: Tuple[int, ...],
    scale: torch.Tensor,
    zero_point: Optional[torch.Tensor],
    output_dtype: torch.dtype,
    quant_min: Optional[Union[int, float]] = None,
    quant_max: Optional[Union[int, float]] = None,
    zero_point_domain: ZeroPointDomain = ZeroPointDomain.INT,
) -> torch.Tensor:
    """

And

@torch.no_grad()
def choose_qparams_affine(
   input: torch.Tensor,
   mapping_type: MappingType,
   block_size: Tuple[int, ...],
   target_dtype: torch.dtype,
   quant_min: Optional[Union[int, float]] = None,
   quant_max: Optional[Union[int, float]] = None,
   eps: Optional[float] = None,
   scale_dtype: Optional[torch.dtype] = None,
   zero_point_dtype: Optional[torch.dtype] = None,
   preserve_zero: bool = True,
   zero_point_domain = ZeroPointDomain.INT,
) -> Tuple[torch.Tensor, torch.Tensor]:

I get the proper memory usage:

Is it okay If I land these as well or is there some other use case I am missing? I imagine all other input_tenosrs have been ints and havent propgated grads before

Perf micro benchmarks

Runner Script:
https://gist.github.com/drisspg/3baf802f8c8631df2a549840f39e7e0d

Trace: internal_link

Max Autotune
Needs this fix: pytorch/pytorch#134765
Trace: internal_link

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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 c771c64 with merge base 05224a9 ():
💚 Looks good so far! There are no failures yet. 💚

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

[jerryzh168]

looks good overall, just had some nit comments inline

[vkuzo]

does this mean AffineQuantizedTensor uses scale in the same way Float8Tensor uses 1 / scale, or something else?

[vkuzo]

it seems like this PR makes ScaledMMConfig a public API. Is this intended? If yes, can we move this object to float8/__init__.py, maybe make a dataclass, ensure it's documented, etc?

[drisspg]

yeah Ill put up a PR