Skip to content

llmcompressor.modifiers.awq.mappings

AWQMapping dataclass

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

Source code in llmcompressor/modifiers/awq/mappings.py
@dataclass
class AWQMapping:
    """
    Dataclass storing config of activation mappings to smooth
    The output activations of smooth_layer are input activations
    into the balance_layers

    `AWQMapping`s are resolved into `ResolvedMapping`s, which
    retain pointers to the actual `torch.nn.Module`s and additional
    metadata at runtime
    """

    smooth_layer: str
    balance_layers: list[str]

ResolvedMapping dataclass

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

Parameters:

Name Type Description Default
smooth_name str

name of the activation layer

required
smooth_layer Module

PyTorch module storing the activation layer

required
balance_layers List[Module]

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

required
balance_names Optional[List[str]]

optional list of names of the balance_layers

None
parent Optional[Module]

parent module of the balance_layers

None
parent_name Optional[str]

name of the parent module

None
Source code in llmcompressor/modifiers/awq/mappings.py
@dataclass
class ResolvedMapping:
    """
    Dataclass for storing the resolved mappings between an activation layer
    and the following weights that must be balanced during smoothing

    :param smooth_name: name of the activation layer
    :param smooth_layer: PyTorch module storing the activation layer
    :param balance_layers: list of PyTorch modules that smooth_layer feeds into, must be
        balanced to offset the smoothing of smooth_layer
    :param balance_names: optional list of names of the balance_layers
    :param parent: parent module of the balance_layers
    :param parent_name: name of the parent module
    """

    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

get_layer_mappings_from_architecture(architecture)

Parameters:

Name Type Description Default
architecture str

str: The architecture of the model

required

Returns:

Type Description
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)