mllam/neural-lam

Add configurable gradient clipping to prevent exploding gradients during autoregressive training

Open

#280 opened on Feb 27, 2026

View on GitHub
 (0 comments) (0 reactions) (1 assignee)Python (276 forks)auto 404
enhancementhelp wanted

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

Contributor guide