compressed_tensors.entrypoints.convert.convert_checkpoint

def convert_checkpoint(
    model_stub: str | os.PathLike,
    save_directory: str | os.PathLike,
    converter: Converter | list[Converter],
    max_workers: int | Literal["auto"] = "auto",
    device: str | torch.device | list[str | torch.device] | None = None,
    job_memory_estimator: Callable[
        [InverseWeightMap, list[Converter]], int
    ] = estimate_job_memory,
):
    """
    Convert a model checkpoint to either:
    - its equivalent quantized format in compressed-tensors
    - the unquantized format
    without loading it up in memory, instead operating directly on the model
    safetensors files. This entrypoint operates on a model stub or folder containing
    weights saved in safetensors files, and updates the corresponding
    quantization_config field in the config.json. All additional files will be
    copied to new checkpoint.

    :param model_stub: huggingface model hub or path to local weights files
    :param save_directory: new checkpoint will be saved in this directory.
    :param max_workers: number of worker threads to process files with. If
        "auto" (default), the number of workers is chosen from the number of
        safetensors files, the estimated memory of each conversion job, the free
        memory of each accelerator, and the number of CPUs available to this
        process. Host memory is not taken into account, so pass a smaller number
        of workers if conversion runs out of host memory
    :param device: device or devices on which to run conversion. When omitted,
        all available accelerator devices are used, falling back to CPU when no
        accelerator is available.
    :param job_memory_estimator: callable returning the estimated memory in bytes
        for a conversion job, given its inverse weight map and the converters to
        apply. Defaults to a meta-tensor profiler; used only for accelerator
        scheduling.
    :param converter: single converter or list of converters to apply
        in order, e.g. a dequantizer followed by a re-quantizer
    """
    if max_workers != "auto" and not (isinstance(max_workers, int) and max_workers > 0):
        raise ValueError(
            f"max_workers must be a positive integer or 'auto', got {max_workers!r}"
        )

    converters = converter if isinstance(converter, list) else [converter]
    devices = _resolve_devices(device)

    # get all model_files for checkpoint
    model_files = get_checkpoint_files(model_stub)

    weight_map = get_weight_map(model_files)

    # Build inverse_weight_maps, so that each job knows how to load up every necessary
    # weight and its dependencies
    inverse_weight_maps = build_inverse_weight_maps(
        weight_map=weight_map,
        model_files=model_files,
        converters=converters,
    )
    Path(save_directory).mkdir(parents=True, exist_ok=True)

    # Build validation/conversion jobs, copy over any other file
    validate_jobs = []
    convert_jobs = []
    for shard_name, resolved_path in model_files.items():
        save_path = Path(save_directory) / shard_name
        save_path.parent.mkdir(parents=True, exist_ok=True)

        if shard_name.endswith("safetensors"):
            if shard_name not in inverse_weight_maps:
                raise ValueError(
                    f"Could not find inverse_weight_map for shard {shard_name}"
                )
            validate_jobs.append(
                (validate_file, inverse_weight_maps[shard_name], converters)
            )
            convert_jobs.append(
                (convert_file, inverse_weight_maps[shard_name], save_path, converters)
            )

        else:
            if is_weights_file(shard_name):
                logger.warning(f"Skip processing for weights file {shard_name}")
            if str(resolved_path) != str(save_path):
                logger.debug(f"Copying {shard_name} -> {save_path}")
                shutil.copyfile(resolved_path, save_path)

    # Validate before long-running procssing job. Validation runs on meta tensors,
    # so it uses as many workers as possible, regardless of max_workers
    num_validate_workers = max(1, min(len(validate_jobs), _max_threads()))
    exec_jobs(validate_jobs, num_validate_workers, desc="Validating")

    # Process weights, accumulating total bytes used and the new weight_map.
    # exec_jobs_dynamic handles CPU natively (running sequentially), so the same
    # scheduler drives both CPU and accelerator runs. Memory estimates are only
    # consulted for accelerator scheduling, so skip the profiler on CPU.
    total_size = 0
    weight_map = dict()
    callable_jobs = [partial(job[0], *job[1:]) for job in convert_jobs]
    if any(dev.type != "cpu" for dev in devices):
        memory_estimates = [
            job_memory_estimator(job[1], converters) for job in convert_jobs
        ]
    else:
        memory_estimates = [0] * len(callable_jobs)
    if max_workers == "auto":
        max_workers = _auto_max_workers(devices, memory_estimates)
    convert_results = exec_jobs_dynamic(
        jobs=callable_jobs,
        devices=devices,
        max_workers=max_workers,
        memory_estimates=memory_estimates,
        desc="Converting",
    )
    for _total_size, _weight_map in convert_results:
        total_size += _total_size
        weight_map.update(_weight_map)

    # Update config and safetensors index
    write_checkpoint_quantization_config(save_directory, converters)
    update_safetensors_index(save_directory, total_size, weight_map)

def exec_jobs(
    jobs: list[tuple[Callable, ...]], max_workers: int = 1, desc: str = "Executing Jobs"
) -> list:
    """
    Execute jobs in parallel, using ThreadPoolExecutor

    :param jobs: list of tuples, the first entry of which is the callable,
        and the remaining elements are the inputs args to the callable
    :param max_workers: number of workers to use
    :param desc: tqdm description
    """
    results = []

    # For easier debugging, don't run single-threaded jobs via ThreadPoolExecutor
    if max_workers == 1:
        for job in tqdm.tqdm(jobs, desc=desc):
            results.append(job[0](*job[1:]))
        return results

    with ThreadPoolExecutor(max_workers) as executor:
        futures = [executor.submit(*job) for job in jobs]
        for future in tqdm.tqdm(as_completed(futures), total=len(futures), desc=desc):
            results.append(future.result())

    return results