diff --git a/CHANGELOG.md b/CHANGELOG.md index dc4e708b8d..8db5f8ce0a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -84,7 +84,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 per step. The solvers share an extended semi-linear callback API (`bias_fn`, `bias_int_fn`, `slope_fn`); users can select all four solvers by string key through `physicsnemo.diffusion.samplers.sample`. - +- Adds flow matching (`training.loss.type: flow_matching`) and the CorrDiff-ERA5-HRRR + downscaling recipe to the `examples/weather/regional_weather_diffusion` example. - `ClusterTree.find_dual_interaction_pairs(source_admissible=...)` marks source nodes that must never be approximated by their aggregate, such as nodes whose normals differ too much. diff --git a/docs/img/corrdiff_era5_hrrr_t2m_downscaling.png b/docs/img/corrdiff_era5_hrrr_t2m_downscaling.png new file mode 100644 index 0000000000..9454a61a3d Binary files /dev/null and b/docs/img/corrdiff_era5_hrrr_t2m_downscaling.png differ diff --git a/examples/weather/regional_weather_diffusion/.gitignore b/examples/weather/regional_weather_diffusion/.gitignore index 8219b834f6..ffb6a5283e 100644 --- a/examples/weather/regional_weather_diffusion/.gitignore +++ b/examples/weather/regional_weather_diffusion/.gitignore @@ -2,4 +2,4 @@ *.png *.pt *.tfevents* -*wandb/ \ No newline at end of file +*wandb/*.nc diff --git a/examples/weather/regional_weather_diffusion/README.md b/examples/weather/regional_weather_diffusion/README.md index 7a790da2f6..cfe0e4dd3a 100644 --- a/examples/weather/regional_weather_diffusion/README.md +++ b/examples/weather/regional_weather_diffusion/README.md @@ -39,6 +39,15 @@ Meanwhile, Stormscope omits the regression model in favor of a diffusion-only ap These models can make longer forecasts (more than one timestep) during inference by feeding their predictions back into the model as input for the next step (autoregressive rollout). Both models can be formulated in terms of a high-resolution state update guided by a low-resolution conditioning. In the code, we refer to the high-resolution training target in generic terms as `state` and to everything the model conditions on as `background`, including any past state history a forecasting model needs as input -- the dataset concatenates that onto `background` as extra channels rather than returning it as a separate tensor. See [Model types](#model-types) and [Adding custom datasets](#adding-custom-datasets) for details. +- [CorrDiff-ERA5-HRRR](https://huggingface.co/nvidia/corrdiff-era5-hrrr), a downscaling model that generates the 3 km HRRR state over the contiguous US +from the 25 km ERA5 reanalysis at the same time. It is a DiT trained with x-prediction rectified +flow matching instead of the EDM denoising score matching of the other models; see +[Flow matching](#training-with-flow-matching-corrdiff-era5-hrrr). + +

+ +

+ ## Getting Started ### Preliminaries @@ -88,7 +97,7 @@ Training is handled by `train.py`, configured using [hydra](https://hydra.cc/doc - `inference`: specifies inference-specific settings, such as which initial condition to run, which model checkpoints to use - `hydra`: specifies basic hydra settings, like where to store outputs (based on the training or inference outputs directories) -Also in the `config` directory are several top-level configs, which show how to train a `regression` model or a `diffusion` model, and run inference (`stormcast-inference`). You can select any of these by specifying it as a config name at the command line (for example, `--config-name=regression`). Optionally, you can also override any specific items of interest uing the command line args, for example: +Also in the `config` directory are several top-level configs, which show how to train a `regression`, a `diffusion`, or a flow matching model (`corrdiff_era5_hrrr`), and run inference (`stormcast-inference`). You can select any of these by specifying it as a config name at the command line (for example, `--config-name=regression`). Optionally, you can also override any specific items of interest uing the command line args, for example: ```bash python train.py --config-name regression training.batch_size=4 @@ -104,7 +113,7 @@ The training script `train.py` will initialize the training experiment and launc training.outdir: 'rundir' # Root path under which to save training outputs training.experiment_name: 'stormcast' # Name for the training experiment training.run_id: '0' # Unique ID to use for this training run -training.rundir: ./${training.outdir}/${training.experiment_name}/${training.run_id} # Path where experiment outputs will be saved +training.rundir: ${training.outdir}/${training.experiment_name}/${training.run_id} # Path where experiment outputs will be saved ``` As you can see, the `training.run_id` setting can be used for distinguishing between different runs of the same configuration. The final training output directory is constructed by composing together the `training.outdir` root path (defaults to `rundir`), the `training.experiment_name`, and the `training.run_id`. For inference runs, equivalent options are available in the `stormcast_inference.yaml` config file used with the `inference.py` script. @@ -190,6 +199,89 @@ training: When `sigma_bin_edges` is not provided, bins are computed automatically so that each bin has equal probability mass under the configured sigma distribution (`lognormal` or `loguniform`). This feature is ignored for regression models. +### Training with Flow Matching (CorrDiff-ERA5-HRRR) + +Three independent settings define a diffusion model in this recipe: + +- `training.loss.type`: the noise schedule and loss. The default `edm` trains with denoising score + matching on the EDM schedule (`EDMNoiseScheduler`, `WeightedMSEDSMLoss`); `flow_matching` trains + with rectified flow matching (`RectifiedFlowNoiseScheduler(t_min, t_max)`, `WeightedFlowMatchingLoss`). +- `training.loss.prediction_type`: what the network outputs, `x0` (the clean state, default), + `epsilon` (the noise), `score`, or `flow` (the velocity). The loss and the sampler convert between + these with the noise scheduler. +- `model.preconditioner`: how the network is wrapped. `edm` (default) applies the EDM input, output, + and noise-level scalings with `EDMPreconditioner`, so the wrapped network predicts `x0`; + `time_scaled` only multiplies the times by `model.time_scale` (`TimeScaledModel` in + [`utils/nn.py`](./utils/nn.py)), so the network output is the prediction itself. The preconditioner + can be combined with either loss type; the EDM preconditioner requires `prediction_type: x0`. + +With `flow_matching`, the training times are drawn uniformly from +`[training.loss.t_min, training.loss.t_max]`, where $t = 0$ is clean data and $t = 1$ is pure noise; +the default `t_max: 1` trains up to pure noise, so sampling starts from the same inputs as training. +With `prediction_type: x0` the network predicts the clean state +([x-prediction](https://arxiv.org/abs/2511.13720)) and the loss is the plain MSE on it. The sampler +is set in `sampler/rectified_flow` (`num_steps`, `solver` as `heun` or `euler`, and `time_shift`). +All components come from `physicsnemo.diffusion`; see [`utils/loss.py`](./utils/loss.py) and +`diffusion_model_forward` in [`utils/nn.py`](./utils/nn.py). + +The top-level config `corrdiff_era5_hrrr.yaml` trains CorrDiff-ERA5-HRRR with this formulation. It +combines the dataset `dataset/corrdiff_era5_hrrr`, the model `model/corrdiff_era5_hrrr` (the DiT below with +`preconditioner: time_scaled` and `time_scale: 999`), x-prediction, and the sampler `sampler/rectified_flow`: + +- **Data**: same-time pairs of ERA5 and HRRR over a 1024 × 1792 pixel CONUS domain, served by + [`datasets/data_loader_corrdiff_era5_hrrr.py`](./datasets/data_loader_corrdiff_era5_hrrr.py). The target is the 99 HRRR + variables of the table [below](#hrrr-variables). The background holds 26 ERA5 variables interpolated + to the HRRR grid, the cosine of the solar zenith angle, and three static fields (mean and standard + deviation of the elevation, land-sea mask). Four scalar conditions, the sine and cosine of the time of + year and presence flags for `tcwv` and `sp`, reach the DiT through its conditioning embedding; the + dataset drops each of these two variables with probability 0.5 during training, so the model also + runs without them. The model trains on random 256 × 256 patches (`dataset.patch_shape`), with the + pure downscaling condition list `["background"]`. +- **Model**: a DiT with 678M parameters (28 blocks, hidden size 1152, 18 heads, 2 × 2 pixel patches) + with 31 × 31 token [NATTEN](https://github.com/SHI-Labs/NATTEN) neighborhood attention, 2D rotary + position embeddings (`attention_backend: natten2d_rope`), and normalized queries and keys. The + normalization uses [NVIDIA Apex](https://github.com/NVIDIA/apex), which the PhysicsNeMo and NGC + PyTorch containers include. +- **Training**: 781,250 steps at batch size 64 (50M samples) with AdamW, learning rate $10^{-4}$, + a 1000-step warmup followed by cosine decay, bf16 mixed precision, and gradient clipping at 1. With + 2 samples per GPU, training uses about 68 GB of memory per GPU. + +To train it with 8 GPUs on one node: + +```bash +torchrun --standalone --nproc_per_node=8 train.py --config-name corrdiff_era5_hrrr dataset.data_path=/path/to/data +``` + +To check the setup on real data before a long run, train for a few steps with frequent outputs, for +example `training.total_train_steps=20 training.checkpoint_freq=10 training.validation_freq=10 +training.validation_steps=1 training.batch_size=16`. A smaller model, such as +`model.hyperparameters.hidden_size=768 model.hyperparameters.depth=16 +model.hyperparameters.num_heads=12`, trains about four times faster. On clusters with a job time +limit, `training.max_run_steps` stops a job after the given number of steps so that the next job of a +chain resumes from its checkpoint. Each checkpoint of this model takes 8 GB, so `training.checkpoint_freq` +is set to 7812 steps (0.5M samples). + +The trained model samples the full 1024 × 1792 domain, which has 28 times more pixels than the training +patches. Larger images need stronger noise to destroy their large-scale structure, so sampling at a +higher resolution than in training shifts the times toward the noise end, +$t \to a t / (1 + (a - 1) t)$ as in [Stable Diffusion 3](https://arxiv.org/abs/2403.03206). The +released model uses `time_shift: 32` with 50 Heun steps on the full domain; the sampler config of the +training keeps `time_shift: 1` because the validation samples the training patches. The recommended +way to run inference is the Earth2Studio wrapper, see [Inference](#inference). The released checkpoint +has the following ensemble-mean RMSE and spread against HRRR for 8 members on 2025-10-24 00 UTC: + +| Variable | Ensemble-mean RMSE | Ensemble spread | +| --- | --- | --- | +| `u10m` | 1.05 m/s | 0.95 m/s | +| `v10m` | 1.15 m/s | 1.01 m/s | +| `t2m` | 0.70 K | 0.58 K | +| `mslp` | 49.2 Pa | 38.9 Pa | +| `refc` | 4.70 dBZ | 4.30 dBZ | + +To test the formulation without data, use `--config-name test_flow_matching` (DiT on mock data), or run +the `test_flow_matching*` and `test_rectified_flow*` tests of `test_training.py`. The training tests also +run with 2 and 4 processes (data and domain parallelism), see [Testing](#testing). + ### Distributed Training Both regression and diffusion training can be distributed easily with data parallelism through `torchrun` or other launchers (for example, SLURM `srun`). As long as GPU memory is sufficient, the same configuration file can be used regardless of the number of GPUs. You just need to ensure the configuration being run has a batch size that is divisible by the number of available GPUs/processes. For example, distributed training of the regression model over 8 GPUs on one node would look something like: @@ -226,7 +318,7 @@ stability. The default configuration uses a batch size of 64 (controlled by `training.batch_size`), as used in the StormCast paper. If you have few GPUs and/or GPUs with limited memory, the default setting of `training.batch_size_per_gpu: 'auto'` may cause you to run out of memory. In that case, you can reduce the per-GPU memory utilization by manually setting `training.batch_size_per_gpu` to an integer value smaller than `training.batch_size` divided by the number of GPUs. The training code will automatically employ gradient accumulation to maintain the desired effective batch size specified by `training.batch_size` while using less memory, at the cost of longer training time. -Another way to reduce memory usage is to enable 16-bit training by setting `training.perf.fp_optimizations` to `amp-bf16`. This is not enabled by default as it was not used in the StormCast paper, but our experience indicates that BF16 training works as well as 32-bit training using U-Nets. Meanwhile, 32-bit training is recommended for DiT as BF16 can introduce training instabilities with that architecture. +Another way to reduce memory usage is to enable 16-bit training by setting `training.perf.fp_optimizations` to `amp-bf16`. This is not enabled by default as it was not used in the StormCast paper, but our experience indicates that BF16 training works as well as 32-bit training using U-Nets. Meanwhile, 32-bit training is recommended for DiT as BF16 can introduce training instabilities with that architecture, unless the attention queries and keys are normalized (`attn_kwargs: {qk_norm: true, norm_layer: apex}` in `model.hyperparameters`, as in `model/corrdiff_era5_hrrr`). If your samples are too big to fit even one of them on a single GPU, consider using [domain parallelism](#domain-parallelism). @@ -251,7 +343,7 @@ python inference.py --config-name This will load regression and diffusion models from directories specified by `inference.regression_checkpoint` and `inference.diffusion_checkpoint` respectively; each of these should be a path to a PhysicsNeMo checkpoint (`.mdlus` file) from your training runs. `inference.py` evaluates `inference.n_steps` independent diagnostic samples starting at `inference.initial_time` -- each step is scored against its own sample straight from the dataset's `background`/`state`, and saves outputs as a `zarr` file along with a few plots saved as `png` files. -**This is not autoregressive rollout.** `inference.py` does not feed its predictions back into the model as input for a following step; it only evaluates single-step accuracy against the dataset's own history. (Prior versions of this recipe carried predictions forward through the two-element `state` convention; that convention has been removed -- see [Adding custom datasets](#adding-custom-datasets).) For autoregressive rollout, custom datasets, or more complex inference workflows, we recommend bringing your checkpoints to [Earth2Studio](https://github.com/NVIDIA/earth2studio) for further analysis and visualizations. The [Earth2Studio wrapper for StormCast](https://github.com/NVIDIA/earth2studio/blob/main/earth2studio/models/px/stormcast.py) can be used as a starting point for custom implementations. +**This is not autoregressive rollout.** `inference.py` does not feed its predictions back into the model as input for a following step; it only evaluates single-step accuracy against the dataset's own history. (Prior versions of this recipe carried predictions forward through the two-element `state` convention; that convention has been removed -- see [Adding custom datasets](#adding-custom-datasets).) For autoregressive rollout, custom datasets, or more complex inference workflows, we recommend bringing your checkpoints to [Earth2Studio](https://github.com/NVIDIA/earth2studio) for further analysis and visualizations. The [Earth2Studio wrapper for StormCast](https://github.com/NVIDIA/earth2studio/blob/main/earth2studio/models/px/stormcast.py) can be used as a starting point for custom implementations. This also applies to the flow matching models: the `TimeScaledModel` checkpoints of CorrDiff-ERA5-HRRR load with `physicsnemo.core.Module.from_checkpoint` and sample with `diffusion_model_forward` in [`utils/nn.py`](./utils/nn.py) (`formulation="flow"`, `prediction_type="x0"`, `time_shift=32` on the full domain). ## Datasets @@ -297,6 +389,17 @@ invariants (the land and water mask and orography). | Orography | - | Surface | | Land/Water Mask | - | Surface | +### ERA5-HRRR CONUS Downscaling Dataset + +CorrDiff-ERA5-HRRR uses a second dataset class, `data_loader_corrdiff_era5_hrrr.CorrDiffEra5HrrrDataset` in +[`datasets/data_loader_corrdiff_era5_hrrr.py`](./datasets/data_loader_corrdiff_era5_hrrr.py), configured in +`config/dataset/corrdiff_era5_hrrr.yaml`. It serves same-time ERA5 and HRRR pairs over the full +1024 × 1792 pixel CONUS domain, or random patches of it, from one `zarr` store per year for each of +HRRR and ERA5, a netCDF file with the static fields, and a JSON file with the normalization +statistics. Its docstring describes the layout. The ERA5 variables are those of the +[table above](#era5-variables) and the HRRR variables those of the [HRRR table](#hrrr-variables) +without the orography and land/water mask, which the dataset appends to the input instead. + ### Adding Custom Datasets While it is possible to train models on custom datasets by formatting them identically to the Zarr datasets used in the ERA5-HRRR example, a more flexible option is to define a custom dataset object. These datasets must follow the `StormCastDataSource`/`StormCastDataset` interface defined in `datasets/dataset.py`. Refer to the docstrings in that file for a full specification of what the functions must accept and return. You can use the `datasets/mock.py` implementation as a minimal synthetic example and the `datasets/data_loader_hrrr_era5.py` implementation as a real-world example. @@ -414,3 +517,7 @@ Consider the following solutions: - [Elucidating the design space of diffusion-based generative models](https://openreview.net/pdf?id=k7FuTOWMOc7) - [Score-Based Generative Modeling through Stochastic Differential Equations](https://arxiv.org/abs/2011.13456) - [Learning Accurate Storm-Scale Evolution from Observations](https://arxiv.org/abs/2601.17268) +- [Back to Basics: Let Denoising Generative Models Denoise (x-prediction)](https://arxiv.org/abs/2511.13720) +- [Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow](https://arxiv.org/abs/2209.03003) +- [Scaling Rectified Flow Transformers for High-Resolution Image Synthesis (Stable Diffusion 3)](https://arxiv.org/abs/2403.03206) +- [Residual Corrective Diffusion Modeling for Km-scale Atmospheric Downscaling (CorrDiff)](https://arxiv.org/abs/2309.15214) diff --git a/examples/weather/regional_weather_diffusion/config/corrdiff_era5_hrrr.yaml b/examples/weather/regional_weather_diffusion/config/corrdiff_era5_hrrr.yaml new file mode 100644 index 0000000000..e77206d1cc --- /dev/null +++ b/examples/weather/regional_weather_diffusion/config/corrdiff_era5_hrrr.yaml @@ -0,0 +1,75 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# Defaults + +# CorrDiff-ERA5-HRRR: kilometer-scale downscaling of ERA5 to the 3 km HRRR grid +# over the contiguous US with x-prediction rectified flow matching: +# - noisy states x_t = (1 - t) x0 + t eps with t drawn uniformly from [0, 1] +# - the DiT predicts the clean HRRR state x0, the loss is ||x0_hat - x0||^2 +# - trained on random 256 x 256 patches of the 1024 x 1792 HRRR domain + +defaults: + + # Dataset + - dataset/corrdiff_era5_hrrr + + # Model + - model/corrdiff_era5_hrrr + + # Training + - training/default + + # Sampler + - sampler/rectified_flow + + # Hydra + - hydra/default + + - _self_ + +training: + experiment_name: 'corrdiff-era5-hrrr' + num_data_workers: 8 + + checkpoint_freq: 7812 # How often to save the checkpoints (8 GB each), measured in number of training steps (7812 steps = 0.5M samples) + validation_freq: 1562 # How often to record the validation loss, measured in number of training steps (1562 steps = 100k samples) + validation_steps: 8 + validation_plot_variables: ["u10m", "v10m", "t2m", "refc"] + validation_plot_background_channels: ["u10m", "t2m"] + + batch_size: 64 + batch_size_per_gpu: 2 # 2 samples per GPU use about 68 GB; gradients are accumulated on fewer than 32 GPUs + total_train_steps: 781250 # 781250 steps with batch size 64 = 50M samples + clip_grad_norm: 1.0 + + perf: + fp_optimizations: amp-bf16 + + optimizer: + name: "adamw" + lr: 1.0E-4 + betas: [0.9, 0.99] + weight_decay: 0.0 + + scheduler: + name: "CosineAnnealingLR" # cosine decay to zero at total_train_steps after the warmup + lr_rampup_steps: 1000 + + loss: + type: 'flow_matching' + prediction_type: 'x0' # the network predicts the clean state (x-prediction) + t_min: 0.0 # Lowest time (clean data) + t_max: 1.0 # Highest time (pure noise), where the sampling starts diff --git a/examples/weather/regional_weather_diffusion/config/dataset/corrdiff_era5_hrrr.yaml b/examples/weather/regional_weather_diffusion/config/dataset/corrdiff_era5_hrrr.yaml new file mode 100644 index 0000000000..46b0a2802c --- /dev/null +++ b/examples/weather/regional_weather_diffusion/config/dataset/corrdiff_era5_hrrr.yaml @@ -0,0 +1,40 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# Defaults +name: data_loader_corrdiff_era5_hrrr.CorrDiffEra5HrrrDataset # Dataset class module and name, in any module in `datasets` folder + +# Main dataset, see datasets/data_loader_corrdiff_era5_hrrr.py for the expected layout +data_path: 'data' # Path to the dataset +years: [2018, 2019, 2020, 2021, 2022, 2023, 2024, 2025] # Years of data to load + +# Temporal selection: data starts at the HRRR version change in 2018, +# missing samples are automatically excluded +time_range_train: ["2018-07-01T00:00:00", "2025-10-01T00:00:00"] +time_range_valid: ["2025-10-01T00:00:00", "2026-01-01T00:00:00"] + +# Domain +crop_window: [[17, 1041], [3, 1795]] # [[y0, y1], [x0, x1]] window of the HRRR grid, 1024 x 1792 pixels +patch_shape: [256, 256] # Shape of the random training patches; null for the full domain + +# Static fields appended to the background channels +invariant_variables: ["elev_mean", "elev_std", "lsm_mean"] + +# Droppable inputs: during training, each listed ERA5 variable is independently +# replaced by its climatological mean with probability drop_prob. A presence +# flag per variable (1 = kept, 0 = dropped) is passed to the model with the +# scalar conditions, so the model can also run without these inputs. +drop_variables: ["tcwv", "sp"] +drop_prob: 0.5 diff --git a/examples/weather/regional_weather_diffusion/config/model/corrdiff_era5_hrrr.yaml b/examples/weather/regional_weather_diffusion/config/model/corrdiff_era5_hrrr.yaml new file mode 100644 index 0000000000..f8cb4d95d0 --- /dev/null +++ b/examples/weather/regional_weather_diffusion/config/model/corrdiff_era5_hrrr.yaml @@ -0,0 +1,49 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# Defaults + +# Diffusion transformer (DiT) of CorrDiff-ERA5-HRRR: 678M parameters, +# neighborhood attention with 2D rotary position embeddings, trained with +# x-prediction rectified flow matching (training.loss.type: flow_matching). +# These settings match the released checkpoint. + +# Model name +model_name: 'diffusion' + +architecture: dit + +# Preconditioner: the times are multiplied by time_scale before the DiT +# timestep embedding, the network output is the prediction itself +preconditioner: time_scaled +time_scale: 999.0 # Flow-matching times in [0, 1] are multiplied by this + +# Model conditions: pure downscaling, the low-resolution ERA5 fields (with the +# static fields appended by the dataset) condition the model +regression_conditions: [] +diffusion_conditions: ["background"] + +# Model hyperparameters +hyperparameters: + hidden_size: 1152 # Number of channels in the transformer blocks + depth: 28 # Number of transformer blocks + num_heads: 18 # Number of attention heads + patch_size: 2 # Size of the patches embedded as tokens + attn_kernel_size: 31 # Neighborhood size of the attention, in tokens + attention_backend: natten2d_rope # Neighborhood attention with 2D rotary position embeddings + layernorm_backend: torch # Layer normalization backend of the blocks, "torch" or "apex" + attn_kwargs: + qk_norm: true # Normalize the attention queries and keys, stabilizes bf16 training + norm_layer: apex # Apex LayerNorm keeps the bf16 dtype that NATTEN needs for queries, keys, and values diff --git a/examples/weather/regional_weather_diffusion/config/model/stormcast.yaml b/examples/weather/regional_weather_diffusion/config/model/stormcast.yaml index 46cc72de63..060dccac8f 100644 --- a/examples/weather/regional_weather_diffusion/config/model/stormcast.yaml +++ b/examples/weather/regional_weather_diffusion/config/model/stormcast.yaml @@ -23,6 +23,9 @@ regression_conditions: ["background", "invariant"] # list consisting of "backgro # The dataset's "background" tensor includes the previous HRRR timestep. diffusion_conditions: ["regression", "invariant"] # list consisting of "regression", "background", "invariant" (default from StormCast paper) +# Preconditioner of the diffusion network: edm preconditioner +preconditioner: edm + # Model hyperparameters model_type: "SongUNetPosEmbd" # Model class to use # Example overrides for architecture hyperparameters (uncomment to customize) diff --git a/examples/weather/regional_weather_diffusion/config/model/stormscope.yaml b/examples/weather/regional_weather_diffusion/config/model/stormscope.yaml index 229624ee82..e027713072 100644 --- a/examples/weather/regional_weather_diffusion/config/model/stormscope.yaml +++ b/examples/weather/regional_weather_diffusion/config/model/stormscope.yaml @@ -25,6 +25,9 @@ regression_conditions: [] diffusion_conditions: ["background", "invariant"] # list consisting of "regression", "background", "invariant". # The dataset's "background" tensor includes any past-state history. +# Preconditioner of the diffusion network: edm preconditioner +preconditioner: edm + # Model hyperparameters (defaults follow Stormscope) hyperparameters: depth: 16 diff --git a/examples/weather/regional_weather_diffusion/config/sampler/rectified_flow.yaml b/examples/weather/regional_weather_diffusion/config/sampler/rectified_flow.yaml new file mode 100644 index 0000000000..edd1f2673e --- /dev/null +++ b/examples/weather/regional_weather_diffusion/config/sampler/rectified_flow.yaml @@ -0,0 +1,30 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# Defaults + +# Sampler args for physicsnemo.diffusion.samplers.sample with a rectified-flow +# model (training.loss.type: flow_matching): integrates the rectified-flow ODE +# dx/dt = (x_t - x0_hat) / t from pure noise (t = 1) to t = 0. + +name: 'Rectified flow' +args: + num_steps: 18 # Number of solver steps + solver: heun # "heun" (2nd order, 2 model evaluations per step) or "euler" (1st order) + time_shift: 1.0 # Time shift a, t -> a t / (1 + (a - 1) t); use a > 1 to sample a larger domain than the training patches + S_churn: 0. # Set > 0 for stochastic sampling (EDM-style churn, see the EDM paper); the churn of the solvers assumes alpha = 1 and is an approximation for rectified flow + S_min: 0. + S_max: .inf + S_noise: 1 diff --git a/examples/weather/regional_weather_diffusion/config/test_flow_matching.yaml b/examples/weather/regional_weather_diffusion/config/test_flow_matching.yaml new file mode 100644 index 0000000000..d87f85e631 --- /dev/null +++ b/examples/weather/regional_weather_diffusion/config/test_flow_matching.yaml @@ -0,0 +1,78 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# Defaults + +# Tiny DiT downscaling model trained with flow matching on mock data, used by +# test_training.py + +# Defaults +defaults: + + # Dataset + - dataset/mock + + # Model + - model/corrdiff_era5_hrrr + + # Training + - training/default + + # Sampler + - sampler/rectified_flow + + # Hydra + - hydra/default + + - _self_ + +training: + num_data_workers: 1 + validation_plot_variables: ["state_0"] + log_to_wandb: True + wandb_mode: "offline" # not using online logging while testing + log_to_tensorboard: True + + print_progress_freq: 5 # How often to print progress, measured in number of training steps + checkpoint_freq: 5 # How often to save the checkpoints, measured in number of training steps + validation_freq: 5 # how often to record the validation loss, measured in number of training steps + batch_size: 1 + total_train_steps: 10 + + optimizer: + name: "adamw" + betas: [0.9, 0.99] + + scheduler: + name: null + + loss: + type: 'flow_matching' + prediction_type: 'x0' + +sampler: + args: + num_steps: 2 + +dataset: + num_samples: 20 + model_type: downscaling + +model: + hyperparameters: + depth: 2 + hidden_size: 32 + num_heads: 1 + attn_kernel_size: 7 diff --git a/examples/weather/regional_weather_diffusion/config/training/default.yaml b/examples/weather/regional_weather_diffusion/config/training/default.yaml index 08cff04adf..0c0bbb13fd 100644 --- a/examples/weather/regional_weather_diffusion/config/training/default.yaml +++ b/examples/weather/regional_weather_diffusion/config/training/default.yaml @@ -19,7 +19,7 @@ outdir: 'rundir' # Root path under which to save training outputs experiment_name: 'stormcast-training' # Name for the training experiment run_id: '0' # Unique ID to use for this training run -rundir: ./${training.outdir}/${training.experiment_name}/${training.run_id} # Path where experiement outputs will be saved +rundir: ${training.outdir}/${training.experiment_name}/${training.run_id} # Path where experiment outputs will be saved num_data_workers: 4 # Number of dataloader worker threads per proc log_to_wandb: False # Whether or not to log to Weights & Biases (requires wandb account) wandb_mode: "online" # logging mode, "online" or "offline" @@ -95,7 +95,10 @@ channel_loss_weights: null # Loss options loss: - type: 'regression' # Loss type; use 'regression' or 'edm' for the regression and diffusion, respectively + type: 'regression' # Loss type: 'regression', 'edm' (denoising score matching), or 'flow_matching' (rectified flow) + prediction_type: 'x0' # Network output of a diffusion model: 'x0', 'epsilon', 'score', or 'flow'; 'x0' with the EDM preconditioner + # t_min: 0.0 # Flow matching: lowest time (clean data) + # t_max: 1.0 # Flow matching: highest time (1 = pure noise) # sigma_distribution: "lognormal" # EDM noise sampling distribution, "lognormal" or "loguniform" # sigma_data: 0.5 # Standard deviation of the data (set to match your normalized data's std) # # Lognormal distribution parameters (used when sigma_distribution: "lognormal") diff --git a/examples/weather/regional_weather_diffusion/datasets/data_loader_corrdiff_era5_hrrr.py b/examples/weather/regional_weather_diffusion/datasets/data_loader_corrdiff_era5_hrrr.py new file mode 100644 index 0000000000..8695c60c8e --- /dev/null +++ b/examples/weather/regional_weather_diffusion/datasets/data_loader_corrdiff_era5_hrrr.py @@ -0,0 +1,320 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Same-time ERA5 -> HRRR downscaling dataset over the contiguous US.""" + +import json +import os +from typing import Any + +import numpy as np +import pandas as pd +import xarray as xr + +from physicsnemo.utils.zenith_angle import cos_zenith_angle + +from .dataset import StormCastDataset + + +class _CorrDiffEra5HrrrDataset(StormCastDataset): + """Same-time ERA5 -> HRRR downscaling samples over the contiguous US, the + dataset of CorrDiff-ERA5-HRRR (``config/corrdiff_era5_hrrr.yaml``). + + Expects the data under ``data_path`` with the following layout: + + | + | -- hrrr-conus/.zarr # "HRRR" (time, channel, y, x), "latitude", + | # "longitude" (y, x), and "valid_data" (time) + | -- era5-conus/.zarr # "data" (time, channel, latitude, longitude) + | -- invariants/invariants_HRRR.nc # static fields on the HRRR grid + | -- stats/hrrr_era5_stats.json # {"hrrr"|"era5"|"invariant": {var: {"mean", "std"}}} + + Each sample is a dictionary with the following normalized arrays, the + "downscaling" model type of this recipe (``model.diffusion_conditions: + ["background"]``): + + - ``state``: the HRRR target state, shape ``(C_state, H, W)`` + - ``background``: the ERA5 fields bilinearly interpolated to the HRRR grid, + followed by the cosine of the solar zenith angle and the invariants, + shape ``(C_background, H, W)`` + - ``scalar_conditions``: sine and cosine of the time of year followed by + one presence flag per droppable ERA5 variable, shape ``(C_scalar,)`` + + The static fields are part of the background channels, so + :meth:`get_invariants` returns ``None``. + + Training samples are random patches of shape ``patch_shape`` cropped out of + the domain, while the position of each validation patch depends only on + the sample index. With ``patch_shape=None`` each sample covers the full + domain. + + During training, the dataset replaces each ERA5 variable in + ``drop_variables`` by its climatological mean (zero in normalized units) + with probability ``drop_prob``, and sets its presence flag to zero. This + lets the model run without these inputs. Validation and inference never + drop inputs. + + Parameters + ---------- + data_path : str + Root directory of the data. + years : list[int] + Years of data to use. + time_range_train : tuple[str, str] + Training samples have ``time_range_train[0] <= time < time_range_train[1]``. + time_range_valid : tuple[str, str] + Validation samples have ``time_range_valid[0] <= time < time_range_valid[1]``. + invariant_variables : list[str] + Static fields appended to the background channels. + crop_window : tuple[tuple[int, int], tuple[int, int]] + ``((y0, y1), (x0, x1))`` window of the HRRR grid used as the domain. + patch_shape : tuple[int, int] | None + Shape of the patches, or ``None`` to use the full domain. + drop_variables : list[str] + ERA5 variables that the dataset randomly drops during training. + drop_prob : float + Probability of dropping each variable in ``drop_variables``. + train : bool + Whether to serve training or validation samples. + """ + + def __init__( + self, + data_path: str, + years: list[int], + time_range_train: tuple[str, str], + time_range_valid: tuple[str, str], + invariant_variables: list[str] = ["elev_mean", "elev_std", "lsm_mean"], + crop_window: tuple[tuple[int, int], tuple[int, int]] = ((17, 1041), (3, 1795)), + patch_shape: tuple[int, int] | None = (256, 256), + drop_variables: list[str] = [], + drop_prob: float = 0.0, + train: bool = True, + ): + self.train = train + self.crop_window = crop_window + self.crop = tuple(slice(i0, i1) for (i0, i1) in crop_window) + self.domain_shape = tuple(i1 - i0 for (i0, i1) in crop_window) + self.patch_shape = None if patch_shape is None else tuple(patch_shape) + self.invariant_variables = list(invariant_variables) + self.drop_variables = list(drop_variables) + self.drop_prob = drop_prob + + self.hrrr_paths = { + year: os.path.join(data_path, "hrrr-conus", f"{year}.zarr") + for year in years + } + self.era5_paths = { + year: os.path.join(data_path, "era5-conus", f"{year}.zarr") + for year in years + } + self.hrrr_datasets = {} # opened on first access in each dataloader worker + self.era5_datasets = {} + + # Valid sample times (valid HRRR data with matching ERA5 data), + # channel names, and HRRR coordinates + times = [] + for year in years: + with xr.open_zarr(self.hrrr_paths[year]) as ds: + hrrr_times = ds["time"].values[ds["valid_data"].values] + with xr.open_zarr(self.era5_paths[year]) as ds: + era5_times = ds["time"].values + times.append(hrrr_times[np.isin(hrrr_times, era5_times)]) + with xr.open_zarr(self.hrrr_paths[years[0]]) as ds: + self.hrrr_channels = [str(c) for c in ds["channel"].values] + self.lat = ds["latitude"].values[self.crop] + self.lon = ds["longitude"].values[self.crop] + with xr.open_zarr(self.era5_paths[years[0]]) as ds: + self.era5_channels = [str(c) for c in ds["channel"].values] + era5_lat = ds["latitude"].values + era5_lon = ds["longitude"].values + + (t0, t1) = ( + np.datetime64(t) for t in (time_range_train if train else time_range_valid) + ) + self.times = np.concatenate(times) + self.times = self.times[(self.times >= t0) & (self.times < t1)] + + # Fractional ERA5 grid indices of each HRRR pixel, used for the bilinear + # interpolation of ERA5 to the HRRR grid + i = (self.lat - era5_lat[0]) / (era5_lat[1] - era5_lat[0]) + j = (self.lon - era5_lon[0]) / (era5_lon[1] - era5_lon[0]) + if not ( + (i.min() >= 0) + and (i.max() < len(era5_lat) - 1) + and (j.min() >= 0) + and (j.max() < len(era5_lon) - 1) + ): + raise ValueError("The ERA5 grid must cover the HRRR domain.") + self.i0 = i.astype(np.int64) + self.di = (i - self.i0).astype(np.float32) + self.j0 = j.astype(np.int64) + self.dj = (j - self.j0).astype(np.float32) + + # Normalization statistics and the normalized invariants + with open(os.path.join(data_path, "stats", "hrrr_era5_stats.json")) as f: + stats = json.load(f) + (self.hrrr_mean, self.hrrr_std) = _load_stats(stats, "hrrr", self.hrrr_channels) + (self.era5_mean, self.era5_std) = _load_stats(stats, "era5", self.era5_channels) + (inv_mean, inv_std) = _load_stats(stats, "invariant", self.invariant_variables) + invariants_path = os.path.join(data_path, "invariants", "invariants_HRRR.nc") + with xr.open_dataset(invariants_path) as ds: + invariants = np.stack( + [ds[v].values[self.crop] for v in self.invariant_variables] + ) + self.invariants = (invariants - inv_mean) / inv_std + + self.drop_indices = [self.era5_channels.index(v) for v in self.drop_variables] + + def __len__(self) -> int: + return len(self.times) + + def __getitem__(self, idx: int) -> dict[str, Any]: + time = pd.Timestamp(self.times[idx]) + (y0, x0) = self._patch_origin(idx) + (h, w) = self.image_shape() + patch = (slice(y0, y0 + h), slice(x0, x0 + w)) + + # HRRR target state; the patch origin is relative to the crop window + hrrr = self._open_dataset(self.hrrr_datasets, self.hrrr_paths, time.year) + (gy0, gx0) = (self.crop_window[0][0] + y0, self.crop_window[1][0] + x0) + state = hrrr["HRRR"].sel(time=time)[:, gy0 : gy0 + h, gx0 : gx0 + w].values + state = (state - self.hrrr_mean) / self.hrrr_std + + # ERA5 interpolated to the HRRR grid, then the zenith angle and invariants + era5 = self._open_dataset(self.era5_datasets, self.era5_paths, time.year) + era5 = era5["data"].sel(time=time).values + era5 = _bilinear_interpolation( + era5, self.i0[patch], self.di[patch], self.j0[patch], self.dj[patch] + ) + era5 = (era5 - self.era5_mean) / self.era5_std + presence_flags = self._drop_variables(era5) + cos_zenith = cos_zenith_angle( + time.to_pydatetime(), lon=self.lon[patch], lat=self.lat[patch] + ) + background = np.concatenate( + [era5, cos_zenith[None], self.invariants[(slice(None),) + patch]] + ) + + # Time of year (one period = 365.25 days) and the input presence flags + start_of_year = pd.Timestamp(year=time.year, month=1, day=1) + year_phase = 2 * np.pi * ((time - start_of_year) / pd.Timedelta(days=365.25)) + scalar_conditions = np.concatenate( + [[np.sin(year_phase), np.cos(year_phase)], presence_flags] + ) + + return { + "background": background.astype(np.float32), + "state": state.astype(np.float32), + "scalar_conditions": scalar_conditions.astype(np.float32), + } + + def _patch_origin(self, idx: int) -> tuple[int, int]: + """Top-left corner of the patch within the domain.""" + if self.patch_shape is None: + return (0, 0) + max_origin = [ + d - p + 1 + for (d, p) in zip(self.domain_shape, self.patch_shape, strict=True) + ] + if self.train: + return tuple(int(np.random.randint(m)) for m in max_origin) + rng = np.random.default_rng(idx) # same validation patches every time + return tuple(int(rng.integers(m)) for m in max_origin) + + def _drop_variables(self, era5: np.ndarray) -> np.ndarray: + """Randomly zero out the droppable ERA5 channels in place (training only) + and return their presence flags (1 = kept, 0 = dropped).""" + flags = np.ones(len(self.drop_indices), dtype=np.float32) + if self.train: + for k, c in enumerate(self.drop_indices): + if np.random.random() < self.drop_prob: + era5[c] = 0.0 + flags[k] = 0.0 + return flags + + @staticmethod + def _open_dataset(datasets: dict, paths: dict, year: int) -> xr.Dataset: + if year not in datasets: + datasets[year] = xr.open_dataset( + paths[year], engine="zarr", chunks=None, mask_and_scale=False + ) + return datasets[year] + + def background_channels(self) -> list[str]: + """ERA5 channels, the cosine of the solar zenith angle, and the invariants.""" + return self.era5_channels + ["cos_zenith"] + self.invariant_variables + + def state_channels(self) -> list[str]: + """HRRR channels.""" + return self.hrrr_channels + + def scalar_condition_channels(self) -> list[str]: + """Time of year and the presence flags of the droppable variables.""" + return ["sin_year", "cos_year"] + [f"{v}_present" for v in self.drop_variables] + + def image_shape(self) -> tuple[int, int]: + """Shape (height, width) of the samples.""" + return self.domain_shape if self.patch_shape is None else self.patch_shape + + def latitude(self) -> np.ndarray: + """Latitude of the full domain. It has the shape ``image_shape()`` only + with ``patch_shape=None``, which inference uses.""" + return self.lat + + def longitude(self) -> np.ndarray: + """Longitude of the full domain. It has the shape ``image_shape()`` only + with ``patch_shape=None``, which inference uses.""" + return self.lon + + def denormalize_background(self, x: np.ndarray) -> np.ndarray: + """Convert the ERA5 channels of the background to physical units, + leaving the zenith angle and invariant channels unchanged.""" + x = x.copy() + n = len(self.era5_channels) + x[..., :n, :, :] = x[..., :n, :, :] * self.era5_std + self.era5_mean + return x + + def denormalize_state(self, x: np.ndarray) -> np.ndarray: + """Convert the state from normalized data to physical units.""" + return x * self.hrrr_std + self.hrrr_mean + + +class CorrDiffEra5HrrrDataset(_CorrDiffEra5HrrrDataset): + """The dataset as the training recipe constructs it, from the ``dataset`` + config (``name: data_loader_corrdiff_era5_hrrr.CorrDiffEra5HrrrDataset``) and a ``train`` flag.""" + + def __init__(self, params: dict, train: bool): + super().__init__(train=train, **params) + + +def _load_stats( + stats: dict, group: str, variables: list[str] +) -> tuple[np.ndarray, np.ndarray]: + """Mean and standard deviation of ``variables``, shaped ``(C, 1, 1)``.""" + mean = np.array([stats[group][v]["mean"] for v in variables], dtype=np.float32) + std = np.array([stats[group][v]["std"] for v in variables], dtype=np.float32) + return (mean[:, None, None], std[:, None, None]) + + +def _bilinear_interpolation( + x: np.ndarray, i0: np.ndarray, di: np.ndarray, j0: np.ndarray, dj: np.ndarray +) -> np.ndarray: + """Bilinearly interpolate ``x`` of shape ``(C, Ny, Nx)`` at the fractional + grid indices ``(i0 + di, j0 + dj)``.""" + f0 = x[:, i0, j0] + dj * (x[:, i0, j0 + 1] - x[:, i0, j0]) + f1 = x[:, i0 + 1, j0] + dj * (x[:, i0 + 1, j0 + 1] - x[:, i0 + 1, j0]) + return f0 + di * (f1 - f0) diff --git a/examples/weather/regional_weather_diffusion/requirements.txt b/examples/weather/regional_weather_diffusion/requirements.txt index 6b56726849..a4ecc88b10 100644 --- a/examples/weather/regional_weather_diffusion/requirements.txt +++ b/examples/weather/regional_weather_diffusion/requirements.txt @@ -6,4 +6,4 @@ matplotlib xarray wandb dask -natten \ No newline at end of file +natten diff --git a/examples/weather/regional_weather_diffusion/test_training.py b/examples/weather/regional_weather_diffusion/test_training.py index c699a1bd5e..d029450e3b 100644 --- a/examples/weather/regional_weather_diffusion/test_training.py +++ b/examples/weather/regional_weather_diffusion/test_training.py @@ -14,24 +14,39 @@ # See the License for the specific language governing permissions and # limitations under the License. +import json import os from contextlib import nullcontext from pathlib import Path from typing import Literal +import numpy as np +import pandas as pd import pytest import torch import train +import xarray as xr from hydra import compose, initialize -from omegaconf import DictConfig +from omegaconf import DictConfig, OmegaConf +from tensordict import TensorDict from torch.distributed.checkpoint.state_dict import StateDictOptions, get_state_dict from torch.distributed.tensor import DTensor from utils import parallel, trainer +from utils.nn import TimeScaledModel, diffusion_model_forward +from physicsnemo.core import Module +from physicsnemo.diffusion.metrics.losses import MSEDSMLoss +from physicsnemo.diffusion.noise_schedulers import RectifiedFlowNoiseScheduler from physicsnemo.distributed import DistributedManager DistributedManager.initialize() +# The rectified-flow scheduler warns about t_max = 1; the x-prediction loss and +# the flow used for sampling stay finite there +pytestmark = pytest.mark.filterwarnings( + "ignore:RectifiedFlowNoiseScheduler was constructed with t_max=1.0" +) + # Retrieve and fixture configs def _load_config(config_name: str) -> DictConfig: @@ -58,6 +73,12 @@ def cfg_diffusion_unet(): return _load_config(config_name="test_diffusion_unet.yaml") +@pytest.fixture +def cfg_flow_matching(): + """Return the flow matching (CorrDiff-ERA5-HRRR style) test configuration.""" + return _load_config(config_name="test_flow_matching.yaml") + + def _setup_rundir(tmp_path, num_procs): """Create and share a temporary run directory across all test ranks.""" # Set up rundir in the temporary directory @@ -866,3 +887,338 @@ def test_datapipes_backend_training( rundir, "checkpoints_diffusion", "EDMPreconditioner.0.10.mdlus" ) assert os.path.isfile(ckpt_path), "Diffusion checkpoint not found" + + +# ============================================================================= +# Flow matching (CorrDiff-ERA5-HRRR) +# ============================================================================= + + +# Launch configurations (world_size, domain_parallel_size, batch_size): a test +# runs the cases matching the number of processes pytest was launched with. +# ``data_parallel`` uses FSDP only, ``domain_parallel`` shards the spatial +# dimension with ShardTensor, ``data_domain_parallel`` combines both on a +# (2, 2) mesh. +_PARALLEL_CONFIGS = pytest.mark.parametrize( + "world_size, domain_parallel_size, batch_size", + [(1, 1, 1), (2, 1, 2), (2, 2, 1), (4, 2, 2)], + ids=["single", "data_parallel", "domain_parallel", "data_domain_parallel"], +) + + +def _skip_unless_world_size(world_size: int) -> DistributedManager: + """Skip the test unless it was launched with ``world_size`` processes.""" + dist = DistributedManager() + if dist.world_size != world_size: + pytest.skip( + f"Skipping: this configuration requires {world_size} process(es), " + f"current: {dist.world_size}." + ) + return dist + + +@_PARALLEL_CONFIGS +@pytest.mark.parametrize("net_architecture", ["dit", "unet"]) +@pytest.mark.parametrize("prediction_type", ["x0", "epsilon", "flow"]) +@pytest.mark.parametrize("fp_optimizations", ["fp32", "amp-bf16"]) +def test_flow_matching_training( + tmp_path: Path, + cfg_flow_matching: DictConfig, + *, + world_size: int, + domain_parallel_size: int, + batch_size: int, + net_architecture: Literal["unet", "dit"], + prediction_type: Literal["x0", "epsilon", "flow"], + fp_optimizations: Literal["fp32", "amp-bf16"], +): + """Test that flow matching training runs with the DiT and the U-Net, with + x0, epsilon, and flow prediction, on 1, 2 and 4 processes (data and + domain parallelism), and writes the TimeScaledModel checkpoint and + validation plots.""" + dist = _skip_unless_world_size(world_size) + rundir = _setup_rundir(tmp_path, dist.world_size) + + cfg = cfg_flow_matching.copy() + cfg.training.rundir = rundir + cfg.training.batch_size = batch_size + cfg.training.domain_parallel_size = domain_parallel_size + cfg.training.perf.fp_optimizations = fp_optimizations + cfg.training.loss.prediction_type = prediction_type + cfg.model.architecture = net_architecture + if net_architecture == "unet": + cfg.dataset.num_scalar_cond_channels = 0 # scalar conditions need the DiT + cfg.model.hyperparameters = {"model_channels": 8, "channel_mult": [1, 1]} + train.main(cfg) + if dist.world_size > 1: + torch.distributed.barrier() + + ckpt_path = os.path.join( + rundir, "checkpoints_diffusion", "TimeScaledModel.0.10.mdlus" + ) + assert os.path.isfile(ckpt_path), "Flow matching checkpoint not found" + assert isinstance(Module.from_checkpoint(ckpt_path), TimeScaledModel) + assert os.path.isdir(os.path.join(rundir, "images", "state_0")) + + +@_PARALLEL_CONFIGS +@pytest.mark.parametrize("prediction_type", ["x0", "epsilon"]) +def test_edm_time_scaled_training( + tmp_path: Path, + cfg_diffusion: DictConfig, + *, + world_size: int, + domain_parallel_size: int, + batch_size: int, + prediction_type: Literal["x0", "epsilon"], +): + """Test that the time-scaled preconditioner also trains with the EDM loss, + predicting the clean state or the noise: the preconditioner, the loss type, + and the prediction type are independent settings. Runs on 1, 2 and 4 + processes (data and domain parallelism).""" + dist = _skip_unless_world_size(world_size) + cfg = cfg_diffusion.copy() + cfg.training.rundir = _setup_rundir(tmp_path, dist.world_size) + cfg.training.batch_size = batch_size + cfg.training.domain_parallel_size = domain_parallel_size + cfg.model.preconditioner = "time_scaled" # with the default time_scale + cfg.training.loss.prediction_type = prediction_type + if "regression" in cfg.model.diffusion_conditions: + cfg.model.diffusion_conditions.remove("regression") + train.main(cfg) + if dist.world_size > 1: + torch.distributed.barrier() + ckpt_path = os.path.join( + cfg.training.rundir, "checkpoints_diffusion", "TimeScaledModel.0.10.mdlus" + ) + assert os.path.isfile(ckpt_path), "Diffusion checkpoint not found" + + +def test_edm_preconditioner_requires_x0(cfg_diffusion: DictConfig): + """Test that the EDM preconditioner rejects other prediction types.""" + cfg = cfg_diffusion.copy() + cfg.training.loss.prediction_type = "epsilon" + with pytest.raises(ValueError, match="predicts x0"): + trainer.MainConfig(**OmegaConf.to_container(cfg, resolve=True)) + + +def test_flow_matching_x0_loss_equals_mse( + tmp_path: Path, cfg_flow_matching: DictConfig +): + """Test that the x0-prediction flow matching loss equals the x0-prediction + loss of MSEDSMLoss, ||x0_hat - x0||^2, for sampled times and at the + endpoints t = 0 and t = 1.""" + dist = DistributedManager() + if dist.world_size > 1: + pytest.skip("Skipping: single-process test.") + cfg = cfg_flow_matching.copy() + cfg.training.rundir = _setup_rundir(tmp_path, dist.world_size) + tr = trainer.Trainer(cfg) + x0_mse = MSEDSMLoss( + tr.net, tr.train_noise_scheduler, prediction_type="x0", reduction="none" + ) + device = tr.device + (h, w) = cfg.dataset.image_size + x0 = torch.randn(2, cfg.dataset.num_state_channels, h, w, device=device) + condition = TensorDict( + { + "cond_concat": torch.randn( + 2, cfg.dataset.num_background_channels, h, w, device=device + ), + "cond_vec": torch.randn( + 2, cfg.dataset.num_scalar_cond_channels, device=device + ), + }, + batch_size=[2], + device=device, + ) + for t in (None, torch.tensor([0.0, 1.0], device=device)): + with torch.no_grad(): + torch.manual_seed(0) + loss = tr.loss_fn(x0, torch.ones_like(x0), t=t, condition=condition) + torch.manual_seed(0) # same times and noise as in the training loss + expected = x0_mse(x0, t=t, condition=condition) + assert torch.isfinite(loss).all() + assert torch.equal(loss, expected) + + +class _GaussianX0Predictor(torch.nn.Module): + """Exact x0-predictor E[x0 | x_t] for data x0 ~ N(mean, std^2) on the + rectified-flow path x_t = (1 - t) x0 + t eps.""" + + def __init__(self, mean: float, std: float): + super().__init__() + self.mean = mean + self.var = std**2 + + def forward(self, x, t, condition=None): + alpha = (1 - t).reshape(-1, 1, 1, 1) + sigma = t.reshape(-1, 1, 1, 1) + gain = alpha * self.var / (alpha**2 * self.var + sigma**2) + return self.mean + gain * (x - alpha * self.mean) + + +@pytest.mark.parametrize( + "solver, num_steps, time_shift", + [("heun", 32, 1.0), ("heun", 64, 4.0), ("euler", 128, 1.0)], +) +def test_rectified_flow_sampling_from_pure_noise( + solver: str, num_steps: int, time_shift: float +): + """Test that rectified-flow sampling from pure noise (t = 1) with the exact + x0-predictor of Gaussian data reproduces the data distribution, with and + without a time shift.""" + device = DistributedManager().device + (mean, std) = (2.0, 0.5) + model = _GaussianX0Predictor(mean, std).to(device) + torch.manual_seed(0) + x = diffusion_model_forward( + model, + torch.zeros(4, 1, 64, 64, device=device), # the predictor ignores the condition + shape=(4, 1, 64, 64), + scheduler=RectifiedFlowNoiseScheduler(t_max=1.0), + sampler_args={ + "num_steps": num_steps, + "solver": solver, + "time_shift": time_shift, + }, + formulation="flow", + prediction_type="x0", + ) + assert torch.isfinite(x).all() + assert abs(x.mean().item() - mean) < 0.02 + assert abs(x.std().item() - std) < 0.02 + + +def _write_corrdiff_era5_hrrr_data(path: Path) -> dict: + """Write a small dataset in the layout that CorrDiffEra5HrrrDataset expects. + The ERA5 fields are linear in latitude and longitude, so their bilinear + interpolation to the HRRR grid is exact.""" + times = pd.date_range("2025-01-01", periods=6, freq="h") + (y, x) = np.meshgrid(np.arange(12), np.arange(16), indexing="ij") + lat = (35.0 + 0.5 * y + 0.1 * x).astype(np.float32) + lon = (-105.0 + 0.6 * x + 0.05 * y).astype(np.float32) + rng = np.random.default_rng(0) + hrrr = rng.normal(size=(6, 2, 12, 16)).astype(np.float32) + xr.Dataset( + { + "HRRR": (("time", "channel", "y", "x"), hrrr), + "latitude": (("y", "x"), lat), + "longitude": (("y", "x"), lon), + "valid_data": ("time", np.array([True] * 5 + [False])), # last hour invalid + }, + coords={"time": times, "channel": np.array(["u10m", "t2m"])}, + ).to_zarr(path / "hrrr-conus" / "2025.zarr") + + # ERA5 fields c0 + c1 * lat + c2 * lon, without the hour at index 4 + (era5_lat, era5_lon) = (np.arange(50.0, 29.0, -1.0), np.arange(-110.0, -89.0)) + coefs = np.array([[1.0, 0.2, -0.3], [2.0, -0.1, 0.4], [3.0, 0.05, 0.1]]) + fields = ( + coefs[:, 0, None, None] + + coefs[:, 1, None, None] * era5_lat[:, None] + + coefs[:, 2, None, None] * era5_lon[None, :] + ) + era5_times = times[[0, 1, 2, 3, 5]] + xr.Dataset( + { + "data": ( + ("time", "channel", "latitude", "longitude"), + np.repeat(fields[None], len(era5_times), axis=0).astype(np.float32), + ) + }, + coords={ + "time": era5_times, + "channel": np.array(["u10m", "tcwv", "sp"]), + "latitude": era5_lat, + "longitude": era5_lon, + }, + ).to_zarr(path / "era5-conus" / "2025.zarr") + + invariant_names = ["elev_mean", "elev_std", "lsm_mean"] + invariants = rng.normal(size=(3, 12, 16)).astype(np.float32) + os.makedirs(path / "invariants") + xr.Dataset( + {v: (("y", "x"), invariants[k]) for (k, v) in enumerate(invariant_names)} + ).to_netcdf(path / "invariants" / "invariants_HRRR.nc") + + stats = { + "hrrr": {"u10m": {"mean": 0.5, "std": 2.0}, "t2m": {"mean": -1.0, "std": 4.0}}, + "era5": {v: {"mean": 1.0, "std": 3.0} for v in ["u10m", "tcwv", "sp"]}, + "invariant": {v: {"mean": 0.1, "std": 2.0} for v in invariant_names}, + } + os.makedirs(path / "stats") + with open(path / "stats" / "hrrr_era5_stats.json", "w") as f: + json.dump(stats, f) + return { + "hrrr": hrrr, + "lat": lat, + "lon": lon, + "coefs": coefs, + "invariants": invariants, + "times": times, + } + + +def test_corrdiff_era5_hrrr_dataset(tmp_path: Path): + """Test time filtering, cropping, interpolation, normalization, patches, + and presence flags of CorrDiffEra5HrrrDataset on a small synthetic dataset.""" + from datasets.data_loader_corrdiff_era5_hrrr import CorrDiffEra5HrrrDataset + + dist = DistributedManager() + if dist.world_size > 1: + pytest.skip("Skipping: single-process test.") + data = _write_corrdiff_era5_hrrr_data(tmp_path) + params = dict( + data_path=str(tmp_path), + years=[2025], + time_range_train=["2025-01-01T00:00:00", "2025-01-01T02:00:00"], + time_range_valid=["2025-01-01T02:00:00", "2025-01-02T00:00:00"], + crop_window=[[1, 11], [2, 14]], + drop_variables=["tcwv", "sp"], + drop_prob=1.0, + ) + crop = (slice(1, 11), slice(2, 14)) + + # Validation on the full domain: the dataset skips the invalid HRRR hour + # and the hour missing in ERA5, and keeps every input + valid = CorrDiffEra5HrrrDataset(dict(params, patch_shape=None), train=False) + assert list(valid.times) == list(data["times"][[2, 3]]) + assert valid.image_shape() == (10, 12) + assert len(valid.background_channels()) == 3 + 1 + 3 + assert valid.get_invariants() is None + item = valid[0] + np.testing.assert_allclose( + item["state"], + (data["hrrr"][2][:, crop[0], crop[1]] - np.array([0.5, -1.0])[:, None, None]) + / np.array([2.0, 4.0])[:, None, None], + rtol=1e-6, + ) + (lat, lon) = (data["lat"][crop], data["lon"][crop]) + coefs = data["coefs"] + era5 = ( + coefs[:, 0, None, None] + + coefs[:, 1, None, None] * lat + + coefs[:, 2, None, None] * lon + ) + np.testing.assert_allclose(item["background"][:3], (era5 - 1.0) / 3.0, atol=1e-4) + np.testing.assert_allclose( + item["background"][4:], + (data["invariants"][:, crop[0], crop[1]] - 0.1) / 2.0, + rtol=1e-6, + ) + year_phase = 2 * np.pi * (2 / 24) / 365.25 # 02:00 on January 1 + np.testing.assert_allclose( + item["scalar_conditions"], + [np.sin(year_phase), np.cos(year_phase), 1, 1], + rtol=1e-5, + ) + + # Training patches: with drop_prob = 1, the dataset replaces tcwv and sp by + # their climatological mean (zero) and sets their presence flags to zero + train_ds = CorrDiffEra5HrrrDataset(dict(params, patch_shape=[4, 4]), train=True) + assert len(train_ds) == 2 + item = train_ds[1] + assert item["state"].shape == (2, 4, 4) + assert item["background"].shape == (7, 4, 4) + assert np.all(item["background"][1:3] == 0) + np.testing.assert_array_equal(item["scalar_conditions"][2:], [0, 0]) diff --git a/examples/weather/regional_weather_diffusion/train.py b/examples/weather/regional_weather_diffusion/train.py index 02c2a86277..8144c4a890 100644 --- a/examples/weather/regional_weather_diffusion/train.py +++ b/examples/weather/regional_weather_diffusion/train.py @@ -44,6 +44,11 @@ def main(cfg: DictConfig) -> None: warnings.filterwarnings( "ignore", message="^.*You are importing from 'physicsnemo.experimental'.*$" ) + # Flow matching trains up to t = 1 (pure noise), which is safe: the loss and + # the flow used for sampling stay finite there + warnings.filterwarnings( + "ignore", message="^RectifiedFlowNoiseScheduler was constructed with t_max=1.0" + ) # Initialize DistributedManager.initialize() diff --git a/examples/weather/regional_weather_diffusion/utils/config.py b/examples/weather/regional_weather_diffusion/utils/config.py index f97770eea0..d7ec4cce5c 100644 --- a/examples/weather/regional_weather_diffusion/utils/config.py +++ b/examples/weather/regional_weather_diffusion/utils/config.py @@ -42,6 +42,13 @@ class ModelConfig: regression_weights: str | None = ( None # file with weights of pre-trained regression model, set to None when training regression or diffusion-only model ) + preconditioner: Literal["edm", "time_scaled"] = ( + "edm" # wrapper of the diffusion network: "edm" (EDM input/output scaling, the network output is x0) + # or "time_scaled" (times multiplied by time_scale, the network output is training.loss.prediction_type) + ) + time_scale: float = Field( + default=999.0, gt=0.0 + ) # factor applied to the times by the "time_scaled" preconditioner hyperparameters: dict[str, Any] = Field( default={} ) # additional parameters to be passed to model @@ -101,7 +108,18 @@ class SchedulerConfig: class LossConfig: """Loss configuration: cfg.training.loss""" - type: Literal["regression", "edm"] = "edm" + type: Literal["regression", "edm", "flow_matching"] = ( + "edm" # "regression", "edm" (denoising score matching), or "flow_matching" (rectified flow) + ) + prediction_type: Literal["x0", "epsilon", "score", "flow"] = ( + "x0" # network output: clean state, noise, score, or flow (velocity); must be "x0" with the EDM preconditioner + ) + # Flow matching parameters (used when type: "flow_matching") + t_min: float = Field(default=0.0, ge=0.0) # Lowest flow-matching time (clean data) + t_max: float = Field( + default=1.0, le=1.0 + ) # Highest flow-matching time (1 = pure noise) + # EDM parameters (used when type: "edm") sigma_distribution: Literal["lognormal", "loguniform"] = "lognormal" sigma_data: float | list[float] = 0.5 P_mean: float = -1.2 # Center of the lognormal noise distribution @@ -115,6 +133,8 @@ class LossConfig: @model_validator(mode="after") def _check_sigmas(self): + if self.t_min >= self.t_max: + raise ValueError("t_min must be smaller than t_max") if self.sigma_min >= self.sigma_max: raise ValueError("sigma_min must be smaller than sigma_max") sigma_data = ( @@ -285,3 +305,17 @@ class MainConfig: model: ModelConfig training: TrainingConfig sampler: SamplerConfig + + @model_validator(mode="after") + def _check_preconditioner(self): + # The EDM preconditioner scales the network output into the clean state + if ( + self.training.loss.type != "regression" + and self.model.preconditioner == "edm" + and self.training.loss.prediction_type != "x0" + ): + raise ValueError( + "model.preconditioner 'edm' predicts x0; set training.loss.prediction_type " + "to 'x0' or use model.preconditioner 'time_scaled'" + ) + return self diff --git a/examples/weather/regional_weather_diffusion/utils/loss.py b/examples/weather/regional_weather_diffusion/utils/loss.py index 86674a3d1a..adab2c021a 100644 --- a/examples/weather/regional_weather_diffusion/utils/loss.py +++ b/examples/weather/regional_weather_diffusion/utils/loss.py @@ -17,9 +17,13 @@ """Loss functions for StormCast training. -Diffusion training uses :class:`~physicsnemo.diffusion.metrics.losses.MSEDSMLoss` +Diffusion training uses :class:`~physicsnemo.diffusion.metrics.losses.WeightedMSEDSMLoss` with an :class:`~physicsnemo.diffusion.noise_schedulers.EDMNoiseScheduler` or -:class:`~physicsnemo.diffusion.noise_schedulers.EDMLogUniformNoiseScheduler`. +:class:`~physicsnemo.diffusion.noise_schedulers.EDMLogUniformNoiseScheduler` +(``training.loss.type: edm``), or +:class:`~physicsnemo.diffusion.metrics.losses.WeightedFlowMatchingLoss` with a +:class:`~physicsnemo.diffusion.noise_schedulers.RectifiedFlowNoiseScheduler` +(``training.loss.type: flow_matching``). When domain parallelism is active, the caller wraps the scheduler with :class:`~physicsnemo.diffusion.DomainParallelNoiseScheduler` (via :meth:`~utils.parallel.ParallelHelper.make_domain_parallel_scheduler`) before @@ -31,23 +35,31 @@ import numpy as np import torch +from utils.nn import flow_conversions +from physicsnemo.diffusion.metrics.losses import ( + WeightedFlowMatchingLoss, + WeightedMSEDSMLoss, +) from physicsnemo.diffusion.noise_schedulers import ( EDMLogUniformNoiseScheduler, EDMNoiseScheduler, + NoiseScheduler, + RectifiedFlowNoiseScheduler, ) def build_noise_scheduler( loss_cfg, logger=None, -) -> EDMNoiseScheduler: - """Create an :class:`~physicsnemo.diffusion.noise_schedulers.EDMNoiseScheduler` from Hydra loss config. +) -> NoiseScheduler: + """Create the training noise scheduler from the Hydra loss config. Parameters ---------- loss_cfg : object - Loss config with ``sigma_distribution``, ``sigma_data``, and + Loss config. With ``type: "flow_matching"`` it must provide ``t_min`` + and ``t_max``; otherwise ``sigma_distribution``, ``sigma_data``, and distribution-specific params (``P_mean``/``P_std`` or ``sigma_min``/``sigma_max``). logger : optional @@ -55,8 +67,13 @@ def build_noise_scheduler( Returns ------- - EDMNoiseScheduler + NoiseScheduler + A :class:`~physicsnemo.diffusion.noise_schedulers.RectifiedFlowNoiseScheduler` + for flow matching, otherwise an EDM scheduler. """ + if loss_cfg.type == "flow_matching": + return RectifiedFlowNoiseScheduler(t_min=loss_cfg.t_min, t_max=loss_cfg.t_max) + sigma_data = loss_cfg.sigma_data if isinstance(sigma_data, Sequence): sigma_data = list(sigma_data) @@ -86,6 +103,72 @@ def build_noise_scheduler( ) +def build_diffusion_loss( + model: torch.nn.Module, noise_scheduler: NoiseScheduler, loss_cfg +) -> WeightedMSEDSMLoss | WeightedFlowMatchingLoss: + """Create the diffusion training loss for the configured loss type. + + Both losses share the call signature + ``loss_fn(x0, weight, t=..., condition=..., **model_kwargs)`` and return + the per-element weighted loss. ``loss_cfg.prediction_type`` tells the loss + what the network outputs (``x0``, ``epsilon``, ``score``, or ``flow``); + the network wrapper (``model.preconditioner``) must produce that quantity. + + Parameters + ---------- + model : torch.nn.Module + The wrapped network, see :func:`~utils.nn.wrap_diffusion_network`. + noise_scheduler : NoiseScheduler + Training scheduler from :func:`build_noise_scheduler`, possibly + wrapped for domain parallelism. + loss_cfg : object + Loss config with ``type`` and ``prediction_type``. + + Returns + ------- + WeightedMSEDSMLoss | WeightedFlowMatchingLoss + """ + prediction_type = loss_cfg.prediction_type + if loss_cfg.type == "edm": + # Predictions other than x0 reach the loss through the scheduler's + # conversion to x0 (score_to_x0, epsilon_to_x0, flow_to_x0). The + # domain-parallel wrapper only implements the sampling and loss + # methods; the element-wise conversions come from the wrapped scheduler + to_x0 = {} + if prediction_type != "x0": + converter = getattr(noise_scheduler, "inner_scheduler", noise_scheduler) + to_x0[f"{prediction_type}_to_x0_fn"] = getattr( + converter, f"{prediction_type}_to_x0" + ) + return WeightedMSEDSMLoss( + model, + noise_scheduler, + prediction_type=prediction_type, + reduction="none", + **to_x0, + ) + if loss_cfg.type != "flow_matching": + raise ValueError("training.loss.type must be 'edm' or 'flow_matching'") + + # The loss converts the prediction and the target to the flow with the + # scheduler, except for x-prediction, which compares the predicted and true + # clean states directly: ||x0_hat - x0||^2 (the scheduler's x0_to_flow would + # weight the error by 1 / t^2) + to_flow = flow_conversions(noise_scheduler) + conversions = {"x0_to_flow_fn": to_flow["x0"]} + if prediction_type == "x0": + conversions["x0_to_flow_fn"] = lambda x0, x_t, t: x0 + elif prediction_type in ("score", "epsilon"): + conversions[f"{prediction_type}_to_flow_fn"] = to_flow[prediction_type] + return WeightedFlowMatchingLoss( + model, + noise_scheduler, + prediction_type=prediction_type, + reduction="none", + **conversions, + ) + + class RegressionLoss: """MSE loss for a regression model. @@ -148,7 +231,10 @@ class SigmaBinTracker: device : torch.device Device for accumulator tensors. loss_type : str - ``"regression"`` or ``"edm"``. Tracking is disabled for regression. + ``"regression"``, ``"edm"``, or ``"flow_matching"``. Tracking is + disabled for regression. For flow matching the bins divide the + uniform time distribution on ``[t_min, t_max]`` instead of the sigma + distribution. """ def __init__(self, loss_cfg, device: torch.device, loss_type: str = "edm"): @@ -162,6 +248,10 @@ def __init__(self, loss_cfg, device: torch.device, loss_type: str = "edm"): if len(loss_cfg.sigma_bin_edges) >= 2: edges = np.asarray(loss_cfg.sigma_bin_edges, dtype=np.float64) + elif loss_type == "flow_matching": + edges = np.linspace( + loss_cfg.t_min, loss_cfg.t_max, int(loss_cfg.sigma_bin_count) + 1 + ) else: n_edges = int(loss_cfg.sigma_bin_count) + 1 if loss_cfg.sigma_distribution == "loguniform": diff --git a/examples/weather/regional_weather_diffusion/utils/nn.py b/examples/weather/regional_weather_diffusion/utils/nn.py index 06986ae7b4..51271b0919 100644 --- a/examples/weather/regional_weather_diffusion/utils/nn.py +++ b/examples/weather/regional_weather_diffusion/utils/nn.py @@ -33,6 +33,94 @@ import utils.apex # do not remove, enables Apex LayerNorm with ShardTensor +class TimeScaledModel(Module): + r"""Scale flow-matching times before passing them to the backbone. + + Flow-matching times live in :math:`[0, 1]`, while the sinusoidal timestep + embedders of the DiT and U-Net resolve inputs spanning roughly + :math:`[0, 1000]`. The scale is part of the saved model, so training and + 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. + + Parameters + ---------- + model : Module + Backbone called as ``model(x, t, condition=condition)``. + time_scale : float, optional + Factor applied to the time before calling the backbone, by default 999. + + Forward + ------- + x : torch.Tensor + Noisy state of shape :math:`(B, C, H, W)`. + t : torch.Tensor + Flow-matching times of shape :math:`(B,)`. + condition : torch.Tensor | TensorDict | None, optional + Condition passed to the backbone. + **model_kwargs : Any + Other keyword arguments passed to the backbone. + + Outputs + ------- + torch.Tensor + Backbone output of shape :math:`(B, C, H, W)`. + """ + + def __init__(self, model: Module, time_scale: float = 999.0): + super().__init__() + self.model = model + self.time_scale = time_scale + + def forward( + self, + x: torch.Tensor, + t: torch.Tensor, + condition: torch.Tensor | TensorDict | None = None, + **model_kwargs: Any, + ) -> torch.Tensor: + if not torch.compiler.is_compiling(): + if x.ndim != 4: + raise ValueError( + f"Expected tensor of shape (B, C, H, W) but got tensor of shape {tuple(x.shape)}" + ) + if t.shape != (x.shape[0],): + raise ValueError( + f"Expected tensor of shape ({x.shape[0]},) but got tensor of shape {tuple(t.shape)}" + ) + # Scale in float32: a reduced precision t would quantize the scaled + # time to a coarser grid than in training + return self.model( + x, t.float() * self.time_scale, condition=condition, **model_kwargs + ) + + +def wrap_diffusion_network( + net: torch.nn.Module, + preconditioner: Literal["edm", "time_scaled"], + time_scale: float = 999.0, +) -> EDMPreconditioner | TimeScaledModel: + """Wrap a network with the configured preconditioner (``model.preconditioner``). + + Args: + net: the backbone, called as ``net(x, t, condition=condition)`` + preconditioner: "edm" applies the EDM input, output, and noise-level + scalings, so the wrapped network predicts the clean state x0; + "time_scaled" only multiplies the times by ``time_scale``, so the + wrapped network predicts ``training.loss.prediction_type`` + time_scale: factor applied to the times by the "time_scaled" preconditioner + Returns: + the wrapped network + """ + wrapped = ConcatConditionWrapper(net) + if preconditioner == "edm": + return EDMPreconditioner(model=wrapped) + elif preconditioner == "time_scaled": + return TimeScaledModel(wrapped, time_scale=time_scale) + raise ValueError("preconditioner must be 'edm' or 'time_scaled'") + + def get_preconditioned_unet( name: str, target_channels: int, @@ -43,8 +131,10 @@ def get_preconditioned_unet( lead_time_channels: int = 4, amp_mode: bool = False, use_apex_gn: bool = False, + preconditioner: Literal["edm", "time_scaled"] = "edm", + time_scale: float = 999.0, **model_kwargs, -) -> EDMPreconditioner | StormCastUNet: +) -> EDMPreconditioner | TimeScaledModel | StormCastUNet: """ Create a preconditioner-wrapped SongUNet network. @@ -58,8 +148,11 @@ def get_preconditioned_unet( lead_time_channels: the number of channels to use for each lead time embedding amp_mode: whether to use automatic mixed precision use_apex_gn: whether to use Apex GroupNorm + preconditioner: "edm" or "time_scaled", see wrap_diffusion_network + time_scale: factor applied to the times by the "time_scaled" preconditioner Returns: - EDMPreconditioner or StormCastUNet: a wrapped torch module net(x+n, sigma, condition) -> x + EDMPreconditioner, TimeScaledModel or StormCastUNet: a wrapped torch + module net(x_t, t, condition) -> prediction """ if model_type is None: @@ -90,7 +183,7 @@ def get_preconditioned_unet( out_channels=out_channels, **model_params, ) - return EDMPreconditioner(model=ConcatConditionWrapper(unet)) + return wrap_diffusion_network(unet, preconditioner, time_scale) elif name == "regression": return StormCastUNet( @@ -111,11 +204,14 @@ def get_preconditioned_natten_dit( patch_size: int = 4, attn_kernel_size: int = 31, lead_time_steps: int = 0, + attention_backend: Literal["natten2d", "natten2d_rope"] = "natten2d", layernorm_backend: Literal["torch", "apex"] = "torch", conditioning_embedder: Literal["dit", "edm", "zero"] = "dit", attn_kwargs: dict[str, Any] = {}, + preconditioner: Literal["edm", "time_scaled"] = "edm", + time_scale: float = 999.0, **model_kwargs, -) -> EDMPreconditioner: +) -> EDMPreconditioner | TimeScaledModel: """ Create a preconditioner-wrapped Diffusion Transformer (DiT) network. @@ -130,10 +226,19 @@ def get_preconditioned_natten_dit( patch_size: the patch size used by the DiT embedder attn_kernel_size: the attention neighborhood size lead_time_steps: the number of possible lead time steps, if 0 lead time embedding will be disabled - attn_kwargs: parameters to be passed to the DiT attention backend + attention_backend: neighborhood attention without ("natten2d") or with + 2D rotary position embeddings ("natten2d_rope") + layernorm_backend: layer normalization implementation of the blocks + conditioning_embedder: the DiT conditioning embedder + attn_kwargs: parameters to be passed to the DiT attention backend, for + example {"qk_norm": True, "norm_layer": "apex"} to normalize the + attention queries and keys + preconditioner: "edm" or "time_scaled", see wrap_diffusion_network + time_scale: factor applied to the times by the "time_scaled" preconditioner **model_kwargs: any additional parameters to the model Returns: - EDMPreconditioner: a wrapped torch module net(x+n, sigma, condition) -> x + EDMPreconditioner or TimeScaledModel: a wrapped torch module + net(x_t, t, condition) -> prediction """ condition_dim = scalar_condition_channels + lead_time_steps @@ -147,14 +252,14 @@ def get_preconditioned_natten_dit( depth=depth, num_heads=num_heads, patch_size=patch_size, - attention_backend="natten2d", + attention_backend=attention_backend, layernorm_backend=layernorm_backend, attn_kwargs=attn_kw, condition_dim=condition_dim, conditioning_embedder=conditioning_embedder, **model_kwargs, ) - return EDMPreconditioner(model=ConcatConditionWrapper(dit)) + return wrap_diffusion_network(dit, preconditioner, time_scale) def build_network_condition_and_target( @@ -216,9 +321,11 @@ def build_network_condition_and_target( if scalar_conditions is not None: condition = TensorDict( - {"cond_concat": condition, "cond_vec": scalar_conditions} - if condition is not None - else {"cond_vec": scalar_conditions}, + ( + {"cond_concat": condition, "cond_vec": scalar_conditions} + if condition is not None + else {"cond_vec": scalar_conditions} + ), device=target.device, ).to(dtype=target.dtype) @@ -264,6 +371,22 @@ def unpack_batch( return (background, state, mask, lead_time_label, scalar_conditions) +def flow_conversions(scheduler: NoiseScheduler) -> dict[str, Any]: + """Scheduler functions ``f(prediction, x_t, t)`` that convert each + prediction type to the flow (velocity) of the rectified-flow ODE.""" + # The domain-parallel wrapper only implements the sampling and loss + # methods; the element-wise conversions come from the wrapped scheduler + scheduler = getattr(scheduler, "inner_scheduler", scheduler) + return { + "x0": scheduler.x0_to_flow, + "score": scheduler.score_to_flow, + "epsilon": lambda eps, x_t, t: scheduler.x0_to_flow( + scheduler.epsilon_to_x0(eps, x_t, t), x_t, t + ), + "flow": lambda flow, x_t, t: flow, + } + + def diffusion_model_forward( model: Module, condition: torch.Tensor, @@ -274,12 +397,17 @@ def diffusion_model_forward( lead_time_label: torch.Tensor | None = None, sampler_args: dict[str, Any] = {}, invalid_mask: torch.Tensor | None = None, + formulation: Literal["edm", "flow"] = "edm", + prediction_type: Literal["x0", "epsilon", "score", "flow"] = "x0", ) -> torch.Tensor: """Run diffusion model sampling using the ``physicsnemo.diffusion`` API. Uses the provided noise scheduler for timestep generation and :func:`~physicsnemo.diffusion.samplers.sample` for the reverse ODE - integration. + integration. With ``formulation="edm"`` the scheduler's denoiser factory + turns the prediction into the ODE right-hand side. For rectified flow + (``formulation="flow"``) the right-hand side is the flow itself, converted + from the prediction with the scheduler. For domain-parallel inference, pass a scheduler that has already been wrapped with @@ -306,10 +434,20 @@ def diffusion_model_forward( sampler_args : dict Sampler configuration. Supported keys: ``num_steps``, ``solver`` (``"heun"`` or ``"euler"``), - ``S_churn``, ``S_min``, ``S_max``, ``S_noise``. + ``S_churn``, ``S_min``, ``S_max``, ``S_noise`` (stochastic sampling, + see the note in the code on rectified flow), and + ``time_shift`` (rectified flow only): the time shift :math:`a` warps + the sampling times as :math:`t \to a t / (1 + (a - 1) t)` to spend + more steps at high noise levels, which helps when sampling a larger + domain than used in training (https://arxiv.org/abs/2403.03206). invalid_mask : torch.Tensor | None Per-sample invalid-region mask ``(B, 1, H, W)`` forwarded to the model's NaN-mask-token path at every denoising step. ``None`` disables masking. + formulation : {"edm", "flow"} + Noise schedule family of ``scheduler``: ``"edm"`` for + ``training.loss.type: edm``, ``"flow"`` for ``flow_matching``. + prediction_type : {"x0", "epsilon", "score", "flow"} + Quantity predicted by ``model`` (``training.loss.prediction_type``). """ if isinstance(condition, TensorDict): ref_tensor = condition.get("cond_concat", condition.get("cond_vec")) @@ -332,9 +470,12 @@ def diffusion_model_forward( S_min = float(sa.get("S_min", 0)) S_max = float(sa.get("S_max", float("inf"))) S_noise = float(sa.get("S_noise", 1)) + time_shift = float(sa.get("time_shift", 1.0)) B = shape[0] t_steps = scheduler.timesteps(num_steps, device=device, dtype=dtype) + if formulation == "flow": + t_steps = time_shift * t_steps / (1.0 + (time_shift - 1.0) * t_steps) tN = t_steps[0].expand(B) xN = scheduler.init_latents(tuple(shape[1:]), tN, device=device, dtype=dtype) @@ -344,11 +485,27 @@ def diffusion_model_forward( if invalid_mask is not None: extra_kwargs["invalid_mask"] = invalid_mask - def x0_predictor(x: torch.Tensor, t: torch.Tensor) -> torch.Tensor: + def predictor(x: torch.Tensor, t: torch.Tensor) -> torch.Tensor: return model(x, t, condition=condition, **extra_kwargs) - denoiser = scheduler.get_denoiser(x0_predictor=x0_predictor) + if formulation == "edm": + denoiser = scheduler.get_denoiser(**{f"{prediction_type}_predictor": predictor}) + elif formulation == "flow": + # The flow is the right-hand side of the rectified-flow ODE; converting + # through scheduler.get_denoiser would divide by alpha(t) = 1 - t, + # which is zero at the pure-noise start t = 1 + to_flow = flow_conversions(scheduler)[prediction_type] + + def denoiser(x: torch.Tensor, t: torch.Tensor) -> torch.Tensor: + return to_flow(predictor(x, t).float(), x, t) + else: + raise ValueError("formulation must be 'edm' or 'flow'") + # Stochastic sampling (S_churn > 0) works with both formulations. The churn + # of the physicsnemo solvers adds noise without rescaling the signal, which + # is exact for the EDM schedule (alpha = 1) and an approximation for + # rectified flow (alpha = 1 - t); sigma(t) = t for both schedulers, so no + # sigma callbacks are needed. solver_options: dict[str, Any] | None = None if S_churn > 0: solver_name = ( diff --git a/examples/weather/regional_weather_diffusion/utils/trainer.py b/examples/weather/regional_weather_diffusion/utils/trainer.py index a48e2575a2..84798dc3d7 100644 --- a/examples/weather/regional_weather_diffusion/utils/trainer.py +++ b/examples/weather/regional_weather_diffusion/utils/trainer.py @@ -14,7 +14,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Trainer class for StormCast/StormScope training.""" +"""Trainer class for Regional Weather Diffusion training.""" import os import time @@ -29,12 +29,20 @@ from physicsnemo.distributed import DistributedManager from physicsnemo.utils import load_checkpoint, load_model_weights, save_checkpoint -from physicsnemo.diffusion.metrics.losses import WeightedMSEDSMLoss -from physicsnemo.diffusion.noise_schedulers import EDMNoiseScheduler, NoiseScheduler +from physicsnemo.diffusion.metrics.losses import ( + WeightedFlowMatchingLoss, + WeightedMSEDSMLoss, +) +from physicsnemo.diffusion.noise_schedulers import ( + EDMNoiseScheduler, + NoiseScheduler, + RectifiedFlowNoiseScheduler, +) from utils.loss import ( RegressionLoss, SigmaBinTracker, + build_diffusion_loss, build_noise_scheduler, ) @@ -56,7 +64,7 @@ class Trainer: r""" - StormCast Trainer class. + Regional Weather Diffusion Trainer class. Encapsulates all training logic including model and optimizer setup, training and validation loops, checkpointing, logging, and validation plotting. @@ -227,6 +235,9 @@ def _parse_config(self): # Model type self.loss_type = cfg.training.loss.type self.net_name = "regression" if self.loss_type == "regression" else "diffusion" + # Noise schedule family and the quantity the network predicts + self.formulation = "flow" if self.loss_type == "flow_matching" else "edm" + self.prediction_type = cfg.training.loss.prediction_type # Invalid-region masking for the DiT NATTEN backend. When enabled, the # per-sample mask from the dataloader is passed to the model's forward @@ -500,6 +511,8 @@ def _setup_model(self) -> Module: lead_time_steps=self.lead_time_steps, amp_mode=self.enable_amp, use_apex_gn=self.use_apex_gn, + preconditioner=model_cfg.preconditioner, + time_scale=model_cfg.time_scale, **model_cfg.hyperparameters, ) elif model_cfg.architecture == "dit": @@ -509,11 +522,14 @@ def _setup_model(self) -> Module: conditional_channels=num_condition_channels, scalar_condition_channels=len(self.scalar_cond_channels), lead_time_steps=self.lead_time_steps, + preconditioner=model_cfg.preconditioner, + time_scale=model_cfg.time_scale, **model_cfg.hyperparameters, ) # Capture patch_size directly from the DiT before FSDP wrapping - # changes the module hierarchy. EDMPreconditioner.model is the - # ConcatConditionWrapper; .model of that is the DiT. + # changes the module hierarchy. EDMPreconditioner.model (or + # TimeScaledModel.model) is the ConcatConditionWrapper; .model of + # that is the DiT. self._dit_patch_size = net.model.model.patch_size else: raise ValueError("model.architecture must be 'unet' or 'dit'") @@ -654,7 +670,9 @@ def _load_regression_net(self) -> Module | None: # Loss and Optimizer Setup # ========================================================================= - def _setup_loss(self) -> WeightedMSEDSMLoss | RegressionLoss: + def _setup_loss( + self, + ) -> WeightedMSEDSMLoss | WeightedFlowMatchingLoss | RegressionLoss: r""" Create the loss function. @@ -662,7 +680,12 @@ def _setup_loss(self) -> WeightedMSEDSMLoss | RegressionLoss: For diffusion models, creates a :class:`~physicsnemo.diffusion.metrics.losses.WeightedMSEDSMLoss` with an - :class:`~physicsnemo.diffusion.noise_schedulers.EDMNoiseScheduler`. + :class:`~physicsnemo.diffusion.noise_schedulers.EDMNoiseScheduler` + (``training.loss.type: edm``), or a + :class:`~physicsnemo.diffusion.metrics.losses.WeightedFlowMatchingLoss` + with a + :class:`~physicsnemo.diffusion.noise_schedulers.RectifiedFlowNoiseScheduler` + (``training.loss.type: flow_matching``). When domain parallelism is active the scheduler is wrapped via :meth:`~utils.parallel.ParallelHelper.make_domain_parallel_scheduler` so that sampled sigmas are broadcast across spatial shards. @@ -676,7 +699,7 @@ def _setup_loss(self) -> WeightedMSEDSMLoss | RegressionLoss: Returns ------- - WeightedMSEDSMLoss | RegressionLoss + WeightedMSEDSMLoss | WeightedFlowMatchingLoss | RegressionLoss The loss function (possibly ``torch.compile``-wrapped). """ self.logger.info("Setting up loss function...") @@ -698,9 +721,13 @@ def _setup_loss(self) -> WeightedMSEDSMLoss | RegressionLoss: return loss_fn loss_params = self.cfg.training.loss - self.logger.info( - f"Using modern diffusion loss: {loss_params.sigma_distribution}" - ) + if self.loss_type == "flow_matching": + self.logger.info( + f"Using flow matching loss with {loss_params.prediction_type} " + f"prediction and t in [{loss_params.t_min}, {loss_params.t_max}]" + ) + else: + self.logger.info(f"Using diffusion loss: {loss_params.sigma_distribution}") noise_scheduler = build_noise_scheduler( loss_params, @@ -709,15 +736,19 @@ def _setup_loss(self) -> WeightedMSEDSMLoss | RegressionLoss: noise_scheduler = self.parallel_helper.make_domain_parallel_scheduler( noise_scheduler, ) - loss_fn = WeightedMSEDSMLoss(self.net, noise_scheduler, reduction="none") - self.train_noise_scheduler = loss_fn.noise_scheduler + loss_fn = build_diffusion_loss(self.net, noise_scheduler, loss_params) + self.train_noise_scheduler = noise_scheduler sa = self.cfg.sampler.args.__dict__ - sampling_scheduler = EDMNoiseScheduler( - sigma_min=sa.get("sigma_min", 0.002), - sigma_max=sa.get("sigma_max", 80.0), - rho=sa.get("rho", 7.0), - ) + if self.loss_type == "flow_matching": + # Sampling starts at the highest training time, pure noise for t_max = 1 + sampling_scheduler = RectifiedFlowNoiseScheduler(t_max=loss_params.t_max) + else: + sampling_scheduler = EDMNoiseScheduler( + sigma_min=sa.get("sigma_min", 0.002), + sigma_max=sa.get("sigma_max", 80.0), + rho=sa.get("rho", 7.0), + ) self.sampling_scheduler = self.parallel_helper.make_domain_parallel_scheduler( sampling_scheduler ) @@ -1089,6 +1120,8 @@ def _get_plot_outputs( sampler_args=self.cfg.sampler.args.__dict__, lead_time_label=lead_time_label, invalid_mask=invalid_mask if self.use_nan_mask_tokens else None, + formulation=self.formulation, + prediction_type=self.prediction_type, ) if "regression" in self.condition_list: outputs += reg_out