[examples] Add selector-based PTQ override policies to quantization YAML configs

[mhs4670go]

Summary

Currently, quantization examples configure PTQ behavior through YAML files under configs. The existing configuration supports coarse-grained quantization policies such as activation, linear_weight, embedding_weight, lm_head_weight, and similar fields.

We want to extend the YAML config format so users can override quantization specs for specific parts of a model without editing Python code.

The main use cases are:

1. Apply an MX quantization spec to activations of a specific op type across all layers.
2. Apply a different spec, such as int16, only to a specific activation observer in a specific layer.
3. Support both text-only models and multi-component models such as VLMs, where text and vision paths may differ.

This issue covers the initial implementation up to Phase 3:

• Phase 1: Add specs, selector-based overrides, and raw_overrides to PTQ YAML configs.
• Phase 2: Add alias-based selectors such as op_type: linear and observer_role: activation.
• Phase 3: Extend selector resolution to model components such as text and vision.

---

Motivation

The current YAML format is good for coarse-grained PTQ configuration, but it is not expressive enough for mixed quantization policies.

Examples of desired configurations:

• Use MX for every Linear activation in all decoder layers.
• Use MX for most activations, but keep one specific layer's mlp.down_proj.act_out as int16.
• Apply different policies to text and vision blocks in a multimodal model.
• Avoid hard-coding model-specific override paths in user configs.

We should provide a user-friendly selector layer in YAML and resolve it internally into the exact PTQConfig.overrides structure used by the quantization wrappers.

Related: #804

---

Goals

• Keep existing YAML configs backward compatible.
• Preserve existing coarse-grained fields such as activation, linear_weight, embedding_weight, and lm_head_weight.
• Add a user-facing override mechanism for fine-grained PTQ control.
• Allow users to target:
  • all layers,
  • a list of layers,
  • a layer range,
  • specific modules,
  • op aliases such as Linear,
  • activation or parameter observer roles,
  • specific observer names.
• Support model components such as text and vision.
• Resolve selector-based rules into exact internal PTQ override paths.
• Allow advanced users to provide exact raw override paths when needed.
• Provide deterministic precedence when multiple rules match the same target.
• Validate selector rules and fail early on invalid or empty matches.

---

Proposed YAML schema

The PTQ stage should support three new optional fields:

pipeline:
  - name: ptq
    enabled: true

    # Existing coarse-grained base policies
    activation: int16
    linear_weight: uint8
    embedding_weight: uint8
    lm_head_weight: uint8

    # New: named quantization specs
    specs:
      mx_fp8_act:
        kind: mx
        elem_format: fp8_e4m3
        axis: -1
        shared_exp_method: max
        round: nearest

      int16_act:
        kind: affine
        dtype: int16

    # New: selector-based user-facing override rules
    overrides:
      - name: all_text_linear_activations_mx
        target:
          component: text
          layers: all
          op_type: linear
          observer_role: activation
        spec: mx_fp8_act

      - name: layer_7_down_proj_output_int16
        target:
          component: text
          layers: [7]
          module: mlp.down_proj
          observers: [act_out]
        spec: int16_act

    # New: exact internal override paths for advanced users
    raw_overrides:
      model.layers.7.mlp.down_proj.act_out: int16

---

Field definitions

specs

specs is a named registry of quantization specs.

The value should use the same grammar already accepted by the existing quant spec parser, for example:

specs:
  int16_act:
    kind: affine
    dtype: int16

  mx_fp8_act:
    kind: mx
    elem_format: fp8_e4m3
    axis: -1

A rule may reference a named spec:

spec: mx_fp8_act

A rule may also define an inline spec:

spec:
  kind: affine
  dtype: int16

overrides

overrides is a list of selector-based rules.

Each rule should have:

- name: string
  enabled: bool        # optional, default true
  target: object
  spec: string|object

name should be required so that rules are easier to debug and can later be overridden from CLI or reported in resolved manifests.

raw_overrides

raw_overrides is an advanced escape hatch.

It maps exact internal PTQ override paths to quant specs:

raw_overrides:
  model.layers.7.mlp.down_proj.act_out: int16
  model.layers.7.self_attn.softmax:
    kind: mx
    elem_format: fp8_e4m3
    axis: -1

This is useful for debugging and highly specific observer control, but it should not be the recommended user-facing API.

---

Target selector schema

The first version should support the following fields:

target:
  component: text | vision | all
  layers: all | int | list[int] | string
  module: string
  modules: list[string]
  op_type: string
  observer_role: string
  observers: list[string]

component

Selects the model component.

Supported values for Phase 3:

component: text
component: vision
component: all

For text-only models, text and all should be valid. vision should raise a validation error if the model does not have a vision component.

For multimodal models, the adapter should resolve each component to the correct internal root path.

Examples:

target:
  component: text
  layers: all
  op_type: linear
  observer_role: activation

target:
  component: vision
  layers: all
  op_type: linear
  observer_role: activation

layers

Selects layer indices.

Supported forms:

layers: all
layers: 7
layers: [0, 1, 7]
layers: "0:8"
layers: "0:32:2"

Layer indices should be zero-based.

Range strings should follow Python slice semantics:

"start:end"
"start:end:step"

Examples:

layers: all
layers: [0, 3, 7]
layers: "0:32:2"

module and modules

Select specific relative module paths inside a layer.

Examples:

module: mlp.down_proj

modules:
  - self_attn.q_proj
  - self_attn.k_proj
  - self_attn.v_proj
  - self_attn.o_proj
  - mlp.gate_proj
  - mlp.up_proj
  - mlp.down_proj

Only one of module, modules, or op_type should be used in a single rule.

op_type

A user-friendly alias for a group of modules.

Phase 2 should support at least:

op_type: linear

For decoder-only LLMs, op_type: linear should expand to the model family's known Linear projection modules, for example:

modules:
  - self_attn.q_proj
  - self_attn.k_proj
  - self_attn.v_proj
  - self_attn.o_proj
  - mlp.gate_proj
  - mlp.up_proj
  - mlp.down_proj

For multimodal models, each adapter should define the appropriate linear module set for each component.

Future aliases may include:

op_type: attention
op_type: mlp
op_type: norm
op_type: embedding
op_type: lm_head

However, only linear is required for this issue.

observer_role

A user-friendly alias for observer groups.

Phase 2 should support:

observer_role: activation
observer_role: parameter
observer_role: input_activation
observer_role: output_activation

Suggested expansion:

observer_role: activation
# expands to activation observers such as act_in and act_out for Linear wrappers

observer_role: input_activation
# expands to act_in

observer_role: output_activation
# expands to act_out

observer_role: parameter
# expands to parameter observers such as weight

For wrappers with custom observer names, aliases should only expand to well-known safe names. Users can use observers for exact names.

observers

Select exact observer names.

Examples:

observers: [act_in, act_out]

observers: [act_out]

observers: [softmax]

observers should be treated as exact names and should not be guessed.

---

Example 1: Apply MX to all Linear activations in text layers

pipeline:
  - name: ptq
    enabled: true

    activation: int16
    linear_weight: uint8

    specs:
      mx_fp8_act:
        kind: mx
        elem_format: fp8_e4m3
        axis: -1
        shared_exp_method: max
        round: nearest

    overrides:
      - name: all_text_linear_activations_mx
        target:
          component: text
          layers: all
          op_type: linear
          observer_role: activation
        spec: mx_fp8_act

Expected behavior:

• The base activation policy is int16.
• All text Linear activation observers matched by op_type: linear and observer_role: activation are overridden to mx_fp8_act.
• Linear weights still use uint8.

---

Example 2: Override one specific activation back to int16

pipeline:
  - name: ptq
    enabled: true

    activation:
      kind: mx
      elem_format: fp8_e4m3
      axis: -1

    linear_weight: uint8

    specs:
      int16_act:
        kind: affine
        dtype: int16

    overrides:
      - name: layer_7_down_proj_output_int16
        target:
          component: text
          layers: [7]
          module: mlp.down_proj
          observers: [act_out]
        spec: int16_act

Expected behavior:

• Most activations use the base MX activation policy.
• Only text.layers.7.mlp.down_proj.act_out is overridden to int16.

The exact internal resolved path may differ by model family, for example:

model.layers.7.mlp.down_proj.act_out

or:

model.language_model.layers.7.mlp.down_proj.act_out

The adapter should be responsible for resolving the logical selector to the correct internal path.

---

Example 3: Different policies for text and vision components

pipeline:
  - name: ptq
    enabled: true

    activation: int16
    linear_weight: uint8

    specs:
      mx_text_act:
        kind: mx
        elem_format: fp8_e4m3
        axis: -1

      mx_vision_act:
        kind: mx
        elem_format: fp8_e5m2
        axis: -1

    overrides:
      - name: text_linear_activations_mx
        target:
          component: text
          layers: all
          op_type: linear
          observer_role: activation
        spec: mx_text_act

      - name: vision_linear_activations_mx
        target:
          component: vision
          layers: all
          op_type: linear
          observer_role: activation
        spec: mx_vision_act

Expected behavior:

• Text Linear activations use mx_text_act.
• Vision Linear activations use mx_vision_act.
• Exact root paths are resolved by the model adapter.

---

Internal design

The user-facing YAML should not directly depend on fragile internal module paths.

Instead, selector-based rules should be compiled by the model adapter or PTQ config builder into exact internal PTQConfig.overrides.

Recommended flow:

YAML PTQ stage config
  ├─ existing base policies
  │    ├─ activation
  │    ├─ linear_weight
  │    ├─ embedding_weight
  │    └─ lm_head_weight
  ├─ specs
  ├─ overrides
  └─ raw_overrides

adapter.build_ptq_config(...)
  ├─ read model family structure
  ├─ read component roots
  ├─ read layer counts
  ├─ build existing default PTQ override tree
  ├─ compile selector-based overrides
  ├─ merge raw_overrides
  └─ return PTQConfig(..., overrides=merged_overrides)

This should happen after the model is loaded, because selector resolution may require:

• number of text layers,
• number of vision layers or blocks,
• model family-specific root paths,
• model family-specific module aliases,
• wrapper-specific observer names.

[Torrero]

Small remark:
For cases where modules or layers overlap in the override rules, we should define a priority of override rules (for example, in the order in which the rules are mentioned).

[mhs4670go]

Good point. Actually I intentionally omit that information as the issue was too long but it should be written down for clarification.

Precedence rules

When multiple policies match the same target, resolution must be deterministic.

Recommended precedence, from weakest to strongest:

1. Existing base PTQ policy
   Example: activation, linear_weight

2. Existing model-family default override tree
   Example: adapter-generated overrides for Linear weights

3. Selector-based overrides

4. More specific selector-based overrides

5. raw_overrides

6. CLI-level config overrides, if applicable

For selector-based rules:

• Higher specificity should win.
• If specificity is equal, the later rule in YAML should win.

Suggested specificity order:

layers: all
  < layers: range/list
  < layers: single index

op_type
  < modules
  < module

observer_role
  < observers

For example:

overrides:
  - name: all_linear_activations_mx
    target:
      component: text
      layers: all
      op_type: linear
      observer_role: activation
    spec: mx_fp8_act

  - name: layer_7_down_proj_output_int16
    target:
      component: text
      layers: [7]
      module: mlp.down_proj
      observers: [act_out]
    spec: int16_act

The second rule should override the first rule only for layer 7 mlp.down_proj.act_out.

[Torrero]

@mhs4670go, could you clarify please,
why mlp.down_proj.act_out should be in int16?

[mhs4670go]

@mhs4670go, could you clarify please,
why mlp.down_proj.act_out should be in int16?

Ah, sorry for confusion. It's just an example. No practical meaning.