Training Stability › Worked Example: Diagnosing a Simulated Instability
Language AI HandbookAll chapters Training Stability › Worked Example: Diagnosing a Simulated InstabilityOne section · Sections in this chapter8 / 15 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Let us walk through a realistic scenario: a small transformer model that encounters instability and recovers. This consolidates the monitoring, detection, and recovery concepts into a concrete workflow that mirrors what practitioners encounter in real large-scale training runs.
The scenario involves a small language model trained on random token sequences, with two "hard batches" injected at known steps to amplify the backward signal. We run two versions, one without gradient clipping and one with clipping. The logged loss is divided by the artificial scale, while PyTorch's clipping function returns the norm measured before rescaling. The comparison therefore shows how similar loss trajectories can hide different update rules at the injected steps.
class MiniLM(nn.Module):
"""Small language model for training stability demonstration."""
def __init__(self, vocab_size=100, d_model=64, n_layers=4, n_heads=4):
super().__init__()
self.embedding = nn.Embedding(vocab_size, d_model)
self.layers = nn.ModuleList(
[SimpleTransformerBlock(d_model, n_heads) for _ in range(n_layers)]
)
self.norm = nn.LayerNorm(d_model)
self.head = nn.Linear(d_model, vocab_size)
def forward(self, x):
h = self.embedding(x)
for layer in self.layers:
h = layer(h)
return self.head(self.norm(h))
def run_stability_scenario(use_clipping=True, clip_val=1.0, seed=42):
"""Run a training scenario with or without gradient clipping."""
torch.manual_seed(seed)
np.random.seed(seed)
model_s = MiniLM(vocab_size=100, d_model=64, n_layers=4, n_heads=4)
opt_s = optim.Adam(model_s.parameters(), lr=5e-3, eps=1e-8)
ce_loss = nn.CrossEntropyLoss()
losses_s = []
grad_norms_s = []
for step in range(200):
# Occasional "hard" batch simulating outlier data
if step in [60, 120]:
seq_len = 64
batch = torch.randint(0, 100, (8, seq_len))
target_tok = torch.randint(0, 100, (8 * seq_len,))
scale = 5.0 # amplify loss signal
else:
seq_len = 32
batch = torch.randint(0, 100, (8, seq_len))
target_tok = torch.randint(0, 100, (8 * seq_len,))
scale = 1.0
opt_s.zero_grad()
logits = model_s(batch).reshape(-1, 100)
loss_s = ce_loss(logits, target_tok) * scale
loss_s.backward()
if use_clipping:
g_norm = torch.nn.utils.clip_grad_norm_(
model_s.parameters(), max_norm=clip_val
)
else:
g_norm = torch.nn.utils.clip_grad_norm_(
model_s.parameters(), max_norm=float("inf")
)
grad_norms_s.append(g_norm.item())
losses_s.append(loss_s.item() / scale) # log true loss, not scaled
opt_s.step()
return losses_s, grad_norms_s
losses_no_clip, gnorms_no_clip = run_stability_scenario(use_clipping=False)
losses_with_clip, gnorms_with_clip = run_stability_scenario(
use_clipping=True, clip_val=1.0
)
Summary: No Clipping vs. Gradient Clipping (threshold=1.0) Metric No Clip With Clip -------------------------------------------------------- Mean loss 4.6297 4.6293 Final loss 4.6135 4.6138 Max gradient norm 1.3862 1.3862 Mean gradient norm 0.4079 0.4095
The results are deliberately more modest than the stylized failure case discussed earlier. The two runs have nearly identical mean and final losses, and both report the same maximum norm because clip_grad_norm_ returns the pre-clipping value. The operational difference is in the update: at steps 60 and 120, the clipped run rescales gradients above 1.0 before optimizer.step(), while the unclipped run applies them at full magnitude. In this short seeded example, that safeguard does not produce a visible loss advantage.