quantize_activation_per_token_absmax use general quant primitives

[jerryzh168]

Summary:
att

Test Plan:
OSS CI

Reviewers:

Subscribers:

Tasks:

Tags:

[jerryzh168]

diff in generated quantized code: https://www.internalfb.com/phabricator/paste/view/P1226948181

[jerryzh168]

Changes to generated quantized code for vit: https://www.internalfb.com/phabricator/paste/view/P1227030036

[cpuhrsch]

Hm those changes seem to indicate that something substantial is still different about this refactor. There must be a slight difference somewhere.

[jerryzh168]

Hm those changes seem to indicate that something substantial is still different about this refactor. There must be a slight difference somewhere.

yeah I don't think the code is exactly the same, it's calling mostly the same ops though, but with different args, for example: https://github.com/pytorch/ao/blob/main/torchao/quantization/quant_primitives.py#L273 has different dim and keepdim args

is there a way to test the speed of generated code?

[jerryzh168]

Hm those changes seem to indicate that something substantial is still different about this refactor. There must be a slight difference somewhere.

yeah I don't think the code is exactly the same, it's calling mostly the same ops though, but with different args, for example: main/torchao/quantization/quant_primitives.py#L273 has different dim and keepdim args

is there a way to test the speed of generated code?

looks like the benchmark code will print the time, let me check

[jerryzh168]

before change: elapsed_time: 1.4610368347167968 milliseconds
after change: elapsed_time: 1.4671536254882813 milliseconds

looks like not much difference

[cpuhrsch]

I'm less worried about speed and more about correctness. I don't think a single benchmark datapoint will be conclusive here either. I saw an additional call to mul in your generated code. Can you try changing the generic primitives until it matches for this one example only? At least so we can see the difference in code needed to match.

[jerryzh168]

I'm less worried about speed and more about correctness. I don't think a single benchmark datapoint will be conclusive here either. I saw an additional call to mul in your generated code. Can you try changing the generic primitives until it matches for this one example only? At least so we can see the difference in code needed to match.

sure, for correctness, are you referring to numerics? so we have regression tests here: https://github.com/pytorch/ao/blob/main/test/integration/test_integration.py#L586 and it passes

[jerryzh168]

I just did a quick test with a 10x10 tensor, the quantized integer values match exactly, but there are some differences for scales, I think it's because it's doing clamping before dividing q_max, instead of after: https://github.com/pytorch/ao/blob/main/torchao/quantization/quant_primitives.py#L335

[jerryzh168]

@vkuzo @HDCharles do you know if this clampping before the scale is calculated is required for quantize_activation_per_token_absmax to work for our use cases?

[vkuzo]

we need an epsilon to handle the case where max(abs(x)) is zero

[jerryzh168]

I see, so can we do this after we divide by q_max? that's where we are doing the clamp typically

[vkuzo]

    if eps is not None:
        scale = torch.clamp(scale, min=eps)

you have this logic in choose_qparams_affine, I would imagine it's for the same purpose. I would recommend:

1. always specify epsilon (not sure why this is optional)
2. ensure test cases include testing for max(abs(x)) == 0

[cpuhrsch]

@vkuzo Is eps best chosen to be (a multiple of / the) machine epsilon of the input tensor's dtype? Or is it a parameter that needs to be searched depending on input data distribution?

[jerryzh168]

OK I'll add a test. will think about if we want to make it required, we didn't do clamping in some of the ops right now I think

[vkuzo]

will think about if we want to make it required

IMO the case of max(abs(x)) == 0 should be handled for every op

Is eps best chosen to be (a multiple of / the) machine epsilon of the input tensor's dtype?

I think that choosing eps based on the dtype makes sense, it should not be data dependent. My educated guess of how to choose this number would be "smallest number so that the resulting scale is finite", I haven't looked into whether expressing that in terms of machine epsilon would make sense.

[jerryzh168]

after a second thought I guess it makes sense to have an eps, otherwise we may have divide by zero during quantization, I'll change it in a separate PR

for eps, I think we can do torch.finfo(input_dtype).eps if it's not provided from user

[HDCharles]

can we verify the torch.compiled code is the same between the old/new primitives in a unit test? It seems pretty important to make sure this is doing what we want after torch.compile on an ongoing basis.

e.g. we use https://github.com/pytorch/ao/blob/main/test/integration/test_integration.py#L1110
to extract the triton code in this test, it'd be good if we could compare the generated triton code.

[jerryzh168]

can we verify the torch.compiled code is the same between the old/new primitives in a unit test? It seems pretty important to make sure this is doing what we want after torch.compile on an ongoing basis.

e.g. we use main/test/integration/test_integration.py#L1110 to extract the triton code in this test, it'd be good if we could compare the generated triton code.

There are some differences for this specific op, see: #193 (comment) (for vit model), but we understand that it's from the order of when we do clampping: #193 (comment)

I feel we probably want to establish high level benchmarks like eval accuracy and performance numbers for these in the future.

[jerryzh168]

can we verify the torch.compiled code is the same between the old/new primitives in a unit test? It seems pretty important to make sure this is doing what we want after torch.compile on an ongoing basis.

e.g. we use main/test/integration/test_integration.py#L1110 to extract the triton code in this test, it'd be good if we could compare the generated triton code.

oh do you mean to verify that mix_mm is called in the compiled code? that's a good idea, I can check that

edit:
it seems that this op is not related to mixed_mm though?

[HDCharles]

can we verify the torch.compiled code is the same between the old/new primitives in a unit test? It seems pretty important to make sure this is doing what we want after torch.compile on an ongoing basis.
e.g. we use main/test/integration/test_integration.py#L1110 to extract the triton code in this test, it'd be good if we could compare the generated triton code.

There are some differences for this specific op, see: #193 (comment) (for vit model), but we understand that it's from the order of when we do clampping: #193 (comment)

I feel we probably want to establish high level benchmarks like eval accuracy and performance numbers for these in the future.

i don't think that's an issue, just reverse the order in the original op and compare the generated code for that, then.

the triton generation test is mostly about verifying perf between two things written in very different ways, functional changes aren't super important to that concern (though still relevant if they are significantly different, this one shouldn't have a large impact).

[jerryzh168]

i don't think that's an issue, just reverse the order in the original op and compare the generated code for that, then.

the triton generation test is mostly about verifying perf between two things written in very different ways, functional changes aren't super important to that concern (though still relevant if they are significantly different, this one shouldn't have a large impact).

so my point is that the generated code may change and the change should be allowed if the generated code has desired op (like int4mm op), similar performance and accuracy