Improve StormScope recipe masking ergonomics

[pzharrington]

PhysicsNeMo Pull Request

Description

Summary

Reworks spatial masking in the StormCast/StormScope recipe into a single, dataloader-owned interface, and fixes a domain-parallelism bug in the loss-mask construction. Previously masking came from two reconciled-at-train-time sources (a dataloader mask key and a NaN-inspection pathway), which was confusing and produced a shape bug under domain parallelism. Changes:

• Single mask source (dataloader-owned). The per-sample "mask" served by the dataset (convention 1=valid, 0=invalid) is now the only mask source. The NaN-inspection pathway is removed entirely. Trainer._build_masks(target, batch_mask) translates that mask into a loss weight and an invalid_mask; use_nan_mask_tokens now solely controls whether invalid_mask is also injected into the DiT forward so it can swap invalid tokens for its learned mask token.
• Token-granularity loss weight for DiT. When the DiT runs with use_nan_mask_tokens, the loss weight is pooled to patch granularity (max_pool2d) and expanded back (interpolate), matching the model's token-level substitution. The patch size is read directly from the constructed DiT (patch_size), and the expansion uses a scalar scale_factor. Pixel-level weight is retained otherwise.
• Domain-parallelism fix. Eliminates a mask shape bug that arose from detecting NaNs under sharding; the mask now stays correctly (B, 1, H, W) as a height-sharded ShardTensor through pooling and expansion.
• Config / docs. Adds use_mask: false (mock dataset) and use_nan_mask_tokens: false (stormscope hyperparameters) defaults so the keys exist in the OmegaConf struct, documents the "mask" key in the dataset docstring and README, and adds an optional use_mask pattern to the mock dataset.

Tests

test_masking now exercises the full masked DiT path across three launch configs — single (1 GPU), domain_parallel (2 GPU, height-sharded), and data_domain_parallel (4 GPU, 2×2 data+domain mesh) — skipping configs that don't match the launched world size.

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• The documentation is up to date with these changes.
• The CHANGELOG.md is up to date with these changes.
• An issue is linked to this pull request.
• If I am implementing a new model or modifying any existing model, I have followed the Models Implementation Coding Standards.

Dependencies

Review Process

All PRs are reviewed by the PhysicsNeMo team before merging.

Depending on which files are changed, GitHub may automatically assign a maintainer for review.

We are also testing AI-based code review tools (e.g., Greptile), which may add automated comments with a confidence score.
This score reflects the AI’s assessment of merge readiness and is not a qualitative judgment of your work, nor is
it an indication that the PR will be accepted / rejected.

AI-generated feedback should be reviewed critically for usefulness.
You are not required to respond to every AI comment, but they are intended to help both authors and reviewers.
Please react to Greptile comments with 👍 or 👎 to provide feedback on their accuracy.

[greptile-apps]

Greptile Summary

This PR consolidates StormScope's spatial masking into a single dataloader-owned interface and fixes a domain-parallelism shape bug. The old dual-source masking (dataset "mask" key + NaN-inspection pathway) is replaced by one canonical path: the dataset provides a (B, 1, H, W) float mask (1=valid), and Trainer._build_masks translates it into both a loss weight and an invalid_mask for the DiT forward. For the DiT with use_nan_mask_tokens, the loss weight is correctly derived at token granularity via max_pool2d + interpolate, matching the model's per-token mask-token substitution.

• Mask unification: NaN detection and sanitization are removed; the dataset now owns the full definition of invalid regions, with use_nan_mask_tokens controlling only whether the mask is also injected into the DiT as invalid_mask.
• Token-granularity loss weight: _build_masks pools the mask to patch level with max_pool2d and expands it back with interpolate(scale_factor=patch_h, mode="nearest") — scalar scale_factor is required for ShardTensor compatibility under domain parallelism.
• Test coverage: test_masking exercises all three launch configs (1 GPU, 2-GPU domain parallel, 4-GPU data+domain parallel) with the mock dataset's right-half valid mask and use_nan_mask_tokens=True.

Important Files Changed

Filename	Overview
examples/weather/stormcast/utils/trainer.py	Core change: replaces NaN-detection masking with a single dataloader-owned mask; introduces _build_masks with token-granularity loss weight for DiT; squareness check fires after max_pool2d rather than at initialisation.
examples/weather/stormcast/datasets/dataset.py	Documents the new mask key in the __getitem__ contract; missing blank-line separator causes the added paragraph to visually merge with the following prose in docstring renderers.
examples/weather/stormcast/datasets/mock.py	Adds optional use_mask flag; mock mask (right-half valid) is correct and simple.
examples/weather/stormcast/test_training.py	Adds test_masking covering single/domain-parallel/data+domain-parallel configs; hardcodes checkpoint filename EDMPreconditioner.0.10.mdlus which would silently miss a checkpoint if the step count in the fixture changes.
examples/weather/stormcast/config/dataset/mock.yaml	Adds use_mask: false default so the OmegaConf struct key exists; straightforward change.
examples/weather/stormcast/config/model/stormscope.yaml	Adds use_nan_mask_tokens: false default with an explanatory comment; straightforward change.
examples/weather/stormcast/README.md	Adds a Spatial mask section documenting the mask key contract, token-level substitution, and caching recommendation; clear and accurate.

_{Reviews (1): Last reviewed commit: "Improve recipe masking ergonomics" | Re-trigger Greptile}

[jleinonen]

With the channel masking added, looks good!