llmcompressor.entrypoints.oneshot

def oneshot(
    # Model arguments
    model: str | PreTrainedModel,
    config_name: str | None = None,
    tokenizer: str | PreTrainedTokenizerBase | None = None,
    processor: str | ProcessorMixin | None = None,
    use_auth_token: bool = False,
    precision: str = "auto",
    tie_word_embeddings: bool = True,
    trust_remote_code_model: bool = False,
    save_compressed: bool = True,
    model_revision: str = "main",
    # Recipe arguments
    recipe: str | list[str] | None = None,
    recipe_args: list[str] | None = None,
    clear_sparse_session: bool = False,
    stage: str | None = None,
    # Dataset arguments
    dataset: str | Dataset | DatasetDict | None = None,
    dataset_config_name: str | None = None,
    dataset_path: str | None = None,
    splits: str | list[str] | dict[str, str] | None = None,
    batch_size: int = 1,
    data_collator: str | Callable = "truncation",
    num_calibration_samples: int = 512,
    shuffle_calibration_samples: bool = True,
    max_seq_length: int = 384,
    pad_to_max_length: bool = True,
    text_column: str = "text",
    concatenate_data: bool = False,
    streaming: bool = False,
    overwrite_cache: bool = False,
    preprocessing_num_workers: int | None = None,
    dataloader_num_workers: int = 0,
    min_tokens_per_module: float | None = None,
    moe_calibrate_all_experts: bool = True,
    pipeline: str | None = "independent",
    tracing_ignore: list[str] = [
        "_update_causal_mask",
        "create_causal_mask",
        "_update_mamba_mask",
        "make_causal_mask",
        "get_causal_mask",
        "mask_interface",
        "mask_function",
        "_prepare_4d_causal_attention_mask",
        "_prepare_fsmt_decoder_inputs",
        "_prepare_4d_causal_attention_mask_with_cache_position",
        "_update_linear_attn_mask",
        "project_per_layer_inputs",
    ],
    sequential_targets: list[str] | None = None,
    sequential_offload_device: str = "cpu",
    quantization_aware_calibration: bool = True,
    sequential_prefetch: bool = False,
    # Miscellaneous arguments
    output_dir: str | None = None,
    log_dir: str | None = None,
    **kwargs,
) -> PreTrainedModel:
    """
    Performs oneshot calibration on a model.

    # Model arguments
    :param model: A pretrained model identifier from huggingface.co/models or a path
        to a local model. Required parameter.
    :param distill_teacher: Teacher model (a trained text generation model)
        for distillation.
    :param config_name: Pretrained config name or path if not the same as
        model_name.
    :param tokenizer: Pretrained tokenizer name or path if not the same as
        model_name.
    :param processor: Pretrained processor name or path if not the same as
        model_name.
    :param use_auth_token: Whether to use Hugging Face auth token for private
        models.
    :param precision: Precision to cast model weights to, default to auto.
    :param tie_word_embeddings: Whether the model's input and output word embeddings
        should be left tied if possible. False means always untie.
    :param trust_remote_code_model: Whether to allow for custom models to execute
        their own modeling files.
    :param save_compressed: Whether to compress sparse models during save.
    :param model_revision: The specific model version to use (can be branch name,
        tag, or commit id).

    # Recipe arguments
    :param recipe: Path to a LLM Compressor recipe, or a list of paths
      to multiple LLM Compressor recipes.
    :param recipe_args: List of recipe arguments to evaluate, in the
        format "key1=value1", "key2=value2".
    :param clear_sparse_session: Whether to clear CompressionSession/
        CompressionLifecycle data between runs.
    :param stage: The stage of the recipe to use for oneshot.

    # Dataset arguments
    :param dataset: The name of the dataset to use (via the datasets
        library).
    :param dataset_config_name: The configuration name of the dataset
        to use.
    :param dataset_path: Path to a custom dataset. Supports json, csv, dvc.
    :param splits: Optional percentages of each split to download.
    :param batch_size: calibration dataset batch size. During calibration,
        LLM Compressor disables lm_head output computations to reduce memory
        usage from large calibration batch sizes. Large batch sizes may result
        excess padding or truncation, depending on the data_collator
    :param data_collator: The function to use to form a batch from the dataset. Can
        also specify 'truncation' or 'padding' to truncate or pad non-uniform sequence
        lengths in a batch. Defaults to 'truncation'.
    :param num_calibration_samples: Number of samples to use for one-shot
        calibration.
    :param shuffle_calibration_samples: Whether to shuffle the dataset before
        calibration.
    :param max_seq_length: Maximum total input sequence length after tokenization.
    :param pad_to_max_length: Whether to pad all samples to `max_seq_length`.
    :param text_column: Key to use as the `text` input to tokenizer/processor.
    :param concatenate_data: Whether to concatenate datapoints to fill
        max_seq_length.
    :param streaming: True to stream data from a cloud dataset.
    :param overwrite_cache: Whether to overwrite the cached preprocessed datasets.
    :param preprocessing_num_workers: Number of processes for dataset preprocessing.
    :param dataloader_num_workers: Number of worker processes for data loading. Default
        is 0 (safe for low CPU/GPU memory). Set to 2 or more for faster calibration if
        you have sufficient RAM. Custom data collators may not work with
        multiprocessing.
    :param min_tokens_per_module: Minimum percentage of tokens per
        module, relevant for MoE models.
    :param moe_calibrate_all_experts: Whether to calibrate all experts during MoE
        model calibration. When True, all experts will see all tokens during
        calibration, ensuring proper quantization statistics. When False, only
        routed experts will be used. Only relevant for MoE models. Default is True.
    :param pipeline: Calibration pipeline used to calibrate model Options:
        ['basic', 'datafree', 'sequential', 'independent']
    :param tracing_ignore: List of functions to ignore during tracing, either
        {module}.{method_name} or {function_name}
    :param sequential_targets: List of layer targets for the sequential pipeline.
        This is typically a single DecoderLayer. Not specifying this argument will
        cause the sequential pipeline to default to using the `no_split_params`
        specified by the HF model definition
    :param sequential_offload_device: Device used to offload intermediate activations
        between sequential layers. It is recommended to use `cuda:1` if using more
        than one gpu. Default is cpu.
    :param quantization_aware_calibration: Whether to enable quantization-aware
        calibration in the sequential pipeline. When True, quantization is applied
        during forward pass in calibration. When False, quantization is disabled
        during forward pass in calibration. Default is set to True.
    :param sequential_prefetch: When using the sequential pipeline, prefetch the
        next batch in a background thread to overlap onload with forward. Default
        False; set True for faster calibration when GPU memory allows.

    # Miscellaneous arguments
    :param output_dir: Path to save the output model after calibration.
        Nothing is saved if None.
    :param log_dir: Path to save logs during oneshot run.
        Nothing is logged to file if None.

    :return: The calibrated PreTrainedModel
    """

    # pass all args directly into Oneshot
    local_args = {
        k: v for k, v in locals().items() if k not in ("local_args", "kwargs")
    }
    one_shot = Oneshot(**local_args, **kwargs)
    one_shot()

    return one_shot.model