mllam/neural-lam

Add configurable gradient clipping to prevent exploding gradients during autoregressive training

Open

#280 ouverte le 27 févr. 2026

Voir sur GitHub
 (0 commentaires) (0 réactions) (1 assigné)Python (276 forks)auto 404
enhancementhelp wanted

Métriques du dépôt

Stars
 (282 stars)
Métriques de merge PR
 (Métriques PR en attente)

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).

Guide contributeur