mllam/neural-lam

Add configurable gradient clipping to prevent exploding gradients during autoregressive training

Open

#280 创建于 2026年2月27日

在 GitHub 查看
 (0 评论) (0 反应) (1 负责人)Python (276 fork)auto 404
enhancementhelp wanted

仓库指标

Star
 (282 star)
PR 合并指标
 (PR 指标待抓取)

描述

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

贡献者指南