Skip to content

llmcompressor.observers.fusion

Classes:

  • FusionHandler –

    Manages fusion bookkeeping for a single Observer.

FusionHandler 

FusionHandler(observer: 'Observer')

Manages fusion bookkeeping for a single Observer.

Every Observer gets a FusionHandler at init time. When observers are fused (e.g. Q/K/V sharing a global_scale), FusionHandler.fuse() links their handlers together and stores weak refs to the modules whose weights may need lazy observation.

Methods:

  • fuse –

    Link all observers in the iterable with each other for shared global_scale.

  • get_fused_statistics –

    Returns statistics for every observer in the fusion group

  • maybe_delete_statistics –

    Cooperative statistics deletion. Each handler calls this when it

Source code in src/llmcompressor/observers/fusion.py 
27
28
29
30
31
def __init__(self, observer: "Observer"):
    self._observer: "Observer" = observer
    self._module_ref: ref[Module] | None = None
    self._group: list[FusionHandler] = []
    self._deletion_called: bool = False

fuse classmethod 

fuse(
    observers_and_modules: Iterable[
        tuple["Observer", Module]
    ],
) -> None

Link all observers in the iterable with each other for shared global_scale. Sets module weak-references and makes handlers aware of each other.

Parameters:

  • observers_and_modules (Iterable[tuple['Observer', Module]]) –

    iterable of (observer, module) tuples

Source code in src/llmcompressor/observers/fusion.py 
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
@classmethod
def fuse(cls, observers_and_modules: Iterable[tuple["Observer", Module]]) -> None:
    """
    Link all observers in the iterable with each other for shared global_scale.
    Sets module weak-references and makes handlers aware of each other.

    :param observers_and_modules: iterable of (observer, module) tuples
    """
    pairs = list(observers_and_modules)
    handlers = []
    for obs, mod in pairs:
        obs.fusion_handler._module_ref = ref(mod)
        handlers.append(obs.fusion_handler)

    for handler in handlers:
        handler._group = [h for h in handlers if h is not handler]

get_fused_statistics 

get_fused_statistics() -> list[dict[str, torch.Tensor]]

Returns statistics for every observer in the fusion group (including this handler's own observer).

If a fused observer lacks statistics, runs observation on its module's weight via the weak ref stored during fuse().

Returns:

  • list[dict[str, Tensor]] –

    list of {"min_vals": Tensor, "max_vals": Tensor} dicts

Source code in src/llmcompressor/observers/fusion.py 
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
def get_fused_statistics(self) -> list[dict[str, torch.Tensor]]:
    """
    Returns statistics for every observer in the fusion group
    (including this handler's own observer).

    If a fused observer lacks statistics, runs observation on its
    module's weight via the weak ref stored during fuse().

    :return: list of {"min_vals": Tensor, "max_vals": Tensor} dicts
    """
    all_handlers = [self] + self._group

    for handler in all_handlers:
        obs = handler._observer
        if not obs.has_statistics:
            mod = handler.module
            assert mod is not None, _msg
            obs(mod.weight)

    return [h._observer.get_statistics() for h in all_handlers]

maybe_delete_statistics 

maybe_delete_statistics() -> None

Cooperative statistics deletion. Each handler calls this when it is done with fused statistics. When ALL handlers in the group have called it, deletes statistics from every observer in the group and resets the bookkeeping.

For unfused observers, deletes statistics immediately.

Source code in src/llmcompressor/observers/fusion.py 
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
def maybe_delete_statistics(self) -> None:
    """
    Cooperative statistics deletion. Each handler calls this when it
    is done with fused statistics. When ALL handlers in the group
    have called it, deletes statistics from every observer in the
    group and resets the bookkeeping.

    For unfused observers, deletes statistics immediately.
    """
    if not self.is_fused:
        self._delete_observer_statistics(self._observer)
        return

    self._deletion_called = True

    if all(h._deletion_called for h in self._group):
        self._deletion_called = False
        for h in self._group:
            h._deletion_called = False
            self._delete_observer_statistics(h._observer)
        self._delete_observer_statistics(self._observer)