Repository navigation
Add Corrdiff-era5-hrrr - #2088
Dibyajyoti-Chakraborty wants to merge 12 commits into
Conversation
CODEOWNERS review mapCurrent for commit No CODEOWNER
Comment |
Dibyajyoti-Chakraborty
changed the title
|
[Low impact] Adds a weather model example with configs and sample code. Satisfy the model interface requirements before merging; the other findings are non-blocking improvements. Findings
SummaryAdds an ERA5-to-HRRR downscaling example using clean-state prediction and rectified flow.
Reviews (1) · Last reviewed commit: "Corrdiff-era5-hrrr initial commit" · Reviewed by Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
|
Only my 2 cents there, and @pzharrington would know better, but I'm not sure we want to have another recipe in addition of the unified one for regional weather models. I haven't looked in detail, but it does feel like a lot of duplication |
|
@CharlelieLrt @pzharrington I can add it to the |
|
I'm not sure it would be so much refactoring. The inference doesn't really matter (and maybe doesn't eevn need to part of the PR), since the recommended way to do inference is the E2S wrapper. |
This was my thinking as well, I didn't know you were making this as a separate recipe. The SAs strongly prefer one centralized recipe to train all these regional weather diffusion models. So if this can be shepherded under that one without too much difficulty/messiness then I think that's a better approach. Sorry I know a bit annoying since you already have this one but I think it makes sense to de-fragment if we can. |
|
Got it. Let me combine it and update the PR. |
| sampling always agree. Unlike the ``EDMPreconditioner``, this wrapper | ||
| does not scale the inputs and outputs: the network output is the | ||
| prediction itself (``training.loss.prediction_type``), with any noise | ||
| schedule. |
There was a problem hiding this comment.
Just curious, wouldn't this be generally useful? Not just for this recipe but across flow-matching use-cases? i.e., why here and not in phyiscsnemo.diffusion?
There was a problem hiding this comment.
Ya, I had a discussion with @CharlelieLrt back in the flow matching PR. The time scaling was found to be too specific to our timestep embedder in DiT/Unet. So not generally applicable. It is put as a note in the FM docstring. @CharlelieLrt ?
For the EDMPreconditioner, it needs to be ablated. FM is variance-preserving and typically more stable in training. But idk, maybe useful in FM too.
Description
Adds flow matching as a configuration option of
examples/weather/regional_weather_diffusionusing the CorrDiff-ERA5-HRRR recipe: ERA5 → HRRR (3 km, CONUS) downscaling with x-prediction rectified flow matching. Inference is with the Earth2Studio wrapper.training.loss.type:edm(unchanged default) orflow_matching, which trains withRectifiedFlowNoiseScheduler(t_min, t_max)andWeightedFlowMatchingLoss.training.loss.prediction_type:x0,epsilon,score, orflow. Both losses take it directly; for rectified flow the sampler converts the prediction to the ODE flow with the scheduler's helpers (finite at the pure-noise startt = 1).model.preconditioner:edm(EDMPreconditioner, requiresprediction_type: x0) ortime_scaled(TimeScaledModel, multiplies the times bymodel.time_scale). Either combines with either loss type.Other changes: the DiT builder accepts
attention_backend(natten2d_rope);sampler/rectified_flow.yamladds a resolution-dependenttime_shiftand the churn options.CorrDiff-ERA5-HRRR recipe (
corrdiff_era5_hrrr.yaml): datasetdata_loader_corrdiff_era5_hrrr.py(same-time ERA5/HRRR pairs over the 1024 × 1792 CONUS domain, random 256 × 256 patches, zenith angle and static fields in the background, droppable inputs with presence flags as scalar conditions), modelmodel/corrdiff_era5_hrrr.yaml(678M-parameter DiT, NATTEN + RoPE, QK normalization,time_scaledpreconditioner), x-prediction witht ~ U[0, 1]in a single stage, AdamW with warmup and cosine decay, bf16. The config matches the released checkpoint, which loads unchanged.Testing: new tests for flow-matching training (DiT and U-Net; x0, epsilon, flow prediction; fp32 and bf16), the EDM loss with the time-scaled preconditioner, the preconditioner/prediction-type validation, pure-noise sampling with and without time shift, and the CONUS dataset on synthetic files. They pass together with the existing EDM tests (
test_model_types, DiT checkpointing, 2-GPU domain-parallel training). A 20-step real-data run ofcorrdiff_era5_hrrron 8 GPUs trains (loss 0.89 → 0.37, 64 GB per GPU at 2 patches per GPU).Checklist
Dependencies
No new library dependencies. The example's
requirements.txtlistsnatten,matplotlib,tensorboard,xarray,zarr,netCDF4, andpandas. The QK normalization uses NVIDIA Apex, which ships in the PhysicsNeMo and NGC PyTorch containers.