What does mlp.output value correspond to in "bigscience/bloom-7b1" model?

When I run the following code:

from nnsight import LanguageModel
import torch
model = LanguageModel("bigscience/bloom-7b1", dispatch='True')

with torch.no_grad():
    with model.trace() as tracer:
        with tracer.invoke('asdjkgsdjgsdkgjs') as invoker:
            a = model.transformer.h[0].mlp.dense_h_to_4h.output.save()
            b = model.transformer.h[0].mlp.gelu_impl.output.save()
            c = model.transformer.h[0].mlp.dense_4h_to_h.output.save()
            d = model.transformer.h[0].mlp.output.save()

print(torch.equal(a, b))
print(torch.equal(a, c))
print(torch.equal(a, d))
print(torch.equal(b, c))
print(torch.equal(b, d))
print(torch.equal(c, d))

Output:

False
False
False
False
False
False

None of the outputs of the MLP (variables a, b, and c) correspond to the value I get with model.transformer.h[0].mlp.output.save() (variable d).
The MLP for “bigscience/bloom-7b1” is defined as:

(mlp): BloomMLP(
          (dense_h_to_4h): Linear(in_features=4096, out_features=16384, bias=True)
          (gelu_impl): BloomGelu()
          (dense_4h_to_h): Linear(in_features=16384, out_features=4096, bias=True)
        )

What does mlp.output correspond to in this case?

1 Like

Following up on this, I realized model.transformer.h[0].mlp.output.save() is equal to model.transformer.h[0].output[0].save(), instead of the output of the down projection. Is that how it is supposed to be? Any insights on how I should interpret the values coming from the hooks are welcome, thanks!

torch.equal(c, d) == False is because the bloom implementation has a non-module dropout layer after the down projection which also takes care of adding the down projection output to the residual stream. See below:

  • MLP

  • dropout

source: transformers/src/transformers/models/bloom/modeling_bloom.py at main · huggingface/transformers · GitHub

1 Like