Day 62: Loss Masking & Target Label Construction

import torch
import torch.nn as nn
import torch.nn.functional as F

logits = torch.randn(1, 4, 3, requires_grad=True)
labels = torch.tensor([[ -100, 1, 2, -100 ]])
criterion = nn.CrossEntropyLoss(ignore_index=-100)
loss = criterion(logits.view(-1, 3), labels.view(-1))

lsm = F.log_softmax(logits[0], dim=-1)
expected = -(lsm[1, 1] + lsm[2, 2]) / 2.0
assert torch.allclose(loss, expected)

import torch
def mask_prompt_labels(input_ids, assistant_header_ids):
labels = input_ids.clone()
seq_len = len(input_ids)
pat_len = len(assistant_header_ids)
for i in range(seq_len - pat_len + 1):
if input_ids[i:i+pat_len].tolist() == assistant_header_ids:
labels[:i+pat_len] = -100
return labels
labels[:] = -100
return labels

import torch
import torch.nn as nn

emb = nn.Embedding(100, 16)
head = nn.Linear(16, 100)

input_ids = torch.tensor([[10, 20, 30, 40], [10, 50, 60, 0]])
labels = torch.tensor([[-100, -100, 30, 40], [-100, -100, 60, -100]])

x = emb(input_ids)
logits = head(x)
loss = nn.CrossEntropyLoss(ignore_index=-100)(logits.view(-1, 100), labels.view(-1))
loss.backward()
# Verified: tokens at prompt/pad positions do not produce gradient in loss targets

class SFTDataCollator:
def __init__(self, response_ids, pad_token_id):
self.response_ids = response_ids
self.pad_token_id = pad_token_id
def __call__(self, batch):
max_len = max(len(x['input_ids']) for x in batch)
padded_inputs, padded_labels = [], []
for x in batch:
ids = list(x['input_ids'])
labels = [-100] * len(ids)
r_len = len(self.response_ids)
for i in range(len(ids) - r_len + 1):
if ids[i:i+r_len] == self.response_ids:
for j in range(i + r_len, len(ids)):
labels[j] = ids[j]
pad_len = max_len - len(ids)
padded_inputs.append(ids + [self.pad_token_id] * pad_len)
padded_labels.append(labels + [-100] * pad_len)
return {
'input_ids': torch.tensor(padded_inputs),
'labels': torch.tensor(padded_labels)
}