llmcompressor.core.session

为 LLM 压缩工作流提供压缩会话管理。

提供主要的 CompressionSession 类来管理压缩工作流，包括生命周期管理、事件处理、回调注册和状态跟踪。

类

• CompressionSession –

  用于压缩的会话，负责保持生命周期

CompressionSession

CompressionSession()

用于压缩的会话，负责保持当前压缩会话的生命周期和状态

方法

• event –

  触发当前 CompressionSession 的事件。

• finalize –

  完成压缩会话。这将运行 finalize 方法

• get_serialized_recipe –

  :return: 当前已编译配方的序列化字符串

• initialize –

  初始化压缩会话。这将运行 initialize 方法

• log –

  记录当前事件类型的模型和损失信息

• reset –

  将会话重置为其初始状态

• reset_stage –

  重置会话以开始新阶段，配方和模型保持不变

属性

• lifecycle (CompressionLifecycle) –

  Lifecycle 用于跟踪压缩的进度

• state (State) –

  当前压缩会话的状态。State 实例

Source code in llmcompressor/core/session.py

def __init__(self):
    self._lifecycle = CompressionLifecycle()

lifecycle property

lifecycle: CompressionLifecycle

Lifecycle 用于跟踪压缩过程中的位置以及哪些修改器处于活动状态。它还提供了在生命周期上触发事件的能力。

返回

• CompressionLifecycle –

  会话的生命周期

state property

state: State

当前压缩会话的状态。State 实例用于存储压缩所需的所有信息，例如配方、模型优化器、数据等。

返回

• State –

  会话的当前状态

event

event(
    event_type: EventType,
    batch_data: Optional[Any] = None,
    loss: Optional[Any] = None,
    **kwargs,
) -> ModifiedState

触发当前 CompressionSession 的事件。

参数

• event_type

  (EventType) –

  要调用的事件类型

• batch_data

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

  用于事件的批数据

• 损失

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

  用于事件的损失（如果存在）

• kwargs

  –

  传递给生命周期事件方法的附加 kwargs

返回

• ModifiedState –

  调用事件后会话的修改状态

Source code in llmcompressor/core/session.py

def event(
    self,
    event_type: EventType,
    batch_data: Optional[Any] = None,
    loss: Optional[Any] = None,
    **kwargs,
) -> ModifiedState:
    """
    Invoke an event for current CompressionSession.

    :param event_type: the event type to invoke
    :param batch_data: the batch data to use for the event
    :param loss: the loss to use for the event if any
    :param kwargs: additional kwargs to pass to the lifecycle's event method
    :return: the modified state of the session after invoking the event
    """
    mod_data = self._lifecycle.event(
        event_type=event_type, batch_data=batch_data, loss=loss, **kwargs
    )
    return ModifiedState(
        model=self.state.model,
        optimizer=self.state.optimizer,
        loss=self.state.loss,  # TODO: is this supposed to be a different type?
        modifier_data=mod_data,
    )

finalize

finalize(**kwargs) -> ModifiedState

完成压缩会话。这将为会话生命周期中的每个修改器运行 finalize 方法。这还将把会话的状态设置为已完成状态。

参数

• kwargs

  –

  传递给生命周期的 finalize 方法的附加 kwargs

返回

• ModifiedState –

  完成后会话的修改状态

Source code in llmcompressor/core/session.py

def finalize(self, **kwargs) -> ModifiedState:
    """
    Finalize the session for compression. This will run the finalize method
    for each modifier in the session's lifecycle. This will also set the session's
    state to the finalized state.

    :param kwargs: additional kwargs to pass to the lifecycle's finalize method
    :return: the modified state of the session after finalizing
    """
    mod_data = self._lifecycle.finalize(**kwargs)

    return ModifiedState(
        model=self.state.model,
        optimizer=self.state.optimizer,
        loss=self.state.loss,
        modifier_data=mod_data,
    )

get_serialized_recipe

get_serialized_recipe() -> Optional[str]

返回

• Optional[str] –

  当前已编译配方的序列化字符串

Source code in llmcompressor/core/session.py

def get_serialized_recipe(self) -> Optional[str]:
    """
    :return: serialized string of the current compiled recipe
    """
    recipe = self.lifecycle.recipe

    if recipe is not None and hasattr(recipe, "yaml"):
        return recipe.yaml()

    logger.warning("Recipe not found in session - it may have been reset")

initialize

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

初始化压缩会话。这将为会话生命周期中的每个修改器运行 initialize 方法。这还将把会话的状态设置为已初始化状态。

参数

• recipe

  (Union[str, List[str], Recipe, List[Recipe], None], default: None ) –

  用于压缩的配方，可以是配方文件的路径、原始配方字符串、配方对象或配方对象列表。

• recipe_stage

  (Union[str, List[str], None], default: None ) –

  要针对的压缩阶段

• recipe_args

  (Union[Dict[str, Any], None], default: None ) –

  用于覆盖配方默认值的参数

• model

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

  要压缩的模型

• teacher_model

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

  用于知识蒸馏的教师模型

• 优化器

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

  用于压缩的优化器

• attach_optim_callbacks

  (bool, 默认值: True ) –

  如果为 True，则将优化器回调附加到压缩生命周期，否则为 False

• train_data

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

  用于压缩的训练数据

• val_data

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

  用于压缩的验证数据

• test_data

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

  用于压缩的测试数据

• calib_data

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

  用于压缩的校准数据

• copy_data

  (bool, 默认值: True ) –

  如果为 True，则复制数据，否则为 False

• start

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

  用于压缩的起始 epoch

• steps_per_epoch

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

  用于压缩的每个 epoch 的步数

• batches_per_step

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

  用于压缩的每步的批次数

• loggers

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

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

• kwargs

  –

  传递给生命周期的 initialize 方法的附加 kwargs

返回

• ModifiedState –

  初始化后会话的修改状态

Source code in llmcompressor/core/session.py

def initialize(
    self,
    recipe: Union[str, List[str], "Recipe", List["Recipe"], None] = None,
    recipe_stage: Union[str, List[str], None] = None,
    recipe_args: Union[Dict[str, Any], None] = None,
    model: Optional[Any] = None,
    teacher_model: Optional[Any] = None,
    optimizer: Optional[Any] = None,
    attach_optim_callbacks: bool = True,
    train_data: Optional[Any] = None,
    val_data: Optional[Any] = None,
    test_data: Optional[Any] = None,
    calib_data: Optional[Any] = None,
    copy_data: bool = True,
    start: Optional[float] = None,
    steps_per_epoch: Optional[int] = None,
    batches_per_step: Optional[int] = None,
    loggers: Union[None, LoggerManager, List[BaseLogger]] = None,
    **kwargs,
) -> ModifiedState:
    """
    Initialize the session for compression. This will run the initialize method
    for each modifier in the session's lifecycle. This will also set the session's
    state to the initialized state.

    :param recipe: the recipe to use for the compression, can be a path to a
        recipe file, a raw recipe string, a recipe object, or a list
        of recipe objects.
    :param recipe_stage: the stage to target for the compression
    :param recipe_args: the args to use for overriding the recipe defaults
    :param model: the model to compress
    :param teacher_model: the teacher model to use for knowledge distillation
    :param optimizer: the optimizer to use for the compression
    :param attach_optim_callbacks: True to attach the optimizer callbacks to the
        compression lifecycle, False otherwise
    :param train_data: the training data to use for the compression
    :param val_data: the validation data to use for the compression
    :param test_data: the testing data to use for the compression
    :param calib_data: the calibration data to use for the compression
    :param copy_data: True to copy the data, False otherwise
    :param start: the start epoch to use for the compression
    :param steps_per_epoch: the number of steps per epoch to use for the
        compression
    :param batches_per_step: the number of batches per step to use for
        compression
    :param loggers: the metrics manager to setup logging important info
        and milestones to, also accepts a list of BaseLogger(s)
    :param kwargs: additional kwargs to pass to the lifecycle's initialize method
    :return: the modified state of the session after initializing
    """
    mod_data = self._lifecycle.initialize(
        recipe=recipe,
        recipe_stage=recipe_stage,
        recipe_args=recipe_args,
        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,
        **kwargs,
    )

    return ModifiedState(
        model=self.state.model,
        optimizer=self.state.optimizer,
        loss=self.state.loss,
        modifier_data=mod_data,
    )

log

log(event_type: EventType, loss: Optional[Any] = None)

记录当前事件类型的模型和损失信息

参数

• event_type

  (EventType) –

  要记录的事件类型

• 损失

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

  要记录的损失（如果存在）

Source code in llmcompressor/core/session.py

def log(self, event_type: EventType, loss: Optional[Any] = None):
    """
    Log model and loss information for the current event type

    :param event_type: the event type to log for
    :param loss: the loss to log if any
    """
    self._log_model_info()
    self._log_loss(event_type=event_type, loss=loss)

reset

reset()

将会话重置为其初始状态

Source code in llmcompressor/core/session.py

def reset(self):
    """
    Reset the session to its initial state
    """
    self._lifecycle.reset()

reset_stage

reset_stage()

重置会话以开始新阶段，配方和模型保持不变

Source code in llmcompressor/core/session.py

def reset_stage(self):
    """
    Reset the session for starting a new stage, recipe and model stays intact
    """
    self.lifecycle.initialized_ = False
    self.lifecycle.finalized = False