update keepdim params for 1D scale factor

[namgyu-youn]

Overview:
In _choose_scale_float8, the per-tensor quantization case (len(block_size) == 0) uses tensor.amax(keepdim=True) while _choose_qparams_affine uses torch.amax(..., keepdim=False) for the same purpose.

This PR aligns _choose_scale_float8 with _choose_qparams_affine by using tensor.amax(keepdim=False) for 1D scale factor.

Related Issue/PR: #3324

Test Plan: test/quantization/test_quant_primitives.py

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

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

[namgyu-youn]

@pytorchbot label "topic: not user facing"

[namgyu-youn]

Local test code:

import torch
from torchao.quantization.quant_primitives import _choose_scale_float8, _choose_scale_float8_old

a = torch.randn(4, 4)
scale = _choose_scale_float8(a, block_size=(4, 1))  # keepdim = False
print(scale)

scale_old = _choose_scale_float8_old(a, block_size=(4, 1))  # keepdim = True
print(scale_old)

And the result is (same):

tensor([[0.0023, 0.0044, 0.0040, 0.0027]])
tensor([[0.0023, 0.0044, 0.0040, 0.0027]])

[jerryzh168]

this is not enough I think there are

ao/torchao/quantization/quant_primitives.py

Lines 2320 to 2323 in 2ff1eb2

	output_shape = [
	input_size // block_size[i] for i, input_size in enumerate(tensor.shape)
	]
	scale = scale.reshape(output_shape)

that will expand the dimension of scale

I actually tried this locally, and doesn't seem to work very well. I think we can leave this for now, but one thing that is useful is to remove the need to do

ao/torchao/quantization/quant_primitives.py

Line 2332 in 2ff1eb2

def _maybe_expand_scale_to_tensor_shape(

, I don't know why we need this if float8 is already using a scale that's matching the dimension of input:

ao/torchao/quantization/quant_primitives.py

Lines 2318 to 2323 in 2ff1eb2

	# Reshape scale back to match the expected output shape
	# The scale tensor should have the same shape as the input divided by block_size
	output_shape = [
	input_size // block_size[i] for i, input_size in enumerate(tensor.shape)
	]
	scale = scale.reshape(output_shape)

maybe try removing calls to _maybe_expand_scale_to_tensor_shape in code and make sure all tests still pass is a good task to work on

also another thing we can do is to simplify the implementation of

ao/torchao/float8/inference.py

Line 130 in 2ff1eb2

def _slice_scale_for_dimension(

since for float8 scale is always matching the dim of input