Fix failing FP6 benchmark

[tobiasvanderwerff]

Summary:
I was playing around with FP6 and noticed that its benchmark script is no longer functional. It seems that it's still using an older API. This PR fixes the script in accordance with the FP6 correctness test in ao/test/test_ops.py:

ao/test/test_ops.py

Lines 62 to 70 in 653efe9

	@parametrize("ebits,mbits", [(3, 2), (2, 2)])
	def test_quant_llm_linear_correctness(self, ebits, mbits, BS, OC, IC, splitK):
	# adapted from https://github.com/usyd-fsalab/fp6_llm/blob/5df6737cca32f604e957e3f63f03ccc2e4d1df0d/tests/python/kernel_test_fpx.py
	floatx_weight, scale, fp16_act = self._create_floatx_inputs(ebits, mbits, BS, OC, IC, "cuda")
	results_floatx = torchao.ops.quant_llm_linear(ebits, mbits, fp16_act, floatx_weight, scale, splitK)
	fp16_weight = from_scaled_tc_floatx(floatx_weight, ebits, mbits, scale).half()
	results_fp16 = fp16_act @ fp16_weight.T

Here is the utility function _create_floatx_inputs:

ao/test/test_ops.py

Lines 36 to 42 in 653efe9

	def _create_floatx_inputs(self, ebits: int, mbits: int, BS: int, OC: int, IC: int, device):
	# Randomly initialize each byte
	nbits = 1 + ebits + mbits
	floatx_weight = torch.randint(256, (OC, IC // 8 * nbits), dtype=torch.uint8)
	scale = torch.rand(OC).half() + 0.5
	fp16_act = torch.rand(BS, IC).half() + 0.5
	return floatx_weight.to(device), scale.to(device), fp16_act.to(device)

[pytorch-bot]

🔗 Helpful Links

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

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

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

[jerryzh168]

thanks, this is using a higher level API, not old API. but it makes sense to keep consistent with test_ops I think

[jerryzh168]

actually, maybe just adding a comment explaining this is enough, it's easier for end users to use the high level quantize_ API and see the perf, since this is what they will be calling for the models in the end. see "Kernel Microbenchmarks" in #391 for how these benchmark script fit into the optimization process.

probably better to just add a comment so it's easier for end users to use these benchmark script

[tobiasvanderwerff]

Thanks for the feedback @jerryzh168. I don't think I made it clear enough, but currently the fp6_benchmark.py script is failing because of this line:

ao/benchmarks/benchmark_fp6.py

Line 12 in 858205a

fp6_weight = to_affine_quantized_floatx(float_data, FloatxTensorCoreLayoutType(3, 2))

The to_affine_quantized_floatx function is not called with the right arguments (it expects 5 of them, but it passes only two), so in that sense the current benchmark script is outdated and worth updating. What do you think?

[jerryzh168]

Thanks for the feedback @jerryzh168. I don't think I made it clear enough, but currently the fp6_benchmark.py script is failing because of this line:

ao/benchmarks/benchmark_fp6.py

Line 12 in 858205a

fp6_weight = to_affine_quantized_floatx(float_data, FloatxTensorCoreLayoutType(3, 2))

The to_affine_quantized_floatx function is not called with the right arguments (it expects 5 of them, but it passes only two), so in that sense the current benchmark script is outdated and worth updating. What do you think?

oh I see, yeah makes sense to update this, it should be to_affine_quantized_fpx I think, the renaming PR broken this..

[tobiasvanderwerff]

That makes sense -- let me update the script with your change, @jerryzh168. Do you think it makes sense to open a new PR and close the current one?

[tobiasvanderwerff]

I tested the change (it works) and created a new PR for it: #931. Closing this PR.