llmcompressor.pytorch.model_load.helpers

函数

• get_completed_stages –

  给定一个分阶段运行的检查点目录，获取已完成的阶段列表

• get_session_model –

  :return: 由活动的 CompressionSession 存储的 pytorch 模块，

• parse_dtype –

  :param dtype_arg: 要解析的 dtype 或字符串

• save_checkpoint –

  保存模型、处理器和配方

• save_completed_stages –

  将已完成的阶段列表保存到检查点目录

get_completed_stages

get_completed_stages(checkpoint_dir: Any) -> List[str]

给定一个分阶段运行的检查点目录，如果 checkpoint_dir 是一个字符串，则获取先前运行中已完成的阶段列表

参数

• checkpoint_dir

  (Any) –

  分阶段检查点的路径

返回

• List[str] –

  已完成的阶段名称列表

源代码在 llmcompressor/pytorch/model_load/helpers.py

def get_completed_stages(checkpoint_dir: Any) -> List[str]:
    """
    Given a checkpoint directory for a staged run, get the list of stages that
    have completed in a prior run if the checkpoint_dir is a string

    :param checkpoint_dir: path to staged checkpoint
    :return: list of completed stage names
    """
    if isinstance(checkpoint_dir, str):
        stage_path = os.path.join(checkpoint_dir, COMPLETED_STAGES_FILENAME)
        if os.path.exists(stage_path):
            with open(stage_path) as stage_file:
                stage_data = json.load(stage_file)
                return stage_data["completed"]

    return []

get_session_model

get_session_model() -> Optional[Module]

返回

• Optional[Module] –

  由活动的 CompressionSession 存储的 pytorch 模块，如果会话未激活则为 None

源代码在 llmcompressor/pytorch/model_load/helpers.py

def get_session_model() -> Optional[Module]:
    """
    :return: pytorch module stored by the active CompressionSession,
        or None if no session is active
    """
    session = active_session()
    if not session:
        return None

    active_model = session.state.model
    return active_model

load_safetensors_state_dict

load_safetensors_state_dict(
    file_path: str,
) -> Dict[str, torch.Tensor]

从磁盘加载 safetensors 文件

参数

• 文件路径

  (str) –

  safetensors 文件的路径

返回

• Dict[str, Tensor] –

  safetensors 数据字典

源代码在 llmcompressor/pytorch/model_load/helpers.py

def load_safetensors_state_dict(file_path: str) -> Dict[str, torch.Tensor]:
    """
    Load a safetensors file from disk

    :param file_path: path to the safetensors file
    :return: dictionary of safetensors data
    """
    with safe_open(file_path, framework="pt", device="cpu") as f:
        return {key: f.get_tensor(key) for key in f.keys()}

parse_dtype

parse_dtype(dtype_arg: Union[str, dtype]) -> torch.dtype

参数

• dtype_arg

  (Union[str, dtype]) –

  要解析的 dtype 或字符串

返回

• dtype –

  从输入字符串解析的 torch.dtype

源代码在 llmcompressor/pytorch/model_load/helpers.py

def parse_dtype(dtype_arg: Union[str, torch.dtype]) -> torch.dtype:
    """
    :param dtype_arg: dtype or string to parse
    :return: torch.dtype parsed from input string
    """
    dtype_arg = str(dtype_arg)
    dtype = "auto"  # get precision from model by default
    if dtype_arg in ("half", "float16", "torch.float16"):
        dtype = torch.float16
    elif dtype_arg in ("torch.bfloat16", "bfloat16"):
        dtype = torch.bfloat16
    elif dtype_arg in ("full", "float32", "torch.float32"):
        dtype = torch.float32

    return dtype

save_checkpoint

save_checkpoint(
    save_path: str,
    model: PreTrainedModel,
    processor: Optional[Processor] = None,
    save_safetensors: bool = True,
    save_compressed: bool = True,
    skip_sparsity_compression_stats: bool = False,
)

保存模型、处理器和配方

参数

• save_path

  (str) –

  用于保存模型和处理器的路径

• model

  (PreTrainedModel) –

  要保存的模型

• processor

  (Optional[Processor], default: None ) –

  要保存的处理器

• save_safetensors

  (bool, 默认值: True ) –

  使用 safetensors 文件类型保存模型检查点

• save_compressed

  (bool, 默认值: True ) –

  使用 compressed-tensors 格式保存模型检查点

源代码在 llmcompressor/pytorch/model_load/helpers.py

def save_checkpoint(
    save_path: str,
    model: PreTrainedModel,
    processor: Optional[Processor] = None,
    save_safetensors: bool = True,
    save_compressed: bool = True,
    skip_sparsity_compression_stats: bool = False,
):
    """
    Save a model, processor, and recipe

    :param save_path: Path used to save model and processor
    :param model: model to save
    :param processor: processor to save
    :param save_safetensors: save model checkpoint using safetensors file type
    :param save_compressed: save model checkpoint using compressed-tensors format
    """
    from llmcompressor.transformers.compression.compressed_tensors_utils import (
        get_model_compressor,  # avoid circular import
    )

    # used for decompression
    # unfortunately, if skip_sparsity_compression_stats==True, sparsity stats
    # are computed twice. In the future, track sparsity from recipe or
    # share recipe between compression and decompression
    compressor = get_model_compressor(
        model=model,
        save_compressed=save_compressed,
        skip_sparsity_compression_stats=skip_sparsity_compression_stats,
    )

    # saving the model also saves the recipe
    model.save_pretrained(
        save_path,
        save_safetensors=save_safetensors,
        save_compressed=save_compressed,
        skip_sparsity_compression_stats=skip_sparsity_compression_stats,
    )
    if processor is not None:
        processor.save_pretrained(save_path)

    # decompression: saving the model modifies the model strcuture
    # as this is only a checkpoint, decompress model to enable future training/oneshot
    if compressor is not None:
        compressor.decompress_model(model)

save_completed_stages

save_completed_stages(
    checkpoint_dir: str, completed_stages: List[str]
)

将已完成的阶段列表保存到检查点目录

参数

• checkpoint_dir

  (str) –

  要将阶段保存到的模型检查点目录

• completed_stages

  (List[str]) –

  已运行阶段的名称列表

源代码在 llmcompressor/pytorch/model_load/helpers.py

def save_completed_stages(checkpoint_dir: str, completed_stages: List[str]):
    """
    Save a list of completed stages to a checkpoint directory

    :param checkpoint_dir: model checkpoint directory to save stages to
    :param completed_stages: list of stage names that have been run
    """
    stage_path = os.path.join(checkpoint_dir, COMPLETED_STAGES_FILENAME)
    with open(stage_path, "w") as out_file:
        json.dump({"completed": completed_stages}, out_file)