Add configurable gradient clipping to prevent exploding gradients during autoregressive training
#280 opened on Feb 27, 2026
Repository metrics
- Stars
- (282 stars)
- PR merge metrics
- (PR metrics pending)
Description
Description
neural-lam's training pipeline has no gradient clipping configured anywhere — neither via PyTorch Lightning's built-in Trainer(gradient_clip_val=...) nor through manual torch.nn.utils.clip_grad_norm_() calls.
This is particularly risky because of how ARModel.unroll_prediction works:
# neural_lam/models/ar_model.py L230-277
def unroll_prediction(self, init_states, forcing_features, true_states):
for i in range(pred_steps): # <-- up to ar_steps sequential passes
pred_state, pred_std = self.predict_step(
prev_state, prev_prev_state, forcing
)
new_state = (
self.boundary_mask * border_state
+ self.interior_mask * pred_state
)
prev_prev_state = prev_state
prev_state = new_state # <-- gradient chain grows each step
Each autoregressive step feeds predictions back as input, creating a computational graph ar_steps layers deep. When --ar_steps_train is set to values like 10–25 (common for multi-day weather forecasting), backpropagation through this chain multiplies gradients at each step — the classic exploding gradient problem from recurrent architectures.
The current Trainer in train_model.py has no clipping:
trainer = pl.Trainer(
max_epochs=args.epochs,
deterministic=True,
strategy="ddp",
# ...
precision=args.precision,
# No gradient_clip_val
# No gradient_clip_algorithm
)
Why it matters
Training instability with longer rollouts: Users increasing --ar_steps_train beyond the default of 1 (e.g., for curriculum learning or multi-step loss) will encounter NaN losses or sudden loss spikes. This is a silent failure — training continues but the model is destroyed.
Standard practice in weather-ML: All major published models use gradient clipping:
GraphCast (Lam et al., 2023): gradient clipping at max norm = 32 Pangu-Weather (Bi et al., 2023): gradient clipping with max norm = 1.0 GenCast (Price et al., 2024): gradient clipping at max norm = 32 FourCastNet (Pathak et al., 2022): gradient clipping with max norm = 1.0 Mixed precision amplifies the risk: neural-lam supports --precision 16 and bf16. Lower precision reduces the representable dynamic range, making gradient overflow happen sooner and at smaller ar_steps values.
Zero-cost safeguard: Gradient clipping adds negligible computational overhead (a single torch.norm() + conditional scale per parameter group per step).