Views
No views yet
nn.Module components:| Component Type | Count | Dimensions |
|---|---|---|
| Attention (self_attn) | 12 | 768d, 12 heads, 64 head_dim |
| Feed-Forward (mlp) | 12 | 768→3072→768 |
| Layer Norm (ln_1, ln_2) | 24 | 768d |
| Full Layers | 12 | 768d |
| Embeddings | 2 | token (50257→768), position (1024→768) |
| Output Head | 1 | 768→50257 |
| Final Norm | 1 | 768d |
1from model_garage.extract.pytorch import PyTorchExtractor
2
3extractor = PyTorchExtractor("gpt2")
4extractor.load_model()
5
6# Extract any component
7attn = extractor.extract_component("self_attention", layer_idx=6)
8ffn = extractor.extract_component("feed_forward", layer_idx=6)
9
10# Test in isolation
11from model_garage.extract.pytorch import ComponentTester
12tester = ComponentTester()
13print(tester.test_attention(attn))pip install model-garage