llmcompressor.entrypoints.utils

用于入口点预处理和后处理操作的实用函数。

提供由一次性入口点使用的通用实用函数。包括模型加载、配置设置、预处理步骤以及压缩工作流的后处理操作。

函数

• post_process –

  如果提供了 model_args，则将模型和分词器/处理器保存到输出目录。

• pre_process –

  为校准准备模型和分词器/处理器。

post_process

post_process(
    model_args: ModelArguments | None = None,
    recipe_args: RecipeArguments | None = None,
    output_dir: str | None = None,
)

如果提供了 model_args 和 output_dir，则将模型和分词器/处理器保存到输出目录。

如果 output_dir 不是默认目录，则该方法会重置生命周期操作。如果 model_args 中指定了压缩格式，则会以压缩格式保存模型。此外，还会保存分词器或处理器（如果可用）。

引发: ValueError: 如果由于 output_dir 无效或其他问题导致保存失败。

源代码位于 llmcompressor/entrypoints/utils.py

def post_process(
    model_args: ModelArguments | None = None,
    recipe_args: RecipeArguments | None = None,
    output_dir: str | None = None,
):
    """
    Saves the model and tokenizer/processor to the output directory if model_args,
    output_dir is provided.

    If the `output_dir` is not the default directory, the method resets lifecycle
    actions. The model is saved in a compressed format if specified in `model_args`.
    Additionally, the tokenizer or processor, if available, is also saved.

    Raises:
        ValueError: If saving fails due to an invalid `output_dir` or other issues.
    """
    # remove any existing dispatches
    if model_args is not None and model_args.model is not None:
        remove_dispatch(model_args.model)

    if model_args is not None and output_dir is not None:
        if recipe_args is not None and getattr(recipe_args, "stage", None) is not None:
            output_dir = os.path.join(output_dir, recipe_args.stage)
            os.makedirs(output_dir, exist_ok=True)
            logger.info(f"[Save] Stage detected. Updating output_dir to {output_dir}")

        # TODO: support general saving parameters, beyond save_compressed
        model_args.model.save_pretrained(
            output_dir, save_compressed=model_args.save_compressed
        )

        if model_args.processor is not None:
            model_args.processor.save_pretrained(output_dir)

    else:
        logger.warning(
            "Optimized model is not saved. To save, please provide"
            "`output_dir` as input arg."
            "Ex. `oneshot(..., output_dir=...)`"
        )

    # Reset the one-time-use session upon completion
    if recipe_args is not None and recipe_args.clear_sparse_session:
        reset_session()

pre_process

pre_process(
    model_args: ModelArguments,
    dataset_args: DatasetArguments,
    output_dir: str | None,
)

为校准准备模型和分词器/处理器。- 如果模型指定为路径或字符串，则初始化模型。- 应用补丁以修复绑定的张量问题并修改 save_pretrained 的行为。- 如果处理器指定为路径或 None，则初始化处理器。- 如果提供了 dataset_args，则设置每个模块的最小 token 数。引发: FileNotFoundError: 如果模型或处理器路径无效。

源代码位于 llmcompressor/entrypoints/utils.py

def pre_process(
    model_args: ModelArguments,
    dataset_args: DatasetArguments,
    output_dir: str | None,
):
    """
    Prepares the model and tokenizer/processor for calibration.
    - Initializes the model if it's specified as a path or string.
    - Applies patches to fix tied tensor issues and modifies `save_pretrained`
        behavior.
    - Initializes the processor if specified as a path or `None`.
    - Sets the minimum tokens per module if `dataset_args` are provided.
    Raises:
        FileNotFoundError: If the model or processor path is invalid.
    """

    # Initialize model
    if isinstance(model_args.model, (str, PosixPath)):
        model = initialize_model_from_path(model_args)
        if is_fsdp_model(model):
            raise NotImplementedError(
                "FSDP models are not supported in the current release but will be "
                "suported in future releases of LLM Compressor."
            )
        model_args.model = model

    # Initialize processor if dataset provided
    if isinstance(model_args.processor, (str, type(None))):
        try:
            model_args.processor = initialize_processor_from_path(
                model_args, model_args.model
            )
        except Exception as e:
            if dataset_args.is_dataset_provided():
                raise RuntimeError(
                    "An error occurred when attempting to initialize "
                    "model processor, which is required when a dataset "
                    "is provided. To resolve, create and pass in a "
                    "processor directly to `oneshot`/`train`."
                ) from e
            elif output_dir:
                logger.warning(
                    "Model processor could not be auto-initialized and "
                    "will not be saved along with the model. To resolve, "
                    "create and pass in a processor directly to "
                    f"`oneshot`/`train`.\nInitialization Error: {e}"
                )

    # untie tie_word_embeddings weights
    if not model_args.tie_word_embeddings:
        untie_word_embeddings(model_args.model)

    # wrap model.save_pretrained
    modify_save_pretrained(model_args.model)