llmcompressor.core.state

用于管理 LLM Compressor 状态的模块。

提供用于保存和更新与数据、硬件和模型压缩相关的状态信息的类。

类

• Data –

  用于保存训练、验证数据的类，

• Hardware –

  用于保存有关所用硬件信息的类。

• ModifiedState –

  用于表示修改后的模型、优化器和损失的类。

• State –

  状态类保存有关当前压缩状态的信息。

Data dataclass

Data(
    train: Optional[Any] = None,
    val: Optional[Any] = None,
    test: Optional[Any] = None,
    calib: Optional[Any] = None,
)

用于保存训练、验证、测试和/或校准的各种数据集的类。每个数据集都是一个 ModifiableData 实例。

参数

• train

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

  训练数据集

• 验证

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

  验证数据集

• test

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

  测试数据集

• calib

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

  校准数据集

Hardware dataclass

Hardware(
    device: Optional[str] = None,
    devices: Optional[List[str]] = None,
    rank: Optional[int] = None,
    world_size: Optional[int] = None,
    local_rank: Optional[int] = None,
    local_world_size: Optional[int] = None,
    distributed: Optional[bool] = None,
    distributed_strategy: Optional[str] = None,
)

用于保存有关所用硬件信息的类。

参数

• device

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

  用于训练的当前设备

• devices

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

  将用于训练的所有设备列表

• 排名

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

  当前设备的 rank

• world_size

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

  使用的设备总数

• local_rank

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

  当前设备的本地 rank

• local_world_size

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

  本地机器上使用的设备总数

• distributed

  (Optional[bool], default: None ) –

  是否正在使用分布式训练

• distributed_strategy

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

  正在使用的分布式策略

ModifiedState dataclass

ModifiedState(model, optimizer, loss, modifier_data)

用于表示修改后的模型、优化器和损失的类。

参数

• model

  (Optional[Any]) –

  修改后的模型

• 优化器

  (Optional[Any]) –

  修改后的优化器

• 损失

  (Optional[Any]) –

  修改后的损失

• modifier_data

  (Optional[List[Dict[str, Any]]]) –

  用于修改模型、优化器和损失的修饰符数据

使用给定的参数初始化 ModifiedState。

参数

• model

  (Any) –

  修改后的模型

• 优化器

  (Any) –

  修改后的优化器

• 损失

  (Any) –

  修改后的损失

• modifier_data

  (List[Dict[str, Any]]) –

  用于修改模型、优化器和损失的修饰符数据

源代码位于 llmcompressor/core/state.py

def __init__(self, model, optimizer, loss, modifier_data):
    """
    Initialize the ModifiedState with the given parameters.

    :param model: The modified model
    :type model: Any
    :param optimizer: The modified optimizer
    :type optimizer: Any
    :param loss: The modified loss
    :type loss: Any
    :param modifier_data: The modifier data used to modify the model, optimizer,
        and loss
    :type modifier_data: List[Dict[str, Any]]
    """
    self.model = model
    self.optimizer = optimizer
    self.loss = loss
    self.modifier_data = modifier_data

State dataclass

State(
    model: Any = None,
    teacher_model: Any = None,
    optimizer: Any = None,
    optim_wrapped: bool = None,
    loss: Any = None,
    batch_data: Any = None,
    data: Data = Data(),
    hardware: Hardware = Hardware(),
    loggers: Optional[LoggerManager] = None,
    model_log_cadence: Optional[float] = None,
    _last_log_step: Union[float, int, None] = None,
)

状态类保存有关当前压缩状态的信息。

参数

• model

  (Any, 默认值: None ) –

  用于压缩的模型

• teacher_model

  (Any, 默认值: None ) –

  用于压缩的教师模型

• 优化器

  (Any, 默认值: None ) –

  用于训练的优化器

• optim_wrapped

  (bool, 默认值: None ) –

  优化器是否已被包装

• 损失

  (Any, 默认值: None ) –

  用于训练的损失函数

• batch_data

  (Any, 默认值: None ) –

  当前用于压缩的批次数据

• data

  (Data, default: Data() ) –

  用于训练、验证、测试和/或校准的数据集，包装在 Data 实例中

• 硬件

  (Hardware, default: Hardware() ) –

  硬件实例，包含有关正在使用的目标硬件的信息

• loggers

  (Optional[LoggerManager], default: None ) –

  LoggerManager 实例，包含所有要记录的日志记录器

• model_log_cadence

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

  记录模型信息相对于 epoch 的节奏。如果为 1，则每个 epoch 记录一次。如果为 2，则每隔一个 epoch 记录一次，依此类推。默认值为 1。

方法

• update –

  使用给定参数更新状态。

属性

• compression_ready (bool) –

  检查模型和优化器是否已设置为进行压缩。

compression_ready property

compression_ready: bool

检查模型和优化器是否已设置为进行压缩。

返回

• bool –

  如果模型和优化器已设置，则为 True，否则为 False

update

update(
    model: Any = None,
    teacher_model: Any = None,
    optimizer: Any = None,
    attach_optim_callbacks: bool = True,
    train_data: Any = None,
    val_data: Any = None,
    test_data: Any = None,
    calib_data: Any = None,
    copy_data: bool = True,
    start: float = None,
    steps_per_epoch: int = None,
    batches_per_step: int = None,
    loggers: Union[
        None, LoggerManager, List[BaseLogger]
    ] = None,
    model_log_cadence: Optional[float] = None,
    **kwargs,
) -> Dict

使用给定参数更新状态。

参数

• model

  (Any, 默认值: None ) –

  用于更新状态的模型

• teacher_model

  (Any, 默认值: None ) –

  用于更新状态的教师模型

• 优化器

  (Any, 默认值: None ) –

  用于更新状态的优化器

• attach_optim_callbacks

  (bool, 默认值: True ) –

  是否附加优化器回调

• train_data

  (Any, 默认值: None ) –

  用于更新状态的训练数据

• val_data

  (Any, 默认值: None ) –

  用于更新状态的验证数据

• test_data

  (Any, 默认值: None ) –

  用于更新状态的测试数据

• calib_data

  (Any, 默认值: None ) –

  用于更新状态的校准数据

• copy_data

  (bool, 默认值: True ) –

  是否复制数据

• start

  (float, 默认值: None ) –

  用于更新状态的起始索引

• steps_per_epoch

  (int, 默认值: None ) –

  用于更新状态的每个 epoch 的步数

• batches_per_step

  (int, 默认值: None ) –

  用于更新状态的每个步骤的批次数量

• loggers

  (Union[None, LoggerManager, List[BaseLogger]], default: None ) –

  用于设置记录重要信息和里程碑的指标管理器，也接受 BaseLogger 列表

• model_log_cadence

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

  记录模型信息相对于 epoch 的节奏。如果为 1，则每个 epoch 记录一次。如果为 2，则每隔一个 epoch 记录一次，依此类推。默认值为 1。

• kwargs

  –

  用于更新状态的额外关键字参数

返回

• Dict –

  更新后的状态字典

源代码位于 llmcompressor/core/state.py

def update(
    self,
    model: Any = None,
    teacher_model: Any = None,
    optimizer: Any = None,
    attach_optim_callbacks: bool = True,
    train_data: Any = None,
    val_data: Any = None,
    test_data: Any = None,
    calib_data: Any = None,
    copy_data: bool = True,
    start: float = None,
    steps_per_epoch: int = None,
    batches_per_step: int = None,
    loggers: Union[None, LoggerManager, List[BaseLogger]] = None,
    model_log_cadence: Optional[float] = None,
    **kwargs,
) -> Dict:
    """
    Update the state with the given parameters.

    :param model: The model to update the state with
    :type model: Any
    :param teacher_model: The teacher model to update the state with
    :type teacher_model: Any
    :param optimizer: The optimizer to update the state with
    :type optimizer: Any
    :param attach_optim_callbacks: Whether or not to attach optimizer callbacks
    :type attach_optim_callbacks: bool
    :param train_data: The training data to update the state with
    :type train_data: Any
    :param val_data: The validation data to update the state with
    :type val_data: Any
    :param test_data: The testing data to update the state with
    :type test_data: Any
    :param calib_data: The calibration data to update the state with
    :type calib_data: Any
    :param copy_data: Whether or not to copy the data
    :type copy_data: bool
    :param start: The start index to update the state with
    :type start: float
    :param steps_per_epoch: The steps per epoch to update the state with
    :type steps_per_epoch: int
    :param batches_per_step: The batches per step to update the state with
    :type batches_per_step: int
    :param loggers: The metrics manager to setup logging important info and
        milestones to, also accepts a list of BaseLogger(s)
    :type loggers: Union[None, LoggerManager, List[BaseLogger]]
    :param model_log_cadence: The cadence to log model information w.r.t epochs.
        If 1, logs every epoch. If 2, logs every other epoch, etc. Default is 1.
    :type model_log_cadence: Optional[float]
    :param kwargs: Additional keyword arguments to update the state with
    :return: The updated state as a dictionary
    :rtype: Dict
    """
    logger.debug(
        "Updating state with provided parameters: {}",
        {
            "model": model,
            "teacher_model": teacher_model,
            "optimizer": optimizer,
            "attach_optim_callbacks": attach_optim_callbacks,
            "train_data": train_data,
            "val_data": val_data,
            "test_data": test_data,
            "calib_data": calib_data,
            "copy_data": copy_data,
            "start": start,
            "steps_per_epoch": steps_per_epoch,
            "batches_per_step": batches_per_step,
            "loggers": loggers,
            "model_log_cadence": model_log_cadence,
            "kwargs": kwargs,
        },
    )

    if model is not None:
        self.model = model
    if teacher_model is not None:
        self.teacher_model = teacher_model
    if optimizer is not None:
        self.optim_wrapped = attach_optim_callbacks
        self.optimizer = optimizer
    if train_data is not None:
        self.data.train = train_data if not copy_data else deepcopy(train_data)
    if val_data is not None:
        self.data.val = val_data if not copy_data else deepcopy(val_data)
    if test_data is not None:
        self.data.test = test_data if not copy_data else deepcopy(test_data)
    if calib_data is not None:
        self.data.calib = calib_data if not copy_data else deepcopy(calib_data)

    if "device" in kwargs:
        self.hardware.device = kwargs["device"]

    loggers = loggers or []
    if isinstance(loggers, list):
        loggers = LoggerManager(loggers)
    self.loggers = loggers

    if model_log_cadence is not None:
        self.model_log_cadence = model_log_cadence

    return kwargs