跳到内容

llmcompressor.core.session_functions

LLM 压缩工作流程的会话管理函数。

提供全局会话管理功能,包括会话创建、激活、重置操作和生命周期回调管理。

函数

  • active_session

    :return: 用于稀疏化的活动会话

  • create_session

    用于创建并生成用于稀疏化的新会话的上下文管理器。

  • reset_session

    将当前活动会话重置为其初始状态

LifecycleCallbacks

用于调用活动会话生命周期事件的类

方法

batch_end classmethod

batch_end(**kwargs) -> ModifiedState

调用活动会话的批次结束事件

参数

  • kwargs

    要传递给当前会话事件方法的其他关键字参数

返回

源自 llmcompressor/core/session_functions.py 的源代码
@classmethod
def batch_end(cls, **kwargs) -> ModifiedState:
    """
    Invoke a batch end event for the active session

    :param kwargs: additional kwargs to pass to the current session's event method
    :return: the modified state of the active session after invoking the event
    """
    active_session()._log_model_info()
    return cls.event(EventType.BATCH_END, **kwargs)

batch_start classmethod

batch_start(
    batch_data: Optional[Any] = None, **kwargs
) -> ModifiedState

调用活动会话的批次开始事件

参数

  • batch_data

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

    用于事件的批数据

  • kwargs

    要传递给当前会话事件方法的其他关键字参数

返回

源自 llmcompressor/core/session_functions.py 的源代码
@classmethod
def batch_start(cls, batch_data: Optional[Any] = None, **kwargs) -> ModifiedState:
    """
    Invoke a batch start event for the active session

    :param batch_data: the batch data to use for the event
    :param kwargs: additional kwargs to pass to the current session's event method
    :return: the modified state of the active session after invoking the event
    """
    return cls.event(EventType.BATCH_START, batch_data=batch_data, **kwargs)

calibration_epoch_end classmethod

calibration_epoch_end(**kwargs) -> ModifiedState

在校准期间调用活动会话的 epoch 结束事件。此事件应在模型校准一个 epoch 后调用

请参阅 src/llmcompressor/pipelines/basic/pipeline.py 了解使用示例

源自 llmcompressor/core/session_functions.py 的源代码
@classmethod
def calibration_epoch_end(cls, **kwargs) -> ModifiedState:
    """
    Invoke a epoch end event for the active session during calibration. This event
    should be called after the model has been calibrated for one epoch

    see `src/llmcompressor/pipelines/basic/pipeline.py` for usage example
    """
    return cls.event(EventType.CALIBRATION_EPOCH_END, **kwargs)

calibration_epoch_start classmethod

calibration_epoch_start(**kwargs) -> ModifiedState

在校准期间调用活动会话的 epoch 开始事件。此事件应在校准开始一个 epoch 前调用

请参阅 src/llmcompressor/pipelines/basic/pipeline.py 了解使用示例

源自 llmcompressor/core/session_functions.py 的源代码
@classmethod
def calibration_epoch_start(cls, **kwargs) -> ModifiedState:
    """
    Invoke a epoch start event for the active session during calibration. This event
    should be called before calibration starts for one epoch

    see `src/llmcompressor/pipelines/basic/pipeline.py` for usage example
    """
    return cls.event(EventType.CALIBRATION_EPOCH_START, **kwargs)

event classmethod

event(event_type: EventType, **kwargs) -> ModifiedState

调用活动会话的事件

参数

  • event_type

    (EventType) –

    要调用的事件类型

  • kwargs

    要传递给当前会话事件方法的其他关键字参数

返回

源自 llmcompressor/core/session_functions.py 的源代码
@classmethod
def event(cls, event_type: EventType, **kwargs) -> ModifiedState:
    """
    Invoke an event for the active session

    :param event_type: the event type to invoke
    :param kwargs: additional kwargs to pass to the current session's event method
    :return: the modified state of the active session after invoking the event
    """
    if event_type in [EventType.INITIALIZE, EventType.FINALIZE]:
        raise ValueError(
            f"Cannot invoke {event_type} event. "
            f"Use the corresponding method instead."
        )

    return active_session().event(event_type, **kwargs)

loss_calculated classmethod

loss_calculated(
    loss: Optional[Any] = None, **kwargs
) -> ModifiedState

调用活动会话的损失计算事件

参数

  • 损失

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

    用于事件的损失

  • kwargs

    要传递给当前会话事件方法的其他关键字参数

返回

源自 llmcompressor/core/session_functions.py 的源代码
@classmethod
def loss_calculated(cls, loss: Optional[Any] = None, **kwargs) -> ModifiedState:
    """
    Invoke a loss calculated event for the active session

    :param loss: the loss to use for the event
    :param kwargs: additional kwargs to pass to the current session's event method
    :return: the modified state of the active session after invoking the event
    """
    # log loss if loss calculated
    active_session()._log_loss(event_type=EventType.LOSS_CALCULATED, loss=loss)
    return cls.event(EventType.LOSS_CALCULATED, loss=loss, **kwargs)

optim_post_step classmethod

optim_post_step(**kwargs) -> ModifiedState

调用活动会话的优化器后步事件

参数

  • kwargs

    要传递给当前会话事件方法的其他关键字参数

返回

源自 llmcompressor/core/session_functions.py 的源代码
@classmethod
def optim_post_step(cls, **kwargs) -> ModifiedState:
    """
    Invoke an optimizer post-step event for the active session

    :param kwargs: additional kwargs to pass to the current session's event method
    :return: the modified state of the active session after invoking the event
    """
    return cls.event(EventType.OPTIM_POST_STEP, **kwargs)

optim_pre_step classmethod

optim_pre_step(**kwargs) -> ModifiedState

调用活动会话的优化器前步事件

参数

  • kwargs

    要传递给当前会话事件方法的其他关键字参数

返回

源自 llmcompressor/core/session_functions.py 的源代码
@classmethod
def optim_pre_step(cls, **kwargs) -> ModifiedState:
    """
    Invoke an optimizer pre-step event for the active session

    :param kwargs: additional kwargs to pass to the current session's event method
    :return: the modified state of the active session after invoking the event
    """
    return cls.event(EventType.OPTIM_PRE_STEP, **kwargs)

sequential_epoch_end classmethod

sequential_epoch_end(
    subgraph: Subgraph, **kwargs
) -> ModifiedState

调用活动会话的顺序 epoch 结束事件。此事件应在单个顺序层校准/训练一个 epoch 后调用

在单个批次校准顺序层后调用此函数,请参阅 src/llmcompressor/pipelines/sequential/pipeline.py 了解使用示例

源自 llmcompressor/core/session_functions.py 的源代码
@classmethod
def sequential_epoch_end(cls, subgraph: "Subgraph", **kwargs) -> ModifiedState:
    """
    Invoke a sequential epoch end event for the active session. This event should be
    called after one sequential layer has been calibrated/trained for one epoch

    This is called after a sequential layer has been calibrated with one batch, see
    `src/llmcompressor/pipelines/sequential/pipeline.py` for usage example
    """
    return cls.event(EventType.SEQUENTIAL_EPOCH_END, subgraph=subgraph, **kwargs)

active_session

active_session() -> CompressionSession

返回

源自 llmcompressor/core/session_functions.py 的源代码
def active_session() -> CompressionSession:
    """
    :return: the active session for sparsification
    """
    global _local_storage
    return getattr(_local_storage, "session", _global_session)

create_session

create_session() -> (
    Generator[CompressionSession, None, None]
)

用于创建并生成用于稀疏化的新会话的上下文管理器。这将设置活动会话为新会话,持续整个上下文。

返回

源自 llmcompressor/core/session_functions.py 的源代码
@contextmanager
def create_session() -> Generator[CompressionSession, None, None]:
    """
    Context manager to create and yield a new session for sparsification.
    This will set the active session to the new session for the duration
    of the context.

    :return: the new session
    """
    global _local_storage
    orig_session = getattr(_local_storage, "session", None)
    new_session = CompressionSession()
    _local_storage.session = new_session
    try:
        yield new_session
    finally:
        _local_storage.session = orig_session

reset_session

reset_session()

将当前活动会话重置为其初始状态

源自 llmcompressor/core/session_functions.py 的源代码
def reset_session():
    """
    Reset the currently active session to its initial state
    """
    session = active_session()
    session._lifecycle.reset()