[Observers] group size + channel wise + per token

[horheynm]

starter script

from copy import deepcopy

import torch
from compressed_tensors.quantization.lifecycle.calibration import set_module_for_calibration
from compressed_tensors.quantization.lifecycle.frozen import freeze_module_quantization
from compressed_tensors.quantization.lifecycle.initialize import (
    initialize_module_for_quantization,
)
from compressed_tensors.quantization.quant_config import QuantizationStatus
from compressed_tensors.quantization.quant_args import QuantizationArgs
from torch.nn import Linear

from compressed_tensors.quantization.quant_scheme import QuantizationScheme

num_bits = 8

quantization_scheme = QuantizationScheme(
    # input_activations=QuantizationArgs(num_bits=num_bits, symmetric=False, group_size = 4), # ADJUST GROUPSIZE HERE FOR CHANNEL WISE, GROUP WISE
    input_activations=QuantizationArgs(num_bits=num_bits, symmetric=False, group_size = -1),
    
    weights=QuantizationArgs(num_bits=num_bits, symmetric=True),
    targets=["*"],
)

layer = Linear(4, 4)
layer.weight.data *= 100

# over write forward pass and register zero_point and scale
initialize_module_for_quantization(layer, quantization_scheme)

set_module_for_calibration(layer)

layer(torch.randn(2, 4, 4))

initalized_layer = deepcopy(layer)

# calibrate the layers with each iteration
for _ in range(10):
    layer(torch.randn(4, 4))

# Freeze, no update after any forward pass
freeze_module_quantization(layer)
for _ in range(10):
    layer(torch.randn(4, 4))

[Satrat]

The main thing I see missing here is that we aren't actually using the strategy field of QuantizationArgs. It makes sense to support group_size=-1 as channelwise but I think the code would be more readable if instead of checking for group size we could just check for the QuantizationArgs.strategy enum. This would make it easier to extend when we add the token strategy too.

Maybe we could add a validator to QuantizationArgs so if the user specifies group_size=-1 we automatically set channel as the strategy

[bfineran]

let's also add at least a simple test for each strategy that validates a forward pass runs and scales/zero points have the expected shape

[Satrat]

LGTM once the test failures are fixed!