Skip to content

llmcompressor.core.session_functions

Session management functions for LLM compression workflows.

Provides global session management functionality including session creation, activation, reset operations, and lifecycle callback management.

Classes:

Functions:

  • active_session –

    :return: the active session for sparsification

  • create_session –

    Context manager to create and yield a new session for sparsification.

  • reset_session –

    Reset the currently active session to its initial state

LifecycleCallbacks

A class for invoking lifecycle events for the active session

Methods:

  • batch_end –

    Invoke a batch end event for the active session

  • batch_start –

    Invoke a batch start event for the active session

  • calibration_epoch_end –

    Invoke a epoch end event for the active session during calibration. This event

  • calibration_epoch_start –

    Invoke a epoch start event for the active session during calibration. This event

  • event –

    Invoke an event for the active session

  • loss_calculated –

    Invoke a loss calculated event for the active session

  • optim_post_step –

    Invoke an optimizer post-step event for the active session

  • optim_pre_step –

    Invoke an optimizer pre-step event for the active session

  • sequential_epoch_end –

    Invoke a sequential epoch end event for the active session. This event should be

batch_end classmethod 

batch_end(**kwargs) -> ModifiedState

Invoke a batch end event for the active session

Parameters:

  • kwargs

    –

    additional kwargs to pass to the current session's event method

Returns:

  • ModifiedState –

    the modified state of the active session after invoking the event

Source code in llmcompressor/core/session_functions.py 
131
132
133
134
135
136
137
138
139
140
@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

Invoke a batch start event for the active session

Parameters:

  • batch_data

    (Optional[Any], default: None ) –

    the batch data to use for the event

  • kwargs

    –

    additional kwargs to pass to the current session's event method

Returns:

  • ModifiedState –

    the modified state of the active session after invoking the event

Source code in llmcompressor/core/session_functions.py 
87
88
89
90
91
92
93
94
95
96
@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

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

Source code in llmcompressor/core/session_functions.py 
163
164
165
166
167
168
169
170
171
@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

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

Source code in llmcompressor/core/session_functions.py 
142
143
144
145
146
147
148
149
150
@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

Invoke an event for the active session

Parameters:

  • event_type

    (EventType) –

    the event type to invoke

  • kwargs

    –

    additional kwargs to pass to the current session's event method

Returns:

  • ModifiedState –

    the modified state of the active session after invoking the event

Source code in llmcompressor/core/session_functions.py 
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
@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

Invoke a loss calculated event for the active session

Parameters:

  • loss

    (Optional[Any], default: None ) –

    the loss to use for the event

  • kwargs

    –

    additional kwargs to pass to the current session's event method

Returns:

  • ModifiedState –

    the modified state of the active session after invoking the event

Source code in llmcompressor/core/session_functions.py 
 98
 99
100
101
102
103
104
105
106
107
108
109
@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

Invoke an optimizer post-step event for the active session

Parameters:

  • kwargs

    –

    additional kwargs to pass to the current session's event method

Returns:

  • ModifiedState –

    the modified state of the active session after invoking the event

Source code in llmcompressor/core/session_functions.py 
121
122
123
124
125
126
127
128
129
@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

Invoke an optimizer pre-step event for the active session

Parameters:

  • kwargs

    –

    additional kwargs to pass to the current session's event method

Returns:

  • ModifiedState –

    the modified state of the active session after invoking the event

Source code in llmcompressor/core/session_functions.py 
111
112
113
114
115
116
117
118
119
@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(**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

Source code in llmcompressor/core/session_functions.py 
152
153
154
155
156
157
158
159
160
161
@classmethod
def sequential_epoch_end(cls, **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, **kwargs)

active_session 

active_session() -> CompressionSession

Returns:

Source code in llmcompressor/core/session_functions.py 
49
50
51
52
53
54
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]
)

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.

Returns:

Source code in llmcompressor/core/session_functions.py 
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
@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()

Reset the currently active session to its initial state

Source code in llmcompressor/core/session_functions.py 
57
58
59
60
61
62
def reset_session():
    """
    Reset the currently active session to its initial state
    """
    session = active_session()
    session._lifecycle.reset()