# Understanding LLM loss calculation


<!-- WARNING: THIS FILE WAS AUTOGENERATED! DO NOT EDIT! -->

``` python
from transformers import AutoModelForImageTextToText, AutoTokenizer
import torch
import torch.nn.functional as F
```

We will download the smaller Qwen 3.5 model with 0.8B(800M) parameters

``` python
model_id = "Qwen/Qwen3.5-0.8B"
model = AutoModelForImageTextToText.from_pretrained(model_id, dtype=torch.bfloat16, device_map="auto", attn_implementation="eager")
```

    Loading weights:   0%|          | 0/473 [00:00<?, ?it/s]

    Warning: You are sending unauthenticated requests to the HF Hub. Please set a HF_TOKEN to enable higher rate limits and faster downloads.

``` python
tok = AutoTokenizer.from_pretrained(model_id)
# let us tokenize the sentece `The capital of France is` and see what is the output
input_seq = "The capital of France is"
inputs = tok(input_seq, return_tensors="pt").to(model.device)

print(inputs)
```

    {'input_ids': tensor([[ 760, 6511,  314, 9338,  369]], device='cuda:0'), 'attention_mask': tensor([[1, 1, 1, 1, 1]], device='cuda:0')}

Our tokenizer output is a dictionary, with the key `input_ids`, which is
the token id for each token in our sequence, and the `attention_mask`,
which says which positions are real tokens and which are padding. We
will use the mask when passing the sequence to the model and when
deciding which labels contribute to the loss

``` python
with torch.no_grad(): # since we are only doing inference, we don't want to accumulate gradients
    out = model(
        **inputs,
        output_hidden_states=True,
        output_attentions=True,
        return_dict=True,
    )
```

    [transformers] `causal_conv1d_fn` is falling back to its reference PyTorch implementation because `causal_conv1d` is not installed. This is correct but much slower; install `causal_conv1d` for the optimized kernel.
    [transformers] `chunk_gated_delta_rule` is falling back to its reference PyTorch implementation because `flash-linear-attention` is not installed. This is correct but much slower; install `flash-linear-attention` for the optimized kernel.

``` python
print(out.keys())
```

    odict_keys(['logits', 'past_key_values', 'hidden_states', 'attentions'])

``` python
print("Shape of Logits(from Head):", tuple(out.logits.shape), "(Batch, SequenceLen, VocabSize)")            # [B, T, vocab]
print("Total Hidden States, including embedding layer:", len(out.hidden_states), "(Embedding Layer + 24 Hidden Layers)")      # embeddings + 24 layers
print("Shape of Embedding Layer Output:", tuple(out.hidden_states[0].shape), "(Batch, SequenceLen, ModelDimension)")  # [B, T, 1024]
print("Shape of 11th Layer Output:", tuple(out.hidden_states[11].shape), "(Batch, SequenceLen, ModelDimension)")  # [B, T, 1024]

# attentions: one tensor per FULL-attention layer (6 of them), not 24
for idx, attn in enumerate(out.attentions, start=1):
    print(f"Shape of the Attention:{idx} -> ",tuple(attn.shape), "(Batch, NumHeads, SeqLen, SeqLen)")
    # typically [B, num_heads=8, T, T]
```

    Shape of Logits(from Head): (1, 5, 248320) (Batch, SequenceLen, VocabSize)
    Total Hidden States, including embedding layer: 25 (Embedding Layer + 24 Hidden Layers)
    Shape of Embedding Layer Output: (1, 5, 1024) (Batch, SequenceLen, ModelDimension)
    Shape of 11th Layer Output: (1, 5, 1024) (Batch, SequenceLen, ModelDimension)
    Shape of the Attention:1 ->  (1, 8, 5, 5) (Batch, NumHeads, SeqLen, SeqLen)
    Shape of the Attention:2 ->  (1, 8, 5, 5) (Batch, NumHeads, SeqLen, SeqLen)
    Shape of the Attention:3 ->  (1, 8, 5, 5) (Batch, NumHeads, SeqLen, SeqLen)
    Shape of the Attention:4 ->  (1, 8, 5, 5) (Batch, NumHeads, SeqLen, SeqLen)
    Shape of the Attention:5 ->  (1, 8, 5, 5) (Batch, NumHeads, SeqLen, SeqLen)
    Shape of the Attention:6 ->  (1, 8, 5, 5) (Batch, NumHeads, SeqLen, SeqLen)

Each position has a logit (a score) for every token in the vocabulary.
Applying softmax would turn these scores into probabilities. Since
softmax keeps the ordering, `argmax` on the logits gives the same token
as `argmax` on the probabilities. So we can directly do `argmax` on
logits and get the predicted token for each position

``` python
ids = out.logits.argmax(dim=-1) # [1, 5]
id_list = ids[0].tolist()
print(id_list)
```

    [2614, 314, 279, 369, 11751]

``` python
# Convert the tokens to string
input_tokens = tok.convert_ids_to_tokens(inputs["input_ids"][0].tolist())
output_tokens = tok.convert_ids_to_tokens(id_list)
print(id_list)
print(output_tokens)
```

    [2614, 314, 279, 369, 11751]
    ['Ġfollowing', 'Ġof', 'Ġthe', 'Ġis', 'ĠParis']

> Note: The `Ġ` character represents a leading space in these tokenizer
> displays. `Ċ` represents a newline.

There is big mis-conception that the model predict the next token given
the `k` tokens as input. In actual the model predicts `k` tokens, each
output token corresponding to the input token. In practice we don’t care
for other tokens and only care for the final token.

During autoregressive generation, we use the final position to choose
one next token, append it to the sequence, and run the model again.
During training, the causal mask lets us calculate all of these
next-token predictions in parallel, and each position is compared with
the token immediately after it

Below we will map each input token and what the model predicted for that
token.

``` python
# input to output mapping
print("Input:", input_seq)
print()
print("position | input token | actual t+1 | head 1 prediction")
print("-" * 86)

for t, input_token in enumerate(input_tokens):
    actual_t1 = input_tokens[t + 1] if t + 1 < len(input_tokens) else "-"
    predicted_t1 = output_tokens[t] if t < len(output_tokens) else None

    print(f"{t:8d} | {input_token:11s} | {actual_t1:11s} | {predicted_t1:18s}")
```

    Input: The capital of France is

    position | input token | actual t+1 | head 1 prediction
    --------------------------------------------------------------------------------------
           0 | The         | Ġcapital    | Ġfollowing        
           1 | Ġcapital    | Ġof         | Ġof               
           2 | Ġof         | ĠFrance     | Ġthe              
           3 | ĠFrance     | Ġis         | Ġis               
           4 | Ġis         | -           | ĠParis            

Here we can see, for the input of `The` the model actually predicted
`following`, whereas our actual second token was `capital`.\
Similarly for the second token, the predicted token was `of` and in
actual too it was `of`.

# Train Model

For a single sequence for a single step

``` python
model.train()
sample_train_seq = "The capital of India is New Delhi."
enc = tok(
    sample_train_seq,
    return_tensors="pt"
).to(model.device)

input_ids = enc["input_ids"]
attn = enc["attention_mask"]
```

``` python
labels = input_ids.clone()

print("tokens: ", tok.convert_ids_to_tokens(input_ids[0].tolist()))
print("ids:    ", input_ids[0].tolist())
print("mask:   ", attn[0].tolist())
print("labels: ", labels[0].tolist())
```

    tokens:  ['The', 'Ġcapital', 'Ġof', 'ĠIndia', 'Ġis', 'ĠNew', 'ĠDelhi', '.']
    ids:     [760, 6511, 314, 6535, 369, 1478, 21318, 13]
    mask:    [1, 1, 1, 1, 1, 1, 1, 1]
    labels:  [760, 6511, 314, 6535, 369, 1478, 21318, 13]

``` python
out = model(**enc)
logits = out.logits
print(logits.shape) # 1, SeqLen, VocabSize
```

    torch.Size([1, 8, 248320])

``` python
# lets see the output string
ids = out.logits.clone().detach().argmax(dim=-1) # [1, 8]
id_list = ids[0].tolist()
print(id_list)
```

    [2614, 314, 279, 369, 279, 21318, 13, 198]

``` python
# 3. per-token strings
input_tokens = tok.convert_ids_to_tokens(input_ids[0].tolist())
output_tokens = tok.convert_ids_to_tokens(id_list)
print(id_list)
print(output_tokens)
```

    [2614, 314, 279, 369, 279, 21318, 13, 198]
    ['Ġfollowing', 'Ġof', 'Ġthe', 'Ġis', 'Ġthe', 'ĠDelhi', '.', 'Ċ']

``` python
# input to output mapping
print("Input:", sample_train_seq)
print()
print("position | input token | actual t+1 | head 1 prediction")
print("-" * 86)

for t, input_token in enumerate(input_tokens):
    actual_t1 = input_tokens[t + 1] if t + 1 < len(input_tokens) else "-"
    predicted_t1 = output_tokens[t] if t < len(output_tokens) else None

    print(f"{t:8d} | {input_token:11s} | {actual_t1:11s} | {predicted_t1:18s}")
```

    Input: The capital of India is New Delhi.

    position | input token | actual t+1 | head 1 prediction
    --------------------------------------------------------------------------------------
           0 | The         | Ġcapital    | Ġfollowing        
           1 | Ġcapital    | Ġof         | Ġof               
           2 | Ġof         | ĠIndia      | Ġthe              
           3 | ĠIndia      | Ġis         | Ġis               
           4 | Ġis         | ĠNew        | Ġthe              
           5 | ĠNew        | ĠDelhi      | ĠDelhi            
           6 | ĠDelhi      | .           | .                 
           7 | .           | -           | Ċ                 

## Causal Shift

For token position `t` in logits, the predicted label lies at position
`t+1` in labels. For example for calculating loss, we have to compare
the token as position `0` of logits to position `1` of labels.

The new token predicted (`8`th token in this case), we actually don’t
know the ground truth for that. So we will ignore the very last token.

For a sequence `[x0, x1, ..., x7]`, the alignment is:

``` text
logits:  z0   z1   z2   z3   z4   z5   z6   [drop z7]
targets: x1   x2   x3   x4   x5   x6   x7
```

``` python
# considering 0th to t-1 tokens for logits
new_logits = logits[:, :-1, :] # [All_Batches, All_Token_But_Last, All_Vocab] [1,7,V]
new_labels = labels[:, 1:] # [All_Bathes, Second_Token_To_End] [1,7]


loss = F.cross_entropy(
    new_logits.reshape(-1, new_logits.size(-1)),  # [7, V]
    new_labels.reshape(-1),                         # [7]
    ignore_index=-100,
)
print("loss:", loss.item())
```

    loss: 2.65625

``` python
# one step
opt = torch.optim.AdamW(model.parameters(), lr=1e-5)
opt.zero_grad()
loss.backward()
opt.step()
```

## Let’s try with padding now

Since in actual we generally have multiple sequences in a batch, and
they all would be different length, so we generally pad them. Our model
support a context window of `262,144` we can have our sequences upto
that length, but for sake of easy demonstration and visualization, let’s
fix our max length to 16

There are two separate ideas here: `attention_mask` tells the model
which input positions are real, while the `-100` labels tell cross
entropy which target positions to ignore. Padding can still produce
logits, but those logits should not contribute to the loss

``` python
model.train()
sample_train_seq = "The capital of India is New Delhi."
enc = tok(
    sample_train_seq,
    return_tensors="pt",
    padding="max_length",
    max_length=16,
    truncation=True,
).to(model.device)

input_ids = enc["input_ids"] # [1, 16]
attn = enc["attention_mask"] # 1 = real token, 0 = pad
input_ids, attn
```

    (tensor([[   760,   6511,    314,   6535,    369,   1478,  21318,     13, 248044,
              248044, 248044, 248044, 248044, 248044, 248044, 248044]],
            device='cuda:0'),
     tensor([[1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0]], device='cuda:0'))

``` python
labels = input_ids.clone()
labels[attn == 0] = -100

print("tokens: ", tok.convert_ids_to_tokens(input_ids[0].tolist()))
print("ids:    ", input_ids[0].tolist())
print("mask:   ", attn[0].tolist())
print("labels: ", labels[0].tolist())
```

    tokens:  ['The', 'Ġcapital', 'Ġof', 'ĠIndia', 'Ġis', 'ĠNew', 'ĠDelhi', '.', '<|endoftext|>', '<|endoftext|>', '<|endoftext|>', '<|endoftext|>', '<|endoftext|>', '<|endoftext|>', '<|endoftext|>', '<|endoftext|>']
    ids:     [760, 6511, 314, 6535, 369, 1478, 21318, 13, 248044, 248044, 248044, 248044, 248044, 248044, 248044, 248044]
    mask:    [1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0]
    labels:  [760, 6511, 314, 6535, 369, 1478, 21318, 13, -100, -100, -100, -100, -100, -100, -100, -100]

``` python
out = model(**enc)
logits = out.logits
print(logits.shape) # 1, SeqLen, VocabSize
```

    torch.Size([1, 16, 248320])

``` python
# lets see the output string
ids = out.logits.clone().detach().argmax(dim=-1) # [1, 16]
id_list = ids[0].tolist()
print(id_list)
```

    [2614, 314, 279, 369, 279, 21318, 13, 198, 271, 271, 271, 271, 271, 271, 271, 271]

``` python
# per-token strings
input_tokens = tok.convert_ids_to_tokens(input_ids[0].tolist())
output_tokens = tok.convert_ids_to_tokens(id_list)
print(id_list)
print(output_tokens)
```

    [2614, 314, 279, 369, 279, 21318, 13, 198, 271, 271, 271, 271, 271, 271, 271, 271]
    ['Ġfollowing', 'Ġof', 'Ġthe', 'Ġis', 'Ġthe', 'ĠDelhi', '.', 'Ċ', 'ĊĊ', 'ĊĊ', 'ĊĊ', 'ĊĊ', 'ĊĊ', 'ĊĊ', 'ĊĊ', 'ĊĊ']

``` python
# input to output mapping
print("Input:", sample_train_seq)
print()
print("position |   input token    |    actual t+1    |  head 1 prediction")
print("-" * 86)

for t, input_token in enumerate(input_tokens):
    actual_t1 = input_tokens[t + 1] if t + 1 < len(input_tokens) else "-"
    predicted_t1 = output_tokens[t] if t < len(output_tokens) else None

    print(f"{t:8d} | {input_token:16s} | {actual_t1:16s} | {predicted_t1:18s}")
```

    Input: The capital of India is New Delhi.

    position |   input token    |    actual t+1    |  head 1 prediction
    --------------------------------------------------------------------------------------
           0 | The              | Ġcapital         | Ġfollowing        
           1 | Ġcapital         | Ġof              | Ġof               
           2 | Ġof              | ĠIndia           | Ġthe              
           3 | ĠIndia           | Ġis              | Ġis               
           4 | Ġis              | ĠNew             | Ġthe              
           5 | ĠNew             | ĠDelhi           | ĠDelhi            
           6 | ĠDelhi           | .                | .                 
           7 | .                | <|endoftext|>    | Ċ                 
           8 | <|endoftext|>    | <|endoftext|>    | ĊĊ                
           9 | <|endoftext|>    | <|endoftext|>    | ĊĊ                
          10 | <|endoftext|>    | <|endoftext|>    | ĊĊ                
          11 | <|endoftext|>    | <|endoftext|>    | ĊĊ                
          12 | <|endoftext|>    | <|endoftext|>    | ĊĊ                
          13 | <|endoftext|>    | <|endoftext|>    | ĊĊ                
          14 | <|endoftext|>    | <|endoftext|>    | ĊĊ                
          15 | <|endoftext|>    | -                | ĊĊ                

You could see two things above, which are interesting, one since we have
already trained on this sequence above, the model is doing better
predictions on each token level.

Here is the previous prediction

<table>
<thead>
<tr>
<th>Position</th>
<th>Input Token</th>
<th>Actual t+1</th>
<th>Head 1 Prediction</th>
</tr>
</thead>
<tbody>
<tr>
<td><strong>0</strong></td>
<td>The</td>
<td>Ġcapital</td>
<td>Ġfollowing</td>
</tr>
<tr>
<td><strong>1</strong></td>
<td>Ġcapital</td>
<td>Ġof</td>
<td>Ġof</td>
</tr>
<tr>
<td><strong>2</strong></td>
<td>Ġof</td>
<td>ĠIndia</td>
<td>Ġthe</td>
</tr>
<tr>
<td><strong>3</strong></td>
<td>ĠIndia</td>
<td>Ġis</td>
<td>Ġis</td>
</tr>
<tr>
<td><strong>4</strong></td>
<td>Ġis</td>
<td>ĠNew</td>
<td>Ġthe</td>
</tr>
<tr>
<td><strong>5</strong></td>
<td>ĠNew</td>
<td>ĠDelhi</td>
<td>ĠDelhi</td>
</tr>
<tr>
<td><strong>6</strong></td>
<td>ĠDelhi</td>
<td>.</td>
<td>.</td>
</tr>
<tr>
<td><strong>7</strong></td>
<td>.</td>
<td>-</td>
<td>Ċ</td>
</tr>
</tbody>
</table>

And now the prediction is

<table>
<thead>
<tr>
<th>Position</th>
<th>Input Token</th>
<th>Actual t+1</th>
<th>Head 1 Prediction</th>
</tr>
</thead>
<tbody>
<tr>
<td><strong>0</strong></td>
<td>The</td>
<td>Ġcapital</td>
<td>Ġfollowing</td>
</tr>
<tr>
<td><strong>1</strong></td>
<td>Ġcapital</td>
<td>Ġof</td>
<td>Ġof</td>
</tr>
<tr>
<td><strong>2</strong></td>
<td>Ġof</td>
<td>ĠIndia</td>
<td>ĠIndia</td>
</tr>
<tr>
<td><strong>3</strong></td>
<td>ĠIndia</td>
<td>Ġis</td>
<td>Ġis</td>
</tr>
<tr>
<td><strong>4</strong></td>
<td>Ġis</td>
<td>ĠNew</td>
<td>ĠNew</td>
</tr>
<tr>
<td><strong>5</strong></td>
<td>ĠNew</td>
<td>ĠDelhi</td>
<td>ĠDelhi</td>
</tr>
<tr>
<td><strong>6</strong></td>
<td>ĠDelhi</td>
<td>.</td>
<td>.</td>
</tr>
<tr>
<td><strong>7</strong></td>
<td>.</td>
<td>&lt;|endoftext|&gt;</td>
<td>Ċ</td>
</tr>
</tbody>
</table>

In the previous case for the input token `of` it was predicting `the`,
but after the previous one step training, it is predicting `India`, and
similarly for the input token of `is` it was predicting `the` and now
it’s predicting `New`. So we have overfit the model fairly well.

Second, the model is predicting `\n\n` (`Ċ` is the token for `\n`) for
each `<|endoftext|>` token. We don’t actually care, since we won’t be
training on that, as out attention mask if you see above have value `0`
for those, and we have marked the labels as `-100` for them.

> Why `-100`? \> We can then ask the cross entropy loss to ignore the
> labels which are `-100` and ask it to not calculate any loss for
> those.

``` python
# considering 0th to t-1 tokens for logits
new_logits = logits[:, :-1, :] # [All_Batches, All_Token_But_Last, All_Vocab] [1,15,V]
new_labels = labels[:, 1:] # [All_Bathes, Second_Token_To_End] [1,15]

# First let us see the loss for each token
loss = F.cross_entropy(
    new_logits.reshape(-1, new_logits.size(-1)),  # [15, V]
    new_labels.reshape(-1),                         # [15]
    ignore_index=-100, # ignore padding
    reduction="none" # don't average/sum the loss
)
print("loss:", loss)
```

    loss: tensor([9.3750, 0.5781, 3.9219, 0.6328, 3.4688, 0.0334, 0.6328, 0.0000, 0.0000,
            0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000], device='cuda:0',
           dtype=torch.bfloat16, grad_fn=<NllLossBackward0>)

As you can see for all the tokens which were padded, we are getting loss
as 0.

Now, lets do the actual loss calculation, which is mean of all the loss
values but ignoring the values which have label as -100

``` python
# considering 0th to t-1 tokens for logits
new_logits = logits[:, :-1, :] # [All_Batches, All_Token_But_Last, All_Vocab] [1,15,V]
new_labels = labels[:, 1:] # [All_Bathes, Second_Token_To_End] [1,15]

# First let us see the loss for each token
loss = F.cross_entropy(
    new_logits.reshape(-1, new_logits.size(-1)),  # [15, V]
    new_labels.reshape(-1),                         # [15]
    ignore_index=-100, # ignore padding
)
print("loss:", loss)
```

    loss: tensor(2.6562, device='cuda:0', dtype=torch.bfloat16,
           grad_fn=<NllLossBackward0>)

``` python
# one step
opt = torch.optim.AdamW(model.parameters(), lr=1e-5)
opt.zero_grad()
loss.backward()
opt.step()
```

## Train with two documents packed in the same sequence.

As mentioned earlier, our context window is around ~262k tokens, and
sometimes our documents may not be that long, we say one way is to do
padding, but that’s wasted compute, since we are not calculating loss
for those tokens even though model is doing some work on those. So what
if we combine multiple documents in the same sequence. But this will
cause an issue, how can let the model know that they are two separate
sentences.

We can simply add a boundary, like padding or an EOS (End of Sequence)
token. In this Qwen tokenizer, `tok.eos_token` is `<|im_end|>`, while
padding uses `<|endoftext|>`. We will see what will be the loss when we
ask the model to predict the boundary and what happens when we mask the
cross-document transition.

``` python
seq_1 = "The capital of India is New Delhi."
seq_2 = "Taj Mahal is NOT situated in New Delhi."

# Adding EOS token as boundary between the two sequences

packed_training_text = seq_1 + tok.eos_token + seq_2 + tok.eos_token

enc = tok(
    packed_training_text,
    return_tensors="pt",
    padding="max_length",
    max_length=24,
    truncation=True,
).to(model.device)

input_ids = enc["input_ids"] # [1, 24]
attn = enc["attention_mask"] # 1 = real token, 0 = pad
input_ids, attn
```

    (tensor([[   760,   6511,    314,   6535,    369,   1478,  21318,     13, 248046,
                  51,   1573,  15695,    278,    369,   4045,  29099,    303,   1478,
               21318,     13, 248046, 248044, 248044, 248044]], device='cuda:0'),
     tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0]],
            device='cuda:0'))

You can see above `248046` corresponds to `<|im_end|>` and `248044`
corresponds to `<|endoftext|>`

``` python
labels_including_boundary = input_ids.clone()
# attn == 0 is only true for padded tokens, since for boundary token the attention mask is 1
labels_including_boundary[attn == 0] = -100

print("tokens: ", tok.convert_ids_to_tokens(input_ids[0].tolist()))
print("ids:    ", input_ids[0].tolist())
print("mask:   ", attn[0].tolist())
print("labels with the boundary also masked: ", labels_including_boundary[0].tolist())
```

    tokens:  ['The', 'Ġcapital', 'Ġof', 'ĠIndia', 'Ġis', 'ĠNew', 'ĠDelhi', '.', '<|im_end|>', 'T', 'aj', 'ĠMah', 'al', 'Ġis', 'ĠNOT', 'Ġsituated', 'Ġin', 'ĠNew', 'ĠDelhi', '.', '<|im_end|>', '<|endoftext|>', '<|endoftext|>', '<|endoftext|>']
    ids:     [760, 6511, 314, 6535, 369, 1478, 21318, 13, 248046, 51, 1573, 15695, 278, 369, 4045, 29099, 303, 1478, 21318, 13, 248046, 248044, 248044, 248044]
    mask:    [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0]
    labels with the boundary also masked:  [760, 6511, 314, 6535, 369, 1478, 21318, 13, 248046, 51, 1573, 15695, 278, 369, 4045, 29099, 303, 1478, 21318, 13, 248046, -100, -100, -100]

``` python
out = model(**enc)
logits = out.logits
print(logits.shape) # 1, SeqLen, VocabSize
```

    torch.Size([1, 24, 248320])

``` python
# lets see the output string
ids = out.logits.clone().detach().argmax(dim=-1) # [1, 24]
id_list = ids[0].tolist()
print(id_list)
```

    [2614, 314, 279, 369, 279, 21318, 13, 198, 198, 198, 1524, 278, 369, 264, 303, 303, 279, 21318, 13, 198, 198, 51, 51, 51]

``` python
# per-token strings
input_tokens = tok.convert_ids_to_tokens(input_ids[0].tolist())
output_tokens = tok.convert_ids_to_tokens(id_list)
print(id_list)
print(output_tokens)
```

    [2614, 314, 279, 369, 279, 21318, 13, 198, 198, 198, 1524, 278, 369, 264, 303, 303, 279, 21318, 13, 198, 198, 51, 51, 51]
    ['Ġfollowing', 'Ġof', 'Ġthe', 'Ġis', 'Ġthe', 'ĠDelhi', '.', 'Ċ', 'Ċ', 'Ċ', 'ik', 'al', 'Ġis', 'Ġa', 'Ġin', 'Ġin', 'Ġthe', 'ĠDelhi', '.', 'Ċ', 'Ċ', 'T', 'T', 'T']

``` python
# input to output mapping
print("Input:", packed_training_text)
print()
print("position |   input token    |    actual t+1    |  head 1 prediction")
print("-" * 86)

for t, input_token in enumerate(input_tokens):
    actual_t1 = input_tokens[t + 1] if t + 1 < len(input_tokens) else "-"
    predicted_t1 = output_tokens[t] if t < len(output_tokens) else None

    print(f"{t:8d} | {input_token:16s} | {actual_t1:16s} | {predicted_t1:18s}")
```

    Input: The capital of India is New Delhi.<|im_end|>Taj Mahal is NOT situated in New Delhi.<|im_end|>

    position |   input token    |    actual t+1    |  head 1 prediction
    --------------------------------------------------------------------------------------
           0 | The              | Ġcapital         | Ġfollowing        
           1 | Ġcapital         | Ġof              | Ġof               
           2 | Ġof              | ĠIndia           | Ġthe              
           3 | ĠIndia           | Ġis              | Ġis               
           4 | Ġis              | ĠNew             | Ġthe              
           5 | ĠNew             | ĠDelhi           | ĠDelhi            
           6 | ĠDelhi           | .                | .                 
           7 | .                | <|im_end|>       | Ċ                 
           8 | <|im_end|>       | T                | Ċ                 
           9 | T                | aj               | Ċ                 
          10 | aj               | ĠMah             | ik                
          11 | ĠMah             | al               | al                
          12 | al               | Ġis              | Ġis               
          13 | Ġis              | ĠNOT             | Ġa                
          14 | ĠNOT             | Ġsituated        | Ġin               
          15 | Ġsituated        | Ġin              | Ġin               
          16 | Ġin              | ĠNew             | Ġthe              
          17 | ĠNew             | ĠDelhi           | ĠDelhi            
          18 | ĠDelhi           | .                | .                 
          19 | .                | <|im_end|>       | Ċ                 
          20 | <|im_end|>       | <|endoftext|>    | Ċ                 
          21 | <|endoftext|>    | <|endoftext|>    | T                 
          22 | <|endoftext|>    | <|endoftext|>    | T                 
          23 | <|endoftext|>    | -                | T                 

``` python
# considering 0th to t-1 tokens for logits
new_logits = logits[:, :-1, :] # [All_Batches, All_Token_But_Last, All_Vocab] [1,23,V]
new_labels = labels_including_boundary[:, 1:] # [All_Bathes, Second_Token_To_End] [1,23]

# First let us see the loss for each token
loss = F.cross_entropy(
    new_logits.reshape(-1, new_logits.size(-1)),  # [23, V]
    new_labels.reshape(-1),                         # [23]
    ignore_index=-100, # ignore padding
    reduction="none" # don't average/sum the loss
)
print("loss:", loss.tolist())
print("calculated mean loss over entire sequence: ", loss.mean().item())


mean_loss = F.cross_entropy(
    new_logits.reshape(-1, new_logits.size(-1)),  # [23, V]
    new_labels.reshape(-1),                         # [23]
    ignore_index=-100, # ignore padding
    reduction="mean" # don't average/sum the loss
)
print("actual mean loss with padded tokens ignored in calculation:", mean_loss.item())
```

    loss: [9.375, 0.57421875, 3.90625, 0.6484375, 3.53125, 0.0341796875, 0.6328125, 4.71875, 27.875, 8.9375, 2.671875, 0.08935546875, 1.1015625, 7.28125, 5.75, 0.388671875, 1.8984375, 0.00927734375, 0.1533203125, 1.5, 0.0, 0.0, 0.0]
    calculated mean loss over entire sequence:  3.53125
    actual mean loss with padded tokens ignored in calculation: 4.0625

### Now let’s also exclude the boundary token in loss calculation

The EOS at the end of the first document is still a valid target. The
position we mask is the first token of the second document, because the
hidden state at the first document’s EOS should not be trained to
predict an unrelated document’s beginning.

``` python
seq_1_ids = tok(
    seq_1,
    add_special_tokens=False,
)["input_ids"]

# The EOS after seq_1 is at index len(seq_1_ids).
# At this position, the model predicts the first token of seq_2.
boundary_index = len(seq_1_ids)

masked_labels = new_labels.clone()
# make the position at which model should predict end of token, as -100, so we can ignore it in loss calculation
masked_labels[0, boundary_index] = -100

# Case 2: boundary transition does not contribute to loss.
loss_with_boundary_mask = F.cross_entropy(
    new_logits.reshape(-1, new_logits.size(-1)),
    masked_labels.reshape(-1),
    ignore_index=-100,
    reduction="none" # don't average/sum the loss
)
print("loss:", loss_with_boundary_mask.tolist())
print("mean over all positions, including ignored zeros: ", loss_with_boundary_mask.mean().item())


mean_loss = F.cross_entropy(
    new_logits.reshape(-1, new_logits.size(-1)),  # [23, V]
    masked_labels.reshape(-1),                         # [23]
    ignore_index=-100, # ignore padding
    reduction="mean" # don't average/sum the loss
)
print("actual mean loss with padded tokens ignored in calculation:", mean_loss.item())
```

    loss: [9.375, 0.57421875, 3.90625, 0.6484375, 3.53125, 0.0341796875, 0.6328125, 4.71875, 0.0, 8.9375, 2.671875, 0.08935546875, 1.1015625, 7.28125, 5.75, 0.388671875, 1.8984375, 0.00927734375, 0.1533203125, 1.5, 0.0, 0.0, 0.0]
    mean over all positions, including ignored zeros:  2.3125
    actual mean loss with padded tokens ignored in calculation: 2.796875

The `reduction="none"` output contains `0` for every `-100` label.
Calling `.mean()` on that tensor divides by all positions, including
ignored padding and the ignored boundary transition. `reduction="mean"`
uses only valid targets in its denominator.

In this example, the cross-document transition has a large loss. Masking
that transition removes it from the average and reduces the mean loss
from about `3.04` to about `1.76`. The model still learns to predict the
EOS at the end of the first document but not the new token of document 2
after EOS

## Computing Perplexity for Untrianed model and our Trained Model

Perplexity is `exp(mean_cross-entropy)`. For a random model, it is
equally likely to predict any token from the vocab, so the loss should
be roughly around `-ln(1/voacb_size)` ~ `12.422473516`. Where as the
perplexity should be somehwere close to the vocab size itself

#### Randomly initialized model

``` python
very_very_long_sample_train_seq = "Elephants are the largest living land animals. Three living species are currently recognised: the African bush elephant (Loxodonta africana), the African forest elephant (L. cyclotis), and the Asian elephant (Elephas maximus). They are the only surviving members of the family Elephantidae and the order Proboscidea; extinct relatives include mammoths and mastodons. Distinctive features of elephants include a long proboscis called a trunk, tusks, large ear flaps, pillar-like legs, and tough but sensitive grey skin. The trunk is prehensile, bringing food and water to the mouth and grasping objects. Tusks, which are derived from the incisor teeth, serve both as weapons and as tools for moving objects and digging. The large ear flaps assist in maintaining a constant body temperature as well as in communication. African elephants have larger ears and concave backs, whereas Asian elephants have smaller ears and convex or level backs. Elephants are scattered throughout sub-Saharan Africa, South Asia, and Southeast Asia. They live in different habitats, including savannahs, forests, deserts, and marshes. They are herbivorous, and they stay near water when it is accessible. They are considered to be keystone species, due to their impact on their environments. Elephants have a fission–fusion society, in which multiple family groups come together to socialise. Females (cows) tend to live in family groups, which can consist of one female with its calves or several related females with offspring. The leader of a female group, usually the oldest cow, is known as the matriarch. Males (bulls) leave their family groups when they reach puberty and may live alone or with other males. Adult bulls mostly interact with family groups when looking for a mate. They enter a state of increased testosterone and aggression known as musth, which helps them gain dominance over other males as well as reproductive success. Calves are the centre of attention in their family groups and rely on their mothers for as long as three years. Elephants can live up to 70 years in the wild. They communicate by touch, sight, smell, and sound; elephants use infrasound and seismic communication over long distances. Elephant intelligence has been compared with that of primates and cetaceans. They appear to have self-awareness, and possibly show concern for dying and dead individuals of their kind. African bush elephants and Asian elephants are listed as endangered and African forest elephants as critically endangered on the IUCN Red Lists. One of the biggest threats to elephant populations is the ivory trade, as the animals are poached for their ivory tusks. Other threats to wild elephants include habitat destruction and conflicts with local people. Elephants are used as working animals in Asia. In the past, they were used in war; today, they are often controversially put on display in zoos, or employed for entertainment in circuses. Elephants have an iconic status in human culture and have been widely featured in art, folklore, religion, literature, and popular culture."
enc = tok(very_very_long_sample_train_seq, return_tensors="pt").to(model.device)

input_ids = enc["input_ids"]
attn = enc["attention_mask"]

# so mask padding using attention_mask, not token also can IDs.
labels = input_ids.clone()
labels[attn == 0] = -100
```

``` python
import math

# Create a fresh randomly initialized model.
untrained_model = AutoModelForImageTextToText.from_config(model.config)
untrained_model = untrained_model.to(device=model.device, dtype=model.dtype)
untrained_model.eval()

with torch.no_grad():
    out = untrained_model(input_ids=input_ids, attention_mask=attn)

logits = out.logits

# Causal next-token shift.
new_logits = logits[:, :-1, :]
new_labels = labels[:, 1:]

# calculate per token loss
loss = F.cross_entropy(
    # converting to float to avoid bfloat16 precision loss
    new_logits.float().reshape(-1, new_logits.size(-1)),
    new_labels.reshape(-1),
    ignore_index=-100,
    reduction='none'
)
print("Per token loss:", loss.tolist()[:5], "...")

# Calculate average cross-entropy over valid tokens.
loss = F.cross_entropy(
    # converting to float to avoid bfloat16 precision loss
    new_logits.float().reshape(-1, new_logits.size(-1)),
    new_labels.reshape(-1),
    ignore_index=-100,
)

perplexity = torch.exp(loss)
vocab_size = new_logits.size(-1)

print("vocabulary size:", vocab_size)
print("expected uniform loss:", math.log(vocab_size))
print("loss:", loss.item())
print("perplexity:", perplexity.item())
print("expected random perplexity:", vocab_size)
```

    Per token loss: [12.649822235107422, 12.856464385986328, 11.195772171020508, 12.947657585144043, 12.372076034545898] ...
    vocabulary size: 248320
    expected uniform loss: 12.422473515976991
    loss: 12.59212589263916
    perplexity: 294232.59375
    expected random perplexity: 248320

#### Already trained Model

``` python
model.eval()

with torch.no_grad():
    out = model(input_ids=input_ids, attention_mask=attn)

logits = out.logits

# Causal next-token shift.
new_logits = logits[:, :-1, :]
new_labels = labels[:, 1:]

# calculate per token loss
loss = F.cross_entropy(
    # converting to float to avoid bfloat16 precision loss
    new_logits.float().reshape(-1, new_logits.size(-1)),
    new_labels.reshape(-1),
    ignore_index=-100,
    reduction='none'
)
print("Per token loss:", loss.tolist()[:5], "...")

# Calculate average cross-entropy over valid tokens.
loss = F.cross_entropy(
    # converting to float to avoid bfloat16 precision loss
    new_logits.float().reshape(-1, new_logits.size(-1)),
    new_labels.reshape(-1),
    ignore_index=-100,
)

perplexity = torch.exp(loss)
vocab_size = new_logits.size(-1)

print("loss:", loss.item())
print("perplexity:", perplexity.item())
```

    Per token loss: [6.58029842376709, 0.8311704397201538, 2.2088489532470703, 0.7635747194290161, 3.1106820106506348] ...
    loss: 2.319715738296509
    perplexity: 10.172781944274902

## Tying Weights

The input embedding table turns token ids into hidden vectors, and the
language-model head turns hidden vectors back into scores for all `V`
tokens. With tied weights, these two operations reuse the same `V x D`
matrix instead of storing two copies. So one single matrix is expected
to do the work of both converting tokens to embeddings and take hidden
dimension vectors (same size as embedding) and convert them to logits
over vocab size. Not ideal but in case of small models, these two
matrices are very huge, as we will see below, hence it make sense to tie
the weights

``` python
input_embeddings = model.get_input_embeddings()
output_embeddings = model.get_output_embeddings()

input_weight = input_embeddings.weight
output_weight = output_embeddings.weight

weights_are_tied = (
    input_weight is output_weight
    or input_weight.data_ptr() == output_weight.data_ptr()
)
# the config can also mention about tying the weights.
config_ties_weights = getattr(model.config, "tie_word_embeddings", None)

# calculate the params
input_embedding_params = sum(p.numel() for p in input_embeddings.parameters())
lm_head_params = sum(p.numel() for p in output_embeddings.parameters())
total_params = sum(p.numel() for p in model.parameters())

# A tied head reuses the input embedding parameters. An untied head adds a
# separate output matrix.
if weights_are_tied:
    tied_total_params = total_params
    untied_total_params = total_params + lm_head_params
else:
    tied_total_params = total_params - lm_head_params
    untied_total_params = total_params

print("config.tie_word_embeddings:", config_ties_weights)
print("weights actually tied:      ", weights_are_tied)
print()
print("input embedding params:     ", f"{input_embedding_params:,}")
print("LM head params if separate:  ", f"{lm_head_params:,}")
print()
print("Total params when tied:           ", f"{tied_total_params:,}")
print("Total params, if output matrix is separate:         ", f"{untied_total_params:,}")
print()
print("extra params when untied:    ", f"{untied_total_params - tied_total_params:,}")
```

    config.tie_word_embeddings: True
    weights actually tied:       True

    input embedding params:      254,279,680
    LM head params if separate:   254,279,680

    Total params when tied:            852,985,920
    Total params, if output matrix is separate:          1,107,265,600

    extra params when untied:     254,279,680

## Chunk Loss Calculation

Until now for the calculation of cross entropy loss, we passed the
complete tensor to cross-entropy loss, where `F.cross_entropy` takes in
`[num_tokens, vocab_size]`, and in our case the vocab size is quite big
~262k, hence the memory used is huge when the `num_tokens` increases. To
avoid it we can simply calculate the cross entropy loss on chunks of
tokens instead of full sequence length at once. This is equivalent in
calculations.

`[num_tokens, vocab_size]` -\> `k * [chunk_size, vocab_size]`

Now since, we calculate loss for smaller chunks, it may take longer
time, but the peak memory is much smaller.

You can restart the notebook, to reset the memory and get better fresh
numbers

``` python
from transformers import AutoModelForImageTextToText, AutoTokenizer
import torch
import torch.nn.functional as F
```

``` python
model_id = "Qwen/Qwen3.5-0.8B"
model = AutoModelForImageTextToText.from_pretrained(model_id, dtype=torch.bfloat16, device_map="auto", attn_implementation="eager")
```

    Loading weights:   0%|          | 0/473 [00:00<?, ?it/s]

``` python
tok = AutoTokenizer.from_pretrained(model_id)
very_very_long_input = """Elephants are the largest living land animals. Three living species are currently recognised: the African bush elephant (Loxodonta africana), the African forest elephant (L. cyclotis), and the Asian elephant (Elephas maximus). They are the only surviving members of the family Elephantidae and the order Proboscidea; extinct relatives include mammoths and mastodons. Distinctive features of elephants include a long proboscis called a trunk, tusks, large ear flaps, pillar-like legs, and tough but sensitive grey skin. The trunk is prehensile, bringing food and water to the mouth and grasping objects. Tusks, which are derived from the incisor teeth, serve both as weapons and as tools for moving objects and digging. The large ear flaps assist in maintaining a constant body temperature as well as in communication. African elephants have larger ears and concave backs, whereas Asian elephants have smaller ears and convex or level backs.
Elephants are scattered throughout sub-Saharan Africa, South Asia, and Southeast Asia. They live in different habitats, including savannahs, forests, deserts, and marshes. They are herbivorous, and they stay near water when it is accessible. They are considered to be keystone species, due to their impact on their environments. Elephants have a fission–fusion society, in which multiple family groups come together to socialise. Females (cows) tend to live in family groups, which can consist of one female with its calves or several related females with offspring. The leader of a female group, usually the oldest cow, is known as the matriarch.
Males (bulls) leave their family groups when they reach puberty and may live alone or with other males. Adult bulls mostly interact with family groups when looking for a mate. They enter a state of increased testosterone and aggression known as musth, which helps them gain dominance over other males as well as reproductive success. Calves are the centre of attention in their family groups and rely on their mothers for as long as three years. Elephants can live up to 70 years in the wild. They communicate by touch, sight, smell, and sound; elephants use infrasound and seismic communication over long distances. Elephant intelligence has been compared with that of primates and cetaceans. They appear to have self-awareness, and possibly show concern for dying and dead individuals of their kind.
African bush elephants and Asian elephants are listed as endangered and African forest elephants as critically endangered on the IUCN Red Lists. One of the biggest threats to elephant populations is the ivory trade, as the animals are poached for their ivory tusks. Other threats to wild elephants include habitat destruction and conflicts with local people. Elephants are used as working animals in Asia. In the past, they were used in war; today, they are often controversially put on display in zoos, or employed for entertainment in circuses. Elephants have an iconic status in human culture and have been widely featured in art, folklore, religion, literature, and popular culture."""
inputs = tok(very_very_long_input, return_tensors="pt").to(model.device)

with torch.no_grad(): # since we are only doing inference, we don't want to accumulate gradients
    out = model(
        **inputs,
        output_hidden_states=True,
        return_dict=True,
    )
```

    [transformers] `causal_conv1d_fn` is falling back to its reference PyTorch implementation because `causal_conv1d` is not installed. This is correct but much slower; install `causal_conv1d` for the optimized kernel.
    [transformers] `chunk_gated_delta_rule` is falling back to its reference PyTorch implementation because `flash-linear-attention` is not installed. This is correct but much slower; install `flash-linear-attention` for the optimized kernel.

``` python
# Final hidden states: [batch, sequence_length, hidden_size]
last_hidden = out.hidden_states[-1]
print("Last hidden state:", tuple(last_hidden.shape))

# hidden at position t predicts token at position t + 1.
# so we are taking all but last one, since we don't know the last token GT
hidden = last_hidden[:, :-1, :]
print("Hidden state for 1st to n-1th tokens:", tuple(hidden.shape))
# similarly for target, we take from second token to the end
targets = inputs["input_ids"][:, 1:]
print("Targets for 2nd to nth tokens:", tuple(targets.shape))

# Flatten batch and sequence dimensions.
hidden = hidden.reshape(-1, hidden.size(-1))
print("Flattened Hidden:", tuple(hidden.shape))
targets = targets.reshape(-1)
print("Flattened Targets:", tuple(targets.shape))
```

    Last hidden state: (1, 632, 1024)
    Hidden state for 1st to n-1th tokens: (1, 631, 1024)
    Targets for 2nd to nth tokens: (1, 631)
    Flattened Hidden: (631, 1024)
    Flattened Targets: (631,)

``` python
# Use the actual output head from the model.
output_head = model.get_output_embeddings()
print("Weight Matrix:", output_head)
```

    Weight Matrix: Linear(in_features=1024, out_features=248320, bias=False)

``` python
def chunked_cross_entropy():
    chunk_size = 32
    total_loss = torch.zeros(
        (),
        device=hidden.device,
        dtype=torch.float32,
    )
    total_tokens = 0

    for start in range(0, hidden.size(0), chunk_size):
        end = start + chunk_size

        chunk_hidden = hidden[start:end]
        chunk_targets = targets[start:end]

        # prediction only for the chunked tokens
        chunk_logits = output_head(chunk_hidden)

        # we sum the loss here, and at the end we can divide by number of tokens
        chunk_loss = F.cross_entropy(chunk_logits.float(), chunk_targets, reduction="sum")

        total_loss += chunk_loss
        total_tokens += chunk_targets.numel()

    return total_loss / total_tokens

def full_cross_entropy():
    # Creates the complete [number_of_tokens, vocabulary_size] tensor.
    logits  = output_head(hidden)
    return F.cross_entropy(logits.float(), targets)
```

``` python
# The measurement function includes CUDA synchronization. Without synchronization, 
# GPU timing can be inaccurate because CUDA operations are asynchronous.

def measure_peak_memory_and_time(function, repeats=10):
    # Warm up once so initialization overhead does not affect the timing.
    with torch.no_grad():
        for _ in range(2):
            warmup_loss = function()

    del warmup_loss

    torch.cuda.synchronize() # Wait for all kernels in all streams on a CUDA device to complete
    torch.cuda.empty_cache() # Release all unoccupied cached memory
    torch.cuda.reset_peak_memory_stats() # Reset the "peak" stats tracked by the CUDA memory allocator

    starting_memory = torch.cuda.memory_allocated() # Return the current GPU memory occupied by tensors in bytes

    # CUDA events are synchronization markers that can be used to monitor the device's progress, to accurately measure timing, and to synchronize CUDA streams.
    start_event = torch.cuda.Event(enable_timing=True)
    end_event = torch.cuda.Event(enable_timing=True)

    start_event.record()

    with torch.no_grad():
        for _ in range(repeats):
            loss = function()

    end_event.record()

    torch.cuda.synchronize()

    elapsed_ms = start_event.elapsed_time(end_event)
    average_time_ms = elapsed_ms / repeats

    peak_memory = torch.cuda.max_memory_allocated()
    peak_memory_used = peak_memory - starting_memory

    return (loss.detach(), peak_memory_used, average_time_ms)
```

``` python
ordinary_loss, ordinary_peak, ordinary_time = measure_peak_memory_and_time(full_cross_entropy)
chunked_loss, chunked_peak, chunked_time = measure_peak_memory_and_time(chunked_cross_entropy)
```

``` python
def mb(number_of_bytes):
    return number_of_bytes / (1024 ** 2)
```

``` python
print("ordinary loss:", ordinary_loss.item())
print("chunked loss: ", chunked_loss.item())
print()

print(f"ordinary CE peak memory: {mb(ordinary_peak):.2f} MB")
print(f"chunked CE peak memory:  {mb(chunked_peak):.2f} MB")
print(f"memory ratio:            {ordinary_peak / max(chunked_peak, 1):.2f}x")
print()

print(f"ordinary CE time:         {ordinary_time:.3f} ms")
print(f"chunked CE time:          {chunked_time:.3f} ms")
print(f"time ratio:               {chunked_time / ordinary_time:.2f}x")
print()

print("losses match:", torch.allclose(ordinary_loss, chunked_loss, atol=1e-2, rtol=1e-3))
```

    ordinary loss: 2.325397253036499
    chunked loss:  2.325397253036499

    ordinary CE peak memory: 1494.86 MB
    chunked CE peak memory:  75.78 MB
    memory ratio:            19.73x

    ordinary CE time:         8.259 ms
    chunked CE time:          20.823 ms
    time ratio:               2.52x

    losses match: True

As you can see the loss is the same upto floating-point rounding, but
the peak memory used for full cross entropy is 20x more than the memory
used for the chunked cross entropy with the trade-off of 2.5x more time
to calculate

# Lets add another head to our model

We will try to do two predictions at once, instead of predicting the
next token, we will add a head, which will try to predict `t+2` token.
Since we don’t want to train it in full, we will train on a small set of
sequences and test it on one of them itself

The hidden state at position `t` will be used by the original head to
predict token `t+1`. The new head sees the same hidden state, but its
target will be shifted by two positions. We will copy the original head
weights for the new head, since we starting from random initialization
will not converge the head with very few examples. Also it will give the
new head a useful vocabulary projection, but we should make sure that
the copy must be independent.\
Directly assigning `head_t2 = head_t1` would share parameters rather
than create a second head.

Then we will train both heads with the sum of their losses. The `t+2`
head has fewer valid positions and is naturally a harder task, because
it must predict two tokens ahead without seeing the token at `t+1` as an
input. But let’s see what will be the actual loss values for each

``` python
from transformers import AutoModelForImageTextToText, AutoTokenizer
import torch
import torch.nn as nn
import torch.nn.functional as F
```

``` python
model_id = "Qwen/Qwen3.5-0.8B"
model = AutoModelForImageTextToText.from_pretrained(model_id, dtype=torch.bfloat16, device_map="auto", attn_implementation="eager")
tok = AutoTokenizer.from_pretrained(model_id)
```

    Loading weights:   0%|          | 0/473 [00:00<?, ?it/s]

``` python
# The existing model output head will be head #1
head_t1 = model.get_output_embeddings()

# new head with the same shape and dtype/device.
head_t2 = nn.Linear(head_t1.in_features, head_t1.out_features, bias=False,
                    device=head_t1.weight.device, dtype=head_t1.weight.dtype
                    )

# copy the weights from head #1 to head #2
with torch.no_grad(): # no grad because during backprop we don't want the grad from 
    # head #2 moving to head #1
    head_t2.weight.copy_(head_t1.weight)

print("head 1 and head 2 share storage:", head_t1.weight.data_ptr() == head_t2.weight.data_ptr())

# check if the weights are same
print("head 1 and head 2 have same weights:", torch.all(head_t1.weight == head_t2.weight).item())
```

    head 1 and head 2 share storage: False
    head 1 and head 2 have same weights: True

The corpus is intentionally tiny, so the result will show memorization
rather than general language ability.

``` python
# A small corpus for training our network
docs = [
    "The cat sat on the mat and watched the birds outside the window.",
    "Birds fly south in winter to find warmer weather and better food.",
    "A young girl opened a book and started to read beside the fireplace.",
    "The river runs through the valley after the rain has filled the streams.",
    "Dogs bark at the moon when the night is still and unusually quiet.",
    "The baker made fresh bread before the sun came up over the village.",
    "Green leaves fall from the trees in late autumn and cover the ground.",
    "The boy kicked the ball across the empty field toward his friends.",
    "A small boat crossed the lake while the wind moved through the reeds.",
    "The scientist recorded each result carefully before repeating the experiment.",
    "The teacher explained the problem slowly and answered every question.",
    "Clouds gathered above the hills before the first drops of rain arrived.",
    "The farmer planted seeds in spring and waited patiently for the harvest.",
    "A bright star appeared above the quiet town as evening turned to night.",
    "The musician practiced the difficult song until every note sounded clear.",
    "Children built a sandcastle near the shore while waves rolled onto the beach.",
    "The old library contained many stories about explorers and distant lands.",
    "Fresh snow covered the road and made the surrounding forest look peaceful.",
    "The engineer designed a bridge that could safely carry heavy trucks.",
    "After breakfast, the family walked through the park and listened to birds.",
    "The doctor examined the patient carefully before recommending a treatment.",
    "A red train traveled across the countryside during the early morning hours.",
]
print("Number of sequences:", len(docs))
```

    Number of sequences: 22

``` python
train_enc = tok(docs, return_tensors="pt", padding="max_length", 
                max_length=32, truncation=True).to(model.device)

print("Train sequence sample tokens")
train_input_ids = train_enc["input_ids"]
print(train_input_ids[0].tolist())
train_attention_mask = train_enc["attention_mask"]
print(train_attention_mask[0].tolist())

# Align each target with the hidden state that should predict it.
print("Head #1 target tokens. Each hidden state at t predicts the input token at t+1")
targets_t1 = train_input_ids[:, 1:].clone()
print(targets_t1[0].tolist())
targets_t1[train_attention_mask[:, 1:] == 0] = -100
print(targets_t1[0].tolist())

# for head 2 targets, we want to check with t+2 tokens for token t
print("Head #2 target tokens. Each hidden state at t predicts the input token at t+2")
targets_t2 = train_input_ids[:, 2:].clone()
print(targets_t2[0].tolist())
targets_t2[train_attention_mask[:, 2:] == 0] = -100
print(targets_t2[0].tolist())
```

    Train sequence sample tokens
    [760, 7993, 7338, 383, 279, 5344, 321, 14936, 279, 19053, 4732, 279, 3136, 13, 248044, 248044, 248044, 248044, 248044, 248044, 248044, 248044, 248044, 248044, 248044, 248044, 248044, 248044, 248044, 248044, 248044, 248044]
    [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
    Head #1 target tokens. Each hidden state at t predicts the input token at t+1
    [7993, 7338, 383, 279, 5344, 321, 14936, 279, 19053, 4732, 279, 3136, 13, 248044, 248044, 248044, 248044, 248044, 248044, 248044, 248044, 248044, 248044, 248044, 248044, 248044, 248044, 248044, 248044, 248044, 248044]
    [7993, 7338, 383, 279, 5344, 321, 14936, 279, 19053, 4732, 279, 3136, 13, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100]
    Head #2 target tokens. Each hidden state at t predicts the input token at t+2
    [7338, 383, 279, 5344, 321, 14936, 279, 19053, 4732, 279, 3136, 13, 248044, 248044, 248044, 248044, 248044, 248044, 248044, 248044, 248044, 248044, 248044, 248044, 248044, 248044, 248044, 248044, 248044, 248044]
    [7338, 383, 279, 5344, 321, 14936, 279, 19053, 4732, 279, 3136, 13, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100]

``` python
# include the original model params along with the new head
opt = torch.optim.AdamW(list(model.parameters()) + list(head_t2.parameters()), lr=1e-5)

model.train()
loss_history = []
num_steps = 100

for step in range(num_steps + 1):
    out = model(
        input_ids=train_input_ids,
        attention_mask=train_attention_mask,
        output_hidden_states=True,
        return_dict=True,
    )
    hidden = out.hidden_states[-1]

    # The model's existing output head predicts t+1.
    logits_t1 = out.logits[:, :-1, :]
    loss_t1 = F.cross_entropy(
        logits_t1.float().reshape(-1, logits_t1.size(-1)),
        targets_t1.reshape(-1),
        ignore_index=-100,
    )

    # The copied, but independent, head predicts t+2.
    logits_t2 = head_t2(hidden[:, :-2, :])
    loss_t2 = F.cross_entropy(
        logits_t2.float().reshape(-1, logits_t2.size(-1)),
        targets_t2.reshape(-1),
        ignore_index=-100,
    )

    total_loss = loss_t1 + loss_t2
    loss_history.append((step, loss_t1.item(), loss_t2.item(), total_loss.item()))

    if step % 50 == 0:
        print(
            f"step {step:3d}  t+1={loss_t1.item():.4f}  "
            f"t+2={loss_t2.item():.4f}  sum={total_loss.item():.4f}"
        )

    if step == num_steps:
        break

    opt.zero_grad(set_to_none=True)
    total_loss.backward()
    opt.step()
```

    step   0  t+1=3.8985  t+2=9.2174  sum=13.1159
    step  50  t+1=0.7429  t+2=0.7333  sum=1.4763
    step 100  t+1=0.6516  t+2=0.6435  sum=1.2951

The `t+2` loss will naturally start higher than the `t+1` loss and
decrease during training, but will usually remain higher. Predicting
`t+2` is less certain than predicting the immediately following token,
and also it has two less possible target positions per sequence.

Now let’s do inference on one of our training samples, and see how well
we have memorized the data, and does our head \#2 predict the output
well

``` python
model.eval()

example = docs[0]
example_enc = tok(example, return_tensors="pt").to(model.device)
example_ids = example_enc["input_ids"][0]
example_tokens = tok.convert_ids_to_tokens(example_ids.tolist())

with torch.no_grad():
    example_out = model(
        **example_enc,
        output_hidden_states=True,
        return_dict=True,
    )
    example_hidden = example_out.hidden_states[-1]

    # Head 1: h_t -> token t+1.
    predictions_t1 = example_out.logits[0, :-1].argmax(dim=-1)

    # Head 2: h_t -> token t+2.
    predictions_t2 = head_t2(example_hidden[:, :-2, :])[0].argmax(dim=-1)

print("Input:", example)
print()
print("position | input token | actual t+1 | head 1 prediction | actual t+2 | head 2 prediction")
print("-" * 86)

for t, input_token in enumerate(example_tokens):
    actual_t1 = example_tokens[t + 1] if t + 1 < len(example_tokens) else "-"
    actual_t2 = example_tokens[t + 2] if t + 2 < len(example_tokens) else "-"
    predicted_t1 = predictions_t1[t].item() if t < len(predictions_t1) else None
    predicted_t2 = predictions_t2[t].item() if t < len(predictions_t2) else None

    predicted_t1 = (
        tok.convert_ids_to_tokens([predicted_t1])[0]
        if predicted_t1 is not None else "-"
    )
    predicted_t2 = (
        tok.convert_ids_to_tokens([predicted_t2])[0]
        if predicted_t2 is not None else "-"
    )

    print(
        f"{t:8d} | {input_token:11s} | {actual_t1:11s} | "
        f"{predicted_t1:18s} | {actual_t2:11s} | {predicted_t2}"
    )
```

    Input: The cat sat on the mat and watched the birds outside the window.

    position | input token | actual t+1 | head 1 prediction | actual t+2 | head 2 prediction
    --------------------------------------------------------------------------------------
           0 | The         | Ġcat        | Ġold               | Ġsat        | Ġexplained
           1 | Ġcat        | Ġsat        | Ġsat               | Ġon         | Ġon
           2 | Ġsat        | Ġon         | Ġon                | Ġthe        | Ġthe
           3 | Ġon         | Ġthe        | Ġmat               | Ġmat        | Ġmat
           4 | Ġthe        | Ġmat        | Ġand               | Ġand        | Ġand
           5 | Ġmat        | Ġand        | Ġand               | Ġwatched    | Ġwatched
           6 | Ġand        | Ġwatched    | Ġwatched           | Ġthe        | Ġthe
           7 | Ġwatched    | Ġthe        | Ġthe               | Ġbirds      | Ġbirds
           8 | Ġthe        | Ġbirds      | Ġoutside           | Ġoutside    | Ġoutside
           9 | Ġbirds      | Ġoutside    | Ġoutside           | Ġthe        | Ġthe
          10 | Ġoutside    | Ġthe        | Ġthe               | Ġwindow     | Ġwindow
          11 | Ġthe        | Ġwindow     | Ġwindow            | .           | .
          12 | Ġwindow     | .           | .                  | -           | -
          13 | .           | -           | -                  | -           | -

Now let’s try it on some different sequence, which we have not trained
on

``` python
model.eval()

example = "The Capital of India is New Delhi."
example_enc = tok(example, return_tensors="pt").to(model.device)
example_ids = example_enc["input_ids"][0]
example_tokens = tok.convert_ids_to_tokens(example_ids.tolist())

with torch.no_grad():
    example_out = model(
        **example_enc,
        output_hidden_states=True,
        return_dict=True,
    )
    example_hidden = example_out.hidden_states[-1]

    # Head 1: h_t -> token t+1.
    predictions_t1 = example_out.logits[0, :-1].argmax(dim=-1)

    # Head 2: h_t -> token t+2.
    predictions_t2 = head_t2(example_hidden[:, :-2, :])[0].argmax(dim=-1)

print("Input:", example)
print()
print("position | input token | actual t+1 | head 1 prediction | actual t+2 | head 2 prediction")
print("-" * 86)

for t, input_token in enumerate(example_tokens):
    actual_t1 = example_tokens[t + 1] if t + 1 < len(example_tokens) else "-"
    actual_t2 = example_tokens[t + 2] if t + 2 < len(example_tokens) else "-"
    predicted_t1 = predictions_t1[t].item() if t < len(predictions_t1) else None
    predicted_t2 = predictions_t2[t].item() if t < len(predictions_t2) else None

    predicted_t1 = (
        tok.convert_ids_to_tokens([predicted_t1])[0]
        if predicted_t1 is not None else "-"
    )
    predicted_t2 = (
        tok.convert_ids_to_tokens([predicted_t2])[0]
        if predicted_t2 is not None else "-"
    )

    print(
        f"{t:8d} | {input_token:11s} | {actual_t1:11s} | "
        f"{predicted_t1:18s} | {actual_t2:11s} | {predicted_t2}"
    )
```

    Input: The Capital of India is New Delhi.

    position | input token | actual t+1 | head 1 prediction | actual t+2 | head 2 prediction
    --------------------------------------------------------------------------------------
           0 | The         | ĠCapital    | Ġold               | Ġof         | Ġexplained
           1 | ĠCapital    | Ġof         | Ġof                | ĠIndia      | Ġof
           2 | Ġof         | ĠIndia      | Ġthe               | Ġis         | Ġthe
           3 | ĠIndia      | Ġis         | Ġthe               | ĠNew        | Ġthe
           4 | Ġis         | ĠNew        | Ġthe               | ĠDelhi      | Ġthe
           5 | ĠNew        | ĠDelhi      | ĠDelhi             | .           | ĠDelhi
           6 | ĠDelhi      | .           | .                  | -           | -
           7 | .           | -           | -                  | -           | -

Our model has officialy gone crazy and is just predicting `the`, since
we have overfit to our very tiny dataset