Skip to content

llmcompressor.modifiers.awq.mappings

Classes:

  • AWQMapping

    Dataclass storing config of activation mappings to smooth

Functions:

AWQMapping dataclass

AWQMapping(smooth_layer: str, balance_layers: list[str])

Dataclass storing config of activation mappings to smooth The output activations of smooth_layer are input activations into the balance_layers

AWQMappings are resolved into ResolvedMappings, which retain pointers to the actual torch.nn.Modules and additional metadata at runtime

ResolvedMapping dataclass

ResolvedMapping(
    smooth_name: str,
    smooth_layer: Module,
    balance_layers: List[Module],
    balance_names: Optional[List[str]] = None,
    parent: Optional[Module] = None,
    parent_name: Optional[str] = None,
)

Dataclass for storing the resolved mappings between an activation layer and the following weights that must be balanced during smoothing

Parameters:

  • smooth_name

    (str) –

    name of the activation layer

  • smooth_layer

    (Module) –

    PyTorch module storing the activation layer

  • balance_layers

    (List[Module]) –

    list of PyTorch modules that smooth_layer feeds into, must be balanced to offset the smoothing of smooth_layer

  • balance_names

    (Optional[List[str]], default: None ) –

    optional list of names of the balance_layers

  • parent

    (Optional[Module], default: None ) –

    parent module of the balance_layers

  • parent_name

    (Optional[str], default: None ) –

    name of the parent module

get_layer_mappings_from_architecture

get_layer_mappings_from_architecture(
    architecture: str,
) -> List[AWQMapping]

Parameters:

  • architecture

    (str) –

    str: The architecture of the model

Returns:

  • List[AWQMapping]

    list: The layer mappings for the given architecture

Source code in llmcompressor/modifiers/awq/mappings.py
def get_layer_mappings_from_architecture(architecture: str) -> List[AWQMapping]:
    """
    :param architecture: str: The architecture of the model
    :return: list: The layer mappings for the given architecture
    """

    if architecture not in AWQ_MAPPING_REGISTRY:
        logger.info(
            f"Architecture {architecture} not found in mappings. "
            f"Using default mappings: {_default_mappings}"
        )

    return AWQ_MAPPING_REGISTRY.get(architecture, _default_mappings)