ai2cm/ace: 2026.5.1

What's Changed

Fine-Tuning & Checkpoint Resume

New config options make it easier to resume or fine-tune from existing checkpoints:

• OptimizationConfig.resume_optimizer_ckpt_path: restore optimizer state when fine-tuning (#1043)
• EMAConfig.resume_ema_ckpt_path: resume from an EMA checkpoint (#1118)
• CheckpointStepperConfig: load stepper config directly from a checkpoint (#1103)
• Optimizer/EMA state is now included in epoch checkpoints (ckpt_{epoch:04d}.tar) (#1104)

Ensemble Inference

• Initial ensemble (IC ensemble) support added to the evaluator and inference aggregators (#709)

New Models & Architecture

• filter_preserves_global_mean option added to SFNO (#1100)
• SecondaryModuleStepConfig / SecondaryModuleStep: compose a secondary module during training steps (#1073)

Coupled Model

• Stochastic CoupledStepper training (#750)
• Randomly sampled LossContributions.n_steps (#869)
• optimize_last_step_only added to coupled LossContributionsConfig (#868)

Downscaling

• Configurable loss weighting for downscaling (#1056)

Diagnostics

• Power spectrum diagnostics logged in the inference entrypoint (#1078, #1079)
• Weather eval entrypoint replaced with a more general additional_inference list (#1096)

Data Processing

• Time subsetting can now be configured prior to time coarsening (#1055)
• PRMSL added to X-SHiELD data processing configurations (#1036)

Bug Fixes

• Clamped SSR calculation that was producing NaNs silently dropped from W&B (#1088)
• Worked around xarray StringDType serialization error (#1086)
• Signal handler now exits with a nonzero code (#1068)
• IceCorrectorConfig correctly registered in CorrectorSelector registry (#1044)

Breaking Changes

• TrainStepperConfig.train_n_forward_steps renamed to TrainStepperConfig.n_forward_steps — all train YAML configs must update this field (#1052)
• TrainConfig.n_forward_steps removed (was deprecated; use stepper_training.n_forward_steps) (#1052)
• TrainConfig.weather_evaluation: WeatherEvaluationConfig | None replaced by TrainConfig.additional_inference: list[AdditionalInferenceConfig] (#1096)
• Sub-aggregator record_batch(time, data) interface replaced by record_batch(data: InferenceBatchData) (#1097)
• StepLoss.forward() now returns LossOutput instead of torch.Tensor; call .total() for the scalar (#1020)
• fme.diffusion package removed (#1084)

Full Changelog: https://github.com/ai2cm/ace/compare/v2026.4.0...v2026.5.1