Here's concise PyTorch (PyTorch-style) pseudocode showing a supervised fine-tuning loop for a decoder-only causal LM with cross-entropy computed from logits, attention_mask, gradient accumulation, and mixed precision via torch.cuda.amp:python
import torch
import torch.nn.functional as F
from torch.cuda.amp import autocast, GradScaler
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
optimizer = ... # assumed provided
dataloader = ... # assumed provided
epochs = 3
grad_accum_steps = 4
pad_token_id = -100 # or model.config.pad_token_id
scaler = GradScaler()
model.train()
for epoch in range(epochs):
optimizer.zero_grad()
for step, batch in enumerate(dataloader):
input_ids = batch['input_ids'].to(device) # (B, L)
attention_mask = batch['attention_mask'].to(device) # (B, L)
with autocast(): # mixed precision context
outputs = model(input_ids=input_ids, attention_mask=attention_mask)
logits = outputs.logits # (B, L, Vocab)
# shift logits and labels for causal LM: predict token t from previous tokens
shift_logits = logits[..., :-1, :].contiguous() # (B, L-1, V)
shift_labels = input_ids[..., 1:].contiguous() # (B, L-1)
# optionally shift attention_mask if needed
loss = F.cross_entropy(
shift_logits.view(-1, shift_logits.size(-1)),
shift_labels.view(-1),
ignore_index=pad_token_id,
reduction='mean'
)
loss = loss / grad_accum_steps # scale for accumulation
# scale + backward for mixed precision
scaler.scale(loss).backward()
# optimizer step when accumulated enough gradients
if (step + 1) % grad_accum_steps == 0:
scaler.step(optimizer) # calls unscale internally
scaler.update()
optimizer.zero_grad()
# optionally: scheduler.step()
