DoMINO Performance Optimizations

CHANGELOG.md

@@ -34,6 +34,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Refactored DiTBlock to be more modular
 - Added NATTEN 2D neighborhood attention backend for DiTBlock
 - Migrated blood flow example to PyTorch Geometric.
+- Refactored DoMINO model code and examples for performance optimizations and improved readability.
 - Migrated HydroGraphNet example to PyTorch Geometric.
 - Support for saving and loading nested `physicsnemo.Module`s. It is now
   possible to create nested modules with `m = Module(submodule, ...)`, and save

examples/cfd/external_aerodynamics/domino/README.md

@@ -77,19 +77,24 @@ please refer to their [paper](https://arxiv.org/pdf/2408.11969).
 
 #### Data Preprocessing
 
-`PhysicsNeMo` has a related project to help with data processing, called [PhysicsNeMo-Curator](https://github.com/NVIDIA/physicsnemo-curator).
+`PhysicsNeMo` has a related project to help with data processing, called
+[PhysicsNeMo-Curator](https://github.com/NVIDIA/physicsnemo-curator).
 Using `PhysicsNeMo-Curator`, the data needed to train a DoMINO model can be setup easily.
-Please refer to [these instructions on getting started](https://github.com/NVIDIA/physicsnemo-curator?tab=readme-ov-file#what-is-physicsnemo-curator)
+Please refer to
+[these instructions on getting started](https://github.com/NVIDIA/physicsnemo-curator?tab=readme-ov-file#what-is-physicsnemo-curator)
 with `PhysicsNeMo-Curator`.
 
-Download the DrivAer ML dataset using the [provided instructions in PhysicsNeMo-Curator](https://github.com/NVIDIA/physicsnemo-curator/blob/main/examples/external_aerodynamics/domino/README.md#download-drivaerml-dataset).
+Download the DrivAer ML dataset using the
+[provided instructions in PhysicsNeMo-Curator](https://github.com/NVIDIA/physicsnemo-curator/blob/main/examples/external_aerodynamics/domino/README.md#download-drivaerml-dataset).
 The first step for running the DoMINO pipeline requires processing the raw data
 (vtp, vtu and stl) into either Zarr or NumPy format for training.
 Each of the raw simulations files are downloaded in `vtp`, `vtu` and `stl` formats.
 For instructions on running data processing to produce a DoMINO training ready dataset,
-please refer to [How-to Curate data for DoMINO Model](https://github.com/NVIDIA/physicsnemo-curator/blob/main/examples/external_aerodynamics/domino/README.md).
+please refer to
+[How-to Curate data for DoMINO Model](https://github.com/NVIDIA/physicsnemo-curator/blob/main/examples/external_aerodynamics/domino/README.md).
 
-Caching is implemented in [`CachedDoMINODataset`](https://github.com/NVIDIA/physicsnemo/blob/main/physicsnemo/datapipes/cae/domino_datapipe.py#L1250).
+Caching is implemented in
+[`CachedDoMINODataset`](https://github.com/NVIDIA/physicsnemo/blob/main/physicsnemo/datapipes/cae/domino_datapipe.py#L1250).
 Optionally, users can run `cache_data.py` to save outputs
 of DoMINO datapipe in the `.npy` files. The DoMINO datapipe is set up to calculate
 Signed Distance Field and Nearest Neighbor interpolations on-the-fly during
@@ -101,6 +106,36 @@ processed files.
 The final processed dataset should be divided and saved into 2 directories,
 for training and validation.
 
+#### Data Scaling factors
+
+DoMINO has several data-specific configuration tools that rely on some
+knowledge of the dataset:
+
+- The output fields (the labels) are normalized during training to a mean
+  of zero and a standard deviation of one, averaged over the dataset.
+  The scaling is controlled by passing the `volume_factors` and
+  `surface_factors` values to the datapipe.
+- The input locations are scaled by, and optionally cropped to, used defined
+  bounding boxes for both surface and volume.  Whether cropping occurs, or not,
+  is controlled by the `sample_in_bbox` value of the datapipe.  Normalization
+  to the bounding box is enabled with `normalize_coordinates`.  By default,
+  both are set to true.  The value of the boxes are configured in the
+  `config.yaml` file, and are configured separately for surface and volume.
+
+> Note: The datapipe module has a helper function `create_domino_dataset`
+> with sensible defaults to help create a Domino Datapipe.
+
+To facilitate setting reasonable values of these, you can use the
+`compute_statistics.py` script.  This will load the core dataset as defined
+in your `config.yaml` file, loop over several events (200, by default), and
+both print and store the surface/volume field statistics as well as the
+coordinate statistics.  
+
+> Note that, for volumetric fields especially, the min/max found may be
+> significantly outside the surface region.  Many simulations extend volumetric
+> sampling to far field, and you may instead want to crop significant amounts
+> of volumetric distance.
+
 #### Training
 
 Specify the training and validation data paths, bounding box sizes etc. in the
@@ -176,9 +211,6 @@ The `domain_size` represents the number of GPUs used for each batch - setting
 but with extra overhead.  `shard_grid` and `shard_points` will enable domain
 parallelism over the latent space and input/output points, respectively.
 
-Please see `src/train_sharded.py` for more details regarding the changes
-from the standard training script required for domain parallel DoMINO training.
-
 As one last note regarding domain-parallel training: in the phase of the DoMINO
 where the output solutions are calculated, the model can used two different
 techniques (numerically identical) to calculate the output.  Due to the
@@ -189,6 +221,114 @@ launch overhead at the cost of more memory use.  For non-sharded
 training, the `two-loop` setting is more optimal. The difference in `one-loop`
 or `two-loop` is purely computational, not algorithmic.
 
+### Performance Optimizations
+
+The training and inference scripts for DoMINO contain several performance
+enhancements to accelerate the training and usage of the model. In this
+section we'll highlight several of them, as well as how to customize them
+if needed.
+
+#### Memory Pool Optimizations
+
+The preprocessor of DoMINO requires a computation of k Nearest Neighbors,
+which is accelerated via the `cuml` Neighbors tool.  By default, `cuml` and
+`torch` both use memory allocation pools to speed up allocating tensors, but
+they do not use the same pool.  This means that during preprocessing, it's
+possible for the kNN operation to spend a significant amount of time in
+memory allocations - and further, it limits the available memory to `torch`.
+
+To mitigate this, by default in DoMINO we use the Rapids Memory Manager
+([`rmm`](https://github.com/rapidsai/rmm)).  If, for some reason, you wish
+to disable this you can do so with an environment variable:
+
+```bash
+export PHYSICSNEMO_DISABLE_RMM=True
+```
+
+Or remove this line from the training script:
+
+```python
+from physicsnemo.utils.memory import unified_gpu_memory
+```
+
+> Note - why not make it configurable?  We have to set up the shared memory
+> pool allocation very early in the program, before the config has even
+> been read.  So, we enable by default and the opt-out path is via the
+> environment.
+
+#### Reduced Volume Reads
+
+The dataset size for volumetric data can be quite substantial - DrivAerML, for
+example, has mesh sizes of 160M points per example.  Even though the models
+do not process all 160M points, in order to down sample dynamically they all
+must be read from disk - which can exceed bandwidth and CPU decoding capacity
+on nodes with multiple GPUs.
+
+As a performance enhancement, DoMINO's data pipeline offers a mitigation: instead
+of reading an entire volumetric mesh, during preprocessing we _shuffle_ the
+volumetric inputs and outputs (in tandem) and subsequent reads choose random
+slices of the volumetric data.  By default, DoMINO will read about 100x more data
+than necessary for the sampling size.  This allows the pipeline to still apply
+cuts for data inside of the bounding box, and further random sampling to improve
+training stability.  To enable/disable this parameter, set
+`data.volume_sample_from_disk=True` (enable) or `False` (disable)
+
+> Note - if you volumetric data is not larger than a few million mesh points,
+> pre-shuffling and sampling from disk is likely not necessary for you.
+
+`physicsnemo-curator` supports shuffling the volumetric data during preprocessing.
+If, however, you've already preprocessed your data and just want to apply
+shuffling, use the script at `src/shuffle_volumetric_curator_output.py`
+
+The shuffling script will also apply sharding to the output files, which
+improves IO performance.  So, `zarr>=3.0` is required to use the outputs from
+curator.  `src/shuffle_volumetric_curator_output.py` is meant to be an example of how
+to apply shuffling, so modify and update as you need for your dataset.
+
+> If you have tensorstore installed (it's in `requirements.txt`), the data reader
+> will work equally well with Zarr 2 or Zarr 3 files.
+
+#### Overall Performance
+
+DoMINO is a computationally complex and challenging workload.  Over the course
+of several releases, we have chipped away at performance bottlenecks to speed
+up the training and inference time (with `inference_on_stl.py`).  Overall
+training performance has decreased from about 5 days to just over 4 hours, with
+eight H100 GPUs.  We hope these optimizations enable you to explore more
+parameters and surrogate models; if there is a performance issue you see,
+please open an issue on GitHub.
+
+![Results from DoMINO for RTWT SC demo](../../../../docs/img/domino_perf.png)
+
+### Example Training Results
+
+To provide an example of what a successful training should look like, we include here
+some example results.  Training curves may look similar to this:
+
+![Combined Training Curve](../../../../docs/img/domino/combined-training-curve.png)
+
+And, when evaluating the results on the validation dataset, this particular
+run had the following L2 and R2 Metrics:
+
+|              Metric | Surface Only | Combined |
+|--------------------:|:------------:|:--------:|
+|          X Velocity |      N/A     |   0.086  |
+|          Y Velocity |      N/A     |   0.185  |
+|          Z Velocity |      N/A     |   0.197  |
+| Volumetric Pressure |      N/A     |   0.106  |
+|             Turb. V |      N/A     |   0.134  |
+|    Surface Pressure |     0.101    |   0.105  |
+|       X-Tau (Shear) |     0.138    |   0.145  |
+|       Y-Tau (Shear) |     0.174    |   0.185  |
+|       Z-Tau (Shear) |     0.198    |   0.207  |
+|             Drag R2 |     0.983    |   0.975  |
+|             Lift R2 |     0.971    |   0.968  |
+
+With the PhysicsNeMo CFD tool, you can create plots of the lift and drag
+forces computed by domino vs. the CFD Solver.  For example, here is the drag force:
+
+![Draf Force R^2](../../../../docs/img/domino/drag-r2.jpg)
+
 ### Training with Physics Losses
 
 DoMINO supports enforcing of PDE residuals as soft constraints. This can be used

examples/cfd/external_aerodynamics/domino/requirements.txt

@@ -3,3 +3,4 @@ warp-lang
 tensorboard
 cuml
 einops
+tensorstore

examples/cfd/external_aerodynamics/domino/src/benchmark_dataloader.py

@@ -0,0 +1,184 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 - 2024 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.
+
+"""
+This code defines a distributed pipeline for training the DoMINO model on
+CFD datasets. It includes the computation of scaling factors, instantiating
+the DoMINO model and datapipe, automatically loading the most recent checkpoint,
+training the model in parallel using DistributedDataParallel across multiple
+GPUs, calculating the loss and updating model parameters using mixed precision.
+This is a common recipe that enables training of combined models for surface and
+volume as well either of them separately. Validation is also conducted every epoch,
+where predictions are compared against ground truth values. The code logs training
+and validation metrics to TensorBoard. The train tab in config.yaml can be used to
+specify batch size, number of epochs and other training parameters.
+"""
+
+import time
+import os
+import re
+import torch
+import torchinfo
+
+from typing import Literal, Any
+
+
+import hydra
+from hydra.utils import to_absolute_path
+from omegaconf import DictConfig, OmegaConf
+
+# This will set up the cupy-ecosystem and pytorch to share memory pools
+from physicsnemo.utils.memory import unified_gpu_memory
+
+
+import torch.distributed as dist
+from torch.cuda.amp import GradScaler, autocast
+from torch.nn.parallel import DistributedDataParallel
+from torch.utils.data import DataLoader
+from torch.utils.data.distributed import DistributedSampler
+from torch.utils.tensorboard import SummaryWriter
+from nvtx import annotate as nvtx_annotate
+import torch.cuda.nvtx as nvtx
+
+
+from physicsnemo.distributed import DistributedManager
+from physicsnemo.launch.utils import load_checkpoint, save_checkpoint
+from physicsnemo.launch.logging import PythonLogger, RankZeroLoggingWrapper
+
+from physicsnemo.datapipes.cae.domino_datapipe import (
+    DoMINODataPipe,
+    compute_scaling_factors,
+    create_domino_dataset,
+)
+from physicsnemo.models.domino.model import DoMINO
+from physicsnemo.utils.domino.utils import *
+
+# This is included for GPU memory tracking:
+from pynvml import nvmlInit, nvmlDeviceGetHandleByIndex, nvmlDeviceGetMemoryInfo
+import time
+
+from utils import (
+    ScalingFactors,
+    get_keys_to_read,
+    coordinate_distributed_environment,
+    load_scaling_factors,
+)
+
+
+from physicsnemo.utils.profiling import profile, Profiler
+
+
+def benchmark_io_epoch(
+    dataloader,
+    logger,
+    gpu_handle,
+    epoch_index,
+    device,
+):
+    dist = DistributedManager()
+
+    # If you tell the dataloader the indices in advance, it will preload
+    # and pre-preprocess data
+    # dataloader.set_indices(indices)
+
+    gpu_start_info = nvmlDeviceGetMemoryInfo(gpu_handle)
+    start_time = time.perf_counter()
+    for i_batch, sample_batched in enumerate(dataloader):
+        # Gather data and report
+        elapsed_time = time.perf_counter() - start_time
+        start_time = time.perf_counter()
+        gpu_end_info = nvmlDeviceGetMemoryInfo(gpu_handle)
+        gpu_memory_used = gpu_end_info.used / (1024**3)
+        gpu_memory_delta = (gpu_end_info.used - gpu_start_info.used) / (1024**3)
+
+        logging_string = f"Device {device}, batch processed: {i_batch + 1}\n"
+        logging_string += f"  GPU memory used: {gpu_memory_used:.3f} Gb\n"
+        logging_string += f"  GPU memory delta: {gpu_memory_delta:.3f} Gb\n"
+        logging_string += f"  Time taken: {elapsed_time:.2f} seconds\n"
+        logger.info(logging_string)
+        gpu_start_info = nvmlDeviceGetMemoryInfo(gpu_handle)
+
+    return
+
+
+@hydra.main(version_base="1.3", config_path="conf", config_name="config")
+def main(cfg: DictConfig) -> None:
+    # initialize distributed manager
+    DistributedManager.initialize()
+    dist = DistributedManager()
+
+    # Initialize NVML
+    nvmlInit()
+
+    gpu_handle = nvmlDeviceGetHandleByIndex(dist.device.index)
+
+    model_type = cfg.model.model_type
+
+    logger = PythonLogger("Train")
+    logger = RankZeroLoggingWrapper(logger, dist)
+
+    logger.info(f"Config summary:\n{OmegaConf.to_yaml(cfg, sort_keys=True)}")
+
+    ################################
+    # Get scaling factors
+    ################################
+    vol_factors, surf_factors = load_scaling_factors(cfg)
+
+    keys_to_read, keys_to_read_if_available = get_keys_to_read(
+        cfg, model_type, get_ground_truth=True
+    )
+
+    domain_mesh, data_mesh, placements = coordinate_distributed_environment(cfg)
+
+    train_dataset = create_domino_dataset(
+        cfg,
+        phase="train",
+        keys_to_read=keys_to_read,
+        keys_to_read_if_available=keys_to_read_if_available,
+        vol_factors=vol_factors,
+        surf_factors=surf_factors,
+        device_mesh=domain_mesh,
+        placements=placements,
+    )
+    train_sampler = DistributedSampler(
+        train_dataset, num_replicas=data_mesh.size(), rank=data_mesh.get_local_rank()
+    )
+
+    for epoch in range(0, cfg.train.epochs):
+        start_time = time.perf_counter()
+        logger.info(f"Device {dist.device}, epoch {epoch}:")
+
+        train_sampler.set_epoch(epoch)
+
+        train_dataset.dataset.set_indices(list(train_sampler))
+
+        epoch_start_time = time.perf_counter()
+        with Profiler():
+            benchmark_io_epoch(
+                dataloader=train_dataset,
+                logger=logger,
+                gpu_handle=gpu_handle,
+                epoch_index=epoch,
+                device=dist.device,
+            )
+        epoch_end_time = time.perf_counter()
+        logger.info(
+            f"Device {dist.device}, Epoch {epoch} took {epoch_end_time - epoch_start_time:.3f} seconds"
+        )
+
+
+if __name__ == "__main__":
+    main()