[Attention] Attention head quantization strategy

[kylesayrs]

Purpose

• Support attention head quantization
  • This will be used by [Transform] Attention/Cache transforms #436

Given an attention state of shape (batch_size, num_heads, seq_len, head_dim), the head attention strategy will generate scales of shape (num_heads, 1, 1).

Prerequisites

• [Tests] Mock Observers, Static Lifecycle Tests #482

Changes

• Add attention head quantization strategy
• Fix shapes of per-tensor attention flattening
• Elaborate on attention calibration tests

Testing

• Added attention head quantization test and validated that generated scales and zero points make sense

[brian-dellabetta]

how are we expecting users to use QuantizationStrategy.ATTN_HEAD in a recipe? If i'm understanding correctly, it would look something like this?

quant_stage:
  quant_modifiers:
    QuantizationModifier:
      config_groups:
        group0:
          targets: ["re:.*self_attn$"]
          weights:
            strategy: attn_head
            ...
        group1:
          targets: ["re:.*(q|k|v)_proj$"]
          weights:
            strategy: group
            ...

[kylesayrs]

@brian-dellabetta I’ve decided that giving per-attention strategy its own strategy (rather than reusing group) makes more sense.

quant_stage:
  quant_modifiers:
    QuantizationModifier:
      config_groups:
        group0:
          targets: ["re:.*self_attn$"]
          input_activations:
            strategy: attn_head
            ...

[brian-dellabetta]

overall format LGTM, but i'm struggling with understanding how we're arriving at some of these expected_shapes

[brian-dellabetta]

Thanks for updating!

The base branch was changed.

[fynnsu]

Looks good! Do we need logic somewhere to reverse flatten_attention_for_quantization? It seems like it would be important to be sure that the unflattening process is implemented in parallel with the flattening function.

[kylesayrs]

@fynnsu The inverse function would require extra metadata (for example, unflattening (batch_size * seq_len) requires knowing either batch_size or seq_len).

Calibration only requires the forward function. Implementing the backwards function would allow us to share the util across calibration and quantization forward. This might be nice for standardization and potentially faster runtime, but isn't high priority right now.