Add weight activation ordering for fp8 block

[HDCharles]

Followup for #2520

to add weight activation ordering since fp8 block is now supported.

The only change would be
g_idx = range / groupsize -> g_idx = range / block_width

this would actually consolidate the existing code because as of now we have 2 ways to map columns, g_idx and the calculation for fp8 block which is doing the same thing as g_idx but as a one off. If we change it to use g_idx then we get actorder almost for free

1. use g_idx for fp8 block indexing
2. add fp8 block + act order

[rk119]

Hi!

Is this issue open to new contributors? I’d love to work on it, and I’m already getting familiar with the codebase :D

[HDCharles]

yeah go ahead, while you do this can you also get rid of the weird has_gidx code in gptq_quantize since its going to be slightly relevent to these changes.

[rk119]

Hi @HDCharles,

I have made the requested changes in PR #2616.

However, there are a few minor issues I still need to address to fully resolve this issue. Using g_idx in GPTQ block quantization is feasible, but unfortunately, when I added a test recipe to tests/llmcompressor/transformers/gptq/test_gptq_oneshot.py in my PR to confirm whether weight activation ordering works as expected for block quantization, I encountered an error related to the compressed-tensors library.

The error:

________________ ERROR collecting tests/llmcompressor/transformers/gptq/test_gptq_oneshot.py _________________
tests/llmcompressor/transformers/gptq/test_gptq_oneshot.py:119: in <module>
    weights=QuantizationArgs(
E   pydantic_core._pydantic_core.ValidationError: 1 validation error for QuantizationArgs
E     Value error, Must use group or tensor_group quantization strategy in order to apply activation ordering [type=value_error, input_value={'num_bits': 8, 'strategy...ering.WEIGHT: 'weight'>}, input_type=dict]
E       For further information visit https://errors.pydantic.dev/2.12/v/value_error

This error is caused specifically by this specifically.

When I make the following changes locally:

# validate activation ordering and strategy
if actorder is not None and strategy not in (
    QuantizationStrategy.GROUP,
    QuantizationStrategy.TENSOR_GROUP,
    QuantizationStrategy.BLOCK,
):
    raise ValueError(
        "Must use group or tensor_group quantization strategy in "
        "order to apply activation ordering"
    )

it works as expected, since I am treating block quantization as a valid quantization technique for activation ordering. This means I need to open a PR in the compressed-tensors library to enable support for actorder with QuantizationStrategy.BLOCK, which I am planning to do soon.

However, another edge case remains.

The solution above also enables ActivationOrdering.GROUP for QuantizationStrategy.BLOCK. If a GPTQ block quantization recipe with actorder=group is used, as in the following example:

GPTQModifier(
    ignore=["lm_head"],
    config_groups={
        "group_0": QuantizationScheme(
            targets=["re:.*model.layers.2.self_attn.q_proj$"],
            weights=QuantizationArgs(
                num_bits=8,
                strategy="block",
                block_structure=[2, 8],
                actorder=ActivationOrdering.GROUP,
            ),
        )
    },
)

this causes another error:

  def update_offload_parameter(module: torch.nn.Module, name: str, data: torch.Tensor):
      """
      Update the offload and onload data of an existing parameter/buffer. Supports both
      parameters of both offloaded modules and non-offloaded modules.
  
      NOTE: This function does not guard against multiple processes writing to offload
      at the same time. It is the responsibility of the caller to ensure that, for any
      parameter/buffer, only one rank calls this function at a time.
  
      NOTE: This function does not update onloaded values across ranks. The caller is
      responsible for broadcasting any updates to other ranks, if they are onloaded.
  
      :param module: module containing the parameter to update
      :param name: name of module parameter to update
      :param data: tensor to update parameter with
      """
      # breakpoint()
      if isinstance(module._parameters, OffloadCache):
          # | Component | Update Implementation       |
          # | --------- | --------------------------- |
          # | CPU       | Copy into shared cpu memory |
          # | Disk      | Write file to disk          |
          # | Device    | Copy into local device      |
          # | --------- | --------------------------- |
          # all implementations update onloaded data if applicable
          if name in module._parameters:
              cache = module._parameters
          elif name in module._buffers:
              cache = module._buffers
          else:
>               raise AttributeError(f"{type(module)} has no attribute {name}")
E               AttributeError: <class 'torch.nn.modules.linear.Linear'> has no attribute weight_g_idx

The source of this error is the lifecycle initialization here where weight_g_idx is not registered as a parameter on the module for block quantization, even when ActivationOrdering.GROUP is passed in the recipe. The following solution works by adding the weight_g_idx parameter to the target module when ActivationOrdering.GROUP is used with block quantization:

elif strategy == QuantizationStrategy.BLOCK:
    assert quantization_args.block_structure is not None
    if len(observed_shape) < 2:
        raise ValueError("Block quant requires at least 2 observed dimensions")

    block_structure = quantization_args.block_structure

    # NOTE: vllm kernels for block-quantization do not require
    # num_rows to be evenly divisible by block_structure[-2],
    # but num_cols does need to be evenly divisible by block_structure[-1]
    num_rows = math.ceil(observed_shape[-2] / block_structure[-2])
    num_cols = strategy_cdiv(observed_shape[-1], block_structure[-1], strategy)
    expected_shape = (num_rows, num_cols)
    
    # initialize activation ordering if applicable
    if actorder == ActivationOrdering.GROUP:
        init_g_idx = Parameter(
            torch.full((observed_shape[-1],), -1, device=device, dtype=torch.int),
            requires_grad=False,
        )
        module.register_parameter(f"{base_name}_g_idx", init_g_idx)

I ran all tests to check for any other bugs or inconsistencies, and everything passes, including specifically comparing the MSE in tests/llmcompressor/transformers/gptq/test_gptq_oneshot.py across all recipes. I can open a PR in compressed-tensors with these changes.

If you do not want to include support for ActivationOrdering.GROUP with block quantization and would prefer to enable only ActivationOrdering.WEIGHT, then I will open a PR in compressed-tensors with only the following changes here:

# validate activation ordering and strategy
if actorder is not None:
    if strategy not in (
        QuantizationStrategy.GROUP,
        QuantizationStrategy.TENSOR_GROUP,
        QuantizationStrategy.BLOCK,
    ):
        raise ValueError(
            "Must use group or tensor_group quantization strategy in "
            "order to apply activation ordering"
        )
    
    if actorder == ActivationOrdering.GROUP and strategy == QuantizationStrategy.BLOCK:
        raise ValueError(
            "Activation ordering by group is not supported for block strategy, "
            "please use actorder=weight or remove actorder"
        )

Do let me know your thoughts on this :D