跳到内容

llmcompressor.recipe

用于定义和管理压缩工作流的 Recipe 系统。

提供 Recipe 框架,用于指定压缩配置,包括元数据跟踪、Recipe 解析和工作流编排。支持基于阶段的执行和灵活的参数管理,以实现复杂的压缩管道。

模块

  • Recipe

    用于表示模型 Recipe 的类。

Recipe

基类:BaseModel

用于表示模型 Recipe 的类。Recipe 编码了修改模型和/或训练过程所需的指令,作为一系列修饰符(modifiers)。

Recipe 可以从文件、字符串或 HuggingFace stub 创建。可接受的文件格式包括 json 和 yaml,但序列化 Recipe 时,默认将使用 yaml。

方法

  • create_instance

    从文件、字符串或 RecipeModifier 对象创建 Recipe 实例

  • dict

    :return: Recipe 的字典表示

  • from_dict

    解析表示 Recipe 的字典并返回 Recipe 实例。

  • from_modifiers

    从修饰符列表创建 Recipe 实例

  • yaml

    返回 Recipe 的 YAML 字符串表示,

create_instance classmethod

create_instance(
    path_or_modifiers: Union[
        str, Modifier, List[Modifier], Recipe
    ],
    modifier_group_name: Optional[str] = None,
    target_stage: Optional[str] = None,
) -> Recipe

从文件、字符串或 RecipeModifier 对象创建 Recipe 实例

支持使用 Recipe 字符串或文件

recipe_str = ''' ... test_stage: ... pruning_modifiers: ... ConstantPruningModifier: ... start: 0.0 ... end: 2.0 ... targets: ['re:.*weight'] ... ''' recipe = Recipe.create_instance(recipe_str, target_stage="test_stage")

参数

  • path_or_modifiers

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

    Recipe 文件的路径或 Recipe 字符串(必须是有效的 json/yaml 文件或有效的 json/yaml 字符串)。还可以接受 RecipeModifier 实例或 RecipeModifiers 列表

  • modifier_group_name

    (Optional[str], 默认值: None ) –

    Recipe 的 stage_name。如果为 oneshottrain,则 run_type 将根据 modifier_group_name 推断;如果为 None,则会分配一个默认的虚拟 group_name。此参数仅在从 Modifier/Modifier 列表实例创建 Recipe 时使用,否则将被忽略。

返回

  • Recipe

    从路径或修饰符创建的 Recipe 实例,或有效的 yaml/json 格式的 Recipe 字符串

源代码位于 llmcompressor/recipe/recipe.py
@classmethod
def create_instance(
    cls,
    path_or_modifiers: Union[str, Modifier, List[Modifier], "Recipe"],
    modifier_group_name: Optional[str] = None,
    target_stage: Optional[str] = None,
) -> "Recipe":
    """
    Create a recipe instance from a file, string, or RecipeModifier objects


    Using a recipe string or file is supported:
    >>> recipe_str = '''
    ... test_stage:
    ...     pruning_modifiers:
    ...         ConstantPruningModifier:
    ...             start: 0.0
    ...             end: 2.0
    ...             targets: ['re:.*weight']
    ... '''
    >>> recipe = Recipe.create_instance(recipe_str, target_stage="test_stage")

    :param path_or_modifiers: The path to the recipe file or
        or the recipe string (must be a valid
        json/yaml file or a valid json/yaml string). Can also
        accept a RecipeModifier instance, or a list of
        RecipeModifiers
    :param modifier_group_name: The stage_name of the recipe,
        if `oneshot` or `train` the run_type of the recipe will be
        inferred from the modifier_group_name, if None, a dummy default
        group_name will be assigned. This argument is only used
        when creating a recipe from a Modifier/list of Modifier(s)
        instance, else it's ignored.
    :return: The Recipe instance created from the path or modifiers,
        or a valid recipe string in yaml/json format
    """
    if isinstance(path_or_modifiers, Recipe):
        # already a recipe
        return path_or_modifiers

    if isinstance(path_or_modifiers, (Modifier, list)):
        return cls.from_modifiers(
            modifiers=path_or_modifiers, modifier_group_name=modifier_group_name
        )

    if not os.path.isfile(path_or_modifiers):
        # not a local file
        # assume it's a string
        logger.debug(
            "Could not initialize recipe as a file path or zoo stub, "
            "attempting to process as a string."
        )
        logger.debug(f"Input string: {path_or_modifiers}")
        obj = _load_json_or_yaml_string(path_or_modifiers)
        return cls.from_dict(filter_dict(obj, target_stage=target_stage))
    else:
        logger.info(f"Loading recipe from file {path_or_modifiers}")

    with open(path_or_modifiers, "r") as file:
        content = file.read().strip()
        if path_or_modifiers.lower().endswith(".md"):
            content = _parse_recipe_from_md(path_or_modifiers, content)

        if path_or_modifiers.lower().endswith(".json"):
            obj = json.loads(content)
        elif path_or_modifiers.lower().endswith(
            ".yaml"
        ) or path_or_modifiers.lower().endswith(".yml"):
            obj = yaml.safe_load(content)
        else:
            try:
                obj = _load_json_or_yaml_string(content)
            except ValueError:
                raise ValueError(
                    f"Could not parse recipe from path {path_or_modifiers}"
                )
        return cls.from_dict(filter_dict(obj, target_stage=target_stage))

dict

dict(*args, **kwargs) -> Dict[str, Any]

返回

  • Dict[str, Any]

    Recipe 的字典表示

源代码位于 llmcompressor/recipe/recipe.py
def dict(self, *args, **kwargs) -> Dict[str, Any]:
    """
    :return: A dictionary representation of the recipe
    """

    return get_yaml_serializable_dict(modifiers=self.modifiers, stage=self.stage)

from_dict classmethod

from_dict(recipe_dict: Dict[str, Any]) -> Recipe

解析表示 Recipe 的字典并返回 Recipe 实例。确保所有修饰符条目都被实例化为 Modifier 对象。

参数

  • recipe_dict

    (Dict[str, Any]) –

    包含 Recipe 结构的字典。

返回

  • Recipe

    实例化了 Modifier 对象的 Recipe 实例。

源代码位于 llmcompressor/recipe/recipe.py
@classmethod
def from_dict(cls, recipe_dict: Dict[str, Any]) -> "Recipe":
    """
    Parses a dictionary representing a recipe and returns a Recipe instance.
    Ensures all modifier entries are instantiated Modifier objects.

    :param recipe_dict: Dictionary containing the recipe structure.
    :return: Recipe instance with instantiated Modifier objects.
    """
    args = recipe_dict.get("args", {})
    modifiers: List[Modifier] = []
    stage = "default"

    if not ModifierFactory._loaded:
        ModifierFactory.refresh()

    for stage_key, stage_val in recipe_dict.items():
        if stage_key.endswith("_stage") and isinstance(stage_val, dict):
            stage = stage_key.replace("_stage", "")
            for group_key, group_val in stage_val.items():
                if group_key.endswith("_modifiers") and isinstance(group_val, dict):
                    inferred_group = group_key.replace("_modifiers", "")
                    for mod_type, mod_args in group_val.items():
                        group = mod_args.get("group", inferred_group)
                        modifier = ModifierFactory.create(
                            mod_type,
                            group=group,
                            allow_registered=True,
                            allow_experimental=True,
                            **mod_args,
                        )
                        modifiers.append(modifier)

    return Recipe(
        args=args,
        stage=stage,
        modifiers=modifiers,
    )

from_modifiers classmethod

from_modifiers(
    modifiers: Union[Modifier, List[Modifier]],
    modifier_group_name: Optional[str] = None,
) -> Recipe

从修饰符列表创建 Recipe 实例

(注意:所有修饰符都被包装到一个阶段中,阶段名称为 modifier_group_name。如果 modifier_group_name 为 None,则默认 run type 为 oneshot)

生命周期: | - 验证修饰符 | - 从修饰符创建 Recipe 字符串 | - 从 Recipe 字符串创建 Recipe 实例

参数

  • 修饰符

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

    RecipeModifier 实例列表

  • modifier_group_name

    (Optional[str], 默认值: None ) –

    Recipe 的 stage_name。如果为 oneshottrain,则 run_type 将根据 modifier_group_name 推断;如果为 None,则会分配一个默认的虚拟 group_name。

返回

  • Recipe

    从修饰符创建的 Recipe 实例

源代码位于 llmcompressor/recipe/recipe.py
@classmethod
def from_modifiers(
    cls,
    modifiers: Union[Modifier, List[Modifier]],
    modifier_group_name: Optional[str] = None,
) -> "Recipe":
    """
    Create a recipe instance from a list of modifiers

    (Note: all modifiers are wrapped into a single stage
    with the modifier_group_name as the stage name. If modifier_group_name is None,
    the default run type is `oneshot`)

    Lfecycle:
    | - Validate Modifiers
    | - Create recipe string from modifiers
    | - Create recipe instance from recipe string

    :param modifiers: The list of RecipeModifier instances
    :param modifier_group_name: The stage_name of the recipe,
        if `oneshot` or `train` the run_type of the recipe will be
        inferred from the modifier_group_name, if None, a dummy default
        group_name will be assigned.
    :return: The Recipe instance created from the modifiers
    """
    logger.info("Creating recipe from modifiers")

    if isinstance(modifiers, Modifier):
        modifiers = [modifiers]

    if any(not isinstance(modifier, Modifier) for modifier in modifiers):
        raise ValueError("modifiers must be a list of Modifier instances")

    group_name = modifier_group_name or "default"

    # assume one stage for modifier instances
    recipe = cls()
    recipe.stage = group_name
    recipe.modifiers = modifiers
    return recipe

yaml

yaml(
    file_path: Optional[str] = None,
    existing_recipe_path: Optional[str] = None,
) -> str

返回 Recipe 的 YAML 字符串表示,可以选择与另一个 YAML 文件合并。

参数

  • 文件路径

    (Optional[str], 默认值: None ) –

    保存 YAML 的可选路径

  • existing_recipe_path

    (Optional[str], 默认值: None ) –

    另一个 recipe.yaml 文件的可选路径

返回

  • str

    合并后的 YAML 字符串

源代码位于 llmcompressor/recipe/recipe.py
def yaml(
    self,
    file_path: Optional[str] = None,
    existing_recipe_path: Optional[str] = None,
) -> str:
    """
    Return a YAML string representation of the recipe,
    optionally merging with another YAML file.

    :param file_path: Optional path to save YAML
    :param existing_recipe_path: Optional path to another recipe.yaml file
    :return: Combined YAML string
    """
    # Load the other recipe from file, if given
    existing_dict = {}
    if existing_recipe_path:
        with open(existing_recipe_path, "r") as f:
            existing_recipe_str = f.read()
        existing_dict = _load_json_or_yaml_string(existing_recipe_str)

    # Serialize current recipe
    self_dict = get_yaml_serializable_dict(
        modifiers=self.modifiers,
        stage=self.stage,
    )

    # Deep merge — keep both recipe contents
    merged_dict = append_recipe_dict(existing_dict, self_dict)

    # Dump YAML
    file_stream = None if file_path is None else open(file_path, "w")
    yaml_str = yaml.dump(
        merged_dict,
        stream=file_stream,
        allow_unicode=True,
        sort_keys=False,
        default_flow_style=None,
        width=88,
    )

    if file_stream:
        file_stream.close()

    return yaml_str