Skip to content

LMOperationalMetric metric

Bases: Metric

Base class for LanguageModel runtime-counter metrics.

Subclasses set _phase to one of "inference", "reward", or "optimizer" to read from the corresponding counter set on each bound LM. Counters are populated by the LM based on the active op_scope (contextvar) the trainer sets for each phase.

The metric binds itself automatically to every LanguageModel reachable from the program (and their .fallback chains) when program.compile() is called, and counters are summed across all.

Example:

program.compile(
    metrics=[
        synalinks.metrics.LMOperationalMetric(),
    ],
)
Source code in synalinks/src/metrics/lm_metrics.py 
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
@synalinks_export(
    [
        "synalinks.metrics.LMOperationalMetric",
        "synalinks.LMOperationalMetric",
    ]
)
class LMOperationalMetric(Metric):
    """Base class for `LanguageModel` runtime-counter metrics.

    Subclasses set `_phase` to one of ``"inference"``, ``"reward"``, or
    ``"optimizer"`` to read from the corresponding counter set on each
    bound LM. Counters are populated by the LM based on the active
    ``op_scope`` (contextvar) the trainer sets for each phase.

    The metric binds itself automatically to every `LanguageModel`
    reachable from the program (and their `.fallback` chains) when
    `program.compile()` is called, and counters are summed across all.

    Example:

    ```python
    program.compile(
        metrics=[
            synalinks.metrics.LMOperationalMetric(),
        ],
    )
    ```
    """

    _phase = "inference"

    def __init__(self, name=None):
        super().__init__(name=name)
        self._language_models = []
        self._baselines = {suffix: 0 for suffix in _TRACKED_SUFFIXES}
        self._wall_baseline = 0.0

    @property
    def language_models(self):
        return list(self._language_models)

    def bind_program(self, program):
        self._language_models = _collect_language_models(program)
        self._snapshot()

    def _attr(self, suffix):
        return f"{self._phase}_cumulated_{suffix}"

    def _read(self, suffix):
        attr = self._attr(suffix)
        return sum(getattr(lm, attr, 0) for lm in self._language_models)

    def _snapshot(self):
        for suffix in _TRACKED_SUFFIXES:
            self._baselines[suffix] = self._read(suffix)
        self._wall_baseline = read_phase_wall_clock_s(self._phase)

    def _delta(self, suffix):
        return self._read(suffix) - self._baselines.get(suffix, 0)

    def _wall_clock_delta(self):
        """Wall-clock seconds the trainer spent in this metric's phase since
        the last snapshot. Used as the throughput denominator so concurrent
        (overlapping) calls don't inflate it the way summed `elapsed_s` does.
        """
        return read_phase_wall_clock_s(self._phase) - self._wall_baseline

    def reset_state(self):
        self._snapshot()

    async def update_state(self, *args, **kwargs):
        return

    def result(self):
        raise NotImplementedError

    def get_config(self):
        return {"name": self.name}