llmcompressor.utils.pytorch.module

实用/辅助函数

函数

• get_layer_by_name –

  按名称获取模块的层。

• get_layers –

  根据目标获取模块的层（也称为子模块）。

• get_matching_layer –

  给定一个目标正则表达式，找到模块中最接近匹配的层名。

• get_no_split_params –

  获取不应在分片时拆分的模块类列表。对于

• qat_active –

  通过检查确定模型中是否有任何层启用了量化。

get_layer_by_name

get_layer_by_name(
    layer_name: str, module: Module
) -> Module

按名称获取模块的层。

参数

• layer_name

  (str) –

  要查找的层的名称。空字符串返回模块本身。

• module

  (Module) –

  要在其中搜索 layer_name 的模块。

返回

• 模块 –

  模块，名为 layer_name 的层。

源文件 llmcompressor/utils/pytorch/module.py

def get_layer_by_name(layer_name: str, module: Module) -> Module:
    """
    Get the layer of a module by name.
    :param layer_name: Name of the layer to find. Empty string returns the
        module itself.
    :param module: Module in which to search for layer_name
    :return: Module, the layer with name layer_name
    """
    if not layer_name:
        return module
    return attrgetter(layer_name)(module)

get_layers

get_layers(
    targets: Union[str, List[str]],
    module: Module,
    exclude_internal_modules: bool = False,
) -> Dict[str, Module]

根据目标获取模块的层（也称为子模块）。

参数

• targets

  (Union[str, List[str]]) –

  要搜索的名称或正则表达式。可以是正则表达式，例如“re:.*input_layernorm$”来查找模块中所有以字符串“input_layernorm”结尾的层。

• module

  (Module) –

  要在其中搜索目标的父模块。

• exclude_internal_modules

  (bool, 默认值: False ) –

  如果为 True，则不包括 llm-compressor 添加的内部模块，例如 Observers 和 Transforms。默认为 False，以保持向后兼容性。

返回

• Dict[str, Module] –

  字典 {层名 -> 模块}，包含模块中所有匹配目标的层。

源文件 llmcompressor/utils/pytorch/module.py

def get_layers(
    targets: Union[str, List[str]],
    module: Module,
    exclude_internal_modules: bool = False,
) -> Dict[str, Module]:
    """
    Get layers (also known as submodules) of module based on targets

    :param targets: names or regexes to search for
        Can be regex, e.g. "re:.*input_layernorm$" to find all layers
        in module whose names end in string "input_layernorm"
    :param module: Parent module in which to search for targets
    :param exclude_internal_modules: If True, don't include internal
        modules added by llm-compressor, e.g. Observers and Transforms.
        Defaults to False to maintain backward compatibility

    :return: dict of {layer name -> module} of all layers in module
        that match targets
    """
    layer_dict = match_layers_params(targets, module)
    if exclude_internal_modules:
        layer_dict = {
            name: layer
            for name, layer in layer_dict.items()
            if not isinstance(layer, InternalModule)
        }

    return layer_dict

get_matching_layer

get_matching_layer(
    target: str, name_to_match: str, module: Module
) -> Optional[Tuple[str, Module]]

给定一个目标正则表达式，找到模块中最接近匹配 name_to_match 字符串的层名。这用于匹配同一层中的子模块，例如将“re.*k_proj”匹配到“model.decoder.layer.0.q_proj”以找到层 0 中存在的 k_proj。

参数

• target

  (str) –

  要搜索的正则表达式。

• name_to_match

  (str) –

  要匹配的完整层名，应该存在于模块中。

• module

  (Module) –

  要在其中搜索目标的模块。

返回

• Optional[Tuple[str, Module]] –

  包含符合目标正则表达式并最匹配 name_to_match 的层名和模块的元组，如果找不到匹配项则为 None。

源文件 llmcompressor/utils/pytorch/module.py

def get_matching_layer(
    target: str, name_to_match: str, module: Module
) -> Optional[Tuple[str, Module]]:
    """
    Given a target regex, find the layer name in the module that most closely matches
    the name_to_match string. This is used to matches submodules in the same layer, for
    instance matching "re.*k_proj" to "model.decoder.layer.0.q_proj" to find the k_proj
    that exists in layer 0.

    :param target: regex to search for
    :param name_to_match: full layer name to match to, should exist in module
    :param module: module to search for target in
    :return: Tuple containing the layer name and module that fits the target regex and
    best matches name_to_match, or None if no match can be found
    """
    potential_matches = get_layers(target, module)
    largest_substring = 0
    match = None
    for name, module in potential_matches.items():
        seq_matcher = difflib.SequenceMatcher(None, name, name_to_match)
        _, _, match_length = seq_matcher.find_longest_match(
            0, len(name), 0, len(name_to_match)
        )
        if match_length > largest_substring:
            match = (name, module)
            largest_substring = match_length

    return match

get_no_split_params

get_no_split_params(
    model: PreTrainedModel,
) -> Union[str, List[str]]

获取不应在分片时拆分的模块类列表。对于 Hugging Face Transformer 模型，这是 decoder 层类型。对于其他类型的模型，它仅返回所有模块名称。

返回

• Union[str, List[str]] –

  不应拆分的类名列表。

源文件 llmcompressor/utils/pytorch/module.py

def get_no_split_params(model: PreTrainedModel) -> Union[str, List[str]]:
    """
    Get list of module classes that shouldn't be split when sharding. For
    Hugging Face Transformer models, this is the decoder layer type. For other
    types of models, this just returns all module names.

    :return: list of class names that shouldn't be split
    """
    # importing here to avoid circular import
    from llmcompressor.utils.fsdp.helpers import maybe_get_wrapped

    model = maybe_get_wrapped(model)
    no_split_modules = model._get_no_split_modules("auto")
    if len(no_split_modules) <= 0:
        return ALL_TARGET

    return no_split_modules

qat_active

qat_active(module: Module) -> bool

通过检查 weight_fake_quant 属性来确定模型中是否有任何层启用了量化。

参数

• module

  (Module) –

  要检查量化的 PyTorch 模型。

返回

• bool –

  如果模型中任何地方的量化都已激活，则为 True，否则为 False。

源文件 llmcompressor/utils/pytorch/module.py

def qat_active(module: Module) -> bool:
    """
    Determines if any layers in the model have quantization enabled by checking for
    weight_fake_quant attributes

    :param module: PyTorch model to check for quantization
    :return: True if quantization is active anywhere in the model, False otherwise
    """
    for _, layer in module.named_modules():
        if isinstance(layer, torch.quantization.FakeQuantize):
            return True
        if is_module_quantized(layer):
            return True

    return False