A learning rate schedule changes the optimizer’s learning rate during training instead of keeping it fixed. In language-model training, scheduling can help control exploding gradients and support faster convergence, while remaining only one part of generalization practice. section 9.2 section 5.2.1
Use scheduling with these practical expectations:
- Pair the schedule with an optimizer; the scheduler adjusts the learning rate used during iterative weight updates. section 9.2 section 5.2.1
- Treat it as a training aid rather than a guarantee: monitoring remains necessary because scheduling is not a silver bullet for overfitting or generalization. section 5.2.1
Cosine annealing in PyTorch
A SimpleLlama training setup uses AdamW with an initial learning rate of 5e-4, then applies PyTorch’s CosineAnnealingLR across 1,000 epochs with a minimum learning rate of 1e-5. section 9.2
MASTER_CONFIG["epochs"] = 1000
MASTER_CONFIG["batch_size"] = 16
MASTER_CONFIG["d_model"] = 768
MASTER_CONFIG["n_layers"] = 8
MASTER_CONFIG["context_window"] = 128
llama = SimpleLlama(MASTER_CONFIG).to(device)
llama_optimizer = torch.optim.AdamW(
llama.parameters(),
betas=(0.9, 0.95),
weight_decay=1e-1,
eps=1e-9,
lr=5e-4,
)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
llama_optimizer, 1000, eta_min=1e-5
) #1
#Epoch 0 | train loss 10.321 | val loss 10.316 | Time 0.622 | ETA:
0:00:12.439990
#lr: [0.0004999987909744553]
#training loss 6.216 | validation loss: 6.046
generate(
llama,
config=MASTER_CONFIG,
temperature=1.0,
top_k=25,
max_new_tokens=50,
)
#'<s> Write a short story. Possible Story: the Story there One.t day. Back
the, went to: her they Possible|. to and a said saw They:. be the She.. a.
to They. they. to and to for He was a in with',',
for i in GLOBAL_KEEP_TRACK:
print(i)
#SimpleFeedForwardNN 18547809 Params | Train: 4.129 | Val: 4.458
#SimpleLlama 187827210 Params | Train: 6.216 | Val: 6.046
Inspect the complete optimizer and scheduler configuration, including beta values, weight decay, initial learning rate, epoch count, and minimum learning rate. section 9.2
