Chapter 25 — The transformers Library Internals

1
2
3
4
5
# What AutoModelForCausalLM.from_pretrained() does internally:
# 1. Download config.json → read "model_type": "llama"
# 2. Look up MODEL_FOR_CAUSAL_LM_MAPPING: "llama" → LlamaForCausalLM
# 3. Instantiate LlamaForCausalLM(config)
# 4. Load safetensors weights into the model

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
LlamaForCausalLM(
  (model): LlamaModel(
    (embed_tokens): Embedding(128256, 4096)           # token embeddings
    (layers): ModuleList(
      (0-31): 32 x LlamaDecoderLayer(
        (self_attn): LlamaSdpaAttention(
          (q_proj): Linear(4096, 4096, bias=False)     # Q projection
          (k_proj): Linear(4096, 1024, bias=False)     # K projection (GQA: fewer heads)
          (v_proj): Linear(4096, 1024, bias=False)     # V projection (GQA)
          (o_proj): Linear(4096, 4096, bias=False)     # output projection
          (rotary_emb): LlamaRotaryEmbedding()         # RoPE
        )
        (mlp): LlamaMLP(
          (gate_proj): Linear(4096, 14336, bias=False) # SwiGLU gate
          (up_proj): Linear(4096, 14336, bias=False)   # SwiGLU up
          (down_proj): Linear(14336, 4096, bias=False) # SwiGLU down
        )
        (input_layernorm): LlamaRMSNorm(4096)          # pre-attention norm
        (post_attention_layernorm): LlamaRMSNorm(4096) # pre-FFN norm
      )
    )
    (norm): LlamaRMSNorm(4096)                         # final norm
  )
  (lm_head): Linear(4096, 128256, bias=False)          # maps to vocab logits
)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
# Simplified forward of LlamaForCausalLM
def forward(self, input_ids, attention_mask=None, labels=None, past_key_values=None):
    # 1. Embeddings
    hidden_states = self.model.embed_tokens(input_ids)  # [B, T] → [B, T, d]
    
    # 2. Decoder layers
    for layer in self.model.layers:
        hidden_states = layer(
            hidden_states,
            attention_mask=attention_mask,
            past_key_values=past_key_values,
        )
    
    # 3. Final norm + lm_head
    hidden_states = self.model.norm(hidden_states)       # [B, T, d]
    logits = self.lm_head(hidden_states)                 # [B, T, vocab_size]
    
    # 4. Loss (if labels provided)
    loss = None
    if labels is not None:
        # Shift: predict next token
        shift_logits = logits[..., :-1, :].contiguous()
        shift_labels = labels[..., 1:].contiguous()
        loss = F.cross_entropy(
            shift_logits.view(-1, self.config.vocab_size),
            shift_labels.view(-1),
        )
    
    return CausalLMOutputWithPast(loss=loss, logits=logits,
                                   past_key_values=past_key_values)

1
2
3
4
5
6
7
8
9
10
11
12
13
# First call (prefill): process full prompt
outputs = model(input_ids, use_cache=True)
past_key_values = outputs.past_key_values  # DynamicCache object

# Subsequent calls (decode): process one new token at a time
for _ in range(max_new_tokens):
    outputs = model(
        next_token_id.unsqueeze(0),         # [1, 1] — single token
        past_key_values=past_key_values,    # reuse cached K, V
        use_cache=True,
    )
    past_key_values = outputs.past_key_values  # updated cache
    next_token_id = outputs.logits[:, -1, :].argmax(dim=-1)

1
2
3
4
5
# Inspect cache structure
cache = outputs.past_key_values
print(len(cache))              # 32 (one per layer)
print(cache[0][0].shape)       # K: [B, n_kv_heads, T_cached, head_dim]
print(cache[0][1].shape)       # V: [B, n_kv_heads, T_cached, head_dim]

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
from transformers import PreTrainedModel, PretrainedConfig

class MyConfig(PretrainedConfig):
    model_type = "my_transformer"
    def __init__(self, d_model=1024, n_layers=12, n_heads=16, vocab_size=32000, **kwargs):
        self.d_model = d_model
        self.n_layers = n_layers
        self.n_heads = n_heads
        super().__init__(vocab_size=vocab_size, **kwargs)

class MyModel(PreTrainedModel):
    config_class = MyConfig
    
    def __init__(self, config):
        super().__init__(config)
        self.embed = nn.Embedding(config.vocab_size, config.d_model)
        self.layers = nn.ModuleList([
            TransformerBlock(config) for _ in range(config.n_layers)
        ])
        self.norm = nn.RMSNorm(config.d_model)
        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
    
    def forward(self, input_ids, **kwargs):
        x = self.embed(input_ids)
        for layer in self.layers:
            x = layer(x)
        return self.lm_head(self.norm(x))

# Register and use
MyConfig.register_for_auto_class()
MyModel.register_for_auto_class("AutoModelForCausalLM")

# Now AutoModelForCausalLM.from_pretrained works with your model

1
2
3
4
5
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.1-8B-Instruct",
    attn_implementation="flash_attention_2",  # or "sdpa", "eager"
    torch_dtype=torch.bfloat16,
)