Modules

We can also apply modules in the model’s module tree at any point during computation, even if they are out of order.

Here, we get the hidden states of the last layer like usual. We also chain apply model.transformer.ln_f and model.lm_head in order to “decode” the hidden states into the vocabulary space. Applying softmax and then argmax then transformz the vocabulary space hidden states into tokens that we can decode with the tokenizer.

from nnsight import LanguageModel
import torch

model = LanguageModel("openai-community/gpt2", device_map='auto')

with model.trace('The Eiffel Tower is in the city of') as tracer:

    hidden_states = model.transformer.h[-1].output[0]

    hidden_states = model.lm_head(model.transformer.ln_f(hidden_states)).save()

    tokens = torch.softmax(hidden_states, dim=2).argmax(dim=2).save()

The output looks like:

print(hidden_states)
print(tokens)
print(model.tokenizer.decode(tokens[0]))

tensor([[[ -36.2874,  -35.0114,  -38.0794,  ...,  -40.5164,  -41.3760,
           -34.9193],
         [ -68.8886,  -70.1562,  -71.8408,  ...,  -80.4195,  -78.2552,
           -71.1205],
         [ -82.2951,  -81.6519,  -83.9940,  ...,  -94.4878,  -94.5194,
           -85.6997],
         ...,
         [-113.8675, -111.8628, -113.6634,  ..., -116.7652, -114.8267,
          -112.3621],
         [ -81.8530,  -83.3006,  -91.8192,  ...,  -92.9943,  -89.8382,
           -85.6898],
         [-103.9307, -102.5053, -105.1563,  ..., -109.3099, -110.4195,
          -103.1395]]], device='mps:0', grad_fn=<LinearBackward0>)
tensor([[ 198,   12,  417, 8765,  318,  257,  262, 3504, 7372, 6342]],
       device='mps:0')

-el Tower is a the middle centre Paris