Skip to content

llmcompressor.observers.helpers

get_observer_token_count(module)

Parse the module and return the number of tokens observed by each module's observer.

Parameters:

Name Type Description Default
module Module

module to parse

 required

Returns:

Type Description
Counter

counter with the number of tokens observed by each observer
Source code in llmcompressor/observers/helpers.py 
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
def get_observer_token_count(module: torch.nn.Module) -> Counter:
    """
    Parse the module and return the number of tokens observed by
    each module's observer.

    :param module: module to parse
    :return: counter with the number of tokens observed by each observer
    """
    token_counts = Counter()
    for name, module in module.named_modules():
        if name.endswith(".input_observer"):
            token_counts[name.replace(".input_observer", "")] = (
                module._num_observed_tokens
            )
    return token_counts